@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d9b4990783

Fork 0

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 3, 19:00 UTC

[MoE Refactor] EPLB refactoring for FusedMoE (#41055)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed May 12, 2026 at 14:16 UTC d9b49907831f878d326ef8b0b854aa007f0ea510

17 files changed +83 -98

tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py

+2 -1

@@ -10,6 +10,7 @@ import torch

11 from tests.kernels.moe.utils import make_test_quant_config

12 from vllm.config import VllmConfig, set_current_vllm_config

13 +from vllm.distributed.eplb.eplb_state import EplbLayerState

14 from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace

15 from vllm.distributed.parallel_state import (

16 ensure_model_parallel_initialized,

@@ -201,7 +202,7 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig):

202 dtype=torch.int32,

203 device=device,

204 )

204 - fml.enable_eplb = True

205 + fml.eplb_state = EplbLayerState()

206 fml.set_eplb_state(

207 lidx,

208 torch.zeros(

tests/kernels/moe/test_moe_layer.py

-3

@@ -1004,7 +1004,6 @@ def make_fake_moe_layer(

1004 activation: str = "silu",

1005 indices_type: torch.dtype | None = None,

1006 expert_map: torch.Tensor | None = None,

1007 - enable_eplb: bool = False,

1007 expert_load_view: torch.Tensor | None = None,

1008 logical_to_physical_map: torch.Tensor | None = None,

1009 logical_replica_count: torch.Tensor | None = None,

@@ -1022,7 +1021,6 @@ def make_fake_moe_layer(

1021 router = create_fused_moe_router(

1022 top_k=top_k,

1023 global_num_experts=global_num_experts,

1025 - # eplb_state=None, # TODO

1024 renormalize=renormalize,

1025 use_grouped_topk=use_grouped_topk,

1026 num_expert_group=num_expert_group,

@@ -1032,7 +1030,6 @@ def make_fake_moe_layer(

1030 routed_scaling_factor=routed_scaling_factor,

1031 e_score_correction_bias=e_score_correction_bias,

1032 num_fused_shared_experts=0, # TODO

1035 - enable_eplb=enable_eplb,

1033 # TODO(bnell): once we can construct the MK at init time, we

1034 # can make this a value.

1035 indices_type_getter=lambda: indices_type,

tests/kernels/moe/test_routing.py

+4 -7

@@ -36,9 +36,11 @@ TOP_KS = [2, 4, 6]

36 NUM_EXPERTS = [8, 16, 64]

39 -def setup_eplb_state(enable_eplb: bool, global_num_experts: int) -> EplbLayerState:

39 +def setup_eplb_state(

40 + enable_eplb: bool, global_num_experts: int

41 +) -> EplbLayerState | None:

42 if not enable_eplb:

41 - return EplbLayerState()

43 + return None

45 # Initialize EPLB state with proper tensors for testing

46 # For testing purposes, we use a simple 1:1 mapping (no redundant experts)

@@ -349,7 +351,6 @@ def test_fused_topk(

351 top_k=top_k,

352 global_num_experts=global_num_experts,

353 renormalize=renormalize,

352 - enable_eplb=enable_eplb,

354 eplb_state=eplb_state,

355 )

356

@@ -400,7 +401,6 @@ def test_fused_topk_bias(

401 top_k=top_k,

402 global_num_experts=global_num_experts,

403 renormalize=renormalize,

403 - enable_eplb=enable_eplb,

404 eplb_state=eplb_state,

405 )

406

@@ -469,7 +469,6 @@ def test_grouped_topk(

469 top_k=top_k,

470 global_num_experts=global_num_experts,

471 renormalize=renormalize,

472 - enable_eplb=enable_eplb,

472 eplb_state=eplb_state,

473 )

474

@@ -540,7 +539,6 @@ def test_custom(

539 global_num_experts=global_num_experts,

540 custom_routing_function=custom_routing_function,

541 renormalize=renormalize,

543 - enable_eplb=enable_eplb,

542 eplb_state=eplb_state,

543 )

544

@@ -580,7 +578,6 @@ def test_custom(

578 # router = create_fused_moe_router(

579 # top_k=top_k,

580 # global_num_experts=global_num_experts,

583 -# enable_eplb=enable_eplb,

581 # eplb_state=eplb_state,

582 # )

583

vllm/distributed/eplb/eplb_state.py

+11

@@ -940,6 +940,17 @@ class EplbLayerState:

940 GPU work.

941 """

942

943 + def set_layer_state(

944 + self,

945 + moe_layer_idx: int,

946 + expert_load_view: torch.Tensor,

947 + logical_to_physical_map: torch.Tensor,

948 + logical_replica_count: torch.Tensor,

949 + ) -> None:

950 + self.expert_load_view = expert_load_view[moe_layer_idx]

951 + self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]

952 + self.logical_replica_count = logical_replica_count[moe_layer_idx]

953 +

954

955 def _node_count_with_rank_mapping(

956 pg: ProcessGroup | StatelessProcessGroup,

vllm/model_executor/layers/fused_moe/layer.py

+29 -20

@@ -203,13 +203,24 @@ class FusedMoE(PluggableLayer):

203 compilation_config.static_all_moe_layers.append(prefix)

204 self.layer_name = prefix

205

206 - self.enable_eplb = enable_eplb

207 - # TODO(bnell): should this be owned by router?

208 - self.eplb_state = EplbLayerState()

206 self.expert_placement_strategy: ExpertPlacementStrategy = (

207 vllm_config.parallel_config.expert_placement_strategy

208 )

209

210 + self.eplb_state: EplbLayerState | None = None

211 + if enable_eplb:

212 + if self.use_ep and self.global_num_experts % self.ep_size != 0:

213 + raise ValueError(

214 + f"EPLB currently only supports even distribution of "

215 + f"experts across ranks. Got {self.global_num_experts} experts "

216 + f"and {self.ep_size} EP ranks."

217 + )

218 + self.eplb_state = EplbLayerState()

219 + else:

220 + assert not self.use_ep or num_redundant_experts == 0, (

221 + "Redundant experts are only supported with EPLB."

222 + )

223 +

224 # ROCm aiter shared experts fusion

225 # AITER only supports gated activations (silu/gelu), so disable it

226 # for non-gated MoE (is_act_and_mul=False)

@@ -237,17 +248,6 @@ class FusedMoE(PluggableLayer):

248 )

249

250 # Determine expert maps

240 - if self.use_ep:

241 - if self.enable_eplb:

242 - assert self.global_num_experts % self.ep_size == 0, (

243 - "EPLB currently only supports even distribution of "

244 - "experts across ranks."

245 - )

246 - else:

247 - assert num_redundant_experts == 0, (

248 - "Redundant experts are only supported with EPLB."

249 - )

250 -

251 max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens

252

253 # Create ExpertMapManager to handle expert mapping and placement for EP.

@@ -261,7 +261,7 @@ class FusedMoE(PluggableLayer):

261 num_expert_group=num_expert_group,

262 moe_parallel_config=self.moe_parallel_config,

263 placement_strategy=self.expert_placement_strategy,

264 - enable_eplb=self.enable_eplb,

264 + enable_eplb=enable_eplb,

265 num_fused_shared_experts=self.num_fused_shared_experts,

266 rocm_aiter_enabled=self.rocm_aiter_fmoe_enabled,

267 )

@@ -313,7 +313,6 @@ class FusedMoE(PluggableLayer):

313 routed_scaling_factor=self.routed_scaling_factor,

314 e_score_correction_bias=e_score_correction_bias,

315 num_fused_shared_experts=self.num_fused_shared_experts,

316 - enable_eplb=enable_eplb,

316 # TODO(bnell): once we can construct the MK at init time, we

317 # can make this a value.

318 indices_type_getter=lambda: self.quant_method.topk_indices_dtype,

@@ -381,7 +380,7 @@ class FusedMoE(PluggableLayer):

380 "is_act_and_mul=False is supported only for CUDA and XPU for now"

381 )

382

384 - if self.enable_eplb and not self.quant_method.supports_eplb:

383 + if enable_eplb and not self.quant_method.supports_eplb:

384 # TODO: Add support for additional quantization methods.

385 # The implementation for other quantization methods does not

386 # contain essential differences, but the current quant API

@@ -1282,10 +1281,20 @@ class FusedMoE(PluggableLayer):

1281

1282 This is used later in forward pass, where we get the expert mapping

1283 and record the load metrics in `expert_load_view`.

1284 +

1285 + Args:

1286 + moe_layer_idx: Index of this MoE layer

1287 + expert_load_view: View into global expert load tracking tensor

1288 + logical_to_physical_map: Mapping from logical to physical expert IDs

1289 + logical_replica_count: Number of replicas for each logical expert

1290 """

1286 - self.eplb_state.expert_load_view = expert_load_view[moe_layer_idx]

1287 - self.eplb_state.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]

1288 - self.eplb_state.logical_replica_count = logical_replica_count[moe_layer_idx]

1291 + if self.eplb_state is not None:

1292 + self.eplb_state.set_layer_state(

1293 + moe_layer_idx,

1294 + expert_load_view,

1295 + logical_to_physical_map,

1296 + logical_replica_count,

1297 + )

1298

1299 def ensure_moe_quant_config_init(self):

1300 if self.quant_method.moe_quant_config is None:

vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py

+1 -3

@@ -33,18 +33,16 @@ class AiterSharedRoutedFusedMoERouter(BaseRouter):

33 self,

34 top_k: int,

35 global_num_experts: int,

36 - eplb_state: EplbLayerState,

36 num_fused_shared_experts: int,

37 + eplb_state: EplbLayerState | None = None,

38 scoring_func: str = "softmax",

39 renormalize: bool = True,

40 - enable_eplb: bool = False,

40 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

41 ):

42 super().__init__(

43 top_k=top_k,

44 global_num_experts=global_num_experts,

45 eplb_state=eplb_state,

47 - enable_eplb=enable_eplb,

46 indices_type_getter=indices_type_getter,

47 )

48 self.renormalize = renormalize

vllm/model_executor/layers/fused_moe/router/base_router.py

+27 -27

@@ -148,8 +148,7 @@ class BaseRouter(FusedMoERouter):

148 self,

149 top_k: int,

150 global_num_experts: int,

151 - eplb_state: EplbLayerState,

152 - enable_eplb: bool = False,

151 + eplb_state: EplbLayerState | None = None,

152 # TODO(bnell): Once the MK is constructed at layer init time, we

153 # can make this a plain value instead of a callback.

154 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

@@ -159,12 +158,17 @@ class BaseRouter(FusedMoERouter):

158 time, so we need to supply a callback to get it at runtime. This is

159 because the indices type is supplied by modular kernels which are

160 created after MoE layer/router construction.

161 +

162 + Args:

163 + top_k: Number of experts to select per token

164 + global_num_experts: Total number of experts

165 + eplb_state: Optional EPLBLayerState for load balancing

166 + indices_type_getter: Optional callback to get indices dtype

167 """

168 super().__init__()

169 self.top_k = top_k

170 self.global_num_experts = global_num_experts

171 self.eplb_state = eplb_state

167 - self.enable_eplb = enable_eplb

172 self.indices_type_getter = indices_type_getter

173 self.capture_fn: Callable[[torch.Tensor], None] | None = None

174

@@ -174,21 +178,16 @@ class BaseRouter(FusedMoERouter):

178

179 def _validate_eplb_state(self) -> None:

180 """Validate that EPLB state is properly initialized if EPLB is enabled."""

177 - if self.enable_eplb:

178 - if self.eplb_state.expert_load_view is None:

179 - raise ValueError("enable_eplb=True requires expert_load_view != None")

180 - if self.eplb_state.logical_to_physical_map is None:

181 - raise ValueError(

182 - "enable_eplb=True requires logical_to_physical_map != None"

183 - )

184 - if self.eplb_state.logical_replica_count is None:

185 - raise ValueError(

186 - "enable_eplb=True requires logical_replica_count != None"

187 - )

188 - if self.eplb_state.should_record_tensor is None:

189 - raise ValueError(

190 - "enable_eplb=True requires should_record_tensor != None"

191 - )

181 + if self.eplb_state is not None:

182 + eplb_state = self.eplb_state

183 + if eplb_state.expert_load_view is None:

184 + raise ValueError("EPLB requires expert_load_view != None")

185 + if eplb_state.logical_to_physical_map is None:

186 + raise ValueError("EPLB requires logical_to_physical_map != None")

187 + if eplb_state.logical_replica_count is None:

188 + raise ValueError("EPLB requires logical_replica_count != None")

189 + if eplb_state.should_record_tensor is None:

190 + raise ValueError("EPLB requires should_record_tensor != None")

191

192 def _get_indices_type(self) -> torch.dtype | None:

193 """Get the desired indices dtype from the getter function."""

@@ -198,17 +197,18 @@ class BaseRouter(FusedMoERouter):

197

198 def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor:

199 """Apply EPLB mapping to convert logical expert IDs to physical expert IDs."""

201 - if self.enable_eplb:

202 - assert self.eplb_state.expert_load_view is not None

203 - assert self.eplb_state.logical_to_physical_map is not None

204 - assert self.eplb_state.logical_replica_count is not None

205 - assert self.eplb_state.should_record_tensor is not None

200 + if self.eplb_state is not None:

201 + eplb_state = self.eplb_state

202 + assert eplb_state.expert_load_view is not None

203 + assert eplb_state.logical_to_physical_map is not None

204 + assert eplb_state.logical_replica_count is not None

205 + assert eplb_state.should_record_tensor is not None

206 return eplb_map_to_physical_and_record(

207 topk_ids=topk_ids,

208 - logical_to_physical_map=self.eplb_state.logical_to_physical_map,

209 - logical_replica_count=self.eplb_state.logical_replica_count,

210 - expert_load_view=self.eplb_state.expert_load_view,

211 - record_enabled=self.eplb_state.should_record_tensor,

208 + logical_to_physical_map=eplb_state.logical_to_physical_map,

209 + logical_replica_count=eplb_state.logical_replica_count,

210 + expert_load_view=eplb_state.expert_load_view,

211 + record_enabled=eplb_state.should_record_tensor,

212 )

213 return topk_ids

214

vllm/model_executor/layers/fused_moe/router/custom_routing_router.py

+1 -3

@@ -16,17 +16,15 @@ class CustomRoutingRouter(BaseRouter):

16 self,

17 top_k: int,

18 global_num_experts: int,

19 - eplb_state: EplbLayerState,

19 custom_routing_function: Callable,

20 + eplb_state: EplbLayerState | None = None,

21 renormalize: bool = True,

22 - enable_eplb: bool = False,

22 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

23 ):

24 super().__init__(

25 top_k=top_k,

26 global_num_experts=global_num_experts,

27 eplb_state=eplb_state,

29 - enable_eplb=enable_eplb,

28 indices_type_getter=indices_type_getter,

29 )

30 self.custom_routing_function = custom_routing_function

vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py

+1 -3

@@ -235,11 +235,10 @@ class FusedTopKBiasRouter(BaseRouter):

235 self,

236 top_k: int,

237 global_num_experts: int,

238 - eplb_state: EplbLayerState,

238 e_score_correction_bias: torch.Tensor | None = None,

239 renormalize: bool = True,

240 routed_scaling_factor: float = 1.0,

242 - enable_eplb: bool = False,

241 + eplb_state: EplbLayerState | None = None,

242 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

243 *,

244 scoring_func: str = "sigmoid",

@@ -249,7 +248,6 @@ class FusedTopKBiasRouter(BaseRouter):

248 top_k=top_k,

249 global_num_experts=global_num_experts,

250 eplb_state=eplb_state,

252 - enable_eplb=enable_eplb,

251 indices_type_getter=indices_type_getter,

252 )

253 self.e_score_correction_bias = e_score_correction_bias

vllm/model_executor/layers/fused_moe/router/fused_topk_router.py

+1 -3

@@ -120,17 +120,15 @@ class FusedTopKRouter(BaseRouter):

120 self,

121 top_k: int,

122 global_num_experts: int,

123 - eplb_state: EplbLayerState,

123 scoring_func: str = "softmax",

124 renormalize: bool = True,

126 - enable_eplb: bool = False,

125 + eplb_state: EplbLayerState | None = None,

126 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

127 ):

128 super().__init__(

129 top_k=top_k,

130 global_num_experts=global_num_experts,

131 eplb_state=eplb_state,

133 - enable_eplb=enable_eplb,

132 indices_type_getter=indices_type_getter,

133 )

134 self.renormalize = renormalize

vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py

+1 -3

@@ -251,7 +251,6 @@ class GroupedTopKRouter(BaseRouter):

251 self,

252 top_k: int,

253 global_num_experts: int,

254 - eplb_state: EplbLayerState,

254 num_expert_group: int,

255 topk_group: int,

256 renormalize: bool = True,

@@ -259,14 +258,13 @@ class GroupedTopKRouter(BaseRouter):

258 routed_scaling_factor: float = 1.0,

259 e_score_correction_bias: torch.Tensor | None = None,

260 num_fused_shared_experts: int = 0,

262 - enable_eplb: bool = False,

261 + eplb_state: EplbLayerState | None = None,

262 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

263 ):

264 super().__init__(

265 top_k=top_k,

266 global_num_experts=global_num_experts,

267 eplb_state=eplb_state,

269 - enable_eplb=enable_eplb,

268 indices_type_getter=indices_type_getter,

269 )

270 self.num_expert_group = num_expert_group

vllm/model_executor/layers/fused_moe/router/router_factory.py

+2 -13

@@ -35,8 +35,6 @@ from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (

35 ZeroExpertRouter,

36 )

38 -EMPTY_EPLB_STATE: EplbLayerState = EplbLayerState()

39 -

39 def create_fused_moe_router(

40 # common parameters

@@ -56,8 +54,7 @@ def create_fused_moe_router(

54 # custom routing parameters

55 custom_routing_function: Callable | None = None,

56 # eplb parameters

59 - enable_eplb: bool = False,

60 - eplb_state: EplbLayerState = EMPTY_EPLB_STATE,

57 + eplb_state: EplbLayerState | None = None,

58 # zero expert parameters

59 zero_expert_type: str | None = None,

60 num_logical_experts: int | None = None,

@@ -98,8 +95,7 @@ def create_fused_moe_router(

95 custom_routing_function: Optional custom routing function

97 EPLB arguments:

101 - enable_eplb: Whether EPLB is enabled

102 - eplb_state: EPLB (Expert Parallelism Load Balancing) state

98 + eplb_state: Optional EplbLayerState, None when EPLB is disabled.

100 Zero expert arguments:

101 zero_expert_type: Type of zero expert (e.g. identity). If not None,

@@ -120,7 +116,6 @@ def create_fused_moe_router(

116 top_k=top_k,

117 global_num_experts=global_num_experts,

118 eplb_state=eplb_state,

123 - enable_eplb=enable_eplb,

119 indices_type_getter=indices_type_getter,

120 )

121

@@ -141,7 +136,6 @@ def create_fused_moe_router(

136 scoring_func=scoring_func,

137 renormalize=renormalize,

138 routed_scaling_factor=routed_scaling_factor,

144 - enable_eplb=enable_eplb,

139 indices_type_getter=indices_type_getter,

140 )

141

@@ -163,7 +157,6 @@ def create_fused_moe_router(

157 routed_scaling_factor=routed_scaling_factor,

158 e_score_correction_bias=e_score_correction_bias,

159 num_fused_shared_experts=num_fused_shared_experts,

166 - enable_eplb=enable_eplb,

160 indices_type_getter=indices_type_getter,

161 )

162 if (

@@ -186,7 +179,6 @@ def create_fused_moe_router(

179 eplb_state=eplb_state,

180 custom_routing_function=custom_routing_function,

181 renormalize=renormalize,

189 - enable_eplb=enable_eplb,

182 indices_type_getter=indices_type_getter,

183 )

184

@@ -200,7 +192,6 @@ def create_fused_moe_router(

192 e_score_correction_bias=e_score_correction_bias,

193 renormalize=renormalize,

194 routed_scaling_factor=routed_scaling_factor,

203 - enable_eplb=enable_eplb,

195 indices_type_getter=indices_type_getter,

196 scoring_func=scoring_func,

197 hash_indices_table=hash_indices_table,

@@ -218,7 +209,6 @@ def create_fused_moe_router(

209 num_fused_shared_experts=num_fused_shared_experts,

210 renormalize=renormalize,

211 scoring_func=scoring_func,

221 - enable_eplb=enable_eplb,

212 indices_type_getter=indices_type_getter,

213 )

214

@@ -228,6 +218,5 @@ def create_fused_moe_router(

218 eplb_state=eplb_state,

219 renormalize=renormalize,

220 scoring_func=scoring_func,

231 - enable_eplb=enable_eplb,

221 indices_type_getter=indices_type_getter,

222 )

vllm/model_executor/layers/fused_moe/router/routing_simulator_router.py

+1 -3

@@ -313,15 +313,13 @@ class RoutingSimulatorRouter(BaseRouter):

313 self,

314 top_k: int,

315 global_num_experts: int,

316 - eplb_state: EplbLayerState,

317 - enable_eplb: bool = False,

316 + eplb_state: EplbLayerState | None = None,

317 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

318 ):

319 super().__init__(

320 top_k=top_k,

321 global_num_experts=global_num_experts,

322 eplb_state=eplb_state,

324 - enable_eplb=enable_eplb,

323 indices_type_getter=indices_type_getter,

324 )

325

vllm/model_executor/layers/fused_moe/router/zero_expert_router.py

+1 -3

@@ -32,21 +32,19 @@ class ZeroExpertRouter(BaseRouter):

32 self,

33 top_k: int,

34 global_num_experts: int,

35 - eplb_state: EplbLayerState,

35 e_score_correction_bias: torch.Tensor,

36 num_logical_experts: int,

37 zero_expert_type: str,

38 scoring_func: str = "softmax",

39 renormalize: bool = False,

40 routed_scaling_factor: float = 1.0,

42 - enable_eplb: bool = False,

41 + eplb_state: EplbLayerState | None = None,

42 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

43 ):

44 super().__init__(

45 top_k=top_k,

46 global_num_experts=global_num_experts,

47 eplb_state=eplb_state,

49 - enable_eplb=enable_eplb,

48 indices_type_getter=indices_type_getter,

49 )

50 self.e_score_correction_bias = e_score_correction_bias

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py

-4

@@ -311,10 +311,6 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

311 shared_experts: SharedExperts | None,

312 shared_experts_input: torch.Tensor | None,

313 ) -> torch.Tensor:

314 - if layer.enable_eplb:

315 - raise NotImplementedError(

316 - "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."

317 - )

314 assert self.moe_quant_config is not None

315

316 from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py

-1

@@ -307,7 +307,6 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):

307 router_logits: torch.Tensor,

308 input_ids: torch.Tensor | None = None,

309 ) -> torch.Tensor:

310 - assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet."

310 assert layer.activation in (

311 MoEActivation.SILU,

312 MoEActivation.SWIGLUOAI,

vllm/model_executor/layers/quantization/modelopt.py

+1 -1

@@ -2101,7 +2101,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):

2101

2102 assert self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

2103

2104 - if layer.enable_eplb:

2104 + if layer.eplb_state is not None:

2105 raise NotImplementedError(

2106 "EPLB is not supported for FlashInfer TRTLLM MXFP8 MoE backend."

2107 )