Jaspery

infra规划

2026-08-12T13:15:46.000Z

AI Infra 学习与研究路线

1. 总体目标

目标定位

未来 2～3 年主攻：

分布式训练与通信优化 + LLM 推理 Infra

最终形成：

                   AI Infra
                      │
         ┌────────────┴────────────┐
         ↓                         ↓
  Distributed Training        LLM Inference
         │                         │
  ┌──────┼──────┐            ┌─────┼──────┐
  ↓      ↓      ↓            ↓     ↓      ↓
CUDA   NCCL  Megatron       vLLM  KVCache  PD分离
         │      │                  │
      DeepSpeed │              Spec Decode
                │                  │
                └────────┬─────────┘
                         ↓
                   GPU / RDMA
                         ↓
                      Network

最终目标不是成为单纯的 CUDA 工程师、网络工程师或大模型应用开发者，而是：

懂 GPU、懂分布式系统、懂 LLM、懂网络的 AI Systems / AI Infra Engineer。

2. 总体技术栈

2.1 GPU 底层

需要掌握：

GPU Architecture
SM
Thread
Warp
Block
Grid
Register
Shared Memory
L1 Cache
L2 Cache
HBM
Tensor Core
CUDA Core
CUDA Stream
CUDA Event
Memory Coalescing
Occupancy
Warp Divergence
Register Pressure

重点理解：

HBM
 ↓
L2 Cache
 ↓
L1 / Shared Memory
 ↓
Register
 ↓
CUDA Core / Tensor Core

核心问题：

为什么一个 CUDA Kernel 会慢？

需要能够从以下几个方面分析：

Compute
Memory Bandwidth
Memory Access
Occupancy
Warp Divergence
Register Pressure
Shared Memory

3. Phase 0：CUDA 与 GPU 基础

目标

建立 GPU 编程和 GPU 执行模型基础。

学习内容

GPU Architecture

掌握：

CPU vs GPU
SM
Warp
Thread
Block
Grid
Register
Shared Memory
L1 / L2
HBM
Tensor Core

CUDA

能够自己实现：

Vector Add
Reduction
Softmax
RMSNorm
LayerNorm
Matrix Multiplication
Attention

重点理解：

Thread
 ↓
Warp
 ↓
Block
 ↓
SM
 ↓
GPU

以及：

Global Memory
 ↓
L2
 ↓
L1 / Shared Memory
 ↓
Register

阶段产出

至少完成：

CUDA Vector Add
CUDA Reduce
CUDA Softmax
CUDA RMSNorm
CUDA MatMul
简单 Attention Kernel

并使用 Benchmark 分析：

Latency
Throughput
Memory Bandwidth
FLOPS
Occupancy

4. Phase 1：集合通信

目标

理解 GPU 之间是如何进行数据交换的。

必须掌握：

AllReduce
AllGather
ReduceScatter
Broadcast
AllToAll
AllToAllv

重点理解：

Ring AllReduce

1
2
3

GPU0 → GPU1 → GPU2 → GPU3
 ↑                       ↓
 └───────────────────────┘

理解：

Latency
Bandwidth
Message Size
GPU Number
Topology

之间的关系。

5. Phase 2：NCCL

这是整个训练通信方向的核心。

需要掌握

NCCL Communicator
NCCL Topology
NCCL Channel
NCCL Ring
NCCL Tree
NCCL Transport
CUDA IPC
NVLink
PCIe
RDMA
InfiniBand
RoCE

重点研究：

GPU 拓扑

例如：

GPU0 ─ NVLink ─ GPU1
 │               │
PCIe            PCIe
 │               │
NIC0            NIC1

需要理解：

GPU 到 GPU 的实际通信路径是什么？

以及：

GPU
 ↓
NVLink / PCIe
 ↓
NIC
 ↓
RDMA
 ↓
Switch
 ↓
NIC
 ↓
GPU

7. Phase 3：Megatron-LM

目标

理解大模型分布式训练。

7.1 Data Parallel

GPU0 → Model Copy
GPU1 → Model Copy
GPU2 → Model Copy
GPU3 → Model Copy

        ↓

Gradient AllReduce

重点：

Data Parallel
Gradient Synchronization
AllReduce

7.2 Tensor Parallel

        Linear
          │
  ┌───────┼───────┐
  ↓       ↓       ↓
GPU0    GPU1    GPU2

重点：

Tensor Parallel
AllGather
ReduceScatter
AllReduce

7.3 Pipeline Parallel

GPU0        GPU1        GPU2        GPU3

Layer1      Layer5      Layer9      Layer13
Layer2      Layer6      Layer10     Layer14
Layer3      Layer7      Layer11     Layer15
Layer4      Layer8      Layer12     Layer16

重点：

Pipeline Parallel
Pipeline Bubble
1F1B
Interleaved Pipeline

7.4 3D Parallelism

最终理解：

1	`DP × TP × PP`

例如：

DP = 4
TP = 4
PP = 4

4 × 4 × 4 = 64 GPUs

核心问题：

为什么 GPU 数量增加后，训练效率不会线性增加？

