gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(cuda): compute_expected_q stride 13 vs denoise_target_q_buf size 12 — OOB writes threads 60-63

Browse Source 
Pre-existing latent bug surfaced by compute-sanitizer after the SP15
NULL-pointer fix at 2e37af29d unblocked the cascade. Threads 60-63
were writing past denoise_target_q_buf's end every batch.

Sanitizer evidence (pre-fix on 2e37af29d):
  Invalid __global__ write of size 4 bytes
    at compute_expected_q+0x2480
    by thread (60..63, 0, 0) in block (0, 0, 0)
    Access at <addr> is out of bounds (45/97/149/201 bytes after
    nearest allocation of size B*12*4 bytes)
  Host backtrace: GpuDqnTrainer::compute_denoise_target_q
    → submit_post_aux_ops → run_full_step

Root cause — semantic mismatch between writer stride and buffer size:
  - compute_expected_q writes q_values[i*total_actions + a] with
    total_actions = b0+b1+b2+b3 = 4+3+3+3 = 13 (4-direction factored)
  - denoise_target_q_buf was allocated b*12 (legacy 3+3+3+3 layout)
  - threads 60..63 wrote slot 12 of samples (60..63 % batch_size)
    past the buffer end every step

The downstream q_denoise_backward + denoise_loss_grad kernels read
Q_target[b * D + i] with hardcoded D=12 because the diffusion denoiser
MLP itself only has 12 output slots (W2[12,24] + b2[12]) — its 12 are
the q_coord_buf's post-cross-branch-attention narrowed output, NOT a
clean subset of the 13 raw Q-actions.

The exact same fix pattern was already applied to the sibling
q_var_buf_trainer allocation in the prior SP4 audit (see comment block
at gpu_dqn_trainer.rs:21620-21630 referencing the identical OOB at
threads 60..63); denoise_target_q_buf 8 lines below was missed because
it was guarded by the SP15 NULL-pointer ILLEGAL_ADDRESS that fault-
stopped the cascade before this OOB could fire — 2e37af29d removed the
upstream ILLEGAL_ADDRESS, surfacing the latent OOB.

Fix architecture (Option A — pad buffer to total_actions, pass stride
to consumer; same pattern as q_var_buf_trainer):
  1. Widen denoise_target_q_buf from b*12 to b*total_actions (= b*13)
     to match compute_expected_q's writer stride.
  2. Add refined_stride / target_stride / input_stride parameters to
     q_denoise_backward kernel; the kernel still computes D=12 per-
     sample (denoiser MLP fixed width) but addresses each input buffer
     at its own per-sample stride.
  3. Add refined_stride / target_stride parameters to denoise_loss_grad
     kernel (same pattern; used by launch_q_denoise_backward_cublas).
  4. Update both Rust launchers (launch_q_denoise_backward,
     launch_q_denoise_backward_cublas) to pass refined_stride=12 (q_coord
     and q_input are post-attention narrowed),
     target_stride=total_actions=13.
  5. denoise_q_input_buf STAYS at b*12 — it's a snapshot of q_coord_buf
     (also b*12) via snapshot_pre_denoise_q's DtoD copy; never written
     by compute_expected_q.
  6. Flat-scan consumers (SP4 target_q_p99_update producer + SP3 slot
     46 threshold-check + dqn_clamp_finite_f32) UNCHANGED — they
     consume the buffer flat via .len(); widening from b*12 to b*13 is
     monotone (one more valid Q-value per sample in the histogram).

Atomic per feedback_no_partial_refactor: 2 kernel signatures + buffer
allocation + 2 launch sites + struct doc comments + SP3 slot 46 doc +
audit doc — all in one commit. Every consumer of the writer's stride
migrates simultaneously.

Verification:
  - SQLX_OFFLINE=true cargo check -p ml --features cuda: clean
  - compute-sanitizer (RTX 3050 Ti): 0 Invalid __global__ errors at
    compute_expected_q post-fix (down from 4 per training step in
    baseline — re-verified on 2e37af29d to confirm the diff)
  - SQLX_OFFLINE=true cargo test -p ml --features cuda --lib: holds
    parent baseline 947 pass / 12 fail (same 12 pre-existing failures)

