fix: port ModelOpt W4A16 NVFP4 rollout to vllm 0.25

- vLLM 0.25's ModelOptNvFp4Config.__init__ installs LinearMethodCls as an
  instance attribute keyed off the quant algo, shadowing the NeMo subclass
  class attribute: the W4A16 config silently instantiated the native W4A4
  linear method, whose process_weights_after_loading reads an input_scale a
  W4A16 checkpoint never loads. from_config now rebinds the instance
  attribute to the NeMo Marlin weight-only method.
- W4A16_NVFP4 is a natively understood algo in 0.25, so the from_config
  normalization to NVFP4 is gone (validate-only); the base FusedMoE __init__
  keys weight-only mode off the algo (activation_key=None), replacing the
  0.20-era duplicated __init__ in NemoModelOptW4A16FusedMoE.
- vLLM 0.25's prepare_nvfp4_moe_layer_for_marlin pads rank-local
  intermediate tiles itself and asserts on unpadded checkpoint shapes, so
  the NeMo-side Marlin pre-padding (_pad_nvfp4_moe_for_marlin) would
  double-pad and trip that assertion; it is removed. Only the E4M3 sign-bit
  canonicalization of ModelOpt scale exports remains NeMo's concern.

Fake-vllm test harness updated to model the 0.25 instance-attribute and
use_a16 behavior so the shadowing bug is covered by a regression test.

Signed-off-by: Terry Kong <terryk@nvidia.com>

Author: terrykong

nemo_rl/modelopt/models/generation/vllm_modelopt.py

@@ -23,7 +23,6 @@
 method-owned MoE kernel references across layerwise reload.
 """
 
-import copy
 from types import MethodType
 from typing import Any
 
@@ -88,17 +87,14 @@ def _load_modelopt_moe_input_scale(
     return True if return_success else None
 
 
-def _normalized_w4a16_config(config: dict[str, Any]) -> dict[str, Any]:
-    normalized = copy.deepcopy(config)
-    quantization = normalized.get("quantization")
-    target = quantization if isinstance(quantization, dict) else normalized
+def _validated_w4a16_config(config: dict[str, Any]) -> dict[str, Any]:
+    quantization = config.get("quantization")
+    target = quantization if isinstance(quantization, dict) else config
     if str(target.get("quant_algo", "")).upper() != _W4A16_ALGO:
         raise ValueError(f"{NEMO_MODELOPT_W4A16} requires quant_algo={_W4A16_ALGO!r}")
-    # vLLM 0.20 validates known ModelOpt algorithms before dispatching to a
-    # custom subclass. Normalize only for its parser; class identity selects
-    # the W4A16 methods below.
-    target["quant_algo"] = _W4A4_ALGO
-    return normalized
+    # vLLM 0.25 understands W4A16_NVFP4 natively, so the algo passes through
+    # unchanged (the base __init__ keys use_a16/LinearMethodCls off it).
+    return config
 
 
 def _canonicalize_nvfp4_scale_(scale: torch.Tensor) -> None:
@@ -107,54 +103,6 @@ def _canonicalize_nvfp4_scale_(scale: torch.Tensor) -> None:
         scale.copy_(scale.to(torch.float32).abs().to(scale.dtype))
 
 
-def _pad_nvfp4_moe_for_marlin(
-    w13: torch.Tensor,
-    w13_scale: torch.Tensor,
-    w2: torch.Tensor,
-    w2_scale: torch.Tensor,
-    *,
-    is_act_and_mul: bool,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
-    """Apply rank-local post-load padding required by the Marlin MoE kernel."""
-    num_experts = w13.shape[0]
-    num_shards = 2 if is_act_and_mul else 1
-    intermediate_size = w13.shape[1] // num_shards
-    hidden_size = w13.shape[2] * 2
-    if hidden_size % 128 == 0:
-        tile_size = 64
-    elif hidden_size % 64 == 0:
-        tile_size = 128
-    else:
-        raise ValueError(
-            f"W4A16 Marlin MoE requires hidden_size divisible by 64, got {hidden_size}"
-        )
-    padded_size = (intermediate_size + tile_size - 1) // tile_size * tile_size
-    if padded_size == intermediate_size:
-        return w13, w13_scale, w2, w2_scale, intermediate_size
-
-    def pad_w13(tensor: torch.Tensor) -> torch.Tensor:
-        tensor = tensor.view(
-            num_experts,
-            num_shards,
-            intermediate_size,
-            tensor.shape[-1],
-        )
-        tensor = torch.nn.functional.pad(
-            tensor,
-            (0, 0, 0, padded_size - intermediate_size),
-        )
-        return tensor.reshape(num_experts, num_shards * padded_size, -1)
-
-    w13 = pad_w13(w13)
-    w13_scale = pad_w13(w13_scale)
-    w2 = torch.nn.functional.pad(w2, (0, (padded_size - intermediate_size) // 2))
-    w2_scale = torch.nn.functional.pad(
-        w2_scale,
-        (0, (padded_size - intermediate_size) // 16),
-    )
-    return w13, w13_scale, w2, w2_scale, padded_size
-
-
 def register_nemo_modelopt_nvfp4() -> None:
     """Register NeMo's two ModelOpt NVFP4 configs through vLLM's public API."""
     global _registered
@@ -165,14 +113,7 @@ def register_nemo_modelopt_nvfp4() -> None:
         MarlinNvFp4LinearKernel,
         NvFp4LinearLayerConfig,
     )
-    from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
-        FusedMoEMethodBase,
-    )
-    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
-        NvFp4MoeBackend,
-        is_global_sf_supported_for_nvfp4_backend,
-        select_nvfp4_moe_backend,
-    )
+    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
     from vllm.model_executor.layers.linear import (
         register_weight_loader_v2_supported_method,
     )
@@ -182,8 +123,6 @@ def register_nemo_modelopt_nvfp4() -> None:
         ModelOptNvFp4FusedMoE,
         ModelOptNvFp4LinearMethod,
     )
-    from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static
-    from vllm.model_executor.utils import replace_parameter
 
     class NemoModelOptNvFp4FusedMoE(ModelOptNvFp4FusedMoE):
         """Native W4A4 MoE plus the vLLM 0.20 input-scale loader fix."""
@@ -297,30 +236,17 @@ def apply(
             return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
 
     class NemoModelOptW4A16FusedMoE(ModelOptNvFp4FusedMoE):
-        """ModelOpt W4A16 MoE using vLLM's NVFP4 Marlin implementation."""
+        """ModelOpt W4A16 MoE using vLLM's NVFP4 Marlin implementation.
+
+        vLLM 0.25's base __init__ keys weight-only mode off
+        quant_config.quant_method == "W4A16_NVFP4" (activation_key=None), so
+        the 0.20-era duplicated __init__ is gone; the algo passes through
+        from_config unchanged.
+        """
 
         moe_kernel: Any
         moe_quant_config: Any
 
