TileLang puts CUDA kernel writing in Python, no C++ required
Curated by the Inblix editorial team
A new Python-native workflow for writing high-performance GPU kernels is taking shape, and it doesn’t require touching a line of C++. TileLang, a domain-specific language built on top of the Apache TVM compiler stack, lets developers design CUDA kernels—including complex tiled matrix multiplications and FlashAttention—entirely in Python. We put it through its paces, and the results are surprising. For a simple vector addition of 16 million elements, the auto-generated kernel hits a memory bandwidth of over 600 GB/s on a modern NVIDIA GPU. That’s a dead heat with PyTorch’s native operation, which is exactly the right outcome. “Both are pure bandwidth, so a tie is the correct outcome,” the tutorial notes matter-of-factly. There’s no magic here—just a compiler that understands thread mapping, shared memory banking, and synchronization well enough to stop humans from shooting themselves in the foot.
The real test, however, comes with workloads that actually stress a GPU’s architecture. TileLang exposes tensor core primitives directly in Python, letting you declare warp-level matrix fragments and orchestrate the pipelined data movement between global, shared, and register memory. The tutorial walks through a tiled GEMM implementation, complete with configurable warp tiling and async copy stages. On an SM80+ GPU with 96 KB of shared memory per block, you can push the number of pipeline stages higher to hide latency. The framework’s autotuning then sweeps through tile sizes and stage counts to find configurations that match cuBLAS within a few percent. That’s a big deal. Writing a competitive GEMM from scratch in CUDA C++ is a multi-week project for an experienced GPU programmer. Here, it’s a Python script that compiles in seconds on the second run thanks to a kernel cache.
Fused operations are where this approach really shines. The tutorial demonstrates a GEMM with an epilogue fusion—a common pattern in transformer inference where you’d normally pay the cost of a separate kernel launch and round-trip through global memory. By expressing the whole thing as a single TileLang function, the compiler stitches the matrix multiply and the pointwise operation together, eliminating the intermediate buffer. The generated CUDA source, which you can inspect directly, shows the fused loop structure. Row-wise softmax and a simplified FlashAttention kernel round out the examples, showing how reductions and online rescaling map onto the language’s primitives without manual warp shuffles or shared memory fence operations.
None of this makes CUDA expertise obsolete. You still need to understand why a 128x128 tile is sensible for your target architecture and when to double-buffer. But TileLang shifts the burden from “how do I implement this correctly?” to “what should the data flow look like?” For teams that need custom kernels but can’t justify a full-time GPU engineer, that’s a genuine shift in who can contribute. The fact that it ships with a bundled TVM and installs via pip—with a nightly channel fallback if the stable wheel acts up—means the barrier to entry is about as low as it gets in this space.
💡 Key Takeaways
- TileLang lets you write tiled tensor-core kernels in pure Python and compiles them to CUDA via TVM, matching cuBLAS within a few percent on autotuned configurations.
- The compiler handles thread mapping, synchronization, and memory layouts automatically, while still exposing explicit control over shared-memory tiles and register fragments.
- Fused operations like GEMM-with-epilogue avoid separate kernel launches and global memory round-trips, a direct win for transformer inference latency.
- Kernel source is inspectable, and autotuning sweeps architecture-dependent parameters like tile size and pipeline stages to find near-optimal configs for your specific GPU.
Keep reading: See related articles below for more coverage on this topic.
Get smarter about AI
The sharpest AI news, curated daily. Delivered free to your inbox.