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, 22:45 UTC
[DeepSeek V4] Move MegaMoE input prep kernel to nvidia/ops (#43632)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 25, 2026 at 21:08 UTC
aa2b56ffb0c1d8b29d6468282a49d9b4a9f0dd1c
3 files changed
+177
-165
tests/models/test_deepseek_v4_mega_moe.py
+2
-2
@@ -8,9 +8,9 @@ import torch
8
9
from vllm.models.deepseek_v4.nvidia.model import (
10
DeepseekV4MegaMoEExperts,
11
- _stage_deepseek_v4_mega_moe_inputs,
11
make_deepseek_v4_expert_params_mapping,
12
)
13
+from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
14
from vllm.platforms import current_platform
15
16
pytestmark = pytest.mark.skipif(
@@ -164,7 +164,7 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
164
fused_topk_idx = torch.empty_like(ref_topk_idx)
165
fused_topk_weights = torch.empty_like(ref_topk_weights)
166
167
- _stage_deepseek_v4_mega_moe_inputs(
167
+ prepare_megamoe_inputs(
168
hidden_states,
169
topk_weights,
170
topk_ids,
vllm/models/deepseek_v4/nvidia/model.py
+2
-163
@@ -59,9 +59,9 @@ from vllm.models.deepseek_v4.nvidia.ops.attention import (
59
DeepseekV4MLAModules,
60
DeepseekV4MultiHeadLatentAttentionWrapper,
61
)
62
+from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
63
from vllm.platforms import current_platform
64
from vllm.sequence import IntermediateTensors
64
-from vllm.triton_utils import tl, triton
65
from vllm.utils.torch_utils import direct_register_custom_op
66
67
@@ -116,167 +116,6 @@ class DeepseekV4MLP(nn.Module):
116
return x
117
118
119
-@triton.jit
120
-def _deepseek_v4_stage_mega_moe_inputs_kernel(
121
- hidden_states,
122
- x_fp8,
123
- x_sf,
124
- topk_ids,
125
- topk_weights,
126
- topk_idx_out,
127
- topk_weights_out,
128
- hidden_stride_m: tl.constexpr,
129
- hidden_stride_k: tl.constexpr,
130
- x_stride_m: tl.constexpr,
131
- x_stride_k: tl.constexpr,
132
- x_sf_stride_m: tl.constexpr,
133
- x_sf_stride_k: tl.constexpr,
134
- topk_ids_stride_m: tl.constexpr,
135
- topk_ids_stride_k: tl.constexpr,
136
- topk_weights_stride_m: tl.constexpr,
137
- topk_weights_stride_k: tl.constexpr,
138
- topk_idx_stride_m: tl.constexpr,
139
- topk_idx_stride_k: tl.constexpr,
140
- topk_weights_out_stride_m: tl.constexpr,
141
- topk_weights_out_stride_k: tl.constexpr,
142
- hidden_size: tl.constexpr,
143
- top_k: tl.constexpr,
144
- BLOCK_K: tl.constexpr,
145
- GROUP_K: tl.constexpr,
146
- BLOCK_TOPK: tl.constexpr,
147
-) -> None:
148
- token_id = tl.program_id(0)
149
- k_block_id = tl.program_id(1)
150
-
151
- k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
152
- k_mask = k_offsets < hidden_size
153
- hidden = tl.load(
154
- hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
155
- mask=k_mask,
156
- other=0.0,
157
- ).to(tl.float32)
158
-
159
- num_groups: tl.constexpr = BLOCK_K // GROUP_K
160
- hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
161
- amax = tl.max(hidden_groups, axis=1)
162
- amax = tl.maximum(amax, 1.0e-4)
163
-
164
- scale = amax / 448.0
165
- scale_bits = scale.to(tl.uint32, bitcast=True)
166
- scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
167
- tl.uint32
168
- )
169
- scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
170
- rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
171
-
172
- hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
173
- scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
174
- scaled = tl.reshape(scaled, [BLOCK_K])
175
- fp8 = scaled.to(tl.float8e4nv)
176
- tl.store(
177
- x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
178
- fp8,
179
- mask=k_mask,
180
- )
181
-
182
- scale_offsets = tl.arange(0, num_groups)
183
- packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
184
- tl.store(
185
- x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
186
- packed_scale,
187
- )
188
-
189
- if k_block_id == 0:
190
- topk_offsets = tl.arange(0, BLOCK_TOPK)
191
- topk_mask = topk_offsets < top_k
192
-
193
- ids = tl.load(
194
- topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
195
- mask=topk_mask,
196
- other=0,
197
- ).to(tl.int64)
198
- tl.store(
199
- topk_idx_out
200
- + token_id * topk_idx_stride_m
201
- + topk_offsets * topk_idx_stride_k,
202
- ids,
203
- mask=topk_mask,
204
- )
205
-
206
- weights = tl.load(
207
- topk_weights
208
- + token_id * topk_weights_stride_m
209
- + topk_offsets * topk_weights_stride_k,
210
- mask=topk_mask,
211
- other=0.0,
212
- )
213
- tl.store(
214
- topk_weights_out
215
- + token_id * topk_weights_out_stride_m
216
- + topk_offsets * topk_weights_out_stride_k,
217
- weights,
218
- mask=topk_mask,
219
- )
220
-
221
-
222
-def _stage_deepseek_v4_mega_moe_inputs(
223
- hidden_states: torch.Tensor,
224
- topk_weights: torch.Tensor,
225
- topk_ids: torch.Tensor,
226
- x_fp8: torch.Tensor,
227
- x_sf: torch.Tensor,
228
- topk_idx_out: torch.Tensor,
229
- topk_weights_out: torch.Tensor,
230
-) -> None:
231
- num_tokens, hidden_size = hidden_states.shape
232
- if num_tokens == 0:
233
- return
234
- if hidden_size % 128 != 0:
235
- raise ValueError(
236
- "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
237
- "a multiple of 128."
238
- )
239
- top_k = topk_ids.shape[1]
240
- if topk_weights.shape != topk_ids.shape:
241
- raise ValueError(
242
- "DeepSeek V4 MegaMoE input staging requires topk_weights and "
243
- "topk_ids to have the same shape."
244
- )
245
-
246
- block_k = 128
247
- grid = (num_tokens, triton.cdiv(hidden_size, block_k))
248
- block_topk = triton.next_power_of_2(top_k)
249
- _deepseek_v4_stage_mega_moe_inputs_kernel[grid](
250
- hidden_states,
251
- x_fp8,
252
- x_sf,
253
- topk_ids,
254
- topk_weights,
255
- topk_idx_out,
256
- topk_weights_out,
257
- hidden_states.stride(0),
258
- hidden_states.stride(1),
259
- x_fp8.stride(0),
260
- x_fp8.stride(1),
261
- x_sf.stride(0),
262
- x_sf.stride(1),
263
- topk_ids.stride(0),
264
- topk_ids.stride(1),
265
- topk_weights.stride(0),
266
- topk_weights.stride(1),
267
- topk_idx_out.stride(0),
268
- topk_idx_out.stride(1),
269
- topk_weights_out.stride(0),
270
- topk_weights_out.stride(1),
271
- hidden_size,
272
- top_k,
273
- BLOCK_K=block_k,
274
- GROUP_K=32,
275
- BLOCK_TOPK=block_topk,
276
- num_warps=4,
277
- )
278
-
279
-
119
def make_deepseek_v4_expert_params_mapping(
120
num_experts: int,
121
) -> list[tuple[str, str, int, str]]:
@@ -542,7 +381,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
381
382
symm_buffer = self.get_symm_buffer()
383
num_tokens = hidden_states.shape[0]
545
- _stage_deepseek_v4_mega_moe_inputs(
384
+ prepare_megamoe_inputs(
385
hidden_states,
386
topk_weights,
387
topk_ids,
vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py
new
+173
@@ -0,0 +1,173 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Triton input-staging kernel for DeepSeek V4 MegaMoE.
4
+
5
+Quantizes hidden states to fp8 with E8M0 group scales and repacks the
6
+routing top-k tensors into the int64/float32 layout that the DeepGEMM
7
+MegaMoE kernels consume.
8
+"""
9
+
10
+import torch
11
+
12
+from vllm.triton_utils import tl, triton
13
+
14
+
15
+@triton.jit
16
+def _prepare_megamoe_inputs_kernel(
17
+ hidden_states,
18
+ x_fp8,
19
+ x_sf,
20
+ topk_ids,
21
+ topk_weights,
22
+ topk_idx_out,
23
+ topk_weights_out,
24
+ hidden_stride_m: tl.constexpr,
25
+ hidden_stride_k: tl.constexpr,
26
+ x_stride_m: tl.constexpr,
27
+ x_stride_k: tl.constexpr,
28
+ x_sf_stride_m: tl.constexpr,
29
+ x_sf_stride_k: tl.constexpr,
30
+ topk_ids_stride_m: tl.constexpr,
31
+ topk_ids_stride_k: tl.constexpr,
32
+ topk_weights_stride_m: tl.constexpr,
33
+ topk_weights_stride_k: tl.constexpr,
34
+ topk_idx_stride_m: tl.constexpr,
35
+ topk_idx_stride_k: tl.constexpr,
36
+ topk_weights_out_stride_m: tl.constexpr,
37
+ topk_weights_out_stride_k: tl.constexpr,
38
+ hidden_size: tl.constexpr,
39
+ top_k: tl.constexpr,
40
+ BLOCK_K: tl.constexpr,
41
+ GROUP_K: tl.constexpr,
42
+ BLOCK_TOPK: tl.constexpr,
43
+) -> None:
44
+ token_id = tl.program_id(0)
45
+ k_block_id = tl.program_id(1)
46
+
47
+ k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
48
+ k_mask = k_offsets < hidden_size
49
+ hidden = tl.load(
50
+ hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
51
+ mask=k_mask,
52
+ other=0.0,
53
+ ).to(tl.float32)
54
+
55
+ num_groups: tl.constexpr = BLOCK_K // GROUP_K
56
+ hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
57
+ amax = tl.max(hidden_groups, axis=1)
58
+ amax = tl.maximum(amax, 1.0e-4)
59
+
60
+ scale = amax / 448.0
61
+ scale_bits = scale.to(tl.uint32, bitcast=True)
62
+ scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
63
+ tl.uint32
64
+ )
65
+ scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
66
+ rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
67
+
68
+ hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
69
+ scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
70
+ scaled = tl.reshape(scaled, [BLOCK_K])
71
+ fp8 = scaled.to(tl.float8e4nv)
72
+ tl.store(
73
+ x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
74
+ fp8,
75
+ mask=k_mask,
76
+ )
77
+
78
+ scale_offsets = tl.arange(0, num_groups)
79
+ packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
80
+ tl.store(
81
+ x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
82
+ packed_scale,
83
+ )
84
+
85
+ if k_block_id == 0:
86
+ topk_offsets = tl.arange(0, BLOCK_TOPK)
87
+ topk_mask = topk_offsets < top_k
88
+
89
+ ids = tl.load(
90
+ topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
91
+ mask=topk_mask,
92
+ other=0,
93
+ ).to(tl.int64)
94
+ tl.store(
95
+ topk_idx_out
96
+ + token_id * topk_idx_stride_m
97
+ + topk_offsets * topk_idx_stride_k,
98
+ ids,
99
+ mask=topk_mask,
100
+ )
101
+
102
+ weights = tl.load(
103
+ topk_weights
104
+ + token_id * topk_weights_stride_m
105
+ + topk_offsets * topk_weights_stride_k,
106
+ mask=topk_mask,
107
+ other=0.0,
108
+ )
109
+ tl.store(
110
+ topk_weights_out
111
+ + token_id * topk_weights_out_stride_m
112
+ + topk_offsets * topk_weights_out_stride_k,
113
+ weights,
114
+ mask=topk_mask,
115
+ )
116
+
117
+
118
+def prepare_megamoe_inputs(
119
+ hidden_states: torch.Tensor,
120
+ topk_weights: torch.Tensor,
121
+ topk_ids: torch.Tensor,
122
+ x_fp8: torch.Tensor,
123
+ x_sf: torch.Tensor,
124
+ topk_idx_out: torch.Tensor,
125
+ topk_weights_out: torch.Tensor,
126
+) -> None:
127
+ num_tokens, hidden_size = hidden_states.shape
128
+ if num_tokens == 0:
129
+ return
130
+ if hidden_size % 128 != 0:
131
+ raise ValueError(
132
+ "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
133
+ "a multiple of 128."
134
+ )
135
+ top_k = topk_ids.shape[1]
136
+ if topk_weights.shape != topk_ids.shape:
137
+ raise ValueError(
138
+ "DeepSeek V4 MegaMoE input staging requires topk_weights and "
139
+ "topk_ids to have the same shape."
140
+ )
141
+
142
+ block_k = 128
143
+ grid = (num_tokens, triton.cdiv(hidden_size, block_k))
144
+ block_topk = triton.next_power_of_2(top_k)
145
+ _prepare_megamoe_inputs_kernel[grid](
146
+ hidden_states,
147
+ x_fp8,
148
+ x_sf,
149
+ topk_ids,
150
+ topk_weights,
151
+ topk_idx_out,
152
+ topk_weights_out,
153
+ hidden_states.stride(0),
154
+ hidden_states.stride(1),
155
+ x_fp8.stride(0),
156
+ x_fp8.stride(1),
157
+ x_sf.stride(0),
158
+ x_sf.stride(1),
159
+ topk_ids.stride(0),
160
+ topk_ids.stride(1),
161
+ topk_weights.stride(0),
162
+ topk_weights.stride(1),
163
+ topk_idx_out.stride(0),
164
+ topk_idx_out.stride(1),
165
+ topk_weights_out.stride(0),
166
+ topk_weights_out.stride(1),
167
+ hidden_size,
168
+ top_k,
169
+ BLOCK_K=block_k,
170
+ GROUP_K=32,
171
+ BLOCK_TOPK=block_topk,
172
+ num_warps=4,
173
+ )