# Agent 228: MAMBA-2 Implementation Gaps

**Date**: 2025-10-15
**Priority**: P0 - CRITICAL (Training Pipeline Blocked)
**Estimated Fix Time**: 3-4 weeks (Waves 229-232)

---

## Gap Summary

| Gap | Priority | Impact | Effort | Status |
|-----|----------|--------|--------|--------|
| 1. SSD Algorithm | P0 | 5x speedup | 2 weeks | ❌ Not Implemented |
| 2. Convolution Layer | P0 | Model accuracy | 2 days | ❌ Missing |
| 3. dt (Time Step) | P0 | Selectivity | 3 days | ❌ Missing |
| 4. D (Skip Connection) | P0 | Training stability | 1 day | ❌ Missing |
| 5. A Matrix Init | P1 | Convergence | 1 day | ❌ Wrong |
| 6. d_state Size | P1 | Model capacity | 1 day | ❌ 8x too small |
| 7. Tensor Core Ops | P1 | GPU utilization | 1 week | ❌ No optimization |
| 8. Memory Access | P1 | Bandwidth | 3 days | ❌ Unoptimized |
| 9. Input Splitting | P2 | Architecture | 2 days | ❌ Missing |
| 10. Mamba2Cache | P2 | Inference speed | 3 days | ❌ Missing |
| 11. RMSNorm | P3 | Training speed | 1 day | ❌ Using LayerNorm |
| 12. State Importance | P3 | Memory | - | ✅ Implemented |

**Total Gaps**: 12 (11 missing, 1 complete)
**Critical Gaps**: 4 (SSD, Conv, dt, D)
**Training Blocked**: YES (cannot train without P0 gaps fixed)

---

## Gap 1: SSD Algorithm (P0 - CRITICAL)

### Current Implementation (WRONG)

```rust
// ml/src/mamba/mod.rs:1108
fn selective_scan_with_gradients(&self, input: &Tensor, A: &Tensor) -> Result<Tensor, MLError> {
    // Sequential scan (Mamba-1 style)
    for t in 0..seq_len {
        let x_t = input.narrow(1, t, 1)?.squeeze(1)?;
        current_state = (current_state.matmul(&A.t()?)? + &x_t)?;  // O(n) sequential
        states.push(current_state.unsqueeze(1)?);
    }
    Tensor::cat(&states, 1)
}
```

**Problem**: This is **Mamba-1 parallel associative scan**, not Mamba-2 SSD.

### Reference Implementation (CORRECT)

```python
# tommyip/mamba2-minimal/mamba2.py
def ssd_chunk_scan(x, dt, A, B, C, chunk_size=256):
    """Structured State Duality chunk scan"""
    batch, seqlen, dim = x.shape
    num_chunks = seqlen // chunk_size

    # Discretize parameters
    dt = F.softplus(dt)
    A_discrete = torch.exp(A * dt)  # Diagonal A
    B_discrete = B * dt

    # Chunk-based processing (tensor core friendly)
    states = []
    for chunk_idx in range(num_chunks):
        # Extract chunk
        chunk_x = x[:, chunk_idx*chunk_size:(chunk_idx+1)*chunk_size]

        # Diagonal block: local state computation (parallel)
        diag_block = compute_diagonal(chunk_x, A_discrete, B_discrete)

        # Off-diagonal block: inter-chunk dependencies
        if chunk_idx > 0:
            offdiag_block = compute_offdiagonal(prev_state, A_discrete, chunk_size)
            diag_block = diag_block + offdiag_block

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

    # Concatenate chunks
    states = torch.cat(states, dim=1)

    # Output transformation: y = states @ C^T
    y = torch.einsum('bld,dc->blc', states, C)
    return y
```

### Required Changes

1. **Replace Sequential Scan**:
   ```rust
   // NEW: ml/src/mamba/ssd_algorithm.rs
   pub fn ssd_chunk_scan(
       input: &Tensor,
       dt: &Tensor,
       A: &Tensor,
       B: &Tensor,
       C: &Tensor,
       chunk_size: usize,
   ) -> Result<Tensor, MLError> {
       // Implement chunk-based SSD algorithm
   }
   ```

2. **Add Diagonal Block Computation**:
   ```rust
   fn compute_diagonal_block(
       chunk_x: &Tensor,
       A_discrete: &Tensor,
       B_discrete: &Tensor,
   ) -> Result<Tensor, MLError> {
       // Matrix multiplication: chunk_x @ B_discrete^T
       // Then scale by A_discrete
   }
   ```

3. **Add Off-Diagonal Block Computation**:
   ```rust
   fn compute_offdiagonal_block(
       prev_state: &Tensor,
       A_discrete: &Tensor,
       chunk_size: usize,
   ) -> Result<Tensor, MLError> {
       // Propagate previous chunk's final state
   }
   ```

**Estimated Effort**: 2 weeks (complex algorithm)
**Blocking**: Training pipeline
**Dependencies**: dt parameter, diagonal A matrix

---

## Gap 2: Convolution Layer (P0)

### Current Implementation (MISSING)

```rust
// ml/src/mamba/mod.rs:570-608
pub fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
    let mut hidden = self.input_projection.forward(input)?;
    // ❌ No convolution here!
    for layer_idx in 0..num_layers {
        let normalized = self.layer_norms[layer_idx].forward(&hidden)?;
        let layer_output = self.forward_ssd_layer(&ssd_layer, &normalized, layer_idx)?;
        ...
    }
}
```

### Reference Implementation (CORRECT)

```python
# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_model, d_conv=4):
        self.conv1d = nn.Conv1d(
            in_channels=d_inner,
            out_channels=d_inner,
            kernel_size=d_conv,
            groups=d_inner,  # Depthwise convolution
            padding=d_conv - 1  # Causal padding
        )

    def forward(self, x):
        x, z = self.in_proj(u).chunk(2, dim=-1)

        # Causal convolution (MISSING in our implementation)
        x = rearrange(x, 'b l d -> b d l')
        x = self.conv1d(x)[:, :, :seqlen]  # Remove extra padding
        x = rearrange(x, 'b d l -> b l d')

        x = F.silu(x)
        y = self.ssm(x)
        ...
```

### Required Changes

1. **Add Conv1d to Mamba2SSM**:
   ```rust
   // ml/src/mamba/mod.rs
   pub struct Mamba2SSM {
       pub input_projection: Linear,
       pub conv1d: Vec<Conv1d>,  // NEW: One per layer
       pub layer_norms: Vec<CudaLayerNorm>,
       ...
   }
   ```

2. **Initialize Convolution**:
   ```rust
   impl Mamba2SSM {
       pub fn new(config: Mamba2Config, device: &Device) -> Result<Self, MLError> {
           let mut conv1d = Vec::new();
           for i in 0..config.num_layers {
               let conv = candle_nn::conv1d(
                   d_inner,           // in_channels
                   d_inner,           // out_channels
                   4,                 // kernel_size (d_conv)
                   candle_nn::Conv1dConfig {
                       groups: d_inner,     // Depthwise
                       padding: 3,          // Causal padding (d_conv - 1)
                       ..Default::default()
                   },
                   vb.pp(&format!("conv1d_{}", i)),
               )?;
               conv1d.push(conv);
           }
           ...
       }
   }
   ```

3. **Apply Convolution in Forward Pass**:
   ```rust
   pub fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
       let mut hidden = self.input_projection.forward(input)?;

       for layer_idx in 0..num_layers {
           // Apply convolution BEFORE layer norm
           hidden = hidden.transpose(1, 2)?;  // [batch, d_inner, seq_len]
           hidden = self.conv1d[layer_idx].forward(&hidden)?;
           hidden = hidden.narrow(2, 0, seq_len)?;  // Remove extra padding
           hidden = hidden.transpose(1, 2)?;  // [batch, seq_len, d_inner]

           hidden = hidden.silu()?;  // SiLU activation

           let normalized = self.layer_norms[layer_idx].forward(&hidden)?;
           ...
       }
   }
   ```

**Estimated Effort**: 2 days
**Blocking**: Training accuracy
**Dependencies**: None (can implement immediately)

---

## Gap 3: dt (Time Step) Parameter (P0)

### Current Implementation (WRONG)

```rust
// ml/src/mamba/mod.rs:261-266
let delta = Tensor::ones((config.d_model,), DType::F64, device)?;  // CONSTANT!
```

