[Perf][DSV4.1] Optimize MegaMoE staging and NVFP4 cache gathers (#57604)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Codex <noreply@openai.com>
Woosuk Kwon committed
Sep 18, 2026 at 12:57 UTC
50d812b66a697c464a9964b3a061b47583b83e34
5 files changed
+130
-62
tests/kernels/test_compressor_kv_cache.py
+68
-27
@@ -473,15 +473,18 @@ _E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
473
474
475
def _decode_nvfp4_row(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
476
- """Unpack one 256-byte e2m1 row and apply its 32 e4m3 tile scales."""
476
+ """Unpack e2m1 rows and apply their e4m3 tile scales."""
477
mags = torch.tensor(_E2M1_MAGNITUDES, device=packed.device)
478
- codes = torch.empty(512, dtype=torch.uint8, device=packed.device)
479
- codes[0::2] = packed & 0xF # even element in the low nibble
480
- codes[1::2] = packed >> 4
478
+ codes = torch.empty(
479
+ (*packed.shape[:-1], 512), dtype=torch.uint8, device=packed.device
480
+ )
481
+ codes[..., 0::2] = packed & 0xF # even element in the low nibble
482
+ codes[..., 1::2] = packed >> 4
483
vals = mags[(codes & 7).long()] * torch.where(codes >= 8, -1.0, 1.0)
484
return (
483
- vals.view(32, 16) * scales.view(torch.float8_e4m3fn).float().view(32, 1)
484
- ).flatten()
485
+ vals.reshape(*packed.shape[:-1], 32, 16)
486
+ * scales.view(torch.float8_e4m3fn).float().reshape(*packed.shape[:-1], 32, 1)
487
+ ).flatten(-2)
488
489
490
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
@@ -543,7 +546,10 @@ def test_v41_rope_insert_nvfp4_record(compress_ratio: int):
546
547
548
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
546
-def test_v41_nvfp4_gather_matches_insert():
549
+@pytest.mark.parametrize("num_reqs", [1, 8, 32])
550
+@pytest.mark.parametrize("num_tokens", [70, 128, 129, 2051])
551
+@pytest.mark.parametrize("partial", [False, True])
552
+def test_v41_nvfp4_gather_matches_insert(num_reqs, num_tokens, partial):
553
"""The NVFP4 gather dequantizes exactly what the insert kernel wrote."""
554
from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
555
from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
@@ -553,38 +559,73 @@ def test_v41_nvfp4_gather_matches_insert():
559
torch.manual_seed(17)
560
device = "cuda"
561
block_size = 64
556
- num_tokens = 70
557
- num_blocks = 4
562
+ num_blocks = math.ceil(num_tokens / block_size)
563
page_bytes = math.ceil(block_size * 288 / 512) * 512
564
565
positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
566
latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device)
562
- angles = torch.randn(128, 32, device=device)
567
+ angles = torch.randn(num_tokens, 32, device=device)
568
cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1)
569
backing = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device)
570
cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1))
571
slots = torch.arange(num_tokens, dtype=torch.int64, device=device)
572
rope_quant_insert(latent, positions, cos_sin, cache, slots, 1)
573
569
- out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device)
570
- dequantize_and_gather_k_cache(
571
- out,
572
- cache,
573
- seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device),
574
- gather_lens=None,
575
- block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view(
576
- 1, -1
577
- ),
578
- block_size=block_size,
579
- offset=0,
574
+ offset = 7 if partial else 0
575
+ seq_lens = num_tokens - torch.arange(num_reqs, dtype=torch.int32, device=device)
576
+ gather_lens = seq_lens - 3 if partial else None
577
+ storage = torch.full(
578
+ (num_reqs, num_tokens + offset + 16, 512),
579
+ -123.0,
580
+ dtype=torch.bfloat16,
581
+ device=device,
582
)
581
- for t in (0, 1, block_size, num_tokens - 1):
582
- page, row = divmod(t, block_size)
583
- expected = _decode_nvfp4_row(
584
- backing[page, row * 256 : (row + 1) * 256],
585
- backing[page, block_size * 256 + row * 32 :][:32],
583
+ out = storage[:, : num_tokens + offset]
584
+ block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).repeat(
585
+ num_reqs, 1
586
+ )
587
+
588
+ def gather():
589
+ dequantize_and_gather_k_cache(
590
+ out,
591
+ cache,
592
+ seq_lens=seq_lens,
593
+ gather_lens=gather_lens,
594
+ block_table=block_table,
595
+ block_size=block_size,
596
+ offset=offset,
597
)
587
- torch.testing.assert_close(out[0, t].float(), expected, rtol=0, atol=0)
598
+
599
+ expected = _decode_nvfp4_row(
600
+ backing[:, : block_size * 256].reshape(-1, 256),
601
+ backing[:, block_size * 256 : block_size * 288].reshape(-1, 32),
602
+ )
603
+
604
+ def check(shorter_by):
605
+ assert (storage[:, num_tokens + offset :] == -123).all()
606
+ for req in range(num_reqs):
607
+ start = 3 if partial else 0
608
+ length = num_tokens - req - start - shorter_by
609
+ assert (out[req, :offset] == -123).all()
610
+ assert (out[req, offset + length :] == -123).all()
611
+ torch.testing.assert_close(
612
+ out[req, offset : offset + length].float(),
613
+ expected[start : start + length],
614
+ rtol=0,
615
+ atol=0,
616
+ )
617
+
618
+ gather()
619
+ check(0)
620
+ graph = torch.cuda.CUDAGraph()
621
+ with torch.cuda.graph(graph):
622
+ gather()
623
+ storage.fill_(-123)
624
+ seq_lens.sub_(1)
625
+ if gather_lens is not None:
626
+ gather_lens.sub_(1)
627
+ graph.replay()
628
+ check(1)
629
630
631
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
tests/models/test_deepseek_v4_mega_moe.py
+13
-9
@@ -686,12 +686,14 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert(
686
not torch.cuda.is_available(),
687
reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
688
)
689
-def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
689
+@pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384])
690
+@pytest.mark.parametrize("hidden_size", [256, 5120])
691
+def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact(
692
+ num_tokens, hidden_size
693
+):
694
from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
695
696
device = torch.device("cuda")
693
- num_tokens = 7
694
- hidden_size = 256
697
top_k = 8
698
699
generator = torch.Generator(device=device)
@@ -846,6 +848,7 @@ def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m):
848
assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
849
assert torch.equal(fused_x_sf, ref_x_sf)
850
assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated])
851
+ assert torch.all(fused_shared_x_sf[~populated] == -1)
852
853
854
def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast():
@@ -884,11 +887,13 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe():
887
reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
888
)
889
@pytest.mark.parametrize("nonfinite_padding", [False, True])
887
-def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(nonfinite_padding):
890
+@pytest.mark.parametrize("num_tokens", [7, 65, 135])
891
+def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(
892
+ nonfinite_padding, num_tokens
893
+):
894
from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
895
896
device = torch.device("cuda")
891
- num_tokens = 7
897
hidden_size = 256
898
top_k = 8
899
@@ -916,10 +921,9 @@ def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(nonfinite_paddin
921
dtype=torch.float32,
922
generator=generator,
923
)
919
- is_padding = torch.tensor(
920
- [False, True, False, False, True, False, True],
921
- device=device,
922
- )
924
+ is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool)
925
+ is_padding[1::3] = True
926
+ is_padding[-1] = True
927
if nonfinite_padding:
928
topk_weights[is_padding] = float("nan")
929
vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py
+43
-24
@@ -12,7 +12,7 @@ import torch
12
from vllm.triton_utils import tl, triton
13
14
15
-@triton.jit
15
+@triton.jit(do_not_specialize=["num_tokens"])
16
def _prepare_megamoe_inputs_kernel(
17
hidden_states,
18
x_fp8,
@@ -40,27 +40,32 @@ def _prepare_megamoe_inputs_kernel(
40
topk_idx_stride_k: tl.constexpr,
41
topk_weights_out_stride_m: tl.constexpr,
42
topk_weights_out_stride_k: tl.constexpr,
43
+ num_tokens,
44
hidden_size: tl.constexpr,
45
top_k: tl.constexpr,
46
+ BLOCK_M: tl.constexpr,
47
BLOCK_K: tl.constexpr,
48
GROUP_K: tl.constexpr,
49
BLOCK_TOPK: tl.constexpr,
50
SHARED_BLOCK_M: tl.constexpr,
51
) -> None:
50
- token_id = tl.program_id(0)
52
+ token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
53
+ token_mask = token_id < num_tokens
54
k_block_id = tl.program_id(1)
55
56
k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
54
- k_mask = k_offsets < hidden_size
57
+ k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size)
58
hidden = tl.load(
56
- hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
59
+ hidden_states
60
+ + token_id[:, None] * hidden_stride_m
61
+ + k_offsets[None, :] * hidden_stride_k,
62
mask=k_mask,
63
other=0.0,
64
).to(tl.float32)
65
66
num_groups: tl.constexpr = BLOCK_K // GROUP_K
62
- hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
63
- amax = tl.max(hidden_groups, axis=1)
67
+ hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K])
68
+ amax = tl.max(hidden_groups, axis=2)
69
amax = tl.maximum(amax, 1.0e-4)
70
71
scale = amax / 448.0
@@ -71,21 +76,24 @@ def _prepare_megamoe_inputs_kernel(
76
scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
77
rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
78
74
- hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
75
- scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
76
- scaled = tl.reshape(scaled, [BLOCK_K])
79
+ hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K])
80
+ scaled = hidden_groups * (1.0 / rounded_scale)[:, :, None]
81
+ scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K])
82
fp8 = scaled.to(tl.float8e4nv)
83
tl.store(
79
- x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
84
+ x_fp8 + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k,
85
fp8,
86
mask=k_mask,
87
)
88
89
scale_offsets = tl.arange(0, num_groups)
85
- packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
90
+ packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to(
91
+ tl.int32
92
+ )
93
tl.store(
94
x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
95
packed_scale,
96
+ mask=token_mask,
97
)
98
99
# DeepGEMM's SM100 shared-expert TMA loads require the activation scales
@@ -105,41 +113,48 @@ def _prepare_megamoe_inputs_kernel(
113
+ shared_row * shared_x_sf_stride_m
114
+ k_block_id * shared_x_sf_stride_k,
115
packed_scale,
116
+ mask=token_mask,
117
)
118
119
if k_block_id == 0:
120
topk_offsets = tl.arange(0, BLOCK_TOPK)
112
- topk_mask = topk_offsets < top_k
113
- token_is_padding = False
121
+ topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k)
122
+ token_is_padding = tl.full((BLOCK_M,), False, tl.int1)
123
if is_padding is not None:
115
- token_is_padding = tl.load(is_padding + token_id * is_padding_stride_m)
124
+ token_is_padding = tl.load(
125
+ is_padding + token_id * is_padding_stride_m,
126
+ mask=token_mask,
127
+ other=True,
128
+ )
129
130
ids = tl.load(
118
- topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
131
+ topk_ids
132
+ + token_id[:, None] * topk_ids_stride_m
133
+ + topk_offsets[None, :] * topk_ids_stride_k,
134
mask=topk_mask,
135
other=0,
136
).to(tl.int64)
122
- ids = tl.where(token_is_padding, -1, ids)
137
+ ids = tl.where(token_is_padding[:, None], -1, ids)
138
tl.store(
139
topk_idx_out
125
- + token_id * topk_idx_stride_m
126
- + topk_offsets * topk_idx_stride_k,
140
+ + token_id[:, None] * topk_idx_stride_m
141
+ + topk_offsets[None, :] * topk_idx_stride_k,
142
ids,
143
mask=topk_mask,
144
)
145
146
weights = tl.load(
147
topk_weights
133
- + token_id * topk_weights_stride_m
134
- + topk_offsets * topk_weights_stride_k,
148
+ + token_id[:, None] * topk_weights_stride_m
149
+ + topk_offsets[None, :] * topk_weights_stride_k,
150
mask=topk_mask,
151
other=0.0,
152
)
138
- weights = tl.where(token_is_padding, 0.0, weights)
153
+ weights = tl.where(token_is_padding[:, None], 0.0, weights)
154
tl.store(
155
topk_weights_out
141
- + token_id * topk_weights_out_stride_m
142
- + topk_offsets * topk_weights_out_stride_k,
156
+ + token_id[:, None] * topk_weights_out_stride_m
157
+ + topk_offsets[None, :] * topk_weights_out_stride_k,
158
weights,
159
mask=topk_mask,
160
)
@@ -195,7 +210,9 @@ def prepare_megamoe_inputs(
210
)
211
212
block_k = 128
198
- grid = (num_tokens, triton.cdiv(hidden_size, block_k))
213
+ # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled.
214
+ block_m = 8 if num_tokens >= 64 else 1
215
+ grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k))
216
block_topk = triton.next_power_of_2(top_k)
217
padding_stride_m = is_padding.stride(0) if is_padding is not None else 0
218
_prepare_megamoe_inputs_kernel[grid](
@@ -225,8 +242,10 @@ def prepare_megamoe_inputs(
242
topk_idx_out.stride(1),
243
topk_weights_out.stride(0),
244
topk_weights_out.stride(1),
245
+ num_tokens,
246
hidden_size,
247
top_k,
248
+ BLOCK_M=block_m,
249
BLOCK_K=block_k,
250
GROUP_K=32,
251
BLOCK_TOPK=block_topk,
vllm/models/deepseek_v41/common/ops/cache_utils.py
+5
-1
@@ -584,7 +584,11 @@ def dequantize_and_gather_k_cache_triton(
584
num_reqs = seq_lens.shape[0]
585
NUM_WORKERS = 128
586
if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN:
587
- _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, NUM_WORKERS)](
587
+ # Keep short gathers unchanged; cap total workers for multi-request chunks.
588
+ num_workers = min(
589
+ max(128, 2048 // max(1, num_reqs)), max(128, out.shape[1] - offset)
590
+ )
591
+ _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, num_workers)](
592
out,
593
out.stride(0),
594
out.stride(1),
vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
+1
-1
@@ -422,7 +422,7 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):
422
if not swa_only:
423
assert flashmla_metadata is not None
424
dequantize_and_gather_k_cache(
425
- kv_ws[:chunk_size],
425
+ kv_ws[:chunk_size, :chunk_n],
426
self._compressed_kv_cache(),
427
seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
428
gather_lens=None,