feat: extended EPLB support for Mistral Large 3 and additional MoE backends (#48355)
Signed-off-by: jdebache <jdebache@nvidia.com>
Julien Debache committed
Aug 7, 2026 at 14:32 UTC
ae934ba8a5577c580c33e3489290ff7d8bf1f83e
16 files changed
+500
-96
tests/distributed/test_eplb_quant_scale_consistency.py
new
+299
@@ -0,0 +1,299 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""EPLB rearrangement consistency tests for quant-method derived state.
5
+
6
+EPLB rearranges every registered Parameter of a RoutedExperts in place,
7
+sliced along dim 0 (see RoutedExperts.get_expert_weights). Quant methods
8
+must therefore register all derived per-expert tensors as Parameters and
9
+alias the same storage in their FusedMoEQuantConfig, so the kernels observe
10
+rearranged values with no extra bookkeeping. These tests verify that
11
+contract: simulating a rearrangement on the registered Parameters must be
12
+indistinguishable from loading a checkpoint with permuted experts.
13
+"""
14
+
15
+import pytest
16
+import torch
17
+
18
+from vllm.config import VllmConfig, set_current_vllm_config
19
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
20
+from vllm.model_executor.layers.fused_moe.config import (
21
+ FusedMoEConfig,
22
+ FusedMoEParallelConfig,
23
+ RoutingMethodType,
24
+)
25
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import ( # noqa: E501
26
+ CompressedTensorsW4A4Nvfp4MoEMethod,
27
+)
28
+from vllm.platforms import current_platform
29
+
30
+EPLB_NVFP4_BACKENDS = ["flashinfer_cutedsl", "flashinfer_trtllm"]
31
+
32
+NUM_EXPERTS = 8
33
+HIDDEN_SIZE = 128
34
+INTERMEDIATE_SIZE = 256
35
+EXPERT_PERMUTATION = [3, 0, 5, 1, 7, 2, 6, 4]
36
+
37
+QUANT_CONFIG_TENSORS = (
38
+ "w1_scale",
39
+ "w2_scale",
40
+ "g1_alphas",
41
+ "g2_alphas",
42
+ "a1_gscale",
43
+ "a2_gscale",
44
+)
45
+
46
+
47
+class _RoutedExpertsStub(torch.nn.Module):
48
+ """Minimal ``RoutedExperts`` stand-in for NVFP4 weight processing."""
49
+
50
+ def __init__(self, moe_config: FusedMoEConfig) -> None:
51
+ super().__init__()
52
+ self.moe_config = moe_config
53
+ self.activation = moe_config.activation
54
+
55
+ def _expert_routing_tables(
56
+ self,
57
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
58
+ return None
59
+
60
+
61
+def _make_moe_config(backend: str) -> FusedMoEConfig:
62
+ parallel_config = FusedMoEParallelConfig(
63
+ tp_size=1,
64
+ pcp_size=1,
65
+ dp_size=1,
66
+ ep_size=1,
67
+ tp_rank=0,
68
+ pcp_rank=0,
69
+ dp_rank=0,
70
+ ep_rank=0,
71
+ sp_size=1,
72
+ use_ep=True,
73
+ all2all_backend="allgather_reducescatter",
74
+ enable_eplb=True,
75
+ )
76
+ return FusedMoEConfig(
77
+ num_experts=NUM_EXPERTS,
78
+ experts_per_token=2,
79
+ hidden_dim=HIDDEN_SIZE,
80
+ intermediate_size=INTERMEDIATE_SIZE,
81
+ num_local_experts=NUM_EXPERTS,
82
+ num_logical_experts=NUM_EXPERTS,
83
+ activation=MoEActivation.SILU,
84
+ device="cuda",
85
+ routing_method=RoutingMethodType.TopK,
86
+ moe_parallel_config=parallel_config,
87
+ in_dtype=torch.bfloat16,
88
+ moe_backend=backend,
89
+ )
90
+
91
+
92
+def _make_raw_weights(
93
+ device: torch.device, generator: torch.Generator
94
+) -> dict[str, torch.Tensor]:
95
+ """Random raw NVFP4 checkpoint tensors, indexable per-expert on dim 0."""
96
+ e, h, i = NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE
97
+
98
+ def packed(*shape: int) -> torch.Tensor:
99
+ return torch.randint(
100
+ 0, 256, shape, dtype=torch.uint8, device=device, generator=generator
101
+ )
102
+
103
+ def block_scale(*shape: int) -> torch.Tensor:
104
+ return torch.randint(
105
+ 0, 128, shape, dtype=torch.uint8, device=device, generator=generator
106
+ ).view(torch.float8_e4m3fn)
107
+
108
+ def global_scale(*shape: int) -> torch.Tensor:
109
+ return torch.rand(shape, device=device, generator=generator) + 0.5
110
+
111
+ return {
112
+ "w13_weight_packed": packed(e, 2 * i, h // 2),
113
+ "w2_weight_packed": packed(e, h, i // 2),
114
+ "w13_weight_scale": block_scale(e, 2 * i, h // 16),
115
+ "w2_weight_scale": block_scale(e, h, i // 16),
116
+ "w13_weight_global_scale": global_scale(e, 2),
117
+ "w2_weight_global_scale": global_scale(e),
118
+ "w13_input_global_scale": global_scale(e, 2),
119
+ "w2_input_global_scale": global_scale(e),
120
+ }
121
+
122
+
123
+def _build_processed_layer(
124
+ backend: str, raw: dict[str, torch.Tensor], device: torch.device
125
+) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]:
126
+ """Create the method + layer, load ``raw`` and run the real
127
+ process_weights_after_loading (which builds the flashinfer kernel)."""
128
+ moe_config = _make_moe_config(backend)
129
+ method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False)
130
+ layer = _RoutedExpertsStub(moe_config).to(device)
131
+ method.create_weights(
132
+ layer,
133
+ num_experts=NUM_EXPERTS,
134
+ hidden_size=HIDDEN_SIZE,
135
+ intermediate_size_per_partition=INTERMEDIATE_SIZE,
136
+ params_dtype=torch.bfloat16,
137
+ weight_loader=lambda *args, **kwargs: None,
138
+ )
139
+ layer = layer.to(device)
140
+ with torch.no_grad():
141
+ for name, value in raw.items():
142
+ getattr(layer, name).copy_(value)
143
+ method.process_weights_after_loading(layer)
144
+ return method, layer
145
+
146
+
147
+def _assert_tensors_equal(
148
+ actual: torch.Tensor, expected: torch.Tensor, name: str
149
+) -> None:
150
+ if actual.dtype == torch.float8_e4m3fn:
151
+ assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), name
152
+ elif actual.dtype == torch.uint8:
153
+ assert torch.equal(actual, expected), name
154
+ else:
155
+ torch.testing.assert_close(actual, expected, msg=lambda m: f"{name}: {m}")
156
+
157
+
158
+def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None:
159
+ """Permute experts of every registered Parameter in place, the way
160
+ EPLB's in-place rearrangement moves expert slices along dim 0."""
161
+ with torch.no_grad():
162
+ for _, param in layer.named_parameters():
163
+ param.copy_(param[perm])
164
+
165
+
166
+def _ensure_world1_distributed() -> None:
167
+ """The enable_eplb config makes weight processing all-reduce activation
168
+ scale amaxes over the EP group, so a (single-rank) EP group must exist."""
169
+ from vllm.distributed.parallel_state import (
170
+ ensure_model_parallel_initialized,
171
+ init_distributed_environment,
172
+ )
173
+
174
+ if not torch.distributed.is_initialized():
175
+ init_distributed_environment(
176
+ world_size=1,
177
+ rank=0,
178
+ distributed_init_method="tcp://127.0.0.1:0",
179
+ local_rank=0,
180
+ )
181
+ ensure_model_parallel_initialized(1, 1)
182
+
183
+
184
+@pytest.mark.parametrize("backend", EPLB_NVFP4_BACKENDS)
185
+def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None:
186
+ if not (
187
+ current_platform.is_cuda() and current_platform.is_device_capability_family(100)
188
+ ):
189
+ pytest.skip("NVFP4 CuteDSL/TRTLLM MoE backends require Blackwell (SM100).")
190
+
191
+ device = torch.device("cuda:0")
192
+ torch.accelerator.set_device_index(device)
193
+ perm = torch.tensor(EXPERT_PERMUTATION, device=device)
194
+
195
+ with set_current_vllm_config(VllmConfig()):
196
+ _ensure_world1_distributed()
197
+ generator = torch.Generator(device=device).manual_seed(1234)
198
+ raw = _make_raw_weights(device, generator)
199
+ raw_permuted = {name: value[perm].contiguous() for name, value in raw.items()}
200
+
201
+ try:
202
+ method, layer = _build_processed_layer(backend, raw, device)
203
+ except ValueError as exc:
204
+ pytest.skip(f"{backend} NVFP4 MoE backend unavailable: {exc}")
205
+ ref_method, ref_layer = _build_processed_layer(backend, raw_permuted, device)
206
+
207
+ # The EPLB contract: every registered Parameter must be an
208
+ # expert-major contiguous tensor so get_expert_weights can view it
209
+ # as (E, -1) and rearrange expert slices in place.
210
+ for name, param in layer.named_parameters():
211
+ assert param.is_contiguous(), f"{name} is not contiguous"
212
+ assert param.shape[0] == NUM_EXPERTS, (
213
+ f"{name} is not expert-major: {tuple(param.shape)}"
214
+ )
215
+
216
+ # Derived per-expert scales must live in registered Parameters that
217
+ # the quant config aliases; anything else goes stale on rearrangement.
218
+ quant_config = method.moe_quant_config
219
+ assert quant_config is not None
220
+ params = dict(layer.named_parameters())
221
+ assert quant_config.g1_alphas.data_ptr() == (
222
+ params["w13_weight_scale_2"].data_ptr()
223
+ )
224
+ assert quant_config.g2_alphas.data_ptr() == (
225
+ params["w2_weight_scale_2"].data_ptr()
226
+ )
227
+ assert quant_config.w1_scale.data_ptr() == (
228
+ params["w13_weight_scale"].data_ptr()
229
+ )
230
+ assert quant_config.w2_scale.data_ptr() == params["w2_weight_scale"].data_ptr()
231
+
232
+ _simulate_eplb_rearrangement(layer, perm)
233
+
234
+ # After rearrangement, all registered state must equal what a fresh
235
+ # load of the permuted experts produces.
236
+ ref_params = dict(ref_layer.named_parameters())
237
+ assert params.keys() == ref_params.keys()
238
+ for name in params:
239
+ _assert_tensors_equal(params[name], ref_params[name], name)
240
+
241
+ # And so must the kernel-visible quant config tensors (the CuteDSL
242
+ # MMA scale views alias the registered Parameters' storage).
243
+ ref_quant_config = ref_method.moe_quant_config
244
+ assert ref_quant_config is not None
245
+ for name in QUANT_CONFIG_TENSORS:
246
+ actual = getattr(quant_config, name)
247
+ expected = getattr(ref_quant_config, name)
248
+ assert (actual is None) == (expected is None), name
249
+ if actual is not None:
250
+ _assert_tensors_equal(actual, expected, f"quant_config.{name}")
251
+
252
+
253
+def test_fp8_per_tensor_alphas_registered_and_aliased() -> None:
254
+ """CPU-only: the fp8 per-tensor oracle must register the fused
255
+ (w_scale * a_scale) products as layer Parameters aliased by the quant
256
+ config, so EPLB rearrangement keeps kernels consistent."""
257
+ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
258
+ Fp8MoeBackend,
259
+ make_fp8_moe_quant_config,
260
+ )
261
+
262
+ e = NUM_EXPERTS
263
+ generator = torch.Generator().manual_seed(1234)
264
+
265
+ def scales() -> tuple[torch.Tensor, torch.Tensor]:
266
+ w_scale = torch.rand((e,), generator=generator) + 0.5
267
+ a_scale = torch.rand((), generator=generator) + 0.5
268
+ return w_scale, a_scale.expand(e).contiguous()
269
+
270
+ w1_scale, a1_scale = scales()
271
+ w2_scale, a2_scale = scales()
272
+
273
+ layer = torch.nn.Module()
274
+ quant_config = make_fp8_moe_quant_config(
275
+ fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS,
276
+ w1_scale=w1_scale,
277
+ w2_scale=w2_scale,
278
+ a1_scale=a1_scale,
279
+ a2_scale=a2_scale,
280
+ layer=layer,
281
+ )
282
+
283
+ params = dict(layer.named_parameters())
284
+ assert quant_config.g1_alphas.data_ptr() == params["g1_alphas"].data_ptr()
285
+ assert quant_config.g2_alphas.data_ptr() == params["g2_alphas"].data_ptr()
286
+
287
+ perm = torch.tensor(EXPERT_PERMUTATION)
288
+ ref_config = make_fp8_moe_quant_config(
289
+ fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS,
290
+ w1_scale=w1_scale[perm].contiguous(),
291
+ w2_scale=w2_scale[perm].contiguous(),
292
+ a1_scale=a1_scale,
293
+ a2_scale=a2_scale,
294
+ layer=None,
295
+ )
296
+ _simulate_eplb_rearrangement(layer, perm)
297
+
298
+ torch.testing.assert_close(quant_config.g1_alphas, ref_config.g1_alphas)
299
+ torch.testing.assert_close(quant_config.g2_alphas, ref_config.g2_alphas)
tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py
+16
-15
@@ -193,7 +193,22 @@ def test_flashinfer_cutedsl_fp4_moe(
193
hidden_states, score, topk, renormalize=False
194
)
195
196
- fake_layer = SimpleNamespace(activation=activation)
196
+ activation = MoEActivation.RELU2_NO_MUL
197
+ moe_config = FusedMoEConfig(
198
+ num_experts=e,
199
+ experts_per_token=topk,
200
+ hidden_dim=k,
201
+ intermediate_size=n,
202
+ num_local_experts=e,
203
+ num_logical_experts=e,
204
+ activation=activation,
205
+ device="cuda",
206
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
207
+ in_dtype=dtype,
208
+ routing_method=RoutingMethodType.TopK,
209
+ max_num_tokens=next_power_of_2(m),
210
+ )
211
+ fake_layer = SimpleNamespace(activation=activation, moe_config=moe_config)
212
a1_scale = torch.ones(1, device="cuda", dtype=torch.float32)
213
a2_scale = torch.ones(1, device="cuda", dtype=torch.float32)
214
(
@@ -230,20 +245,6 @@ def test_flashinfer_cutedsl_fp4_moe(
245
gemm1_beta=beta,
246
gemm1_clamp_limit=limit,
247
)
233
- moe_config = FusedMoEConfig(
234
- num_experts=e,
235
- experts_per_token=topk,
236
- hidden_dim=k,
237
- intermediate_size=n,
238
- num_local_experts=e,
239
- num_logical_experts=e,
240
- activation=activation,
241
- device="cuda",
242
- moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
243
- in_dtype=dtype,
244
- routing_method=RoutingMethodType.TopK,
245
- max_num_tokens=next_power_of_2(m),
246
- )
248
249
cutedsl_experts = mk.FusedMoEKernel(
250
maybe_make_prepare_finalize(
tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py
+6
-1
@@ -19,7 +19,12 @@ def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
19
monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
20
21
num_experts = 3
22
- layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False))
22
+ layer = SimpleNamespace(
23
+ activation=SimpleNamespace(is_gated=False),
24
+ moe_config=SimpleNamespace(
25
+ moe_parallel_config=SimpleNamespace(enable_eplb=False)
26
+ ),
27
+ )
28
w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
29
w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
30
w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+6
-4
@@ -108,8 +108,10 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch):
108
)
109
monkeypatch.setattr(
110
eplb,
111
- "is_mixture_of_experts",
112
- lambda loaded_model: getattr(loaded_model, "is_moe", False),
111
+ "get_mixture_of_experts_model",
112
+ lambda loaded_model: (
113
+ loaded_model if getattr(loaded_model, "is_moe", False) else None
114
+ ),
115
)
116
117
runner = _make_runner(is_last_pp_rank=False)
@@ -138,7 +140,7 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch):
140
"init_model_state",
141
lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1),
142
)
141
- monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True)
143
+ monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model)
144
145
runner = _make_runner(is_last_pp_rank=False)
146
mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True)
@@ -153,7 +155,7 @@ def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch):
155
FakeEplbState.instances.clear()
156
FakeEplbState.from_mapping_kwargs = None
157
monkeypatch.setattr(eplb, "EplbState", FakeEplbState)
156
- monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True)
158
+ monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model)
159
160
runner = _make_runner(model=SimpleNamespace(is_moe=True))
161
mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+13
-2
@@ -628,6 +628,17 @@ def make_fp8_moe_quant_config(
628
and block_shape is None
629
):
630
assert a1_scale is not None and a2_scale is not None
631
+ g1_alphas = w1_scale * a1_scale
632
+ g2_alphas = w2_scale * a2_scale
633
+ if layer is not None:
634
+ layer.register_parameter(
635
+ "g1_alphas", torch.nn.Parameter(g1_alphas, requires_grad=False)
636
+ )
637
+ layer.register_parameter(
638
+ "g2_alphas", torch.nn.Parameter(g2_alphas, requires_grad=False)
639
+ )
640
+ g1_alphas = layer.g1_alphas
641
+ g2_alphas = layer.g2_alphas
642
return fp8_w8a8_moe_quant_config(
643
w1_scale=w1_scale,
644
w2_scale=w2_scale,
@@ -637,8 +648,8 @@ def make_fp8_moe_quant_config(
648
a2_scale=a2_scale,
649
a1_gscale=(1.0 / a1_scale),
650
a2_gscale=(1.0 / a2_scale),
640
- g1_alphas=(w1_scale * a1_scale).squeeze(),
641
- g2_alphas=(w2_scale * a2_scale).squeeze(),
651
+ g1_alphas=g1_alphas,
652
+ g2_alphas=g2_alphas,
653
gemm1_clamp_limit=swiglu_limit,
654
)
655
# MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+5
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import (
19
)
20
from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
21
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
22
+ nvfp4_swizzled_scale_to_cutedsl_mma_view,
23
prepare_nvfp4_moe_layer_for_fi_or_cutlass,
24
prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
25
)
@@ -506,6 +507,10 @@ def make_nvfp4_moe_quant_config(
507
gemm1_clamp_limit=swiglu_limit,
508
)
509
510
+ if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
511
+ w13_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w13_scale)
512
+ w2_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w2_scale)
513
+
514
# Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas.
515
# The expert's process_weights_after_loading will fuse activation
516
# scales in-place. Since the quant config references the same tensor
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+16
-2
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEQuantConfig,
16
)
17
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
18
+ NvFp4MoeBackend,
19
convert_to_nvfp4_moe_kernel_format,
20
is_global_sf_supported_for_nvfp4_backend,
21
make_nvfp4_moe_kernel,
@@ -30,6 +31,15 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
)
32
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
33
34
+# NVFP4 backends that have been verified to support EPLB for this quantization recipe.
35
+_EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset(
36
+ {
37
+ NvFp4MoeBackend.FLASHINFER_CUTEDSL,
38
+ NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
39
+ NvFp4MoeBackend.FLASHINFER_TRTLLM,
40
+ }
41
+)
42
+
43
logger = init_logger(__name__)
44
45
@@ -54,6 +64,10 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
64
self.nvfp4_backend
65
)
66
67
+ @property
68
+ def supports_eplb(self) -> bool:
69
+ return self.nvfp4_backend in _EPLB_SUPPORTED_NVFP4_BACKENDS
70
+
71
def create_weights(
72
self,
73
layer: torch.nn.Module,
@@ -223,8 +237,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
237
replace_parameter(layer, "w13_weight_scale", w13_scale)
238
replace_parameter(layer, "w2_weight", w2)
239
replace_parameter(layer, "w2_weight_scale", w2_scale)
226
- layer.w13_weight_scale_2 = w13_scale_2
227
- layer.w2_weight_scale_2 = w2_scale_2
240
+ replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
241
+ replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
242
layer.w13_input_scale = a13_scale
243
layer.w2_input_scale = a2_scale
244
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
+3
-1
@@ -290,7 +290,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
290
assert self.input_quant.strategy == QuantizationStrategy.TENSOR
291
assert w13_input_scale is not None and w2_input_scale is not None
292
w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
293
- w13_input_scale, w2_input_scale
293
+ w13_input_scale,
294
+ w2_input_scale,
295
+ layer.moe_config.moe_parallel_config.enable_eplb,
296
)
297
replace_parameter(layer, "w13_input_scale", w13_input_scale)
298
replace_parameter(layer, "w2_input_scale", w2_input_scale)
vllm/model_executor/layers/quantization/fp8.py
+3
-1
@@ -709,7 +709,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
709
assert not self.block_quant
710
assert w13_input_scale is not None and w2_input_scale is not None
711
w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
712
- w13_input_scale, w2_input_scale
712
+ w13_input_scale,
713
+ w2_input_scale,
714
+ layer.moe_config.moe_parallel_config.enable_eplb,
715
)
716
replace_parameter(layer, "w13_input_scale", w13_input_scale)
717
replace_parameter(layer, "w2_input_scale", w2_input_scale)
vllm/model_executor/layers/quantization/modelopt.py
+3
-1
@@ -899,7 +899,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
899
900
# Per tensor kernels require single activation scale. Use the max.
901
w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
902
- w13_input_scale, w2_input_scale
902
+ w13_input_scale,
903
+ w2_input_scale,
904
+ layer.moe_config.moe_parallel_config.enable_eplb,
905
)
906
replace_parameter(layer, "w13_input_scale", w13_input_scale)
907
replace_parameter(layer, "w2_input_scale", w2_input_scale)
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+40
-27
@@ -15,6 +15,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
15
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
16
swizzle_blockscale,
17
)
18
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
+ amax_for_moe_activation_quant,
20
+)
21
22
if TYPE_CHECKING:
23
from vllm.model_executor.layers.fused_moe import RoutedExperts
@@ -129,12 +132,13 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
132
and interleaves w13 gate/linear rows for gated activations. Non-gated
133
activations use a single w13 projection and keep its row order unchanged.
134
"""
132
- from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
133
-
135
# Global scaling factors (same as other FlashInfer backends).
136
num_experts = w13.shape[0]
136
- a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
137
- a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
137
+ enable_eplb = layer.moe_config.moe_parallel_config.enable_eplb
138
+ a13_scale = amax_for_moe_activation_quant(a13_scale, enable_eplb).repeat(
139
+ num_experts
140
+ )
141
+ a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(num_experts)
142
143
if layer.activation.is_gated:
144
w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl(
@@ -145,29 +149,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
149
w13 = interleave_linear_and_gate(w13, group_size=64, dim=1)
150
w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1)
151
148
- # Convert w13 scale factors: linear → swizzled → MMA layout.
152
w13_scale = swizzle_blockscale(w13_scale)
150
- E, M_padded, K_sf_padded = w13_scale.shape
151
- w13_scale_flat = w13_scale.reshape(E * M_padded, K_sf_padded)
152
- w13_scale = convert_sf_to_mma_layout(
153
- w13_scale_flat,
154
- m=M_padded,
155
- k=K_sf_padded * 16,
156
- num_groups=E,
157
- sf_vec_size=16,
158
- )
159
-
160
- # Convert w2 scale factors: linear → swizzled → MMA layout.
153
w2_scale = swizzle_blockscale(w2_scale)
162
- E, M_padded, K_sf_padded = w2_scale.shape
163
- w2_scale_flat = w2_scale.reshape(E * M_padded, K_sf_padded)
164
- w2_scale = convert_sf_to_mma_layout(
165
- w2_scale_flat,
166
- m=M_padded,
167
- k=K_sf_padded * 16,
168
- num_groups=E,
169
- sf_vec_size=16,
170
- )
154
155
return (
156
w13,
@@ -181,6 +164,31 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
164
)
165
166
167
+def nvfp4_swizzled_scale_to_cutedsl_mma_view(scale: torch.Tensor) -> torch.Tensor:
168
+ """View a swizzled (E, M_padded, K_sf_padded) block-scale tensor in the
169
+ MMA layout expected by the CuteDSL MoE kernel.
170
+
171
+ The returned tensor aliases `scale`'s storage, so in-place updates of the
172
+ registered Parameter (weight reloads, EPLB rearrangement) are visible to
173
+ the kernel with no extra bookkeeping.
174
+ """
175
+ from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
176
+
177
+ num_experts, m_padded, k_sf_padded = scale.shape
178
+ mma_view = convert_sf_to_mma_layout(
179
+ scale.reshape(num_experts * m_padded, k_sf_padded),
180
+ m=m_padded,
181
+ k=k_sf_padded * 16,
182
+ num_groups=num_experts,
183
+ sf_vec_size=16,
184
+ )
185
+ assert mma_view.data_ptr() == scale.data_ptr(), (
186
+ "convert_sf_to_mma_layout no longer returns a view of its input; "
187
+ "the quant config would go stale after weight updates."
188
+ )
189
+ return mma_view
190
+
191
+
192
def prepare_static_weights_for_trtllm_fp4_moe(
193
# args_dequant,
194
# args,
@@ -363,8 +371,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
371
# For some FI kernels, the input scales are shared by all experts.
372
if is_global_sf_supported_for_nvfp4_backend(backend):
373
num_experts = w13.shape[0]
366
- a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
367
- a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
374
+ enable_eplb = layer.moe_config.moe_parallel_config.enable_eplb
375
+ a13_scale = amax_for_moe_activation_quant(a13_scale, enable_eplb).repeat(
376
+ num_experts
377
+ )
378
+ a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(
379
+ num_experts
380
+ )
381
else:
382
a13_scale = a13_scale.max(dim=1).values.to(torch.float32)
383
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+6
-1
@@ -14,6 +14,7 @@ import vllm.envs as envs
14
from vllm import _custom_ops as ops
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
+ amax_for_moe_activation_quant,
18
get_fp8_min_max,
19
)
20
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -1437,6 +1438,7 @@ def process_fp8_weight_tensor_strategy_moe(
1438
def process_fp8_input_tensor_strategy_moe(
1439
w13_input_scale: torch.Tensor,
1440
w2_input_scale: torch.Tensor,
1441
+ enable_eplb: bool,
1442
) -> tuple[torch.Tensor, torch.Tensor]:
1443
"""Process moe input scales for tensor-wise quantization strategy."""
1444
@@ -1447,4 +1449,7 @@ def process_fp8_input_tensor_strategy_moe(
1449
"for each layer."
1450
)
1451
1450
- return w13_input_scale.max(), w2_input_scale.max()
1452
+ return (
1453
+ amax_for_moe_activation_quant(w13_input_scale, enable_eplb),
1454
+ amax_for_moe_activation_quant(w2_input_scale, enable_eplb),
1455
+ )
vllm/model_executor/layers/quantization/utils/quant_utils.py
+22
@@ -62,6 +62,28 @@ def amax_for_moe_weight_quant(amax: torch.Tensor, moe_tp_size: int) -> torch.Ten
62
return amax
63
64
65
+def amax_for_moe_activation_quant(
66
+ a_scale: torch.Tensor, enable_eplb: bool
67
+) -> torch.Tensor:
68
+ """Reduce a per-expert activation scale to one value shared by all experts.
69
+
70
+ Note: when EPLB is enabled and since this quantization scales get
71
+ folded into the per-expert dequantization alphas, we can only
72
+ ensure that the quant/dequant scales match by having a single
73
+ quantization scale shared across all ranks.
74
+ """
75
+ a_max = a_scale.max().to(torch.float32)
76
+ if enable_eplb:
77
+ from vllm.distributed.parallel_state import get_ep_group
78
+
79
+ torch.distributed.all_reduce(
80
+ a_max,
81
+ op=torch.distributed.ReduceOp.MAX,
82
+ group=get_ep_group().device_group,
83
+ )
84
+ return a_max
85
+
86
+
87
def get_fp8_min_max() -> tuple[float, float]:
88
"""Get the min and max values for FP8 quantization."""
89
# Using the default value (240.0) from pytorch will cause accuracy
vllm/model_executor/models/interfaces.py
+32
-2
@@ -5,7 +5,6 @@ import asyncio
5
from collections.abc import (
6
AsyncGenerator,
7
Callable,
8
- Iterable,
8
Mapping,
9
MutableSequence,
10
Sequence,
@@ -949,7 +948,7 @@ class MixtureOfExperts(Protocol):
948
num_redundant_experts: int
949
"""Number of redundant experts in this model."""
950
952
- moe_layers: Iterable["MoERunner"]
951
+ moe_layers: Sequence["MoERunner"]
952
"""List of MoE layers in this model."""
953
954
def set_eplb_state(
@@ -992,6 +991,37 @@ class MixtureOfExperts(Protocol):
991
) -> None: ...
992
993
994
+def get_mixture_of_experts_model(model: object) -> MixtureOfExperts | None:
995
+ """Return the MixtureOfExperts contained within an arbitrary model.
996
+
997
+ - If the model itself is a MixtureOfExperts, return the model directly.
998
+ - If the model is a multi-modal model, and its `language_model` is a
999
+ MixtureOfExperts, return the `language_model`.
1000
+ - If neither, return None.
1001
+
1002
+ Args:
1003
+ model: Model being served.
1004
+
1005
+ Returns:
1006
+ The MixtureOfExperts instance contained within the model, or None.
1007
+ """
1008
+
1009
+ if is_mixture_of_experts(model):
1010
+ return model
1011
+
1012
+ if isinstance(model, SupportsMultiModal):
1013
+ try:
1014
+ mm_language_model = model.get_language_model()
1015
+ return (
1016
+ mm_language_model if is_mixture_of_experts(mm_language_model) else None
1017
+ )
1018
+ except NotImplementedError:
1019
+ logger.info_once("Cannot fetch language_model from MultiModal model")
1020
+ return None
1021
+
1022
+ return None
1023
+
1024
+
1025
def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]:
1026
return (
1027
isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0
vllm/v1/worker/gpu/eplb_utils.py
+13
-21
@@ -12,23 +12,12 @@ from vllm.config import ModelConfig
12
from vllm.distributed.eplb.eplb_state import EplbState
13
from vllm.logger import init_logger
14
from vllm.model_executor.models.interfaces import (
15
- SupportsMultiModal,
16
- is_mixture_of_experts,
15
+ get_mixture_of_experts_model,
16
)
17
18
logger = init_logger(__name__)
19
20
22
-def _unwrap_moe(model: nn.Module) -> nn.Module:
23
- # VLM wrappers (e.g. KimiK25ForConditionalGeneration) hold the MoE
24
- # language model under `.language_model` but don't implement
25
- # MixtureOfExperts themselves. Mirror the V1 path
26
- # (see vllm/v1/worker/gpu_model_runner.py, PR #39805).
27
- if not is_mixture_of_experts(model) and isinstance(model, SupportsMultiModal):
28
- return model.get_language_model()
29
- return model
30
-
31
-
21
def step_eplb_after(*, is_dummy: bool = False) -> Callable:
22
"""Step EPLB after a model runner method completes successfully."""
23
@@ -78,7 +67,8 @@ class EPLBController:
67
return False
68
69
draft_model = speculator.model
81
- if not is_mixture_of_experts(draft_model):
70
+ draft_moe_model = get_mixture_of_experts_model(draft_model)
71
+ if draft_moe_model is None:
72
return False
73
74
assert not self.parallel_config.enable_elastic_ep, (
@@ -88,7 +78,7 @@ class EPLBController:
78
assert speculative_config.draft_model_config is not None
79
assert self.state is not None
80
self.state.add_model(
91
- draft_model,
81
+ draft_moe_model,
82
speculative_config.draft_model_config,
83
)
84
speculator.set_eplb_state(self.state)
@@ -104,13 +94,15 @@ class EPLBController:
94
if not self.parallel_config.enable_eplb or load_dummy_weights:
95
return False
96
107
- model = _unwrap_moe(model)
108
- if not is_mixture_of_experts(model):
97
+ moe_model = get_mixture_of_experts_model(model)
98
+ if moe_model is None:
99
return False
100
111
- logger.info_once("EPLB is enabled for model %s.", model_config.model)
101
+ logger.info_once(
102
+ "EPLB is enabled for MoE part of model %s.", model_config.model
103
+ )
104
assert self.state is not None
113
- self.state.add_model(model, model_config)
105
+ self.state.add_model(moe_model, model_config)
106
self._has_registered_models = True
107
return True
108
@@ -154,11 +146,11 @@ class EPLBController:
146
expanded_physical_to_logical: torch.Tensor,
147
old_num_physical_experts: int,
148
) -> None:
157
- model = _unwrap_moe(model)
158
- assert is_mixture_of_experts(model)
149
+ moe_model = get_mixture_of_experts_model(model)
150
+ assert moe_model is not None
151
152
self.state = EplbState.from_mapping(
161
- model=model,
153
+ model=moe_model,
154
model_config=model_config,
155
device=self.device,
156
parallel_config=self.parallel_config,
vllm/v1/worker/gpu_model_runner.py
+17
-18
@@ -82,7 +82,7 @@ from vllm.model_executor.models.interfaces import (
82
SupportsMRoPE,
83
SupportsMultiModal,
84
SupportsXDRoPE,
85
- is_mixture_of_experts,
85
+ get_mixture_of_experts_model,
86
supports_eagle3,
87
supports_mrope,
88
supports_multimodal_pruning,
@@ -5364,9 +5364,14 @@ class GPUModelRunner(
5364
if hasattr(self.drafter, "load_model"):
5365
self.drafter.load_model(self.model)
5366
if (
5367
- hasattr(self.drafter, "model")
5368
- and is_mixture_of_experts(self.drafter.model)
5369
- and self.parallel_config.enable_eplb
5367
+ self.parallel_config.enable_eplb
5368
+ and hasattr(self.drafter, "model")
5369
+ and (
5370
+ drafter_moe_model := get_mixture_of_experts_model(
5371
+ self.drafter.model
5372
+ )
5373
+ )
5374
+ is not None
5375
):
5376
assert not self.parallel_config.enable_elastic_ep, (
5377
"Elastic EP is not supported with drafter model."
@@ -5375,7 +5380,7 @@ class GPUModelRunner(
5380
assert spec_config is not None
5381
assert spec_config.draft_model_config is not None
5382
logger.info_once(
5378
- "EPLB is enabled for drafter model %s.",
5383
+ "EPLB is enabled for MoE part of drafter model %s.",
5384
spec_config.draft_model_config.model,
5385
)
5386
if self.eplb_state is None:
@@ -5383,7 +5388,7 @@ class GPUModelRunner(
5388
self.parallel_config, self.device
5389
)
5390
self.eplb_state.add_model(
5386
- self.drafter.model,
5391
+ drafter_moe_model,
5392
spec_config.draft_model_config,
5393
)
5394
assert hasattr(self.drafter, "set_eplb_state")
@@ -5396,21 +5401,15 @@ class GPUModelRunner(
5401
# VLM models (e.g. KimiK25ForConditionalGeneration) wrap the
5402
# actual MoE language model but don't implement
5403
# MixtureOfExperts themselves.
5399
- moe_candidate = self.model
5400
- if not is_mixture_of_experts(moe_candidate) and isinstance(
5401
- moe_candidate, SupportsMultiModal
5402
- ):
5403
- moe_candidate = moe_candidate.get_language_model()
5404
- if is_mixture_of_experts(moe_candidate):
5405
- self._moe_model = moe_candidate
5404
+ self._moe_model = get_mixture_of_experts_model(self.model)
5405
5406
if (
5408
- self._moe_model is not None
5409
- and self.parallel_config.enable_eplb
5407
+ self.parallel_config.enable_eplb
5408
and not load_dummy_weights
5409
+ and self._moe_model is not None
5410
):
5411
logger.info_once(
5413
- "EPLB is enabled for model %s.",
5412
+ "EPLB is enabled for MoE part of model %s.",
5413
self.model_config.model,
5414
)
5415
assert self.eplb_state is not None
@@ -5450,8 +5449,8 @@ class GPUModelRunner(
5449
) # Temporary hack for dynamic res video w/o support for bs>1 yet
5450
5451
if (
5453
- self._moe_model is not None
5454
- and self.parallel_config.enable_eplb
5452
+ self.parallel_config.enable_eplb
5453
+ and self._moe_model is not None
5454
and not load_dummy_weights
5455
and self.eplb_state is not None
5456
and self.eplb_state.is_async