A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 09:45 UTC
[MoE Refactor] Remove supports_expert_map (#43108)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
May 29, 2026 at 17:26 UTC
7b98f498cdf0bf9cf0ecc37b5c0c994cb94513c3
27 files changed
+26
-148
tests/kernels/moe/modular_kernel_tools/common.py
+10
-5
@@ -224,10 +224,6 @@ class Config:
224
info = expert_info(self.fused_experts_type)
225
return info.blocked_quantization_support
226
227
- def supports_expert_map(self):
228
- info = expert_info(self.fused_experts_type)
229
- return info.supports_expert_map
230
-
227
def supports_apply_weight_on_input(self):
228
info = prepare_finalize_info(self.prepare_finalize_type)
229
return info.supports_apply_weight_on_input
@@ -326,6 +322,15 @@ class Config:
322
if self.needs_mori() and not has_mori(): # noqa: SIM103
323
return False, "Needs MoRI, but MoRI not available."
324
325
+ try:
326
+ if not self.fused_experts_type._supports_current_device():
327
+ return (
328
+ False,
329
+ f"{self.fused_experts_type} not supported on the current device.",
330
+ )
331
+ except NotImplementedError:
332
+ pass
333
+
334
return True, None
335
336
@@ -471,7 +476,7 @@ class RankTensors:
476
topk_ids = topk_ids.to(device=device)
477
478
expert_map = None
474
- if config.world_size > 1 and config.supports_expert_map():
479
+ if config.world_size > 1:
480
expert_map = torch.full(
481
(global_num_experts,), fill_value=-1, dtype=torch.int32
482
)
tests/kernels/moe/modular_kernel_tools/mk_objects.py
-15
@@ -67,7 +67,6 @@ class ExpertInfo:
67
activation_format: mk.FusedMoEActivationFormat
68
supported_dtypes: list[torch.dtype | str]
69
blocked_quantization_support: bool
70
- supports_expert_map: bool
70
needs_matching_quant: bool = False
71
needs_deep_gemm: bool = False
72
needs_aiter: bool = False
@@ -129,7 +128,6 @@ def register_experts(
128
activation_format: mk.FusedMoEActivationFormat,
129
supported_dtypes: list[torch.dtype | str],
130
blocked_quantization_support: bool,
132
- supports_expert_map: bool,
131
needs_matching_quant: bool = False,
132
needs_deep_gemm: bool = False,
133
needs_aiter: bool = False,
@@ -142,7 +140,6 @@ def register_experts(
140
activation_format,
141
supported_dtypes,
142
blocked_quantization_support,
145
- supports_expert_map,
143
needs_matching_quant,
144
needs_deep_gemm,
145
needs_aiter,
@@ -176,7 +173,6 @@ register_experts(
173
batched_format,
174
common_float_types,
175
blocked_quantization_support=True,
179
- supports_expert_map=False,
176
needs_matching_quant=True,
177
)
178
@@ -185,7 +181,6 @@ register_experts(
181
standard_format,
182
common_float_and_int_types,
183
blocked_quantization_support=True,
188
- supports_expert_map=True,
184
needs_matching_quant=True,
185
)
186
@@ -194,7 +189,6 @@ register_experts(
189
batched_format,
190
common_float_and_int_types,
191
blocked_quantization_support=True,
197
- supports_expert_map=True,
192
)
193
194
# Disable on blackwell for now
@@ -260,7 +254,6 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
254
nvfp4_types + fp8_types,
255
blocked_quantization_support=True,
256
# Note: this is a hack to get it to run for now
263
- supports_expert_map=True,
257
)
258
else:
259
FlashInferCutlassMoEPrepareAndFinalize = None
@@ -294,7 +287,6 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
287
standard_format,
288
nvfp4_types,
289
blocked_quantization_support=False,
297
- supports_expert_map=True,
290
)
291
292
if has_aiter():
@@ -307,7 +299,6 @@ if has_aiter():
299
standard_format,
300
fp8_types,
301
blocked_quantization_support=True,
310
- supports_expert_map=True,
302
needs_aiter=True,
303
)
304
else:
@@ -319,7 +310,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
310
batched_format,
311
fp8_types,
312
blocked_quantization_support=True,
322
- supports_expert_map=False,
313
needs_matching_quant=False,
314
needs_deep_gemm=True,
315
)
@@ -328,7 +318,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
318
standard_format,
319
fp8_types,
320
blocked_quantization_support=True,
331
- supports_expert_map=True,
321
needs_matching_quant=False,
322
needs_deep_gemm=True,
323
)
@@ -337,7 +326,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
326
standard_format,
327
common_float_and_int_types,
328
blocked_quantization_support=True,
340
- supports_expert_map=True,
329
needs_matching_quant=True,
330
needs_deep_gemm=True,
331
)
@@ -353,14 +341,12 @@ if cutlass_fp8_supported():
341
standard_format,
342
fp8_types,
343
blocked_quantization_support=False,
356
- supports_expert_map=False,
344
)
345
register_experts(
346
CutlassBatchedExpertsFp8,
347
batched_format,
348
fp8_types,
349
blocked_quantization_support=False,
363
- supports_expert_map=False,
350
)
351
else:
352
CutlassBatchedExpertsFp8 = None
@@ -376,7 +362,6 @@ if cutlass_fp4_supported():
362
standard_format,
363
nvfp4_types,
364
blocked_quantization_support=True,
379
- supports_expert_map=False,
365
)
366
else:
367
CutlassExpertsFp4 = None
tests/kernels/moe/test_modular_kernel_combinations.py
+1
-1
@@ -227,7 +227,7 @@ def is_nyi_config(config: Config) -> bool:
227
) == 1
228
return unsupported_quant_config
229
230
- return not info.supports_expert_map
230
+ return False
231
232
233
def generate_valid_test_cases(
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
-3
@@ -248,9 +248,6 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
248
) -> bool:
249
return True
250
251
- def supports_expert_map(self) -> bool:
252
- return False # Expert parallelism not yet supported
253
-
251
@property
252
def expects_unquantized_inputs(self) -> bool:
253
return True
vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
-3
@@ -316,9 +316,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular):
316
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
317
return True
318
319
- def supports_expert_map(self) -> bool:
320
- return False
321
-
319
def supports_packed_ue8m0_act_scales(self) -> bool:
320
"""
321
DeepGemm supports packed ue8m0 activation scales format in devices == sm100
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
-6
@@ -100,9 +100,6 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
100
) -> bool:
101
return True
102
103
- def supports_expert_map(self) -> bool:
104
- return False
105
-
103
def apply(
104
self,
105
hidden_states: torch.Tensor,
@@ -256,9 +253,6 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
253
) -> bool:
254
return True
255
259
- def supports_expert_map(self) -> bool:
260
- return False
261
-
256
def apply(
257
self,
258
hidden_states: torch.Tensor,
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+2
-16
@@ -378,7 +378,8 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
378
topk_ids,
379
activation,
380
global_num_experts,
381
- expert_map,
381
+ # the fp8 cutlass experts use their own expert map.
382
+ None,
383
self.w1_scale,
384
self.w2_scale,
385
a1q_scale,
@@ -418,9 +419,6 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
419
or moe_parallel_config.use_fi_nvl_one_sided_kernels
420
)
421
421
- def supports_expert_map(self) -> bool:
422
- return False
423
-
422
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
423
# topk weights and reduction are fused in moe_unpermute cuda kernel
424
return TopKWeightAndReduceNoOP()
@@ -460,9 +458,6 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
458
def activation_format() -> mk.FusedMoEActivationFormat:
459
return mk.FusedMoEActivationFormat.BatchedExperts
460
463
- def supports_expert_map(self) -> bool:
464
- return False
465
-
461
def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
462
return self.out_dtype if self.out_dtype is not None else act_dtype
463
@@ -741,9 +736,6 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
736
def activation_format() -> mk.FusedMoEActivationFormat:
737
return mk.FusedMoEActivationFormat.Standard
738
744
- def supports_expert_map(self) -> bool:
745
- return False
746
-
739
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
740
return TopKWeightAndReduceNoOP()
741
@@ -1038,9 +1030,6 @@ class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
1030
def activation_format() -> mk.FusedMoEActivationFormat:
1031
return mk.FusedMoEActivationFormat.Standard
1032
1041
- def supports_expert_map(self) -> bool:
1042
- return False
1043
-
1033
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1034
return TopKWeightAndReduceNoOP()
1035
@@ -1340,9 +1329,6 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
1329
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
1330
return True
1331
1343
- def supports_expert_map(self) -> bool:
1344
- return True
1345
-
1332
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1333
# topk weights and reduction are fused in moe_unpermute cuda kernel
1334
return TopKWeightAndReduceNoOP()
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
-6
@@ -164,9 +164,6 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
164
or moe_parallel_config.use_fi_nvl_one_sided_kernels
165
)
166
167
- def supports_expert_map(self) -> bool:
168
- return True
169
-
167
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
168
return TopKWeightAndReduceNoOP()
169
@@ -388,9 +385,6 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
385
or moe_parallel_config.use_fi_nvl_one_sided_kernels
386
)
387
391
- def supports_expert_map(self) -> bool:
392
- return True
393
-
388
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
389
return TopKWeightAndReduceNoOP()
390
vllm/model_executor/layers/fused_moe/experts/fallback.py
-10
@@ -92,16 +92,6 @@ class FallbackExperts(mk.FusedMoEExpertsModular, ABC):
92
moe_parallel_config
93
) and fallback_cls._supports_parallel_config(moe_parallel_config)
94
95
- def supports_expert_map(self) -> bool:
96
- assert (
97
- self.experts.supports_expert_map()
98
- == self.fallback_experts.supports_expert_map()
99
- )
100
- return (
101
- self.experts.supports_expert_map()
102
- and self.fallback_experts.supports_expert_map()
103
- )
104
-
95
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
96
e_war = self.experts.finalize_weight_and_reduce_impl()
97
fbe_war = self.fallback_experts.finalize_weight_and_reduce_impl()
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py
-3
@@ -89,9 +89,6 @@ class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
89
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
90
return True
91
92
- def supports_expert_map(self) -> bool:
93
- return False
94
-
92
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
93
# Let PrepareAndFinalize::finalize() decide the impl.
94
return TopKWeightAndReduceDelegate()
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
-3
@@ -98,9 +98,6 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
98
) -> bool:
99
return True
100
101
- def supports_expert_map(self) -> bool:
102
- return False
103
-
101
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
102
return TopKWeightAndReduceNoOP()
103
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
-3
@@ -207,9 +207,6 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
207
def activation_format() -> mk.FusedMoEActivationFormat:
208
return mk.FusedMoEActivationFormat.Standard
209
210
- def supports_expert_map(self) -> bool:
211
- return False
212
-
210
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
211
return TopKWeightAndReduceNoOP()
212
vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
-6
@@ -555,9 +555,6 @@ class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
555
"This method should not be called."
556
)
557
558
- def supports_expert_map(self) -> bool:
559
- return False
560
-
558
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
559
# Let PrepareAndFinalize::finalize() decide the impl.
560
return TopKWeightAndReduceDelegate()
@@ -799,9 +796,6 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
796
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
797
return True
798
802
- def supports_expert_map(self) -> bool:
803
- return False
804
-
799
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
800
# Let PrepareAndFinalize::finalize() decide the impl.
801
return TopKWeightAndReduceDelegate()
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
-3
@@ -156,9 +156,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
156
) -> bool:
157
return True
158
159
- def supports_expert_map(self) -> bool:
160
- return True
161
-
159
@staticmethod
160
def _supports_current_device() -> bool:
161
platform = current_platform
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
-6
@@ -608,9 +608,6 @@ class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
608
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
609
return True
610
611
- def supports_expert_map(self) -> bool:
612
- return True
613
-
611
def moe_problem_size(
612
self,
613
a1: torch.Tensor,
@@ -1036,9 +1033,6 @@ class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
1033
) -> bool:
1034
return True
1035
1039
- def supports_expert_map(self) -> bool:
1040
- return True
1041
-
1036
@property
1037
def expects_unquantized_inputs(self) -> bool:
1038
return True
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
-6
@@ -686,9 +686,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
686
class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
687
"""Marlin-based fused MoE expert implementation."""
688
689
- def supports_expert_map(self) -> bool:
690
- return True
691
-
689
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
690
return TopKWeightAndReduceNoOP()
691
@@ -920,9 +917,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
917
is_k_full=is_k_full,
918
)
919
923
- def supports_expert_map(self) -> bool:
924
- return True
925
-
920
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
921
return TopKWeightAndReduceDelegate()
922
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
-3
@@ -441,9 +441,6 @@ class AiterExperts(mk.FusedMoEExpertsModular):
441
or moe_parallel_config.use_fi_nvl_one_sided_kernels
442
)
443
444
- def supports_expert_map(self):
445
- return True
446
-
444
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
445
return TopKWeightAndReduceNoOP()
446
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
-3
@@ -128,9 +128,6 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
128
def _supports_batch_invariance():
129
return True
130
131
- def supports_expert_map(self) -> bool:
132
- return True
133
-
131
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
132
return TopKWeightAndReduceNoOP()
133
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
-3
@@ -99,9 +99,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
99
) -> bool:
100
return True
101
102
- def supports_expert_map(self) -> bool:
103
- return False
104
-
102
@property
103
def expects_unquantized_inputs(self) -> bool:
104
return True
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
-3
@@ -88,9 +88,6 @@ class TrtLlmFp8ExpertsBase:
88
or moe_parallel_config.use_ag_rs_all2all_kernels
89
) and not moe_parallel_config.enable_eplb
90
91
- def supports_expert_map(self) -> bool:
92
- return False
93
-
91
92
class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
93
"""
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
-6
@@ -113,9 +113,6 @@ class TrtLlmMxfp4ExpertsBase:
113
def activation_format() -> mk.FusedMoEActivationFormat:
114
return mk.FusedMoEActivationFormat.Standard
115
116
- def supports_expert_map(self) -> bool:
117
- return False
118
-
116
@property
117
def expects_unquantized_inputs(self) -> bool:
118
return False
@@ -248,9 +245,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
245
# routing is done externally, so accept any routing method.
246
return True
247
251
- def supports_expert_map(self) -> bool:
252
- return True
253
-
248
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
249
return TopKWeightAndReduceNoOP()
250
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
-3
@@ -179,9 +179,6 @@ class TrtLlmNvFp4ExpertsBase:
179
300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts)
180
)
181
182
- def supports_expert_map(self) -> bool:
183
- return False
184
-
182
183
class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular):
184
"""
vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
-3
@@ -107,9 +107,6 @@ class XPUExperts(mk.FusedMoEExpertsModular):
107
]
108
return (weight_key, activation_key) in SUPPORTED_W_A
109
110
- def supports_expert_map(self) -> bool:
111
- return True
112
-
110
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
111
return TopKWeightAndReduceNoOP()
112
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+1
-6
@@ -34,11 +34,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
34
super().__init__(moe_kernel.moe_config)
35
self.moe_quant_config = old_quant_method.moe_quant_config
36
self.moe_kernel = moe_kernel
37
- self.disable_expert_map = getattr(
38
- old_quant_method,
39
- "disable_expert_map",
40
- not self.moe_kernel.supports_expert_map(),
41
- )
37
self.old_quant_method = old_quant_method
38
logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)
39
@@ -103,7 +98,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
98
activation=layer.activation,
99
global_num_experts=layer.global_num_experts,
100
apply_router_weight_on_input=layer.apply_router_weight_on_input,
106
- expert_map=None if self.disable_expert_map else layer.expert_map,
101
+ expert_map=layer.expert_map,
102
shared_experts=shared_experts,
103
shared_experts_input=shared_experts_input,
104
)
vllm/model_executor/layers/fused_moe/modular_kernel.py
-13
@@ -751,13 +751,6 @@ class FusedMoEExperts(ABC):
751
"""
752
return False
753
754
- @abstractmethod
755
- def supports_expert_map(self) -> bool:
756
- """
757
- A flag indicating whether or not this class supports expert maps
758
- """
759
- raise NotImplementedError
760
-
754
def supports_packed_ue8m0_act_scales(self) -> bool:
755
"""
756
A flag indicating whether or not this class can process packed ue8m0
@@ -1567,12 +1560,6 @@ class FusedMoEKernel:
1560
== self.fused_experts.activation_format()
1561
)
1562
1570
- def supports_expert_map(self) -> bool:
1571
- """
1572
- A flag indicating whether or not this class supports expert maps.
1573
- """
1574
- return self.fused_experts.supports_expert_map()
1575
-
1563
def output_is_reduced(self) -> bool:
1564
"""
1565
Indicates whether or not the output of fused MoE kernel
vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
+12
-8
@@ -17,7 +17,7 @@ def _quantize_and_setup_dispatch(
17
a1: torch.Tensor,
18
quant_config: FusedMoEQuantConfig,
19
defer_input_quant: bool = False,
20
-) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
20
+) -> tuple[torch.Tensor, list[torch.Tensor] | None, torch.Tensor | None]:
21
# Defer input quantization to the MoE kernel.
22
if defer_input_quant:
23
a1q = a1
@@ -33,7 +33,7 @@ def _quantize_and_setup_dispatch(
33
# which makes the scales tensor different shape than
34
# the hidden states, breaking the A2A kernel. So, we
35
# delay the swizzling until after the A2A.
36
- a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input(
36
+ a1q, a1q_scale = moe_kernel_quantize_input(
37
a1,
38
input_sf,
39
quant_dtype=quant_config.quant_dtype,
@@ -49,7 +49,7 @@ def _quantize_and_setup_dispatch(
49
skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0
50
scales = None if skip_gather_scales else [a1q_scale]
51
52
- return a1q, scales
52
+ return a1q, scales, a1q_scale
53
54
55
def _unwrap_scale_and_prepare_for_moe(
@@ -129,7 +129,9 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
129
)
130
a1 = a1 * topk_weights.to(a1.dtype)
131
132
- a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant)
132
+ a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch(
133
+ a1, quant_config, defer_input_quant
134
+ )
135
136
# When LoRA is active, dispatch the per-token LoRA id along with
137
# hidden_states so every rank receives the correct mapping for the
@@ -164,7 +166,7 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
166
if extra_tensors is None:
167
assert len(res) == 3
168
a1q, topk_weights, topk_ids = res
167
- a1q_scale = None
169
+ a1q_scale = a1q_scale_orig
170
else:
171
assert len(res) == 4
172
a1q, topk_weights, topk_ids, gathered_extras = res
@@ -178,7 +180,7 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular
180
gathered_extras, quant_config
181
)
182
else:
181
- a1q_scale = None
183
+ a1q_scale = a1q_scale_orig
184
185
return a1q, a1q_scale, None, topk_ids, topk_weights
186
@@ -249,7 +251,9 @@ class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMono
251
) -> mk.PrepareMonolithicResultType:
252
"""Quantize and Dispatch Router Logits."""
253
252
- a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant)
254
+ a1q, scales, a1q_scale_orig = _quantize_and_setup_dispatch(
255
+ a1, quant_config, defer_input_quant
256
+ )
257
258
res = get_ep_group().dispatch_router_logits(
259
a1q,
@@ -261,7 +265,7 @@ class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMono
265
if scales is None:
266
assert len(res) == 2
267
a1q, router_logits = res
264
- a1q_scale = None
268
+ a1q_scale = a1q_scale_orig
269
else:
270
assert len(res) == 3
271
a1q, router_logits, scales = res
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
-2
@@ -405,8 +405,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
405
topk_ids,
406
activation=layer.activation,
407
global_num_experts=layer.global_num_experts,
408
- # TODO(rob): investigate the disable_expert_map introduced by:
409
- # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501
408
expert_map=layer.expert_map,
409
apply_router_weight_on_input=layer.apply_router_weight_on_input,
410
shared_experts=shared_experts,