@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 275e0d2a99
Fork 0

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 8, 11:15 UTC

[NVIDIA][test] Tests for flashinfer TRTLLM BF16 MoE (#33715)

Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com> Co-authored-by: Pavani Majety <pmajety@nvidia.com>

Linda committed Feb 11, 2026 at 13:38 UTC 275e0d2a993b271cfaec9da87711868719d50d8c
7 files changed +296 -1
tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-fi-cutlass.yaml
+2
@@ -5,3 +5,5 @@ num_fewshot: 5
5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
6 env:
7 VLLM_USE_FLASHINFER_MOE_FP16: "1"
8 + VLLM_FLASHINFER_MOE_BACKEND: "throughput"
9 +
tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-fi-cutlass.yaml
+1
@@ -5,3 +5,4 @@ num_fewshot: 5
5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel"
6 env:
7 VLLM_USE_FLASHINFER_MOE_FP16: "1"
8 + VLLM_FLASHINFER_MOE_BACKEND: "throughput"
tests/kernels/moe/test_flashinfer.py
+41
@@ -318,3 +318,44 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
318 torch.testing.assert_close(
319 output, flashinfer_cutlass_output, atol=5.5e-2, rtol=1e-2
320 )
321 +
322 +
323 +@pytest.mark.parametrize(
324 + "num_experts,intermediate,hidden",
325 + [
326 + (8, 2048, 1536),
327 + (64, 4096, 4096),
328 + ],
329 +)
330 +def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
331 + num_experts, intermediate, hidden
332 +):
333 + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
334 + convert_moe_weights_to_flashinfer_trtllm_block_layout,
335 + )
336 +
337 + w13 = torch.randn(
338 + (num_experts, 2 * intermediate, hidden), dtype=torch.bfloat16, device="cuda"
339 + )
340 + w2 = torch.randn(
341 + (num_experts, hidden, intermediate), dtype=torch.bfloat16, device="cuda"
342 + )
343 +
344 + cache: dict[torch.Size, torch.Tensor] = {}
345 + w13_converted, w2_converted = convert_moe_weights_to_flashinfer_trtllm_block_layout(
346 + cache, w13, w2
347 + )
348 +
349 + assert w13_converted.ndim == 4, (
350 + f"Expected 4D tensor, got shape {w13_converted.shape}"
351 + )
352 + assert w2_converted.ndim == 4, f"Expected 4D tensor, got shape {w2_converted.shape}"
353 +
354 + assert w13_converted.numel() == w13.numel(), "W13 element count should be preserved"
355 + assert w2_converted.numel() == w2.numel(), "W2 element count should be preserved"
356 +
357 + assert w13_converted.dtype == torch.bfloat16
358 + assert w2_converted.dtype == torch.bfloat16
359 +
360 + assert w13_converted.shape[0] == num_experts
361 + assert w2_converted.shape[0] == num_experts
tests/kernels/moe/test_moe.py
+100
@@ -1558,3 +1558,103 @@ def test_batched_fused_marlin_moe(
1558 marlin_output = br.run(a, kwargs)
1559
1560 torch.testing.assert_close(marlin_output, ref_marlin_output, atol=1e-3, rtol=0)
1561 +
1562 +
1563 +@pytest.mark.parametrize("m,n,k", [(32, 1024, 1024)])
1564 +@pytest.mark.parametrize("e,topk", [(8, 2)])
1565 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
1566 +@pytest.mark.skipif(
1567 + not current_platform.is_device_capability_family(100),
1568 + reason="TRTLLM backend test only runs on Blackwell GPUs (SM10x).",
1569 +)
1570 +def test_unquantized_bf16_flashinfer_trtllm_backend(
1571 + m: int,
1572 + n: int,
1573 + k: int,
1574 + e: int,
1575 + topk: int,
1576 + dtype: torch.dtype,
1577 + monkeypatch,
1578 + workspace_init,
1579 +):
1580 + """
1581 + Test BF16 unquantized MoE with FlashInfer TRTLLM backend.
1582 + """
1583 + set_random_seed(7)
1584 +
1585 + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
1586 +
1587 + from vllm.model_executor.layers.fused_moe.config import (
1588 + FusedMoEConfig,
1589 + FusedMoEParallelConfig,
1590 + RoutingMethodType,
1591 + )
1592 + from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1593 + UnquantizedMoeBackend,
1594 + )
1595 + from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
1596 + UnquantizedFusedMoEMethod,
1597 + )
1598 +
1599 + # Setup test data
1600 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
1601 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
1602 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
1603 + router_logits = torch.randn((m, e), device="cuda", dtype=dtype)
1604 +
1605 + moe_config = FusedMoEConfig(
1606 + num_experts=e,
1607 + experts_per_token=topk,
1608 + hidden_dim=k,
1609 + intermediate_size_per_partition=n,
1610 + num_local_experts=e,
1611 + activation="silu",
1612 + device="cuda",
1613 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1614 + in_dtype=dtype,
1615 + is_act_and_mul=True,
1616 + routing_method=RoutingMethodType.Renormalize,
1617 + max_num_tokens=m,
1618 + )
1619 +
1620 + with set_current_vllm_config(vllm_config):
1621 + quant_method = UnquantizedFusedMoEMethod(moe_config)
1622 +
1623 + # Verify TRTLLM backend was selected
1624 + assert (
1625 + quant_method.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
1626 + ), f"Expected FLASHINFER_TRTLLM backend, got {quant_method.unquantized_backend}"
1627 +
1628 + # Verify it's using monolithic path
1629 + assert quant_method.is_monolithic, (
1630 + "FLASHINFER_TRTLLM backend should use monolithic forward"
1631 + )
1632 + layer = torch.nn.Module()
1633 + layer.w13_weight = Parameter(w1.clone(), requires_grad=False)
1634 + layer.w2_weight = Parameter(w2.clone(), requires_grad=False)
1635 + layer.global_num_experts = e
1636 + layer.local_num_experts = e
1637 + layer.top_k = topk
1638 + layer.num_expert_group = 1
1639 + layer.topk_group = 1
1640 + layer.intermediate_size_per_partition = n
1641 + layer.ep_rank = 0
1642 + layer.activation = "silu"
1643 + layer.e_score_correction_bias = None
1644 + layer.routing_method_type = RoutingMethodType.Renormalize
1645 +
1646 + quant_method.process_weights_after_loading(layer)
1647 +
1648 + trtllm_output = quant_method.forward_monolithic_cuda(
1649 + layer=layer,
1650 + x=a,
1651 + router_logits=router_logits,
1652 + )
1653 +
1654 + # Compute torch baseline
1655 + w1_original = w1.clone()
1656 + w2_original = w2.clone()
1657 + baseline_output = torch_moe(a, w1_original, w2_original, router_logits, topk)
1658 +
1659 + close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85)
1660 + assert close.float().mean() > 0.925
tests/kernels/moe/test_unquantized_backend_selection.py new
+132
@@ -0,0 +1,132 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from unittest.mock import patch
4 +
5 +import pytest
6 +
7 +from tests.kernels.moe.utils import make_dummy_moe_config
8 +from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
9 + UnquantizedMoeBackend,
10 + select_unquantized_moe_backend,
11 +)
12 +
13 +
14 +@pytest.mark.parametrize(
15 + "platform_method,expected_backend",
16 + [
17 + ("is_cuda", UnquantizedMoeBackend.TRITON), # Default CUDA without FlashInfer
18 + ("is_rocm", UnquantizedMoeBackend.TRITON),
19 + ("is_cpu", UnquantizedMoeBackend.CPU),
20 + ("is_xpu", UnquantizedMoeBackend.XPU),
21 + ("is_tpu", UnquantizedMoeBackend.TPU),
22 + ("is_out_of_tree", UnquantizedMoeBackend.OOT),
23 + ],
24 +)
25 +@patch(
26 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
27 + return_value=False,
28 +)
29 +def test_select_default_backend_by_platform(
30 + mock_has_flashinfer,
31 + monkeypatch,
32 + platform_method,
33 + expected_backend,
34 +):
35 + """Test backend selection for different platforms."""
36 + with patch(
37 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
38 + ) as mock_platform:
39 + # Set all platform checks to False
40 + mock_platform.is_cuda.return_value = False
41 + mock_platform.is_rocm.return_value = False
42 + mock_platform.is_cpu.return_value = False
43 + mock_platform.is_xpu.return_value = False
44 + mock_platform.is_tpu.return_value = False
45 + mock_platform.is_out_of_tree.return_value = False
46 +
47 + # Set only the specified platform to True
48 + getattr(mock_platform, platform_method).return_value = True
49 +
50 + moe_config = make_dummy_moe_config()
51 + selected_backend = select_unquantized_moe_backend(
52 + moe_config=moe_config,
53 + use_ep=False,
54 + use_dp=False,
55 + )
56 +
57 + assert selected_backend == expected_backend
58 +
59 +
60 +@patch(
61 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
62 + return_value=True,
63 +)
64 +@patch(
65 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
66 + return_value=(True, None),
67 +)
68 +def test_select_cuda_flashinfer_trtllm_backend(
69 + mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
70 +):
71 + """Test CUDA backend selection when FlashInfer TRTLLM is available and enabled."""
72 + with patch(
73 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
74 + ) as mock_platform:
75 + # Set as CUDA platform
76 + mock_platform.is_cuda.return_value = True
77 + mock_platform.is_rocm.return_value = False
78 + mock_platform.is_cpu.return_value = False
79 + mock_platform.is_xpu.return_value = False
80 + mock_platform.is_tpu.return_value = False
81 + mock_platform.is_out_of_tree.return_value = False
82 +
83 + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
84 +
85 + moe_config = make_dummy_moe_config()
86 +
87 + selected_backend = select_unquantized_moe_backend(
88 + moe_config=moe_config,
89 + use_ep=True,
90 + use_dp=False,
91 + )
92 +
93 + assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
94 +
95 +
96 +@patch(
97 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
98 + return_value=True,
99 +)
100 +@patch(
101 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
102 + return_value=(False, None),
103 +)
104 +def test_select_cuda_flashinfer_cutlass_backend(
105 + mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
106 +):
107 + """Test CUDA backend selection when FlashInfer TRTLLM is not available
108 + and FlashInfer CUTLASS is available."""
109 + with patch(
110 + "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
111 + ) as mock_platform:
112 + # Set as CUDA platform with Hopper capability
113 + mock_platform.is_cuda.return_value = True
114 + mock_platform.is_rocm.return_value = False
115 + mock_platform.is_cpu.return_value = False
116 + mock_platform.is_xpu.return_value = False
117 + mock_platform.is_tpu.return_value = False
118 + mock_platform.is_out_of_tree.return_value = False
119 + mock_platform.has_device_capability.return_value = True # SM90+
120 +
121 + # Enable FlashInfer via env var
122 + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
123 +
124 + moe_config = make_dummy_moe_config()
125 +
126 + selected_backend = select_unquantized_moe_backend(
127 + moe_config=moe_config,
128 + use_ep=True, # CUTLASS requires EP
129 + use_dp=False, # CUTLASS doesn't support DP
130 + )
131 +
132 + assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS
tests/quantization/test_blackwell_moe.py
+8
@@ -178,3 +178,11 @@ def test_gptoss_eager(monkeypatch: pytest.MonkeyPatch):
178 hf_overrides=HF_OVERRIDE_TEXT,
179 extra_args=["--enforce-eager"],
180 )
181 +
182 +
183 +## Qwen3 Next ##
184 +
185 +
186 +def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
187 + monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
188 + can_initialize("Qwen/Qwen3-Next-80B-A3B-Instruct", hf_overrides=HF_OVERRIDE_TEXT)
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+12 -1
@@ -78,7 +78,10 @@ def select_unquantized_moe_backend(
78 activation_format=activation_format,
79 )
80 flashinfer_trtllm_moe_enabled = (
81 - has_flashinfer() and envs.VLLM_USE_FLASHINFER_MOE_FP16 and trtllm_supported
81 + has_flashinfer()
82 + and envs.VLLM_USE_FLASHINFER_MOE_FP16
83 + and trtllm_supported
84 + and envs.VLLM_FLASHINFER_MOE_BACKEND == "latency"
85 )
86 # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
87 flashinfer_cutlass_moe_enabled = (
@@ -98,11 +101,19 @@ def select_unquantized_moe_backend(
101 backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM
102 elif flashinfer_cutlass_moe_enabled:
103 backend = UnquantizedMoeBackend.FLASHINFER_CUTLASS
104 + if trtllm_supported:
105 + logger.info_once(
106 + "FlashInfer TRTLLM MoE is available but not enabled, "
107 + "consider setting VLLM_FLASHINFER_MOE_BACKEND=latency "
108 + "to enable it for better performance.",
109 + scope="local",
110 + )
111 else:
112 if not envs.VLLM_USE_FLASHINFER_MOE_FP16 and trtllm_supported:
113 logger.info_once(
114 "FlashInfer TRTLLM MoE is available but not enabled, "
115 "consider setting VLLM_USE_FLASHINFER_MOE_FP16=1 "
116 + "and VLLM_FLASHINFER_MOE_BACKEND=latency "
117 "to enable it for better performance.",
118 scope="local",
119 )