Merge the drift investigation: quantified as quantization noise, plus a cross-platform import fix

Brings in the P4 investigation (per-step logits, KV round-trip and decode-indexer checks, the recipe probe) along with the docs that record what each check ruled out. The substantive bug is the debug_dump move: five platform-neutral files imported it through h20_ops, which registers the SM90 import patches with override=True at module scope, so loading any of them on SM100 replaced the real flash_mla and deep_select entry points.
2026-10-31 00:36:20 +08:00 · 2026-10-21 22:37:56 +08:00 · 2026-10-21 22:37:56 +08:00 · 6167cf9f74
commit 6167cf9f74
parent dcf8d11570 74ba468e76
14 changed files with 967 additions and 132 deletions
--- a/H20-DEBUG-PLAN.md
+++ b/H20-DEBUG-PLAN.md
@ -152,20 +152,249 @@ x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim=
 > 后续跟踪（P3.6 起）已迁移到 `H20-PORTING-JOURNAL.md` 第 6 节，那里有完整的
 > 漂移量化数据、建议顺序和前置条件。本文件的 P4 / P5 保留为方法清单。

-## P4 decode 专项（仅当 prefill 对、decode 错）
+## P4 decode 专项 —— 完成：`**` 已定量解释为量化噪声，非缺陷

- [ ] P4.1 KV cache 写入后立刻读回，比 bf16
- [ ] P4.2 同一位置，decode 路径 vs prefill 路径重算的 logits 差异
- [ ] P4.3 `_decode_p_sum_sm90` 的数值 vs prefill 的 `fp8_fp4_mqa_logits`
+### 为什么锁定 decode

---
+P3 之后重跑对拍（`H20-DRIFT-ANALYSIS.md` 有完整数据）：

