Ch16 SIMD 与 GPU 编程 · 讲解稿¶

章节定位: Ch16 性能内功。教"如何让 AI 模型跑得快": SIMD (AVX/NEON) + CUDA + Triton + RISC-V + Vulkan。承接: Ch13 §02 架构 → Ch16 编程实现走向: Ch17 §02 LLM 推理 / Ch18 §04 ML 推理篇幅: ~4500 字 / 阅读 13 分钟 / 讲解 30-35 分钟

一 · 一句话总览¶

SIMD/GPU = AI 性能加速的"内功": CPU SIMD (AVX-512 / NEON) 一条指令多数据, GPU CUDA / Triton 大规模并行, TPU 脉动阵列, RISC-V 边缘部署, Vulkan 跨平台。AI 工程师需理解至少一种。

二 · 心智模型¶

CPU SIMD: 一心多用, 处理连续数据快
GPU SIMT: 千军万马, 数据并行
TPU: 流水线工厂, matmul 极快
RISC-V: 嵌入式, 低功耗

三 · 章节主线 (5 段)¶

第 1 段 · C++ 与 ML 框架基础 (10 分钟)¶

3.1.1 为什么 AI 框架用 C++¶

Python 慢, 性能关键用 C++/CUDA。PyTorch 栈:

Python API: 用户友好
ATen (C++): Tensor 操作核心
autograd: 自动微分
cuDNN / cuBLAS: GPU 加速

3.1.2 SIMD 原理¶

SIMD = Single Instruction Multiple Data: 一条指令同时操作多个数据。

e.g., AVX-512 一次算 16 个 FP32, 理论 16× 加速 (实际受内存限制)。

3.1.3 SIMT 原理¶

SIMT = Single Instruction Multiple Threads: GPU 上 32 个 thread (warp) 同指令, 但每 thread 有独立数据。

类比: SIMD 像 16 个工人做同一动作; SIMT 像 32 个学生做同一作业。

第 2 段 · 硬件基础与 Roofline (10 分钟)¶

3.2.1 内存层级¶

层级	延迟	带宽
寄存器	< 1 ns	TB/s
L1/L2	~1-30 ns	TB/s
DRAM	~100 ns	50 GB/s
HBM (GPU)	~400 ns	3 TB/s

3.2.2 Roofline 模型¶

\[P = \min(\pi_{\text{peak}}, I \times B_{\text{mem}})\]

\(I\)	含义
计算强度 (ops / byte)	算法特性

两类 kernel: - 计算受限 (compute-bound): \(I\) 高, matmul - 带宽受限 (memory-bound): \(I\) 低, attention, embedding

3.2.3 内存墙¶

CPU / GPU 算力增长 > 内存带宽增长。

LLM 推理就是 memory-bound — 加载 70 GB 权重, GPU 带宽 3 TB/s → 至少 23 ms。

3.2.4 Flash Attention 思路¶

不存中间 \(O(n^2)\) attention matrix, 分块计算 → 减少 IO。

第 3 段 · CPU SIMD: AVX 与 NEON (10 分钟)¶

3.3.1 x86 AVX¶

指令集	宽度	寄存器
SSE	128	XMM
AVX	256	YMM
AVX-512	512	ZMM (32 个)

#include <immintrin.h>

void vec_add(float* a, float* b, float* c, int n) {
    for (int i = 0; i < n; i += 8) {
        __m256 va = _mm256_load_ps(a + i);
        __m256 vb = _mm256_load_ps(b + i);
        __m256 vc = _mm256_add_ps(va, vb);
        _mm256_store_ps(c + i, vc);
    }
}

gcc -mavx2 -O3 启用 AVX2。

3.3.2 ARM NEON¶

128 bit, 16 个 Q 寄存器。

#include <arm_neon.h>

void vec_add(float* a, float* b, float* c, int n) {
    for (int i = 0; i < n; i += 4) {
        float32x4_t va = vld1q_f32(a + i);
        float32x4_t vb = vld1q_f32(b + i);
        float32x4_t vc = vaddq_f32(va, vb);
        vst1q_f32(c + i, vc);
    }
}

Apple Silicon + 手机普遍用。

3.3.3 FMA¶

Fused Multiply-Add: a*b+c 一条指令, 提升精度 + 性能。

第 4 段 · GPU CUDA 编程 (12 分钟)¶

3.4.1 编程模型¶

