brin in meaningful changes from child PR

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

Author: fxmarty-amd

tests/quantization/test_online_mxfp4.py

@@ -350,7 +350,7 @@ def make_layer(prefix: str) -> RoutedExperts:
             input_config={"dtype": "fp4", "is_dynamic": True},
             moe=checkpoint_layer.moe_config,
         )
-        online_method = Mxfp4OnlineMoEMethod(layer=online_layer)
+        online_method = Mxfp4OnlineMoEMethod(moe=online_layer.moe_config)
 
         # `RoutedExperts.__init__` applies this round-up in production; these
         # layers are built without a quant config, so it is applied explicitly.

vllm/model_executor/layers/quantization/experts_int8.py

@@ -56,5 +56,5 @@ def get_quant_method(
         if isinstance(layer, LinearBase):
             return UnquantizedLinearMethod()
         elif isinstance(layer, RoutedExperts):
-            return Int8OnlineMoEMethod(layer=layer)
+            return Int8OnlineMoEMethod(moe=layer.moe_config)
         return None

vllm/model_executor/layers/quantization/fp8.py

@@ -216,7 +216,7 @@ def get_quant_method(
                     Fp8PerTensorOnlineMoEMethod,
                 )
 
-                return Fp8PerTensorOnlineMoEMethod(layer=layer)
+                return Fp8PerTensorOnlineMoEMethod(moe=layer.moe_config)
         elif isinstance(layer, Attention):
             return Fp8KVCacheMethod(self)
         return None

vllm/model_executor/layers/quantization/online/base.py

@@ -1,13 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from typing import Any
+from typing import Any, cast
 
 import torch
 
 from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
+    FusedMoEMethodBase,
     RoutedExperts,
 )
 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
@@ -32,6 +33,7 @@
     Fp8PerTensorOnlineMoEMethod,
     Fp8PtpcOnlineLinearMethod,
     Fp8PtpcOnlineMoEMethod,
+    OnlineLinearBase,
 )
 from vllm.model_executor.layers.quantization.online.int8 import (
     Int8OnlineMoEMethod,
@@ -223,8 +225,11 @@ def get_quant_method(
             self.quantized_layers[prefix] = (source, str(spec), None)
             cls = target[2]
             if isinstance(layer, RoutedExperts):
-                return cls(layer=layer)
-            return cls()
+                assert issubclass(cls, FusedMoEMethodBase)
+                return cls(moe=layer.moe_config)
+            assert issubclass(cls, OnlineLinearBase)
+            linear_method_cls = cast(type[OnlineLinearBase], cls)
+            return linear_method_cls()
 
         if isinstance(layer, LinearBase):
             return UnquantizedLinearMethod()

vllm/model_executor/layers/quantization/online/fp8.py

@@ -23,6 +23,7 @@
     MarlinFP8ScaledMMLinearKernel,
 )
 from vllm.model_executor.layers.fused_moe import RoutedExperts
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
     select_fp8_moe_backend,
 )
@@ -111,7 +112,7 @@ def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool:
     return is_row_parallel or (reduces_output_dim and is_column_parallel)
 
 
-class _Fp8OnlineLinearBase(LinearMethodBase):
+class OnlineLinearBase(LinearMethodBase):
     """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
     weights onto meta device and materializes them just-in-time."""
 
@@ -155,7 +156,7 @@ def create_weights(
         initialize_online_processing(layer)
 
 
-class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
+class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase):
     """Online tensorwise FP8 linear quantization.
     Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
 
@@ -256,7 +257,7 @@ def apply(
         return self.fp8_linear.apply_weights(layer, x, bias)
 
 
-class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
+class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase):
     """Online blockwise FP8 linear quantization.
     Loads fp16/bf16 weights and quantizes them per-block during loading."""
 
@@ -336,7 +337,7 @@ def apply(
         )
 
 
-class Fp8PtpcOnlineLinearMethod(_Fp8OnlineLinearBase):
+class Fp8PtpcOnlineLinearMethod(OnlineLinearBase):
     """Online PTPC FP8 linear quantization.
 
     Per-output-channel weight scale + dynamic per-token activation scale. The
@@ -443,12 +444,12 @@ def __init__(
         self,
         *,
         weight_block_size: list[int] | None,
-        layer: torch.nn.Module,
+        moe: FusedMoEConfig,
         weight_key: "QuantKey | None" = None,
         activation_key: "QuantKey | None" = None,
         allow_vllm_cutlass: bool = False,
     ):