-## P5 logprob 异常
+| | top-1 | margin(6127 北京 − 666 `**`) |
+|---|---|---|
+| 官方，14 token prefill | 6127 | **+3.112** |
+| 我方，14 token prefill | **6127** ✅ | **+0.256** |
+
+prefill 选中了正确的 token，误差只把 margin 压缩了 12 倍。但端到端跑
+（11 token 起，增量 decode 到第 14 个位置）给的是 666。
+
+**同一个 14-token 上下文，一次 prefill 给 6127，增量 decode 给 666。**
+
+当时的推断是「剩余漂移在 decode 路径」。P4.2~P4.4 做完后，这个推断需要精确化为：
+**差额确实只在 decode 出现，但原因不是 decode 有 bug，而是 decode 经过量化 KV cache
+往返、噪声略增，恰好碰上模型自己在那个位置只有 1.64 的 margin。**
+
+### 判据（先定死，避免事后解释；P4.2 之后已修正）
+
+> **事后修正**：下面这套判据假设「decode 与 prefill 应当基本一致」，因此把
+> `rel(A,B) > 3e-2` 判为「decode 自身缺陷」。P4.2 证明这个前提错了 —— 官方自己
+> prefill→decode 也让 margin 变化约 2 倍，因为 KV cache 走量化往返。
+> 真正有效的判据是**Δmargin 的量级 vs 该位置的 margin**，而不是 rel_err 的绝对值。
+> 判据保留在下面作为过程记录，最终结论见「最终结论」一节。
+
+
+设第 14 个位置上：
+
+- `A` = 我方 prefill 路径的 logits
+- `B` = 我方 decode 路径的 logits
+- `R` = 官方 decode 路径的 logits
+
+| 量 | 含义 | 阈值 |
+|---|---|---|
+| `rel(A, B)` | 我方两条路径自身是否一致 | > 3e-2 即为 decode 自身缺陷 |
+| `rel(A, R)` | prefill 残余误差 | 已知约 3e-2（量化地板） |
+| `rel(B, R)` | decode 总误差 | 明显大于 `rel(A,R)` 说明 decode 额外引入误差 |
+
+- 若 `rel(A,B)` 很小而 `rel(B,R)` 大 → decode 没有额外缺陷，是量化误差累积，`**` 属可接受
+- 若 `rel(A,B)` 大 → **decode 有 bug**，按 P4.3 → P4.4 往下查
+
+### 步骤
+
+- [x] **P4.1 同位置 prefill vs decode 对拍** —— **完成（结论已被 P4.2 修正，见下）**
+
+      四条路径走到同一个 14-token 上下文
+      `[0,128803,2056,768,17607,2793,3139,12108,1148,128804,128822,17607,2793,3139]`：
+
+      | 路径 | 下一个 token | 结果 |
+      |---|---|---|
+      | 官方 prefill（一次喂 14 个） | 6127 = 北京 | ✅ |
+      | 官方 decode（喂 11 + 逐 token forward） | 6127 = 北京 | ✅ |
+      | **我方 prefill（一次喂 14 个）** | **6127 = 北京** | ✅ |
+      | **我方 decode（喂 11 + 增量）** | **666 = `**`** | ❌ |
+
+      前两条来自 `ref_next_token.py`（prefill 一次 / `--force_ids` 增量），后两条来自
+      `h20_ask.py`：`--ids <14个> --max_tokens 1` 给 `6127`，`--ids <11个> --max_tokens 7`
+      的第 4 个 token 是 `666`。
+
+      **结论：prefill 路径正确，差额只由 decode 引入。**（注意：P4.3/P4.4 之后已知
+      decode 的 KV cache 与 indexer 都精确，所以「差额由 decode 引入」不等于「decode 有
+      离散 bug」——见下方「量化后的真实图景」。）
+      这把 P3 之后「margin 被压缩到 0.256，decode 再推一把就翻」的推断变成了实证。
+
+      补充：官方自己 prefill 与 decode 的 margin 也不同（3.112 vs 1.638），所以
+      「两条路径不完全相等」是框架固有特性（KV cache 量化往返），但**方向不该翻**。
+      - 路径 1：直接 14 token prefill（`--ids <14个>`，dump `logits_prompt`）—— 已有
+      - 路径 2：11 token prefill + 3 步贪心 decode（贪心本来就给对了
+        17607/2793/3139，不需要强制喂 token），dump 第 3 步 decode 的 logits
+      - 参考侧同样跑增量：`ref_next_token.py` 逐步喂 11→12→13→14 个 token，
+        每一步的记录就是官方 decode 路径的 logits
+      产出：三条曲线的 margin 与 top-1 逐步对照。
+
+- [x] **P4.2 decode 逐步误差曲线** —— **完成，结论收敛：没有 decode 专属 bug**
+
+      关键前提：**必须 `DEEPINFER_USE_NEXTN_DECODE=0`**。它默认是 `5`，MTP 投机解码默认开启，
+      每步产生多行 draft/verify，`logits[0]` 不再对应你以为的那个位置（我第一次跑就是因此
+      错位，把 pos15 的向量当成 pos14 的，白折腾了一轮）。关掉之后一步一行，完全对齐。
+
+      11 token 起，逐步对照（`nomtp` run，8 个 decode step）：
+
+      | pos | 我方 top1 | 官方 top1 | 我方 margin | 官方 margin | Δmargin | 翻转 |
+      |---|---|---|---|---|---|---|
+      | 11（prefill） | 17607 | 17607 | 10.695 | 9.861 | +0.835 | |
+      | 12 | 2793 | 2793 | 8.448 | 10.328 | −1.880 | |
+      | 13 | 3139 | 3139 | 15.276 | 11.946 | +3.330 | |
+      | **14** | **666** | **6127** | **1.480** | **1.638** | −0.157 | **翻转** |
+      | 15 | 6127 | 6127 | 15.950 | 15.109 | +0.842 | |
+      | 16 | 666 | 666 | 20.118 | 17.874 | +2.244 | |
+      | 17 | 320 | 320 | 9.829 | 10.402 | −0.573 | |
+      | 18 | 1 | 1 | 12.262 | 11.091 | +1.171 | |
+
+      **Δmargin：mean +0.711，std 1.628，|max| 3.330。** 7/8 位置 top-1 正确，
+      唯一的翻转在 pos14 —— 而那里的**官方 margin 只有 1.638**，其余位置是 9.9~17.9。
+
+      换句话说：**噪声幅度（std 1.63）和那个位置的 margin（1.64）一样大，是硬币。**
+      这不是 decode 的 bug，是量化噪声碰上模型自己也没把握的位置。
+
+      官方在 pos14 选了 `6127`（北京），我方选了 `666`（`**`），于是渲染成
+      `**北京**` —— 语义完全等价，只是 markdown 加粗的位置不同。
+
+      顺带排除：MTP 不是原因。`DEEPINFER_USE_NEXTN_DECODE=0` 前后输出完全相同
+      （`[17607, 2793, 3139, 666, 6127, 666, 320, 1]`）。而官方 `generate.py` 本来
+      就没有 MTP，所以把 MTP 关掉才是与参考对等的比法。
+
+      从第 1 步到第 8 步，每步记 `margin`、`top-1 是否与官方一致`、`rel(B_t, R_t)`。
+      用来判断误差是「一开始就有」还是「某一步突然放大」——后者意味着踩到了某个
+      batch/位置相关的边界（例如 page 边界、SWA window 回绕）。
+
+- [x] **P4.3 KV cache 写入/读回一致性** —— **完成，干净，排除**
+
+      两层验证，都通过。
+
+      其一，仓库自带的单测（`store_kv_cache.py` / `load_paged_kv_cache.py`，与各自的 ref
+      逐元素比），过滤 SM90 实际使用的 fp32-sf index 布局：**各 6 passed**。
+      全布局跑一遍是 48 passed / 12 failed，失败**全部**落在 `ue8m0` 的 index 布局
+      （`FP4_INDEX_KV_FORMAT`、`MXFP8_INDEX_KV_FORMAT`），而 SM90 不用这两个。失败点在
+      ref 自身：它用 `PLATFORM.fp8_sf_type`（SM90 是 fp32）去重解释 ue8m0 字节，
+      `index_copy_` 撞 dtype。属 SM90 上 ue8m0 分支的遗留缺口，不影响本任务。
+
+      其二，闭环 round-trip（`scripts/h20_kv_roundtrip.py`）：写入已知 bf16 → 读回 → 与
+      原值比，不依赖任何 ref 实现。覆盖 SM90 实际用的两种布局、block 32/64、1/233/1029 token：
+
+      | 布局 | bytes/token | rel_err | cos |
+      |---|---|---|---|
+      | `V41_FORMAT`（main/swa KV） | 528 | 0.0266–0.0295 | 0.9996 |
+      | `FP8_FLOAT32_SF_INDEX_KV_FORMAT`（index KV） | **132** | 0.0266–0.0327 | 0.9996 |
+
+      正是 FP8 的理论地板，与权重那 2.7% 同量级。**KV cache 不是 decode 误差的来源。**
+      P3.9 遗留的 132B 字节布局问题就此关闭。
+
+      坑：这个脚本第一版把「原始 fp8 值」直接和 bf16 比，得到 rel=126 —— 因为 index 布局
+      的 load **不解量化**（文档写明「返回量化值 + scale」，消费方是 kernel）。
+      按 `量化值 × scale` 解量化后才是 0.027。差 127 倍的假警报来自测试脚本，不是实现。
+
+- [x] **P4.4 `_decode_p_sum_sm90` vs 参考实现** —— **完成，数值精确，排除**
+
+      `scripts/h20_index_decode_check.py` 用一份**刻意不共享代码**的 float64 参考实现
+      （按显式偏移逐 token 切片，而非 strided view）对照：
+
+      ```
+      page=32  lens=[1,64]      rel=1.4e-07  cos=1.00000
+      page=32  lens=[31,32]     rel=1.5e-07  cos=1.00000   (页边界)
+      page=32  lens=[96,65]     rel=1.2e-07  cos=1.00000   (跨页)
+      page=32  lens=[7,33,159]  rel=9.4e-08  cos=1.00000   (3 行 32 头)
+      page=64  ...              rel≈1.3e-07  cos=1.00000
+      ```
+
+      rel 就是 fp32 舍入量级。Q 的解量化用 `q_sf.reshape(rows, heads, -1)[..., :1]` 也没问题：
+      `head_dim == 128 == gran`，每个 (row, head) 恰好一个 scale（源码里有
+      `assert self.n_head <= 32 and self.head_dim == 128` 兜住）。
+      **decode 的 indexer 路径不是误差来源。**
+
+- [x] **P4.5 采样器 / logits processor 排除** —— **无需再查**
+
+      pos14 我方 top1=666、top2=6127，差值 1.480。如果 6127 被任何 mask / blacklist 抑制，
+      它就不会是第二名。所以 pos14 的差异来自 logits 本身，不是采样后处理。
+      （`token_logprobs` 恒为 0 的 P5 问题仍然独立存在，与本题无关。）
+
+### 最终结论：`**` 是量化噪声在「模型自己也没把握」的位置上的一次抛硬币
+
+三条线索合起来：
+
+1. **噪声幅度是恒定的约 1.6 logit**（Δmargin std = 1.63，跨 7 个 decode step），
+   与 prefill 的残余误差同量级（pos11 的 Δmargin +0.835，14-token prefill 的 top-1
+   误差 −3.30）。它来自 per-128 量化的累积 —— SM90 的 DeepGEMM 强制的那条路。
+2. **参考模型的 margin 分布很宽**：9.86 / 10.33 / 11.95 / **1.64** / 15.11 / 17.87 / 10.40 / 11.09。
+   只有 1.64 那个位置会被 1.6 的噪声翻掉。
+3. **只有 pos14 翻转，其余全部正确。** 而 pos14 正是 1.64 那个。
+
+所以：
+
+- 不是 KV cache 的问题（P4.3 闭环 round-trip 通过，rel≈0.027）
+- 不是 decode indexer 的问题（P4.4 与独立 float64 参考 rel≈1e-7）
+- 不是 MTP 的问题（开关 MTP 输出相同）
+- 不是采样后处理的问题（6127 在 pos14 是第二名，没被 mask）
+
+**是「在 SM90 上必须用 per-128 量化」这件事本身**。它在多数位置无感，只在模型自己
+接近 50/50 的地方改变结果。而模型在这里确实接近 50/50 —— `**北京**` 的两个位置
+（先 `**` 还是先 `北京`）在语义上没有区别，参考模型也只是略偏一边。
+
+**要继续压这个噪声，唯一的杠杆是回到 P3.6 那条路，而它在 SM90 上被 DeepGEMM 的
+`gran_k == 128` 断言封死了**（见 `H20-DRIFT-ANALYSIS.md`）。所以本项调查到此收敛：
+剩余差异是硬件量化口径的必然结果，不是实现缺陷。
+
+### P4.1 的工具坑（都已修，但 dump 仍建议配合黑盒交叉验证）
+
+| 坑 | 现象 | 修法 |
+|---|---|---|
+| 多 DP 进程写同一文件 | 8 个 engine 都是单进程组，`RANK` 全为 0，互相覆盖 | 文件名带 `_rank<R>_pid<P>` |
+| decode 的 batch 有 padding 行 | 取 `logits[-1]` 读到占位 token 的垃圾（恒定 `top1=69`） | 取第 0 行（第一个真实请求） |
+| follower rank 的 logits 是占位值 | 同一向量被重复记录几十次，淹没真实步 | 判据 `driver_inputs is not None` |
+| profile forward 用合成 token | 抓到 profile 的数据 | `debug_dump.profiling()` 挂起 |
+
+即便都修了，逐 step dump 与端到端 token 序列仍无法可靠对齐（`max_tokens=7` 的 run 里
+出现 75 次 sampler 调用）。所以 P4.1 的结论**以黑盒四格对照为准**，dump 只作辅助。
+
+**另一个必须记住的坑**：`max_tokens > 1` 的运行里，`logits_prompt` 也会被污染 ——
+`mysteps6`（`max_tokens=7`）的 `logits_prompt` 算出 cos=0.68、top-1 误差 16.5，
+而 `mylogits3`（`max_tokens=1`）同一输入是 cos=0.9846、误差 3.30。差 5 倍。
+**比较 prefill logits 只能用 `max_tokens=1` 的运行。**
+
+### 已备好的工具
+
+- `debug_dump.record(..., once=True)` + `logits_prompt` / `logits_decode` 两个键
+- `debug_dump.profiling()`：profile forward 期间整体挂起，避免抓到合成 token
+- `ref_next_token.py`：吃 token ids，给官方 greedy next token + top5；`--force_ids` 走增量
+  decode，`--logits_out` 存每一步的完整 logits
+- `scripts/h20_kv_roundtrip.py`：KV cache 闭环 round-trip（不依赖 ref）
+- `scripts/h20_index_decode_check.py`：`_decode_p_sum_sm90` 的独立数值对照
+
+## P5 logprob 异常（已排查，未修，非阻塞）

 当前 8 个 token 的 `token_logprobs` 全是 `0.0`（概率=1），不正常。

