Conversation
ameynaik-hub
requested review from
bkryu,
kahyunnam,
yongwww and
yzh119
as code owners
June 2, 2026 01:50
Contributor
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThreads per-batch ChangesMTP Decode Output State Indexing
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
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 |
![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 introduces support for separate read and write indices (output_state_indices) in the Gated Delta Rule MTP kernels, routes fp32 state with T > 1 through the MTP kernel, and fixes an indexing bug in the BF16 state MTP kernel where intermediate states were incorrectly indexed by the pool slot instead of the batch index. Feedback suggests using h0_out_indices.to(initial_state_indices) to safely align both device and dtype, preventing potential device mismatch issues.
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.
Comment on lines
+2395
to
+2396
| if h0_out_indices.dtype != initial_state_indices.dtype: | ||
| h0_out_indices = h0_out_indices.to(initial_state_indices.dtype) |
Contributor
There was a problem hiding this comment.
Using h0_out_indices.to(initial_state_indices) is safer and more idiomatic than only casting the dtype. This automatically handles both device and dtype alignment, preventing potential device mismatch issues if output_state_indices is on a different device (e.g., CPU).
h0_out_indices = h0_out_indices.to(initial_state_indices)![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/gdn_decode.py (1)
723-741:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRequire a contiguous batch-sized intermediate-state buffer.
The cache is now documented and indexed per batch, but this block still only validates
cache_steps. With repeated pool indices,buffer_sizecan be smaller thanB, which makes the kernel write past the end. And ifreshape()/.contiguous()materializes a temporary, those writes never reachintermediate_states_buffer.Suggested validation
if cache_intermediate_states: buffer_size = intermediate_states_buffer.shape[0] cache_steps = intermediate_states_buffer.shape[1] + assert buffer_size >= B, ( + f"intermediate_states_buffer first dimension ({buffer_size}) must be >= B={B}" + ) assert cache_steps >= T, ( f"intermediate_states_buffer second dimension (cache_steps={cache_steps}) must be at least T={T} to prevent out-of-bounds indexing" ) assert intermediate_states_buffer.dtype == torch.float32, ( f"intermediate_states_buffer must be float32, " f"got {intermediate_states_buffer.dtype}" ) + assert intermediate_states_buffer.is_contiguous(), ( + "intermediate_states_buffer must be contiguous; otherwise kernel writes land in a temporary" + ) - - intermediate_states = intermediate_states_buffer.reshape( + intermediate_states = intermediate_states_buffer.view( buffer_size * cache_steps * HV, V, K ) - if not intermediate_states.is_contiguous(): - intermediate_states = intermediate_states.contiguous()🤖 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_decode.py` around lines 723 - 741, The code only checks cache_steps and dtype but must also ensure the provided intermediate_states_buffer is large and contiguous enough for per-batch indexing: validate that buffer_size (intermediate_states_buffer.shape[0]) is at least B (the batch size) and that intermediate_states_buffer.is_contiguous() is true (or call intermediate_states_buffer = intermediate_states_buffer.contiguous() before reshaping) so writes to the kernel land in the original buffer; additionally assert the total number of elements matches or exceeds B * cache_steps * HV * V * K (use intermediate_states_buffer.numel()) before performing the reshape into intermediate_states to prevent out-of-bounds or temporary-materialization issues when calling reshape()/contiguous().
🧹 Nitpick comments (1)
tests/gdn/test_decode_delta_rule.py (1)
1374-1447: ⚡ Quick win
cache_intermediate_states=Trueis exercised but never verified.When
cache_intermediate_statesisTrue,intermediate_bufferis allocated and passed to the pool-path kernel, but its contents are never asserted against a reference (the reference run usesintermediate_states_buffer=None). The[False, True]parametrization therefore only confirms the kernel doesn't crash with a buffer present — it does not validate that intermediate states are written correctly. Given the PR's focus on pool read/write indexing, this is the most valuable property to check here.Consider following the pattern in
_test_verify_kernel_mtp/_test_gdn_decode_bf16_state_mtp_kernel: gather the per-step reference states and compare againstintermediate_bufferafter the pool-path run.🤖 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 `@tests/gdn/test_decode_delta_rule.py` around lines 1374 - 1447, The test enables cache_intermediate_states and allocates intermediate_buffer but never verifies its contents; add assertions that intermediate_buffer contains the same per-step intermediate states produced by a reference run. Run the reference "gather→direct→scatter" path (the existing gated_delta_rule_mtp call that returns out_direct/updated_direct) while collecting per-step states (as done in _test_verify_kernel_mtp / _test_gdn_decode_bf16_state_mtp_kernel), then compare those per-step reference states to intermediate_buffer from the pool-path run (pool_under_test/out_pool) using torch.testing.assert_close with the same atol/rtol; ensure you only perform this check when cache_intermediate_states is True and reuse existing symbols: intermediate_buffer, cache_intermediate_states, gated_delta_rule_mtp, pool_under_test, out_pool, out_direct.
🤖 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_decode.py`:
- Around line 717-720: The code currently does h0_source =
initial_state.reshape(...), which can create a temporary when initial_state is
non-contiguous so mutations to h0_source are lost; before reshaping, ensure
initial_state is a contiguous tensor (e.g. replace initial_state with
initial_state.contiguous()) so reshape/view produces a real view and mutations
applied by run_mtp_decode() to h0_source persist; update the code around
h0_source / initial_state (and any callers expecting the returned state) to use
the contiguous copy.
In `@flashinfer/gdn_kernels/gdn_decode_bf16_state.py`:
- Around line 1879-1880: The code computes flat_idx using i_n (batch index) when
cache_intermediate_states is true but the BF16 wrapper still reshapes
intermediate_states_buffer assuming its first dimension equals the pool size,
which can be < B; to fix, ensure the buffer is batch-sized or index by
cache_idx: either (preferred) change the BF16 wrapper and any buffer allocation
for intermediate_states_buffer to allocate/reshape its first dimension to at
least B and add an assertion that intermediate_states_buffer.shape[0] >= B
before using i_n, or (alternative) change the indexing in the decode path to use
cache_idx instead of i_n (update flat_idx calculation). Reference symbols:
flat_idx, i_n, cache_idx, cache_intermediate_states, intermediate_states_buffer,
and the BF16 wrapper in gdn_decode_bf16_state.py.
---
Outside diff comments:
In `@flashinfer/gdn_decode.py`:
- Around line 723-741: The code only checks cache_steps and dtype but must also
ensure the provided intermediate_states_buffer is large and contiguous enough
for per-batch indexing: validate that buffer_size
(intermediate_states_buffer.shape[0]) is at least B (the batch size) and that
intermediate_states_buffer.is_contiguous() is true (or call
intermediate_states_buffer = intermediate_states_buffer.contiguous() before
reshaping) so writes to the kernel land in the original buffer; additionally
assert the total number of elements matches or exceeds B * cache_steps * HV * V
* K (use intermediate_states_buffer.numel()) before performing the reshape into
intermediate_states to prevent out-of-bounds or temporary-materialization issues
when calling reshape()/contiguous().
---
Nitpick comments:
In `@tests/gdn/test_decode_delta_rule.py`:
- Around line 1374-1447: The test enables cache_intermediate_states and
allocates intermediate_buffer but never verifies its contents; add assertions
that intermediate_buffer contains the same per-step intermediate states produced
by a reference run. Run the reference "gather→direct→scatter" path (the existing
gated_delta_rule_mtp call that returns out_direct/updated_direct) while
collecting per-step states (as done in _test_verify_kernel_mtp /
_test_gdn_decode_bf16_state_mtp_kernel), then compare those per-step reference
states to intermediate_buffer from the pool-path run (pool_under_test/out_pool)
using torch.testing.assert_close with the same atol/rtol; ensure you only
perform this check when cache_intermediate_states is True and reuse existing
symbols: intermediate_buffer, cache_intermediate_states, gated_delta_rule_mtp,
pool_under_test, out_pool, out_direct.
🪄 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: d462a7c6-b72c-483f-814f-7fc2dd33a8c2
📥 Commits
Reviewing files that changed from the base of the PR and between fc12ef2 and dc528bc5a25fe70876a789e2fdeffd59b5cc3089.
📒 Files selected for processing (4)
flashinfer/gdn_decode.pyflashinfer/gdn_kernels/gdn_decode_bf16_state.pyflashinfer/gdn_kernels/gdn_decode_mtp.pytests/gdn/test_decode_delta_rule.py
Comment on lines
+1879
to
+1880
| if cutlass.const_expr(cache_intermediate_states): | ||
| flat_idx = cache_idx * T * HV + i_t * HV + i_hv | ||
| flat_idx = i_n * T * HV + i_t * HV + i_hv |
Contributor
There was a problem hiding this comment.
Batch-indexed caching needs a batch-sized buffer.
flat_idx is now keyed by i_n, not cache_idx, but the BF16 wrapper still reshapes intermediate_states_buffer from its first dimension without asserting it is at least B. Pool mode legitimately allows pool_size < B when multiple batch rows share a state slot, so a pool-sized cache buffer will now index past the end here.
🤖 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 1879 - 1880,
The code computes flat_idx using i_n (batch index) when
cache_intermediate_states is true but the BF16 wrapper still reshapes
intermediate_states_buffer assuming its first dimension equals the pool size,
which can be < B; to fix, ensure the buffer is batch-sized or index by
cache_idx: either (preferred) change the BF16 wrapper and any buffer allocation
for intermediate_states_buffer to allocate/reshape its first dimension to at
least B and add an assertion that intermediate_states_buffer.shape[0] >= B
before using i_n, or (alternative) change the indexing in the decode path to use
cache_idx instead of i_n (update flat_idx calculation). Reference symbols:
flat_idx, i_n, cache_idx, cache_intermediate_states, intermediate_states_buffer,
and the BF16 wrapper in gdn_decode_bf16_state.py.
ameynaik-hub
force-pushed
the
ameyn/fix-fp32-mtp-pool-out-indices
branch
from
June 2, 2026 16:34
dc528bc to
665999d
Compare
Contributor
Author
|
/bot run |
Collaborator
ameynaik-hub
force-pushed
the
ameyn/fix-fp32-mtp-pool-out-indices
branch
2 times, most recently
from
June 3, 2026 00:34
a05ef49 to
fe43db2
Compare
Collaborator
|
[FAILED] Pipeline #53431684: 11/20 passed |
ameynaik-hub
force-pushed
the
ameyn/fix-fp32-mtp-pool-out-indices
branch
from
June 4, 2026 00:08
fe43db2 to
6340e6b
Compare
Contributor
Author
|
/bot run |
Collaborator
Collaborator
|
[FAILED] Pipeline #53598579: 9/20 passed |
…ty with bf16)
This PR fixes the two issues vLLM hit with the fp32 GDN MTP decode path:
Correctness: the wrapper's `.reshape(pool*HV, V, K)` silently densifies a
non-contiguous (page-strided) pool. The kernel then writes
that throwaway copy, dropping updates for vLLM-style pools.
Perf: the densification copy runs every call, regardless of whether
state actually changed.
The fix is in two layers:
1. Native 4D-pool support in both fp32 MTP kernels (gdn_decode_mtp.py):
- `gdn_verify_kernel_mtp` (warp-spec, B*HV > 128) and
`gdn_verify_kernel_mtp_inline` (small batch) each gain a
`use_pool_indexing: cutlass.Constexpr[bool]` switch.
- Once per CTA the kernel builds a 2D (V, K) view onto the pool slot.
The constexpr branch is the only site that knows the actual layout:
* True : 4D `[pool, HV, V, K]` — slice with (cache_idx, i_hv, :, :);
works for non-contiguous strided pools (vLLM).
* False: 3D `[pool*HV, V, K]` — slice with (flat_state_idx, :, :);
free reshape view of a contiguous pool (existing fast path).
- All ~43 `cute.local_tile(h0_source, (1, 1, vec_size), (flat_*_idx, X,
lane))` call sites are replaced with the view-based form
`cute.local_tile(h_*_view, (1, vec_size), (X, lane))`. Same memory
accesses, same instruction stream for the contiguous fast path.
- `flat_write_idx` and `write_cache_idx` are pre-declared / clamped to
satisfy CuTe DSL's "no variable out of control flow" rule. The
original-sign signal `write_cache_idx_raw` drives the per-site
write-skip gates so negative output indices still suppress the
writeback (preserving fp32 padding-skip semantics).
- Launchers extract `v_dim` / `k_dim` from the correct layout axes
depending on `use_pool_indexing`.
- `run_mtp_decode` cache key includes `use_pool_indexing` plus
`tuple(h0_source.stride())` (only when use_pool_indexing=True) so
different page-stride patterns each get their own compile and don't
alias to a stale binary.
2. Wrapper parity with the bf16 MTP path (gdn_decode.py:gated_delta_rule_mtp):
- Add `output_state_indices` parameter (mirrors the bf16 wrapper).
Defaults to `initial_state_indices`. Negative write indices skip the
writeback for that batch slot.
- Drop redundant `.to(torch.float32)` casts (state was already asserted
fp32). Validate `intermediate_states_buffer.dtype == float32`.
- Make `.contiguous()` on the intermediate buffer conditional, matching
the bf16 wrapper.
- Dispatch: when `initial_state.is_contiguous()`, take the existing 3D
fast path (free reshape view, `use_pool_indexing=False`). Else, pass
the 4D tensor through unchanged with `use_pool_indexing=True` — the
kernel writes the strided pool in place, no densification, no
scatter step.
- `intermediate_states_buffer` is still flat-indexed by batch (i_n), so
a non-contiguous buffer still triggers a staging copy + scatter back.
Native 4D for the intermediate buffer is a separate follow-up.
Additionally, `gated_delta_rule_decode_pretranspose` now routes
fp32 + T>1 (pool mode) through `gated_delta_rule_mtp` so the dispatcher
has a single entry point.
Tests (test_decode_delta_rule.py):
- `test_mtp_fp32_state_pool` (24 parametric variants): non-trivial
indices, optional separate output_state_indices, optional intermediate
caching. Verifies gather→direct reference parity, write destination,
and that non-targeted pool slots are bit-exactly unchanged.
- `test_mtp_fp32_state_pool_non_contiguous` (8 parametric variants:
B in {1, 4} x T in {2, 4} x stride_multiplier in {2, 3}). Allocates an
oversized HV-stride backing tensor and slices every Nth head-slot to
produce a strided 4D pool. Verifies output parity with a contiguous
reference, that the strided pool itself receives the updates (the
exact regression guard), and that interleaved non-selected backing
slots are bit-exactly unchanged (proves no densification copy).
Validation:
- Correctness: 149/149 pass across the new non-contig sweep, existing
contiguous fp32 MTP sweep (B 1..512 x T 2..8), bf16 verify,
pretranspose pool, negative_indices, and all_padding regressions.
- Perf (HV=64, B200, --update-state --cache-intermediate-states, 100
iters / 20 warmup, contiguous-pool path): median delta = 0%, mean
delta = +0.05% across 72 (BS, T) cells vs the pre-edit baseline. The
constexpr branch + view indirection compile away on the contiguous
fast path. Cell-level deltas within +/-5%, within run-to-run noise.
- Perf (contig vs strided pool, same B200): kernel-level delta is in
the noise (+/-2% for B >= 16). The strided path's win is in *not*
doing the per-call densification copy the old code would have done.
Review feedback addressed (CodeRabbit + Gemini):
- Device+dtype alignment of write indices: use
`output_state_indices.to(initial_state_indices)` (tensor target,
not just dtype) so a CPU-side output_state_indices is realigned to
the kernel's device automatically. No-op if already aligned.
- Strict intermediate-buffer validation in the wrapper:
* assert buffer_size >= B (kernel indexes by batch i_n in [0, B);
a smaller buffer caused silent OOB writes — the bf16 wrapper
already had this assert via PR flashinfer-ai#3145; fp32 wrapper was missing
the parallel).
* assert intermediate_states_buffer.is_contiguous() (caller
contract, removes the silent staging-copy fallback).
* Replace .reshape() with .view() so the no-copy contract is
enforced at runtime — raises if the layout ever doesn't support
a view.
* Removed the now-unused post-kernel scatter-back block.
- Test coverage: test_mtp_fp32_state_pool now passes a parallel
intermediate buffer to the reference path and verifies the cached
intermediate states match cell-for-cell when
cache_intermediate_states=True.
Re-verified: 34/34 spot-check tests pass; HV=64 BS×T perf sweep shows
mean Δ = +0.10%, median 0% vs pre-review-fix (kernel code byte-identical;
review fixes are wrapper-only).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
ameynaik-hub
force-pushed
the
ameyn/fix-fp32-mtp-pool-out-indices
branch
from
June 10, 2026 00:32
47135b6 to
e253324
Compare
Contributor
Author
|
/bot run tests/gdn |
Collaborator
Collaborator
|
[FAILED] Pipeline #54260359: 7/20 passed |
Contributor
Author
|
/bot run tests/gdn |
Collaborator
Collaborator
|
[FAILED] Pipeline #54716955: 9/20 passed |
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>
Contributor
Author
|
/bot run tests/gdn |
Collaborator
Contributor
Author
|
/bot run tests/gdn |
Collaborator
Collaborator
|
[FAILED] Pipeline #55856439: 8/20 passed |
kahyunnam
approved these changes
Jun 26, 2026
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
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:
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 thecaller'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.Separate read/write slots (new capability). New optional
output_state_indicesarg (shape[B]): where to write the updated state, separate from where it is read (initial_state_indices). Defaults toinitial_state_indices(write back to the read slot); negative entries skip thewriteback for that batch slot (matching the read-side
-1padding semantics). This is what speculative-decoding / MTP-verify needs: read prior state from one pool slot, write the verified state to another.The rest of the signature is unchanged; the standard contiguous single-token/MTP path behaves identically.
Rebased on current
mainThis branch was merged up to current
main, which required reconciling the feature with two landed changes:B/pool_sizeare no longer part of the kernel cache key, so the cache keys here were updated to match the batch-agnostic kernel stubs.per_token_pool_scatter.The non-contiguous 4D pool is passed through
from_dlpackwithoutmark_compact_shape_dynamic(which assumes a compact 3D layout); strides are keyed viapool_strides_keyand the pool-dim stride is batch-size-independent, so a single compiled kernel is reused across batch sizes — compilation staysbatch-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