[None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel by nv-guomingz · Pull Request #15975 · NVIDIA/TensorRT-LLM

nv-guomingz · 2026-07-06T09:59:20Z

Summary by CodeRabbit

New Features
- Added a FlashInfer-backed speculative decode verification path for compatible GPU, dtype, and shape combinations.
- Expanded runtime dispatch to automatically choose the faster supported verification kernel when available.
Bug Fixes
- Preserved existing behavior as a fallback when FlashInfer support is unavailable or inputs are out of range.
- Ensured intermediate decode state remains isolated for each batch during verification.
Tests
- Added coverage for FlashInfer parity, buffer scoping, and environment-based dispatch controls.

Description

Test Coverage

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

nv-guomingz · 2026-07-06T10:01:55Z

/bot run

coderabbitai · 2026-07-06T10:06:23Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f90ddbb1-039d-4933-8f1c-9bec10340ac3

📥 Commits

Reviewing files that changed from the base of the PR and between 0044d5b and 76ea781.

📒 Files selected for processing (3)

tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py
tensorrt_llm/_torch/modules/mamba/gdn_mixer.py
tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py

📝 Walkthrough

Walkthrough

This PR adds a FlashInfer-backed GDN verify kernel path for speculative decoding, introducing eligibility gating (_can_use_flashinfer_gdn_verify) and a dispatch wrapper (_flashinfer_gdn_verify). Qwen3NextGatedDeltaNet's decode/extend paths conditionally use this kernel with a Triton fallback. New unit tests validate parity and gating behavior.

Changes

FlashInfer GDN Verify Integration

Layer / File(s)	Summary
Eligibility gating and verify kernel wrapper `tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py`	Imports MTP bf16-state kernel, adds `_FI_GDN_MAX_MTP_T` constant, `_can_use_flashinfer_gdn_verify` eligibility checks (env killswitches, arch, dtype, head dims, draft length), and `_flashinfer_gdn_verify` dispatch wrapper that writes intermediate states with state updates disabled.
Speculative decode/extend wiring `tensorrt_llm/_torch/modules/mamba/gdn_mixer.py`	`forward_decode` and `forward_extend` now check eligibility and route to `_flashinfer_gdn_verify` when supported, otherwise fall back to the existing Triton `fused_recurrent_gated_delta_rule_update` path.
Parity and gating tests `tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py`	New tests compare FlashInfer verify outputs/intermediate states against Triton reference, assert batch-scope buffer validation, and check gating behavior under killswitches and invalid shapes.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant GatedDeltaNet as Qwen3NextGatedDeltaNet
  participant Gate as _can_use_flashinfer_gdn_verify
  participant FIVerify as _flashinfer_gdn_verify
  participant Kernel as _fi_gdn_decode_bf16_state_mtp
  participant Triton as fused_recurrent_gated_delta_rule_update

  GatedDeltaNet->>Gate: check shapes, dtype, draft_token_num
  alt eligible
    GatedDeltaNet->>FIVerify: q,k,v,a,b,ssm_states,cache_indices
    FIVerify->>Kernel: dispatch, disable_state_update=True
    Kernel-->>FIVerify: write intermediate_states_buffer, output
    FIVerify-->>GatedDeltaNet: attention output
  else not eligible
    GatedDeltaNet->>Triton: beta,g,initial_state,disable_state_update=True
    Triton-->>GatedDeltaNet: attention output
  end

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description keeps the template placeholders and does not provide real Description or Test Coverage details.	Fill in the Description and Test Coverage sections with a brief problem/solution summary and the tests covering the new FlashInfer verify path.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly states the main change: dispatching GDN MTP target-verify to the FlashInfer bf16 kernel.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

tensorrt-cicd · 2026-07-06T10:07:24Z

PR_Github #57749 [ run ] triggered by Bot. Commit: 76ea781 Link to invocation

nv-guomingz · 2026-07-06T13:15:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T13:23:46Z

PR_Github #57768 [ run ] triggered by Bot. Commit: 418050f Link to invocation

tensorrt-cicd · 2026-07-06T13:28:26Z

PR_Github #57749 [ run ] completed with state ABORTED. Commit: 76ea781

Link to invocation

tensorrt-cicd · 2026-07-06T16:35:26Z

PR_Github #57768 [ run ] completed with state SUCCESS. Commit: 418050f
/LLM/main/L0_MergeRequest_PR pipeline #46471 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

nv-guomingz · 2026-07-07T02:37:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T02:43:26Z

PR_Github #57898 [ run ] triggered by Bot. Commit: 61dbc4f Link to invocation

tensorrt-cicd · 2026-07-07T06:37:44Z

PR_Github #57898 [ run ] completed with state SUCCESS. Commit: 61dbc4f
/LLM/main/L0_MergeRequest_PR pipeline #46592 completed with status: 'SUCCESS'

CI Report

Link to invocation

…NVIDIA#15975) Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

…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>

nv-guomingz requested review from a team as code owners July 6, 2026 09:59

nv-guomingz requested review from symphonylyh, tomeras91 and yuxianq July 6, 2026 09:59

github-actions Bot assigned nv-guomingz Jul 6, 2026

nv-guomingz force-pushed the user/guomingz/update_qwen3.5_mtp_path branch from 76ea781 to 418050f Compare July 6, 2026 13:14

[None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel

61dbc4f

Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

nv-guomingz force-pushed the user/guomingz/update_qwen3.5_mtp_path branch from 418050f to 61dbc4f Compare July 7, 2026 02:37

leslie-fang25 approved these changes Jul 7, 2026

View reviewed changes

Wanli-Jiang reviewed Jul 7, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py

Wanli-Jiang approved these changes Jul 7, 2026

View reviewed changes

nv-guomingz merged commit c09fbbd into NVIDIA:main Jul 7, 2026
7 checks passed

Navjot10 mentioned this pull request Jul 29, 2026

[Bug]: GDN FlashInfer decode: initial_state_indices.int() is a no-op for int32 and leaves misaligned index slices (Misaligned Tensor data, alignment=32 bytes) #17027

Open

2 tasks

Navjot10 mentioned this pull request Jul 29, 2026

[#17027][fix] Realign misaligned int32 index slices in FlashInfer GDN decode #17028

Open

1 task

nv-guomingz deleted the user/guomingz/update_qwen3.5_mtp_path branch August 28, 2026 02:44

Navigation Menu

Uh oh!

[None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel - #15975

[None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel#15975
nv-guomingz merged 1 commit into
NVIDIA:mainfrom
nv-guomingz:user/guomingz/update_qwen3.5_mtp_path

nv-guomingz commented Jul 6, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

nv-guomingz commented Jul 6, 2026

Uh oh!

coderabbitai Bot commented Jul 6, 2026

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

nv-guomingz commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

nv-guomingz commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

nv-guomingz commented Jul 6, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

nv-guomingz commented Jul 6, 2026

Uh oh!

coderabbitai Bot commented Jul 6, 2026

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

nv-guomingz commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

nv-guomingz commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

nv-guomingz commented Jul 6, 2026 •

edited by coderabbitai Bot

Loading