diff --git a/H20-DEBUG-PLAN.md b/H20-DEBUG-PLAN.md index c307a69..df65624 100644 --- a/H20-DEBUG-PLAN.md +++ b/H20-DEBUG-PLAN.md @@ -152,20 +152,249 @@ x = _unpack_fp4(packed) * sf.to(torch.float32).repeat_interleave(FP4_BLOCK, dim= > 后续跟踪（P3.6 起）已迁移到 `H20-PORTING-JOURNAL.md` 第 6 节，那里有完整的 > 漂移量化数据、建议顺序和前置条件。本文件的 P4 / P5 保留为方法清单。 -## P4 decode 专项（仅当 prefill 对、decode 错） +## P4 decode 专项 —— 完成：`**` 已定量解释为量化噪声，非缺陷 -- [ ] P4.1 KV cache 写入后立刻读回，比 bf16 -- [ ] P4.2 同一位置，decode 路径 vs prefill 路径重算的 logits 差异 -- [ ] P4.3 `_decode_p_sum_sm90` 的数值 vs prefill 的 `fp8_fp4_mqa_logits` +### 为什么锁定 decode ---- +P3 之后重跑对拍（`H20-DRIFT-ANALYSIS.md` 有完整数据）： -## P5 logprob 异常 +| | top-1 | margin(6127 北京 − 666 `**`) | +|---|---|---| +| 官方，14 token prefill | 6127 | **+3.112** | +| 我方，14 token prefill | **6127** ✅ | **+0.256** | + +prefill 选中了正确的 token，误差只把 margin 压缩了 12 倍。但端到端跑 +（11 token 起，增量 decode 到第 14 个位置）给的是 666。 + +**同一个 14-token 上下文，一次 prefill 给 6127，增量 decode 给 666。** + +当时的推断是「剩余漂移在 decode 路径」。P4.2~P4.4 做完后，这个推断需要精确化为： +**差额确实只在 decode 出现，但原因不是 decode 有 bug，而是 decode 经过量化 KV cache +往返、噪声略增，恰好碰上模型自己在那个位置只有 1.64 的 margin。** + +### 判据（先定死，避免事后解释；P4.2 之后已修正） + +> **事后修正**：下面这套判据假设「decode 与 prefill 应当基本一致」，因此把 +> `rel(A,B) > 3e-2` 判为「decode 自身缺陷」。P4.2 证明这个前提错了 —— 官方自己 +> prefill→decode 也让 margin 变化约 2 倍，因为 KV cache 走量化往返。 +> 真正有效的判据是**Δmargin 的量级 vs 该位置的 margin**，而不是 rel_err 的绝对值。 +> 判据保留在下面作为过程记录，最终结论见「最终结论」一节。 + + +设第 14 个位置上： + +- `A` = 我方 prefill 路径的 logits +- `B` = 我方 decode 路径的 logits +- `R` = 官方 decode 路径的 logits + +| 量 | 含义 | 阈值 | +|---|---|---| +| `rel(A, B)` | 我方两条路径自身是否一致 | > 3e-2 即为 decode 自身缺陷 | +| `rel(A, R)` | prefill 残余误差 | 已知约 3e-2（量化地板） | +| `rel(B, R)` | decode 总误差 | 明显大于 `rel(A,R)` 说明 decode 额外引入误差 | + +- 若 `rel(A,B)` 很小而 `rel(B,R)` 大 → decode 没有额外缺陷，是量化误差累积，`**` 属可接受 +- 若 `rel(A,B)` 大 → **decode 有 bug**，按 P4.3 → P4.4 往下查 + +### 步骤 + +- [x] **P4.1 同位置 prefill vs decode 对拍** —— **完成（结论已被 P4.2 修正，见下）** + + 四条路径走到同一个 14-token 上下文 + `[0,128803,2056,768,17607,2793,3139,12108,1148,128804,128822,17607,2793,3139]`： + + | 路径 | 下一个 token | 结果 | + |---|---|---| + | 官方 prefill（一次喂 14 个） | 6127 = 北京 | ✅ | + | 官方 decode（喂 11 + 逐 token forward） | 6127 = 北京 | ✅ | + | **我方 prefill（一次喂 14 个）** | **6127 = 北京** | ✅ | + | **我方 decode（喂 11 + 增量）** | **666 = `**`** | ❌ | + + 前两条来自 `ref_next_token.py`（prefill 一次 / `--force_ids` 增量），后两条来自 + `h20_ask.py`：`--ids <14个> --max_tokens 1` 给 `6127`，`--ids <11个> --max_tokens 7` + 的第 4 个 token 是 `666`。 + + **结论：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），所以 + 「两条路径不完全相等」是框架固有特性（KV cache 量化往返），但**方向不该翻**。 + - 路径 1：直接 14 token prefill（`--ids <14个>`，dump `logits_prompt`）—— 已有 + - 路径 2：11 token prefill + 3 步贪心 decode（贪心本来就给对了 + 17607/2793/3139，不需要强制喂 token），dump 第 3 步 decode 的 logits + - 参考侧同样跑增量：`ref_next_token.py` 逐步喂 11→12→13→14 个 token， + 每一步的记录就是官方 decode 路径的 logits + 产出：三条曲线的 margin 与 top-1 逐步对照。 + +- [x] **P4.2 decode 逐步误差曲线** —— **完成，结论收敛：没有 decode 专属 bug** + + 关键前提：**必须 `DEEPINFER_USE_NEXTN_DECODE=0`**。它默认是 `5`，MTP 投机解码默认开启， + 每步产生多行 draft/verify，`logits[0]` 不再对应你以为的那个位置（我第一次跑就是因此 + 错位，把 pos15 的向量当成 pos14 的，白折腾了一轮）。关掉之后一步一行，完全对齐。 + + 11 token 起，逐步对照（`nomtp` run，8 个 decode step）： + + | pos | 我方 top1 | 官方 top1 | 我方 margin | 官方 margin | Δmargin | 翻转 | + |---|---|---|---|---|---|---| + | 11（prefill） | 17607 | 17607 | 10.695 | 9.861 | +0.835 | | + | 12 | 2793 | 2793 | 8.448 | 10.328 | −1.880 | | + | 13 | 3139 | 3139 | 15.276 | 11.946 | +3.330 | | + | **14** | **666** | **6127** | **1.480** | **1.638** | −0.157 | **翻转** | + | 15 | 6127 | 6127 | 15.950 | 15.109 | +0.842 | | + | 16 | 666 | 666 | 20.118 | 17.874 | +2.244 | | + | 17 | 320 | 320 | 9.829 | 10.402 | −0.573 | | + | 18 | 1 | 1 | 12.262 | 11.091 | +1.171 | | + + **Δmargin：mean +0.711，std 1.628，|max| 3.330。** 7/8 位置 top-1 正确， + 唯一的翻转在 pos14 —— 而那里的**官方 margin 只有 1.638**，其余位置是 9.9~17.9。 + + 换句话说：**噪声幅度（std 1.63）和那个位置的 margin（1.64）一样大，是硬币。** + 这不是 decode 的 bug，是量化噪声碰上模型自己也没把握的位置。 + + 官方在 pos14 选了 `6127`（北京），我方选了 `666`（`**`），于是渲染成 + `**北京**` —— 语义完全等价，只是 markdown 加粗的位置不同。 + + 顺带排除：MTP 不是原因。`DEEPINFER_USE_NEXTN_DECODE=0` 前后输出完全相同 + （`[17607, 2793, 3139, 666, 6127, 666, 320, 1]`）。而官方 `generate.py` 本来 + 就没有 MTP，所以把 MTP 关掉才是与参考对等的比法。 + + 从第 1 步到第 8 步，每步记 `margin`、`top-1 是否与官方一致`、`rel(B_t, R_t)`。 + 用来判断误差是「一开始就有」还是「某一步突然放大」——后者意味着踩到了某个 + batch/位置相关的边界（例如 page 边界、SWA window 回绕）。 + +- [x] **P4.3 KV cache 写入/读回一致性** —— **完成，干净，排除** + + 两层验证，都通过。 + + 其一，仓库自带的单测（`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 路径不是误差来源。** + +- [x] **P4.5 采样器 / logits processor 排除** —— **无需再查** + + pos14 我方 top1=666、top2=6127，差值 1.480。如果 6127 被任何 mask / blacklist 抑制， + 它就不会是第二名。所以 pos14 的差异来自 logits 本身，不是采样后处理。 + （`token_logprobs` 恒为 0 的 P5 问题仍然独立存在，与本题无关。） + +### 最终结论：`**` 是量化噪声在「模型自己也没把握」的位置上的一次抛硬币 + +三条线索合起来： + +1. **噪声幅度是恒定的约 1.6 logit**（Δmargin std = 1.63，跨 7 个 decode step）， + 与 prefill 的残余误差同量级（pos11 的 Δmargin +0.835，14-token prefill 的 top-1 + 误差 −3.30）。它来自 per-128 量化的累积 —— SM90 的 DeepGEMM 强制的那条路。 +2. **参考模型的 margin 分布很宽**：9.86 / 10.33 / 11.95 / **1.64** / 15.11 / 17.87 / 10.40 / 11.09。 + 只有 1.64 那个位置会被 1.6 的噪声翻掉。 +3. **只有 pos14 翻转，其余全部正确。** 而 pos14 正是 1.64 那个。 + +所以： + +- 不是 KV cache 的问题（P4.3 闭环 round-trip 通过，rel≈0.027） +- 不是 decode indexer 的问题（P4.4 与独立 float64 参考 rel≈1e-7） +- 不是 MTP 的问题（开关 MTP 输出相同） +- 不是采样后处理的问题（6127 在 pos14 是第二名，没被 mask） + +**是「在 SM90 上必须用 per-128 量化」这件事本身**。它在多数位置无感，只在模型自己 +接近 50/50 的地方改变结果。而模型在这里确实接近 50/50 —— `**北京**` 的两个位置 +（先 `**` 还是先 `北京`）在语义上没有区别，参考模型也只是略偏一边。 + +**要继续压这个噪声，唯一的杠杆是回到 P3.6 那条路，而它在 SM90 上被 DeepGEMM 的 +`gran_k == 128` 断言封死了**（见 `H20-DRIFT-ANALYSIS.md`）。所以本项调查到此收敛： +剩余差异是硬件量化口径的必然结果，不是实现缺陷。 + +### P4.1 的工具坑（都已修，但 dump 仍建议配合黑盒交叉验证） + +| 坑 | 现象 | 修法 | +|---|---|---| +| 多 DP 进程写同一文件 | 8 个 engine 都是单进程组，`RANK` 全为 0，互相覆盖 | 文件名带 `_rank_pid

` | +| decode 的 batch 有 padding 行 | 取 `logits[-1]` 读到占位 token 的垃圾（恒定 `top1=69`） | 取第 0 行（第一个真实请求） | +| follower rank 的 logits 是占位值 | 同一向量被重复记录几十次，淹没真实步 | 判据 `driver_inputs is not None` | +| profile forward 用合成 token | 抓到 profile 的数据 | `debug_dump.profiling()` 挂起 | + +即便都修了，逐 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；`--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 异常（已排查，未修，非阻塞） 当前 8 个 token 的 `token_logprobs` 全是 `0.0`（概率=1），不正常。 -- [ ] P5.1 确认是采样器 bug 还是 logits 极端（打印 logits 的 max/median） -- [ ] P5.2 检查 `DEEPINFER_SAMPLER_FORCE_TOPK=4096` 与 logprob 回填的关系 +### 已排除（约 20 分钟静态排查 + 单测） + +| 嫌疑 | 结论 | +|---|---| +| `map_lookup` kernel | 无罪。单测 6 passed，`by_key` + float32 + `force_topk=4096` 全覆盖，与 `map_lookup_ref` 逐元素相等 | +| sampler 全部 kernel | 无罪。`sampler.py -m "not benchmark"` → **74 passed / 0 failed** | +| `SampledInfo.from_reqcnt` 的 `torch.zeros` | 死代码，全仓库无调用点；`update(is_final, candidate)` 也无调用点 | +| sampler→engine 字段传递 | 直通，`llm_engine.py:747` 原样传 `output.token_logprobs` | + +**一个反向线索**：`map_lookup` 里 `out = torch.empty`，键未命中时本应写 `-inf` 而**不是 0**。 +所以 0.0 更像「上层根本没把值带下来」，而不是「查表查空了」。 + +### 剩余候选与判定路径 + +- `process_sampling` 返回的 `logprobs` 本身为 0（`fused_sample` 的 logprob 输出路径； + 现有单测只覆盖 kernel 本身，没有覆盖「整个 forward 的 logprob 输出」） +- `build_request_outputs` 里逐 token 对象的属性回填 + +判定只需跑 1 次整机（~5 min）：在 sampler 里打印 `logprobs.max()` 与 `token_logprobs_cur`。 +sampler 层就是 0 → 定位在 `fused_sample`；sampler 层正常 → 定位在 output 回填。 + +### 为什么留着不修 + +不影响推理正确性，只影响 logprob 字段的可用性（AIME 打分、perplexity 这类评估用途）。 +预估典型 1 h、最坏 2 h（若牵涉 MTP 异步 logprob 回填），与本次 SM90 支持的目标无关。 --- diff --git a/H20-DRIFT-ANALYSIS.md b/H20-DRIFT-ANALYSIS.md new file mode 100644 index 0000000..e81baec --- /dev/null +++ b/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 比较。 diff --git a/deepinfer/debug_dump.py b/deepinfer/debug_dump.py new file mode 100644 index 0000000..12f96a9 --- /dev/null +++ b/deepinfer/debug_dump.py @@ -0,0 +1,183 @@ +"""Opt-in tensor dumps for H20 bring-up comparisons. + +Enabled with `DEEPINFER_DUMP_LAYERS=

`. Each call appends one row to +`/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] = {} +# 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) +# Per-tag call counter for `record_step`. +_STEPS: dict[str, int] = {} +# Number of rows in the most recent logits tensor, used to skip warmup passes. +_ROWS = 0 +# A decode step that serves real requests has at most this many rows. Anything wider is warmup. +_REAL_ROWS = int(os.getenv("DEEPINFER_DUMP_MAX_ROWS", "1") or 1) + + +_SUSPENDED = False + + +def enabled() -> bool: + return bool(_DUMP_DIR) and not _SUSPENDED + + +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 not once and tensor.shape[0] != _DUMP_TOKENS: + return + if once and tag in _CAPTURED: + 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) + # 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 set_rows(rows: int) -> None: + """Tell the dump how many rows the sampler just produced, so warmup can be filtered out.""" + global _ROWS + _ROWS = rows + + +def record_step(tag: str, tensor: torch.Tensor) -> None: + """Append `tensor`'s last row under `tag_`, numbering the calls. + + Decode needs every step kept, not just the first: the point of the decode comparison is to + watch the margin evolve and find the step where it collapses. + """ + if not _DUMP_DIR: + return + # Decode with one request produces exactly `_REAL_ROWS` rows; warmup batches are larger. + if _ROWS > _REAL_ROWS: + return + index = _STEPS.get(tag, 0) + _STEPS[tag] = index + 1 + # Row 0, not the last row: during decode the batch is padded (CUDA graph slots, or just + # `max_num_seqs`), and the padding rows hold logits for placeholder tokens. Row 0 is the + # first real request's current token, which is the one to compare against a batch-1 reference. + row = tensor[0] if tensor.dim() >= 2 else tensor + from deepinfer.utils import allow_cuda_sync + + with allow_cuda_sync(): + _CAPTURED[f"{tag}_{index}"] = 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 _rank_suffix() -> str: + """`_rank` when running distributed, so ranks do not overwrite each other's dump. + + With data parallelism every rank registers `atexit(flush)` and writes the same filename; the + last one to exit wins, which silently mixes one rank's tensors with another's. Only rank 0 + holds the real prompt anyway, so the suffix makes that explicit instead of accidental. + """ + rank = os.getenv("RANK") + if rank is None: + try: + import torch.distributed as dist + + if dist.is_initialized(): + rank = str(dist.get_rank()) + except Exception: # noqa: BLE001 - torch may not be usable at exit + rank = None + # With data parallelism each engine is its own single-process group, so every rank reports + # RANK=0 and they all target the same file. The pid disambiguates them. + return f"_rank{rank if rank is not None else 'x'}_pid{os.getpid()}" + + +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 / f"deepinfer_layers{_rank_suffix()}.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) diff --git a/deepinfer/h20_ops/debug_dump.py b/deepinfer/h20_ops/debug_dump.py deleted file mode 100644 index 6629906..0000000 --- a/deepinfer/h20_ops/debug_dump.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Opt-in tensor dumps for H20 bring-up comparisons. - -Enabled with `DEEPINFER_DUMP_LAYERS=`. Each call appends one row to -`/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] = {} -# 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: - 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 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) - # 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 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: - 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) diff --git a/deepinfer/model_executor/layers/sampler.py b/deepinfer/model_executor/layers/sampler.py index 384a588..d81488a 100644 --- a/deepinfer/model_executor/layers/sampler.py +++ b/deepinfer/model_executor/layers/sampler.py @@ -11,6 +11,7 @@ import torch import torch.nn as nn from deep_select.interface import topk +from deepinfer import debug_dump as _debug_dump from deepinfer import envs from deepinfer.logger import init_logger from deepinfer.model_executor.input_metadata import InputMetadata @@ -140,6 +141,19 @@ 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) + # `driver_inputs is None` means this rank only participates in the embedding all-gather and + # never sees the real result, so its logits are placeholders -- identical on every call. + # Recording those buries the real steps under dozens of copies of the same vector. + if _debug_dump.enabled() and driver_inputs is not None: + # 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. + if is_prompt: + _debug_dump.record("logits_prompt", logits, once=True) + else: + # Every decode step, so the comparison can watch the margin collapse step by step. + _debug_dump.record_step("logits_decode", logits) del hidden_states if driver_inputs is None: diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index 6aa27d4..59213c3 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/deepinfer/model_executor/models/decoder_layer.py @@ -5,7 +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 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 diff --git a/deepinfer/model_executor/models/deepseek.py b/deepinfer/model_executor/models/deepseek.py index 36eb614..4e01ac9 100644 --- a/deepinfer/model_executor/models/deepseek.py +++ b/deepinfer/model_executor/models/deepseek.py @@ -9,8 +9,8 @@ import torch from torch import nn from deepinfer import AttrDict, envs +from deepinfer import debug_dump as _debug_dump 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 diff --git a/deepinfer/model_executor/models/v4_modules/mlp.py b/deepinfer/model_executor/models/v4_modules/mlp.py index 2edcc4c..3c8b2e4 100644 --- a/deepinfer/model_executor/models/v4_modules/mlp.py +++ b/deepinfer/model_executor/models/v4_modules/mlp.py @@ -21,8 +21,8 @@ from tile_kernels.quant import swiglu_forward_and_per_token_cast from torch import nn from deepinfer import PLATFORM, AttrDict, envs +from deepinfer import debug_dump as _debug_dump 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 diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index 9a03065..f97f67b 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -11,6 +11,7 @@ import numpy as np import torch import torch.distributed as dist +from deepinfer import debug_dump as _debug_dump 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 @@ -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( 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() diff --git a/scripts/h20_recipe_probe.py b/scripts/h20_recipe_probe.py new file mode 100644 index 0000000..e8e4bef --- /dev/null +++ b/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() diff --git a/scripts/ref_layer_dump.py b/scripts/ref_layer_dump.py index 76ba49a..3990a18 100644 --- a/scripts/ref_layer_dump.py +++ b/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: diff --git a/scripts/ref_next_token.py b/scripts/ref_next_token.py index 25ee053..1d93de9 100644 --- a/scripts/ref_next_token.py +++ b/scripts/ref_next_token.py @@ -14,6 +14,7 @@ import json import os import sys +import numpy as np import torch import torch.distributed as dist @@ -30,6 +31,15 @@ def main() -> None: parser.add_argument("--ids", required=True, help="comma-separated prompt token ids") parser.add_argument("--out", default="/data/deepinfer-h20/ref/next_token.json") parser.add_argument("--topk", type=int, default=5) + parser.add_argument( + "--force_ids", + default="", + help="Comma-separated ids fed one at a time after the prompt, i.e. incremental decode. " + "This mirrors `generate.py`: prefill the prompt, then forward a single token per step with " + "prev_pos advancing, so the resulting logits come from the decode path rather than from " + "recomputing the whole prefix.", + ) + parser.add_argument("--logits_out", default="", help="optional .npz to store every step's full logits") args = parser.parse_args() ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x] @@ -59,21 +69,45 @@ def main() -> None: load_model(model, os.path.join(args.ckpt_path, f"model{rank}-mp{world_size}.safetensors")) torch.set_default_device("cuda") + forced = [int(x) for x in args.force_ids.replace(" ", "").split(",") if x] + tokens = torch.tensor([ids], dtype=torch.long) + steps: list[tuple[str, torch.Tensor]] = [] with torch.inference_mode(): # `Transformer.forward` returns (output_ids, logits, main_hidden). output_ids, logits, _ = model.forward(tokens, 0) + steps.append(("prefill", logits[0].float().clone())) - row = logits[0, -1].float() + # Same loop shape as `generate.py`: one token per forward, prev_pos advancing. The token + # fed in is the ground truth, which is exactly what the engine does when it accepts a + # greedy token it already agrees with. + prev_pos = tokens.shape[1] + for token in forced: + output_ids, logits, _ = model.forward(torch.tensor([[token]], dtype=torch.long), prev_pos) + steps.append((f"decode_{prev_pos}", logits[0].float().clone())) + prev_pos += 1 + + row = steps[-1][1] top = torch.topk(row, args.topk) - assert int(top.indices[0]) == int(output_ids[0, -1]), "argmax disagrees with the sampled id" payload = { "prompt_ids": ids, + "force_ids": forced, "next_token": int(top.indices[0]), "top_tokens": [int(x) for x in top.indices], "top_logits": [round(float(x), 4) for x in top.values], + "steps": [ + { + "name": name, + "top_tokens": [int(x) for x in torch.topk(v, args.topk).indices], + "top_logits": [round(float(x), 4) for x in torch.topk(v, args.topk).values], + } + for name, v in steps + ], } + if args.logits_out and rank == 0: + os.makedirs(os.path.dirname(args.logits_out), exist_ok=True) + np.savez(args.logits_out, **{f"logits_{name}": v.cpu().numpy() for name, v in steps}) if rank == 0: print(json.dumps(payload, ensure_ascii=False), flush=True) os.makedirs(os.path.dirname(args.out), exist_ok=True)