foxhunt/docs/superpowers/specs/2026-04-18-multi-stream-phase2-design.md

# Phase 2: Multi-Stream Graph Parallelism

## Goal

Reduce per-step GPU time from 4.0s to ~2.0s by parallelizing independent operations
within the aux_child and forward_child graphs. Target: ~350s epochs (from ~718s).

## Measured Baseline (2026-04-18, H100 PCIe, batch=16384)

```
Per-step graph replay breakdown (last batch of epoch 1):
  spectral =   0.2ms  (<1%)
  forward  = 130.3ms  (31%)   ← branches already multi-stream, backward is NOT
  ddqn     =   1.0ms  (<1%)
  aux      = 295.5ms  (69%)   ← 4 independent trainers running serial
  adam     = TBD (event fixed, awaiting measurement)
  total    = ~427ms measured + adam
```

Forward branches already use 4 streams with fork-join events (`batched_forward.rs:503-567`).
Backward branches are sequential (`batched_backward.rs` — no branch_stream infrastructure).

## Constraints (from NVIDIA research)

1. **SM saturation**: At B=16384, IQL GEMMs (M=128, N=16384, K=128) produce ~128 CTAs
   filling 97% of 132 SMs. Multi-stream GEMM-vs-GEMM parallelism provides minimal benefit.
2. **Element-wise overlap**: Small kernels (bias_add, SiLU, reduce — 10-20µs) leave SMs idle.
   A concurrent GEMM from another trainer CAN fill those gaps.
3. **Workspace isolation**: Each concurrent cublasLtMatmul stream needs its own workspace (32MB).
   Same `cublasLtHandle_t` is safe (stateless per-call).
4. **Grouped GEMM**: `cublasLtMatmulGroupedBatched` packs multiple independent GEMMs with
   different M/N/K into ONE kernel launch — eliminates inter-kernel gaps.

## Design

### 2A. Aux Child: Hybrid Parallelism

**IQL high + IQL low** have identical GEMM shapes (same architecture, different weights).
Use grouped GEMM to fuse both trainers' forward+backward into single kernel launches.
This halves IQL kernel count (16 GEMMs → 8 grouped GEMMs).

**IQN + Attention** run on 2 separate streams, forked from the main aux stream after
IQL completes (IQL must run first because IQL advantage weights are read by CQL later).

```
aux_child capture (RELAXED mode):
  main_stream: HER relabel → EMA target → IQL grouped GEMM (fwd+bwd+adam) → [FORK]
                                                                                │
  iqn_stream:  ─── IQN (fwd+bwd+adam) ──────────────────────────────────────────┤
  attn_stream: ─── Attention (fwd+bwd+adam) ────────────────────────────────────┤
                                                                                │
  main_stream: ─── [JOIN] → CQL gradient → ensemble → recursive_conf_bwd → regime_scale
```

**Workspace**: 2 extra 32MB workspaces (IQN stream + Attention stream). Total: 64MB additional.
IQL grouped GEMM uses the existing shared workspace (sequential on main stream).

**Streams**: 2 new `CudaStream` objects, forked from the main stream context.
Pre-allocated events: `iql_done_event`, `iqn_done_event`, `attn_done_event`.

**Expected speedup on aux_child**:
- IQL: ~50% faster (grouped GEMM eliminates 8 kernel launch gaps)
- IQN + Attention: overlap with each other (~40% faster on combined time)
- CQL + ensemble + conf_bwd: unchanged (serial, depends on IQL/IQN/Attention outputs)
- Overall aux: ~295ms → ~180ms

### 2B. Forward Child: Backward Branch Parallelism

The forward pass already uses 4 branch streams. The backward pass does NOT.
Add the same fork-join pattern to the backward: after the trunk backward (serial),
fork 4 branch backward chains to their respective streams.

```
backward in forward_child:
  main_stream: d_logits_clamp → [4 branch dW + dX GEMMs sequential] → trunk backward

  ↓ becomes:

  main_stream: d_logits_clamp → [FORK]
                                    │
  branch_stream[0]: ── branch0 dW + dX ──┤
  branch_stream[1]: ── branch1 dW + dX ──┤
  branch_stream[2]: ── branch2 dW + dX ──┤
  branch_stream[3]: ── branch3 dW + dX ──┤
                                    │
  main_stream: ──── [JOIN] → trunk backward (d_h_s2 reduction → d_h_s1 → d_states)
```

**Reuses existing infrastructure**: `branch_streams[4]`, `trunk_done_event`, `branch_done_events[4]`
from `batched_forward.rs`. The backward just needs to use them too.

**Expected speedup on forward_child**: ~130ms → ~100ms (backward branches ~30ms of the 130ms)

### 2C. Grouped GEMM for IQL

Use `cublasLtMatmulGroupedBatched` (cuBLAS 12.5+) to fuse IQL high-tau and IQL low-tau
GEMMs. Both trainers have identical shapes:
- Forward: 3 GEMMs (h1, h2, v_out) × 2 trainers = 6 GEMMs → 3 grouped GEMMs
- Backward: 5 GEMMs (dW3, dh2, dW2, dh1, dW1) × 2 trainers = 10 GEMMs → 5 grouped GEMMs
- Adam: 1 kernel × 2 trainers = 2 kernels → 1 grouped kernel (or keep separate)

Each grouped GEMM takes arrays of pointers (A_ptrs, B_ptrs, C_ptrs) and sizes.
The cuBLAS runtime dispatches both GEMMs in a single kernel launch.

**Fallback**: If `cublasLtMatmulGroupedBatched` is not available (requires cuBLAS 12.5+),
fall back to sequential execution (current behavior). Check at runtime.

## VRAM Budget

| Component | Current | Phase 2 |
|-----------|---------|---------|
| IQN workspace | shared (0 extra) | 32MB dedicated |
| Attention workspace | shared (0 extra) | 32MB dedicated |
| Branch streams | already allocated | 0 extra |
| Aux streams (2) | 0 | ~1KB (stream objects) |
| Aux events (3) | 0 | ~1KB |
| **Total extra** | | **64MB** |

On H100 with 80GB VRAM, 64MB is <0.1%.

## Files Modified

- `crates/ml/src/cuda_pipeline/batched_backward.rs` — add branch stream fork-join
- `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs` — grouped GEMM support
- `crates/ml/src/cuda_pipeline/shared_cublas_handle.rs` — aux workspace allocation
- `crates/ml/src/trainers/dqn/fused_training.rs` — aux_child multi-stream capture, RELAXED mode

## Success Criteria

- `GPU child graph breakdown` shows aux < 200ms (from 295ms)
- `GPU child graph breakdown` shows forward < 110ms (from 130ms)
- Epoch time < 400s (from ~718s)
- All 899 unit tests pass
- No cuStreamSynchronize in training loop
- Zero atomicAdd on gradient paths (maintained)