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

Phase 2: Multi-Stream Graph Parallelism Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Reduce per-step GPU time from 4.0s to ~2.0s by parallelizing aux_child (295ms, 69%) and forward_child backward branches (part of 130ms).

Architecture: Aux trainers (IQN, Attention) run on 2 dedicated streams forked from main. IQL high+low stay sequential (grouped GEMM future optimization). Forward backward branches use existing branch_streams infrastructure. Capture mode switches to RELAXED for multi-stream graph capture.

Tech Stack: CUDA 13, cudarc (vendored), cublasLtMatmul, CU_STREAM_CAPTURE_MODE_RELAXED, cuEventRecord fork-join

---

Task 1: Aux Stream Infrastructure — Allocate 2 streams + workspaces + events

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs (FusedTrainingCtx struct + constructor)

• Modify: crates/ml/src/cuda_pipeline/shared_cublas_handle.rs (workspace allocation helper)

• Step 1: Add aux stream fields to FusedTrainingCtx

In fused_training.rs, add to the struct after line ~272 (gpu_attention):

// ── Phase 2: Aux multi-stream ──
/// Dedicated stream for IQN training (parallel with Attention).
iqn_stream: Option<Arc<CudaStream>>,
/// Dedicated stream for Attention training (parallel with IQN).
attn_stream: Option<Arc<CudaStream>>,
/// cublasLt workspace for IQN stream (32MB, separate from shared handle).
iqn_workspace: Option<CudaSlice<u8>>,
/// cublasLt workspace for Attention stream (32MB, separate from shared handle).
attn_workspace: Option<CudaSlice<u8>>,
/// Fork event: IQL done on main stream, IQN/Attention can start.
iql_done_event: Option<CudaEvent>,
/// Join event: IQN done on iqn_stream.
iqn_done_event: Option<CudaEvent>,
/// Join event: Attention done on attn_stream.
attn_done_event: Option<CudaEvent>,

• Step 2: Allocate in constructor

After IQL/IQN/Attention initialization, allocate the streams and workspaces:

let iqn_stream = Arc::new(stream.context().new_stream()
    .map_err(|e| anyhow::anyhow!("iqn_stream: {e}"))?);
let attn_stream = Arc::new(stream.context().new_stream()
    .map_err(|e| anyhow::anyhow!("attn_stream: {e}"))?);
let iqn_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)
    .map_err(|e| anyhow::anyhow!("iqn workspace: {e}"))?;
let attn_workspace = stream.alloc_zeros::<u8>(32 * 1024 * 1024)
    .map_err(|e| anyhow::anyhow!("attn workspace: {e}"))?;
let iql_done_event = stream.context().new_event(None)
    .map_err(|e| anyhow::anyhow!("iql_done_event: {e}"))?;
let iqn_done_event = stream.context().new_event(None)
    .map_err(|e| anyhow::anyhow!("iqn_done_event: {e}"))?;
let attn_done_event = stream.context().new_event(None)
    .map_err(|e| anyhow::anyhow!("attn_done_event: {e}"))?;

• Step 3: Initialize fields in struct initializer and add to reset_for_fold

• Step 4: Build and test

Run: SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib

• Step 5: Commit

git commit -m "feat: allocate aux streams, workspaces, events for Phase 2 multi-stream"

---

Task 2: IQN/Attention Accept Per-Call Workspace + Stream Override

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
• Modify: crates/ml/src/cuda_pipeline/gpu_attention.rs

Currently IQN and Attention use the shared handle's workspace and stream for all cublasLtMatmul calls. They need to accept an optional override workspace pointer and stream reference so they can run on dedicated streams with dedicated workspaces.

• Step 1: Add workspace/stream override to IQN's execute_training_pipeline

Add optional parameters to execute_training_pipeline:

