DeepSeek 的 GPU 内核帝国：从 TileLang DSL 到 Engram 条件记忆，拆解下一代 LLM 的基础设施

DeepSeek 在 2026 年 4 月开源了 TileKernels——用 TileLang DSL 编写的 GPU 内核库，覆盖 MoE 路由、FP8/FP4 量化、Engram 条件记忆门控、流形超连接（mHC）等核心操作。这不是一个普通的内核库，而是 DeepSeek V4 架构的基础设施层。本文综合 3 篇论文 + 4 个代码仓库，从编译器设计到模型架构，完整拆解这套系统。

全景图：五个项目的关系

┌─────────────────────────────────────────────────────────┐
│                    DeepSeek V4 架构                       │
│                                                          │
│  ┌──────────┐  ┌──────────┐  ┌──────────┐               │
│  │   mHC    │  │  Engram  │  │   MoE    │   模型架构层    │
│  │ 残差连接  │  │ 条件记忆  │  │ 专家路由  │               │
│  └────┬─────┘  └────┬─────┘  └────┬─────┘               │
│       │              │              │                     │
│  ┌────┴──────────────┴──────────────┴─────┐              │
│  │           TileKernels (Python/TileLang) │   内核实现层  │
│  │  MoE路由 · FP8/FP4量化 · Engram门控      │              │
│  │  mHC Sinkhorn · Transpose               │              │
│  └──────────────────┬──────────────────────┘              │
│                     │                                     │
│  ┌──────────────────┴──────────────────────┐              │
│  │           TileLang (Python DSL)          │   编译器层   │
│  │  Tiled编程模型 · Layout推断 · Pipeline    │              │
│  │  Tensor Core映射 · 多后端(NVIDIA/AMD)    │              │
│  └──────────────────┬──────────────────────┘              │
│                     │                                     │
│  ┌──────────────────┴──────────────────────┐              │
│  │           DeepGEMM (CUDA C++)           │   底层内核层  │
│  │  GEMM · Mega MoE · HC PreNorm · FP8/FP4 │              │
│  │  JIT编译 · Symmetric Buffer              │              │
│  └──────────────────────────────────────────┘              │
└─────────────────────────────────────────────────────────┘

关键洞察：DeepSeek 同时维护了两套内核栈——TileLang（高层 Python DSL）和 DeepGEMM（底层 CUDA C++）。这不是冗余，而是分层策略：TileLang 用于快速迭代新架构（Engram、mHC），DeepGEMM 用于极致性能的关键路径（GEMM、Mega MoE）。

第一层：TileLang——让 Python 成为 GPU 编程的一等公民

论文核心：TileLang: A Composable Tiled Programming Model for AI Systems

作者：Lei Wang, Yu Cheng 等（北大 + 微软研究院） 发表：arXiv:2504.17577, 2025 年 4 月

核心思想：Tile 是一等公民

TileLang 的核心主张可以用一句话概括：所有高性能 GPU 内核都遵循相同的数据流模式——在 DRAM 和 SRAM 之间搬运 tile，在 tile 上执行计算。 既然模式相同，为什么不把这个模式变成语言原语？

TileLang 是一个 Python DSL，底层基于 TVM 编译器。它不是"Python 绑定"（像 PyTorch 那样调用 C++ 后端），而是在 Python 中描述 GPU 内核逻辑，编译器负责翻译成 PTX/CUDA。

五维调度空间

TileLang 的关键创新是将 GPU 内核优化拆解为五个独立的调度维度，用户只需关注数据流，编译器处理其余：

一个 GEMM 的完整生命周期

import tilelang.language as T

@tilelang.jit
def matmul(M, N, K, block_M, block_N, block_K):
    @T.prim_func
    def main(A, B, C):
        with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by):
            # 1. 显式声明内存层级
            A_shared = T.alloc_shared((block_M, block_K), "float16")  # → Shared Memory
            B_shared = T.alloc_shared((block_K, block_N), "float16")  # → Shared Memory
            C_local  = T.alloc_fragment((block_M, block_N), "float32") # → Register Files

            T.clear(C_local)

            # 2. 软件流水线：3 阶段 overlap 数据搬运和计算
            for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
                T.copy(A[by * block_M, k * block_K], A_shared)
                T.copy(B[k * block_K, bx * block_N], B_shared)
                T.gemm(A_shared, B_shared, C_local)  # → 自动映射到 Tensor Core

            T.copy(C_local, C[by * block_M, bx * block_N])
    return main

15 行 Python，编译器自动处理：

• Shared memory bank conflict elimination（通过 Layout Swizzling）
• Thread-to-tile binding（通过 Layout Inference）
• Tensor Core MMA 指令映射（通过 Tensorization）
• Multi-stage software pipeline（通过 T.Pipelined）
• Warp-level memory access coalescing

等价的 CUTLASS C++ 代码需要 300-500 行。

Layout Inference：TileLang 的杀手锏

TileLang 论文中最有技术深度的贡献是 Layout Inference 机制。这是 Triton 做不到的。

问题：在 GPU 编程中，数据在 shared memory 中的布局直接影响 bank conflict。传统的做法是手写 swizzle pattern，比如 A_shared[i, k] 实际映射到 A_shared[i * 32 + ((k + (i % 8) / 2) ^ (i / 8 % 4)) * 8 + k % 8]。这种代码几乎不可读。

TileLang 的解法：用户写 A_shared = T.alloc_shared((128, 32))，编译器自动推断出 bank conflict free 的布局。具体流程：

1. 用户层：声明 buffer 和操作（T.copy, T.gemm）
2. Layout Inference Pass：分析数据流，推断每个 buffer 的最优布局
3. Thread Binding Inference：根据布局推断线程映射
4. Code Generation：生成带 swizzle 的 CUDA 代码

论文中展示了这个过程的可视化（Fig. 8）：从简单的 A[tid // 4, tid % 4 * 8 + v % 8] 到带 swizzle 的 A[tid // 4, (((tid % 4 * 8 + v % 8) // 8) ^ ((tid // 4) % 8 // 2)) * 8 + (tid % 4 * 8 + v % 8) % 8]。

性能数据

TileLang 论文提供了详尽的 benchmark：

• GEMM：在 A100 上匹配 cuBLAS，在 H100 上达到 cuBLAS 的 97-99%
• FlashAttention：在 A100 上达到 FlashAttention-2 的 99.3%，在 H100 上达到 98.7%
• Dequant GEMM：相比 cuBLAS-WFP16AFP16，最高 7.65× 加速（WINT2AINT8 配置）
• Mixed Precision GEMV：相比 Marlin 平均 1.04×，相比 BitsAndBytes 平均 1.62×

论文特别指出：Triton 在 mixed precision 场景下难以实现 TileLang 的性能，因为 Triton 缺乏对 tile 内部数据布局的细粒度控制。

多后端支持

TileLang 不仅支持 NVIDIA GPU，还支持：

• AMD GPU（HIP/ROCm）
• Metal（Apple Silicon）
• CuTeDSL（NVIDIA 下一代线性代数库）
• CPU（通用处理器）

第二层：DeepGEMM——极致性能的 CUDA 内核库

DeepGEMM: clean and efficient BLAS kernel library on GPU

作者：Chenggang Zhao, Zhean Xu 等（DeepSeek） Star：~10K+（2026 年 4 月） 协议：MIT

定位：DeepSeek 的"手写 CUDA"层

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

• DeepGEMM：手写 CUDA C++，JIT 编译，极致性能。用于 GEMM、Mega MoE 等计算密集型操作
• TileLang：Python DSL，快速开发。用于 Engram、mHC 等需要频繁迭代的架构创新

DeepGEMM 的 README 明确说："leverages some concepts from CUTLASS and CuTe, but avoids heavy reliance on their templates or algebras. The library is designed for simplicity, with only a limited number of core kernel functions."

Mega MoE：通信-计算重叠的极致

DeepGEMM 2026 年 4 月新增的 Mega MoE 是最值得关注的功能：

def fp8_fp4_mega_moe(y, l1_weights, l2_weights, sym_buffer, ...):
    _C.fp8_fp4_mega_moe(y, l1_weights, l2_weights, ...)

核心创新：使用 PyTorch 的 Symmetric Buffer（torch.distributed.symmetric_memory）实现跨 GPU 的零拷贝 MoE 计算。

传统 MoE 推理流程：

1. All-to-All 通信分发 token
2. 等待通信完成
3. 执行 GEMM
4. All-to-All 通信回收结果
5. 等待通信完成

Mega MoE 的流程：

1. 通信和计算完全重叠——token 在 GPU 间传输的同时，当前 GPU 已经在处理已到达的 token
2. 使用 symmetric buffer 避免 GPU 间的数据拷贝
3. FP8 dispatch + FP4 权重，最大化带宽利用率

DeepGEMM 的内核清单

从 __init__.py 可以看到 DeepGEMM 提供的完整内核列表：

与 TileKernels 的关系

DeepGEMM 的 third-party/tilelang_ops/ 目录包含 TileLang 编写的辅助内核（如 swiglu_apply_weight_to_fp8.py），说明DeepGEMM 在某些场景下也会调用 TileLang 内核。两套栈不是隔离的，而是互相嵌入。

第三层：TileKernels——DeepSeek V4 的内核实现

TileKernels: Optimized GPU kernels for LLM operations

作者：Xiangwen Wang, Chenhao Xu 等（DeepSeek） Star：1,078（3 天内） 协议：MIT 创建：2026-04-22

架构总览

tile_kernels/
├── moe/        # MoE 路由：TopK 门控、token-expert 映射、融合扩展/归约
├── quant/      # 量化：Per-token/Per-block/Per-channel FP8/FP4/E5M6
├── engram/     # Engram 条件记忆：门控前向/反向、权重融合、哈希
├── mhc/        # 流形超连接：Sinkhorn 归一化、混合/应用、重计算
├── transpose/  # 批量转置
├── modeling/   # 高层 autograd 封装（EngramGateFn, mHC pipeline）
├── torch/      # PyTorch 参考实现（用于正确性验证）
└── testing/    # 测试和 benchmark 工具

Engram 门控内核：融合的艺术

Engram 的前向传播公式（从论文和代码综合）：

TileKernels 的实现把这个多步计算融合成单个 GPU 内核：

# engram_gate_kernel.py 核心结构
@tilelang.jit
def get_engram_gate_fwd_kernel(hidden_size, eps, scalar, ...):
    threads = 32
    vec_size = 8
    # 单个 kernel 完成：
    # 1. RMSNorm(h) 和 RMSNorm(k)
    # 2. 点积 + signed_sqrt + sigmoid 门控
    # 3. 门控值 × value
    # 4. 残差加法
    # 5. 保存 backward 所需的中间值

关键优化：

• 参数融合：weight_hidden 和 weight_embed 的 RMSNorm 权重被预融合到一个连续 buffer 中
• 分支共享：4 个 mHC 分支共享同一个 Value 投影矩阵 \(\mathbf{W}_V\)，只有 Key 投影矩阵 \(\mathbf{W}_K^{(m)}\) 是分支特定的。这允许将线性投影融合为单个 FP8 矩阵乘法
• Block-wise 处理：hidden_size 被分成 1024/768/512/256 的 block，每个 thread 处理 8 个元素（vec_size=8），最大化内存带宽利用率

MoE TopK 门控：40 行 Python 替代 300 行 CUDA

@tilelang.jit
def get_topk_gate_kernel(num_experts, num_topk):
    num_threads = 32
    @T.prim_func
    def topk_gate_kernel(scores, topk_idx):
        with T.Kernel(num_tokens, threads=num_threads) as pid:
            scores_fragment = T.alloc_fragment((num_aligned_experts,), T.float32)
            amax_fragment = T.alloc_fragment((1,), T.float32)
            idx_fragment = T.alloc_fragment((num_aligned_experts,), T.int32)
            idx_reducer = T.alloc_reducer((1,), T.int32, 'min', replication='all')
            topk_idx_shared = T.alloc_shared((num_topk,), T.int32)

            # 加载 + padding
            for i in T.Parallel(num_aligned_experts):
                scores_fragment[i] = scores[pid, i] if i < num_experts else -INFINITY

            # num_topk 轮迭代选择
            for _ in range(num_topk):
                T.reduce_max(scores_fragment, amax_fragment)
                T.reduce_argmax(scores_fragment, idx_reducer)
                topk_idx_shared[_] = idx_reducer[0]
                scores_fragment[idx_reducer[0]] = -INFINITY  # 标记已选

            T.copy(topk_idx_shared, topk_idx[pid, 0])
    return topk_gate_kernel

这个内核使用 T.alloc_reducer 实现 warp-level 的 argmax 归约，避免了全局同步。

Sinkhorn 归一化内核：mHC 的数学核心

mHC 的核心约束是将残差映射 \(\mathcal{H}_l^{\text{res}}\) 投影到双随机矩阵流形（Birkhoff polytope）上。这通过 Sinkhorn-Knopp 算法实现：

TileKernels 的 Sinkhorn 内核在单个 GPU kernel 中完成 20 次迭代的行列归一化：

@tilelang.jit
def _mhc_sinkhorn_fwd(hidden_size, token_block_size, repeat, eps):
    @T.prim_func
    def mhc_sinkhorn_kernel(comb_res_mix, comb_res_mix_out):
        with T.Kernel(T.ceildiv(num_tokens, token_block_size)) as pid_x:
            comb_frag = T.alloc_fragment((token_block_size, hidden_size, hidden_size), T.float32)
            row_sum = T.alloc_fragment((token_block_size, hidden_size), T.float32)
            col_sum = T.alloc_fragment((token_block_size, hidden_size), T.float32)

            T.copy(comb_res_mix[pid_x * token_block_size, 0, 0], comb_frag)

            # softmax(-1) + eps
            row_max = T.alloc_fragment((token_block_size, hidden_size), T.float32)
            T.reduce_max(comb_frag, row_max, dim=2)
            # ... exp + row_sum + normalize

            # repeat 轮双向归一化
            for step in range(repeat):
                T.reduce_sum(comb_frag, col_sum, dim=1)
                # ... col normalize
                T.reduce_sum(comb_frag, row_sum, dim=2)
                # ... row normalize

            T.copy(comb_frag, comb_res_mix_out[pid_x * token_block_size, 0, 0])
    return mhc_sinkhorn_kernel

内存优化：前向需要保存所有中间结果用于反向传播。TileKernels 的反向内核在单个 kernel 中重新计算所有中间值，避免存储 \(O(T \times H \times H \times \text{repeat})\) 的中间激活。

第四层：Engram——条件记忆，LLM 的"新轴"

论文核心：Conditional Memory via Scalable Lookup

作者：Xin Cheng, Wangding Zeng 等（北大 + DeepSeek） 发表：arXiv:2601.07372, 2026 年 1 月

核心论点：LLM 缺少"知识查找"原语

Engram 论文提出了一个根本性的观察：当前 LLM 用"计算"来模拟"记忆查找"，这是低效的。

论文用 Table 3 的例子说明：为了识别实体 "Diana, Princess of Wales"，LLM 需要消耗 6 层 Attention + FFN 来逐步组合特征：

• Layer 1-2: "Country in the United Kingdom" → Wales
• Layer 3: "Country in Europe" → Wales
• Layer 4: "Title held by female sovereigns" → Princess of Wales (unspecific)
• Layer 5: "Wife of Prince Charles" → Princess of Wales (unspecific)
• Layer 6: 完整实体 → Diana, Princess of Wales

Engram 的主张：这种"通过深度计算重建静态知识"的过程，本质上是一个 \(O(1)\) 的查找操作。为什么不直接查表？

架构设计

Engram 的完整流程：

Phase 1: 稀疏检索（O(1) 查找）

1. Tokenizer 压缩：将 128K 词表通过 NFKC + lowercasing 投影到更小的规范词表（压缩率 23.43%）
2. Multi-Head Hashing：对每个 token 位置，提取后缀 N-gram（N=2,3），通过 K=8 个不同的哈希头映射到嵌入表
3. 确定性查找：\(z_{t,n,k} = \varphi_{n,k}(g_{t,n})\)，\(\mathbf{e}_{t,n,k} = \mathbf{E}_{n,k}[z_{t,n,k}]\)

Phase 2: 上下文感知门控

门控值 \(\alpha_t \in (0, 1)\) 的作用：如果检索到的记忆与当前上下文矛盾，门控趋向 0，抑制噪声。

Phase 3: 轻量卷积 + 残差

使用 kernel_size=4, dilation=max(N-gram order) 的 depthwise causal convolution。

U 形稀疏分配定律

Engram 论文最深刻的发现是 Sparsity Allocation 定律：

给定固定参数预算 \(P_{\text{tot}}\) 和固定激活参数 \(P_{\text{act}}\)，定义分配比 \(\rho \in [0, 1]\) 为分配给 MoE 的非激活参数比例：

实验结果：验证损失与 \(\rho\) 呈 U 形关系。

• \(\rho = 100\%\)（纯 MoE）：缺乏静态记忆，浪费深度重建固定模式
• \(\rho \to 0\%\)（纯 Engram）：缺乏条件计算能力，推理任务退化
• 最优 \(\rho \approx 75\%-80\%\)：将 20-25% 的稀疏参数预算分配给 Engram

在 10B 规模下，最优分配将验证损失从 1.7248 降到 1.7109（\(\Delta = 0.0139\)）。

"有效深度"增加：CKA 分析

Engram 论文用 CKA（Centered Kernel Alignment）分析证明了一个惊人的结论：Engram 在功能上等价于增加了模型深度。

具体发现：

• Engram-27B 第 5 层的表示与 MoE-27B 第 12 层的表示最相似
• LogitLens 分析显示 Engram 的预测收敛速度显著快于 MoE
• 结论：Engram 通过显式查找绕过了早期的特征组合，释放了网络深度用于复杂推理

系统效率：100B 参数表卸载到 CPU，开销 < 3%

Engram 的确定性寻址（基于 token 序列的哈希）使得预取成为可能：

• 训练时：嵌入表分片到多个 GPU，使用 All-to-All 通信
• 推理时：嵌入表卸载到 host memory，利用前面层的计算时间异步预取

实验结果：在 H800 上，100B 参数的 Engram 表完全驻留在 CPU 内存中，推理吞吐量仅下降 2.8%。

性能数据

最令人惊讶的发现：Engram 在推理任务（BBH +5.0, ARC +3.7）上的增益大于知识任务（MMLU +3.0）。这颠覆了"记忆模块只帮助知识检索"的直觉。

消融实验：关闭 Engram 后

这证实了 Engram 是参数知识的主要存储库，而阅读理解主要依赖 backbone 的注意力机制。

第五层：mHC——让超连接稳定可扩展

论文核心：mHC: Manifold-Constrained Hyper-Connections

作者：Zhenda Xie 等（DeepSeek） 发表：arXiv:2512.24880, 2025 年 12 月

问题：Hyper-Connections 的不稳定性

标准残差连接：\(\mathbf{x}_{l+1} = \mathbf{x}_l + \mathcal{F}(\mathbf{x}_l, \mathcal{W}_l)\)

Hyper-Connections（HC）将残差流扩展为 \(n\) 个并行分支：

其中 \(\mathbf{x}_l \in \mathbb{R}^{n \times C}\)，\(\mathcal{H}_l^{\text{res}} \in \mathbb{R}^{n \times n}\)。

问题：当扩展到多层时，复合映射 \(\prod_{i=1}^{L-l} \mathcal{H}_{L-i}^{\text{res}}\) 不保持恒等映射性质。实验中观察到：

• 27B 模型在 ~12K 步出现 loss surge
• 复合映射的 Amax Gain Magnitude 峰值达到 3000（理想值应为 1）

解法：投影到双随机矩阵流形

mHC 约束 \(\mathcal{H}_l^{\text{res}}\) 为双随机矩阵（行和 = 列和 = 1，所有元素 ≥ 0）：

通过 Sinkhorn-Knopp 算法实现投影：\(\mathcal{H}_l^{\text{res}} = \text{Sinkhorn-Knopp}(\tilde{\mathcal{H}}_l^{\text{res}})\)，其中 \(\tilde{\mathcal{H}}_l^{\text{res}}\) 是原始未约束的映射。

三个理论保证：

1. 范数保持：\(\|\mathcal{H}_l^{\text{res}}\|_2 \leq 1\)，防止梯度爆炸
2. 组合闭包：双随机矩阵的乘积仍然是双随机的，多层复合映射保持稳定
3. 几何解释：Birkhoff polytope 是置换矩阵的凸包，残差映射本质上是置换的凸组合

实际效果

基础设施优化

mHC 的 6.7% 额外开销是通过三项基础设施优化实现的：

1. Kernel Fusion：将 RMSNorm + 线性投影 + Sigmoid/Sinkhorn 融合为单个 kernel，使用 TileLang 实现。将内存读取从 \((3n+1)C\) 减少到 \((n+1)C\)
2. 选择性重计算：丢弃 mHC 中间激活，反向时重新计算。最优块大小 \(L_r^* \approx \sqrt{nL/(n+2)}\)
3. DualPipe 通信重叠：扩展 DualPipe 调度，在 pipeline stage 边界处重叠 mHC 的重计算和跨 stage 通信

综合分析：DeepSeek 的技术哲学

1. "条件计算 + 条件记忆"的双轴稀疏

DeepSeek V4 的架构可以用两个轴来理解：

U 形分配定律表明，两个轴的最优比例约为 75:25。这不是偶然的——它反映了语言信号的双重本质：组合推理（需要计算）和固定模式（需要查找）。

2. "高层 DSL + 底层 CUDA"的双栈策略

DeepSeek 同时维护 TileLang 和 DeepGEMM 两套内核栈，这不是资源浪费，而是工程成熟度的体现：

• 架构探索阶段（Engram、mHC）：用 TileLang 快速迭代，几天内完成内核原型
• 性能优化阶段（GEMM、Mega MoE）：用 DeepGEMM 手写 CUDA，榨干硬件性能
• 稳定阶段：TileLang 内核可能被移植到 DeepGEMM（如 TileKernels 中的内核）

3. "算法-系统协同设计"的方法论

Engram 的设计体现了 DeepSeek 的核心方法论——算法和系统同时设计：

• 算法层：Engram 的确定性寻址不是偶然的，而是为了支持系统层的预取
• 系统层：host memory offload 不是事后优化，而是架构设计的一等公民
• 结果：100B 参数表卸载到 CPU，开销 < 3%

同样，mHC 的 Sinkhorn 约束不仅是为了数学优雅，更是为了让复合映射保持有界，从而允许大规模训练。

4. 对 AI 基础设施的启示

这套系统对 AI 基础设施社区有几个重要启示：

a) DSL 是 GPU 编程的未来方向

