DeepInfer/H20-TRITON-PLAN.md

H20 Triton 改造计划

把 DeepInfer 在 H20（SM90）上推理所需的 TileLang kernel 全部换成 Triton。

分支：h20-sm90-triton，基于 h20-sm90-support（PR #180）。

零、现状（2026-10-22 09:30 实测）

能说人话。 8×H20、DEEPINFER_OPS_BACKEND=triton、MTP 默认开（nextn_predict=3）：

=== prompt: '问题：中国的首都是哪里？\n答案：'
=== continuation: '北京。\n问题：中国的首都是哪里？…'        # 贪心 64 token 后开始复读提示，正常退化

=== prompt: '请用两句话解释什么是注意力机制。\n答：'
=== continuation: '注意力机制是一种让模型在处理信息时，能够“聚焦”于相关部分而“忽略”无关部分的
技术。它通过计算查询（Query）与键（Key）的相似度来分配权重，再根据权重对值（Value）进行加权求和，
从而得到最终的输出。<｜end▁of▁sentence｜>…'

8 个 DP rank 逐 token 相同，回退提示 5 条（MTP 那 4 个 + 采样器 verify 变体），0 报错。

但「直接这个状态」= 默认配置跑不起来，必须两个 env（两个都不是 Triton 引入的）：

env	为什么必须
EP_DISABLE_GIN=1	容器没挂 /dev/infiniband（host 上有 4 个 uverbs），DeepEP 查 props.ginType 得 NONE 就 assert。SM90 单机本来走 scale-up，不需要 GIN。
DEEPINFER_ENFORCE_EAGER=1	见下面 P0：CUDA graph capture 里 engram 本地 gather 会 D2H。

已知问题清单（按优先级）

#	问题	状态	影响
P0	CUDA graph capture 挂掉：engram_local.py:108 torch.bincount(owner)（owner 在 GPU）在 capture 里触发 CPU↔CUDA 拷贝。且 _exchange_sizes 用 all_to_all_single + 动态 split sizes，本身不可 capture。	未修，用 eager 绕过	默认配置（graph on）warmup 直接失败，8 rank 全挂
P1	MTP 4 个 kernel 未移植：remap_draft_topk / get_accept_result / finalize_verification / get_fused_dflash_get_input	未做	MTP 路径仍有 TileLang 依赖，5 条回退
P1	采样器 rejection-sampling 变体（MTP verify）未实现	显式委托 cuda	同上
P2	导入期报 lacks 20 kernel(s)：这些是可选路径（structure_output 引导解码、image_embed_ops 多模态、prepare_cp_mirror CP>1、schedule_mtp_varlen、state_cache、candidate_blocks、sum_smaller_probs、dflash_copy、outlier_detection 等），当前配置不执行，但一旦打开就回退	未做	功能可用（回退），但「技术栈全 Triton」不成立
P3	窄 cut 采样器 token 差 913%（top_k=128 / top_p=0.5 之类）；引擎配置（top_p=0.99）100%	已定量解释，不修	只影响极端参数，分布本身差 1e-5
P3	token_logprobs 恒 0	未查	与 Triton 无关，历史 TileLang 日志同样为 0；影响 logprob 类接口
P4	性能没测	未做	没有任何 Triton vs TileLang 吞吐数据
—	2.7% logits 漂移	SM90 forced per-128 + fp32 scale 的量化地板	非 Triton 问题，答案稳健性已验（margin 1.638）

一、范围是怎么定下来的

不靠读代码猜，而是实测。ops.py 里加了一个 DEEPINFER_TRACE_KERNELS=<path> 探针 （_KernelTrace 代理每个 *_tl 工厂，第一次被调用时记名），然后跑一次完整文本生成， 得到真实执行面。

24 个 kernel，就是本计划要换的全部：

