""" ============================================================================= CS349E · 第 1 讲共享模块：llm_math.py —— LLM 推理的"账本"：参数、FLOPs、KV cache、每一步多久 ============================================================================= 【这个模块要解决什么问题】 第 1 讲要把 roofline 用到 LLM 推理上，需要回答四个问题： 1. 模型有多少参数？权重占多少字节？（决定每一步至少要搬多少字节） 2. 生成一个 token 要做多少次运算？（决定 T_math） 3. 每个 token 的 KV cache 占多少字节？（决定 decode 时每个序列还要额外读多少字节， 以及显存里能同时放下多少个序列） 4. 给定芯片、批大小、上下文长度，一步 decode / 一次 prefill 至少要多久？ 这里只用"公式"回答，不真的跑模型。第 2 讲会用一个能真正运行的 numpy Transformer 逐项验证这些公式（数一数每个矩阵乘法的运算次数）。 【模型结构假设：Llama 风格的 decoder-only Transformer】 每一层 = 注意力（GQA：n_heads 个 query 头共享 n_kv_heads 组 K/V）+ SwiGLU MLP（3 个矩阵） + 2 个 RMSNorm。输入是词嵌入表（embedding），输出是 lm_head（把向量映射回词表）。 逐项参数（d = d_model，dh = d_head，H = n_heads，G = n_kv_heads，f = d_ff）： 注意力：Wq [d, H·dh] + Wk [d, G·dh] + Wv [d, G·dh] + Wo [H·dh, d] MLP ：W_gate [d, f] + W_up [d, f] + W_down [f, d] RMSNorm：每层 2·d 个缩放系数；最后还有 1 个 final norm 词嵌入：vocab·d；lm_head：vocab·d（若与词嵌入共享权重 tied，则不重复计） 【三个最重要的近似（务必记住）】 - 每个 token 的矩阵乘法运算量 ≈ 2 × "参与矩阵乘法的参数量"（每个参数做一次乘、一次加）。 词嵌入只是"查表"，不做乘法，所以不算。 - 注意力打分额外的运算：每个 token 对上下文里每个位置做 QKᵀ 和 PV 两次点积， 每层 2·2·H·dh·ctx 次。上下文很长时这一项不能忽略。 - decode 每一步要把**全部权重读一遍**（被批里的 B 个 token 共享），再加上**每个序列各自的** KV cache（不能共享）。激活值的读写很小，忽略。 【运行】python3 llm_math.py # 自检：三个 Llama 3.1 模型的参数量和公开数字对比 【依赖】只用 Python 标准库；从同目录导入 hw_specs。 """ import math from dataclasses import dataclass from hw_specs import BYTES_PER_ELEM @dataclass(frozen=True) class ModelConfig: """ 一个 Llama 风格 decoder-only Transformer 的结构参数。 【字段】 n_layers ：层数 d_model ：隐藏维度（每个 token 的向量长度） n_heads ：query 头数 n_kv_heads：K/V 头数（= n_heads 是普通多头注意力 MHA；< n_heads 是 GQA；= 1 是 MQA） d_head ：每个头的维度 d_ff ：MLP 中间层维度 vocab ：词表大小 tied ：词嵌入和 lm_head 是否共享同一个矩阵 """ name: str n_layers: int d_model: int n_heads: int n_kv_heads: int d_head: int d_ff: int vocab: int tied: bool = False # ---------------- 参数量 ---------------- def attn_params_per_layer(self): """每层注意力的参数：Wq + Wk + Wv + Wo。""" d, dh = self.d_model, self.d_head return d * self.n_heads * dh + 2 * d * self.n_kv_heads * dh + self.n_heads * dh * d def mlp_params_per_layer(self): """每层 SwiGLU MLP 的参数：gate、up、down 三个矩阵。""" return 3 * self.d_model * self.d_ff def n_params(self): """ 总参数量（和模型卡上报的数字可以直接比较）。 【边界】tied=True 时 lm_head 不重复计数。 """ per_layer = self.attn_params_per_layer() + self.mlp_params_per_layer() + 2 * self.d_model embed = self.vocab * self.d_model head = 0 if self.tied else self.vocab * self.d_model return embed + self.n_layers * per_layer + self.d_model + head def matmul_params(self): """ 参与矩阵乘法的参数量 = 各层的注意力和 MLP 矩阵 + lm_head。 【为什么不含词嵌入】词嵌入是"按 token id 取一行"，不做乘法；decode 时也只读 B 行， 相比整张表可以忽略。（tied 时 lm_head 用的就是词嵌入那张表，仍要整张读、整张乘。） 【为什么不含 RMSNorm】它们是逐元素缩放，参数量和运算量都很小。 """ return self.n_layers * (self.attn_params_per_layer() + self.mlp_params_per_layer()) \ + self.vocab * self.d_model # ---------------- KV cache ---------------- def kv_bytes_per_token(self, kv_elem_bytes=2.0): """ 每个 token 的 KV cache 字节数 = 2（K 和 V）× 层数 × KV 头数 × 每头维度 × 每个数的字节数。 【例子】Llama-3.1-8B、BF16：2 × 32 × 8 × 128 × 2 = 131,072 字节 = 128 KiB / token。 【系统联系】GQA 把 n_kv_heads 从 32 降到 8，KV cache 就小了 4 倍——这是为推理做的结构设计。 """ return 2 * self.n_layers * self.n_kv_heads * self.d_head * kv_elem_bytes # ---------------- 运算量 ---------------- def linear_flops_per_token(self): """每个 token 的矩阵乘法运算量 = 2 × matmul_params（和上下文长度无关）。""" return 2 * self.matmul_params() def attn_flops_per_token(self, ctx): """ 一个 token 对 ctx 个位置做注意力打分的额外运算量 = 层数 × 2·2·H·dh·ctx。 【推导】每个 query 头：和 ctx 个 key 各做一次长度 dh 的点积（2·dh·ctx 次运算）， 再用 ctx 个权重对 ctx 个 value 加权求和（又是 2·dh·ctx 次）。 【边界】ctx = 0 时为 0。 """ return self.n_layers * 4 * self.n_heads * self.d_head * ctx def flops_per_token(self, ctx): """生成一个 token（上下文长度 ctx）的总运算量。""" return self.linear_flops_per_token() + self.attn_flops_per_token(ctx) def weight_bytes(self, w_prec="bf16"): """全部参数按精度 w_prec 存储时占多少字节（显存账用这个）。""" return self.n_params() * BYTES_PER_ELEM[w_prec] # --------------------------------------------------------------------------- # 公开模型配置（Meta Llama 3.1 模型卡 / config.json） # --------------------------------------------------------------------------- LLAMA31_8B = ModelConfig("Llama-3.1-8B", n_layers=32, d_model=4096, n_heads=32, n_kv_heads=8, d_head=128, d_ff=14336, vocab=128256) LLAMA31_70B = ModelConfig("Llama-3.1-70B", n_layers=80, d_model=8192, n_heads=64, n_kv_heads=8, d_head=128, d_ff=28672, vocab=128256) LLAMA31_405B = ModelConfig("Llama-3.1-405B", n_layers=126, d_model=16384, n_heads=128, n_kv_heads=8, d_head=128, d_ff=53248, vocab=128256) MODELS = {m.name: m for m in (LLAMA31_8B, LLAMA31_70B, LLAMA31_405B)} OFFICIAL_PARAMS = {"Llama-3.1-8B": 8.03e9, "Llama-3.1-70B": 70.6e9, "Llama-3.1-405B": 405e9} # --------------------------------------------------------------------------- # 一步 decode、一次 prefill 至少要多久（roofline 下界） # --------------------------------------------------------------------------- def decode_step(model, chip, batch, ctx, w_prec="bf16", compute_prec=None, kv_elem_bytes=2.0, n_chips=1, compute_eff=1.0, mem_eff=1.0): """ 一步 decode（批里 batch 个序列各生成 1 个 token，每个序列已有 ctx 个 token 的上下文）的 roofline 估计。 【参数】 model, chip ：ModelConfig 和 hw_specs.Chip batch ：批大小 B（这一步一共处理 B 个 token） ctx ：每个序列的上下文长度（这里假设批里所有序列一样长） w_prec ：权重存储精度（决定要读多少字节） compute_prec ：矩阵乘法用哪种精度的算力；None 表示和 w_prec 一样 kv_elem_bytes：KV cache 每个数几个字节 n_chips ：把模型切到几块芯片上（理想切分：算力和带宽都乘 n_chips，**不计通信**——第 4 讲再算） compute_eff, mem_eff：能达到峰值的比例（默认 1.0 = 理想下界） 【返回】字典：flops、weight_bytes、kv_bytes、t_math、t_mem、t_step（= 两者较大者）、bound、 intensity（FLOPs / 总字节） 【边界】batch = 0 时返回的 t_step 只含读权重的时间（没有意义，调用方应避免）。 """ compute_prec = compute_prec or w_prec w_bytes = model.matmul_params() * BYTES_PER_ELEM[w_prec] kv_bytes = batch * ctx * model.kv_bytes_per_token(kv_elem_bytes) flops = batch * model.flops_per_token(ctx) peak = chip.peak_flops(compute_prec) * n_chips * compute_eff bw = chip.mem_bw * n_chips * mem_eff t_math = flops / peak t_mem = (w_bytes + kv_bytes) / bw return { "flops": flops, "weight_bytes": w_bytes, "kv_bytes": kv_bytes, "t_math": t_math, "t_mem": t_mem, "t_step": max(t_math, t_mem), "bound": "compute" if t_math >= t_mem else "memory", "intensity": flops / (w_bytes + kv_bytes), } def prefill_time(model, chip, n_tokens, w_prec="bf16", compute_prec=None, kv_elem_bytes=2.0, n_chips=1, compute_eff=1.0, mem_eff=1.0): """ 一次 prefill（一个序列的 n_tokens 个提示 token 一起算）的 roofline 估计。 【运算量】线性层：n_tokens × 2·matmul_params；注意力（因果掩码，第 i 个 token 看前 i 个位置）： Σ_{i=1..n} attn_flops_per_token(i) = attn_flops_per_token(1) · n(n+1)/2。 【搬运量】权重读一遍 + 把 n_tokens 个 token 的 KV 写回显存（假设用了 FlashAttention 一类的融合 内核，n×n 的分数矩阵不落到显存，第 11 讲）。 【返回】同 decode_step 的字典（没有 kv 读取项，kv_bytes 表示写入量） """ compute_prec = compute_prec or w_prec w_bytes = model.matmul_params() * BYTES_PER_ELEM[w_prec] kv_bytes = n_tokens * model.kv_bytes_per_token(kv_elem_bytes) flops = n_tokens * model.linear_flops_per_token() \ + model.attn_flops_per_token(1) * n_tokens * (n_tokens + 1) / 2 peak = chip.peak_flops(compute_prec) * n_chips * compute_eff bw = chip.mem_bw * n_chips * mem_eff t_math = flops / peak t_mem = (w_bytes + kv_bytes) / bw return { "flops": flops, "weight_bytes": w_bytes, "kv_bytes": kv_bytes, "t_math": t_math, "t_mem": t_mem, "t_step": max(t_math, t_mem), "bound": "compute" if t_math >= t_mem else "memory", "intensity": flops / (w_bytes + kv_bytes), } def max_batch(model, chip, ctx, w_prec="bf16", kv_elem_bytes=2.0, n_chips=1, reserve_frac=0.0): """ 显存容量允许的最大批大小：(总显存 × (1 − 预留比例) − 权重) / (ctx × 每 token 的 KV 字节)。 【参数】reserve_frac：留给激活、工作区、碎片等的比例（真实引擎通常预留 5%–10%） 【返回】整数；权重本身就放不下时返回 0 【Python 小知识】math.floor 向下取整：放不下半个序列。 """ free = n_chips * chip.mem_cap * (1.0 - reserve_frac) - model.weight_bytes(w_prec) if free <= 0: return 0 return int(math.floor(free / (ctx * model.kv_bytes_per_token(kv_elem_bytes)))) def serving_metrics(step, batch, chip, n_chips=1, price_per_chip_hour=None): """ 把一步 decode 的时间换算成用户和运营者关心的指标。 【参数】step：decode_step 的返回值；batch：批大小；price_per_chip_hour：每块芯片每小时的价格（美元，可选） 【返回】字典： tok_s_user ：每个用户每秒看到几个 token = 1 / t_step（TPOT 的倒数） tok_s_chip ：每块芯片每秒产出几个 token = batch / t_step / n_chips j_per_tok ：每个 token 的能耗（焦耳）= 功耗 × 芯片数 × t_step / batch（按最大功耗算，偏保守） tok_s_mw ：每兆瓦每秒产出几个 token（NVIDIA 用的 "TPS/MW"） usd_per_mtok：每百万 token 的芯片成本（只算芯片租金） 【边界】芯片没有功耗数据时能耗相关项为 None；没给价格时成本为 None。 """ t = step["t_step"] out = { "tok_s_user": 1.0 / t, "tok_s_chip": batch / t / n_chips, "j_per_tok": None, "tok_s_mw": None, "usd_per_mtok": None, } if chip.power_w is not None: watts = chip.power_w * n_chips out["j_per_tok"] = watts * t / batch out["tok_s_mw"] = batch / t / (watts / 1e6) if price_per_chip_hour is not None: out["usd_per_mtok"] = price_per_chip_hour * n_chips / 3600.0 * t / batch * 1e6 return out # --------------------------------------------------------------------------- # 自检 # --------------------------------------------------------------------------- if __name__ == "__main__": from console import print_table, fmt_bytes, fmt_num, fmt_time from hw_specs import H100 rows = [] for m in MODELS.values(): n = m.n_params() rows.append([m.name, f"{n:,}", fmt_num(OFFICIAL_PARAMS[m.name], 2), f"{(n / OFFICIAL_PARAMS[m.name] - 1) * 100:+.2f}%", fmt_bytes(m.weight_bytes("bf16")), fmt_bytes(m.kv_bytes_per_token(), binary=True)]) print_table(["模型", "公式算出的参数量", "官方", "偏差", "BF16 权重", "每 token KV"], rows) s = decode_step(LLAMA31_8B, H100, batch=1, ctx=1024) print(f"\nLlama-3.1-8B 在 H100 上 batch=1、ctx=1024 的一步 decode：T_math={fmt_time(s['t_math'])}，" f"T_mem={fmt_time(s['t_mem'])}，下界 {fmt_time(s['t_step'])}（{s['bound']}-bound），强度 {s['intensity']:.2f}")