feat(jit): JitSpec ABC + disk cache for JIT-compiled CuTe-DSL kernels

[bkryu]

📌 Description

FlashInfer's nvcc-compiled kernels are cached on disk and load instantly in new processes; CuTe-DSL kernels were memoized only in-process, so every new process recompiled every kernel.

The PR adds the core disk caching functionality and applies it to nvfp4_quantize(backend='cute-dsl') by:

• JitSpec becomes an abstract base class (ABC) with a shared build_and_load() template method (caching, locking, FLASHINFER_DISABLE_JIT); the former dataclass is now JitSpecNvcc (behavior unchanged), and the new
• JitSpecCuteDsl persists CuTe-DSL kernels as export_to_c() object files under cached_ops/, reloaded via JITLink in 3–30 ms. Future DSLs (e.g. cutile) implement the same three methods.
• Wires up nvfp4_quantize as the first user. Example kernel cache path:

  .../cached_ops/
  ├── fp4_quantization_100/                  # nvcc module (unchanged)
  ├── nvfp4_quantize_sm100a_cute_dsl/        # CuTe-DSL module (new)
  │   ├── meta.json                          #   one per module (invalidation)
  │   ├── swizzled_bfloat16_k4096_sf0_pdl0.o #   one .o per specialization (~28 KB)
  │   └── ...
  

• Design doc: docs/design_docs/cute_dsl_kernel_cache.md covers the lifecycle, cache layout, invalidation keys, concurrency/crash safety, and alternatives considered (including the measured-and-rejected K-agnostic compilation). Please start there for review.

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

• New Features
  • Added persistent on-disk caching for CuTe-DSL JIT kernels to reduce repeated compile time, including automatic cache invalidation based on relevant build inputs.
  • Added FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 to bypass the disk cache and force recompilation.
  • Re-exported additional JIT specification classes for direct use.
• Bug Fixes
  • Improved cache correctness with strict metadata validation and safer concurrent, crash-tolerant cache writes/updates.
• Documentation
  • Added design docs describing the CuTe-DSL cache layout, validity rules, and locking behavior.
• Tests
  • Updated the JIT spec construction in the C++ extension test to match the new JIT lifecycle.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds persistent CuTe-DSL kernel caching, refactors the JIT lifecycle into a shared base template, routes NVFP4 CuTe-DSL compilation through the cache wrapper, and updates related documentation and exports.

Changes

CuTe-DSL JIT and kernel caching

Layer / File(s)	Summary
JitSpec lifecycle refactor flashinfer/jit/core.py, flashinfer/jit/__init__.py, tests/test_jit_cpp_ext.py	Introduces the abstract JitSpec template, updates JitSpecNvcc behavior, adjusts registry/status handling, and updates exports and tests.
CuTe-DSL cache core flashinfer/jit/cute_dsl_core.py	Adds CuTe-DSL cache disablement, module naming, metadata hashing, cache lookup, invalidation, crash-safe persistence, and the wrapper used to compile or reuse kernels.
NVFP4 CuTe-DSL integration flashinfer/quantization/kernels/nvfp4_quantize.py	Adds source-based invalidation inputs and routes the NVFP4 CuTe-DSL compile paths through the cache wrapper.
Caching docs CLAUDE.md, docs/design_docs/cute_dsl_kernel_cache.md	Expands the JIT/CuTe-DSL documentation and adds a design page for the cache layout and behavior.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant NVFP4 as nvfp4_quantize compile path
  participant Cache as build_and_load_cute_dsl_kernel
  participant JitSpec as JitSpecCuteDsl
  participant Disk as cached_ops
  participant Compiler as cute.compile

  NVFP4->>Cache: request kernel(name, extra_key_files)
  Cache->>JitSpec: build_and_load()
  JitSpec->>Disk: check .o and meta.json
  alt cache hit
    Disk-->>JitSpec: load exported symbol
  else cache miss or stale
    JitSpec->>Compiler: compile_fn()
    Compiler-->>JitSpec: compiled kernel
    JitSpec->>Disk: write .o then meta.json
  end
  JitSpec-->>NVFP4: kernel handle

