« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2025-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"""Layer's blocks submodule"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
18from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
19from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
21if TYPE_CHECKING:
22 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
23 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
25mb = EvalUtils.mb
28class EvalAttn:
29 """Attention formulas class"""
31 @staticmethod
32 def num_params_mla(ccfg: CostModelConfig, _) -> float:
33 """Parameters count for Multi-Head Latent Attention"""
34 # W_up_q = ccfg.dc_q * ccfg.dh * ccfg.a
35 # W_up_k = ccfg.dc_kv * ccfg.dh * ccfg.n_kv
36 # W_up_v = ccfg.dc_kv * ccfg.dh * ccfg.n_kv
37 # W_down_q = ccfg.dc_q * ccfg.h
38 # W_down_kv = ccfg.dc_kv * ccfg.h
39 # W_q_rope = ccfg.a * ccfg.dhr * ccfg.dc_q
40 # W_k_rope = ccfg.dhr * ccfg.h
41 # Wo = ccfg.h * ccfg.a * ccfg.dh
43 c_kv_fact = ccfg.dc_kv * (ccfg.n_kv * ccfg.dh + ccfg.h)
44 c_q_fact = ccfg.dc_q * (ccfg.a * ccfg.dh + ccfg.h + ccfg.a * ccfg.dhr)
45 rest_fact = (ccfg.h * ccfg.a * ccfg.dh) + (ccfg.h * ccfg.dhr)
46 res = (
47 0.5 * ccfg.n_attMM * c_kv_fact
48 + 0.25 * ccfg.n_attMM * c_q_fact
49 + 0.25 * ccfg.n_attMM * rest_fact
50 )
51 return res
53 @staticmethod
54 def num_params_attn(ccfg: CostModelConfig, ctx: Context) -> float:
55 """Parameters count for Multi-Head/Grouped-Q./Multi-Q. Attention"""
56 if ccfg.dc_kv == 0:
57 # Q,O and K,V have distinct shapes
58 return 0.5 * ccfg.n_attMM * (
59 ccfg.h * ccfg.h + ccfg.h
60 ) + 0.5 * ccfg.n_attMM * (ccfg.h * ccfg.n_kv * ccfg.dh + ccfg.h)
61 return EvalAttn.num_params_mla(ccfg, ctx)
63 @staticmethod
64 def attn_qkv_activations(ccfg: CostModelConfig, ctx: Context) -> float:
65 """QKV linear Activations"""
66 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
67 att_qkv_size = 0
68 if ccfg.dc_kv == 0:
69 n_op = ccfg.n_attMM + ccfg.n_attParamCast
70 att_qkv_size = (
71 ccfg.s
72 * ccfg.b
73 * ccfg.bytes_compute
74 * (
75 0.25 * n_op * ccfg.h
76 + 0.5 * n_op * ccfg.dh * ccfg.n_kv
77 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.attBMM)
78 * ccfg.n_attBMM
79 * ccfg.dh
80 )
81 )
82 else:
83 q_size = (
84 0.25
85 * (ccfg.n_attMM + ccfg.n_attParamCast)
86 * (ccfg.dc_q + 2 * ccfg.a * (ccfg.dh + ccfg.dhr))
87 )
88 k_size = (
89 0.25
90 * (ccfg.n_attMM + ccfg.n_attParamCast)
91 * (ccfg.dhr + ccfg.n_kv * (2 * ccfg.dh + ccfg.dhr))
92 )
93 v_size = (
94 0.25
95 * (ccfg.n_attMM + ccfg.n_attParamCast)
96 * (ccfg.n_kv * ccfg.dh + ccfg.dc_kv)
97 )
98 att_qkv_size = (
99 ccfg.s
100 * ccfg.b
101 * ccfg.bytes_compute
102 * (
103 q_size
104 + k_size
105 + v_size
106 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.attBMM)
107 * ccfg.n_attBMM
108 * ccfg.dh
109 )
110 )
111 micro_factor = ctx.micro_factor
112 return micro_factor * att_qkv_size / (ccfg.t * ccfg.cp)
114 @staticmethod
115 def attn_score_activations(ccfg: CostModelConfig, ctx: Context) -> float:
116 """Score/Softmax Activations"""
117 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
118 att_score = (
119 ccfg.s_fa
120 * ccfg.b
121 * ccfg.a
122 * ccfg.s
123 * (
124 ccfg.n_softmax
125 * (
126 ccfg.rec_op.softmax * ccfg.bytes_softmax
127 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.dropout)
128 * ccfg.bytes_dropout
129 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.headCast)
130 * ccfg.bytes_compute
131 )
132 )
133 )
134 micro_factor = ctx.micro_factor
135 # cp_sq_div=1: Ring CP shards Q along seq and all-gathers KV, so the S²
136 # score tensor has only one S dim divided → B·H·S²/cp (slope -1).
137 # Only an unimplemented blockwise-ring path (shard both Q and KV along
138 # seq → (s/cp)² scores) would need cp_sq_div=cp; do not re-add the
139 # conditional for the current Ring/Ulysses algorithms.
140 cp_sq_div = 1
141 return micro_factor * att_score / (ccfg.t * ccfg.cp * cp_sq_div)
143 @staticmethod
144 def attn_proj_activations(ccfg: CostModelConfig, ctx: Context) -> float:
145 """Output projection Activations"""
146 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
147 att_proj = (
148 ccfg.s
149 * ccfg.b
150 * ccfg.h
151 * ccfg.bytes_compute
152 * (
153 0.25 * (ccfg.n_attMM + ccfg.n_attParamCast)
154 + EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.dropout)
155 * ccfg.n_dropout
156 * ccfg.bytes_dropout
157 )
158 )
159 micro_factor = ctx.micro_factor
160 return micro_factor * att_proj / max(ccfg.sp, ccfg.cp)
163class EvalFFn:
164 """Feed-forward formulas class"""
166 @staticmethod
167 def num_params_ffn(ccfg: CostModelConfig, _) -> float:
168 """Parameters count"""
169 experts_param_size = (
170 (ccfg.n_exp + ccfg.n_shared_exp)
171 * max(ccfg.n_ffMM, ccfg.n_ffBMM)
172 * (ccfg.hff * ccfg.h + ccfg.hff)
173 )
174 return experts_param_size
176 @staticmethod
177 def num_params_routed_expert(ccfg: CostModelConfig, _) -> float:
178 """Routed expert parameters count (with ETP correction)"""
179 hff_sliced = ccfg.hff_exp / max(ccfg.etp, 1)
180 return ccfg.n_exp * max(ccfg.n_ffMM, ccfg.n_ffBMM) * (hff_sliced * ccfg.h + hff_sliced)
182 @staticmethod
183 def num_params_shared_expert(ccfg: CostModelConfig, _) -> float:
184 """Shared expert parameters count"""
185 return ccfg.n_shared_exp * max(ccfg.n_ffMM, ccfg.n_ffBMM) * (ccfg.hff * ccfg.h + ccfg.hff)
187 @staticmethod
188 def ffn_activations(ccfg: CostModelConfig, ctx: Context) -> float:
189 """ "Activations count"""
190 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
191 tok_size = ccfg.s * ccfg.b
192 n_mm = max(ccfg.n_ffMM, ccfg.n_ffBMM)
193 if n_mm % 2 == 0:
194 matmul = 0.5 * ccfg.h + 0.5 * ccfg.hff
195 else:
196 matmul = 1 / 3 * ccfg.h + 2 / 3 * ccfg.hff
197 matmul *= ccfg.bytes_compute * n_mm
198 activ_fun = ccfg.bytes_compute * ccfg.hff
199 activ_fun *= EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.ffAct)
200 pcast = ccfg.bytes_compute * ccfg.hff * ccfg.n_ffParamCast
201 activ_size = matmul + pcast + activ_fun
202 micro_factor = ctx.micro_factor
203 return micro_factor * tok_size * activ_size / (ccfg.t * ccfg.cp)
205 @staticmethod
206 def ffn_router_and_concat_activations(
207 ccfg: CostModelConfig, ctx: Context
208 ) -> float:
209 """MoE router and output activations"""
210 # Router activations (logits, probs, mask)
211 r = ccfg.s * ccfg.b * ccfg.bytes_compute
212 r *= 2 * ccfg.n_exp + ccfg.n_chosen_exp
213 # Concat all exp output
214 c = ccfg.s * ccfg.b * ccfg.bytes_compute * ccfg.h
215 micro_factor = ctx.micro_factor
216 return micro_factor * (r + c) / (ccfg.t * ccfg.cp)
218 @staticmethod
219 def shared_exp_activations(ccfg: CostModelConfig, ctx: Context) -> float:
220 """Shared expert activations"""
221 return ccfg.n_shared_exp * EvalFFn.ffn_activations(ccfg, ctx)
223 @staticmethod
224 def routed_exp_activations(ccfg: CostModelConfig, ctx: Context) -> float:
225 """MoE topK activations"""
226 tok_size = ccfg.s * ccfg.b
227 activ_size = EvalFFn.ffn_activations(ccfg, ctx) / tok_size
228 avg_num_toks = tok_size * ccfg.n_chosen_exp / ccfg.n_exp
229 if not ccfg.gmm: # Capacity mode
230 expert_capacity = avg_num_toks * ccfg.cap_fact * ccfg.n_exp
231 routed_activ = activ_size * expert_capacity
232 else: # Dropless mode
233 load = avg_num_toks * ccfg.n_exp * ctx.dropless_tok_factor
234 routed_activ = load * activ_size
235 return routed_activ
237 @staticmethod
238 def ffn_moe_activations(ccfg: CostModelConfig, ctx: Context) -> float:
239 """Sum of Activations"""
240 return (
241 EvalFFn.routed_exp_activations(ccfg, ctx)
242 + EvalFFn.shared_exp_activations(ccfg, ctx)
243 + EvalFFn.ffn_router_and_concat_activations(ccfg, ctx)
244 )
247class EvalNorm:
248 """Normalization formulas class"""
250 @staticmethod
251 def num_params_norm(ccfg: CostModelConfig, _) -> float:
252 """Parameters count"""
253 return ccfg.n_normOp * 2 * ccfg.h
255 @staticmethod
256 def norm_activations(ccfg: CostModelConfig, ctx: Context) -> float:
257 """Activations"""
258 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
259 norm = (
260 ccfg.s
261 * ccfg.b
262 * ccfg.bytes_norm
263 * ccfg.h
264 * ccfg.n_normOp
265 * EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.normOp)
266 )
267 micro_factor = ctx.micro_factor
268 return micro_factor * norm / (ccfg.t * ccfg.cp)