Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_mhc_grad / graph.py: 98%
160 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"""Token-tiled fused graph for shifted Single-Pass HyperMegaMhcGrad."""
17from __future__ import annotations
19from dataclasses import dataclass
21from hyper_multicore.scheduler.config import (
22 ATOMIC_ADD_VALUE_LEN,
23 EVENT_INVALID_ID,
24 FAST_DEPENDENCY_POLL_INTERVAL_US,
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_GRAD_TOKEN_TILE = 32
43NATIVE_GRAD_VECTOR_TILE = 32
44NATIVE_GRAD_CUBE_TILE = 64
45INIT_DONE_EVENT = 0
48@dataclass(frozen=True)
49class MegaMhcGradTiling:
50 """Outer scheduler tiling decoupled from the native 32/64-row micro-tiles."""
52 token_count: int
53 token_tile: int
54 cube_tile: int
55 tile_count: int
56 macro_count: int
58 def macro_for_tile(self, task_index: int) -> int:
59 """Return the AIC macro-tile that consumes one AIV tile.
61 Args:
62 task_index: AIV tile index.
64 Returns:
65 Corresponding AIC macro-tile index.
66 """
67 return task_index * self.token_tile // self.cube_tile
69 def vector_tiles_in_macro(self, macro_index: int) -> int:
70 """Return the number of AIV tiles contributing to one AIC macro-tile.
72 Args:
73 macro_index: AIC macro-tile index.
75 Returns:
76 Number of contributing AIV tiles.
77 """
78 return sum(
79 self.macro_for_tile(task_index) == macro_index
80 for task_index in range(self.tile_count)
81 )
84@dataclass(frozen=True)
85class MegaMhcGradEventLayout:
86 """Contiguous event ranges for the fused backward pipeline."""
88 init_done: int
89 rms_ready_base: int
90 macro_ready_base: int
91 phi_rms_base: int
92 final: int
93 event_count: int
96def build_event_layout(tile_count: int, macro_count: int) -> MegaMhcGradEventLayout:
97 """Return event IDs for configurable AIV tiles and AIC macro-tiles.
99 Args:
100 tile_count: Number of AIV token tiles.
101 macro_count: Number of AIC macro-tiles.
103 Returns:
104 Contiguous event ranges for the backward pipeline.
105 """
106 if tile_count <= 0 or macro_count <= 0:
107 raise ValueError(
108 f"tile_count and macro_count must be positive, got {(tile_count, macro_count)}."
109 )
110 rms_ready_base = INIT_DONE_EVENT + 1
111 macro_ready_base = rms_ready_base + tile_count
112 phi_rms_base = macro_ready_base + macro_count
113 final = phi_rms_base + macro_count
114 event_count = final + ATOMIC_ADD_VALUE_LEN
115 if event_count > MIN_EVENT_CAPACITY:
116 raise ValueError(
117 f"HyperMegaMhcGrad needs {event_count} events for {tile_count} AIV tiles "
118 f"and {macro_count} AIC tiles, but RuntimeConfig supports only {MIN_EVENT_CAPACITY}."
119 )
120 return MegaMhcGradEventLayout(
121 init_done=INIT_DONE_EVENT,
122 rms_ready_base=rms_ready_base,
123 macro_ready_base=macro_ready_base,
124 phi_rms_base=phi_rms_base,
125 final=final,
126 event_count=event_count,
127 )
130def resolve_grad_token_tile(token_count: int, requested_tile: int) -> int:
131 """Choose a 32-row-aligned outer tile that fits the event-counter table.
133 Args:
134 token_count: Total number of flattened tokens.
135 requested_tile: Requested AIV token tile.
137 Returns:
138 Resolved token tile that fits the runtime event capacity.
139 """
140 if token_count <= 0 or requested_tile <= 0:
141 raise ValueError(
142 "token_count and requested_tile must both be positive, "
143 f"got {(token_count, requested_tile)}."
144 )
145 token_tile = (
146 (requested_tile + NATIVE_GRAD_VECTOR_TILE - 1)
147 // NATIVE_GRAD_VECTOR_TILE
148 * NATIVE_GRAD_VECTOR_TILE
149 )
150 while True:
151 tiling = make_grad_tiling(token_count, token_tile)
152 if 1 + tiling.tile_count + 2 * tiling.macro_count + ATOMIC_ADD_VALUE_LEN <= MIN_EVENT_CAPACITY:
153 return token_tile
154 token_tile += NATIVE_GRAD_VECTOR_TILE
157def make_grad_tiling(token_count: int, token_tile: int) -> MegaMhcGradTiling:
158 """Construct outer AIV/AIC tiling for one resolved token tile.
160 Args:
161 token_count: Total number of flattened tokens.
162 token_tile: Resolved AIV token tile.
164 Returns:
165 Derived AIV and AIC tiling metadata.
166 """
167 cube_tile = max(token_tile, NATIVE_GRAD_CUBE_TILE)
168 tile_count = (token_count + token_tile - 1) // token_tile
169 macro_count = (token_count + cube_tile - 1) // cube_tile
170 return MegaMhcGradTiling(token_count, token_tile, cube_tile, tile_count, macro_count)
173def _append_task(
174 cfg: RuntimeConfigC,
175 tsv: TaskSplitValue,
176 *,
177 task_type: TaskType,
178 core_type: TaskAiCoreType,
179 dependent_event: int,
180 trigger_event: int,
181 task_index: int,
182 task_num: int,
183 split_value: int,
184 dependency_poll_interval_us: int = 0,
185) -> None:
186 """Append one pure AIC or AIV task to RuntimeConfig."""
187 task = TaskDescC()
188 task.task_type = task_type
189 task.task_aicore_type = core_type
190 task.dependent_event = dependent_event
191 task.trigger_event = trigger_event
192 task.task_index = task_index
193 task.task_split_num = task_num
194 task.task_split_value = split_value
195 # Reuse the current runtime ABI's third task-specific value for the
196 # dependency polling interval required by the 4-task backward pipeline.
197 task.extra_value_2 = dependency_poll_interval_us
198 task_id = tsv.pre_task_num
199 cfg.all_tasks[task_id] = task
200 if core_type == TaskAiCoreType.TASK_AICORE_CUBE:
201 cfg.cube_task_indices[tsv.pre_cube_task_num] = task_id
202 cfg.task_index_num[0] += 1
203 tsv.pre_cube_task_num += 1
204 elif core_type == TaskAiCoreType.TASK_AICORE_VECTOR:
205 cfg.vector_task_indices[tsv.pre_vector_task_num] = task_id
206 cfg.task_index_num[1] += 1
207 tsv.pre_vector_task_num += 1
208 else:
209 raise ValueError(f"HyperMegaMhcGrad task must be pure AIC or AIV, got {core_type!r}.")
210 tsv.pre_task_num += 1
213@dataclass(frozen=True)
214class _InitFillConfig(FillConfig):
215 """Zero atomic-reduction outputs before any producer is released."""
217 tiling: MegaMhcGradTiling
219 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
220 """Append the reduction-output initialization task.
222 Args:
223 cfg: Runtime configuration receiving the task.
224 op: Initialization graph node.
225 tsv: Mutable task counters.
226 """
227 layout = build_event_layout(self.tiling.tile_count, self.tiling.macro_count)
228 cfg.all_event_num_triggers[layout.init_done] = op.task_num
229 for task_index in range(op.task_num):
230 _append_task(
231 cfg,
232 tsv,
233 task_type=TaskType.TASK_MHC_GRAD_PREV_A,
234 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
235 dependent_event=EVENT_INVALID_ID,
236 trigger_event=layout.init_done,
237 task_index=task_index,
238 task_num=op.task_num,
239 split_value=0,
240 )
243@dataclass(frozen=True)
244class _RmsNormGradFillConfig(FillConfig):
245 """Emit one standalone RMSNormGrad task for each AIV token tile."""
247 tiling: MegaMhcGradTiling
249 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
250 """Append one RMSNorm gradient task per AIV token tile.
252 Args:
253 cfg: Runtime configuration receiving the tasks.
254 op: RMSNorm gradient graph node.
255 tsv: Mutable task counters.
256 """
257 layout = build_event_layout(self.tiling.tile_count, self.tiling.macro_count)
258 for task_index in range(op.task_num):
259 ready_event = layout.rms_ready_base + task_index
260 cfg.all_event_num_triggers[ready_event] = 1
261 _append_task(
262 cfg,
263 tsv,
264 task_type=TaskType.TASK_RMS_NORM_GRAD,
265 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
266 dependent_event=layout.init_done,
267 trigger_event=ready_event,
268 task_index=task_index,
269 task_num=op.task_num,
270 split_value=self.tiling.token_tile,
271 dependency_poll_interval_us=FAST_DEPENDENCY_POLL_INTERVAL_US,
272 )
275@dataclass(frozen=True)
276class _PrepareFillConfig(FillConfig):
277 """Emit fused shifted previous-A and mapping tasks."""
279 tiling: MegaMhcGradTiling
281 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
282 """Append fused previous-A and mapping tasks.
284 Args:
285 cfg: Runtime configuration receiving the tasks.
286 op: Fused preparation graph node.
287 tsv: Mutable task counters.
288 """
289 layout = build_event_layout(self.tiling.tile_count, self.tiling.macro_count)
290 for macro_index in range(self.tiling.macro_count):
291 producer_count = self.tiling.vector_tiles_in_macro(macro_index)
292 cfg.all_event_num_triggers[layout.macro_ready_base + macro_index] = producer_count
293 for task_index in range(op.task_num):
294 _append_task(
295 cfg,
296 tsv,
297 task_type=TaskType.TASK_MHC_GRAD_PREV_A_AND_MAPPING,
298 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
299 dependent_event=layout.rms_ready_base + task_index,
300 trigger_event=layout.macro_ready_base + self.tiling.macro_for_tile(task_index),
301 task_index=task_index,
302 task_num=op.task_num,
303 split_value=self.tiling.token_tile,
304 dependency_poll_interval_us=FAST_DEPENDENCY_POLL_INTERVAL_US,
305 )
308@dataclass(frozen=True)
309class _PhiRmsFillConfig(FillConfig):
310 """Emit AIC ProcessMatmul1/2 tasks over configurable macro-tiles."""
312 tiling: MegaMhcGradTiling
314 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
315 """Append projection-gradient tasks for all AIC macro-tiles.
317 Args:
318 cfg: Runtime configuration receiving the tasks.
319 op: Projection-gradient graph node.
320 tsv: Mutable task counters.
321 """
322 layout = build_event_layout(self.tiling.tile_count, self.tiling.macro_count)
323 for macro_index in range(op.task_num):
324 cfg.all_event_num_triggers[layout.phi_rms_base + macro_index] = 1
325 _append_task(
326 cfg,
327 tsv,
328 task_type=TaskType.TASK_MHC_GRAD_PHI_RMS,
329 core_type=TaskAiCoreType.TASK_AICORE_CUBE,
330 dependent_event=layout.macro_ready_base + macro_index,
331 trigger_event=layout.phi_rms_base + macro_index,
332 task_index=macro_index,
333 task_num=op.task_num,
334 split_value=self.tiling.cube_tile,
335 )
338@dataclass(frozen=True)
339class _PreviousXFillConfig(FillConfig):
340 """Emit fused ComputeGradX1 and MhcPost backward tasks."""
342 tiling: MegaMhcGradTiling
344 def fill(self, cfg: RuntimeConfigC, op: OperatorNode, tsv: TaskSplitValue) -> None:
345 """Append fused previous-X and MhcPost gradient tasks.
347 Args:
348 cfg: Runtime configuration receiving the tasks.
349 op: Fused final-stage graph node.
350 tsv: Mutable task counters.
351 """
352 layout = build_event_layout(self.tiling.tile_count, self.tiling.macro_count)
353 cfg.all_event_num_triggers[layout.final] = op.task_num
354 for task_index in range(op.task_num):
355 _append_task(
356 cfg,
357 tsv,
358 task_type=TaskType.TASK_MHC_GRAD_PREV_X_AND_POST,
359 core_type=TaskAiCoreType.TASK_AICORE_VECTOR,
360 dependent_event=(layout.phi_rms_base + self.tiling.macro_for_tile(task_index)),
361 trigger_event=layout.final,
362 task_index=task_index,
363 task_num=op.task_num,
364 split_value=self.tiling.token_tile,
365 )
368def order_vector_tasks_by_stage(config: RuntimeConfigC) -> None:
369 """Queue all tiles of one AIV stage before advancing to the next stage.
371 Args:
372 config: Mutable runtime configuration whose AIV queue is reordered.
373 """
374 vector_count = int(config.task_index_num[1])
375 task_ids = [int(config.vector_task_indices[index]) for index in range(vector_count)]
376 init_tasks = []
377 stage_types = (
378 TaskType.TASK_RMS_NORM_GRAD,
379 TaskType.TASK_MHC_GRAD_PREV_A_AND_MAPPING,
380 TaskType.TASK_MHC_GRAD_PREV_X_AND_POST,
381 )
382 by_type: dict[TaskType, dict[int, int]] = {task_type: {} for task_type in stage_types}
383 for task_id in task_ids:
384 task = config.all_tasks[task_id]
385 task_type = TaskType(task.task_type)
386 if task_type == TaskType.TASK_MHC_GRAD_PREV_A and task.task_split_value == 0:
387 init_tasks.append(task_id)
388 else:
389 by_type[task_type][int(task.task_index)] = task_id
391 reordered = list(init_tasks)
392 for task_type in stage_types:
393 reordered.extend(by_type[task_type][index] for index in sorted(by_type[task_type]))
394 if len(reordered) != vector_count or set(reordered) != set(task_ids):
395 raise RuntimeError("HyperMegaMhcGrad stage ordering lost or duplicated an AIV task.")
396 for index, task_id in enumerate(reordered):
397 config.vector_task_indices[index] = task_id
400def build_mega_mhc_grad_graph(
401 token_count: int,
402 hidden_size: int,
403 token_tile: int = DEFAULT_GRAD_TOKEN_TILE,
404 *,
405 num_vector_cores: int,
406) -> tuple[ComputeGraph, TaskSplitValue]:
407 """Build the RMS -> PrevAAndMapping -> PhiRms -> PrevXAndPost token pipeline.
409 Args:
410 token_count: Total number of flattened tokens.
411 hidden_size: Hidden dimension of every token.
412 token_tile: Requested AIV token tile.
413 num_vector_cores: Number of Vector cores.
415 Returns:
416 Compute graph and its single-rank topology description.
417 """
418 if token_count <= 0 or hidden_size <= 0 or num_vector_cores <= 0:
419 raise ValueError(
420 "token_count, hidden_size, and num_vector_cores must be positive, "
421 f"got {(token_count, hidden_size, num_vector_cores)}."
422 )
423 token_tile = resolve_grad_token_tile(token_count, token_tile)
424 tiling = make_grad_tiling(token_count, token_tile)
425 build_event_layout(tiling.tile_count, tiling.macro_count)
426 topology = TaskSplitValue(tp=1, ep=1, seq_size=token_count, all_expert_num=1, top_k=1)
427 data = TensorSpec("grad_data", [token_count, hidden_size], 0)
429 def make_op(
430 name: str,
431 task_num: int,
432 split_value: int,
433 fill_config: FillConfig,
434 diagnostic_name: str,
435 ) -> OperatorNode:
436 """Construct one backward pipeline operator node.
438 Args:
439 name: Internal graph-node name.
440 task_num: Number of logical tasks.
441 split_value: Tokens processed by each task.
442 fill_config: Runtime task materializer.
443 diagnostic_name: Profiler-visible stage name.
445 Returns:
446 Configured backward operator node.
447 """
448 return OperatorNode(
449 name=name,
450 op_type=OpType.MHC_PIPELINE,
451 inputs=[data],
452 outputs=[],
453 param_positions=[],
454 split_value=split_value,
455 split_spec=SplitSpec(None, lambda _tsv, count=task_num: count, []),
456 tiling_position=34,
457 fill_config=fill_config,
458 diagnostic_name=diagnostic_name,
459 )
461 init = make_op(
462 "grad_output_init",
463 num_vector_cores,
464 0,
465 _InitFillConfig(tiling),
466 "GradOutputInit",
467 )
468 rms_norm_grad = make_op(
469 "rms_norm_grad",
470 tiling.tile_count,
471 token_tile,
472 _RmsNormGradFillConfig(tiling),
473 "RmsNormGrad",
474 )
475 prev_a_and_mapping = make_op(
476 "mhc_grad_prev_a_and_mapping",
477 tiling.tile_count,
478 token_tile,
479 _PrepareFillConfig(tiling),
480 "MhcGradPrevAAndMapping",
481 )
482 phi_rms = make_op(
483 "mhc_grad_phi_rms",
484 tiling.macro_count,
485 tiling.cube_tile,
486 _PhiRmsFillConfig(tiling),
487 "MhcGradPhiRms",
488 )
489 prev_x_and_post = make_op(
490 "mhc_grad_prev_x_and_post",
491 tiling.tile_count,
492 token_tile,
493 _PreviousXFillConfig(tiling),
494 "MhcGradPrevXAndPost",
495 )
497 graph = ComputeGraph()
498 for operator in (init, rms_norm_grad, prev_a_and_mapping, phi_rms, prev_x_and_post):
499 graph.add_op(operator)
500 graph.add_edge(init, rms_norm_grad).add_edge(rms_norm_grad, prev_a_and_mapping)
501 graph.add_edge(prev_a_and_mapping, phi_rms)
502 graph.add_edge(phi_rms, prev_x_and_post)
503 graph.propagate_splits(topology)
504 return graph, topology