debug: capture final logits, and keep the profile forward out of the dumps

Comparing final logits was impossible before: `ref_layer_dump.py` read `logits[0, -1]`, but the
head already collapsed the sequence axis (`x = x[:, -1]`), so logits are `[b, vocab]` and that
index silently stored a 0-dimensional scalar. Every earlier dump had an unusable
`final_row_logits`.

With that fixed, the comparison reframes the remaining drift. On the same 14 tokens:

    prefill   top-1 = 6127 (北京)   margin(6127 - 666) = +0.256
    reference top-1 = 6127 (北京)   margin             = +3.112

Our prefill picks the right token; the error only compresses the margin, by ~12x. But the
end-to-end run, which reaches that context via decode steps, picks 666 (`**`) instead. Same
context, different path, different answer -- so the remaining drift is in **decode**, not in the
per-128 quantization. The quantization floor is still 2.7e-2 and still irreducible
(`H20-DRIFT-ANALYSIS.md` has the DeepGEMM assertion that forces `gran_k == 128` on SM90), but it
is not what flips the token.

Two traps had to be handled to get numbers that mean anything:
- `execute_model_profile_run` runs a full forward + sampler on synthetic tokens. The layer dumps
  were filtered by their prompt-length gate; logits are not, because their row count is the
  number of *selected* tokens (1). The first attempt captured that step and produced cos = 0.34.
  `debug_dump.profiling()` now suspends capture for the whole profile step.
- With `max_tokens=1` both a prompt and a decode step run, and only the prompt one is comparable
  with a prefill-only reference. They are dumped as `logits_prompt` / `logits_decode`.

Also re-runs the layer comparison now that the expert-scale fix has landed: layer 0 went from
1.27e-1 to 2.88e-2 and layer 39 sits at 2.00e-1, so the error curve flattened by an order of
magnitude.

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 比较。

deepinfer/h20_ops/debug_dump.py

@@ -27,19 +27,47 @@ _CAPTURED: dict[str, np.ndarray] = {}
 _MOE_LAYERS = int(os.getenv("DEEPINFER_DUMP_MOE_LAYERS", "2") or 0)
 
 
+_SUSPENDED = False
+
+
 def enabled() -> bool:
-    return bool(_DUMP_DIR)
+    return bool(_DUMP_DIR) and not _SUSPENDED
 
 
-def record(tag: str, tensor: torch.Tensor) -> None:
+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 tensor.shape[0] != _DUMP_TOKENS:
+    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

deepinfer/model_executor/layers/sampler.py

@@ -12,6 +12,7 @@ import torch.nn as nn
 from deep_select.interface import topk
 
 from deepinfer import envs
+from deepinfer.h20_ops import debug_dump as _debug_dump
 from deepinfer.logger import init_logger
 from deepinfer.model_executor.input_metadata import InputMetadata
 from deepinfer.model_executor.layers.mtp_info_pool import MTPInfoPool
@@ -140,6 +141,12 @@ 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():
+            # 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)
         del hidden_states
 
         if driver_inputs is None:

deepinfer/worker/model_runner.py

@@ -15,6 +15,7 @@ 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
 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.input_metadata import InputMetadata
 from deepinfer.model_executor.layers.mtp_info_pool import get_nextn_pad_token_ids
@@ -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(

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()

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: