• 人工智能
• 大模型
• 模型推理服务
• 推理引擎
• DeepSeek

🚀DeepInfer是 DeepSeek 官方开源的高性能大规模推理框架，支撑其生产级大模型服务。它的极致性能不只靠调度器，更来自一套用TileLang编写的自定义算子内核（TileLang custom kernels）。本文将带你快速读懂 DeepInfer 中 TileLang 内核的架构、写法与验证方式，是面向新手的完整教程。

为什么性能优化要下沉到自定义内核？

大模型推理的吞吐瓶颈往往不在单个大矩阵乘，而在成百上千个"小而碎"的操作：RMSNorm、KV cache 读写、索引 gather、采样、量化转换……这些操作如果用通用库逐个调用，CPU 端调度开销和显存往返会蚕食掉大量时间。

DeepInfer 的答案是：用 TileLang 这类 DSL 把这些"碎操作"手写成本地化、融合化的 GPU/NPU 内核，直接嵌入推理热路径。DeepInfer 的优化内核组合包括 FlashMLA、DeepGEMM、DeepEP、TileKernels、DeepSelect，以及项目自研的 TileLang kernels。

上图是 DeepInfer 在 64 卡 SM100 集群上跑 DeepSeek V4.1 Flash 的 decode 阶段 Pareto 图：单用户速度最高约 320 tokens/s，批量吞吐可达 26,000+ tokens/s/GPU——这正是自定义内核堆出来的结果。

目录结构速览：三层分离架构

TileLang 内核全部集中在 deepinfer/model_executor/layers/tilelang_ops/，采用清晰的三层分离：

这种"一份接口、双硬件实现"的设计，让上层模型代码完全无感切换硬件平台。

内核分发器：ops.py 的巧思

ops.py 只做一件事：读取平台配置PLATFORM.tilelang_ops_backend（NVIDIA 为cuda，Ascend 为asc，见 platform.py），然后自动扫描后端包内所有以_tl结尾的内核工厂函数并重新导出。调用方只需：

from deepinfer.model_executor.layers.tilelang_ops import ops kernel = ops.gather_tl(D=4096, dtype=..., ...)

如果当前后端缺少某个内核，KernelUnavailableError会给出清晰的错误提示，测试框架还会自动将其转为 skip。

内核清单：DeepInfer 里有哪些 TileLang 算子？

cuda/目录下约 28 个内核文件（合计近 5000 行 DSL），覆盖了推理热路径的关键环节：

• 基础运算：norm.py（RMSNorm，可融合 residual 加法与 FP8 量化）、gather.py、quant.py、outlier_detection.py
• KV cache 管理：store_kv_cache.py、load_paged_kv_cache.py、get_kv_index.py
• 采样与投机解码：sampler.py、random_sampling.py、fused_sample.py、dflash.py（置信度头融合内核）
• 调度辅助：schedule_mtp_varlen.py、swa_index.py、candidate_blocks.py

以store_kv_cache为例，接口层会根据 KV cache 的实际 dtype 自动推断是 BF16 还是 FP4 量化布局，再调用对应的ops.store_kv_cache_tl(...)工厂——调用方无需关心底层细节。

上手写内核：gather 算子拆解

TileLang 内核的骨架由三个部分组成：@tilelang.jit装饰器、T.dynamic动态形状、@T.prim_func定义的函数体。下面以 cuda/gather.py 为例拆解：

@tilelang.jit(pass_configs={ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, }) def gather_tl(D: int, dtype, index_dtype, x_stride, threads=256): n = T.dynamic("n") # 行数运行时才知道，编译期留动态 m = T.dynamic("m") ... @T.prim_func def _gather(X, INDEX, Y): with T.Kernel(nbd * nbm, threads=threads) as (bx): index = T.alloc_fragment([buf_m], index_dtype) # 寄存器级缓冲 values = T.alloc_fragment([buf_m, buf_n], dtype) for m_i, d_i in T.Parallel(buf_m, buf_n): # 并行映射到线程 if index[m_i] >= 0 and index[m_i] < n and ...: values[m_i, d_i] = X[index[m_i], ...] else: values[m_i, d_i] = 0 T.copy(values, Y[...]) return _gather

新手可以从中提炼出四条通用技巧：

1. 动态形状：T.dynamic("n")把 batch 维度声明为运行时变量，一次编译适配所有长度；
2. 按问题规模选块：D <= 4用 C=1024，D > 128时放大到 C=16384 以最大化向量化（int4 加载）；
3. 寄存器暂存：T.alloc_fragment把中间结果放进寄存器，减少全局显存往返；
4. 边界安全：对越界索引显式写 0，保证 padding 语义正确。

而norm.py展示了融合的价值：一次内核同时完成 RMSNorm + residual 相加 + FP8 cast + scale factor 写出，比串起来调 3~4 个库函数少一半以上的显存读写。

双硬件移植：cuda 与 asc 怎么共存？

asc/目录与cuda/目录文件一一对应（如dflash.py、norm.py、sampler.py），各自针对 NVIDIA SM100 与 Ascend 950 NPU 的硬件特性调优。分发层只认xxx_tl后缀的符号，并在加载时做符号冲突断言（见 ops.py），确保两份后端不会悄悄覆盖彼此。

配合这些内核，DeepInfer 在 4 卡 SM100 上跑 prefill 时，非缓存输入吞吐可超过 90,000 tokens/s/GPU，延迟保持在数百毫秒级——这是调度优化与自定义内核协同的结果。

如何验证内核正确性与性能不回退？

DeepInfer 用 conftest.py 给整个内核目录装配了tile_kernels测试插件，形成完整的质量闭环：

• 正确性：每个内核文件内都带 pytest 用例，与 PyTorch 参考实现逐元素对比（如 dflash.py 中的confidence_head_fused_ref）；
• 性能基线：benchmark 结果保存在benchmark_baselines/，回归超过 15%（TK_BENCHMARK_REGRESSION_THRESHOLD）即报警，最小判定粒度 3 微秒；
• 显存画像：gpu_mem_profiles/跟踪内核显存占用，防止隐性泄漏；
• 分布式并行：xdist 自动为每个 worker 绑定独立 GPU，多卡并行跑测试。

想体验完整的性能压测流程，可以跑 examples/sweep/run_sweep.sh，它会自动扫描批量大小并产出上面的 Pareto 曲线（绘图脚本见 plot_pareto.py）。

总结：DeepInfer 性能方法论清单

✅算子融合——Norm、量化、KV 读写一步完成，消灭中间显存读写 ✅动态形状编译——一次编译，适配任意 batch 与序列长度 ✅平台无关接口——ops.xxx_tl一个入口，GPU/NPU 双后端透明切换 ✅基线守护——性能回归阈值 15%，杜绝"越优化越慢" ✅面向 DeepSeek 架构深度定制——DSpark 验证、SWA 有界重放、MTP 变长调度均有专属内核

对想深入推理引擎内核优化的新手，建议的阅读路线是：先读 ops.py 理解分发机制，再挑gather.py、norm.py两个内核精读 DSL 写法，最后对照conftest.py看懂质量门禁。掌握这套模式后，你就能像 DeepInfer 团队一样，把每一微秒都榨干 ⚡

• 人工智能
• 大模型
• 模型推理服务
• 推理引擎
• DeepSeek