diff --git a/H20-TRITON-PLAN.md b/H20-TRITON-PLAN.md index 3892c51..bf73baf 100644 --- a/H20-TRITON-PLAN.md +++ b/H20-TRITON-PLAN.md @@ -178,46 +178,59 @@ argmax 不受影响(token-ids 那次 8 个 token 全对),但 T=1.0 的采 | kernel | 状态 | |---|---| -| `fused_sample_tl` | **未移植**。它自带一套 radix top-k + top-p + 定点累积采样,随机数来自 TileLang 的 `T.rng_init/T.rng_rand`(C++ codegen 里的 Philox)。详见下一节。 | +| ~~`fused_sample_tl`~~ | **已移植**(见下节),引擎配置下 token 100% 一致。 | | `compute_seqlens_k_zos_tl`、`confidence_head_fused_tl`、`dflash_input_store_tl`、`get_accept_result_tl`、`finalize_verification_tl`、`remap_draft_topk_tl`、`get_fused_dflash_get_input_tl` | 未移植,均在 MTP 路径上(`DEEPINFER_USE_NEXTN_DECODE != 0` 时才会构造) | -### `fused_sample_tl` 为什么没有移植(有实测依据,不是「没时间」) +### `fused_sample_tl` 已移植,但「token 一致」有明确边界 -先用探针拿到引擎真实传的参数(`DEEPINFER_PROBE_SAMPLING=1`,异步 D2H,不在 forward 里同步): +先解掉了唯一真正的障碍:**随机流**。`T.rng_init/T.rng_rand` 降级成 `curand_init`/`curand`, +generator 默认 `curandStatePhilox4_32_10_t`。Philox 的混合函数是公开的,但 `curand_init` 的 +seeding 没有文档,所以是**实测反推**的(`scripts/h20_rng_dump.py` dump 真值): ``` -temperatures: [1.0, 1.0] -top_ps: [0.99, 0.99] -top_ks: [-1, -1] -sample_positions: [4097, 2048] +key = (seed & 0xFFFFFFFF, seed >> 32) +ctr = (offset & 0xFFFFFFFF, offset >> 32, subsequence & 0xFFFFFFFF, subsequence >> 32) +每次 draw 依次取 block 的 4 个 word,然后 ctr.x += 1 ``` -也就是说这是**温度 1.0 + top-p 0.99 的真实核采样**,不是 argmax:`force_topk` 先把 logits 截到 4096, -然后 `r = rng() % total_mass`,取定点累积质量第一个超过 `r` 的那一列。所以每次抽样都真的用了随机数。 +**subsequence 落在第 3、4 个 counter word,不是第 1 个**(cuRAND 的 subsequence 跨 2^64 个 block)。 +这个细节在 subsequence=0 时完全看不出来、其他情况全错,所以对拍同时覆盖 0 和 4097。 +`scripts/h20_philox_check.py`:5 个 seed(含 2^64-1)× 4 个 subsequence(0/1/2048/4097) +→ **20/20 bit-identical**。 -有两个障碍,都属 TileLang 内部实现: +采样器本体(单 tile,vocab=4096,与引擎实际构造一致)已移植,A/B 结果: -1. **随机流**。`T.rng_init(seed, seq)` / `T.rng_rand()` 是 Philox,实现在 TileLang 的 C++ codegen 里 - (`src/cuda/codegen/codegen_cuda.cc`),Python 侧没有等价物。Triton 当然可以自带一个生成器, - 但那是**另一条流**,温度 > 0 时每一个采样出的 token 都会变。这是行为改变,不是移植。 -2. **归约顺序**。基线专门用 `annotate_layout` 钉死 LSE 累加器的布局,注释写明是为了 - "preserve the original reduction order for bitwise alignment";随后概率被量化成 2^31 定标的 - `uint32` CDF,一个 ulp 的概率差就可能让某一列跨过边界。 +| 配置 | token 一致率 | logprobs | +|---|---|---| +| **engine_cfg(top_p=0.99,2644 保留列)** | **128/128 = 100%** | 1e-5 | +| engine_cfg batch32 | 32/32 = 100% | 1e-5 | +| top_k=1 / no_cut | 100% | 1e-5 | +| top_k=128 | 91.4% (117/128) | 1e-5 | +| top_p=0.5 | 86.7% (111/128) | 1e-5 | +| top_p=0.9 + top_k=256 | 87.5% (112/128) | 1e-5 | -两者叠加的结果是:Triton 版采样器会**静默产生不同的文本**。所以让 `fused_sample` 留在 TileLang、 -其余全部换成 Triton,是正确的工程选择 —— 每 kernel 的回退机制本来就是为这种情况准备的。运行日志里的 +**保留集合逐列一致,logprobs 差 1e-5(= 行归约顺序不同)。token 在窄 cut 下有 ~10% 不同, +原因已定量定位,不是实现缺陷:** ``` -[ops] kernel 'fused_sample_tl' used from 'cuda': 'triton' does not implement it +r = draw % total, total = sum(floor(p_i * 2^31)) ``` -就是它在自我报告,而且**这就是完整清单**(非 MTP 路径 19 个 kernel,只有它一个)。 +baseline 的拒采样上界是 `(2^32-1) // total * total`,这个表达式在 `total` 跨过 **2^31** 时 +从 `2*total` **阶跃**到 `1*total` —— 而 `total` 因为概率和为 1 就贴着 2^31。台阶这一侧几乎不拒绝、 +另一侧拒绝约一半,每次拒绝就换到下一个 draw,token 直接变。所以 `total` 差几个单位(归约顺序 +决定的,不可能一致)就会换 draw。**把 total clamp 到 2^31 以下没用**:baseline 自己也会跨过去 +(两个方向都有),clamp 只是把不一致换到别的行(11 → 30)。 -要移植它,顺序是:先做一个与 TileLang 流完全一致的 Philox(用固定 seed 对拍 `T.rng_rand()` 的输出), -再做 radix top-k/top-p 过滤,最后做定点 CDF 与两级搜索。 +结论:采样器本身是对的(分布一致到 1e-5),只是它的随机性与 baseline 在窄 cut 上对不齐。 -### 回退提示改成了「使用期」报告 +### 移植过程中踩到的四个坑(每一个第一眼都像别的问题) + +1. **Triton 对 uint32 的 `//`、`%`、`<=` 走有符号运算**,4294967295 当 -1。这让 `limit` 变成 + `2^32 - total`,拒采样决策反了 → 整条 draw 序列错位 → 23% 的行 token 全错。整个 CDF 改用 int64。 +2. **嵌套的运行期 `if` 会产出「分支内赋的值不等于分支外读到的值」**。cut 的组合改用 `tl.where`。 +3. **threshold 取自概率数组却拿去和 logit 比**。排序必须对 logit 做(等价,因为映射单调), + 再从排好序的 logit 重算质量。 +4. **用除法重新归一化,而不是 baseline 的 log 域**(`lse0 -= log(mass)` 再 `exp(logit*inv + lse0)`)。 + 同数学不同舍入,改成 log 域后一致率从 76% 抬到 91%。 -原来的提示在 import 时打印「缺 N 个 kernel」,那是静态计数,看不出运行到底走没走到。现在每个未实现的 -kernel 被包成 `_FallbackNotice`,第一次被**调用**时打印一行。这样「哪些 kernel 还在 TileLang 上」 -是端到端运行自己报出来的,不需要从名字表里推断。