Lecture 11: Programming Systems for Specialized Hardware（面向专用硬件的编程系统）（日期：Oct 28）

目录 · ← l10 · l12 →

Lecture 11: Programming Systems for Specialized Hardware（面向专用硬件的编程系统）（日期：Oct 28）

概述：第 10 讲讲了”为什么需要专用硬件”，本讲回答”怎么给专用硬件编程“。专用硬件（TPU 脉动阵列、H100/B100 tensor core、SambaNova SN40L 数据流架构）的能效来自消除指令流开销与异步化，但这也让编程变得极其复杂。本讲以三条主线展开：(1) GPU 侧——NVIDIA 引入 TMA、TMEM、tcgen05 等异步机制后，裸 CUDA 已难以驾驭，需要用 ThunderKittens 这类嵌入式 DSL 把”16×16 tile + asynchrony + producer-consumer 流水”封装起来；(2) 数据流侧——SambaNova 用”数据并行模式（map/zip/reduce）+ metapipelining”这种以数据为中心的模型，让程序员以简单方式获得极致异步；(3) 对比——GPU 上 Llama 3.1 8B 每 token 约 800 次 kernel 调用，而 RDU 上一个 kernel 融合整个 decoder，每 token 仅 3 次调用，凸显 kernel fusion 与同步开销的重要性。

注意：本讲对应 Assignment 4（在 AWS Trainium2 加速器上编写 fused Conv+MaxPool kernel）。Trainium2 与 H100/B100 类似，拥有专用计算引擎（Tensor Engine 做 128×128 矩阵运算、Vector Engine 做向量运算）与软件管理的片上存储（SBUF/PSUM），需要你用”tile 分块 + 显式数据搬运 + 计算/搬运重叠”的方式编程——正是本讲的核心思想。

一、核心概念与定义

1. Programmability vs Efficiency（可编程性与能效的权衡）

定义：能效与可编程性成反比——”Programmability adds overhead ⇒ reduces efficiency”（可编程性带来开销，从而降低能效）。第 10 讲的能效梯度再次出现：Energy-optimized CPU（最好编程）→ GPU ~10× → 领域专用加速器 ~20×（TPU）→ FPGA ~50×？→ ASIC 100–1000×（几乎不可编程）。
现实类比：请一个”什么都会的全能秘书”（通用 CPU）效率低但省心；请一群只会干一件事的专家（专用硬件）效率高，但你必须精确告诉每个专家干什么、怎么衔接（编程难）。
公式/图示：能效 vs 可编程性是光谱的两端；本讲所有编程系统的目标都是：在不损失太多能效的前提下，让专用硬件变得可编程。

2. Asynchronous (nonblocking) execution（异步/非阻塞执行）

定义：在较早的操作完成之前就启动较晚的操作。第 4 页对比了同步执行（LD0→ST0→AO0 串行等待，每步等上一步完成）与异步执行（LD_a0 发出后不等完成就发 ST_a0、AO_a0，多个 LD/ST/AO 同时进行）。实现方式：软件+硬件配合的异步指令与同步原语，以及硬件乱序执行（out-of-order execution）。
现实类比：同步点菜 = 等第一道菜吃完才点第二道；异步 = 一次性把整桌菜都下单，厨房并行做，先上的先吃。吞吐量大幅提升，但你要学会”不等结果就先干别的，之后再来收结果”。

公式/图示：

同步：LD0 →(等)→ ST0 →(等)→ AO0 →(等)→ LD1 →(等)→ ST1 →(等)→ AO1 …
异步：LD0, ST0, AO0, LD1, ST1, AO1, LD2, ST2, AO2 全部发出，互不等待

3. Tiled programming model（分块编程模型）

定义：以 16×16、32×32 等形状的”张量块（tile）”为基本编程原语，而不是标量或向量。为什么？因为 GEMM 要打满 tensor core 峰值、指令开销要低。第 10 讲列出 tiled 编程模型：CUTLASS、Triton、Thunderkittens。
现实类比：集装箱运输 vs 散货运输——散货（标量）一件件装卸慢且贵；集装箱（tile）整箱吊装，效率高一个量级，但你需要标准化的装箱方案（tile 布局与搬运计划）。
关键点：tile 的形状要匹配硬件计算单元的形状（如 H100 tensor core 是 16×16 输入块、TMEM 16×16×16）；”Use 16 x 16 tiles of fp16 data ⇒ matches Tensor core compute”。

