A hackable compiler to generate efficient fused GPU kernels for AI models [P]
The modern ML (LLM) compiler stack is brutal. TVM is 500K+ lines of C++. PyTorch piles Dynamo, Inductor, and Triton on top of each other. I built a hackable LLM compiler from scratch and am documenting the process. It takes a small model (TinyLlama, Qwen2.5-7B) and lowers it to a sequence of CUDA kernels through six IRs.
Currently, on RTX 5090, the emitted FP32 kernels run at geomean 1.11× vs PyTorch eager and 1.20× vs torch.compile, with full-block parity on TinyLlama-128 and Qwen2.5-7B at seq=128. Wins on small reductions / SDPA / kv-projections (up to 4.7×); losses on dense matmul at seq=512.
Part 1 took an RMSNorm layer end-to-end and walked the upper half of that pipeline in detail. This second part closes the gap and explains Tile IR, Kernel IR, and associated lowering rules in depth.
Full article: A Principled ML Compiler Stack in 5,000 Lines of Python
The article focuses on producing a GPU schedule for an operation written in loop-nest form (Loop IR). Example for RMSNorm:
v0 = reciprocal(2048)
for a0 in 0..32: # free
for a1 in 0..2048: # reduce
in2 = load x[0, a0, a1]
v1 = multiply(in2, in2)
acc0 <- add(acc0, v1)
v2 = multiply(acc0, v0)
v3 = add(v2, 1e-06)
v4 = rsqrt(v3)
for a2 in 0..2048: # free
in3 = load x[0, a0, a2]
in4 = load p_weight[a2]
v5 = multiply(in3, v4)
v6 = multiply(v5, in4)
merged_n0[0, a0, a2] = v6
The stack mimics a sequence of optimization steps a CUDA engineer would perform when optimizing kernels: stage inputs to smem, reduce bank conflicts, increase occupancy, and so on.
LoopOp
│
▼
[001] tileify — lift outer free Loops to thread axes
[002] chunk_matmul_k — chunk the K reduce into K-outer × K-inner (intra-CTA)
[003] split_matmul_k — promote the K-outer chunk loop into a grid dimension
[004] cooperative_reduce — let multiple threads share one reduce; tree-merge with Combine
[005] blockify_launch — pick block extents; partition free axes into BLOCK and THREAD
[006] chunk_reduce — chunk non-matmul reduces so their Loads fit in shared memory
[007] stage_inputs — hoist hot input slabs into Stage nodes
[008] register_tile — replicate the inner tile so each thread owns a register block
[009] permute_register_tile — reorder the register strip so bank-conflicting loads land on far columns
[010] double_buffer — promote K-outer Stages to BufferedStage (ping-pong)
[011] tma_copy — narrow eligible BufferedStages to TmaBufferedStage (sm_90+)
[012] split_inner_for_swizzle — split the inner cache axis of a TmaBufferedStage for swizzle
[013] async_copy — narrow the rest to AsyncBufferedStage (cp.async, sm_80+)
[014] pad_smem — pad shared-memory strides to break bank conflicts
[015] pipeline_k_outer — rotate the K-outer loop into prologue/steady-state/epilogue (cp.async + TMA)
[016] mark_unroll — annotate small inner loops for #pragma unroll
│
▼
TileOp (fully scheduled)
Each stage can be reproduced with a CLI command. For example, the stage_inputs pass stages input buffers into smem if possible and if there is a benefit in doing that (inputs are being read multiple times within CTA). To see it, the following command can be used:
deplodock compile \
-c "torch.nn.RMSNorm(2048)(torch.randn(1,32,2048))" \
--ir tile -vv \
| awk '/^>>> t:007/,/^<<< t:007/'
>>> t:007_stage_inputs
@@ matched at rms_norm (in-place) @@
@@ -2,6 +2,7 @@
v0 = reciprocal(2048)
Tile(axes=(a0:256=THREAD, a1:32=BLOCK)):
+ x_smem = Stage(x, origin=(0, a1, 0), slab=(a2:2048@2))
StridedLoop(a2 = a0; < 2048; += 256): # reduce
- in2 = load x[0, a1, a2]
+ in2 = load x_smem[a2]
v1 = multiply(in2, in2)
acc0 <- add(acc0, v1)
@@ -11,5 +12,5 @@
v4 = rsqrt(v3)
StridedLoop(a2 = a0; < 2048; += 256): # free
- in3 = load x[0, a1, a2]
+ in3 = load x_smem[a2]
in4 = load p_weight[a2]
v5 = multiply(in3, v4)
<<< t:007_stage_inputs
The final CUDA kernel for the RMSNorm layer:
deplodock compile \
-c "torch.nn.RMSNorm(2048)(torch.randn(1,32,2048))" \
--target sm_120 --ir cuda
extern "C" __global__
__launch_bounds__(256) void k_rms_norm_reduce(
const float* x, const float* p_weight, float* rms_norm) {
float v0 = 1.0f / 2048.0f;
int a1 = blockIdx.x;
int a0 = threadIdx.x;
int lane = threadIdx.x & 31;
int warp = threadIdx.x >> 5;
float acc0 = 0.0f;
__shared__ float x_smem[2048];
for (int x_smem_flat = a0; x_smem_flat < 2048; x_smem_flat += 256) {
float x_smem_v = x[a1 * 2048 + x_smem_flat];
x_smem[x_smem_flat] = x_smem_v;
}
__syncthreads();
for (int a2 = a0; a2 < 2048; a2 += 256) {
float in2 = x_smem[a2];
float v1 = in2 * in2;
acc0 += v1;
}
float acc0_w = acc0;
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 16);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 8);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 4);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 2);
acc0_w = acc0_w + __shfl_xor_sync(0xffffffff, acc0_w, 1);
__shared__ float acc0_smem[8];
if (lane == 0) {
acc0_smem[warp] = acc0_w;
}
__syncthreads();
for (int s = 4; s > 0; s >>= 1) {
if (warp < s) {
acc0_smem[warp] = acc0_smem[warp] + acc0_smem[warp + s];
}
__syncthreads();
}
float acc0_b = acc0_smem[0];
float v2 = acc0_b * v0;
float v3 = v2 + 1e-06f;
float v4 = rsqrtf(v3);
for (int a2 = a0; a2 < 2048; a2 += 256) {
float in3 = x_smem[a2];
float in4 = p_weight[a2];
float v5 = in3 * v4;
float v6 = v5 * in4;
rms_norm[a1 * 2048 + a2] = v6;
}
}