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"],

                    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} "

    @staticmethod
    def __adjust_framework_overhead(fw_oh: int, d_shard: int, d_replicate: int, tp: int) -> int:
        """Adjust framework overhead based on FSDP/HSDP/TP configuration."""
        if fw_oh > 0 and d_shard > 1:
            if d_replicate > 1:
                fw_oh = int(fw_oh * (1.0 + 0.03 * (d_shard - 1)))
            else:
                fw_oh = int(fw_oh * (1.0 + 0.01 * (d_shard - 1)))
        if fw_oh > 0 and d_replicate > 1:
            fw_oh = int(fw_oh / (1.0 + 0.025 * (d_replicate - 1)))
        if fw_oh > 0 and d_shard > 1 and d_replicate > 1:
            fw_oh = int(fw_oh / (1.0 + 0.03 * (d_shard - 1) * (d_replicate - 1) / d_shard))
        if fw_oh > 0 and d_replicate > 1 and tp >= 8:
            tp_hsdp_extra = int(fw_oh * 0.28)
            fw_oh += tp_hsdp_extra
        return fw_oh

    def __build_stage_insight(
        self, stage_id, raw_dyn, total_raw, total_fw_oh,

            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["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
        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)

        num_stages = self._ccfg.p
        safety_buffer = 1024 * 1024 * 1024  # 1 GB
        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"
                )

        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)

                        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

        """
        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

        """
        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:

        Use this in ``stat_*_p`` (static param memory) functions to add
        the FSDP all-gather buffer to the model-parameter memory budget.
        """
        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 in **bytes**.

        embedding gradient fragments.
        """
        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:
        """FSDP/HSDP gradient reduce-scatter buffer for embedding params (bytes).

    def hsdp_comm_mtp(ccfg: CostModelConfig, ctx: Context) -> float:
        """HSDP inter-node reduce-scatter buffer size (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"""

        Returns 0 for pure FSDP or non-FSDP configs.
        """
        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"""

    def __eq__(self, other: object) -> bool:
        """Check equality based on pp and vpp values."""
        if not isinstance(other, Pipeline):
            return NotImplemented
        return self.pp == other.pp and self.vpp == other.vpp


def infer_pp_and_vpp(offset: Any) -> Any:

            self.__maybe_set_int(target_ccfg, attr, value)
        if target_ccfg.use_seq_parallel:
            target_ccfg.sp = target_ccfg.t if target_ccfg.t > 1 else 1
        else:
            target_ccfg.sp = 1
        if target_ccfg.recompute_slice_activation and tp is not None:
            target_ccfg.shard_recompute_input = target_ccfg.t
            target_ccfg.shard_output_activ = target_ccfg.t
        if op and isinstance(op, int):

        self.ccfg.bytes_p = self.ccfg.fp_bytes(
            self.config.model.model_config.param_init_type,
        )
        if not self.ccfg.bytes_p:
            self.ccfg.bytes_p = self.ccfg.fp_bytes(
                self.config.model.model_config.params_dtype
            )
        self.ccfg.bytes_compute = self.ccfg.fp_bytes(
            self.config.model.model_config.compute_dtype

        self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
            self.config.model.model_config.softmax_compute_type
        )
        if not self.ccfg.bytes_softmax:
            self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
                self.config.model.model_config.softmax_compute_dtype
            )
        if not self.ccfg.bytes_p:
            raise AttributeError("bytes_p not positive")
        if not self.ccfg.bytes_compute:
            raise AttributeError("bytes_compute not positive")
        self.ccfg.bytes_grad = 4
        self.ccfg.bytes_os = 4
        self.ccfg.bytes_norm = 4

            self.ccfg.framework_overhead = int(
                10.5 * 1024 ** 3 * (cap_gb / 32.0) ** 1.74
            )
        else:
            self.ccfg.framework_overhead = getattr(
                self.ccfg, "framework_overhead", 0
            )

    def __config_parse_yaml_optimizer_shard(self):

        parts[RealParts.DP_WAIT][-1] += real["op_wait"]

    tp = "tp_wait"
    if tp in real.keys():
        parts[RealParts.MP_WAIT][-1] += real["tp_wait"]

    sp = "sp_wait"
    if sp in real.keys():
        parts[RealParts.MP_WAIT][-1] += real["sp_wait"]

            valid = dims_val[MBN] >= dims_val[PP]
            valid = valid and not (dims_val[PP] == 1 and dims_val[MBN] > 1)
            if not valid:
                return False
        return True

    @staticmethod
    def _check_power_of_two(dim, value):
        """Check whether *value* is a power of 2 for the given dimension.

        """Validate that OP is a power of 2."""
        if OP in self.all_dims:
            op = self.dims_val[OP]
            if not (op & (op - 1)) == 0:
                logger.warning("OP %d must be a power of 2", op)
                return False
        return True

    def _validate_fsdp_constraints(self):
        """Validate FSDP coexistence with OP, EP and the DP lower bound."""

        # If MB (microbatch number) is not in the config dimensions,
        # inject it so set_strategy can compute m from gbs/d/b.
        injected_mbn = False
        if Dim.MBN not in config.dims_val and Dim.MBS in config.dims_val and Dim.DP in config.dims_val:
            dp_val = config.dims_val[Dim.DP]
            mbs_val = config.dims_val[Dim.MBS]
            if isinstance(dp_val, int) and isinstance(mbs_val, int) and dp_val > 0 and mbs_val > 0:
                mbn = self.global_batch_size // dp_val // mbs_val
                config.dims_val[Dim.MBN] = mbn
                injected_mbn = True
        self.config.set_parallel_config(config)
        if injected_mbn:
            del config.dims_val[Dim.MBN]
        peak_mem = self.memory_estim()
        est_kwargs = {
            "debugger": debugger,
            "device_type": self.machine.device,

        if cfg.n_exp > 1:
            fsdp_intra_vol += exp * bc * 2
        d_shard_local = cfg.d_shard_or_d
        if getattr(cfg, "comm_hsdp", 0) > 0 and d_shard_local < cfg.d:
            sharded_non_exp = non_exp / (d_shard_local * cfg.cp * cfg.t)
            sharded_exp = exp / (d_shard_local * cfg.cp * cfg.t_exp) if cfg.n_exp > 1 else 0
            hsdp_inter_vol += (sharded_non_exp + sharded_exp) * bc
    fsdp_layer_count = 1 if is_fsdp_layer else 0

    comm[Dim.TP] += EvalLayerComm.tp_comm_layer(cfg, ctx, 1)
    tp_layer_count += 1

        comm[Dim.TP] = _flop_mode_tp_comm(
            cfg, d_shard_val, param["device_type"], mb=mb,
        )
    else:
        tp_scaling = (_FLOP_TP_A + _FLOP_TP_B * tp_val + _FLOP_TP_C / tp_val) / max(d_val, 1)
        comm[Dim.TP] *= max(1, tp_val // dev_per_node) * tp_scaling

    comm[Dim.EP] *= max(1, param["cfg"].ep // dev_per_node)
    comm[Dim.CP] *= max(1, param["cfg"].cp // dev_per_node)

    comm[Dim.DP] = _flop_mode_dp_comm(comm[Dim.DP], cfg, d_shard_val, param["device_type"], mb=mb)

    if (cfg.fsdp or d_shard_val > 1) and fsdp_layer_count > 0:
        comm[Dim.FSDP] = _flop_mode_fsdp_comm(
            cfg, fsdp_layer_count, d_shard_val, param["device_type"],
            pp=pp, mb=mb,
        )

        return 300
    if level == NetworkLevel.CLUSTER:
        return 25
    if device is not None:
        idx = level.value - 1 if isinstance(level, NetworkLevel) else level - 1
        if 0 <= idx < len(device.level_bandwidth):
            return device.level_bandwidth[idx]
    raise ValueError


def level_latency(level: Union[int, NetworkLevel], device: Any = None) -> Any:

    FSDP/HSDP all-gather/reduce-scatter happens across d_shard*cp*t ranks.
    This function returns a list of n_level per hierarchy level for that
    shard group, independent of how DP/TP/CP are individually assigned.
    """
    remaining = shard_size
    levels = []
    for level in range(device.levels):
        bound = device.level_bound_number[level]
        if bound:
            n = min(remaining, bound)
            remaining = remaining // n if n > 0 else remaining
        else:
            n = remaining
        levels.append(n)
    return levels


def estimate_comm_score(
    cfg: Any,

    size (not total).  Skip the /n_level division.  Use for FSDP/HSDP
    where the volume is pre-computed as a per-rank shard.
    """
    if comm_volume <= 0:
        return 0

    d_shard = cfg.d_shard_or_d

    if shard_group_size > 0:
        n_levels = _shard_group_levels(device, shard_group_size)
    else:
        assignment = device.level_assign(
            dp=cfg.d, tp=cfg.t, cp=cfg.cp, pp=cfg.p, d_shard=d_shard
        )

        n_level = n_levels[level]
        if n_level <= 1:
            continue
        if per_rank_msg:
            msg_per_rank = comm_volume
        elif shard_group_size > 0:
            prod_from_l = 1
            for k in range(level, len(n_levels)):
                prod_from_l *= n_levels[k]
            msg_per_rank = comm_volume / prod_from_l
        else:
            msg_per_rank = comm_volume / n_level
        if isinstance(a2a_efficiency, dict):
            eff = a2a_efficiency.get(level, level_efficiency(NetworkLevel(level + 1), device=device))
        elif a2a_efficiency is not None:
            eff = a2a_efficiency
        else:
            eff = level_efficiency(NetworkLevel(level + 1), device=device)
        bw_bps = device.level_bandwidth[level] * 1e9 * eff
        lat = level_latency(NetworkLevel(level + 1), device=device)

            eff = level_efficiency(NetworkLevel(level + 1), device=device)
        bw_bps = device.level_bandwidth[level] * 1e9 * eff
        lat = level_latency(NetworkLevel(level + 1), device=device)
        if per_rank_msg:
            collective_time = (
                (n_level - 1) * msg_per_rank / bw_bps
                + lat * 2 * (n_level - 1)
            )
        else:

        mb = cfg.m if hasattr(cfg, 'm') and cfg.m > 0 else 1
        comp_score = _flop_mode_comp_comm(cfg, d_shard_val, device_type=device_type, mb=mb)
        flops = [comp_score for _ in flops]
    else:
        comp_tp_scale = (_FLOP_COMP_A + _FLOP_COMP_B / tp_val) * tp_val / d_val
        flops = [f * comp_tp_scale for f in flops]

    return flops

def _estimate_non_hsdp_perf(cfg, ccfg, stage_perfs, debugger, device_type, stage_focused):
    """Estimate performance using the standard stage-based model."""
    perf = estimate_perf(
        cfg, ccfg, stage_perfs, stage_focused=stage_focused, debugger=debugger
    )
    perf += estimate_p2p(cfg, ccfg, stage_perfs, debugger=debugger, device_type=device_type)
    logger.info("PerfEst: perf %s", perf)
    return perf


# performance estimation
def estimate_performance(*args: Any, **kwargs: Any) -> Any:

    if is_hsdp:
        perf = _estimate_hsdp_perf(cfg, d_shard_val, device_type, debugger)
    else:
        stage_focused = kwargs.get("stage_focused", None)
        perf = _estimate_non_hsdp_perf(
            cfg, ccfg, stage_perfs, debugger, device_type, stage_focused
        )

    cache_file = kwargs.get("cache_file")