Triton and MLA: portable attention and multi-head latent attention

源码版本v0.25.1

Responsibilities

TritonAttentionBackend (triton_attn.py:271) is vLLM's pure-Python attention kernel written in Triton. It does not depend on Dao Labs' flash_attn package or on FlashInfer. Its purpose is to serve as the fallback: when GPU compute is insufficient, when head_size is outside FA's supported range, or when the runtime platform is CPU/XPU, the selector still has a working implementation. It supports fp16/bf16/fp32, multiple FP8/INT4 KV cache quantization formats, and per-token-per-head inline scale layouts such as fp8_per_token_head / int8_per_token_head.

MLA (Multi-head Latent Attention, DeepSeek family) is a different path. It is not "swap the kernel" but "swap the KV cache semantics": K/V are not expanded, only the compressed latent kv_c (typically kv_lora_rank=512) plus a short RoPE segment k_pe are cached. So the MLA backend base class is not AttentionBackend but MLACommonBackend (mla_attention.py:1206); the impl base class is not AttentionImpl but MLAAttentionImpl (backend.py:941), and the interface is forward_mqa / forward_mha, not the ordinary forward. TritonMLABackend (triton_mla.py:81) is one such implementation: a pure-Triton decode kernel (decode_attention_fwd) that runs on any hardware. Several other MLA backends — FlashAttn / FlashInfer / FlashMLA / Cutlass / Tokenspeed / Aiter (ROCm) — live under vllm/v1/attention/backends/mla/.

Design motivation

• Triton fallback covers a wide matrix: TritonAttentionBackend.supported_dtypes includes float32 (triton_attn.py:271-287) and supported_kv_cache_dtypes includes int4_per_token_head / int8_per_token_head / fp8_per_token_head. Edge combinations that FA/FI reject run here.
• Non-causal attention supported: supports_non_causal returns True (triton_attn.py:301-303), so Prefix LM / ViT style bidirectional attention can run on Triton.
• No cascade: TritonAttentionBackend.use_cascade_attention explicitly returns False (triton_attn.py:368-370); it only takes the naive varlen path.
• MLA uses a single KV channel: MLACommonBackend.get_kv_cache_shape returns (num_blocks, block_size, head_size) (mla_attention.py:1215-1223). There are no separate K/V slots, because what is cached is the concatenated latent + RoPE vector, fetched as MQA during decode.
• MLA restricts head_size: get_supported_head_sizes only accepts [320, 576] (mla_attention.py:1236-1238), matching the kv_lora_rank + qk_rope_head_dim combinations of DeepSeek-V2/V3. Other head_sizes are rejected outright.
• lse must be returned: TritonMLAImpl.can_return_lse_for_decode = True (triton_mla.py:134-135), because cross-shard merging in DCP relies on the LSE. lse_base_on_e defaults to True, and the cross-shard merge kernel uses cp_lse_ag_out_rs.
• Multiple backends compete: AttentionBackendEnum lists TRITON_MLA / FLASH_ATTN_MLA / FLASHINFER_MLA / FLASHMLA / CUTLASS_MLA / TOKENSPEED_MLA (registry.py:80-82). The selector picks the best one based on GPU compute.

Key files

• TritonAttentionBackend:271 — Triton backend class; supports fp32 and several per_token_head quantized KV caches.
• Triton get_kv_cache_shape:317-345 — in per_token_head mode, pads head_dim to accommodate the inline scale.
• TritonAttentionImpl.forward:560 — calls the paged_attention Triton kernel; explicitly rejects output_block_scale.
• MLAAttention:339 — model-layer entry; holds kv_b_proj, qk_nope_head_dim, qk_rope_head_dim, kv_lora_rank.
• MLACommonBackend:1206 — MLA backend base class; defines is_mla() = True and the 3D KV cache shape.
• MLACommonImpl:1988 — MLA impl base class; declares the abstract forward_mqa / forward_mha.
• TritonMLABackend:81 — pure-Triton MLA backend; runs on any GPU.
• TritonMLAImpl.forward_mqa:189 — calls decode_attention_fwd Triton kernel with the is_mla=True branch.
• _compute_num_kv_splits:41-47 — picks KV split count from max_seq_len / 512, capped at sm_count * 2.
• FlashAttnMLABackend:43 — FA-route MLA, using flash_attn_varlen_func + MLA-specific metadata.

Data flow

