Guest Lecture:Dan Fu - LLM 推理系统从零构建
Stanford CS336 Language Modeling from Scratch | Spring 2026 | Guest Lecture: Dan Fu - LLM Inference from Scratch,时长 01:11:40。
推理系统全景:一个 Token 的生命周期#
大语言模型训练完成后,如何将它变成一个能持续响应请求的在线服务?这个问题的答案就是推理引擎(Inference Engine)。推理引擎是将电力转化为智能的核心组件,正如内燃机将石油转化为有用动能一样。GPU 本身只是沙子加工而成的芯片,是推理引擎赋予了它将数学运算映射为可用智能的能力。

一个完整的推理系统包含以下关键环节:
- Request In(请求接入):用户发送一段文本,系统首先对其进行 tokenization。
- Scheduling(调度):调度器决定将请求分配到哪些 GPU 上执行。这里涉及 continuous batching(连续批处理)和 chunked prefill(分块预填充)等技术,用于在多个请求之间高效共享 GPU 资源。
- KV Cache 查询:在真正执行计算前,系统会检查是否已经见过这个请求的部分或全部 token。如果能命中缓存,就可以跳过大量重复计算。
- Execution(执行):核心的机器学习计算,分为 prefill(预填充)和 decode(解码)两个截然不同的阶段。Prefill 处理大量输入 token,decode 逐个生成输出 token。这两个阶段的计算特性差异巨大,后文会详细展开。
- Parallelism(并行化):当模型过大无法装入单个 GPU 时,需要 tensor parallelism(TP)、expert parallelism(EP)、attention DP 等策略将计算分散到多个 GPU 上。
- Disaggregation(分离):将 prefill 和 decode 放到不同的机器集群上执行,通过 async RDMA 传输 KV Cache,让两种计算各自在最适合的硬件上运行。
- Tokens Out(输出):生成的 token 经过 detokenization、stop token 检测、安全检查等后处理,最终返回给用户。
整个推理引擎在一个循环中持续运行:等待请求到来,执行调度、计算、采样,然后重复。这个看似简单的循环背后隐藏着巨大的系统复杂度。
如果理解了推理引擎和 GPU kernel 的工作原理,就能在机器学习算法层面实现全栈创新。这是贯穿整篇笔记的核心观点。
工作负载特征与 SLA 目标#
不同应用场景下,推理系统面对的流量模式差异极大。理解这些差异是做出正确系统设计决策的前提。
工作负载的三个维度#

输入/输出 token 长度分布 + Cache 命中率:以编程助手(如 Cursor)为例,用户通常会把整个代码仓库作为上下文传入,输入端可达数万 token;输出端则取决于模型是否产生 thinking token,可能很短也可能很长。而一个普通的聊天对话,输入输出都相对简短。具体的分位数分布为:p1 约 32/64 token(cache hit 5%),p50 约 512/256 token(cache hit 60%),p99 可达 4096/1024 token(cache hit 98%)。
每个 session 的对话轮数:当前 LLM 的使用模式高度 turn-based。编程 agent 会与模型反复交互,调用工具(grep、搜索)后再将结果喂回模型。中位数大约 7 轮,但 p99 可以到 50 轮。这决定了 KV Cache 的生命周期管理策略。
轮间间隔时间:交互式 voice 模式的响应间隔可能只有 0.5 秒,而有人与 ChatGPT 讨论健身计划的对话可能两周才交互一次。中位数约 10 秒,p99 到 300 秒。这个分布直接影响 KV Cache 的保留策略。
三大 SLA 指标#
- TTFT(Time to First Token)< 500 ms:首个 token 的响应时间,反映 prefill 速度。在交互式场景下,用户期望不到一秒就看到模型「开始思考」。
- TBT(Time Between Tokens)< 30 ms:相邻 token 之间的生成间隔,反映 decode 速度。30 ms 对应约 33 token/s 的生成速率。
- Target QPS per GPU:单 GPU 的吞吐量,是衡量推理成本效率的核心指标。
关键认识是:同一个延迟目标在不同工作负载下的难度可能天差地别。一个短对话的 500 ms TTFT 轻而易举,但对一个 4096 token 输入、98% cache hit 的编程请求,同样的目标可能极具挑战。SLA 必须结合工作负载特征来理解。
Continuous Batching 与请求调度#
当推理系统同时处理多个请求时,如何高效利用 GPU 资源就成了核心问题。Continuous Batching 是当前主流推理引擎(如 vLLM、SGLang)采用的调度策略。

