diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index 55cfc7498..3b6953507 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -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, batch_size: usize) -> Result<&CudaSlice, MLError> { + pub fn forward( + &mut self, + states: &CudaSlice, + batch_size: usize, + override_stream: Option<&Arc>, + override_workspace: Option<(u64, usize)>, + ) -> Result<&CudaSlice, 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::(); 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, batch_size: usize, + override_stream: Option<&Arc>, + 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>, + 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 diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index f62b75d07..bfece2484 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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(); diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 24471c787..1c03d514e 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -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, next_states_buf: &CudaSlice, target_dueling: &DuelingWeightSet, + override_stream: Option<&Arc>, + override_workspace: Option<(u64, usize)>, ) -> Result { + 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::(); 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::(), self.stream.cu_stream(), + self.total_loss_dev_ptr, 0, std::mem::size_of::(), 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) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 340d735a0..e93e17599 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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}"))?; }