GemLite is a collection of Triton kernels designed for efficient low-bit matrix multiplication, emphasizing simplicity and reusability. It provides a practical solution for achieving significant performance gains, delivering up to 7-8x faster prefill and 3-6x faster decoding compared to default Torch AO kernels. For more detailed benchmarks, check the Performance section.
GemLite strikes the perfect balance between flexibility and performance, allowing users to easily use and modify the codebase to develop high-performance kernels optimized for their specific hardware. We have included multiple versions of the kernels to maximize performance across different matrix shapes.
The project started with CUDA kernels, but we have switched to Triton for enhanced flexibility. For the old CUDA version, please refer to this branch.
| End-to-end Performance (Llama3 8-bit) | Matmul Performance (A16W8) |
|---|---|
 |  |
Extensive performance results across different bitwidths, batch sizes, and devices are available in the Performance section below.
- Improved performance with a focus on
sm_120. - GemLite now supports MXFP4/NVFP4 for Blackwell.
- GemLite now supports vLLM V1 and is
torch.compilecompatible. - GemLite now supports
bfloat16. - GemLite is now available in vLLM via the HQQ library.
- GemLite is now integrated with TorchAO/SGLang for 4-bit quantization. Check out the blog post.
- Major performance improvements, especially on the A100 and H100.
- Flexible bit packing: use 8-bit packing for improved batched performance on the A100 and H100 with packed data.
- Autotune caching: save and load the best autotune configs across all kernels with a single line of code.
- Helper functions: make it easier to get started, especially for dynamic quantization.
- New GEMV RevSplit-K algorithm: outperforms GEMM Split-K and GEMV for batch size = 1 with packed data.
- Channel-wise scaling: added support for channel-wise scaling for weights, activations, or both.
- Precision support: includes FP16 × Wn, FP8 × FP8, FP8 × Wn, INT8 × INT8, INT8 × Wn, and MXFPn × MXFPn.
torch.compile()support.
pip install git+https://github.com/dropbox/gemlite/pip install gemliteimportgemlitefromgemliteimportDType, GemLiteLineargemlite_linear=GemLiteLinear(
W_nbits, # weight quantization bit width. supported: [8, 4, 2, 1]group_size=group_size, # any group_size divisible by 32 - enable autotune for group_size < 128 (!)in_features=in_features, # input sizeout_features=out_features, # output sizeinput_dtype=DType.FP16, # FP16, BF16, FP8, INT8output_dtype=DType.FP16, # FP16, BF16, FP32, FP8, INT32scaled_activations=False, # whether the activations are scaled
)
# Packing: we follow the HQQ format (W_q - zeros) * scales ~= W# https://github.com/dropbox/hqq/gemlite_linear.pack(W_q, scales, zeros, bias)
# Forwardout=gemlite_linear(x)Settings
# Set packing width for packed data - recommended to leave this at the default valuegemlite.set_packing_bitwidth(int)
# Set the accumulation dtype - this is configured automatically.# On consumer GPUs, fp16 is used by default.gemlite.set_acc_dtype(DType)
# Enable TMA - disabled by default. Only supported for MXFP/NVFP kernelsgemlite.enable_tma(True)
# Enable Triton warp specialization on the k-loop - disabled by defaultgemlite.enable_warp_specialize(True)
# Enable/disable native bfp16 atomic addition - recommended to leave this at the default valuegemlite.set_native_atomic_bfp16(True)
# Enable optimized PTX FP4 packing in the MXFP4/NVFP4 activation quant kernel - requires CUDA 13 ptxasgemlite.set_ptx_fp4_pack(True)
# Experimental fast mode for NVFP4, using a static meta scale for activationsgemlite.set_fast_nvfp4(True)
# Use CUDA graphs for autotuning - this will slow down autotuninggemlite.enable_cudagraph_autotune(True)
# Enable activation quantization only from a specified batch size onward.# Smaller batch sizes will use weight-only quantization.gemlite.enable_activation_scaling(int)
# Enable kernel caching: makes some GEMV kernels faster,# but might break with some torch.compile settingsgemlite.set_kernel_caching(True)Additionally, we offer helper functions that operate as follows:
fromgemlite.helperimport*device, dtype='cuda:0', torch.float16# AxWy: x = activation precision in bits, y = weight precision in bits.# Weight-onlygemlite_linear=A16W8_INT8(device=device, dtype=dtype).from_linear(layer)
gemlite_linear=A16W8_FP8(device=device, dtype=dtype).from_linear(layer)
gemlite_linear=A16W8_HQQ_INT(device=device, dtype=dtype).from_hqqlinear(hqq_layer)
gemlite_linear=A16W4_HQQ_INT(device=device, dtype=dtype).from_hqqlinear(hqq_layer)
gemlite_linear=A16W2_HQQ_INT(device=device, dtype=dtype).from_hqqlinear(hqq_layer)
gemlite_linear=A16W158_INT(device=device, dtype=dtype).from_bitlinear(bitlinear_layer)
# 8-bit activation dynamic quant (channelwise; pass block_quant=True for DeepSeek-style 128x128 block quant)gemlite_linear=A8W8_INT8_dynamic(device=device, dtype=dtype).from_linear(layer)
gemlite_linear=A8W8_FP8_dynamic(device=device, dtype=dtype).from_linear(layer)
gemlite_linear=A8W4_HQQ_INT_dynamic(device=device, dtype=dtype).from_hqqlinear(hqq_layer)
gemlite_linear=A8W158_INT_dynamic(device=device, dtype=dtype).from_bitlinear(bitlinear_layer)
# MXFP weight-onlygemlite_linear=A16W8_MXFP(device=device, dtype=dtype).from_linear(layer)
gemlite_linear=A16W4_MXFP(device=device, dtype=dtype).from_linear(layer)
# MXFP/NVFP dynamic quant - if post_scale=True, uses channel-wise activation quantization.# Support depends on Triton's ability to support native MXFP/NVFP MMA.gemlite_linear=A8W8_MXFP_dynamic(device=device, dtype=dtype, post_scale=False).from_linear(layer)
gemlite_linear=A8W8_MXFP_dynamic(device=device, dtype=dtype, post_scale=True).from_linear(layer)
gemlite_linear=A8W4_MXFP_dynamic(device=device, dtype=dtype, post_scale=False).from_linear(layer)
gemlite_linear=A8W4_MXFP_dynamic(device=device, dtype=dtype, post_scale=True).from_linear(layer)
gemlite_linear=A4W4_MXFP_dynamic(device=device, dtype=dtype).from_linear(layer)
gemlite_linear=A4W4_NVFP_dynamic(device=device, dtype=dtype).from_linear(layer)You can also patch the whole model, even from CPU, as follows:
fromgemlite.helperimport*patch_model(model, device=device, processor=A8W8_INT8_dynamic())Triton autotuning can be time-consuming. To accelerate this process, we provide tools to automatically cache and load the optimal autotuning configurations for all kernels:
importgemlitegemlite.reset_config() # resets cached configs for all kernelsgemlite.cache_config('gemlite_config.json') # cachegemlite.load_config('gemlite_config.json') # loadEnsure that you use one JSON cache file per GPU model. When the cache is loaded, the kernels will skip autotuning, leading to faster startup times.
You can warm up specific shapes using the following helper function:
importgemlite# Ignore pre-loaded configs if you want to start from scratch (optional)# gemlite.reset_config()# Set autotune mode: fast or max# gemlite.set_autotune("max")# Autotune with the default batch sizeswarmup(A8W8_INT8_dynamic(), shapes=[(4096, 4096), (2048, 4096)])
# You can specify batch sizes toowarmup(A8W8_INT8_dynamic(), shapes=[(4096, 4096), (2048, 4096)], batch_sizes=[1, 8, 64, 128])
# If you want to specify the group size for HQQ-style quantizationwarmup(A16W4_HQQ_INT(), shapes=[(4096, 4096), (2048, 4096)], group_size=64)
# Cache your new configgemlite.cache_config('new_config.json')You can use GemLite with vLLM via TorchAO or HQQ as follows:
fromhqq.utils.vllmimportset_vllm_onthefly_hqq_quantskip_modules= ['lm_head', 'visual', 'vision']
# Select one of the following modes:# INT/FP formatset_vllm_onthefly_hqq_quant(weight_bits=8, group_size=None, quant_mode='int8_weightonly', skip_modules=skip_modules) # A16W8 - INT8 weight-onlyset_vllm_onthefly_hqq_quant(weight_bits=4, group_size=128, quant_mode='int4_weightonly', skip_modules=skip_modules) # A16W4 - HQQ weight-onlyset_vllm_onthefly_hqq_quant(weight_bits=8, quant_mode='int8_dynamic', skip_modules=skip_modules) # A8W8 - INT8 x INT8 dynamicset_vllm_onthefly_hqq_quant(weight_bits=8, quant_mode='fp8_dynamic', skip_modules=skip_modules) # A8W8 - FP8 x FP8 dynamic# MXFP formatset_vllm_onthefly_hqq_quant(weight_bits=8, group_size=None, quant_mode='mxfp8_dynamic', skip_modules=skip_modules) # A8W8 - MXFP8 x MXFP8 - post_scale=Trueset_vllm_onthefly_hqq_quant(weight_bits=8, group_size=32, quant_mode='mxfp8_dynamic', skip_modules=skip_modules) # A8W8 - MXFP8 x MXFP8 - post_scale=Falseset_vllm_onthefly_hqq_quant(weight_bits=4, quant_mode='mxfp4_weightonly', skip_modules=skip_modules) # A16W4 - MXFP4 weight-onlyset_vllm_onthefly_hqq_quant(weight_bits=4, quant_mode='mxfp8_dynamic', skip_modules=skip_modules) # A8W4 - MXFP8 x MXFP4 dynamicset_vllm_onthefly_hqq_quant(weight_bits=4, quant_mode='mxfp4_dynamic', skip_modules=skip_modules) # A4W4 - MXFP4 x MXFP4 dynamicset_vllm_onthefly_hqq_quant(weight_bits=4, quant_mode='nvfp4_dynamic', skip_modules=skip_modules) # A4W4 - NVFP4 x NVFP4 dynamic# Load your vLLM modelllm=LLM(model="meta-llama/Llama-3.1-8B-Instruct", max_model_len=4096, gpu_memory_utilization=0.80, dtype=torch.float16)We implement various versions of Triton kernels:
GEMM: This GEMM kernel is implemented similarly to GPTQ-triton. Since it uses tensor cores, activations must be padded with zeros along the batch dimension to at least 16 rows. It supports both float32 and float16 accumulation for fp16 inputs, but only float32 accumulation for bfloat16.
GEMM Split-K: This Split-K GEMM kernel is implemented similarly to the GPTQ Split-K version. We build on the GEMM version above and add another grid dimension that splits the K dimension into multiple jobs that calculate partial sums, which are atomically added and then stored. Split-K performs particularly well for batched LLM decoding (batch sizes between 2 and 32).
GEMV: This GEMV kernel splits activations into 1D chunks, performs the dot product using
tl.sum, and accumulates via atomic addition. It is primarily intended for use with small batch sizes (M == 1).GEMV RevSplit-K: This algorithm, newly introduced in GemLite, operates in contrast to the GEMM Split-K approach, but within a GEMV context. By doubling the workload per Triton program launched in the GEMV kernel, it reduces the frequency of loading scales/zeros and lowers the number of threads needed. As a result, this method delivers the best performance for batch size = 1 decoding.
All kernels are flexible, supporting 8-, 4-, 2-, and 1-bit weight precision, as well as float16, bfloat16, and int8/fp8 activations.
Make sure to use CUDA 13 ptxas for Blackwell:
export TRITON_PTXAS_BLACKWELL_PATH=/usr/local/cuda-13.0/bin/ptxas| Batch Size | FP16 | GemLite FP8 | RedHat FP8 | GemLite MXFP4 | GemLite NVFP4 | RedHat NVFP4 |
|---|---|---|---|---|---|---|
| 1 | 15.4 ms | 10.3 ms | 9.9 ms | 7.3 ms | 8.3 ms | 10.5 ms |
| 8 | 32.4 ms | 23.3 ms | 23.6 ms | 20.5 ms | 20.4 ms | 22.2 ms |
| 16 | 36.9 ms | 29.8 ms | 29.2 ms | 27.1 ms | 27.7 ms | 28.5 ms |
| 32 | 56.7 ms | 48.1 ms | 48.0 ms | 42.6 ms | 43.9 ms | 44.4 ms |
| 64 | 104.0 ms | 86.6 ms | 93.7 ms | 87.6 ms | 87.1 ms | 75.3 ms |
| 128 | 198.4 ms | 164.9 ms | 153.5 ms | 151.4 ms | 143.1 ms | 141.2 ms |
| Batch Size | FP16 | GemLite FP8 | RedHat FP8 | GemLite MXFP4 | GemLite NVFP4 | RedHat NVFP4 |
|---|---|---|---|---|---|---|
| 1 | 11.75s | 6.75s | 8.00s | 4.84s | 5.94s | 8.19s |
| 8 | 11.92s | 7.41s | 7.78s | 5.19s | 6.32s | 8.40s |
| 16 | 12.44s | 7.89s | 8.23s | 5.66s | 6.77s | 8.76s |
| 32 | 13.83s | 8.74s | 9.53s | 6.68s | 7.71s | 9.38s |
| 64 | 15.69s | 10.41s | 11.08s | 8.96s | 9.24s | 10.62s |
| 128 | 19.32s | 14.71s | 14.65s | 12.39s | 13.34s | 13.81s |
Check out the talk by lead author Dr. Hicham Badri about GemLite at GPU MODE. You can also find the slides here.
Please note that GemLite is under active development, and the content discussed in the talk may evolve as the library continues to improve.
Contributions are always welcome. Please feel free to raise issues, submit pull requests, or start a discussion.
If you're looking to integrate GemLite with major inference and AI libraries, we'd love to hear from you!