Skip to content

Latest commit

 

History

1 Commit

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

llm-fusion-compiler

Python CUDA SM75 License Status Tests Correctness TFLOPs Efficiency Patterns Hardware

A compiler MVP that detects fusion patterns in Transformer operations, generates optimized CUDA kernels, compiles them with nvcc, and executes them on real GPU hardware with validated correctness.

IR graph -> pattern fusion -> CUDA codegen -> nvcc compile -> GPU execute -> validate

This is a research and educational project demonstrating compiler infrastructure for LLM operator fusion. It exposes the full compilation pipeline as independently testable phases, designed for learning and demonstration rather than as a production replacement for existing tools (TVM, MLIR, TensorRT).


What This Project Demonstrates

This is not a simulation. The compiler generates real CUDA kernels that:

  • Compile via nvcc for SM75 (Turing)
  • Execute on a real RTX 2070 GPU
  • Produce correct output validated against PyTorch (max error < 0.004)
  • Achieve 10.5 TFLOPs (36.7% of SM75 FP16/FP32-acc Tensor Core peak)

Hardware and Precision

Target:              SM75 (Turing) -- RTX 2070
Tensor Cores:        FP16 input, FP32 accumulate (WMMA API)
Peak (FP16/FP32acc): 28.5 TFLOPs
Memory BW:           448 GB/s
CUDA Toolkit:        11.5

Benchmarks

Generated GEMM+Bias+GeLU fused kernel, compiled with nvcc, measured on RTX 2070 with CUDA event timing, warmup, and synchronization.

Config: BM=64 BN=64 BK=32 threads=256
WMMA API, vectorized float4 global loads, register-resident accumulators

Shape      Gen kernel   cuBLAS ref   Ratio (gen/cuBLAS)   Max error
M=512       0.48ms       0.30ms            1.60x            0.002
M=1024      3.35ms       1.54ms            2.18x            0.004
M=2048      6.59ms       3.20ms            2.06x            0.004
M=4096     13.30ms       5.58ms            2.38x            0.004

Peak throughput: 10.5 TFLOPs at M=4096 (36.7% of SM75 peak)
FLOPs formula: 2*M*N*K + elementwise ops (standard convention)
Correctness: PASS at all shapes (max abs error < 0.004)

Note: ratio > 1.0 means the generated kernel is slower than cuBLAS. cuBLAS uses software pipelining, register tiling, and per-GPU tuned assembly. The generated kernel is a first-generation WMMA template.

Optimization journey

Version                          TFLOPs    % Peak    vs V1
V1 (scalar loads, smem accum)     3.0      10.5%     baseline
V2 (register accum)              3.1      10.9%     1.03x
V3 (vectorized float4 loads)    10.5      36.7%     3.50x

Vectorized global memory loads (float4 = 8 halfs per instruction) were the single highest-impact optimization, providing 3.5x speedup. This confirms that on Turing without cp.async, memory load efficiency dominates kernel performance over compute scheduling.


Architecture

IR (ops.py, graph.py, tensor.py, builder.py)
       |
Pattern Matcher (passes/fusion.py)
  - topological traversal of dataflow graph
  - linear chain detection
  - 9 patterns with priority-based selection
  - no double-fusion (each op used at most once)
       |
Code Generator (codegen/generator.py)
  - Jinja2 templates -> CUDA source
  - WMMA API (FP16 in, FP32 acc)
  - vectorized float4 global loads
  - register-resident accumulators
  - fused epilogue (bias + activation inline)
       |
NVCC Wrapper (codegen/nvcc_wrapper.py)
  - programmatic compilation to .so
  - error handling and PTX generation
       |
Auto-Tuner (tuner/autotuner.py)
  - search space: BM/BN/BK/threads
  - hardware constraints: 48KB smem, 1024 threads
  - grid search with profiling
  - heuristic fallback when nvcc unavailable
  - persistent JSON cache per (pattern, shape, arch)

Supported Fusion Patterns

Pattern               Operations fused
gemm_bias_gelu        GEMM + Bias + GeLU (fast approximation)
gemm_bias_silu        GEMM + Bias + SiLU
gemm_bias_relu        GEMM + Bias + ReLU
gemm_bias             GEMM + Bias
gemm_bias_residual    GEMM + Bias + Residual add
gemm_gelu             GEMM + GeLU
gemm_silu             GEMM + SiLU
layernorm_gemm        LayerNorm + GEMM
rmsnorm_gemm          RMSNorm + GEMM (two-kernel with warp reduction)

End-to-End Pipeline

FFN block (SwiGLU style):
  IR:      7 ops (3 GEMMs, 1 activation, 1 elementwise, 1 bias, 1 residual)
  Fusion:  4 ops (gemm_silu + gemm_bias_residual fused)
  Codegen: 2 CUDA kernels generated
  Compile: 2 .so files via nvcc
  Tuner:   2 tile configs selected

