foxhunt/AGENT_168_SUMMARY.md

AGENT 168: MAMBA-2 E2E Test Execution Validation

Status: 🟡 IN PROGRESS - Critical bugs fixed, final shape mismatch remaining

Date: 2025-10-15 Mission: Execute MAMBA-2 E2E test suite after Agent 167 compilation success Test File: /home/jgrusewski/Work/foxhunt/ml/tests/e2e_mamba2_training.rs

---

Executive Summary

Executed comprehensive MAMBA-2 E2E test suite (7 tests) and identified/fixed critical bugs:

1. ✅ F32/F64 dtype mismatch in cuda_compat.rs layer normalization
2. ✅ B/C matrix dimension mismatch in MAMBA-2 state initialization
3. ✅ matmul transpose in prepare_scan_input
4. 🟡 Final shape mismatch in output transformation (investigation ongoing)

Test Results: 0/7 passing (down from 7/7 failures due to dtype, now 7/7 failures due to shape)

---

Test Execution Timeline

Attempt 1: Initial Run (Dtype Mismatch)

Command:

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

Result: 7/7 FAILED - F32/F64 dtype mismatch

Error:

Error: Model error: Candle error: dtype mismatch in add, lhs: F64, rhs: F32
Location: ml::cuda_compat::cuda_layer_norm (line 106)

Root Cause: cuda_compat.rs line 105 hardcoded eps as F32, but MAMBA-2 uses F64 throughout.

Fix Applied (Agent 168):

// BEFORE (line 105):
let eps_tensor = Tensor::new(&[eps as f32], x.device())?;

// AFTER (lines 106-110):
let eps_tensor = match x.dtype() {
    candle_core::DType::F32 => Tensor::new(&[eps as f32], x.device())?,
    candle_core::DType::F64 => Tensor::new(&[eps], x.device())?,
    _ => return Err(MLError::ModelError(format!("Unsupported dtype for layer norm: {:?}", x.dtype()))),
};

Impact: Dtype errors eliminated ✅

---

Attempt 2: After Dtype Fix (Shape Mismatch - B/C Matrices)

Result: 7/7 FAILED - Matrix dimension mismatch

Error:

Error: Model error: Candle error: shape mismatch in matmul, lhs: [8, 60, 1024], rhs: [16, 256]
Location: ml::mamba::Mamba2SSM::forward (prepare_scan_input line 699)

Analysis:

• LHS: [batch=8, seq=60, d_inner=1024] (input after input_projection expansion)
• RHS: [d_state=16, d_model=256] (B matrix)
• Problem: B initialized with d_model instead of d_inner

Architecture Flow:

Input: [batch, seq, d_model=256]
  ↓ input_projection (linear: d_model → d_inner)
  ↓ [batch, seq, d_inner=1024] (expansion_factor=4)
  ↓ SSM layers (B matrix must match d_inner!)
  ↓ output_projection (linear: d_inner → 1)
Output: [batch, seq, 1]

Root Cause: SSM matrices B and C initialized with d_model instead of d_inner in Mamba2State::zeros().

Fix Applied (Agent 168):

// BEFORE (lines 243, 250):
let B = Tensor::randn(0.0, 1.0, (config.d_state, config.d_model), device)...
let C = Tensor::randn(0.0, 1.0, (config.d_model, config.d_state), device)...

// AFTER (lines 225, 245, 253):
let d_inner = config.d_model * config.expand; // CRITICAL: Use d_inner

let B = Tensor::randn(0.0, 1.0, (config.d_state, d_inner), device)...
// [16, 1024] instead of [16, 256]

let C = Tensor::randn(0.0, 1.0, (d_inner, config.d_state), device)...
// [1024, 16] instead of [256, 16]

Impact: B/C dimensions now match expanded input ✅

---

Attempt 3: After B/C Fix (Shape Mismatch - Missing Transpose)

Result: 7/7 FAILED - Matrix dimension mismatch (different error!)

Error:

Error: Model error: Candle error: shape mismatch in matmul, lhs: [8, 60, 1024], rhs: [16, 1024]
Location: ml::mamba::Mamba2SSM::forward (prepare_scan_input line 702)

Analysis:

• LHS: [batch=8, seq=60, d_inner=1024]
• RHS: [d_state=16, d_inner=1024] (B matrix)
• Problem: matmul requires compatible dimensions: [..., M, K] @ [K, N]

Expected Dimensions:

input: [8, 60, 1024]
B.t(): [1024, 16]  (transposed from [16, 1024])
Result: [8, 60, 16]

Fix Applied (Agent 168):

// BEFORE (line 702):
let Bu = input.matmul(B)?;

// AFTER (lines 701-704):
// FIXED: Transpose B to match matmul dimensions
// input: [batch, seq, d_inner], B: [d_state, d_inner]
// B.t(): [d_inner, d_state] → result: [batch, seq, d_state]
let Bu = input.matmul(&B.t()?)?;

