几个月前，我们组接到一个让人头大的任务：一台 8 卡 H20 推理服务器，要跑一套 600B 级别的 MoE 大模型，项目内部代号 DeepSeek-V4.1（下面都简称 V4.1），长上下文能力要推到 512K token，首字延迟压到 3 秒以内。需求来自几个企业客户，他们要处理超大型文档问答和代码仓库级分析，业务方的耐心上限基本就是两三秒，超过这个数字产品体验约等于零。结果基线一测直接凉透：TP=8 硬塞，512K 上下文进来之后 prefill 阶段把显存几乎吃满，首 token 出现在 33 秒量级。最终我们做到了 2.6 秒上下，正好 12.7 倍。
这篇实录把整个过程完整记录下来：为什么 512K 首字延迟这么难降、8 卡 H20 的显存算力账怎么算、按什么顺序做了哪些优化、每一步实测数字是多少，以及踩过的那些改一下就崩、折腾很久才搞明白的坑。适合正在做长上下文推理优化的工程同学参考，也适合对 MoE 推理、PD 分离、稀疏注意力这些概念有兴趣但还没实际动过手的读者。
1. 先认清问题：512K 首字延迟为什么难啃
1.1 模型与硬件基线
先说基线配置。硬件是 8 卡 H20，单卡 96GB HBM3，显存带宽 4TB/s 级别。模型是 600B 级别总参数、40B 左右激活参数的 MoE 架构，attention 采用了 MLA（Multi-head Latent Attention），也就是对 KV Cache 做了低秩压缩的那种结构。训练时的最长序列是 128K，但我们要推理 512K，上下文长度直接外扩 4 倍。
这里需要先理解一个关键事实：长上下文推理最难受的不是 decode，而是 prefill。decode 阶段一次只生成一个 token，瓶颈是显存带宽；prefill 阶段要把用户给的 512K token 全部前向计算一遍，同时还要为每个 token 生成 KV Cache，瓶颈是算力加带宽，而且 attention 的计算量随序列长度平方增长。所以 prefill 阶段的首字延迟，本质上就是从"海量并行计算"里抢时间的问题。
1.2 目标拆解：不是把 TTFT 降下来，而是把整条流水线节奏调对
我们的目标指标是首 token 时间（TTFT）。业务方给的需求是：512K 输入，首字不超过 3 秒。单独看这个目标很吓人——哪怕按最优状态估算，512K token 要在 3 秒内完成 prefill，相当于每秒要处理 17 万 token 以上的 prefill 吞吐，这对 8 卡 H20 来说不是靠一两个 kernel 优化能做到的，必须把整条链路从上到下重排一遍。
所以我拿到任务后做的第一件事不是写代码，而是把整条推理链路拆开：显存能装下多少、prefill 和 decode 要不要分开、并行策略怎么切、attention 的计算能不能稀疏化。后面你会发现，真正把 TTFT 打下来的，不是某一个"魔法技巧"，而是一条链路上五六个环节的累积收益。
2. 先算账：显存、算力与通信三条硬约束
2.1 KV Cache 与权重显存估算：MLA 救了半条命
先算显存。600B 参数如果用 FP16 存权重，光权重就是 1.2TB，8 张卡 768GB 完全装不下，所以权重必须压缩。我们最终用 FP8 存权重，实际占用约 600GB，摊到每卡 75GB，剩余每卡只有 21GB 多一点，非常紧张。
然后算 KV Cache。V4.1 因为用了 MLA，每个 token 的 KV Cache 被压缩成一个低维向量，和传统 MHA 那种"头数 × 头维度 × 2"的存法完全不是一个量级。以 60 层、每层每 token 约 576 维 latent 计算，512K token 的 KV Cache 大约是 512K × 60 × 576 × 2 字节，约 36GB。这个量级对 768GB 显存来说完全可接受。如果是传统 MHA 结构，这个数字至少要乘几十倍，512K 长度基本就是不可能任务。所以模型架构本身就是长上下文的第一层优化，这个优势我们只是"保住"而不是"创造"。
但显存装得下不代表能跑得快。KV Cache 虽然只有 36GB，prefill 阶段的中间激活、临时张量才是真正吃显存的大户。我们用 4bit 对 KV Cache 再做了一层压缩后，KV 降到 9GB 左右，这才给 prefill 和 decode 的调度留出了操作空间。
2.2 算力与带宽：为什么"硬算"打不进 3 秒
H20 的显存带宽大概在 4TB/s 级别，8 卡合计 32TB/s 左右；FP16 算力按单卡 140T 左右估算，8 卡约 1.1P。512K token 的 prefill 阶段，模型要跑几十层 Transformer，每层 attention 的计算量都和序列长度平方挂钩。粗暴估算下来，512K 长度下所有 attention 的总 FLOPs 在几十到上百 PFLOPS 量级，哪怕算力打满也要 10 秒以上，这还不算通信、kernel 启动、显存搬运的损耗。
所以 3 秒目标在"老老实实全量 attention + 单机并行"的框架下根本不现实。必须从两个方向同时下手：减少计算量（稀疏 attention）和让计算更高效（PD 分离、内核融合、通信重叠）。算完这笔账，我心里大概有数了：这项目不是优化某个 kernel 能解决的，而是要从调度架构层面重做。
3. 核心优化点：按收益排序逐个拆开
3.1 权重 FP8 与 KV Cache 4bit：先把显存水位打下来
第一件事是压缩权重到 FP8。600B 参数 FP8 后约 600GB，8 卡每卡约 75GB，剩余 21GB 给中间激活和 KV。这里有一个易错点：MoE 模型的 expert 权重在 8 卡之间天然不均匀，FP8 切分时必须按 expert 维度重新做负载均衡，否则某一刻某些卡会先 OOM。
KV Cache 从 FP16 压缩到 4bit。MLA 的低秩 latent 向量数值分布比较稳定，直接按 tensor 维度做 min-max 量化问题不大，但粗暴量化会带来明显的 benchmark 掉点。我们最终按 per-head 分组量化，把极少数的 outlier 维度分离出来单独存 FP16，KV 占用从 36GB 降到 9GB，掉点控制在了可接受范围。单独看这一步对 TTFT 提升很有限，但没有它，后面 PD 分离和长序列调度都玩不转。
3.2 PD 分离：让 Prefill 和 Decode 各干各的
PD 分离是我们的核心策略之一。传统模式下，同一条 GPU pipeline 既要处理 prefill，又要在 decode 阶段逐 token 生成，两者对资源的需求完全相反：prefill 是计算密集的"大爆炸"，decode 是访存密集的"细水长流"。混合跑的时候，prefill 会把 decode 的小请求堵死，decode 又拖慢 prefill 的吞吐。
我们把 8 卡拆成两个角色：6 卡专门跑 prefill，2 卡专门跑 decode，中间通过 KV Cache 共享池衔接。prefill 阶段算完的 KV 直接写入共享显存池，decode 卡按需拉取。这里要注意：decode 阶段要访问的是 512K 的全量 KV，跨卡拉取带宽压力非常大。我们最终让 prefill 侧在写 KV 时按 sequence 维度做分片，decode 侧用环形方式循环读取，避免单一 all-to-all 通信变成热点。
PD 分离后，prefill 不再被 decode 的 token 步长拖住，TTFT 直接掉了一截，更重要的是 decode 阶段的单 token 延迟也稳定了，这是产品体验里另一个容易被忽视的隐性指标。
3.3 序列并行与环形注意力：让 512K 分布到多卡
单卡处理 512K 的 attention 不现实，必须把序列维度切开。我们用序列并行（Context Parallel）把 512K 的序列切成 4 段，每卡处理 128K。切法上做了两层：第一层切 sequence，每卡只算自己的 128K 段；第二层加环形 all-gather 通信，因为 attention 的 softmax 需要全局归一化信息，每卡必须拿到其他卡的 K/V 的一部分才能算完整结果。
为了避免每次 attention 都做全量 all-gather，我们实现的是分块稀疏注意力协议：本地段只对本地 attention 做全量计算，跨卡的远距离依赖通过稀疏窗口处理。这个方案的通信量从 O(L²) 降到接近 O(L)，环形通信的延迟也能被计算重叠掉。还有一个细节：序列并行切分后，每个 micro-batch 的真实长度不再一致，动态 padding 到 128K 会浪费大量算力，我们改成按真实长度做不齐 pad 的 batch，pad 率从 25% 降到了不足 5%。
3.4 稀疏注意力与锚点 token：从"全看"到"挑着看"
这是把 TTFT 打下来最狠的一刀。512K 全量 attention 的计算量不可接受，但分析业务场景后发现：长文档问答里，大部分 token 的真正依赖是局部的，加上少数全局的"锚点"。所以我们把 attention 模式改成两部分叠加：局部窗口注意力加全局锚点注意力。窗口大小设成 16K，配合 128 个均匀分布的锚点 token，每个位置都保持对全局锚点的注意力，保留了对全局信息的感知。
这个改动把单卡 attention 计算量砍到原来的约四分之一。实现时要注意：不能简单地把非窗口位置的 attention 置零就完事，否则模型在长上下文上的表现会崩塌。我们用训练阶段做过的稀疏 attention 位置规划来对齐推理稀疏模式，并在窗口边界做了位置编码的平滑过渡。实测在长文档问答 benchmark 上掉点控制在一个点以内，TTFT 收益却是接近 1.9 倍。
3.5 RoPE 改造：高频外推与局部位置解耦
模型训练只到 128K，直接推到 512K，最核心的问题不是显存，而是位置编码。标准 RoPE 在超出训练长度后，相邻位置的高频旋转变量会出现混淆，具体表现是长距离检索和文档内引用全部失灵。我们采用的方法是把 RoPE 分成两路：局部窗口内用高频位置编码，负责精确的相邻位置关系；全局锚点之间用低频位置编码，负责长距离结构信息。这样避免同一个维度既要负责局部微调又要负责全局定位的矛盾。
实现上，对关键 attention 层做了可学习的温度缩放，把位置空间的频率上限压低，保证 512K 范围内不会溢出到频率混淆区。这一步对 TTFT 没有直接收益，但没有它，前面所有稀疏化、并行化的优化都会变成"生成一堆幻觉的加速器"。长上下文项目的数值稳定性永远是第一位。
3.6 投机式 Prefill 与动态批量调度：把峰值压力削掉
到这一步，显存和计算量都降下来了，prefill 仍有最后一个浪费点：动态 batch 里的长度差异。512K 的请求进来，如果 batch 里还有几个 100K 的小任务，GPU 需要等最长的 512K 算完才能推进。我们用投机式调度：把输入按长度分成多个 bin，每个 bin 内部再按 prefill 计算量做动态拆分，短任务可以插进长任务的空隙里执行。
具体方式是把过长的 prompt 按语义边界切成若干 chunk，并行 prefill 之后再合并 KV。这里有个看着简单但容易翻车的细节：KV 合并要按 chunk 的位置顺序串起来，不能简单 concat，否则 decode 阶段的位置编码会错乱。我们在合并时对每个 chunk 重新计算了 RoPE offset，确保 decode 阶段读到的位置信息是连续的。这一步对单个超长请求的 TTFT 收益不算最大，但把整个服务的吞吐稳定度提升了约 1.4 倍。
3.7 内核融合、CUDA Graph 与通信重叠：榨干最后一层效率
前面几板斧把大方向的复杂度砍下来了，剩下的就是 kernel 级和系统级的水分。我们做了三个比较重要的优化：第一是 attention 内核融合，把 QK 内积、softmax、PV 加权三段融合成一个 kernel，减少中间量在显存里的往返；第二是 CUDA Graph 捕获，把 prefill 阶段可能出现的动态 shape 全部用静态 shape 桶化，捕获一次后重复回放，避免每步 kernel 启动的开销；第三是通信重叠，把 all-gather 和计算交错排布，让显卡在等网络数据时先把本地段算完。
这三个优化听着不如前几个"性感"，但累积起来大概贡献了 1.3 倍的整体加速。尤其是在 8 卡小集群上，kernel 启动开销占比比想象中大得多，CUDA Graph 的收益非常明显。这一轮做完之后，profiler 里能看到的主要时间已经不再是显存拷贝，而是真正在算 attention。
4. 实操实录：从 33 秒到 2.6 秒的分阶段记录
4.1 阶段一：基线，33.1 秒
直接用常见推理框架的默认配置，TP=8 跑全量 attention，KV 保持 FP16。512K 输入进来之后，显存峰值逼近单卡 96GB，出现轻微 swap。首 token 时间在 33.1 秒左右，而且服务在 decode 阶段还不稳定，batch 稍大就 OOM。这个阶段的结论很明确：先把显存水位降下来，是后面所有优化的大前提。
4.2 阶段二：FP8 权重 + KV 4bit，26.9 秒
权重换成 FP8，KV Cache 改 4bit 并按 per-head 量化。显存峰值从顶满降到 82% 左右，不再 swap。但这一轮 TTFT 只从 33.1 秒降到 26.9 秒，提升约 1.23 倍，因为计算量没有变，只是内存压力变小了。真正的价值是给后面的调度方案留出了操作空间，同时验证了量化掉点可控。
4.3 阶段三：PD 分离 + prefill 专用 6 卡，15.7 秒
把 8 卡分成 6 prefill + 2 decode。prefill 阶段不再被 decode 的小步长拖累，同时把 batch 内的长请求优先调度。TTFT 降到 15.7 秒左右，提升约 1.71 倍。这里有个反直觉的现象：decode 只分到 2 卡后，单 token 延迟反而降低了，因为混合调度时的等待消失了，每个阶段都在按自己的节奏跑。
4.4 阶段四：序列并行 + 稀疏注意力 + 锚点 token，8.2 秒
序列并行切成 4 段，每卡处理 128K，同时引入 16K 窗口加 128 锚点的稀疏 attention。TTFT 从 15.7 秒降到 8.2 秒，提升约 1.91 倍。这一时期开始重点盯长文 benchmark 的掉点，逐层检查稀疏 mask 与训练位置的匹配度。如果掉点超过预期，问题多半不是 attention 本身，而是切分方式破坏了某些全局依赖。
4.5 阶段五：RoPE 改造 + 内核融合 + CUDA Graph，5.1 秒
RoPE 分频改造后，512K 的长文检索正确率才真正稳下来；同时融合 kernel、静态 shape 桶化、CUDA Graph 回放，TTFT 降到 5.1 秒，提升约 1.61 倍。这个阶段工作量最大，但单项收益都相对小，属于"细节堆出来的胜利"。过程中反复调了位置编码频率的缩放参数，最终是拿一组真实业务 prompt 做基准测出来的。
4.6 阶段六：投机式 Prefill + 动态长度调度 + 通信重叠，2.6 秒
最后把 prefill 的请求按 chunk 并行处理，decode 侧做环形 KV 读取，配合通信重叠后，TTFT 稳定在 2.6 秒左右。相比初版 33.1 秒，提升约 12.7 倍，正好摸到了业务目标。加了双份 batch 压测时仍能稳定在 2.8 秒以内，没有再出现显存峰值突刺。
以下是各阶段累计收益的汇总：
| 阶段 | 核心改动 | TTFT | 相对上一阶段提升 | 累计倍数 |
|---|---|---|---|---|
| 基线 | TP=8，全量 attention，KV FP16 | 33.1s | - | 1x |
| 阶段二 | FP8 权重 + KV 4bit | 26.9s | 1.23x | 1.23x |
| 阶段三 | PD 分离（6 prefill + 2 decode） | 15.7s | 1.71x | 2.11x |
| 阶段四 | 序列并行 + 稀疏注意力 | 8.2s | 1.91x | 4.04x |
| 阶段五 | RoPE 改造 + kernel 融合 + CUDA Graph | 5.1s | 1.61x | 6.49x |
| 阶段六 | 投机 prefill + 动态调度 + 通信重叠 | 2.6s | 1.96x | 12.73x |
数字看起来比较整齐，其实每个阶段都反复调试了很多轮。要特别提醒的是：累计倍数不是简单的乘法叠加，后面阶段的优化之所以还能有近 2 倍收益，是因为前几阶段把显存水位和调度灵活性都腾出来了。如果一开始就直奔稀疏 attention，前面几个步骤的收益会互相打架，工程上很难收敛。
5. 踩坑实录：过程中最想删掉的 7 个问题
5.1 RoPE 外扩后长文一致性直接崩盘
第一次把位置频率简单调低后，短文本 benchmark 完全正常，一上 512K 长文，文档内引用全乱套。原因就是高频维度开始混淆。后面改成"局部高频 + 全局低频"的双路方案，才真正稳住。这个坑的教训是：长上下文基准测试必须用带"文档内精确引用"的数据集，只看困惑度会被骗。
5.2 KV Cache 4bit 量化掉点 0.8，怎么修回来的
直接用 tensor-level min-max 量化，评测掉点接近 0.8 个点，对业务影响很大。后来发现 MLA 的 latent 向量里有少量特征维度数值特别大，不按比例压缩会毁掉整组向量。解法是把 outlier 维度挑出来单独存 FP16，其余维度继续 4bit。实测掉点降到 0.15 以内，显存增加幅度不到 0.5GB。
5.3 CUDA Graph 捕获失败：动态 shape 的锅
prefill 阶段因为请求长度不定，第一版动态 shape 根本捕获不出来。后来把所有可能的 shape 分桶（16K、32K、64K 等），桶内做静态 padding，超出桶的就走 fallback 路径。代价是浪费一小部分算力，但换来的是 Graph 回放稳定。排查这类问题的经验：先确认是不是 shape 变化触发的捕获失败，用固定 shape 单独测一遍就能快速定位。
5.4 环形通信热点：all-gather 卡在最长段
序列并行做环形 attention 时，通信量由最长的那一段决定。我们的 prompt 里有几个特别长的 chunk，导致 all-gather 时间被拉满。把 chunk 长度做截断合并后，通信热点才消失。这个问题用性能分析器看通信时间占比就能发现，超过 30% 基本就是负载不均衡，不用怀疑别的。
5.5 Prefill 峰值显存突刺：小幅 OOM 反复出现
显存峰值在阶段二之后仍然偶发 OOM，后来定位到是临时激活张量在 batch 合并瞬间暴增。解法是给 prefill 的 batch 加显存预算检查，超过单卡安全水位就自动切分请求。这个小检查只花了半天时间，但之后连续两周压测再没出过 OOM。
5.6 工具链版本不一致：半天的"灵异"复现
有一段时间同样的配置在不同机器上结果差很多。排查一圈，发现是 tokenizer 和权重切分逻辑的版本不一致导致的。通用解法：所有环境统一锁定 engine 版本，配置文件里加版本 hash 检测，不一致直接启动失败。这算是最没有技术含量但最容易消耗团队时间的坑。
5.7 长上下文评测集不够"长"
最初评测集最长只有 200K，所有优化看着都正常。等到 512K 实测，稀疏 attention 的锚点位置不够深，部分长距离依赖漏掉。后面自己构造了一份带"跨段引用"的 480K-512K 评测集，专门压锚点覆盖。给所有做长上下文的朋友一个建议：评测集长度必须和产品实际长度对齐，否则你会被自己的乐观数据骗过去。
6. 如果重来一次：优化顺序与工具链建议
按这次实战经验，做同类项目时不要一上来就堆高级技巧，先用性能分析器找到瓶颈分布。我们的实际顺序是：显存水位 → 调度结构（PD 分离） → 计算压缩（稀疏 attention） → 数值稳定性（RoPE） → 内核级压榨。这个顺序的好处是每一步都有明确的量化收益，而且不会因为前一阶段的遗留问题让后一阶段优化失效。
工具链方面，我习惯盯住三个指标：显存峰值、kernel 时间占比、通信时间占比。只要这三项数据是清晰可见的，绝大多数优化方向都能被数据逼出来。很多团队拿着别人的"最佳实践"直接抄，却没先看自己的瓶颈分布，结果收益率很低。
我个人的体会是：TTFT 这种指标像接力赛，每个环节快 20%，最后就是几倍差距，但前提是先把交接棒握稳——也就是显存、调度和数值稳定性这些基础盘。如果后续还要继续压首字延迟，我会优先考虑把 decode 侧的模型并行度再调一调，再看能不能把 decode 的 KV 读取进一步流水式处理，而不是一次性拉全量。这次项目给我的最终教训是：长上下文优化的每一点收益都来自对整条链路的精确理解，少做任何一门功课，最终的 12.7 倍都可能只是纸面数字。
