TLX documentation

Tutorial kernels

GEMM and attention examples implemented with TLX.

GEMM kernels

Pipelined GEMM on Hopper

Warp-specialized GEMM on Hopper

Warp-specialized GEMM on Blackwell

Warp-specialized MXFP8 GEMM on Blackwell

Pipelined GEMM on Blackwell

CLC GEMM on Blackwell

2-CTA GEMM on Blackwell

Attention kernels

Warp-specialized pipelined persistent FA fwd/bwd on Blackwell

Warp-Specialized computation-pipelined pingpong FA fwd on Hopper

AMD kernels (gfx950 / CDNA4)

LDS-pipelined GEMM

Warp-pipelined GEMM

Async-DMA Flash Attention fwd — simple / prefetch

Persistent Flash Attention fwd — XCD zig-zag, cross-attention / decode

Rotated 4-cluster Flash Attention fwd/bwd

Fused addmm + GLU (Gated Linear Unit: out = x + x*y, x = A@B + bias)

IKBO Flash Attention (In-Kernel Broadcast Optimization, candidate/user broadcast)

IKBO LCE (logit cross-entropy over candidate/user embeddings — not attention)

AMD kernels (gfx1250)

TDM-pipelined GEMM

MXFP TDM-pipelined GEMM