**Problem**: Time step is **constant**, not input-dependent. This breaks selectivity.

### Reference Implementation (CORRECT)

```python
# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_model, dt_rank=None):
        dt_rank = dt_rank or ceil(d_model / 16)

        # Project input to dt space
        self.x_proj = nn.Linear(d_inner, dt_rank + 2*d_state)

        # Project dt_rank to d_inner (learnable)
        self.dt_proj = nn.Linear(dt_rank, d_inner, bias=True)

        # Initialize dt bias
        dt_init_std = dt_rank**-0.5
        dt = torch.exp(
            torch.rand(d_inner) * (math.log(dt_max) - math.log(dt_min))
            + math.log(dt_min)
        )
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        self.dt_bias = nn.Parameter(inv_dt)

    def forward(self, x):
        # Extract dt from input
        x_proj = self.x_proj(x)
        dt, B, C = torch.split(x_proj, [dt_rank, d_state, d_state], dim=-1)

        # Project and activate
        dt = self.dt_proj(dt)
        dt = F.softplus(dt + self.dt_bias)
        dt = dt.clamp(min=dt_min, max=dt_max)

        # Use dt in discretization
        A_discrete = torch.exp(self.A_log * dt)
        B_discrete = B * dt
        ...
```

### Required Changes

1. **Add dt Parameters to Config**:
   ```rust
   // ml/src/mamba/mod.rs:68-107
   pub struct Mamba2Config {
       pub dt_rank: usize,        // NEW: ceil(d_model / 16)
       pub dt_min: f64,           // NEW: 0.001
       pub dt_max: f64,           // NEW: 0.1
       pub dt_init_floor: f64,    // NEW: 1e-4
       ...
   }
   ```

2. **Add dt Projection Layers**:
   ```rust
   pub struct Mamba2SSM {
       pub x_proj: Vec<Linear>,      // NEW: input → (dt_rank + 2*d_state)
       pub dt_proj: Vec<Linear>,     // NEW: dt_rank → d_inner
       pub dt_bias: Vec<Tensor>,     // NEW: learnable bias
       ...
   }
   ```

3. **Initialize dt Parameters**:
   ```rust
   impl Mamba2SSM {
       pub fn new(config: Mamba2Config, device: &Device) -> Result<Self, MLError> {
           let dt_rank = (config.d_model as f64 / 16.0).ceil() as usize;
           let mut x_proj = Vec::new();
           let mut dt_proj = Vec::new();
           let mut dt_bias = Vec::new();

           for i in 0..config.num_layers {
               // x_proj: d_inner → (dt_rank + 2*d_state)
               let x_p = candle_nn::linear(
                   d_inner,
                   dt_rank + 2 * config.d_state,
                   vb.pp(&format!("x_proj_{}", i)),
               )?;
               x_proj.push(x_p);

               // dt_proj: dt_rank → d_inner
               let dt_p = candle_nn::linear(
                   dt_rank,
                   d_inner,
                   vb.pp(&format!("dt_proj_{}", i)),
               )?;
               dt_proj.push(dt_p);

               // dt_bias initialization (log-uniform)
               let dt_init_std = (dt_rank as f64).powf(-0.5);
               let dt_init = Tensor::rand(0.0f32, 1.0f32, (d_inner,), device)?
                   .mul(&Tensor::new(&[(config.dt_max.ln() - config.dt_min.ln()) as f32], device)?)?
                   .add(&Tensor::new(&[config.dt_min.ln() as f32], device)?)?
                   .exp()?;

               // Inverse softplus transformation
               let inv_dt = dt_init.clone().add(&dt_init.neg()?.expm1()?.neg()?.log()?)?;
               dt_bias.push(inv_dt);
           }

           Ok(Self { x_proj, dt_proj, dt_bias, ... })
       }
   }
   ```