Loading

Possibly related PRs

• flashinfer-ai/flashinfer#2339: Also changes flashinfer/jit/core.py and the JitSpec lifecycle.
• flashinfer-ai/flashinfer#2904: Also touches NVFP4 CuTe-DSL kernel compilation and selection in flashinfer/quantization/kernels/nvfp4_quantize.py.
• flashinfer-ai/flashinfer#3205: Also builds on the JIT spec framework in flashinfer/jit/core.py.

Suggested reviewers: yzh119, cyx-6

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title is concise and accurately captures the main change: introducing a JitSpec ABC and CuTe-DSL disk caching.
Description check	✅ Passed	The description matches the template structure and includes the required sections, with a substantive summary and checklist/reviewer notes present.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces an on-disk caching mechanism for JIT-compiled CuTe-DSL kernels, mirroring the two-level caching used for nvcc-compiled modules. It adds a core module for managing the cache and integrates it with the nvfp4 quantization kernels. The review feedback is highly constructive, focusing on making the cache architecture-aware on mixed multi-GPU systems by passing a device parameter, and adding defensive checks to prevent potential FileNotFoundError or AttributeError when hashing source files.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[coderabbitai]

Actionable comments posted: 3

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

docs/design_docs/cute_dsl_kernel_cache.md (1)

123-129: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Broken internal cross-reference: "see 3.2" doesn't exist.

Line 127 references section "3.2" for linking .os into .sos, but "## 3. Alternatives considered" only contains subsection "3.1" (Line 99). This should point to 3.1.

📝 Proposed fix

-- **No `JitSpecRegistry` / AOT integration**: cute-dsl kernels are invisible to `flashinfer aot` and `flashinfer-jit-cache` packaging, and do not honor `FLASHINFER_DISABLE_JIT`. AOT support would prebuild the module directories (or link them into `.so`s, see 3.2).
+- **No `JitSpecRegistry` / AOT integration**: cute-dsl kernels are invisible to `flashinfer aot` and `flashinfer-jit-cache` packaging, and do not honor `FLASHINFER_DISABLE_JIT`. AOT support would prebuild the module directories (or link them into `.so`s, see §3.1).

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/design_docs/cute_dsl_kernel_cache.md` around lines 123 - 129, The
cross-reference in the “No `JitSpecRegistry` / AOT integration” bullet points to
a non-existent section. Update the mention of linking into `.so`s in the
`cute_dsl_kernel_cache.md` design doc so it references the existing “3.1”
subsection instead of “3.2”, keeping the wording aligned with the surrounding
discussion in the limitations section.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/design_docs/cute_dsl_kernel_cache.md`:
- Around line 53-63: The fenced code block in the Cute DSL kernel cache docs is
missing a language hint, which triggers markdownlint MD040. Update the fenced
block in the documented cache layout section to use a language tag such as text
so the block remains readable and lint-compliant.

In `@flashinfer/jit/cute_dsl_core.py`:
- Around line 66-73: The _hash_source_files helper in cute_dsl_core.py should
not be cached solely by the tuple of paths, because that makes source_sha256
stale when file contents change during the process lifetime. Remove the
`@functools.cache` on _hash_source_files or change the invalidation strategy so it
reflects file freshness/content changes, and ensure the caller path through
source_sha256 still produces a hash that updates when any extra_key_files
content changes.
- Around line 161-164: The module invalidation check in the CuTe-DSL module
write path should treat any existing module directory as stale when its metadata
is missing or does not match the expected value. Update the logic around the
meta check in the module directory handling code so that _read_meta(meta_path)
is compared against expected_meta even if meta_path does not exist, and delete
the directory whenever the metadata is absent or mismatched before writing new
metadata. Use the existing module_dir, meta_path, _read_meta, expected_meta, and
logger flow to locate and adjust the invalidation behavior.

---

