# GPU Regime Stratification + H100 Training Hang Fix

**Status**: Design approved  
**Priority**: Critical — blocks H100 production training  
**Date**: 2026-04-05

---

## Problem Statement

Three issues block production H100 training:

1. **H100 training hang**: Process freezes after CUDA graph capture with `gpu_n_episodes=4096`. Root cause: `cuBLAS::set_stream()` and event synchronization calls execute during CUDA graph capture but are NOT recorded in the graph. On replay, the cuBLAS handle is in an undefined stream state → deadlock.

2. **Regime stratification bug**: `train_walk_forward()` at `trainer/mod.rs:570` calls `generate_folds()` (non-stratified) instead of `generate_folds_stratified()`. Walk-forward folds are not regime-balanced.

3. **CPU regime computation**: Regime classification and distribution computation run on CPU. For 1.1M bars this is fast, but the architecture should be GPU-native for scaling to 10M+ bars with real-time MBP-10 feeds.

---

## Fix 1: Dual-Stream CUDA Graph Architecture

### Root Cause

In `gpu_dqn_trainer.rs`, `submit_forward_ops()` is called INSIDE `stream.begin_capture()` / `stream.end_capture()`. This function contains:

- **Line 4702**: `cublas.set_stream(&self.double_dqn_stream)` — executes but NOT captured
- **Line 4712**: `cublas.set_stream(&self.stream)` — executes but NOT captured  
- **Line 4683**: `self.pass1_event.record(&self.stream)` — events disallowed in capture
- **Line 4687**: `self.double_dqn_stream.wait(&self.pass1_event)` — not captured

On graph replay, the cuBLAS handle's internal stream pointer is stale. The multi-stream event synchronization is missing. Result: deadlock.

### Design

Split the single `graph_forward` into two independent graphs:

**`graph_forward_main`**: Captures Pass 1 (online forward on states) + Pass 2 (target forward on next_states) on the main stream. No stream switching inside.

**`graph_forward_ddqn`**: Captures Pass 3 (online forward on next_states for Double DQN action selection) on `double_dqn_stream`. Uses a separate cuBLAS handle already bound to that stream.

**Replay sequence** (ungraphed orchestration):
```
1. graph_forward_main.launch()           // main stream
2. pass1_event.record(main_stream)       // ungraphed event
3. double_dqn_stream.wait(pass1_event)   // ungraphed sync
4. graph_forward_ddqn.launch()           // double_dqn_stream
5. pass3_event.record(double_dqn_stream) // ungraphed event
6. main_stream.wait(pass3_event)         // ungraphed sync — join before loss kernels
7. Loss + gradient kernels (main stream, part of graph_forward_main or separate)
```

Pass 2 and Pass 3 execute concurrently on separate SM partitions (H100 has 132 SMs, each GEMM uses ~30-40 SMs).

### Changes

| File | Change |
|------|--------|
| `gpu_dqn_trainer.rs` | Split `submit_forward_ops()` into `submit_pass1_pass2()` and `submit_pass3_ddqn()` |
| `gpu_dqn_trainer.rs` | `capture_training_graphs()` captures two forward graphs instead of one |
| `gpu_dqn_trainer.rs` | `replay_forward()` orchestrates both graphs with ungraphed event sync |
| `gpu_dqn_trainer.rs` | Loss/gradient kernels stay on main stream (after join) |
| `gpu_dqn_trainer.rs` | Remove `set_stream()` calls from inside any captured function |
| `gpu_dqn_trainer.rs` | Create second `CublasForward` instance for `double_dqn_stream` (currently shares one handle via `set_stream`) |
| `batched_forward.rs` | No changes — `CublasForward::new()` already takes a stream and creates an independent handle |

### Constraints

- No `cuBLAS::set_stream()` inside ANY function called during graph capture
- No CUDA event operations inside graph capture
- Each graph uses exactly ONE stream with ONE cuBLAS handle
- The `double_dqn_stream` cuBLAS handle must be a SEPARATE instance (not shared)
- Event synchronization happens ONLY in ungraphed replay orchestration code

### Validation

- `test_no_nan_after_graph_capture`: 200/200 passes (existing)
- `test_no_hang_single_epoch`: passes with `gpu_n_episodes=4096` (new parameter override)
- H100 Argo workflow completes epoch 1 within 60 seconds

---

## Fix 2: Regime Stratification Bug

### Root Cause

`trainer/mod.rs` line 570:
```rust
let folds = wf_config.generate_folds(features.len());
```

Should be:
```rust
let folds = wf_config.generate_folds_stratified(training_data);
```

The non-stratified `generate_folds()` ignores regime distribution entirely. Validation sets may be dominated by a single regime (e.g., 90% trending), making walk-forward evaluation unreliable.

