[#17027][fix] Realign misaligned int32 index slices in FlashInfer GDN decode - #17028
[#17027][fix] Realign misaligned int32 index slices in FlashInfer GDN decode#17028Navjot10 wants to merge 4 commits into
Conversation
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. WalkthroughThe FlashInfer GDN decode path now realigns misaligned ChangesGDN decode alignment
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to The change realigns misaligned decode index slices before FlashInfer dispatch and adds focused regression coverage. No concrete current-head merge-blocking risk remains. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py`:
- Around line 28-36: Rename the module-level pytest marker skip_unsupported to
an UPPER_SNAKE_CASE constant name, then update its use as the decorator on
test_fi_decode_misaligned_index_slice. Add an explicit None return annotation to
that test function, preserving its existing behavior.
- Around line 21-24: Update the FlashInfer availability check around
gated_delta_rule to catch only the expected missing-module import exception,
allowing initialization and runtime errors to propagate and fail the regression
test instead of being reported as skips.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3c1489e2-a13d-4dc4-be04-b0ae5c02cacd
📒 Files selected for processing (2)
tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.pytests/unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py
87589ed to
5a060cc
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py`:
- Line 22: Wrap the gated_delta_rule import in the FlashInfer probe test across
multiple lines to satisfy the line-length limit, while preserving the F401
suppression for the intentionally unused import.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a1e5ab65-89cc-410f-90c3-85ed743212ff
📒 Files selected for processing (2)
tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.pytests/unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py
|
LGTM |
|
/bot run |
|
|
||
| pool_ref = state_pool.clone() | ||
| out_ref = _flashinfer_gdn_decode( | ||
| A_log=A_log, |
There was a problem hiding this comment.
Nit: you could define commonkwargs that then get passed **kwargs to both _flashinfer_gdn_decode calls along with each call's specific args (seems like only initial_state_source and initial_state_indices).
There was a problem hiding this comment.
Done in 9e04434 — the shared arguments now live in a single common_kwargs dict passed as **common_kwargs to both _flashinfer_gdn_decode calls; each call keeps only its own initial_state_source (per-call state-pool copy) and initial_state_indices (misaligned view vs aligned clone).
9e04434 to
889e1d5
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
…er GDN decode _flashinfer_gdn_decode passed initial_state_indices to the FlashInfer CuTe-DSL kernel through .int(), which is a no-op on an already-int32 tensor and keeps the pointer of a sliced view (e.g. the decode half state_indices[num_prefills:] of a mixed batch, offset 4*num_prefills bytes). The kernel asserts 32-byte data alignment on every tensor argument and rejects such views at runtime with 'Misaligned Tensor data on argument ... expected data alignment=32 bytes'; the dispatch gate does not check alignment, so there is no Triton fallback. Clone the index tensor into fresh, allocator-aligned storage when (and only when) it is misaligned, exactly like _flashinfer_gdn_verify (NVIDIA#15975) and the a/b activation guards in this same function (NVIDIA#15194). Adds a CUDA regression test mirroring test_fi_mtp_verify_misaligned_index_slice for the decode entry point. Signed-off-by: Navjot10 <83138142+Navjot10@users.noreply.github.com>
…annotation - Catch (ImportError, RuntimeError) instead of Exception in the FI availability probe. This mirrors the guard on the same import in fused_sigmoid_gating_recurrent.py and FlashInfer's own gdn_kernels/__init__.py: a missing build raises ImportError and a CuTe/CUTLASS mismatch raises RuntimeError at module import; in both cases production falls back to Triton, so the FI path is unavailable and the test should skip rather than error. Anything else still propagates as a failure. - Rename skip_unsupported -> SKIP_UNSUPPORTED (module-level constant). - Annotate the test with -> None. Signed-off-by: Navjot10 <83138142+Navjot10@users.noreply.github.com>
…e kwargs - The FI probe import line exceeded reviewers' preferred length, but the parenthesized wrap suggested in review cannot survive this repo's ruff config (line-length=100 with isort split-on-trailing-comma=false collapses any import that fits on one line). Import the module with an alias and probe the symbol via hasattr instead: the line drops to 79 chars, the '# noqa: F401' suppression becomes unnecessary, and the skip/fail semantics are unchanged (a broken submodule import still raises ImportError/RuntimeError inside the try, and a build without gated_delta_rule still reports unavailable). - Factor the arguments shared by both _flashinfer_gdn_decode calls into common_kwargs; each call now passes only its own state pool copy and index tensor (misaligned view vs aligned clone). Signed-off-by: Navjot10 <83138142+Navjot10@users.noreply.github.com> Assisted-By: devx/848d9b7d-732b-4283-aa05-274b4dc8e390
889e1d5 to
5afcb16
Compare
4 participants
Description
Fixes #17027.
_flashinfer_gdn_decode(tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py) passesinitial_state_indicesto the FlashInfer CuTe-DSL kernel through.int(), which is a no-op on an already-int32 tensor: it returns the same storage and keeps the data pointer of a sliced view (e.g. the decode halfstate_indices[num_prefills:]of a mixed batch, offset4 * num_prefillsbytes). The kernel asserts 32-byte data alignment on every tensor argument, so such a view is rejected at runtime:_can_use_flashinfer_gdn_decodedoes not check alignment, so this is a hard error with no Triton fallback. The hazard is dormant for the current in-tree callers (they pass zero-offset views) but fires for any caller handing this public dispatch API an offset slice at an element offset not divisible by 8.The fix mirrors the two guards that already exist for this exact bug class:
_flashinfer_gdn_verifyrealigns the same argument with a clone-when-misaligned guard ([None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel #15975);% 32guards for thea/bactivation slices in_flashinfer_gdn_decodeitself, but left the index tensor unguarded.The clone only happens when the pointer is actually misaligned, so the common aligned case stays zero-copy.
Test Coverage
tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py::test_fi_decode_misaligned_index_slice— feeds an int32 index slice with a 4-byte storage offset to_flashinfer_gdn_decodeand checks output and final state pool against an aligned-clone baseline. Without this fix, the call raises the misaligned-tensorValueError. Mirrors the existingtest_fi_mtp_verify_misaligned_index_slice. Requires SM90/SM100/SM103 + a FlashInfer build withgdn_kernels.gdn_decode_bf16_state(skips otherwise), so it needs GPU CI.tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py(verify-path twin) is unaffected.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
Dev Engineer Review
_flashinfer_gdn_decodeconvertsinitial_state_indicestoint32.QA Engineer Review
test_fi_decode_misaligned_index_slice().tests/integration/test_lists/for CI or manual QA.