这张图的时间轴从上往下流动,展示了调度器如何动态管理多个请求:
- Step 1:一个长请求 Req A 和一个短请求 Req B 同时在 GPU 上执行。
- Step 2:Req B 完成后释放资源,新请求 Req C 立刻加入执行队列,无需等待 Req A 结束。这就是 continuous batching 的核心思想:请求可以随时加入和退出,不需要等一个 batch 中所有请求都完成。
- Step 3:系统进行容量检查(是否有足够的 KV block)和 microbatch 检查(是否超过最大 token 数)。
- Step 4:Req A 完成,Req D 加入。但如果此时 GPU 内存不足以存放 Req D 的 KV Cache,它需要在队列中等待。
- Step 5-6:剩余请求依次完成。
与传统的 static batching 相比,continuous batching 的优势在于:GPU 不会因为某个请求提前结束而空转,新请求可以随时填补空出的计算槽位。底部注释还提到了 chunked prefill:当一个很长的 prompt 到来时,不必一次性完成所有 prefill 计算,而是将其拆分成多个 chunk,与正在进行的 decode 操作交错执行,避免长 prefill 阻塞其他请求的 token 生成。
这里涉及两类资源竞争:计算资源(每个 step 中需要在多个请求之间分配 GPU 算力)和内存资源(每个活跃请求都需要在 GPU 上维护自己的 KV Cache)。当 KV Cache 填满 GPU 内存时,新请求就不得不排队。
KV Cache:前缀共享、分层存储与预取#
KV Cache 是推理系统中最重要的优化之一。其核心思想是:如果多个请求共享相同的 token 前缀,或者同一个 session 的后续轮次复用了之前的上下文,就没有必要重复计算这些 token 的 attention 激活值。
Radix Tree 前缀共享#

系统使用一棵 Radix Tree 来管理 KV Cache。树的根节点存储 system prompt(如 “You are a helpful assistant…”),这部分所有用户共享,只需计算一次。每个 session 在树中形成自己的分支:Session 1、Session 2、Session 3 各自维护不同的对话历史。
当一个新请求到来时,系统沿着树逐级匹配 token hash:已经匹配到的节点直接复用 KV Cache(cache hit),只有未匹配的新 token 才需要进入 prefill 计算(cache miss)。图中每个方块代表 32 个 token 的 KV 数据,形状为 [layers, 32, hidden_dim],用 block hash = f(content, parent) 做唯一标识。
这个设计让大量请求可以共享 system prompt 的计算结果,并且多轮对话只需要对新增的 token 做 prefill。
三级存储层次#

