[Attention] Overlap sparse MLA indexer with native streams - #2
[Attention] Overlap sparse MLA indexer with native streams#2LucasWilkinson wants to merge 2 commits into
Conversation
Run the DSA sparse-MLA indexer on a shared high-priority CUDA stream so it overlaps the main stream path to attention. Build on the PyTorch 2.13 update and use its native torch.compile stream annotations instead of vLLM side-stream begin/end custom ops. Keep the indexer output as an ordering dependency and join before sparse attention consumes the top-k indices. Also defer the fused indexer Q-scale multiplication so Q RoPE and quantization can issue before the WK GEMM. Tests: .venv/bin/python -m pytest tests/compile/test_side_stream.py -q; focused ruff, torch CUDA API, and mypy 3.12 pre-commit hooks. Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
Make PyTorch 2.13 native stream objects survive vLLM AOT serialization and reload. Keep stream-bearing pieces out of independent piecewise CUDA graph capture so fork and join operations remain balanced within the enclosing full graph. Tests: tests/compile/test_side_stream.py; focused Ruff and mypy; staged-file pre-commit; GLM-5.2-NVFP4 TP4 FULL_AND_PIECEWISE throughput and GSM8K-64. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
LucasWilkinson
changed the title
|
Closing after full lifecycle audit. The native-stream version passes isolated torch.compile graphs but fails vLLM full-model AOT/piecewise execution because PyTorch clears the external-object registry after vLLM can safely restore it. Fixed-slot registration and the global getter monkeypatch are not acceptable. A clean one-op boundary implementation works on the PyTorch 2.13 base, but publishing it here would duplicate vllm-project#47355. The native and scoped-registry prototypes remain preserved locally for an upstream PyTorch reproducer. |
1 participant
Status: blocked — do not merge
This draft tested whether PyTorch 2.13 native stream operations could replace the stream-control custom ops in vllm-project#47355.
The review concerns are valid. In the local corrective revision I removed the global PyTorch code monkeypatch, unconditional external-object slot registration, AOT mapping machinery, duplicate stream join, duplicated AOT test setup, and unrelated FlashAttention type-ignore change. The focused native-stream graph tests still pass, but full-model validation exposed two blockers:
Therefore the native implementation is not currently viable in vLLM without either unsafe private-registry manipulation or upstream PyTorch support for preserving external stream objects through compilation and serialization.
No corrective commit has been pushed because the remaining native patch is nonfunctional end to end. The earlier AOT/GSM8K/throughput results depended on the unsafe registry workaround and must not be treated as validation of a mergeable implementation.
The existing custom-op implementation remains vllm-project#47355. Converting this draft back to that approach would duplicate the open PR, so this draft is retained only as an investigation record pending closure.