Files
foxhunt/crates
jgrusewski 2542fbabd4 perf: Phase 1 mega-graph — capture CQL + C51 clip + pruning in graph_adam
Moved 3 per-step operations into CUDA Graph capture:

1. CQL gradient (submit_cql_ops): CQL logit grad kernel + f32→bf16 cast
   + cuBLAS backward into cql_grad_scratch + grad_norm + clipped SAXPY.
   Full cuBLAS backward is graph-capturable. Budget fractions baked as
   literals (all auxiliaries always active: CQL=25%, C51=60%).

2. C51 gradient clip (submit_c51_clip_ops): grad_norm + finalize + clip.
   Budget 60% baked at capture time.

3. Pruning mask (submit_pruning_mask_ops): element-wise grad *= mask.

graph_adam now contains: CQL backward + C51 clip + pruning + grad_norm
+ Adam + unflatten. Eliminates ~10 ungraphed kernel launches per step.

Per-step kernel launches reduced from ~38 to ~28.

Remaining ungraphed (Phase 2 targets):
- spectral_norm (~3 launches)
- EMA (~1 launch)
- attention fwd+bwd+adam (~6 launches)
- IQL train_value_step (~5 launches)
- IQN train + trunk grad (~10 launches)
- HER relabel (~2 launches)
- causal (1/100 steps)
- vaccine (1/10 steps)

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