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 5, 09:45 UTC
[Bugfix] FlashInfer MXINT4 MoE crashes, missing do_finalize (#39315)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Benjamin Chislett committed
Apr 8, 2026 at 20:36 UTC
8332078cfdbd5e44e527893b695e79052d008172
2 files changed
+98
-1
tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py
+93
@@ -269,3 +269,96 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group
269
# Note: Different quantization schemes (UINT4b8 vs signed MXINT4) cause
270
# some differences
271
torch.testing.assert_close(marlin_output, trtllm_output, atol=0.3, rtol=6.0)
272
+
273
+
274
+@pytest.mark.skipif(not TRTLLM_GEN_AVAILABLE, reason="Skip for non SM100")
275
+@pytest.mark.parametrize("m", [1, 33])
276
+@pytest.mark.parametrize("n", [7168])
277
+@pytest.mark.parametrize("k", [512])
278
+@pytest.mark.parametrize("e", [384])
279
+@pytest.mark.parametrize("topk", [8])
280
+@torch.inference_mode()
281
+def test_flashinfer_trtllm_mxint4_moe_wrapper(m, n, k, e, topk):
282
+ """Test that the flashinfer_trtllm_mxint4_moe wrapper matches the raw
283
+ trtllm_mxint4_block_scale_moe kernel call."""
284
+ pytest.importorskip("flashinfer")
285
+ from flashinfer import RoutingMethodType
286
+ from flashinfer.fused_moe import trtllm_mxint4_block_scale_moe
287
+
288
+ from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
289
+ flashinfer_trtllm_mxint4_moe,
290
+ )
291
+
292
+ torch.cuda.manual_seed(0)
293
+ dtype = torch.bfloat16
294
+
295
+ a = torch.randn((m, k), device="cuda", dtype=dtype) * 0.5
296
+ router_logits = torch.randn((m, e), device="cuda", dtype=torch.float32) * 1.5
297
+ routing_bias = torch.randn(e, device="cuda", dtype=torch.float32) * 0.8
298
+
299
+ std_w1 = (2.0 / (k + 2 * n)) ** 0.5
300
+ std_w2 = (2.0 / (n + k)) ** 0.5
301
+ w1_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) * std_w1
302
+ w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) * std_w2
303
+
304
+ w1_int4, w1_scales = mxint4_quantize_moe_weights(w1_bf16)
305
+ w2_int4, w2_scales = mxint4_quantize_moe_weights(w2_bf16)
306
+
307
+ prepared = prepare_static_weights_for_trtllm_mxint4_moe(
308
+ gemm1_weights=w1_int4,
309
+ gemm1_scales=w1_scales,
310
+ gemm2_weights=w2_int4,
311
+ gemm2_scales=w2_scales,
312
+ )
313
+
314
+ # Raw kernel call (reference)
315
+ raw_out = trtllm_mxint4_block_scale_moe(
316
+ routing_logits=router_logits.to(torch.float32),
317
+ routing_bias=routing_bias.to(torch.bfloat16),
318
+ hidden_states=a,
319
+ gemm1_weights=prepared["gemm1_weights"].data,
320
+ gemm1_weights_scale=prepared["gemm1_scales"].data,
321
+ gemm1_alpha=None,
322
+ gemm1_beta=None,
323
+ gemm1_clamp_limit=None,
324
+ gemm2_weights=prepared["gemm2_weights"].data,
325
+ gemm2_weights_scale=prepared["gemm2_scales"].data,
326
+ num_experts=e,
327
+ top_k=topk,
328
+ n_group=1,
329
+ topk_group=1,
330
+ intermediate_size=n,
331
+ local_expert_offset=0,
332
+ local_num_experts=e,
333
+ routed_scaling_factor=None,
334
+ routing_method_type=RoutingMethodType.DeepSeekV3,
335
+ enable_pdl=None,
336
+ output=None,
337
+ tune_max_num_tokens=8192,
338
+ )
339
+ if not isinstance(raw_out, torch.Tensor):
340
+ raw_out = raw_out[0]
341
+ raw_out = raw_out.to(dtype)
342
+
343
+ # Wrapper call
344
+ wrapper_out = flashinfer_trtllm_mxint4_moe(
345
+ x=a,
346
+ router_logits=router_logits,
347
+ w13_weight_packed=prepared["gemm1_weights"],
348
+ w13_weight_scale=prepared["gemm1_scales"],
349
+ w2_weight_packed=prepared["gemm2_weights"],
350
+ w2_weight_scale=prepared["gemm2_scales"],
351
+ global_num_experts=e,
352
+ top_k=topk,
353
+ intermediate_size_per_partition=n,
354
+ local_num_experts=e,
355
+ ep_rank=0,
356
+ num_expert_group=1,
357
+ topk_group=1,
358
+ e_score_correction_bias=routing_bias,
359
+ routing_method_type=RoutingMethodType.DeepSeekV3,
360
+ )
361
+
362
+ assert wrapper_out.shape == (m, k)
363
+ assert wrapper_out.dtype == dtype
364
+ torch.testing.assert_close(wrapper_out, raw_out, atol=0.0, rtol=0.0)
vllm/model_executor/layers/quantization/utils/flashinfer_mxint4_moe.py
+5
-1
@@ -259,8 +259,12 @@ def flashinfer_trtllm_mxint4_moe(
259
routed_scaling_factor=None,
260
routing_method_type=routing_method_type,
261
enable_pdl=None,
262
+ do_finalize=True,
263
output=None,
264
tune_max_num_tokens=8192,
264
- ).to(x.dtype)
265
+ )
266
+ if isinstance(out, (tuple, list)):
267
+ out = out[0]
268
+ out = out.to(x.dtype)
269
270
return out