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**:
```bash
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):
```rust
// 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):
```rust
// 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):
```rust
// 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**:
```rust
let eps_tensor = Tensor::new(&[eps as f32], x.device())?;
```

**After**:
```rust
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):
```rust
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):
```rust
// 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):
```rust
// 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):
```rust
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**:
   ```rust
   // 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**:
   ```rust
   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**:
```rust
// 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**:
```rust
// 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**:
```rust
// 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