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

feat: IQN/Attention accept workspace+stream override for multi-stream

Browse Source 
Add optional override_stream and override_workspace parameters to
execute_training_pipeline (GpuIqnHead), and forward/backward/adam_step
(GpuAttention). All existing callers pass None, None — behaviour is
unchanged. The parallel path can now dispatch these ops to dedicated
CUDA streams with separate cuBLAS workspaces.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 16:20:31 +02:00
parent dc9f25df63
commit 16f5147356
4 changed files with 170 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

174
crates/ml/src/cuda_pipeline/gpu_attention.rs
View File

@@ -456,7 +456,16 @@ impl GpuAttention {
/// Returns a reference to the output buffer (attended states).
pub fn stream_handle(&self) -> u64 { self.stream.cu_stream() as u64 }
pub fn forward(&mut self, states: &CudaSlice<f32>, batch_size: usize) -> Result<&CudaSlice<f32>, MLError> {
pub fn forward(
&mut self,
states: &CudaSlice<f32>,
batch_size: usize,
override_stream: Option<&Arc<CudaStream>>,
override_workspace: Option<(u64, usize)>,
) -> Result<&CudaSlice<f32>, MLError> {
let stream = override_stream.unwrap_or(&self.stream);
let (ws_ptr, ws_size) = override_workspace
.unwrap_or((self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size));
let b = batch_size;
let d = self.config.state_dim;
let num_heads = self.config.num_heads;
@@ -465,7 +474,7 @@ impl GpuAttention {
let n_bytes = b * d * std::mem::size_of::<f32>();
unsafe {
cudarc::driver::result::memcpy_dtod_async(
self.saved_input.raw_ptr(), states.raw_ptr(), n_bytes, self.stream.cu_stream()
self.saved_input.raw_ptr(), states.raw_ptr(), n_bytes, stream.cu_stream()
).map_err(|e| MLError::ModelError(format!("attention save input DtoD: {e}")))?;
}
@@ -482,38 +491,38 @@ impl GpuAttention {
let params_base = self.params.raw_ptr();
// 1. Q = W_Q^T @ states
self.lt_matmul(&self.gemm_fwd_q,
self.lt_matmul_ex(&self.gemm_fwd_q,
params_base + (w_q_off * f32_sz) as u64,
states.raw_ptr(),
self.proj_q_buf.raw_ptr(),
1.0, 0.0, "fwd_q")?;
1.0, 0.0, "fwd_q", stream, ws_ptr, ws_size)?;
// + bias_add Q
self.launch_bias_add(self.proj_q_buf.raw_ptr(), params_base + (b_q_off * f32_sz) as u64, d, b)?;
self.launch_bias_add_ex(self.proj_q_buf.raw_ptr(), params_base + (b_q_off * f32_sz) as u64, d, b, stream)?;
// 2. K = W_K^T @ states
self.lt_matmul(&self.gemm_fwd_k,
self.lt_matmul_ex(&self.gemm_fwd_k,
params_base + (w_k_off * f32_sz) as u64,
states.raw_ptr(),
self.proj_k_buf.raw_ptr(),
1.0, 0.0, "fwd_k")?;
1.0, 0.0, "fwd_k", stream, ws_ptr, ws_size)?;
// + bias_add K
self.launch_bias_add(self.proj_k_buf.raw_ptr(), params_base + (b_k_off * f32_sz) as u64, d, b)?;
self.launch_bias_add_ex(self.proj_k_buf.raw_ptr(), params_base + (b_k_off * f32_sz) as u64, d, b, stream)?;
// 3. V = W_V^T @ states
self.lt_matmul(&self.gemm_fwd_v,
self.lt_matmul_ex(&self.gemm_fwd_v,
params_base + (w_v_off * f32_sz) as u64,
states.raw_ptr(),
self.proj_v_buf.raw_ptr(),
1.0, 0.0, "fwd_v")?;
1.0, 0.0, "fwd_v", stream, ws_ptr, ws_size)?;
// + bias_add V
self.launch_bias_add(self.proj_v_buf.raw_ptr(), params_base + (b_v_off * f32_sz) as u64, d, b)?;
self.launch_bias_add_ex(self.proj_v_buf.raw_ptr(), params_base + (b_v_off * f32_sz) as u64, d, b, stream)?;
// 4. SDP: attn_out = SDP_fwd(Q, K, V)
let b_i32 = b as i32;
let d_i32 = d as i32;
let num_heads_i32 = num_heads as i32;
unsafe {
self.stream
stream
.launch_builder(&self.sdp_fwd_kernel)
.arg(&self.proj_q_buf)
.arg(&self.proj_k_buf)
@@ -532,19 +541,19 @@ impl GpuAttention {
}
// 5. projected = W_O^T @ attn_out
self.lt_matmul(&self.gemm_fwd_o,
self.lt_matmul_ex(&self.gemm_fwd_o,
params_base + (w_o_off * f32_sz) as u64,
self.attn_out_buf.raw_ptr(),
self.projected_buf.raw_ptr(),
1.0, 0.0, "fwd_o")?;
1.0, 0.0, "fwd_o", stream, ws_ptr, ws_size)?;
// + bias_add O
self.launch_bias_add(self.projected_buf.raw_ptr(), params_base + (b_o_off * f32_sz) as u64, d, b)?;
self.launch_bias_add_ex(self.projected_buf.raw_ptr(), params_base + (b_o_off * f32_sz) as u64, d, b, stream)?;
// 6. LayerNorm(projected + residual=states) -> output
let gamma_off = 4 * d * d + 4 * d;
let beta_off = gamma_off + d;
unsafe {
self.stream
stream
.launch_builder(&self.layer_norm_fwd_kernel)
.arg(&self.projected_buf)
.arg(states)
@@ -585,7 +594,12 @@ impl GpuAttention {
&mut self,
d_output: &CudaSlice<f32>,
batch_size: usize,
override_stream: Option<&Arc<CudaStream>>,
override_workspace: Option<(u64, usize)>,
) -> Result<(), MLError> {
let stream = override_stream.unwrap_or(&self.stream);
let (ws_ptr, ws_size) = override_workspace
.unwrap_or((self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size));
let b = batch_size;
let d = self.config.state_dim;
let num_heads = self.config.num_heads;
@@ -593,7 +607,7 @@ impl GpuAttention {
let inv_batch = 1.0_f32 / b as f32;
// Zero grad_buf before accumulation
self.stream.memset_zeros(&mut self.d_params)
stream.memset_zeros(&mut self.d_params)
.map_err(|e| MLError::ModelError(format!("attention zero d_params: {e}")))?;
// Weight layout offsets
@@ -616,7 +630,7 @@ impl GpuAttention {
// 1. LayerNorm backward -> d_projected (also writes d_gamma, d_beta into grad_buf)
unsafe {
self.stream
stream
.launch_builder(&self.layer_norm_bwd_kernel)
.arg(d_output)
.arg(&self.saved_input) // x (pre-LN input)
@@ -638,28 +652,28 @@ impl GpuAttention {
}
// 2. dW_O = d_projected[D,B] @ attn_out^T[B,D] -> [D,D], alpha=1/B for mean
self.lt_matmul(&self.gemm_bwd_dw_o,
self.lt_matmul_ex(&self.gemm_bwd_dw_o,
self.d_projected_buf.raw_ptr(),
self.attn_out_buf.raw_ptr(),
grad_base + (w_o_off * f32_sz) as u64,
inv_batch, 0.0, "bwd_dw_o")?;
inv_batch, 0.0, "bwd_dw_o", stream, ws_ptr, ws_size)?;
// 3. Bias gradient reduce for b_O
self.launch_bias_grad_reduce(
self.launch_bias_grad_reduce_ex(
self.d_projected_buf.raw_ptr(),
grad_base + (b_o_off * f32_sz) as u64,
d, b, inv_batch)?;
d, b, inv_batch, stream)?;
// 4. d_attn_out = W_O[D,D] @ d_projected[D,B] -> [D,B]
self.lt_matmul(&self.gemm_bwd_dx_o,
self.lt_matmul_ex(&self.gemm_bwd_dx_o,
params_base + (w_o_off * f32_sz) as u64,
self.d_projected_buf.raw_ptr(),
self.d_attn_out_buf.raw_ptr(),
1.0, 0.0, "bwd_dx_o")?;
1.0, 0.0, "bwd_dx_o", stream, ws_ptr, ws_size)?;
// 5. SDP backward: d_attn_out -> d_Q, d_K, d_V
unsafe {
self.stream
stream
.launch_builder(&self.sdp_bwd_kernel)
.arg(&self.d_attn_out_buf)
.arg(&self.proj_q_buf)
@@ -681,59 +695,59 @@ impl GpuAttention {
}
// 6. dW_Q = d_Q[D,B] @ states^T[B,D], alpha=1/B
self.lt_matmul(&self.gemm_bwd_dw_q,
self.lt_matmul_ex(&self.gemm_bwd_dw_q,
self.d_proj_q_buf.raw_ptr(),
self.saved_input.raw_ptr(),
grad_base + (w_q_off * f32_sz) as u64,
inv_batch, 0.0, "bwd_dw_q")?;
inv_batch, 0.0, "bwd_dw_q", stream, ws_ptr, ws_size)?;
// dW_K = d_K[D,B] @ states^T[B,D], alpha=1/B
self.lt_matmul(&self.gemm_bwd_dw_k,
self.lt_matmul_ex(&self.gemm_bwd_dw_k,
self.d_proj_k_buf.raw_ptr(),
self.saved_input.raw_ptr(),
grad_base + (w_k_off * f32_sz) as u64,
inv_batch, 0.0, "bwd_dw_k")?;
inv_batch, 0.0, "bwd_dw_k", stream, ws_ptr, ws_size)?;
// dW_V = d_V[D,B] @ states^T[B,D], alpha=1/B
self.lt_matmul(&self.gemm_bwd_dw_v,
self.lt_matmul_ex(&self.gemm_bwd_dw_v,
self.d_proj_v_buf.raw_ptr(),
self.saved_input.raw_ptr(),
grad_base + (w_v_off * f32_sz) as u64,
inv_batch, 0.0, "bwd_dw_v")?;
inv_batch, 0.0, "bwd_dw_v", stream, ws_ptr, ws_size)?;
// 7. Bias gradient reduce for b_Q, b_K, b_V
self.launch_bias_grad_reduce(
self.launch_bias_grad_reduce_ex(
self.d_proj_q_buf.raw_ptr(),
grad_base + (b_q_off * f32_sz) as u64,
d, b, inv_batch)?;
self.launch_bias_grad_reduce(
d, b, inv_batch, stream)?;
self.launch_bias_grad_reduce_ex(
self.d_proj_k_buf.raw_ptr(),
grad_base + (b_k_off * f32_sz) as u64,
d, b, inv_batch)?;
self.launch_bias_grad_reduce(
d, b, inv_batch, stream)?;
self.launch_bias_grad_reduce_ex(
self.d_proj_v_buf.raw_ptr(),
grad_base + (b_v_off * f32_sz) as u64,
d, b, inv_batch)?;
d, b, inv_batch, stream)?;
// 8. d_input = W_Q @ d_Q + W_K @ d_K + W_V @ d_V
// First: d_input = W_Q @ d_Q (beta=0, overwrites)
self.lt_matmul(&self.gemm_bwd_dx_q,
self.lt_matmul_ex(&self.gemm_bwd_dx_q,
params_base + (w_q_off * f32_sz) as u64,
self.d_proj_q_buf.raw_ptr(),
self.d_input_scratch.raw_ptr(),
1.0, 0.0, "bwd_dx_q")?;
1.0, 0.0, "bwd_dx_q", stream, ws_ptr, ws_size)?;
// Accumulate: d_input += W_K @ d_K (beta=1)
self.lt_matmul(&self.gemm_bwd_dx_k,
self.lt_matmul_ex(&self.gemm_bwd_dx_k,
params_base + (w_k_off * f32_sz) as u64,
self.d_proj_k_buf.raw_ptr(),
self.d_input_scratch.raw_ptr(),
1.0, 1.0, "bwd_dx_k")?;
1.0, 1.0, "bwd_dx_k", stream, ws_ptr, ws_size)?;
// Accumulate: d_input += W_V @ d_V (beta=1)
self.lt_matmul(&self.gemm_bwd_dx_v,
self.lt_matmul_ex(&self.gemm_bwd_dx_v,
params_base + (w_v_off * f32_sz) as u64,
self.d_proj_v_buf.raw_ptr(),
self.d_input_scratch.raw_ptr(),
1.0, 1.0, "bwd_dx_v")?;
1.0, 1.0, "bwd_dx_v", stream, ws_ptr, ws_size)?;
Ok(())
}
@@ -748,7 +762,17 @@ impl GpuAttention {
/// Must be called after `backward()`. Uses the same grad_norm + Adam
/// kernel pattern as IQN (separate from DQN trunk Adam).
pub fn adam_step(&mut self, lr: f32, max_grad_norm: f32) -> Result<(), MLError> {
pub fn adam_step(
&mut self,
lr: f32,
max_grad_norm: f32,
override_stream: Option<&Arc<CudaStream>>,
override_workspace: Option<(u64, usize)>,
) -> Result<(), MLError> {
let stream = override_stream.unwrap_or(&self.stream);
// override_workspace is accepted for API symmetry; adam_step does not use
// cublasLt workspace, but we accept it so callers pass a uniform tuple.
let _override_workspace = override_workspace;
let tp = self.total_params;
let tp_i32 = tp as i32;
@@ -761,7 +785,7 @@ impl GpuAttention {
};
// Phase 1: per-block partial sums
unsafe {
self.stream
stream
.launch_builder(&self.grad_norm_phase1_kernel)
.arg(&self.d_params)
.arg(&mut self.grad_norm_block_sums)
@@ -772,7 +796,7 @@ impl GpuAttention {
// Phase 2: deterministic reduction -> L2 norm
let norm_blocks_i32 = norm_blocks as i32;
unsafe {
self.stream
stream
.launch_builder(&self.grad_norm_phase2_kernel)
.arg(&self.grad_norm_block_sums)
.arg(&mut self.grad_norm_buf)
@@ -798,7 +822,7 @@ impl GpuAttention {
let t_ptr = self.t_dev_ptr;
unsafe {
self.stream
stream
.launch_builder(&self.adam_kernel)
.arg(&mut self.params)
.arg(&mut self.d_params)
@@ -832,8 +856,10 @@ impl GpuAttention {
// ── Private helpers ─────────────────────────────────────────────────────
/// Execute cublasLtMatmul with pre-cached GEMM descriptor.
fn lt_matmul(
/// Execute cublasLtMatmul with explicit stream and workspace (used by public
/// methods that accept override_stream / override_workspace).
#[allow(clippy::too_many_arguments)]
fn lt_matmul_ex(
&self,
desc: &AttnGemmDesc,
a_ptr: u64,
@@ -842,9 +868,12 @@ impl GpuAttention {
alpha: f32,
beta: f32,
label: &str,
stream: &CudaStream,
ws_ptr: u64,
ws_size: usize,
) -> Result<(), MLError> {
unsafe {
let cu_stream = self.stream.cu_stream() as cublaslt_sys::cudaStream_t;
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
let matmul_status = cublaslt_sys::cublasLtMatmul(
self.shared_handle.lt_handle.0,
desc.matmul_desc,
@@ -859,8 +888,8 @@ impl GpuAttention {
c_ptr as *mut std::ffi::c_void,
desc.d_layout,
&desc.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
self.shared_handle.lt_workspace_ptr as *mut std::ffi::c_void,
self.shared_handle.lt_workspace_size,
ws_ptr as *mut std::ffi::c_void,
ws_size,
cu_stream,
);
if matmul_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
@@ -873,13 +902,32 @@ impl GpuAttention {
Ok(())
}
/// Launch bias_add kernel: x[idx] += bias[idx % out_dim]
fn launch_bias_add(&self, out_ptr: u64, bias_ptr: u64, out_dim: usize, batch: usize) -> Result<(), MLError> {
/// Execute cublasLtMatmul with the default stream and workspace from shared_handle.
fn lt_matmul(
&self,
desc: &AttnGemmDesc,
a_ptr: u64,
b_ptr: u64,
c_ptr: u64,
alpha: f32,
beta: f32,
label: &str,
) -> Result<(), MLError> {
self.lt_matmul_ex(
desc, a_ptr, b_ptr, c_ptr, alpha, beta, label,
&self.stream,
self.shared_handle.lt_workspace_ptr,
self.shared_handle.lt_workspace_size,
)
}
/// Launch bias_add kernel with an explicit stream.
fn launch_bias_add_ex(&self, out_ptr: u64, bias_ptr: u64, out_dim: usize, batch: usize, stream: &CudaStream) -> Result<(), MLError> {
let total = (batch * out_dim) as i32;
let out_dim_i32 = out_dim as i32;
let blocks = ((batch * out_dim + 255) / 256) as u32;
unsafe {
self.stream
stream
.launch_builder(&self.shared_handle.add_bias_f32_kernel)
.arg(&out_ptr)
.arg(&bias_ptr)
@@ -895,13 +943,18 @@ impl GpuAttention {
Ok(())
}
/// Launch bias gradient reduction kernel: d_bias[j] = sum(d_pre[j + b*out_dim]) * inv_batch
fn launch_bias_grad_reduce(&self, d_pre_ptr: u64, d_bias_ptr: u64, out_dim: usize, batch: usize, inv_batch: f32) -> Result<(), MLError> {
/// Launch bias_add kernel with the default stream.
fn launch_bias_add(&self, out_ptr: u64, bias_ptr: u64, out_dim: usize, batch: usize) -> Result<(), MLError> {
self.launch_bias_add_ex(out_ptr, bias_ptr, out_dim, batch, &self.stream)
}
/// Launch bias gradient reduction kernel with an explicit stream.
fn launch_bias_grad_reduce_ex(&self, d_pre_ptr: u64, d_bias_ptr: u64, out_dim: usize, batch: usize, inv_batch: f32, stream: &CudaStream) -> Result<(), MLError> {
let out_dim_i32 = out_dim as i32;
let batch_i32 = batch as i32;
let blocks = ((out_dim + 255) / 256) as u32;
unsafe {
self.stream
stream
.launch_builder(&self.bias_grad_reduce_kernel)
.arg(&d_pre_ptr)
.arg(&d_bias_ptr)
@@ -917,6 +970,11 @@ impl GpuAttention {
}
Ok(())
}
/// Launch bias gradient reduction kernel with the default stream.
fn launch_bias_grad_reduce(&self, d_pre_ptr: u64, d_bias_ptr: u64, out_dim: usize, batch: usize, inv_batch: f32) -> Result<(), MLError> {
self.launch_bias_grad_reduce_ex(d_pre_ptr, d_bias_ptr, out_dim, batch, inv_batch, &self.stream)
}
}
// SAFETY: GpuAttention is only accessed from the training thread that owns

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

@@ -3917,7 +3917,7 @@ impl GpuDqnTrainer {
}
// Run attention forward: f32 scratch → attention output buffer (f32)
let attn_output = attention.forward(&self.attn_bw_scratch, batch_size)?;
let attn_output = attention.forward(&self.attn_bw_scratch, batch_size, None, None)?;
// Copy attention output (f32) back → save_h_s2 (f32)
let attn_src = attn_output.raw_ptr();

95
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -645,7 +645,7 @@ impl GpuIqnHead {
dtod_copy(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?;
// 3-9. Execute core training pipeline
self.execute_training_pipeline(online_h_s2, next_states_buf, target_dueling)
self.execute_training_pipeline(online_h_s2, next_states_buf, target_dueling, None, None)
}
/// Set a pre-computed target h_s2 buffer from the DQN trainer's Pass 2.
@@ -658,12 +658,21 @@ impl GpuIqnHead {
/// Core IQN training pipeline: trunk forward → cuBLAS forward → loss → backward → Adam.
///
/// Assumes `branch_actions`, `rewards_buf`, `dones_buf` are already populated.
fn execute_training_pipeline(
///
/// `override_stream` and `override_workspace` allow a dedicated CUDA stream and
/// cuBLAS workspace to be used instead of the IQN head's defaults. Pass `None`
/// for both to preserve the existing behaviour (used by all current callers).
pub fn execute_training_pipeline(
&mut self,
online_h_s2: &CudaSlice<f32>,
next_states_buf: &CudaSlice<f32>,
target_dueling: &DuelingWeightSet,
override_stream: Option<&Arc<CudaStream>>,
override_workspace: Option<(u64, usize)>,
) -> Result<f32, MLError> {
let effective_stream = override_stream.unwrap_or(&self.stream);
let (lt_ws_ptr, lt_ws_size) = override_workspace
.unwrap_or((self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size));
let b = self.config.batch_size;
let n = self.config.num_quantiles;
let h = self.config.hidden_dim;
@@ -687,7 +696,7 @@ impl GpuIqnHead {
let n_bytes = b * h * std::mem::size_of::<f32>();
let dst = self.target_h_s2.raw_ptr();
unsafe {
let r = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(dst, cached_ptr, n_bytes, self.stream.cu_stream());
let r = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(dst, cached_ptr, n_bytes, effective_stream.cu_stream());
if r != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
return Err(MLError::ModelError(format!("IQN cached h_s2 DtoD: {r:?}")));
}
@@ -701,7 +710,7 @@ impl GpuIqnHead {
};
let state_dim_i32 = self.config.state_dim as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.trunk_forward_kernel)
.arg(next_states_buf)
.arg(&target_dueling.w_s1)
@@ -737,7 +746,7 @@ impl GpuIqnHead {
// Tile online h_s2
unsafe {
self.stream
effective_stream
.launch_builder(&self.h_s2_tile_kernel)
.arg(online_h_s2)
.arg(&mut self.h_s2_tiled_buf)
@@ -750,7 +759,7 @@ impl GpuIqnHead {
// Tile cos_features [D, N] → [D, B*Q]
unsafe {
self.stream
effective_stream
.launch_builder(&self.cos_tile_kernel)
.arg(&self.cos_features)
.arg(&mut self.cos_features_tiled_buf)
@@ -780,13 +789,11 @@ impl GpuIqnHead {
}
let lt_handle_raw = self.shared_handle.lt_handle.0;
let lt_ws_ptr = self.shared_handle.lt_workspace_ptr;
let lt_ws_size = self.shared_handle.lt_workspace_size;
// ── Step 3: Forward embedding GEMM ──────────────────────────────
// embed_pre[H, B*Q] = W_embed^T @ cos_features_tiled
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_fwd_embed,
effective_stream, lt_handle_raw, &self.gemm_fwd_embed,
1.0, 0.0,
w_embed_ptr, self.cos_features_tiled_buf.raw_ptr(),
self.embed_pre_buf.raw_ptr(),
@@ -797,7 +804,7 @@ impl GpuIqnHead {
let hbq = h * bq;
let hbq_i32 = hbq as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.bias_add_kernel)
.arg(&mut self.embed_pre_buf)
.arg(&b_embed_ptr)
@@ -807,7 +814,7 @@ impl GpuIqnHead {
.map_err(|e| MLError::ModelError(format!("IQN bias_add embed: {e}")))?;
}
unsafe {
self.stream
effective_stream
.launch_builder(&self.relu_fwd_kernel)
.arg(&self.embed_pre_buf)
.arg(&mut self.embed_buf)
@@ -818,7 +825,7 @@ impl GpuIqnHead {
// ── Step 4: Hadamard product: combined = h_s2_tiled * sigmoid(embed) ──
unsafe {
self.stream
effective_stream
.launch_builder(&self.hadamard_sigmoid_kernel)
.arg(&self.h_s2_tiled_buf)
.arg(&self.embed_buf)
@@ -836,7 +843,7 @@ impl GpuIqnHead {
let bs = branch_sizes[i];
let out_ptr = self.branch_logits_buf.raw_ptr() + (row_off * bq) as u64 * f32_sz;
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_fwd_branch[i],
effective_stream, lt_handle_raw, &self.gemm_fwd_branch[i],
1.0, 0.0,
branch_w_ptrs[i], self.combined_buf.raw_ptr(),
out_ptr,
@@ -847,7 +854,7 @@ impl GpuIqnHead {
let bs_i32 = bs as i32;
let total_i32 = (bs * bq) as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.bias_add_kernel)
.arg(&out_ptr)
.arg(&branch_b_ptrs[i])
@@ -866,7 +873,7 @@ impl GpuIqnHead {
let src = self.branch_logits_buf.raw_ptr();
let dst = self.save_q_online.raw_ptr();
unsafe {
let r = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(dst, src, n_bytes, self.stream.cu_stream());
let r = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(dst, src, n_bytes, effective_stream.cu_stream());
if r != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
return Err(MLError::ModelError(format!("IQN save_q_online DtoD: {r:?}")));
}
@@ -881,7 +888,7 @@ impl GpuIqnHead {
// Tile target_h_s2 into h_s2_tiled_buf (reuse buffer)
unsafe {
self.stream
effective_stream
.launch_builder(&self.h_s2_tile_kernel)
.arg(&self.target_h_s2)
.arg(&mut self.h_s2_tiled_buf)
@@ -894,14 +901,14 @@ impl GpuIqnHead {
// Target embedding forward
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_fwd_embed,
effective_stream, lt_handle_raw, &self.gemm_fwd_embed,
1.0, 0.0,
tw_embed_ptr, self.cos_features_tiled_buf.raw_ptr(),
self.embed_pre_buf.raw_ptr(), // reuse for target
lt_ws_ptr, lt_ws_size, "iqn_tgt_fwd_embed",
)?;
unsafe {
self.stream
effective_stream
.launch_builder(&self.bias_add_kernel)
.arg(&mut self.embed_pre_buf)
.arg(&tb_embed_ptr)
@@ -912,7 +919,7 @@ impl GpuIqnHead {
}
// ReLU into a scratch — reuse embed_buf (we saved online logits already)
unsafe {
self.stream
effective_stream
.launch_builder(&self.relu_fwd_kernel)
.arg(&self.embed_pre_buf)
.arg(&mut self.embed_buf)
@@ -922,7 +929,7 @@ impl GpuIqnHead {
}
// Target combined
unsafe {
self.stream
effective_stream
.launch_builder(&self.hadamard_sigmoid_kernel)
.arg(&self.h_s2_tiled_buf)
.arg(&self.embed_buf)
@@ -944,7 +951,7 @@ impl GpuIqnHead {
let out_ptr = self.target_branch_logits_buf.raw_ptr() + (row_off * bq) as u64 * f32_sz;
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_fwd_branch[i],
effective_stream, lt_handle_raw, &self.gemm_fwd_branch[i],
1.0, 0.0,
tw_ptr, self.combined_buf.raw_ptr(),
out_ptr,
@@ -954,7 +961,7 @@ impl GpuIqnHead {
let bs_i32 = bs as i32;
let total_i32 = (bs * bq) as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.bias_add_kernel)
.arg(&out_ptr)
.arg(&tb_ptr)
@@ -969,7 +976,7 @@ impl GpuIqnHead {
// ── Step 6: Quantile Huber loss + dq gradients ──────────────────
// Zero d_branch_logits_buf before loss kernel writes to it
self.stream.memset_zeros(&mut self.d_branch_logits_buf)
effective_stream.memset_zeros(&mut self.d_branch_logits_buf)
.map_err(|e| MLError::ModelError(format!("IQN zero d_branch_logits: {e}")))?;
let kappa = self.config.kappa;
@@ -977,7 +984,7 @@ impl GpuIqnHead {
let batch_size_i32 = b as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.quantile_huber_loss_kernel)
.arg(&self.branch_logits_buf)
.arg(&self.target_branch_logits_buf)
@@ -1004,10 +1011,10 @@ impl GpuIqnHead {
// Loss reduce → pinned device-mapped buffer (zero-copy readback, no sync)
unsafe {
cudarc::driver::sys::cuMemsetD8Async(
self.total_loss_dev_ptr, 0, std::mem::size_of::<f32>(), self.stream.cu_stream(),
self.total_loss_dev_ptr, 0, std::mem::size_of::<f32>(), effective_stream.cu_stream(),
);
let loss_ptr = self.total_loss_dev_ptr;
self.stream
effective_stream
.launch_builder(&self.loss_reduce_kernel)
.arg(&self.per_sample_loss)
.arg(&loss_ptr)
@@ -1021,7 +1028,7 @@ impl GpuIqnHead {
// from the ONLINE forward pass. But we overwrote them with target.
// Re-tile online h_s2 and redo online embedding forward to recover saves.
unsafe {
self.stream
effective_stream
.launch_builder(&self.h_s2_tile_kernel)
.arg(online_h_s2)
.arg(&mut self.h_s2_tiled_buf)
@@ -1033,14 +1040,14 @@ impl GpuIqnHead {
}
// Redo online embedding forward to recover embed_pre_buf + embed_buf + combined_buf
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_fwd_embed,
effective_stream, lt_handle_raw, &self.gemm_fwd_embed,
1.0, 0.0,
w_embed_ptr, self.cos_features_tiled_buf.raw_ptr(),
self.embed_pre_buf.raw_ptr(),
lt_ws_ptr, lt_ws_size, "iqn_bwd_recover_embed",
)?;
unsafe {
self.stream
effective_stream
.launch_builder(&self.bias_add_kernel)
.arg(&mut self.embed_pre_buf)
.arg(&b_embed_ptr)
@@ -1050,7 +1057,7 @@ impl GpuIqnHead {
.map_err(|e| MLError::ModelError(format!("IQN bwd recover bias_add: {e}")))?;
}
unsafe {
self.stream
effective_stream
.launch_builder(&self.relu_fwd_kernel)
.arg(&self.embed_pre_buf)
.arg(&mut self.embed_buf)
@@ -1059,7 +1066,7 @@ impl GpuIqnHead {
.map_err(|e| MLError::ModelError(format!("IQN bwd recover relu: {e}")))?;
}
unsafe {
self.stream
effective_stream
.launch_builder(&self.hadamard_sigmoid_kernel)
.arg(&self.h_s2_tiled_buf)
.arg(&self.embed_buf)
@@ -1073,10 +1080,10 @@ impl GpuIqnHead {
let inv_bq = 1.0_f32 / bq as f32;
// Zero grad_buf before accumulating weight gradients
self.stream.memset_zeros(&mut self.grad_buf)
effective_stream.memset_zeros(&mut self.grad_buf)
.map_err(|e| MLError::ModelError(format!("IQN zero grad_buf: {e}")))?;
// Zero d_combined_buf for accumulation
self.stream.memset_zeros(&mut self.d_combined_buf)
effective_stream.memset_zeros(&mut self.d_combined_buf)
.map_err(|e| MLError::ModelError(format!("IQN zero d_combined: {e}")))?;
{
@@ -1089,7 +1096,7 @@ impl GpuIqnHead {
// dW_bk = dq_bk @ combined^T (alpha=inv_bq for mean reduction)
let dw_ptr = self.grad_buf.raw_ptr() + grad_off as u64 * f32_sz;
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_bwd_dw_branch[i],
effective_stream, lt_handle_raw, &self.gemm_bwd_dw_branch[i],
inv_bq, 0.0,
dq_ptr, self.combined_buf.raw_ptr(),
dw_ptr,
@@ -1103,7 +1110,7 @@ impl GpuIqnHead {
let bs_i32 = bs as i32;
let bq_i32 = bq as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.bias_grad_reduce_kernel)
.arg(&dq_ptr)
.arg(&db_ptr)
@@ -1117,7 +1124,7 @@ impl GpuIqnHead {
// d_combined += W_bk @ dq_bk (beta=1.0 for accumulation)
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_bwd_dx_branch[i],
effective_stream, lt_handle_raw, &self.gemm_bwd_dx_branch[i],
1.0, 1.0, // beta=1.0 to accumulate across branches
branch_w_ptrs[i], dq_ptr,
self.d_combined_buf.raw_ptr(),
@@ -1131,7 +1138,7 @@ impl GpuIqnHead {
// Step 7b: Backward through hadamard_sigmoid → d_embed, d_h_s2_tiled
unsafe {
self.stream
effective_stream
.launch_builder(&self.hadamard_sigmoid_bwd_kernel)
.arg(&self.h_s2_tiled_buf)
.arg(&self.embed_buf)
@@ -1145,7 +1152,7 @@ impl GpuIqnHead {
// Step 7c: Backward through ReLU
unsafe {
self.stream
effective_stream
.launch_builder(&self.relu_bwd_kernel)
.arg(&self.embed_pre_buf)
.arg(&self.d_embed_buf)
@@ -1158,7 +1165,7 @@ impl GpuIqnHead {
// Step 7d: dW_embed = d_embed_pre @ cos_features_tiled^T (alpha=inv_bq)
let dw_embed_ptr = self.grad_buf.raw_ptr(); // offset 0 in grad_buf
iqn_lt_matmul(
&self.stream, lt_handle_raw, &self.gemm_bwd_dw_embed,
effective_stream, lt_handle_raw, &self.gemm_bwd_dw_embed,
inv_bq, 0.0,
self.d_embed_pre_buf.raw_ptr(), self.cos_features_tiled_buf.raw_ptr(),
dw_embed_ptr,
@@ -1169,7 +1176,7 @@ impl GpuIqnHead {
let db_embed_ptr = self.grad_buf.raw_ptr() + (h * d) as u64 * f32_sz;
let bq_i32 = bq as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.bias_grad_reduce_kernel)
.arg(&self.d_embed_pre_buf.raw_ptr())
.arg(&db_embed_ptr)
@@ -1182,7 +1189,7 @@ impl GpuIqnHead {
// Step 7f: Reduce d_h_s2_tiled [H, B*Q] → d_h_s2 [H, B]
unsafe {
self.stream
effective_stream
.launch_builder(&self.d_h_s2_reduce_kernel)
.arg(&self.d_h_s2_tiled_buf)
.arg(&mut self.d_h_s2_buf)
@@ -1198,7 +1205,7 @@ impl GpuIqnHead {
let norm_blocks = (self.total_params + 255) / 256;
unsafe {
self.stream
effective_stream
.launch_builder(&self.grad_norm_phase1_kernel)
.arg(&self.grad_buf)
.arg(&mut self.grad_norm_block_sums)
@@ -1212,7 +1219,7 @@ impl GpuIqnHead {
}
let norm_blocks_i32 = norm_blocks as i32;
unsafe {
self.stream
effective_stream
.launch_builder(&self.grad_norm_phase2_kernel)
.arg(&self.grad_norm_block_sums)
.arg(&mut self.grad_norm_buf)
@@ -1240,7 +1247,7 @@ impl GpuIqnHead {
let t_ptr = self.t_dev_ptr;
unsafe {
self.stream
effective_stream
.launch_builder(&self.adam_kernel)
.arg(&mut self.online_params)
.arg(&mut self.grad_buf)

4
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -1311,11 +1311,11 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("Attention forward: {e}"))?;
let d_h_s2_ref = self.trainer.bw_d_h_s2()
.map_err(|e| anyhow::anyhow!("bw_d_h_s2 ref: {e}"))?;
attn.backward(d_h_s2_ref, self.batch_size)
attn.backward(d_h_s2_ref, self.batch_size, None, None)
.map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?;
let lr = self.trainer.config().lr;
let mgn = self.trainer.config().max_grad_norm;
attn.adam_step(lr, mgn)
attn.adam_step(lr, mgn, None, None)
.map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?;
}