DeepInfer 源码拆解（三十八）：采样内核的微观账——基数筛选、两级定点采样与残差并行的 TileLang 实现

归约顺序为什么要被锁死

在线 log-sum-exp 的合并公式是：

这个递归对分块的合并顺序敏感：浮点加法不满足结合律，不同的合并顺序得到不同的最后几位。源码在 NUM_STAGES 分支里显式给 lf0 指定了 fragment 布局：

1
if NUM_STAGES:
2
    # Preserve the original reduction order for bitwise
3
    # alignment with the reference cross-entropy implementation.
4
    T.annotate_layout(
5
        {lf0: T.Fragment((LSE_BD,), forward_fn=lambda i: ((i % (THREADS * 8)) // 8, i // (THREADS * 8) * 8 + i % 8))}
6
    )

这个 forward_fn 把逻辑下标 $i$ 映射成 (thread_id, register_index)：每一轮 8 个连续元素落在同一个线程的相邻寄存器里，下 8 个落到下一个线程。也就是说，每个线程持有的是一组长度为 8、跨度为 THREADS * 8 的连续片段。这样每个线程内部按固定步长顺序累加，跨线程再经过 T.finalize_reducer 汇总，整个规约树就与 PyTorch 的参考实现一致。

对采样来说，logprob 的绝对误差通常不会改变采出的 token；但只要引擎把 logprob 回传给用户（logprobs 输出字段），那么”同一条请求在任意 batch 下给出逐 bit 相同的 logprob”就是一项对外承诺。归约顺序被锁死是为了兑现这条承诺，而不是为了数学上的精确性。

最后一块的 pdl_trigger

1
if i0 == N_LSE_CHUNKS - 1:
2
    T.pdl_trigger()

T.pdl_sync() 与 T.pdl_trigger() 是一对：前者让本 CTA 等待所有依赖的前序 launch 到达，后者提前向硬件宣告”我即将写完输出、后续 launch 可以开始依赖同步”。第二段 kernel 开头先 T.pdl_sync() 再读 lse_workspace，而第一段选在”最后一次块归约 finalize 之前”就触发，含义是”最后一块计算尚未收敛，但你要的输入数据已经在途了”。硬件会把第二段的一部分启动开销与第一段最后一块的归约重叠。这是 Hopper/Blackwell 上 PDL 的典型用法：提前 1～2 微秒，在几百行的 batch 上就是可见的收益。

可排序整数键

top-p 要按概率降序排，top-k 要按logit降序排。softmax 是单调函数，$p_i > p_j \iff z_i > z_j$（温度为正时），所以两者的降序关系完全一致。内核把浮点 logit 映射成一个 u32 整数键，使得：

1
raw = T.reinterpret(T.uint32, x)
2
return T.Select(raw > SIGN_BIT, ~raw, raw | SIGN_BIT)

这是 IEEE-754 浮点数被当作有符号整数比较时的一个经典修正：

• 正数（符号位为 0）的位模式作为无符号整数本身就是单调的，只需翻转符号位（raw | SIGN_BIT），让它们的键落在 u32 的高半区；
• 负数（符号位为 1）的位模式大小与数值大小相反——数值越大，位模式越小，所以需要取反（~raw）；取反后符号位自然变成 0，键落进低半区。

注释写着”Map +0 and -0 to the same key for stable ties”。这个效果来自分支条件的边界：T.Select 判断的是 raw > SIGN_BIT。$+0$ 的 raw 是 0x00000000，不满足大于，走 raw | SIGN_BIT，得到 0x80000000；$-0$ 的 raw 恰好等于 0x80000000，同样不满足严格大于，也走 raw | SIGN_BIT，还是 0x80000000。两者因此落在同一个键上。如果条件写成 >=，$-0$ 就会走取反分支，得到 0x7FFFFFFF，两个零的键被拆开，并列判定会随符号位而抖动。

$-0$ 在 logits 里不常见，但在被 mask 的 -inf 两侧容易出现各种等价的零；把它归一化成一个键后，并列判定就不再因为符号位闪烁而随机。

11 位基数直方图

筛选的核心是一个”前 $k$ 个”的判定问题，可以统一写成：

• top-p：求最小的前缀 $k$，使前 $k-1$ 项的累计质量 $\le p_{top}$、而前 $k$ 项的累计质量 $> p_{top}$（即前 $k$ 个是首次严格跨过阈值的前缀）；
• top-k：求 $k = \text{top\_k}$ 处第 $k$ 大 logit 的键。

两者都对排序后的顺序有要求，但真正需要的其实只是”键的分布”——第 $k$ 大的键是多少、前 $k$ 个键的（分数形式）概率和是多少。于是内核按位从高到低做多次直方图统计：

1
RADIX_BITS = 8 if LSE_BD <= 1024 else 11
2
RADIX_BUCKETS = 1 << RADIX_BITS
3
ROUNDS = T.ceildiv(32, RADIX_BITS)

一个 u32 键被切成 ROUNDS 段，每段 RADIX_BITS 位。每一轮：

1. 清空 hist（top-p 的双定点累加）与 topk_hist（计数）；
2. 扫一遍整行 logits，只统计”前缀与已确定前缀一致”的键：prefix_match = ((bits ^ desired) & dm) == 0；
3. 对命中前缀的项，分别往 hist 和 topk_hist 的对应桶里累加；
4. 对两个直方图做反向累积和（block_cumsum(..., reverse=True)），得到”键大于等于本桶的项的总质量”；
5. 在每个桶上检查”从最高桶往下的累计概率首次达到 $p_{top}$“或”从最高桶往下的计数首次达到 $k$“——找到的那一段位就写进 desired，并把剩余需要的质量/数量记进 ps_s/ks_s，供下一轮使用；
6. dm 追加掩码，下一轮只在确定的这一位段内继续细分。

RADIX_BITS 的取值有一个精确的权衡：

1
# Small LSE tiles use 256 buckets to reduce histogram clearing/scanning;
2
# larger tiles use 2048 buckets to save one radix pass.

• 桶越多，每一轮需要扫描的桶越多，清空与累积和的开销随桶数线性增长；
• 桶越多，每轮能确定的位数越多，总轮数 ROUNDS = ceil(32 / RADIX_BITS) 越少。

键是 u32，因此 8 位一轮要 4 轮，11 位一轮只要 3 轮（最后一轮只用到剩余的 10 位）。而 LSE_BD 在 vocab_size 很大时是 8192，直方图扫描的成本相对整行扫两次 logits 的成本并不显著，所以选 11 位。小词表（LSE_BD <= 1024）时整行本来就很小，扫描桶反而成了主要成本，于是退回 8 位。

为什么用反向累积和？因为”保留前 $k$ 个”是保留键最大的那些。反向累积和 hist[bucket] 表示”键 $\ge$ 本桶下界的项总和”，它直接对应”从高到低保留了多少”。第 5 步的比较条件 hist[bucket] >= ps_cur 是”累计保留质量已经达到目标”，prev < ps_cur（prev = hist[bucket+1]）是”上一格还没达到”。两个条件一起保证了唯一确定一个桶：这轮选择的桶是”跨过目标的那一格”，而 ps_cur - prev 就是”下一轮还要在更低位继续找的剩余质量”。

T.atomic_add 被用于往直方图桶里累加：一个 CTA 内多个线程可能命中同一个桶。这里没有使用 T.reduce_sum，因为直方图是散射写入，桶与桶之间的累加无法用规约表达。

双量同树：为什么 top-p 和 top-k 能在一棵树上收敛

当 top_p_active and top_k_active 同时成立时，find_threshold(True, True, ...) 在一轮直方图里同时维护两个量：

• hist：按分数形式累积概率，定点编码为 uint32，scale 为 $2^{31}$；
• topk_hist：计数，uint32。

两者共享同一个”前缀”路径（desired、dm），因为在键空间里两个条件说的是同一件事：保留前 $k$ 个最大的键。top-p 的 $k$ 由概率和质量决定，top-k 的 $k$ 由计数和参数决定，但两者都在”键降序”这个共同序上定义。

所以”先 top-p 后 top-k”的真实语义是：

由于键的排序与 logit 排序一致，最终阈值就是两个阈值里更高的那个。树结构天然解决了”谁的截断更紧”：desired 在逐位细化时，只要有一个条件没达到，就会继续细分到更低位的位段。逐位收敛的最大值就是最终阈值。

并列截点又受两个条件共同约束——既不能保留超过 top-p 允许的并列项数量，也不能保留超过 top-k 允许的并列项数量。sel 的公式正是这两个约束的交：ks_s 是 top-k 剩下还能保留的个数，ps_s / ps_th 是 top-p 剩下还能保留的”以阈值概率计的份数”，两者取小后上取整，得到”最多还能保留几个等于阈值的 token”。后续所有阶段一致地使用”col > cut 即剔除”。

注意 softmax 的浮点舍入：内核在键空间里做的是”logit 的精确比较”，而参考实现是在已经经 softmax 变成概率之后再比较累积和。由于 softmax 严格单调（不考虑舍入），两者在只开 top-p 时选取的边界一致；但在并列键上，概率可能因为指数运算的舍入而不完全相等，导致边界在极少的情况下多退一步或少退一步。这也是为什么测试里同时跑了多组温度与 vocab_size，而不是只测一组。

还有一个更结构性的差异：top-k 与 top-p 同时开启时，内核在同一棵键树上分别求两个阈值、再取更严格的那个，$p_{top}$ 的质量口径是完整词表的 1；而把 top-k 与 top-p 作为两个先后 warper 串联的实现（HF 的 LogitsProcessorList、vLLM 的 apply_top_k_top_p_pytorch、llama.cpp 的采样器链）会先按 top-k 裁掉尾部，再在重新归一化的 top-k 子集上计算 top-p，因此同样参数下候选集可能更窄。两套语义都成立，但跨框架对照采样结果时必须先确认属于哪一种。内核注释里”与参考实现一致”指的是单截断路径与归约顺序，不覆盖合取语义。

并列阈值：分组扫描

找到了 threshold 之后，整行里可能有大量 token 的键恰好等于 threshold——比如 logits 被结构化输出的 mask 大量置成同一个 -inf，或者模型在某个位置给出多个完全相等的 logit。ps_s[0] 与 ks_s[0] 记录的是”达到阈值后还保留多少”：

ps_th 是阈值处的概率 threshold_p 按同一个 31 位定点 scale 编码后的整数：

1
threshold = sortable_to_f32(encoded)
2
threshold_p = T.exp((threshold + lse1) * inv_temperature + lse0)

接着 sel_f 表示”只保留 sel 个等于阈值的项”。扫描是分块进行的：

• 若当前块之前累计的 tie 数量已经达到 sel_f，本块的 ties 全部丢弃（连同严格小于阈值的项一起置 -inf）；
• 否则先数本块的 ties：若加完后仍未达到 sel_f，保留全部 ties、把严格小于阈值的项置 -inf；
• 若这个块跨过了 sel_f（cnt > sel_f），说明边界就落在这个块里，且最多发生一次。此时只对这个块做分组扫描：每个 256 项组计数 → 组计数的前缀和 → 锁定”第 sel_f 个 tie”落在哪个组 → 在组内再前缀和定位第几个元素 → 得到 cut；
• 之后的块里，所有 tie 都被置成 -inf（col > cut），同时把剩余的有效项概率质量累加进 blk_mass。

为什么只做一次分组扫描？因为 sel_f 是确定的、而块是按顺序处理的，“跨过 sel_f”的块唯一。分组扫描成本是 $O(\text{TIE\_BD})$ 的两次 cumsum，不随整行长度增长。整体复杂度仍是 $O(V)$ 一次遍历加一次局部扫描。

blk_mass 是排除后剩余候选的质量和，最后：

1
mass = T.alloc_var(T.float, blk_mass[0])
2
cut = cut_s[0]
3
lse0 -= T.log(mass)

把 lse0 减去 $\log(\text{mass})$，等于把归一化因子从”全词表的 $\log \sum$“换成”候选集内部的 $\log \sum$“。之后采样阶段用的 lf[j] = (lf[j] + lse1) * t + lse0 就自动落在重归一化后的候选分布上，不需要再单独算一次候选集的归一化。

有一个精度细节：blk_mass 被累加成浮点，而直方图用的 31 位定点质量。两者之间的取舍是”候选集质量需要参与一次 log 运算（浮点）“与”排序判定需要可靠的前缀和（整数）“。threshold 处的概率 threshold_p 会带来另一层舍入：当 threshold_p 极小（远距离尾部）时，$p\_s / p_{th}$ 会很大，sel_f 被截到 vocab_size；当 threshold_p <= 0 时（下溢为 0），sel_f 设为 0，即”一个 tie 都不保留”，把边界推到严格不等号。

定点数与”零概率永不采中”

为什么不用浮点 CDF？参考实现使用过浮点前缀和，但有一个隐蔽的问题：浮点累加在长前缀下误差会逐步放大，而”样本落在哪个区间”是一个不连续的选择——误差跨过阈值时，采出的 token 就变了。定点数把这个问题变成可分析的量化误差：

• 每个概率向下截断到 $2^{-31}$ 的网格，量化误差在 $[-2^{-31}, 0)$ 之间；
• $V$ 项累计误差落在 $[-V \cdot 2^{-31}, 0)$，$V = 129280$ 时最多约 $-6 \times 10^{-5}$；
• 所有 $p_i < 2^{-31}$ 的项都量化成 0，累积区间长度为 0，永远不会被选中。

第二条让”小于量化底线的 token 永不采出”成为硬性保证，而不只是”几乎不会”。这与测试里”采样结果的概率必须大于 0”的断言一致。float_to_fixed 的实现对 uint64 有专门路径：

1
bits = T.reinterpret(x, "uint32")
2
exp = T.cast((bits >> 23) & 0xFF, T.int32)
3
mant = T.cast((bits & 0x7FFFFF) | 0x800000, dtype)
4
top = mant << T.cast(scale_bits - 23, dtype)
5
rshift = 127 - exp
6
return T.if_then_else(rshift < 64, top >> T.cast(rshift, dtype), T.cast(0, dtype))

把 float32 直接按位拆成阶码与尾数，尾数补上隐含位 1 后左移到 64 位空间，再按阶码右移，从而避免先转浮点再乘以 $2^{31}$ 时可能产生的第二重舍入。rshift >= 64 的情形（指数很小、下溢）返回 0，这就是”极小概率量化成 0”的来源。scale_dtype 是 uint32 时走的是简单乘法路径（x * 2^{scale_bits} 再截断），因为此时不存在 scale_bits - 23 > 63 的移位越界问题，不需要按位拆。

第一级：拒取法与无分支二分

第一级的输入是 lane_cdf = cumsum(acc_frag)（反向不需要，这里用正向），输出是”随机数 $r$ 落在哪个车道 $k_0$，以及它在车道内的余量 $r_2$”。

随机数怎么来：

1
total = T.max(rt_s[1], 1)
2
limit = (~T.cast(0, SCALE_DTYPE)) // total * total
3
rnd = T.alloc_var(SCALE_DTYPE, limit)
4
while rnd >= limit:
5
    rnd = T.rng_rand()
6
r_s[0] = rnd % total

三段：

1. limit 是 $\lfloor (2^{32}-1) / \text{total} \rfloor \cdot \text{total}$，即 $2^{32}-1$ 向下取整到 total 的倍数。这是拒取法（rejection sampling）的经典做法：均匀随机数范围是 $[0, 2^{32})$，取模后落在 $[0, \text{total})$，但 $\text{total}$ 不能整除 $2^{32}$，导致一部分余数区间被多采到一点点。把大于 limit 的样本拒绝并重抽，得到的模值就是严格均匀的。while 循环的期望迭代次数小于 2。
2. 为什么需要严格均匀？因为采样的”正确性”（与目标分布一致）依赖于均匀随机数。定点采样对随机数偏差的容忍度比浮点低：浮点误差可以被吸收，但整数模偏差在固定点网格上会成为真实的概率偏移。
3. T.rng_init(seed_val, verify_seq) 出现在这个之前。随机数流由 (seed, seq) 初始化，其中 seq = sample_positions[row] * 3 + case。这是 第八篇 讲过的三路分流：case 0 是正常采样、case 1 是残差重采样、case 2 是投机接受判定。每一个需要随机数的环节都只消耗自己的那条流，互不干扰。

定位车道用的是一个无分支二分：

1
search_threads = min(cfg.threads, BD)
2
slots = BD // search_threads
3
for lane in T.Parallel(search_threads):
4
    slot = 0
5
    for level in T.unroll(slots.bit_length() - 1):
6
        step = 1 << (slots.bit_length() - 2 - level)
7
        probe = (slot + step - 1) * search_threads + lane
8
        if lane_cdf[probe] <= r:
9
            slot += step
10
    k = slot * search_threads + lane
11
    prev = T.if_then_else(k > 0, lane_cdf[k - 1], 0)
12
    if prev <= r and lane_cdf[k] > r:
13
        offset_s[0] = k
14
        rt_s[0] = r - prev

每个线程负责一条单调子序列 lane_cdf[lane], lane_cdf[lane + T], lane_cdf[lane + 2T], ...，在这条子序列上做二分。由于整条 CDF 是单调的，任意线程的子序列也单调。各线程找到”本子序列中第一个 $\ge r$ 的位置”后，再检查一次相邻项，只有唯一一个线程会同时满足 prev <= r && lane_cdf[k] > r。

这里的选择比较有意思：一般做法是单线程串行二分 2048 项（11 次），或者整块 CTA 用 one-hot 前缀和。这个实现用了 256 个线程、每线程 8 项的子序列二分（3 次），总代价是 3 次读取（每次读 search_threads 个并行地址），不用同步、不用归约，也没有分支发散——循环体是固定次数的 if 更新。

slots & (slots - 1) == 0 是一个断言：子序列长度必须是 2 的幂，这样 slots.bit_length() - 1 恰好是二分步数。assert BD % search_threads == 0 保证均匀分配。

第二级：槽位精定位

第一级给出车道 k_0 与余量 $r_2 = r - \text{lane\_cdf}[k_0-1]$。第二级只需要在”车道 $k_0$ 的 64 个槽位”中定位：

1
for m in T.Parallel(N_SLOTS):
2
    col = offset + m * BD
3
    slot_cdf[m] = T.if_then_else(
4
        col < vocab_size,
5
        float_to_fixed(cfg.scale_dtype, cfg.scale_bits, slot_probs[m]),
6
        0,
7
    )
8
T.cumsum(slot_cdf, dim=0)
9
for m in T.Parallel(N_SLOTS):
10
    col = offset + m * BD
11
    prev = T.if_then_else(m > 0, slot_cdf[m - 1], 0)
12
    cum = slot_cdf[m]
13
    if col < vocab_size and prev <= r2 and cum > r2:
14
        sel_s[0] = col

slot_probs[m] 是 probs[row, k_0 + m*BD]，即”车道 $k_0$ 的那一列”。这是精确到 token 的概率，不是车道和——所以第二级需要在采样阶段重新从 logprobs 里取出这些值：

1
if col < vocab_size:
2
    slot_probs[m] = T.exp(logprobs[row, col])
3
else:
4
    slot_probs[m] = 0.0

logprobs 在采样阶段被写入（logprobs[row, i0 * SAMPLE_BD + j] = lf[j]），所以候选集合的重归一化已经存在列方向上了。取一列 = 64 次跨步访问，对每个 CTA 的 256 个线程而言并行度足够。

T.cumsum(slot_cdf, dim=0) 对 64 项做前缀和，然后每个槽位各自检查一次区间。sel_s 只被一个线程写入（唯一匹配），所以不需要原子。第二阶段并不做二分：每个线程检查自己的槽位，因为只有 64 项，T.cumsum 一次就把前缀和算完了。

两级采样的总复杂度：第一级每个线程 3 次读、共 $3 \times 256$ 个地址；第二级 64 项 cumsum + 64 次检查。整行的采样成本与 $V$ 无关，只取决于 BD=2048 与 N_SLOTS=64。

为什么”第一级用定点、第二级用浮点”

sample_acc_frag 是定点累加，而 slot_probs 是浮点概率。原因在于两级承担的不同需求：

• 第一级需要全局单调、可精确可复现的比较。一个跨 64 个槽位的和，如果每次用浮点重算都可能有最后一位的波动，那么”哪个车道”就不稳定，结果会随 batch 或编译配置变化。
• 第二级只需要在一列 64 项内做前缀和，误差来源单一、数量级小，浮点足以保证正确性；而且 float_to_fixed 也在这时把浮点转回定点，所以最终比较仍然在定点网格上。

也就是说，“定点”是采样核心的数学域，浮点只是装载与重归一化过程中的表达形式。这个设计让”同一请求任何 batch 下采出同一 token”成为一条编译期可以推理的性质，而不是需要统计验证的性质。

接受判定

1
if is_reject_sampling:
2
    draft_tid = T.cast(last_next_n_token_ids[draft_row], T.int32)
3
    accepted_s[0] = False
4
    if draft_tid >= 0:
5
        draft_lp = last_nextn_token_logprobs[draft_row, draft_tid]
6
        draft_logit = logits[row, draft_tid]
7
        target_lp = (draft_logit + lse1) * inv_temperature + lse0
8
        if (top_p_active or top_k_active) and (
9
            draft_logit < threshold or (draft_logit == threshold and draft_tid > cut)
10
        ):
11
            target_lp = -inf
12
        accept_lp = T.min(target_lp - draft_lp, 0)
13
        verify_seq = sample_positions[row] * 3 + 2
14
        T.rng_init(seed_val, verify_seq)
15
        accepted_s[0] = T.rng_rand_float(32, "uniform") < T.exp(accept_lp)
16
    accept_result_verify[row] = accepted_s[0]

有几个细节需要逐条解释：

draft_tid >= 0 的提前返回。草稿 token 为 -1 表示”这一轮没有草稿”（在变长验证中，某些行没有对应的草稿位置）。此时直接拒绝，不参与判定。这与 第五篇 讲的变长验证一致：不是每一行都有草稿。

target_lp 在阈值判定中被置成 -inf。这一点很重要：如果被截断后的候选集里没有草稿 token，那么草稿 token 在目标分布里的概率是 0，它无论如何都不该被接受——即便原始 logit 很高。因为”目标分布”在这里的定义是截断并重归一化后的分布，不是原始 softmax。

残差路径则不是”直接套阈值”，而是用 (threshold, cut) 这个边界再判一次”是否落在候选集内”：采样阶段里每个被截断的项 lf[j] 被置成 -inf，pf[j] = exp(-inf) = 0，所以 max(pf - α·draft, 0) 对被截断的 token 恒为 0。边界一次算好，两个分支复用——这正是前文反复强调”阈值要与采样阶段共用同一份 logits”的原因。

accept_lp = min(target_lp - draft_lp, 0)。草稿概率用 log 存，所以目标/草稿比值 $q/p$ 的对数是 $\log q - \log p = \texttt{target\_lp} - \texttt{draft\_lp}$；封顶 0 后就是 $\min(1, q/p)$。draft_lp 是草稿模型自己的 logprob，取自 last_nextn_token_logprobs——即草稿头的输出分布。

accept_lp.exp() 与均匀数比较。这是”以概率 $a$ 接受”的标准写法。随机数用的是 case=2 流，与正在进行的采样（可能还没做）和残差重采样都不同流。

残差并行收敛

用统一的记号：设目标分布为 $q$、草稿分布为 $p$（源码 draft_lf 的 p 对应草稿），被拒绝后应从 $\text{norm}(\max(q - p, 0))$ 重采样。最直接的做法是先把残差写成一个 $V$ 维向量，再像普通采样那样走两级 CDF。但这样要多写多读一遍 $V$ 项，而且第二级采样本来就是”取一列”，残差也取一列就行。

内核的做法是：在原本用来写 logprobs 的同一趟遍历里，同时把残差概率累进 sample_acc_frag：

1
for reject in T.unroll(2 if do_reject_sampling else 1, explicit=True):
2
    if is_reject_sampling == (reject != 0):
3
        for i0 in T.Pipelined(N_SAMPLE_CHUNKS, num_stages=3 if NUM_STAGES else 0):
4
            ...
5
            for j in T.Parallel(SAMPLE_BD):
6
                pf[j] = T.Select(i0 * SAMPLE_BD + j < vocab_size, T.exp(lf[j]), 0.0)
7
            for j in T.Parallel(SAMPLE_BD):
8
                logprobs[row, i0 * SAMPLE_BD + j] = lf[j]
9
            if reject != 0:
10
                ...
11
                for j in T.Parallel(SAMPLE_BD):
12
                    pf[j] = T.Select(
13
                        i0 * SAMPLE_BD + j < vocab_size,
14
                        T.max(pf[j] - T.exp(draft_lf[j] - DRAFT_PROB_PENALTY), 0.0),
15
                        0.0,
16
                    )
17
            if do_sample:
18
                accumulate(sample_cfg, pf, sample_acc_frag)

for reject in T.unroll(2 if do_reject_sampling else 1, explicit=True) 是一个按行分流的技巧：每一行要么走普通采样（reject == 0 分支），要么走残差路径（reject == 1 分支），两者的代码几乎相同，区别只在 pf[j] 的计算和是否读取草稿分布 draft_lf。T.unroll + explicit=True 让编译器把两个分支展开成两段（未流水化时可能被合并成同一段遍历），运行时由 is_reject_sampling == (reject != 0) 的谓词决定哪些行进入哪一段。这样残差行不需要额外一趟 kernel，而普通行不会为残差计算付任何代价。

这里 accumulate 累加的是未归一化的相对质量：

其中 $\alpha = e^{-\text{DRAFT\_PROB\_PENALTY}}$。代码里并没有显式除以归一化因子 $Z = \sum_i \max(q_i - \alpha p_i, 0)$。那第二级采样怎么”归一化”的？答案在 sample_offset：

1
total = T.max(rt_s[1], 1)
2
...
3
r_s[0] = rnd % total

rt_s[1] 就是 lane_cdf[BD-1]，若候选集归一化到 1，它就是 1；若残差只累加了 $Z < 1$ 的相对质量，那么 total 就是 $Z$，随机数对 total 取模，相当于从 $[0, Z)$ 均匀取数，再与未归一化的 CDF 比较——自动完成了重归一化。这是一步很巧的替换：

等价于在归一化分布上采样。同样的技巧也用于 lane_cdf[k] > r 的边界判定：边界比较的尺度与累加器一致，不需要显式除法。

而 allow_zero=True 时，若 total 恰好为 0（残差质量全为 0，例如草稿分布处处 $\ge$ 目标），则 total = max(0, 1) = 1、limit = (2^{32}-1)//1*1 = 2^{32}-1，所以循环条件 rnd >= limit 几乎立即为假，rnd % 1 = 0，sample_select 选不出任何槽位、最终返回 0。真正兜底的是后面的 rt_s_rej[1] == 0 判定：

1
if rt_s_rej[1] == 0:
2
    max_prob_tid = T.alloc_reducer((1,), T.int32, op="min", replication="all")
3
    T.fill(max_prob_tid, vocab_size)
4
    for i0 in T.serial(N_SAMPLE_CHUNKS):
5
        fallback_logits = T.alloc_fragment((SAMPLE_BD,), T.float)
6
        T.copy(logits[row, i0 * SAMPLE_BD], fallback_logits)
7
        for j in T.Parallel(SAMPLE_BD):
8
            col = i0 * SAMPLE_BD + j
9
            if col < vocab_size and fallback_logits[j] == max_v:
10
                max_prob_tid[0] = T.min(max_prob_tid[0], col)
11
    T.finalize_reducer(max_prob_tid)
12
    idx_rej = max_prob_tid[0]

rt_s[1] 在 allow_zero 时由 sample_offset 以 lane_cdf[BD-1] 的形式返回——真实的未归一化总质量。当它等于 0 时，内核回退到”取最大 logit 的最小下标”（op="min" 保证并列时取最左边的一个）。也就是说，零残差时的回退是确定性的 argmax，而不是”随便采一个”。这与第八篇讲的”同输入同输出”性质一致。

再补一处：allow_zero 时的 total 计算 T.max(rt_s[1], 1) 是”避免除以零”的守卫——但 rnd % 1 = 0，所以所有车道都会被”选中”，rt_s[0] 为 0，最后全靠 rt_s_rej[1] == 0 分支的结果。这属于”先让程序跑通、再在调用方处纠正”的写法；关键是纠正路径本身是确定性的。

accept_probs：残差总质量即接受概率

1
if accepted_s[0]:
2
    residual_mass = T.alloc_reducer((1,), sample_cfg.scale_dtype, replication="all")
3
    T.fill(residual_mass, 0)
4
    for j in T.Parallel(SAMPLE_BD):
5
        residual_mass[0] += sample_acc_frag[j]
6
    T.finalize_reducer(residual_mass)
7
    if T.get_thread_binding() == 0:
8
        accept_probs[row] = T.max(1.0 - T.cast(residual_mass[0], T.float) / sample_cfg.scale, 0.0)

accept_probs 回传到上游的是 $1 - \sum_i \max(q_i - \alpha p_i, 0)$ 的定点近似。从数学上，若忽略惩罚项 $\alpha$，则 $1 - \sum_i \max(q_i - p_i, 0) = \sum_i \min(p_i, q_i)$，这正是随机采样接受概率。它被上层用来做投机深度控制。

一个容易误读的点：这段只在 accepted_s[0] 为真时计算，即”草稿被接受”之后。这看起来很反直觉——接受率不应该由是否接受这次来决定。实际上 accept_probs 在这里不是”这一次的接受结果”，而是”这一行在当前分布下的期望接受概率”，它对所有验证行都要计算，而 residual_mass 恰好是已经算好的残差总质量。accepted_s[0] 只是控制流上的分支：拒绝的那一支已经有了 rt_s_rej[1]（同样表示总质量），也能直接算出 1 - 残差质量。所以两个分支都产出 accept_probs，只是数据来源不同。源码里拒绝分支的写法：

1
accept_probs[row] = T.max(1.0 - T.cast(rt_s_rej[1], T.float) / sample_cfg.scale, 0.0)

注意 sample_cfg.scale = 2^{scale_bits} = 2^{31}，所以 residual_mass / scale 是”残差质量占全概率的比例”。

random_sampling_tl：非融合采样路径与参考基线

random_sampling_tl(vocab_size, is_nextn_predict) 处理的是”已经算出 probs 张量”的场景，不包含温度缩放、top-p/top-k 与 log-softmax。它的结构与 fused_sample 的采样阶段几乎一样：

• accumulate 把每一块概率折叠成车道和；
• sample_offset 取第一级的车道与余量；
• sample_select 在选中车道的一列中精定位。

它与融合版的差别在”输入是概率还是 logits”和”是否带验证分支”。主前向路径不再调用它：fused_sample_ref 需要一条独立于融合内核的可执行路径来做逐行对照，测试也需要用它验证定点 CDF 的分布性质；此外它还充当 benchmark 的基线，用来衡量”融合”相对于”先算 probs 再采样”到底省了多少。这三件事决定了它必须保留一份可执行的实现，而不是只留在文档里。

sample_case = 1 if is_nextn_predict else 0 区分了”目标采样”与”残差采样”的随机流，这与融合内核里 sample_positions * 3 + case 的三路编码对齐：单独调 random_sampling 时没有 case=2 的接受判定，因为它不参与验证。

sum_smaller_probs_tl：把”升序 CDF 值”变成一个标量

1
if val < logprob_limit or (val == logprob_limit and col <= sample_result):
2
    thread_sum[0] += T.exp(val)

sum_smaller_probs 的定义是：把一行的 token 按概率升序排列，求”直到采样 token（含）“的概率和。它不输出 token id，而是输出一个标量——采样 token 在升序 CDF 上的累计概率值（不是位置，也不是排名）。

这个量由 sampler.py 里的 DEEPINFER_COMPUTE_SUM_SMALLER_PROBS 控制，默认关闭；开启后随 RequestOutput 回传给上层。它的用途是给采样结果附加一个”这条 token 在分布里的分位位置”式的标量：例如离线分析可以据此判断一次采样是落在头部还是尾部，回放工具可以把它作为确定性校验的附加指纹。它不是生成 token 的必需数据，而是可观测性接口的一部分。

内核的并列规则与阈值筛选共用同一套比较语义：val < limit or (val == limit and col <= sample_result) 表示”概率小于它，或概率相等但下标不大于它”。这等价于把并列 token 按原始下标排序后求前缀和——原始下标是唯一、稳定的排序键，不依赖浮点相等。sum_smaller_probs 因此与融合内核的阈值判定使用同一套并列规则，避免两套实现给出不一致的”升序位置”。

detect_outlier_and_abort：把静默的坏值变成可定位的 abort

1
if T.isnan(value) or value <= -max_abs or value >= max_abs:
2
    T.evaluate(T.call_extern("int32", "printf", "outlier detected: id=%ld, value=%.9g, offset=%d\n", ...))
3
    T.device_assert(False, no_stack_info=True)

outlier_detection.py 把 MAX_ABS 定在 $10^{11}$：这不是一个”业务阈值”，而是用来拦住异常值——NaN、无穷，或者数量级已经超出任何正常激活范围的死值。它扫描任意形状的扁平张量，遇到坏值就打印 id（调用方传入的运行期标识）、fp32 数值和扁平偏移，然后让设备断言失败。

它守的是采样的输入前提。采样内核第一段里已经有一道守卫：

1
T.device_assert(T.isfinite(sum_v) and sum_v > 0, msg="Invalid logits: contains NaN or +inf, or all entries are -inf")

但这条守卫的粒度是”整行”：它只能告诉你”这一行的 logits 不合法”，不会告诉你是哪一层、哪个请求、哪个偏移出了问题。更常见的麻烦是部分位置出现 NaN，而不是整行全坏——此时 sum_v 可能仍是有限值（NaN 参与 max / 求和时行为取决于归约路径），softmax 被污染，采样随机数比较退化成全 false，最终静默地落在一个固定的 fallback token 上。生产上这类故障最难查：不崩、不报错，只是输出分布悄悄漂移。

detect_outlier_and_abort 的价值就在这里：它把”静默的坏值”变成一条带 id 和偏移的 printf，然后显式 abort。id 是调用方给的运行期标识（例如请求序号、层号或步号），用来在日志里指认”是哪一份张量的哪一个位置”；偏移则是扁平数组里的下标，配合形状就能反推出高维坐标。断言失败后 CUDA 上下文不可继续使用，进程需要退出——这是有意的选择：宁可让一个副本下线，也不要让它继续吐出看起来正常、实则可复现性已丧失的 token。

它与采样内核里的 device assert 构成了两层守卫：内层保证不会用坏 logits 产出 token；外层的诊断工具帮你定位坏值的来源。两道守卫都不保护”数值精度”——采样阶段对 logits 的量化是另一套问题，见后文。

动态形状与 JIT 开销

fused_sample_tl(vocab_size, do_sample, do_reject_sampling) 只有三个参数，但 total_q 是动态的。这意味着同一个 (vocab_size, True, True) 编译出来的内核可以服务任意批大小，不需要按批大小重新编译。代价是 T.Kernel(total_q) 的 CTA 数量在运行时才确定，编译器无法做”按批大小调整分块”的优化。

这也解释了为什么采样阶段要把 vocab_size 编译进模板：不同模型的词表不同（常见的 8192、32768、129280），而同一个模型在整个服务生命周期里词表固定。按词表编译只发生一次，之后所有请求复用同一份二进制。

PDL 的语义边界

T.pdl_sync() / T.pdl_trigger() 在 Hopper 与 Blackwell 上有硬件支持；在更老的架构上，tilelang lower 成普通的 kernel 边界同步（没有重叠），功能正确但收益消失。代码依赖 PDL 的地方有三处：两段 kernel 之间、kv_index_cat 等周边内核的输入准备、以及 rgb_pad_to_patches 的图片预处理。PDL 不是正确性前提，而是延迟优化；如果某平台不支持，内核仍然能跑，只是两段之间的启动延迟没法叠起来。

定点量化的硬上界

采样阶段的定点是 31 位。词表越大，累计误差越大：

$V = 65{,}536$ 时约 $3 \times 10^{-5}$；$V = 129{,}280$ 时约 $6 \times 10^{-5}$；$V = 1{,}000{,}000$（百万级词表）时约 $4.7 \times 10^{-4}$。这个误差是分布层面的，不是单个 token 的——它影响的是”采样分布与目标分布的接近程度”。

若将来出现更大的词表，要么增加定点位宽（uint64 路径已在 float_to_fixed 中实现，scale_bits 最高可取 63 位），要么把定点改为”定点 + 残差补偿”。但 uint64 定点会把累加器的寄存器压力翻倍，这又是 sample_acc_frag 在 BD=2048 下能否驻留的关键。这是一个典型的”精度—占用”折中。

T.unroll 的双分支与代码膨胀

1
for reject in T.unroll(2 if do_reject_sampling else 1, explicit=True):
2
    if is_reject_sampling == (reject != 0):
3
        for i0 in T.Pipelined(N_SAMPLE_CHUNKS, num_stages=3 if NUM_STAGES else 0):

两个分支展开后，每个分支都包含完整的 N_SAMPLE_CHUNKS 循环与流水化。若 N_SAMPLE_CHUNKS = 64、SAMPLE_BD = 2048，整个 kernel 的指令数会相当可观。这是为了”同一趟遍历里同时处理普通行与残差行”的代价。好在 T.Pipelined 与分块让每段循环体积可控，而 is_reject_sampling 的谓词让每个分支最多只有一半的行进入循环体。

并列与浮点相等的边界

lf[j] == threshold 是浮点相等比较。threshold 是从 u32 键反解出浮点值的：

1
def sortable_to_f32(e: T.uint32) -> T.float:
2
    return T.reinterpret(T.float, e ^ (T.cast(T.cast(~e, T.int32) >> 31, T.uint32) | SIGN_BIT))

它是 f32_to_sortable 的逆映射：对键做同样的符号翻转，还原出原始浮点。由于键是从原始 logit 精确映射出来的，反解得到的是完全相同的浮点值，所以 == 比较是精确的——只要 logits 本身没有被中途改写。这就是为什么采样阶段的所有 lf[j] 都来自重新读的 logits 或经 (lf + lse1)*t + lse0 计算——只有在”筛选与采样读同一份 logits”时，并列比较才可靠。如果中途有别的 kernel 改写了 logits，阈值比较就会失效。生产路径里没有这类改写，但这是调用方必须遵守的约定。

极端值：+inf、-inf、并行空行

• f32_to_sortable 里正无穷的 raw 是 0x7F800000，翻转符号位后键为 0xFF800000；负无穷的 raw 是 0xFF800000，取反后键为 0x007FFFFF。SORTABLE_POS_INF 与 SORTABLE_NEG_INF 正是这两个值。desired 在细化过程中被 T.clamp(desired_s[0], SORTABLE_NEG_INF, SORTABLE_POS_INF) 夹住——这一段对全 -inf 行尤其重要：如果一整行的 logits 都是 -inf，第一段 kernel 的 device_assert 会先失败，根本走不到这里。但如果只有一部分是 -inf、其余很小，threshold 可能正好落在 -inf 键上，此时 threshold_p = exp(-inf) = 0，代码把 sel_f 置为 0，即”不保留任何 -inf”。
• T.Kernel(total_q) 是”每行一个 CTA”，不存在空行；但 pool_rows（投机信息池行数）可以大于 total_q，残差阶段从 mtp_info_read_rows[row] 取到的 draft_row 是池里的真实行号，可能覆盖整个池。
• sample_positions 为 0 是合法的（序列的第一步），三路分流得到 seq = 0/1/2，与 seed 组合后仍是一个有效的哈希输入。

与 lse_workspace 的生命周期

lse_workspace 是 [total_q, 2] 的临时张量，由调用方分配。它只活在两段 kernel 之间（同一调用内的两次 launch），不会被跨步持有。这与引擎里其他”借用缓冲”的约束一致——比如 SchedulerOutputs 里的 clone_blocks 是调度器内部 vector 的零视图，必须在下次 schedule() 前消费完。lse_workspace 更安全一些，因为它在内核内部的两次 launch 之间传递，调用方看到的只是一个会被覆写的张量。