Building an LLM Inference Engine from ScratchPart II / Module 4
Part II · GPU Compute

Module 4：为什么需要 GPU，什么是 Shader？

从 CPU/GPU 执行模型、内存层次、coalesced access 到 roofline model，建立写高性能 LLM kernel 的第一张地图。

学习目标

  • 解释 CPU 与 GPU 的执行模型差异：低延迟 vs 高吞吐，SIMT / wavefront / occupancy。
  • 理解 GPU 内存层次：register、groupshared/LDS、global VRAM、host memory。
  • 说明 coalesced memory access 为什么常常比“多做一点算术”更重要。
  • 理解 compute shader 的基本概念：thread、threadgroup、dispatch grid。
  • 用 arithmetic intensity 与 roofline model 判断一个 kernel 是 memory-bound 还是 compute-bound。

4.1CPU vs GPU：两种不同的并行哲学

CPU 与 GPU 面向不同的执行目标。CPU 依靠少量复杂核心、乱序执行、分支预测和多级缓存降低单个任务的延迟，适合解析、调度、采样和系统调用等控制流较复杂的工作。GPU 使用大量较轻量的执行单元，把同一段程序同时作用于大量数据，以吞吐量换取单线程延迟。LLM 推理中的矩阵乘、归约、attention 和逐元素算子满足这种数据并行形式，因此适合放到 GPU 上执行。

维度CPUGPU
目标降低单个任务延迟提高大量相似任务的总吞吐
并行粒度几十个硬件线程成千上万个轻量线程
控制流复杂分支也能处理分支发散会浪费 lanes
内存访问缓存隐藏很多随机访问强依赖 coalesced access
适合任务调度、解析、采样、系统逻辑matmul、attention、归约、逐元素算子
定义：SIMT

SIMT（Single Instruction, Multiple Threads）指一组线程按同一指令流执行，但每个线程拥有独立的寄存器、索引和谓词。NVIDIA 通常称这组线程为 warp，AMD 常称 wavefront 或 wave；本教程统一称为 wave。分支发散并不会停止程序运行，但同一 wave 内不同分支路径需要分批执行，空闲 lane 会降低有效吞吐。

CPU：少量强核心 Core 0 Core 1 Core 2 Core 3 强控制流、低延迟、大缓存 GPU：大量轻线程 高吞吐、同构计算、依赖内存访问模式
图 4-1：CPU 与 GPU 不是“谁更强”的关系，而是为不同工作负载设计。

4.2GPU 内存层次：越近越快，越远越大

GPU kernel 的性能通常受数据移动限制。寄存器、groupshared/LDS、global VRAM 与 host memory 的容量、可见范围和访问延迟不同；正确的设计会把频繁复用的数据放在更近的位置，并尽量减少 CPU 与 GPU 之间的传输。简化的内存层次如下：

Registers 每线程私有；最快；容量小 groupshared / LDS 同一 threadgroup 共享；需要 barrier Global VRAM 大；高带宽；访问模式决定效率 Host memory / CPU readback（最慢，跨 PCIe/系统边界）
图 4-2：GPU 内存层次。LLM kernel 的优化经常是“少去远处拿数据”。
位置谁能访问典型用途常见错误
Register单个线程临时累加、Q 向量片段用太多导致 occupancy 降低
groupshared / LDS同一 threadgroup归约、tile 缓存、共享 activation忘记 barrier 或 bank conflict
Global VRAM所有 GPU 线程权重、KV cache、activation buffer非 coalesced 访问
Host memoryCPU加载权重、最终 token id每 token 读回大量 logits

4.3Coalesced memory access：GEMV 性能的生命线

显存访问以 cache line 或 memory transaction 为单位完成，而不是按单个线程的单个 float 独立计费。同一 wave 中相邻线程读取连续地址时，硬件可把访问合并为少量事务，这称为 coalesced access。若相邻线程访问相距很远的地址，同样数量的标量读取会拆成更多事务，实际带宽下降。

Coalesced：相邻线程读相邻地址 一次或少数几次 memory transaction Strided：相邻线程读分散地址 多次事务，带宽浪费 为什么这会影响 LLM？ decode 阶段的线性层通常是 GEMV：一个 activation vector × 巨大的权重矩阵。 如果每个输出由一个线程独自读一整行权重，相邻线程读到的地址相隔 K，往往不 coalesced。 优化后的 xinfer 让一个 threadgroup 计算一个输出元素，线程沿 K 方向读取相邻权重并归约。
图 4-3：coalesced access 决定显存带宽能否被真正用起来。

4.4什么是 compute shader？

Compute shader 是由 CPU 录入命令、在 GPU 上执行的程序。CPU 侧提交一次 dispatch，指定 threadgroup 的三维数量；shader 代码中的 [numthreads] 指定每个 group 内的线程数量。每个线程通过系统语义取得全局或组内编号，并据此选择要处理的数据元素。

// HLSL: 每个 thread 计算 out[i] = in[i] * 2 + 1
StructuredBuffer<float> In  : register(t0);
RWStructuredBuffer<float> Out : register(u0);

cbuffer Params : register(b0) { uint count; };

[numthreads(64, 1, 1)]
void main(uint3 tid : SV_DispatchThreadID) {
    uint i = tid.x;
    if (i < count) {
        Out[i] = In[i] * 2.0f + 1.0f;
    }
}

这里 [numthreads(64,1,1)] 表示一个 threadgroup 有 64 个线程。CPU 侧如果 dispatch GG 个 groups，那么最多会启动 64G64G 个线程。通常：

G=⌈N64⌉G = \left\lceil \frac{N}{64} \right\rceil
Dispatch grid：groups × threads Group 0 64 threads Group 1 64 threads Group 2 64 threads 每个线程根据 SV_DispatchThreadID / SV_GroupID / SV_GroupThreadID 决定自己处理哪一份数据。
图 4-4：compute shader 的执行结构：dispatch grid 由多个 threadgroup 组成。

4.5Arithmetic intensity 与 roofline model

判断 kernel 瓶颈时，可先计算 arithmetic intensity（算术强度）：

I=FLOPsbytes movedI = \frac{\text{FLOPs}}{\text{bytes moved}}

II 低表示每读取或写入一个 byte 只完成少量计算，kernel 往往受内存带宽限制，即 memory-bound。II 足够高时，kernel 才可能接近计算单元的峰值吞吐。

roofline model 可以写成：

attainable FLOP/s=min⁡(peak FLOP/s,  I⋅memory bandwidth)\text{attainable FLOP/s} = \min\left(\text{peak FLOP/s},\; I\cdot \text{memory bandwidth}\right)
图 4-5：roofline 示意：左侧受带宽限制，右侧受计算峰值限制。

例子：decode GEMV 为什么常常 memory-bound？

一个线性层 y=xW⊤y=xW^\top，其中 x∈RKx\in\mathbb{R}^{K}，W∈RN×KW\in\mathbb{R}^{N\times K}。 对每个输出元素，需要 KK 次乘加，约 2K2K FLOPs；同时至少要读取 KK 个权重。 如果权重是 f16，则权重读取约 2K2K bytes。因此单个输出的算术强度粗略为：

I≈2K2K=1  FLOP/byteI \approx \frac{2K}{2K} = 1\;\text{FLOP/byte}

这个强度很低。decode GEMV 中每个权重通常只参与一次乘加，性能主要取决于权重读取带宽，而不是乘法器数量。f16 weights、coalesced loads、量化和 KV cache 的意义都可从这一点理解：它们减少传输量或改善访问效率。

小结

本章建立了后续编写 HLSL kernel 所需的性能模型。CPU 负责低延迟控制流，GPU 负责大量同构数据并行；GPU 线程按 SIMT/wave 执行，分支发散和低 occupancy 都会降低吞吐。内存层次决定了许多 LLM kernel 的上限，尤其是 decode 阶段的 GEMV：linear_f16.hlsl 用 64 个线程组成一个 threadgroup 计算一个输出元素，线程沿 KK 维做连续的 f16 读取，并在 groupshared memory 中完成树形归约。D3D12 单轴最多 65535 个 group；LM head 的输出维度超过该限制时，xinfer 使用二维 dispatch grid，并在 shader 中用 idx = gid.y * grid_x + gid.x 还原线性输出索引。

