From 6d4b662bab60b5bce993c7c35adfc2f50e978cb8 Mon Sep 17 00:00:00 2001 From: SumanthRH Date: Sat, 20 Sep 2025 00:17:15 +0000 Subject: [PATCH] pass extra weight loader args; Signed-off-by: SumanthRH --- .../compressed_tensors_moe.py | 4 +++- .../layers/quantization/modelopt.py | 5 ++++- .../layers/quantization/mxfp4.py | 9 ++++++++- .../quantization/utils/marlin_utils_fp4.py | 18 +++++++++++++++++- 4 files changed, 32 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 42c43cbc03e5..271ea3b13fd5 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -118,6 +118,7 @@ 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, @@ -125,6 +126,7 @@ def create_weights(self, layer: torch.nn.Module, num_experts: int, 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 diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e0f462b36976..f42ceb63e18e 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/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 diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 3c5d83037cde..bd69b321c6fc 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/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( diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 94ffdcd26ecd..3e09b3298387 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/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)