#	kernel	文件	阶段
1	gather_tl	gather.py	P1
2	gather_input_ids_zos_tl	gather_input_ids_zos.py	P1
3	gather_engram_tokens_zos_tl	gather_engram_tokens_zos.py	P1
4	expand_swa_index_tl	swa_index.py	P1
5	kv_index_cat_tl	kv_index_cat.py	P1
6	convert_kv_index_tl	flash_mla_utils.py	P1
7	load_kv_index_tl	get_kv_index.py	P1
8	map_lookup_tl	sampler.py	P1
9	norm_tl	norm.py	P1
10	copy_block_tables_tl	prepare_decode_metadata.py	P1
11	prepare_v4_decode_metadata_tl	prepare_decode_metadata.py	P1
12	expand_sampling_params_tl	sampler.py	P1
13	store_mtp_info_tl	store_mtp_info.py	P1
14	store_kv_cache_tl	store_kv_cache.py	P2
15	load_paged_kv_cache_tl	load_paged_kv_cache.py	P2
16	fused_norm_rope_hadamard_cast_store_kvcache_tl	fused_transform.py	P2
17	prepare_verification_metadata_tl	sampler.py	P3
18	get_accept_result_tl	sampler.py	P3
19	finalize_verification_tl	sampler.py	P3
20	remap_draft_topk_tl	sampler.py	P3
21	confidence_head_fused_tl	dflash.py	P3
22	dflash_input_store_tl	dflash.py	P3
23	get_fused_dflash_get_input_tl	get_kv_index.py	P3
24	fused_sample_tl	fused_sample.py	P4

不在此列的（vision、structured output、engram 主表、compress、candidate_blocks 等） 在这条路径上没被执行，不在本次范围。

二、关键架构决策：并行后端，而不是原地替换

不做原地改写。 新建 tilelang_ops/triton/ 包，与 cuda/ 平级，暴露同名 *_tl 工厂； ops.py 按 DEEPINFER_OPS_BACKEND 环境变量选后端（默认仍是 cuda，保持现状不回归）。

理由：

1. 每个 kernel 都能 A/B 对照。同一输入跑 TileLang 版和 Triton 版，比输出。这比只和 Python ref 比强得多 —— TileLang 版是已经在整机验证过的基线。
2. 失败可回退。某个 kernel 没达标时，可以单独把那个名字指回 cuda 后端（见下）， 其余 23 个照样走 Triton，端到端仍然能跑。这样交付是渐进的，不是全有或全无。
3. review 友好。PR 里能直接看到「同一契约、两套实现」，而不是一大片删除。

ops.py 的解析顺序：DEEPINFER_OPS_BACKEND（整包） → DEEPINFER_OPS_FALLBACK（逗号分隔的 kernel 名，回退到 cuda）。

三、验证协议（每个 kernel 都必须过三关）

1. 契约一致：Triton 版与 TileLang 版在真实 shape 上，输出逐元素一致（浮点按 dtype 容差；索引/计数类要求精确相等）。
2. 边界覆盖：空输入、-1 padding、跨页、多段、奇偶长度、超出范围的值 —— 按各 kernel 自己的语义，从 TileLang 版的测试用例里继承（cuda/*.py 里已有 test_* 参数化矩阵，直接复用）。
3. 不退化：不能因为它们跑通就悄悄降级语义（例如把排序改成不排序、把原子累加 改成串行），必须保持与 TileLang 版相同的契约。

四、阶段与顺序

依赖顺序决定阶段顺序：底层 gather/metadata 在前，cache 在后，采样最后。

• P0 基础设施：triton/ 包骨架、后端选择与 fallback、A/B 对照脚本。
• P1 轻量 kernel（13 个）：纯搬运/索引/逐元素，无矩阵乘。先易后难建立信心。
• P2 KV cache（3 个）：量化写、分页读、fused norm+rope+hadamard+cast+store。 store_kv_cache / load_paged_kv_cache 已有闭环 round-trip 脚本可复用 （scripts/h20_kv_roundtrip.py），直接拿来验 Triton 版。
• P3 MTP 链路（7 个）：验证元数据、接受判定、draft top-k 重映射、confidence head。 这条链路只在 MTP 开启时执行，需要额外用默认配置跑。
• P4 采样（1 个）：fused_sample_tl，最复杂的一个（top-k + 温度 + 拒绝采样）。
• P5 端到端：h20_ask.py 输出与 TileLang 版逐 token 一致，两种配置都过 （DEEPINFER_USE_NEXTN_DECODE=0 和默认开启）。

五、已知风险与对策

风险	对策
T.pdl_sync()（Hopper PDL）在 Triton 没有直接对应	该同步保证后续 kernel 读到已写数据；Triton 侧去掉依赖 PDL 的假设，改用 stream 顺序语义。若某 kernel 依赖 PDL 才能正确，就在 Python 侧显式 wait_event，不靠隐式
TileLang T.dynamic(...) 动态 shape	Triton 用 constexpr 分块 + 运行期循环边界，shape 从 tensor 推导；避免为每个 shape 编译
T.StridedTensor 非连续视图	Triton 传入 stride 作为参数，或先 contiguous()（会记录在注释里并计入性能）
原子操作	atomic_add 在 Triton 可用，但需注意非确定性；若原 kernel 有确定性保证，需复现（例如按固定顺序 reduce）
fp8 / fp4 位级操作	Triton 支持 tl.float8e4nv 等；fp4 打包需手工位运算，参考 quant.py 的现有实现
编译期开销	Triton 首次编译慢；用 tl.constexpr 收敛变体数量，必要时预编译缓存

六、完成判据

• 24 个 kernel 全部有 triton/ 实现，且 A/B 对照通过；
• DEEPINFER_OPS_BACKEND=triton 下 h20_ask.py 端到端输出与 TileLang 版逐 token 相同 （MTP 开/关两种配置）；
• ruff check --no-cache 与 ruff format --check --no-cache 通过；
• 性能不要求超过 TileLang，但需测量并记录差距（附 benchmark 表）。

七、进度

22 个 kernel 已移植，A/B 全过

验证方式：python scripts/h20_ops_ab.py（可带 kernel 名子集）。当前 22 个 kernel、约 140 个用例， 最近一次全量运行是 AB_OK。

kernel	用例数	关键覆盖
gather_tl	9	int8/bf16 载荷、负索引、strided 行、空输入
gather_input_ids_zos_tl	5	remap 重放 vs fresh
gather_engram_tokens_zos_tl	4	n-gram 窗口、image sentinel
expand_swa_index_tl	5	两段区间、溢出 width
kv_index_cat_tl	3	拼接 + -1 保持
convert_kv_index_tl	3	逻辑→物理页表
load_kv_index_tl	5	prefill 连续 / decode 分页
store_mtp_info_tl	5	has_index/has_confidence、slot 0
norm_tl	8	bf16/fp8/residual、d=512..7168
map_lookup_tl	3	by_index / by_key
expand_sampling_params_tl	3	5 个输出逐一比对
prepare_verification_metadata_tl	6	含 num_steps=0、超单 tile
copy_block_tables_tl	3	多 chunk、非 2 幂
store_kv_cache_tl	5	ue8m0 + fp32 sf、-1 槽位
load_paged_kv_cache_tl	4	do_cast、index 直通、洞
fused_norm_rope_hadamard_cast_store_kvcache_tl	18	4 种输出形态 + 主/索引 KV 写入
compress_and_update_state_prefill_tl	5	跨 chunk 窗口 + 环形 cache 尾写
compress_and_update_state_decode_tl	4	逐 token 写 + 块收口
prepare_v4_decode_metadata_tl	5	全部 14 个输出，含 ratio 分桶
fused_sample_tl	7	logprobs + token，engine 配置 100% 一致（见下节）
compute_seqlens_k_zos_tl	6	fresh / 续接 / 小 prev batch
confidence_head_fused_tl	5	有/无 markov、低温 clamp
dflash_input_store_tl	4	hc_mult 1/2/4

端到端：MTP 关闭时只剩 1 个 kernel 走 TileLang

DEEPINFER_OPS_BACKEND=triton + DEEPINFER_USE_NEXTN_DECODE=0 跑 8×H20：

token ids: [17607, 2793, 3139, 666, 6127, 666, 320, 1]

与纯 TileLang 逐 token 相同，8 rank 一致。更关键的是回退提示现在是使用期报告的：

[ops] kernel 'fused_sample_tl' used from 'cuda': 'triton' does not implement it

即非 MTP 路径 19 个 kernel 里 18 个已经真的跑在 Triton 上，唯一剩下的是 fused_sample_tl。

但要说清楚一件事：随机采样会把 1 ulp 的差异放大成不同的续写

同一个自然语言 prompt 跑两次，一次纯 TileLang、一次 Triton 后端：

TileLang: 北京。\n问题：中国的首都是哪里
Triton  : 北京。\n分析：这是一个基础地理常识

前两个 token 相同（北京、。），第三个开始分叉。两边各自 8 个 rank 都是内部一致的，所以不是 不确定性，而是两个后端的 CDF 有一点点不同，于是抽到了不同的分支。

原因在 P4 已经量化过：H20 的 per-128 量化让 logits 有 3e-2 量级的差异；而 kernel 换实现之后， 连 bf16 归约顺序都会变（norm 的 A/B 里 655360 个元素有 10 个差 1 ulp，fp8 字节和 scale 完全一致）。 argmax 不受影响（token-ids 那次 8 个 token 全对），但 T=1.0 的采样相当于把这个 ulp 级差异 放大成一次不同的硬币。

所以「端到端一致」的准确说法是：

输入	结果
官方 chat-template token ids，8 token	逐 token 相同，8 rank 一致
自然语言 prompt	首个 token 相同，续写在采样分支上分叉

这不是移植错误，也不能靠「再多测几次」消掉 —— 只要采样是随机的、且两个实现的浮点归约顺序不同， 就会有分支差异。若要连采样结果都一致，必须同时做到：bit-exact 的归约顺序（Triton 做不到，它有 自己的 reduction tree）+ 同一条 PRNG 流（见下一节）。

剩余：MTP 路径还有 4 个

MTP（DEEPINFER_USE_NEXTN_DECODE != 0，默认开启）路径上一共 7 个 kernel，已移植 3 个。

kernel	状态	说明
compute_seqlens_k_zos_tl	✅	6 用例
confidence_head_fused_tl	✅	5 用例
dflash_input_store_tl	✅	4 用例
remap_draft_topk_tl	⬜	重。要做 draft top-k 空间 → target top-k 空间的 logprob 重映射，baseline 用一块 25856/44032 项的 shared 池按 vocab 分块扫。Triton 没有共享内存池原语，照搬要换实现（位图或 global scratch），不能机械翻译。
get_accept_result_tl	⬜	接受判定，带 has_golden 变体
finalize_verification_tl	⬜	验证收尾，带 has_sum_smaller_probs / has_confidence / num_steps 三个变体维度
get_fused_dflash_get_input_tl	⬜	MTP draft 输入组装

这 4 个当前自动回退到 TileLang（_FallbackNotice 会各自打印一行），所以 MTP 能跑，只是路径上 还有 TileLang 依赖。要彻底达成「技术栈换 Triton」，这 4 个是剩下的全部工作。

另外采样器的 rejection sampling 变体（MTP verify 路径）也没实现，需要 accept/reject 比较、 residual mass、拒绝后的第二次 draw。它是显式委托而不是静默回退 —— fused_sample 是本模块 导出的，_FallbackNotice 覆盖不到，所以工厂自己打印一行说明。

MTP 开启时的实测回退清单（完整且自洽）

[ops] fused_sample: rejection sampling (MTP verify) is not implemented in the triton backend, ...
[ops] kernel 'finalize_verification_tl' used from 'cuda'
[ops] kernel 'get_accept_result_tl' used from 'cuda'
[ops] kernel 'get_fused_dflash_get_input_tl' used from 'cuda'
[ops] kernel 'remap_draft_topk_tl' used from 'cuda'

就是上面那 4 个加上采样器的 verify 变体，没有第六项。且该配置下 token 输出仍然正确：

token ids (8): [17607, 2793, 3139, 666, 6127, 666, 320, 1]   8 rank 一致，0 报错

fused_sample_tl 已移植，但「token 一致」有明确边界

先解掉了唯一真正的障碍：随机流。T.rng_init/T.rng_rand 降级成 curand_init/curand， generator 默认 curandStatePhilox4_32_10_t。Philox 的混合函数是公开的，但 curand_init 的 seeding 没有文档，所以是实测反推的（scripts/h20_rng_dump.py dump 真值）：

key = (seed & 0xFFFFFFFF, seed >> 32)
ctr = (offset & 0xFFFFFFFF, offset >> 32, subsequence & 0xFFFFFFFF, subsequence >> 32)
每次 draw 依次取 block 的 4 个 word，然后 ctr.x += 1

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。

采样器本体（单 tile，vocab=4096，与引擎实际构造一致）已移植，A/B 结果：

配置	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

保留集合逐列一致，logprobs 差 1e-5（= 行归约顺序不同）。token 在窄 cut 下有 ~10% 不同， 原因已定量定位，不是实现缺陷：

r = draw % total,  total = sum(floor(p_i * 2^31))

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）。

结论：采样器本身是对的（分布一致到 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%。