随着服务规模增长,GPU HBM 很快就不够用了。KV Cache 的存储形成了一个三级层次结构:
- GPU HBM(热数据):存放当前活跃 session 的 KV Cache,访问速度最快。
- CPU pinned memory(温数据):当 GPU 内存不足时,通过 LRU(最近最少使用)策略将不活跃的 KV block 驱逐到 CPU 内存。
- NVMe SSD(冷数据):更久没有访问的数据继续下沉到磁盘。
这个层级与经典操作系统中的虚拟内存管理(70-80 年代的 OS 教科书级内容)几乎完全同构:打开太多应用耗尽物理内存时,OS 会将不活跃进程的页面换出到磁盘。推理系统面对的是同样的问题,只是「应用」变成了「session」,「物理内存」变成了「GPU HBM」。
最理想的情况是能预测未来:当用户在 ChatGPT 中打开一个月前的对话时,这是一个强信号——他马上要在这个 session 中提问。系统可以提前将对应的 KV Cache 从 SSD 加载到 GPU。退而求其次的启发式方法是 LRU。
右侧还展示了 cache sharing:多个 Prefill Worker 之间通过 RDMA 共享 KV block,避免对热门前缀(如 system prompt)做重复 prefill。
Jensen Huang 最近开始痴迷于 CPU 性能,原因之一就是上一代 CPU 速度太慢,成了 KV Cache offload 的瓶颈。一台 50 万美金的 GPU 机器被 1000 美金的 CPU 拖慢,是一个非常糟糕的局面。OpenAI 大量采购 SSD 和 DRAM,部分原因也在于 KV Cache 存储的需求。
Prefill/Decode 分离与异构硬件#
Prefill 和 decode 是推理计算中两个性质截然不同的阶段,理解它们的差异是推理系统设计的基础。
Prefill vs Decode 的本质区别#
Prefill(预填充):处理用户输入的全部 prompt token。假设输入 10,000 个 token,一次性计算所有 attention 和 FFN,输出一个 token。这是一个 compute-bound(计算密集型)操作,与训练时的前向传播非常类似,区别仅在于不需要反向传播。GPU 的 FLOPS 利用率高。
Decode(解码):逐个生成输出 token。每生成一个 token,都需要将其送回模型做一次前向传播。虽然每次前向传播的 FLOPS 很少(只处理一个 token),但必须从 HBM 中完整加载一次模型权重。这是一个 memory-bandwidth-bound(带宽密集型)操作。
两者在时间消耗上的模式也不同:prefill 对一个 prompt 只执行一次,但耗时较长;decode 对每个输出 token 各执行一次,单次很快但次数多。如果使用了 speculative decoding,每次可能验证 3-4 个 token。
分离部署架构#

由于 prefill 和 decode 的计算瓶颈完全不同,将它们部署在不同的机器集群上成为标准做法:
- Prefill Worker:使用 TP4(4 路 tensor parallelism),优化目标是最大化 FLOPS,异步执行 prefill loop。计算完成后将 KV Cache 放入 holding buffer。
- Router Node:路由节点将请求和 KV Cache 来源信息转发到 decode 侧。
- Decode Worker:拥有更多 GPU(图中展示了 16 个),使用 wide EP(expert parallelism,将 MoE 专家分散到多 GPU)和 attention DP(将 attention 沿 batch 维度分片)。Decode worker 通过 RDMA 异步获取 KV Cache,收到后开始 decode loop。
这种分离让每侧都能独立优化:prefill 侧追求计算吞吐,decode 侧追求带宽效率。
异构硬件的兴起#
Prefill 和 decode 的瓶颈差异大到可以使用不同类型的芯片分别处理:
- NVIDIA 收购 Groq:NVIDIA 计划下一代产品中用自家 GPU 做 prefill,用 Groq 的 LPU 芯片做 decode。LPU 擅长的正是 decode 这种带宽密集型操作。
- OpenAI + Cerebras:Cerebras 的 wafer-scale 芯片同样更适合 decode 工作负载。
- SambaNova 等公司也在这个领域押注。
这意味着未来的推理集群可能不再是清一色的 GPU,而是 prefill GPU + decode ASIC 的异构组合。
大规模部署中的 Bug 故事#
当推理系统扩展到每天处理数万亿 token 时,概率极低的 bug(0.001% 甚至更低的触发率)就会变成必然事件。以下是 2025 年底几个开源推理引擎中发生的真实案例。
NaN 循环#
某个 kernel 存在极其细微的数值错误,触发条件非常苛刻。一旦触发,部分 logit 在计算过程中变成 NaN。NaN 经过 softmax 后,模型的采样行为被彻底破坏,开始反复输出同一个 token,表现为无限循环的 “hi hi hi hi…” 或连串感叹号。
Tool Call 死循环#
有人修改了处理 tool call 的代码逻辑。正常情况下,模型发出 tool call(如 “请做一次网络搜索”)后会生成一个 stop token,将控制权交还给外部 harness 去执行实际搜索。但修改后,stop token 没有被正确处理。模型会说 “请做网络搜索”,然后困惑于为什么没有收到搜索结果,接着再次请求搜索,如此循环数万个 token。外部观测到的症状是 completion length 急剧飙升。
Off-by-One 导致的语言切换#
这个 bug 同时影响了多家推理服务商。表面症状是:用英文提问,模型突然开始用中文回答。最初被归咎于量化精度问题,实际根因是 kernel 中一个 off-by-one 错误。某些情况下,kernel 会读取一小段未初始化的 GPU 内存,将其作为 attention 输入处理。经过 attention 计算后,这段垃圾数据碰巧映射到中文字符的 token ID。模型看到自己「思考」出了一个中文字符,就推断用户可能在用中文提问,于是切换到中文继续生成。
这类 bug 在小规模测试中几乎不可能复现。只有在生产级流量下,那些 0.001% 概率的边界条件才会被触发。这也解释了为什么推理引擎的可观测性(observability)如此重要。
NVL72、容错与下一代推理基础设施#
随着模型规模增长到万亿参数,单个 GPU 甚至单台机器已经无法容纳完整模型。NVIDIA 最新一代 Blackwell 架构推出了 NVL72 Grace Blackwell 系统,将 72 块 GPU 通过高速 NVLink 互联组成一个紧密耦合的计算单元。
Wide Expert Parallelism#

对于 Mixture-of-Experts(MoE)模型,Wide EP 是一种将专家分散到大量 GPU 上的并行策略。图中上半部分展示了前向传播的数据流:token 经过 learned router 后,通过 all-to-all 通信被路由到持有对应专家的 GPU 上,计算完成后结果再通过 all-to-all 收集回来。E0-E3、E4-E7 等专家组分布在不同的 GPU 节点上。
Wide EP 的优势在于:将专家分散到更多 GPU 上可以获得更高的 batch size 和更高的 arithmetic intensity,从而提升 decode 阶段的效率。
故障处理#
图中下半部分展示了节点故障场景。当 72 个 GPU 中的一个节点宕机时:
- 取消受影响的请求,从 cache 或从头重新 prefill
- 重新分配孤立专家到剩余 GPU + 热备节点
- 排空失败节点上的 in-flight 请求
32 个以上 GPU 组成的 worker 中,故障是常态而非异常。硬件层面,NVLink 连接器是塑料材质而非金属,如果安装时用力过猛导致线缆弯曲,NVLink 连接就会变得不稳定。这需要自动化故障检测、专家迁移和请求重放机制。
新挑战#
- 百万 token 上下文:当前模型已经支持百万级上下文窗口。如何将这么长的上下文跨多个 GPU 并行处理?是否需要将 context 本身分片到不同 GPU?
- 容错设计:在 64 个 GPU 上拆分模型并服务数百万用户的生产流量时,单个 GPU 故障的影响范围和恢复策略都需要仔细设计。
Cache-aware Prefill-Decode Disaggregation#
理解了推理系统的各个组件后,可以从系统层面做一些简单但有效的优化。Cache-aware routing 是一个典型案例:两行路由代码,带来高达 40% 的服务加速。
问题:冷热请求混跑#
在一个 prefill-decode 分离的系统中,所有 prefill 请求被发送到同一组 prefill worker。但这些请求的计算量差异极大:
- 冷请求:用户开启一个全新对话,paste 进一整本书。这是一个完整的、零 cache hit 的大 prefill,可能有数万个 token 需要从头计算。
- 热请求:用户在一个已有 10 轮对话的 session 中追问一个简短问题。绝大部分 token 已经在 KV Cache 中,只需要对几十个新 token 做 prefill。
如果一个平均 10 轮的对话,那么大约 10% 的请求是全新的冷请求,90% 是增量式的热请求。将这两类请求混在同一组 GPU 上处理,冷请求的长时间 prefill 会拖慢热请求的响应。
解法:按 cache hit rate 分流#
做法非常简单:在路由层加两行代码,根据请求的 cache hit rate 分流——
- 低 cache hit rate 的冷请求 → 发往一组专用 prefill GPU
- 高 cache hit rate 的热请求 → 发往另一组 prefill GPU
这样冷请求之间互相「伤害」,不会影响到热请求的低延迟响应。这个优化可以带来高达 40% 的服务速度提升。
这种系统级优化不涉及任何 kernel 改动或算法创新,纯粹是对流量特征的观察和利用。类似的机会在推理系统中还有很多,因为我们对这些生产级流量模式的理解还处于非常早期的阶段。10-20 年后回看,这些优化可能显而易见,但在当下它们代表了对新型流量模式的首批系统性回应。
Megakernel:从操作级到模型级的 GPU 融合#
Decode 阶段的根本瓶颈在于:为了生成单个 token,必须加载整个模型的权重。GPU 的大规模并行计算能力在这里被浪费了,整个设备退化为一个「高级内存搬运工」。
传统 kernel 方式的开销#
通常的做法是为模型中的每种操作各写一个 kernel:Norm kernel、MatMul kernel、Attention kernel,依次启动和执行。这种方式易于编程,但引入了大量空闲时间。

这张图的 x 轴是时间,y 轴是 GPU 上的各个 Streaming Multiprocessor(SM)。H100 有 132 个 SM,B200 有 148 个。彩色条表示有用计算,空白区域是浪费。开销来源包括:
- Kernel Launch / Teardown:每个 kernel 的启动和清理需要时间,图中红色和黄色区域之间的大间隙就是 kernel 切换开销。
- Tail Effects(尾部效应):一个 batch 中的不同请求长度不一,短请求先完成后对应的 SM 只能空等长请求结束。这与 continuous batching 中长短请求共存的问题是同一个问题在更底层的映射。
- Memory I/O:权重加载和结果写回的等待时间。
当这些开销跨多个 kernel 累积时,GPU 的实际利用率远低于理论峰值。
Megakernel 思想#
Megakernel 的核心思路是:不再为每个操作写单独的 kernel,而是将多个操作融合为一个大 kernel。这类似于 Flash Attention 的融合思想,但更加激进,跨越了更多操作。
关键的观念转变是:将 GPU 看作一个大型分布式系统。132 个 SM 就是 132 个 worker,每个 worker 可以被分配不同的任务。哪些任务之间有依赖关系、哪些可以并行,由一个统一的调度器管理。
具体的优化手段包括操作间重叠:
- QKV + RoPE 与 KV Cache 加载重叠:在 QKV 投影和 RoPE 编码还在计算时,就开始加载 KV Cache 到 shared memory。等 QKV 完成后,KV Cache 已经就位,attention 可以立刻开始。
- Attention 与 O-Projection 权重加载重叠:attention 计算尚未结束时,就开始加载 O-Projection 的权重。
Llama-1B 整层 Megakernel#

这张图展示了将 Llama-1B 模型的整个 transformer layer 融合为单个 kernel 后的 SM 利用率。图中不同颜色代表不同的 sub-kernel:LayerNorm+QKV+RoPE(蓝色)、Attention Partial(橙色)、Attention Reduction(绿色)、O-Projection+Residual(红色)、FFN(粉紫色)等。各操作以复杂的方式交错执行:权重加载与计算重叠、前一步的 reduction 与下一步的矩阵乘法重叠。
实现框架:ThunderKittens#
整个 Megakernel 系统基于 ThunderKittens 库实现。ThunderKittens 可以类比为更底层的 Triton:它提供对 GPU 硬件的细粒度控制,允许开发者精确管理 SM 级别的调度和 shared memory 分配。框架采用 instruction-based abstraction,每个 sub-kernel 可以在独立文件中实现,通过虚拟化的 shared memory 系统协调它们的执行。
性能结果#
仅将 attention inference kernel 替换为 Megakernel 版本,就能获得 30%-70% 的加速。进一步将整层融合后,在 Llama-1B BF16 batch-size 1 的 decode 吞吐测试中,Megakernel 在 H100 上达到了 72% 的内存带宽利用率,接近硬件的物理速度极限(speed of light)。在 H100 和 B200 上,Megakernel 的吞吐量均大幅超过 vLLM 和 SGLang。
代价#
Megakernel 的代价完全在工程侧。一位全职的高水平 kernel 工程师,一年时间大约能为一种硬件上的 2-3 个模型、batch size 1-16 完成 Megakernel 开发。换到 batch size 17?从头来过。换到新硬件?从头来过。Together AI 正在开发编译器来自动化部分流程,但目前这仍然是一项极其劳动密集的工作。
Megakernel 的基本思想在 GPU 编程的几十年历史中经历过多次起伏。当前这一轮的驱动力是 decode 的带宽瓶颈和 LLM 推理的巨大商业价值。如果能做到,速度一定是最快的;问题只在于做到它需要付出多大的工程代价。
Parcae:用状态空间理论稳定循环 Transformer#
前面的讨论都围绕如何更快地运行现有的 Transformer 模型。Parcae 提出了一个不同的问题:参数量是通向更高质量的唯一途径吗?
Loop Transformer 的动机#
Loop Transformer 的核心思想是:取 Transformer 的若干层,让激活值在这些层中循环多次,而不是只通过一次。传入 → 前几层 → 循环块(重复 T 次)→ 后几层 → 输出。
这样做有几个理论上的好处:
- 参数量不变,FLOPS 可调:循环次数是一个旋钮,增大它就增加了计算量,但不增加参数。如果 more FLOPs = higher quality 的假设成立,这提供了一种不增加模型体积就提升质量的途径。
- 更高的表达能力:理论工作表明,循环模型能表达一些同等参数量的非循环模型无法表达的函数。
- 推理效率潜力:更少的参数意味着更少的 GPU 内存占用,可以留更多空间给 KV Cache,或者减少跨 GPU 的通信。如果循环块足够小,甚至可以写一个 Megakernel 让它在 on-chip memory 中就完成所有循环。
Maryland 大学 Tom Goldstein 组的 ARC 任务实验已经显示出循环模型的优势。Twitter 上还出现过一波关于 “Claude Mythos 是循环模型” 的猜测,发起者是一位 OpenAI 员工。Dan Fu 认为这个说法不可信。后来该员工发了博客文章承认是自己编造的。
不稳定性的根源#
然而,所有已有的 Loop Transformer 都面临一个严重问题:训练极其不稳定。做一个 learning rate sweep,9/10 的配置会导致 loss 爆炸。即使找到一个能训练的 learning rate(如 2e-4),稍作调整就崩溃。频繁的 loss spike 意味着训练过程出了根本性的问题。
之前的解决方案要么是在每一层加 norm(治标不治本),要么是严格锁定某个 learning rate 不做任何调整。Parcae 团队认为 loss spike 的存在暗示着更深层的结构性问题。
状态空间模型视角的分析#

