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