diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 99ccff694..52c5a1fa1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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, // [B, K, SH2] rolling history + mamba2_h_history: CudaSlice, // [B, K, SH2+OFI_EMBED_DIM] rolling history (widened for OFI) mamba2_h_enriched: CudaSlice, // [B, SH2] output mamba2_update_kernel: CudaFunction, mamba2_copy_enriched_kernel: CudaFunction, - mamba2_params: CudaSlice, // [3 * SH2 * STATE_D] W_A, W_B, W_C concatenated - mamba2_grad: CudaSlice, // gradient buffer - mamba2_adam_m: CudaSlice, // Adam first moment - mamba2_adam_v: CudaSlice, // Adam second moment + mamba2_params: CudaSlice, // W_A[(SH2+OFI)*SD] + W_B[(SH2+OFI)*SD] + W_C[SH2*SD] + mamba2_grad: CudaSlice, // gradient buffer (same layout as params) + mamba2_adam_m: CudaSlice, // Adam first moment (same size as params) + mamba2_adam_v: CudaSlice, // 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, // [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::(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::(b * MAMBA2_HISTORY_K * h_width) .map_err(|e| MLError::ModelError(format!("alloc mamba2_h_history: {e}")))?; let mamba2_h_enriched = stream.alloc_zeros::(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 = (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. diff --git a/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu b/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu index c60f9ee0f..fc4c011b9 100644 --- a/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu +++ b/crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu @@ -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)]; + } } /* ================================================================== */