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

[Qwen3.8] Enable NVIDIA NVFP4 on DGX Spark with file-backed PLE and PDL router fix (#39126)

Browse Source 
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: rdxa <rdxa@rdxa-int-spark-01.yvb.moe>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: Manrique <nanomlm@gmail.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>
This commit is contained in:
Jimmy Shong
2026-09-13 16:23:41 +08:00
committed by GitHub
co-authored by Claude Fable 5.1 rdxa Yangmin Li Manrique yhyang201
parent d6fabb74b4
commit cebca698e2
21 changed files with 1376 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernel/embeddings/test_qwen4_ple_offload.py
+76 -2
View File
@@ -1,3 +1,5 @@
import os
import tempfile
from types import SimpleNamespace
import pytest
@@ -79,11 +81,11 @@ def _make_source_embedding(
)
def _load_rows(offloaded, rows):
def _load_rows(offloaded, rows, *, pinned=True):
pointer = offloaded.weight.data_ptr()
offloaded.weight_loader(offloaded.weight, rows)
assert offloaded.weight.data_ptr() == pointer
assert offloaded.weight.is_pinned()
assert offloaded.weight.is_pinned() == pinned
assert offloaded.weight.weight_loader.__self__ is offloaded
assert offloaded.quant_method is None
@@ -189,6 +191,78 @@ def test_qwen4_ple_prefetch_buffer_lifecycle(monkeypatch):
assert set(layer._graph_prefetch_buffers) == {3, 5}
def _file_backend_supported() -> bool:
from sglang.srt.models.qwen4_exp_ple_table import device_uses_host_page_tables
return (
torch.cuda.is_available()
and device_uses_host_page_tables(torch.cuda.current_device()) is True
)
@pytest.mark.skipif(
not _file_backend_supported(),
reason="the file backend needs pageable host memory reachable through host page tables",
)
@pytest.mark.parametrize("embedding_dim", [7, 160])
def test_qwen4_ple_file_backend_matches_pinned(embedding_dim):
with tempfile.TemporaryDirectory() as table_dir:
pinned = Qwen4ExpPinnedHostEmbedding(
_make_source_embedding(embedding_dim=embedding_dim)
)
filed = Qwen4ExpPinnedHostEmbedding(
_make_source_embedding(embedding_dim=embedding_dim),
backend="file",
table_dir=table_dir,
)
assert pinned._file_prefetcher is None and filed._file_prefetcher is not None
(name,) = os.listdir(table_dir)
assert "rows0-8" in name # this rank's vocabulary shard
rows = torch.arange(
8 * embedding_dim, dtype=torch.bfloat16, device="cuda"
).reshape(8, embedding_dim)
_load_rows(pinned, rows)
_load_rows(filed, rows, pinned=False)
ids = torch.tensor([[0, 7, 3], [4, 1, 6]], dtype=torch.int64, device="cuda")
torch.testing.assert_close(filed(ids), pinned(ids), rtol=0, atol=0)
# A prefill-sized gather goes through the page-cache hint path.
big = torch.randint(0, 8, (4096,), device="cuda")
torch.testing.assert_close(
filed(big), rows.index_select(0, big), rtol=0, atol=0
)
@pytest.mark.skipif(
not _file_backend_supported(),
reason="the file backend needs pageable host memory reachable through host page tables",
)
def test_qwen4_ple_file_backend_fp8_table():
embedding_dim = 160
with tempfile.TemporaryDirectory() as table_dir:
filed = Qwen4ExpPinnedHostEmbedding(
_make_source_embedding(
embedding_dim=embedding_dim, dtype=torch.float8_e4m3fn
),
backend="file",
table_dir=table_dir,
)
assert filed.weight.dtype == torch.float8_e4m3fn
rows = (
torch.arange(8 * embedding_dim, dtype=torch.float32, device="cuda").reshape(
8, embedding_dim
)
/ 64
).to(torch.float8_e4m3fn)
_load_rows(filed, rows, pinned=False)
ids = torch.tensor([[0, 7, 3]], dtype=torch.int64, device="cuda")
expected = (
rows.index_select(0, ids.flatten())
.to(torch.bfloat16)
.reshape(1, 3, embedding_dim)
)
torch.testing.assert_close(filed(ids), expected, rtol=0, atol=0)
if __name__ == "__main__":
import sys
test/registered/unit/layers/quantization/test_fp8_moe_runner_fallback.py
+38
View File
@@ -0,0 +1,38 @@
"""Fp8MoEMethod builds a triton runner when the global MoE runner backend is
flashinfer_cutlass or flashinfer_cutedsl, which have no fp8 MoE path."""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
from unittest.mock import patch
import sglang.srt.layers.quantization.fp8 as fp8
from sglang.srt.layers.moe import MoeRunnerBackend, MoeRunnerConfig
from sglang.test.test_utils import CustomTestCase
class TestFp8MoeRunnerFallback(CustomTestCase):
def _runner_backend_for(self, global_backend):
method = fp8.Fp8MoEMethod.__new__(fp8.Fp8MoEMethod)
with patch.object(fp8, "get_moe_runner_backend", return_value=global_backend):
method.create_moe_runner(layer=None, moe_runner_config=MoeRunnerConfig())
return method.runner.runner_backend
def test_flashinfer_cutlass_falls_back_to_triton(self):
self.assertTrue(
self._runner_backend_for(MoeRunnerBackend.FLASHINFER_CUTLASS).is_triton()
)
def test_flashinfer_cutedsl_falls_back_to_triton(self):
self.assertTrue(
self._runner_backend_for(MoeRunnerBackend.FLASHINFER_CUTEDSL).is_triton()
)
def test_triton_is_kept(self):
self.assertTrue(self._runner_backend_for(MoeRunnerBackend.TRITON).is_triton())
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_loader/test_modelopt_loader.py
+56 -1
View File
@@ -22,7 +22,12 @@ from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method
from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.quantization.fp8 import (
Fp8Config,
Fp8LinearMethod,
Fp8MoEMethod,
)
from sglang.srt.layers.quantization.modelopt_quant import (
ModelOptFp4Config,
ModelOptFp4LinearMethod,
@@ -1188,6 +1193,56 @@ class TestModelOptMixedPrecisionConfig(CustomTestCase):
"FP8",
)
def test_mixed_precision_resolves_vl_language_model_keys(self):
# nvidia/Qwen3.8-Flash-Next-NVFP4 keys the text stack as
# `model.language_model.*` while Qwen4-Exp modules are `model.*`.
quant_config = ModelOptMixedPrecisionConfig.from_config(
{
"quant_algo": "MIXED_PRECISION",
"quantized_layers": {
"model.language_model.layers.3.mlp.experts": {
"quant_algo": "NVFP4",
"group_size": 16,
},
"model.language_model.layers.1.ple.ple_embedding.ngram_embedding": {
"quant_algo": "FP8"
},
"mtp.layers.0.mlp.experts": {
"quant_algo": "FP8_BLOCK_SCALES",
"group_size": 128,
},
},
}
)
self.assertEqual(quant_config.exclude_modules, [])
moe = FusedMoE.__new__(FusedMoE)
self.assertIsInstance(
quant_config.get_quant_method(moe, "mtp.layers.0.mlp.experts"),
Fp8MoEMethod,
)
self.assertEqual(
quant_config.get_quant_method(
moe, "mtp.layers.0.mlp.experts"
).quant_config.weight_block_size,
[128, 128],
)
self.assertEqual(
quant_config.resolve_quant_algo("model.layers.3.mlp.experts"), "NVFP4"
)
self.assertEqual(
quant_config.resolve_quant_algo(
"model.layers.1.ple.ple_embedding.ngram_embedding"
),
"FP8",
)
self.assertIsNone(
quant_config.resolve_quant_algo("model.layers.1.ple.key_proj")
)
self.assertIsNone(
quant_config.resolve_quant_algo("model.layers.3.mlp.shared_expert")
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/models/test_qwen4_exp_ple_table.py
+332
View File
@@ -0,0 +1,332 @@
"""File-backed host storage for the offloaded Qwen4-Exp PLE table.
CPU part: the allocator builds a sparse file of exactly the table's size, hands
back a tensor with the requested shape/dtype whose writes land in the file and
survive a re-open, reuses the file across calls, replaces one of the wrong size,
and the prefetcher computes the right page set and honours its size floor. The
resident-set trimmer measures only its own mapping, drops its pages once over
budget without losing what was written through them, and is off when the budget
is zero or the mapping is pinned.
GPU part (skipped unless the device reads pageable host memory through the host
page tables, i.e. unified-memory parts such as GB10): the production Triton
gather kernel reading from the file-backed table matches a torch gather.
"""
import os
import tempfile
import unittest
from unittest import mock
import torch
from sglang.srt.models.qwen4_exp_ple_table import (
PleFilePrefetcher,
PleFileRssTrimmer,
_mapping_rss_bytes,
allocate_ple_host_table,
default_ple_table_dir,
device_uses_host_page_tables,
make_ple_file_prefetcher,
make_ple_file_rss_trimmer,
ple_table_file_name,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
class TestPleFileTableAllocator(CustomTestCase):
def test_file_is_sparse_and_sized_exactly(self):
with tempfile.TemporaryDirectory() as d:
table = allocate_ple_host_table((1000, 160), torch.float8_e4m3fn, "file", d)
path = os.path.join(
d, ple_table_file_name((1000, 160), torch.float8_e4m3fn)
)
self.assertTrue(os.path.exists(path))
self.assertEqual(os.path.getsize(path), 1000 * 160)
self.assertEqual(tuple(table.shape), (1000, 160))
self.assertEqual(table.dtype, torch.float8_e4m3fn)
# Sparse: nothing written yet, so (almost) no blocks allocated.
self.assertLess(os.stat(path).st_blocks * 512, 64 * 1024)
def test_writes_persist_and_file_is_reused(self):
with tempfile.TemporaryDirectory() as d:
shape, dtype = (64, 32), torch.bfloat16
table = allocate_ple_host_table(shape, dtype, "file", d)
row = torch.arange(32, dtype=torch.float32).to(dtype)
table[7].copy_(row) # what the weight loader does, row by row
del table
again = allocate_ple_host_table(shape, dtype, "file", d)
self.assertTrue(torch.equal(again[7].float(), row.float()))
self.assertEqual(len(os.listdir(d)), 1)
def test_wrong_sized_file_is_replaced(self):
with tempfile.TemporaryDirectory() as d:
path = os.path.join(d, ple_table_file_name((8, 8), torch.bfloat16))
with open(path, "wb") as f:
f.write(b"\x01" * 10)
table = allocate_ple_host_table((8, 8), torch.bfloat16, "file", d)
self.assertEqual(os.path.getsize(path), 8 * 8 * 2)
self.assertEqual(tuple(table.shape), (8, 8))
def test_tag_separates_tensor_parallel_shards(self):
with tempfile.TemporaryDirectory() as d:
a = allocate_ple_host_table(
(8, 8), torch.bfloat16, "file", d, tag="rows0-8"
)
b = allocate_ple_host_table(
(8, 8), torch.bfloat16, "file", d, tag="rows8-16"
)
a.fill_(1.0)
self.assertEqual(len(os.listdir(d)), 2)
self.assertTrue(torch.all(b.float() == 0.0))
self.assertIn(
"rows8-16", ple_table_file_name((8, 8), torch.bfloat16, "rows8-16")
)
def test_default_dir_is_per_checkpoint(self):
with mock.patch.dict(os.environ, {"SGLANG_QWEN4_PLE_FILE_DIR": "/cache/ple"}):
a = default_ple_table_dir("RadixArk/Qwen3.8-Flash-Next-NVFP4")
b = default_ple_table_dir("/root/.cache/huggingface/flashnext-fp8/")
self.assertEqual(a, "/cache/ple/RadixArk_Qwen3.8-Flash-Next-NVFP4")
self.assertEqual(b, "/cache/ple/root_.cache_huggingface_flashnext-fp8")
self.assertNotEqual(a, b)
def test_unknown_backend_rejected(self):
with self.assertRaises(ValueError):
allocate_ple_host_table((4, 4), torch.bfloat16, "nvme", None)
def test_pinned_backend_unchanged(self):
if not torch.cuda.is_available():
self.skipTest("pinned memory needs a CUDA runtime")
table = allocate_ple_host_table((4, 4), torch.bfloat16, "pinned", None)
self.assertTrue(table.is_pinned())
self.assertIsNone(make_ple_file_prefetcher(table))
class TestPleFilePrefetcher(CustomTestCase):
def test_page_set_covers_row_start_and_end(self):
# 160-byte rows: row 25 spans bytes 4000-4159, i.e. pages 0 and 1.
pages = PleFilePrefetcher.pages_for_rows(torch.tensor([25, 0]), 160)
self.assertEqual(pages, [0, 1])
pages = PleFilePrefetcher.pages_for_rows(torch.tensor([1000, 1000]), 160)
self.assertEqual(pages, [39]) # dedup, single page
def test_enqueue_uses_local_tp_offsets_and_ignores_other_shards(self):
with tempfile.TemporaryDirectory() as d:
path = os.path.join(d, "t.bin")
with open(path, "wb") as f:
f.truncate(8192)
pf = PleFilePrefetcher(path, row_bytes=160, min_rows=1)
try:
with mock.patch("os.posix_fadvise") as fadvise:
self.assertFalse(
pf.enqueue(
torch.tensor([999, 1032]), vocab_start=1000, vocab_end=1032
)
)
self.assertTrue(
pf.enqueue(
torch.tensor([999, 1000, 1025, 1032]),
vocab_start=1000,
vocab_end=1032,
)
)
pf._pool.shutdown(wait=True)
self.assertEqual(
sorted(c.args[1] for c in fadvise.call_args_list), [0, 4096]
)
finally:
pf.close()
def test_enqueue_respects_min_rows_and_advises_pages(self):
with tempfile.TemporaryDirectory() as d:
path = os.path.join(d, "t.bin")
with open(path, "wb") as f:
f.truncate(1 << 20)
pf = PleFilePrefetcher(path, row_bytes=160, min_rows=4)
try:
self.assertFalse(pf.enqueue(torch.tensor([1, 2, 3])))
with mock.patch("os.posix_fadvise") as fadvise:
self.assertTrue(pf.enqueue(torch.tensor([0, 1, 2, 30])))
pf._pool.shutdown(wait=True)
offsets = sorted(c.args[1] for c in fadvise.call_args_list)
# rows 0-2 live in page 0; row 30 (bytes 4800-4959) in page 1
self.assertEqual(offsets, [0, 4096])
finally:
pf.close()
class TestSmapsParsing(CustomTestCase):
"""The parser runs anywhere; only the live mapping needs Linux."""
SMAPS = """00400000-00401000 r--p 00000000 08:01 1 /usr/bin/x
Size: 4 kB
Rss: 4 kB
7f0000000000-7f0004000000 rw-s 00000000 08:01 2 /cache/ple/table.bin
Size: 65536 kB
Rss: 32768 kB
7f0004000000-7f0008000000 rw-s 00000000 08:01 2 /cache/ple/table.bin
Size: 65536 kB
Rss: 1024 kB
7f0100000000-7f0100001000 rw-p 00000000 00:00 0
Size: 4 kB
Rss: 4 kB
"""
def _rss(self, addr, nbytes):
with tempfile.NamedTemporaryFile("w", suffix=".smaps", delete=False) as f:
f.write(self.SMAPS)
path = f.name
try:
return _mapping_rss_bytes(addr, nbytes, smaps_path=path)
finally:
os.unlink(path)
def test_sums_every_vma_of_the_table_and_nothing_else(self):
# The table spans both of its VMAs; the unrelated ones must not count.
self.assertEqual(self._rss(0x7F0000000000, 0x8000000), (32768 + 1024) * 1024)
def test_counts_a_partially_overlapping_vma(self):
# A range ending inside the first VMA still needs that VMA's pages.
self.assertEqual(self._rss(0x7F0000000000, 0x1000), 32768 * 1024)
def test_ignores_unrelated_mappings(self):
self.assertEqual(self._rss(0x7F0200000000, 0x1000), 0)
def test_missing_smaps_is_reported_as_unknown(self):
self.assertIsNone(
_mapping_rss_bytes(0x1000, 0x1000, smaps_path="/nonexistent/smaps")
)
class TestPleFileRssTrimmerConfig(CustomTestCase):
def test_budget_zero_disables_the_trimmer(self):
with tempfile.TemporaryDirectory() as d:
table = allocate_ple_host_table((64, 32), torch.bfloat16, "file", d)
with mock.patch.dict(
os.environ, {"SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB": "0"}
):
self.assertIsNone(make_ple_file_rss_trimmer(table))
def test_pinned_table_has_no_trimmer(self):
if not torch.cuda.is_available():
self.skipTest("pinned memory needs a CUDA runtime")
table = allocate_ple_host_table((4, 4), torch.bfloat16, "pinned", None)
self.assertIsNone(make_ple_file_rss_trimmer(table))
@unittest.skipUnless(
os.path.exists("/proc/self/smaps"),
"the resident set of a mapping is only readable on Linux",
)
class TestPleFileRssTrimmer(CustomTestCase):
# 64 MiB: large enough that the Rss of the mapping stands out, small
# enough to write in a CPU test.
SHAPE = (32768, 1024)
NBYTES = 32768 * 1024 * 2
def _trimmer(self, table, budget_bytes):
return PleFileRssTrimmer(
addr=table.data_ptr(),
nbytes=self.NBYTES,
budget_bytes=budget_bytes,
interval_s=3600.0,
chunk_bytes=16 << 20, # several chunks, as in production
)
def test_measures_its_own_mapping_only(self):
with tempfile.TemporaryDirectory() as d:
table = allocate_ple_host_table(self.SHAPE, torch.bfloat16, "file", d)
trimmer = self._trimmer(table, 0)
empty = trimmer.mapping_rss_bytes()
table.fill_(1.0) # touches every page
touched = trimmer.mapping_rss_bytes()
self.assertIsNotNone(touched)
self.assertGreater(touched, empty)
self.assertGreater(touched, self.NBYTES // 2)
# Never the whole process: only the VMAs backing this table.
self.assertLessEqual(touched, self.NBYTES + (16 << 20))
def test_over_budget_drops_pages_and_keeps_the_data(self):
with tempfile.TemporaryDirectory() as d:
table = allocate_ple_host_table(self.SHAPE, torch.bfloat16, "file", d)
table.fill_(1.0)
table[7][3] = 2.0
trimmer = self._trimmer(table, budget_bytes=1 << 20)
before = trimmer.mapping_rss_bytes()
freed = trimmer.trim_once()
after = trimmer.mapping_rss_bytes()
self.assertGreater(freed, 0)
self.assertLess(after, before // 2)
# MADV_DONTNEED on a shared file mapping drops the page-table
# entries, not the page cache: the writes are still there.
self.assertEqual(table[7][3].item(), 2.0)
self.assertEqual(table[9][3].item(), 1.0)
def test_under_budget_is_a_no_op(self):
with tempfile.TemporaryDirectory() as d:
table = allocate_ple_host_table(self.SHAPE, torch.bfloat16, "file", d)
table.fill_(1.0)
trimmer = self._trimmer(table, budget_bytes=self.NBYTES * 4)
before = trimmer.mapping_rss_bytes()
self.assertEqual(trimmer.trim_once(), 0)
self.assertEqual(trimmer.mapping_rss_bytes(), before)
def test_factory_starts_and_stops_a_thread(self):
with tempfile.TemporaryDirectory() as d:
table = allocate_ple_host_table((64, 32), torch.bfloat16, "file", d)
with mock.patch.dict(
os.environ,
{
"SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB": "1",
"SGLANG_QWEN4_PLE_FILE_RSS_INTERVAL_S": "3600",
},
):
trimmer = make_ple_file_rss_trimmer(table)
self.assertIsNotNone(trimmer)
try:
self.assertTrue(trimmer._thread.is_alive())
finally:
trimmer.close()
trimmer._thread.join(timeout=5)
self.assertFalse(trimmer._thread.is_alive())
@unittest.skipUnless(
torch.cuda.is_available() and device_uses_host_page_tables(0) is True,
"needs a device that reads pageable host memory through the host page tables",
)
class TestPleFileTableGatherOnDevice(CustomTestCase):
def test_triton_gather_reads_file_backed_table(self):
import triton
from sglang.srt.models.qwen4_exp import (
_gather_ple_embedding_from_pinned_kernel,
)
rows, dim = 4096, 160
with tempfile.TemporaryDirectory() as d:
table = allocate_ple_host_table((rows, dim), torch.bfloat16, "file", d)
table.copy_(torch.randn(rows, dim).to(torch.bfloat16))
ids = torch.randint(0, rows, (2048,), device="cuda")
out = torch.empty(2048, dim, dtype=torch.bfloat16, device="cuda")
_gather_ple_embedding_from_pinned_kernel[(ids.numel(),)](
table.data_ptr(),
ids,
out,
embedding_dim=dim,
tp_vocab_start=0,
tp_vocab_end=rows,
is_fp8=False,
BLOCK_D=triton.next_power_of_2(dim),
)
torch.cuda.synchronize()
expected = table[ids.cpu()].to("cuda")
self.assertTrue(torch.equal(out, expected))
if __name__ == "__main__":
unittest.main()
test/registered/unit/models/test_shared_experts_fusion_gates.py
+14 -1
View File
@@ -656,7 +656,20 @@ class TestWrapperEntryClassGates(_FusionGateCase):
)
# The normalization the constructor applies, shared with the gate.
self.assertIsNone(_mtp_quant_config(_quant("modelopt_mixed")))
mixed_bf16_mtp = SimpleNamespace(
get_name=lambda: "modelopt_mixed",
quantized_layers={"model.layers.0.mlp.experts": {"quant_algo": "NVFP4"}},
)
self.assertIsNone(_mtp_quant_config(mixed_bf16_mtp))
# MIXED_PRECISION checkpoints that quantize the MTP head keep it.
mixed_fp8_mtp = SimpleNamespace(
get_name=lambda: "modelopt_mixed",
quantized_layers={
"model.layers.0.mlp.experts": {"quant_algo": "NVFP4"},
"mtp.layers.0.mlp.experts": {"quant_algo": "FP8_BLOCK_SCALES"},
},
)
self.assertIs(_mtp_quant_config(mixed_fp8_mtp), mixed_fp8_mtp)
serialized = SimpleNamespace(
get_name=lambda: "modelopt_fp4", is_checkpoint_nvfp4_serialized=True
)
test/registered/unit/test_model_overrides.py
+103
View File
@@ -637,6 +637,22 @@ class TestGoldenModelOverrides(_IsolatedPublish):
pp_size=2,
)
def test_qwen4_ple_file_requires_offload(self):
qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
with override_platform(is_cuda=True):
sa = self._construct(
*qwen4,
ple_offload_embedding=True,
ple_offload_backend="file",
ple_offload_dir="/tmp/ple",
)
self.assertEqual(self._resolved(sa, "ple_offload_backend"), "file")
self.assertEqual(self._resolved(sa, "ple_offload_dir"), "/tmp/ple")
with self.assertRaisesRegex(ValueError, "requires --ple-offload-embedding"):
self._construct(
*qwen4, ple_offload_embedding=False, ple_offload_backend="file"
)
def test_qwen4_ple_offload_default(self):
qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
with override_platform(is_cuda=True):
@@ -3072,6 +3088,93 @@ class TestGoldenModelOverrides(_IsolatedPublish):
with override_platform(is_sm100=False):
self.assertEqual(_qwen3_moe_family_overrides(None, None), {})
def test_qwen3_moe_family_mixed_precision_moe_runner(self):
from sglang.srt.arg_groups.model_overrides.qwen3_moe import (
_qwen3_moe_family_overrides,
)
def _mixed(expert_algo):
return SimpleNamespace(
architectures=["Qwen4ExpForConditionalGeneration"],
quantization_config={
"quant_method": "modelopt_mixed",
"quantized_layers": {
"model.language_model.layers.0.mlp.experts": {
"quant_algo": expert_algo
}
},
},
)
args = SimpleNamespace(
quantization="modelopt_mixed",
_quantization_explicitly_unset=False,
moe_a2a_backend="none",
moe_runner_backend="auto",
)
with override_platform(is_sm100=True):
# W4A4 experts take trtllm-gen like modelopt_fp4; W4A16 has no
# trtllm-gen kernel and goes to marlin.
self.assertEqual(
_qwen3_moe_family_overrides(args, _mixed("NVFP4")),
{"moe_runner_backend": "flashinfer_trtllm"},
)
self.assertEqual(
_qwen3_moe_family_overrides(args, _mixed("W4A16_NVFP4")),
{"moe_runner_backend": "marlin"},
)
def test_qwen3_moe_family_w4a16_explicit_runner(self):
"""Keep opted-in CuTe DSL v2 W4A16 accepted and auto routed to Marlin."""
from sglang.srt.arg_groups.model_overrides.qwen3_moe import (
_qwen3_moe_family_overrides,
)
hf_config = SimpleNamespace(
architectures=["Qwen4ExpForConditionalGeneration"],
quantization_config={
"quant_method": "modelopt_mixed",
"quantized_layers": {
"model.language_model.layers.0.mlp.experts": {
"quant_algo": "W4A16_NVFP4"
}
},
},
)
cases = [
("auto", "none", False, {"moe_runner_backend": "marlin"}),
("auto", "none", True, {"moe_runner_backend": "marlin"}),
("marlin", "none", False, {}),
("marlin", "none", True, {}),
("flashinfer_cutedsl", "none", True, {}),
("flashinfer_cutedsl", "flashinfer", True, {}),
("flashinfer_cutedsl", "none", False, None),
("flashinfer_cutedsl", "flashinfer", False, None),
("flashinfer_cutedsl", "deepep", True, None),
("flashinfer_cutlass", "none", True, None),
("flashinfer_trtllm", "none", True, None),
]
for runner, a2a, w4a16_enabled, expected in cases:
with (
self.subTest(runner=runner, a2a=a2a, w4a16=w4a16_enabled),
override_platform(is_sm100=True),
envs.SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16.override(w4a16_enabled),
):
args = SimpleNamespace(
quantization=None,
_quantization_explicitly_unset=False,
moe_a2a_backend=a2a,
moe_runner_backend=runner,
)
if expected is None:
with self.assertRaisesRegex(ValueError, "W4A16_NVFP4"):
_qwen3_moe_family_overrides(args, hf_config)
else:
self.assertEqual(
_qwen3_moe_family_overrides(args, hf_config),
{"quantization": "modelopt_mixed", **expected},
)
def test_step3p_declarations_at_callable_level(self):
from sglang.srt.arg_groups.overrides import _step3p_overrides