feat(bench): add GDN routines to flashinfer_benchmark.py - #3572
Conversation
📝 WalkthroughWalkthroughAdds GDN (Gated Delta Rule) benchmarking: Triton reference kernels, CLI wiring, a new GDN routine module with FLOPs/bytes estimators and tests, benchmark script updates (caching/accounting and optional FLA), docs, and sample invocations. ChangesGDN Benchmarking Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds Gated Delta Net (GDN) linear attention benchmarking support (decode, MTP, and prefill) to the FlashInfer benchmark suite, extracting Triton reference kernels into a shared module and updating the benchmark infrastructure. Feedback on the changes highlights two critical issues: first, importing Triton reference functions will raise an ImportError on systems without Triton installed, breaking non-Triton benchmarks; second, BF16 state decode benchmarking will fail at runtime when pool_split is False because it incorrectly falls through to a non-pool path instead of routing to the pool-only path required by BF16 state kernels.
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.
| BV=BV, | ||
| ) | ||
|
|
||
| return output, initial_state |
There was a problem hiding this comment.
If triton is not installed, importing triton_gdn_decode, triton_gdn_decode_pretranspose, or triton_gdn_mtp from gdn_triton_reference will raise an ImportError because these names are only defined inside the if TRITON_AVAILABLE: block. Since other scripts (like bench_gdn_decode.py and routines/gdn.py) import these functions at the top level, this completely breaks benchmarking on systems without Triton, even when only benchmarking the flashinfer backend.
Please add an else block at the end of the file to define these functions as None when Triton is not available.
| return output, initial_state | |
| return output, initial_state | |
| else: | |
| triton_gdn_decode = None | |
| triton_gdn_decode_pretranspose = None | |
| triton_gdn_mtp = None |
There was a problem hiding this comment.
Agreed. Adding the else-block None fallbacks in gdn_triton_reference.py in next commit
| return gated_delta_rule_decode( | ||
| q, k, v, state, A_log, a, dt_bias, b, scale, output, use_qk_l2norm | ||
| )[0] | ||
| if pool_split: |
There was a problem hiding this comment.
The BF16 state decode/MTP kernels are pool-only and do not support the standard in-place non-pool path (where state is passed as the 4th positional argument).
When is_mtp is False (decode T=1), use_bf16_state is True, and pool_split is False (i.e., pool_mode is single), the code currently falls through to the non-pool path on line 623. This will cause a runtime failure when benchmarking the BF16 state decode kernel.
To fix this, we should route the call to the pool-only path whenever use_bf16_state is True, even if pool_split is False.
| if pool_split: | |
| if pool_split or use_bf16_state: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
benchmarks/README.md (1)
80-83: 💤 Low valueOptional: Clarify SM90+ support scope.
The description states "SM90+" which correctly indicates the minimum compute capability, but the support matrix (lines 548-550) shows FlashInfer backend support is limited to SM 9.0–10.3, with only Triton reference available for SM 11.0+ (and no prefill support for SM 11.0+). Consider adding a brief note like:
-- GDN (Gated Delta Net linear attention, SM90+): +- GDN (Gated Delta Net linear attention, SM90+ minimum; FlashInfer kernels available on SM90–SM10.3):This would set clearer expectations for users on newer architectures (SM 11.0–12.1) who will only have Triton reference kernels for decode/MTP.
🤖 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 `@benchmarks/README.md` around lines 80 - 83, Add a brief clarification in the GDN section explaining the SM90+ support scope: state that while the header uses "SM90+" as the minimum compute capability, FlashInfer backend support is limited to SM 9.0–10.3 (including prefill/preftransposed kernels), whereas for SM 11.0+ only the Triton reference kernels are available (no FlashInfer prefill support for SM 11.0+), and note that BF16 state kernels or pretranspose variants (referenced by gated_delta_rule_decode, gated_delta_rule_decode_pretranspose, gated_delta_rule_mtp, and chunk_gated_delta_rule) may therefore be backend-limited (flashinfer vs triton vs fla) on newer architectures to set clear expectations for users.
🤖 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 `@benchmarks/README.md`:
- Around line 80-83: Add a brief clarification in the GDN section explaining the
SM90+ support scope: state that while the header uses "SM90+" as the minimum
compute capability, FlashInfer backend support is limited to SM 9.0–10.3
(including prefill/preftransposed kernels), whereas for SM 11.0+ only the Triton
reference kernels are available (no FlashInfer prefill support for SM 11.0+),
and note that BF16 state kernels or pretranspose variants (referenced by
gated_delta_rule_decode, gated_delta_rule_decode_pretranspose,
gated_delta_rule_mtp, and chunk_gated_delta_rule) may therefore be
backend-limited (flashinfer vs triton vs fla) on newer architectures to set
clear expectations for users.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: cbc113f9-07de-452b-bd34-77750657e7be
📒 Files selected for processing (8)
benchmarks/README.mdbenchmarks/bench_gdn_decode.pybenchmarks/bench_gdn_prefill.pybenchmarks/flashinfer_benchmark.pybenchmarks/gdn_triton_reference.pybenchmarks/routines/flashinfer_benchmark_utils.pybenchmarks/routines/gdn.pybenchmarks/samples/sample_testlist.txt
2 participants
📌 Description
Adds Gated Delta Net (GDN) support to
flashinfer_benchmark.pyas a newgdnroutine category:Added Routines:
gated_delta_rule_decode— T=1 decode;--state_layout {pretranspose,nontranspose},--state_dtype {float32,bfloat16}(bf16 state kernels),--pool_mode {single,split}gated_delta_rule_mtp— T>=2 multi-token processing with state pool + indiceschunk_gated_delta_rule— varlen chunked prefillBug fixes in the standalone benches
bench_gdn_decode.py: all-layouts mode called the BF16-state kernel without pool indices (column was always N/A); bytes model counted intermediate-state traffic for T>1 even with caching disabled, inflating reported TB/sbench_gdn_prefill.py: fed log-spacegto the FlashInfer kernel, which takes a linear-space alpha — outputs/state were NaN; FLA is now optional instead of a hard exit🔍 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
Bug Fixes
Documentation