Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / expert_parallel / expert_parallel.py: 85%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Expert Parallelism distributed strategies. 

16 

17Provides token permutation helpers and four parallel styles that compose with 

18:class:`~hyper_parallel.core.expert_parallel.moe.GroupedExperts`: 

19 

20- :class:`BaseExpertParallel` — abstract base for EP strategies with 

21 all-to-all token dispatch/combine. 

22- :class:`ExpertParallel` — standard EP: each rank owns a shard of experts; 

23 tokens are routed via differentiable all-to-all. 

24- :class:`TensorParallel` — TP-only weight sharding for experts with no token 

25 dispatch; for use when EP degree = 1. 

26- :class:`ExpertTensorParallel` — combined EP + TP on a 2-D mesh ``[ep, tp]``; 

27 weights are doubly sharded, dispatch uses the EP sub-mesh. 

28""" 

29__all__ = [ 

30 "AllToAllTokenDispatcher", 

31 "DeredundencyTokenDispatcher", 

32 "BaseExpertParallel", 

33 "ExpertParallel", 

34 "TensorParallel", 

35 "ExpertTensorParallel", 

36] 

37 

38from abc import ABC, abstractmethod 

39from dataclasses import dataclass 

40from typing import Any, List, Optional, Tuple, Union 

41 

42from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

43from hyper_parallel.core.dtensor.dtensor import ( 

44 distribute_module, 

45 distribute_tensor, 

46 _distribute_module_iter_params, 

47 _distribute_module_new_parameter, 

48 _distribute_module_param_source, 

49 _distribute_module_set_param, 

50) 

51from hyper_parallel.core.dtensor.placement_types import Shard 

52from hyper_parallel.core.tensor_parallel.style import ParallelStyle 

53from hyper_parallel.platform import AsyncHandle, get_platform 

54 

55platform = get_platform() 

56Module = platform.Module 

57 

58 

59# --------------------------------------------------------------------------- 

60# Token permutation helpers 

61# --------------------------------------------------------------------------- 

62 

63def _generate_permute_indices( 

64 tokens_per_expert_group, 

65 experts_per_rank: int, 

66 num_ranks: int, 

67): 

68 """Generate permutation indices for rank-major → expert-major reordering. 

69 

70 After all-to-all, received tokens are laid out in rank-major order:: 

71 

72 [rank0·expert0 tokens | rank0·expert1 tokens | ... | 

73 rank1·expert0 tokens | rank1·expert1 tokens | ...] 

74 

75 Expert computation requires expert-major order:: 

76 

77 [all tokens for local expert 0 | all tokens for local expert 1 | ...] 

78 

79 Args: 

80 tokens_per_expert_group: 1-D integer tensor of shape 

81 ``[num_ranks * experts_per_rank]``. Entry ``[r * E + e]`` is the 

82 number of tokens received from rank ``r`` for local expert ``e``. 

83 experts_per_rank: Number of experts owned by each rank. 

84 num_ranks: EP degree (total number of ranks in the EP group). 

85 

86 Returns: 

87 Tuple of: 

88 

89 - ``permuted_indices``: 1-D long tensor of length 

90 ``total_received_tokens``. ``permuted_indices[i]`` is the source 

91 position in the rank-major buffer for destination position ``i`` in 

92 the expert-major buffer. 

93 - ``num_tokens_per_expert``: 1-D integer tensor of length 

94 ``experts_per_rank`` with the token count per local expert. 

95 """ 

96 counts = tokens_per_expert_group # [num_ranks * experts_per_rank] 

97 

98 # num_tokens_per_expert[e] = Σ_r counts[r * E + e] 

99 counts_2d = counts.view(num_ranks, experts_per_rank) # [R, E] 

100 num_tokens_per_expert = counts_2d.sum(dim=0) # [E] 

101 

102 # ``total`` must be a host int because ``arange`` needs a scalar size. 

103 # That single D2H drain is unavoidable. Everything else stays on 

104 # device — no per-block ``.item()`` in a loop. 

105 total = int(num_tokens_per_expert.sum()) 

106 if total == 0: 

107 return counts.new_zeros(0, dtype=counts.dtype), num_tokens_per_expert 

108 

109 # ---- Vectorized expert-major permutation, no host stalls ----------- 

110 # Source offsets in the rank-major receive buffer for each (r, e) block. 

111 src_offsets_rm = counts.cumsum(0) - counts # [R*E], starts of each block 

112 # Reorder src offsets to expert-major iteration order: block (e, r). 

113 src_offsets_em = ( 

114 src_offsets_rm.view(num_ranks, experts_per_rank).T.contiguous().view(-1) 

115 ) # [E*R] 

116 # Counts in expert-major iteration order. 

117 counts_em = counts_2d.T.contiguous().view(-1) # [E*R] 

118 

119 # ``repeat_interleave`` expands each block's src start to one entry per 

120 # token in that block — gives the source position of each output token. 

121 block_src_starts = src_offsets_em.repeat_interleave(counts_em) # [total] 

122 

123 # Destination block starts in expert-major order, then expanded. The 

124 # ``arange(total) - dst_block_starts_per_token`` produces 0..n-1 within 

125 # each block, i.e. the intra-block offset. 

126 dst_block_starts = counts_em.cumsum(0) - counts_em # [E*R] 

127 dst_block_starts_per_token = dst_block_starts.repeat_interleave(counts_em) 

128 intra = platform.arange(0, total, device=counts.device) - dst_block_starts_per_token 

129 

130 permuted_indices = (block_src_starts + intra).long() 

131 return permuted_indices, num_tokens_per_expert 

132 

133 

134def _permute(x, tokens_per_expert_group, ep_degree: int, num_local_experts: int): 

135 """Apply rank-major → expert-major permutation to routed tokens. 

