perf(gdn): reuse pretranspose kernels across pool capacity and stride by leonardHONG · Pull Request #4444 · flashinfer-ai/flashinfer

leonardHONG · 2026-08-10T14:26:55Z

📌 Description

Make the leading pool dimension of the GDN pretranspose path dynamic:

Treat pool capacity and stride[0] as runtime values.
Keep HV/V/K and inner strides compile-time static.
Remove pool capacity and leading stride from the compilation cache key.
Preserve 64-bit pool-offset arithmetic.

This allows one compiled kernel to be reused across non-contiguous state pools with different capacities and leading strides.

🔍 Related Issues

Partially addresses #4110.

🚀 Pull Request Checklist

✅ Pre-commit Checks

I have installed pre-commit.
I have installed the hooks.
I have run pre-commit run --all-files.

🧪 Tests

Tests have been added or updated as needed.
All relevant tests are passing.

Validated on H800/SM90a and RTX 5090/SM120, including non-contiguous pools, compile reuse, and large 64-bit offsets. No steady-state regression was observed; cold compilation time was reduced by approximately 62%.

Reviewer Notes

Only the outer pool mode is dynamic. Inner dimensions and strides remain static for kernel specialization.

Summary by CodeRabbit

Performance
- Improved pretranspose kernel reuse across different pool capacities and outer memory layouts, reducing unnecessary recompilation.
- Preserved separate handling for layouts with differing inner strides.
Reliability
- Added validation to reject pool slot strides that are not properly aligned for efficient data transfer.
Tests
- Expanded coverage for cache reuse, layout variations, alignment validation, and environments where the optional kernel is unavailable.

Signed-off-by: 梁厚宏 <2695316095@qq.com>

coderabbitai · 2026-08-10T14:27:35Z

📝 Walkthrough

Walkthrough

The pretranspose kernel now treats pool size and outer stride as runtime-symbolic inputs while keeping inner strides compilation-static. Tests cover cache reuse, distinct inner-stride kernels, output and state correctness, optional module loading, and misaligned stride rejection.

Changes

Pool pretranspose compilation

Layer / File(s)	Summary
Symbolic pool inputs and stride validation `flashinfer/gdn_kernels/gdn_decode_pretranspose.py`	Pool-indexed state tensors use symbolic pool size and stride-0 values. Runtime inner strides remain available. Pool slot stride-0 must be divisible by four FP32 elements.
Compilation-cache regression coverage `tests/gdn/test_decode_pretranspose_noncontiguous_pool.py`	Tests verify reuse across pool capacities and outer strides, separate compilation for different inner strides, correct outputs and state updates, optional module loading, and rejection of misaligned strides.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: bkryu, kahyunnam

🚥 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 clearly and concisely describes kernel reuse across pool capacity and stride changes, which is the main pull request objective.
Description check	✅ Passed	The description includes the required sections, explains the changes, references an issue, reports checklist completion, and summarizes test coverage and results.

✨ 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.}

coderabbitai

🧹 Nitpick comments (1)

flashinfer/gdn_kernels/gdn_decode_pretranspose.py (1)
1002-1013: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document the outer-layout specialization alternative.

The comments explain why pool capacity and stride0 are dynamic. They do not document the alternative. State that specializing these values would create one compiled callable per outer layout without changing the local tile logic. State why stride1 through stride3 remain static.
Proposed comment update
-            # Pool capacity and the distance between slots do not affect codegen.
-            # Keep the inner state layout static while accepting arbitrary pool
-            # sizes and padded slot strides through the same compiled callable.
+            # Keep pool capacity and stride0 dynamic. Specializing either value
+            # would create one compiled callable per outer layout without changing
+            # local tile logic. Keep inner strides static because they define
+            # in-slot addressing used by the compiled kernel.
As per coding guidelines, “For performance-critical hot paths, document the rationale for special algorithmic choices and relevant alternatives in comments.”
🤖 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 `@flashinfer/gdn_kernels/gdn_decode_pretranspose.py` around lines 1002 - 1013,
Update the comments in the use_pool_indexing branch near h0_source_tensor to
document that specializing pool capacity and stride0 would produce one compiled
callable per outer layout without changing local tile logic, while stride1
through stride3 remain static because they define the inner tile layout.
Source: Coding guidelines

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

