From 72933526297d3cb9d33c51b00a20af577d8b057f Mon Sep 17 00:00:00 2001 From: loopyt Date: Thu, 22 Oct 2026 08:58:10 +0800 Subject: [PATCH] docs: record the sampler's measured split and the four traps behind it The plan note still said the sampler was deliberately left on TileLang, which is no longer true. It now carries the measured numbers instead: 100% token agreement on the engine's own configuration, ~90% on narrow cuts, log-probabilities within 1e-5 everywhere, and the reason the token cannot do better -- the baseline's rejection limit steps at `total == 2^31`, which the fixed-point total sits next to, so a few units of rounding error move to a different draw. Plus the four traps that each looked like something else: signed arithmetic on uint32 (`/`, `%` and `<=`), nested runtime `if`s losing the assigned value, a threshold taken from probabilities and compared against logits, and renormalising by division instead of through the log domain. --- H20-TRITON-PLAN.md | 65 +++++++++++++++++++++++++++------------------- 1 file changed, 39 insertions(+), 26 deletions(-) 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 上」 -是端到端运行自己报出来的，不需要从名字表里推断。