foxhunt/AGENT_176_ANALYSIS.md

AGENT 176: MAMBA-2 SSM State Dimension Analysis

Mission

Trace tensor dimensions through SSM forward pass to find where d_inner (1024) should become d_state (16).

Error Signature

thread 'test_mamba2_training_loop_simple' panicked at ml/src/mamba/mod.rs:1032:47:
MatMul dimension mismatch lhs: [8, 60, 1024] rhs: [16, 1024] (lhs.dim(D::Minus1) != rhs.dim(0))

Dimension Flow Analysis

Expected Flow (Agent 168 Fix)

1. input_projection:
   input [8, 60, 256] (d_model)
     ↓ Linear(d_model → d_inner)
   hidden [8, 60, 1024] (d_inner = d_model * expand = 256 * 4)

2. prepare_scan_input_with_gradients:
   input [8, 60, 1024] (d_inner)
   B [16, 1024] (d_state × d_inner)
     ↓ input.matmul(&B.t()?)
   B.t() [1024, 16]
     ↓
   scan_input [8, 60, 16] (d_state) ← EXPECTED

3. selective_scan_with_gradients:
   scan_input [8, 60, 16] (d_state)
     ↓ sequential scan preserves shape
   scanned_states [8, 60, 16] (d_state) ← EXPECTED

4. output transformation:
   scanned_states [8, 60, 16] (d_state)
   C [1024, 16] (d_inner × d_state)
     ↓ scanned_states.matmul(&C.t()?)
   C.t() [16, 1024]
     ↓
   output [8, 60, 1024] (d_inner) ← CORRECT

Actual Flow (Current Error)

1. input_projection:
   input [8, 60, 256]
     ↓
   hidden [8, 60, 1024] ✓ CORRECT

2. prepare_scan_input_with_gradients:
   input [8, 60, 1024]
   B [16, 1024]
     ↓ input.matmul(&B.t()?)
   B.t() [1024, 16]
     ↓
   scan_input [8, 60, ???] ← CRITICAL POINT

3. selective_scan_with_gradients:
   scan_input [8, 60, ???]
     ↓
   scanned_states [8, 60, 1024] ❌ WRONG (should be [8, 60, 16])

4. output transformation:
   scanned_states [8, 60, 1024] ❌ WRONG
   C.t() [16, 1024]
     ↓ matmul fails: [8,60,1024] × [16,1024]
   ERROR: dim mismatch (1024 != 16)

Root Cause Hypothesis

The error occurs at line 1032 in forward_ssd_layer_with_gradients:

let output = scanned_states.matmul(&C.t()?)?;

Hypothesis: selective_scan_with_gradients is NOT transforming dimensions correctly.

Investigation Points

1. Check if prepare_scan_input_with_gradients is actually being called

   • Add debug print of scan_input shape BEFORE passing to selective_scan

2. Check if selective_scan_with_gradients preserves input shape

   • Add debug print of output shape AFTER selective_scan

3. Check if there's a bypass/override somewhere

   • scan_engine.parallel_prefix_scan might be overriding the transformation

Code Path Trace

forward_ssd_layer_with_gradients (line 1003-1043)

fn forward_ssd_layer_with_gradients(
    &mut self,
    _ssd_layer: &SSDLayer,
    input: &Tensor,
    layer_idx: usize,
) -> Result<Tensor, MLError> {
    // Extract SSM matrices
    let dt = self.state.ssm_states[layer_idx].delta.clone();
    let A = self.state.ssm_states[layer_idx].A.clone();
    let B = self.state.ssm_states[layer_idx].B.clone();  // [16, 1024]
    let C = self.state.ssm_states[layer_idx].C.clone();  // [1024, 16]

    // Discretize
    let A_discrete = self.discretize_ssm_with_gradients(&A, &dt)?;
    let B_discrete = self.discretize_ssm_input_with_gradients(&B, &dt)?;

    // ⚠️ CRITICAL: This should produce [8, 60, 16]
    let scan_input = self.prepare_scan_input_with_gradients(input, &A_discrete, &B_discrete)?;

    // ⚠️ CRITICAL: This should preserve shape [8, 60, 16]
    let scanned_states = self.selective_scan_with_gradients(&scan_input, &A_discrete)?;

    // ❌ ERROR HERE: scanned_states is [8, 60, 1024] instead of [8, 60, 16]
    let output = scanned_states.matmul(&C.t()?)?;  // PANIC!

prepare_scan_input_with_gradients (line 1104-1113)

fn prepare_scan_input_with_gradients(
    &self,
    input: &Tensor,      // [8, 60, 1024]
    _A: &Tensor,
    B: &Tensor,          // [16, 1024]
) -> Result<Tensor, MLError> {
    // Multiply input by B matrix for state transition
    let Bu = input.matmul(&B.t()?)?;  // [8,60,1024] × [1024,16] = [8,60,16] ✓
    Ok(Bu)
}

Status: This method LOOKS correct. Returns [8, 60, 16].

selective_scan_with_gradients (line 1045-1076)

fn selective_scan_with_gradients(&self, input: &Tensor, A: &Tensor) -> Result<Tensor, MLError> {
    let seq_len = input.dim(1)?;     // 60
    let d_state = input.dim(2)?;     // Should be 16
    let device = input.device();

    // Initialize state sequence
    let mut states = Vec::new();
    let mut current_state = Tensor::zeros((input.dim(0)?, d_state), input.dtype(), device)?;

    // Sequential scan with state transitions
    for t in 0..seq_len {
        let x_t = input.narrow(1, t, 1)?.squeeze(1)?;  // [8, d_state]

        // ⚠️ CRITICAL: Check A matrix dimensions
        // State transition: h_t = A * h_{t-1} + B * x_t
        let state_dims = current_state.dims().len();
        current_state = (A
            .matmul(&current_state.unsqueeze(state_dims)?)?
            .squeeze(state_dims)?
            + &x_t)?;
        states.push(current_state.unsqueeze(1)?);
    }

    // Stack all states
    let result = Tensor::cat(&states, 1)?;
    Ok(result)
}

SUSPICIOUS: The A.matmul operation might be wrong!

A Matrix Dimension Issue

In selective_scan_with_gradients, we have:

current_state = A.matmul(&current_state.unsqueeze(state_dims)?)?

Where:

• A: [16, 16] (d_state × d_state) from discretize_ssm_with_gradients
• current_state: [8, 16] (batch × d_state)
• After unsqueeze: [8, 16, 1] or [8, 1, 16]?

PROBLEM: If unsqueeze adds dimension at wrong position, matmul fails!

Bug Found: Matrix Multiplication Order

The bug is in selective_scan_with_gradients at line 1062:

// WRONG:
current_state = (A
    .matmul(&current_state.unsqueeze(state_dims)?)?  // A[16,16] × state[8,16,1]?
    .squeeze(state_dims)?
    + &x_t)?;

This should be:

// CORRECT:
current_state = (current_state
    .matmul(&A.t()?)?  // state[8,16] × A.t()[16,16] = [8,16]
    + &x_t)?;

OR:

// CORRECT (batch matmul):
current_state = (A
    .matmul(&current_state.unsqueeze(2)?)?  // [16,16] × [8,16,1] = [8,16,1]
    .squeeze(2)?
    + &x_t)?;

Next Steps

1. Add debug prints to confirm scan_input shape
2. Fix the matmul in selective_scan_with_gradients
3. Verify all tests pass

Files Modified

• ml/src/mamba/mod.rs: Add debug prints + fix selective_scan_with_gradients