Harden DeepSeek V3.2 fused kernel grids (#52381)
Signed-off-by: yimdev <5779256+yimdev@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
yimdev committed
Aug 18, 2026 at 16:43 UTC
eab1cff5b0ca87ec415c3ca555f060933d0a6b72
2 files changed
+79
-6
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+73
@@ -362,6 +362,46 @@ def test_fused_norm_rope_ds_mla(num_tokens: int):
362
assert (topk == 7).all(), "topk buffer should be untouched (no indexer)"
363
364
365
+def test_fused_norm_rope_supports_large_token_count():
366
+ """Keep the token count off CUDA grid-y at its 65,536-block boundary."""
367
+ num_tokens = 65536
368
+ dev = "cuda"
369
+ dtype = torch.bfloat16
370
+ positions = torch.zeros(num_tokens, device=dev, dtype=torch.int64)
371
+ q_c = torch.ones((num_tokens, 1), device=dev, dtype=dtype)
372
+ kv_c = torch.ones((num_tokens, 1), device=dev, dtype=dtype)
373
+ k_pe = torch.ones((num_tokens, 2), device=dev, dtype=dtype)
374
+ norm_w = torch.ones(1, device=dev, dtype=dtype)
375
+ cos_sin = torch.tensor([[1.0, 0.0]], device=dev, dtype=torch.float32)
376
+ topk = torch.empty((num_tokens, 1), device=dev, dtype=torch.int32)
377
+ slot_mapping = torch.arange(num_tokens, device=dev, dtype=torch.int64)
378
+ mla_cache = torch.empty((1, num_tokens, 3), device=dev, dtype=dtype)
379
+
380
+ q_out = K.fused_norm_rope(
381
+ positions,
382
+ q_c,
383
+ norm_w,
384
+ EPS,
385
+ kv_c,
386
+ norm_w,
387
+ EPS,
388
+ k_pe,
389
+ cos_sin,
390
+ None,
391
+ None,
392
+ None,
393
+ EPS,
394
+ None,
395
+ topk,
396
+ slot_mapping=slot_mapping,
397
+ mla_kv_cache=mla_cache,
398
+ has_indexer=False,
399
+ )
400
+
401
+ rows = torch.tensor([0, num_tokens - 1], device=dev)
402
+ assert_bf16(q_out[rows], rms_norm(q_c[rows], norm_w), "large-token q norm")
403
+
404
+
405
# ── fused_q ──────────────────────────────────────────────────────────────────
406
407
@@ -539,6 +579,39 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool):
579
torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3)
580
581
582
+def test_fused_q_triton_supports_large_token_count():
583
+ """Keep the token count off CUDA grid-y in the Triton fallback.
584
+
585
+ The minimal dimensions also bypass CuTeDSL on SM100.
586
+ """
587
+ num_tokens = 65536
588
+ dev = "cuda"
589
+ dtype = torch.bfloat16
590
+ positions = torch.zeros(num_tokens, device=dev, dtype=torch.int64)
591
+ q_pe = torch.ones((num_tokens, 1, 2), device=dev, dtype=dtype)
592
+ ql_nope = torch.ones((num_tokens, 1, 1), device=dev, dtype=dtype)
593
+ cos_sin = torch.tensor([[1.0, 0.0]], device=dev, dtype=torch.float32)
594
+ q_scale = torch.ones(1, device=dev, dtype=torch.float32)
595
+
596
+ _, _, mqa_q = K.fused_q(
597
+ positions,
598
+ q_pe,
599
+ cos_sin,
600
+ None,
601
+ None,
602
+ ql_nope,
603
+ q_scale,
604
+ None,
605
+ 0.0,
606
+ 0.0,
607
+ has_indexer=False,
608
+ )
609
+
610
+ rows = torch.tensor([0, num_tokens - 1], device=dev)
611
+ ref = torch.cat([ql_nope, q_pe], dim=-1).to(FP8)
612
+ assert_fp8(mqa_q[rows], ref[rows], "large-token fused Q")
613
+
614
+
615
# ── fused_eh_norm (MTP) ──────────────────────────────────────────────────────
616
617
vllm/models/deepseek_v32/common/kernels.py
+6
-6
@@ -152,8 +152,8 @@ def _fused_norm_rope_kernel(
152
INDEX_ROPE_INTERLEAVE: tl.constexpr,
153
USE_PDL: tl.constexpr,
154
):
155
- pid = tl.program_id(0)
156
- tok_idx = tl.program_id(1)
155
+ tok_idx = tl.program_id(0).to(tl.int64)
156
+ pid = tl.program_id(1)
157
if USE_PDL:
158
tl.extra.cuda.gdc_wait()
159
tl.extra.cuda.gdc_launch_dependents()
@@ -479,7 +479,7 @@ def fused_norm_rope(
479
if q_c_out is None:
480
q_c_out = torch.empty_like(q_c)
481
use_pdl = current_platform.is_arch_support_pdl()
482
- _fused_norm_rope_kernel[(4, num_tokens)](
482
+ _fused_norm_rope_kernel[(num_tokens, 4)](
483
positions,
484
# Q RMS norm
485
q_c,
@@ -594,8 +594,8 @@ def _fused_q_kernel(
594
QUANTIZE_MQA: tl.constexpr,
595
USE_PDL: tl.constexpr,
596
):
597
- pid = tl.program_id(0)
598
- tok_idx = tl.program_id(1)
597
+ tok_idx = tl.program_id(0).to(tl.int64)
598
+ pid = tl.program_id(1)
599
head_idx = tl.program_id(2)
600
if USE_PDL:
601
tl.extra.cuda.gdc_wait()
@@ -869,7 +869,7 @@ def fused_q(
869
return index_q_fp8, index_weights_out, mqa_q
870
871
use_pdl = current_platform.is_arch_support_pdl()
872
- _fused_q_kernel[(3, num_tokens, grid_heads)](
872
+ _fused_q_kernel[(num_tokens, 3, grid_heads)](
873
positions,
874
q_pe,
875
q_pe.stride(0),