Parcae 的关键洞察是将循环过程建模为一个离散动态系统。虽然 Transformer block 内部有 attention、GELU、RoPE 等大量非线性组件,但实证观测发现:每次循环对激活值的改变量(residual)其实很小。
将非线性部分统一记为 Rˉ(ht,e),整个循环可以写成:
ht+1=Aˉht+Bˉe+Rˉ(ht,e)
其中 ht 是第 t 次循环后的激活值,e 是初始注入(来自前几层的输出),Aˉ 控制激活值的线性变换,Bˉ 控制初始注入的权重。
由于 Rˉ 的量级在实验中被 Aˉ 和 Bˉ 项主导,暂时忽略它后得到一个可以用高中数学求解的线性系统。闭式解中出现了 AˉT 项(Aˉ 的 T 次幂)。如果 Aˉ 的 spectral radius(谱半径,可以理解为矩阵的「范数」)大于 1,比如 Aˉ 等效于标量 2,循环 16 次后激活值就被放大了 216 倍。这就是 loss spike 的数学根源。
回顾之前的 Loop Transformer 对 Aˉ 和 Bˉ 的选择:
| 方法 | Aˉ | Bˉ | 稳定性 |
|---|---|---|---|
| Addition(直接相加) | I(单位矩阵) | I | 边际稳定 |
| Concatenation(拼接) | Rdh×dh(全可学习) | Rdh×de | 不稳定 |
单位矩阵 Aˉ=I 意味着谱半径恰好为 1,处于稳定边界上(marginally stable);全可学习矩阵则很容易学出谱半径 > 1 的参数,直接不稳定。
Parcae 的稳定化方案#
Parcae 对 Aˉ 和 Bˉ 施加约束:
- Aˉ:参数化为负对角矩阵。对角元素为负意味着 AˉT 随 T 增大趋向于零而非爆炸,系统天然稳定。
- Bˉ:加一个简单的线性 norm。Bˉ 只被应用一次(初始注入),不会像 Aˉ 那样被反复取幂,因此轻量约束就足够了。
这个方案的思路直接来自控制论和状态空间模型(SSM)理论:要让一个线性递推系统稳定,控制其状态转移矩阵的谱半径小于 1 就够了。
实验结果验证了这个分析:
- 未约束的模型(baseline):激活值 norm 爆炸到 1019
- 加 norm 的模型(之前的做法):norm 被控制住了,但 loss 依然有剧烈 spike。原因是模型内部「想扩大」激活空间以获得更好的表示,而 norm 层又把它压回来,两股力量持续对抗
- Parcae:激活值和 loss 都稳定。即使使用 6e-4 这个让其他模型崩溃的 learning rate,训练也能顺利完成
在 benchmark 比较中,Parcae 在多个评测上超越了之前的 Loop Transformer(Recurrent Depth Model),也超越了社区高度优化的 nano-scale Transformer baseline(社区竞赛中众多参与者反复调优的版本)。
循环 Transformer 的 Scaling Law#
Parcae 证明了循环 Transformer 可以稳定训练且质量更好。下一个自然问题是:循环次数应该随训练规模如何变化?
经典 Scaling Law 的回顾#
Chinchilla 等工作建立了参数量和训练数据之间的 scaling law:给定一个 FLOP 预算,最优策略是同时增大模型参数和训练 token 数。在对数坐标下,最优配置的轨迹向右下方倾斜,意味着两者应该联合增长。如果轨迹垂直向下,说明只需增加数据;如果水平向右,说明只需增加参数。实际观测到的斜率介于两者之间,所以 Llama 4 这样的模型会用万亿参数在 35 万亿 token 上训练。
循环次数作为第三个缩放维度#
Parcae 在经典的「参数 vs 数据」二维空间中引入了第三个维度:循环次数(recurrence)。
实验设计:固定参数量(140M 和 370M 两个规模),变化训练数据量和循环次数,观察 loss 的等高线走向。
关键发现是:在固定参数量下,随着训练数据增加,最优循环次数也应该增加。等高线同样向右下方倾斜,与 Chinchilla 在参数-数据空间中观测到的模式完全一致。