8. Phase 4：DeepSpeed

重点掌握：

1
2
3

ZeRO-1
ZeRO-2
ZeRO-3

理解：

1
2
3

Parameters
Gradients
Optimizer States

如何在 GPU 之间切分。

进一步学习：

ZeRO
Offload
Activation Checkpointing
Communication Overlap

9. 第一阶段科研方向：训练通信优化

不建议一开始直接研究：

万卡集群通信优化

而应该从小规模实验开始。

9.1 拓扑感知通信优化

研究：

GPU Topology
      ↓
Communication Path
      ↓
Collective Algorithm
      ↓
Communication Performance

例如比较：

1
2
3

PCIe
NVLink
RDMA

不同拓扑下：

1
2
3

Latency
Bandwidth
Scaling Efficiency

9.2 计算与通信重叠

目标：

Compute
   │
   ├───────────────┐
   ↓               ↓
  GEMM          AllReduce
   │               │
   └──────┬────────┘
          ↓
        Next

通过：

CUDA Stream
NCCL
Pipeline
Asynchronous Communication

实现：

Communication Overlap

重点指标：

Training Throughput
Communication Time
GPU Utilization
Scaling Efficiency

10. Phase 5：LLM Inference

训练方向达到一定程度之后进入推理。

推荐学习顺序：

vLLM
 ↓
PagedAttention
 ↓
KV Cache
 ↓
Continuous Batching
 ↓
Tensor Parallel
 ↓
Speculative Decoding
 ↓
PD Disaggregation
 ↓
KV Cache Transfer
 ↓
Inference Network

11. Phase 6：vLLM

首先把 vLLM 的整体执行链路搞懂：

Request
   ↓
Tokenizer
   ↓
Scheduler
   ↓
Continuous Batch
   ↓
Model Runner
   ↓
Attention
   ↓
KV Cache
   ↓
GPU

重点研究：

Scheduler
Continuous Batching
KV Cache Manager
PagedAttention
GPU Worker
Model Runner
Memory Management

12. Phase 7：KV Cache

KV Cache 是推理 Infra 的核心。

Transformer Attention：

1 2	`K = XWk V = XWv`

生成新 Token：

1 2	`K_new V_new`

不断追加：

KV Cache

Token 1
Token 2
Token 3
...
Token N

Context 越长：

KV Cache
    ↓
GPU HBM
    ↓
Memory Pressure
    ↓
OOM

因此需要研究：

PagedAttention
KV Cache Paging
KV Cache Offloading
KV Cache Compression
KV Cache Sharing
KV Cache Eviction
KV Cache Placement

13. Phase 8：Continuous Batching

理解 LLM Decode 的动态 Batch。

例如：

Request A → Token
Request B → Token
Request C → Token
Request D → EOS
Request E → Token

下一轮：

A B C E

重新形成 Batch。

核心：

每一轮 Decode 都可能由不同请求组成动态 Batch。

重点指标：

Throughput
TTFT
TPOT
GPU Utilization
Batch Size
Scheduling Overhead

14. Phase 9：Speculative Decoding

基本结构：

  Draft Model
       ↓
t1 t2 t3 t4
       ↓
 Target Model
       ↓
   Verify
       ↓
Accept / Reject

传统生成：

Target
 ↓
1 Token
 ↓
Target
 ↓
1 Token
 ↓
Target

Speculative Decoding：

Draft
 ↓
多个 Token
 ↓
Target 一次验证

重点研究：

Acceptance Rate
Draft Latency
Target Latency
Memory Overhead
Throughput
TTFT
TPOT

15. Phase 10：Disaggregated Inference

这是重点方向。

传统：

1
2
3

Prefill + Decode
       ↓
    同一 GPU

分离式推理：

             Network
                │
       ┌────────┴────────┐
       ↓                 ↓
Prefill Cluster     Decode Cluster
       │                 │
    GPU GPU GPU        GPU GPU GPU
       │
       ↓
    KV Cache
       │
       └──────→ Network
                    ↓
                 Decode

即：

Prefill 和 Decode 分离部署。

16. 为什么 PD Disaggregation 特别适合通信背景

Prefill：

输入 Prompt
    ↓
Prefill GPU
    ↓
产生 KV Cache

然后：

KV Cache
    ↓
Network
    ↓
Decode GPU

这会产生大量系统问题：

网络带宽

KV Cache Size
      ↓
Transfer Bandwidth
      ↓
Transfer Latency

网络拥塞

多个 Prefill
      ↓
KV Transfer
      ↓
Network Congestion

GPU 调度

1
2
3

Prefill GPU
     ↓
哪个 Decode GPU？

KV Cache Placement

KV Cache
   ↓
放在哪里？
   ↓
哪个 GPU？
   ↓
什么时候迁移？

因此：

PD Disaggregation = LLM + GPU + Network + RDMA + Distributed Systems

非常适合网络背景。

17. 第二阶段科研方向：Inference + Network

重点考虑：

Network-aware LLM Inference

整体结构：

         LLM Inference
              │
   ┌──────────┴──────────┐
   ↓                     ↓