Outside diff comments:
In `@docs/design_docs/cute_dsl_kernel_cache.md`:
- Around line 123-129: The cross-reference in the “No `JitSpecRegistry` / AOT
integration” bullet points to a non-existent section. Update the mention of
linking into `.so`s in the `cute_dsl_kernel_cache.md` design doc so it
references the existing “3.1” subsection instead of “3.2”, keeping the wording
aligned with the surrounding discussion in the limitations section.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: b1a048dd-0e0c-43ca-94c2-a340472fa50a

📥 Commits

Reviewing files that changed from the base of the PR and between fe9523d and 80accc9.

📒 Files selected for processing (4)

• CLAUDE.md
• docs/design_docs/cute_dsl_kernel_cache.md
• flashinfer/jit/cute_dsl_core.py
• flashinfer/quantization/kernels/nvfp4_quantize.py

[waynehacking8]

Ran a cold/warm check of this branch on SM120 (RTX PRO 6000, CUDA 13.0) with nvfp4_quantize(backend='cute-dsl'), M=640 K=7168 bf16: current main (f2f9646) pays ~0.45s first-call in every fresh process, this branch pays 0.61s once (compile + persist) then 0.11s per fresh process from the cached .o, with outputs bitwise-equal to backend='cuda' in all runs. From the consumer side this is very welcome -- the heavier DSL specializations we lean on for SM120 MoE take far longer than this quantize kernel to compile, and today they pay it on every server start.

[aleozlx]

let's check if the error/exception handling is uniform and document it in the abc

[bkryu]

Good idea. Did two things in the latest commit:

1. Added a try-except block in try_load()
2. Added docstrings to ABC methods: try_load, build, load to state the uniform expected behavior.

[qiching]

some points:

1. This MVP wires only nvfp4_quantize. The autotuning cold start vLLM hit was mm_fp4 (fp4_gemm) compiling ~20–30 cute-dsl tactics sequentially, and vllm #48268 already handles that today by skipping to the CuTe DSL FP4 GEMM Heuristic #2940 heuristic. Where feat(jit): JitSpec ABC + disk cache for JIT-compiled CuTe-DSL kernels #3874 comes in is the next step that wiring mm_fp4 lets compiled kernels survive restarts, which a) removes even the residual single-kernel compile skip_ops still pays per process, and b) lets us re-enable mm_fp4 autotuning without reintroducing cold start. Could you post the rollout order and rough timeline for wiring mm_fp4? (The trtllm MoE cost is profiling-bound / prebuilt, feat(jit): JitSpec ABC + disk cache for JIT-compiled CuTe-DSL kernels #3874 correctly does not apply there.)
2. Selection rides entirely on the _nvfp4_kernel_name string; meta.json guards only arch / DSL-version / source-SHA, not per-kernel params. Since rollout means other authors write these names, could we add a test asserting the name is a function of every codegen arg?
3. Could we document the full preconditions that same flashinfer version + nvcc arch-list (in the path) and same compile arch + nvidia-cutlass-dsl stack (in meta.json).
4. cute-dsl kernels are not in flashinfer-jit-cache yet, so ship prebuilt kernels won't cover cute-dsl , first serve on a fresh image still compiles. Flagging to scope separately.

+1 on @aleozlx comment LGTM.
Please correct me if i understand wrong.

Thank @qiching

1. Right the plan is to incrementally roll out to mm_fp4 and others cute-dsl kernels. Due to the issues you raised, mm_fp4 should be first, followed by others.
2. Agreed and added a test_cute_dsl_cache.py file in latest commit
3. Agreed and documented in cute_dsl_kernel_cache.md section 2.3
4. Right and not all kernels will be shippable in the flashinfer-jit-cache due to some input parameters being compile-time constant but this is a direction to explore in the future. Especially for kernels like mm_fp4 where it takes a long time to compile and is in fact possible to put into flashinfer-jit-cache AFAIK.

@bkryu Cool! thanks! LGTM

[aleozlx]

lgtm