Skip to content

CLAUDE.md

Latest commit

 

History

History
938 lines (733 loc) · 69.4 KB

CLAUDE.md

File metadata and controls

938 lines (733 loc) · 69.4 KB
Raw

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

FlashInfer is a GPU kernel library for LLM serving that uses JIT (Just-In-Time) compilation by default. This means kernel code changes are automatically picked up without reinstalling the package - extremely convenient for development.

Quick Reference

Task Command
Install for development pip install --no-build-isolation -e . -v
Initialize submodules git submodule update --init --recursive
Install CUPTI for benchmarking pip install -U cupti-python
Run regular tests pytest tests/
Run full parameter matrices pytest tests/ --full
Run specific test pytest tests/path/test_file.py::test_function
Run multi-GPU test mpirun -np 4 pytest tests/comm/test_allreduce_unified_api.py
Run benchmark python benchmarks/flashinfer_benchmark.py --routine <name> <flags>
Run linting pre-commit run -a
Dump environment report (bug reports) python -m flashinfer.collect_env (or flashinfer collect-env [--json])
Install pre-commit hooks pre-commit install
Clear JIT cache rm -rf ~/.cache/flashinfer/
Inspect devcontainer compiler cache sccache --show-stats
Enable API logging (basic) export FLASHINFER_LOGLEVEL=1
Enable API logging (detailed) export FLASHINFER_LOGLEVEL=3
Enable API logging (with stats) export FLASHINFER_LOGLEVEL=5
Set API log destination export FLASHINFER_LOGDEST=mylog.txt
Enable verbose JIT logging export FLASHINFER_JIT_VERBOSE=1
Enable debug build export FLASHINFER_JIT_DEBUG=1
Set target architectures export FLASHINFER_CUDA_ARCH_LIST="8.0 9.0a"
Set parallel compilation export FLASHINFER_NVCC_THREADS=4
Limit parallel ninja jobs export MAX_JOBS=4
Enable GDN native short-T path export FLASHINFER_GDN_WY_NATIVE_T=1
Enable GDN strided QKV path export FLASHINFER_GDN_WY_STRIDED_QKV=1
Enable GDN native A/B tensors export FLASHINFER_GDN_WY_NATIVE_AB=1
Let backend="auto" pick experimental backends export FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1
Override ragged-prefill auto backend order (Blackwell) export FLASHINFER_RAGGED_AUTO_BACKEND_ORDER=cutlass,cudnn
Override CuTe-DSL prefill scheduling export FLASHINFER_CUTE_PREFILL_PERSISTENT=0 (non-persistent) or 1 (persistent)
Skip MoE EP CuTe-DSL import/version guard export FLASHINFER_MOE_EP_SKIP_DSL_CHECK=1
Override MoE EP knob-cache path export FLASHINFER_MOE_EP_KNOB_CACHE=/path/to/knobs.json
Disable MoE EP fused staging kernel export FLASHINFER_MEGA_FUSED_STAGE=0
Enable MoE EP EXPERT_MAJOR padding-row exclusion (opt-in; shape-dependent) export FLASHINFER_MOE_EP_EXCLUDE_PADDING_ROWS=1
Enable distribution-aware MoE autotune and kernel dispatch (experimental; TRT-LLM MoE only) export FLASHINFER_DIST_AWARE_AUTOTUNE=1
Prims-TS debug checks / task-manager verify export FLASHINFER_PRIMS_TS_DEBUG_CHECKS=1
Override Prims-TS CuTe-DSL compile options export FLASHINFER_PRIMS_TS_COMPILE_OPTIONS="--opt-level 2"
Configure NVFP4 4over6 quantization (preferred over the env vars) nvfp4_quantize(..., nvfp4_4over6=NVFP44Over6Config(e4m3_max=256))

The minimum Python version used by CI and build tooling is defined in .python-version. CI Docker images use a stable Conda environment name and derive their interpreter and site-packages paths from that file. When raising the minimum, also update project.requires-python and tool.mypy.python_version in pyproject.toml.

Quick Start for Development

Installation

git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
cd flashinfer
pip install --no-build-isolation -e . -v

Important: The --recursive flag is required to initialize submodules in 3rdparty/ (cutlass, spdlog).

If you forgot --recursive when cloning:

git submodule update --init --recursive

That's it! You can now:

  • Run all benchmarks and unit tests
  • Modify kernel source code in include/ without reinstalling
  • Changes are JIT-compiled on next use

The --no-build-isolation flag prevents pip from pulling incompatible PyTorch/CUDA versions from PyPI.

How JIT Compilation Works

When you call a FlashInfer API:

  1. First call: Generates specialized CUDA code based on parameters (dtype, head_dim, etc.), compiles it with ninja, caches the .so file
  2. Subsequent calls: Uses cached compiled module
  3. After kernel changes: Automatically detects changes and recompiles

No manual rebuild step needed - just edit .cuh files and run your code again.

Pre-compiled Packages (Optional)

FlashInfer provides optional pre-compiled packages for users who want faster initialization:

  • flashinfer-jit-cache: Pre-built kernel cache
  • flashinfer-cubin: Pre-compiled kernel binaries

For development, you typically DON'T need these. JIT compilation is fast enough and gives you live code reload.

Compiler Caching in Devcontainers

The devcontainers run every nvcc and host C++ compile through sccache by setting FLASHINFER_NVCC_LAUNCHER and FLASHINFER_CXX_LAUNCHER. Cached objects live in a per-devcontainer Docker volume mounted at ~/.cache/sccache, so they outlive a container rebuild — which wipes ~/.cache/flashinfer and would otherwise force a full recompile.

sccache --show-stats                  # hit rate, cache size, uncacheable reasons
export FLASHINFER_NVCC_LAUNCHER=""    # disable for the current shell

Caveats:

  • Toggling a launcher changes the compile commands in build.ninja, so the first build after switching it on or off recompiles objects that already exist.
  • Link steps, cubin embedding, and CuTe-DSL kernels never invoke nvcc, so they are not cached.
  • The cache is per-devcontainer by design: sccache's local storage supports a single server at a time, so two containers must not share one cache directory.

cuDNN linear attention

Use nvidia-cudnn-frontend>=1.31.0; KDA and GDN normalization use additive 1e-6 without an epsilon argument. Auto selection also requires CuTe DSL 4.7+. The bounded SM100 policies live in kda.py and gdn_prefill.py. Validate both output and final state before changing a policy. Explicit state-pool calls require unique, in-range slots; auto keeps pools on their native route.

Use benchmarks/bench_cudnn_linear_attention.py for public-call comparisons. Its raw-gate GDN baseline includes Torch gate transforms; it is not a vLLM fused-producer baseline. Keep host CPU, completed wall and reset-inclusive CUDA graph times separate.

Testing

Run the regular test suite:

pytest tests/

Run full parameter matrices (used by nightly testing):

pytest tests/ --full

Run specific test file:

pytest tests/attention/test_hopper.py

Run specific test function:

pytest tests/attention/test_hopper.py::test_single_prefill

Skipping Tests Based on CUDA Architecture

Use flashinfer.utils functions to skip tests on unsupported GPU architectures:

Available check functions:

  • get_compute_capability(device) - Returns (major, minor) tuple
  • is_sm90a_supported() - Hopper (requires CUDA 12.3+)
  • is_sm100a_supported() - Blackwell (requires CUDA 12.8+)
  • is_sm100f_supported() - Blackwell feature-set (requires CUDA 12.9+)
  • is_sm110a_supported(), is_sm120a_supported(), is_sm120f_supported(), is_sm121a_supported()

APIs decorated with @backend_requirement also provide:

  • api_name.is_compute_capability_supported(cc) - e.g., mm_fp4.is_compute_capability_supported(100)
  • api_name.is_backend_supported("backend") - e.g., mm_fp4.is_backend_supported("cudnn")
  • api_name.is_backend_supported("backend", cc) - e.g., mm_fp4.is_backend_supported("cudnn", 80)

Example:

from flashinfer.utils import is_sm90a_supported

def test_hopper_attention():
    if not is_sm90a_supported(torch.device("cuda")):
        pytest.skip("Requires SM90a")
    # Test code...

Common requirements:

Feature Min SM Check Function
FlashAttention-3 SM90a is_sm90a_supported()
MLA Attention SM100a is_sm100a_supported()
FP8 GEMM SM89+ get_compute_capability()[0] >= 9

Note: tests/conftest.py auto-skips tests that trigger OOM, but tests should be written to avoid OOM by using appropriate problem sizes.

Benchmarking

For SM90 NVFP4 fused/split and MXFP4 fused MegaMoE, use benchmarks/bench_megamoe_fp4_sm90.py for timing and tests/moe/test_megamoe_fp4_sm90.py for correctness checks.

