[Qwen3.8-Flash-Next] Enable FP8 TP with FlashInfer TRTLLM MoE (#55867)
Signed-off-by: Canlin <canlinguosdu@gmail.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Canlin Guo committed
Sep 17, 2026 at 16:56 UTC
08633cb5cd77ea315351ae531945a9a253016788
5 files changed
+434
-34
tests/kernels/moe/test_fp8_tp_loading.py
new
+311
@@ -0,0 +1,311 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+
5
+import pytest
6
+import torch
7
+
8
+import vllm.distributed as distributed
9
+from tests.kernels.moe.utils import make_dummy_moe_config
10
+from vllm.model_executor.layers.fused_moe import (
11
+ unquantized_fused_moe_method as unquantized_module,
12
+)
13
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
14
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
15
+from vllm.model_executor.layers.quantization import fp8 as fp8_module
16
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
17
+
18
+
19
+def _make_fp8_tp_experts(
20
+ monkeypatch,
21
+ tp_size,
22
+ tp_rank,
23
+ backend="flashinfer_trtllm",
24
+ *,
25
+ mock_backend=True,
26
+ num_experts=2,
27
+ hidden_dim=256,
28
+ intermediate_size=640,
29
+ quant_config=None,
30
+ config_overrides=None,
31
+):
32
+ # Exercise allocation and the public weight loader without selecting a GPU kernel.
33
+ if mock_backend:
34
+ monkeypatch.setattr(
35
+ fp8_module, "select_fp8_moe_backend", lambda **kwargs: (None, None)
36
+ )
37
+ monkeypatch.setattr(
38
+ unquantized_module,
39
+ "select_unquantized_moe_backend",
40
+ lambda **kwargs: (None, None),
41
+ )
42
+ monkeypatch.setattr(
43
+ fp8_module, "get_tensor_model_parallel_world_size", lambda: tp_size
44
+ )
45
+ # The block-shape validator imports this function from vllm.distributed.
46
+ monkeypatch.setattr(
47
+ distributed, "get_tensor_model_parallel_world_size", lambda: tp_size
48
+ )
49
+ config = make_dummy_moe_config(
50
+ num_experts=num_experts,
51
+ hidden_dim=hidden_dim,
52
+ intermediate_size=intermediate_size,
53
+ experts_per_token=min(num_experts, 10),
54
+ )
55
+ config.routing_method = RoutingMethodType.RenormalizeNaive
56
+ config.moe_parallel_config.tp_size = tp_size
57
+ config.moe_parallel_config.tp_rank = tp_rank
58
+ config.intermediate_size_per_partition = intermediate_size // tp_size
59
+ config.intermediate_size_per_partition_unpadded = intermediate_size // tp_size
60
+ config.moe_backend = backend
61
+ for field, value in (config_overrides or {}).items():
62
+ setattr(config, field, value)
63
+ return RoutedExperts(
64
+ "model.layers.0.mlp.experts",
65
+ torch.bfloat16,
66
+ config,
67
+ quant_config
68
+ or Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]),
69
+ expert_map_manager=SimpleNamespace(
70
+ local_num_experts=num_experts,
71
+ placement_strategy="linear",
72
+ expert_map=None,
73
+ expert_mask=None,
74
+ routing_tables=None,
75
+ map_global_to_local=lambda index: index,
76
+ ),
77
+ )
78
+
79
+
80
+@pytest.mark.parametrize("tp_size", [2, 4, 8])
81
+@pytest.mark.parametrize("batched", [False, True])
82
+@pytest.mark.parametrize("intermediate_size", [640, 896])
83
+def test_fp8_block_aligned_tp_preserves_checkpoint(
84
+ monkeypatch, tp_size, batched, intermediate_size
85
+):
86
+ """All ranks reconstruct the original dequantized projections, including
87
+ padding-only ranks. Reloading must clear stale weights and scales.
88
+ """
89
+ generator = torch.Generator().manual_seed(42)
90
+ weights = {
91
+ "w1": torch.randn(2, intermediate_size, 256, generator=generator).to(
92
+ torch.float8_e4m3fn
93
+ ),
94
+ "w3": torch.randn(2, intermediate_size, 256, generator=generator).to(
95
+ torch.float8_e4m3fn
96
+ ),
97
+ "w2": torch.randn(2, 256, intermediate_size, generator=generator).to(
98
+ torch.float8_e4m3fn
99
+ ),
100
+ }
101
+ scales = {
102
+ name: torch.rand(2, w.shape[1] // 128, w.shape[2] // 128, generator=generator)
103
+ + 0.1
104
+ for name, w in weights.items()
105
+ }
106
+ reconstructed: dict[str, list[torch.Tensor]] = {name: [] for name in weights}
107
+ num_blocks = intermediate_size // 128
108
+ for rank in range(tp_size):
109
+ layer = _make_fp8_tp_experts(
110
+ monkeypatch, tp_size, rank, intermediate_size=intermediate_size
111
+ )
112
+ assert layer.quant_method.weight_scale_refine is None
113
+ assert layer.quant_method.moe_block_shape == [128, 128]
114
+ width = layer.moe_config.intermediate_size_per_partition
115
+ assert width == ((num_blocks + tp_size - 1) // tp_size) * 128
116
+ for name in weights:
117
+ prefix = "w2" if name == "w2" else "w13"
118
+ for suffix, checkpoint in (
119
+ ("weight", weights[name]),
120
+ ("weight_scale_inv", scales[name]),
121
+ ):
122
+ param_name = f"{prefix}_{suffix}"
123
+ param = getattr(layer, param_name)
124
+ # Also reload after dirtying the destination, to catch padding leaks.
125
+ for _ in range(2):
126
+ if name != "w3":
127
+ param.data.fill_(7)
128
+ if batched:
129
+ assert param.weight_loader(
130
+ param, checkpoint, param_name, name, 0, return_success=True
131
+ )
132
+ else:
133
+ for expert in range(2):
134
+ assert param.weight_loader(
135
+ param,
136
+ checkpoint[expert],
137
+ param_name,
138
+ name,
139
+ expert,
140
+ return_success=True,
141
+ )
142
+ w = getattr(layer, f"{prefix}_weight").float()
143
+ s = getattr(layer, f"{prefix}_weight_scale_inv")
144
+ if name in ("w1", "w3"):
145
+ half = 0 if name == "w1" else 1
146
+ w = w.chunk(2, dim=1)[half]
147
+ s = s.chunk(2, dim=1)[half]
148
+ dim = 2 if name == "w2" else 1
149
+ valid = max(0, min(width, intermediate_size - rank * width))
150
+ assert torch.count_nonzero(w.narrow(dim, valid, width - valid)) == 0
151
+ assert torch.all(s.narrow(dim, valid // 128, (width - valid) // 128) == 1)
152
+ dequant = w * s.repeat_interleave(128, dim=1).repeat_interleave(128, dim=2)
153
+ reconstructed[name].append(dequant.narrow(dim, 0, valid))
154
+ for name, parts in reconstructed.items():
155
+ expected = weights[name].float() * scales[name].repeat_interleave(
156
+ 128, dim=1
157
+ ).repeat_interleave(128, dim=2)
158
+ assert torch.equal(torch.cat(parts, dim=2 if name == "w2" else 1), expected)
159
+
160
+
161
+@pytest.mark.parametrize("backend", ["auto", "triton"])
162
+@pytest.mark.parametrize("tp_size,block", [(2, 64), (4, 32)])
163
+def test_fp8_tp_default_keeps_refined_layout(monkeypatch, backend, tp_size, block):
164
+ layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0, backend)
165
+ assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size
166
+ assert layer.quant_method.moe_block_shape == [block, block]
167
+
168
+
169
+@pytest.mark.parametrize(
170
+ "field,value",
171
+ [
172
+ ("hidden_dim", 192),
173
+ ("intermediate_size", 672),
174
+ ("is_lora_enabled", True),
175
+ ("has_bias", True),
176
+ ],
177
+)
178
+def test_fp8_block_aligned_tp_rejects_unsupported_layout(monkeypatch, field, value):
179
+ with pytest.raises(ValueError, match="Block-aligned FP8 TP sharding requires"):
180
+ _make_fp8_tp_experts(monkeypatch, 4, 0, config_overrides={field: value})
181
+
182
+
183
+@pytest.mark.parametrize("tp_size", [1, 5])
184
+def test_fp8_aligned_tp_keeps_original_layout(monkeypatch, tp_size):
185
+ layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0)
186
+ assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size
187
+ assert not layer.moe_config.tp_shard_with_padding
188
+ assert layer.quant_method.weight_scale_refine is None
189
+
190
+
191
+def test_fp8_skipped_layer_keeps_original_tp_layout(monkeypatch, default_vllm_config):
192
+ layer = _make_fp8_tp_experts(
193
+ monkeypatch,
194
+ 4,
195
+ 0,
196
+ quant_config=Fp8Config(
197
+ is_checkpoint_fp8_serialized=True,
198
+ weight_block_size=[128, 128],
199
+ ignored_layers=["model.layers.0.mlp.experts"],
200
+ ),
201
+ )
202
+ assert isinstance(layer.quant_method, unquantized_module.UnquantizedFusedMoEMethod)
203
+ assert not layer.moe_config.tp_shard_with_padding
204
+ assert layer.moe_config.intermediate_size_per_partition == 160
205
+ assert layer.w13_weight.shape == (2, 320, 256)
206
+ assert layer.w2_weight.shape == (2, 256, 160)
207
+
208
+
209
+def test_fp8_block_aligned_tp_rejects_presharded_weights(monkeypatch):
210
+ layer = _make_fp8_tp_experts(monkeypatch, 4, 0)
211
+ with pytest.raises(ValueError, match="unsharded checkpoint"):
212
+ layer.w13_weight.weight_loader(
213
+ layer.w13_weight,
214
+ torch.zeros(160, 256, dtype=torch.float8_e4m3fn),
215
+ "w13_weight",
216
+ "w1",
217
+ 0,
218
+ )
219
+
220
+
221
+@pytest.mark.parametrize("tp_size", [2, 4, 8])
222
+@pytest.mark.parametrize("num_tokens", [1, 17])
223
+@torch.inference_mode()
224
+def test_fp8_block_aligned_tp_flashinfer_matches_unsharded(
225
+ monkeypatch,
226
+ tp_size,
227
+ num_tokens,
228
+ dist_init,
229
+ workspace_init,
230
+):
231
+ """Compare sharded FI with unsharded FI, reporting the cross-backend
232
+ difference separately because activation quantization can differ.
233
+ """
234
+ from vllm.config import VllmConfig, set_current_vllm_config
235
+ from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
236
+ from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
237
+ from vllm.platforms import current_platform
238
+
239
+ if not current_platform.is_device_capability_family(100):
240
+ pytest.skip("Requires datacenter Blackwell and FlashInfer TRTLLM")
241
+ torch.manual_seed(42)
242
+ e, h, n, topk = 16, 2560, 640, 10
243
+ with torch.device("cuda"), set_current_vllm_config(VllmConfig()):
244
+ x = torch.randn(num_tokens, h, dtype=torch.bfloat16) / 10
245
+ logits = torch.randn(num_tokens, e, dtype=torch.float32)
246
+ probabilities = logits.softmax(dim=-1)
247
+ topk_weights, topk_ids = probabilities.topk(topk, dim=-1)
248
+ topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
249
+ weights = {
250
+ "w1": torch.randn(e, n, h).to(torch.float8_e4m3fn),
251
+ "w3": torch.randn(e, n, h).to(torch.float8_e4m3fn),
252
+ "w2": torch.randn(e, h, n).to(torch.float8_e4m3fn),
253
+ }
254
+ scales = {
255
+ name: (torch.rand(e, w.shape[1] // 128, w.shape[2] // 128) + 0.5) / 32
256
+ for name, w in weights.items()
257
+ }
258
+ ref = fused_experts(
259
+ x,
260
+ torch.cat([weights["w1"], weights["w3"]], dim=1),
261
+ weights["w2"],
262
+ topk_weights,
263
+ topk_ids,
264
+ quant_config=fp8_w8a8_moe_quant_config(
265
+ w1_scale=torch.cat([scales["w1"], scales["w3"]], dim=1),
266
+ w2_scale=scales["w2"],
267
+ block_shape=[128, 128],
268
+ ),
269
+ )
270
+ total = torch.zeros_like(ref, dtype=torch.float32)
271
+ fi_reference = None
272
+ for rank in range(-1, tp_size):
273
+ layer = _make_fp8_tp_experts(
274
+ monkeypatch,
275
+ 1 if rank == -1 else tp_size,
276
+ max(rank, 0),
277
+ mock_backend=False,
278
+ num_experts=e,
279
+ hidden_dim=h,
280
+ )
281
+ for name in weights:
282
+ prefix = "w2" if name == "w2" else "w13"
283
+ for suffix, checkpoint in (
284
+ ("weight", weights[name]),
285
+ ("weight_scale_inv", scales[name]),
286
+ ):
287
+ param_name = f"{prefix}_{suffix}"
288
+ param = getattr(layer, param_name)
289
+ param.weight_loader(param, checkpoint, param_name, name, 0)
290
+ method = layer.quant_method
291
+ method.process_weights_after_loading(layer)
292
+ assert method.is_monolithic
293
+ out = method.apply_monolithic(layer, x, logits)
294
+ assert torch.isfinite(out).all()
295
+ if rank == -1:
296
+ fi_reference = out.float().clone()
297
+ else:
298
+ total += out.float()
299
+ # Different FP8 GEMM reduction orders need not be bitwise identical.
300
+ assert fi_reference is not None
301
+ relative_l2 = (total - fi_reference).norm() / fi_reference.norm().clamp_min(
302
+ 1e-8
303
+ )
304
+ cross_backend_l2 = (
305
+ fi_reference - ref.float()
306
+ ).norm() / ref.float().norm().clamp_min(1e-8)
307
+ print(
308
+ f"TP{tp_size} M={num_tokens}: shard_l2={relative_l2.item():.6f}, "
309
+ f"unsharded_fi_vs_triton_l2={cross_backend_l2.item():.6f}"
310
+ )
311
+ assert relative_l2 < 0.01, relative_l2.item()
vllm/model_executor/layers/fused_moe/config.py
+2
@@ -1309,6 +1309,8 @@ class FusedMoEConfig:
1309
1310
# Set by __post_init__
1311
intermediate_size_per_partition: int = -1
1312
+ # Use the allocated width as the checkpoint TP stride, including for scales.
1313
+ tp_shard_with_padding: bool = False
1314
rocm_aiter_fmoe_enabled: bool = False
1315
aiter_fmoe_shared_expert_enabled: bool = False
1316
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+77
@@ -31,11 +31,14 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
31
prepare_fp8_moe_layer_for_marlin,
32
)
33
from vllm.model_executor.layers.quantization.utils.quant_utils import (
34
+ GroupShape,
35
QuantKey,
36
+ create_fp8_quant_key,
37
kFp8Dynamic128Sym,
38
kFp8Static128BlockSym,
39
)
40
from vllm.platforms import current_platform
41
+from vllm.utils.math_utils import round_up
42
43
logger = init_logger(__name__)
44
@@ -299,6 +302,80 @@ def refine_fp8_moe_block_shape(
302
return [refine, refine]
303
304
305
+def pad_tp_shard_to_weight_blocks(
306
+ config: FusedMoEConfig,
307
+ weight_block_size: list[int],
308
+) -> bool:
309
+ """Pad the TP shard to whole checkpoint blocks, keeping scales rank-local."""
310
+ block_n, block_k = weight_block_size
311
+ if (
312
+ block_n != block_k
313
+ or config.tp_size == 1
314
+ or config.intermediate_size_per_partition % block_n == 0
315
+ ):
316
+ return False
317
+ if (
318
+ config.intermediate_size % block_n != 0
319
+ or config.hidden_dim % block_n != 0
320
+ or config.ep_size != 1
321
+ or config.is_lora_enabled
322
+ or config.has_bias
323
+ ):
324
+ raise ValueError(
325
+ f"Block-aligned FP8 TP sharding requires {block_n}-aligned "
326
+ "global expert dimensions, pure TP, and no LoRA or "
327
+ "expert bias."
328
+ )
329
+ config.tp_shard_with_padding = True
330
+ config.intermediate_size_per_partition = round_up(
331
+ config.intermediate_size_per_partition, block_n
332
+ )
333
+ return True
334
+
335
+
336
+def resolve_fp8_moe_weight_block_shape(
337
+ config: FusedMoEConfig,
338
+ weight_block_size: list[int],
339
+ activation_key: QuantKey,
340
+ is_checkpoint_fp8_serialized: bool,
341
+) -> tuple[list[int], tuple[int, int] | None]:
342
+ """Return the TP-adapted block shape and refine factor:
343
+ refine if kernels allow, else pad to the TP shard."""
344
+ refined_shape = refine_fp8_moe_block_shape(config, weight_block_size)
345
+ if is_checkpoint_fp8_serialized and config.moe_backend != "auto":
346
+ kernel_classes = backend_to_kernel_cls(map_fp8_backend(config.moe_backend))
347
+ can_refine = refined_shape is not None and any(
348
+ k_cls._supports_quant_scheme(
349
+ create_fp8_quant_key(
350
+ static=True, group_shape=GroupShape(*refined_shape)
351
+ ),
352
+ activation_key,
353
+ )
354
+ for k_cls in kernel_classes
355
+ )
356
+ if not can_refine and pad_tp_shard_to_weight_blocks(config, weight_block_size):
357
+ logger.info_once(
358
+ "FP8 %s TP loading uses complete checkpoint blocks: "
359
+ "local allocation %d, without weight requantization.",
360
+ config.moe_backend,
361
+ config.intermediate_size_per_partition,
362
+ )
363
+ return weight_block_size, None
364
+ if refined_shape is None:
365
+ return weight_block_size, None
366
+ logger.info_once(
367
+ "FP8 MoE block scales refined from %s to %s to fit "
368
+ "the TP-sharded intermediate size %d.",
369
+ str(weight_block_size),
370
+ str(refined_shape),
371
+ config.intermediate_size_per_partition,
372
+ )
373
+ return refined_shape, (
374
+ weight_block_size[0] // refined_shape[0],
375
+ weight_block_size[1] // refined_shape[1],
376
+ )
377
+
378
+
379
def select_fp8_moe_backend(
380
config: FusedMoEConfig,
381
weight_key: QuantKey | None,
vllm/model_executor/layers/fused_moe/routed_experts.py
+33
-2
@@ -332,6 +332,7 @@ class RoutedExperts(PluggableLayer):
332
shard_id: str,
333
loaded_weight: torch.Tensor,
334
tp_rank: int,
335
+ is_scale: bool = False,
336
):
337
"""Load grouped weight scales for group quantization or model weights
338
@@ -341,14 +342,41 @@ class RoutedExperts(PluggableLayer):
342
shard_id: either w1, w2, or w3
343
loaded_weight: checkpoint weight to load into the param
344
tp_rank: tensor parallel rank
345
+ is_scale: whether padding should use unit scales instead of zero weights.
346
347
"""
348
+ padded_tp = self.moe_config.tp_shard_with_padding
349
+ if padded_tp:
350
+ destination = expert_data
351
+ if shard_id in ("w1", "w3") and self.moe_config.is_act_and_mul:
352
+ half = destination.shape[shard_dim] // 2
353
+ destination = destination.narrow(
354
+ shard_dim, 0 if shard_id == "w1" else half, half
355
+ )
356
+ shard_size = destination.shape[shard_dim]
357
+ expected_size = (
358
+ self.moe_config.intermediate_size
359
+ * shard_size
360
+ // self.moe_config.intermediate_size_per_partition
361
+ )
362
+ if loaded_weight.shape[shard_dim] != expected_size:
363
+ raise ValueError(
364
+ "Block-aligned TP loading expects an unsharded checkpoint "
365
+ f"projection of size {expected_size}, got "
366
+ f"{loaded_weight.shape[shard_dim]}."
367
+ )
368
+ start = min(tp_rank * shard_size, expected_size)
369
+ size = min(shard_size, expected_size - start)
370
+ loaded_weight = loaded_weight.narrow(shard_dim, start, size)
371
+ destination.fill_(1 if is_scale else 0)
372
+
373
if shard_id == "w2":
374
self._load_w2(
375
shard_dim=shard_dim,
376
loaded_weight=loaded_weight,
377
expert_data=expert_data,
378
tp_rank=tp_rank,
379
+ load_full=padded_tp,
380
)
381
elif shard_id in ("w1", "w3"):
382
self._load_w13(
@@ -357,6 +385,7 @@ class RoutedExperts(PluggableLayer):
385
loaded_weight=loaded_weight,
386
expert_data=expert_data,
387
tp_rank=tp_rank,
388
+ load_full=padded_tp,
389
)
390
391
def _load_per_channel_weight_scale(
@@ -511,11 +540,12 @@ class RoutedExperts(PluggableLayer):
540
shard_dim: int,
541
loaded_weight: torch.Tensor,
542
tp_rank: int,
543
+ load_full: bool = False,
544
):
545
# Index the loaded weight for tp sharding.
546
# down_proj: "RowParallel" so tp sharding on input_dim
517
- # Only narrow if the loaded_weight is not a scalar (0-dim tensor).
518
- if loaded_weight.ndim > 0:
547
+ # Padded TP weights have already been sliced by the grouped loader.
548
+ if not load_full and loaded_weight.ndim > 0:
549
# Same padding fix as _load_w13: use unpadded per-rank size.
550
tp_size = self.moe_config.moe_parallel_config.tp_size
551
loaded_per_rank = loaded_weight.shape[shard_dim] // tp_size
@@ -800,6 +830,7 @@ class RoutedExperts(PluggableLayer):
830
loaded_weight=loaded_weight,
831
expert_data=expert_data,
832
tp_rank=self.moe_config.tp_rank,
833
+ is_scale=True,
834
)
835
elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value:
836
self._load_per_tensor_weight_scale(
vllm/model_executor/layers/quantization/fp8.py
+11
-32
@@ -8,7 +8,6 @@ import torch
8
import vllm.envs as envs
9
from vllm.config import get_current_vllm_config
10
from vllm.distributed import get_tensor_model_parallel_world_size
11
-from vllm.logger import init_logger
11
from vllm.model_executor.kernels.linear import (
12
init_fp8_linear_kernel,
13
)
@@ -24,14 +23,12 @@ from vllm.model_executor.layers.fused_moe import (
23
SharedExperts,
24
UnquantizedFusedMoEMethod,
25
)
27
-from vllm.model_executor.layers.fused_moe.config import (
28
- FusedMoEQuantConfig,
29
-)
26
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
27
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
28
convert_to_fp8_moe_kernel_format,
29
make_fp8_moe_kernel,
30
make_fp8_moe_quant_config,
34
- refine_fp8_moe_block_shape,
31
+ resolve_fp8_moe_weight_block_shape,
32
select_fp8_moe_backend,
33
)
34
from vllm.model_executor.layers.linear import (
@@ -90,8 +87,6 @@ if TYPE_CHECKING:
87
88
ACTIVATION_SCHEMES = ["static", "dynamic"]
89
93
-logger = init_logger(__name__)
94
-
90
91
class Fp8Config(QuantizationConfig):
92
"""Config class for FP8."""
@@ -511,32 +506,16 @@ class Fp8MoEMethod(FusedMoEMethodBase):
506
# Set weight key and activation key for kernel compatibility
507
if self.block_quant:
508
assert self.weight_block_size is not None
514
- # TP shards the intermediate dim of the expert weights, so a
515
- # per-shard size that is not a multiple of the checkpoint's block
516
- # size makes the checkpoint's block scales impossible to shard
517
- # exactly. When a finer block size (>= 32) divides both the
518
- # checkpoint blocks and all involved dims, the weight scales are
519
- # refined to that granularity at load time (a lossless upsampling,
520
- # since the refined block divides the checkpoint block). The
521
- # refined block shape is encoded in the weight key, so the oracle
522
- # only selects kernels that support it (e.g. Triton, which takes
523
- # the block shape as a runtime argument).
524
- refined_shape = refine_fp8_moe_block_shape(self.moe, self.weight_block_size)
525
- if refined_shape is not None:
526
- block_n, block_k = self.weight_block_size
527
- self.weight_scale_refine = (
528
- block_n // refined_shape[0],
529
- block_k // refined_shape[1],
530
- )
531
- self.moe_block_shape = refined_shape
532
- logger.info_once(
533
- "FP8 MoE block scales refined from %s to %s to fit "
534
- "the TP-sharded intermediate size %d.",
535
- str(self.weight_block_size),
536
- str(refined_shape),
537
- self.moe.intermediate_size_per_partition,
509
+ # Adapt the checkpoint block shape to TP sharding before
510
+ # building the weight key.
511
+ self.moe_block_shape, self.weight_scale_refine = (
512
+ resolve_fp8_moe_weight_block_shape(
513
+ self.moe,
514
+ self.weight_block_size,
515
+ kFp8Dynamic128Sym,
516
+ self.quant_config.is_checkpoint_fp8_serialized,
517
)
539
- assert self.moe_block_shape is not None
518
+ )
519
weight_key = create_fp8_quant_key(
520
static=True, group_shape=GroupShape(*self.moe_block_shape)
521
)