5000 Lines of Python to Build a Hackable LLM Compiler: From TinyLlama to Qwen2.5 Generating CUDA Kernels Directly

The modern ML compiler stack has become a behemoth. TVM has more than 500,000 lines of C++ code, and PyTorch piles Dynamo, Inductor, and Triton on top of one another. One developer got tired of this complexity and built a hackable LLM compiler from scratch in 5000 lines of Python that can directly lower small models like TinyLlama and Qwen2.5-7B into efficient CUDA kernel sequences.

The project has recently drawn attention on Reddit’s r/MachineLearning. The author completes the full compilation pipeline through six layers of intermediate representation (IR). On an RTX 5090, the generated FP32 kernels achieve 1.11× performance compared to PyTorch eager mode, and 1.20× compared to torch.compile. It achieves full block-level alignment on TinyLlama-128 and Qwen2.5-7B (sequence length 128).

Performance: Crushing Small Ops, Losing on Big Matrices

The compiler’s performance profile is quite interesting. For small reductions, SDPA (Scaled Dot-Product Attention), and kv-projection operators, performance gains reach up to 4.7×. But for dense matrix multiplications with sequence length 512, it loses to PyTorch.

This performance behavior actually reflects the essence of compiler optimization: for compute-intensive large GEMMs, hand-optimized libraries like cuBLAS are already near hardware limits, so compilers can hardly beat them. But for irregularly shaped, memory-bound small ops, automatically fused kernels can significantly improve performance by reducing kernel launch overhead and optimizing memory access patterns.

Six Layers of IR: From High-Level Semantics to CUDA Code

The author explains the compiler’s architecture in detail. The entire pipeline includes six IR layers, each with a clear responsibility:

Upper IR (covered in Part 1):

• From model definition to computation graph
• Computation graph to Loop IR (loop-nest representation)

Lower IR (focus of Part 2):

• Loop IR to Tile IR: introduces tiling and parallelization strategies
• Tile IR to Kernel IR: maps to the GPU thread hierarchy
• Kernel IR to CUDA code generation

The author uses the RMSNorm layer as an end-to-end example, showing how a high-level operation is progressively lowered to an executable CUDA kernel. The layered design allows independent optimization and debugging at each stage, unlike TVM which mixes everything together.

Loop IR: Describing Computation via Nested Loops

Loop IR is the core abstraction of the compiler. It describes tensor computations in nested loop form, similar to Halide or TVM schedules. For an RMSNorm operation:

# Pseudocode example
for b in range(batch_size):
    for s in range(seq_len):
        # Compute root-mean-square
        sum_sq = 0
        for d in range(hidden_dim):
            sum_sq += x[b, s, d] ** 2
        rms = sqrt(sum_sq / hidden_dim + eps)
        # Normalize
        for d in range(hidden_dim):
            y[b, s, d] = x[b, s, d] / rms * weight[d]

This representation is intuitive and easy to transform. The compiler can reorder, fuse, or tile loops without understanding specific mathematical semantics.

Tile IR: Introducing GPU Parallelization

Tile IR adds tiling and parallelization annotations on top of Loop IR. It splits loops into outer (block-level) and inner (thread-level) parts, deciding tile sizes for each dimension.

For RMSNorm, the compiler might:

• Map the batch and seq_len dimensions to blocks
• Handle hidden_dim within threads
• Use shared memory to cache intermediate results

The key here is to generate appropriate tiling strategies automatically. The author mentions using heuristics—determining tile sizes from data reuse patterns and deciding whether to use shared memory based on register pressure.

Kernel IR: Mapping to the CUDA Thread Model

Kernel IR is the layer closest to CUDA. It explicitly defines each thread’s workload, shared-memory layout, and synchronization points. This level handles many low-level details:

• Thread index calculation: mapping threadIdx, blockIdx to data indices
• Memory hierarchy: managing global, shared, and register memory
• Synchronization primitives: inserting __syncthreads()
• Boundary handling: dealing with non-divisible tile sizes

The author emphasizes the "lowering" rules in Kernel IR. These define how to convert higher-level Tile IR operations into concrete CUDA operations—such as transforming a reduction loop into one that uses warp shuffles or atomic operations.

Why “Hackable”

The “hackable” label is deliberate. The author’s design philosophy is to keep each IR layer simple enough that developers can easily plug in custom optimizations.

