« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2025 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"""Expert compute FLOPs estimation"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
19if TYPE_CHECKING:
20 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
21 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
24class EvalExpertCompute:
25 """Expert compute (FLOPs) estimation formulas.
27 All methods return FLOPs (not bytes), so they must NOT go through
28 __wrap_mem_counter — they are not memory quantities.
29 """
31 @staticmethod
32 def router_compute_cost(ccfg: CostModelConfig, ctx: Context) -> float: # pylint: disable=unused-argument
33 """Router (gate) FLOPs — replicated, not scaled by EP.
35 Router computes a score vector of length n_exp for each token,
36 then selects top-K. The cost is dominated by the linear projection:
37 FLOPs = 2 * s * b * h * n_exp (per microbatch, per layer)
39 Router is replicated across EP ranks because topK selection needs
40 the full expert score vector.
41 """
42 return 2 * ccfg.s * ccfg.b * ccfg.h * ccfg.n_exp / ccfg.cp
44 @staticmethod
45 def expert_compute_cost_balanced(
46 ccfg: CostModelConfig, ctx: Context # pylint: disable=unused-argument
47 ) -> float:
48 """Routed expert compute FLOPs for balanced token distribution.
50 Per-rank FLOPs for routed experts only.
51 Each rank holds n_exp/ep experts and processes 1/ep of the tokens
52 (balanced assumption):
53 FLOPs = 2 * n_ffMM * s * b * h * hff_exp * n_chosen_exp
54 / (ep * cp * t_exp)
56 t_exp = etp if etp > 1 else tp (alternative, not multiplicative).
57 Factor 2 accounts for multiply-add (MAC = 2 FLOPs).
58 n_ffMM is the number of feedforward linear layers per expert
59 (SwiGLU: gate+up+down = 3; standard MLP: 2), set in arch_hooks.
60 """
61 ep = max(ccfg.ep, 1)
62 t_exp = ccfg.etp if ccfg.etp > 1 else ccfg.t
63 n_ff = max(getattr(ccfg, "n_ffMM", 1), 1)
64 return (
65 2 * n_ff * ccfg.s * ccfg.b * ccfg.h * ccfg.hff_exp
66 * ccfg.n_chosen_exp / (ep * max(t_exp, 1) * ccfg.cp)
67 )
69 @staticmethod
70 def expert_compute_cost_imbalanced(
71 ccfg: CostModelConfig, ctx: Context
72 ) -> float:
73 """Routed expert compute FLOPs for imbalanced token distribution.
75 Uses max(tokens_per_rank) to bound per-rank compute (bucket effect).
76 Falls back to balanced when tokens_per_expert is None or
77 n_exp not divisible by ep.
79 t_exp = etp if etp > 1 else tp (alternative, not multiplicative).
80 n_ffMM is the number of feedforward linear layers per expert
81 (SwiGLU: 3, standard MLP: 2), set in arch_hooks.
83 tokens_per_expert: global per-expert token count per microbatch
84 (all EP ranks combined, before all-to-all dispatch; None = balanced).
85 """
86 ep = max(ccfg.ep, 1)
87 tokens = ccfg.tokens_per_expert
88 if not tokens:
89 return EvalExpertCompute.expert_compute_cost_balanced(ccfg, ctx)
90 if ccfg.n_exp % ep != 0:
91 return EvalExpertCompute.expert_compute_cost_balanced(ccfg, ctx)
92 if len(tokens) < ccfg.n_exp:
93 return EvalExpertCompute.expert_compute_cost_balanced(ccfg, ctx)
94 t_exp = ccfg.etp if ccfg.etp > 1 else ccfg.t
95 n_ff = max(getattr(ccfg, "n_ffMM", 1), 1)
96 experts_per_rank = ccfg.n_exp // ep
97 rank_tokens = []
98 for r in range(ep):
99 rank_sum = sum(
100 tokens[r * experts_per_rank + i] for i in range(experts_per_rank)
101 )
102 rank_tokens.append(rank_sum)
103 max_etp = max(*rank_tokens, 1)
104 return (
105 2 * n_ff * max_etp * ccfg.h * ccfg.hff_exp
106 / (max(t_exp, 1) * ccfg.cp)
107 )
109 @staticmethod
110 def shared_expert_compute_cost(
111 ccfg: CostModelConfig, ctx: Context # pylint: disable=unused-argument
112 ) -> float:
113 """Shared expert compute FLOPs — replicated, not scaled by EP.
115 Shared experts process ALL tokens (not dispatched via EP),
116 so their compute is the same regardless of EP degree.
117 Shared experts use hff (not hff_exp) for their hidden dimension.
118 FLOPs = 2 * n_ffMM * s * b * h * hff * n_shared_exp / (t_exp * cp)
120 Factor t_exp because shared expert is TP/ETP-sharded (not EP-sharded).
121 t_exp = etp if etp > 1 else tp (alternative, not multiplicative).
122 n_ffMM is the number of feedforward linear layers per expert
123 (SwiGLU: 3, standard MLP: 2), set in arch_hooks.
124 """
125 t_exp = ccfg.etp if ccfg.etp > 1 else ccfg.t
126 n_ff = max(getattr(ccfg, "n_ffMM", 1), 1)
127 return (
128 2 * n_ff * ccfg.s * ccfg.b * ccfg.h * ccfg.hff
129 * ccfg.n_shared_exp / (max(t_exp, 1) * ccfg.cp)
130 )
132 @staticmethod
133 def expert_layer_compute(ccfg: CostModelConfig, ctx: Context) -> float:
134 """Dispatcher: routes to balanced or imbalanced compute based on tokens_per_expert."""
135 if ccfg.n_exp <= 1:
136 return 0
137 router = EvalExpertCompute.router_compute_cost(ccfg, ctx)
138 if ccfg.tokens_per_expert is not None:
139 expert = EvalExpertCompute.expert_compute_cost_imbalanced(ccfg, ctx)
140 else:
141 expert = EvalExpertCompute.expert_compute_cost_balanced(ccfg, ctx)
142 shared = EvalExpertCompute.shared_expert_compute_cost(ccfg, ctx)
143 return router + expert + shared