LLM System: Transformer Engine 01 - 在 AI Infra 技术栈中的位置

本篇目标:了解 Transformer Engine 的技术定位,搞清楚它为什么存在,以及它和 PyTorch、cuBLAS、Megatron-LM 的边界。 基本接口 Layer 类定义接口非常直接,最表层的使用方式就是把 torch.nn 模块替换成 transformer_engine.pytorch 模块。 普通 PyTorch 写法: self.linear = torch.nn.Linear(hidden_size, 4 * hidden_size) TE 写法: import transformer_engine.pytorch as te self.linear = te.Linear(hidden_size, 4 * hidden_size) 量化上下文: from transformer_engine.pytorch import fp8_autocast with fp8_autocast(enabled=True): y = module(x) 进入 TE 的 FP8 上下文之后,TE 会围绕量化、反量化、fused path、tensor cache 和 backend 选择做一系列处理。相较于纯 PyTorch 计算图优化,TE 会拿到更多信息 tensor parallel、sequence parallel、FP8 recipe 等。这些额外信息给底层算子优化留下了空间。 这一点目前还是比较 general 层面 的认知,后面要继续顺着源码和 profiler trace 去验证。 TE 和 Megatron 的边界 Megatron-Core 负责模型并行、训练 loop、optimizer、activation checkpoint、MoE routing、pipeline schedule 和 config。 ...

May 21, 2026 · 1 min

LLM System: KV Cache 查询 01 - PagedAttention 原理

TODO: 这里写 PagedAttention 的核心抽象:block/page、block table、逻辑 token 到物理 KV block 的映射。 基础:tensor 级拆请求的形状(大量细节) 定义符号:$B$ 是 batch size,$T$ 是 seq_len,$D$ 是 token_dim,$d_q$ 是把 embedding token 投影到 $Q$ 后的维度。 推理框架拿到的请求是:\(R \in \mathbb{R}^{B \times T}\)。 $R_{b,t}$ 是一个最最基本的 token id 标量。 raw 请求经过 embedding lookup,做的操作是把这个 token 标量映射成一个高维向量。假设原先 token 是 1234 这个标量,现在就把 token 映射成 [0.1, 0.2, 0.3, 0.4] 这样的向量。 所以 $R$ 经过 embedding lookup 之后,得到:\(X \in \mathbb{R}^{B \times T \times D}\)。 因为我们目前只考虑推理场景,所以把 $W_Q$、$W_K$、$W_V$ 之类的矩阵当成固定的模型参数。 然后很多博客会直接写:\(Q = XW_Q\)。 ...

May 10, 2026 · 5 min

LLM Theory: MuP 02 - 推广到更 general 的 init setup

本文从训练模型要考虑的第一性原理(稳定性和速度)出发,探讨了 LLM 预训练中的初始化设定问题。 参考文献 MuP之上:1. 好模型的三个特征 https://spaces.ac.cn/archives/11605 https://spaces.ac.cn/archives/11647 https://spaces.ac.cn/archives/11729

May 6, 2026 · 1 min

gemm和alltoall通算融合

1. 总体思想 这次做的是单机八卡 H200、NVLink、Ulysses CP 下的 GEMM 和 AllToAll 融合。先把 forward 写清楚: A2A → QKV projection → QK → PV → A2A 需要接起来的主边界有两个。输入侧是 A2A→QKV projection,通信先把各个 peer 的输入 tile 搬到本地最终布局,GEMM 拿到一块就算一块。输出侧是 batched PV→A2A,每个本地 head 都有一组独立的 P×V,GEMM 算完一个 tile,通信 CTA 立刻把它送到目标 rank 的最终 Ulysses 布局。 如果 GEMM 和 NCCL 顺序执行,端到端时间接近两段时间相加。这里把通信 CTA 和 GEMM CTA 放进同一个 cooperative persistent grid,两种 CTA 常驻在不同的 SM 上,用 tile 级 ready epoch 接力。A2A→GEMM 由通信生产、GEMM 消费;GEMM→A2A 交换生产消费关系。这样首批 tile 到达后就能启动计算,前面的 tile 也可以在后续 GEMM 还在跑时发出去。 ...

August 17, 2026 · 3 min

LLM Theory: MuP 03 - 为什么MoE模型训练起来有难度?MuP视角

众所周知MoE有一个Router,也就是门控单元,基于topk选择路由到的专家。训练稳定的基本思想是希望随着宽度d增大,模型的activation、grad、delta loss、delta weight变化都比较小。 \[ \operatorname{RMS}(h_l)=\Theta(1),\quad \operatorname{RMS}(\Delta h_l)=\Theta(1),\quad \Delta L=\Theta(1). \] 不好分析的根源是因为topk不光滑,所以找不到一个有界的全局利普希茨常数。 也就是假设两个输入的x很接近,但是在router这里发生了跳变,导致L也就是$\lVert f(x_2)-f(x_1)\rVert/(x_2-x_1)$趋近于无穷大。 所以说topk没有全局的LipschitzBound。那么为什么lipschitz常数对mup会有影响,原因是对于某个token的一次backward里面参数更新,如果weight更新后导致路由切换,delta y的计算就不只是简单的$\Delta y\approx J_\theta y\Delta\theta$,而是两个expert相减。$\Delta y=F_j(h)-F_i(h)$. 这样可能会导致,$\lVert\Delta x\rVert\to0,\lVert\Delta y\rVert=\Theta(1)$.这个导致的最直接问题就是我们没有办法再去用一阶泰勒近似delta y了,也就是用dense的方式,没有办法用局部梯度去正确描述某个token的路由表随参数更新导致变化的这个事件,也就不能囊括所有的梯度更新类型。 从另一个角度,如果强行进行泰勒展开(当然肯定有各种不严谨,我是工程民科,凑合看看算了),就有$y(\theta+\Delta\theta)-y(\theta)=J_\theta y\Delta\theta+o(\lVert\Delta\theta\rVert)$. 但最大的问题是$\lVert y(\theta+\Delta\theta)-y(\theta)\rVert=\Theta(1)$,$\frac{\lVert y(\theta+\Delta\theta)-y(\theta)\rVert}{\lVert\Delta\theta\rVert}$所以分子分母不同阶。taylor余项也不是无穷小的。 要是想做MoE的mup,就得想办法去控制它由于router切换导致的跳变大小,也就是说必须要建模触发router跳变的条件,以及给跳变以后专家输出的diff加上一个bound。路由跳变对均方特征的贡献可以如下描述: \[ \mathbb{E}\left[\lVert\Delta y_{\mathrm{switch}}\rVert^2\right] \lesssim \int_0^{2\varepsilon}q(u)f_\gamma(u)\,\mathrm{d}u. \] f描述有多少token距离路由边界为u,q表征输出跳变的距离,2sigma代表参数更新能跨越多宽的边界。那么这件事容易吗? 我们之前只是假设了输入不变,在mup前面的系列里面,mup的目标之一是要让中间的特征发生一些非退化的显著变化,$\Delta h=\Theta(1)$。所以即使router参数冻结了,h的变化也会导致路由切换。这个问题会让特征更新的尺度推导变得困难。 总而言之,传统 Dense μP 能闭环,是因为给定每层的: \[ \operatorname{RMS}(h),\quad \operatorname{RMS}(\delta),\quad \operatorname{RMS}(\Delta W), \] 就能递推出下一步的特征更新: \[ \Delta h\approx J_W h\Delta W, \] 而这个更新尺度仍只由矩阵维度、初始化尺度和学习率决定。 MoE 不闭环,是因为即使这些 RMS 都已知,仍然无法确定输出更新: \[ \Delta y = \Delta y_{\mathrm{smooth}} + \Delta y_{\mathrm{switch}}. \] 第二项还取决于两个额外变量: ...

July 26, 2026 · 1 min

LLM Theory: SMD 01 - 权重更新的拆解

本文是 LLM Theory 下 SMD 专题的第一篇,是关于 Spherical Motion Dynamics(SMD)的学习笔记,主要参考 Wan et al. 的 Spherical Motion Dynamics: Learning Dynamics of Neural Network with Normalization, Weight Decay, and SGD 以及 NeurIPS 2021 版本 Spherical Motion Dynamics: Learning Dynamics of Normalized Neural Network using SGD and Weight Decay。 朴素的描述模型更新 对于本科一二年级时候的笔者而言,如果想要描述模型的更新量,那么只会考虑这个非常直接的东西:$\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$。 这确实是很直白的,相当于计算了模型权重的欧氏距离。是什么让我们必须放弃这种非常直观的欧氏距离呢?源于常用的归一化。 加上归一化之后,对于模型的输出,$y = \operatorname{BN}(x, k\boldsymbol{W}) = \operatorname{BN}(x, \boldsymbol{W})$,看不出模型尺度对 $y$ 的影响。但是却会从 $\lVert k\boldsymbol{W}_{t+1} - k\boldsymbol{W}_{t} \rVert$ 这一度量手段上产生 $k$ 倍的差距。 这就发现一个很显然的问题了:我们是希望通过观测类似于 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 这样的东西来控制模型的训练,而不是直接看最后的 $y$。因为只通过看 $y$ 的变化,并基于这种变化的规律,去指导如何进行训练前的设置,这件事并不容易。我们是希望找到一个更可以写出明确表达式,更容易观测,意义更明确的指标去观测,而这个指标恰好还要有一些规律和 $y$ 的规律“趋同”,这样我们就得到了一个 $y$ 的近似物,而且这个近似物比 $y$ 更好分析。所以我们会很直接的想到 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$。但是现在,当我们给模型加上BN(为什么加BN不赘述了)以后发现的问题是 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 会随着模型尺度产生成倍的差距,而 $y$ 却不会被尺度影响。这说明 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 不是一个好的指标,因为它错误地预测了 $y$ 随着尺度变化的行为规律。 当然,原则上,我们其实可以发现,上面的地方有一个逻辑漏洞:为什么权重变化一定能预测 $y$ 的变化,这是错误的呀,权重变了 $y$ 不一定变。所以 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 这个东西失效是情理之中的,不能先入为主的认为 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 的规律能和 $y$ 趋同,那么出现上面的问题十分合理。我们一个理想中的目标是通过权重变化去预测 $y$ 的变化,这样我们就可以通过设计一套权重变化的方案去精确得到想要的 $y$。理想是很远大的,现实会比较难做,但是可以在最朴素的方案上向前不断推进预测的细化程度。SMD就是这样一种东西。 ...

