DeepSeek 开源了 TileKernels——一个完全用 Python TileLang DSL 编写的 GPU 内核库，覆盖 MoE 路由、FP8/FP4 量化、Engram 条件记忆、流形超连接（mHC）等核心操作。所有内核都逼近硬件极限性能。这背后是一个更大的野心：让 Python 成为 GPU 编程的一等公民。

两个仓库的关系

TileLang 是语言和编译器，TileKernels 是用这个语言写的应用。理解 TileKernels 的关键，是先理解 TileLang 的设计哲学。

TileLang：Pythonic 的 GPU 内核编程

TileLang 的核心主张很简单：用 Python 写 GPU 内核，性能不打折。

它不是一个"Python 绑定"（像 PyTorch 那样调用 C++ 后端），而是一个领域特定语言（DSL）——你在 Python 中写的是 GPU 内核的逻辑，TileLang 编译器把它翻译成高效的 PTX/CUDA 代码。

一个 GEMM 的例子

一个典型的 TileLang GEMM 内核长这样：

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 架构特性：

还支持 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（量化） 这是最丰富的模块，覆盖了多种量化格式：

关键设计：

• 融合 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 的架构：

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 (5.7K ⭐)
• TileKernels: https://github.com/deepseek-ai/TileKernels (1K ⭐, MIT)

#TileLang #TileKernels #DeepSeek #GPU #CUDA #DSL #MoE #量化 #Engram #流形超连接 #系统编程