DOCSEARCHbyAlgolia
Submit Docs

GemLite Documentation

repository·master·Indexed 19 days ago

https://github.com/dropbox/gemlite

A collection of Triton kernels optimized for efficient low-bit matrix multiplication to improve LLM inference prefill and decoding performance. GemLite provides GemLiteLinear for quantized linear layers, helper functions for various quantization schemes (Weight-only, Dynamic, and MXFP/NVFP), and deep integration with vLLM for routing quantized forward paths or performing on-the-fly quantization of fp16/bf16 checkpoints.

Tokens
6.4K
Snippets
18
Records
21
Agent score
17%

What's inside GemLite

  2. Integrate GemLite with vLLM (Overview)

    master

    GemLite provides an integration for vLLM that routes quantized forward paths through GemLite's Triton kernels or enables on-the-fly quantization of fp16/bf16 checkpoints at load time.

    Key Entry Points:

    • enable_gemlite(names=None): Routes pre-quantized checkpoints through GemLite.
    • set_onthefly_quant(...): Quantizes fp16/bf16 checkpoints at load time.
    • patch_vllm(): Idempotent application of both methods via environment variables.
    • register(): vLLM plugin entry point.

    Important: The swap must occur before vLLM resolves the quantization configuration for the model (i.e., before LLM(...) or vllm serve finishes importing the engine).

  3. Understand GemLite Triton Kernels

    master

    GemLite implements several specialized Triton kernels optimized for different workloads and batch sizes. All kernels support 8-, 4-, 2-, and 1-bit weight precision, and can handle float16, bfloat16, and int8/fp8 activations.

    Available Kernels:

    • GEMM: A General Matrix Multiplication kernel. Because it utilizes tensor cores, activations must be padded with zeros along the batch dimension to at least 16 rows. It supports float32 and float16 accumulation for fp16 inputs, and float32 accumulation for bfloat16 inputs.
    • GEMM Split-K: An extension of the GEMM kernel that splits the K dimension into multiple jobs to calculate partial sums via atomic addition. This is optimized for batched LLM decoding with batch sizes between 2 and 32.
    • GEMV: A General Matrix-Vector multiplication kernel that splits activations into 1D chunks. It is primarily intended for small batch sizes where M == 1.
    • GEMV RevSplit-K: A specialized algorithm for GEMV that doubles the workload per Triton program. This reduces the frequency of loading scales/zeros and lowers thread requirements, providing the best performance for batch size = 1 decoding.

  4. Configure GGUF models in vLLM with GemLite

    master

    When serving GGUF checkpoints via vllm serve, note the following requirements:

    1. Dtype Requirement: vLLM rejects bfloat16 for GGUF. You must explicitly pass --dtype float16.
    2. Config Path: If the GGUF repository does not contain a config.json (common in unsloth/*-GGUF repositories), you must provide the path to the original unquantized Hugging Face repository using the --hf-config-path <hf-repo> flag.
    vllm serve <model_path> --dtype float16 --hf-config-path <unquantized-repo>

  5. Use GemLite in Interactive Python / Offline LLM

    master

    To use GemLite with an offline LLM instance, you must import and call enable_gemlite() before constructing the LLM object. You can optionally pass a list of specific scheme names to restrict the supported quantization types.

    from gemlite.vllm import enable_gemlite
enable_gemlite()                          # enables all supported schemes

# Or restrict to a subset:
# enable_gemlite(["A8W8_FP8_DYNAMIC", "A16W4_HQQ_INT"])

from vllm import LLM, SamplingParams
llm = LLM(model="Qwen/Qwen3-4B-Instruct-2507-FP8", dtype="bfloat16")
out = llm.generate(["What is 2+2?"], SamplingParams(max_tokens=16))
print(out[0].outputs[0].text)

  6. Use GemLite helper functions for quantization

    master

    GemLite provides high-level helper functions to quickly apply different quantization schemes to existing layers or entire models. These helpers follow the AxWy pattern where x is activation precision and y is weight precision.

    Common Quantization Types:

    • Weight-only: A16W8_INT8, A16W4_HQQ_INT, etc.
    • Dynamic Quantization (Activation + Weight): A8W8_INT8_dynamic, A8W4_HQQ_INT_dynamic, etc.
    • MXFP/NVFP Quantization: A16W4_MXFP, A4W4_NVFP_dynamic, etc.

    Patching a Model

    You can patch an entire model (even from CPU) using patch_model.

    from gemlite.helper import *
import torch

device, dtype = 'cuda:0', torch.float16

# Example: Weight-only quantization
gemlite_linear = A16W8_INT8(device=device, dtype=dtype).from_linear(layer)

# Example: 8-bit activation dynamic quant (channelwise)
gemlite_linear = A8W8_INT8_dynamic(device=device, dtype=dtype).from_linear(layer)

# Example: Patching a whole model
patch_model(model, device=device, processor=A8W8_INT8_dynamic())

  7. Use GemLite with `vllm serve` (OpenAI-compatible server)

    master

    GemLite registers as a vllm.general_plugins entry point, so vllm serve will automatically discover and apply the patch at startup if you set the VLLM_GEMLITE_ENABLE environment variable.

    To restrict the supported schemes, use VLLM_GEMLITE_ENABLE_LIST with a comma-separated list of names.

    Bootstrap Fallback: If the plugin is not discovered (e.g., in certain editable installs), you can use a wrapper to pre-import gemlite.vllm before running the server.

    # Standard usage
export VLLM_GEMLITE_ENABLE=1
vllm serve Qwen/Qwen3-4B-Instruct-2507-FP8 --dtype bfloat16 --port 8000

# Restricting schemes
export VLLM_GEMLITE_ENABLE_LIST=A8W8_FP8_DYNAMIC,A16W4_HQQ_INT

# Bootstrap fallback if plugin discovery fails
export VLLM_GEMLITE_ENABLE=1
python3 -c "
import sys, gemlite.vllm            # triggers patch_vllm() via env var
sys.argv = ['vllm', 'serve', 'Qwen/Qwen3-4B-Instruct-2507-FP8', 
            '--dtype', 'bfloat16', '--port', '8000']
from vllm.entrypoints.cli.main import main
main()
"

  8. Cache and load autotune configurations

    master

    Triton autotuning can be time-consuming. GemLite allows you to save and load the best autotuning configurations to a JSON file to speed up subsequent startups.

    Important: Use one JSON cache file per GPU model.

    To accelerate the process, you can use the warmup function to pre-run specific shapes and batch sizes before caching.

    import gemlite
from gemlite.helper import warmup

# 1. Reset, Warmup, and Cache
gemlite.reset_config()
# Set autotune mode: "fast" or "max"
# gemlite.set_autotune("max")

# Warmup with specific shapes and batch sizes
warmup(gemlite.A8W8_INT8_dynamic(), shapes=[(4096, 4096), (2048, 4096)], batch_sizes=[1, 8, 64, 128])

# Cache the new config
gemlite.cache_config('gemlite_config.json')

# 2. Later, in a different session, load the config to skip autotuning
gemlite.load_config('gemlite_config.json')

  9. Install GemLite

    master

    You can install GemLite using pip.

    To install the latest recommended version (from GitHub):

    pip install git+https://github.com/dropbox/gemlite/

    To install the latest stable version:

    pip install gemlite
    pip install git+https://github.com/dropbox/gemlite/

  10. Configure ptxas for Blackwell GPUs

    master

    When running benchmarks or using GemLite on Blackwell architecture, ensure you are using the CUDA 13 ptxas compiler by setting the TRITON_PTXAS_BLACKWELL_PATH environment variable.

    export TRITON_PTXAS_BLACKWELL_PATH=/usr/local/cuda-13.0/bin/ptxas

  11. Integrate GemLite with vLLM

    master

    GemLite can be used with vLLM via the set_vllm_onthefly_hqq_quant utility. This allows you to apply quantization on-the-fly during model loading.

    Supported modes include:

    • INT/FP formats: int8_weightonly, int4_weightonly, int8_dynamic, fp8_dynamic.
    • MXFP formats: mxfp8_dynamic (with or without post-scale), mxfp4_weightonly, mxfp4_dynamic, nvfp4_dynamic.
    from hqq.utils.vllm import set_vllm_onthefly_hqq_quant
from vllm import LLM
import torch

skip_modules = ['lm_head', 'visual', 'vision']

# Example: Apply A8W8 - INT8 x INT8 dynamic quantization
set_vllm_onthefly_hqq_quant(
    weight_bits=8, 
    group_size=None, 
    quant_mode='int8_dynamic', 
    skip_modules=skip_modules
)

# Load the vLLM model
llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct", 
    max_model_len=4096, 
    gpu_memory_utilization=0.80, 
    dtype=torch.float16
)

  12. Perform On-the-fly Quantization

    master

    On-the-fly quantization converts fp16 / bf16 checkpoints to a quantized format at load time. This is a no-op for already-quantized models (e.g., FP8, AWQ, GPTQ).

    Recommended Method (Environment Variables): Use environment variables for vllm serve or multi-process LLM instances to ensure the configuration propagates to worker subprocesses.

    Programmatic Method (Offline LLM only): Use set_onthefly_quant(...) for offline LLM instances where VLLM_USE_V1=0 or for code paths that do not fork workers. This method will not propagate to workers under the v1 engine.

    Note: int4_weightonly requires pip install hqq.

    # Recommended for vllm serve
export VLLM_GEMLITE_ONTHEFLY_QUANT=A16W8_INT8
export VLLM_GEMLITE_SKIP_MODULES=lm_head,visual,vision
vllm serve Qwen/Qwen3-4B --dtype bfloat16 --port 8000
    # Programmatic (Offline LLM only, VLLM_USE_V1=0)
from gemlite.vllm import set_onthefly_quant
set_onthefly_quant(
    weight_bits=8, group_size=None, quant_mode="int8_weightonly",
    skip_modules=["lm_head", "vision", "visual"],
)

from vllm import LLM
llm = LLM(model="Qwen/Qwen3-4B", dtype="bfloat16")