Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_mhc / graph.py: 95%
176 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"""Pure AIC/AIV token-tile task graph for shifted Single-Pass HyperMegaMhc."""
17from __future__ import annotations
19from dataclasses import dataclass
20from typing import Callable
22from hyper_multicore.scheduler.config import (
23 ATOMIC_ADD_VALUE_LEN,
24 EVENT_INVALID_ID,
25 MIN_EVENT_CAPACITY,
26 RuntimeConfigC,
27 TaskAiCoreType,
28 TaskDescC,
29 TaskSplitValue,
30 TaskType,
31)
32from hyper_multicore.scheduler.graph import (
33 ComputeGraph,
34 OperatorNode,
35 OpType,
36 SplitSpec,
37 TensorSpec,
38)
39from hyper_multicore.tasks.task_base import FillConfig
42DEFAULT_TOKEN_TILE = 32
43_STAGE1_BS_FACTOR = 32
44POST_DONE_EVENT = 0
45NORM_EVENT_BASE = 1
48@dataclass(frozen=True)
49class MegaMhcEventLayout:
50 """Contiguous event ranges for one shape-bound pipeline."""
52 post_done: int
53 norm_base: int
54 input_mix_base: int
55 projection_base: int
56 final: int
57 event_count: int
60def build_event_layout(tile_count: int) -> MegaMhcEventLayout:
61 """Return event IDs for a six-stage, per-token-tile pipeline.
63 Args:
64 tile_count: Number of logical token tiles.
65 """
66 if tile_count <= 0:
67 raise ValueError(f"tile_count must be positive, got {tile_count}.")
68 input_mix_base = NORM_EVENT_BASE + tile_count
69 projection_base = input_mix_base + tile_count
70 final = projection_base + tile_count
71 event_count = final + ATOMIC_ADD_VALUE_LEN
72 if event_count > MIN_EVENT_CAPACITY:
73 raise ValueError(
74 f"HyperMegaMhc needs {event_count} events for {tile_count} token tiles, "
75 f"but RuntimeConfig supports only {MIN_EVENT_CAPACITY}."
76 )
77 return MegaMhcEventLayout(
78 post_done=POST_DONE_EVENT,
79 norm_base=NORM_EVENT_BASE,
80 input_mix_base=input_mix_base,
81 projection_base=projection_base,
82 final=final,
83 event_count=event_count,
84 )
87def resolve_token_tile(token_count: int, requested_tile: int, num_cube_cores: int) -> int:
88 """Choose an aligned tile that fits event and native Stage1 workspace limits.
90 Args:
91 token_count: Total flattened token count.
92 requested_tile: Preferred token count per task.
93 num_cube_cores: Number of Cube cores.
94 """
95 if token_count <= 0 or requested_tile <= 0 or num_cube_cores <= 0:
96 raise ValueError(
97 "token_count, requested_tile, and num_cube_cores must all be positive, "
98 f"got {(token_count, requested_tile, num_cube_cores)}."
99 )
100 max_tiles = (MIN_EVENT_CAPACITY - NORM_EVENT_BASE - ATOMIC_ADD_VALUE_LEN) // 3
101 minimum_tile = (token_count + max_tiles - 1) // max_tiles
102 minimum_tile = (minimum_tile + _STAGE1_BS_FACTOR - 1) // _STAGE1_BS_FACTOR * _STAGE1_BS_FACTOR
103 token_tile = max(requested_tile, minimum_tile)
104 token_tile = (token_tile + _STAGE1_BS_FACTOR - 1) // _STAGE1_BS_FACTOR * _STAGE1_BS_FACTOR
105 stage1_workspace_rows = 4 * num_cube_cores * _STAGE1_BS_FACTOR
106 if token_tile > stage1_workspace_rows:
107 raise ValueError(
108 f"token tile {token_tile} exceeds the native Stage1 X-cast workspace capacity "
109 f"of {stage1_workspace_rows} rows."
110 )
111 return token_tile
114def _append_task(
115 cfg: RuntimeConfigC,
116 tsv: TaskSplitValue,
117 *,
118 task_type: TaskType,
119 core_type: TaskAiCoreType,
120 dependent_event: int,
121 trigger_event: int,
122 task_index: int,
123 task_num: int,
124 split_value: int,
125) -> None:
126 """Append one task to the matching static AIC or AIV schedule."""
127 task = TaskDescC()
128 task.task_type = task_type
129 task.task_aicore_type = core_type
130 task.dependent_event = dependent_event
131 task.trigger_event = trigger_event
132 task.task_index = task_index
133 task.task_split_num = task_num
134 task.task_split_value = split_value
135 task_id = tsv.pre_task_num
136 cfg.all_tasks[task_id] = task
137 if core_type == TaskAiCoreType.TASK_AICORE_CUBE:
138 cfg.cube_task_indices[tsv.pre_cube_task_num] = task_id
139 cfg.task_index_num[0] += 1
140 tsv.pre_cube_task_num += 1
141 elif core_type == TaskAiCoreType.TASK_AICORE_VECTOR:
142 cfg.vector_task_indices[tsv.pre_vector_task_num] = task_id
143 cfg.task_index_num[1] += 1
144 tsv.pre_vector_task_num += 1
145 else:
146 raise ValueError(f"HyperMegaMhc task must be pure AIC or AIV, got {core_type!r}.")
147 tsv.pre_task_num += 1
150_AIV_PIPELINE_TYPES = (
151 TaskType.TASK_MHC_NORM_CAST,
152 TaskType.TASK_MHC_INPUT_MIX,
153 TaskType.TASK_RMS_NORM,
154 TaskType.TASK_MHC_MAPPING,
155)
158def _reorder_vector_task_waves(
159 cfg: RuntimeConfigC,
160 tsv: TaskSplitValue,
161 tile_count: int,
162 wave_size: int,
163) -> None:
164 """Fill Projection gaps with NormCast/InputMix/RMSNorm waves before Mapping."""
165 if wave_size <= 0:
166 raise ValueError(f"wave_size must be positive, got {wave_size}.")
167 post_task_ids = []
168 pipeline_task_ids = {task_type: [] for task_type in _AIV_PIPELINE_TYPES}
169 for schedule_index in range(tsv.pre_vector_task_num):
170 task_id = cfg.vector_task_indices[schedule_index]
171 task_type = TaskType(cfg.all_tasks[task_id].task_type)
172 if task_type == TaskType.TASK_MHC_POST:
173 post_task_ids.append(task_id)
174 elif task_type in pipeline_task_ids:
175 pipeline_task_ids[task_type].append(task_id)
176 else:
177 raise ValueError(f"Unexpected AIV task type in HyperMegaMhc schedule: {task_type!r}.")
179 if len(post_task_ids) != tile_count:
180 raise ValueError(f"Expected {tile_count} MhcPost tasks, got {len(post_task_ids)}.")
181 for task_type, task_ids in pipeline_task_ids.items():
182 if len(task_ids) != tile_count:
183 raise ValueError(f"Expected {tile_count} {task_type.name} tasks, got {len(task_ids)}.")
185 reordered_task_ids = list(post_task_ids)
186 for wave_start in range(0, tile_count, wave_size):
187 wave_end = min(wave_start + wave_size, tile_count)
188 for task_type in _AIV_PIPELINE_TYPES[:3]:
189 reordered_task_ids.extend(pipeline_task_ids[task_type][wave_start:wave_end])
190 reordered_task_ids.extend(pipeline_task_ids[TaskType.TASK_MHC_MAPPING])
191 for schedule_index, task_id in enumerate(reordered_task_ids):
192 cfg.vector_task_indices[schedule_index] = task_id
195@dataclass(frozen=True)
196class _PostFillConfig(FillConfig):
197 """Emit token-local residual update tasks on AIV."""
199 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
200 """Append the MhcPost task slice to the runtime configuration.
202 Args:
203 cfg: Runtime configuration receiving tasks.
204 op: Scheduled operator node.
205 tsv: Mutable task split counters.
206 """
207 layout = build_event_layout(op.task_num)
208 cfg.all_event_num_triggers[layout.post_done] = op.task_num
209 for task_index in range(op.task_num):
210 _append_task(
211 cfg,
212 tsv,
213 task_type=TaskType.TASK_MHC_POST,
214 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
215 dependent_event=EVENT_INVALID_ID,
216 trigger_event=layout.post_done,
217 task_index=task_index,
218 task_num=op.task_num,
219 split_value=op.split_value,
220 )
223@dataclass(frozen=True)
224class _NormCastFillConfig(FillConfig):
225 """Emit AIV normalization/cast tasks with safe X-cast ring-slot reuse."""
227 ring_slots: int
229 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
230 """Append NormCast tasks and ring-buffer reuse dependencies.
232 Args:
233 cfg: Runtime configuration receiving tasks.
234 op: Scheduled operator node.
235 tsv: Mutable task split counters.
236 """
237 layout = build_event_layout(op.task_num)
238 for task_index in range(op.task_num):
239 dependent_event = layout.post_done
240 if task_index >= self.ring_slots:
241 dependent_event = layout.projection_base + task_index - self.ring_slots
242 trigger_event = layout.norm_base + task_index
243 cfg.all_event_num_triggers[trigger_event] = 1
244 _append_task(
245 cfg,
246 tsv,
247 task_type=TaskType.TASK_MHC_NORM_CAST,
248 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
249 dependent_event=dependent_event,
250 trigger_event=trigger_event,
251 task_index=task_index,
252 task_num=op.task_num,
253 split_value=op.split_value,
254 )
257@dataclass(frozen=True)
258class _ProjectionFillConfig(FillConfig):
259 """Emit projection tasks on AIC."""
261 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
262 """Append token-tiled AIC projection tasks.
264 Args:
265 cfg: Runtime configuration receiving tasks.
266 op: Scheduled operator node.
267 tsv: Mutable task split counters.
268 """
269 layout = build_event_layout(op.task_num)
270 for task_index in range(op.task_num):
271 trigger_event = layout.projection_base + task_index
272 cfg.all_event_num_triggers[trigger_event] = 1
273 _append_task(
274 cfg,
275 tsv,
276 task_type=TaskType.TASK_MHC_PROJECTION,
277 core_type=TaskAiCoreType.TASK_AICORE_CUBE,
278 dependent_event=layout.norm_base + task_index,
279 trigger_event=trigger_event,
280 task_index=task_index,
281 task_num=op.task_num,
282 split_value=op.split_value,
283 )
286@dataclass(frozen=True)
287class _MappingFillConfig(FillConfig):
288 """Emit per-tile mapping/Sinkhorn tasks after the corresponding Projection."""
290 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
291 """Append mapping and Sinkhorn tasks after their projection tiles.
293 Args:
294 cfg: Runtime configuration receiving tasks.
295 op: Scheduled operator node.
296 tsv: Mutable task split counters.
297 """
298 layout = build_event_layout(op.task_num)
299 cfg.all_event_num_triggers[layout.final] = 2 * op.task_num
300 for task_index in range(op.task_num):
301 _append_task(
302 cfg,
303 tsv,
304 task_type=TaskType.TASK_MHC_MAPPING,
305 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
306 dependent_event=layout.projection_base + task_index,
307 trigger_event=layout.final,
308 task_index=task_index,
309 task_num=op.task_num,
310 split_value=op.split_value,
311 )
314@dataclass(frozen=True)
315class _InputMixFillConfig(FillConfig):
316 """Emit input-mix tasks after NormCast to fill the Projection AIV gap."""
318 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
319 """Append shifted input-mix tasks after the matching NormCast tiles.
321 Args:
322 cfg: Runtime configuration receiving tasks.
323 op: Scheduled operator node.
324 tsv: Mutable task split counters.
325 """
326 layout = build_event_layout(op.task_num)
327 for task_index in range(op.task_num):
328 trigger_event = layout.input_mix_base + task_index
329 cfg.all_event_num_triggers[trigger_event] = 1
330 _append_task(
331 cfg,
332 tsv,
333 task_type=TaskType.TASK_MHC_INPUT_MIX,
334 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
335 dependent_event=layout.norm_base + task_index,
336 trigger_event=trigger_event,
337 task_index=task_index,
338 task_num=op.task_num,
339 split_value=op.split_value,
340 )
343@dataclass(frozen=True)
344class _RmsNormFillConfig(FillConfig):
345 """Emit shifted input RMSNorm tasks on AIV."""
347 wave_size: int
349 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
350 """Append shifted RMSNorm tasks and arrange the AIV task waves.
352 Args:
353 cfg: Runtime configuration receiving tasks.
354 op: Scheduled operator node.
355 tsv: Mutable task split counters.
356 """
357 layout = build_event_layout(op.task_num)
358 for task_index in range(op.task_num):
359 _append_task(
360 cfg,
361 tsv,
362 task_type=TaskType.TASK_RMS_NORM,
363 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
364 dependent_event=layout.input_mix_base + task_index,
365 trigger_event=layout.final,
366 task_index=task_index,
367 task_num=op.task_num,
368 split_value=op.split_value,
369 )
370 _reorder_vector_task_waves(cfg, tsv, op.task_num, self.wave_size)
373def _assemble_forward_graph(
374 topology: TaskSplitValue,
375 post: OperatorNode,
376 norm_cast: OperatorNode,
377 projection: OperatorNode,
378 input_mix: OperatorNode,
379 rms_norm: OperatorNode,
380 mapping: OperatorNode,
381) -> ComputeGraph:
382 """Connect the forward stages and resolve their task splits."""
383 graph = ComputeGraph()
384 for operator in (post, norm_cast, projection, input_mix, rms_norm, mapping):
385 graph.add_op(operator)
386 graph.add_edge(post, norm_cast).add_edge(norm_cast, projection).add_edge(projection, mapping)
387 graph.add_edge(norm_cast, input_mix).add_edge(input_mix, rms_norm)
388 graph.propagate_splits(topology)
389 return graph
392def _make_post_and_norm_cast(
393 residual: TensorSpec,
394 updated_residual: TensorSpec,
395 normalized: TensorSpec,
396 token_tile: int,
397 task_num_fn: Callable[[TaskSplitValue], int],
398 ring_slots: int,
399) -> tuple[OperatorNode, OperatorNode]:
400 """Create the two opening AIV stages."""
401 post = OperatorNode(
402 name="mhc_post",
403 op_type=OpType.MHC_POST,
404 inputs=[residual],
405 outputs=[updated_residual],
406 param_positions=[1, 12],
407 split_value=token_tile,
408 split_spec=SplitSpec(None, task_num_fn, [0]),
409 tiling_position=18,
410 fill_config=_PostFillConfig(),
411 diagnostic_name="MhcPost",
412 )
413 norm_cast = OperatorNode(
414 name="mhc_norm_cast",
415 op_type=OpType.MHC_NORM_CAST,
416 inputs=[updated_residual],
417 outputs=[normalized],
418 param_positions=[12, 17],
419 split_value=token_tile,
420 split_spec=SplitSpec(None, task_num_fn, [0]),
421 tiling_position=18,
422 fill_config=_NormCastFillConfig(ring_slots),
423 diagnostic_name="MhcNormCast",
424 )
425 return post, norm_cast
428def build_mega_mhc_graph(
429 token_count: int,
430 hidden_size: int,
431 num_cube_cores: int,
432 token_tile: int = DEFAULT_TOKEN_TILE,
433) -> tuple[ComputeGraph, TaskSplitValue]:
434 """Build the overlapped Projection and shifted InputMix/RMSNorm token pipeline.
436 Args:
437 token_count: Total flattened token count.
438 hidden_size: Hidden dimension.
439 num_cube_cores: Number of Cube cores.
440 token_tile: Preferred token count per task.
441 """
442 if token_count <= 0 or hidden_size <= 0 or num_cube_cores <= 0:
443 raise ValueError(
444 "token_count, hidden_size, and num_cube_cores must all be positive, "
445 f"got {(token_count, hidden_size, num_cube_cores)}."
446 )
447 token_tile = resolve_token_tile(token_count, token_tile, num_cube_cores)
448 tile_count = (token_count + token_tile - 1) // token_tile
449 build_event_layout(tile_count)
450 stage1_workspace_rows = 4 * num_cube_cores * _STAGE1_BS_FACTOR
451 ring_slots = min(tile_count, stage1_workspace_rows // token_tile)
453 topology = TaskSplitValue(tp=1, ep=1, seq_size=token_count, all_expert_num=1, top_k=1)
454 residual = TensorSpec("residual", [token_count, 4, hidden_size], 1)
455 updated_residual = TensorSpec("updated_residual", [token_count, 4, hidden_size], 12)
456 normalized = TensorSpec("normalized", [token_tile, 4, hidden_size], 17, dtype_size=4)
457 projection = TensorSpec("projection", [token_count, 24], 17, dtype_size=4)
458 mappings = TensorSpec("mappings", [token_count, 24], 13, dtype_size=4)
459 mixed_input = TensorSpec("mixed_input", [token_count, hidden_size], 17)
460 block_input = TensorSpec("block_input", [token_count, hidden_size], 16)
462 def task_num_fn(tsv: TaskSplitValue) -> int:
463 """Return the token tile count shared by all forward stages.
465 Args:
466 tsv: Unused topology value required by the scheduler callback.
467 """
468 del tsv
469 return tile_count
471 post, norm_cast = _make_post_and_norm_cast(
472 residual,
473 updated_residual,
474 normalized,
475 token_tile,
476 task_num_fn,
477 ring_slots,
478 )
479 projection_op = OperatorNode(
480 name="mhc_projection",
481 op_type=OpType.MHC_PROJECTION,
482 inputs=[normalized],
483 outputs=[projection],
484 param_positions=[17],
485 split_value=token_tile,
486 split_spec=SplitSpec(None, task_num_fn, [0]),
487 tiling_position=18,
488 fill_config=_ProjectionFillConfig(),
489 diagnostic_name="MhcProjection",
490 )
491 input_mix = OperatorNode(
492 name="mhc_input_mix",
493 op_type=OpType.MHC_INPUT_MIX,
494 inputs=[updated_residual],
495 outputs=[mixed_input],
496 param_positions=[12, 2, 17],
497 split_value=token_tile,
498 split_spec=SplitSpec(None, task_num_fn, [0]),
499 tiling_position=18,
500 fill_config=_InputMixFillConfig(),
501 diagnostic_name="MhcInputMix",
502 )
503 rms_norm = OperatorNode(
504 name="mhc_rms_norm",
505 op_type=OpType.RMS_NORM,
506 inputs=[mixed_input],
507 outputs=[block_input],
508 param_positions=[16],
509 split_value=token_tile,
510 split_spec=SplitSpec(None, task_num_fn, [0]),
511 tiling_position=18,
512 fill_config=_RmsNormFillConfig(2 * num_cube_cores),
513 diagnostic_name="ShiftedRmsNorm",
514 )
515 mapping = OperatorNode(
516 name="mhc_mapping",
517 op_type=OpType.MHC_MAPPING,
518 inputs=[projection, updated_residual],
519 outputs=[mappings],
520 param_positions=[17, 12, 13],
521 split_value=token_tile,
522 split_spec=SplitSpec(None, task_num_fn, [0]),
523 tiling_position=18,
524 fill_config=_MappingFillConfig(),
525 diagnostic_name="MhcMapping/Sinkhorn",
526 )
527 graph = _assemble_forward_graph(
528 topology, post, norm_cast, projection_op, input_mix, rms_norm, mapping
529 )
530 return graph, topology