DeepSeek 开源 TileKernels：用 Python 写出逼近硬件极限的 GPU 内核，Engram 和 mHC 暴露 V4 架构野心

> DeepSeek 开源了 TileKernels——完全用 Python（TileLang DSL）编写的高性能 GPU 内核库。Engram 和 mHC 模块的暴露，揭示了 DeepSeek V4 的架构野心。

一个反直觉的事实

写一个高性能 GPU 内核，你需要什么？

传统答案是：C++、CUDA、模板元编程、几百行甚至上千行的代码、数周的调试时间。

DeepSeek 的答案是：Python，70 行以内。

这不是玩笑。TileKernels 是 DeepSeek 最新开源的 GPU 内核库，所有内核都用 TileLang——一个 Python DSL（领域特定语言）编写。在 H800 上跑出了接近硬件极限的性能，部分内核已经用于 DeepSeek 内部的训练和推理。

更关键的是，这个库里藏着两个模块——Engram 和 mHC——它们直接暴露了 DeepSeek V4 的架构设计。

为什么放弃 CUDA？

要理解 TileKernels 的意义，先要理解传统 GPU 内核开发的痛点。

写一个高性能的矩阵乘法（GEMM）内核，用纯 CUDA + C++ 模板，代码量通常在 1000-3000 行。这些代码充满了：

手动的 shared memory 分配和 bank conflict 处理
warp 级别的同步原语
循环展开和寄存器分块
针对不同 GPU 架构的条件编译

改一个参数？重新编译，重新调优。换一个 GPU 架构？可能要重写一半。

TileLang 的核心思想是：把"瓦片"（Tile）作为唯一的编程抽象。

在 GPU 编程中，"瓦片"是把大矩阵切成小块在 shared memory 中处理的基本策略。几乎所有高性能内核——GEMM、FlashAttention、LayerNorm——本质上都是"分块-加载-计算-存储"的循环。

TileLang 把这个模式抽象成了 DSL：

@tilelang.jit
def my_kernel(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K):
    @T.prim_func
    def kernel(A, B, C):
        with T.Kernel(T.ceildiv(M, BLOCK_M), T.ceildiv(N, BLOCK_N)) as (by, bx):
            A_frag = T.alloc_fragment((BLOCK_M, BLOCK_K), T.float16)
            B_frag = T.alloc_fragment((BLOCK_K, BLOCK_N), T.float16)
            C_frag = T.alloc_fragment((BLOCK_M, BLOCK_N), T.float32)
            
            for k in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
                T.copy(A[by * BLOCK_M, k * BLOCK_K], A_frag)
                T.copy(B[k * BLOCK_K, bx * BLOCK_N], B_frag)
                T.gemm(A_frag, B_frag, C_frag)
            
            T.copy(C_frag, C[by * BLOCK_M, bx * BLOCK_N])
    return kernel

没有 shared memory 手动管理，没有 bank conflict，没有 warp 同步。TileLang 的编译器自动处理这些底层细节。

TileLang 论文（arXiv:2504.17577）报告：相比手写 CUDA，代码量减少 最高 85.5%，而性能持平甚至更好。

TileKernels 的六大模块

TileKernels 不是 TileLang 的 demo，而是 DeepSeek 真正在生产中使用的内核库。它包含六个模块：

1. Gating — MoE 路由门控

Mixture of Experts 的核心：Top-k 专家选择和评分。

包含 topk_gate（Top-k 选择）、top2_sum_gate（Top-2 求和门控）、aux_fi（负载均衡损失计算）等内核。这些内核直接服务于 DeepSeek 的 MoE 架构——每个 token 需要被路由到最合适的专家。

2. MoE Routing — 专家路由

Token 到专家的映射、fused expansion/reduction、权重归一化。

expand_to_fused 把 token 分配到对应专家，reduce_fused 把专家输出合并回来。关键优化：fused 操作——把多个小 kernel 合并成一个大 kernel，减少全局内存访问次数。

3. Quantization — 量化

这是最丰富的模块，支持 FP8、FP4、E5M6 三种低精度格式，以及 per-token、per-block、per-channel 三种粒度。

特别值得注意的是：

SwiGLU + 量化融合内核：把激活函数和量化合并成一次 kernel launch
Cast-back 内核：推理时把低精度权重转回高精度计算
E5M6 支持：一种 5-bit 指数 + 6-bit 尾数的格式，比 FP8 更激进

这些量化内核直接服务于 DeepSeek 在算力受限条件下的训练策略——用更低的精度换取更大的模型和更多的数据。

4. Transpose — 批量转置

看似简单的矩阵转置，在 GPU 上要做到高性能并不容易。TileKernels 的转置内核针对 batched 场景优化，支持不同的内存布局。

5. Engram — 条件记忆门控

这是最令人兴奋的模块。

Engram 是 DeepSeek 在 2026 年 1 月提出的条件记忆架构（arXiv:2601.07372），核心思想是给 Transformer 加一个"第二大脑"——一个基于 N-gram 哈希的查找表，实现 O(1) 的知识检索。

TileKernels 中的 Engram 模块包含：

engram_hash：N-gram 哈希内核。用 XOR 和乘法组合计算 token 序列的哈希值，映射到嵌入表索引。支持 2 层 N-gram、最多 3-gram、每层 8 个嵌入表。
engram_gate_fwd/bwd：Engram 门控的前向和反向传播。融合了 RMSNorm、signed-sqrt 激活、sigmoid 门控。
engram_fused_weight：融合权重预处理。
engram_grad_w_reduce：梯度归约。

门控的计算公式是：

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