Impact: Transpose added to prepare_scan_input() ✅

---

Attempt 4: Current Status (Output Transformation Shape Mismatch)

Result: 7/7 FAILED - Matrix dimension mismatch in final output

Error:

Error: Model error: Candle error: shape mismatch in matmul, lhs: [8, 60, 1024], rhs: [1024, 16]
Location: ml::mamba::Mamba2SSM::forward (line 633)

Analysis:

• Error Location: scanned_states.matmul(&C.t()?)
• Expected: [8, 60, 16] @ [16, 1024] = [8, 60, 1024]
• Actual Error: [8, 60, 1024] @ [1024, 16] (dimensions inverted!)

Hypothesis: The error message suggests scanned_states is [8, 60, 1024] instead of [8, 60, 16]. This indicates:

1. Option A: parallel_prefix_scan() is returning the wrong shape
2. Option B: prepare_scan_input() is NOT reducing dimensions as expected
3. Option C: The scan_engine is passing through input unchanged (placeholder behavior)

Investigation Needed:

• Check ParallelScanEngine::parallel_prefix_scan() implementation
• Verify ScanOperator::SSMScan behavior
• Trace tensor shapes through scan pipeline

---

Files Modified

1. /home/jgrusewski/Work/foxhunt/ml/src/cuda_compat.rs

Lines Changed: 105-111 (+6 lines)

Change: Dynamic dtype handling for epsilon tensor in layer normalization

Before:

let eps_tensor = Tensor::new(&[eps as f32], x.device())?;

After:

let eps_tensor = match x.dtype() {
    candle_core::DType::F32 => Tensor::new(&[eps as f32], x.device())?,
    candle_core::DType::F64 => Tensor::new(&[eps], x.device())?,
    _ => return Err(MLError::ModelError(format!("Unsupported dtype for layer norm: {:?}", x.dtype()))),
};

Impact: Fixes F32/F64 dtype mismatch for MAMBA-2 (uses F64) and other models (DQN/PPO use F32)

---

2. /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs

Lines Changed: 222-258 (+4 lines, modified 3 lines)

Change 1: Added d_inner calculation in Mamba2State::zeros()

Line 225 (added):

let d_inner = config.d_model * config.expand; // CRITICAL: Use d_inner after input_projection

Change 2: Updated B matrix dimensions

Lines 244-250 (modified):

// FIXED: B must be [d_state, d_inner] to match expanded input dimension after input_projection
let B = Tensor::randn(0.0, 1.0, (config.d_state, d_inner), device).map_err(
    |e| MLError::TensorCreationError {
        operation: format!("SSM B matrix creation for layer {}", layer_idx),
        reason: e.to_string(),
    },
)?;

Change 3: Updated C matrix dimensions

Lines 252-258 (modified):

// FIXED: C must be [d_inner, d_state] to match expanded hidden dimension after input_projection
let C = Tensor::randn(0.0, 1.0, (d_inner, config.d_state), device).map_err(
    |e| MLError::TensorCreationError {
        operation: format!("SSM C matrix creation for layer {}", layer_idx),
        reason: e.to_string(),
    },
)?;

Change 4: Fixed prepare_scan_input() transpose

Lines 695-706 (modified):

fn prepare_scan_input(
    &self,
    input: &Tensor,
    _A: &Tensor,
    B: &Tensor,
) -> Result<Tensor, MLError> {
    // FIXED: Transpose B to match matmul dimensions
    // input: [batch, seq, d_inner], B: [d_state, d_inner]
    // B.t(): [d_inner, d_state] → result: [batch, seq, d_state]
    let Bu = input.matmul(&B.t()?)?;
    Ok(Bu)
}

Impact:

• B matrix: [16, 256] → [16, 1024] ✅
• C matrix: [256, 16] → [1024, 16] ✅
• Transpose added to prepare_scan_input ✅

---

Test Results Summary

Test Name	Status	Error Type
test_mamba2_simple_forward_pass	❌ FAILED	Shape mismatch in output transformation
test_mamba2_batch_shapes	❌ FAILED	Shape mismatch in output transformation
test_mamba2_cuda_device	❌ FAILED	Shape mismatch in output transformation
test_mamba2_sequence_lengths	❌ FAILED	Shape mismatch in output transformation
test_mamba2_gradient_flow	❌ FAILED	Shape mismatch in output transformation
test_mamba2_training_loop_simple	❌ FAILED	Shape mismatch in output transformation
test_mamba2_config_variations	❌ FAILED	Shape mismatch in output transformation

Pass Rate: 0/7 (0%) Duration: 0.39-0.56 seconds

---

Bugs Fixed

Bug #1: F32/F64 Dtype Mismatch in Layer Normalization ✅