By contrast, TVM’s schedule system, while powerful, has a steep learning curve. To add a new optimization pass, you must understand the whole compiler internals. This 5000-line compiler keeps each layer straightforward—you can directly modify lowering rules or add your own transformations.

Such design is valuable for researchers. If you want to test a new kernel fusion or memory optimization technique, you don’t need to navigate through 500,000 lines of code—just adjust the relevant IR layer.

Comparison with Mainstream Systems

vs PyTorch Inductor:
Inductor, PyTorch 2.0’s default compile backend, uses Triton for GPU code generation. Its advantage is deep integration with the PyTorch ecosystem and support for dynamic shapes and complex control flow. But its abstractions are very high-level, making fine-grained kernel optimization difficult. By generating CUDA directly, this project achieves better performance on small ops.

vs TVM:
TVM is widely used academically and industrially and supports many hardware backends. But its codebase is huge, making modification and extension costly. This project implements TVM’s core ideas in just 1% of the code. While less feature-complete, it’s sufficient for specific scenarios like small-model inference.

vs Triton:
Triton (by OpenAI) is a GPU programming language providing higher-level abstractions than CUDA. Some might ask: why not just use Triton? The author’s answer: Triton still requires writing each kernel by hand. This compiler instead automatically generates all kernels from the model definition.

Implementation Details: RMSNorm End-to-End Example

Part 1 of the article explains the full compilation process for the RMSNorm layer, which replaces LayerNorm in Llama models:

y = x / sqrt(mean(x^2) + eps) * weight

The compiler breaks this operation into three stages:

1. Reduction: compute the squared sum per token
2. Normalization: normalize via square root
3. Scaling: multiply by the learnable weight

At the Tile IR level, it decides:

• How many tokens each block processes
• Whether to use shared memory for intermediate results
• Whether reduction uses warp shuffle or atomics

At the Kernel IR level, it generates concrete CUDA code including:

• Thread index computations
• Shared memory declarations
• Loop unrolling and vectorized load/store
• Synchronization inserts

The resulting CUDA kernel runs 2–3× faster than PyTorch’s for small batch sizes, thanks to reduced launch overhead and optimized memory access.

Current Limitations and Future Directions

The author is candid about existing limitations:

Performance:

• Worse at large GEMMs than cuBLAS
• FP32 only; no mixed precision
• No advanced optimizations like FlashAttention

Functionality:

• Static shapes only
• No dynamic control flow
• Tested only on TinyLlama and Qwen2.5-7B

Engineering:

• No auto-tuning
• Error messages not friendly
• Lacks full test coverage

Future improvements include:

• Integrate cuBLAS for large GEMMs
• Add FP16/BF16 support
• Implement FlashAttention and advanced optimizations
• Support more model architectures
• Add performance-model-based auto-tuning

Takeaways for Developers

The project’s greatest value isn’t raw performance—it’s demonstrating how complex compiler functionality can be achieved with minimal, clean code. 5000 lines of Python prove that much of modern compiler complexity is avoidable.

For developers exploring LLM inference optimization, this project offers a full learning path:

1. Understand computation graph representation
2. Learn loop-nest optimization techniques
3. Master GPU parallelization strategies
4. Study CUDA programming best practices

For teams customizing inference engines, this shows you don’t need heavyweight frameworks like TVM or Triton. If your use case is narrow (specific models and hardware), building a lightweight compiler from scratch may be a better option.

Significance for the Open-Source Ecosystem

With mature inference engines like vLLM and TGI, why do we need this project? Because it addresses a different layer.

vLLM focuses on high-throughput production inference, using techniques such as PagedAttention and continuous batching. Its kernels are mainly hand-tuned CUDA code or calls to cuBLAS/FlashAttention.

This compiler instead focuses on automatically generating efficient kernels. It doesn’t aim to replace vLLM, but to demonstrate an alternate approach: if auto-generated kernels can come close to hand-optimized performance, then supporting new models or hardware becomes far easier—no manual CUDA coding.

Broadly speaking, projects like this advance the whole ecosystem. When someone demonstrates that 5000 lines of code can implement the core functionality, others are prompted to reflect on whether their complexity is necessary. This healthy competition benefits everyone.

Technical Details: Example of a Lowering Rule

To help readers understand how the compiler works internally, here’s a simplified example of lowering a reduction from Loop IR to Kernel IR:

# Loop IR
for i in range(N):
    sum += data[i]