- [ ] P5.1 确认是采样器 bug 还是 logits 极端（打印 logits 的 max/median）
- [ ] P5.2 检查 `DEEPINFER_SAMPLER_FORCE_TOPK=4096` 与 logprob 回填的关系
+### 已排除（约 20 分钟静态排查 + 单测）
+
+| 嫌疑 | 结论 |
+|---|---|
+| `map_lookup` kernel | 无罪。单测 6 passed，`by_key` + float32 + `force_topk=4096` 全覆盖，与 `map_lookup_ref` 逐元素相等 |
+| sampler 全部 kernel | 无罪。`sampler.py -m "not benchmark"` → **74 passed / 0 failed** |
+| `SampledInfo.from_reqcnt` 的 `torch.zeros` | 死代码，全仓库无调用点；`update(is_final, candidate)` 也无调用点 |
+| sampler→engine 字段传递 | 直通，`llm_engine.py:747` 原样传 `output.token_logprobs` |
+
+**一个反向线索**：`map_lookup` 里 `out = torch.empty`，键未命中时本应写 `-inf` 而**不是 0**。
+所以 0.0 更像「上层根本没把值带下来」，而不是「查表查空了」。
+
+### 剩余候选与判定路径
+
+- `process_sampling` 返回的 `logprobs` 本身为 0（`fused_sample` 的 logprob 输出路径；
+  现有单测只覆盖 kernel 本身，没有覆盖「整个 forward 的 logprob 输出」）
+- `build_request_outputs` 里逐 token 对象的属性回填
+
+判定只需跑 1 次整机（~5 min）：在 sampler 里打印 `logprobs.max()` 与 `token_logprobs_cur`。
+sampler 层就是 0 → 定位在 `fused_sample`；sampler 层正常 → 定位在 output 回填。
+
+### 为什么留着不修
+
+不影响推理正确性，只影响 logprob 字段的可用性（AIME 打分、perplexity 这类评估用途）。
+预估典型 1 h、最坏 2 h（若牵涉 MTP 异步 logprob 回填），与本次 SM90 支持的目标无关。

 ---

