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： ```python @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](https://github.com/deepseek-ai/TileKernels)（MIT） - **TileLang**: [github.com/tile-ai/tilelang](https://github.com/tile-ai/tilelang) - **TileLang 论文**: [arXiv:2504.17577](https://arxiv.org/abs/2504.17577) - **DeepGEMM**: [github.com/deepseek-ai/DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) - **Engram 论文**: [arXiv:2601.07372](https://arxiv.org/abs/2601.07372) - **mHC 论文**: [arXiv:2512.24880](https://arxiv.org/abs/2512.24880) - **作者**: Xiangwen Wang, Chenhao Xu, Huanqi Cao, Rui Tian, Weilin Zhao, Kuai Yu, Chenggang Zhao