4. Tensor core programming（B100：tcgen05 与 TMEM）

定义：B100 上 tensor core 的编程模型发生了革命性变化（”Not your father’s CUDA”）：寄存器带宽限制 tensor core，张量数据放 SMEM 与 TMEM（tensor memory），单个线程即可执行 MMA ⇒ 不再有 warp 的概念。编程步骤：tcgen05.alloc（分配 TMEM 与描述符）→ cp.async.bulk.tensor（用 TMA 异步预取/流式搬 tile，配 mbarrier 协调）→ tcgen05.mma batch + tcgen05.commit（启动异步 MMA 批量）→ tcgen05.fence（排序与回收）。
现实类比：以前是”几十个工人（warp 线程）一起抬一件货（矩阵）”；现在是”一个工人按下一个按钮，专用起重机（单线程 + 硬件单元）把整箱货吊到位”。控制更简单，但你要学会操作起重机（TMEM 分配、描述符、异步提交）。
公式/图示：B100 tensor core 数据流：Global(HBM) –TMA–> SMEM –cp.async.bulk.tensor–> TMEM –tcgen05.mma–> 结果。

5. TMA（Tensor Memory Accelerator，张量内存加速器）

定义：NVIDIA 的专用数据搬运单元，用专用指令高效搬运数据：异步地从 global memory 加载/存储一个张量区域到 shared memory。由拷贝描述符（copy descriptor）描述区域，单个线程发起 TMA 操作（cuda::memcpy_async），拷贝完成时信号屏障（mbarrier），硬件完成地址生成与数据搬运。A100 有 LDGSTS（bypass L1），H100 的 TMA 更进一步。
为什么省电（第 29 页）：消除数千条指令与内存寻址开销；消除不必要地经过 L1 与寄存器的数据搬运。
现实类比：以前搬货要一箱箱人工登记地址（每箱一条指令、每个地址一次寻址）；TMA 是”整批快递单一次性生成、仓库自动分拣直送”，省掉中间的人工中转站（L1/寄存器）。
相关术语：Warpgroup = 128 个连续线程；PTX（Parallel Thread Execution）= NVIDIA 的虚拟指令集架构（virtual ISA）。

6. Embedded DSL（嵌入式领域专用语言：ThunderKittens）

定义：ThunderKittens（TK）是”Embedded CUDA DSL template library”——一个嵌入 CUDA 的模板库，而不是独立语言。它提供模板化的数据类型（register tiles：寄存器上的 2D 张量，含 height/width/layout；register vectors：1D；shared memory tiles/vectors）与操作（initializer、unary op 如 exp、binary op 如 mul、row/column op 如 row_sum）。
设计三原则（第 36 页）：① 16×16 tile 为基本数据类型（TK 管理 layout、提供基本操作）；② 处处异步（暴露原语让用户管理，追求顶级性能）；③ 高层 GPU 协调模式（如 producer-consumer 处理）。
现实类比：TK 是给 GPU 的”标准集装箱物流系统”——你只需说明”我要运哪些 16×16 的箱子、从哪到哪”，箱子规格、吊装顺序、交接点（barrier）由系统帮你标准化。

7. Producer-consumer pipeline（生产者-消费者流水线）

定义：把 kernel 执行组织成流水：Producer 从 global memory 异步加载 tile 到 shared memory → 中间寄存器暂存 → Consumer（tensor core）计算 → 结果经 shared memory 写回 global memory。幻灯片第 37 页给出 tile 处理流水：Global Memory(HBM/L2) → Shared Memory → Registers → Tensor cores → Shared Memory → Global Memory，并标注 Producer / Consumer / Finish 阶段。
现实类比：餐厅流水线——备菜员（producer）不断把食材（tile）摆到备餐台（shared memory），厨师（consumer）只管炒（compute），收盘员（finish）把成品端走。备菜员不能等厨师炒完才备下一份（异步、多级缓冲）。
关键点：让 tensor core 永不空闲（>90% TFLOPS）的核心手段：重叠访存与计算 ⇒ 使用 asynchrony。

8. Dataflow architecture / RDA（数据流架构 / 可重构数据流架构）

