iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Online quantization] Extend re-quantization API to MOE (showcase NVFP4 modelopt/CT/quark to any online MOE method) (#60276)

Browse Source 
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
Signed-off-by: Tan Pin Siang <tanpinsiang@gmail.com>
Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com>
Co-authored-by: Tan Pin Siang <tanpinsiang@gmail.com>
This commit is contained in:
fxmarty-amd
2026-10-14 15:49:37 +08:00
committed by GitHub
co-authored by vllmellm Tan Pin Siang
parent e1726d9142
commit c9bfa574ce
16 changed files with 1353 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/features/quantization/online.md
+14 -3
View File
@@ -111,7 +111,10 @@ Currently supported combinations are:
| Checkpoint layer | Online target | Requirements |
| ---------------- | ------------- | ------------ |
| Serialized ModelOpt MXFP8 linear | any | Target-specific hardware requirements |
| Serialized ModelOpt MXFP8 linear | any online linear method | Target-specific hardware requirements |
| Serialized compressed-tensors NVFP4 MoE | any online MoE method | Target-specific hardware requirements |
| Serialized ModelOpt NVFP4 MoE | any online MoE method | Target-specific hardware requirements |
| Serialized Quark NVFP4 MoE | any online MoE method | Target-specific hardware requirements |
For example, serialized ModelOpt MXFP8 linears can be converted to
per-channel-weight/per-token-activation FP8 at load time:
@@ -122,10 +125,18 @@ vllm serve MiniMaxAI/MiniMax-M3-MXFP8 \
--quantization-config.linear fp8_per_channel
```
Similarly, for example, NVFP4 MoE checkpoints can be re-quantized to any online MoE method,
for example compressed-tensors NVFP4 experts to online MXFP4:
```bash
vllm serve RedHatAI/Qwen3-30B-A3B-NVFP4 \
--quantization-config '{"targets":{"*mlp.experts*":"mxfp4"}}'
```
!!! info
This is only available for linear layers at the moment.
Re-quantization is available for both linear and MoE layers.
Checkpoint-exclusion policy is not enforced by online quantization/re-quantization. Only exclusion through `--quantization-config.targets`/`--quantization-config.ignore` is enforced.
### Online quantization on unquantized layers from partially-quantized checkpoints
tests/quantization/test_online.py
+925 -23
View File
File diff suppressed because it is too large Load Diff
vllm/model_executor/layers/fused_moe/routed_experts.py
+21 -4
View File
@@ -700,12 +700,29 @@ class RoutedExperts(PluggableLayer):
param.data[:, :dim1, :dim2].copy_(loaded_weight)
return True if return_success else None
quant_method_name = self.quant_method.__class__.__name__
from vllm.model_executor.layers.quantization.online.moe_base import (
OnlineMoEMethodBase,
)
base_quant_method = self.quant_method
# In case online MOE re-quantization is used, loading is handled
# by the original quant method.
if (
isinstance(base_quant_method, OnlineMoEMethodBase)
and base_quant_method.requantization_source is not None
):
assert isinstance(
base_quant_method.requantization_source, FusedMoEMethodBase
)
base_quant_method = base_quant_method.requantization_source
quant_method_name = base_quant_method.__class__.__name__
global_expert_id = expert_id
expert_id = self._map_global_expert_id_to_local_expert_id(global_expert_id)
use_global_sf = (
getattr(self.quant_method, "use_global_sf", False)
getattr(base_quant_method, "use_global_sf", False)
and "input_scale" in weight_name
)
@@ -842,7 +859,7 @@ class RoutedExperts(PluggableLayer):
if "ModelOpt" in quant_method_name:
# Determine per-tensor weight scale patterns based on variant
# Use the dedicated method instead of brittle string matching
uses_weight_scale_2 = self.quant_method.uses_weight_scale_2_pattern()
uses_weight_scale_2 = base_quant_method.uses_weight_scale_2_pattern()
quant_method = getattr(param, "quant_method", None)
# Call _load_per_tensor_weight_scale() to load per-tensor (scalar)
@@ -921,7 +938,7 @@ class RoutedExperts(PluggableLayer):
FusedMoeWeightScaleSupported.GROUP.value,
FusedMoeWeightScaleSupported.BLOCK.value,
]:
scale_refine = getattr(self.quant_method, "weight_scale_refine", None)
scale_refine = getattr(base_quant_method, "weight_scale_refine", None)
if (
quant_method == FusedMoeWeightScaleSupported.BLOCK.value
and scale_refine is not None
vllm/model_executor/layers/linear.py
+30 -10
View File
@@ -516,6 +516,20 @@ class ColumnParallelLinear(LinearBase):
self._maybe_allow_fp8_block_shape_mismatch()
self.gather_output = gather_output
from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase
quant_method = self.quant_method
if (
isinstance(quant_method, OnlineLinearBase)
and quant_method.requantization_source is not None
):
quant_method = quant_method.requantization_source
weight_loader = (
self.weight_loader_v2
if quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
else self.weight_loader
)
self.quant_method.create_weights(
layer=self,
input_size_per_partition=self.input_size_per_partition,
@@ -523,11 +537,7 @@ class ColumnParallelLinear(LinearBase):
input_size=self.input_size,
output_size=self.output_size,
params_dtype=self.params_dtype,
weight_loader=(
self.weight_loader_v2
if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
else self.weight_loader
),
weight_loader=weight_loader,
)
if bias:
@@ -1697,6 +1707,20 @@ class RowParallelLinear(LinearBase):
self.input_is_parallel = input_is_parallel
self.reduce_results = reduce_results
from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase
quant_method = self.quant_method
if (
isinstance(quant_method, OnlineLinearBase)
and quant_method.requantization_source is not None
):
quant_method = quant_method.requantization_source
weight_loader = (
self.weight_loader_v2
if quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
else self.weight_loader
)
self.quant_method.create_weights(
layer=self,
input_size_per_partition=self.input_size_per_partition,
@@ -1704,11 +1728,7 @@ class RowParallelLinear(LinearBase):
input_size=self.input_size,
output_size=self.output_size,
params_dtype=self.params_dtype,
weight_loader=(
self.weight_loader_v2
if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
else self.weight_loader
),
weight_loader=weight_loader,
)
if not reduce_results and (bias and not skip_bias_add):
raise ValueError(
vllm/model_executor/layers/quantization/base_config.py
+9 -8
View File
@@ -84,11 +84,14 @@ class QuantizeMethodBase(ABC):
"""
return
def dequantize_weight(self, layer: nn.Module) -> torch.Tensor:
def dequantize_weight(
self, layer: nn.Module
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Materialize a serialized quantized weight for requantization."""
raise NotImplementedError(
f"The quantization method {type(self)} does not implement "
"dequantize_weight. Please open an issue."
"dequantize_weight, so re-quantization from this checkpoint is "
"not supported. Please open an issue."
)
@@ -302,6 +305,9 @@ def resolve_quant_method(
UnquantizedLinearMethod,
)
from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase
from vllm.model_executor.layers.quantization.online.moe_base import (
OnlineMoEMethodBase,
)
base_quant_method = quant_config.get_quant_method(layer, prefix)
if quant_config.online_quantization_config is None:
@@ -328,11 +334,6 @@ def resolve_quant_method(
# The checkpoint quant method is applied as there is no online override.
return base_quant_method
if isinstance(layer, RoutedExperts):
raise NotImplementedError(
"Requantizing checkpoint-quantized MoE layers is not supported."
)
online_quant_method = quant_config.online_quantization_config.get_quant_method(
layer, prefix
)
@@ -341,7 +342,7 @@ def resolve_quant_method(
base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod)
)
assert isinstance(online_quant_method, OnlineLinearBase)
assert isinstance(online_quant_method, (OnlineLinearBase, OnlineMoEMethodBase))
online_quant_method.set_requantization_source(base_quant_method)
# The online method dequantizes the checkpoint method before requantizing.
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+19
View File
@@ -26,6 +26,9 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
CompressedTensorsMoEMethod,
)
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
dequantize_nvfp4_moe_weights,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
kNvfp4Dynamic,
kNvfp4Static,
@@ -185,6 +188,22 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
)
set_weight_attrs(w2_input_scale, extra_weight_attrs)
def dequantize_weight(
self, layer: torch.nn.Module
) -> tuple[torch.Tensor, torch.Tensor]:
"""Materialize serialized NVFP4 expert weights for requantization."""
return dequantize_nvfp4_moe_weights(
layer.w13_weight_packed,
layer.w13_weight_scale,
layer.w13_weight_global_scale,
layer.w2_weight_packed,
layer.w2_weight_scale,
layer.w2_weight_global_scale,
layer.params_dtype,
group_size=self.group_size,
invert_global_scales=True,
)
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
"""Convert NVFP4 MoE weights into kernel format and setup the kernel."""
# NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod
vllm/model_executor/layers/quantization/modelopt.py
+19
View File
@@ -81,6 +81,9 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
MXFP8_VALUE_DTYPE,
dequant_mxfp8_to_bf16,
)
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
dequantize_nvfp4_moe_weights,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
FP4_DTYPE,
QuantKey,
@@ -889,6 +892,22 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
"""FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales."""
return True
def dequantize_weight(
self, layer: torch.nn.Module
) -> tuple[torch.Tensor, torch.Tensor]:
"""Materialize serialized ModelOpt NVFP4 expert weights."""
return dequantize_nvfp4_moe_weights(
layer.w13_weight,
layer.w13_weight_scale,
layer.w13_weight_scale_2,
layer.w2_weight,
layer.w2_weight_scale,
layer.w2_weight_scale_2,
layer.params_dtype,
group_size=self.quant_config.group_size,
invert_global_scales=False,
)
def create_weights(
self,
layer: RoutedExperts,
vllm/model_executor/layers/quantization/online/fp8.py
+56 -60
View File
@@ -34,10 +34,12 @@ from vllm.model_executor.layers.fusion.quant_activation import (
from vllm.model_executor.layers.linear import (
LinearMethodBase,
)
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.model_executor.layers.quantization.online.moe_base import (
OnlineMoEMethodBase,
)
from vllm.model_executor.layers.quantization.online.requantization import (
OnlineRequantizationMixin,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
GroupShape,
QuantKey,
@@ -118,7 +120,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 OnlineLinearBase(LinearMethodBase):
class OnlineLinearBase(OnlineRequantizationMixin, LinearMethodBase):
"""Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
weights onto meta device and materializes them just-in-time."""
@@ -126,15 +128,17 @@ class OnlineLinearBase(LinearMethodBase):
activation_quant_key: QuantKey | None
def __init__(self):
super().__init__()
self.out_dtype = torch.get_default_dtype()
self.input_dtype = get_current_vllm_config().model_config.dtype
self.requantization_source: QuantizeMethodBase | None = None
self.requantization_source_parameters: dict[str, torch.nn.Parameter] = {}
def set_requantization_source(self, source_method: QuantizeMethodBase) -> None:
"""Configure serialized-weight conversion before online quantization."""
self.requantization_source = source_method
self.uses_meta_device = False
def get_weight_for_quantization(self, layer: torch.nn.Module) -> torch.Tensor:
if self.requantization_source is None:
return layer.weight
weight = self.requantization_source.dequantize_weight(layer)
if not isinstance(weight, torch.Tensor):
raise TypeError("Linear requantization requires a single weight tensor.")
return weight
def create_weights(
self,
@@ -146,22 +150,15 @@ class OnlineLinearBase(LinearMethodBase):
params_dtype: torch.dtype,
**extra_weight_attrs,
):
if self.requantization_source is not None:
existing_parameter_names = set(layer._parameters)
self.requantization_source.create_weights(
layer,
input_size_per_partition,
output_partition_sizes,
input_size,
output_size,
params_dtype,
**extra_weight_attrs,
)
self.requantization_source_parameters = {
name: parameter
for name, parameter in layer._parameters.items()
if name not in existing_parameter_names and parameter is not None
}
if self.create_requantization_source_weights(
layer,
input_size_per_partition,
output_partition_sizes,
input_size,
output_size,
params_dtype,
**extra_weight_attrs,
):
return
output_size_per_partition = sum(output_partition_sizes)
@@ -187,19 +184,6 @@ class OnlineLinearBase(LinearMethodBase):
initialize_online_processing(layer)
def release_requantization_source_weights(self, layer: torch.nn.Module) -> None:
"""Release checkpoint parameters after successful requantization."""
for name, source_parameter in self.requantization_source_parameters.items():
if layer._parameters.get(name) is source_parameter:
delattr(layer, name)
self.requantization_source_parameters.clear()
def get_weight_for_quantization(self, layer: Module) -> torch.Tensor:
"""Return checkpoint weights materialized for online quantization."""
if self.requantization_source is None:
return layer.weight
return self.requantization_source.dequantize_weight(layer)
class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase):
"""Online tensorwise FP8 linear quantization.
@@ -268,6 +252,7 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase):
if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):
self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
self.fp8_linear.process_weights_after_loading(layer)
self.release_requantization_source_weights(layer)
# Prevent duplicate processing (e.g., during weight reload)
layer._already_called_process_weights_after_loading = True
@@ -368,6 +353,7 @@ class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase):
del layer.weight_scale
self.fp8_linear.process_weights_after_loading(layer)
self.release_requantization_source_weights(layer)
# Prevent duplicate processing (e.g., during weight reload)
layer._already_called_process_weights_after_loading = True
@@ -624,27 +610,29 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
w13_weight, w2_weight = self.get_weights_for_quantization(layer)
# If checkpoint is fp16, quantize in place.
fp8_dtype = current_platform.fp8_dtype()
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
w13 = torch.empty_like(w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(w2_weight, dtype=fp8_dtype)
layer.w13_input_scale = None
layer.w2_input_scale = None
moe_tp_size = self.moe.tp_size
w13_amax = weight_amax(layer.w13_weight.flatten(1), dim=-1)
w13_amax = weight_amax(w13_weight.flatten(1), dim=-1)
w13_amax = amax_for_moe_weight_quant(w13_amax, moe_tp_size)
w13_scale = _fp8_scale(w13_amax)
w2_amax = weight_amax(layer.w2_weight.flatten(1), dim=-1)
w2_amax = weight_amax(w2_weight.flatten(1), dim=-1)
w2_amax = amax_for_moe_weight_quant(w2_amax, moe_tp_size)
w2_scale = _fp8_scale(w2_amax)
for expert in range(layer.local_num_experts):
w13[expert, :, :], _ = ops.scaled_fp8_quant(
layer.w13_weight[expert, :, :], scale=w13_scale[expert]
w13_weight[expert, :, :], scale=w13_scale[expert]
)
w2[expert, :, :], _ = ops.scaled_fp8_quant(
layer.w2_weight[expert, :, :], scale=w2_scale[expert]
w2_weight[expert, :, :], scale=w2_scale[expert]
)
# Shuffle weights to runtime format and setup kernel.
@@ -657,6 +645,7 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
w13_input_scale=layer.w13_input_scale,
w2_input_scale=layer.w2_input_scale,
)
self.release_requantization_source_weights(layer)
# Prevent duplicate processing (e.g., during weight reload)
layer._already_called_process_weights_after_loading = True
@@ -700,11 +689,12 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
self._zero_padding(layer)
w13_weight, w2_weight = self.get_weights_for_quantization(layer)
self._zero_padding(layer, w13_weight, w2_weight)
fp8_dtype = current_platform.fp8_dtype()
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
w13 = torch.empty_like(w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(w2_weight, dtype=fp8_dtype)
block_size = self.weight_block_size
assert block_size is not None
@@ -713,8 +703,8 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
# Create block-shaped scales (computed here rather than in
# create_weights because online quant doesn't need them until now).
num_experts = layer.local_num_experts
_, w13_out, w13_in = layer.w13_weight.shape
_, w2_out, w2_in = layer.w2_weight.shape
_, w13_out, w13_in = w13_weight.shape
_, w2_out, w2_in = w2_weight.shape
w13_scale = torch.ones(
num_experts,
@@ -733,18 +723,22 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
for expert in range(num_experts):
w13[expert], w13_scale[expert] = per_block_cast_to_fp8(
layer.w13_weight[expert],
w13_weight[expert],
block_size=block_size,
use_ue8m0=False,
)
w2[expert], w2_scale[expert] = per_block_cast_to_fp8(
layer.w2_weight[expert],
w2_weight[expert],
block_size=block_size,
use_ue8m0=False,
)
layer.weight_block_size = block_size
# Activation quantization is dynamic (including in re-quantization)
layer.w13_input_scale = None
layer.w2_input_scale = None
# Shuffle weights to runtime format and setup kernel.
self._setup_kernel(
layer,
@@ -752,9 +746,10 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
w2,
w13_scale,
w2_scale,
layer.w13_input_scale,
layer.w2_input_scale,
w13_input_scale=None,
w2_input_scale=None,
)
self.release_requantization_source_weights(layer)
# Prevent duplicate processing (e.g., during weight reload)
layer._already_called_process_weights_after_loading = True
@@ -806,32 +801,32 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
w13_weight, w2_weight = self.get_weights_for_quantization(layer)
fp8_dtype = current_platform.fp8_dtype()
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
w13 = torch.empty_like(w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(w2_weight, dtype=fp8_dtype)
# Scale's leading dim is taken from the fp8 weight tensor by
# construction, so it cannot drift from the weight's expert count
# under EP / padded MoE.
n_w13 = layer.w13_weight.shape[1]
n_w13 = w13_weight.shape[1]
w13_scale = torch.ones(
w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32
)
layer.w13_input_scale = None
layer.w2_input_scale = None
w2_amax = weight_amax(layer.w2_weight, dim=-1, keepdim=True)
w2_amax = weight_amax(w2_weight, dim=-1, keepdim=True)
w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size)
w2_scale = _fp8_channel_scale(w2_amax)
for expert in range(layer.local_num_experts):
w13[expert], w13_scale[expert] = ops.scaled_fp8_quant(
layer.w13_weight[expert],
w13_weight[expert],
scale=None,
use_per_token_if_dynamic=True,
)
w2[expert] = _fp8_quant_per_channel(
layer.w2_weight[expert], w2_scale[expert]
)
w2[expert] = _fp8_quant_per_channel(w2_weight[expert], w2_scale[expert])
self._setup_kernel(
layer,
@@ -842,5 +837,6 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
w13_input_scale=None,
w2_input_scale=None,
)
self.release_requantization_source_weights(layer)
layer._already_called_process_weights_after_loading = True
vllm/model_executor/layers/quantization/online/int8.py
+18 -9
View File
@@ -52,42 +52,51 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
self._quantize_weights(layer)
w13_weight, w2_weight = self.get_weights_for_quantization(layer)
self._quantize_weights(layer, w13_weight, w2_weight)
self._setup_kernel(layer)
self.release_requantization_source_weights(layer)
layer._already_called_process_weights_after_loading = True
def _quantize_weights(self, layer: Module) -> None:
def _quantize_weights(
self,
layer: Module,
w13_weight: torch.Tensor | None = None,
w2_weight: torch.Tensor | None = None,
) -> None:
w13_weight = layer.w13_weight if w13_weight is None else w13_weight
w2_weight = layer.w2_weight if w2_weight is None else w2_weight
vmax = torch.iinfo(torch.int8).max
w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8)
w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8)
w13 = torch.empty_like(w13_weight, dtype=torch.int8)
w2 = torch.empty_like(w2_weight, dtype=torch.int8)
w13_scale = torch.zeros(
layer.num_experts,
layer.w13_weight.shape[1],
w13_weight.shape[1],
device=w13.device,
dtype=torch.float32,
)
w2_scale = torch.zeros(
layer.num_experts,
layer.w2_weight.shape[1],
w2_weight.shape[1],
device=w2.device,
dtype=torch.float32,
)
w2_amax = weight_amax(layer.w2_weight, dim=-1)
w2_amax = weight_amax(w2_weight, dim=-1)
w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size)
for expert in range(layer.local_num_experts):
# w13: per-row quantization over hidden_size dim
w = layer.w13_weight[expert, :, :]
w = w13_weight[expert, :, :]
scales = w.abs().amax(dim=1) / vmax
q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
w13[expert, :, :] = q.to(torch.int8)
w13_scale[expert, :] = scales
# w2: per-row quantization over intermediate_size dim
w = layer.w2_weight[expert, :, :]
w = w2_weight[expert, :, :]
scales = w2_amax[expert] / vmax
q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
w2[expert, :, :] = q.to(torch.int8)
vllm/model_executor/layers/quantization/online/moe_base.py
+101 -16
View File
@@ -10,26 +10,26 @@ from vllm.model_executor.layers.fused_moe import (
RoutedExperts,
SharedExperts,
)
from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.model_executor.layers.quantization.online.requantization import (
OnlineRequantizationMixin,
)
from vllm.model_executor.model_loader.reload.layerwise import (
initialize_online_processing,
)
from vllm.model_executor.utils import set_weight_attrs
class OnlineMoEMethodBase(FusedMoEMethodBase):
class OnlineMoEMethodBase(OnlineRequantizationMixin, FusedMoEMethodBase):
"""Base for MoE methods that load full-precision weights on meta device
and quantize them after loading via the QeRL layerwise processing system.
"""
uses_meta_device: bool = True
def set_requantization_source(self, source_method: QuantizeMethodBase) -> None:
"""Reject requantization from a checkpoint-quantized MoE method."""
raise NotImplementedError(
"Requantizing checkpoint-quantized MoE layers is not supported."
)
def __init__(self, moe: FusedMoEConfig):
super().__init__(moe)
def create_weights(
self,
@@ -40,10 +40,22 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
params_dtype: torch.dtype,
**extra_weight_attrs,
):
# Target conversion needs this metadata even when the source creates weights.
layer.num_experts = num_experts
layer.orig_dtype = params_dtype
layer.weight_block_size = None
if self.requantization_source is not None:
self.create_requantization_source_weights(
layer,
num_experts,
layer.moe_config.hidden_dim_unpadded,
layer.moe_config.intermediate_size_per_partition_unpadded,
params_dtype,
**extra_weight_attrs,
)
return
# Fused gate_up_proj (column parallel) — full precision on meta device
w13_weight = torch.nn.Parameter(
torch.empty(
@@ -103,22 +115,95 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
initialize_online_processing(layer)
def _zero_padding(self, layer: torch.nn.Module) -> None:
def get_weights_for_quantization(
self, layer: torch.nn.Module
) -> tuple[torch.Tensor, torch.Tensor]:
if self.requantization_source is None:
return layer.w13_weight, layer.w2_weight
weights = self.requantization_source.dequantize_weight(layer)
if not isinstance(weights, tuple):
raise TypeError("MoE requantization requires w13 and w2 weight tensors.")
w13, w2 = weights
config = layer.moe_config
hidden_size = config.hidden_dim_unpadded
intermediate_size = config.intermediate_size_per_partition_unpadded
target_hidden_size = config.hidden_dim
target_intermediate_size = config.intermediate_size_per_partition
num_experts = layer.local_num_experts
num_shards = self.moe.w13_num_shards
if w13.shape != (num_experts, num_shards * intermediate_size, hidden_size):
raise ValueError("Requantization requires unpadded source w13 weights.")
if w2.shape != (num_experts, hidden_size, intermediate_size):
raise ValueError("Requantization requires unpadded source w2 weights.")
if (
target_hidden_size < hidden_size
or target_intermediate_size < intermediate_size
):
raise ValueError("Requantization target dimensions cannot shrink weights.")
# No padding required.
if (hidden_size, intermediate_size) == (
target_hidden_size,
target_intermediate_size,
):
return w13, w2
if any(
getattr(layer, name, None) is not None for name in ("w13_bias", "w2_bias")
):
raise NotImplementedError(
"Padded requantization with biases is unsupported."
)
def pad_moe_weights(
w13: torch.Tensor, w2: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
padded_w13 = w13.new_zeros(
(num_experts, num_shards * target_intermediate_size, target_hidden_size)
)
padded_w2 = w2.new_zeros(
(num_experts, target_hidden_size, target_intermediate_size)
)
for shard in range(num_shards):
source_start = shard * intermediate_size
target_start = shard * target_intermediate_size
padded_w13[
:, target_start : target_start + intermediate_size, :hidden_size
] = w13[:, source_start : source_start + intermediate_size, :]
padded_w2[:, :hidden_size, :intermediate_size] = w2
return padded_w13, padded_w2
return pad_moe_weights(w13, w2)
def _zero_padding(
self,
layer: torch.nn.Module,
w13_weight: torch.Tensor | None = None,
w2_weight: torch.Tensor | None = None,
) -> None:
w13_weight = layer.w13_weight if w13_weight is None else w13_weight
w2_weight = layer.w2_weight if w2_weight is None else w2_weight
hidden_size = layer.moe_config.hidden_dim_unpadded
intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded
w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards
w13_shard = w13_weight.shape[1] // self.moe.w13_num_shards
if w13_shard > intermediate_size:
for shard in range(self.moe.w13_num_shards):
start = shard * w13_shard + intermediate_size
layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0
if layer.w13_weight.shape[2] > hidden_size:
layer.w13_weight[:, :, hidden_size:] = 0
w13_weight[:, start : (shard + 1) * w13_shard, :] = 0
if w13_weight.shape[2] > hidden_size:
w13_weight[:, :, hidden_size:] = 0
if layer.w2_weight.shape[1] > hidden_size:
layer.w2_weight[:, hidden_size:, :] = 0
if layer.w2_weight.shape[2] > intermediate_size:
layer.w2_weight[:, :, intermediate_size:] = 0
if w2_weight.shape[1] > hidden_size:
w2_weight[:, hidden_size:, :] = 0
if w2_weight.shape[2] > intermediate_size:
w2_weight[:, :, intermediate_size:] = 0
if getattr(layer, "w13_bias", None) is not None:
w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards
vllm/model_executor/layers/quantization/online/mxfp4.py
+6 -3
View File
@@ -127,6 +127,7 @@ class Mxfp4OnlineLinearMethod(OnlineLinearBase):
replace_parameter(layer, "weight_scale", weight_scale.data)
self.kernel.process_weights_after_loading(layer)
self.release_requantization_source_weights(layer)
layer._already_called_process_weights_after_loading = True
@@ -295,7 +296,8 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
self._zero_padding(layer)
w13_weight, w2_weight = self.get_weights_for_quantization(layer)
self._zero_padding(layer, w13_weight, w2_weight)
if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
layer._already_called_process_weights_after_loading = True
@@ -304,8 +306,8 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):
layer.w13_input_scale = None
layer.w2_input_scale = None
w13, w13_scale = _quantize_mxfp4_moe_weight(layer.w13_weight)
w2, w2_scale = _quantize_mxfp4_moe_weight(layer.w2_weight)
w13, w13_scale = _quantize_mxfp4_moe_weight(w13_weight)
w2, w2_scale = _quantize_mxfp4_moe_weight(w2_weight)
self._setup_kernel(
layer,
@@ -316,5 +318,6 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):
getattr(layer, "w13_bias", None),
getattr(layer, "w2_bias", None),
)
self.release_requantization_source_weights(layer)
layer._already_called_process_weights_after_loading = True
vllm/model_executor/layers/quantization/online/mxfp8.py
+5 -6
View File
@@ -34,7 +34,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
mxfp8_e4m3_quantize,
)
from vllm.model_executor.utils import replace_parameter
from vllm.platforms import current_platform
class Mxfp8OnlineLinearMethod(OnlineLinearBase):
@@ -86,6 +85,7 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase):
replace_parameter(layer, "weight_scale", weight_scale.data)
self.kernel.process_weights_after_loading(layer)
self.release_requantization_source_weights(layer)
layer._already_called_process_weights_after_loading = True
@@ -237,14 +237,12 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
fp8_dtype = current_platform.fp8_dtype()
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
w13_weight, w2_weight = self.get_weights_for_quantization(layer)
layer.w13_input_scale = None
layer.w2_input_scale = None
w13, w13_scale = self._quantize_mxfp8_moe_weight(layer.w13_weight)
w2, w2_scale = self._quantize_mxfp8_moe_weight(layer.w2_weight)
w13, w13_scale = self._quantize_mxfp8_moe_weight(w13_weight)
w2, w2_scale = self._quantize_mxfp8_moe_weight(w2_weight)
self._setup_kernel(
layer,
@@ -255,5 +253,6 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
layer.w13_input_scale,
layer.w2_input_scale,
)
self.release_requantization_source_weights(layer)
layer._already_called_process_weights_after_loading = True
vllm/model_executor/layers/quantization/online/nvfp4.py
+13 -6
View File
@@ -116,19 +116,26 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
if getattr(layer, "_already_called_process_weights_after_loading", False):
return
self._quantize_weights(layer)
w13_weight, w2_weight = self.get_weights_for_quantization(layer)
self._quantize_weights(layer, w13_weight, w2_weight)
self._setup_kernel(layer)
self.release_requantization_source_weights(layer)
layer._already_called_process_weights_after_loading = True
def _quantize_weights(self, layer: Module) -> None:
def _quantize_weights(
self,
layer: Module,
w13_weight: torch.Tensor | None = None,
w2_weight: torch.Tensor | None = None,
) -> None:
w13_weight = layer.w13_weight if w13_weight is None else w13_weight
w2_weight = layer.w2_weight if w2_weight is None else w2_weight
moe_tp_size = self.moe.tp_size
w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(
layer.w13_weight, moe_tp_size
)
w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(
layer.w2_weight, moe_tp_size
w13_weight, moe_tp_size
)
w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(w2_weight, moe_tp_size)
replace_parameter(layer, "w13_weight", w13)
replace_parameter(layer, "w13_weight_scale", w13_scale)
vllm/model_executor/layers/quantization/online/requantization.py
+45
View File
@@ -0,0 +1,45 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
class OnlineRequantizationMixin:
"""Shared lifecycle for converting serialized weights before quantization."""
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.requantization_source: QuantizeMethodBase | None = None
self.requantization_source_parameters: dict[str, torch.nn.Parameter] = {}
def set_requantization_source(self, source_method: QuantizeMethodBase) -> None:
"""Configure serialized-weight conversion before online quantization."""
self.requantization_source = source_method
self.uses_meta_device = False
def create_requantization_source_weights(
self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
) -> bool:
"""Delegate weight creation to the source and track its parameters."""
if self.requantization_source is None:
return False
existing_parameter_names = set(layer._parameters)
self.requantization_source.create_weights(
layer, *weight_args, **extra_weight_attrs
)
self.requantization_source_parameters = {
name: parameter
for name, parameter in layer._parameters.items()
if name not in existing_parameter_names and parameter is not None
}
return True
def release_requantization_source_weights(self, layer: torch.nn.Module) -> None:
"""Release serialized parameters after successful requantization."""
for name, source_parameter in self.requantization_source_parameters.items():
if layer._parameters.get(name) is source_parameter:
delattr(layer, name)
self.requantization_source_parameters.clear()
vllm/model_executor/layers/quantization/quark/quark_moe.py
+19
View File
@@ -69,6 +69,9 @@ from vllm.model_executor.layers.quantization.quark.utils import (
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
validate_fp8_block_shape_moe,
)
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
dequantize_nvfp4_moe_weights,
)
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
_ACTIVATION_QUANT_KEY_MAP,
_WEIGHT_QUANT_KEY_MAP,
@@ -2110,6 +2113,22 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):
layer.register_parameter("w2_input_scale_2", w2_input_scale_2)
set_weight_attrs(w2_input_scale_2, extra_weight_attrs)
def dequantize_weight(
self, layer: torch.nn.Module
) -> tuple[torch.Tensor, torch.Tensor]:
"""Dequantize Quark NVFP4 expert weights for requantization."""
return dequantize_nvfp4_moe_weights(
layer.w13_weight,
layer.w13_weight_scale,
layer.w13_weight_scale_2,
layer.w2_weight,
layer.w2_weight_scale,
layer.w2_weight_scale_2,
layer.params_dtype,
group_size=self.group_size,
invert_global_scales=False,
)
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
"""Convert NVFP4 MoE weights into kernel format and setup the kernel."""
# Match existing NVFP4 MoE paths: fused w13 uses the w1 global scale.
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+53
View File
@@ -11,6 +11,7 @@ from vllm.utils.torch_utils import direct_register_custom_op
__all__ = [
"break_fp4_bytes",
"dequantize_nvfp4_moe_weights",
"dequantize_to_dtype",
"nvfp4_gathered_bias",
"ref_nvfp4_quant",
@@ -629,6 +630,58 @@ def dequantize_to_dtype(
return out.to(dtype)
def dequantize_nvfp4_moe_weights(
w13_weight_packed: torch.Tensor,
w13_weight_scale: torch.Tensor,
w13_weight_global_scale: torch.Tensor,
w2_weight_packed: torch.Tensor,
w2_weight_scale: torch.Tensor,
w2_weight_global_scale: torch.Tensor,
dtype: torch.dtype,
group_size: int = 16,
*,
invert_global_scales: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Materialize serialized NVFP4 MoE weights for requantization.
Compressed-tensors stores global encode scales, while ModelOpt stores
``weight_scale_2`` as the corresponding decode multiplier. Set
``invert_global_scales`` according to the source checkpoint convention.
"""
if invert_global_scales:
w13_weight_global_scale = 1.0 / w13_weight_global_scale
w2_weight_global_scale = 1.0 / w2_weight_global_scale
w13_num_shards = w13_weight_global_scale.shape[1]
w13_weights = torch.chunk(w13_weight_packed, w13_num_shards, dim=1)
w13_scales = torch.chunk(w13_weight_scale, w13_num_shards, dim=1)
w13 = torch.cat(
[
dequantize_to_dtype(
weight,
scale,
w13_weight_global_scale[:, shard_idx].contiguous(),
dtype,
block_size=group_size,
swizzle=False,
)
for shard_idx, (weight, scale) in enumerate(
zip(w13_weights, w13_scales, strict=True)
)
],
dim=1,
)
w2 = dequantize_to_dtype(
w2_weight_packed,
w2_weight_scale,
w2_weight_global_scale,
dtype,
block_size=group_size,
swizzle=False,
)
return w13, w2
def get_reciprocal(x):
if isinstance(x, torch.Tensor):
# torch.where yields operation not permitted when stream is capturing.