| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepSeek-V4.1 FlashMLA sparse backend, metadata, and metadata builders.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | from typing import Any, ClassVar |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.config.cache import CacheDType |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.platforms.interface import DeviceCapability |
| 14 | from vllm.v1.attention.backend import ( |
| 15 | AttentionBackend, |
| 16 | AttentionCGSupport, |
| 17 | AttentionMetadata, |
| 18 | AttentionMetadataBuilder, |
| 19 | CommonAttentionMetadata, |
| 20 | MultipleOf, |
| 21 | ) |
| 22 | from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping |
| 23 | from vllm.v1.attention.backends.mla.sparse_swa import ( |
| 24 | _LAYER_TYPE_C1A, |
| 25 | _LAYER_TYPE_C2A, |
| 26 | _LAYER_TYPE_SWAONLY, |
| 27 | DeepseekSparseSWAMetadataBuilder, |
| 28 | ) |
| 29 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 30 | |
| 31 | # v4.1 per-layer compress ratios: 0 = pure sliding window, 1 = full-length |
| 32 | # compressed cache, 2 = ratio-2 compressed cache. Ratio-1 and ratio-2 layers |
| 33 | # both attend over indexer topk indices into a shared compressed cache but |
| 34 | # differ in compressed page block size (block_size // ratio), so each needs its |
| 35 | # own FlashMLA tile-scheduler plan. |
| 36 | _V41_LAYER_TYPES: dict[int, str] = { |
| 37 | 0: _LAYER_TYPE_SWAONLY, |
| 38 | 1: _LAYER_TYPE_C1A, |
| 39 | 2: _LAYER_TYPE_C2A, |
| 40 | } |
| 41 | |
| 42 | |
| 43 | def deepseek_v41_layer_type(compress_ratio: int) -> str: |
| 44 | layer_type = _V41_LAYER_TYPES.get(compress_ratio) |
| 45 | if layer_type is None: |
| 46 | raise ValueError( |
| 47 | f"Unsupported DeepSeek V4.1 compress_ratio={compress_ratio}; " |
| 48 | "expected 0, 1, or 2." |
| 49 | ) |
| 50 | return layer_type |
| 51 | |
| 52 | |
| 53 | class DeepseekV41SparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): |
| 54 | """SWA metadata builder base for v4.1. |
| 55 | |
| 56 | The shared builder classifies decode layer types by the v4.0 ratios |
| 57 | (1 = SWA-only, 4, 128). v4.1 uses 0 / 1 / 2, so recompute the set of |
| 58 | tile-scheduler plans from the v4.1 topology. |
| 59 | """ |
| 60 | |
| 61 | def __init__(self, *args, **kwargs): |
| 62 | super().__init__(*args, **kwargs) |
| 63 | compress_ratios = getattr( |
| 64 | self.vllm_config.model_config.hf_config, "compress_ratios", None |
| 65 | ) or [0] |
| 66 | self._layer_types = { |
| 67 | deepseek_v41_layer_type(int(ratio)) for ratio in compress_ratios |
| 68 | } |
| 69 | |
| 70 | |
| 71 | class DeepseekV4SparseMLABackend(AttentionBackend): |
| 72 | """DeepSeek-V4.1 sparse-MLA backend base. |
| 73 | |
| 74 | Subclasses ``AttentionBackend`` directly (not the V3.2 |
| 75 | ``FlashMLASparseBackend``): DeepSeek-V4.1 runs its own attention layer |
| 76 | (``DeepseekV4Attention``), so it does not reuse the V3.2 builder or impl, and |
| 77 | only needs to declare its own metadata builder, KV-cache layout, and the |
| 78 | sparse-MLA capability flags. |
| 79 | """ |
| 80 | |
| 81 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] |
| 82 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 83 | "auto", |
| 84 | "fp8_ds_mla", |
| 85 | "fp8", # alias for fp8_ds_mla |
| 86 | ] |
| 87 | |
| 88 | @staticmethod |
| 89 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 90 | return [64 if current_platform.is_device_capability_family(90) else 128] |
| 91 | |
| 92 | @staticmethod |
| 93 | def get_builder_cls() -> type["DeepseekV4SparseMLAMetadataBuilder"]: |
| 94 | return DeepseekV4SparseMLAMetadataBuilder |
| 95 | |
| 96 | @staticmethod |
| 97 | def get_impl_cls() -> type[Any]: |
| 98 | # DeepSeek-V4.1 runs its attention through ``DeepseekV4Attention.forward``, |
| 99 | # not the generic ``Attention``/``MLAAttention`` layer, so the backend's |
| 100 | # impl class is never instantiated. |
| 101 | raise NotImplementedError( |
| 102 | "DeepseekV4SparseMLABackend has no separate impl class; DeepSeek-V4.1 " |
| 103 | "attention runs through DeepseekV4Attention." |
| 104 | ) |
| 105 | |
| 106 | @classmethod |
| 107 | def get_supported_head_sizes(cls) -> list[int]: |
| 108 | # DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512. |
| 109 | return [512] |
| 110 | |
| 111 | @classmethod |
| 112 | def is_mla(cls) -> bool: |
| 113 | return True |
| 114 | |
| 115 | @classmethod |
| 116 | def is_sparse(cls) -> bool: |
| 117 | return True |
| 118 | |
| 119 | @classmethod |
| 120 | def supports_sink(cls) -> bool: |
| 121 | return True |
| 122 | |
| 123 | @classmethod |
| 124 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 125 | return capability.major in [9, 10] |
| 126 | |
| 127 | |
| 128 | @dataclass |
| 129 | class DeepseekV4FlashMLAMetadata(AttentionMetadata): |
| 130 | num_reqs: int |
| 131 | max_query_len: int |
| 132 | max_seq_len: int |
| 133 | |
| 134 | num_actual_tokens: int # Number of tokens excluding padding. |
| 135 | query_start_loc: torch.Tensor |
| 136 | slot_mapping: torch.Tensor |
| 137 | |
| 138 | block_table: torch.Tensor |
| 139 | req_id_per_token: torch.Tensor |
| 140 | block_size: int |
| 141 | topk_tokens: int |
| 142 | |
| 143 | |
| 144 | class DeepseekV4SparseMLAMetadataBuilder( |
| 145 | AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata] |
| 146 | ): |
| 147 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 148 | |
| 149 | def __init__( |
| 150 | self, |
| 151 | kv_cache_spec: AttentionSpec, |
| 152 | layer_names: list[str], |
| 153 | vllm_config: VllmConfig, |
| 154 | device: torch.device, |
| 155 | ) -> None: |
| 156 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 157 | self.model_config = vllm_config.model_config |
| 158 | # Classify single-token queries (plus num_speculative_tokens via |
| 159 | # supports_spec_as_decode=True) as decodes; longer queries go to prefill. |
| 160 | self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) |
| 161 | self.topk_tokens = self.model_config.hf_config.index_topk |
| 162 | |
| 163 | max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 164 | self.req_id_per_token_buffer = torch.empty( |
| 165 | (max_num_batched_tokens,), dtype=torch.int32, device=device |
| 166 | ) |
| 167 | |
| 168 | assert isinstance(self.kv_cache_spec.tokens_per_state, int) |
| 169 | self.compress_ratio = self.kv_cache_spec.tokens_per_state |
| 170 | # Only kv-source layers own a compressed cache, and v4.1 compresses at |
| 171 | # ratio 1 or 2. |
| 172 | if self.compress_ratio not in (1, 2): |
| 173 | raise ValueError( |
| 174 | "DeepSeek V4.1 compressed-KV caches use compress_ratio 1 or 2; " |
| 175 | f"got {self.compress_ratio}." |
| 176 | ) |
| 177 | |
| 178 | # Pre-allocate compressed slot mapping buffer for CUDA graph address |
| 179 | # stability when compress_ratio > 1. |
| 180 | if self.compress_ratio > 1: |
| 181 | self.compressed_slot_mapping_buffer = torch.empty( |
| 182 | max_num_batched_tokens, dtype=torch.int64, device=device |
| 183 | ) |
| 184 | |
| 185 | def build( |
| 186 | self, |
| 187 | common_prefix_len: int, |
| 188 | common_attn_metadata: CommonAttentionMetadata, |
| 189 | fast_build: bool = False, |
| 190 | ) -> DeepseekV4FlashMLAMetadata: |
| 191 | cm = common_attn_metadata |
| 192 | req_id_per_token = cm.token_to_req_indices(self.req_id_per_token_buffer) |
| 193 | |
| 194 | slot_mapping = cm.slot_mapping |
| 195 | if self.compress_ratio > 1: |
| 196 | slot_mapping = get_compressed_slot_mapping( |
| 197 | cm.num_actual_tokens, |
| 198 | cm.slot_mapping, |
| 199 | cm.query_start_loc, |
| 200 | cm.seq_lens, |
| 201 | cm.block_table_tensor.clamp_(min=0), |
| 202 | int(self.kv_cache_spec.num_states), |
| 203 | self.compress_ratio, |
| 204 | out=self.compressed_slot_mapping_buffer, |
| 205 | ) |
| 206 | |
| 207 | return DeepseekV4FlashMLAMetadata( |
| 208 | num_reqs=cm.num_reqs, |
| 209 | max_query_len=cm.max_query_len, |
| 210 | max_seq_len=cm.max_seq_len, |
| 211 | num_actual_tokens=cm.num_actual_tokens, |
| 212 | query_start_loc=cm.query_start_loc, |
| 213 | slot_mapping=slot_mapping, |
| 214 | block_table=cm.block_table_tensor, |
| 215 | req_id_per_token=req_id_per_token, |
| 216 | block_size=self.kv_cache_spec.block_size, |
| 217 | topk_tokens=self.topk_tokens, |
| 218 | ) |
| 219 | |
| 220 | |
| 221 | class DeepseekV4FlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): |
| 222 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 223 | |
| 224 | |
| 225 | class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend): |
| 226 | @staticmethod |
| 227 | def get_name() -> str: |
| 228 | return "FLASHMLA_SPARSE_DSV41" |
| 229 | |
| 230 | @staticmethod |
| 231 | def get_builder_cls() -> type[DeepseekV4FlashMLAMetadataBuilder]: |
| 232 | return DeepseekV4FlashMLAMetadataBuilder |
| 233 | |
| 234 | |
| 235 | class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend): |
| 236 | """FlashMLA's mega-attention kernel: Q RoPE + sparse attention + inverse |
| 237 | RoPE + FP8 cast of the output, in one launch. |
| 238 | |
| 239 | Same metadata and KV-cache geometry as ``FLASHMLA_SPARSE_DSV41`` -- what |
| 240 | differs is the attention layer's interface contract (it takes an unnormed, |
| 241 | unroped Q and returns an already-inverse-RoPE'd, quantized output) and the |
| 242 | extra ``nvfp4_ds_mla`` compressed-cache record only this kernel can read. |
| 243 | SM100 only; the kernel has no SM90 instantiation. |
| 244 | """ |
| 245 | |
| 246 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 247 | "auto", |
| 248 | "fp8_ds_mla", |
| 249 | "fp8", # alias for fp8_ds_mla |
| 250 | "nvfp4_ds_mla", # V4.1 fp8 SWA cache + NVFP4 compressed cache |
| 251 | ] |
| 252 | |
| 253 | @staticmethod |
| 254 | def get_name() -> str: |
| 255 | return "FLASHMLA_MEGA_ATTN_DSV41" |
| 256 | |
| 257 | @staticmethod |
| 258 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 259 | return [128] |
| 260 | |
| 261 | @classmethod |
| 262 | def supports_compute_capability(cls, capability: DeviceCapability) -> bool: |
| 263 | return capability.major == 10 |
| 264 | |
| 265 | @classmethod |
| 266 | def supports_combination( |
| 267 | cls, |
| 268 | head_size: int, |
| 269 | dtype: torch.dtype, |
| 270 | kv_cache_dtype: CacheDType | None, |
| 271 | block_size: int | None, |
| 272 | use_mla: bool, |
| 273 | has_sink: bool, |
| 274 | use_sparse: bool, |
| 275 | use_mm_prefix: bool, |
| 276 | device_capability: DeviceCapability, |
| 277 | ) -> str | None: |
| 278 | # Imported here: the layer module imports this backend class. |
| 279 | from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( |
| 280 | is_flashmla_mega_attn_supported, |
| 281 | ) |
| 282 | |
| 283 | return is_flashmla_mega_attn_supported()[1] |