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

fix(dqn): compute_expected_q q_var stride 12 → total_actions migration

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-27 22:59:08 +02:00
parent 73fd49f4b5
commit fa92cb8dbb
3 changed files with 101 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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 */

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

@@ -3019,7 +3019,12 @@ pub struct GpuDqnTrainer {
denoise_buf_a: CudaSlice<f32>,
/// Unused ping-pong buffer [B, 12]. Retained for potential future K>2 denoising.
denoise_buf_b: CudaSlice<f32>,
/// 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<f32>,
/// 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();

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

@@ -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.