DeepSeek TileKernels: A Complete Guide to GPU Kernel Development with Tilelang

Forum topic · ✨步子哥 · 2026-04-25

Summary

This in-depth tutorial series from zhichai.net explores DeepSeek's open-source TileKernels library, a collection of high-performance GPU kernels built on the Tilelang DSL (requiring Tilelang >= 0.1.9). The guide positions Tilelang as an escape from the maintenance burden of raw CUDA C++ development, replacing template metaprogramming and macro-heavy configs with Python dataclasses and JIT compilation. Key topics include: a hands-on introduction to Tilelang configuration, dynamic shapes, and hardware-aware vectorization (16-byte on Ampere/Hopper, 32-byte on Blackwell); the swiglu_forward_and_per_token_cast_kernel, which fuses SwiGLU activation, MoE top-k weight scaling, clamping, and FP8 per-token quantization into a single kernel launch; deterministic MoE gating with index-based tie-breaking and a fused reduction kernel that cuts memory traffic from two reads plus two writes to one read plus one write; memory alchemy via multilayer recompute (trading compute for VRAM) and Engram XOR-hash N-gram embedding compression; and a pytest-benchmark-based CI performance framework using CUPTI timing and per-worker GPU/memory isolation for multi-process testing. The series argues that deep operator fusion and register-level compute are the practical answers to the memory wall in LLM inference.

DeepSeek TileKernels: From Beginner to Expert

This is a structured English summary of a multi-chapter Chinese tutorial on DeepSeek's open-source TileKernels library, built on the Tilelang DSL (version >=0.1.9).

Key points

Motivation — escaping the "CUDA tar pit": Traditional CUDA C++ kernel development exposes hardware details (block/warp/thread mapping, shared memory bank conflicts, register spilling) and suffers from high maintenance costs: small changes like switching FP16 to FP8 (e4m3) or non-multiple-of-128 shapes trigger template instantiation explosions and long compile times.
Tilelang as the escape hatch: A domain-specific language letting developers write GPU kernels in Python with JIT compilation, while retaining low-level control (explicit memory layout via T.Fragment / T.annotate_layout).
Operator fusion beats the memory wall: In LLM inference, kernels are memory-bound. TileKernels' core idea is deep compute/memory fusion — minimal code, maximal overlap of computation with data movement.

Chapter highlights

1. Tilelang first experience

Operator configs use Python dataclasses (e.g., CastInputConfig with frozen=True and @property-derived fields) instead of C++ macros.
get_per_token_cast_kernel acts as a kernel factory: JIT selects the best vectorization width per architecture — 16 bytes on SM80/SM90 (Ampere/Hopper), 32 bytes on SM100 (Blackwell).
T.dynamic('num_tokens') enables dynamic shapes; @tilelang.jit compiles Python descriptions to optimized GPU code at runtime. Setting TK_PRINT_KERNEL_SOURCE prints generated kernel source for debugging.

2. SwiGLU + quantization in a single kernel

swiglu_forward_and_per_token_cast_kernel fuses SwiGLU activation, MoE topk_weights scaling, optional clamping (with atomic-add overflow counting), and FP8 per-token quantization (scale-factor computation from row-wise abs-max) into one launch — no intermediate HBM round-trips.
Input layout (num_expanded_tokens, hidden * 2) lets one memory access fetch both x_L and x_R via a simple offset.
Persistent kernel design (num_blocks = num_sms * 4) keeps all SMs busy for variable-length token streams.

3. MoE gating and reduction

Deterministic top-k selection: ties are broken by expert index (i < lane_idx), guaranteeing bitwise-consistent routing across distributed ranks despite floating-point drift.
reduce_fused_kernel performs weight application, quantization-factor and global-scale handling in registers, compressing the classic 2-read/2-write reduction into 1 read + 1 write — roughly 50% bandwidth savings.
Dynamic address mapping (token_topk_to_pos) handles imbalanced, variable-length expert workloads, with L2 cache smoothing load skew.

4. VRAM alchemy: MHC recompute and Engram hashing

tile_kernels/mhc/multilayer_recompute_kernel.py recomputes per-layer inputs in-kernel from an initial residual state (compute-for-memory trade), using shared-memory double buffering and T.ptr pointer tables to avoid dynamic addressing overhead.
tile_kernels/engram/engram_hash_kernel.py maps N-grams into compact learnable embedding tables via XOR-multiply hashing with per-table modulo and offsets, avoiding the exponential blow-up of raw N-gram vocabularies.
Guiding principles: registers first, async-copy compute/latency hiding, and explicit precision control (FP32 compute, BF16 storage).

5. Benchmarking discipline (pytest-benchmark)

tests/pytest_benchmark_plugin.py integrates performance tests into CI:
With pytest-xdist, each worker is pinned to a GPU (CUDA_VISIBLE_DEVICES from worker ID) and memory-capped via torch.cuda.set_per_process_memory_fraction, reserving 10 GB for the framework to prevent OOM collisions.
Timing uses the CUPTI backend via tilelang.profiler.bench.do_bench (30 reps, converted to microseconds), avoiding both cold-start error and CPU-side async dispatch mismeasurement.

Conclusion

The series frames TileKernels as a benchmark of "how to do deep memory fusion correctly with a DSL": performance is extracted through extreme memory-layout control, deep operator stitching, and hardware awareness — demonstrating that near-peak GPU throughput is achievable from Python-level code without sacrificing maintainability.

References cited in the original post

1. NVIDIA. *CUDA C++ Programming Guide* 2. Tillet, P., Kung, H. T., & Cox, D. (2019). *Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations* 3. TileLang: https://github.com/tile-ai/tilelang 4. DeepSeek-V3: https://github.com/deepseek-ai/DeepSeek-V3 5. Shazeer, N. (2020). *GLU Variants Improve Transformer*. arXiv:2002.05202 6. Fedus, W., et al. (2022). *Switch Transformers* 7. Lepikhin, D., et al. (2020). *GShard*. arXiv:2006.16668 8. Dao, T. (2022). *FlashAttention*. NeurIPS