| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import TYPE_CHECKING, ClassVar, Optional |
| 6 | |
| 7 | import numpy as np |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.config.cache import CacheDType |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 14 | get_mla_dims, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.attention.sparse_mla_attention import ( |
| 17 | SharedTopkIndicesBuffer, |
| 18 | ) |
| 19 | from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor |
| 20 | from vllm.v1.attention.backend import ( |
| 21 | AttentionBackend, |
| 22 | AttentionCGSupport, |
| 23 | AttentionLayer, |
| 24 | AttentionMetadata, |
| 25 | AttentionMetadataBuilder, |
| 26 | CommonAttentionMetadata, |
| 27 | MLAAttentionImpl, |
| 28 | ) |
| 29 | from vllm.v1.attention.backends.mla.sparse_utils import ( |
| 30 | flat_kv_row_view, |
| 31 | triton_convert_req_index_to_global_index, |
| 32 | ) |
| 33 | from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface |
| 34 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 35 | |
| 36 | if TYPE_CHECKING: |
| 37 | from vllm.model_executor.models.deepseek_v2 import Indexer |
| 38 | logger = init_logger(__name__) |
| 39 | |
| 40 | |
| 41 | class XPUMLASparseBackend(AttentionBackend): |
| 42 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] |
| 43 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 44 | "auto", |
| 45 | "float16", |
| 46 | "bfloat16", |
| 47 | ] |
| 48 | |
| 49 | @staticmethod |
| 50 | def get_name() -> str: |
| 51 | return "XPU_MLA_SPARSE" |
| 52 | |
| 53 | @staticmethod |
| 54 | def get_metadata_cls() -> type["XPUMLASparseMetadata"]: |
| 55 | return XPUMLASparseMetadata |
| 56 | |
| 57 | @staticmethod |
| 58 | def get_builder_cls() -> type["XPUMLASparseMetadataBuilder"]: |
| 59 | return XPUMLASparseMetadataBuilder |
| 60 | |
| 61 | @staticmethod |
| 62 | def get_impl_cls() -> type["XPUMLASparseImpl"]: |
| 63 | return XPUMLASparseImpl |
| 64 | |
| 65 | @classmethod |
| 66 | def is_mla(cls) -> bool: |
| 67 | return True |
| 68 | |
| 69 | @classmethod |
| 70 | def is_sparse(cls) -> bool: |
| 71 | return True |
| 72 | |
| 73 | @classmethod |
| 74 | def get_supported_head_sizes(cls) -> list[int]: |
| 75 | return [576] |
| 76 | |
| 77 | |
| 78 | @dataclass |
| 79 | class XPUMLASparseMetadata(AttentionMetadata): |
| 80 | num_reqs: int |
| 81 | max_query_len: int |
| 82 | max_seq_len: int |
| 83 | |
| 84 | num_actual_tokens: int # Number of tokens excluding padding. |
| 85 | query_start_loc: torch.Tensor |
| 86 | slot_mapping: torch.Tensor |
| 87 | |
| 88 | block_table: torch.Tensor |
| 89 | req_id_per_token: torch.Tensor |
| 90 | |
| 91 | block_size: int = 1 |
| 92 | topk_tokens: int = 2048 |
| 93 | |
| 94 | # The shared MLA layer (`mla_attention.py::forward_impl`) reads these |
| 95 | # decode/prefill counts unconditionally for every MLA metadata (it asserts |
| 96 | # `num_decodes`/`num_prefills`/`num_decode_tokens is not None` and uses |
| 97 | # `num_decode_tokens` to split MQA vs dense-MHA tokens). The CUDA sparse |
| 98 | # backends carry them via `SparseMLACommonMetadataBuilder`; this XPU backend |
| 99 | # builds its own metadata and previously omitted them, so a sparse-MLA |
| 100 | # (DeepSeek / GLM DSA) run on XPU crashed with |
| 101 | # `'XPUMLASparseMetadata' object has no attribute 'num_decode_tokens'`. |
| 102 | # This backend serves both prefill and decode through the top-k sparse MQA |
| 103 | # path (see `forward_mqa`), so all tokens are routed as "decode" |
| 104 | # (`num_decode_tokens == num_actual_tokens`, `num_prefills == 0`); that keeps |
| 105 | # the shared layer's `num_mha_tokens` at 0 and never enters the dense-MHA |
| 106 | # prefill branch (which needs prefill-only fields this backend lacks). |
| 107 | num_decodes: int = 0 |
| 108 | num_prefills: int = 0 |
| 109 | num_decode_tokens: int = 0 |
| 110 | |
| 111 | |
| 112 | @dataclass |
| 113 | class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]): |
| 114 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER |
| 115 | |
| 116 | def __init__( |
| 117 | self, |
| 118 | kv_cache_spec: AttentionSpec, |
| 119 | layer_names: list[str], |
| 120 | vllm_config: VllmConfig, |
| 121 | device: torch.device, |
| 122 | ): |
| 123 | self.kv_cache_spec = kv_cache_spec |
| 124 | self.model_config = vllm_config.model_config |
| 125 | parallel_config = vllm_config.parallel_config |
| 126 | self.device = device |
| 127 | max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 128 | |
| 129 | self.num_heads = self.model_config.get_num_attention_heads(parallel_config) |
| 130 | self.mla_dims = get_mla_dims(self.model_config) |
| 131 | self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk |
| 132 | self.topk_tokens_tensor = torch.tensor( |
| 133 | [self.topk_tokens], device=device, dtype=torch.int32 |
| 134 | ) |
| 135 | self.max_model_len_tensor = torch.tensor( |
| 136 | [self.model_config.max_model_len], device=device, dtype=torch.int32 |
| 137 | ) |
| 138 | # this is ignored by `flash_mla_with_kvcache` if indices not None |
| 139 | self.dummy_block_table = torch.empty( |
| 140 | (1, 1), dtype=torch.int32, device=self.device |
| 141 | ) |
| 142 | |
| 143 | self.req_id_per_token_buffer = torch.empty( |
| 144 | (max_num_batched_tokens,), |
| 145 | dtype=torch.int32, |
| 146 | device=device, |
| 147 | ) |
| 148 | |
| 149 | def build( |
| 150 | self, |
| 151 | common_prefix_len: int, |
| 152 | common_attn_metadata: CommonAttentionMetadata, |
| 153 | fast_build: bool = False, |
| 154 | ) -> XPUMLASparseMetadata: |
| 155 | num_tokens = common_attn_metadata.num_actual_tokens |
| 156 | starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32) |
| 157 | seg_lengths = np.diff(starts) |
| 158 | req_id_per_token = np.repeat( |
| 159 | np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths |
| 160 | ) |
| 161 | # Zero-fill for cudagraphs |
| 162 | self.req_id_per_token_buffer.fill_(0) |
| 163 | self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_( |
| 164 | np_to_pinned_tensor(req_id_per_token), non_blocking=True |
| 165 | ) |
| 166 | |
| 167 | req_id_per_token = self.req_id_per_token_buffer[:num_tokens] |
| 168 | |
| 169 | metadata = XPUMLASparseMetadata( |
| 170 | num_reqs=common_attn_metadata.num_reqs, |
| 171 | max_query_len=common_attn_metadata.max_query_len, |
| 172 | max_seq_len=common_attn_metadata.max_seq_len, |
| 173 | num_actual_tokens=common_attn_metadata.num_actual_tokens, |
| 174 | query_start_loc=common_attn_metadata.query_start_loc, |
| 175 | slot_mapping=common_attn_metadata.slot_mapping, |
| 176 | block_table=common_attn_metadata.block_table_tensor, |
| 177 | req_id_per_token=req_id_per_token, |
| 178 | block_size=self.kv_cache_spec.block_size, |
| 179 | topk_tokens=self.topk_tokens, |
| 180 | # Route every token through the sparse MQA path (see the field |
| 181 | # definitions above); this backend has no dense-MHA prefill. |
| 182 | num_decodes=common_attn_metadata.num_reqs, |
| 183 | num_prefills=0, |
| 184 | num_decode_tokens=common_attn_metadata.num_actual_tokens, |
| 185 | ) |
| 186 | return metadata |
| 187 | |
| 188 | |
| 189 | class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata], SharedTopkIndicesBuffer): |
| 190 | is_sparse = True |
| 191 | |
| 192 | def __init__( |
| 193 | self, |
| 194 | num_heads: int, |
| 195 | head_size: int, |
| 196 | scale: float, |
| 197 | num_kv_heads: int, |
| 198 | alibi_slopes: list[float] | None, |
| 199 | sliding_window: int | None, |
| 200 | kv_cache_dtype: str, |
| 201 | logits_soft_cap: float | None, |
| 202 | attn_type: str, |
| 203 | kv_sharing_target_layer_name: str | None, |
| 204 | # MLA Specific Arguments |
| 205 | topk_indices_buffer: torch.Tensor | None = None, |
| 206 | indexer: Optional["Indexer"] = None, |
| 207 | **mla_args, |
| 208 | ) -> None: |
| 209 | self.num_heads = num_heads |
| 210 | self.head_size = head_size |
| 211 | self.scale = float(scale) |
| 212 | self.num_kv_heads = num_kv_heads |
| 213 | self.kv_cache_dtype = kv_cache_dtype |
| 214 | self.kv_lora_rank: int = mla_args["kv_lora_rank"] |
| 215 | self.softmax_scale = scale |
| 216 | self.init_topk_indices_buffer(indexer, topk_indices_buffer) |
| 217 | |
| 218 | def _forward_bf16_kv( |
| 219 | self, |
| 220 | q: torch.Tensor, # [sq, heads, d_qk] |
| 221 | kv_c_and_k_pe_cache: torch.Tensor, # [blocks, heads, d_qk] |
| 222 | topk_indices: torch.Tensor, # [sq, topk] |
| 223 | attn_metadata: XPUMLASparseMetadata, |
| 224 | ) -> torch.Tensor: |
| 225 | num_tokens = q.shape[0] |
| 226 | kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view( |
| 227 | -1, 1, kv_c_and_k_pe_cache.shape[-1] |
| 228 | ) |
| 229 | |
| 230 | topk_indices = topk_indices.view(num_tokens, 1, -1) |
| 231 | |
| 232 | output, _, _ = triton_bf16_mla_sparse_interface( |
| 233 | q, |
| 234 | kv_c_and_k_pe_cache, |
| 235 | topk_indices, |
| 236 | sm_scale=self.softmax_scale, |
| 237 | ) |
| 238 | |
| 239 | return output[:, : self.num_heads, :] |
| 240 | |
| 241 | def forward_mqa( |
| 242 | self, |
| 243 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 244 | kv_c_and_k_pe_cache: torch.Tensor, |
| 245 | attn_metadata: XPUMLASparseMetadata, |
| 246 | layer: AttentionLayer, |
| 247 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 248 | # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use |
| 249 | # MQA 576/512 approach for both prefill and decode |
| 250 | |
| 251 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 252 | raise NotImplementedError("FP8 kv is not supported with XPU MLA Sparse yet") |
| 253 | |
| 254 | # Concatenate q if it's a tuple (ql_nope, q_pe) |
| 255 | if isinstance(q, tuple): |
| 256 | q = torch.cat(q, dim=-1) |
| 257 | |
| 258 | num_actual_toks = q.shape[0] |
| 259 | |
| 260 | assert self.topk_indices_buffer is not None |
| 261 | topk_indices = self.topk_indices_buffer[:num_actual_toks] |
| 262 | |
| 263 | kv_rows, block_stride_rows = flat_kv_row_view( |
| 264 | kv_c_and_k_pe_cache, attn_metadata.block_size |
| 265 | ) |
| 266 | topk_indices_global = triton_convert_req_index_to_global_index( |
| 267 | attn_metadata.req_id_per_token, |
| 268 | attn_metadata.block_table, |
| 269 | topk_indices, |
| 270 | BLOCK_SIZE=attn_metadata.block_size, |
| 271 | BLOCK_STRIDE_ROWS=block_stride_rows, |
| 272 | NUM_TOPK_TOKENS=attn_metadata.topk_tokens, |
| 273 | ) |
| 274 | |
| 275 | attn_out = self._forward_bf16_kv(q, kv_rows, topk_indices_global, attn_metadata) |
| 276 | |
| 277 | return attn_out, None |