CLAUDE.md

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
GDN ReplaySSM u-cache flush kernel with an fp16 state pool (on SM100 served by the UMMA backend; gated_delta_rule_mtp_ucache_flush(backend=...))	export GDN_UCACHE_STATE_DTYPE=fp16
Opt in to multi-wave KV split in trtllm-gen GQA decode (long or mixed-length KV; 2 = wave-fill split, see the table below)	export FLASHINFER_TRTLLM_GEN_MULTI_WAVE_KV_SPLIT=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
Disable SM107 fused routing tail masking (enabled by default; use a separate fill)	export FLASHINFER_SM107_FUSED_TAIL_MASK=0
Log MegaMoE compilation and workspace diagnostics	export FLASHINFER_MEGAMOE_LOG_COMPILE_SPEC=1
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"
Plan cake_fmha small-M decode splits on the device's SM count without reading the launch stream (on by default: splits follow the stream's SM partition; turn it off only when no small-M launch runs on a green context; flashinfer.green_ctx helpers turn it back on)	export FLASHINFER_CAKE_FMHA_SM_PARTITIONS=0
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.

GitHub PR Test GPU jobs keep a different, per-runner cache: they set CI_COMPILE_CACHE_DIR to a directory under the runner's work directory, and ci/bash.sh namespaces it by runner and the pulled image's immutable ID. CI_COMPILE_CACHE_MODE=read-only (the default) mounts a read-only seed and runs against a private writable copy, discarded even after test failures. The PR Test workflow selects read-write only for scheduled runs on main. A stable per-runner lock protects maintenance and execution, including accidentally shared workspaces. Other ci/bash.sh callers leave caching unset and run cold. Before each trusted writable job its runner's cache is cleared when it exceeds CI_COMPILE_CACHE_MAX_MB (default 40960) or its volume has less than CI_COMPILE_CACHE_MIN_FREE_MB (default 20480) free. Only content-keyed artifacts survive between jobs: CuTe-DSL (with FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1), TIRx, KDA PTX, Triton and downloaded cubins. nvcc modules, tuning results and the generated-source tree are pruned, since a fresh checkout rebuilds nvcc modules anyway. Each test job's FLASHINFER_JIT_BUILD_LOG is summarized per toolchain by scripts/summarize_jit_builds.py in the job summary. See ci/COMPILE_CACHE.md for lifecycle tests and trust boundaries.

A shared S3 layer sits under it (ci/compile_cache_s3.sh), so new EC2 leases and new runners do not start cold: one tarball per image ID, GPU and lane under s3://<SCCACHE_BUCKET>/compile-cache/main/, readable anonymously by every job. Pull request jobs keep a per-runner mirror of it (refreshed by ETag); the job's entrypoint copies the mirror first and fills its gaps from the seed, taking a CuTe-DSL module whole when absent and adding objects only when its meta.json is identical. Only scheduled runs on main write, through the OIDC role in the COMPILE_CACHE_S3_ROLE_ARN repository variable, whose IAM trust policy accepts only refs/heads/main; until it is set they update seeds as before. A writer retires its runner's seeds, compiles in a new runner.temp directory on top of the S3 tarball alone (without the shared pip cache), then uploads it as one object, so concurrent writers never interleave a module. It skips the upload when it built nothing new or its restore failed. A tarball over CI_COMPILE_CACHE_S3_MAX_MB (default 4096) is cut down to that job's own builds, dropping kernels it reused. Restore and upload failures only warn.

cuDNN linear attention

Use nvidia-cudnn-frontend>=1.32.0; KDA and GDN normalization use additive 1e-6 without an epsilon argument. Auto selection also requires CuTe DSL 4.7+. The bounded KDA SM100/SM103 and GDN SM100 policies live in kda.py and gdn_prefill.py. KDA's head/length regions are documented in docs/api/kda.rst. Auto-selected cuDNN ignores the native KDA workspace; native fallback uses it. Validate 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.

cuDNN query sparse attention

flashinfer.qsa_forward accepts packed BF16 QSA block bitmaps; flashinfer.qsa_indexed_forward accepts existing causal token IDs. Both call the public cudnn.ops API and return BF16 output, FP32 output and natural-log LSE. They do not register autograd. Keep cuDNN imports lazy, preserve its input contract, and test changed selection during CUDA graph replay. Indexed forward is not bitwise repeatable because cuDNN's index compaction can change accumulation order. Cover both sides of the bitmap dispatch boundary at 512/513 tokens when testing cuDNN Frontend 1.32.

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

The sharded test runner defaults to one worker per visible GPU. H100 PR CI opts into --workers-per-gpu 2: eight workers on four GPUs. Other CI lanes retain the default. Locally, scripts/task_run_unit_tests.sh --workers-per-gpu 2 enables sharing; --workers N sets the total instead. Workers share the existing JIT compilation budget.

Scheduling markers apply to the whole source file:

• node_exclusive: run alone across the allocation after parallel work, with all GPUs visible. Use for multi-GPU requirements or host-memory OOM risk.
• gpu_exclusive: run alone on the assigned GPU; other GPUs can keep working.
• host_memory_heavy: prioritize this file and serialize its units with other host-heavy work, while lighter files can overlap.

Host and GPU reservations cover setup, execution, and process cleanup, and are acquired together. Blocked units stay queued so workers can take runnable work from their own queue or steal from another. Pending GPU-exclusive work can drain one GPU without reserving the host-heavy slot or blocking other GPUs. RESOURCE WAIT reports worker idle time of at least one second while work is queued but cannot run on that worker's resources.

