foxhunt/AGENT_228_REFERENCE_IMPLEMENTATIONS.md

Agent 228: MAMBA-2 Reference Implementations Comparison

Date: 2025-10-15 Agent: Agent 228 Mission: Compare authoritative MAMBA-2 implementations against Foxhunt implementation Status: ✅ COMPLETE (3 reference implementations analyzed)

---

Executive Summary

Analyzed 3 authoritative MAMBA-2 implementations and identified 12 critical gaps in our Rust implementation. The reference implementations (state-spaces/mamba, tommyip/mamba2-minimal, Hugging Face Transformers) all implement the Structured State Duality (SSD) algorithm with hardware-optimized matrix operations, while our implementation uses a simplified SSM approach without true SSD.

Critical Finding: Our implementation is MAMBA-1 style, not MAMBA-2. We're missing the core SSD algorithm that provides 5x speedup.

---

1. Reference Implementations Found

1.1 Official state-spaces/mamba (PRIMARY REFERENCE)

• Repository: https://github.com/state-spaces/mamba
• Language: Python/CUDA
• Trust Score: 10/10 (official implementation by authors Tri Dao & Albert Gu)
• Key Features:
  • Mamba-2 SSD layer with tensor core optimization
  • Chunk-based parallel processing
  • Hardware-aware memory access patterns
  • Custom CUDA kernels for A100/H100 GPUs

Code Snippet (from Context7):

from mamba_ssm import Mamba2

model = Mamba2(
    d_model=dim,      # Model dimension
    d_state=64,       # SSM state expansion factor (64 or 128 for Mamba-2)
    d_conv=4,         # Local convolution width
    expand=2,         # Block expansion factor
).to("cuda")
y = model(x)

1.2 tommyip/mamba2-minimal (EDUCATIONAL REFERENCE)

• Repository: https://github.com/tommyip/mamba2-minimal
• Language: Pure PyTorch (single file)
• Trust Score: 9/10 (minimal, readable implementation)
• Key Features:
  • Structured State Duality (SSD) algorithm
  • Chunk-based matrix operations
  • Diagonal and off-diagonal block computation
  • Inference-optimized forward pass

Architecture Overview:

Input → In-Projection → Convolution → SSD Layer → Out-Projection → Output
                                     ↓
                            Chunk-wise Processing
                            ├─ Diagonal Blocks
                            └─ Off-Diagonal Blocks

1.3 Hugging Face Transformers Mamba2

• Repository: https://github.com/huggingface/transformers/tree/main/src/transformers/models/mamba2
• Language: Python/PyTorch
• Trust Score: 10/10 (production-grade implementation)
• Key Features:
  • Mamba2Cache for efficient state management
  • Sequence parallel and tensor parallel support
  • Dynamic time step scaling
  • Selective state normalization

Official Mamba2 Block:

class Mamba2Block:
    - in_proj: Linear(d_model, d_inner * 2)
    - conv1d: Conv1d(d_inner, d_inner, kernel_size=d_conv)
    - x_proj: Linear(d_inner, dt_rank + d_state * 2)
    - dt_proj: Linear(dt_rank, d_inner)
    - A_log: Parameter(d_inner, d_state)
    - D: Parameter(d_inner)
    - out_proj: Linear(d_inner, d_model)

---

2. Structured State Duality (SSD) Algorithm

2.1 What is SSD?

Key Insight from Perplexity AI:

"Structured State Duality (SSD) refers to a theoretical and practical equivalence between a special class of structured state-space models (SSMs) and masked attention mechanisms, enabling efficient sequence modeling with both recurrent (linear-time) and attention-like (quadratic-time) algorithms."

2.2 SSD Mathematical Formulation

For a sequence input X ∈ ℝ^(T×d), the SSD block computes:

Y = diag(p) · M · diag(q) · X

Where:

• M is a 1-semiseparable mask matrix (causal mask)
• p, q are vectors derived from input and parameter projections
• This is equivalent to masked attention with specific structure

State Matrix Constraint: A must be scalar-times-identity or diagonal (this is the "duality")

2.3 SSD vs Traditional SSM

Feature	Traditional SSM (Mamba-1)	SSD (Mamba-2)
State Matrix A	General structured matrix	Diagonal or scalar×I
Computation	Parallel associative scan	Matrix multiplication (tensor cores)
Hardware	Limited GPU optimization	Tensor core optimized
Complexity	O(n) work, O(log n) depth	O(n) work, hardware-efficient
Speed	Baseline	5x faster
State Size	Limited by memory	8x larger for same memory

---

3. Key Architectural Differences

3.1 Forward Pass Comparison

Reference Implementation (tommyip/mamba2-minimal)

def forward(self, x):
    # 1. Input projection + split
    z, x = self.in_proj(x).chunk(2, dim=-1)

    # 2. Convolution (causal padding)
    x = self.conv1d(x.transpose(1, 2)).transpose(1, 2)
    x = F.silu(x)

    # 3. SSM parameters projection
    x_proj = self.x_proj(x)
    dt, B, C = torch.split(x_proj, [self.dt_rank, self.d_state, self.d_state], dim=-1)
    dt = self.dt_proj(dt)

    # 4. SSD algorithm (chunk-based)
    y = self.selective_scan(x, dt, A, B, C)

    # 5. Output projection
    y = y * F.silu(z)
    output = self.out_proj(y)
    return output

Our Implementation (ml/src/mamba/mod.rs)

pub fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
    // 1. Input projection (no split)
    let mut hidden = self.input_projection.forward(input)?;

    // 2. Process through layers
    for layer_idx in 0..num_layers {
        let normalized = self.layer_norms[layer_idx].forward(&hidden)?;

        // 3. SIMPLIFIED SSM (not true SSD!)
        let layer_output = self.forward_ssd_layer(&ssd_layer, &normalized, layer_idx)?;

        hidden = (&hidden + &layer_output)?;  // Residual
        hidden = self.dropouts[layer_idx].forward(&hidden, true)?;
    }

    // 4. Output projection
    let output = self.output_projection.forward(&hidden)?;
    Ok(output)
}

Gap: We're missing the convolution, parameter splitting, and true SSD chunk-based algorithm.

3.2 Selective Scan Algorithm

Reference (Mamba-2 SSD Scan)

def selective_scan(x, dt, A, B, C, chunk_size=256):
    """
    Chunk-based SSD scan with tensor core optimization
    """
    batch, seqlen, dim = x.shape

    # Discretize continuous parameters
    dt = F.softplus(dt + dt_bias)
    A_discrete = torch.exp(A * dt)  # Elementwise for diagonal A
    B_discrete = B * dt

    # Process in chunks for hardware efficiency
    chunks = seqlen // chunk_size
    states = []

    for chunk_idx in range(chunks):
        # 1. Compute diagonal blocks (local chunk)
        chunk_x = x[:, chunk_idx*chunk_size:(chunk_idx+1)*chunk_size]
        chunk_state = compute_diagonal_block(chunk_x, A_discrete, B_discrete)

        # 2. Compute off-diagonal blocks (inter-chunk)
        if chunk_idx > 0:
            chunk_state = combine_with_prev_state(prev_state, chunk_state, A_discrete)

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

    # 3. Output transformation
    states = torch.cat(states, dim=1)
    y = torch.einsum('bld,bdc->blc', states, C)
    return y

Our Implementation (simplified SSM)

fn selective_scan_with_gradients(&self, input: &Tensor, A: &Tensor) -> Result<Tensor, MLError> {
    let seq_len = input.dim(1)?;
    let mut states = Vec::new();
    let mut current_state = Tensor::zeros(...)?;

    // Sequential scan (NO chunking, NO tensor cores)
    for t in 0..seq_len {
        let x_t = input.narrow(1, t, 1)?.squeeze(1)?;

        // Simple recurrence: h_t = h_{t-1} @ A^T + x_t
        current_state = (current_state.matmul(&A.t()?)? + &x_t)?;
        states.push(current_state.unsqueeze(1)?);
    }

    Tensor::cat(&states, 1)
}

Gaps:

1. ❌ No chunk-based processing
2. ❌ No diagonal A matrix constraint
3. ❌ No tensor core optimization
4. ❌ Sequential scan instead of parallel prefix scan
5. ❌ Missing dt (time step) parameter
6. ❌ No discretization of continuous parameters

---

4. Missing Hardware Optimizations

4.1 Tensor Core Utilization (CRITICAL GAP)

Reference Insight from Perplexity AI:

"MAMBA-2 achieves hardware-aware optimization by leveraging tensor cores for matrix multiplication, optimizing memory access patterns, and adopting parallelization strategies that map efficiently to modern GPU architectures. By structuring the algorithm to use block-wise matrix multiplications, MAMBA-2 achieves substantial speedups—up to 16x on A100 and H100 GPUs."

What We're Missing:

• Block-wise matrix multiplication layout
• Tensor core-friendly dimensions (multiples of 8/16)
• FP16/BF16 mixed precision
• WMMA (Warp Matrix Multiply-Accumulate) operations

4.2 Memory Access Patterns

Reference (Optimized)

# Coalesced memory access
x_chunks = x.view(batch, num_chunks, chunk_size, dim)  # Contiguous
states_chunks = compute_chunked_states(x_chunks)       # Block-wise

# Minimize HBM transfers
intermediate = recompute_on_the_fly()  # FlashAttention-style

Our Implementation (Suboptimal)

// Random access pattern
for t in 0..seq_len {
    let x_t = input.narrow(1, t, 1)?;  // Non-contiguous memory access
    current_state = current_state.matmul(&A.t()?)?;  // No tensor core usage
}

4.3 Parallel Scan Implementation

Our scan_algorithms.rs implements block-wise parallel prefix scan, which is closer to Mamba-1's approach. Mamba-2 SSD uses matrix multiplication instead:

// Our approach (Mamba-1 style)
pub fn block_parallel_scan(&self, input: &Tensor, op: ScanOperator) -> Result<Tensor, MLError> {
    // Phase 1: Process blocks independently
    for block_idx in 0..num_blocks {
        let block_result = self.sequential_scan(&block_input, op)?;
        block_carries.push(carry);
    }

    // Phase 2: Prefix scan of carries
    let carry_scan = self.sequential_scan(&carries_tensor, op)?;

    // Phase 3: Combine with carry propagation
    ...
}

Gap: This is correct for Mamba-1 but not the SSD algorithm for Mamba-2.

---

5. Parameter Differences

5.1 Default Configuration Comparison

Parameter	state-spaces/mamba	tommyip/mamba2-minimal	Foxhunt Implementation
d_state	64-128	64	16 (8x smaller!)
d_conv	4	4	❌ Missing
expand	2	2	1-2
dt_rank	Automatic	ceil(d_model / 16)	❌ Missing
dt_min	0.001	0.001	❌ Missing
dt_max	0.1	0.1	❌ Missing
A_init	log(range(1, d_state+1))	log(range(1, d_state+1))	Random normal
D parameter	✅ Present	✅ Present	❌ Missing

5.2 Missing Parameters

dt (Time Step) Parameter

# Reference
dt = F.softplus(dt_proj(x) + dt_bias)  # Learned, input-dependent
dt = dt.clamp(dt_min, dt_max)

# Our implementation
let delta = Tensor::ones((config.d_model,), DType::F64, device)?;  # Constant!

D (Skip Connection) Parameter

# Reference
y = D * x + ssm_output  # Learnable residual weight

# Our implementation
# ❌ Missing entirely

Convolution Layer

# Reference
x = self.conv1d(x.transpose(1, 2))  # Causal convolution

# Our implementation
# ❌ Missing entirely

---

6. SSD Layer Implementation Gap

6.1 Reference SSD Layer (tommyip)

Key Components:

1. In-projection: Splits into z (gate) and x (input to SSM)
2. Causal Convolution: Local context aggregation
3. SSM Parameter Projection: Generates dt, B, C from input
4. SSD Algorithm: Chunk-based state computation
5. Output Gating: y = y * silu(z)

6.2 Our SSD Layer (ml/src/mamba/ssd_layer.rs)

What We Have:

• ✅ Multi-head attention-like projections (Q, K, V)
• ✅ Linear attention mechanism (O(n) complexity)
• ✅ State space transformation
• ✅ Gating mechanism

What We're Missing:

• ❌ Input splitting (z, x)
• ❌ Causal convolution
• ❌ Dynamic SSM parameter projection
• ❌ Chunk-based SSD algorithm
• ❌ Proper silu(z) gating

Conclusion: Our SSD layer is actually a linear attention layer, not a true SSD layer.

---

7. Implementation Gaps Summary

7.1 Critical Gaps (High Priority)

1. SSD Algorithm (P0 - CRITICAL)

   • Replace sequential scan with chunk-based SSD algorithm
   • Implement diagonal A matrix constraint
   • Add tensor core-friendly matrix operations

2. Convolution Layer (P0)

   • Add 1D causal convolution before SSM
   • Kernel size: d_conv = 4

3. dt (Time Step) Parameter (P0)

   • Add learnable dt_proj layer
   • Implement softplus activation + clamping
   • Make dt input-dependent

4. D (Skip Connection) (P0)

   • Add learnable D parameter
   • Implement y = D * x + ssm_output

5. Parameter Initialization (P1)

   • Fix A initialization: A_log = log(arange(1, d_state+1))
   • Increase d_state from 16 to 64-128
   • Add dt_bias initialization

7.2 Hardware Optimization Gaps (Medium Priority)

6. Tensor Core Optimization (P1)

   • Restructure matrix ops for tensor cores
   • Use FP16/BF16 mixed precision
   • Align dimensions to 8/16

7. Memory Access Patterns (P1)

   • Implement chunk-based processing
   • Coalesce memory accesses
   • Reduce HBM transfers

8. Parallel Scan (P2)

   • Replace scan with SSD matrix multiplication
   • Remove scan_algorithms.rs dependency for Mamba-2

7.3 Architecture Gaps (Low Priority)

9. Input/Output Projection (P2)

   • Split input projection: in_proj → (z, x)
   • Add output gating: y * silu(z)

10. Cache Management (P2)

    • Implement Mamba2Cache for inference
    • Add KV cache for generation

11. Normalization (P3)

    • Add RMSNorm before final projection (reference uses this)
    • Replace LayerNorm with RMSNorm

12. Selective State Mechanism (P3)

    • Make B, C input-dependent (already doing this)
    • Add importance-based state compression (already have this)

---

8. Code Architecture Comparison

8.1 Reference Module Hierarchy

mamba_ssm/
├── ops/                    # CUDA kernels
│   ├── selective_scan.py
│   └── ssd_combined.py
├── modules/
│   ├── mamba_simple.py     # Mamba-1
│   └── mamba2.py           # Mamba-2 with SSD
└── models/
    └── mixer_seq_simple.py  # Full model

8.2 Our Module Hierarchy

ml/src/mamba/
├── mod.rs                  # Mamba2SSM (main model)
├── ssd_layer.rs            # Linear attention (NOT true SSD)
├── scan_algorithms.rs      # Parallel prefix scan (Mamba-1 style)
├── selective_state.rs      # State compression
└── hardware_aware.rs       # SIMD optimizations

Gap: We need to restructure to match reference architecture.

---

9. Rust/Candle Implementation Challenges

9.1 Candle Limitations

1. No Custom CUDA Kernels: Reference uses custom CUDA for SSD

   • Workaround: Implement SSD using Candle primitives (matmul, elementwise ops)

2. No Tensor Cores API: Candle doesn't expose tensor core control

   • Workaround: Use FP16/BF16 dtype, align dimensions to 8/16

3. No FlashAttention: Reference uses FlashAttention-style recomputation

   • Workaround: Manual gradient checkpointing

9.2 Candle Mamba Implementation

Found flawedmatrix/mamba-ssm (Rust/Candle implementation):

// Inference-only Mamba in Rust
// Uses CPU/Apple Silicon (no CUDA dependency)
// Generates at ~6.5 tokens/s with FP32 on M3 Max

