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

feat: TLOB microstructure injection — OFI features to order/urgency branches (3,072 params)

Browse Source 
Order branch gets [bid_ask_spread, depth_imbalance, queue_pressure].
Urgency branch gets [spread_velocity, depth_change_rate, trade_arrival_rate].
Same concat pattern as Layer 3 direction→magnitude conditioning.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 09:26:26 +02:00
parent e328bcdb13
commit 278eb152c0
3 changed files with 281 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -800,6 +800,12 @@ impl CublasBackwardSet {
// 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)
// 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)
// Urgency branch OFI conditioning: saved forward concat and dX output
urg_concat_ptr: u64, // [B, SH2+3] saved forward concat (for dW)
d_urg_concat_ptr: u64, // [B, SH2+3] dX output for branch 3 (caller accumulates)
// GLU saved forward activations for backward pass
glu_gate_pre: &[u64; 4], // [B, AH] each — saved pre-sigmoid gate activations
glu_value: &[u64; 4], // [B, AH] each — saved value path outputs
@@ -859,11 +865,11 @@ impl CublasBackwardSet {
let goff_w_b1out: u64 = goff_b_b1fc + ah * f32;
let goff_b_b1out: u64 = goff_w_b1out + b1 * na64 * ah * f32;
let goff_w_b2fc: u64 = goff_b_b1out + b1 * na64 * f32;
let goff_b_b2fc: u64 = goff_w_b2fc + ah * sh2 * f32;
let goff_b_b2fc: u64 = goff_w_b2fc + ah * (sh2 + 3) * f32; // order uses wider input: SH2+3
let goff_w_b2out: u64 = goff_b_b2fc + ah * f32;
let goff_b_b2out: u64 = goff_w_b2out + b2 * na64 * ah * f32;
let goff_w_b3fc: u64 = goff_b_b2out + b2 * na64 * f32;
let goff_b_b3fc: u64 = goff_w_b3fc + ah * sh2 * f32;
let goff_b_b3fc: u64 = goff_w_b3fc + ah * (sh2 + 3) * f32; // urgency uses wider input: SH2+3
let goff_w_b3out: u64 = goff_b_b3fc + ah * f32;
let goff_b_b3out: u64 = goff_w_b3out + b3 * na64 * ah * f32;
@@ -974,9 +980,13 @@ impl CublasBackwardSet {
// ── Step 3: Value path weight gradients ────────────────────
// dW_bdf += d_value^T @ vsn_input, db_bdf += sum(d_value)
// Branch 1 (magnitude) uses wider concat input [B, SH2+3] for dW.
// Branches 1,2,3 use wider concat input [B, SH2+3] for dW.
let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 {
(mag_concat_ptr, self.shared_h2 + 3)
} else if d == 2 && ord_concat_ptr != 0 {
(ord_concat_ptr, self.shared_h2 + 3)
} else if d == 3 && urg_concat_ptr != 0 {
(urg_concat_ptr, self.shared_h2 + 3)
} else {
(save_h_s2, self.shared_h2)
};
@@ -1033,6 +1043,56 @@ impl CublasBackwardSet {
)?;
// Caller will call accumulate_d_h_s2_from_concat(d_mag_concat, d_h_s2, B, 1.0)
// after backward_full returns.
} else if d == 2 && ord_concat_ptr != 0 {
// Order: dX writes to d_ord_concat [B, SH2+3].
// Value path: beta=0 (overwrite)
self.launch_dx_only(
stream,
scratch_d_glu_value,
w_fc,
d_ord_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32,
)?;
// Gate path: beta=1 (accumulate)
self.launch_dx_only(
stream,
scratch_d_glu_gate,
w_gate,
d_ord_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32,
)?;
// Caller will call accumulate_d_h_s2_from_concat(d_ord_concat, d_h_s2, B, 1.0)
} else if d == 3 && urg_concat_ptr != 0 {
// Urgency: dX writes to d_urg_concat [B, SH2+3].
// Value path: beta=0 (overwrite)
self.launch_dx_only(
stream,
scratch_d_glu_value,
w_fc,
d_urg_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32,
)?;
// Gate path: beta=1 (accumulate)
self.launch_dx_only(
stream,
scratch_d_glu_gate,
w_gate,
d_urg_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32,
)?;
// Caller will call accumulate_d_h_s2_from_concat(d_urg_concat, d_h_s2, B, 1.0)
} else {
// Other branches: accumulate into d_h_s2 directly.
let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 };

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

@@ -189,6 +189,10 @@ pub struct CublasGemmSet {
glu_value_ptrs: [u64; 4],
/// VSN bottleneck rank (R=16).
vsn_rank: usize,
/// Order branch (d=2) OFI concat buffer ptr [B, SH2+3] (0 if not wired).
ord_concat_ptr: u64,
/// Urgency branch (d=3) OFI concat buffer ptr [B, SH2+3] (0 if not wired).
urg_concat_ptr: u64,
}
/// Backward-compatible type alias — existing callers (`GpuDqnTrainer`,
@@ -349,6 +353,8 @@ impl CublasGemmSet {
glu_gate_pre_ptrs: [0; 4],
glu_value_ptrs: [0; 4],
vsn_rank: 0,
ord_concat_ptr: 0,
urg_concat_ptr: 0,
})
}
@@ -381,6 +387,15 @@ impl CublasGemmSet {
self.vsn_rank = vsn_rank;
}
/// Wire OFI concat buffer pointers for order (d=2) and urgency (d=3) branches.
///
/// Called once after construction from `GpuDqnTrainer`, after `wire_vsn_glu`.
/// These pointers point to [B, SH2+3] buffers pre-built by `launch_concat_ofi`.
pub fn wire_ofi_concat(&mut self, ord_concat_ptr: u64, urg_concat_ptr: u64) {
self.ord_concat_ptr = ord_concat_ptr;
self.urg_concat_ptr = urg_concat_ptr;
}
/// Rebind the cuBLAS handle to a (potentially new) stream.
///
// Accessors for diagnostic test
@@ -985,22 +1000,23 @@ impl CublasGemmSet {
)))?;
}
// ── 2. Rebuild mag_concat from VSN-masked output for d==1 ──
// The caller's launch_mag_concat used raw h_s2 (pre-VSN). Now that VSN
// has produced vsn_masked, re-run mag_concat_qdir with vsn_masked as
// source so magnitude sees feature-selected [vsn_masked; Q_dir].
// The Q_dir values (last 3 cols) are recomputed from the same logits.
// ── 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==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.
//
// 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 scatter_blocks = ((total_scatter as u32 + 255) / 256).max(1);
let (vsn_input, fc_k) = if d == 1 && mag_concat_ptr != 0 {
// Re-launch mag_concat with vsn_masked as source instead of raw h_s2
// The mag_concat_kernel reads source[B, SH2] + computes Q_dir from logits
// Scatter vsn_masked [B, SH2] → first SH2 cols of mag_concat [B, SH2+3].
// Q_dir (last 3 cols) preserved from the pre-VSN launch_mag_concat_from call.
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 scatter_blocks = ((total_scatter as u32 + 255) / 256).max(1);
unsafe {
stream
.launch_builder(scatter_k)
@@ -1019,6 +1035,44 @@ impl CublasGemmSet {
)))?;
}
(mag_concat_ptr, self.mag_concat_dim)
} else if d == 2 && self.ord_concat_ptr != 0 {
unsafe {
stream
.launch_builder(scatter_k)
.arg(&self.vsn_masked_ptr)
.arg(&self.ord_concat_ptr)
.arg(&src_stride_scatter)
.arg(&dst_stride_scatter)
.arg(&total_scatter)
.launch(LaunchConfig {
grid_dim: (scatter_blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"strided_scatter vsn→ord_concat: {e}"
)))?;
}
(self.ord_concat_ptr, self.mag_concat_dim)
} else if d == 3 && self.urg_concat_ptr != 0 {
unsafe {
stream
.launch_builder(scatter_k)
.arg(&self.vsn_masked_ptr)
.arg(&self.urg_concat_ptr)
.arg(&src_stride_scatter)
.arg(&dst_stride_scatter)
.arg(&total_scatter)
.launch(LaunchConfig {
grid_dim: (scatter_blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"strided_scatter vsn→urg_concat: {e}"
)))?;
}
(self.urg_concat_ptr, self.mag_concat_dim)
} else {
(self.vsn_masked_ptr, sh2)
};

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

