foxhunt/docs/archive/agents/AGENT_228_IMPLEMENTATION_GAPS.md

Agent 228: MAMBA-2 Implementation Gaps

Date: 2025-10-15 Priority: P0 - CRITICAL (Training Pipeline Blocked) Estimated Fix Time: 3-4 weeks (Waves 229-232)

---

Gap Summary

Gap	Priority	Impact	Effort	Status
1. SSD Algorithm	P0	5x speedup	2 weeks	❌ Not Implemented
2. Convolution Layer	P0	Model accuracy	2 days	❌ Missing
3. dt (Time Step)	P0	Selectivity	3 days	❌ Missing
4. D (Skip Connection)	P0	Training stability	1 day	❌ Missing
5. A Matrix Init	P1	Convergence	1 day	❌ Wrong
6. d_state Size	P1	Model capacity	1 day	❌ 8x too small
7. Tensor Core Ops	P1	GPU utilization	1 week	❌ No optimization
8. Memory Access	P1	Bandwidth	3 days	❌ Unoptimized
9. Input Splitting	P2	Architecture	2 days	❌ Missing
10. Mamba2Cache	P2	Inference speed	3 days	❌ Missing
11. RMSNorm	P3	Training speed	1 day	❌ Using LayerNorm
12. State Importance	P3	Memory	-	✅ Implemented

Total Gaps: 12 (11 missing, 1 complete) Critical Gaps: 4 (SSD, Conv, dt, D) Training Blocked: YES (cannot train without P0 gaps fixed)

---

Gap 1: SSD Algorithm (P0 - CRITICAL)

Current Implementation (WRONG)

// ml/src/mamba/mod.rs:1108
fn selective_scan_with_gradients(&self, input: &Tensor, A: &Tensor) -> Result<Tensor, MLError> {
    // Sequential scan (Mamba-1 style)
    for t in 0..seq_len {
        let x_t = input.narrow(1, t, 1)?.squeeze(1)?;
        current_state = (current_state.matmul(&A.t()?)? + &x_t)?;  // O(n) sequential
        states.push(current_state.unsqueeze(1)?);
    }
    Tensor::cat(&states, 1)
}

Problem: This is Mamba-1 parallel associative scan, not Mamba-2 SSD.

Reference Implementation (CORRECT)

# tommyip/mamba2-minimal/mamba2.py
def ssd_chunk_scan(x, dt, A, B, C, chunk_size=256):
    """Structured State Duality chunk scan"""
    batch, seqlen, dim = x.shape
    num_chunks = seqlen // chunk_size

    # Discretize parameters
    dt = F.softplus(dt)
    A_discrete = torch.exp(A * dt)  # Diagonal A
    B_discrete = B * dt

    # Chunk-based processing (tensor core friendly)
    states = []
    for chunk_idx in range(num_chunks):
        # Extract chunk
        chunk_x = x[:, chunk_idx*chunk_size:(chunk_idx+1)*chunk_size]

        # Diagonal block: local state computation (parallel)
        diag_block = compute_diagonal(chunk_x, A_discrete, B_discrete)

        # Off-diagonal block: inter-chunk dependencies
        if chunk_idx > 0:
            offdiag_block = compute_offdiagonal(prev_state, A_discrete, chunk_size)
            diag_block = diag_block + offdiag_block

        states.append(diag_block)
        prev_state = diag_block[:, -1]  # Last state as carry

    # Concatenate chunks
    states = torch.cat(states, dim=1)

    # Output transformation: y = states @ C^T
    y = torch.einsum('bld,dc->blc', states, C)
    return y

Required Changes

1. Replace Sequential Scan:

   // NEW: ml/src/mamba/ssd_algorithm.rs
   pub fn ssd_chunk_scan(
       input: &Tensor,
       dt: &Tensor,
       A: &Tensor,
       B: &Tensor,
       C: &Tensor,
       chunk_size: usize,
   ) -> Result<Tensor, MLError> {
       // Implement chunk-based SSD algorithm
   }
   

2. Add Diagonal Block Computation:

   fn compute_diagonal_block(
       chunk_x: &Tensor,
       A_discrete: &Tensor,
       B_discrete: &Tensor,
   ) -> Result<Tensor, MLError> {
       // Matrix multiplication: chunk_x @ B_discrete^T
       // Then scale by A_discrete
   }
   

3. Add Off-Diagonal Block Computation:

   fn compute_offdiagonal_block(
       prev_state: &Tensor,
       A_discrete: &Tensor,
       chunk_size: usize,
   ) -> Result<Tensor, MLError> {
       // Propagate previous chunk's final state
   }
   

Estimated Effort: 2 weeks (complex algorithm) Blocking: Training pipeline Dependencies: dt parameter, diagonal A matrix

---

Gap 2: Convolution Layer (P0)

Current Implementation (MISSING)

// ml/src/mamba/mod.rs:570-608
pub fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
    let mut hidden = self.input_projection.forward(input)?;
    // ❌ No convolution here!
    for layer_idx in 0..num_layers {
        let normalized = self.layer_norms[layer_idx].forward(&hidden)?;
        let layer_output = self.forward_ssd_layer(&ssd_layer, &normalized, layer_idx)?;
        ...
    }
}

Reference Implementation (CORRECT)

# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_model, d_conv=4):
        self.conv1d = nn.Conv1d(
            in_channels=d_inner,
            out_channels=d_inner,
            kernel_size=d_conv,
            groups=d_inner,  # Depthwise convolution
            padding=d_conv - 1  # Causal padding
        )

    def forward(self, x):
        x, z = self.in_proj(u).chunk(2, dim=-1)

        # Causal convolution (MISSING in our implementation)
        x = rearrange(x, 'b l d -> b d l')
        x = self.conv1d(x)[:, :, :seqlen]  # Remove extra padding
        x = rearrange(x, 'b d l -> b l d')

        x = F.silu(x)
        y = self.ssm(x)
        ...

Required Changes

1. Add Conv1d to Mamba2SSM:

   // ml/src/mamba/mod.rs
   pub struct Mamba2SSM {
       pub input_projection: Linear,
       pub conv1d: Vec<Conv1d>,  // NEW: One per layer
       pub layer_norms: Vec<CudaLayerNorm>,
       ...
   }
   

2. Initialize Convolution:

   impl Mamba2SSM {
       pub fn new(config: Mamba2Config, device: &Device) -> Result<Self, MLError> {
           let mut conv1d = Vec::new();
           for i in 0..config.num_layers {
               let conv = candle_nn::conv1d(
                   d_inner,           // in_channels
                   d_inner,           // out_channels
                   4,                 // kernel_size (d_conv)
                   candle_nn::Conv1dConfig {
                       groups: d_inner,     // Depthwise
                       padding: 3,          // Causal padding (d_conv - 1)
                       ..Default::default()
                   },
                   vb.pp(&format!("conv1d_{}", i)),
               )?;
               conv1d.push(conv);
           }
           ...
       }
   }
   

3. Apply Convolution in Forward Pass:

   pub fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
       let mut hidden = self.input_projection.forward(input)?;
   
       for layer_idx in 0..num_layers {
           // Apply convolution BEFORE layer norm
           hidden = hidden.transpose(1, 2)?;  // [batch, d_inner, seq_len]
           hidden = self.conv1d[layer_idx].forward(&hidden)?;
           hidden = hidden.narrow(2, 0, seq_len)?;  // Remove extra padding
           hidden = hidden.transpose(1, 2)?;  // [batch, seq_len, d_inner]
   
           hidden = hidden.silu()?;  // SiLU activation
   
           let normalized = self.layer_norms[layer_idx].forward(&hidden)?;
           ...
       }
   }
   

Estimated Effort: 2 days Blocking: Training accuracy Dependencies: None (can implement immediately)

---

Gap 3: dt (Time Step) Parameter (P0)

Current Implementation (WRONG)

// ml/src/mamba/mod.rs:261-266
let delta = Tensor::ones((config.d_model,), DType::F64, device)?;  // CONSTANT!

Problem: Time step is constant, not input-dependent. This breaks selectivity.

Reference Implementation (CORRECT)

# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_model, dt_rank=None):
        dt_rank = dt_rank or ceil(d_model / 16)

        # Project input to dt space
        self.x_proj = nn.Linear(d_inner, dt_rank + 2*d_state)

        # Project dt_rank to d_inner (learnable)
        self.dt_proj = nn.Linear(dt_rank, d_inner, bias=True)

        # Initialize dt bias
        dt_init_std = dt_rank**-0.5
        dt = torch.exp(
            torch.rand(d_inner) * (math.log(dt_max) - math.log(dt_min))
            + math.log(dt_min)
        )
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        self.dt_bias = nn.Parameter(inv_dt)

    def forward(self, x):
        # Extract dt from input
        x_proj = self.x_proj(x)
        dt, B, C = torch.split(x_proj, [dt_rank, d_state, d_state], dim=-1)

        # Project and activate
        dt = self.dt_proj(dt)
        dt = F.softplus(dt + self.dt_bias)
        dt = dt.clamp(min=dt_min, max=dt_max)

        # Use dt in discretization
        A_discrete = torch.exp(self.A_log * dt)
        B_discrete = B * dt
        ...

Required Changes

1. Add dt Parameters to Config:

   // ml/src/mamba/mod.rs:68-107
   pub struct Mamba2Config {
       pub dt_rank: usize,        // NEW: ceil(d_model / 16)
       pub dt_min: f64,           // NEW: 0.001
       pub dt_max: f64,           // NEW: 0.1
       pub dt_init_floor: f64,    // NEW: 1e-4
       ...
   }
   

2. Add dt Projection Layers:

   pub struct Mamba2SSM {
       pub x_proj: Vec<Linear>,      // NEW: input → (dt_rank + 2*d_state)
       pub dt_proj: Vec<Linear>,     // NEW: dt_rank → d_inner
       pub dt_bias: Vec<Tensor>,     // NEW: learnable bias
       ...
   }
   

3. Initialize dt Parameters:

   impl Mamba2SSM {
       pub fn new(config: Mamba2Config, device: &Device) -> Result<Self, MLError> {
           let dt_rank = (config.d_model as f64 / 16.0).ceil() as usize;
           let mut x_proj = Vec::new();
           let mut dt_proj = Vec::new();
           let mut dt_bias = Vec::new();
   
           for i in 0..config.num_layers {
               // x_proj: d_inner → (dt_rank + 2*d_state)
               let x_p = candle_nn::linear(
                   d_inner,
                   dt_rank + 2 * config.d_state,
                   vb.pp(&format!("x_proj_{}", i)),
               )?;
               x_proj.push(x_p);
   
               // dt_proj: dt_rank → d_inner
               let dt_p = candle_nn::linear(
                   dt_rank,
                   d_inner,
                   vb.pp(&format!("dt_proj_{}", i)),
               )?;
               dt_proj.push(dt_p);
   
               // dt_bias initialization (log-uniform)
               let dt_init_std = (dt_rank as f64).powf(-0.5);
               let dt_init = Tensor::rand(0.0f32, 1.0f32, (d_inner,), device)?
                   .mul(&Tensor::new(&[(config.dt_max.ln() - config.dt_min.ln()) as f32], device)?)?
                   .add(&Tensor::new(&[config.dt_min.ln() as f32], device)?)?
                   .exp()?;
   
               // Inverse softplus transformation
               let inv_dt = dt_init.clone().add(&dt_init.neg()?.expm1()?.neg()?.log()?)?;
               dt_bias.push(inv_dt);
           }
   
           Ok(Self { x_proj, dt_proj, dt_bias, ... })
       }
   }
   

