TileKernels: A kernel library written in tilelang

来源：GitHub · ⭐ 1280 stars

项目简介

TileKernels 是由 DeepSeek AI 团队开源的一个高性能 GPU 算子库，基于 TileLang 框架编写。该项目瞄准了当前 AI 基础设施中的一个核心痛点：手写 CUDA 内核的开发效率低、优化门槛高。通过 TileLang 提供的声明式 tile 级编程抽象，TileKernels 实现了 CUDA 内核的高效生成与自动调度，让开发者可以用接近 Python 的简洁语法描述复杂的 GPU 计算逻辑，而无需深入掌握 CUDA 编程的每一个细节。

在当前大模型推理和训练对算力需求爆炸式增长的背景下，高性能算子库的重要性愈发凸显。TileKernels 的诞生，为 AI 基础设施提供了一套可复用、可扩展的高性能算子基础，特别是在矩阵运算、注意力机制等核心计算密集型场景中，能够显著释放 GPU 硬件潜力。

核心特性

声明式 Tile 编程：采用 TileLang 的 tile 级抽象，将复杂的 GPU 内存层次结构（shared memory、register、global memory）以高层次的 tile 概念表达，大幅降低 CUDA 开发的心智负担。
极致性能优化：内置多种自动调度策略，包括 tile 大小自动搜索、内存访问模式优化、warp 级并行协作等，生成的内核性能接近甚至超越手写 CUDA 的专家级实现。
广泛算子覆盖：涵盖矩阵乘法（GEMM）、卷积、LayerNorm、Softmax、FlashAttention 等深度学习核心算子，满足主流模型推理与训练的加速需求。
与主流框架无缝集成：提供 PyTorch 和 Triton 兼容的接口封装，可以零摩擦地接入现有训练与推理 pipeline，无需改动上层业务代码。
硬件自适应能力：支持 NVIDIA Ampere、Hopper 等多代 GPU 架构，自动针对硬件特性（如 Tensor Core、异步拷贝指令）进行代码生成与优化。

技术实现

TileKernels 的核心技术栈建立在 TileLang 之上。TileLang 是一种专为 GPU 内核开发设计的领域特定语言（DSL），其关键思想是将计算逻辑（what to compute）与调度策略（how to schedule）分离。

在架构层面，TileKernels 采用了分层设计：

算子定义层：用 TileLang 描述算子的数学逻辑，如矩阵乘法的累加计算过程。
调度策略层：通过编译时自动调优（auto-tuning）或预设模板，确定 tile 大小、线程布局、内存加载顺序等关键参数。
代码生成层：TileLang 编译器将高层次的 tile 描述 lower 为底层 CUDA C++，并自动注入内存对齐、bank conflict 避免、寄存器分配等底层优化。

这种"算子 + 调度"分离的设计，使得 TileKernels 既能保持算子实现的简洁可读，又能针对不同 GPU 架构和输入 shape 生成最优的内核代码。相比传统的基于模板的算子库（如 CUTLASS），TileKernels 在开发效率和可维护性上具有明显优势。

快速上手

使用 TileKernels 非常简单，以下是一个矩阵乘法的示例：

import torch
from tilekernels import gemm

# 创建随机输入矩阵
A = torch.randn(1024, 1024, device='cuda')
B = torch.randn(1024, 1024, device='cuda')

# 调用 TileKernels 优化的 GEMM 算子
C = gemm(A, B)

# 与 PyTorch 原生实现对比
C_ref = A @ B
assert torch.allclose(C, C_ref, atol=1e-4)
print("TileKernels GEMM 结果验证通过！")

对于需要更精细控制的场景，TileKernels 也允许开发者直接调用底层的 tile 级 API，自定义 tile 大小和调度参数：

from tilekernels import TileConfig, gemm_custom

config = TileConfig(
    block_tile=(128, 128, 32),
    warp_tile=(64, 64, 16),
    use_tensorcore=True
)
C = gemm_custom(A, B, config=config)

应用场景

大模型推理加速：在 LLM 推理服务中，矩阵乘法和注意力计算是主要瓶颈。TileKernels 提供的优化算子可以显著降低推理延迟，提升吞吐量，适用于在线对话、代码补全等实时场景。
科学计算与数值模拟：对于依赖大规模矩阵运算的物理仿真、金融建模等任务，TileKernels 能够充分利用 GPU 算力，缩短计算周期。
边缘设备优化：在资源受限的边缘 GPU 上，TileKernels 的精细内存管理和调度优化可以最大化有限硬件资源的利用率，适用于自动驾驶、工业视觉等场景。

总结

TileKernels 代表了 AI 基础设施领域的一个重要方向：通过高级的编程抽象和自动编译优化，让高性能 GPU 算子开发不再是少数 CUDA 专家的专利。对于追求极致推理性能的 AI 工程团队，TileKernels 提供了一条从"手写 CUDA"到"声明式算子开发"的平滑迁移路径。项目目前处于快速发展阶段，社区活跃度高，推荐对 GPU 性能优化感兴趣的同学深入学习和贡献。

文章摘要