定义：把计算”铺”在芯片上：数据在计算单元（PCU）与存储单元（PMU）之间通过开关（switch）直接流动，无指令 ⇒ 无取指/译码开销，极端异步（无顺序指令执行）。SambaNova SN40L RDU 是商用实现：1040 个 PCU+PMU、638 TFLOPS (bf16)、520 MB 片上 SRAM、64 GB HBM、1.5 TB DDR；PCU 做 systolic + SIMD 计算（16×8 bf16）；PMU 高地址生成灵活性与带宽（0.5 MB）；AGCU（Address Generator and Coalescing Unit）是访问片外内存与 IO 的”门户”。
现实类比：河流（数据）流经一串水车（PCU）——每个水车做固定的一步工作，水车间用渠道（switch/PMU）直连，没有”调度员逐条下令”。
公式/图示：AI 模型（数据流图：GEMM + map/filter/reduce 等并行模式）⇒ 数据流架构 = 把图直接映射到 PCU/PMU 网格。

9. Data parallel patterns（数据并行模式）

定义：可组合的计算原语：MM（矩阵乘）、Map（逐元素映射）、Zip（按位置合并）、Reduce（归约）、Gather、Scatter…… 程序 = 模式组合，编译器负责 Tiling、Parallelization、Metapipelining、Place & Route、Codegen。关键特性：”Flexible scheduling in space and time ⇒ spatial execution”（在空间和时间上灵活调度 ⇒ 空间执行）。
现实类比：乐高积木——标准件（模式）种类有限，但组合方式无穷；拼装说明书（程序）由编译器自动翻译成”每个积木放哪”（布局布线）。
示例：简化 Softmax = Map(exp) → Reduce(+) → Zip(÷)，三个模式组合即可表达；GPU 上要写三个 kernel，RDU 上可在片上流水式执行。

10. Metapipelining（元流水线）

定义：层级化的粗粒度流水线：一个”流水线的流水线”。它利用嵌套循环并行性：把并行模式（循环）转换成流式流水线——在循环体内插入流水级（pipe stages），各流水级并行执行，重叠多个循环迭代的执行；级间中间数据用双缓冲（double buffers）存储；能处理执行时间不均衡的流水级；与 tiling 配合良好，缓冲可改变访问模式（如转置数据）；metapipelining 在 fusion 失效时仍然有效。
现实类比：工厂车间里的流水线，每条流水线上又有多个工位——”流水线的流水线”。外层流水线处理”批次”，内层流水线处理”批次内的零件”，两层同时并行。
公式/图示：METAPIPE(M/MM) { LOAD_TILE(A); METAPIPE(N/NN) { LOAD_TILE(B); MAT_MUL; BUFFER; STORE_TILE } }——外层沿 M 分块、内层沿 N 分块，形成两级流水。

11. Kernel fusion（kernel 融合）

定义：把多个计算步骤（多个 kernel）融合成一个 kernel，中间结果不落回 off-chip 内存，从而：提高数据局部性（high data locality）、消除 kernel 启动与同步开销（zero extra launch overheads）。Llama 3.1 8B 案例：GPU 上 Tensor-RT LLM 每 token 约 10 个 kernel、融合有限（K1..K10，low kernel fusion、low data locality、high launch & synchronization overheads）；RDU 上一个 kernel 融合整个 decoder（K0），每 decoder 一次 kernel 调用，3 次调用/token vs GPU ~800 次调用/token，kernel 调用减少约 100×。
现实类比：去食堂点”一荤一素一汤”分开三次排队（三个 kernel、三次启动）；还是”套餐窗口”一次取齐（融合 kernel）——省掉三次排队时间，菜也不用端回桌上再端出来（中间结果留在片上）。
关键数据：SN40L 520 MB 片上 SRAM vs H100 100 MB（5× 优势）；数据流融合消除了 GB 级 off-chip 中间结果流量。

12. Compute-communication overlap（计算-通信重叠）

定义：让通信（如 AllReduce）与计算（如 GEMM、权重加载）同时进行，互不阻塞。RDU 上”Fully overlap allreduce with weight load and compute；Allreduce does not consume HBM capacity or bandwidth”——把 AllReduce 与 Down GEMM 等流水化（第 59 页、第 39 页的 Pipelined AllReduce with Compute, no HBM traffic!）。
现实类比：一边开车一边听导航播报（通信与驾驶并行），而不是”停车听完导航再开”。
公式/图示：RDU 0 与 RDU 1 之间：Down GEMM → Add → AllReduce，三者重叠，通信不再占用 HBM 带宽。

