diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 7f5fd0f23..94f2ca3f8 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -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, config: &GpuDqnTrainConfig, - glu_backward_kernel: CudaFunction, + kan_gate_backward_kernel: CudaFunction, ) -> Result { 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, - 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 ──────────────────── diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 6ea0accd7..719de730b 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -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, - /// GLU combine kernel: output = sigmoid(gate_pre) * value. - glu_combine_kernel: Option, + /// KAN spline gate kernel: output = kan_spline(gate_pre, coeffs, resid_w) * value. + kan_gate_combine_kernel: Option, /// Strided scatter: vsn_masked [B, SH2] → first SH2 cols of mag_concat [B, SH2+3]. strided_scatter_kernel: Option, /// 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}" )))?; } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 372bf687b..ab6778b81 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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; 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, // [B, NUM_BRANCHES(4), NUM_ATOMS] save_projected: CudaSlice, // [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(¶m_sizes, 8); - // End offset: after last branch weight (index 41, b_gate_3) - let end_byte = padded_byte_offset(¶m_sizes, 42); // one past last + // End offset: after last branch weight (index 49, kan_resid_3) + let end_byte = padded_byte_offset(¶m_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).