Prefill                Decode
   │                     │
KV Cache                 │
   │                     │
   └──────→ Network ─────┘
               │
              RDMA
               │
           Scheduler

可以研究：

Network-aware Scheduling
KV Cache Placement
KV Cache Routing
RDMA KV Transfer
Congestion-aware Inference
Topology-aware Inference
GPU-aware Scheduling
Prefill/Decode Load Balancing

18. 推荐的 4 个科研项目

项目 1：NCCL Topology Benchmark

目标

研究不同 GPU / 网络拓扑下集合通信性能。

测试：

1
2
3

2 GPU
4 GPU
8 GPU

以及：

1
2
3

PCIe
NVLink
RDMA

测量：

1
2
3

Latency
Bandwidth
Scaling Efficiency

最终产出

NCCL Benchmark
+
Topology Analysis
+
Performance Model

19. 项目 2：Megatron 通信与计算 Overlap

Baseline

GEMM
 ↓
AllReduce
 ↓
GEMM
 ↓
AllReduce

Optimization

1
2
3

GEMM ────────────────
      AllReduce ────────
            GEMM ─────────

利用：

CUDA Stream
NCCL
Async Communication

优化：

1	`Communication / Computation Overlap`

最终测：

Training Throughput
GPU Utilization
Communication Ratio
Scaling Efficiency

20. 项目 3：KV Cache 优化

研究：

KV Cache
    ↓
Memory Pressure
    ↓
Eviction
    ↓
Offload / Migration

重点：

KV Cache Placement
KV Cache Eviction
KV Cache Offloading
KV Cache Compression

测试：

1
2
3

Short Context
Medium Context
Long Context

指标：

TTFT
TPOT
Throughput
GPU Memory Usage
Cache Hit Rate

21. 项目 4：PD Disaggregation + RDMA

最推荐作为核心科研项目。

架构：

  ┌──────────────┐
  │    Client    │
  └──────┬───────┘
         ↓
┌─────────────────┐
│ Prefill Cluster │
└────────┬────────┘
         ↓
     KV Cache
         ↓
   RDMA Network
         ↓
┌─────────────────┐
│ Decode Cluster  │
└────────┬────────┘
         ↓
     Output

测试：

1
2
3

TCP
vs
RDMA

不同网络：

1 Gbps
10 Gbps
25 Gbps
100 Gbps

研究：

KV Transfer Latency
TTFT
TPOT
Throughput
GPU Utilization
Network Utilization

进一步做：

Network-aware KV Cache Scheduling

这个方向非常适合作为论文/毕业课题。

22. 12 个月学习规划

时间	主攻内容	目标
0～2 月	CUDA + GPU Architecture	掌握 GPU 执行模型
2～4 月	Collective Communication + NCCL	掌握 GPU 通信
4～6 月	Megatron + DeepSpeed	掌握分布式训练
6～8 月	vLLM + KV Cache	掌握 LLM 推理
8～10 月	Speculative Decoding	掌握推理优化
10～12 月	PD Disaggregation + RDMA	形成科研方向

23. 每个阶段的学习方法

不要：

看论文
 ↓
看源码
 ↓
继续看论文
 ↓
继续看源码

而应该：

理论
 ↓
最小 Demo
 ↓
源码
 ↓
Benchmark
 ↓
性能分析
 ↓
修改
 ↓
实验

例如学习 NCCL：

AllReduce 原理
      ↓
自己实现简单 AllReduce
      ↓
NCCL API
      ↓
NCCL 源码
      ↓
NCCL Benchmark
      ↓
改变 Message Size
      ↓
改变 GPU 数量
      ↓
分析 Scaling

24. 实验资源要求

不需要一开始就有万卡集群。

建议最低：

2 GPU
+
CUDA
+
NCCL

比较理想：

4～8 GPU
+
NVLink / PCIe
+
RDMA
+
25/100GbE

研究万卡问题时可以通过：

Small-scale Benchmark
+
Performance Model
+
Simulation
+
Extrapolation

研究大规模系统。

25. 未来需要掌握的技术栈

GPU

CUDA
CUDA Graph
Triton
Tensor Core
FlashAttention
Kernel Optimization

Distributed Training

PyTorch Distributed
NCCL
Megatron-LM
DeepSpeed
ZeRO
FSDP
DP
TP
PP
EP

LLM Inference

vLLM
SGLang
TensorRT-LLM
PagedAttention
KV Cache
Continuous Batching
Speculative Decoding
PD Disaggregation

Network

RDMA
RoCE
InfiniBand
NVLink
PCIe
NCCL Network
Topology
Congestion Control

System

Linux
C/C++
Python
CUDA
Docker
Kubernetes
Distributed Systems
Performance Profiling

26. 就业方向

最终可以对应：

AI Infra

LLM Inference Engineer
AI Infra Engineer
Distributed Training Engineer
GPU Cluster Engineer

GPU / Systems

1
2
3

CUDA Engineer
GPU Optimization Engineer
Distributed Systems Engineer

Network + AI

AI Network Engineer
Inference Network Engineer
RDMA Engineer
GPU Network Engineer