二、代码示例与详细解说（本讲重点）

示例 1：ThunderKittens 编写 GEMM kernel（Step 1：定义 layouts）

// ThunderKittens（TK）—— Embedded CUDA DSL template library
// 目标：在 H100 上打满 tensor core 的 GEMM
// Step 1: Define layouts（定义数据布局）
#include "kittens.cuh"
#include "prototype.cuh"
using namespace kittens;
using namespace kittens::prototype;
using namespace kittens::prototype::lcf;

struct matmul_layout {
   using  a_global_layout = gl<bf16, 1, 1, -1, -1, st_bf<64, 64>>;  // A 的 TMA 描述符：64x64 tile
   using  b_global_layout = gl<bf16, 1, 1, -1, -1, st_bf<64, 256>>; // B 的 TMA 描述符：64x256 tile
   using  c_global_layout = gl<bf16, 1, 1, -1, -1>;                 // C 不需要 TMA 描述符
   struct globals        { a_global_layout A; b_global_layout B; c_global_layout C; };
   struct input_block    { st_bf<64, 64> a[2]; st_bf<64, 256> b; } // 输入 tile：共享内存
   struct finish_block   { st_bf<64, 256> c[2]; };                 // 结果 tile：共享内存
   struct consumer_state { rt_fl<16, 256> accum; };                // 累加器：寄存器 tile
};

【代码做了什么？】 这是 TK GEMM 的第一步：声明整个 kernel 的数据布局。gl<bf16, ...> 是 global memory 布局（同时生成 TMA 描述符，TMA 需要知道从 HBM 的哪里、以什么步幅搬运多大的块）；st_bf<64,64> 是 shared memory 中的 tile（bf16、64 行 64 列）；rt_fl<16,256> 是寄存器文件中的累加器 tile（fp32 累加）。注意 A 是 64×64 tile、B 是 64×256 tile、C 是 64×256——这些形状与后续 consumer warpgroup 的分工（8 个 consumer warpgroup 各负责一部分）严格对应。程序员写的是”数据长什么样”，而不是”每个线程怎么搬数据”——这正是 TK 作为 DSL 的价值：布局与描述符由模板类型自动管理。

【并行机制解说】 对应本讲概念 3（tiled programming）、6（embedded DSL）、5（TMA）。并行性体现在：(1) 以 64×64/64×256 的 tile 为搬运与计算单位，匹配 tensor core 的 16×16 MMA 形状（由 16×16 基本 tile 组合而来）；(2) gl<...> 直接生成 TMA 描述符，意味着访存由 TMA 硬件异步完成，无需每元素一条加载指令——消除了数千条指令与寻址开销（第 29 页 “Eliminates 1000’s of instructions and memory addressing overhead”）；(3) 布局类型化后，编译器/模板能在编译期推导 tile 在 shared memory 与寄存器间的排布，程序员不再手工计算 bank conflict 与地址偏移。这就是”低指令开销、高 tile 复用”的软件化表达。

示例 2：ThunderKittens 的 producer-consumer 流水（Step 2 & Step 3）

// Step 2: Define pipeline and producers（定义流水线与生产者）
struct matmul_template {
   using layout = matmul_layout;
   static constexpr int NUM_CONSUMER_WARPS=8, INPUT_PIPE_STAGES=4;
   // 8 个活跃 consumer warp，4 级输入流水
   static constexpr int PRODUCER_BARRIER_ARRIVALS=1, CONSUMER_BARRIER_ARRIVALS=2;

   struct producer {
      __device__ static void setup(producer_setup_args<layout> args) {
         warpgroup::decrease_registers<40>(); // 生产者少占寄存器，留给消费者
      }
      __device__ static void load(producer_load_args<layout> args) {
         if(warpgroup::warpid() == 0) {   // 只需一个 warp（其实一个线程）发起 TMA
            tma::expect(args.inputs_arrived, args.input); // 告诉 mbarrier 期待多少字节
            for(int i = 0; i < 2; i++) {  // 加载两个 A tile（每个 consumer warpgroup 一个）
               tma::load_async(args.input.a[i], args.globals.A,
                               {blockIdx.x*2+i, args.iter}, args.inputs_arrived);
            }
            // 加载 B tile：一个 64x256 tile，所有 consumer warpgroup 共享
            tma::load_async(args.input.b, args.globals.B,
                            {args.iter, blockIdx.y}, args.inputs_arrived);
         }
      }
   };