### Change

One line in `trainer/mod.rs:570`. The `generate_folds_stratified()` function already exists in `gpu_walk_forward.rs` and is tested.

### Validation

- Existing walk-forward tests pass
- Log output shows regime distribution per fold with max deviation < 10%

---

## Fix 3: GPU Regime Classification + Prefix Sum

### Current State

`gpu_walk_forward.rs:322-338` classifies regimes on CPU by reading features[40] (ADX) and features[41] (CUSUM) for each bar. The boundary adjustment search at lines 240-264 recomputes regime distribution per shift candidate at O(val_size) per candidate.

### Design

**CUDA kernel: `regime_classify_kernel`**

```cuda
// One thread per bar. Reads features[40] and features[41], writes u8 regime label.
// Input:  features [total_bars, 42] f64 (host) or bf16 (GPU)
// Output: regimes  [total_bars] u8
extern "C" __global__ void regime_classify_kernel(
    const double* features,  // [total_bars × 42]
    unsigned char* regimes,  // [total_bars]
    int total_bars,
    int feat_stride           // 42
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_bars) return;
    double adx   = features[i * feat_stride + 40];
    double cusum = features[i * feat_stride + 41];
    unsigned char regime;
    if (adx > 0.25)            regime = 0;  // Trending
    else if (fabs(cusum) > 0.7) regime = 2;  // Volatile
    else                        regime = 1;  // Ranging
    regimes[i] = regime;
}
```

**CUDA kernel: `regime_prefix_sum_kernel`**

Compute per-regime cumulative counts using parallel prefix sum (scan). Output: three arrays of length `total_bars`, where `prefix[r][i]` = count of regime `r` in bars `[0..i)`.

With prefix sums, any range query `regime_count(r, start, end) = prefix[r][end] - prefix[r][start]` is O(1). The boundary adjustment search becomes O(num_shifts) instead of O(num_shifts × val_size).

### Data Flow

```
fxcache features [1.1M, 42] f64
    ↓ upload to GPU as f64
regime_classify_kernel → regimes [1.1M] u8 (GPU)
    ↓
regime_prefix_sum_kernel → prefix_counts [3, 1.1M] u32 (GPU)
    ↓ DtoH transfer (3 × 4.4MB = 13.2MB)
CPU boundary search using O(1) range queries
    ↓
fold_ranges: Vec<FoldRange>
```

### Changes

| File | Change |
|------|--------|
| `gpu_walk_forward.rs` | Add `regime_classify_kernel` and `regime_prefix_sum_kernel` CUDA source |
| `gpu_walk_forward.rs` | New `classify_regimes_gpu()` function: upload features, run kernel, return GPU buffer |
| `gpu_walk_forward.rs` | New `compute_regime_prefix_sums_gpu()`: run prefix sum, transfer to CPU |
| `gpu_walk_forward.rs` | Update `generate_folds_stratified()` to use GPU prefix sums for O(1) range queries |
| `gpu_walk_forward.rs` | `regime_distribution_from_labels()` replaced by prefix sum lookup |

### Performance

| Operation | CPU (current) | GPU (new) |
|-----------|--------------|-----------|
| Classify 1.1M bars | 2ms | 0.05ms |
| Distribution per range query | O(val_size) = 0.5ms | O(1) = 0.001ms |
| Boundary search (4 folds × 500 shifts) | 1000ms | 2ms |
| Prefix sum computation | N/A | 0.5ms |
| DtoH transfer | N/A | 0.1ms |
| **Total** | **~1000ms** | **~3ms** |

### Validation

- `generate_folds_stratified()` produces identical fold ranges as CPU version on same input
- Unit test: prefix sum correctness on known regime sequence
- Walk-forward tests pass with GPU regime path

---

## Files Modified (Summary)

| File | Changes |
|------|---------|
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Split forward graphs, fix stream binding, dual-graph replay |
| `crates/ml/src/cuda_pipeline/gpu_walk_forward.rs` | GPU regime kernels, prefix sum, updated stratification |
| `crates/ml/src/trainers/dqn/trainer/mod.rs` | Fix `generate_folds()` → `generate_folds_stratified()` |
| `crates/ml/src/trainers/dqn/smoke_tests/regression.rs` | Add gpu_n_episodes=4096 hang regression test |

## Success Criteria

1. H100 training completes 200 epochs without hang
2. Zero NaN (existing 200/200 regression holds)
3. Walk-forward folds are regime-stratified (max deviation < 10%)
4. 19/19 smoke tests pass
5. Epoch time on H100 ≤ 30s (1.1M bars, batch=8192)
