From 4020bc95a7b5b88b8de5f354f4850a9b1f881298 Mon Sep 17 00:00:00 2001 From: danielafrimi <45691845+danielafrimi@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:00:28 +0300 Subject: [PATCH] Fix Nemotron W4A16 NVFP4 MoE backend (#33543) Signed-off-by: dafrimi --- python/sglang/srt/arg_groups/overrides.py | 24 +++- .../layers/quantization/marlin_utils_fp4.py | 64 ++++++--- test/registered/unit/test_model_overrides.py | 134 +++++++++++++++++- 3 files changed, 202 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 7ef959181..fcd310d81 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -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": diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py index 4ade14f15..ad6ee8332 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py +++ b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py @@ -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) diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index c3241e5c5..32229f3b8 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -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,