{"owner":"flashinfer-ai","repo":"flashinfer","hasSkills":true,"totalSkillsCount":5,"totalTokensCount":13736,"categories":["root-instruction","claude-rule","mcp-config","marketplace","plugin-manifest"],"hasMcp":true,"mcpConfig":{"mcpServers":{"flashinfer":{"command":"npx","args":["-y","@modelcontextprotocol/server-flashinfer"]}}},"found":["AGENTS.md","CLAUDE.md",".cursor/mcp.json",".claude-plugin/marketplace.json",".claude-plugin/plugin.json"],"skills":{"AGENTS.md":"This repository’s agent instructions live in [CLAUDE.md](./CLAUDE.md).\n\nBefore making changes, read `CLAUDE.md` and follow its guidance for build, test, style, and contribution workflows.\n","CLAUDE.md":"# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashInfer 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.\n\n## Quick Reference\n\n| Task | Command |\n|------|---------|\n| Install for development | `pip install --no-build-isolation -e . -v` |\n| Initialize submodules | `git submodule update --init --recursive` |\n| Install CUPTI for benchmarking | `pip install -U cupti-python` |\n| Run regular tests | `pytest tests/` |\n| Run full parameter matrices | `pytest tests/ --full` |\n| Run specific test | `pytest tests/path/test_file.py::test_function` |\n| Run multi-GPU test | `mpirun -np 4 pytest tests/comm/test_allreduce_unified_api.py` |\n| Run benchmark | `python benchmarks/flashinfer_benchmark.py --routine <name> <flags>` |\n| Run linting | `pre-commit run -a` |\n| Dump environment report (bug reports) | `python -m flashinfer.collect_env` (or `flashinfer collect-env [--json]`) |\n| Install pre-commit hooks | `pre-commit install` |\n| Clear JIT cache | `rm -rf ~/.cache/flashinfer/` |\n| Inspect devcontainer compiler cache | `sccache --show-stats` |\n| Enable API logging (basic) | `export FLASHINFER_LOGLEVEL=1` |\n| Enable API logging (detailed) | `export FLASHINFER_LOGLEVEL=3` |\n| Enable API logging (with stats) | `export FLASHINFER_LOGLEVEL=5` |\n| Set API log destination | `export FLASHINFER_LOGDEST=mylog.txt` |\n| Enable verbose JIT logging | `export FLASHINFER_JIT_VERBOSE=1` |\n| Enable debug build | `export FLASHINFER_JIT_DEBUG=1` |\n| Set target architectures | `export FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"` |\n| Set parallel compilation | `export FLASHINFER_NVCC_THREADS=4` |\n| Limit parallel ninja jobs | `export MAX_JOBS=4` |\n| Enable GDN native short-T path | `export FLASHINFER_GDN_WY_NATIVE_T=1` |\n| Enable GDN strided QKV path | `export FLASHINFER_GDN_WY_STRIDED_QKV=1` |\n| Enable GDN native A/B tensors | `export FLASHINFER_GDN_WY_NATIVE_AB=1` |\n| 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` |\n| Opt in to multi-wave KV split in trtllm-gen GQA decode (long or mixed-length KV) | `export FLASHINFER_TRTLLM_GEN_MULTI_WAVE_KV_SPLIT=1` |\n| Let `backend=\"auto\"` pick experimental backends | `export FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1` |\n| Override ragged-prefill `auto` backend order (Blackwell) | `export FLASHINFER_RAGGED_AUTO_BACKEND_ORDER=cutlass,cudnn` |\n| Override CuTe-DSL prefill scheduling | `export FLASHINFER_CUTE_PREFILL_PERSISTENT=0` (non-persistent) or `1` (persistent) |\n| Skip MoE EP CuTe-DSL import/version guard | `export FLASHINFER_MOE_EP_SKIP_DSL_CHECK=1` |\n| Override MoE EP knob-cache path | `export FLASHINFER_MOE_EP_KNOB_CACHE=/path/to/knobs.json` |\n| Disable MoE EP fused staging kernel | `export FLASHINFER_MEGA_FUSED_STAGE=0` |\n| Log MegaMoE compilation and workspace diagnostics | `export FLASHINFER_MEGAMOE_LOG_COMPILE_SPEC=1` |\n| Enable MoE EP EXPERT_MAJOR padding-row exclusion (opt-in; shape-dependent) | `export FLASHINFER_MOE_EP_EXCLUDE_PADDING_ROWS=1` |\n| Enable distribution-aware MoE autotune and kernel dispatch (experimental; TRT-LLM MoE only) | `export FLASHINFER_DIST_AWARE_AUTOTUNE=1` |\n| Prims-TS debug checks / task-manager verify | `export FLASHINFER_PRIMS_TS_DEBUG_CHECKS=1` |\n| Override Prims-TS CuTe-DSL compile options | `export FLASHINFER_PRIMS_TS_COMPILE_OPTIONS=\"--opt-level 2\"` |\n| Configure NVFP4 4over6 quantization (preferred over the env vars) | `nvfp4_quantize(..., nvfp4_4over6=NVFP44Over6Config(e4m3_max=256))` |\n\nThe minimum Python version used by CI and build tooling is defined in\n`.python-version`. CI Docker images use a stable Conda environment name and\nderive their interpreter and site-packages paths from that file. When raising\nthe minimum, also update `project.requires-python` and `tool.mypy.python_version`\nin `pyproject.toml`.\n\n## Quick Start for Development\n\n### Installation\n\n```bash\ngit clone https://github.com/flashinfer-ai/flashinfer.git --recursive\ncd flashinfer\npip install --no-build-isolation -e . -v\n```\n\n**Important**: The `--recursive` flag is required to initialize submodules in `3rdparty/` (cutlass, spdlog).\n\nIf you forgot `--recursive` when cloning:\n```bash\ngit submodule update --init --recursive\n```\n\nThat's it! You can now:\n\n- Run all benchmarks and unit tests\n- Modify kernel source code in `include/` without reinstalling\n- Changes are JIT-compiled on next use\n\nThe `--no-build-isolation` flag prevents pip from pulling incompatible PyTorch/CUDA versions from PyPI.\n\n### How JIT Compilation Works\n\nWhen you call a FlashInfer API:\n\n1. **First call**: Generates specialized CUDA code based on parameters (dtype, head_dim, etc.), compiles it with ninja, caches the .so file\n2. **Subsequent calls**: Uses cached compiled module\n3. **After kernel changes**: Automatically detects changes and recompiles\n\n**No manual rebuild step needed** - just edit `.cuh` files and run your code again.\n\n### Pre-compiled Packages (Optional)\n\nFlashInfer provides optional pre-compiled packages for users who want faster initialization:\n\n- `flashinfer-jit-cache`: Pre-built kernel cache\n- `flashinfer-cubin`: Pre-compiled kernel binaries\n\n**For development, you typically DON'T need these.** JIT compilation is fast enough and gives you live code reload.\n\n### Compiler Caching in Devcontainers\n\nThe devcontainers run every nvcc and host C++ compile through\n[sccache](https://github.com/mozilla/sccache) by setting `FLASHINFER_NVCC_LAUNCHER`\nand `FLASHINFER_CXX_LAUNCHER`. Cached objects live in a per-devcontainer Docker\nvolume mounted at `~/.cache/sccache`, so they outlive a container rebuild — which\nwipes `~/.cache/flashinfer` and would otherwise force a full recompile.\n\n```bash\nsccache --show-stats                  # hit rate, cache size, uncacheable reasons\nexport FLASHINFER_NVCC_LAUNCHER=\"\"    # disable for the current shell\n```\n\nCaveats:\n\n- Toggling a launcher changes the compile commands in `build.ninja`, so the first\n  build after switching it on or off recompiles objects that already exist.\n- Link steps, cubin embedding, and CuTe-DSL kernels never invoke nvcc, so they are\n  not cached.\n- The cache is per-devcontainer by design: sccache's local storage supports a\n  single server at a time, so two containers must not share one cache directory.\n\nGitHub PR Test GPU jobs keep a different, per-runner cache: they set\n`CI_COMPILE_CACHE_DIR` to a directory under the runner's work directory, and\n`ci/bash.sh` namespaces it by runner and the pulled image's immutable ID.\n`CI_COMPILE_CACHE_MODE=read-only` (the default) mounts a read-only seed and runs\nagainst a private writable copy, discarded even after test failures. The PR Test\nworkflow selects `read-write` only for pushes to `main`. A stable per-runner lock\nprotects maintenance and execution, including accidentally shared workspaces.\nOther `ci/bash.sh` callers leave caching unset and run cold.\nBefore each trusted writable job its runner's cache is cleared when it exceeds\n`CI_COMPILE_CACHE_MAX_MB`\n(default 40960) or its volume has less than `CI_COMPILE_CACHE_MIN_FREE_MB`\n(default 20480) free. Only content-keyed artifacts survive between jobs: CuTe-DSL (with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1`), TIRx, KDA PTX, Triton and downloaded\ncubins. nvcc modules, tuning results and the generated-source tree are pruned,\nsince a fresh checkout rebuilds nvcc modules anyway. Each test job's\n`FLASHINFER_JIT_BUILD_LOG` is summarized per toolchain by\n`scripts/summarize_jit_builds.py` in the job summary.\nSee [`ci/COMPILE_CACHE.md`](ci/COMPILE_CACHE.md) for lifecycle tests and trust\nboundaries.\n\n## cuDNN linear attention\n\nUse `nvidia-cudnn-frontend>=1.31.0`; KDA and GDN normalization use additive\n`1e-6` without an epsilon argument. Auto selection also requires CuTe DSL 4.7+.\nThe bounded KDA SM100/SM103 and GDN SM100 policies live in `kda.py` and\n`gdn_prefill.py`. KDA's head/length regions are documented in `docs/api/kda.rst`.\nAuto-selected cuDNN ignores the native KDA workspace; native fallback uses it.\nValidate output and final state before changing a policy. Explicit state-pool calls\nrequire unique, in-range slots; auto keeps pools on their native route.\n\nUse `benchmarks/bench_cudnn_linear_attention.py` for public-call comparisons.\nIts raw-gate GDN baseline includes Torch gate transforms; it is not a vLLM\nfused-producer baseline. Keep host CPU, completed wall and reset-inclusive\nCUDA graph times separate.\n\n## Testing\n\nRun the regular test suite:\n\n```bash\npytest tests/\n```\n\nRun full parameter matrices (used by nightly testing):\n\n```bash\npytest tests/ --full\n```\n\nRun specific test file:\n\n```bash\npytest tests/attention/test_hopper.py\n```\n\nRun specific test function:\n\n```bash\npytest tests/attention/test_hopper.py::test_single_prefill\n```\n\nThe sharded test runner defaults to one worker per visible GPU. H100 PR CI opts\ninto `--workers-per-gpu 2`: eight workers on four GPUs. Other CI lanes retain\nthe default. Locally, `scripts/task_run_unit_tests.sh --workers-per-gpu 2` enables\nsharing; `--workers N` sets the total instead. Workers share the existing JIT\ncompilation budget.\n\nScheduling markers apply to the whole source file:\n\n- `node_exclusive`: run alone across the allocation after parallel work, with\n  all GPUs visible. Use for multi-GPU requirements or host-memory OOM risk.\n- `gpu_exclusive`: run alone on the assigned GPU; other GPUs can keep working.\n- `host_memory_heavy`: prioritize this file and serialize its units with other\n  host-heavy work, while lighter files can overlap.\n\nHost and GPU reservations cover setup, execution, and process cleanup, and are\nacquired together. Blocked units stay queued so workers can take runnable work\nfrom their own queue or steal from another. Pending GPU-exclusive work can\ndrain one GPU without reserving the host-heavy slot or blocking other GPUs.\n`RESOURCE WAIT` reports worker idle time of at least one second while work is\nqueued but cannot run on that worker's resources.\n\nThese markers enforce overlap rules, not memory limits. Most H100 annotations\nuse measured peaks above 128 GiB host RSS or 32 GiB GPU memory. Attention sink\nalso uses both narrower markers: its single-GPU H100 run peaked at 89.2 GiB host\nRSS and 4.5 GiB GPU memory, and other GPUs can keep working while it runs.\nCold JIT compilation and other GPU families can change the peaks. Validate\ncompletion rate, OOM failures, and repeated wall times with matching scope and\ncache conditions when changing a CI lane's concurrency policy.\n\nCapacity reports include the serial floor from host-heavy and node-exclusive\nwork. They remain optimistic about GPU contention because the estimator uses\nworker counts without GPU topology or measured memory capacity.\n\nRun the CPU-only coordinator suite with\n`pytest --confcutdir=tests/test_sharding tests/test_sharding/`.\n\nCommunication-test wrappers optionally emit cache-publication completion reports\nwhen `UNIT_TEST_RUN_STATUS_DIR` is set to a fresh directory unique to that wrapper\nand container/rank. `test_run_status.sh` records each pytest invocation, and\n`test_run_status.py` atomically writes `script-result.json` (schema version 1,\n`complete`, `exit_code`, `tests_executed`, and `infrastructure_errors`). A small\npytest plugin records execution and collection/internal errors inside each MPI\nor torchrun worker; launcher exit 1 alone is not evidence of an assertion failure.\nAll local worker ranks must finish, and at least one test must execute, to allow\npublication. A completed run with assertion failures remains eligible, while\ncrashes, timeouts, and unfinished invocations do not. With the variable unset,\nreporting is disabled. The CPU suite above covers this contract without GPU\ndependencies.\n\n### Skipping Tests Based on CUDA Architecture\n\nUse `flashinfer.utils` functions to skip tests on unsupported GPU architectures:\n\n**Available check functions:**\n- `get_compute_capability(device)` - Returns `(major, minor)` tuple\n- `is_sm90a_supported()` - Hopper (requires CUDA 12.3+)\n- `is_sm100a_supported()` - Blackwell (requires CUDA 12.8+)\n- `is_sm100f_supported()` - Blackwell feature-set (requires CUDA 12.9+)\n- `is_sm110a_supported()`, `is_sm120a_supported()`, `is_sm120f_supported()`, `is_sm121a_supported()`\n\n**APIs decorated with `@backend_requirement`** also provide:\n- `api_name.is_compute_capability_supported(cc)` - e.g., `mm_fp4.is_compute_capability_supported(100)`\n- `api_name.is_backend_supported(\"backend\")` - e.g., `mm_fp4.is_backend_supported(\"cudnn\")`\n- `api_name.is_backend_supported(\"backend\", cc)` - e.g., `mm_fp4.is_backend_supported(\"cudnn\", 80)`\n\n**Example:**\n```python\nfrom flashinfer.utils import is_sm90a_supported\n\ndef test_hopper_attention():\n    if not is_sm90a_supported(torch.device(\"cuda\")):\n        pytest.skip(\"Requires SM90a\")\n    # Test code...\n```\n\n**Common requirements:**\n\n| Feature | Min SM | Check Function |\n|---------|--------|----------------|\n| FlashAttention-3 | SM90a | `is_sm90a_supported()` |\n| MLA Attention | SM100a | `is_sm100a_supported()` |\n| FP8 GEMM | SM89+ | `get_compute_capability()[0] >= 9` |\n\n**Note:** `tests/conftest.py` auto-skips tests that trigger OOM, but tests should be written to avoid OOM by using appropriate problem sizes.\n\n## Benchmarking\n\nFor SM90 NVFP4 fused/split and MXFP4 fused MegaMoE, use\n[`benchmarks/bench_megamoe_fp4_sm90.py`](benchmarks/bench_megamoe_fp4_sm90.py) for timing\nand [`tests/moe/test_megamoe_fp4_sm90.py`](tests/moe/test_megamoe_fp4_sm90.py) for correctness checks.\n\nFlashInfer provides a unified benchmarking framework in `benchmarks/flashinfer_benchmark.py`.\n\n**Key features:**\n- Supports attention, GEMM, and MOE kernels\n- Multiple backends: FlashAttention2/3, cuDNN, CUTLASS, TensorRT-LLM, cuBLAS\n- **CUPTI timing (recommended)**: Hardware-level profiling for accurate GPU kernel time\n  - Automatically falls back to CUDA events if CUPTI unavailable\n  - Install: `pip install -U cupti-python` (requires CUDA 13+)\n- Batch testing, reference checking, CSV output\n\n**Quick example:**\n```bash\npython benchmarks/flashinfer_benchmark.py \\\n    --routine BatchDecodeWithPagedKVCacheWrapper \\\n    --backends fa2 cudnn \\\n    --batch_size 32 --s_kv 2048 \\\n    --num_qo_heads 32 --num_kv_heads 8 \\\n    --head_dim_qk 128 --head_dim_vo 128 \\\n    --page_size 16 --refcheck -vv\n```\n\n**Python API:**\n```python\nfrom flashinfer.testing import bench_gpu_time\n\n# CUPTI preferred, auto-fallback to CUDA events.\n# Returns per-iteration times in milliseconds (a list).\ntimes = bench_gpu_time(\n    my_kernel, input_args=(x, y), enable_cupti=True,\n    dry_run_iters=10, repeat_iters=30,\n)\nmedian_time_ms = statistics.median(times)\n```\n\n→ **For complete benchmarking guide, see [`.claude/skills/benchmark-kernel/SKILL.md`](.claude/skills/benchmark-kernel/SKILL.md)**\n\n`flashinfer.gemm.group_gemm_fp8_nt_groupwise_contiguous` is the standalone\nCuTe-DSL grouped FP8 API; `group_deepgemm_fp8_nt_groupwise` keeps its original\nDeepGEMM-only signature. Both use per-row expert indices, unlike the indptr-based\n`group_gemm_fp8_nt_groupwise` API.\n\nFor contiguous grouped FP8 GEMM, `python benchmarks/bench_grouped_fp8.py\n--production-shapes` compares DeepGEMM and CuTe-DSL with preallocated outputs;\n`--cache-probe` measures first-call and new-token-count compilation separately.\nCompiled grouped CuTe-DSL kernels are reused across token counts with the same\n128-row alignment class, with at most two variants per device and weight shape.\nConcurrent misses share one compilation; warm hits do not acquire the compile lock.\nThe standalone API's optional `validate_indices=True` checks expert-index values\nand synchronizes with the CPU. Validate new routing data before CUDA graph\ncapture; the default path assumes valid indices and performs metadata checks only.\n\nPrims-TS decode keeps raw storage dtypes (`q_dtype`, `k_dtype`, `v_dtype`)\nseparate from effective MMA dtypes (`qk_dtype`, `pv_dtype`). Transformed\nFP8/NVFP4 K/V uses Q's compute dtype; BF16 Q/K with FP8 V uses FP8 PV without\nthat transform. Kernel predicates, P packing, and attention-sink scaling must\nfollow effective compute precision; TMA descriptors and raw allocations must\nfollow storage precision. The public plan accepts separate `k_data_type` and\n`v_data_type`, with `kv_data_type` retained as a common-KV compatibility alias.\n\n## Code Linting\n\nRun all pre-commit hooks:\n\n```bash\npre-commit run -a\n```\n\nInstall hooks to run on every commit:\n\n```bash\npre-commit install\n```\n\n## Opening a Pull Request\n\nFollow [`CONTRIBUTING.md`](CONTRIBUTING.md) when opening a PR. The two requirements most often\nmissed:\n\n- **Keep the default PR template.** Fill in `.github/pull_request_template.md` (Description,\n  Related Issues, Checklist, Tests, Reviewer Notes). Do **not** overwrite or replace it with a\n  custom or tool-generated description format. The title and description normally become the\n  commit title and message on squash-merge, and are what a `git bisect` surfaces months later\n  when someone is hunting the owner of a regression.\n- **Report before/after numbers for performance work.** A perf PR must include measurements from\n  a reproducible benchmark (e.g. `benchmarks/flashinfer_benchmark.py`), naming the GPU and the\n  problem sizes. Speedup ratios without absolute numbers, or numbers without a named GPU, are\n  not enough.\n\nWhen preparing changes for a PR, a self-review is strongly recommended: follow\n[`.claude/skills/self-review/SKILL.md`](.claude/skills/self-review/SKILL.md), which walks the\ndiff through the repository's review guidance and PR rules. It is informal and not required —\nits purpose is to make review easier for reviewers and contributors alike by catching common\nissues before a human looks. Changes made only to test something locally, with no PR intended,\ndo not need one.\n\n→ **For the complete contribution rules, see [`CONTRIBUTING.md`](CONTRIBUTING.md)**\n\n## Code Review\n\nWhen reviewing a diff (as an agent or a human), follow the shared focus areas, kernel-review\npolicy, and effort calibration in [`docs/code_review_guidance.md`](docs/code_review_guidance.md).\nNote: unlike human review, agents keep **kernel implementation details in scope** — read the\nkernel logic and report bugs, labeling findings by confidence.\n\n→ **For the complete review rule set, see [`docs/code_review_guidance.md`](docs/code_review_guidance.md)**\n\n## Architecture: JIT Compilation System\n\nFlashInfer's JIT system has three layers:\n\n### Layer 1: JitSpec (flashinfer/jit/core.py)\n\n`JitSpec` is an abstract base class defining the kernel-module lifecycle\n(`try_load()` / `build()` / `load()`, with the shared `build_and_load()`\ntemplate method handling caching, locking, and `FLASHINFER_DISABLE_JIT`).\nOne subclass per compilation toolchain:\n\n- `JitSpecNvcc` (nvcc/ninja modules, returned by `gen_jit_spec()`) defines:\n  - `name`: Unique identifier (URI hash from parameters)\n  - `sources`: List of .cu/.cpp files to compile\n  - `extra_cuda_cflags`, `extra_cflags`, `extra_ldflags`: Compiler flags\n- `JitSpecCuteDsl` (flashinfer/jit/cute_dsl_core.py) caches CuTe-DSL kernels\n  (see \"CuTe-DSL kernels\" under Module Caching below)\n- The optional TIRx KDA backend uses a private `JitSpec` in\n  `flashinfer/kda_kernels/tirx/cache.py` to cache TVM modules, with source/compiler\n  fingerprints and artifact checksums. Its TVM/TIRx dependencies are loaded only\n  by explicit `recurrent_kda(..., backend=\"tirx\")`; see `docs/api/kda.rst`.\n\n### JIT Directory Rules\n\n**NEVER write to package directories** - they may be read-only after installation.\n\n| Directory | Writable | Use for |\n|-----------|----------|---------|\n| `FLASHINFER_GEN_SRC_DIR` | ✓ Yes | Generated source files (Jinja output, copied .cu files) |\n| `FLASHINFER_JIT_DIR` | ✓ Yes | Compiled `.so` outputs |\n| `FLASHINFER_CSRC_DIR` | ✗ No | Read-only source templates |\n| `FLASHINFER_AOT_DIR` | ✗ No | Read-only pre-compiled binaries |\n\n### Compilation Context: Architecture-Specific Compilation\n\nFlashInfer uses `CompilationContext` to manage CUDA architecture targets. Some kernels only work on specific GPU architectures (e.g., Hopper SM90, Blackwell SM100/SM12x).\n\n**How it works:**\n- Auto-detects GPUs in system or reads `FLASHINFER_CUDA_ARCH_LIST` environment variable\n- CuTe-DSL FP4 MM compilation parallelism can be controlled with `FLASHINFER_MM_FP4_CUTE_DSL_COMPILE_WORKERS`\n- JIT modules specify `supported_major_versions=[9, 10, 11, 12]` to limit compilation to specific SM versions\n- If GPU not supported → `RuntimeError: No supported CUDA architectures found`\n\n→ **See [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md) for usage examples**\n\n### Layer 2: Code Generation\n\nEvery `gen_*_module()` function in `flashinfer/jit/` follows this pattern:\n\n```python\ndef gen_some_module(dtype_in, dtype_out, ...):\n    # 1. Compute unique identifier from parameters\n    uri = get_some_uri(dtype_in, dtype_out, ...)\n\n    # 2. Create generation directory\n    gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri\n\n    # 3. (Optional) Render Jinja template to generate type-specialized config\n    # Skip this step if you don't need type specialization\n    with open(jit_env.FLASHINFER_CSRC_DIR / \"some_customize_config.jinja\") as f:\n        template = jinja2.Template(f.read())\n    config_content = template.render(\n        dtype_in=dtype_map[dtype_in],\n        dtype_out=dtype_map[dtype_out],\n        # ... more parameters\n    )\n    write_if_different(gen_directory / \"some_config.inc\", config_content)\n\n    # 4. Copy source files to gen directory\n    sources = []\n    for fname in [\"some_kernel.cu\", \"some_jit_binding.cu\"]:\n        shutil.copy(jit_env.FLASHINFER_CSRC_DIR / fname, gen_directory / fname)\n        sources.append(gen_directory / fname)\n\n    # 5. Return JitSpec\n    return gen_jit_spec(uri, sources, extra_cuda_cflags=[...])\n```\n\n**Note**: If your operation doesn't need type specialization, you can skip step 3 entirely and just copy the source files directly.\n\n### Layer 3: Compilation and Loading\n\n`JitSpecNvcc` methods:\n\n- `write_ninja()` - Generates `build.ninja` file\n- `build()` - Executes `ninja` to compile sources\n- `build_and_load()` - Compiles and loads via TVM-FFI\n\nThe generated `build.ninja` file uses nvcc to compile .cu → .cuda.o → .so, then loads via TVM-FFI.\n\n### Jinja Templates (Optional)\n\n**Note: Jinja templates are NOT required.** You can write C++ code directly without templating.\n\nFor operations that need type specialization, templates in `csrc/*.jinja` can generate C++ code:\n\n```jinja\n// Input template\nusing DTypeIn = {{ dtype_in }};\nusing DTypeOut = {{ dtype_out }};\nconstexpr int PARAM = {{ param_value }};\n\n// After render\nusing DTypeIn = float16;\nusing DTypeOut = float16;\nconstexpr int PARAM = 128;\n```\n\nThis 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.\n\n## Directory Structure\n\n```\nflashinfer/\n├── include/flashinfer/           # Header-only CUDA kernel templates\n│   ├── attention/                # Attention kernels\n│   ├── gemm/                     # GEMM kernels\n│   ├── comm/                     # Communication kernels\n│   ├── mma.cuh                   # Matrix multiply utilities\n│   ├── utils.cuh                 # Common utilities\n│   └── [...]\n│\n├── csrc/                          # Framework bindings (via TVM-FFI)\n│   ├── *.cu                       # Kernel launcher implementations\n│   ├── *_jit_binding.cu           # TVM-FFI exports\n│   ├── *_customize_config.jinja   # Type config templates (optional)\n│   └── [...]\n│\n├── flashinfer/                    # Python package\n│   ├── jit/\n│   │   ├── core.py                # JitSpec, compilation infrastructure\n│   │   ├── cpp_ext.py             # Ninja build generation\n│   │   ├── env.py                 # Workspace paths\n│   │   ├── attention/             # Attention module generators\n│   │   ├── gemm/                  # GEMM module generators\n│   │   ├── fused_moe/             # MOE module generators\n│   │   └── [...]\n│   ├── gemm/                      # GEMM Python APIs\n│   ├── fused_moe/                 # MOE Python APIs\n│   ├── comm/                      # Communication Python APIs\n│   ├── *.py                       # Other high-level Python APIs\n│   ├── aot.py                     # AOT compilation for pre-built packages\n│   └── [...]\n│\n├── tests/                         # Test suite\n│   ├── attention/                 # Attention kernel tests\n│   ├── gemm/                      # GEMM kernel tests\n│   ├── moe/                       # MOE kernel tests\n│   ├── comm/                      # Communication tests\n│   ├── utils/                     # Utility tests\n│   └── conftest.py                # Pytest configuration\n│\n└── build_backend.py               # PEP 517 build backend\n```\n\n### Critical Rule: Framework Separation\n\n**Torch headers MUST NOT be included in `include/` directory files.**\n\n- `include/`: Framework-agnostic CUDA kernels (accept raw pointers)\n- `csrc/`: Framework bindings via TVM-FFI (currently PyTorch, but can support other frameworks)\n\n## Adding a New Operation\n\n→ **For complete step-by-step tutorial, see [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md)**\n\n**Quick overview of the process:**\n1. Write kernel in `include/flashinfer/new_op.cuh` (framework-agnostic, raw pointers)\n2. Write launcher in `csrc/new_op.cu` (PyTorch tensor handling)\n3. Create TVM-FFI bindings in `csrc/new_op_jit_binding.cu`\n4. (Optional) Create Jinja template for type specialization\n5. Write JIT module generator in `flashinfer/jit/new_op.py`\n6. Write Python API in `flashinfer/new_op.py` with `@functools.cache`\n7. Write tests in `tests/`\n8. Register in `flashinfer/aot.py` for AOT compilation\n9. Export in `flashinfer/__init__.py`\n10. Add a `TraceTemplate` in `flashinfer/trace/templates/` and wire it via `@flashinfer_api(trace=...)` (see below)\n11. Add an example call in `tests/trace/example.py`, re-run to regenerate `fi_trace_out/`, and commit the new JSON files\n\n### Trace Template Checklist (for new or updated APIs)\n\nEvery 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.\n\n1. **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.\n2. **Wire the template** to the API: `@flashinfer_api(trace=my_trace)` on the Python function (or class method's `run()`).\n3. **Add an example call** in `tests/trace/example.py` that exercises the new trace with realistic shapes.\n4. **Regenerate examples**: `rm -rf tests/trace/fi_trace_out && python tests/trace/example.py` — verify the expected JSON appears.\n5. **Update the docstring** in `tests/trace/example.py` to list the new file(s).\n6. **Run tests**: `pytest tests/trace/ -v` — all template-consistency and end-to-end tests must pass.\n7. **Commit the new JSON files** under `tests/trace/fi_trace_out/` alongside the code changes.\n\n**Example implementations:**\n- **Simple**: `flashinfer/norm/__init__.py` (RMSNorm) - no Jinja, good starting point\n- **Moderate**: `flashinfer/sampling.py` - with Jinja templating\n- **Complex**: `flashinfer/decode.py` - plan-run pattern, advanced workspace\n\n## Experimental Code\n\nExperimental backends live under `flashinfer/experimental/`; experimental\npublic APIs live in core marked with `@flashinfer_experimental_api` (defined in\n`flashinfer/api_logging.py`). Using them is an explicit opt-in: calling an\nexperimental API, or naming an experimental backend with `backend=\"<name>\"`,\nneeds no environment variable (both warn once). Only automatic selection is\ngated: `backend=\"auto\"` skips backends marked `@experimental_backend` unless\n`FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. Experimental code is JIT-only\n(never registered in `flashinfer/aot.py`) and tested under\n`tests/experimental/`.\n\n→ **When working under `flashinfer/experimental/`, follow\n[`flashinfer/experimental/CLAUDE.md`](flashinfer/experimental/CLAUDE.md); the\nfull policy is [`flashinfer/experimental/README.md`](flashinfer/experimental/README.md).**\n\n## Key Architectural Patterns\n\n### Module Caching\n\nFlashInfer uses two-level caching to avoid recompilation:\n\n1. **Python-level** (`@functools.cache`): In-memory cache of loaded modules\n2. **File-level** (`~/.cache/flashinfer/`): Compiled `.so` files on disk\n\n**Cache invalidation** (automatic):\n- Source file changes (SHA256 hash)\n- Compilation flags change\n- CUDA architecture change\n- FlashInfer version change\n\nURI computed as: `hash(operation_type + parameters + source_hashes + flags + cuda_arch)`\n\n**Cache management:**\n- Clear cache: `rm -rf ~/.cache/flashinfer/`\n- Override location: `export FLASHINFER_WORKSPACE_BASE=\"/scratch\"`\n\n**CuTe-DSL kernels** (`flashinfer/jit/cute_dsl_core.py`): `cute.compile()` has no\npersistent cache, so `JitSpecCuteDsl` (a `JitSpec` subclass, wrapped by the\n`build_and_load_cute_dsl_kernel()` helper) exports compiled kernels as\nobject files (`export_to_c()`) and reloads them with `cute.runtime.load_module(...,\nenable_tvm_ffi=True)`. Artifacts live in `cached_ops/` next to the nvcc modules, one\ndirectory per op family — `cached_ops/<module>_<arch>_cute_dsl/` (e.g.\n`nvfp4_quantize_sm100a_cute_dsl/`) holding one `meta.json` plus one `.o` per\nspecialization. The arch comes from the DSL's compile target (`CUTE_DSL_ARCH` or the\ncurrent device) since the artifacts are single-arch, unlike nvcc fatbins.\nInvalidation is module-granular: changed DSL/Torch/TVM-FFI distributions, Python\nABI, or\nkernel-source SHA256 wipes and lazily rebuilds the module. Reference usage:\n`flashinfer/quantization/kernels/nvfp4_quantize.py`. Disable with\n`FLASHINFER_CUTE_DSL_DISABLE_CACHE=1`. The SHA256 covers only the files a caller\npasses as `extra_key_files` (often just `__file__`); with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1` it also covers every flashinfer module\nthose files can run. Source analysis and hashing live in\n`flashinfer/jit/_cute_dsl_source.py`; `cute_dsl_core.py` manages compilation and\ncached artifacts. Plain modules are keyed by content with everything they\nimport, including literal `importlib.import_module` calls. Package initializers\nare resolved name by name: a name looked up in a package follows the binding\nthat defines it (a re-export to its source, a function into its body, a module\nthrough the attributes read from it) plus the submodule of that name. The code\nan initializer runs at import, including evaluated function annotations, is\nfollowed for every keyed file in its package. Star imports remain covered when\nthey can replace an earlier binding; lazy fallbacks are followed when a binding\nmay be absent.\nEach initializer is keyed by that code and the import lines of used names, so a\nnew re-export or `__all__` entry invalidates nothing. Not keyed: import-time\nside effects of modules that keyed code neither imports nor uses, nonliteral\ndynamic imports, and runtime file reads; list the latter in `extra_key_files`.\n`version.py` and `_build_meta.py` are excluded to avoid keying on the checkout\ncommit. This is intended for long-lived caches shared across commits.\nTask-scheduled attention shares this cache through\n`flashinfer/attention/prims_ts/_compile_cache.py`; its keys cover compile specs,\noptions, and embedded device-capacity constants. It hashes the attention source\ntree, including shared preparation, masks, and reductions. Occupancy queries\nare cached per device and cluster size by `_hardware.py`.\n\n### Dispatch Macros\n\nHandle combinatorial parameter spaces:\n\n```cpp\nDISPATCH_DTYPE(input_dtype, DTypeIn, {\n  DISPATCH_DTYPE(output_dtype, DTypeOut, {\n    DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {\n      LaunchKernel<DTypeIn, DTypeOut, BLOCK_SIZE>(...);\n    });\n  });\n});\n```\n\nDefined in `.jinja` files and expanded after rendering.\n\n## API Logging with @flashinfer_api\n\nFlashInfer provides the `@flashinfer_api` decorator for debugging API calls.\n\n**Key features:**\n- **Crash-safe**: Logs inputs BEFORE execution (preserves info even if kernel crashes)\n- **Zero overhead when disabled**: `FLASHINFER_LOGLEVEL=0` (default)\n- **Multiple verbosity levels**: 0 (off), 1 (names), 3 (inputs/outputs), 5 (+ statistics)\n- **CUDA graph compatible**: Auto-skips stats during graph capture\n\n**Quick usage:**\n```bash\n# Enable detailed logging\nexport FLASHINFER_LOGLEVEL=3              # 0, 1, 3, or 5\nexport FLASHINFER_LOGDEST=debug.log       # stdout, stderr, or file path\n\npython my_script.py\n```\n\n**Why use this?**\n- Debug CUDA crashes (see inputs that caused crash)\n- Track tensor shapes/dtypes through pipeline\n- Detect NaN/Inf issues (level 5)\n\n→ **For complete debugging guide, see [`.claude/skills/debug-cuda-crash/SKILL.md`](.claude/skills/debug-cuda-crash/SKILL.md)**\n\n## Debugging\n\n### Enable Logging\n\n```bash\nexport FLASHINFER_JIT_VERBOSE=1      # Verbose JIT output\nexport FLASHINFER_JIT_DEBUG=1        # Debug symbols, -O0\nexport FLASHINFER_LOGLEVEL=3         # API logging (0=off, 1=basic, 3=detailed)\nexport FLASHINFER_LOGDEST=stdout\n```\n\n### Inspect Generated Code\n\n```bash\n# Generated sources\nls -la ~/.cache/flashinfer/0.6.0/*/generated/\n\n# Compiled modules\nls -la ~/.cache/flashinfer/0.6.0/*/cached_ops/\n\n# Build files\ncat ~/.cache/flashinfer/0.6.0/*/cached_ops/*/build.ninja\n```\n\n### Environment Variables\n\n```bash\n# Compilation\nexport FLASHINFER_NVCC_THREADS=4              # Threads per nvcc process (--threads=N)\nexport MAX_JOBS=4                             # Parallel ninja jobs (nvcc processes)\nexport FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"  # Target architectures\n# Memory note: total compilation memory ≈ MAX_JOBS × FLASHINFER_NVCC_THREADS × per-thread mem.\n\n# Behavior\nexport FLASHINFER_WORKSPACE_BASE=\"/scratch\"   # Custom cache directory\n```\n\n#### Full Environment Variable Reference\n\nAlready covered above (Quick Reference / Debugging): `FLASHINFER_LOGLEVEL`,\n`FLASHINFER_LOGDEST`, `FLASHINFER_JIT_VERBOSE`, `FLASHINFER_JIT_DEBUG`,\n`FLASHINFER_CUDA_ARCH_LIST`, `FLASHINFER_NVCC_THREADS`,\n`FLASHINFER_WORKSPACE_BASE`, `MAX_JOBS`.\n\nThe remaining `FLASHINFER_*` knobs read by the code are listed below. Defaults\nmatch what the code uses today; values are strings unless noted.\n\n##### JIT / Build Toolchain\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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`). |\n| `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. |\n| `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`. |\n| `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. |\n| `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. |\n| `FLASHINFER_JIT_LINEINFO` | `0` | `flashinfer/jit/core.py` | `1` adds `-lineinfo` to nvcc so profiler / `cuda-gdb` can map PTX back to CUDA source. |\n| `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. |\n| `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). |\n| `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`. |\n| `FLASHINFER_CXX_LAUNCHER` | `\"\"` (`sccache` in the devcontainers) | `flashinfer/jit/cpp_ext.py` | Same idea as `FLASHINFER_NVCC_LAUNCHER` but for the host C++ compiler. |\n| `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. |\n| `FLASHINFER_EXTRA_CFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to the host C++ compiler. |\n| `FLASHINFER_EXTRA_CUDAFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to `nvcc`. |\n| `FLASHINFER_EXTRA_LDFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra linker flags passed to the linker. |\n| `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. |\n| `FLASHINFER_CAKE_DSA_CHECK_DST_MAP` | `0` | `flashinfer/dsa_sparse_attention.py` | Fallback of the `check_dst_map` keyword of `dsa_sparse_attention` / `dsa_sparse_attention_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. |\n| `FLASHINFER_CAKE_DSA_DKV_DIRECT` | `auto` | `flashinfer/dsa_sparse_attention.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. |\n| `FLASHINFER_CAKE_DSA_TRAIN_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/dsa_sparse_attention.py` | Capacity of the CAKE DSA training binding cache (`cake_backend.BINDING_CACHE`, least recently used first out): read once when `flashinfer.dsa_sparse_attention` 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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_GEMM_TUNING` | unset | `flashinfer/chunked_lm_head.py` | Fallback of the `gemm_tuning` keyword of `chunked_lm_head_loss` / `chunked_lm_head_logprob` 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`. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_FUSE_DW_CAST` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DX_FINALIZE` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DW_STREAM` | `auto` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_HIDDEN_COUNT` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_COMPACT_ROWS` | `1` | `flashinfer/chunked_lm_head.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`). |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE` | `1` | `flashinfer/chunked_lm_head.py` | Read once when `flashinfer.chunked_lm_head` is imported and applied through `cake_backend.BINDING_CACHE.configure` when the wrapper first reaches the backend (the first call of either entry point): `0` disables the per-process LRU cache of compiled bindings. Direct users of the package call `configure` themselves. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/chunked_lm_head.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. |\n| `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. |\n| `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. |\n\nRelease/nightly JIT-cache provider builds run sccache in server-side mode and\nwrap the complete wheel build with a no-output watchdog. After 90 minutes with\nno build output, the watchdog captures system, process, compiler, and sccache\nstate; terminates the wrapped build process group (TERM, then KILL after a\ntwo-minute grace period); and exits with status 124. Compiler discovery scans\nthe whole build container, so diagnostics include compiler processes owned by\nthe sccache daemon even though they are outside the wrapped process group. The\nCI `docker run --rm` boundary provides final cleanup for those processes.\n\n##### Cubin / Artifact Loader\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_CUBIN_CHECKSUM_DISABLED` | unset | `flashinfer/jit/cubin_loader.py` | If set, skip SHA checksum verification of downloaded cubins. Debug aid only. |\n| `FLASHINFER_CUBIN_DOWNLOAD_THREADS` | `4` | `flashinfer/artifacts.py` | Thread-pool size used by `flashinfer artifacts download`. |\n| `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. |\n| `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. |\n| `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. |\n| `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`. |\n| `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. |\n\n##### API Dump / Logging Extensions\n\nThese complement `FLASHINFER_LOGLEVEL`/`FLASHINFER_LOGDEST` and are all read in `flashinfer/api_logging.py`.\n\n| Variable | Default | Effect |\n|----------|---------|--------|\n| `FLASHINFER_DUMP_DIR` | `flashinfer_dumps` | Directory where `@flashinfer_api` writes per-call tensor dumps when logging is enabled. |\n| `FLASHINFER_DUMP_INCLUDE` | `\"\"` | Comma-separated allow-list of API names; only matching calls are dumped. Wildcards (`*`) supported. |\n| `FLASHINFER_DUMP_EXCLUDE` | `\"\"` | Comma-separated deny-list of API names; matching calls skip the (expensive) stats / dump path. |\n| `FLASHINFER_DUMP_MAX_COUNT` | `1000` | Hard cap on number of dumped events; once reached, further dumps are dropped with a warning. |\n| `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`). |\n| `FLASHINFER_DUMP_SAFETENSORS` | `0` | `1` writes dumps as `.safetensors` (portable / Hugging Face-style); otherwise plain `.pt`. |\n\n##### Trace Capture\n\nUsed by `flashinfer.trace` / `fi_trace`.\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `FLASHINFER_TRACE_DUMP` | unset | `flashinfer/fi_trace.py` | If set, decorated APIs auto-dump benchmark-definition JSON for each call (the \"fi_trace\" feature). |\n| `FLASHINFER_TRACE_DUMP_DIR` | cwd | `flashinfer/fi_trace.py` | Directory where the trace JSON files are written. |\n| `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. |\n| `FLASHINFER_TRACE_APPLY_PATH` | unset | `flashinfer/trace_apply/config.py` | Directory from which deployment-configured solutions are loaded for Trace Apply. |\n\n##### Validation / Autotuning / Routing / Kernel Selection\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_AUTOTUNER_LOAD_FROM_FILE` | `0` | `flashinfer/autotuner/autotuner.py` | `1` loads previously serialized autotune results from disk instead of re-running the search. |\n| `FLASHINFER_DIST_AWARE_AUTOTUNE` | `0` | `flashinfer/fuse",".cursor/mcp.json":"# Model Context Protocol (MCP) Configuration\n\nPath: `.cursor/mcp.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.cursor/mcp.json)",".claude-plugin/marketplace.json":"# Claude Plugin Marketplace Catalog\n\nPath: `.claude-plugin/marketplace.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/marketplace.json)",".claude-plugin/plugin.json":"# claude-plugin Plugin Manifest\n\nPath: `.claude-plugin/plugin.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/plugin.json)"},"files":{"AGENTS.md":"This repository’s agent instructions live in [CLAUDE.md](./CLAUDE.md).\n\nBefore making changes, read `CLAUDE.md` and follow its guidance for build, test, style, and contribution workflows.\n","CLAUDE.md":"# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashInfer 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.\n\n## Quick Reference\n\n| Task | Command |\n|------|---------|\n| Install for development | `pip install --no-build-isolation -e . -v` |\n| Initialize submodules | `git submodule update --init --recursive` |\n| Install CUPTI for benchmarking | `pip install -U cupti-python` |\n| Run regular tests | `pytest tests/` |\n| Run full parameter matrices | `pytest tests/ --full` |\n| Run specific test | `pytest tests/path/test_file.py::test_function` |\n| Run multi-GPU test | `mpirun -np 4 pytest tests/comm/test_allreduce_unified_api.py` |\n| Run benchmark | `python benchmarks/flashinfer_benchmark.py --routine <name> <flags>` |\n| Run linting | `pre-commit run -a` |\n| Dump environment report (bug reports) | `python -m flashinfer.collect_env` (or `flashinfer collect-env [--json]`) |\n| Install pre-commit hooks | `pre-commit install` |\n| Clear JIT cache | `rm -rf ~/.cache/flashinfer/` |\n| Inspect devcontainer compiler cache | `sccache --show-stats` |\n| Enable API logging (basic) | `export FLASHINFER_LOGLEVEL=1` |\n| Enable API logging (detailed) | `export FLASHINFER_LOGLEVEL=3` |\n| Enable API logging (with stats) | `export FLASHINFER_LOGLEVEL=5` |\n| Set API log destination | `export FLASHINFER_LOGDEST=mylog.txt` |\n| Enable verbose JIT logging | `export FLASHINFER_JIT_VERBOSE=1` |\n| Enable debug build | `export FLASHINFER_JIT_DEBUG=1` |\n| Set target architectures | `export FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"` |\n| Set parallel compilation | `export FLASHINFER_NVCC_THREADS=4` |\n| Limit parallel ninja jobs | `export MAX_JOBS=4` |\n| Enable GDN native short-T path | `export FLASHINFER_GDN_WY_NATIVE_T=1` |\n| Enable GDN strided QKV path | `export FLASHINFER_GDN_WY_STRIDED_QKV=1` |\n| Enable GDN native A/B tensors | `export FLASHINFER_GDN_WY_NATIVE_AB=1` |\n| 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` |\n| Opt in to multi-wave KV split in trtllm-gen GQA decode (long or mixed-length KV) | `export FLASHINFER_TRTLLM_GEN_MULTI_WAVE_KV_SPLIT=1` |\n| Let `backend=\"auto\"` pick experimental backends | `export FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1` |\n| Override ragged-prefill `auto` backend order (Blackwell) | `export FLASHINFER_RAGGED_AUTO_BACKEND_ORDER=cutlass,cudnn` |\n| Override CuTe-DSL prefill scheduling | `export FLASHINFER_CUTE_PREFILL_PERSISTENT=0` (non-persistent) or `1` (persistent) |\n| Skip MoE EP CuTe-DSL import/version guard | `export FLASHINFER_MOE_EP_SKIP_DSL_CHECK=1` |\n| Override MoE EP knob-cache path | `export FLASHINFER_MOE_EP_KNOB_CACHE=/path/to/knobs.json` |\n| Disable MoE EP fused staging kernel | `export FLASHINFER_MEGA_FUSED_STAGE=0` |\n| Log MegaMoE compilation and workspace diagnostics | `export FLASHINFER_MEGAMOE_LOG_COMPILE_SPEC=1` |\n| Enable MoE EP EXPERT_MAJOR padding-row exclusion (opt-in; shape-dependent) | `export FLASHINFER_MOE_EP_EXCLUDE_PADDING_ROWS=1` |\n| Enable distribution-aware MoE autotune and kernel dispatch (experimental; TRT-LLM MoE only) | `export FLASHINFER_DIST_AWARE_AUTOTUNE=1` |\n| Prims-TS debug checks / task-manager verify | `export FLASHINFER_PRIMS_TS_DEBUG_CHECKS=1` |\n| Override Prims-TS CuTe-DSL compile options | `export FLASHINFER_PRIMS_TS_COMPILE_OPTIONS=\"--opt-level 2\"` |\n| Configure NVFP4 4over6 quantization (preferred over the env vars) | `nvfp4_quantize(..., nvfp4_4over6=NVFP44Over6Config(e4m3_max=256))` |\n\nThe minimum Python version used by CI and build tooling is defined in\n`.python-version`. CI Docker images use a stable Conda environment name and\nderive their interpreter and site-packages paths from that file. When raising\nthe minimum, also update `project.requires-python` and `tool.mypy.python_version`\nin `pyproject.toml`.\n\n## Quick Start for Development\n\n### Installation\n\n```bash\ngit clone https://github.com/flashinfer-ai/flashinfer.git --recursive\ncd flashinfer\npip install --no-build-isolation -e . -v\n```\n\n**Important**: The `--recursive` flag is required to initialize submodules in `3rdparty/` (cutlass, spdlog).\n\nIf you forgot `--recursive` when cloning:\n```bash\ngit submodule update --init --recursive\n```\n\nThat's it! You can now:\n\n- Run all benchmarks and unit tests\n- Modify kernel source code in `include/` without reinstalling\n- Changes are JIT-compiled on next use\n\nThe `--no-build-isolation` flag prevents pip from pulling incompatible PyTorch/CUDA versions from PyPI.\n\n### How JIT Compilation Works\n\nWhen you call a FlashInfer API:\n\n1. **First call**: Generates specialized CUDA code based on parameters (dtype, head_dim, etc.), compiles it with ninja, caches the .so file\n2. **Subsequent calls**: Uses cached compiled module\n3. **After kernel changes**: Automatically detects changes and recompiles\n\n**No manual rebuild step needed** - just edit `.cuh` files and run your code again.\n\n### Pre-compiled Packages (Optional)\n\nFlashInfer provides optional pre-compiled packages for users who want faster initialization:\n\n- `flashinfer-jit-cache`: Pre-built kernel cache\n- `flashinfer-cubin`: Pre-compiled kernel binaries\n\n**For development, you typically DON'T need these.** JIT compilation is fast enough and gives you live code reload.\n\n### Compiler Caching in Devcontainers\n\nThe devcontainers run every nvcc and host C++ compile through\n[sccache](https://github.com/mozilla/sccache) by setting `FLASHINFER_NVCC_LAUNCHER`\nand `FLASHINFER_CXX_LAUNCHER`. Cached objects live in a per-devcontainer Docker\nvolume mounted at `~/.cache/sccache`, so they outlive a container rebuild — which\nwipes `~/.cache/flashinfer` and would otherwise force a full recompile.\n\n```bash\nsccache --show-stats                  # hit rate, cache size, uncacheable reasons\nexport FLASHINFER_NVCC_LAUNCHER=\"\"    # disable for the current shell\n```\n\nCaveats:\n\n- Toggling a launcher changes the compile commands in `build.ninja`, so the first\n  build after switching it on or off recompiles objects that already exist.\n- Link steps, cubin embedding, and CuTe-DSL kernels never invoke nvcc, so they are\n  not cached.\n- The cache is per-devcontainer by design: sccache's local storage supports a\n  single server at a time, so two containers must not share one cache directory.\n\nGitHub PR Test GPU jobs keep a different, per-runner cache: they set\n`CI_COMPILE_CACHE_DIR` to a directory under the runner's work directory, and\n`ci/bash.sh` namespaces it by runner and the pulled image's immutable ID.\n`CI_COMPILE_CACHE_MODE=read-only` (the default) mounts a read-only seed and runs\nagainst a private writable copy, discarded even after test failures. The PR Test\nworkflow selects `read-write` only for pushes to `main`. A stable per-runner lock\nprotects maintenance and execution, including accidentally shared workspaces.\nOther `ci/bash.sh` callers leave caching unset and run cold.\nBefore each trusted writable job its runner's cache is cleared when it exceeds\n`CI_COMPILE_CACHE_MAX_MB`\n(default 40960) or its volume has less than `CI_COMPILE_CACHE_MIN_FREE_MB`\n(default 20480) free. Only content-keyed artifacts survive between jobs: CuTe-DSL (with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1`), TIRx, KDA PTX, Triton and downloaded\ncubins. nvcc modules, tuning results and the generated-source tree are pruned,\nsince a fresh checkout rebuilds nvcc modules anyway. Each test job's\n`FLASHINFER_JIT_BUILD_LOG` is summarized per toolchain by\n`scripts/summarize_jit_builds.py` in the job summary.\nSee [`ci/COMPILE_CACHE.md`](ci/COMPILE_CACHE.md) for lifecycle tests and trust\nboundaries.\n\n## cuDNN linear attention\n\nUse `nvidia-cudnn-frontend>=1.31.0`; KDA and GDN normalization use additive\n`1e-6` without an epsilon argument. Auto selection also requires CuTe DSL 4.7+.\nThe bounded KDA SM100/SM103 and GDN SM100 policies live in `kda.py` and\n`gdn_prefill.py`. KDA's head/length regions are documented in `docs/api/kda.rst`.\nAuto-selected cuDNN ignores the native KDA workspace; native fallback uses it.\nValidate output and final state before changing a policy. Explicit state-pool calls\nrequire unique, in-range slots; auto keeps pools on their native route.\n\nUse `benchmarks/bench_cudnn_linear_attention.py` for public-call comparisons.\nIts raw-gate GDN baseline includes Torch gate transforms; it is not a vLLM\nfused-producer baseline. Keep host CPU, completed wall and reset-inclusive\nCUDA graph times separate.\n\n## Testing\n\nRun the regular test suite:\n\n```bash\npytest tests/\n```\n\nRun full parameter matrices (used by nightly testing):\n\n```bash\npytest tests/ --full\n```\n\nRun specific test file:\n\n```bash\npytest tests/attention/test_hopper.py\n```\n\nRun specific test function:\n\n```bash\npytest tests/attention/test_hopper.py::test_single_prefill\n```\n\nThe sharded test runner defaults to one worker per visible GPU. H100 PR CI opts\ninto `--workers-per-gpu 2`: eight workers on four GPUs. Other CI lanes retain\nthe default. Locally, `scripts/task_run_unit_tests.sh --workers-per-gpu 2` enables\nsharing; `--workers N` sets the total instead. Workers share the existing JIT\ncompilation budget.\n\nScheduling markers apply to the whole source file:\n\n- `node_exclusive`: run alone across the allocation after parallel work, with\n  all GPUs visible. Use for multi-GPU requirements or host-memory OOM risk.\n- `gpu_exclusive`: run alone on the assigned GPU; other GPUs can keep working.\n- `host_memory_heavy`: prioritize this file and serialize its units with other\n  host-heavy work, while lighter files can overlap.\n\nHost and GPU reservations cover setup, execution, and process cleanup, and are\nacquired together. Blocked units stay queued so workers can take runnable work\nfrom their own queue or steal from another. Pending GPU-exclusive work can\ndrain one GPU without reserving the host-heavy slot or blocking other GPUs.\n`RESOURCE WAIT` reports worker idle time of at least one second while work is\nqueued but cannot run on that worker's resources.\n\nThese markers enforce overlap rules, not memory limits. Most H100 annotations\nuse measured peaks above 128 GiB host RSS or 32 GiB GPU memory. Attention sink\nalso uses both narrower markers: its single-GPU H100 run peaked at 89.2 GiB host\nRSS and 4.5 GiB GPU memory, and other GPUs can keep working while it runs.\nCold JIT compilation and other GPU families can change the peaks. Validate\ncompletion rate, OOM failures, and repeated wall times with matching scope and\ncache conditions when changing a CI lane's concurrency policy.\n\nCapacity reports include the serial floor from host-heavy and node-exclusive\nwork. They remain optimistic about GPU contention because the estimator uses\nworker counts without GPU topology or measured memory capacity.\n\nRun the CPU-only coordinator suite with\n`pytest --confcutdir=tests/test_sharding tests/test_sharding/`.\n\nCommunication-test wrappers optionally emit cache-publication completion reports\nwhen `UNIT_TEST_RUN_STATUS_DIR` is set to a fresh directory unique to that wrapper\nand container/rank. `test_run_status.sh` records each pytest invocation, and\n`test_run_status.py` atomically writes `script-result.json` (schema version 1,\n`complete`, `exit_code`, `tests_executed`, and `infrastructure_errors`). A small\npytest plugin records execution and collection/internal errors inside each MPI\nor torchrun worker; launcher exit 1 alone is not evidence of an assertion failure.\nAll local worker ranks must finish, and at least one test must execute, to allow\npublication. A completed run with assertion failures remains eligible, while\ncrashes, timeouts, and unfinished invocations do not. With the variable unset,\nreporting is disabled. The CPU suite above covers this contract without GPU\ndependencies.\n\n### Skipping Tests Based on CUDA Architecture\n\nUse `flashinfer.utils` functions to skip tests on unsupported GPU architectures:\n\n**Available check functions:**\n- `get_compute_capability(device)` - Returns `(major, minor)` tuple\n- `is_sm90a_supported()` - Hopper (requires CUDA 12.3+)\n- `is_sm100a_supported()` - Blackwell (requires CUDA 12.8+)\n- `is_sm100f_supported()` - Blackwell feature-set (requires CUDA 12.9+)\n- `is_sm110a_supported()`, `is_sm120a_supported()`, `is_sm120f_supported()`, `is_sm121a_supported()`\n\n**APIs decorated with `@backend_requirement`** also provide:\n- `api_name.is_compute_capability_supported(cc)` - e.g., `mm_fp4.is_compute_capability_supported(100)`\n- `api_name.is_backend_supported(\"backend\")` - e.g., `mm_fp4.is_backend_supported(\"cudnn\")`\n- `api_name.is_backend_supported(\"backend\", cc)` - e.g., `mm_fp4.is_backend_supported(\"cudnn\", 80)`\n\n**Example:**\n```python\nfrom flashinfer.utils import is_sm90a_supported\n\ndef test_hopper_attention():\n    if not is_sm90a_supported(torch.device(\"cuda\")):\n        pytest.skip(\"Requires SM90a\")\n    # Test code...\n```\n\n**Common requirements:**\n\n| Feature | Min SM | Check Function |\n|---------|--------|----------------|\n| FlashAttention-3 | SM90a | `is_sm90a_supported()` |\n| MLA Attention | SM100a | `is_sm100a_supported()` |\n| FP8 GEMM | SM89+ | `get_compute_capability()[0] >= 9` |\n\n**Note:** `tests/conftest.py` auto-skips tests that trigger OOM, but tests should be written to avoid OOM by using appropriate problem sizes.\n\n## Benchmarking\n\nFor SM90 NVFP4 fused/split and MXFP4 fused MegaMoE, use\n[`benchmarks/bench_megamoe_fp4_sm90.py`](benchmarks/bench_megamoe_fp4_sm90.py) for timing\nand [`tests/moe/test_megamoe_fp4_sm90.py`](tests/moe/test_megamoe_fp4_sm90.py) for correctness checks.\n\nFlashInfer provides a unified benchmarking framework in `benchmarks/flashinfer_benchmark.py`.\n\n**Key features:**\n- Supports attention, GEMM, and MOE kernels\n- Multiple backends: FlashAttention2/3, cuDNN, CUTLASS, TensorRT-LLM, cuBLAS\n- **CUPTI timing (recommended)**: Hardware-level profiling for accurate GPU kernel time\n  - Automatically falls back to CUDA events if CUPTI unavailable\n  - Install: `pip install -U cupti-python` (requires CUDA 13+)\n- Batch testing, reference checking, CSV output\n\n**Quick example:**\n```bash\npython benchmarks/flashinfer_benchmark.py \\\n    --routine BatchDecodeWithPagedKVCacheWrapper \\\n    --backends fa2 cudnn \\\n    --batch_size 32 --s_kv 2048 \\\n    --num_qo_heads 32 --num_kv_heads 8 \\\n    --head_dim_qk 128 --head_dim_vo 128 \\\n    --page_size 16 --refcheck -vv\n```\n\n**Python API:**\n```python\nfrom flashinfer.testing import bench_gpu_time\n\n# CUPTI preferred, auto-fallback to CUDA events.\n# Returns per-iteration times in milliseconds (a list).\ntimes = bench_gpu_time(\n    my_kernel, input_args=(x, y), enable_cupti=True,\n    dry_run_iters=10, repeat_iters=30,\n)\nmedian_time_ms = statistics.median(times)\n```\n\n→ **For complete benchmarking guide, see [`.claude/skills/benchmark-kernel/SKILL.md`](.claude/skills/benchmark-kernel/SKILL.md)**\n\n`flashinfer.gemm.group_gemm_fp8_nt_groupwise_contiguous` is the standalone\nCuTe-DSL grouped FP8 API; `group_deepgemm_fp8_nt_groupwise` keeps its original\nDeepGEMM-only signature. Both use per-row expert indices, unlike the indptr-based\n`group_gemm_fp8_nt_groupwise` API.\n\nFor contiguous grouped FP8 GEMM, `python benchmarks/bench_grouped_fp8.py\n--production-shapes` compares DeepGEMM and CuTe-DSL with preallocated outputs;\n`--cache-probe` measures first-call and new-token-count compilation separately.\nCompiled grouped CuTe-DSL kernels are reused across token counts with the same\n128-row alignment class, with at most two variants per device and weight shape.\nConcurrent misses share one compilation; warm hits do not acquire the compile lock.\nThe standalone API's optional `validate_indices=True` checks expert-index values\nand synchronizes with the CPU. Validate new routing data before CUDA graph\ncapture; the default path assumes valid indices and performs metadata checks only.\n\nPrims-TS decode keeps raw storage dtypes (`q_dtype`, `k_dtype`, `v_dtype`)\nseparate from effective MMA dtypes (`qk_dtype`, `pv_dtype`). Transformed\nFP8/NVFP4 K/V uses Q's compute dtype; BF16 Q/K with FP8 V uses FP8 PV without\nthat transform. Kernel predicates, P packing, and attention-sink scaling must\nfollow effective compute precision; TMA descriptors and raw allocations must\nfollow storage precision. The public plan accepts separate `k_data_type` and\n`v_data_type`, with `kv_data_type` retained as a common-KV compatibility alias.\n\n## Code Linting\n\nRun all pre-commit hooks:\n\n```bash\npre-commit run -a\n```\n\nInstall hooks to run on every commit:\n\n```bash\npre-commit install\n```\n\n## Opening a Pull Request\n\nFollow [`CONTRIBUTING.md`](CONTRIBUTING.md) when opening a PR. The two requirements most often\nmissed:\n\n- **Keep the default PR template.** Fill in `.github/pull_request_template.md` (Description,\n  Related Issues, Checklist, Tests, Reviewer Notes). Do **not** overwrite or replace it with a\n  custom or tool-generated description format. The title and description normally become the\n  commit title and message on squash-merge, and are what a `git bisect` surfaces months later\n  when someone is hunting the owner of a regression.\n- **Report before/after numbers for performance work.** A perf PR must include measurements from\n  a reproducible benchmark (e.g. `benchmarks/flashinfer_benchmark.py`), naming the GPU and the\n  problem sizes. Speedup ratios without absolute numbers, or numbers without a named GPU, are\n  not enough.\n\nWhen preparing changes for a PR, a self-review is strongly recommended: follow\n[`.claude/skills/self-review/SKILL.md`](.claude/skills/self-review/SKILL.md), which walks the\ndiff through the repository's review guidance and PR rules. It is informal and not required —\nits purpose is to make review easier for reviewers and contributors alike by catching common\nissues before a human looks. Changes made only to test something locally, with no PR intended,\ndo not need one.\n\n→ **For the complete contribution rules, see [`CONTRIBUTING.md`](CONTRIBUTING.md)**\n\n## Code Review\n\nWhen reviewing a diff (as an agent or a human), follow the shared focus areas, kernel-review\npolicy, and effort calibration in [`docs/code_review_guidance.md`](docs/code_review_guidance.md).\nNote: unlike human review, agents keep **kernel implementation details in scope** — read the\nkernel logic and report bugs, labeling findings by confidence.\n\n→ **For the complete review rule set, see [`docs/code_review_guidance.md`](docs/code_review_guidance.md)**\n\n## Architecture: JIT Compilation System\n\nFlashInfer's JIT system has three layers:\n\n### Layer 1: JitSpec (flashinfer/jit/core.py)\n\n`JitSpec` is an abstract base class defining the kernel-module lifecycle\n(`try_load()` / `build()` / `load()`, with the shared `build_and_load()`\ntemplate method handling caching, locking, and `FLASHINFER_DISABLE_JIT`).\nOne subclass per compilation toolchain:\n\n- `JitSpecNvcc` (nvcc/ninja modules, returned by `gen_jit_spec()`) defines:\n  - `name`: Unique identifier (URI hash from parameters)\n  - `sources`: List of .cu/.cpp files to compile\n  - `extra_cuda_cflags`, `extra_cflags`, `extra_ldflags`: Compiler flags\n- `JitSpecCuteDsl` (flashinfer/jit/cute_dsl_core.py) caches CuTe-DSL kernels\n  (see \"CuTe-DSL kernels\" under Module Caching below)\n- The optional TIRx KDA backend uses a private `JitSpec` in\n  `flashinfer/kda_kernels/tirx/cache.py` to cache TVM modules, with source/compiler\n  fingerprints and artifact checksums. Its TVM/TIRx dependencies are loaded only\n  by explicit `recurrent_kda(..., backend=\"tirx\")`; see `docs/api/kda.rst`.\n\n### JIT Directory Rules\n\n**NEVER write to package directories** - they may be read-only after installation.\n\n| Directory | Writable | Use for |\n|-----------|----------|---------|\n| `FLASHINFER_GEN_SRC_DIR` | ✓ Yes | Generated source files (Jinja output, copied .cu files) |\n| `FLASHINFER_JIT_DIR` | ✓ Yes | Compiled `.so` outputs |\n| `FLASHINFER_CSRC_DIR` | ✗ No | Read-only source templates |\n| `FLASHINFER_AOT_DIR` | ✗ No | Read-only pre-compiled binaries |\n\n### Compilation Context: Architecture-Specific Compilation\n\nFlashInfer uses `CompilationContext` to manage CUDA architecture targets. Some kernels only work on specific GPU architectures (e.g., Hopper SM90, Blackwell SM100/SM12x).\n\n**How it works:**\n- Auto-detects GPUs in system or reads `FLASHINFER_CUDA_ARCH_LIST` environment variable\n- CuTe-DSL FP4 MM compilation parallelism can be controlled with `FLASHINFER_MM_FP4_CUTE_DSL_COMPILE_WORKERS`\n- JIT modules specify `supported_major_versions=[9, 10, 11, 12]` to limit compilation to specific SM versions\n- If GPU not supported → `RuntimeError: No supported CUDA architectures found`\n\n→ **See [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md) for usage examples**\n\n### Layer 2: Code Generation\n\nEvery `gen_*_module()` function in `flashinfer/jit/` follows this pattern:\n\n```python\ndef gen_some_module(dtype_in, dtype_out, ...):\n    # 1. Compute unique identifier from parameters\n    uri = get_some_uri(dtype_in, dtype_out, ...)\n\n    # 2. Create generation directory\n    gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri\n\n    # 3. (Optional) Render Jinja template to generate type-specialized config\n    # Skip this step if you don't need type specialization\n    with open(jit_env.FLASHINFER_CSRC_DIR / \"some_customize_config.jinja\") as f:\n        template = jinja2.Template(f.read())\n    config_content = template.render(\n        dtype_in=dtype_map[dtype_in],\n        dtype_out=dtype_map[dtype_out],\n        # ... more parameters\n    )\n    write_if_different(gen_directory / \"some_config.inc\", config_content)\n\n    # 4. Copy source files to gen directory\n    sources = []\n    for fname in [\"some_kernel.cu\", \"some_jit_binding.cu\"]:\n        shutil.copy(jit_env.FLASHINFER_CSRC_DIR / fname, gen_directory / fname)\n        sources.append(gen_directory / fname)\n\n    # 5. Return JitSpec\n    return gen_jit_spec(uri, sources, extra_cuda_cflags=[...])\n```\n\n**Note**: If your operation doesn't need type specialization, you can skip step 3 entirely and just copy the source files directly.\n\n### Layer 3: Compilation and Loading\n\n`JitSpecNvcc` methods:\n\n- `write_ninja()` - Generates `build.ninja` file\n- `build()` - Executes `ninja` to compile sources\n- `build_and_load()` - Compiles and loads via TVM-FFI\n\nThe generated `build.ninja` file uses nvcc to compile .cu → .cuda.o → .so, then loads via TVM-FFI.\n\n### Jinja Templates (Optional)\n\n**Note: Jinja templates are NOT required.** You can write C++ code directly without templating.\n\nFor operations that need type specialization, templates in `csrc/*.jinja` can generate C++ code:\n\n```jinja\n// Input template\nusing DTypeIn = {{ dtype_in }};\nusing DTypeOut = {{ dtype_out }};\nconstexpr int PARAM = {{ param_value }};\n\n// After render\nusing DTypeIn = float16;\nusing DTypeOut = float16;\nconstexpr int PARAM = 128;\n```\n\nThis 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.\n\n## Directory Structure\n\n```\nflashinfer/\n├── include/flashinfer/           # Header-only CUDA kernel templates\n│   ├── attention/                # Attention kernels\n│   ├── gemm/                     # GEMM kernels\n│   ├── comm/                     # Communication kernels\n│   ├── mma.cuh                   # Matrix multiply utilities\n│   ├── utils.cuh                 # Common utilities\n│   └── [...]\n│\n├── csrc/                          # Framework bindings (via TVM-FFI)\n│   ├── *.cu                       # Kernel launcher implementations\n│   ├── *_jit_binding.cu           # TVM-FFI exports\n│   ├── *_customize_config.jinja   # Type config templates (optional)\n│   └── [...]\n│\n├── flashinfer/                    # Python package\n│   ├── jit/\n│   │   ├── core.py                # JitSpec, compilation infrastructure\n│   │   ├── cpp_ext.py             # Ninja build generation\n│   │   ├── env.py                 # Workspace paths\n│   │   ├── attention/             # Attention module generators\n│   │   ├── gemm/                  # GEMM module generators\n│   │   ├── fused_moe/             # MOE module generators\n│   │   └── [...]\n│   ├── gemm/                      # GEMM Python APIs\n│   ├── fused_moe/                 # MOE Python APIs\n│   ├── comm/                      # Communication Python APIs\n│   ├── *.py                       # Other high-level Python APIs\n│   ├── aot.py                     # AOT compilation for pre-built packages\n│   └── [...]\n│\n├── tests/                         # Test suite\n│   ├── attention/                 # Attention kernel tests\n│   ├── gemm/                      # GEMM kernel tests\n│   ├── moe/                       # MOE kernel tests\n│   ├── comm/                      # Communication tests\n│   ├── utils/                     # Utility tests\n│   └── conftest.py                # Pytest configuration\n│\n└── build_backend.py               # PEP 517 build backend\n```\n\n### Critical Rule: Framework Separation\n\n**Torch headers MUST NOT be included in `include/` directory files.**\n\n- `include/`: Framework-agnostic CUDA kernels (accept raw pointers)\n- `csrc/`: Framework bindings via TVM-FFI (currently PyTorch, but can support other frameworks)\n\n## Adding a New Operation\n\n→ **For complete step-by-step tutorial, see [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md)**\n\n**Quick overview of the process:**\n1. Write kernel in `include/flashinfer/new_op.cuh` (framework-agnostic, raw pointers)\n2. Write launcher in `csrc/new_op.cu` (PyTorch tensor handling)\n3. Create TVM-FFI bindings in `csrc/new_op_jit_binding.cu`\n4. (Optional) Create Jinja template for type specialization\n5. Write JIT module generator in `flashinfer/jit/new_op.py`\n6. Write Python API in `flashinfer/new_op.py` with `@functools.cache`\n7. Write tests in `tests/`\n8. Register in `flashinfer/aot.py` for AOT compilation\n9. Export in `flashinfer/__init__.py`\n10. Add a `TraceTemplate` in `flashinfer/trace/templates/` and wire it via `@flashinfer_api(trace=...)` (see below)\n11. Add an example call in `tests/trace/example.py`, re-run to regenerate `fi_trace_out/`, and commit the new JSON files\n\n### Trace Template Checklist (for new or updated APIs)\n\nEvery 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.\n\n1. **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.\n2. **Wire the template** to the API: `@flashinfer_api(trace=my_trace)` on the Python function (or class method's `run()`).\n3. **Add an example call** in `tests/trace/example.py` that exercises the new trace with realistic shapes.\n4. **Regenerate examples**: `rm -rf tests/trace/fi_trace_out && python tests/trace/example.py` — verify the expected JSON appears.\n5. **Update the docstring** in `tests/trace/example.py` to list the new file(s).\n6. **Run tests**: `pytest tests/trace/ -v` — all template-consistency and end-to-end tests must pass.\n7. **Commit the new JSON files** under `tests/trace/fi_trace_out/` alongside the code changes.\n\n**Example implementations:**\n- **Simple**: `flashinfer/norm/__init__.py` (RMSNorm) - no Jinja, good starting point\n- **Moderate**: `flashinfer/sampling.py` - with Jinja templating\n- **Complex**: `flashinfer/decode.py` - plan-run pattern, advanced workspace\n\n## Experimental Code\n\nExperimental backends live under `flashinfer/experimental/`; experimental\npublic APIs live in core marked with `@flashinfer_experimental_api` (defined in\n`flashinfer/api_logging.py`). Using them is an explicit opt-in: calling an\nexperimental API, or naming an experimental backend with `backend=\"<name>\"`,\nneeds no environment variable (both warn once). Only automatic selection is\ngated: `backend=\"auto\"` skips backends marked `@experimental_backend` unless\n`FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. Experimental code is JIT-only\n(never registered in `flashinfer/aot.py`) and tested under\n`tests/experimental/`.\n\n→ **When working under `flashinfer/experimental/`, follow\n[`flashinfer/experimental/CLAUDE.md`](flashinfer/experimental/CLAUDE.md); the\nfull policy is [`flashinfer/experimental/README.md`](flashinfer/experimental/README.md).**\n\n## Key Architectural Patterns\n\n### Module Caching\n\nFlashInfer uses two-level caching to avoid recompilation:\n\n1. **Python-level** (`@functools.cache`): In-memory cache of loaded modules\n2. **File-level** (`~/.cache/flashinfer/`): Compiled `.so` files on disk\n\n**Cache invalidation** (automatic):\n- Source file changes (SHA256 hash)\n- Compilation flags change\n- CUDA architecture change\n- FlashInfer version change\n\nURI computed as: `hash(operation_type + parameters + source_hashes + flags + cuda_arch)`\n\n**Cache management:**\n- Clear cache: `rm -rf ~/.cache/flashinfer/`\n- Override location: `export FLASHINFER_WORKSPACE_BASE=\"/scratch\"`\n\n**CuTe-DSL kernels** (`flashinfer/jit/cute_dsl_core.py`): `cute.compile()` has no\npersistent cache, so `JitSpecCuteDsl` (a `JitSpec` subclass, wrapped by the\n`build_and_load_cute_dsl_kernel()` helper) exports compiled kernels as\nobject files (`export_to_c()`) and reloads them with `cute.runtime.load_module(...,\nenable_tvm_ffi=True)`. Artifacts live in `cached_ops/` next to the nvcc modules, one\ndirectory per op family — `cached_ops/<module>_<arch>_cute_dsl/` (e.g.\n`nvfp4_quantize_sm100a_cute_dsl/`) holding one `meta.json` plus one `.o` per\nspecialization. The arch comes from the DSL's compile target (`CUTE_DSL_ARCH` or the\ncurrent device) since the artifacts are single-arch, unlike nvcc fatbins.\nInvalidation is module-granular: changed DSL/Torch/TVM-FFI distributions, Python\nABI, or\nkernel-source SHA256 wipes and lazily rebuilds the module. Reference usage:\n`flashinfer/quantization/kernels/nvfp4_quantize.py`. Disable with\n`FLASHINFER_CUTE_DSL_DISABLE_CACHE=1`. The SHA256 covers only the files a caller\npasses as `extra_key_files` (often just `__file__`); with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1` it also covers every flashinfer module\nthose files can run. Source analysis and hashing live in\n`flashinfer/jit/_cute_dsl_source.py`; `cute_dsl_core.py` manages compilation and\ncached artifacts. Plain modules are keyed by content with everything they\nimport, including literal `importlib.import_module` calls. Package initializers\nare resolved name by name: a name looked up in a package follows the binding\nthat defines it (a re-export to its source, a function into its body, a module\nthrough the attributes read from it) plus the submodule of that name. The code\nan initializer runs at import, including evaluated function annotations, is\nfollowed for every keyed file in its package. Star imports remain covered when\nthey can replace an earlier binding; lazy fallbacks are followed when a binding\nmay be absent.\nEach initializer is keyed by that code and the import lines of used names, so a\nnew re-export or `__all__` entry invalidates nothing. Not keyed: import-time\nside effects of modules that keyed code neither imports nor uses, nonliteral\ndynamic imports, and runtime file reads; list the latter in `extra_key_files`.\n`version.py` and `_build_meta.py` are excluded to avoid keying on the checkout\ncommit. This is intended for long-lived caches shared across commits.\nTask-scheduled attention shares this cache through\n`flashinfer/attention/prims_ts/_compile_cache.py`; its keys cover compile specs,\noptions, and embedded device-capacity constants. It hashes the attention source\ntree, including shared preparation, masks, and reductions. Occupancy queries\nare cached per device and cluster size by `_hardware.py`.\n\n### Dispatch Macros\n\nHandle combinatorial parameter spaces:\n\n```cpp\nDISPATCH_DTYPE(input_dtype, DTypeIn, {\n  DISPATCH_DTYPE(output_dtype, DTypeOut, {\n    DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {\n      LaunchKernel<DTypeIn, DTypeOut, BLOCK_SIZE>(...);\n    });\n  });\n});\n```\n\nDefined in `.jinja` files and expanded after rendering.\n\n## API Logging with @flashinfer_api\n\nFlashInfer provides the `@flashinfer_api` decorator for debugging API calls.\n\n**Key features:**\n- **Crash-safe**: Logs inputs BEFORE execution (preserves info even if kernel crashes)\n- **Zero overhead when disabled**: `FLASHINFER_LOGLEVEL=0` (default)\n- **Multiple verbosity levels**: 0 (off), 1 (names), 3 (inputs/outputs), 5 (+ statistics)\n- **CUDA graph compatible**: Auto-skips stats during graph capture\n\n**Quick usage:**\n```bash\n# Enable detailed logging\nexport FLASHINFER_LOGLEVEL=3              # 0, 1, 3, or 5\nexport FLASHINFER_LOGDEST=debug.log       # stdout, stderr, or file path\n\npython my_script.py\n```\n\n**Why use this?**\n- Debug CUDA crashes (see inputs that caused crash)\n- Track tensor shapes/dtypes through pipeline\n- Detect NaN/Inf issues (level 5)\n\n→ **For complete debugging guide, see [`.claude/skills/debug-cuda-crash/SKILL.md`](.claude/skills/debug-cuda-crash/SKILL.md)**\n\n## Debugging\n\n### Enable Logging\n\n```bash\nexport FLASHINFER_JIT_VERBOSE=1      # Verbose JIT output\nexport FLASHINFER_JIT_DEBUG=1        # Debug symbols, -O0\nexport FLASHINFER_LOGLEVEL=3         # API logging (0=off, 1=basic, 3=detailed)\nexport FLASHINFER_LOGDEST=stdout\n```\n\n### Inspect Generated Code\n\n```bash\n# Generated sources\nls -la ~/.cache/flashinfer/0.6.0/*/generated/\n\n# Compiled modules\nls -la ~/.cache/flashinfer/0.6.0/*/cached_ops/\n\n# Build files\ncat ~/.cache/flashinfer/0.6.0/*/cached_ops/*/build.ninja\n```\n\n### Environment Variables\n\n```bash\n# Compilation\nexport FLASHINFER_NVCC_THREADS=4              # Threads per nvcc process (--threads=N)\nexport MAX_JOBS=4                             # Parallel ninja jobs (nvcc processes)\nexport FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"  # Target architectures\n# Memory note: total compilation memory ≈ MAX_JOBS × FLASHINFER_NVCC_THREADS × per-thread mem.\n\n# Behavior\nexport FLASHINFER_WORKSPACE_BASE=\"/scratch\"   # Custom cache directory\n```\n\n#### Full Environment Variable Reference\n\nAlready covered above (Quick Reference / Debugging): `FLASHINFER_LOGLEVEL`,\n`FLASHINFER_LOGDEST`, `FLASHINFER_JIT_VERBOSE`, `FLASHINFER_JIT_DEBUG`,\n`FLASHINFER_CUDA_ARCH_LIST`, `FLASHINFER_NVCC_THREADS`,\n`FLASHINFER_WORKSPACE_BASE`, `MAX_JOBS`.\n\nThe remaining `FLASHINFER_*` knobs read by the code are listed below. Defaults\nmatch what the code uses today; values are strings unless noted.\n\n##### JIT / Build Toolchain\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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`). |\n| `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. |\n| `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`. |\n| `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. |\n| `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. |\n| `FLASHINFER_JIT_LINEINFO` | `0` | `flashinfer/jit/core.py` | `1` adds `-lineinfo` to nvcc so profiler / `cuda-gdb` can map PTX back to CUDA source. |\n| `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. |\n| `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). |\n| `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`. |\n| `FLASHINFER_CXX_LAUNCHER` | `\"\"` (`sccache` in the devcontainers) | `flashinfer/jit/cpp_ext.py` | Same idea as `FLASHINFER_NVCC_LAUNCHER` but for the host C++ compiler. |\n| `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. |\n| `FLASHINFER_EXTRA_CFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to the host C++ compiler. |\n| `FLASHINFER_EXTRA_CUDAFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to `nvcc`. |\n| `FLASHINFER_EXTRA_LDFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra linker flags passed to the linker. |\n| `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. |\n| `FLASHINFER_CAKE_DSA_CHECK_DST_MAP` | `0` | `flashinfer/dsa_sparse_attention.py` | Fallback of the `check_dst_map` keyword of `dsa_sparse_attention` / `dsa_sparse_attention_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. |\n| `FLASHINFER_CAKE_DSA_DKV_DIRECT` | `auto` | `flashinfer/dsa_sparse_attention.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. |\n| `FLASHINFER_CAKE_DSA_TRAIN_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/dsa_sparse_attention.py` | Capacity of the CAKE DSA training binding cache (`cake_backend.BINDING_CACHE`, least recently used first out): read once when `flashinfer.dsa_sparse_attention` 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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_GEMM_TUNING` | unset | `flashinfer/chunked_lm_head.py` | Fallback of the `gemm_tuning` keyword of `chunked_lm_head_loss` / `chunked_lm_head_logprob` 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`. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_FUSE_DW_CAST` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DX_FINALIZE` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DW_STREAM` | `auto` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_HIDDEN_COUNT` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_COMPACT_ROWS` | `1` | `flashinfer/chunked_lm_head.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`). |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE` | `1` | `flashinfer/chunked_lm_head.py` | Read once when `flashinfer.chunked_lm_head` is imported and applied through `cake_backend.BINDING_CACHE.configure` when the wrapper first reaches the backend (the first call of either entry point): `0` disables the per-process LRU cache of compiled bindings. Direct users of the package call `configure` themselves. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/chunked_lm_head.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. |\n| `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. |\n| `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. |\n\nRelease/nightly JIT-cache provider builds run sccache in server-side mode and\nwrap the complete wheel build with a no-output watchdog. After 90 minutes with\nno build output, the watchdog captures system, process, compiler, and sccache\nstate; terminates the wrapped build process group (TERM, then KILL after a\ntwo-minute grace period); and exits with status 124. Compiler discovery scans\nthe whole build container, so diagnostics include compiler processes owned by\nthe sccache daemon even though they are outside the wrapped process group. The\nCI `docker run --rm` boundary provides final cleanup for those processes.\n\n##### Cubin / Artifact Loader\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_CUBIN_CHECKSUM_DISABLED` | unset | `flashinfer/jit/cubin_loader.py` | If set, skip SHA checksum verification of downloaded cubins. Debug aid only. |\n| `FLASHINFER_CUBIN_DOWNLOAD_THREADS` | `4` | `flashinfer/artifacts.py` | Thread-pool size used by `flashinfer artifacts download`. |\n| `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. |\n| `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. |\n| `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. |\n| `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`. |\n| `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. |\n\n##### API Dump / Logging Extensions\n\nThese complement `FLASHINFER_LOGLEVEL`/`FLASHINFER_LOGDEST` and are all read in `flashinfer/api_logging.py`.\n\n| Variable | Default | Effect |\n|----------|---------|--------|\n| `FLASHINFER_DUMP_DIR` | `flashinfer_dumps` | Directory where `@flashinfer_api` writes per-call tensor dumps when logging is enabled. |\n| `FLASHINFER_DUMP_INCLUDE` | `\"\"` | Comma-separated allow-list of API names; only matching calls are dumped. Wildcards (`*`) supported. |\n| `FLASHINFER_DUMP_EXCLUDE` | `\"\"` | Comma-separated deny-list of API names; matching calls skip the (expensive) stats / dump path. |\n| `FLASHINFER_DUMP_MAX_COUNT` | `1000` | Hard cap on number of dumped events; once reached, further dumps are dropped with a warning. |\n| `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`). |\n| `FLASHINFER_DUMP_SAFETENSORS` | `0` | `1` writes dumps as `.safetensors` (portable / Hugging Face-style); otherwise plain `.pt`. |\n\n##### Trace Capture\n\nUsed by `flashinfer.trace` / `fi_trace`.\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `FLASHINFER_TRACE_DUMP` | unset | `flashinfer/fi_trace.py` | If set, decorated APIs auto-dump benchmark-definition JSON for each call (the \"fi_trace\" feature). |\n| `FLASHINFER_TRACE_DUMP_DIR` | cwd | `flashinfer/fi_trace.py` | Directory where the trace JSON files are written. |\n| `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. |\n| `FLASHINFER_TRACE_APPLY_PATH` | unset | `flashinfer/trace_apply/config.py` | Directory from which deployment-configured solutions are loaded for Trace Apply. |\n\n##### Validation / Autotuning / Routing / Kernel Selection\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_AUTOTUNER_LOAD_FROM_FILE` | `0` | `flashinfer/autotuner/autotuner.py` | `1` loads previously serialized autotune results from disk instead of re-running the search. |\n| `FLASHINFER_DIST_AWARE_AUTOTUNE` | `0` | `flashinfer/fuse",".cursor/mcp.json":"# Model Context Protocol (MCP) Configuration\n\nPath: `.cursor/mcp.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.cursor/mcp.json)",".claude-plugin/marketplace.json":"# Claude Plugin Marketplace Catalog\n\nPath: `.claude-plugin/marketplace.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/marketplace.json)",".claude-plugin/plugin.json":"# claude-plugin Plugin Manifest\n\nPath: `.claude-plugin/plugin.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/plugin.json)"},"items":[{"name":"AGENTS.md","path":"AGENTS.md","rawUrl":"https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/AGENTS.md","title":"AI Agent Protocol & Instructions","category":"root-instruction","format":"markdown","content":"This repository’s agent instructions live in [CLAUDE.md](./CLAUDE.md).\n\nBefore making changes, read `CLAUDE.md` and follow its guidance for build, test, style, and contribution workflows.\n","isInternal":false,"tokens":45,"sizeBytes":190},{"name":"CLAUDE.md","path":"CLAUDE.md","rawUrl":"https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/CLAUDE.md","title":"Claude Agent Guidelines & System Prompt","category":"claude-rule","format":"markdown","content":"# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashInfer 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.\n\n## Quick Reference\n\n| Task | Command |\n|------|---------|\n| Install for development | `pip install --no-build-isolation -e . -v` |\n| Initialize submodules | `git submodule update --init --recursive` |\n| Install CUPTI for benchmarking | `pip install -U cupti-python` |\n| Run regular tests | `pytest tests/` |\n| Run full parameter matrices | `pytest tests/ --full` |\n| Run specific test | `pytest tests/path/test_file.py::test_function` |\n| Run multi-GPU test | `mpirun -np 4 pytest tests/comm/test_allreduce_unified_api.py` |\n| Run benchmark | `python benchmarks/flashinfer_benchmark.py --routine <name> <flags>` |\n| Run linting | `pre-commit run -a` |\n| Dump environment report (bug reports) | `python -m flashinfer.collect_env` (or `flashinfer collect-env [--json]`) |\n| Install pre-commit hooks | `pre-commit install` |\n| Clear JIT cache | `rm -rf ~/.cache/flashinfer/` |\n| Inspect devcontainer compiler cache | `sccache --show-stats` |\n| Enable API logging (basic) | `export FLASHINFER_LOGLEVEL=1` |\n| Enable API logging (detailed) | `export FLASHINFER_LOGLEVEL=3` |\n| Enable API logging (with stats) | `export FLASHINFER_LOGLEVEL=5` |\n| Set API log destination | `export FLASHINFER_LOGDEST=mylog.txt` |\n| Enable verbose JIT logging | `export FLASHINFER_JIT_VERBOSE=1` |\n| Enable debug build | `export FLASHINFER_JIT_DEBUG=1` |\n| Set target architectures | `export FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"` |\n| Set parallel compilation | `export FLASHINFER_NVCC_THREADS=4` |\n| Limit parallel ninja jobs | `export MAX_JOBS=4` |\n| Enable GDN native short-T path | `export FLASHINFER_GDN_WY_NATIVE_T=1` |\n| Enable GDN strided QKV path | `export FLASHINFER_GDN_WY_STRIDED_QKV=1` |\n| Enable GDN native A/B tensors | `export FLASHINFER_GDN_WY_NATIVE_AB=1` |\n| 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` |\n| Opt in to multi-wave KV split in trtllm-gen GQA decode (long or mixed-length KV) | `export FLASHINFER_TRTLLM_GEN_MULTI_WAVE_KV_SPLIT=1` |\n| Let `backend=\"auto\"` pick experimental backends | `export FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1` |\n| Override ragged-prefill `auto` backend order (Blackwell) | `export FLASHINFER_RAGGED_AUTO_BACKEND_ORDER=cutlass,cudnn` |\n| Override CuTe-DSL prefill scheduling | `export FLASHINFER_CUTE_PREFILL_PERSISTENT=0` (non-persistent) or `1` (persistent) |\n| Skip MoE EP CuTe-DSL import/version guard | `export FLASHINFER_MOE_EP_SKIP_DSL_CHECK=1` |\n| Override MoE EP knob-cache path | `export FLASHINFER_MOE_EP_KNOB_CACHE=/path/to/knobs.json` |\n| Disable MoE EP fused staging kernel | `export FLASHINFER_MEGA_FUSED_STAGE=0` |\n| Log MegaMoE compilation and workspace diagnostics | `export FLASHINFER_MEGAMOE_LOG_COMPILE_SPEC=1` |\n| Enable MoE EP EXPERT_MAJOR padding-row exclusion (opt-in; shape-dependent) | `export FLASHINFER_MOE_EP_EXCLUDE_PADDING_ROWS=1` |\n| Enable distribution-aware MoE autotune and kernel dispatch (experimental; TRT-LLM MoE only) | `export FLASHINFER_DIST_AWARE_AUTOTUNE=1` |\n| Prims-TS debug checks / task-manager verify | `export FLASHINFER_PRIMS_TS_DEBUG_CHECKS=1` |\n| Override Prims-TS CuTe-DSL compile options | `export FLASHINFER_PRIMS_TS_COMPILE_OPTIONS=\"--opt-level 2\"` |\n| Configure NVFP4 4over6 quantization (preferred over the env vars) | `nvfp4_quantize(..., nvfp4_4over6=NVFP44Over6Config(e4m3_max=256))` |\n\nThe minimum Python version used by CI and build tooling is defined in\n`.python-version`. CI Docker images use a stable Conda environment name and\nderive their interpreter and site-packages paths from that file. When raising\nthe minimum, also update `project.requires-python` and `tool.mypy.python_version`\nin `pyproject.toml`.\n\n## Quick Start for Development\n\n### Installation\n\n```bash\ngit clone https://github.com/flashinfer-ai/flashinfer.git --recursive\ncd flashinfer\npip install --no-build-isolation -e . -v\n```\n\n**Important**: The `--recursive` flag is required to initialize submodules in `3rdparty/` (cutlass, spdlog).\n\nIf you forgot `--recursive` when cloning:\n```bash\ngit submodule update --init --recursive\n```\n\nThat's it! You can now:\n\n- Run all benchmarks and unit tests\n- Modify kernel source code in `include/` without reinstalling\n- Changes are JIT-compiled on next use\n\nThe `--no-build-isolation` flag prevents pip from pulling incompatible PyTorch/CUDA versions from PyPI.\n\n### How JIT Compilation Works\n\nWhen you call a FlashInfer API:\n\n1. **First call**: Generates specialized CUDA code based on parameters (dtype, head_dim, etc.), compiles it with ninja, caches the .so file\n2. **Subsequent calls**: Uses cached compiled module\n3. **After kernel changes**: Automatically detects changes and recompiles\n\n**No manual rebuild step needed** - just edit `.cuh` files and run your code again.\n\n### Pre-compiled Packages (Optional)\n\nFlashInfer provides optional pre-compiled packages for users who want faster initialization:\n\n- `flashinfer-jit-cache`: Pre-built kernel cache\n- `flashinfer-cubin`: Pre-compiled kernel binaries\n\n**For development, you typically DON'T need these.** JIT compilation is fast enough and gives you live code reload.\n\n### Compiler Caching in Devcontainers\n\nThe devcontainers run every nvcc and host C++ compile through\n[sccache](https://github.com/mozilla/sccache) by setting `FLASHINFER_NVCC_LAUNCHER`\nand `FLASHINFER_CXX_LAUNCHER`. Cached objects live in a per-devcontainer Docker\nvolume mounted at `~/.cache/sccache`, so they outlive a container rebuild — which\nwipes `~/.cache/flashinfer` and would otherwise force a full recompile.\n\n```bash\nsccache --show-stats                  # hit rate, cache size, uncacheable reasons\nexport FLASHINFER_NVCC_LAUNCHER=\"\"    # disable for the current shell\n```\n\nCaveats:\n\n- Toggling a launcher changes the compile commands in `build.ninja`, so the first\n  build after switching it on or off recompiles objects that already exist.\n- Link steps, cubin embedding, and CuTe-DSL kernels never invoke nvcc, so they are\n  not cached.\n- The cache is per-devcontainer by design: sccache's local storage supports a\n  single server at a time, so two containers must not share one cache directory.\n\nGitHub PR Test GPU jobs keep a different, per-runner cache: they set\n`CI_COMPILE_CACHE_DIR` to a directory under the runner's work directory, and\n`ci/bash.sh` namespaces it by runner and the pulled image's immutable ID.\n`CI_COMPILE_CACHE_MODE=read-only` (the default) mounts a read-only seed and runs\nagainst a private writable copy, discarded even after test failures. The PR Test\nworkflow selects `read-write` only for pushes to `main`. A stable per-runner lock\nprotects maintenance and execution, including accidentally shared workspaces.\nOther `ci/bash.sh` callers leave caching unset and run cold.\nBefore each trusted writable job its runner's cache is cleared when it exceeds\n`CI_COMPILE_CACHE_MAX_MB`\n(default 40960) or its volume has less than `CI_COMPILE_CACHE_MIN_FREE_MB`\n(default 20480) free. Only content-keyed artifacts survive between jobs: CuTe-DSL (with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1`), TIRx, KDA PTX, Triton and downloaded\ncubins. nvcc modules, tuning results and the generated-source tree are pruned,\nsince a fresh checkout rebuilds nvcc modules anyway. Each test job's\n`FLASHINFER_JIT_BUILD_LOG` is summarized per toolchain by\n`scripts/summarize_jit_builds.py` in the job summary.\nSee [`ci/COMPILE_CACHE.md`](ci/COMPILE_CACHE.md) for lifecycle tests and trust\nboundaries.\n\n## cuDNN linear attention\n\nUse `nvidia-cudnn-frontend>=1.31.0`; KDA and GDN normalization use additive\n`1e-6` without an epsilon argument. Auto selection also requires CuTe DSL 4.7+.\nThe bounded KDA SM100/SM103 and GDN SM100 policies live in `kda.py` and\n`gdn_prefill.py`. KDA's head/length regions are documented in `docs/api/kda.rst`.\nAuto-selected cuDNN ignores the native KDA workspace; native fallback uses it.\nValidate output and final state before changing a policy. Explicit state-pool calls\nrequire unique, in-range slots; auto keeps pools on their native route.\n\nUse `benchmarks/bench_cudnn_linear_attention.py` for public-call comparisons.\nIts raw-gate GDN baseline includes Torch gate transforms; it is not a vLLM\nfused-producer baseline. Keep host CPU, completed wall and reset-inclusive\nCUDA graph times separate.\n\n## Testing\n\nRun the regular test suite:\n\n```bash\npytest tests/\n```\n\nRun full parameter matrices (used by nightly testing):\n\n```bash\npytest tests/ --full\n```\n\nRun specific test file:\n\n```bash\npytest tests/attention/test_hopper.py\n```\n\nRun specific test function:\n\n```bash\npytest tests/attention/test_hopper.py::test_single_prefill\n```\n\nThe sharded test runner defaults to one worker per visible GPU. H100 PR CI opts\ninto `--workers-per-gpu 2`: eight workers on four GPUs. Other CI lanes retain\nthe default. Locally, `scripts/task_run_unit_tests.sh --workers-per-gpu 2` enables\nsharing; `--workers N` sets the total instead. Workers share the existing JIT\ncompilation budget.\n\nScheduling markers apply to the whole source file:\n\n- `node_exclusive`: run alone across the allocation after parallel work, with\n  all GPUs visible. Use for multi-GPU requirements or host-memory OOM risk.\n- `gpu_exclusive`: run alone on the assigned GPU; other GPUs can keep working.\n- `host_memory_heavy`: prioritize this file and serialize its units with other\n  host-heavy work, while lighter files can overlap.\n\nHost and GPU reservations cover setup, execution, and process cleanup, and are\nacquired together. Blocked units stay queued so workers can take runnable work\nfrom their own queue or steal from another. Pending GPU-exclusive work can\ndrain one GPU without reserving the host-heavy slot or blocking other GPUs.\n`RESOURCE WAIT` reports worker idle time of at least one second while work is\nqueued but cannot run on that worker's resources.\n\nThese markers enforce overlap rules, not memory limits. Most H100 annotations\nuse measured peaks above 128 GiB host RSS or 32 GiB GPU memory. Attention sink\nalso uses both narrower markers: its single-GPU H100 run peaked at 89.2 GiB host\nRSS and 4.5 GiB GPU memory, and other GPUs can keep working while it runs.\nCold JIT compilation and other GPU families can change the peaks. Validate\ncompletion rate, OOM failures, and repeated wall times with matching scope and\ncache conditions when changing a CI lane's concurrency policy.\n\nCapacity reports include the serial floor from host-heavy and node-exclusive\nwork. They remain optimistic about GPU contention because the estimator uses\nworker counts without GPU topology or measured memory capacity.\n\nRun the CPU-only coordinator suite with\n`pytest --confcutdir=tests/test_sharding tests/test_sharding/`.\n\nCommunication-test wrappers optionally emit cache-publication completion reports\nwhen `UNIT_TEST_RUN_STATUS_DIR` is set to a fresh directory unique to that wrapper\nand container/rank. `test_run_status.sh` records each pytest invocation, and\n`test_run_status.py` atomically writes `script-result.json` (schema version 1,\n`complete`, `exit_code`, `tests_executed`, and `infrastructure_errors`). A small\npytest plugin records execution and collection/internal errors inside each MPI\nor torchrun worker; launcher exit 1 alone is not evidence of an assertion failure.\nAll local worker ranks must finish, and at least one test must execute, to allow\npublication. A completed run with assertion failures remains eligible, while\ncrashes, timeouts, and unfinished invocations do not. With the variable unset,\nreporting is disabled. The CPU suite above covers this contract without GPU\ndependencies.\n\n### Skipping Tests Based on CUDA Architecture\n\nUse `flashinfer.utils` functions to skip tests on unsupported GPU architectures:\n\n**Available check functions:**\n- `get_compute_capability(device)` - Returns `(major, minor)` tuple\n- `is_sm90a_supported()` - Hopper (requires CUDA 12.3+)\n- `is_sm100a_supported()` - Blackwell (requires CUDA 12.8+)\n- `is_sm100f_supported()` - Blackwell feature-set (requires CUDA 12.9+)\n- `is_sm110a_supported()`, `is_sm120a_supported()`, `is_sm120f_supported()`, `is_sm121a_supported()`\n\n**APIs decorated with `@backend_requirement`** also provide:\n- `api_name.is_compute_capability_supported(cc)` - e.g., `mm_fp4.is_compute_capability_supported(100)`\n- `api_name.is_backend_supported(\"backend\")` - e.g., `mm_fp4.is_backend_supported(\"cudnn\")`\n- `api_name.is_backend_supported(\"backend\", cc)` - e.g., `mm_fp4.is_backend_supported(\"cudnn\", 80)`\n\n**Example:**\n```python\nfrom flashinfer.utils import is_sm90a_supported\n\ndef test_hopper_attention():\n    if not is_sm90a_supported(torch.device(\"cuda\")):\n        pytest.skip(\"Requires SM90a\")\n    # Test code...\n```\n\n**Common requirements:**\n\n| Feature | Min SM | Check Function |\n|---------|--------|----------------|\n| FlashAttention-3 | SM90a | `is_sm90a_supported()` |\n| MLA Attention | SM100a | `is_sm100a_supported()` |\n| FP8 GEMM | SM89+ | `get_compute_capability()[0] >= 9` |\n\n**Note:** `tests/conftest.py` auto-skips tests that trigger OOM, but tests should be written to avoid OOM by using appropriate problem sizes.\n\n## Benchmarking\n\nFor SM90 NVFP4 fused/split and MXFP4 fused MegaMoE, use\n[`benchmarks/bench_megamoe_fp4_sm90.py`](benchmarks/bench_megamoe_fp4_sm90.py) for timing\nand [`tests/moe/test_megamoe_fp4_sm90.py`](tests/moe/test_megamoe_fp4_sm90.py) for correctness checks.\n\nFlashInfer provides a unified benchmarking framework in `benchmarks/flashinfer_benchmark.py`.\n\n**Key features:**\n- Supports attention, GEMM, and MOE kernels\n- Multiple backends: FlashAttention2/3, cuDNN, CUTLASS, TensorRT-LLM, cuBLAS\n- **CUPTI timing (recommended)**: Hardware-level profiling for accurate GPU kernel time\n  - Automatically falls back to CUDA events if CUPTI unavailable\n  - Install: `pip install -U cupti-python` (requires CUDA 13+)\n- Batch testing, reference checking, CSV output\n\n**Quick example:**\n```bash\npython benchmarks/flashinfer_benchmark.py \\\n    --routine BatchDecodeWithPagedKVCacheWrapper \\\n    --backends fa2 cudnn \\\n    --batch_size 32 --s_kv 2048 \\\n    --num_qo_heads 32 --num_kv_heads 8 \\\n    --head_dim_qk 128 --head_dim_vo 128 \\\n    --page_size 16 --refcheck -vv\n```\n\n**Python API:**\n```python\nfrom flashinfer.testing import bench_gpu_time\n\n# CUPTI preferred, auto-fallback to CUDA events.\n# Returns per-iteration times in milliseconds (a list).\ntimes = bench_gpu_time(\n    my_kernel, input_args=(x, y), enable_cupti=True,\n    dry_run_iters=10, repeat_iters=30,\n)\nmedian_time_ms = statistics.median(times)\n```\n\n→ **For complete benchmarking guide, see [`.claude/skills/benchmark-kernel/SKILL.md`](.claude/skills/benchmark-kernel/SKILL.md)**\n\n`flashinfer.gemm.group_gemm_fp8_nt_groupwise_contiguous` is the standalone\nCuTe-DSL grouped FP8 API; `group_deepgemm_fp8_nt_groupwise` keeps its original\nDeepGEMM-only signature. Both use per-row expert indices, unlike the indptr-based\n`group_gemm_fp8_nt_groupwise` API.\n\nFor contiguous grouped FP8 GEMM, `python benchmarks/bench_grouped_fp8.py\n--production-shapes` compares DeepGEMM and CuTe-DSL with preallocated outputs;\n`--cache-probe` measures first-call and new-token-count compilation separately.\nCompiled grouped CuTe-DSL kernels are reused across token counts with the same\n128-row alignment class, with at most two variants per device and weight shape.\nConcurrent misses share one compilation; warm hits do not acquire the compile lock.\nThe standalone API's optional `validate_indices=True` checks expert-index values\nand synchronizes with the CPU. Validate new routing data before CUDA graph\ncapture; the default path assumes valid indices and performs metadata checks only.\n\nPrims-TS decode keeps raw storage dtypes (`q_dtype`, `k_dtype`, `v_dtype`)\nseparate from effective MMA dtypes (`qk_dtype`, `pv_dtype`). Transformed\nFP8/NVFP4 K/V uses Q's compute dtype; BF16 Q/K with FP8 V uses FP8 PV without\nthat transform. Kernel predicates, P packing, and attention-sink scaling must\nfollow effective compute precision; TMA descriptors and raw allocations must\nfollow storage precision. The public plan accepts separate `k_data_type` and\n`v_data_type`, with `kv_data_type` retained as a common-KV compatibility alias.\n\n## Code Linting\n\nRun all pre-commit hooks:\n\n```bash\npre-commit run -a\n```\n\nInstall hooks to run on every commit:\n\n```bash\npre-commit install\n```\n\n## Opening a Pull Request\n\nFollow [`CONTRIBUTING.md`](CONTRIBUTING.md) when opening a PR. The two requirements most often\nmissed:\n\n- **Keep the default PR template.** Fill in `.github/pull_request_template.md` (Description,\n  Related Issues, Checklist, Tests, Reviewer Notes). Do **not** overwrite or replace it with a\n  custom or tool-generated description format. The title and description normally become the\n  commit title and message on squash-merge, and are what a `git bisect` surfaces months later\n  when someone is hunting the owner of a regression.\n- **Report before/after numbers for performance work.** A perf PR must include measurements from\n  a reproducible benchmark (e.g. `benchmarks/flashinfer_benchmark.py`), naming the GPU and the\n  problem sizes. Speedup ratios without absolute numbers, or numbers without a named GPU, are\n  not enough.\n\nWhen preparing changes for a PR, a self-review is strongly recommended: follow\n[`.claude/skills/self-review/SKILL.md`](.claude/skills/self-review/SKILL.md), which walks the\ndiff through the repository's review guidance and PR rules. It is informal and not required —\nits purpose is to make review easier for reviewers and contributors alike by catching common\nissues before a human looks. Changes made only to test something locally, with no PR intended,\ndo not need one.\n\n→ **For the complete contribution rules, see [`CONTRIBUTING.md`](CONTRIBUTING.md)**\n\n## Code Review\n\nWhen reviewing a diff (as an agent or a human), follow the shared focus areas, kernel-review\npolicy, and effort calibration in [`docs/code_review_guidance.md`](docs/code_review_guidance.md).\nNote: unlike human review, agents keep **kernel implementation details in scope** — read the\nkernel logic and report bugs, labeling findings by confidence.\n\n→ **For the complete review rule set, see [`docs/code_review_guidance.md`](docs/code_review_guidance.md)**\n\n## Architecture: JIT Compilation System\n\nFlashInfer's JIT system has three layers:\n\n### Layer 1: JitSpec (flashinfer/jit/core.py)\n\n`JitSpec` is an abstract base class defining the kernel-module lifecycle\n(`try_load()` / `build()` / `load()`, with the shared `build_and_load()`\ntemplate method handling caching, locking, and `FLASHINFER_DISABLE_JIT`).\nOne subclass per compilation toolchain:\n\n- `JitSpecNvcc` (nvcc/ninja modules, returned by `gen_jit_spec()`) defines:\n  - `name`: Unique identifier (URI hash from parameters)\n  - `sources`: List of .cu/.cpp files to compile\n  - `extra_cuda_cflags`, `extra_cflags`, `extra_ldflags`: Compiler flags\n- `JitSpecCuteDsl` (flashinfer/jit/cute_dsl_core.py) caches CuTe-DSL kernels\n  (see \"CuTe-DSL kernels\" under Module Caching below)\n- The optional TIRx KDA backend uses a private `JitSpec` in\n  `flashinfer/kda_kernels/tirx/cache.py` to cache TVM modules, with source/compiler\n  fingerprints and artifact checksums. Its TVM/TIRx dependencies are loaded only\n  by explicit `recurrent_kda(..., backend=\"tirx\")`; see `docs/api/kda.rst`.\n\n### JIT Directory Rules\n\n**NEVER write to package directories** - they may be read-only after installation.\n\n| Directory | Writable | Use for |\n|-----------|----------|---------|\n| `FLASHINFER_GEN_SRC_DIR` | ✓ Yes | Generated source files (Jinja output, copied .cu files) |\n| `FLASHINFER_JIT_DIR` | ✓ Yes | Compiled `.so` outputs |\n| `FLASHINFER_CSRC_DIR` | ✗ No | Read-only source templates |\n| `FLASHINFER_AOT_DIR` | ✗ No | Read-only pre-compiled binaries |\n\n### Compilation Context: Architecture-Specific Compilation\n\nFlashInfer uses `CompilationContext` to manage CUDA architecture targets. Some kernels only work on specific GPU architectures (e.g., Hopper SM90, Blackwell SM100/SM12x).\n\n**How it works:**\n- Auto-detects GPUs in system or reads `FLASHINFER_CUDA_ARCH_LIST` environment variable\n- CuTe-DSL FP4 MM compilation parallelism can be controlled with `FLASHINFER_MM_FP4_CUTE_DSL_COMPILE_WORKERS`\n- JIT modules specify `supported_major_versions=[9, 10, 11, 12]` to limit compilation to specific SM versions\n- If GPU not supported → `RuntimeError: No supported CUDA architectures found`\n\n→ **See [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md) for usage examples**\n\n### Layer 2: Code Generation\n\nEvery `gen_*_module()` function in `flashinfer/jit/` follows this pattern:\n\n```python\ndef gen_some_module(dtype_in, dtype_out, ...):\n    # 1. Compute unique identifier from parameters\n    uri = get_some_uri(dtype_in, dtype_out, ...)\n\n    # 2. Create generation directory\n    gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri\n\n    # 3. (Optional) Render Jinja template to generate type-specialized config\n    # Skip this step if you don't need type specialization\n    with open(jit_env.FLASHINFER_CSRC_DIR / \"some_customize_config.jinja\") as f:\n        template = jinja2.Template(f.read())\n    config_content = template.render(\n        dtype_in=dtype_map[dtype_in],\n        dtype_out=dtype_map[dtype_out],\n        # ... more parameters\n    )\n    write_if_different(gen_directory / \"some_config.inc\", config_content)\n\n    # 4. Copy source files to gen directory\n    sources = []\n    for fname in [\"some_kernel.cu\", \"some_jit_binding.cu\"]:\n        shutil.copy(jit_env.FLASHINFER_CSRC_DIR / fname, gen_directory / fname)\n        sources.append(gen_directory / fname)\n\n    # 5. Return JitSpec\n    return gen_jit_spec(uri, sources, extra_cuda_cflags=[...])\n```\n\n**Note**: If your operation doesn't need type specialization, you can skip step 3 entirely and just copy the source files directly.\n\n### Layer 3: Compilation and Loading\n\n`JitSpecNvcc` methods:\n\n- `write_ninja()` - Generates `build.ninja` file\n- `build()` - Executes `ninja` to compile sources\n- `build_and_load()` - Compiles and loads via TVM-FFI\n\nThe generated `build.ninja` file uses nvcc to compile .cu → .cuda.o → .so, then loads via TVM-FFI.\n\n### Jinja Templates (Optional)\n\n**Note: Jinja templates are NOT required.** You can write C++ code directly without templating.\n\nFor operations that need type specialization, templates in `csrc/*.jinja` can generate C++ code:\n\n```jinja\n// Input template\nusing DTypeIn = {{ dtype_in }};\nusing DTypeOut = {{ dtype_out }};\nconstexpr int PARAM = {{ param_value }};\n\n// After render\nusing DTypeIn = float16;\nusing DTypeOut = float16;\nconstexpr int PARAM = 128;\n```\n\nThis 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.\n\n## Directory Structure\n\n```\nflashinfer/\n├── include/flashinfer/           # Header-only CUDA kernel templates\n│   ├── attention/                # Attention kernels\n│   ├── gemm/                     # GEMM kernels\n│   ├── comm/                     # Communication kernels\n│   ├── mma.cuh                   # Matrix multiply utilities\n│   ├── utils.cuh                 # Common utilities\n│   └── [...]\n│\n├── csrc/                          # Framework bindings (via TVM-FFI)\n│   ├── *.cu                       # Kernel launcher implementations\n│   ├── *_jit_binding.cu           # TVM-FFI exports\n│   ├── *_customize_config.jinja   # Type config templates (optional)\n│   └── [...]\n│\n├── flashinfer/                    # Python package\n│   ├── jit/\n│   │   ├── core.py                # JitSpec, compilation infrastructure\n│   │   ├── cpp_ext.py             # Ninja build generation\n│   │   ├── env.py                 # Workspace paths\n│   │   ├── attention/             # Attention module generators\n│   │   ├── gemm/                  # GEMM module generators\n│   │   ├── fused_moe/             # MOE module generators\n│   │   └── [...]\n│   ├── gemm/                      # GEMM Python APIs\n│   ├── fused_moe/                 # MOE Python APIs\n│   ├── comm/                      # Communication Python APIs\n│   ├── *.py                       # Other high-level Python APIs\n│   ├── aot.py                     # AOT compilation for pre-built packages\n│   └── [...]\n│\n├── tests/                         # Test suite\n│   ├── attention/                 # Attention kernel tests\n│   ├── gemm/                      # GEMM kernel tests\n│   ├── moe/                       # MOE kernel tests\n│   ├── comm/                      # Communication tests\n│   ├── utils/                     # Utility tests\n│   └── conftest.py                # Pytest configuration\n│\n└── build_backend.py               # PEP 517 build backend\n```\n\n### Critical Rule: Framework Separation\n\n**Torch headers MUST NOT be included in `include/` directory files.**\n\n- `include/`: Framework-agnostic CUDA kernels (accept raw pointers)\n- `csrc/`: Framework bindings via TVM-FFI (currently PyTorch, but can support other frameworks)\n\n## Adding a New Operation\n\n→ **For complete step-by-step tutorial, see [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md)**\n\n**Quick overview of the process:**\n1. Write kernel in `include/flashinfer/new_op.cuh` (framework-agnostic, raw pointers)\n2. Write launcher in `csrc/new_op.cu` (PyTorch tensor handling)\n3. Create TVM-FFI bindings in `csrc/new_op_jit_binding.cu`\n4. (Optional) Create Jinja template for type specialization\n5. Write JIT module generator in `flashinfer/jit/new_op.py`\n6. Write Python API in `flashinfer/new_op.py` with `@functools.cache`\n7. Write tests in `tests/`\n8. Register in `flashinfer/aot.py` for AOT compilation\n9. Export in `flashinfer/__init__.py`\n10. Add a `TraceTemplate` in `flashinfer/trace/templates/` and wire it via `@flashinfer_api(trace=...)` (see below)\n11. Add an example call in `tests/trace/example.py`, re-run to regenerate `fi_trace_out/`, and commit the new JSON files\n\n### Trace Template Checklist (for new or updated APIs)\n\nEvery 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.\n\n1. **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.\n2. **Wire the template** to the API: `@flashinfer_api(trace=my_trace)` on the Python function (or class method's `run()`).\n3. **Add an example call** in `tests/trace/example.py` that exercises the new trace with realistic shapes.\n4. **Regenerate examples**: `rm -rf tests/trace/fi_trace_out && python tests/trace/example.py` — verify the expected JSON appears.\n5. **Update the docstring** in `tests/trace/example.py` to list the new file(s).\n6. **Run tests**: `pytest tests/trace/ -v` — all template-consistency and end-to-end tests must pass.\n7. **Commit the new JSON files** under `tests/trace/fi_trace_out/` alongside the code changes.\n\n**Example implementations:**\n- **Simple**: `flashinfer/norm/__init__.py` (RMSNorm) - no Jinja, good starting point\n- **Moderate**: `flashinfer/sampling.py` - with Jinja templating\n- **Complex**: `flashinfer/decode.py` - plan-run pattern, advanced workspace\n\n## Experimental Code\n\nExperimental backends live under `flashinfer/experimental/`; experimental\npublic APIs live in core marked with `@flashinfer_experimental_api` (defined in\n`flashinfer/api_logging.py`). Using them is an explicit opt-in: calling an\nexperimental API, or naming an experimental backend with `backend=\"<name>\"`,\nneeds no environment variable (both warn once). Only automatic selection is\ngated: `backend=\"auto\"` skips backends marked `@experimental_backend` unless\n`FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. Experimental code is JIT-only\n(never registered in `flashinfer/aot.py`) and tested under\n`tests/experimental/`.\n\n→ **When working under `flashinfer/experimental/`, follow\n[`flashinfer/experimental/CLAUDE.md`](flashinfer/experimental/CLAUDE.md); the\nfull policy is [`flashinfer/experimental/README.md`](flashinfer/experimental/README.md).**\n\n## Key Architectural Patterns\n\n### Module Caching\n\nFlashInfer uses two-level caching to avoid recompilation:\n\n1. **Python-level** (`@functools.cache`): In-memory cache of loaded modules\n2. **File-level** (`~/.cache/flashinfer/`): Compiled `.so` files on disk\n\n**Cache invalidation** (automatic):\n- Source file changes (SHA256 hash)\n- Compilation flags change\n- CUDA architecture change\n- FlashInfer version change\n\nURI computed as: `hash(operation_type + parameters + source_hashes + flags + cuda_arch)`\n\n**Cache management:**\n- Clear cache: `rm -rf ~/.cache/flashinfer/`\n- Override location: `export FLASHINFER_WORKSPACE_BASE=\"/scratch\"`\n\n**CuTe-DSL kernels** (`flashinfer/jit/cute_dsl_core.py`): `cute.compile()` has no\npersistent cache, so `JitSpecCuteDsl` (a `JitSpec` subclass, wrapped by the\n`build_and_load_cute_dsl_kernel()` helper) exports compiled kernels as\nobject files (`export_to_c()`) and reloads them with `cute.runtime.load_module(...,\nenable_tvm_ffi=True)`. Artifacts live in `cached_ops/` next to the nvcc modules, one\ndirectory per op family — `cached_ops/<module>_<arch>_cute_dsl/` (e.g.\n`nvfp4_quantize_sm100a_cute_dsl/`) holding one `meta.json` plus one `.o` per\nspecialization. The arch comes from the DSL's compile target (`CUTE_DSL_ARCH` or the\ncurrent device) since the artifacts are single-arch, unlike nvcc fatbins.\nInvalidation is module-granular: changed DSL/Torch/TVM-FFI distributions, Python\nABI, or\nkernel-source SHA256 wipes and lazily rebuilds the module. Reference usage:\n`flashinfer/quantization/kernels/nvfp4_quantize.py`. Disable with\n`FLASHINFER_CUTE_DSL_DISABLE_CACHE=1`. The SHA256 covers only the files a caller\npasses as `extra_key_files` (often just `__file__`); with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1` it also covers every flashinfer module\nthose files can run. Source analysis and hashing live in\n`flashinfer/jit/_cute_dsl_source.py`; `cute_dsl_core.py` manages compilation and\ncached artifacts. Plain modules are keyed by content with everything they\nimport, including literal `importlib.import_module` calls. Package initializers\nare resolved name by name: a name looked up in a package follows the binding\nthat defines it (a re-export to its source, a function into its body, a module\nthrough the attributes read from it) plus the submodule of that name. The code\nan initializer runs at import, including evaluated function annotations, is\nfollowed for every keyed file in its package. Star imports remain covered when\nthey can replace an earlier binding; lazy fallbacks are followed when a binding\nmay be absent.\nEach initializer is keyed by that code and the import lines of used names, so a\nnew re-export or `__all__` entry invalidates nothing. Not keyed: import-time\nside effects of modules that keyed code neither imports nor uses, nonliteral\ndynamic imports, and runtime file reads; list the latter in `extra_key_files`.\n`version.py` and `_build_meta.py` are excluded to avoid keying on the checkout\ncommit. This is intended for long-lived caches shared across commits.\nTask-scheduled attention shares this cache through\n`flashinfer/attention/prims_ts/_compile_cache.py`; its keys cover compile specs,\noptions, and embedded device-capacity constants. It hashes the attention source\ntree, including shared preparation, masks, and reductions. Occupancy queries\nare cached per device and cluster size by `_hardware.py`.\n\n### Dispatch Macros\n\nHandle combinatorial parameter spaces:\n\n```cpp\nDISPATCH_DTYPE(input_dtype, DTypeIn, {\n  DISPATCH_DTYPE(output_dtype, DTypeOut, {\n    DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {\n      LaunchKernel<DTypeIn, DTypeOut, BLOCK_SIZE>(...);\n    });\n  });\n});\n```\n\nDefined in `.jinja` files and expanded after rendering.\n\n## API Logging with @flashinfer_api\n\nFlashInfer provides the `@flashinfer_api` decorator for debugging API calls.\n\n**Key features:**\n- **Crash-safe**: Logs inputs BEFORE execution (preserves info even if kernel crashes)\n- **Zero overhead when disabled**: `FLASHINFER_LOGLEVEL=0` (default)\n- **Multiple verbosity levels**: 0 (off), 1 (names), 3 (inputs/outputs), 5 (+ statistics)\n- **CUDA graph compatible**: Auto-skips stats during graph capture\n\n**Quick usage:**\n```bash\n# Enable detailed logging\nexport FLASHINFER_LOGLEVEL=3              # 0, 1, 3, or 5\nexport FLASHINFER_LOGDEST=debug.log       # stdout, stderr, or file path\n\npython my_script.py\n```\n\n**Why use this?**\n- Debug CUDA crashes (see inputs that caused crash)\n- Track tensor shapes/dtypes through pipeline\n- Detect NaN/Inf issues (level 5)\n\n→ **For complete debugging guide, see [`.claude/skills/debug-cuda-crash/SKILL.md`](.claude/skills/debug-cuda-crash/SKILL.md)**\n\n## Debugging\n\n### Enable Logging\n\n```bash\nexport FLASHINFER_JIT_VERBOSE=1      # Verbose JIT output\nexport FLASHINFER_JIT_DEBUG=1        # Debug symbols, -O0\nexport FLASHINFER_LOGLEVEL=3         # API logging (0=off, 1=basic, 3=detailed)\nexport FLASHINFER_LOGDEST=stdout\n```\n\n### Inspect Generated Code\n\n```bash\n# Generated sources\nls -la ~/.cache/flashinfer/0.6.0/*/generated/\n\n# Compiled modules\nls -la ~/.cache/flashinfer/0.6.0/*/cached_ops/\n\n# Build files\ncat ~/.cache/flashinfer/0.6.0/*/cached_ops/*/build.ninja\n```\n\n### Environment Variables\n\n```bash\n# Compilation\nexport FLASHINFER_NVCC_THREADS=4              # Threads per nvcc process (--threads=N)\nexport MAX_JOBS=4                             # Parallel ninja jobs (nvcc processes)\nexport FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"  # Target architectures\n# Memory note: total compilation memory ≈ MAX_JOBS × FLASHINFER_NVCC_THREADS × per-thread mem.\n\n# Behavior\nexport FLASHINFER_WORKSPACE_BASE=\"/scratch\"   # Custom cache directory\n```\n\n#### Full Environment Variable Reference\n\nAlready covered above (Quick Reference / Debugging): `FLASHINFER_LOGLEVEL`,\n`FLASHINFER_LOGDEST`, `FLASHINFER_JIT_VERBOSE`, `FLASHINFER_JIT_DEBUG`,\n`FLASHINFER_CUDA_ARCH_LIST`, `FLASHINFER_NVCC_THREADS`,\n`FLASHINFER_WORKSPACE_BASE`, `MAX_JOBS`.\n\nThe remaining `FLASHINFER_*` knobs read by the code are listed below. Defaults\nmatch what the code uses today; values are strings unless noted.\n\n##### JIT / Build Toolchain\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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`). |\n| `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. |\n| `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`. |\n| `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. |\n| `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. |\n| `FLASHINFER_JIT_LINEINFO` | `0` | `flashinfer/jit/core.py` | `1` adds `-lineinfo` to nvcc so profiler / `cuda-gdb` can map PTX back to CUDA source. |\n| `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. |\n| `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). |\n| `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`. |\n| `FLASHINFER_CXX_LAUNCHER` | `\"\"` (`sccache` in the devcontainers) | `flashinfer/jit/cpp_ext.py` | Same idea as `FLASHINFER_NVCC_LAUNCHER` but for the host C++ compiler. |\n| `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. |\n| `FLASHINFER_EXTRA_CFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to the host C++ compiler. |\n| `FLASHINFER_EXTRA_CUDAFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to `nvcc`. |\n| `FLASHINFER_EXTRA_LDFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra linker flags passed to the linker. |\n| `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. |\n| `FLASHINFER_CAKE_DSA_CHECK_DST_MAP` | `0` | `flashinfer/dsa_sparse_attention.py` | Fallback of the `check_dst_map` keyword of `dsa_sparse_attention` / `dsa_sparse_attention_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. |\n| `FLASHINFER_CAKE_DSA_DKV_DIRECT` | `auto` | `flashinfer/dsa_sparse_attention.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. |\n| `FLASHINFER_CAKE_DSA_TRAIN_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/dsa_sparse_attention.py` | Capacity of the CAKE DSA training binding cache (`cake_backend.BINDING_CACHE`, least recently used first out): read once when `flashinfer.dsa_sparse_attention` 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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_GEMM_TUNING` | unset | `flashinfer/chunked_lm_head.py` | Fallback of the `gemm_tuning` keyword of `chunked_lm_head_loss` / `chunked_lm_head_logprob` 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`. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_FUSE_DW_CAST` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DX_FINALIZE` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DW_STREAM` | `auto` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_HIDDEN_COUNT` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_COMPACT_ROWS` | `1` | `flashinfer/chunked_lm_head.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`). |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE` | `1` | `flashinfer/chunked_lm_head.py` | Read once when `flashinfer.chunked_lm_head` is imported and applied through `cake_backend.BINDING_CACHE.configure` when the wrapper first reaches the backend (the first call of either entry point): `0` disables the per-process LRU cache of compiled bindings. Direct users of the package call `configure` themselves. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/chunked_lm_head.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. |\n| `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. |\n| `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. |\n\nRelease/nightly JIT-cache provider builds run sccache in server-side mode and\nwrap the complete wheel build with a no-output watchdog. After 90 minutes with\nno build output, the watchdog captures system, process, compiler, and sccache\nstate; terminates the wrapped build process group (TERM, then KILL after a\ntwo-minute grace period); and exits with status 124. Compiler discovery scans\nthe whole build container, so diagnostics include compiler processes owned by\nthe sccache daemon even though they are outside the wrapped process group. The\nCI `docker run --rm` boundary provides final cleanup for those processes.\n\n##### Cubin / Artifact Loader\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_CUBIN_CHECKSUM_DISABLED` | unset | `flashinfer/jit/cubin_loader.py` | If set, skip SHA checksum verification of downloaded cubins. Debug aid only. |\n| `FLASHINFER_CUBIN_DOWNLOAD_THREADS` | `4` | `flashinfer/artifacts.py` | Thread-pool size used by `flashinfer artifacts download`. |\n| `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. |\n| `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. |\n| `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. |\n| `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`. |\n| `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. |\n\n##### API Dump / Logging Extensions\n\nThese complement `FLASHINFER_LOGLEVEL`/`FLASHINFER_LOGDEST` and are all read in `flashinfer/api_logging.py`.\n\n| Variable | Default | Effect |\n|----------|---------|--------|\n| `FLASHINFER_DUMP_DIR` | `flashinfer_dumps` | Directory where `@flashinfer_api` writes per-call tensor dumps when logging is enabled. |\n| `FLASHINFER_DUMP_INCLUDE` | `\"\"` | Comma-separated allow-list of API names; only matching calls are dumped. Wildcards (`*`) supported. |\n| `FLASHINFER_DUMP_EXCLUDE` | `\"\"` | Comma-separated deny-list of API names; matching calls skip the (expensive) stats / dump path. |\n| `FLASHINFER_DUMP_MAX_COUNT` | `1000` | Hard cap on number of dumped events; once reached, further dumps are dropped with a warning. |\n| `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`). |\n| `FLASHINFER_DUMP_SAFETENSORS` | `0` | `1` writes dumps as `.safetensors` (portable / Hugging Face-style); otherwise plain `.pt`. |\n\n##### Trace Capture\n\nUsed by `flashinfer.trace` / `fi_trace`.\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `FLASHINFER_TRACE_DUMP` | unset | `flashinfer/fi_trace.py` | If set, decorated APIs auto-dump benchmark-definition JSON for each call (the \"fi_trace\" feature). |\n| `FLASHINFER_TRACE_DUMP_DIR` | cwd | `flashinfer/fi_trace.py` | Directory where the trace JSON files are written. |\n| `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. |\n| `FLASHINFER_TRACE_APPLY_PATH` | unset | `flashinfer/trace_apply/config.py` | Directory from which deployment-configured solutions are loaded for Trace Apply. |\n\n##### Validation / Autotuning / Routing / Kernel Selection\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_AUTOTUNER_LOAD_FROM_FILE` | `0` | `flashinfer/autotuner/autotuner.py` | `1` loads previously serialized autotune results from disk instead of re-running the search. |\n| `FLASHINFER_DIST_AWARE_AUTOTUNE` | `0` | `flashinfer/fuse","isInternal":false,"tokens":12791,"sizeBytes":50336},{"name":"marketplace.json","path":".claude-plugin/marketplace.json","rawUrl":"https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/marketplace.json","title":"Claude Plugin Marketplace Catalog","category":"marketplace","format":"json","content":"# Claude Plugin Marketplace Catalog\n\nPath: `.claude-plugin/marketplace.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/marketplace.json)","tokens":300,"sizeBytes":203},{"name":"plugin.json","path":".claude-plugin/plugin.json","rawUrl":"https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/plugin.json","title":"claude-plugin Plugin Manifest","category":"plugin-manifest","format":"json","content":"# claude-plugin Plugin Manifest\n\nPath: `.claude-plugin/plugin.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/plugin.json)","tokens":300,"sizeBytes":189},{"name":"mcp.json","path":".cursor/mcp.json","rawUrl":"https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.cursor/mcp.json","title":"Model Context Protocol (MCP) Configuration","category":"mcp-config","format":"json","content":"# Model Context Protocol (MCP) Configuration\n\nPath: `.cursor/mcp.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.cursor/mcp.json)","tokens":300,"sizeBytes":182}],"systemPromptSnippet":"<yakaai_skills repo=\"flashinfer-ai/flashinfer\">\n<!-- File: AGENTS.md (Tokens: ~45 | Category: root-instruction) -->\nThis repository’s agent instructions live in [CLAUDE.md](./CLAUDE.md).\n\nBefore making changes, read `CLAUDE.md` and follow its guidance for build, test, style, and contribution workflows.\n\n<!-- File: CLAUDE.md (Tokens: ~12791 | Category: claude-rule) -->\n# CLAUDE.md\n\nThis file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.\n\n## Project Overview\n\nFlashInfer 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.\n\n## Quick Reference\n\n| Task | Command |\n|------|---------|\n| Install for development | `pip install --no-build-isolation -e . -v` |\n| Initialize submodules | `git submodule update --init --recursive` |\n| Install CUPTI for benchmarking | `pip install -U cupti-python` |\n| Run regular tests | `pytest tests/` |\n| Run full parameter matrices | `pytest tests/ --full` |\n| Run specific test | `pytest tests/path/test_file.py::test_function` |\n| Run multi-GPU test | `mpirun -np 4 pytest tests/comm/test_allreduce_unified_api.py` |\n| Run benchmark | `python benchmarks/flashinfer_benchmark.py --routine <name> <flags>` |\n| Run linting | `pre-commit run -a` |\n| Dump environment report (bug reports) | `python -m flashinfer.collect_env` (or `flashinfer collect-env [--json]`) |\n| Install pre-commit hooks | `pre-commit install` |\n| Clear JIT cache | `rm -rf ~/.cache/flashinfer/` |\n| Inspect devcontainer compiler cache | `sccache --show-stats` |\n| Enable API logging (basic) | `export FLASHINFER_LOGLEVEL=1` |\n| Enable API logging (detailed) | `export FLASHINFER_LOGLEVEL=3` |\n| Enable API logging (with stats) | `export FLASHINFER_LOGLEVEL=5` |\n| Set API log destination | `export FLASHINFER_LOGDEST=mylog.txt` |\n| Enable verbose JIT logging | `export FLASHINFER_JIT_VERBOSE=1` |\n| Enable debug build | `export FLASHINFER_JIT_DEBUG=1` |\n| Set target architectures | `export FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"` |\n| Set parallel compilation | `export FLASHINFER_NVCC_THREADS=4` |\n| Limit parallel ninja jobs | `export MAX_JOBS=4` |\n| Enable GDN native short-T path | `export FLASHINFER_GDN_WY_NATIVE_T=1` |\n| Enable GDN strided QKV path | `export FLASHINFER_GDN_WY_STRIDED_QKV=1` |\n| Enable GDN native A/B tensors | `export FLASHINFER_GDN_WY_NATIVE_AB=1` |\n| 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` |\n| Opt in to multi-wave KV split in trtllm-gen GQA decode (long or mixed-length KV) | `export FLASHINFER_TRTLLM_GEN_MULTI_WAVE_KV_SPLIT=1` |\n| Let `backend=\"auto\"` pick experimental backends | `export FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1` |\n| Override ragged-prefill `auto` backend order (Blackwell) | `export FLASHINFER_RAGGED_AUTO_BACKEND_ORDER=cutlass,cudnn` |\n| Override CuTe-DSL prefill scheduling | `export FLASHINFER_CUTE_PREFILL_PERSISTENT=0` (non-persistent) or `1` (persistent) |\n| Skip MoE EP CuTe-DSL import/version guard | `export FLASHINFER_MOE_EP_SKIP_DSL_CHECK=1` |\n| Override MoE EP knob-cache path | `export FLASHINFER_MOE_EP_KNOB_CACHE=/path/to/knobs.json` |\n| Disable MoE EP fused staging kernel | `export FLASHINFER_MEGA_FUSED_STAGE=0` |\n| Log MegaMoE compilation and workspace diagnostics | `export FLASHINFER_MEGAMOE_LOG_COMPILE_SPEC=1` |\n| Enable MoE EP EXPERT_MAJOR padding-row exclusion (opt-in; shape-dependent) | `export FLASHINFER_MOE_EP_EXCLUDE_PADDING_ROWS=1` |\n| Enable distribution-aware MoE autotune and kernel dispatch (experimental; TRT-LLM MoE only) | `export FLASHINFER_DIST_AWARE_AUTOTUNE=1` |\n| Prims-TS debug checks / task-manager verify | `export FLASHINFER_PRIMS_TS_DEBUG_CHECKS=1` |\n| Override Prims-TS CuTe-DSL compile options | `export FLASHINFER_PRIMS_TS_COMPILE_OPTIONS=\"--opt-level 2\"` |\n| Configure NVFP4 4over6 quantization (preferred over the env vars) | `nvfp4_quantize(..., nvfp4_4over6=NVFP44Over6Config(e4m3_max=256))` |\n\nThe minimum Python version used by CI and build tooling is defined in\n`.python-version`. CI Docker images use a stable Conda environment name and\nderive their interpreter and site-packages paths from that file. When raising\nthe minimum, also update `project.requires-python` and `tool.mypy.python_version`\nin `pyproject.toml`.\n\n## Quick Start for Development\n\n### Installation\n\n```bash\ngit clone https://github.com/flashinfer-ai/flashinfer.git --recursive\ncd flashinfer\npip install --no-build-isolation -e . -v\n```\n\n**Important**: The `--recursive` flag is required to initialize submodules in `3rdparty/` (cutlass, spdlog).\n\nIf you forgot `--recursive` when cloning:\n```bash\ngit submodule update --init --recursive\n```\n\nThat's it! You can now:\n\n- Run all benchmarks and unit tests\n- Modify kernel source code in `include/` without reinstalling\n- Changes are JIT-compiled on next use\n\nThe `--no-build-isolation` flag prevents pip from pulling incompatible PyTorch/CUDA versions from PyPI.\n\n### How JIT Compilation Works\n\nWhen you call a FlashInfer API:\n\n1. **First call**: Generates specialized CUDA code based on parameters (dtype, head_dim, etc.), compiles it with ninja, caches the .so file\n2. **Subsequent calls**: Uses cached compiled module\n3. **After kernel changes**: Automatically detects changes and recompiles\n\n**No manual rebuild step needed** - just edit `.cuh` files and run your code again.\n\n### Pre-compiled Packages (Optional)\n\nFlashInfer provides optional pre-compiled packages for users who want faster initialization:\n\n- `flashinfer-jit-cache`: Pre-built kernel cache\n- `flashinfer-cubin`: Pre-compiled kernel binaries\n\n**For development, you typically DON'T need these.** JIT compilation is fast enough and gives you live code reload.\n\n### Compiler Caching in Devcontainers\n\nThe devcontainers run every nvcc and host C++ compile through\n[sccache](https://github.com/mozilla/sccache) by setting `FLASHINFER_NVCC_LAUNCHER`\nand `FLASHINFER_CXX_LAUNCHER`. Cached objects live in a per-devcontainer Docker\nvolume mounted at `~/.cache/sccache`, so they outlive a container rebuild — which\nwipes `~/.cache/flashinfer` and would otherwise force a full recompile.\n\n```bash\nsccache --show-stats                  # hit rate, cache size, uncacheable reasons\nexport FLASHINFER_NVCC_LAUNCHER=\"\"    # disable for the current shell\n```\n\nCaveats:\n\n- Toggling a launcher changes the compile commands in `build.ninja`, so the first\n  build after switching it on or off recompiles objects that already exist.\n- Link steps, cubin embedding, and CuTe-DSL kernels never invoke nvcc, so they are\n  not cached.\n- The cache is per-devcontainer by design: sccache's local storage supports a\n  single server at a time, so two containers must not share one cache directory.\n\nGitHub PR Test GPU jobs keep a different, per-runner cache: they set\n`CI_COMPILE_CACHE_DIR` to a directory under the runner's work directory, and\n`ci/bash.sh` namespaces it by runner and the pulled image's immutable ID.\n`CI_COMPILE_CACHE_MODE=read-only` (the default) mounts a read-only seed and runs\nagainst a private writable copy, discarded even after test failures. The PR Test\nworkflow selects `read-write` only for pushes to `main`. A stable per-runner lock\nprotects maintenance and execution, including accidentally shared workspaces.\nOther `ci/bash.sh` callers leave caching unset and run cold.\nBefore each trusted writable job its runner's cache is cleared when it exceeds\n`CI_COMPILE_CACHE_MAX_MB`\n(default 40960) or its volume has less than `CI_COMPILE_CACHE_MIN_FREE_MB`\n(default 20480) free. Only content-keyed artifacts survive between jobs: CuTe-DSL (with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1`), TIRx, KDA PTX, Triton and downloaded\ncubins. nvcc modules, tuning results and the generated-source tree are pruned,\nsince a fresh checkout rebuilds nvcc modules anyway. Each test job's\n`FLASHINFER_JIT_BUILD_LOG` is summarized per toolchain by\n`scripts/summarize_jit_builds.py` in the job summary.\nSee [`ci/COMPILE_CACHE.md`](ci/COMPILE_CACHE.md) for lifecycle tests and trust\nboundaries.\n\n## cuDNN linear attention\n\nUse `nvidia-cudnn-frontend>=1.31.0`; KDA and GDN normalization use additive\n`1e-6` without an epsilon argument. Auto selection also requires CuTe DSL 4.7+.\nThe bounded KDA SM100/SM103 and GDN SM100 policies live in `kda.py` and\n`gdn_prefill.py`. KDA's head/length regions are documented in `docs/api/kda.rst`.\nAuto-selected cuDNN ignores the native KDA workspace; native fallback uses it.\nValidate output and final state before changing a policy. Explicit state-pool calls\nrequire unique, in-range slots; auto keeps pools on their native route.\n\nUse `benchmarks/bench_cudnn_linear_attention.py` for public-call comparisons.\nIts raw-gate GDN baseline includes Torch gate transforms; it is not a vLLM\nfused-producer baseline. Keep host CPU, completed wall and reset-inclusive\nCUDA graph times separate.\n\n## Testing\n\nRun the regular test suite:\n\n```bash\npytest tests/\n```\n\nRun full parameter matrices (used by nightly testing):\n\n```bash\npytest tests/ --full\n```\n\nRun specific test file:\n\n```bash\npytest tests/attention/test_hopper.py\n```\n\nRun specific test function:\n\n```bash\npytest tests/attention/test_hopper.py::test_single_prefill\n```\n\nThe sharded test runner defaults to one worker per visible GPU. H100 PR CI opts\ninto `--workers-per-gpu 2`: eight workers on four GPUs. Other CI lanes retain\nthe default. Locally, `scripts/task_run_unit_tests.sh --workers-per-gpu 2` enables\nsharing; `--workers N` sets the total instead. Workers share the existing JIT\ncompilation budget.\n\nScheduling markers apply to the whole source file:\n\n- `node_exclusive`: run alone across the allocation after parallel work, with\n  all GPUs visible. Use for multi-GPU requirements or host-memory OOM risk.\n- `gpu_exclusive`: run alone on the assigned GPU; other GPUs can keep working.\n- `host_memory_heavy`: prioritize this file and serialize its units with other\n  host-heavy work, while lighter files can overlap.\n\nHost and GPU reservations cover setup, execution, and process cleanup, and are\nacquired together. Blocked units stay queued so workers can take runnable work\nfrom their own queue or steal from another. Pending GPU-exclusive work can\ndrain one GPU without reserving the host-heavy slot or blocking other GPUs.\n`RESOURCE WAIT` reports worker idle time of at least one second while work is\nqueued but cannot run on that worker's resources.\n\nThese markers enforce overlap rules, not memory limits. Most H100 annotations\nuse measured peaks above 128 GiB host RSS or 32 GiB GPU memory. Attention sink\nalso uses both narrower markers: its single-GPU H100 run peaked at 89.2 GiB host\nRSS and 4.5 GiB GPU memory, and other GPUs can keep working while it runs.\nCold JIT compilation and other GPU families can change the peaks. Validate\ncompletion rate, OOM failures, and repeated wall times with matching scope and\ncache conditions when changing a CI lane's concurrency policy.\n\nCapacity reports include the serial floor from host-heavy and node-exclusive\nwork. They remain optimistic about GPU contention because the estimator uses\nworker counts without GPU topology or measured memory capacity.\n\nRun the CPU-only coordinator suite with\n`pytest --confcutdir=tests/test_sharding tests/test_sharding/`.\n\nCommunication-test wrappers optionally emit cache-publication completion reports\nwhen `UNIT_TEST_RUN_STATUS_DIR` is set to a fresh directory unique to that wrapper\nand container/rank. `test_run_status.sh` records each pytest invocation, and\n`test_run_status.py` atomically writes `script-result.json` (schema version 1,\n`complete`, `exit_code`, `tests_executed`, and `infrastructure_errors`). A small\npytest plugin records execution and collection/internal errors inside each MPI\nor torchrun worker; launcher exit 1 alone is not evidence of an assertion failure.\nAll local worker ranks must finish, and at least one test must execute, to allow\npublication. A completed run with assertion failures remains eligible, while\ncrashes, timeouts, and unfinished invocations do not. With the variable unset,\nreporting is disabled. The CPU suite above covers this contract without GPU\ndependencies.\n\n### Skipping Tests Based on CUDA Architecture\n\nUse `flashinfer.utils` functions to skip tests on unsupported GPU architectures:\n\n**Available check functions:**\n- `get_compute_capability(device)` - Returns `(major, minor)` tuple\n- `is_sm90a_supported()` - Hopper (requires CUDA 12.3+)\n- `is_sm100a_supported()` - Blackwell (requires CUDA 12.8+)\n- `is_sm100f_supported()` - Blackwell feature-set (requires CUDA 12.9+)\n- `is_sm110a_supported()`, `is_sm120a_supported()`, `is_sm120f_supported()`, `is_sm121a_supported()`\n\n**APIs decorated with `@backend_requirement`** also provide:\n- `api_name.is_compute_capability_supported(cc)` - e.g., `mm_fp4.is_compute_capability_supported(100)`\n- `api_name.is_backend_supported(\"backend\")` - e.g., `mm_fp4.is_backend_supported(\"cudnn\")`\n- `api_name.is_backend_supported(\"backend\", cc)` - e.g., `mm_fp4.is_backend_supported(\"cudnn\", 80)`\n\n**Example:**\n```python\nfrom flashinfer.utils import is_sm90a_supported\n\ndef test_hopper_attention():\n    if not is_sm90a_supported(torch.device(\"cuda\")):\n        pytest.skip(\"Requires SM90a\")\n    # Test code...\n```\n\n**Common requirements:**\n\n| Feature | Min SM | Check Function |\n|---------|--------|----------------|\n| FlashAttention-3 | SM90a | `is_sm90a_supported()` |\n| MLA Attention | SM100a | `is_sm100a_supported()` |\n| FP8 GEMM | SM89+ | `get_compute_capability()[0] >= 9` |\n\n**Note:** `tests/conftest.py` auto-skips tests that trigger OOM, but tests should be written to avoid OOM by using appropriate problem sizes.\n\n## Benchmarking\n\nFor SM90 NVFP4 fused/split and MXFP4 fused MegaMoE, use\n[`benchmarks/bench_megamoe_fp4_sm90.py`](benchmarks/bench_megamoe_fp4_sm90.py) for timing\nand [`tests/moe/test_megamoe_fp4_sm90.py`](tests/moe/test_megamoe_fp4_sm90.py) for correctness checks.\n\nFlashInfer provides a unified benchmarking framework in `benchmarks/flashinfer_benchmark.py`.\n\n**Key features:**\n- Supports attention, GEMM, and MOE kernels\n- Multiple backends: FlashAttention2/3, cuDNN, CUTLASS, TensorRT-LLM, cuBLAS\n- **CUPTI timing (recommended)**: Hardware-level profiling for accurate GPU kernel time\n  - Automatically falls back to CUDA events if CUPTI unavailable\n  - Install: `pip install -U cupti-python` (requires CUDA 13+)\n- Batch testing, reference checking, CSV output\n\n**Quick example:**\n```bash\npython benchmarks/flashinfer_benchmark.py \\\n    --routine BatchDecodeWithPagedKVCacheWrapper \\\n    --backends fa2 cudnn \\\n    --batch_size 32 --s_kv 2048 \\\n    --num_qo_heads 32 --num_kv_heads 8 \\\n    --head_dim_qk 128 --head_dim_vo 128 \\\n    --page_size 16 --refcheck -vv\n```\n\n**Python API:**\n```python\nfrom flashinfer.testing import bench_gpu_time\n\n# CUPTI preferred, auto-fallback to CUDA events.\n# Returns per-iteration times in milliseconds (a list).\ntimes = bench_gpu_time(\n    my_kernel, input_args=(x, y), enable_cupti=True,\n    dry_run_iters=10, repeat_iters=30,\n)\nmedian_time_ms = statistics.median(times)\n```\n\n→ **For complete benchmarking guide, see [`.claude/skills/benchmark-kernel/SKILL.md`](.claude/skills/benchmark-kernel/SKILL.md)**\n\n`flashinfer.gemm.group_gemm_fp8_nt_groupwise_contiguous` is the standalone\nCuTe-DSL grouped FP8 API; `group_deepgemm_fp8_nt_groupwise` keeps its original\nDeepGEMM-only signature. Both use per-row expert indices, unlike the indptr-based\n`group_gemm_fp8_nt_groupwise` API.\n\nFor contiguous grouped FP8 GEMM, `python benchmarks/bench_grouped_fp8.py\n--production-shapes` compares DeepGEMM and CuTe-DSL with preallocated outputs;\n`--cache-probe` measures first-call and new-token-count compilation separately.\nCompiled grouped CuTe-DSL kernels are reused across token counts with the same\n128-row alignment class, with at most two variants per device and weight shape.\nConcurrent misses share one compilation; warm hits do not acquire the compile lock.\nThe standalone API's optional `validate_indices=True` checks expert-index values\nand synchronizes with the CPU. Validate new routing data before CUDA graph\ncapture; the default path assumes valid indices and performs metadata checks only.\n\nPrims-TS decode keeps raw storage dtypes (`q_dtype`, `k_dtype`, `v_dtype`)\nseparate from effective MMA dtypes (`qk_dtype`, `pv_dtype`). Transformed\nFP8/NVFP4 K/V uses Q's compute dtype; BF16 Q/K with FP8 V uses FP8 PV without\nthat transform. Kernel predicates, P packing, and attention-sink scaling must\nfollow effective compute precision; TMA descriptors and raw allocations must\nfollow storage precision. The public plan accepts separate `k_data_type` and\n`v_data_type`, with `kv_data_type` retained as a common-KV compatibility alias.\n\n## Code Linting\n\nRun all pre-commit hooks:\n\n```bash\npre-commit run -a\n```\n\nInstall hooks to run on every commit:\n\n```bash\npre-commit install\n```\n\n## Opening a Pull Request\n\nFollow [`CONTRIBUTING.md`](CONTRIBUTING.md) when opening a PR. The two requirements most often\nmissed:\n\n- **Keep the default PR template.** Fill in `.github/pull_request_template.md` (Description,\n  Related Issues, Checklist, Tests, Reviewer Notes). Do **not** overwrite or replace it with a\n  custom or tool-generated description format. The title and description normally become the\n  commit title and message on squash-merge, and are what a `git bisect` surfaces months later\n  when someone is hunting the owner of a regression.\n- **Report before/after numbers for performance work.** A perf PR must include measurements from\n  a reproducible benchmark (e.g. `benchmarks/flashinfer_benchmark.py`), naming the GPU and the\n  problem sizes. Speedup ratios without absolute numbers, or numbers without a named GPU, are\n  not enough.\n\nWhen preparing changes for a PR, a self-review is strongly recommended: follow\n[`.claude/skills/self-review/SKILL.md`](.claude/skills/self-review/SKILL.md), which walks the\ndiff through the repository's review guidance and PR rules. It is informal and not required —\nits purpose is to make review easier for reviewers and contributors alike by catching common\nissues before a human looks. Changes made only to test something locally, with no PR intended,\ndo not need one.\n\n→ **For the complete contribution rules, see [`CONTRIBUTING.md`](CONTRIBUTING.md)**\n\n## Code Review\n\nWhen reviewing a diff (as an agent or a human), follow the shared focus areas, kernel-review\npolicy, and effort calibration in [`docs/code_review_guidance.md`](docs/code_review_guidance.md).\nNote: unlike human review, agents keep **kernel implementation details in scope** — read the\nkernel logic and report bugs, labeling findings by confidence.\n\n→ **For the complete review rule set, see [`docs/code_review_guidance.md`](docs/code_review_guidance.md)**\n\n## Architecture: JIT Compilation System\n\nFlashInfer's JIT system has three layers:\n\n### Layer 1: JitSpec (flashinfer/jit/core.py)\n\n`JitSpec` is an abstract base class defining the kernel-module lifecycle\n(`try_load()` / `build()` / `load()`, with the shared `build_and_load()`\ntemplate method handling caching, locking, and `FLASHINFER_DISABLE_JIT`).\nOne subclass per compilation toolchain:\n\n- `JitSpecNvcc` (nvcc/ninja modules, returned by `gen_jit_spec()`) defines:\n  - `name`: Unique identifier (URI hash from parameters)\n  - `sources`: List of .cu/.cpp files to compile\n  - `extra_cuda_cflags`, `extra_cflags`, `extra_ldflags`: Compiler flags\n- `JitSpecCuteDsl` (flashinfer/jit/cute_dsl_core.py) caches CuTe-DSL kernels\n  (see \"CuTe-DSL kernels\" under Module Caching below)\n- The optional TIRx KDA backend uses a private `JitSpec` in\n  `flashinfer/kda_kernels/tirx/cache.py` to cache TVM modules, with source/compiler\n  fingerprints and artifact checksums. Its TVM/TIRx dependencies are loaded only\n  by explicit `recurrent_kda(..., backend=\"tirx\")`; see `docs/api/kda.rst`.\n\n### JIT Directory Rules\n\n**NEVER write to package directories** - they may be read-only after installation.\n\n| Directory | Writable | Use for |\n|-----------|----------|---------|\n| `FLASHINFER_GEN_SRC_DIR` | ✓ Yes | Generated source files (Jinja output, copied .cu files) |\n| `FLASHINFER_JIT_DIR` | ✓ Yes | Compiled `.so` outputs |\n| `FLASHINFER_CSRC_DIR` | ✗ No | Read-only source templates |\n| `FLASHINFER_AOT_DIR` | ✗ No | Read-only pre-compiled binaries |\n\n### Compilation Context: Architecture-Specific Compilation\n\nFlashInfer uses `CompilationContext` to manage CUDA architecture targets. Some kernels only work on specific GPU architectures (e.g., Hopper SM90, Blackwell SM100/SM12x).\n\n**How it works:**\n- Auto-detects GPUs in system or reads `FLASHINFER_CUDA_ARCH_LIST` environment variable\n- CuTe-DSL FP4 MM compilation parallelism can be controlled with `FLASHINFER_MM_FP4_CUTE_DSL_COMPILE_WORKERS`\n- JIT modules specify `supported_major_versions=[9, 10, 11, 12]` to limit compilation to specific SM versions\n- If GPU not supported → `RuntimeError: No supported CUDA architectures found`\n\n→ **See [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md) for usage examples**\n\n### Layer 2: Code Generation\n\nEvery `gen_*_module()` function in `flashinfer/jit/` follows this pattern:\n\n```python\ndef gen_some_module(dtype_in, dtype_out, ...):\n    # 1. Compute unique identifier from parameters\n    uri = get_some_uri(dtype_in, dtype_out, ...)\n\n    # 2. Create generation directory\n    gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri\n\n    # 3. (Optional) Render Jinja template to generate type-specialized config\n    # Skip this step if you don't need type specialization\n    with open(jit_env.FLASHINFER_CSRC_DIR / \"some_customize_config.jinja\") as f:\n        template = jinja2.Template(f.read())\n    config_content = template.render(\n        dtype_in=dtype_map[dtype_in],\n        dtype_out=dtype_map[dtype_out],\n        # ... more parameters\n    )\n    write_if_different(gen_directory / \"some_config.inc\", config_content)\n\n    # 4. Copy source files to gen directory\n    sources = []\n    for fname in [\"some_kernel.cu\", \"some_jit_binding.cu\"]:\n        shutil.copy(jit_env.FLASHINFER_CSRC_DIR / fname, gen_directory / fname)\n        sources.append(gen_directory / fname)\n\n    # 5. Return JitSpec\n    return gen_jit_spec(uri, sources, extra_cuda_cflags=[...])\n```\n\n**Note**: If your operation doesn't need type specialization, you can skip step 3 entirely and just copy the source files directly.\n\n### Layer 3: Compilation and Loading\n\n`JitSpecNvcc` methods:\n\n- `write_ninja()` - Generates `build.ninja` file\n- `build()` - Executes `ninja` to compile sources\n- `build_and_load()` - Compiles and loads via TVM-FFI\n\nThe generated `build.ninja` file uses nvcc to compile .cu → .cuda.o → .so, then loads via TVM-FFI.\n\n### Jinja Templates (Optional)\n\n**Note: Jinja templates are NOT required.** You can write C++ code directly without templating.\n\nFor operations that need type specialization, templates in `csrc/*.jinja` can generate C++ code:\n\n```jinja\n// Input template\nusing DTypeIn = {{ dtype_in }};\nusing DTypeOut = {{ dtype_out }};\nconstexpr int PARAM = {{ param_value }};\n\n// After render\nusing DTypeIn = float16;\nusing DTypeOut = float16;\nconstexpr int PARAM = 128;\n```\n\nThis 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.\n\n## Directory Structure\n\n```\nflashinfer/\n├── include/flashinfer/           # Header-only CUDA kernel templates\n│   ├── attention/                # Attention kernels\n│   ├── gemm/                     # GEMM kernels\n│   ├── comm/                     # Communication kernels\n│   ├── mma.cuh                   # Matrix multiply utilities\n│   ├── utils.cuh                 # Common utilities\n│   └── [...]\n│\n├── csrc/                          # Framework bindings (via TVM-FFI)\n│   ├── *.cu                       # Kernel launcher implementations\n│   ├── *_jit_binding.cu           # TVM-FFI exports\n│   ├── *_customize_config.jinja   # Type config templates (optional)\n│   └── [...]\n│\n├── flashinfer/                    # Python package\n│   ├── jit/\n│   │   ├── core.py                # JitSpec, compilation infrastructure\n│   │   ├── cpp_ext.py             # Ninja build generation\n│   │   ├── env.py                 # Workspace paths\n│   │   ├── attention/             # Attention module generators\n│   │   ├── gemm/                  # GEMM module generators\n│   │   ├── fused_moe/             # MOE module generators\n│   │   └── [...]\n│   ├── gemm/                      # GEMM Python APIs\n│   ├── fused_moe/                 # MOE Python APIs\n│   ├── comm/                      # Communication Python APIs\n│   ├── *.py                       # Other high-level Python APIs\n│   ├── aot.py                     # AOT compilation for pre-built packages\n│   └── [...]\n│\n├── tests/                         # Test suite\n│   ├── attention/                 # Attention kernel tests\n│   ├── gemm/                      # GEMM kernel tests\n│   ├── moe/                       # MOE kernel tests\n│   ├── comm/                      # Communication tests\n│   ├── utils/                     # Utility tests\n│   └── conftest.py                # Pytest configuration\n│\n└── build_backend.py               # PEP 517 build backend\n```\n\n### Critical Rule: Framework Separation\n\n**Torch headers MUST NOT be included in `include/` directory files.**\n\n- `include/`: Framework-agnostic CUDA kernels (accept raw pointers)\n- `csrc/`: Framework bindings via TVM-FFI (currently PyTorch, but can support other frameworks)\n\n## Adding a New Operation\n\n→ **For complete step-by-step tutorial, see [`.claude/skills/add-cuda-kernel/SKILL.md`](.claude/skills/add-cuda-kernel/SKILL.md)**\n\n**Quick overview of the process:**\n1. Write kernel in `include/flashinfer/new_op.cuh` (framework-agnostic, raw pointers)\n2. Write launcher in `csrc/new_op.cu` (PyTorch tensor handling)\n3. Create TVM-FFI bindings in `csrc/new_op_jit_binding.cu`\n4. (Optional) Create Jinja template for type specialization\n5. Write JIT module generator in `flashinfer/jit/new_op.py`\n6. Write Python API in `flashinfer/new_op.py` with `@functools.cache`\n7. Write tests in `tests/`\n8. Register in `flashinfer/aot.py` for AOT compilation\n9. Export in `flashinfer/__init__.py`\n10. Add a `TraceTemplate` in `flashinfer/trace/templates/` and wire it via `@flashinfer_api(trace=...)` (see below)\n11. Add an example call in `tests/trace/example.py`, re-run to regenerate `fi_trace_out/`, and commit the new JSON files\n\n### Trace Template Checklist (for new or updated APIs)\n\nEvery 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.\n\n1. **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.\n2. **Wire the template** to the API: `@flashinfer_api(trace=my_trace)` on the Python function (or class method's `run()`).\n3. **Add an example call** in `tests/trace/example.py` that exercises the new trace with realistic shapes.\n4. **Regenerate examples**: `rm -rf tests/trace/fi_trace_out && python tests/trace/example.py` — verify the expected JSON appears.\n5. **Update the docstring** in `tests/trace/example.py` to list the new file(s).\n6. **Run tests**: `pytest tests/trace/ -v` — all template-consistency and end-to-end tests must pass.\n7. **Commit the new JSON files** under `tests/trace/fi_trace_out/` alongside the code changes.\n\n**Example implementations:**\n- **Simple**: `flashinfer/norm/__init__.py` (RMSNorm) - no Jinja, good starting point\n- **Moderate**: `flashinfer/sampling.py` - with Jinja templating\n- **Complex**: `flashinfer/decode.py` - plan-run pattern, advanced workspace\n\n## Experimental Code\n\nExperimental backends live under `flashinfer/experimental/`; experimental\npublic APIs live in core marked with `@flashinfer_experimental_api` (defined in\n`flashinfer/api_logging.py`). Using them is an explicit opt-in: calling an\nexperimental API, or naming an experimental backend with `backend=\"<name>\"`,\nneeds no environment variable (both warn once). Only automatic selection is\ngated: `backend=\"auto\"` skips backends marked `@experimental_backend` unless\n`FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. Experimental code is JIT-only\n(never registered in `flashinfer/aot.py`) and tested under\n`tests/experimental/`.\n\n→ **When working under `flashinfer/experimental/`, follow\n[`flashinfer/experimental/CLAUDE.md`](flashinfer/experimental/CLAUDE.md); the\nfull policy is [`flashinfer/experimental/README.md`](flashinfer/experimental/README.md).**\n\n## Key Architectural Patterns\n\n### Module Caching\n\nFlashInfer uses two-level caching to avoid recompilation:\n\n1. **Python-level** (`@functools.cache`): In-memory cache of loaded modules\n2. **File-level** (`~/.cache/flashinfer/`): Compiled `.so` files on disk\n\n**Cache invalidation** (automatic):\n- Source file changes (SHA256 hash)\n- Compilation flags change\n- CUDA architecture change\n- FlashInfer version change\n\nURI computed as: `hash(operation_type + parameters + source_hashes + flags + cuda_arch)`\n\n**Cache management:**\n- Clear cache: `rm -rf ~/.cache/flashinfer/`\n- Override location: `export FLASHINFER_WORKSPACE_BASE=\"/scratch\"`\n\n**CuTe-DSL kernels** (`flashinfer/jit/cute_dsl_core.py`): `cute.compile()` has no\npersistent cache, so `JitSpecCuteDsl` (a `JitSpec` subclass, wrapped by the\n`build_and_load_cute_dsl_kernel()` helper) exports compiled kernels as\nobject files (`export_to_c()`) and reloads them with `cute.runtime.load_module(...,\nenable_tvm_ffi=True)`. Artifacts live in `cached_ops/` next to the nvcc modules, one\ndirectory per op family — `cached_ops/<module>_<arch>_cute_dsl/` (e.g.\n`nvfp4_quantize_sm100a_cute_dsl/`) holding one `meta.json` plus one `.o` per\nspecialization. The arch comes from the DSL's compile target (`CUTE_DSL_ARCH` or the\ncurrent device) since the artifacts are single-arch, unlike nvcc fatbins.\nInvalidation is module-granular: changed DSL/Torch/TVM-FFI distributions, Python\nABI, or\nkernel-source SHA256 wipes and lazily rebuilds the module. Reference usage:\n`flashinfer/quantization/kernels/nvfp4_quantize.py`. Disable with\n`FLASHINFER_CUTE_DSL_DISABLE_CACHE=1`. The SHA256 covers only the files a caller\npasses as `extra_key_files` (often just `__file__`); with\n`FLASHINFER_CUTE_DSL_STRICT_CACHE_KEY=1` it also covers every flashinfer module\nthose files can run. Source analysis and hashing live in\n`flashinfer/jit/_cute_dsl_source.py`; `cute_dsl_core.py` manages compilation and\ncached artifacts. Plain modules are keyed by content with everything they\nimport, including literal `importlib.import_module` calls. Package initializers\nare resolved name by name: a name looked up in a package follows the binding\nthat defines it (a re-export to its source, a function into its body, a module\nthrough the attributes read from it) plus the submodule of that name. The code\nan initializer runs at import, including evaluated function annotations, is\nfollowed for every keyed file in its package. Star imports remain covered when\nthey can replace an earlier binding; lazy fallbacks are followed when a binding\nmay be absent.\nEach initializer is keyed by that code and the import lines of used names, so a\nnew re-export or `__all__` entry invalidates nothing. Not keyed: import-time\nside effects of modules that keyed code neither imports nor uses, nonliteral\ndynamic imports, and runtime file reads; list the latter in `extra_key_files`.\n`version.py` and `_build_meta.py` are excluded to avoid keying on the checkout\ncommit. This is intended for long-lived caches shared across commits.\nTask-scheduled attention shares this cache through\n`flashinfer/attention/prims_ts/_compile_cache.py`; its keys cover compile specs,\noptions, and embedded device-capacity constants. It hashes the attention source\ntree, including shared preparation, masks, and reductions. Occupancy queries\nare cached per device and cluster size by `_hardware.py`.\n\n### Dispatch Macros\n\nHandle combinatorial parameter spaces:\n\n```cpp\nDISPATCH_DTYPE(input_dtype, DTypeIn, {\n  DISPATCH_DTYPE(output_dtype, DTypeOut, {\n    DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, {\n      LaunchKernel<DTypeIn, DTypeOut, BLOCK_SIZE>(...);\n    });\n  });\n});\n```\n\nDefined in `.jinja` files and expanded after rendering.\n\n## API Logging with @flashinfer_api\n\nFlashInfer provides the `@flashinfer_api` decorator for debugging API calls.\n\n**Key features:**\n- **Crash-safe**: Logs inputs BEFORE execution (preserves info even if kernel crashes)\n- **Zero overhead when disabled**: `FLASHINFER_LOGLEVEL=0` (default)\n- **Multiple verbosity levels**: 0 (off), 1 (names), 3 (inputs/outputs), 5 (+ statistics)\n- **CUDA graph compatible**: Auto-skips stats during graph capture\n\n**Quick usage:**\n```bash\n# Enable detailed logging\nexport FLASHINFER_LOGLEVEL=3              # 0, 1, 3, or 5\nexport FLASHINFER_LOGDEST=debug.log       # stdout, stderr, or file path\n\npython my_script.py\n```\n\n**Why use this?**\n- Debug CUDA crashes (see inputs that caused crash)\n- Track tensor shapes/dtypes through pipeline\n- Detect NaN/Inf issues (level 5)\n\n→ **For complete debugging guide, see [`.claude/skills/debug-cuda-crash/SKILL.md`](.claude/skills/debug-cuda-crash/SKILL.md)**\n\n## Debugging\n\n### Enable Logging\n\n```bash\nexport FLASHINFER_JIT_VERBOSE=1      # Verbose JIT output\nexport FLASHINFER_JIT_DEBUG=1        # Debug symbols, -O0\nexport FLASHINFER_LOGLEVEL=3         # API logging (0=off, 1=basic, 3=detailed)\nexport FLASHINFER_LOGDEST=stdout\n```\n\n### Inspect Generated Code\n\n```bash\n# Generated sources\nls -la ~/.cache/flashinfer/0.6.0/*/generated/\n\n# Compiled modules\nls -la ~/.cache/flashinfer/0.6.0/*/cached_ops/\n\n# Build files\ncat ~/.cache/flashinfer/0.6.0/*/cached_ops/*/build.ninja\n```\n\n### Environment Variables\n\n```bash\n# Compilation\nexport FLASHINFER_NVCC_THREADS=4              # Threads per nvcc process (--threads=N)\nexport MAX_JOBS=4                             # Parallel ninja jobs (nvcc processes)\nexport FLASHINFER_CUDA_ARCH_LIST=\"8.0 9.0a\"  # Target architectures\n# Memory note: total compilation memory ≈ MAX_JOBS × FLASHINFER_NVCC_THREADS × per-thread mem.\n\n# Behavior\nexport FLASHINFER_WORKSPACE_BASE=\"/scratch\"   # Custom cache directory\n```\n\n#### Full Environment Variable Reference\n\nAlready covered above (Quick Reference / Debugging): `FLASHINFER_LOGLEVEL`,\n`FLASHINFER_LOGDEST`, `FLASHINFER_JIT_VERBOSE`, `FLASHINFER_JIT_DEBUG`,\n`FLASHINFER_CUDA_ARCH_LIST`, `FLASHINFER_NVCC_THREADS`,\n`FLASHINFER_WORKSPACE_BASE`, `MAX_JOBS`.\n\nThe remaining `FLASHINFER_*` knobs read by the code are listed below. Defaults\nmatch what the code uses today; values are strings unless noted.\n\n##### JIT / Build Toolchain\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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`). |\n| `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. |\n| `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`. |\n| `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. |\n| `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. |\n| `FLASHINFER_JIT_LINEINFO` | `0` | `flashinfer/jit/core.py` | `1` adds `-lineinfo` to nvcc so profiler / `cuda-gdb` can map PTX back to CUDA source. |\n| `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. |\n| `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). |\n| `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`. |\n| `FLASHINFER_CXX_LAUNCHER` | `\"\"` (`sccache` in the devcontainers) | `flashinfer/jit/cpp_ext.py` | Same idea as `FLASHINFER_NVCC_LAUNCHER` but for the host C++ compiler. |\n| `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. |\n| `FLASHINFER_EXTRA_CFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to the host C++ compiler. |\n| `FLASHINFER_EXTRA_CUDAFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra compiler flags passed to `nvcc`. |\n| `FLASHINFER_EXTRA_LDFLAGS` | unset | `flashinfer/jit/cpp_ext.py` | Extra linker flags passed to the linker. |\n| `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. |\n| `FLASHINFER_CAKE_DSA_CHECK_DST_MAP` | `0` | `flashinfer/dsa_sparse_attention.py` | Fallback of the `check_dst_map` keyword of `dsa_sparse_attention` / `dsa_sparse_attention_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. |\n| `FLASHINFER_CAKE_DSA_DKV_DIRECT` | `auto` | `flashinfer/dsa_sparse_attention.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. |\n| `FLASHINFER_CAKE_DSA_TRAIN_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/dsa_sparse_attention.py` | Capacity of the CAKE DSA training binding cache (`cake_backend.BINDING_CACHE`, least recently used first out): read once when `flashinfer.dsa_sparse_attention` 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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_GEMM_TUNING` | unset | `flashinfer/chunked_lm_head.py` | Fallback of the `gemm_tuning` keyword of `chunked_lm_head_loss` / `chunked_lm_head_logprob` 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`. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_FUSE_DW_CAST` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DX_FINALIZE` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_DW_STREAM` | `auto` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_HIDDEN_COUNT` | `1` | `flashinfer/chunked_lm_head.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. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_COMPACT_ROWS` | `1` | `flashinfer/chunked_lm_head.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`). |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE` | `1` | `flashinfer/chunked_lm_head.py` | Read once when `flashinfer.chunked_lm_head` is imported and applied through `cake_backend.BINDING_CACHE.configure` when the wrapper first reaches the backend (the first call of either entry point): `0` disables the per-process LRU cache of compiled bindings. Direct users of the package call `configure` themselves. |\n| `FLASHINFER_CAKE_LM_HEAD_LOSS_BINDING_CACHE_CAPACITY` | `64` | `flashinfer/chunked_lm_head.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. |\n| `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. |\n| `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. |\n\nRelease/nightly JIT-cache provider builds run sccache in server-side mode and\nwrap the complete wheel build with a no-output watchdog. After 90 minutes with\nno build output, the watchdog captures system, process, compiler, and sccache\nstate; terminates the wrapped build process group (TERM, then KILL after a\ntwo-minute grace period); and exits with status 124. Compiler discovery scans\nthe whole build container, so diagnostics include compiler processes owned by\nthe sccache daemon even though they are outside the wrapped process group. The\nCI `docker run --rm` boundary provides final cleanup for those processes.\n\n##### Cubin / Artifact Loader\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_CUBIN_CHECKSUM_DISABLED` | unset | `flashinfer/jit/cubin_loader.py` | If set, skip SHA checksum verification of downloaded cubins. Debug aid only. |\n| `FLASHINFER_CUBIN_DOWNLOAD_THREADS` | `4` | `flashinfer/artifacts.py` | Thread-pool size used by `flashinfer artifacts download`. |\n| `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. |\n| `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. |\n| `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. |\n| `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`. |\n| `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. |\n\n##### API Dump / Logging Extensions\n\nThese complement `FLASHINFER_LOGLEVEL`/`FLASHINFER_LOGDEST` and are all read in `flashinfer/api_logging.py`.\n\n| Variable | Default | Effect |\n|----------|---------|--------|\n| `FLASHINFER_DUMP_DIR` | `flashinfer_dumps` | Directory where `@flashinfer_api` writes per-call tensor dumps when logging is enabled. |\n| `FLASHINFER_DUMP_INCLUDE` | `\"\"` | Comma-separated allow-list of API names; only matching calls are dumped. Wildcards (`*`) supported. |\n| `FLASHINFER_DUMP_EXCLUDE` | `\"\"` | Comma-separated deny-list of API names; matching calls skip the (expensive) stats / dump path. |\n| `FLASHINFER_DUMP_MAX_COUNT` | `1000` | Hard cap on number of dumped events; once reached, further dumps are dropped with a warning. |\n| `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`). |\n| `FLASHINFER_DUMP_SAFETENSORS` | `0` | `1` writes dumps as `.safetensors` (portable / Hugging Face-style); otherwise plain `.pt`. |\n\n##### Trace Capture\n\nUsed by `flashinfer.trace` / `fi_trace`.\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `FLASHINFER_TRACE_DUMP` | unset | `flashinfer/fi_trace.py` | If set, decorated APIs auto-dump benchmark-definition JSON for each call (the \"fi_trace\" feature). |\n| `FLASHINFER_TRACE_DUMP_DIR` | cwd | `flashinfer/fi_trace.py` | Directory where the trace JSON files are written. |\n| `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. |\n| `FLASHINFER_TRACE_APPLY_PATH` | unset | `flashinfer/trace_apply/config.py` | Directory from which deployment-configured solutions are loaded for Trace Apply. |\n\n##### Validation / Autotuning / Routing / Kernel Selection\n\n| Variable | Default | Read in | Effect |\n|----------|---------|---------|--------|\n| `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. |\n| `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. |\n| `FLASHINFER_AUTOTUNER_LOAD_FROM_FILE` | `0` | `flashinfer/autotuner/autotuner.py` | `1` loads previously serialized autotune results from disk instead of re-running the search. |\n| `FLASHINFER_DIST_AWARE_AUTOTUNE` | `0` | `flashinfer/fuse\n\n<!-- File: .claude-plugin/marketplace.json (Tokens: ~300 | Category: marketplace) -->\n# Claude Plugin Marketplace Catalog\n\nPath: `.claude-plugin/marketplace.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/marketplace.json)\n\n<!-- File: .claude-plugin/plugin.json (Tokens: ~300 | Category: plugin-manifest) -->\n# claude-plugin Plugin Manifest\n\nPath: `.claude-plugin/plugin.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.claude-plugin/plugin.json)\n\n<!-- File: .cursor/mcp.json (Tokens: ~300 | Category: mcp-config) -->\n# Model Context Protocol (MCP) Configuration\n\nPath: `.cursor/mcp.json`\n\n[View Raw Content on GitHub](https://raw.githubusercontent.com/flashinfer-ai/flashinfer/HEAD/.cursor/mcp.json)\n</yakaai_skills>"}