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
[CI] Enable Blackwell Llama4 MoE tests (#26731)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Oct 15, 2025 at 23:02 UTC
f8a0acbdbeb7751cf2a25f5d5191cf1ae2ce6e22
2 files changed
+36
-22
.buildkite/test-pipeline.yaml
+1
-1
@@ -529,7 +529,7 @@ steps:
529
# we can only upgrade after this is resolved
530
# TODO(jerryzh168): resolve the above comment
531
- uv pip install --system torchao==0.13.0
532
- - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/
532
+ - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
533
534
- label: LM Eval Small Models # 53min
535
timeout_in_minutes: 75
tests/quantization/test_blackwell_moe.py
+35
-21
@@ -3,6 +3,7 @@
3
4
import json
5
import os
6
+from typing import Any
7
8
import pytest
9
@@ -24,12 +25,21 @@ def set_test_environment():
25
os.environ["FLASHINFER_NVCC_THREADS"] = "16"
26
27
27
-# dummy_hf_overrides = {"num_layers": 4, "num_hidden_layers": 4,
28
-# "text_config": {"num_layers": 4, "num_hidden_layers": 4}}
29
-dummy_hf_overrides = {"num_layers": 4, "num_hidden_layers": 4}
28
+# Overide the backbone layers to 4 for faster startup
29
+HF_OVERRIDE_TEXT = {
30
+ "num_layers": 4,
31
+ "num_hidden_layers": 4,
32
+}
33
+HF_OVERRIDE_MM = {
34
+ "text_config": {"num_layers": 4, "num_hidden_layers": 4},
35
+}
36
37
32
-def can_initialize(model: str, extra_args: list[str] | None = None):
38
+def can_initialize(
39
+ model: str,
40
+ hf_overrides: dict[str, Any] | None = None,
41
+ extra_args: list[str] | None = None,
42
+):
43
# Server arguments
44
extra_args = extra_args if extra_args is not None else []
45
server_args = [
@@ -50,7 +60,7 @@ def can_initialize(model: str, extra_args: list[str] | None = None):
60
model,
61
server_args,
62
max_wait_seconds=1500, # Due to FlashInfer compile
53
- override_hf_configs=dummy_hf_overrides,
63
+ override_hf_configs=hf_overrides,
64
) as server:
65
client = server.get_client()
66
# Make a simple request to verify the server works
@@ -77,28 +87,33 @@ def can_initialize(model: str, extra_args: list[str] | None = None):
87
def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
88
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
89
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
80
- can_initialize("nvidia/Llama-4-Scout-17B-16E-Instruct-FP8")
90
+ can_initialize(
91
+ "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM
92
+ )
93
94
83
-@pytest.mark.skip(reason="Works, but takes too long to run")
95
def test_llama4_fp8_tensor_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
96
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
97
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
87
- can_initialize("nvidia/Llama-4-Scout-17B-16E-Instruct-FP8")
98
+ can_initialize(
99
+ "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM
100
+ )
101
102
90
-@pytest.mark.skip(reason="Works, but takes too long to run")
103
def test_llama4_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
104
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
105
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
94
- can_initialize("nvidia/Llama-4-Scout-17B-16E-Instruct-FP4")
106
+ can_initialize(
107
+ "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", hf_overrides=HF_OVERRIDE_MM
108
+ )
109
110
97
-@pytest.mark.skip(reason="RuntimeError: No kernel found for the given options")
111
def test_llama4_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
112
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
113
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
101
- can_initialize("nvidia/Llama-4-Scout-17B-16E-Instruct-FP4")
114
+ can_initialize(
115
+ "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", hf_overrides=HF_OVERRIDE_MM
116
+ )
117
118
119
## DeepSeekV3 ##
@@ -106,7 +121,7 @@ def test_llama4_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
121
122
def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch):
123
monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1")
109
- can_initialize("deepseek-ai/DeepSeek-V3.1")
124
+ can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT)
125
126
127
@pytest.mark.skip(
@@ -118,26 +133,25 @@ def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch):
133
def test_deepseek_fp8_block_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
134
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
135
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
121
- can_initialize("deepseek-ai/DeepSeek-V3.1")
136
+ can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT)
137
138
139
def test_deepseek_fp8_block_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
140
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1")
141
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
127
- can_initialize("deepseek-ai/DeepSeek-V3.1")
142
+ can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT)
143
144
145
def test_deepseek_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
146
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
147
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput")
133
- can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2")
148
+ can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2", hf_overrides=HF_OVERRIDE_TEXT)
149
150
136
-@pytest.mark.skip(reason="RuntimeError: No kernel found for the given options")
151
def test_deepseek_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
152
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1")
153
monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency")
140
- can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2")
154
+ can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2", hf_overrides=HF_OVERRIDE_TEXT)
155
156
157
## GPT-OSS ##
@@ -145,14 +159,14 @@ def test_deepseek_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
159
160
def test_gptoss_mxfp4bf16_moe_flashinfer(monkeypatch: pytest.MonkeyPatch):
161
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16", "1")
148
- can_initialize("openai/gpt-oss-20b")
162
+ can_initialize("openai/gpt-oss-20b", hf_overrides=HF_OVERRIDE_TEXT)
163
164
165
def test_gptoss_mxfp4mxfp8_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):
166
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS", "1")
153
- can_initialize("openai/gpt-oss-20b")
167
+ can_initialize("openai/gpt-oss-20b", hf_overrides=HF_OVERRIDE_TEXT)
168
169
170
def test_gptoss_mxfp4mxfp8_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
171
monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
158
- can_initialize("openai/gpt-oss-20b")
172
+ can_initialize("openai/gpt-oss-20b", hf_overrides=HF_OVERRIDE_TEXT)