Correctness and Testing

39 tests passing across 6 test suites
GeLU approximation: max error 0.000473 vs torch.gelu
Generated kernel:   max error 0.004 vs PyTorch at all shapes
Reference suite:    24/24 pattern x shape combinations validated

Test suites:
  test_ir.py          7 tests    IR, builder, validation, topological sort
  test_fusion.py     11 tests    pattern matching, rewrite, priority, edge cases
  test_codegen.py     7 tests    template rendering, nvcc compilation
  test_tuner.py       7 tests    search space, constraints, cache, heuristic
  test_correctness.py 5 tests    PyTorch reference, GeLU, LayerNorm, RMSNorm
  test_pipeline.py    2 tests    end-to-end IR -> compile -> tune

Comparison with Existing Tools

Tool             Scope                                    Relationship to this project
Torch Inductor   Production compiler, kernel library      Full-featured; no fusion IR exposed
TVM              Full compiler stack                      Requires hardware-specific templates
Triton           Python DSL for GPU kernels               Higher-level than CUDA C++
CUTLASS          C++ template library for GEMM            Requires manual composition
This project     Minimal compiler MVP                     Full pipeline as testable phases

This project exposes the complete compilation pipeline (IR, pattern matching, code generation, compilation, auto-tuning) as independently testable Python modules. It is designed for learning and demonstration, not as a production replacement for the tools above.


Quick Start

pip install -r requirements.txt

# Full pipeline demo
python examples/fusion_demo.py
python examples/codegen_demo.py
python examples/tuning_demo.py

# All tests
python tests/test_ir.py
python tests/test_fusion.py
python tests/test_codegen.py
python tests/test_tuner.py
python tests/test_correctness.py
python tests/test_pipeline.py

# Benchmark (requires CUDA GPU + nvcc)
python benchmarks/benchmark_suite.py

Example: GEMM + Bias + GeLU

from src.ir.builder import GraphBuilder
from src.passes.fusion import find_fusion_matches, apply_fusion
from src.codegen.generator import CodeGenerator

b = GraphBuilder('example')
x    = b.input('x',    (1024, 4096))
w    = b.input('w',    (4096, 4096))
bias = b.input('bias', (4096,))
y    = b.gemm(x, w)
y2   = b.bias(y, bias)
y3   = b.activation(y2, 'gelu')
b.output(y3)
graph = b.build()

matches = find_fusion_matches(graph)
_, fused_ops = apply_fusion(graph, matches)
# -> [gemm_bias_gelu]: 3 ops fused into 1 kernel

gen = CodeGenerator('build/kernels')
cu_path = gen.generate_and_save(fused_ops[0])
# -> build/kernels/gemm_bias_gelu_64x64x32.cu

# With nvcc: compile and execute on real GPU
from src.codegen.nvcc_wrapper import compile_kernel
so_path = compile_kernel(cu_path, arch='sm_75')
# -> gemm_bias_gelu_64x64x32.so (loadable via ctypes)

Auto-Tuning

Search space:   BM in [32,64] x BN in [32,64] x BK in [16,32] x threads in [128,256]
Total configs:  16 per pattern
Constraints:    shared memory <= 48KB (including epilogue buffer)
                threads <= 1024, all dims multiples of 16
Best config:    BM=64 BN=64 BK=32 threads=256
Cache:          JSON per (pattern, M, N, K, arch)

Project Structure

llm-fusion-compiler/
+-- src/
|   +-- ir/              TensorDesc, Op, Graph, GraphBuilder
|   +-- passes/          Pattern matcher, validation
|   +-- codegen/         CodeGenerator, CUDA templates, NVCC wrapper
|   +-- tuner/           Search space, profiler, AutoTuner
+-- tests/               39 tests across 6 suites
+-- benchmarks/          Real kernel benchmarks
+-- examples/            Runnable demos
+-- DESIGN.md            Architecture decisions
+-- summary.txt          Project summary

Limitations

  • Forward pass only (no backward)
  • FP16 input, FP32 accumulate only (no INT8 or BF16)
  • SM75 target (no cp.async, no persistent kernels)
  • Programmatic IR only (no PyTorch or ONNX graph parser)
  • Generated kernels at 37% peak (vs cuBLAS at 95%)
  • 48KB shared memory limit constrains tile sizes to BM=BN=64
  • No software pipelining or double buffering

Documentation


License

MIT License -- Copyright (c) 2026 Joao Felipe De Souza


Author

Joao Felipe De Souza 2026

About

Compiler MVP that detects Transformer fusion patterns, generates optimized CUDA kernels with WMMA Tensor Cores, and executes them on real GPU hardware — 10.5 TFLOPs on RTX 2070, correctness validated against PyTorch.

Topics

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages