第 7 章　decode 固定成本压缩 ​

decode 的瓶颈不是一个，而是一层叠着一层。每压掉一层，下面那层才会浮出来。

7.1　问题：非 duck 那段时间去哪里了 ​

先算清楚 ​

2026-03-12 前后，在调试 Qwen3-30B-A3B-FP8 的 MoE decode 时，发现一个值得单独追查的现象：即便 MoE duck forward 本体已经明显变快，端到端的 decode TPS 并没有按比例提升。

做了一批粗估后，数字很清楚：

  • 30B tp=12, multiple_of=64：decode 约 77ms/token，duck 侧单层 judge_wait 约 0.995ms，48 层合计约 47.8ms
  • 剩下的约 29ms/token 不在 duck 里

这约 29ms 包括：

  1. attention：老的实现在 decode 热路径里用 torch.cat 追加 KV cache，每层每 token 都要重新 materialize 一次 KV 矩阵
  2. 每层 duck 前后的 D2H / H2D：hidden_states.to("cpu") 进 duck，result.to("cuda") 回来
  3. MoE gate/topk（仍在 GPU 上）
  4. final norm、lm_head、argmax

其中 attention 的 torch.cat 是最直接可优化的一项——它会随着上下文长度增长而越来越重，而且是在纯 eager Python 里发生的，每一层每一 token 都走一次。

7.2　第一步：flash attention 2 + 静态 KV cache ​

静态 KV cache ​

解决 torch.cat 问题的方案很直接：预分配。

改为在 decode 开始前按 max_position_embeddings 一次性分配 KV cache buffer，prompt 时 copy_ 写入前缀，decode 时写入固定槽位。这样 decode 热路径里不再有动态分配，也不再有 torch.cat 的内存重新 materialize。

这一步的主要价值不是立刻带来数字上的大提升，而是为后续 flash_attn_with_kvcache 和 CUDA Graph 提供了稳定的、形状固定的 cache 结构。

flash attention 2 ​

2026-03-14，flash_attn_2 接进 attention 路径：

  • prefill：用 flash_attn_func
  • decode：用 flash_attn_with_kvcache（接受 static cache、不再需要 seq 维度上的 cat）

这两处变化一起提交（commit 对应 attention backend / static KV），作为后续 flattened decode 的基础。

7.3　flattened decode：从模块树到分阶段 runtime ​

为什么模块树不够用 ​

flash_attn_2 和静态 KV cache 解决了 attention 内部的问题，但还没解决更大的问题：每层的 forward 仍然是模块树递归调用，Python 解释器在每一步都有 host 侧的开销。

在 mixed compute 系统里，每层 decode 的结构是：

input_layernorm  (GPU)
→ attention      (GPU)
→ residual add   (GPU)
→ post-attn norm (GPU)
→ duck MLP/MoE   (CPU ← GPU → CPU → GPU)
→ residual add   (GPU)

每一次从 GPU 进 duck、再从 duck 回 GPU，都有 D2H + H2D 的 to() 调用，加上前后的 CUDA 同步。48 层的 30B 就是 48 次这样的往返。

对于这类结构，如果想把 GPU 侧的固定开销压下来，就必须把 decode 路径从“树状模块调用”改成“分阶段的显式 runtime”——图只 capture GPU-only 的那段，duck 那段留在图外。

flattened runtime 的设计 ​

2026-03-14，用了整整一天时间重新设计并实现了 flattened decode path。

核心思路是新增专门的 runtime 类（放在 qwen3_flattened_runtime.py），不修改原有模块树，而是绕过它：

  • runtime 持有对每层子模块的直接引用
  • decode 时不再递归调用 .forward()，而是手动按阶段逐步推进
  • 每个“GPU 前缀”（pre-duck GPU 段）单独 capture 成一张 CUDA Graph
  • duck 部分维持 eager 执行，图外
  • duck 返回后，“GPU 后缀”（post-duck GPU 段）再 replay 另一张图

最终形成的层次结构：

qwen3_dense.py / qwen3_moe.py        ← 权重持有者 + 模型语义（不动）
qwen3_flattened_runtime.py           ← 通用 graph/state runtime
qwen3_dense_flattened.py             ← dense connector
qwen3_moe_flattened.py               ← MoE connector

Rust 侧也同步改成了 out-buffer 语义：dense duck MLP 和 MoE duck expert 现在可以直接把结果写进长期存活的 pinned host buffer，不再每次返回新的 pageable tensor。

CUDA Graph 粒度 ​

值得明确的是：这里的 CUDA Graph 是每层一张图，而不是整条 token path 的总图。

原因是 duck forward 横亘在每层中间，打断了 GPU 的连续执行。如果想要“一次 token 一张大图”，就必须把 duck forward 也封进图——但 duck forward 本质上是 Python → Rust → C++ DPDK 的跨层调用，最终要把数据通过网络发到 duck 节点、等结果回来，整个路径涉及 IO 和 host 侧阻塞等待，根本无法被 CUDA Graph 捕获。

每层一张图的好处已经足够明显：图内的小 kernel launch（layernorm、residual add、rotary 等）不再各自走一次 Python host launch，而是合并进一次 cudaGraphLaunch。

从 30B 的 CUDA Graph 结果可以直接验证这一点：

cudaGraphLaunch   = 480 次
cudaLaunchKernel  = 15022 次

# 480 ≈ 48层 × 10个 profiled decode step
# 每层约 480/480=1 次 graph launch，替代了多次 kernel launch

7.4　数字：flattened decode 的实际收益 ​

2026-03-14 19:14，flattened decode runtime 提交进仓库，当天下午的阶段性测量结果如下（flash_attn_2 + bf16 + cuda，multiple_of=32）：

模型tpeager decodeflattened decode提升
Qwen3-4B-FP82921 tok/s24 tok/s+14%
Qwen3-30B-A3B-FP82418 tok/s24 tok/s+33%
Qwen3-235B-A22B-FP8244.3 tok/s4.7 tok/s+9%

日志里对应的具体数字：

# 235B, eager
Generation stats: decode 4.351 tok/s, tpot 0.230s

# 235B, flattened + cudagraph
Generation stats: decode 4.563 tok/s, tpot 0.219s

对 235B 来说，flattened 把每 token 的 non-duck 时间从约 0.44ms/layer 压到约 0.20ms/layer，降幅约 55%。

30B 的 tp=12 版本有更早的对照数字：

# 30B tp=12, fa2, eager
Generation stats: decode 13.577 tok/s, tpot 0.074s

# 30B tp=12, fa2 + cudagraph
Generation stats: decode 15.964 tok/s, tpot 0.063s

改善约 11ms/token，相对降幅约 16%。

这批数字说明 flattened decode 的收益是真实的，也是非线性的——duck 时间越重的模型（235B），GPU 侧剩余时间占比越低，flattened 能额外拿到的相对收益自然也越小。

跨模型验证：Qwen3.5-35B ​

flattened decode 不是 Qwen3 专属优化。2026-03-20，在 Qwen3.5-35B-A3B-FP8 上接入 flattened decode 后，首轮实测（tp=16，每层 Rust 约 470µs，共 40 层）：

模式TPOTtok/s非 duck 时间
eager fastpath41.3 ms24.222.5 ms
flattened decode26.4 ms37.97.6 ms

TPOT 下降约 14.9ms，tok/s 提升约 56%，非 duck GPU 时间从 22.5ms 压到 7.6ms。这表明同一套 flattened runtime 和 CUDA Graph 机制，在结构差异显著的 Qwen3.5-MoE 上同样有效——并非 Qwen3 特例。

7.5　state hygiene：第二轮对话的乱码 bug ​

症状 ​

flattened decode 上线后，发现了一个典型的“第二次才坏”的 bug：

  • 第一轮对话：输出正常
  • 第二轮对话：输出乱码
  • 关掉 flattened_decode=True：不复现

根因 ​

问题出在两个层面：

第一层：graph 复用时绕过了 decode state 刷新。 flattened runtime 的 _ensure_first_graph() 等方法，如果 graph 已经存在且 shape 没变，会直接 fast-path 返回。旧代码在这条快路径上，没有重新刷新 token_position_gpu、token_cache_seqlens_gpu 和 attention 的 cache_seqlens binding。结果是：第二轮对话的 decode 仍在用第一轮留下来的 decode 位置状态。

第二层：prompt 变长触发 KV cache 扩容时，graph 绑定的 buffer 指针已经过期。prepare_generation_cache(required_total_length) 可能让 attention 层重新分配更大的 cache。但旧 graph 仍然 replay 到第一次 capture 时绑定的 buffer 地址——这才是输出彻底乱掉的直接原因。

修法 ​

修法收在三处：

  1. 在 generate() 入口统一调 reset_generation_state()，确保每次新生成从干净状态开始。
  2. graph 复用命中时，也必须先刷新 decode position 和 cache seqlens binding，再 replay。
  3. KV cache 扩容时，通知 flattened runtime 丢弃所有已 capture 的 graph，下次 decode 重新 capture。

用户验证后：第二轮对话恢复正常。

这个 bug 是 flattened runtime 独有的——普通 eager 路径的状态管理更简单，每次 generate() 都会覆盖旧状态，不会触发 graph 绑定地址过期这类问题。它提醒了一件事：graph capture 不只是性能优化，它还在不同轮对话之间引入了额外的状态依赖，必须在会话边界处显式清理。

