Launch sm90 attention and fused quant kernels on the current torch stream (CUDA-graph capture correctness)#371
Open
shanai13 wants to merge 1 commit into
Open
Conversation
…am (CUDA-graph-capture correctness)
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
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.
All kernel launches in
csrc/qattn/qk_int_sv_f8_cuda_sm90.cu(2 sites) andcsrc/fused/fused.cu(8 sites) use<<<grid, block[, smem]>>>with no stream argument, i.e. the legacy default stream, rather than PyTorch's current stream.In ordinary eager use this is invisible (PyTorch's default current stream aliases the legacy stream). But under CUDA Graph capture — e.g.
torch.compile(mode="reduce-overhead"), or a manualtorch.cuda.graph()around a transformer that uses sageattention — the capture runs on a non-default stream, so these launches escape the capture: the graph records everything except the attention/quant kernels, and replays return stale output buffers (silent corruption, no error). In our Wan2.1-1.3B DiT the replayed forward was wrong by ~80% relative magnitude and appeared to run 2.3× faster, because the attention work was simply missing. The same defect would order incorrectly with any multi-stream pipeline.Fix: take
at::cuda::getCurrentCUDAStream()in the dispatch functions and pass it as the launch stream (plus theATen/cuda/CUDAContext.hinclude). After this change, a manual CUDA-graph capture of a Wan2.1-1.3B forward with sage self-attention replays bit-identically to the direct call (verified at 5 timesteps on an H100 / sm90, with no measurable performance change in normal eager use).The sm80/sm89 qattn files have the same pattern and could be fixed identically; this PR touches only the paths we could test on hardware (sm90 + fused).