4. **Use dt in Forward Pass**:
   ```rust
   fn forward_ssd_layer_with_gradients(
       &mut self,
       input: &Tensor,
       layer_idx: usize,
   ) -> Result<Tensor, MLError> {
       // Project input to get dt, B, C
       let x_proj = self.x_proj[layer_idx].forward(input)?;
       let dt_rank = (self.config.d_model as f64 / 16.0).ceil() as usize;

       let dt = x_proj.narrow(2, 0, dt_rank)?;
       let B = x_proj.narrow(2, dt_rank, self.config.d_state)?;
       let C = x_proj.narrow(2, dt_rank + self.config.d_state, self.config.d_state)?;

       // Project dt and apply softplus
       let dt = self.dt_proj[layer_idx].forward(&dt)?;
       let dt = (dt + &self.dt_bias[layer_idx])?;
       let dt = dt.softplus()?;  // softplus(x) = log(1 + exp(x))
       let dt = dt.clamp(self.config.dt_min, self.config.dt_max)?;

       // Discretize A and B using dt
       let A = &self.state.ssm_states[layer_idx].A;
       let A_discrete = (A * &dt)?.exp()?;  // exp(A_log * dt)
       let B_discrete = (B * &dt)?;

       // Use in SSD algorithm
       let y = ssd_chunk_scan(input, &dt, &A_discrete, &B_discrete, &C)?;
       Ok(y)
   }
   ```

**Estimated Effort**: 3 days
**Blocking**: Model selectivity
**Dependencies**: None (can implement immediately)

---

## Gap 4: D (Skip Connection) Parameter (P0)

### Current Implementation (MISSING)

```rust
// ml/src/mamba/mod.rs:1090
let output = scanned_states.matmul(&C_broadcasted)?;
// ❌ No D * x skip connection!
```

### Reference Implementation (CORRECT)

```python
# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_model):
        self.D = nn.Parameter(torch.ones(d_inner))  # Learnable skip weight

    def forward(self, x):
        # SSM computation
        y_ssm = self.ssm(x, A, B, C)

        # Skip connection with learnable weight
        y = self.D * x + y_ssm

        # Output projection
        output = self.out_proj(y)
        return output
```

### Required Changes

1. **Add D Parameter**:
   ```rust
   // ml/src/mamba/mod.rs:193-211
   pub struct SSMState {
       pub A: Tensor,
       pub B: Tensor,
       pub C: Tensor,
       pub delta: Tensor,
       pub D: Tensor,  // NEW: Skip connection weight
       pub hidden: Tensor,
   }
   ```

2. **Initialize D**:
   ```rust
   impl Mamba2State {
       pub fn zeros(config: &Mamba2Config, device: &Device) -> Result<Self, MLError> {
           for layer_idx in 0..config.num_layers {
               // Initialize D to ones (identity skip connection)
               let D = Tensor::ones((d_inner,), DType::F64, device)?;

               ssm_states.push(SSMState { A, B, C, delta, D, hidden });
           }
           ...
       }
   }
   ```

3. **Apply D in Forward Pass**:
   ```rust
   fn forward_ssd_layer_with_gradients(
       &mut self,
       input: &Tensor,
       layer_idx: usize,
   ) -> Result<Tensor, MLError> {
       let D = &self.state.ssm_states[layer_idx].D;

       // SSM output
       let y_ssm = ssd_chunk_scan(input, &dt, &A_discrete, &B_discrete, &C)?;

       // Skip connection: D * x + y_ssm
       let D_broadcasted = D.unsqueeze(0)?.unsqueeze(0)?;  // [1, 1, d_inner]
       let skip = (input * &D_broadcasted)?;
       let output = (skip + y_ssm)?;

       Ok(output)
   }
   ```

4. **Add D to Optimizer**:
   ```rust
   fn optimizer_step(&mut self) -> Result<(), MLError> {
       for layer_idx in 0..num_layers {
           // Update D parameter (like B, C)
           if let Some(ref D_grad) = self.gradients.get(&format!("D_{}", layer_idx)) {
               let mut D_param = self.state.ssm_states[layer_idx].D.clone();
               self.apply_adam_update(&mut D_param, D_grad, layer_idx, "D", ...)?;
               self.state.ssm_states[layer_idx].D = D_param;
           }
       }
   }
   ```

**Estimated Effort**: 1 day
**Blocking**: Training stability
**Dependencies**: None (can implement immediately)

---

## Gap 5: A Matrix Initialization (P1)

### Current Implementation (WRONG)

