DeepInfer 源码完全拆解（七）：TileLang 特化算子与双平台——一套接口如何同时落到 CUDA 与 Ascend

前六篇讲到哪里，这一篇补哪一块

DeepInfer 是 DeepSeek 官方开源的高性能推理引擎，这一系列沿它的源码逐模块拆解。前六篇已经把”引擎骨架”讲完了：

• （一）整体架构与模块全景：仓库地形、entrypoints/llm.py 的装配线、一次请求如何穿过引擎，以及 DeepInfer 相对 vLLM/SGLang/TensorRT-LLM 最有辨识度的六条设计。
• （二）异步调度与 ZOS：Worker.execute_model_step 的步编排、seqlens_q 的 GPU→CPU 修正、运行时耗时估计器如何把实测耗时喂回调度器。
• （三）KV 内存与 SWA 有界重放：cache_engine.py 的两本块账、YoCo 截断的 mask/compact 双模式与等价性验证方法。
• （四）MoE 专家并行与 EPLB：logical/physical expert 的放置算法、冗余专家的取整规则与在线重排。
• （五）投机解码：DSpark 变长验证、DFlash 草稿头与置信度调度。
• （六）结构化输出与确定性：语法树异步编译、投机回滚与 batch-invariant。

这六篇里反复出现同一类动作：在 GPU 上把一批形状不规则的索引算出来、把一批 token 搬来搬去、把一批 logits 归约成 token。前六篇一直把它们当作”某个函数”一笔带过。这一篇把这层打开：DeepInfer 把这类算子集中放在 deepinfer/model_executor/layers/tilelang_ops/ 下，用 TileLang 写了一套自定义 kernel，而且同一套 Python 接口要同时落到 NVIDIA CUDA 与 华为 Ascend 两套后端上。

这一层有多少东西？tilelang_ops/ 顶层有 24 个模块，加上 cuda/ 与 asc/ 两个子目录，合计约 33 个文件、23145 行 kernel 代码。本篇不追求讲完每一个（那需要一整本书的篇幅），而是先把分发机制、模块分类、双平台对应关系三件事讲透，再挑几个最有代表性的 kernel 逐行拆解。

为什么推理引擎需要”自家的一套小算子”

TileLang 是什么，为什么选它

TileLang 是 tile 级 kernel DSL（官方仓库、官方文档），DeepInfer 依赖的是 tilelang==0.1.15（见 requirements.txt）。它的官方文档把编译流程描述得很清楚：

（图源：TileLang 官方文档 The Tile Language: A Brief Introduction）

关键在于这份描述与硬件后端解耦：从 IRModule 开始，TileLang 生成的是 “target-specific source code”，官方文档写的是 “C/CUDA/HIP/LLVM/…”，对目标平台的表述是 “tuned for the desired backends or GPU architectures (e.g., NVIDIA, AMD)“。文档首页的 API 索引里直接存在一整套 tilelang.ascend.* 模块（ascend/backend、ascend/codegen、ascend/language/dma、ascend/language/gemm_op 等），说明 Ascend 后端是 TileLang 的一等公民而不是外挂。

这正是 DeepInfer 能同时支持 “NVIDIA SM100 GPUs” 与 “Huawei Ascend 950 NPUs” 的技术基础。

TileLang 还提供了三个层次的编程接口，这个分层在本篇后面讲”为什么两个后端写法不同”时很重要：

层次	面向	提供的抽象
Beginner（硬件无关）	只关心计算逻辑	纯 tile 级描述，不管内存层次（官方标注尚未完全实现）
Developer（硬件感知 + Tile Library）	了解 GPU 内存层次的开发者	预置的 tile 级原语与优化模式，不碰线程细节
Expert（硬件感知 + 线程原语）	需要精细控制的专家	直接给出线程原语与内存原语，可手工控制布局与同步

官方文档明确说这三种接口可以在同一个 kernel 里混用。这一点在 DeepInfer 的代码里体现得非常明显——下面会看到同一个 kernel 里既有 T.copy 这种 Tile Library 级原语，也有 T.Parallel 这种线程级描述。

分发机制：ops.py 如何把后端差异关进一个文件

先看整个子系统最核心的设计——它只有 38 行，但决定了 33 个 kernel 模块如何被调用。

1
"""Kernel dispatcher: interface modules do
2
``from deepinfer.model_executor.layers.tilelang_ops import ops`` and call
3
``ops.xxx_tl``; backend selection happens only in this file.
4

5
Re-exports all ``xxx_tl`` factories of the selected backend.
6
"""
7

8
import importlib
9
import pkgutil
10
from types import ModuleType
11
from typing import Any
12

13
from deepinfer import PLATFORM
14

15
_pkg = importlib.import_module(f"{__package__}.{PLATFORM.tilelang_ops_backend}")
16
for _info in pkgutil.iter_modules(_pkg.__path__):
17
    _mod = importlib.import_module(f"{_pkg.__name__}.{_info.name}")
