feat(gdn): adopt the CuTe-DSL disk cache for TVM-FFI GDN kernels by elwhyjay · Pull Request #4912 · flashinfer-ai/flashinfer

elwhyjay · 2026-09-03T03:59:49Z

📌 Description

Migrate the GDN CuTe-DSL call sites that already compile with --enable-tvm-ffi to build_and_load_cute_dsl_kernel, so each specialization is exported to the on-disk kernel cache and later processes JITLink it instead of recompiling.

Covered call sites (8 compile sites in 5 files):

File	Sites	Module
`gdn_decode_nontranspose.py`	small/big batch	`gdn_decode_nontranspose`
`gdn_decode_pretranspose.py`	1	`gdn_decode_pretranspose`
`gdn_decode_mtp.py`	inline, warp	`gdn_decode_mtp`
`gdn_decode_bf16_state.py`	wide_vec, wide_vec_t1, mtp_ilp4	`gdn_decode_bf16_state`
`blackwell/gdn_prefill.py`	1	`gdn_blackwell_prefill`

How it works:

Only the cute.compile(...) call is wrapped in a compile_fn closure. Every in-process cache structure (getter dicts, defaults_by_B, aux buffers) is unchanged, and on a disk hit the closure never runs.
Specialization names mirror the in-process cache keys, which fix(gdn): convert fp16 decode operands and fix cache identity #4219 and feat(gdn): add pooled state and state checkpointing and dtype support for feature parity #4436 already audited to cover every codegen parameter.
Where several entry points share one module, the name adds a variant tag. The Blackwell prefill name also adds num_sm, which the compile bakes in as max_active_clusters.
A small shared formatter (gdn_kernels/cute_dsl_cache_naming.py) keeps names symbol-safe and caps their length.

Measurements with fresh processes and an isolated cache dir:

GPU	Selection	Cold (compile + export)	Warm (disk hit)
RTX PRO 6000 (SM120)	bf16_state + fp32 MTP verify subsets (56 tests)	604s	2.3s
RTX PRO 6000 (SM120)	pretranspose + nontranspose basics (7 tests)	2.9s	1.4s
B200 (SM100)	prefill suites (2835 tests, includes GPU test execution)	217s	110s

🔍 Related Issues

#4214 (GDN-P1, PR 4 of the suggested implementation sequence)

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

🧪 Tests

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

Ran on RTX PRO 6000 (SM120) and B200 (SM100):

tests/gdn/test_cute_dsl_kernel_cache.py (new, modeled on tests/jit/test_cute_dsl_cache.py): 68 passed on both GPUs. Naming contracts plus a disk-cache round trip (reload with cute.compile forbidden, bit-identical results).
tests/jit/test_cute_dsl_cache.py: 42 passed
decode parity subsets from tests/gdn/test_decode_delta_rule.py (SM120): passed cold, warm, and with FLASHINFER_CUTE_DSL_DISABLE_CACHE=1
tests/gdn/test_multistream_overlap.py: 4 passed on B200 (1 of them skipped on SM120)
prefill suites on B200: 2835 passed, 864 skipped, 24 failed, identical cold/warm/cache-disabled. The 24 failures (fp8-state test_prefill_kernel_state_dtype cases) fail identically on unmodified main on the same machine, i.e. pre-existing in my environment.

Reviewer Notes

Out of scope, as follow-ups:

WY kernels: off the gdn_decode.py dispatch and not TVM-FFI yet (persisting them needs a calling-convention change first).
delta_rule_dsl kernels: use explicit cute.compile[gpu_arch] targets, which the disk cache's arch keying does not reflect yet (documented limitation).
GDN-H2/H3 (explicit compile targets, PR 2 of the sequence) is unchanged.

Open to adjusting the split if a different scope was intended for GDN-P1.

Summary by CodeRabbit

Performance
- Added persistent caching for GDN GPU kernel builds, reducing compilation overhead for repeated workloads.
- Cached kernels are reused across sessions when the relevant configuration and source remain unchanged.
- Kernel builds automatically refresh when source code or compilation settings change, helping ensure correct results across updates.
- Improved cache handling across multiple GDN decode and prefill kernel variants for more consistent build reuse.

coderabbitai · 2026-09-03T03:59:57Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Team

Run ID: 65759bba-41e3-43a0-8da5-afa2df907999

📥 Commits

Reviewing files that changed from the base of the PR and between 58609a0 and ed98c5a.

📒 Files selected for processing (2)

tests/gdn/test_decode_delta_rule.py
tests/gdn/test_decode_pretranspose_noncontiguous_pool.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

Walkthrough

GDN CuTe DSL decode and Blackwell prefill compilation now uses named, source-invalidated disk caching. A shared naming utility encodes specialization parameters into safe, bounded names. Tests cover naming coverage, uniqueness, cache reuse, and compile-count behavior.

