[inference] Use Triton MLA latent KV append during CUDA Graph capture
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
The MLA latent KV-cache append path currently uses the same Triton path whenever Triton is available. However, the append path has different performance characteristics in eager execution and CUDA Graph capture.
On an A100, the Triton path is slower than the existing PyTorch indexed assignment in eager execution, but is significantly faster during CUDA Graph replay.
This issue proposes a capture-aware dispatch:
- use the Triton MLA append kernel during CUDA Graph capture;
- preserve the existing PyTorch indexed-assignment fallback for eager execution;
- keep the normal KV-cache append path unchanged.
## Scope
- MLA latent KV-cache append only.
- No cache-layout change.
- No allocator or block-table change.
- No dtype or quantization change.
- No MLA attention-dispatch change.
- No FlashInfer or FlashMLA dependency change.
- No behavior change for the normal KV-cache path.
## Current validation
- Real local CudaGraphManager capture/replay smoke test passes.
- Eager MLA append remains on the PyTorch fallback.
- CUDA Graph append results are bitwise equal to the PyTorch reference.
- A100 CUDA Graph append microbenchmark shows approximately 2.1x–3.0x speedup.
- Full MLA attention/decode validation is currently blocked because the available GPU is A100/SM80 and the upstream FlashMLA implementation requires SM90/SM100.
Related MLA PRs #4918, #4919, and #4920 are orthogonal to this append-dispatch change.
Contributor guide
Assessment
This issue has not been assessed yet.