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

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 13, 20:45 UTC

[Kernels] Isolate modular kernel code from FusedMoEMethodBase subclasses. (#27123)

bnellnm committed Nov 4, 2025 at 08:59 UTC 938772af03ce01590c7e92b0d3fd0a5bdc899d19

16 files changed +271 -311

vllm/distributed/device_communicators/base_device_communicator.py

+2 -2

@@ -266,14 +266,14 @@ class DeviceCommunicatorBase:

266 module

267 for module in model.modules()

268 # TODO(bnell): Should use isinstance but can't. Maybe search for

269 - # presence of quant_method.init_prepare_finalize?

269 + # presence of quant_method.maybe_init_modular_kernel?

270 if (

271 module.__class__.__name__ == "FusedMoE"

272 or module.__class__.__name__ == "SharedFusedMoE"

273 )

274 ]

275 for module in moe_modules:

276 - module.quant_method.init_prepare_finalize(module)

276 + module.maybe_init_modular_kernel()

277

278 def dispatch(

279 self,

vllm/model_executor/layers/fused_moe/layer.py

+199 -62

@@ -117,10 +117,8 @@ class FusedMoeWeightScaleSupported(Enum):

117 class FusedMoEMethodBase(QuantizeMethodBase):

118 def __init__(self, moe: FusedMoEConfig):

119 super().__init__()

120 - self.moe = moe

120 + self.moe: FusedMoEConfig = moe

121 self.moe_quant_config: FusedMoEQuantConfig | None = None

122 - self.fused_experts: FusedMoEModularKernel | None = None

123 - self.topk_indices_dtype = None

122

123 @abstractmethod

124 def create_weights(

@@ -245,9 +243,9 @@ class FusedMoEMethodBase(QuantizeMethodBase):

243 else:

244 return None

245

248 - # Note: init_prepare_finalize should only be called by

249 - # prepare_communication_buffer_for_model.

250 - def init_prepare_finalize(self, layer: torch.nn.Module):

246 + def maybe_init_modular_kernel(

247 + self, layer: torch.nn.Module

248 + ) -> FusedMoEModularKernel | None:

249 assert self.moe is not None

250

251 # We must get the quant config here so that the layer is

@@ -261,17 +259,14 @@ class FusedMoEMethodBase(QuantizeMethodBase):

259 logger.debug(

260 "%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)

261 )

264 - assert self.topk_indices_dtype is None

265 - assert self.fused_experts is None, (

266 - f"Attempt to override experts for {id(self)}!"

267 - )

268 - self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()

262 experts = self.select_gemm_impl(prepare_finalize, layer)

270 - self.fused_experts = FusedMoEModularKernel(

263 + return FusedMoEModularKernel(

264 prepare_finalize,

265 experts,

266 layer.shared_experts,

267 )

268 + else:

269 + return None

270

271 def select_gemm_impl(

272 self,

@@ -292,8 +287,16 @@ class FusedMoEMethodBase(QuantizeMethodBase):

287 raise NotImplementedError

288

289 @property

295 - def using_modular_kernel(self) -> bool:

296 - return self.fused_experts is not None

290 + def topk_indices_dtype(self) -> torch.dtype | None:

291 + return None

292 +

293 + @property

294 + def supports_eplb(self) -> bool:

295 + return False

296 +

297 + @property

298 + def allow_inplace(self) -> bool:

299 + return False

300

301 @abstractmethod

302 def apply(

@@ -322,6 +325,138 @@ class FusedMoEMethodBase(QuantizeMethodBase):

325 raise NotImplementedError

326

327

328 +@CustomOp.register("modular_fused_moe")

329 +class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

330 + def __init__(

331 + self,

332 + old_quant_method: FusedMoEMethodBase,

333 + fused_experts: FusedMoEModularKernel,

334 + ):

335 + super().__init__(old_quant_method.moe)

336 + # Find better way to copy attributes? Should we even copy attributes?

337 + # self.__dict__.update(old_quant_method.__dict__)

338 + self.moe_quant_config = old_quant_method.moe_quant_config

339 + self.fused_experts = fused_experts

340 + self.disable_expert_map = getattr(

341 + old_quant_method,

342 + "disable_expert_map",

343 + not fused_experts.supports_expert_map(),

344 + )

345 + self.old_quant_method = old_quant_method

346 + logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)

347 +

348 + @property

349 + def topk_indices_dtype(self) -> torch.dtype | None:

350 + return self.fused_experts.prepare_finalize.topk_indices_dtype()

351 +

352 + @property

353 + def supports_eplb(self) -> bool:

354 + return self.old_quant_method.supports_eplb

355 +

356 + @property

357 + def allow_inplace(self) -> bool:

358 + return self.old_quant_method.allow_inplace

359 +

360 + def create_weights(

361 + self,

362 + layer: torch.nn.Module,

363 + num_experts: int,

364 + hidden_size: int,

365 + intermediate_size_per_partition: int,

366 + params_dtype: torch.dtype,

367 + **extra_weight_attrs,

368 + ):

369 + raise NotImplementedError

370 +

371 + def get_fused_moe_quant_config(

372 + self, layer: torch.nn.Module

373 + ) -> FusedMoEQuantConfig | None:

374 + return self.moe_quant_config

375 +

376 + def apply(

377 + self,

378 + layer: torch.nn.Module,

379 + x: torch.Tensor,

380 + router_logits: torch.Tensor,

381 + top_k: int,

382 + renormalize: bool,

383 + use_grouped_topk: bool = False,

384 + topk_group: int | None = None,

385 + num_expert_group: int | None = None,

386 + global_num_experts: int = -1,

387 + expert_map: torch.Tensor | None = None,

388 + custom_routing_function: Callable | None = None,

389 + scoring_func: str = "softmax",

390 + routed_scaling_factor: float = 1.0,

391 + e_score_correction_bias: torch.Tensor | None = None,

392 + apply_router_weight_on_input: bool = False,

393 + activation: str = "silu",

394 + enable_eplb: bool = False,

395 + expert_load_view: torch.Tensor | None = None,

396 + logical_to_physical_map: torch.Tensor | None = None,

397 + logical_replica_count: torch.Tensor | None = None,

398 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

399 + # Is getattr needed?

400 + zero_expert_num = getattr(layer, "zero_expert_num", 0)

401 + zero_expert_type = getattr(layer, "zero_expert_type", None)

402 +

403 + if enable_eplb:

404 + if self.supports_eplb:

405 + assert expert_load_view is not None

406 + assert logical_to_physical_map is not None

407 + assert logical_replica_count is not None

408 + assert isinstance(layer, FusedMoE)

409 + else:

410 + raise NotImplementedError(

411 + "EPLB is not supported for "

412 + f"{self.old_quant_method.__class__.__name__}."

413 + )

414 +

415 + topk_weights, topk_ids, zero_expert_result = FusedMoE.select_experts(

416 + hidden_states=x,

417 + router_logits=router_logits,

418 + use_grouped_topk=use_grouped_topk,

419 + top_k=top_k,

420 + renormalize=renormalize,

421 + topk_group=topk_group,

422 + num_expert_group=num_expert_group,

423 + custom_routing_function=custom_routing_function,

424 + scoring_func=scoring_func,

425 + routed_scaling_factor=routed_scaling_factor,

426 + e_score_correction_bias=e_score_correction_bias,

427 + indices_type=self.topk_indices_dtype,

428 + enable_eplb=enable_eplb,

429 + expert_map=expert_map,

430 + expert_load_view=expert_load_view,

431 + logical_to_physical_map=logical_to_physical_map,

432 + logical_replica_count=logical_replica_count,

433 + global_num_experts=global_num_experts,

434 + zero_expert_num=zero_expert_num,

435 + zero_expert_type=zero_expert_type,

436 + )

437 +

438 + result = self.fused_experts(

439 + hidden_states=x,

440 + w1=layer.w13_weight,

441 + w2=layer.w2_weight,

442 + topk_weights=topk_weights,

443 + topk_ids=topk_ids,

444 + inplace=self.allow_inplace,

445 + activation=activation,

446 + global_num_experts=global_num_experts,

447 + apply_router_weight_on_input=apply_router_weight_on_input,

448 + expert_map=None if self.disable_expert_map else expert_map,

449 + )

450 +

451 + if zero_expert_num != 0 and zero_expert_type is not None:

452 + assert not isinstance(result, tuple), (

453 + "Shared + zero experts are mutually exclusive not yet supported"

454 + )

455 + return result, zero_expert_result

456 + else:

457 + return result

458 +

459 +

460 @CustomOp.register("unquantized_fused_moe")

461 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

462 """MoE method without quantization."""

@@ -378,6 +513,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

513 )

514 self.flashinfer_cutlass_moe = None # type: ignore

515

516 + @property

517 + def supports_eplb(self) -> bool:

518 + return True

519 +

520 + @property

521 + def allow_inplace(self) -> bool:

522 + return True

523 +

524 def maybe_make_prepare_finalize(self) -> FusedMoEPrepareAndFinalize | None:

525 if self.rocm_aiter_moe_enabled:

526 return None

@@ -650,7 +793,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

793 )

794

795 if self.rocm_aiter_moe_enabled:

653 - assert self.fused_experts is None

796 result = self.rocm_aiter_fused_experts(

797 hidden_states=x,

798 w1=layer.w13_weight,

@@ -671,21 +813,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

813 activation=activation,

814 apply_router_weight_on_input=apply_router_weight_on_input,

815 )

674 - elif self.fused_experts is not None:

675 - result = self.fused_experts(

676 - hidden_states=x,

677 - w1=layer.w13_weight,

678 - w2=layer.w2_weight,

679 - topk_weights=topk_weights,

680 - topk_ids=topk_ids,

681 - inplace=True,

682 - activation=activation,

683 - apply_router_weight_on_input=apply_router_weight_on_input,

684 - global_num_experts=global_num_experts,

685 - expert_map=expert_map,

686 - )

816 else:

688 - assert fused_experts is not None

817 result = fused_experts(

818 hidden_states=x,

819 w1=layer.w13_weight,

@@ -1267,7 +1395,7 @@ class FusedMoE(CustomOp):

1395 "Only softmax scoring function is supported for non-grouped topk."

1396 )

1397

1270 - moe = FusedMoEConfig(

1398 + self.moe_config: FusedMoEConfig = FusedMoEConfig(

1399 num_experts=self.global_num_experts,

1400 experts_per_token=top_k,

1401 hidden_dim=hidden_size,

@@ -1279,24 +1407,26 @@ class FusedMoE(CustomOp):

1407 is_act_and_mul=is_act_and_mul,

1408 is_lora_enabled=vllm_config.lora_config is not None,

1409 )

1282 - self.moe_config: FusedMoEConfig = moe

1410 +

1411 self.moe_quant_config: FusedMoEQuantConfig | None = None

1412 self.quant_config = quant_config

1413

1414 + def _get_quant_method() -> FusedMoEMethodBase:

1415 + """

1416 + Helper method to ensure self.quant_method is never None and

1417 + of the proper type.

1418 + """

1419 + quant_method = None

1420 + if self.quant_config is not None:

1421 + quant_method = self.quant_config.get_quant_method(self, prefix)

1422 + if quant_method is None:

1423 + quant_method = UnquantizedFusedMoEMethod(self.moe_config)

1424 + assert isinstance(quant_method, FusedMoEMethodBase)

1425 + return quant_method

1426 +

1427 # Note: get_quant_method will look at the layer's local_num_experts

1428 # for heuristic purposes, so it must be initialized first.

1288 - quant_method: QuantizeMethodBase | None = None

1289 - quant_method = (

1290 - UnquantizedFusedMoEMethod(moe)

1291 - if quant_config is None

1292 - else quant_config.get_quant_method(self, prefix)

1293 - )

1294 - if quant_method is None:

1295 - quant_method = UnquantizedFusedMoEMethod(moe)

1296 -

1297 - assert quant_method is not None

1298 - assert isinstance(quant_method, FusedMoEMethodBase)

1299 - self.quant_method = quant_method

1429 + self.quant_method: FusedMoEMethodBase = _get_quant_method()

1430

1431 if not self.moe_config.is_act_and_mul:

1432 # Avoid circular import

@@ -1305,7 +1435,7 @@ class FusedMoE(CustomOp):

1435 )

1436

1437 if not isinstance(

1308 - quant_method, (UnquantizedFusedMoEMethod, ModelOptFp8MoEMethod)

1438 + self.quant_method, (UnquantizedFusedMoEMethod, ModelOptFp8MoEMethod)

1439 ):

1440 raise NotImplementedError(

1441 "is_act_and_mul=False is supported only for unquantized "

@@ -1316,20 +1446,18 @@ class FusedMoE(CustomOp):

1446 "is_act_and_mul=False is supported only for CUDA for now"

1447 )

1448

1319 - if self.enable_eplb:

1320 - from vllm.model_executor.layers.quantization.fp8 import Fp8MoEMethod

1321 -

1322 - if not isinstance(quant_method, (Fp8MoEMethod, UnquantizedFusedMoEMethod)):

1323 - # TODO: Add support for additional quantization methods.

1324 - # The implementation for other quantization methods does not

1325 - # contain essential differences, but the current quant API

1326 - # design causes duplicated work when extending to new

1327 - # quantization methods, so I'm leaving it for now.

1328 - # If you plan to add support for more quantization methods,

1329 - # please refer to the implementation in `Fp8MoEMethod`.

1330 - raise NotImplementedError(

1331 - "EPLB is only supported for FP8 quantization for now."

1332 - )

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

1450 + # TODO: Add support for additional quantization methods.

1451 + # The implementation for other quantization methods does not

1452 + # contain essential differences, but the current quant API

1453 + # design causes duplicated work when extending to new

1454 + # quantization methods, so I'm leaving it for now.

1455 + # If you plan to add support for more quantization methods,

1456 + # please refer to the implementation in `Fp8MoEMethod`.

1457 + raise NotImplementedError(

1458 + f"EPLB is not supported {self.quant_method.__class__.__name__}. "

1459 + "EPLB is only supported for FP8 quantization for now."

1460 + )

1461

1462 moe_quant_params = {

1463 "num_experts": self.local_num_experts,

@@ -1353,6 +1481,15 @@ class FusedMoE(CustomOp):

1481 self.batched_hidden_states: torch.Tensor | None = None

1482 self.batched_router_logits: torch.Tensor | None = None

1483

1484 + # Note: maybe_init_modular_kernel should only be called by

1485 + # prepare_communication_buffer_for_model.

1486 + # This is called after all weight loading and post-processing, so it

1487 + # should be safe to swap out the quant_method.

1488 + def maybe_init_modular_kernel(self) -> None:

1489 + mk = self.quant_method.maybe_init_modular_kernel(self)

1490 + if mk is not None:

1491 + self.quant_method = FusedMoEModularMethod(self.quant_method, mk)

1492 +

1493 @property

1494 def shared_experts(self) -> torch.nn.Module | None:

1495 return None

@@ -2167,7 +2304,7 @@ class FusedMoE(CustomOp):

2304 """

2305 assert self.quant_method is not None

2306 return (

2170 - self.quant_method.fused_experts is not None

2307 + isinstance(self.quant_method, FusedMoEModularMethod)

2308 and self.quant_method.fused_experts.output_is_reduced()

2309 )

2310

@@ -2403,7 +2540,7 @@ class FusedMoE(CustomOp):

2540 self.ensure_dp_chunking_init()

2541

2542 has_separate_shared_experts = (

2406 - not isinstance(self.quant_method.fused_experts, FusedMoEModularKernel)

2543 + not isinstance(self.quant_method, FusedMoEModularMethod)

2544 and self.shared_experts is not None

2545 )

2546

@@ -2430,8 +2567,8 @@ class FusedMoE(CustomOp):

2567 hidden_states, router_logits, has_separate_shared_experts

2568 )

2569

2433 - do_naive_dispatch_combine: bool = (

2434 - self.dp_size > 1 and not self.quant_method.using_modular_kernel

2570 + do_naive_dispatch_combine: bool = self.dp_size > 1 and not isinstance(

2571 + self.quant_method, FusedMoEModularMethod

2572 )

2573

2574 # If there are shared experts but we are not using a modular kernel, the

vllm/model_executor/layers/fused_moe/modular_kernel.py

@@ -707,6 +707,12 @@ class FusedMoEModularKernel(torch.nn.Module):

707 f"{fused_experts.activation_formats[0]}"

708 )

709

710 + def supports_expert_map(self) -> bool:

711 + """

712 + A flag indicating whether or not this class supports expert maps.

713 + """

714 + return self.fused_experts.supports_expert_map()

715 +

716 def output_is_reduced(self) -> bool:

717 """

718 Indicates whether or not the output of fused MoE kernel

vllm/model_executor/layers/quantization/awq_marlin.py

-2

@@ -617,8 +617,6 @@ class AWQMoEMethod(FusedMoEMethodBase):

617 logical_to_physical_map: torch.Tensor | None = None,

618 logical_replica_count: torch.Tensor | None = None,

619 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

620 - assert self.fused_experts is None

621 -

620 if enable_eplb:

621 raise NotImplementedError("EPLB not supported for `AWQMoEMethod` yet.")

622

vllm/model_executor/layers/quantization/bitsandbytes.py

+1 -2

@@ -518,12 +518,11 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

518 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

519 from vllm.model_executor.layers.fused_moe import fused_experts

520

521 - assert self.fused_experts is None

522 -

521 if enable_eplb:

522 raise NotImplementedError(

523 "EPLB not supported for `BitsAndBytesMoEMethod` yet."

524 )

525 +

526 topk_weights, topk_ids, _ = FusedMoE.select_experts(

527 hidden_states=x,

528 router_logits=router_logits,

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

-47

@@ -462,12 +462,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

462 indices_type=self.topk_indices_dtype,

463 )

464

465 - #

466 - # Note: the order here is important. self.fused_experts can override

467 - # flashinfer cutlass, cutlass fp4 or fused_experts but not marlin.

468 - #

465 if self.use_marlin:

470 - assert self.fused_experts is None

466 return fused_marlin_moe(

467 x,

468 layer.w13_weight,

@@ -488,24 +483,6 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

483 workspace=layer.workspace,

484 )

485

491 - elif self.fused_experts is not None:

492 - assert is_valid_flashinfer_cutlass_fused_moe(

493 - x, layer.w13_weight, layer.w2_weight

494 - ), "Flashinfer CUTLASS Fused MoE not applicable!"

495 -

496 - return self.fused_experts(

497 - hidden_states=x,

498 - w1=layer.w13_weight,

499 - w2=layer.w2_weight,

500 - topk_weights=topk_weights,

501 - topk_ids=topk_ids,

502 - inplace=False, # TODO(shuw): fix later, now output is high prec

503 - activation=activation,

504 - global_num_experts=global_num_experts,

505 - expert_map=expert_map,

506 - apply_router_weight_on_input=apply_router_weight_on_input,

507 - )

508 -

486 # FlashInfer fused experts path

487 elif self.allow_flashinfer:

488 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501

@@ -1066,13 +1043,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1043 per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN

1044 per_channel_quant = self.weight_quant.strategy == QuantizationStrategy.CHANNEL

1045

1069 - #

1070 - # Note: the order here is important. self.fused_experts can override

1071 - # cutlass fp8 or fused_experts but not marlin or rocm.

1072 - #

1046 if self.use_marlin:

1047 assert activation == "silu", f"{activation} not supported for Marlin MoE."

1075 - assert self.fused_experts is None

1048 return fused_marlin_moe(

1049 x,

1050 layer.w13_weight,

@@ -1098,7 +1070,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1070

1071 assert per_act_token == per_channel_quant

1072 assert self.moe_quant_config is not None

1101 - assert self.fused_experts is None

1073 return rocm_aiter_fused_experts(

1074 hidden_states=x,

1075 w1=layer.w13_weight,

@@ -1111,18 +1082,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1082 quant_config=self.moe_quant_config,

1083 )

1084

1114 - elif self.fused_experts is not None:

1115 - return self.fused_experts(

1116 - x,

1117 - layer.w13_weight,

1118 - layer.w2_weight,

1119 - topk_weights,

1120 - topk_ids,

1121 - activation=activation,

1122 - global_num_experts=global_num_experts,

1123 - expert_map=None if self.disable_expert_map else expert_map,

1124 - )

1125 -

1085 # cutlass path

1086 elif self.use_cutlass:

1087 assert self.moe_quant_config is not None

@@ -1318,8 +1277,6 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1277 logical_to_physical_map: torch.Tensor | None = None,

1278 logical_replica_count: torch.Tensor | None = None,

1279 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1321 - assert self.fused_experts is None

1322 -

1280 if enable_eplb:

1281 raise NotImplementedError(

1282 "EPLB not supported for `CompressedTensorsW8A8Int8MoEMethod` yet."

@@ -1636,8 +1593,6 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1593 logical_to_physical_map: torch.Tensor | None = None,

1594 logical_replica_count: torch.Tensor | None = None,

1595 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1639 - assert self.fused_experts is None

1640 -

1596 if enable_eplb:

1597 raise NotImplementedError(

1598 "EPLB not supported for `CompressedTensorsWNA16MarlinMoEMethod` yet."

@@ -1901,8 +1856,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1856 logical_to_physical_map: torch.Tensor | None = None,

1857 logical_replica_count: torch.Tensor | None = None,

1858 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1904 - assert self.fused_experts is None

1905 -

1859 if enable_eplb:

1860 raise NotImplementedError(

1861 "EPLB not supported for `CompressedTensorsWNA16MoEMethod` yet."

vllm/model_executor/layers/quantization/experts_int8.py

-2

@@ -158,8 +158,6 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

158 logical_to_physical_map: torch.Tensor | None = None,

159 logical_replica_count: torch.Tensor | None = None,

160 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

161 - assert self.fused_experts is None

162 -

161 if enable_eplb:

162 raise NotImplementedError(

163 "EPLB not supported for `ExpertsInt8MoEMethod` yet."

vllm/model_executor/layers/quantization/fp8.py

+9 -26

@@ -703,9 +703,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

703 self.quant_config = quant_config

704 self.weight_block_size = self.quant_config.weight_block_size

705 self.block_quant: bool = self.weight_block_size is not None

706 -

707 - self.fused_experts: mk.FusedMoEModularKernel | None = None # type: ignore

708 -

706 self.fp8_backend = get_fp8_moe_backend(self.block_quant)

707

708 self.use_marlin = self.fp8_backend == Fp8MoeBackend.MARLIN

@@ -1181,6 +1178,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1178 block_shape=self.weight_block_size,

1179 )

1180

1181 + @property

1182 + def supports_eplb(self) -> bool:

1183 + return True

1184 +

1185 + @property

1186 + def allow_inplace(self) -> bool:

1187 + return True

1188 +

1189 def apply(

1190 self,

1191 layer: torch.nn.Module,

@@ -1210,10 +1215,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1215 assert logical_replica_count is not None

1216 assert isinstance(layer, FusedMoE)

1217

1213 - if (

1214 - self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1215 - and self.fused_experts is None

1216 - ):

1218 + if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

1219 assert activation == "silu", (

1220 f"Expected 'silu' activation but got {activation}"

1221 )

@@ -1290,10 +1292,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1292 num_fused_shared_experts=layer.num_fused_shared_experts,

1293 )

1294

1293 - #

1294 - # Note: the order of checks is important since self.fused_experts

1295 - # can override fused_experts or cutlass but not rocm or marlin.

1296 - #

1295 topk_weights, topk_ids, zero_expert_result = select_result

1296

1297 if self.rocm_aiter_moe_enabled:

@@ -1301,7 +1299,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1299 rocm_aiter_fused_experts,

1300 )

1301

1304 - assert self.fused_experts is None

1302 result = rocm_aiter_fused_experts(

1303 x,

1304 layer.w13_weight,

@@ -1315,7 +1312,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1312 )

1313 elif self.use_marlin:

1314 assert activation == "silu", f"{activation} not supported for Marlin MoE."

1318 - assert self.fused_experts is None

1315 result = fused_marlin_moe(

1316 x,

1317 layer.w13_weight,

@@ -1333,19 +1329,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1329 expert_map=expert_map,

1330 workspace=layer.workspace,

1331 )

1336 - elif self.fused_experts:

1337 - result = self.fused_experts(

1338 - hidden_states=x,

1339 - w1=layer.w13_weight,

1340 - w2=layer.w2_weight,

1341 - topk_weights=topk_weights,

1342 - topk_ids=topk_ids,

1343 - inplace=True,

1344 - activation=activation,

1345 - global_num_experts=global_num_experts,

1346 - apply_router_weight_on_input=apply_router_weight_on_input,

1347 - expert_map=expert_map,

1348 - )

1332 elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1333 assert not self.block_quant

1334 assert not renormalize and custom_routing_function is not None

vllm/model_executor/layers/quantization/gguf.py

-2

@@ -585,8 +585,6 @@ class GGUFMoEMethod(FusedMoEMethodBase):

585 logical_to_physical_map: torch.Tensor | None = None,

586 logical_replica_count: torch.Tensor | None = None,

587 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

588 - assert self.fused_experts is None

589 -

588 if enable_eplb:

589 raise NotImplementedError("EPLB not supported for `GGUFMoEMethod` yet.")

590

vllm/model_executor/layers/quantization/gptq_marlin.py

-2

@@ -742,8 +742,6 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

742 logical_to_physical_map: torch.Tensor | None = None,

743 logical_replica_count: torch.Tensor | None = None,

744 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

745 - assert self.fused_experts is None

746 -

745 if enable_eplb:

746 raise NotImplementedError(

747 "EPLB not supported for `GPTQMarlinMoEMethod` yet."

vllm/model_executor/layers/quantization/modelopt.py

+11 -39

@@ -18,9 +18,6 @@ from vllm.model_executor.layers.fused_moe.config import (

18 fp8_w8a8_moe_quant_config,

19 nvfp4_moe_quant_config,

20 )

21 -from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

22 - is_valid_flashinfer_cutlass_fused_moe,

23 -)

21 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

22 from vllm.model_executor.layers.fused_moe.layer import (

23 FusedMoE,

@@ -605,7 +602,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

602 )

603

604 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

608 - assert self.fused_experts is None

605 assert activation == "silu", (

606 f"Expected 'silu' activation but got {activation}"

607 )

@@ -638,24 +634,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

634 indices_type=self.topk_indices_dtype,

635 )

636

641 - #

642 - # Note: the order here is important. self.fused_experts can override

643 - # cutlass or fused_experts.

644 - #

645 - if self.fused_experts is not None:

646 - return self.fused_experts(

647 - x,

648 - layer.w13_weight,

649 - layer.w2_weight,

650 - topk_weights,

651 - topk_ids,

652 - inplace=False,

653 - activation=activation,

654 - global_num_experts=global_num_experts,

655 - expert_map=expert_map,

656 - apply_router_weight_on_input=apply_router_weight_on_input,

657 - )

658 - elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

637 + if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

638 assert not renormalize

639 assert activation == "silu", (

640 f"Expected 'silu' activation but got {activation}"

@@ -1647,8 +1626,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1626

1627 from vllm.model_executor.models.llama4 import Llama4MoE

1628

1650 - assert self.fused_experts is None

1651 -

1629 a1_gscale = layer.w13_input_scale_quant

1630 (hidden_states_fp4, hidden_states_scale_linear_fp4) = (

1631 flashinfer.fp4_quantize(

@@ -1720,13 +1697,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1697 indices_type=self.topk_indices_dtype,

1698 )

1699

1723 - #

1724 - # Note: the order here is important. self.fused_experts can override

1725 - # flashinfer cutlass, cutlass fp4 or fused_experts but not marlin or

1726 - # trtllm.

1727 - #

1700 if self.use_marlin:

1729 - assert self.fused_experts is None

1701 return fused_marlin_moe(

1702 x,

1703 layer.w13_weight,

@@ -1747,23 +1718,24 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1718 workspace=layer.workspace,

1719 )

1720

1750 - elif self.fused_experts is not None:

1751 - assert (

1752 - self.allow_flashinfer

1753 - and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1721 + elif (

1722 + self.allow_flashinfer

1723 + and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1724 + ):

1725 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501

1726 + flashinfer_cutlass_moe_fp4,

1727 )

1728

1756 - assert is_valid_flashinfer_cutlass_fused_moe(

1757 - x, layer.w13_weight, layer.w2_weight

1758 - ), "Flashinfer CUTLASS Fused MoE not applicable!"

1729 + assert self.moe_quant_config is not None

1730

1760 - return self.fused_experts(

1731 + return flashinfer_cutlass_moe_fp4(

1732 hidden_states=x,

1733 w1=layer.w13_weight,

1734 w2=layer.w2_weight,

1735 topk_weights=topk_weights,

1736 topk_ids=topk_ids,

1766 - inplace=False, # TODO(shuw): fix later, now output is high prec

1737 + quant_config=self.moe_quant_config,

1738 + inplace=False,

1739 activation=activation,

1740 global_num_experts=global_num_experts,

1741 expert_map=expert_map,

vllm/model_executor/layers/quantization/moe_wna16.py

-2

@@ -226,7 +226,6 @@ class MoeWNA16Method(FusedMoEMethodBase):

226 params_dtype: torch.dtype,

227 **extra_weight_attrs,

228 ):

229 - self.moe = layer

229 layer.quant_config = self.quant_config

230 bit8_pack_factor = self.quant_config.bit8_pack_factor

231 group_size = self.quant_config.group_size

@@ -381,7 +380,6 @@ class MoeWNA16Method(FusedMoEMethodBase):

380 logical_to_physical_map: torch.Tensor | None = None,

381 logical_replica_count: torch.Tensor | None = None,

382 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

384 - assert self.fused_experts is None

383 if enable_eplb:

384 raise NotImplementedError("EPLB not supported for `MoeWNA16Method` yet.")

385

vllm/model_executor/layers/quantization/mxfp4.py

+15 -90

@@ -197,8 +197,6 @@ class Mxfp4Config(QuantizationConfig):

197 class Mxfp4MoEMethod(FusedMoEMethodBase):

198 def __init__(self, moe: FusedMoEConfig):

199 super().__init__(moe)

200 - self.topk_indices_dtype = None

201 - self.moe = moe

200 self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)

201 self.max_capture_size = (

202 get_current_vllm_config().compilation_config.max_cudagraph_capture_size

@@ -815,6 +813,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

813 "EP batched experts format"

814 )

815 else:

816 + layer.w13_weight = (

817 + self.w13_weight_triton_tensor

818 + if layer.w13_weight is None

819 + else layer.w13_weight

820 + )

821 + layer.w2_weight = (

822 + self.w2_weight_triton_tensor

823 + if layer.w2_weight is None

824 + else layer.w2_weight

825 + )

826 + assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]])

827 +

828 assert self.moe_quant_config is not None

829 if (

830 self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM

@@ -838,71 +848,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

848 f"Incompatible Mxfp4 backend ({self.mxfp4_backend}) for EP"

849 )

850

841 - def _route_and_experts(

842 - self,

843 - layer: torch.nn.Module,

844 - x: torch.Tensor,

845 - router_logits: torch.Tensor,

846 - top_k: int,

847 - renormalize: bool,

848 - use_grouped_topk: bool = False,

849 - topk_group: int | None = None,

850 - num_expert_group: int | None = None,

851 - global_num_experts: int = -1,

852 - expert_map: torch.Tensor | None = None,

853 - custom_routing_function: Callable | None = None,

854 - scoring_func: str = "softmax",

855 - e_score_correction_bias: torch.Tensor | None = None,

856 - apply_router_weight_on_input: bool = False,

857 - activation: str = "silu",

858 - enable_eplb: bool = False,

859 - expert_load_view: torch.Tensor | None = None,

860 - logical_to_physical_map: torch.Tensor | None = None,

861 - logical_replica_count: torch.Tensor | None = None,

862 - ) -> torch.Tensor:

863 - assert isinstance(self.fused_experts, mk.FusedMoEModularKernel)

864 -

865 - topk_weights, topk_ids, _ = FusedMoE.select_experts(

866 - hidden_states=x,

867 - router_logits=router_logits,

868 - use_grouped_topk=use_grouped_topk,

869 - top_k=top_k,

870 - renormalize=renormalize,

871 - topk_group=topk_group,

872 - num_expert_group=num_expert_group,

873 - custom_routing_function=custom_routing_function,

874 - scoring_func=scoring_func,

875 - e_score_correction_bias=e_score_correction_bias,

876 - indices_type=self.topk_indices_dtype,

877 - enable_eplb=enable_eplb,

878 - expert_map=expert_map,

879 - expert_load_view=expert_load_view,

880 - logical_to_physical_map=logical_to_physical_map,

881 - logical_replica_count=logical_replica_count,

882 - )

883 -

884 - w13_weight = (

885 - self.w13_weight_triton_tensor

886 - if layer.w13_weight is None

887 - else layer.w13_weight

888 - )

889 - w2_weight = (

890 - self.w2_weight_triton_tensor if layer.w2_weight is None else layer.w2_weight

891 - )

892 - assert all([w is not None for w in [w13_weight, w2_weight]])

893 -

894 - return self.fused_experts(

895 - hidden_states=x,

896 - w1=w13_weight,

897 - w2=w2_weight,

898 - topk_weights=topk_weights,

899 - topk_ids=topk_ids,

900 - inplace=True,

901 - activation=activation,

902 - global_num_experts=global_num_experts,

903 - expert_map=expert_map,

904 - apply_router_weight_on_input=apply_router_weight_on_input,

905 - )

851 + @property

852 + def allow_inplace(self) -> bool:

853 + return True

854

855 def apply(

856 self,

@@ -930,29 +878,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

878 if enable_eplb:

879 raise NotImplementedError("EPLB is not supported for mxfp4")

880

933 - if self.fused_experts is not None:

934 - return self._route_and_experts(

935 - layer,

936 - x,

937 - router_logits,

938 - top_k,

939 - renormalize,

940 - use_grouped_topk,

941 - topk_group,

942 - num_expert_group,

943 - global_num_experts,

944 - expert_map,

945 - custom_routing_function,

946 - scoring_func,

947 - e_score_correction_bias,

948 - apply_router_weight_on_input,

949 - activation,

950 - enable_eplb,

951 - expert_load_view,

952 - logical_to_physical_map,

953 - logical_replica_count,

954 - )

955 -

881 if self.mxfp4_backend == Mxfp4Backend.MARLIN:

882 topk_weights, topk_ids, _ = FusedMoE.select_experts(

883 hidden_states=x,

vllm/model_executor/layers/quantization/quark/quark_moe.py

+25 -28

@@ -310,7 +310,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

310 # Property to determine if AITER is used

311 if self.rocm_aiter_moe_enabled:

312 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa E501

313 - rocm_aiter_fused_experts,

313 shuffle_weights,

314 )

315

@@ -322,17 +321,11 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

321 layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)

322 layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)

323

325 - self.rocm_aiter_fused_experts_func = rocm_aiter_fused_experts

324 elif self.use_marlin:

325 prepare_moe_fp8_layer_for_marlin(layer, False)

326 # Activations not quantized for marlin.

327 del layer.w13_input_scale

328 del layer.w2_input_scale

331 - self.fused_experts_func = None

332 - else:

333 - from vllm.model_executor.layers.fused_moe import fused_experts

334 -

335 - self.fused_experts_func = fused_experts

329

330 def get_fused_moe_quant_config(

331 self, layer: torch.nn.Module

@@ -369,8 +362,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

362 logical_to_physical_map: torch.Tensor | None = None,

363 logical_replica_count: torch.Tensor | None = None,

364 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

372 - assert self.fused_experts is None

373 -

365 if enable_eplb:

366 raise NotImplementedError(

367 "EPLB not supported for `QuarkW8A8Fp8MoEMethod` yet."

@@ -392,7 +383,11 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

383 )

384

385 if self.rocm_aiter_moe_enabled:

395 - return self.rocm_aiter_fused_experts_func(

386 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

387 + rocm_aiter_fused_experts,

388 + )

389 +

390 + return rocm_aiter_fused_experts(

391 hidden_states=x,

392 w1=layer.w13_weight,

393 w2=layer.w2_weight,

@@ -403,7 +398,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

398 quant_config=self.moe_quant_config,

399 expert_map=expert_map,

400 )

406 - if self.use_marlin:

401 + elif self.use_marlin:

402 assert activation == "silu", f"{activation} not supported for Marlin MoE."

403 return fused_marlin_moe(

404 x,

@@ -421,22 +416,22 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

416 global_num_experts=global_num_experts,

417 expert_map=expert_map,

418 )

419 + else:

420 + from vllm.model_executor.layers.fused_moe import fused_experts

421

425 - assert self.fused_experts_func is not None

426 -

427 - return self.fused_experts_func(

428 - hidden_states=x,

429 - w1=layer.w13_weight,

430 - w2=layer.w2_weight,

431 - topk_weights=topk_weights,

432 - topk_ids=topk_ids,

433 - inplace=True,

434 - activation=activation,

435 - apply_router_weight_on_input=apply_router_weight_on_input,

436 - global_num_experts=global_num_experts,

437 - expert_map=expert_map,

438 - quant_config=self.moe_quant_config,

439 - )

422 + return fused_experts(

423 + hidden_states=x,

424 + w1=layer.w13_weight,

425 + w2=layer.w2_weight,

426 + topk_weights=topk_weights,

427 + topk_ids=topk_ids,

428 + inplace=True,

429 + activation=activation,

430 + apply_router_weight_on_input=apply_router_weight_on_input,

431 + global_num_experts=global_num_experts,

432 + expert_map=expert_map,

433 + quant_config=self.moe_quant_config,

434 + )

435

436

437 class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

@@ -601,6 +596,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

596 block_shape=None,

597 )

598

599 + @property

600 + def allow_inplace(self) -> bool:

601 + return True

602 +

603 def apply(

604 self,

605 layer: torch.nn.Module,

@@ -624,8 +623,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

623 logical_to_physical_map: torch.Tensor | None = None,

624 logical_replica_count: torch.Tensor | None = None,

625 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

627 - assert self.fused_experts is None

628 -

626 if enable_eplb:

627 raise NotImplementedError(

628 "EPLB not supported for `QuarkOCP_MX_MoEMethod` yet."

vllm/model_executor/layers/quantization/rtn.py

-2

@@ -377,8 +377,6 @@ class RTNMoEMethod(FusedMoEMethodBase):

377 logical_to_physical_map: torch.Tensor | None = None,

378 logical_replica_count: torch.Tensor | None = None,

379 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

380 - assert self.fused_experts is None

381 -

380 if enable_eplb:

381 raise NotImplementedError("EPLB not supported for `RTNMoEMethod` yet.")

382

vllm/model_executor/warmup/deep_gemm_warmup.py

+3 -3

@@ -13,7 +13,7 @@ import vllm.envs as envs

13 from vllm.distributed.parallel_state import get_dp_group

14 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts

15 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M

16 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE

16 +from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEModularMethod

17 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel

18 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

19 TritonOrDeepGemmExperts,

@@ -160,8 +160,8 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:

160 ):

161 return False

162

163 - if not isinstance(module.quant_method.fused_experts, FusedMoEModularKernel):

164 - # fused_experts could invoke deep_gemm_moe_fp8

163 + if not isinstance(module.quant_method, FusedMoEModularMethod):

164 + # modular kernels could invoke deep_gemm_moe_fp8

165 return True

166

167 mk: FusedMoEModularKernel = module.quant_method.fused_experts