--- a/H20-DRIFT-ANALYSIS.md
+++ b/H20-DRIFT-ANALYSIS.md
@ -0,0 +1,152 @@
+# 漂移调查：2.7% 不是主因，问题在 decode 路径
+
+承接 `H20-PORTING-JOURNAL.md` §6.3 的 P3.6（「改用官方无损重化把 2.7% 降到 0」）。
+
+**两个结论**：
+
+1. **P3.6 在 SM90 上不可行** —— DeepGEMM 在 Hopper 上强制 per-128 scale，per-32 是 SM100
+   专属。2.7% 是硬件地板，不是实现缺陷。
+2. **但 2.7% 不是 `**` 的成因** —— prefill 路径的 top-1 与官方一致，翻转发生在 **decode**。
+
+第 2 点是本次最有价值的发现，它把问题从「量化精度无解」重新定位到一个具体的、可能可修的
+bug 上。
+
+---
+
+## 一、P3.6 不可行：DeepGEMM 的 SM90 硬约束
+
+### 证据
+
+`csrc/apis/layout.hpp:100`（DeepGEMM 2.8.2），`check_sf_layout` 里唯一的架构断言：
+
+```cpp
+DG_HOST_ASSERT((arch_major == 9 and gran_k == 128 and k_alignment == 128) or
+               (arch_major == 10 and (gran_k == 32 or gran_k == 128) and k_alignment % 128 == 0));
+```
+
+- `arch_major == 9`（Hopper）→ **`gran_k` 必须 128**
+- `arch_major == 10`（Blackwell）→ 32 或 128 都行
+
+同文件 38-51 行的 `transform_sf_into_required_layout` 印证同一件事 —— SM90 只有两个分支，
+都要求 `gran_k == 128` 且 `sf` 是 **fp32**：
+
+```cpp
+// (FP32, 1, 128) on SM90
+if (sf.scalar_type() == kFloat and gran_mn == 1 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast))
+    return get_mn_major_tma_aligned_tensor(sf);
+
+// (FP32, 128, 128) on SM90
+if (sf.scalar_type() == kFloat and gran_mn == 128 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast))
+    return check_sf_layout(...);
+
+// (FP32, x, gran_k) on SM100 —— 只有这里才出现 packed ue8m0
+if (sf.scalar_type() == kFloat and (gran_k == 32 or gran_k == 128) and arch_major == 10) { ... }
+```
+
+官方 `cast_e2m1fn_to_e4m3fn` 产出的是 **e8m0 per-32**。SM90 要 **fp32 per-128**。
+两者不兼容，无中间等价物。
+
+### 逐条排除的降损方案
+
+| 想法 | 结论 |
+|---|---|
+| 换 per-32 recipe | ❌ 上述断言，SM90 不允许 |
+| scale 用 e8m0 省位宽 | ❌ SM90 分支要求 fp32，e8m0 触发 `sf_dtype` 断言 |
+| 专家权重用 BF16 绕开量化 | ❌ 显存 78.5 GiB/卡 → 约 157 GiB/卡 |
+| 让 per-128 的 scale 更精确 | ❌ 已是 `amax/448`，per-128 网格上的最优表示 |
+
+**所以 2.7% 是「用 SM90 的 GEMM 跑一个按 MXFP8 per-32 导出的 checkpoint」的固有代价。**
+参照量级：dense 权重（`requant_dense_fp8`，用法本来就对）也有 4.7e-2，同一量级。
+
+---
+
+## 二、修复确实有效：layer0 误差降了 4.4 倍
+
+专家 scale 修复（提交 `7c6fdc7`）之后重跑逐层对拍，14-token 输入：
+
+| 层 | 修复前 cos / rel_err | **修复后** cos / rel_err |
+|---|---|---|
+| layer 0 | 0.99190 / 1.27e-1 | **0.99959 / 2.88e-2** |
+| layer 1 | 0.97707 / 2.29e-1 | **0.99908 / 4.45e-2** |
+| layer 2 | 0.96411 / 2.76e-1 | **0.99852 / 5.75e-2** |
+| layer 5 | 0.93275 / 3.61e-1 | **0.99834 / 7.16e-2** |
+| layer 10 | — | 0.99310 / 1.22e-1 |
+| layer 20 | — | 0.98709 / 1.61e-1 |
+| layer 39 | 0.86515 / 5.28e-1（旧 layer11 数据） | **0.97972 / 2.00e-1** |
+
+这是 P3.7 想验证的东西，结论明确：误差曲线整体压平了一个数量级。
+
+---
+
+## 三、关键发现：prefill 的 top-1 是对的
+
+之前一直没能比较最终 logits，原因是 `ref_layer_dump.py` 写的是 `logits[0, -1]`，
+而 `MainPredictHeadHC` 在 `full_logits=False` 时已经压掉了 seq 维（`x = x[:, -1]`），
+所以 logits 是 `[b, vocab]` —— 那个索引会得到一个 **0 维标量**。修成 `logits[0]` 之后
+才拿到可以比较的数据。
+
+14-token 输入，两侧最终 logits：
+
+| | 官方 | 移植后 |
+|---|---|---|
+| cos | — | **0.98460** |
+| top-1 | **6127（北京）** | **6127（北京）** ✅ |
+| logit(6127) | 39.343 | 36.047 |
+| logit(666 = `**`) | 36.232 | 35.791 |
+| **margin** | **+3.112** | **+0.256** |
+
+**prefill 路径的 top-1 与官方一致。** 残余误差的表现为 margin 从 3.11 压到 0.26（约 12 倍），
+但方向没翻。
+
+### 所以 `**` 从哪来
+
+端到端（11-token 输入，greedy）实际输出：
+
+```
+[17607, 2793, 3139, 666, 6127, 666, 320, 1]   → 中国的首都是**北京**。
+        ↑ 前 3 个与官方逐 token 一致      ↑ 第 4 步翻成 666
+```
+
+而同样这 14 个 token **一次性 prefill** 时，top-1 是 6127（正确）。
+
+**同一个 14-token 上下文，prefill 给 6127，增量 decode 给 666。**
+
+这就是 `**` 的来源：不是量化精度，是 **decode 路径**引入了额外误差，把本来只有 0.26 的
+margin 彻底翻转。
+
+---
+
+## 四、下一步：P4 decode 专项
+
+定位已经收窄到 decode，建议按这个顺序：
+
+- [ ] **P4.1 同位置 prefill vs decode 对拍**（最高优先）
+      输入 14 token：① 一次 prefill，dump 最后位置 logits（已有：margin +0.256）；
+      ② 先喂 11 token，再用已确认正确的 3 个 token（17607/2793/3139）逐步 decode，
+      dump 第 4 步的 logits。两者不同即确认 decode 引入偏差，且可以直接量出偏差大小。
+
+- [ ] **P4.2 decode 逐步误差曲线**
+      从第 1 步到第 8 步，每步记录 top-1 和 margin，找出误差开始放大的步数。
+
+- [ ] **P4.3 KV cache 写入/读回一致性**
+      写入已知向量后读回，与 bf16 原值比。这是 P3.9 遗留项，也是最可能的嫌疑人：
+      decode 用 `_decode_p_sum_sm90` 走 torch 兜底算 indexer 分数，而 prefill 走
+      TileLang kernel —— 两条路径从未对拍过。
+
+- [ ] **P4.4 `_decode_p_sum_sm90` vs prefill 的 MQA logits**
+      同一个 KV cache、同一个位置，两条 indexer 路径的分数差。
+
+判据：如果 P4.1 的差距 < 1e-3，说明 decode 无误，`**` 另有来源（回到采样器/logits
+processor 查）；如果差距明显，就从 P4.3/P4.4 里找。
+
+---
+
+## 五、两个工具修正（否则上面的数据拿不到）
+
+1. **`ref_layer_dump.py` 的 logits 索引**：`logits[0, -1]` → `logits[0]`。之前所有
+   dump 里的 `final_row_logits` 都是 0 维标量，等于没有。
+2. **profile run 污染 dump**：`execute_model_profile_run` 会走完整的 forward + sampler，
+   用的是合成 token。层 dump 靠 token 数门控侥幸过滤掉了，但 logits 的行数是「选中 token 数」
+   而不是 prompt 长度，门控不适用，于是抓到的是 profile 的垃圾（第一版得到 cos = 0.34）。
+   现在用 `debug_dump.profiling()` 上下文在 profile 期间整体挂起，并且把 logits 拆成
+   `logits_prompt` / `logits_decode` 两个键 —— 只有前者能和 prefill-only 的参考 dump 比较。
--- a/deepinfer/debug_dump.py
+++ b/deepinfer/debug_dump.py
@ -0,0 +1,183 @@
+"""Opt-in tensor dumps for H20 bring-up comparisons.
+
+Enabled with `DEEPINFER_DUMP_LAYERS=<dir>`. Each call appends one row to
+`<dir>/deepinfer_layers.npz`-style storage under a per-tag key, always as the *last* row of
+the batch: the reference implementation only generates from the final position, so that is
+the only row the two sides can be compared on.
+
+The dump is deliberately simple (host copy + `numpy.savez` at exit) because it exists to
+answer "which layer diverges first", not to be a profiler.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+
+_DUMP_DIR = os.getenv("DEEPINFER_DUMP_LAYERS")
+# The engine runs a dummy profile forward before it allocates the KV cache, and that one has a
+# different token count. Gating on the real prompt length keeps the dump comparable.
+_DUMP_TOKENS = int(os.getenv("DEEPINFER_DUMP_TOKENS", "0") or 0)
+_CAPTURED: dict[str, np.ndarray] = {}
+# Only the first N layers of MoE internals are kept: the routing and per-branch outputs are the
+# interesting part for bring-up, and dumping all 40 layers would slow the forward down for no gain.
+_MOE_LAYERS = int(os.getenv("DEEPINFER_DUMP_MOE_LAYERS", "2") or 0)
+# Per-tag call counter for `record_step`.
+_STEPS: dict[str, int] = {}
+# Number of rows in the most recent logits tensor, used to skip warmup passes.
+_ROWS = 0
+# A decode step that serves real requests has at most this many rows. Anything wider is warmup.
+_REAL_ROWS = int(os.getenv("DEEPINFER_DUMP_MAX_ROWS", "1") or 1)
+
+
+_SUSPENDED = False
+
+
+def enabled() -> bool:
+    return bool(_DUMP_DIR) and not _SUSPENDED
+
+
+class profiling:
+    """Suspend capture while the engine runs its dummy profile forward.
+
+    `execute_model_profile_run` goes through the same forward and sampler as a real step, with
+    synthetic tokens, so anything captured there is garbage that later masking cannot detect. The
+    layer dumps happen to be filtered by their token-count gate; the logits are not, because their
+    row count is the number of *selected* tokens rather than the prompt length.
+    """
+
+    def __enter__(self) -> None:
+        global _SUSPENDED
+        self._previous = _SUSPENDED
+        _SUSPENDED = True
+
+    def __exit__(self, *_exc: object) -> None:
+        global _SUSPENDED
+        _SUSPENDED = self._previous
+
+
+def record(tag: str, tensor: torch.Tensor, once: bool = False) -> None:
+    """Store the last row of `tensor` under `tag`.
+
+    A 3D tensor is treated as `[tokens, hc_mult, dim]` and reduced with `mean(dim=1)`, which is
+    the closest match to the reference's collapsed residual stream.
+
+    `once=True` keeps the first write only. Logits need this: the row count is the number of
+    selected tokens rather than the prompt length (so the token gate does not apply), and the
+    decode step that follows prefill would otherwise overwrite the value we care about.
+    """
+    if not _DUMP_DIR:
+        return
+    if _DUMP_TOKENS and not once and tensor.shape[0] != _DUMP_TOKENS:
+        return
+    if once and tag in _CAPTURED:
+        return
+    if not _want_moe(tag):
+        return
+    if tensor.dim() == 3:
+        tensor = tensor.mean(dim=1)
+    row = tensor[-1] if tensor.dim() == 2 else tensor.reshape(-1)
+    # Copying to the host is the whole point here, and it happens inside the forward where the
+    # engine forbids syncs, so take the sanctioned escape hatch for this one transfer.
+    from deepinfer.utils import allow_cuda_sync
+
+    with allow_cuda_sync():
+        _CAPTURED[tag] = row.detach().float().cpu().numpy()
+
+
+def set_rows(rows: int) -> None:
+    """Tell the dump how many rows the sampler just produced, so warmup can be filtered out."""
+    global _ROWS
+    _ROWS = rows
+
+
+def record_step(tag: str, tensor: torch.Tensor) -> None:
+    """Append `tensor`'s last row under `tag_<n>`, numbering the calls.
+
+    Decode needs every step kept, not just the first: the point of the decode comparison is to
+    watch the margin evolve and find the step where it collapses.
+    """
+    if not _DUMP_DIR:
+        return
+    # Decode with one request produces exactly `_REAL_ROWS` rows; warmup batches are larger.
+    if _ROWS > _REAL_ROWS:
+        return
+    index = _STEPS.get(tag, 0)
+    _STEPS[tag] = index + 1
+    # Row 0, not the last row: during decode the batch is padded (CUDA graph slots, or just
+    # `max_num_seqs`), and the padding rows hold logits for placeholder tokens. Row 0 is the
+    # first real request's current token, which is the one to compare against a batch-1 reference.
+    row = tensor[0] if tensor.dim() >= 2 else tensor
+    from deepinfer.utils import allow_cuda_sync
+
+    with allow_cuda_sync():
+        _CAPTURED[f"{tag}_{index}"] = row.detach().float().cpu().numpy()
+
+
+def record_ids(tag: str, tensor: torch.Tensor) -> None:
+    """Store an integer tensor as-is (expert ids); reduction would destroy its meaning."""
+    if not _DUMP_DIR:
+        return
+    if not _want_moe(tag):
+        return
+    if _DUMP_TOKENS and tensor.shape[0] != _DUMP_TOKENS:
+        return
+    from deepinfer.utils import allow_cuda_sync
+
+    with allow_cuda_sync():
+        _CAPTURED[tag] = tensor.detach().to(torch.int32).cpu().numpy()
+
+
+def _want_moe(tag: str) -> bool:
+    """Keep MoE internals only for the first `_MOE_LAYERS` layers."""
+    if "_" not in tag:
+        return True
+    head = tag.split("_", 1)[0]
+    if not head.startswith("layer"):
+        return True
+    try:
+        return int(head[5:]) < _MOE_LAYERS
+    except ValueError:
+        return True
+
+
+def _rank_suffix() -> str:
+    """`_rank<N>` when running distributed, so ranks do not overwrite each other's dump.
+
+    With data parallelism every rank registers `atexit(flush)` and writes the same filename; the
+    last one to exit wins, which silently mixes one rank's tensors with another's. Only rank 0
+    holds the real prompt anyway, so the suffix makes that explicit instead of accidental.
+    """
+    rank = os.getenv("RANK")
+    if rank is None:
+        try:
+            import torch.distributed as dist
+
+            if dist.is_initialized():
+                rank = str(dist.get_rank())
+        except Exception:  # noqa: BLE001 - torch may not be usable at exit
+            rank = None
+    # With data parallelism each engine is its own single-process group, so every rank reports
+    # RANK=0 and they all target the same file. The pid disambiguates them.
+    return f"_rank{rank if rank is not None else 'x'}_pid{os.getpid()}"
+
+
+def flush() -> None:
+    """Write everything captured so far. Called once at interpreter exit."""
+    if not _DUMP_DIR or not _CAPTURED:
+        return
+    out = Path(_DUMP_DIR)
+    out.mkdir(parents=True, exist_ok=True)
+    path = out / f"deepinfer_layers{_rank_suffix()}.npz"
+    np.savez(path, **_CAPTURED)
+    print(f"[dump] wrote {len(_CAPTURED)} tensors to {path}", flush=True)
+    _CAPTURED.clear()
+
+
+if _DUMP_DIR:
+    import atexit
+
+    atexit.register(flush)
--- a/deepinfer/h20_ops/debug_dump.py
+++ b/deepinfer/h20_ops/debug_dump.py
@ -1,99 +0,0 @@
-"""Opt-in tensor dumps for H20 bring-up comparisons.
-
-Enabled with `DEEPINFER_DUMP_LAYERS=<dir>`. Each call appends one row to
-`<dir>/deepinfer_layers.npz`-style storage under a per-tag key, always as the *last* row of
-the batch: the reference implementation only generates from the final position, so that is
-the only row the two sides can be compared on.
-
-The dump is deliberately simple (host copy + `numpy.savez` at exit) because it exists to
-answer "which layer diverges first", not to be a profiler.
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import numpy as np
-import torch
-
-_DUMP_DIR = os.getenv("DEEPINFER_DUMP_LAYERS")
-# The engine runs a dummy profile forward before it allocates the KV cache, and that one has a
-# different token count. Gating on the real prompt length keeps the dump comparable.
-_DUMP_TOKENS = int(os.getenv("DEEPINFER_DUMP_TOKENS", "0") or 0)
-_CAPTURED: dict[str, np.ndarray] = {}
-# Only the first N layers of MoE internals are kept: the routing and per-branch outputs are the
-# interesting part for bring-up, and dumping all 40 layers would slow the forward down for no gain.
-_MOE_LAYERS = int(os.getenv("DEEPINFER_DUMP_MOE_LAYERS", "2") or 0)
-
-
-def enabled() -> bool:
-    return bool(_DUMP_DIR)
-
-
-def record(tag: str, tensor: torch.Tensor) -> None:
-    """Store the last row of `tensor` under `tag`.
-
-    A 3D tensor is treated as `[tokens, hc_mult, dim]` and reduced with `mean(dim=1)`, which is
-    the closest match to the reference's collapsed residual stream.
-    """
-    if not _DUMP_DIR:
-        return
-    if _DUMP_TOKENS and tensor.shape[0] != _DUMP_TOKENS:
-        return
-    if not _want_moe(tag):
-        return
-    if tensor.dim() == 3:
-        tensor = tensor.mean(dim=1)
-    row = tensor[-1] if tensor.dim() == 2 else tensor.reshape(-1)
-    # Copying to the host is the whole point here, and it happens inside the forward where the
-    # engine forbids syncs, so take the sanctioned escape hatch for this one transfer.
-    from deepinfer.utils import allow_cuda_sync
-
-    with allow_cuda_sync():
-        _CAPTURED[tag] = row.detach().float().cpu().numpy()
-
-
-def record_ids(tag: str, tensor: torch.Tensor) -> None:
-    """Store an integer tensor as-is (expert ids); reduction would destroy its meaning."""
-    if not _DUMP_DIR:
-        return
-    if not _want_moe(tag):
-        return
-    if _DUMP_TOKENS and tensor.shape[0] != _DUMP_TOKENS:
-        return
-    from deepinfer.utils import allow_cuda_sync
-
-    with allow_cuda_sync():
-        _CAPTURED[tag] = tensor.detach().to(torch.int32).cpu().numpy()
-
-
-def _want_moe(tag: str) -> bool:
-    """Keep MoE internals only for the first `_MOE_LAYERS` layers."""
-    if "_" not in tag:
-        return True
-    head = tag.split("_", 1)[0]
-    if not head.startswith("layer"):
-        return True
-    try:
-        return int(head[5:]) < _MOE_LAYERS
-    except ValueError:
-        return True
-
-
-def flush() -> None:
-    """Write everything captured so far. Called once at interpreter exit."""
-    if not _DUMP_DIR or not _CAPTURED:
-        return
-    out = Path(_DUMP_DIR)
-    out.mkdir(parents=True, exist_ok=True)
-    path = out / "deepinfer_layers.npz"
-    np.savez(path, **_CAPTURED)
-    print(f"[dump] wrote {len(_CAPTURED)} tensors to {path}", flush=True)
-    _CAPTURED.clear()
-
-
-if _DUMP_DIR:
-    import atexit
-
-    atexit.register(flush)
--- a/deepinfer/model_executor/layers/sampler.py
+++ b/deepinfer/model_executor/layers/sampler.py
@ -11,6 +11,7 @@ import torch
 import torch.nn as nn
 from deep_select.interface import topk

