⚡ LLMBoost

Compiler-level kernel fusion for LLM inference

---

🎯 What is this?

LLMBoost is an MLIR compiler pass that automatically detects and fuses the RMSNorm → Linear pattern found in every transformer decoder layer — eliminating one full HBM round-trip. Validated on a 4× NVIDIA A30 cluster with real hardware benchmarks.

Bottom line: 2.20× faster fp16 inference, 1.55× bf16. No model changes required. The pass fires automatically on any IR containing the pattern.

---

📊 Benchmark Results

Device  :  4× NVIDIA A30  (SM80, 24 GB HBM2, CUDA 12.3)
Shape   :  [512, 4096] × [4096, 4096]   fp16

Kernel	Latency	Speedup	Notes
PyTorch unfused	0.340 ms	1.00×	Two separate HBM passes
LLMBoost fused (CuTe)	0.208 ms	2.20×	Single HBM pass + CUTLASS CuTe

Correctness (vs. fp32 PyTorch reference):

Metric	Value	Status
max_abs_err	1.07e-02	✅ within fp16 GEMM tolerance
mean_abs_err	9.27e-04	✅ within fp16 GEMM tolerance
mean_rel_err	1.48e-02	✅ within fp16 GEMM tolerance

---

💡 The problem: two HBM passes where one is enough

Every transformer decoder layer runs RMSNorm immediately before a linear projection. The unfused path:

x ──► [HBM read] ──► RMSNorm ──► [HBM write] ──► [HBM read] ──► Linear ──► y
                                        ↑
                             eliminated by LLMBoost

The normalised activation is written to HBM then immediately read back for the GEMM. At 4096-wide hidden dimensions this round-trip is the dominant memory bottleneck in autoregressive decode. LLMBoost fuses both ops — normalised values stay on-chip.

Design decisions

Why MLIR and not Triton? An MLIR pass is composable — it fires automatically on any IR containing the pattern, across any model, without the caller opting in. A Triton kernel requires manual dispatch per call site and does not integrate into compiler pipelines like TensorRT-LLM.

Why MLIR and not torch.compile? torch.compile fuses pointwise elementwise ops via Triton but does not cross the RMSNorm/GEMM boundary — the normalised tensor still materialises in HBM. LLMBoost operates at IR level where the full dataflow graph is visible, enabling cross-op fusion that Triton-level backends cannot attempt.

Why cuBLAS for the GEMM? cuBLAS uses vendor-tuned Tensor Core configurations per GPU generation. Hand-written GEMMs at 4096×4096 are rarely competitive without a search. The TVM MetaSchedule path (tvm_tune.py) replaces cuBLAS with an auto-tuned kernel found via Bayesian optimisation over tile sizes, loop orderings, and vectorisation widths.

Why the one-user safety guard? If the normalised tensor feeds two downstream consumers — two projections in a MoE gate, for example — fusing would require recomputing RMSNorm twice or buffering the result, potentially making things worse. The pass skips conservatively. The negative lit test (fuse_negative.mlir) locks this invariant automatically.

---

🏗️ Architecture

mlir-pass/
├── include/Transforms/
│   ├── LLMOps.td              ← TableGen: fused op declaration + verifier
│   └── LLMTransforms.h        ← Pass factory declarations
├── lib/Transforms/
│   ├── LLMDialect.cpp         ← Dialect init + shape verifier
│   ├── FuseRMSNormLinear.cpp  ← Pattern matcher + rewrite pass
│   └── LLMLowering.cpp        ← Lower fused op → LLVM external call
├── kernels/
│   └── fused_rmsnorm_linear.cu  ← CUDA kernel (RMSNorm + cuBLAS HGEMM)
└── test/
    ├── fuse_positive.mlir     ← lit: fusion fires correctly
    └── fuse_negative.mlir     ← lit: two-user safety guard holds

tvm-tuning/
├── scripts/
│   ├── build_relax_module.py  ← TVM Relax IRModule + TIR lowering
│   └── tvm_tune.py            ← MetaSchedule 1000-trial Bayesian tuning
└── benchmarks/
    ├── benchmark.py           ← 3-way latency benchmark
    └── test_correctness.py    ← pytest correctness suite

---

📈 Speedup across shapes

Batch Size	D=1024	D=2048	D=4096
1	6.31×	5.14×	1.76×
32	5.68×	5.43×	1.83×
128	4.95×	4.17×	1.91×
512	2.64×	1.78×	1.76×
1024	2.73×	1.93×	1.58×

Speedup is highest at small batch sizes where the RMSNorm overhead dominates. At large batch + large dim the GEMM saturates HBM bandwidth and both kernels converge.

🔢 dtype support

dtype	Speedup	Notes
fp16	2.20×	cuBLAS HGEMM + CuTe RMSNorm
bf16	1.55×	cublasGemmEx + CuTe RMSNorm

🔧 Technical Deep-Dive

1. Custom MLIR Dialect (LLMOps.td)

Defined llm.fused_rmsnorm_linear in TableGen with four typed inputs, a shape verifier, and custom assembly format. The verifier fires at IR construction time — bad shapes are caught before execution:

// BEFORE: three HBM transactions
%sum_sq = linalg.generic { iterator_types = ["parallel", "reduction"] } ...
%normed  = linalg.generic { iterator_types = ["parallel", "parallel"]  } ...
%result  = linalg.matmul ins(%normed, %w_proj : ...)

