[GPT-OSS] Save extra_weight_attrs and use at load_weights time for Marlin kernel

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -118,13 +118,15 @@ def __init__(self, moe: FusedMoEConfig, layer: torch.nn.Module):
         self.use_marlin = _nvfp4.use_marlin
         self.group_size = 16
         self.layer = layer
+        self._extra_weight_attrs = {}
 
     def create_weights(self, layer: torch.nn.Module, num_experts: int,
                        hidden_size: int, intermediate_size_per_partition: int,
                        params_dtype: torch.dtype, **extra_weight_attrs):
 
         layer.num_experts = num_experts
         layer.params_dtype = params_dtype
+        self._extra_weight_attrs = extra_weight_attrs
 
         w13_weight = torch.nn.Parameter(
             torch.empty(
@@ -242,7 +244,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             1 / layer.w2_weight_global_scale.data, requires_grad=False)
 
         if self.use_marlin:
-            prepare_moe_fp4_layer_for_marlin(layer)
+            prepare_moe_fp4_layer_for_marlin(layer, self._extra_weight_attrs)
             return
 
         # swizzle weight scales

vllm/model_executor/layers/quantization/modelopt.py

@@ -969,6 +969,8 @@ def __init__(
         self.fused_experts: Optional[
             mk.FusedMoEModularKernel] = None  # type: ignore[assignment]
 
+        self._extra_weight_attrs = {}
+
     def maybe_make_prepare_finalize(
         self,
         moe: FusedMoEConfig,
@@ -1013,6 +1015,7 @@ def create_weights(self, layer: torch.nn.Module, num_experts: int,
                              " dynamic quantization is not supported.")
 
         layer.num_experts = num_experts
+        self._extra_weight_attrs = extra_weight_attrs
         layer.params_dtype = params_dtype
         layer.quant_config = self.quant_config
         weight_dtype = torch.uint8
@@ -1287,7 +1290,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                                         requires_grad=False)
 
         if self.use_marlin:
-            prepare_moe_fp4_layer_for_marlin(layer)
+            prepare_moe_fp4_layer_for_marlin(layer, self._extra_weight_attrs)
             del layer.g1_alphas
             del layer.g2_alphas
             del layer.w13_input_scale_quant

vllm/model_executor/layers/quantization/mxfp4.py

@@ -87,6 +87,12 @@ def __init__(self, moe: FusedMoEConfig):
         self.moe = moe
         self.use_marlin = self._should_use_marlin()
 
+        # NOTE: with `use_marlin`, we repack weights and create a
+        # new nn.Parameter object in `process_weights_after_loading`.
+        # To preserve any additional attributes passed during `create_weights`,
+        # we store them as a class attribute here.
+        self._extra_weight_attrs = {}
+
     def _should_use_marlin(self):
         if envs.VLLM_MXFP4_USE_MARLIN is not None:
             return envs.VLLM_MXFP4_USE_MARLIN
@@ -105,6 +111,7 @@ def create_weights(self, layer: torch.nn.Module, num_experts: int,
                        hidden_size: int, intermediate_size_per_partition: int,
                        params_dtype: torch.dtype, **extra_weight_attrs):
         self.num_experts = num_experts
+        self._extra_weight_attrs = extra_weight_attrs
         weight_dtype = torch.uint8
         scale_dtype = torch.uint8
 
@@ -229,7 +236,7 @@ def create_weights(self, layer: torch.nn.Module, num_experts: int,
 
     def process_weights_after_loading(self, layer):
         if self.use_marlin:
-            prepare_moe_fp4_layer_for_marlin(layer)
+            prepare_moe_fp4_layer_for_marlin(layer, self._extra_weight_attrs)
         elif (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
               or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
             layer.gemm1_alpha = Parameter(torch.tensor(

vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py

@@ -192,7 +192,9 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
     return
 
 
-def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
+def prepare_moe_fp4_layer_for_marlin(
+        layer: torch.nn.Module,
+        extra_weight_attrs: Optional[dict] = None) -> None:
     logger.warning_once(
         "Your GPU does not have native support for FP4 computation but "
         "FP4 quantization is being used. Weight-only FP4 compression will "
@@ -237,6 +239,10 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
         weight = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
         weight = torch.nn.Parameter(weight, requires_grad=False)
 
+        if extra_weight_attrs:
+            for key, value in extra_weight_attrs.items():
+                setattr(weight, key, value)
+
         setattr(layer, name, weight)
 
     # WEIGHT SCALES
@@ -271,12 +277,19 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
 
         scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
         scales = torch.nn.Parameter(scales, requires_grad=False)
+
+        if extra_weight_attrs:
+            for key, value in extra_weight_attrs.items():
+                setattr(scales, key, value)
         setattr(layer, name + "_weight_scale", scales)
 
         if is_nvfp4:
             global_scale = nvfp4_marlin_process_global_scale(global_scale)
             global_scale = torch.nn.Parameter(global_scale,
                                               requires_grad=False)
+            if extra_weight_attrs:
+                for key, value in extra_weight_attrs.items():
+                    setattr(scales, key, value)
             setattr(layer, name + "_weight_scale_2", global_scale)
 
     # BIAS
@@ -294,6 +307,9 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
 
         bias = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
         bias = torch.nn.Parameter(bias, requires_grad=False)
+        if extra_weight_attrs:
+            for key, value in extra_weight_attrs.items():
+                setattr(bias, key, value)
         setattr(layer, name, bias)