4. Use dt in Forward Pass:

   fn forward_ssd_layer_with_gradients(
       &mut self,
       input: &Tensor,
       layer_idx: usize,
   ) -> Result<Tensor, MLError> {
       // Project input to get dt, B, C
       let x_proj = self.x_proj[layer_idx].forward(input)?;
       let dt_rank = (self.config.d_model as f64 / 16.0).ceil() as usize;
   
       let dt = x_proj.narrow(2, 0, dt_rank)?;
       let B = x_proj.narrow(2, dt_rank, self.config.d_state)?;
       let C = x_proj.narrow(2, dt_rank + self.config.d_state, self.config.d_state)?;
   
       // Project dt and apply softplus
       let dt = self.dt_proj[layer_idx].forward(&dt)?;
       let dt = (dt + &self.dt_bias[layer_idx])?;
       let dt = dt.softplus()?;  // softplus(x) = log(1 + exp(x))
       let dt = dt.clamp(self.config.dt_min, self.config.dt_max)?;
   
       // Discretize A and B using dt
       let A = &self.state.ssm_states[layer_idx].A;
       let A_discrete = (A * &dt)?.exp()?;  // exp(A_log * dt)
       let B_discrete = (B * &dt)?;
   
       // Use in SSD algorithm
       let y = ssd_chunk_scan(input, &dt, &A_discrete, &B_discrete, &C)?;
       Ok(y)
   }
   

Estimated Effort: 3 days Blocking: Model selectivity Dependencies: None (can implement immediately)

---

Gap 4: D (Skip Connection) Parameter (P0)

Current Implementation (MISSING)

// ml/src/mamba/mod.rs:1090
let output = scanned_states.matmul(&C_broadcasted)?;
// ❌ No D * x skip connection!

Reference Implementation (CORRECT)

# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_model):
        self.D = nn.Parameter(torch.ones(d_inner))  # Learnable skip weight

    def forward(self, x):
        # SSM computation
        y_ssm = self.ssm(x, A, B, C)

        # Skip connection with learnable weight
        y = self.D * x + y_ssm

        # Output projection
        output = self.out_proj(y)
        return output

Required Changes

1. Add D Parameter:

   // ml/src/mamba/mod.rs:193-211
   pub struct SSMState {
       pub A: Tensor,
       pub B: Tensor,
       pub C: Tensor,
       pub delta: Tensor,
       pub D: Tensor,  // NEW: Skip connection weight
       pub hidden: Tensor,
   }
   

2. Initialize D:

   impl Mamba2State {
       pub fn zeros(config: &Mamba2Config, device: &Device) -> Result<Self, MLError> {
           for layer_idx in 0..config.num_layers {
               // Initialize D to ones (identity skip connection)
               let D = Tensor::ones((d_inner,), DType::F64, device)?;
   
               ssm_states.push(SSMState { A, B, C, delta, D, hidden });
           }
           ...
       }
   }
   

3. Apply D in Forward Pass:

   fn forward_ssd_layer_with_gradients(
       &mut self,
       input: &Tensor,
       layer_idx: usize,
   ) -> Result<Tensor, MLError> {
       let D = &self.state.ssm_states[layer_idx].D;
   
       // SSM output
       let y_ssm = ssd_chunk_scan(input, &dt, &A_discrete, &B_discrete, &C)?;
   
       // Skip connection: D * x + y_ssm
       let D_broadcasted = D.unsqueeze(0)?.unsqueeze(0)?;  // [1, 1, d_inner]
       let skip = (input * &D_broadcasted)?;
       let output = (skip + y_ssm)?;
   
       Ok(output)
   }
   

4. Add D to Optimizer:

   fn optimizer_step(&mut self) -> Result<(), MLError> {
       for layer_idx in 0..num_layers {
           // Update D parameter (like B, C)
           if let Some(ref D_grad) = self.gradients.get(&format!("D_{}", layer_idx)) {
               let mut D_param = self.state.ssm_states[layer_idx].D.clone();
               self.apply_adam_update(&mut D_param, D_grad, layer_idx, "D", ...)?;
               self.state.ssm_states[layer_idx].D = D_param;
           }
       }
   }
   

Estimated Effort: 1 day Blocking: Training stability Dependencies: None (can implement immediately)

