# lab02_calculator.py —— H3与Qwen3.8 128B 显存及方案性价比计算器

def calc_model_size(param_billion, bits):
    """计算模型权重体积 (GB)"""
    return (param_billion * 1e9 * bits) / 8 / (1024**3)

def calc_kv_cache(context_len, batch_size=1):
    """粗略估算KV Cache占用 (GB)，按经验值折算"""
    # 假设 32K 上下文单请求约 1GB，按比例和并发数放大
    base_32k_gb = 1.0
    return (context_len / 32768) * base_32k_gb * batch_size

def calc_value(price_wan, speed_tok_s):
    """计算性价比 (tok/s/万元)"""
    return speed_tok_s / price_wan

if __name__ == '__main__':
    print("=== 1. 模型体积与显存需求推导 ===")
    params = 128 # 128B
    size_int4 = calc_model_size(params, 4)
    size_int8 = calc_model_size(params, 8)
    size_fp16 = calc_model_size(params, 16)
    
    print(f"128B 模型体积：")
    print(f"  INT4: {size_int4:.1f} GB")
    print(f"  INT8: {size_int8:.1f} GB")
    print(f"  FP16: {size_fp16:.1f} GB")
    
    kv_cache = calc_kv_cache(context_len=32768, batch_size=4)
    h3_size = 12.0
    print(f"\nKV Cache (32K上下文, 4并发): 约 {kv_cache:.1f} GB")
    print(f"MiniMax H3 (pruned INT8): {h3_size:.1f} GB")
    
    print("\n=== 2. 三方案性价比对比 ===")
    # 方案数据: (名称, 显存+内存, 出词速度, H3耗时, 总价, 性价比)
    plans = [
        {"name": "主推: 5090 32G + 128G DDR5", "speed": 20, "price": 8.54},
        {"name": "对照A: 128G 统一内存", "speed": 10, "price": 2.50},
        {"name": "对照B: 双卡 64G 工作站", "speed": 48, "price": 14.00}
    ]
    
    for p in plans:
        cpb = calc_value(p["price"], p["speed"])
        print(f"{p['name']}:")
        print(f"  价格: {p['price']} 万元 | 出词: {p['speed']} tok/s")
        print(f"  性价比: {cpb:.2f} tok/s/万元\n")

    print("=== 3. 预算分配核对 ===")
    total_budget = 8.54
    gpu_price = 5.24
    mem_price = 2.20
    print(f"总预算: {total_budget} 万")
    print(f"GPU占比: {(gpu_price/total_budget)*100:.1f}% (目标55-60%)")
    print(f"内存占比: {(mem_price/total_budget)*100:.1f}% (目标10%，因近期内存暴涨偏高)")