Note: This is Mamba-1, not Mamba-2.

---

10. Recommendations

10.1 Immediate Actions (Wave 229)

1. Study Reference Code:

   • Clone tommyip/mamba2-minimal (single file, easiest to understand)
   • Read SSD algorithm implementation line-by-line
   • Map PyTorch ops to Candle equivalents

2. Implement dt Parameter:

   • Add dt_proj: Linear layer
   • Add dt_bias: Tensor parameter
   • Implement softplus + clamping

3. Add Convolution Layer:

   • Use Candle's Conv1d with causal padding
   • Kernel size: 4, groups: d_inner

4. Fix A Matrix Initialization:

   • Change from random normal to log(arange(1, d_state+1))
   • Make A diagonal (not full matrix)

10.2 Medium-term Refactor (Wave 230-232)

5. Implement SSD Algorithm:

   • Replace selective_scan_with_gradients with chunk-based SSD
   • Use matrix multiplication instead of sequential scan
   • Implement diagonal/off-diagonal block computation

6. Add D Parameter:

   • Initialize as torch.ones(d_inner)
   • Add skip connection: y = D * x + ssm_output

7. Restructure SSD Layer:

   • Split ssd_layer.rs into input/SSM/output components
   • Remove linear attention (not needed for Mamba-2)
   • Add proper input splitting (z, x)

10.3 Long-term Optimization (Wave 233-235)

8. Hardware Optimization:

   • Profile matrix operations
   • Use FP16 for training, BF16 for inference
   • Align dimensions to tensor core sizes

9. Benchmark Against Reference:

   • Run official Mamba-2 benchmark
   • Compare our implementation speed
   • Target: <2x slowdown vs PyTorch + CUDA

10. Integration Testing:

    • Test with real ES.FUT data
    • Validate gradient flow
    • Check numerical stability

---

11. Reference Documentation

11.1 Papers

1. Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Gu & Dao, 2023)

   • https://arxiv.org/abs/2312.00752

2. Transformers are SSMs: Generalized Models and Efficient Algorithms through Structured State Space Duality (Dao & Gu, 2024)

   • https://arxiv.org/abs/2405.21060

11.2 Blog Posts (Excellent Explanations)

1. Tri Dao's Blog (Author of Mamba-2):

   • Part I (Model): https://tridao.me/blog/2024/mamba2-part1-model/
   • Part II (Theory): https://tridao.me/blog/2024/mamba2-part2-theory/
   • Part III (Algorithm): https://tridao.me/blog/2024/mamba2-part3-algorithm/
   • Part IV (Systems): https://tridao.me/blog/2024/mamba2-part4-systems/

2. Princeton PLI Blog:

   • Mamba-2 Algorithms and Systems: https://pli.princeton.edu/blog/2024/mamba-2-algorithms-and-systems

3. From Mamba to Mamba-2 (n1o.github.io):

   • https://n1o.github.io/posts/from-mamba-to-mamba2/

11.3 Code Repositories

1. Official: https://github.com/state-spaces/mamba
2. Minimal: https://github.com/tommyip/mamba2-minimal
3. Hugging Face: https://github.com/huggingface/transformers/tree/main/src/transformers/models/mamba2
4. Rust (Mamba-1): https://github.com/flawedmatrix/mamba-ssm
5. Candle Examples: https://github.com/huggingface/candle/tree/main/candle-examples/examples/mamba

---

12. Conclusion

Our current implementation is functionally a Mamba-1 model with linear attention, not true Mamba-2 with SSD. To achieve the advertised 5x speedup and 8x larger state size, we must implement:

1. Structured State Duality (SSD) algorithm with chunk-based processing
2. Convolution layer for local context
3. Dynamic time step (dt) parameter
4. Skip connection (D) parameter
5. Diagonal A matrix constraint
6. Tensor core-friendly matrix operations

Estimated Effort: 3-4 weeks (Waves 229-232) for complete Mamba-2 implementation.

Next Agent: Agent 229 should start with dt parameter implementation (easiest, high impact).

---

Agent 228 Out 🎯