# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import urllib.error
import urllib.request
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple


GiB = 1024**3


def _to_bool(value: Any, default: bool) -> bool:
    if value is None:
        return default
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        if value.lower() in ("true", "1", "yes", "y"):
            return True
        if value.lower() in ("false", "0", "no", "n"):
            return False
    return default


_DTYPE_BYTES_TABLE: Dict[str, float] = {
    # 16-bit
    "float16": 2,
    "fp16": 2,
    "half": 2,
    "f16": 2,
    "bfloat16": 2,
    "bf16": 2,
    # 32-bit
    "float32": 4,
    "fp32": 4,
    "f32": 4,
    # 8-bit (raw storage bytes; real impl may need scales)
    "float8": 1,
    "fp8": 1,
    "float8_e4m3fn": 1,
    "float8_e4m3fnuz": 1,
    "float8_e5m2": 1,
    "int8": 1,
    "uint8": 1,
    # 4-bit packed storage (bytes/value; excludes scale/zero-point overhead)
    "int4": 0.5,
    "fp4": 0.5,
}


_DTYPE_HELP = (
    "float16/fp16/half, bfloat16/bf16, float32/fp32, float8/fp8, int8, uint8, int4, fp4 "
    "(also accepts 'torch.<dtype>')"
)


def dtype_bytes_from_str(dtype: Optional[str]) -> float:
    if dtype is None:
        raise ValueError("dtype is required")
    raw = str(dtype)
    key = raw.strip().lower()
    if not key:
        raise ValueError("dtype cannot be empty")
    if key.startswith("torch."):
        key = key[len("torch.") :]
    b = _DTYPE_BYTES_TABLE.get(key)
    if b is None:
        raise ValueError(f"Unsupported dtype: {raw!r}. Supported: {_DTYPE_HELP}.")
    return b


def format_count(x: float, *, unit: str = "") -> str:
    """Human readable big numbers (T/B/M/K)."""
    if x >= 1e12:
        return f"{x/1e12:.4f}T{unit}"
    if x >= 1e9:
        return f"{x/1e9:.4f}B{unit}"
    if x >= 1e6:
        return f"{x/1e6:.4f}M{unit}"
    if x >= 1e3:
        return f"{x/1e3:.4f}K{unit}"
    return f"{x:.0f}{unit}"


def format_gib(x: float) -> str:
    return f"{x:.3f} GiB"


def divisors(n: int) -> List[int]:
    if n <= 0:
        return []
    small: List[int] = []
    large: List[int] = []
    i = 1
    while i * i <= n:
        if n % i == 0:
            small.append(i)
            if i * i != n:
                large.append(n // i)
        i += 1
    return small + large[::-1]


@dataclass(frozen=True)
class ModelDims:
    dtype: str
    dtype_bytes: float
    hidden_size: int
    num_layers: int
    vocab_size: int
    tie_word_embeddings: bool

    # MLA
    num_heads: int
    q_lora_rank: int
    kv_lora_rank: int
    qk_nope_head_dim: int
    qk_rope_head_dim: int
    v_head_dim: int

    # MoE
    first_k_dense_replace: int
    intermediate_size: int
    moe_intermediate_size: int
    n_routed_experts: int
    n_shared_experts: int
    num_experts_per_tok: int

    # NSA
    index_n_heads: int
    index_head_dim: int
    index_topk: int
    max_position_embeddings: int

    @property
    def qk_head_dim(self) -> int:
        return self.qk_nope_head_dim + self.qk_rope_head_dim

    @property
    def moe_layers(self) -> int:
        return max(0, self.num_layers - self.first_k_dense_replace)

    @property
    def dense_layers(self) -> int:
        return min(self.first_k_dense_replace, self.num_layers)


@dataclass(frozen=True)
class MhcSpec:
    """mHC (multi-stream HyperConnections) spec (approximation).

    Assumption (mhc-torch style):
    - residual streams expand (B,T,D) -> (B*S,T,D) where S = streams.
    - each HyperConnections module has parameters:
        H_res_logits:  (S,S)
        H_pre_logits:  (S,)   (num_input_views=1)
        H_post_logits: (S,)
      so params/module ~= S^2 + 2S.
    - typical transformer integration uses 2 modules per layer (attn + mlp).
    """

    streams: int = 1
    branches_per_layer: int = 2
    branch_mode: str = "single"  # "single" | "per_stream"

    @property
    def enabled(self) -> bool:
        return self.streams > 1 and self.branches_per_layer > 0

    @property
    def branch_streams(self) -> int:
        """Effective number of streams that run attention/MLP branches.

        - single: branch stays single-stream (mhc-torch style) => 1
        - per_stream: each stream runs its own attention/MLP => streams
        """
        if str(self.branch_mode).lower() == "per_stream":
            return max(1, int(self.streams))
        return 1


def estimate_mhc_params_total(num_layers: int, mhc: Optional[MhcSpec]) -> int:
    if mhc is None or not mhc.enabled:
        return 0
    s = int(mhc.streams)
    k = int(mhc.branches_per_layer)
    if s <= 1 or k <= 0 or num_layers <= 0:
        return 0
    per_module = s * s + 2 * s
    return int(num_layers) * k * per_module


def load_dims(config: Dict[str, Any]) -> ModelDims:
    dtype = str(config.get("dtype") or config.get("torch_dtype") or "bfloat16")
    dtype_bytes = dtype_bytes_from_str(dtype)

    hidden_size = int(config["hidden_size"])
    num_layers = int(config["num_hidden_layers"])
    vocab_size = int(config.get("vocab_size", 0))
    tie_word_embeddings = _to_bool(config.get("tie_word_embeddings"), default=True)

    num_heads = int(config.get("num_attention_heads", 0))
    q_lora_rank = int(config.get("q_lora_rank", 0))
    kv_lora_rank = int(config.get("kv_lora_rank", 0))
    qk_nope_head_dim = int(config.get("qk_nope_head_dim", 0))
    qk_rope_head_dim = int(config.get("qk_rope_head_dim", 0))
    v_head_dim = int(config.get("v_head_dim", 0))

    intermediate_size = int(config.get("intermediate_size", 0))
    moe_intermediate_size = int(config.get("moe_intermediate_size", 0))
    n_routed_experts = int(config.get("n_routed_experts", 0))
    n_shared_experts = int(config.get("n_shared_experts", 0))
    num_experts_per_tok = int(config.get("num_experts_per_tok", 0))

    first_k_dense_replace_raw = config.get("first_k_dense_replace", None)
    if first_k_dense_replace_raw is None:
        has_moe = (n_routed_experts > 0) or (moe_intermediate_size > 0)
        # Default behavior:
        # - Dense-only models (no MoE dims) => all layers are dense
        # - MoE models without explicit split => assume all layers are MoE
        first_k_dense_replace = 0 if has_moe else num_layers
    else:
        first_k_dense_replace = int(first_k_dense_replace_raw)

    index_n_heads = int(config.get("index_n_heads", 0))
    index_head_dim = int(config.get("index_head_dim", 0))
    index_topk = int(config.get("index_topk", 0))
    max_position_embeddings = int(config.get("max_position_embeddings", 0))

    return ModelDims(
        dtype=dtype,
        dtype_bytes=dtype_bytes,
        hidden_size=hidden_size,
        num_layers=num_layers,
        vocab_size=vocab_size,
        tie_word_embeddings=tie_word_embeddings,
        num_heads=num_heads,
        q_lora_rank=q_lora_rank,
        kv_lora_rank=kv_lora_rank,
        qk_nope_head_dim=qk_nope_head_dim,
        qk_rope_head_dim=qk_rope_head_dim,
        v_head_dim=v_head_dim,
        first_k_dense_replace=first_k_dense_replace,
        intermediate_size=intermediate_size,
        moe_intermediate_size=moe_intermediate_size,
        n_routed_experts=n_routed_experts,
        n_shared_experts=n_shared_experts,
        num_experts_per_tok=num_experts_per_tok,
        index_n_heads=index_n_heads,
        index_head_dim=index_head_dim,
        index_topk=index_topk,
        max_position_embeddings=max_position_embeddings,
    )


@dataclass(frozen=True)
class ParamBreakdown:
    embedding: int
    lm_head: int

    # Attention
    attn_fused_qkv_a_total: int
    attn_q_b_total: int
    attn_kv_b_total: int
    attn_o_proj_total: int

    indexer_total: int
    dense_ffn_total: int

    # MoE
    moe_routed_expert_total: int
    moe_shared_expert_total: int
    moe_router_total: int
    mhc_total: int

    @property
    def attention_total(self) -> int:
        return (
            self.attn_fused_qkv_a_total
            + self.attn_q_b_total
            + self.attn_kv_b_total
            + self.attn_o_proj_total
        )

    @property
    def moe_expert_total(self) -> int:
        return self.moe_routed_expert_total + self.moe_shared_expert_total

    @property
    def total(self) -> int:
        return (
            self.embedding
            + self.lm_head
            + self.attention_total
            + self.indexer_total
            + self.dense_ffn_total
            + self.moe_expert_total
            + self.moe_router_total
            + self.mhc_total
        )

def estimate_params(d: ModelDims, *, mhc: Optional[MhcSpec] = None) -> ParamBreakdown:
    embedding = d.vocab_size * d.hidden_size
    lm_head = 0 if d.tie_word_embeddings else embedding

    need = (
        d.num_heads,
        d.q_lora_rank,
        d.kv_lora_rank,
        d.qk_nope_head_dim,
        d.qk_rope_head_dim,
        d.v_head_dim,
    )
    if not all(x > 0 for x in need):
        raise ValueError(
            "Config missing MLA attention dims; expected: "
            "num_attention_heads, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim"
        )

    attn_fused_qkv_a_per_layer = d.hidden_size * (
        d.q_lora_rank + d.kv_lora_rank + d.qk_rope_head_dim
    )
    attn_q_b_per_layer = d.q_lora_rank * (d.num_heads * d.qk_head_dim)
    attn_kv_b_per_layer = d.kv_lora_rank * (
        d.num_heads * (d.qk_nope_head_dim + d.v_head_dim)
    )
    attn_o_proj_per_layer = (d.num_heads * d.v_head_dim) * d.hidden_size

    attn_fused_qkv_a_total = attn_fused_qkv_a_per_layer * d.num_layers
    attn_q_b_total = attn_q_b_per_layer * d.num_layers
    attn_kv_b_total = attn_kv_b_per_layer * d.num_layers
    attn_o_proj_total = attn_o_proj_per_layer * d.num_layers

    # Indexer weights scale with index_n_heads, but index-cache size does not (pool stores a fixed head).
    if d.index_n_heads > 0 and d.index_head_dim > 0:
        wq_b = d.q_lora_rank * (d.index_n_heads * d.index_head_dim)
        wk = d.hidden_size * d.index_head_dim
        weights_proj = d.hidden_size * d.index_n_heads
        indexer_per_layer = wq_b + wk + weights_proj
    else:
        indexer_per_layer = 0
    indexer_total = indexer_per_layer * d.num_layers

    # Dense FFN
    if d.dense_layers > 0:
        if d.intermediate_size <= 0:
            raise ValueError("Config missing intermediate_size for dense layers")
        dense_ffn_per_layer = 3 * d.hidden_size * d.intermediate_size
        dense_ffn_total = dense_ffn_per_layer * d.dense_layers
    else:
        dense_ffn_total = 0

    # MoE experts
    if d.moe_layers > 0:
        if d.moe_intermediate_size <= 0 or d.n_routed_experts <= 0:
            raise ValueError(
                "Config missing MoE dims; expected moe_intermediate_size and n_routed_experts"
            )
        per_expert = 3 * d.hidden_size * d.moe_intermediate_size
        moe_routed_expert_total = per_expert * d.n_routed_experts * d.moe_layers
        moe_shared_expert_total = (
            per_expert * max(0, d.n_shared_experts) * d.moe_layers
        )
        moe_router_total = d.hidden_size * d.n_routed_experts * d.moe_layers
    else:
        moe_routed_expert_total = 0
        moe_shared_expert_total = 0
        moe_router_total = 0

    mhc_total = estimate_mhc_params_total(d.num_layers, mhc)

    return ParamBreakdown(
        embedding=embedding,
        lm_head=lm_head,
        attn_fused_qkv_a_total=attn_fused_qkv_a_total,
        attn_q_b_total=attn_q_b_total,
        attn_kv_b_total=attn_kv_b_total,
        attn_o_proj_total=attn_o_proj_total,
        indexer_total=indexer_total,
        dense_ffn_total=dense_ffn_total,
        moe_routed_expert_total=moe_routed_expert_total,
        moe_shared_expert_total=moe_shared_expert_total,
        moe_router_total=moe_router_total,
        mhc_total=mhc_total,
    )


def estimate_activated_params_per_token(d: ModelDims, p: ParamBreakdown) -> Dict[str, int]:
    per_expert = 3 * d.hidden_size * d.moe_intermediate_size
    active_experts = d.num_experts_per_tok + max(0, d.n_shared_experts)

    active_moe = per_expert * active_experts * d.moe_layers
    active_dense = 3 * d.hidden_size * d.intermediate_size * d.dense_layers
    active_mhc = p.mhc_total
    always_attn_indexer = p.attention_total + p.indexer_total
    always_active = always_attn_indexer + active_mhc

    total_excluding_embed_router = active_moe + active_dense + always_active
    total_including_router = total_excluding_embed_router + p.moe_router_total
    total_including_embed = total_including_router + p.embedding + p.lm_head
    return {
        "active_moe": active_moe,
        "active_dense_ffn": active_dense,
        "active_mhc": active_mhc,
        "always_active(attn+indexer)": always_attn_indexer,
        "always_active(attn+indexer+mhc)": always_active,
        "total_excluding_embed_router": total_excluding_embed_router,
        "total_including_router": total_including_router,
        "total_including_embed": total_including_embed,
    }


def estimate_kv_cache_gib(
    d: ModelDims, *, batch: float, seq_len: int, kv_dtype_bytes: float
) -> float:
    kv_dim = d.kv_lora_rank + d.qk_rope_head_dim
    elems = batch * seq_len * d.num_layers * kv_dim
    return (elems * kv_dtype_bytes) / GiB


def estimate_nsa_index_cache_gib(
    d: ModelDims, *, batch: float, seq_len: int
) -> Optional[float]:
    if d.index_head_dim <= 0 or d.index_n_heads <= 0:
        return None
    quant_block = 128
    blocks = math.ceil(d.index_head_dim / quant_block)
    bytes_per_token_per_layer = d.index_head_dim + blocks * 4
    total_bytes = batch * seq_len * d.num_layers * bytes_per_token_per_layer
    return total_bytes / GiB


@dataclass(frozen=True)
class Parallelism:
    tp: int = 1
    dp: int = 1
    sp: int = 1
    ep: int = 1
    moe_dp: int = 1

    @property
    def attn_tp(self) -> int:
        if self.tp <= 0:
            raise ValueError(
                f"Invalid parallelism: require tp >= 1, got {self.tp=}, {self.dp=}, {self.sp=}"
            )
        denom = self.dp * self.sp
        if denom <= 0 or self.tp % denom != 0:
            raise ValueError(
                f"Invalid parallelism: require tp % (dp*sp) == 0, got {self.tp=}, {self.dp=}, {self.sp=}"
            )
        return self.tp // denom

    @property
    def moe_tp(self) -> int:
        if self.tp <= 0:
            raise ValueError(
                f"Invalid parallelism: require tp >= 1, got {self.tp=}, {self.ep=}, {self.moe_dp=}"
            )
        denom = self.ep * self.moe_dp
        if denom <= 0 or self.tp % denom != 0:
            raise ValueError(
                f"Invalid parallelism: require tp % (ep*moe_dp) == 0, got {self.tp=}, {self.ep=}, {self.moe_dp=}"
            )
        return self.tp // denom


@dataclass(frozen=True)
class ShardedWeightBreakdown:
    embedding: float
    lm_head: float

    attn_fused_qkv_a: float
    attn_q_b: float
    attn_kv_b: float
    attn_o_proj: float

    indexer: float
    dense_ffn: float

    moe_routed_expert: float
    moe_shared_expert: float
    moe_router: float
    mhc: float

    @property
    def total(self) -> float:
        return (
            self.embedding
            + self.lm_head
            + self.attn_fused_qkv_a
            + self.attn_q_b
            + self.attn_kv_b
            + self.attn_o_proj
            + self.indexer
            + self.dense_ffn
            + self.moe_routed_expert
            + self.moe_shared_expert
            + self.moe_router
            + self.mhc
        )


def estimate_weight_params_per_rank(
    d: ModelDims, p: ParamBreakdown, parallel: Parallelism
) -> ShardedWeightBreakdown:
    tp_size = parallel.tp
    attn_tp = parallel.attn_tp
    moe_tp = parallel.moe_tp

    # Common inference implementation: vocab-parallel embedding / lm_head.
    embedding = p.embedding / float(tp_size) if tp_size > 0 else float(p.embedding)
    lm_head = p.lm_head / float(tp_size) if tp_size > 0 else float(p.lm_head)

    # DeepSeek-MLA in SGLang: fused_qkv_a is replicated; others are sharded by attention TP.
    attn_fused_qkv_a = float(p.attn_fused_qkv_a_total)
    attn_q_b = p.attn_q_b_total / float(attn_tp)
    attn_kv_b = p.attn_kv_b_total / float(attn_tp)
    attn_o_proj = p.attn_o_proj_total / float(attn_tp)

    # NSA indexer is replicated in SGLang.
    indexer = float(p.indexer_total)

    # Dense FFN sharded by TP.
    dense_ffn = p.dense_ffn_total / float(tp_size) if tp_size > 0 else float(p.dense_ffn_total)

    # MoE expert sharding:
    moe_routed = p.moe_routed_expert_total / float(parallel.ep * moe_tp)
    moe_shared = p.moe_shared_expert_total / float(moe_tp)
    moe_router = float(p.moe_router_total)
    mhc = float(p.mhc_total)

    return ShardedWeightBreakdown(
        embedding=embedding,
        lm_head=lm_head,
        attn_fused_qkv_a=attn_fused_qkv_a,
        attn_q_b=attn_q_b,
        attn_kv_b=attn_kv_b,
        attn_o_proj=attn_o_proj,
        indexer=indexer,
        dense_ffn=dense_ffn,
        moe_routed_expert=moe_routed,
        moe_shared_expert=moe_shared,
        moe_router=moe_router,
        mhc=mhc,
    )


@dataclass(frozen=True)
class ActivationWorkingSet:
    tokens: int
    dtype_bytes: float
    mhc_streams: int
    mhc_branch_streams: int
    local_heads: int
    dense_intermediate_per_partition: int
    intermediate_per_partition: int
    active_experts_per_tok: int

    hidden_elems: int
    attn_elems: int
    dense_elems: int
    moe_elems: int
    peak_elems: int

    hidden_gib: float
    attn_gib: float
    dense_gib: float
    moe_gib: float
    peak_gib: float


def estimate_activation_working_set(
    d: ModelDims,
    parallel: Parallelism,
    *,
    tokens: int,
    act_dtype_bytes: float,
    mhc_streams: int = 1,
    mhc_branch_streams: int = 1,
) -> ActivationWorkingSet:
    tokens = max(0, int(tokens))
    mhc_streams = max(1, int(mhc_streams))
    mhc_branch_streams = max(1, int(mhc_branch_streams))
    local_heads = int(math.ceil(d.num_heads / float(parallel.attn_tp)))
    dense_intermediate_per_partition = (
        int(math.ceil(d.intermediate_size / float(max(1, parallel.tp))))
        if d.dense_layers > 0 and d.intermediate_size > 0
        else 0
    )
    intermediate_per_partition = (
        int(math.ceil(d.moe_intermediate_size / float(max(1, parallel.moe_tp))))
        if d.moe_intermediate_size > 0
        else 0
    )
    active_experts_per_tok = d.num_experts_per_tok + max(0, d.n_shared_experts)

    hidden_elems = tokens * (2 * d.hidden_size) * mhc_streams
    attn_elems = tokens * (
        d.hidden_size
        + local_heads * d.qk_head_dim
        + local_heads * d.kv_lora_rank
    ) * mhc_branch_streams
    dense_gate_up = (
        tokens * (2 * dense_intermediate_per_partition) * mhc_branch_streams
        if d.dense_layers > 0
        else 0
    )
    dense_elems = (
        tokens * d.hidden_size * mhc_branch_streams + dense_gate_up
        if d.dense_layers > 0
        else 0
    )
    moe_gate_up = (
        tokens
        * active_experts_per_tok
        * (2 * intermediate_per_partition)
        * mhc_branch_streams
    )
    moe_elems = (
        tokens * d.hidden_size * mhc_branch_streams + moe_gate_up
        if d.moe_layers > 0
        else 0
    )

    peak_elems = max(hidden_elems, attn_elems, dense_elems, moe_elems)

    hidden_gib = (hidden_elems * act_dtype_bytes) / GiB
    attn_gib = (attn_elems * act_dtype_bytes) / GiB
    dense_gib = (dense_elems * act_dtype_bytes) / GiB
    moe_gib = (moe_elems * act_dtype_bytes) / GiB
    peak_gib = (peak_elems * act_dtype_bytes) / GiB

    return ActivationWorkingSet(
        tokens=tokens,
        dtype_bytes=act_dtype_bytes,
        mhc_streams=mhc_streams,
        mhc_branch_streams=mhc_branch_streams,
        local_heads=local_heads,
        dense_intermediate_per_partition=dense_intermediate_per_partition,
        intermediate_per_partition=intermediate_per_partition,
        active_experts_per_tok=active_experts_per_tok,
        hidden_elems=hidden_elems,
        attn_elems=attn_elems,
        dense_elems=dense_elems,
        moe_elems=moe_elems,
        peak_elems=peak_elems,
        hidden_gib=hidden_gib,
        attn_gib=attn_gib,
        dense_gib=dense_gib,
        moe_gib=moe_gib,
        peak_gib=peak_gib,
    )


@dataclass(frozen=True)
class MemoryPoint:
    dp: int
    sp: int
    tp: int
    ep: int
    moe_dp: int
    attn_tp: int
    moe_tp: int

    weights_per_rank_gib: float
    kv_per_rank_gib: float
    index_per_rank_gib: float
    act_per_rank_gib: float
    misc_per_rank_gib: float

    total_per_rank_gib: float
    budget_gib: float

    fits: bool
    intra_node_attn: bool

    # --- perf (optional; filled when perf spec provided) ---
    decode_ms: Optional[float]
    prefill_ms: Optional[float]
    decode_tok_s: Optional[float]
    prefill_tok_s: Optional[float]

    score: float


@dataclass(frozen=True)
class HardwareSpec:
    nodes: int
    gpus_per_node: int
    gpu_mem_gib: float
    mem_fraction: float = 0.90
    misc_gib: float = 6.0

    # --- performance modeling (optional) ---
    # Peak compute throughput for the main GEMMs (e.g., BF16 TensorCore), per GPU.
    gpu_tflops: Optional[float] = None
    # Peak HBM bandwidth, per GPU.
    gpu_hbm_gibps: Optional[float] = None
    # Effective collective bandwidth within a node (e.g., NVLink/xGMI), per GPU.
    intra_gibps: Optional[float] = None
    # Effective collective bandwidth across nodes (e.g., InfiniBand), per GPU.
    inter_gibps: Optional[float] = None
    # Latency per collective step (rough), in microseconds.
    intra_us: float = 2.0
    inter_us: float = 8.0
    # Efficiency scalars (0~1). The perf model is a lower bound; these make it less optimistic.
    flops_eff: float = 0.50
    hbm_eff: float = 0.80
    comm_eff: float = 0.85
    # Simple "small-M penalty" for decode GEMMs: tokens below this reduce effective TFLOPS.
    gemm_m_saturation: int = 256

    @property
    def world_size(self) -> int:
        return self.nodes * self.gpus_per_node

    @property
    def budget_gib(self) -> float:
        return self.gpu_mem_gib * self.mem_fraction


def _has_perf(hw: HardwareSpec) -> bool:
    need = (hw.gpu_tflops, hw.gpu_hbm_gibps, hw.intra_gibps, hw.inter_gibps)
    return all(x is not None and float(x) > 0 for x in need)


def _effective_tflops(hw: HardwareSpec, *, tokens_m: int) -> float:
    assert hw.gpu_tflops is not None
    base = float(hw.gpu_tflops) * 1e12 * float(max(1e-6, hw.flops_eff))
    sat = max(1, int(hw.gemm_m_saturation))
    scale = min(1.0, float(max(0, tokens_m)) / float(sat))
    return base * scale


def _effective_hbm_bw_Bps(hw: HardwareSpec) -> float:
    assert hw.gpu_hbm_gibps is not None
    return float(hw.gpu_hbm_gibps) * GiB * float(max(1e-6, hw.hbm_eff))


def _effective_link_bw_Bps(hw: HardwareSpec, *, inter_node: bool) -> float:
    bw = hw.inter_gibps if inter_node else hw.intra_gibps
    if bw is None:
        return 0.0
    return float(bw) * GiB * float(max(1e-6, hw.comm_eff))


def _link_latency_s(hw: HardwareSpec, *, inter_node: bool) -> float:
    us = hw.inter_us if inter_node else hw.intra_us
    return float(max(0.0, us)) * 1e-6


def _comm_crosses_nodes(hw: HardwareSpec, *, group_size: int, stride: int = 1) -> bool:
    """Best-effort: assume ranks are contiguous per node, and a group spans
    `(group_size-1)*stride+1` ranks in the global rank space.
    """
    if hw.nodes <= 1:
        return False
    if group_size <= 1:
        return False
    stride = max(1, int(stride))
    span = (int(group_size) - 1) * stride + 1
    return span > hw.gpus_per_node


def _collective_comm_bytes(kind: str, *, message_bytes: float, group_size: int) -> float:
    if group_size <= 1 or message_bytes <= 0:
        return 0.0
    frac = (float(group_size - 1)) / float(group_size)
    if kind in ("all_reduce",):
        return 2.0 * frac * message_bytes
    if kind in ("all_gather", "reduce_scatter"):
        return 1.0 * frac * message_bytes
    if kind in ("all_to_all", "a2a"):
        # A2A is backend-dependent; treat as "send + recv" with ring-like scaling.
        return 2.0 * frac * message_bytes
    return 2.0 * frac * message_bytes


def _collective_steps(kind: str, *, group_size: int) -> int:
    if group_size <= 1:
        return 0
    # Ring-like estimate.
    return int(group_size - 1)


@dataclass(frozen=True)
class OpTime:
    name: str
    flops: float
    hbm_bytes: float
    comm_bytes: float
    comm_kind: str
    comm_group: str
    comm_group_size: int
    comm_inter_node: bool
    comm_steps: int

    t_compute_s: float
    t_mem_s: float
    t_comm_s: float
    t_total_s: float
    bound: str


def _estimate_op_time(
    hw: HardwareSpec,
    *,
    name: str,
    tokens_m: int,
    flops: float,
    hbm_bytes: float,
    comm_kind: str = "",
    comm_group: str = "",
    comm_group_size: int = 1,
    comm_message_bytes: float = 0.0,
    comm_stride: int = 1,
) -> OpTime:
    if not _has_perf(hw):
        raise ValueError("Perf spec missing; cannot estimate op time")

    eff_tflops = _effective_tflops(hw, tokens_m=tokens_m)
    eff_hbm = _effective_hbm_bw_Bps(hw)

    t_compute = 0.0 if flops <= 0 else float(flops) / eff_tflops
    t_mem = 0.0 if hbm_bytes <= 0 else float(hbm_bytes) / eff_hbm

    comm_group_size = max(1, int(comm_group_size))
    comm_message_bytes = float(max(0.0, comm_message_bytes))
    comm_bytes = 0.0
    comm_steps = 0
    t_comm = 0.0
    comm_inter = False
    if comm_kind and comm_group_size > 1 and comm_message_bytes > 0:
        comm_bytes = _collective_comm_bytes(
            comm_kind, message_bytes=comm_message_bytes, group_size=comm_group_size
        )
        comm_steps = _collective_steps(comm_kind, group_size=comm_group_size)
        comm_inter = _comm_crosses_nodes(
            hw, group_size=comm_group_size, stride=int(comm_stride)
        )
        bw = _effective_link_bw_Bps(hw, inter_node=comm_inter)
        lat = _link_latency_s(hw, inter_node=comm_inter)
        t_comm = (comm_bytes / bw) + (comm_steps * lat) if bw > 0 else float("inf")

    t_total = max(t_compute, t_mem, t_comm)
    if t_total == t_comm:
        bound = "comm"
    elif t_total == t_mem:
        bound = "memory"
    else:
        bound = "compute"

    return OpTime(
        name=name,
        flops=float(flops),
        hbm_bytes=float(hbm_bytes),
        comm_bytes=float(comm_bytes),
        comm_kind=str(comm_kind),
        comm_group=str(comm_group),
        comm_group_size=int(comm_group_size),
        comm_inter_node=bool(comm_inter),
        comm_steps=int(comm_steps),
        t_compute_s=float(t_compute),
        t_mem_s=float(t_mem),
        t_comm_s=float(t_comm),
        t_total_s=float(t_total),
        bound=bound,
    )


@dataclass(frozen=True)
class StepTime:
    phase: str
    tokens_per_rank: int
    tokens_global: int
    total_s: float
    ops: Tuple[OpTime, ...]

    @property
    def ms(self) -> float:
        return self.total_s * 1e3

    @property
    def tok_s(self) -> float:
        return (float(self.tokens_global) / self.total_s) if self.total_s > 0 else 0.0


def _expected_unique_experts(
    *,
    total_experts: int,
    selections: float,
) -> float:
    """Expected number of unique experts touched when sampling `selections` times
    uniformly among `total_experts` buckets.
    """
    total_experts = max(0, int(total_experts))
    if total_experts <= 0:
        return 0.0
    if selections <= 0:
        return 0.0
    if total_experts == 1:
        return 1.0
    # E * (1 - (1-1/E)^k)
    return float(total_experts) * (1.0 - (1.0 - 1.0 / float(total_experts)) ** float(selections))


def estimate_step_time(
    d: ModelDims,
    parallel: Parallelism,
    *,
    hw: HardwareSpec,
    phase: str,
    tokens_per_rank: int,
    seq_len_local: int,
    global_batch: int,
    weight_dtype_bytes: float,
    act_dtype_bytes: float,
    kv_dtype_bytes: float,
    dp_padding_mode: str = "auto",
) -> StepTime:
    """Estimate time of one forward 'step' (either prefill chunk or decode 1 token).

    The result is a rough lower bound using a roofline model + simplified collectives.
    It is intended for *comparing configs*, not predicting exact latency.
    """
    if not _has_perf(hw):
        raise ValueError("Perf spec missing; cannot estimate step time")

    phase = str(phase)
    tokens_per_rank = max(0, int(tokens_per_rank))
    seq_len_local = max(1, int(seq_len_local))

    attn_tp = parallel.attn_tp
    moe_tp = parallel.moe_tp

    local_heads = int(math.ceil(d.num_heads / float(attn_tp)))
    kv_dim = d.kv_lora_rank + d.qk_rope_head_dim
    topk = int(d.index_topk) if d.index_topk > 0 else min(2048, seq_len_local)
    topk = min(topk, seq_len_local)

    # Global tokens processed by this step across the tp group.
    if phase == "prefill":
        tokens_global = max(0, tokens_per_rank * parallel.dp * parallel.sp)
    else:
        tokens_global = int(max(0, global_batch))

    ops: List[OpTime] = []

    # --- Attention (all layers) ---
    out_qkv_a = d.q_lora_rank + d.kv_lora_rank + d.qk_rope_head_dim
    ops.append(
        _estimate_op_time(
            hw,
            name="attn_fused_qkv_a",
            tokens_m=tokens_per_rank,
            flops=2.0 * tokens_per_rank * d.hidden_size * out_qkv_a,
            hbm_bytes=(
                tokens_per_rank * d.hidden_size * act_dtype_bytes
                + d.hidden_size * out_qkv_a * weight_dtype_bytes
                + tokens_per_rank * out_qkv_a * act_dtype_bytes
            ),
        )
    )

    # q_b: q_lora -> (local_heads * qk_head_dim)
    q_out = local_heads * d.qk_head_dim
    ops.append(
        _estimate_op_time(
            hw,
            name="attn_q_b",
            tokens_m=tokens_per_rank,
            flops=2.0 * tokens_per_rank * d.q_lora_rank * q_out,
            hbm_bytes=(
                tokens_per_rank * d.q_lora_rank * act_dtype_bytes
                + d.q_lora_rank * q_out * weight_dtype_bytes
                + tokens_per_rank * q_out * act_dtype_bytes
            ),
        )
    )

    # kv_b: kv_lora -> (local_heads * (qk_nope + v_head_dim))
    kv_out = local_heads * (d.qk_nope_head_dim + d.v_head_dim)
    ops.append(
        _estimate_op_time(
            hw,
            name="attn_kv_b",
            tokens_m=tokens_per_rank,
            flops=2.0 * tokens_per_rank * d.kv_lora_rank * kv_out,
            hbm_bytes=(
                tokens_per_rank * d.kv_lora_rank * act_dtype_bytes
                + d.kv_lora_rank * kv_out * weight_dtype_bytes
                + tokens_per_rank * kv_out * act_dtype_bytes
            ),
        )
    )

    # Attention compute (sparse topk)
    attn_flops = 2.0 * tokens_per_rank * local_heads * topk * (
        d.qk_head_dim + d.v_head_dim
    )
    attn_bytes = (
        tokens_per_rank * local_heads * d.qk_head_dim * act_dtype_bytes
        + tokens_per_rank * local_heads * topk * (d.qk_head_dim + d.v_head_dim) * kv_dtype_bytes
        + tokens_per_rank * local_heads * d.v_head_dim * act_dtype_bytes
    )
    ops.append(
        _estimate_op_time(
            hw,
            name=f"attn_topk_compute(k={topk})",
            tokens_m=tokens_per_rank,
            flops=attn_flops,
            hbm_bytes=attn_bytes,
        )
    )

    # o_proj: (local_heads*v) -> hidden
    o_in = local_heads * d.v_head_dim
    ops.append(
        _estimate_op_time(
            hw,
            name="attn_o_proj",
            tokens_m=tokens_per_rank,
            flops=2.0 * tokens_per_rank * o_in * d.hidden_size,
            hbm_bytes=(
                tokens_per_rank * o_in * act_dtype_bytes
                + o_in * d.hidden_size * weight_dtype_bytes
                + tokens_per_rank * d.hidden_size * act_dtype_bytes
            ),
        )
    )
    # row-parallel reduction for attention output
    ops.append(
        _estimate_op_time(
            hw,
            name="attn_tp_all_reduce(hidden)",
            tokens_m=tokens_per_rank,
            flops=0.0,
            hbm_bytes=0.0,
            comm_kind="all_reduce",
            comm_group="attn_tp",
            comm_group_size=attn_tp,
            comm_message_bytes=float(tokens_per_rank * d.hidden_size * act_dtype_bytes),
        )
    )

    # KV cache write (latent)
    ops.append(
        _estimate_op_time(
            hw,
            name="kv_cache_write(latent)",
            tokens_m=tokens_per_rank,
            flops=0.0,
            hbm_bytes=float(tokens_per_rank * kv_dim * kv_dtype_bytes),
        )
    )

    # --- NSA / Indexer (optional) ---
    if d.index_n_heads > 0 and d.index_head_dim > 0:
        idx_out = d.index_n_heads * d.index_head_dim
        ops.append(
            _estimate_op_time(
                hw,
                name="indexer_wq_b",
                tokens_m=tokens_per_rank,
                flops=2.0 * tokens_per_rank * d.q_lora_rank * idx_out,
                hbm_bytes=(
                    tokens_per_rank * d.q_lora_rank * act_dtype_bytes
                    + d.q_lora_rank * idx_out * weight_dtype_bytes
                    + tokens_per_rank * idx_out * act_dtype_bytes
                ),
            )
        )
        ops.append(
            _estimate_op_time(
                hw,
                name="indexer_wk",
                tokens_m=tokens_per_rank,
                flops=2.0 * tokens_per_rank * d.hidden_size * d.index_head_dim,
                hbm_bytes=(
                    tokens_per_rank * d.hidden_size * act_dtype_bytes
                    + d.hidden_size * d.index_head_dim * weight_dtype_bytes
                    + tokens_per_rank * d.index_head_dim * act_dtype_bytes
                ),
            )
        )

    # --- MLP: dense layers ---
    if d.dense_layers > 0:
        local_inter = int(math.ceil(d.intermediate_size / float(max(1, parallel.tp))))
        # gate + up
        ops.append(
            _estimate_op_time(
                hw,
                name="dense_ffn_gate_up",
                tokens_m=tokens_per_rank,
                flops=4.0 * tokens_per_rank * d.hidden_size * local_inter,
                hbm_bytes=(
                    # inputs read twice
                    2.0 * tokens_per_rank * d.hidden_size * act_dtype_bytes
                    # weights: 2 matrices
                    + 2.0 * d.hidden_size * local_inter * weight_dtype_bytes
                    # outputs: 2 intermediates
                    + 2.0 * tokens_per_rank * local_inter * act_dtype_bytes
                ),
            )
        )
        # down
        ops.append(
            _estimate_op_time(
                hw,
                name="dense_ffn_down",
                tokens_m=tokens_per_rank,
                flops=2.0 * tokens_per_rank * local_inter * d.hidden_size,
                hbm_bytes=(
                    tokens_per_rank * local_inter * act_dtype_bytes
                    + local_inter * d.hidden_size * weight_dtype_bytes
                    + tokens_per_rank * d.hidden_size * act_dtype_bytes
                ),
            )
        )
        ops.append(
            _estimate_op_time(
                hw,
                name="tp_all_reduce(hidden) (dense_ffn)",
                tokens_m=tokens_per_rank,
                flops=0.0,
                hbm_bytes=0.0,
                comm_kind="all_reduce",
                comm_group="tp",
                comm_group_size=parallel.tp,
                comm_message_bytes=float(tokens_per_rank * d.hidden_size * act_dtype_bytes),
            )
        )

    # --- MoE layers ---
    if d.moe_layers > 0:
        local_moe_inter = int(
            math.ceil(d.moe_intermediate_size / float(max(1, moe_tp)))
        )
        # MoE dispatch / combine (only routed experts; shared expert stays local)
        if parallel.ep > 1 and d.num_experts_per_tok > 0:
            message = float(tokens_per_rank * d.hidden_size * act_dtype_bytes * d.num_experts_per_tok)
            ops.append(
                _estimate_op_time(
                    hw,
                    name="moe_dispatch_a2a",
                    tokens_m=tokens_per_rank,
                    flops=0.0,
                    hbm_bytes=0.0,
                    comm_kind="all_to_all",
                    comm_group="moe_ep",
                    comm_group_size=parallel.ep,
                    comm_message_bytes=message,
                    comm_stride=max(1, moe_tp),
                )
            )
            ops.append(
                _estimate_op_time(
                    hw,
                    name="moe_combine_a2a",
                    tokens_m=tokens_per_rank,
                    flops=0.0,
                    hbm_bytes=0.0,
                    comm_kind="all_to_all",
                    comm_group="moe_ep",
                    comm_group_size=parallel.ep,
                    comm_message_bytes=message,
                    comm_stride=max(1, moe_tp),
                )
            )

        # Expert compute (expected balanced routing)
        routed_per_tok = (
            float(d.num_experts_per_tok) / float(max(1, parallel.ep))
            if d.num_experts_per_tok > 0
            else 0.0
        )
        experts_per_tok_local = float(max(0, d.n_shared_experts)) + routed_per_tok

        moe_flops = (
            float(tokens_per_rank)
            * experts_per_tok_local
            * (6.0 * float(d.hidden_size) * float(local_moe_inter))
        )

        # Memory bytes: estimate how many unique routed experts are touched per step on this rank.
        per_expert_params = 3 * d.hidden_size * d.moe_intermediate_size
        per_expert_params_rank = float(per_expert_params) / float(max(1, moe_tp))

        selections_per_rank = float(tokens_per_rank) * float(d.num_experts_per_tok) / float(
            max(1, parallel.ep)
        )
        local_experts = int(math.ceil(float(d.n_routed_experts) / float(max(1, parallel.ep))))
        unique_routed = _expected_unique_experts(
            total_experts=local_experts, selections=selections_per_rank
        )
        moe_weight_bytes = (
            unique_routed * per_expert_params_rank
            + float(max(0, d.n_shared_experts)) * per_expert_params_rank
        ) * float(weight_dtype_bytes)

        moe_act_bytes = (
            # gate+up read inputs twice
            2.0 * selections_per_rank * d.hidden_size * act_dtype_bytes
            # gate+up outputs
            + 2.0 * selections_per_rank * local_moe_inter * act_dtype_bytes
            # down input + output (approx)
            + selections_per_rank * local_moe_inter * act_dtype_bytes
            + float(tokens_per_rank) * d.hidden_size * act_dtype_bytes
        )

        ops.append(
            _estimate_op_time(
                hw,
                name="moe_expert_mlp",
                tokens_m=tokens_per_rank,
                flops=moe_flops,
                hbm_bytes=float(moe_weight_bytes + moe_act_bytes),
            )
        )
        ops.append(
            _estimate_op_time(
                hw,
                name="moe_tp_all_reduce(hidden)",
                tokens_m=tokens_per_rank,
                flops=0.0,
                hbm_bytes=0.0,
                comm_kind="all_reduce",
                comm_group="moe_tp",
                comm_group_size=moe_tp,
                comm_message_bytes=float(tokens_per_rank * d.hidden_size * act_dtype_bytes),
            )
        )

    # --- DP-attn gather (rough) ---
    if parallel.dp > 1:
        # A simplified view of dp-attn gather:
        # - MAX_LEN: all_gather (+ optional reduce_scatter on attn_tp when attn_tp>1)
        # - SUM_LEN: all_reduce
        mode = str(dp_padding_mode)
        if mode == "auto":
            mode = "sum_len" if phase == "prefill" else "max_len"

        if mode == "max_len":
            ops.append(
                _estimate_op_time(
                    hw,
                    name="dp_attn_all_gather(hidden)",
                    tokens_m=tokens_per_rank,
                    flops=0.0,
                    hbm_bytes=0.0,
                    comm_kind="all_gather",
                    comm_group="attn_dp",
                    comm_group_size=parallel.dp,
                    comm_message_bytes=float(tokens_per_rank * parallel.dp * d.hidden_size * act_dtype_bytes),
                )
            )
            if attn_tp > 1:
                ops.append(
                    _estimate_op_time(
                        hw,
                        name="dp_attn_attn_tp_reduce_scatter(hidden)",
                        tokens_m=tokens_per_rank,
                        flops=0.0,
                        hbm_bytes=0.0,
                        comm_kind="reduce_scatter",
                        comm_group="attn_tp",
                        comm_group_size=attn_tp,
                        comm_message_bytes=float(tokens_per_rank * d.hidden_size * act_dtype_bytes),
                    )
                )
        else:
            ops.append(
                _estimate_op_time(
                    hw,
                    name="dp_attn_all_reduce(hidden)",
                    tokens_m=tokens_per_rank,
                    flops=0.0,
                    hbm_bytes=0.0,
                    comm_kind="all_reduce",
                    comm_group="attn_dp",
                    comm_group_size=parallel.dp,
                    comm_message_bytes=float(tokens_per_rank * parallel.dp * d.hidden_size * act_dtype_bytes),
                )
            )

    # --- Prefill CP (sp) gather/scatter (rough) ---
    if phase == "prefill" and parallel.sp > 1:
        # Modeled as one reduce_scatter + one all_gather on hidden states.
        msg = float(tokens_per_rank * parallel.sp * d.hidden_size * act_dtype_bytes)
        ops.append(
            _estimate_op_time(
                hw,
                name="prefill_cp_reduce_scatter(hidden)",
                tokens_m=tokens_per_rank,
                flops=0.0,
                hbm_bytes=0.0,
                comm_kind="reduce_scatter",
                comm_group="attn_cp",
                comm_group_size=parallel.sp,
                comm_message_bytes=msg,
                comm_stride=max(1, attn_tp),
            )
        )
        ops.append(
            _estimate_op_time(
                hw,
                name="prefill_cp_all_gather(hidden)",
                tokens_m=tokens_per_rank,
                flops=0.0,
                hbm_bytes=0.0,
                comm_kind="all_gather",
                comm_group="attn_cp",
                comm_group_size=parallel.sp,
                comm_message_bytes=msg,
                comm_stride=max(1, attn_tp),
            )
        )

    # --- aggregate over layers ---
    # Most ops above are per-layer; now scale by layer counts.
    scaled_ops: List[OpTime] = []
    for op in ops:
        if op.name.startswith("dense_ffn") or "(dense_ffn)" in op.name or op.name.startswith("tp_all_reduce") and "dense_ffn" in op.name:
            mult = d.dense_layers
        elif op.name.startswith("moe_") or op.comm_group in ("moe_ep", "moe_tp"):
            mult = d.moe_layers
        else:
            mult = d.num_layers
        if mult <= 1:
            scaled_ops.append(op)
            continue
        scaled_ops.append(
            OpTime(
                **{
                    **op.__dict__,
                    "flops": op.flops * mult,
                    "hbm_bytes": op.hbm_bytes * mult,
                    "comm_bytes": op.comm_bytes * mult,
                    "t_compute_s": op.t_compute_s * mult,
                    "t_mem_s": op.t_mem_s * mult,
                    "t_comm_s": op.t_comm_s * mult,
                    "t_total_s": op.t_total_s * mult,
                }
            )
        )

    total_s = sum(op.t_total_s for op in scaled_ops)
    return StepTime(
        phase=phase,
        tokens_per_rank=tokens_per_rank,
        tokens_global=int(tokens_global),
        total_s=float(total_s),
        ops=tuple(scaled_ops),
    )


def score_config(point: MemoryPoint, hw: HardwareSpec, *, preference: str) -> float:
    if not point.fits:
        return -1e9 - point.total_per_rank_gib

    headroom = point.budget_gib - point.total_per_rank_gib
    kv_cost = point.kv_per_rank_gib + point.index_per_rank_gib
    cross_node_penalty = 0.0 if point.intra_node_attn else 1.0
    sp_penalty = 0.05 * float(max(0, point.sp - 1))
    dp_penalty = 0.01 * float(max(0, point.dp - 1))

    # If perf metrics exist, use them as the primary signal.
    if point.decode_ms is not None and point.decode_tok_s is not None:
        decode_ms = float(point.decode_ms)
        decode_tok_s = float(point.decode_tok_s)
        prefill_ms = float(point.prefill_ms or 0.0)

        if preference == "latency":
            return (
                -decode_ms
                - 0.1 * prefill_ms
                + 0.05 * headroom
                - 50.0 * cross_node_penalty
                - 10.0 * sp_penalty
                - 2.0 * dp_penalty
            )
        if preference == "throughput":
            return (
                +decode_tok_s
                - 10.0 * kv_cost
                + 0.02 * headroom
                - 5.0 * cross_node_penalty
                - 10.0 * sp_penalty
            )
        # balanced
        return (
            +0.5 * decode_tok_s
            - 0.2 * decode_ms
            - 5.0 * kv_cost
            + 0.02 * headroom
            - 20.0 * cross_node_penalty
            - 10.0 * sp_penalty
            - 1.0 * dp_penalty
        )

    if preference == "latency":
        return (
            headroom
            - 100.0 * cross_node_penalty
            - 10.0 * sp_penalty
            - 2.0 * dp_penalty
        )
    if preference == "throughput":
        return (
            -10.0 * kv_cost
            + 1.0 * headroom
            - 10.0 * sp_penalty
            - 1.0 * cross_node_penalty
        )
    return (
        0.5 * headroom
        - 5.0 * kv_cost
        - 30.0 * cross_node_penalty
        - 10.0 * sp_penalty
        - 1.0 * dp_penalty
    )


def _choose_min_divisor_at_least(divs: Sequence[int], target: int) -> Optional[int]:
    for d in divs:
        if d >= target:
            return d
    return None


def recommend_parallelism(
    d: ModelDims,
    p: ParamBreakdown,
    *,
    hw: HardwareSpec,
    seq_len: int,
    global_batch: int,
    weight_dtype_bytes: float,
    kv_dtype_bytes: float,
    act_dtype_bytes: float,
    chunked_prefill_size: int,
    mhc_streams: int = 1,
    mhc_branch_streams: int = 1,
    ep: int,
    moe_dp: int,
    tp: Optional[int] = None,
    sp_candidates: Sequence[int] = (1,),
    preference: str = "balanced",
    dp_padding_mode: str = "auto",
) -> Tuple[List[MemoryPoint], Dict[str, Optional[int]]]:
    if tp is None:
        tp = hw.world_size

    mhc_streams = max(1, int(mhc_streams))
    mhc_branch_streams = max(1, int(mhc_branch_streams))

    tp_divs = divisors(tp)
    points: List[MemoryPoint] = []

    for sp in sorted(set(int(x) for x in sp_candidates if int(x) > 0)):
        for dp in tp_divs:
            if dp <= 0:
                continue
            if dp * sp > tp:
                continue
            if tp % (dp * sp) != 0:
                continue

            parallel = Parallelism(tp=tp, dp=dp, sp=sp, ep=ep, moe_dp=moe_dp)
            sharded = estimate_weight_params_per_rank(d, p, parallel)
            weights_per_rank_gib = (sharded.total * weight_dtype_bytes) / GiB

            local_batch = float(global_batch) / float(dp)
            local_seq = int(math.ceil(seq_len / float(sp)))

            kv_batch = local_batch * float(max(1, mhc_branch_streams))
            kv_per_rank_gib = estimate_kv_cache_gib(
                d, batch=kv_batch, seq_len=local_seq, kv_dtype_bytes=kv_dtype_bytes
            )
            idx_per_rank_gib = (
                estimate_nsa_index_cache_gib(d, batch=kv_batch, seq_len=local_seq)
                or 0.0
            )

            local_chunk = int(max(1, chunked_prefill_size // dp))
            local_prefill_tokens = int(
                min(local_chunk, math.ceil(local_batch * local_seq))
            )
            act_ws = estimate_activation_working_set(
                d,
                parallel,
                tokens=local_prefill_tokens,
                act_dtype_bytes=act_dtype_bytes,
                mhc_streams=mhc_streams,
                mhc_branch_streams=mhc_branch_streams,
            )
            act_per_rank_gib = act_ws.peak_gib

            decode_ms: Optional[float] = None
            prefill_ms: Optional[float] = None
            decode_tok_s: Optional[float] = None
            prefill_tok_s: Optional[float] = None
            if _has_perf(hw):
                local_decode_tokens = int(max(1, math.ceil(local_batch)))
                try:
                    decode = estimate_step_time(
                        d,
                        parallel,
                        hw=hw,
                        phase="decode",
                        tokens_per_rank=local_decode_tokens,
                        seq_len_local=local_seq,
                        global_batch=int(global_batch),
                        weight_dtype_bytes=weight_dtype_bytes,
                        act_dtype_bytes=act_dtype_bytes,
                        kv_dtype_bytes=kv_dtype_bytes,
                        dp_padding_mode=str(dp_padding_mode),
                    )
                    prefill = estimate_step_time(
                        d,
                        parallel,
                        hw=hw,
                        phase="prefill",
                        tokens_per_rank=local_prefill_tokens,
                        seq_len_local=local_seq,
                        global_batch=int(global_batch),
                        weight_dtype_bytes=weight_dtype_bytes,
                        act_dtype_bytes=act_dtype_bytes,
                        kv_dtype_bytes=kv_dtype_bytes,
                        dp_padding_mode=str(dp_padding_mode),
                    )
                    decode_ms = decode.ms
                    prefill_ms = prefill.ms
                    decode_tok_s = float(global_batch) / decode.total_s if decode.total_s > 0 else 0.0
                    global_prefill_tokens = int(local_prefill_tokens * parallel.dp * parallel.sp)
                    prefill_tok_s = (
                        float(global_prefill_tokens) / prefill.total_s if prefill.total_s > 0 else 0.0
                    )
                except Exception:
                    # Keep perf fields None; recommendation falls back to memory-only heuristic.
                    decode_ms = None

            misc = hw.misc_gib
            budget = hw.budget_gib
            total = (
                weights_per_rank_gib
                + kv_per_rank_gib
                + idx_per_rank_gib
                + act_per_rank_gib
                + misc
            )
            fits = total <= budget
            intra_node_attn = parallel.attn_tp <= hw.gpus_per_node and (
                hw.gpus_per_node % max(1, parallel.attn_tp) == 0
            )

            point = MemoryPoint(
                dp=dp,
                sp=sp,
                tp=tp,
                ep=ep,
                moe_dp=moe_dp,
                attn_tp=parallel.attn_tp,
                moe_tp=parallel.moe_tp,
                weights_per_rank_gib=weights_per_rank_gib,
                kv_per_rank_gib=kv_per_rank_gib,
                index_per_rank_gib=idx_per_rank_gib,
                act_per_rank_gib=act_per_rank_gib,
                misc_per_rank_gib=misc,
                total_per_rank_gib=total,
                budget_gib=budget,
                fits=fits,
                intra_node_attn=intra_node_attn,
                decode_ms=decode_ms,
                prefill_ms=prefill_ms,
                decode_tok_s=decode_tok_s,
                prefill_tok_s=prefill_tok_s,
                score=0.0,  # filled below
            )
            point = MemoryPoint(
                **{**point.__dict__, "score": score_config(point, hw, preference=preference)}
            )
            points.append(point)

    points.sort(key=lambda x: x.score, reverse=True)

    # --- critical points (sp=1 baseline when possible) ---
    critical: Dict[str, Optional[int]] = {"dp_min_fit": None, "dp_min_intra_node": None}
    sp0 = 1
    if sp0 in sp_candidates:
        dp_divs = divisors(tp)
        target = int(math.ceil(tp / float(max(1, hw.gpus_per_node * sp0))))
        critical["dp_min_intra_node"] = _choose_min_divisor_at_least(dp_divs, target)

        dp_fit: Optional[int] = None
        for dp in dp_divs:
            parallel = Parallelism(tp=tp, dp=dp, sp=sp0, ep=ep, moe_dp=moe_dp)
            sharded = estimate_weight_params_per_rank(d, p, parallel)
            weights_per_rank_gib = (sharded.total * weight_dtype_bytes) / GiB
            local_batch = float(global_batch) / float(dp)
            local_seq = int(math.ceil(seq_len / float(sp0)))
            kv_batch = local_batch * float(max(1, mhc_branch_streams))
            kv_per_rank_gib = estimate_kv_cache_gib(
                d, batch=kv_batch, seq_len=local_seq, kv_dtype_bytes=kv_dtype_bytes
            )
            idx_per_rank_gib = (
                estimate_nsa_index_cache_gib(d, batch=kv_batch, seq_len=local_seq)
                or 0.0
            )
            local_chunk = int(max(1, chunked_prefill_size // dp))
            local_prefill_tokens = int(
                min(local_chunk, math.ceil(local_batch * local_seq))
            )
            act_ws = estimate_activation_working_set(
                d,
                parallel,
                tokens=local_prefill_tokens,
                act_dtype_bytes=act_dtype_bytes,
                mhc_streams=mhc_streams,
                mhc_branch_streams=mhc_branch_streams,
            )
            total = (
                weights_per_rank_gib
                + kv_per_rank_gib
                + idx_per_rank_gib
                + act_ws.peak_gib
                + hw.misc_gib
            )
            if total <= hw.budget_gib:
                dp_fit = dp
                break
        critical["dp_min_fit"] = dp_fit

    return points, critical


def _svg_escape(text: str) -> str:
    return (
        text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
        .replace("'", "&#39;")
    )


def write_memory_scan_svg(
    path: Path,
    *,
    title: str,
    x: Sequence[int],
    series: Dict[str, Sequence[float]],
    hlines: Dict[str, float],
    vlines: Dict[str, Optional[int]],
    y_label: str = "GiB per rank",
) -> None:
    width, height = 960, 540
    margin = dict(l=70, r=30, t=60, b=60)

    if not x:
        path.write_text(
            "<svg xmlns='http://www.w3.org/2000/svg' width='400' height='80'></svg>"
        )
        return

    x_min, x_max = min(x), max(x)
    y_values: List[float] = []
    for ys in series.values():
        y_values.extend(list(ys))
    y_values.extend(list(hlines.values()))
    y_min = 0.0
    y_max = max(y_values) if y_values else 1.0
    if y_max <= 0:
        y_max = 1.0

    def map_x(xv: float) -> float:
        if x_max == x_min:
            return margin["l"] + (width - margin["l"] - margin["r"]) / 2.0
        return margin["l"] + (xv - x_min) * (width - margin["l"] - margin["r"]) / (
            x_max - x_min
        )

    def map_y(yv: float) -> float:
        return (
            height
            - margin["b"]
            - (yv - y_min) * (height - margin["t"] - margin["b"]) / (y_max - y_min)
        )

    colors = [
        "#1f77b4",
        "#ff7f0e",
        "#2ca02c",
        "#d62728",
        "#9467bd",
        "#8c564b",
    ]

    lines: List[str] = []
    lines.append(
        f"<svg xmlns='http://www.w3.org/2000/svg' width='{width}' height='{height}'>"
    )
    lines.append("<rect width='100%' height='100%' fill='white'/>")
    lines.append(
        f"<text x='{width/2:.1f}' y='{margin['t']/2:.1f}' text-anchor='middle' font-size='18' font-family='sans-serif'>{_svg_escape(title)}</text>"
    )

    x0 = margin["l"]
    y0 = height - margin["b"]
    x1 = width - margin["r"]
    y1 = margin["t"]
    lines.append(f"<line x1='{x0}' y1='{y0}' x2='{x1}' y2='{y0}' stroke='black'/>")
    lines.append(f"<line x1='{x0}' y1='{y0}' x2='{x0}' y2='{y1}' stroke='black'/>")

    for i in range(6):
        yv = y_min + (y_max - y_min) * i / 5.0
        py = map_y(yv)
        lines.append(
            f"<line x1='{x0}' y1='{py:.1f}' x2='{x1}' y2='{py:.1f}' stroke='#e5e5e5'/>"
        )
        lines.append(
            f"<text x='{x0-8}' y='{py+4:.1f}' text-anchor='end' font-size='12' font-family='sans-serif'>{yv:.1f}</text>"
        )

    max_ticks = 12
    if len(x) <= max_ticks:
        xticks = x
    else:
        step = max(1, len(x) // max_ticks)
        xticks = [x[i] for i in range(0, len(x), step)]
        if xticks[-1] != x[-1]:
            xticks.append(x[-1])

    for xv in xticks:
        px = map_x(xv)
        lines.append(f"<line x1='{px:.1f}' y1='{y0}' x2='{px:.1f}' y2='{y0+6}' stroke='black'/>")
        lines.append(
            f"<text x='{px:.1f}' y='{y0+22}' text-anchor='middle' font-size='12' font-family='sans-serif'>{xv}</text>"
        )

    lines.append(
        f"<text x='{(x0+x1)/2:.1f}' y='{height-20}' text-anchor='middle' font-size='12' font-family='sans-serif'>dp</text>"
    )
    lines.append(
        f"<text x='18' y='{(y0+y1)/2:.1f}' text-anchor='middle' font-size='12' font-family='sans-serif' transform='rotate(-90 18 {(y0+y1)/2:.1f})'>{_svg_escape(y_label)}</text>"
    )

    legend_y = margin["t"] + 10
    legend_x = x1 - 220
    idx = 0
    for name, ys in series.items():
        color = colors[idx % len(colors)]
        idx += 1
        pts = []
        for xv, yv in zip(x, ys):
            pts.append(f"{map_x(xv):.1f},{map_y(yv):.1f}")
        lines.append(
            f"<polyline fill='none' stroke='{color}' stroke-width='2' points='{' '.join(pts)}'/>"
        )
        lines.append(
            f"<rect x='{legend_x}' y='{legend_y-10}' width='12' height='12' fill='{color}'/>"
        )
        lines.append(
            f"<text x='{legend_x+18}' y='{legend_y}' font-size='12' font-family='sans-serif'>{_svg_escape(name)}</text>"
        )
        legend_y += 18

    for name, yv in hlines.items():
        py = map_y(yv)
        lines.append(
            f"<line x1='{x0}' y1='{py:.1f}' x2='{x1}' y2='{py:.1f}' stroke='#444' stroke-dasharray='6,4'/>"
        )
        lines.append(
            f"<text x='{x0+6}' y='{py-6:.1f}' font-size='12' font-family='sans-serif' fill='#444'>{_svg_escape(name)}={yv:.1f} GiB</text>"
        )

    for name, xv in vlines.items():
        if xv is None:
            continue
        px = map_x(xv)
        lines.append(
            f"<line x1='{px:.1f}' y1='{y0}' x2='{px:.1f}' y2='{y1}' stroke='#444' stroke-dasharray='4,4'/>"
        )
        lines.append(
            f"<text x='{px+4:.1f}' y='{y1+14}' font-size='12' font-family='sans-serif' fill='#444'>{_svg_escape(name)}={xv}</text>"
        )

    lines.append("</svg>")
    path.write_text("\n".join(lines))


def write_memory_breakdown_svg(
    path: Path,
    *,
    title: str,
    bars: Sequence[Tuple[str, Dict[str, float]]],
    budget_gib: Optional[float] = None,
    x_label: str = "GiB per rank",
) -> None:
    """Write a simple stacked horizontal bar chart for memory breakdown.

    `bars`: [(label, {"weights":..,"kv":..,"index":..,"act":..,"misc":..}), ...]
    """
    width, height = 960, 320
    margin = dict(l=160, r=40, t=70, b=60)
    bar_h = 26
    bar_gap = 16

    if not bars:
        path.write_text(
            "<svg xmlns='http://www.w3.org/2000/svg' width='400' height='80'></svg>"
        )
        return

    parts = ("weights", "kv", "index", "act", "misc")
    colors = {
        "weights": "#1f77b4",
        "kv": "#ff7f0e",
        "index": "#2ca02c",
        "act": "#d62728",
        "misc": "#9467bd",
    }

    totals = []
    for _, m in bars:
        total = 0.0
        for k in parts:
            total += float(m.get(k, 0.0))
        totals.append(total)

    x_max = max(totals)
    if budget_gib is not None:
        x_max = max(x_max, float(budget_gib))
    if x_max <= 0:
        x_max = 1.0
    x_max *= 1.10  # headroom for labels

    plot_w = width - margin["l"] - margin["r"]
    plot_h = height - margin["t"] - margin["b"]

    needed_h = len(bars) * bar_h + (len(bars) - 1) * bar_gap
    if needed_h > plot_h:
        height = margin["t"] + needed_h + margin["b"]
        plot_h = height - margin["t"] - margin["b"]

    def map_x(xv: float) -> float:
        return margin["l"] + (float(xv) / x_max) * plot_w

    lines: List[str] = []
    lines.append(
        f"<svg xmlns='http://www.w3.org/2000/svg' width='{width}' height='{height}'>"
    )
    lines.append("<rect width='100%' height='100%' fill='white'/>")
    lines.append(
        f"<text x='{width/2:.1f}' y='{margin['t']/2:.1f}' text-anchor='middle' font-size='18' font-family='sans-serif'>{_svg_escape(title)}</text>"
    )

    x0 = margin["l"]
    y0 = height - margin["b"]
    x1 = width - margin["r"]
    y1 = margin["t"]

    # axis
    lines.append(f"<line x1='{x0}' y1='{y0}' x2='{x1}' y2='{y0}' stroke='black'/>")
    lines.append(f"<line x1='{x0}' y1='{y0}' x2='{x0}' y2='{y1}' stroke='black'/>")

    # x ticks
    ticks = 6
    for i in range(ticks + 1):
        xv = x_max * i / ticks
        px = map_x(xv)
        lines.append(
            f"<line x1='{px:.1f}' y1='{y0}' x2='{px:.1f}' y2='{y0+6}' stroke='black'/>"
        )
        lines.append(
            f"<text x='{px:.1f}' y='{y0+22}' text-anchor='middle' font-size='12' font-family='sans-serif'>{xv:.0f}</text>"
        )

    lines.append(
        f"<text x='{(x0+x1)/2:.1f}' y='{height-20}' text-anchor='middle' font-size='12' font-family='sans-serif'>{_svg_escape(x_label)}</text>"
    )

    if budget_gib is not None:
        px = map_x(float(budget_gib))
        lines.append(
            f"<line x1='{px:.1f}' y1='{y0}' x2='{px:.1f}' y2='{y1}' stroke='#444' stroke-dasharray='6,4'/>"
        )
        lines.append(
            f"<text x='{px+4:.1f}' y='{y1+14}' font-size='12' font-family='sans-serif' fill='#444'>budget={float(budget_gib):.1f}</text>"
        )

    # bars
    y = y1
    for (label, m), total in zip(bars, totals):
        y_bar = y
        # label
        lines.append(
            f"<text x='{x0-10:.1f}' y='{y_bar + bar_h*0.70:.1f}' text-anchor='end' font-size='12' font-family='sans-serif'>{_svg_escape(label)}</text>"
        )
        # stacked rects
        cur = 0.0
        for k in parts:
            v = float(m.get(k, 0.0))
            if v <= 0:
                continue
            x_left = map_x(cur)
            x_right = map_x(cur + v)
            w = max(0.0, x_right - x_left)
            lines.append(
                f"<rect x='{x_left:.1f}' y='{y_bar:.1f}' width='{w:.1f}' height='{bar_h:.1f}' fill='{colors[k]}'/>"
            )
            cur += v

        # outline + total label
        total_w = map_x(total) - map_x(0.0)
        outline = "#d62728" if budget_gib is not None and total > float(budget_gib) else "#222"
        lines.append(
            f"<rect x='{x0:.1f}' y='{y_bar:.1f}' width='{total_w:.1f}' height='{bar_h:.1f}' fill='none' stroke='{outline}' stroke-width='1'/>"
        )
        lines.append(
            f"<text x='{map_x(total)+6:.1f}' y='{y_bar + bar_h*0.70:.1f}' text-anchor='start' font-size='12' font-family='sans-serif'>{total:.1f}</text>"
        )
        y += bar_h + bar_gap

    # legend
    legend_y = height - margin["b"] + 18
    legend_x = margin["l"]
    for k in parts:
        lines.append(
            f"<rect x='{legend_x:.1f}' y='{legend_y-10:.1f}' width='12' height='12' fill='{colors[k]}'/>"
        )
        lines.append(
            f"<text x='{legend_x+18:.1f}' y='{legend_y:.1f}' font-size='12' font-family='sans-serif'>{_svg_escape(k)}</text>"
        )
        legend_x += 110

    lines.append("</svg>")
    path.write_text("\n".join(lines))


@dataclass(frozen=True)
class SingleSeqEstimate:
    seq_len: int
    local_seq_len: int
    global_batch: float
    local_batch: float

    kv_unique_gib: float
    kv_per_rank_gib: float
    idx_unique_gib: Optional[float]
    idx_per_rank_gib: Optional[float]

    prefill_tokens_per_rank: int
    decode_tokens_per_rank: int
    prefill_ws: ActivationWorkingSet
    decode_ws: ActivationWorkingSet

    decode_step: Optional[StepTime]
    prefill_step: Optional[StepTime]


def render_single_config_report_md(
    *,
    args: argparse.Namespace,
    config_src: str,
    ts: str,
    out_md: Path,
    d: ModelDims,
    p: ParamBreakdown,
    active: Dict[str, int],
    mhc: MhcSpec,
    parallel: Parallelism,
    weight_dtype: str,
    kv_dtype: str,
    act_dtype: str,
    weight_dtype_bytes: float,
    kv_dtype_bytes: float,
    act_dtype_bytes: float,
    sharded_weights: ShardedWeightBreakdown,
    weights_unique_gib: float,
    weights_per_rank_gib: float,
    budget_gib: Optional[float],
    misc_gib: float,
    seq_estimates: Sequence[SingleSeqEstimate],
    out_svg: Optional[Path],
) -> str:
    lines: List[str] = []
    if args.obsidian:
        model_title = _sanitize_for_title(config_src.replace("path:", ""))
        lines.append(
            _render_obsidian_frontmatter(
                title=f"LLM Memory Estimate - {model_title}",
                created=ts,
                tags=("llm", "vram", "kv-cache", "activation", "parallelism"),
                model=config_src,
            )
        )

    lines.append("# LLM Memory / Perf Estimate Report\n\n")
    lines.append(f"Generated: {ts}\n\n")

    lines.append("## Inputs\n")
    lines.append(f"- config: `{config_src}`\n")
    lines.append(f"- weight dtype: `{weight_dtype}` ({weight_dtype_bytes} B)\n")
    lines.append(f"- kv dtype: `{kv_dtype}` ({kv_dtype_bytes} B)\n")
    lines.append(f"- act dtype: `{act_dtype}` ({act_dtype_bytes} B)\n")
    lines.append(f"- seq_len: `{args.seq_len}`\n")
    lines.append(f"- batch(scope={args.batch_scope}): `{args.batch}`\n")
    lines.append(f"- chunked_prefill_size: `{args.chunked_prefill_size}`\n")
    lines.append(
        f"- parallelism: `tp={parallel.tp}`, `dp={parallel.dp}`, `sp={parallel.sp}`, `ep={parallel.ep}`, `moe_dp={parallel.moe_dp}`\n"
    )
    lines.append(
        f"- derived: `attn_tp={parallel.attn_tp}`, `moe_tp={parallel.moe_tp}`\n"
    )
    lines.append(
        f"- mHC: `streams={mhc.streams}`, `branches_per_layer={mhc.branches_per_layer}`, `branch_mode={mhc.branch_mode}`, `branch_streams={mhc.branch_streams}`\n"
    )
    if budget_gib is not None:
        lines.append("\n## Hardware Budget (per GPU)\n")
        lines.append(f"- gpu_mem: `{args.gpu_mem_gib} GiB`\n")
        lines.append(f"- mem_fraction: `{args.mem_fraction}` ⇒ budget `{budget_gib:.2f} GiB`\n")
        lines.append(f"- misc reserved: `{misc_gib:.2f} GiB`\n")

    if out_svg is not None and out_svg.exists():
        lines.append("\n## Plot\n")
        if args.obsidian:
            rel_svg = _relative_path_for_embed(out_svg, out_md.parent)
            lines.append(f"- breakdown: ![[{rel_svg}]]\n")
        else:
            lines.append(f"- breakdown svg: `{out_svg}`\n")

    lines.append("\n## Model Summary (from config)\n")
    lines.append(f"- hidden_size: `{d.hidden_size}`\n")
    lines.append(f"- layers: `{d.num_layers}`\n")
    lines.append(f"- heads: `{d.num_heads}`\n")
    lines.append(f"- vocab_size: `{d.vocab_size}`\n")
    lines.append(f"- MoE layers: `{d.moe_layers}` (dense first `{d.dense_layers}`)\n")
    lines.append(
        "- MLA dims: "
        f"`q_lora_rank={d.q_lora_rank}`, `kv_lora_rank={d.kv_lora_rank}`, "
        f"`qk_nope={d.qk_nope_head_dim}`, `qk_rope={d.qk_rope_head_dim}`, `v={d.v_head_dim}`\n"
    )
    if d.index_n_heads and d.index_head_dim:
        lines.append(
            f"- NSA indexer: `index_n_heads={d.index_n_heads}`, `index_head_dim={d.index_head_dim}`, `index_topk={d.index_topk}`\n"
        )

    lines.append("\n## Activated Params / Token (compute metric)\n")
    lines.append(
        f"- total_excluding_embed_router: `{format_count(active['total_excluding_embed_router'])}` params\n"
    )
    lines.append(
        f"- breakdown: moe `{format_count(active['active_moe'])}` + dense_ffn `{format_count(active['active_dense_ffn'])}` + always(attn+indexer) `{format_count(active['always_active(attn+indexer)'])}` + mhc `{format_count(active['active_mhc'])}`\n"
    )

    lines.append("\n## Weights (unique / per-rank)\n")
    lines.append(f"- weights_unique: `{weights_unique_gib:.3f} GiB`\n")
    lines.append(f"- weights_per_rank: `{weights_per_rank_gib:.3f} GiB`\n")

    def _wgib(x_params: float) -> float:
        return (x_params * weight_dtype_bytes) / GiB

    lines.append("\n### Weights Breakdown (per rank, approx)\n")
    lines.append("| part | params | GiB |\n|---|---:|---:|\n")
    lines.append(
        f"| embedding | {format_count(sharded_weights.embedding)} | {_wgib(sharded_weights.embedding):.3f} |\n"
    )
    lines.append(
        f"| lm_head | {format_count(sharded_weights.lm_head)} | {_wgib(sharded_weights.lm_head):.3f} |\n"
    )
    lines.append(
        f"| attn_fused_qkv_a | {format_count(sharded_weights.attn_fused_qkv_a)} | {_wgib(sharded_weights.attn_fused_qkv_a):.3f} |\n"
    )
    lines.append(
        f"| attn_q_b | {format_count(sharded_weights.attn_q_b)} | {_wgib(sharded_weights.attn_q_b):.3f} |\n"
    )
    lines.append(
        f"| attn_kv_b | {format_count(sharded_weights.attn_kv_b)} | {_wgib(sharded_weights.attn_kv_b):.3f} |\n"
    )
    lines.append(
        f"| attn_o_proj | {format_count(sharded_weights.attn_o_proj)} | {_wgib(sharded_weights.attn_o_proj):.3f} |\n"
    )
    lines.append(
        f"| indexer | {format_count(sharded_weights.indexer)} | {_wgib(sharded_weights.indexer):.3f} |\n"
    )
    lines.append(
        f"| dense_ffn | {format_count(sharded_weights.dense_ffn)} | {_wgib(sharded_weights.dense_ffn):.3f} |\n"
    )
    lines.append(
        f"| moe_routed | {format_count(sharded_weights.moe_routed_expert)} | {_wgib(sharded_weights.moe_routed_expert):.3f} |\n"
    )
    lines.append(
        f"| moe_shared | {format_count(sharded_weights.moe_shared_expert)} | {_wgib(sharded_weights.moe_shared_expert):.3f} |\n"
    )
    lines.append(
        f"| moe_router | {format_count(sharded_weights.moe_router)} | {_wgib(sharded_weights.moe_router):.3f} |\n"
    )
    lines.append(
        f"| mhc | {format_count(sharded_weights.mhc)} | {_wgib(sharded_weights.mhc):.3f} |\n"
    )

    lines.append("\n## KV / Index / Activations (per rank)\n")
    lines.append(
        "> Notes: KV cache uses the MLA latent cache shape: `kv_dim = kv_lora_rank + qk_rope_head_dim`.\n\n"
    )

    lines.append(
        "| seq_len | local_seq | local_batch | kv | index | act_peak(prefill) | misc | total | budget | headroom | fits |\n"
        "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|:---:|\n"
    )
    for est in seq_estimates:
        idx = float(est.idx_per_rank_gib or 0.0)
        act_peak = float(est.prefill_ws.peak_gib)
        misc = float(misc_gib)
        total = weights_per_rank_gib + float(est.kv_per_rank_gib) + idx + act_peak + misc
        if budget_gib is None:
            budget_text = "N/A"
            headroom_text = "N/A"
            fits_text = "N/A"
        else:
            headroom = float(budget_gib) - total
            budget_text = f"{float(budget_gib):.1f}"
            headroom_text = f"{headroom:.1f}"
            fits_text = "yes" if headroom >= 0 else "no"
        lines.append(
            f"| {est.seq_len} | {est.local_seq_len} | {est.local_batch:.4f} | "
            f"{est.kv_per_rank_gib:.1f} | {idx:.1f} | {act_peak:.2f} | {misc:.1f} | {total:.1f} | "
            f"{budget_text} | {headroom_text} | {fits_text} |\n"
        )

    # perf section
    has_perf = any(est.decode_step is not None for est in seq_estimates)
    if has_perf:
        lines.append("\n## Perf Estimate (roofline, per step)\n")
        lines.append(
            "| seq_len | local_seq | decode_ms | decode_tok/s | prefill_ms | prefill_tok/s |\n"
            "|---:|---:|---:|---:|---:|---:|\n"
        )
        for est in seq_estimates:
            if est.decode_step is None or est.prefill_step is None:
                continue
            global_prefill_tokens = int(
                est.prefill_tokens_per_rank * parallel.dp * parallel.sp
            )
            lines.append(
                f"| {est.seq_len} | {est.local_seq_len} | {est.decode_step.ms:.2f} | {est.decode_step.tok_s:.1f} | "
                f"{est.prefill_step.ms:.2f} | {(float(global_prefill_tokens)/est.prefill_step.total_s if est.prefill_step.total_s>0 else 0.0):.1f} |\n"
            )

        # top ops for the first perf-enabled estimate
        first = next((x for x in seq_estimates if x.decode_step is not None), None)
        if first is not None and first.decode_step is not None and first.prefill_step is not None:
            def _render_op_table(label: str, step: StepTime) -> None:
                lines.append(f"\n### Top Ops ({label})\n")
                lines.append(
                    "| op | total_ms | bound | compute_ms | mem_ms | comm_ms | comm | inter-node |\n"
                    "|---|---:|:---:|---:|---:|---:|---|:---:|\n"
                )
                for op in sorted(step.ops, key=lambda x: x.t_total_s, reverse=True)[:12]:
                    comm = (
                        f"{op.comm_kind}:{op.comm_group}[{op.comm_group_size}]"
                        if op.comm_kind
                        else ""
                    )
                    lines.append(
                        f"| {op.name} | {op.t_total_s*1e3:.2f} | {op.bound} | {op.t_compute_s*1e3:.2f} | {op.t_mem_s*1e3:.2f} | {op.t_comm_s*1e3:.2f} | {comm} | {'yes' if op.comm_inter_node else 'no'} |\n"
                    )

            _render_op_table("decode", first.decode_step)
            _render_op_table("prefill", first.prefill_step)

        lines.append(
            "\n> Notes: This is a simplified model (roofline + collectives). "
            "Use it to compare configs / find knees, not as an exact latency predictor.\n"
        )

    return "".join(lines)


def _iter_seq_lens(base: int, also_max_pos: bool, max_pos: int) -> Iterable[int]:
    yield base
    if also_max_pos and max_pos > 0 and max_pos != base:
        yield max_pos


def _guess_hf_token() -> Optional[str]:
    return (
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
        or os.environ.get("HF_ACCESS_TOKEN")
    )


def fetch_hf_config(
    repo_id: str,
    *,
    revision: str = "main",
    token: Optional[str] = None,
    timeout_s: int = 30,
) -> Tuple[Dict[str, Any], str]:
    token = token or _guess_hf_token()

    try:
        from huggingface_hub import hf_hub_download  # type: ignore
    except Exception:
        hf_hub_download = None

    if hf_hub_download is not None:
        path = hf_hub_download(
            repo_id=repo_id,
            filename="config.json",
            revision=revision,
            token=token,
        )
        cfg = json.loads(Path(path).read_text())
        return cfg, f"hf:{repo_id}@{revision}"

    url = f"https://huggingface.co/{repo_id}/raw/{revision}/config.json"
    req = urllib.request.Request(url)
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    with urllib.request.urlopen(req, timeout=timeout_s) as resp:
        text = resp.read().decode("utf-8")
    cfg = json.loads(text)
    return cfg, f"hf:{repo_id}@{revision}"


def _sanitize_for_title(text: str) -> str:
    return "".join(
        c if c.isalnum() or c in ("-", "_", "/", ".", "@") else "_" for c in text
    )


def _render_obsidian_frontmatter(
    *,
    title: str,
    created: str,
    tags: Sequence[str],
    model: str,
) -> str:
    tags_text = ", ".join(tags)
    return (
        "---\n"
        f"title: \"{title}\"\n"
        f"created: {created}\n"
        f"tags: [{tags_text}]\n"
        f"model: {model}\n"
        "---\n\n"
    )


def _relative_path_for_embed(target: Path, base: Path) -> str:
    try:
        rel = os.path.relpath(str(target), start=str(base))
    except Exception:
        rel = str(target)
    return rel.replace("\\", "/")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="llm-mem",
        description="Estimate model VRAM usage and recommend tp/dp/sp configs.",
    )
    try:
        from llm_mem_planner import __version__
    except Exception:
        __version__ = "unknown"
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")

    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument("--config", type=str, help="Path to config.json")
    src.add_argument("--model", type=str, help="HF repo id (e.g. deepseek-ai/DeepSeek-V3.2)")
    parser.add_argument(
        "--revision",
        type=str,
        default="main",
        help="HF revision for --model (branch/tag/commit). Default: main",
    )
    parser.add_argument(
        "--hf-token",
        type=str,
        default=None,
        help="HF token for private repos (or set HF_TOKEN/HUGGINGFACE_HUB_TOKEN).",
    )

    parser.add_argument("--seq-len", type=int, default=20480)
    parser.add_argument("--batch", type=int, default=1)
    parser.add_argument(
        "--batch-scope",
        choices=["global", "local"],
        default="global",
        help=(
            "Interpretation of --batch when estimating KV/cache: "
            "'global' means total sequences across the DP-attention dimension "
            "(local_batch = global_batch / dp). "
            "'local' means per-rank batch (global_batch ~= local_batch * dp)."
        ),
    )
    parser.add_argument(
        "--also-max-pos",
        action="store_true",
        help="Also estimate at max_position_embeddings",
    )
    parser.add_argument(
        "--weight-dtype",
        type=str,
        default=None,
        help="Weights dtype override (default: config dtype).",
    )
    parser.add_argument(
        "--kv-dtype",
        type=str,
        default=None,
        help="KV cache dtype override (default: config dtype).",
    )
    parser.add_argument(
        "--act-dtype",
        type=str,
        default=None,
        help="Activation dtype override (default: config dtype).",
    )
    parser.add_argument(
        "--chunked-prefill-size",
        type=int,
        default=8192,
        help=(
            "Approx prefill tokens per forward step. "
            "If DP-attention is enabled, some runtimes divide this by dp."
        ),
    )

    # mHC (multi-stream residuals) modeling (optional)
    parser.add_argument(
        "--mhc-streams",
        "--hc-num-streams",
        type=int,
        default=1,
        help=(
            "mHC residual streams (S). "
            "Use with --mhc-branch-mode to choose whether branches are single-stream or per-stream. "
            "S<=1 disables. Default: 1."
        ),
    )
    parser.add_argument(
        "--mhc-branches-per-layer",
        type=int,
        default=2,
        help=(
            "Number of HyperConnections modules per transformer layer. "
            "Typical integration is 2 (attn + mlp). Default: 2."
        ),
    )
    parser.add_argument(
        "--mhc-branch-mode",
        choices=["single", "per_stream"],
        default="single",
        help=(
            "mHC branch compute model: "
            "'single' assumes mhc-torch style (branch stays single-stream; KV/index cache unchanged); "
            "'per_stream' assumes each stream runs its own attention/MLP (KV/index + act scale with streams)."
        ),
    )

    parser.add_argument(
        "--tp",
        type=int,
        default=None,
        help=(
            "Tensor-parallel world size (tp). Default: 1; in --recommend defaults to nodes*gpus-per-node."
        ),
    )
    parser.add_argument("--dp", type=int, default=1)
    parser.add_argument("--sp", type=int, default=1)
    parser.add_argument("--ep", type=int, default=1)
    parser.add_argument("--moe-dp", type=int, default=1)

    # recommend mode
    parser.add_argument("--recommend", action="store_true")
    parser.add_argument("--nodes", type=int, default=1)
    parser.add_argument("--gpus-per-node", type=int, default=8)
    parser.add_argument("--gpu-mem-gib", type=float, default=None)
    parser.add_argument("--mem-fraction", type=float, default=0.90)
    parser.add_argument("--misc-gib", type=float, default=6.0)
    # perf modeling (optional)
    parser.add_argument(
        "--gpu-tflops",
        type=float,
        default=None,
        help="Peak GPU throughput in TFLOPS (for main GEMMs, per GPU).",
    )
    parser.add_argument(
        "--gpu-hbm-gibps",
        type=float,
        default=None,
        help="Peak GPU HBM bandwidth in GiB/s (per GPU).",
    )
    parser.add_argument(
        "--intra-gibps",
        type=float,
        default=None,
        help="Effective intra-node collective bandwidth in GiB/s.",
    )
    parser.add_argument(
        "--inter-gibps",
        type=float,
        default=None,
        help="Effective inter-node collective bandwidth in GiB/s.",
    )
    parser.add_argument(
        "--intra-us",
        type=float,
        default=2.0,
        help="Per-step latency for intra-node collectives (microseconds).",
    )
    parser.add_argument(
        "--inter-us",
        type=float,
        default=8.0,
        help="Per-step latency for inter-node collectives (microseconds).",
    )
    parser.add_argument(
        "--flops-eff",
        type=float,
        default=0.50,
        help="Compute efficiency scalar (0~1) applied to peak TFLOPS.",
    )
    parser.add_argument(
        "--hbm-eff",
        type=float,
        default=0.80,
        help="HBM bandwidth efficiency scalar (0~1).",
    )
    parser.add_argument(
        "--comm-eff",
        type=float,
        default=0.85,
        help="Network bandwidth efficiency scalar (0~1).",
    )
    parser.add_argument(
        "--gemm-m-saturation",
        type=int,
        default=256,
        help="Small-M GEMM penalty knee (tokens per rank).",
    )
    parser.add_argument(
        "--dp-padding-mode",
        choices=["auto", "max_len", "sum_len"],
        default="auto",
        help="dp-attn padding mode model (auto/max_len/sum_len).",
    )
    parser.add_argument(
        "--preference",
        choices=["latency", "throughput", "balanced"],
        default="balanced",
    )
    parser.add_argument("--sp-candidates", type=str, default="1")
    parser.add_argument("--topk", type=int, default=8)
    parser.add_argument("--out-md", type=str, default=None)
    parser.add_argument("--out-svg", type=str, default=None)
    parser.add_argument("--out-perf-svg", type=str, default=None)
    parser.add_argument("--obsidian", action="store_true")
    return parser


def _load_config(args: argparse.Namespace) -> Tuple[Dict[str, Any], str]:
    if args.config:
        path = Path(args.config)
        cfg = json.loads(path.read_text())
        return cfg, f"path:{path}"
    try:
        cfg, src = fetch_hf_config(
            str(args.model),
            revision=str(args.revision),
            token=args.hf_token,
        )
    except json.JSONDecodeError:
        raise
    except Exception as exc:
        raise urllib.error.URLError(f"{type(exc).__name__}: {exc}") from exc
    return cfg, src


def _user_die(message: str, *, exit_code: int = 1) -> None:
    print(f"error: {message}", file=sys.stderr)
    raise SystemExit(exit_code)


def _parse_csv_ints(value: str, *, arg_name: str) -> List[int]:
    raw = str(value)
    parts = [p.strip() for p in raw.split(",") if p.strip()]
    if not parts:
        raise ValueError(f"{arg_name} cannot be empty")
    out: List[int] = []
    for p in parts:
        try:
            out.append(int(p))
        except Exception as exc:
            raise ValueError(
                f"Invalid {arg_name}: {raw!r}. Expected a comma-separated list of integers."
            ) from exc
    return out


def _run(args: argparse.Namespace) -> None:

    if args.recommend and args.batch_scope != "global":
        raise SystemExit("--recommend currently requires --batch-scope global")

    mhc_streams = int(args.mhc_streams)
    mhc_branches_per_layer = int(args.mhc_branches_per_layer)
    mhc_branch_mode = str(args.mhc_branch_mode)
    if mhc_streams < 1:
        raise SystemExit("--mhc-streams must be >= 1")
    if mhc_branches_per_layer < 0:
        raise SystemExit("--mhc-branches-per-layer must be >= 0")
    mhc = MhcSpec(
        streams=mhc_streams,
        branches_per_layer=mhc_branches_per_layer,
        branch_mode=mhc_branch_mode,
    )

    config, config_src = _load_config(args)
    d = load_dims(config)
    p = estimate_params(d, mhc=mhc)
    active = estimate_activated_params_per_token(d, p)

    weight_dtype = args.weight_dtype or d.dtype
    weight_dtype_bytes = dtype_bytes_from_str(weight_dtype)
    kv_dtype = args.kv_dtype or d.dtype
    kv_dtype_bytes = dtype_bytes_from_str(kv_dtype)
    act_dtype = args.act_dtype or d.dtype
    act_dtype_bytes = dtype_bytes_from_str(act_dtype)

    if args.recommend:
        if args.gpu_mem_gib is None:
            raise SystemExit("--recommend requires --gpu-mem-gib")

        hw = HardwareSpec(
            nodes=max(1, int(args.nodes)),
            gpus_per_node=max(1, int(args.gpus_per_node)),
            gpu_mem_gib=float(args.gpu_mem_gib),
            mem_fraction=float(args.mem_fraction),
            misc_gib=float(args.misc_gib),
            gpu_tflops=args.gpu_tflops,
            gpu_hbm_gibps=args.gpu_hbm_gibps,
            intra_gibps=args.intra_gibps,
            inter_gibps=args.inter_gibps,
            intra_us=float(args.intra_us),
            inter_us=float(args.inter_us),
            flops_eff=float(args.flops_eff),
            hbm_eff=float(args.hbm_eff),
            comm_eff=float(args.comm_eff),
            gemm_m_saturation=int(args.gemm_m_saturation),
        )
        sp_candidates = _parse_csv_ints(str(args.sp_candidates), arg_name="--sp-candidates")
        if any(x <= 0 for x in sp_candidates):
            raise ValueError(
                f"Invalid --sp-candidates: {args.sp_candidates!r}. Values must be >= 1."
            )
        tp_arg = int(args.tp) if args.tp is not None else None
        if tp_arg is not None and tp_arg <= 0:
            raise ValueError(
                f"Invalid --tp: {args.tp!r}. Expected >= 1. "
                "Omit --tp in --recommend mode to default to nodes*gpus-per-node."
            )
        tp_effective = tp_arg if tp_arg is not None else hw.world_size
        points, critical = recommend_parallelism(
            d,
            p,
            hw=hw,
            seq_len=int(args.seq_len),
            global_batch=int(args.batch),
            weight_dtype_bytes=weight_dtype_bytes,
            kv_dtype_bytes=kv_dtype_bytes,
            act_dtype_bytes=act_dtype_bytes,
            chunked_prefill_size=int(args.chunked_prefill_size),
            mhc_streams=mhc.streams,
            mhc_branch_streams=mhc.branch_streams,
            ep=int(args.ep),
            moe_dp=int(args.moe_dp),
            tp=tp_effective,
            sp_candidates=sp_candidates,
            preference=str(args.preference),
            dp_padding_mode=str(args.dp_padding_mode),
        )

        best_point = next((pt for pt in points if pt.fits), None)
        best_weight_breakdown: Optional[ShardedWeightBreakdown] = None
        if best_point is not None:
            best_parallel = Parallelism(
                tp=best_point.tp,
                dp=best_point.dp,
                sp=best_point.sp,
                ep=best_point.ep,
                moe_dp=best_point.moe_dp,
            )
            best_weight_breakdown = estimate_weight_params_per_rank(d, p, best_parallel)

        out_svg = Path(args.out_svg) if args.out_svg else Path("parallelism_dp_scan.svg")

        scan_sp = 1
        scan_points = [pt for pt in points if pt.sp == scan_sp]
        scan_points.sort(key=lambda x: x.dp)
        if scan_points:
            x = [pt.dp for pt in scan_points]
            series = {
                "total": [pt.total_per_rank_gib for pt in scan_points],
                "weights": [pt.weights_per_rank_gib for pt in scan_points],
                "kv+index": [
                    pt.kv_per_rank_gib + pt.index_per_rank_gib for pt in scan_points
                ],
                "act": [pt.act_per_rank_gib for pt in scan_points],
            }
            write_memory_scan_svg(
                out_svg,
                title=f"DP scan (sp={scan_sp}) | seq={args.seq_len}, batch={args.batch}",
                x=x,
                series=series,
                hlines={"budget": hw.budget_gib},
                vlines={
                    "dp_fit": critical.get("dp_min_fit"),
                    "dp_intra_node": critical.get("dp_min_intra_node"),
                },
            )

        out_perf_svg = (
            Path(args.out_perf_svg)
            if args.out_perf_svg
            else Path("parallelism_dp_scan_perf.svg")
        )
        if _has_perf(hw) and scan_points and all(pt.decode_ms is not None for pt in scan_points):
            x = [pt.dp for pt in scan_points]
            series = {
                "decode_ms": [float(pt.decode_ms or 0.0) for pt in scan_points],
                "prefill_ms": [float(pt.prefill_ms or 0.0) for pt in scan_points],
            }
            write_memory_scan_svg(
                out_perf_svg,
                title=f"DP scan (perf, sp={scan_sp}) | seq={args.seq_len}, batch={args.batch}",
                x=x,
                series=series,
                hlines={},
                vlines={
                    "dp_fit": critical.get("dp_min_fit"),
                    "dp_intra_node": critical.get("dp_min_intra_node"),
                },
                y_label="ms per step (roofline est.)",
            )

        ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        out_md = Path(args.out_md) if args.out_md else Path("parallelism_report.md")

        lines: List[str] = []
        if args.obsidian:
            model_title = _sanitize_for_title(config_src.replace("path:", ""))
            lines.append(
                _render_obsidian_frontmatter(
                    title=f"Parallelism Recommendation - {model_title}",
                    created=ts,
                    tags=("llm", "vram", "parallelism", "dp-attention"),
                    model=config_src,
                )
            )

        lines.append("# Parallelism Recommendation Report\n\n")
        lines.append(f"Generated: {ts}\n\n")
        lines.append("## Inputs\n")
        lines.append(f"- config: `{config_src}`\n")
        lines.append(f"- weight dtype: `{weight_dtype}`\n")
        lines.append(f"- kv dtype: `{kv_dtype}`\n")
        lines.append(f"- act dtype: `{act_dtype}`\n")
        lines.append(f"- seq_len: `{args.seq_len}`\n")
        lines.append(f"- batch(scope=global): `{args.batch}`\n")
        lines.append(f"- chunked_prefill_size: `{args.chunked_prefill_size}`\n")
        lines.append(
            f"- mHC: `streams={mhc.streams}`, `branches_per_layer={mhc.branches_per_layer}`, "
            f"`branch_mode={mhc.branch_mode}`, `branch_streams={mhc.branch_streams}`\n"
        )
        if mhc.enabled:
            lines.append(
                f"- mHC params_total: `{format_count(p.mhc_total)}` params (replicated)\n"
            )
        lines.append(f"- tp={tp_effective}, ep={args.ep}, moe_dp={args.moe_dp}\n")
        lines.append(f"- sp candidates: `{args.sp_candidates}`\n")
        lines.append(f"- preference: `{args.preference}`\n")

        lines.append("\n## Hardware\n")
        lines.append(f"- nodes: `{hw.nodes}`\n")
        lines.append(f"- gpus_per_node: `{hw.gpus_per_node}`\n")
        lines.append(f"- gpu_mem: `{hw.gpu_mem_gib:.1f} GiB`\n")
        lines.append(
            f"- mem_fraction: `{hw.mem_fraction:.2f}` ⇒ budget `{hw.budget_gib:.1f} GiB`\n"
        )
        lines.append(f"- misc reserved: `{hw.misc_gib:.1f} GiB`\n")
        if _has_perf(hw):
            lines.append("\n### Perf Model (roofline)\n")
            lines.append(f"- gpu_tflops: `{hw.gpu_tflops}`\n")
            lines.append(f"- gpu_hbm_gibps: `{hw.gpu_hbm_gibps}`\n")
            lines.append(f"- intra_gibps: `{hw.intra_gibps}`\n")
            lines.append(f"- inter_gibps: `{hw.inter_gibps}`\n")
            lines.append(f"- intra_us: `{hw.intra_us}`\n")
            lines.append(f"- inter_us: `{hw.inter_us}`\n")
            lines.append(
                f"- eff: flops `{hw.flops_eff:.2f}`, hbm `{hw.hbm_eff:.2f}`, comm `{hw.comm_eff:.2f}`\n"
            )
            lines.append(f"- gemm_m_saturation: `{hw.gemm_m_saturation}`\n")
            lines.append(f"- dp_padding_mode: `{args.dp_padding_mode}`\n")
            if mhc.enabled:
                lines.append(
                    "\n> [!warning] Perf model currently ignores mHC overhead; treat perf estimates as optimistic.\n"
                    if args.obsidian
                    else "\nNote: perf model currently ignores mHC overhead; treat perf estimates as optimistic.\n"
                )
        else:
            lines.append(
                "\n> [!warning] Perf model disabled (missing --gpu-tflops/--gpu-hbm-gibps/--intra-gibps/--inter-gibps)\n"
                if args.obsidian
                else "\nNote: perf model disabled (missing gpu/comm specs).\n"
            )

        if args.obsidian:
            lines.append("\n> [!info] Critical Points (sp=1 baseline)\n")
            lines.append(f"> - dp_min_fit: `{critical.get('dp_min_fit')}`\n")
            lines.append(
                f"> - dp_min_intra_node(attn_tp<=gpus_per_node): `{critical.get('dp_min_intra_node')}`\n"
            )
        else:
            lines.append("\n## Critical Points (sp=1 baseline)\n")
            lines.append(f"- dp_min_fit: `{critical.get('dp_min_fit')}`\n")
            lines.append(
                f"- dp_min_intra_node(attn_tp<=gpus_per_node): `{critical.get('dp_min_intra_node')}`\n"
            )

        if out_svg.exists() or out_perf_svg.exists():
            lines.append("\n## Plots\n")
            if out_svg.exists():
                if args.obsidian:
                    rel_svg = _relative_path_for_embed(out_svg, out_md.parent)
                    lines.append(f"- memory scan: ![[{rel_svg}]]\n")
                else:
                    lines.append(f"- memory scan svg: `{out_svg}`\n")
            if out_perf_svg.exists():
                if args.obsidian:
                    rel_svg = _relative_path_for_embed(out_perf_svg, out_md.parent)
                    lines.append(f"- perf scan: ![[{rel_svg}]]\n")
                else:
                    lines.append(f"- perf scan svg: `{out_perf_svg}`\n")

        lines.append("\n## Model Summary (from config)\n")
        lines.append(f"- hidden_size: `{d.hidden_size}`\n")
        lines.append(f"- layers: `{d.num_layers}`\n")
        lines.append(f"- heads: `{d.num_heads}`\n")
        lines.append(f"- MoE layers: `{d.moe_layers}` (dense first `{d.dense_layers}`)\n")
        lines.append(
            "- MLA dims: "
            f"`q_lora_rank={d.q_lora_rank}`, `kv_lora_rank={d.kv_lora_rank}`, "
            f"`qk_nope={d.qk_nope_head_dim}`, `qk_rope={d.qk_rope_head_dim}`, `v={d.v_head_dim}`\n"
        )
        if d.index_n_heads and d.index_head_dim:
            lines.append(
                f"- NSA indexer: `index_n_heads={d.index_n_heads}`, `index_head_dim={d.index_head_dim}`, `index_topk={d.index_topk}`\n"
            )

        lines.append("\n## Activated Params / Token (compute metric)\n")
        lines.append(
            f"- total_excluding_embed_router: `{format_count(active['total_excluding_embed_router'])}` params\n"
        )
        lines.append(
            f"- breakdown: moe `{format_count(active['active_moe'])}` + dense_ffn `{format_count(active['active_dense_ffn'])}` + always(attn+indexer) `{format_count(active['always_active(attn+indexer)'])}` + mhc `{format_count(active['active_mhc'])}`\n"
        )

        lines.append("\n## Top Candidates\n")
        has_perf_points = any(pt.decode_ms is not None for pt in points)
        if has_perf_points:
            lines.append(
                "| rank | dp | sp | attn_tp | weights | kv | index | act | total | headroom | decode_ms | tok/s | prefill_ms | intra-node |\n"
                "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|:---:|\n"
            )
        else:
            lines.append(
                "| rank | dp | sp | attn_tp | weights | kv | index | act | total | headroom | intra-node |\n"
                "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|:---:|\n"
            )

        shown = 0
        for pt in points:
            if shown >= int(args.topk):
                break
            if not pt.fits:
                continue
            headroom = pt.budget_gib - pt.total_per_rank_gib
            if has_perf_points:
                decode_ms = f"{pt.decode_ms:.2f}" if pt.decode_ms is not None else "N/A"
                decode_tok_s = (
                    f"{pt.decode_tok_s:.1f}" if pt.decode_tok_s is not None else "N/A"
                )
                prefill_ms = f"{pt.prefill_ms:.2f}" if pt.prefill_ms is not None else "N/A"
                lines.append(
                    f"| {shown+1} | {pt.dp} | {pt.sp} | {pt.attn_tp} | "
                    f"{pt.weights_per_rank_gib:.1f} | {pt.kv_per_rank_gib:.1f} | {pt.index_per_rank_gib:.1f} | {pt.act_per_rank_gib:.1f} | "
                    f"{pt.total_per_rank_gib:.1f} | {headroom:.1f} | "
                    f"{decode_ms} | {decode_tok_s} | {prefill_ms} | "
                    f"{'yes' if pt.intra_node_attn else 'no'} |\n"
                )
            else:
                lines.append(
                    f"| {shown+1} | {pt.dp} | {pt.sp} | {pt.attn_tp} | "
                    f"{pt.weights_per_rank_gib:.1f} | {pt.kv_per_rank_gib:.1f} | {pt.index_per_rank_gib:.1f} | {pt.act_per_rank_gib:.1f} | "
                    f"{pt.total_per_rank_gib:.1f} | {headroom:.1f} | {'yes' if pt.intra_node_attn else 'no'} |\n"
                )
            shown += 1

        if shown == 0:
            lines.append("\nNo feasible configs found under given constraints.\n")

        if best_point is not None and best_weight_breakdown is not None:
            lines.append("\n## Best-Config Weights Breakdown (per rank, approx)\n")
            lines.append(
                f"- best: `dp={best_point.dp}`, `sp={best_point.sp}`, `attn_tp={best_point.attn_tp}`, `moe_tp={best_point.moe_tp}`\n"
            )

            def _wgib(x_params: float) -> float:
                return (x_params * weight_dtype_bytes) / GiB

            lines.append("| part | params | GiB |\n|---|---:|---:|\n")
            lines.append(
                f"| embedding | {format_count(best_weight_breakdown.embedding)} | {_wgib(best_weight_breakdown.embedding):.3f} |\n"
            )
            lines.append(
                f"| lm_head | {format_count(best_weight_breakdown.lm_head)} | {_wgib(best_weight_breakdown.lm_head):.3f} |\n"
            )
            lines.append(
                f"| attn_fused_qkv_a | {format_count(best_weight_breakdown.attn_fused_qkv_a)} | {_wgib(best_weight_breakdown.attn_fused_qkv_a):.3f} |\n"
            )
            lines.append(
                f"| attn_q_b | {format_count(best_weight_breakdown.attn_q_b)} | {_wgib(best_weight_breakdown.attn_q_b):.3f} |\n"
            )
            lines.append(
                f"| attn_kv_b | {format_count(best_weight_breakdown.attn_kv_b)} | {_wgib(best_weight_breakdown.attn_kv_b):.3f} |\n"
            )
            lines.append(
                f"| attn_o_proj | {format_count(best_weight_breakdown.attn_o_proj)} | {_wgib(best_weight_breakdown.attn_o_proj):.3f} |\n"
            )
            lines.append(
                f"| indexer | {format_count(best_weight_breakdown.indexer)} | {_wgib(best_weight_breakdown.indexer):.3f} |\n"
            )
            lines.append(
                f"| dense_ffn | {format_count(best_weight_breakdown.dense_ffn)} | {_wgib(best_weight_breakdown.dense_ffn):.3f} |\n"
            )
            lines.append(
                f"| moe_routed | {format_count(best_weight_breakdown.moe_routed_expert)} | {_wgib(best_weight_breakdown.moe_routed_expert):.3f} |\n"
            )
            lines.append(
                f"| moe_shared | {format_count(best_weight_breakdown.moe_shared_expert)} | {_wgib(best_weight_breakdown.moe_shared_expert):.3f} |\n"
            )
            lines.append(
                f"| moe_router | {format_count(best_weight_breakdown.moe_router)} | {_wgib(best_weight_breakdown.moe_router):.3f} |\n"
            )
            lines.append(
                f"| mhc | {format_count(best_weight_breakdown.mhc)} | {_wgib(best_weight_breakdown.mhc):.3f} |\n"
            )

        # --- perf breakdown for the best config ---
        if best_point is not None and _has_perf(hw):
            best_parallel = Parallelism(
                tp=best_point.tp,
                dp=best_point.dp,
                sp=best_point.sp,
                ep=best_point.ep,
                moe_dp=best_point.moe_dp,
            )
            best_local_batch = float(args.batch) / float(max(1, best_point.dp))
            best_local_seq = int(math.ceil(int(args.seq_len) / float(max(1, best_point.sp))))
            best_local_decode_tokens = int(max(1, math.ceil(best_local_batch)))
            best_local_chunk = int(max(1, int(args.chunked_prefill_size) // max(1, best_point.dp)))
            best_local_prefill_tokens = int(
                min(best_local_chunk, math.ceil(best_local_batch * best_local_seq))
            )
            try:
                decode = estimate_step_time(
                    d,
                    best_parallel,
                    hw=hw,
                    phase="decode",
                    tokens_per_rank=best_local_decode_tokens,
                    seq_len_local=best_local_seq,
                    global_batch=int(args.batch),
                    weight_dtype_bytes=weight_dtype_bytes,
                    act_dtype_bytes=act_dtype_bytes,
                    kv_dtype_bytes=kv_dtype_bytes,
                    dp_padding_mode=str(args.dp_padding_mode),
                )
                prefill = estimate_step_time(
                    d,
                    best_parallel,
                    hw=hw,
                    phase="prefill",
                    tokens_per_rank=best_local_prefill_tokens,
                    seq_len_local=best_local_seq,
                    global_batch=int(args.batch),
                    weight_dtype_bytes=weight_dtype_bytes,
                    act_dtype_bytes=act_dtype_bytes,
                    kv_dtype_bytes=kv_dtype_bytes,
                    dp_padding_mode=str(args.dp_padding_mode),
                )

                lines.append("\n## Perf Estimate (roofline, best config)\n")
                lines.append(
                    f"- decode: `{decode.ms:.2f} ms/step` ⇒ `{(float(args.batch)/decode.total_s if decode.total_s>0 else 0.0):.2f} tok/s` (global)\n"
                )
                global_prefill_tokens = int(best_local_prefill_tokens * best_point.dp * best_point.sp)
                lines.append(
                    f"- prefill(chunk): `{prefill.ms:.2f} ms/step` ⇒ `{(float(global_prefill_tokens)/prefill.total_s if prefill.total_s>0 else 0.0):.2f} tok/s` (global)\n"
                )

                def _sum_ms(items: Sequence[OpTime], field: str) -> float:
                    return 1e3 * sum(getattr(x, field) for x in items)

                lines.append("\n### Lower Bounds (by resource)\n")
                lines.append("| phase | total_ms | compute_lb_ms | mem_lb_ms | comm_lb_ms |\n")
                lines.append("|---|---:|---:|---:|---:|\n")
                lines.append(
                    f"| decode | {decode.ms:.2f} | {_sum_ms(decode.ops, 't_compute_s'):.2f} | {_sum_ms(decode.ops, 't_mem_s'):.2f} | {_sum_ms(decode.ops, 't_comm_s'):.2f} |\n"
                )
                lines.append(
                    f"| prefill | {prefill.ms:.2f} | {_sum_ms(prefill.ops, 't_compute_s'):.2f} | {_sum_ms(prefill.ops, 't_mem_s'):.2f} | {_sum_ms(prefill.ops, 't_comm_s'):.2f} |\n"
                )

                def _render_op_table(phase_name: str, step: StepTime) -> None:
                    lines.append(f"\n### Op Breakdown ({phase_name})\n")
                    lines.append(
                        "| op | total_ms | bound | compute_ms | mem_ms | comm_ms | comm | inter |\n"
                    )
                    lines.append("|---|---:|:---:|---:|---:|---:|---|:---:|\n")
                    top_ops = sorted(step.ops, key=lambda x: x.t_total_s, reverse=True)[:20]
                    for op in top_ops:
                        comm = (
                            f"{op.comm_kind}:{op.comm_group}[{op.comm_group_size}]"
                            if op.comm_kind
                            else ""
                        )
                        lines.append(
                            f"| {op.name} | {op.t_total_s*1e3:.2f} | {op.bound} | {op.t_compute_s*1e3:.2f} | {op.t_mem_s*1e3:.2f} | {op.t_comm_s*1e3:.2f} | {comm} | {'yes' if op.comm_inter_node else 'no'} |\n"
                        )

                _render_op_table("decode", decode)
                _render_op_table("prefill", prefill)

                lines.append(
                    "\n> Notes: This is a simplified model (roofline + collectives). "
                    "Use it to compare configs / find knees, not as an exact latency predictor.\n"
                )
            except Exception as exc:
                lines.append(
                    f"\n> [!warning] Perf estimate failed for best config: `{type(exc).__name__}: {exc}`\n"
                    if args.obsidian
                    else f"\nPerf estimate failed for best config: {type(exc).__name__}: {exc}\n"
                )

        out_md.write_text("".join(lines))

        print("== Recommendation Output ==")
        print(f"report: {out_md}")
        if out_svg.exists():
            print(f"plot:   {out_svg}")
        if out_perf_svg.exists():
            print(f"perf:   {out_perf_svg}")
        print(f"critical(dp_min_fit): {critical.get('dp_min_fit')}")
        print(f"critical(dp_min_intra_node): {critical.get('dp_min_intra_node')}")
        print("\nTop feasible configs:")
        remaining = int(args.topk)
        for pt in points:
            if not pt.fits:
                continue
            decode_ms_text = f"{pt.decode_ms:.2f}" if pt.decode_ms is not None else "N/A"
            print(
                f"- dp={pt.dp}, sp={pt.sp}, attn_tp={pt.attn_tp}: "
                f"total_per_rank={pt.total_per_rank_gib:.2f} GiB (budget={pt.budget_gib:.2f}), "
                f"decode_ms={decode_ms_text}"
            )
            remaining -= 1
            if remaining <= 0:
                break
        return

    tp = int(args.tp) if args.tp is not None else 1
    parallel = Parallelism(
        tp=tp,
        dp=int(args.dp),
        sp=int(args.sp),
        ep=int(args.ep),
        moe_dp=int(args.moe_dp),
    )
    sharded_weights = estimate_weight_params_per_rank(d, p, parallel)

    # --- non-recommend mode: print details ---
    want_report = bool(args.out_md) or bool(args.obsidian)
    want_plot = bool(args.out_svg) or bool(args.obsidian)
    out_md = Path(args.out_md) if args.out_md else Path("mem_report.md")
    out_svg: Optional[Path]
    if not want_plot:
        out_svg = None
    elif args.out_svg:
        out_svg = Path(args.out_svg)
    elif want_report:
        out_svg = out_md.with_suffix(".svg")
    else:
        out_svg = Path("mem_breakdown.svg")

    budget_gib: Optional[float] = None
    misc_gib = 0.0
    if args.gpu_mem_gib is not None and float(args.gpu_mem_gib) > 0:
        budget_gib = float(args.gpu_mem_gib) * float(args.mem_fraction)
        misc_gib = float(args.misc_gib)

    print("== Model Summary ==")
    print(f"config: {config_src}")
    print(f"dtype: {d.dtype} ({d.dtype_bytes} bytes)")
    print(f"weight_dtype: {weight_dtype} ({weight_dtype_bytes} bytes)")
    print(f"hidden_size: {d.hidden_size}")
    print(f"layers: {d.num_layers}")
    print(f"heads: {d.num_heads}")
    print(f"vocab_size: {d.vocab_size}")
    print(f"MoE layers: {d.moe_layers} (dense first {d.dense_layers})")
    print(
        "MLA dims: "
        f"q_lora_rank={d.q_lora_rank}, kv_lora_rank={d.kv_lora_rank}, "
        f"qk_nope={d.qk_nope_head_dim}, qk_rope={d.qk_rope_head_dim}, v={d.v_head_dim}"
    )
    print(
        "mHC: "
        f"streams={mhc.streams}, branches_per_layer={mhc.branches_per_layer}, "
        f"branch_mode={mhc.branch_mode}, branch_streams={mhc.branch_streams}, "
        f"params_total={format_count(p.mhc_total)}"
    )
    if d.index_n_heads and d.index_head_dim:
        extra = f", index_topk={d.index_topk}" if d.index_topk else ""
        print(
            f"NSA indexer: index_n_heads={d.index_n_heads}, index_head_dim={d.index_head_dim}{extra}"
        )

    print("\n== Parallelism ==")
    print(f"tp={parallel.tp}, dp={parallel.dp}, sp={parallel.sp}, ep={parallel.ep}, moe_dp={parallel.moe_dp}")
    print(f"derived: attn_tp={parallel.attn_tp}, moe_tp={parallel.moe_tp}")

    print("\n== Parameter Breakdown (approx) ==")
    print(f"embedding:        {format_count(p.embedding)} params")
    print(f"lm_head:          {format_count(p.lm_head)} params")
    print(f"attention_total:  {format_count(p.attention_total)} params")
    print(f"  - fused_qkv_a:  {format_count(p.attn_fused_qkv_a_total)} params")
    print(f"  - q_b_proj:     {format_count(p.attn_q_b_total)} params")
    print(f"  - kv_b_proj:    {format_count(p.attn_kv_b_total)} params")
    print(f"  - o_proj:       {format_count(p.attn_o_proj_total)} params")
    print(f"indexer_total:    {format_count(p.indexer_total)} params")
    print(f"dense_ffn_total:  {format_count(p.dense_ffn_total)} params")
    print(f"moe_expert_total: {format_count(p.moe_expert_total)} params")
    print(f"  - routed:       {format_count(p.moe_routed_expert_total)} params")
    print(f"  - shared:       {format_count(p.moe_shared_expert_total)} params")
    print(f"moe_router_total: {format_count(p.moe_router_total)} params")
    print(f"mhc_total:        {format_count(p.mhc_total)} params")
    print(f"TOTAL:            {format_count(p.total)} params")

    weight_gib = (p.total * weight_dtype_bytes) / GiB
    print("\n== Weight Memory ==")
    print(f"weights_unique: {weight_gib:.3f} GiB (dtype={weight_dtype})")

    per_rank_weight_gib = (sharded_weights.total * weight_dtype_bytes) / GiB
    print("\n== Weight Memory (per rank, approx) ==")
    print(f"weights_per_rank: {per_rank_weight_gib:.3f} GiB (dtype={weight_dtype})")

    print("\n== Weight Memory Breakdown (per rank, approx) ==")

    def _mem(x_params: float) -> float:
        return (x_params * weight_dtype_bytes) / GiB

    print(f"embedding:        {format_gib(_mem(sharded_weights.embedding))}")
    print(f"lm_head:          {format_gib(_mem(sharded_weights.lm_head))}")
    print(f"attn_fused_qkv_a: {format_gib(_mem(sharded_weights.attn_fused_qkv_a))}")
    print(f"attn_q_b:         {format_gib(_mem(sharded_weights.attn_q_b))}")
    print(f"attn_kv_b:        {format_gib(_mem(sharded_weights.attn_kv_b))}")
    print(f"attn_o_proj:      {format_gib(_mem(sharded_weights.attn_o_proj))}")
    print(f"indexer:          {format_gib(_mem(sharded_weights.indexer))}")
    print(f"dense_ffn:        {format_gib(_mem(sharded_weights.dense_ffn))}")
    print(f"moe_routed:       {format_gib(_mem(sharded_weights.moe_routed_expert))}")
    print(f"moe_shared:       {format_gib(_mem(sharded_weights.moe_shared_expert))}")
    print(f"moe_router:       {format_gib(_mem(sharded_weights.moe_router))}")
    print(f"mhc:              {format_gib(_mem(sharded_weights.mhc))}")

    print("\n== Activated Params / Token (compute metric) ==")
    for k, v in active.items():
        print(f"{k}: {format_count(v)} params")

    print("\n== KV Cache (MLA latent cache) ==")
    print(f"kv_dtype: {kv_dtype} ({kv_dtype_bytes} bytes)")
    kv_streams = int(mhc.branch_streams)
    if kv_streams > 1 or mhc.streams > 1:
        print(
            f"mHC kv/index scaling: branch_mode={mhc.branch_mode}, kv_streams={kv_streams} (batch*=kv_streams)"
        )
    if args.batch_scope == "global":
        global_batch = float(args.batch)
        local_batch = global_batch / float(max(1, parallel.dp))
    else:
        local_batch = float(args.batch)
        global_batch = local_batch * float(max(1, parallel.dp))

    print(
        "batch mapping: "
        f"scope={args.batch_scope}, global_batch≈{global_batch:.4f}, local_batch≈{local_batch:.4f}"
    )

    seq_estimates: List[SingleSeqEstimate] = []
    for s in _iter_seq_lens(int(args.seq_len), bool(args.also_max_pos), d.max_position_embeddings):
        local_seq_len = int(math.ceil(s / float(max(1, parallel.sp))))

        unique_kv_gib = estimate_kv_cache_gib(
            d, batch=global_batch * float(kv_streams), seq_len=s, kv_dtype_bytes=kv_dtype_bytes
        )
        per_rank_kv_gib = estimate_kv_cache_gib(
            d,
            batch=local_batch * float(kv_streams),
            seq_len=local_seq_len,
            kv_dtype_bytes=kv_dtype_bytes,
        )
        tp_total_kv_gib = per_rank_kv_gib * float(max(1, parallel.tp))

        unique_idx_gib = estimate_nsa_index_cache_gib(
            d, batch=global_batch * float(kv_streams), seq_len=s
        )
        per_rank_idx_gib = estimate_nsa_index_cache_gib(
            d, batch=local_batch * float(kv_streams), seq_len=local_seq_len
        )
        tp_total_idx_gib = (
            None if per_rank_idx_gib is None else per_rank_idx_gib * float(max(1, parallel.tp))
        )

        if unique_idx_gib is None:
            print(
                f"seq_len={s} (local_seq_len≈{local_seq_len}): "
                f"kv_unique={unique_kv_gib:.3f} GiB, kv_per_rank={per_rank_kv_gib:.3f} GiB, kv_tp_total={tp_total_kv_gib:.3f} GiB"
            )
        else:
            assert per_rank_idx_gib is not None and tp_total_idx_gib is not None
            print(
                f"seq_len={s} (local_seq_len≈{local_seq_len}): "
                f"kv_unique={unique_kv_gib:.3f} GiB, kv_per_rank={per_rank_kv_gib:.3f} GiB, kv_tp_total={tp_total_kv_gib:.3f} GiB; "
                f"index_unique={unique_idx_gib:.3f} GiB, index_per_rank={per_rank_idx_gib:.3f} GiB, index_tp_total={tp_total_idx_gib:.3f} GiB"
            )

        # Activation working-set: prefill + decode
        local_chunk = int(max(1, int(args.chunked_prefill_size) // max(1, parallel.dp)))
        local_prefill_tokens = int(
            min(local_chunk, math.ceil(local_batch * local_seq_len))
        )
        local_decode_tokens = int(max(1, math.ceil(local_batch)))
        prefill_ws = estimate_activation_working_set(
            d,
            parallel,
            tokens=local_prefill_tokens,
            act_dtype_bytes=act_dtype_bytes,
            mhc_streams=mhc.streams,
            mhc_branch_streams=mhc.branch_streams,
        )
        decode_ws = estimate_activation_working_set(
            d,
            parallel,
            tokens=local_decode_tokens,
            act_dtype_bytes=act_dtype_bytes,
            mhc_streams=mhc.streams,
            mhc_branch_streams=mhc.branch_streams,
        )
        print("\n== Activation Working Set (heuristic, per rank) ==")
        print(
            f"act_dtype: {act_dtype} ({act_dtype_bytes} bytes), "
            f"mhc_streams={prefill_ws.mhc_streams}, mhc_branch_streams={prefill_ws.mhc_branch_streams} "
            f"(branch_mode={mhc.branch_mode})"
        )
        print(
            f"prefill(tokens≈{local_prefill_tokens}, local_heads={prefill_ws.local_heads}, "
            f"dense_inter_per_part={prefill_ws.dense_intermediate_per_partition}, "
            f"moe_inter_per_part={prefill_ws.intermediate_per_partition}): "
            f"hidden={prefill_ws.hidden_gib:.3f} GiB, attn={prefill_ws.attn_gib:.3f} GiB, "
            f"dense={prefill_ws.dense_gib:.3f} GiB, moe={prefill_ws.moe_gib:.3f} GiB, peak≈{prefill_ws.peak_gib:.3f} GiB"
        )
        print(
            f"decode(tokens≈{local_decode_tokens}, local_heads={decode_ws.local_heads}, "
            f"dense_inter_per_part={decode_ws.dense_intermediate_per_partition}, "
            f"moe_inter_per_part={decode_ws.intermediate_per_partition}): "
            f"hidden={decode_ws.hidden_gib:.3f} GiB, attn={decode_ws.attn_gib:.3f} GiB, "
            f"dense={decode_ws.dense_gib:.3f} GiB, moe={decode_ws.moe_gib:.3f} GiB, peak≈{decode_ws.peak_gib:.3f} GiB"
        )

        # Perf breakdown (optional)
        decode_step: Optional[StepTime] = None
        prefill_step: Optional[StepTime] = None
        hw_perf = HardwareSpec(
            nodes=max(1, int(args.nodes)),
            gpus_per_node=max(1, int(args.gpus_per_node)),
            gpu_mem_gib=float(args.gpu_mem_gib or 0.0),
            mem_fraction=1.0,
            misc_gib=0.0,
            gpu_tflops=args.gpu_tflops,
            gpu_hbm_gibps=args.gpu_hbm_gibps,
            intra_gibps=args.intra_gibps,
            inter_gibps=args.inter_gibps,
            intra_us=float(args.intra_us),
            inter_us=float(args.inter_us),
            flops_eff=float(args.flops_eff),
            hbm_eff=float(args.hbm_eff),
            comm_eff=float(args.comm_eff),
            gemm_m_saturation=int(args.gemm_m_saturation),
        )
        if _has_perf(hw_perf):
            print("\n== Perf Estimate (roofline, per step) ==")
            decode_step = estimate_step_time(
                d,
                parallel,
                hw=hw_perf,
                phase="decode",
                tokens_per_rank=local_decode_tokens,
                seq_len_local=local_seq_len,
                global_batch=int(math.ceil(global_batch)),
                weight_dtype_bytes=weight_dtype_bytes,
                act_dtype_bytes=act_dtype_bytes,
                kv_dtype_bytes=kv_dtype_bytes,
                dp_padding_mode=str(args.dp_padding_mode),
            )
            prefill_step = estimate_step_time(
                d,
                parallel,
                hw=hw_perf,
                phase="prefill",
                tokens_per_rank=local_prefill_tokens,
                seq_len_local=local_seq_len,
                global_batch=int(math.ceil(global_batch)),
                weight_dtype_bytes=weight_dtype_bytes,
                act_dtype_bytes=act_dtype_bytes,
                kv_dtype_bytes=kv_dtype_bytes,
                dp_padding_mode=str(args.dp_padding_mode),
            )
            print(
                f"decode:  {decode_step.ms:.2f} ms/step, throughput≈{(float(global_batch)/decode_step.total_s if decode_step.total_s>0 else 0.0):.2f} tok/s"
            )
            global_prefill_tokens = int(local_prefill_tokens * parallel.dp * parallel.sp)
            print(
                f"prefill: {prefill_step.ms:.2f} ms/step, throughput≈{(float(global_prefill_tokens)/prefill_step.total_s if prefill_step.total_s>0 else 0.0):.2f} tok/s"
            )

            def _print_top_ops(label: str, step: StepTime) -> None:
                print(f"\n-- top ops ({label}) --")
                for op in sorted(step.ops, key=lambda x: x.t_total_s, reverse=True)[:12]:
                    comm = (
                        f"{op.comm_kind}:{op.comm_group}[{op.comm_group_size}]"
                        if op.comm_kind
                        else ""
                    )
                    print(
                        f"{op.name:36s}  total={op.t_total_s*1e3:8.2f} ms  bound={op.bound:6s}  comm={comm}"
                    )

            _print_top_ops("decode", decode_step)
            _print_top_ops("prefill", prefill_step)

        seq_estimates.append(
            SingleSeqEstimate(
                seq_len=int(s),
                local_seq_len=int(local_seq_len),
                global_batch=float(global_batch),
                local_batch=float(local_batch),
                kv_unique_gib=float(unique_kv_gib),
                kv_per_rank_gib=float(per_rank_kv_gib),
                idx_unique_gib=None if unique_idx_gib is None else float(unique_idx_gib),
                idx_per_rank_gib=None if per_rank_idx_gib is None else float(per_rank_idx_gib),
                prefill_tokens_per_rank=int(local_prefill_tokens),
                decode_tokens_per_rank=int(local_decode_tokens),
                prefill_ws=prefill_ws,
                decode_ws=decode_ws,
                decode_step=decode_step,
                prefill_step=prefill_step,
            )
        )

    # --- optional file outputs (single-config) ---
    if want_report or want_plot:
        ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        weights_unique_gib = (p.total * weight_dtype_bytes) / GiB
        weights_per_rank_gib = (sharded_weights.total * weight_dtype_bytes) / GiB

        if want_plot and out_svg is not None:
            out_svg.parent.mkdir(parents=True, exist_ok=True)
            bars: List[Tuple[str, Dict[str, float]]] = []
            for est in seq_estimates:
                label = f"seq={est.seq_len} (local={est.local_seq_len})"
                bars.append(
                    (
                        label,
                        {
                            "weights": float(weights_per_rank_gib),
                            "kv": float(est.kv_per_rank_gib),
                            "index": float(est.idx_per_rank_gib or 0.0),
                            "act": float(est.prefill_ws.peak_gib),
                            "misc": float(misc_gib),
                        },
                    )
                )
            write_memory_breakdown_svg(
                out_svg,
                title=(
                    f"VRAM breakdown | tp={parallel.tp},dp={parallel.dp},sp={parallel.sp} "
                    f"| batch={args.batch} ({args.batch_scope})"
                ),
                bars=bars,
                budget_gib=budget_gib,
            )

        if want_report:
            out_md.parent.mkdir(parents=True, exist_ok=True)
            text = render_single_config_report_md(
                args=args,
                config_src=config_src,
                ts=ts,
                out_md=out_md,
                d=d,
                p=p,
                active=active,
                mhc=mhc,
                parallel=parallel,
                weight_dtype=weight_dtype,
                kv_dtype=kv_dtype,
                act_dtype=act_dtype,
                weight_dtype_bytes=weight_dtype_bytes,
                kv_dtype_bytes=kv_dtype_bytes,
                act_dtype_bytes=act_dtype_bytes,
                sharded_weights=sharded_weights,
                weights_unique_gib=float(weights_unique_gib),
                weights_per_rank_gib=float(weights_per_rank_gib),
                budget_gib=budget_gib,
                misc_gib=float(misc_gib),
                seq_estimates=seq_estimates,
                out_svg=out_svg if want_plot else None,
            )
            out_md.write_text(text)

            print("\n== Single-Config Output ==")
            print(f"report: {out_md}")
            if out_svg is not None and out_svg.exists():
                print(f"plot:   {out_svg}")


def main(argv: Optional[Sequence[str]] = None) -> None:
    parser = build_parser()
    args: Optional[argparse.Namespace] = None
    try:
        args = parser.parse_args(list(argv) if argv is not None else None)
        _run(args)
    except SystemExit:
        raise
    except json.JSONDecodeError as exc:
        cfg = getattr(args, "config", None)
        if cfg:
            _user_die(f"Failed to parse JSON in --config file {cfg!r}: {exc}")
        _user_die(f"Failed to parse JSON config: {exc}")
    except FileNotFoundError as exc:
        _user_die(f"Config file not found: {exc.filename}")
    except urllib.error.URLError as exc:
        model = getattr(args, "model", None)
        if model:
            _user_die(f"Failed to fetch HuggingFace config for --model {model!r}: {exc}")
        _user_die(f"Failed to fetch HuggingFace config: {exc}")
    except (PermissionError, OSError) as exc:
        _user_die(f"Failed to read config file: {exc}")
    except KeyError as exc:
        missing = exc.args[0] if exc.args else "<unknown>"
        _user_die(f"Config missing required field: {missing}")
    except ValueError as exc:
        _user_die(str(exc))


if __name__ == "__main__":
    main()
