A lightweight, from-scratch CUDA inference engine for transformer models. No PyTorch, No TensorRT. Just raw CUDA kernels, hand-optimized for performance.
Built as a deep-dive into GPU computing fundamentals — custom GEMM, Flash Attention, INT8 quantization, and a full GPT-2 inference pipeline with Python bindings.
| Getting Started | Installation, build, quick start |
| API Reference | Full Python API docs |
| Architecture | Kernel internals and design decisions |
| Benchmarks | Results, methodology, how to reproduce |
| Contributing | Adding kernels, models, or fixes |
Most inference libraries are black boxes. NanoInfer is the opposite: every kernel is written by hand, every optimization is explicit, and every benchmark tells you exactly why the code is fast.
This project covers the same ground as NVIDIA's TensorRT team works on daily:
- Tiled matrix multiplication with shared memory.
- IO-aware attention (Flash Attention v1/v2).
- INT8 quantization with
dp4ainstructions. - KV-cache with paged allocation.
- Custom GEMM kernel - tiled implementation with share memory, reaching ~75% of cuBLAS peak TFLOPs on A100.
- Online Softmax - single-pass numerically stable softmax using warp-level
__shfl_down_syncreduction. - Flash Attention v1 - IO-aware exact attention, O(N) memory vs naive O(N²)
- INT8 GEMM - quantized matrix multiply using
dp4a, 2–3× throughput vs FP32 - GPT-2 inference pipeline - full autoregressive generation with KV-cache.
- Python bindings — pybind11 interface with a HuggingFace-compatible API
- Benchmark suite - TFLOPS, latency, memory reports vs cuBLAS and llama.cpp
Tested on NVIDIA A100 80GB SXM4. All results averaged over 100 runs after 10 warmup iterations.
| Size | NanoInfer | cuBLAS | % Peak TFLOPS |
|---|---|---|---|
| 512×512 | 8.1 TFLOPS | 10.2 TFLOPS | 79% |
| 1024×1024 | 11.4 TFLOPS | 14.1 TFLOPS | 81% |
| 4096×4096 | 13.8 TFLOPS | 17.2 TFLOPS | 80% |
| Rows × Cols | Naive 3-pass | Online 1-pass | Speedup |
|---|---|---|---|
| 1024 × 1024 | 0.41 ms | 0.14 ms | 2.9× |
| 4096 × 4096 | 6.2 ms | 2.1 ms | 3.0× |
| Batch | Naive mem | Flash Attention mem | Speedup |
|---|---|---|---|
| 1 | 1.07 GB | 48 MB | 22× mem |
| 8 | 8.59 GB | 384 MB | 22× mem |
| 8 | 18.4 ms | 6.1 ms | 3.0× latency |
| Precision | Tokens/sec | Latency/token |
|---|---|---|
| FP32 | 312 | 3.2 ms |
| FP16 | 641 | 1.6 ms |
| INT8 | 1,180 | 0.85 ms |
| Engine | Tokens/sec | Latency/token | vs TensorRT FP16 |
|---|---|---|---|
| PyTorch eager FP32 | 198 | 5.1 ms | 0.17× |
| NanoInfer FP32 | 312 | 3.2 ms | 0.26× |
| NanoInfer FP16 | 641 | 1.6 ms | 0.54× |
| NanoInfer INT8 | 1,180 | 0.85 ms | ~1.0× |
| TensorRT FP16 | 1,190 | 0.84 ms | baseline |
| TensorRT INT8 | 2,340 | 0.43 ms | 1.97× |
NanoInfer INT8 reaches parity with TensorRT FP16 — without using TensorRT. The remaining gap to TensorRT INT8 is primarily due to layer fusion and kernel auto-tuning, which TensorRT performs at compile time.
nanoinfer/
├── nanoinfer/
│ ├── __init__.py
│ └── nanoinfer.pyi # Show autocompletion and type hints for C++ extension
├── src/
│ ├── kernels/
│ │ ├── gemm.cu # Tiled GEMM, TILE_SIZE=16 and 32
| | ├── gemm.h
│ │ ├── softmax.cu # Naive, safe, and online softmax
│ │ ├── softmax.h
│ │ ├── attention.cu # Flash Attention v1/v2
│ │ ├── attention.h
│ │ ├── quantize.cu # INT8 GEMM with dp4a
│ │ └── quantize.h
│ ├── layers/
│ │ ├── linear.cu # Linear layer using custom GEMM
│ │ ├── layernorm.cu # Fused LayerNorm kernel
│ │ └── embedding.cu # Token + positional embedding
│ ├── engine/
│ │ ├── model.cpp # GPT-2 model assembly
│ │ ├── model.h
│ │ ├── kvcache.cpp # Paged KV-cache allocator
│ │ ├── kvcache.h
│ │ ├── sampler.cpp # Top-k / top-p sampling
│ │ └── sampler.h
│ └── bindings/
│ ├── CMakeLists.txt
│ └── python.cpp # pybind11 Python interface
├── benchmarks/
│ ├── bench_gemm.cu # GEMM vs cuBLAS
│ ├── bench_softmax.cu # Softmax variants
│ ├── bench_attention.cu # Flash vs naive attention
│ └── bench_e2e.py # End-to-end GPT-2 benchmark
├── tests/
│ ├── test_gemm.py # Correctness vs PyTorch reference
│ ├── test_softmax.py
│ ├── test_attention.py
│ └── test_quantize.py
├── tools/
│ └── profile.sh # Nsight Systems profiling script
├── CMakeLists.txt
└── setup.py
- CUDA 12.0+
- NVIDIA GPU with compute capability ≥ 8.0 (Ampere or newer)
- CMake 3.20+
- Python 3.9+ (for bindings and tests)
- pybind11
Optional for benchmarking:
- cuBLAS (for comparison baselines)
- NVIDIA Nsight Systems (for profiling)
git clone https://github.com/lecuong1502/NanoInfer.git
cd NanoInfer
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release -DCUDA_ARCHITECTURES=80
make -j8Build Python bindings:
pip install -e .# GEMM benchmark
./build/benchmarks/bench_gemm
# Full benchmark suite
python benchmarks/bench_e2e.py --model gpt2 --precision fp16
# With Nsight profiling
bash tools/profile.sh ./build/benchmarks/bench_attention# Build CUDA tests
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release -DCUDA_ARCHITECTURES=80
make -j$(nproc)
# Chạy từng test
./build/tests/test_gemm
./build/tests/test_softmax
./build/tests/test_attention
./build/tests/test_quantize
# Chạy tất cả qua CTest
ctest --test-dir build -V
# Python tests (sau khi build .so)
pip install -e ".[dev]"
pytest tests/test_kernels.py -vAll correctness tests compare output against PyTorch reference with tolerance atol=1e-4, rtol=1e-4.
from nanoinfer import NanoInfer
model = NanoInfer.from_pretrained("gpt2", precision="int8")
output = model.generate(
prompt="The quick brown fox",
max_tokens=50,
temperature=0.8,
top_p=0.95
)
print(output)Matrix multiplication is tiled across shared memory to exploit data reuse. Each thread block loads a TILE_SIZE × TILE_SIZE tile of A and B into shared memory, computes a partial dot product, then advances to the next tile.
Bank conflicts in As[][] are eliminated by adding 1 element of padding to each row (As[TILE_SIZE][TILE_SIZE + 1]), offsetting each row by one bank and preventing 16-way serialization on column reads.
__shared__ float As[TILE_SIZE][TILE_SIZE + 1]; // +1 eliminates bank conflict
__shared__ float Bs[TILE_SIZE][TILE_SIZE];Standard softmax requires three passes: find max, compute exp(x - max), normalize. Online softmax merges these into one pass by maintaining a running max m and a rescaled denominator d:
m_new = max(m_old, x_i)
d_new = d_old * exp(m_old - m_new) + exp(x_i - m_new)
This is numerically identical to the safe 3-pass version but reads data only once. Final normalization uses warp-level reduction via __shfl_down_sync to avoid shared memory overhead.
Naive attention materializes the full N×N score matrix in HBM, giving O(N²) memory. Flash Attention tiles the computation - loading small blocks of Q, K, V into SRAM, computing partial attention scores, and accumulating results with an online softmax correction factor.
Memory usage drops from O(N²) to O(N), and HBM reads/writes drop from O(N²) to O(N·d), yielding 2–4× end-to-end speedup on long sequences.
One of TensorRT's core optimizations is fusing consecutive operations into a single kernel to eliminate redundant HBM round-trips. NanoInfer implements this manually for the most common pattern in transformer layers: LayerNorm → Linear.
Without fusion, the unfused sequence looks like this:
[HBM] → LayerNorm kernel → [HBM] → Linear kernel → [HBM]
Each arrow is a full read + write of the activation tensor. For a hidden size of 768 at batch=8, seq=512, that's 3 × (8 × 512 × 768 × 4 bytes) = 75 MB of HBM traffic just for one sub-layer — traffic that exists only to pass data between kernels.
The fused kernel collapses both operations into one:
[HBM] → LayerNorm+Linear kernel → [HBM]
The normalized activations live in registers and shared memory — they never touch HBM. HBM traffic drops to 1 read + 1 write, a 2× reduction in memory operations for this pattern.
/ Fused LayerNorm + Linear: one kernel, no intermediate HBM write
__global__ void layernorm_linear_fused(
const float* __restrict__ input, // [B, S, H]
const float* __restrict__ gamma, // [H]
const float* __restrict__ beta, // [H]
const float* __restrict__ weight, // [H_out, H]
float* __restrict__ output, // [B, S, H_out]
int H, int H_out, float eps)
{
// Step 1: Compute LayerNorm — result stays in registers/smem
extern __shared__ float smem[];
int row = blockIdx.x;
float mean = 0.f, var = 0.f;
// Warp-reduce mean
for (int i = threadIdx.x; i < H; i += blockDim.x)
mean += input[row * H + i];
mean = block_reduce_sum(mean) / H;
// Warp-reduce variance
for (int i = threadIdx.x; i < H; i += blockDim.x) {
float diff = input[row * H + i] - mean;
var += diff * diff;
}
var = block_reduce_sum(var) / H;
float inv_std = rsqrtf(var + eps);
// Store normalized values in shared memory (not HBM)
for (int i = threadIdx.x; i < H; i += blockDim.x)
smem[i] = (input[row * H + i] - mean) * inv_std * gamma[i] + beta[i];
__syncthreads();
// Step 2: Linear projection — reads from smem, writes to HBM once
for (int j = threadIdx.x; j < H_out; j += blockDim.x) {
float acc = 0.f;
for (int i = 0; i < H; i++)
acc += smem[i] * weight[j * H + i];
output[row * H_out + j] = acc;
}
}Measured speedup: 1.4× on GPT-2 hidden layers (H=768→3072) at batch=8, seq=512 on A100. The gain scales with sequence length — longer sequences mean more redundant HBM traffic to eliminate.
Importance: Layer fusion is the primary technique behind TensorRT's performance advantage. Every fused kernel in TensorRT's graph optimizer follows the same principle: keep intermediate results on-chip, minimize HBM touches. Writing this from scratch help me understand why TensorRT is fast, not just how to use it.
Post-training quantization maps FP32 weights and activations to INT8 using per-tensor scale and zero-point calibration:
x_q = round(x / scale) + zero_point
The INT8 GEMM kernel uses CUDA's dp4a instruction — a 4-element dot product in a single clock cycle — giving 4× theoretical throughput over FP32 on integer tensor cores.
Use Nsight Systems to see the kernel timeline:
nsys profile --stats=true ./build/benchmarks/bench_attentionKey metrics to watch:
- SM utilization — should be >80% for large GEMMs
- Memory throughput — compare vs peak HBM bandwidth (2 TB/s on A100)
- Shared memory bank conflicts — should be 0 after padding fix
- Tiled GEMM (FP32)
- Online Softmax with warp reduction
- Flash Attention v1
- Fused LayerNorm + Linear kernel
- Flash Attention v2 (work partitioning across warps)
- INT8 GEMM with dp4a
- GPT-2 full pipeline
- LLaMA-3 support
- Multi-GPU (tensor parallel)
- CUTLASS-backed GEMM comparison
- FP8 support (H100+ only, requires compute capability ≥ 9.0)
FP8 (__nv_fp8_e4m3 for weights, __nv_fp8_e5m2 for activations/gradients) is supported from Hopper (H100) onwards via cuda_fp8.h. It offers better dynamic range than INT8 for LLM activations — particularly useful for handling outlier values that cause INT8 to lose accuracy.
The current implementation targets Ampere (A100, compute capability 8.0) where INT8 with dp4a is the right choice. FP8 support is on the roadmap for H100 targets. If you have H100 access, the migration path is: swap int8_t buffers for __nv_fp8_e4m3, replace dp4a with __hmma_fp8 Tensor Core calls, and re-run calibration — the quantization math stays identical.
- Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022)
- Dao et al., FlashAttention-2 (2023)
- NVIDIA, CUDA C++ Best Practices Guide
- NVIDIA, Efficient Matrix Transpose in CUDA
- Leimberger, Online normalizer calculation for softmax (2018)
MIT