fix(sm90): expand expert scales on the contraction axis only

`requant_expert_fp4` expanded the UE8M0 scale with `_expand_blocks`, which repeats *both*
axes. That is right for dense weights, whose scales are `(n_out/32, k/32)`, but expert scales
are already per-output-channel: `(n_out, k/32)`. Expanding both axes therefore treated
`(2304, 160)` as `(72, 5)` and handed every one of the 2304 output channels the scales of the
first 72 rows, scrambling each expert's weights.

This is what the layer comparison was pointing at: `shared_expert` matched the reference
(cos 0.9987) because it goes through `requant_dense_fp8`, which uses `_expand_blocks`
correctly, while `combined` -- routed experts plus shared -- was off by 3.6e-1.

Verified against the authoritative implementation, `inference/convert.py`'s
`cast_e2m1fn_to_e4m3fn`, on layer 0 / expert 370:

    w1  cos 0.8936 -> 1.0000    rel_err 5.35e-1 -> 2.84e-2
    w3  cos 0.8960 -> 0.9997    rel_err 5.26e-1 -> 2.69e-2
    w2  cos 0.8964 -> 0.9996    rel_err 5.20e-1 -> 2.65e-2

The residual ~2.7% is the per-32 -> per-128 requantization loss, not a defect; the reference
is lossless because it folds the relative per-32 offset into the FP8 values.

End to end on the 11-token chat prompt:

    before:  [17607, 41175, 389, 666]                     "中国的首都是**"
    after:   [17607, 2793, 3139, 666, 6127, 666, 320, 1]  "中国的首都是**北京**。"
    ref:     [17607, 2793, 3139, 6127, 320]               "中国的首都是北京。"

The answer is now correct and the first three tokens match the reference exactly. The remaining
difference is markdown emphasis (`**`), so some numerical drift still moves the top-1 on the
occasional position.

Also adds two checks that would have caught this immediately:
`scripts/h20_expert_weight_check.py` diffs our requantization against DeepSeek's own converter,
and `scripts/h20_quant_check.py` pins the E2M1 table, nibble order and UE8M0 bias. The nibble
order and table turn out to match the reference; only the scale broadcast was wrong.

H20-DEBUG-PLAN.md

@@ -93,21 +93,61 @@ length : max_new_tokens = 6
 shared expert 输出、以及 swiglu/route 权重，看是路由（top-6 选中了不同专家）还是
 专家计算本身偏。`engram` 在第 1/14 层，本轮已可排除其为**首个**偏离点。
 
-## P3 层内细分（按嫌疑排序）
+## P3 已完成：根因是专家权重量化的 scale 广播错误 —— 已修复，答案正确
 
-- [ ] P3.1 **indexer 选出的 top-k 索引**（离散量，最容易判断）：同一层、同一输入，
-      比较索引集合的重合率。低于 ~90% 就说明 indexer 错了。
-- [ ] P3.2 **index KV 字节布局**：SM90 用 `FP8_FLOAT32_SF_INDEX_KV_FORMAT`
-      （128 fp8 + 1 f32 = 132B/token）。store kernel 是否按「整页 values 后接整页 scales」
-      写、`_decode_p_sum_sm90` 的读法是否一致 —— 用一个 round-trip 测试验证：
-      写入已知向量 → 读出 → 比 bf16 原值。
-- [ ] P3.3 **主 KV cache 的 ue8m0 解量化**：SM90 用 fp32 scale，但 cache 仍存 ue8m0
-      字节，`store_kv_cache` / `load_paged_kv_cache` 两侧口径要对齐。
-- [ ] P3.4 **engram**：表 96M 行、per-32 scale，本地 gather 已单测过，但和官方
-      对拍还没做。
-- [ ] P3.5 **MoE**：路由（sqrtsoftplus + route_scale 1.5）+ swiglu_limit 10.0 是否一致。
+### 根因
 
