"""
=============================================================================
CS349E · 前置讲 P3 · 05_quantization_basics.py —— 量化入门：缩放因子、零点、截断、粒度与块缩放格式
=============================================================================

【这个脚本要回答的问题】
    量化（quantization）= 用一个缩放因子把一组实数"挪"到低精度格式的表示范围里，再舍入。
    第 1 讲已经看到 NVFP4 的块缩放比整张量缩放好得多；本脚本把背后的几个基本问题拆开来看：

      演示 1：仿射量化：scale 和 zero point；对称 vs 非对称；以及"只用整数"怎样算矩阵乘法（再乘回缩放因子）。
      演示 2：截断（clipping）：故意让最大的几个数溢出，换来其余数更精细——最优截断点在哪里？
      演示 3：粒度：整张量 / 每通道 / 每组。权重的"每行尺度不同"和激活的"离群通道"分别需要哪种粒度？
      演示 4：块缩放浮点格式：MXFP8 / MXFP6 / MXFP4 / NVFP4 / INT4 分组——同样的比特，误差差多少？
              NVFP4 比 MXFP4 好，是因为块更小，还是因为缩放因子更精细？
      演示 5：逐元素误差和矩阵乘法输出误差是什么关系？以及"每多一位，SQNR 多约 6 dB"。

【怎么读输出】全部是确定性的计算（固定种子），每次运行结果相同。数据都是合成的（高斯、拉普拉斯、Student-t、
    人为加入的"离群通道"），用来展示机制；真实模型上的效果是第 3 讲的内容。

【运行】python3 05_quantization_basics.py      （约 3 秒）
【依赖】numpy；从 ../lecture01 导入 console；同目录的 numfmt。
"""

import os
import sys

import numpy as np

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

from console import print_header, print_observations, print_table  # noqa: E402
from numfmt import (FP4_E2M1, FP6_E2M3, FP6_E3M2, FP8_E4M3, mx_quantize,  # noqa: E402
                    nvfp4_quantize, quantize_int, rel_error, sqnr_db, tensor_scaled)


# ---------------------------------------------------------------------------
# 演示 1：仿射量化与整数矩阵乘法
# ---------------------------------------------------------------------------
def demo1_affine_and_integer_matmul(seed=0):
    """
    第一部分：ReLU 之后的激活（全部 ≥ 0），用 8 位对称（int8，zero = 0）和 8 位非对称（uint8，带 zero point）量化。
    第二部分：y = A·Wᵀ，A 用非对称 uint8、W 用对称 int8。整数流水线：
        y_ij = s_a·s_w · [ Σ_k qa_ik·qw_jk  −  z_a · Σ_k qw_jk ]
        方括号里全是整数运算（累加到 int32 就够，见 01 脚本演示 3），第二项每个输出通道只算一次。
        最后乘一次浮点的 s_a·s_w。验证它和"先反量化再用浮点乘"完全相同。

    【为什么权重用对称、激活用非对称】权重以 0 为中心，对称最简单（zero = 0，第二项消失）；
        激活常常偏向一侧（ReLU、SiLU 之后），非对称能用满 256 个编码。代价是多一项修正。
    """
    print_header("演示 1：仿射量化——scale、zero point，以及只用整数的矩阵乘法")
    rng = np.random.default_rng(seed)
    a = np.maximum(rng.standard_normal(100_000), 0.0)            # ReLU 之后：一半是 0，一半是正数
    rows = []
    for sym, label in ((True, "对称 int8（zero = 0）"), (False, "非对称 uint8（带 zero point）")):
        d, q, s, z = quantize_int(a, 8, symmetric=sym)
        used = len(np.unique(q))
        rows.append([label, f"{float(s.ravel()[0]):.5f}", int(z.ravel()[0]), f"{q.min()} … {q.max()}", used,
                     f"{rel_error(d, a) * 100:.3f}%"])
    print_table(["方式", "scale", "zero point", "编码范围", "实际用到几个编码", "相对误差"], rows)

    B, D, F = 64, 1024, 256
    A = np.maximum(rng.standard_normal((B, D)), 0.0)
    W = rng.standard_normal((F, D)) * 0.05
    _, qa, sa, za = quantize_int(A, 8, symmetric=False)
    _, qw, sw, _ = quantize_int(W, 8, symmetric=True, granularity="channel")
    sa, za = float(sa.ravel()[0]), float(za.ravel()[0])
    acc = qa.astype(np.int64) @ qw.T.astype(np.int64)             # 纯整数累加
    corr = za * qw.sum(axis=1)                                     # 每个输出通道一个修正项（也是整数）
    y_int = sa * sw.ravel()[None, :] * (acc - corr[None, :])
    deqA = sa * (qa - za)
    deqW = sw * qw
    y_deq = deqA @ deqW.T
    y_true = A @ W.T
    print(f"\n整数流水线：|A·Wᵀ 的整数部分| 最大 {np.abs(acc).max():,}（int32 上限 {2 ** 31 - 1:,}）")
    print(f"  整数流水线 vs '反量化后浮点乘'：最大差 {np.abs(y_int - y_deq).max():.2e}（只有浮点舍入误差）")
    print(f"  量化后的结果 vs 原始浮点结果：相对误差 {rel_error(y_int, y_true) * 100:.3f}%")
    print_observations([
        "对称量化把一半的编码（负数）浪费在永远不会出现的值上：ReLU 输出只能用 0 … 127（实际用到 107 个）；"
        "非对称量化能用 0 … 255（实际用到 203 个），scale 减半，相对误差也小约一半——相当于多了 1 位。"
        "（这组数的最小值正好是 0，所以 zero point 也是 0；数据范围不含 0 时 zero point 就不为 0 了。）",
        "整数流水线和'先反量化再浮点乘'的结果一致到浮点舍入误差：乘加全部在整数里做，最后只乘一次缩放因子。"
        "这就是 INT8 推理的标准做法——Tensor Core 的 INT8 算力是 BF16 的 2 倍，功耗也低得多（第 1 讲 §2）。",
        "zero point 带来的修正项 z_a·Σ_k qw_jk 只和权重有关，可以提前算好；如果两边都用非对称，还会多出更多交叉项，"
        "所以实践中常见'权重对称、激活非对称'或者干脆都用对称。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：截断
# ---------------------------------------------------------------------------
def demo2_clipping(n=1 << 18, seed=0):
    """
    INT4 对称量化（编码 −7 … 7），把范围从 max|x| 缩到 c × max|x|：超过的数被截断，其余的数间隔变小。

    【两种误差的拉锯】
        舍入误差 ∝ 间隔² ∝ (c·max)²：c 越小越好；
        截断误差：超过 c·max 的数被截到边界，c 越小越多。
        最优 c 取决于分布的"尾巴"：尾巴越重，最大值离大多数数越远，越值得截断。
    【数据】高斯（轻尾）、拉普拉斯（指数尾）、Student-t（自由度 5，重尾，像权重）。
    """
    print_header("演示 2：截断——INT4 对称量化，范围取 c × max|x| 时的相对误差")
    rng = np.random.default_rng(seed)
    dists = {"高斯": rng.standard_normal(n), "拉普拉斯": rng.laplace(size=n), "Student-t（ν = 5）": rng.standard_t(5, n)}
    cs = (1.0, 0.8, 0.6, 0.5, 0.4, 0.3, 0.25, 0.2, 0.15, 0.1, 0.07, 0.05)
    rows, best = [], {}
    for name, z in dists.items():
        errs = [rel_error(quantize_int(z, 4, clip=c)[0], z) for c in cs]
        i = int(np.argmin(errs))
        best[name] = (cs[i], errs[i], errs[0], cs[i] * np.abs(z).max() / z.std())
        rows.append([name, f"{np.abs(z).max() / z.std():.1f}"] + [f"{e * 100:.1f}" for e in errs])
    print_table(["分布", "max/σ"] + [f"c={c:g}" for c in cs], rows)
    print("（表中是相对误差，单位 %）")
    lines = [f"{k}：最优 c = {v[0]:g}（截断点 ≈ {v[3]:.1f}σ），误差从 {v[2] * 100:.1f}% 降到 {v[1] * 100:.1f}%"
             for k, v in best.items()]
    print_observations(lines + [
        "尾巴越重，最大值离'大多数数'越远（Student-t 的 max/σ 比高斯大得多），按最大值缩放就越浪费，截断的收益越大。",
        "但截断点不能太小：c 过小时被截断的数太多，误差又迅速上升。这个 U 形曲线是所有'选缩放因子'的方法"
        "（MSE 搜索、百分位数校准、GPTQ / AWQ 的缩放搜索）背后的共同权衡。",
        "截断扔掉的恰恰是最大的那些数——如果它们很重要（LLM 激活里的离群特征），截断就会出事。"
        "这引出了另一条路：不截断，而是把离群值单独处理（LLM.int8，第 3 讲）或者缩小缩放的粒度（演示 3、4）。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：粒度
# ---------------------------------------------------------------------------
def demo3_granularity(seed=1):
    """
    两个合成矩阵：
        W（权重，1024 × 4096）：每一行（输出通道）有自己的尺度（对数正态，σ = 0.5）；
        A（激活，256 个 token × 4096 个通道）：随机挑 8 个通道，放大 60 倍——模仿 LLM 激活里的"离群通道"
            （它们在**每个 token** 上都很大，LLM.int8 论文的观察）。
    粒度：整张量、每行（W 的每个输出通道 / A 的每个 token）、每 128 个、每 32 个（沿最后一维分组）。
    "每元素额外比特"：每组一个 FP16 的缩放因子，摊到组内每个元素上 = 16 / 组大小。
    """
    print_header("演示 3：粒度——整张量 / 每行 / 每组，INT8 与 INT4 的相对误差")
    rng = np.random.default_rng(seed)
    F, D, T = 1024, 4096, 256
    W = rng.standard_normal((F, D)) * 0.02 * np.exp(rng.normal(0, 0.5, size=(F, 1)))
    A = rng.standard_normal((T, D))
    out_ch = rng.choice(D, 8, replace=False)
    A[:, out_ch] *= 60
    grans = [("整张量", "tensor", None, 0.0), ("每行", "channel", None, 16 / D), ("每 128 个一组", "group", 128, 16 / 128),
             ("每 32 个一组", "group", 32, 16 / 32)]
    rows = []
    res = {}
    for mname, M in (("W：每行尺度不同", W), ("A：8 个离群通道", A)):
        for bits in (8, 4):
            cells = [mname, f"INT{bits}"]
            for gname, g, size, extra in grans:
                e = rel_error(quantize_int(M, bits, granularity=g, group=size)[0], M)
                res[(mname[0], bits, gname)] = e
                cells.append(f"{e * 100:.2f}%")
            rows.append(cells)
    print_table(["矩阵", "位数"] + [f"{g[0]}（+{g[3]:.3g} bit）" for g in grans], rows)
    print_observations([
        f"权重 W：按行（输出通道）缩放把 INT4 的误差从 {res[('W', 4, '整张量')] * 100:.0f}% 降到 "
        f"{res[('W', 4, '每行')] * 100:.0f}%——每一行的尺度不同，整张量的缩放因子只照顾得了最大的那一行。"
        "per-channel 几乎不增加开销（每行一个缩放因子），是权重量化的默认做法。",
        f"激活 A：按 token（每行）缩放几乎没用（INT4 {res[('A', 4, '整张量')] * 100:.0f}% → {res[('A', 4, '每行')] * 100:.0f}%），"
        "因为离群通道在每个 token 上都很大，每一行的最大值都被它决定。沿通道方向分组才有效："
        f"每 128 个一组 {res[('A', 4, '每 128 个一组')] * 100:.0f}%、每 32 个 {res[('A', 4, '每 32 个一组')] * 100:.0f}%"
        "——离群通道只'连累'和它同组的通道。",
        "问题是：矩阵乘法 A·Wᵀ 沿通道方向（内维 k）求和，如果缩放因子也沿 k 变化，就不能在累加完以后再统一乘缩放因子，"
        "硬件得在累加过程中处理缩放（DeepSeek-V3 的每 128 个提升一次、MX / NVFP4 的硬件块缩放就是在解决这个问题）。"
        "LLM.int8 走的是另一条路：把离群通道拆出去用 FP16 单独算（第 3 讲）。",
        "组越小，误差越小，但缩放因子的开销越大：每 32 个一个 FP16 缩放因子，相当于每个数多 0.5 bit。",
    ])


# ---------------------------------------------------------------------------
# 演示 4：块缩放浮点格式
# ---------------------------------------------------------------------------
def block_scaled(x, elem, block, scale_kind):
    """
    通用的块缩放：每 block 个数一个缩放因子，scale_kind 决定缩放因子的格式：
        "e8m0"：2 的幂（OCP MX 规范的算法，见 numfmt.mx_quantize）；
        "e4m3"：FP8 E4M3 + 张量级 FP32（NVFP4 的做法，见 numfmt.nvfp4_quantize）；
        "fp32"：精确的缩放因子 = 块内最大值 / 元素最大值（理想情况，用来看"缩放因子的精度"值多少）。
    返回 (近似值, 被截断到元素最大值的比例)。
    """
    xb = x.reshape(-1, block)
    amax = np.abs(xb).max(axis=1, keepdims=True)
    if scale_kind == "e8m0":
        _, k = np.frexp(np.where(amax > 0, amax, 1.0))
        X = np.ldexp(1.0, k - 1 - elem.emax)
    elif scale_kind == "e4m3":
        g = np.abs(x).max() / (elem.max_value * FP8_E4M3.max_value)
        X = FP8_E4M3.round(amax / elem.max_value / g, saturate=True) * g
    else:
        X = amax / elem.max_value
    X = np.where(X > 0, X, 1.0)
    y = xb / X
    clipped = np.mean(np.abs(y) > elem.max_value)
    return (elem.round(y, saturate=True) * X).ravel(), clipped


def demo4_block_formats(n=1 << 20, seed=0):
    """
    在第 1 讲同样的数据（100 万个 Student-t(5) × 0.02，同样的种子）上比较各种格式。

    【每元素比特】元素位数 + 缩放因子位数 / 块大小（张量级的 FP32 因子忽略不计）。
    【第二张表：拆开 NVFP4 vs MXFP4 的两个差别】块大小（16 vs 32）和缩放因子格式（E4M3 vs E8M0）。
    """
    print_header("演示 4：块缩放格式——同一组'像权重'的数（与第 1 讲相同的数据）")
    rng = np.random.default_rng(seed)
    x = rng.standard_t(df=5, size=n) * 0.02
    rows = []
    entries = [
        ("FP8 E4M3，整张量缩放", 8.0, tensor_scaled(x, FP8_E4M3)),
        ("INT8，整张量（对称）", 8.0, quantize_int(x, 8)[0]),
        ("MXFP8（E4M3，每 32 个，E8M0）", 8.25, mx_quantize(x, FP8_E4M3)[0]),
        ("MXFP6（E2M3，每 32 个，E8M0）", 6.25, mx_quantize(x, FP6_E2M3)[0]),
        ("MXFP6（E3M2，每 32 个，E8M0）", 6.25, mx_quantize(x, FP6_E3M2)[0]),
        ("INT4，每 32 个一组（FP16 缩放）", 4.5, quantize_int(x, 4, granularity="group", group=32)[0]),
        ("MXFP4（E2M1，每 32 个，E8M0）", 4.25, mx_quantize(x, FP4_E2M1)[0]),
        ("NVFP4（E2M1，每 16 个，E4M3 + FP32）", 4.5, nvfp4_quantize(x)[0]),
        ("FP4 E2M1，整张量缩放", 4.0, tensor_scaled(x, FP4_E2M1)),
    ]
    errs = {}
    for label, bits, q in entries:
        errs[label] = rel_error(q, x)
        rows.append([label, f"{bits:g}", f"{rel_error(q, x) * 100:.2f}%", f"{sqnr_db(q, x):.1f}"])
    print_table(["格式", "每元素比特", "相对误差", "SQNR (dB)"], rows)

    _, clip8 = block_scaled(x, FP8_E4M3, 32, "e8m0")
    print(f"\nMXFP8 的块内最大值除以 2 的幂缩放后落在 [256, 512)，超过 448 的被截断：占全部元素的 {clip8 * 100:.3f}%。")

    print("\n拆开 NVFP4 和 MXFP4 的两个差别（元素都是 FP4 E2M1）：")
    rows = []
    abl = {}
    for block in (32, 16):
        cells = [f"每 {block} 个"]
        for kind in ("e8m0", "e4m3", "fp32"):
            q, clipped = block_scaled(x, FP4_E2M1, block, kind)
            abl[(block, kind)] = rel_error(q, x)
            cells.append(f"{rel_error(q, x) * 100:.2f}%（截断 {clipped * 100:.2f}%）")
        rows.append(cells)
    print_table(["块大小", "E8M0 缩放（2 的幂）", "E4M3 缩放（NVFP4）", "精确缩放（理想）"], rows)
    print_observations([
        f"4.5 bit 的预算下，NVFP4（{errs['NVFP4（E2M1，每 16 个，E4M3 + FP32）'] * 100:.1f}%）比 INT4 每 32 个一组"
        f"（{errs['INT4，每 32 个一组（FP16 缩放）'] * 100:.1f}%）更准：重尾数据里，浮点元素的'非均匀'刻度（0 附近密、远处疏）"
        "比整数的均匀刻度更合适。",
        f"缩放因子的格式比块大小更重要：同样每 32 个，E8M0 → E4M3 把误差从 {abl[(32, 'e8m0')] * 100:.1f}% 降到 "
        f"{abl[(32, 'e4m3')] * 100:.1f}%；同样用 E4M3，块从 32 缩到 16 只再降到 {abl[(16, 'e4m3')] * 100:.1f}%。"
        "2 的幂的缩放因子太粗：块内最大值除完以后可能落在 4–8 之间的任何地方，平均浪费了不少刻度，还要截断超过 6 的数。",
        f"E4M3 缩放已经非常接近'精确缩放'（{abl[(16, 'fp32')] * 100:.2f}%）：8 位的缩放因子足够了。",
        f"MXFP8（{errs['MXFP8（E4M3，每 32 个，E8M0）'] * 100:.2f}%）反而比整张量缩放的 FP8（{errs['FP8 E4M3，整张量缩放'] * 100:.2f}%）略差："
        "E4M3 自己的动态范围已经够覆盖这组数，块缩放帮不上忙，而 2 的幂缩放还会截断少数块内最大值。块缩放主要是为窄格式（FP4）准备的。",
        f"MXFP6 的 E2M3（{errs['MXFP6（E2M3，每 32 个，E8M0）'] * 100:.2f}%）比 E3M2（{errs['MXFP6（E3M2，每 32 个，E8M0）'] * 100:.2f}%）准得多，"
        "甚至和 MXFP8 差不多：块内的数大小相近，指数位用处不大，尾数位才决定精度。"
        "这和 DeepSeek-V3 '细粒度缩放后所有张量都用 E4M3（尾数优先）'是同一个道理。",
    ])


# ---------------------------------------------------------------------------
# 演示 5：从逐元素误差到输出误差
# ---------------------------------------------------------------------------
def demo5_error_propagation(seed=0):
    """
    第一部分：W（Student-t 权重，1024 × 4096）量化后，Y = W·X（X 是 64 列随机激活）的相对误差，
        和 W 本身的逐元素相对误差比较。
    第二部分：INT2 … INT8（每行缩放）量化高斯数据的 SQNR，看每多一位增加多少 dB。

    【预期】误差 E = Ŵ − W，输出误差 = E·X。当量化误差和激活不相关时，
        ||E·X|| / ||W·X|| ≈ ||E|| / ||W||——相对误差**原样传递**，既不会因为"点积里有正有负"而变小，也不会变大。
    """
    print_header("演示 5：逐元素误差 → 矩阵乘法的输出误差；每多一位多几 dB")
    rng = np.random.default_rng(seed)
    F, D, B = 1024, 4096, 64
    W = rng.standard_t(df=5, size=(F, D)) * 0.02
    X = rng.standard_normal((D, B))
    Y = W @ X
    rows = []
    for label, Wq in (("INT8 每行", quantize_int(W, 8, granularity="channel")[0]),
                      ("NVFP4", nvfp4_quantize(W.ravel())[0].reshape(W.shape)),
                      ("MXFP4", mx_quantize(W.ravel(), FP4_E2M1)[0].reshape(W.shape)),
                      ("INT4 每 128 个一组", quantize_int(W, 4, granularity="group", group=128)[0])):
        rows.append([label, f"{rel_error(Wq, W) * 100:.2f}%", f"{rel_error(Wq @ X, Y) * 100:.2f}%"])
    print_table(["权重格式", "权重的相对误差", "输出 Y = W·X 的相对误差"], rows)

    z = rng.standard_normal((256, 4096))
    rows = []
    prev = None
    for bits in (2, 3, 4, 5, 6, 7, 8):
        s = sqnr_db(quantize_int(z, bits, granularity="channel")[0], z)
        rows.append([f"INT{bits}", f"{s:.1f}", "—" if prev is None else f"+{s - prev:.1f}"])
        prev = s
    print()
    print_table(["位数（每行缩放，高斯数据）", "SQNR (dB)", "比少一位多"], rows)
    print_observations([
        "输出的相对误差和权重的相对误差几乎完全相同：量化噪声在点积里**没有**被'抵消'掉——误差和信号一样按 √K 增长，"
        "比值不变。所以'逐元素误差 9%'就意味着'每个矩阵乘法的输出也有约 9% 的噪声'。",
        "那为什么 NVFP4 的模型精度还能和 FP8 相差不到 1%（第 1 讲 §5.4）？答案不在这一层：模型对噪声有一定的容忍度"
        "（最终看的是 argmax、是排序），不同层、不同张量的敏感度差别很大，量化方法还会针对性地保护重要的部分（第 3 讲）。",
        "每多 1 位，SQNR 增加约 6 dB（理论值 20·log10(2) ≈ 6.02 dB：间隔减半，误差减半）。位数很少时（INT2–INT3）"
        "截断和'0 附近的刻度太粗'让增量偏离 6 dB。",
    ])


if __name__ == "__main__":
    demo1_affine_and_integer_matmul()
    demo2_clipping()
    demo3_granularity()
    demo4_block_formats()
    demo5_error_propagation()