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 3, 00:30 UTC
[XPU] Enable multiple key kernels for sparse attention (#37888)
Signed-off-by: Xiaochang Wu <xiaochang.wu@intel.com> Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Xiaochang Wu committed
May 22, 2026 at 12:02 UTC
8c8b1825eb26c1ffae776baaab16f2eebf92b7d3
3 files changed
+200
-310
vllm/_custom_ops.py
+44
@@ -2826,6 +2826,50 @@ def indexer_k_quant_and_cache(
2826
)
2827
2828
2829
+def top_k_per_row_prefill(
2830
+ logits: torch.Tensor,
2831
+ cu_seqlen_ks: torch.Tensor,
2832
+ cu_seqlen_ke: torch.Tensor,
2833
+ raw_topk_indices: torch.Tensor,
2834
+ num_rows: int,
2835
+ stride0: int,
2836
+ stride1: int,
2837
+ topk_tokens: int,
2838
+) -> None:
2839
+ torch.ops._C.top_k_per_row_prefill(
2840
+ logits,
2841
+ cu_seqlen_ks,
2842
+ cu_seqlen_ke,
2843
+ raw_topk_indices,
2844
+ num_rows,
2845
+ stride0,
2846
+ stride1,
2847
+ topk_tokens,
2848
+ )
2849
+
2850
+
2851
+def top_k_per_row_decode(
2852
+ logits: torch.Tensor,
2853
+ next_n: int,
2854
+ seq_lens: torch.Tensor,
2855
+ raw_topk_indices: torch.Tensor,
2856
+ num_rows: int,
2857
+ stride0: int,
2858
+ stride1: int,
2859
+ topk_tokens: int,
2860
+) -> None:
2861
+ torch.ops._C.top_k_per_row_decode(
2862
+ logits,
2863
+ next_n,
2864
+ seq_lens,
2865
+ raw_topk_indices,
2866
+ num_rows,
2867
+ stride0,
2868
+ stride1,
2869
+ topk_tokens,
2870
+ )
2871
+
2872
+
2873
def cp_gather_indexer_k_quant_cache(
2874
kv_cache: torch.Tensor,
2875
dst_k: torch.Tensor,
vllm/_xpu_ops.py
+82
-245
@@ -185,6 +185,76 @@ def _xpu_ops_deepseek_scaling_rope_fake(
185
return query, key
186
187
188
+def _xpu_fp8_mqa_logits_impl(
189
+ q: torch.Tensor,
190
+ k_quant: torch.Tensor,
191
+ k_scale: torch.Tensor,
192
+ weights: torch.Tensor,
193
+ cu_seqlen_ks: torch.Tensor,
194
+ cu_seqlen_ke: torch.Tensor,
195
+) -> torch.Tensor:
196
+ return torch.ops._xpu_C.fp8_mqa_logits(
197
+ q,
198
+ k_quant,
199
+ k_scale,
200
+ weights,
201
+ cu_seqlen_ks,
202
+ cu_seqlen_ke,
203
+ )
204
+
205
+
206
+def _xpu_fp8_mqa_logits_fake(
207
+ q: torch.Tensor,
208
+ k_quant: torch.Tensor,
209
+ k_scale: torch.Tensor,
210
+ weights: torch.Tensor,
211
+ cu_seqlen_ks: torch.Tensor,
212
+ cu_seqlen_ke: torch.Tensor,
213
+) -> torch.Tensor:
214
+ return torch.empty(
215
+ (q.shape[0], k_quant.shape[0]),
216
+ dtype=torch.float32,
217
+ device=q.device,
218
+ )
219
+
220
+
221
+def _xpu_fp8_paged_mqa_logits_impl(
222
+ q: torch.Tensor,
223
+ kv_cache: torch.Tensor,
224
+ weights: torch.Tensor,
225
+ context_lens: torch.Tensor,
226
+ block_tables: torch.Tensor,
227
+ schedule_metadata: torch.Tensor,
228
+ max_model_len: int,
229
+) -> torch.Tensor:
230
+ return torch.ops._xpu_C.fp8_paged_mqa_logits(
231
+ q,
232
+ kv_cache,
233
+ weights,
234
+ context_lens,
235
+ block_tables,
236
+ schedule_metadata,
237
+ max_model_len,
238
+ )
239
+
240
+
241
+def _xpu_fp8_paged_mqa_logits_fake(
242
+ q: torch.Tensor,
243
+ kv_cache: torch.Tensor,
244
+ weights: torch.Tensor,
245
+ context_lens: torch.Tensor,
246
+ block_tables: torch.Tensor,
247
+ schedule_metadata: torch.Tensor,
248
+ max_model_len: int,
249
+) -> torch.Tensor:
250
+ batch_size, next_n = q.shape[:2]
251
+ return torch.empty(
252
+ (batch_size * next_n, max_model_len),
253
+ dtype=torch.float32,
254
+ device=q.device,
255
+ )
256
+
257
+
258
def _topk_topp_sample_impl(
259
random_sampled: torch.Tensor,
260
logits_to_return: torch.Tensor | None,
@@ -438,251 +508,6 @@ class xpu_ops:
508
)
509
return None
510
441
- @staticmethod
442
- def indexer_k_quant_and_cache(
443
- k: torch.Tensor,
444
- kv_cache: torch.Tensor,
445
- slot_mapping: torch.Tensor,
446
- quant_block_size: int,
447
- scale_fmt: str | None,
448
- ) -> None:
449
- head_dim = k.shape[-1]
450
- k = k.view(-1, head_dim) # [total_tokens, head_dim]
451
-
452
- def group_quant_torch(
453
- x: torch.Tensor,
454
- group_size: int,
455
- eps: float = 1e-10,
456
- dtype: torch.dtype | None = None,
457
- column_major_scales: bool = False,
458
- out_q: torch.Tensor | None = None,
459
- use_ue8m0: bool | None = None,
460
- ) -> tuple[torch.Tensor, torch.Tensor]:
461
- if use_ue8m0 is None:
462
- # Default fallback - could import is_deep_gemm_e8m0_used if needed
463
- use_ue8m0 = False
464
-
465
- if dtype is None:
466
- dtype = current_platform.fp8_dtype()
467
-
468
- # Validate inputs
469
- assert x.shape[-1] % group_size == 0, (
470
- f"Last dimension {x.shape[-1]} must be divisible by "
471
- f"group_size {group_size}"
472
- )
473
- assert x.stride(-1) == 1, "Input tensor groups must be contiguous"
474
-
475
- # Prepare output tensor
476
- if out_q is None:
477
- x_q = torch.empty_like(x, dtype=dtype)
478
- else:
479
- assert out_q.shape == x.shape
480
- x_q = out_q
481
-
482
- # Reshape input for group processing
483
- # Original shape: (..., last_dim)
484
- # Target shape: (..., num_groups, group_size)
485
- original_shape = x.shape
486
- num_groups = original_shape[-1] // group_size
487
-
488
- # Reshape to separate groups
489
- group_shape = original_shape[:-1] + (num_groups, group_size)
490
- x_grouped = x.view(group_shape)
491
-
492
- # Compute per-group absolute maximum values
493
- # Shape: (..., num_groups)
494
- abs_max = torch.amax(torch.abs(x_grouped), dim=-1, keepdim=False)
495
- abs_max = torch.maximum(
496
- abs_max, torch.tensor(eps, device=x.device, dtype=x.dtype)
497
- )
498
-
499
- # Compute scales
500
- FP8_MAX = torch.finfo(dtype).max
501
- FP8_MIN = torch.finfo(dtype).min
502
- scale_raw = abs_max / FP8_MAX
503
-
504
- if use_ue8m0:
505
- # For UE8M0 format, scales must be powers of 2
506
- scales = torch.pow(2.0, torch.ceil(torch.log2(scale_raw)))
507
- else:
508
- scales = scale_raw
509
-
510
- # Expand scales for broadcasting with grouped data
511
- # Shape: (..., num_groups, 1)
512
- scales_expanded = scales.unsqueeze(-1)
513
-
514
- # Quantize the grouped data
515
- x_scaled = x_grouped / scales_expanded
516
- x_clamped = torch.clamp(x_scaled, FP8_MIN, FP8_MAX)
517
- x_quantized = x_clamped.to(dtype)
518
-
519
- # Reshape back to original shape
520
- x_q.copy_(x_quantized.view(original_shape))
521
-
522
- # Prepare scales tensor in requested format
523
- if column_major_scales:
524
- # Column-major: (num_groups,) + batch_dims
525
- # Transpose the scales to put group dimension first
526
- scales_shape = (num_groups,) + original_shape[:-1]
527
- x_s = scales.permute(-1, *range(len(original_shape) - 1))
528
- x_s = x_s.contiguous().view(scales_shape)
529
- else:
530
- # Row-major: batch_dims + (num_groups,)
531
- x_s = scales.contiguous()
532
-
533
- # Ensure scales are float32
534
- return x_q, x_s.float()
535
-
536
- k_fp8, k_scale = group_quant_torch(
537
- k,
538
- group_size=quant_block_size,
539
- column_major_scales=False,
540
- use_ue8m0=(scale_fmt == "ue8m0"),
541
- )
542
-
543
- k_fp8_bytes = k_fp8.view(-1, head_dim).view(torch.uint8)
544
- scale_bytes = k_scale.view(torch.uint8).view(-1, 4)
545
- k = torch.cat(
546
- [k_fp8_bytes, scale_bytes], dim=-1
547
- ) # [total_tokens, head_dim + 4]
548
-
549
- slot_mapping = slot_mapping.flatten()
550
- # kv_cache: [num_block, block_size, head_dim + 4]
551
- kv_cache.view(-1, kv_cache.shape[-1]).index_copy_(0, slot_mapping, k)
552
-
553
- @staticmethod
554
- def cp_gather_indexer_k_quant_cache(
555
- kv_cache: torch.Tensor,
556
- dst_k: torch.Tensor,
557
- dst_scale: torch.Tensor,
558
- block_table: torch.Tensor,
559
- cu_seq_lens: torch.Tensor,
560
- ) -> None:
561
- """
562
- Args:
563
- kv_cache: [num_blocks, block_size, cache_stride] - quantized KV cache
564
- Layout per block: [k_values, scale_values]
565
- - k_values: [block_size * head_dim]
566
- - scale_values: [block_size * head_dim * 4 / quant_block_size]
567
- dst_k: [num_tokens, head_dim] - output tensor for K values
568
- dst_scale: [num_tokens, head_dim / quant_block_size * 4]
569
- - output tensor for scale values
570
- block_table: [batch_size, num_blocks] - block table for indexing
571
- cu_seq_lens: [batch_size + 1] - cumulative sequence lengths
572
- """
573
- batch_size = block_table.size(0)
574
- num_tokens = dst_k.size(0)
575
- head_dim = dst_k.size(1)
576
- cache_block_size = kv_cache.size(1)
577
- quant_block_size = head_dim * 4 // dst_scale.size(1)
578
-
579
- # For each token, find which batch it belongs to using searchsorted
580
- token_indices = torch.arange(num_tokens, device=dst_k.device) + 1
581
- # cu_seq_lens is [batch_size + 1], we need to find which interval each
582
- # token belongs to
583
- batch_indices = torch.searchsorted(cu_seq_lens, token_indices) - 1
584
- batch_indices = torch.clamp(batch_indices, 0, batch_size - 1)
585
-
586
- # Calculate the in-batch sequence index for each token
587
- inbatch_seq_indices = token_indices - cu_seq_lens[batch_indices]
588
-
589
- # Find which block each token belongs to
590
- block_indices_in_table = inbatch_seq_indices // cache_block_size
591
- physical_block_indices = block_table[batch_indices, block_indices_in_table]
592
-
593
- # Calculate the offset within each block
594
- inblock_offsets = (inbatch_seq_indices - 1) % cache_block_size
595
-
596
- # Calculate strides
597
- block_stride = kv_cache.stride(0) # stride for each block
598
-
599
- # Flatten kv_cache for easier indexing
600
- kv_cache_flat = kv_cache.view(-1)
601
-
602
- # Calculate source offset for K values for all tokens (vectorized)
603
- src_block_offsets = physical_block_indices * block_stride
604
- src_k_offsets = src_block_offsets + inblock_offsets * head_dim
605
-
606
- # Gather K values using advanced indexing
607
- # Create indices for all elements we need to gather
608
- k_indices = src_k_offsets.unsqueeze(1) + torch.arange(
609
- head_dim, device=dst_k.device
610
- )
611
- dst_k[:] = kv_cache_flat[k_indices]
612
-
613
- # Calculate source offset for scale values (vectorized)
614
- # Scales are stored after all K values for each block
615
- scale_size = head_dim * 4 // quant_block_size
616
- src_scale_offsets = src_block_offsets + head_dim + inblock_offsets * scale_size
617
-
618
- # Gather scale values
619
- scale_indices = src_scale_offsets.unsqueeze(1) + torch.arange(
620
- scale_size, device=dst_scale.device
621
- )
622
- dst_scale[:] = kv_cache_flat[scale_indices]
623
-
624
- @staticmethod
625
- def top_k_per_row_prefill(
626
- logits: torch.Tensor,
627
- cu_seqlen_ks: torch.Tensor,
628
- cu_seqlen_ke: torch.Tensor,
629
- raw_topk_indices: torch.Tensor,
630
- num_rows: int,
631
- stride0: int,
632
- strdide1: int,
633
- topk_tokens: int,
634
- ) -> torch.Tensor:
635
- real_topk = min(topk_tokens, logits.shape[-1])
636
- topk_indices = logits.topk(real_topk, dim=-1)[1].to(torch.int32)
637
- topk_indices -= cu_seqlen_ks[:, None]
638
- mask_lo = topk_indices >= 0
639
- mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0
640
- mask = torch.full_like(
641
- topk_indices, False, dtype=torch.bool, device=topk_indices.device
642
- )
643
- mask = mask_lo & mask_hi
644
- topk_indices.masked_fill_(~mask, -1)
645
- raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
646
- topk_indices
647
- )
648
-
649
- @staticmethod
650
- def top_k_per_row_decode(
651
- logits: torch.Tensor,
652
- next_n: int,
653
- seq_lens: torch.Tensor,
654
- raw_topk_indices: torch.Tensor,
655
- num_rows: int,
656
- stride0: int,
657
- stride1: int,
658
- topk_tokens: int,
659
- ) -> torch.Tensor:
660
- device = logits.device
661
- batch_size = seq_lens.size(0)
662
- # padded query len
663
- padded_num_tokens = batch_size * next_n
664
- positions = (
665
- torch.arange(logits.shape[-1], device=device)
666
- .unsqueeze(0)
667
- .expand(batch_size * next_n, -1)
668
- )
669
- row_indices = torch.arange(padded_num_tokens, device=device) // next_n
670
- next_n_offset = torch.arange(padded_num_tokens, device=device) % next_n
671
- index_end_pos = (seq_lens[row_indices] - next_n + next_n_offset).unsqueeze(1)
672
- # index_end_pos: [B * N, 1]
673
- mask = positions <= index_end_pos
674
- # mask: [B * N, L]
675
- logits = logits.masked_fill(~mask, float("-inf"))
676
- real_topk = min(topk_tokens, logits.shape[-1])
677
- topk_indices = logits.topk(real_topk, dim=-1)[1].to(torch.int32) # [B * N, K]
678
- # ensure we don't set indices for the top k
679
- # that is out of range(masked already)
680
- # this will happen if context length is shorter than K
681
- topk_indices[topk_indices > index_end_pos] = -1
682
- raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
683
- topk_indices
684
- )
685
-
511
@staticmethod
512
def register_ops_once() -> None:
513
global _OPS_REGISTERED
@@ -708,6 +533,18 @@ class xpu_ops:
533
fake_impl=_xpu_mxfp4_quantize_fake,
534
)
535
536
+ direct_register_custom_op(
537
+ op_name="xpu_fp8_mqa_logits",
538
+ op_func=_xpu_fp8_mqa_logits_impl,
539
+ fake_impl=_xpu_fp8_mqa_logits_fake,
540
+ )
541
+
542
+ direct_register_custom_op(
543
+ op_name="xpu_fp8_paged_mqa_logits",
544
+ op_func=_xpu_fp8_paged_mqa_logits_impl,
545
+ fake_impl=_xpu_fp8_paged_mqa_logits_fake,
546
+ )
547
+
548
direct_register_custom_op(
549
op_name="gdn_attention_core_xpu",
550
op_func=_gdn_attention_core_xpu_impl,
vllm/model_executor/layers/sparse_attn_indexer.py
+74
-65
@@ -5,6 +5,7 @@
5
import torch
6
7
import vllm.envs as envs
8
+from vllm import _custom_ops as ops
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11
from vllm.forward_context import get_forward_context
@@ -28,11 +29,6 @@ from vllm.v1.attention.backends.mla.indexer import (
29
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
30
from vllm.v1.worker.workspace import current_workspace_manager
31
31
-if current_platform.is_cuda_alike():
32
- from vllm import _custom_ops as ops
33
-elif current_platform.is_xpu():
34
- from vllm._xpu_ops import xpu_ops
35
-
32
logger = init_logger(__name__)
33
34
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
@@ -222,42 +218,42 @@ def sparse_attn_indexer(
218
q_slice_cast = q_slice
219
k_quant_cast = k_quant
220
k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
225
- logits = fp8_fp4_mqa_logits(
226
- (q_slice_cast, q_scale_slice),
227
- (k_quant_cast, k_scale_cast),
228
- weights[chunk.token_start : chunk.token_end],
229
- chunk.cu_seqlen_ks,
230
- chunk.cu_seqlen_ke,
231
- clean_logits=False,
232
- )
233
- num_rows = logits.shape[0]
234
-
235
- topk_indices = topk_indices_buffer[
236
- chunk.token_start : chunk.token_end, :topk_tokens
237
- ]
238
-
221
if current_platform.is_xpu():
240
- xpu_ops.top_k_per_row_prefill( # type: ignore[attr-defined]
241
- logits,
222
+ if q_scale_slice is not None:
223
+ raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
224
+ logits = torch.ops.vllm.xpu_fp8_mqa_logits(
225
+ q_slice_cast,
226
+ k_quant_cast,
227
+ k_scale_cast,
228
+ weights[chunk.token_start : chunk.token_end],
229
chunk.cu_seqlen_ks,
230
chunk.cu_seqlen_ke,
244
- topk_indices,
245
- num_rows,
246
- logits.stride(0),
247
- logits.stride(1),
248
- topk_tokens,
231
)
232
else:
251
- torch.ops._C.top_k_per_row_prefill(
252
- logits,
233
+ logits = fp8_fp4_mqa_logits(
234
+ (q_slice_cast, q_scale_slice),
235
+ (k_quant_cast, k_scale_cast),
236
+ weights[chunk.token_start : chunk.token_end],
237
chunk.cu_seqlen_ks,
238
chunk.cu_seqlen_ke,
255
- topk_indices,
256
- num_rows,
257
- logits.stride(0),
258
- logits.stride(1),
259
- topk_tokens,
239
+ clean_logits=False,
240
)
241
+ num_rows = logits.shape[0]
242
+
243
+ topk_indices = topk_indices_buffer[
244
+ chunk.token_start : chunk.token_end, :topk_tokens
245
+ ]
246
+
247
+ ops.top_k_per_row_prefill(
248
+ logits,
249
+ chunk.cu_seqlen_ks,
250
+ chunk.cu_seqlen_ke,
251
+ topk_indices,
252
+ num_rows,
253
+ logits.stride(0),
254
+ logits.stride(1),
255
+ topk_tokens,
256
+ )
257
258
if has_decode:
259
decode_metadata = attn_metadata_narrowed.decode
@@ -309,16 +305,32 @@ def sparse_attn_indexer(
305
if use_fp4_cache
306
else padded_q_quant_decode_tokens
307
)
312
- logits = fp8_fp4_paged_mqa_logits(
313
- (padded_q_quant_cast, padded_q_scale),
314
- kv_cache,
315
- weights[:num_padded_tokens],
316
- seq_lens,
317
- decode_metadata.block_table,
318
- decode_metadata.schedule_metadata,
319
- max_model_len=max_model_len,
320
- clean_logits=False,
321
- )
308
+ if current_platform.is_xpu():
309
+ if padded_q_scale is not None:
310
+ raise RuntimeError("XPU fp8_paged_mqa_logits does not support FP4 Q")
311
+ seq_lens_xpu = (
312
+ seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
313
+ )
314
+ logits = torch.ops.vllm.xpu_fp8_paged_mqa_logits(
315
+ padded_q_quant_cast,
316
+ kv_cache,
317
+ weights[:num_padded_tokens],
318
+ seq_lens_xpu,
319
+ decode_metadata.block_table,
320
+ decode_metadata.schedule_metadata,
321
+ max_model_len,
322
+ )
323
+ else:
324
+ logits = fp8_fp4_paged_mqa_logits(
325
+ (padded_q_quant_cast, padded_q_scale),
326
+ kv_cache,
327
+ weights[:num_padded_tokens],
328
+ seq_lens,
329
+ decode_metadata.block_table,
330
+ decode_metadata.schedule_metadata,
331
+ max_model_len=max_model_len,
332
+ clean_logits=False,
333
+ )
334
num_rows = logits.shape[0]
335
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
336
@@ -336,28 +348,16 @@ def sparse_attn_indexer(
348
attn_metadata_narrowed.max_seq_len,
349
)
350
else:
339
- if current_platform.is_xpu():
340
- xpu_ops.top_k_per_row_decode( # type: ignore[attr-defined]
341
- logits,
342
- next_n,
343
- seq_lens,
344
- topk_indices,
345
- num_rows,
346
- logits.stride(0),
347
- logits.stride(1),
348
- topk_tokens,
349
- )
350
- else:
351
- torch.ops._C.top_k_per_row_decode(
352
- logits,
353
- next_n,
354
- seq_lens,
355
- topk_indices,
356
- num_rows,
357
- logits.stride(0),
358
- logits.stride(1),
359
- topk_tokens,
360
- )
351
+ ops.top_k_per_row_decode(
352
+ logits,
353
+ next_n,
354
+ seq_lens,
355
+ topk_indices,
356
+ num_rows,
357
+ logits.stride(0),
358
+ logits.stride(1),
359
+ topk_tokens,
360
+ )
361
362
if decode_metadata.requires_padding:
363
# if padded, we need to unpack
@@ -494,6 +494,15 @@ class SparseAttnIndexer(CustomOp):
494
self.use_fp4_cache,
495
)
496
497
+ def forward_xpu(
498
+ self,
499
+ hidden_states: torch.Tensor,
500
+ q_fp8: torch.Tensor,
501
+ k: torch.Tensor,
502
+ weights: torch.Tensor,
503
+ ):
504
+ return self.forward_cuda(hidden_states, q_fp8, k, weights)
505
+
506
def forward_hip(
507
self,
508
hidden_states: torch.Tensor,