----
+`h20_ops/quant.py::requant_expert_fp4` 里用了 `_expand_blocks(sf, n_out, k, 32)` 去展开
+UE8M0 scale。但两者的布局前提不同：
+
+| | scale 形状 | 需要展开的维度 |
+|---|---|---|
+| dense 权重（`requant_dense_fp8`） | `(n_out/32, k/32)` | **两个**维度 |
+| 专家权重（`requant_expert_fp4`） | `(n_out, k/32)` | **只有 k 维度** |
+
+`_expand_blocks` 对两个维度都做 `repeat_interleave(32)` 再截断，于是
+`(2304, 160)` 被当成 `(72, 5)` 来展开 —— **全部 2304 个输出通道都拿到了前 72 行的 scale**，
+整个专家权重被打乱。这精确解释了 P2 观察到的「shared 对、routed 错」的非对称：
+shared expert 走 `requant_dense_fp8`（用法正确），routed experts 走这条错误路径。
+
+修复：只沿 contraction 轴展开，并加断言锁住布局。
+
+```python
+assert tuple(sf.shape) == (n_out, k // FP4_BLOCK)
+x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim=1)
+```
+
+### 验证（`scripts/h20_expert_weight_check.py`，与官方 `inference/convert.py`
+### 的 `cast_e2m1fn_to_e4m3fn` 逐元素对拍）
+
+| 权重 | 修复前 cos / rel_err | 修复后 cos / rel_err |
+|---|---|---|
+| w1 | 0.8936 / 5.35e-1 | **1.0000 / 2.84e-2** |
+| w3 | 0.8960 / 5.26e-1 | **0.9997 / 2.69e-2** |
+| w2 | 0.8964 / 5.20e-1 | **0.9996 / 2.65e-2** |
+
+剩余 2.7% 是 per-32 → per-128 的量化损失（FP8 e4m3 三位尾数，理论量级 ~3%），
+不是 bug。官方 `cast_e2m1fn_to_e4m3fn` 是**无损**的（把 per-32 的相对 offset 折进 FP8 数值），
+如果要进一步压掉这 2.7% 可以照搬，见下方 P3.6。
+
+### 端到端结果
+
+```
+修复前:  [17607, 41175, 389, 666]                    -> 中国的首都是**
+修复后:  [17607, 2793, 3139, 666, 6127, 666, 320, 1] -> 中国的首都是**北京**。
+官方:    [17607, 2793, 3139, 6127, 320]              -> 中国的首都是北京。
+```
+
+**语义已正确**（「北京」），前 3 个 token 与官方逐 token 一致。差异只剩 markdown 加粗
+`**`（我方多出 666/666 两个 token），说明仍有残余数值差异让 top-1 在个别位置偏移。
+
+### 下一步（可选，按收益排序）
+
+- [ ] **P3.6 改用官方无损重化**：照 `cast_e2m1fn_to_e4m3fn` 的做法，把 per-32 e8m0 相对
+      offset 折进 FP8、scale 保持 e8m0 per-32，理论上可把 2.7% 降到 0。需要确认 SM90
+      DeepGEMM 是否接受 e8m0 per-32 的 expert scale（当前配方是 fp32 per-128）。
+- [ ] **P3.7 重跑 P2 分层对拍**：确认 layer0 整层 rel_err 从 1.27e-1 降到 ~3e-2 量级。
+- [ ] **P3.8 `**` 的来源**：定位是哪个位置翻的 top-1，判断是否还有别的量化口径差异。
 
 ## P4 decode 专项（仅当 prefill 对、decode 错）
 

deepinfer/h20_ops/debug_dump.py

@@ -22,6 +22,9 @@ _DUMP_DIR = os.getenv("DEEPINFER_DUMP_LAYERS")
 # 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:
@@ -38,6 +41,8 @@ def record(tag: str, tensor: torch.Tensor) -> None:
         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)
@@ -49,6 +54,33 @@ def record(tag: str, tensor: torch.Tensor) -> None:
         _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:

deepinfer/h20_ops/quant.py

