MoE: close out the tuning space in NOTES.md
Re-priced stage 1’s epilogue on the graded layout rather than on func_a: staging it through shared memory is within 0.4% either way (6.522 base, 6.531 merged, 6.506 unmerged), because the whole epilogue is only 2.6% of the scope. Stage 2 is the opposite case and keeps its staged store when unpaired.
Also states the wall explicitly: registers, not shared memory. Stage 1’s two 128x128 fp32 accumulators are 64 VGPRs at 512 threads and the kernel needs <=128 for four waves per SIMD, i.e. two blocks per CU, so bm*bn cannot exceed ~19456. Dropping to one accumulator by splitting gate and up into two passes would double bm but move 85 flops/byte to 87 plus an intermediate round trip – a net loss. Shared memory is not the constraint: exp_occupancy.py puts it at >=192 KiB per CU.
Co-Authored-By: Claude Fable 5 noreply@anthropic.com
版权所有:中国计算机学会技术支持:开源发展技术委员会
京ICP备13000930号-9
京公网安备 11010802047560号
Tile Language
Tile Language (tile-lang) is a concise domain-specific language designed to streamline the development of high-performance GPU/CPU kernels (e.g., GEMM, Dequant GEMM, FlashAttention, LinearAttention). By employing a Pythonic syntax with an underlying compiler infrastructure on top of TVM, tile-lang allows developers to focus on productivity without sacrificing the low-level optimizations necessary for state-of-the-art performance.
Latest News
T.gemm_spfor 2:4 sparse tensor core support, check out Pull Request #526 for details.T.printfor printing variables/buffers (docs) and a memory layout plotter (examples/plot_layout).Tested Devices
Although tile-lang aims to be portable across a range of Devices, it has been specifically tested and validated on the following devices: for MetaX GPUs, it includes the C500; for NVIDIA GPUs, this includes the H100 (with Auto TMA/WGMMA support), A100, V100, RTX 4090, RTX 3090, and RTX A6000; for AMD GPUs, it includes the MI250 (with Auto MatrixCore support) and the MI300X (with Async Copy support).
OP Implementation Examples
tile-lang provides the building blocks to implement a wide variety of operators. Some examples include:
Within the
examplesdirectory, you will also find additional complex kernels—such as convolutions, forward/backward passes for FlashAttention, more operators will continuously be added.Benchmark Summary
TileLang achieves exceptional performance across a variety of computational patterns. Comprehensive benchmark scripts and settings are available at tilelang-benchmark. Below are selected results showcasing its capabilities:
MLA Decoding Performance on H100
Flash Attention Performance on H100
Matmul Performance on GPUs (RTX 4090, A100, H100, MI300X)
Dequantize Matmul Performance on A100
Installation
Build from Source
We currently provide several ways to install tile-lang from source:
Method 3: Install with Nightly Version
For users who want access to the latest features and improvements before official releases, we provide nightly builds of tile-lang.
Quick Start
In this section, you’ll learn how to write and execute a straightforward GEMM (matrix multiplication) kernel using tile-lang, followed by techniques for layout optimizations, pipelining, and L2-cache–friendly swizzling.
GEMM Example with Annotations (Layout, L2 Cache Swizzling, and Pipelining, etc.)
Below is an example that demonstrates more advanced features: layout annotation, parallelized copy, and swizzle for improved L2 cache locality. This snippet shows how to adapt your kernel to maximize performance on complex hardware.
Dive Deep into TileLang Beyond GEMM
In addition to GEMM, we provide a variety of examples to showcase the versatility and power of TileLang, including:
Upcoming Features
Check our tilelang v0.2.0 release plan for upcoming features.
TileLang has now been used in project BitBLAS and AttentionEngine.
Join the Discussion
Welcome to join our Discord community for discussions, support, and collaboration!
Acknowledgments
We would like to express our gratitude to the TVM community for their invaluable contributions. The initial version of this project was mainly developed by LeiWang1999, chengyupku and nox-410 with supervision from Prof. Zhi Yang at Peking University. Part of this work was carried out during an internship at Microsoft Research, where Dr. Lingxiao Ma, Dr. Yuqing Xia, Dr. Jilong Xue, and Dr. Fan Yang offered valuable advice and support. We deeply appreciate their mentorship and contributions.