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

[NVIDIA] Support flashinfer Mega Moe (#31470)

Browse Source 
Co-authored-by: djns99 <40156487+djns99@users.noreply.github.com>
Co-authored-by: 云挚 <ningyunxiao.nyx@antgroup.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
Shu Wang
2026-09-10 00:22:47 -07:00
committed by GitHub
co-authored by djns99 云挚 Yangmin Li Po-Han Huang
parent c0b790cf7f
commit 1b77f498a0
31 changed files with 2021 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/ep/test_flashinfer_dispatcher.py
+181 -5
View File
@@ -4,13 +4,15 @@ import torch
from sglang.srt.distributed import init_distributed_environment
from sglang.srt.distributed.parallel_state import (
destroy_distributed_environment,
destroy_model_parallel,
get_tp_group,
initialize_model_parallel,
)
from sglang.srt.layers.dp_attention import set_dp_buffer_len
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.utils import initialize_moe_config
from sglang.srt.runtime_context import publish
from sglang.srt.runtime_context import get_context, publish
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.test_utils import CustomTestCase
@@ -21,6 +23,7 @@ class TestFlashinferDispatcher(CustomTestCase):
server_args = ServerArgs(model_path="dummy")
server_args.moe_runner_backend = "flashinfer_cutlass"
server_args.moe_a2a_backend = "flashinfer"
cls.server_args = server_args
set_global_server_args_for_scheduler(server_args)
publish(server_args, role="scheduler")
initialize_moe_config()
@@ -41,9 +44,18 @@ class TestFlashinferDispatcher(CustomTestCase):
@classmethod
def tearDownClass(cls):
# Clean up distributed environment
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()
try:
from flashinfer.comm.trtllm_moe_alltoall import MoeAlltoAll
for workspace in MoeAlltoAll._WORKSPACE_CACHE.values():
mnnvl_mem = workspace.get("mnnvl_mem")
if mnnvl_mem is not None and "ptr" in vars(mnnvl_mem):
del mnnvl_mem.ptr
MoeAlltoAll._WORKSPACE_CACHE.clear()
except ImportError:
pass
destroy_model_parallel()
destroy_distributed_environment()
def create_dispatcher(
self, router_topk=2, num_experts=8, num_local_experts=4, hidden_size=128
@@ -58,8 +70,40 @@ class TestFlashinferDispatcher(CustomTestCase):
params_dtype=torch.bfloat16,
)
def set_dispatch_type(self, dispatch_type):
get_context().override(
"test_flashinfer_dispatcher",
flashinfer_a2a_dispatch_type=dispatch_type,
)
def _zero_moe_a2a_dispatch_payloads(self):
# Shared MoeAlltoAll workspaces keep stale recv payloads across tests.
# Zero only the payload region so unused-source == 0 asserts stay valid.
try:
from flashinfer.comm.trtllm_moe_alltoall import (
MoeAlltoAll,
get_moe_alltoall_module,
)
except ImportError:
return
module = get_moe_alltoall_module()
for ws in MoeAlltoAll._WORKSPACE_CACHE.values():
workspace = ws["workspace"]
aux = int(
module.moe_a2a_get_aux_data_size(
ws["ep_size"],
ws["max_num_tokens"],
ws["eplb_stats_num_experts"],
)
)
aux = ((aux + 127) // 128) * 128
if aux < workspace.shape[1]:
workspace[:, aux:].zero_()
def test_dispatch_basic(self):
"""Test basic dispatch functionality"""
self.set_dispatch_type("bf16")
num_tokens = 16
hidden_size = 128
router_topk = 1 # Single expert per token for simplicity
@@ -143,9 +187,10 @@ class TestFlashinferDispatcher(CustomTestCase):
def test_dispatch_with_empty_tokens(self):
"""Test dispatch when there are no tokens (edge case)"""
self.set_dispatch_type("bf16")
# This tests the dummy token handling
num_tokens = 16
hidden_size = 1
hidden_size = 128
router_topk = 1 # Single expert per token for simplicity
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
@@ -195,6 +240,9 @@ class TestFlashinferDispatcher(CustomTestCase):
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
)
self._zero_moe_a2a_dispatch_payloads()
torch.distributed.barrier()
dispatcher = self.create_dispatcher(
router_topk=router_topk,
num_experts=num_experts,
@@ -250,6 +298,7 @@ class TestFlashinferDispatcher(CustomTestCase):
def test_dispatch_with_fp4_quantization(self):
"""Test dispatch with FP4 quantization enabled"""
self.set_dispatch_type("nvfp4")
num_tokens = 128
hidden_size = 128
router_topk = 1 # Single expert per token for simplicity
@@ -312,6 +361,133 @@ class TestFlashinferDispatcher(CustomTestCase):
)
self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
def test_dispatch_with_mxfp8_quantization(self):
"""Test dispatch with MXFP8 quantization enabled"""
self.set_dispatch_type("mxfp8")
num_tokens = 128
hidden_size = 128
router_topk = 1
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
num_experts = world_size
num_local_experts = 1
set_dp_buffer_len(
global_dp_buffer_len=num_tokens * world_size,
local_dp_buffer_len=num_tokens,
dp_max_padding=True,
global_num_tokens=None,
)
hidden_states = torch.randn(
(num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
)
target_rank = (rank + 1) % world_size
target_expert = target_rank
topk_ids = torch.full(
(num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
)
topk_weights = torch.ones(
(num_tokens, router_topk), dtype=torch.float32, device="cuda"
)
from sglang.srt.layers.moe.topk import StandardTopKOutput
topk_output = StandardTopKOutput(
topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
)
dispatcher = self.create_dispatcher(
router_topk=router_topk,
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
)
dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})
dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
self.assertEqual(
dispatch_output.hidden_states.shape,
(num_tokens * world_size, hidden_size),
)
self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
self.assertEqual(dispatch_output.output_dtype, torch.bfloat16)
self.assertIsNotNone(dispatch_output.hidden_states_scale)
self.assertEqual(
dispatch_output.hidden_states_scale.shape,
(num_tokens * world_size, hidden_size // 32),
)
self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
self.assertEqual(
dispatch_output.topk_output.topk_ids.shape,
(num_tokens * world_size, router_topk),
)
self.assertEqual(dispatch_output.topk_output.topk_ids.dtype, torch.int32)
def test_dispatch_with_mxfp8_quantization_and_empty_rank(self):
"""All ranks must contribute the same payload dtypes, including empty ranks."""
self.set_dispatch_type("mxfp8")
num_tokens = 16
hidden_size = 128
router_topk = 1
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
empty_rank = 1
global_num_tokens = [num_tokens] * world_size
global_num_tokens[empty_rank] = 0
set_dp_buffer_len(
global_dp_buffer_len=num_tokens * world_size,
local_dp_buffer_len=num_tokens,
dp_max_padding=False,
global_num_tokens=global_num_tokens,
)
local_tokens = 0 if rank == empty_rank else num_tokens
hidden_states = torch.randn(
(local_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
)
target_expert = (rank + 1) % world_size
topk_ids = torch.full(
(local_tokens, router_topk),
target_expert,
dtype=torch.int32,
device="cuda",
)
topk_weights = torch.ones(
(local_tokens, router_topk), dtype=torch.float32, device="cuda"
)
from sglang.srt.layers.moe.topk import StandardTopKOutput
dispatcher = self.create_dispatcher(
router_topk=router_topk,
num_experts=world_size,
num_local_experts=1,
hidden_size=hidden_size,
)
dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})
self._zero_moe_a2a_dispatch_payloads()
torch.distributed.barrier()
dispatch_output = dispatcher.dispatch(
hidden_states,
StandardTopKOutput(
topk_weights=topk_weights,
topk_ids=topk_ids,
router_logits=None,
),
)
self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
self.assertEqual(
dispatch_output.hidden_states.shape,
(num_tokens * world_size, hidden_size),
)
if __name__ == "__main__":
"""
test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py
+56
View File
@@ -155,6 +155,56 @@ class FlashinferTrtllmGenMoeBackendMXFP8Base:
self.assertGreater(metrics["score"], 0.93)
class FlashinferTrtllmGenMoeBackendMXFP8A2ABase:
backend = "flashinfer_trtllm_routed"
@classmethod
def setUpClass(cls):
cls.model = "zianglih/Qwen3-30B-A3B-Instruct-2507-MXFP8"
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, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"},
other_args=[
"--quantization",
"mxfp8",
"--enable-dp-attention",
"--dp-size",
"4",
"--tp-size",
"4",
"--moe-a2a-backend",
"flashinfer",
"--moe-runner-backend",
cls.backend,
"--flashinfer-a2a-dispatch-type",
"mxfp8",
"--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",
api="completion",
max_tokens=512,
num_examples=200,
num_threads=128,
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["score"], 0.93)
class FlashinferTrtllmGenMoeBackendMXFP8MixedBF16Base:
backend = None
@@ -261,6 +311,12 @@ class TestFlashinferTrtllmGenMoeBackendMXFP8Routed(
backend = "flashinfer_trtllm_routed"
class TestFlashinferTrtllmGenMoeBackendMXFP8A2A(
FlashinferTrtllmGenMoeBackendMXFP8A2ABase, CustomTestCase
):
pass
class TestFlashinferTrtllmRoutedMxfp8MixedBF16(
FlashinferTrtllmGenMoeBackendMXFP8MixedBF16Base, CustomTestCase
):
test/registered/unit/layers/moe/test_flashinfer_dispatcher.py
+62
View File
@@ -0,0 +1,62 @@
import sys
from unittest.mock import patch
import torch
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.topk import StandardTopKOutput
from sglang.srt.layers.moe.utils import FlashinferA2ADispatchType
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
def test_empty_mxfp8_dispatch_uses_same_payload_dtype_as_nonempty_rank():
class FakeMoeAlltoAll:
def dispatch(self, _topk_ids, payloads, *_args, **_kwargs):
self.payload_dtypes = [payload.dtype for payload in payloads]
return payloads
dispatcher = object.__new__(FlashinferDispatcher)
dispatcher.dispatch_type = FlashinferA2ADispatchType.MXFP8
dispatcher.hidden_size = 128
dispatcher.max_num_tokens = 0
dispatcher.ep_size = 1
dispatcher.invalid_token_expert_id = 8
dispatcher.payload_in_workspace = False
dispatcher.quant_config = {"use_mxfp8": True}
dispatcher.moe_a2a = FakeMoeAlltoAll()
hidden_states = torch.empty((0, 128), dtype=torch.bfloat16)
topk_output = StandardTopKOutput(
topk_weights=torch.empty((0, 1), dtype=torch.float32),
topk_ids=torch.empty((0, 1), dtype=torch.int32),
router_logits=None,
)
with (
patch(
"sglang.srt.layers.moe.token_dispatcher.flashinfer.get_dp_global_num_tokens",
return_value=None,
),
patch(
"sglang.srt.layers.moe.token_dispatcher.flashinfer.is_dp_attention_enabled",
return_value=False,
),
):
output = dispatcher.dispatch(hidden_states, topk_output)
assert dispatcher.moe_a2a.payload_dtypes == [
torch.float8_e4m3fn,
torch.uint8,
torch.int32,
torch.float32,
]
assert output.hidden_states.dtype == torch.float8_e4m3fn
assert output.hidden_states_scale.dtype == torch.uint8
if __name__ == "__main__":
import pytest
sys.exit(pytest.main([__file__, "-v"]))
test/registered/unit/layers/moe/test_flashinfer_megamoe.py
+220
View File
@@ -0,0 +1,220 @@
import importlib.util
import sys
import types
from pathlib import Path
import torch
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _load_megamoe_module(monkeypatch):
"""Load the adapter with only its small import-time dependencies stubbed."""
class MoeQuantInfo:
pass
class MoeRunnerConfig:
pass
def register_fused_func(*_args, **_kwargs):
return lambda fn: fn
fake_modules = {
"sglang": types.ModuleType("sglang"),
"sglang.srt": types.ModuleType("sglang.srt"),
"sglang.srt.environ": types.ModuleType("sglang.srt.environ"),
"sglang.srt.layers": types.ModuleType("sglang.srt.layers"),
"sglang.srt.layers.moe": types.ModuleType("sglang.srt.layers.moe"),
"sglang.srt.layers.moe.moe_runner": types.ModuleType(
"sglang.srt.layers.moe.moe_runner"
),
"sglang.srt.layers.moe.moe_runner.base": types.ModuleType(
"sglang.srt.layers.moe.moe_runner.base"
),
"sglang.srt.layers.moe.token_dispatcher": types.ModuleType(
"sglang.srt.layers.moe.token_dispatcher"
),
"sglang.srt.runtime_context": types.ModuleType("sglang.srt.runtime_context"),
"deep_gemm": types.ModuleType("deep_gemm"),
"deep_gemm.utils": types.ModuleType("deep_gemm.utils"),
"deep_gemm.utils.math": types.ModuleType("deep_gemm.utils.math"),
}
fake_modules["sglang.srt.environ"].envs = types.SimpleNamespace(
SGLANG_FLASHINFER_MEGAMOE_MAX_TOKENS_PER_RANK=types.SimpleNamespace(
get=lambda: 0
),
SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE=types.SimpleNamespace(
get=lambda: "bf16"
),
SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE=types.SimpleNamespace(
get=lambda: False
),
)
runtime_context = fake_modules["sglang.srt.runtime_context"]
runtime_context.cutedsl_moe_max_num_tokens = lambda: 2048
base = fake_modules["sglang.srt.layers.moe.moe_runner.base"]
base.MoeQuantInfo = MoeQuantInfo
base.MoeRunnerConfig = MoeRunnerConfig
base.register_fused_func = register_fused_func
token_dispatcher = fake_modules["sglang.srt.layers.moe.token_dispatcher"]
class StandardCombineInput:
def __init__(self, *, hidden_states):
self.hidden_states = hidden_states
token_dispatcher.StandardCombineInput = StandardCombineInput
for name, module in fake_modules.items():
monkeypatch.setitem(sys.modules, name, module)
module_path = (
Path(__file__).resolve().parents[5]
/ "python/sglang/srt/layers/moe/flashinfer_megamoe.py"
)
module_name = "sglang_flashinfer_megamoe_adapter_test"
spec = importlib.util.spec_from_file_location(module_name, module_path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, module_name, module)
spec.loader.exec_module(module)
return module
def test_max_tokens_uses_runtime_context_accessor(monkeypatch):
module = _load_megamoe_module(monkeypatch)
assert module._resolve_max_tokens_per_rank() == 2048
runtime_context = sys.modules["sglang.srt.runtime_context"]
runtime_context.cutedsl_moe_max_num_tokens = lambda: 0
assert module._resolve_max_tokens_per_rank() == 1024
def test_adapter_keeps_router_ids_int32(monkeypatch):
module = _load_megamoe_module(monkeypatch)
class FakeMoEEpTensors:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
fake_moe_ep = types.ModuleType("flashinfer.moe_ep")
fake_moe_ep.MoEEpTensors = FakeMoEEpTensors
fake_flashinfer = types.ModuleType("flashinfer")
fake_flashinfer.moe_ep = fake_moe_ep
monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer)
monkeypatch.setitem(sys.modules, "flashinfer.moe_ep", fake_moe_ep)
hidden_states = torch.randn((3, 4), dtype=torch.bfloat16)
topk_ids = torch.tensor([[0, 1], [1, 0], [0, 1]], dtype=torch.int32)
topk_weights = torch.randn((3, 2), dtype=torch.float32)
output = torch.randn_like(hidden_states)
class Mega:
_workspace = object()
def forward(self, tensors):
self.tensors = tensors
return output
mega = Mega()
dispatch_output = types.SimpleNamespace(
hidden_states=hidden_states,
topk_output=types.SimpleNamespace(
topk_ids=topk_ids,
topk_weights=topk_weights,
),
)
quant_info = module.FlashInferMegaMoeQuantInfo(mega=mega)
runner_config = types.SimpleNamespace(routed_scaling_factor=1.0)
result = module.run_flashinfer_megamoe(
dispatch_output,
quant_info,
runner_config,
)
assert mega.tensors.topk_ids.data_ptr() == topk_ids.data_ptr()
assert mega.tensors.topk_ids.dtype == torch.int32
assert result.hidden_states is output
def test_adapter_requests_workspace_output_view(monkeypatch):
module = _load_megamoe_module(monkeypatch)
class FakeMoEEpTensors:
def __init__(self, **kwargs):
self.__dict__.update(kwargs)
fake_moe_ep = types.ModuleType("flashinfer.moe_ep")
fake_moe_ep.MoEEpTensors = FakeMoEEpTensors
fake_flashinfer = types.ModuleType("flashinfer")
fake_flashinfer.moe_ep = fake_moe_ep
monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer)
monkeypatch.setitem(sys.modules, "flashinfer.moe_ep", fake_moe_ep)
hidden_states = torch.randn((2, 4), dtype=torch.bfloat16)
topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32)
topk_weights = torch.ones((2, 2), dtype=torch.float32)
output = torch.randn_like(hidden_states)
class Mega:
supports_output_view = True
_workspace = object()
def forward(self, tensors, *, return_workspace_view=False):
self.tensors = tensors
self.return_workspace_view = return_workspace_view
return output
mega = Mega()
dispatch_output = types.SimpleNamespace(
hidden_states=hidden_states,
topk_output=types.SimpleNamespace(
topk_ids=topk_ids,
topk_weights=topk_weights,
),
)
result = module.run_flashinfer_megamoe(
dispatch_output,
module.FlashInferMegaMoeQuantInfo(mega=mega),
types.SimpleNamespace(routed_scaling_factor=1.0),
)
assert result.hidden_states is output
assert mega.tensors.topk_ids.data_ptr() == topk_ids.data_ptr()
assert mega.tensors.topk_ids.dtype == torch.int32
assert mega.return_workspace_view is True
def test_capture_safe_ue8m0_pack_is_scoped(monkeypatch):
module = _load_megamoe_module(monkeypatch)
dgm = sys.modules["deep_gemm.utils.math"]
def original(value):
return value
dgm.pack_ue8m0_to_int = original
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
with module._capture_safe_ue8m0_pack():
assert dgm.pack_ue8m0_to_int is original
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
with module._capture_safe_ue8m0_pack():
assert dgm.pack_ue8m0_to_int is not original
packed = dgm.pack_ue8m0_to_int(torch.ones(4, dtype=torch.float32))
assert packed.dtype == torch.int32
assert dgm.pack_ue8m0_to_int is original
if __name__ == "__main__":
import pytest
sys.exit(pytest.main([__file__, "-v"]))
test/registered/unit/lora/test_mem_pool_ep_unit.py
+6 -4
View File
@@ -906,17 +906,19 @@ class TestModuleLevelHelpers(unittest.TestCase):
# Without a specific flashinfer backend selected, default is False.
self.assertFalse(_moe_runner_keeps_global_expert_ids())
def test_real_backend_predicate_matches_dispatcher_and_pool(self):
def test_real_backend_predicates_match_supported_id_contracts(self):
backends = _load_moe_backend_enum()
expected_global = {
dispatcher_global_ids = {
backends.FLASHINFER_TRTLLM,
backends.EXPERIMENTAL_SGL_TRTLLM,
backends.FLASHINFER_TRTLLM_ROUTED,
backends.FLASHINFER_CUTLASS,
backends.FLASHINFER_MXFP4,
backends.FLASHINFER_CUTEDSL,
backends.FLASHINFER_MEGAMOE,
backends.HPC_OPS,
}
lora_global_ids = dispatcher_global_ids - {backends.FLASHINFER_MEGAMOE}
config = types.SimpleNamespace(
num_experts=8,
num_local_experts=2,
@@ -938,11 +940,11 @@ class TestModuleLevelHelpers(unittest.TestCase):
dispatcher = standard_dispatcher(config)
self.assertEqual(
dispatcher.skip_local_expert_mapping,
backend in expected_global,
backend in dispatcher_global_ids,
)
self.assertEqual(
_moe_runner_keeps_global_expert_ids(),
backend in expected_global,
backend in lora_global_ids,
)
test/registered/unit/server_args/test_server_args.py
+262
View File
@@ -76,6 +76,10 @@ from sglang.srt.entrypoints.sidecar import (
)
from sglang.srt.environ import envs
from sglang.srt.layers.cp.base import is_cp_enabled, is_interleave
from sglang.srt.layers.moe.utils import (
FlashinferA2ADispatchType,
get_flashinfer_a2a_dispatch_type,
)
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
CudaGraphConfig,
@@ -1123,6 +1127,264 @@ class TestContextParallelServerArgs(CustomTestCase):
self.assertTrue(is_interleave())
class TestFlashinferA2ADispatchType(CustomTestCase):
def setUp(self):
self._nvfp4_env_backup = os.environ.get("SGLANG_MOE_NVFP4_DISPATCH")
envs.SGLANG_MOE_NVFP4_DISPATCH.clear()
def tearDown(self):
if self._nvfp4_env_backup is None:
envs.SGLANG_MOE_NVFP4_DISPATCH.clear()
else:
os.environ["SGLANG_MOE_NVFP4_DISPATCH"] = self._nvfp4_env_backup
def _make_args(
self,
quantization=None,
dispatch_type=None,
runner_backend="flashinfer_trtllm_routed",
):
server_args = ServerArgs(
model_path="dummy",
quantization=quantization,
moe_a2a_backend="flashinfer",
moe_runner_backend=runner_backend,
flashinfer_a2a_dispatch_type=dispatch_type,
enable_dp_attention=True,
dp_size=4,
tp_size=4,
)
server_args._model_config = SimpleNamespace(nvfp4_moe_meta=None)
return server_args
def test_auto_resolves_mxfp8_and_normalizes_trtllm(self):
server_args = self._make_args(
quantization="mxfp8",
dispatch_type="auto",
runner_backend="flashinfer_trtllm",
)
handle_a2a_moe(server_args)
self.assertEqual(
resolution_result(server_args, "moe_runner_backend"),
"flashinfer_trtllm_routed",
)
self.assertEqual(
resolution_result(server_args, "flashinfer_a2a_dispatch_type"), "mxfp8"
)
def test_auto_resolves_modelopt_fp4_to_nvfp4(self):
server_args = self._make_args(quantization="modelopt_fp4", dispatch_type="auto")
handle_a2a_moe(server_args)
self.assertEqual(
resolution_result(server_args, "flashinfer_a2a_dispatch_type"), "nvfp4"
)
def test_auto_resolves_hybrid_nvfp4_metadata_to_nvfp4(self):
server_args = self._make_args(quantization="fp8", dispatch_type="auto")
server_args._model_config = SimpleNamespace(nvfp4_moe_meta={})
handle_a2a_moe(server_args)
self.assertEqual(
resolution_result(server_args, "flashinfer_a2a_dispatch_type"), "nvfp4"
)
def test_unspecified_preserves_legacy_nvfp4_auto_enable(self):
server_args = self._make_args(quantization="modelopt_fp4")
handle_a2a_moe(server_args)
self.assertIsNone(server_args.flashinfer_a2a_dispatch_type)
self.assertTrue(envs.SGLANG_MOE_NVFP4_DISPATCH.get())
def test_unspecified_getter_preserves_legacy_bf16_fallback(self):
with get_context().override_server_args(
flashinfer_a2a_dispatch_type=None,
quantization="mxfp8",
):
self.assertEqual(
get_flashinfer_a2a_dispatch_type(),
FlashinferA2ADispatchType.BF16,
)
def test_runtime_getter_rejects_unresolved_auto(self):
with get_context().override_server_args(
flashinfer_a2a_dispatch_type="auto",
):
with self.assertRaisesRegex(RuntimeError, "must resolve it"):
get_flashinfer_a2a_dispatch_type()
def test_explicit_nvfp4_checks_hybrid_metadata_for_mxfp8_quantization(self):
server_args = self._make_args(quantization="mxfp8", dispatch_type="nvfp4")
server_args._model_config = SimpleNamespace(nvfp4_moe_meta={})
handle_a2a_moe(server_args)
self.assertEqual(
resolution_result(server_args, "flashinfer_a2a_dispatch_type"), "nvfp4"
)
def test_explicit_bf16_overrides_auto(self):
server_args = self._make_args(quantization="modelopt_fp4", dispatch_type="bf16")
handle_a2a_moe(server_args)
self.assertEqual(
resolution_result(server_args, "flashinfer_a2a_dispatch_type"), "bf16"
)
def test_legacy_env_maps_to_dispatch_type(self):
with envs.SGLANG_MOE_NVFP4_DISPATCH.override("1"):
server_args = self._make_args(quantization="modelopt_fp4")
handle_a2a_moe(server_args)
self.assertIsNone(server_args.flashinfer_a2a_dispatch_type)
with envs.SGLANG_MOE_NVFP4_DISPATCH.override("0"):
server_args = self._make_args(quantization="modelopt_fp4")
handle_a2a_moe(server_args)
self.assertIsNone(server_args.flashinfer_a2a_dispatch_type)
def test_legacy_env_conflicts_with_explicit_cli(self):
with envs.SGLANG_MOE_NVFP4_DISPATCH.override("1"):
server_args = self._make_args(
quantization="modelopt_fp4", dispatch_type="bf16"
)
with self.assertRaisesRegex(
ValueError, "SGLANG_MOE_NVFP4_DISPATCH cannot be set"
):
handle_a2a_moe(server_args)
def test_mxfp8_dispatch_requires_mxfp8_quantization(self):
server_args = self._make_args(quantization="fp8", dispatch_type="mxfp8")
with self.assertRaisesRegex(ValueError, "requires --quantization mxfp8"):
handle_a2a_moe(server_args)
def test_explicit_dispatch_type_requires_flashinfer_a2a(self):
server_args = ServerArgs(
model_path="dummy",
moe_a2a_backend="none",
flashinfer_a2a_dispatch_type="bf16",
)
with self.assertRaisesRegex(ValueError, "requires --moe-a2a-backend"):
handle_a2a_moe(server_args)
class TestFlashinferMegaMoeConfig(CustomTestCase):
def setUp(self):
self._combine_dtype_backup = os.environ.get(
"SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE"
)
self._ikr_backup = os.environ.get(
"SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE"
)
envs.SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE.clear()
envs.SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE.clear()
def tearDown(self):
if self._combine_dtype_backup is None:
envs.SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE.clear()
else:
os.environ["SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE"] = (
self._combine_dtype_backup
)
if self._ikr_backup is None:
envs.SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE.clear()
else:
os.environ["SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE"] = (
self._ikr_backup
)
def _make_args(
self,
architecture="DeepseekV4ForCausalLM",
quantization="modelopt_fp4",
*,
is_fp4_experts=False,
nvfp4_moe_meta=None,
):
server_args = ServerArgs(
model_path="dummy",
quantization=quantization,
moe_a2a_backend="flashinfer_megamoe",
moe_runner_backend="flashinfer_megamoe",
enable_dp_attention=True,
dp_size=4,
tp_size=4,
)
server_args._model_config = SimpleNamespace(
hf_config=SimpleNamespace(architectures=[architecture]),
is_fp4_experts=is_fp4_experts,
nvfp4_moe_meta=nvfp4_moe_meta,
)
return server_args
@patch("sglang.srt.arg_groups.moe_hook.is_sm100_supported", return_value=True)
def test_megamoe_accepts_audited_model_architectures(self, _):
supported = (
"DeepseekV2ForCausalLM",
"DeepseekV3ForCausalLM",
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
"Glm4MoeForCausalLM",
"NemotronHForCausalLM",
"NemotronHPuzzleForCausalLM",
"Qwen2MoeForCausalLM",
"Qwen3MoeForCausalLM",
)
for architecture in supported:
with self.subTest(architecture=architecture):
handle_a2a_moe(self._make_args(architecture))
def test_megamoe_rejects_unaudited_model_architecture(self):
with self.assertRaisesRegex(
ValueError,
"not validated for model architectures.*UnsupportedMoeForCausalLM",
):
handle_a2a_moe(self._make_args("UnsupportedMoeForCausalLM"))
@patch("sglang.srt.arg_groups.moe_hook.is_sm100_supported", return_value=True)
def test_megamoe_accepts_supported_quantization_formats(self, _):
supported = (
{"quantization": "modelopt_fp4"},
{"quantization": "mxfp8"},
{"quantization": "fp8", "is_fp4_experts": True},
{"quantization": "modelopt_mixed", "nvfp4_moe_meta": {}},
)
for config in supported:
with self.subTest(config=config):
handle_a2a_moe(self._make_args(**config))
def test_megamoe_rejects_standard_fp8(self):
with self.assertRaisesRegex(ValueError, "Standard FP8 MoE checkpoints"):
handle_a2a_moe(self._make_args(quantization="fp8"))
@patch("sglang.srt.arg_groups.moe_hook.is_sm100_supported", return_value=False)
def test_megamoe_requires_sm100_for_all_quantization_formats(self, _):
with self.assertRaisesRegex(ValueError, "requires an SM100-family"):
handle_a2a_moe(self._make_args(quantization="fp8", is_fp4_experts=True))
@patch("sglang.srt.arg_groups.moe_hook.is_sm100_supported", return_value=True)
def test_megamoe_combine_dtype_accepts_quantized_values(self, _):
with envs.SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE.override("nvfp4"):
handle_a2a_moe(self._make_args())
with envs.SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE.override("mxfp8"):
handle_a2a_moe(self._make_args())
@patch("sglang.srt.arg_groups.moe_hook.is_sm100_supported", return_value=True)
def test_megamoe_combine_dtype_rejects_invalid_value(self, _):
with envs.SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE.override("fp8"):
with self.assertRaisesRegex(
ValueError, "SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE"
):
handle_a2a_moe(self._make_args())
@patch("sglang.srt.arg_groups.moe_hook.is_sm100_supported", return_value=True)
def test_megamoe_combine_dtype_conflicts_with_ikr(self, _):
with envs.SGLANG_FLASHINFER_MEGAMOE_COMBINE_DTYPE.override("nvfp4"):
with envs.SGLANG_FLASHINFER_MEGAMOE_IN_KERNEL_FC2_REDUCE.override("1"):
with self.assertRaisesRegex(ValueError, "incompatible"):
handle_a2a_moe(self._make_args())
class TestPortArgs(unittest.TestCase):
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
def test_init_new_standard_case(self, mock_temp_file):
test/registered/unit/test_model_overrides.py
+5 -1
View File
@@ -2466,7 +2466,11 @@ class TestGoldenModelOverrides(_IsolatedPublish):
)
def _view(**kw):
defaults = dict(quantization=None, moe_runner_backend="auto")
defaults = dict(
quantization=None,
moe_runner_backend="auto",
moe_a2a_backend="none",
)
defaults.update(kw)
return ResolvedView(SimpleNamespace(**defaults))