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

feat: Mamba2 history enrichment SH2→SH2+10 — temporal scan learns OFI momentum

Browse Source 
h_history widened from [B, K, SH2] to [B, K, SH2+10]. OFI embed (10-dim)
appended per timestep. W_A/W_B projections grow to [SH2+10, STATE_D].
W_C stays [SH2, STATE_D] — operates on scan output, not history input.

The SSM temporal scan now sees order flow momentum, book aggression, and
bar duration alongside trunk features. Enables learning sequential
patterns in microstructure dynamics.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 00:56:00 +02:00
parent 6575eea8da
commit 227d5cb5fb
2 changed files with 67 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -79,6 +79,7 @@ pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_
/// Mamba2 temporal scan configuration.
const MAMBA2_HISTORY_K: usize = 8; // Rolling history length
const MAMBA2_STATE_DIM: usize = 16; // SSM state dimension
const OFI_EMBED_DIM: usize = 10; // OFI embedding width appended to history rows
/// Introspective State Vector (ISV) configuration.
const ISV_K: usize = 4; // Temporal ISV history length
@@ -1474,14 +1475,14 @@ pub struct GpuDqnTrainer {
prev_q_mean_dev_ptr: u64,
// ── Mamba2 temporal scan ──
mamba2_h_history: CudaSlice<f32>, // [B, K, SH2] rolling history
mamba2_h_history: CudaSlice<f32>, // [B, K, SH2+OFI_EMBED_DIM] rolling history (widened for OFI)
mamba2_h_enriched: CudaSlice<f32>, // [B, SH2] output
mamba2_update_kernel: CudaFunction,
mamba2_copy_enriched_kernel: CudaFunction,
mamba2_params: CudaSlice<f32>, // [3 * SH2 * STATE_D] W_A, W_B, W_C concatenated
mamba2_grad: CudaSlice<f32>, // gradient buffer
mamba2_adam_m: CudaSlice<f32>, // Adam first moment
mamba2_adam_v: CudaSlice<f32>, // Adam second moment
mamba2_params: CudaSlice<f32>, // W_A[(SH2+OFI)*SD] + W_B[(SH2+OFI)*SD] + W_C[SH2*SD]
mamba2_grad: CudaSlice<f32>, // gradient buffer (same layout as params)
mamba2_adam_m: CudaSlice<f32>, // Adam first moment (same size as params)
mamba2_adam_v: CudaSlice<f32>, // Adam second moment (same size as params)
mamba2_adam_step: i32, // Adam step counter
// ── Mamba2 cuBLAS projection buffers ──
@@ -1493,8 +1494,8 @@ pub struct GpuDqnTrainer {
mamba2_d_tw: CudaSlice<f32>, // [B, SH2] scaled d_h_enriched * temporal_weight
// ── Mamba2 cuBLAS GEMM descriptors ──
mamba2_gemm_proj: Mamba2GemmDesc, // fwd: W^T @ h, m=STATE_D, n=B*K, k=SH2
mamba2_gemm_dw: Mamba2GemmDesc, // bwd dW_A/dW_B: d^T @ h, m=STATE_D, n=SH2, k=B*K
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 cuBLAS scan kernels ──
@@ -3222,10 +3223,11 @@ impl GpuDqnTrainer {
/// New: cuBLAS GEMM (SH2=256 -> STATE_D=16) + scan grid=(B, ceil(SH2/256)).
pub(crate) fn mamba2_forward(&self, batch_size: usize) -> Result<(), MLError> {
let sh2 = self.config.shared_h2;
let h_width = sh2 + OFI_EMBED_DIM;
let param_ptr = self.mamba2_params.raw_ptr();
let w_a_ptr = param_ptr;
let w_b_ptr = param_ptr + (sh2 * MAMBA2_STATE_DIM * 4) as u64;
let w_c_ptr = param_ptr + (2 * sh2 * MAMBA2_STATE_DIM * 4) as u64;
let w_b_ptr = param_ptr + (h_width * MAMBA2_STATE_DIM * 4) as u64;
let w_c_ptr = param_ptr + (2 * h_width * MAMBA2_STATE_DIM * 4) as u64;
let lt_handle = self.shared_cublas.lt_handle.0;
let lt_ws_ptr = self.shared_cublas.lt_workspace_ptr;
@@ -3234,8 +3236,8 @@ impl GpuDqnTrainer {
let alpha: f32 = 1.0;
let beta: f32 = 0.0;
// GEMM 1: A_proj[B*K, STATE_D] = h_history[B*K, SH2] @ W_A[SH2, STATE_D]
// Col-major: C[STATE_D, B*K] = W_A_cm[STATE_D, SH2] @ h_cm[SH2, B*K]
// GEMM 1: A_proj[B*K, STATE_D] = h_history[B*K, SH2+OFI] @ W_A[SH2+OFI, STATE_D]
// Col-major: C[STATE_D, B*K] = W_A_cm[STATE_D, SH2+OFI] @ h_cm[SH2+OFI, B*K]
unsafe {
cublaslt_sys::cublasLtMatmul(
lt_handle,
@@ -3257,7 +3259,7 @@ impl GpuDqnTrainer {
);
}
// GEMM 2: B_proj[B*K, STATE_D] = h_history[B*K, SH2] @ W_B[SH2, STATE_D]
// GEMM 2: B_proj[B*K, STATE_D] = h_history[B*K, SH2+OFI] @ W_B[SH2+OFI, STATE_D]
unsafe {
cublaslt_sys::cublasLtMatmul(
lt_handle,
@@ -3311,20 +3313,25 @@ impl GpuDqnTrainer {
Ok(())
}
/// Update rolling history: shift left, insert current h_s2 at K-1.
/// Update rolling history: shift left, insert [h_s2; ofi_embed] at K-1.
pub(crate) fn mamba2_update_history(&self, batch_size: usize) -> Result<(), MLError> {
let sh2 = self.config.shared_h2;
let total = batch_size * sh2;
let embed_dim = OFI_EMBED_DIM as i32;
let total_width = sh2 + OFI_EMBED_DIM;
let total = batch_size * total_width;
let mamba2_h_history_ptr = self.mamba2_h_history.raw_ptr();
let save_h_s2_ptr = self.save_h_s2.raw_ptr();
let ofi_embed_ptr = self.ofi_embed_output_buf.raw_ptr();
unsafe {
self.stream.launch_builder(&self.mamba2_update_kernel)
.arg(&mamba2_h_history_ptr)
.arg(&save_h_s2_ptr)
.arg(&ofi_embed_ptr)
.arg(&(batch_size as i32))
.arg(&(MAMBA2_HISTORY_K as i32))
.arg(&(sh2 as i32))
.arg(&embed_dim)
.launch(LaunchConfig::for_num_elems(total as u32))
.map_err(|e| MLError::ModelError(format!("mamba2_update_history: {e}")))?;
}
@@ -3365,13 +3372,14 @@ impl GpuDqnTrainer {
/// New: B*STATE_D threads with K=8 + SH2=256 inner loop, then 3 cuBLAS GEMMs.
pub(crate) fn mamba2_backward(&mut self, batch_size: usize) -> Result<(), MLError> {
let sh2 = self.config.shared_h2;
let h_width = sh2 + OFI_EMBED_DIM;
let param_ptr = self.mamba2_params.raw_ptr();
let w_c_ptr = param_ptr + (2 * sh2 * MAMBA2_STATE_DIM * 4) as u64;
let w_c_ptr = param_ptr + (2 * h_width * MAMBA2_STATE_DIM * 4) as u64;
let grad_ptr = self.mamba2_grad.raw_ptr();
let d_w_a = grad_ptr;
let d_w_b = grad_ptr + (sh2 * MAMBA2_STATE_DIM * 4) as u64;
let d_w_c = grad_ptr + (2 * sh2 * MAMBA2_STATE_DIM * 4) as u64;
let d_w_b = grad_ptr + (h_width * MAMBA2_STATE_DIM * 4) as u64;
let d_w_c = grad_ptr + (2 * h_width * MAMBA2_STATE_DIM * 4) as u64;
let lt_handle = self.shared_cublas.lt_handle.0;
let lt_ws_ptr = self.shared_cublas.lt_workspace_ptr;
@@ -3501,7 +3509,8 @@ impl GpuDqnTrainer {
/// Standalone SGD step for Mamba2 parameters (separate from main params_buf Adam).
/// params += -lr * grad via dqn_saxpy_f32_kernel.
pub(crate) fn step_mamba2_adam(&mut self) -> Result<(), MLError> {
let n = 3 * self.config.shared_h2 * MAMBA2_STATE_DIM;
let h_width = self.config.shared_h2 + OFI_EMBED_DIM;
let n = 2 * h_width * MAMBA2_STATE_DIM + self.config.shared_h2 * MAMBA2_STATE_DIM;
self.mamba2_adam_step += 1;
let base_lr: f32 = 1e-4;
let lr = base_lr * self.new_component_lr_scale();
@@ -6142,9 +6151,15 @@ impl GpuDqnTrainer {
info!("GpuDqnTrainer: liquid_tau_rk4_step kernel loaded (RK4/Euler adaptive ODE, 4 branches)");
// ── Mamba2 temporal scan: rolling history + selective SSM ──────────
// W_A and W_B project from h_history width (SH2+OFI_EMBED_DIM) to STATE_D.
// W_C projects from scan output (STATE_D) to SH2, so it stays at SH2 width.
let sh2 = config.shared_h2;
let mamba2_param_count = 3 * sh2 * MAMBA2_STATE_DIM;
let mamba2_h_history = stream.alloc_zeros::<f32>(b * MAMBA2_HISTORY_K * sh2)
let h_width = sh2 + OFI_EMBED_DIM; // history row width: trunk + OFI embedding
let wa_size = h_width * MAMBA2_STATE_DIM;
let wb_size = h_width * MAMBA2_STATE_DIM;
let wc_size = sh2 * MAMBA2_STATE_DIM;
let mamba2_param_count = wa_size + wb_size + wc_size;
let mamba2_h_history = stream.alloc_zeros::<f32>(b * MAMBA2_HISTORY_K * h_width)
.map_err(|e| MLError::ModelError(format!("alloc mamba2_h_history: {e}")))?;
let mamba2_h_enriched = stream.alloc_zeros::<f32>(b * sh2)
.map_err(|e| MLError::ModelError(format!("alloc mamba2_h_enriched: {e}")))?;
@@ -6158,7 +6173,7 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("alloc mamba2_adam_v: {e}")))?;
// Xavier init for mamba2_params
{
let scale = (2.0_f32 / (sh2 + MAMBA2_STATE_DIM) as f32).sqrt();
let scale = (2.0_f32 / (h_width + MAMBA2_STATE_DIM) as f32).sqrt();
let init_data: Vec<f32> = (0..mamba2_param_count).map(|j| {
let hash = (j as u32).wrapping_mul(2654435761).wrapping_add(0xDEADBEEF);
let u = (hash as f32) / (u32::MAX as f32) * 2.0 - 1.0;
@@ -6199,28 +6214,29 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("alloc mamba2_d_tw: {e}")))?;
// ── Mamba2 cuBLAS GEMM descriptors ──
// W_A/W_B/W_C are [SH2, STATE_D] row-major = [STATE_D, SH2] col-major with ld=STATE_D
// h_history is [B*K, SH2] row-major = [SH2, B*K] col-major with ld=SH2
// Forward: C[STATE_D, B*K] = W_cm[STATE_D, SH2] @ h_cm[SH2, B*K]
// W_A/W_B are [SH2+OFI, STATE_D] row-major = [STATE_D, SH2+OFI] col-major with ld=STATE_D
// h_history is [B*K, SH2+OFI] row-major = [SH2+OFI, B*K] col-major with ld=SH2+OFI
// Forward: C[STATE_D, B*K] = W_cm[STATE_D, SH2+OFI] @ h_cm[SH2+OFI, B*K]
let lt_handle = shared_cublas.lt_handle.0;
let lt_ws_size = shared_cublas.lt_workspace_size;
let mamba2_gemm_proj = create_mamba2_gemm_desc(
lt_handle, 0, 0,
MAMBA2_STATE_DIM, bk, sh2,
MAMBA2_STATE_DIM, sh2, MAMBA2_STATE_DIM,
MAMBA2_STATE_DIM, bk, h_width,
MAMBA2_STATE_DIM, h_width, MAMBA2_STATE_DIM,
lt_ws_size, "mamba2_fwd_proj",
)?;
// Backward dW: C[STATE_D, SH2] = d_cm[STATE_D, B*K] @ h_cm^T[B*K, SH2]
// Backward dW_A/dW_B: C[STATE_D, SH2+OFI] = d_cm[STATE_D, B*K] @ h_cm^T[B*K, SH2+OFI]
let mamba2_gemm_dw = create_mamba2_gemm_desc(
lt_handle, 0, 1,
MAMBA2_STATE_DIM, sh2, bk,
MAMBA2_STATE_DIM, sh2, MAMBA2_STATE_DIM,
MAMBA2_STATE_DIM, h_width, bk,
MAMBA2_STATE_DIM, h_width, MAMBA2_STATE_DIM,
lt_ws_size, "mamba2_bwd_dw",
)?;
// Backward dW_C: C[STATE_D, SH2] = xK_cm[STATE_D, B] @ dtw_cm^T[B, SH2]
// W_C stays at SH2 width — it projects from scan output, not h_history
let mamba2_gemm_dwc = create_mamba2_gemm_desc(
lt_handle, 0, 1,
MAMBA2_STATE_DIM, sh2, b,
@@ -6229,11 +6245,11 @@ impl GpuDqnTrainer {
)?;
info!(
bk, state_dim = MAMBA2_STATE_DIM, sh2,
bk, state_dim = MAMBA2_STATE_DIM, sh2, h_width,
a_proj_bytes = bk * MAMBA2_STATE_DIM * 4,
"GpuDqnTrainer: Mamba2 cuBLAS projections initialized (6 scratch buffers, 3 GEMM descs)"
);
info!("GpuDqnTrainer: Mamba2 temporal scan loaded ({} params, K={}, state_dim={})", mamba2_param_count, MAMBA2_HISTORY_K, MAMBA2_STATE_DIM);
info!("GpuDqnTrainer: Mamba2 temporal scan loaded ({} params, K={}, state_dim={}, h_width={})", mamba2_param_count, MAMBA2_HISTORY_K, MAMBA2_STATE_DIM, h_width);
// v8: Pessimistic Q-value initialization — shift value head bias to -0.1
// Pessimistic Q-init REMOVED — incompatible with per-sample support.

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

@@ -1,28 +1,39 @@
/**
* Update rolling history buffer: shift left by 1, insert current h_s2 at position K-1.
* Update rolling history buffer: shift left by 1, insert [h_s2(SH2); ofi_embed(embed_dim)]
* as one concatenated row at position K-1.
*
* Grid: ceil(N * SH2 / 256), Block: 256.
* Grid: ceil(N * (SH2 + embed_dim) / 256), Block: 256.
*/
extern "C" __global__ void mamba2_update_history(
float* __restrict__ h_history,
const float* __restrict__ h_s2,
const float* __restrict__ ofi_embed, /* [N, embed_dim] */
int N,
int K,
int sh2
int sh2,
int embed_dim /* 10 */
) {
int total_width = sh2 + embed_dim;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = N * sh2;
int total = N * total_width;
if (idx >= total) return;
int i = idx / sh2;
int j = idx % sh2;
int i = idx / total_width; /* sample index */
int j = idx % total_width; /* feature index within [SH2+embed_dim] */
float* base = h_history + (long long)i * K * sh2;
float* base = h_history + (long long)i * K * total_width;
/* Shift: copy slot t+1 -> t for t=0..K-2 */
for (int t = 0; t < K - 1; t++) {
base[(long long)t * sh2 + j] = base[(long long)(t + 1) * sh2 + j];
base[(long long)t * total_width + j] = base[(long long)(t + 1) * total_width + j];
}
base[(long long)(K - 1) * sh2 + j] = h_s2[(long long)i * sh2 + j];
/* Insert at K-1: first SH2 cols from h_s2, last embed_dim cols from ofi_embed */
if (j < sh2) {
base[(long long)(K - 1) * total_width + j] = h_s2[(long long)i * sh2 + j];
} else {
base[(long long)(K - 1) * total_width + j] = ofi_embed[i * embed_dim + (j - sh2)];
}
}
/* ================================================================== */