把 flattened decode 扩展到 DeepSeek-V3.2 时，也出现了类似的问题：第一轮对话正常，第二轮 decode 从 token #1 开始持续输出 0。排查后确认原因和 Qwen3 的 state hygiene bug 同构——新会话复用了旧 graph，但 graph 内绑定的 buffer 已经失效。修法是在 reset_generation_state() 里统一调用 _invalidate_graphs()，让 first_graph、middle_graphs、final_graph 在每轮新会话开始时全部失效，下一轮 decode 重新 warmup 并 capture。代价是每次新对话多付一次 graph capture 的固定成本，但 steady-state decode 不受影响。

7.6　Triton fp8 路径：从 BF16 反量化到原位 fp8 GEMM ​

老路径的问题 ​

duck-llm 最初的 Linear 模块在加载时会检查权重是否是 FP8：如果是，且前缀不含 "mlp" 字样，就直接乘上 weight_scale_inv，把权重膨胀成 BF16 常驻显存。

这个策略在 Qwen3 时期能接受，但到 DeepSeek-V3.2 的 61 层、每层含大量 attention 投影和 indexer 权重时，显存压力变得很大——只能靠 embed_tokens、lm_head、shared_experts 临时 CPU offload 来止血。

根本原因很清楚：不是模型太大，而是不该膨胀的权重都提前膨胀成了 BF16。

Triton bf16 @ fp8 kernel ​

2026-03-20，在 duck-llm 里新增了 triton_kernels/ 子包，其中最关键的是 bf16_block_scaled_mm.py：

python
# 核心 dispatch 语义：
# activation: bf16
# weight: fp8 (e4m3), 128×128 block scale
# 输出: bf16
# → 权重保持 fp8 常驻，不在 load 阶段膨胀

通过全局环境变量 DUCK_TRITON_FP8_LINEAR=1 开启：启用后，CUDA 上且未被显式强制反量化的量化 Linear，会保留原始 FP8 权重，forward 时用 Triton kernel 做 bf16 @ fp8 → bf16 的 GEMM。

V3.2 的一个 bug：kv_b_proj 直读权重绕过了 scale ​

Triton path 开启后，DeepSeek-V3.2 出现乱码。排查过程很清楚：

DeepseekV32MLA 的 decode fast path 有一段手写 einsum 逻辑直接读取 kv_b_proj.weight，但不经过通用 Linear.forward()——也不会把 weight_scale_inv 乘回去。

开启 Triton 之后，Linear 把 kv_b_proj.weight 保留成 FP8 + 转置状态（而不是反量化 BF16）。那段 einsum 于是把原始 FP8 量化值当成稠密 BF16 权重在用，输出自然乱码。

修法是给 kv_b_proj 新增一个独立的 DeepseekV32DecodeDirectLinear 类，强制保持 BF16，不走 Triton backend。代价是约 +16 MiB/layer，全 61 层合计约 +976 MiB——但这是必要的代价，且已经在注释里明确写清楚。

显存收益 ​

修完 kv_b_proj 之后，V32 把 shared_experts 也重新放回 Triton fp8 路径（原本是临时 CPU offload）。真实日志里的显存对比：

# shared_experts 临时 BF16 offload 时
resident_total_mib ≈ 19345.0 MiB

# shared_experts 恢复 Triton fp8 后  
resident_total_mib ≈ 16909.6 MiB

# 节约约 2435 MiB

CUDA Graph 下的 Triton 开销 ​

Triton kernel 在 eager 模式下每次 launch 都要走一段 Python wrapper（取 device/stream、算 specialization key、查缓存等），即使 JIT cache 命中，也有约 5.5µs 的固定开销。

但一旦这些 Triton kernel 被 capture 进 CUDA Graph，replay 时就不再经过 Python wrapper——直接由 CUDA driver 重放 graph 节点。验证实验表明，capture 时 pre-run hook 触发 N 次，后续 5000 次 replay 时 hook 增量为 0。

这意味着 Triton + flattened decode + CUDA Graph 的组合是自洽的：Triton 的 Python-side launch 开销会被 CUDA Graph 完全吸收，不会堆叠成额外的固定成本。

V3.2 decode context 内核：split128 与设备计数器 ​

Triton fp8 Linear 解决了注意力投影和共享专家的显存问题，但 DeepSeek-V3.2 的 decode 路径上还有一颗独立的内核：decode_context——负责在逐 token 生成时拼接 KV cache、管理 active length、并更新 rotary position。

早期尝试把 cache append 和 decode context 作为两颗独立内核提交，但在真实模型上反复出现输出异常。2026-03-22 的排查把问题收敛到了两个关键点：