Nitpick comments:
In `@flashinfer/gdn_kernels/gdn_decode_pretranspose.py`:
- Around line 1002-1013: Update the comments in the use_pool_indexing branch
near h0_source_tensor to document that specializing pool capacity and stride0
would produce one compiled callable per outer layout without changing local tile
logic, while stride1 through stride3 remain static because they define the inner
tile layout.

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 75ac65f5-279d-4134-b6aa-e36a04cd5a88

📥 Commits

Reviewing files that changed from the base of the PR and between 2ab910c and 91702c3.

📒 Files selected for processing (2)

flashinfer/gdn_kernels/gdn_decode_pretranspose.py
tests/gdn/test_decode_pretranspose_noncontiguous_pool.py

kahyunnam · 2026-08-11T22:30:33Z

/bot run tests/gdn

flashinfer-bot · 2026-08-11T22:32:41Z

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

kahyunnam

Thanks for the contribution! LGTM, I will help merge pending test results.

flashinfer-bot · 2026-08-12T09:37:45Z

[SUCCESS] Pipeline #62224473: 18/18 executed test jobs passed

…ions (#4513) ## 📌 Description `tests/gdn/test_decode_delta_rule.py` re-ran the same compiled cubins many times over. Two independent reasons, both verified against the cache keys rather than assumed: **1. Batch size is not a compile key.** It has been dynamic since #3649, so it reaches the cache only through coarse buckets: | Path | How `B` enters the key | |------|------------------------| | pretranspose (`gdn_decode_pretranspose.py:964`) | not at all | | bf16-state wide-vec (`gdn_decode_bf16_state.py:3440`) | not at all — the tests build a contiguous pool, so `pool_size_key = -1` and `pool_slot_stride = (-1,)` are B-independent sentinels | | nontranspose (`gdn_decode_nontranspose.py:725`) | only via `use_small_batch = B < 32` | | fp32 / bf16 MTP (`gdn_decode_mtp.py:2500`, `gdn_decode_bf16_state.py:3813`) | only via `get_mtp_config` / `_get_bf16_mtp_config` | The clearest case was `test_gdn_decode_bf16_state_wide_vec_mtp_kernel`: 378 of the file's 817 parametrized cases (46%) but only **42** compile keys, because `tile_v` is an explicit monkeypatched axis and the 9 batch sizes contribute nothing. The first commit keeps one batch size per bucket. I verified each kept set reproduces the *full* key set, at `NUM_SMS` 108/132/148 — this caught a real mistake, where `[1,8,16,32]` for `test_gdn_decode_bf16_state_t1_kernel` silently dropped the `HV=64, tile_v=32` key (that test sweeps `HV` ∈ {32,64}). **2. Intermediate `seq_len` values only re-specialize on `T`.** `get_mtp_config` returns an identical `(tile_v, vec_size, ilp_rows, use_smem_v)` set for every `T >= 3`, so T=3/5/6/7 compile fresh MTP cubins without covering a tile config that T=4 or T=8 does not already cover. `T=2` is kept as the one structurally distinct case — it alone reaches the `ilp=8` and `tile_v=16 / ilp=2` branches. ### Effect | | cases | compile keys (retuned tests) | |---|---|---| | before | 817 | 102 | | after batch collapse | 501 | 102 | | after `seq_len` trim | 409 | 62 | Collected tests go 838 → 416 (409 parametrized + 7 non-parametrized). ### Coverage cost The first commit costs nothing in specialization coverage — the same cubins still run, just at fewer runtime batch sizes. The second commit is a deliberate reduction: T=3/5/6/7 still exercise distinct unrolled loop counts, so a T-specific off-by-one would no longer be caught. It is a separate commit so it can be dropped if reviewers would rather keep the full sweep. ## 🔍 Related Issues Refs #4110 (GDN cold-compile CI time). Complements #4128 and #4444, which remove key entries that provably do not reach codegen; this removes test cases that map onto keys already covered. ## 🚀 Pull Request Checklist ### ✅ 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. ## 🧪 Tests - [x] `pytest tests/gdn/test_decode_delta_rule.py -q` on H100: **416 passed in 32m25s** - [ ] GPU CI for the timing comparison ## Reviewer Notes - I did not measure a clean before/after wall clock: the ~38 min baseline I was working from comes from #4219's description rather than the same machine, so I'd rather let CI provide the comparison than quote a number I can't stand behind. The case and key counts above are exact and static. - Worth noting for #4110 more broadly: cutting 50% of the cases bought substantially less than 50% of the wall clock, which suggests the remaining cost is dominated by compilation and fixed overhead rather than per-case execution. That points at persistent/AOT CuTe-DSL artifacts (GDN-P1 in #4214) as the larger lever. - The batch sizes kept per test are load-bearing, not arbitrary — each set is one representative per config bucket. I left a one-line comment at each site so they don't get "restored" later.  ## Summary by CodeRabbit * **Tests** * Streamlined decode and MTP test coverage to use representative batch sizes and sequence lengths. * Preserved coverage for key thresholds, tile configurations, transposition modes, precision variants, and sequence-length scenarios.  --------- Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>

…flashinfer-ai#4444) ## 📌 Description Make the leading pool dimension of the GDN pretranspose path dynamic: - Treat pool capacity and `stride[0]` as runtime values. - Keep `HV/V/K` and inner strides compile-time static. - Remove pool capacity and leading stride from the compilation cache key. - Preserve 64-bit pool-offset arithmetic. This allows one compiled kernel to be reused across non-contiguous state pools with different capacities and leading strides. ## 🔍 Related Issues Partially addresses flashinfer-ai#4110. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit`. - [x] I have installed the hooks. - [x] I have run `pre-commit run --all-files`. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All relevant tests are passing. Validated on H800/SM90a and RTX 5090/SM120, including non-contiguous pools, compile reuse, and large 64-bit offsets. No steady-state regression was observed; cold compilation time was reduced by approximately 62%. ## Reviewer Notes Only the outer pool mode is dynamic. Inner dimensions and strides remain static for kernel specialization.  ## Summary by CodeRabbit * **Performance** * Improved pretranspose kernel reuse across different pool capacities and outer memory layouts, reducing unnecessary recompilation. * Preserved separate handling for layouts with differing inner strides. * **Reliability** * Added validation to reject pool slot strides that are not properly aligned for efficient data transfer. * **Tests** * Expanded coverage for cache reuse, layout variations, alignment validation, and environments where the optional kernel is unavailable.  Signed-off-by: 梁厚宏 <2695316095@qq.com>

…ions (flashinfer-ai#4513) ## 📌 Description `tests/gdn/test_decode_delta_rule.py` re-ran the same compiled cubins many times over. Two independent reasons, both verified against the cache keys rather than assumed: **1. Batch size is not a compile key.** It has been dynamic since flashinfer-ai#3649, so it reaches the cache only through coarse buckets: | Path | How `B` enters the key | |------|------------------------| | pretranspose (`gdn_decode_pretranspose.py:964`) | not at all | | bf16-state wide-vec (`gdn_decode_bf16_state.py:3440`) | not at all — the tests build a contiguous pool, so `pool_size_key = -1` and `pool_slot_stride = (-1,)` are B-independent sentinels | | nontranspose (`gdn_decode_nontranspose.py:725`) | only via `use_small_batch = B < 32` | | fp32 / bf16 MTP (`gdn_decode_mtp.py:2500`, `gdn_decode_bf16_state.py:3813`) | only via `get_mtp_config` / `_get_bf16_mtp_config` | The clearest case was `test_gdn_decode_bf16_state_wide_vec_mtp_kernel`: 378 of the file's 817 parametrized cases (46%) but only **42** compile keys, because `tile_v` is an explicit monkeypatched axis and the 9 batch sizes contribute nothing. The first commit keeps one batch size per bucket. I verified each kept set reproduces the *full* key set, at `NUM_SMS` 108/132/148 — this caught a real mistake, where `[1,8,16,32]` for `test_gdn_decode_bf16_state_t1_kernel` silently dropped the `HV=64, tile_v=32` key (that test sweeps `HV` ∈ {32,64}). **2. Intermediate `seq_len` values only re-specialize on `T`.** `get_mtp_config` returns an identical `(tile_v, vec_size, ilp_rows, use_smem_v)` set for every `T >= 3`, so T=3/5/6/7 compile fresh MTP cubins without covering a tile config that T=4 or T=8 does not already cover. `T=2` is kept as the one structurally distinct case — it alone reaches the `ilp=8` and `tile_v=16 / ilp=2` branches. ### Effect | | cases | compile keys (retuned tests) | |---|---|---| | before | 817 | 102 | | after batch collapse | 501 | 102 | | after `seq_len` trim | 409 | 62 | Collected tests go 838 → 416 (409 parametrized + 7 non-parametrized). ### Coverage cost The first commit costs nothing in specialization coverage — the same cubins still run, just at fewer runtime batch sizes. The second commit is a deliberate reduction: T=3/5/6/7 still exercise distinct unrolled loop counts, so a T-specific off-by-one would no longer be caught. It is a separate commit so it can be dropped if reviewers would rather keep the full sweep. ## 🔍 Related Issues Refs flashinfer-ai#4110 (GDN cold-compile CI time). Complements flashinfer-ai#4128 and flashinfer-ai#4444, which remove key entries that provably do not reach codegen; this removes test cases that map onto keys already covered. ## 🚀 Pull Request Checklist ### ✅ 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. ## 🧪 Tests - [x] `pytest tests/gdn/test_decode_delta_rule.py -q` on H100: **416 passed in 32m25s** - [ ] GPU CI for the timing comparison ## Reviewer Notes - I did not measure a clean before/after wall clock: the ~38 min baseline I was working from comes from flashinfer-ai#4219's description rather than the same machine, so I'd rather let CI provide the comparison than quote a number I can't stand behind. The case and key counts above are exact and static. - Worth noting for flashinfer-ai#4110 more broadly: cutting 50% of the cases bought substantially less than 50% of the wall clock, which suggests the remaining cost is dominated by compilation and fixed overhead rather than per-case execution. That points at persistent/AOT CuTe-DSL artifacts (GDN-P1 in flashinfer-ai#4214) as the larger lever. - The batch sizes kept per test are load-bearing, not arbitrary — each set is one representative per config bucket. I left a one-line comment at each site so they don't get "restored" later.  ## Summary by CodeRabbit * **Tests** * Streamlined decode and MTP test coverage to use representative batch sizes and sequence lengths. * Preserved coverage for key thresholds, tile configurations, transposition modes, precision variants, and sequence-length scenarios.  --------- Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>

perf(gdn): reuse pretranspose kernels across pool capacity and stride

91702c3

Signed-off-by: 梁厚宏 <2695316095@qq.com>

leonardHONG requested review from bkryu, jiahanc, kahyunnam, yongwww and yzh119 as code owners August 10, 2026 14:26

coderabbitai Bot reviewed Aug 10, 2026

View reviewed changes

kahyunnam added op: linear attention KDA, mamba, GDN, etc. review filtering. run-ci labels Aug 11, 2026

kahyunnam assigned kahyunnam and leonardHONG Aug 11, 2026

coderabbitai Bot mentioned this pull request Aug 11, 2026

fix(gdn): convert fp16 decode operands and fix cache identity #4219

Merged

4 tasks

kahyunnam approved these changes Aug 11, 2026

View reviewed changes

kahyunnam merged commit ba6bf4e into flashinfer-ai:main Aug 13, 2026
40 of 42 checks passed

This was referenced Aug 13, 2026

perf(gdn): reuse MTP decode kernels across cache modes #4128

Merged

perf(gdn): reuse FP32 MTP kernels across pool capacity and stride #4512

Open

kahyunnam mentioned this pull request Aug 13, 2026

test(gdn): trim GDN decode test matrix to distinct kernel specializations #4513

Merged

5 tasks

coderabbitai Bot mentioned this pull request Aug 17, 2026

feat(gdn): add pooled state and state checkpointing and dtype support for feature parity #4436

Merged

Navigation Menu

Uh oh!

perf(gdn): reuse pretranspose kernels across pool capacity and stride - #4444

perf(gdn): reuse pretranspose kernels across pool capacity and stride#4444
kahyunnam merged 1 commit into
flashinfer-ai:mainfrom
leonardHONG:feat/gdn-pretranspose-dynamic-pool

leonardHONG commented Aug 10, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Aug 10, 2026 •

edited

Loading

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Uh oh!

kahyunnam commented Aug 11, 2026

Uh oh!

flashinfer-bot commented Aug 11, 2026

Uh oh!

kahyunnam left a comment

Uh oh!

flashinfer-bot commented Aug 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

leonardHONG commented Aug 10, 2026 • edited 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 Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

kahyunnam commented Aug 11, 2026

Uh oh!

flashinfer-bot commented Aug 11, 2026

Uh oh!

kahyunnam left a comment

Choose a reason for hiding this comment

Uh oh!

flashinfer-bot commented Aug 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

leonardHONG commented Aug 10, 2026 •

edited

Loading

coderabbitai Bot commented Aug 10, 2026 •

edited

Loading