TileLang 证明了 Python DSL 可以达到手写 CUDA 的性能，同时大幅降低开发门槛。Triton 在这方面开了头，但 TileLang 的 Layout Inference 机制走得更远。

b) "条件记忆"可能成为下一代 LLM 的标准组件

Engram 的 U 形分配定律和有效深度增加的 CKA 分析表明，条件记忆不是锦上添花，而是架构层面的必要补充。预计未来会有更多模型采用类似的设计。

c) 开源内核库是架构创新的催化剂

DeepSeek 开源 TileKernels（MIT 协议）的时机（V4 发布前后）和内容（核心计算内核）表明：内核本身不是壁垒，架构设计才是。开源可以吸引社区贡献优化，同时展示技术实力。

d) 系统优化是架构可扩展性的前提

mHC 如果没有 kernel fusion + selective recomputation + DualPipe overlapping，6.7% 的开销会变成 30-50%，使得 n=4 的扩展在实践中不可行。每一个架构创新都需要配套的系统优化才能落地。

参考资源

论文

• TileLang: arXiv:2504.17577
• Engram: arXiv:2601.07372
• mHC: arXiv:2512.24880

代码仓库

• TileLang: github.com/tile-ai/tilelang (5.7K ⭐)
• TileKernels: github.com/deepseek-ai/TileKernels (1K ⭐, MIT)
• DeepGEMM: github.com/deepseek-ai/DeepGEMM (~10K ⭐, MIT)

#DeepSeek #TileLang #TileKernels #DeepGEMM #Engram #mHC #GPU #CUDA #DSL #MoE #量化 #系统编程 #LLM架构