These markers enforce overlap rules, not memory limits. Most H100 annotations use measured peaks above 128 GiB host RSS or 32 GiB GPU memory. Attention sink also uses both narrower markers: its single-GPU H100 run peaked at 89.2 GiB host RSS and 4.5 GiB GPU memory, and other GPUs can keep working while it runs. Cold JIT compilation and other GPU families can change the peaks. Validate completion rate, OOM failures, and repeated wall times with matching scope and cache conditions when changing a CI lane's concurrency policy.

Capacity reports include the serial floor from host-heavy and node-exclusive work. They remain optimistic about GPU contention because the estimator uses worker counts without GPU topology or measured memory capacity.

Allocation telemetry defaults to a 5-second interval (RESOURCE_MONITOR_INTERVAL); per-batch sampling defaults to 2 seconds (MEMORY_MONITOR_INTERVAL). The sampler uses resource_monitor.py, with sampling, CSV persistence, and summary aggregation kept separate. Allocation and batch monitors share host CPU readers in scripts/test_sharding/host_cpu.py. Unavailable per-process GPU memory stays unknown in snapshots and marks source resource summaries partial. memory_summary.csv leaves GPU maxima blank when no valid observation exists. Host CSV rows are flushed before GPU queries; query costs are in sampler.csv. Nested runners reuse a sampler only while its exported owner PID is alive. Namespace-relative cgroup paths require membership evidence before attribution.

Run the CPU-only coordinator suite with pytest --confcutdir=tests/test_sharding tests/test_sharding/.