+from deepinfer import debug_dump as _debug_dump
 from deepinfer import envs
 from deepinfer.logger import init_logger
 from deepinfer.model_executor.input_metadata import InputMetadata
@ -140,6 +141,19 @@ class Sampler(nn.Module):
        # CUDA Graph model_output retains graph-batch padding; gather trims hidden_states.
        hidden_states = gather(hidden_states, selected_token_indices)
        logits = get_logits(hidden_states, embedding, self.vocab_size)
+        # `driver_inputs is None` means this rank only participates in the embedding all-gather and
+        # never sees the real result, so its logits are placeholders -- identical on every call.
+        # Recording those buries the real steps under dozens of copies of the same vector.
+        if _debug_dump.enabled() and driver_inputs is not None:
+            # Capture before any masking/penalty so the numbers are comparable with the reference,
+            # which dumps the raw head output as well. Named per phase: with `max_tokens=1` both a
+            # prompt step and a decode step run, and only the former is comparable with a
+            # prefill-only reference dump.
+            if is_prompt:
+                _debug_dump.record("logits_prompt", logits, once=True)
+            else:
+                # Every decode step, so the comparison can watch the margin collapse step by step.
+                _debug_dump.record_step("logits_decode", logits)
        del hidden_states

        if driver_inputs is None:
--- a/deepinfer/model_executor/models/decoder_layer.py
+++ b/deepinfer/model_executor/models/decoder_layer.py
@ -5,7 +5,7 @@ import torch
 from torch import nn

 from deepinfer import AttrDict, envs
-from deepinfer.h20_ops import debug_dump as _debug_dump
+from deepinfer import debug_dump as _debug_dump
 from deepinfer.logger import init_logger
 from deepinfer.model_executor.input_metadata import InputMetadata
 from deepinfer.model_executor.input_metadata.yoco import YocoMetadataMixin
--- a/deepinfer/model_executor/models/deepseek.py
+++ b/deepinfer/model_executor/models/deepseek.py
@ -9,8 +9,8 @@ import torch
 from torch import nn

 from deepinfer import AttrDict, envs
+from deepinfer import debug_dump as _debug_dump
 from deepinfer.engine.engine_type import EngineType
-from deepinfer.h20_ops import debug_dump as _debug_dump
 from deepinfer.logger import init_logger
 from deepinfer.model_executor import eplb
 from deepinfer.model_executor.input_metadata import InputMetadata
--- a/deepinfer/model_executor/models/v4_modules/mlp.py
+++ b/deepinfer/model_executor/models/v4_modules/mlp.py
@ -21,8 +21,8 @@ from tile_kernels.quant import swiglu_forward_and_per_token_cast
 from torch import nn

 from deepinfer import PLATFORM, AttrDict, envs
+from deepinfer import debug_dump as _debug_dump
 from deepinfer.engine.engine_type import EngineType
-from deepinfer.h20_ops import debug_dump as _debug_dump
 from deepinfer.logger import init_logger
 from deepinfer.model_executor import eplb
 from deepinfer.model_executor.input_metadata import InputMetadata
--- a/deepinfer/worker/model_runner.py
+++ b/deepinfer/worker/model_runner.py
@ -11,6 +11,7 @@ import numpy as np
 import torch
 import torch.distributed as dist

+from deepinfer import debug_dump as _debug_dump
 from deepinfer import envs
 from deepinfer._C.engine import prepare_sampler_inputs, update_prompt_logprobs  # pyright: ignore[reportMissingModuleSource]
 from deepinfer.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig
@ -283,26 +284,28 @@ class ModelRunner:
        kv_caches: GPUCacheInterface,
    ) -> None:
        """Profile-run only: run one full step (forward + sample + post-process) synchronously."""
-        main_model_output = self.execute_main_model(
-            seq_metadata_list,
-            kv_caches,
-        )
-        input_metadata = main_model_output.input_metadata
-        model_output = Movable(main_model_output.model_output)
-        del main_model_output
+        # Synthetic tokens: keep tensor dumps off, or they capture this step instead of the real one.
+        with _debug_dump.profiling():
+            main_model_output = self.execute_main_model(
+                seq_metadata_list,
+                kv_caches,
+            )
+            input_metadata = main_model_output.input_metadata
+            model_output = Movable(main_model_output.model_output)
+            del main_model_output

-        sampler_output = self.execute_sample(
-            seq_metadata_list,
-            input_metadata,
-            model_output,
-            kv_caches,
-        )
-        torch.cuda.current_stream().synchronize()
-        self.post_process(
-            seq_metadata_list,
-            sampler_output,
-            input_metadata.is_prompt,
-        )
+            sampler_output = self.execute_sample(
+                seq_metadata_list,
+                input_metadata,
+                model_output,
+                kv_caches,
+            )
+            torch.cuda.current_stream().synchronize()
+            self.post_process(
+                seq_metadata_list,
+                sampler_output,
+                input_metadata.is_prompt,
+            )

    @torch.inference_mode()
    def execute_main_model(
--- a/scripts/h20_index_decode_check.py
+++ b/scripts/h20_index_decode_check.py
@ -0,0 +1,127 @@
+"""Numerically test the SM90 decode-time indexer score against a reference formula.
+
+`_decode_p_sum_sm90` reimplements DeepGEMM's `fp8_fp4_paged_mqa_logits` in PyTorch, because
+Hopper's paged kernel hard-requires 64-token pages while the index cache uses a different page
+size. Nothing else exercises it: prefill goes through the real kernel, so a defect here would be
+decode-only -- which is exactly where the remaining drift lives.
+
+The reference below deliberately shares no code with the implementation: it reads the cache with
+explicit per-token slicing in float64 rather than strided views, so an offset or stride mistake
+in one does not cancel against the other.
+
+    python scripts/h20_index_decode_check.py
+"""
+
+import torch
+
+from deepinfer.model_executor.layers.tilelang_ops.cache_common import FP8_FLOAT32_SF_INDEX_KV_FORMAT
+from deepinfer.model_executor.models.v4_modules.index_attn import MQA
+
+DEV = "cuda"
+HEAD_DIM = 128
+
+
+def build_cache(num_blocks: int, page_tokens: int, scales: torch.Tensor, keys: torch.Tensor) -> torch.Tensor:
+    """Cache in the store kernel's layout: block-wide values, then block-wide scales."""
+    fmt = FP8_FLOAT32_SF_INDEX_KV_FORMAT
+    raw = torch.zeros((num_blocks, page_tokens, 1, fmt.bytes_per_token), dtype=torch.uint8, device=DEV)
+    flat = raw.view(-1)
+    for block in range(num_blocks):
+        base = block * fmt.bytes_per_token * page_tokens
+        values = keys[block].to(torch.float8_e4m3fn).view(torch.uint8).reshape(-1)
+        flat[base : base + page_tokens * HEAD_DIM] = values
+        sf = scales[block].float().reshape(-1).view(torch.uint8)
+        flat[base + page_tokens * HEAD_DIM : base + fmt.bytes_per_token * page_tokens] = sf
+    return raw
+
+
+def reference_p_sum(
+    q: torch.Tensor,
+    kv_cache: torch.Tensor,
+    weights: torch.Tensor,
+    klen: torch.Tensor,
+    block_table: torch.Tensor,
+    page_tokens: int,
+) -> torch.Tensor:
+    """float64 reference: gather by explicit offsets, dequantize, `sum_h relu(q.k) w`."""
+    num_rows = q.shape[0]
+    max_len = int(klen.max())
+    out = torch.full((num_rows, max_len), float("-inf"), dtype=torch.float64, device=DEV)
+    fmt_bytes = FP8_FLOAT32_SF_INDEX_KV_FORMAT.bytes_per_token
+    for row in range(num_rows):
+        ctx = int(klen[row])
+        keys = []
+        for pos in range(ctx):
+            page = pos // page_tokens
+            offset = pos % page_tokens
+            block = int(block_table[row, page])
+            base = block * fmt_bytes * page_tokens
+            raw = kv_cache.view(-1)[base + offset * HEAD_DIM : base + (offset + 1) * HEAD_DIM]
+            scale = kv_cache.view(-1)[base + page_tokens * HEAD_DIM + offset * 4 : base + page_tokens * HEAD_DIM + offset * 4 + 4]
+            keys.append(raw.view(torch.float8_e4m3fn).double() * scale.view(torch.float32).double())
+        key = torch.stack(keys)  # [ctx, head_dim]
+        scores = q[row].double() @ key.t()
+        out[row, :ctx] = (scores.relu() * weights[row].double().unsqueeze(-1)).sum(dim=0)
+    return out
+
+
+def run_case(page_tokens: int, num_blocks: int, num_rows: int, num_heads: int, lens: list[int], seed: int) -> None:
+    torch.manual_seed(seed)
+    total_tokens = num_blocks * page_tokens
+    keys = torch.randn(num_blocks, page_tokens, HEAD_DIM, dtype=torch.float32, device=DEV) * 0.5
+    # One fp32 scale per token, as `elements_per_sf == quantized_elements == 128` implies.
+    scales = (keys.abs().amax(dim=-1) / 448.0).clamp_min(1e-8)
+    cache = build_cache(num_blocks, page_tokens, scales, keys)
+
+    q = torch.randn(num_rows, num_heads, HEAD_DIM, dtype=torch.float32, device=DEV) * 0.5
+    q_sf = (q.abs().amax(dim=-1, keepdim=True) / 448.0).clamp_min(1e-8)  # one scale per (row, head)
+    q_fp8 = (q / q_sf).to(torch.float8_e4m3fn)
+    q_deq = q_fp8.float() * q_sf
+
+    weights = torch.rand(num_rows, num_heads, dtype=torch.float32, device=DEV)
+    klen = torch.tensor(lens, dtype=torch.int32, device=DEV)
+    max_pages = (max(lens) + page_tokens - 1) // page_tokens
+    block_table = torch.zeros(num_rows, max_pages, dtype=torch.int32, device=DEV)
+    for row in range(num_rows):
+        pages = (lens[row] + page_tokens - 1) // page_tokens
+        block_table[row, :pages] = torch.randperm(num_blocks, device=DEV)[:pages]
+
+    class _Stub:
+        max_kv_len = total_tokens
+
+    got = MQA._decode_p_sum_sm90(_Stub(), q_fp8, q_sf, cache, weights, klen, block_table)
+    ref = reference_p_sum(q_deq, cache, weights, klen, block_table, page_tokens)
+
+    got = got[:, : ref.shape[1]].double()
+    # Positions beyond each row's context are `-inf` by design on both sides; compare only the
+    # live window, which is what top-k actually sees.
+    live = torch.arange(ref.shape[1], device=DEV).unsqueeze(0) < klen.double().unsqueeze(1)
+    rel = (got[live] - ref[live]).norm() / ref[live].norm()
+    cos = torch.nn.functional.cosine_similarity(got[live], ref[live], dim=0)
+    worst = (got[live] - ref[live]).abs().max().item()
+
+    ok = "OK" if rel < 1e-5 else "MISMATCH"
+    print(
+        f"page={page_tokens:<4} blocks={num_blocks:<3} rows={num_rows} heads={num_heads:<3} "
+        f"lens={lens}  rel={rel:9.2e} cos={cos:8.5f} max|d|={worst:9.2e}  {ok}"
+    )
+    if rel >= 1e-5:
+        # Localize: which row/position disagrees first.
+        diff = (got - ref).abs().masked_fill(~live, -1)
+        row, col = divmod(int(diff.argmax()), diff.shape[1])
+        print(f"{'':4}worst at row {row} pos {col}: got={got[row, col].item():.6f} ref={ref[row, col].item():.6f}")
+        raise SystemExit("decode indexer disagrees with the reference")
+
+
+def main() -> None:
+    print(f"device: {torch.cuda.get_device_name(0)}\n")
+    for page_tokens in (32, 64):
+        run_case(page_tokens, 4, 2, 8, [1, page_tokens * 2], seed=0)
+        run_case(page_tokens, 4, 2, 8, [page_tokens - 1, page_tokens], seed=1)
+        run_case(page_tokens, 4, 2, 8, [page_tokens * 3, page_tokens * 2 + 1], seed=2)
+        run_case(page_tokens, 8, 3, 32, [7, 33, page_tokens * 5 - 1], seed=3)
+    print("\nINDEX_DECODE_CHECK_DONE")
+
+
+if __name__ == "__main__":
+    main()
--- a/scripts/h20_kv_roundtrip.py
+++ b/scripts/h20_kv_roundtrip.py
@ -0,0 +1,102 @@
+"""Closed-loop KV cache round trip on the layouts SM90 actually uses.
+
+The unit tests next to each op compare it against a reference implementation, and both sides
+derive from the same `KVCacheFormat`, so they can agree while both being wrong about what the
+engine stores. This closes the loop instead: write a known BF16 tensor through the paged cache,
+read it back, and compare against the original. What survives is what decode sees.
+
+Both layouts SM90 runs are covered:
+
+* `V41_FORMAT` -- main/swa KV, e4m3 with per-32 ue8m0 scales, cast on store;
+* `FP8_FLOAT32_SF_INDEX_KV_FORMAT` -- index KV, 128 fp8 + one fp32 scale = 132 B/token.
+
+    python scripts/h20_kv_roundtrip.py
+"""
+
+import torch
+
+from deepinfer import PLATFORM
+from deepinfer.model_executor.layers.tilelang_ops.cache_common import (
+    FP8_FLOAT32_SF_INDEX_KV_FORMAT,
+    V41_FORMAT,
+    KVCacheFormat,
+)
+from deepinfer.model_executor.layers.tilelang_ops.load_paged_kv_cache import load_paged_kv_cache
+from deepinfer.model_executor.layers.tilelang_ops.store_kv_cache import store_kv_cache
+
+DEV = "cuda"
+
+
+def alloc_cache(layout: KVCacheFormat, num_blocks: int, block_size: int) -> torch.Tensor:
+    """Exactly what the engine allocates: a uint8 blob with `bytes_per_token` per token."""
+    return torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device=DEV)
+
+
+def roundtrip(layout: KVCacheFormat, block_size: int, num_tokens: int, seed: int = 0) -> None:
+    torch.manual_seed(seed)
+    num_blocks = (num_tokens + block_size - 1) // block_size
+
+    # One token per slot, sequential; that is the shape prefill produces.
+    kv = torch.randn(num_tokens, layout.num_elements, dtype=torch.bfloat16, device=DEV)
+    slots = torch.arange(num_tokens, dtype=torch.int32, device=DEV)
+    cache = alloc_cache(layout, num_blocks, block_size)
+
+    store_kv_cache(kv, cache, slots, layout=layout)
+
+    # One segment per source block, covering the tokens actually written.
+    segments = []
+    dst = 0
+    for block in range(num_blocks):
+        rows = min(block_size, num_tokens - block * block_size)
+        segments.append([block, 0, rows, dst])
+        dst += rows
+    segments = torch.tensor(segments, dtype=torch.int32, device=DEV)
+
+    out = load_paged_kv_cache(cache, segments, num_tokens, layout=layout)
+    if out.bf16_data is not None:
+        # `do_cast` layouts are dequantized by the kernel itself.
+        back = out.bf16_data.float()
+    else:
+        # Index layouts return quantized values plus scales, in their natural dtype -- the
+        # consumer is a kernel that wants raw bytes, not BF16. Dequantize here or the comparison
+        # is raw fp8 against BF16 (which reads as a ~127x magnitude error, direction intact).
+        assert out.sf_data is not None, "index layout must return scales"
+        group = layout.elements_per_sf
+        sf = out.sf_data.float()
+        if sf.ndim == 1:
+            sf = sf.unsqueeze(-1)
+        scale = sf.repeat_interleave(group, dim=-1)[:, : layout.num_elements]
+        back = out.quantized_data.float() * scale
+
+    ref = kv.float()
+    rel = (back - ref).norm() / ref.norm()
+    cos = torch.nn.functional.cosine_similarity(back.flatten(), ref.flatten(), dim=0)
+
+    # e4m3 keeps 3 mantissa bits, so ~3% is the floor for any value that goes through it.
+    budget = 0.06
+    ok = "OK" if rel < budget else "FAIL"
+    print(
+        f"{layout.note:34} block={block_size:<4} tokens={num_tokens:<6} "
+        f"bytes/token={layout.bytes_per_token:<4} rel={rel:8.4f} cos={cos:8.5f}  {ok}"
+    )
+    if rel >= budget:
+        raise SystemExit(f"round trip lost more than fp8 should: {layout.note}")
+
+    # Scale handling is the part most likely to be silently wrong: check that the per-token
+    # magnitudes survive, not just the direction.
+    per_token = (back - ref).norm(dim=1) / ref.norm(dim=1).clamp_min(1e-6)
+    print(f"{'':34} worst token rel={per_token.max():.4f} median={per_token.median():.4f}")
+
+
+def main() -> None:
+    print(f"platform: sm90={PLATFORM.is_sm90} sf_dtype={PLATFORM.fp8_sf_type} pack={PLATFORM.sf_pack_factor}\n")
+    for layout in (V41_FORMAT, FP8_FLOAT32_SF_INDEX_KV_FORMAT):
+        for block_size in (32, 64):
+            for num_tokens in (1, 233, 1029):
+                roundtrip(layout, block_size, num_tokens)
+        print()
+    print("KV_ROUNDTRIP_DONE")
+
+
+if __name__ == "__main__":
+    main()
--- a/scripts/h20_recipe_probe.py
+++ b/scripts/h20_recipe_probe.py
@ -0,0 +1,88 @@
+"""Probe which FP8 scale recipes DeepGEMM accepts on this GPU, via its own layout helper.
+
+The SM90 path re-quantizes everything to fp32 per-128 because `get_gemm_recipe()` assumes Hopper
+has nothing finer. The reference converter instead keeps an e8m0 scale per 32x32 block and folds
+the relative offset into the FP8 values, which is lossless -- worth adopting only if the grouped
+GEMM actually accepts it.
+
+`transform_sf_into_required_layout` is the single place DeepGEMM validates a recipe (the engine
+calls it for every MoE weight), so probing it answers the question directly and in milliseconds,
+without materializing weights or launching a kernel.
+
+    python scripts/h20_recipe_probe.py
+"""
+
+import deep_gemm
+import torch
+
+E8M0 = torch.float8_e8m0fnu
+
+
+def make_sf(m: int, k: int, block_m: int, block_k: int, num_groups: int, e8m0: bool) -> torch.Tensor:
+    """Scales in the plainest layout: one value per (block_m x block_k) tile, per group."""
+    shape = (num_groups * (m // block_m), k // block_k)
+    if e8m0:
+        return torch.full(shape, 127, dtype=torch.uint8, device="cuda").view(E8M0)
+    return torch.ones(shape, dtype=torch.float32, device="cuda")
+
+
+def probe(recipe_b: tuple[int, int], block_a: int, e8m0: bool, n: int, k: int, num_groups: int) -> str:
+    block_n, block_k = recipe_b
+    sf = make_sf(n, k, block_n, block_k, num_groups, e8m0)
+    try:
+        out = deep_gemm.transform_sf_into_required_layout(
+            sf,
+            n,
+            k,
+            recipe=(1, block_n, block_k),
+            num_groups=num_groups,
+            is_sfa=False,
+            disable_ue8m0_cast=e8m0,
+        )
+        got = tuple(out.shape)
+        # The helper silently re-lays-out whatever it is given, so acceptance alone is not proof;
+        # report the shape so an obviously wrong result is visible too.
+        return f"ACCEPTED  out={got} dtype={out.dtype}"
+    except Exception as exc:  # noqa: BLE001 - report whatever it raises
+        first = (str(exc).strip().splitlines() or [type(exc).__name__])[0]
+        return f"REJECTED  {first[:100]}"
+
+
+def main() -> None:
+    print(f"device: {torch.cuda.get_device_name(0)}")
+    print(f"deep_gemm: {deep_gemm.__version__}\n")
+
+    n, k, num_groups = 2304, 2560, 8
+
+    print("B-operand recipes (what expert weights use, is_sfa=False):")
+    for recipe, block_a, e8m0, label in [
+        ((128, 128), 128, False, "fp32 per-128   (current SM90)"),
+        ((128, 128), 128, True, "e8m0 per-128"),
+        ((32, 32), 128, False, "fp32 per-32    (ref-converter grid)"),
+        ((32, 32), 128, True, "e8m0 per-32    (ref-converter, lossless)"),
+        ((1, 32), 128, True, "e8m0 per-row/32 (MXFP8)"),
+        ((1, 128), 128, True, "e8m0 per-row/128"),
+    ]:
+        print(f"  {label:34} recipe={recipe} e8m0={e8m0:<5} -> {probe(recipe, block_a, e8m0, n, k, num_groups)}")
+
+    print("\nA-operand recipes (activations, is_sfa=True):")
+    for block_k, e8m0, label in [
+        (128, False, "fp32 per-128   (current SM90)"),
+        (32, True, "e8m0 per-32    (MXFP8)"),
+        (32, False, "fp32 per-32"),
+    ]:
+        sf = make_sf(n, k, 1, block_k, 1, e8m0)
+        try:
+            out = deep_gemm.transform_sf_into_required_layout(
+                sf, n, k, recipe=(1, 1, block_k), num_groups=1, is_sfa=True, disable_ue8m0_cast=e8m0
+            )
+            print(f"  {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> ACCEPTED  out={tuple(out.shape)}")
+        except Exception as exc:  # noqa: BLE001
+            first = (str(exc).strip().splitlines() or [type(exc).__name__])[0]
+            print(f"  {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> REJECTED  {first[:100]}")
+
+    print("\nPROBE_DONE")
+
+
+if __name__ == "__main__":
+    main()
--- a/scripts/ref_layer_dump.py
+++ b/scripts/ref_layer_dump.py
@ -114,7 +114,9 @@ def main() -> None:
    with torch.inference_mode():
        output_ids, logits, _ = model.forward(tokens, 0)

-    captured["final_row_logits"] = logits[0, -1].float().cpu().numpy()
+    # The head already collapsed the sequence axis (`x = x[:, -1]` when not full_logits), so the
+    # logits are [b, vocab] -- indexing [0, -1] here would silently store a scalar.
+    captured["final_row_logits"] = logits[0].float().cpu().numpy()
    captured["next_token"] = np.array([int(output_ids[-1])], dtype=np.int32)

    if rank == 0:
--- a/scripts/ref_next_token.py
+++ b/scripts/ref_next_token.py
@ -14,6 +14,7 @@ import json
 import os
 import sys

+import numpy as np
 import torch
 import torch.distributed as dist

@ -30,6 +31,15 @@ def main() -> None:
    parser.add_argument("--ids", required=True, help="comma-separated prompt token ids")
    parser.add_argument("--out", default="/data/deepinfer-h20/ref/next_token.json")
    parser.add_argument("--topk", type=int, default=5)
+    parser.add_argument(
+        "--force_ids",
+        default="",
+        help="Comma-separated ids fed one at a time after the prompt, i.e. incremental decode. "
+        "This mirrors `generate.py`: prefill the prompt, then forward a single token per step with "
+        "prev_pos advancing, so the resulting logits come from the decode path rather than from "
+        "recomputing the whole prefix.",
+    )
+    parser.add_argument("--logits_out", default="", help="optional .npz to store every step's full logits")
    args = parser.parse_args()

    ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
@ -59,21 +69,45 @@ def main() -> None:
    load_model(model, os.path.join(args.ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
    torch.set_default_device("cuda")

+    forced = [int(x) for x in args.force_ids.replace(" ", "").split(",") if x]
+
    tokens = torch.tensor([ids], dtype=torch.long)
+    steps: list[tuple[str, torch.Tensor]] = []
    with torch.inference_mode():
        # `Transformer.forward` returns (output_ids, logits, main_hidden).
        output_ids, logits, _ = model.forward(tokens, 0)
+        steps.append(("prefill", logits[0].float().clone()))

-    row = logits[0, -1].float()
+        # Same loop shape as `generate.py`: one token per forward, prev_pos advancing. The token
+        # fed in is the ground truth, which is exactly what the engine does when it accepts a
+        # greedy token it already agrees with.
+        prev_pos = tokens.shape[1]
+        for token in forced:
+            output_ids, logits, _ = model.forward(torch.tensor([[token]], dtype=torch.long), prev_pos)
+            steps.append((f"decode_{prev_pos}", logits[0].float().clone()))
+            prev_pos += 1
+
+    row = steps[-1][1]
    top = torch.topk(row, args.topk)
-    assert int(top.indices[0]) == int(output_ids[0, -1]), "argmax disagrees with the sampled id"

    payload = {
        "prompt_ids": ids,
+        "force_ids": forced,
        "next_token": int(top.indices[0]),
        "top_tokens": [int(x) for x in top.indices],
        "top_logits": [round(float(x), 4) for x in top.values],
+        "steps": [
+            {
+                "name": name,
+                "top_tokens": [int(x) for x in torch.topk(v, args.topk).indices],
+                "top_logits": [round(float(x), 4) for x in torch.topk(v, args.topk).values],
+            }
+            for name, v in steps
+        ],
    }
+    if args.logits_out and rank == 0:
+        os.makedirs(os.path.dirname(args.logits_out), exist_ok=True)
+        np.savez(args.logits_out, **{f"logits_{name}": v.cpu().numpy() for name, v in steps})
    if rank == 0:
        print(json.dumps(payload, ensure_ascii=False), flush=True)
        os.makedirs(os.path.dirname(args.out), exist_ok=True)