From 10d88e1b2aaa24b2aad444a4bfad9d84b2dfcb43 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 16 Apr 2026 00:27:17 +0200 Subject: [PATCH] fix: mamba2_backward used grad_buf (params) not bw_d_h_s2 (trunk activation gradient) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mamba2_scan_backward kernel reads d_h_enriched [B, SH2] but was passed self.grad_buf [TOTAL_PARAMS] — wrong buffer, wrong size. At batch_size=4096 the kernel read 1M floats from a 582K buffer → 2749 OOB reads. Fixed: use self.bw_d_h_s2 [B, SH2] which is the actual trunk activation gradient from the cuBLAS backward pass. Also increased smoke test batch_size to 4096 to catch scale-dependent OOB. compute-sanitizer: 0 errors at batch_size=4096. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 2 +- crates/ml/src/trainers/dqn/smoke_tests/generalization.rs | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 1b492be30..e3b7ebaec 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2276,7 +2276,7 @@ impl GpuDqnTrainer { self.stream.launch_builder(&self.mamba2_backward_kernel) .arg(&self.mamba2_h_history) .arg(&self.save_h_s2) - .arg(&self.grad_buf) // d_h_enriched = trunk gradient + .arg(&self.bw_d_h_s2) // d_h_enriched = trunk activation gradient [B, SH2] .arg(&w_a_ptr) .arg(&w_b_ptr) .arg(&w_c_ptr) diff --git a/crates/ml/src/trainers/dqn/smoke_tests/generalization.rs b/crates/ml/src/trainers/dqn/smoke_tests/generalization.rs index e386d9aff..8f6daad8b 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/generalization.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/generalization.rs @@ -35,6 +35,8 @@ fn test_generalization_components_smoke() -> anyhow::Result<()> { let mut params = smoke_params(); params.epochs = 3; + params.batch_size = 4096; // Larger batch to catch OOB that only trigger at scale + params.buffer_size = 8192; params.weight_decay = 1e-4; // G1: AdamW active let mut trainer = smoke_trainer_with(params)?;