[XPU][DeepSeekV4]Add DeepSeek-V4 fuse_index_q SYCL kernel path (#45991)
Signed-off-by: Avery Miao <avery.miao@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Avery Miao committed
Jul 21, 2026 at 21:22 UTC
c67650f04bc9241073130a7f8f245ce470bb69ac
2 files changed
+150
vllm/_xpu_ops.py
+127
@@ -375,6 +375,115 @@ def _topk_topp_sample_fake(
375
return
376
377
378
+def _xpu_deepseek_fused_indexer_q_rope_fp8_impl(
379
+ index_q: torch.Tensor,
380
+ positions: torch.Tensor,
381
+ index_q_cos_sin_cache: torch.Tensor,
382
+ index_weights: torch.Tensor,
383
+ index_weights_softmax_scale: float,
384
+ index_weights_head_scale: float,
385
+ index_q_fp8: torch.Tensor,
386
+ index_weights_out: torch.Tensor,
387
+) -> None:
388
+ """Fused RoPE + FP8 quant of the DeepSeek-V4 sparse-indexer Q (XPU).
389
+ Writes ``index_q_fp8`` and ``index_weights_out`` in place (no return).
390
+ Requires head_dim=128, rope_dim=64, and num_heads divisible by 2.
391
+
392
+ Args:
393
+ index_q: (T, H, 128) bfloat16 Q before RoPE. Contiguous.
394
+ positions: (T,) int64 absolute token positions.
395
+ index_q_cos_sin_cache: (max_pos, 64) float32 RoPE cos/sin cache.
396
+ index_weights: (T, H) bfloat16 raw indexer weights.
397
+ index_weights_softmax_scale: scalar softmax scale.
398
+ index_weights_head_scale: scalar per-head scale.
399
+ index_q_fp8: (T, H, 128) fp8 (e4m3) output, preallocated. [written]
400
+ index_weights_out: (T, H) float32 output, preallocated;
401
+ = index_weights * q_scale * softmax_scale * head_scale (the
402
+ per-(token, head) q_scale is folded in here). [written]
403
+ """
404
+ torch.ops._xpu_C.deepseek_fused_indexer_q_rope_fp8(
405
+ index_q,
406
+ positions,
407
+ index_q_cos_sin_cache,
408
+ index_weights,
409
+ index_weights_softmax_scale,
410
+ index_weights_head_scale,
411
+ index_q_fp8,
412
+ index_weights_out,
413
+ )
414
+
415
+
416
+def _xpu_deepseek_fused_indexer_q_rope_fp8_fake(
417
+ index_q: torch.Tensor,
418
+ positions: torch.Tensor,
419
+ index_q_cos_sin_cache: torch.Tensor,
420
+ index_weights: torch.Tensor,
421
+ index_weights_softmax_scale: float,
422
+ index_weights_head_scale: float,
423
+ index_q_fp8: torch.Tensor,
424
+ index_weights_out: torch.Tensor,
425
+) -> None:
426
+ return
427
+
428
+
429
+def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl(
430
+ index_q: torch.Tensor,
431
+ positions: torch.Tensor,
432
+ index_q_cos_sin_cache: torch.Tensor,
433
+ index_weights: torch.Tensor,
434
+ index_weights_softmax_scale: float,
435
+ index_weights_head_scale: float,
436
+ index_q_packed: torch.Tensor,
437
+ index_q_scale: torch.Tensor,
438
+ index_weights_out: torch.Tensor,
439
+) -> None:
440
+ """Fused RoPE + MXFP4 quant of the DeepSeek-V4 sparse-indexer Q (XPU).
441
+ Writes ``index_q_packed``, ``index_q_scale`` and ``index_weights_out`` in
442
+ place (no return). Requires head_dim=128, rope_dim=64, and num_heads
443
+ divisible by 4.
444
+
445
+ Args:
446
+ index_q: (T, H, 128) bfloat16 Q before RoPE. Contiguous.
447
+ positions: (T,) int64 absolute token positions.
448
+ index_q_cos_sin_cache: (max_pos, 64) float32 RoPE cos/sin cache.
449
+ index_weights: (T, H) bfloat16 raw indexer weights.
450
+ index_weights_softmax_scale: scalar softmax scale.
451
+ index_weights_head_scale: scalar per-head scale.
452
+ index_q_packed: (T, H, 64) uint8 packed E2M1 nibbles (2 per byte),
453
+ preallocated. [written]
454
+ index_q_scale: (T, H, 4) uint8 ue8m0 per-block (32-elem) scales,
455
+ preallocated. [written]
456
+ index_weights_out: (T, H) float32 output, preallocated;
457
+ = index_weights * softmax_scale * head_scale (no q_scale folded;
458
+ per-block scales live in index_q_scale). [written]
459
+ """
460
+ torch.ops._xpu_C.deepseek_fused_indexer_q_rope_mxfp4(
461
+ index_q,
462
+ positions,
463
+ index_q_cos_sin_cache,
464
+ index_weights,
465
+ index_weights_softmax_scale,
466
+ index_weights_head_scale,
467
+ index_q_packed,
468
+ index_q_scale,
469
+ index_weights_out,
470
+ )
471
+
472
+
473
+def _xpu_deepseek_fused_indexer_q_rope_mxfp4_fake(
474
+ index_q: torch.Tensor,
475
+ positions: torch.Tensor,
476
+ index_q_cos_sin_cache: torch.Tensor,
477
+ index_weights: torch.Tensor,
478
+ index_weights_softmax_scale: float,
479
+ index_weights_head_scale: float,
480
+ index_q_packed: torch.Tensor,
481
+ index_q_scale: torch.Tensor,
482
+ index_weights_out: torch.Tensor,
483
+) -> None:
484
+ return
485
+
486
+
487
def _xpu_mxfp8_quantize_impl(
488
x: torch.Tensor, dtype: torch.dtype | None = None
489
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1141,6 +1250,24 @@ class xpu_ops:
1250
fake_impl=_topk_topp_sample_fake,
1251
)
1252
1253
+ direct_register_custom_op(
1254
+ op_name="xpu_deepseek_fused_indexer_q_rope_fp8",
1255
+ op_func=_xpu_deepseek_fused_indexer_q_rope_fp8_impl,
1256
+ mutates_args=["index_q_fp8", "index_weights_out"],
1257
+ fake_impl=_xpu_deepseek_fused_indexer_q_rope_fp8_fake,
1258
+ )
1259
+
1260
+ direct_register_custom_op(
1261
+ op_name="xpu_deepseek_fused_indexer_q_rope_mxfp4",
1262
+ op_func=_xpu_deepseek_fused_indexer_q_rope_mxfp4_impl,
1263
+ mutates_args=[
1264
+ "index_q_packed",
1265
+ "index_q_scale",
1266
+ "index_weights_out",
1267
+ ],
1268
+ fake_impl=_xpu_deepseek_fused_indexer_q_rope_mxfp4_fake,
1269
+ )
1270
+
1271
_OPS_REGISTERED = True
1272
1273
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+23
@@ -367,6 +367,18 @@ def fused_indexer_q_rope_quant(
367
index_q_scale,
368
index_weights_out,
369
)
370
+ elif current_platform.is_xpu():
371
+ torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4(
372
+ index_q,
373
+ positions,
374
+ index_q_cos_sin_cache,
375
+ index_weights,
376
+ index_weights_softmax_scale,
377
+ index_weights_head_scale,
378
+ index_q_packed,
379
+ index_q_scale,
380
+ index_weights_out,
381
+ )
382
else:
383
_fused_indexer_q_rope_mxfp4_kernel[(num_tokens, num_index_q_heads)](
384
positions,
@@ -423,6 +435,17 @@ def fused_indexer_q_rope_quant(
435
index_q_fp8,
436
index_weights_out,
437
)
438
+ elif current_platform.is_xpu():
439
+ torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8(
440
+ index_q,
441
+ positions,
442
+ index_q_cos_sin_cache,
443
+ index_weights,
444
+ index_weights_softmax_scale,
445
+ index_weights_head_scale,
446
+ index_q_fp8,
447
+ index_weights_out,
448
+ )
449
else:
450
_fused_indexer_q_rope_quant_kernel[(num_tokens, num_index_q_heads)](
451
positions,