思考探索:高效 AI 的状态账本:少搬状态,比少算一次更重要

思考探索:高效 AI 的状态账本:少搬状态,比少算一次更重要

Charles Lv8

过去我理解高效 AI,首先会数 FLOPs:矩阵乘少了多少,attention block 跳过多少,模型调用减少几次。连续读完这一组训练与推理论文后,我更在意另一张账:一次训练或生成里,系统保存了哪些状态,这些状态被复制几份,何时跨设备搬运,又在什么时刻才真正需要完整物化。

这里的“状态”范围很宽。参数、梯度和 Adam 一二阶矩是状态;attention 的 Q/K/V、反向重算中间量和长上下文 KV cache 也是状态;投机解码的候选树、stereo 的 cost volume、量化 scale 与 packed weight,同样会占据内存、带宽和调度位置。

我现在会把高效系统的目标写成一句更具体的话:

少搬、少复制、延后物化,只在真正的消费者到来时保留足够精度的状态。

“少算”仍然重要。只是很多系统真正省下来的时间,来自计算前后那条更长的数据路径。反过来,一项算法即使显著减少理论 FLOPs,只要制造了更多格式转换、全量同步、临时缓存或无效候选,收益就可能停在局部 benchmark 里。

15 篇材料分别在移动什么

下表先把 15 篇文章放进同一张状态账本。表中的实验数字和系统结论属于论文作者或官方材料报告;“我的读法”是这里的跨文综合,不是这些论文共同证明的定律。

分组 材料 主要状态 怎样减少无效移动或复制 新增代价 / 我的读法
训练状态 ZeRO 参数、梯度、optimizer states 在 DP 组内分片,当前层需要时再 gather 常驻副本变少,通信、checkpoint 和状态生命周期更复杂
训练计算 Megatron-LM MLP、attention、vocab 的 tensor shards 按 column/row/head 切分,把同步推迟到必要边界 切分是否有效取决于 GEMM shape、非线性位置和通信域
长上下文 RingAttention 长序列 Q/K/V blocks query 留在本地,KV 沿 ring 轮转,避免每卡持有全序列 需要足够大的 block 和高带宽互联,长上下文本身仍要付 FLOPs
异构注意力 MagiAttention heterogeneous mask、KV/dKV、dispatch metadata 按有效 mask area 分配计算,只向需要者 GroupCast KV solver、metadata、group collectives 与 overlap 调参进入系统
低精注意力 Attn-QAT Q/K/V/P、反向重算量、量化 scale 让 forward、backward recomputation 与 inference kernel 共用低精语义 为了正确部署,训练时反而要保留高精辅助状态 (O’)
优化器 Muon 矩阵动量、正交化更新 矩阵参数只保留一份 momentum,按矩阵几何更新 完整矩阵正交化需要 DP gather;少状态不等于零通信
稀疏计算 SLA / SLA2 attention blocks、router、linear compensation critical 精确算,marginal 线性补,negligible 跳过 路由、近似误差、微调和低比特 kernel 共同决定收益
性能观测 GPU Utilization kernel timeline、SM、tensor pipe、通信等待 不直接搬状态;负责识别哪些搬运和 kernel 没产生有效 FLOPs 100% utilization 只能证明 GPU 忙,不能证明状态流动有效
KV 组织 PagedAttention 每个请求持续增长的 KV cache logical block 映射到 physical block,按需分配并共享前缀 block table、copy-on-write 和 scheduler 成为正确性边界
KV 内容 KVSlimmer Key 与 Value cache entries 合并更同质的 Key,尽量保留异质 Value;权重只用 forward 变量 平均分之外仍要检查关键证据是否被局部合并抹掉
Draft 接口 EAGLE target hidden feature、draft tokens、candidate KV 复用 target 已算出的 feature,不再运行一个完整小 LLM target 必须暴露稳定 feature,draft 开销仍要由接受长度摊销
Draft 调度 EAGLE-2 draft tree nodes 与 confidence 只扩展 path value 高的节点,再动态重排验证树 confidence 失准时会把验证预算投错分支
Draft 训练 EAGLE-3 draft 自己生成的递归状态 训练时就把自生成状态喂回模型,减少深层候选浪费 仍需 tree mask、临时 KV 和 target verification 共同正确
实时视觉 Fast-FoundationStereo feature pyramid、cost volume、ConvGRU hidden state backbone 蒸馏、cost filtering 搜索、refinement 剪枝分别处理 搜索成本、伪标签门禁、相机标定与 TensorRT 路径仍在账内
数值格式 Low-bit LLM Survey packed weights、activation、KV、scale、zero point 用低 bit 减少常驻字节与 HBM / 通信流量 Q/DQ、scale 读取、layout 和硬件 kernel 可能吃掉节省的时间

把这些论文并排后,可以看到一个很稳定的分工:ZeRO、PagedAttention 处理“存几份”;Megatron、Ring、Magi 处理“在哪算、往哪传”;Attn-QAT、low-bit survey 处理“以什么格式传”;SLA、KVSlimmer、EAGLE-2 处理“哪些内容值得传”;EAGLE-3 和 Fast-FoundationStereo 继续追问“训练出来的状态是否真能沿部署路径被消费”。

先给状态分工,再谈压缩

同样占 1GB 的张量,职责可能完全不同。对它们使用同一种压缩策略,通常会把系统问题伪装成一个统一 bitwidth 或统一 sparsity 参数。

状态职责 典型对象 真正的消费者 常见误判
长期知识 模型参数、embedding、expert weights 每次 forward 的算子 只看模型文件大小,不看每 token 激活和权重搬运
训练记忆 gradient、momentum、variance、master weight optimizer step 与故障恢复 只看参数 dtype,不算 optimizer state 和 checkpoint
临时计算证据 activation、Q/K/V、softmax 统计、cost volume backward、refinement 或下一融合算子 以为临时张量寿命短,所以移动成本不重要
历史记忆 LLM KV、长视频 KV、recurrent hidden state 后续 decode 或 rollout 只看单请求容量,不看并发、共享和释放时机
候选状态 draft tree、beam、检索结果、规划分支 昂贵 verifier 或 target model 候选越多看似覆盖越好,却可能放大验证与临时 KV
数值元数据 scale、zero point、block map、router mask kernel、allocator、scheduler 元数据很小,就误以为它不会改变热路径

SLA 和 KVSlimmer 给了两个很直观的反例。SLA 观察到,高权重 attention blocks 仍可能是高秩结构,应该精确计算;大量中间权重单个不大,却会累计影响输出,适合 linear compensation;真正可跳过的是贡献极小的部分。KVSlimmer 则把 K 和 V 拆开:Key 更像寻址,Value 更像内容。相邻 Key 可合并,不代表相邻 Value 也能平均。

压缩前,我最先问“这个状态一旦出错,谁会受影响”,然后才问“它能压到几 bit”。如果消费者是 softmax 反向、长文档证据检索或机器人深度边界,一次局部误差可能沿后续状态链放大;如果消费者只是低价值候选排序,近似空间会大得多。

训练系统先消除副本,再安排同步边界

ZeRO 最清楚地展示了“副本”有多贵。mixed-precision Adam 的 model states 粗略可到每参数 16 bytes;普通数据并行让每张卡都长期保存完整参数、梯度和 optimizer states。ZeRO 依次分片 optimizer、gradient 和 parameter,把“每卡完整拥有”改成“DP 组共同拥有,计算时短暂恢复”。

Megatron-LM 做的是另一层重排。MLP 第一层按输出维切,GeLU 留在本地;第二层按输入维切,到 block 输出才 all-reduce。attention 按 head 切,也把同步放在 output projection 之后。这里没有删掉 Transformer 的数学计算,收益来自中间 activation 不必在每个非线性前来回聚合。

RingAttention 把同一思路推到序列维度:query block 留在设备上,KV blocks 在设备环上逐站经过。每个 query 最终仍看到全序列,省掉的是“先把全部 KV 复制到每张卡再算”的冲动。

RingAttention method 原图

图源:Ring Attention with Blockwise Transformers for Near-Infinite Context,Figure 2。原图上半部分展示 query block 留在本设备、KV block 沿 ring 传递,下半部分展示 blockwise attention 与 feedforward 的组织。这里的重点是“流动状态”可以替代“复制状态”;这张图不能证明长上下文数据一定有用,也不能证明慢互联下通信总能被隐藏。

MagiAttention 又给 Ring 的规则通信加了一层 workload awareness。heterogeneous mask 下,并非每个 rank 都需要每段 KV。AttnSlice、dispatch solver 和 GroupCast/GroupReduce 让数据移动感知有效 mask area,减少把无用 KV 绕完整个 ring 的情况。

Muon 提供了一个有意思的反例:它为了对完整矩阵动量做正交化,需要额外 gather。单看通信字节,它未必总比 AdamW 简单;但矩阵参数只保存一份 momentum,并尝试用更好的更新几何减少达到同 loss 所需的训练 FLOPs。必要的搬运可以保留,应该先删掉的是那些没有产生足够训练价值的搬运。

推理系统把缓存和候选变成可调度资源

PagedAttention 的贡献很像一次内存观念更新。KV cache 不再是一条必须连续增长的大数组,而是一组 logical blocks。物理 block 可以分散存放、按需分配、请求结束后立即回收;共享 prompt 的多个候选先共用 block,真正分叉时才 copy-on-write。

这项设计没有减少每个历史 token 的 K/V 内容,也没有降低 attention 读取历史的理论成本。它先提高同一显存里可容纳的 active tokens,再让 continuous batching 把容量变成吞吐。内存抽象和 scheduler 必须一起工作。

KVSlimmer 继续处理“block 里面装什么”。它保留 Value 的异质信息,主要合并相邻 Key,并把原本需要 Hessian/gradient 的权重化成 forward-only 计算。PagedAttention 管物理组织,KVSlimmer 改内容表示;两者作用层不同,也可以叠加。叠加后必须增加 lost-evidence 测试,因为容量增长不应以丢掉少数关键 token 为代价。

EAGLE 系列处理的是另一类临时状态:尚未被 target 接受的未来。

1
2
3
EAGLE:复用 target feature,降低提出一个候选的成本
EAGLE-2:按 confidence / path value,只物化更值得验证的树节点
EAGLE-3:训练时见到自己的递归状态,提高深层候选被接受的概率

三篇连起来后,投机解码更像一个状态分配器。draft 负责便宜地产生候选状态,tree policy 决定物化哪些分支,target verification 决定哪些 token 和 KV 可以提交。拒绝分支必须被清理,兄弟节点不能互相偷看,临时 KV 也要进入显存账本。平均接受长度只是算法指标,P95 latency、峰值 KV 和 batch throughput 才能说明系统是否闭合。

Fast-FoundationStereo 把这种按职责分配方法带到实时视觉。feature backbone 保存的是视觉先验,适合蒸馏;cost filtering 操作小通道 4D cost volume,直接剪枝收益有限,所以用分块搜索;ConvGRU refinement 有重复迭代和通道冗余,更适合结构化剪枝。论文作者报告 3090 上 runtime 从 496ms 降到 49ms,TensorRT 路径约 21ms。这个结果来自三段热路径共同缩短,也依赖 rectified 输入、伪标签过滤和部署 backend。

状态格式必须和部署计算语义一致

状态移动不只有“多少字节”,还包括 dtype、layout 和重算语义。低比特 survey 反复提醒:INT4 checkpoint 变小,可能只是 host 或 HBM 中少搬了权重;如果 kernel 先把权重 dequant 成 FP16,再进入普通 MatMul,算力路径并没有一起变短。Weight-only、W&A 和 KV quantization 应该分开验收。

Attn-QAT 把这个问题推进到 fused attention 内部。FlashAttention 为省显存,会在 backward 重算概率并使用高精度 softmax 恒等式。forward 一旦换成 FP4,backward 若仍假设原来的高精度中间量,梯度语义就会错位。Attn-QAT 因此在训练中额外保留高精度辅助输出 O',并让重算的量化路径与 inference kernel 对齐。

这个例子改变了我对“省状态”的直觉。为了最终少搬 FP4 attention,训练时可能要暂时多保存一个高精度状态。多出来的 buffer 不是浪费,它在保护训练和部署对同一算子的理解。真正该删的是没有消费者、重复保存或无法转成端到端收益的状态,不是所有高精度中间量。

SLA2 的 low-bit sparse branch 也需要放在这条边界上读。SLA2 专题页描述了低比特 forward 与高精 backward 的 QAT 路径;Attn-QAT 则表明,某些 FlashAttention-style backward 不能简单沿用高精假设。两者模型、格式和 kernel 不同,不能据此判定冲突已经发生;如果未来把 SLA2 routing、FP4 attention 和分布式长序列 kernel 合并,backward recomputation 应单独做数值回归,而不是默认组件可直接拼装。

EAGLE-3 和 Fast-FoundationStereo 也在处理同一种接口问题。EAGLE-3 让 draft 训练时看到部署中的自生成状态;Fast-FoundationStereo 的搜索目标直接记录候选 block 的真实 runtime,并给出 TensorRT 路径。算法图里更小的张量,只有进入真实 runtime 后才算完成加速。

每次加速都会把瓶颈推给下一层

作者报告的几个数字很能说明问题。

局部改进 作者报告 下一处压力
GPU 一直有工作 Trainy 案例中 GPU utilization 为 100%,MFU 仍约 20%;优化后约 38% softmax、memory-bound kernel、通信和并行布局
SLA attention kernel forward attention 相对 FlashAttention2 约 13.7x 端到端视频生成约 2.2x,VAE、MLP 和采样循环成为大头
ZeRO state sharding 每卡 model-state memory 可随 DP degree 大幅下降 all-gather、reduce-scatter、checkpoint 与恢复
Ring / Magi 长上下文 单卡不必持有全部序列 KV interconnect、block size、mask solver 和最慢 rank
PagedAttention 碎片和前缀复制减少,可放入更多 active tokens scheduler、block table、临时 speculative KV;历史读取仍然存在
KVSlimmer cache 和 decoder 开销低于反传式 merge baseline 关键证据丢失、layer/head 压缩比例与非局部冗余
Low-bit 权重、activation 或 KV 字节下降 scale、packing、Q/DQ、寄存器压力和 kernel availability

所以我不把“瓶颈迁移”理解成优化失败。它反而是成功优化后的正常现象。问题在于,团队是否提前知道下一层是谁,并用端到端指标确认局部收益确实传了下去。

GPU utilization 那篇文章给了一个很实用的顺序:先看设备是否 idle,再看 MFU、step time 或 tokens/s,然后用 SM、tensor pipe、memory bandwidth 和 timeline 找具体阶段。推理侧也应采用相同思路:先看 TTFT、TPOT、吞吐和 P95/P99,再追 KV block、draft、Q/DQ 或 attention kernel。只看一张“100%”表盘,很容易把忙于搬运误读成高效计算。

我会怎样审计一个高效 AI 方案

下面是我从这 15 篇材料里整理出的状态审计表。它不是标准 benchmark,只是比单独数 FLOPs 更接近系统真实路径。

问题 要查的对象 失败信号
谁拥有状态 replicated、sharded、shared、copy-on-write 每个 rank / candidate 长期保存同一份内容
状态活多久 常驻、按层、按 step、按请求、按 phase 为极少使用的最大长度或分支提前预留
谁消费状态 backward、下一 token、verifier、planner、refinement 压缩指标正常,但关键消费者的任务质量下降
以什么格式存在 BF16、FP4、INT4、packed block、linear compensation Q/DQ 与 layout 转换占据热路径,或训练/部署语义不同
沿哪条路径移动 HBM、NVLink、InfiniBand、PCIe、CPU、磁盘 通信无法被计算覆盖,allocator / scheduler 出现空洞
失败后怎样恢复 rejected KV、checkpoint shard、cache block、视觉伪标签 状态污染下一个请求,或恢复后 loss / 输出不连续
端到端是否获益 MFU、step time、TTFT、TPOT、P95、任务成功率 kernel/FLOPs 很漂亮,用户侧延迟或质量没有改善

我会再补一张很朴素的账:

1
2
状态账本 = 常驻状态 + 临时状态 + 缓存状态 + 共享状态 + 重算状态
移动账本 = HBM 读写 + 设备互联 + CPU/GPU 搬运 + 格式转换 + 分配/回收

如果一定要用一个式子概括我的判断,我会写成:

有效效率被任务消费者真正使用的状态更新计算+搬运+复制+临时物化+恢复成本.\text{有效效率} \approx \frac{\text{被任务消费者真正使用的状态更新}} {\text{计算}+\text{搬运}+\text{复制}+\text{临时物化}+\text{恢复成本}}.

这不是论文里的正式指标。它只是提醒我:一次更少的矩阵乘,如果换来更多无效状态、错误精度或昂贵恢复,未必更高效;一次额外 gather 或高精 buffer,如果显著提高收敛、接受率或部署正确性,也可能很划算。

证据边界

事实底稿来自站内 15 篇专题页及其对应论文或官方材料。文中的 16 bytes/parameter、MFU、kernel speedup、runtime、context length 和 benchmark 结果都属于作者报告,不等于独立复现。不同论文使用的模型、硬件、数据、runtime 和 workload 不一致,数字不能直接横向排名。

“高效 AI 更像状态移动问题”“按需物化优于长期完整持有”“瓶颈会沿状态链迁移”是我的跨文综合。它们适合用来检查系统设计,还不是已经有统一实验证明的理论。具体项目仍应测自己的形状、流量、互联、任务质量与恢复路径。

  • Title: 思考探索:高效 AI 的状态账本:少搬状态,比少算一次更重要
  • Author: Charles
  • Created at : 2026-07-12 09:00:00
  • Updated at : 2026-07-12 09:00:00
  • Link: https://charles2530.github.io/2026/07/12/ai-files-thinking-exploration-efficient-ai-is-state-movement/
  • License: This work is licensed under CC BY-NC-SA 4.0.
Comments