18
    for _name, _obj in vars(_mod).items():
19
        if _name.startswith("_") or isinstance(_obj, ModuleType) or not _name.endswith("_tl"):
20
            continue
21
        assert globals().get(_name, _obj) is _obj, f"backend symbol collision: {_name}"
22
        globals()[_name] = _obj
23

24

25
class KernelUnavailableError(AttributeError):
26
    """Selected backend lacks this kernel factory; conftest.py turns it into
27
    a test skip on non-cuda backends. Subclasses AttributeError to keep the
28
    module ``__getattr__`` / ``hasattr`` protocol intact."""

逐点拆开：

（1）后端目录名来自平台对象。 PLATFORM.tilelang_ops_backend 是一个字符串，在 platform.py 里定义：PlatformGPU.tilelang_ops_backend = "cuda"，NPU 平台对应 "asc"。ops.py 自己不知道有哪些后端，它只负责 importlib.import_module("...tilelang_ops.cuda") 或 "...tilelang_ops.asc"。新增一个后端只需要在 platform.py 里加一个类、在目录下加一个子包，ops.py 一行不改。

（2）发现是自动的。 pkgutil.iter_modules 遍历选定后端目录下的所有模块并逐个 import。这意味着”这个后端支持哪些 kernel”不是维护出来的清单，而是目录内容本身——加一个文件就多一个 kernel，删一个就少一个。

（3）命名约定就是接口契约。 筛选条件是 _name.endswith("_tl") 且不是下划线开头、不是模块对象。所以每个后端模块必须把自己对外的 kernel 工厂命名为 xxx_tl，其余符号（辅助函数、常量）自动被跳过。后缀 _tl 显然指 TileLang，它同时起两个作用：在调用侧一眼看出”这是 TileLang kernel 不是普通 PyTorch 函数”，在分发侧充当命名空间过滤。

（4）碰撞检测写在导入期。 assert globals().get(_name, _obj) is _obj, f"backend symbol collision: {_name}"。如果两个模块定义了同名 xxx_tl，导入时立刻断言失败并报出名字，而不是等到运行时某个调用静默走到错误的实现。这是一个很划算的设计：把”重名”这个低级错误从运行时的疑难杂症变成导入时的明确崩溃。

（5）缺失也有明确语义。 KernelUnavailableError 继承 AttributeError，注释点明了两个用途：一是保持 hasattr / __getattr__ 协议正常（属性不存在就该抛 AttributeError 的子类）；二是让测试框架 conftest.py 能把”当前后端没这个 kernel”转成 skip 而不是 fail。这解释了为什么 asc/ 目录下模块数量与 cuda/ 不完全一致——Ascend 后端缺少某些 kernel 是预期状态，不是 bug。

调用侧因此长这样，完全没有平台分支：

1
from deepinfer.model_executor.layers.tilelang_ops import ops
2

3
def some_layer_forward(...):
4
    ops.gather_input_ids_zos_tl(...)   # CUDA 上走 cuda/gather_input_ids_zos.py
5
    ops.load_paged_kv_cache_tl(...)    # Ascend 上走 asc/load_paged_kv_cache.py

对比另一种常见写法——在每个调用点写 if platform.is_gpu: ... else: ...。那种写法的问题不只是丑：它把平台知识散布到几十个调用点，每加一个后端就要改几十处，而且容易漏。DeepInfer 的做法是把平台知识收敛到一个文件加一个目录结构，调用点只认 ops.xxx_tl 这一个契约。这是典型的”间接层换可维护性”。

模块分类：33 个 kernel 在干什么

打开后端目录（以 cuda/ 为例），模块可以按职责分成六族。下表按族列出代表模块与它解决的问题——这张表就是这一层的源码地图。

几个观察：

带 _zos 后缀的模块直接对应第二篇的零开销调度。 gather_input_ids_zos.py、gather_engram_tokens_zos.py、compute_seqlens_k_zos.py 这三个是”延后一拍”编排的直接产物——原本需要 CPU 先算好再喂给 GPU 的东西，现在改成 GPU 上自己算，从而把 CPU 后处理移出关键路径。后缀是它在代码里的标记。

fused_* 前缀是融合的信号。 fused_sample.py（CUDA 524 行 / Ascend 673 行）和 fused_transform.py（CUDA 619 行 / Ascend 899 行）是把”变换 + 过滤 + 采样”多个步骤塞进一个 kernel。这是小算子层最重要的价值来源：不是让每个小步骤更快，而是让多个小步骤共享一次 launch 和一次显存往返。

