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

[AMD][Quantization] Online MXFP4 quantization 4/N - NVFP4 to MXFP4 Online Requantization on AMD GPUs (#29328)

Browse Source
This commit is contained in:
Colin Z
2026-08-14 21:59:39 -07:00
committed by GitHub
parent 5afdb1caea
commit bc7e3ba66c
14 changed files with 1218 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/quant/test_quark_mxfp4.py
+61 -22
View File
@@ -6,6 +6,7 @@ import unittest
from sglang.test.ci.ci_register import register_amd_ci
register_amd_ci(est_time=106, suite="stage-b-test-1-gpu-small-amd-mi35x")
import os
import time
from types import SimpleNamespace
@@ -87,20 +88,10 @@ class TestOnlineQuantizationMemoryLoad(CustomTestCase):
raise RuntimeError(f"Server {url} failed to start in {timeout}s")
time.sleep(1)
# # Extract and display peak GPU memory from logs
combined_output = cls.stdout.getvalue() + cls.stderr.getvalue()
peak_memory_before_load = cls._extract_peak_memory_before_load(combined_output)
if is_cuda_alike() and not peak_memory_before_load:
raise ValueError("Should have found peak memory")
cls.peak_memory_before_load = float(peak_memory_before_load)
memory_increase_load_weights = cls._extract_memory_increase_load_weights(
combined_output
)
if is_cuda_alike() and not memory_increase_load_weights:
raise ValueError("Should have found memory increase in load_weights")
cls.memory_increase_load_weights = float(memory_increase_load_weights)
# Keep the raw server for memory numbers, which are parsed lazily by
# _test_peak_memory so subclasses that don't test memory (e.g. the
# NVFP4->MXFP4 accuracy-only class) don't require these log lines.
cls.combined_output = cls.stdout.getvalue() + cls.stderr.getvalue()
@classmethod
def _extract_peak_memory_before_load(cls, log_output):
@@ -115,8 +106,11 @@ class TestOnlineQuantizationMemoryLoad(CustomTestCase):
@classmethod
def _extract_memory_increase_load_weights(cls, log_output):
"""Extract memory increase during load_weights call."""
# Search for the log message pattern
pattern = r"Memory increase during load_weights:\s+([\d.]+)\s+GiB"
# Signed: the value is (free_before - free_after) around load_weights.
# When the on-device source representation is larger than the loaded
# result (e.g. requantizing to a more compact format), loading frees
# net memory and the reported increase is negative.
pattern = r"Memory increase during load_weights:\s+(-?[\d.]+)\s+GiB"
match = re.search(pattern, log_output)
if match:
return match.group(1)
@@ -138,21 +132,33 @@ class TestOnlineQuantizationMemoryLoad(CustomTestCase):
if not is_cuda_alike():
self.skipTest("not is_cuda_alike")
peak_memory_before_load = self._extract_peak_memory_before_load(
self.combined_output
)
if not peak_memory_before_load:
raise ValueError("Should have found peak memory")
peak_memory_before_load = float(peak_memory_before_load)
memory_increase_load_weights = self._extract_memory_increase_load_weights(
self.combined_output
)
if not memory_increase_load_weights:
raise ValueError("Should have found memory increase in load_weights")
memory_increase_load_weights = float(memory_increase_load_weights)
# NOTE: We can not simply rely on peak memory after `load_weights` as functions used
# in-between (e.g. FP8->MXFP4 requantization) during weight loading may have a higher peak memory footprint
# in-between (e.g. NVFP4->MXFP4 requantization) during weight loading may have a higher peak memory footprint
# than simply the allocated weights.
if add_peak_memory_before_load:
reference_gib = (
self.memory_increase_load_weights + self.peak_memory_before_load
)
reference_gib = memory_increase_load_weights + peak_memory_before_load
else:
reference_gib = self.memory_increase_load_weights
reference_gib = memory_increase_load_weights
assert reference_gib < threshold
if test_start:
# Weights initialized on meta device (not for dense BF16->MXFP4)
assert self.peak_memory_before_load < 5
assert peak_memory_before_load < 5
def _test_gsm8k(self, accuracy_threshold):
"""Helper method to test GSM8K accuracy against a threshold."""
@@ -205,6 +211,39 @@ class TestOnlineQuantizationMemoryLoadMOE(TestOnlineQuantizationMemoryLoad):
self._test_gsm8k(accuracy_threshold=0.89)
class TestNVFP4ToMXFP4MOETP1(TestOnlineQuantizationMemoryLoad):
# ModelOpt NVFP4 export (quant_method="modelopt", quant_algo="NVFP4") =>
# Nvfp4SourceConfig(). Exercises the NVFP4 -> MXFP4 MoE requantization path:
# the per-expert dequantize_nvfp4 + dynamic_mxfp4_quant requant, and the w13
# gate/up weight_scale_2 split in _requantize_nvfp4_to_mxfp4.
model = "nvidia/Qwen3-30B-A3B-NVFP4" # NVFP4 model
tp = 1
def test_gsm8k(self):
# Requantized NVFP4 -> MXFP4 observed accuracy: ~0.88
# (BF16 Qwen/Qwen3-30B-A3B reference: ~0.94).
self._test_gsm8k(accuracy_threshold=0.85)
@unittest.skipIf(is_in_ci(), "local test only")
class TestDeepSeekR10528NVFP4ToMXFP4(TestOnlineQuantizationMemoryLoad):
# NVFP4 to MXFP4 online requantization for DeepSeek-R1-0528-NVFP4 on TP=8.
# Exercises the MLA attention path (attention_backend=aiter), multi-threaded
# weight loading, and the per-expert NVFP4 MoE requantization path.
model = "nvidia/DeepSeek-R1-0528-NVFP4" # NVFP4 model
tp = 8
runner_args = [
"--attention-backend",
"aiter",
"--model-loader-extra-config",
'{"enable_multithread_load": true}',
]
def test_gsm8k(self):
# Requantized NVFP4 -> MXFP4 observed accuracy: ~0.95.
self._test_gsm8k(accuracy_threshold=0.90)
class TestFP8ToMXFP4DenseTP1(TestOnlineQuantizationMemoryLoad):
tp = 1
model = "Qwen/Qwen3-8B-FP8"
test/registered/unit/layers/quantization/test_quark_config.py
+125 -1
View File
@@ -7,7 +7,15 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
from unittest.mock import patch
from sglang.srt.layers.quantization.quark.quark import QuarkConfig
import torch
from sglang.srt.layers.quantization.quark.quark import (
QuarkConfig,
_build_mixed_precision_layer_quant_config,
_mixed_precision_layer_map,
_parse_nvfp4_excludes,
)
from sglang.srt.layers.quantization.quark.utils import check_equal_or_regex_match
from sglang.test.test_utils import CustomTestCase
_GET_CAP = "sglang.srt.layers.quantization.quark.quark.get_device_capability"
@@ -83,5 +91,121 @@ class TestCheckSchemeSupportedError(CustomTestCase):
self.assertFalse(ok)
class TestMixedPrecisionLayerConfig(CustomTestCase):
"""NVFP4-only-experts + FP8-elsewhere online requant (quark_mxfp4).
A MIXED_PRECISION NVFP4 checkpoint (e.g. nvidia/Qwen3.5-397B-A17B-NVFP4-V2)
keeps some layers in NVFP4 while others in FP8. Online requant must send
only the NVFP4 layers through the dequant->MXFP4 path and load the FP8 layers
as FP8.
"""
_LAYER_MAP_SRC = {
"quant_algo": "MIXED_PRECISION",
"quantized_layers": {
"model.language_model.layers.0.self_attn.q_proj": {"quant_algo": "FP8"},
"model.language_model.layers.0.self_attn.k_proj": {"quant_algo": "FP8"},
"model.language_model.layers.0.self_attn.v_proj": {"quant_algo": "FP8"},
"model.language_model.layers.0.self_attn.o_proj": {"quant_algo": "FP8"},
"model.language_model.layers.0.mlp.shared_expert.gate_proj": {
"quant_algo": "FP8"
},
"model.language_model.layers.0.mlp.shared_expert.down_proj": {
"quant_algo": "FP8"
},
"model.language_model.layers.0.mlp.experts": {
"quant_algo": "NVFP4",
"group_size": 16,
},
"model.language_model.layers.1.mlp.experts": {
"quant_algo": "NVFP4",
"group_size": 16,
},
"model.language_model.layers.1.self_attn.q_proj": {"quant_algo": "FP8"},
},
}
def _build_bare_config(self) -> QuarkConfig:
layer_map = _mixed_precision_layer_map(self._LAYER_MAP_SRC)
layer_quant_config, has_nvfp4 = _build_mixed_precision_layer_quant_config(
layer_map
)
self.assertTrue(has_nvfp4)
synth_config = QuarkConfig._create_online_mxfp4_config(
model_type="qwen3_5_moe",
layer_quant_config=layer_quant_config,
)
synth_config["packed_modules_mapping"] = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
}
quark_config = _bare_config()
quark_config.quant_config = synth_config
quark_config.packed_modules_mapping = synth_config["packed_modules_mapping"]
quark_config.exclude_layers = synth_config["exclude"]
return quark_config
def test_experts_route_to_mxfp4_requant(self):
# fnmatch keys (not `re:`) must match the sglang module path so experts
# hit the fp4 target, not fall through to the global config
quark_config = self._build_bare_config()
matched = quark_config._find_matched_config(
"model.layers.0.mlp.experts", torch.nn.Module()
)
self.assertEqual(matched["weight"]["dtype"], "fp4")
self.assertEqual(matched["weight"]["group_size"], 32)
def test_fp8_layers_not_requantized(self):
quark_config = self._build_bare_config()
for name in (
"model.layers.0.self_attn.o_proj",
"model.layers.0.mlp.shared_expert.gate_proj",
"model.layers.0.mlp.shared_expert.down_proj",
):
matched = quark_config._find_matched_config(name, torch.nn.Module())
self.assertEqual(matched["weight"]["dtype"], "fp8_e4m3", msg=name)
self.assertEqual(matched["weight"]["qscheme"], "per_tensor", msg=name)
def test_fused_qkv_shards_share_fp8_scheme(self):
# _find_matched_config expands qkv_proj -> q/k/v shards and requires a
# consistent scheme; all three are FP8 so this must resolve
quark_config = self._build_bare_config()
matched = quark_config._find_matched_config(
"model.layers.0.self_attn.qkv_proj", torch.nn.Module()
)
self.assertEqual(matched["weight"]["dtype"], "fp8_e4m3")
def test_shared_expert_fusion_disabled_on_precision_mismatch(self):
quark_config = self._build_bare_config()
self.assertFalse(quark_config.can_fuse_shared_expert())
def test_mixed_precision_skips_model_type_default_excludes(self):
quark_config = self._build_bare_config()
self.assertNotIn("re:.*shared_expert", quark_config.exclude_layers)
self.assertNotIn("re:.*o_proj", quark_config.exclude_layers)
def test_non_mixed_config_returns_none(self):
self.assertIsNone(_mixed_precision_layer_map({"quant_algo": "NVFP4"}))
class TestParseNvfp4Excludes(CustomTestCase):
"""ModelOpt `ignore` lists mix `re:`-prefixed regexes with fnmatch globs."""
def test_already_regex_entries_pass_through_and_match(self):
# wrapping an already-`re:` entry with another `re:` +
# fnmatch.translate produced `re:(?s:re:\\..*...)` which never matches,
excludes = _parse_nvfp4_excludes(
{"ignore": [r"re:.*linear_attn\.in_proj_a$", "mtp*"]}
)
self.assertTrue(
check_equal_or_regex_match("model.layers.0.linear_attn.in_proj_a", excludes)
)
# fnmatch glob still translated and matches.
self.assertTrue(check_equal_or_regex_match("mtp.layers.0.foo", excludes))
# A quantized layer stays un-excluded.
self.assertFalse(
check_equal_or_regex_match("model.layers.0.mlp.experts", excludes)
)
if __name__ == "__main__":
unittest.main()