这个设计的精妙之处在于：用门控机制决定"什么时候查记忆"，而不是"查什么记忆"。 记忆的检索是确定性的（哈希），但是否使用检索结果是学习出来的。

6. mHC — 流形超连接

这是另一个直接暴露 V4 架构的模块。

mHC（Manifold-Constrained Hyper-Connections）是 DeepSeek 在 2025 年 12 月提出的架构创新（arXiv:2512.24880），已被引用 33 次。

传统 Transformer 的残差连接是 output = x + fn(x)——一条直线。mHC 把它升级成了 多条超连接：

# 传统残差连接
residual = residual + fn(x)  # 1 条连接

# mHC：多头残差
residual = residual + mix * fn(x)  # mhc_mult 条连接（默认 4 条）

TileKernels 中的 mHC 模块包含完整的流水线：

expand_to_mhc：把标准嵌入扩展为多头格式 (batch, seq, hidden) → (batch, seq, mhc_mult, hidden)
mhc_pre_norm_fn：带 RMSNorm 的前置归一化
mhc_head_compute_mix：用 sigmoid 计算混合权重 mix = sigmoid(input * scale + base) + eps
mhc_pre_apply_mix：应用混合权重到残差
mhc_post：后处理归约
sinkhorn_normalize：Sinkhorn 归一化——这是 mHC 的核心算法

Sinkhorn 归一化 是什么？它是一种把矩阵变成"双随机矩阵"（行和列都归一化为 1）的迭代算法。在 mHC 中，它确保多头连接的混合权重形成一个有效的概率分布，防止某些头被过度激活或抑制。

代码中的实现非常直接：交替进行行归一化和列归一化，迭代 repeat 次（默认 3 次），加上 epsilon 防止除零。

与 DeepGEMM 的关系

DeepSeek 之前开源的 DeepGEMM 是一个 FP8 GEMM 内核库，专注于矩阵乘法这一单一操作，在 H800 上达到 1550 TFLOPS。

TileKernels 和 DeepGEMM 的关系是互补而非替代：

维度	DeepGEMM	TileKernels
聚焦	纯 GEMM（矩阵乘法）	GEMM 之外的所有操作
语言	C++ / CUDA	Python (TileLang)
覆盖	FP8 矩阵乘法	MoE 路由、量化、Engram、mHC、转置
定位	基础计算原语	上层算子

一个完整的 LLM 推理/训练流水线需要两者配合：DeepGEMM 负责密集矩阵乘法，TileKernels 负责其他所有操作。

工程师的实用指南

如果你想在自己的项目中使用 TileKernels：

1. 环境要求：NVIDIA Hopper 架构 GPU（H100/H800），TileLang，PyTorch 2. 安装：pip install tilelang，然后 clone TileKernels 3. 快速验证：pytest tests/ 运行正确性测试 4. 性能测试：pytest tests/ --run-benchmark 运行基准测试 5. 查看生成的 CUDA 代码：设置 TK_PRINT_KERNEL_SOURCE=1

注意事项：

Engram 门控内核目前只针对 hidden_size in {4096, 7168} 做了性能调优
mHC 目前只保证 mhc_mult=4 能正常工作
许可证是 MIT，可以自由使用和修改
代码注释明确说"不代表最佳实践"，团队在持续改进

DeepSeek V4 的架构拼图

TileKernels 的开源，让我们可以拼出 DeepSeek V4 架构的更完整图景：

已确认的 V4 特性（从 TileKernels 代码推断）：

1. Engram 条件记忆：N-gram 哈希查找表，O(1) 知识检索，2 层 N-gram，最多 3-gram 2. mHC 流形超连接：多头残差连接（4 头），Sinkhorn 归一化 3. FP8/FP4 混合精度：per-token、per-block、per-channel 三种量化粒度 4. SwiGLU + 量化融合：激活函数和量化在一次 kernel 中完成 5. MoE 架构：Top-k 路由 + fused expansion/reduction

这些特性组合在一起，描绘了一个比 V3 更激进的设计：用条件记忆分担推理负担，用多头连接增强信息流动，用极低精度压缩计算成本。

我的思考

TileKernels 最重要的意义不在于性能数字，而在于它代表的一种范式转移。

过去，GPU 内核开发是"手艺人"的工作——需要深厚的 CUDA 经验、对硬件架构的深刻理解、大量的试错。这导致了一个问题：只有少数顶尖团队能写出高性能内核。

TileLang 正在改变这一点。当内核可以用 70 行 Python 写出来，当编译器自动处理 shared memory 和 warp 同步，GPU 优化的门槛将大幅降低。

DeepSeek 选择开源这些内核，而不是藏着掖着，说明他们判断：在算力受限的条件下，工程能力的差距比算法创新更致命。 与其独占内核优化技术，不如让整个社区一起推动工具链的成熟。

这也解释了为什么 DeepSeek 在被芯片禁令限制的情况下，依然能训练出世界级的模型。极致的工程化能力——从 DeepGEMM 到 TileKernels 到 TileLang——才是打破封锁的真正武器。

---

项目信息

TileKernels: github.com/deepseek-ai/TileKernels（MIT）
TileLang: github.com/tile-ai/tilelang
TileLang 论文: arXiv:2504.17577
DeepGEMM: github.com/deepseek-ai/DeepGEMM
Engram 论文: arXiv:2601.07372
mHC 论文: arXiv:2512.24880
作者: Xiangwen Wang, Chenhao Xu, Huanqi Cao, Rui Tian, Weilin Zhao, Kuai Yu, Chenggang Zhao