[Quantization] Share online weight scales across TP (#49764)

Signed-off-by: S1ro1 <matej.sirovatka@gmail.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Roger Wang <hey@rogerw.io>

tests/quantization/test_online.py

@@ -2,13 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests online quantization."""
 
+from types import SimpleNamespace
+from typing import cast
+
 import pytest
 import torch
+from torch.distributed import ProcessGroup
 
 from tests.quantization.utils import (
     _test_online_quant_peak_mem_impl,
     is_quant_method_supported,
 )
+from vllm import _custom_ops as ops
 from vllm._custom_ops import scaled_fp4_quant
 from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.online.fp8 import (
@@ -16,11 +21,22 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
     Fp8PerBlockOnlineMoEMethod,
     Fp8PerTensorOnlineLinearMethod,
     Fp8PerTensorOnlineMoEMethod,
+    _fp8_channel_scale,
+    _fp8_quant_per_channel,
+    _fp8_scale,
+    _is_tp_sharded,
 )
+from vllm.model_executor.layers.quantization.online.int8 import Int8OnlineMoEMethod
 from vllm.model_executor.layers.quantization.online.nvfp4 import (
     Nvfp4OnlineMoEMethod,
     _quantize_moe_weight_to_nvfp4,
 )
+from vllm.model_executor.layers.quantization.utils import quant_utils
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    amax_for_moe_weight_quant,
+    amax_for_tp_weight_quant,
+    weight_amax,
+)
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 
@@ -186,6 +202,169 @@ def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None:
         print(outputs[0][1])
 
 
+def _patch_max_reduce(monkeypatch, full_amax) -> None:
+    """Stand in for the TP/EP MAX all-reduce, returning the unsharded amax."""
+    expected = cast(ProcessGroup, object())
+    stub = SimpleNamespace(device_group=expected)
+    monkeypatch.setattr(quant_utils, "get_tp_group", lambda: stub)
+    monkeypatch.setattr(quant_utils, "get_ep_group", lambda: stub)
+
+    def fake_all_reduce(tensor, op, group):
+        assert op == torch.distributed.ReduceOp.MAX
+        assert group is expected
+        tensor.copy_(full_amax)
+
+    monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce)
+
+
+def test_is_tp_sharded_false_when_scale_is_already_global() -> None:
+    """Replicated and column-parallel-with-channel-scales need no collective."""
+    replicated = SimpleNamespace(
+        tp_size=4,
+        input_size=64,
+        output_size=32,
+        input_size_per_partition=64,
+        output_size_per_partition=32,
+    )
+    assert not _is_tp_sharded(replicated)
+
+    column = SimpleNamespace(
+        tp_size=4,
+        input_size=64,
+        output_size=32,
+        input_size_per_partition=64,
+        output_size_per_partition=8,
+    )
+    assert not _is_tp_sharded(column, reduces_output_dim=False)
+    assert _is_tp_sharded(column)
+
+
+def _quantize_linear(weight, scheme, is_sharded):
+    if scheme == "per_tensor":
+        amax = weight_amax(weight).reshape(1)
+        scale = _fp8_scale(amax_for_tp_weight_quant(amax, is_sharded))
+        return ops.scaled_fp8_quant(weight, scale=scale)[0], scale
+    amax = weight_amax(weight, dim=-1, keepdim=True)
+    scale = _fp8_channel_scale(amax_for_tp_weight_quant(amax, is_sharded))
+    return _fp8_quant_per_channel(weight, scale), scale
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+@pytest.mark.parametrize("scheme", ["per_tensor", "per_channel"])
+@pytest.mark.parametrize("shard_dim", [0, 1])
+def test_online_linear_tp_weight_quant_matches_unsharded(
+    monkeypatch, scheme: str, shard_dim: int
+) -> None:
+    """TP shards pack the same FP8 values and scales as the unsharded weight."""
+    torch.manual_seed(0)
+    weight = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
+    weight[-1, -1] = 64.0
+
+    full_weight, full_scale = _quantize_linear(weight, scheme, False)
+
+    # Per-channel scales reduce only the input dim, so a column (dim 0) shard
+    # already matches without a collective.
+    is_sharded = scheme == "per_tensor" or shard_dim == 1
+    if is_sharded:
+        full_amax = (
+            weight_amax(weight).reshape(1)
+            if scheme == "per_tensor"
+            else weight_amax(weight, dim=-1, keepdim=True)
+        )
+        _patch_max_reduce(monkeypatch, full_amax)
+
+    shard_size = weight.shape[shard_dim] // 2
+    shard = weight.narrow(shard_dim, 0, shard_size).contiguous()
+    tp_weight, tp_scale = _quantize_linear(shard, scheme, is_sharded)
+
+    assert torch.equal(tp_weight, full_weight.narrow(shard_dim, 0, shard_size))
+    if scheme == "per_channel" and shard_dim == 0:
+        assert torch.equal(tp_scale, full_scale.narrow(0, 0, shard_size))
+    else:
+        assert torch.equal(tp_scale, full_scale)
+
+
+def _quantize_moe(weight, scheme, moe_tp_size):
+    if scheme == "nvfp4":
+        return _quantize_moe_weight_to_nvfp4(weight, moe_tp_size)
+    if scheme == "per_tensor":
+        amax = weight_amax(weight.flatten(1), dim=-1)
+        scale = _fp8_scale(amax_for_moe_weight_quant(amax, moe_tp_size))
+        quant = lambda w, s: ops.scaled_fp8_quant(w, scale=s)[0]  # noqa: E731
+    else:
+        amax = weight_amax(weight, dim=-1, keepdim=True)
+        scale = _fp8_channel_scale(amax_for_moe_weight_quant(amax, moe_tp_size))
+        quant = _fp8_quant_per_channel
+    qweight = torch.stack([quant(w, s) for w, s in zip(weight, scale)])
+    return qweight, scale
+
+
+@pytest.mark.parametrize("scheme", ["per_tensor", "per_channel", "nvfp4"])
+def test_online_moe_tp_weight_quant_matches_ep(monkeypatch, scheme: str) -> None:
+    """TP shards of w2 pack the same values and scales as full experts."""
+    if scheme == "nvfp4":
+        if not (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(100)
+        ):
+            pytest.skip("NVFP4 weight quantization needs a Blackwell (SM100) GPU.")
+    elif not is_quant_method_supported("fp8"):
+        pytest.skip("FP8 is not supported on this GPU type.")
+
+    torch.manual_seed(0)
+    weight = torch.randn(2, 32, 32, device="cuda", dtype=torch.bfloat16)
+    weight[:, -1, -1] = torch.tensor([32.0, 64.0], device="cuda")
+
+    ep_out = _quantize_moe(weight, scheme, 1)
+
+    full_amax = (
+        weight_amax(weight, dim=-1, keepdim=True)
+        if scheme == "per_channel"
+        else weight_amax(weight.flatten(1), dim=-1).to(torch.float32)
+    )
+    _patch_max_reduce(monkeypatch, full_amax)
+
+    # w2 is sharded along its last (intermediate) dim.
+    shard_size = weight.shape[2] // 2
+    tp_out = _quantize_moe(weight[:, :, :shard_size], scheme, 2)
+
+    packing = 2 if scheme == "nvfp4" else 1
+    assert torch.equal(tp_out[0], ep_out[0][:, :, : shard_size // packing])
+    if scheme == "nvfp4":
+        assert torch.equal(tp_out[1], ep_out[1][:, :, : shard_size // 16])
+    assert torch.equal(tp_out[-1], ep_out[-1])
+
+
+def test_online_int8_moe_w2_scale_matches_unsharded(monkeypatch) -> None:
+    """Int8 MoE w2 reduces over the sharded intermediate dim."""
+    torch.manual_seed(0)
+    w13 = torch.randn(2, 16, 8, dtype=torch.bfloat16)
+    w2 = torch.randn(2, 8, 16, dtype=torch.bfloat16)
+    w2[:, -1, -1] = 64.0
+
+    def quantize(w2_in, moe_tp_size):
+        layer = torch.nn.Module()
+        layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False)
+        layer.w2_weight = torch.nn.Parameter(w2_in, requires_grad=False)
+        layer.num_experts = layer.local_num_experts = w13.shape[0]
+        method = SimpleNamespace(moe=SimpleNamespace(tp_size=moe_tp_size))
+        Int8OnlineMoEMethod._quantize_weights(method, layer)
+        return layer.w2_weight, layer.w2_scale
+
+    full_weight, full_scale = quantize(w2, 1)
+
+    _patch_max_reduce(monkeypatch, weight_amax(w2, dim=-1))
+
+    shard_size = w2.shape[2] // 2
+    tp_weight, tp_scale = quantize(w2[:, :, :shard_size].contiguous(), 2)
+
+    assert torch.equal(tp_weight, full_weight[:, :, :shard_size])
+    assert torch.equal(tp_scale, full_scale)
+
+
 @pytest.mark.skipif(
     not (
         current_platform.is_cuda() and current_platform.is_device_capability_family(100)

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

@@ -34,13 +34,17 @@ from vllm.model_executor.layers.quantization.online.moe_base import (
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
+    amax_for_moe_weight_quant,
+    amax_for_tp_weight_quant,
     create_fp8_quant_key,
+    get_fp8_min_max,
     kFp8Dynamic128Sym,
     kFp8DynamicTensorSym,
     kFp8DynamicTokenSym,
     kFp8Static128BlockSym,
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
+    weight_amax,
 )
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
     cutlass_fp8_supported,
@@ -59,6 +63,54 @@ from vllm.utils.math_utils import round_up
 # ---------------------------------------------------------------------------
 
 
+def _fp8_max(device: torch.device) -> torch.Tensor:
+    # A 0-d tensor divisor keeps the true division the quant kernels do;
+    # dividing by a Python float lowers to multiply-by-reciprocal instead.
+    return torch.tensor(get_fp8_min_max()[1], device=device, dtype=torch.float32)
+
+
+def _fp8_scale(amax: torch.Tensor) -> torch.Tensor:
+    amax = amax.to(torch.float32)
+    return amax / _fp8_max(amax.device)
+
+
+def _fp8_channel_scale(amax: torch.Tensor) -> torch.Tensor:
+    fp8_max = _fp8_max(amax.device)
+    return _fp8_scale(amax).clamp_min(1.0 / (fp8_max * 512))
+
+
+# Chunk so the fp32 intermediate from the divide stays under ~64 MB, regardless
+# of how wide the weight is.
+_QUANT_CHUNK_ELEMS = 16 * 1024 * 1024
+
+
+def _fp8_quant_per_channel(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    # True division, matching the dynamic per-token kernel; the static kernel
+    # multiplies by the reciprocal instead and can round a value that sits on a
+    # representable midpoint the other way.
+    fp8_min, fp8_max = get_fp8_min_max()
+    out = torch.empty_like(weight, dtype=current_platform.fp8_dtype())
+    chunk_rows = max(1, _QUANT_CHUNK_ELEMS // weight.shape[-1])
+    for start in range(0, weight.shape[0], chunk_rows):
+        rows = slice(start, start + chunk_rows)
+        out[rows] = (weight[rows] / scale[rows]).clamp_(fp8_min, fp8_max).to(out.dtype)
+    return out
+
+
+def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool:
+    """Whether the weight is sharded along a dim its ``amax`` reduces over.
+
+    Row parallel shards the input dim, column parallel the output dim.
+    Per-output-channel scales only reduce the input dim, so they are already
+    unsharded-equivalent on column-parallel layers.
+    """
+    if layer.tp_size == 1:
+        return False
+    is_row_parallel = layer.input_size != layer.input_size_per_partition
+    is_column_parallel = layer.output_size != layer.output_size_per_partition
+    return is_row_parallel or (reduces_output_dim and is_column_parallel)
+
+
 class _Fp8OnlineLinearBase(LinearMethodBase):
     """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
     weights onto meta device and materializes them just-in-time."""
@@ -156,7 +208,10 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
             return
 
         layer.input_scale = None
-        qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
+        amax = weight_amax(layer.weight).reshape(1)
+        amax = amax_for_tp_weight_quant(amax, _is_tp_sharded(layer))
+        weight_scale = _fp8_scale(amax)
+        qweight, _ = ops.scaled_fp8_quant(layer.weight, scale=weight_scale)
 
         # Update layer with new values.
         replace_parameter(layer, "weight", qweight.t().data)
@@ -335,9 +390,12 @@ class Fp8PtpcOnlineLinearMethod(_Fp8OnlineLinearBase):
             return
 
         layer.input_scale = None
-        qweight, weight_scale = ops.scaled_fp8_quant(
-            layer.weight, scale=None, use_per_token_if_dynamic=True
+        amax = weight_amax(layer.weight, dim=-1, keepdim=True)
+        amax = amax_for_tp_weight_quant(
+            amax, _is_tp_sharded(layer, reduces_output_dim=False)
         )
+        weight_scale = _fp8_channel_scale(amax)
+        qweight = _fp8_quant_per_channel(layer.weight, weight_scale)
 
         replace_parameter(layer, "weight", qweight.t())
         replace_parameter(layer, "weight_scale", weight_scale)
@@ -514,19 +572,23 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
         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_scale = torch.ones(
-            layer.num_experts, device=w13.device, dtype=torch.float32
-        )
-        w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)
         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 = 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 = 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, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
-                layer.w13_weight[expert, :, :]
+            w13[expert, :, :], _ = ops.scaled_fp8_quant(
+                layer.w13_weight[expert, :, :], scale=w13_scale[expert]
             )
-            w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
-                layer.w2_weight[expert, :, :]
+            w2[expert, :, :], _ = ops.scaled_fp8_quant(
+                layer.w2_weight[expert, :, :], scale=w2_scale[expert]
             )
 
         # Shuffle weights to runtime format and setup kernel.
@@ -723,26 +785,24 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
         # construction, so it cannot drift from the weight's expert count
         # under EP / padded MoE.
         n_w13 = layer.w13_weight.shape[1]
-        n_w2 = layer.w2_weight.shape[1]
         w13_scale = torch.ones(
             w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32
         )
-        w2_scale = torch.ones(
-            w2.shape[0], n_w2, 1, device=w2.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 = 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],
                 scale=None,
                 use_per_token_if_dynamic=True,
             )
-            w2[expert], w2_scale[expert] = ops.scaled_fp8_quant(
-                layer.w2_weight[expert],
-                scale=None,
-                use_per_token_if_dynamic=True,
+            w2[expert] = _fp8_quant_per_channel(
+                layer.w2_weight[expert], w2_scale[expert]
             )
 
         self._setup_kernel(

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

@@ -22,8 +22,10 @@ from vllm.model_executor.layers.quantization.online.moe_base import (
     OnlineMoEMethodBase,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    amax_for_moe_weight_quant,
     kInt8DynamicTokenSym,
     kInt8StaticChannelSym,
+    weight_amax,
 )
 from vllm.model_executor.utils import replace_parameter
 
@@ -72,6 +74,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
             dtype=torch.float32,
         )
 
+        w2_amax = weight_amax(layer.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, :, :]
@@ -82,7 +87,7 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
 
             # w2: per-row quantization over intermediate_size dim
             w = layer.w2_weight[expert, :, :]
-            scales = w.abs().amax(dim=1) / vmax
+            scales = w2_amax[expert] / vmax
             q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
             w2[expert, :, :] = q.to(torch.int8)
             w2_scale[expert, :] = scales

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

@@ -20,8 +20,10 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
     FLOAT4_E2M1_MAX,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    amax_for_moe_weight_quant,
     kNvfp4Dynamic,
     kNvfp4Static,
+    weight_amax,
 )
 from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
@@ -31,6 +33,7 @@ FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
 
 def _quantize_moe_weight_to_nvfp4(
     weight: torch.Tensor,
+    moe_tp_size: int = 1,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Quantize stacked MoE expert weights ``(E, N, K)`` to NVFP4.
 
@@ -43,7 +46,8 @@ def _quantize_moe_weight_to_nvfp4(
     k = weight.shape[-1]
     assert k % 16 == 0, f"last dim must be a multiple of 16, got {k}"
 
-    amax = weight.abs().amax(dim=(1, 2)).to(torch.float32).clamp_min(1e-8)
+    amax = weight_amax(weight.flatten(1), dim=-1).to(torch.float32)
+    amax = amax_for_moe_weight_quant(amax, moe_tp_size).clamp_min(1e-8)
     global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax
     weight_scale_2 = (1.0 / global_scale).to(torch.float32)
 
@@ -106,8 +110,13 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
         layer._already_called_process_weights_after_loading = True
 
     def _quantize_weights(self, layer: Module) -> None:
-        w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w13_weight)
-        w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(layer.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
+        )
 
         replace_parameter(layer, "w13_weight", w13)
         replace_parameter(layer, "w13_weight_scale", w13_scale)

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

@@ -11,6 +11,7 @@ import numpy
 import torch
 from torch import fx
 
+from vllm.distributed.parallel_state import get_ep_group, get_tp_group
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
 
@@ -24,6 +25,43 @@ INT4_DTYPE = scalar_types.uint4b8
 INT8_DTYPE = scalar_types.uint8b128
 
 
+def weight_amax(
+    weight: torch.Tensor, *, dim: int | None = None, keepdim: bool = False
+) -> torch.Tensor:
+    """``max(|weight|)``, without materializing a full-size ``abs()``."""
+    lo, hi = weight.aminmax(dim=dim, keepdim=keepdim)
+    return torch.maximum(lo.abs(), hi.abs())
+
+
+def amax_for_tp_weight_quant(amax: torch.Tensor, is_sharded: bool) -> torch.Tensor:
+    """Reduce a weight ``amax`` over the TP group when the weight is sharded
+    along a dim the ``amax`` reduces over, so each shard derives the same scale
+    it would as part of the whole weight.
+    """
+    if is_sharded:
+        torch.distributed.all_reduce(
+            amax,
+            op=torch.distributed.ReduceOp.MAX,
+            group=get_tp_group().device_group,
+        )
+    return amax
+
+
+def amax_for_moe_weight_quant(amax: torch.Tensor, moe_tp_size: int) -> torch.Tensor:
+    """Reduce a per-expert weight ``amax`` over the ranks that tensor-shard the
+    MoE weights. That sharding is flattened over DP x PCP x TP, exactly the EP
+    group's span. Under EP ``moe_tp_size`` is 1 and each rank owns whole
+    experts, so no reduction is needed.
+    """
+    if moe_tp_size > 1:
+        torch.distributed.all_reduce(
+            amax,
+            op=torch.distributed.ReduceOp.MAX,
+            group=get_ep_group().device_group,
+        )
+    return amax
+
+
 def get_fp8_min_max() -> tuple[float, float]:
     """Get the min and max values for FP8 quantization."""
     # Using the default value (240.0) from pytorch will cause accuracy