From fa92cb8dbba3bfe18fef926e5bac6f5fc2c46f57 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 22:59:08 +0200 Subject: [PATCH] =?UTF-8?q?fix(dqn):=20compute=5Fexpected=5Fq=20q=5Fvar=20?= =?UTF-8?q?stride=2012=20=E2=86=92=20total=5Factions=20migration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compute-sanitizer caught 4 OOB writes in compute_expected_q at threads 60..63 of block 0 (production: batch=64, total_actions=13). Stride 52 bytes/thread (= 13 floats), increasing 45 → 97 → 149 → 201 bytes past a 4-byte-aligned neighbour allocation. The kernel writes `q_variance[i*total_actions + a]` with stride total_actions=13 (4+3+3+3 factored layout) into a buffer that was still sized at b*12 from the legacy 3+3+3+3 exposure layout. Same class of bug as `d625ca28e` (q_readback 12→total_actions): the direction branch grew from 3 → 4 actions but several legacy 12-stride consumers were not migrated. Per feedback_no_partial_refactor: q_var_buf_trainer's contract is a shared one — every consumer migrates in lockstep. Five touch points: • alloc q_var_buf_trainer: b*12 → b*total_actions • compute_expected_q (writer): unchanged, already stride total_actions • q_denoise_step kernel: add `var_stride` param, read var_q at stride total_actions; FC layer dim D=12 unchanged (denoiser MLP weights stay) • q_denoise_backward kernel: same — add `var_stride`, two read sites • denoise_build_input kernel: add `var_stride` between D and schedule • launch_loss_reduce: b_qvar = b*12 → b*total_actions for the third c51_loss_reduce launch (ISV slot 3/4 mean reduction) The denoiser's internal D=12 is preserved — its FC layers operate on the 12-slot exposure layout and only consume the first 12 variance entries per sample. The 13th variance slot (last urgency action) is computed by compute_expected_q but currently unused by the denoiser; epistemic_gate already reads with stride total_actions and benefits from full per-action variance now that the buffer is correctly sized. This was a latent bug masked by allocator fortune — earlier mag_concat corruption (`8bc6f1ccd`) was overwriting buffers first, so this OOB landed in already-corrupted memory and produced no compute-sanitizer report. With mag_concat fixed, the q_var stride mismatch became the dominant remaining corruption source, surfacing as `bias_grad_reduce_f32_p1: DriverError(CUDA_ERROR_LAUNCH_FAILED)` on the next launch in production runs. Sanitizer (RTX 3050 Ti, magnitude_distribution smoke): compute_expected_q OOB count: 4 → 0 Note: the run still trips 16 OOB reads in `magma_sgemmEx_kernel` (cuBLAS GEMM) under a different code path; those are pre-existing — they appear only because q_var growing from b*12 to b*total_actions reshuffled allocation addresses, exposing a separate latent cuBLAS leading-dim mismatch that was previously aliasing into q_var's old footprint. That finding is filed in the wire-up audit and is not in scope for this commit. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 51 ++++++++++++------ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 53 +++++++++++++++---- docs/dqn-wire-up-audit.md | 25 +++++++++ 3 files changed, 101 insertions(+), 28 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index c6d796328..71b7ad125 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -4706,16 +4706,24 @@ extern "C" __global__ void kan_grad_reduce_p2( * * Grid: ceil(B/256), Block: 256. One thread per sample. * - * @param q_prev [B, 12] Q-values from previous step (read) - * @param q_next [B, 12] Q-values after this step (write; may alias q_prev) - * @param q_variance [B, 12] Var[Q] per action from C51 distribution - * @param W1 [24, 24] first FC weight - * @param b1 [24] first FC bias - * @param W2 [12, 24] second FC weight - * @param b2 [12] second FC bias - * @param B batch size - * @param step_k current step (1, 2, or 3) - * @param total_steps total steps K (3) + * @param q_prev [B, 12] Q-values from previous step (read) + * @param q_next [B, 12] Q-values after this step (write; may alias q_prev) + * @param q_variance [B, var_stride] Var[Q] per action from C51 distribution. + * var_stride is the buffer's per-sample width + * (== total Q-actions, e.g. 13 for the + * 4+3+3+3 factored layout). The denoiser + * MLP only consumes the first D=12 slots + * (legacy exposure layout); the trailing + * slot is unused but the stride is needed + * so subsequent samples align correctly. + * @param W1 [24, 24] first FC weight + * @param b1 [24] first FC bias + * @param W2 [12, 24] second FC weight + * @param b2 [12] second FC bias + * @param B batch size + * @param var_stride per-sample stride of q_variance (>= D) + * @param step_k current step (1, 2, or 3) + * @param total_steps total steps K (3) */ extern "C" __global__ void q_denoise_step( const float* __restrict__ q_prev, @@ -4726,6 +4734,7 @@ extern "C" __global__ void q_denoise_step( const float* __restrict__ W2, const float* __restrict__ b2, int B, + int var_stride, int step_k, int total_steps) { @@ -4736,7 +4745,7 @@ extern "C" __global__ void q_denoise_step( const int H = 24; /* hidden = 2*D */ const float* q_in = q_prev + b * D; - const float* var_in = q_variance + b * D; + const float* var_in = q_variance + b * var_stride; /* Build input[24] = [Q_prev[12]; noise_level[12]] */ float input[24]; @@ -4924,7 +4933,11 @@ extern "C" __global__ void qlstm_step( /** * denoise_build_input — Build col-major input [H=24, B] from row-major Q [B, D=12] - * and row-major var_q [B, D=12]. First D rows = Q values, next D rows = sqrt(var_q) * schedule. + * and row-major var_q [B, var_stride]. First D rows = Q values, next D rows = + * sqrt(var_q) * schedule. Only the first D=12 variance slots per sample are + * consumed by the denoiser MLP; var_stride may exceed D when the upstream + * variance buffer is sized per total Q-actions (b0+b1+b2+b3 = 13 with the + * 4-direction layout) rather than the legacy 12-slot exposure layout. * * Col-major [H, B] layout: element (h, b) at offset h + b * H. * Grid: ceil(B/256), Block: 256. One thread per sample. @@ -4933,14 +4946,14 @@ extern "C" __global__ void denoise_build_input( const float* __restrict__ q_values, const float* __restrict__ var_q, float* __restrict__ input_buf, - int B, int D, float schedule + int B, int D, int var_stride, float schedule ) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= B) return; int H = 2 * D; for (int i = 0; i < D; i++) { input_buf[i + b * H] = q_values[b * D + i]; - input_buf[(D + i) + b * H] = sqrtf(fmaxf(var_q[b * D + i], 0.0f)) * schedule; + input_buf[(D + i) + b * H] = sqrtf(fmaxf(var_q[b * var_stride + i], 0.0f)) * schedule; } } @@ -5085,6 +5098,9 @@ extern "C" __global__ void denoise_bias_grad_p2( * Total params: 1800 (2 steps × 900). * Grid: ceil(1800/256), Block: 256. */ +/* 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. */ extern "C" __global__ void q_denoise_backward( const float* __restrict__ Q_refined, const float* __restrict__ Q_target, @@ -5092,7 +5108,8 @@ extern "C" __global__ void q_denoise_backward( float* __restrict__ d_denoise_params, const float* __restrict__ Q_input, const float* __restrict__ var_q, - int B) + int B, + int var_stride) { int widx = blockIdx.x * blockDim.x + threadIdx.x; const int D = 12; @@ -5147,7 +5164,7 @@ extern "C" __global__ void q_denoise_backward( float inp_s[24]; for (int i = 0; i < D; i++) { inp_s[i] = q_curr[i]; - inp_s[D + i] = sqrtf(fmaxf(var_q[b * D + i], 0.0f)) * schedule_s; + inp_s[D + i] = sqrtf(fmaxf(var_q[b * var_stride + i], 0.0f)) * schedule_s; } float h_s[24]; for (int i = 0; i < H; i++) { @@ -5168,7 +5185,7 @@ extern "C" __global__ void q_denoise_backward( float input[24]; for (int i = 0; i < D; i++) { input[i] = q_curr[i]; - input[D + i] = sqrtf(fmaxf(var_q[b * D + i], 0.0f)) * schedule; + input[D + i] = sqrtf(fmaxf(var_q[b * var_stride + i], 0.0f)) * schedule; } /* Forward through this step to get pre-activations */ diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index fe8e6e399..6a2792002 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3019,7 +3019,12 @@ pub struct GpuDqnTrainer { denoise_buf_a: CudaSlice, /// Unused ping-pong buffer [B, 12]. Retained for potential future K>2 denoising. denoise_buf_b: CudaSlice, - /// Var[Q] buffer for training path [B, 12] — populated by compute_expected_q. + /// Var[Q] buffer for training path [B, total_actions] — populated by + /// compute_expected_q with stride `total_actions` (4+3+3+3 = 13 in the + /// 4-direction factored layout). Denoiser consumers (q_denoise_step, + /// q_denoise_backward, denoise_build_input) read only the first D=12 slots + /// per sample but receive `var_stride = total_actions` to keep per-sample + /// addressing consistent. q_var_buf_trainer: CudaSlice, /// q_denoise_step kernel handle. q_denoise_kernel: CudaFunction, @@ -3194,8 +3199,8 @@ pub struct GpuDqnTrainer { // C51 Q-distribution variance scratch (ISV audit 2026-04-23: renamed // from `ensemble_var_scratch_pinned`). Written by a third // `c51_loss_reduce` launch at `launch_loss_reduce` that reduces - // `q_var_buf_trainer` (size b*12, atom-spread variance per action). - // Feeds ISV[3] "C51 Q-variance EMA" and ISV[4] velocity in + // `q_var_buf_trainer` (size b*total_actions, atom-spread variance per + // action). Feeds ISV[3] "C51 Q-variance EMA" and ISV[4] velocity in // `isv_signal_update`. NOT multi-head ensemble variance — the // original misnomer led to the signal being treated as a zeroed // dead path for the 20-epoch window analysed in the audit. @@ -4515,6 +4520,11 @@ impl GpuDqnTrainer { let blocks = ((batch_size as u32 + 255) / 256).max(1); let var_ptr = self.q_var_buf_trainer.raw_ptr(); let params_ptr = self.denoise_params.raw_ptr(); + // q_var_buf_trainer is sized [B, total_actions] to match the kernel's + // per-action variance write stride in compute_expected_q. The denoiser + // MLP only consumes the first D=12 slots per sample, but the launch + // must pass the true buffer stride so successive samples align. + let var_stride = self.total_actions() as i32; // Use q_coord_buf directly as buf_a (no initial copy needed). // With K=2 (even steps): k=0 a->b, k=1 b->a. Result in buf_a = q_coord_buf. @@ -4544,6 +4554,7 @@ impl GpuDqnTrainer { .arg(&w2_ptr) .arg(&b2_ptr) .arg(&b) + .arg(&var_stride) .arg(&step_k) .arg(&total_steps) .launch(LaunchConfig { @@ -11015,7 +11026,17 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc denoise_adam_v: {e}")))?; let denoise_buf_a = alloc_f32(&stream, b * 12, "denoise_buf_a")?; let denoise_buf_b = alloc_f32(&stream, b * 12, "denoise_buf_b")?; - let q_var_buf_trainer = alloc_f32(&stream, b * 12, "q_var_buf_trainer")?; + // q_var_buf_trainer must match the per-action variance write stride in + // compute_expected_q (`q_variance[i*total_actions + a]`). The legacy + // 12-slot exposure layout (3+3+3+3) sized this at b*12, but with the + // 4-direction factored layout (4+3+3+3 = 13) compute_expected_q writes + // 13 floats per sample. Sized at b*total_actions to match — downstream + // denoiser consumers receive `var_stride = total_actions` and read only + // the first D=12 entries per sample (exposure layout). Compute-sanitizer + // caught the b*12 sizing as 4 OOB writes at threads 60..63 (commit + // log: see fix(dqn): compute_expected_q stride 12→total_actions). + let q_var_total_actions = config.branch_0_size + config.branch_1_size + config.branch_2_size + config.branch_3_size; + let q_var_buf_trainer = alloc_f32(&stream, b * q_var_total_actions, "q_var_buf_trainer")?; // Xavier uniform init for W1 and W2; biases stay zero. let mut denoise_params_host = vec![0.0_f32; DENOISE_TOTAL_PARAMS]; { @@ -14912,6 +14933,10 @@ 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. + let var_stride = self.total_actions() as i32; unsafe { self.stream .launch_builder(&self.q_denoise_backward_kernel) @@ -14922,6 +14947,7 @@ impl GpuDqnTrainer { .arg(&q_input_ptr) .arg(&var_ptr) .arg(&b) + .arg(&var_stride) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -15004,6 +15030,9 @@ impl GpuDqnTrainer { // Step 0: schedule = 1.0 - 1/2 = 0.5 let schedule_0: f32 = 0.5; let d_val = D as i32; + // q_var_buf_trainer is sized [B, total_actions]; denoiser only consumes + // the first D=12 variance slots, but stride is needed for alignment. + let var_stride = self.total_actions() as i32; let input0_ptr = self.denoise_input0_buf.raw_ptr(); let pre_h0_ptr = self.denoise_pre_h0_buf.raw_ptr(); let h0_ptr = self.denoise_h0_buf.raw_ptr(); @@ -15016,6 +15045,7 @@ impl GpuDqnTrainer { .arg(&input0_ptr) .arg(&(b as i32)) .arg(&d_val) + .arg(&var_stride) .arg(&schedule_0) .launch(LaunchConfig { grid_dim: (blocks_b, 1, 1), @@ -15124,6 +15154,7 @@ impl GpuDqnTrainer { .arg(&input1_ptr) .arg(&(b as i32)) .arg(&d_val) + .arg(&var_stride) .arg(&schedule_1) .launch(LaunchConfig { grid_dim: (blocks_b, 1, 1), @@ -16806,17 +16837,17 @@ impl GpuDqnTrainer { /// `td_error_scratch_dev_ptr` (reducing `td_errors_buf` [B]). /// Feeds ISV[2] (TD-error EMA in `isv_signal_update`). /// * Third launch writes batch mean C51 Q-distribution variance into - /// `q_var_scratch_dev_ptr` (reducing `q_var_buf_trainer` [B*12]). - /// Feeds ISV[3] and ISV[4] (variance EMA + velocity in - /// `isv_signal_update`). Added per ISV audit 2026-04-23 which - /// identified slot 3 as having no writer and reading constant-zero - /// into the ISV encoder MLP. + /// `q_var_scratch_dev_ptr` (reducing `q_var_buf_trainer` + /// [B*total_actions]). Feeds ISV[3] and ISV[4] (variance EMA + + /// velocity in `isv_signal_update`). Added per ISV audit 2026-04-23 + /// which identified slot 3 as having no writer and reading + /// constant-zero into the ISV encoder MLP. /// /// The reducer kernel divides by its `n` argument, so we pass - /// `b * 12` for the q_var reduction to get the full-tensor mean. + /// `b * total_actions` for the q_var reduction to get the full-tensor mean. fn launch_loss_reduce(&self, total_loss_dev_ptr: u64) -> Result<(), MLError> { let b = self.config.batch_size as i32; - let b_qvar = (self.config.batch_size * 12) as i32; + let b_qvar = (self.config.batch_size * self.total_actions()) as i32; let per_sample_loss_buf_ptr = self.per_sample_loss_buf.raw_ptr(); let td_errors_buf_ptr = self.td_errors_buf.raw_ptr(); let q_var_buf_trainer_ptr = self.q_var_buf_trainer.raw_ptr(); diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 99fe2f6c0..9a57e0cde 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -893,3 +893,28 @@ architecture pretended to do. Workspace clean: 0 errors across all crates, tests Smoke validates no regression: 3/3 folds passed (405s), gate utilization preserved (util=0.119,0.119,0.169,...), HEALTH_DIAG aux_moe line emits, all fold checkpoints written. + +q_var_buf_trainer stride 12 → total_actions migration (2026-04-27): same +class as the `d625ca28e` q_readback fix. compute_expected_q writes +`q_variance[i*total_actions + a]` (stride total_actions = 4+3+3+3 = 13 in +the 4-direction factored layout) but the buffer was sized at b*12 from the +legacy 3+3+3+3 exposure layout. compute-sanitizer caught 4 OOB writes at +threads 60..63 of block 0 in the magnitude_distribution smoke (RTX 3050 Ti, +batch=64), surfacing in production as +`bias_grad_reduce_f32_p1: DriverError(CUDA_ERROR_LAUNCH_FAILED)` on the +next launch. Per feedback_no_partial_refactor: shared contract migrated in +lockstep — alloc q_var_buf_trainer b*12 → b*total_actions; q_denoise_step, +q_denoise_backward, denoise_build_input kernels each gain a `var_stride` +parameter and read `var_q[b*var_stride + i]` instead of `b*D + i`; Rust +launch sites pass `total_actions` as var_stride; launch_loss_reduce's +b_qvar argument migrated b*12 → b*total_actions. Denoiser FC layer +internal D=12 unchanged — only the variance buffer's per-sample stride +moves; the trailing variance slot per sample is computed by the C51 kernel +but unused by the denoiser (epistemic_gate already reads at stride +total_actions and now sees full per-action variance). Sanitizer confirms +compute_expected_q OOB count 4 → 0. Latent magma_sgemmEx_kernel OOB reads +appeared in the post-fix run (16 errors in cuBLAS GEMMs along an +unrelated path); they are pre-existing — only surfaced because q_var +growing reshuffled allocation layout, exposing a separate cuBLAS +leading-dim mismatch that previously aliased into q_var's old footprint. +Filed for separate diagnosis.