---

Gap 5: A Matrix Initialization (P1)

Current Implementation (WRONG)

// ml/src/mamba/mod.rs:237-242
let A = Tensor::randn(0.0, 1.0, (config.d_state, config.d_state), device)?;

Problem: Random initialization, full matrix (not diagonal).

Reference Implementation (CORRECT)

# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_state):
        # A_log initialization: log(range(1, d_state+1))
        A_log = torch.log(torch.arange(1, d_state + 1, dtype=torch.float32))
        self.A_log = nn.Parameter(A_log)  # Shape: (d_state,) - DIAGONAL!

    def forward(self, x):
        A = -torch.exp(self.A_log)  # Negative for stability
        ...

Required Changes

1. Make A Diagonal:

   // ml/src/mamba/mod.rs:237-242
   // OLD: let A = Tensor::randn(0.0, 1.0, (config.d_state, config.d_state), device)?;
   
   // NEW: Initialize as diagonal with log(range(1, d_state+1))
   let A_diag: Vec<f64> = (1..=config.d_state)
       .map(|i| (i as f64).ln())
       .collect();
   let A = Tensor::from_vec(A_diag, (config.d_state,), device)?;
   

2. Use Diagonal A in Discretization:

   fn discretize_ssm_with_gradients(&self, A_log: &Tensor, dt: &Tensor) -> Result<Tensor, MLError> {
       // A is now 1D (diagonal), not 2D matrix
       let A = A_log.neg()?.exp()?;  // Negative exponential for stability
   
       // Discretize: A_discrete = exp(-A * dt)
       let A_discrete = (A * dt)?.neg()?.exp()?;  // Elementwise
   
       Ok(A_discrete)
   }
   

3. Update SSMState Structure:

   pub struct SSMState {
       pub A: Tensor,  // Shape: (d_state,) - DIAGONAL ONLY
       pub B: Tensor,  // Shape: (d_state, d_inner)
       pub C: Tensor,  // Shape: (d_inner, d_state)
       pub D: Tensor,  // Shape: (d_inner,)
       pub delta: Tensor,
       pub hidden: Tensor,
   }
   

Estimated Effort: 1 day Blocking: Convergence speed Dependencies: None (can implement immediately)

---

Gap 6: d_state Size (P1)

Current Implementation (WRONG)

// ml/src/mamba/mod.rs:139
d_state: 16,  // 8x too small!

Reference Implementations (CORRECT)

• state-spaces/mamba: d_state = 64-128
• tommyip/mamba2-minimal: d_state = 64
• Hugging Face: d_state = 128

Required Changes

// ml/src/mamba/mod.rs:134-158
impl Mamba2Config {
    pub fn emergency_safe_defaults() -> Self {
        Self {
            d_state: 64,  // FIXED: Was 16, now 64 (minimum for Mamba-2)
            ...
        }
    }
}

Impact:

• Larger state → more model capacity
• 4x memory increase (16 → 64)
• Better long-range dependencies

Estimated Effort: 1 day (just change constant + retrain) Blocking: Model capacity Dependencies: None

---

Gap 7-12: See Detailed Implementation Plan

(Remaining gaps documented in AGENT_228_REFERENCE_IMPLEMENTATIONS.md sections 7-12)

---

Implementation Priority Queue

Wave 229 (This Week - 3 days)

1. dt Parameter (3 days, P0)

   • Add x_proj, dt_proj, dt_bias
   • Implement softplus + clamping
   • Use in discretization

2. D Parameter (1 day, P0)

   • Add to SSMState
   • Initialize to ones
   • Apply skip connection

3. Convolution Layer (2 days, P0)

   • Add Conv1d to model
   • Apply before SSM
   • Causal padding

Wave 230 (Next Week - 5 days)

4. A Matrix Fix (1 day, P1)

   • Change to diagonal
   • Log initialization
   • Update discretization

5. d_state Increase (1 day, P1)

   • Change from 16 to 64
   • Test memory usage

