dropbox/gemlite

Fast low-bit matmul kernels in Triton

What it solves

GemLite addresses the performance bottlenecks of low-bit matrix multiplication in LLM inference. It provides a set of highly optimized kernels that significantly speed up the prefill and decoding stages of model execution compared to default Torch AO kernels, reducing latency and increasing throughput for quantized models.

How it works

GemLite uses Triton to implement several specialized matrix multiplication kernels tailored for different workloads:

  • GEMM: A standard tensor-core based kernel for general matrix multiplication.
  • GEMM Split-K: Optimizes batched LLM decoding (batch sizes 2-32) by splitting the K dimension into multiple jobs that calculate partial sums.
  • GEMV: Designed for small batch sizes (M=1) by splitting activations into 1D chunks.
  • GEMV RevSplit-K: A new algorithm that optimizes batch size 1 decoding by doubling the workload per program to reduce the overhead of loading scales and zeros.

It supports a wide range of precisions, including FP16, BF16, FP8, INT8, and low-bit weights (8, 4, 2, 1-bit), as well as MXFP and NVFP formats for newer hardware like Blackwell.

Who it’s for

GemLite is intended for developers and researchers working on LLM inference optimization, specifically those using quantization to reduce model size and memory bandwidth requirements while maintaining high performance on NVIDIA GPUs.

Highlights

  • Significant Speedups: Delivers up to 7-8x faster prefill and 3-6x faster decoding than default Torch AO kernels.
  • Broad Precision Support: Compatible with FP16, BF16, FP8, INT8, and MXFP/NVFP4.
  • Integration: Works with vLLM (via HQQ or TorchAO), SGLang, and is torch.compile compatible.
  • Autotune Caching: Allows saving and loading optimal kernel configurations to eliminate repeated autotuning time during startup.
  • Flexible Packing: Supports HQQ-style weight packing and channel-wise scaling for weights and activations.

Related

  • Project
    linkedin/Liger-Kernel

    A collection of efficient Triton kernels for LLM training that increases throughput by 20% and reduces memory usage by 60% through kernel fusion and chunking.

  • Project
    deepseek-ai/DeepGEMM

    A high-performance tensor core kernel library for NVIDIA GPUs that provides optimized GEMM, MoE, and MQA kernels to accelerate LLM training and inference.

  • Project
    ServiceNow/Fast-LLM

    A high-performance library for training large language models that uses 3D parallelism and optimized kernels to reduce training time and cost.

  • Dispatch
    vLLM Multi-LoRA Serving for MoE Models

    vLLM version 0.15.0 introduces optimized Multi-LoRA serving for Mixture of Experts (MoE) models, enabling multiple fine-tuned adapters to share a single GPU to reduce idle compute capacity.

  • Project
    NVIDIA/TensorRT-LLM

    An open-source library from NVIDIA that optimizes inference for LLMs and visual generation models using specialized kernels and an efficient runtime to maximize throughput and latency on NVIDIA GPUs.