Pre-compiled cubins for FlashInfer

0.6.7.post3 · active · verified Sun Apr 12

FlashInfer-cubin provides pre-compiled kernel binaries for FlashInfer, supporting a wide range of GPU architectures. This optional package for `flashinfer-python` eliminates JIT compilation and downloading overhead at runtime, leading to faster initialization and enabling offline usage. The FlashInfer project focuses on delivering high-performance LLM GPU kernels for serving and inference, maintaining an active development cycle with frequent nightly builds and regular patch releases.

Warnings

breaking FlashInfer, and by extension `flashinfer-cubin`, has strict compatibility requirements for CUDA and PyTorch versions. Incompatible versions can lead to runtime failures due to mismatches in precompiled kernels (e.g., CUDA 12 vs 13 toolkits) or Python library dependencies.
Fix: Always check the official FlashInfer documentation for supported PyTorch and CUDA versions. Use `flashinfer show-config` to verify your environment post-installation. It is recommended to install `flashinfer-python` and `flashinfer-cubin` together from PyPI to ensure compatible versions.
gotcha `flashinfer-cubin` might not always contain all necessary pre-compiled cubins for every kernel or newer GPU architectures, especially for specific components like TRTLLM FMHA kernels. In such cases, `flashinfer-python` may attempt to download missing cubins at runtime, which can fail in isolated network environments or lead to unexpected JIT compilation.
Fix: For air-gapped or restricted environments, consider using `flashinfer-jit-cache` (if available for your specific CUDA version) or pre-downloading kernels if `flashinfer-cubin` is insufficient. Monitor GitHub issues for updates on comprehensive cubin inclusion or use FlashInfer's source build for full control over compilation.
gotcha The `FLASHINFER_CUBIN_DIR` environment variable, intended to specify a custom path for cubin files, may be ignored when `flashinfer-cubin` is installed via pip. This can lead to issues in containerized or non-root environments where explicit control over artifact paths is required.
Fix: If deploying in environments where `FLASHINFER_CUBIN_DIR` is critical, verify its behavior. It may be necessary to either prevent `flashinfer-cubin` from being installed by pip and manage cubins manually, or use `flashinfer-jit-cache` with a specified index URL for specific CUDA versions.
gotcha While FlashInfer supports a wide range of NVIDIA GPU architectures (SM 7.5 'Turing' and later, up to SM 12.1 'Blackwell'), not all advanced features (e.g., FP8/FP4 operations, certain attention types) are supported across all compute capabilities. Performance can also vary significantly.
Fix: Consult the FlashInfer documentation for a detailed breakdown of feature support per GPU architecture. Test critical workloads on your target hardware to ensure expected performance and functionality.

Install

pip install flashinfer-python flashinfer-cubin Recommended installation alongside core FlashInfer
pip install flashinfer-cubin Direct installation of cubin package

Imports

flashinfer-cubin
```
N/A
```
The `flashinfer-cubin` package is not designed for direct import by end-users. It provides pre-compiled CUDA kernel files (`.cubin` files) that the `flashinfer-python` library loads and utilizes internally for optimized performance. Users interact with the `flashinfer` library directly.

Quickstart

This quickstart demonstrates the usage of the core `flashinfer` library for a single-request decode attention operation. When `flashinfer-cubin` is installed, it transparently provides pre-compiled CUDA kernels to `flashinfer-python`, significantly speeding up operations like this by avoiding runtime compilation overhead.

import torch
import flashinfer

# Example of FlashInfer's single-request decode attention
# (flashinfer-cubin provides the underlying kernels for optimal performance)

kv_len = 2048
num_kv_heads = 32
head_dim = 128

q = torch.randn(1, head_dim, dtype=torch.float16, device='cuda')
k_tensor = torch.randn(kv_len, num_kv_heads, head_dim, dtype=torch.float16, device='cuda')
v_tensor = torch.randn(kv_len, num_kv_heads, head_dim, dtype=torch.float16, device='cuda')

# Prepare FlashInfer attention wrapper
wrapper = flashinfer.to_flashinfer_decode_wrapper(
    kv_len,
    num_kv_heads,
    head_dim,
    0 # page_size, use 0 for single request
)

# Allocate KV cache
k_cache, v_cache = wrapper.alloc_kv_cache(torch.float16, device='cuda')

# Append K/V to cache (simulates historical tokens)
wrapper.begin_forward(k_cache, v_cache)
wrapper.end_forward()

# Perform decode attention
output = flashinfer.single_decode_with_kv_cache(
    q,
    k_cache,
    v_cache,
    wrapper.kv_layout,
    wrapper.kv_indices,
    wrapper.kv_indptr,
    wrapper.last_page_len,
    num_kv_heads,
    num_kv_heads, # num_query_heads == num_kv_heads for single decode
    head_dim,
    True # casual
)

print(output.shape)

view raw JSON →