Changes

GDN CuTe DSL kernel caching

Layer / File(s)	Summary
Specialization name generation `flashinfer/gdn_kernels/cute_dsl_cache_naming.py`	Adds formatting for supported cache-key values and bounded names with SHA-256 suffixes for long names.
Decode kernel cache integration `flashinfer/gdn_kernels/gdn_decode_bf16_state.py`, `flashinfer/gdn_kernels/gdn_decode_mtp.py`, `flashinfer/gdn_kernels/gdn_decode_nontranspose.py`, `flashinfer/gdn_kernels/gdn_decode_pretranspose.py`	Routes decode compilation through `build_and_load_cute_dsl_kernel` with variant-specific names and source-file invalidation keys.
Blackwell prefill cache integration `flashinfer/gdn_kernels/blackwell/gdn_prefill.py`	Builds a prefill cache key and uses it with `num_sm` and kernel source files for cached compilation.
Cache naming and reuse validation `tests/gdn/test_cute_dsl_kernel_cache.py`, `tests/gdn/test_decode_delta_rule.py`, `tests/gdn/test_decode_pretranspose_noncontiguous_pool.py`	Tests parameter coverage, name variation, symbol safety, length limits, disk-cache reuse, and compile counts with disk caching disabled.

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

Merge Risk: ⚪ Minimal · up to ed98c

GDN CuTe-DSL kernels now reuse source-invalidated disk-cached specializations across processes. Compile-count tests retain their intended behavior with persistent caching disabled, leaving no current merge-blocking risk.

Sequence Diagram(s)

sequenceDiagram
  participant GDNKernel
  participant build_and_load_cute_dsl_kernel
  participant DiskCache
  participant cute_compile
  GDNKernel->>build_and_load_cute_dsl_kernel: provide module and specialization name
  build_and_load_cute_dsl_kernel->>DiskCache: check source-invalidated artifact
  alt cache miss
    DiskCache->>cute_compile: invoke compilation callback
    cute_compile-->>DiskCache: return compiled kernel
    DiskCache-->>build_and_load_cute_dsl_kernel: store artifact
  else cache hit
    DiskCache-->>build_and_load_cute_dsl_kernel: load artifact
  end
  build_and_load_cute_dsl_kernel-->>GDNKernel: return compiled kernel

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 58.97% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 39 functions across 9 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: adopting the CuTe-DSL disk cache for GDN kernels.
Description check	✅ Passed	The description is complete and relevant. It explains the scope and implementation, lists the related issue, completes the applicable checklist items, documents test coverage and known pre-existing fa…
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.

Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

Create stacked PR
Commit on current branch

🧪 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

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

bkryu · 2026-09-04T23:38:12Z

@flashinfer-bot run

bkryu · 2026-09-04T23:38:19Z

/bot run tests/gdn

flashinfer-bot · 2026-09-04T23:39:06Z

GitLab MR !1422 has been created, and the CI pipeline #66311518 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-09-05T11:07:18Z

[FAILED] Pipeline #66311518 — 6/17 executed test jobs passed

Compared with nightly #66202435.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Other	Notes
B200	❔ Unknown	❔ Unknown	—	Not compared: `tests.gdn.test_decode_pretranspose_noncontiguous_pool` (1 failure; CUDA 12.9) Unknown: `script failed before producing a JUnit report` (1 job; CUDA 13.0)
GB200	❌ New	❌ New	—	New: `tests.gdn.test_decode_pretranspose_noncontiguous_pool` (2 failures; CUDA 12.9, CUDA 13.0)
GB300	❌ New	❌ New	—	New: `tests.gdn.test_decode_pretranspose_noncontiguous_pool` (2 failures; CUDA 12.9, CUDA 13.0)
H100	❌ New	❌ New	—	New: `tests.gdn.test_decode_pretranspose_noncontiguous_pool` (2 failures; CUDA 12.9, CUDA 13.0)
RTX Pro 6000 Blackwell	❌ New	❌ New	—	New: `tests.gdn.test_decode_pretranspose_noncontiguous_pool` (2 failures; CUDA 12.9, CUDA 13.0)
VR200 CU134	—	—	❔ Unknown	Unknown: `script failed before producing a JUnit report` (1 job)

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 6/6 passed

GPU	CUDA 12.9	CUDA 13.0	Other	Notes
B300 (multi-GPU)	✅ Pass	✅ Pass	—	—
GB200 (multi-node)	✅ Pass	✅ Pass	—	—
GB300 (multi-node)	✅ Pass	✅ Pass	—	—

Failure details

New relative to nightly (attribution uncertain)

tests.gdn.test_decode_pretranspose_noncontiguous_pool — 8 failures on GB200 / CUDA 12.9, GB200 / CUDA 13.0, GB300 / CUDA 12.9, GB300 / CUDA 13.0, H100 / CUDA 12.9, H100 / CUDA 13.0, RTX Pro 6000 Blackwell / CUDA 12.9, RTX Pro 6000 Blackwell / CUDA 13.0
- assert 0 == 1

Could not compare

tests.gdn.test_decode_pretranspose_noncontiguous_pool — 1 failure on B200 / CUDA 12.9
- assert 0 == 1

Timeouts, infrastructure, or incomplete jobs

Unknown: script failed before producing a JUnit report — B200 / CUDA 13.0, VR200 CU134
- Jobs: unit_test_vr200_cu134, unit_test_b200: [cu130]

elwhyjay · 2026-09-07T08:59:11Z

Hi @bkryu Thanks for running the pipeline. The test_decode_pretranspose_noncontiguous_pool failures are caused by this PR. That test counts cute.compile calls and asserts exactly one, but with the disk cache a specialization exported earlier in the same job is reloaded from disk, so the observed count is zero. I pushed a fix that sets FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 inside the two compile-count tests (the MTP one had the same hazard), which keeps them asserting in-process reuse as before. In my testing they now pass both with a fresh cache directory and with a populated one.

bkryu · 2026-09-08T18:36:06Z

Thanks @elwhyjay, was going to followup on the failure but seems like you got to it first. I'll retrigger the pipeline

bkryu · 2026-09-08T18:36:16Z

@flashinfer-bot run

bkryu · 2026-09-08T18:36:20Z

/bot run tests/gdn

flashinfer-bot · 2026-09-08T18:37:14Z

GitLab MR !1422 has been updated with latest changes, and the CI pipeline #66822742 is currently running. I'll report back once the pipeline job completes.

