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)?;