27. 技术能力优先级

建议优先级：

★★★★★
LLM Inference / AI Infra

★★★★★
Distributed Training / NCCL

★★★★☆
CUDA / GPU Optimization

★★★★☆
RDMA / GPU Network

★★★☆☆
DeepSpeed / Megatron 深度源码

★★★☆☆
Speculative Decoding

★★☆☆☆
纯 LLM Application / Agent

对于个人背景而言，CUDA + NCCL + RDMA + LLM Inference 的组合尤其值得建立。

28. 最终能力画像

最终希望形成：

             AI Systems Engineer
                    │
       ┌────────────┴────────────┐
       ↓                         ↓
Distributed Training        LLM Inference
       │                         │
   Megatron                  vLLM/SGLang
       │                         │
  DeepSpeed                  KV Cache
       │                         │
     NCCL                  Spec Decode
       │                         │
       └────────────┬────────────┘
                    ↓
                  CUDA
                    ↓
                   GPU
                    ↓
             NVLink / PCIe
                    ↓
                RDMA/NIC
                    ↓
                 Network

核心能力：

能够从 LLM 上层 workload 一直追到底层 GPU、通信库、RDMA 和网络，并定位系统瓶颈、设计优化方案、完成实验验证。

29. 最终主线

整个研究生阶段可以始终围绕下面这条主线：

CUDA
 ↓
GPU Architecture
 ↓
Collective Communication
 ↓
NCCL
 ↓
Megatron / DeepSpeed
 ↓
Distributed Training
 ↓
vLLM
 ↓
KV Cache
 ↓
Speculative Decoding
 ↓
PD Disaggregation
 ↓
RDMA
 ↓
Inference Network
 ↓
Network-aware LLM Inference

最终形成一个明确的研究标签：

LLM Systems / AI Infra + GPU Communication + RDMA

LLM中的数学原理及infra优化

2026-08-09T11:05:17.000Z

引言

本文将拆解大模型中几个核心操作（RMSNorm、Softmax、Causal Mask、Sampling）背后的数学与 Infra 优化逻辑。看完你会发现，Infra 优化，本质上就是在用数学上的等价变换，或者对精度的适度妥协，去换取更高的硬件利用率和极致的推理速度。

1.RMSNorm - 均方根归一化

大语言模型（Transformer 结构）通常包含数十甚至上百个堆叠的隐藏层（如 Transformer 结构）。输入张量（Tensor）在经过连续的矩阵乘法和加法操作后，其数值的分布范围会发生剧烈的变化。

这种数值大小的不可控会导致两个严重的工程和算法问题：

算法收敛困难：数值变得过大或过小会影响训练的稳定性：极端数值要么会落入激活函数的饱和区导致梯度消失，要么会顺着网络不受控地放大引发梯度爆炸或硬件溢出。
硬件层面的溢出与截断：在当前主流的 GPU 推理和训练中，为了追求极致的吞吐，底层计算会使用低精度浮点格式（如 FP16 或 BF16）。FP16的数值范围小，容易溢出，BF数值范围大，但是精度低（大数与小数加和时容易因为四舍五入将小数截断）。

这两者都会导致模型输出乱码或训练彻底崩溃。为了保证大模型在深层网络中的数值稳定性，研究人员在架构中引入了特征归一化（Normalization）机制（例如 LayerNorm、RMSNorm）。其核心目的，是在数据的层间传递过程中，对其数值分布进行强制的缩放与平移，将其约束在一个标准、安全的物理尺度内, 防止方差膨胀引发的溢出。

神经网络真正关心的并非数值的绝对大小，而是特征之间的相对差异，故可以平移。

1.1 LayerNorm

LayerNorm 的核心思想是对同一个 Token 内的所有特征维度（hidden size, d）进行标准化，使其均值为 0，方差为 1。

与LayerNorm对应的是BatchNorm, LayerNorm/RMSNorm 是严格在 Token 级别（Hidden Size 维度） 闭环的。无论外部的 Batch Size 怎么变，无论旁边并行的请求是写诗还是写代码，每一个 Token 自身的归一化结果不会发生变化。训推一致。

git add .git commit -m "新文章"git pushbash

数学公式

$y = \frac{x-\mu}{\sqrt{\sigma^2+\epsilon}} \odot \gamma + \beta$

符号说明：

$\odot$ ：哈达玛积（逐元素相乘）
$\mu$ ：均值， $\sigma^2$ ：方差， $\epsilon$ ：防止分母为 0 的极小常数
$\gamma$ ：缩放参数， $\beta$ ：偏移参数

标准化并做仿射变换（乘以可学习参数 γ，加上偏置 β ，为了增强特征表达能力）。

Infra视角

LayerNorm 是一个典型的 Memory-bound（访存密集型） 算子。它的计算包含了两次 Global Reduction（全局规约）操作。

最致命的是数据依赖：你必须先完整遍历一次数据算出均值 $\mu$ ，然后才能用 $\mu$ 去遍历第二次算方差 $\sigma^2$ 。

在 GPU 上，这意味着：

