vllm.v1.attention.backends.mla.flashattn_mla_sparse

Classes:

• FlashAttnMLASparseFA4Backend –

  FA4 for is_decode kernel calls; trtllm-gen for the rest.

• FlashAttnMLASparseFA4MetadataBuilder –

  The trtllm-gen builder, for the prefill lane's DCP workspace pre-size.

FlashAttnMLASparseFA4Backend

Bases: FlashInferMLASparseTRTLLMBackend

FA4 for is_decode kernel calls; trtllm-gen for the rest.

Source code in vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

class FlashAttnMLASparseFA4Backend(FlashInferMLASparseTRTLLMBackend):
    """FA4 for is_decode kernel calls; trtllm-gen for the rest."""

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
    # The qv kernel asserts every descale is None: BF16 KV cache only.
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]

    @staticmethod
    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
        return [MultipleOf(64)]

    @staticmethod
    def get_name() -> str:
        return "FLASH_ATTN_MLA_SPARSE_FA4"

    @staticmethod
    def get_builder_cls() -> type["FlashAttnMLASparseFA4MetadataBuilder"]:
        return FlashAttnMLASparseFA4MetadataBuilder

    @staticmethod
    def get_impl_cls() -> type[MLAAttentionImpl[Any]]:
        return FlashAttnMLASparseFA4Impl

    @classmethod
    def get_supported_head_sizes(cls) -> list[int]:
        # 512 latent + 64 RoPE; split back apart for the qv kernel.
        return [576]

    @classmethod
    def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None:
        from vllm.config import get_current_vllm_config_or_none
        from vllm.utils.flashinfer import has_flashinfer

        if not has_flashinfer():
            return "FA4 sparse MLA runs its prefill batches on FlashInfer's kernel"
        if (reason := _fa4_cute_mla_available()) is not None:
            return reason
        vllm_config = get_current_vllm_config_or_none()
        if vllm_config is None:
            return None
        if vllm_config.model_config is not None:
            hf_config = vllm_config.model_config.hf_text_config
            topk = getattr(hf_config, "index_topk", None)
            # gather_kv_indices' last dim must be a whole number of n-tiles.
            if topk is None or topk % 128 != 0:
                return f"FA4 sparse MLA requires index_topk % 128 == 0, got {topk}"
            dims = (
                getattr(hf_config, "kv_lora_rank", None),
                getattr(hf_config, "qk_rope_head_dim", None),
            )
            if dims != (512, 64):
                return (
                    "FA4 sparse MLA requires (kv_lora_rank, qk_rope_head_dim) "
                    f"== (512, 64), got {dims}"
                )
            # Under DCP the query is all-gathered: the kernel sees heads * dcp_size.
            dcp_size = vllm_config.parallel_config.decode_context_parallel_size
            num_heads = vllm_config.model_config.get_num_attention_heads(
                vllm_config.parallel_config
            )
            # 128 gathered heads take FA's prefill kernel, unvalidated under DCP.
            head_counts = (8, 16, 32, 64) if dcp_size > 1 else (8, 16, 32, 64, 128)
            if num_heads * dcp_size not in head_counts:
                return (
                    f"FA4 sparse MLA requires {head_counts} gathered query heads, "
                    f"got num_heads={num_heads} * dcp_size={dcp_size}"
                )
        return super().supports_combination(*args, **kwargs)

FlashAttnMLASparseFA4MetadataBuilder

Bases: FlashInferMLASparseTRTLLMMetadataBuilder

The trtllm-gen builder, for the prefill lane's DCP workspace pre-size.

Its UNIFORM_BATCH support and varlen decode bound also fit FA4, which runs one query row per token and bakes the row count into its varlen scalars.

Source code in vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

class FlashAttnMLASparseFA4MetadataBuilder(FlashInferMLASparseTRTLLMMetadataBuilder):
    """The trtllm-gen builder, for the prefill lane's DCP workspace pre-size.

    Its UNIFORM_BATCH support and varlen decode bound also fit FA4, which runs
    one query row per token and bakes the row count into its varlen scalars.
    """

_fa4_cute_mla_available()

Reason the FA4 cute-DSL MLA entry point is unusable, or None.

Source code in vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

def _fa4_cute_mla_available() -> str | None:
    """Reason the FA4 cute-DSL MLA entry point is unusable, or None."""
    from vllm.vllm_flash_attn import flash_attn_interface as fa

    if not fa.is_fa_version_supported(4):
        return f"FA4 unavailable: {fa.fa_version_unsupported_reason(4)}"
    try:
        from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd  # noqa: F401
    except Exception as e:  # cute-DSL / cutlass import chain
        return f"FA4 cute-DSL entry point failed to import: {e!r}"
    return None

_fa4_varlen_scalars(num_tokens, num_kv_rows, device) cached

Per-token varlen scalars; never evict, CUDA graphs bake in the pointers.

Source code in vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

@cache
def _fa4_varlen_scalars(
    num_tokens: int, num_kv_rows: int, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Per-token varlen scalars; never evict, CUDA graphs bake in the pointers."""
    cu_seqlens_q = torch.arange(num_tokens + 1, dtype=torch.int32, device=device)
    cu_seqlens_k = torch.zeros(num_tokens + 1, dtype=torch.int32, device=device)
    seqused_k = torch.full((num_tokens,), num_kv_rows, dtype=torch.int32, device=device)
    return cu_seqlens_q, cu_seqlens_k, seqused_k