diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 5875ad526..a17971240 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -498,6 +498,66 @@ extern "C" __global__ void bn_tanh_concat_kernel( } } +/* ══════════════════════════════════════════════════════════════════════ + * CAUSAL INTERVENTION KERNEL (#34) + * + * Copies states to scratch, zeros feature k, all on GPU. + * Used for causal sensitivity analysis — measures how Q-values change + * when a specific feature is set to 0 (do-calculus intervention). + * + * Grid: ceil(B / 256), Block: 256. One thread per batch sample. + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void causal_intervene_feature( + const __nv_bfloat16* __restrict__ states, /* [B, state_dim_padded] source */ + __nv_bfloat16* __restrict__ scratch, /* [B, state_dim_padded] output */ + int feature_k, /* which feature to zero */ + int state_dim_padded, /* row stride */ + int batch_size +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + + long long base = (long long)b * state_dim_padded; + for (int j = 0; j < state_dim_padded; j++) { + scratch[base + j] = (j == feature_k) ? bf16_zero() : states[base + j]; + } +} + +/** + * Compute |Q_original - Q_intervened|² summed over batch for a single feature. + * Reads two f32 logit buffers, reduces to a single scalar per feature. + * + * Grid: (1, 1, 1), Block: (256, 1, 1). + */ +extern "C" __global__ void causal_q_delta_reduce( + const float* __restrict__ q_original, /* [B, num_atoms] */ + const float* __restrict__ q_intervened, /* [B, num_atoms] */ + float* __restrict__ sensitivity_out, /* [market_dim] — atomicAdd per feature */ + int feature_k, + int total_elements /* B * num_atoms */ +) { + __shared__ float s_sum[256]; + int tid = threadIdx.x; + float local_sum = 0.0f; + + for (int i = tid; i < total_elements; i += 256) { + float d = q_original[i] - q_intervened[i]; + local_sum += d * d; + } + + s_sum[tid] = local_sum; + __syncthreads(); + + for (int s = 128; s > 0; s >>= 1) { + if (tid < s) s_sum[tid] += s_sum[tid + s]; + __syncthreads(); + } + + if (tid == 0) { + atomicAdd(&sensitivity_out[feature_k], s_sum[0] / (float)total_elements); + } +} + /* ══════════════════════════════════════════════════════════════════════ * GRADIENT VACCINE KERNELS (#32) * @@ -637,13 +697,42 @@ extern "C" __global__ void bn_bias_grad_kernel( } /* ══════════════════════════════════════════════════════════════════════ - * STOCHASTIC DEPTH KERNEL (#21) + * STOCHASTIC DEPTH RNG KERNEL (#21) + * + * Generates per-layer drop/keep scales ENTIRELY on GPU using LCG RNG. + * Replaces CPU rand::thread_rng() + memcpy_htod that ran every training step. + * + * For each of 3 layers: draw uniform [0,1), compare with drop_prob. + * If uniform > p: scale = 1/(1-p) (keep, expected-value correction). + * If uniform <= p: scale = 0 (drop this layer). + * + * Single thread — 3 layers is trivially fast, no parallelism needed. + * Grid: (1, 1, 1), Block: (1, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void stochastic_depth_rng( + float* __restrict__ scale_buf, /* [3] output: per-layer scales */ + unsigned int* __restrict__ rng_state, /* [1] LCG state (read-write) */ + float drop_prob /* probability of dropping each layer */ +) { + unsigned int rng = rng_state[0]; + float keep_scale = 1.0f / (1.0f - drop_prob); + + for (int layer = 0; layer < 3; layer++) { + rng = rng * 1664525u + 1013904223u; + float u = (float)(rng & 0x00FFFFFFu) / 16777216.0f; + scale_buf[layer] = (u > drop_prob) ? keep_scale : 0.0f; + } + rng_state[0] = rng; +} + +/* ══════════════════════════════════════════════════════════════════════ + * STOCHASTIC DEPTH SCALE KERNEL (#21) * * Scales a bf16 hidden activation buffer by a per-layer scalar. * Used for stochastic depth: each hidden layer is either kept (scaled * by 1/(1-p) for expected-value correction) or dropped (scaled by 0). * - * The scale buffer [num_layers] is written by the host BEFORE each + * The scale buffer [3] is written by stochastic_depth_rng BEFORE each * CUDA Graph replay. The graph reads the updated values — this is legal * because CUDA Graphs capture buffer addresses, not contents. * diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 2f3e56b07..5792e7d09 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -620,6 +620,10 @@ pub struct GpuDqnTrainer { causal_q_scratch: Option>, /// #34 Causal intervention config. pub(crate) causal_config: CausalInterventionConfig, + /// #34 GPU-native causal intervention kernel. + causal_intervene_kernel: Option, + /// #34 GPU-native causal Q-delta reduction kernel. + causal_reduce_kernel: Option, /// #32 Gradient vaccine: saved g_train gradient [total_params] f32. vaccine_grad_save: CudaSlice, @@ -648,10 +652,14 @@ pub struct GpuDqnTrainer { bn_bias_grad_kernel: Option, /// #21 Stochastic depth: per-layer scale buffer [3] (h_s1, h_s2, h_v). - /// Written by host before each graph replay. 0.0=drop, 1/(1-p)=keep. + /// Written by GPU RNG kernel before each graph replay. stochastic_depth_scale_buf: CudaSlice, - /// #21 Stochastic depth kernel function. + /// #21 Stochastic depth scale kernel function (applies scale to activations). stochastic_depth_kernel: CudaFunction, + /// #21 Stochastic depth RNG kernel (generates scales on GPU, zero CPU RNG). + stochastic_depth_rng_kernel: CudaFunction, + /// #21 GPU-resident RNG state for stochastic depth [1] u32. + stochastic_depth_rng_state: CudaSlice, /// #21 Stochastic depth drop probability (0.0=disabled, 0.2=20% drop). stochastic_depth_prob: f32, @@ -2038,7 +2046,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─ - let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, stochastic_depth_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel) = + let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, causal_intervene_kernel_fn, causal_reduce_kernel_fn) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace). @@ -2384,11 +2392,17 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("alloc drawdown_depths_buf: {e}")))?; let ensemble_std_buf = stream.alloc_zeros::(b) .map_err(|e| MLError::ModelError(format!("alloc ensemble_std_buf: {e}")))?; - // #21 Stochastic depth: 3 layer scales (h_s1, h_s2, h_v). Init to 1.0 (keep all). + // #21 Stochastic depth: 3 layer scales + GPU RNG state let mut stochastic_depth_scale_buf = stream.alloc_zeros::(3) .map_err(|e| MLError::ModelError(format!("alloc stochastic_depth_scale: {e}")))?; stream.memcpy_htod(&[1.0_f32, 1.0, 1.0], &mut stochastic_depth_scale_buf) .map_err(|e| MLError::ModelError(format!("init stochastic_depth_scale: {e}")))?; + let mut stochastic_depth_rng_state = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("alloc sd_rng_state: {e}")))?; + // Seed with process ID + timestamp for uniqueness + let sd_seed = (std::process::id() as u32).wrapping_mul(2654435761); + stream.memcpy_htod(&[sd_seed], &mut stochastic_depth_rng_state) + .map_err(|e| MLError::ModelError(format!("seed sd_rng: {e}")))?; let curiosity_inference_func = { static CURIOSITY_INFERENCE_CUBIN: &[u8] = include_bytes!( concat!(env!("OUT_DIR"), "/curiosity_inference_kernel.cubin") @@ -2635,6 +2649,8 @@ impl GpuDqnTrainer { causal_states_scratch, causal_sensitivity_buf, causal_q_scratch, + causal_intervene_kernel: if causal_enabled { Some(causal_intervene_kernel_fn) } else { None }, + causal_reduce_kernel: if causal_enabled { Some(causal_reduce_kernel_fn) } else { None }, causal_config: CausalInterventionConfig { enabled: causal_enabled, weight: causal_wt, @@ -2655,6 +2671,8 @@ impl GpuDqnTrainer { bn_bias_grad_kernel: if bn_dim > 0 { Some(bn_bias_grad_kernel) } else { None }, stochastic_depth_scale_buf, stochastic_depth_kernel, + stochastic_depth_rng_kernel, + stochastic_depth_rng_state, stochastic_depth_prob: sd_prob, ensemble_std_buf, ensemble_disagreement_weight: ens_weight, @@ -2918,13 +2936,15 @@ impl GpuDqnTrainer { /// Replay graph_forward only (zero → forward → loss → grad → backward). /// After this call, grad_buf contains C51 gradients. The caller can inject /// auxiliary gradients (IQN, attention, ensemble) before calling replay_adam(). - /// #34 Causal intervention: compute per-feature sensitivity by running - /// intervened forward passes. For each feature k, set it to 0 in a copy of - /// the states, run forward, and measure how much Q-values change. - /// Features with high sensitivity are causal; low sensitivity = noise. + /// #34 Causal intervention: compute per-feature sensitivity on GPU. + /// Zero CPU-side computation — all interventions, forward passes, and + /// reductions happen via CUDA kernels. /// - /// Returns the mean sensitivity across features (for logging). - /// Modifies grad_buf to add causal regularization gradient. + /// For each feature k: GPU kernel zeros feature k in scratch states, + /// cuBLAS forward produces intervened Q-values, GPU reduction kernel + /// computes |Q_original - Q_intervened|² and accumulates sensitivity. + /// + /// Returns mean sensitivity (read back once at the end, not per-feature). pub fn run_causal_intervention(&mut self, step: usize) -> Result { if !self.causal_config.enabled || step % self.causal_config.interval != 0 { return Ok(0.0); @@ -2933,10 +2953,9 @@ impl GpuDqnTrainer { let b = self.config.batch_size; let market_dim = self.causal_config.market_dim; let pad_sd = (self.config.state_dim + 127) & !127; - let total_actions = self.config.branch_0_size + self.config.branch_1_size + self.config.branch_2_size; + let na = self.config.num_atoms; + let val_size = (b * na) as i32; - let scratch_states = self.causal_states_scratch.as_mut() - .ok_or_else(|| MLError::ModelError("causal_states_scratch not allocated".into()))?; let sensitivity = self.causal_sensitivity_buf.as_mut() .ok_or_else(|| MLError::ModelError("causal_sensitivity_buf not allocated".into()))?; @@ -2944,123 +2963,90 @@ impl GpuDqnTrainer { self.stream.memset_zeros(sensitivity) .map_err(|e| MLError::ModelError(format!("zero causal_sens: {e}")))?; - // Get baseline Q-values from the current forward pass (already computed by graph). - // The on_v_logits_buf + on_b_logits_buf contain the online network's output. - // We use these as the "original" Q-values. - // - // For each feature k in [0..market_dim): - // 1. Copy states_buf → scratch_states - // 2. Zero feature k in scratch_states (do(X_k = 0)) - // 3. Forward pass on scratch_states - // 4. Compute |Q_original - Q_intervened|² summed over batch and actions - // 5. sensitivity[k] = mean |delta_Q|² - let param_sizes = compute_param_sizes(&self.config); let on_w_ptrs = bf16_weight_ptrs_from_base(self.ptrs.params_buf, ¶m_sizes); - - // Baseline expected Q: already in the logits buffers from the graph forward. - // We only need the expected Q-values, not the full distributional output. - // For causal sensitivity, we compare E[Q] between original and intervened. - let states_ptr = self.ptrs.states_buf; - let scratch_ptr = scratch_states.raw_ptr(); - let state_bytes = b * pad_sd * std::mem::size_of::(); + let scratch_ptr = self.causal_states_scratch.as_ref().unwrap().raw_ptr(); + let sensitivity_ptr = sensitivity.raw_ptr(); - // Scratch activation buffers (reuse existing ones — causal intervention runs - // AFTER the graph forward, so saved activations are no longer needed). - let h_s1_ptr = self.ptrs.save_h_s1; - let h_s2_ptr = self.ptrs.save_h_s2; - let h_v_ptr = self.ptrs.save_h_v; - let h_b0_ptr = self.ptrs.save_h_b0; - let h_b1_ptr = self.ptrs.save_h_b1; - let h_b2_ptr = self.ptrs.save_h_b2; + let intervene_kernel = self.causal_intervene_kernel.as_ref() + .ok_or_else(|| MLError::ModelError("causal_intervene_kernel not loaded".into()))?; + let reduce_kernel = self.causal_reduce_kernel.as_ref() + .ok_or_else(|| MLError::ModelError("causal_reduce_kernel not loaded".into()))?; - let mut total_sensitivity = 0.0_f32; + let b_i32 = b as i32; + let pad_sd_i32 = pad_sd as i32; + let blocks = ((b as u32 + 255) / 256) as u32; - for k in 0..market_dim.min(14) { // Only active features (14 of 42) - // Step 1: Copy states → scratch - dtod_copy(scratch_ptr, states_ptr, state_bytes, &self.stream, k, "causal_copy")?; + // Scratch activation buffers (reuse — causal runs AFTER graph forward) + let h_s1 = self.ptrs.save_h_s1; + let h_s2 = self.ptrs.save_h_s2; + let h_v = self.ptrs.save_h_v; + let h_b0 = self.ptrs.save_h_b0; + let h_b1 = self.ptrs.save_h_b1; + let h_b2 = self.ptrs.save_h_b2; - // Step 2: Zero feature k in scratch (set column k to 0 across all B rows) - // Each row is pad_sd bf16 elements. Feature k is at offset k in each row. - // Use a tiny kernel or just set individual elements. - // For simplicity, use the zero kernel on a strided pattern. - // Actually, we need a custom kernel for strided zeroing. - // Simpler: use cuMemsetD16Async on each row's feature k. - // Even simpler: launch a 1-thread-per-row kernel. - // Simplest correct: use the existing stochastic_depth_scale kernel - // with a scale of 0 on the specific feature column. - // - // Actually, the simplest: set the feature to a fixed value via a small loop. - // Since this runs outside the graph and is infrequent (every 10 steps, 14 features), - // a host-side loop setting individual bf16 values is acceptable. - let zero_val = [half::bf16::ZERO; 1]; - for row in 0..b { - let byte_offset = (row * pad_sd + k) * std::mem::size_of::(); - let dst_ptr = scratch_ptr + byte_offset as u64; - unsafe { - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - dst_ptr, - zero_val.as_ptr().cast(), - std::mem::size_of::(), - self.stream.cu_stream(), - ); - } + for k in 0..market_dim.min(14) { + let k_i32 = k as i32; + + // Step 1+2: GPU kernel copies states and zeros feature k + unsafe { + self.stream + .launch_builder(intervene_kernel) + .arg(&states_ptr) + .arg(&scratch_ptr) + .arg(&k_i32) + .arg(&pad_sd_i32) + .arg(&b_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "causal_intervene k={k}: {e}" + )))?; } - // Step 3: Forward pass on intervened states - // Reuse existing cuBLAS forward (non-graph, direct launch) - let cublas = &self.cublas_forward; - - // Use scratch logit buffers to avoid overwriting the graph's outputs. - // Reuse on_next_v/b_logits (Double-DQN scratch, not needed after graph). - cublas.forward_online_raw( + // Step 3: Forward pass on intervened states (cuBLAS, non-graph) + self.cublas_forward.forward_online_raw( &self.stream, scratch_ptr, &on_w_ptrs, - h_s1_ptr, h_s2_ptr, h_v_ptr, - h_b0_ptr, h_b1_ptr, h_b2_ptr, - self.ptrs.on_next_v_logits_buf, // reuse as causal Q scratch + h_s1, h_s2, h_v, h_b0, h_b1, h_b2, + self.ptrs.on_next_v_logits_buf, self.ptrs.on_next_b_logits_buf, )?; - // Step 4: Compute |Q_original - Q_intervened| for this feature - // For now, compute a simple L1 distance on the value logits. - // Download both value logits and compute on CPU (14 features × small download). - // This is the ONE allowed GPU→CPU transfer for causal intervention — - // it happens every 10 steps × 14 features = 140 tiny downloads per 10 steps. - let na = self.config.num_atoms; - let val_size = b * na; - let mut original_vals = vec![0.0_f32; val_size]; - let mut intervened_vals = vec![0.0_f32; val_size]; - - self.stream.synchronize() - .map_err(|e| MLError::ModelError(format!("causal sync: {e}")))?; - + // Step 4: GPU reduction — |Q_orig - Q_interv|² → sensitivity[k] unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - original_vals.as_mut_ptr().cast(), - self.ptrs.on_v_logits_buf, - val_size * std::mem::size_of::(), - ); - cudarc::driver::sys::cuMemcpyDtoH_v2( - intervened_vals.as_mut_ptr().cast(), - self.ptrs.on_next_v_logits_buf, - val_size * std::mem::size_of::(), - ); + self.stream + .launch_builder(reduce_kernel) + .arg(&self.ptrs.on_v_logits_buf) + .arg(&self.ptrs.on_next_v_logits_buf) + .arg(&sensitivity_ptr) + .arg(&k_i32) + .arg(&val_size) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "causal_reduce k={k}: {e}" + )))?; } - - let mut delta_sq_sum = 0.0_f32; - for i in 0..val_size { - let d = original_vals[i] - intervened_vals[i]; - delta_sq_sum += d * d; - } - let sensitivity_k = delta_sq_sum / val_size as f32; - total_sensitivity += sensitivity_k; } - let mean_sensitivity = total_sensitivity / market_dim.min(14) as f32; - Ok(mean_sensitivity) + // Single DtoH readback at the end — read mean sensitivity + let n_features = market_dim.min(14); + let mut host_sens = vec![0.0_f32; market_dim]; + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("causal final sync: {e}")))?; + self.stream.memcpy_dtoh(sensitivity, &mut host_sens[..market_dim]) + .map_err(|e| MLError::ModelError(format!("causal readback: {e}")))?; + let total: f32 = host_sens[..n_features].iter().sum(); + Ok(total / n_features as f32) } /// #32 Gradient Vaccine: project out overfitting gradient components. @@ -3150,24 +3136,30 @@ impl GpuDqnTrainer { Ok(()) } - /// #21 Stochastic depth: write random per-layer drop/keep scales before graph replay. + /// #21 Stochastic depth: generate per-layer drop/keep scales on GPU. + /// Uses GPU-resident LCG RNG — zero CPU-side randomness or HtoD transfers. /// Must be called BEFORE replay_forward() each training step. /// CUDA Graphs capture buffer addresses, not contents — this is safe. pub fn update_stochastic_depth_mask(&mut self) -> Result<(), MLError> { if self.stochastic_depth_prob <= 0.0 { - return Ok(()); // Disabled — mask stays [1.0, 1.0, 1.0] + return Ok(()); } - use rand::Rng; - let mut rng = rand::thread_rng(); let p = self.stochastic_depth_prob; - let keep_scale = 1.0 / (1.0 - p); // Expected-value correction - let mask: [f32; 3] = [ - if rng.gen::() > p { keep_scale } else { 0.0 }, - if rng.gen::() > p { keep_scale } else { 0.0 }, - if rng.gen::() > p { keep_scale } else { 0.0 }, - ]; - self.stream.memcpy_htod(&mask, &mut self.stochastic_depth_scale_buf) - .map_err(|e| MLError::ModelError(format!("stochastic_depth mask upload: {e}")))?; + let scale_ptr = self.stochastic_depth_scale_buf.raw_ptr(); + let rng_ptr = self.stochastic_depth_rng_state.raw_ptr(); + unsafe { + self.stream + .launch_builder(&self.stochastic_depth_rng_kernel) + .arg(&scale_ptr) + .arg(&rng_ptr) + .arg(&p) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("stochastic_depth_rng: {e}")))?; + } Ok(()) } @@ -5221,7 +5213,7 @@ impl GpuDqnTrainer { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -5265,6 +5257,12 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("pad_states_kernel load: {e}")))?; let stochastic_depth = module.load_function("stochastic_depth_scale") .map_err(|e| MLError::ModelError(format!("stochastic_depth_scale load: {e}")))?; + let stochastic_depth_rng = module.load_function("stochastic_depth_rng") + .map_err(|e| MLError::ModelError(format!("stochastic_depth_rng load: {e}")))?; + let causal_intervene = module.load_function("causal_intervene_feature") + .map_err(|e| MLError::ModelError(format!("causal_intervene load: {e}")))?; + let causal_reduce = module.load_function("causal_q_delta_reduce") + .map_err(|e| MLError::ModelError(format!("causal_reduce load: {e}")))?; let vaccine_dot = module.load_function("gradient_dot_and_norm") .map_err(|e| MLError::ModelError(format!("gradient_dot_and_norm load: {e}")))?; let vaccine_project = module.load_function("gradient_project") @@ -5276,8 +5274,8 @@ fn compile_training_kernels( let bn_tanh_concat = module.load_function("bn_tanh_concat_kernel") .map_err(|e| MLError::ModelError(format!("bn_tanh_concat_kernel load: {e}")))?; - info!("GpuDqnTrainer: 21 utility kernels loaded from precompiled cubin"); - Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32, stochastic_depth, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project)) + info!("GpuDqnTrainer: 24 utility kernels loaded from precompiled cubin"); + Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states, saxpy_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, causal_intervene, causal_reduce)) } /// Load the standalone Polyak EMA kernel from precompiled cubin.