Skip to content

About

Eight GPU-kernel projects for modern LLM serving

Topics

Resources

Contributing

Stars

1 star

Watchers

0 watching

Forks

Repository files navigation

LLM Serving Kernels

Eight compact GPU projects that move from fused tensor operations to one real vLLM integration.

The engineering rule is simple:

Define the PyTorch reference, implement the GPU path, prove correctness, benchmark against a named baseline, inspect generated code, and profile before making performance claims.

The primary validation target is one NVIDIA L4 (SM89, 24 GB). Source is written for NVIDIA GPUs supported by the installed CUDA, PyTorch, Triton, CUTLASS, and vLLM releases; every architecture still requires its own correctness and performance run.

Verified NVIDIA L4 release

The first verified release passed 51 GPU correctness tests across all eight projects on an NVIDIA L4.

Selected results:

  • fused RMSNorm: 1.71× faster than torch.compile;
  • paged decode: 5.12× faster than the materializing PyTorch reference;
  • FP8 KV cache: 49.9% less storage with 0.00110 decode mean absolute error;
  • vLLM custom operator: fully integrated and correct, currently 9.7% slower than stock vLLM.

See RESULTS.md for the complete environment, benchmark table, limitations, and raw evidence.

Projects

# Project Primary implementation What it teaches
01 RMSNorm Triton Row reductions, fusion, memory bandwidth
02 SwiGLU Triton Elementwise fusion and serving-layout contracts
03 FP16 GEMM Triton + CuTe DSL Tiling, Tensor Cores, layout reasoning
04 Paged decode attention CUDA C++ Page tables, GQA/MQA, online softmax
05 FP8 KV cache CUDA C++ E4M3 storage, page scales, fused dequantization
06 Exact top-k sampling Triton Hierarchical selection and categorical sampling
07 Tiny MoE Triton Top-2 routing, dispatch, grouped GEMM, combine
08 vLLM custom operator Python + Triton Out-of-tree serving-engine integration

Repository map

llm-serving-kernels/
├── kernel_lab/                 shared benchmark and extension utilities
├── projects/                   eight self-contained projects
├── tools/                      environment, PTX/SASS, Nsight, and suite tools
├── docs/                       design, methods, execution, and publishing
├── tests/                      CPU reference and repository tests
├── results/                    verified measurements only
├── modal_app.py                optional reproducible L4 runner
├── pyproject.toml
└── Makefile

Install the main environment

Requirements:

  • Linux x86-64;
  • Python 3.12;
  • an NVIDIA driver compatible with the selected CUDA build;
  • CUDA development toolkit and nvcc for Projects 04 and 05;
  • uv.
uv venv --python 3.12
source .venv/bin/activate
uv pip install --torch-backend=auto -e '.[dev]'
python tools/environment.py
make check
make test

CPU checks validate references, packaging policy, and Python syntax. They do not compile or execute GPU kernels. Run the GPU suite on the L4:

make gpu-test
make benchmark

Run one project

Each project contains a focused README, correctness suite, and benchmark:

pytest -q projects/01_rmsnorm/test_correctness.py
python projects/01_rmsnorm/benchmark.py

Follow docs/l4_runbook.md for the complete order, including CuTe DSL, Nsight, and result capture.

Run Project 08 in its own environment

vLLM controls its own PyTorch, Triton, and CUDA compatibility matrix, so use a separate environment:

cd projects/08_vllm_plugin
uv venv --python 3.12 --seed .venv
source .venv/bin/activate
uv pip install --torch-backend=auto -e . pytest
VLLM_PLUGINS=llm_serving_swiglu pytest -q test_operator.py
VLLM_PLUGINS=llm_serving_swiglu \
  python benchmark_operator.py > ../../results/08_vllm_operator.json

The plugin uses vLLM's general-plugin entry point, replaces SiluAndMul through CustomOp.register_oot, and exposes the Triton kernel through torch.library.triton_op. It does not fork vLLM or compile a native extension.

Optional L4 execution with Modal

uv pip install modal
modal setup
modal run modal_app.py --task test
modal run modal_app.py --task benchmark
modal run modal_app.py --task vllm

Inspect PTX, SASS, and profiler evidence

python tools/inspect_triton.py swiglu --stage ptx \
  > results/02_swiglu.ptx
bash tools/profile_nsys.sh 04_paged_decode \
  python projects/04_paged_decode/benchmark.py
bash tools/profile_ncu.sh 01_rmsnorm '.*rmsnorm.*' \
  python projects/01_rmsnorm/benchmark.py

Read docs/ptx_sass.md before interpreting generated instructions and docs/benchmarking.md before publishing results.

Evidence standard

A committed result must record:

  • GPU model and compute capability;
  • driver, CUDA, PyTorch, Triton, CUTLASS, and vLLM versions where relevant;
  • input shapes, layouts, dtypes, and numerical tolerances;
  • warm-up and measurement method;
  • median, p20, p80, and minimum latency;
  • exact baseline;
  • profiler evidence for every bottleneck claim.

There are no invented L4 numbers in this repository. Commit measurements only after the matching correctness suite passes in the same environment. The release gate is defined in docs/validation.md.

Scope

This is a focused kernel-engineering laboratory, not a replacement for vLLM, SGLang, TensorRT-LLM, FlashInfer, cuBLAS, or CUTLASS. It demonstrates selected mechanisms with understandable code and finishes with one supported out-of-tree vLLM operator replacement.

See docs/architecture.md for language decisions, docs/references.md for primary sources, and docs/github.md for exact publishing commands.

License

Apache-2.0. Third-party dependencies retain their own licenses.

About

Eight GPU-kernel projects for modern LLM serving

Topics

Resources

Contributing

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages