Overview
Multi-head Latent Attention (MLA) is how DeepSeek-V2, V3 and R1 keep long contexts affordable: instead of storing full keys and values for every head, it compresses each token into one small latent vector. That shrinks the KV cache, but torch.compile still runs MLA decode as several separate GPU operations and writes the full attention-score matrix to memory between them. triton-mla does the whole step in a single pass over the compressed cache.
How it works
- Absorbed form. At decode time the up-projections are folded into the query and output, so all 128 heads attend directly over one shared 576-dim vector per token. That makes decode purely bandwidth bound on reading the cache.
- Load once, reuse everywhere. Because every head shares the same cached vector, each tile is read from memory once and reused by up to 64 heads through tensor-core matmuls, with an online softmax kept in registers.
- Split-KV. Long sequences are split across the GPU so small batches still fill every SM, then a second stage merges the partial results with a log-sum-exp reduction. The split count is chosen to fill about one wave of SMs.
- Tuning. Block sizes, warps and pipeline stages were swept on an H100 through Modal, and only configs that won at least one shape are kept. Tuning alone cut some shapes by more than a third.
- Paged cache. It reads from a paged KV cache, the layout inference servers actually use, and every benchmark case is checked against a PyTorch reference before it is timed.
Results
On an H100 with 128 heads in bf16, the kernel beats torch.compile in max-autotune mode in 15 of 16 shapes and peaks at 249 TFLOPS. The gap grows with batch size and context length, reaching 3.6x at batch 128 and 32K tokens. The one loss is batch 8 at 1K tokens, where launch overhead dominates.
Open source
triton-mla is MIT licensed and installs with pip install triton-mla. The package includes the kernel, a full MLA layer, and a triton-mla-bench command for running the benchmarks yourself. It currently covers single-token decode only, without an FP8 cache or Blackwell support.