debug: add layer-wise dump + reference comparison for H20 bring-up

Narrows the SM90 numerical divergence to layer 0's FFN. The sparse MLA attention
kernel is not the problem: at layer 0 its output matches the reference with
cos=0.99897 / rel_err=4.7e-2, while the full layer output is already at 12.7% and the
error compounds to cos=0.865 by layer 11, which finally flips the top-1 token.

Tooling
- `h20_ops/debug_dump.py`: opt-in capture (`DEEPINFER_DUMP_LAYERS=<dir>`), last
  position only, reduced over `hc_mult`, written to `.npz` at exit.
  `DEEPINFER_DUMP_TOKENS=<n>` gates on the real prompt length, because the engine
  runs a dummy profile forward with a different token count before it allocates the
  KV cache and that dump would otherwise be mistaken for the real one.
- `scripts/ref_layer_dump.py`: the same anchors from the reference model via forward
  hooks (block output, attention sub-layer output, embedding), plus the final
  logits and greedy token.
- `scripts/compare_layers.py`: cosine similarity, relative L2 error and norm ratio
  per layer, reporting the first layer above a tolerance.

Two traps this had to fix before the numbers meant anything
- `decoder_layer` returns the *sub-layer* output; the residual stream lives in
  `residual` (`HCLazyPostData`, hc_mult=4). Comparing without merging them through
  `hc_post` produced a bogus cos of 0.02 at every layer.
- `DEEPINFER_DEBUG_NAN` and the dump both need a host copy inside the forward, so
  both take `allow_cuda_sync()` rather than fighting the engine's sync guard.

Also documents the reproduction case and the findings in H20-DEBUG-PLAN.md.

H20-DEBUG-PLAN.md

@@ -55,21 +55,43 @@ length : max_new_tokens = 6
 
 ---
 
-## P2 分层对拍（下一步，进行中）
+## P2 分层对拍 —— 已完成，定位到第 0 层的 FFN 阶段
 
-输入固定为 14-token 那条（官方 `6127` / 我方 `666`），这是最短的可复现分歧。
+**方法**：`scripts/ref_layer_dump.py`（官方，forward hook）+ `DEEPINFER_DUMP_LAYERS`（我方，
+`h20_ops/debug_dump.py`），两侧都取最后位置、都按 hc_mult 求均值，用
+`scripts/compare_layers.py` 对比。输入固定为那条 14-token ids。
 
-- [ ] P2.1 官方 `model.py` 挂 forward hook，dump 每层 `h`（最后一行）到 `.npy`
-- [ ] P2.2 DeepInfer 侧加同样 dump（`DEEPINFER_DEBUG_DUMP=<dir>`）
-- [ ] P2.3 逐层算 `rel_err = ||x_mine - x_ref|| / ||x_ref||`，找**首个**超 1e-2 的层
+**踩过的两个坑（重要，否则会得出错误结论）**：
 
-**判据**：拿到具体层号。若是第 0 层就已经偏，说明输入侧（embedding / engram）有问题；
-若是中间层开始偏，则集中查该层的 attention 或 MoE。
+1. 我方 `decoder_layer` 返回的 `hidden_states` 只是**子层输出**（2D），完整残差流在 `residual`
+   （`HCLazyPostData`，hc_mult=4）里。必须 `hc_post(hidden_states, residual)` 之后再比，
+   否则第一版对比得到的 cos≈0.02 是假的。
+2. 引擎在分配 KV cache 前会跑一次 **dummy profile forward**，token 数与真实 prompt 不同。
+   必须用 `DEEPINFER_DUMP_TOKENS=<n>` 门控，否则 dump 到的是 dummy 输入。
 
-**注意**：engram 在第 1~2 层就有贡献，而 engram 表是我自己 gather 的，属于高嫌疑区，
-所以「第 0/1 层就偏」的可能性不小。
+**结果（14-token 输入，cos / rel_err / 范数比）**：
 