Collection stdout/stderr, including best-effort JIT prewarm warnings, is retained under junit/collection-*/*.log and uploaded with the public CI telemetry. Each collection gets a separate directory so concurrent collectors or later attempts preserve earlier diagnostics. Collection metadata remains temporary.

Communication-test wrappers optionally emit cache-publication completion reports when UNIT_TEST_RUN_STATUS_DIR is set to a fresh directory unique to that wrapper and container/rank. test_run_status.sh records each pytest invocation, and test_run_status.py atomically writes script-result.json (schema version 1, complete, exit_code, tests_executed, and infrastructure_errors). A small pytest plugin records execution and collection/internal errors inside each MPI or torchrun worker; launcher exit 1 alone is not evidence of an assertion failure. All local worker ranks must finish, and at least one test must execute, to allow publication. A completed run with assertion failures remains eligible, while crashes, timeouts, and unfinished invocations do not. With the variable unset, reporting is disabled. Missing JUnit output and file timeouts still fail the wrapper when reporting is disabled. The CPU suite above covers this contract without GPU dependencies.

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.

Attention plan and CUDA Graph lifetime

When auto changes providers, an older captured call can still hold the previous provider's metadata pointers. Keep wrapper-owned tables and length buffers alive and refresh them in plan() whenever replay can still use them; run() must not rebuild or repack metadata. Validate transitions with changed page IDs as well as changed lengths, including replacing an owned table with a caller-supplied table of a different width. Shortening a KV prefix without changing page IDs can accidentally pass while the captured table is stale. A reserved page table can be wider than one graph's declared KV range: declare only that range while preserving the storage row stride and pointer. Include a backend-only provider control in capacity-lifetime tests, since FROST can accept an oversized descriptor that the backend rejects. Table capacity must cover initialized indices and rounded graph bounds; also test a fallback beyond cuDNN admission. An explicit KV maximum may be looser than the caller table's capacity. Intersect these two upper bounds when declaring the graph, using the same effective bound for the attention attribute and table slice; actual device lengths still mask the live tokens, without a host readback. Replay an old capture only with lengths inside its declared bounds and compatible buffers. Keep numerical and lifetime tests independent of heuristic winners and timing thresholds.

Also capture with a caller-supplied table from the first plan, then change Q partitions while auto falls back. Table ownership does not tell whether an old cuDNN capture still needs its wrapper-owned length buffers refreshed.

Valid attention lengths and physical ragged offsets are separate operands: seq_lens_q can be shorter than its query slot. Infer lengths from adjacent offsets only when lengths were omitted. Test explicit/default length transitions against a compacted reference while replaying the original graph; switching mask representations must not strand a captured pointer. Keep token-offset units independent of whether the mask uses lengths or cumulative lengths, and include that distinction in graph descriptors and override shapes. Provider controls need an isolated real graph cache and an actual selected-engine check. For graph-mode ragged wrappers, stage both mask lengths into owned storage from the first cuDNN plan, even when both are inferred. Replan with fresh caller tensors and alternate Q-only, KV-only, both and omitted lengths while replaying the same capture; checking only the latest eager run misses stale bindings.

FA2 captures also bind CPU-generated scheduling metadata, including output and merge prefixes. Updating only Q/KV indptr tensors after switching to cuDNN leaves that schedule stale when requests exchange query rows. In graph mode, keep FA2 once it has been planned; eager auto selection can still change providers. Test FA2-first as well as cuDNN-first with changed Q partitions and page IDs.

Auto dispatch must preserve run-time calibration scales for half-precision as well as FP8 inputs. Test non-unit Q/K scales with nontrivial logits, V scales against O and unchanged LSE, and GPU scalar changes before replay. A plan-time predicate cannot exclude arguments supplied only to run().

Refining auto admission must also preserve the explicit-bound, device-length no-readback path. Declining cuDNN can itself cause FA2 to stage KV metadata back to the host. Use per-request work only when its lengths are already on CPU; otherwise retain the existing bound-based policy. Check with a busy GPU stream and synchronization debugging, not just a predicate-level test.

Full-prefill admission requires each request's Q length to equal its KV length; equal batch maxima do not establish that contract. Prove it using existing CPU metadata, and decline the full-prefill extension when independent query-length overrides or device-only KV lengths prevent that proof. Preserve the existing prefix policy without adding a device readback just for dispatch.

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.

Triggering CI: targeted tests are the default. Public CI: @flashinfer-bot run tests/<dir-or-file> [tests/...]; internal CI: /bot run tests/<dir-or-file> [tests/...]. Targets are whitespace-separated paths under tests/ (files or directories, no :: selectors), written as real paths relative to the repo root, e.g. tests/gemm/test_mm_fp4.py tests/utils/; name the test files or directories the change touches. Bare run asks AI for a recommendation; focused or inconclusive results only post advice. run all starts the expensive full suite, appropriate for a cross-cutting change or final pre-merge check. The PR template has a checkbox to declare and justify that full-suite request.

→ 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.

For an init shared by multiple Const shapes, list the keyword-only parameters in init._trace_init_const_axes to bind their defaults to each dumped definition's resolved Const values. Test the exported snippet with only Var axes supplied so standalone replay does not depend on callers overriding Const defaults.

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: changed DSL/Torch/TVM-FFI distributions, Python ABI, 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. The SHA256 covers only the files a caller passes as extra_key_files (often just __file__); with FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1 it also covers every flashinfer module those files can run. Source analysis and hashing live in flashinfer/jit/_cute_dsl_source.py; cute_dsl_core.py manages compilation and cached artifacts. Plain modules are keyed by content with everything they import, including literal importlib.import_module calls. Package initializers are resolved name by name: a name looked up in a package follows the binding that defines it (a re-export to its source, a function into its body, a module through the attributes read from it) plus the submodule of that name. The code an initializer runs at import, including evaluated function annotations, is followed for every keyed file in its package. Star imports remain covered when they can replace an earlier binding; lazy fallbacks are followed when a binding may be absent. Each initializer is keyed by that code and the import lines of used names, so a new re-export or __all__ entry invalidates nothing. Not keyed: import-time side effects of modules that keyed code neither imports nor uses, nonliteral dynamic imports, and runtime file reads; list the latter in extra_key_files. version.py and _build_meta.py are excluded to avoid keying on the checkout commit. This is intended for long-lived caches shared across commits. Task-scheduled attention shares this cache through flashinfer/attention/prims_ts/_compile_cache.py; its keys cover compile specs, options, and embedded device-capacity constants. It hashes the attention source tree, including shared preparation, masks, and reductions. Occupancy queries are cached per device and cluster size by _hardware.py.

The experimental b12x backend uses flashinfer/experimental/b12x/_lib/compiler.py and stores compiled objects in $B12X_COMPILE_CACHE_DIR, or $XDG_CACHE_HOME/b12x/compile (defaulting to ~/.cache/b12x/compile). Its disk identity includes the callable's transitive local source files, compiler support files, resolved specialization, effective code-generation options, toolchain, and device architecture. AST import tracing includes nested and conditional imports, inherited callables, decorators, and callable defaults. External packages participate through toolchain identity. Raw environment variables and diagnostic output destinations do not participate; environment-selected code generation must be represented by resolved kernel parameters or compiler options. Computed imports in implementation modules declare local targets in __b12x_source_dependencies__. Shared helpers live in leaf modules under _lib/intrinsics/, _lib/utils/, _lib/dense_gemm/, and moe/_shared/kernels/; consumers import the defining module to keep dependency closures narrow. Callable source discovery is cached by function or class, including the owner class for bound methods, and cleared by clear_compile_cache(). Source and compiler configuration are process-immutable: restart after editing code or changing compiler defaults. Warm in-memory hits and prepared CUDA graph replay do not read or parse source files.

b12x autotuning selection keys include the actual compile keys of every eligible candidate, including previously losing candidates. Preparation describes their programs without lowering or launching kernels and reuses those descriptions for compilation. Changes to relevant source, specialization, compiler options, or toolchain require a fresh race. Unrelated kernels retain their saved selections. The tuning contract's schema versions and B12X_TUNING_CACHE_VERSION remain explicit controls for changes to tuning semantics and manual retuning.

Experimental b12x preparation requests can opt into FinalistPolicy to retain measured configurations by tuning-parameter family, independently of the standalone winner. The policy has a separate selection-cache identity. Sharded finalist requests pass every rank's contribution, including empty shards, to PreparationJob.advance(tuning=...) in rank order; ordinary requests pass one consolidated winner per key. Callers prepare alternative plans and capture complete computations as ContextCandidate objects, then call PreparationSession.race_contexts before freezing the session. The returned contextual winner does not replace standalone selections. Input reset and production occur outside the timed replay; producer work that belongs in the measurement must be captured inside the graph. Cold cache eviction follows input production. Multi-rank contexts agree on candidate catalogs, rank membership, sample schedules, producer counts, and cache policy before replay, and score each sample by its slowest rank. Trial graph cleanup callbacks run after synchronization; callers release losing plans before final capture. Contextual races require an unfrozen session with preparation drained.

Every completed b12x tuning race appends a self-describing record to the ledger_v1 table of the race-progress database beside the selection cache: component, encoded query, effort, coverage, measurement protocol, and for every candidate its assignment, lowered config, program keys, the lowest autotune effort that admits it, and per-round timings, plus the batches raced, rejected candidates, and wall seconds per phase. Records are append-only, so retuning accumulates replications. python -m b12x.tools.tuning_ledger [--root DIR] [--records] [--replications] merges tensor-parallel shard records into runs and reports default regret, pruning regret (the default tier's best candidate against an audit winner), runner-up margins, round-to-round noise, and replication agreement.

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_CUTE_DSL_STRICT_CACHE_KEY	0	flashinfer/jit/cute_dsl_core.py	1 extends each CuTe-DSL disk-cache key from the caller's extra_key_files to their in-package import closure, so editing an imported helper invalidates the kernel. Changes every key, so the first run after enabling recompiles. Set by CI's persistent compile cache.
FLASHINFER_JIT_BUILD_LOG	unset	flashinfer/jit/core.py	Path of a JSONL file that gets one record (toolchain, name, seconds) per build made through JitSpec.build_and_load() or a bulk build_jit_specs() prebuild: nvcc modules (ninja no-ops included), CuTe-DSL kernels and TIRx modules. Triton, cuTile and the cake nvcc/NVRTC side paths are not recorded. Summarize with scripts/summarize_jit_builds.py.
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/cake_dsa_sparse_attention_nondet_train.py	Fallback of the check_dst_map keyword of cake_dsa_sparse_attention_nondet / cake_dsa_sparse_attention_nondet_varlen when it is None: unset, empty or 0 leaves the values of a caller-provided dkv_dst_map unchecked; any other value makes the backward validate them (every destination row in [0, S_dst) of dkv_acc) at the cost of one device synchronization and raise ValueError on a violation. The kernel itself does not range-check the map; the package backend reads no environment variable.
FLASHINFER_CAKE_DSA_DKV_DIRECT	auto	flashinfer/cake_dsa_sparse_attention_nondet_train.py	Fallback of the dkv_direct keyword of the two public entries when it is None, read on every call: auto -> None (the record's size rule: at least dkv_direct.min_keys_per_query keys per query token), 1 -> True (direct dK / dV accumulation whenever the natural-layout stages are registered), 0 -> False (the accumulate-and-cast path); any other value raises ValueError. The resolved value is part of the backend's binding key.
FLASHINFER_CAKE_DSA_TRAIN_BINDING_CACHE_CAPACITY	64	flashinfer/cake_dsa_sparse_attention_nondet_train.py	Capacity of the CAKE DSA training binding cache (cake_backend.BINDING_CACHE, least recently used first out): read once when flashinfer.cake_dsa_sparse_attention_nondet_train is imported (a value that is not a positive integer raises ValueError; unset or blank keeps the default) and applied through BINDING_CACHE.configure(capacity=...) when the wrapper first reaches the backend. Direct users of the package call configure themselves.
FLASHINFER_CAKE_LM_HEAD_LOSS_GEMM_TUNING	unset	flashinfer/_cake_lm_head_train_knobs.py	Fallback of the gemm_tuning keyword of flashinfer.cake_lm_head_det_train.cake_lm_head_det_train_loss / cake_lm_head_det_train_logprob (and of the non-deterministic facade's loss / log-probability entries; every row of this table applies to both packages, read through either facade) when it is None: a JSON object of explicit per-GEMM knobs, e.g. {"dx": {"tile_n": 256, "k_slices": 2}} (the package backend takes the same dict as a keyword and reads no environment variable). Unset / empty: the per-geometry rules choose. Invalid JSON or a non-object is a ValueError.
FLASHINFER_CAKE_LM_HEAD_LOSS_FUSE_DW_CAST	1	flashinfer/_cake_lm_head_train_knobs.py	Fallback of the fuse_dw_cast keyword when it is None: 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). dW is bitwise the same either way.
FLASHINFER_CAKE_LM_HEAD_LOSS_DX_FINALIZE	1	flashinfer/_cake_lm_head_train_knobs.py	Fallback of the dx_finalize keyword when it is None: 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 by the slab_sum and scale_cast_scatter_bf16 kernels). Bitwise the same dX; resolved once per forward.
FLASHINFER_CAKE_LM_HEAD_LOSS_DW_STREAM	auto	flashinfer/_cake_lm_head_train_knobs.py	Fallback of the dw_stream keyword when it is None (0 / false / off / no -> never, 1 / true / on / yes -> every multi-chunk call, anything else -> auto = calls of three or more chunks): 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. Same kernels in the same per-kernel order; bitwise the same outputs.
FLASHINFER_CAKE_LM_HEAD_LOSS_HIDDEN_COUNT	1	flashinfer/_cake_lm_head_train_knobs.py	Fallback of the hidden_count keyword when it is None: 0 turns off the hidden valid-row count of the eager entry points (served to calls of three or more chunks): a compacted call then counts its valid rows and forms their index on the host before its first launch instead of on the device with chunk 0's gather kernel and device-count logits GEMM queued ahead of the read-back. Bitwise the same outputs.
FLASHINFER_CAKE_LM_HEAD_LOSS_COMPACT_ROWS	1	flashinfer/_cake_lm_head_train_knobs.py	Fallback of the compact_rows keyword when it is None: 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/_cake_lm_head_train_knobs.py	Read once when each LM-head facade (flashinfer.cake_lm_head_det_train, and the non-deterministic facade) is imported and applied through its package's cake_backend.BINDING_CACHE.configure when that facade first reaches the backend (its first call): 0 disables the per-process LRU cache of compiled bindings. Direct users of the package call configure themselves.
FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE_CAPACITY	64	flashinfer/_cake_lm_head_train_knobs.py	Capacity of that binding cache, read once at the same import (a value that is not a positive integer raises ValueError; unset or blank keeps the default) and applied through configure on the same first use.
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_WHEEL_EXCLUDE_EXPERIMENTAL_SOURCES	unset	build_backend.py (_slim_wheel)	1 leaves the experimental packages' generated kernel sources (flashinfer/experimental/*/csrc/**, csrc/experimental/**) out of the built flashinfer-python wheel and adds flashinfer/experimental/KERNEL_SOURCES_EXCLUDED.txt; the release workflow sets it so the PyPI wheel stays under PyPI's 100 MB file limit (#7977). Nightlies and source installs keep the sources. Experimental APIs still import; their first kernel load raises. Independent of the flag, every built wheel is re-deflated at level 9 and the vendored 3rdparty trees outside the header directories are dropped.
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. They build in the container's torch-pinned interpreter rather than a PEP 517 isolated environment (scripts/build_jit_cache_wheel.py): an isolated environment's random path reaches every compile command through the tvm_ffi include directory and changes the sccache key of each translation unit that includes tvm_ffi on every run. Before compiling, the same script gives every installed header one fixed mtime: GCC's #pragma once skips a header whose size, contents and mtime match an already included file, and torch ships byte-identical #pragma once headers (the ATen/ops/*_native.h stubs), so otherwise the preprocessed output of each torch-including translation unit depends on pip's install timing. 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.

PR GPU test lanes run their test script through scripts/run_with_test_sccache.sh, which routes JIT compiles through the same S3 cache and prefix as the AOT lanes and uploads sccache-stats-<kind>-<cuda>-<index> artifacts. Set the repository variable FLASHINFER_GPU_TEST_SCCACHE=0 to turn this off for A/B timing. If the sccache install fails, the wrapper runs the lane uncached. Compiles run in the sccache server, so a timed-out unit's process-group kill does not reach them, and per-unit memory telemetry excludes nvcc. Locally: scripts/run_with_test_sccache.sh scripts/task_run_unit_tests.sh (read-only without AWS credentials).

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.
FLASHINFER_DISABLE_TINYGEMM2_TCGEN05	0	flashinfer/gemm/routergemm.py	Set to 1 to disable the generated SM100/SM103/SM107 tcgen05 tinygemm2 programs and let the existing tinygemm2 routes handle eligible shapes. Useful as an escape hatch when debugging backend-selection or kernel issues on Blackwell systems.
FLASHINFER_DISABLE_TINYGEMM2_WGMMA	0	flashinfer/gemm/routergemm.py	Set to 1 to disable the generated SM90 wgmma tinygemm2 programs and keep the reference tinygemm2 kernel for every shape. Useful as an escape hatch when debugging backend-selection or kernel issues on Hopper systems.
FLASHINFER_DISABLE_CAKE_BF16_BMM_SWAP_AB	0	flashinfer/gemm/gemm_base.py	Set to 1 to keep bmm_bf16(..., backend="cake") on its mma.sync dispatcher for every shape, instead of the generated SM100/SM103 tcgen05 swap-AB programs that serve the decode shapes of its route table. 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.
FLASHINFER_KERNEL_PACKAGES	unset	flashinfer/__init__.py, flashinfer/trace_apply/packages.py	Comma-separated sha256 digests of sealed kernel packages to verify, import and enable when FlashInfer is imported (best effort: a bad package warns).
FLASHINFER_KERNEL_PACKAGE_ROOT	/opt/flashinfer/kernel_packages	flashinfer/__init__.py, flashinfer/trace_apply/packages.py	Directory holding the digest-named sealed package directories.
FLASHINFER_KERNEL_PLUGINS_DISABLE	unset	flashinfer/trace_apply/packages.py	Kill switch for the environment-triggered loading above: when truthy, the import-time hook refuses with a warning even if the two variables are set. A program that calls load_packages() itself is not affected.

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_TRTLLM_GEN_MULTI_WAVE_KV_SPLIT	unset	include/flashinfer/trtllm/fmha/fmhaKernels.cuh	KV split of trtllm-gen GQA decode (not MLA), read once per process. 1: when one wave of CTAs gives at most 2 KV splits, split KV over up to 4 waves. 2 (for long, mixed-length KV): when one wave gives 1..15 splits, take the largest split <= min(16, tuned maximum) within 8 waves and with ceil(attention window / split) >= 4096, whose waves are at least 4/5 as full as the fullest candidate's; one-wave splits of 9..15 take the 4-wave split capped the same way; base CTAs above the SM count use the 1 rule. Any other value keeps the stock one-wave split. The attention window comes from max_seq_len (or the sliding window), not from the per-request KV lengths, so short or uniform KV under a large bound can be slower. A raised split that does not fit the workspace falls back to the stock split.
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_SM107_MEGA_CHECK_ROUTES	0	flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py	1 enables routing value checks. By default, the caller is responsible for valid routing. Metadata checks and staging remain active. Set before warmup and graph capture; changing it requires recapturing graphs.
FLASHINFER_SM107_FUSED_ROUTE_CHECK	1	flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py	0 uses Torch routing checks instead of the fused checker. When routing checks are enabled, both paths check expert IDs, duplicates, and finite scores, including during graph replay.
FLASHINFER_SM107_PREP_TVM_FFI	1	flashinfer/moe_ep/backends/mega/kernel/sm107/routing_validation.py	0 launches the fused routing checker through CuTe's DLPack bindings instead of TVM-FFI. Set before warmup; each mode needs its own compiled kernel.
FLASHINFER_MEGAMOE_LOG_COMPILE_SPEC	unset	flashinfer/moe_ep/kernel_src/sm120/split_cutedsl_megakernel/shim/runtime.py	Exactly 1 prints MegaMoE split-kernel compile choices, workspace memory diagnostics, and frontend compile/capture timings; other values leave these diagnostics off.
FLASHINFER_MEGAMOE_SPIN_DEBUG	0	flashinfer/jit/megamoe.py	Exactly 1 passes spin_debug=True to the SM90 NVFP4 fused MegaMoE JIT template; other values leave it false.
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_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_SM90_CAKE_BF16_FUSED_TAIL	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's fused_combine_tail keyword when it is None: 0 / false / off / no runs the vendored wait, reduce and ack combine-tail kernels with the per-round combine-inbox fill instead of the one-kernel fused tail. Bit-identical outputs; ranks may differ. The pre-reduced combine wires require the fused tail.
FLASHINFER_SM90_CAKE_BF16_FUSED_DISPATCH	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's fused_dispatch keyword when it is None: 0 / false / off / no keeps the vendored three-kernel dispatch for small rounds instead of the single cooperative dispatch launch. Bit-identical outputs; ranks may differ.
FLASHINFER_SM90_CAKE_BF16_FUSED_DISPATCH_MAX_TOKENS	512	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's fused_dispatch_max_tokens keyword when it is None: dedup rounds with at most this many tokens per rank (and T * top_k <= 4096) use the fused dispatch launch. Must be in [1, 512].
FLASHINFER_SM90_CAKE_BF16_GRID_DISPATCH_MIN_TOKENS	unset (512 at EP 8, 513 at other EP sizes)	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's grid_dispatch_min_tokens keyword when it is None: dedup rounds with at least this many tokens per rank use the grid-strided single-launch dispatch instead of the vendored kernels; 0 disables it. Must be >= 0. Bit-identical outputs; ranks may differ.
FLASHINFER_SM90_CAKE_BF16_PDL	6 (FC1 and FC2)	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's pdl keyword when it is None: which kernels of the round get a programmatic dependent launch. 0 / false / off / no none, 1 / true / on / yes all, or an integer mask of compact 1, FC1 2, FC2 4, combine publish 8, combine tail 16.
FLASHINFER_SM90_CAKE_BF16_PDL_LATE	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's pdl_late keyword when it is None: 1 lets the PDL dependents of the package's small kernels be scheduled only when those kernels exit; 0 / false / off / no triggers them right after their own grid-dependency wait.
FLASHINFER_SM90_CAKE_BF16_PER_ROUTE_PUBLISH_PDL	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's per_route_publish_pdl keyword when it is None: with the per-route combine wire, 1 launches the combine publish with PDL as if bit 8 were in the FLASHINFER_SM90_CAKE_BF16_PDL mask; 0, false, off, no or empty does not. Pipes on the pre-reduced wires are unaffected.
FLASHINFER_SM90_CAKE_BF16_FUSED_BEGIN	2	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's fused_begin keyword when it is None: 0 (or false / off / no) keeps the round opening in its own launches, 1 folds the tag bump and the ack wait into the single-launch dispatch, 2 also folds the receiver prefix. The fused dispatch takes the level when the fused tail, the fused dispatch and dedup dispatch are on (2 falls back to 1 there when E * ep + 2 * ep exceeds 1024 keys); the grid dispatch takes it when the fused tail and dedup dispatch are on. Other values raise ValueError.
FLASHINFER_SM90_CAKE_BF16_COMPACT_FORM	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's compact_form keyword when it is None: 1 the latency-pipelined compact, 0 the one-block-per-row compact. Both write the same bytes; other values raise ValueError.
FLASHINFER_SM90_CAKE_BF16_PUBLISH_FORM	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's publish_form keyword when it is None: 1 the load-hoisted pre-reduced combine publish, 0 the route-serial publish. Both write the same bytes; other values raise ValueError.
FLASHINFER_SM90_CAKE_BF16_FC2_EARLY_MAX_TOKENS	128	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's fc2_early_max_tokens keyword when it is None: rounds with 0 < T <= value tokens per rank start FC2 inside FC1's tail wave (FC1 releases per-m-tile readiness counters, FC2 claims ready tiles; same bits). 0 disables it. It is off when the layer has more than 1023 experts, when min(hidden_size, intermediate_size) < 256, or when the GEMM implementation does not support early start.
FLASHINFER_SM90_CAKE_BF16_FC2_EARLY_WINDOWS	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's fc2_early_windows keyword when it is None: 1 also starts FC2 inside FC1's tail wave in rounds whose token count per rank lies in the window (lo, hi] that FC2_EARLY_WINDOWS lists for the layer's (hidden size, intermediate size, total experts, top_k, EP size); 0, false, off, no or empty keeps only the FLASHINFER_SM90_CAKE_BF16_FC2_EARLY_MAX_TOKENS window. Same bits; off under the same conditions as that window.
FLASHINFER_SM90_CAKE_BF16_WCX_MAX_TOKENS	0	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's wcx_max_tokens keyword when it is None (an int in [0, 128]; unset or empty = 0, off): early-start rounds with 0 < T <= value tokens per rank read the layer's lossless compressed (WCX) expert-weight copies through the swap-AB FC1/FC2 pair; same bits. A non-zero value makes the backend attach the compressed copies (about 0.77x the bf16 expert bytes per layer and rank) and build the swap-AB modules. A non-zero value also makes rounds in the layer shape's WCX_WINDOWS window (lo, hi] above 128 tokens per rank read the copies through the swap-AB early-start pair. Off when the FLASHINFER_SM90_CAKE_BF16_FC2_EARLY_MAX_TOKENS window is off; a layer that cannot be encoded keeps the bf16 pair.
FLASHINFER_SM90_CAKE_BF16_PREFIX_FORM	1	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's prefix_form keyword when it is None: 1 computes the receiver prefix with a whole-block integer scan in the grid dispatch and in the fused dispatch's small-route instance (T * top_k <= 1024); the large-route instance (T * top_k > 1024) walks the keys with one thread for both values. 0 walks the keys with one thread everywhere. Both write the same values; other values raise ValueError.
FLASHINFER_SM90_CAKE_BF16_DIRECT_PLACE	unset (on at EP 8 and EP 4)	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's direct_place keyword when it is None; unset = the per-EP default DIRECT_PLACE_DEFAULTS (on at EP sizes 8 and 4, off at the others): 1 lets grid-dispatch rounds store each token's first row for a rank straight into that rank's (symmetric) activation matrix, so the compact copies only the token's later routes to the same rank. Effective only when every rank requests it and its pipe can run it (the AND over ranks; dedup dispatch, grid dispatch on, compact form 1, at most 1024 (expert, rank) keys, a token capacity of at least the grid-dispatch bound and the direct-placement bound; otherwise the regular activation matrix); grid rounds only, and a round is direct only when every rank dispatches it that way. Bit-identical outputs. Values other than 0 / 1 raise ValueError.
FLASHINFER_SM90_CAKE_BF16_DIRECT_PLACE_MIN_TOKENS	unset (per EP size and shape)	flashinfer/moe_ep/kernel_src/sm90/cake_bf16_megamoe/shim/cake_runner.py	Fallback of the runner's direct_place_min_tokens keyword when it is None (unset = at the EP sizes in DIRECT_PLACE_DEFAULTS the larger of the listed bound 513 and ceil(DIRECT_PLACE_MIN_ROUTE_ELEMS / (top_k * hidden_size)) with DIRECT_PLACE_MIN_ROUTE_ELEMS = 2**25: 586 at hidden 7168 / top-k 8, 781 at 7168 / 6, 1366 at 4096 / 6; else the grid-dispatch bound 513): the smallest local token count of a grid-dispatch round that takes the direct placement (with FLASHINFER_SM90_CAKE_BF16_DIRECT_PLACE); a rank below it makes the whole round take the inbox path. Must be >= 0. Bit-identical outputs.
FLASHINFER_GDN_SM107_PREFILL_KERNEL	sm107	flashinfer/gdn_kernels/blackwell/gdn_prefill.py	sm100 makes chunk_gated_delta_rule on SM107 (Rubin) run the Blackwell chunked-prefill kernel class instead of the Rubin subclass. For A/B measurements and bisection.
FLASHINFER_QKV_EXCHANGE_FOLD_BARRIER	auto	flashinfer/comm/cake_qkv_norm_rope_exchange_overlap.py	Launch path of the fused QKV exchange session: 1 the producer with the folded cross-rank rendezvous (two launches per call), 0 the separate barrier kernel before the producer (three launches), auto the architecture default (folded on SM90 / SM100 / SM103 / SM107). The fold_barrier keyword takes precedence. Every rank of a group must use the same value.
FLASHINFER_QKV_EXCHANGE_SHAPE_GATES	auto	flashinfer/comm/cake_qkv_norm_rope_exchange_overlap.py	Row-count gates of the session's host plan: auto the per-(architecture, world size) default, off every gate off, or gate=value pairs over words_capacity_small_m, reserved_sms_large_m, per_dest_large_m and pair_off_large_m (0 = never, auto = that gate's default, all = every row for words_capacity_small_m). The first three change the host plan only; pair_off_large_m (off everywhere) runs the rows it fires on with the single-CTA producer, which no delivered Blackwell route ships today; SM90 with 8 ranks and SM107 with 2 or 4 ranks ship words_capacity_small_m=all. per_dest_k_table (0, 1 or auto) switches the per-hidden-size per_dest_large_m thresholds of SM100 / SM103 (an explicit per_dest_large_m replaces them); xband_min_m (a row count or auto) sets the row count from which an eligible call runs the 256-column pair tile across band edges (SM100 and SM103 with 8 ranks ship 4096, 0 = off). No gate changes the stored bits. Every rank of a group must use the same value.
FLASHINFER_QKV_EXCHANGE_PRODUCER_VARIANT	auto	flashinfer/comm/cake_qkv_norm_rope_exchange_overlap.py	Producer build of the session: auto the per-(architecture, world size) default (bulk_ow256_fma on SM100 with 8 ranks, ow256_fma on SM100 with 4 ranks and on SM103 with 8 or 4 ranks, fma on SM100 and SM103 with 2 ranks, bulk on SM107 with 4 ranks, the baseline build otherwise), shipped the baseline build, or a delivered build name; a default build that is not delivered for the session's tile width runs the baseline build instead, a named build that is not delivered for it raises and names the delivered ones. The builds without the fma epilogue store identical bits; an fma build contracts the norm and RoPE multiply-adds into single-rounding fused multiply-adds, so its stored bits can differ from theirs within the kernel's accuracy tolerance. Every rank of a group must use the same value.
FLASHINFER_QKV_EXCHANGE_GROUP_M_SCALE	auto	flashinfer/comm/cake_qkv_norm_rope_exchange_overlap.py	Group scaling of the session's host plan: auto the per-(architecture, world size) default (2 on SM100 with 4 or 2 ranks and on SM103, 1 otherwise), or 1 / 2 / 4; applies to grouped rows with a hidden size of at least 6144 and changes the tile order only. Every rank of a group 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, flashinfer/experimental/mhc_pre_big_fuse_with_norm/backend.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, and makes the experimental mhc_pre_big_fuse_with_norm skip its SM100/SM103 cluster kernel and run its composable PyTorch implementation. Read at call time; the kill switch for A/B testing the fused paths.
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, SM100/SM107 packed prefixes (BF16 D192/V128 with equal heads, or FP16/BF16 D128/V128 with integral GQA1/2/4/8/16; 4–64 query heads, B1–4, Q2–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. Q2–128 share the Q128 class; Q=1 retains the separate decode class. 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, satisfied by the package minimum of 4.8.0; 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.

Ulysses PCIe / RDMA Transport (experimental)

Read when UlyssesCommunicator uses the PCIe backend -- named explicitly, or selected by backend="auto" where NVLink is unavailable and FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1 is set. Every rank must set these identically (rank-ordered lists take one comma-separated value per rank).

Variable	Default	Read in	Effect
FLASHINFER_ULYSSES_PCIE_NICS	unset (auto)	flashinfer/comm/ulysses_topology.py	Override automatic PCI-distance NIC routing for the RDMA routes: mlx5 device names, one per rank in rank order.
FLASHINFER_ULYSSES_PCIE_GID_INDICES	unset (auto)	flashinfer/comm/ulysses_topology.py	Pick one GID table index per rank when a chosen NIC has several usable IPv4 RoCE v2 entries.
FLASHINFER_ULYSSES_PCIE_ROUTE	auto	flashinfer/comm/ulysses_topology.py	p2p forces the all-P2P route; rdma forces all-RDMA (per-rank mlx5 to every peer) at world size 2/4/8; hybrid forces the eight-rank 4+4 NUMA hybrid; forced RDMA routes fall back to all-P2P with a RuntimeWarning when their requirements are unmet; auto prefers the all-RDMA route at world sizes 4 and 8.

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

• Portable descriptor locks: _flashinfer_filelock uses stdlib fcntl.flock on Unix and Portalocker's Win32Locker on Windows for shared/exclusive locks. Callers own file opening, timeout policy and cleanup; nonblocking contention raises BlockingIOError. The top-level module keeps standalone tools free of GPU imports. Existing filelock.FileLock users remain unchanged.
• 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
• Wheel post-processing: build_backend.build_wheel rewrites the finished wheel (_slim_wheel): level-9 deflate, the vendored 3rdparty Python trees that packages.find sweeps in are dropped, and with FLASHINFER_WHEEL_EXCLUDE_EXPERIMENTAL_SOURCES=1 (release workflow only) the experimental kernel sources are left out; RECORD is regenerated. Covered by tests/cpu/test_build_backend_slim_wheel.py.

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

• TVM-FFI: Apache TVM's Foreign Function Interface

  • Documentation: https://tvm.apache.org/ffi/
  • Package: apache-tvm-ffi (https://pypi.org/project/apache-tvm-ffi/)
  • Use for: Understanding FFI export syntax, cross-language bindings

• CUTLASS: NVIDIA's CUDA Templates for Linear Algebra Subroutines

  • Recommended: Read source code directly in 3rdparty/cutlass/ (documentation is often outdated)
  • Repository: https://github.com/NVIDIA/cutlass
  • Use for: GEMM kernel implementations, tensor core operations

• CuTe (CUTE DSL): CUTLASS's Cute Layout and Tensor DSL

  • Documentation: https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl.html
  • Tip: Add .md to get Markdown format: https://docs.nvidia.com/cutlass/media/docs/pythonDSL/cute_dsl.html.md
  • The Cute DSL kernels rely on Python modules from the nvidia-cutlass-dsl pip package, not to be confused with Python modules in the 3rdparty/cutlass submodule
  • Tutorial: https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL

• PTX ISA (Parallel Thread Execution): NVIDIA's PTX instruction set documentation

  • Documentation: https://docs.nvidia.com/cuda/parallel-thread-execution/
  • Index/Table of Contents: https://docs.nvidia.com/cuda/parallel-thread-execution/index.html.md
  • Tip: Add .md to any page URL to get Markdown format
  • Use for: Low-level instruction details, new GPU architecture features, inline PTX assembly

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