Host (CPU) → launch kernel<<<grid, block>>>(...) → Device (GPU)

__global__ void matmul(float* A, float* B, float* C, int M, int N, int K) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < M && col < N) {
        float sum = 0;
        for (int k = 0; k < K; k++) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

// Host
dim3 block(16, 16);
dim3 grid((N + 15) / 16, (M + 15) / 16);
matmul<<<grid, block>>>(A, B, C, M, N, K);

3.4.2 内存层级¶

类型	速度	用途
Register	极快	局部变量
Shared	快 (~100ns)	block 内通信
L1/L2	中	cache
Global (HBM)	慢	主存

__shared__ float tile[16][16];
// 加载
tile[ty][tx] = A[row * K + k * 16 + tx];
__syncthreads();
// 通信
sum += tile[ty][k] * tile[k][tx];

3.4.3 优化技巧¶

优化	做法
coalesced	32 thread 访问连续 128 字节
bank conflict	shared mem 分 32 bank
occupancy	调整 block 大小
tensor core	用 mma 指令

3.4.4 Flash Attention CUDA 优化¶

__global__ void flash_attn_kernel(
    float* Q, float* K, float* V, float* O,
    int N, int D, int Br, int Bc
) {
    // 分块 Q, K, V
    // 在线 softmax
    // 不存完整 attention matrix
}

第 5 段 · Triton / TPU / RISC-V / Vulkan (8 分钟)¶

3.5.1 Triton¶

OpenAI 推出的 Python-like DSL:

import triton
import triton.language as tl

@triton.jit
def vec_add(A, B, C, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    a = tl.load(A + offsets, mask=offsets < N)
    b = tl.load(B + offsets, mask=offsets < N)
    tl.store(C + offsets, a + b, mask=offsets < N)

优势: Python 简洁 + 自动优化 + 性能接近 CUDA。

3.5.2 TPU¶

Google 自研, 脉动阵列 (systolic):

数据从左流入, 从上流入, 一周期算 128×128 矩阵乘

代	FLOPS	内存
v3	123 TFLOPS	32 GB HBM
v4	275 TFLOPS	32 GB
v5e	197 TFLOPS	16 GB

适用: 大模型训练 (Gemini, PaLM)。

3.5.3 RISC-V¶

开放 ISA, 中国 + 美国生态。

V 扩展: 128-2048 bit 向量
厂商: SiFive, 平头哥玄铁
AI 应用: 边缘部署 (Llama.cpp 量化)

3.5.4 Vulkan Compute¶

跨平台 GPU API:

#version 450
layout(local_size_x = 64) in;
layout(set = 0, binding = 0) buffer A { float data[]; };

void main() {
    uint idx = gl_GlobalInvocationID.x;
    data[idx] *= 2.0;
}

优势: Windows / Linux / Android / macOS (via MoltenVK) 统一。

四 · 易错点 (5 条)¶

AVX 宽度: AVX-2 是 256, AVX-512 是 512, 别混。
CUDA warp: 32 thread 固定, block 32-1024。
shared memory bank: 32 bank, 同 bank 冲突。
Triton 不万能: 极端优化仍需 CUDA。
量化精度: INT4 节省内存但损失精度。

五 · 跨章连接¶

章节	关联
Ch13 §02 架构	硬件原理
Ch13 §04 CUDA	GPU 编程
Ch17 §02 LLM 推理	性能优化
Ch15 §05 K8s	GPU 调度

六 · 自测 5 题¶

SIMD vs SIMT 区别?
AVX-512 寄存器?
CUDA warp 大小?
Flash Attention 优化点?
Triton 与 CUDA 区别?

七 · 一句话总结¶

SIMD/GPU = AI 性能内功。CPU SIMD (AVX/NEON) 适用于单核优化, GPU CUDA + Triton 适用于并行 matmul, TPU 适用于大模型训练, RISC-V 适用于边缘部署, Vulkan 适用于跨平台。Ch16 是 AI 性能工程师的必修课。

八 · 板书建议¶

[板书 1: SIMD]
   AVX-512: 512 bit, 16×FP32
   NEON: 128 bit, 4×FP32

[板书 2: Roofline]
   P = min(π, I × B)
   计算 vs 带宽受限

[板书 3: CUDA]
   warp (32) / block / grid
   shared mem + coalesced

[板书 4: Triton]
   Python-like
   JIT + autotune

[板书 5: TPU]
   脉动阵列
   128×128 MXU

九 · 2024-2026 硬件 + 编译器前沿 (前沿补丁)¶

9.1 NVIDIA H100 / B200 / GB200 (Grace Blackwell)¶

H100 (2023-2024): - 80GB HBM3, 3.35 TB/s 带宽 - FP8 tensor core - Transformer Engine - 989 TFLOPs FP16

B200 (2024-2025): - 192GB HBM3e, 8 TB/s - FP4 tensor core (新增) - 2.5× 训练 / 30× 推理 vs H100 - 208 TFLOPs FP4

GB200 (Grace Blackwell): - 72 颗 CPU + 36 颗 B200 - NVLink 900 GB/s - 1.4 ExaFLOPS FP4

9.2 Async Tensor Core / WGMMA¶

WGMMA (Warp Group Matrix Multiply Accumulate): - Hopper 异步 tensor core - 1 个 warp group (4 warp) 算大 matmul - Flash Attention 3 核心: 利用异步重叠数据搬运

PTX 7.0+ 指令: - wgmma.mma_async.sync - cp.async.bulk (TMA: Tensor Memory Accelerator) - cluster (跨 SM 共享)

9.3 Triton 3.0 (2024)¶

主要更新: - Persistent kernels: 自动持久化 - TMA 支持: 直接 Hopper 内存拷贝 - 自动 warp specialization - 高效 FP8 kernels - 描述符化: 显式 shape

@triton.jit
def attn_kernel(Q, K, V, O, ...):
    # TMA 加载
    q = tl.make_block_ptr(Q, ...)
    q = tl.load(q, ...)
    # 异步计算
    ...

9.4 CUTLASS 3.x (2024)¶

CUTLASS (CUDA Templates for Linear Algebra Subroutines): - CUTLASS 3.0: Blackwell 优化 - 组合格式 (mxfp8, mxfp4) - Python DSL (CUTLASS DSL) - Python 端到端调用

9.5 内核编程新范式¶

范式	优势
Persistent Threads	减少启动开销
Warp Specialization	Producer-consumer
Stream-K	负载均衡
CUDA Graphs	减少 CPU-GPU 同步
TMA (Tensor Memory Accelerator)	Hopper 异步加载
Cluster Launch	跨 SM 同步

9.6 Apple Silicon (MLX)¶

MLX / MLX-LM (Apple, 2023-2024): - Apple Silicon 优化 - Unified Memory Architecture (UMA) - 与 PyTorch 互操作 - Llama / Mistral 一键部署

9.7 国产硬件¶

芯片	公司	算力
昇腾 910B/910C	华为	FP16 780 TFLOPs
寒武纪 MLU370	寒武纪	INT8 256 TFLOPs
壁仞 BR104	壁仞	FP16 1024 TFLOPs
摩尔线程 MTT S5000	摩尔线程	FP32 50 TFLOPs
天数智芯天垓 100	天数智芯	FP16 295 TFLOPs

9.8 2024-2026 时间线¶

timeline
    2023 : H100 大规模部署, Mamba 1
    2024Q1 : Triton 3.0, WGMMA 普及
    2024Q2 : PyTorch 2.4, CUDA 12.5
    2024Q3 : B200 发布, 中国禁 H100/B200
    2024Q4 : CUTLASS 3.0, MLX 1.0
    2025Q1 : GB200 出货, 国产替代
    2025Q2 : AMD MI300X 主流
    2025Q3 : 推理 FPGA 加速
    2025Q4 : 4PF 桌面 GPU
    2026H1 : 端侧 4bit 推理普及

9.9 一句话总结¶

2024-2026 硬件战 = 4 个战场: NVIDIA H100/B200 主导高端 + Apple MLX 端侧 + 国产替代 (昇腾/寒武纪) + AMD MI300X 竞争。软件栈: Triton 3 + CUTLASS 3 + WGMMA/TMA 把硬件用满。

十 · 参考资料¶

"Programming Massively Parallel Processors" (CUDA 经典)
NVIDIA CUDA Programming Guide, 2024
Triton 3.0 文档, 2024
"Computer Architecture" (H&P)
NVIDIA, "B200 / GB200", 2024
Apple, "MLX", 2024
NVIDIA, "PTX 7.0 ISA", 2024
Khronos, "Vulkan 1.3 / 1.4", 2024
CUTLASS 3.0 文档, 2024
Ch13 §02 架构 (本套书)
Ch17 §02 LLM 推理 (本套书)