----
+| 位置 | cos | rel_err | ‖mine‖/‖ref‖ |
+|---|---|---|---|
+| layer0 attention 子层输出 | **0.99897** | **4.71e-02** | 0.986 |
+| layer0 整层输出 | 0.99190 | 1.27e-01 | 1.000 |
+| layer1 整层 | 0.97707 | 2.29e-01 | 1.062 |
+| layer1 attention 子层 | 0.99324 | 1.21e-01 | 0.959 |
+| layer2 | 0.96411 | 2.76e-01 | 1.039 |
+| layer5 | 0.93275 | 3.61e-01 | 0.943 |
+| layer11 | 0.86515 | 5.28e-01 | 1.030 |
+
+**结论**：
+
+1. **稀疏 MLA attention 本身是对的**：第 0 层 attention 子层 cos=0.999，误差 4.7%，
+   在 FP8/量化口径下属于合理量级。之前怀疑的 TileLang kernel 可以排除。
+2. **误差主要在第 0 层的 FFN/MoE 阶段引入**：attention 4.7% → 整层 12.7%，
+   即 FFN 贡献了主要偏差。
+3. **误差单调累积**，到第 11 层 cos 已降到 0.865，最终把 top-1 从「北京」翻成「**」。
+
+**下一步**：在 layer0 的 FFN 内再切一刀 —— 分别 dump `routed_experts_forward` 的输出、
+shared expert 输出、以及 swiglu/route 权重，看是路由（top-6 选中了不同专家）还是
+专家计算本身偏。`engram` 在第 1/14 层，本轮已可排除其为**首个**偏离点。
 
 ## P3 层内细分（按嫌疑排序）
 

deepinfer/h20_ops/debug_dump.py

