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

feat(iql): new buffers and launcher methods for full integration

Browse Source 
17 kernels, per-sample support seeding, adv sigma EMA,
branch scales, expectile gap, per-sample epsilon launchers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 15:12:25 +02:00
parent b1cf06d48e
commit 5fbb39b0b5
1 changed files with 607 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files

741
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -3,12 +3,13 @@
//! GPU-accelerated IQL (Implicit Q-Learning) value network trainer.
//!
//! Implements the value network component of IQL (Kostrikov et al., 2021)
//! as a fused CUDA kernel pipeline:
//! as a fused CUDA kernel pipeline — zero atomicAdd, fully deterministic:
//!
//! 1. **Forward + Expectile Loss** -- V(s) MLP forward + asymmetric loss
//! 2. **Gradient Norm** -- L2 norm for gradient clipping
//! 3. **Backward** -- backprop through SiLU MLP, atomicAdd gradients
//! 4. **Adam** -- AdamW update with gradient clipping
//! 1. **Forward + Expectile Loss** — V(s) MLP forward + asymmetric loss
//! 2. **Backward (per-sample)** — backprop into per-sample gradient buffer
//! 3. **Weight Grad Reduce** — deterministic sum across samples
//! 4. **Grad Norm Phase 1+2** — two-phase L2 norm (no atomicAdd)
//! 5. **Adam** — AdamW update with gradient clipping
//!
//! ## Architecture
//!
@@ -17,9 +18,9 @@
//!
//! ## Integration
//!
//! Called after the main DQN training step in `FusedTrainingCtx::run_full_step()`.
//! The value loss trains V(s) to approximate the expectile of Q(s,a), and the
//! advantage weights `exp(beta * (Q(s,a) - V(s)))` can modulate the DQN policy.
//! Called after the main DQN training step in `FusedTrainingCtx::submit_aux_ops()`.
//! The V(s) output drives advantage weights `exp(beta * (Q(s,a) - V(s)))` that
//! modulate PER priorities for advantage-weighted replay.
use std::sync::Arc;
@@ -57,6 +58,18 @@ pub struct GpuIqlConfig {
pub weight_decay: f32,
/// Maximum gradient L2 norm for clipping.
pub max_grad_norm: f32,
/// Number of C51 atoms.
pub num_atoms: usize,
/// Total factored actions (b0*b1*b2*b3).
pub total_actions: usize,
/// Branch sizes [b0, b1, b2, b3] for advantage decomposition.
pub branch_sizes: [usize; 4],
/// Discount factor for Bellman headroom in per-sample support.
pub gamma: f32,
/// Staleness decay rate for PER modulation.
pub staleness_lambda: f32,
/// Staleness age normalizer (steps).
pub staleness_tau: f32,
}
impl Default for GpuIqlConfig {
@@ -73,6 +86,12 @@ impl Default for GpuIqlConfig {
epsilon: 1e-8,
weight_decay: 1e-5,
max_grad_norm: 1.0,
num_atoms: 51,
total_actions: 81,
branch_sizes: [3, 3, 3, 3],
gamma: 0.99,
staleness_lambda: 0.001,
staleness_tau: 10000.0,
}
}
}
@@ -94,9 +113,11 @@ impl GpuIqlConfig {
/// GPU-accelerated IQL value network trainer.
///
/// Owns pre-allocated GPU buffers for the V(s) network weights, Adam state,
/// activation saves, and the compiled CUDA kernels. Operates entirely on GPU
/// with no per-step CPU-GPU data transfers (states and Q-values are already
/// on device from the DQN training step).
/// per-sample gradients, and compiled CUDA kernels. Operates entirely on GPU
/// with no per-step CPU-GPU data transfers.
///
/// Zero atomicAdd — all gradient accumulation uses per-sample buffers with
/// deterministic cross-sample reduction.
#[allow(missing_debug_implementations)]
pub struct GpuIqlTrainer {
config: GpuIqlConfig,
@@ -104,19 +125,33 @@ pub struct GpuIqlTrainer {
// ── Compiled kernels ────────────────────────────────────────────
forward_loss_kernel: CudaFunction,
backward_kernel: CudaFunction,
grad_norm_kernel: CudaFunction,
backward_per_sample_kernel: CudaFunction,
weight_grad_reduce_kernel: CudaFunction,
loss_reduce_kernel: CudaFunction,
grad_norm_phase1_kernel: CudaFunction,
grad_norm_phase2_kernel: CudaFunction,
adam_kernel: CudaFunction,
forward_kernel: CudaFunction,
gather_q_taken_kernel: CudaFunction,
advantage_weight_kernel: CudaFunction,
modulate_td_kernel: CudaFunction,
adv_variance_kernel: CudaFunction,
per_sample_support_kernel: CudaFunction,
branch_advantage_kernel: CudaFunction,
expectile_gap_kernel: CudaFunction,
gap_mean_kernel: CudaFunction,
per_sample_epsilon_kernel: CudaFunction,
// ── V network parameters (flat f32 on GPU) ─────────────────────
params_buf: CudaSlice<f32>,
// ── Adam optimizer state ────────────────────────────────────────
m_buf: CudaSlice<f32>, // first moment
v_buf: CudaSlice<f32>, // second moment
grad_buf: CudaSlice<f32>, // gradient accumulator
grad_norm_buf: CudaSlice<f32>, // [1] gradient L2 norm
m_buf: CudaSlice<f32>, // first moment
v_buf: CudaSlice<f32>, // second moment
grad_buf: CudaSlice<f32>, // reduced gradient [total_params]
grads_per_sample: CudaSlice<f32>, // per-sample gradients [B * total_params]
grad_norm_buf: CudaSlice<f32>, // [1] gradient L2 norm (sum-of-squares)
grad_norm_partials: CudaSlice<f32>, // [num_blocks] partial sums for phase1
// ── Activation save buffers (forward -> backward) ───────────────
save_pre1: CudaSlice<f32>, // [B, H] pre-activation layer 1
@@ -125,22 +160,33 @@ pub struct GpuIqlTrainer {
save_h2: CudaSlice<f32>, // [B, H] post-activation layer 2
// ── Output buffers ──────────────────────────────────────────────
v_out_buf: CudaSlice<f32>, // [B] V(s) predictions
loss_buf: CudaSlice<f32>, // [B] per-sample loss
total_loss_buf: CudaSlice<f32>, // [1] batch-mean loss
v_out_buf: CudaSlice<f32>, // [B] V(s) predictions
loss_buf: CudaSlice<f32>, // [B] per-sample loss
total_loss_buf: CudaSlice<f32>, // [1] batch-mean loss
q_taken_buf: CudaSlice<f32>, // [B] Q(s, a_taken) gathered from q_out
advantage_weights_buf: CudaSlice<f32>, // [B] advantage weights
// ── New integration buffers ─────────────────────────────────────
adv_stats_buf: CudaSlice<f32>, // [2] mean, variance
adv_sigma_ema: f32, // host-side EMA
per_sample_support_buf: CudaSlice<f32>, // [B, 3]
branch_scales_buf: CudaSlice<f32>, // [B, 4]
expectile_gap_buf: CudaSlice<f32>, // [B]
gap_mean_buf: CudaSlice<f32>, // [1]
per_sample_epsilon_buf: CudaSlice<f32>, // [B]
// ── Training state ──────────────────────────────────────────────
adam_step: i32,
t_buf: CudaSlice<i32>,
total_params: usize,
grad_norm_blocks: usize,
}
impl GpuIqlTrainer {
/// Create a new GPU IQL trainer.
///
/// Compiles the 5 CUDA kernels (forward+loss, backward, grad_norm, adam,
/// forward-only), initializes V(s) weights with Xavier/Glorot uniform,
/// and pre-allocates all GPU buffers.
/// Compiles 9 CUDA kernels, initializes V(s) weights with Xavier/Glorot,
/// and pre-allocates all GPU buffers including per-sample gradient storage.
pub fn new(
stream: Arc<CudaStream>,
config: GpuIqlConfig,
@@ -149,9 +195,8 @@ impl GpuIqlTrainer {
let b = config.batch_size;
let h = config.value_hidden_dim;
// Compile all 5 kernels
let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_kernel, forward_kernel) =
compile_iql_kernels(&stream, &config)?;
// Compile all 9 kernels
let kernels = compile_iql_kernels(&stream, &config)?;
// Allocate parameter buffer and initialize with Xavier/Glorot
let params_buf = init_xavier_weights(&stream, &config)?;
@@ -160,7 +205,12 @@ impl GpuIqlTrainer {
let m_buf = alloc_f32(&stream, total_params, "iql_m")?;
let v_buf = alloc_f32(&stream, total_params, "iql_v")?;
let grad_buf = alloc_f32(&stream, total_params, "iql_grad")?;
let grads_per_sample = alloc_f32(&stream, b * total_params, "iql_grads_per_sample")?;
// Grad norm buffers (two-phase)
let grad_norm_blocks = (total_params + 255) / 256;
let grad_norm_buf = alloc_f32(&stream, 1, "iql_grad_norm")?;
let grad_norm_partials = alloc_f32(&stream, grad_norm_blocks, "iql_grad_norm_partials")?;
// Allocate activation save buffers
let save_pre1 = alloc_f32(&stream, b * h, "iql_save_pre1")?;
@@ -172,11 +222,33 @@ impl GpuIqlTrainer {
let v_out_buf = alloc_f32(&stream, b, "iql_v_out")?;
let loss_buf = alloc_f32(&stream, b, "iql_loss")?;
let total_loss_buf = alloc_f32(&stream, 1, "iql_total_loss")?;
let q_taken_buf = alloc_f32(&stream, b, "iql_q_taken")?;
let advantage_weights_buf = alloc_f32(&stream, b, "iql_adv_weights")?;
let vram_bytes = (total_params * 4 // params + m + v + grad = 4 copies
+ b * h * 4 // 4 activation buffers
+ b * 2 + 2 // v_out + loss + total_loss + grad_norm
) * 4;
// New integration buffers
let adv_stats_buf = alloc_f32(&stream, 2, "iql_adv_stats")?;
let mut per_sample_support_buf = alloc_f32(&stream, b * 3, "iql_per_sample_support")?;
let branch_scales_buf = alloc_f32(&stream, b * 4, "iql_branch_scales")?;
let expectile_gap_buf = alloc_f32(&stream, b, "iql_expectile_gap")?;
let gap_mean_buf = alloc_f32(&stream, 1, "iql_gap_mean")?;
let per_sample_epsilon_buf = alloc_f32(&stream, b, "iql_per_sample_epsilon")?;
// Seed per_sample_support with defaults for step 0
let na = config.num_atoms.max(2) as f32;
let default_delta_z = 2.0 / (na - 1.0);
let mut default_support = vec![0.0_f32; b * 3];
for i in 0..b {
default_support[i * 3] = -1.0;
default_support[i * 3 + 1] = 1.0;
default_support[i * 3 + 2] = default_delta_z;
}
super::htod_f32(&stream, &default_support, &mut per_sample_support_buf)?;
let per_sample_vram = b * total_params * 4;
let vram_bytes = total_params * 4 * 4 // params + m + v + grad
+ per_sample_vram // per-sample gradients
+ b * h * 4 * 4 // 4 activation buffers
+ (b * 3 + 2) * 4; // v_out + loss + total_loss + adv_weights + grad_norm
info!(
state_dim = config.state_dim,
@@ -185,8 +257,9 @@ impl GpuIqlTrainer {
total_params,
expectile_tau = config.expectile_tau,
advantage_temperature = config.advantage_temperature,
per_sample_grad_mb = per_sample_vram / (1024 * 1024),
vram_kb = vram_bytes / 1024,
"GpuIqlTrainer initialized: 5 kernels compiled, V(s) weights Xavier-initialized"
"GpuIqlTrainer initialized: 17 kernels, zero atomicAdd, V(s) Xavier-initialized"
);
let t_buf = stream.alloc_zeros::<i32>(1)
@@ -194,16 +267,30 @@ impl GpuIqlTrainer {
Ok(Self {
config,
stream,
forward_loss_kernel,
backward_kernel,
grad_norm_kernel,
adam_kernel,
forward_kernel,
forward_loss_kernel: kernels.forward_loss,
backward_per_sample_kernel: kernels.backward_per_sample,
weight_grad_reduce_kernel: kernels.weight_grad_reduce,
loss_reduce_kernel: kernels.loss_reduce,
grad_norm_phase1_kernel: kernels.grad_norm_phase1,
grad_norm_phase2_kernel: kernels.grad_norm_phase2,
adam_kernel: kernels.adam,
forward_kernel: kernels.forward,
gather_q_taken_kernel: kernels.gather_q_taken,
advantage_weight_kernel: kernels.advantage_weight,
modulate_td_kernel: kernels.modulate_td,
adv_variance_kernel: kernels.adv_variance,
per_sample_support_kernel: kernels.per_sample_support,
branch_advantage_kernel: kernels.branch_advantage,
expectile_gap_kernel: kernels.expectile_gap,
gap_mean_kernel: kernels.gap_mean,
per_sample_epsilon_kernel: kernels.per_sample_epsilon,
params_buf,
m_buf,
v_buf,
grad_buf,
grads_per_sample,
grad_norm_buf,
grad_norm_partials,
save_pre1,
save_pre2,
save_h1,
@@ -211,122 +298,156 @@ impl GpuIqlTrainer {
v_out_buf,
loss_buf,
total_loss_buf,
q_taken_buf,
advantage_weights_buf,
adv_stats_buf,
adv_sigma_ema: 0.0,
per_sample_support_buf,
branch_scales_buf,
expectile_gap_buf,
gap_mean_buf,
per_sample_epsilon_buf,
adam_step: 0,
t_buf,
total_params,
grad_norm_blocks,
})
}
/// Run one IQL value network training step.
/// Run one IQL value network training step — fully deterministic.
///
/// Executes the full forward + expectile loss + backward + Adam cycle:
/// 1. Zero total_loss and grad_norm scalars
/// 2. Forward + loss kernel (8 warps per sample)
/// 3. Backward kernel (8 warps per sample, atomicAdd gradients)
/// 4. Grad norm kernel (parallel reduction)
/// 5. Adam update kernel (one thread per parameter)
/// Pipeline: forward+loss → backward_per_sample → weight_grad_reduce
/// → loss_reduce → grad_norm phase1+2 → Adam.
///
/// Returns the batch-mean expectile loss (read back from GPU).
/// Zero atomicAdd. All cross-sample reduction is deterministic.
///
/// # Arguments
///
/// * `states_f32` - F32 CudaSlice `[B * STATE_DIM]` on CUDA device
/// * `q_values_f32` - F32 CudaSlice `[B]` on CUDA device (target Q-values from DQN)
///
/// Takes pre-converted F32 CudaSlice buffers directly — zero Candle
/// tensor involvement.
/// `gather_q_taken()` MUST be called first to populate `q_taken_buf`.
/// The expectile regression target is Q(s, a_taken) from the DQN.
pub fn train_value_step(
&mut self,
states_f32: &CudaSlice<f32>,
q_values_f32: &CudaSlice<f32>,
) -> Result<f32, MLError> {
) -> Result<(), MLError> {
let b = self.config.batch_size;
// Zero total_loss and grad_norm before kernel launches
let zero_f32 = [0.0_f32];
super::htod_f32(&self.stream, &zero_f32, &mut self.total_loss_buf)?;
super::htod_f32(&self.stream, &zero_f32, &mut self.grad_norm_buf)?;
let state_dim_i32 = self.config.state_dim as i32;
let batch_size_i32 = b as i32;
let total_params_i32 = self.total_params as i32;
let expectile_tau = self.config.expectile_tau;
// 1. Forward + loss kernel (256 threads: 8 warps per sample)
// Shared memory: warp_sums[8] in f32 for block-level reduction
// 1. Forward + loss kernel (256 threads per sample)
let fwd_shmem = 8 * std::mem::size_of::<f32>();
let fwd_config = LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: fwd_shmem as u32,
};
// Safety: states_f32 and q_values_f32 are valid F32 CudaSlice buffers
// on the same CUDA context. All internal buffers are pre-allocated.
let state_dim_i32 = self.config.state_dim as i32;
unsafe {
self.stream
.launch_builder(&self.forward_loss_kernel)
.arg(states_f32)
.arg(q_values_f32)
.arg(&self.q_taken_buf)
.arg(&self.params_buf)
.arg(&mut self.v_out_buf)
.arg(&mut self.loss_buf)
.arg(&mut self.total_loss_buf)
.arg(&mut self.save_pre1)
.arg(&mut self.save_pre2)
.arg(&mut self.save_h1)
.arg(&mut self.save_h2)
.arg(&batch_size_i32)
.arg(&state_dim_i32)
.launch(fwd_config)
.map_err(|e| MLError::ModelError(format!("IQL forward+loss kernel: {e}")))?;
.arg(&expectile_tau)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: fwd_shmem as u32,
})
.map_err(|e| MLError::ModelError(format!("IQL forward+loss: {e}")))?;
}
// 2. Backward kernel (256 threads: 8 warps per sample)
let bwd_config = LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
// Safety: all pointers are valid GPU allocations on the same stream.
// 2. Backward per-sample (256 threads per sample, no atomicAdd)
unsafe {
self.stream
.launch_builder(&self.backward_kernel)
.launch_builder(&self.backward_per_sample_kernel)
.arg(states_f32)
.arg(q_values_f32)
.arg(&self.q_taken_buf)
.arg(&self.v_out_buf)
.arg(&self.params_buf)
.arg(&self.save_pre1)
.arg(&self.save_pre2)
.arg(&self.save_h1)
.arg(&self.save_h2)
.arg(&mut self.grad_buf)
.arg(&mut self.grads_per_sample)
.arg(&batch_size_i32)
.arg(&state_dim_i32)
.launch(bwd_config)
.map_err(|e| MLError::ModelError(format!("IQL backward kernel: {e}")))?;
.arg(&total_params_i32)
.arg(&expectile_tau)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL backward_per_sample: {e}")))?;
}
// 3. Grad norm kernel
let norm_blocks = (self.total_params + 255) / 256;
let norm_config = LaunchConfig {
grid_dim: (norm_blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
// Safety: grad_buf and grad_norm_buf are valid allocations.
// 3. Weight grad reduce (deterministic sum across samples)
let reduce_blocks = (self.total_params + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.grad_norm_kernel)
.arg(&self.grad_buf)
.arg(&mut self.grad_norm_buf)
.launch_builder(&self.weight_grad_reduce_kernel)
.arg(&self.grads_per_sample)
.arg(&mut self.grad_buf)
.arg(&batch_size_i32)
.arg(&total_params_i32)
.launch(norm_config)
.map_err(|e| MLError::ModelError(format!("IQL grad_norm kernel: {e}")))?;
.launch(LaunchConfig {
grid_dim: (reduce_blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL weight_grad_reduce: {e}")))?;
}
// 4. Adam update kernel — adam_step on GPU (async HtoD)
// 4. Loss reduce (deterministic sequential sum)
unsafe {
self.stream
.launch_builder(&self.loss_reduce_kernel)
.arg(&self.loss_buf)
.arg(&mut self.total_loss_buf)
.arg(&batch_size_i32)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL loss_reduce: {e}")))?;
}
// 5. Grad norm phase 1 (per-block partial sums)
let norm_blocks = self.grad_norm_blocks;
unsafe {
self.stream
.launch_builder(&self.grad_norm_phase1_kernel)
.arg(&self.grad_buf)
.arg(&mut self.grad_norm_partials)
.arg(&total_params_i32)
.launch(LaunchConfig {
grid_dim: (norm_blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 256 * std::mem::size_of::<f32>() as u32,
})
.map_err(|e| MLError::ModelError(format!("IQL grad_norm_phase1: {e}")))?;
}
// 6. Grad norm phase 2 (sequential sum of partials)
let norm_blocks_i32 = norm_blocks as i32;
unsafe {
self.stream
.launch_builder(&self.grad_norm_phase2_kernel)
.arg(&self.grad_norm_partials)
.arg(&mut self.grad_norm_buf)
.arg(&norm_blocks_i32)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL grad_norm_phase2: {e}")))?;
}
// 7. Adam update — adam_step on GPU (async HtoD)
self.adam_step += 1;
unsafe {
cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
@@ -337,12 +458,6 @@ impl GpuIqlTrainer {
);
}
let adam_blocks = (self.total_params + 255) / 256;
let adam_config = LaunchConfig {
grid_dim: (adam_blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
let lr = self.config.lr;
let beta1 = self.config.beta1;
let beta2 = self.config.beta2;
@@ -367,13 +482,109 @@ impl GpuIqlTrainer {
.arg(&mgn)
.arg(&t_ptr)
.arg(&total_params_i32)
.launch(adam_config)
.map_err(|e| MLError::ModelError(format!("IQL adam kernel: {e}")))?;
.launch(LaunchConfig {
grid_dim: (adam_blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL adam: {e}")))?;
}
// Loss stays on GPU — no per-step DtoH. Caller discards the value.
// Removing the synchronous readback unblocks CUDA Graph capture.
Ok(0.0)
Ok(())
}
/// Gather Q(s, a_taken) from the full Q-value buffer into q_taken_buf.
///
/// Must be called BEFORE `train_value_step` so IQL trains on Q-values,
/// not raw rewards.
pub fn gather_q_taken(
&mut self,
q_out_buf: &CudaSlice<f32>,
actions_buf: &CudaSlice<i32>,
total_actions: usize,
) -> Result<&CudaSlice<f32>, MLError> {
let b = self.config.batch_size;
let batch_size_i32 = b as i32;
let total_actions_i32 = total_actions as i32;
let blocks = (b + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.gather_q_taken_kernel)
.arg(q_out_buf)
.arg(actions_buf)
.arg(&mut self.q_taken_buf)
.arg(&batch_size_i32)
.arg(&total_actions_i32)
.launch(LaunchConfig {
grid_dim: (blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL gather_q_taken: {e}")))?;
}
Ok(&self.q_taken_buf)
}
/// Reference to the gathered Q(s, a_taken) buffer [B].
pub fn q_taken_buf(&self) -> &CudaSlice<f32> {
&self.q_taken_buf
}
/// Compute IQL advantage weights: w[b] = exp(beta * (Q_taken - V(s))).
///
/// Must be called AFTER `train_value_step` (v_out_buf is populated)
/// and AFTER `gather_q_taken` (q_taken_buf is populated).
///
/// Advantage weights are written to `advantage_weights_buf [B]` and
/// can be used to modulate PER priorities.
pub fn compute_advantage_weights(
&mut self,
q_out_buf: &CudaSlice<f32>,
actions_buf: &CudaSlice<i32>,
total_actions: usize,
) -> Result<(), MLError> {
let b = self.config.batch_size;
let beta = self.config.advantage_temperature;
let batch_size_i32 = b as i32;
let total_actions_i32 = total_actions as i32;
let blocks = (b + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.advantage_weight_kernel)
.arg(q_out_buf)
.arg(actions_buf)
.arg(&self.v_out_buf)
.arg(&mut self.advantage_weights_buf)
.arg(&beta)
.arg(&batch_size_i32)
.arg(&total_actions_i32)
.launch(LaunchConfig {
grid_dim: (blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL advantage_weights: {e}")))?;
}
Ok(())
}
/// Raw device pointer to V(s) output buffer [B].
pub fn v_out_ptr(&self) -> u64 {
self.v_out_buf.raw_ptr()
}
/// Raw device pointer to advantage weights buffer [B].
pub fn advantage_weights_ptr(&self) -> u64 {
self.advantage_weights_buf.raw_ptr()
}
/// Raw device pointer to total loss buffer [1].
pub fn total_loss_ptr(&self) -> u64 {
self.total_loss_buf.raw_ptr()
}
/// Current Adam step count.
@@ -386,6 +597,260 @@ impl GpuIqlTrainer {
pub fn increment_adam_step(&mut self) {
self.adam_step += 1;
}
// ── New integration launchers ────────────────────────────────────
/// Modulate td_errors in-place with advantage weights and staleness decay.
pub fn modulate_td_errors(
&self,
td_errors: &mut CudaSlice<f32>,
indices_ptr: u64,
write_pos: i32,
capacity: i32,
) -> Result<(), MLError> {
let b = self.config.batch_size;
let batch_i32 = b as i32;
let beta = self.config.advantage_temperature;
let sigma = self.adv_sigma_ema;
let lambda = self.config.staleness_lambda;
let tau = self.config.staleness_tau;
let blocks = (b + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.modulate_td_kernel)
.arg(td_errors)
.arg(&self.advantage_weights_buf)
.arg(&indices_ptr)
.arg(&sigma)
.arg(&beta)
.arg(&lambda)
.arg(&tau)
.arg(&write_pos)
.arg(&capacity)
.arg(&batch_i32)
.launch(LaunchConfig {
grid_dim: (blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL modulate_td_errors: {e}")))?;
}
Ok(())
}
/// Update the EMA of advantage standard deviation from the current batch.
pub fn update_adv_sigma(&mut self) -> Result<(), MLError> {
let batch_i32 = self.config.batch_size as i32;
unsafe {
self.stream
.launch_builder(&self.adv_variance_kernel)
.arg(&self.advantage_weights_buf)
.arg(&mut self.adv_stats_buf)
.arg(&batch_i32)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL adv_variance_reduce: {e}")))?;
}
let mut stats = [0.0_f32; 2];
self.stream.memcpy_dtoh(&self.adv_stats_buf, &mut stats)
.map_err(|e| MLError::ModelError(format!("IQL adv_stats DtoH: {e}")))?;
let sigma = stats[1].max(0.0).sqrt();
const EMA_BETA: f32 = 0.99;
if self.adv_sigma_ema < 1e-8 {
self.adv_sigma_ema = sigma;
} else {
self.adv_sigma_ema = EMA_BETA * self.adv_sigma_ema + (1.0 - EMA_BETA) * sigma;
}
Ok(())
}
/// Compute per-sample C51 atom support centered on V(s).
pub fn compute_per_sample_support(
&mut self,
q_out_buf: &CudaSlice<f32>,
) -> Result<(), MLError> {
let b = self.config.batch_size;
let batch_i32 = b as i32;
let total_actions_i32 = self.config.total_actions as i32;
let num_atoms_i32 = self.config.num_atoms as i32;
let gamma = self.config.gamma;
let blocks = (b + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.per_sample_support_kernel)
.arg(&self.v_out_buf)
.arg(q_out_buf)
.arg(&mut self.per_sample_support_buf)
.arg(&gamma)
.arg(&batch_i32)
.arg(&total_actions_i32)
.arg(&num_atoms_i32)
.launch(LaunchConfig {
grid_dim: (blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL per_sample_support: {e}")))?;
}
Ok(())
}
/// Compute per-branch gradient scales from Q-value marginalization.
pub fn compute_branch_scales(
&mut self,
q_out_buf: &CudaSlice<f32>,
actions_buf: &CudaSlice<i32>,
) -> Result<(), MLError> {
let b = self.config.batch_size;
let batch_i32 = b as i32;
let ta = self.config.total_actions as i32;
let bs = self.config.branch_sizes;
let b0 = bs[0] as i32;
let b1 = bs[1] as i32;
let b2 = bs[2] as i32;
let b3 = bs[3] as i32;
let blocks = (b + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.branch_advantage_kernel)
.arg(q_out_buf)
.arg(&self.v_out_buf)
.arg(actions_buf)
.arg(&mut self.branch_scales_buf)
.arg(&batch_i32)
.arg(&ta)
.arg(&b0)
.arg(&b1)
.arg(&b2)
.arg(&b3)
.launch(LaunchConfig {
grid_dim: (blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL branch_advantage: {e}")))?;
}
Ok(())
}
/// Compute expectile gap between two V(s) estimates (high-tau minus low-tau).
pub fn compute_expectile_gap(
&mut self,
v_low_ptr: u64,
) -> Result<(), MLError> {
let b = self.config.batch_size;
let batch_i32 = b as i32;
let blocks = (b + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.expectile_gap_kernel)
.arg(&self.v_out_buf)
.arg(&v_low_ptr)
.arg(&mut self.expectile_gap_buf)
.arg(&batch_i32)
.launch(LaunchConfig {
grid_dim: (blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL expectile_gap: {e}")))?;
}
unsafe {
self.stream
.launch_builder(&self.gap_mean_kernel)
.arg(&self.expectile_gap_buf)
.arg(&mut self.gap_mean_buf)
.arg(&batch_i32)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL gap_mean_reduce: {e}")))?;
}
Ok(())
}
/// Compute per-sample epsilon from expectile gap.
pub fn compute_per_sample_epsilon(
&mut self,
base_epsilon: f32,
) -> Result<(), MLError> {
let b = self.config.batch_size;
let batch_i32 = b as i32;
let blocks = (b + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.per_sample_epsilon_kernel)
.arg(&self.expectile_gap_buf)
.arg(&self.gap_mean_buf)
.arg(&mut self.per_sample_epsilon_buf)
.arg(&base_epsilon)
.arg(&batch_i32)
.launch(LaunchConfig {
grid_dim: (blocks as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQL per_sample_epsilon: {e}")))?;
}
Ok(())
}
/// Raw pointer to per-sample support buffer [B, 3].
pub fn per_sample_support_ptr(&self) -> u64 {
self.per_sample_support_buf.raw_ptr()
}
/// Raw pointer to branch scales buffer [B, 4].
pub fn branch_scales_ptr(&self) -> u64 {
self.branch_scales_buf.raw_ptr()
}
/// Raw pointer to per-sample epsilon buffer [B].
pub fn per_sample_epsilon_ptr(&self) -> u64 {
self.per_sample_epsilon_buf.raw_ptr()
}
/// Reference to per-sample epsilon buffer.
pub fn per_sample_epsilon_buf(&self) -> &CudaSlice<f32> {
&self.per_sample_epsilon_buf
}
}
// ---------------------------------------------------------------------------
// Compiled kernel set
// ---------------------------------------------------------------------------
struct IqlKernels {
forward_loss: CudaFunction,
backward_per_sample: CudaFunction,
weight_grad_reduce: CudaFunction,
loss_reduce: CudaFunction,
grad_norm_phase1: CudaFunction,
grad_norm_phase2: CudaFunction,
adam: CudaFunction,
forward: CudaFunction,
gather_q_taken: CudaFunction,
advantage_weight: CudaFunction,
modulate_td: CudaFunction,
adv_variance: CudaFunction,
per_sample_support: CudaFunction,
branch_advantage: CudaFunction,
expectile_gap: CudaFunction,
gap_mean: CudaFunction,
per_sample_epsilon: CudaFunction,
}
// ---------------------------------------------------------------------------
@@ -395,17 +860,17 @@ impl GpuIqlTrainer {
/// Precompiled IQL value kernel cubin, embedded at compile time by build.rs.
static IQL_VALUE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/iql_value_kernel.cubin"));
/// Load all 5 IQL CUDA kernels from precompiled cubin.
/// Load all 17 IQL CUDA kernels from precompiled cubin.
fn compile_iql_kernels(
stream: &Arc<CudaStream>,
config: &GpuIqlConfig,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
) -> Result<IqlKernels, MLError> {
info!(
state_dim = config.state_dim,
hidden_dim = config.value_hidden_dim,
expectile_tau = config.expectile_tau,
total_params = config.total_params(),
"GpuIqlTrainer: loading precompiled IQL kernels"
"GpuIqlTrainer: loading 17 deterministic IQL kernels (zero atomicAdd)"
);
let context = stream.context();
@@ -413,24 +878,33 @@ fn compile_iql_kernels(
MLError::ModelError(format!("iql_value module load: {e}"))
})?;
let forward_loss = module
.load_function("iql_forward_loss_kernel")
.map_err(|e| MLError::ModelError(format!("iql_forward_loss_kernel load: {e}")))?;
let backward = module
.load_function("iql_backward_kernel")
.map_err(|e| MLError::ModelError(format!("iql_backward_kernel load: {e}")))?;
let grad_norm = module
.load_function("iql_grad_norm_kernel")
.map_err(|e| MLError::ModelError(format!("iql_grad_norm_kernel load: {e}")))?;
let adam = module
.load_function("iql_adam_kernel")
.map_err(|e| MLError::ModelError(format!("iql_adam_kernel load: {e}")))?;
let forward = module
.load_function("iql_forward_kernel")
.map_err(|e| MLError::ModelError(format!("iql_forward_kernel load: {e}")))?;
let load = |name: &str| -> Result<CudaFunction, MLError> {
module.load_function(name)
.map_err(|e| MLError::ModelError(format!("{name} load: {e}")))
};
info!("GpuIqlTrainer: 5 kernels compiled and loaded");
Ok((forward_loss, backward, grad_norm, adam, forward))
let kernels = IqlKernels {
forward_loss: load("iql_forward_loss_kernel")?,
backward_per_sample: load("iql_backward_per_sample")?,
weight_grad_reduce: load("iql_weight_grad_reduce")?,
loss_reduce: load("iql_loss_reduce")?,
grad_norm_phase1: load("iql_grad_norm_phase1")?,
grad_norm_phase2: load("iql_grad_norm_phase2")?,
adam: load("iql_adam_kernel")?,
forward: load("iql_forward_kernel")?,
gather_q_taken: load("iql_gather_q_taken")?,
advantage_weight: load("iql_compute_advantage_weights")?,
modulate_td: load("iql_modulate_td_errors")?,
adv_variance: load("iql_adv_variance_reduce")?,
per_sample_support: load("iql_compute_per_sample_support")?,
branch_advantage: load("iql_per_branch_advantage")?,
expectile_gap: load("iql_expectile_gap")?,
gap_mean: load("iql_gap_mean_reduce")?,
per_sample_epsilon: load("iql_compute_per_sample_epsilon")?,
};
info!("GpuIqlTrainer: 17 kernels loaded");
Ok(kernels)
}
// ---------------------------------------------------------------------------
@@ -503,4 +977,3 @@ fn alloc_f32(
MLError::ModelError(format!("GpuIql alloc {name} ({n} f32): {e}"))
})
}