[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>
Matej Sirovatka committed
Aug 6, 2026 at 16:34 UTC
8170c23c4fa36ffdc5890e5df46b4825fd9d0745
5 files changed
+314
-23
tests/quantization/test_online.py
+179
@@ -2,13 +2,18 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Tests online quantization."""
4
5
+from types import SimpleNamespace
6
+from typing import cast
7
+
8
import pytest
9
import torch
10
+from torch.distributed import ProcessGroup
11
12
from tests.quantization.utils import (
13
_test_online_quant_peak_mem_impl,
14
is_quant_method_supported,
15
)
16
+from vllm import _custom_ops as ops
17
from vllm._custom_ops import scaled_fp4_quant
18
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
19
from vllm.model_executor.layers.quantization.online.fp8 import (
@@ -16,11 +21,22 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
21
Fp8PerBlockOnlineMoEMethod,
22
Fp8PerTensorOnlineLinearMethod,
23
Fp8PerTensorOnlineMoEMethod,
24
+ _fp8_channel_scale,
25
+ _fp8_quant_per_channel,
26
+ _fp8_scale,
27
+ _is_tp_sharded,
28
)
29
+from vllm.model_executor.layers.quantization.online.int8 import Int8OnlineMoEMethod
30
from vllm.model_executor.layers.quantization.online.nvfp4 import (
31
Nvfp4OnlineMoEMethod,
32
_quantize_moe_weight_to_nvfp4,
33
)
34
+from vllm.model_executor.layers.quantization.utils import quant_utils
35
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
36
+ amax_for_moe_weight_quant,
37
+ amax_for_tp_weight_quant,
38
+ weight_amax,
39
+)
40
from vllm.platforms import current_platform
41
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
42
@@ -186,6 +202,169 @@ def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None:
202
print(outputs[0][1])
203
204
205
+def _patch_max_reduce(monkeypatch, full_amax) -> None:
206
+ """Stand in for the TP/EP MAX all-reduce, returning the unsharded amax."""
207
+ expected = cast(ProcessGroup, object())
208
+ stub = SimpleNamespace(device_group=expected)
209
+ monkeypatch.setattr(quant_utils, "get_tp_group", lambda: stub)
210
+ monkeypatch.setattr(quant_utils, "get_ep_group", lambda: stub)
211
+
212
+ def fake_all_reduce(tensor, op, group):
213
+ assert op == torch.distributed.ReduceOp.MAX
214
+ assert group is expected
215
+ tensor.copy_(full_amax)
216
+
217
+ monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce)
218
+
219
+
220
+def test_is_tp_sharded_false_when_scale_is_already_global() -> None:
221
+ """Replicated and column-parallel-with-channel-scales need no collective."""
222
+ replicated = SimpleNamespace(
223
+ tp_size=4,
224
+ input_size=64,
225
+ output_size=32,
226
+ input_size_per_partition=64,
227
+ output_size_per_partition=32,
228
+ )
229
+ assert not _is_tp_sharded(replicated)
230
+
231
+ column = SimpleNamespace(
232
+ tp_size=4,
233
+ input_size=64,
234
+ output_size=32,
235
+ input_size_per_partition=64,
236
+ output_size_per_partition=8,
237
+ )
238
+ assert not _is_tp_sharded(column, reduces_output_dim=False)
239
+ assert _is_tp_sharded(column)
240
+
241
+
242
+def _quantize_linear(weight, scheme, is_sharded):
243
+ if scheme == "per_tensor":
244
+ amax = weight_amax(weight).reshape(1)
245
+ scale = _fp8_scale(amax_for_tp_weight_quant(amax, is_sharded))
246
+ return ops.scaled_fp8_quant(weight, scale=scale)[0], scale
247
+ amax = weight_amax(weight, dim=-1, keepdim=True)
248
+ scale = _fp8_channel_scale(amax_for_tp_weight_quant(amax, is_sharded))
249
+ return _fp8_quant_per_channel(weight, scale), scale
250
+
251
+
252
+@pytest.mark.skipif(
253
+ not is_quant_method_supported("fp8"),
254
+ reason="FP8 is not supported on this GPU type.",
255
+)
256
+@pytest.mark.parametrize("scheme", ["per_tensor", "per_channel"])
257
+@pytest.mark.parametrize("shard_dim", [0, 1])
258
+def test_online_linear_tp_weight_quant_matches_unsharded(
259
+ monkeypatch, scheme: str, shard_dim: int
260
+) -> None:
261
+ """TP shards pack the same FP8 values and scales as the unsharded weight."""
262
+ torch.manual_seed(0)
263
+ weight = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
264
+ weight[-1, -1] = 64.0
265
+
266
+ full_weight, full_scale = _quantize_linear(weight, scheme, False)
267
+
268
+ # Per-channel scales reduce only the input dim, so a column (dim 0) shard
269
+ # already matches without a collective.
270
+ is_sharded = scheme == "per_tensor" or shard_dim == 1
271
+ if is_sharded:
272
+ full_amax = (
273
+ weight_amax(weight).reshape(1)
274
+ if scheme == "per_tensor"
275
+ else weight_amax(weight, dim=-1, keepdim=True)
276
+ )
277
+ _patch_max_reduce(monkeypatch, full_amax)
278
+
279
+ shard_size = weight.shape[shard_dim] // 2
280
+ shard = weight.narrow(shard_dim, 0, shard_size).contiguous()
281
+ tp_weight, tp_scale = _quantize_linear(shard, scheme, is_sharded)
282
+
283
+ assert torch.equal(tp_weight, full_weight.narrow(shard_dim, 0, shard_size))
284
+ if scheme == "per_channel" and shard_dim == 0:
285
+ assert torch.equal(tp_scale, full_scale.narrow(0, 0, shard_size))
286
+ else:
287
+ assert torch.equal(tp_scale, full_scale)
288
+
289
+
290
+def _quantize_moe(weight, scheme, moe_tp_size):
291
+ if scheme == "nvfp4":
292
+ return _quantize_moe_weight_to_nvfp4(weight, moe_tp_size)
293
+ if scheme == "per_tensor":
294
+ amax = weight_amax(weight.flatten(1), dim=-1)
295
+ scale = _fp8_scale(amax_for_moe_weight_quant(amax, moe_tp_size))
296
+ quant = lambda w, s: ops.scaled_fp8_quant(w, scale=s)[0] # noqa: E731
297
+ else:
298
+ amax = weight_amax(weight, dim=-1, keepdim=True)
299
+ scale = _fp8_channel_scale(amax_for_moe_weight_quant(amax, moe_tp_size))
300
+ quant = _fp8_quant_per_channel
301
+ qweight = torch.stack([quant(w, s) for w, s in zip(weight, scale)])
302
+ return qweight, scale
303
+
304
+
305
+@pytest.mark.parametrize("scheme", ["per_tensor", "per_channel", "nvfp4"])
306
+def test_online_moe_tp_weight_quant_matches_ep(monkeypatch, scheme: str) -> None:
307
+ """TP shards of w2 pack the same values and scales as full experts."""
308
+ if scheme == "nvfp4":
309
+ if not (
310
+ current_platform.is_cuda()
311
+ and current_platform.is_device_capability_family(100)
312
+ ):
313
+ pytest.skip("NVFP4 weight quantization needs a Blackwell (SM100) GPU.")
314
+ elif not is_quant_method_supported("fp8"):
315
+ pytest.skip("FP8 is not supported on this GPU type.")
316
+
317
+ torch.manual_seed(0)
318
+ weight = torch.randn(2, 32, 32, device="cuda", dtype=torch.bfloat16)
319
+ weight[:, -1, -1] = torch.tensor([32.0, 64.0], device="cuda")
320
+
321
+ ep_out = _quantize_moe(weight, scheme, 1)
322
+
323
+ full_amax = (
324
+ weight_amax(weight, dim=-1, keepdim=True)
325
+ if scheme == "per_channel"
326
+ else weight_amax(weight.flatten(1), dim=-1).to(torch.float32)
327
+ )
328
+ _patch_max_reduce(monkeypatch, full_amax)
329
+
330
+ # w2 is sharded along its last (intermediate) dim.
331
+ shard_size = weight.shape[2] // 2
332
+ tp_out = _quantize_moe(weight[:, :, :shard_size], scheme, 2)
333
+
334
+ packing = 2 if scheme == "nvfp4" else 1
335
+ assert torch.equal(tp_out[0], ep_out[0][:, :, : shard_size // packing])
336
+ if scheme == "nvfp4":
337
+ assert torch.equal(tp_out[1], ep_out[1][:, :, : shard_size // 16])
338
+ assert torch.equal(tp_out[-1], ep_out[-1])
339
+
340
+
341
+def test_online_int8_moe_w2_scale_matches_unsharded(monkeypatch) -> None:
342
+ """Int8 MoE w2 reduces over the sharded intermediate dim."""
343
+ torch.manual_seed(0)
344
+ w13 = torch.randn(2, 16, 8, dtype=torch.bfloat16)
345
+ w2 = torch.randn(2, 8, 16, dtype=torch.bfloat16)
346
+ w2[:, -1, -1] = 64.0
347
+
348
+ def quantize(w2_in, moe_tp_size):
349
+ layer = torch.nn.Module()
350
+ layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False)
351
+ layer.w2_weight = torch.nn.Parameter(w2_in, requires_grad=False)
352
+ layer.num_experts = layer.local_num_experts = w13.shape[0]
353
+ method = SimpleNamespace(moe=SimpleNamespace(tp_size=moe_tp_size))
354
+ Int8OnlineMoEMethod._quantize_weights(method, layer)
355
+ return layer.w2_weight, layer.w2_scale
356
+
357
+ full_weight, full_scale = quantize(w2, 1)
358
+
359
+ _patch_max_reduce(monkeypatch, weight_amax(w2, dim=-1))
360
+
361
+ shard_size = w2.shape[2] // 2
362
+ tp_weight, tp_scale = quantize(w2[:, :, :shard_size].contiguous(), 2)
363
+
364
+ assert torch.equal(tp_weight, full_weight[:, :, :shard_size])
365
+ assert torch.equal(tp_scale, full_scale)
366
+
367
+
368
@pytest.mark.skipif(
369
not (
370
current_platform.is_cuda() and current_platform.is_device_capability_family(100)
vllm/model_executor/layers/quantization/online/fp8.py
+79
-19
@@ -34,13 +34,17 @@ from vllm.model_executor.layers.quantization.online.moe_base import (
34
)
35
from vllm.model_executor.layers.quantization.utils.quant_utils import (
36
GroupShape,
37
+ amax_for_moe_weight_quant,
38
+ amax_for_tp_weight_quant,
39
create_fp8_quant_key,
40
+ get_fp8_min_max,
41
kFp8Dynamic128Sym,
42
kFp8DynamicTensorSym,
43
kFp8DynamicTokenSym,
44
kFp8Static128BlockSym,
45
kFp8StaticChannelSym,
46
kFp8StaticTensorSym,
47
+ weight_amax,
48
)
49
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
50
cutlass_fp8_supported,
@@ -59,6 +63,54 @@ from vllm.utils.math_utils import round_up
63
# ---------------------------------------------------------------------------
64
65
66
+def _fp8_max(device: torch.device) -> torch.Tensor:
67
+ # A 0-d tensor divisor keeps the true division the quant kernels do;
68
+ # dividing by a Python float lowers to multiply-by-reciprocal instead.
69
+ return torch.tensor(get_fp8_min_max()[1], device=device, dtype=torch.float32)
70
+
71
+
72
+def _fp8_scale(amax: torch.Tensor) -> torch.Tensor:
73
+ amax = amax.to(torch.float32)
74
+ return amax / _fp8_max(amax.device)
75
+
76
+
77
+def _fp8_channel_scale(amax: torch.Tensor) -> torch.Tensor:
78
+ fp8_max = _fp8_max(amax.device)
79
+ return _fp8_scale(amax).clamp_min(1.0 / (fp8_max * 512))
80
+
81
+
82
+# Chunk so the fp32 intermediate from the divide stays under ~64 MB, regardless
83
+# of how wide the weight is.
84
+_QUANT_CHUNK_ELEMS = 16 * 1024 * 1024
85
+
86
+
87
+def _fp8_quant_per_channel(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
88
+ # True division, matching the dynamic per-token kernel; the static kernel
89
+ # multiplies by the reciprocal instead and can round a value that sits on a
90
+ # representable midpoint the other way.
91
+ fp8_min, fp8_max = get_fp8_min_max()
92
+ out = torch.empty_like(weight, dtype=current_platform.fp8_dtype())
93
+ chunk_rows = max(1, _QUANT_CHUNK_ELEMS // weight.shape[-1])
94
+ for start in range(0, weight.shape[0], chunk_rows):
95
+ rows = slice(start, start + chunk_rows)
96
+ out[rows] = (weight[rows] / scale[rows]).clamp_(fp8_min, fp8_max).to(out.dtype)
97
+ return out
98
+
99
+
100
+def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool:
101
+ """Whether the weight is sharded along a dim its ``amax`` reduces over.
102
+
103
+ Row parallel shards the input dim, column parallel the output dim.
104
+ Per-output-channel scales only reduce the input dim, so they are already
105
+ unsharded-equivalent on column-parallel layers.
106
+ """
107
+ if layer.tp_size == 1:
108
+ return False
109
+ is_row_parallel = layer.input_size != layer.input_size_per_partition
110
+ is_column_parallel = layer.output_size != layer.output_size_per_partition
111
+ return is_row_parallel or (reduces_output_dim and is_column_parallel)
112
+
113
+
114
class _Fp8OnlineLinearBase(LinearMethodBase):
115
"""Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
116
weights onto meta device and materializes them just-in-time."""
@@ -156,7 +208,10 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
208
return
209
210
layer.input_scale = None
159
- qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
211
+ amax = weight_amax(layer.weight).reshape(1)
212
+ amax = amax_for_tp_weight_quant(amax, _is_tp_sharded(layer))
213
+ weight_scale = _fp8_scale(amax)
214
+ qweight, _ = ops.scaled_fp8_quant(layer.weight, scale=weight_scale)
215
216
# Update layer with new values.
217
replace_parameter(layer, "weight", qweight.t().data)
@@ -335,9 +390,12 @@ class Fp8PtpcOnlineLinearMethod(_Fp8OnlineLinearBase):
390
return
391
392
layer.input_scale = None
338
- qweight, weight_scale = ops.scaled_fp8_quant(
339
- layer.weight, scale=None, use_per_token_if_dynamic=True
393
+ amax = weight_amax(layer.weight, dim=-1, keepdim=True)
394
+ amax = amax_for_tp_weight_quant(
395
+ amax, _is_tp_sharded(layer, reduces_output_dim=False)
396
)
397
+ weight_scale = _fp8_channel_scale(amax)
398
+ qweight = _fp8_quant_per_channel(layer.weight, weight_scale)
399
400
replace_parameter(layer, "weight", qweight.t())
401
replace_parameter(layer, "weight_scale", weight_scale)
@@ -514,19 +572,23 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
572
fp8_dtype = current_platform.fp8_dtype()
573
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
574
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
517
- w13_scale = torch.ones(
518
- layer.num_experts, device=w13.device, dtype=torch.float32
519
- )
520
- w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)
575
layer.w13_input_scale = None
576
layer.w2_input_scale = None
577
578
+ moe_tp_size = self.moe.tp_size
579
+ w13_amax = weight_amax(layer.w13_weight.flatten(1), dim=-1)
580
+ w13_amax = amax_for_moe_weight_quant(w13_amax, moe_tp_size)
581
+ w13_scale = _fp8_scale(w13_amax)
582
+ w2_amax = weight_amax(layer.w2_weight.flatten(1), dim=-1)
583
+ w2_amax = amax_for_moe_weight_quant(w2_amax, moe_tp_size)
584
+ w2_scale = _fp8_scale(w2_amax)
585
+
586
for expert in range(layer.local_num_experts):
525
- w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
526
- layer.w13_weight[expert, :, :]
587
+ w13[expert, :, :], _ = ops.scaled_fp8_quant(
588
+ layer.w13_weight[expert, :, :], scale=w13_scale[expert]
589
)
528
- w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
529
- layer.w2_weight[expert, :, :]
590
+ w2[expert, :, :], _ = ops.scaled_fp8_quant(
591
+ layer.w2_weight[expert, :, :], scale=w2_scale[expert]
592
)
593
594
# Shuffle weights to runtime format and setup kernel.
@@ -723,26 +785,24 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase):
785
# construction, so it cannot drift from the weight's expert count
786
# under EP / padded MoE.
787
n_w13 = layer.w13_weight.shape[1]
726
- n_w2 = layer.w2_weight.shape[1]
788
w13_scale = torch.ones(
789
w13.shape[0], n_w13, 1, device=w13.device, dtype=torch.float32
790
)
730
- w2_scale = torch.ones(
731
- w2.shape[0], n_w2, 1, device=w2.device, dtype=torch.float32
732
- )
791
layer.w13_input_scale = None
792
layer.w2_input_scale = None
793
794
+ w2_amax = weight_amax(layer.w2_weight, dim=-1, keepdim=True)
795
+ w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size)
796
+ w2_scale = _fp8_channel_scale(w2_amax)
797
+
798
for expert in range(layer.local_num_experts):
799
w13[expert], w13_scale[expert] = ops.scaled_fp8_quant(
800
layer.w13_weight[expert],
801
scale=None,
802
use_per_token_if_dynamic=True,
803
)
742
- w2[expert], w2_scale[expert] = ops.scaled_fp8_quant(
743
- layer.w2_weight[expert],
744
- scale=None,
745
- use_per_token_if_dynamic=True,
804
+ w2[expert] = _fp8_quant_per_channel(
805
+ layer.w2_weight[expert], w2_scale[expert]
806
)
807
808
self._setup_kernel(
vllm/model_executor/layers/quantization/online/int8.py
+6
-1
@@ -22,8 +22,10 @@ from vllm.model_executor.layers.quantization.online.moe_base import (
22
OnlineMoEMethodBase,
23
)
24
from vllm.model_executor.layers.quantization.utils.quant_utils import (
25
+ amax_for_moe_weight_quant,
26
kInt8DynamicTokenSym,
27
kInt8StaticChannelSym,
28
+ weight_amax,
29
)
30
from vllm.model_executor.utils import replace_parameter
31
@@ -72,6 +74,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
74
dtype=torch.float32,
75
)
76
77
+ w2_amax = weight_amax(layer.w2_weight, dim=-1)
78
+ w2_amax = amax_for_moe_weight_quant(w2_amax, self.moe.tp_size)
79
+
80
for expert in range(layer.local_num_experts):
81
# w13: per-row quantization over hidden_size dim
82
w = layer.w13_weight[expert, :, :]
@@ -82,7 +87,7 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
87
88
# w2: per-row quantization over intermediate_size dim
89
w = layer.w2_weight[expert, :, :]
85
- scales = w.abs().amax(dim=1) / vmax
90
+ scales = w2_amax[expert] / vmax
91
q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
92
w2[expert, :, :] = q.to(torch.int8)
93
w2_scale[expert, :] = scales
vllm/model_executor/layers/quantization/online/nvfp4.py
+12
-3
@@ -20,8 +20,10 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
20
FLOAT4_E2M1_MAX,
21
)
22
from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
+ amax_for_moe_weight_quant,
24
kNvfp4Dynamic,
25
kNvfp4Static,
26
+ weight_amax,
27
)
28
from vllm.model_executor.utils import replace_parameter
29
from vllm.platforms import current_platform
@@ -31,6 +33,7 @@ FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
33
34
def _quantize_moe_weight_to_nvfp4(
35
weight: torch.Tensor,
36
+ moe_tp_size: int = 1,
37
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
38
"""Quantize stacked MoE expert weights ``(E, N, K)`` to NVFP4.
39
@@ -43,7 +46,8 @@ def _quantize_moe_weight_to_nvfp4(
46
k = weight.shape[-1]
47
assert k % 16 == 0, f"last dim must be a multiple of 16, got {k}"
48
46
- amax = weight.abs().amax(dim=(1, 2)).to(torch.float32).clamp_min(1e-8)
49
+ amax = weight_amax(weight.flatten(1), dim=-1).to(torch.float32)
50
+ amax = amax_for_moe_weight_quant(amax, moe_tp_size).clamp_min(1e-8)
51
global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax
52
weight_scale_2 = (1.0 / global_scale).to(torch.float32)
53
@@ -106,8 +110,13 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
110
layer._already_called_process_weights_after_loading = True
111
112
def _quantize_weights(self, layer: Module) -> None:
109
- w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w13_weight)
110
- w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w2_weight)
113
+ moe_tp_size = self.moe.tp_size
114
+ w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(
115
+ layer.w13_weight, moe_tp_size
116
+ )
117
+ w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(
118
+ layer.w2_weight, moe_tp_size
119
+ )
120
121
replace_parameter(layer, "w13_weight", w13)
122
replace_parameter(layer, "w13_weight_scale", w13_scale)
vllm/model_executor/layers/quantization/utils/quant_utils.py
+38
@@ -11,6 +11,7 @@ import numpy
11
import torch
12
from torch import fx
13
14
+from vllm.distributed.parallel_state import get_ep_group, get_tp_group
15
from vllm.platforms import current_platform
16
from vllm.scalar_type import ScalarType, scalar_types
17
@@ -24,6 +25,43 @@ INT4_DTYPE = scalar_types.uint4b8
25
INT8_DTYPE = scalar_types.uint8b128
26
27
28
+def weight_amax(
29
+ weight: torch.Tensor, *, dim: int | None = None, keepdim: bool = False
30
+) -> torch.Tensor:
31
+ """``max(|weight|)``, without materializing a full-size ``abs()``."""
32
+ lo, hi = weight.aminmax(dim=dim, keepdim=keepdim)
33
+ return torch.maximum(lo.abs(), hi.abs())
34
+
35
+
36
+def amax_for_tp_weight_quant(amax: torch.Tensor, is_sharded: bool) -> torch.Tensor:
37
+ """Reduce a weight ``amax`` over the TP group when the weight is sharded
38
+ along a dim the ``amax`` reduces over, so each shard derives the same scale
39
+ it would as part of the whole weight.
40
+ """
41
+ if is_sharded:
42
+ torch.distributed.all_reduce(
43
+ amax,
44
+ op=torch.distributed.ReduceOp.MAX,
45
+ group=get_tp_group().device_group,
46
+ )
47
+ return amax
48
+
49
+
50
+def amax_for_moe_weight_quant(amax: torch.Tensor, moe_tp_size: int) -> torch.Tensor:
51
+ """Reduce a per-expert weight ``amax`` over the ranks that tensor-shard the
52
+ MoE weights. That sharding is flattened over DP x PCP x TP, exactly the EP
53
+ group's span. Under EP ``moe_tp_size`` is 1 and each rank owns whole
54
+ experts, so no reduction is needed.
55
+ """
56
+ if moe_tp_size > 1:
57
+ torch.distributed.all_reduce(
58
+ amax,
59
+ op=torch.distributed.ReduceOp.MAX,
60
+ group=get_ep_group().device_group,
61
+ )
62
+ return amax
63
+
64
+
65
def get_fp8_min_max() -> tuple[float, float]:
66
"""Get the min and max values for FP8 quantization."""
67
# Using the default value (240.0) from pytorch will cause accuracy