foxhunt/AGENT_182_QUICK_FIX.md

AGENT 182 QUICK FIX: parallel_prefix_scan Shape Bug

Mission: Fix parallel_prefix_scan to preserve [batch, seq, d_state] shape

Priority: 🔴 CRITICAL - Blocking all MAMBA-2 training (0/7 tests passing)

---

🎯 The Bug

File: /home/jgrusewski/Work/foxhunt/ml/src/mamba/scan_algorithms.rs:111

Function: parallel_prefix_scan

Problem: Returns [batch, seq, d_inner] instead of [batch, seq, d_state]

Impact: Causes shape mismatch at line 633 of /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs:

let output = scanned_states.matmul(&C.t()?)?;
// ERROR: [8, 60, 1024] @ [1024, 16] - dimension mismatch!

---

🔍 Root Cause

Expected Behavior

Input to parallel_prefix_scan:  [8, 60, 16] (d_state)
Output from parallel_prefix_scan: [8, 60, 16] (preserve shape)

Actual Behavior

Input to parallel_prefix_scan:  [8, 60, 16] (d_state)
Output from parallel_prefix_scan: [8, 60, 1024] (d_inner) ❌ WRONG!

Where the Bug Occurs

The scan algorithm is likely using the wrong tensor in one of these functions:

1. sequential_scan (line 148)
2. block_parallel_scan (called from line 124)

Hypothesis: One of these functions is using the original input ([*, *, d_inner]) instead of the scan input ([*, *, d_state]).

---

🔧 Investigation Steps

Step 1: Check sequential_scan

# Search for where the result tensor is created in sequential_scan
grep -A 30 "fn sequential_scan" ml/src/mamba/scan_algorithms.rs

Look for:

• Result tensor creation
• Shape used for result allocation
• Which tensor is being scanned (should be input parameter, not anything else)

Step 2: Check block_parallel_scan

# Search for block_parallel_scan implementation
grep -A 50 "fn block_parallel_scan" ml/src/mamba/scan_algorithms.rs

Look for:

• Block size calculations using wrong dimensions
• Result tensor shape allocation
• Concatenation operations that might expand dimensions

Step 3: Look for d_inner references

# Check if scan_algorithms.rs incorrectly references d_inner
grep -n "d_inner\|1024" ml/src/mamba/scan_algorithms.rs

Expected: NO references to d_inner or hardcoded 1024 in scan_algorithms.rs

---

🎯 Likely Fix

Scenario A: Using Wrong Tensor

If the scan is using self.state.hidden or input_projection output instead of the input parameter:

// WRONG:
let result = self.scan(self.hidden_state)?; // Uses d_inner dimension

// CORRECT:
let result = self.scan(input)?; // Uses d_state dimension from parameter

Scenario B: Wrong Result Shape Allocation

If the result tensor is allocated with wrong dimensions:

// WRONG:
let result = Tensor::zeros((batch_size, seq_len, d_inner), ...)?;

// CORRECT:
let result = Tensor::zeros((batch_size, seq_len, input.dim(2)?), ...)?;

Scenario C: Accumulator Shape Bug

If the accumulator in sequential_scan is using wrong shape:

// WRONG:
let mut accumulator = Tensor::zeros((batch_size, 1, d_inner), ...)?;

// CORRECT:
let mut accumulator = input.narrow(0, 0, 1)?.narrow(1, 0, 1)?; // Use input shape

---

📝 Files to Modify

Primary:

• /home/jgrusewski/Work/foxhunt/ml/src/mamba/scan_algorithms.rs

Verify:

• /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs (no changes needed, already correct)

---

✅ Success Criteria

After fix, run:

cargo test -p ml --test e2e_mamba2_training --features cuda

Expected:

test result: ok. 7 passed; 0 failed

Test that will pass first: test_mamba2_simple_forward_pass

Shape trace should show:

scan_input:      [8, 60, 16] ✅
scanned_states:  [8, 60, 16] ✅ (not [8, 60, 1024])
output:          [8, 60, 1024] ✅

---

🚨 Critical Notes

1. DO NOT modify B/C matrix shapes - They are already correct!
2. DO NOT modify prepare_scan_input - It's working correctly!
3. ONLY fix the scan algorithm - Shape should be preserved

---

📊 Test Configuration

d_model: 256
d_state: 16
expand: 4
d_inner: 1024 (256 * 4)

B: [16, 1024] (d_state × d_inner) ✅
C: [1024, 16] (d_inner × d_state) ✅
scan_input: [8, 60, 16] ✅
scanned_states: [8, 60, 16] ← FIX THIS (currently [8, 60, 1024])

---

🔬 Debugging Commands

# Run single test with full output
cargo test -p ml test_mamba2_simple_forward_pass --features cuda -- --nocapture

# Check scan_algorithms.rs for dimension bugs
rg "d_inner|1024" ml/src/mamba/scan_algorithms.rs

# Look for tensor shape allocations
rg "Tensor::zeros|Tensor::ones" ml/src/mamba/scan_algorithms.rs

---

⏱️ Estimated Fix Time

30-60 minutes (scan algorithm is isolated module)

Confidence: ✅ High - Root cause clearly identified, fix is localized

---

End of Quick Fix Guide