@@ -81,7 +81,12 @@ def requant_expert_fp4(
     """
     n_out, half_k = packed.shape
     k = half_k * 2
-    x = _unpack_fp4(packed) * _expand_blocks(sf.to(torch.float32), n_out, k, FP4_BLOCK)
+    # Expert scales are per-output-channel already: shape `(n_out, k / 32)`. Only the contraction
+    # axis is blocked, so only that axis gets expanded -- running this through `_expand_blocks`
+    # (which repeats *both* axes) would hand every output channel the scales of the first
+    # `n_out / 32` rows, scrambling the whole expert.
+    assert tuple(sf.shape) == (n_out, k // FP4_BLOCK), f"unexpected expert scale layout: {tuple(sf.shape)}"
+    x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim=1)
     q, scale = _requant(x, out_block)
     if contraction_major_sf:
         scale = scale.t().contiguous()

deepinfer/model_executor/models/v4_modules/mlp.py

@@ -22,6 +22,7 @@ from torch import nn
 
 from deepinfer import PLATFORM, 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
@@ -623,6 +624,14 @@ class DeepseekMoE(nn.Module):
             combine_pre_evt = self.ep_buffer.capture()
             with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu):
                 shared_out = self.shared_forward(hidden_states)
+            if _debug_dump.enabled():
+                # Capture the routing decision as well as each branch's contribution: the ids are
+                # a discrete quantity, so a set mismatch proves routing is wrong while an equal
+                # set with different outputs points at the expert weights instead.
+                _debug_dump.record_ids(f"layer{self.layer_idx}_topk_idx", topk_idx)
+                _debug_dump.record(f"layer{self.layer_idx}_topk_w", topk_weights)
+                _debug_dump.record(f"layer{self.layer_idx}_shared_out", shared_out)
+                _debug_dump.record(f"layer{self.layer_idx}_moe_out", moe_out_hidden)
             event_overlap = self.combine(
                 moe_out_hidden,
                 shared_out,
@@ -630,6 +639,8 @@ class DeepseekMoE(nn.Module):
                 combine_pre_evt,
             )
             moe_combined_hidden, _combined_topk_weights = event_overlap.current_stream_wait()
+            if _debug_dump.enabled():
+                _debug_dump.record(f"layer{self.layer_idx}_combined", moe_combined_hidden)
 
         except Exception:
             logger.exception(f"[rank {PG.get_rank()}] Exception during MoE forward, topk_idx.shape = {topk_idx.shape}")

scripts/compare_layers.py

@@ -39,6 +39,24 @@ def main() -> None:
     if missing:
         print(f"WARNING: no dump from DeepInfer for layers {missing}")
 
+    # Routing first: expert ids are discrete, so a set mismatch is a definitive verdict and the
+    # weights/branch comparisons below only make sense once routing is known to agree.
+    if args.ref and args.mine:
+        for layer in range(4):
+            key = f"layer{layer}_topk_idx"
+            if key not in ref.files or key not in mine.files:
+                continue
+            r_idx, m_idx = ref[key].astype(np.int64), mine[key].astype(np.int64)
+            r_set = set(r_idx.reshape(-1).tolist())
+            m_set = set(m_idx.reshape(-1).tolist())
+            inter = r_set & m_set
+            print(
+                f"{key:>18}  ref={sorted(r_set)}\n{'':>18}  mine={sorted(m_set)}"
+                f"\n{'':>18}  overlap={len(inter)}/{len(r_set)}"
+                f"  {'IDENTICAL' if r_idx.shape == m_idx.shape and (r_idx == m_idx).all() else 'DIFFERS'}"
+            )
+        print()
+
     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)
@@ -74,6 +92,15 @@ def main() -> None:
         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()
+    for key in sorted(k for k in ref.files if k.endswith(("_shared_out", "_combined", "_moe_out"))):
+        if key in mine.files and ref[key].shape == mine[key].shape:
+            r, m = ref[key].astype(np.float64), mine[key].astype(np.float64)
+            print(
+                f"{key:>18}  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()
     if first_bad is None:
         print(f"RESULT: every dumped layer is within rel_err <= {args.tol}; divergence is downstream")

scripts/h20_expert_weight_check.py

@@ -0,0 +1,119 @@
+"""Compare our FP4->FP8 expert requantization against DeepSeek's own reference conversion.
+
+`inference/convert.py::cast_e2m1fn_to_e4m3fn` is the authoritative implementation: it is what
+produced the TP8 checkpoint that the reference model actually runs on. Two things are worth
+checking there, and both matter because a wrong answer scrambles every routed expert while
+leaving the (already-FP8) shared expert intact -- exactly the asymmetry the layer comparison
+showed.
+
+1. Equality of the *decoded* weights. Both paths should represent the same real values; any
+   difference is accuracy we are throwing away.
+2. The scale grid. The reference keeps a per-32 e8m0 scale (relative offsets folded into the FP8
+   values), while the SM90 path re-quantizes on a per-128 fp32 grid. That is a legitimate design
+   difference only if the magnitude error stays small.
+
+    python scripts/h20_expert_weight_check.py --layer 0 --expert 370
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+from deepinfer.h20_ops.quant import requant_expert_fp4, ue8m0_to_fp32
+
+# Copied verbatim from inference/convert.py so the comparison uses the reference algorithm
+# rather than a paraphrase of it.
+FP4_TABLE = torch.tensor(
+    [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
+    dtype=torch.float32,
+)
+
+
+def cast_e2m1fn_to_e4m3fn(x: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """DeepSeek's lossless FP4 -> FP8 cast (inference/convert.py)."""
+    assert x.dtype == torch.int8 and x.ndim == 2
+    out_dim, in_dim = x.size()
+    in_dim *= 2
+    fp8_block_size = 32
+    fp4_block_size = 32
+    assert in_dim % fp8_block_size == 0 and out_dim % fp8_block_size == 0
+    assert scale.size(0) == out_dim and scale.size(1) == in_dim // fp4_block_size
+
+    x = x.view(torch.uint8)
+    low = x & 0x0F
+    high = (x >> 4) & 0x0F
+    x = torch.stack([FP4_TABLE[low.long()], FP4_TABLE[high.long()]], dim=-1).flatten(2)
+
+    MAX_OFFSET_BITS = 6
+    b_out = out_dim // fp8_block_size
+    b_in = in_dim // fp8_block_size
+    x = x.view(b_out, fp8_block_size, b_in, fp8_block_size).transpose(1, 2)
+    scale = scale.float().view(b_out, fp8_block_size, b_in, -1).transpose(1, 2).flatten(2)
+    scale_max_offset_bits = scale.amax(dim=-1, keepdim=True) / (2**MAX_OFFSET_BITS)
+    offset = scale / scale_max_offset_bits
+    offset = offset.unflatten(-1, (fp8_block_size, -1)).repeat_interleave(fp4_block_size, dim=-1)
+    x = (x * offset).transpose(1, 2).reshape(out_dim, in_dim)
+    return x.to(torch.float8_e4m3fn), scale_max_offset_bits.squeeze(-1).to(torch.float8_e8m0fnu)
+
+
+def read_tensor(hf_dir: Path, index: dict, name: str) -> torch.Tensor:
+    shard = index["weight_map"][name]
+    with safe_open(str(hf_dir / shard), framework="pt", device="cpu") as f:
+        return f.get_tensor(name)
+
+
+def dequant_ref(fp8: torch.Tensor, sf_e8m0: torch.Tensor, block: int = 32) -> torch.Tensor:
+    """Reference layout: (out, in) FP8, scale (out/32, in/32) e8m0 on a 32x32 grid."""
+    s = ue8m0_to_fp32(sf_e8m0)
+    s = s.repeat_interleave(block, 0).repeat_interleave(block, 1)[: fp8.shape[0], : fp8.shape[1]]
+    return fp8.float() * s
+
+
+def dequant_ours(fp8: torch.Tensor, sf_fp32: torch.Tensor, block: int = 128) -> torch.Tensor:
+    """SM90 layout: (out, in) FP8, scales (in/block, out/block) fp32, contraction-major."""
+    s = sf_fp32.t().repeat_interleave(block, 0).repeat_interleave(block, 1)[: fp8.shape[0], : fp8.shape[1]]
+    return fp8.float() * s
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--hf-dir", default="/data/models/DeepSeek-V4.1-Flash-1017")
+    parser.add_argument("--layer", type=int, default=0)
+    parser.add_argument("--expert", type=int, default=370)
+    args = parser.parse_args()
+
+    hf_dir = Path(args.hf_dir)
+    index = json.loads((hf_dir / "model.safetensors.index.json").read_text())
+    prefix = f"layers.{args.layer}.ffn.experts.{args.expert}."
+
+    for proj in ("w1", "w3", "w2"):
+        packed = read_tensor(hf_dir, index, prefix + proj + ".weight")
+        scale = read_tensor(hf_dir, index, prefix + proj + ".scale")
+        print(f"\n{proj}: packed {tuple(packed.shape)} {packed.dtype}, scale {tuple(scale.shape)} {scale.dtype}")
+
+        ref_fp8, ref_sf = cast_e2m1fn_to_e4m3fn(packed, scale)
+        ours_fp8, ours_sf = requant_expert_fp4(packed, scale, out_block=128)
+
+        # `requant_expert_fp4` returns (n_out, k) FP8 and a contraction-major scale; the reference
+        # returns (out, in) too, so the value grids line up directly.
+        ref_vals = dequant_ref(ref_fp8, ref_sf)
+        ours_vals = dequant_ours(ours_fp8, ours_sf)
+
+        assert ref_vals.shape == ours_vals.shape, f"{ref_vals.shape} vs {ours_vals.shape}"
+        num = (ours_vals - ref_vals).norm().item()
+        den = ref_vals.norm().item()
+        cos = torch.nn.functional.cosine_similarity(
+            ours_vals.reshape(-1), ref_vals.reshape(-1), dim=0
+        ).item()
+        print(f"  ref scale: e8m0 {tuple(ref_sf.shape)}   ours: fp32 {tuple(ours_sf.shape)}")
+        print(f"  |ref|={den:.4f}  |ours|={ours_vals.norm().item():.4f}  ratio={ours_vals.norm().item() / den:.5f}")
+        print(f"  rel_err={num / den:.5f}  cos={cos:.7f}")
+
+    print("\nEXPERT_WEIGHT_CHECK_DONE")
+
+
+if __name__ == "__main__":
+    main()

scripts/h20_quant_check.py

@@ -0,0 +1,105 @@
+"""Check `h20_ops.quant` against PyTorch's native FP4/FP8 semantics.
+
+The SM90 path re-quantizes packed E2M1 experts into FP8 while loading, and that code makes three
+assumptions the checkpoint does not spell out: which nibble holds the even element, the E2M1 code
+table, and the UE8M0 bias. If any of them is wrong every expert weight is scrambled, which is
+exactly the symptom the layer comparison shows (shared expert matches, routed experts do not).
+
+    python scripts/h20_quant_check.py
+"""
+
+import torch
+
+from deepinfer.h20_ops.quant import _E2M1_VALUES, _unpack_fp4, requant_expert_fp4, ue8m0_to_fp32
+
+# PyTorch's FP4 cast kernel does not exist for SM90 (device-side assert on H20), and the decode
+# under test is pure integer/table work, so everything here runs on the CPU.
+DEV = "cpu"
+
+
+def check_e2m1_table() -> None:
+    """Every 4-bit code must decode to the value PyTorch assigns it."""
+    codes = torch.arange(16, dtype=torch.uint8, device=DEV)
+    # PyTorch's own decoder: pack two codes per byte and view as the FP4 dtype.
+    packed = torch.stack([codes, torch.zeros_like(codes)], dim=-1).reshape(-1).contiguous()
+    ref = packed.view(torch.float4_e2m1fn_x2).to(torch.float32)[0::2]
+    table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=DEV)
+    ok = torch.equal(ref, table)
+    print(f"E2M1 table: {'OK' if ok else 'MISMATCH'}")
+    if not ok:
+        for code in range(16):
+            if ref[code] != table[code]:
+                print(f"  code {code:04b}: torch={ref[code].item():+.1f} ours={table[code].item():+.1f}")
+        raise SystemExit("E2M1 table disagrees with PyTorch")
+
+
+def check_nibble_order() -> None:
+    """`_unpack_fp4` must interleave in the same order as the packed dtype."""
+    # Distinct, exactly representable values so an order swap is unambiguous.
+    values = torch.tensor([1.0, 2.0, 4.0, 6.0, 0.5, 1.5, 3.0, -1.0], dtype=torch.float32, device=DEV)
+    packed = values.to(torch.float4_e2m1fn_x2).view(torch.uint8)  # (4,) bytes
+    ref = packed.view(torch.float4_e2m1fn_x2).to(torch.float32)
+    ours = _unpack_fp4(packed.view(torch.int8)).reshape(-1)
+    ok = torch.equal(ref, ours)
+    print(f"nibble order: {'OK' if ok else 'MISMATCH'}")
+    if not ok:
+        print(f"  torch: {ref.tolist()}")
+        print(f"  ours : {ours.tolist()}")
+        swapped = ours.view(-1, 2).flip(-1).reshape(-1)
+        print(f"  ours (nibbles swapped): {swapped.tolist()}")
+        print(f"  -> swapping would {'fix it' if torch.equal(ref, swapped) else 'NOT fix it'}")
+        raise SystemExit("nibble order disagrees with PyTorch")
+
+
+def check_ue8m0() -> None:
+    """`ue8m0_to_fp32` must match the checkpoint's scale dtype."""
+    bits = torch.tensor([127, 128, 126, 130, 124], dtype=torch.uint8, device=DEV)
+    scales = bits.view(torch.float8_e8m0fnu)
+    ours = ue8m0_to_fp32(scales)
+    ref = torch.pow(2.0, bits.to(torch.float32) - 127)
+    ok = torch.equal(ours, ref)
+    print(f"UE8M0 bias: {'OK' if ok else 'MISMATCH'}  ours={ours.tolist()} ref={ref.tolist()}")
+    if not ok:
+        raise SystemExit("UE8M0 decode disagrees")
+
+
+def check_requant_roundtrip() -> None:
+    """End to end: dequantize FP4 on the CPU, requantize, and measure the error."""
+    torch.manual_seed(0)
+    n_out, k = 256, 512
+    x = torch.randn(n_out, k, device=DEV, dtype=torch.float32) * 0.05
+
+    # Quantize to FP4 with the checkpoint's per-32 UE8M0 grid, so the round trip starts from a
+    # realistic representation rather than an exact float. FP8 casts also stay on the CPU here.
+    groups = x.view(n_out, k // 32, 32)
+    amax = groups.abs().amax(dim=-1).clamp_min(torch.finfo(torch.float32).tiny)
+    exp = torch.ceil(torch.log2(amax / 6.0)).clamp(-127, 127)
+    scale = torch.pow(2.0, exp)
+    q4 = (groups / scale.unsqueeze(-1)).to(torch.float4_e2m1fn_x2)
+    packed = q4.view(torch.int8)
+    sf_bits = (exp.to(torch.int32) + 127).to(torch.uint8)
+    sf = sf_bits.view(torch.float8_e8m0fnu)
+
+    deq_ref = q4.to(torch.float32) * scale.repeat_interleave(32, dim=1)
+
+    weight, sfb = requant_expert_fp4(packed, sf, out_block=128)
+    assert weight.dtype == torch.float8_e4m3fn, weight.dtype
+    assert sfb.dtype == torch.float32, sfb.dtype
+    assert sfb.shape == (k // 128, n_out // 128), f"scale layout {tuple(sfb.shape)}"
+
+    deq_ours = weight.to(torch.float32) * sfb.t().repeat_interleave(128, 0).repeat_interleave(128, 1)[:n_out, :k]
+    err = (deq_ours - deq_ref).norm() / deq_ref.norm()
+    print(f"requant roundtrip: rel_err={err.item():.4f} (fp8 e4m3 floor ~0.01, fp4 source itself is ~0.1)")
+    assert err < 0.05, "requantization lost more precision than FP8 should"
+
+
+def main() -> None:
+    check_e2m1_table()
+    check_nibble_order()
+    check_ue8m0()
+    check_requant_roundtrip()
+    print("QUANT_CHECK_DONE")
+
+
+if __name__ == "__main__":
+    main()

scripts/ref_layer_dump.py

@@ -32,6 +32,7 @@ def main() -> None:
     parser.add_argument("--config", required=True)
     parser.add_argument("--ids", required=True)
     parser.add_argument("--out", default="/data/deepinfer-h20/ref")
+    parser.add_argument("--moe_layers", type=int, default=2, help="dump MoE internals for this many layers")
     args = parser.parse_args()
 
     ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
@@ -77,10 +78,36 @@ def main() -> None:
     def embed_hook(_module, _inputs, output):
         captured["embed"] = output[0, -1].float().cpu().numpy()
 
+    def make_gate_hook(layer_id: int):
+        def hook(_module, _inputs, output):
+            # Gate.forward returns (weights [n, topk], indices [n, topk]).
+            weights, indices = output
+            captured[f"layer{layer_id}_topk_idx"] = indices[-1].to(torch.int32).cpu().numpy()
+            captured[f"layer{layer_id}_topk_w"] = weights[-1].float().cpu().numpy()
+
+        return hook
+
+    def make_shared_hook(layer_id: int):
+        def hook(_module, _inputs, output):
+            captured[f"layer{layer_id}_shared_out"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()
+
+        return hook
+
+    def make_moe_hook(layer_id: int):
+        def hook(_module, _inputs, output):
+            # The MoE output already includes the shared expert, matching our `_combined`.
+            captured[f"layer{layer_id}_combined"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()
+
+        return hook
+
     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))
+        if layer_id < args.moe_layers:
+            layer.ffn.gate.register_forward_hook(make_gate_hook(layer_id))
+            layer.ffn.shared_experts.register_forward_hook(make_shared_hook(layer_id))
+            layer.ffn.register_forward_hook(make_moe_hook(layer_id))
 
     tokens = torch.tensor([ids], dtype=torch.long)
     with torch.inference_mode():