Severity: 🔴 CRITICAL (blocks all MAMBA-2 inference)

Root Cause: cuda_compat.rs line 105 hardcoded epsilon as F32, but MAMBA-2 uses F64 for all tensors.

Error Message:

Candle error: dtype mismatch in add, lhs: F64, rhs: F32

Fix: Added dtype matching logic to dynamically create F32 or F64 epsilon tensor based on input dtype.

Test Coverage: Affects all 7 E2E tests Status: ✅ FIXED (Agent 168)

---

Bug #2: B/C Matrix Dimension Mismatch ✅

Severity: 🔴 CRITICAL (architectural design flaw)

Root Cause: SSM matrices B and C initialized with d_model=256 instead of d_inner=1024, causing shape mismatch after input_projection expansion.

Error Message:

Candle error: shape mismatch in matmul, lhs: [8, 60, 1024], rhs: [16, 256]

Fix: Updated Mamba2State::zeros() to use d_inner = d_model * expand for B/C matrix dimensions.

Dimension Changes:

• B: [d_state, d_model] → [d_state, d_inner] (e.g., [16, 256] → [16, 1024])
• C: [d_model, d_state] → [d_inner, d_state] (e.g., [256, 16] → [1024, 16])

Test Coverage: Affects all 7 E2E tests Status: ✅ FIXED (Agent 168)

---

Bug #3: Missing Transpose in prepare_scan_input ✅

Severity: 🔴 CRITICAL (matmul incompatibility)

Root Cause: prepare_scan_input() performed input.matmul(B) without transposing B, causing incompatible matmul dimensions.

Error Message:

Candle error: shape mismatch in matmul, lhs: [8, 60, 1024], rhs: [16, 1024]

Fix: Added B.t()? to transpose B matrix before matmul.

Dimension Flow:

input: [batch=8, seq=60, d_inner=1024]
B: [d_state=16, d_inner=1024]
B.t(): [d_inner=1024, d_state=16]
Result: [batch=8, seq=60, d_state=16] ✅

Test Coverage: Affects all 7 E2E tests Status: ✅ FIXED (Agent 168)

---

Bug #4: Output Transformation Shape Mismatch 🟡

Severity: 🔴 CRITICAL (blocks all E2E tests)

Root Cause: UNDER INVESTIGATION

Error Message:

Candle error: shape mismatch in matmul, lhs: [8, 60, 1024], rhs: [1024, 16]

Hypothesis: parallel_prefix_scan() may be returning input unchanged (placeholder behavior) instead of producing [batch, seq, d_state] output.

Expected Flow:

prepare_scan_input: [8, 60, 1024] @ [1024, 16] → [8, 60, 16]
parallel_prefix_scan: [8, 60, 16] → [8, 60, 16] (scan operation)
output transformation: [8, 60, 16] @ [16, 1024] → [8, 60, 1024]

Actual Flow (suspected):

prepare_scan_input: [8, 60, 1024] @ [1024, 16] → [8, 60, 16]
parallel_prefix_scan: [8, 60, 16] → [8, 60, 1024] (⚠️ wrong shape!)
output transformation: [8, 60, 1024] @ [1024, 16] → ❌ SHAPE MISMATCH

Next Steps:

1. Instrument parallel_prefix_scan() to log input/output shapes
2. Check ScanOperator::SSMScan implementation
3. Verify scan_engine is NOT passing through input unchanged
4. Add shape assertions between pipeline stages

Test Coverage: Affects all 7 E2E tests Status: 🟡 IN PROGRESS (Agent 168)

---

Performance Metrics

Compilation Time: 3m 11s (first run), ~20-40s (subsequent) Test Duration: 0.39-0.56 seconds (all 7 tests) GPU Devices Used: CUDA devices 1-7 (parallel test execution)

Compilation Warnings:

• 17 warnings in ml crate (unused imports, unsafe blocks, missing Debug)
• 66 warnings in test binary (unused dependencies)

Test Execution Speed: Very fast (~80ms per test), indicating early failure in forward pass.

---

Next Actions

IMMEDIATE (Agent 169 or continuation)

1. Debug parallel_prefix_scan shape issue:

   // Add to line 627 in mod.rs:
   println!("scan_input shape: {:?}", scan_input.dims());
   println!("scanned_states shape: {:?}", scanned_states.dims());
   

2. Verify ParallelScanEngine implementation:

   • Check scan_algorithms.rs line 111-130
   • Confirm SSMScan operator behavior
   • Ensure output shape matches input shape for d_state dimension

3. Add shape assertions:

   assert_eq!(scan_input.dims(), &[batch_size, seq_len, d_state]);
   assert_eq!(scanned_states.dims(), &[batch_size, seq_len, d_state]);
   

4. Fix output transformation:

   • If scanned_states is [8, 60, 16]: Use scanned_states.matmul(&C.t()?) ✅
   • If scanned_states is [8, 60, 1024]: Investigate why scan didn't reduce dimensions

