diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 4bbbe99f1..83fd1adcb 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -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 { diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 136f5a9d1..0ce26908b 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -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, 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) }; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index ec01617cc..fe8e6e399 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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, // [B, ADV_H] save_h_b3: CudaSlice, // [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, - /// 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, /// 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, )?; } diff --git a/crates/ml/src/cuda_pipeline/value_decoder.rs b/crates/ml/src/cuda_pipeline/value_decoder.rs index a6d828ecc..7df6ecc5a 100644 --- a/crates/ml/src/cuda_pipeline/value_decoder.rs +++ b/crates/ml/src/cuda_pipeline/value_decoder.rs @@ -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`, diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs index 398af6569..a07c6d026 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs @@ -67,15 +67,19 @@ fn alloc_branching(stream: &Arc, cfg: &GpuDqnTrainCo let bf16_data: Vec = 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 = 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), ], }; diff --git a/docs/dqn-named-dims.md b/docs/dqn-named-dims.md index e265f87b3..ffb7aebc7 100644 --- a/docs/dqn-named-dims.md +++ b/docs/dqn-named-dims.md @@ -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