diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 795e691af..ec01617cc 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3090,14 +3090,6 @@ pub struct GpuDqnTrainer { denoise_bias_grad_p1_kernel: CudaFunction, /// Kernel handle: denoise_bias_grad_p2. denoise_bias_grad_p2_kernel: CudaFunction, - /// Dedicated CUfunction handles for OFI bias-grad reduction. - /// Loaded from a SEPARATE CUmodule from the denoise variant — each - /// child graph (post_aux, adam_grad) must hold its own CUfunction - /// reference, otherwise the driver's per-node handle validation on - /// Ada/Hopper trips CUDA_ERROR_ILLEGAL_ADDRESS at kernel launch - /// (same pattern as graph_safe_copy_kernel below). - ofi_embed_bias_grad_p1_kernel: CudaFunction, - ofi_embed_bias_grad_p2_kernel: CudaFunction, // ── xLSTM temporal context (mLSTM cell) ──────────────────────────── /// [528] 6 weight matrices × [11, 8]: W_k, W_v, W_q, W_i, W_f, W_o @@ -5188,7 +5180,7 @@ impl GpuDqnTrainer { let grad_b_ptr = self.ofi_embed_grad_b.raw_ptr(); let dim_i32 = OFI_EMBED_DIM as i32; unsafe { - self.stream.launch_builder(&self.ofi_embed_bias_grad_p1_kernel) + self.stream.launch_builder(&self.denoise_bias_grad_p1_kernel) .arg(&d_combined_ptr) .arg(&partials_ptr) .arg(&dim_i32) @@ -5201,7 +5193,7 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad_p1: {e}")))?; } unsafe { - self.stream.launch_builder(&self.ofi_embed_bias_grad_p2_kernel) + self.stream.launch_builder(&self.denoise_bias_grad_p2_kernel) .arg(&partials_ptr) .arg(&grad_b_ptr) .arg(&dim_i32) @@ -11105,16 +11097,6 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p1 load: {e}")))?; let denoise_bias_grad_p2_kernel = cpbi_module_denoise_cublas.load_function("denoise_bias_grad_p2") .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p2 load: {e}")))?; - // Dedicated CUmodule for OFI bias-grad reduction (Ada/Hopper child-graph - // CUfunction isolation — same pattern as graph_safe_copy below). - let ofi_embed_bias_grad_module = stream.context().load_cubin(EXPECTED_Q_CUBIN.to_vec()) - .map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad module load: {e}")))?; - let ofi_embed_bias_grad_p1_kernel = ofi_embed_bias_grad_module - .load_function("denoise_bias_grad_p1") - .map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad_p1 load: {e}")))?; - let ofi_embed_bias_grad_p2_kernel = ofi_embed_bias_grad_module - .load_function("denoise_bias_grad_p2") - .map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad_p2 load: {e}")))?; // cuBLAS GEMM descriptors for denoise MLP backward let denoise_lt_handle = shared_cublas.lt_handle.0; @@ -12055,8 +12037,6 @@ impl GpuDqnTrainer { denoise_skip_conn_kernel, denoise_bias_grad_p1_kernel, denoise_bias_grad_p2_kernel, - ofi_embed_bias_grad_p1_kernel, - ofi_embed_bias_grad_p2_kernel, qlstm_weights, qlstm_c, qlstm_n, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index d1dd14685..e04d13b73 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -858,15 +858,3 @@ architecture pretended to do. Workspace clean: 0 errors across all crates, tests Smoke validates no regression: 3/3 folds passed (405s), gate utilization preserved (util=0.119,0.119,0.169,...), HEALTH_DIAG aux_moe line emits, all fold checkpoints written. - -OFI bias-grad CUfunction isolation (2026-04-27): fixed -CUDA_ERROR_ILLEGAL_ADDRESS crash in launch_ofi_embed_backward on L40S -(workflow train-multi-seed-tdjtr, fold 0). Root cause: shared -CUfunction handles for denoise_bias_grad_p1/p2 across two captured -child graphs (post_aux, adam_grad) trip Ada/Hopper driver per-node -handle validation when the partials buffers land far apart in VRAM. -Pre-existing latent bug, surfaced by MoE allocator-state shift. -Fix: dedicated CUmodule for OFI bias-grad kernels, mirroring the -graph_safe_copy_kernel precedent (line ~11523). Kernel bytecode -identical; constructor adds ~2ms; zero runtime cost. Smoke validates -no regression at B=64.