Skip to content

Attention parity: run every attention type on the best available path on both targets #23

Description

@randomvariable

Tracking issue for the standing requirement that both targets — CUDA sm_121a and ROCm gfx1151 — run every attention type on the most optimized available path, rather than merely running. A model that only serves by falling back to a slower generic backend counts as a gap here, not as support.

Backend selection is capability-based: each candidate is filtered by validate_configuration() (vllm/v1/attention/backend.py:320), so a model can silently land on a generic path when a probe rejects the specialised one. That is invisible without reading the selected-backend logs, which is what makes these gaps worth tracking explicitly.

Priority order is FLASH_ATTN → FLASHINFER → TRITON_ATTN on CUDA (vllm/platforms/cuda.py:145) and ROCM_ATTN → ROCM_AITER_FA → ROCM_AITER_UNIFIED_ATTN → TRITON_ATTN on ROCm (vllm/platforms/rocm.py:450).

Current state

Attention type CUDA sm_121a ROCm gfx1151
Dense causal GQA/MQA FLASH_ATTN, FA2 (FA4 is gated to major 10, fa_utils.py:154) ROCM_ATTN; native paged kernel only within a narrow shape envelope (rocm.py:413), otherwise internal Triton
Sliding-window / global FLASH_ATTN with native window metadata ROCM_ATTN Triton paged decode; out-of-window block skip landed, hardware-validated, not yet benchmarked
Attention sinks TRITON_ATTN — FA2 needs FA3/FA4 (fa_utils.py:315), FlashInfer allows sinks only on SM100 (flashinfer.py:472) ROCM_AITER_UNIFIED_ATTN when AITER is present (rocm_aiter_unified_attn.py:78), else TRITON_ATTN
Dense MLA TRITON_MLA (cuda.py:130) TRITON_MLA; AITER MLA now gated by architecture
Sparse / compressed MLA specialised DeepseekV4FlashInferSM120Attention ROCM_AITER_MLA_SPARSE; shuffled-layout indexer decode restored in dfff80ad6, pending hardware validation (#26)
Gated DeltaNet GDN_ATTN, in-tree Triton (FlashInfer selection recognises Hopper/SM10.x, not SM12) Triton two-kernel prefill; the fused path was removed in 3c9a8d20d as inaccurate and uncompilable
Mamba2 Triton SSU + chunk scan same
Kimi Delta Attention in-tree FLA/Triton same
Encoder / cross attention FLASH_ATTN where supported ViT gets Triton FlashAttention (rocm.py:715); cross-attention falls to TRITON_ATTN (rocm.py:560)

Open gaps

Ranked by cost of a miss, not by effort. Original numbering retained so earlier references stay valid.

  1. sm_121a attention sinks drop to generic Triton. Affects sink-based models. Closing this means porting sink support into an SM12-capable path, not relaxing a gate.
  2. sm_121a dense MLA is generic Triton. Distinct from the sparse DeepSeek-V4 path, which is already specialised.
  3. ROCm cross-attention is generic Triton. Encoder self-attention being optimized does not imply cross-attention is.
  4. Dense gfx1151 shapes outside the native envelope silently use Triton while still reporting ROCM_ATTN. Profile before extending the kernel.

Closed gaps

  1. gfx1151 sparse-MLA indexer — was a correctness bug, not a performance gap. The DSV4 indexer cache is always written in a shuffled 16x16 packing: get_supported_kernel_block_sizes() returns [256], the indexer exists only at compress_ratio == 4, and storage_block_size = block_size // compress_ratio gives 64, so both writers select SHUFFLE. Neither the AITER stage-1 route nor the Torch reference understood that layout, so the indexer read tiled bytes as row-major and returned plausible but wrong top-k selections — silently, with degraded output quality as the only symptom. Prefill was unaffected (cp_gather_indexer_k_quant_cache_triton takes a LAYOUT argument). Restored in dfff80ad6 with a regression oracle that independently packs and reads the 16x16 layout and checks exact top-k selection. Restore sparse-MLA indexer decode on gfx1151 #26 remains open for hardware end-to-end long-context validation.

  2. gfx1151 dense MLA selected a backend unsupported there. AiterMLABackend had no architecture rejection, so automatic selection could pick it over TRITON_MLA. Fixed by adding a capability gate; auto-selection now falls through correctly on gfx1151.

  3. gfx1151 Gated DeltaNet fused prefill — investigated, path removed. Framed as "a faster path exists but is never auto-selected". Both halves were wrong. On hardware: it did not compile as shipped (BV=8 autotune configs emit an 8-wide tl.dot, below RDNA3.5 WMMA's 16-element minimum); forced to compile, it was less accurate, 9.5x worse at seq_len 63 rising to ~76x at 192, with final_state bit-exact throughout — isolating the defect to the inter-chunk o-GEMM rather than the recurrence; and it was unreachable anyway, since a NotImplementedError guard fired on every prefill. The fused backend was removed in 3c9a8d20d; gfx1151 GDN prefill stays on the accurate two-kernel Triton path. Fused GDN prefill produces inaccurate output on gfx1151 #24 and gdn_prefill_backend="fla_fused" cannot serve a request #25 closed.

Non-gap: block-sparse GQA

Recorded as a positive case so the table above is not read as exhaustive. vllm/models/minimax_m3/ implements block-sparse GQA selection with a working path on both targets: common/indexer.py scores KV blocks with index heads and selects top-k plus fixed init/local blocks, common/sparse_attention.py consumes only the selected blocks, and amd/ops/sparse_pa.py compacts selected logical blocks into page-16 tables on ROCm. The kernel implementation is chosen via select_indexer_impl_cls, resolved by current_platform.is_rocm().

The mechanism is model-specific, not generic: registered as MINIMAX_M3_SPARSE, hard-wired to 128-token pages with trained index heads, deliberately mirroring DeepSeek-V4's separate-indexer structure. A model that would benefit from block-sparse GQA but lacks trained index heads has no generic path — the same category of gap this issue tracks, one layer above backend selection.

Notes

gfx1151 has no FP8 tensor cores, so gaps involving FP8 compute need BF16 or integer RDNA kernels — relaxing a capability gate would produce a slow or incorrect path, not acceleration.

Acceptance for any item here is the selected backend recorded from logs, correctness against the previous path, and a measurement on the target. Architecture support is not a performance claim.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area/ideaOptimization idea candidate for evaluationkimirocm

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions