Skip to content

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
thu-ml:mainfrom
shanai13:fix/stream-correct-kernel-launches
Open

Launch sm90 attention and fused quant kernels on the current torch stream (CUDA-graph capture correctness)#371
shanai13 wants to merge 1 commit into
thu-ml:mainfrom
shanai13:fix/stream-correct-kernel-launches

Conversation

@shanai13

Copy link
Copy Markdown

All kernel launches in csrc/qattn/qk_int_sv_f8_cuda_sm90.cu (2 sites) and csrc/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 manual torch.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 the ATen/cuda/CUDAContext.h include). 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).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant