test(sm90): rule out the KV cache and the decode indexer

Two closed-loop checks, both clean, which finally retires the two standing suspects.

P4.3 KV cache. The repo's own store/load tests pass for the fp32-scale index layout SM90 uses
(6 each). Running every layout gives 48 passed / 12 failed, and every failure is a ue8m0 index
layout -- `FP4_INDEX_KV_FORMAT` or `MXFP8_INDEX_KV_FORMAT`, neither of which SM90 selects. Those
failures are in the test's reference, not the kernel: it reinterprets ue8m0 bytes as
`PLATFORM.fp8_sf_type`, which is float32 here, so `index_copy_` hits a dtype mismatch.

`scripts/h20_kv_roundtrip.py` then closes the loop without any reference at all -- write known
BF16 through the paged cache, read it back, compare. Both SM90 layouts, block 32 and 64, 1/233/
1029 tokens:

    V41_FORMAT (main/swa KV)              rel 0.0266-0.0295   cos 0.9996
    FP8_FLOAT32_SF_INDEX_KV_FORMAT (132B) rel 0.0266-0.0327   cos 0.9996

That is the fp8 floor, the same order as the 2.7e-2 on weights, so the cache is not where the
decode error comes from and the 132-byte layout question from P3.9 is closed.

One trap: the first version of the script compared raw fp8 values against BF16 and reported
rel=126. Index layouts are deliberately *not* dequantized by `load_paged_kv_cache` -- the
consumer is a kernel that wants bytes plus scales. Multiplying by the scale gives 0.027. The
127x was the test, not the implementation.

P4.4 the decode indexer. `_decode_p_sum_sm90` reimplements DeepGEMM's paged MQA logits in torch
because Hopper's kernel requires 64-token pages; nothing else exercised it. Compared against a
float64 reference that shares no code -- explicit per-token offsets instead of strided views --
it agrees to rel~1e-7 / cos 1.0 across page sizes 32 and 64, boundary-crossing lengths,
multi-page contexts and 32 heads. The `q_sf[..., :1]` dequantization is also sound: head_dim is
128, equal to the granularity, so one scale per (row, head) is correct.

What this changes. `logits_prompt` from the max_tokens=1 run puts the prefill error at cos
0.9846 with an absolute error of 3.30 on the winning token, against a reference margin of 3.112:
the margin is thin because the error is the same size as the margin, not by coincidence. The
reference's own prefill-to-decode transition also costs about 2x margin (3.112 -> 1.638) purely
from the quantized KV round trip. Our 0.256 halved would still be positive, yet it flips, so
roughly 0.5-1.0 of logit-level error in decode remains unexplained -- now narrowed to attention
over the cached main KV and the SWA-window / compressed-position mapping.

Also corrects an earlier overstatement: "the difference is introduced by decode" does not mean
decode has a discrete bug, and P4.3/P4.4 show its two main suspects do not.

One more dump trap worth recording: with max_tokens > 1, `logits_prompt` is contaminated too.
mysteps6 (max_tokens=7) gives cos 0.68 and a 16.5 top-1 error for the same input where mylogits3
(max_tokens=1) gives 0.9846 and 3.30. Prefill logits are only trustworthy from a
max_tokens=1 run.

H20-DEBUG-PLAN.md

@@ -152,7 +152,7 @@ x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim=
 > 后续跟踪（P3.6 起）已迁移到 `H20-PORTING-JOURNAL.md` 第 6 节，那里有完整的
 > 漂移量化数据、建议顺序和前置条件。本文件的 P4 / P5 保留为方法清单。
 
-## P4 decode 专项 —— 已确认 decode 是 `**` 的来源
+## P4 decode 专项 —— 差额确认由 decode 引入，但 decode 的两个嫌疑点都已排除
 
 ### 为什么锁定 decode
 
@@ -188,7 +188,7 @@ prefill 选中了正确的 token，误差只把 margin 压缩了 12 倍。但端
 
 ### 步骤
 
