Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / spec.py: 67%
101 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"""Validated shape, topology, and hardware specification for MegaMoe."""
17from __future__ import annotations
19import math
20from collections.abc import Mapping
21from dataclasses import dataclass
22from typing import Any
24import torch
25import torch.distributed as dist
27from hyper_multicore._core_spec import get_device_core_spec
30_COMMUNICATION_SPLIT = 128
31_DEFAULT_CAPACITY_FACTOR = 1.25
34@dataclass(frozen=True)
35class MegaMoeSpec:
36 """Shape, capacity, topology, and Device core limits bound on first use."""
38 local_num_tokens: int
39 hidden_size: int
40 intermediate_size: int
41 num_experts: int
42 top_k: int
43 initial_capacity_factor: float | None
44 receive_capacity: int
45 ep_size: int
46 ep_group: Any | None
47 rank_id: int
48 num_cube_cores: int
49 dispatch_mode: str = "push"
50 capacity_growth_factor: float | None = _DEFAULT_CAPACITY_FACTOR
51 dispatch_split: int = _COMMUNICATION_SPLIT
52 combine_split: int = _COMMUNICATION_SPLIT
53 swiglu_split: int = _COMMUNICATION_SPLIT
54 swiglu_limit: float | None = None
55 max_local_num_tokens: int | None = None
56 schedule_tokens: int | None = None
57 expert_weight_layout: str = "EIO"
58 auto_tokens: bool = False
60 @property
61 def dynamic_tokens(self) -> bool:
62 """Report variable-length execution with automatic or bounded storage."""
63 return self.auto_tokens or self.max_local_num_tokens is not None
65 @property
66 def plan_tokens(self) -> int:
67 """Return schedule coverage independently of reserved source storage."""
68 return self.local_num_tokens if self.schedule_tokens is None else self.schedule_tokens
70 @property
71 def local_experts(self) -> int:
72 """Return experts owned by this expert-parallel rank."""
73 return self.num_experts // self.ep_size
75 @property
76 def routed_slots(self) -> int:
77 """Return routed rows produced by local Top-K expansion."""
78 return self.local_num_tokens * self.top_k
81def _align_capacity(capacity: int) -> int:
82 """Align receive capacity to the baseline compute tile, independent of communication tiles."""
83 return (
84 (capacity + _COMMUNICATION_SPLIT - 1)
85 // _COMMUNICATION_SPLIT
86 * _COMMUNICATION_SPLIT
87 )
90def _resolve_receive_capacity(
91 initial_capacity_factor: float | None,
92 routed_slots: int,
93 ep_size: int,
94) -> int:
95 """Resolve initial receive rows without exceeding the lossless route bound."""
96 factor = 1.0 if initial_capacity_factor is None else initial_capacity_factor
97 requested = math.ceil(min(factor, ep_size) * routed_slots)
98 return _align_capacity(requested)
101def _grow_receive_capacity(spec: MegaMoeSpec, current: int, required: int) -> tuple[int, int]:
102 """Return geometric and minimum capacities without exceeding the lossless bound."""
103 upper = _align_capacity(spec.ep_size * spec.routed_slots)
104 if required > upper:
105 raise ValueError("MegaMoe route exceeds the lossless token bound")
106 requested = max(required, math.ceil(min(upper, current * spec.capacity_growth_factor)))
107 return min(upper, _align_capacity(requested)), _align_capacity(required)
110def _resolve_capacity_factors(
111 dispatch_mode: str,
112 initial_capacity_factor: float | None,
113 capacity_growth_factor: float | None,
114) -> tuple[float | None, float | None]:
115 """Resolve push defaults and reject capacity knobs for pull before allocation."""
116 if dispatch_mode not in ("push", "pull"):
117 raise ValueError("dispatch_mode must be push or pull")
118 if dispatch_mode == "pull":
119 if initial_capacity_factor is not None or capacity_growth_factor is not None:
120 raise ValueError("initial_capacity_factor and capacity_growth_factor are only supported for push")
121 return None, None
122 factors = {
123 "initial_capacity_factor": (
124 _DEFAULT_CAPACITY_FACTOR if initial_capacity_factor is None else initial_capacity_factor),
125 "capacity_growth_factor": (
126 _DEFAULT_CAPACITY_FACTOR if capacity_growth_factor is None else capacity_growth_factor),
127 }
128 for name, value in factors.items():
129 valid_type = isinstance(value, (int, float)) and not isinstance(value, bool)
130 try:
131 valid = valid_type and math.isfinite(value) and value >= 1.0
132 except OverflowError:
133 valid = False
134 if not valid:
135 raise ValueError(f"{name} must be a finite number at least 1.0, got {value!r}")
136 return float(factors["initial_capacity_factor"]), float(factors["capacity_growth_factor"])
139def bind_mega_moe_spec(
140 specification: Mapping[str, Any],
141 tensor: torch.Tensor,
142) -> MegaMoeSpec:
143 """Bind hardware and process-group values to the first NPU input.
145 Args:
146 specification: Constructor-validated static resource specification.
147 tensor: First input tensor used to resolve device-specific limits.
149 Returns:
150 Bound specification cached for subsequent forwards.
152 Raises:
153 ValueError: If the device or EP process group does not match the specification.
154 TypeError: If the activation dtype is unsupported.
155 RuntimeError: If the device does not report a valid Cube-core count.
156 """
157 device = torch.device(tensor.device)
158 if device.type != "npu":
159 raise ValueError(f"MegaMoeExperts requires an NPU input, got {device}.")
160 if tensor.dtype != torch.bfloat16:
161 raise TypeError(
162 f"MegaMoeExperts supports BF16 activations, got {tensor.dtype}."
163 )
165 ep_group = specification["ep_group"]
166 ep_size = specification["ep_size"]
167 if dist.is_initialized():
168 if ep_group is None and dist.get_world_size() > 1:
169 raise ValueError("Pass ep_group explicitly for distributed MegaMoeExperts")
170 group_size = dist.get_world_size(ep_group)
171 rank_id = dist.get_rank(ep_group)
172 else:
173 if ep_group is not None:
174 raise RuntimeError(
175 "ep_group cannot be used before torch.distributed initialization."
176 )
177 group_size = 1
178 rank_id = 0
179 if group_size != ep_size:
180 raise ValueError(
181 f"ep_group world size ({group_size}) does not match ep_size ({ep_size})."
182 )
184 device_index = (
185 device.index if device.index is not None else torch.npu.current_device()
186 )
187 core_spec = get_device_core_spec(device_index)
189 local_num_tokens = specification["local_num_tokens"]
190 top_k = specification["top_k"]
191 initial_capacity_factor = specification["initial_capacity_factor"]
192 swiglu_limit = specification["swiglu_limit"]
193 receive_capacity = _resolve_receive_capacity(
194 initial_capacity_factor,
195 local_num_tokens * top_k,
196 ep_size,
197 )
198 return MegaMoeSpec(
199 local_num_tokens=local_num_tokens,
200 hidden_size=specification["hidden_size"],
201 intermediate_size=specification["intermediate_size"],
202 num_experts=specification["num_experts"],
203 top_k=top_k,
204 initial_capacity_factor=initial_capacity_factor,
205 receive_capacity=receive_capacity,
206 ep_size=ep_size,
207 ep_group=ep_group,
208 rank_id=rank_id,
209 num_cube_cores=core_spec.num_cube_cores,
210 swiglu_limit=swiglu_limit,
211 dispatch_mode=specification.get("dispatch_mode", "push"),
212 expert_weight_layout=specification.get("expert_weight_layout", "EIO"),
213 capacity_growth_factor=specification.get("capacity_growth_factor", _DEFAULT_CAPACITY_FACTOR),
214 dispatch_split=specification.get("dispatch_split", _COMMUNICATION_SPLIT),
215 combine_split=specification.get("combine_split", _COMMUNICATION_SPLIT),
216 max_local_num_tokens=specification.get("max_local_num_tokens"),
217 auto_tokens=specification.get("auto_tokens", False),
218 )