-        super().__init__(layer.moe_config)
+        super().__init__(moe)
         self.weight_block_size = weight_block_size
         self.block_quant: bool = self.weight_block_size is not None
         self.weight_scale_name = (
@@ -555,11 +556,11 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
     def __init__(
         self,
         *,
-        layer: torch.nn.Module,
+        moe: FusedMoEConfig,
     ):
         super().__init__(
             weight_block_size=None,
-            layer=layer,
+            moe=moe,
         )
 
     def process_weights_after_loading(self, layer: Module) -> None:
@@ -612,11 +613,11 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
     def __init__(
         self,
         *,
-        layer: torch.nn.Module,
+        moe: FusedMoEConfig,
     ):
         super().__init__(
             weight_block_size=[128, 128],
-            layer=layer,
+            moe=moe,
         )
 
     def maybe_roundup_sizes(
@@ -716,13 +717,13 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
     def __init__(
         self,
         *,
-        layer: torch.nn.Module,
+        moe: FusedMoEConfig,
     ):
         from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
 
         super().__init__(
             weight_block_size=None,
-            layer=layer,
+            moe=moe,
             weight_key=kFp8StaticChannelSym,
             activation_key=kFp8DynamicTokenSym,
             allow_vllm_cutlass=True,

vllm/model_executor/layers/quantization/online/int8.py

@@ -12,6 +12,7 @@
     )
 
 from vllm.model_executor.layers.fused_moe import RoutedExperts
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 from vllm.model_executor.layers.fused_moe.oracle.int8 import (
     convert_to_int8_moe_kernel_format,
     make_int8_moe_kernel,
@@ -38,9 +39,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
     def __init__(
         self,
         *,
-        layer: torch.nn.Module,
+        moe: FusedMoEConfig,
     ):
-        super().__init__(layer.moe_config)
+        super().__init__(moe)
         self.int8_backend, self.experts_cls = select_int8_moe_backend(
             config=self.moe,
             weight_key=kInt8StaticChannelSym,

vllm/model_executor/layers/quantization/online/mxfp4.py

@@ -17,6 +17,7 @@
     )
 
 from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
     TRITON_BACKENDS,
     Mxfp4MoeBackend,
@@ -27,7 +28,7 @@
     select_mxfp4_moe_backend,
 )
 from vllm.model_executor.layers.quantization.online.fp8 import (
-    _Fp8OnlineLinearBase,
+    OnlineLinearBase,
 )
 from vllm.model_executor.layers.quantization.online.moe_base import (
     OnlineMoEMethodBase,
@@ -73,7 +74,7 @@ def _quantize_mxfp4_moe_weight(
     return w_quant, w_scales
 
 
-class Mxfp4OnlineLinearMethod(_Fp8OnlineLinearBase):
+class Mxfp4OnlineLinearMethod(OnlineLinearBase):
     """Online MXFP4 linear method.
     Loads bf16/fp16 checkpoints and quantizes weights to MXFP4 (microscaling
     FP4 with block-32 scales) during weight loading.
@@ -144,8 +145,8 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):
     experts_cls: "type[mk.FusedMoEExperts] | None"
     activation_quant_key = kMxfp4Dynamic
 
-    def __init__(self, *, layer: torch.nn.Module):
-        super().__init__(layer.moe_config)
+    def __init__(self, *, moe: FusedMoEConfig):
+        super().__init__(moe)
         self.weight_block_size: list[int] = [1, MXFP4_BLOCK_SIZE]
         self.weight_scale_name = "weight_scale"
 

vllm/model_executor/layers/quantization/online/mxfp8.py

@@ -8,6 +8,8 @@
 import torch
 from torch.nn import Module
 
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
+
 if TYPE_CHECKING:
     import vllm.model_executor.layers.fused_moe.modular_kernel as mk
     from vllm.model_executor.layers.fused_moe import (
@@ -21,7 +23,7 @@
     select_mxfp8_moe_backend,
 )
 from vllm.model_executor.layers.quantization.online.fp8 import (
-    _Fp8OnlineLinearBase,
+    OnlineLinearBase,
 )
 from vllm.model_executor.layers.quantization.online.moe_base import (
     OnlineMoEMethodBase,
@@ -34,7 +36,7 @@
 from vllm.platforms import current_platform
 
 
-class Mxfp8OnlineLinearMethod(_Fp8OnlineLinearBase):
+class Mxfp8OnlineLinearMethod(OnlineLinearBase):
     """Online MXFP8 linear method.
     Loads bf16/fp16 checkpoints and quantizes weights to MXFP8 (microscaling
     FP8 with block-32 scales) during weight loading.
@@ -100,8 +102,8 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
     fp8_backend: "Fp8MoeBackend"
     experts_cls: "type[mk.FusedMoEExperts] | None"
 
-    def __init__(self, *, layer: torch.nn.Module):
-        super().__init__(layer.moe_config)
+    def __init__(self, *, moe: FusedMoEConfig):
+        super().__init__(moe)
         self.weight_block_size: list[int] = [1, MXFP8_BLOCK_SIZE]
         self.weight_scale_name = "weight_scale"
 

vllm/model_executor/layers/quantization/online/nvfp4.py

@@ -6,7 +6,10 @@
 
 from vllm._custom_ops import scaled_fp4_quant
 from vllm.model_executor.layers.fused_moe import RoutedExperts
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+)
 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
     convert_to_nvfp4_moe_kernel_format,
     make_nvfp4_moe_kernel,
@@ -87,13 +90,13 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
     def __init__(
         self,
         *,
-        layer: torch.nn.Module,
+        moe: FusedMoEConfig,
     ):
         if not current_platform.is_device_capability_family(100):
             raise ValueError(
                 "nvfp4_per_token online quantization requires a Blackwell (SM100) GPU."
             )
-        super().__init__(layer.moe_config)
+        super().__init__(moe)
         self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
             config=self.moe,
             weight_key=kNvfp4Static,