图中展示了四组幂律拟合:
- 左两列:μrec Power Law:最优循环次数随 FLOPs 的增长关系。140M 模型为 μrec∝C0.40,370M 模型为 μrec∝C0.38。幂指数约 0.4 意味着 FLOPs 增加 10 倍,最优循环次数大约增加 2.5 倍。
- 右两列:Token Power Law:最优训练 token 数随 FLOPs 的增长关系。D∝C0.77 和 D∝C0.78,与已知的 Chinchilla scaling law 一致。
这些幂律拟合质量很好,数据点紧密分布在拟合线附近。
三维 Scaling 的含义#
团队还做了一个三维实验(同时变化参数量、数据量和循环次数),结论指向同一个方向:三者应该联合增长。三维图因为难以可视化没有详细展示,但定性结论一致。
将固定参数量的实验用另一种方式呈现:橙色曲线是传统 Transformer(固定深度,无循环),蓝色曲线是在相同 FLOP 预算下通过增加循环次数得到的最优 Parcae 模型。在相同 FLOP 预算下,蓝色曲线持续低于橙色曲线,意味着将一部分 FLOPs 分配给循环(而非纯增加数据)可以获得更低的 validation loss。
当前所有主流 LLM 的循环次数都是零,位于这些 scaling 曲线的最左端。而 scaling law 的趋势显示,随着训练数据量增加,最优策略应该包含越来越多的循环。这暗示着当前的大规模预训练可能存在改进空间。
预训练模型上的循环#
一个有趣的观察是:在已训练好的模型上直接做循环(不做任何额外训练),有时也能提升质量。有人将 Qwen 模型的 2-3 层做简单循环,在某些数学 benchmark 上获得了更好的结果并登上了排行榜,整个过程没有训练任何参数。这个现象令人困惑,Parcae 团队计划深入研究权重层面的机制解释。
推理效率视角#
循环模型的推理效率优势体现在非线性收益上:
- 更少的参数 → 更少的 GPU 内存占用 → 可以存更大的 KV Cache → 可以服务更多并发 session
- 更少的参数 → 可以用更少的 GPU 跑同一个模型 → 减少跨 GPU 通信
- 如果循环块足够小,可以写 Megakernel 让权重常驻 on-chip memory,循环在片上完成
下一代 Groq LPU 芯片只有约 250 MB 的片上内存。常规模型放不进去,但一个足够小的循环块也许可以。如果能跨过这个阈值,decode 速度会有质的飞跃。
Q&A 精选#
循环 vs 增大参数:何时该选循环?#
给定固定 FLOP 预算,compute-optimal 策略通常是同时扩大模型和数据。如果模型大小已经被限定(比如要在笔记本上部署),那就增加循环次数。如果数据已经用完,也可以通过增加循环来继续利用 FLOP 预算。最终的选择还要考虑实际部署因素:模型能不能被开源社区方便地使用?能跑在什么硬件上?这些工程约束会反过来影响参数量和循环次数的平衡点。
硬件-架构协同设计#
为特定推理芯片设计模型时,内存是首要约束。做法是先查目标芯片的内存容量,确保模型放得下且有足够空间留给 KV Cache。
一个值得关注的信号是:中国模型团队的最新作品在架构选择上呈现出某些特征,这些特征暗示他们可能在为华为芯片做协同设计。量化格式的选择也反映了硬件绑定:NVIDIA 的 Nemotron 模型在 NVFP4(NVIDIA 专有的 FP4 格式)下训练;AMD 阵营则使用 MXFP4。两种格式各有优劣,选择哪种直接取决于目标部署硬件。
不同使用场景下的最优架构差异#
Agent 工作流和批处理任务对架构的需求差异巨大。Agent 场景最在意的是 KV Cache 保持热状态,因为用户会反复与同一个 session 交互。DeepSeek 的 MLA(Multi-Latent Attention)是一种激进的 KV Cache 压缩方案,对 agent 场景特别有价值。FP4 量化的 KV Cache 也能显著压缩存储。
另一个大分界线是 causal attention vs bidirectional attention。批处理任务(如搜索排序)不需要逐 token 生成,Google 至今可能仍然在搜索中使用 BERT 模型(bidirectional),因为只需做一次前向传播得到一个向量表示。T5 曾经是一个折中方案:bidirectional encoder + causal decoder。而 chat/agent 场景必须使用 causal attention 以支持自回归生成。
多 GPU 通信与 Megakernel#
NCCL 通信调用可以被融合进 Megakernel。DeepSeek 4 发布了一个 MoE 推理层的 Megakernel,实际融合了部分跨 GPU 通信。但目前还没有找到一个「杀手级用例」,因为有时候瓶颈就在 NCCL 调用本身的延迟上,融合并不能消除它。
未来的趋势可能是:模型的某些计算密集部分各自有一个 Megakernel,而不是整个模型一个 Megakernel。除非愿意付出巨大的工程代价(blood, sweat, and tears),否则全模型 Megakernel 不太现实。
部分内容可能已过时
评论