136 

137 Args: 

138 x: Received token tensor of shape 

139 ``[sum(tokens_per_expert_group), *feature_dims]``. 

140 tokens_per_expert_group: 1-D integer tensor of shape 

141 ``[ep_degree * num_local_experts]`` (output of the first 

142 all-to-all that exchanges token counts). 

143 ep_degree: EP group size (number of ranks). 

144 num_local_experts: Number of experts owned by this rank. 

145 

146 Returns: 

147 Tuple of: 

148 

149 - ``original_shape``: shape of *x* before permutation. 

150 - ``permuted_x``: tokens reordered to expert-major layout. 

151 - ``permuted_indices``: permutation indices (needed for 

152 :func:`_unpermute`). 

153 - ``num_tokens_per_expert``: token count per local expert. 

154 """ 

155 original_shape = x.shape 

156 permuted_indices, num_tokens_per_expert = _generate_permute_indices( 

157 tokens_per_expert_group, num_local_experts, ep_degree 

158 ) 

159 # ``x[permuted_indices]`` works for empty indices too (returns a 

160 # shape-0 tensor with a real grad_fn). Avoid the early-return with 

161 # ``new_zeros`` which would produce a leaf tensor without grad_fn and 

162 # silently break autograd for ranks that happen to receive zero tokens. 

163 permuted_x = x[permuted_indices] 

164 return original_shape, permuted_x, permuted_indices, num_tokens_per_expert 

165 

166 

167def _unpermute(out, original_shape, permuted_indices): 

168 """Reverse the permutation applied by :func:`_permute`. 

169 

170 Args: 

171 out: Expert-major output tensor of shape 

172 ``[sum(num_tokens_per_expert), *feature_dims]``. 

173 original_shape: Shape before permutation (from :func:`_permute`). 

174 permuted_indices: Permutation indices from :func:`_permute`. 

175 

176 Returns: 

177 Token tensor restored to the rank-major layout received after 

178 all-to-all, with shape ``original_shape``. 

179 """ 

180 # ``result[permuted_indices] = out`` is a differentiable scatter that 

181 # also handles the empty-index case (no-op assignment, but autograd 

182 # still connects ``result`` back to ``out``). Do NOT short-circuit 

183 # with a bare ``new_zeros`` — that returns a leaf tensor without 

184 # grad_fn and the downstream combine a2a loses its backward path, 

185 # which manifests as "element 0 of tensors does not require grad". 

186 result = out.new_zeros(*original_shape) 

187 result[permuted_indices] = out 

188 return result 

189 

190 

191# --------------------------------------------------------------------------- 

192# DispatchContext — state shared between token dispatch and combine 

193# --------------------------------------------------------------------------- 

194 

195@dataclass 

196class DispatchContext: 

197 """Intermediate state between token dispatch and combine. 

198 

199 Stored in ``module._ep_dispatch_ctx`` for a single forward pass. 

200 This solves the instance sharing problem when the same ExpertParallel 

201 style object is applied to multiple layers: 

202 

203 Example problem (before this fix): 

204 ep_style = ExpertParallel() 

205 ep_style.apply(layer1.experts, mesh) # registers hooks 

206 ep_style.apply(layer2.experts, mesh) # reuses same ep_style 

207 

208 # During forward: 

209 # layer1.dispatch writes to ep_style._state_stack 

210 # layer2.dispatch pushes to same stack ← INTERLEAVING 

211 # layer1.combine pops wrong state (LIFO violation) 

212 

213 Solution: Store context per-module, not per-style-instance. 

214 

215 Built by :meth:`AllToAllTokenDispatcher.dispatch` and consumed by 

216 :meth:`AllToAllTokenDispatcher.combine`. The caller 

217 (e.g. :class:`ExpertParallel`) stores this on the module between the 

218 paired dispatch/combine calls. 

219 """ 

220 

221 input_splits: List[int] 

222 output_splits: List[int] 

223 input_shape: Tuple[int, ...] 

224 permuted_indices: Any 

225 

226 

227@dataclass 

228class DeredundencyDispatchContext(DispatchContext): 

229 """State shared by deredundency dispatch and combine. 

230 

231 The inherited split and permutation fields describe the inner-EP 

232 all-to-all. The extra fields describe the OEP shared token view and the 

233 whiteboard scatter used before the final reduce-scatter combine. 

234 """ 

235 

236 dispatch_indices: Optional[object] = None 

237 router_coeff: Optional[object] = None 

238 gathered_shape: Optional[tuple] = None 

239 oep_size: int = 1 

240 

241 

242 

243@dataclass(frozen=True) 

244class _DeredundencyMeshInfo: 

245 """Resolved mesh metadata for two-stage deredundency token exchange.""" 

246 

247 oep_group: object 

248 iep_group: object 

249 oep_size: int 

250 iep_size: int 

251 outer_rank: int 

252 inner_rank: int 

253 

254 

255def _get_deredundency_mesh_info(device_mesh: DeviceMesh) -> _DeredundencyMeshInfo: 

256 """Resolve ``oep`` / ``iep`` groups from a 1-D or 2-D EP mesh.""" 

257 ndim = getattr(device_mesh, "ndim", 1) 

258 if not isinstance(ndim, int): 

259 ndim = 1 

260 if ndim == 1: 

261 return _DeredundencyMeshInfo( 

262 oep_group=None, 

263 iep_group=device_mesh.get_group(), 

264 oep_size=1, 

265 iep_size=device_mesh.size(), 

266 outer_rank=0, 

267 inner_rank=device_mesh.get_local_rank(), 

268 ) 

269 if ndim != 2: 

270 raise ValueError( 

271 "DeredundencyTokenDispatcher expects a 1-D EP mesh or a 2-D " 

272 f"[oep, iep] EP mesh, but got ndim={ndim}." 

273 ) 

274 

275 mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None) or () 

276 oep_dim = mesh_dim_names.index("oep") if "oep" in mesh_dim_names else 0 

277 iep_dim = mesh_dim_names.index("iep") if "iep" in mesh_dim_names else 1 

278 if oep_dim == iep_dim: 

279 raise ValueError("DeredundencyTokenDispatcher requires distinct oep and iep mesh dimensions.") 

280 

281 return _DeredundencyMeshInfo( 

282 oep_group=device_mesh.get_group(oep_dim), 

283 iep_group=device_mesh.get_group(iep_dim), 

284 oep_size=device_mesh.size(oep_dim), 

285 iep_size=device_mesh.size(iep_dim), 

286 outer_rank=device_mesh.get_local_rank(oep_dim), 

287 inner_rank=device_mesh.get_local_rank(iep_dim), 

288 ) 

289 

290 

291def _generate_deredundency_dispatch_indices( 

292 tokens_per_expert_by_source, 

293 expert_start: int, 

294 iep_size: int, 

295 num_local_experts: int, 

296): 

297 """Generate gather-view indices ordered by IEP destination rank. 

298 

299 ``tokens_per_expert_by_source`` is shaped ``[oep_size, num_experts]`` and 

300 describes each source rank's expert-major routed buffer after the OEP 

301 all-gather. The returned indices select the current outer expert range 

302 and order it as ``[iep_dst, local_expert, oep_source]`` so each IEP 

303 destination chunk keeps local-expert blocks contiguous for the later 

304 rank-major → expert-major permutation. 

305 """ 

306 oep_size = tokens_per_expert_by_source.shape[0] 

307 experts_per_outer = iep_size * num_local_experts 

308 expert_end = expert_start + experts_per_outer 

309 

310 source_totals = tokens_per_expert_by_source.sum(dim=1) 

311 source_offsets = source_totals.cumsum(0) - source_totals 

312 expert_offsets = ( 

313 tokens_per_expert_by_source.cumsum(dim=1) 

314 - tokens_per_expert_by_source 

315 + source_offsets.view(oep_size, 1) 

316 ) 

317 

318 selected_counts = tokens_per_expert_by_source[:, expert_start:expert_end].view( 

319 oep_size, iep_size, num_local_experts, 

320 ) 

321 selected_offsets = expert_offsets[:, expert_start:expert_end].view( 

322 oep_size, iep_size, num_local_experts, 

323 ) 

324 counts_by_destination = selected_counts.permute(1, 2, 0).contiguous() 

325 offsets_by_destination = selected_offsets.permute(1, 2, 0).contiguous() 

326 

327 block_counts = counts_by_destination.view(-1) 

328 token_counts_by_destination_expert = selected_counts.sum(dim=0).contiguous().view(-1) 

329 total = int(block_counts.sum()) 

330 if total == 0: 

331 return block_counts.new_zeros(0, dtype=block_counts.dtype).long(), token_counts_by_destination_expert 

332 

333 block_starts = offsets_by_destination.view(-1).repeat_interleave(block_counts) 

334 block_offsets = block_counts.cumsum(0) - block_counts 

335 block_offsets_per_token = block_offsets.repeat_interleave(block_counts) 

336 intra = platform.arange(0, total, device=tokens_per_expert_by_source.device) - block_offsets_per_token 

337 

338 return (block_starts + intra).long(), token_counts_by_destination_expert 

339 

340 

341def _scale_by_router_coeff(tokens, router_coeff): 

342 """Scale routed expert outputs by optional router coefficients.""" 

343 if router_coeff is None: 

344 return tokens 

345 if router_coeff.shape[0] != tokens.shape[0]: 

346 raise ValueError( 

347 "router_coeff length must match routed token count, got " 

348 f"{router_coeff.shape[0]} and {tokens.shape[0]}." 

349 ) 

350 coeff = router_coeff 

351 if len(coeff.shape) == 1 and len(tokens.shape) > 1: 

352 coeff = coeff.reshape((-1,) + (1,) * (len(tokens.shape) - 1)) 

353 return tokens * coeff 

354 

355 

356def _scatter_add_first_dim(src, indices, output_shape): 

357 """Scatter-add rows of ``src`` into a zero tensor along dim 0.""" 

358 result = src.new_zeros(*output_shape) 

359 if len(src.shape) == 1: 

360 scatter_indices = indices 

361 else: 

362 scatter_indices = indices.reshape((-1,) + (1,) * (len(src.shape) - 1)).expand( 

363 -1, *src.shape[1:], 

364 ) 

365 if hasattr(result, "scatter_add"): 

366 return result.scatter_add(0, scatter_indices, src) 

367 if hasattr(result, "index_add"): 

368 return result.index_add(0, indices, src) 

369 raise RuntimeError( 

370 "DeredundencyTokenDispatcher.combine requires tensor scatter_add or " 

371 "index_add support for exdispatch_idx accumulation." 

372 ) 

373 

374 

375class _DeredundencyCombineHandle(AsyncHandle): 

376 """Async handle that finishes deredundency combine post-processing.""" 

377 

378 def __init__( 

379 self, 

380 async_tensor: object, 

381 mesh_info: _DeredundencyMeshInfo, 

382 ctx: DeredundencyDispatchContext, 

383 ) -> None: 

384 super().__init__(async_tensor) 

385 self._mesh_info = mesh_info 

386 self._ctx = ctx 

387 self._combined: Optional[object] = None 

388 

389 def wait(self) -> object: 

390 """Wait for IEP a2a, then finish OEP scatter/reduce combine once.""" 

391 if self._combined is None: 

392 outer_output = super().wait() 

393 weighted_output = _scale_by_router_coeff(outer_output, self._ctx.router_coeff) 

394 combine_whiteboard = _scatter_add_first_dim( 

395 weighted_output, 

396 self._ctx.dispatch_indices, 

397 self._ctx.gathered_shape, 

398 ) 

399 if self._ctx.oep_size == 1: 

400 self._combined = combine_whiteboard 

401 else: 

402 self._combined = platform.differentiable_reduce_scatter( 

403 combine_whiteboard, 

404 self._ctx.oep_size, 

405 0, 

406 "sum", 

407 self._mesh_info.oep_group, 

408 ) 

409 return self._combined 

410 

411 

412# --------------------------------------------------------------------------- 

413# BaseExpertParallel — abstract base for all-to-all EP strategies 

414# --------------------------------------------------------------------------- 

415 

416class BaseExpertParallel(ParallelStyle, ABC): 

417 """Abstract base class for Expert Parallel strategies with token dispatch. 

