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

fix(dqn): mag_concat_qdir SH2+b0 vs SH2+3 off-by-one — buffer + weight migration

Browse Source 
compute-sanitizer pinpointed mag_concat_qdir at experience_kernels.cu:3590
writing 260 floats/state (SH2 + b0_size where b0=4) into a buffer
allocated for 259 floats/state (SH2 + 3, legacy 3-direction layout).
Off-by-one corrupted the next row's first column on every write and
overran past the buffer end on the final row, surfacing as
CUDA_ERROR_ILLEGAL_ADDRESS in downstream kernels (denoise_bias_grad_p1,
cublasLt h_v matmul) on L40S production batch sizes.

The `+3` constant was overloaded:
- direction-conditioned (mag_concat, w_b1fc, w_gate_1) — incorrectly
  hardcoded SH2+3 instead of SH2+branch_0_size when the kernel migrated
  to 4-direction (S/H/L/F).
- OFI-conditioned (ord_concat, urg_concat, w_b2fc, w_b3fc, w_gate_2,
  w_gate_3) — correctly SH2+3 for 3 OFI features per branch
  (concat_ofi_features).

Migrated all direction-conditioned consumers in lockstep
(feedback_no_partial_refactor):

- gpu_dqn_trainer.rs: w_b1fc, w_gate_1 use shared_h2+branch_0_size
- gpu_dqn_trainer.rs: split mag_concat_dim (SH2+b0) from
  ofi_concat_dim (SH2+3) for buffer alloc
- gpu_dqn_trainer.rs: accumulate_d_h_s2_from_concat takes src_stride
  param so mag callers pass SH2+b0, ord/urg callers pass SH2+3
- batched_forward.rs: split mag_concat_dim (SH2+b0) from ofi_concat_dim
  (SH2+3); strided_scatter dst_stride and fc_k now diverge between mag
  (d==1) and ord/urg (d==2,3); add separate (SH2+3) GEMM cache shape
- batched_backward.rs: d==1 (magnitude) dX/dW dims use SH2+b0;
  d==2/3 (order/urgency) keep SH2+3; add separate (SH2+3) GEMM cache
  shape for OFI branches
- gradient_budget.rs: smoke test now allocates b1 with SH2+b0
  and b2/b3 with SH2+3 (was buggy SH2+3 for all three)
