Unverified Commit 20db99cc authored by Michael Goin's avatar Michael Goin Committed by GitHub
Browse files

[CI Bugfix] Make sure TRTLLM attention is available in test_blackwell_moe (#26188)


Signed-off-by: default avatarmgoin <mgoin64@gmail.com>
Signed-off-by: default avatarMichael Goin <mgoin64@gmail.com>
Co-authored-by: default avatargemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
parent 6431be80
......@@ -15,7 +15,15 @@ if not current_platform.is_device_capability(100):
"This test only runs on Blackwell GPUs (SM100).", allow_module_level=True
)
os.environ["FLASHINFER_NVCC_THREADS"] = "16"
@pytest.fixture(scope="module", autouse=True)
def set_test_environment():
"""Sets environment variables required for this test module."""
# Make sure TRTLLM attention is available
os.environ["VLLM_HAS_FLASHINFER_CUBIN"] = "1"
# Set compilation threads to 16 to speed up startup
os.environ["FLASHINFER_NVCC_THREADS"] = "16"
# dummy_hf_overrides = {"num_layers": 4, "num_hidden_layers": 4,
# "text_config": {"num_layers": 4, "num_hidden_layers": 4}}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment