gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: Mamba2 d_h_history backward — gradient flows to OFI embed MLP

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 01:03:31 +02:00
parent 227d5cb5fb
commit 36d02b8f3c
2 changed files with 146 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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<f32>, // [(SH2+OFI) * B * K] backward scratch
d_ofi_embed_mamba2: CudaSlice<f32>, // [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<f32> {
&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::<f32>(h_width * bk)
.map_err(|e| MLError::ModelError(format!("alloc d_h_history_buf: {e}")))?;
let d_ofi_embed_mamba2 = stream.alloc_zeros::<f32>(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,

33
crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu
View File

@@ -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;
}