FlashInfer provides a unified benchmarking framework in benchmarks/flashinfer_benchmark.py.

Key features:

  • Supports attention, GEMM, and MOE kernels
  • Multiple backends: FlashAttention2/3, cuDNN, CUTLASS, TensorRT-LLM, cuBLAS
  • CUPTI timing (recommended): Hardware-level profiling for accurate GPU kernel time
    • Automatically falls back to CUDA events if CUPTI unavailable
    • Install: pip install -U cupti-python (requires CUDA 13+)
  • Batch testing, reference checking, CSV output

Quick example:

python benchmarks/flashinfer_benchmark.py \
    --routine BatchDecodeWithPagedKVCacheWrapper \
    --backends fa2 cudnn \
    --batch_size 32 --s_kv 2048 \
    --num_qo_heads 32 --num_kv_heads 8 \
    --head_dim_qk 128 --head_dim_vo 128 \
    --page_size 16 --refcheck -vv

Python API:

from flashinfer.testing import bench_gpu_time

# CUPTI preferred, auto-fallback to CUDA events.
# Returns per-iteration times in milliseconds (a list).
times = bench_gpu_time(
    my_kernel, input_args=(x, y), enable_cupti=True,
    dry_run_iters=10, repeat_iters=30,
)
median_time_ms = statistics.median(times)

→ For complete benchmarking guide, see .claude/skills/benchmark-kernel/SKILL.md

flashinfer.gemm.group_gemm_fp8_nt_groupwise_contiguous is the standalone CuTe-DSL grouped FP8 API; group_deepgemm_fp8_nt_groupwise keeps its original DeepGEMM-only signature. Both use per-row expert indices, unlike the indptr-based group_gemm_fp8_nt_groupwise API.

For contiguous grouped FP8 GEMM, python benchmarks/bench_grouped_fp8.py --production-shapes compares DeepGEMM and CuTe-DSL with preallocated outputs; --cache-probe measures first-call and new-token-count compilation separately. Compiled grouped CuTe-DSL kernels are reused across token counts with the same 128-row alignment class, with at most two variants per device and weight shape. Concurrent misses share one compilation; warm hits do not acquire the compile lock. The standalone API's optional validate_indices=True checks expert-index values and synchronizes with the CPU. Validate new routing data before CUDA graph capture; the default path assumes valid indices and performs metadata checks only.

Prims-TS decode keeps raw storage dtypes (q_dtype, k_dtype, v_dtype) separate from effective MMA dtypes (qk_dtype, pv_dtype). Transformed FP8/NVFP4 K/V uses Q's compute dtype; BF16 Q/K with FP8 V uses FP8 PV without that transform. Kernel predicates, P packing, and attention-sink scaling must follow effective compute precision; TMA descriptors and raw allocations must follow storage precision. The public plan accepts separate k_data_type and v_data_type, with kv_data_type retained as a common-KV compatibility alias.

Code Linting

Run all pre-commit hooks:

pre-commit run -a

Install hooks to run on every commit:

pre-commit install

Opening a Pull Request

Follow CONTRIBUTING.md when opening a PR. The two requirements most often missed:

  • Keep the default PR template. Fill in .github/pull_request_template.md (Description, Related Issues, Checklist, Tests, Reviewer Notes). Do not overwrite or replace it with a custom or tool-generated description format. The title and description normally become the commit title and message on squash-merge, and are what a git bisect surfaces months later when someone is hunting the owner of a regression.
  • Report before/after numbers for performance work. A perf PR must include measurements from a reproducible benchmark (e.g. benchmarks/flashinfer_benchmark.py), naming the GPU and the problem sizes. Speedup ratios without absolute numbers, or numbers without a named GPU, are not enough.

When preparing changes for a PR, a self-review is strongly recommended: follow .claude/skills/self-review/SKILL.md, which walks the diff through the repository's review guidance and PR rules. It is informal and not required — its purpose is to make review easier for reviewers and contributors alike by catching common issues before a human looks. Changes made only to test something locally, with no PR intended, do not need one.

→ For the complete contribution rules, see CONTRIBUTING.md

Code Review

When reviewing a diff (as an agent or a human), follow the shared focus areas, kernel-review policy, and effort calibration in docs/code_review_guidance.md. Note: unlike human review, agents keep kernel implementation details in scope — read the kernel logic and report bugs, labeling findings by confidence.

→ For the complete review rule set, see docs/code_review_guidance.md

Architecture: JIT Compilation System

FlashInfer's JIT system has three layers:

Layer 1: JitSpec (flashinfer/jit/core.py)

JitSpec is an abstract base class defining the kernel-module lifecycle (try_load() / build() / load(), with the shared build_and_load() template method handling caching, locking, and FLASHINFER_DISABLE_JIT). One subclass per compilation toolchain:

  • JitSpecNvcc (nvcc/ninja modules, returned by gen_jit_spec()) defines:
    • name: Unique identifier (URI hash from parameters)
    • sources: List of .cu/.cpp files to compile
    • extra_cuda_cflags, extra_cflags, extra_ldflags: Compiler flags
  • JitSpecCuteDsl (flashinfer/jit/cute_dsl_core.py) caches CuTe-DSL kernels (see "CuTe-DSL kernels" under Module Caching below)
  • The optional TIRx KDA backend uses a private JitSpec in flashinfer/kda_kernels/tirx/cache.py to cache TVM modules, with source/compiler fingerprints and artifact checksums. Its TVM/TIRx dependencies are loaded only by explicit recurrent_kda(..., backend="tirx"); see docs/api/kda.rst.

JIT Directory Rules

NEVER write to package directories - they may be read-only after installation.

Directory Writable Use for
FLASHINFER_GEN_SRC_DIR ✓ Yes Generated source files (Jinja output, copied .cu files)
FLASHINFER_JIT_DIR ✓ Yes Compiled .so outputs
FLASHINFER_CSRC_DIR ✗ No Read-only source templates
FLASHINFER_AOT_DIR ✗ No Read-only pre-compiled binaries

Compilation Context: Architecture-Specific Compilation

FlashInfer uses CompilationContext to manage CUDA architecture targets. Some kernels only work on specific GPU architectures (e.g., Hopper SM90, Blackwell SM100/SM12x).

How it works:

  • Auto-detects GPUs in system or reads FLASHINFER_CUDA_ARCH_LIST environment variable
  • CuTe-DSL FP4 MM compilation parallelism can be controlled with FLASHINFER_MM_FP4_CUTE_DSL_COMPILE_WORKERS
  • JIT modules specify supported_major_versions=[9, 10, 11, 12] to limit compilation to specific SM versions
  • If GPU not supported → RuntimeError: No supported CUDA architectures found

→ See .claude/skills/add-cuda-kernel/SKILL.md for usage examples

Layer 2: Code Generation

Every gen_*_module() function in flashinfer/jit/ follows this pattern:

def gen_some_module(dtype_in, dtype_out, ...):
    # 1. Compute unique identifier from parameters
    uri = get_some_uri(dtype_in, dtype_out, ...)

    # 2. Create generation directory
    gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri

    # 3. (Optional) Render Jinja template to generate type-specialized config
    # Skip this step if you don't need type specialization
    with open(jit_env.FLASHINFER_CSRC_DIR / "some_customize_config.jinja") as f:
        template = jinja2.Template(f.read())
    config_content = template.render(
        dtype_in=dtype_map[dtype_in],
        dtype_out=dtype_map[dtype_out],
        # ... more parameters
    )
    write_if_different(gen_directory / "some_config.inc", config_content)

    # 4. Copy source files to gen directory
    sources = []
    for fname in ["some_kernel.cu", "some_jit_binding.cu"]:
        shutil.copy(jit_env.FLASHINFER_CSRC_DIR / fname, gen_directory / fname)
        sources.append(gen_directory / fname)

    # 5. Return JitSpec
    return gen_jit_spec(uri, sources, extra_cuda_cflags=[...])

Note: If your operation doesn't need type specialization, you can skip step 3 entirely and just copy the source files directly.

Layer 3: Compilation and Loading

JitSpecNvcc methods:

  • write_ninja() - Generates build.ninja file
  • build() - Executes ninja to compile sources
  • build_and_load() - Compiles and loads via TVM-FFI

The generated build.ninja file uses nvcc to compile .cu → .cuda.o → .so, then loads via TVM-FFI.

Jinja Templates (Optional)

Note: Jinja templates are NOT required. You can write C++ code directly without templating.

For operations that need type specialization, templates in csrc/*.jinja can generate C++ code:

// Input template
using DTypeIn = {{ dtype_in }};
using DTypeOut = {{ dtype_out }};
constexpr int PARAM = {{ param_value }};

// After render
using DTypeIn = float16;
using DTypeOut = float16;
constexpr int PARAM = 128;

This allows the same CUDA template code to be compiled with different concrete types. However, if your operation doesn't need this, you can skip Jinja and write the .cu files directly.

Directory Structure

flashinfer/
├── include/flashinfer/           # Header-only CUDA kernel templates
│   ├── attention/                # Attention kernels
│   ├── gemm/                     # GEMM kernels
│   ├── comm/                     # Communication kernels
│   ├── mma.cuh                   # Matrix multiply utilities
│   ├── utils.cuh                 # Common utilities
│   └── [...]
│
├── csrc/                          # Framework bindings (via TVM-FFI)
│   ├── *.cu                       # Kernel launcher implementations
│   ├── *_jit_binding.cu           # TVM-FFI exports
│   ├── *_customize_config.jinja   # Type config templates (optional)
│   └── [...]
│
├── flashinfer/                    # Python package
│   ├── jit/
│   │   ├── core.py                # JitSpec, compilation infrastructure
│   │   ├── cpp_ext.py             # Ninja build generation
│   │   ├── env.py                 # Workspace paths
│   │   ├── attention/             # Attention module generators
│   │   ├── gemm/                  # GEMM module generators
│   │   ├── fused_moe/             # MOE module generators
│   │   └── [...]
│   ├── gemm/                      # GEMM Python APIs
│   ├── fused_moe/                 # MOE Python APIs
│   ├── comm/                      # Communication Python APIs
│   ├── *.py                       # Other high-level Python APIs
│   ├── aot.py                     # AOT compilation for pre-built packages
│   └── [...]
│
├── tests/                         # Test suite
│   ├── attention/                 # Attention kernel tests
│   ├── gemm/                      # GEMM kernel tests
│   ├── moe/                       # MOE kernel tests
│   ├── comm/                      # Communication tests
│   ├── utils/                     # Utility tests
│   └── conftest.py                # Pytest configuration
│
└── build_backend.py               # PEP 517 build backend

Critical Rule: Framework Separation

Torch headers MUST NOT be included in include/ directory files.

  • include/: Framework-agnostic CUDA kernels (accept raw pointers)
  • csrc/: Framework bindings via TVM-FFI (currently PyTorch, but can support other frameworks)

Adding a New Operation

→ For complete step-by-step tutorial, see .claude/skills/add-cuda-kernel/SKILL.md

Quick overview of the process:

  1. Write kernel in include/flashinfer/new_op.cuh (framework-agnostic, raw pointers)
  2. Write launcher in csrc/new_op.cu (PyTorch tensor handling)
  3. Create TVM-FFI bindings in csrc/new_op_jit_binding.cu
  4. (Optional) Create Jinja template for type specialization
  5. Write JIT module generator in flashinfer/jit/new_op.py
  6. Write Python API in flashinfer/new_op.py with @functools.cache
  7. Write tests in tests/
  8. Register in flashinfer/aot.py for AOT compilation
  9. Export in flashinfer/__init__.py
  10. Add a TraceTemplate in flashinfer/trace/templates/ and wire it via @flashinfer_api(trace=...) (see below)
  11. Add an example call in tests/trace/example.py, re-run to regenerate fi_trace_out/, and commit the new JSON files

Trace Template Checklist (for new or updated APIs)

Every public API decorated with @flashinfer_api should also carry a trace= argument so that fi_trace() works and auto-dump produces a benchmark definition JSON.

  1. Create or update a TraceTemplate in flashinfer/trace/templates/<category>.py (e.g., norm.py, activation.py, cascade.py, gdn.py). Define axes, inputs, outputs, and optionally a reference function.
  2. Wire the template to the API: @flashinfer_api(trace=my_trace) on the Python function (or class method's run()).
  3. Add an example call in tests/trace/example.py that exercises the new trace with realistic shapes.
  4. Regenerate examples: rm -rf tests/trace/fi_trace_out && python tests/trace/example.py — verify the expected JSON appears.
  5. Update the docstring in tests/trace/example.py to list the new file(s).
  6. Run tests: pytest tests/trace/ -v — all template-consistency and end-to-end tests must pass.
  7. Commit the new JSON files under tests/trace/fi_trace_out/ alongside the code changes.

Example implementations:

  • Simple: flashinfer/norm/__init__.py (RMSNorm) - no Jinja, good starting point
  • Moderate: flashinfer/sampling.py - with Jinja templating
  • Complex: flashinfer/decode.py - plan-run pattern, advanced workspace

Experimental Code

Experimental backends live under flashinfer/experimental/; experimental public APIs live in core marked with @flashinfer_experimental_api (defined in flashinfer/api_logging.py). Using them is an explicit opt-in: calling an experimental API, or naming an experimental backend with backend="<name>", needs no environment variable (both warn once). Only automatic selection is gated: backend="auto" skips backends marked @experimental_backend unless FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. Experimental code is JIT-only (never registered in flashinfer/aot.py) and tested under tests/experimental/.

→ When working under flashinfer/experimental/, follow flashinfer/experimental/CLAUDE.md; the full policy is flashinfer/experimental/README.md.

Key Architectural Patterns

Module Caching

FlashInfer uses two-level caching to avoid recompilation:

  1. Python-level (@functools.cache): In-memory cache of loaded modules
  2. File-level (~/.cache/flashinfer/): Compiled .so files on disk

Cache invalidation (automatic):

  • Source file changes (SHA256 hash)
  • Compilation flags change
  • CUDA architecture change
  • FlashInfer version change

URI computed as: hash(operation_type + parameters + source_hashes + flags + cuda_arch)

Cache management:

  • Clear cache: rm -rf ~/.cache/flashinfer/
  • Override location: export FLASHINFER_WORKSPACE_BASE="/scratch"

CuTe-DSL kernels (flashinfer/jit/cute_dsl_core.py): cute.compile() has no persistent cache, so JitSpecCuteDsl (a JitSpec subclass, wrapped by the build_and_load_cute_dsl_kernel() helper) exports compiled kernels as object files (export_to_c()) and reloads them with cute.runtime.load_module(..., enable_tvm_ffi=True). Artifacts live in cached_ops/ next to the nvcc modules, one directory per op family — cached_ops/<module>_<arch>_cute_dsl/ (e.g. nvfp4_quantize_sm100a_cute_dsl/) holding one meta.json plus one .o per specialization. The arch comes from the DSL's compile target (CUTE_DSL_ARCH or the current device) since the artifacts are single-arch, unlike nvcc fatbins. Invalidation is module-granular: a changed nvidia-cutlass-dsl version or kernel-source SHA256 wipes and lazily rebuilds the module. Reference usage: flashinfer/quantization/kernels/nvfp4_quantize.py. Disable with FLASHINFER_CUTE_DSL_DISABLE_CACHE=1.

Dispatch Macros

Handle combinatorial parameter spaces:

DISPATCH_DTYPE(input_dtype, DTypeIn, {
  DISPATCH_DTYPE(output_dtype, DTypeOut, {
    DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {
      LaunchKernel<DTypeIn, DTypeOut, BLOCK_SIZE>(...);
    });
  });
});

Defined in .jinja files and expanded after rendering.

API Logging with @flashinfer_api

FlashInfer provides the @flashinfer_api decorator for debugging API calls.

Key features:

  • Crash-safe: Logs inputs BEFORE execution (preserves info even if kernel crashes)
  • Zero overhead when disabled: FLASHINFER_LOGLEVEL=0 (default)
  • Multiple verbosity levels: 0 (off), 1 (names), 3 (inputs/outputs), 5 (+ statistics)
  • CUDA graph compatible: Auto-skips stats during graph capture

Quick usage:

# Enable detailed logging
export FLASHINFER_LOGLEVEL=3              # 0, 1, 3, or 5
export FLASHINFER_LOGDEST=debug.log       # stdout, stderr, or file path

python my_script.py

Why use this?

  • Debug CUDA crashes (see inputs that caused crash)
  • Track tensor shapes/dtypes through pipeline
  • Detect NaN/Inf issues (level 5)

→ For complete debugging guide, see .claude/skills/debug-cuda-crash/SKILL.md

Debugging

Enable Logging

export FLASHINFER_JIT_VERBOSE=1      # Verbose JIT output
export FLASHINFER_JIT_DEBUG=1        # Debug symbols, -O0
export FLASHINFER_LOGLEVEL=3         # API logging (0=off, 1=basic, 3=detailed)
export FLASHINFER_LOGDEST=stdout

Inspect Generated Code

# Generated sources
ls -la ~/.cache/flashinfer/0.6.0/*/generated/