6. SSD Algorithm (3 days, P0)

   • Implement chunk-based scan
   • Diagonal/off-diagonal blocks
   • Matrix multiplication approach

Wave 231 (Week 3 - 5 days)

7. Input Splitting (2 days, P2)

   • Split in_proj → (z, x)
   • Add SiLU gating
   • Update output

8. Tensor Core Optimization (3 days, P1)

   • Use FP16/BF16
   • Align dimensions
   • Profile performance

Wave 232 (Week 4 - 5 days)

9. Memory Access Patterns (3 days, P1)

   • Coalesce memory operations
   • Reduce HBM transfers
   • Profile bandwidth

10. Integration Testing (2 days)

    • Test with ES.FUT data
    • Validate gradients
    • Benchmark vs PyTorch

---

Success Criteria

Functional Requirements

• ✅ Model trains without errors
• ✅ Gradients flow correctly
• ✅ Validation loss decreases
• ✅ Inference latency <5μs

Performance Requirements

• ✅ Training speed ≥50% of PyTorch Mamba-2
• ✅ Memory usage ≤2x PyTorch
• ✅ GPU utilization >70%

Architectural Requirements

• ✅ SSD algorithm implemented
• ✅ All parameters present (dt, D, A, B, C)
• ✅ Convolution layer working
• ✅ Tensor core optimization enabled

---

Testing Strategy

Unit Tests

#[test]
fn test_dt_parameter() {
    // Test dt projection and clamping
    let config = Mamba2Config::default();
    let model = Mamba2SSM::new(config, &Device::Cpu)?;

    let input = Tensor::randn(0.0, 1.0, (1, 10, 64), &Device::Cpu)?;
    let dt = model.compute_dt(&input, 0)?;

    assert!(dt.min()? >= config.dt_min);
    assert!(dt.max()? <= config.dt_max);
}

#[test]
fn test_d_skip_connection() {
    // Test D parameter skip connection
    let model = Mamba2SSM::new(config, &Device::Cpu)?;
    let input = Tensor::ones((1, 10, 64), &Device::Cpu)?;

    let output = model.forward(&input)?;

    // Output should include skip connection
    assert!(output.dims() == input.dims());
}

#[test]
fn test_conv1d_causal() {
    // Test causal convolution (no future leakage)
    let model = Mamba2SSM::new(config, &Device::Cpu)?;
    let input = Tensor::zeros((1, 10, 64), &Device::Cpu)?;
    input.narrow(1, 5, 1)?.fill_(1.0)?;  // Set t=5 to 1

    let output = model.forward(&input)?;

    // Positions t<5 should be zero (no future info)
    assert!(output.narrow(1, 0, 5)?.abs().sum()? < 1e-6);
}

Integration Tests

#[test]
fn test_e2e_mamba2_training() {
    // Test end-to-end training loop
    let mut model = Mamba2SSM::new(config, &Device::cuda_if_available(0)?)?;
    let train_data = load_es_fut_data()?;

    let history = model.train(&train_data, &val_data, epochs=10).await?;

    // Loss should decrease
    assert!(history.last().unwrap().loss < history[0].loss);
}

Performance Benchmarks

# Run GPU training benchmark
cargo run --release -p ml --example gpu_training_benchmark

# Compare against PyTorch
python benchmarks/compare_mamba2.py --model foxhunt --baseline pytorch

---

Risk Mitigation

Risk 1: SSD Algorithm Complexity

• Probability: HIGH
• Impact: CRITICAL
• Mitigation: Start with tommyip/mamba2-minimal (simplest implementation)
• Fallback: Use Mamba-1 associative scan temporarily

Risk 2: Candle Limitations

• Probability: MEDIUM
• Impact: HIGH
• Mitigation: Implement SSD using primitive ops (matmul, elementwise)
• Fallback: Request custom CUDA kernel support from Candle team

Risk 3: Memory Increase

• Probability: LOW
• Impact: MEDIUM
• Mitigation: Profile memory usage, optimize batch size
• Fallback: Reduce d_state if GPU OOM

---

Agent 228 Out 🎯