perf(gdn): make GDN kernels compilation batch-size agnostic (support dynamic batch shapes for vLLM integration) - #3649
Merged
Conversation
The CuTe DSL BF16 GDN decode kernels (wide_vec and MTP ILP4) baked the batch size B into the compiled cubin via a Constexpr kernel arg and the compile cache key, forcing a fresh cute.compile() for every new running batch size at inference time (reported during vLLM integration). B is never used inside the kernels (i_n is decoded from block_idx), so demote it to a runtime grid extent (cute.size(q.shape[0])), mark the batch dim of the per-batch tensors dynamic before cute.compile, and drop B from the cache keys. One cubin per (tile_v, T, HV, ...) now serves all batch sizes.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
💤 Files with no reviewable changes (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThis PR modifies all four GDN decode kernel paths (BF16 state, MTP, nontranspose, pretranspose) to treat the batch dimension ChangesBF16 State: Dynamic-B Decode Kernel Compilation
MTP Decode: Dynamic-B Kernel and Caching
Nontranspose Decode: Dynamic-B Kernel Compilation
Pretranspose Decode: Dynamic-B Kernel and Compilation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
kahyunnam
marked this pull request as ready for review
June 16, 2026 00:05
kahyunnam
requested review from
bkryu,
jiahanc,
yongwww and
yzh119
as code owners
June 16, 2026 00:05
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request implements dynamic batching for the GDN decode BF16 kernels by removing the batch dimension B from the compile-time constants and resolving it dynamically at runtime. This optimization prevents kernel recompilation for different batch sizes during inference. No review comments were provided, and the implementation is clean and well-structured.
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.
Member
Author
|
/bot run tests/gdn |
Collaborator
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
flashinfer/gdn_kernels/gdn_decode_bf16_state.py (2)
54-60: 💤 Low valueConsider adding a docstring for clarity.
This new helper encapsulates non-trivial TVM-FFI interop logic. A brief docstring would aid future maintainability.
📝 Suggested docstring
def _mark_batch_dynamic(torch_t: torch.Tensor, *, assumed_align: int = 32): + """Mark a tensor's batch dimension (mode 0) as dynamically compactable. + + Required for batch-dependent tensors passed to cute.compile() when the + compiled cubin should be reusable across different batch sizes. + """ # stride_order is explicit because auto-deduction is ambiguous when the # batch dim is size 1 (B=1, T=1 -> several dims share equal strides). stride_order = tuple(sorted(range(torch_t.dim()), key=lambda d: -torch_t.stride(d)))🤖 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_bf16_state.py` around lines 54 - 60, The function _mark_batch_dynamic lacks a docstring despite implementing non-trivial TVM-FFI interop logic. Add a docstring to this function that documents its purpose (marking the batch dimension as dynamic in TVM), explains the parameters (torch_t as the input tensor, assumed_align as alignment assumption, and the stride_order computation logic), and describes the return value from the mark_compact_shape_dynamic call.
1767-1770: 💤 Low valueComment placement is misleading.
The comment "Dummy [1,1,1] tensor (caching off) has no unique stride-1 dim" is placed inside the
if cache_intermediate_states:branch (caching ON), but it explains why we skip marking in the opposite case. This could confuse future readers.✏️ Suggested clarification
inter_ = from_dlpack(intermediate_states, assumed_align=32, enable_tvm_ffi=True) if cache_intermediate_states: - # Dummy [1,1,1] tensor (caching off) has no unique stride-1 dim. + # Mark batch-shaped intermediate_states as dynamic. When caching + # is off, intermediate_states is a dummy [1,1,1] tensor that + # cannot be marked (no unique stride-1 dim), so skip in that case. inter_ = _mark_batch_dynamic(intermediate_states)🤖 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_bf16_state.py` around lines 1767 - 1770, The comment referencing the dummy tensor and caching off state is placed inside the if cache_intermediate_states branch, but the comment describes the opposite case (caching off). Move this comment to the appropriate location that actually corresponds to the case it describes. Either restructure the comment to clarify what happens when caching is enabled in the current if branch, or relocate it to where the caching off case is handled to avoid confusing future readers about when this constraint applies.
🤖 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_bf16_state.py`:
- Around line 54-60: The function _mark_batch_dynamic lacks a docstring despite
implementing non-trivial TVM-FFI interop logic. Add a docstring to this function
that documents its purpose (marking the batch dimension as dynamic in TVM),
explains the parameters (torch_t as the input tensor, assumed_align as alignment
assumption, and the stride_order computation logic), and describes the return
value from the mark_compact_shape_dynamic call.
- Around line 1767-1770: The comment referencing the dummy tensor and caching
off state is placed inside the if cache_intermediate_states branch, but the
comment describes the opposite case (caching off). Move this comment to the
appropriate location that actually corresponds to the case it describes. Either
restructure the comment to clarify what happens when caching is enabled in the
current if branch, or relocate it to where the caching off case is handled to
avoid confusing future readers about when this constraint applies.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a17cd87e-0d32-4222-82a0-e288bdba0d87
📒 Files selected for processing (1)
flashinfer/gdn_kernels/gdn_decode_bf16_state.py
Collaborator
|
[FAILED] Pipeline #54876489: 10/20 passed |
…ze agnostic Extends the BF16-state batch-agnostic fix to the remaining GDN decode kernels (float32 pretranspose, nontranspose, and MTP), so a new running batch size no longer triggers a fresh ~45s cute.compile at inference time. B was a Constexpr in every compile key but unused in all kernel bodies (i_n is decoded from block_idx). B is now a runtime grid extent (q.shape[0]); per-batch tensors are marked dynamic before cute.compile; B (and pool_size for MTP) are dropped from the cache keys; batch-sized aux tensors (h0_indices / cu_seqlens) move to per-(B, device) maps. nontranspose keys on use_small_batch instead of B. Padded/strided pools keep baked strides; T remains constexpr. Mirrors the GDN-kernel portion of Brian Ryu's flashinfer-ai#3601 so that PR can drop its kernel changes and stay test-only.
kahyunnam
changed the title
Member
Author
|
/bot run tests/gdn |
Collaborator
Aligns the BF16 decode pool handling with the other GDN kernels and with Brian Ryu's flashinfer-ai#3601: a contiguous state pool now marks its slot dimension dynamic (with sentinel -1 pool_size/stride cache keys) so a single cubin serves all pool sizes, while padded/strided pools (vLLM's packed conv+ssm page) keep their real pool_size/stride baked in via the key. Previously B was dropped from the key but pool_size/stride were always retained, so callers with pool_size == batch_size (e.g. the GDN unit tests) still recompiled per batch size. Now compilation is bounded by the tile_v/kernel variant only.
Only mark intermediate_states' leading dim dynamic when caching is on; the caching-off dummy ([1,1,1]) is never read by the kernel and need not be marked. Matches the BF16 path and avoids relying on the DSL tolerating a degenerate all-strides-equal tensor on the common (caching-off) path. Also corrects a misleading pool-scoping comment. (Review follow-up.)
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
flashinfer/gdn_kernels/gdn_decode_bf16_state.py (2)
1700-1710:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftUse
cache_stepsfor intermediate-state batch stride.Both paths accept
cache_steps >= T, but the kernels flatten writes withi_n * T * HV; whencache_steps > T, batches after 0 write into the wrong buffer slots. Either threadcache_stepsthrough the launcher/kernel and usei_n * cache_steps * HV, or tighten these assertions to requirecache_steps == T.Also applies to: 1931-1943
🤖 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_bf16_state.py` around lines 1700 - 1710, The kernel buffer stride calculation uses T in the flattening logic (i_n * T * HV) but the intermediate_states_buffer is reshaped using cache_steps (B_val * cache_steps * HV_val), causing batch writes to target incorrect buffer slots when cache_steps > T. Fix this by either: (1) threading the cache_steps parameter through the kernel launchers and updating the stride calculation to use i_n * cache_steps * HV instead of i_n * T * HV, or (2) tightening the assertions at the current location and the sibling locations (lines 1931-1943) to require cache_steps == T instead of allowing cache_steps >= T. Choose one approach and apply it consistently across all affected sites to ensure stride calculations match the buffer layout.
1789-1790:⚠️ Potential issue | 🟠 MajorCompile split-pool write indices from the write-index tensor.
h0_out_idx_is always marked frominitial_state_indices, even whensame_pool=Falseand the runtime call passesoutput_state_indices. This causes the kernel to be compiled with a placeholder for the write index tensor using the read index's layout, then launched with a potentially different write index tensor, resulting in a layout mismatch. Build the placeholder fromoutput_state_indicesfor split-pool calls.Proposed localized fix
- h0_idx_ = _mark_batch_dynamic(initial_state_indices) - h0_out_idx_ = _mark_batch_dynamic(initial_state_indices) + h0_idx_ = _mark_batch_dynamic(initial_state_indices) + h0_out_idx_source = initial_state_indices if same_pool else output_state_indices + h0_out_idx_ = _mark_batch_dynamic(h0_out_idx_source)Apply the same change in both compile blocks.
🤖 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_bf16_state.py` around lines 1789 - 1790, The `h0_out_idx_` variable is currently always initialized using `initial_state_indices`, but when `same_pool=False`, the kernel is launched with `output_state_indices` instead, causing a layout mismatch. Fix this by making the source of the mark conditional: when `same_pool=False`, use `_mark_batch_dynamic(output_state_indices)` for `h0_out_idx_` instead of always using `initial_state_indices`. Apply this conditional logic in both compile blocks where this pattern appears to ensure consistency.
🤖 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 `@flashinfer/gdn_kernels/gdn_decode_bf16_state.py`:
- Around line 1741-1743: The non-contiguous branch uses only pool_size and
initial_state_source.stride(0) for cache key generation, which can cause cache
collisions when tensors have identical pool_size and stride(0) but different
inner strides. Update the cache key generation in the non-contiguous branches to
use the full stride tuple tuple(initial_state_source.stride()) instead of just
stride(0), matching the approach used in the contiguous branch. This change must
be applied at two locations: where pool_size_key and pool_slot_stride are set
(around lines 1741-1743) and at the corresponding non-contiguous branch location
(around lines 2004-2006).
---
Outside diff comments:
In `@flashinfer/gdn_kernels/gdn_decode_bf16_state.py`:
- Around line 1700-1710: The kernel buffer stride calculation uses T in the
flattening logic (i_n * T * HV) but the intermediate_states_buffer is reshaped
using cache_steps (B_val * cache_steps * HV_val), causing batch writes to target
incorrect buffer slots when cache_steps > T. Fix this by either: (1) threading
the cache_steps parameter through the kernel launchers and updating the stride
calculation to use i_n * cache_steps * HV instead of i_n * T * HV, or (2)
tightening the assertions at the current location and the sibling locations
(lines 1931-1943) to require cache_steps == T instead of allowing cache_steps >=
T. Choose one approach and apply it consistently across all affected sites to
ensure stride calculations match the buffer layout.
- Around line 1789-1790: The `h0_out_idx_` variable is currently always
initialized using `initial_state_indices`, but when `same_pool=False`, the
kernel is launched with `output_state_indices` instead, causing a layout
mismatch. Fix this by making the source of the mark conditional: when
`same_pool=False`, use `_mark_batch_dynamic(output_state_indices)` for
`h0_out_idx_` instead of always using `initial_state_indices`. Apply this
conditional logic in both compile blocks where this pattern appears to ensure
consistency.
🪄 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: c1bcdfa2-bc1e-4957-a502-7373aa2b4de1
📒 Files selected for processing (1)
flashinfer/gdn_kernels/gdn_decode_bf16_state.py
Remove comments that restate adjacent code; keep only non-obvious rationale (explicit stride_order for B=1/T=1, contiguous-pool -1 sentinels, caching-off dummy-skip guards).
The non-contiguous (padded pool) branch keyed only on pool_size and stride(0); two pools sharing those but differing in inner strides would reuse a wrong cubin (cute.compile bakes the full layout), causing silent mis-addressing. Key on the full stride tuple instead (contiguous branch uses a (-1,) sentinel). Not reachable with vLLM's slot-padded-but-inner- contiguous pool, but strictly safer.
bkryu
approved these changes
Jun 16, 2026
Member
Author
|
/bot run tests/gdn |
Collaborator
Collaborator
|
[FAILED] Pipeline #55046740: 10/20 passed |
This was referenced Jun 22, 2026
ameynaik-hub
added a commit
to ameynaik-hub/flashinfer
that referenced
this pull request
Jun 25, 2026
Resolves conflicts in flashinfer/gdn_decode.py and flashinfer/gdn_kernels/gdn_decode_mtp.py between this PR's fp32 4D-pool feature (output_state_indices, use_pool_indexing, in-place 4D writeback) and main's flashinfer-ai#3649 (batch-size-agnostic compilation) and flashinfer-ai#3502 (BF16 recovery / per-request K, FLA per-token scatter). Two integration fixes beyond the mechanical conflict resolution: - Drop B and pool_size from the inline/warp cache keys: flashinfer-ai#3649 made kernel compilation batch-size-agnostic and removed them from the compiled kernel stubs, so the resolved cache-key tuples must match (was passing 20 args to an 18-arg stub). - Make the h0_source dlpack marking conditional on use_pool_indexing: main's mark_compact_shape_dynamic(stride_order=(0,1,2)) assumes a contiguous 3D flat pool and fails on the PR's 4D strided pool. Use plain from_dlpack for the 4D path (strides keyed via pool_strides_key; the pool-dim stride is batch-size-independent, so kernel reuse is safe). Verified on GB200 (sm_100): 32/32 fp32 MTP pool tests pass; benchmark shows no regression vs main on the standard contiguous path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
kahyunnam
added a commit
to kahyunnam/flashinfer
that referenced
this pull request
Jun 25, 2026
Integrate latest main (incl. flashinfer-ai#3502/flashinfer-ai#3649 GDN refactors) and resolve conflicts in the GDN decode kernels and tests. Per-batch q/k/v/a/b/o use mark_layout_dynamic to accept non-compact packed Q/K/V; pool/intermediate tensors keep mark_compact_shape_dynamic(mode=0) so launchers can derive constexpr tile counts.
kahyunnam
added a commit
that referenced
this pull request
Jun 26, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description #3649 made the GDN decode kernels batch-size agnostic by marking the per-batch tensors (q/k/v/a/b/o) with CuTe's `mark_compact_shape_dynamic`, which requires a compact layout. This was a functional layout coverage regression, and flagged as being incompatible with an e2e sglang test. The fix is to mark the per-batch tensors with `mark_layout_dynamic` instead of `mark_compact_shape_dynamic`. To verify this fix, this PR adds packed-QKV coverage for unit tests: q/k/v built as head-dim slices of a fused buffer (the SGLang layout), asserting results to the contiguous equivalents. These fail on current main (reproducing the regression) and pass with the fix. ## 🔍 Related Issues none ## 🚀 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 - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Improved dynamic-shape/layout annotations used for cached and compiled decoding kernels across pretranspose, nontranspose, and MTP paths, better aligning with fused/packed QKV layouts and ensuring consistent handling of cached intermediates. * **Bug Fixes** * Fixed decoding behavior for packed, non-contiguous Q/K/V tensor views to produce consistent results (including bit-identical outputs vs contiguous tensors). * **Tests** * Added regression tests covering pretranspose, nontranspose, and MTP decode with packed/non-contiguous Q/K/V views, including output/state equivalence checks. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
kahyunnam
pushed a commit
that referenced
this pull request
Jun 26, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description This is the **fp32 sibling of #3268** — the same 4D-pool support, now for the **fp32 MTP decode path** (`gated_delta_rule_mtp`). It does two things: 1. **In-place 4D pool writeback (correctness fix).** Previously a 4D state pool `[pool_size, HV, V, K]` was reshaped to 3D internally. For a **non-contiguous** pool (e.g. a strided slice of an oversized backing buffer) `.reshape()` silently materializes a **copy**, the kernel updates the copy, and the caller's pool is left untouched → **state updates lost**. The kernel now reads/writes the pool **in place via native 4D indexing** (`use_pool_indexing`), so writes land in the caller's tensor with no silent copy. 2. **Separate read/write slots (new capability).** New optional `output_state_indices` arg (shape `[B]`): where to **write** the updated state, separate from where it is **read** (`initial_state_indices`). Defaults to `initial_state_indices` (write back to the read slot); **negative entries skip the writeback** for that batch slot (matching the read-side `-1` padding semantics). This is what speculative-decoding / MTP-verify needs: read prior state from one pool slot, write the verified state to another. ```python out, state = gated_delta_rule_mtp( q, k, v, initial_state=pool, # 4D pool, contiguous OR strided initial_state_indices=read_idx, output_state_indices=write_idx, # NEW (optional; defaults to read_idx; -1 = skip) ... ) ``` The rest of the signature is unchanged; the standard contiguous single-token/MTP path behaves identically. ### Rebased on current `main` This branch was merged up to current `main`, which required reconciling the feature with two landed changes: - **#3649** (batch-size-agnostic GDN compilation): `B`/`pool_size` are no longer part of the kernel cache key, so the cache keys here were updated to match the batch-agnostic kernel stubs. - **#3502** (BF16 recovery / per-request K, FLA per-token scatter): the writeback guards now compose with `per_token_pool_scatter`. The non-contiguous 4D pool is passed through `from_dlpack` without `mark_compact_shape_dynamic` (which assumes a compact 3D layout); strides are keyed via `pool_strides_key` and the pool-dim stride is batch-size-independent, so a single compiled kernel is reused across batch sizes — compilation stays batch-size-agnostic. ## 🧪 Tests Added to `tests/gdn/test_decode_delta_rule.py`: - `test_mtp_fp32_state_pool` — pool read/write across seq lengths and batch sizes, optional separate output indices, intermediate-state caching. Params: `batch_size ∈ {1,4,16} × seq_len ∈ {2,4} × use_separate_output_indices ∈ {F,T} × cache_intermediate_states ∈ {F,T}` (24 cases). - `test_mtp_fp32_state_pool_non_contiguous` — strided (non-contiguous) pools. Params: `batch_size ∈ {1,4} × seq_len ∈ {2,4} × stride_multiplier ∈ {2,3}` (8 cases). This is the fp32 sibling of PR #3268 — same 4D-pool support, now for the fp32 MTP decode path. ### Before vs after — wrapper-side ```python # Before out, state = gated_delta_rule_mtp( q, k, v, initial_state=pool, # 4D pool initial_state_indices=read_idx, ... ) # Internally: pool is reshaped to 3D. If pool is non-contiguous, # this silently materializes a copy → kernel updates the COPY, # original pool is left untouched. Updates lost. # After out, state = gated_delta_rule_mtp( q, k, v, initial_state=pool, # 4D pool, contiguous OR strided initial_state_indices=read_idx, output_state_indices=write_idx, # NEW (optional, defaults to read_idx) ... ) # Kernel reads/writes pool in place via native 4D indexing. # No silent copy. Writes land in the caller's tensor. Before vs after — memory layout the kernel sees Contiguous pool (unchanged fast path): caller's [pool, HV, V, K] ─reshape view→ kernel's [pool*HV, V, K] (free, no copy) Non-contiguous pool (the new path): Before: caller's strided [pool, HV, V, K] ─.reshape() silent copy→ scratch [pool*HV, V, K] kernel writes scratch ❌ (caller's pool unchanged) After: caller's strided [pool, HV, V, K] ─pass through, use_pool_indexing=True→ kernel kernel writes caller's pool in place ✓ API change — one new argument output_state_indices : torch.Tensor, optional (shape [B]) Where to WRITE the updated state, separate from where you READ. Defaults to initial_state_indices (write back to the read slot). Negative entries skip the writeback for that batch. Everything else in the signature behaves identically. ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 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](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Multi-token FP32 decoding: supports pool-backed initial states, per-batch configurable output-state writeback indices, and correct handling for non-contiguous pooled layouts. Non-pool single-token decode behavior remains unchanged. * **Tests** * New FP32 tests cover pool read/write across sequence lengths and batch sizes, optional separate output indices, intermediate-state caching, non-contiguous pools, and stride variations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 tasks
5 tasks
kahyunnam
added a commit
that referenced
this pull request
Aug 15, 2026
…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. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## 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. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…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. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## 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. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
Addresses part 2 of #3602 , also brings over overlapping effort from @bkryu 's work in #3601
🔍 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