"""
=============================================================================
CS349E · 第 1 讲 · 05_codesign_blackwell_ultra.py —— 案例：用 roofline 读懂 Blackwell Ultra 的设计取舍
=============================================================================

【这个脚本要回答的问题】
    本讲的阅读材料之一是 NVIDIA 的博客 *Inside NVIDIA Blackwell Ultra*。厂商博客里满是"更快、更大、更多"，
    怎样读出它**为什么**这样设计？本脚本挑三个设计决定，用前面学的工具逐个检验：

      演示 1：代际趋势——算力、带宽、容量、互连各涨了多少？谁跑得快、谁掉了队？
      演示 2：NVFP4——"两级缩放的 4 位浮点"为什么比普通 FP4 准？用 numpy 真的量化一组权重看误差。
      演示 3：softmax 的指数运算——为什么 Blackwell Ultra 要把特殊函数单元（SFU）的指数吞吐翻倍？
              算一算注意力里"每做一次 exp 要配多少次矩阵乘法"，就会发现 Blackwell 的 Tensor Core 已经快到
              让 exp 成了瓶颈。

【怎么读输出】演示 1、3 是规格上的算术；演示 2 用固定随机种子生成数据，每次运行完全相同。

【运行】python3 05_codesign_blackwell_ultra.py      （约 1 秒）
【依赖】numpy（演示 2）。
"""

import numpy as np

from console import print_header, print_observations, print_table, fmt_num, fmt_bytes
from hw_specs import H100, B200, B300, RUBIN, BYTES_PER_ELEM


# ---------------------------------------------------------------------------
# 演示 1：代际趋势
# ---------------------------------------------------------------------------
def demo1_trends():
    """
    把 Hopper → Blackwell → Blackwell Ultra → Rubin（公告）四代的关键规格并排比较。

    【怎么比】每一代取它支持的"最低精度"的稠密算力（Hopper 是 FP8，之后是 NVFP4）——这是厂商宣传
        "代际提升"时的口径；同时列出 FP8 的算力做同口径对比。所有比值都相对 H100。
    【边界】Rubin 的 50 PFLOPS 是"up to"，未说明是否稠密；它没有 FP8 数字（"-"）。
    """
    print_header("演示 1：四代 NVIDIA 数据中心 GPU 的关键规格（相对 H100 的倍数）")
    gens = [
        (H100, "fp8", "fp8"),
        (B200, "nvfp4", "fp8"),
        (B300, "nvfp4", "fp8"),
        (RUBIN, "nvfp4", None),
    ]
    rows = []
    for chip, lowp, fp8 in gens:
        low = chip.peak_flops(lowp)
        f8 = chip.peak_flops(fp8) if fp8 else None
        rows.append([
            chip.name, f"{fmt_num(low)}FLOP/s ({lowp})", f"{low / H100.peak_flops('fp8'):.1f}×",
            f"{f8 / H100.peak_flops('fp8'):.1f}×" if f8 else "-",
            f"{chip.mem_bw / H100.mem_bw:.1f}×", f"{chip.mem_cap / H100.mem_cap:.1f}×",
            f"{chip.link_bw / H100.link_bw:.1f}×",
            f"{chip.power_w / H100.power_w:.1f}×" if chip.power_w else "-",
            f"{low / chip.mem_bw:,.0f}",
        ])
    print_table(["芯片", "最低精度稠密算力", "算力倍数", "FP8 倍数", "带宽倍数", "容量倍数", "NVLink 倍数",
                 "功耗倍数", "屋脊点 FLOP/B"], rows)
    print_observations([
        "按厂商口径（各代最低精度），Blackwell Ultra 的算力是 H100 的 7.6 倍（博客写 7.5 倍）；但同为 FP8 只有 2.5 倍——"
        "一半以上的'提升'来自换了更低的数值格式。读厂商数字时，先确认比较的是不是同一种精度。",
        "带宽只涨了 2.4 倍，容量涨了 3.6 倍，NVLink 涨了 2 倍：算力增长远快于带宽和互连。屋脊点从 591 涨到 1,875——"
        "芯片越来越'算力过剩、带宽不足'，这就是常说的'内存墙'（memory wall）。",
        "Ma & Patterson（2026）给过一个更长期的数字：2012–2022 年 NVIDIA GPU 的 64 位算力涨了 80 倍，带宽只涨了 17 倍。",
        "Rubin 换成 HBM4，带宽一下涨到 22 TB/s（Blackwell 的 2.75 倍）：当 decode 受限于带宽，下一代最该补的就是带宽。",
        "功耗也在涨（700 W → 1400 W）：每块芯片更强，但每瓦的产出（tok/s/MW）才是数据中心真正在乎的（04 脚本演示 5）。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：NVFP4 的两级缩放
# ---------------------------------------------------------------------------
E2M1_GRID = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])  # FP4（E2M1）能表示的非负值


def round_e4m3(x):
    """
    把非负数组 x 舍入到 FP8 E4M3 能表示的最近的数（NVFP4 的块缩放因子用这个格式存）。

    【E4M3 是什么】1 位符号 + 4 位指数 + 3 位尾数。正规数是 (1 + m/8) × 2^e，e ∈ [-6, 8]，最大 448；
        比 2^-6 更小的数是"次正规数"，间隔固定为 2^-9。也就是说：在每个 2 的幂区间 [2^e, 2^(e+1)) 里，
        能表示的数是等间隔的 8 个，间隔 2^(e-3)。
    【做法】先算出 x 所在区间的指数 e（小于 -6 时按 -6 处理，即次正规区），间隔 step = 2^(e-3)，
        再把 x / step 四舍五入到整数（np.round 是"四舍六入五成双"，和硬件的 round-to-nearest-even 一致）。
    【边界】超过 448 的截到 448；0 保持 0。
    【Python 小知识】np.maximum(x, 2**-9) 逐元素取较大值，避免对 0 取 log2 得到 -inf。
    """
    x = np.minimum(np.asarray(x, dtype=np.float64), 448.0)
    e = np.floor(np.log2(np.maximum(x, 2.0 ** -9)))
    e = np.maximum(e, -6.0)
    step = 2.0 ** (e - 3)
    return np.minimum(np.round(x / step) * step, 448.0)


def to_e2m1(y):
    """
    把数组 y（已经除过缩放因子）舍入到最近的 E2M1 值，保留符号；绝对值超过 6 的截到 6。

    【做法】对 |y| 和 8 个网格点逐一求距离，取最近的那个（数组不大时这样写最清楚）。
    【Python 小知识】y[..., None] - grid 利用广播（broadcasting）一次算出每个元素到每个网格点的距离，
        得到形状 (..., 8) 的数组；argmin(axis=-1) 沿最后一维取最小值的下标。
    """
    a = np.minimum(np.abs(y), 6.0)
    idx = np.abs(a[..., None] - E2M1_GRID).argmin(axis=-1)
    return np.sign(y) * E2M1_GRID[idx]


def quantize(x, scheme):
    """
    用四种方案把一维数组 x 量化再反量化，返回近似值（和 x 同形状）。

    【方案】
        "fp8_tensor"  ：FP8 E4M3，整个张量一个缩放因子 s = max|x| / 448（对照组）
        "fp4_tensor"  ：FP4 E2M1，整个张量一个缩放因子 s = max|x| / 6（"标准 FP4"）
        "mxfp4"       ：每 32 个数一块，块缩放因子是 2 的幂（E8M0）：2^(floor(log2(块内 max)) − 2)
                        （OCP MX 规范；块内最大值除以缩放后可能略超 6，按规范截断到 6）
        "nvfp4"       ：每 16 个数一块，块缩放因子用 FP8 E4M3 存；外加一个 FP32 的张量级缩放 g，
                        把块缩放因子挪到 E4M3 的表示范围里：g = max|x| / (6 × 448)，
                        块因子 s_b = round_e4m3(块内 max / 6 / g)，实际缩放 = s_b × g
    【边界】x 的长度必须能被块大小整除；全零的块缩放因子为 0 时，直接输出 0。
    """
    x = np.asarray(x, dtype=np.float64)
    amax = np.abs(x).max()
    if scheme == "fp8_tensor":
        s = amax / 448.0
        return np.sign(x) * round_e4m3(np.abs(x) / s) * s
    if scheme == "fp4_tensor":
        s = amax / 6.0
        return to_e2m1(x / s) * s
    if scheme == "mxfp4":
        blocks = x.reshape(-1, 32)
        bmax = np.abs(blocks).max(axis=1, keepdims=True)
        s = 2.0 ** (np.floor(np.log2(np.maximum(bmax, 1e-30))) - 2)
        return (to_e2m1(blocks / s) * s).reshape(x.shape)
    if scheme == "nvfp4":
        g = amax / (6.0 * 448.0)
        blocks = x.reshape(-1, 16)
        bmax = np.abs(blocks).max(axis=1, keepdims=True)
        sb = round_e4m3(bmax / 6.0 / g)
        scale = sb * g
        safe = np.where(scale > 0, scale, 1.0)
        q = np.where(scale > 0, to_e2m1(blocks / safe) * scale, 0.0)
        return q.reshape(x.shape)
    raise ValueError(f"未知方案 {scheme}")