注意顶层模块与后端模块的分工。 structure_output.py、state_cache.py、schedule_mtp_varlen.py、gather_input_ids_zos.py 等文件在顶层也有一份，同时 cuda/ 下也有一份。看顶层那份的内容就明白了：它不是 kernel，而是Python 侧的封装函数 + 参考实现 + pytest 用例。以 schedule_mtp_varlen.py 为例，它定义了 compute_seqlens_q_by_confidence_rates()（调 ops.xxx_tl）、compute_seqlens_q_by_confidence_rates_ref()（纯 PyTorch 参考实现），以及参数化的 test_* 用例断言两者 torch.equal。顶层管”这个算子的语义是什么、怎么验证它正确”，后端目录管”在当前硬件上怎么把它算出来”。 这个分层贯穿整个 tilelang_ops/。

双平台对应关系：接口层与实现层怎么分家

先看目录结构本身传递的信息：

1
tilelang_ops/
2
├── ops.py                 ← 分发器（唯二与平台相关的地方）
3
├── conftest.py            ← 测试框架：GPU 绑定、benchmark 基线配置
4
├── common.py              ← 共享：量化常量与数学辅助
5
├── cache_common.py        ← 共享：KVCacheFormat 与 KV 布局定义
6
├── schedule_mtp_varlen.py ← 共享：算子语义 + 参考实现 + 测试
7
├── gather_input_ids_zos.py← 共享：同上
8
├── structure_output.py    ← 共享：同上
9
├── ...（其余语义/测试模块）
10
├── cuda/                  ← NVIDIA 后端实现（28 个模块）
11
│   ├── gather_input_ids_zos.py
12
│   ├── load_paged_kv_cache.py
13
│   ├── fused_sample.py
14
│   ├── quant.py           ← CUDA 特有
15
│   └── ...
16
└── asc/                   ← Ascend 后端实现（27 个模块）
17
    ├── gather_input_ids_zos.py
18
    ├── load_paged_kv_cache.py
19
    ├── fused_sample.py
20
    ├── kv_quant.py        ← Ascend 特有
21
    ├── rstd_1ulp.py       ← Ascend 特有
22
    └── ...

两个后端目录的差异精确到三个文件：

1
=== 只在 cuda/ 有 ===         === 只在 asc/ 有 ===
2
outlier_detection.py          kv_quant.py
3
prepare_decode_metadata.py    prepare_v4_decode_metadata.py
4
quant.py                      rstd_1ulp.py

对上层的”接口模块”（比如模型里的某个 layer）来说，它 import 的是 ops，调用 ops.xxx_tl。它既不知道后端是 CUDA 还是 Ascend，也不需要知道。

但接口的粒度是关键设计决策。看行数对比：cuda/fused_sample.py 524 行，asc/fused_sample.py 673 行；cuda/fused_transform.py 619 行，asc/fused_transform.py 899 行。同名文件在两边的行数差 20%~45%，说明它们不是同一份代码的复制粘贴，而是各自按硬件的并行模型与 intrinsics 写了实现，只对外暴露同样的 _tl 工厂名和同样的参数签名。

对同一个算子 gather_input_ids_zos，两边的写法差异是理解这套抽象的最佳样本。CUDA 版：

1
with T.Kernel(n_seq, threads=128) as seq_idx:
2
    T.pdl_sync()
3
    sq = seqlens_q[seq_idx]
4
    q_base = T.if_then_else(seq_idx == 0, 0, cu_seqlens_q[seq_idx - 1])
5
    pidx = remap[seq_idx]
6

7
    for j in T.Parallel(max_seqlen_q):
8
        if j < sq:
9
            if pidx >= 0:
10
                if j == 0:
11
                    out_input_ids[q_base + j] = prev_sampled_token_ids[pidx]
12
                else:
13
                    out_input_ids[q_base + j] = prev_predict_token_ids[pidx, j - 1]
14
            else:
15
                out_input_ids[q_base + j] = fresh_tokens[seq_idx, j]

一个 block 一条序列（T.Kernel(n_seq, threads=128)），线程在 max_seqlen_q 上并行，直接读写全局内存。这依赖 GPU 的 SIMT 模型：几百个线程各自算自己的下标、各自访存，靠硬件做合并。

Ascend 版：

1
num_cores = tile_kernels.get_num_vec_cores()
2
BR = 8
3
BS = 8
4
threads = BR * BS
5
num_stages = 4
6
ALIGN_I32 = 32 // 4
7
out_len_aligned = T.ceildiv(BR * max_seqlen_q, ALIGN_I32) * ALIGN_I32
8

9
with T.Kernel(num_cores) as pid:
10
    sq_ub = T.alloc_shared((ALIGN_I32,), "int32")
11
    remap_ub = T.alloc_shared((ALIGN_I32,), "int32")
12
    qbase_ub = T.alloc_shared((ALIGN_I32,), "int32")
13
    out_ub = T.alloc_shared((out_len_aligned,), "int32")
14
    T.annotate_buffer_versions({sq_ub: num_stages, remap_ub: num_stages, ...})
15

16
    for block in T.Persistent([T.ceildiv(n_seq, BR)], num_cores, pid, group_size=1, num_stages=num_stages):
17
        ...
18
        T.copy(seqlens_q[blk_start : blk_start + blk_len], sq_ub[:blk_len])
19
        ...
20
        with T.SimtVF(threads=threads):
21
            for j1 in T.serial(0, max_seqlen_q, BS):
22
                for i, j in T.Parallel(BR, BS):
23
                    ...
24
                    out_ub[out_off] = ...
25
        T.copy(out_ub[:blk_total], out_input_ids[block_qbase : block_qbase + blk_total])

差异是结构性的，不是风格性的：

为什么会差这么多？ 因为两种硬件的存储层次不同。GPU 的 shared memory 是可选的优化手段，不显式管理也能跑（走 L1/L2）；Ascend 的 AI Core 则要求计算前先把数据搬进 UB，这是硬约束，不搬就跑不起来。于是”多缓冲 + 分段拷贝”在 Ascend 版里成了必需，在 CUDA 版里可以省略。

这就是为什么接口必须停在”算法级”。 如果用更低的抽象统一两边（比如强制两边都写显式 UB 拷贝），CUDA 侧会白白损失性能；如果用更高的抽象（一份代码两端编译），UB 大小、核数、对齐这些硬件约束根本表达不出来。三种方案对比：

DeepInfer 选了 C，证据有三处：ops.py 注释明确写 “backend selection happens only in this file”；common.py 的量化常量、cache_common.py 的 KVCacheFormat 放顶层共享；KernelUnavailableError 专门处理后端缺 kernel——如果两个后端必然一一对应，就不需要这套机制。

代价也很明确：同名文件必须行为一致。这正是顶层那些 *_ref 参考实现与 pytest 用例存在的意义：语义在顶层定义一次、验证一次，两端的 kernel 都要对得上。

逐 kernel 拆解：三类典型算子

下面挑三类最有代表性的 kernel，把”索引算术”到底难在哪讲清楚。选择标准是：它们分别对应第二篇（ZOS）、第三篇（KV 分页）、第五篇（投机解码）讲过的机制，正好把前几篇的抽象层落到具体代码。

1
fresh_tokens          # (n_seq, max_seqlen_q) int32  本步新算出的 token id
2
prev_sampled_token_ids# (prev_n_seq,)        int32  上一步采样出的 token
3
prev_predict_token_ids# (prev_n_seq, nextn)  int32  上一步投机预测的草稿 token
4
remap                 # (n_seq,)             int32  -1 或上一步批次索引
5
seqlens_q             # (n_seq,)             int32  本步每条序列的 query 长度
6
cu_seqlens_q          # (n_seq,)             int32  累积长度
7
out_input_ids         # (total_q,)           int32  扁平化输出

1
pidx = remap[seq_idx]
2
if pidx >= 0:                       # 这条序列上一步存在，且投机解码开了
3
    if j == 0:
4
        out_input_ids[q_base + j] = prev_sampled_token_ids[pidx]      # 第一个 token 是上一步采出来的
5
    else:
6
        out_input_ids[q_base + j] = prev_predict_token_ids[pidx, j-1] # 后面都是上一步草稿预测的
7
else:
8
    out_input_ids[q_base + j] = fresh_tokens[seq_idx, j]              # 新序列，用现成的 token

1
with T.Kernel(T.ceildiv(n_seq, BLOCK_SIZE), threads=BLOCK_SIZE) as bx, T.Parallel(BLOCK_SIZE) as tx:
2
    sk = T.alloc_var("int32")
3
    seq_idx = bx * BLOCK_SIZE + tx
4
    T.pdl_sync()
5
    if seq_idx < n_seq:
6
        sk = seqlens_k[seq_idx] + seqlens_q[seq_idx]
7
        prev_idx = remap[seq_idx]
8
        if prev_idx != -1:
9
            sk = sk - nextn + prev_num_accepted_tokens[prev_idx]
10
        seqlens_k[seq_idx] = sk

1
def load_paged_kv_cache_tl(
2
    layout: KVCacheFormat,
3
    block_size: int,
4
    fp4_value_block_stride: int = 0,
5
    fp4_value_stride: int = 0,
6
    fp4_scale_stride: int = 0,
7
) -> object:
8
    """Build a compact-range paged-cache load kernel from ``layout``."""