@@ -630,6 +630,16 @@ pub struct GpuDqnTrainer {
strided_accumulate_kernel: CudaFunction,
/// Kernel: strided_scatter — writes vsn_masked into first SH2 cols of mag_concat.
strided_scatter_kernel: CudaFunction,
/// Order branch concat input [B, SH2+3]: [vsn_masked; OFI_order] for microstructure conditioning.
ord_concat_buf: CudaSlice<f32>,
/// Backward scratch for order concat dX [B, SH2+3].
d_ord_concat_buf: CudaSlice<f32>,
/// Urgency branch concat input [B, SH2+3]: [vsn_masked; OFI_urgency] for microstructure conditioning.
urg_concat_buf: CudaSlice<f32>,
/// Backward scratch for urgency concat dX [B, SH2+3].
d_urg_concat_buf: CudaSlice<f32>,
/// Kernel: concat_ofi_features — builds [vsn_masked; OFI] for order/urgency branches.
concat_ofi_kernel: CudaFunction,
// ── Cross-Branch Q Attention ──
q_attn_params: CudaSlice<f32>, // [624]
@@ -1324,6 +1334,44 @@ impl GpuDqnTrainer {
Ok(())
}
/// Build OFI concat for order (d=2) or urgency (d=3) branch.
///
/// Reads vsn_masked[B, SH2] and raw feature vector states[B, SD], appends 3 OFI
/// features at ofi_base offset, producing [B, SH2+3] in ord_concat_buf / urg_concat_buf.
///
/// For order (branch_idx=2): OFI features at indices 42, 43, 44
/// For urgency (branch_idx=3): OFI features at indices 45, 46, 47
pub(crate) fn launch_concat_ofi(&self, branch_idx: usize, states_ptr: u64) -> Result<(), MLError> {
let (concat_ptr, ofi_base) = match branch_idx {
2 => (self.ptrs.ord_concat_buf, 42_i32),
3 => (self.ptrs.urg_concat_buf, 45_i32),
_ => return Ok(()),
};
let b = self.config.batch_size as i32;
let sh2 = self.config.shared_h2 as i32;
let sd = self.config.state_dim as i32;
let total = b * (sh2 + 3);
let blocks = ((total as u32 + 255) / 256).max(1);
let vsn_masked_ptr = self.vsn_masked_buf.raw_ptr();
unsafe {
self.stream
.launch_builder(&self.concat_ofi_kernel)
.arg(&concat_ptr)
.arg(&vsn_masked_ptr)
.arg(&states_ptr)
.arg(&b)
.arg(&sh2)
.arg(&sd)
.arg(&ofi_base)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("concat_ofi branch {branch_idx}: {e}")))?;
}
Ok(())
}
/// Run 2-head cross-branch Q attention over the 12 raw Q-values in q_out_buf.
/// Writes coordinated Q-values to q_coord_buf [batch_size, 12].
@@ -2349,6 +2397,10 @@ impl GpuDqnTrainer {
0, // CQL backward: no bottleneck dX needed (separate gradient budget)
self.ptrs.mag_concat_buf,
self.ptrs.d_mag_concat_buf,
self.ptrs.ord_concat_buf,
self.d_ord_concat_buf.raw_ptr(),
self.ptrs.urg_concat_buf,
self.d_urg_concat_buf.raw_ptr(),
&glu_gate_pre_ptrs,
&glu_value_ptrs,
self.ptrs.bw_d_glu_value,
@@ -2364,6 +2416,22 @@ impl GpuDqnTrainer {
1.0, // beta=1: d==0 already wrote to scratch_d_h_s2
)?;
}
if self.ptrs.ord_concat_buf != 0 {
self.accumulate_d_h_s2_from_concat(
self.d_ord_concat_buf.raw_ptr(),
scratch_d_h_s2,
self.config.batch_size,
1.0,
)?;
}
if self.ptrs.urg_concat_buf != 0 {
self.accumulate_d_h_s2_from_concat(
self.d_urg_concat_buf.raw_ptr(),
scratch_d_h_s2,
self.config.batch_size,
1.0,
)?;
}
Ok(true)
}
@@ -2743,6 +2811,10 @@ impl GpuDqnTrainer {
let mag_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 save_current_lp = alloc_f32(
&stream,
b * num_branches * config.num_atoms,
@@ -2957,7 +3029,9 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("strided_accumulate load: {e}")))?;
let strided_scatter_kernel = exp_module_for_mag.load_function("strided_scatter")
.map_err(|e| MLError::ModelError(format!("strided_scatter load: {e}")))?;
info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter kernels loaded");
let concat_ofi_kernel = exp_module_for_mag.load_function("concat_ofi_features")
.map_err(|e| MLError::ModelError(format!("concat_ofi_features load: {e}")))?;
info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi kernels loaded");
// ── Cross-Branch Q Attention buffers ──────────────────────────
let q_attn_params = stream.alloc_zeros::<f32>(624)
@@ -3281,7 +3355,11 @@ impl GpuDqnTrainer {
],
16, // VSN bottleneck rank R=16
);
info!("GpuDqnTrainer: cuBLAS batched forward initialized (VSN+KAN wired)");
cublas_forward.wire_ofi_concat(
ord_concat_buf.raw_ptr(),
urg_concat_buf.raw_ptr(),
);
info!("GpuDqnTrainer: cuBLAS batched forward initialized (VSN+KAN+OFI wired)");
// DDQN forward shares the same cuBLAS handle — runs sequentially
// on the main stream for NVIDIA bit-wise reproducibility.
@@ -3317,7 +3395,11 @@ impl GpuDqnTrainer {
],
16,
);
info!("GpuDqnTrainer: cuBLAS batched forward (DDQN) initialized (VSN+KAN wired)");
cublas_forward_ddqn.wire_ofi_concat(
ord_concat_buf.raw_ptr(),
urg_concat_buf.raw_ptr(),
);
info!("GpuDqnTrainer: cuBLAS batched forward (DDQN) initialized (VSN+KAN+OFI wired)");
// ── Initialize cuBLAS backward context (required) ──────────
let cublas_backward = CublasBackwardSet::new(
@@ -3414,6 +3496,8 @@ impl GpuDqnTrainer {
save_h_b3: save_h_b3.raw_ptr(),
mag_concat_buf: mag_concat_buf.raw_ptr(),
d_mag_concat_buf: d_mag_concat_buf.raw_ptr(),
ord_concat_buf: ord_concat_buf.raw_ptr(),
urg_concat_buf: urg_concat_buf.raw_ptr(),
bw_d_h_s1: bw_d_h_s1.raw_ptr(),
bw_d_h_s2: bw_d_h_s2.raw_ptr(),
bw_d_h_v: bw_d_h_v.raw_ptr(),
@@ -3653,6 +3737,11 @@ impl GpuDqnTrainer {
mag_concat_kernel,
strided_accumulate_kernel,
strided_scatter_kernel,
ord_concat_buf,
d_ord_concat_buf,
urg_concat_buf,
d_urg_concat_buf,
concat_ofi_kernel,
q_attn_params,
q_attn_adam_m,
q_attn_adam_v,
@@ -4883,6 +4972,8 @@ impl GpuDqnTrainer {
let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, &param_sizes);
// Build mag_concat from previous step's logits (one-step lag)
self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.on_v_logits_buf, self.ptrs.on_b_logits_buf)?;
self.launch_concat_ofi(2, self.ptrs.states_buf)?;
self.launch_concat_ofi(3, self.ptrs.states_buf)?;
self.cublas_forward.forward_online_raw(
&self.stream, self.ptrs.states_buf, &on_w_ptrs,
self.ptrs.save_h_s1, self.ptrs.save_h_s2, self.ptrs.save_h_v,
@@ -4946,6 +5037,8 @@ impl GpuDqnTrainer {
let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, &param_sizes);
// Build mag_concat from previous step's logits (zeros at step 0)
self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.on_v_logits_buf, self.ptrs.on_b_logits_buf)?;
self.launch_concat_ofi(2, self.ptrs.states_buf)?;
self.launch_concat_ofi(3, self.ptrs.states_buf)?;
self.cublas_forward.forward_online_raw(
&self.stream, self.ptrs.states_buf, &on_w_ptrs,
self.ptrs.save_h_s1, self.ptrs.save_h_s2, self.ptrs.save_h_v,
@@ -5846,6 +5939,11 @@ impl GpuDqnTrainer {
// This call is captured in the CUDA graph — on replay it reads the logits
// from the previous graph execution (still in the buffer).
self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.on_v_logits_buf, self.ptrs.on_b_logits_buf)?;
// ── Build OFI concat for order (d=2) and urgency (d=3) branches (one-step lag).
// Uses prev-step vsn_masked (one-step lag, same as mag_concat). The vsn_masked
// portion is updated inside launch_vsn_glu_branch after VSN runs.
self.launch_concat_ofi(2, self.ptrs.states_buf)?;
self.launch_concat_ofi(3, self.ptrs.states_buf)?;
// ── Pass 1: Online forward on STATES (graph-safe: CachedPtrs, no device_ptr)
// When bottleneck is active, s1_input_ptr = bn_concat (compressed features).
@@ -5899,6 +5997,10 @@ impl GpuDqnTrainer {
// Build target mag_concat from previous step's target logits (one-step lag).
// Reuses the same mag_concat_buf — target forward runs after online forward.
self.launch_mag_concat_from(self.ptrs.save_h_s2, self.ptrs.tg_v_logits_buf, self.ptrs.tg_b_logits_buf)?;
// OFI concat for target forward (branches 2 and 3) — reuses ord/urg concat buffers
// with next_states features (target forward sees next_states).
self.launch_concat_ofi(2, self.ptrs.next_states_buf)?;
self.launch_concat_ofi(3, self.ptrs.next_states_buf)?;
cublas.forward_target_raw(
&self.stream, self.ptrs.next_states_buf, &tg_w_ptrs,
self.ptrs.tg_h_s1_scratch, self.ptrs.tg_h_s2_buf,
@@ -6485,6 +6587,10 @@ impl GpuDqnTrainer {
s1_dx_output,
self.ptrs.mag_concat_buf,
self.ptrs.d_mag_concat_buf,
self.ptrs.ord_concat_buf,
self.d_ord_concat_buf.raw_ptr(),
self.ptrs.urg_concat_buf,
self.d_urg_concat_buf.raw_ptr(),
&glu_gate_pre_ptrs,
&glu_value_ptrs,
self.ptrs.bw_d_glu_value,
@@ -6500,6 +6606,24 @@ impl GpuDqnTrainer {
1.0, // beta=1: d==0 already wrote to d_h_s2 in backward_full
)?;
}
// Accumulate order branch dX (first SH2 columns) into d_h_s2
if self.ptrs.ord_concat_buf != 0 {
self.accumulate_d_h_s2_from_concat(
self.d_ord_concat_buf.raw_ptr(),
d_h_s2_ptr,
self.config.batch_size,
1.0,
)?;
}
// Accumulate urgency branch dX (first SH2 columns) into d_h_s2
if self.ptrs.urg_concat_buf != 0 {
self.accumulate_d_h_s2_from_concat(
self.d_urg_concat_buf.raw_ptr(),
d_h_s2_ptr,
self.config.batch_size,
1.0,
)?;
}
// ── #31 Bottleneck backward: chain rule through tanh + GEMM ──
// After backward_full, d_bn_concat_buf has dL/d(bn_concat) [B, concat_dim] f32.
@@ -6845,11 +6969,11 @@ impl GpuDqnTrainer {
(0, 0), // [13] b_b1fc
(cfg.branch_1_size * cfg.num_atoms, cfg.adv_h), // [14] w_b1out
(0, 0), // [15] b_b1out
(cfg.adv_h, cfg.shared_h2), // [16] w_b2fc
(cfg.adv_h, cfg.shared_h2 + 3), // [16] w_b2fc (OFI-conditioned: SH2+3 input)
(0, 0), // [17] b_b2fc
(cfg.branch_2_size * cfg.num_atoms, cfg.adv_h), // [18] w_b2out
(0, 0), // [19] b_b2out
(cfg.adv_h, cfg.shared_h2), // [20] w_b3fc
(cfg.adv_h, cfg.shared_h2 + 3), // [20] w_b3fc (OFI-conditioned: SH2+3 input)
(0, 0), // [21] b_b3fc
(cfg.branch_3_size * cfg.num_atoms, cfg.adv_h), // [22] w_b3out
(0, 0), // [23] b_b3out
@@ -6906,6 +7030,27 @@ impl GpuDqnTrainer {
offset += align4(n);
}
// ── Zero the last 3 OFI columns of w_b2fc [16] and w_b3fc [20] ──
// xavier_init fills the full tensor (adv_h × (shared_h2+3)) but the OFI
// columns should start at zero so the network learns to use them gradually.
// (gate weights w_gate_2/w_gate_3 are already zero via (0,0) fan_dims.)
for &fc_idx in &[16_usize, 20_usize] {
let base: usize = sizes[..fc_idx].iter().map(|&s| align4(s)).sum();
let sh2 = cfg.shared_h2;
let sh2p3 = sh2 + 3;
let ah = cfg.adv_h;
// w_b2fc/w_b3fc layout: [AH, SH2+3] row-major.
// Zero columns [SH2, SH2+1, SH2+2] for all AH rows.
for row in 0..ah {
for col in sh2..sh2p3 {
let idx = base + row * sh2p3 + col;
if idx < host_buf.len() {
host_buf[idx] = 0.0_f32;
}
}
}
}
// ── KAN spline coefficient init: uniform 0.125 (flat 0.5 gate ≈ sigmoid(0)) ──
// Indices 42, 44, 46, 48 are kan_coeff_d [AH, 8]. Fill with 1/8 so the
// spline sums to 0.5 regardless of input (mimics sigmoid(0) = 0.5).