Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / forward / graph.py: 100%
44 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"""MoE-FFN forward compute graph."""
17from hyper_multicore._core_spec import DEFAULT_CORE_SPEC
18from hyper_multicore.scheduler.graph import (
19 ComputeGraph, OperatorNode, TensorSpec, SplitSpec, OpType,
20)
21from hyper_multicore.tasks.alltoall import AllToAllFillConfig, AllToAllType
22from hyper_multicore.tasks.gmm import GmmFillConfig
23from hyper_multicore.tasks.swiglu import SwiGLUFillConfig
25# Tiling buffer positions in the C++ runtime descriptor array.
26_TILING_POS_UP_PROJ = 17
27_TILING_POS_SWIGLU = 18
28_TILING_POS_DN_PROJ = 19
31def _build_fwd_tensor_specs(
32 tsv, hidden_size, intermediate_size, dtype_size, expert_weight_layout
33):
34 """Create all TensorSpec objects for the forward graph."""
35 sre = tsv.single_rank_expert_num
36 weight_is_eoi = expert_weight_layout == "EOI"
37 target = TensorSpec(
38 "target", [tsv.per_rank_seq, hidden_size],
39 param_position=0, dtype_size=dtype_size, is_dynamic=True)
40 target_offset = TensorSpec(
41 "target_offset", [tsv.all_expert_num],
42 param_position=1, dtype_size=8)
43 src = TensorSpec(
44 "src", [tsv.per_rank_seq, hidden_size],
45 param_position=2, dtype_size=dtype_size)
46 src_offset = TensorSpec(
47 "src_offset", [tsv.all_expert_num],
48 param_position=3, dtype_size=8)
49 size_d = TensorSpec(
50 "size_d", [tsv.all_expert_num],
51 param_position=4, dtype_size=4)
52 up_proj_weight = TensorSpec(
53 "up_proj_weight",
54 ([sre, intermediate_size * 2, hidden_size]
55 if weight_is_eoi else [sre, hidden_size, intermediate_size * 2]),
56 param_position=5, dtype_size=dtype_size, tensor_type=1,
57 transpose=weight_is_eoi)
58 up_proj_glist = TensorSpec(
59 "up_proj_glist", [sre],
60 param_position=6, dtype_size=8)
61 up_proj_y = TensorSpec(
62 "up_proj_y", [tsv.per_rank_seq, intermediate_size * 2],
63 param_position=7, dtype_size=dtype_size, is_dynamic=True)
64 swiglu_out = TensorSpec(
65 "swiglu_out", [tsv.per_rank_seq, intermediate_size],
66 param_position=8, dtype_size=dtype_size, tensor_type=1, is_dynamic=True)
67 down_proj_weight = TensorSpec(
68 "down_proj_weight",
69 ([sre, hidden_size, intermediate_size]
70 if weight_is_eoi else [sre, intermediate_size, hidden_size]),
71 param_position=9, dtype_size=dtype_size, tensor_type=1,
72 transpose=weight_is_eoi)
73 down_proj_glist = TensorSpec(
74 "down_proj_glist", [sre],
75 param_position=10, dtype_size=8)
76 down_proj_y = TensorSpec(
77 "down_proj_y", [tsv.per_rank_seq, hidden_size],
78 param_position=11, dtype_size=dtype_size, tensor_type=1, is_dynamic=True)
79 combine_out = TensorSpec(
80 "combine_out", [tsv.per_rank_seq, hidden_size],
81 param_position=12, dtype_size=dtype_size)
82 target_offset_c = TensorSpec(
83 "target_offset_c", [tsv.all_expert_num],
84 param_position=13, dtype_size=8)
85 src_offset_c = TensorSpec(
86 "src_offset_c", [tsv.all_expert_num],
87 param_position=14, dtype_size=8)
88 size_c = TensorSpec(
89 "size_c", [tsv.all_expert_num],
90 param_position=15, dtype_size=4)
91 return (target, target_offset, src, src_offset, size_d,
92 up_proj_weight, up_proj_glist, up_proj_y, swiglu_out,
93 down_proj_weight, down_proj_glist, down_proj_y,
94 combine_out, target_offset_c, src_offset_c, size_c)
97def _build_fwd_ops(tsv, specs, *, dispatch_sv, up_proj_sv, swiglu_sv,
98 down_proj_sv, combine_sv, num_cube_cores, swiglu_limit=None):
99 """Create all OperatorNode objects for the forward graph."""
100 (target, target_offset, src, src_offset, size_d,
101 up_proj_weight, up_proj_glist, up_proj_y, swiglu_out,
102 down_proj_weight, down_proj_glist, down_proj_y,
103 combine_out, target_offset_c, src_offset_c, size_c) = specs
104 dispatch = OperatorNode(
105 name="dispatch", op_type=OpType.ALLTOALL,
106 diagnostic_name="Dispatch",
107 inputs=[target_offset, src, src_offset, size_d],
108 outputs=[target],
109 param_positions=[1, 2, 3, 4, 0],
110 split_value=dispatch_sv,
111 split_spec=SplitSpec(
112 split_inputs=None, split_output_dims=[0],
113 task_num_fn=lambda tsv: tsv.all_expert_num * (tsv.per_expert_seq_to_other // dispatch_sv),
114 ),
115 tiling_position=-1,
116 fill_config=AllToAllFillConfig(moe_type=AllToAllType.DISPATCH, advance="vector",
117 event_group=tsv.all_expert_num),
118 )
119 up_proj = OperatorNode(
120 name="up_proj", op_type=OpType.GMM,
121 diagnostic_name="GMM1",
122 inputs=[target, up_proj_weight, up_proj_glist], outputs=[up_proj_y],
123 param_positions=[0, 5, 6, 7], split_value=up_proj_sv,
124 split_spec=SplitSpec(
125 split_inputs=[(0, 0)], split_output_dims=[0],
126 task_num_fn=lambda tsv, _ncc=num_cube_cores: _ncc * tsv.single_rank_expert_num,
127 ),
128 tiling_position=_TILING_POS_UP_PROJ,
129 fill_config=GmmFillConfig(offset_inputs={0}, rank_in_event=True, global_trigger=False,
130 out_offset=True, advance="cube", num_cube_cores=num_cube_cores),
131 )
132 swiglu = OperatorNode(
133 name="swiglu", op_type=OpType.SWIGLU,
134 diagnostic_name="SwiGLU",
135 inputs=[up_proj_y], outputs=[swiglu_out],
136 param_positions=[7, 8], split_value=swiglu_sv,
137 split_spec=SplitSpec(
138 split_inputs=[(0, 0)], split_output_dims=[0],
139 task_num_fn=lambda tsv: (tsv.per_expert_seq // swiglu_sv) * tsv.single_rank_expert_num,
140 ),
141 tiling_position=_TILING_POS_SWIGLU,
142 fill_config=SwiGLUFillConfig(clamp_limit=swiglu_limit),
143 )
144 down_proj = OperatorNode(
145 name="down_proj", op_type=OpType.GMM,
146 diagnostic_name="GMM2",
147 inputs=[swiglu_out, down_proj_weight, down_proj_glist], outputs=[down_proj_y],
148 param_positions=[8, 9, 10, 11], split_value=down_proj_sv,
149 split_spec=SplitSpec(
150 split_inputs=[(0, 0)], split_output_dims=[0],
151 task_num_fn=lambda tsv, _ncc=num_cube_cores: _ncc * tsv.single_rank_expert_num,
152 ),
153 tiling_position=_TILING_POS_DN_PROJ,
154 fill_config=GmmFillConfig(offset_inputs={0}, rank_in_event=False, global_trigger=False,
155 out_offset=True, advance="cube", num_cube_cores=num_cube_cores),
156 )
157 combine = OperatorNode(
158 name="combine", op_type=OpType.ALLTOALL,
159 diagnostic_name="Combine",
160 inputs=[target_offset_c, down_proj_y, src_offset_c, size_c], outputs=[combine_out],
161 param_positions=[13, 11, 14, 15, 12], split_value=combine_sv,
162 split_spec=SplitSpec(
163 split_inputs=[(1, 0)], split_output_dims=[0],
164 task_num_fn=lambda tsv: tsv.all_expert_num * (tsv.per_expert_seq_to_other // combine_sv),
165 ),
166 tiling_position=-1,
167 fill_config=AllToAllFillConfig(moe_type=AllToAllType.COMBINE, advance="vector", event_group=1),
168 )
169 return dispatch, up_proj, swiglu, down_proj, combine
172def build_forward_graph(tsv, *,
173 dispatch_sv: int = 128,
174 up_proj_sv: int = 4096,
175 swiglu_sv: int = 128,
176 down_proj_sv: int = 4096,
177 combine_sv: int = 128,
178 hidden_size: int = 7168,
179 intermediate_size: int = 2048,
180 dtype_size: int = 2,
181 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores,
182 swiglu_limit: float | None = None,
183 expert_weight_layout: str = "EIO") -> ComputeGraph:
184 """Build the MoE-FFN forward DAG: dispatch -> up_proj -> swiglu -> down_proj -> combine.
186 Operator execution order and param_positions (C++ memory slots):
187 dispatch: [target_offset=1, src=2, src_offset=3, size=4 | target=0]
188 up_proj: [x=0, weight=5, glist=6 | y=7]
189 swiglu: [x=7 | out=8]
190 down_proj: [x=8, weight=9, glist=10 | y=11]
191 combine: [target_offset=13, src=11, src_offset=14, size=15 | target=12]
193 Args:
194 tsv: TaskSplitValue carrying TP/EP/seq partition metadata.
195 dispatch_sv: tile size for dispatch AllToAll (vector cores).
196 up_proj_sv: tile size for up-projection GMM (cube cores).
197 swiglu_sv: tile size for SwiGLU (vector cores).
198 down_proj_sv: tile size for down-projection GMM (cube cores).
199 combine_sv: tile size for combine AllToAll (vector cores).
200 hidden_size: model hidden dimension.
201 intermediate_size: FFN intermediate dimension after SwiGLU halving.
202 dtype_size: bytes per activation element (2=bf16, 4=fp32).
203 num_cube_cores: number of AIC cube cores on the target device.
204 swiglu_limit: optional positive clamp limit for SwiGLU; ``None`` keeps
205 the original unclamped operator path.
206 expert_weight_layout: Logical expert-weight axis order, either ``"EIO"`` or ``"EOI"``.
208 Returns:
209 A fully-connected ComputeGraph ready for propagate_splits().
210 """
211 if expert_weight_layout not in ("EIO", "EOI"):
212 raise ValueError(f"expert_weight_layout must be 'EIO' or 'EOI', got {expert_weight_layout!r}.")
213 specs = _build_fwd_tensor_specs(
214 tsv, hidden_size, intermediate_size, dtype_size, expert_weight_layout
215 )
216 dispatch, up_proj, swiglu, down_proj, combine = _build_fwd_ops(
217 tsv, specs,
218 dispatch_sv=dispatch_sv, up_proj_sv=up_proj_sv, swiglu_sv=swiglu_sv,
219 down_proj_sv=down_proj_sv, combine_sv=combine_sv, num_cube_cores=num_cube_cores,
220 swiglu_limit=swiglu_limit,
221 )
222 graph = ComputeGraph()
223 (graph.add_op(dispatch).add_op(up_proj).add_op(swiglu).add_op(down_proj).add_op(combine)
224 .add_edge(dispatch, up_proj)
225 .add_edge(up_proj, swiglu)
226 .add_edge(swiglu, down_proj)
227 .add_edge(down_proj, combine))
228 return graph