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

fix(generalization): eliminate ALL CPU hot-path violations

Browse Source 
Stochastic depth (#21): replaced rand::thread_rng() + memcpy_htod
with GPU-native stochastic_depth_rng kernel. Single-thread LCG
generates 3 per-layer drop/keep scales directly in GPU memory.
Zero CPU RNG, zero HtoD transfers per training step.

Causal intervention (#34): replaced per-feature CPU zeroing loop
(B × cuMemcpyHtoDAsync per feature) and DtoH Q-value comparison with:
- causal_intervene_feature kernel: GPU-native state copy + feature zero
  One thread per batch sample, handles entire row copy + selective zero
- causal_q_delta_reduce kernel: GPU reduction of |Q_orig - Q_interv|²
  Block-level reduction with atomicAdd into sensitivity buffer
Only 1 DtoH of 168 bytes at the end (sensitivity readback for logging).

Audit result — remaining CPU paths in our code:
- Feature mask RNG: 42 floats, once per epoch (config generation)
- Domain randomization RNG: ~10 scalars, once per epoch (config generation)
- Causal sensitivity readback: 168 bytes, every 10th step (logging only)
All per-step computation is now 100% GPU-native.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-30 11:11:39 +02:00
parent 8bf1cdca13
commit a24fb94bbd
2 changed files with 215 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

93
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -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.
*

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

@@ -620,6 +620,10 @@ pub struct GpuDqnTrainer {
causal_q_scratch: Option<CudaSlice<half::bf16>>,
/// #34 Causal intervention config.
pub(crate) causal_config: CausalInterventionConfig,
/// #34 GPU-native causal intervention kernel.
causal_intervene_kernel: Option<CudaFunction>,
/// #34 GPU-native causal Q-delta reduction kernel.
causal_reduce_kernel: Option<CudaFunction>,
/// #32 Gradient vaccine: saved g_train gradient [total_params] f32.
vaccine_grad_save: CudaSlice<f32>,
@@ -648,10 +652,14 @@ pub struct GpuDqnTrainer {
bn_bias_grad_kernel: Option<CudaFunction>,
/// #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<f32>,
/// #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<u32>,
/// #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::<f32>(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::<f32>(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::<u32>(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<f32, MLError> {
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, &param_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::<half::bf16>();
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::<half::bf16>();
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::<half::bf16>(),
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::<f32>(),
);
cudarc::driver::sys::cuMemcpyDtoH_v2(
intervened_vals.as_mut_ptr().cast(),
self.ptrs.on_next_v_logits_buf,
val_size * std::mem::size_of::<f32>(),
);
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::<f32>() > p { keep_scale } else { 0.0 },
if rng.gen::<f32>() > p { keep_scale } else { 0.0 },
if rng.gen::<f32>() > 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<CudaStream>,
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.