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 Bugfix] Make sure TRTLLM attention is available in test_blackwell_moe (#26188)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Michael Goin committed
Oct 6, 2025 at 13:50 UTC
20db99cc692afcb2f5397a9805acb259c2564d28
1 file changed
+9
-1
tests/quantization/test_blackwell_moe.py
+9
-1
@@ -15,7 +15,15 @@ if not current_platform.is_device_capability(100):
15
"This test only runs on Blackwell GPUs (SM100).", allow_module_level=True
16
)
17
18
-os.environ["FLASHINFER_NVCC_THREADS"] = "16"
18
+
19
+@pytest.fixture(scope="module", autouse=True)
20
+def set_test_environment():
21
+ """Sets environment variables required for this test module."""
22
+ # Make sure TRTLLM attention is available
23
+ os.environ["VLLM_HAS_FLASHINFER_CUBIN"] = "1"
24
+ # Set compilation threads to 16 to speed up startup
25
+ os.environ["FLASHINFER_NVCC_THREADS"] = "16"
26
+
27
28
# dummy_hf_overrides = {"num_layers": 4, "num_hidden_layers": 4,
29
# "text_config": {"num_layers": 4, "num_hidden_layers": 4}}