From 36d02b8f3cf4fefb6bcf7bc8272a6805bf102a21 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 20 Apr 2026 01:03:31 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20Mamba2=20d=5Fh=5Fhistory=20backward=20?= =?UTF-8?q?=E2=80=94=20gradient=20flows=20to=20OFI=20embed=20MLP?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 2 new cuBLAS GEMMs: d_h_history = W_A^T @ d_gate + W_B^T @ d_x. Extracts last 10 dims as d_ofi_embed_mamba2, accumulated across K=8 timesteps. This enables the OFI embed MLP to learn from temporal patterns discovered by the Mamba2 SSM. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 113 ++++++++++++++++++ .../cuda_pipeline/mamba2_temporal_kernel.cu | 33 +++++ 2 files changed, 146 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 52c5a1fa1..443147874 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1497,6 +1497,12 @@ pub struct GpuDqnTrainer { mamba2_gemm_proj: Mamba2GemmDesc, // fwd: W^T @ h, m=STATE_D, n=B*K, k=SH2+OFI mamba2_gemm_dw: Mamba2GemmDesc, // bwd dW_A/dW_B: d^T @ h, m=STATE_D, n=SH2+OFI, k=B*K mamba2_gemm_dwc: Mamba2GemmDesc, // bwd dW_C: xK^T @ dtw, m=STATE_D, n=SH2, k=B + mamba2_gemm_dh: Mamba2GemmDesc, // bwd d_h_history: W^T @ d, m=SH2+OFI, n=B*K, k=STATE_D + + // ── Mamba2 input gradient buffers ── + d_h_history_buf: CudaSlice, // [(SH2+OFI) * B * K] backward scratch + d_ofi_embed_mamba2: CudaSlice, // [OFI_EMBED_DIM * B] accumulated OFI embed gradient + extract_ofi_embed_grad_kernel: CudaFunction, // ── Mamba2 cuBLAS scan kernels ── mamba2_scan_proj_fwd_kernel: CudaFunction, @@ -3503,6 +3509,82 @@ impl GpuDqnTrainer { ); } + // Step 6: d_h_history = W_A^T @ d_gate + W_B^T @ d_x + // This backpropagates through the W_A/W_B projections to get the input gradient. + // W_A col-major is [STATE_D, h_width] with ld=STATE_D. + // TRANSA=T gives [h_width, STATE_D], multiplied by d_gate[STATE_D, B*K]. + // Result: d_h_history[h_width, B*K] col-major. + let w_a_ptr = param_ptr; + let w_b_ptr = param_ptr + (h_width * MAMBA2_STATE_DIM * 4) as u64; + let d_h_ptr = self.d_h_history_buf.raw_ptr(); + + // GEMM 6a: d_h_history = W_A^T @ d_gate (beta=0, overwrite) + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.mamba2_gemm_dh.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w_a_ptr as *const std::ffi::c_void, + self.mamba2_gemm_dh.a_layout, + d_gate_ptr as *const std::ffi::c_void, + self.mamba2_gemm_dh.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + d_h_ptr as *mut std::ffi::c_void, + self.mamba2_gemm_dh.c_layout, + d_h_ptr as *mut std::ffi::c_void, + self.mamba2_gemm_dh.d_layout, + &self.mamba2_gemm_dh.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // GEMM 6b: d_h_history += W_B^T @ d_x (beta=1, accumulate) + let beta_one: f32 = 1.0; + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.mamba2_gemm_dh.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w_b_ptr as *const std::ffi::c_void, + self.mamba2_gemm_dh.a_layout, + d_x_ptr as *const std::ffi::c_void, + self.mamba2_gemm_dh.b_layout, + &beta_one as *const f32 as *const std::ffi::c_void, + d_h_ptr as *mut std::ffi::c_void, + self.mamba2_gemm_dh.c_layout, + d_h_ptr as *mut std::ffi::c_void, + self.mamba2_gemm_dh.d_layout, + &self.mamba2_gemm_dh.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // Step 7: Extract d_ofi_embed from d_h_history (last OFI_EMBED_DIM rows, sum over K) + let d_ofi_ptr = self.d_ofi_embed_mamba2.raw_ptr(); + let sh2_i32 = sh2 as i32; + let embed_dim_i32 = OFI_EMBED_DIM as i32; + let b_i32 = batch_size as i32; + let k_i32 = MAMBA2_HISTORY_K as i32; + unsafe { + self.stream.launch_builder(&self.extract_ofi_embed_grad_kernel) + .arg(&d_h_ptr) + .arg(&d_ofi_ptr) + .arg(&b_i32) + .arg(&k_i32) + .arg(&sh2_i32) + .arg(&embed_dim_i32) + .launch(LaunchConfig { + grid_dim: (((batch_size + 255) / 256) as u32, OFI_EMBED_DIM as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("extract_ofi_embed_grad: {e}")))?; + } + Ok(()) } @@ -3532,6 +3614,15 @@ impl GpuDqnTrainer { Ok(()) } + /// Reference to the Mamba2 OFI embed gradient buffer on GPU. + /// + /// Shape: `[OFI_EMBED_DIM, B]` col-major — accumulated d_ofi_embed from `mamba2_backward()`. + /// Written by Step 7 (extract_ofi_embed_grad kernel) after d_h_history GEMMs. + /// Read by Task 8 (OFI embed MLP backward) to accumulate with the attention gradient. + pub(crate) fn d_ofi_embed_mamba2(&self) -> &CudaSlice { + &self.d_ofi_embed_mamba2 + } + /// Reference to the states buffer on GPU. /// /// Shape: `[B, STATE_DIM]` — contains the batch's states after `upload_batch_gpu()`. @@ -6244,6 +6335,24 @@ impl GpuDqnTrainer { lt_ws_size, "mamba2_bwd_dwc", )?; + // Backward d_h_history: C[h_width, B*K] = W_cm^T[h_width, STATE_D] @ d[STATE_D, B*K] + // W_A/W_B col-major layout is [STATE_D, h_width] with ld=STATE_D. + // TRANSA=T transposes to [h_width, STATE_D]. TRANSB=N on d[STATE_D, B*K]. + let mamba2_gemm_dh = create_mamba2_gemm_desc( + lt_handle, 1, 0, + h_width, bk, MAMBA2_STATE_DIM, + MAMBA2_STATE_DIM, MAMBA2_STATE_DIM, h_width, + lt_ws_size, "mamba2_bwd_dh", + )?; + + // ── Mamba2 input gradient buffers ── + let d_h_history_buf = stream.alloc_zeros::(h_width * bk) + .map_err(|e| MLError::ModelError(format!("alloc d_h_history_buf: {e}")))?; + let d_ofi_embed_mamba2 = stream.alloc_zeros::(OFI_EMBED_DIM * b) + .map_err(|e| MLError::ModelError(format!("alloc d_ofi_embed_mamba2: {e}")))?; + let extract_ofi_embed_grad_kernel = mamba2_module.load_function("extract_ofi_embed_grad") + .map_err(|e| MLError::ModelError(format!("extract_ofi_embed_grad load: {e}")))?; + info!( bk, state_dim = MAMBA2_STATE_DIM, sh2, h_width, a_proj_bytes = bk * MAMBA2_STATE_DIM * 4, @@ -6812,6 +6921,10 @@ impl GpuDqnTrainer { mamba2_gemm_proj, mamba2_gemm_dw, mamba2_gemm_dwc, + mamba2_gemm_dh, + d_h_history_buf, + d_ofi_embed_mamba2, + extract_ofi_embed_grad_kernel, mamba2_scan_proj_fwd_kernel, mamba2_scan_proj_bwd_kernel, mamba2_scale_d_enriched_kernel, diff --git a/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu b/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu index fc4c011b9..745fae7cb 100644 --- a/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu @@ -246,3 +246,36 @@ extern "C" __global__ void mamba2_scale_d_enriched( float tw = (temporal_weight != NULL) ? temporal_weight[j] : 1.0f; d_tw[idx] = d_h_enriched[idx] * tw; } + + +/* ================================================================== */ +/* Kernel: extract_ofi_embed_grad */ +/* ================================================================== */ + +/** + * Extract last embed_dim rows of d_h_history and accumulate across K + * timesteps to produce d_ofi_embed[embed_dim, B]. + * + * d_h_history is col-major [h_width, B*K] where h_width = SH2 + embed_dim. + * The OFI embed gradient lives in the last embed_dim rows (rows SH2..SH2+embed_dim-1). + * We sum over the K timestep dimension for each sample: + * d_ofi_embed[d, b] = sum_{k=0}^{K-1} d_h_history[(SH2+d), b*K+k] + * + * Grid: (ceil(B/256), embed_dim), Block: 256. + */ +extern "C" __global__ void extract_ofi_embed_grad( + const float* __restrict__ d_h_history, /* [h_width, B*K] col-major */ + float* __restrict__ d_ofi_embed, /* [embed_dim, B] col-major output */ + int B, int K, int sh2, int embed_dim +) { + int d = blockIdx.y; /* which embed dimension */ + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B || d >= embed_dim) return; + int h_width = sh2 + embed_dim; + float sum = 0.0f; + for (int k = 0; k < K; k++) { + /* d_h_history col-major [h_width, B*K]: element (row, col) = row + col * h_width */ + sum += d_h_history[(sh2 + d) + (long long)(b * K + k) * h_width]; + } + d_ofi_embed[d + b * embed_dim] = sum; +}