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

feat: KAN spline gates replace sigmoid — learned activation per neuron (NUM_WEIGHT_TENSORS 42->50)

Browse Source 
8 cubic B-spline basis functions per gate neuron + residual connection.
Spline coefficients initialized to uniform 0.125 (flat 0.5 gate).
kan_gate_combine replaces glu_combine in forward.
kan_gate_backward replaces glu_backward in backward.
9,216 new parameters (4 x AH x 9).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 09:13:49 +02:00
parent 454612989d
commit cd7c67f541
3 changed files with 137 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -155,10 +155,11 @@ pub struct CublasBackwardSet {
/// `bias_grad_reduce_f32_kernel(dy, db, out_dim, batch_size)` — reduce f32 dY over batch.
bias_grad_kernel: CudaFunction,
/// `glu_gate_backward(d_output, gate_pre, value, d_gate_pre, d_value, n)` —
/// GLU backward pass: splits d_h_bd into d_value and d_gate_pre components.
/// Replaces ReLU mask in the branch backward pass.
glu_backward_kernel: CudaFunction,
/// `kan_gate_backward(d_output, gate_pre, value, spline_coeff, residual_w,
/// d_gate_pre, d_value, d_spline_coeff, d_residual_w, n, adv_h)` —
/// KAN spline backward pass: splits d_h_bd into d_value and d_gate_pre,
/// and accumulates coefficient/residual gradients via atomicAdd.
kan_gate_backward_kernel: CudaFunction,
// ── Network dimensions (baked at construction) ──
batch_size: usize,
@@ -177,7 +178,7 @@ pub struct CublasBackwardSet {
branch_2_size: usize,
branch_3_size: usize,
/// Pre-computed parameter sizes for all 42 weight tensors.
/// Pre-computed parameter sizes for all 50 weight tensors.
/// Used to compute `padded_byte_offset` for gradient buffer offsets.
param_sizes: [usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
@@ -201,7 +202,7 @@ impl CublasBackwardSet {
pub fn new(
shared: Arc<super::shared_cublas_handle::SharedCublasHandle>,
config: &GpuDqnTrainConfig,
glu_backward_kernel: CudaFunction,
kan_gate_backward_kernel: CudaFunction,
) -> Result<Self, MLError> {
let stream = &shared.stream;
let lt_raw_handle = shared.lt_handle.0;
@@ -286,7 +287,7 @@ impl CublasBackwardSet {
handle: shared,
relu_mask_kernel,
bias_grad_kernel,
glu_backward_kernel,
kan_gate_backward_kernel,
batch_size: config.batch_size,
state_dim: config.state_dim,
state_dim_padded: sd_pad,
@@ -671,44 +672,58 @@ impl CublasBackwardSet {
Ok(())
}
/// GLU backward: split upstream gradient `d_output` into `d_value` and
/// `d_gate_pre` using saved forward activations.
/// KAN spline gate backward: split upstream gradient `d_output` into `d_value`
/// and `d_gate_pre`, and accumulate coefficient/residual gradients.
///
/// ```text
/// d_value[i] = d_output[i] * sigmoid(gate_pre[i])
/// d_gate_pre[i] = d_output[i] * value[i] * sigmoid(gate_pre[i]) * (1 - sigmoid(gate_pre[i]))
/// gate = clamp(sum_k coeff[k]*B_k(pre) + resid_w*pre, 0, 1)
/// d_value[i] = d_output[i] * gate
/// d_gate_pre[i] = d_output[i] * value[i] * clamp_mask * (dB/dx + resid_w)
/// d_coeff[k] += d_output[i] * value[i] * clamp_mask * B_k(pre) (atomicAdd)
/// d_resid_w += d_output[i] * value[i] * clamp_mask * pre (atomicAdd)
/// ```
///
/// Replaces `relu_mask` for branch FC layers that use GLU gating.
/// Grid: `ceil(n / 256)`, Block: 256.
pub fn launch_glu_backward(
#[allow(clippy::too_many_arguments)]
pub fn launch_kan_gate_backward(
&self,
stream: &Arc<CudaStream>,
d_output: u64, // [B, AH] upstream gradient from branch output backward
gate_pre: u64, // [B, AH] saved pre-sigmoid gate activation
value: u64, // [B, AH] saved value path output
d_gate_pre: u64, // [B, AH] output: gradient for gate path
d_value: u64, // [B, AH] output: gradient for value path
d_output: u64, // [B, AH] upstream gradient from branch output backward
gate_pre: u64, // [B, AH] saved pre-activation gate input
value: u64, // [B, AH] saved value path output
spline_coeff: u64, // [AH, 8] current spline coefficients (from params_buf)
residual_w: u64, // [AH] current residual weights (from params_buf)
d_gate_pre: u64, // [B, AH] output: gradient for gate path (upstream to W_gate)
d_value: u64, // [B, AH] output: gradient for value path (upstream to W_bdf)
d_spline_coeff: u64, // [AH, 8] output: accumulated coeff gradients (in grad_buf)
d_residual_w: u64, // [AH] output: accumulated residual weight gradients (in grad_buf)
n: usize,
adv_h: usize,
) -> Result<(), MLError> {
let n_i32 = n as i32;
let adv_h_i32 = adv_h as i32;
let blocks = ((n + 255) / 256) as u32;
unsafe {
stream
.launch_builder(&self.glu_backward_kernel)
.launch_builder(&self.kan_gate_backward_kernel)
.arg(&d_output)
.arg(&gate_pre)
.arg(&value)
.arg(&spline_coeff)
.arg(&residual_w)
.arg(&d_gate_pre)
.arg(&d_value)
.arg(&d_spline_coeff)
.arg(&d_residual_w)
.arg(&n_i32)
.arg(&adv_h_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("glu_backward_kernel: {e}")))?;
.map_err(|e| MLError::ModelError(format!("kan_gate_backward_kernel: {e}")))?;
}
Ok(())
@@ -747,7 +762,7 @@ impl CublasBackwardSet {
/// - `save_h_s2` : `[B, SH2]` saved shared layer 2 ReLU output
/// - `save_h_v` : `[B, VH]` saved value FC ReLU output
/// - `save_h_b` : four `[B, AH]` saved GLU output (post-gating activation)
/// - `w_ptrs` : 42-element array of raw F32 device pointers into
/// - `w_ptrs` : 50-element array of raw F32 device pointers into
/// `params_buf` (online weights, read-only)
/// - `grad_buf` : flat gradient accumulator `[TOTAL_PARAMS]` — must be
/// zeroed before this call
@@ -873,6 +888,8 @@ impl CublasBackwardSet {
// [20]=w_b3fc, [21]=b_b3fc, [22]=w_b3out, [23]=b_b3out
// [34]=w_gate_0, [35]=b_gate_0, [36]=w_gate_1, [37]=b_gate_1,
// [38]=w_gate_2, [39]=b_gate_2, [40]=w_gate_3, [41]=b_gate_3
// [42]=kan_coeff_0, [43]=kan_resid_0, [44]=kan_coeff_1, [45]=kan_resid_1,
// [46]=kan_coeff_2, [47]=kan_resid_2, [48]=kan_coeff_3, [49]=kan_resid_3
let w_bfc_idx = [8_usize, 12, 16, 20];
let w_bout_idx = [10_usize, 14, 18, 22];
// GLU gate weight indices: w_gate_d at 34+d*2, b_gate_d at 34+d*2+1
@@ -892,8 +909,22 @@ impl CublasBackwardSet {
padded_byte_offset(&self.param_sizes, 41),
];
// KAN spline coefficient gradient offsets (in grad_buf)
let goff_kan_coeff = [
padded_byte_offset(&self.param_sizes, 42),
padded_byte_offset(&self.param_sizes, 44),
padded_byte_offset(&self.param_sizes, 46),
padded_byte_offset(&self.param_sizes, 48),
];
let goff_kan_resid = [
padded_byte_offset(&self.param_sizes, 43),
padded_byte_offset(&self.param_sizes, 45),
padded_byte_offset(&self.param_sizes, 47),
padded_byte_offset(&self.param_sizes, 49),
];
// ══════════════════════════════════════════════════════════════════
// BRANCHES (all 4, with GLU backward replacing ReLU mask)
// BRANCHES (all 4, with KAN spline backward replacing ReLU mask)
// ══════════════════════════════════════════════════════════════════
for d in 0..4 {
@@ -902,9 +933,13 @@ impl CublasBackwardSet {
let w_fc = w_ptrs[w_bfc_idx[d]]; // W_bdk_fc[AH, SH2] (value path weight)
let w_gate = w_ptrs[w_gate_idx[d]]; // W_gate_d[AH, SH2] (gate path weight)
// KAN spline weight pointers (from params_buf, read-only)
let kan_coeff_idx = 42 + d * 2; // kan_coeff_d
let kan_resid_idx = 42 + d * 2 + 1; // kan_resid_d
// ── Step 1: Branch output layer (no activation — logit layer) ──
// dY = d_adv_logits[d] [B, n_d]
// X = save_h_b[d] [B, AH] (GLU output)
// X = save_h_b[d] [B, AH] (KAN gate output)
// W = W_bdk_out [n_d, AH]
// dW += dY^T @ X, db += sum(dY), dX_b = dY @ W^T → scratch_d_h_b[d]
self.backward_fc_layer(
@@ -914,24 +949,27 @@ impl CublasBackwardSet {
w_out, // W [n_d, AH]
grad_buf_base + goff_w_bout[d], // dW
grad_buf_base + goff_b_bout[d], // db
scratch_d_h_b[d], // dX [B, AH] → d_h_bd (gradient into GLU output)
scratch_d_h_b[d], // dX [B, AH] → d_h_bd (gradient into KAN gate output)
n_d, // out_dim
self.adv_h, // in_dim
b,
)?;
// ── Step 2: GLU backward (replaces ReLU mask) ──────────────
// d_h_bd → d_value + d_gate_pre using saved gate_pre and value
// d_value[i] = d_h_bd[i] * sigmoid(gate_pre[i])
// d_gate_pre[i] = d_h_bd[i] * value[i] * sigmoid(gate_pre[i]) * (1 - sigmoid(gate_pre[i]))
self.launch_glu_backward(
// ── Step 2: KAN spline backward (replaces GLU backward) ──────
// d_h_bd → d_value + d_gate_pre + d_coeff + d_resid_w
self.launch_kan_gate_backward(
stream,
scratch_d_h_b[d], // d_output = d_h_bd [B, AH]
glu_gate_pre[d], // saved gate_pre [B, AH]
glu_value[d], // saved value [B, AH]
w_ptrs[kan_coeff_idx], // spline_coeff [AH, 8] (current weights)
w_ptrs[kan_resid_idx], // residual_w [AH] (current weights)
scratch_d_glu_gate, // d_gate_pre output [B, AH]
scratch_d_glu_value, // d_value output [B, AH]
grad_buf_base + goff_kan_coeff[d], // d_spline_coeff [AH, 8] (accumulated in grad_buf)
grad_buf_base + goff_kan_resid[d], // d_residual_w [AH] (accumulated in grad_buf)
b * self.adv_h,
self.adv_h,
)?;
// ── Step 3: Value path weight gradients ────────────────────

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

@@ -177,8 +177,8 @@ pub struct CublasGemmSet {
// ── VSN bottleneck + GLU gating (wired post-construction via wire_vsn_glu) ──
/// Variable Selection Network bottleneck kernel (SH2→R=16→SH2).
vsn_kernel: Option<CudaFunction>,
/// GLU combine kernel: output = sigmoid(gate_pre) * value.
glu_combine_kernel: Option<CudaFunction>,
/// KAN spline gate kernel: output = kan_spline(gate_pre, coeffs, resid_w) * value.
kan_gate_combine_kernel: Option<CudaFunction>,
/// Strided scatter: vsn_masked [B, SH2] → first SH2 cols of mag_concat [B, SH2+3].
strided_scatter_kernel: Option<CudaFunction>,
/// VSN masked output scratch [B, SH2].
@@ -341,9 +341,9 @@ impl CublasGemmSet {
branch_done_events,
gemm_cache,
gemm_cache_relu_bias,
// VSN/GLU: not wired until wire_vsn_glu() is called
// VSN/KAN: not wired until wire_vsn_glu() is called
vsn_kernel: None,
glu_combine_kernel: None,
kan_gate_combine_kernel: None,
strided_scatter_kernel: None,
vsn_masked_ptr: 0,
glu_gate_pre_ptrs: [0; 4],
@@ -365,7 +365,7 @@ impl CublasGemmSet {
pub fn wire_vsn_glu(
&mut self,
vsn_kernel: CudaFunction,
glu_combine_kernel: CudaFunction,
kan_gate_combine_kernel: CudaFunction,
strided_scatter_kernel: CudaFunction,
vsn_masked_ptr: u64,
glu_gate_pre_ptrs: [u64; 4],
@@ -373,7 +373,7 @@ impl CublasGemmSet {
vsn_rank: usize,
) {
self.vsn_kernel = Some(vsn_kernel);
self.glu_combine_kernel = Some(glu_combine_kernel);
self.kan_gate_combine_kernel = Some(kan_gate_combine_kernel);
self.strided_scatter_kernel = Some(strided_scatter_kernel);
self.vsn_masked_ptr = vsn_masked_ptr;
self.glu_gate_pre_ptrs = glu_gate_pre_ptrs;
@@ -936,7 +936,7 @@ impl CublasGemmSet {
/// For others: vsn_input = vsn_masked
/// 3. Value GEMM: W_bdf @ vsn_input → glu_value [B, AH], add bias (no ReLU)
/// 4. Gate GEMM: W_gate @ vsn_input → glu_gate_pre [B, AH], add gate bias
/// 5. GLU combine: h_bd = sigmoid(gate_pre) * value
/// 5. KAN gate: h_bd = kan_spline(gate_pre, coeffs, resid_w) * value
///
/// `ws_ptr`/`ws_size`: cublasLt workspace for this stream (branch or main).
/// `is_branch_stream`: if true, uses `sgemm_f32_branch`; if false, `sgemm_f32`.
@@ -955,7 +955,7 @@ impl CublasGemmSet {
_label_prefix: &str,
) -> Result<(), MLError> {
let vsn_k = self.vsn_kernel.as_ref().unwrap();
let glu_k = self.glu_combine_kernel.as_ref().unwrap();
let kan_k = self.kan_gate_combine_kernel.as_ref().unwrap();
let b = self.batch_size;
let sh2 = self.shared_h2;
let ah = self.adv_h;
@@ -1048,23 +1048,29 @@ impl CublasGemmSet {
}
self.launch_add_bias_f32_raw(stream, gate_ptr, w_ptrs[b_gate_idx], ah, b)?;
// ── 5. GLU combine: h_bd = sigmoid(gate_pre) * value ──
// ── 5. KAN gate: h_bd = kan_spline(gate_pre, coeffs, resid_w) * value ──
let kan_coeff_idx = 42 + d * 2; // kan_coeff_d
let kan_resid_idx = 42 + d * 2 + 1; // kan_resid_d
let n_total = (b * ah) as i32;
let num_neurons = ah as i32;
let blocks = ((b * ah + 255) / 256) as u32;
unsafe {
stream
.launch_builder(glu_k)
.launch_builder(kan_k)
.arg(&branch_h_ptr)
.arg(&gate_ptr)
.arg(&value_ptr)
.arg(&w_ptrs[kan_coeff_idx])
.arg(&w_ptrs[kan_resid_idx])
.arg(&n_total)
.arg(&num_neurons)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"{_label_prefix} glu_combine branch {d}: {e}"
"{_label_prefix} kan_gate_combine branch {d}: {e}"
)))?;
}

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

@@ -351,8 +351,8 @@ impl Default for CausalInterventionConfig {
/// Number of weight tensors in the flat parameter buffer.
/// 20 original (DQN network) + 4 branch 3 (magnitude) + 2 bottleneck (w_bn, b_bn)
/// + 8 VSN bottleneck (R=16) + 8 GLU gate = 42.
pub(crate) const NUM_WEIGHT_TENSORS: usize = 42;
/// + 8 VSN bottleneck (R=16) + 8 GLU gate + 8 KAN spline (coeff+resid per branch) = 50.
pub(crate) const NUM_WEIGHT_TENSORS: usize = 50;
/// Compute the size (element count) of each weight tensor.
///
@@ -361,6 +361,7 @@ pub(crate) const NUM_WEIGHT_TENSORS: usize = 42;
/// Tensors 24-25: temporal causal bottleneck (0 elements when bottleneck_dim=0).
/// Tensors 26-33: Variable Selection Network bottleneck (R=16).
/// Tensors 34-41: GLU gate weights and biases.
/// Tensors 42-49: KAN spline coefficients and residual weights (per branch).
///
/// When bottleneck is active, w_s1 input dimension changes from state_dim to
/// (bottleneck_dim + portfolio_dim) where portfolio_dim = state_dim - market_dim.
@@ -421,6 +422,15 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT
cfg.adv_h, // [39] b_gate_2
cfg.adv_h * cfg.shared_h2, // [40] w_gate_3
cfg.adv_h, // [41] b_gate_3
// ── KAN spline coefficients (8 bases per neuron) + residual weight ──
cfg.adv_h * 8, // [42] kan_coeff_0 [AH, 8]
cfg.adv_h, // [43] kan_resid_0 [AH]
cfg.adv_h * 8, // [44] kan_coeff_1 [AH, 8]
cfg.adv_h, // [45] kan_resid_1 [AH]
cfg.adv_h * 8, // [46] kan_coeff_2 [AH, 8]
cfg.adv_h, // [47] kan_resid_2 [AH]
cfg.adv_h * 8, // [48] kan_coeff_3 [AH, 8]
cfg.adv_h, // [49] kan_resid_3 [AH]
]
}
@@ -662,6 +672,8 @@ pub struct GpuDqnTrainer {
glu_value_buf: [CudaSlice<f32>; 4], // 4 × [B, AH]
glu_combine_kernel: CudaFunction,
glu_backward_kernel: CudaFunction,
kan_gate_combine_kernel: CudaFunction,
kan_gate_backward_kernel: CudaFunction,
save_current_lp: CudaSlice<f32>, // [B, NUM_BRANCHES(4), NUM_ATOMS]
save_projected: CudaSlice<f32>, // [B, NUM_BRANCHES(4), NUM_ATOMS]
@@ -1076,15 +1088,15 @@ impl GpuDqnTrainer {
Ok(())
}
/// Reset Adam momentum (m/v) for branch weights only (indices 8-41).
/// Reset Adam momentum (m/v) for branch weights only (indices 8-49).
/// Trunk momentum (indices 0-7) preserved — carries valuable convergence history.
/// Lighter than full rewind — tries to break Adam fixed point without changing weights.
pub(crate) fn reset_branch_adam_momentum(&mut self) -> Result<(), MLError> {
let param_sizes = compute_param_sizes(&self.config);
// Start offset: index 8 (w_b0fc)
let start_byte = padded_byte_offset(&param_sizes, 8);
// End offset: after last branch weight (index 41, b_gate_3)
let end_byte = padded_byte_offset(&param_sizes, 42); // one past last
// End offset: after last branch weight (index 49, kan_resid_3)
let end_byte = padded_byte_offset(&param_sizes, 50); // one past last
let range_bytes = (end_byte - start_byte) as usize;
// Zero the m_buf and v_buf ranges for branch weights only
@@ -2960,7 +2972,11 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("glu_combine load: {e}")))?;
let glu_backward_kernel = cpbi_module.load_function("glu_backward")
.map_err(|e| MLError::ModelError(format!("glu_backward load: {e}")))?;
info!("GpuDqnTrainer: Q-attn + selectivity + VSN + GLU kernels loaded");
let kan_gate_combine_kernel = cpbi_module.load_function("kan_gate_combine")
.map_err(|e| MLError::ModelError(format!("kan_gate_combine load: {e}")))?;
let kan_gate_backward_kernel = cpbi_module.load_function("kan_gate_backward")
.map_err(|e| MLError::ModelError(format!("kan_gate_backward load: {e}")))?;
info!("GpuDqnTrainer: Q-attn + selectivity + VSN + GLU + KAN kernels loaded");
// ── Compile CQL penalty kernel (if enabled) ──────────────────────
let cql_logit_grad_kernel = if config.cql_alpha > 0.0 {
@@ -3164,10 +3180,10 @@ impl GpuDqnTrainer {
config.branch_3_size,
s1_input_dim,
)?;
// Wire VSN bottleneck + GLU gating into branch forward paths
// Wire VSN bottleneck + KAN spline gating into branch forward paths
cublas_forward.wire_vsn_glu(
vsn_kernel.clone(),
glu_combine_kernel.clone(),
kan_gate_combine_kernel.clone(),
strided_scatter_kernel.clone(),
vsn_masked_buf.raw_ptr(),
[
@@ -3180,7 +3196,7 @@ impl GpuDqnTrainer {
],
16, // VSN bottleneck rank R=16
);
info!("GpuDqnTrainer: cuBLAS batched forward initialized (VSN+GLU wired)");
info!("GpuDqnTrainer: cuBLAS batched forward initialized (VSN+KAN wired)");
// DDQN forward shares the same cuBLAS handle — runs sequentially
// on the main stream for NVIDIA bit-wise reproducibility.
@@ -3199,11 +3215,11 @@ impl GpuDqnTrainer {
config.branch_3_size,
s1_input_dim,
)?;
// DDQN argmax pass also uses VSN+GLU — same kernels, same scratch buffers
// DDQN argmax pass also uses VSN+KAN — same kernels, same scratch buffers
// (DDQN runs sequentially on main stream, no concurrent access).
cublas_forward_ddqn.wire_vsn_glu(
vsn_kernel.clone(),
glu_combine_kernel.clone(),
kan_gate_combine_kernel.clone(),
strided_scatter_kernel.clone(),
vsn_masked_buf.raw_ptr(),
[
@@ -3216,13 +3232,13 @@ impl GpuDqnTrainer {
],
16,
);
info!("GpuDqnTrainer: cuBLAS batched forward (DDQN) initialized (VSN+GLU wired)");
info!("GpuDqnTrainer: cuBLAS batched forward (DDQN) initialized (VSN+KAN wired)");
// ── Initialize cuBLAS backward context (required) ──────────
let cublas_backward = CublasBackwardSet::new(
Arc::clone(&shared_cublas), &config, glu_backward_kernel.clone(),
Arc::clone(&shared_cublas), &config, kan_gate_backward_kernel.clone(),
)?;
info!("GpuDqnTrainer: cuBLAS batched backward initialized (GLU backward wired)");
info!("GpuDqnTrainer: cuBLAS batched backward initialized (KAN backward wired)");
// ── Backward scratch buffers ────────────────────────────────
// Pre-allocate inter-layer gradient buffers for the cuBLAS backward
@@ -3534,6 +3550,8 @@ impl GpuDqnTrainer {
glu_value_buf,
glu_combine_kernel,
glu_backward_kernel,
kan_gate_combine_kernel,
kan_gate_backward_kernel,
save_current_lp,
save_projected,
per_sample_loss_buf,
@@ -6709,6 +6727,14 @@ impl GpuDqnTrainer {
(0, 0), // [39] b_gate_2
(0, 0), // [40] w_gate_3
(0, 0), // [41] b_gate_3
(0, 0), // [42] kan_coeff_0 — special init (sigmoid approx)
(0, 0), // [43] kan_resid_0 — zero
(0, 0), // [44] kan_coeff_1
(0, 0), // [45] kan_resid_1
(0, 0), // [46] kan_coeff_2
(0, 0), // [47] kan_resid_2
(0, 0), // [48] kan_coeff_3
(0, 0), // [49] kan_resid_3
];
// Build flat host buffer: Xavier init for weights, zeros for biases + padding.
@@ -6735,6 +6761,19 @@ impl GpuDqnTrainer {
// Biases and zero-size tensors stay zero from vec init
offset += align4(n);
}
// ── 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).
// Residual weights (43, 45, 47, 49) stay zero.
for &coeff_idx in &[42_usize, 44, 46, 48] {
let coeff_offset: usize = sizes[..coeff_idx].iter().map(|&s| align4(s)).sum();
let n_coeff = sizes[coeff_idx];
for w in &mut host_buf[coeff_offset..coeff_offset + n_coeff] {
*w = 0.125_f32;
}
}
debug_assert_eq!(offset, total, "xavier_init: offset mismatch with total_params");
// HtoD to both params_buf and target_params_buf (same initial weights).