pub fn execute_training_pipeline(
    &mut self,
    // ... existing params ...
    override_stream: Option<&Arc<CudaStream>>,
    override_workspace: Option<(u64, usize)>,  // (ptr, size)
) -> Result<(), MLError> {
    let stream = override_stream.unwrap_or(&self.stream);
    let (ws_ptr, ws_size) = override_workspace
        .unwrap_or((self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size));
    // Use `stream` and `ws_ptr/ws_size` for all cublasLtMatmul calls in this function

• Step 2: Same for Attention forward/backward/adam_step

Add override_stream and override_workspace parameters to forward(), backward(), adam_step().

• Step 3: Update all call sites to pass None, None (no behavior change)

All existing callers in submit_aux_ops and elsewhere pass None, None.

• Step 4: Build and test — behavior unchanged

• Step 5: Commit

git commit -m "feat: IQN/Attention accept workspace+stream override for multi-stream"

---

Task 3: Fork-Join in submit_aux_ops

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs (submit_aux_ops)

• Step 1: Fork IQN and Attention to dedicated streams after IQL

In submit_aux_ops, after the IQL pipeline (gather → train → advantages → support → branch_scales → gap → epsilon), fork IQN and Attention to their dedicated streams:

// IQL pipeline runs on main stream (sequential) ...
// ... existing IQL code ...

// Fork: IQN and Attention run in parallel on dedicated streams.
// Record IQL completion on main stream.
if let Some(ref event) = self.iql_done_event {
    event.record(&self.stream)
        .map_err(|e| anyhow::anyhow!("iql_done record: {e}"))?;
}

// IQN on iqn_stream
if let (Some(ref mut iqn), Some(ref iqn_stream), Some(ref iqn_ws), Some(ref iql_event)) =
    (&mut self.gpu_iqn, &self.iqn_stream, &self.iqn_workspace, &self.iql_done_event)
{
    iqn_stream.wait(iql_event)
        .map_err(|e| anyhow::anyhow!("iqn wait iql: {e}"))?;
    // Execute IQN training pipeline on iqn_stream with iqn workspace
    iqn.execute_training_pipeline(
        /* ... existing args ... */
        Some(iqn_stream),
        Some((iqn_ws.raw_ptr(), 32 * 1024 * 1024)),
    ).map_err(|e| anyhow::anyhow!("IQN parallel: {e}"))?;
    // Record IQN completion
    if let Some(ref event) = self.iqn_done_event {
        event.record(iqn_stream)
            .map_err(|e| anyhow::anyhow!("iqn_done record: {e}"))?;
    }
}

// Attention on attn_stream (same pattern)
if let (Some(ref mut attn), Some(ref attn_stream), Some(ref attn_ws), Some(ref iql_event)) =
    (&mut self.gpu_attention, &self.attn_stream, &self.attn_workspace, &self.iql_done_event)
{
    attn_stream.wait(iql_event)
        .map_err(|e| anyhow::anyhow!("attn wait iql: {e}"))?;
    self.trainer.apply_attention_forward_on_stream(attn, self.batch_size, attn_stream, attn_ws)
        .map_err(|e| anyhow::anyhow!("Attention parallel fwd: {e}"))?;
    // ... backward + adam on attn_stream ...
    if let Some(ref event) = self.attn_done_event {
        event.record(attn_stream)
            .map_err(|e| anyhow::anyhow!("attn_done record: {e}"))?;
    }
}

// Join: main stream waits for IQN and Attention to complete
if let Some(ref event) = self.iqn_done_event {
    self.stream.wait(event)
        .map_err(|e| anyhow::anyhow!("main wait iqn: {e}"))?;
}
if let Some(ref event) = self.attn_done_event {
    self.stream.wait(event)
        .map_err(|e| anyhow::anyhow!("main wait attn: {e}"))?;
}

// CQL + ensemble + recursive_confidence_bwd run on main stream after join

• Step 2: Move IQN code out of the sequential section

Remove the existing IQN training from the sequential section (currently after IQL).

• Step 3: Move Attention code out of the sequential section

Remove the existing Attention forward+backward+adam from the sequential section.

• Step 4: Build and test

• Step 5: Commit

git commit -m "feat: fork IQN+Attention to parallel streams in submit_aux_ops"

---

Task 4: Switch Capture Mode to RELAXED

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs (capture_child_graph)

• Step 1: Change capture mode from GLOBAL to RELAXED

In capture_child_graph, change:

// Before:
let begin_result = self.stream.begin_capture(
    cuda_sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_GLOBAL,
);

// After:
let begin_result = self.stream.begin_capture(
    cuda_sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED,
);

RELAXED mode is required for multi-stream capture — it allows other streams to join the capture graph via cuStreamWaitEvent. GLOBAL mode would reject non-captured streams.

• Step 2: Build and test

• Step 3: Deploy to H100 and verify child graph breakdown shows aux < 200ms

Run: ./scripts/argo-train.sh dqn --baseline --epochs 2

Check logs for: GPU child graph breakdown: spectral=X forward=Y ddqn=Z aux=W adam=V

Expected: aux drops from ~295ms to ~180ms.

• Step 4: Commit

git commit -m "perf: switch to RELAXED capture mode — enables multi-stream graph capture"

---

Task 5: Forward Backward Branch Parallelism

Files:

• Modify: crates/ml/src/cuda_pipeline/batched_backward.rs (backward_full)

The forward pass already uses branch_streams[4] with fork-join. The backward needs the same pattern. However, branches share scratch buffers (scratch_d_glu_value, scratch_d_glu_gate) that prevent parallel execution.

• Step 1: Add per-branch GLU scratch buffers to CublasBackwardSet

// In CublasBackwardSet struct:
/// Per-branch GLU backward scratch (Phase 2: parallel branch backward).
/// 4 × [B, AH] for d_value and 4 × [B, AH] for d_gate_pre.
branch_d_glu_value: [CudaSlice<f32>; 4],
branch_d_glu_gate: [CudaSlice<f32>; 4],

Allocate in constructor: 4 × 2 × B × AH × 4 bytes = 4 × 2 × 16384 × 32 × 4 = 16MB.

• Step 2: Add branch_streams + events to CublasBackwardSet

Reuse the same branch_streams and events from CublasForward (passed at construction or shared via Arc).

• Step 3: Fork-join in backward_full branch loop

Replace the sequential for d in 0..4 loop with fork-join:

// Record pre-branch state on main stream
self.trunk_done_event.record(stream)?;

for d in 0..4 {
    let bs = &self.branch_streams[d];
    bs.wait(&self.trunk_done_event)?;
    
    // Run branch d backward on branch_stream[d]
    // Use per-branch scratch: branch_d_glu_value[d], branch_d_glu_gate[d]
    self.backward_branch_on_stream(bs, d, /* ... args with per-branch scratch ... */)?;
    
    self.branch_done_events[d].record(bs)?;
    stream.wait(&self.branch_done_events[d])?;
}

// Trunk backward runs on main stream after all branches complete

• Step 4: Extract single-branch backward into backward_branch_on_stream

Move the body of the for d in 0..4 loop into a method that takes a stream parameter.

• Step 5: Build and test

• Step 6: Commit

git commit -m "perf: parallel backward branches using fork-join on branch_streams"

---

Task 6: Fix Adam Child Event Timing

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs

This was already done in commit 8ebb65a94 — added adam_child_end event. Verify the output shows a real value instead of -1.0ms.

• Step 1: Deploy and check logs

Expected: GPU child graph breakdown: spectral=X forward=Y ddqn=Z aux=W adam=V with V > 0.

• Step 2: If adam is unexpectedly large (>100ms), investigate

The adam_child contains: mamba2_backward (12K params), step_mamba2_adam (saxpy), apply_pruning_mask, compute_grad_norm (2 kernels), launch_adam_update (660K params), unflatten (24 DtoD or zero-copy skip), update_isv_signals. Total should be <50ms.

---

Task 7: Verify and Profile

• Step 1: Deploy with all changes

./scripts/argo-train.sh dqn --baseline --epochs 5 --watch

• Step 2: Check child graph breakdown

Expected:

GPU child graph breakdown:
  spectral =   ~0.2ms
  forward  =  ~100ms  (backward branches parallel)
  ddqn     =    ~1ms
  aux      =  ~180ms  (IQN + Attention parallel)
  adam     =   ~30ms
  total    =  ~310ms  (down from ~430ms)

• Step 3: Check epoch time

Expected: ~310ms/step × 178 steps = ~55s training + 3.5s step0 + 2s experience + 1s validation = ~62s per epoch (down from ~718s).

• Step 4: Commit final timing results to spec

Update docs/superpowers/specs/2026-04-18-multi-stream-phase2-design.md with measured results.