Files
foxhunt/crates
jgrusewski fb05631997 perf: capture spectral norm in graph_forward + EMA in graph_ema
Spectral norm (10 weight matrices):
- Moved into graph_forward capture (submit_spectral_norm_ops)
- Runs before cuBLAS forward — graph ensures normalized weights
- All device_ptr() → raw_ptr() (no cudarc events)
- ~10 kernel launches eliminated from per-step ungraphed ops

EMA target update:
- Captured as graph_ema (EMA kernel + f32→bf16 cast + 20 unflatten DtoD)
- tau read from GPU-resident tau_buf (async HtoD before replay)
- ~22 kernel launches (EMA + cast + 20 DtoD) → 1 graph replay

Per-step kernel launches now: 7 graph replays + ~4 ungraphed
(was: 5 graph replays + ~14 ungraphed)

Remaining ungraphed:
- upload_batch_gpu (8 ops — batch ptrs change per step)
- HER relabel (2 ops — donor indices change)
- IQN trunk gradient (2 ops — tau changes)
- IQN→PER DtoD (1 op)
- causal/vaccine (amortized, rare)
- regime_scale + PER update (2 ops)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-02 09:58:47 +02:00
..