NVIDIA / NVIDIA/Megatron-LM

[inference] Use Triton MLA latent KV append during CUDA Graph capture

Open
#5,823 1 comment 0 reactions 1 assignee Claimed by @sidsingh-nvidia View on GitHub
community-request enhancement waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.