TileLang + TileKernels：DeepSeek 的 GPU 内核开发新范式，70 行 Python 替代 3000 行 CUDA

> DeepSeek 开源了 TileKernels——一个完全用 Python TileLang DSL 编写的 GPU 内核库，覆盖 MoE 路由、FP8/FP4 量化、Engram 条件记忆、流形超连接（mHC）等核心操作。所有内核都逼近硬件极限性能。这背后是一个更大的野心：让 Python 成为 GPU 编程的一等公民。 ## 两个仓库的关系 | | TileLang | TileKernels | |---|---------|-------------| | **是什么** | GPU/CPU 内核 DSL + 编译器基础设施 | 用 TileLang 写的高性能 LLM 内核库 | | **Star** | 5,710 | 1,078 | | **作者** | 北大 + 微软研究院 | DeepSeek | | **底层** | 基于 TVM 编译器 | 基于 TileLang | | **类比** | "GPU 编程的 Rust 编译器" | "用 Rust 写的 Web 框架" | TileLang 是**语言和编译器**，TileKernels 是**用这个语言写的应用**。理解 TileKernels 的关键，是先理解 TileLang 的设计哲学。 ## TileLang：Pythonic 的 GPU 内核编程 TileLang 的核心主张很简单：**用 Python 写 GPU 内核，性能不打折。** 它不是一个"Python 绑定"（像 PyTorch 那样调用 C++ 后端），而是一个**领域特定语言（DSL）**——你在 Python 中写的是 GPU 内核的逻辑，TileLang 编译器把它翻译成高效的 PTX/CUDA 代码。 ### 一个 GEMM 的例子 一个典型的 TileLang GEMM 内核长这样： ```python import tilelang from tilelang import language as T @tilelang.jit def matmul(M, N, K, block_M, block_N, block_K): @T.prim_func def main( A: T.Tensor[(M, K), "float16"], B: T.Tensor[(K, N), "float16"], C: T.Tensor[(M, N), "float16"], ): with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by): A_shared = T.alloc_shared((block_M, block_K), "float16") B_shared = T.alloc_shared((block_K, block_N), "float16") C_local = T.alloc_fragment((block_M, block_N), "float16") for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): T.copy(A[bx * block_M, k * block_K], A_shared) T.copy(B[k * block_K, by * block_N], B_shared) T.gemm(A_shared, B_shared, C_local) T.copy(C_local, C[bx * block_M, by * block_N]) return main ``` 关键设计元素： - **`@tilelang.jit`**：JIT 编译装饰器，首次调用时编译内核 - **`T.prim_func`**：声明式内核定义，描述计算逻辑而非执行步骤 - **`T.Kernel`**：block/grid 维度映射 - **`T.alloc_shared` / `T.alloc_fragment`**：显式内存层次控制（shared memory vs. register） - **`T.Pipelined`**：软件流水线，自动插入异步拷贝和屏障同步 - **`T.gemm`**：张量核心指令映射，自动选择 MMA/WMMA/wgmma ### 编译器架构 TileLang 的编译器栈基于 TVM 构建，但做了大量定制： ``` Python DSL → TileLang IR → TVM IR → 优化 Passes → PTX/CUDA/HIP ``` 关键的编译 Pass 包括： - **Layout Inference**：自动推断数据布局（行优先/列优先/banked swizzle） - **Software Pipeline**：自动插入 TMA 异步拷贝 + mbarrier 同步 - **Warp Specialization**：自动将 producer/consumer 操作分配到不同 warp - **Auto-tuning**：基于实测性能自动搜索最优 tile 大小和流水线参数 - **Multi-backend**：支持 CUDA (sm70-sm100)、HIP (AMD)、Metal (Apple)、CuTeDSL ### 支持的硬件特性 TileLang 覆盖了从 Volta 到 Blackwell 的全部 NVIDIA 架构特性： | 架构 | Compute Capability | 关键特性 | |------|-------------------|---------| | Volta | sm_70 | Tensor Core v1 (MMA 884) | | Ampere | sm_80 | MMA 16816, TMA, Async Copy | | Hopper | sm_90 | wgmma, TMA, Cluster | | Blackwell | sm_100 | TCGen05, FP4, 2xAcc FP8 | 还支持 AMD CDNA (MFMA) 和 Apple Metal。 ## TileKernels：DeepSeek V4 的内核武器库 TileKernels 是 DeepSeek 用 TileLang 写的 GPU 内核库，MIT 协议开源。它暴露了 DeepSeek 下一代模型（疑似 V4）的核心计算组件。 ### 六大模块 **1. Gating（门控选择）** - Top-k 专家选择：给定 `[num_tokens, num_experts]` 的门控分数，选出每个 token 的 top-k 专家 - 实现方式：每个 token 一个 warp（32 线程），对齐到 32 的倍数后并行扫描 - 关键细节：平局时返回更小的索引（确定性保证） **2. MoE Routing（MoE 路由）** - Token-to-expert 映射 - 融合的 expand/reduce 操作 - 权重归一化 - 这是 DeepSeek MoE 架构的核心——每个 token 只激活部分专家，路由效率直接决定训练速度 **3. Quantization（量化）** 这是最丰富的模块，覆盖了多种量化格式： | 格式 | 精度 | 粒度 | 用途 | |------|------|------|------| | FP8 (E4M3) | 8-bit | per-token / per-block / per-channel | 权重/激活量化 | | FP4 (E2M1) | 4-bit | per-token / per-block | 极致压缩 | | E5M6 | 8-bit | per-token | 缩放因子存储 | | UE8M0 | 8-bit | packed | 缩放因子压缩 | 关键设计： - **融合 SwiGLU + 量化**：`swiglu_forward_and_per_token_cast_kernel.py` 把激活函数和量化合并为一个内核，减少一次全局内存读写 - **per-block vs per-token**：per-block 量化把 hidden_size 分成多个 block，每个 block 独立计算缩放因子，精度更高但计算量更大 - **TMA 对齐的列优先缩放因子**：利用 Hopper 的 TMA 硬件加速缩放因子的加载 **4. Engram（条件记忆门控）** 这是 DeepSeek V4 架构中最创新的部分之一。Engram Gate 的计算公式： ``` gate = sigmoid(signed_sqrt(dot(RMSNorm(x, wh), RMSNorm(k, we)) * scalar)) output = hidden_states + gate * v ``` 其中 `signed_sqrt(x) = sign(x) * sqrt(|x|)`，`scalar = 1/sqrt(hidden_size)`。 内核实现的关键优化： - **融合 RMSNorm + 点积 + 门控 + 残差连接**：全部在一个内核中完成 - **前向保存中间结果**：dot、gate_score、rstd_x、rstd_k 保存用于反向传播 - **反向梯度归约**：权重梯度按 SM 数量分片归约，最后一步原子加法合并 - **性能调优**：只针对 hidden_size ∈ {4096, 7168} 调优（DeepSeek 的标准维度） **5. Manifold HyperConnection (mHC，流形超连接)** 这是 DeepSeek V4 的另一个核心创新——替代传统 Transformer 的残差连接。 mHC 的前向计算流程： 1. **RMSNorm**：对残差做归一化 2. **Split Mixes**：将归一化结果分成 pre_mix、post_mix、comb_mix 三部分 3. **Sinkhorn 归一化**：对 comb_mix 做双向归一化（行归一化 → 列归一化 → 重复 N 次） 4. **Apply Mix**：用 pre_mix 对残差做线性变换，得到子层输入 Sinkhorn 归一化的内核实现： - 前向：交替做行 softmax + eps 和列 softmax + eps，重复 `repeat` 次（默认 10） - 反向：需要保存所有中间 softmax 结果（`repeat * 2` 个），反向时逐层回传梯度 - 内存优化：token_block_size 可配置，平衡 shared memory 使用和并行度 **6. Modeling（高层建模）** 将底层内核封装为 `torch.autograd.Function`，可以直接嵌入 PyTorch 训练循环： - `EngramGateFn`：Engram 门控的完整前向+反向 - `mhc_pre` / `mhc_post` / `mhc_head`：mHC 的完整流水线 - 支持 `main_grad`（ZeRO 优化器的 fp32 梯度缓冲区） ### 代码风格分析 TileKernels 的代码有一个显著特点：**每个内核文件都很短**。 以 `topk_gate_kernel.py` 为例，核心内核定义只有约 40 行 Python。如果用纯 CUDA 写同样的功能，大约需要 300-500 行（包含 warp-level 原语、shared memory 管理、边界检查等）。TileLang 的编译器自动处理了这些底层细节。 但代码中也有"不 Pythonic"的部分： - 大量 `T.Parallel`、`T.alloc_fragment`、`T.alloc_reducer` 等 TileLang 特有的 API - 需要手动管理 shared memory 大小和线程布局 - 性能调优参数（block_k、num_threads、vec_size）需要根据硬件特性手动选择 这说明 TileLang 的抽象层次是**刻意选择的**——它不是要完全隐藏 GPU 编程的复杂性，而是把"正确但繁琐"的部分自动化（内存管理、同步、指令映射），同时保留"需要人工判断"的部分（tile 大小、流水线深度、内存布局）。 ## 关键洞察 ### 1. DeepSeek V4 架构拼图 TileKernels 的模块结构清晰地揭示了 DeepSeek V4 的架构： | TileKernels 模块 | 对应架构组件 | |-----------------|-------------| | Gating + MoE Routing | MoE 层（多专家混合） | | Quantization (FP8/FP4) | 混合精度训练/推理 | | Engram | 条件记忆系统（动态选择性记忆） | | mHC | 新型残差连接（替代标准 residual） | | Transpose | 数据布局转换 | **Engram + mHC 的组合尤其值得关注。** Engram 是一个"条件门控"机制——根据输入动态决定从持久记忆中检索什么。mHC 是一种"流形超连接"——用 Sinkhorn 归一化的混合矩阵替代简单的残差加法。两者结合，意味着 DeepSeek V4 的每一层都有**动态的、数据依赖的信息流**，而不是固定的前馈路径。 ### 2. Python DSL 的工程价值 TileLang + TileKernels 的组合证明了一个重要观点：**GPU 内核开发的瓶颈不是"写代码"，而是"调参数"。** 用 CUDA 写一个 GEMM 内核，80% 的时间花在调试 shared memory bank conflict、warp divergence、pipeline bubble 上。TileLang 的编译器自动处理了这些，让开发者专注于**算法层面的优化**（tile 大小、融合策略、数据布局）。 DeepSeek 团队选择 TileLang 而不是 Triton，可能是因为： - TileLang 基于 TVM，有更成熟的编译器基础设施 - TileLang 的 layout inference 和 auto-tuning 更强大 - TileLang 支持更多硬件后端（AMD、Metal、CuTeDSL） - TileLang 的 warp specialization 和 software pipeline 更成熟 ### 3. 开源策略 DeepSeek 开源 TileKernels 的时机（2026-04-22，V4 发布前后）和内容（核心计算内核、MIT 协议）表明： - **内核本身不是核心壁垒**——架构设计才是 - 开源内核可以吸引社区贡献优化，降低自己的维护成本 - MIT 协议意味着任何人都可以自由使用和修改 ## 我的思考 TileKernels 让我看到了 GPU 编程的一个可能未来：**Python 成为 GPU 内核的"源语言"，编译器负责把高级描述翻译成硬件最优的机器码。** 这和编译器领域的历史趋势一致：从汇编 → C → C++ → Rust，每一代新语言都在提高抽象层次，同时保持（或逼近）底层性能。TileLang 代表了下一层抽象：**从"描述怎么做"到"描述做什么"**。 但 TileLang 也不是银弹。它的学习曲线不比 CUDA 低多少——你需要理解 shared memory、warp、tensor core、软件流水线等概念，只是不需要手写 PTX 指令。它更像是一个"带自动优化的 CUDA"，而不是"不需要懂 GPU 的魔法工具"。 对从业者来说，TileKernels 的最大价值可能不是可以直接拿来用（毕竟它是为 DeepSeek V4 定制的），而是它展示了**如何用高层 DSL 表达复杂的 GPU 内核**。如果你在开发自己的 LLM 推理引擎，TileKernels 的代码是极好的参考实现。 --- **代码仓库** - TileLang: [https://github.com/tile-ai/tilelang](https://github.com/tile-ai/tilelang) (5.7K ⭐) - TileKernels: [https://github.com/deepseek-ai/TileKernels](https://github.com/deepseek-ai/TileKernels) (1K ⭐, MIT) #TileLang #TileKernels #DeepSeek #GPU #CUDA #DSL #MoE #量化 #Engram #流形超连接 #系统编程