```rust
// ml/src/mamba/mod.rs:237-242
let A = Tensor::randn(0.0, 1.0, (config.d_state, config.d_state), device)?;
```

**Problem**: Random initialization, full matrix (not diagonal).

### Reference Implementation (CORRECT)

```python
# state-spaces/mamba/mamba_ssm/modules/mamba2.py
class Mamba2Block:
    def __init__(self, d_state):
        # A_log initialization: log(range(1, d_state+1))
        A_log = torch.log(torch.arange(1, d_state + 1, dtype=torch.float32))
        self.A_log = nn.Parameter(A_log)  # Shape: (d_state,) - DIAGONAL!

    def forward(self, x):
        A = -torch.exp(self.A_log)  # Negative for stability
        ...
```

### Required Changes

1. **Make A Diagonal**:
   ```rust
   // ml/src/mamba/mod.rs:237-242
   // OLD: let A = Tensor::randn(0.0, 1.0, (config.d_state, config.d_state), device)?;

   // NEW: Initialize as diagonal with log(range(1, d_state+1))
   let A_diag: Vec<f64> = (1..=config.d_state)
       .map(|i| (i as f64).ln())
       .collect();
   let A = Tensor::from_vec(A_diag, (config.d_state,), device)?;
   ```

2. **Use Diagonal A in Discretization**:
   ```rust
   fn discretize_ssm_with_gradients(&self, A_log: &Tensor, dt: &Tensor) -> Result<Tensor, MLError> {
       // A is now 1D (diagonal), not 2D matrix
       let A = A_log.neg()?.exp()?;  // Negative exponential for stability

       // Discretize: A_discrete = exp(-A * dt)
       let A_discrete = (A * dt)?.neg()?.exp()?;  // Elementwise

       Ok(A_discrete)
   }
   ```

3. **Update SSMState Structure**:
   ```rust
   pub struct SSMState {
       pub A: Tensor,  // Shape: (d_state,) - DIAGONAL ONLY
       pub B: Tensor,  // Shape: (d_state, d_inner)
       pub C: Tensor,  // Shape: (d_inner, d_state)
       pub D: Tensor,  // Shape: (d_inner,)
       pub delta: Tensor,
       pub hidden: Tensor,
   }
   ```

**Estimated Effort**: 1 day
**Blocking**: Convergence speed
**Dependencies**: None (can implement immediately)

---

## Gap 6: d_state Size (P1)

### Current Implementation (WRONG)

```rust
// ml/src/mamba/mod.rs:139
d_state: 16,  // 8x too small!
```

### Reference Implementations (CORRECT)

- state-spaces/mamba: `d_state = 64-128`
- tommyip/mamba2-minimal: `d_state = 64`
- Hugging Face: `d_state = 128`

### Required Changes

```rust
// ml/src/mamba/mod.rs:134-158
impl Mamba2Config {
    pub fn emergency_safe_defaults() -> Self {
        Self {
            d_state: 64,  // FIXED: Was 16, now 64 (minimum for Mamba-2)
            ...
        }
    }
}
```

**Impact**:
- Larger state → more model capacity
- 4x memory increase (16 → 64)
- Better long-range dependencies

**Estimated Effort**: 1 day (just change constant + retrain)
**Blocking**: Model capacity
**Dependencies**: None

---

## Gap 7-12: See Detailed Implementation Plan

(Remaining gaps documented in `AGENT_228_REFERENCE_IMPLEMENTATIONS.md` sections 7-12)

---

## Implementation Priority Queue

### Wave 229 (This Week - 3 days)
1. **dt Parameter** (3 days, P0)
   - Add `x_proj`, `dt_proj`, `dt_bias`
   - Implement softplus + clamping
   - Use in discretization

2. **D Parameter** (1 day, P0)
   - Add to `SSMState`
   - Initialize to ones
   - Apply skip connection

3. **Convolution Layer** (2 days, P0)
   - Add `Conv1d` to model
   - Apply before SSM
   - Causal padding

### Wave 230 (Next Week - 5 days)
4. **A Matrix Fix** (1 day, P1)
   - Change to diagonal
   - Log initialization
   - Update discretization

5. **d_state Increase** (1 day, P1)
   - Change from 16 to 64
   - Test memory usage