…ackend (#5104)  ## 📌 Description  `tests/gdn/test_cute_dsl_kernel_cache.py::test_nontranspose_disk_cache_round_trip` fails deterministically on SM100/SM103 (B200, GB200, GB300) with `AssertionError: expected one exported artifact, got []`. Root cause is a cross-PR dispatch interaction, not a disk-cache bug: the test (added in #4912, merged Sept 8) calls `gated_delta_rule_decode` with the default `backend="auto"` and asserts the CuTe-DSL nontranspose kernel exports one artifact. #4581 (merged Sept 9 — the day the failures began) added the Cake GDN backend, which `auto` now tries *first* on SM100/103; it handles the test's BF16 decode and returns, so the CuTe-DSL kernel under test never compiles and nothing is exported. This also explains the arch fingerprint: the test fails on B200/GB200/GB300 (Cake-eligible) and passes on H100 (SM90, where Cake falls through to the CuTe-DSL kernel). Each PR was green alone; composed they are red. Fix: pin both `gated_delta_rule_decode` calls in the test to `backend="flashinfer"` (the dispatcher's explicit Cake bypass), so the test exercises the backend whose cache it verifies — robust to future `auto`-dispatch changes. ## 🔍 Related Issues  Fixes #5091. Interaction between #4912 and #4581. ## 🚀 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 - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] 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](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## 🔬 Experimental Track  - [ ] This PR is **experimental**: it adds or changes code under `flashinfer/experimental/` and/or an `@flashinfer_experimental_api`. Tracking issue: # - [ ] The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release. - [ ] Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff). - [ ] Tests live in `tests/experimental/` and were validated on the intended hardware; a runnable example is included. - [ ] Nothing is registered in `flashinfer/aot.py`, and no experimental backend is reachable from `backend="auto"` without `FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. (Calling an `@flashinfer_experimental_api` or naming a backend explicitly is itself the opt-in and needs no environment variable.) - [ ] **Test scope declared below.** The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.  ```experimental-tests # One target per line: a directory or a file. (A pytest ::selector is not # supported -- the sharding runner cannot consume one.) Must be under # tests/experimental/ and must exist. Delete these comment lines and add yours, e.g. # # tests/experimental/test_my_backend.py # tests/experimental/my_backend/ # # Declaring the whole tree (tests/experimental/) is allowed but means every # experimental PR pays for every other feature's tests, in every matrix cell. ``` ## Reviewer Notes   ## Summary by CodeRabbit * **Tests** * Updated kernel cache round-trip coverage to explicitly use the FlashInfer backend for both decode operations.

…ackend (#5104)  ## 📌 Description  `tests/gdn/test_cute_dsl_kernel_cache.py::test_nontranspose_disk_cache_round_trip` fails deterministically on SM100/SM103 (B200, GB200, GB300) with `AssertionError: expected one exported artifact, got []`. Root cause is a cross-PR dispatch interaction, not a disk-cache bug: the test (added in #4912, merged Sept 8) calls `gated_delta_rule_decode` with the default `backend="auto"` and asserts the CuTe-DSL nontranspose kernel exports one artifact. #4581 (merged Sept 9 — the day the failures began) added the Cake GDN backend, which `auto` now tries *first* on SM100/103; it handles the test's BF16 decode and returns, so the CuTe-DSL kernel under test never compiles and nothing is exported. This also explains the arch fingerprint: the test fails on B200/GB200/GB300 (Cake-eligible) and passes on H100 (SM90, where Cake falls through to the CuTe-DSL kernel). Each PR was green alone; composed they are red. Fix: pin both `gated_delta_rule_decode` calls in the test to `backend="flashinfer"` (the dispatcher's explicit Cake bypass), so the test exercises the backend whose cache it verifies — robust to future `auto`-dispatch changes. ## 🔍 Related Issues  Fixes #5091. Interaction between #4912 and #4581. ## 🚀 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 - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] 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](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## 🔬 Experimental Track  - [ ] This PR is **experimental**: it adds or changes code under `flashinfer/experimental/` and/or an `@flashinfer_experimental_api`. Tracking issue: # - [ ] The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release. - [ ] Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff). - [ ] Tests live in `tests/experimental/` and were validated on the intended hardware; a runnable example is included. - [ ] Nothing is registered in `flashinfer/aot.py`, and no experimental backend is reachable from `backend="auto"` without `FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. (Calling an `@flashinfer_experimental_api` or naming a backend explicitly is itself the opt-in and needs no environment variable.) - [ ] **Test scope declared below.** The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.  ```experimental-tests # One target per line: a directory or a file. (A pytest ::selector is not # supported -- the sharding runner cannot consume one.) Must be under # tests/experimental/ and must exist. Delete these comment lines and add yours, e.g. # # tests/experimental/test_my_backend.py # tests/experimental/my_backend/ # # Declaring the whole tree (tests/experimental/) is allowed but means every # experimental PR pays for every other feature's tests, in every matrix cell. ``` ## Reviewer Notes   ## Summary by CodeRabbit * **Tests** * Updated kernel cache round-trip coverage to explicitly use the FlashInfer backend for both decode operations.  (cherry picked from commit 82090ee)

flashinfer-bot added the op: linear attention KDA, mamba, GDN, etc. review filtering. label Sep 3, 2026

elwhyjay mentioned this pull request Sep 3, 2026

GDN CuTe-DSL cache architecture improvements #4214

Open

feat(gdn): adopt the CuTe-DSL disk cache for TVM-FFI GDN kernels

58609a0

elwhyjay force-pushed the feat/gdn-p1 branch from c575ed9 to 58609a0 Compare September 3, 2026 05:57

elwhyjay marked this pull request as ready for review September 3, 2026 06:03

elwhyjay requested review from bkryu, jiahanc, kahyunnam, yongwww and yzh119 as code owners September 3, 2026 06:03

flashinfer-bot added the run-ci label Sep 4, 2026

test(gdn): pin the disk cache off in compile-count tests

ed98c5a

flashinfer-bot added run-ci and removed run-ci labels Sep 8, 2026

bkryu approved these changes Sep 8, 2026

View reviewed changes

bkryu merged commit 6b3d8b4 into flashinfer-ai:main Sep 8, 2026
25 of 26 checks passed

bkryu mentioned this pull request Sep 10, 2026

fix(test): pin the GDN disk-cache round-trip test to the flashinfer backend #5104

Merged

11 tasks

XFDG mentioned this pull request Sep 12, 2026

GDN prefill on SM100 is nondeterministic: test_prefill_kernel_state_dtype flips between all-pass and all-fail across identical runs #4704

Open

1 task

bkryu mentioned this pull request Sep 14, 2026

Add NVFP4 paged-KV MSA sparse decode and prefill for SM100/SM103 #4982

Open

11 tasks

kahyunnam mentioned this pull request Sep 15, 2026

[Compilation] GDN decode cold kernel compilation time is too long on SM90/100/120/121 #4110

Open

Uh oh!

Conversation

elwhyjay commented Sep 3, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

bkryu commented Sep 4, 2026

Uh oh!

bkryu commented Sep 4, 2026

Uh oh!

flashinfer-bot commented Sep 4, 2026

Uh oh!

flashinfer-bot commented Sep 5, 2026

Unit Tests

Multi-GPU and Multi-Node Tests — 6/6 passed

New relative to nightly (attribution uncertain)

Could not compare

Timeouts, infrastructure, or incomplete jobs

Uh oh!

elwhyjay commented Sep 7, 2026

Uh oh!

bkryu commented Sep 8, 2026

Uh oh!

bkryu commented Sep 8, 2026

Uh oh!

bkryu commented Sep 8, 2026

Uh oh!

flashinfer-bot commented Sep 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

elwhyjay commented Sep 3, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Sep 3, 2026 •

edited

Loading