« prev ^ index » next coverage.py v7.13.1, created at 2026-10-26 05:43 +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, Any, Dict, Mapping, Optional
18from hyper_parallel.auto_parallel._op_records import load_op_records
19from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
20from hyper_parallel.auto_parallel.sapp_nd.nd.common.cp_types import CPAlgo, _resolve_cp_algo
21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.utils import EvalUtils
23if TYPE_CHECKING:
24 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
25 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
27mb = EvalUtils.mb
30def _bias(ccfg: CostModelConfig, stated: str, width: float) -> float:
31 """A projection's bias: *width* where the model states one, none where it states none.
33 Where the model states neither, as a parser that reads no bias does,
34 each projection is counted a bias of the hidden width, the formulas'
35 convention.
36 """
37 has_bias = getattr(ccfg, stated, None)
38 if has_bias is None:
39 return ccfg.h
40 return width if has_bias else 0
43def _mlp_biases(ccfg: CostModelConfig, width: float) -> float:
44 """The biases of one feed-forward of *width*, or one expert's, as :func:`_bias` counts a projection's.
46 Stated, each projection into the width has one of it and the last,
47 back to the hidden width, one of that; unstated, each projection one of
48 the width.
49 """
50 has_bias = getattr(ccfg, "mlp_bias", None)
51 if has_bias is None:
52 return ccfg.n_ffMM * width
53 return (ccfg.n_ffMM - 1) * width + ccfg.h if has_bias else 0
56class EvalAttn:
57 """Attention formulas class"""
59 @staticmethod
60 def num_params_mla(ccfg: CostModelConfig, _) -> float:
61 """Parameters count for Multi-Head Latent Attention.
63 The queries: a down-projection to their latent, its norm and an
64 up-projection to every head's non-rotary and rotary part, or one
65 projection where the model has no query latent. The keys and
66 values: one down-projection to their shared latent beside the
67 rotary key, its norm, and an up-projection to every head's
68 non-rotary key and value. The output projection from the values.
69 Each head's non-rotary key is ``qk_nope_head_dim`` wide, ``dh``,
70 its value head's width, unless the model states it.
71 """
72 d_qk = getattr(ccfg, "qk_nope_head_dim", None) or ccfg.dh
73 heads_q = ccfg.a * (d_qk + ccfg.dhr)
74 if ccfg.dc_q:
75 query = ccfg.h * ccfg.dc_q + ccfg.dc_q + ccfg.dc_q * heads_q
76 else:
77 query = ccfg.h * heads_q
78 key_value = (ccfg.h * (ccfg.dc_kv + ccfg.dhr) + ccfg.dc_kv
79 + ccfg.dc_kv * ccfg.n_kv * (d_qk + ccfg.dh))
80 output = ccfg.a * ccfg.dh * ccfg.h
81 return 0.25 * ccfg.n_attMM * (query + key_value + output)
83 @staticmethod
84 def num_params_attn(ccfg: CostModelConfig, ctx: Context) -> float:
85 """Parameters count for Multi-Head/Grouped-Q./Multi-Q. Attention"""
86 if ccfg.dc_kv == 0:
87 # Q,O and K,V have distinct shapes. Q and O are h x (a*dh),
88 # which equals h x h only when head_dim = h/a; a fused output
89 # gate (Qwen3.5) doubles the Q projection.
90 d_h = ccfg.dh or (ccfg.h / ccfg.a if ccfg.a else 0)
91 d_q = ccfg.a * d_h
92 q_fact = 2 if ccfg.attn_output_gate else 1
93 return 0.25 * ccfg.n_attMM * (
94 q_fact * ccfg.h * d_q + _bias(ccfg, "qkv_bias", q_fact * d_q)
95 ) + 0.25 * ccfg.n_attMM * (
96 ccfg.h * d_q + _bias(ccfg, "o_bias", ccfg.h)
97 ) + 0.5 * ccfg.n_attMM * (
98 ccfg.h * ccfg.n_kv * d_h + _bias(ccfg, "qkv_bias", ccfg.n_kv * d_h)
99 ) + ccfg.attn_extra_p
100 return EvalAttn.num_params_mla(ccfg, ctx)
102 @staticmethod
103 def kv_shards(ccfg: CostModelConfig) -> float:
104 """Over how many CP ranks a layer's keys and values are split.
106 Colossal-AI and hybrid CP all-gather the keys and values, and the
107 attention kernel keeps them whole for the backward: none. Ulysses
108 CP gives each rank its heads' share, and a linear-attention layer
109 passes its state rather than gathering them: all of CP.
110 """
111 if _resolve_cp_algo(ccfg) == CPAlgo.ULYSSES_CP or getattr(ccfg, "n_linrec", 0):
112 return ccfg.cp
113 return 1
115 @staticmethod
116 def gathered_kv_bytes(ccfg: CostModelConfig) -> float:
117 """The bytes per token of a rank's share of the sequence it keeps of the other ranks' keys and values.
119 Where CP gathers the keys and values (:meth:`kv_shards`), the
120 attention keeps the whole sequence's, the rank's own share and
121 ``cp - 1`` others, each split over TP.
122 """
123 if ccfg.cp <= 1 or EvalAttn.kv_shards(ccfg) > 1:
124 return 0
125 return (ccfg.cp - 1) * 2 * ccfg.n_kv * ccfg.dh * ccfg.bytes_compute / max(1, ccfg.t)
127 @staticmethod
128 def attn_qkv_activations(ccfg: CostModelConfig, ctx: Context) -> float:
129 """QKV linear Activations: the op records' ``qkv`` slot"""
130 return EvalRecords.slot_bytes(ccfg, ctx, "qkv")
132 @staticmethod
133 def attn_score_activations(ccfg: CostModelConfig, ctx: Context) -> float:
134 """Score/Softmax Activations: the op records' ``score`` slot.
136 Ring CP shards the queries along the sequence and gathers the keys and
137 values, so the score tensor has one sequence dimension divided, B H S^2
138 / cp; only a blockwise ring, which shards both, would divide it again.
139 """
140 return EvalRecords.slot_bytes(ccfg, ctx, "score")
142 @staticmethod
143 def attn_proj_activations(ccfg: CostModelConfig, ctx: Context) -> float:
144 """Output projection Activations: the op records' ``proj`` slot"""
145 return EvalRecords.slot_bytes(ccfg, ctx, "proj")
148class EvalFFn:
149 """Feed-forward formulas class"""
151 @staticmethod
152 def num_params_ffn(ccfg: CostModelConfig, _) -> float:
153 """Parameters count"""
154 return (ccfg.n_exp + ccfg.n_shared_exp) * (
155 ccfg.n_ffMM * ccfg.hff * ccfg.h + _mlp_biases(ccfg, ccfg.hff))
157 @staticmethod
158 def num_params_routed_expert(ccfg: CostModelConfig, _) -> float:
159 """Routed expert parameters count (with ETP correction)"""
160 hff_sliced = ccfg.hff_exp / max(ccfg.etp, 1)
161 return ccfg.n_exp * (ccfg.n_ffMM * hff_sliced * ccfg.h + _mlp_biases(ccfg, hff_sliced))
163 @staticmethod
164 def num_params_router(ccfg: CostModelConfig, _: Any) -> float:
165 """The router's parameters: a weight per routed expert, the hidden width wide, where the layer routes."""
166 return ccfg.h * ccfg.n_exp if ccfg.n_exp > 1 else 0
168 @staticmethod
169 def num_params_shared_expert(ccfg: CostModelConfig, _) -> float:
170 """Shared expert parameters count, and the weight that gates its output where the model has one.
172 Every producer states the shared experts as a count of experts of
173 the routed ones' width, ``hff_exp``: one wide shared expert as that
174 many.
175 """
176 gate = ccfg.h if ccfg.n_shared_exp and getattr(ccfg, "shared_expert_gate", None) else 0
177 width = ccfg.hff_exp
178 return ccfg.n_shared_exp * (ccfg.n_ffMM * width * ccfg.h + _mlp_biases(ccfg, width)) + gate
180 @staticmethod
181 def ffn_activations(ccfg: CostModelConfig, ctx: Context, width: Optional[float] = None) -> float:
182 """Activations of a feed-forward *width* wide, the model's dense width unless given: the ``ffn`` slot"""
183 return EvalRecords.slot_bytes(ccfg, ctx, "ffn", overrides=None if width is None else {"hff": width})
185 @staticmethod
186 def ffn_router_and_concat_activations(
187 ccfg: CostModelConfig, ctx: Context
188 ) -> float:
189 """MoE router and output activations: the ``router`` slot"""
190 return EvalRecords.slot_bytes(ccfg, ctx, "router")
192 @staticmethod
193 def shared_exp_activations(ccfg: CostModelConfig, ctx: Context) -> float:
194 """Shared expert activations, each shared expert of the routed ones' width: the ``shared`` slot"""
195 return EvalRecords.slot_bytes(ccfg, ctx, "shared")
197 @staticmethod
198 def routed_exp_activations(ccfg: CostModelConfig, ctx: Context) -> float:
199 """MoE topK activations, each expert at its width: the ``routed`` slot"""
200 return EvalRecords.slot_bytes(ccfg, ctx, "routed")
202 @staticmethod
203 def ffn_moe_activations(ccfg: CostModelConfig, ctx: Context) -> float:
204 """Sum of Activations"""
205 return (
206 EvalFFn.routed_exp_activations(ccfg, ctx)
207 + EvalFFn.shared_exp_activations(ccfg, ctx)
208 + EvalFFn.ffn_router_and_concat_activations(ccfg, ctx)
209 )
212class EvalNorm:
213 """Normalization formulas class"""
215 @staticmethod
216 def head_dim(ccfg: CostModelConfig) -> float:
217 """The width a QK-norm normalizes over: one attention head's"""
218 return ccfg.dh or (ccfg.h / ccfg.a if ccfg.a else 0)
220 @staticmethod
221 def num_params_norm(ccfg: CostModelConfig, _) -> float:
222 """Parameters count: the layer's norms, and a QK-norm's query and key weights.
224 A model that states its norms holds a weight in each, and a bias
225 beside it in a LayerNorm; one that does not is counted two vectors
226 per norm op.
227 """
228 layer_norms = getattr(ccfg, "layer_norms", None)
229 if layer_norms is None:
230 vectors = ccfg.n_normOp * 2
231 else:
232 vectors = layer_norms * (2 if getattr(ccfg, "norm_bias", None) else 1)
233 return vectors * ccfg.h + getattr(ccfg, "n_qknorm", 0) * 2 * EvalNorm.head_dim(ccfg)
235 @staticmethod
236 def norm_activations(ccfg: CostModelConfig, ctx: Context) -> float:
237 """Activations: the norms' inputs, a QK-norm's every head's queries and keys: the ``norm`` slot"""
238 return EvalRecords.slot_bytes(ccfg, ctx, "norm")
241class EvalRecords:
242 """The op records (shared decision S1), priced on a layer config.
244 The activation formulas of :class:`EvalAttn`, :class:`EvalFFn` and
245 :class:`EvalNorm` are the records' slots, so a caller prices a layer
246 under any setting of the recompute switches, or one op alone, without
247 setting anything on its config.
248 """
250 @staticmethod
251 def slot_bytes(ccfg: CostModelConfig, ctx: Context, slot: str,
252 switches: Optional[Mapping[str, Any]] = None,
253 overrides: Optional[Mapping[str, Any]] = None, only: Optional[str] = None) -> float:
254 """The bytes a layer of *ccfg* keeps in *slot*, one rank's share of a micro-batch.
256 Args:
257 ccfg: The layer's config.
258 ctx: The evaluation: its micro factor, and its node, a selective
259 layer's (``SEL_REC_LAYER``) dropping what its switches drop.
260 slot: The slot.
261 switches: A selective layer's switches, 1 to keep an op's
262 activations and 0 to drop them; where omitted, those the
263 context carries, else the config's own.
264 overrides: Values that stand for the config's fields.
265 only: An op to price alone, its parts with it, or a part, whatever
266 the switches say.
267 """
268 records = load_op_records()
269 rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
270 stated = EvalUtils.switches(ccfg, ctx) if switches is None else switches
272 def keep(op: str) -> Any:
273 """Whether this layer keeps *op*'s activation: 1 keeps it, 0 recomputes it."""
274 switch = records.record(op).switch
275 if switch is None:
276 return 1
277 return EvalUtils.rec_coeff(rec_layer, EvalUtils.state(stated, switch))
279 return records.evaluate(slot, EvalRecords.values(ccfg, ctx, overrides), keep, only)
281 @staticmethod
282 def values(ccfg: CostModelConfig, ctx: Context, overrides: Optional[Mapping[str, Any]] = None) -> Any:
283 """The value of each name a record reads: an override, the evaluation's, or the config's field."""
284 records = load_op_records()
285 overrides = overrides or {}
287 def value(name: str) -> Any:
288 """The value of *name*: an override, the evaluation's, a record's default or the config's field."""
289 if name in overrides:
290 return overrides[name]
291 if name in ("micro_factor", "dropless_tok_factor"):
292 return getattr(ctx, name)
293 if name == "kv_shards":
294 return EvalAttn.kv_shards(ccfg)
295 if name in records.defaults:
296 return getattr(ccfg, name, records.defaults[name])
297 return getattr(ccfg, name)
299 return value
301 @staticmethod
302 def layer_slots(ccfg: CostModelConfig, ctx: Context) -> Dict[str, str]:
303 """The slots a layer of *ccfg* holds, each with its part."""
304 value = EvalRecords.values(ccfg, ctx)
305 return {name: slot.part for name, slot in load_op_records().slots.items() if slot.holds(value)}
307 @staticmethod
308 def op_bytes(ccfg: CostModelConfig, ctx: Context, op: str) -> Dict[str, float]:
309 """What a layer of *ccfg* keeps for *op* where nothing drops it, by slot of the layer's."""
310 held = EvalRecords.layer_slots(ccfg, ctx)
311 return {slot: EvalRecords.slot_bytes(ccfg, ctx, slot, only=op)
312 for slot in load_op_records().keeps(op) if slot in held}
314 @staticmethod
315 def layer_bytes(ccfg: CostModelConfig, ctx: Context,
316 switches: Optional[Mapping[str, Any]] = None) -> Dict[str, float]:
317 """What a layer of *ccfg* keeps, by part, under *switches* or its config's own."""
318 parts: Dict[str, float] = {}
319 for slot, part in EvalRecords.layer_slots(ccfg, ctx).items():
320 parts[part] = parts.get(part, 0.0) + EvalRecords.slot_bytes(ccfg, ctx, slot, switches)
321 return parts