// AFTER: one HBM transaction
%result = llm.fused_rmsnorm_linear(%x, %w_norm, %w_proj, epsilon = 1.0e-5)
          : (tensor<512x4096xf16>, tensor<4096xf16>, tensor<4096x4096xf16>)
         -> tensor<512x4096xf16>

2. Pattern Matcher (FuseRMSNormLinear.cpp)

OpRewritePattern<linalg::MatmulOp> fires when a matmul's A-input comes from a linalg.generic matching the RMSNorm normalize signature. Two checks: iterator types AND block body structure — catching only the exact pattern, not false positives like L2 norm or softmax denominators:

static bool isRMSNormReduceBody(linalg::GenericOp op) {
    // iterator_types = ["parallel", "reduction"]
    // block body: mulf → addf  (sum-of-squares)
}
static bool isRMSNormNormalizeBody(linalg::GenericOp op) {
    // iterator_types = ["parallel", "parallel"]
    // block body: contains math::RsqrtOp
}

3. CUDA Kernel

Two-level warp/block reduction for RMSNorm (zero global memory traffic for the intermediate), then cuBLAS HGEMM with Tensor Core math mode:

// Warp reduce → shared memory → block reduce → rsqrt
for (int m = 16; m > 0; m >>= 1)
    ss += __shfl_xor_sync(0xffffffff, ss, m);
if (lane == 0) smem[wid] = ss;
__syncthreads();
if (wid == 0) {
    ss = (lane < (THREADS/32)) ? smem[lane] : 0.f;
    for (int m = 16; m > 0; m >>= 1)
        ss += __shfl_xor_sync(0xffffffff, ss, m);
    if (lane == 0) smem[0] = rsqrtf(ss / d_in + epsilon);
}

4. MLIR Lowering Pipeline

mlir-opt input.mlir \
  --fuse-rmsnorm-linear \      # detects pattern, emits llm.fused_rmsnorm_linear
  --convert-llm-to-llvm        # lowers to LLVM external call → libLLMKernels.so

5. TVM MetaSchedule Tuning (tvm_tune.py)

Bayesian optimisation over TIR transformations across all 4 GPUs in parallel:

• Search space: tile sizes, loop orderings, shared memory configs, vectorisation widths, thread/block dimensions
• Cost model: XGBoost surrogate trained on measured latencies
• Runner: 4 RPCRunners (one per A30), 1000 trials, 64 candidates/iteration
• Target: a30_cluster.yaml — SM80, 49 KB shared memory, 1024 max threads/block

---

🛠️ Tech Stack

Layer	Technology
Compiler IR	MLIR 17 (linalg, func, llvm dialects)
Op definition	TableGen (-gen-op-decls, -gen-op-defs)
Pattern rewriting	OpRewritePattern + GreedyPatternRewriteDriver
GPU kernel	CUDA C++ (SM80/SM86) + cuBLAS HGEMM
Auto-tuning	TVM MetaSchedule + XGBoost cost model
Tiling	CUTLASS CuTe 4.2 — Layout + Tensor abstractions for RMSNorm kernel (SM80/SM86)
Build	CMake 4.3 + Ninja, gcc 11.4, nvcc 12.3
Testing	MLIR lit + FileCheck, pytest
Cluster	4× NVIDIA A30, 24 GB HBM2, CUDA 12.3

---

🚀 Build & Run

# 1. Download prebuilt LLVM 17 (no compile needed)
wget https://github.com/llvm/llvm-project/releases/download/llvmorg-17.0.6/clang+llvm-17.0.6-x86_64-linux-gnu-ubuntu-22.04.tar.xz
tar -xf clang+llvm-17.0.6-x86_64-linux-gnu-ubuntu-22.04.tar.xz
export LLVM_BUILD=$PWD/clang+llvm-17.0.6-x86_64-linux-gnu-ubuntu-22.04

# 2. Build pass library + CUDA kernel
mkdir -p mlir-pass/build && cd mlir-pass/build
cmake -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_C_COMPILER=/usr/bin/gcc \
    -DCMAKE_CXX_COMPILER=/usr/bin/g++ \
    -DMLIR_DIR=$LLVM_BUILD/lib/cmake/mlir \
    -DLLVM_DIR=$LLVM_BUILD/lib/cmake/llvm \
    ..
ninja

# 3. Benchmark
cd ../..
python tvm-tuning/benchmarks/benchmark.py \
    --mlir_so mlir-pass/build/lib/libLLMKernels.so \
    --batch_seq 512 --d_in 4096 --d_out 4096

# 4. Lit tests (2/2 expected to pass)
cd mlir-pass/build && ninja check-mlir-llm

# 5. All 4 GPUs in parallel
cd ../..
python scripts/run_all_gpus.py \
    --mlir_so mlir-pass/build/lib/libLLMKernels.so \
    --batch_seq 512 --d_in 4096 --d_out 4096

---

📁 Key Files

File	What it does
FuseRMSNormLinear.cpp	Core pattern matcher + rewrite pass
fused_rmsnorm_linear.cu	RMSNorm CUDA kernel + cuBLAS HGEMM wrapper
LLMOps.td	TableGen op definition with shape verifier
LLMLowering.cpp	MLIR → LLVM external call lowering
tvm_tune.py	MetaSchedule tuning pipeline
benchmark.py	3-way benchmarking harness

---

Sai Teja Srivilli · GitHub @saitejasrivilli

Built on a 4× NVIDIA A30 cluster.