"""
=============================================================================
CS349E · 第 16 讲 · 01_full_stack_ledger.py —— 一个 token 的全栈账与"优化阶梯"
=============================================================================

【这个脚本要回答的问题】
    Llama-3.1-70B 在 8 块 H100 上（张量并行 8）生成一个 token，时间花在哪里？每一讲的技术各管哪一项、能省多少？
      演示 1：decode 一步的时间分解（读权重 / 读 KV / 计算 / 张量并行的通信 / 内核启动），批 1 和批 64，对应到各讲。
      演示 2：优化阶梯：从"最朴素的部署"出发，逐个加上 CUDA Graphs、融合、连续批处理、FP8、投机解码，
              每个用户每秒几个 token、每块 GPU 每秒几个 token、每百万 token 的成本和能耗。
      演示 3：同一个模型的 prefill：为什么它的账完全不同。

【关于数字】roofline 模型 + 效率折扣（60% 峰值算力、80% 峰值带宽）+ 固定开销（eager 每个内核 4 µs，CUDA Graphs 后 1 µs——假设值）；
    每层的内核个数用第 12 讲的 graph_lab 数出来；all-reduce 用 P5 的 α–β 公式（NVLink 4）；投机解码用第 6 讲的式 1，
    草稿模型 Llama-3.1-8B 也在同样的 8 块 GPU 上，接受率 α = 0.8 是假设值；价格 2 美元 / GPU·小时、功耗 700 W / 块 是假设值。
【运行】python3 01_full_stack_ledger.py      （约 1 秒）
【依赖】numpy；本目录的 planner；../lecture01 的 console。
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from planner import H100, LLAMA31_70B, PRICES, Deploy, decode_kernels_per_layer, ladder_point, tp_step  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402

M, TP, CTX = LLAMA31_70B, 8, 4096


def ms(x):
    return f"{x * 1e3:.2f} ms"


def demo1_breakdown(k):
    print_header(f"演示 1：Llama-3.1-70B decode 一步（8 × H100，张量并行 8，上下文 {CTX}）的时间分解")
    rows = []
    for name, b, kern, launch in (("批 1，eager", 1, k["eager"], 4e-6), ("批 1，融合 + CUDA Graphs", 1, k["fused"], 1e-6),
                                  ("批 64，融合 + CUDA Graphs", 64, k["fused"], 1e-6)):
        s = tp_step(M, H100, TP, b, CTX, kernels=kern, launch=launch)
        rows.append([name, ms(s["weights"]), ms(s["kv"]), ms(s["compute"]), ms(s["comm"]), ms(s["launch"]), ms(s["t"]), s["bound"]])
    print_table(["配置", "读权重", "读 KV", "计算", "通信", "内核启动", "一步", "访存 vs 计算"], rows)
    print(f"每层的内核数（第 12 讲 graph_lab 数的 decode 层）：eager {k['eager']} 个，分解 + 融合 + FlashAttention {k['fused']} 个；共 {M.n_layers} 层。")
    print("\n这一步里的每一项，归哪一讲管：")
    print_table(["项", "由什么决定", "对策（讲次）"], [
        ["读权重", "参数量 × 每个数的字节 / 带宽", "低精度（3）、批处理摊薄（1、6）、更大带宽的硬件（8）"],
        ["读 KV", "批 × 上下文 × 每 token 的 KV", "GQA / MLA（2）、KV 量化（3）、分页与前缀缓存（7）"],
        ["计算", "FLOPs / 峰值", "Tensor Core（8、10）、FlashAttention（11）、低精度的算力（3）"],
        ["通信", "每层 2 次 all-reduce 的 α–β", "切分方式（4、P5）、计算和通信重叠"],
        ["内核启动", "层数 × 每层内核数 × 每个的开销", "融合（12）、CUDA Graphs / megakernel（11）"],
        ["排队与干扰（这里没算）", "到达过程、prefill / decode 混跑", "连续批处理（6）、分块 prefill / 解耦（5）"],
    ], aligns=["<", "<", "<"])
    b1 = tp_step(M, H100, TP, 1, CTX, kernels=k["eager"], launch=4e-6)
    print_observations([
        f"批 1、eager 时，内核启动（{ms(b1['launch'])}）比读权重（{ms(b1['weights'])}）还多：{M.n_layers} 层 × {k['eager']} 个内核 × 4 µs。"
        "第 11 讲的 megakernel、第 12 讲的融合和 CUDA Graphs 针对的就是这一项——小批量的 decode 里它是头号开销。",
        f"批 1 时，每块 GPU 每步要读 1/8 的权重（{M.matmul_params() * 2 / TP / 1e9:.1f} GB），计算只占读权重时间的不到 1%：典型的带宽受限（第 1 讲）。"
        "批 64 时读 KV 已经和读权重是同一个量级（64 × 4096 个 token 的 KV），计算也涨到 2 ms，但仍然是带宽受限。",
        "张量并行的通信每步约 1–2 ms：批 1 时消息很小，几乎全是 α（每次 all-reduce 几微秒的延迟 × 160 次）；批 64 时带宽项开始变大（第 4 讲）。",
    ])


def demo2_ladder(k):
    print_header("演示 2：优化阶梯（同一个模型、同样 8 块 H100，逐个加上各讲的技术）")
    steps = [
        Deploy("0. 朴素：批 1、BF16、eager", 1, 2, 2, k["eager"], 4e-6),
        Deploy("1. + CUDA Graphs（第 11、12 讲）", 1, 2, 2, k["eager"], 1e-6),
        Deploy("2. + 融合（第 12 讲）", 1, 2, 2, k["fused"], 1e-6),
        Deploy("3. + 连续批处理到批 64（第 6、7 讲）", 64, 2, 2, k["fused"], 1e-6),
        Deploy("4. + FP8 权重和 KV（第 3 讲）", 64, 1, 1, k["fused"], 1e-6),
        Deploy("5. + 投机解码 γ = 4（第 6 讲）", 64, 1, 1, k["fused"], 1e-6, (0.8, 4)),
    ]
    alt = Deploy("（对照）批 1 + FP8 + 投机解码", 1, 1, 1, k["fused"], 1e-6, (0.8, 4))
    rows, res = [], []
    for d in steps + [alt]:
        r = ladder_point(M, H100, TP, CTX, d)
        j_tok = 700 * TP / r["tok_s_gpu"] / TP          # 每 token 的能耗 = 总功率 / 总吞吐
        res.append(r)
        rows.append([d.name, f"{r['tok_s_user']:.0f}", f"{r['tok_s_gpu']:.0f}", f"{r['usd_per_mtok']:.2f}", f"{j_tok:.2f}"])
    print_table(["部署", "每个用户 token/s", "每块 GPU token/s", "美元 / 百万 token", "焦耳 / token"], rows)
    base, last = res[0], res[5]
    print_observations([
        f"从朴素到第 5 级，每块 GPU 的吞吐提高了 {last['tok_s_gpu'] / base['tok_s_gpu']:.0f} 倍，成本降到 1/{base['usd_per_mtok'] / last['usd_per_mtok']:.0f}。"
        "最大的一步是批处理（第 3 级）：decode 受限于读权重，批 64 让 64 个用户分摊同一次读取——这是第 1 讲 roofline 的直接推论。",
        "前几级（CUDA Graphs、融合）主要改善**单个用户的延迟**（每个用户每秒的 token 数从约 60 到约 120），吞吐的绝对值仍然很小；"
        "批处理提高了吞吐，但每个用户的速度反而下降（每步要读更多 KV）——延迟和吞吐的取舍（第 1、4 讲）。",
        "FP8 同时减少读权重和读 KV 的字节，把第 3 级失去的单用户速度补了回来；投机解码在批 64 时还有收益，但不如批 1 时大"
        f"（对照行：批 1 时每个用户 {res[6]['tok_s_user']:.0f} token/s）——批越大，验证 γ+1 个 token 越接近计算受限，\"免费的验证\"就越少（第 6 讲 §5）。",
        "这条阶梯上的每一级都只用了前面某一讲的模型。真实系统里它们会互相影响（显存是共享的、投机解码和批处理争算力），"
        "所以最后一步总是端到端地测（或者至少端到端地模拟，见 03 脚本）。",
    ])


def demo3_prefill(k):
    print_header("演示 3：同一个模型的 prefill（一次 2,048 个 token 的提示）")
    s = tp_step(M, H100, TP, 1, 0, n_tokens=2048, kernels=k["fused"], launch=1e-6)
    print_table(["读权重", "读 KV", "计算", "通信", "内核启动", "一次 prefill", "访存 vs 计算"],
                [[ms(s["weights"]), ms(s["kv"]), ms(s["compute"]), ms(s["comm"]), ms(s["launch"]), ms(s["t"]), s["bound"]]])
    print_observations([
        f"prefill 一次处理 2,048 个 token，权重只读一遍，运算量是 decode 一步（批 1）的 2,048 倍：计算 {ms(s['compute'])} 远大于读权重 "
        f"{ms(s['weights'])}，**计算受限**。decode 和 prefill 是同一个模型的两种\"性格\"（第 1、5 讲）。",
        f"通信（{ms(s['comm'])}）在 prefill 里也变大了：每次 all-reduce 的消息是 2,048 × 8,192 × 2 字节。"
        "这就是解耦（第 5 讲）的出发点：两种性格对硬件、并行方式、批大小的要求不同，放在一起会互相干扰。",
    ])


if __name__ == "__main__":
    kk = decode_kernels_per_layer()
    demo1_breakdown(kk)
    demo2_ladder(kk)
    demo3_prefill(kk)
    assert PRICES[H100.name] == 2.0