May 17, 2026 · 10 min

LLM Theory: SMD 02 - Angular Update 才更接近真实改变量

本文是 LLM Theory 下 SMD 专题的第二篇。上一篇主要把 SGD、SGD + weight decay、Adam、AdamW 的径向/切向更新拆开了;这一篇继续沿着这个拆法,讨论几个笔者读 SMD 时觉得更关键的 insight。 Insight 1:为什么只看 Effective Learning Rate 不够 在讨论 normalization 和 weight decay 的联合作用时,有一句非常关键的话: 对于 scale-invariant weight,任务梯度 $\partial L / \partial \boldsymbol{w}$ 总是与权重 $\boldsymbol{w}$ 垂直。因此,梯度分量总是倾向于增大权重范数,而 weight decay 提供的分量总是倾向于减小权重范数。 这句话其实就是 SMD 的基本物理图像: \[ \text{任务梯度} \Rightarrow \text{切向更新,改变方向,同时二阶增大范数} \] \[ \text{Weight Decay} \Rightarrow \text{径向更新,缩小范数} \] 也就是说,在带 normalization 的网络中,SGD 梯度并不是在普通欧氏空间里随便走,而是在球面的切线方向上推动权重转动;weight decay 则像一个向心力,把权重往原点拉。两者共同作用,最终可能让 weight norm 进入一个稳定状态。 但是这里有一个很容易被忽略的问题: Weight decay 可以对冲掉梯度更新导致的 weight norm 增长,但这并不意味着 gradient norm 本身也被稳定住了。 ...

May 17, 2026 · 9 min

为什么 vLLM 和 SGLang 在生成式搜推场景如此废柴?

workload的区别 推理的类型是非常多的,RL rollout,在线服务,推荐……这些workload都不一样,虽然说推理框架这一层原则上是一个runtime的workload调度层,但是随着不断的发展,每个推理框架也并不能做到对任意的workload都有一套好的runtime策略,最终这几个框架的优化还是固化到了某几个场景,比如高贵的RL后训练,比如最主流的Serving,然后一些类似于广告推荐的边缘场景就被一脚踢死了(笑。 这他吗就很尴尬呀,现状就是典型场景已经穷尽了可能性,没东西做了,非典型的场景又没好的实现,导致必须从头开始吃屎,不过这里不吐槽这种极度的不均衡,只谈谈为什么知名的v和s两家推理框架换个场景就不行了,本质上是在讨论workload。 v和s两家框架面向的是在线序列的批处理,也就是一下QPS来了很多,然后以一个很高效的调度方式让这些任务都能被调度起来。框架层只是下发,实际的工作都在算子队列里面。而GR显然不一样,作为推荐输入的embedding是共享的,推理出来的东西是一个多级索引,这个多级索引被硬编码为token,输出token只能从这里挑选,一般叫Sematic ID(所以也得小小扩充一下词表)。然后一般的GR还会做beam search,类似于给非top的输出一次参与后续自回归生成机会,以增加结果的diversity。这个操作对于传统llm来说是很少见的。 做beam的时候虽然kvcache解决了重计算的问题,prefix sharing或copyonwrite解决了重复存储的问题,但是beam阶段load同一kv的操作依然是问题。算子上这是很多次重复没必要的G2S。如果要改这个算子的话,会更麻烦。 beam导致的kv重排 beam因为在自回归这一层做了分叉,原先设计的那种纯顺序化的kvcache就很容易出问题了,beam会主动更新自己候选的beam,也就是说他的batch来自于beam主动的而不是外界来的很多共同前缀的请求。beam能主动产生batch也能修改他们的顺序关系(熵之源说是),一旦修改了,就和顺序化的自回归假设不一样了,就得重新换block/radix treenode顺序。beam的候选+淘汰对传统自回归的范式设计影响太大了,必须得侵入式修改算子+框架。 其他调度 gr还会有很多奇奇怪怪的前处理后处理,这些都没被v和s优化过,所以自然性能也不行。

July 16, 2026 · 1 min