| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Sink-capable FlashMLA sparse backend for HY V4 (NVIDIA). |
| 4 | |
| 5 | HY V4 adds a per-head learnable attention sink on top of sparse MLA. The |
| 6 | vendored FlashMLA kernels already accept an ``attn_sink`` argument, but vLLM's |
| 7 | shared ``FLASHMLA_SPARSE`` backend neither advertises sink support nor forwards |
| 8 | the tensor, so the bias would be silently dropped. |
| 9 | |
| 10 | This module supplies the missing wiring inside the model package, mirroring how |
| 11 | `vllm.models.deepseek_v4.nvidia.flashinfer_sparse` hands ``sinks`` to the |
| 12 | FlashInfer sparse MLA kernels: subclass the platform backend, declare |
| 13 | `supports_sink`, and thread the sink into every kernel call. |
| 14 | |
| 15 | The subclass intentionally keeps the inherited ``get_name()`` |
| 16 | (``"FLASHMLA_SPARSE"``). Several shared code paths key off that exact string — |
| 17 | `_canonicalize_sparse_mla_kv_cache_dtype` promotes a quantized KV cache to |
| 18 | ``fp8_ds_mla`` for it, and `FlashMLASparseImpl` asserts that layout — so a new |
| 19 | name would silently change KV cache behaviour. Only``supports_sink`` and the |
| 20 | two kernel wrappers differ from the parent. |
| 21 | """ |
| 22 | |
| 23 | from typing import TYPE_CHECKING |
| 24 | |
| 25 | import torch |
| 26 | |
| 27 | from vllm.logger import init_logger |
| 28 | from vllm.v1.attention.backends.mla.flashmla_sparse import ( |
| 29 | FlashMLASparseBackend, |
| 30 | FlashMLASparseImpl, |
| 31 | FlashMLASparseMetadata, |
| 32 | ) |
| 33 | from vllm.v1.attention.ops.flashmla import ( |
| 34 | flash_mla_sparse_fwd, |
| 35 | flash_mla_with_kvcache, |
| 36 | ) |
| 37 | from vllm.v1.worker.workspace import current_workspace_manager |
| 38 | |
| 39 | if TYPE_CHECKING: |
| 40 | from vllm.model_executor.models.deepseek_v2 import Indexer |
| 41 | |
| 42 | logger = init_logger(__name__) |
| 43 | |
| 44 | |
| 45 | class HYV4FlashMLASparseImpl(FlashMLASparseImpl): |
| 46 | """FlashMLA sparse impl that applies HY V4's per-head learnable sink. |
| 47 | |
| 48 | The sink enters as the ``sinks`` impl kwarg of |
| 49 | `vllm.model_executor.layers.attention.MLAAttention` and is consumed by the |
| 50 | FlashMLA kernels, which fold it into the softmax denominator: |
| 51 | ``out *= exp(lse) / (exp(lse) + exp(sink))``. |
| 52 | """ |
| 53 | |
| 54 | def __init__( |
| 55 | self, |
| 56 | num_heads: int, |
| 57 | head_size: int, |
| 58 | scale: float, |
| 59 | num_kv_heads: int, |
| 60 | alibi_slopes: list[float] | None, |
| 61 | sliding_window: int | None, |
| 62 | kv_cache_dtype: str, |
| 63 | logits_soft_cap: float | None, |
| 64 | attn_type: str, |
| 65 | kv_sharing_target_layer_name: str | None, |
| 66 | # MLA Specific Arguments |
| 67 | topk_indices_buffer: torch.Tensor | None = None, |
| 68 | indexer: "Indexer | None" = None, |
| 69 | **mla_args, |
| 70 | ) -> None: |
| 71 | # ``SparseMLACommonImpl`` takes explicit keyword arguments only, so the |
| 72 | # sink has to be removed before the base classes see``mla_args``. |
| 73 | sinks: torch.Tensor | None = mla_args.pop("sinks", None) |
| 74 | super().__init__( |
| 75 | num_heads, |
| 76 | head_size, |
| 77 | scale, |
| 78 | num_kv_heads, |
| 79 | alibi_slopes, |
| 80 | sliding_window, |
| 81 | kv_cache_dtype, |
| 82 | logits_soft_cap, |
| 83 | attn_type, |
| 84 | kv_sharing_target_layer_name, |
| 85 | topk_indices_buffer=topk_indices_buffer, |
| 86 | indexer=indexer, |
| 87 | **mla_args, |
| 88 | ) |
| 89 | self._validate_sinks(sinks, num_heads) |
| 90 | self.sinks = sinks |
| 91 | |
| 92 | @staticmethod |
| 93 | def _validate_sinks(sinks: torch.Tensor | None, num_heads: int) -> None: |
| 94 | """Reject sink tensors the FlashMLA kernels cannot consume. |
| 95 | |
| 96 | Args: |
| 97 | sinks: Candidate sink tensor, or None when the layer has no sink. |
| 98 | num_heads: Local (TP-sharded) query head count of this layer. |
| 99 | |
| 100 | Raises: |
| 101 | ValueError: If the dtype is not float32 or the shape is not |
| 102 | ``(num_heads,)``. |
| 103 | |
| 104 | """ |
| 105 | if sinks is None: |
| 106 | return |
| 107 | if sinks.dtype != torch.float32: |
| 108 | raise ValueError( |
| 109 | "HYV4 FlashMLA sparse attention sinks must have dtype " |
| 110 | f"torch.float32, but got {sinks.dtype}." |
| 111 | ) |
| 112 | if sinks.ndim != 1 or sinks.shape[0] != num_heads: |
| 113 | raise ValueError( |
| 114 | "HYV4 FlashMLA sparse attention sinks must have shape " |
| 115 | f"({num_heads},), but got {tuple(sinks.shape)}." |
| 116 | ) |
| 117 | |
| 118 | def _sinks_for_query( |
| 119 | self, |
| 120 | q: torch.Tensor, |
| 121 | head_dim: int, |
| 122 | kernel_heads: int, |
| 123 | ) -> torch.Tensor | None: |
| 124 | """Return the sink laid out for the kernel's query head count. |
| 125 | |
| 126 | Args: |
| 127 | q: Query tensor, before any head padding. |
| 128 | head_dim: Axis of ``q`` holding the query heads. |
| 129 | kernel_heads: Head count the kernel is invoked with, which may |
| 130 | exceed the query head count because of padding. |
| 131 | |
| 132 | Returns: |
| 133 | The sink tensor padded to ``kernel_heads`` with ``-inf`` (a no-op |
| 134 | sink) for the padded lanes, or None when the layer has no sink. |
| 135 | |
| 136 | Raises: |
| 137 | ValueError: If the sink and query head layouts disagree, or if they |
| 138 | live on different devices. |
| 139 | |
| 140 | """ |
| 141 | sinks = self.sinks |
| 142 | if sinks is None: |
| 143 | return None |
| 144 | |
| 145 | query_heads = q.shape[head_dim] |
| 146 | if sinks.shape[0] != query_heads: |
| 147 | raise ValueError( |
| 148 | "HYV4 FlashMLA sparse attention sink head count must match the " |
| 149 | f"runtime query layout: sinks={sinks.shape[0]}, " |
| 150 | f"query_heads={query_heads}. The sink must use the same " |
| 151 | "unpadded head layout as the query." |
| 152 | ) |
| 153 | if sinks.device != q.device: |
| 154 | raise ValueError( |
| 155 | "HYV4 FlashMLA sparse attention sinks and query must be on the " |
| 156 | f"same device, but got sinks={sinks.device}, query={q.device}." |
| 157 | ) |
| 158 | if kernel_heads < query_heads: |
| 159 | raise ValueError( |
| 160 | "HYV4 FlashMLA sparse kernel head count cannot be smaller than " |
| 161 | f"the runtime query layout: query_heads={query_heads}, " |
| 162 | f"kernel_heads={kernel_heads}." |
| 163 | ) |
| 164 | if kernel_heads == query_heads: |
| 165 | return sinks |
| 166 | |
| 167 | # Mirror the query padding the kernels require. Reading ``sinks`` here |
| 168 | # (rather than caching a padded copy at construction time) keeps the |
| 169 | # values correct for weights loaded after the module is built, and the |
| 170 | # allocation plus copy are captured in the CUDA graph. |
| 171 | padded_sinks = sinks.new_full((kernel_heads,), float("-inf")) |
| 172 | padded_sinks[:query_heads] = sinks |
| 173 | return padded_sinks |
| 174 | |
| 175 | def _fp8_flash_mla_kernel( |
| 176 | self, |
| 177 | q: torch.Tensor, |
| 178 | kv_c_and_k_pe_cache: torch.Tensor, |
| 179 | topk_indices: torch.Tensor, |
| 180 | kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata, |
| 181 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 182 | # q shape: (batch, seq_len, num_heads, head_dim) |
| 183 | actual_num_heads = q.size(2) |
| 184 | padded_num_heads = self.fp8_decode_padded_heads |
| 185 | attn_sink = self._sinks_for_query(q, head_dim=2, kernel_heads=padded_num_heads) |
| 186 | |
| 187 | # Pad query if needed (kernel only supports h_q = 64 or 128) |
| 188 | if actual_num_heads < padded_num_heads: |
| 189 | logger.warning_once( |
| 190 | f"Padding num_heads from {actual_num_heads} to " |
| 191 | f"{padded_num_heads} for FP8 sparse decode kernel" |
| 192 | ) |
| 193 | q_padded = q.new_zeros((q.size(0), q.size(1), padded_num_heads, q.size(3))) |
| 194 | q_padded[:, :, :actual_num_heads, :] = q |
| 195 | q = q_padded |
| 196 | |
| 197 | out, lse = flash_mla_with_kvcache( |
| 198 | q=q, |
| 199 | k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2), |
| 200 | block_table=kernel_metadata.dummy_block_table, |
| 201 | head_dim_v=512, |
| 202 | cache_seqlens=kernel_metadata.cache_lens, |
| 203 | tile_scheduler_metadata=kernel_metadata.scheduler_metadata, |
| 204 | is_fp8_kvcache=True, |
| 205 | indices=topk_indices, |
| 206 | softmax_scale=self.softmax_scale, |
| 207 | attn_sink=attn_sink, |
| 208 | ) |
| 209 | |
| 210 | # Slice output and lse back to actual head count if we padded |
| 211 | if actual_num_heads < padded_num_heads: |
| 212 | out = out[:, :, :actual_num_heads, :] |
| 213 | lse = lse[:, :actual_num_heads, :] |
| 214 | |
| 215 | return out, lse |
| 216 | |
| 217 | def _bf16_flash_mla_kernel( |
| 218 | self, |
| 219 | q: torch.Tensor, |
| 220 | kv_c_and_k_pe_cache: torch.Tensor, |
| 221 | topk_indices: torch.Tensor, |
| 222 | topk_length: torch.Tensor | None = None, |
| 223 | actual_num_heads: int | None = None, |
| 224 | out: torch.Tensor | None = None, |
| 225 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 226 | num_tokens = q.shape[0] |
| 227 | kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( |
| 228 | -1, 1, kv_c_and_k_pe_cache.shape[-1] |
| 229 | ) |
| 230 | |
| 231 | # NOTE(Chen): kernel requires num_local_head to be a multiple of |
| 232 | # 64 on hopper and 128 on blackwell. Pad from q's head count, not |
| 233 | # self.num_heads: under DCP the heads are all-gathered before this. |
| 234 | if actual_num_heads is None: |
| 235 | actual_num_heads = q.shape[1] |
| 236 | padded_num_heads = ( |
| 237 | (actual_num_heads + self.prefill_padding - 1) |
| 238 | // self.prefill_padding |
| 239 | * self.prefill_padding |
| 240 | ) |
| 241 | # q_concat_buffer may already include the kernel-required padding. |
| 242 | # The sink remains defined only for the real query heads. |
| 243 | attn_sink = self._sinks_for_query( |
| 244 | q[:, :actual_num_heads], head_dim=1, kernel_heads=padded_num_heads |
| 245 | ) |
| 246 | |
| 247 | if q.shape[1] < padded_num_heads: |
| 248 | logger.warning_once( |
| 249 | f"Padding num_heads from {actual_num_heads} to " |
| 250 | f"{padded_num_heads} for BF16 sparse prefill kernel" |
| 251 | ) |
| 252 | # Zero (not new_empty) the padded lanes: topk_indices is shared by |
| 253 | # all heads, so the kernel reduces across the head group and NaNs |
| 254 | # from uninitialized memory would leak into the real heads. |
| 255 | if out is None: |
| 256 | q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2])) |
| 257 | else: |
| 258 | *_, q_padded, _, _ = current_workspace_manager().get_simultaneous( |
| 259 | *self.workspace_specs |
| 260 | ) |
| 261 | q_padded = q_padded[:num_tokens] |
| 262 | q_padded[:, actual_num_heads:].zero_() |
| 263 | q_padded[:, :actual_num_heads, :] = q |
| 264 | q = q_padded |
| 265 | |
| 266 | topk_indices = topk_indices.view(num_tokens, 1, -1) |
| 267 | output, _, lse = flash_mla_sparse_fwd( |
| 268 | q, |
| 269 | kv_c_and_k_pe_cache, |
| 270 | topk_indices, |
| 271 | self.softmax_scale, |
| 272 | attn_sink=attn_sink, |
| 273 | topk_length=topk_length, |
| 274 | out=out, |
| 275 | ) |
| 276 | |
| 277 | output = output[:, :actual_num_heads, :] |
| 278 | lse = lse[:, :actual_num_heads] |
| 279 | return output, lse |
| 280 | |
| 281 | |
| 282 | class HYV4FlashMLASparseBackend(FlashMLASparseBackend): |
| 283 | """``FLASHMLA_SPARSE`` with attention-sink support for HY V4. |
| 284 | |
| 285 | Keeps the parent's name, metadata and builder; only the impl class and the |
| 286 | sink capability differ. See the module docstring for why the name is reused. |
| 287 | """ |
| 288 | |
| 289 | @staticmethod |
| 290 | def get_impl_cls() -> type[HYV4FlashMLASparseImpl]: |
| 291 | return HYV4FlashMLASparseImpl |
| 292 | |
| 293 | @classmethod |
| 294 | def supports_sink(cls) -> bool: |
| 295 | return True |