更复杂的线程同步逻辑；
在 Kernel 未极致融合时，需要多次往返读写 HBM（显存），极大地浪费了宝贵的内存带宽。

1.2 RMSNorm (Root Mean Square Normalization)

RMSNorm 的作者（Biao Zhang 等人，2019）通过实验发现：LayerNorm 之所以有效，主要是因为缩放（Scaling，即除以标准差）的作用，而平移（Mean-centering，即减去均值 μ）对模型收敛的贡献微乎其微。

既然均值没用，那就直接砍掉它。放弃了计算均值，只保留对向量 RMS 尺度的归一化。当然也有说法：LayerNorm 强行减均值（取平均归零）这个动作，其实是在人为地阉割模型的表达能力。

数学公式

其中均方根（RMS）的计算：

$\text{RMS}(x) = \sqrt{\frac{1}{n}\sum_{i=1}^{n}x_i^2}$

$\hat{x}_i = \frac{x_i}{\text{RMS}(x) + \epsilon} \cdot \gamma$

其中 $\epsilon = 10^{-6}$ ，即代码默认值：eps = 1e-6。实际运行值取决于具体模型的 config.rms_norm_eps 配置，通常为 1e-5 或 1e-6。 $\epsilon$ 的作用是防止除以零，保证数值稳定性，其值非常小，对最终计算结果的影响微乎其微。

Infra 视角

打破数据依赖：RMSNorm 不需要算 $\mu$ ，直接计算每个元素的平方和即可。这意味着只需要一次单向的 Reduction 操作。
极致的访存优化：在编写 Triton 或 CUDA Kernel 时，RMSNorm 可以非常丝滑地在一个 Block 内完成「数据加载 (SRAM) → 平方求和 → 广播 → 缩放」的流水线，中间变量极少。
计算量减少：省去了大量减法操作。

主流模型在使用 RMSNorm 时，通常连后处理的偏置项 $\beta$ 也一并去掉了（即无 Bias 线性层），进一步减少了参数加载和 element-wise 加法的开销。

值得注意的是：相对 LayerNorm，Bias 也被去掉了：RMSNorm 常常只保留 $\gamma$ 而去掉 $\beta$ ，这不是数学上必然要求，而是现代 LLM 架构中的经验选择。它通常与无 bias Linear、Pre-Norm 残差结构、SwiGLU 等设计共同出现，整体上减少参数与访存，同时保持效果。

也有类似说法：LayerNorm 包含减均值，因此后面紧跟的线性层加 Bias 是有意义的。而 RMSNorm 砍掉了平移，只做纯粹的尺度缩放，如果它后面紧跟的 Linear 层仍保留 Bias，就破坏了 RMSNorm 抛弃绝对中心、只维持相对尺度的初衷。

当今的主流开源大模型不仅仅 RMSNorm 去掉了 $\beta$ ，而是几乎所有的 Linear 层都去掉了 Bias：

q_proj, k_proj, v_proj, o_proj 没有 Bias
MLP 的 gate_proj, up_proj, down_proj 也没有 Bias

对此有很多解释：

训练更稳定：
“No biases were used in any of the dense kernels or layer norms. We found this to result in increased training stability for large models.” —— PaLM 论文
架构冗余：RMSNorm 后紧接 Linear 时，bias 的位移作用会被下一个 Norm 的 $\gamma/\beta$ 吸收；SwiGLU 的门控本身也提供了类似 bias 的自由度。bias 在现代架构里已被架构本身替代。
Infra 友好：少一次 add 与 bias load。

但要我说，就是实验有效 + Infra 友好——专业的说法，哈哈：从工程视角看，这类设计往往并非单一数学原则的必然结果，而是效果、稳定性、实现成本和硬件效率共同权衡后的经验选择。

LayerNorm vs RMSNorm（底层实现视角）

LayerNorm 可借助 $\mathbb{E}[X^2] - (\mathbb{E}[X])^2$ 公式与 Kernel 融合，实现只需访问一次 HBM。而 RMSNorm 进一步斩断了均值计算，压缩了 SRAM 占用和 ALU 指令周期。

LayerNorm 虽然理论上需要均值和方差两步，但在高性能实现中并不必然需要两次 HBM 访问：可以在一次 HBM load 中同时累计 $\sum x$ 和 $\sum x^2$ ，甚至使用 Welford 算法提升数值稳定性。但相较 RMSNorm，LayerNorm 仍需要维护均值相关统计量，并在归一化阶段执行额外的减均值操作，因此寄存器压力、规约状态、ALU 指令数都更高。

RMSNorm 相对于 LayerNorm 的收益：

减少了寄存器 / SRAM 的占用
节省了大量的 ALU（逻辑运算单元）指令，特别是消除了对全部元素的减法（减去均值，element-wise）操作

其实在当今主流的 Fused CUDA / Triton Kernel 中，LayerNorm 也是可以做到单次 HBM 访存（1-Pass）的。在数学上，方差可以等价展开为：