   // Step 3: Compute!（消费者）
   struct consumer {
      __device__ static void setup(consumer_setup_args<layout> args) {
         warpgroup::increase_registers<232>(); // 消费者多占寄存器（累加器）
         zero(args.state.accum);               // 累加器清零
      }
      __device__ static void compute(consumer_compute_args<layout> args) {
         // 模板会等输入 tile 就绪后再调用本函数
         warpgroup::mma_AB(args.state.accum, args.input.a[warpgroup::groupid()],
                           args.input.b);           // tensor core 执行 MMA
         warpgroup::mma_async_wait();               // 等待异步 MMA 完成
         if(warpgroup::laneid() == 0) arrive(args.inputs_finished); // 标记内存可复用
      }
      __device__ static void finish(consumer_finish_args<layout> args) {
         int wg = warpgroup::groupid();
         warpgroup::store(args.finish.c[wg], args.state.accum);  // 先存到共享内存
         warpgroup::sync();  // 在 SMEM 中重组，便于对 HBM 合并（coalescing）写
         warpgroup::store(args.globals.C, args.finish.c[wg],
                          args.state.accum, {blockIdx.x*2+wg, blockIdx.y});
      }
   };
};

【代码做了什么？】 这是 TK GEMM 的”流水线骨架”：生产者（producer）用 TMA 把 A、B 的 tile 异步搬入 shared memory（只用一个 warp 里的一个线程发起，tma::load_async 配 tma::expect 与 mbarrier inputs_arrived）；消费者（consumer）等输入就绪后调用 warpgroup::mma_AB 让 tensor core 做矩阵乘加，mma_async_wait 等异步 MMA 完成，然后 arrive 通知”这块 shared memory 可以复用了”；finish 阶段把累加结果先写回 shared memory 重组（提高对 HBM 的合并写效率），再写回 global memory。常量 INPUT_PIPE_STAGES=4 表示 4 级输入流水（多缓冲），NUM_CONSUMER_WARPS=8 表示 8 个 consumer warp 并行。生产者主动 decrease_registers<40>()、消费者 increase_registers<232>()——寄存器在生产者与消费者之间动态分配。

【并行机制解说】 对应概念 2（asynchrony）、7（producer-consumer pipeline）、4（tensor core 编程）。并行机制是流水线并行（pipelining）+ 数据并行：(1) 4 级输入流水让”搬运 tile k+1”与”计算 tile k”重叠，tensor core 永远不空等访存（第 35 页：Make sure compute is never idle; Overlap memory access and compute ⇒ use asynchrony）；(2) 8 个 consumer warp 各自负责一部分输出列（a[warpgroup::groupid()]），是数据并行；(3) TMA + mbarrier 让同步原语极轻量（arrive/wait 而非 lock）；(4) MMA 由 tensor core 硬件执行（16×16 输入、fp32 累加），单条指令完成整个矩阵块运算，摊销指令开销。TK 的贡献是：把这些原本散落在 CUDA/PTX 深处的机制（mbarrier、TMA descriptor、async MMA、寄存器分配）封装成类型化、可组合的模板原语，让程序员用”声明布局 + 写 producer/consumer 回调”的方式表达，而不是手写 PTX——”A Simple Embedded DSL for AI kernels”。

示例 3：SambaNova 的 Matmul Metapipeline（数据流 DSL）

// SambaNova RDU 上的矩阵乘 metapipeline（数据流编程模型）
// 以数据为中心：程序描述"数据如何分块、如何在片上流动"
auto format = DataFormat::kBF16;
int64_t M = args::M.getValue();   // 输出行数
int64_t N = args::N.getValue();   // 输出列数
int64_t K = args::K.getValue();   // 归约维

auto A = INPUT_REGION("A", (M, K), format);    // 声明输入张量区域
auto B = INPUT_REGION("B", (K, N), format);
auto C = OUTPUT_REGION("C", (M, N), format);

auto MM = 256;  // M 方向 tile 大小（假设能整除 M）
auto NN = 64;   // N 方向 tile 大小（假设能整除 N）

