@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e584dce52b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 21:45 UTC

Add XPU MLA Sparse backend for DeepSeek v3.2 (#33230)

Signed-off-by: Zhang, Wuxun <wuxun.zhang@intel.com>

Wuxun Zhang committed Mar 11, 2026 at 19:19 UTC e584dce52b9584ffb0fc4a1a4cd31163d4257a41
9 files changed +940 -24
docs/design/attention_backends.md
+1
@@ -214,3 +214,4 @@ configuration.
214 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
215 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
216 | `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
217 +| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
tests/kernels/attention/test_xpu_mla_sparse.py new
+118
@@ -0,0 +1,118 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
8 +
9 +
10 +# https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L7
11 +def _merge_two_lse(
12 + lse0: torch.Tensor, lse1: torch.Tensor | None, s_q: int, h_q: int
13 +) -> torch.Tensor:
14 + if lse1 is None:
15 + return lse0
16 + else:
17 + return torch.logsumexp(
18 + torch.stack([lse0.view(s_q, h_q), lse1.broadcast_to(s_q, h_q)], dim=0),
19 + dim=0,
20 + )
21 +
22 +
23 +# Adapted from https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L19
24 +def reference_mla_sparse_prefill(
25 + q: torch.Tensor,
26 + kv: torch.Tensor,
27 + indices: torch.Tensor,
28 + sm_scale: float,
29 + d_v: int,
30 + topk_length: torch.Tensor | None = None,
31 + attn_sink: torch.Tensor | None = None,
32 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
33 + """
34 + Returns:
35 + - o: [s_q, h_q, dv]
36 + - o_fp32: [s_q, h_q, dv]
37 + - max_logits: [s_q, h_q]
38 + - lse: [s_q, h_q]
39 + """
40 + s_q, h_q, d_qk = q.shape
41 + s_kv, _, _ = kv.shape
42 + _, _, topk = indices.shape
43 +
44 + indices = indices.clone().squeeze(1)
45 + if topk_length is not None:
46 + mask = torch.arange(topk, device=topk_length.device).unsqueeze(0).broadcast_to(
47 + s_q, topk
48 + ) >= topk_length.unsqueeze(1) # [s_q, topk]
49 + indices[mask] = -1
50 + invalid_mask = (indices < 0) | (indices >= s_kv) # [s_q, topk]
51 + indices[invalid_mask] = 0
52 +
53 + q = q.float()
54 + gathered_kv = (
55 + kv.index_select(dim=0, index=indices.flatten()).reshape(s_q, topk, d_qk).float()
56 + ) # [s_q, topk, d_qk]
57 + P = q @ gathered_kv.transpose(1, 2) # [s_q, h_q, topk]
58 + P *= sm_scale
59 + P[invalid_mask.unsqueeze(1).broadcast_to(P.shape)] = float("-inf")
60 +
61 + orig_lse = torch.logsumexp(P, dim=-1) # [s_q, h_q]
62 + max_logits = P.max(dim=-1).values # [s_q, h_q]
63 +
64 + lse_for_o = _merge_two_lse(orig_lse, attn_sink, s_q, h_q)
65 + if not torch.is_inference_mode_enabled():
66 + lse_for_o = lse_for_o.clone()
67 + lse_for_o[lse_for_o == float("-inf")] = float(
68 + "+inf"
69 + ) # So that corresponding O will be 0
70 + s_for_o = torch.exp(P - lse_for_o.unsqueeze(-1))
71 + out = s_for_o @ gathered_kv[..., :d_v] # [s_q, h_q, dv]
72 +
73 + lonely_q_mask = orig_lse == float("-inf") # [s_q, h_q]
74 + orig_lse[lonely_q_mask] = float("+inf")
75 + return (out.to(kv.dtype), out, max_logits, orig_lse)
76 +
77 +
78 +@pytest.mark.parametrize("device_str", ["xpu"])
79 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
80 +@pytest.mark.skipif(
81 + not torch.xpu.is_available(),
82 + reason="XPU is required",
83 +)
84 +def test_bf16_triton_sparse_mla(device_str, dtype):
85 + device = torch.device(device_str)
86 + s_q = 1
87 + s_kv = 256
88 + h_q = 64 # kernel expects multiple of 64
89 + h_kv = 1
90 + d_qk = 576
91 + d_v = 512
92 + topk = 128
93 +
94 + torch.random.manual_seed(1234)
95 +
96 + q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device)
97 + kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device)
98 + indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device)
99 + for t in range(s_q):
100 + for h in range(h_kv):
101 + i_i = torch.randperm(max(1, t))[:topk]
102 + indices[t, h, : len(i_i)] = i_i
103 +
104 + sm_scale = d_qk**-0.5
105 +
106 + out, max_logits, lse = triton_bf16_mla_sparse_interface(
107 + q, kv, indices, sm_scale, d_v
108 + )
109 + assert out.shape == (s_q, h_q, d_v)
110 + assert max_logits.shape == (s_q, h_q)
111 + assert lse.shape == (s_q, h_q)
112 +
113 + ref_out, ref_out_fp32, ref_max_logits, ref_lse = reference_mla_sparse_prefill(
114 + q, kv, indices, sm_scale, d_v
115 + )
116 + assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2)
117 + assert torch.allclose(max_logits, ref_max_logits, atol=1e-3, rtol=1e-3)
118 + assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3)
vllm/_xpu_ops.py
+245
@@ -7,6 +7,7 @@ import torch
7 from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
8
9 from vllm.logger import init_logger
10 +from vllm.platforms import current_platform
11
12 logger = init_logger(__name__)
13
@@ -157,3 +158,247 @@ class xpu_ops:
158 "get_scheduler_metadata is not implemented for xpu_ops, returning None."
159 )
160 return None
161 +
162 + @staticmethod
163 + def indexer_k_quant_and_cache(
164 + k: torch.Tensor,
165 + kv_cache: torch.Tensor,
166 + slot_mapping: torch.Tensor,
167 + quant_block_size: int,
168 + scale_fmt: str | None,
169 + ) -> None:
170 + head_dim = k.shape[-1]
171 + k = k.view(-1, head_dim) # [total_tokens, head_dim]
172 +
173 + def group_quant_torch(
174 + x: torch.Tensor,
175 + group_size: int,
176 + eps: float = 1e-10,
177 + dtype: torch.dtype | None = None,
178 + column_major_scales: bool = False,
179 + out_q: torch.Tensor | None = None,
180 + use_ue8m0: bool | None = None,
181 + ) -> tuple[torch.Tensor, torch.Tensor]:
182 + if use_ue8m0 is None:
183 + # Default fallback - could import is_deep_gemm_e8m0_used if needed
184 + use_ue8m0 = False
185 +
186 + if dtype is None:
187 + dtype = current_platform.fp8_dtype()
188 +
189 + # Validate inputs
190 + assert x.shape[-1] % group_size == 0, (
191 + f"Last dimension {x.shape[-1]} must be divisible by "
192 + f"group_size {group_size}"
193 + )
194 + assert x.stride(-1) == 1, "Input tensor groups must be contiguous"
195 +
196 + # Prepare output tensor
197 + if out_q is None:
198 + x_q = torch.empty_like(x, dtype=dtype)
199 + else:
200 + assert out_q.shape == x.shape
201 + x_q = out_q
202 +
203 + # Reshape input for group processing
204 + # Original shape: (..., last_dim)
205 + # Target shape: (..., num_groups, group_size)
206 + original_shape = x.shape
207 + num_groups = original_shape[-1] // group_size
208 +
209 + # Reshape to separate groups
210 + group_shape = original_shape[:-1] + (num_groups, group_size)
211 + x_grouped = x.view(group_shape)
212 +
213 + # Compute per-group absolute maximum values
214 + # Shape: (..., num_groups)
215 + abs_max = torch.amax(torch.abs(x_grouped), dim=-1, keepdim=False)
216 + abs_max = torch.maximum(
217 + abs_max, torch.tensor(eps, device=x.device, dtype=x.dtype)
218 + )
219 +
220 + # Compute scales
221 + FP8_MAX = torch.finfo(dtype).max
222 + FP8_MIN = torch.finfo(dtype).min
223 + scale_raw = abs_max / FP8_MAX
224 +
225 + if use_ue8m0:
226 + # For UE8M0 format, scales must be powers of 2
227 + scales = torch.pow(2.0, torch.ceil(torch.log2(scale_raw)))
228 + else:
229 + scales = scale_raw
230 +
231 + # Expand scales for broadcasting with grouped data
232 + # Shape: (..., num_groups, 1)
233 + scales_expanded = scales.unsqueeze(-1)
234 +
235 + # Quantize the grouped data
236 + x_scaled = x_grouped / scales_expanded
237 + x_clamped = torch.clamp(x_scaled, FP8_MIN, FP8_MAX)
238 + x_quantized = x_clamped.to(dtype)
239 +
240 + # Reshape back to original shape
241 + x_q.copy_(x_quantized.view(original_shape))
242 +
243 + # Prepare scales tensor in requested format
244 + if column_major_scales:
245 + # Column-major: (num_groups,) + batch_dims
246 + # Transpose the scales to put group dimension first
247 + scales_shape = (num_groups,) + original_shape[:-1]
248 + x_s = scales.permute(-1, *range(len(original_shape) - 1))
249 + x_s = x_s.contiguous().view(scales_shape)
250 + else:
251 + # Row-major: batch_dims + (num_groups,)
252 + x_s = scales.contiguous()
253 +
254 + # Ensure scales are float32
255 + return x_q, x_s.float()
256 +
257 + k_fp8, k_scale = group_quant_torch(
258 + k,
259 + group_size=quant_block_size,
260 + column_major_scales=False,
261 + use_ue8m0=(scale_fmt == "ue8m0"),
262 + )
263 +
264 + k_fp8_bytes = k_fp8.view(-1, head_dim).view(torch.uint8)
265 + scale_bytes = k_scale.view(torch.uint8).view(-1, 4)
266 + k = torch.cat(
267 + [k_fp8_bytes, scale_bytes], dim=-1
268 + ) # [total_tokens, head_dim + 4]
269 +
270 + slot_mapping = slot_mapping.flatten()
271 + # kv_cache: [num_block, block_size, head_dim + 4]
272 + kv_cache.view(-1, kv_cache.shape[-1]).index_copy_(0, slot_mapping, k)
273 +
274 + @staticmethod
275 + def cp_gather_indexer_k_quant_cache(
276 + kv_cache: torch.Tensor,
277 + dst_k: torch.Tensor,
278 + dst_scale: torch.Tensor,
279 + block_table: torch.Tensor,
280 + cu_seq_lens: torch.Tensor,
281 + ) -> None:
282 + """
283 + Args:
284 + kv_cache: [num_blocks, block_size, cache_stride] - quantized KV cache
285 + Layout per block: [k_values, scale_values]
286 + - k_values: [block_size * head_dim]
287 + - scale_values: [block_size * head_dim * 4 / quant_block_size]
288 + dst_k: [num_tokens, head_dim] - output tensor for K values
289 + dst_scale: [num_tokens, head_dim / quant_block_size * 4]
290 + - output tensor for scale values
291 + block_table: [batch_size, num_blocks] - block table for indexing
292 + cu_seq_lens: [batch_size + 1] - cumulative sequence lengths
293 + """
294 + batch_size = block_table.size(0)
295 + num_tokens = dst_k.size(0)
296 + head_dim = dst_k.size(1)
297 + cache_block_size = kv_cache.size(1)
298 + quant_block_size = head_dim * 4 // dst_scale.size(1)
299 +
300 + # For each token, find which batch it belongs to using searchsorted
301 + token_indices = torch.arange(num_tokens, device=dst_k.device) + 1
302 + # cu_seq_lens is [batch_size + 1], we need to find which interval each
303 + # token belongs to
304 + batch_indices = torch.searchsorted(cu_seq_lens, token_indices) - 1
305 + batch_indices = torch.clamp(batch_indices, 0, batch_size - 1)
306 +
307 + # Calculate the in-batch sequence index for each token
308 + inbatch_seq_indices = token_indices - cu_seq_lens[batch_indices]
309 +
310 + # Find which block each token belongs to
311 + block_indices_in_table = inbatch_seq_indices // cache_block_size
312 + physical_block_indices = block_table[batch_indices, block_indices_in_table]
313 +
314 + # Calculate the offset within each block
315 + inblock_offsets = (inbatch_seq_indices - 1) % cache_block_size
316 +
317 + # Calculate strides
318 + block_stride = kv_cache.stride(0) # stride for each block
319 +
320 + # Flatten kv_cache for easier indexing
321 + kv_cache_flat = kv_cache.view(-1)
322 +
323 + # Calculate source offset for K values for all tokens (vectorized)
324 + src_block_offsets = physical_block_indices * block_stride
325 + src_k_offsets = src_block_offsets + inblock_offsets * head_dim
326 +
327 + # Gather K values using advanced indexing
328 + # Create indices for all elements we need to gather
329 + k_indices = src_k_offsets.unsqueeze(1) + torch.arange(
330 + head_dim, device=dst_k.device
331 + )
332 + dst_k[:] = kv_cache_flat[k_indices]
333 +
334 + # Calculate source offset for scale values (vectorized)
335 + # Scales are stored after all K values for each block
336 + scale_size = head_dim * 4 // quant_block_size
337 + src_scale_offsets = src_block_offsets + head_dim + inblock_offsets * scale_size
338 +
339 + # Gather scale values
340 + scale_indices = src_scale_offsets.unsqueeze(1) + torch.arange(
341 + scale_size, device=dst_scale.device
342 + )
343 + dst_scale[:] = kv_cache_flat[scale_indices]
344 +
345 + @staticmethod
346 + def top_k_per_row_prefill(
347 + logits: torch.Tensor,
348 + cu_seqlen_ks: torch.Tensor,
349 + cu_seqlen_ke: torch.Tensor,
350 + raw_topk_indices: torch.Tensor,
351 + num_rows: int,
352 + stride0: int,
353 + strdide1: int,
354 + topk_tokens: int,
355 + ) -> torch.Tensor:
356 + real_topk = min(topk_tokens, logits.shape[-1])
357 + topk_indices = logits.topk(real_topk, dim=-1)[1].to(torch.int32)
358 + topk_indices -= cu_seqlen_ks[:, None]
359 + mask_lo = topk_indices >= 0
360 + mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0
361 + mask = torch.full_like(
362 + topk_indices, False, dtype=torch.bool, device=topk_indices.device
363 + )
364 + mask = mask_lo & mask_hi
365 + topk_indices.masked_fill_(~mask, -1)
366 + raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
367 + topk_indices
368 + )
369 +
370 + @staticmethod
371 + def top_k_per_row_decode(
372 + logits: torch.Tensor,
373 + next_n: int,
374 + seq_lens: torch.Tensor,
375 + raw_topk_indices: torch.Tensor,
376 + num_rows: int,
377 + stride0: int,
378 + stride1: int,
379 + topk_tokens: int,
380 + ) -> torch.Tensor:
381 + device = logits.device
382 + batch_size = seq_lens.size(0)
383 + # padded query len
384 + padded_num_tokens = batch_size * next_n
385 + positions = (
386 + torch.arange(logits.shape[-1], device=device)
387 + .unsqueeze(0)
388 + .expand(batch_size * next_n, -1)
389 + )
390 + row_indices = torch.arange(padded_num_tokens, device=device) // next_n
391 + next_n_offset = torch.arange(padded_num_tokens, device=device) % next_n
392 + index_end_pos = (seq_lens[row_indices] - next_n + next_n_offset).unsqueeze(1)
393 + # index_end_pos: [B * N, 1]
394 + mask = positions <= index_end_pos
395 + # mask: [B * N, L]
396 + logits = logits.masked_fill(~mask, float("-inf"))
397 + topk_indices = logits.topk(topk_tokens, dim=-1)[1].to(torch.int32) # [B * N, K]
398 + # ensure we don't set indices for the top k
399 + # that is out of range(masked already)
400 + # this will happen if context length is shorter than K
401 + topk_indices[topk_indices > index_end_pos] = -1
402 + raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
403 + topk_indices
404 + )
vllm/model_executor/layers/sparse_attn_indexer.py
+47 -22
@@ -135,16 +135,29 @@ def sparse_attn_indexer(
135 topk_indices = topk_indices_buffer[
136 chunk.token_start : chunk.token_end, :topk_tokens
137 ]
138 - torch.ops._C.top_k_per_row_prefill(
139 - logits,
140 - chunk.cu_seqlen_ks,
141 - chunk.cu_seqlen_ke,
142 - topk_indices,
143 - num_rows,
144 - logits.stride(0),
145 - logits.stride(1),
146 - topk_tokens,
147 - )
138 +
139 + if current_platform.is_xpu():
140 + ops.top_k_per_row_prefill(
141 + logits,
142 + chunk.cu_seqlen_ks,
143 + chunk.cu_seqlen_ke,
144 + topk_indices,
145 + num_rows,
146 + logits.stride(0),
147 + logits.stride(1),
148 + topk_tokens,
149 + )
150 + else:
151 + torch.ops._C.top_k_per_row_prefill(
152 + logits,
153 + chunk.cu_seqlen_ks,
154 + chunk.cu_seqlen_ke,
155 + topk_indices,
156 + num_rows,
157 + logits.stride(0),
158 + logits.stride(1),
159 + topk_tokens,
160 + )
161
162 # Compute lengths from row spans
163 # lengths = (chunk.cu_seqlen_ke - chunk.cu_seqlen_ks).to(torch.int32)
@@ -220,16 +233,28 @@ def sparse_attn_indexer(
233 None,
234 )
235 else:
223 - torch.ops._C.top_k_per_row_decode(
224 - logits,
225 - next_n,
226 - decode_metadata.seq_lens,
227 - topk_indices,
228 - num_rows,
229 - logits.stride(0),
230 - logits.stride(1),
231 - topk_tokens,
232 - )
236 + if current_platform.is_xpu():
237 + ops.top_k_per_row_decode(
238 + logits,
239 + next_n,
240 + decode_metadata.seq_lens,
241 + topk_indices,
242 + num_rows,
243 + logits.stride(0),
244 + logits.stride(1),
245 + topk_tokens,
246 + )
247 + else:
248 + torch.ops._C.top_k_per_row_decode(
249 + logits,
250 + next_n,
251 + decode_metadata.seq_lens,
252 + topk_indices,
253 + num_rows,
254 + logits.stride(0),
255 + logits.stride(1),
256 + topk_tokens,
257 + )
258
259 if decode_metadata.requires_padding:
260 # if padded, we need to unpack
@@ -320,14 +345,14 @@ class SparseAttnIndexer(CustomOp):
345 k: torch.Tensor,
346 weights: torch.Tensor,
347 ):
323 - if current_platform.is_cuda():
348 + if current_platform.is_cuda() or current_platform.is_xpu():
349 return self.forward_cuda(hidden_states, q_fp8, k, weights)
350 elif current_platform.is_rocm():
351 return self.forward_hip(hidden_states, q_fp8, k, weights)
352 else:
353 raise NotImplementedError(
354 "SparseAttnIndexer native forward is only implemented for "
330 - "CUDA and ROCm platform."
355 + "CUDA, ROCm and XPU platforms."
356 )
357
358 def forward_cuda(
vllm/platforms/xpu.py
+2 -1
@@ -61,7 +61,8 @@ class XPUPlatform(Platform):
61
62 dtype = attn_selector_config.dtype
63 if attn_selector_config.use_sparse:
64 - raise NotImplementedError("Sparse Attention is not supported on XPU.")
64 + logger.info_once("Using XPU MLA Sparse backend.")
65 + return AttentionBackendEnum.XPU_MLA_SPARSE.get_path()
66 if attn_selector_config.use_mla:
67 logger.info_once("Using Triton MLA backend on V1 engine.")
68 return AttentionBackendEnum.TRITON_MLA.get_path()
vllm/triton_utils/__init__.py
+4 -1
@@ -17,4 +17,7 @@ else:
17 tl = TritonLanguagePlaceholder()
18 tldevice = TritonLanguagePlaceholder()
19
20 -__all__ = ["HAS_TRITON", "triton", "tl", "tldevice"]
20 +LOG2E = 1.4426950408889634
21 +LOGE2 = 0.6931471805599453
22 +
23 +__all__ = ["HAS_TRITON", "triton", "tl", "tldevice", "LOG2E", "LOGE2"]
vllm/v1/attention/backends/mla/xpu_mla_sparse.py new
+257
@@ -0,0 +1,257 @@
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.v1.attention.backend import (
17 + AttentionBackend,
18 + AttentionCGSupport,
19 + AttentionLayer,
20 + AttentionMetadata,
21 + AttentionMetadataBuilder,
22 + CommonAttentionMetadata,
23 + SparseMLAAttentionImpl,
24 +)
25 +from vllm.v1.attention.backends.mla.flashmla_sparse import (
26 + triton_convert_req_index_to_global_index,
27 +)
28 +from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
29 +from vllm.v1.kv_cache_interface import AttentionSpec
30 +
31 +if TYPE_CHECKING:
32 + from vllm.model_executor.models.deepseek_v2 import Indexer
33 +logger = init_logger(__name__)
34 +
35 +
36 +class XPUMLASparseBackend(AttentionBackend):
37 + accept_output_buffer: bool = True
38 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
39 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
40 + "auto",
41 + "bfloat16",
42 + ]
43 +
44 + @staticmethod
45 + def get_name() -> str:
46 + return "XPU_MLA_SPARSE"
47 +
48 + @staticmethod
49 + def get_metadata_cls() -> type["XPUMLASparseMetadata"]:
50 + return XPUMLASparseMetadata
51 +
52 + @staticmethod
53 + def get_builder_cls() -> type["XPUMLASparseMetadataBuilder"]:
54 + return XPUMLASparseMetadataBuilder
55 +
56 + @staticmethod
57 + def get_impl_cls() -> type["XPUMLASparseImpl"]:
58 + return XPUMLASparseImpl
59 +
60 + @classmethod
61 + def is_mla(cls) -> bool:
62 + return True
63 +
64 + @classmethod
65 + def is_sparse(cls) -> bool:
66 + return True
67 +
68 + @staticmethod
69 + def get_kv_cache_shape(
70 + num_blocks: int,
71 + block_size: int,
72 + num_kv_heads: int, # assumed to be 1 for MLA
73 + head_size: int,
74 + cache_dtype_str: str = "auto",
75 + ) -> tuple[int, ...]:
76 + return (num_blocks, block_size, head_size)
77 +
78 + @classmethod
79 + def get_supported_head_sizes(cls) -> list[int]:
80 + return [576]
81 +
82 +
83 +@dataclass
84 +class XPUMLASparseMetadata(AttentionMetadata):
85 + num_reqs: int
86 + max_query_len: int
87 + max_seq_len: int
88 +
89 + num_actual_tokens: int # Number of tokens excluding padding.
90 + query_start_loc: torch.Tensor
91 + slot_mapping: torch.Tensor
92 +
93 + block_table: torch.Tensor
94 + req_id_per_token: torch.Tensor
95 +
96 + block_size: int = 1
97 + topk_tokens: int = 2048
98 +
99 +
100 +@dataclass
101 +class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]):
102 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
103 +
104 + def __init__(
105 + self,
106 + kv_cache_spec: AttentionSpec,
107 + layer_names: list[str],
108 + vllm_config: VllmConfig,
109 + device: torch.device,
110 + ):
111 + self.kv_cache_spec = kv_cache_spec
112 + self.model_config = vllm_config.model_config
113 + parallel_config = vllm_config.parallel_config
114 + self.device = device
115 + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
116 +
117 + self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
118 + self.mla_dims = get_mla_dims(self.model_config)
119 + self.topk_tokens = vllm_config.model_config.hf_config.index_topk
120 + self.topk_tokens_tensor = torch.tensor(
121 + [self.topk_tokens], device=device, dtype=torch.int32
122 + )
123 + self.max_model_len_tensor = torch.tensor(
124 + [self.model_config.max_model_len], device=device, dtype=torch.int32
125 + )
126 + # this is ignored by `flash_mla_with_kvcache` if indices not None
127 + self.dummy_block_table = torch.empty(
128 + (1, 1), dtype=torch.int32, device=self.device
129 + )
130 +
131 + self.req_id_per_token_buffer = torch.empty(
132 + (max_num_batched_tokens,),
133 + dtype=torch.int32,
134 + device=device,
135 + )
136 +
137 + def build(
138 + self,
139 + common_prefix_len: int,
140 + common_attn_metadata: CommonAttentionMetadata,
141 + fast_build: bool = False,
142 + ) -> XPUMLASparseMetadata:
143 + num_tokens = common_attn_metadata.num_actual_tokens
144 + starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
145 + seg_lengths = np.diff(starts)
146 + req_id_per_token = np.repeat(
147 + np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
148 + )
149 + # Zero-fill for cudagraphs
150 + self.req_id_per_token_buffer.fill_(0)
151 + self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
152 + torch.from_numpy(req_id_per_token), non_blocking=True
153 + )
154 +
155 + req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
156 +
157 + metadata = XPUMLASparseMetadata(
158 + num_reqs=common_attn_metadata.num_reqs,
159 + max_query_len=common_attn_metadata.max_query_len,
160 + max_seq_len=common_attn_metadata.max_seq_len,
161 + num_actual_tokens=common_attn_metadata.num_actual_tokens,
162 + query_start_loc=common_attn_metadata.query_start_loc,
163 + slot_mapping=common_attn_metadata.slot_mapping,
164 + block_table=common_attn_metadata.block_table_tensor,
165 + req_id_per_token=req_id_per_token,
166 + block_size=self.kv_cache_spec.block_size,
167 + topk_tokens=self.topk_tokens,
168 + )
169 + return metadata
170 +
171 +
172 +class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):
173 + def __init__(
174 + self,
175 + num_heads: int,
176 + head_size: int,
177 + scale: float,
178 + num_kv_heads: int,
179 + alibi_slopes: list[float] | None,
180 + sliding_window: int | None,
181 + kv_cache_dtype: str,
182 + logits_soft_cap: float | None,
183 + attn_type: str,
184 + kv_sharing_target_layer_name: str | None,
185 + # MLA Specific Arguments
186 + topk_indice_buffer: torch.Tensor | None = None,
187 + indexer: Optional["Indexer"] = None,
188 + **mla_args,
189 + ) -> None:
190 + self.num_heads = num_heads
191 + self.head_size = head_size
192 + self.scale = float(scale)
193 + self.num_kv_heads = num_kv_heads
194 + self.kv_cache_dtype = kv_cache_dtype
195 + self.kv_lora_rank: int = mla_args["kv_lora_rank"]
196 + self.softmax_scale = scale
197 + assert indexer is not None
198 + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
199 +
200 + def _forward_bf16_kv(
201 + self,
202 + q: torch.Tensor, # [sq, heads, d_qk]
203 + kv_c_and_k_pe_cache: torch.Tensor, # [blocks, heads, d_qk]
204 + topk_indices: torch.Tensor, # [sq, topk]
205 + attn_metadata: XPUMLASparseMetadata,
206 + ) -> torch.Tensor:
207 + num_tokens = q.shape[0]
208 + kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
209 + -1, 1, kv_c_and_k_pe_cache.shape[-1]
210 + )
211 +
212 + topk_indices = topk_indices.view(num_tokens, 1, -1)
213 +
214 + output, _, _ = triton_bf16_mla_sparse_interface(
215 + q,
216 + kv_c_and_k_pe_cache,
217 + topk_indices,
218 + sm_scale=self.softmax_scale,
219 + )
220 +
221 + return output[:, : self.num_heads, :]
222 +
223 + def forward_mqa(
224 + self,
225 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
226 + kv_c_and_k_pe_cache: torch.Tensor,
227 + attn_metadata: XPUMLASparseMetadata,
228 + layer: AttentionLayer,
229 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
230 + # NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use
231 + # MQA 576/512 approach for both prefill and decode
232 +
233 + if self.kv_cache_dtype.startswith("fp8"):
234 + raise NotImplementedError("FP8 kv is not supported with XPU MLA Sparse yet")
235 +
236 + # Concatenate q if it's a tuple (ql_nope, q_pe)
237 + if isinstance(q, tuple):
238 + q = torch.cat(q, dim=-1)
239 +
240 + num_actual_toks = q.shape[0]
241 +
242 + assert self.topk_indices_buffer is not None
243 + topk_indices = self.topk_indices_buffer[:num_actual_toks]
244 +
245 + topk_indices_global = triton_convert_req_index_to_global_index(
246 + attn_metadata.req_id_per_token,
247 + attn_metadata.block_table,
248 + topk_indices,
249 + BLOCK_SIZE=attn_metadata.block_size,
250 + NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
251 + )
252 +
253 + attn_out = self._forward_bf16_kv(
254 + q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata
255 + )
256 +
257 + return attn_out, None
vllm/v1/attention/backends/registry.py
+1
@@ -57,6 +57,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
57 ROCM_AITER_MLA_SPARSE = (
58 "vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse.ROCMAiterMLASparseBackend"
59 )
60 + XPU_MLA_SPARSE = "vllm.v1.attention.backends.mla.xpu_mla_sparse.XPUMLASparseBackend"
61 TORCH_SDPA = "" # this tag is only used for ViT
62 FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
63 FLASHINFER_MLA = (
vllm/v1/attention/ops/xpu_mla_sparse.py new
+265
@@ -0,0 +1,265 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.triton_utils import LOG2E, LOGE2, tl, triton
7 +
8 +
9 +@triton.jit
10 +def _bf16_mla_sparse_kernel(
11 + q_buffer,
12 + k_buffer,
13 + v_buffer,
14 + indices_ptr,
15 + out_ptr,
16 + softmax_lse_ptr,
17 + max_logits_ptr,
18 + seq_q,
19 + seq_kv,
20 + h_q,
21 + dim_qk,
22 + dim_v,
23 + stride_q_token,
24 + stride_q_head,
25 + stride_k_token,
26 + stride_k_head,
27 + stride_v_token,
28 + stride_v_head,
29 + stride_out_token,
30 + stride_out_head,
31 + stride_lse,
32 + stride_indices_token,
33 + stride_indices_head,
34 + sm_scale,
35 + kv_group_num: tl.constexpr,
36 + index_topk: tl.constexpr,
37 + BLOCK_H: tl.constexpr, # block size for num heads
38 + BLOCK_M: tl.constexpr, # block size for num tokens
39 + BLOCK_N: tl.constexpr, # block size for indices
40 + BLOCK_DV: tl.constexpr, # block size for dim_v
41 + BLOCK_DMODEL: tl.constexpr, # block size for dim_nope
42 + BLOCK_DPE: tl.constexpr, # block size for positional embedding
43 + LOGE2: tl.constexpr,
44 +):
45 + cur_q = tl.program_id(0)
46 + cur_head_id = tl.program_id(1)
47 + cur_kv_head_id = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
48 +
49 + VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num
50 + cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
51 + mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
52 + mask_h = mask_h & (cur_head < h_q)
53 +
54 + offs_d = tl.arange(0, BLOCK_DMODEL)
55 + offs_dv = tl.arange(0, BLOCK_DV)
56 +
57 + off_q = cur_q * stride_q_token + cur_head[:, None] * stride_q_head + offs_d[None, :]
58 + mask_dmodel = offs_d < BLOCK_DMODEL
59 + q = tl.load(
60 + q_buffer + off_q, mask=(mask_h[:, None]) & (mask_dmodel[None, :]), other=0.0
61 + )
62 +
63 + if BLOCK_DPE > 0:
64 + offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
65 + off_qpe = (
66 + cur_q * stride_q_token
67 + + cur_head[:, None] * stride_q_head
68 + + offs_dpe[None, :]
69 + )
70 + # assume dim_qk == BLOCK_DMODEL + BLOCK_DPE
71 + mask_dpe = offs_dpe < dim_qk
72 + qpe = tl.load(
73 + q_buffer + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
74 + )
75 +
76 + e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
77 + e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
78 + acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
79 +
80 + for start_indice in range(0, index_topk, BLOCK_N):
81 + offs_indice = start_indice + tl.arange(0, BLOCK_N)
82 + mask_indice = offs_indice < index_topk
83 + indices = tl.load(
84 + indices_ptr
85 + + (
86 + cur_q * stride_indices_token
87 + + cur_kv_head_id * stride_indices_head
88 + + offs_indice
89 + ),
90 + mask=mask_indice,
91 + other=-1,
92 + )
93 +
94 + mask_kv = (indices >= 0) & (indices < seq_kv)
95 + mask_kv_d = mask_dmodel
96 + offs_k = (
97 + indices[None, :] * stride_k_token
98 + + cur_kv_head_id * stride_k_head
99 + + offs_d[:, None]
100 + )
101 +
102 + # q_nope @ k_nope
103 + k = tl.load(
104 + k_buffer + offs_k, mask=(mask_kv[None, :]) & (mask_kv_d[:, None]), other=0.0
105 + )
106 + qk = tl.dot(q, k.to(q.dtype))
107 +
108 + if BLOCK_DPE > 0:
109 + # q_rope @ k_rope
110 + offs_kpe = (
111 + indices[None, :] * stride_k_token
112 + + cur_kv_head_id * stride_k_head
113 + + offs_dpe[:, None]
114 + )
115 + mask_k_dpe = offs_dpe < dim_qk
116 + kpe = tl.load(
117 + k_buffer + offs_kpe,
118 + mask=(mask_kv[None, :]) & (mask_k_dpe[:, None]),
119 + other=0.0,
120 + )
121 + qk += tl.dot(qpe, kpe.to(q.dtype))
122 +
123 + # apply scaling
124 + qk *= sm_scale
125 + qk = tl.where((mask_h[:, None]) & (mask_kv[None, :]), qk, -float("inf"))
126 +
127 + # load v
128 + mask_v_d = offs_dv < dim_v
129 + offs_v = (
130 + indices[:, None] * stride_v_token
131 + + cur_kv_head_id * stride_v_head
132 + + offs_dv[None, :]
133 + )
134 + v = tl.load(
135 + v_buffer + offs_v, mask=(mask_kv[:, None]) & (mask_v_d[None, :]), other=0.0
136 + )
137 +
138 + # online softmax
139 + n_e_max = tl.maximum(tl.max(qk, 1), e_max)
140 + re_scale = tl.exp2(e_max - n_e_max)
141 + p = tl.exp2(qk - n_e_max[:, None])
142 + acc *= re_scale[:, None]
143 +
144 + # score @ v
145 + acc += tl.dot(p.to(v.dtype), v)
146 +
147 + # update global sum and max
148 + e_sum = e_sum * re_scale + tl.sum(p, 1)
149 + e_max = n_e_max
150 +
151 + # rescaling
152 + acc /= e_sum[:, None]
153 +
154 + max_logits = e_max * LOGE2
155 + # calculate lse
156 + lse = max_logits + tl.log2(e_sum) * LOGE2
157 +
158 + # write output
159 + offs_o = (
160 + cur_q * stride_out_token
161 + + cur_head[:, None] * stride_out_head
162 + + offs_dv[None, :]
163 + )
164 + mask_out_d = offs_dv < dim_v
165 + tl.store(
166 + out_ptr + offs_o,
167 + acc.to(tl.bfloat16),
168 + mask=(mask_h[:, None]) & (mask_out_d[None, :]),
169 + )
170 +
171 + offs_lse = cur_q * stride_lse + cur_head
172 + tl.store(softmax_lse_ptr + offs_lse, lse, mask=mask_h)
173 + tl.store(max_logits_ptr + offs_lse, max_logits, mask=mask_h)
174 +
175 +
176 +# reference implementation of bf16 sparse prefill kernel
177 +def triton_bf16_mla_sparse_interface(
178 + q: torch.Tensor, # [num_tokens, num_heads_q, dim_qk]
179 + kv: torch.Tensor, # [num_tokens, num_heads_kv, dim_qk]
180 + indices: torch.Tensor, # [num_tokens, num_heads_kv, topk]
181 + sm_scale: float,
182 + d_v: int = 512,
183 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
184 + """
185 + out : [num_tokens, num_heads_q, d_v]
186 + max_logits : [num_tokens, num_heads_q]
187 + lse : logsumexp, [num_tokens, num_heads_q]
188 + """
189 + num_tokens, num_heads_q, dim_qk = q.shape
190 + _, num_heads_kv, _ = kv.shape
191 + assert dim_qk == kv.shape[2], "q and kv have different head dimensions"
192 +
193 + # for deepseek v3.2, index topk should be 2048
194 + _, _, index_topk = indices.shape
195 +
196 + BLOCK_H = 16
197 + BLOCK_DMODEL = 512
198 + BLOCK_DPE = 64
199 + BLOCK_M = 32
200 + BLOCK_N = 16
201 + BLOCK_DV = 512
202 + assert d_v == BLOCK_DV, "only support d_v = 512"
203 +
204 + assert dim_qk == BLOCK_DMODEL + BLOCK_DPE, (
205 + "dim_qk does not match BLOCK_DMODEL + BLOCK_DPE"
206 + )
207 + assert num_heads_kv == 1, "only support kv head = 1 for now"
208 + assert index_topk % BLOCK_N == 0, "index_topk must be multiple of BLOCK_N"
209 +
210 + sm_scale *= LOG2E
211 +
212 + kv_group_num = num_heads_q // num_heads_kv
213 + grid = (
214 + num_tokens,
215 + triton.cdiv(num_heads_q, min(BLOCK_H, kv_group_num)),
216 + )
217 +
218 + out = torch.zeros((num_tokens, num_heads_q, d_v), dtype=q.dtype, device=q.device)
219 + softmax_lse = torch.zeros(
220 + (num_tokens, num_heads_q), dtype=torch.float32, device=q.device
221 + )
222 + max_logits = torch.zeros(
223 + (num_tokens, num_heads_q), dtype=torch.float32, device=q.device
224 + )
225 +
226 + k = kv
227 + v = kv[..., :d_v]
228 +
229 + _bf16_mla_sparse_kernel[grid](
230 + q_buffer=q,
231 + k_buffer=k,
232 + v_buffer=v,
233 + indices_ptr=indices,
234 + out_ptr=out,
235 + softmax_lse_ptr=softmax_lse,
236 + max_logits_ptr=max_logits,
237 + seq_q=num_tokens,
238 + seq_kv=kv.shape[0],
239 + h_q=num_heads_q,
240 + dim_qk=dim_qk,
241 + dim_v=d_v,
242 + stride_q_token=q.stride(0),
243 + stride_q_head=q.stride(1),
244 + stride_k_token=k.stride(0),
245 + stride_k_head=k.stride(1),
246 + stride_v_token=v.stride(0),
247 + stride_v_head=v.stride(1),
248 + stride_out_token=out.stride(0),
249 + stride_out_head=out.stride(1),
250 + stride_lse=softmax_lse.stride(0),
251 + stride_indices_token=indices.stride(0),
252 + stride_indices_head=indices.stride(1),
253 + sm_scale=sm_scale,
254 + kv_group_num=kv_group_num,
255 + index_topk=index_topk,
256 + BLOCK_H=BLOCK_H,
257 + BLOCK_M=BLOCK_M,
258 + BLOCK_N=BLOCK_N,
259 + BLOCK_DV=BLOCK_DV,
260 + BLOCK_DMODEL=BLOCK_DMODEL,
261 + BLOCK_DPE=BLOCK_DPE,
262 + LOGE2=LOGE2,
263 + )
264 +
265 + return out, max_logits, softmax_lse