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 14, 05:30 UTC
add tg-mxfp4-moe-test (#22540)
Signed-off-by: siyuanf <siyuanf@nvidia.com> Signed-off-by: Siyuan Fu <siyuanf@nvidia.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Lain committed
Aug 21, 2025 at 10:05 UTC
f8ce022948873a84e6c857c9fc6ac06c9dedc56f
2 files changed
+420
-1
.buildkite/test-pipeline.yaml
+1
@@ -655,6 +655,7 @@ steps:
655
- pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
656
- pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
657
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
658
+ - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
659
# Fusion
660
- pytest -v -s tests/compile/test_fusion_all_reduce.py
661
- pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
tests/kernels/moe/test_mxfp4_moe.py
+419
-1
@@ -4,15 +4,27 @@
4
import importlib
5
import importlib.metadata
6
from dataclasses import dataclass
7
+from typing import Optional
8
9
import pytest
10
import torch
11
from packaging import version
12
13
+from vllm.platforms import current_platform
14
+
15
QUARK_MXFP4_AVAILABLE = importlib.util.find_spec(
16
"quark") is not None and version.parse(
17
importlib.metadata.version("amd-quark")) >= version.parse('0.8.99')
18
19
+TRTLLM_GEN_MXFP4_AVAILABLE = current_platform.is_cuda(
20
+) and current_platform.is_device_capability(100)
21
+
22
+if TRTLLM_GEN_MXFP4_AVAILABLE:
23
+ from flashinfer import (fp4_quantize, mxfp8_quantize,
24
+ next_positive_power_of_2,
25
+ reorder_rows_for_gated_act_gemm, shuffle_matrix_a,
26
+ shuffle_matrix_sf_a, trtllm_fp4_block_scale_moe)
27
+
28
29
@dataclass
30
class ModelCase:
@@ -54,4 +66,410 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
66
67
output = llm.generate_greedy("Today I am in the French Alps and",
68
max_tokens=20)
57
- assert output
\ No newline at end of file
69
+ assert output
70
+
71
+
72
+def swiglu(x,
73
+ alpha: float = 1.702,
74
+ beta: float = 1.0,
75
+ limit: Optional[float] = None):
76
+ # Note we add an extra bias of 1 to the linear layer
77
+ x_glu, x_linear = torch.chunk(x, 2, dim=-1)
78
+ if limit is not None:
79
+ x_glu = x_glu.clamp(max=limit)
80
+ x_linear = x_linear.clamp(min=-limit, max=limit)
81
+ out_glu = x_glu * torch.sigmoid(alpha * x_glu)
82
+ return out_glu * (x_linear + beta)
83
+
84
+
85
+fp4_lookup_table = [
86
+ 0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6
87
+]
88
+
89
+
90
+def mxfp4_dequantize(x, scale):
91
+ assert x.dtype == torch.uint8
92
+ x = x.view(torch.uint8).to(torch.int32)
93
+ x_unpacked = torch.zeros(*x.shape[:-1],
94
+ x.shape[-1] * 2,
95
+ dtype=torch.int32,
96
+ device=x.device)
97
+ x_unpacked[..., 0::2].copy_(x & 0xF)
98
+ x_unpacked[..., 1::2].copy_((x >> 4) & 0xF)
99
+
100
+ x_float = torch.zeros(x_unpacked.shape,
101
+ dtype=torch.float32,
102
+ device=x.device)
103
+ for i, val in enumerate(fp4_lookup_table):
104
+ x_float[x_unpacked == i] = val
105
+
106
+ scale = scale.view(torch.uint8).to(torch.int32)
107
+ scale = (scale << 23).view(torch.float32)
108
+ scale = scale.reshape(*x.shape[:-1], -1)
109
+ scale = torch.stack([scale] * 32, dim=-1).reshape(*x_float.shape)
110
+
111
+ return x_float * scale
112
+
113
+
114
+def mxfp8_dequantize(x, scale):
115
+ assert x.dtype == torch.float8_e4m3fn
116
+ x_float = x.to(torch.float32)
117
+
118
+ scale = scale.view(torch.uint8).to(torch.int32)
119
+ scale = (scale << 23).view(torch.float32)
120
+ scale = scale.reshape(*x.shape[:-1], -1)
121
+ scale = torch.stack([scale] * 32, dim=-1).reshape(*x_float.shape)
122
+
123
+ return x_float * scale
124
+
125
+
126
+def reference_moe(
127
+ roouting_logits,
128
+ topk,
129
+ num_experts,
130
+ hidden_states,
131
+ w13,
132
+ bias13,
133
+ w2,
134
+ bias2,
135
+ alpha,
136
+ beta,
137
+ limit,
138
+ act_type,
139
+):
140
+ # renormalize routing
141
+ experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True)
142
+ expert_weights = torch.nn.functional.softmax(experts.values, dim=1)
143
+ expert_indices = experts.indices
144
+ t = hidden_states.clone()
145
+ # MLP #1
146
+ mlp1_weight = w13[expert_indices, ...]
147
+ mlp1_bias = bias13[expert_indices, ...]
148
+ t = torch.einsum("beck,bk->bec", mlp1_weight, t) + mlp1_bias
149
+ t = swiglu(t, alpha=alpha, beta=beta, limit=limit)
150
+
151
+ if act_type == 'mxfp8':
152
+ t_quantized, t_scale = mxfp8_quantize(t.to(torch.bfloat16),
153
+ is_sf_swizzled_layout=False)
154
+ t = mxfp8_dequantize(t_quantized, t_scale)
155
+ # MLP #2
156
+ mlp2_weight = w2[expert_indices, ...]
157
+ mlp2_bias = bias2[expert_indices, ...]
158
+ t = torch.einsum("beck,bek->bec", mlp2_weight, t) + mlp2_bias
159
+ # Weighted sum of experts
160
+ t = torch.einsum("bec,be->bc", t, expert_weights)
161
+ assert t.shape == hidden_states.shape
162
+ return t.to(torch.bfloat16)
163
+
164
+
165
+def get_tile_tokens_dim(x: torch.Tensor, top_k: int, num_experts: int):
166
+ # Number of tokens in the input tensor.
167
+ num_tokens = x.shape[0]
168
+ # Factor to account for the imbalance of the experts.
169
+ # factor equals to the
170
+ # max_real_num_tokens_per_expert / perfect_num_tokens_per_expert
171
+ # - 1.0 means perfect expert distribution.
172
+ # - > 1.0 means some experts have more
173
+ # tokens than the perfect distribution.
174
+ # - < 1.0 does not make sense.
175
+ imbalance_factor = 1.3
176
+ # Calculate the number of tokens per expert
177
+ # assuming perfect distribution.
178
+ num_tokens_per_expert = (num_tokens * top_k) // num_experts
179
+ # Apply the imbalance factor.
180
+ num_tokens_per_expert = int(num_tokens_per_expert * imbalance_factor)
181
+ # And pad the number to the next power of 2.
182
+ tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert)
183
+ # Cap to 8-64 tokens per CTA tile
184
+ # as it's the range supported by the kernel.
185
+ tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
186
+ return tile_tokens_dim
187
+
188
+
189
+def tg_mxfp4_moe(
190
+ router_logits,
191
+ topk,
192
+ num_experts,
193
+ intermediate_size,
194
+ hidden_size,
195
+ hidden_states,
196
+ hidden_states_scale,
197
+ w13_weight,
198
+ w13_weight_scale,
199
+ w13_bias,
200
+ w2_weight,
201
+ w2_weight_scale,
202
+ w2_bias,
203
+ act_type,
204
+ alpha,
205
+ beta,
206
+ limit,
207
+) -> torch.Tensor:
208
+ sf_block_size = 32
209
+ assert (w13_weight.dim() == 3 and w13_weight.shape[0] == num_experts
210
+ and w13_weight.shape[1] == intermediate_size * 2
211
+ and w13_weight.shape[2] == hidden_size // 2)
212
+ assert (w13_weight_scale.dim() == 3
213
+ and w13_weight_scale.shape[0] == num_experts
214
+ and w13_weight_scale.shape[1] == intermediate_size * 2
215
+ and w13_weight_scale.shape[2] == hidden_size // sf_block_size)
216
+ assert (w2_weight.dim() == 3 and w2_weight.shape[0] == num_experts
217
+ and w2_weight.shape[1] == hidden_size
218
+ and w2_weight.shape[2] == intermediate_size // 2)
219
+ assert (w2_weight_scale.dim() == 3
220
+ and w2_weight_scale.shape[1] == hidden_size
221
+ and w2_weight_scale.shape[2] == intermediate_size // sf_block_size)
222
+ assert (w13_bias.dim() == 2 and w13_bias.shape[0] == num_experts
223
+ and w13_bias.shape[1] == intermediate_size * 2)
224
+ assert (w2_bias.dim() == 2 and w2_bias.shape[0] == num_experts
225
+ and w2_bias.shape[1] == hidden_size)
226
+
227
+ # Swap w1 and w3 as the defenition of
228
+ # swiglu is different in the trtllm-gen
229
+ w13_weight_scale_ = w13_weight_scale.clone()
230
+ w13_weight_ = w13_weight.clone()
231
+ w13_bias_ = w13_bias.clone()
232
+ w13_weight[:, :intermediate_size, :].copy_(
233
+ w13_weight_[:, intermediate_size:, :])
234
+ w13_weight[:, intermediate_size:, :].copy_(
235
+ w13_weight_[:, :intermediate_size, :])
236
+ w13_weight_scale[:, :intermediate_size, :].copy_(
237
+ w13_weight_scale_[:, intermediate_size:, :])
238
+ w13_weight_scale[:, intermediate_size:, :].copy_(
239
+ w13_weight_scale_[:, :intermediate_size, :])
240
+ w13_bias[:, :intermediate_size].copy_(w13_bias_[:, intermediate_size:])
241
+ w13_bias[:, intermediate_size:].copy_(w13_bias_[:, :intermediate_size])
242
+
243
+ # Interleave the weights and scaling factors for activation
244
+ w13_weight_interleaved = []
245
+ w13_weight_scale_interleaved = []
246
+ w13_bias_interleaved = []
247
+ for i in range(num_experts):
248
+ w13_weight_interleaved.append(
249
+ reorder_rows_for_gated_act_gemm(w13_weight[i].clone()))
250
+ w13_weight_scale_interleaved.append(
251
+ reorder_rows_for_gated_act_gemm(w13_weight_scale[i].clone()))
252
+ w13_bias_interleaved.append(
253
+ reorder_rows_for_gated_act_gemm(w13_bias[i].clone().reshape(-1,
254
+ 1)))
255
+ w13_weight = torch.stack(w13_weight_interleaved).reshape(
256
+ num_experts, 2 * intermediate_size, hidden_size // 2)
257
+ w13_weight_scale = torch.stack(w13_weight_scale_interleaved).reshape(
258
+ num_experts, 2 * intermediate_size, hidden_size // 32)
259
+ w13_bias = torch.stack(w13_bias_interleaved).reshape(
260
+ num_experts, 2 * intermediate_size)
261
+
262
+ # Shuffle weights and scaling factors for transposed mma output
263
+ gemm1_weights_shuffled = []
264
+ gemm1_scales_shuffled = []
265
+ gemm2_weights_shuffled = []
266
+ gemm2_scales_shuffled = []
267
+ gemm1_bias_shuffled = []
268
+ gemm2_bias_shuffled = []
269
+ epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
270
+ for i in range(num_experts):
271
+ gemm1_weights_shuffled.append(
272
+ shuffle_matrix_a(w13_weight[i].view(torch.uint8), epilogue_tile_m))
273
+ gemm1_scales_shuffled.append(
274
+ shuffle_matrix_sf_a(w13_weight_scale[i].view(torch.uint8),
275
+ epilogue_tile_m))
276
+
277
+ gemm2_weights_shuffled.append(
278
+ shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m))
279
+ gemm2_scales_shuffled.append(
280
+ shuffle_matrix_sf_a(w2_weight_scale[i].view(torch.uint8),
281
+ epilogue_tile_m))
282
+ gemm1_bias_shuffled.append(
283
+ shuffle_matrix_a(w13_bias[i].reshape(-1, 1), epilogue_tile_m))
284
+ gemm2_bias_shuffled.append(
285
+ shuffle_matrix_a(w2_bias[i].reshape(-1, 1), epilogue_tile_m))
286
+
287
+ w13_weight = torch.stack(gemm1_weights_shuffled)
288
+ w13_weight_scale = torch.stack(gemm1_scales_shuffled).reshape(
289
+ num_experts, 2 * intermediate_size,
290
+ hidden_size // sf_block_size).view(torch.float8_e4m3fn)
291
+ w13_bias = torch.stack(gemm1_bias_shuffled).reshape(num_experts, -1)
292
+
293
+ w2_weight = torch.stack(gemm2_weights_shuffled)
294
+ w2_weight_scale = torch.stack(gemm2_scales_shuffled).reshape(
295
+ num_experts, hidden_size,
296
+ intermediate_size // sf_block_size).view(torch.float8_e4m3fn)
297
+ w2_bias = torch.stack(gemm2_bias_shuffled).reshape(num_experts, -1)
298
+
299
+ tg_result = trtllm_fp4_block_scale_moe(
300
+ routing_logits=router_logits.to(torch.bfloat16),
301
+ routing_bias=None,
302
+ hidden_states=hidden_states,
303
+ hidden_states_scale=hidden_states_scale,
304
+ gemm1_weights=w13_weight,
305
+ gemm1_weights_scale=w13_weight_scale,
306
+ gemm1_bias=w13_bias,
307
+ gemm1_alpha=alpha,
308
+ gemm1_beta=beta,
309
+ gemm1_clamp_limit=limit,
310
+ gemm2_weights=w2_weight,
311
+ gemm2_weights_scale=w2_weight_scale,
312
+ gemm2_bias=w2_bias,
313
+ output1_scale_scalar=None,
314
+ output1_scale_gate_scalar=None,
315
+ output2_scale_scalar=None,
316
+ num_experts=num_experts,
317
+ top_k=topk,
318
+ n_group=None,
319
+ topk_group=None,
320
+ intermediate_size=intermediate_size,
321
+ local_expert_offset=0,
322
+ local_num_experts=num_experts,
323
+ routed_scaling_factor=None,
324
+ tile_tokens_dim=get_tile_tokens_dim(hidden_states, topk, num_experts),
325
+ routing_method_type=1, # renormalize
326
+ do_finalize=True)[0]
327
+ return tg_result
328
+
329
+
330
+def check_accuracy(a, b, atol, rtol, percent):
331
+ """Allow a mismatch percentage of 1 - percent."""
332
+ if torch.any(torch.isnan(a)):
333
+ raise Exception("NaN in reference output")
334
+ if torch.any(torch.isnan(b)):
335
+ raise Exception("NaN in actual output")
336
+ if torch.any(torch.isinf(a)):
337
+ raise Exception("Inf in reference output")
338
+ if torch.any(torch.isinf(b)):
339
+ raise Exception("Inf in actual output")
340
+ assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}"
341
+
342
+ left = torch.abs(a - b)
343
+ right = atol + rtol * torch.abs(b)
344
+ count = torch.sum(left > right)
345
+ mismatch_percent = count / a.numel()
346
+ if mismatch_percent > 1 - percent:
347
+ raise Exception(
348
+ f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
349
+ f"(threshold: {1-percent:.4f})")
350
+
351
+
352
+@pytest.mark.parametrize("topk", [1, 4])
353
+@pytest.mark.parametrize("num_experts", [32, 128])
354
+@pytest.mark.parametrize("num_tokens", [1, 128, 1024])
355
+@pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)])
356
+@pytest.mark.parametrize("alpha,beta,limit", [(1.0, 1.0, None),
357
+ (1.702, 1.0, 7.0)])
358
+@pytest.mark.parametrize("act_type", ['mxfp8', 'bf16'])
359
+@pytest.mark.skipif(
360
+ not TRTLLM_GEN_MXFP4_AVAILABLE,
361
+ reason="nvidia gpu and compute capability sm100 is required for this test")
362
+def test_trtllm_gen_mxfp4_fused_moe(
363
+ topk: int,
364
+ num_experts: int,
365
+ num_tokens: int,
366
+ intermediate_size: int,
367
+ hidden_size: int,
368
+ alpha: float,
369
+ beta: float,
370
+ limit: Optional[float],
371
+ act_type: str,
372
+):
373
+ seed = 42
374
+ torch.manual_seed(seed)
375
+ hidden_states = torch.randn(num_tokens,
376
+ hidden_size,
377
+ device="cuda:0",
378
+ dtype=torch.bfloat16)
379
+ w13 = (torch.randn(num_experts,
380
+ intermediate_size * 2,
381
+ hidden_size,
382
+ device="cuda:0",
383
+ dtype=torch.bfloat16))
384
+ w2 = (torch.randn(num_experts,
385
+ hidden_size,
386
+ intermediate_size,
387
+ device="cuda:0",
388
+ dtype=torch.bfloat16))
389
+ bias13 = torch.randn(num_experts, intermediate_size * 2,
390
+ device="cuda:0") * 10
391
+ bias2 = torch.randn(num_experts, hidden_size, device="cuda:0") * 10
392
+ router_logits = torch.rand(num_tokens, num_experts,
393
+ dtype=torch.float32).cuda()
394
+
395
+ w13, w13_scale = fp4_quantize(w13,
396
+ torch.tensor(1.0, device="cuda:0"),
397
+ 32,
398
+ sf_use_ue8m0=True,
399
+ is_sf_swizzled_layout=False)
400
+ w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
401
+ num_experts, intermediate_size * 2, hidden_size // 32)
402
+ w2, w2_scale = fp4_quantize(w2,
403
+ torch.tensor(1.0, device="cuda:0"),
404
+ 32,
405
+ sf_use_ue8m0=True,
406
+ is_sf_swizzled_layout=False)
407
+ w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
408
+ num_experts, hidden_size, intermediate_size // 32)
409
+ if act_type == 'mxfp8':
410
+ hidden_states, hidden_states_scale = mxfp8_quantize(
411
+ hidden_states, is_sf_swizzled_layout=False)
412
+ hidden_states_scale = hidden_states_scale.view(
413
+ torch.float8_e4m3fn).reshape(-1)
414
+ else:
415
+ hidden_states_scale = None
416
+
417
+ # reference result
418
+ ref_result = torch.empty_like(hidden_states, dtype=torch.bfloat16)
419
+ w13_ref = mxfp4_dequantize(w13.clone(), w13_scale.clone())
420
+ w2_ref = mxfp4_dequantize(w2.clone(), w2_scale.clone())
421
+ bias13_ref = bias13
422
+ bias2_ref = bias2
423
+ if act_type == 'mxfp8':
424
+ hidden_states_ref = mxfp8_dequantize(
425
+ hidden_states, hidden_states_scale).to(torch.float32)
426
+ else:
427
+ hidden_states_ref = hidden_states.to(torch.float32)
428
+ # Process tokens in chunks of 32 to reduce memory usage
429
+ chunk_size = 32
430
+ num_chunks = (num_tokens + chunk_size - 1) // chunk_size
431
+ for i in range(num_chunks):
432
+ start_idx = i * chunk_size
433
+ end_idx = min(start_idx + chunk_size, num_tokens)
434
+ chunk_result = reference_moe(
435
+ router_logits[start_idx:end_idx].to(torch.float32),
436
+ topk,
437
+ num_experts,
438
+ hidden_states_ref[start_idx:end_idx],
439
+ w13_ref,
440
+ bias13_ref,
441
+ w2_ref,
442
+ bias2_ref,
443
+ alpha,
444
+ beta,
445
+ limit,
446
+ act_type,
447
+ )
448
+ ref_result[start_idx:end_idx].copy_(chunk_result)
449
+
450
+ # trtllm-gen result
451
+ if alpha is not None:
452
+ alpha = torch.full((num_experts, ), alpha, device=hidden_states.device)
453
+ if limit is not None:
454
+ limit = torch.full((num_experts, ), limit, device=hidden_states.device)
455
+ if beta is not None:
456
+ beta = torch.full((num_experts, ), beta, device=hidden_states.device)
457
+ tg_result = tg_mxfp4_moe(router_logits,
458
+ topk,
459
+ num_experts,
460
+ intermediate_size,
461
+ hidden_size,
462
+ hidden_states,
463
+ hidden_states_scale,
464
+ w13,
465
+ w13_scale,
466
+ bias13,
467
+ w2,
468
+ w2_scale,
469
+ bias2,
470
+ act_type,
471
+ alpha=alpha,
472
+ beta=beta,
473
+ limit=limit)
474
+ # relatively loose check since the mxfp4 quantization is less accurate
475
+ check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8)