MEDIUM PRIORITY

5. Fix compilation warnings (technical debt):

   • Remove unused imports (Device, ModelVote, TradingAction)
   • Add #[derive(Debug)] to 6 structs
   • Prefix unused variables with _
   • Add #[allow(unsafe_code)] justification comments

6. Add integration tests for scan_engine:

   • Test parallel_prefix_scan with various input shapes
   • Verify SSMScan operator correctness
   • Benchmark scan performance vs sequential

LOW PRIORITY

7. Documentation updates:
   • Update ml/src/mamba/mod.rs docstrings with d_inner clarifications
   • Add architecture diagram showing dimension flow
   • Document B/C matrix dimension requirements

---

Recommendations

For Next Agent (Agent 169)

Focus: Fix final shape mismatch in output transformation (Bug #4)

Approach:

1. Add debug logging to trace tensor shapes through SSM pipeline
2. Verify parallel_prefix_scan() implementation in scan_algorithms.rs
3. Check if SSMScan operator is a placeholder (just returns input)
4. Fix scan logic to maintain [batch, seq, d_state] shape

Expected Outcome: 7/7 tests passing after scan shape fix

For MAMBA-2 Training

Once tests pass:

1. Execute GPU training benchmark (30-60 min) to determine training platform
2. Download 90 days ES/NQ/ZN/6E data (~$2, 180K bars)
3. Begin 4-6 week ML training pipeline
4. Validate trained models with production data

Current Blocker: E2E tests must pass before production training begins

---

Lessons Learned

1. Architecture Misalignment (B/C Matrix Bug)

Issue: SSM matrices initialized with d_model but used after input_projection expansion to d_inner.

Lesson: Always trace tensor dimensions through the ENTIRE pipeline, especially when projection layers change dimensions.

Prevention:

// Add shape assertions after each transformation:
assert_eq!(hidden.dims()[2], d_inner, "Expected d_inner after input_projection");
assert_eq!(Bu.dims()[2], d_state, "Expected d_state after B projection");

2. Dtype Consistency (F32/F64 Bug)

Issue: Hardcoded F32 epsilon in generic layer norm function, breaking F64 models.

Lesson: Always use x.dtype() to match input tensor dtype for scalar operations.

Prevention:

// Pattern for dtype-agnostic operations:
let scalar = match input.dtype() {
    DType::F32 => Tensor::new(&[value as f32], device)?,
    DType::F64 => Tensor::new(&[value], device)?,
    _ => return Err(...),
};

3. Transpose Assumptions (prepare_scan_input Bug)

Issue: Assumed B matrix orientation without verifying matmul compatibility.

Lesson: ALWAYS check matmul dimension compatibility: [..., M, K] @ [K, N] → [..., M, N]

Prevention:

// Annotate expected shapes in comments:
// input: [batch, seq, d_inner]
// B: [d_state, d_inner]
// B.t(): [d_inner, d_state]
// Result: [batch, seq, d_state]
let Bu = input.matmul(&B.t()?)?;

4. Test-Driven Development Value

Impact: Agent 146's comprehensive E2E tests caught ALL these bugs before production.

Without E2E Tests:

• Bugs would appear during training (wasted 4-6 weeks)
• Debugging would be harder (production data, GPU constraints)
• Risk of corrupted checkpoints/wasted compute

With E2E Tests:

• Bugs caught in <5 minutes
• Fixed in isolation with clear error messages
• Training can proceed with confidence

---

References

Related Agents:

• Agent 146: Created 7 comprehensive MAMBA-2 E2E tests
• Agent 155: Fixed F32→F64 dtype in MAMBA-2 core
• Agent 167: Fixed compilation errors in MAMBA-2
• Agent 168: This agent - executed tests and fixed critical bugs

Files:

• /home/jgrusewski/Work/foxhunt/ml/src/cuda_compat.rs (layer norm dtype fix)
• /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs (B/C matrix fix, transpose fix)
• /home/jgrusewski/Work/foxhunt/ml/src/mamba/scan_algorithms.rs (investigation target)
• /home/jgrusewski/Work/foxhunt/ml/tests/e2e_mamba2_training.rs (test suite)

Documentation:

• AGENT_146_MAMBA2_TESTS.md (E2E test design)
• AGENT_155_SUMMARY.md (F64 dtype standardization)
• AGENT_167_SUMMARY.md (compilation fixes)
• GPU_TRAINING_BENCHMARK.md (next milestone)

---

End of Report

Agent 168 Status: 🟡 PARTIAL SUCCESS - Fixed 3/4 critical bugs, 1 remaining Next Agent: Agent 169 - Fix parallel_prefix_scan shape mismatch Production Readiness: 🔴 BLOCKED - E2E tests must pass before ML training