From 2ea88a343117784cca3316a59e3eadcf0d105c40 Mon Sep 17 00:00:00 2001 From: loopyt Date: Wed, 21 Oct 2026 20:02:19 +0800 Subject: [PATCH] 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 | 102 ++++++++++++++++++++---- scripts/h20_index_decode_check.py | 127 ++++++++++++++++++++++++++++++ scripts/h20_kv_roundtrip.py | 102 ++++++++++++++++++++++++ 3 files changed, 315 insertions(+), 16 deletions(-) create mode 100644 scripts/h20_index_decode_check.py create mode 100644 scripts/h20_kv_roundtrip.py diff --git a/H20-DEBUG-PLAN.md b/H20-DEBUG-PLAN.md index dfbd227..4506e8f 100644 --- a/H20-DEBUG-PLAN.md +++ b/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 异常 diff --git a/scripts/h20_index_decode_check.py b/scripts/h20_index_decode_check.py new file mode 100644 index 0000000..9cbe320 --- /dev/null +++ b/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() diff --git a/scripts/h20_kv_roundtrip.py b/scripts/h20_kv_roundtrip.py new file mode 100644 index 0000000..b0e9bb1 --- /dev/null +++ b/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()