@@ -0,0 +1,67 @@
+"""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] = {}
+
+
+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 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 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)

deepinfer/model_executor/models/decoder_layer.py

@@ -5,6 +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.logger import init_logger
 from deepinfer.model_executor.input_metadata import InputMetadata
 from deepinfer.model_executor.input_metadata.yoco import YocoMetadataMixin
@@ -179,6 +180,7 @@ class DeepseekDecoderLayer(HCLayerMixin):
                 **attn_kwargs,
             )
             _check_finite(f"layer{self.layer_idx}.attn_out", hidden_states)
+            _debug_dump.record(f"layer{self.layer_idx}_attn", hidden_states)
             hidden_states, hidden_states_bf16, residual = self.ffn_norm_forward(hidden_states, residual)
 
         mlp_timer_scale = self.num_prefill_mlp_layers if input_metadata.is_prompt else self.num_mlp_layers
@@ -193,6 +195,14 @@ class DeepseekDecoderLayer(HCLayerMixin):
             )
 
         _check_finite(f"layer{self.layer_idx}.mlp_out", hidden_states)
+        if _debug_dump.enabled():
+            # `hidden_states` is only this layer's sub-layer output; the residual stream lives in
+            # `residual` and is merged by `hc_post` ([n, hc_mult, hidden]). The reference dumps the
+            # merged stream, so merge it here too -- otherwise the two are not comparable.
+            _debug_dump.record(
+                f"layer{self.layer_idx}",
+                HCLayerMixin.hc_post(hidden_states, residual) if self.hc_enabled else hidden_states,
+            )
         return hidden_states, residual
 
 

deepinfer/model_executor/models/deepseek.py

@@ -10,6 +10,7 @@ from torch import nn
 
 from deepinfer import AttrDict, envs
 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
@@ -160,6 +161,7 @@ class DeepseekModel(nn.Module):
     ) -> CausalLMForwardOutput:
         if hidden_states is None:
             hidden_states = self.embed_tokens(input_metadata.input_ids)
+            _debug_dump.record("embed", hidden_states)
 
         residual = None
         # 跨层持有的运行时值 (shared_hidden / pre_rope_kv / topk): 逐层传同一引用，由各层原地写入。
@@ -237,6 +239,7 @@ class DeepseekModel(nn.Module):
             hidden_states = hidden_states.unflatten(-1, (self.hc_head.hc_mult, -1))
             hidden_states = self.hc_head(hidden_states)
         hidden_states = self.output_norm(hidden_states)
+        _debug_dump.record("final_norm", hidden_states)
 
         if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
             with userbuffer.cp_session():

scripts/compare_layers.py

@@ -0,0 +1,85 @@
+"""Diff per-layer hidden states between DeepInfer and the reference implementation.
+
+Both dumps keep only the final position of the prompt, reduced over the `hc_mult` copies, so
+the comparison is one vector per layer. Reports cosine similarity, relative L2 error and the
+ratio of norms, which together distinguish "slightly off" from "different direction".
+
+    python scripts/compare_layers.py --ref ref_layers.npz --mine deepinfer_layers.npz
+"""
+
+import argparse
+
+import numpy as np
+
+
+def rel_err(a: np.ndarray, b: np.ndarray) -> float:
+    return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-12))
+
+
+def cosine(a: np.ndarray, b: np.ndarray) -> float:
+    na, nb = np.linalg.norm(a), np.linalg.norm(b)
+    if na == 0 or nb == 0:
+        return float("nan")
+    return float(np.dot(a, b) / (na * nb))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ref", required=True)
+    parser.add_argument("--mine", required=True)
+    parser.add_argument("--tol", type=float, default=1e-2, help="first layer above this rel_err is the culprit")
+    args = parser.parse_args()
+
+    ref = np.load(args.ref)
+    mine = np.load(args.mine)
+
+    layers = sorted(int(k[5:]) for k in ref.files if k.startswith("layer") and k[5:].isdigit())
+    mine_layers = {int(k[5:]) for k in mine.files if k.startswith("layer") and k[5:].isdigit()}
+    missing = [i for i in layers if i not in mine_layers]
+    if missing:
+        print(f"WARNING: no dump from DeepInfer for layers {missing}")
+
+    for key in ("embed", "layer0_attn", "layer1_attn"):
+        if key in ref.files and key in mine.files:
+            r, m = ref[key].astype(np.float64), mine[key].astype(np.float64)
+            if r.shape == m.shape:
+                print(
+                    f"{key:>12}  cos={cosine(m, r):9.5f}  rel_err={rel_err(m, r):9.2e}"
+                    f"  |mine|/|ref|={np.linalg.norm(m) / (np.linalg.norm(r) + 1e-12):8.4f}"
+                )
+    print()
+    print(f"{'layer':>5}  {'cos':>9}  {'rel_err':>9}  {'|mine|/|ref|':>12}")
+    first_bad = None
+    for i in layers:
+        if i not in mine_layers:
+            continue
+        r, m = ref[f"layer{i}"].astype(np.float64), mine[f"layer{i}"].astype(np.float64)
+        if r.shape != m.shape:
+            print(f"{i:>5}  shape mismatch ref={r.shape} mine={m.shape}")
+            continue
+        err, cos = rel_err(m, r), cosine(m, r)
+        ratio = float(np.linalg.norm(m) / (np.linalg.norm(r) + 1e-12))
+        flag = ""
+        if first_bad is None and err > args.tol:
+            first_bad = i
+            flag = "   <-- first > tol"
+        print(f"{i:>5}  {cos:9.5f}  {err:9.2e}  {ratio:12.4f}{flag}")
+
+    ref_tok = int(ref["next_token"][0]) if "next_token" in ref.files else None
+    print(f"\nreference next token: {ref_tok}")
+
+    if "final_row_logits" in ref.files:
+        rl = ref["final_row_logits"].astype(np.float64)
+        top = np.argsort(-rl)[:5]
+        print("reference top-5 tokens:", [int(x) for x in top], "logits:", [round(float(rl[x]), 3) for x in top])
+        # DeepInfer's logits are not dumped; derive its choice from the harness output instead.
+
+    print()
+    if first_bad is None:
+        print(f"RESULT: every dumped layer is within rel_err <= {args.tol}; divergence is downstream")
+    else:
+        print(f"RESULT: first layer above rel_err {args.tol} is layer {first_bad}")
+
+
+if __name__ == "__main__":
+    main()

scripts/ref_layer_dump.py

@@ -0,0 +1,105 @@
+"""Dump the reference model's per-layer hidden state for one prompt.
+
+Counterpart to DeepInfer's `DEEPINFER_DUMP_LAYERS`, so the two can be diffed layer by layer.
+The reference residual stream is `hc_mult` parallel copies, so each layer is reduced with
+`mean(dim=1)` -- the same convention `debug_dump.record` uses -- and only the final position
+is kept, because that is the only row generation depends on.
+
+    torchrun --nproc-per-node 8 ref_layer_dump.py --ids 0,128803,... --out /data/deepinfer-h20/ref
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from safetensors.torch import load_model  # noqa: E402
+from transformers import AutoTokenizer  # noqa: E402
+
+from model import ModelArgs, Transformer  # noqa: E402
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt-path", required=True)
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--ids", required=True)
+    parser.add_argument("--out", default="/data/deepinfer-h20/ref")
+    args = parser.parse_args()
+
+    ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
+    rank = int(os.getenv("RANK", "0"))
+    local_rank = int(os.getenv("LOCAL_RANK", "0"))
+    world_size = int(os.getenv("WORLD_SIZE", "1"))
+    torch.cuda.set_device(local_rank)
+    if world_size > 1:
+        dist.init_process_group("nccl")
+
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_num_threads(8)
+    with open(args.config) as f:
+        model_args = ModelArgs(**json.load(f))
+    model_args.temperature = 0.0
+    model_args.max_batch_size = 1
+    model_args.max_seq_len = 4096
+
+    tokenizer = AutoTokenizer.from_pretrained(args.ckpt_path)
+    with torch.device("cuda"):
+        model = Transformer(model_args, tokenizer)
+    load_model(model, os.path.join(args.ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
+    torch.set_default_device("cuda")
+
+    captured: dict[str, np.ndarray] = {}
+
+    def make_hook(layer_id: int):
+        def hook(_module, _inputs, output):
+            # Block.forward returns (x, pre_mix); x is [b, s, hc_mult, dim].
+            x = output[0] if isinstance(output, tuple) else output
+            captured[f"layer{layer_id}"] = x.mean(dim=2)[0, -1].float().cpu().numpy()
+
+        return hook
+
+    def make_attn_hook(layer_id: int):
+        def hook(_module, _inputs, output):
+            # The attention sub-layer output is already collapsed: [b, s, dim].
+            x = output[0] if isinstance(output, tuple) else output
+            captured[f"layer{layer_id}_attn"] = x[0, -1].float().cpu().numpy()
+
+        return hook
+
+    def embed_hook(_module, _inputs, output):
+        captured["embed"] = output[0, -1].float().cpu().numpy()
+
+    model.embed.register_forward_hook(embed_hook)
+    for layer_id, layer in enumerate(model.layers):
+        layer.register_forward_hook(make_hook(layer_id))
+        layer.attn.register_forward_hook(make_attn_hook(layer_id))
+
+    tokens = torch.tensor([ids], dtype=torch.long)
+    with torch.inference_mode():
+        output_ids, logits, _ = model.forward(tokens, 0)
+
+    captured["final_row_logits"] = logits[0, -1].float().cpu().numpy()
+    captured["next_token"] = np.array([int(output_ids[-1])], dtype=np.int32)
+
+    if rank == 0:
+        out = Path(args.out)
+        out.mkdir(parents=True, exist_ok=True)
+        path = out / "ref_layers.npz"
+        np.savez(path, **captured)
+        print(f"wrote {len(captured)} tensors to {path}", flush=True)
+        print(json.dumps({"next_token": int(output_ids[-1]), "layers": len(model.layers)}), flush=True)
+
+    if world_size > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()