#!/usr/bin/env python3
"""Byte-exact VRAM budget for GLM-5.2 hybrid quant on 4x RTX PRO 6000 (96GB, SM120).

Model constants from zai-org/GLM-5.2 config.json (verified against the
lukealonso/GLM-5.2-NVFP4 safetensors index: predicted 450.99 GB vs actual
450.93 GB total_size).
"""
GB = 1e9
GiB = 2**30

# --- architecture ---
H = 6144                 # hidden_size
MOE_I = 2048             # moe_intermediate_size
DENSE_I = 12288          # intermediate_size (dense layers 0-2)
VOCAB = 154880
N_LAYERS = 78            # main layers
N_MTP = 1                # layer 78
N_DENSE = 3              # layers 0,1,2
N_EXPERT_LAYERS = 75 + 1 # layers 3..77 + MTP layer
N_EXPERTS = 256
KV_LORA = 512
QK_ROPE = 64
Q_LORA = 2048
N_HEADS = 64
QK_NOPE = 192
V_HEAD = 256
IDX_HEADS, IDX_HDIM = 32, 128
N_FULL_INDEXERS = 22     # 21 main (IndexShare freq=4) + 1 MTP
CTX = 1_048_576

# --- per-component parameter counts ---
expert_params_down = MOE_I * H            # 12.58M
expert_params_gateup = 2 * H * MOE_I      # 25.17M
expert_params = expert_params_down + expert_params_gateup   # 37.75M
layer_expert_params = N_EXPERTS * expert_params             # 9.664B

attn_params = (H*Q_LORA + Q_LORA*(N_HEADS*(QK_NOPE+QK_ROPE))
               + H*(KV_LORA+QK_ROPE) + KV_LORA*(N_HEADS*(QK_NOPE+V_HEAD))
               + (N_HEADS*V_HEAD)*H)      # 165.03M per layer
indexer_params = Q_LORA*IDX_HEADS*IDX_HDIM + H*IDX_HDIM + H*IDX_HEADS  # 9.37M
shared_expert_params = expert_params
dense_mlp_params = 3 * H * DENSE_I
gate_params = H * N_EXPERTS
embed_params = VOCAB * H

TOTAL_LAYER_ENTRIES = N_LAYERS + N_MTP    # 79

nonexpert = {
    'attention':      attn_params * TOTAL_LAYER_ENTRIES,
    'indexers':       indexer_params * N_FULL_INDEXERS,
    'dense_mlp':      dense_mlp_params * N_DENSE,
    'shared_experts': shared_expert_params * N_EXPERT_LAYERS,
    'gates_fp32':     gate_params * N_EXPERT_LAYERS,      # router kept fp32
    'embed+lm_head':  2 * embed_params,
    'mtp_extras':     2 * H * H + 4 * H,                  # eh_proj etc
}

# --- bytes/param for each quant recipe (routed experts) ---
def gguf_bpw(down_bpw, gateup_bpw):
    return (expert_params_down*down_bpw + expert_params_gateup*gateup_bpw) / expert_params

RECIPES = {
    # NVFP4: 4b weight + fp8 scale per 16 = 4.5 bpw (+ negligible per-tensor scales)
    'nvfp4':  4.5,
    # unsloth UD-Q3_K_XL donor: down=IQ4_XS(4.25), gate/up=IQ3_XXS(3.0625)
    'ud_q3':  gguf_bpw(4.25, 3.0625),
    # unsloth UD-Q2_K_XL donor: down=IQ3_XXS(3.0625), gate/up=IQ2_XS(2.3125)
    'ud_q2':  gguf_bpw(3.0625, 2.3125),
}

def layer_expert_bytes(recipe):
    return layer_expert_params * RECIPES[recipe] / 8

# --- KV cache at 1M ---
kv_latent_per_tok = (KV_LORA + QK_ROPE) * TOTAL_LAYER_ENTRIES        # values (MTP has KV too)
def kv_bytes(bytes_per_val, scale_overhead=0.02):
    return kv_latent_per_tok * bytes_per_val * CTX * (1 + scale_overhead)
idx_cache = N_FULL_INDEXERS * IDX_HDIM * 1 * CTX   # fp8 indexer k-cache

# --- GPU budget ---
GPU_MEM = 97887 * 2**20        # RTX PRO 6000 Blackwell, nvidia-smi reported
N_GPU = 4
RESERVE_PER_GPU = 6 * GiB      # cuda ctx + nccl + cudagraphs + activations + frag

def report(attn_dtype_bytes=2, n4=0, n3=0, n2=0, kv='fp8'):
    assert n4+n3+n2 == 75
    ne = dict(nonexpert)
    scale = {'attention': attn_dtype_bytes/2, 'shared_experts': attn_dtype_bytes/2,
             'dense_mlp': attn_dtype_bytes/2}
    ne_bytes = 0
    for k, p in ne.items():
        b = p * (4 if k=='gates_fp32' else 2)
        b *= scale.get(k, 1)
        ne_bytes += b
    expert_bytes = (n4*layer_expert_bytes('nvfp4') + n3*layer_expert_bytes('ud_q3')
                    + n2*layer_expert_bytes('ud_q2') + layer_expert_bytes('nvfp4'))  # +MTP nvfp4
    weights = ne_bytes + expert_bytes
    kvb = kv_bytes(1 if kv=='fp8' else 0.5) + idx_cache
    total_need = weights + kvb
    total_have = N_GPU * (GPU_MEM - RESERVE_PER_GPU)
    print(f'--- attn={16*attn_dtype_bytes//2}bit  experts: {n4}xNVFP4 {n3}xUD-Q3 {n2}xUD-Q2 (+MTP nvfp4), KV={kv} ---')
    print(f'  non-expert weights: {ne_bytes/GB:8.2f} GB')
    print(f'  expert weights:     {expert_bytes/GB:8.2f} GB   (avg {expert_bytes*8/ (N_EXPERT_LAYERS*layer_expert_params):.2f} bpw)')
    print(f'  KV latent @1M:      {kv_bytes(1 if kv=="fp8" else 0.5)/GB:8.2f} GB')
    print(f'  indexer cache @1M:  {idx_cache/GB:8.2f} GB')
    print(f'  TOTAL NEED:         {total_need/GB:8.2f} GB = {total_need/GiB:.1f} GiB')
    print(f'  TOTAL HAVE (4x{GPU_MEM/GiB:.1f}GiB - 4x{RESERVE_PER_GPU/GiB:.0f}GiB): {total_have/GB:.2f} GB = {total_have/GiB:.1f} GiB')
    margin = total_have - total_need
    print(f'  MARGIN:             {margin/GB:+8.2f} GB   {"FITS" if margin>0 else "DOES NOT FIT"}')
    # PP=4 per-stage worst case: stage holding embed+lm_head + its layers
    return margin

print('kv latent values/token:', kv_latent_per_tok, '->', kv_latent_per_tok, 'B/tok fp8,',
      f'{kv_bytes(1)/GB:.1f} GB @1M fp8')
print('nonexpert param count:', sum(nonexpert.values())/1e9, 'B params')
print('checkpoint check: nvfp4-all =',
      (sum(v*2 for v in nonexpert.values()) + 2*nonexpert['gates_fp32'] +
       76*layer_expert_bytes('nvfp4'))/GB, 'GB vs actual 450.93')
print()
print('=== v1: BF16 attention (as shipped in NVFP4 repo) ===')
report(2, 10, 30, 35)
report(2, 8, 26, 41)
report(2, 6, 22, 47)
print()
print('=== v2: FP8 attention/shared/dense (unsloth uses Q8_0 there too) ===')
report(1, 14, 40, 21)
report(1, 12, 36, 27)
report(1, 10, 30, 35)