Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / experts.py: 86%
395 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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# ============================================================================
16"""expert_parallel.experts: EP expert execution and binding.
18- Routed-experts pipeline (``ep_routed_forward``, 05 §6.4.8): router (local
19 chunk) -> a2a dispatch (extended EP group) -> local SwiGLU (complete
20 expert weights, no internal communication) -> a2a combine -> weighted
21 aggregation. SP-in -> SP-out, all communication cohesive inside. NO
22 shared-expert/gate branch — that composition is the caller's job.
23- Expert entry point (``bind_local_expert_forward`` /
24 ``resolve_swiglu_weights``): installs ``experts.forward`` (via the forward
25 rewriter's bound-forward install point, 05 §15.2.3) so nested FSDP hooks
26 unshard/reshard around the local SwiGLU.
27- Interface helpers (``require_attrs`` / ``describe_moe_module``):
28 build-time interface assertions with teaching errors, and a structural
29 diagnostic for mapping a concrete MoE module to an archetype.
31Nothing in this module probes model structure with getattr fallback chains.
32Split out of components/distributed/ep_utils.py in stage 4e.
33"""
35import logging
36import math
37import os
38from dataclasses import dataclass
39from typing import Any, Callable, NamedTuple, Optional
40import torch
41import torch.distributed as dist
42import torch.nn.functional as F
44from hyper_parallel.components.functional.npu_grouped_swiglu import (
45 npu_grouped_swiglu,
46)
47from hyper_parallel.distributed._builder.forward_rewriter import (
48 _install_bound_forward,
49)
50from hyper_parallel.core.expert_parallel.static_splits import (
51 get_static_plan,
52 probe_mode,
53 static_plan_key,
54 store_static_plan,
55)
56from hyper_parallel.distributed.expert_parallel.routing import (
57 _resolve_capacity_order,
58 apply_capacity_limit,
59)
60from hyper_parallel.distributed.expert_parallel.collectives import (
61 ep_all_to_all,
62 ep_all_to_all_async,
63 wait_ep_all_to_all,
64)
65from hyper_parallel.distributed.expert_parallel.routing import (
66 MOE_ROUTER_ADAPTERS,
67)
70logger = logging.getLogger(__name__)
73@dataclass(frozen=True)
74class _EPDispatch:
75 """Prepared tensors and split sizes for one routed expert exchange."""
77 source_indices: torch.Tensor
78 expert_weights: torch.Tensor
79 dispatch_order: torch.Tensor
80 states: torch.Tensor
81 expert_indices: torch.Tensor
82 send_counts: list[int]
83 receive_counts: list[int]
86def resolve_swiglu_weights(
87 experts: Any,
88) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
89 """Resolve SwiGLU weights from the stacked holder (three layouts).
91 Returns (w_gate, w_up, w_down):
92 - Separate naming: gate_proj/up_proj/down_proj or w1/w3/w2 (all
93 [E, I, H]/[E, H, I]);
94 - **fused layout** (after the HF 2025 refactor, D-11): gate_up_proj
95 [E, 2I, H] + down_proj [E, H, I] -> returns (gate_up_proj, None,
96 down_proj); w_up=None marks the fused case (the compute side chunks
97 out gate/up).
98 """
99 w_fused = getattr(experts, "gate_up_proj", None)
100 if w_fused is None:
101 w_fused = getattr(experts, "gate_and_up_projs", None)
102 w_down = getattr(experts, "down_proj", None)
103 w_down = w_down if w_down is not None else getattr(experts, "down_projs", None)
104 w_down = w_down if w_down is not None else getattr(experts, "w2", None)
105 if w_fused is not None and w_down is not None:
106 return w_fused, None, w_down
108 w_gate = getattr(experts, "gate_proj", None)
109 w_gate = w_gate if w_gate is not None else getattr(experts, "w1", None)
110 w_up = getattr(experts, "up_proj", None)
111 w_up = w_up if w_up is not None else getattr(experts, "w3", None)
112 if w_gate is None or w_up is None or w_down is None:
113 raise NotImplementedError(
114 f"{type(experts).__name__}: only SwiGLU experts are supported "
115 "(fused gate_up_proj, or separate gate_proj/up_proj/down_proj, "
116 "w1/w2/w3 three-matrix layouts); use an EP-aware MoE module instead"
117 )
118 return w_gate, w_up, w_down
121# ``ArgSort`` has no int32/int64 AICore kernel and silently falls back to
122# AI_CPU (runtime warning), which is expensive on the EP dispatch path where it
123# sits before the all-to-all. Every key sorted here is a small non-negative
124# integer and float32 represents integers up to 2**24 exactly, so sorting a
125# float32 key preserves the order and lets the kernel run on AICore. Opt-in
126# while it is being measured (HP_EP_SORT_FP32=1).
127_SORT_KEY_FP32_LIMIT = 2 ** 24
128_SORT_FP32_ENABLED = os.environ.get("HP_EP_SORT_FP32", "0") == "1"
130# The routed dispatch issues two ragged all-to-all exchanges per pass (hidden
131# states, then expert indices) whose token counts are identical and whose
132# payloads differ only in width. A profile of the 256-card step shows the
133# exchanges strictly serialized (zero overlapping pairs across 1799 pairs), so
134# carrying both payloads in ONE exchange shortens a serial chain instead of
135# adding a third one. Opt-in while it is being measured
136# (HP_EP_FUSED_DISPATCH=1).
137_FUSED_DISPATCH_ENABLED = os.environ.get("HP_EP_FUSED_DISPATCH", "0") == "1"
139# The dispatch is followed by an expert GEMM that hard-depends on the arrived
140# tokens, so the routed all-to-all is the one collective in the step that has no
141# independent work to hide behind: a 512-card one-step profile of the K2.6 MoE
142# step puts hcom_alltoall + hcom_alltoallv at 5.87 s of union with only 0.375 s
143# of it overlapping AI-core kernels, while the same step's all-gather (2270
144# calls) and reduce-scatter (921 calls) already hide 46% / 54%. The work has to
145# be *created*, and the routed stream is what makes it possible: split the
146# expert-major token order into HP_EP_DISPATCH_CHUNKS contiguous chunks and
147# software-pipeline them, so chunk c's expert GEMM runs against chunk c+1's
148# dispatch exchange and chunk c's combine against chunk c+1's GEMM. Only the
149# first dispatch and the last combine stay exposed.
150#
151# Default 1 = the original single-exchange schedule, bit-identical: the chunked
152# path is never entered and every count/list it derives is unused.
153_DISPATCH_CHUNKS_RAW = os.environ.get("HP_EP_DISPATCH_CHUNKS", "1")
155# The dispatch tags every routed slot with its source token as
156# ``arange(token_count).repeat_interleave(top_k)``. That eager call lowers to a
157# *single-block* vector kernel: on one NPU it costs 9.25 ms per call at
158# T=8192 / K=8 (one vector core out of 48, ~58 MB/s of writes), i.e. 36 x 9.25 ms
159# = 0.33 s per step on the 18-layer testbed (16% of a 2.05 s step) and 1.07 s per
160# step on the 256-card shape.
161#
162# ``expand`` (broadcast copy) and ``div`` (integer floor division) build the
163# numerically identical int64 tensor in 0.034 / 0.041 ms -- 270x cheaper, all
164# cores. Measured end to end on the 18-layer testbed, 3 interleaved rounds each
165# (min-of-run step time, order rotated): legacy 2.049 / 2.057 / 2.065 s vs
166# expand 2.016 / 2.016 / 2.026 s, i.e. -1.6% step time and +2.4% throughput.
167# The gap between the kernel's 0.33 s and the 0.04 s actually recovered is the
168# usual one: a device kernel span is not the same as recoverable step time, and
169# the eager kernel spent ~half of its span overlapped with communication.
170_SOURCE_INDEX_MODES = ("legacy", "expand", "div")
171_SOURCE_INDEX_MODE = os.environ.get("HP_EP_SOURCE_INDEX", "expand").lower()
172# Dispatch counter for the sparse capacity-limit log (see ep_routed_dispatch).
173_CAPACITY_SYNC_COUNT = 0
176def _argsort_keys(keys: torch.Tensor, *, bound: int) -> torch.Tensor:
177 """Sort non-negative integer keys, on AICore when the range allows.
179 Args:
180 keys: Non-negative integer sort keys.
181 bound: Exclusive upper bound on every key, known statically by the
182 caller so no device sync is needed to check the float32 range.
184 Returns:
185 The int64 permutation that sorts ``keys``.
186 """
187 if _SORT_FP32_ENABLED and 0 < bound <= _SORT_KEY_FP32_LIMIT:
188 return keys.to(torch.float32).argsort()
189 return keys.argsort()
192def _expert_token_counts(indices: torch.Tensor, num_experts: int) -> torch.Tensor:
193 """Count routed tokens per expert with device ops only.
195 ``torch.bincount`` on NPU reads the input's min and max back to the host
196 (input validation and output sizing), so every call drains the device queue
197 and blocks the host. The counts here are consumed either fully on device
198 (grouped GEMM group list) or -- where the ragged all-to-all genuinely needs
199 them -- as a single host list, so accumulating them with ``scatter_add_``
200 keeps the same values at one device kernel and no host stall.
202 Args:
203 indices: Expert indices, any shape; values must be in ``[0, num_experts)``.
204 num_experts: Length of the returned histogram (``minlength``).
206 Returns:
207 int64 counts of shape ``[num_experts]`` on ``indices.device``.
208 """
209 flat = indices.to(torch.int64).reshape(-1)
210 counts = torch.zeros(num_experts, dtype=torch.int64, device=flat.device)
211 return counts.scatter_add_(0, flat, torch.ones_like(flat))
214def _local_swiglu_expert_forward(experts, dispatched_states, local_expert_indices):
215 """Compute dispatched tokens with the local stacked SwiGLU experts.
217 This function is installed as ``experts.forward`` for the HF-native EP
218 path. Calling the expert module, instead of indexing its parameters from
219 the parent MoE forward, allows nested FSDP forward hooks to unshard and
220 reshard expert parameters around the local computation.
221 """
222 token_order = _argsort_keys(
223 local_expert_indices, bound=experts.local_expert_count)
224 sorted_states = dispatched_states[token_order]
225 local_expert_counts = _expert_token_counts(
226 local_expert_indices, experts.local_expert_count)
227 if getattr(experts, "_ep_use_grouped_gemm", False):
228 grouped_forward = getattr(experts, "forward_expert_major", None)
229 if callable(grouped_forward):
230 sorted_output = grouped_forward(sorted_states, local_expert_counts)
231 else:
232 gate_weight, up_weight, down_weight = resolve_swiglu_weights(experts)
233 if up_weight is not None:
234 raise ValueError(
235 "EP grouped GEMM currently requires packed gate_up_proj weights"
236 )
237 sorted_output = npu_grouped_swiglu(
238 sorted_states,
239 gate_weight,
240 down_weight,
241 local_expert_counts,
242 )
243 output = torch.empty_like(sorted_output)
244 output[token_order] = sorted_output
245 return output
247 gate_weight, up_weight, down_weight = resolve_swiglu_weights(experts)
248 # The eager path needs host ints: read them in one transfer rather than
249 # draining the queue once per local expert.
250 expert_token_counts = local_expert_counts.tolist()
251 sorted_outputs = []
252 token_start = 0
253 for local_expert_index in range(experts.local_expert_count):
254 expert_token_count = expert_token_counts[local_expert_index]
255 expert_states = sorted_states[token_start:token_start + expert_token_count]
256 if up_weight is None:
257 gate_up_states = F.linear( # pylint: disable=not-callable
258 expert_states,
259 gate_weight[local_expert_index],
260 )
261 apply_gate = getattr(experts, "_ep_apply_gate", None)
262 if apply_gate is None:
263 gate_states, up_states = gate_up_states.chunk(2, dim=-1)
264 activation = getattr(experts, "_ep_act_fn", F.silu)
265 activated_states = activation(gate_states) * up_states
266 else:
267 activated_states = apply_gate(gate_up_states)
268 else:
269 gate_states = F.linear( # pylint: disable=not-callable
270 expert_states, gate_weight[local_expert_index]
271 )
272 up_states = F.linear( # pylint: disable=not-callable
273 expert_states, up_weight[local_expert_index]
274 )
275 activation = getattr(experts, "_ep_act_fn", F.silu)
276 activated_states = activation(gate_states) * up_states
277 sorted_outputs.append(
278 F.linear( # pylint: disable=not-callable
279 activated_states,
280 down_weight[local_expert_index],
281 )
282 )
283 token_start += expert_token_count
285 sorted_output = torch.cat(sorted_outputs)
286 output = torch.empty_like(sorted_output)
287 output[token_order] = sorted_output
288 return output
291def _get_global_expert_count(module):
292 """Return the model-level routed expert count for an MoE module."""
293 if hasattr(module.experts, "num_experts"):
294 return module.experts.num_experts
295 if hasattr(module, "num_experts"):
296 return module.num_experts
297 if hasattr(module, "config") and hasattr(module.config, "num_experts"):
298 return module.config.num_experts
299 if hasattr(module, "config") and hasattr(module.config, "n_routed_experts"):
300 return module.config.n_routed_experts
301 raise ValueError(
302 f"{type(module).__name__}: cannot determine the global routed expert count"
303 )
306def bind_local_expert_forward(
307 module: Any,
308 ep_size: int,
309 use_grouped_gemm: bool = False,
310 apply_gate: Optional[Callable] = None,
311) -> None:
312 """Install the local expert compute entry used by TP-extend-EP.
314 Called by the EP compute factory (archetype or user-written) at apply
315 time: sets ``module.experts.local_expert_count`` and installs
316 ``experts.forward`` (via the forward rewriter's bound-forward install
317 point) so nested FSDP hooks unshard/reshard around the local SwiGLU
318 computation. ``apply_gate`` supplies model-specific fused gate/up
319 semantics when the default activation-times-up rule is insufficient.
320 """
321 global_expert_count = _get_global_expert_count(module)
322 if global_expert_count % ep_size != 0:
323 raise ValueError(
324 f"num_experts ({global_expert_count}) must be divisible by ep_size ({ep_size})"
325 )
326 module.experts.local_expert_count = global_expert_count // ep_size
327 activation = getattr(module.experts, "act_fn", None)
328 if activation is None:
329 hidden_act = getattr(getattr(module, "config", None), "hidden_act", "silu")
330 activation = {
331 "gelu": F.gelu,
332 "relu": F.relu,
333 "silu": F.silu,
334 "swish": F.silu,
335 }.get(hidden_act)
336 if activation is None:
337 raise ValueError(
338 f"{type(module).__name__}: unsupported expert activation {hidden_act!r}; "
339 "provide experts.act_fn or extend the EP activation registry"
340 )
341 if apply_gate is not None and use_grouped_gemm:
342 raise ValueError("custom expert gate activation is not supported by grouped GEMM")
343 # ``_ep_*`` is the contract the local compute path reads back off the module's
344 # experts child; setting them here is the point of the binder.
345 # pylint: disable=protected-access
346 module.experts._ep_act_fn = activation
347 module.experts._ep_apply_gate = apply_gate
348 module.experts._ep_use_grouped_gemm = use_grouped_gemm
349 # The forward write itself lives in the forward rewriter (05 §15.2.3:
350 # the single MethodType/assignment site); this binder only sets the
351 # companion attributes above.
352 _install_bound_forward(module.experts, _local_swiglu_expert_forward)
355def _resolve_capacity_factor(module: Any) -> Optional[float]:
356 """Return the configured expert capacity factor, or ``None`` when it is off.
358 Resolution order: the ``HP_EP_CAPACITY_FACTOR`` env override (handy for a sweep without
359 editing YAML), then the MoE block's own ``capacity_factor`` attribute (set by the EP recipe
360 from ``plan_overrides[].local_compute_fn``), then off. Off is the default, so an
361 unconfigured run keeps the previous behaviour exactly.
363 Args:
364 module: MoE block (``None`` is tolerated for callers without one).
366 Returns:
367 A positive float, or ``None`` when the capacity limit is disabled.
368 """
369 override = os.environ.get("HP_EP_CAPACITY_FACTOR", "").strip()
370 if override:
371 try:
372 value = float(override)
373 except ValueError:
374 logger.warning(
375 "HP_EP_CAPACITY_FACTOR=%r is not a number; capacity limit stays off", override)
376 return None
377 return value if value > 0 else None
378 value = getattr(getattr(module, "experts", None), "capacity_factor", None)
379 if value is None:
380 return None
381 value = float(value)
382 return value if value > 0 else None
385def _routed_slot_token_ids(token_count: int, experts_per_token: int, device: Any) -> torch.Tensor:
386 """Return the source token of every routed slot, as an index tensor.
388 ``arange(token_count).repeat_interleave(experts_per_token)`` and both fast
389 paths produce the same int64 tensor: slot ``token * K + i`` holds ``token``.
391 Args:
392 token_count: Number of local tokens ``T`` in this dispatch.
393 experts_per_token: Routed experts per token ``K``.
394 device: Device to build the tensor on.
396 Returns:
397 Contiguous int64 tensor of shape ``[T * K]``.
399 Raises:
400 ValueError: If ``HP_EP_SOURCE_INDEX`` names an unknown mode.
401 """
402 mode = _SOURCE_INDEX_MODE
403 if mode not in _SOURCE_INDEX_MODES:
404 raise ValueError(
405 f"HP_EP_SOURCE_INDEX must be one of {_SOURCE_INDEX_MODES}, but got {mode!r}"
406 )
407 if mode == "expand" and experts_per_token > 0:
408 # ``reshape`` may hand back a stride-0 view when the token dim is 1, so
409 # materialize: the callers index with the result and expect a flat buffer.
410 return (
411 torch.arange(token_count, device=device)
412 .unsqueeze(1)
413 .expand(token_count, experts_per_token)
414 .reshape(-1)
415 .contiguous()
416 )
417 if mode == "div" and experts_per_token > 0:
418 return torch.arange(token_count * experts_per_token, device=device).div(
419 experts_per_token, rounding_mode="floor"
420 )
421 return torch.arange(token_count, device=device).repeat_interleave(experts_per_token)
424def _prepare_ep_dispatch(
425 hidden_states: torch.Tensor,
426 topk_indices: torch.Tensor,
427 topk_weights: torch.Tensor,
428 *,
429 local_expert_count: int,
430 global_expert_count: int,
431 ep_size: int,
432 ep_group: Any,
433 keep: Optional[torch.Tensor] = None,
434) -> _EPDispatch:
435 """Sort routed tokens and exchange per-rank dispatch counts.
437 ``keep`` is an optional ``[T, K]`` bool mask (see
438 :func:`routing.apply_capacity_limit`). It compacts the routed slots *before* anything is
439 exchanged: dropped slots never become rows, which is the entire point of a capacity limit,
440 because a rank's expert buffers are sized by the rows it receives. The rest of the
441 pipeline is already ragged-safe -- the expert compute is row-based, and the combine
442 accumulates with ``index_add_`` over ``source_indices`` -- so a token with fewer than ``K``
443 surviving slots simply receives fewer contributions.
444 """
445 flattened_states = hidden_states.reshape(-1, hidden_states.shape[-1])
446 token_count = flattened_states.shape[0]
447 experts_per_token = topk_indices.shape[1]
448 if keep is None:
449 expert_indices = topk_indices.reshape(-1)
450 expert_weights = topk_weights.reshape(-1).to(flattened_states.dtype)
451 source_indices = _routed_slot_token_ids(
452 token_count, experts_per_token, flattened_states.device)
453 else:
454 flat_keep = keep.reshape(-1)
455 expert_indices = topk_indices.reshape(-1)[flat_keep]
456 expert_weights = topk_weights.reshape(-1)[flat_keep].to(flattened_states.dtype)
457 source_indices = _routed_slot_token_ids(
458 token_count, experts_per_token, flattened_states.device)[flat_keep]
459 destination_ranks = torch.div(expert_indices, local_expert_count, rounding_mode="floor")
460 dispatch_order = _argsort_keys(
461 destination_ranks * global_expert_count + expert_indices,
462 bound=ep_size * global_expert_count,
463 )
464 dispatched_states = flattened_states[source_indices[dispatch_order]].contiguous()
465 dispatched_indices = expert_indices[dispatch_order].unsqueeze(-1).contiguous()
466 # The counts exchange is a host-serialising step: a synchronous all-to-all followed
467 # by two ``.tolist()`` drains, and a drain waits for everything already enqueued on
468 # the stream, so the token all-to-all cannot be enqueued until the counts land. When
469 # the routing plan is static (``fix_router: true``) the answer is the same in every
470 # layer, so it is computed once and reused -- that removes the extra collective and
471 # both syncs from every later dispatch. See ``static_splits`` for the caveats.
472 plan_key = static_plan_key(
473 token_count, experts_per_token, ep_size, local_expert_count, global_expert_count,
474 device=flattened_states.device,
475 )
476 cached_counts = get_static_plan(plan_key)
477 if cached_counts is not None:
478 send_counts, receive_counts, send_counts_tensor, receive_counts_tensor = cached_counts
479 # Bisect probes (see ``probe_mode``): run one half of the original work so the
480 # 4/4 OOM can be attributed. Both branches are no-ops in a normal run.
481 mode = probe_mode()
482 if mode == "a2a":
483 dist.all_to_all_single(
484 torch.empty_like(send_counts_tensor), send_counts_tensor, group=ep_group,
485 )
486 elif mode == "drains":
487 send_counts_tensor.tolist()
488 receive_counts_tensor.tolist()
489 else:
490 send_counts_tensor = _expert_token_counts(destination_ranks, ep_size)
491 receive_counts_tensor = torch.empty_like(send_counts_tensor)
492 dist.all_to_all_single(receive_counts_tensor, send_counts_tensor, group=ep_group)
493 send_counts = send_counts_tensor.tolist()
494 receive_counts = receive_counts_tensor.tolist()
495 store_static_plan(
496 plan_key, (send_counts, receive_counts, send_counts_tensor, receive_counts_tensor),
497 f"(ep={ep_size}, tokens={token_count}, topk={experts_per_token}, "
498 f"tokens/peer={send_counts[0] if send_counts else 0})",
499 )
500 return _EPDispatch(
501 source_indices=source_indices,
502 expert_weights=expert_weights,
503 dispatch_order=dispatch_order,
504 states=dispatched_states,
505 expert_indices=dispatched_indices,
506 send_counts=send_counts,
507 receive_counts=receive_counts,
508 )
511def _resolve_dispatch_chunks(ep_size: int) -> int:
512 """Resolve ``HP_EP_DISPATCH_CHUNKS`` for an EP group of size ``ep_size``.
514 Args:
515 ep_size: Size of the EP group. The chunk count is clamped to it, since
516 a chunk covers a non-empty contiguous group of EP ranks.
518 Returns:
519 The number of chunks to split the routed exchange into; 1 (also the
520 default) keeps the original single-exchange schedule.
522 Raises:
523 ValueError: If the knob is not a positive integer.
524 """
525 raw = _DISPATCH_CHUNKS_RAW.strip()
526 try:
527 chunk_count = int(raw, 10)
528 except ValueError:
529 raise ValueError(
530 "HP_EP_DISPATCH_CHUNKS must be a positive integer, but got "
531 f"{raw!r}; leave it unset (or 1) for the unchunked schedule"
532 ) from None
533 if chunk_count < 1:
534 raise ValueError(
535 f"HP_EP_DISPATCH_CHUNKS must be >= 1, but got {chunk_count}"
536 )
537 if chunk_count > ep_size:
538 if logger.isEnabledFor(logging.DEBUG):
539 logger.debug(
540 "HP_EP_DISPATCH_CHUNKS=%d exceeds the EP group size %d; clamping to %d",
541 chunk_count, ep_size, ep_size,
542 )
543 chunk_count = ep_size
544 return chunk_count
547def _chunk_rank_range(group: int, ep_size: int, chunk_count: int) -> tuple[int, int]:
548 """Half-open EP rank range covered by chunk group ``group``.
550 The ``chunk_count`` groups partition the EP ranks into contiguous,
551 as-equal-as-possible ranges, so a group is always non-empty when
552 ``chunk_count <= ep_size`` (which :func:`_resolve_dispatch_chunks`
553 guarantees).
555 Args:
556 group: Chunk group index.
557 ep_size: Size of the EP group.
558 chunk_count: Number of chunk groups.
560 Returns:
561 ``(first_rank, last_rank)``, the first inclusive / last exclusive rank.
562 """
563 return group * ep_size // chunk_count, (group + 1) * ep_size // chunk_count
566def _chunk_rank_group(ep_rank: int, ep_size: int, chunk_count: int) -> int:
567 """Chunk group owning ``ep_rank`` (the inverse of :func:`_chunk_rank_range`).
569 Args:
570 ep_rank: This rank's index in the EP group.
571 ep_size: Size of the EP group.
572 chunk_count: Number of chunk groups.
574 Returns:
575 The group index whose range contains ``ep_rank``.
576 """
577 return ((ep_rank + 1) * chunk_count - 1) // ep_size
580class _EPDispatchChunks(NamedTuple):
581 """Chunk plan of one routed exchange (see :func:`_ep_dispatch_chunks`).
583 ``send_counts[c]`` / ``recv_counts[c]`` are full-length (``ep_size``) count
584 lists for chunk ``c``; ``row_ranges[c]`` is the chunk's half-open slice of
585 the expert-major stream. That slice is the same for the dispatch input and
586 for the combined output, because both are ordered by destination rank.
587 """
589 send_counts: list
590 recv_counts: list
591 row_ranges: list
594def _ep_dispatch_chunks(
595 send_counts: list,
596 receive_counts: list,
597 ep_size: int,
598 chunk_count: int,
599 ep_rank: int,
600) -> _EPDispatchChunks:
601 """Split the routed exchange into ``chunk_count`` self-consistent slices.
603 Chunk ``c`` of rank ``r`` carries exactly the rows ``r`` dispatches to rank
604 group ``(c + group(r)) % chunk_count`` -- a contiguous slice of the sorted
605 expert-major stream -- and expects its reply from rank group
606 ``(group(r) - c) % chunk_count``. Every count is derived from the counts
607 the caller already exchanged in :func:`_prepare_ep_dispatch`; no extra
608 collective is needed to learn a chunk's traffic.
610 The per-chunk rotation is what makes the pipeline meaningful: pinned to the
611 rank's own destination group, a chunk would receive *all* of the rank's rows
612 (a rank receives exactly the rows addressed to it) and the other chunks
613 would have nothing to compute on. Rotating by the rank group spreads both
614 the sends and the receives of every rank evenly over the chunks.
616 Args:
617 send_counts: Rows this rank dispatches to each EP rank.
618 receive_counts: Rows this rank receives from each EP rank.
619 ep_size: Size of the EP group.
620 chunk_count: Number of chunks (``2 <= chunk_count <= ep_size``).
621 ep_rank: This rank's index in the EP group.
623 Returns:
624 The :class:`_EPDispatchChunks` plan, in chunk order.
625 """
626 row_offsets = [0]
627 for count in send_counts:
628 row_offsets.append(row_offsets[-1] + count)
629 own_group = _chunk_rank_group(ep_rank, ep_size, chunk_count)
630 send_plan, recv_plan, row_ranges = [], [], []
631 for chunk in range(chunk_count):
632 first, last = _chunk_rank_range(
633 (chunk + own_group) % chunk_count, ep_size, chunk_count)
634 send_plan.append(
635 [send_counts[rank] if first <= rank < last else 0 for rank in range(ep_size)])
636 row_ranges.append((row_offsets[first], row_offsets[last]))
637 first, last = _chunk_rank_range(
638 (own_group - chunk) % chunk_count, ep_size, chunk_count)
639 recv_plan.append(
640 [receive_counts[rank] if first <= rank < last else 0 for rank in range(ep_size)])
641 return _EPDispatchChunks(send_plan, recv_plan, row_ranges)
644def _run_ep_local_experts(
645 module: Any,
646 dispatched_states: torch.Tensor,
647 dispatched_indices: torch.Tensor,
648 send_counts: list[int],
649 receive_counts: list[int],
650 ep_group: Any,
651 expert_offset: int,
652) -> torch.Tensor:
653 """Dispatch tokens, run local experts, and return combined expert outputs."""
654 received_states = ep_all_to_all(dispatched_states, send_counts, receive_counts, ep_group)
655 received_indices = ep_all_to_all(
656 dispatched_indices, send_counts, receive_counts, ep_group
657 ).squeeze(-1)
658 local_outputs = module.experts(received_states, received_indices - expert_offset)
659 return ep_all_to_all(local_outputs.contiguous(), receive_counts, send_counts, ep_group)
662def _fused_row_layout(
663 hidden_size: int,
664 state_dtype: torch.dtype,
665 index_dtype: torch.dtype,
666) -> tuple[int, int]:
667 """Slot geometry of one fused dispatch row.
669 The row is ``[hidden states | head pad | index bytes]`` read in the states'
670 dtype: ``index_slots`` trailing state elements carry the index, and the head
671 pad moves that region to a byte offset which is a multiple of the index
672 element size -- the alignment ``Tensor.view(dtype)`` needs to reinterpret it
673 without a copy.
675 Args:
676 hidden_size: Width of the hidden states in the row.
677 state_dtype: Dtype of the hidden states (the row's storage dtype).
678 index_dtype: Dtype of the expert indices carried in the tail slots.
680 Returns:
681 ``(head_slots, index_slots)`` in state elements; the row is
682 ``hidden_size + head_slots + index_slots`` state elements wide.
683 """
684 index_slots = math.lcm(state_dtype.itemsize, index_dtype.itemsize) // state_dtype.itemsize
685 return (-hidden_size) % index_slots, index_slots
688def _pack_fused_dispatch(
689 dispatched_states: torch.Tensor,
690 dispatched_expert_indices: torch.Tensor,
691) -> torch.Tensor:
692 """Pack the hidden states and the expert indices of a token into one row.
694 Byte-view packing: the index bytes sit in the tail slots reinterpreted in
695 the states' dtype, so nothing is converted numerically and both payloads
696 stay bit-exact (a bf16 buffer cannot hold an expert index above 256
697 exactly, so the index must travel as bytes). The row is one tensor, so the
698 whole fused exchange stays a differentiable ``cat`` on the send side and a
699 pair of views on the receive side.
701 Args:
702 dispatched_states: Expert-major hidden states, shape ``[tokens, H]``.
703 dispatched_expert_indices: Matching indices, shape ``[tokens, 1]``.
705 Returns:
706 ``[tokens, H + head_slots + index_slots]`` in the states' dtype.
707 """
708 token_count, hidden_size = dispatched_states.shape
709 head_slots, index_slots = _fused_row_layout(
710 hidden_size, dispatched_states.dtype, dispatched_expert_indices.dtype)
711 index_bytes = dispatched_expert_indices.reshape(token_count, 1).view(torch.uint8)
712 index_slots_buffer = F.pad(
713 index_bytes, (0, index_slots * dispatched_states.element_size() - index_bytes.shape[1]))
714 head_buffer = torch.zeros(
715 token_count, head_slots, dtype=dispatched_states.dtype, device=dispatched_states.device)
716 return torch.cat(
717 [dispatched_states, head_buffer, index_slots_buffer.view(dispatched_states.dtype)],
718 dim=1,
719 )
722def _unpack_fused_dispatch(
723 packed: torch.Tensor,
724 *,
725 hidden_size: int,
726 index_dtype: torch.dtype,
727) -> tuple[torch.Tensor, torch.Tensor]:
728 """Split a received fused row back into ``(states, indices)``.
730 Both halves are views of the exchanged buffer, so the received bytes are
731 handed over as they arrived -- no decode step can change a value. The index
732 half is detached: it is payload on its way into integer indexing, while the
733 states half stays the differentiable path that carries the routed gradient.
734 Staying view-only also keeps the exchange's lazy wait intact: a
735 materializing op here would enqueue the wait before the caller reaches its
736 independent work.
738 Args:
739 packed: Exchanged rows, shape ``[tokens, row]`` in the states' dtype.
740 hidden_size: Width of the hidden states in the row.
741 index_dtype: Dtype of the expert indices carried in the tail slots.
743 Returns:
744 ``(states [tokens, H], indices [tokens])``, both views of ``packed``.
745 """
746 head_slots, index_slots = _fused_row_layout(hidden_size, packed.dtype, index_dtype)
747 index_start = hidden_size + head_slots
748 index_bytes = packed.detach()[:, index_start:index_start + index_slots].view(torch.uint8)
749 return (
750 packed[:, :hidden_size],
751 index_bytes[:, :index_dtype.itemsize].view(index_dtype).squeeze(-1),
752 )
755def _fused_dispatch_exchange(
756 dispatched_states: torch.Tensor,
757 dispatched_expert_indices: torch.Tensor,
758 send_counts: list[int],
759 receive_counts: list[int],
760 ep_group: Any,
761) -> tuple[torch.Tensor, torch.Tensor]:
762 """Dispatch the hidden states and the expert indices in ONE all-to-all.
764 Same handle/wait contract as the two exchanges it replaces: the exchange is
765 issued through :func:`ep_all_to_all_async` and both returned tensors are
766 views of it, so the wait still lands on their first non-view consumer (the
767 shared-expert overlap window is untouched). ``allow_pending=False`` is what
768 keeps that possible: the unpack below is a view chain, which the split-free
769 exchange's pending handle cannot carry, so this exchange stays on the
770 async-tensor path even with ``HP_EP_EQUAL_A2A`` on.
771 """
772 packed = _pack_fused_dispatch(dispatched_states, dispatched_expert_indices)
773 # ``allow_pending=False``: the unpack below reads the exchanged buffer
774 # through a view chain (slice / detach / view-as-dtype), which a pending
775 # handle cannot express, so the fused exchange keeps the async-tensor lazy
776 # path and with it its own overlap window.
777 exchanged = ep_all_to_all_async(
778 packed, send_counts, receive_counts, ep_group, allow_pending=False)
779 return _unpack_fused_dispatch(
780 exchanged,
781 hidden_size=dispatched_states.shape[-1],
782 index_dtype=dispatched_expert_indices.dtype,
783 )
786def _dispatch_chunk(
787 dispatched_states: torch.Tensor,
788 dispatched_expert_indices: torch.Tensor,
789 chunk_send_counts: list[int],
790 chunk_recv_counts: list[int],
791 ep_group: Any,
792 fused: bool = False,
793) -> tuple[Any, Any]:
794 """Issue one chunk's dispatch exchange, without waiting for it.
796 Every mode goes through :func:`ep_all_to_all_async`, and neither its fused
797 unpack nor the ``squeeze`` below reads the payload, so nothing materializes
798 the exchange here: the wait is enqueued by the chunk's first consumer (the
799 expert GEMM, through :func:`wait_ep_all_to_all`), which is what lets that
800 GEMM overlap the next chunk's exchange. ``fused`` mirrors the caller's
801 switch: the packed states+indices exchange belongs to the split dispatch
802 (:func:`ep_routed_dispatch`), not to the cohesive routed branch, which
803 issues the two exchanges itself.
805 Args:
806 dispatched_states: This chunk's contiguous slice of the expert-major
807 hidden states, ``[rows, H]``.
808 dispatched_expert_indices: Matching slice, ``[rows, 1]``.
809 chunk_send_counts: Rows of this chunk sent to each EP rank.
810 chunk_recv_counts: Rows of this chunk received from each EP rank.
811 ep_group: Extended EP process group.
812 fused: Send the states and the indices in one packed exchange.
814 Returns:
815 ``(received_states, received_indices)``, pending their exchange -- a
816 tensor or a :class:`_PendingEqualA2A` handle, depending on the knob, so
817 a caller must pass them through :func:`wait_ep_all_to_all`.
818 """
819 if fused:
820 return _fused_dispatch_exchange(
821 dispatched_states,
822 dispatched_expert_indices,
823 chunk_send_counts,
824 chunk_recv_counts,
825 ep_group,
826 )
827 received_states = ep_all_to_all_async(
828 dispatched_states, chunk_send_counts, chunk_recv_counts, ep_group)
829 # squeeze is a view (the handle's own, or a tensor's), so the wait stays
830 # deferred until the experts read it.
831 received_indices = ep_all_to_all_async(
832 dispatched_expert_indices, chunk_send_counts, chunk_recv_counts, ep_group).squeeze(-1)
833 return received_states, received_indices
836def _aggregate_ep_outputs(
837 combined_outputs: torch.Tensor,
838 expert_weights: torch.Tensor,
839 source_indices: torch.Tensor,
840 dispatch_order: torch.Tensor,
841 output_shape: tuple[int, int, int],
842) -> torch.Tensor:
843 """Undo expert-major routing and aggregate weighted top-k outputs."""
844 weighted_outputs = combined_outputs * expert_weights[dispatch_order].unsqueeze(-1)
845 output = torch.zeros(
846 output_shape[0] * output_shape[1],
847 output_shape[-1],
848 dtype=weighted_outputs.dtype,
849 device=weighted_outputs.device,
850 )
851 output.index_add_(0, source_indices[dispatch_order], weighted_outputs)
852 return output.view(*output_shape)
855def _aggregate_ep_chunks(
856 combined_chunks: list[torch.Tensor],
857 row_ranges: list[tuple[int, int]],
858 expert_weights: torch.Tensor,
859 source_indices: torch.Tensor,
860 dispatch_order: torch.Tensor,
861 output_shape: tuple[int, int, int],
862) -> torch.Tensor:
863 """Aggregate the per-chunk combine outputs into the routed output.
865 Same arithmetic as :func:`_aggregate_ep_outputs`, applied to each chunk's
866 contiguous slice of the expert-major stream. Every chunk accumulates into
867 the one output buffer through ``index_add_`` (which sums duplicate source
868 slots), so the routed output is assembled once, after the last chunk, and no
869 stream-sized copy is needed to concatenate the chunks back into stream
870 order. Reading a chunk is also what waits for its combine exchange, so the
871 wait of chunk ``c`` lands after chunk ``c + 1``'s experts have been issued.
873 Args:
874 combined_chunks: Per-chunk combine outputs, in plan order (tensors, or
875 pending handles with ``HP_EP_EQUAL_A2A=1``).
876 row_ranges: Matching half-open slices of the expert-major stream.
877 expert_weights: Flattened routing weights, ``[tokens * top_k]``.
878 source_indices: Source token of every routed slot, ``[tokens * top_k]``.
879 dispatch_order: The dispatch permutation, ``[tokens * top_k]``.
880 output_shape: Shape of the routed output.
882 Returns:
883 The routed output, shape ``output_shape``.
884 """
885 reference = combined_chunks[0]
886 output = torch.zeros(
887 output_shape[0] * output_shape[1],
888 output_shape[-1],
889 dtype=reference.dtype,
890 device=reference.device,
891 )
892 for combined, (row_start, row_end) in zip(combined_chunks, row_ranges):
893 rows = dispatch_order[row_start:row_end]
894 # Reading a chunk is what orders its stream on the pending combine.
895 weighted_outputs = wait_ep_all_to_all(combined) * expert_weights[rows].unsqueeze(-1)
896 output.index_add_(0, source_indices[rows], weighted_outputs)
897 return output.view(*output_shape)
900def _run_ep_local_experts_chunked(
901 module: Any,
902 plan: _EPDispatchChunks,
903 dispatched_states: torch.Tensor,
904 dispatched_indices: torch.Tensor,
905 ep_group: Any,
906 expert_offset: int,
907) -> list[torch.Tensor]:
908 """Dispatch, run the local experts and combine, one chunk at a time.
910 The chunks are software-pipelined: the next chunk's dispatch exchange is
911 issued *before* the current chunk's experts run, so the GEMM overlaps the
912 in-flight exchange, and each chunk's combine is issued right after its GEMM,
913 so it overlaps the next chunk's GEMM. Only chunk 0's dispatch and the last
914 chunk's combine stay fully exposed.
916 Neither wait is forced here -- the experts' first read materializes the
917 dispatch, and the aggregation is what reads a combine -- so the chunked path
918 keeps the lazy-wait property of :func:`ep_all_to_all_async`.
920 Args:
921 module: MoE block exposing ``experts`` (with ``local_expert_count``).
922 plan: Chunk plan from :func:`_ep_dispatch_chunks`.
923 dispatched_states: Expert-major hidden states, ``[rows, H]``.
924 dispatched_indices: Matching expert indices, ``[rows, 1]``.
925 ep_group: Extended EP process group.
926 expert_offset: First global expert index owned by this rank.
928 Returns:
929 One combine output per chunk, in plan order (their wait is pending).
930 """
931 chunk_count = len(plan.row_ranges)
933 def issue(chunk: int) -> tuple[torch.Tensor, torch.Tensor]:
934 """Issue chunk ``chunk``'s dispatch exchange over its stream slice.
936 Args:
937 chunk: Chunk index into ``plan``.
939 Returns:
940 The chunk's pending ``(received_states, received_indices)``.
941 """
942 row_start, row_end = plan.row_ranges[chunk]
943 return _dispatch_chunk(
944 dispatched_states[row_start:row_end],
945 dispatched_indices[row_start:row_end],
946 plan.send_counts[chunk],
947 plan.recv_counts[chunk],
948 ep_group,
949 )
951 incoming = issue(0)
952 combined_chunks = []
953 for chunk in range(chunk_count):
954 # Issued before the experts of this chunk run: that exchange is what
955 # their GEMM overlaps.
956 following = issue(chunk + 1) if chunk + 1 < chunk_count else None
957 received_states, received_indices = incoming
958 local_outputs = module.experts(
959 wait_ep_all_to_all(received_states),
960 wait_ep_all_to_all(received_indices) - expert_offset)
961 combined_chunks.append(ep_all_to_all_async(
962 local_outputs.contiguous(),
963 plan.recv_counts[chunk],
964 plan.send_counts[chunk],
965 ep_group,
966 ))
967 incoming = following
968 return combined_chunks
971def _chunked_ep_forward(
972 module: Any,
973 dispatch: "_EPDispatch",
974 ep_group: Any,
975 ep_size: int,
976 ep_rank: int,
977 chunk_count: int,
978 expert_offset: int,
979 output_shape: tuple[int, int, int],
980) -> torch.Tensor:
981 """Run the routed forward as ``HP_EP_DISPATCH_CHUNKS`` software-pipelined slices.
983 Split out of :func:`ep_routed_forward` so that entry point stays readable; the schedule
984 itself is unchanged (see :func:`_run_ep_local_experts_chunked`).
985 """
986 plan = _ep_dispatch_chunks(
987 dispatch.send_counts, dispatch.receive_counts, ep_size, chunk_count, ep_rank)
988 combined_chunks = _run_ep_local_experts_chunked(
989 module,
990 plan,
991 dispatch.states,
992 dispatch.expert_indices,
993 ep_group,
994 expert_offset,
995 )
996 return _aggregate_ep_chunks(
997 combined_chunks,
998 plan.row_ranges,
999 dispatch.expert_weights,
1000 dispatch.source_indices,
1001 dispatch.dispatch_order,
1002 output_shape,
1003 )
1006def ep_routed_forward(
1007 module: Any,
1008 hidden_states: torch.Tensor,
1009 *,
1010 router_fn: Callable,
1011 ep_group: Any,
1012) -> torch.Tensor:
1013 """Routed-experts pipeline: SP-in (local chunk) -> all communication
1014 inside -> SP-out. **Routed branch only.**
1016 This primitive deliberately does NOT handle shared experts / scalar
1017 gates / branch merging — that composition is model semantics and belongs
1018 to the caller (an ep_compute.py archetype or a user-written factory,
1019 accuracy_fix_plan.md §3). There is no ``tp_group`` parameter: if the
1020 caller invokes a nested-boundary submodule (e.g. ``module.shared_expert``),
1021 that submodule's own boundary performs its TP communication — the
1022 **nested-boundary call contract**:
1024 1. the input is the parent local region's current logical local layout;
1025 2. the nested boundary exclusively owns its parameter layout and its TP
1026 communication (entry/exit via its own PrecompiledBoundary);
1027 3. the return value is already the nested boundary's out_dst logical
1028 layout (e.g. under SP: the complete per-token values of the local
1029 sequence chunk);
1030 4. the caller MUST NOT repeat any compensating collective
1031 (all-reduce / reduce-scatter / all-gather) on the returned value
1032 over the nested boundary's mesh.
1034 Communication flow (isomorphic to Megatron token_dispatcher.py
1035 MoEAlltoAllTokenDispatcher):
1036 router (local chunk, no communication) -> a2a dispatch (extended EP
1037 group, including TP ranks) -> local SwiGLU (complete expert weights, no
1038 internal communication, no Partial) -> a2a combine (returns over the
1039 same group) -> weighted aggregation.
1041 Input hidden [B, S/tp, H] (local sequence chunk, boundary identity);
1042 output [B, S/tp, H] (complete, boundary identity).
1044 ``router_fn`` is supplied BY THE CALLER (explicit choice, e.g. an entry
1045 of MOE_ROUTER_ADAPTERS picked by name in the factory code).
1047 With ``HP_EP_DISPATCH_CHUNKS > 1`` the same exchange is split into that many
1048 contiguous chunks of the expert-major order and software-pipelined (see
1049 :func:`_run_ep_local_experts_chunked`): the math is the same up to the
1050 grouping of the experts' GEMM, and the routed all-to-all -- the one
1051 collective of the step with no independent work to hide behind -- gets the
1052 chunk GEMMs to overlap with. Default 1 keeps the unchunked schedule.
1054 Extended EP group = the ep axis of the derived expert mesh (flatten
1055 ep_size consecutive ranks: first span the TP group, then extend to
1056 adjacent dp/cp ranks; MindSpeed TP-extend-EP / Megatron etp=1 + ep
1057 homogeneous across TP). Expert weights are only Shard(0) along the
1058 expert dim -- each rank holds num_experts/ep_size complete experts, so
1059 there is no all_gather/reduce_scatter pair.
1060 """
1061 ep_size = ep_group.size()
1062 ep_rank = dist.get_rank(group=ep_group)
1063 local_expert_count = module.experts.local_expert_count
1064 global_expert_count = local_expert_count * ep_size
1065 expert_offset = ep_rank * local_expert_count
1066 # Resolved before the dispatch preparation so a mistyped knob fails before
1067 # any collective is issued.
1068 chunk_count = _resolve_dispatch_chunks(ep_size)
1070 output_shape = tuple(hidden_states.shape)
1071 topk_indices, topk_weights = router_fn(module, hidden_states) # [T, K]
1072 dispatch = _prepare_ep_dispatch(
1073 hidden_states,
1074 topk_indices,
1075 topk_weights,
1076 local_expert_count=local_expert_count,
1077 global_expert_count=global_expert_count,
1078 ep_size=ep_size,
1079 ep_group=ep_group,
1080 )
1081 if chunk_count > 1:
1082 return _chunked_ep_forward(
1083 module, dispatch, ep_group, ep_size, ep_rank, chunk_count, expert_offset,
1084 output_shape,
1085 )
1086 combined_expert_outputs = _run_ep_local_experts(
1087 module,
1088 dispatch.states,
1089 dispatch.expert_indices,
1090 dispatch.send_counts,
1091 dispatch.receive_counts,
1092 ep_group,
1093 expert_offset,
1094 )
1095 return _aggregate_ep_outputs(
1096 combined_expert_outputs,
1097 dispatch.expert_weights,
1098 dispatch.source_indices,
1099 dispatch.dispatch_order,
1100 output_shape,
1101 )
1104class _EPRoutedChunks(NamedTuple):
1105 """Pending per-chunk dispatch state of a chunked routed exchange.
1107 ``received_states[c]`` / ``received_indices[c]`` are chunk ``c``'s issued but
1108 unwritten exchange results (the wait lands on the chunk's expert GEMM, via
1109 :func:`wait_ep_all_to_all`); the chunk's traffic and stream slice are in
1110 ``plan``.
1111 """
1113 plan: _EPDispatchChunks
1114 received_states: list
1115 received_indices: list
1118class EPRoutedState(NamedTuple):
1119 """Pending routed-experts state between dispatch and experts+combine.
1121 Produced by :func:`ep_routed_dispatch` and consumed by
1122 :func:`ep_routed_experts_and_combine`. ``received_states`` /
1123 ``received_indices`` may still be in flight (the exchange is issued
1124 asynchronously); they materialize on first non-view use, or -- with
1125 ``HP_EP_EQUAL_A2A=1`` -- on the explicit :func:`wait_ep_all_to_all` the
1126 consumer must call before reading them. A consumer outside
1127 :func:`ep_routed_experts_and_combine` therefore has to pass them through
1128 that helper rather than reading them directly.
1130 With ``HP_EP_DISPATCH_CHUNKS > 1`` the exchange is split into per-chunk
1131 slices, which cannot be expressed as the flat ``received_states`` /
1132 ``received_indices`` tensors: those two stay ``None`` and the pending state
1133 travels in ``chunks`` instead. ``send_counts`` / ``receive_counts`` keep the
1134 full (unchunked) counts either way.
1135 """
1137 source_token_indices: torch.Tensor
1138 flattened_expert_weights: torch.Tensor
1139 dispatch_order: torch.Tensor
1140 received_states: Optional[Any]
1141 received_indices: Optional[Any]
1142 send_counts: list
1143 receive_counts: list
1144 output_shape: tuple
1145 chunks: Optional[_EPRoutedChunks] = None
1148def _capacity_limited_slots(
1149 topk_indices: torch.Tensor,
1150 topk_weights: torch.Tensor,
1151 capacity_factor: float,
1152 global_expert_count: int,
1153) -> torch.Tensor:
1154 """Bound every expert's routed slots and report the drop rate sparsely.
1156 Split out of :func:`ep_routed_dispatch` to keep that entry point readable; the reporting
1157 cadence is part of the contract (see the comment inside).
1158 """
1159 keep, dropped = apply_capacity_limit(
1160 topk_indices, capacity_factor, global_expert_count, topk_weights)
1161 # Draining the drop count is a host sync, so pay it on the first dispatch and then
1162 # sparsely -- never once per layer, which is what made the first sweep look +90%.
1163 global _CAPACITY_SYNC_COUNT # pylint: disable=global-statement
1164 _CAPACITY_SYNC_COUNT += 1
1165 if _CAPACITY_SYNC_COUNT == 1 or _CAPACITY_SYNC_COUNT % 200 == 0:
1166 total_slots = topk_indices.numel()
1167 dropped_slots = int(dropped.item())
1168 logger.info(
1169 "EP capacity limit: factor=%s order=%s dropped %.2f%% of routed slots "
1170 "(%d experts, call %d)",
1171 capacity_factor, _resolve_capacity_order(),
1172 100.0 * dropped_slots / max(total_slots, 1),
1173 global_expert_count, _CAPACITY_SYNC_COUNT,
1174 )
1175 return keep
1178def _issue_chunked_dispatch(
1179 dispatch: "_EPDispatch",
1180 ep_group: Any,
1181 ep_size: int,
1182 ep_rank: int,
1183 chunk_count: int,
1184 batch_size: int,
1185 sequence_length: int,
1186 hidden_size: int,
1187) -> "EPRoutedState":
1188 """Issue every chunk's split exchange up front and keep the handles pending.
1190 Split out of :func:`ep_routed_dispatch`; the pending-handle contract is unchanged.
1191 """
1192 plan = _ep_dispatch_chunks(
1193 dispatch.send_counts, dispatch.receive_counts, ep_size, chunk_count, ep_rank)
1194 chunk_states, chunk_indices = [], []
1195 for chunk in range(chunk_count):
1196 row_start, row_end = plan.row_ranges[chunk]
1197 chunk_state, chunk_index = _dispatch_chunk(
1198 dispatch.states[row_start:row_end],
1199 dispatch.expert_indices[row_start:row_end],
1200 plan.send_counts[chunk],
1201 plan.recv_counts[chunk],
1202 ep_group,
1203 fused=_FUSED_DISPATCH_ENABLED,
1204 )
1205 chunk_states.append(chunk_state)
1206 chunk_indices.append(chunk_index)
1207 return EPRoutedState(
1208 dispatch.source_indices,
1209 dispatch.expert_weights,
1210 dispatch.dispatch_order,
1211 None,
1212 None,
1213 dispatch.send_counts,
1214 dispatch.receive_counts,
1215 (batch_size, sequence_length, hidden_size),
1216 _EPRoutedChunks(plan, chunk_states, chunk_indices),
1217 )
1220def ep_routed_dispatch(
1221 module: Any,
1222 hidden_states: torch.Tensor,
1223 *,
1224 router_fn: Callable,
1225 ep_group: Any,
1226) -> EPRoutedState:
1227 """Route locally and **launch** the token exchange without waiting.
1229 Same routing/dispatch preparation as :func:`ep_routed_forward`, but the two
1230 token exchanges are issued through :func:`ep_all_to_all_async`, whose wait
1231 is deferred to the first non-view consumer. A caller that runs independent
1232 work between this call and :func:`ep_routed_experts_and_combine` (for
1233 DeepSeek-V3 style blocks: the shared-expert MLP, which depends only on
1234 ``hidden_states``) therefore overlaps that work with the in-flight
1235 exchange instead of serializing behind it.
1237 With ``HP_EP_FUSED_DISPATCH=1`` the states and the indices travel in one
1238 packed exchange instead of two (see :func:`_fused_dispatch_exchange`); the
1239 pending state and the wait point are the same either way. With
1240 ``HP_EP_DISPATCH_CHUNKS > 1`` every chunk's exchange is issued here in the
1241 same way, so a caller's independent work still overlaps the whole dispatch
1242 rather than only its first exchange (see :func:`_ep_dispatch_chunks`).
1244 The pending values in the returned state are *not* necessarily tensors: with
1245 ``HP_EP_EQUAL_A2A=1`` and a uniform plan they are pending handles, and a
1246 consumer that drives this state itself must read them through
1247 :func:`wait_ep_all_to_all`.
1249 Args:
1250 module: MoE block exposing ``experts`` (with ``local_expert_count``).
1251 hidden_states: Local sequence chunk, shape ``[B, S, H]``.
1252 router_fn: Router adapter, called as ``router_fn(module, hidden_states)``.
1253 ep_group: Extended EP process group.
1255 Returns:
1256 The pending :class:`EPRoutedState` to hand to
1257 :func:`ep_routed_experts_and_combine`.
1258 """
1259 ep_size = ep_group.size()
1260 local_expert_count = module.experts.local_expert_count
1261 global_expert_count = local_expert_count * ep_size
1262 # Resolved before the dispatch preparation so a mistyped knob fails before
1263 # any collective is issued.
1264 chunk_count = _resolve_dispatch_chunks(ep_size)
1266 batch_size, sequence_length, hidden_size = hidden_states.shape
1267 topk_indices, topk_weights = router_fn(module, hidden_states)
1268 capacity_factor = _resolve_capacity_factor(module)
1269 keep = None
1270 if capacity_factor is not None:
1271 keep = _capacity_limited_slots(
1272 topk_indices, topk_weights, capacity_factor, global_expert_count)
1273 dispatch = _prepare_ep_dispatch(
1274 hidden_states,
1275 topk_indices,
1276 topk_weights,
1277 local_expert_count=local_expert_count,
1278 global_expert_count=global_expert_count,
1279 ep_size=ep_size,
1280 ep_group=ep_group,
1281 keep=keep,
1282 )
1283 source_token_indices = dispatch.source_indices
1284 flattened_expert_weights = dispatch.expert_weights
1285 dispatch_order = dispatch.dispatch_order
1286 dispatched_states = dispatch.states
1287 dispatched_expert_indices = dispatch.expert_indices
1288 send_counts = dispatch.send_counts
1289 receive_counts = dispatch.receive_counts
1290 if chunk_count > 1:
1291 return _issue_chunked_dispatch(
1292 dispatch, ep_group, ep_size, dist.get_rank(group=ep_group), chunk_count,
1293 batch_size, sequence_length, hidden_size,
1294 )
1295 if _FUSED_DISPATCH_ENABLED:
1296 received_states, received_indices = _fused_dispatch_exchange(
1297 dispatched_states,
1298 dispatched_expert_indices,
1299 send_counts,
1300 receive_counts,
1301 ep_group,
1302 )
1303 else:
1304 received_states = ep_all_to_all_async(
1305 dispatched_states, send_counts, receive_counts, ep_group)
1306 # squeeze is a view, so the wait stays deferred until the experts read it.
1307 received_indices = ep_all_to_all_async(
1308 dispatched_expert_indices, send_counts, receive_counts, ep_group).squeeze(-1)
1309 return EPRoutedState(
1310 source_token_indices,
1311 flattened_expert_weights,
1312 dispatch_order,
1313 received_states,
1314 received_indices,
1315 send_counts,
1316 receive_counts,
1317 (batch_size, sequence_length, hidden_size),
1318 )
1321def ep_routed_experts_and_combine(
1322 module: Any,
1323 state: EPRoutedState,
1324 ep_group: Any,
1325) -> torch.Tensor:
1326 """Run the local experts on a dispatched state and exchange the outputs back.
1328 The dispatch exchange reached ``state`` from :func:`ep_routed_dispatch` and
1329 is waited here -- at the first expert read, i.e. after whatever work the
1330 caller ran in between -- and the combine exchange is issued asynchronously
1331 as well, so its wait lands on the weighted aggregation instead of blocking
1332 right after the local experts. Both waits go through
1333 :func:`wait_ep_all_to_all`, which is what a pending exchange (with
1334 ``HP_EP_EQUAL_A2A=1``) needs before it can be read.
1336 With ``HP_EP_DISPATCH_CHUNKS > 1`` (``state.chunks``) the experts and the
1337 combines run one chunk at a time, keeping the same schedule property: chunk
1338 ``c``'s combine is issued before chunk ``c + 1``'s experts run, so it
1339 overlaps their GEMM, and every combine wait lands on the aggregation.
1341 Args:
1342 module: MoE block exposing ``experts`` (with ``local_expert_count``).
1343 state: State returned by :func:`ep_routed_dispatch`.
1344 ep_group: Extended EP process group.
1346 Returns:
1347 The routed branch output, shape ``state.output_shape``.
1348 """
1349 expert_offset = dist.get_rank(group=ep_group) * module.experts.local_expert_count
1350 if state.chunks is not None:
1351 plan = state.chunks.plan
1352 combined_chunks = []
1353 for chunk, (received_states, received_indices) in enumerate(
1354 zip(state.chunks.received_states, state.chunks.received_indices)):
1355 local_outputs = module.experts(
1356 wait_ep_all_to_all(received_states),
1357 wait_ep_all_to_all(received_indices) - expert_offset)
1358 combined_chunks.append(ep_all_to_all_async(
1359 local_outputs.contiguous(),
1360 plan.recv_counts[chunk],
1361 plan.send_counts[chunk],
1362 ep_group,
1363 ))
1364 return _aggregate_ep_chunks(
1365 combined_chunks,
1366 plan.row_ranges,
1367 state.flattened_expert_weights,
1368 state.source_token_indices,
1369 state.dispatch_order,
1370 state.output_shape,
1371 )
1372 local_outputs = module.experts(
1373 wait_ep_all_to_all(state.received_states),
1374 wait_ep_all_to_all(state.received_indices) - expert_offset)
1375 combined_expert_outputs = ep_all_to_all_async(
1376 local_outputs.contiguous(), state.receive_counts, state.send_counts, ep_group)
1377 return _aggregate_ep_outputs(
1378 wait_ep_all_to_all(combined_expert_outputs),
1379 state.flattened_expert_weights,
1380 state.source_token_indices,
1381 state.dispatch_order,
1382 state.output_shape,
1383 )
1384def _merge_shared_none(
1385 module: Any,
1386 hidden_states: torch.Tensor,
1387 routed: torch.Tensor,
1388) -> torch.Tensor:
1389 """Routed branch only (no shared expert)."""
1390 del module, hidden_states
1391 return routed
1394def _merge_shared_additive(
1395 module: Any,
1396 hidden_states: torch.Tensor,
1397 routed: torch.Tensor,
1398) -> torch.Tensor:
1399 """Compose ``routed + module.shared_experts(x)`` (DeepSeek-V3 / GLM-4-MoE).
1401 ``module.shared_experts(...)`` is a planned nested TP boundary: its exit
1402 already performs the TP reduction (nested-boundary call contract) — do NOT
1403 all-reduce its return value.
1404 """
1405 return routed + module.shared_experts(hidden_states)
1408def _merge_shared_gated(
1409 module: Any,
1410 hidden_states: torch.Tensor,
1411 routed: torch.Tensor,
1412) -> torch.Tensor:
1413 """Compose ``routed + sigmoid(shared_expert_gate(x)) * shared_expert(x)``.
1415 Qwen2-MoE layout. ``module.shared_expert(...)`` is a planned nested TP
1416 boundary — the reduction contract is documented on
1417 :func:`_merge_shared_additive`.
1418 """
1419 shared = module.shared_expert(hidden_states)
1420 gate = torch.sigmoid(module.shared_expert_gate(hidden_states))
1421 return routed + gate * shared
1424# The three shared-branch merges for the codegen inline EP body
1425# (``moe_ep_forward``), keyed by STRUCTURAL mode. They mirror the combine
1426# closures the native archetype factories in ``recipes.py`` define inline;
1427# codegen imports callable-free plans, so the generated body selects the
1428# composition here by string key.
1429SHARED_MERGE_MODES: dict[str, Callable[[Any, torch.Tensor, torch.Tensor], torch.Tensor]] = {
1430 "none": _merge_shared_none,
1431 "additive": _merge_shared_additive,
1432 "gated": _merge_shared_gated,
1433}
1436def moe_ep_forward(
1437 module: Any,
1438 hidden_states: torch.Tensor,
1439 *,
1440 router_kind: str,
1441 shared: str = "none",
1442 ep_group: Any,
1443) -> torch.Tensor:
1444 """Shared routed-MoE forward surfaced to generated inline EP artifacts.
1446 The single framework-generic EP body used across model families in
1447 place of per-family strategy strings: runs the routed branch and merges
1448 the shared branch according to ``shared``. Router semantics are selected
1449 by ``router_kind`` — a STRUCTURAL key (``softmax_topk`` /
1450 ``topk_router_module`` / ``sigmoid_group``), never a model name.
1452 ``shared`` selects an entry of :data:`SHARED_MERGE_MODES`, mirroring
1453 the combine closures the native archetype factories define inline:
1454 * ``"none"`` — routed output only.
1455 * ``"additive"`` — ``routed + module.shared_experts(x)``
1456 (DeepSeek-V3 / GLM-4-MoE).
1457 * ``"gated"`` — ``routed + sigmoid(module.shared_expert_gate(x)) *
1458 module.shared_expert(x)`` (Qwen2-MoE).
1460 Keeps the same numeric contract as ``ep_routed_forward`` plus the
1461 ``combine`` step of the runtime EP recipes.
1462 """
1463 merge = SHARED_MERGE_MODES.get(shared)
1464 if merge is None:
1465 raise ValueError(
1466 f"moe_ep_forward: unknown shared merge mode {shared!r} "
1467 f"(expected one of {sorted(SHARED_MERGE_MODES)})"
1468 )
1469 routed = ep_routed_forward(
1470 module,
1471 hidden_states,
1472 router_fn=MOE_ROUTER_ADAPTERS[router_kind],
1473 ep_group=ep_group,
1474 )
1475 return merge(module, hidden_states, routed)
1478def require_attrs(module: Any, *names: str, owner: str = "") -> None:
1479 """Assert that ``module`` has every attribute in ``names``; raise a
1480 teaching ValueError listing the module's ACTUAL children otherwise.
1482 Used by EP compute factories (archetype or user-written) at apply time:
1483 the factory body runs ONCE at apply time (before the wrapped forward
1484 ever executes), so an interface mismatch fails the model build in
1485 seconds instead of surfacing as a runtime AttributeError at step N.
1486 The check is structural (the names the implementation will call exist);
1487 it does not prove semantic correctness — that is vouched by numeric
1488 verification.
1489 """
1490 missing = [n for n in names if not hasattr(module, n)]
1491 if not missing:
1492 return
1493 children = [n for n, _ in module.named_children()]
1494 who = f"{owner} " if owner else ""
1495 raise ValueError(
1496 f"{who}expects MoE module attribute(s) {missing} on "
1497 f"{type(module).__name__}, but they do not exist; the module's "
1498 f"actual children are {children}. Pick the matching EP archetype "
1499 f"(see the archetype table in ep_compute.py), or write your own "
1500 f"factory (reference: examples/distributed/ep_factories.py) — the "
1501 f"names your compute_fn calls on module.<child> must match the "
1502 f"model's actual attribute names"
1503 )
1506def describe_moe_module(module: Any) -> str:
1507 """Structural diagnostic for a MoE module: child submodules, direct
1508 parameter shapes, and expert-related attributes — the facts needed to
1509 pick an EP archetype or write a custom factory. Returns the report;
1510 also logged at INFO."""
1511 lines = [f"MoE module: {type(module).__name__}"]
1512 children = list(module.named_children())
1513 lines.append(f"children ({len(children)}):")
1514 for name, child in children:
1515 lines.append(f" - {name}: {type(child).__name__}")
1516 direct_params = list(module.named_parameters(recurse=False))
1517 if direct_params:
1518 lines.append("direct parameters:")
1519 for name, p in direct_params:
1520 lines.append(f" - {name}: shape={tuple(p.shape)}")
1521 experts = getattr(module, "experts", None)
1522 if experts is not None:
1523 lines.append(
1524 f"experts: {type(experts).__name__}, "
1525 f"local_expert_count={getattr(experts, 'local_expert_count', '<unset>')}"
1526 )
1527 report = "\n".join(lines)
1528 return report