tilelang

What it is and what it does

Tilelang is a Python-embedded domain-specific language that lets you write GPU and CPU kernels at a higher level of abstraction than raw CUDA or HIP, while still accessing low-level optimizations. It compiles to efficient machine code via TVM, supporting multiple backends including NVIDIA CUDA, AMD HIP, Apple Metal, and emerging targets like NVIDIA CuTe DSL and Huawei Ascend. The language provides Pythonic syntax for declaring tensor shapes, allocating shared memory, managing data movement, and expressing computation patterns like matrix multiplication, attention mechanisms, and sparse operations.

You write a kernel as a Python function decorated with `@tilelang.jit`, declare tensor shapes and dtypes, allocate buffers, and express the computation using tile-level operations. Tilelang handles the compilation to optimized code for your target hardware. It integrates z3-solver for symbolic reasoning and automatic correctness verification, and recently migrated to apache-tvm-ffi to reduce CPU overhead. The package includes examples for GEMM, dequantization, FlashAttention, and MLA decoding, with tested support for NVIDIA (H100, A100, V100, RTX series), AMD (MI250, MI300X), and Apple Metal devices.

Use it for

• Implement custom GEMM kernels with layout optimization and L2-cache swizzling without writing raw CUDA
• Build high-performance attention kernels (FlashAttention, MLA) in ~80 lines of Python code
• Compile sparse tensor operations and dequantization kernels targeting multiple GPU architectures
• Prototype and optimize AI workload kernels while maintaining portability across NVIDIA, AMD, and Apple hardware
• Generate WebGPU or CuTe DSL code for emerging deployment targets without rewriting kernel logic

Worth the install?

AI-flagged interpretation of the facts on this page. Verify before relying on it.

Install

tilelang on PyPI

Before you install

Medium install friction due to 11 runtime dependencies including torch, apache-tvm-ffi, and z3-solver. Active maintenance with recent release (11 days old). Prebuilt wheels available for common platforms (macOS ARM64, Linux x86_64/aarch64, Windows), reducing compilation overhead.

Requires Python 3.10 or later; torch and apache-tvm-ffi are mandatory runtime dependencies with their own system requirements (CUDA/HIP toolchain for GPU targets).

License in practice

MIT license permits commercial and private use with minimal restrictions, making it suitable for production deployment in proprietary projects.

Quickstart

pip install tilelang

import tilelang as T

@T.jit
def matmul(A, B, block_M: int = 64, block_N: int = 64, block_K: int = 64):
    M, N, K = T.const('M, N, K')
    A: T.Tensor[[M, K], T.float16]
    B: T.Tensor[[K, N], T.float16]
    C = T.empty([M, N], T.float16)
    # kernel implementation follows

Package facts

Evidence: tilelang-0.1.13-cp38-abi3-macosx_11_0_arm64.whl; tilelang-0.1.13-cp38-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl; tilelang-0.1.13-cp38-abi3-manylinux_2_34_aarch64.whl; tilelang-0.1.13-cp38-abi3-win_amd64.whl

Tags

More Artificial Intelligence packages

See also nvidia-cuda-tileiras · cuda-tile · sgl-deep-gemm · flashinfer-python · flydsl · humming-kernels · nvidia-cutlass-dsl · helion · nvidia-cutlass-dsl-libs-core · sglang-kernel