vllm.models.deepseek_v41.sparse_mla

DeepSeek-V4.1 FlashMLA sparse backend, metadata, and metadata builders.

Classes:

• DeepseekV41SparseSWAMetadataBuilder –

  SWA metadata builder base for v4.1.

• DeepseekV4SparseMLABackend –

  DeepSeek-V4.1 sparse-MLA backend base.

• FlashMLAMegaAttnBackend –

  FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse

DeepseekV41SparseSWAMetadataBuilder

Bases: DeepseekSparseSWAMetadataBuilder

SWA metadata builder base for v4.1.

The shared builder classifies decode layer types by the v4.0 ratios (1 = SWA-only, 4, 128). v4.1 uses 0 / 1 / 2, so recompute the set of tile-scheduler plans from the v4.1 topology.

Source code in vllm/models/deepseek_v41/sparse_mla.py

class DeepseekV41SparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
    """SWA metadata builder base for v4.1.

    The shared builder classifies decode layer types by the v4.0 ratios
    (1 = SWA-only, 4, 128). v4.1 uses 0 / 1 / 2, so recompute the set of
    tile-scheduler plans from the v4.1 topology.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        compress_ratios = getattr(
            self.vllm_config.model_config.hf_config, "compress_ratios", None
        ) or [0]
        self._layer_types = {
            deepseek_v41_layer_type(int(ratio)) for ratio in compress_ratios
        }

DeepseekV4SparseMLABackend

Bases: AttentionBackend

DeepSeek-V4.1 sparse-MLA backend base.

Subclasses AttentionBackend directly (not the V3.2 FlashMLASparseBackend): DeepSeek-V4.1 runs its own attention layer (DeepseekV4Attention), so it does not reuse the V3.2 builder or impl, and only needs to declare its own metadata builder, KV-cache layout, and the sparse-MLA capability flags.

Source code in vllm/models/deepseek_v41/sparse_mla.py

class DeepseekV4SparseMLABackend(AttentionBackend):
    """DeepSeek-V4.1 sparse-MLA backend base.

    Subclasses ``AttentionBackend`` directly (not the V3.2
    ``FlashMLASparseBackend``): DeepSeek-V4.1 runs its own attention layer
    (``DeepseekV4Attention``), so it does not reuse the V3.2 builder or impl, and
    only needs to declare its own metadata builder, KV-cache layout, and the
    sparse-MLA capability flags.
    """

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "fp8_ds_mla",
        "fp8",  # alias for fp8_ds_mla
    ]

    @staticmethod
    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
        return [64 if current_platform.is_device_capability_family(90) else 128]

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

    @staticmethod
    def get_impl_cls() -> type[Any]:
        # DeepSeek-V4.1 runs its attention through ``DeepseekV4Attention.forward``,
        # not the generic ``Attention``/``MLAAttention`` layer, so the backend's
        # impl class is never instantiated.
        raise NotImplementedError(
            "DeepseekV4SparseMLABackend has no separate impl class; DeepSeek-V4.1 "
            "attention runs through DeepseekV4Attention."
        )

    @classmethod
    def get_supported_head_sizes(cls) -> list[int]:
        # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512.
        return [512]

    @classmethod
    def is_mla(cls) -> bool:
        return True

    @classmethod
    def is_sparse(cls) -> bool:
        return True

    @classmethod
    def supports_sink(cls) -> bool:
        return True

    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        return capability.major in [9, 10]

FlashMLAMegaAttnBackend

Bases: DeepseekV4FlashMLABackend

FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse RoPE + FP8 cast of the output, in one launch.

Same metadata and KV-cache geometry as FLASHMLA_SPARSE_DSV41 -- what differs is the attention layer's interface contract (it takes an unnormed, unroped Q and returns an already-inverse-RoPE'd, quantized output) and the extra nvfp4_ds_mla compressed-cache record only this kernel can read. SM100 only; the kernel has no SM90 instantiation.

Source code in vllm/models/deepseek_v41/sparse_mla.py

class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend):
    """FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse
    RoPE + FP8 cast of the output, in one launch.

    Same metadata and KV-cache geometry as ``FLASHMLA_SPARSE_DSV41`` -- what
    differs is the attention layer's interface contract (it takes an unnormed,
    unroped Q and returns an already-inverse-RoPE'd, quantized output) and the
    extra ``nvfp4_ds_mla`` compressed-cache record only this kernel can read.
    SM100 only; the kernel has no SM90 instantiation.
    """

    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "fp8_ds_mla",
        "fp8",  # alias for fp8_ds_mla
        "nvfp4_ds_mla",  # V4.1 fp8 SWA cache + NVFP4 compressed cache
    ]

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

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

    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        return capability.major == 10

    @classmethod
    def supports_combination(
        cls,
        head_size: int,
        dtype: torch.dtype,
        kv_cache_dtype: CacheDType | None,
        block_size: int | None,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
        use_mm_prefix: bool,
        device_capability: DeviceCapability,
    ) -> str | None:
        # Imported here: the layer module imports this backend class.
        from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
            is_flashmla_mega_attn_supported,
        )

        return is_flashmla_mega_attn_supported()[1]