vllm.models.deepseek_v4_1.sparse_mla ¶

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
        }

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() -> 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]