minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[FlashInfer v0.6.16] Support FlashInfer CuTe DSL NVFP4 MoE quantization (#28354)

Browse Source
This commit is contained in:
Ziang Li
2026-08-13 17:33:46 -07:00
committed by GitHub
parent c4271c3fe1
commit 9d34c2809f
14 changed files with 405 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/backends/test_flashinfer_nvfp4_online_moe_backend.py
+136
View File
@@ -0,0 +1,136 @@
import os
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=800, stage="nightly", runner_config="4-gpu-b200")
class FlashinferNvFp4OnlineMoeBackendBase:
backend = None
model = None
extra_args = []
extra_env = {}
eval_args = {}
spec_accept_length_threshold = None
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
env={**os.environ, **cls.extra_env, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"},
other_args=[
*cls.extra_args,
"--moe-runner-backend",
cls.backend,
"--cuda-graph-max-bs-decode",
"128",
"--tp-size",
"4",
"--ep-size",
"4",
"--quantization",
"nvfp4_online",
"--mem-fraction-static",
"0.7",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
num_examples=200,
num_threads=128,
**self.eval_args,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.90)
if self.spec_accept_length_threshold is not None:
server_info = requests.get(self.base_url + "/server_info").json()
avg_spec_accept_length = server_info["internal_states"][0][
"avg_spec_accept_length"
]
print(f"{avg_spec_accept_length=}")
self.assertGreater(
avg_spec_accept_length, self.spec_accept_length_threshold
)
# Only this class is affected, but the file runs with failfast, so leaving it
# enabled also cuts off the class that sorts after it.
@unittest.skip(
"flashinfer-ai/flashinfer#4486: on SM100/SM103 the TRTLLM_GEN tile-192 BMM "
"path returns non-finite MoE output from FlashInfer 0.6.16.post4 on, so the "
"first real prefill trips the sampler NaN assert and gsm8k scores 0.0. "
"See #34629 for the package bisect."
)
class TestFlashinferTrtllmGenMoeBackendNvFp4Online(
FlashinferNvFp4OnlineMoeBackendBase, CustomTestCase
):
backend = "flashinfer_trtllm"
model = "Qwen/Qwen3-30B-A3B-Instruct-2507-FP8"
eval_args = {"api": "completion", "max_tokens": 512}
extra_env = {
"FLASHINFER_NVFP4_4OVER6": "1",
"FLASHINFER_NVFP4_4OVER6_ERR_MODE": "MSE",
"FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": "1",
"FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": "1",
"SGLANG_FP4_IGNORED_LAYERS": ",".join(
["shared_expert"]
+ [f"model.layers.{layer_id}" for layer_id in range(40, 48)]
),
}
class TestFlashinferCuteDSLMoeBackendNvFp4Online(
FlashinferNvFp4OnlineMoeBackendBase, CustomTestCase
):
backend = "flashinfer_cutedsl"
model = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8"
extra_args = [
"--reasoning-parser",
"nemotron_3",
"--tool-call-parser",
"qwen3_coder",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
]
eval_args = {"max_tokens": 16000, "temperature": 1.0, "top_p": 0.95}
spec_accept_length_threshold = 2.5
extra_env = {
"FLASHINFER_NVFP4_4OVER6": "1",
"FLASHINFER_NVFP4_4OVER6_ERR_MODE": "MSE",
"FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": "1",
"FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": "1",
}
if __name__ == "__main__":
unittest.main()
test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py
+6 -76
View File
@@ -243,58 +243,6 @@ class FlashinferTrtllmGenMoeBackendNVFP4Base:
self.assertGreater(metrics["score"], 0.89)
class FlashinferTrtllmGenMoeBackendNvFp4OnlineBase:
backend = None
extra_env = {}
@classmethod
def setUpClass(cls):
cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
env={**os.environ, **cls.extra_env, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"},
other_args=[
"--attention-backend",
"triton",
"--moe-runner-backend",
cls.backend,
"--cuda-graph-max-bs-decode",
"128",
"--tp-size",
"4",
"--ep-size",
"2",
"--quantization",
"nvfp4_online",
"--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(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.90)
class TestFlashinferTrtllmGenMoeBackendFP8(
FlashinferTrtllmGenMoeBackendFP8Base, CustomTestCase
):
@@ -325,6 +273,12 @@ class TestFlashinferTrtllmGenMoeBackendBF16Routed(
backend = "flashinfer_trtllm_routed"
@unittest.skip(
"flashinfer-ai/flashinfer#4486: on SM100/SM103 the TRTLLM_GEN tile-192 BMM "
"path returns non-finite MoE output from FlashInfer 0.6.16.post4 on, so the "
"first real prefill trips the sampler NaN assert and gsm8k scores 0.0. "
"See #34629 for the package bisect."
)
class TestFlashinferTrtllmGenMoeBackendNvFp4PerTokenActivationRouted(
FlashinferTrtllmGenMoeBackendNVFP4Base, CustomTestCase
):
@@ -332,29 +286,5 @@ class TestFlashinferTrtllmGenMoeBackendNvFp4PerTokenActivationRouted(
backend = "flashinfer_trtllm_routed"
# Only this class is affected, but the file runs with failfast, so leaving it
# enabled also cuts off the two classes that sort after it.
@unittest.skip(
"flashinfer-ai/flashinfer#4486: on SM100/SM103 the TRTLLM_GEN tile-192 BMM "
"path returns non-finite MoE output from FlashInfer 0.6.16.post4 on, so the "
"first real prefill trips the sampler NaN assert and gsm8k scores 0.0. "
"See #34629 for the package bisect."
)
class TestFlashinferTrtllmGenMoeBackendNvFp4Online(
FlashinferTrtllmGenMoeBackendNvFp4OnlineBase, CustomTestCase
):
extra_env = {
"FLASHINFER_NVFP4_4OVER6": "1",
"FLASHINFER_NVFP4_4OVER6_ERR_MODE": "MSE",
"FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": "1",
"FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": "1",
"SGLANG_FP4_IGNORED_LAYERS": ",".join(
["shared_expert"]
+ [f"model.layers.{layer_id}" for layer_id in range(40, 48)]
),
}
backend = "flashinfer_trtllm"
if __name__ == "__main__":
unittest.main()
test/registered/rl/test_update_weights_from_disk_blackwell.py
+63 -9
View File
@@ -7,6 +7,11 @@ import unittest
import requests
from sglang.srt.constants import (
GPU_MEMORY_TYPE_CUDA_GRAPH,
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
)
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -44,10 +49,17 @@ class UpdateWeightsFromDiskBase:
)
def _launch_server(self, backend_test_suite):
launch_kwargs = {}
if self.launch_env is not None:
launch_kwargs["env"] = self.launch_env
other_args = backend_test_suite.get("other_args")
launch_kwargs = {
"env": {
"SGLANG_MEMORY_SAVER_CUDA_GRAPH": "1",
**(self.launch_env or {}),
}
}
other_args = (
*backend_test_suite.get("other_args", ()),
"--enable-memory-saver",
"--cuda-graph-backend-prefill=disabled",
)
return popen_launch_server(
self.model,
self.base_url,
@@ -140,6 +152,18 @@ class UpdateWeightsFromDiskBase:
timeout=self.update_timeout,
)
def _offload_engine_and_resume_weights(self):
self._post_json("/release_memory_occupation", {})
self._post_json(
"/resume_memory_occupation", {"tags": [GPU_MEMORY_TYPE_WEIGHTS]}
)
def _resume_kv_cache_and_cuda_graph(self):
self._post_json(
"/resume_memory_occupation",
{"tags": [GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH]},
)
def test_parameterized_update_weights_from_disk(self):
for backend_test_suite in self.backend_test_suites:
case_name = backend_test_suite.get("name", "default")
@@ -154,6 +178,7 @@ class UpdateWeightsFromDiskBase:
for update_test_suite in self.update_test_suites:
with self.subTest(case_name=case_name, **update_test_suite):
self._wait_until_idle()
self._offload_engine_and_resume_weights()
ret = self._run_update_weights(
self.model,
flush_cache=update_test_suite["flush_cache"],
@@ -161,6 +186,7 @@ class UpdateWeightsFromDiskBase:
"abort_all_requests"
],
)
self._resume_kv_cache_and_cuda_graph()
self.assertTrue(ret.get("success"), f"{ret=}")
self.assertEqual(self._get_model_info(), self.model)
self._assert_non_empty_decode()
@@ -169,7 +195,7 @@ class UpdateWeightsFromDiskBase:
baseline_sig, updated_sig
)
finally:
kill_process_tree(process.pid)
kill_process_tree(process.pid, wait_timeout=60)
class TestServerUpdateWeightsFromDiskMXFP8(UpdateWeightsFromDiskBase, CustomTestCase):
@@ -179,8 +205,6 @@ class TestServerUpdateWeightsFromDiskMXFP8(UpdateWeightsFromDiskBase, CustomTest
{
"name": "flashinfer_trtllm_routed_mxfp8",
"other_args": (
"--base-gpu-id",
"0",
"--tp-size",
"4",
"--dp-size",
@@ -202,8 +226,6 @@ class TestServerUpdateWeightsFromDiskNVFP4(UpdateWeightsFromDiskBase, CustomTest
{
"name": "flashinfer_trtllm_nvfp4",
"other_args": (
"--base-gpu-id",
"0",
"--tp-size",
"4",
"--fp4-gemm-backend",
@@ -215,5 +237,37 @@ class TestServerUpdateWeightsFromDiskNVFP4(UpdateWeightsFromDiskBase, CustomTest
)
class TestServerUpdateWeightsFromDiskNVFP4CuteDSL(
UpdateWeightsFromDiskBase, CustomTestCase
):
model = "nvidia/Qwen3-30B-A3B-NVFP4"
launch_env = {
"SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1",
"SGLANG_FLASHINFER_MOE_FUSED_FINALIZE": "1",
"FLASHINFER_NVFP4_4OVER6": "1",
"FLASHINFER_NVFP4_4OVER6_ERR_MODE": "MSE",
"FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": "1",
"FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": "1",
}
backend_test_suites = (
{
"name": "flashinfer_cutedsl_nvfp4",
"other_args": (
"--tp-size",
"4",
"--ep-size",
"4",
"--fp4-gemm-backend",
"flashinfer_cutedsl",
"--moe-runner-backend",
"flashinfer_cutedsl",
"--moe-a2a-backend",
"none",
"--enable-deterministic-inference",
),
},
)
if __name__ == "__main__":
unittest.main()