-        def __init__(self, quant_config: object, moe_config: object) -> None:
-            # Duplicates vLLM v0.20.0 ModelOptNvFp4FusedMoE.__init__ except
-            # activation_key=None (weight-only); the base hard-wires
-            # kNvfp4Dynamic and offers no hook:
-            # https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/layers/quantization/modelopt.py#L1218-L1234
-            # Intentionally calls FusedMoEMethodBase.__init__ to skip the
-            # parent __init__; do not replace it with super().__init__().
-            # Re-sync on vLLM bumps.
-            FusedMoEMethodBase.__init__(self, moe_config)
-            self.quant_config = quant_config
-            self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
-                config=self.moe,
-                weight_key=kNvfp4Static,
-                activation_key=None,
-            )
-            self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
-                self.nvfp4_backend
-            )
-
         def create_weights(
             self,
             layer: Any,
@@ -344,35 +270,20 @@ def create_weights(
         def process_weights_after_loading(self, layer: Any) -> None:
             reload_kernel = self.moe_kernel
             reload_quant_config = self.moe_quant_config
-            original_intermediate_size = (
-                layer.moe_config.intermediate_size_per_partition
-            )
             if self.nvfp4_backend == NvFp4MoeBackend.MARLIN:
-                w13, w13_scale, w2, w2_scale, padded_size = _pad_nvfp4_moe_for_marlin(
-                    layer.w13_weight,
-                    layer.w13_weight_scale,
-                    layer.w2_weight,
-                    layer.w2_weight_scale,
-                    is_act_and_mul=self.moe.is_act_and_mul,
-                )
-                replace_parameter(layer, "w13_weight", w13)
-                replace_parameter(layer, "w13_weight_scale", w13_scale)
-                replace_parameter(layer, "w2_weight", w2)
-                replace_parameter(layer, "w2_weight_scale", w2_scale)
+                # vLLM 0.25's prepare_nvfp4_moe_layer_for_marlin pads the
+                # rank-local intermediate tiles itself (and asserts on the
+                # unpadded checkpoint shapes), so no NeMo-side pre-padding.
+                # Only the E4M3 sign-bit canonicalization of the ModelOpt
+                # export remains our concern.
                 _canonicalize_nvfp4_scale_(layer.w13_weight_scale)
                 _canonicalize_nvfp4_scale_(layer.w2_weight_scale)
-                layer.moe_config.intermediate_size_per_partition = padded_size
             # W4A16 checkpoint metadata deliberately omits activation scales so
             # layerwise reload never waits for tensors that do not exist. The
             # native Marlin converter accepts None and removes these attributes.
             layer.w13_input_scale = None
             layer.w2_input_scale = None
-            try:
-                super().process_weights_after_loading(layer)
-            finally:
-                layer.moe_config.intermediate_size_per_partition = (
-                    original_intermediate_size
-                )
+            super().process_weights_after_loading(layer)
             if reload_kernel is not None:
                 self.moe_kernel = reload_kernel
                 self.moe_quant_config = reload_quant_config
@@ -401,7 +312,13 @@ def override_quantization_method(
 
         @classmethod
         def from_config(cls, config: dict[str, Any]) -> Any:
-            return super().from_config(_normalized_w4a16_config(config))
+            instance = super().from_config(_validated_w4a16_config(config))
+            # vLLM 0.25's ModelOptNvFp4Config.__init__ selects LinearMethodCls
+            # from the quant algo as an *instance* attribute, which shadows
+            # the class attribute above; rebind the NeMo method explicitly so
+            # W4A16 linears keep the refit-friendly Marlin implementation.
+            instance.LinearMethodCls = NemoModelOptW4A16LinearMethod
+            return instance
 
     register_quantization_config(NEMO_MODELOPT_W4A4)(NemoModelOptNvFp4Config)
     register_quantization_config(NEMO_MODELOPT_W4A16)(NemoModelOptW4A16Config)

tests/unit/models/generation/test_vllm_modelopt_real_quant_config.py

@@ -27,7 +27,6 @@
 from nemo_rl.modelopt.models.generation.vllm_modelopt import (
     NEMO_MODELOPT_W4A4,
     NEMO_MODELOPT_W4A16,
-    _pad_nvfp4_moe_for_marlin,
     quantization_method_for_mode,
     register_nemo_modelopt_nvfp4,
 )
@@ -334,6 +333,13 @@ def __init__(self, quant_config, moe_config):
             self.moe = moe_config
             self.moe_kernel = None
             self.moe_quant_config = None
+            # Mirror vLLM 0.25: weight-only mode and backend selection are
+            # keyed off the config's quant_method in the base __init__.
+            self.use_a16 = (
+                getattr(quant_config, "quant_method", "NVFP4") == "W4A16_NVFP4"
+            )
+            self.nvfp4_backend = "marlin"
+            self.use_global_sf = False
 
         def create_weights(self, layer, *args, **kwargs):
             events.append(("native_create_weights", layer, args, kwargs))
@@ -372,17 +378,35 @@ def process_weights_after_loading(self, layer):
             else:
                 self.moe_quant_config = types.SimpleNamespace(source="native")
 
+    class FakeModelOptNvFp4W4A16LinearMethod(FakeModelOptNvFp4LinearMethod):
+        pass
+
     class FakeModelOptNvFp4Config:
         LinearMethodCls = FakeModelOptNvFp4LinearMethod
         FusedMoEMethodCls = FakeModelOptNvFp4FusedMoE
 
-        def __init__(self, group_size=16):
+        def __init__(self, quant_method="NVFP4", group_size=16):
+            self.quant_method = quant_method
             self.group_size = group_size
+            # Mirror vLLM 0.25: __init__ installs LinearMethodCls as an
+            # *instance* attribute keyed off quant_method, shadowing any
+            # subclass class attribute.
+            if quant_method == "NVFP4":
+                self.LinearMethodCls = FakeModelOptNvFp4LinearMethod
+            elif quant_method == "W4A16_NVFP4":
+                self.LinearMethodCls = FakeModelOptNvFp4W4A16LinearMethod
+            else:
+                raise ValueError(
+                    f"Unsupported ModelOpt NVFP4 quant_algo: {quant_method}"
+                )
 
         @classmethod
         def from_config(cls, config):
             target = config.get("quantization", config)
-            instance = cls(group_size=target.get("group_size", 16))
+            instance = cls(
+                quant_method=str(target.get("quant_algo", "NVFP4")).upper(),
+                group_size=target.get("group_size", 16),
+            )
             instance.parsed_config = config
             return instance
 
@@ -2325,8 +2349,12 @@ def test_register_nemo_modelopt_nvfp4_uses_public_vllm_registry(monkeypatch):
     source_config = {"quant_algo": "W4A16_NVFP4", "group_size": 16}
     w4a16_config = fake_vllm.registry[NEMO_MODELOPT_W4A16].from_config(source_config)
     assert source_config["quant_algo"] == "W4A16_NVFP4"
-    assert w4a16_config.parsed_config["quant_algo"] == "NVFP4"
+    # vLLM 0.25 understands W4A16_NVFP4 natively; the algo passes through.
+    assert w4a16_config.parsed_config["quant_algo"] == "W4A16_NVFP4"
     assert w4a16_config.get_name() == NEMO_MODELOPT_W4A16
+    # The base __init__ installs its own LinearMethodCls instance attribute;
+    # the NeMo config must rebind it to the refit-friendly Marlin method.
+    assert w4a16_config.LinearMethodCls.__name__ == "NemoModelOptW4A16LinearMethod"
 
     with pytest.raises(ValueError, match="requires quant_algo='W4A16_NVFP4'"):
         fake_vllm.registry[NEMO_MODELOPT_W4A16].from_config({"quant_algo": "NVFP4"})
@@ -2562,63 +2590,6 @@ def test_registered_w4a16_dense_method_uses_marlin_weight_only(monkeypatch):
     assert kernel_args["bias"] is None
 
 
-@pytest.mark.parametrize(
-    ("is_act_and_mul", "packed_hidden_size", "expected_padded_size"),
-    [
-        (False, 64, 192),
-        (True, 32, 256),
-    ],
-)
-def test_pad_nvfp4_moe_for_marlin_uses_hidden_size_tile_alignment(
-    is_act_and_mul,
-    packed_hidden_size,
-    expected_padded_size,
-):
-    num_shards = 2 if is_act_and_mul else 1
-    intermediate_size = 144
-    w13 = torch.ones(
-        1,
-        num_shards * intermediate_size,
-        packed_hidden_size,
-    )
-    w13_scale = torch.ones(1, num_shards * intermediate_size, 2)
-    w2 = torch.ones(1, 2, intermediate_size // 2)
-    w2_scale = torch.ones(1, 2, intermediate_size // 16)
-
-    padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_size = (
-        _pad_nvfp4_moe_for_marlin(
-            w13,
-            w13_scale,
-            w2,
-            w2_scale,
-            is_act_and_mul=is_act_and_mul,
-        )
-    )
-
-    assert padded_size == expected_padded_size
-    assert padded_w13.shape == (
-        1,
-        num_shards * expected_padded_size,
-        packed_hidden_size,
-    )
-    assert padded_w13_scale.shape == (1, num_shards * expected_padded_size, 2)
-    assert padded_w2.shape == (1, 2, expected_padded_size // 2)
-    assert padded_w2_scale.shape == (1, 2, expected_padded_size // 16)
-
-    padded_w13 = padded_w13.view(
-        1, num_shards, expected_padded_size, packed_hidden_size
-    )
-    padded_w13_scale = padded_w13_scale.view(1, num_shards, expected_padded_size, 2)
-    assert torch.all(padded_w13[:, :, :intermediate_size] == 1)
-    assert torch.count_nonzero(padded_w13[:, :, intermediate_size:]) == 0
-    assert torch.all(padded_w13_scale[:, :, :intermediate_size] == 1)
-    assert torch.count_nonzero(padded_w13_scale[:, :, intermediate_size:]) == 0
-    assert torch.all(padded_w2[..., : intermediate_size // 2] == 1)
-    assert torch.count_nonzero(padded_w2[..., intermediate_size // 2 :]) == 0
-    assert torch.all(padded_w2_scale[..., : intermediate_size // 16] == 1)
-    assert torch.count_nonzero(padded_w2_scale[..., intermediate_size // 16 :]) == 0
-
-
 def test_registered_w4a16_moe_create_weights_keeps_checkpoint_layout(monkeypatch):
     fake_vllm = _install_fake_registered_vllm_modelopt(monkeypatch)
     monkeypatch.setattr(vllm_modelopt, "_registered", False)
@@ -2683,11 +2654,14 @@ def test_registered_w4a16_moe_preserves_kernel_during_reload(monkeypatch):
 
     assert quant_method.moe_kernel is original_kernel
     assert quant_method.moe_quant_config is original_quant_config
+    # vLLM 0.25's native Marlin converter owns tile padding, so the NeMo
+    # override leaves shapes and moe_config untouched and only canonicalizes
+    # the ModelOpt sign-carrying scales in place.
     assert layer.moe_config.intermediate_size_per_partition == 80
-    assert layer.w13_weight.shape == (1, 128, 32)
-    assert layer.w13_weight_scale.shape == (1, 128, 2)
-    assert layer.w2_weight.shape == (1, 2, 64)
-    assert layer.w2_weight_scale.shape == (1, 2, 8)
+    assert layer.w13_weight.shape == (1, 80, 32)
+    assert layer.w13_weight_scale.shape == (1, 80, 2)
+    assert layer.w2_weight.shape == (1, 2, 40)
+    assert layer.w2_weight_scale.shape == (1, 2, 5)
     assert torch.all(layer.w13_weight_scale >= 0)
     assert torch.all(layer.w2_weight_scale >= 0)
-    assert fake_vllm.events == [("native_process_moe", 128)]
+    assert fake_vllm.events == [("native_process_moe", 80)]