def demo2_nvfp4(n=1 << 20, seed=0):
    """
    在一组"像权重"的随机数上比较四种量化方案的误差。

    【数据】自由度为 5 的 Student-t 分布：大部分数集中在 0 附近，但有比高斯分布更多的"大值"（重尾）——
        真实模型的权重和激活都有这种离群值（outlier），第 3 讲（LLM.int8）会专门讲。
    【指标】相对误差 = ||q − x|| / ||x||；SQNR（信号量化噪声比）= 20·log10(1 / 相对误差)，单位 dB，越大越好。
    【每个数平均几比特】fp4 = 4；mxfp4 = 4 + 8/32 = 4.25；nvfp4 = 4 + 8/16 = 4.5（张量级 FP32 因子可忽略）；fp8 = 8。
    """
    print_header("演示 2：NVFP4 的两级缩放——同样 4 位，误差差多少？（真的量化 100 万个数）")
    rng = np.random.default_rng(seed)
    x = rng.standard_t(df=5, size=n) * 0.02
    schemes = [
        ("fp8_tensor", "FP8 E4M3，整个张量一个缩放", 8.0),
        ("fp4_tensor", "FP4 E2M1，整个张量一个缩放", 4.0),
        ("mxfp4", "MXFP4：每 32 个数一个 2 的幂缩放", 4.25),
        ("nvfp4", "NVFP4：每 16 个数一个 FP8 缩放 + 张量级 FP32", 4.5),
    ]
    rows, errs = [], {}
    for key, label, bits in schemes:
        q = quantize(x, key)
        rel = np.linalg.norm(q - x) / np.linalg.norm(x)
        errs[key] = rel
        zero_frac = np.mean((q == 0) & (x != 0))
        rows.append([label, f"{bits:g}", f"{rel * 100:.2f}%", f"{20 * np.log10(1 / rel):.1f}", f"{zero_frac * 100:.1f}%"])
    print_table(["方案", "比特/数", "相对误差", "SQNR (dB)", "被量化成 0 的比例"], rows)

    w300 = 300e9 * BYTES_PER_ELEM["nvfp4"]
    print(f"\n存储账：3000 亿参数 × 4.5 bit = {fmt_bytes(w300)}，放进 288 GB 后还剩 {fmt_bytes(288e9 - w300)} 给 KV cache；"
          f"288 GB 最多能放 {288e9 / BYTES_PER_ELEM['nvfp4'] / 1e9:.0f}B 参数（不留 KV）。")
    print_observations([
        f"整个张量只用一个缩放因子时，FP4 的相对误差是 {errs['fp4_tensor'] * 100:.0f}%：为了照顾少数离群的大值，"
        "缩放因子被拉得很大，大部分'普通'的数都被量化成了 0。",
        f"把缩放因子缩小到每 16 个数一个（NVFP4），误差降到 {errs['nvfp4'] * 100:.1f}%——每一小块都按自己的最大值缩放，"
        "离群值只'连累'同一块里的 15 个邻居。代价是每个数多 0.5 bit。",
        f"MXFP4 的块更大（32）、缩放因子只能是 2 的幂（最多浪费将近一半的表示范围），误差 {errs['mxfp4'] * 100:.1f}%，"
        "介于两者之间。NVIDIA 说 NVFP4 'markedly lower error than standard FP4'，在这组数据上可以复现这个方向。",
        f"作为对照，8 位的 FP8 误差只有 {errs['fp8_tensor'] * 100:.1f}%。博客说 NVFP4 的模型精度'通常与 FP8 相差不到约 1%'——"
        "那是**模型任务精度**，不是这里的逐元素误差；两者的关系要靠第 3 讲的实验来理解。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：softmax 的指数瓶颈
# ---------------------------------------------------------------------------
def demo3_exp_bottleneck():
    """
    注意力里每个分数 s = q·k 都要做一次 exp（softmax），同时伴随 QKᵀ 和 PV 两次长度 d 的点积。

    【每次 exp 配多少 FLOP】QKᵀ：每个分数 2·d_qk 次运算；PV：每个分数 2·d_v 次运算。
        所以 flops_per_exp = 2·d_qk + 2·d_v（普通注意力 d_qk = d_v = d_head，就是 4·d_head）。
    【判断谁是瓶颈】Tensor Core 做完这些矩阵乘法要 t_mma = flops_per_exp / peak（每个分数）；
        SFU 做一次 exp 要 t_exp = 1 / exp_per_s。比值 t_exp / t_mma > 1 就说明 exp 单元跟不上 Tensor Core
        （除非能把 exp 挪到别的单元上去做，见观察）。
    【边界】只看"exp 吞吐 vs 矩阵乘法吞吐"这一对；softmax 还需要求最大值、减法、乘法等普通运算，
        这里不计，所以结论是偏乐观的。
    """
    print_header("演示 3：softmax 的指数运算会不会拖慢注意力？（t_exp / t_mma > 100% 表示 exp 是瓶颈）")
    shapes = [("d_head = 64", 64, 64), ("d_head = 128", 128, 128), ("MLA 式（d_qk=192, d_v=128）", 192, 128)]
    rows = []
    for chip in (H100, B200, B300):
        for prec in ("bf16", "fp8"):
            peak = chip.peak_flops(prec)
            cells = [chip.name, prec, f"{chip.exp_per_s / 1e12:.1f} T/s"]
            for _, dqk, dv in shapes:
                fpe = 2 * dqk + 2 * dv
                need = peak / fpe                  # 让 Tensor Core 满速需要的 exp/s
                cells.append(f"{need / chip.exp_per_s * 100:.0f}%")
            rows.append(cells)
    print_table(["芯片", "注意力矩阵乘法精度", "SFU exp 吞吐"] + [s[0] for s in shapes], rows)
    print_observations([
        "Hopper 上 exp 基本不是瓶颈：d_head = 128 时，BF16 下 exp 只需要矩阵乘法时间的 43%，FP8 下 86%。",
        "Blackwell 的 Tensor Core 比 Hopper 快了约 2.5 倍，exp 吞吐却只从 4.5 涨到 5 T/s：FP8 注意力、d_head = 128 时，"
        "exp 需要的时间是矩阵乘法的约 2 倍——注意力的速度被 SFU 卡住了，再快的 Tensor Core 也没用（Amdahl 定律）。",
        "Blackwell Ultra 把 exp 吞吐翻倍到 10.7 T/s，正好把 d_head = 128 的比值拉回 100% 以下。这就是博客里"
        "'accelerated softmax'的来历：一个看似不起眼的单元，因为别的单元变快而成了瓶颈。",
        "软件也在想办法：FlashAttention-4 等内核把一部分 exp 改用普通 FMA 单元上的多项式来近似计算，把负载从 SFU 分出去；"
        "Rubin 则继续提高 exp 吞吐（博客：FP32 2 倍、BF16 4 倍于 Blackwell）。硬件和软件在同一个瓶颈上一起使劲——这就是协同设计。",
        "d_head 越小，每次 exp 配的矩阵乘法越少，越容易被 exp 卡住（d_head = 64 在三代芯片的 FP8 下都超过 100%）："
        "模型结构的一个'小参数'，直接决定了硬件的哪个单元成为瓶颈。",
    ])


if __name__ == "__main__":
    demo1_trends()
    demo2_nvfp4()
    demo3_exp_bottleneck()