The MLA oddity: Q is not computed together with K/V before attention. Instead it is split into q_pe (goes through RoPE) and q_nope (projected directly against the KV latent). The KV cache stores a single concatenated vector [kv_c (kv_lora_rank), k_pe (qk_rope_head_dim)]; decode fetches it once MQA-style and then runs o_proj over the head dimension. When TritonMLAImpl.forward_mqa calls decode_attention_fwd, it passes kv_c_and_k_pe_cache as a "single-head K cache":

# vllm/v1/attention/backends/mla/triton_mla.py L189-L256
def forward_mqa(
    self,
    q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    kv_c_and_k_pe_cache: torch.Tensor,
    attn_metadata: MLACommonMetadata,
    layer: AttentionLayer,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    assert kv_c_and_k_pe_cache.numel() > 0
    assert attn_metadata.decode is not None

    if type(q) is tuple:
        q = torch.cat(q, dim=-1)

    assert isinstance(q, torch.Tensor)
    B = q.shape[0]
    q_num_heads = q.shape[1]
    o = torch.zeros(
        B, q_num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
    )
    lse = torch.zeros(B, q_num_heads, dtype=q.dtype, device=q.device)

    # For batch invariance, use only 1 split to ensure deterministic reduction
    if envs.VLLM_BATCH_INVARIANT:
        num_kv_splits = 1
    else:
        num_kv_splits = _compute_num_kv_splits(
            attn_metadata.max_seq_len, self._sm_count
        )

    # NOTE: the +1 stores the LogSumExp (LSE) that the stage2 kernel uses
    # to merge partial attention outputs across splits.
    logits_shape = (B, q_num_heads, num_kv_splits, self.kv_lora_rank + 1)
    if is_workspace_manager_initialized():
        (attn_logits,) = current_workspace_manager().get_simultaneous(
            (logits_shape, torch.float32),
        )
    else:
        attn_logits = torch.empty(
            logits_shape, dtype=torch.float32, device=q.device
        )

    # Add a head dim of 1
    kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)
    kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]
    PAGE_SIZE = kv_c_and_k_pe_cache.size(1)

    # Run MQA — always pass layer scales. When KV cache is
    # BF16 the kernel's `if dtype.is_fp8()` check is a no-op.
    decode_attention_fwd(
        q,
        kv_c_and_k_pe_cache,
        kv_c_cache,
        o,
        lse,
        attn_metadata.decode.block_table,
        attn_metadata.decode.seq_lens,
        attn_logits,
        num_kv_splits,
        self.scale,
        PAGE_SIZE,
        k_scale=layer._k_scale,
        v_scale=layer._k_scale,
        is_mla=True,
    )

    return o, lse

num_kv_splits is not hard-coded: it is computed dynamically from the longest seq_len in the batch. Long sequences get more splits so more SMs can run in parallel; the merge stage uses lse for LogSumExp reduction. The workspace is fetched from workspace_manager to avoid torch.empty on the decode hot path.

Boundaries and failure modes

• Triton does not support fused block_scale output: when TritonAttentionImpl.forward sees output_block_scale is not None it raises NotImplementedError (triton_attn.py:584-588).
• MLA does not support alibi/sliding_window/logits_soft_cap: TritonMLAImpl.__init__ explicitly checks and raises NotImplementedError (triton_mla.py:166-170).
• MLA head_size is restricted to 320 or 576: MLACommonBackend.get_supported_head_sizes (mla_attention.py:1236-1238) returns only these two, so anything outside DeepSeek-V2/V3 cannot get in.
• FP8 KV cache requires BF16 query: under FP8 KV cache, TritonMLAImpl sets supports_quant_query_input = False (triton_mla.py:183-185). The Triton kernel does the dequant internally, so the upper layer cannot also compress Q to FP8.
• workspace_manager fallback when uninitialized: forward_mqa falls back to torch.empty when is_workspace_manager_initialized() returns False (triton_mla.py:225-232). This works in unit tests, but on the production path GPUModelRunner must have pre-allocated the workspace.
• num_kv_splits affects LSE merge correctness: a change in splits changes the shape of attn_logits, and the stage2 merge kernel must use the same split value. That is why _compute_num_kv_splits (triton_mla.py:41) is the same function called both when the builder reserves the workspace and during forward.

Summary

Triton is vLLM's fallback backend when FA/FI are unavailable, covering fp32 and edge-case quantized KV caches; MLA is a dedicated sub-abstraction for DeepSeek-style multi-head latent attention, where the KV cache stores the compressed latent rather than expanded K/V. Both hang underneath the /attention/backend abstraction and are selected by get_attn_backend at initialization. Triton and MLA KV cache shapes are ultimately managed by /kv-cache/kv-cache-manager; the sampling hand-off is in /sampling/sampler.

See official docs: vLLM docs · README