$\text{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2$

在 GPU 寄存器 / SRAM 层面，我们在单一的一个 Block 遍历输入张量时，可以同时累计 $\sum x$ 和 $\sum x^2$ 。由于大模型的 Hidden Size（如 4096 或 8192）对应的字节数（约 8–16 KB）完全可以被塞进单个 SM 的 Shared Memory 中，因此无论 LayerNorm 还是 RMSNorm，现代算子在 HBM 层面都是只读一遍、写一遍。

虽然 $\text{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2$ 能实现 1-pass，但在 FP16 或 BF16 精度下，如果 $\mathbb{E}[X^2]$ 与 $(\mathbb{E}[X])^2$ 的值非常接近，相减容易引发灾难性抵消（Catastrophic Cancellation），导致方差精度丢失甚至计算出负数（最后开根号出 NaN）。因此在实际的 Kernel（如 Apex 或 Triton 内部）中，有时会采用 Welford 算法来兼顾 1-pass 和数值稳定性，或者在累加时强制转换到 FP32 进行计算。

$\text{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2$ 的推导逻辑

假设数据的平均值为 $\mu$ （即 $\mu = \mathbb{E}[X]$ ）。方差的原始定义是：每个数减去平均值的平方，再求平均。即：

$\text{Var}(X) = \mathbb{E}[(X - \mu)^2]$

根据 $(a - b)^2 = a^2 - 2ab + b^2$ ，展开得：

$\begin{aligned}\text{Var}(X)&= \mathbb{E}[X^2 - 2X\mu + \mu^2] \\&= \mathbb{E}[X^2] - \mathbb{E}[2X\mu] + \mathbb{E}[\mu^2]\end{aligned}$

$\mathbb{E}[2X\mu] = 2\mu \cdot \mathbb{E}[X]$ ，而 $\mathbb{E}[X] = \mu$ ，所以 $\mathbb{E}[2X\mu] = 2\mu^2$
$\mathbb{E}[\mu^2]$ ：因为 $\mu^2$ 是常数，常数的期望仍是它本身，即 $\mathbb{E}[\mu^2] = \mu^2$

代入得：

$\text{Var}(X) = \mathbb{E}[X^2] - 2\mu^2 + \mu^2 = \mathbb{E}[X^2] - \mu^2$

而 $\mu = \mathbb{E}[X]$ ，于是得到最终公式：

$\text{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2$

Triton_learning

2026-08-06T14:13:33.000Z

GPU简介

1. CPU vs GPU 的比喻

CPU：像一个数学教授，非常聪明，能做复杂的计算，但只有十几双手（核心数少，比如8核、16核）。
GPU：像1000个小学生，每个只会做简单的加法，但人数超多（几千个核心），可以同时处理大量简单任务。

向量加法就是典型的"简单任务重复无数遍"，特别适合GPU。

2. GPU 的结构

GPU 上有成千上万个 线程（thread），每个线程做一点点工作。这些线程被组织成：

线程（Thread）：最小的执行单元，就像一个小学生。
线程块（Block）：一组线程的集合，就像一个小队。
网格（Grid）：所有线程块组成的整体，就像整个学校。

FlashAttention

2026-08-04T07:17:24.000Z

FlashAttention

一、Introduction

为了加快LLM的训练和推理速度，针对transformer注意力机制的特点和GPU等硬件结构，通过分块和重计算来减少HBM读写次数,进而加快注意力计算的一种优化算法。

二、 Background

2.1 硬件特性

GPU的内存结构如上图左侧所示：

SRAM是GPU的片上内存，GPU计算时必须先把数据搬运到SRAM上才能进行计算。内存最小，但是IO速度最快

HBM是高带宽内存，存放模型参数权重、KVcache以及其它计算中间变量。内存较大，IO速度较快。(是否正确)

DRAM是主存，cpu进程存放数据的地方。内存大，IO速度慢。

模型运行时，GPUs拥有海量线程来执行一个操作/函数（kernel），每个kernel将输入从HBM加载到寄存器和SRAM上，然后计算，最后将输出写回HBM。

2.2 操作特性

依据计算和内存访问的平衡，操作可以分为计算密集型操作和内存密集型操作

计算密集型：操作所需时间主要取决于算术运算次数，访问内存只占很少的时间。典型示例为维度很大的矩阵乘法、通道数很多的卷积。

内存密集型：操作所需时间主要取决于内存访问次数，计算花费的时间很少。典型示例为逐元素操作（激活函数，dropout）、reduce（求和、Softmax、Batch Normalization、layer Normalization）。

2.3 标准Attention实现（简化的核心部分）

输入 $\mathbf{Q},\mathbf{K},\mathbf{V} \in \mathbb{R}^{N \times d}$

先计算注意力分数矩阵S（每个token的Q与所有token的K进行计算）

$\mathbf{S} = \mathbf{Q}\mathbf{K}^T \in \mathbb{R}^{N \times N}$

再计算注意力权重矩阵P(对S按行进行softmax)

$\quad \mathbf{P} = \mathrm{softmax}(\mathbf{S}) \in \mathbb{R}^{N \times N}$

最后计算输出矩阵O

$\quad \mathbf{O} = \mathbf{P}\mathbf{V} \in \mathbb{R}^{N \times d}$

标准的Attention实现需要将S和P存储到HBM中，占用 $\mathbf{O(N^2)}$ 的内存

部分操作例如softmax操作是内存密集型操作，大量内存访问导致较慢的运行时间。

还有其它的逐元素操作如掩码（按行）和P的丢弃也加剧了这一现象。

三、 FlashAttention算法的设计与分析

3.1 设计目标

给定 HBM 中的输入 $\mathbf{Q},\mathbf{K},\mathbf{V} \in \mathbb{R}^{N \times d}$ ，计算注意力输出 $\mathbf{O} \in \mathbb{R}^{N \times d}$ 并将其写入 HBM。目标是减少 HBM 访问量，使访问量少于 $\mathbf{\Theta(N^2)}$ 。

3.2 分块

使用了分块（tiling）和重计算（recomputation）两种方法来实现。

分块计算的详细流程：

如图1右侧所示，先从行维度将Q分成 $T_r$ 个块，将K和V都分成 $T_c$ 个块。

关于块大小的选取，原则是使SRAM能同时放下 $\mathbf{Q_i}$ 、 $\mathbf{K_j}$ 、 $\mathbf{V_j}$ 、 $\mathbf{O_i}$ 、 $\ell_i$ 、 $m_i$ 的前提下，让分块的数量尽可能少。按上述分块大小，所需的总空间为 $(2B_c \times d + 2B_r \times d +2B_r)$ 。上述理论推导时忽略了 $\ell_i$ 、 $m_i$ 等低阶空间复杂度变量，实际工程时会采用保守缩小块大小的策略。

遍历 $\mathbf{Q}$ 的分块，从HBM取一个块 $\mathbf{Q_i}$
遍历 $\mathbf{K}$ 、 $\mathbf{V}$ 的分块，从HBM取[ $\mathbf{K_j}$ , $\mathbf{V_j}$ ]，注意二者分块的行序列要一致。
在SRAM上计算注意力分数矩阵 $\mathbf{S}_{ij} = \mathbf{Q}_i \mathbf{K}_j^T \in \mathbb{R}^{B_r \times B_c}$
在SRAM上计算 $\mathbf{Q_i}$ 中选中的行的每行最大值，注意力矩阵逐元素取e指数 $\tilde{\mathbf{P}}_{ij} = \exp(\mathbf{S}_{ij} - \tilde{m}_{ij}) \in \mathbb{R}^{B_r \times B_c}$ ，进行softmax操作的分母 $\tilde{\ell}_{ij} = \mathrm{rowsum}(\tilde{\mathbf{P}}_{ij}) \in \mathbb{R}^{B_r}$ 。
更新 $\mathbf{Q_i}$ 中选中行的每行行最大值 $m_i$ ，每行各元素的e指数之和 $\ell_i$ 。
增量式计算更新 $\mathbf{Q_i}$ 中选中的行的注意力输出 $\mathbf{O}_i \leftarrow \mathrm{diag}(\ell_i^{\mathrm{new}})^{-1}\big(\mathrm{diag}(\ell_i)e^{m_i - m_i^{\mathrm{new}}}\mathbf{O}_i + e^{\tilde{m}_{ij} - m_i^{\mathrm{new}}}\tilde{\mathbf{P}}_{ij}\mathbf{V}_j\big)$ ，写入HBM。
将更新后的 $m_i$ 和 $\ell_i$ 写回HBM。

与标准Attention访问次数对比分析：

二者 $\mathbf{Q},\mathbf{K},\mathbf{V}$ 读取次数相同，区别在于FlashAttention不用存 $\mathbf{S}$ 和 $\mathbf{P}$ ，为什么传统attention需要存呢？——因为一般N很大，并且N>>d，SRAM存 $\mathbf{Nd}$ 空间没问题，存不下S和P这种 $\mathbf{N^2}$ 空间的。但是每次多了 $\ell_i$ 、 $m_i$ 的读写，总共需要 $4 \times B_r \times T_r \times T_c = 4NT_c$ 次读写，同时多了 $2 \times T_r \times (T_c-1) \times B_r \times d = 2Nd(T_c-1)$ 次读写，复杂度为 $\mathbf{\Theta(N^2d^2/M)}$ ，而标准的Attention的HBM读写次数复杂度为 $\mathbf{\Theta(N^2)}$ 。

对于典型的d（64-128）和M（约100KB）值， $d^2$ 远小于M，因此FlashAttention所需的HBM访问次数比标准Attention实现的少很多倍。

3.3 重计算

通过分块，无需存储S和P,但是训练时反向传播需要使用 $\mathbf{S}$ 和 $\mathbf{P}$ 来计算梯度。解决办法是重新计算。

通过存储输出 $\mathbf{O}$ 和 softmax 归一化统计量 $(m,\ell)$ ，我们可以在SRAM中轻松地重新计算注意力矩阵 $\mathbf{S}$ 和 $\mathbf{P}$ 。

相比之下，即使FlashAttention有更多的 FLOPs，重新计算通过减少 HBM 访问来加速反向传播。

3.4 扩展：Block-Sparse FlashAttention

用一个掩码矩阵 $M \in \{0, 1\}^{N / B_r \times N / B_c}$ 来表示块与块之间是否需要计算，

欢迎来到我的博客

2026-08-03T16:00:00.000Z

你好呀，欢迎来到我的博客 👋

这里会用来记录我的学习笔记和日常生活：

📚 学习：技术文章、读书笔记、学习心得
🌱 生活：日常随笔、旅行记录、成长感悟

✍️ 如何写一篇新文章

方式一：命令行创建（适合纯文本文章）

在项目根目录运行：

1	`npx hexo new post 文章标题`

然后编辑 source/_posts/ 下生成的 Markdown 文件，文章会自动出现在首页。

方式二：Typora + PicGo 写文章（推荐，含图片）

在 Typora 中正常写作，插入图片时会自动调用 PicGo 上传到云端图床，md 里会写入 https://picgocloud.com/... 这样的外链；
写完后把文件保存到 source/_posts/ 目录下（文件名建议用英文，页面标题由 front-matter 的 title 决定）；
文章最前面需要加上 front-matter，至少包含 title 和 date，示例：

---
title: 文章标题
date: 2026-08-14
categories: [学习]
tags: [标签1, 标签2]
---

⚠️ 提示：外链图片依赖图床服务，图床失效图片就会丢失。如果担心，可以执行下面的「图片本地化」步骤，把图片下载到仓库里。

📥 图片本地化（可选，推荐）

把文章里的外链图片下载到本站仓库，由自己的站点直接提供，不再依赖第三方图床：

1 2	`node scripts/localize-images.mjs # 处理全部文章 node scripts/localize-images.mjs 文件名.md # 只处理指定文章`

脚本会自动：

下载外链图片到该文章同名的资源目录（如 source/_posts/文章名/）；
把 md 里的 https://... 链接改写为相对文件名（如 ![](xxx.png)），Typora 和网站都能正常显示；
已处理过的图片自动跳过，可放心重复运行。

🖼️ 文章封面图（可选）

首页卡片默认使用统一的渐变封面。想给文章配专属封面：

在 source/img/covers/ 下放一张 xxx.svg（可参照现有的模板修改文字和配色）；
在文章 front-matter 中加一行：

1	`index_img: /img/covers/xxx.svg`

🖥️ 本地预览

1	`npx hexo server`

浏览器打开 http://localhost:4000 就能看到效果。

🚀 发布上线

写完文章后，提交代码并推送到 GitHub，Vercel 会自动重新部署：

1
2
3

git add .
git commit -m "新文章"
git push

完整流程回顾：Typora 写作 → （可选）图片本地化 → git 提交推送 → Vercel 自动部署 ✅

Jaspery

infra规划

AI Infra 学习与研究路线

1. 总体目标

目标定位

2. 总体技术栈

2.1 GPU 底层

3. Phase 0：CUDA 与 GPU 基础

目标

学习内容

GPU Architecture

CUDA

阶段产出

4. Phase 1：集合通信

目标

Ring AllReduce

5. Phase 2：NCCL

需要掌握

GPU 拓扑

7. Phase 3：Megatron-LM

目标

7.1 Data Parallel

7.2 Tensor Parallel

7.3 Pipeline Parallel

7.4 3D Parallelism

8. Phase 4：DeepSpeed

9. 第一阶段科研方向：训练通信优化

9.1 拓扑感知通信优化

9.2 计算与通信重叠

10. Phase 5：LLM Inference

11. Phase 6：vLLM

12. Phase 7：KV Cache

13. Phase 8：Continuous Batching

14. Phase 9：Speculative Decoding

15. Phase 10：Disaggregated Inference

16. 为什么 PD Disaggregation 特别适合通信背景

网络带宽

网络拥塞

GPU 调度

KV Cache Placement

17. 第二阶段科研方向：Inference + Network

18. 推荐的 4 个科研项目

项目 1：NCCL Topology Benchmark

目标

最终产出

19. 项目 2：Megatron 通信与计算 Overlap

Baseline

Optimization

20. 项目 3：KV Cache 优化

21. 项目 4：PD Disaggregation + RDMA

22. 12 个月学习规划

23. 每个阶段的学习方法

24. 实验资源要求

25. 未来需要掌握的技术栈

GPU

Distributed Training

LLM Inference

Network

System

26. 就业方向

AI Infra

GPU / Systems

Network + AI

27. 技术能力优先级

28. 最终能力画像

29. 最终主线

LLM中的数学原理及infra优化

引言

1.RMSNorm - 均方根归一化

1.1 LayerNorm

数学公式

Infra视角

1.2 RMSNorm (Root Mean Square Normalization)

数学公式

Infra 视角

LayerNorm vs RMSNorm（底层实现视角）

Var(X)=E[X2]−(E[X])2\text{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2Var(X)=E[X2]−(E[X])2 的推导逻辑

Triton_learning

GPU简介

1. CPU vs GPU 的比喻

$\text{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2$ 的推导逻辑