6. **SSD Algorithm** (3 days, P0)
   - Implement chunk-based scan
   - Diagonal/off-diagonal blocks
   - Matrix multiplication approach

### Wave 231 (Week 3 - 5 days)
7. **Input Splitting** (2 days, P2)
   - Split `in_proj → (z, x)`
   - Add SiLU gating
   - Update output

8. **Tensor Core Optimization** (3 days, P1)
   - Use FP16/BF16
   - Align dimensions
   - Profile performance

### Wave 232 (Week 4 - 5 days)
9. **Memory Access Patterns** (3 days, P1)
   - Coalesce memory operations
   - Reduce HBM transfers
   - Profile bandwidth

10. **Integration Testing** (2 days)
    - Test with ES.FUT data
    - Validate gradients
    - Benchmark vs PyTorch

---

## Success Criteria

### Functional Requirements
- ✅ Model trains without errors
- ✅ Gradients flow correctly
- ✅ Validation loss decreases
- ✅ Inference latency <5μs

### Performance Requirements
- ✅ Training speed ≥50% of PyTorch Mamba-2
- ✅ Memory usage ≤2x PyTorch
- ✅ GPU utilization >70%

### Architectural Requirements
- ✅ SSD algorithm implemented
- ✅ All parameters present (dt, D, A, B, C)
- ✅ Convolution layer working
- ✅ Tensor core optimization enabled

---

## Testing Strategy

### Unit Tests
```rust
#[test]
fn test_dt_parameter() {
    // Test dt projection and clamping
    let config = Mamba2Config::default();
    let model = Mamba2SSM::new(config, &Device::Cpu)?;

    let input = Tensor::randn(0.0, 1.0, (1, 10, 64), &Device::Cpu)?;
    let dt = model.compute_dt(&input, 0)?;

    assert!(dt.min()? >= config.dt_min);
    assert!(dt.max()? <= config.dt_max);
}

#[test]
fn test_d_skip_connection() {
    // Test D parameter skip connection
    let model = Mamba2SSM::new(config, &Device::Cpu)?;
    let input = Tensor::ones((1, 10, 64), &Device::Cpu)?;

    let output = model.forward(&input)?;

    // Output should include skip connection
    assert!(output.dims() == input.dims());
}

#[test]
fn test_conv1d_causal() {
    // Test causal convolution (no future leakage)
    let model = Mamba2SSM::new(config, &Device::Cpu)?;
    let input = Tensor::zeros((1, 10, 64), &Device::Cpu)?;
    input.narrow(1, 5, 1)?.fill_(1.0)?;  // Set t=5 to 1

    let output = model.forward(&input)?;

    // Positions t<5 should be zero (no future info)
    assert!(output.narrow(1, 0, 5)?.abs().sum()? < 1e-6);
}
```

### Integration Tests
```rust
#[test]
fn test_e2e_mamba2_training() {
    // Test end-to-end training loop
    let mut model = Mamba2SSM::new(config, &Device::cuda_if_available(0)?)?;
    let train_data = load_es_fut_data()?;

    let history = model.train(&train_data, &val_data, epochs=10).await?;

    // Loss should decrease
    assert!(history.last().unwrap().loss < history[0].loss);
}
```

### Performance Benchmarks
```bash
# Run GPU training benchmark
cargo run --release -p ml --example gpu_training_benchmark

# Compare against PyTorch
python benchmarks/compare_mamba2.py --model foxhunt --baseline pytorch
```

---

## Risk Mitigation

### Risk 1: SSD Algorithm Complexity
- **Probability**: HIGH
- **Impact**: CRITICAL
- **Mitigation**: Start with tommyip/mamba2-minimal (simplest implementation)
- **Fallback**: Use Mamba-1 associative scan temporarily

### Risk 2: Candle Limitations
- **Probability**: MEDIUM
- **Impact**: HIGH
- **Mitigation**: Implement SSD using primitive ops (matmul, elementwise)
- **Fallback**: Request custom CUDA kernel support from Candle team

### Risk 3: Memory Increase
- **Probability**: LOW
- **Impact**: MEDIUM
- **Mitigation**: Profile memory usage, optimize batch size
- **Fallback**: Reduce d_state if GPU OOM

---

**Agent 228 Out** 🎯