debug: per-step decode dumps, and prove the flip is decode-only

Adds `record_step` so every decode step is kept instead of just the first, a `--force_ids` mode
to `ref_next_token.py` that mirrors `generate.py`'s incremental loop, and four fixes to the dump
path without which the data was unusable:

- filenames carry `_rank<R>_pid<P>`: with data parallelism every engine is its own
  single-process group and reports RANK=0, so eight processes wrote the same file and the last
  one won;
- decode reads row 0, not the last row: the batch is padded and the padding rows hold logits for
  placeholder tokens, so the last row was a constant `top1=69, margin=0.0528`;
- only the driver rank records: a follower only joins the embedding all-gather, so its logits are
  placeholders that got recorded dozens of times and buried the real steps;
- `DEEPINFER_DUMP_TOKENS` gating plus `profiling()` now cover both profile paths.

The conclusion itself comes from a black-box four-way comparison rather than the dumps, because
per-step dumps still do not line up with the emitted token sequence (`max_tokens=7` produced 75
sampler calls). At the same 14-token context:

    reference prefill   6127 (北京)   correct
    reference decode    6127 (北京)   correct
    our prefill         6127 (北京)   correct
    our decode           666 (**)     wrong

So prefill is right and the remaining flip is introduced by decode alone. Note the reference's
own prefill and decode margins differ (3.112 vs 1.638), so the two paths are not expected to be
identical -- KV cache goes through a quantized round trip -- but they should not disagree on the
direction.

Remaining suspects are therefore narrowed to decode-only state: the KV cache round trip (P4.3)
and the indexer scored in torch versus the TileLang prefill kernel (P4.4), which have never
been diffed against each other.

H20-DEBUG-PLAN.md

@@ -152,13 +152,109 @@ 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 专项 —— 已确认 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` 有完整数据）：
+
+| | 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 路径，而且 margin 只剩 0.256，decode 稍微再加一点误差就会翻。
+
+### 判据（先定死，避免事后解释）
+
+设第 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 对拍** —— **完成，确认 decode 有 bug**
+
+      四条路径走到同一个 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 路径错误，差额只由 decode 引入。**
+      这把 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 逐步对照。
+
+- [ ] **P4.2 decode 逐步误差曲线**
+      从第 1 步到第 8 步，每步记 `margin`、`top-1 是否与官方一致`、`rel(B_t, R_t)`。
+      用来判断误差是「一开始就有」还是「某一步突然放大」——后者意味着踩到了某个
+      batch/位置相关的边界（例如 page 边界、SWA window 回绕）。
+
+- [ ] **P4.3 KV cache 写入/读回一致性**（最高嫌疑之一）
+      写入已知向量后立刻读回，与 bf16 原值比。decode 依赖 cache 能正确回读，
+      而 prefill 的 KV 是当场算的、几乎不读 cache。若这一步就错，P4.1 的差距可解释。
+      同时验证 index KV 的字节布局（`FP8_FLOAT32_SF_INDEX_KV_FORMAT`：128 fp8 + 1 f32
+      = 132 B/token），这是 P3.9 遗留项，**从未做过 round-trip 验证**。
+
+- [ ] **P4.4 `_decode_p_sum_sm90` vs prefill 的 indexer 分数**（最高嫌疑之二）
+      decode 的 indexer 走 torch 兜底（SM90 paged kernel 要求 64-token page，cache 是 128），
+      prefill 走 TileLang 的 `fp8_fp4_mqa_logits`。**两条路径从未对拍过。**
+      给同一个 KV cache、同一个位置，比较两边算出的 top-k 索引集合与分数。
+      P3 已证明「路由 top-k 一致」时误差仍存在，所以这里要查的是**稀疏 attention 选哪些
+      KV**，而不是 MoE 选哪些专家。
+
+- [ ] **P4.5 采样器 / logits processor 排除**
+      若 P4.3、P4.4 都干净，则查采样侧：`DEEPINFER_SAMPLER_FORCE_TOPK=4096`
+      会不会在 topk=4096 时改变数值，以及 logits 是否被额外 mask（blacklist、EOS 抑制）。
+
+### 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 只作辅助。
+
+### 已备好的工具
+
+- `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
 
 ## P5 logprob 异常
 

deepinfer/h20_ops/debug_dump.py

@@ -25,6 +25,12 @@ _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
@@ -82,6 +88,35 @@ def record(tag: str, tensor: torch.Tensor, once: bool = False) -> None:
         _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:
@@ -109,13 +144,34 @@ def _want_moe(tag: str) -> bool:
         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 / "deepinfer_layers.npz"
+    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()

deepinfer/model_executor/layers/sampler.py

@@ -141,12 +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)
-        if _debug_dump.enabled():
+        # `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.
-            _debug_dump.record("logits_prompt" if is_prompt else "logits_decode", logits, once=True)
+            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:

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)