Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / plan.py: 96%
79 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"""Build the static graph and tiling plan for managed MegaMoe execution."""
17from __future__ import annotations
19from dataclasses import dataclass
20from typing import Any
22import numpy as np
23import torch
24import torch.distributed as dist
26from hyper_multicore.diagnostics import MegaKernelDiagnostics, prepare_mega_kernel_diagnostics
27from hyper_multicore.modules.mega_moe.backward.gen_runtime_data import (
28 build_config_for_rank as build_backward_config,
29)
30from hyper_multicore.modules.mega_moe.backward.graph import (
31 build_backward_graph,
32)
33from hyper_multicore.modules.mega_moe.backward.storage import can_reuse_backward_dispatch
34from hyper_multicore.modules.mega_moe.backward.tiling_tables import (
35 get_act_grad_tiling_bytes,
36 get_gate_grad_tiling_bytes,
37 get_swiglu_grad_tiling_bytes,
38 get_w1_grad_tiling_bytes,
39 get_w2_grad_tiling_bytes,
40)
41from hyper_multicore.modules.mega_moe.forward.gen_runtime_data import (
42 build_config_for_rank as build_forward_config,
43)
44from hyper_multicore.modules.mega_moe.forward.graph import (
45 build_forward_graph,
46)
47from hyper_multicore.modules.mega_moe.forward.tiling_tables import (
48 get_down_proj_tiling_bytes,
49 get_swiglu_tiling_bytes,
50 get_up_proj_tiling_bytes,
51)
52from hyper_multicore.scheduler.config import TaskSplitValue
53from hyper_multicore.scheduler.graph import ComputeGraph
55from .diagnostics import _mega_moe_owner_options
56from .spec import MegaMoeSpec
57from .tensor_dump import attach_mega_moe_l1_view_resolvers
60@dataclass(frozen=True)
61class MegaMoePlan:
62 """Rank-local schedules and tiling tensors for forward and backward."""
64 spec: MegaMoeSpec
65 forward_graph: ComputeGraph
66 fwd_diagnostics: MegaKernelDiagnostics
67 up_proj_tiling: Any
68 swiglu_tiling: Any
69 down_proj_tiling: Any
70 backward_graph: ComputeGraph
71 bwd_diagnostics: MegaKernelDiagnostics
72 act_grad_tiling: Any
73 gate_grad_tiling: Any
74 w1_grad_tiling: Any
75 w2_grad_tiling: Any
76 swiglu_grad_tiling: Any
77 reuse_backward_dispatch: bool = False
79 @property
80 def fwd_runtime_config(self) -> Any:
81 """Return the original ordinary forward RuntimeConfig."""
82 return self.fwd_diagnostics.runtime_config
84 @property
85 def bwd_runtime_config(self) -> Any:
86 """Return the original ordinary backward RuntimeConfig."""
87 return self.bwd_diagnostics.runtime_config
90def record_plan_stream(plan: MegaMoePlan, runtime_config: Any) -> None:
91 """Protect cached tiling and runtime storage from cross-stream eviction."""
92 tensors = [value for value in vars(plan).values() if isinstance(value, torch.Tensor)]
93 tensors.append(runtime_config)
94 for tensor in tensors:
95 if isinstance(tensor, torch.Tensor) and tensor.device.type == "npu":
96 tensor.record_stream(torch.npu.current_stream(tensor.device))
99def _tensor_from_bytes(data: bytes, device: Any) -> torch.Tensor:
100 """Copy serialized runtime data into a device uint8 tensor."""
101 array = np.frombuffer(bytearray(data), dtype=np.uint8).copy()
102 return torch.from_numpy(array).to(device=device, dtype=torch.uint8)
105def _resize_swiglu_tiling(data: bytes, num_cube_cores: int) -> bytes:
106 """Resize the baseline 24-core repeated struct for the active workers."""
107 baseline_workers = 49
108 if len(data) % baseline_workers:
109 raise RuntimeError("invalid baseline SwiGLU tiling byte length.")
110 struct_size = len(data) // baseline_workers
111 return data[:struct_size] * (2 * num_cube_cores + 1)
114def _build_task_values(spec: MegaMoeSpec) -> TaskSplitValue:
115 """Translate rank-local token shape into legacy scheduler coordinates.
117 Args:
118 spec: Validated local-token and expert topology specification.
120 Returns:
121 Scheduler values whose single partition is the rank-local input.
122 """
123 return TaskSplitValue(
124 tp=1,
125 ep=spec.ep_size,
126 seq_size=spec.plan_tokens,
127 all_expert_num=spec.num_experts,
128 top_k=spec.top_k,
129 dispatch_mode=spec.dispatch_mode,
130 )
133def _build_runtime_artifacts(spec: MegaMoeSpec) -> tuple[Any, Any, Any, Any]:
134 """Build forward/backward graphs and their serialized RuntimeConfig objects."""
135 task_values = _build_task_values(spec)
136 forward_graph = build_forward_graph(
137 task_values,
138 dispatch_sv=spec.dispatch_split,
139 swiglu_sv=spec.swiglu_split,
140 combine_sv=spec.combine_split,
141 hidden_size=spec.hidden_size,
142 intermediate_size=spec.intermediate_size,
143 num_cube_cores=spec.num_cube_cores,
144 swiglu_limit=spec.swiglu_limit,
145 expert_weight_layout=spec.expert_weight_layout,
146 )
147 forward_graph.propagate_splits(task_values)
148 attach_mega_moe_l1_view_resolvers(forward_graph)
149 forward_data = build_forward_config(
150 forward_graph,
151 task_values,
152 spec.rank_id,
153 spec.num_cube_cores,
154 )
155 backward_graph = build_backward_graph(
156 task_values,
157 dispatch_sv=spec.dispatch_split,
158 swiglu_sv=spec.swiglu_split,
159 combine_sv=spec.combine_split,
160 hidden_size=spec.hidden_size,
161 intermediate_size=spec.intermediate_size,
162 num_cube_cores=spec.num_cube_cores,
163 swiglu_limit=spec.swiglu_limit,
164 expert_weight_layout=spec.expert_weight_layout,
165 )
166 backward_graph.propagate_splits(task_values)
167 attach_mega_moe_l1_view_resolvers(backward_graph)
168 backward_data = build_backward_config(
169 backward_graph,
170 task_values,
171 spec.rank_id,
172 spec.num_cube_cores,
173 )
174 return forward_graph, forward_data, backward_graph, backward_data
177def _prepare_diagnostics(
178 spec: MegaMoeSpec,
179 forward_graph: ComputeGraph,
180 forward_config: Any,
181 backward_graph: ComputeGraph,
182 backward_config: Any,
183 device: Any,
184) -> tuple[MegaKernelDiagnostics, MegaKernelDiagnostics]:
185 """Bind forward/backward plans to the public diagnostics integration API."""
186 topology = _build_task_values(spec)
187 topology.rank_id = spec.rank_id
188 options = {
189 "kernel_name": "mega_moe",
190 "device": device,
191 "rank": dist.get_rank() if dist.is_initialized() else spec.rank_id,
192 **_mega_moe_owner_options(topology, num_cube_cores=spec.num_cube_cores),
193 }
194 return (
195 prepare_mega_kernel_diagnostics(forward_graph, forward_config, parameter_count=16, **options),
196 prepare_mega_kernel_diagnostics(backward_graph, backward_config, parameter_count=20, **options),
197 )
200def build_mega_moe_plan(spec: MegaMoeSpec, device: Any) -> MegaMoePlan:
201 """Build trimmed dense forward/backward runtime images without fusion slots.
203 Args:
204 spec: Validated local-token and expert topology specification.
205 device: NPU device receiving serialized descriptors and tiling tensors.
207 Returns:
208 Rank-local forward and backward runtime resources.
209 """
210 forward_graph, forward_config, backward_graph, backward_config = _build_runtime_artifacts(spec)
211 fwd_diagnostics, bwd_diagnostics = _prepare_diagnostics(
212 spec, forward_graph, forward_config, backward_graph, backward_config, device,
213 )
214 gmm_options = {
215 "hidden_size": spec.hidden_size,
216 "intermediate_size": spec.intermediate_size,
217 "num_groups": spec.local_experts,
218 "num_cube_cores": spec.num_cube_cores,
219 }
220 return MegaMoePlan(
221 spec=spec,
222 reuse_backward_dispatch=can_reuse_backward_dispatch(
223 backward_config, spec.local_experts, spec.num_cube_cores,
224 ),
225 forward_graph=forward_graph,
226 fwd_diagnostics=fwd_diagnostics,
227 up_proj_tiling=_tensor_from_bytes(
228 get_up_proj_tiling_bytes(forward_graph.get_op("up_proj").split_value, **gmm_options),
229 device,
230 ),
231 swiglu_tiling=_tensor_from_bytes(
232 _resize_swiglu_tiling(
233 get_swiglu_tiling_bytes(
234 forward_graph.get_op("swiglu").split_value,
235 intermediate_size=spec.intermediate_size,
236 ),
237 spec.num_cube_cores,
238 ),
239 device,
240 ),
241 down_proj_tiling=_tensor_from_bytes(
242 get_down_proj_tiling_bytes(forward_graph.get_op("down_proj").split_value, **gmm_options),
243 device,
244 ),
245 backward_graph=backward_graph,
246 bwd_diagnostics=bwd_diagnostics,
247 act_grad_tiling=_tensor_from_bytes(
248 get_act_grad_tiling_bytes(backward_graph.get_op("act_grad").split_value, **gmm_options),
249 device,
250 ),
251 gate_grad_tiling=_tensor_from_bytes(
252 get_gate_grad_tiling_bytes(backward_graph.get_op("gate_grad").split_value, **gmm_options),
253 device,
254 ),
255 w1_grad_tiling=_tensor_from_bytes(
256 get_w1_grad_tiling_bytes(backward_graph.get_op("w1_grad").split_value,
257 expert_weight_layout=spec.expert_weight_layout, **gmm_options),
258 device,
259 ),
260 w2_grad_tiling=_tensor_from_bytes(
261 get_w2_grad_tiling_bytes(backward_graph.get_op("w2_grad").split_value,
262 expert_weight_layout=spec.expert_weight_layout, **gmm_options),
263 device,
264 ),
265 swiglu_grad_tiling=_tensor_from_bytes(
266 _resize_swiglu_tiling(
267 get_swiglu_grad_tiling_bytes(
268 backward_graph.get_op("swiglu_grad").split_value,
269 intermediate_size=spec.intermediate_size,
270 ),
271 spec.num_cube_cores,
272 ),
273 device,
274 ),
275 )