# Tile IR (tile_size=256)
for block in range(N // 256):
    partial_sum = 0
    for thread in range(256):
        partial_sum += data[block * 256 + thread]
    # Reduction within block
    sum += partial_sum

# Kernel IR
__shared__ float shared_data[256];
int tid = threadIdx.x;
int bid = blockIdx.x;

// Each thread loads one element
shared_data[tid] = data[bid * 256 + tid];
__syncthreads();

// Tree reduction
for (int stride = 128; stride > 0; stride >>= 1) {
    if (tid < stride) {
        shared_data[tid] += shared_data[tid + stride];
    }
    __syncthreads();
}

// Thread 0 writes back result
if (tid == 0) {
    atomicAdd(&sum, shared_data[0]);
}

This example shows how the compiler represents a single computation at different IR levels, progressively introducing GPU-specific optimizations like shared memory, tree reduction, and atomic operations.

Why Small Ops Run Faster

There are several reasons for small-operator speedups:

Kernel fusion:
The compiler fuses multiple small ops into one kernel, reducing global memory reads/writes. For example, it can fuse RMSNorm with the following matmul to avoid writing intermediate results to VRAM.

Memory access optimization:
For irregular memory patterns, the compiler generates custom load/store code using vectorized instructions and coalesced memory access. PyTorch’s general kernels can’t optimize to this degree.

Register allocation:
Small ops allow more intermediates to stay in registers, minimizing shared/global memory traffic. The compiler can fine-tune register usage per computation.

Launch overhead:
Merging multiple small kernels into one large kernel greatly reduces the kernel launch overhead—which on an RTX 5090 can reach 10–20% of total runtime.

Links to Academic Research

Many of the project’s ideas stem from academic compiler research. Projects like Halide, TVM, and Tiramisu explore ways to describe tensor computations with high-level abstractions and generate efficient code automatically.

However, academic projects often aim for generality and completeness, leading to massive codebases. This project takes a more pragmatic route: focus solely on LLM inference and implement only the core functionality with the smallest amount of code.

This "subtractive" philosophy is valuable in engineering practice—often, 80% of performance gains come from 20% of techniques. Identifying and focusing on that 20% can yield great results with minimal complexity.

Real-World Use Cases

What’s this compiler good for?

Edge-device inference:
On resource-constrained hardware, hand-optimizing every kernel is too expensive. Automatically generated kernels may not be perfect but are good enough and adaptable to new models quickly.

Model architecture exploration:
Researchers designing new architectures need quick performance validation. This compiler auto-generates inference code—no manual CUDA kernel writing needed.

Custom inference engines:
If your product only supports a few models, you can build a lightweight inference engine on top of this compiler and avoid heavy dependencies like vLLM.

Teaching and learning:
For developers learning compiler technology, this is a complete and graspable example. You can read 5000 lines of code in a few days—TVM might take months.

Open-Source Project Info

The project is open-source on GitHub under the MIT license. The author mentioned on Reddit that they will continue updating it, aiming to add more model support and performance improvements soon.

The repository includes:

• Full compiler implementation (≈5000 lines of Python)
• Test cases for TinyLlama and Qwen2.5-7B
• Detailed documentation and tutorials
• Benchmarking scripts

The author encourages community contributions—particularly:

• Support for new model architectures
• Performance optimizations
• Bug reports/fixes
• Documentation improvements

Conclusion

This project demonstrates a key principle: modern ML compilers don’t need to be complex. With clear layering and a focused scope, 5000 lines of Python can implement an end-to-end pipeline from model definition to CUDA kernel generation.

It won’t replace mature projects like TVM or vLLM, but it offers developers a new option: if you need a lightweight, customizable, and understandable LLM compiler, it’s worth a look.

More importantly, it illustrates the essence of compiler design—progressively lowering high-level abstractions to low-level code, performing the right optimizations at each layer. This mindset applies far beyond LLM inference to any high-performance computing domain.

For AI infrastructure practitioners, the project offers inspiration: while we pursue universality and completeness, we shouldn’t forget the value of simplicity and maintainability. Sometimes, less is more.

---

References

• Reddit thread: A hackable compiler to generate efficient fused GPU kernels for AI models – original author post with technical discussion
• vLLM blog: New features and roadmap for vLLM – overview of vLLM’s CUDA kernel optimizations and caching techniques for comparison