Qwen3.8-Flash-Next on Jetson AGX Thor：FP8 KV Cache 移植记录
日期：2026-09-22（跨 9/23 凌晨）
机器：懒猫 AIPod X5-T5000 = Jetson AGX Thor T5000（sm_110，122.86 GiB 统一内存，L4T R39.2 / JetPack 7.2，CUDA 13.2）
引擎：NVIDIA Jetson 专用 vLLM 构建 0.28.1rc1.dev580+g385dce36b（包路径 vllm/models/qwen4_exp/）
模型：Qwen3.8-Flash-Next 180B MoE（NVFP4 · 消审查版）· 36 层 GDN 线性注意力 + 12 层 QSA 稀疏注意力 · 内嵌 MTP
镜像：vllm-thor-plemmap:latest（未重建镜像；所有改动以 只读挂载挂入）
1. 做了什么
目标：让这台机器上的 vLLM 支持 --kv-cache-dtype fp8_e4m3，把 KV 池做大。
镜像里 fp8 KV 本来被五道门拒绝，且拒绝理由各不相同。逐道拆：
补丁 1 · vllm/models/qwen4_exp/nvidia/qsa.py（6 处）
| # | 原代码 | 改动 |
|---|---|---|
| 1 | supported_kv_cache_dtypes = ["auto","bfloat16"] |
加入 "fp8", "fp8_e4m3" |
| 2 | if self.kv_cache_dtype not in ("auto","bfloat16"): raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache") |
放行 fp8 |
| 3 | if cache_config.cache_dtype not in ("auto","bfloat16"): raise NotImplementedError(同上) |
放行 fp8 |
| 4 | if self.kv_cache_torch_dtype != torch.bfloat16: raise NotImplementedError("requires BF16 cache storage") |
放行 torch.float8_e4m3fn 与 torch.uint8（vLLM 的 fp8 缓存以 uint8 字节分配；拒 uint8 等于拒掉内核唯一会拿到的存储） |
| 5 | 主 KV cache 送进内核前 | key_cache.view(torch.float8_e4m3fn) / value_cache.view(...)：位重解释（不是数值转换） |
| 6 | 内核调用 | 传入 k_scale / v_scale / kv_quant_mode |
补丁 2 · vllm/models/qwen4_exp/nvidia/ops/qsa.py（9 处）
- Triton 内核签名加
k_scale/v_scale/KV_QUANT_MODE: tl.constexpr（不开 fp8 时编译期惰性，零性能影响） - fp8 tile 提升为 bf16 再做点积（不整块反量化 → 这是消掉早期 −6% 的那版做法）
scores *= k_scale（乘在 score 上）、normalized_output *= v_scale（乘在累加器上）- 算子入口断言放行 fp8；封装 / launch / warmup 三处签名同步
- 量化时
block_n = max(16, block_n // 2)：避免_cast_kv_tile物化 fp32 tile 后 shared memory 翻倍超限
补丁 3 · vllm/v1/attention/backends/fa_utils.py（1 处）
原： return (fa_version == 3 and family(90)) or (fa_version == 4 and family(100))
改： 追加 or (fa_version in (3,4) and family(110))
依据：vLLM 官方文档的注意力后端表写着 FA4 的 fp8 KV 支持 compute capability ≥ 10.0； 镜像里的判定却只认 major 9 / major 10，注释自承 “Blackwell (SM100+, restrict to SM100 for now)”。 即：这是上游的暂时性限制，sm_110 本应在内。
补丁 4 · vllm/v1/attention/backends/flash_attn.py（1 处）
_fp8kv_thor = current_platform.is_device_capability_family(110)
if is_quantized_kv_cache(self.kv_cache_dtype) and not _fp8kv_thor and not flash_attn_supports_kv_cache_dtype(...):
raise NotImplementedError("FlashAttention does not support ... kv-cache on this device.")
原因：本模型的 QSA 后端类名是 Qwen4ExpQSAFlashAttentionBackend —— 它继承自 FlashAttention 后端，
于是白继承了 FA 的设备门。但真正执行 KV 读写的是 QSA 自己的 Triton 内核，不经过 FA，这道门属于误伤。
旁证：参考实现（DGX Spark / GB10）的镜像里根本没有 flash_attn 模块，它的 QSA 是独立 Triton 后端，因此不存在这道门。
启动参数
--kv-cache-dtype fp8_e4m3
--attention-config.flash_attn_version=4
--long-prefill-token-threshold 1024
（`--enable-prefix-caching` 亦于同日打开）
挂载（只读、可秒级回退）：
fp8kv/patch/qsa.nvidia.py → .../vllm/models/qwen4_exp/nvidia/qsa.py
fp8kv/patch/qsa.ops.py → .../vllm/models/qwen4_exp/nvidia/ops/qsa.py
fp8kv/patch/fa_utils.py → .../vllm/v1/attention/backends/fa_utils.py
fp8kv/patch/flash_attn.py → .../vllm/v1/attention/backends/flash_attn.py
一键回退：FP8KV=0 KV_DTYPE=auto bash start_seqs.sh
2. 跑分（全部实测 · 同机同镜像）
KV 池（--gpu-memory-utilization 0.70，同模型同参数）
| KV dtype | GPU KV cache size | 可用 KV 显存 |
|---|---|---|
| bf16 | 496,770 token | 14.04 GiB |
| fp8_e4m3 | 833,617 ~ 872,939 token（多次启动有摆动，profiler headroom 受页缓存影响） | 13.25 ~ 13.69 GiB |
| 增益 | ×1.68 ~ 1.76 |
prefill（sparkDash 协议复刻 · 同 sizes · 每发新随机长度）
| ctx | bf16 | fp8_e4m3 | 差 |
|---|---|---|---|
| 16.4k | 1833 | 1690 | −7.8% |
| 32.8k | 1863 | 1708 | −8.3% |
| 65.6k | 1867 | 1713 | −8.3% |
| 131k | 1863 | 1703 | −8.6% |
| 250k | 1836 | 1674 | −8.8% |
decode · 短上下文（thorbench · 与 bf16 基线同 flags）
| prompt | 并发 | bf16 | fp8 | 差 |
|---|---|---|---|---|
| prose | c8 | 174.6 | 169.0 | −3.2% |
| code | c8 | 235.2 | 210.3 | −10.6% |
| json | c8 | 242.3 | 251.5 | +3.8% |
decode · 长上下文（单发 · ITL 增量口径）
| ctx | bf16 ms/step | fp8 ms/step | 差 |
|---|---|---|---|
| 82k | 43.77 | 44.73 | −2.2% |
| 156k | 44.29 | 45.50 | −2.7% |
8 路并发 × ~16k 上下文（20 秒纯解码窗口 · 等全部 prefill 完成后计时）
| 配置 | 聚合 | 每流中位 | 快慢比 | 最长 token 间隔 |
|---|---|---|---|---|
| fp8 KV · 8 路 | 39.2 tok/s | 4.8 | 2.2× | 0.11~0.21 s |
| bf16 KV · 8 路 | 36.0 tok/s | 4.7 | 1.9× | 0.12~0.22 s |
| bf16 KV · 4 路 | 19.6 tok/s | 5.0 | 1.3× | 0.08 s |
每流速度与流数几乎无关（4 路 5.0 / 8 路 4.7 tok/s）⇒ 每步成本 ≈ 0.20 s，只随上下文变长而变贵。
数值（8 道 GPQA Diamond 固定题 · temp0 seed42 · 关思考）
- 答案字母：fp8 与 bf16 8/8 完全一致
- 自由推理文本逐字一致：2/8（其余为措辞差异；「empirical parity, not bitwise identity」）
- 批量（3~6 并发）下的旧记录与单发记录答案有偏移 ⇒ fp8 在批量下有质量漂移，需要按场景取舍
3. 同日顺带完成的四项改动
| 项目 | 改动 | 实测 |
|---|---|---|
| PLE 抓行 | 单线程 → 线程池（24 workers · 2048 行/任务 · 小批量也进池 · 抓前去重+排序） | 每行 277 → 10 µs（社区同口径最优约 32 µs/页，我们更好） |
| 写回 | out.copy_(from_numpy(blob)) → 固定内存 staging + 非阻塞 H2D |
H2D 段降至 ~0.2 ms/次 |
| 计时 | 拆成 D2H / gather / H2D 三段 | 定位依据 |
| 前缀缓存 | 打开 --enable-prefix-caching（本修订已含上游 block-size 修复：_block_size_resolved=True, block_size=3184） |
12k 前缀重复请求 TTFT 8.18 → 3.87 s |
| 长 prefill | --long-prefill-token-threshold 1024 |
解码流被 prefill 饿死的现象显著缓解（失败题 169 → 127） |
盘层基准（4 KB 随机读 · O_DIRECT · 铺满 71 GiB 模型目录）
| 队列深度 | 平均延迟 | 聚合吞吐 | IOPS |
|---|---|---|---|
| QD1 | 143 µs | 26.9 MB/s | 6,558 |
| QD16 | 307 µs | 94 MB/s | 23,026 |
| QD64 | 600 µs | 89 MB/s | 21,709 |
同口径复刻参考实现探针（100 次随机 64 页 · 14 线程）：我们 2.78 ms/64 行（43.5 µs/页） vs 参考机 2.06 ms（32.2 µs/页） ⇒ 仅 1.4×，盘不是瓶颈。
4. 结论
- fp8 KV 在这台机器上可用：池子 ×1.7，数值在单发场景与 bf16 一致，批量下有漂移。
- 它买的是容量，不是速度：prefill 稳定 −8%，decode −2~11%。原因是 QSA 是稀疏注意力——每步只读被选中的一小块 KV，KV 压成 fp8 省不到带宽，只多付反量化。
- 真正限制多 agent 并发的是"长上下文下的每步注意力/索引开销"，不是 PLE、不是磁盘、不是 KV 精度、不是调度：每步 ~0.20 s 且与流数无关（每流恒定 ~5 tok/s），而短上下文 c8 可达 169~192 tok/s。
- 五道门里，四道是过期的保守检查，一道（FA 设备矩阵）是上游的暂时性架构限制；没有一处需要改内核实现本身。
短上下文聚合的"200+ 成立条件"（c2/c4/c8 × 三种 prompt · 384 token 生成）
| prompt | 每步 token（MTP 接受) | c2 | c4 | c8 |
|---|---|---|---|---|
| json | 2.79 | 95.5 | 141.3 | 230.2 |
| code | 2.50 | 87.4 | 141.5 | 231.7 |
| prose | 2.14 | 77.8 | 114.7 | 170.6 |
聚合 ≈ 并发 × 每步 token ÷ 每步耗时。每步 token 由 prompt 类型决定（json/code 的结构化输出接受率高）， 每步耗时由上下文长度决定（短 ctx ≈ 85 ms；16k ctx ≈ 200 ms，约 40 倍）。 ⇒ "200+ tok/s 聚合"成立的条件是：短上下文 + code/json + c8，三者缺一不可。
★ MTP 在"多路 + 长上下文"下是净亏（当晚最大发现）
| 配置（8 路 × 16k · 20 秒纯解码窗口 · 同旋钮 bf16） | 聚合 | 每流中位 | 快慢比 | 无饿死 |
|---|---|---|---|---|
| MTP=2 | 36.0 / 35.9 tok/s | 4.7 | 1.9~2.2× | ✓ |
| MTP 关 | 68.3 tok/s | 8.3 | 1.3× | ✓ |
关掉投机解码后吞吐 +90%。机制：16k 上下文 + 8 路时草稿接受率只有 39.9% （短上下文基线 json 2.79 / code 2.50 / prose 2.14 tok/step ⇒ 约 57%）， 每步多一轮 draft+verify 的固定延迟却只换来 +0.8 个 token ⇒ 净亏，且这笔固定成本 与上下文长度、与流数都无关——这正是前面"16k→32k 每流只掉 16%""4 路≈8 路"两条异常的来源。
操作结论：
- 多路 / 长上下文（≥16k）⇒ 关 MTP（MTP_OFF=1，收益 +90%，KV 池还多 24%：497k → 617,593）
- 单流 / 短上下文 ⇒ 保留 MTP（c8 短 ctx 169~192 tok/s 是开着 MTP 测的）
- MTP=0 会被 pydantic 校验拒绝（要求 >0）⇒ 必须整段去掉 --speculative-config
★ MTP 的 K × 上下文长度完整对照（当晚补齐三档）
| 场景（同机同镜像 · 20 秒纯解码窗口/thorbench） | K=2 | K=1 | K=0 |
|---|---|---|---|
| 短 ctx c1 json（384 tok 生成） | — | 49.4 | 29.2 |
| 短 ctx c2 json | 95.5 | 80.9 | 52.3 |
| 短 ctx c4 json | 141.3 | 127.0 | 90.5 |
| 短 ctx c4 prose | 114.7 | 113.2 | 87.5 |
| 短 ctx c4 code | 141.5 | 130.0 | 96.3 |
| 4 × 16k 聚合 | 19.6 | 24.6 | 37.2 |
| 4 × 16k 每流中位 | 5.0 | 6.1 | 9.5 |
| 4 × 16k 快慢比 / 饿死 | 1.3× / 无 | 1.2× / 无 | 1.2× / 无 |
| 8 × 16k 聚合 | 36.0 | — | 68.3 |
| KV 池（GMU 0.70） | ~497k | ~483k | 642,771 |
机制（spec-decode 计数器实测，不是推测） - K=1 时单个草稿的命中率：prose 0.70 / code 0.86 / json 0.93（每步接受数） - K=2 相对 K=1 只多 0.44 / 0.64 / 0.86 个 token/步 ⇒ 第二个草稿命中率腰斩 - 上下文越长 ⇒ 每步 draft+verify 的固定成本越高，而收益（多 0.4 token）不变 ⇒ 在 16k 处彻底翻转
可落地分档
- 长 ctx（≥8k 提示，agent/coding 场景）⇒ MTP=0：4 路 37.2（= K=2 的 1.9×）· 8 路 68.3（同 1.9×）
- 短 ctx（几百 token 的短问）⇒ MTP=2：c4 json 141.3（K=0 只有 90.5，关掉反而掉 36%）
- 引擎侧限制：num_speculative_tokens_per_batch_size 的定义是 (range_start, range_end, K) 闭区间、
按当前 batch size（= 并发度）每步查表（speculative.py:479-484 + scheduler.py:275），
看不到上下文长度 ⇒ 想按 ctx 分档只能做两套 profile，不能靠一张调度表
- 教训：早期试动态 K 时手写的那版 [[1,2,3],[3,5,2],[6,8,1]] 是凭直觉编的 ✗，
实测形状完全不是那样（长 ctx 下连 4 路都该 K=0）⇒ 调度表必须用实测填
★ 用户客户端实测（Windows bench_gui.py 手动测试 · 2026-09-23 当天档位）
| 场景 | 聚合 prefill | 聚合 decode |
|---|---|---|
| 单流 · 200k ctx | — | 50 tok/s |
| 双流 · 100k ctx | — | 70 tok/s |
| 3 流 · 100k ctx | — | 60 tok/s |
| 4 流 · 16k ctx | — | 73 tok/s |
| 6 流 · 32k ctx | — | 90 tok/s |
| 8 流 · 32k ctx | — | 100 tok/s |
| 各档聚合 prefill | ~1600 tok/s | — |
口径说明（重要，别和上面的表混着比）
- 本表用 bench_gui.py（tkinter 并发流式观察窗）手动逐档跑：并排 N 路流式窗口，肉眼确认是否同时出首 token、哪个流在卡；聚合值取整轮墙钟口径。
- 本表是客户端整轮聚合口径：一轮请求从发起到结束的全部生成 token ÷ 墙钟时间，含爬坡段。
- 上面那些「20 秒纯解码窗口」是稳态 + 每流口径（例：4×16k K=0 = 聚合 37.2 / 每流 9.5）。
- 两套数字不能直接相减或换算；但方向性结论互相印证：
- 聚合 prefill ~1600 tok/s 与本文单发 sparkDash 测得的 1674~1713 tok/s 同量级 ✓
- 「8 路难用」说的是单流交互延迟✗，不是聚合吞吐 ✗ —— 8 路 32k 聚合 100 tok/s ⇒ 每流约 12.5 tok/s，是能用的。
- 单流 200k ctx 仍有 50 tok/s ⇒ 长上下文单路体验明显好于 20 秒窗口口径给人的印象。
- 多路（6~8 流）能跑到 90~100 tok/s ⇒ MTP 关掉的长 ctx 档在多路下是对的（投机在这种负载下净亏，见上一节）。
5. 复现
# 起服务（FP8 默认档）
SEQS=8 GPU_MEM=0.70 MTP=2 PORT=8001 BATCHED=4096 LPT=1024 bash ~/mtpvocab/start_seqs.sh
# 回退 bf16
FP8KV=0 KV_DTYPE=auto SEQS=8 GPU_MEM=0.70 MTP=2 PORT=8001 BATCHED=4096 bash ~/mtpvocab/start_seqs.sh
补丁生成器：~/mtpvocab/fp8kv/patch_fp8kv.py（锚点唯一命中校验 + AST 校验，任一不命中即整体不写）
备份与回滚：~/mtpvocab/fp8kv/backup/（镜像 ID、两文件原版、sha256）
附件
「用户客户端实测」一节用的工具：bench_gui.py —— tkinter 并发流式观察窗，无额外依赖（tkinter 标准库 + requests）。
⬇️ 下载 bench_gui.py（26 KB）
python bench_gui.py --url http://<host>:8001/v1 --model qwen3.8-flash-next \
--ctx 32768 --concurrency 8
| 参数 | 默认 | 说明 |
|---|---|---|
--url |
http://localhost:1234/v1 |
OpenAI 兼容端点 |
--model |
local-model |
模型名 |
--api-key |
空 | 需要认证时填 |
--ctx |
4096 | 每路 prompt 目标 token 数（点「随机」按此生成，每路内容不同 → 避开 prefix cache） |
--max-tokens |
256 | 每路生成上限 |
--concurrency |
2 | 并发路数 |
测量特点：
- 精确优先：响应带
usage就用服务端真值算 TTFT / prefill / decode；没有才走估算 - 自动校准：用服务端真值对估算系数（字符/token、每块 token）做滑动平均，越跑越准
- 卡顿检测：任意路超过 0.5 s 没有新 token 就标记，并统计最大间隔
- 看门狗兜底：停止后 2 s 强制恢复按钮，防线程未收尾卡死