第一，decode_context 必须保留 fp32 中间量。 在 exact segment probe 已经证明 flattened runtime 全层对齐之后，只要把 v32_fp8_decode_context.py 的中间计算从 fp32 降回 bf16，flattened + triton 的真实聊天输出就会立刻变坏；改回 fp32 中间量后恢复正常。这说明 decode context 里的某些累积操作（如 scale 合并、位置编码合成）对精度足够敏感，不能为了省显存或带宽牺牲中间格式。当前工作版本保留了 correctness-first 的 fp32 中间量。

第二，active_len 的 off-by-one。 旧版 decode_context 吃到的 active_len 是 append 前的长度，导致 decode attention 会把刚 append 的当前 token 排除在 context 外，表现为“降智 + 复读”。修法是让 decode-context 单独吃 cached_len_device + 1，而 append offset 和 rotary position 继续保持旧长度语义。

split128 + device-owned counter 定稿 ​

2026-03-21，split128 版本的 decode context 内核正式提交。这条路线做了两件事：

  1. 把 decode context 确定为主线上唯一的 split128 路径，不再保留多套 kernel 变体。
  2. 把 cache append 的 row_offset_tensor 和 decode context 的 active_len_tensor 接到同一个 device-owned counter 上，避免 host 侧反复构造和传输这两个张量。

真实模型结果：V32 clean_fp8_full_attn_ref + Triton split128 decode context 的 TPOT 约为 161 ms。这条线可以视为当前里程碑已完成。

7.7　整条演化线的阶段性总结 ​

把这几步优化放在一张时间线上。第一张表记录的是 Qwen 模型上能复用的 decode 优化成果：

日期变化235B tp24 decode30B tp24 decode
2026-03-12eager，旧 attention~4.253 tok/s~18 tok/s
2026-03-14flash_attn_2 + static KV cache中间态中间态
2026-03-14flattened decode + CUDA Graph4.7 tok/s24 tok/s
2026-03-16（MoE SMP kernel，见第 8 章）fp9 MoE 多核~9 tok/s—

flattened decode 这一步本身的收益在 235B 上看起来不大（4.3 → 4.7），但它的意义不只是这个数字——它把 non-duck GPU 时间从约 0.44ms/layer 压到约 0.20ms/layer，让后续优化的重心可以完全回到 duck 本身。后续 fp9 MoE 多核把 235B 拉到约 9 tok/s，正是在这个 flattened 基础上发生的。

第二张表记录的是 V3.2 专属的 Triton 路径和 decode context 工作线，这些优化不在 Qwen 模型上发生：

日期变化V3.2 decode
2026-03-20Triton fp8 Linear（显存优化，CPU offload 止血）可跑通
2026-03-21split128 decode context + device counter~6.2 tok/s（TPOT ≈ 161 ms）
2026-03-24act_quant 收尾 + fp8 dispatch~9.9 tok/s

Triton fp8 路径的价值主要体现在显存侧，而不是速度侧（在未 graph 化的 eager 模式下，Triton launch 有额外固定成本）。但它是 V3.2 能在 32GB 显存上稳定运行的关键：没有它，embed_tokens、lm_head、shared_experts 的 CPU offload 只能算临时止血，不是长期方案。

Triton decode context 内核则是另一条独立的工作线。它不解决显存问题，而是解决 decode 路径上每层都必须执行的 KV cache 拼接和 active length 管理。这条线最终确定为 split128 单一内核变体 + device-owned counter，并且确认了一个关键约束：中间计算必须保留 fp32，否则真实输出会损坏。

项目声明 / Project Disclaimer

本项目为作者以个人身份、利用业余时间推进的个人娱乐项目；除非另有明确说明，它与作者的任何雇主、客户、学校、单位或其他组织均无合作、雇佣、委托、赞助或背书关系，也不代表任何该等主体的立场。除普通个人捐赠外，本项目未获得任何资金支持。

This project is a personal hobby project developed by the author in a personal capacity and in personal time. Unless explicitly stated otherwise, it has no collaboration, employment, commission, sponsorship, endorsement, or institutional affiliation with any employer, client, school, partner organization, or other entity, and it does not represent any such party's views. No funding was received for this project except ordinary personal donations.

许可 / License

除非另有说明，本页原创文字、本站原创图片与本站原创图表采用 CC BY-NC-ND 4.0 发布。
转载时请保留原文标题、署名“JudgeDuck AI”、发布日期与原始链接；禁止商业转载、改写、摘编、翻译或基于原文创作演绎作品。

第三方商标、外部链接内容，以及文中另有标注的材料，不在上述许可范围内。

Unless otherwise noted, the original text, original images, and original figures on this page are licensed under the Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International License.