Diff Coverage

                real_id = self._ctx.real_lay_ids[chunk_id][i][lay_id]
                name = f"Lay_{real_id}_{lay_type.name[0]}"
                idx = chunk_id * len(dyn_mem_i) + lay_id
                stat = self.mb(stat_mem_i[chunk_id][lay_id])
                dyn = self.mb(sum(dyn_mem_i[chunk_id][lay_id]))
                ax.bar(
                    0.1,
                    [dyn],
                    bottom=bottoms["dyn"],

                mt: 0 for mt in list(MemType)
            }
        for mem_type in list(MemType):
            val = self._ctx.accu_mem_type[mem_type]
            if mem_type in (MemType.AG_COMM, MemType.A2A_COMM):
                stage_logs[stage_id].accu_mem_type[mem_type] = max(
                    stage_logs[stage_id].accu_mem_type[mem_type], val
                )
            else:
                stage_logs[stage_id].accu_mem_type[mem_type] += val

    def __preprocess_layer_custom_config_list(self, stages: list) -> list:
        """flatten layer_custom_config for backbone estimation"""
        flatten = sum(

                    self._ctx.current_lay_id = lay_id
                    self._ctx.current_node = node
                    static_mem = self._inner_static_mem()
                    sm["stat"][stage_id][chunk_id][lay_id] = static_mem
                    dyn_val = self._inner_dynamic_mem()
                    sm["dyn"][stage_id][chunk_id][lay_id] = dyn_val
                    if os.environ.get("_SAPP_DEBUG_DYN") == "1":
                        _sys.stderr.write(
                            f"DEBUG_DYN stage={stage_id} chunk={chunk_id} lay={lay_id} "
                            f"node={node} active={dyn_val[0] if isinstance(dyn_val, tuple) else dyn_val} "
                            f"comm={dyn_val[1] if isinstance(dyn_val, tuple) else 0} "
                            f"mf={self._ctx.micro_factor} h={self._ccfg.h} "

                self.__update_stage_logs(sm["logs"], stage_id)

    def __compute_raw_dynamic(self, num_stages, stages, record_lay_types, sm, safety_buffer):
        """Compute raw dynamic memory per stage."""
        raw_dyn = []
        for stage_id in range(num_stages):
            act_sum = sum(
                sum(aval for aval, _ in cell) for cell in sm["dyn"][stage_id]
            )
            comm_peak = max(
                max(comm for _, comm in cell) for cell in sm["dyn"][stage_id]
            )
            dyn = act_sum + comm_peak
            self._ctx.init_tmp_buff()
            if not self._ccfg.freeze:
                dyn += self._overhead_obj.estimate(
                    stages, stage_id, record_lay_types
                )
                self.__update_stage_logs(sm["logs"], stage_id)
            if dyn > 0:
                dyn += safety_buffer
            raw_dyn.append(dyn)
        return raw_dyn

    @staticmethod
    def __adjust_framework_overhead(fw_oh: int, d_shard: int, d_replicate: int, tp: int,
                                     ccfg=None) -> int:

    @staticmethod
    def __adjust_framework_overhead(fw_oh: int, d_shard: int, d_replicate: int, tp: int,
                                     ccfg=None) -> int:
        """Adjust framework overhead based on FSDP/HSDP/TP configuration."""
        hsdp_coeff = getattr(ccfg, 'fw_oh_hsdp_shard_coeff', 0.03) if ccfg else 0.03
        fsdp_coeff = getattr(ccfg, 'fw_oh_fsdp_shard_coeff', 0.01) if ccfg else 0.01
        rep_discount = getattr(ccfg, 'fw_oh_replicate_discount', 0.025) if ccfg else 0.025
        shard_rep_interact = getattr(ccfg, 'fw_oh_shard_replicate_interact', 0.03) if ccfg else 0.03
        tp_threshold = getattr(ccfg, 'fw_oh_tp_hsdp_threshold', 8) if ccfg else 8
        tp_extra_ratio = getattr(ccfg, 'fw_oh_tp_hsdp_extra_ratio', 0.28) if ccfg else 0.28

        if fw_oh > 0 and d_shard > 1:
            if d_replicate > 1:
                fw_oh = int(fw_oh * (1.0 + hsdp_coeff * (d_shard - 1)))
            else:
                fw_oh = int(fw_oh * (1.0 + fsdp_coeff * (d_shard - 1)))
        if fw_oh > 0 and d_replicate > 1:
            fw_oh = int(fw_oh / (1.0 + rep_discount * (d_replicate - 1)))
        if fw_oh > 0 and d_shard > 1 and d_replicate > 1:
            fw_oh = int(fw_oh / (1.0 + shard_rep_interact * (d_shard - 1) * (d_replicate - 1) / d_shard))
        if fw_oh > 0 and d_replicate > 1 and tp >= tp_threshold:
            tp_hsdp_extra = int(fw_oh * tp_extra_ratio)
            fw_oh += tp_hsdp_extra
        return fw_oh

    def __build_stage_insight(
        self, stage_id, raw_dyn, total_raw, total_fw_oh,
        num_stages, sm, verbose, spec_stage_id, insights,

        self, stage_id, raw_dyn, total_raw, total_fw_oh,
        num_stages, sm, verbose, spec_stage_id, insights,
    ):
        """Build memory insight dict for a single stage."""
        ins = {}  # Mem Insights purpose
        ins["Static"] = sum(
            sum(mem for mem in c) for c in sm["stat"][stage_id]
        )
        ins["Dynamic"] = raw_dyn[stage_id]
        if total_raw > 0 and ins["Dynamic"] > 0:
            fw_oh_share = total_fw_oh * (
                0.5 / num_stages + 0.5 * raw_dyn[stage_id] / total_raw
            )
            ins["Dynamic"] += int(fw_oh_share)
        elif ins["Dynamic"] > 0:
            ins["Dynamic"] += total_fw_oh // num_stages
        stage_accu = sm["logs"][stage_id].accu_mem_type
        ins["ModelParameters"] = self.mb(stage_accu[MemType.MODEL_PARAM])
        ins["OptimizerStates"] = self.mb(stage_accu[MemType.OPTIM_STATE])
        ins["AccumulGradients"] = self.mb(stage_accu[MemType.ACCU_GRAD])
        ins["Attn"] = self.mb(stage_accu[MemType.ATTN_ACTIV])
        ins["FFn"] = self.mb(stage_accu[MemType.FFN_ACTIV])
        ins["Norm"] = self.mb(stage_accu[MemType.NORM_ACTIV])
        ins["AllGather Comm"] = self.mb(stage_accu[MemType.AG_COMM])
        ins["All2All Comm"] = self.mb(stage_accu[MemType.A2A_COMM])

        if self._ccfg.cp > 1:
            cp_memory = EvalBody.act_cp_layer(self._ccfg, self._ctx)
            cp_comm_buffer = EvalLayerComm.cp_comm_buffer(self._ccfg, self._ctx)

            ins["CP KV Cache"] = self.mb(cp_memory.kv_cache_memory)
            ins["CP Attn Scores"] = self.mb(cp_memory.attention_scores_memory)
            ins["CP Softmax"] = self.mb(cp_memory.softmax_outputs_memory)
            ins["CP Comm Buffer"] = self.mb(cp_comm_buffer)
            ins["CP Reduction"] = self.mb(cp_memory.total_reduction)

        ins["Node Log"] = sm["logs"][stage_id].node_compute_log
        # VERBOSE
        if verbose and spec_stage_id in (-1, stage_id):
            self.__verbose_insights(sm, stage_id, ins)
        ins["Static"] = self.mb(ins["Static"])
        ins["Dynamic"] = self.mb(ins["Dynamic"])
        insights.append(ins)

    def __postprocess_stages(self, *args):
        """Build memory insights from raw stage evaluation buffers."""
        stages, record_lay_types = args[0], args[1]
        sm = args[2]
        verbose, spec_stage_id = args[3], args[4]
        insights = args[5]
        num_stages = self._ccfg.p
        safety_buffer = int(getattr(self._ccfg, 'safety_buffer_gb', 1)) * 1024 * 1024 * 1024
        raw_dyn = self.__compute_raw_dynamic(num_stages, stages, record_lay_types, sm, safety_buffer)
        if os.environ.get("_SAPP_DEBUG_DYN") == "1":
            for sid in range(num_stages):
                act_s = sum(
                    sum(aval for aval, _ in cell) for cell in sm["dyn"][sid]
                )
                comm_p = max(
                    max(comm for _, comm in cell) for cell in sm["dyn"][sid]
                )
                _sys.stderr.write(
                    f"DEBUG_POST stage={sid} act_sum={act_s/(1024**2):.1f}MB "
                    f"comm_peak={comm_p/(1024**2):.1f}MB "
                    f"raw_dyn={raw_dyn[sid]/(1024**2):.1f}MB\n"
                )
        total_raw = sum(raw_dyn)
        fw_oh = self._ccfg.framework_overhead
        d_shard = self._ccfg.d_shard_or_d
        d_replicate = self._ccfg.d // d_shard if d_shard > 0 else 1
        tp = self._ccfg.t
        fw_oh = self.__adjust_framework_overhead(fw_oh, d_shard, d_replicate, tp, ccfg=self._ccfg)
        total_fw_oh = fw_oh * num_stages
        for stage_id in range(num_stages):
            self.__build_stage_insight(
                stage_id, raw_dyn, total_raw, total_fw_oh, num_stages,
                sm, verbose, spec_stage_id, insights,
            )

        """potential overhead due to imbalanced chunks"""

        def dyn(chunk_id):
            """chunk total mem"""
            return sum(sum(v) for v in dyn_mem_i[chunk_id])

        # less mem
        micro_left = self._ccfg.m - self._ccfg.p
        vpp = self._ccfg.vp

        return max_overhead

    @staticmethod
    def _is_full_rec(node: Any) -> bool:
        return node == LayerType.FULL_REC_LAYER

    def __bwd_dyn_mem(self, is_recomputed: bool, **kwargs: Any) -> float:
        """Compute backward overhead dynamic memory."""
        activation, comm = self._inner_dynamic_mem(**kwargs)
        if is_recomputed:
            return activation + comm
        bwd_workspace = 0.0
        current_node = getattr(self._ctx, "current_node", None)
        if current_node == LayerType.OUTPUT_LAYER:
            t = getattr(self._ccfg, "t", 1)
            if t > 1:
                bwd_workspace = (
                    (1.0 - 1.0 / t)
                    * getattr(self._ccfg, "s", 0)
                    * getattr(self._ccfg, "b", 0)
                    * getattr(self._ccfg, "v", 0)

                    * getattr(self._ccfg, "b", 0)
                    * getattr(self._ccfg, "v", 0)
                    * getattr(self._ccfg, "bytes_compute", 0)
                )
        return comm + bwd_workspace

    def __stage_bwd_overhead_1f1b(
        self, stages: list, stage_id: int, record_lay_types: dict
    ) -> float:

        if stages[stage_id][self._ccfg.vp - 1]:
            last_node = self._fetch_node_and_switch_env(
                stages, record_lay_types, stage_id, -1, -1
            )
            if self._is_full_rec(last_node):
                # full rec -> not rec + grad
                # not rec -> grad
                self._ctx.current_lay_id = f"rec_{self._ctx.current_lay_id}"
                self._ctx.current_node = LayerType.NOT_REC_LAYER
                res = self.__bwd_dyn_mem(
                    is_recomputed=True, default_micro_factor=1
                )
            else:
                self._ctx.current_node = last_node

                )
            else:
                self._ctx.current_node = last_node
                self._ctx.current_lay_id = f"G_{self._ctx.current_lay_id}"
                res = self.__bwd_dyn_mem(
                    is_recomputed=False, default_micro_factor=1
                )
                if (
                    last_node == LayerType.OUTPUT_LAYER

                ):
                    last_mtp = self._fetch_node_and_switch_env(
                        stages, record_lay_types, stage_id, -1, -2
                    )
                    if self._is_full_rec(last_mtp):
                        self._ctx.current_lay_id = (
                            f"rec_{self._ctx.current_lay_id}"
                        )
                        self._ctx.current_node = LayerType.NOT_REC_LAYER
                        res += self.__bwd_dyn_mem(
                            is_recomputed=True, default_micro_factor=1
                        )
                    else:
                        self._ctx.current_lay_id = (

                        self._ctx.current_lay_id = (
                            f"G_{self._ctx.current_lay_id}"
                        )
                        self._ctx.current_node = last_mtp
                        res += self.__bwd_dyn_mem(
                            is_recomputed=False, default_micro_factor=1
                        )
        return res

    ep_imbalanced: Any = None
    fsdp_grad: Any = None

    def __repr__(self):
        parts = [
            f"dyn.comm.dp={_qname(self.dp)}",
            f"dyn.comm.tp={_qname(self.tp)}",
            f"dyn.comm.cp={_qname(self.cp)}",
            f"dyn.comm.ep={_qname(self.ep)}",

            f"dyn.comm.fsdp={_qname(self.fsdp)}",
            f"dyn.comm.hsdp={_qname(self.hsdp)}",
            f"dyn.comm.fsdp_grad={_qname(self.fsdp_grad)}",
        ]
        return ", ".join(parts)


@dataclass
class NodeComputeEval:

        dyn_cp_comm = kwargs.get("dyn_cp_comm", c_comm)
        dyn_ep_comm = kwargs.get("dyn_ep_comm", c_comm)
        dyn_ep_comm_balanced = kwargs.get("dyn_ep_comm_balanced", None)
        dyn_ep_comm_imbalanced = kwargs.get("dyn_ep_comm_imbalanced", None)
        dyn_fsdp_comm = kwargs.get("dyn_fsdp_comm", c_comm)
        dyn_hsdp_comm = kwargs.get("dyn_hsdp_comm", c_comm)
        dyn_fsdp_grad_comm = kwargs.get("dyn_fsdp_grad_comm", c_comm)
        if not self.__is_valid_eval_func(dyn_dp_comm):
            dyn_dp_comm = self._ctx.node_eval[target_node].dyn.comm.dp
        if not self.__is_valid_eval_func(dyn_tp_comm):
            dyn_tp_comm = self._ctx.node_eval[target_node].dyn.comm.tp

        if not self.__is_valid_eval_func(dyn_cp_comm):
            dyn_cp_comm = self._ctx.node_eval[target_node].dyn.comm.cp
        if not self.__is_valid_eval_func(dyn_ep_comm):
            dyn_ep_comm = self._ctx.node_eval[target_node].dyn.comm.ep
        if not self.__is_valid_eval_func(dyn_fsdp_comm):
            if target_node in self._ctx.node_eval:
                dyn_fsdp_comm = self._ctx.node_eval[target_node].dyn.comm.fsdp
            else:
                dyn_fsdp_comm = 0
        if not self.__is_valid_eval_func(dyn_hsdp_comm):
            if target_node in self._ctx.node_eval:
                dyn_hsdp_comm = self._ctx.node_eval[target_node].dyn.comm.hsdp
            else:
                dyn_hsdp_comm = 0
        if not self.__is_valid_eval_func(dyn_fsdp_grad_comm):
            if target_node in self._ctx.node_eval:
                dyn_fsdp_grad_comm = self._ctx.node_eval[target_node].dyn.comm.fsdp_grad
            if not self.__is_valid_eval_func(dyn_fsdp_grad_comm):
                dyn_fsdp_grad_comm = 0
        comm_cls_obj = cls_obj
        if self.is_regular_layer(target_node):
            comm_cls_obj = EvalLayerComm
        ep_balanced = (

                "dyn_ep_comm": comm_cfg.ep,
                "dyn_fsdp_comm": comm_cfg.fsdp,
                "dyn_hsdp_comm": comm_cfg.hsdp,
            }
            if hasattr(comm_cfg, "fsdp_grad"):
                comm_kwargs["dyn_fsdp_grad_comm"] = comm_cfg.fsdp_grad
            if hasattr(comm_cfg, "ep_balanced"):
                comm_kwargs["dyn_ep_comm_balanced"] = comm_cfg.ep_balanced
            if hasattr(comm_cfg, "ep_imbalanced"):
                comm_kwargs["dyn_ep_comm_imbalanced"] = comm_cfg.ep_imbalanced

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._child_cls = self
        self.ppb = None

    def reset_config(self) -> None:
        """reset current config"""
        self.update_config(self.config_path)

        non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
        # FSDP stores sharded parameters in compute dtype (e.g. bf16=2 bytes),
        # not param_init_type (fp32=4 bytes).  Without this, param memory is
        # overcounted by bytes_p/bytes_compute (typically 2x).
        bytes_param = ccfg.bytes_compute if getattr(ccfg, "fsdp", False) else ccfg.bytes_p
        # Routed experts: EP sharding
        routed_mem = routed_p / ccfg.ep * bytes_param / ccfg.shard_p_os_exp
        # Shared experts: partial DP sharding
        shared_mem = shared_p * bytes_param / ccfg.shard_p_os_exp_partial
        # Non expert
        non_exp_mem = non_exp_p * bytes_param / ccfg.shard_p_os_non_exp_partial
        # NOTE: FSDP all-gather buffer was previously included here but is a
        # dynamic (temporary) buffer, not static param memory.  It is already
        # accounted for in the dynamic comm path via ``fsdp_buffer_comm`` in
        # the YAML config (``comm_expr: max(fsdp + hsdp, fsdp_grad)

        """gradients"""
        non_exp_p, routed_p, shared_p = ctx.eval.num_p(ccfg, ctx)
        # FSDP stores sharded gradients in compute dtype (e.g. bf16=2 bytes),
        # not fp32 (4 bytes).  Without this, gradient memory is overcounted.
        bytes_g = ccfg.bytes_compute if getattr(ccfg, "fsdp", False) else ccfg.bytes_grad
        # Routed experts
        routed_mem = routed_p / ccfg.ep * bytes_g / ccfg.shard_grad_exp
        # Shared experts: use shard_grad_exp_partial (independent of os sharding)
        shared_mem = shared_p * bytes_g / ccfg.shard_grad_exp_partial
        # Non expert
        non_exp_mem = non_exp_p * bytes_g / ccfg.shard_grad_non_exp
        return non_exp_mem + routed_mem + shared_mem

    # No recompute and select recompute

            dp_comm_non_exp += non_exp / (ccfg.cp * ccfg.t)
            dp_comm_non_exp += non_exp / ccfg.t
        # Non expert ZeRO LvL 3 / FSDP
        if ccfg.comm_d_non_exp == 3:
            dp_comm_non_exp += non_exp / (ccfg.cp * ccfg.t)
            if not ccfg.fsdp:
                dp_comm_non_exp += non_exp / ccfg.t
                dp_comm_non_exp += non_exp / ccfg.t
        return dp_comm_non_exp

    @staticmethod
    def dp_comm_exp(ccfg: CostModelConfig, ctx: Context) -> float:

            dp_comm_exp += exp_param_size / max(ccfg.ep, ccfg.t_exp)
        # Expert ZeRO LvL 3 / FSDP
        if ccfg.comm_d_exp == 3:
            dp_comm_exp += exp_param_size / (ccfg.cp * ccfg.t_exp * ccfg.ep)
            if not ccfg.fsdp:
                dp_comm_exp += exp_param_size / max(ccfg.ep, ccfg.t_exp)
                dp_comm_exp += exp_param_size / max(ccfg.ep, ccfg.t_exp)
        return dp_comm_exp

    @staticmethod
    def dp_comm_layer(ccfg: CostModelConfig, ctx: Context) -> float:

        per layer, routed through TP bandwidth.
        Result is in bytes (like TP activation comm), unlike CP/DP which are
        in element counts (parameter comm).
        """
        if ccfg.comm_ep == 0:
            return 0
        if ccfg.n_exp <= 1:
            return 0
        if ccfg.ep <= 1:
            return 0
        t_local = mb * ccfg.n_chosen_exp * ccfg.s * ccfg.b / ccfg.cp
        t_cross = t_local * (ccfg.ep - 1) / ccfg.ep
        return t_cross * ccfg.h * ccfg.bytes_compute * 2 * ccfg.comm_ep

        Result is in bytes (like TP activation comm), unlike CP/DP which are
        in element counts (parameter comm).
        """
        if ccfg.comm_ep == 0:
            return 0
        if ccfg.n_exp <= 1:
            return 0
        if ccfg.ep <= 1:
            return 0
        tokens = ccfg.tokens_per_expert
        if not tokens:
            return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)

    @staticmethod
    def ep_comm_layer(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
        """EP comm dispatcher: balanced or imbalanced based on tokens_per_expert."""
        if ccfg.comm_ep == 0:
            return 0
        if ccfg.n_exp <= 1:
            return 0
        if ccfg.ep <= 1:
            return 0
        if ccfg.tokens_per_expert is not None:
            return EvalLayerComm.ep_comm_layer_imbalanced(ccfg, ctx, mb)
        return EvalLayerComm.ep_comm_layer_balanced(ccfg, ctx, mb)

    @staticmethod
    def fsdp_comm_layer(ccfg: CostModelConfig, ctx: Context) -> float:
        """FSDP/HSDP communication volume estimation."""
        non_exp, routed, shared = ctx.eval.num_p(ccfg, ctx)
        exp = routed + shared
        d_shard = ccfg.d_shard_or_d

        non_exp_comm = (
            ccfg.comm_fsdp
            * non_exp / (d_shard * ccfg.cp * ccfg.t)
            * ccfg.bytes_compute
            * 2

            * non_exp / (d_shard * ccfg.cp * ccfg.t)
            * ccfg.bytes_compute
            * 2
        )
        exp_comm = (
            ccfg.comm_fsdp
            * exp / (d_shard * ccfg.ep * ccfg.cp * ccfg.t_exp)
            * ccfg.bytes_compute
            * 2

            if ccfg.n_exp > 1
            else 0
        )

        if getattr(ccfg, "comm_hsdp", 0) > 0:
            d_replicate = ccfg.d // d_shard
            sharded_non_exp = non_exp / (d_shard * ccfg.cp * ccfg.t)
            sharded_exp = exp / (d_shard * ccfg.cp * ccfg.t_exp) if ccfg.n_exp > 1 else 0
            hsdp_comm = ccfg.comm_hsdp / d_replicate * (sharded_non_exp + sharded_exp) * ccfg.bytes_compute
            non_exp_comm += hsdp_comm

        return non_exp_comm + exp_comm

    @staticmethod
    def fsdp_buffer_layer(ccfg: CostModelConfig, ctx: Context) -> float:
        """FSDP/HSDP all-gather buffer memory (bytes)."""
        non_exp, routed, shared = ctx.eval.num_p(ccfg, ctx)
        non_exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * non_exp * ccfg.bytes_p / (ccfg.cp * ccfg.t)
        )
        exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * (routed + shared) * ccfg.bytes_p
            / (ccfg.ep * ccfg.cp * ccfg.t_exp)
            if ccfg.n_exp > 1

            / (ccfg.ep * ccfg.cp * ccfg.t_exp)
            if ccfg.n_exp > 1
            else 0
        )
        return non_exp_buf + exp_buf

    @staticmethod
    def fsdp_buffer_comm(ccfg: CostModelConfig, ctx: Context) -> float:
        """FSDP/HSDP all-gather buffer size (bytes)."""
        non_exp, routed, shared = ctx.eval.num_p(ccfg, ctx)
        non_exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * non_exp * ccfg.bytes_compute / (ccfg.cp * ccfg.t)
        )
        exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * (routed + shared) * ccfg.bytes_compute
            / (ccfg.ep * ccfg.cp * ccfg.t_exp)
            if ccfg.n_exp > 1

            / (ccfg.ep * ccfg.cp * ccfg.t_exp)
            if ccfg.n_exp > 1
            else 0
        )
        return non_exp_buf + exp_buf

    @staticmethod
    def fsdp_grad_buffer_comm(ccfg: CostModelConfig, ctx: Context) -> float:
        """FSDP/HSDP gradient reduce-scatter buffer size (bytes)."""
        non_exp, routed, shared = ctx.eval.num_p(ccfg, ctx)
        non_exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * non_exp * ccfg.bytes_grad / (ccfg.cp * ccfg.t)
        )
        exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * (routed + shared) * ccfg.bytes_grad
            / (ccfg.ep * ccfg.cp * ccfg.t_exp)
            if ccfg.n_exp > 1

            / (ccfg.ep * ccfg.cp * ccfg.t_exp)
            if ccfg.n_exp > 1
            else 0
        )
        return non_exp_buf + exp_buf

    @staticmethod
    def hsdp_inter_buffer_comm(ccfg: CostModelConfig, ctx: Context) -> float:
        """HSDP inter-node reduce-scatter buffer size (bytes)."""
        if getattr(ccfg, "comm_hsdp", 0) <= 0:
            return 0.0
        non_exp, routed, shared = ctx.eval.num_p(ccfg, ctx)
        exp = routed + shared
        d_shard = ccfg.d_shard_or_d
        sharded_non_exp = non_exp / (d_shard * ccfg.cp * ccfg.t)
        sharded_exp = (
            exp / (d_shard * ccfg.cp * ccfg.t_exp)
            if ccfg.n_exp > 1
            else 0
        )
        return (
            ccfg.comm_hsdp
            * (sharded_non_exp + sharded_exp)
            * ccfg.bytes_compute
        )

        )

    @staticmethod
    def dp_buffer_comm(ccfg: CostModelConfig, ctx: Context) -> float:
        del ccfg, ctx
        return 0.0

    @staticmethod
    def tp_buffer_comm(ccfg: CostModelConfig, ctx: Context, mb: int) -> float:
        """Peak HBM buffer for TP all-gather (bytes)."""
        if ccfg.t <= 1 or ccfg.comm_t == 0:
            return 0
        rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
        rec_coeff = EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather)

        attn_buf = ccfg.s * ccfg.b * ccfg.h * ccfg.bytes_compute
        ffn_buf = ccfg.s * ccfg.b * 2 * ccfg.hff * ccfg.bytes_compute

        peak_buf = max(attn_buf, ffn_buf)
        return rec_coeff * ccfg.comm_t * mb * peak_buf / ccfg.cp

        if ccfg.tie_emb_out:
            return 0
        param_size = ctx.eval.num_p(ccfg, ctx)
        param_size /= ccfg.shard_embed
        b_p = ccfg.bytes_compute if getattr(ccfg, "fsdp", False) else ccfg.bytes_p
        b_p /= ccfg.cp
        return param_size * b_p

    @staticmethod

        if ccfg.tie_emb_out:
            return 0
        param_size = ctx.eval.num_p(ccfg, ctx)
        param_size /= ccfg.shard_embed
        b_grad = ccfg.bytes_compute if getattr(ccfg, "fsdp", False) else ccfg.bytes_grad
        b_grad /= ccfg.cp
        return param_size * b_grad

    @staticmethod

        return param_size * b_grad

    @staticmethod
    def dp_comm_embed(ccfg: CostModelConfig, ctx: Context) -> float:
        if getattr(ccfg, "fsdp", False):
            return 0.0
        return (
            ccfg.comm_d_non_exp
            * ctx.eval.num_p(ccfg, ctx)
            / (ccfg.shard_embed * ccfg.cp)

        )

    @staticmethod
    def tp_comm_embed(ccfg: CostModelConfig, _) -> float:
        if ccfg.t <= 1 or ccfg.comm_t == 0:
            return 0
        return (
            ccfg.comm_t
            * ccfg.s
            * ccfg.b

        )

    @staticmethod
    def fsdp_comm_embed(ccfg: CostModelConfig, ctx: Context) -> float:
        param_size = ctx.eval.num_p(ccfg, ctx)
        non_exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * param_size * ccfg.bytes_compute / (ccfg.cp * ccfg.t)
        )
        return non_exp_buf

    @staticmethod
    def hsdp_comm_embed(ccfg: CostModelConfig, ctx: Context) -> float:
        if getattr(ccfg, "comm_hsdp", 0) <= 0:
            return 0.0
        param_size = ctx.eval.num_p(ccfg, ctx)
        d_shard = ccfg.d_shard_or_d
        sharded_size = param_size / (d_shard * ccfg.cp * ccfg.t)
        return ccfg.comm_hsdp * sharded_size * ccfg.bytes_compute

    @staticmethod
    def fsdp_grad_comm_embed(ccfg: CostModelConfig, ctx: Context) -> float:
        param_size = ctx.eval.num_p(ccfg, ctx)
        non_exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * param_size * ccfg.bytes_grad / (ccfg.cp * ccfg.t)
        )
        return non_exp_buf

    @staticmethod
    def activ_embed(ccfg: CostModelConfig, ctx: Context) -> float:
        """activations"""

    @staticmethod
    def fsdp_comm_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
        """FSDP/HSDP all-gather buffer size for MTP (bytes)."""
        if not ccfg.n_mtp:
            return 0
        mtp_fsdp_comm = 0.0
        param_size = EvalMTP.num_params_mtp(ccfg, ctx)
        mtp_fsdp_comm += (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * ccfg.n_mtp * param_size * ccfg.bytes_compute / (ccfg.cp * ccfg.t)
        )
        ctx.current_node = LayerType.EMBEDDING_LAYER
        param_size = ctx.eval.num_p(ccfg, ctx)
        mtp_fsdp_comm += (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * ccfg.n_mtp * param_size * ccfg.bytes_compute / (ccfg.cp * ccfg.t)
        )
        mtp_fsdp_comm += ccfg.n_mtp * ctx.eval.dyn.comm.fsdp(ccfg, ctx)
        ctx.current_node = LayerType.OUTPUT_LAYER
        param_size = ctx.eval.num_p(ccfg, ctx)
        mtp_fsdp_comm += (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * ccfg.n_mtp * param_size * ccfg.bytes_compute / (ccfg.cp * ccfg.t)
        )
        mtp_fsdp_comm += ccfg.n_mtp * EvalTailSingle.fsdp_comm_out_single(
            ccfg, ctx
        )
        return mtp_fsdp_comm

    @staticmethod
    def fsdp_grad_comm_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
        """FSDP/HSDP gradient reduce-scatter buffer size for MTP (bytes)."""
        if not ccfg.n_mtp:
            return 0
        mtp_fsdp_grad_comm = 0.0
        param_size = EvalMTP.num_params_mtp(ccfg, ctx)
        mtp_fsdp_grad_comm += (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * ccfg.n_mtp * param_size * ccfg.bytes_grad / (ccfg.cp * ccfg.t)
        )
        ctx.current_node = LayerType.EMBEDDING_LAYER
        param_size = ctx.eval.num_p(ccfg, ctx)
        mtp_fsdp_grad_comm += (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * ccfg.n_mtp * param_size * ccfg.bytes_grad / (ccfg.cp * ccfg.t)
        )
        mtp_fsdp_grad_comm += ccfg.n_mtp * ctx.eval.dyn.comm.fsdp_grad(ccfg, ctx)
        ctx.current_node = LayerType.OUTPUT_LAYER
        param_size = ctx.eval.num_p(ccfg, ctx)
        mtp_fsdp_grad_comm += (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * ccfg.n_mtp * param_size * ccfg.bytes_grad / (ccfg.cp * ccfg.t)
        )
        mtp_fsdp_grad_comm += ccfg.n_mtp * EvalTailSingle.fsdp_grad_comm_out_single(
            ccfg, ctx
        )
        return mtp_fsdp_grad_comm

    @staticmethod
    def hsdp_comm_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
        """HSDP inter-node reduce-scatter buffer size for MTP (bytes)."""
        if not ccfg.n_mtp or getattr(ccfg, "comm_hsdp", 0) <= 0:
            return 0
        d_shard = ccfg.d_shard_or_d
        d_replicate = ccfg.d // d_shard
        mtp_hsdp_comm = 0.0
        param_size = EvalMTP.num_params_mtp(ccfg, ctx)
        sharded_size = param_size / (d_shard * ccfg.cp * ccfg.t)
        mtp_hsdp_comm += (
            ccfg.comm_hsdp * ccfg.n_mtp * sharded_size
            * ccfg.bytes_compute / d_replicate
        )
        ctx.current_node = LayerType.EMBEDDING_LAYER
        param_size = ctx.eval.num_p(ccfg, ctx)
        sharded_size = param_size / (d_shard * ccfg.cp * ccfg.t)
        mtp_hsdp_comm += (
            ccfg.comm_hsdp * ccfg.n_mtp * sharded_size
            * ccfg.bytes_compute / d_replicate
        )
        ctx.current_node = LayerType.OUTPUT_LAYER
        param_size = ctx.eval.num_p(ccfg, ctx)
        sharded_size = param_size / (d_shard * ccfg.cp * ccfg.t)
        mtp_hsdp_comm += (
            ccfg.comm_hsdp * ccfg.n_mtp * sharded_size
            * ccfg.bytes_compute / d_replicate
        )
        return mtp_hsdp_comm


class EvalTailSingle:
    """Single tail layer formulas class"""

    @staticmethod
    def stat_output_single_p(ccfg: CostModelConfig, ctx: Context) -> float:
        """static mem for model param (lmhead)"""
        param_size = ctx.eval.num_p(ccfg, ctx)
        b_p = ccfg.bytes_compute if getattr(ccfg, "fsdp", False) else ccfg.bytes_p
        b_p /= ccfg.shard_p_os_non_exp_partial
        return param_size * b_p

    @staticmethod

    @staticmethod
    def stat_output_single_grad(ccfg: CostModelConfig, ctx: Context) -> float:
        """static mem for gradient (lmhead)"""
        param_size = ctx.eval.num_p(ccfg, ctx)
        b_grad = ccfg.bytes_compute if getattr(ccfg, "fsdp", False) else ccfg.bytes_grad
        b_grad /= ccfg.shard_grad_non_exp
        return param_size * b_grad

    @staticmethod

        return micro_factor * activ_size

    @staticmethod
    def comm_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
        del ccfg, ctx
        return 0.0

    @staticmethod
    def fsdp_comm_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
        param_size = ctx.eval.num_p(ccfg, ctx)
        non_exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * param_size * ccfg.bytes_compute / (ccfg.cp * ccfg.t)
        )
        return non_exp_buf

    @staticmethod
    def fsdp_grad_comm_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
        param_size = ctx.eval.num_p(ccfg, ctx)
        non_exp_buf = (
            ccfg.comm_fsdp * ccfg.fsdp_all_gather_buffer
            * param_size * ccfg.bytes_grad / (ccfg.cp * ccfg.t)
        )
        return non_exp_buf

    @staticmethod
    def hsdp_comm_out_single(ccfg: CostModelConfig, ctx: Context) -> float:
        if getattr(ccfg, "comm_hsdp", 0) <= 0:
            return 0.0
        param_size = ctx.eval.num_p(ccfg, ctx)
        d_shard = ccfg.d_shard_or_d
        sharded_size = param_size / (d_shard * ccfg.cp * ccfg.t)
        return ccfg.comm_hsdp * sharded_size * ccfg.bytes_compute


class EvalTail:
    """Single tail layer formulas class"""

        )

    @staticmethod
    def fsdp_comm_output(ccfg: CostModelConfig, ctx: Context) -> float:
        return sum(
            [
                EvalTailSingle.fsdp_comm_out_single(ccfg, ctx),
                EvalMTP.fsdp_comm_mtp(ccfg, ctx),
            ]

        )

    @staticmethod
    def hsdp_comm_output(ccfg: CostModelConfig, ctx: Context) -> float:
        return sum(
            [
                EvalTailSingle.hsdp_comm_out_single(ccfg, ctx),
                EvalMTP.hsdp_comm_mtp(ccfg, ctx),
            ]

        )

    @staticmethod
    def fsdp_grad_comm_output(ccfg: CostModelConfig, ctx: Context) -> float:
        return sum(
            [
                EvalTailSingle.fsdp_grad_comm_out_single(ccfg, ctx),
                EvalMTP.fsdp_grad_comm_mtp(ccfg, ctx),
            ]

    @staticmethod
    def __log_ast_mem(ctx, cat, mem):
        """Save AST memory evaluation result in the current context."""
        ctx.save2log(cat, mem)
        if cat is not None and cat.name in ("AG_COMM", "A2A_COMM"):
            ctx.accu_mem_type[cat] = max(ctx.accu_mem_type[cat], mem)
        else:
            ctx.accu_mem_type[cat] += mem

    @classmethod
    def __eval_ast_name(cls, n: ast.Name, depth: int, wait: bool, **kwargs):
        """Evaluate a named memory term."""