[NVIDIA] Enable TRTLLM BF16 MoE on Blackwell GPUs (#13798)

This commit is contained in:
Sam
2025-12-11 22:56:13 -08:00
committed by GitHub
parent c05d3afb5d
commit d7ed8a8c24
6 changed files with 198 additions and 41 deletions
@@ -15,7 +15,7 @@ from sglang.test.test_utils import (
register_cuda_ci(est_time=300, suite="nightly-4-gpu-b200", nightly=True)
class TestFlashinferTrtllmGenMoeBackend(CustomTestCase):
class TestFlashinferTrtllmGenMoeBackendFP8(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
@@ -60,5 +60,51 @@ class TestFlashinferTrtllmGenMoeBackend(CustomTestCase):
self.assertGreater(metrics["accuracy"], 0.93)
class TestFlashinferTrtllmGenMoeBackendBF16(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct"
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--attention-backend",
"triton",
"--moe-runner-backend",
"flashinfer_trtllm",
"--cuda-graph-max-bs",
"512",
"--tp-size",
"4",
"--ep-size",
"4",
"--mem-fraction-static",
"0.7",
"--mamba-ssm-dtype",
"bfloat16",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.93)
if __name__ == "__main__":
unittest.main()