""" ============================================================================= CS349E · 第 9 讲 · 05_nki_perf_model.py —— 在真实的芯片上建模：AWS NKI 性能指南的几条规则 ============================================================================= 【这个脚本要回答的问题】 Interstellar 是给"还没造出来的芯片"建模。造出来以后，写内核的人也需要一个心里的性能模型。 第二篇阅读材料是 AWS 的 NKI（Neuron Kernel Interface）性能指南：在 Trainium / Inferentia 的 NeuronCore 上写内核时， 十条优化各自对应一个简单的代价模型。这里把其中几条写成可以算的公式： 演示 1：指令的固定开销（Opt #5a）：每条指令约 100 个周期的固定开销，一条指令只处理一个元素时效率极低。 演示 2：用满 128 个分区（Opt #5b）：只用 64 个分区的指令和用 128 个一样慢。 演示 3：矩阵 × 向量（decode）该把谁放在"驻留"一侧（Opt #7）：短的那个应该做 moving。 演示 4：DMA 每次要搬够大的块（Opt #9）：一次 4 字节和一次 32 KiB 的效率差多少。 演示 5：融合与片上存储（Opt #1、#2）：注意力的中间结果放不下 SBUF 就要溢出到 HBM，算术强度掉一个数量级。 【NeuronCore 的模型（务必看）】 - 时钟 1.4 GHz：指南里"189 ns = NC-v2 上 264 个周期"推出来的。 - 128 个分区：片上存储 SBUF / PSUM 按 128 个分区组织，每个分区对应计算引擎的一条并行通道。 - Vector / Scalar 引擎：每条指令固定开销约 100 个周期（指南："on the order of ~100 cycles"），之后每个周期每个分区处理 1 个元素（假设）。 - Tensor 引擎：128×128 的脉动阵列。LoadStationary 装入驻留的张量（"快速装载"比正常快 4 倍）； MultiplyMoving 让 moving 张量流过阵列，每周期一列，最少 64 个周期（指南：trn1 上 MM_INIT_LATENCY = 64）。 - DMA：16 个引擎；每次传输有固定开销（假设 0.33 µs）、每个引擎 25 GB/s（假设）。指南只说"要大于 32 KiB 才能达到理想带宽"。 - 峰值、带宽这些绝对数字是为了演示取的量级，**不是 Trainium 的官方规格**。比较的是写法之间的倍数。 【运行】python3 05_nki_perf_model.py （< 1 秒） 【依赖】../lecture01 的 console。 """ import math import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 CLK = 264 / 189e-9 # ≈ 1.397 GHz PARTS = 128 INSTR_OVERHEAD = 100 # 周期 MM_INIT = 64 # MultiplyMoving 的最少周期 DMA_ENGINES = 16 DMA_T0 = 0.33e-6 # 每次 DMA 传输的固定开销（假设） DMA_BW = 25e9 # 每个 DMA 引擎的带宽（假设） def vector_instr_cycles(free, partitions=PARTS): """ Vector / Scalar 引擎上一条指令的周期数：固定开销 + 每个分区要处理的元素数（free 维的大小）。 【参数】free：每个分区的元素个数；partitions：用到几个分区（不影响时间——128 条通道并行） 【返回】(周期数, 有用的"元素·通道"占比) """ cycles = INSTR_OVERHEAD + free useful = free * partitions / (cycles * PARTS) return cycles, useful def tensor_tile_cycles(stat_free, mov_free, fast_load=True): """ Tensor 引擎上一次"装驻留张量 + 流过 moving 张量"的周期数（K = 128 的一块）。 【参数】stat_free：驻留张量的自由维大小（≤ 128，对应阵列的 128 列）；mov_free：moving 张量的自由维（≤ 512） 【模型】LoadStationary 的间隔 = stat_free（快速装载时 / 4）；MultiplyMoving 的间隔 = max(mov_free, 64)。 两条指令流水起来，吞吐由较慢的那个决定（指南图 72）。 """ ls = stat_free / 4 if fast_load else stat_free mm = max(mov_free, MM_INIT) return max(ls, mm) def matvec_cycles(M, N, weights_stationary): """ [M, 128] × [128, N] 的矩阵乘法（M 个 token、权重的一块 K = 128 行）在 Tensor 引擎上要多少周期。 【两种映射】 weights_stationary=True ：权重做驻留（每次装 128 列）、激活做 moving（每次最多 512 列）——"短的做 moving"； weights_stationary=False：激活做驻留（每次最多 128 个 token）、权重做 moving——"短的做驻留"。 【返回】周期数 """ if weights_stationary: n_load = math.ceil(N / 128) n_mov = math.ceil(M / 512) return n_load * n_mov * tensor_tile_cycles(min(N, 128), min(M, 512)) n_load = math.ceil(M / 128) n_mov = math.ceil(N / 512) return n_load * n_mov * tensor_tile_cycles(min(M, 128), min(N, 512)) def dma_time(total_bytes, per_transfer): """ 用 16 个 DMA 引擎搬 total_bytes 字节，每次传输 per_transfer 字节：传输数平均分到各个引擎上， 每次 = 固定开销 + 字节 / 带宽。返回 (时间, 达到的带宽占峰值的比例)。 """ n = math.ceil(total_bytes / per_transfer) per_engine = math.ceil(n / DMA_ENGINES) t = per_engine * (DMA_T0 + per_transfer / DMA_BW) return t, total_bytes / t / (DMA_ENGINES * DMA_BW) # ----------------------------------------------------------------------------- def demo1_free_dim(): print_header("演示 1：指令的固定开销（Opt #5a）——处理 128 × 4096 个元素，每条指令的自由维大小不同") rows = [] total = 4096 for free in (1, 16, 128, 512, 2048, 4096): cyc, useful = vector_instr_cycles(free) n_instr = total // free rows.append([free, n_instr, cyc, fmt_time(n_instr * cyc / CLK), f"{useful:.0%}"]) print_table(["每条指令的自由维", "指令数", "每条指令的周期", "总时间", "引擎的有效利用率"], rows) seq = 4096 naive = seq * 189e-9 scan = (INSTR_OVERHEAD + seq) / CLK print(f"\n指南里 Mamba 的例子：前缀和（scan）写成 {seq} 条单元素指令，每条间隔 189 ns → {fmt_time(naive)}；" f"换成一条 tensor_tensor_scan 指令处理整个序列 → 约 {fmt_time(scan)}（{naive / scan:.0f} 倍）。") print_observations([ "每条指令约 100 个周期的固定开销：一条指令只处理 1 个元素时，引擎 99% 的时间在付开销。指南：每个分区至少要 128 个元素才算高效。", "这和 GPU 的内核启动开销（前置讲 P2 §8）、Python 每次调用 numpy 的开销是同一种问题：固定开销要用足够大的工作量摊薄。", "但自由维也不是越大越好：块太大会让引擎之间的流水不畅（Opt #3）、占太多 SBUF（Opt #2）——要在两者之间取舍。", ]) # ----------------------------------------------------------------------------- def demo2_partitions(): print_header("演示 2：用满 128 个分区（Opt #5b）——两个 64 分区的归约 vs 一个 128 分区的归约（每个分区 512 个元素）") c64, u64 = vector_instr_cycles(512, 64) c128, u128 = vector_instr_cycles(512, 128) rows = [["两条 64 分区的指令（串行）", 2 * c64, f"{u64:.0%}"], ["一条 128 分区的指令（\"分区向量化\"）", c128, f"{u128:.0%}"]] print_table(["写法", "周期", "通道利用率"], rows) print_observations([ "每个分区对应一条并行通道：只用 64 个分区，另外 64 条通道空着，指令却和用满 128 个一样慢。", "指南的办法：让两个矩阵乘法把结果写到 PSUM 的不同分区（0–63、64–127），再用一条 128 分区的指令一起处理——快 2 倍。" "Tensor 引擎也一样：输出只占 96 个分区，128×128 的阵列就只用了 96 列。和第 8 讲脉动阵列\"维度太小填不满\"、" "Interstellar 的\"复制\"（03 脚本演示 3）是同一个问题。", ]) # ----------------------------------------------------------------------------- def demo3_short_moving(): print_header("演示 3：矩阵 × 向量（Opt #7）——M 个 token 乘一块 128 × 4096 的权重，谁做驻留？") rows = [] N = 4096 for M in (1, 8, 64, 128, 512, 2048): a = matvec_cycles(M, N, True) b = matvec_cycles(M, N, False) util = M * 128 * N / (min(a, b) * 128 * 128) rows.append([M, f"{a:,.0f}", f"{b:,.0f}", f"{b / a:.2f}×", f"{util:.1%}"]) print_table(["token 数 M", "权重驻留、激活 moving（周期）", "激活驻留、权重 moving（周期）", "后者慢几倍", "Tensor 引擎利用率（取较快的）"], rows) print_observations([ "decode（M = 1）时，把短的激活放在 moving 一侧快 2 倍：权重块用\"快速装载\"（每 4 个元素 1 个周期），" "激活只有 1 列，MultiplyMoving 只付 64 个周期的最小开销；反过来，权重做 moving 要一列一列地流过 4096 列。", "token 多到 128 以上，两种映射一样快——都在满负荷地流数据。这又是\"不同的数据流在资源用满时差不多\"（03 脚本）。", "但不管怎么映射，M = 1 时 Tensor 引擎的利用率都只有约 1.6%：decode 的矩阵 × 向量本来就受限于内存，" "这条优化只是不让 Tensor 引擎**额外**拖后腿。", ]) # ----------------------------------------------------------------------------- def demo4_dma(): print_header("演示 4：DMA 每次搬多大（Opt #9）——用 16 个 DMA 引擎搬 64 MiB") total = 64 * 1024 * 1024 rows = [] for per in (4, 512, 4 * 1024, 32 * 1024, 128 * 1024, 1024 * 1024): t, eff = dma_time(total, per) rows.append([fmt_bytes(per, binary=True), f"{math.ceil(total / per):,}", fmt_time(t), f"{eff:.1%}"]) print_table(["每次传输", "传输次数", "总时间", "达到峰值带宽的"], rows) print_observations([ "每次传输有固定开销：一次 4 字节时几乎全是开销（指南的例子：16 个 DMA 引擎上全是 4 字节的小传输，空档比传输还长）。", "指南：Trainium / Inferentia2 上每次传输要大于 32 KiB 才能达到理想带宽——在我们假设的开销下，32 KiB 约 80%、128 KiB 约 94%。" "让 nl.load 的分区维和自由维都尽量大（比如 128 × 1024 的 fp32 块 = 16 次 32 KiB 的传输）。", "这和第 8 讲 TMA 一次搬一整块、前置讲 P2 的合并访存是同一个道理：搬数据的固定开销要被足够多的字节摊薄。", ]) # ----------------------------------------------------------------------------- def demo5_fusion(): print_header("演示 5：融合与片上存储（Opt #1、#2）——一个注意力头（d = 128，BF16）的 HBM 流量") d = 128 peak = 2 * 128 * 128 * CLK hbm = DMA_ENGINES * DMA_BW ridge = peak / hbm sbuf = 24 * 1024 * 1024 rows = [] for S in (1024, 4096, 16384): flops = 4 * S * S * d qkvo = 4 * S * d * 2 scores = S * S * 2 unfused = qkvo + 4 * scores # 写分数、softmax 读写一次、第二个矩阵乘法再读一次 fused = qkvo rows.append([S, fmt_bytes(scores, binary=True), "放得下" if 2 * scores <= sbuf else "放不下", f"{flops / unfused:.0f}", f"{flops / fused:.0f}", f"{min(1.0, flops / unfused / ridge):.0%}", f"{min(1.0, flops / fused / ridge):.0%}"]) print(f"（按假设的 NeuronCore：Tensor 引擎峰值 {peak / 1e12:.1f} TFLOP/s、HBM {hbm / 1e9:.0f} GB/s → 拐点约 {ridge:.0f} FLOP/字节；" f"SBUF 按 24 MiB。）") print_table(["序列长 S", "S×S 的分数矩阵", "整块放得进 SBUF 吗", "不融合：FLOP/字节", "融合：FLOP/字节", "不融合能跑到峰值的", "融合能跑到峰值的"], rows) print_observations([ "注意力 = 矩阵乘法 → softmax → 矩阵乘法。分成几个内核写时，S×S 的分数矩阵要写到 HBM、读回来做 softmax、再写出去、" "再读进第二个矩阵乘法，算术强度只有约 60 FLOP/字节，受限于内存（只能跑到峰值的一半左右）。", "融合成一个内核，HBM 流量只剩 Q、K、V、O，算术强度随 S 线性增长，变成受限于计算。S = 1024 时整个分数矩阵（2 MiB）" "还放得进 SBUF；S = 4096 起放不下，融合的内核必须一块一块地算、只在片上保留分数矩阵的一小块（在线 softmax）——" "这就是 FlashAttention（第 11 讲）；指南的 Fused Self Attention 教程也是这么做的。", "指南的 Opt #1（别重复加载输入）和 Opt #2（融合以避免溢出）都是 01 脚本那条规律在真实芯片上的样子：" "在有限的片上容量里，把复用做大。", ]) if __name__ == "__main__": demo1_free_dim() demo2_partitions() demo3_short_moving() demo4_dma() demo5_fusion()