# Compiled modules
ls -la ~/.cache/flashinfer/0.6.0/*/cached_ops/

# Build files
cat ~/.cache/flashinfer/0.6.0/*/cached_ops/*/build.ninja

Environment Variables

# Compilation
export FLASHINFER_NVCC_THREADS=4              # Threads per nvcc process (--threads=N)
export MAX_JOBS=4                             # Parallel ninja jobs (nvcc processes)
export FLASHINFER_CUDA_ARCH_LIST="8.0 9.0a"  # Target architectures
# Memory note: total compilation memory ≈ MAX_JOBS × FLASHINFER_NVCC_THREADS × per-thread mem.

# Behavior
export FLASHINFER_WORKSPACE_BASE="/scratch"   # Custom cache directory

Full Environment Variable Reference

Already covered above (Quick Reference / Debugging): FLASHINFER_LOGLEVEL, FLASHINFER_LOGDEST, FLASHINFER_JIT_VERBOSE, FLASHINFER_JIT_DEBUG, FLASHINFER_CUDA_ARCH_LIST, FLASHINFER_NVCC_THREADS, FLASHINFER_WORKSPACE_BASE, MAX_JOBS.

The remaining FLASHINFER_* knobs read by the code are listed below. Defaults match what the code uses today; values are strings unless noted.

JIT / Build Toolchain
Variable Default Read in Effect
FLASHINFER_DISABLE_JIT unset flashinfer/jit/core.py If set (any non-empty value), JIT compilation is refused and modules must already exist in the cache or be provided via AOT packages.
FLASHINFER_CUTE_DSL_DISABLE_CACHE 0 flashinfer/jit/cute_dsl_core.py 1 disables the on-disk cache for JIT-compiled CuTe-DSL kernels (every process recompiles via cute.compile).
FLASHINFER_CUDNN_FROST_PTXAS unset flashinfer/fused_moe/backends/cudnn_frost/compiler.py Select an external PTXAS executable path for all Frost dtypes; unset or empty uses the bundled assembler. The executable path, version, and SHA-256 participate in compiler and tactic cache identities. MoE engines snapshot this configuration at build time; call layer.refresh_backends() or reconstruct the layer to adopt changes.
FLASHINFER_DISABLE_VERSION_CHECK unset flashinfer/jit/env.py Skip the check that pins flashinfer-cubin and flashinfer-jit-cache to the installed flashinfer-python. Without it, mismatched kernel wheels (typically left behind by an upgrade) are ignored with a warning at import, falling back to on-demand cubin downloads and JIT compilation until flashinfer download-kernels installs matching ones. Set it only when you intentionally mix versions.
FLASHINFER_JIT_LINEINFO 0 flashinfer/jit/core.py 1 adds -lineinfo to nvcc so profiler / cuda-gdb can map PTX back to CUDA source.
FLASHINFER_JIT_PREBUILD_MAX_JOBS unset (uses MAX_JOBS) flashinfer/jit/core.py Maximum parallel Ninja jobs for a bulk build_jit_specs() prebuild. This does not control ordinary on-demand module builds. The sharded test runner sets it to the host-wide automatic build budget before reducing each worker's MAX_JOBS; an explicit value is preserved.
FLASHINFER_NVCC $cuda_home/bin/nvcc flashinfer/jit/cpp_ext.py Override the nvcc binary used by the JIT (useful for sccache wrappers or non-default CUDA installs).
FLASHINFER_NVCC_LAUNCHER "" (sccache in the devcontainers) flashinfer/jit/cpp_ext.py Optional launcher prefix for nvcc (e.g. ccache, sccache). Combined with FLASHINFER_NVCC.
FLASHINFER_CXX_LAUNCHER "" (sccache in the devcontainers) flashinfer/jit/cpp_ext.py Same idea as FLASHINFER_NVCC_LAUNCHER but for the host C++ compiler.
FLASHINFER_FMHA_V2_VERBOSE unset flashinfer/jit/attention/fmha_v2/fmha_library.py (FMHA v2 codegen, C++ side) When set, the FMHA-v2 codegen / runtime prints verbose dispatcher diagnostics. Leave unset for normal runs.
FLASHINFER_EXTRA_CFLAGS unset flashinfer/jit/cpp_ext.py Extra compiler flags passed to the host C++ compiler.
FLASHINFER_EXTRA_CUDAFLAGS unset flashinfer/jit/cpp_ext.py Extra compiler flags passed to nvcc.
FLASHINFER_EXTRA_LDFLAGS unset flashinfer/jit/cpp_ext.py Extra linker flags passed to the linker.
FLASHINFER_CAKE_GDN_VALIDATE_SLOTS 0 flashinfer/gdn_decode.py 1 enables extra slot-validation checks in the CAKE GDN path for debugging or invariant validation. Leave disabled in normal runs to avoid extra validation overhead.
FLASHINFER_CAKE_DSA_CHECK_DST_MAP 0 flashinfer/experimental/cake_dsa_train/cake_backend.py 1 makes the CAKE DSA training backward validate the values of a caller-provided dkv_dst_map (every destination row must lie in [0, S_dst) of dkv_acc) on each call, at the cost of one device synchronization; the kernel itself does not range-check the map. Leave disabled in normal runs.
FLASHINFER_CAKE_DSA_DKV_DIRECT auto flashinfer/experimental/cake_dsa_train/cake_backend.py Selects how the CAKE DSA training backward accumulates dK / dV when the caller passes dkv_acc: auto takes the direct path (the main stage reduces into the caller's packed fp32 rows; no workspace accumulators, zero fill or cast launch) whenever the program registers the natural-layout stages and the row has at least dkv_direct.min_keys_per_query (4) keys per query token, 1 forces the direct path whenever the stages are registered, 0 always uses the accumulator-and-cast path. Any other value raises ValueError. Read on every plan.
FLASHINFER_CAKE_DSA_TRAIN_BINDING_CACHE 1 flashinfer/experimental/cake_dsa_train/cake_backend.py 0 disables the CAKE DSA training binding cache (every call re-validates and re-binds its inputs instead of launching from the remembered argument plans). Read once at import.
FLASHINFER_CAKE_DSA_TRAIN_BINDING_CACHE_CAPACITY 256 flashinfer/experimental/cake_dsa_train/cake_backend.py Number of forward / backward bindings the CAKE DSA training binding cache keeps before evicting the least recently used one.
FLASHINFER_CAKE_LM_HEAD_LOSS_GEMM_TUNING unset flashinfer/experimental/cake_lm_head_loss/cake_backend.py JSON object of explicit per-GEMM knobs for the eager chunked LM-head + loss entry points (the gemm_tuning argument of prepare_lm_head_loss takes the same object), e.g. {"dx": {"tile_n": 256, "k_slices": 2}}; a knob given as null pins the default form. Unset / empty: the per-geometry rules choose. A variant the compiled record does not register fails closed.
FLASHINFER_CAKE_LM_HEAD_LOSS_FUSE_DW_CAST 1 flashinfer/experimental/cake_lm_head_loss/cake_backend.py 0 turns off the fused weight-gradient cast: every chunk then accumulates into the FP32 dW_acc and a separate kernel applies the upstream scale and the output cast (by default the last chunk's weight-gradient GEMM fuses both into its epilogue). dW is bitwise the same either way.
FLASHINFER_CAKE_LM_HEAD_LOSS_DX_FINALIZE 1 flashinfer/experimental/cake_lm_head_loss/cake_backend.py 0 turns off the fused dX finalize: a K-sliced dX GEMM's FP32 slabs are then added by in-place add_ launches and a compacted call's dX is cast, zero-filled and index_copy_-scattered on the host instead of the slab_sum and scale_cast_scatter_bf16 kernels (bitwise the same dX; read once per forward).
FLASHINFER_CAKE_LM_HEAD_LOSS_DW_STREAM auto flashinfer/experimental/cake_lm_head_loss/cake_backend.py Launches each chunk's weight-gradient accumulate GEMM on a per-device side stream (forked after the chunk's row gradients, joined before the chunk buffer is reused and before the call returns) so it overlaps the tail of the chunk's dX GEMM: auto for calls of three or more chunks, 1 for every multi-chunk call, 0 never. Same kernels in the same order per kernel; bitwise the same outputs.
FLASHINFER_CAKE_LM_HEAD_LOSS_HIDDEN_COUNT 1 flashinfer/experimental/cake_lm_head_loss/cake_backend.py 0 turns off the hidden valid-row count of the eager entry points (served to calls of three or more chunks; one- and two-chunk calls take the host-count path on every setting): a compacted call then counts its valid rows and forms their index on the host before its first launch, instead of forming the index on the device and queueing chunk 0's row gather kernel and device-count logits GEMM before the count is read back through a pinned cell and a CUDA event. Bitwise the same outputs. On CUDA 13.0 the sm_103a calls take the host-count path regardless (cake_jit.toolchain_runs_hidden_count; the README's toolchain note).
FLASHINFER_CAKE_LM_HEAD_LOSS_COMPACT_ROWS 1 flashinfer/experimental/cake_lm_head_loss/cake_backend.py 0 turns off valid-row compaction: the chunk loop then covers every row instead of only the rows whose label is not -100.
FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE 1 flashinfer/experimental/cake_lm_head_loss/cake_backend.py 0 disables the per-process LRU cache of compiled bindings (read at import).
FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE_CAPACITY 64 flashinfer/experimental/cake_lm_head_loss/cake_backend.py Capacity of that binding cache (a positive integer).
FLASHINFER_JIT_CACHE_PROVIDER_ARCHS required flashinfer-jit-cache/build_backend.py Space-separated provider architectures added to a shim wheel's exact Requires-Dist metadata.
FLASHINFER_JIT_CACHE_PROVIDER_ARCH required flashinfer-jit-cache-provider/package_config.py Select exactly one architecture, such as 9.0a or sm120f, when building a binary provider wheel.

Release/nightly JIT-cache provider builds run sccache in server-side mode and wrap the complete wheel build with a no-output watchdog. After 90 minutes with no build output, the watchdog captures system, process, compiler, and sccache state; terminates the wrapped build process group (TERM, then KILL after a two-minute grace period); and exits with status 124. Compiler discovery scans the whole build container, so diagnostics include compiler processes owned by the sccache daemon even though they are outside the wrapped process group. The CI docker run --rm boundary provides final cleanup for those processes.

Cubin / Artifact Loader
Variable Default Read in Effect
FLASHINFER_CUBIN_DIR <FLASHINFER_WORKSPACE_BASE>/.cache/flashinfer/cubins flashinfer/jit/env.py (exposed via flashinfer/__main__.py show-config) Local directory used to cache downloaded cubins. Override to share a cache between users.
FLASHINFER_CUBINS_REPOSITORY https://edge.urm.nvidia.com/artifactory/sw-kernelinferencelibrary-public-generic-local flashinfer/jit/cubin_loader.py Base URL the loader downloads cubins from. Point to a mirror for offline or air-gapped setups.
FLASHINFER_CUBIN_CHECKSUM_DISABLED unset flashinfer/jit/cubin_loader.py If set, skip SHA checksum verification of downloaded cubins. Debug aid only.
FLASHINFER_CUBIN_DOWNLOAD_THREADS 4 flashinfer/artifacts.py Thread-pool size used by flashinfer artifacts download.
FLASHINFER_CUBIN_RETRY_WINDOW_SECONDS 0 flashinfer/artifacts.py Total wall-clock retry window for download_artifacts(). A failed artifact keeps re-attempting, with capped exponential backoff, until the window closes (e.g. 86400 for a 24-hour nightly window). 0 keeps only download_file()'s own per-call retry budget.
FLASHINFER_NO_DOWNLOAD unset flashinfer/jit/cubin_loader.py Hard-fail if a cubin is missing locally instead of attempting to download. Useful in CI / locked-down environments.
FLASHINFER_DSL_FMHA_LOCAL_DIR unset flashinfer/attention/cute_dsl/fmha.py Path to a local checkout of the CuTe-DSL FMHA kernel sources. The loader checks here before downloading.
FLASHINFER_LOGGING_LEVEL INFO flashinfer/artifacts.py, flashinfer/jit/core.py Python logging level for the artifacts/cubin loader and the JIT compiler (DEBUG/INFO/WARNING/ERROR). Distinct from FLASHINFER_LOGLEVEL.
FLASHINFER_DISABLE_TINYGEMM2_SM100 0 flashinfer/gemm/routergemm.py Set to 1 to disable the generated SM100/SM103 tinygemm2 backend and force the dispatcher to fall back to the legacy implementation. Useful as an escape hatch when debugging backend-selection or kernel issues on Blackwell systems.
API Dump / Logging Extensions

These complement FLASHINFER_LOGLEVEL/FLASHINFER_LOGDEST and are all read in flashinfer/api_logging.py.

Variable Default Effect
FLASHINFER_DUMP_DIR flashinfer_dumps Directory where @flashinfer_api writes per-call tensor dumps when logging is enabled.
FLASHINFER_DUMP_INCLUDE "" Comma-separated allow-list of API names; only matching calls are dumped. Wildcards (*) supported.
FLASHINFER_DUMP_EXCLUDE "" Comma-separated deny-list of API names; matching calls skip the (expensive) stats / dump path.
FLASHINFER_DUMP_MAX_COUNT 1000 Hard cap on number of dumped events; once reached, further dumps are dropped with a warning.
FLASHINFER_DUMP_MAX_SIZE_GB 20 Hard cap on total dump size in gigabytes (GB) written to FLASHINFER_DUMP_DIR (parsed as float; see _DUMP_MAX_SIZE_GB in flashinfer/api_logging.py).
FLASHINFER_DUMP_SAFETENSORS 0 1 writes dumps as .safetensors (portable / Hugging Face-style); otherwise plain .pt.
Trace Capture

Used by flashinfer.trace / fi_trace.

Variable Default Read in Effect
FLASHINFER_TRACE_DUMP unset flashinfer/fi_trace.py If set, decorated APIs auto-dump benchmark-definition JSON for each call (the "fi_trace" feature).
FLASHINFER_TRACE_DUMP_DIR cwd flashinfer/fi_trace.py Directory where the trace JSON files are written.
FLASHINFER_TRACE_APPLY 0 flashinfer/__init__.py, flashinfer/trace_apply/config.py Set to 1 to enable Trace Apply (runtime kernel substitution) when FlashInfer is imported.
FLASHINFER_TRACE_APPLY_PATH unset flashinfer/trace_apply/config.py Directory from which deployment-configured solutions are loaded for Trace Apply.
Validation / Autotuning / Routing / Kernel Selection
Variable Default Read in Effect
FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS unset flashinfer/utils.py (backend_requirement), flashinfer/api_logging.py 1 lets backend="auto" (dispatch heuristics and autotuning) select backends marked @experimental_backend (see the "Experimental Code" section). Explicit use never needs it: calling an @flashinfer_experimental_api function or passing backend="<experimental>" always works and emits an ExperimentalWarning once. Without it, automatic selection considers only stable backends.
FLASHINFER_VALIDATE_INPUTS 0 flashinfer/mla/_core.py (MLA wrapper) Non-zero / non-empty value enables defensive input validation inside the MLA wrapper. Device-synchronizing checks are skipped during CUDA Graph capture. Adds host-side overhead; intended for debugging.
FLASHINFER_AUTOTUNER_LOAD_FROM_FILE 0 flashinfer/autotuner/autotuner.py 1 loads previously serialized autotune results from disk instead of re-running the search.
FLASHINFER_DIST_AWARE_AUTOTUNE 0 flashinfer/fused_moe/da_config.py 1 enables experimental distribution-aware autotune and kernel dispatch (TRT-LLM MoE only).
FLASHINFER_DA_DISTRIBUTIONS built-in distribution catalog flashinfer/fused_moe/da_config.py Comma-separated training distributions used by the experimental TRT-LLM distribution-aware MoE autotuner.
FLASHINFER_PRIMS_TS_DEBUG_CHECKS 0 flashinfer/prims_ts/batched_gemm/batched_gemm_kernel.py, batched_gemm_run.py, flashinfer/prims_ts/gemm/kernel.py 1/true/yes/on enables exhaustive deadlock/race checks and TaskManager verification in Prims-TS batched GEMM and dense GEMM. Production JIT leaves both off. tests/prims_ts turns the variable on for test_batched_gemm_* and test_gemm_*.
FLASHINFER_PRIMS_TS_COMPILE_OPTIONS --opt-level 2 flashinfer/prims_ts/utils.py, batched_gemm_run.py Extra CuTe-DSL / nvc compile options passed when building Prims-TS kernels.
FLASHINFER_PRIMS_TS_BF16_CONSTANT_WEIGHTS unset flashinfer/prims_ts/batched_gemm/batched_gemm_run.py 1 fills BF16 autotune/cast-A weight tensors with constant 1.0 (debug / deterministic fill).
FLASHINFER_PRIMS_TS_CASTA_CONSTANT_ACT unset flashinfer/prims_ts/batched_gemm/batched_gemm_run.py 1 fills cast-A / BF16 activation tensors with constant 1.0 during autotune buffer setup.
FLASHINFER_PRIMS_TS_CASTA_PATTERN_ACT unset flashinfer/prims_ts/batched_gemm/batched_gemm_run.py 1 fills activations with a quarter-wise pattern (1,2,4,6) instead of random data.
FLASHINFER_PRIMS_TS_CASTA_PATTERN_WEIGHTS unset flashinfer/prims_ts/batched_gemm/batched_gemm_run.py 1 fills cast-A weights with a deterministic pattern (takes precedence over constant fill).
FLASHINFER_PRIMS_TS_CASTA_CONSTANT_WEIGHTS unset flashinfer/prims_ts/batched_gemm/batched_gemm_run.py 1 fills cast-A weights with constant values when pattern fill is not set.
FLASHINFER_AUTOTUNE_DIR unset flashinfer/mla/_sparse_mla_sm120/_calibration.py, flashinfer/comm/pcie_ipc_tuning.py Override the disk path for tuning cache files (sparse-MLA SM120 cpb calibration constants, and the PCIe IPC all-reduce). Falls back to FLASHINFER_WORKSPACE_DIR when unset.
FLASHINFER_AUTOTUNE_TIMER unset (auto) flashinfer/autotuner/autotuner.py Selects the autotuner's per-tactic timer: globaltimer forces the GPU %globaltimer register, cuda_event forces cudaEvent, unset/anything-else auto-detects (uses %globaltimer only when Confidential Computing is detected). Under CC cudaEventElapsedTime is unreliable (can go negative), so the globaltimer path keeps tactic ranking stable.
FLASHINFER_CUTILE_AUTOTUNE_DISABLED 0 flashinfer/quantization/kernels/cutile/rope_quantize_fp8_cutile.py Non-zero skips exhaustive cuTile RoPE-FP8 tuning and uses the built-in token-count heuristic.
FLASHINFER_CONFIDENTIAL_COMPUTE unset flashinfer/utils.py Override NVIDIA Confidential Computing (CC) auto-detection used by is_confidential_compute() (which drives the autotuner timer above): 1 forces CC, 0 forces non-CC. Useful for CI or hosts without pynvml.
FLASHINFER_MSA_PREFILL_SCHEDULE unset flashinfer/msa_ops/_blackwell_sm100.py Set to m64 to force the eligible M64 Blackwell MSA prefill schedule; any other non-empty value is rejected. Leave unset for automatic routing.
FLASHINFER_MSA_FP8_Q1_SCHEDULE unset flashinfer/msa_ops/_blackwell_sm100.py Force an eligible FP8 Q1 MSA decode route: batch_attention, q1_exact, q1_flat_xform2, q1_paged_xform2, or paged_uniform_fp8. Leave unset for automatic routing.
FLASHINFER_AUTOTUNE_CACHE_DIR FLASHINFER_CACHE_DIR/autotune flashinfer/autotune_cache.py Root directory of the managed v2 autotune store used by autotune_v2() (placement only; distinct from the MLA-specific FLASHINFER_AUTOTUNE_DIR above).
FLASHINFER_RAGGED_AUTO_BACKEND_ORDER unset flashinfer/prefill.py Comma-separated backend order for BatchPrefillWithRaggedKVCacheWrapper under backend="auto" on Blackwell (e.g. cutlass,cudnn, or a single name to pin one). Overrides both the per-shape table and the global default; eligibility is still enforced, so a backend that cannot serve the problem is skipped rather than forced. For benchmarking and regression bisection on hardware whose ranking differs from the measured B200 defaults.
FLASHINFER_NVFP4_QUANTIZE_USE_TMA unset flashinfer/quantization/kernels/nvfp4_quantize.py Overrides the NVFP4 CuTe-DSL TMA-kernel heuristic: 0 forces the vectorized-load kernel, 1 enables the shape heuristic on every architecture and input dtype. Unset, TMA is used only where it was measured to win (SM107 with FP16/BF16 input, layout-specific M/K crossovers). Never taken for fp8 input or an active 4over6 recipe, whatever the setting. For benchmarking and re-tuning on new hardware.
FLASHINFER_NVFP4_PER_TOKEN_THREADS shape-dependent (128-512) flashinfer/quantization/kernels/nvfp4_quantize.py Override the CTA thread count for per-token NVFP4 quantization. Intended for benchmarking; invalid or unsupported values can make kernel compilation or launch fail.
FLASHINFER_KDA_AFFINE_POLICY model flashinfer/cake_kda_tf32_runtime.py Set to legacy to restore the pre-2026-09-23 BF16 affine-split gate for A/B measurements. model, packed, and unrecognized values use the measured cost model.
FLASHINFER_KDA_AFFINE_WINDOW_WAVES auto flashinfer/cake_kda_tf32_runtime.py Override resident affine-composite windows in CTA waves. Integer values are clamped to 1-3; unset, auto, and invalid values let the cost model choose.
FLASHINFER_KDA_AFFINE_FUSED_EPILOGUE 1 flashinfer/cake_kda_tf32_runtime.py 0 restores the equivalent PyTorch checkpoint/tail/final-state epilogue for A/B testing; any other value keeps the fused epilogue.
FLASHINFER_KDA_PLAN_CACHE_MAX_BYTES 8589934592 (8 GiB) flashinfer/kda_prefill.py Maximum workspace bytes retained by the KDA prefill plan cache. Must parse as a positive integer.
FLASHINFER_KDA_PTXAS auto-discovered flashinfer/kda_kernels/ptx/static_runtime.py Path to a ptxas executable at version 13.2 or newer for assembling retained PTX 9.2 kernels. When unset, searches the nvidia-cuda-nvcc wheel, CUDA_HOME/CUDA_PATH, PATH, then /usr/local/cuda/bin/ptxas.
FLASHINFER_KDA_PTX_CACHE_DIR FLASHINFER_JIT_DIR/kda_ptx flashinfer/kda_kernels/ptx/static_runtime.py Override the cache directory for cubins assembled from retained KDA PTX. Defaults to FLASHINFER_JIT_DIR/kda_ptx; the selected directory is created on demand.
FLASHINFER_SM90_CAKE_BF16_DEV_LAUNCHER unset flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_gemm.py Development-only factory for building the SM90 BF16 CAKE grouped-GEMM kernels in process. Set to module or module:attribute (default attribute: create_grouped_gemm); unset, empty, or 0 selects the shipped launcher.
FLASHINFER_SM90_CAKE_BF16_COMBINE_WIRE shape-dependent flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py When combine_wire=None, select per_route, prereduced, or prereduced_hilo; unset chooses per_route for token capacities up to 8 and prereduced above. Every rank in a pipe must use the same value.
FLASHINFER_TOPK_ALGO unset flashinfer/topk.py Force a specific top-k backend (otherwise the dispatcher chooses based on shape/dtype/mode via benefit gates): default (radix), clusters (SM100), cub (cub::DeviceBatchedTopK; bypasses the benefit gates). Used for benchmarking / regression bisection.
FLASHINFER_USE_CUDA_NORM 0 flashinfer/norm/__init__.py 1 switches the norm path from the default backend to the legacy CUDA-only kernels. Diagnostic toggle.
FLASHINFER_DISABLE_EAGLE_VERIFY_FP8KV_SM90 0 flashinfer/prefill.py 1 keeps BatchPrefillWithPagedKVCacheWrapper.run on the FA2 kernel for the speculative-decoding target-verify problems that the specialized SM90 fp8-KV split-KV kernel (csrc/eagle_verify_fp8kv_sm90.cu) would otherwise serve: plan(..., uniform_q_len=4) with a custom mask, 4 query heads per KV head, head_dim 256, page size 1, NHD, bf16 q/o, e4m3 KV, compute capability 9.0. Read at plan and run time. The uniform_q_len plan hint itself is unaffected.
FLASHINFER_SPECIALIZED_KERNEL_DISABLE unset flashinfer/mla_kv_pack.py 1 forces concat_mla_kv_quant_fp8 onto its composable PyTorch fallback (two fp8 casts + two strided copies, same bytes) instead of the fused MLA K/V pack kernel. Read at call time; the kill switch for A/B testing the fused path.
FLASHINFER_ROUTING_FORCE_BLOCK_PER_TOKEN unset csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cuh Forces the TRT-LLM MoE custom-routing kernel into "one-block-per-token" mode regardless of the active routing policy. Mainly used to reproduce specific perf points.
FLASHINFER_B12X_MICRO_SHARE_INPUT 1 flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py 0 disables the B12x MoE micro-batch input-sharing optimization. Internal/experimental — leave at the default unless investigating an SM12x MoE regression.
FLASHINFER_B12X_FORCE_MOE_W4A16 unset flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py 1 forces the SM12x MoE dispatcher onto the W4A16 path with BF16 activations. Internal/experimental — used to reproduce W4A16-specific issues.
FLASHINFER_B12X_STATIC_COMPACT_CUTOVER_PAIRS shape/device-dependent flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py Override the routed-pair threshold between static and dynamic MoE for device/workload tuning. Set before first use; the decision is cached. A nonnegative integer, with 0 disabling the static band. Legacy aliases, in decreasing precedence: B12X_STATIC_COMPACT_CUTOVER_PAIRS, B12X_DYNAMIC_STATIC_CUTOVER_PAIRS, B12X_LEVEL10_STATIC_CUTOVER_PAIRS.
FLASHINFER_TACTICS_BLOCKLIST unset flashinfer/autotuner/autotuner.py Path to a JSON tactics-blocklist file generated by flashinfer tactics-blocklist generate (or python -m flashinfer tactics-blocklist generate). When set, the autotuner loads the file at startup and skips any kernel tactics listed as invalid for the current GPU/driver environment, preventing hang or crash on known-bad tactics.
FLASHINFER_CUDNN_PREFILL_SHAPE_OVERRIDE 1 flashinfer/cudnn/prefill.py (ragged token-indptr path) 0 makes the cuDNN ragged prefill graph declare the exact (batch, max_seq_len) per call instead of building one override-enabled graph per length class (declared max_len 128 or 65536+, batch 4096+) and passing the real shapes at execute time via cuDNN's shape override. With matching FE 1.31+ Python/native support, BF16 SM100 MLA prefixes (D192/V128, equal 4–64 heads, B1–4, Q64–1024, KV2K–32K at least 4x Q) use bounded power-of-two batch/Q/KV classes and a packed-Q capacity instead; this preserves reuse within each class while exposing useful plan-time occupancy and workspace bounds. Crossing a class builds one cached plan. Override exists because cuDNN finalizes one execution plan per declared shape (55-70 ms, ~1 s with a kernel compile) and serving changes the shape every step; needs cuDNN 9.22+ and cudnn-frontend 1.29+, otherwise the exact declaration is used. Disable only to bisect a suspected override regression.
FLASHINFER_ALLOW_CUDNN_PLAN_BUILD_IN_CAPTURE 0 flashinfer/gemm/gemm_base.py (read once at import) 1 downgrades the refusal to build a cuDNN execution plan during CUDA graph capture from a CudnnPlanBuildInCaptureError to a warning. Set it only if you have verified your cuDNN version tolerates plan build under capture; the supported alternative is to run each shape once eagerly outside the capture region, which populates the per-shape plan cache the captured call then reuses.
FLASHINFER_DECODE_AUTO_CUDNN unset flashinfer/decode.py (_auto_decode_prefers_cudnn) BatchDecodeWithPagedKVCacheWrapper(backend="auto"): unset applies the measured rule (SM100/SM103 with cudnn-frontend 1.30+, fp16/bf16 head_dim 128, GQA group dividing 128, 2 <= q_len_per_req <= 4 with q_len_per_req * group in [32, 128] rows per CTA and batch_size * num_kv_heads >= 64 CTAs, no RoPE / soft-cap / window resolve to cudnn, whose default SM100 engine runs those multi-token rows on its FROST decode tile at 0.35-0.95x fa2; single-token decode stays on fa2 except head_dim-256 GQA models with group 4 / 8 / 16 and batch_size * num_kv_heads in [96, 256], where the d256 decode tile measures at 0.5-0.95x fa2, from 64 under use_cuda_graph=True with a caller-owned table on cudnn-frontend 1.31+, which the wrapper tells the graph is replayed (is_cuda_graph_replay_expected) so the tile splits KV; multi-token rows without tensor cores, which have no fa2 kernel, take cudnn whenever its decode path can run them); 0 never resolves to cudnn; 1 applies the rule on an older frontend too (benchmarking; the backend engine is ~8x slower there). Under CUDA graphs auto takes cudnn only with a caller-owned block_tables (the auto-built table cannot grow once captured) or for multi-token rows without tensor cores, and the resolution is frozen after the first plan. A caller-owned table narrower than 2048 tokens keeps fa2 on cudnn-frontend 1.30 (its placement ranks the backend's prefill-class engine first below that; 1.31+ lifts the limit). auto takes cudnn only where the frontend's FROST engines can run (its own CuTe-DSL check: nvidia-cutlass-dsl >= 4.7.0, above the 4.6.2 the dependency metadata admits; the 1 override and the no-fa2 route are not gated), and callers passing fixed_split_size / disable_split_kv (fa2's batch-invariance controls) keep fa2 wherever fa2 can serve the rows. wrapper.resolved_backend reports the choice.
FLASHINFER_CAKE_SSD_CHUNK_PARALLEL auto flashinfer/mamba/cake_ssd_combined.py (SSDCombined(backend="cake"), read on every call) Chooses between the two exact Mamba2 SSD prefill program families of the Cake backend: auto applies the calibrated rule (chunk_parallel_selected: the chunk-parallel family only when the call has fewer (sequence, head) work items than SMs, its per-tile S/H workspace fits in 256 MiB and its predicted main-kernel time beats the exact scan's by the selection margin), always forces the chunk-parallel family, never forces the exact scan. Both families run the same exact-scan arithmetic. Any other value raises ValueError.
FLASHINFER_NVFP4_4OVER6 unset flashinfer/quantization/nvfp4_quantization_utils.py, csrc/nv_internal/cpp/common/envUtils.cpp Legacy. 1 enables the NVFP4 "4over6" block-scale search. Standard NVFP4 gives each 16-element block the single E4M3 scale implied by its absolute max (amax / 6, because 6 is the E2M1 max). 4over6 additionally tries the 1.5x tighter amax / 4, quantizes the block both ways, dequantizes both, and keeps whichever reconstructs the block with less error — a little more quantizer work for lower quantization error. Requires fp16/bf16 input and E4M3 (not UE8M0) block scales. Consulted only when the public nvfp4_4over6= parameter is omitted (its default); a FutureWarning is emitted only when this variable actually enables 4over6 (unset or 0 stays silent). An explicit None (off) or NVFP44Over6Config(...) ignores this variable and the three below, with no per-field merge. Prefer the parameter.
FLASHINFER_NVFP4_4OVER6_ERR_MODE MAE flashinfer/quantization/nvfp4_quantization_utils.py, csrc/nv_internal/cpp/common/envUtils.cpp Legacy. Metric used to choose between the two 4over6 candidates: MAE (sum of absolute errors) or MSE (sum of squared errors). The amax / 4 candidate buys resolution for the bulk of the block by clipping its largest magnitudes, and MSE penalizes those few large residuals much harder than MAE does, so MSE selects it less often — prefer MSE when a handful of outlier activations dominate downstream accuracy. Read only when FLASHINFER_NVFP4_4OVER6=1 and nvfp4_4over6= is omitted. Parameter spelling: NVFP44Over6Config(err_mode="MSE").
FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH unset flashinfer/quantization/nvfp4_quantization_utils.py, csrc/nv_internal/cpp/common/envUtils.cpp Legacy. 1 evaluates the 4over6 candidate error in fp16 (a packed f16x2 dequant of each candidate) instead of exactly in fp32. Cheaper, and on near-ties it can pick the other candidate, so it is a distinct recipe rather than a pure speed knob. Same gating as FLASHINFER_NVFP4_4OVER6_ERR_MODE. Parameter spelling: NVFP44Over6Config(err_use_fast_math=True).
FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 unset flashinfer/quantization/nvfp4_quantization_utils.py, csrc/nv_internal/cpp/common/envUtils.cpp Legacy. 1 treats 256, not the full-range 448, as the top of the E4M3 block-scale range: the per-tensor global scale becomes 256 * 6 / amax (per-token: 1 / (256 * 6)) and the kernel's error normalization follows. This is what keeps the 1.5x-larger amax / 4 candidate representable — 256 * 1.5 = 384 still fits E4M3, whereas a 448-based scale would need 672 and saturate for blocks near the tensor amax. The global scale must be built from the same recipe (make_nvfp4_global_scale(..., nvfp4_4over6_config=...)); a mismatch silently rescales the whole tensor. Same gating as FLASHINFER_NVFP4_4OVER6_ERR_MODE. Parameter spelling: NVFP44Over6Config(e4m3_max=256).
FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH unset flashinfer/quantization/kernels/nvfp4_quantize.py, csrc/nv_internal/cpp/common/envUtils.cpp 1 makes the NVFP4 quantizer derive its block/global scales with IEEE round-to-nearest arithmetic (fdiv_rn, nvfp4_scale_from_amax_rn) instead of the default fast approximate reciprocal (rcp_approx_ftz). Slower, but reproduces a reference implementation bit-exactly — use it when bisecting an NVFP4 accuracy mismatch. The flag is part of the CuTe-DSL specialization name (_nofastmath suffix), so toggling it compiles a separate kernel rather than reusing the cached one. The C++ side (getEnvDisableFP4QuantFastMath()) also accepts TRT-LLM's TRTLLM_DISABLE_FP4_QUANT_FAST_MATH as an alias.
Experimental KDA Output-Only Decode Tuning

Low-level overrides for the output-only / RecoverSSM-verify KDA decode kernels (flashinfer/kda_kernels/kda_decode_wy_output_only.py); intended for benchmarking and bring-up.

Variable Default Effect
FLASHINFER_KDA_OO_REC_MAX_BH 192 WY-vs-recurrent auto-dispatch threshold on B * HV at T <= 2 (a quarter of it at T in (2, 4]).
FLASHINFER_KDA_OO_REC_KS batch policy Override the grouped recurrent backend's K-slices per state column.
FLASHINFER_KDA_OO_REC_VSPLIT batch policy Override the grouped recurrent backend's value-column split per CTA.
FLASHINFER_KDA_OO_MBP grid policy Override the WY kernel's minimum-blocks-per-SM launch-bound hint.
Experimental Packed-KDA Decode Tuning

These low-level overrides are for benchmarking and kernel bring-up. Production users should normally leave the batch-size policy at its defaults.

Variable Default Effect
FLASHINFER_KDA_T1_FAST_PATH 1 0 disables routing eligible recurrent KDA T=1 calls to the packed CuTe-DSL decode kernel.
FLASHINFER_PACKED_KDA_TILE_V batch policy Override the V-dimension tile width with a supported packed-KDA tile size.
FLASHINFER_PACKED_KDA_ILP batch policy Override the number of rows processed per instruction-level-parallel group.
FLASHINFER_PACKED_KDA_THREADS batch policy Override CTA threads; the value divided by 16 must select a valid group count.
FLASHINFER_PACKED_KDA_EVICT_FIRST batch policy Set to 1 or 0 to override the policy's first-use cache-eviction hint.
FLASHINFER_PACKED_KDA_CHUNKR policy-derived Override rows per staged state-pool chunk.
FLASHINFER_PACKED_KDA_STAGES batch policy Override cp.async pipeline depth; use 0 for register prefetch or at least 2 for pipelining.
FLASHINFER_PACKED_KDA_NO_CPASYNC 0 1 forces the register-prefetch kernel instead of cp.async staging.
FLASHINFER_PACKED_KDA_PERSIST 0 1 enables the persistent-CTA schedule when alignment and pipeline constraints permit it.
FLASHINFER_PACKED_KDA_TMA 0 1 uses TMA state-pool reads when the staged kernel is active; mutually exclusive with bulk stores.
FLASHINFER_PACKED_KDA_BULK 0 1 uses bulk stores for an eligible staged kernel; mutually exclusive with TMA reads.
FLASHINFER_PACKED_KDA_PRIVRING 1 1 selects private-ring synchronization when neither TMA nor bulk-store mode is active.
FLASHINFER_PACKED_KDA_MAXRREG 0 Override the generated kernel's maximum-register tuning hint (0 leaves it unspecified).
FLASHINFER_PACKED_KDA_L2HINT 256 Override the generated kernel's L2-prefetch-size hint.
FLASHINFER_PACKED_KDA_HPC 1 Override heads combined per CTA for eligible staged whole-head tiles.
FLASHINFER_PACKED_KDA_L2POL 0 Override the experimental L2 cache-policy selector.
FLASHINFER_PACKED_KDA_MINBLOCKS 0 Override the generated kernel's minimum-blocks-per-SM launch-bound hint.

Development Workflow

Typical Development Loop

  1. Edit kernel code in include/flashinfer/some_kernel.cuh
  2. Run test: pytest tests/test_some_kernel.py::test_specific_case
  3. FlashInfer detects changes and recompiles automatically
  4. No pip install needed!

Modifying Existing Kernels

  • Kernel templates: include/flashinfer/**/*.cuh - Changes picked up on next JIT compile
  • Launcher code: csrc/*.cu - May need changes if adding new template parameters
  • Jinja templates: csrc/*.jinja - Update if adding new config parameters
  • Python API: flashinfer/*.py - Update if changing function signatures

Creating Pre-compiled Packages

When ready to distribute:

# Build one binary provider for the target GPU.
export FLASHINFER_JIT_CACHE_PROVIDER_ARCH=9.0a
python -m build --no-isolation --wheel flashinfer-jit-cache-provider

# Build the matching shim with an exact dependency on that provider.
export FLASHINFER_JIT_CACHE_PROVIDER_ARCHS="9.0a"
python -m build --no-isolation --wheel flashinfer-jit-cache

The provider build runs flashinfer/aot.py, which calls the registered gen_*_module() functions and pre-compiles them for exactly one architecture. The shim contains metadata and exact dependencies, not compiled kernels. Use the same FLASHINFER_LOCAL_VERSION and FLASHINFER_DEV_RELEASE_SUFFIX values for every provider and its shim when producing CUDA-specific or nightly wheels.

Build System Details

  • Build backend: Custom PEP 517 backend in build_backend.py
  • Data directories: Build creates symlinks for editable installs:
    • 3rdparty/cutlass → flashinfer/data/cutlass
    • csrc → flashinfer/data/csrc
    • include → flashinfer/data/include
  • Version: Generated in flashinfer/_build_meta.py from version.txt

External Integrations

TVM-FFI: Cross-Language Unified ABI

FlashInfer uses TVM-FFI (Apache TVM's Foreign Function Interface) for bindings, which provides a cross-language unified ABI. This means:

  • Not limited to PyTorch: The same compiled kernels can be used from multiple frameworks
  • Language agnostic: Bindings can be created for Python, C++, Rust, etc.
  • Type-safe marshaling: Automatic tensor/array conversion between languages
  • Export syntax: Use TVM_FFI_DLL_EXPORT_TYPED_FUNC(name, func) to expose C++ functions

While FlashInfer currently provides PyTorch bindings, the underlying kernels are framework-agnostic thanks to TVM-FFI.

Other Integrations

  • PyTorch Custom Ops: torch.library for torch.compile() and CUDA graph support
  • Ninja Build: Direct ninja generation, no CMake complexity

Supported GPU Architectures

FlashInfer supports NVIDIA SM75, SM80, SM86, SM89, SM90, SM100, SM103, SM110, SM120, and SM121.

Release Versioning

FlashInfer follows a "right-shifted" versioning scheme (major.minor.patch[.post1]):

  • major: Architectural milestone and/or incompatible API changes (similar to PyTorch 2.0)
  • minor: Significant backwards-compatible new features
  • patch: Small backwards-compatible features (new kernels, new SM support) and backwards-compatible bug fixes
  • post1: Optional suffix for quick follow-up release with just backwards-compatible bug fixes

External Documentation Resources

When working with FlashInfer's dependencies and tools, refer to these official documentation sources:

Core Dependencies

When to Consult These Docs

  • Understanding new GPU architecture features → Check PTX ISA documentation for latest instruction details
  • Working on FFI bindings → Check TVM-FFI docs for export patterns and type marshaling
  • Implementing Tensor-Core kernels using CUTLASS → Read source code in 3rdparty/cutlass/
  • Using tensor layouts or warp-level operations → Refer to CuTe documentation
  • Writing inline PTX assembly → Consult PTX ISA for instruction syntax and semantics

These dependencies are included in FlashInfer's 3rdparty/ directory or requirements.txt.

Some final suggestions for all AI agents

Because practical engineering involves the accumulated experience of trial and error, match the coding style, efficiency, complexity, verbosity, and defensiveness by learning from existing code as much as possible—this document contains many pointers on where to find examples. Document intentional departures with rationale. Mentioning "AI-assisted" in the git commit message is good transparency. For performance-critical hot paths, leave justification for the special algorithmic choices and other potential alternatives in a comment for review.

Keep documentation in sync with code changes: When modifying code that is referenced in this document or in .claude/skills/, update the corresponding documentation immediately. This includes:

  • Important infrastructure changes (e.g., @flashinfer_api, @backend_requirement, TVM-FFI macros) → Update examples in CLAUDE.md and relevant skill files
  • New patterns or conventions → Document them for future reference
  • Deprecated approaches → Remove or mark as deprecated in docs
  • New error handling patterns, macros, or utilities → Add to relevant skill tutorials