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

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 12, 23:30 UTC

[Feat][Perf] Enable deepep-low-latency with round-robin expert placement. (#28449)

Signed-off-by: bruceszchen <bruceszchen@tencent.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Chen Bruce committed Nov 19, 2025 at 20:46 UTC da2f6800e0d6ac768c6f63b95f7c0755407f4263

8 files changed +208 -37

vllm/model_executor/layers/fused_moe/all2all_utils.py

+11

@@ -67,6 +67,7 @@ def maybe_roundup_layer_hidden_size(

67 def maybe_make_prepare_finalize(

68 moe: FusedMoEConfig,

69 quant_config: FusedMoEQuantConfig | None,

70 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

71 ) -> FusedMoEPrepareAndFinalize | None:

72 if not moe.moe_parallel_config.use_all2all_kernels:

73 return None

@@ -134,6 +135,13 @@ def maybe_make_prepare_finalize(

135

136 elif moe.use_deepep_ll_kernels:

137 assert quant_config is not None

138 + global_to_physical = physical_to_global = local_expert_global_ids = None

139 + if routing_tables is not None:

140 + (

141 + global_to_physical,

142 + physical_to_global,

143 + local_expert_global_ids,

144 + ) = routing_tables

145 all_to_all_args = dict(

146 max_num_tokens_per_dp_rank=moe.max_num_tokens,

147 token_hidden_size=moe.hidden_dim,

@@ -155,6 +163,9 @@ def maybe_make_prepare_finalize(

163 max_tokens_per_rank=moe.max_num_tokens,

164 num_dispatchers=all2all_manager.world_size,

165 use_fp8_dispatch=use_fp8_dispatch,

166 + global_to_physical=global_to_physical,

167 + physical_to_global=physical_to_global,

168 + local_expert_global_ids=local_expert_global_ids,

169 )

170

171 return prepare_finalize

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+28 -2

@@ -85,6 +85,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

85 max_tokens_per_rank: int,

86 num_dispatchers: int,

87 use_fp8_dispatch: bool = False,

88 + global_to_physical: torch.Tensor | None = None,

89 + physical_to_global: torch.Tensor | None = None,

90 + local_expert_global_ids: torch.Tensor | None = None,

91 ):

92 super().__init__()

@@ -97,6 +100,17 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

100 self.handles: list[tuple | None] = [None, None]

101 self.num_dispatchers_ = num_dispatchers

102

103 + topk_indices_dtype = self.topk_indices_dtype()

104 +

105 + def _maybe_cast(tensor: torch.Tensor | None) -> torch.Tensor | None:

106 + if tensor is None or topk_indices_dtype is None:

107 + return tensor

108 + return tensor.to(dtype=topk_indices_dtype)

109 +

110 + self.global_to_physical = _maybe_cast(global_to_physical)

111 + self.physical_to_global = _maybe_cast(physical_to_global)

112 + self.local_expert_global_ids = _maybe_cast(local_expert_global_ids)

113 +

114 # We don't have enough information to determine if we should dispatch

115 # activation scales in a packed ue8m0 format during object construction

116 # time. This setting is handled by post_init_setup.

@@ -136,6 +150,16 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

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

151 return torch.int64

152

153 + def _map_global_to_physical_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:

154 + if self.global_to_physical is None:

155 + return topk_ids

156 + return self.global_to_physical[topk_ids]

157 +

158 + def _map_local_to_global_ids(self, expert_topk_ids: torch.Tensor) -> torch.Tensor:

159 + if self.local_expert_global_ids is None:

160 + return expert_topk_ids

161 + return self.local_expert_global_ids[expert_topk_ids]

162 +

163 def _do_quant(

164 self,

165 x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

@@ -226,9 +250,10 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

250 a1 = a1 * topk_weights.to(a1.dtype)

251

252 # Dispatch

253 + dispatch_topk_ids = self._map_global_to_physical_ids(topk_ids)

254 expert_x, expert_num_tokens, handle, _, hook = self.buffer.low_latency_dispatch(

255 a1,

231 - topk_ids,

256 + dispatch_topk_ids,

257 self.max_tokens_per_rank,

258 num_experts,

259 use_fp8=self.use_fp8_dispatch,

@@ -313,11 +338,12 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

338 # weights have already been applied.

339 combine_topk_weights = torch.ones_like(topk_weights)

340

341 + combine_topk_ids = self._map_global_to_physical_ids(topk_ids)

342 # TODO (varun) : Enable zero copy mode

343 dbo_maybe_run_recv_hook()

344 _, _, recv_hook = self.buffer.low_latency_combine(

345 fused_expert_output,

320 - topk_ids,

346 + combine_topk_ids,

347 combine_topk_weights,

348 handle,

349 async_finish=False,

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

+7 -2

@@ -50,10 +50,15 @@ class FusedMoEMethodBase(QuantizeMethodBase):

50 """

51 return False

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

53 + def maybe_make_prepare_finalize(

54 + self,

55 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

56 + ) -> FusedMoEPrepareAndFinalize | None:

57 from .all2all_utils import maybe_make_prepare_finalize

56 - return maybe_make_prepare_finalize(self.moe, self.moe_quant_config)

59 + return maybe_make_prepare_finalize(

60 + self.moe, self.moe_quant_config, routing_tables

61 + )

63 def select_gemm_impl(

64 self,

vllm/model_executor/layers/fused_moe/layer.py

+135 -22

@@ -5,7 +5,7 @@ from collections.abc import Callable, Iterable

5 from contextlib import nullcontext

6 from enum import Enum

7 from functools import partial

8 -from typing import Literal, get_args, overload

8 +from typing import Literal, cast, get_args, overload

10 import torch

11 import torch.nn.functional as F

@@ -192,6 +192,42 @@ def determine_expert_map(

192 return (local_num_experts, expert_map, expert_mask)

193

194

195 +def determine_expert_placement_strategy(

196 + expert_placement_strategy: ExpertPlacementStrategy,

197 + moe_parallel_config: FusedMoEParallelConfig,

198 + num_expert_group: int | None,

199 + num_redundant_experts: int,

200 + enable_eplb: bool,

201 +) -> ExpertPlacementStrategy:

202 + if expert_placement_strategy == "round_robin":

203 + round_robin_supported = (

204 + (num_expert_group is not None and num_expert_group > 1)

205 + and num_redundant_experts == 0

206 + and not enable_eplb

207 + )

208 +

209 + if not round_robin_supported:

210 + logger.warning(

211 + "Round-robin expert placement is only supported for "

212 + "models with multiple expert groups and no redundant "

213 + "experts. Falling back to linear expert placement."

214 + )

215 + return "linear"

216 + if (

217 + moe_parallel_config.use_all2all_kernels

218 + and not moe_parallel_config.use_deepep_ll_kernels

219 + ):

220 + logger.warning(

221 + "Round-robin expert placement currently only supports "

222 + "the DeepEP low-latency backend, but '%s' was configured. "

223 + "Falling back to linear expert placement.",

224 + moe_parallel_config.all2all_backend,

225 + )

226 + return "linear"

227 +

228 + return expert_placement_strategy

229 +

230 +

231 def get_compressed_expert_map(expert_map: torch.Tensor) -> str:

232 """

233 Compresses the expert map by removing any -1 entries.

@@ -400,6 +436,9 @@ class FusedMoE(CustomOp):

436 self.expert_load_view: torch.Tensor | None = None

437 self.logical_to_physical_map: torch.Tensor | None = None

438 self.logical_replica_count: torch.Tensor | None = None

439 + self.expert_placement_strategy: ExpertPlacementStrategy = (

440 + vllm_config.parallel_config.expert_placement_strategy

441 + )

442

443 # ROCm aiter shared experts fusion

444 self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled()

@@ -433,38 +472,27 @@ class FusedMoE(CustomOp):

472 "Redundant experts are only supported with EPLB."

473 )

474

436 - expert_placement_strategy = (

437 - vllm_config.parallel_config.expert_placement_strategy

475 + self.expert_placement_strategy = determine_expert_placement_strategy(

476 + expert_placement_strategy=self.expert_placement_strategy,

477 + moe_parallel_config=self.moe_parallel_config,

478 + num_expert_group=num_expert_group,

479 + num_redundant_experts=num_redundant_experts,

480 + enable_eplb=self.enable_eplb,

481 )

439 - if expert_placement_strategy == "round_robin":

440 - # TODO(Bruce): will support round robin expert placement with

441 - # EPLB enabled in the future.

442 - round_robin_supported = (

443 - (num_expert_group is not None and num_expert_group > 1)

444 - and num_redundant_experts == 0

445 - and not self.enable_eplb

446 - )

447 -

448 - if not round_robin_supported:

449 - logger.warning(

450 - "Round-robin expert placement is only supported for "

451 - "models with multiple expert groups and no redundant "

452 - "experts. Falling back to linear expert placement."

453 - )

454 - expert_placement_strategy = "linear"

482

483 self.expert_map: torch.Tensor | None

484 local_num_experts, expert_map, expert_mask = determine_expert_map(

485 ep_size=self.ep_size,

486 ep_rank=self.ep_rank,

487 global_num_experts=self.global_num_experts,

461 - expert_placement_strategy=expert_placement_strategy,

488 + expert_placement_strategy=self.expert_placement_strategy,

489 num_fused_shared_experts=self.num_fused_shared_experts,

490 return_expert_mask=self.rocm_aiter_fmoe_enabled,

491 )

492 self.local_num_experts = local_num_experts

493 self.register_buffer("expert_map", expert_map)

494 self.register_buffer("expert_mask", expert_mask)

495 + self._maybe_init_expert_routing_tables()

496 logger.info_once(

497 "[EP Rank %s/%s] Expert parallelism is enabled. Expert "

498 "placement strategy: %s. Local/global"

@@ -472,7 +500,7 @@ class FusedMoE(CustomOp):

500 " %s.",

501 self.ep_rank,

502 self.ep_size,

475 - expert_placement_strategy,

503 + self.expert_placement_strategy,

504 self.local_num_experts,

505 self.global_num_experts,

506 get_compressed_expert_map(self.expert_map),

@@ -621,7 +649,12 @@ class FusedMoE(CustomOp):

649 # should be safe to swap out the quant_method.

650 def maybe_init_modular_kernel(self) -> None:

651 self.ensure_moe_quant_config_init()

624 - prepare_finalize = self.quant_method.maybe_make_prepare_finalize()

652 + # routing_tables only needed for round-robin expert placement with

653 + # DeepEP all2all backend.

654 + routing_tables = self._maybe_init_expert_routing_tables()

655 + prepare_finalize = self.quant_method.maybe_make_prepare_finalize(

656 + routing_tables=routing_tables

657 + )

658 if prepare_finalize is not None:

659 logger.debug(

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

@@ -703,6 +736,84 @@ class FusedMoE(CustomOp):

736 # By default, router/gate is called before FusedMoE forward pass

737 return False

738

739 + def _maybe_init_expert_routing_tables(

740 + self,

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

742 + # Currently routing_tables only needed for round-robin expert placement

743 + # with DeepEP-ll all2all backend.

744 + if (

745 + self.expert_placement_strategy != "round_robin"

746 + or not self.use_deepep_ll_kernels

747 + ):

748 + return None

749 +

750 + if hasattr(self, "expert_global_to_physical"):

751 + return cast(

752 + tuple[torch.Tensor, torch.Tensor, torch.Tensor],

753 + (

754 + self.expert_global_to_physical,

755 + self.expert_physical_to_global,

756 + self.expert_local_to_global,

757 + ),

758 + )

759 +

760 + if self.expert_map is None:

761 + return None

762 +

763 + routing_tables = self.ensure_round_robin_expert_routing_tables(

764 + global_num_experts=self.global_num_experts,

765 + ep_size=self.ep_size,

766 + ep_rank=self.ep_rank,

767 + local_num_experts=self.local_num_experts,

768 + device=self.expert_map.device,

769 + )

770 +

771 + global_to_physical, physical_to_global, local_global = routing_tables

772 + self.register_buffer("expert_global_to_physical", global_to_physical)

773 + self.register_buffer("expert_physical_to_global", physical_to_global)

774 + self.register_buffer("expert_local_to_global", local_global)

775 +

776 + return routing_tables

777 +

778 + @staticmethod

779 + def ensure_round_robin_expert_routing_tables(

780 + global_num_experts: int,

781 + ep_size: int,

782 + ep_rank: int,

783 + local_num_experts: int,

784 + device: torch.device | None = None,

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

786 + device_kwargs = {"device": device} if device is not None else {}

787 + global_indices = torch.arange(

788 + global_num_experts, dtype=torch.long, **device_kwargs

789 + )

790 + owner = torch.remainder(global_indices, ep_size)

791 + local_index = torch.div(global_indices, ep_size, rounding_mode="floor")

792 + base = global_num_experts // ep_size

793 + remainder = global_num_experts % ep_size

794 + physical_offset = owner * base

795 + if remainder > 0:

796 + remainder_tensor = torch.tensor(

797 + remainder, dtype=torch.long, **device_kwargs

798 + )

799 + physical_offset = physical_offset + torch.minimum(owner, remainder_tensor)

800 +

801 + global_to_physical = physical_offset + local_index

802 + physical_to_global = torch.empty_like(global_to_physical)

803 + physical_to_global[global_to_physical] = global_indices

804 +

805 + local_global = torch.arange(

806 + ep_rank,

807 + global_num_experts,

808 + ep_size,

809 + dtype=torch.long,

810 + **device_kwargs,

811 + )

812 + if local_global.numel() != local_num_experts:

813 + local_global = local_global[:local_num_experts]

814 +

815 + return (global_to_physical, physical_to_global, local_global)

816 +

817 def update_expert_map(self):

818 # ep_size and ep_rank should already be updated

819 assert self.expert_map is not None

@@ -711,12 +822,14 @@ class FusedMoE(CustomOp):

822 ep_size=self.ep_size,

823 ep_rank=self.ep_rank,

824 global_num_experts=self.global_num_experts,

825 + expert_placement_strategy=self.expert_placement_strategy,

826 num_fused_shared_experts=self.num_fused_shared_experts,

827 return_expert_mask=self.rocm_aiter_fmoe_enabled,

828 )

829 self.local_num_experts = local_num_experts

830 self.register_buffer("expert_map", expert_map)

831 self.register_buffer("expert_mask", expert_mask)

832 + self._maybe_init_expert_routing_tables()

833 if self.aiter_fmoe_shared_expert_enabled:

834 self._init_aiter_shared_experts_topK_buffer(

835 vllm_config=get_current_vllm_config(),

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+5 -2

@@ -108,11 +108,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

108 def allow_inplace(self) -> bool:

109 return True

110

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

111 + def maybe_make_prepare_finalize(

112 + self,

113 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

114 + ) -> FusedMoEPrepareAndFinalize | None:

115 if self.rocm_aiter_moe_enabled:

116 return None

117 else:

115 - return super().maybe_make_prepare_finalize()

118 + return super().maybe_make_prepare_finalize(routing_tables)

119

120 def select_gemm_impl(

121 self,

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

+10 -4

@@ -380,11 +380,14 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

380 (layer.w2_input_global_scale), requires_grad=False

381 )

382

383 - def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

383 + def maybe_make_prepare_finalize(

384 + self,

385 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

386 + ) -> mk.FusedMoEPrepareAndFinalize | None:

387 if self.use_marlin:

388 return None

389 elif not self.allow_flashinfer:

387 - return super().maybe_make_prepare_finalize()

390 + return super().maybe_make_prepare_finalize(routing_tables)

391

392 prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(self.moe)

393 logger.debug_once("%s", prepare_finalize.__class__.__name__)

@@ -890,11 +893,14 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

893 layer.w2_weight_scale

894 )

895

893 - def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

896 + def maybe_make_prepare_finalize(

897 + self,

898 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

899 + ) -> mk.FusedMoEPrepareAndFinalize | None:

900 if self.use_marlin or self.rocm_aiter_moe_enabled:

901 return None

902 else:

897 - return super().maybe_make_prepare_finalize()

903 + return super().maybe_make_prepare_finalize(routing_tables)

904

905 def select_gemm_impl(

906 self,

vllm/model_executor/layers/quantization/fp8.py

+5 -2

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

1018 del layer.w13_input_scale

1019 del layer.w2_input_scale

1020

1021 - def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

1021 + def maybe_make_prepare_finalize(

1022 + self,

1023 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

1024 + ) -> mk.FusedMoEPrepareAndFinalize | None:

1025 if (

1026 self.rocm_aiter_moe_enabled

1027 or self.use_marlin

@@ -1039,7 +1042,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1042 logger.debug_once("%s", prepare_finalize.__class__.__name__)

1043 return prepare_finalize

1044 else:

1042 - return super().maybe_make_prepare_finalize()

1045 + return super().maybe_make_prepare_finalize(routing_tables)

1046

1047 def select_gemm_impl(

1048 self,

vllm/model_executor/layers/quantization/modelopt.py

+7 -3

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

373

374 def maybe_make_prepare_finalize(

375 self,

376 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

377 ) -> mk.FusedMoEPrepareAndFinalize | None:

378 # TRT LLM not supported with all2all yet.

379 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

@@ -384,7 +385,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

385 logger.debug_once("%s", prepare_finalize.__class__.__name__)

386 return prepare_finalize

387 else:

387 - return super().maybe_make_prepare_finalize()

388 + return super().maybe_make_prepare_finalize(routing_tables)

389

390 def select_gemm_impl(

391 self,

@@ -1179,7 +1180,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1180 " for ModelOptNvFp4FusedMoE."

1181 )

1182

1182 - def maybe_make_prepare_finalize(self) -> mk.FusedMoEPrepareAndFinalize | None:

1183 + def maybe_make_prepare_finalize(

1184 + self,

1185 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

1186 + ) -> mk.FusedMoEPrepareAndFinalize | None:

1187 if self.use_marlin or (

1188 self.allow_flashinfer

1189 and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

@@ -1196,7 +1200,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1200 logger.debug_once("%s", prepare_finalize.__class__.__name__)

1201 return prepare_finalize

1202 else:

1199 - return super().maybe_make_prepare_finalize()

1203 + return super().maybe_make_prepare_finalize(routing_tables)

1204

1205 def select_gemm_impl(

1206 self,