418 

419 Subclasses implement :meth:`_partition_fn`, :meth:`_token_dispatch`, and 

420 :meth:`_token_combine`; this class wires them into :func:`distribute_module`. 

421 """ 

422 

423 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

424 """Apply EP sharding and dispatch/combine hooks to *module*. 

425 

426 Args: 

427 module: A :class:`~hyper_parallel.core.expert_parallel.moe.GroupedExperts` 

428 instance to shard. 

429 device_mesh: Device mesh for this EP strategy. 

430 

431 Returns: 

432 The module with distributed parameters and dispatch/combine hooks. 

433 """ 

434 return distribute_module( 

435 module, 

436 device_mesh, 

437 self._partition_fn, 

438 self._token_dispatch, 

439 self._token_combine, 

440 ) 

441 

442 @abstractmethod 

443 def _partition_fn( 

444 self, name: str, module: Module, device_mesh: DeviceMesh 

445 ) -> None: 

446 """Shard module parameters according to this strategy. 

447 

448 Args: 

449 name: Submodule name. 

450 module: The module whose parameters are being sharded. 

451 device_mesh: Device mesh for this EP strategy. 

452 """ 

453 

454 @abstractmethod 

455 def _token_dispatch(self, module: Module, inputs, device_mesh: DeviceMesh): 

456 """Pre-hook: route input tokens to their assigned ranks. 

457 

458 Args: 

459 module: The ``GroupedExperts`` module. 

460 inputs: Forward inputs tuple. 

461 device_mesh: Device mesh for this EP strategy. 

462 

463 Returns: 

464 Transformed inputs for local expert computation. 

465 """ 

466 

467 @abstractmethod 

468 def _token_combine(self, module: Module, routed_output, device_mesh: DeviceMesh): 

469 """Post-hook: gather expert outputs back to the originating ranks. 

470 

471 Args: 

472 module: The ``GroupedExperts`` module. 

473 routed_output: Expert output tensor in expert-major order. 

474 device_mesh: Device mesh for this EP strategy. 

475 

476 Returns: 

477 Token tensor in the original token-major layout. 

478 """ 

479 

480 

481# --------------------------------------------------------------------------- 

482# AllToAllTokenDispatcher — token dispatch/combine via all-to-all 

483# --------------------------------------------------------------------------- 

484 

485class AllToAllTokenDispatcher: 

486 """Token dispatch and combine via all-to-all for expert parallelism. 

487 

488 Provides :meth:`dispatch` and :meth:`combine` as static methods that 

489 receive and return a :class:`DispatchContext` object. This decouples 

490 the all-to-all token routing logic from the parallel style class so 

491 that it can be reused or tested independently. 

492 

493 Callers (e.g. :class:`ExpertParallel`) are responsible for storing the 

494 context between the paired dispatch/combine calls. 

495 """ 

496 

497 @staticmethod 

498 def dispatch(module: Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple: 

499 """Dispatch tokens to their assigned ranks via all-to-all. 

500 

501 Called as an ``input_fn`` hook by :func:`distribute_module`. Receives 

502 the module's forward inputs and returns transformed inputs. 

503 

504 Args: 

505 module: The ``GroupedExperts`` module. 

