TileKernels from Beginner to Master: DeepSeek's TileLang-Powered GPU Kernel Guide

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

Summary

This comprehensive Chinese forum post is a multi-chapter tutorial on DeepSeek's open-source TileKernels library, which is built on the TileLang DSL for writing high-performance GPU kernels in Python. It opens by contrasting traditional CUDA C++ development — with its macro hell, long compile times, and fragile template metaprogramming — against TileLang's dataclass-based configuration and JIT compilation, which dynamically specialize kernels per hardware (16-byte vectorization on SM80/SM90, 32-byte on SM100). Chapter two dissects the swiglu_forward_and_per_token_cast_kernel, showing how SwiGLU activation, MoE top-k weight scaling, and FP8 per-token quantization are fused into a single kernel launch with clamping and persistent-kernel scheduling (num_blocks = num_sms * 4). Chapter three covers MoE gating and reduction: deterministic top-k tie-breaking via index comparison for cross-device consistency, and a fused reduce kernel that cuts memory traffic from two-read/two-write to one-read/one-write. Chapter four explains VRAM-saving techniques: MHC multilayer recomputation (trading compute for memory via pointer tables, shared memory double buffering, and async copy) and Engram XOR-based multi-table N-gram hashing. The final chapter describes a pytest-benchmark-based performance framework using pytest-xdist GPU binding, per-process memory fraction limits, and CUPTI-based microsecond timing. The post frames TileKernels as a benchmark for achieving deep memory fusion with minimal code.

TileKernels from Beginner to Master: A DeepSeek TileLang Tutorial

This post is a five-chapter tutorial on TileKernels, DeepSeek's open-source GPU kernel library built on the TileLang DSL (version >= 0.1.9). It tells the story of escaping the "tar pit" of traditional CUDA C++ development using modern, Python-first kernel engineering.

Key Points

Escaping the CUDA tar pit: Hand-written CUDA C++ exposes hardware details (block/warp/thread mapping, shared memory bank conflicts, register spilling), and small changes like switching from FP16 to FP8 trigger template-instantiation explosions and long compile times. DSLs like OpenAI's Triton and DeepSeek's TileLang offer a JIT-compiled, Python-based alternative that keeps low-level control (explicit memory layout) while hiding build complexity.
Why fusion matters — the memory wall: GPU compute grows faster than HBM bandwidth. In frameworks like PyTorch, an operation sequence such as Linear -> Swish -> Mul writes intermediates back to HBM between steps. TileKernels is positioned as an open-source showcase of deep memory fusion: computing and memory access are maximally overlapped with minimal code.

Chapter 1: TileLang First Experience (Configuration and JIT)

Operator configs (e.g., CastInputConfig) use Python dataclasses with frozen=True and @property-derived fields instead of C macros, giving type safety, immutability, and dynamic configuration for formats like FP8_E4M3 and FP4_E2M1.
get_per_token_cast_kernel acts as a kernel factory: get_best_vectorize_size selects a 16-byte vectorization baseline on Ampere/Hopper (SM80/SM90) and 32 bytes on Blackwell (SM100), letting JIT emit higher-throughput load/store instructions without code changes.
T.dynamic('num_tokens') removes fixed-shape dependence; T.Fragment and T.annotate_layout let developers describe logical layouts while the compiler handles physical register mapping, avoiding manual index math and bank conflicts.
The pipeline: perceive input type/shape → build config → JIT-compile to NVCC source → launch. Setting TK_PRINT_KERNEL_SOURCE exposes the generated low-level source for debugging and transparency.

Chapter 2: Anatomy of SwiGLU + Quant Single-Kernel Fusion

The swiglu_forward_and_per_token_cast_kernel fuses three stages into one launch: (1) SwiGLU activation on two input halves packed as one tensor of shape (num_expanded_tokens, hidden * 2) — a "conjoined twin" layout so x_L and x_R are fetched in a single memory access via a hidden offset; (2) scaling by MoE topk_weights in-register; (3) FP8 per-token quantization, computing the absmax scale factor before data ever leaves registers.
A clamping guard (use_clamp / count_clamp with T.atomic_add) caps runaway activations before quantization and counts how many values hit the threshold.
A persistent kernel strategy (num_blocks = num_sms * 4) keeps all SMs saturated across variable-length token streams — a major scheduling win for irregular MoE workloads.

Chapter 3: The Art of MoE Gating and Reduction

Deterministic top-k gating: floating-point drift across GPUs (e.g., in 8-GPU distributed runs) can flip expert selection. TileKernels breaks ties strictly by index — (other_top2_sum == topk_sum_var and i < lane_idx) — guaranteeing bit-identical routing everywhere.
Fused reduction: a textbook reduction does 2 reads + 2 writes of expert outputs (multiply weights, write back, re-read, sum). reduce_fused_kernel holds routing weights, quantization scale factors, and global scale in registers, compressing this to 1 read + 1 write — roughly a 50% bandwidth saving for DeepSeek-scale MoE models.
Dynamic addressing: a token_topk_to_pos mapping table lets the reduction kernel gather data regardless of load imbalance across experts, using parallel thread blocks and L2 cache to smooth bursty token distributions.

Chapter 4: VRAM Alchemy — MHC Recompute and Engram Hashing

MHC multilayer recomputation (tile_kernels/mhc/multilayer_recompute_kernel.py): only the initial residual is stored; each layer's input is recomputed in-kernel from pre_mix projections using register fragments, shared-memory double buffering, and async copy — trading a small amount of compute for large memory savings. T.ptr pointer tables are prebuilt on CPU to avoid dynamic addressing inside the kernel. Explicit casts (FP32 compute, BF16 storage) balance stability and bandwidth.
Engram hashing (tile_kernels/engram/engram_hash_kernel.py): full N-gram vocabularies are astronomically large (e.g., 50,000³ for trigrams). Engram instead uses XOR-multiplier hashing to map N-grams into a small set of learnable embedding tables, with per-table modulo and offsets. XOR is chosen for speed and bit-diffusion properties that reduce collisions.

Chapter 5: Performance Benchmarks with pytest_benchmark

A custom tests/pytest_benchmark_plugin.py integrates microsecond-level benchmarking into CI: with pytest-xdist, each worker binds to a GPU via CUDA_VISIBLE_DEVICES and is capped by torch.cuda.set_per_process_memory_fraction, reserving 10 GB for the system and preventing OOM collisions between concurrent workers.
Timing uses TileLang's do_bench with the CUPTI backend (warmup=0, rep=30), correctly accounting for CUDA's asynchronous execution rather than naive CPU-side timers.

Conclusion

The post closes by framing TileKernels as a victory of extreme engineering: performance is not requested from the hardware but squeezed out of it through memory-layout design, deep operator fusion, deterministic scheduling, and rigorous benchmarking — all in concise, JIT-compiled Python instead of CUDA macro-laden C++.

References (cited in the original post)

1. TileLang: <https://github.com/tile-ai/tilelang> 2. DeepSeek-V3 Technical Report: <https://github.com/deepseek-ai/DeepSeek-V3> 3. NVIDIA CUDA Programming Guide: <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html> 4. Shazeer, N. (2020). *GLU Variants Improve Transformer*. arXiv:2002.05202 5. Fedus, W., et al. (2022). *Switch Transformers*. JMLR 6. Tillet, P., Kung, H. T., & Cox, D. (2019). *Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations*