llm-speed: A High-Performance CUDA Kernel Library for LLM Inference

llm-speed is a CUDA kernel library optimized specifically for LLM inference scenarios, designed to address performance bottlenecks in large language model inference (such as memory bandwidth, computational efficiency, and memory usage issues). It offers high-performance implementations of FlashAttention, HGEMM (Half-precision Matrix Multiplication), and Tensor Core GEMM, with Python bindings via pybind11, helping developers significantly improve inference performance without sacrificing precision.

The inference process of large language models involves extensive matrix operations (attention computation, feed-forward network computation), which faces multiple challenges when executed on GPUs:

1. Memory Bandwidth Bottleneck: The Transformer attention mechanism frequently accesses the KV Cache, leading to linear growth in memory access as sequence length increases;
2. Computational Efficiency Issue: Standard matrix multiplication cannot fully utilize GPU Tensor Core units, resulting in idle resources;
3. Memory Usage Issue: Activations and intermediate results during inference occupy a large amount of VRAM, limiting batch size and sequence length. These challenges require targeted optimization solutions, which is exactly what llm-speed is designed for.

llm-speed implements three key compute kernels:

FlashAttention Implementation

Using block-wise computation and online softmax technology, it avoids storing the complete attention matrix, reducing memory overhead and improving efficiency. Optimized for CUDA architecture, it uses block strategies to reduce global memory access and fine-grained thread-level parallelism design to maximize GPU compute unit utilization, making it suitable for long-sequence inference.

HGEMM (Half-precision Matrix Multiplication)

It fully utilizes the Tensor Core units of NVIDIA GPUs, adopting Warp-level matrix multiplication primitives (WMMA), block strategies (tuned based on shared memory capacity and Tensor Core dimensions), double buffering, and pipelining techniques to hide memory access latency.

Tensor Core GEMM

It provides a general matrix multiplication interface, supporting multiple data types (e.g., FP16 input with FP32 accumulation) and matrix layouts, allowing users to tune block size parameters based on hardware, balancing flexibility and performance.

llm-speed provides Python bindings via pybind11 for easy integration into the Python ecosystem:

• Concise API design: Users do not need to write CUDA code; kernel functions can be called with just a few lines of code;
• Data compatibility: Handles data type conversion and memory management, supporting mainstream libraries like PyTorch and NumPy;
• Flexible integration: Can be used as an independent library, or embedded into custom inference engines or for researching new attention variants.

The performance improvement of llm-speed comes from multi-level optimizations:

• Algorithm Level: FlashAttention's online computation strategy reduces memory complexity from quadratic to linear, avoiding bandwidth bottlenecks;
• Implementation Level: Tuning thread block partitioning, shared memory usage (maximizing reuse + avoiding bank conflicts), and register allocation (balancing parallelism and pressure) for the CUDA execution model;
• Hardware Level: Fully utilizing Tensor Core capabilities, optimizing data layout and memory access patterns (coalesced global memory access).

Applicable Scenarios:

• Online services: Reduce latency and improve concurrency;
• Batch processing: Increase throughput and shorten task time;
• Edge deployment: Support larger models or longer sequences under limited computing power. Integration Methods:
• PyTorch users: Integrate via custom CUDA extensions;
• TensorRT/other framework users: Adapt kernel implementations;
• Custom inference systems: Directly call the C++ API.

Similar projects in the LLM inference optimization field have their own positioning:

• vLLM: Focuses on service layer optimization and provides a complete inference framework;
• TensorRT-LLM: NVIDIA's official solution for comprehensive model optimization;
• DeepSpeed: Focuses on training optimization, with inference support as a secondary feature. The advantages of llm-speed lie in focus and customizability: It focuses on underlying compute kernel optimization, provides fine-grained control interfaces, and is suitable as a building block for integration into custom systems to meet the needs of deeply customized inference workflows.

Summary: llm-speed helps developers improve LLM inference performance through carefully implemented FlashAttention, HGEMM, and Tensor Core GEMM kernels. Its modular design and Python bindings make it easy to adopt, serving as an important tool for AI application developers pursuing extreme inference performance. Future Directions:

1. Support more attention variants (sliding window, sparse attention);
2. Adapt to new hardware features (NVIDIA Blackwell architecture, AMD GPUs);
3. Add support for low-precision quantization (INT8, INT4).