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

Fix Nemotron W4A16 NVFP4 MoE backend (#33543)

Browse Source 
Signed-off-by: dafrimi <dafrimi@nvidia.com>
This commit is contained in:
danielafrimi
2026-08-07 12:00:28 +00:00
committed by GitHub
parent 5ca734fc3d
commit 4020bc95a7
3 changed files with 202 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
+23 -1
View File
@@ -1237,7 +1237,29 @@ def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict:
quantization = model_config.quantization
overrides["quantization"] = quantization
if (is_modelopt or model_config.quantization is None) and (
has_w4a16_moe_layers = False
if is_modelopt and quantization == "modelopt_mixed":
has_w4a16_moe_layers = any(
info.get("quant_algo") == "W4A16_NVFP4" and ".experts." in name
for name, info in hf_config.quantization_config.get(
"quantized_layers", {}
).items()
)
if has_w4a16_moe_layers:
if server_args.moe_a2a_backend != "none":
raise ValueError("W4A16_NVFP4 MoE layers require --moe-a2a-backend=none.")
if server_args.moe_runner_backend not in ("auto", "marlin"):
raise ValueError(
"W4A16_NVFP4 MoE layers require --moe-runner-backend=marlin."
)
if server_args.moe_runner_backend == "auto":
overrides["moe_runner_backend"] = "marlin"
logger.info(
"Use marlin as MoE runner backend for "
f"{model_arch} with W4A16_NVFP4 MoE layers"
)
elif (is_modelopt or model_config.quantization is None) and (
server_args.moe_runner_backend == "auto"
):
if is_sm100_supported() and server_args.moe_a2a_backend == "none":
python/sglang/srt/layers/quantization/marlin_utils_fp4.py
+48 -16
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
import logging
import torch
from sglang.srt.layers.quantization.marlin_utils import (
@@ -20,15 +22,14 @@ if _is_cuda:
from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack
ScalarType, scalar_types = get_scalar_types()
logger = logging.getLogger(__name__)
def nvfp4_marlin_process_scales(marlin_scales: torch.Tensor) -> torch.Tensor:
if not (marlin_scales >= 0).all():
# NVFP4 ModelOpt scales are expected to be non-negative. Keep this as
# a warning so unusual checkpoints can still load for diagnosis.
import logging
logging.getLogger(__name__).warning_once(
logger.warning_once(
"NVFP4 Marlin assumes non-negative scales, but negative scales "
"were found. Accuracy may be degraded."
)
@@ -90,11 +91,18 @@ def apply_fp4_marlin_linear(
reshaped_x = input.reshape(-1, input.shape[-1])
out_shape = input.shape[:-1] + (size_n,)
# Recover the physical Marlin tile dimensions from the repacked weight.
# They can exceed the logical TP shard dimensions when preparation padded a
# misaligned shard (e.g. N=928 -> 960 for TP=4).
padded_size_k = weight.size(0) * 16
padded_size_n = weight.size(1) * 8 // 16
if padded_size_k != size_k:
reshaped_x = torch.nn.functional.pad(reshaped_x, (0, padded_size_k - size_k))
use_atomic_add = should_use_atomic_add_reduce(
m=reshaped_x.size(0),
n=size_n,
k=size_k,
n=padded_size_n,
k=padded_size_k,
device=input.device,
dtype=input.dtype,
)
@@ -111,8 +119,8 @@ def apply_fp4_marlin_linear(
workspace=workspace,
b_q_type=scalar_types.float4_e2m1f,
size_m=reshaped_x.size(0),
size_n=size_n,
size_k=size_k,
size_n=padded_size_n,
size_k=padded_size_k,
is_k_full=True,
use_atomic_add=use_atomic_add,
use_fp32_reduce=use_fp32_reduce,
@@ -121,6 +129,10 @@ def apply_fp4_marlin_linear(
if bias is not None:
output.add_(bias)
# A narrowed N dimension has the padded row stride, so materialize it
# before reshaping. This is only needed for a TP shard that was padded.
if padded_size_n != size_n:
output = output[:, :size_n].contiguous()
return output.reshape(out_shape)
@@ -138,10 +150,29 @@ def prepare_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
assert layer.weight.shape == (part_size_n, part_size_k // 2)
if part_size_n % 64 != 0:
raise ValueError(
f"NVFP4 Marlin requires output_size_per_partition to be a multiple of 64, "
f"got {part_size_n}."
# Marlin accepts either N%64/K%128 or N%128/K%64. Select the smaller
# padded shape, matching vLLM's marlin_padded_nk helper.
padded_size_n, padded_size_k = min(
(
((part_size_n + 63) // 64 * 64, (part_size_k + 127) // 128 * 128),
((part_size_n + 127) // 128 * 128, (part_size_k + 63) // 64 * 64),
),
key=lambda nk: (nk[0] * nk[1], nk[0] + nk[1]),
)
if (padded_size_n, padded_size_k) != (part_size_n, part_size_k):
pad_rows = padded_size_n - part_size_n
pad_cols = (padded_size_k - part_size_k) // 2
scale_pad_cols = (padded_size_k - part_size_k) // 16
layer.weight = torch.nn.Parameter(
torch.nn.functional.pad(layer.weight, (0, pad_cols, 0, pad_rows)),
requires_grad=False,
)
layer.weight_scale = torch.nn.Parameter(
torch.nn.functional.pad(
layer.weight_scale, (0, scale_pad_cols, 0, pad_rows)
),
requires_grad=False,
)
device = layer.weight.device
@@ -152,8 +183,8 @@ def prepare_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
marlin_qweight = gptq_marlin_repack(
b_q_weight=qweight,
perm=perm,
size_k=part_size_k,
size_n=part_size_n,
size_k=padded_size_k,
size_n=padded_size_n,
num_bits=4,
)
layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False)
@@ -161,8 +192,8 @@ def prepare_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
weight_scale = layer.weight_scale.T.contiguous().to(param_dtype)
weight_scale = marlin_permute_scales(
s=weight_scale,
size_k=part_size_k,
size_n=part_size_n,
size_k=padded_size_k,
size_n=padded_size_n,
group_size=16,
)
weight_scale = nvfp4_marlin_process_scales(weight_scale)
@@ -176,7 +207,8 @@ def prepare_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
if hasattr(layer, "bias") and layer.bias is not None:
assert layer.bias.shape == (part_size_n,)
bias = marlin_permute_bias(layer.bias)
bias = torch.nn.functional.pad(layer.bias, (0, padded_size_n - part_size_n))
bias = marlin_permute_bias(bias)
layer.bias = torch.nn.Parameter(bias, requires_grad=False)
test/registered/unit/test_model_overrides.py
+131 -3
View File
@@ -405,6 +405,116 @@ class TestGoldenModelOverrides(_IsolatedPublish):
[("_mimo_v2_overrides", {"enable_multi_layer_eagle": True})],
)
def _nemotron_h_args(self, *, quantized_layers):
hf_config = SimpleNamespace(
architectures=["NemotronHForCausalLM"],
mlp_hidden_act="relu2",
quantization_config={
"quant_algo": "MIXED_PRECISION",
"quant_method": "modelopt_mixed",
"quantized_layers": quantized_layers,
},
)
model_config = SimpleNamespace(
quantization="modelopt_mixed", hf_config=hf_config
)
return (
SimpleNamespace(
quantization="modelopt_fp4",
moe_runner_backend="auto",
moe_a2a_backend="none",
attention_backend=None,
get_model_config=lambda: model_config,
),
hf_config,
)
def test_nemotron_h_w4a16_moe_uses_marlin_on_sm100(self):
from sglang.srt.arg_groups.overrides import _nemotron_h_overrides
server_args, hf_config = self._nemotron_h_args(
quantized_layers={
"backbone.layers.1.mixer.experts.0.up_proj": {
"quant_algo": "W4A16_NVFP4",
"group_size": 16,
},
"backbone.layers.1.mixer.experts.0.down_proj": {
"quant_algo": "W4A16_NVFP4",
"group_size": 16,
},
"backbone.layers.0.mixer.in_proj": {"quant_algo": "FP8"},
}
)
with patch.object(overrides_module, "is_sm100_supported", return_value=True):
self.assertEqual(
_nemotron_h_overrides(server_args, hf_config),
{
"quantization": "modelopt_mixed",
"moe_runner_backend": "marlin",
"attention_backend": "flashinfer",
},
)
def test_nemotron_h_nvfp4_moe_keeps_flashinfer_trtllm_on_sm100(self):
from sglang.srt.arg_groups.overrides import _nemotron_h_overrides
server_args, hf_config = self._nemotron_h_args(
quantized_layers={
"backbone.layers.1.mixer.experts.0.up_proj": {
"quant_algo": "NVFP4",
"group_size": 16,
},
"backbone.layers.1.mixer.experts.0.down_proj": {
"quant_algo": "NVFP4",
"group_size": 16,
},
"backbone.layers.0.mixer.in_proj": {"quant_algo": "FP8"},
}
)
with patch.object(overrides_module, "is_sm100_supported", return_value=True):
self.assertEqual(
_nemotron_h_overrides(server_args, hf_config),
{
"quantization": "modelopt_mixed",
"moe_runner_backend": "flashinfer_trtllm",
"attention_backend": "flashinfer",
},
)
def test_nemotron_h_w4a16_moe_rejects_a2a_backend(self):
from sglang.srt.arg_groups.overrides import _nemotron_h_overrides
server_args, hf_config = self._nemotron_h_args(
quantized_layers={
"backbone.layers.1.mixer.experts.0.up_proj": {
"quant_algo": "W4A16_NVFP4",
"group_size": 16,
}
}
)
server_args.moe_a2a_backend = "deepep"
with self.assertRaisesRegex(ValueError, "moe-a2a-backend=none"):
_nemotron_h_overrides(server_args, hf_config)
def test_nemotron_h_w4a16_moe_rejects_non_marlin_runner(self):
from sglang.srt.arg_groups.overrides import _nemotron_h_overrides
server_args, hf_config = self._nemotron_h_args(
quantized_layers={
"backbone.layers.1.mixer.experts.0.up_proj": {
"quant_algo": "W4A16_NVFP4",
"group_size": 16,
}
}
)
server_args.moe_runner_backend = "flashinfer_trtllm"
with self.assertRaisesRegex(ValueError, "moe-runner-backend=marlin"):
_nemotron_h_overrides(server_args, hf_config)
def test_step3p_hierarchical_cache_golden(self):
# SWA-hybrid arch: the mini config needs layer_types/sliding_window.
config_extra = {
@@ -878,12 +988,14 @@ class TestGoldenModelOverrides(_IsolatedPublish):
def test_nemotron_h_overrides_at_callable_level(self):
from sglang.srt.arg_groups.overrides import _nemotron_h_overrides
def _hf(quant_algo="NVFP4"):
return SimpleNamespace(
def _hf(quant_algo="NVFP4", *, include_quantization_config=True):
hf = SimpleNamespace(
architectures=["NemotronHForCausalLM"],
mlp_hidden_act="relu2",
quantization_config={"quant_algo": quant_algo},
)
if include_quantization_config:
hf.quantization_config = {"quant_algo": quant_algo}
return hf
def _args(mc_quant, hf, **kw):
mc = SimpleNamespace(quantization=mc_quant, hf_config=hf)
@@ -940,6 +1052,22 @@ class TestGoldenModelOverrides(_IsolatedPublish):
{},
)
hf_without_quant_cfg = _hf(include_quantization_config=False)
with patch.object(overrides_module, "is_sm100_supported", return_value=True):
for modelopt_quantization in ("modelopt_fp8", "modelopt_fp4"):
with self.subTest(modelopt_quantization=modelopt_quantization):
self.assertEqual(
_nemotron_h_overrides(
_args(modelopt_quantization, hf_without_quant_cfg),
hf_without_quant_cfg,
),
{
"quantization": modelopt_quantization,
"moe_runner_backend": "flashinfer_trtllm",
"attention_backend": "flashinfer",
},
)
def test_speculative_moe_runner_default_pass(self):
from sglang.srt.arg_groups.overrides import (
ResolvedView,