506 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` where 

507 ``routed_input`` has shape ``[total_tokens, dim]`` and 

508 ``num_tokens_per_expert`` has shape ``[num_experts]``. 

509 device_mesh: EP device mesh (1-D). 

510 

511 Returns: 

512 Tuple ``(permuted_local_input, local_token_counts, ctx)`` — 

513 the first two elements are the transformed inputs for local 

514 expert computation; *ctx* is a :class:`DispatchContext` 

515 carrying the updated state to be stored by the caller. 

516 """ 

517 del module # module unused, kept for API consistency 

518 routed_input, num_tokens_per_expert = inputs[0], inputs[1] 

519 ep_group = device_mesh.get_group() 

520 ep_size = device_mesh.size() 

521 num_local_experts = num_tokens_per_expert.shape[0] // ep_size 

522 

523 # --- Step 1: exchange token counts (no gradient needed) --- 

524 # Each rank needs to know how many tokens it will receive from every 

525 # other rank (for each local expert). Uses ``async_op=True`` + an 

526 # explicit ``handle.wait()`` rather than ``async_op=False`` because 

527 # the implicit cross-stream sync is NCCL-only; on HCCL the compute 

528 # stream may read ``counts_out`` before the collective write is 

529 # visible, producing garbage values that blow up the downstream 

530 # ``torch.empty(sum(output_splits), ...)`` allocation. 

531 counts_out, handle = platform.all_to_all_single( 

532 num_tokens_per_expert, 

533 output_shape=[num_tokens_per_expert.shape[0]], 

534 group=ep_group, 

535 async_op=True, 

536 ) 

537 if handle is not None: 

538 handle.wait() 

539 # counts_out shape: [ep_size * num_local_experts] 

540 # counts_out[r * num_local_experts + e] = tokens from rank r for expert e 

541 

542 # --- Step 2: compute input / output splits --- 

543 # input_splits[r] = tokens this rank sends to rank r 

544 # output_splits[r] = tokens this rank receives from rank r 

545 # Reshape to [ep_size, num_local_experts] and sum per rank on device; 

546 # a single ``tolist()`` drains the rank-sum vector to host, replacing 

547 # ``2 * ep_size`` scalar ``int()`` D2H syncs with 2. 

548 input_splits = num_tokens_per_expert.view(ep_size, num_local_experts).sum(dim=1).tolist() 

549 output_splits = counts_out.view(ep_size, num_local_experts).sum(dim=1).tolist() 

550 

551 # --- Step 3: exchange actual tokens (differentiable) --- 

552 dispatched = platform.differentiable_all_to_all_single( 

553 routed_input, input_splits, output_splits, group=ep_group, 

554 ) 

555 

556 # --- Step 4: rank-major → expert-major permutation --- 

557 input_shape, permuted, permuted_indices, local_counts = _permute( 

558 dispatched, counts_out, ep_size, num_local_experts 

559 ) 

560 

561 # Build dispatch context for combine step. 

562 # Caller (e.g., ExpertParallel._token_dispatch) is responsible for storing 

563 # this context and passing it to combine(). This decouples dispatch/combine 

564 # from module state and solves the instance sharing problem. 

565 ctx = DispatchContext( 

566 input_splits=input_splits, 

567 output_splits=output_splits, 

568 input_shape=input_shape, 

569 permuted_indices=permuted_indices, 

570 ) 

571 

572 return permuted, local_counts, ctx 

573 

574 @staticmethod 

575 def combine(module: Module, routed_output: object, device_mesh: DeviceMesh, ctx: DispatchContext) -> object: 

576 """Gather expert outputs back to the originating ranks via all-to-all. 

577 

578 Called as an ``output_fn`` hook by :func:`distribute_module`. 

579 Receives dispatch context from the caller (previously returned by dispatch). 

580 

581 Args: 

582 module: The ``GroupedExperts`` module (unused, for API consistency). 

583 routed_output: Expert output tensor in expert-major order, 

584 shape ``[sum(local_counts), dim]``. 

585 device_mesh: EP device mesh (1-D). 

586 ctx: :class:`DispatchContext` previously returned by 

587 :meth:`dispatch`. 

588 

589 Returns: 

590 Token tensor in the original token-major layout, 

591 shape ``[sum(input_splits), dim]``. 

592 """ 

593 del module # module not used, kept for API consistency 

594 ep_group = device_mesh.get_group() 

595 

596 # expert-major → rank-major 

597 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices) 

598 

599 # reverse all-to-all (output/input splits are swapped) 

600 combined = platform.differentiable_all_to_all_single( 

601 unpermuted, 

602 ctx.output_splits, # was output, now becomes input 

603 ctx.input_splits, # was input, now becomes output 

604 group=ep_group, 

605 ) 

606 return combined 

607 

608 @staticmethod 

609 def combine_start(routed_output, device_mesh, ctx): 

610 """Launch async combine all-to-all without waiting for completion. 

611 

612 Splits the combine into two phases so that the caller can overlap 

613 the a2a communication with independent computation (e.g. a shared 

614 expert forward pass). The caller must later call 

615 :meth:`combine_wait` or ``handle.wait()`` to obtain the final 

616 result. 

617 

618 Step 1 (synchronous, local): expert-major → rank-major unpermute. 

619 Step 2 (asynchronous, cross-rank): reverse all-to-all. 

620 

621 Args: 

622 routed_output: Expert output tensor in expert-major order, 

623 shape ``[sum(local_counts), dim]``. 

624 device_mesh: EP device mesh (1-D). 

625 ctx: :class:`DispatchContext` previously returned by 

626 :meth:`dispatch`. 

627 

628 Returns: 

629 :class:`AsyncHandle` carrying the state needed by 

630 :meth:`combine_wait`. 

631 """ 

632 ep_group = device_mesh.get_group() 

633 

634 # expert-major → rank-major (local, no communication) 

635 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices) 

636 

637 # async reverse all-to-all (output/input splits are swapped) 

638 combined_async = platform.differentiable_all_to_all_single_async( 

639 unpermuted, 

640 ctx.output_splits, 

641 ctx.input_splits, 

642 group=ep_group, 

643 ) 

644 

645 return AsyncHandle(combined_async) 

646 

647 @staticmethod 

648 def combine_wait(handle): 

649 """Wait for the async combine all-to-all to complete. 

650 

651 Args: 

652 handle: :class:`AsyncHandle` returned by :meth:`combine_start`. 

653 

654 Returns: 

655 Combined tensor in the original token-major layout. 

656 """ 

657 return handle.wait() 

658 

659 

660# --------------------------------------------------------------------------- 

661# DeredundencyTokenDispatcher — token dispatch via OEP all-gather + IEP all-to-all 

662# --------------------------------------------------------------------------- 

663 

664class DeredundencyTokenDispatcher: 

665 """Token dispatch/combine via OEP all-gather plus IEP all-to-all. 

666 

667 This dispatcher keeps the same public contract as 

668 :class:`AllToAllTokenDispatcher`, but decomposes the global EP all-to-all 

669 into the deredundency flow described in 

670 ``docs/moe_alltoall_deredundency_token_permutation.md``: 

671 

672 1. Form a shared token/count view across the OEP group. 

673 2. Select only the current outer expert range. 

674 3. Send selected tokens to concrete local-expert ranks inside the IEP 

675 group. 

676 4. Sort received tokens into local expert-major order. 

677 

678 For a 2-D mesh, dimension ``"oep"`` / ``0`` is the outer group and 

679 ``"iep"`` / ``1`` is the inner group. A 1-D mesh is treated as 

680 ``oep_size == 1`` and degenerates to the standard all-to-all data flow. 

681 """ 

682 

683 @staticmethod 

684 def _oep_gather_for_dispatch( 

685 num_tokens_per_expert, 

686 routed_input, 

687 router_coeff, 

688 mesh_info: _DeredundencyMeshInfo, 

689 ) -> tuple: 

690 """All-gather token counts and routed input across the OEP group. 

691 

692 Args: 

693 num_tokens_per_expert: Token count per expert ``[num_experts]``. 

694 routed_input: Routed token tensor ``[total_tokens, dim]``. 

695 router_coeff: Optional router coefficients ``[total_tokens]``. 

696 mesh_info: Resolved OEP/IEP mesh descriptor. 

697 

698 Returns: 

699 Tuple ``(gathered_counts, gathered_routed, gathered_router_coeff)`` 

700 where ``gathered_counts`` has shape ``[oep_size, num_experts]``, 

701 ``gathered_routed`` has shape ``[oep_size * total_tokens, dim]``, 

702 and ``gathered_router_coeff`` is the gathered coefficients or None. 

703 

704 Raises: 

705 ValueError: If routed token counts differ across OEP ranks. 

706 """ 

707 if mesh_info.oep_size == 1: 

708 gathered_counts = num_tokens_per_expert.view(1, num_tokens_per_expert.shape[0]) 

709 return gathered_counts, routed_input, router_coeff 

710 

711 gathered_counts, handle = platform.all_gather_single( 

712 num_tokens_per_expert, 

713 output_shape=[mesh_info.oep_size * num_tokens_per_expert.shape[0]], 

714 group=mesh_info.oep_group, 

715 async_op=True, 

716 ) 

717 if handle is not None: 

718 handle.wait() 

719 gathered_counts = gathered_counts.view(mesh_info.oep_size, num_tokens_per_expert.shape[0]) 

720 source_token_totals = gathered_counts.sum(dim=1).tolist() 

721 if any(total != routed_input.shape[0] for total in source_token_totals): 

722 raise ValueError( 

723 "DeredundencyTokenDispatcher requires equal routed token " 

724 "counts within each OEP group because the shared token view " 

725 f"uses all-gather, got totals {source_token_totals}." 

726 ) 

727 gathered_routed = platform.differentiable_all_gather_concat( 

728 routed_input, mesh_info.oep_group, mesh_info.oep_size, 0, 

729 ) 

730 if router_coeff is None: 

731 gathered_router_coeff = None 

732 else: 

733 gathered_router_coeff = platform.differentiable_all_gather_concat( 

734 router_coeff, mesh_info.oep_group, mesh_info.oep_size, 0, 

735 ) 

736 return gathered_counts, gathered_routed, gathered_router_coeff 

737 

738 @staticmethod 

739 def dispatch(module: Module, inputs: tuple, device_mesh: DeviceMesh) -> tuple: 

740 """Dispatch tokens using OEP all-gather and IEP all-to-all. 

741 

742 Args: 

743 module: The ``GroupedExperts`` module (unused here). 

744 inputs: Tuple ``(routed_input, num_tokens_per_expert)`` where 

745 ``routed_input`` has shape ``[total_tokens, dim]`` and 

746 ``num_tokens_per_expert`` has shape ``[num_experts]``. 

747 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh. 

748 

749 Returns: 

750 Tuple ``(permuted_local_input, local_token_counts, ctx)`` with the 

751 same meaning as :meth:`AllToAllTokenDispatcher.dispatch`. 

752 

753 Raises: 

754 ValueError: If the expert count is not divisible by the full EP 

755 size represented by the deredundency mesh. 

756 """ 

757 del module 

758 routed_input, num_tokens_per_expert = inputs[0], inputs[1] 

759 router_coeff = inputs[2] if len(inputs) > 2 else None 

760 if router_coeff is not None and router_coeff.shape[0] != routed_input.shape[0]: 

761 raise ValueError( 

762 "router_coeff length must match routed_input token count, got " 

763 f"{router_coeff.shape[0]} and {routed_input.shape[0]}." 

764 ) 

765 mesh_info = _get_deredundency_mesh_info(device_mesh) 

766 ep_size = mesh_info.oep_size * mesh_info.iep_size 

767 if num_tokens_per_expert.shape[0] % ep_size != 0: 

768 raise ValueError( 

769 "num_tokens_per_expert length must be divisible by the full " 

770 f"EP size {ep_size}, got {num_tokens_per_expert.shape[0]}." 

771 ) 

772 num_local_experts = num_tokens_per_expert.shape[0] // ep_size 

773 experts_per_outer = mesh_info.iep_size * num_local_experts 

774 expert_start = mesh_info.outer_rank * experts_per_outer 

775 

776 gathered_counts, gathered_routed, gathered_router_coeff = ( 

777 DeredundencyTokenDispatcher._oep_gather_for_dispatch( 

778 num_tokens_per_expert, routed_input, router_coeff, mesh_info, 

779 ) 

780 ) 

781 

782 dispatch_indices, node_counts_per_expert = _generate_deredundency_dispatch_indices( 

783 gathered_counts, 

784 expert_start, 

785 mesh_info.iep_size, 

786 num_local_experts, 

787 ) 

788 iep_input_splits = node_counts_per_expert.view(mesh_info.iep_size, num_local_experts).sum(dim=1).tolist() 

789 

790 iep_counts_out, handle = platform.all_to_all_single( 

791 node_counts_per_expert, 

792 output_shape=[node_counts_per_expert.shape[0]], 

793 group=mesh_info.iep_group, 

794 async_op=True, 

795 ) 

796 if handle is not None: 

797 handle.wait() 

798 iep_output_splits = iep_counts_out.view(mesh_info.iep_size, num_local_experts).sum(dim=1).tolist() 

799 

800 outer_routed_input = gathered_routed[dispatch_indices] 

801 outer_router_coeff = ( 

802 None if gathered_router_coeff is None else gathered_router_coeff[dispatch_indices] 

803 ) 

804 dispatched = platform.differentiable_all_to_all_single( 

805 outer_routed_input, 

806 iep_input_splits, 

807 iep_output_splits, 

808 group=mesh_info.iep_group, 

809 ) 

810 

811 input_shape, permuted, permuted_indices, local_counts = _permute( 

812 dispatched, iep_counts_out, mesh_info.iep_size, num_local_experts, 

813 ) 

814 ctx = DeredundencyDispatchContext( 

815 input_splits=iep_input_splits, 

816 output_splits=iep_output_splits, 

817 input_shape=input_shape, 

818 permuted_indices=permuted_indices, 

819 dispatch_indices=dispatch_indices, 

820 router_coeff=outer_router_coeff, 

821 gathered_shape=gathered_routed.shape, 

822 oep_size=mesh_info.oep_size, 

823 ) 

824 return permuted, local_counts, ctx 

825 

826 @staticmethod 

827 def combine(module: Module, routed_output: object, device_mesh: DeviceMesh, 

828 ctx: DeredundencyDispatchContext) -> object: 

829 """Gather expert outputs back to the originating ranks. 

830 

831 Args: 

832 module: The ``GroupedExperts`` module (unused). 

833 routed_output: Expert output tensor in expert-major order. 

834 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh. 

835 ctx: Context returned by :meth:`dispatch`. 

836 

837 Returns: 

838 Token tensor in the original source-rank routed order. 

839 """ 

840 del module 

841 mesh_info = _get_deredundency_mesh_info(device_mesh) 

842 DeredundencyTokenDispatcher._validate_combine_mesh(mesh_info, ctx) 

843 

844 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices) 

845 outer_output = platform.differentiable_all_to_all_single( 

846 unpermuted, 

847 ctx.output_splits, 

848 ctx.input_splits, 

849 group=mesh_info.iep_group, 

850 ) 

851 

852 weighted_output = _scale_by_router_coeff(outer_output, ctx.router_coeff) 

853 combine_whiteboard = _scatter_add_first_dim( 

854 weighted_output, ctx.dispatch_indices, ctx.gathered_shape, 

855 ) 

856 if ctx.oep_size == 1: 

857 return combine_whiteboard 

858 

859 return platform.differentiable_reduce_scatter( 

860 combine_whiteboard, 

861 ctx.oep_size, 

862 0, 

863 "sum", 

864 mesh_info.oep_group, 

865 ) 

866 

867 @staticmethod 

868 def _validate_combine_mesh( 

869 mesh_info: _DeredundencyMeshInfo, 

870 ctx: DeredundencyDispatchContext, 

871 ) -> None: 

872 """Validate that dispatch context and combine mesh are compatible.""" 

873 if mesh_info.oep_size != ctx.oep_size: 

874 raise ValueError( 

875 "DeredundencyTokenDispatcher.combine received a context for " 

876 f"oep_size={ctx.oep_size}, but the mesh resolves to oep_size={mesh_info.oep_size}." 

877 ) 

878 

879 @staticmethod 

880 def combine_start( 

881 routed_output: object, 

882 device_mesh: DeviceMesh, 

883 ctx: DeredundencyDispatchContext, 

884 ) -> AsyncHandle: 

885 """Launch async IEP combine all-to-all and defer deredundency post-processing. 

886 

887 The local expert-major → rank-major unpermute is performed 

888 synchronously. The reverse IEP all-to-all is launched asynchronously, 

889 and :meth:`combine_wait` finishes router weighting, whiteboard 

890 scatter-add, and optional OEP reduce-scatter after the async output is 

891 materialised. 

892 

893 Args: 

894 routed_output: Expert output tensor in expert-major order. 

895 device_mesh: 1-D EP mesh or 2-D ``[oep, iep]`` EP mesh. 

896 ctx: Context returned by :meth:`dispatch`. 

897 

898 Returns: 

899 :class:`AsyncHandle` carrying the pending IEP a2a and deredundency 

900 combine state. 

901 """ 

902 mesh_info = _get_deredundency_mesh_info(device_mesh) 

903 DeredundencyTokenDispatcher._validate_combine_mesh(mesh_info, ctx) 

904 

905 unpermuted = _unpermute(routed_output, ctx.input_shape, ctx.permuted_indices) 

906 outer_output_async = platform.differentiable_all_to_all_single_async( 

907 unpermuted, 

908 ctx.output_splits, 

909 ctx.input_splits, 

910 group=mesh_info.iep_group, 

911 ) 

912 return _DeredundencyCombineHandle(outer_output_async, mesh_info, ctx) 

913 

914 @staticmethod 

915 def combine_wait(handle: AsyncHandle) -> object: 

916 """Wait for async deredundency combine and return the final tensor. 

917 

918 Args: 

919 handle: :class:`AsyncHandle` returned by :meth:`combine_start`. 

920 

921 Returns: 

922 Token tensor in the original source-rank routed order. 

923 """ 

924 return handle.wait() 

925 

926 

927_TOKEN_DISPATCHERS = { 

928 "all_to_all": AllToAllTokenDispatcher, 

929 "deredundency": DeredundencyTokenDispatcher, 

930} 

931 

932 

933def _resolve_token_dispatcher(token_dispatcher: str): 

934 """Resolve a token dispatcher name to its implementation class.""" 

935 try: 

936 return _TOKEN_DISPATCHERS[token_dispatcher] 

937 except KeyError as exc: 

938 supported = "', '".join(sorted(_TOKEN_DISPATCHERS)) 

939 raise ValueError( 

940 f"token_dispatcher must be one of '{supported}', got {token_dispatcher!r}." 

941 ) from exc 

942 

943 

944def _get_flattened_ep_mesh(device_mesh: DeviceMesh) -> DeviceMesh: 

945 """Return a 1-D EP mesh, flattening a 2-D deredundency mesh if needed.""" 

946 if getattr(device_mesh, "ndim", 1) == 1: 

947 return device_mesh 

948 mesh_dim_names = getattr(device_mesh, "mesh_dim_names", None) or () 

949 if "ep" in mesh_dim_names or "ep" in device_mesh.get_flatten_mapping(): 

950 return device_mesh["ep"] 

951 if set(mesh_dim_names) == {"oep", "iep"}: 

952 return device_mesh.flatten("ep") 

953 raise ValueError( 

954 "Deredundency ExpertParallel expects a 1-D EP mesh or a 2-D " 

955 "[oep, iep] mesh when partitioning expert weights." 

956 ) 

957 

958 

959# --------------------------------------------------------------------------- 

960# ExpertParallel — standard all-to-all EP 

961# --------------------------------------------------------------------------- 

962 

963class ExpertParallel(BaseExpertParallel): 

964 """Expert Parallel: shard experts across ranks via all-to-all token routing. 

965 

966 Applies :meth:`apply` to a :class:`GroupedExperts` module: 

967 

968 1. **Partition** — distributes expert weights on dim 0 (``Shard(0)``) so 

969 each rank holds ``num_experts // ep_degree`` local experts. 

970 2. **Token dispatch** (forward pre-hook) — two-step all-to-all: 

971 a. Exchange token counts (non-differentiable). 

972 b. Exchange actual tokens (differentiable, gradient flows back). 

973 Followed by rank-major → expert-major permutation. 

974 3. **Token combine** (forward post-hook) — expert-major → rank-major 

975 unpermute, then reverse all-to-all (differentiable). 

976 

977 All collectives use ``platform.differentiable_all_to_all_single`` / 

978 ``platform.all_to_all_single`` — no direct ``torch.distributed`` calls. 

979 

980 The token dispatcher is selectable. ``"all_to_all"`` uses 

981 :class:`AllToAllTokenDispatcher`; ``"deredundency"`` uses 

982 :class:`DeredundencyTokenDispatcher`.