A High Performance LLM Inference Engine
A custom-built inference engine designed to minimize memory bandwidth bottlenecks in large language model deployment. Implemented entirely in C++ and CUDA, bypassing Python interpreter overhead to achieve near-theoretical memory bandwidth utilization.
THE PROBLEM
In modern Large Language Model (LLM) deployment, computation is rarely the primary bottleneck during sequence generation. The principal constraint is memory bandwidth—specifically, loading massive parameter matrices from High Bandwidth Memory (HBM) into the Streaming Multiprocessors (SMs) on the GPU for every single token generated.
Standard Python-based frameworks add significant kernel dispatch latency and host-to-device memory allocation overhead, masking true hardware capabilities. The objective of this project was to eliminate intermediate abstractions and address the memory wall directly at the CUDA kernel level.
TECHNICAL ARCHITECTURE
Custom CUDA Kernels
Developed fused matrix-vector multiplication kernels to maximize register reuse and minimize L2 cache misses. Implemented tiled memory access patterns specifically tuned for NVIDIA Ampere and Hopper GPU architectures.
Paginated KV Cache
Implemented a paginated Key-Value cache system in native C++, enabling zero-copy memory management, non-contiguous page allocation, and fragmentation-free batching across variable-length sequences.
// Fused FlashAttention-v2 style CUDA kernel implementation
__global__ void flash_attention_v2_kernel(
const half* __restrict__ Q,
const half* __restrict__ K,
const half* __restrict__ V,
half* __restrict__ O,
const int seq_len,
const int head_dim
) {
// Shared memory allocations for tiling
extern __shared__ half s_mem[];
half* s_Q = s_mem;
half* s_K = &s_mem[BLOCK_SIZE * head_dim];
// Grid-stride loop structure for coalesced memory loading
int tid = threadIdx.x;
int bid = blockIdx.x;
// Load Q tile into shared memory (SRAM)
#pragma unroll
for (int i = tid; i < BLOCK_SIZE * head_dim; i += blockDim.x) {
s_Q[i] = Q[bid * BLOCK_SIZE * head_dim + i];
}
__syncthreads();
// Perform online softmax and gemm reduction
// ... (Tiled computation omitted for brevity)
}WHAT I LEARNED
Building this engine from raw primitives reinforced that high performance in modern AI is predominantly a memory management problem, not a compute capability issue. Navigating the CUDA memory hierarchy (Global Memory → L2 Cache → Shared Memory SRAM → Warp Registers) effectively makes the difference between achieving 25% versus 90%+ theoretical hardware utilization.