1
V41_FORMAT = KVCacheFormat(
2
    quantized_dtype="e4m3",         # FP8 E4M3
3
    quantized_elements=512,         # 每个 token 512 个量化元素
4
    elements_per_sf=32,             # 每 32 个元素共享一个 scale
5
    bf16_elements=0,                # 不再支持 FP8/BF16 混合缓存
6
    do_cast=True,                   # 读取时需要反量化
7
    sf_format="ue8m0",              # scale 用 UE8M0（8 位指数）
8
    note="V41, main/swa kv.",
9
    token_interleaved=True,         # 每个 token 内部按 [数据 | scale] 交错存放
10
)
11
V41_FP4_FORMAT = KVCacheFormat(
12
    quantized_dtype="e2m1",         # FP4 E2M1
13
    quantized_elements=512,
14
    elements_per_sf=16,
15
    do_cast=True, sf_format="e4m3",
16
    note="V41, FP4 main kv.", token_interleaved=True,
17
)
18
MXFP8_INDEX_KV_FORMAT = KVCacheFormat(
19
    quantized_dtype="e4m3", quantized_elements=128,
20
    elements_per_sf=32, do_cast=False,   # 索引缓存不反量化，直接当索引读
21
    sf_format="ue8m0", note="V41, index kv.",
22
)

1
is_quantized = layout.quantized_dtype is not None
2
is_index     = is_quantized and not layout.do_cast
3
is_plain_bf16= not is_quantized
4
is_fp4_cast  = layout.do_cast and layout.quantized_dtype == "e2m1"
5

6
vector_size = min(group_size, 32, value_token_alignment)
7
token_tile  = block_size if is_index else min(block_size, 8 if is_plain_bf16 else 32)
8
threads     = 128 if is_index and layout.quantized_dtype == "e4m3" and block_size == 32 else 256

1
load_segments: T.Tensor((num_segments, 4), T.int32)
2
# 每一行 = (src_block_id, src_begin, src_end, dst_offset)
3

4
with T.Kernel(num_segments, num_token_tiles, threads=threads) as (segment_idx, token_tile_idx):
5
    T.pdl_sync()
6
    src_block_id = load_segments[segment_idx, 0]
7
    src_begin    = load_segments[segment_idx, 1]
8
    src_end      = load_segments[segment_idx, 2]
9
    dst_offset   = load_segments[segment_idx, 3]
10
    ...
11
    for token_offset, group_chunk in T.Parallel(token_tile, valid_groups * channel_chunks):
12
        block_offset = token_base + token_offset
13
        if block_offset < block_size and src_begin <= block_offset and block_offset < src_end:
14
            dst = dst_offset + block_offset - src_begin
15
            ...
16
            if src_block_id == -1:
17
                for channel in T.vectorized(vector_size):
18
                    kv[dst, channel_base + channel] = 0     # 无效块填 0

1
if is_quantized:
2
    value_dim, value_dtype = quantized_elements, T.float4_e2m1fn
3
    if is_fp8:
4
        value_dtype = T.float8_e4m3fn
5
    valid_groups = (quantized_elements + layout.elements_per_sf - 1) // layout.elements_per_sf
6
    sf_dim = layout.sf_storage_elements
7
    sf_dtype = {"ue8m0": "uint8", "e4m3": "float8_e4m3fn", "fp32": "float32"}[layout.sf_format]