-- [x] **P4.1 同位置 prefill vs decode 对拍** —— **完成，确认 decode 有 bug**
+- [x] **P4.1 同位置 prefill vs decode 对拍** —— **完成，确认差额只在 decode 引入**
 
       四条路径走到同一个 14-token 上下文
       `[0,128803,2056,768,17607,2793,3139,12108,1148,128804,128822,17607,2793,3139]`：
@@ -204,7 +204,9 @@ prefill 选中了正确的 token，误差只把 margin 压缩了 12 倍。但端
       `h20_ask.py`：`--ids <14个> --max_tokens 1` 给 `6127`，`--ids <11个> --max_tokens 7`
       的第 4 个 token 是 `666`。
 
-      **结论：prefill 路径正确，decode 路径错误，差额只由 decode 引入。**
+      **结论：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），所以
@@ -216,28 +218,88 @@ prefill 选中了正确的 token，误差只把 margin 压缩了 12 倍。但端
         每一步的记录就是官方 decode 路径的 logits
       产出：三条曲线的 margin 与 top-1 逐步对照。
 
-- [ ] **P4.2 decode 逐步误差曲线**
+- [~] **P4.2 decode 逐步误差曲线**（受阻）
+      需要一次干净的 `max_tokens=4` 单请求运行，但 GPU 被同事任务占满（`flagos-deepseek`
+      在跑 DeepSeek-R1-1.5B，2 小时），未强停。工具已就绪，见 P4.1 的坑表。
+
       从第 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 验证**。
+- [x] **P4.3 KV cache 写入/读回一致性** —— **完成，干净，排除**
 
-- [ ] **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 选哪些专家。
+      两层验证，都通过。
+
+      其一，仓库自带的单测（`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 路径不是误差来源。**
 
 - [ ] **P4.5 采样器 / logits processor 排除**
       若 P4.3、P4.4 都干净，则查采样侧：`DEEPINFER_SAMPLER_FORCE_TOPK=4096`
       会不会在 topk=4096 时改变数值，以及 logits 是否被额外 mask（blacklist、EOS 抑制）。
 
+### 量化后的真实图景：prefill 本身就已经很紧
+
+把 P4.1 之后的可信数据合起来（`max_tokens=1` 那次运行的 `logits_prompt`，14 token）：
+
+| 量 | 官方 | 移植后 |
+|---|---|---|
+| logits cos | — | **0.9846** |
+| logits rel_err | — | 0.1863 |
+| 绝对误差 mean / p99 / max | — | 0.597 / 1.93 / 4.63 |
+| top-1（6127 北京）logit | 39.343 | 36.047（**-3.30**） |
+| top-2（666 `**`）logit | 36.232 | 35.791（-0.44） |
+| **margin** | **+3.112** | **+0.256** |
+
+**top-1 的绝对误差 3.30 就把 3.11 的 margin 吃光了。** 所以 margin 只剩 0.256 不是
+「恰好接近」，而是误差与 margin 同量级的必然结果。
+
+修正之前的一个说法：P4.1 证明了「差额只在 decode 引入」，但 P4.3/P4.4 证明 KV cache 与
+decode indexer 都精确。把两件事合起来，现在的判断是：
+
+1. **prefill 的 logit 误差（3.3）已经把 margin 压到临界**，主因仍是 per-128 量化的累积；
+2. **decode 只贡献最后一点**（官方自己 prefill→decode 也让 margin 从 3.112 掉到 1.638，
+   约 2 倍，这是 KV cache 量化往返的固有代价）；
+3. 但按 0.256 再掉 2 倍应该还剩 0.13，实际却翻成负 —— **仍有约 0.5~1.0 的 logit 级误差
+   未被解释**。范围已收窄到「decode 时对已缓存 main KV 做 attention」与「SWA window 边界 /
+   compressed 位置映射」这一段，这是 P4.2 / P4.5 要继续查的地方。
+
 ### P4.1 的工具坑（都已修，但 dump 仍建议配合黑盒交叉验证）
 
 | 坑 | 现象 | 修法 |
@@ -250,11 +312,19 @@ prefill 选中了正确的 token，误差只把 margin 压缩了 12 倍。但端
 即便都修了，逐 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
+- `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 异常
 

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

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