From e54b9e5abddfc100d2e030ce9551cfb98a421b80 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 18 Apr 2026 00:09:40 +0200 Subject: [PATCH] =?UTF-8?q?spec:=20Phase=202=20multi-stream=20graph=20para?= =?UTF-8?q?llelism=20=E2=80=94=20designed=20for=20future=20implementation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit forward_child: 4 branch streams fork after h_s2, join before loss aux_child: IQL/IQN/attention on parallel streams Uses CU_STREAM_CAPTURE_MODE_RELAXED with fork-join events as graph edges Existing branch_streams[4] + events in batched_forward.rs ready to activate Target: <40s epochs (from <80s Phase 1) Co-Authored-By: Claude Opus 4.6 (1M context) --- ...17-unified-cublas-training-graph-design.md | 106 +++++++++++++++++- 1 file changed, 103 insertions(+), 3 deletions(-) diff --git a/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md b/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md index 110d90d4d..758867363 100644 --- a/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md +++ b/docs/superpowers/specs/2026-04-17-unified-cublas-training-graph-design.md @@ -300,13 +300,113 @@ Modified files: - Epoch timing: target <80s (from 685s) - Determinism: two consecutive runs must produce identical val_Sharpe_raw -### Performance Target +### Performance Target (Phase 1: Single-stream) -| Metric | Before | After | -|--------|--------|-------| +| Metric | Before | Phase 1 | +|--------|--------|---------| | Epoch time | 685s | <80s | | Per-step aux ops | ~3.5s | ~5ms | | Graph launches/step | 4 | 1 | | Temporal ops | disabled | active | | aux_frequency | 4 (workaround) | 1 (every step) | | Total kernel launches/step | ~71 (per-sample) | ~60 (cuBLAS batched) | + +--- + +## Phase 2: Multi-stream Graph Parallelism (designed now, implemented after Phase 1 validates) + +### Problem + +Single-stream execution wastes H100's 132 SMs. When a 1-thread ISV signal +update runs, 131 SMs sit idle. The branch head GEMMs (direction, magnitude, +order, urgency) are independent but execute sequentially. IQL high/low are +independent but serialized. + +### Architecture + +The child graph architecture enables this naturally. Each child sub-graph +manages its own internal parallelism via fork-join events captured as graph +edges. The parent graph replays the full multi-stream topology in one launch. + +``` +forward_child (internal multi-stream): + main_stream: states → h_s1 → h_s2 → [FORK EVENT] + │ + branch_stream[0]: ───── h_v → v_logits ───┤ + branch_stream[1]: ───── h_b0 → b0_logits ─┤ + branch_stream[2]: ───── h_b1 → b1_logits ─┤ + branch_stream[3]: ───── h_b2 → b2_logits ─┤ + │ + main_stream: ──────── [JOIN EVENT] → loss → grad → backward + +aux_child (internal multi-stream): + stream_a: ── IQL high-tau (forward + backward + adam) ──┐ + stream_b: ── IQL low-tau (forward + backward + adam) ──┤ + ├── [JOIN] → CQL → conf_bwd + stream_c: ── IQN (forward + backward + adam) ──┤ + stream_d: ── Attention (forward + backward + adam) ──┘ +``` + +### CUDA Graph Multi-stream Capture + +Use `CU_STREAM_CAPTURE_MODE_RELAXED` instead of `GLOBAL`. Fork-join events +within the capture are recorded as graph dependency edges: + +```rust +// Inside forward_child capture: +stream.begin_capture(CU_STREAM_CAPTURE_MODE_RELAXED)?; + +// Trunk (main stream) +submit_trunk_forward()?; // states → h_s1 → h_s2 + +// Fork: record event on main stream, wait on branch streams +let fork_event = stream.record_event()?; +for branch_stream in &branch_streams { + branch_stream.wait_event(&fork_event)?; +} + +// Branch GEMMs (parallel on 4 streams) +for (i, branch_stream) in branch_streams.iter().enumerate() { + submit_branch_forward(branch_stream, i)?; // h_s2 → h_bi → logits_bi + let done = branch_stream.record_event()?; + stream.wait_event(&done)?; // Join back to main +} + +// Loss + backward (main stream, after all branches joined) +submit_loss_and_backward()?; + +stream.end_capture()?; // Graph captures the fork-join topology +``` + +The captured graph replays all 4 branches in parallel automatically — the +graph scheduler sees the dependency edges and dispatches to multiple SMs. + +### Expected Performance (Phase 2) + +| Metric | Phase 1 | Phase 2 | +|--------|---------|---------| +| Forward pass | ~3ms (sequential branches) | ~1.2ms (parallel branches) | +| Aux ops | ~5ms (sequential IQL+IQN+attn) | ~2ms (parallel trainers) | +| SM utilization | ~40% (sequential) | ~85% (parallel) | +| **Epoch time** | **<80s** | **<40s** | + +### What Already Exists + +`batched_forward.rs` already has the infrastructure: +- `branch_streams: [Arc; 4]` — 4 dedicated branch streams +- `branch_workspace_ptrs: [u64; 4]` — per-branch cuBLAS workspace (no contention) +- `trunk_done_event: CudaEvent` — fork event after trunk +- `branch_done_events: [CudaEvent; 4]` — join events per branch + +These are currently unused during graph capture (single-stream mode). Phase 2 +activates them inside the `forward_child` capture with RELAXED mode. + +### Implementation Order + +1. Phase 1: single-stream child graphs (this spec) — validate correctness +2. Phase 2: multi-stream forward_child — parallelize 4 branch heads +3. Phase 2b: multi-stream aux_child — parallelize IQL/IQN/attention +4. Phase 2c: double-buffered experience collection — overlap with training + +Each phase is additive — Phase 1 child graphs don't change, they just gain +internal parallelism. The parent graph composition is unchanged.