Files touched:
  crates/ml/src/cuda_pipeline/experience_kernels.cu
    — q_denoise_backward + denoise_loss_grad kernel signatures
  crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
    — alloc widen + 2 launcher updates + 4 doc comment updates
  docs/dqn-wire-up-audit.md
    — audit entry

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 11:39:34 +02:00
parent 2e37af29d4
commit 23e9a1f78c
3 changed files with 95 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -6245,8 +6245,15 @@ extern "C" __global__ void denoise_silu_bwd(
/**
* denoise_loss_grad — MSE gradient with 0.1x residual scaling.
* Reads row-major [B, D], writes col-major [D, B].
* d_res[d, b] = 0.1 * 2.0 * (Q_refined[b*D+d] - Q_target[b*D+d]) * inv_B.
* Reads row-major [B, refined_stride] and [B, target_stride], writes col-major [D, B].
* d_res[d, b] = 0.1 * 2.0 * (Q_refined[b*refined_stride+d] - Q_target[b*target_stride+d]) * inv_B.
*
* refined_stride / target_stride MAY exceed D when the upstream buffers are
* sized per total Q-actions (e.g. 13 for the 4+3+3+3 factored layout) instead
* of the legacy D=12 exposure layout. The denoiser MSE only consumes the
* first D entries per sample; the trailing slot(s) are ignored. Per-sample
* stride must be passed so addressing matches the buffer layout — otherwise
* sample b reads slot 12 of the previous sample.
*
* Col-major [D, B] output: element (d, b) at offset d + b * D.
* Grid: ceil(D*B/256), Block: 256.
@@ -6255,14 +6262,16 @@ extern "C" __global__ void denoise_loss_grad(
const float* __restrict__ Q_refined,
const float* __restrict__ Q_target,
float* __restrict__ d_res,
int B, int D, float inv_B
int B, int D,
int refined_stride, int target_stride,
float inv_B
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = D * B;
if (idx >= total) return;
int d = idx % D;
int b = idx / D;
d_res[idx] = 0.2f * (Q_refined[b * D + d] - Q_target[b * D + d]) * inv_B;
d_res[idx] = 0.2f * (Q_refined[b * refined_stride + d] - Q_target[b * target_stride + d]) * inv_B;
}
/**
@@ -6344,7 +6353,16 @@ extern "C" __global__ void denoise_bias_grad_p2(
*/
/* var_q stride matches the variance buffer's per-sample width (>= D). The
* denoiser only uses the first D entries per sample; pass var_stride from the
* Rust launch site to keep alignment when q_var is sized per total Q-actions. */
* Rust launch site to keep alignment when q_var is sized per total Q-actions.
*
* refined_stride / target_stride / input_stride mirror this pattern for the
* Q buffers themselves: each may be >= D when the upstream buffer is sized
* per total Q-actions (13 for 4+3+3+3 factored layout). The denoiser MSE
* + replay only consumes the first D entries per sample. Compute-sanitizer
* caught the legacy hard-coded `b * D` indexing as 4 OOB writes per batch
* (threads 60..63) when compute_expected_q wrote 13 floats per sample into
* a `b*12` target buffer (commit log: see fix(cuda) compute_expected_q
* stride 13 vs denoise_target_q_buf size 12). */
extern "C" __global__ void q_denoise_backward(
const float* __restrict__ Q_refined,
const float* __restrict__ Q_target,
@@ -6353,7 +6371,10 @@ extern "C" __global__ void q_denoise_backward(
const float* __restrict__ Q_input,
const float* __restrict__ var_q,
int B,
int var_stride)
int var_stride,
int refined_stride,
int target_stride,
int input_stride)
{
int widx = blockIdx.x * blockDim.x + threadIdx.x;
const int D = 12;
@@ -6385,7 +6406,7 @@ extern "C" __global__ void q_denoise_backward(
/* Compute loss gradient for this sample */
float d_out_b[12];
for (int i = 0; i < D; i++) {
d_out_b[i] = 2.0f * (Q_refined[b * D + i] - Q_target[b * D + i]) * inv_B;
d_out_b[i] = 2.0f * (Q_refined[b * refined_stride + i] - Q_target[b * target_stride + i]) * inv_B;
}
/* Process steps in reverse. The gradient for the current step is needed. */
@@ -6395,7 +6416,7 @@ extern "C" __global__ void q_denoise_backward(
/* Reconstruct Q at start of `step` by running forward from Q_input
* through steps 0..(step-1). */
float q_curr[12];
for (int i = 0; i < D; i++) q_curr[i] = Q_input[b * D + i];
for (int i = 0; i < D; i++) q_curr[i] = Q_input[b * input_stride + i];
for (int s = 0; s < step; s++) {
int soff = s * STEP_PARAMS;

78
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -6203,7 +6203,7 @@ pub struct GpuDqnTrainer {
// ── SP4 Layer A Task A5: target_q_p99 producer ─────────────────────
/// SP4 first end-to-end producer kernel. Single-block kernel that reads
/// `denoise_target_q_buf` (target-network Q-values [B, 12] flattened),
/// `denoise_target_q_buf` (target-network Q-values [B, total_actions] flattened),
/// computes p99(|target_q|) via the `sp4_histogram_p99<256>` header-only
/// device function, writes the step observation to
/// `producer_step_scratch_buf[0]`. Loaded from
@@ -7001,7 +7001,11 @@ pub struct GpuDqnTrainer {
q_denoise_backward_kernel: CudaFunction,
/// Gradient accumulator for denoise_params [1800]. Zeroed before each backward.
denoise_grad: CudaSlice<f32>,
/// Target network Q-values [B, 12] for denoiser training signal.
/// Target network Q-values [B, total_actions] for denoiser training signal.
/// Sized at `total_actions = b0+b1+b2+b3` (= 13 for the 4-direction
/// factored layout) to match the `compute_expected_q` writer stride.
/// Denoiser consumers (q_denoise_backward, denoise_loss_grad) read only
/// the first D=12 slots per sample via the `target_stride` parameter.
denoise_target_q_buf: CudaSlice<f32>,
/// Pre-denoise Q-values snapshot [B, 12] — saved before denoiser forward for backward.
denoise_q_input_buf: CudaSlice<f32>,
@@ -11020,15 +11024,23 @@ impl GpuDqnTrainer {
(end - start) / std::mem::size_of::<f32>()
}
/// SP3 slot 46: device pointer to `denoise_target_q_buf` `[B, 12]`,
/// the target-network Q-values written by the C51 target-q compute pass.
/// The threshold-check kernel scans this for post-clip outliers (Task B2
/// extends the clipping pipeline to feed this buffer).
/// SP3 slot 46: device pointer to `denoise_target_q_buf` `[B, total_actions]`,
/// the target-network Q-values written by the C51 target-q compute pass via
/// `compute_expected_q` (which strides by `total_actions = b0+b1+b2+b3 = 13`
/// in the 4-direction factored layout). The threshold-check kernel scans
/// the full buffer flat — all entries are valid Q-values, so widening from
/// the legacy `b*12` exposure-layout sizing to `b*total_actions` is
/// monotone (more samples in the histogram, no semantic change). Per-sample
/// consumers (q_denoise_backward, denoise_loss_grad) receive
/// `target_stride = total_actions` and read only the first D=12 entries —
/// matching the cross-branch-attention narrowed denoiser pipeline.
pub(crate) fn target_q_ptr(&self) -> u64 {
self.denoise_target_q_buf.raw_ptr()
}
/// Element count of the target-q buffer (= `B * 12`).
/// Element count of the target-q buffer (= `B * total_actions`, e.g. `B * 13`
/// for the 4-direction factored layout). Used by SP4 p99 producer +
/// SP3 threshold-check kernels for flat scans of the full buffer.
pub(crate) fn target_q_len(&self) -> usize {
self.denoise_target_q_buf.len()
}
@@ -12777,7 +12789,7 @@ impl GpuDqnTrainer {
/// SP4 Layer A Task A5: producer for ISV[TARGET_Q_BOUND_INDEX=131].
///
/// Reads `denoise_target_q_buf` (target-network Q-values [B, 12]), launches
/// Reads `denoise_target_q_buf` (target-network Q-values [B, total_actions]), launches
/// the single-block `target_q_p99_update` kernel which computes
/// p99(|target_q|) via the header-only `sp4_histogram_p99<256>` device
/// function and writes the step observation to
@@ -21670,7 +21682,21 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("q_denoise_backward load: {e}")))?;
let denoise_grad = stream.alloc_zeros::<f32>(DENOISE_TOTAL_PARAMS)
.map_err(|e| MLError::ModelError(format!("alloc denoise_grad: {e}")))?;
let denoise_target_q_buf = alloc_f32(&stream, b * 12, "denoise_target_q_buf")?;
// denoise_target_q_buf is written by `compute_expected_q` which strides
// by `total_actions` (= 4+3+3+3 = 13 in the 4-direction factored layout).
// The legacy `b*12` sizing was the 12-slot exposure layout (3+3+3+3) and
// produced 4 OOB writes per batch (threads 60..63) under compute-sanitizer.
// Sized at `b*total_actions` to match the kernel writer; downstream
// denoiser consumers (q_denoise_backward, denoise_loss_grad) receive
// `target_stride = total_actions` and read only the first D=12 entries
// per sample. Same pattern as q_var_buf_trainer (see comment above
// line 21629).
// denoise_q_input_buf is a snapshot of q_coord_buf which stays at b*12
// (the cross-branch attention output narrows 13→12 before the denoiser
// consumes it), so its sizing is unchanged. The kernel input_stride
// parameter still threads through for symmetry with target/refined.
let q_var_total_actions_for_target = config.branch_0_size + config.branch_1_size + config.branch_2_size + config.branch_3_size;
let denoise_target_q_buf = alloc_f32(&stream, b * q_var_total_actions_for_target, "denoise_target_q_buf")?;
let denoise_q_input_buf = alloc_f32(&stream, b * 12, "denoise_q_input_buf")?;
let denoise_norm_partials = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc denoise_norm_partials: {e}")))?;
@@ -26158,8 +26184,12 @@ impl GpuDqnTrainer {
Ok(())
}
/// Compute target network Q-values [B, 12] into denoise_target_q_buf.
/// Compute target network Q-values [B, total_actions] into denoise_target_q_buf.
/// Uses compute_expected_q on target logits (tg_v_logits_buf, tg_b_logits_buf).
/// `compute_expected_q` writes `total_actions = b0+b1+b2+b3 = 13` floats per
/// sample (4-direction factored layout); the buffer is sized to match.
/// Per-sample consumers thread `target_stride = total_actions` to read only
/// the first D=12 entries.
pub(crate) fn compute_denoise_target_q(&self, batch_size: usize) -> Result<(), MLError> {
let n = batch_size as i32;
let na = self.config.num_atoms as i32;
@@ -26230,10 +26260,19 @@ impl GpuDqnTrainer {
let d_params_ptr = self.denoise_grad.raw_ptr();
let q_input_ptr = self.denoise_q_input_buf.raw_ptr();
let var_ptr = self.q_var_buf_trainer.raw_ptr();
// q_var_buf_trainer is sized [B, total_actions]; denoiser only reads
// the first D=12 slots per sample, but the kernel needs the true
// stride so per-sample addressing matches the buffer layout.
// Per-sample strides for the four [B, ?] inputs. The denoiser MLP only
// reads the first D=12 slots per sample, but each upstream buffer may
// be sized wider — the kernel needs the true stride to align per-sample
// addressing.
//
// q_var_buf_trainer: [B, total_actions=13] (compute_expected_q writes)
// denoise_target_q_buf: [B, total_actions=13] (compute_expected_q writes)
// q_coord_buf: [B, 12] (cross-branch attention output)
// denoise_q_input_buf: [B, 12] (snapshot of q_coord_buf)
let var_stride = self.total_actions() as i32;
let target_stride = self.total_actions() as i32;
let refined_stride = 12_i32;
let input_stride = 12_i32;
unsafe {
self.stream
.launch_builder(&self.q_denoise_backward_kernel)
@@ -26245,6 +26284,9 @@ impl GpuDqnTrainer {
.arg(&var_ptr)
.arg(&b)
.arg(&var_stride)
.arg(&refined_stride)
.arg(&target_stride)
.arg(&input_stride)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
@@ -26508,9 +26550,15 @@ impl GpuDqnTrainer {
// BACKWARD
// ═══════════════════════════════════════════════════════
// 10. Loss gradient: d_res = 0.1 * 2 * (Q_refined - Q_target) / B
// 10. Loss gradient: d_res = 0.1 * 2 * (Q_refined - Q_target) / B.
// Q_refined = q_coord_buf [B, 12] (post-cross-branch-attention).
// Q_target = denoise_target_q_buf [B, total_actions=13] (post-compute_expected_q).
// Kernel reads only the first D=12 entries from each per sample but
// needs per-sample strides to handle the asymmetric layout.
let inv_b: f32 = 1.0 / b as f32;
let d_res_ptr = self.denoise_d_res_buf.raw_ptr();
let refined_stride: i32 = 12;
let target_stride: i32 = self.total_actions() as i32;
unsafe {
self.stream.launch_builder(&self.denoise_loss_grad_kernel)
.arg(&q_refined_ptr)
@@ -26518,6 +26566,8 @@ impl GpuDqnTrainer {
.arg(&d_res_ptr)
.arg(&(b as i32))
.arg(&d_val)
.arg(&refined_stride)
.arg(&target_stride)
.arg(&inv_b)
.launch(LaunchConfig {
grid_dim: (((D * b) as u32 + 255) / 256, 1, 1),

2
docs/dqn-wire-up-audit.md
View File

File diff suppressed because one or more lines are too long