Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / module.py: 89%
277 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Model-facing Torch MegaMoe expert module."""
17from __future__ import annotations
19__all__ = ["MegaMoeExperts"]
21import math
22import struct
23from collections.abc import Mapping
24from dataclasses import replace
25from typing import Any, Literal
27import torch
28import torch.distributed as dist
30from hyper_multicore import _collective, shmem
32from ..module import MulticoreModule
33from .function import execute_mega_moe_with_permutation
34from .heap_manager import _get_heap_manager, _workspace_layout
35from .plan import build_mega_moe_plan
36from .route import prepare_topk_route, restore_topk_output
37from .spec import (
38 _COMMUNICATION_SPLIT,
39 _resolve_capacity_factors,
40 bind_mega_moe_spec,
41)
42from .workspace import MegaMoeWorkspace
45ExpertWeightLayout = Literal["EIO", "EOI"]
46_EXPERT_WEIGHT_LAYOUTS = frozenset(("EIO", "EOI"))
49def _validate_swiglu_limit(swiglu_limit: float | None) -> None:
50 """Validate a positive clamp value that survives float32 serialization."""
51 if swiglu_limit is None:
52 return
53 valid_type = isinstance(swiglu_limit, (int, float)) and not isinstance(
54 swiglu_limit, bool
55 )
56 try:
57 encoded_limit = (
58 struct.unpack("<f", struct.pack("<f", float(swiglu_limit)))[0]
59 if valid_type
60 else 0.0
61 )
62 valid_value = (
63 valid_type
64 and math.isfinite(swiglu_limit)
65 and math.isfinite(encoded_limit)
66 and encoded_limit > 0
67 )
68 except (OverflowError, TypeError, ValueError, struct.error):
69 valid_value = False
70 if not valid_value:
71 raise ValueError(
72 "swiglu_limit must be None or a finite positive float32-representable number, "
73 f"got {swiglu_limit!r}."
74 )
77def _create_mega_moe_parameters(
78 local_experts: int,
79 hidden_size: int,
80 intermediate_size: int,
81 expert_weight_layout: ExpertWeightLayout,
82) -> tuple[torch.nn.Parameter, torch.nn.Parameter]:
83 """Create independently initialized local expert parameters."""
84 if expert_weight_layout == "EOI":
85 gate_up_shape = (local_experts, intermediate_size * 2, hidden_size)
86 down_shape = (local_experts, hidden_size, intermediate_size)
87 else:
88 gate_up_shape = (local_experts, hidden_size, intermediate_size * 2)
89 down_shape = (local_experts, intermediate_size, hidden_size)
90 gate_up_weight = torch.nn.Parameter(
91 torch.empty(gate_up_shape)
92 )
93 down_weight = torch.nn.Parameter(
94 torch.empty(down_shape)
95 )
96 for expert_index in range(local_experts):
97 torch.nn.init.xavier_uniform_(gate_up_weight[expert_index])
98 torch.nn.init.xavier_uniform_(down_weight[expert_index])
99 return gate_up_weight, down_weight
102def _resolve_ep_size(ep_size: int | None, ep_group: Any | None) -> int:
103 """Infer size from an explicit group without silently choosing WORLD."""
104 if ep_size is not None and (not isinstance(ep_size, int) or isinstance(ep_size, bool) or ep_size <= 0):
105 raise ValueError("ep_size must be a positive integer or None")
106 if ep_group is not None:
107 if not dist.is_initialized():
108 raise ValueError("Initialize torch.distributed before passing ep_group")
109 group_size = dist.get_world_size(ep_group)
110 if ep_size is not None and ep_size != group_size:
111 raise ValueError(f"ep_size={ep_size} does not match ep_group size={group_size}")
112 return group_size
113 if dist.is_initialized() and dist.get_world_size() > 1:
114 raise ValueError("Pass ep_group explicitly for distributed MegaMoeExperts")
115 return 1 if ep_size is None else ep_size
118class _MegaMoeExecutionResources:
119 """Own one shape-bound plan and workspace in the shared SHMEM lifecycle."""
121 def __init__(
122 self,
123 specification: dict[str, Any],
124 tensor: torch.Tensor,
125 *,
126 shared: bool,
127 active_specifications: tuple[Mapping[str, Any], ...],
128 ) -> None:
129 """Bind resources once to the first NPU tensor."""
130 self.spec = bind_mega_moe_spec(specification, tensor)
131 self._token_capacity = self.spec.local_num_tokens
132 self._plans = {}
133 if dist.is_initialized():
134 self._validate_distributed_configuration(active_specifications)
135 reservations = tuple((item, _workspace_layout(item, tensor.dtype)) for item in active_specifications)
136 self.heap_manager = _get_heap_manager(self.spec.ep_group, tensor)
137 added_specifications = self.heap_manager.reserve(
138 reservations,
139 allow_resize=any(item.get("dispatch_mode", "push") == "push" or item.get("auto_tokens", False)
140 for item in active_specifications),
141 )
142 attached = False
143 self.plan = None
144 self.workspace = None
145 try:
146 attached = self.heap_manager.attach_runtime()
148 def construct() -> None:
149 """Build and bind local resources before converging construction results."""
150 self.plan = (build_mega_moe_plan(self.spec, tensor.device)
151 if not self.spec.dynamic_tokens else None)
152 self.workspace = MegaMoeWorkspace(shared=shared)
153 self.workspace.capacity_floor = self.heap_manager.bind(self.workspace, self.spec, specification)
154 self.workspace.heap_manager = self.heap_manager
156 _collective.collective_call(
157 construct, stage="MegaMoe resource construction", group=self.heap_manager.root_group,
158 )
159 except Exception:
160 self.heap_manager.rollback_reservations(added_specifications, self.workspace)
161 if attached and not self.heap_manager.workspaces:
162 self.heap_manager.detach_runtime()
163 raise
164 self._closed = False
165 self._workspace_closed = False
167 def ensure_capacity(self, maximum_received_slots: int) -> None:
168 """Submit a larger layout, falling back to minimum capacity within an explicit budget.
170 Args:
171 maximum_received_slots: Maximum destination load from the root count exchange.
172 """
173 self.heap_manager.ensure_capacity(self.workspace, maximum_received_slots)
175 def lifecycle_signature(self) -> tuple[Any, ...]:
176 """Identify allocations independently of process-local group addresses."""
177 return (
178 "mega_moe", self.spec.local_num_tokens, self.spec.hidden_size,
179 self.spec.intermediate_size, self.spec.num_experts, self.spec.top_k,
180 self.spec.receive_capacity, self.spec.ep_size, self.spec.dispatch_mode,
181 self.spec.expert_weight_layout,
182 self.spec.initial_capacity_factor, self.spec.capacity_growth_factor,
183 self.spec.dispatch_split, self.spec.combine_split, self.spec.max_local_num_tokens, self.spec.auto_tokens,
184 )
186 def prepare_auto_workspace(self, tokens: int, received: int, tensor: torch.Tensor) -> None:
187 """Commit physical capacity before taking a pull lease and publishing source rows.
189 Args:
190 tokens: Aligned group-maximum actual input length.
191 received: Group-maximum received route rows.
192 tensor: Input supplying device and dtype.
193 """
194 physical_spec = self.heap_manager.ensure_token_capacity(self.workspace, tokens, received)
195 self.workspace.ensure(physical_spec, tensor.dtype, tensor.device)
196 self._token_capacity = physical_spec.local_num_tokens
197 if self.spec.dispatch_mode == "pull":
198 self.workspace.claim()
200 def can_close(self) -> bool:
201 """Report whether neither an active call nor a backward graph needs buffers."""
202 return self._workspace_closed or self.workspace.can_close()
204 def _validate_distributed_configuration(self, specifications: tuple[Mapping[str, Any], ...]) -> None:
205 """Reject different allocation manifests before acquiring symmetric memory."""
206 if self.spec.ep_size == 1:
207 return
208 fields = ("local_num_tokens", "max_local_num_tokens", "hidden_size", "intermediate_size",
209 "num_experts", "top_k", "initial_capacity_factor", "capacity_growth_factor", "dispatch_mode",
210 "dispatch_split", "combine_split", "swiglu_limit", "expert_weight_layout", "ep_size", "auto_tokens")
211 # BF16 is the sole supported dtype; the literal tags also identify the protocol.
212 manifest = ("dynamic_tokens_v3", "bfloat16",
213 tuple(tuple(spec.get(key) for key in fields) for spec in specifications))
214 # Resource groups must have the same order as their symmetric allocations.
215 manifests = [None] * self.spec.ep_size
216 dist.all_gather_object(manifests, manifest, group=self.spec.ep_group)
217 if any(value != manifests[0] for value in manifests):
218 raise ValueError("MegaMoe resource configurations differ across EP ranks")
220 def plan_for_tokens(self, tokens: int, device: Any) -> Any:
221 """Reuse one covering plan per geometric bucket within committed capacity.
223 Args:
224 tokens: EP group maximum input rows, already aligned to 128.
225 device: Execution device used for the parameters and workspace.
227 Returns:
228 Immutable plan covering the requested rows without enlarging storage.
229 """
230 if tokens > self._token_capacity:
231 raise RuntimeError(f"MegaMoe schedule exceeds committed token capacity: "
232 f"requested={tokens}, capacity={self._token_capacity}")
233 bucket = 1 << (max(128, tokens) - 1).bit_length()
234 plan = self._plans.get(bucket)
235 if plan is None or plan.spec.plan_tokens < tokens:
236 schedule_tokens = min(bucket, self._token_capacity)
237 plan = build_mega_moe_plan(replace(self.spec, schedule_tokens=schedule_tokens), device)
238 # INT32 route bounds allow at most 25 buckets starting at 128 rows.
239 # A growing partial bucket replaces its plan; backward keeps the old object.
240 self._plans[bucket] = plan
241 return plan
243 def close(self) -> None:
244 """Release the workspace and leave the shared SHMEM lifecycle."""
245 if self._closed:
246 return
247 if not self._workspace_closed:
248 self.heap_manager.close_workspace(self.workspace)
249 self._workspace_closed = True
250 self._plans.clear()
251 self._closed = True
254class MegaMoeExperts(MulticoreModule):
255 """Execute Router-selected local experts with the Torch MegaMoe kernel.
257 Serial model layers can call :meth:`share_execution_resources` before
258 their first forward to share one automatically growing workspace while keeping
259 independent parameters and optimizer state.
260 """
262 @classmethod
263 def from_module(
264 cls, module: torch.nn.Module, *, top_k: int, ep_group: Any | None = None, **execution_options: Any,
265 ) -> MegaMoeExperts:
266 """Convert materialized local GroupedExperts SwiGLU weights before optimizer construction.
268 Args:
269 module: Module exposing w1/w3 [local_E, I, H], w2 [local_E, H, I].
270 DTensor local EP shards are accepted; FSDP weights must be materialized first.
271 top_k: Experts selected by the existing router for each token.
272 ep_group: Explicit EP group owning these local experts, or None for a standalone rank.
273 **execution_options: Advanced MegaMoe options, including num_experts and expert_weight_layout.
275 Returns:
276 Experts with copied, converted Parameters. Optimizer state is not migrated.
277 """
278 values = []
279 for name in ("w1", "w3", "w2"):
280 value = getattr(module, name, None)
281 if hasattr(value, "to_local"):
282 value = value.to_local()
283 if not isinstance(value, torch.Tensor) or value.ndim != 3:
284 raise ValueError(f"source {name} must be a materialized local rank-three expert tensor")
285 if value.is_meta:
286 raise ValueError(f"source {name} is meta; materialize and initialize/load weights before from_module")
287 values.append(value)
288 gate, up, down = values[0], values[1], values[2]
289 if (up.shape != gate.shape or down.shape != (gate.shape[0], gate.shape[2], gate.shape[1])
290 or any(value.device != gate.device or value.dtype != gate.dtype for value in values)
291 or gate.requires_grad != up.requires_grad):
292 raise ValueError("source SwiGLU expert shapes, dtypes, devices and gate/up gradient policies must match")
293 ep_size = _resolve_ep_size(execution_options.get("ep_size"), ep_group)
294 experts = execution_options.pop("num_experts", getattr(module, "num_experts", gate.shape[0] * ep_size))
295 if experts != gate.shape[0] * ep_size:
296 raise ValueError("source local expert count does not match num_experts and ep_group")
297 result = cls(hidden_size=gate.shape[2], intermediate_size=gate.shape[1], num_experts=experts,
298 top_k=top_k, ep_group=ep_group, create_parameters=False, **execution_options)
299 gate_up = torch.cat((gate.detach(), up.detach()), dim=1)
300 down_data = down.detach()
301 if result.expert_weight_layout == "EIO":
302 gate_up = gate_up.transpose(1, 2)
303 down_data = down_data.transpose(1, 2)
304 result.gate_up_weight = torch.nn.Parameter(
305 gate_up.contiguous(), requires_grad=gate.requires_grad)
306 result.down_weight = torch.nn.Parameter(
307 down_data.clone(memory_format=torch.contiguous_format),
308 requires_grad=down.requires_grad)
309 return result
311 def _resource_signature(self, tensor: Any) -> Any:
312 """Compare static configuration without process-local group identities."""
313 configuration = tuple(sorted((name, value) for name, value in self._resource_group.specification.items()
314 if name != "ep_group"))
315 return (super()._resource_signature(tensor), configuration,
316 self._resource_group.shared)
318 def _register_runtime_release(self, manager: Any, resources: Any) -> None:
319 """Register the Heap Manager's sole Runtime reference for final cleanup."""
320 manager.runtime_release = resources.heap_manager.detach_runtime
322 def __init__(
323 self,
324 *,
325 hidden_size: int,
326 intermediate_size: int,
327 num_experts: int,
328 top_k: int,
329 swiglu_limit: float | None = None,
330 ep_size: int | None = None,
331 ep_group: Any | None = None,
332 create_parameters: bool = True,
333 expert_weight_layout: ExpertWeightLayout = "EIO",
334 dispatch_mode: str = "push",
335 initial_capacity_factor: float | None = None,
336 capacity_growth_factor: float | None = None,
337 dispatch_split: int = _COMMUNICATION_SPLIT,
338 combine_split: int = _COMMUNICATION_SPLIT,
339 ) -> None:
340 """Initialize local expert parameters and a lazy execution owner.
342 Args:
343 hidden_size: Input and output hidden dimension.
344 intermediate_size: SwiGLU intermediate dimension per expert.
345 num_experts: Global routed-expert count.
346 top_k: Experts selected for every token.
347 initial_capacity_factor: Push initial receive rows as a multiple of local routed rows.
348 Defaults to 1.25 for push; omitted for pull. Must be finite and at least 1.0.
349 capacity_growth_factor: Push capacity multiplier on overflow, defaulting to 1.25.
350 Must be finite and at least 1.0; 1.0 grows only to the current route demand.
351 The resulting capacity is capped by the lossless route bound. Pull rejects explicit factors.
352 dispatch_mode: Dispatch transport, either "push" (default) or "pull".
353 Construct separate modules to switch modes; sharing requires equal modes.
354 dispatch_split: Dispatch tile rows in forward and backward, default 128.
355 Must be a positive multiple of 128.
356 combine_split: Combine tile rows in forward and backward, with the same
357 constraints and default as dispatch_split. Sharing requires equal splits.
358 swiglu_limit: Optional positive, finite float32-representable clamp
359 limit for SwiGLU. The gate branch uses ``min(gate, limit)`` and
360 the up branch is clamped to ``[-limit, limit]``. ``None``
361 preserves the legacy unclamped path.
362 ep_size: Optional compatibility check; normally inferred from explicit ep_group.
363 create_parameters: Allocate owned weights; False requires explicit expert_weights each forward.
364 expert_weight_layout: Logical expert-weight axis order. ``"EIO"`` means
365 ``[experts, in_features, out_features]`` and preserves the historical layout. ``"EOI"`` means
366 ``[experts, out_features, in_features]`` and accepts Torch linear storage directly. ``EOI`` weights
367 must use dense contiguous storage for zero-copy execution.
368 ep_group: Explicit EP group; expert ownership follows group-local rank order.
369 Disjoint PP/DP groups bootstrap independently. Only one ordered EP membership
370 can be active in SHMEM per process.
371 """
372 specification = {
373 "local_num_tokens": _COMMUNICATION_SPLIT,
374 "hidden_size": hidden_size,
375 "intermediate_size": intermediate_size,
376 "num_experts": num_experts,
377 "top_k": top_k,
378 "initial_capacity_factor": initial_capacity_factor,
379 "swiglu_limit": swiglu_limit,
380 "ep_size": _resolve_ep_size(ep_size, ep_group),
381 "ep_group": ep_group,
382 "expert_weight_layout": expert_weight_layout,
383 "dispatch_mode": dispatch_mode,
384 "capacity_growth_factor": capacity_growth_factor,
385 "dispatch_split": dispatch_split,
386 "combine_split": combine_split,
387 }
388 specification["auto_tokens"] = True
389 specification = self._resolve_specification(specification)
390 compatibility_key = tuple(
391 id(value) if name == "ep_group" else value for name, value in specification.items())
392 super().__init__(
393 resource_specification=specification,
394 resource_compatibility_key=compatibility_key,
395 resource_scope_key=("mega_moe", shmem.root_members(ep_group) if dist.is_initialized() else id(ep_group)),
396 )
397 self._initialize_module_state(specification, create_parameters=create_parameters)
399 def _initialize_module_state(
400 self, specification: Mapping[str, Any], *, create_parameters: bool,
401 ) -> None:
402 """Bind normalized configuration and initialize expert parameters."""
403 self.hidden_size = specification["hidden_size"]
404 self.intermediate_size = specification["intermediate_size"]
405 self.num_experts = specification["num_experts"]
406 self.top_k = specification["top_k"]
407 self.expert_weight_layout = specification["expert_weight_layout"]
408 self.initial_capacity_factor = specification["initial_capacity_factor"]
409 self.capacity_growth_factor = specification["capacity_growth_factor"]
410 self.dispatch_mode = specification["dispatch_mode"]
411 self.dispatch_split = specification["dispatch_split"]
412 self.combine_split = specification["combine_split"]
413 self.swiglu_limit = specification["swiglu_limit"]
414 self.ep_size = specification["ep_size"]
415 self.local_experts = self.num_experts // self.ep_size
416 self._ep_group = specification["ep_group"]
417 if create_parameters:
418 self.gate_up_weight, self.down_weight = _create_mega_moe_parameters(
419 self.local_experts, self.hidden_size, self.intermediate_size, self.expert_weight_layout,
420 )
421 else:
422 self.register_parameter("gate_up_weight", None)
423 self.register_parameter("down_weight", None)
425 @classmethod
426 def _resolve_specification(cls, specification: dict[str, Any]) -> dict[str, Any]:
427 """Validate and normalize configuration before registering an execution owner."""
428 expert_weight_layout = specification["expert_weight_layout"]
429 if expert_weight_layout not in _EXPERT_WEIGHT_LAYOUTS:
430 raise ValueError(
431 "expert_weight_layout must be 'EIO' ([experts, in_features, out_features]) "
432 "or 'EOI' ([experts, out_features, in_features]), "
433 f"got {expert_weight_layout!r}."
434 )
435 initial_capacity_factor, capacity_growth_factor = _resolve_capacity_factors(
436 specification["dispatch_mode"],
437 specification["initial_capacity_factor"],
438 specification["capacity_growth_factor"],
439 )
440 cls._validate_topology(
441 local_num_tokens=specification["local_num_tokens"],
442 hidden_size=specification["hidden_size"],
443 intermediate_size=specification["intermediate_size"],
444 num_experts=specification["num_experts"],
445 top_k=specification["top_k"],
446 swiglu_limit=specification["swiglu_limit"],
447 ep_size=specification["ep_size"],
448 dispatch_split=specification["dispatch_split"],
449 combine_split=specification["combine_split"],
450 auto_tokens=specification.get("auto_tokens", False),
451 )
452 swiglu_limit = specification["swiglu_limit"]
453 return {
454 **specification,
455 "initial_capacity_factor": initial_capacity_factor,
456 "capacity_growth_factor": capacity_growth_factor,
457 "swiglu_limit": float(swiglu_limit) if swiglu_limit is not None else None,
458 }
460 @staticmethod
461 def _validate_topology(
462 *,
463 local_num_tokens: int,
464 hidden_size: int,
465 intermediate_size: int,
466 num_experts: int,
467 top_k: int,
468 swiglu_limit: float | None,
469 ep_size: int,
470 dispatch_split: int,
471 combine_split: int,
472 auto_tokens: bool = False,
473 ) -> None:
474 """Validate static shape and topology values before allocation."""
475 values = {
476 "local_num_tokens": local_num_tokens,
477 "hidden_size": hidden_size,
478 "intermediate_size": intermediate_size,
479 "num_experts": num_experts,
480 "top_k": top_k,
481 "ep_size": ep_size,
482 }
483 for name, value in values.items():
484 if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
485 raise ValueError(f"{name} must be a positive integer, got {value!r}.")
486 if top_k > num_experts:
487 raise ValueError(
488 f"top_k ({top_k}) cannot exceed num_experts ({num_experts})."
489 )
490 if num_experts % ep_size:
491 raise ValueError(
492 f"num_experts ({num_experts}) must be divisible by ep_size ({ep_size})."
493 )
494 _validate_swiglu_limit(swiglu_limit)
495 for name, split in (("dispatch_split", dispatch_split), ("combine_split", combine_split)):
496 if (not isinstance(split, int) or isinstance(split, bool)
497 or split <= 0 or split % _COMMUNICATION_SPLIT):
498 raise ValueError(
499 f"{name} must be a positive integer multiple of {_COMMUNICATION_SPLIT}, got {split!r}.")
500 if not auto_tokens and local_num_tokens % split:
501 raise ValueError(f"local_num_tokens ({local_num_tokens}) must be divisible by {name} ({split}).")
503 def _root_group(self) -> Any:
504 """Keep lifecycle coordination within this module's EP group."""
505 return self._ep_group
507 def _validate_tensors(self, hidden_states: torch.Tensor, expert_weights: tuple) -> None:
508 """Validate activation and parameter metadata before acquiring resources."""
509 if hidden_states.dtype != torch.bfloat16 or not getattr(hidden_states, "is_npu", False):
510 raise TypeError("MegaMoeExperts requires BF16 NPU hidden states.")
511 weight1, weight2 = expert_weights
512 if weight1 is None or weight2 is None:
513 raise ValueError("Parameterless MegaMoe requires explicit expert_weights")
514 if self.expert_weight_layout == "EOI":
515 expected_weight1 = (self.local_experts, self.intermediate_size * 2, self.hidden_size)
516 expected_weight2 = (self.local_experts, self.hidden_size, self.intermediate_size)
517 else:
518 expected_weight1 = (self.local_experts, self.hidden_size, self.intermediate_size * 2)
519 expected_weight2 = (self.local_experts, self.intermediate_size, self.hidden_size)
520 if tuple(weight1.shape) != expected_weight1:
521 raise ValueError(
522 f"gate_up_weight must have shape {expected_weight1}, got {tuple(weight1.shape)}."
523 )
524 if tuple(weight2.shape) != expected_weight2:
525 raise ValueError(
526 f"down_weight must have shape {expected_weight2}, got {tuple(weight2.shape)}."
527 )
528 if (
529 weight1.device != hidden_states.device
530 or weight2.device != hidden_states.device
531 ):
532 raise ValueError(
533 "expert weights and hidden states must be on the same NPU."
534 )
535 if weight1.dtype != hidden_states.dtype or weight2.dtype != hidden_states.dtype:
536 raise TypeError("expert weights and hidden states must all use BF16.")
537 if self.expert_weight_layout == "EOI" and (
538 not weight1.is_contiguous() or not weight2.is_contiguous()
539 ):
540 raise ValueError("EOI expert weights must be contiguous for zero-copy execution.")
542 def _validate_forward_inputs(
543 self,
544 hidden_states: torch.Tensor,
545 topk_ids: torch.Tensor,
546 topk_weights: torch.Tensor,
547 tokens_per_expert: torch.Tensor | None,
548 expert_weights: tuple,
549 ) -> torch.Tensor:
550 """Validate input shapes and return flattened local token states."""
551 if hidden_states.ndim < 2 or hidden_states.shape[-1] != self.hidden_size:
552 raise ValueError(
553 "hidden_states must have at least two dimensions and configured "
554 f"hidden size {self.hidden_size}, got {tuple(hidden_states.shape)}."
555 )
556 hidden_flat = hidden_states.reshape(-1, self.hidden_size)
557 route_shape = (hidden_flat.shape[0], self.top_k)
558 if tuple(topk_ids.shape) != route_shape:
559 raise ValueError(
560 f"topk_ids must have shape {route_shape}, got {tuple(topk_ids.shape)}."
561 )
562 if tuple(topk_weights.shape) != route_shape:
563 raise ValueError(
564 f"topk_weights must have shape {route_shape}, "
565 f"got {tuple(topk_weights.shape)}."
566 )
567 if tokens_per_expert is not None and tuple(tokens_per_expert.shape) != (
568 self.num_experts,
569 ):
570 raise ValueError(
571 f"tokens_per_expert must have shape ({self.num_experts},), "
572 f"got {tuple(tokens_per_expert.shape)}."
573 )
574 self._validate_tensors(hidden_flat, expert_weights)
575 return hidden_flat
577 def forward(
578 self,
579 hidden_states: torch.Tensor,
580 topk_ids: torch.Tensor,
581 topk_weights: torch.Tensor,
582 *,
583 tokens_per_expert: torch.Tensor | None = None,
584 expert_weights: tuple[torch.Tensor, torch.Tensor] | None = None,
585 ) -> torch.Tensor:
586 """Route token states through the configured experts.
588 Args:
589 hidden_states: Tensor ending in ``hidden_size``.
590 topk_ids: Global expert IDs shaped ``[tokens, top_k]`` for the current input.
591 topk_weights: Router weights with the same shape as ``topk_ids``.
592 expert_weights: Current unsharded weights owned by an outer module.
593 Noncontiguous EIO weights are copied with autograd tracking, never cached
594 by this executor. EOI weights must be contiguous for zero-copy execution.
595 tokens_per_expert: Optional trusted exact global-expert histogram.
596 Supplying Router-produced counts skips histogram recomputation.
598 Returns:
599 Tensor matching the shape and dtype of ``hidden_states``.
601 Raises:
602 ValueError: EOI expert weights are noncontiguous.
603 RuntimeError: The current NPU stream is being captured.
605 Note:
606 Automatic token management requires eager execution for both push
607 and pull. NPU graph capture is unsupported, even with stable token
608 counts and capacity.
610 ``tokens_per_expert`` must describe the current ``topk_ids`` exactly.
611 The steady-state path validates metadata but deliberately does not
612 rebuild and compare the histogram.
613 """
614 weights = ((self.gate_up_weight, self.down_weight) if expert_weights is None else expert_weights)
615 hidden_flat = self._validate_forward_inputs(
616 hidden_states,
617 topk_ids,
618 topk_weights,
619 tokens_per_expert,
620 weights,
621 )
622 if torch.npu.is_current_stream_capturing():
623 raise RuntimeError("MegaMoe push heap growth and dynamic tokens do not support graph capture")
624 resources = self._get_execution_resources(hidden_flat)
625 route_options = {"autograd_mask": (
626 int(torch.is_grad_enabled()) | (int(hidden_flat.requires_grad) << 1)
627 | (int(weights[0].requires_grad) << 2) | (int(weights[1].requires_grad) << 3)
628 | (int(topk_weights.requires_grad) << 4)
629 )}
630 pull = self.dispatch_mode == "pull"
631 automatic = resources.spec.auto_tokens
632 claimed = False
633 def prepare_workspace(tokens: int, received: int) -> None:
634 """Grow before publishing pull source rows into symmetric memory.
636 Args:
637 tokens: Current local source-token count.
638 received: Current routed receive-row count.
639 """
640 nonlocal claimed
641 resources.prepare_auto_workspace(tokens, received, hidden_flat)
642 claimed = pull
644 if automatic:
645 route_options["prepare_workspace"] = prepare_workspace
646 elif pull:
647 resources.workspace.ensure(resources.spec, hidden_flat.dtype, hidden_flat.device)
648 resources.workspace.claim()
649 claimed = True
650 try:
651 # Pull source writes and execution must share one workspace lease.
652 with torch.no_grad():
653 route = prepare_topk_route(
654 hidden_flat,
655 topk_ids,
656 topk_weights,
657 resources.spec,
658 tokens_per_expert,
659 workspace=resources.workspace,
660 **route_options,
661 )
662 if not pull and not automatic:
663 resources.ensure_capacity(route.maximum_received_slots)
664 plan = resources.plan_for_tokens(route.metadata.plan_tokens, hidden_flat.device)
665 expert_output = execute_mega_moe_with_permutation(
666 hidden_flat,
667 topk_ids,
668 weights[0],
669 weights[1],
670 route,
671 plan,
672 resources.workspace,
673 topk_weights=topk_weights if pull else None,
674 workspace_claimed=pull,
675 )
676 finally:
677 if claimed:
678 resources.workspace.release()
679 if pull:
680 return expert_output.reshape_as(hidden_states)
681 output = restore_topk_output(
682 expert_output,
683 route.unpermute_mapping,
684 topk_weights,
685 )
686 return output.reshape(hidden_states.shape)
688 def _create_execution_resources(
689 self,
690 tensor: torch.Tensor,
691 *,
692 shared: bool,
693 active_specifications: tuple[Mapping[str, Any], ...],
694 ) -> _MegaMoeExecutionResources:
695 """Create the plan and workspace owned by this resource group."""
696 return _MegaMoeExecutionResources(
697 self._resource_group.specification,
698 tensor,
699 shared=shared,
700 active_specifications=active_specifications,
701 )