Lab 4纸上估算 matmul 带宽

本 lab 不要求编写代码，目标是用数量级估算判断性能瓶颈。

  1. 设 LM head 是 K=896,N=151936K=896, N=151936，decode 时 M=1M=1。估算一次 forward 的 FLOPs。
  2. 若权重为 f32，估算至少要读取多少 MB 权重；若为 f16 呢？
  3. 用 I=FLOPs/bytesI=\text{FLOPs}/\text{bytes} 估算 f32 与 f16 的 arithmetic intensity。
  4. 假设 GPU 有 800 GB/s 实际有效带宽，估算这个 GEMV 的理论下限耗时。
  5. 解释为什么 GPU argmax 可以显著提升 decode：它减少的是计算、显存访问，还是 CPU/GPU readback？
参考方向

你不需要得到完全精确的数值。系统性能分析的第一步是判断数量级：瓶颈在计算、显存带宽、同步，还是跨设备传输。

思考与练习

基础用一句话解释 thread 与 threadgroup 的区别。

thread 是执行 kernel 的最小单元（一份独立的索引与寄存器）；threadgroup 是一组一起调度、可共享 groupshared（LDS）内存并通过 barrier 同步的线程集合。

基础为什么 GPU 上相邻线程读取相邻地址通常更快？

因为内存事务是按 cache line / 对齐的内存段为单位完成的。当 warp/wave 内相邻线程访问连续地址时，硬件可以把它们合并（coalesce）成少数几次宽内存事务；否则每个线程触发各自的事务，带宽利用率骤降。这正是 kernel 设计要让 W 读取连续的原因。

进阶解释为什么 decode GEMV 的 arithmetic intensity 远低于大 batch GEMM。

arithmetic intensity = FLOPs / bytes。decode 时 M=1M=1，y=xW⊤y=xW^\top 要读入整个 N×KN\times K 权重却只做 N×KN\times K 次乘加（每个权重只用一次），强度约 O(1)O(1)，纯带宽受限。大 batch GEMM 中每个权重被 MM 行复用，FLOPs 增长 MM 倍而权重读取不变，强度约 O(M)O(M)，可进入计算受限区。这就是 decode 慢、需要 f16/coalescing 优化的根因。

进阶如果一个 kernel 需要 threadgroup 内部所有线程先写 LDS 再读 LDS，为什么必须使用 barrier？

因为 threadgroup 内的线程并非严格同步执行（不同 wave 进度不同）。没有 barrier 时，某线程可能在别的线程还没写入它需要的 LDS 槽位前就去读，读到旧值/未定义值，产生竞态。GroupMemoryBarrierWithGroupSync() 保证“所有写已完成且所有线程都到达此点”后再继续读。

挑战为 y=xW⊤y=xW^\top 设计一个 threadgroup 布局，使得权重读取 coalesced，并说明每个线程负责哪些 KK 维元素。

布局：一个 threadgroup 负责一个输出元素 yny_n。以 linear_f16.hlsl 为例，TG=64，线程 tt 处理 f16 pair 索引 t,t+64,t+128,…t,t+64,t+128,\dots，即沿 KK 维跨步累加 xk⋅Wn,kx_k\cdot W_{n,k}。同一轮中相邻线程读取相邻权重 pair，因此 W 读取可 coalesce。每个线程得到 partial sum 后写入 groupshared float partials[64]，再用 GroupMemoryBarrierWithGroupSync() 做树形归约，由 t==0 写出最终 dot product。二维 dispatch grid 只改变输出索引的生成方式，不改变组内归约结构。

← Module 3：自回归生成、KV Cache 与采样教程目录 ↑Module 5：GPU Compute API 版图 →