- value_decoder.rs: doc updated
- docs/dqn-named-dims.md: new "Branch FC input strides" section
  documenting the direction- vs OFI-conditioning invariant

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-27 22:19:41 +02:00
parent 5857f5147a
commit 8bc6f1ccd3
6 changed files with 121 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -333,7 +333,9 @@ impl CublasBackwardSet {
}
// Branch FC layers: (adv_h, shared_h2)
add_fc_shapes(&mut unique_shapes, ah, sh2);
// Branch 1 FC (magnitude) wider input: (adv_h, shared_h2 + 3)
// Branch 1 FC (magnitude) wider input (direction-conditioned): (adv_h, shared_h2 + branch_0_size)
add_fc_shapes(&mut unique_shapes, ah, sh2 + config.branch_0_size);
// Branches 2/3 FC (order/urgency) wider input (OFI-conditioned, 3 features per branch): (adv_h, shared_h2 + 3)
add_fc_shapes(&mut unique_shapes, ah, sh2 + 3);
// Value output: (NA, value_h)
add_fc_shapes(&mut unique_shapes, na, vh);
@@ -1457,14 +1459,14 @@ impl CublasBackwardSet {
let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream);
if d == 1 && mag_concat_ptr != 0 {
// Magnitude: dX writes to d_mag_concat [B, SH2+3].
// Magnitude (direction-conditioned): dX writes to d_mag_concat [B, SH2+branch_0_size].
self.launch_dx_only(
stream,
d_glu_value_ptr,
w_fc,
d_mag_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
self.shared_h2 + self.branch_0_size,
b,
0.0_f32,
)?;
@@ -1474,7 +1476,7 @@ impl CublasBackwardSet {
w_gate,
d_mag_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
self.shared_h2 + self.branch_0_size,
b,
1.0_f32,
)?;
@@ -1686,8 +1688,9 @@ impl CublasBackwardSet {
// When 0, dX is not computed (states not trainable — no bottleneck).
s1_dx_output: u64,
// Magnitude branch conditioning: saved forward concat and dX output
mag_concat_ptr: u64, // [B, SH2+3] saved forward concat (for dW)
d_mag_concat_ptr: u64, // [B, SH2+3] dX output for branch 1 (caller accumulates)
// (direction-conditioned: SH2 + branch_0_size cols).
mag_concat_ptr: u64, // [B, SH2+branch_0_size] saved forward concat (for dW)
d_mag_concat_ptr: u64, // [B, SH2+branch_0_size] dX output for branch 1 (caller accumulates)
// Order branch OFI conditioning: saved forward concat and dX output
ord_concat_ptr: u64, // [B, SH2+3] saved forward concat (for dW)
d_ord_concat_ptr: u64, // [B, SH2+3] dX output for branch 2 (caller accumulates)
@@ -1848,9 +1851,11 @@ impl CublasBackwardSet {
let n_d = branch_n[d];
let w_out = w_ptrs[w_bout_idx[d]];
// FC input for dW (branches 1,2,3 may use wider concat buffers)
// FC input for dW (branches 1,2,3 may use wider concat buffers).
// Magnitude (d==1) is direction-conditioned (SH2 + branch_0_size);
// order/urgency (d==2,3) are OFI-conditioned (SH2 + 3).
let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 {
(mag_concat_ptr, self.shared_h2 + 3)
(mag_concat_ptr, self.shared_h2 + self.branch_0_size)
} else if d == 2 && ord_concat_ptr != 0 {
(ord_concat_ptr, self.shared_h2 + 3)
} else if d == 3 && urg_concat_ptr != 0 {

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

@@ -145,10 +145,18 @@ pub struct CublasGemmSet {
branch_2_size: usize,
branch_3_size: usize,
/// Magnitude branch (d==1) FC input dimension: shared_h2 + 3
/// (concat of h_s2 and Q_dir). Used when mag_concat_ptr != 0.
/// Magnitude branch (d==1) FC input dimension: shared_h2 + branch_0_size
/// (concat of h_s2 and per-direction Q_dir). Direction-conditioned, so
/// the trailing-stride width derives from `branch_0_size` (4 in production:
/// S/H/L/F). Used when mag_concat_ptr != 0.
mag_concat_dim: usize,
/// Order/urgency branch (d∈{2,3}) FC input dimension: shared_h2 + 3
/// (concat of vsn_masked and 3 OFI features per branch). OFI-conditioned,
/// so the trailing-stride width is fixed at 3 (matches `concat_ofi_features`
/// kernel writes). Used when ord_concat_ptr/urg_concat_ptr != 0.
ofi_concat_dim: usize,
// ── Multi-stream branch dispatch ──
/// 4 forked CUDA streams for parallel advantage branch execution.
/// Each branch (exposure, order, urgency, ...) submits GEMMs to its own stream,
@@ -435,7 +443,11 @@ impl CublasGemmSet {
unique_shapes.push((num_atoms, batch_size, value_h, value_h));
// h_bd (×4): M=adv_h, N=batch, K=shared_h2, ldb=shared_h2
unique_shapes.push((adv_h, batch_size, shared_h2, shared_h2));
// h_bd magnitude (d==1) wider input: M=adv_h, N=batch, K=shared_h2+3, ldb=shared_h2+3
// h_bd magnitude (d==1) wider input (direction-conditioned):
// M=adv_h, N=batch, K=shared_h2+branch_0_size, ldb=shared_h2+branch_0_size
unique_shapes.push((adv_h, batch_size, shared_h2 + branch_0_size, shared_h2 + branch_0_size));
// h_bd order/urgency (d∈{2,3}) wider input (OFI-conditioned, 3 features per branch):
// M=adv_h, N=batch, K=shared_h2+3, ldb=shared_h2+3
unique_shapes.push((adv_h, batch_size, shared_h2 + 3, shared_h2 + 3));
// adv_logits (×4): M=branch_k*num_atoms, N=batch, K=adv_h, ldb=adv_h
for &bs in &branch_sizes {
@@ -457,7 +469,8 @@ impl CublasGemmSet {
(shared_h2, batch_size, shared_h1, shared_h1), // h_s2
(value_h, batch_size, shared_h2, shared_h2), // h_v
(adv_h, batch_size, shared_h2, shared_h2), // h_bd (×4)
(adv_h, batch_size, shared_h2 + 3, shared_h2 + 3), // h_bd magnitude wider input
(adv_h, batch_size, shared_h2 + branch_0_size, shared_h2 + branch_0_size), // h_bd magnitude wider input (direction-conditioned)
(adv_h, batch_size, shared_h2 + 3, shared_h2 + 3), // h_bd order/urgency wider input (OFI-conditioned)
];
let mut gemm_cache_relu_bias = HashMap::new();
for &(n, b, k, ldb) in &relu_bias_shapes {
@@ -623,7 +636,8 @@ impl CublasGemmSet {
branch_1_size,
branch_2_size,
branch_3_size,
mag_concat_dim: shared_h2 + 3,
mag_concat_dim: shared_h2 + branch_0_size,
ofi_concat_dim: shared_h2 + 3,
branch_streams,
_branch_workspace_bufs: branch_workspace_bufs,
branch_workspace_ptrs,
@@ -1778,7 +1792,8 @@ impl CublasGemmSet {
/// `branch_idx ∈ [0, 4)` selects which branch's weights/buffers to use.
/// `h_s2_ptr` is the encoder output (read-only). For `branch_idx == 1`
/// (magnitude), pass the pre-built mag_concat pointer to use the wider
/// `[B, SH2 + 3]` input; pass 0 for the legacy `[B, SH2]` input.
/// `[B, SH2 + branch_0_size]` input (direction-conditioned); pass 0 for
/// the legacy `[B, SH2]` input.
/// `branch_h_ptr` is the per-branch hidden activation buffer; the
/// `adv_logits_ptr` is the per-branch slice into the flat
/// `b_logits_buf` at byte offset
@@ -2464,7 +2479,7 @@ impl CublasGemmSet {
stream: &Arc<CudaStream>,
d: usize, // branch index 0..3
h_s2_ptr: u64, // [B, SH2] trunk output
mag_concat_ptr: u64, // [B, SH2+3] magnitude concat (0 if unused)
mag_concat_ptr: u64, // [B, SH2+branch_0_size] magnitude concat (0 if unused)
w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
branch_h_ptr: u64, // [B, AH] output (save_h_bd)
_ws_ptr: u64,
@@ -2505,18 +2520,22 @@ impl CublasGemmSet {
// ── 2. Rebuild concat from VSN-masked output ──
//
// d==1 (magnitude): Scatter vsn_masked → first SH2 cols of mag_concat [B, SH2+3].
// Q_dir (last 3 cols) were written by the pre-VSN launch_mag_concat_from call.
// d==1 (magnitude): Scatter vsn_masked → first SH2 cols of mag_concat
// [B, SH2+branch_0_size]. Q_dir (last branch_0_size cols) was written
// by the pre-VSN launch_mag_concat_from call.
//
// d==2 (order) / d==3 (urgency): Scatter vsn_masked → first SH2 cols of
// ord_concat / urg_concat [B, SH2+3].
// OFI (last 3 cols) were written by the pre-VSN launch_concat_ofi call.
// ord_concat / urg_concat [B, SH2+3]. OFI (last 3 cols) was written
// by the pre-VSN launch_concat_ofi call.
//
// Note: dst_stride differs between mag and ord/urg branches when
// branch_0_size != 3 (production: 4 vs 3).
// This is one kernel launch (B threads, negligible cost).
let scatter_k = self.strided_scatter_kernel.as_ref().unwrap();
let total_scatter = (b * sh2) as i32;
let src_stride_scatter = sh2 as i32;
let dst_stride_scatter = self.mag_concat_dim as i32;
let mag_dst_stride = self.mag_concat_dim as i32;
let ofi_dst_stride = self.ofi_concat_dim as i32;
let scatter_blocks = ((total_scatter as u32 + 255) / 256).max(1);
let (vsn_input, fc_k) = if d == 1 && mag_concat_ptr != 0 {
@@ -2526,7 +2545,7 @@ impl CublasGemmSet {
.arg(&self.vsn_masked_ptr)
.arg(&mag_concat_ptr)
.arg(&src_stride_scatter)
.arg(&dst_stride_scatter)
.arg(&mag_dst_stride)
.arg(&total_scatter)
.launch(LaunchConfig {
grid_dim: (scatter_blocks, 1, 1),
@@ -2545,7 +2564,7 @@ impl CublasGemmSet {
.arg(&self.vsn_masked_ptr)
.arg(&self.ord_concat_ptr)
.arg(&src_stride_scatter)
.arg(&dst_stride_scatter)
.arg(&ofi_dst_stride)
.arg(&total_scatter)
.launch(LaunchConfig {
grid_dim: (scatter_blocks, 1, 1),
@@ -2556,7 +2575,7 @@ impl CublasGemmSet {
"strided_scatter vsn→ord_concat: {e}"
)))?;
}
(self.ord_concat_ptr, self.mag_concat_dim)
(self.ord_concat_ptr, self.ofi_concat_dim)
} else if d == 3 && self.urg_concat_ptr != 0 {
unsafe {
stream
@@ -2564,7 +2583,7 @@ impl CublasGemmSet {
.arg(&self.vsn_masked_ptr)
.arg(&self.urg_concat_ptr)
.arg(&src_stride_scatter)
.arg(&dst_stride_scatter)
.arg(&ofi_dst_stride)
.arg(&total_scatter)
.launch(LaunchConfig {
grid_dim: (scatter_blocks, 1, 1),
@@ -2575,7 +2594,7 @@ impl CublasGemmSet {
"strided_scatter vsn→urg_concat: {e}"
)))?;
}
(self.urg_concat_ptr, self.mag_concat_dim)
(self.urg_concat_ptr, self.ofi_concat_dim)
} else {
(self.vsn_masked_ptr, sh2)
};

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

@@ -1467,7 +1467,7 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT
cfg.adv_h, // [18] b_b0fc
cfg.branch_0_size * cfg.num_atoms * cfg.adv_h, // [19] w_b0out
cfg.branch_0_size * cfg.num_atoms, // [20] b_b0out
cfg.adv_h * (cfg.shared_h2 + 3), // [21] w_b1fc (direction-conditioned: SH2+3 input)
cfg.adv_h * (cfg.shared_h2 + cfg.branch_0_size), // [21] w_b1fc (direction-conditioned: SH2+b0 input)
cfg.adv_h, // [22] b_b1fc
cfg.branch_1_size * cfg.num_atoms * cfg.adv_h, // [23] w_b1out
cfg.branch_1_size * cfg.num_atoms, // [24] b_b1out
@@ -1494,7 +1494,7 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT
// ── GLU gate weights ──
cfg.adv_h * cfg.shared_h2, // [43] w_gate_0 [AH, SH2]
cfg.adv_h, // [44] b_gate_0
cfg.adv_h * (cfg.shared_h2 + 3), // [45] w_gate_1 [AH, SH2+3]
cfg.adv_h * (cfg.shared_h2 + cfg.branch_0_size), // [45] w_gate_1 [AH, SH2+b0] (direction-conditioned)
cfg.adv_h, // [46] b_gate_1
cfg.adv_h * (cfg.shared_h2 + 3), // [47] w_gate_2 (OFI-conditioned)
cfg.adv_h, // [48] b_gate_2
@@ -1896,9 +1896,10 @@ pub struct GpuDqnTrainer {
save_h_b2: CudaSlice<f32>, // [B, ADV_H]
save_h_b3: CudaSlice<f32>, // [B, ADV_H]
/// Magnitude branch concat input [B, SH2+3]: [h_s2; Q_dir] for direction conditioning.
/// Magnitude branch concat input [B, SH2+branch_0_size]: [h_s2; Q_dir]
/// for direction conditioning (4 directions: S/H/L/F).
mag_concat_buf: CudaSlice<f32>,
/// Backward scratch for magnitude concat dX [B, SH2+3].
/// Backward scratch for magnitude concat dX [B, SH2+branch_0_size].
d_mag_concat_buf: CudaSlice<f32>,
/// Kernel: mag_concat_qdir — builds [h_s2; Q_dir] concat.
mag_concat_kernel: CudaFunction,
@@ -5703,25 +5704,29 @@ impl GpuDqnTrainer {
Ok(())
}
/// Accumulate first SH2 columns of d_mag_concat [B, SH2+3] into d_h_s2 [B, SH2].
/// Accumulate first SH2 columns of a concat dX buffer into d_h_s2 [B, SH2]
/// from a wider buffer with `src_stride` columns. Magnitude callers pass
/// `SH2 + branch_0_size` (direction-conditioned); order/urgency callers pass
/// `SH2 + 3` (OFI-conditioned).
pub(crate) fn accumulate_d_h_s2_from_concat(
&self,
d_concat_ptr: u64,
d_h_s2_ptr: u64,
batch: usize,
src_stride: usize,
beta: f32,
) -> Result<(), MLError> {
let total = (batch * self.config.shared_h2) as i32;
let blocks = ((total as u32 + 255) / 256).max(1);
let sh2 = self.config.shared_h2 as i32;
let sh2_plus_3 = (self.config.shared_h2 + 3) as i32;
let src_stride_i32 = src_stride as i32;
unsafe {
self.stream
.launch_builder(&self.strided_accumulate_kernel)
.arg(&d_concat_ptr)
.arg(&d_h_s2_ptr)
.arg(&sh2)
.arg(&sh2_plus_3)
.arg(&src_stride_i32)
.arg(&total)
.arg(&beta)
.launch(LaunchConfig {
@@ -7339,6 +7344,7 @@ impl GpuDqnTrainer {
self.ptrs.d_mag_concat_buf,
scratch_d_h_s2,
self.config.batch_size,
self.config.shared_h2 + self.config.branch_0_size, // mag: direction-conditioned
1.0, // beta=1: d==0 already wrote to scratch_d_h_s2
)?;
}
@@ -7347,6 +7353,7 @@ impl GpuDqnTrainer {
self.d_ord_concat_buf.raw_ptr(),
scratch_d_h_s2,
self.config.batch_size,
self.config.shared_h2 + 3, // ord: OFI-conditioned (3 OFI features)
1.0,
)?;
}
@@ -7355,6 +7362,7 @@ impl GpuDqnTrainer {
self.d_urg_concat_buf.raw_ptr(),
scratch_d_h_s2,
self.config.batch_size,
self.config.shared_h2 + 3, // urg: OFI-conditioned (3 OFI features)
1.0,
)?;
}
@@ -8785,13 +8793,19 @@ impl GpuDqnTrainer {
let save_h_b1 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b1")?;
let save_h_b2 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b2")?;
let save_h_b3 = alloc_f32(&stream, b * config.adv_h + kt, "save_h_b3")?;
let mag_concat_dim = config.shared_h2 + 3;
// Magnitude branch is direction-conditioned (mag_concat_qdir kernel writes
// SH2 + branch_0_size floats per state). Order/urgency branches are
// OFI-conditioned (concat_ofi_features kernel writes SH2 + 3 floats per
// state — exactly 3 OFI features per branch). These two strides differ
// when branch_0_size != 3 (production: 4 = S/H/L/F).
let mag_concat_dim = config.shared_h2 + config.branch_0_size;
let ofi_concat_dim = config.shared_h2 + 3;
let mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "mag_concat_buf")?;
let d_mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_mag_concat_buf")?;
let ord_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "ord_concat_buf")?;
let d_ord_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_ord_concat_buf")?;
let urg_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "urg_concat_buf")?;
let d_urg_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_urg_concat_buf")?;
let ord_concat_buf = alloc_f32(&stream, b * ofi_concat_dim + kt, "ord_concat_buf")?;
let d_ord_concat_buf = alloc_f32(&stream, b * ofi_concat_dim + kt, "d_ord_concat_buf")?;
let urg_concat_buf = alloc_f32(&stream, b * ofi_concat_dim + kt, "urg_concat_buf")?;
let d_urg_concat_buf = alloc_f32(&stream, b * ofi_concat_dim + kt, "d_urg_concat_buf")?;
let save_current_lp = alloc_f32(
&stream,
b * num_branches * config.num_atoms,
@@ -17542,6 +17556,7 @@ impl GpuDqnTrainer {
self.ptrs.d_mag_concat_buf,
d_h_s2_ptr,
self.config.batch_size,
self.config.shared_h2 + self.config.branch_0_size, // mag: direction-conditioned
1.0, // beta=1: d==0 already wrote to d_h_s2 in backward_full
)?;
}
@@ -17551,6 +17566,7 @@ impl GpuDqnTrainer {
self.d_ord_concat_buf.raw_ptr(),
d_h_s2_ptr,
self.config.batch_size,
self.config.shared_h2 + 3, // ord: OFI-conditioned (3 OFI features)
1.0,
)?;
}
@@ -17560,6 +17576,7 @@ impl GpuDqnTrainer {
self.d_urg_concat_buf.raw_ptr(),
d_h_s2_ptr,
self.config.batch_size,
self.config.shared_h2 + 3, // urg: OFI-conditioned (3 OFI features)
1.0,
)?;
}

6
crates/ml/src/cuda_pipeline/value_decoder.rs
View File

@@ -114,8 +114,10 @@ impl<'a> ValueDecoder<'a> {
///
/// Pointer ownership (caller-supplied):
/// - `h_s2_dev_ptr` : `[B, shared_h2]` encoder output (read-only)
/// - `mag_concat_dev_ptr` : `[B, shared_h2 + 3]` magnitude branch
/// wider input; pass `0` for legacy `[B, shared_h2]` input
/// - `mag_concat_dev_ptr` : `[B, shared_h2 + branch_0_size]` magnitude
/// branch wider input (direction-conditioned: SH2 + b0 = trunk activation
/// concatenated with per-direction Q values); pass `0` for legacy
/// `[B, shared_h2]` input
/// - `branch_h_dev_ptr` : `[B, adv_h]` per-branch hidden scratch
/// - `q_per_action_dev_ptr`: `[B, branch_q_dim * num_atoms]` output
/// - `v_short_dev_ptr`,

25
crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs
View File

@@ -67,15 +67,19 @@ fn alloc_branching(stream: &Arc<cudarc::driver::CudaStream>, cfg: &GpuDqnTrainCo
let bf16_data: Vec<f32> = vec![0.1; n];
stream.clone_htod(&bf16_data).expect("alloc branching weight")
};
// Branches 1-3 are direction-conditioned: input dim = shared_h2 + 3 (not shared_h2)
let b1_input = cfg.shared_h2 + 3;
// Branch 1 (magnitude) is direction-conditioned: input dim = SH2 + branch_0_size.
// Branches 2-3 (order, urgency) are OFI-conditioned: input dim = SH2 + 3
// (3 OFI features per branch — see concat_ofi_features kernel).
let b1_input = cfg.shared_h2 + cfg.branch_0_size;
let b2_input = cfg.shared_h2 + 3;
let b3_input = cfg.shared_h2 + 3;
let backing = BranchingWeightBacking {
slices: [
alloc(cfg.adv_h * b1_input), alloc(cfg.adv_h),
alloc(cfg.branch_1_size * na * cfg.adv_h), alloc(cfg.branch_1_size * na),
alloc(cfg.adv_h * b1_input), alloc(cfg.adv_h),
alloc(cfg.adv_h * b2_input), alloc(cfg.adv_h),
alloc(cfg.branch_2_size * na * cfg.adv_h), alloc(cfg.branch_2_size * na),
alloc(cfg.adv_h * b1_input), alloc(cfg.adv_h),
alloc(cfg.adv_h * b3_input), alloc(cfg.adv_h),
alloc(cfg.branch_3_size * na * cfg.adv_h), alloc(cfg.branch_3_size * na),
],
};
@@ -119,7 +123,7 @@ fn test_spectral_norm_all_heads_no_panic() -> anyhow::Result<()> {
"W_v1 has non-finite values after spectral norm"
);
let mut w_bo1_bf16 = vec![0.0_f32; cfg.adv_h * (cfg.shared_h2 + 3)];
let mut w_bo1_bf16 = vec![0.0_f32; cfg.adv_h * (cfg.shared_h2 + cfg.branch_0_size)];
stream.memcpy_dtoh(&b_backing.slices[0], &mut w_bo1_bf16)
.map_err(|e| anyhow::anyhow!("{e}"))?;
let w_bo1_host: Vec<f32> = w_bo1_bf16.to_vec();
@@ -158,15 +162,18 @@ fn test_spectral_norm_constrains_operator_norm() -> anyhow::Result<()> {
],
};
let dueling = d_backing.weight_set();
// Branches 1-3 are direction-conditioned: input dim = shared_h2 + 3
let b1_input = cfg.shared_h2 + 3;
// Branch 1 (magnitude) is direction-conditioned: input dim = SH2 + branch_0_size.
// Branches 2-3 (order, urgency) are OFI-conditioned: input dim = SH2 + 3.
let b1_input = cfg.shared_h2 + cfg.branch_0_size;
let b2_input = cfg.shared_h2 + 3;
let b3_input = cfg.shared_h2 + 3;
let b_backing = BranchingWeightBacking {
slices: [
alloc_large(cfg.adv_h * b1_input), alloc_large(cfg.adv_h),
alloc_large(cfg.branch_1_size * na * cfg.adv_h), alloc_large(cfg.branch_1_size * na),
alloc_large(cfg.adv_h * b1_input), alloc_large(cfg.adv_h),
alloc_large(cfg.adv_h * b2_input), alloc_large(cfg.adv_h),
alloc_large(cfg.branch_2_size * na * cfg.adv_h), alloc_large(cfg.branch_2_size * na),
alloc_large(cfg.adv_h * b1_input), alloc_large(cfg.adv_h),
alloc_large(cfg.adv_h * b3_input), alloc_large(cfg.adv_h),
alloc_large(cfg.branch_3_size * na * cfg.adv_h), alloc_large(cfg.branch_3_size * na),
],
};

23
docs/dqn-named-dims.md
View File

@@ -117,6 +117,29 @@ From the trade_plan MLP output.
| 1 | `MAG_HALF` | 0.50× max_position |
| 2 | `MAG_FULL` | 1.00× max_position |
## Branch FC input strides (direction- vs OFI-conditioned)
The four branch FC heads (`w_b{0,1,2,3}fc`) and their gate twins
(`w_gate_{0,1,2,3}`) have different input strides depending on what each
branch concatenates onto the trunk activation:
| Branch (d) | Constant role | Input stride | Conditioning |
|---|---|---|---|
| 0 (Direction) | trunk only | `shared_h2` | none |
| 1 (Magnitude) | mag_concat | `shared_h2 + branch_0_size` | direction-conditioned (concat of `h_s2` + per-direction Q values, 4 in production) |
| 2 (Order) | ord_concat | `shared_h2 + 3` | OFI-conditioned (concat of `vsn_masked` + 3 OFI features) |
| 3 (Urgency) | urg_concat | `shared_h2 + 3` | OFI-conditioned (concat of `vsn_masked` + 3 OFI features) |
**Invariant (off-by-one trap).** The `+3` constant is overloaded: in the OFI
branches it's a literal (3 features per branch from `concat_ofi_features`), but
in the magnitude branch it must derive from `branch_0_size` (since 4-direction
S/H/L/F was added). Buffers, weight tensors, GEMM cache shapes, accumulator
strides, and dX backward dims all share this contract — see commit fixing
mag_concat_qdir OOB (compute-sanitizer caught 1679 errors). Direction-conditioned
sites: `w_b1fc`, `w_gate_1`, `mag_concat_buf`, `d_mag_concat_buf`. OFI-conditioned
sites: `w_b{2,3}fc`, `w_gate_{2,3}`, `{ord,urg}_concat_buf`,
`d_{ord,urg}_concat_buf`.
## Commit history
- Task 4A (ps[0..PS_STRIDE) constants): commit 144c85b85