perf(cuda-graphs): capture-clean decoder + CUDA-graph AR forward (stacked on #3)#4
Open
trvachov wants to merge 1 commit into
Open
Conversation
Add a capture-clean decoder with explicit is_causal and a per-shape CUDA-graph cache for batched AR forwards. Preserve eager/graph layout parity, bind capture and workers to the selected device, bound independent graph pools with an LRU and one-worker enforcement, make GPU lock release exception-safe, guard unsupported modes, and copy public logits to owned FP32 storage. Add checkpoint-backed CUDA-graph parity, device, guard, lock, ownership, and memory regression tests. Stacked on trvachov/batched-ar-decode-tested; this commit depends on that batched-AR change and must merge after it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds CUDA-graph acceleration for the AR decoder forward (~2.25× on the forward), plus the capture-clean decoder change that makes it possible.
decoder.py): passtgt_is_causal/memory_is_causalexplicitly sonn.TransformerDecoderskips_detect_is_causal_mask— a per-forward host sync (bool((mask==...).all())) that both added latency and made the forward uncapturable by CUDA graphs. Numerically identical.sampler.py, opt-inREASYN_CUDAGRAPH): captures one graph per bucketed(batch, seq)shape and replays it; static token/gather buffers updated per call. 2.25× on the AR forward, output bit-identical to eager.This PR is based on
trvachov/batched-ar-decode-tested(PR #3), notreasyn_v2, because CUDA graphs are built on the batched/bucketed AR forward. Its diff shows only the decoder + CUDA-graph changes. Merge #3 first, then this. (Independent of #2.)Correctness / safety (from adversarial review + fixes)
REASYN_CUDAGRAPH=0vs1is numerically equivalent (verified over many varied token/gather inputs under identical RNG).self.device; workers callset_device(cuda:<gpu_id>).num_workers_per_gpu==1+ batched+bucket+bf16 + a CUDA model device..to(float32, copy=True)output.Tests
tests/test_cuda_graphs.py— env-off/env-on parity, identical-RNG sampling, device/guard/lock checks, bounded memory. 19 passed on A100 (this suite + the inherited batched-AR suite).Toggle
REASYN_CUDAGRAPH(opt-in, default off; requires batching+bucketing+bf16 + CUDA).Testing environment⚠️
Results were produced on Python 3.12 / PyTorch 2.11 / CUDA 13.2 — not the repo-pinned Python 3.10 / PyTorch 2.7 / CUDA 11.8. CUDA-graph capture and SDPA kernel selection are especially version-sensitive; re-confirm on the pinned stack / CI before merge.
Independently reviewed, fixed, and re-verified (rebased onto the fixed batched-AR; both test suites green) before opening.