METAPIPE(M / MM, [&]() {                      // 外层流水：沿 M 分块
    auto a_tile = LOAD_TILE(A, a_tile_shape); // 从 AGCU/HBM 加载 A 的 tile
    METAPIPE(N / NN, [&]() {                  // 内层流水：沿 N 分块
        auto b_tile = LOAD_TILE(B, b_tile_shape, row_par = 4); // 4 路行并行加载
        auto c = MAT_MUL(a_tile, b_tile);     // PCU 执行矩阵乘
        auto c_tile = BUFFER(c);              // 结果进片上缓冲（可转置/改变访问模式）
        STORE_TILE(C, c_tile);                // 写回 HBM
    });
});

【代码做了什么？】 这是 SambaNova 数据流编程模型下的 GEMM：程序员声明输入/输出张量区域（INPUT_REGION/OUTPUT_REGION），然后写一个两层嵌套的 METAPIPE——外层沿 M 分块加载 A 的 tile，内层沿 N 分块加载 B 的 tile、执行 MAT_MUL、把结果 BUFFER 后 STORE_TILE 写回。与普通 C++/CUDA 循环不同，这里的”循环”不是指令级的，而是声明式的流水结构：编译器会把外层/内层循环转换成片上的流式流水线（第 51-53 页）：LOAD_TILE 映射到 AGCU（地址生成与合并单元）→ a_tile/b_tile 进 PMU（片上存储）→ MAT_MUL 在 PCU 上执行 → c_tile 经 BUFFER 缓冲（双缓冲、可改变数据布局）→ STORE_TILE 经 AGCU 写回。

【并行机制解说】 对应概念 8（dataflow）、9（data parallel patterns）、10（metapipelining）。并行机制是层级流水（pipeline of pipelines）：(1) 外层 METAPIPE(M/MM) 与内层 METAPIPE(N/NN) 各自成为流水级，多个迭代重叠执行（第 49 页：Overlap execution of multiple loop iterations）；(2) 中间数据（a_tile、b_tile、c_tile）用双缓冲存储，且 row_par=4 指定 B 的加载按 4 路行并行；(3) BUFFER 不只是存储——缓冲可以改变访问模式（如转置），这是融合做不到的（第 49 页：Buffers can be used to change access pattern; Metapipelining can work when fusion does not）；(4) 数据在 AGCU/PMU/PCU 之间直接流动，无指令流、无 kernel 启动开销。程序员写的是”以数据为中心的流水描述”，硬件把计算按空间布局（spatial execution）铺开——这就是”Can we have asynchrony with a simpler programming model?（数据中心的视角）”的答案。

示例 4：简化 Softmax 的数据并行模式表达（Map/Reduce/Zip）

# SambaNova 数据流编程：用可组合的并行模式表达计算
# 简化 Softmax：softmax(x)[i] = exp(x[i]) / sum_j exp(x[j])
# 模式：Map(exp) → Reduce(+) → Zip(/)

def simplified_softmax(x):
    # 1) Map：逐元素求 exp（每个元素独立）
    e = map(exp, x)                # e[i] = exp(x[i])
    # 2) Reduce：归约求和（得到分母）
    s = reduce(add, e)             # s = Σ e[i]
    # 3) Zip：按位置两两相除（每对元素独立）
    y = zip(div, e, broadcast(s))  # y[i] = e[i] / s
    return y

# 编译流水：Tiling → Parallelization → Metapipelining → Place&Route → Codegen
# 片上执行：三个模式被布局到不同 PCU，数据流式流过：
#   PMU(x) → PCU(Map:exp) → PCU(Reduce:+) → PMU(s) → PCU(Zip:/) → 输出

【代码做了什么？】 这是一个”以数据为中心”的编程示例：把 softmax 拆成三个可组合的并行模式——Map（逐元素 exp）、Reduce（求和）、Zip（逐元素相除）。程序员只描述”数据变换”，不描述”哪个线程/哪个周期做什么”。编译器负责后续的 Tiling（分块）、Parallelization（并行化）、Metapipelining（流水化）、Place & Route（布局布线）与 Codegen（代码生成），最终把模式图映射到 RDU 的 PCU/PMU 网格上。第 48 页的”SIMPLIFIED SOFTMAX”图正是这个例子的硬件映射示意。

【并行机制解说】 对应概念 9（data parallel patterns）、8（dataflow）。并行机制：(1) 每个模式内部是数据并行——Map/Reduce/Zip 天然适合大量 PCU 并行处理不同数据元素；(2) 模式之间是流式流水并行——exp 的输出直接流入 reduce 的输入、再流入 zip，中间结果不落 HBM（kernel fusion 效果）；(3) 空间执行（spatial execution）——计算按”在芯片上布局”来调度（第 48 页：Flexible scheduling in space and time ⇒ spatial execution），而 GPU 需要多次 kernel launch（每次启动 + 同步开销）。这解释了第 57 页的对比：同样一个模型，GPU 上”Low kernel fusion, Low data locality, High Launch and Synchronization Overheads”，RDU 上”High kernel fusion: One kernel call for per decoder ⇒ High data locality, Zero Kernel extra launch overheads”。

三、关键要点

能效来自专用化，但专用化让编程变难，DSL/编程系统是桥梁：H100 的异步机制（TMA、TMEM、tcgen05、mbarrier）让裸 CUDA 编程变得”不是父辈的 CUDA”（Not your father’s CUDA）；ThunderKittens 这类嵌入式 DSL 把复杂性封装为”16×16 tile + asynchrony + producer-consumer”原语。
异步（重叠）是打满专用硬件的必要条件：tensor core 要 >90% TFLOPS，就必须”计算永远不空闲、访存与计算重叠”——同步执行（每步等待）会让访存延迟直接变成空闲时间。
数据流（dataflow）编程模型以数据为中心，天然获得极致异步：SambaNova 用”并行模式（MM/Map/Zip/Reduce/Gather/Scatter）+ metapipelining”编程，无指令流、无 kernel 启动、无 lock 同步（token 控制）；Metapipeline 是”流水线的流水线”，能处理融合（fusion）处理不了的不规则/需要改布局的场景。
Kernel fusion 的价值量化：Llama 3.1 8B 推理，GPU 上每 token ~800 次 kernel 调用，RDU 上融合整个 decoder 后仅 3 次——100× 更少的 kernel 调用，消除了 GB 级 off-chip 中间结果流量与大量启动/同步开销。
GPU kernel 的重要性：2025 年 NVIDIA 单季营收 >$47B，AI kernel 在价值数亿美元的 GPU 集群上运行数月；FlashAttention-2 在 A100 上约 70% 效率、到 H100 掉到 ~35%，两年后 FlashAttention-3 才回到 ~65%——低质量 kernel 会浪费数十亿美元的计算资源。
通信与计算必须重叠：多 socket 扩展时通信时间占比上升，不做重叠时通信成为瓶颈（GPU 需要巨大互连带宽）；RDU 把 AllReduce 与权重加载/计算完全重叠，通信不消耗 HBM 容量与带宽。

四、常见陷阱与注意事项

把异步当同步用：写 TK/CUDA 时若在每个异步操作后立刻等待（sync/wait），流水线退化为同步执行，访存延迟全部暴露，tensor core 利用率暴跌。正确做法是让 producer 提前多级预取（如 INPUT_PIPE_STAGES=4），消费者尽量”晚等”。
tile 形状与硬件不匹配：tensor core 吃 16×16（fp16）的块，如果你用与硬件形状不匹配的 tile（如 8×8 或非对齐尺寸），MMA 无法全速执行，指令开销摊销失败。TK 中 16×16 tile 是基本类型，CUTLASS 中 tile 形状要与 SM 资源严格核算。
忽略同步/启动开销的放大效应：GPU 上每个 kernel launch、每次 kernel 间同步都有固定开销；模型级联几十个 kernel 时，这些开销与中间结果落 HBM 的流量会被放大（Llama 8B 案例：~800 calls/token）。融合 kernel 或使用持久化 kernel（如 RDU 的”一个 kernel 跑所有 decoder”）是正解。
数据搬运仍是最贵的：即使有 TMA，数据一旦必须落 HBM 就付出 1200 pJ/64bit 量级的代价。别把中间结果随便写回 global memory；利用 fusion、片上缓冲（BUFFER）、双缓冲把数据留在 SRAM 里流动（SN40L 520 MB vs H100 100 MB 的 5× SRAM 优势就是这么用的）。
寄存器/shared memory 预算失衡：TK 示例中生产者主动降寄存器（40）、消费者占用（232）——如果所有 warp 都贪心占寄存器，SM 上能驻留的 warp 变少，流水深度下降。资源分配是性能的一部分。
误以为 metapipelining 万能：metapipelining 能处理 fusion 失效的场景（缓冲改布局、不均衡流水级），但片上有存储/单元是有限资源（PCU/PMU 数量、SRAM 容量），无限流水会导致布局布线失败或利用率下降。

五、思考题（带答案）

Q1：为什么 H100 需要 ThunderKittens 这类 DSL，而早期 CUDA（如 V100 时代）相对”好写”？请从硬件机制变化的角度回答，并说明 DSL 具体封装了哪些复杂性。

答案：因为硬件越来越专用化（第 33 页：Nvidia chips becoming more specialized）。V100 只有 tensor core；到 A100 加了 sparsity、FP8、Transformer Engine、异步执行、distributed SHMEM；到 H100 加了第四代 tensor core、TMA、异步拷贝；到 B100 甚至引入 TMEM、tcgen05、FP4、解压引擎，且”单个线程执行 MMA ⇒ 不再有 warp”。程序员要直接驾驭这些机制（TMA 描述符、mbarrier、异步 MMA 提交/回收、寄存器分配、TMEM 分配）极其繁琐且易错。TK 封装了：(1) 布局管理（gl/st_bf/rt_fl 类型自动生成 TMA 描述符与内存排布）；(2) 异步原语（load_async/mma_async/mbarrier arrive-wait）；(3) 高层协调模式（producer-consumer 流水、4 级多缓冲）；(4) 资源分配策略（生产者/消费者寄存器配比）。程序员写”布局声明 + producer/consumer 回调”，复杂机制由模板与编译器处理。

Q2：SambaNova 的 metapipelining 与 GPU 上的 kernel fusion 有什么本质区别？为什么说”metapipelining 在 fusion 失效时仍然有效”？

答案：kernel fusion 是把多个算子合并成一个 kernel、中间结果留在片上，本质仍是”指令流 + 集中同步”的执行（GPU 上仍有 kernel launch、grid/sync、L1/L2 往返）。metapipelining 是把嵌套循环转换成层级化的数据流流水线（pipeline of pipelines）：各级是物理上独立的 PCU/PMU，级间通过双缓冲与 token 控制流动，无锁同步、无指令取指、无启动开销。fusion 失效的场景（第 49 页）：(1) 算子间访问模式需要改变（如转置）——fusion 无法在中间插入布局转换，而 metapipelining 的 BUFFER 可以改变访问模式；(2) 流水级执行时间不均衡——双缓冲与流水结构天然吸收不均衡；(3) 需要重叠多个循环迭代——融合只是”减少中间落盘”，metapipelining 还重叠了迭代间执行。一句话：fusion 优化”一次执行内的数据局部性”，metapipelining 把”整个计算变成一条永不停止的流水线”。

Q3：RDU 上”一个 kernel 融合整个 decoder（Llama 3.1 8B）、3 calls/token”，而 GPU 上约 800 calls/token。请分析：这是否意味着 RDU 一定比 GPU 快？GPU 该如何缩小这一差距？

答案：不一定。kernel 调用次数少只消除了”启动/同步开销”与”中间结果 HBM 往返”，但推理性能还受峰值算力、HBM 带宽、片上 SRAM 容量、能效等因素约束（第 58 页也承认 “HBM BW limits inference performance”，关键是 overlap 让 HBM 始终忙碌）。RDU 的 520 MB SRAM（5× H100）使其能容纳整个 decoder 的中间状态，这是”一个 kernel 融合”可行的前提；GPU 缩小差距的路径：(1) 用 CUTLASS/TK/Triton 写融合 kernel（如 FlashAttention 系列把 attention 融合，FlashAttention-3 把效率从 35% 拉回 ~65%）；(2) 使用 CUDA Graphs / persistent kernels 减少 launch 开销；(3) 用 producer-consumer 流水 + TMA 多缓冲让权重加载与计算重叠（GPU 也可以做到”HBM 始终忙碌”）；(4) 增大 L2 驻留（L2 Cache Residency 机制）。所以核心结论是：kernel 次数是表象，本质是”中间数据是否留在片上 + 计算/访存/通信是否重叠”——这正是第 57 页两种实现的真正差异。

Share on

Bluesky Facebook LinkedIn Mastodon X (formerly Twitter)

Runguo Li