8
    blk, threads = min(max(4096 // input_dim, 1), 16), 64
9
else:
10
    value_dim, value_dtype = input_dim, "bfloat16"
11
    blk, threads = max(4096 // input_dim, 1), 256

1
@T.macro
2
def _amax_to_ue8m0_scales(blk, groups, amax, scale_local, sf_ue8m0):
3
    """Convert bf16 amax values to fp8 quant multipliers and ue8m0 scales."""
4
    for i, j in T.Parallel(blk, groups):
5
        scale = T.max(T.cast(amax[i, j], "float32"), FP8_AMAX_MARGIN) / FLOAT8E4M3_MAX
6
        exp_scale = ue8m0_exp_scale(scale)
7
        scale_local[i, j] = T.reinterpret("float32", (exp_scale + 127) << 23)
8
        sf_ue8m0[i, j] = T.cast(-exp_scale + 127, "uint8")

1
FP4_AMAX_MARGIN = 6.0 * (2**-126)
2
FLOAT4_E2M1_MAX = 6.0
3
FP8_AMAX_MARGIN = 1e-4
4
FLOAT8E4M3_MAX = 448.0
5
FLOAT8E4M3_MIN_POSITIVE = 2**-9

1
def compute_seqlens_q_by_confidence_rates(
2
    total_q: int,
3
    remap: torch.Tensor,             # [num_seqs] int32，上一步批次索引或 -1
4
    prev_confidence_rates: torch.Tensor,  # [n_prev_seq, nextn] float32，上一步 MTP 头的逐位置置信度
5
    seqlens_q: torch.Tensor,         # [num_seqs] int32 (out)，本步每条序列的 query 长度
6
) -> None:
7
    """Distribute total_q tokens across sequences based on MTP confidence rates.
8

9
    Fuses remap scatter: confidence_rates[i] = prev_confidence_rates[remap[i]] if remapped else 0.0.
10
    """

1
# 阶段 1：按 remap 把上一步置信度捞过来（新序列视为 0），再沿 nextn 维做累积乘积
2
mask = remap != -1
3
safe_remap = remap.clamp(min=0)
4
gathered_rates = prev_confidence_rates[safe_remap]
5
gathered_rates = torch.where(mask.unsqueeze(1), gathered_rates, torch.zeros_like(gathered_rates))
6
cum_rates = torch.cumprod(gathered_rates, dim=1)     # [num_seqs, nextn]
7

8
# 阶段 2：把 (序列, 位置) 全部展平，按累积置信度全局排序，前 extra_budget 名获胜
9
extra_budget = total_q - num_seqs                     # 每条序列先保底 1 个 token
10
flat_rates = cum_rates.reshape(-1)
11
sorted_indices = torch.argsort(flat_rates, descending=True, stable=True)
12
ranks = torch.empty(len(flat_rates), device=...); ranks[sorted_indices] = torch.arange(...)
13
wins = (ranks.reshape(num_seqs, nextn) < extra_budget).to(torch.int32).sum(dim=1)
14
seqlens_q.copy_(1 + wins)                             # 每条序列 = 1 个保底 + 赢到的名额

1
def fused_sample_tl(vocab_size: int, do_sample: bool, do_reject_sampling: bool) -> Any:
2
    assert vocab_size >= 32, "fused_sample requires vocab_size >= 32"
3
    assert not do_reject_sampling or do_sample, "reject sampling requires sampling"
4
    ...
5
    if vocab_size >= 8192:
6
        THREADS, LSE_BD, SAMPLE_BD = 256, 8192, 2048
7
    else:
8
        LSE_BD = tilelang.next_power_of_2(vocab_size)
9
        THREADS = min(128, LSE_BD)
10
        SAMPLE_BD = LSE_BD
11

12
    FILTER_BD = min(LSE_BD, 4096)
13
    TIE_BD = min(FILTER_BD, 256)
14
    TIE_GROUPS = FILTER_BD // TIE_BD
15
    N_FILTER_CHUNKS = T.ceildiv(vocab_size, FILTER_BD)
16
    N_LSE_CHUNKS    = T.ceildiv(vocab_size, LSE_BD)
17
    N_SAMPLE_CHUNKS = T.ceildiv(vocab_size, SAMPLE_BD)
18
    NUM_STAGES = 2 if vocab_size > LSE_BD else 0
19

20
    RADIX_BITS = 8 if LSE_BD <= 1024 else 11

1
关于 DRAFT_PROB_PENALTY:
2
设 p 为 draft 分布，q 为 target 分布。先从 p 采样；若采到 token j 且 p[j] <= q[j]，则直接接受；
3
否则从归一化后的 reject 权重中采样：r[i] = max(q[i] - alpha * p[i], 0),  alpha = exp(-DRAFT_PROB_PENALTY).
4
...
5
所以对任意 token j，有
6
    q[j] - eps / (1 + eps) <= A[j] <= (1 + eps) * q[j],
7
其中 eps = 1 - exp(-DRAFT_PROB_PENALTY)。

工程细节：从源码里能读出的约束

最后把这些散落在代码里的工程约束归拢一下——它们是”写这个 kernel”和”用这个 kernel”之间真正的接口。

内存层次必须显式声明，布局由编译器推导而不是手写。 TileLang 要求显式 alloc_shared / alloc_fragment / alloc_local / alloc_var，而不是让编译器猜。官方文档把这一步讲得很细：T.alloc_shared 对应 GPU 的 shared memory（片上、快、同一线程块内共享），T.alloc_fragment 对应寄存器文件（保存累加器与部分和，进一步降低延迟）；而”每个线程实际拿多少寄存器”由编译期的一个 Layout Inference Pass 推导出来的 T.Fragment 布局对象决定。

（图源：TileLang 官方文档 The Tile Language: A Brief Introduction）

这就是”写 TileLang 而不是写 CUDA”的取舍所在：布局推理让同一份 tile 级描述能适配不同线程组织与不同硬件，代价是布局不再由你亲手指定——所以 DeepInfer 那些对访存形态敏感的地方（块内连续、块间不连续的分页 KV）选择用线程原语 + 段表把访存形态写成确定的，而不是交给推导。TileLang 把这两条路放在同一门语言里，正是它被选中的原因。

具体到 Ascend 版，“数据必须先搬进 UB”是硬约束（对应上面的 T.alloc_shared 在这个后端落到 Unified Buffer）；CUDA 版里显式内存层次则是性能决策——把不规则索引先算进寄存器，再从连续区间做向量化访存，是这类算子的标准手法。

T.assume 把不变量告诉编译器。 整个 cuda/ 目录下有 87 处 T.assume。例如 load_paged_kv_cache 里：

1
T.assume(-1 <= src_block_id < num_blocks)
2
T.assume(0 <= src_begin <= src_end)
3
T.assume(src_end <= block_size)
4
T.assume(value_block_stride % 32 == 0)
5
T.assume(value_token_stride % value_token_alignment == 0)

这些断言在正确性上是冗余的——调用方本来就该保证。它们的作用是让编译器知道边界，从而省掉它自己加的保守检查（比如对每个访存做越界防护），同时也能据此做向量化。所以这不是防御性编程，是性能契约：写错了不是”安全地慢一点”，而是未定义行为。这解释了为什么 T.assume 和 T.ASSERT 是两种不同的东西，本文代码里用的是前者。

StridedTensor 把物理布局写进类型。 cuda/ 目录下出现 56 次 StridedTensor。以 KV 缓存参数为例：

1
kv_cache_value: T.StridedTensor((num_blocks, block_size, value_dim), value_strides, value_dtype)
2
value_strides = (value_block_stride, value_token_stride, 1)

value_stride 是运行时的动态变量（value_token_stride = T.dynamic("stride_kvcache_value")），不是编译期常量。也就是说逻辑形状是编译期定的、物理步长是运行时传的。这个区分很实际：逻辑形状（块数、块大小、维度）决定了 thread 组织与向量化策略，必须编译期定；而物理步长可能因为内存对齐要求、页表实现、图捕获时的缓冲复用而变化，编译期定不下来。把”形状”和”步长”分开处理，是这套 kernel 能复用在不同内存布局上的前提。

整数溢出与 dtype。 索引算术用 int32 还是 int64 是个真实决策：int32 访存更快、寄存器压力更小，但总 token 位置在长上下文下会超过 $2^{31}$。DeepInfer 的 KV 池采用块索引而非全局线性索引，一个直接好处就是索引值域被压回 int32 安全区——num_blocks 而不是 total_k 才是索引的上界。

PDL 是整层的启动策略，不是局部技巧。 28 个后端模块用到 T.pdl_sync()。它的含义是：本 kernel 要读的元数据可能由前一个 kernel 刚写出，与其让 GPU 等整个前序 kernel 结束（end-of-kernel 同步），不如让本 kernel 提前启动、在真正需要数据时才同步。这一层小 kernel 密集，正好是 PDL 收益最大的场景——它们彼此依赖链短、单个耗时小，启动延迟的占比很高。

Ascend 侧的文件不对称说明硬件差异不是被抽象掉的。 asc/kv_quant.py、asc/rstd_1ulp.py 与 CUDA 侧的 outlier_detection.py、quant.py 各自存在，说明某些数值处理必须跟着硬件走。“1ulp”（一个最小精度单位）这类命名透露的是归约精度策略——不同硬件在浮点归约路径上不同，需要不同的误差控制手段，硬套同一份实现反而会损失性能或精度。

测试结构说明契约强度。 顶层每个算子模块都带着 _ref 参考实现与参数化 pytest；conftest.py 把 KernelUnavailableError 转成 skip，并配有 benchmark 插件（TK_BENCHMARK_BASELINES_DIR、TK_BENCHMARK_REGRESSION_THRESHOLD=0.15、TK_BENCHMARK_MIN_DELTA_US=3.0 ——回归阈值 15% 且小于 3 微秒的差异不计，这两个数说明”多小的抖动可以忽略”是被明确定义过的）。这意味着测试用例是按”算法”写的、按后端跑的：同一个语义测试在两个后端都跑，谁都跑不了就 skip。

这一层在整个引擎里的位置

把本篇与前六篇接起来看，DeepInfer 在这一层的分层大致是这样：

1
调用层   模型 layer / worker / sampler   ——  只认 ops.xxx_tl
2
   ↓
3
分发层   ops.py + platform.py           ——  唯二知道后端的地方
4
   ↓
5
语义层   顶层模块：封装 + _ref 参考实现 + pytest 用例  ——  与硬件无关
6
   ↓
7
实现层   cuda/  |  asc/                 ——  按硬件各写，共享接口与算法结构
8
   ↓
9
共享层   common.py（量化常量）/ cache_common.py（KVCacheFormat）  ——  公式与常量
10
   ↓
11
底层库   FlashMLA / DeepGEMM / DeepEP / TileKernels / DeepSelect / TileLang

这里有个值得记住的判断：DeepInfer 没有试图用一套代码适配所有硬件。它做的是”把平台知识收敛到极小面积”（一个分发文件 + 一个平台类），再把”语义”从”实现”里剥出来放到顶层，最后用参考实现与测试约束两端行为一致。这是工程可维护性与”一份代码到处跑”的理想之间的一次明确取舍。

底层库里还有两个本篇没展开的：TileKernels 是 DeepSeek 用 TileLang 写的一批通用 kernel（Ascend 版 kernel 里直接 import tile_kernels 使用它的 get_num_vec_cores() 与 quant 模块），DeepSelect 从命名看与选择/路由相关。它们与 tilelang_ops 的分工（哪些算子放引擎内、哪些放外部库）是后续可以单独展开的话题。

小结

• tilelang_ops 存在的理由是主干大 kernel 之间的小算子缝隙：形状不规则、索引算术为主、单次开销小但数量多，用 PyTorch 拼会引入大量 launch 与同步开销。官方 README 把它列为 “custom TileLang kernels”，与 FlashMLA/DeepGEMM/DeepEP/TileKernels/DeepSelect 并列。
• ops.py 用 38 行完成分发：后端目录名来自 PLATFORM，模块靠 pkgutil 自动发现，接口契约是 _tl 后缀命名，导入期断言拦重名，KernelUnavailableError 让后端缺失变成测试 skip。平台知识被收敛到 ops.py + platform.py + 两个后端目录。
• 模块分六族：输入准备、KV 索引与元数据、KV 读写与量化、MoE 与路由辅助、采样与输出、投机与状态。带 _zos 的后缀对应零开销调度，带 fused_ 的对应融合优化。
• 顶层与后端是”语义 vs 实现”的分层：顶层模块放封装函数、_ref 参考实现与 pytest 用例，后端目录放具体 kernel。语义定义一次、验证一次，两端 kernel 都要对得上。
• 双平台是”接口停在算法级、实现跟着硬件走”：同一个 gather_input_ids_zos，CUDA 版一个 block 一条序列直接读写全局内存，Ascend 版按核持久化、先拷贝进 UB、用 num_stages=4 做多缓冲——因为 Ascend 的 “数据必须先搬进 UB” 是硬约束。两后端目录只差三个文件（quant.py/outlier_detection.py/prepare_decode_metadata.py 对 kv_quant.py/rstd_1ulp.py/prepare_v4_decode_metadata.py）。
• KV 读写是”一个被参数化的 kernel”：KVCacheFormat 描述一个 token 的逻辑与物理内容，工厂在构建期就把线程数、tile、向量宽度算成常量，kernel 内没有运行时布局分支。寻址靠 (src_block_id, src_begin, src_end, dst_offset) 四元组的段表，-1 编码无效块，段内向量化。
• 真难点是索引算术不是浮点算术：分页寻址的跨块拆分与对齐、变长验证的全局排序装箱、seqlens_k 的回滚修正、gather 的元数据 GPU 侧自产。
• 工程契约集中在三处：T.assume（把不变量告诉编译器，省掉保守检查）、StridedTensor（形状编译期定、步长运行时定）、T.pdl_sync（28 个模块共享的提前启动策略）。

下一篇回到模型侧：DeepInfer 的注意力层与 MLA 在引擎里的组织方式，以及 V4.1-Flash 的模型模块（model_executor/models/v4_modules/）如何把 FlashMLA、YoCo 与 DSpark 三个机制拼成一次前向。更早的机制背景可以回看 DeepInfer 源码完全拆解（一）：整体架构与模块全景，KV 侧的两本块账见 （三）：KV 内存与 SWA 有界重放。

参考资料

1. DeepInfer 官方仓库（README 功能清单与支持硬件）
2. DeepInfer tilelang_ops 目录（本文拆解的 kernel 集合）
3. DeepInfer ops.py（后端分发器源码）
4. DeepInfer platform.py（后端目录名与平台能力定义）
5. DeepInfer cache_common.py（KVCacheFormat 与 KV 布局定义）
6. DeepInfer common.py（量化常量与 DRAFT_PROB_PENALTY 推导）
7. DeepInfer gather_input_ids_zos.py（ZOS 输入 gather kernel，CUDA 版）
8. DeepInfer compute_seqlens_k_zos.py（seqlens_k 修正 kernel）
9. DeepInfer load_paged_kv_cache.py（分页 KV 读取与反量化 kernel）
10. DeepInfer store_kv_cache.py（分页 KV 写入与量化 kernel）
11. DeepInfer schedule_mtp_varlen.py（变长验证调度：语义、参考实现与测试）
12. DeepInfer fused_sample.py（采样融合 kernel，CUDA 版）
13. TileLang 官方仓库
14. TileLang 官方文档：The Tile Language: A Brief Introduction
15. TileLang 官方文档：编程指南总览
16. TileKernels 官方仓库（DeepSeek 的 TileLang 通用 kernel 集）
17. DeepSelect 官方仓库
18. FlashMLA 官方仓库（MLA 注意力 kernel）
19. DeepGEMM 官方仓库（FP8 矩阵乘内核）
20. DeepEP 官方仓库（MoE dispatch/combine）
21. DeepSeek-V4.1-Flash 模型卡