[TRTLLM-14551][perf] avoid GDN state reset host synchronization by liji-nv · Pull Request #16716 · NVIDIA/TensorRT-LLM

liji-nv · 2026-07-22T08:26:13Z

GDN prefill reset currently selects cache slots with a CUDA boolean mask before assigning zeros. The boolean-indexed selection has a data-dependent output shape, which forces the host to obtain the selected element count and introduces a stream synchronization for every GDN layer.

Replace the two indexed assignments with a fixed-shape Triton launch. Each program reads the request's cache index and initialization flag on device, validates the cache slot, and clears both the SSM and convolution state rows in one kernel. Requests with initialized state or invalid sentinel indices remain untouched.

Add a CUDA unit test covering reset slots, preserved initialized slots, untouched unrelated slots, and the negative invalid-index sentinel.

Dev Engineer Review

Replaced CUDA masked host-side/state reset logic in Qwen3NextGatedDeltaNet.forward_core with an on-device fixed-shape Triton kernel reset path:
- Computes state_indices = mamba_metadata.state_indices[: num_prefills + num_decodes], splits into state_indices_p (prefills) and decode portion.
- When num_prefills > 0, calls _reset_gdn_states(ssm_states, conv_states, state_indices_p, has_initial_states_p) where has_initial_states_p = mamba_metadata.has_initial_states[:num_prefills].
Added Triton JIT kernel _reset_gdn_states_kernel and wrapper _reset_gdn_states in tensorrt_llm/_torch/modules/mamba/gdn_mixer.py to clear GDN state cache rows without host synchronization:
- For each request_idx:
  - Loads state_idx = state_indices[request_idx] and has_initial_states[request_idx].
  - Computes needs_reset = ~has_initial_states (reset only when initial state is not present).
  - Validates state_idx with valid_state = (state_idx >= 0) & (state_idx < NUM_CACHE_LINES) so sentinel -1 is treated as invalid.
  - Zeros ssm_states and conv_states rows via masked tl.store:
    - Store mask includes needs_reset & valid_state
    - Additionally bounds by per-buffer sizes: (offsets < SSM_STATE_SIZE) / (offsets < CONV_STATE_SIZE).
- Wrapper _reset_gdn_states:
  - Derives ssm_state_size = ssm_states.numel() // ssm_states.shape[0] and conv_state_size similarly.
  - Launches 2D grid (num_requests, triton.cdiv(max(ssm_state_size, conv_state_size), block_size)) with block_size=256.
  - Passes row addressing strides ssm_states.stride(0) / conv_states.stride(0) and NUM_CACHE_LINES = ssm_states.shape[0].
CI note from comments: L0 merge-request pipeline failed on multiple reruns, with NVIDIA team requesting review/fixes and rerun; no code changes beyond the reset-kernel/test are reflected here.

QA Engineer Review

Added CUDA-only unit test test_reset_gdn_states_preserves_initialized_and_invalid_slots in tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py.
- Coverage assertions:
  - Resets targeted non-initialized slot (state_indices=[1,...], has_initial_states=False → state_pool[1] becomes zero).
  - Preserves initialized slot (state_indices=[3,...], has_initial_states=True → state_pool[3] unchanged).
  - Preserves unrelated slots (state_pool[0], state_pool[2] unchanged).
  - Handles invalid sentinel index (state_indices includes -1 → no reset performed).
- Verdict: needs follow-up (no verification that this unit test is included in tests/integration/test_lists/ / test-db / qa coverage inputs).

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.

liji-nv · 2026-07-22T08:27:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-22T08:33:48Z

PR_Github #60933 [ run ] triggered by Bot. Commit: f8026d7 Link to invocation

coderabbitai · 2026-07-22T08:35:31Z

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

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: dbeed570-8cac-4015-8235-16cf03e22012

📥 Commits

Reviewing files that changed from the base of the PR and between 5699bbc and 3a00456.

📒 Files selected for processing (2)

tensorrt_llm/_torch/modules/mamba/gdn_mixer.py
tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py

🚧 Files skipped from review as they are similar to previous changes (2)

tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py
tensorrt_llm/_torch/modules/mamba/gdn_mixer.py

Walkthrough

The GDN mixer adds a Triton helper to reset selected SSM and convolution cache states, replaces inline tensor assignments in forward_core, and adds a CUDA test covering valid, invalid, initialized, and uninitialized slots.

Changes

GDN state reset

Layer / File(s)	Summary
Triton state reset helper `tensorrt_llm/_torch/modules/mamba/gdn_mixer.py`	Adds `_reset_gdn_states_kernel` and `_reset_gdn_states` with masks for state flags, valid indices, and state-size bounds.
Forward integration and validation `tensorrt_llm/_torch/modules/mamba/gdn_mixer.py`, `tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py`	Routes `forward_core` through the helper and verifies that only the targeted valid slot is zeroed while other slots retain their values.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Qwen3NextGatedDeltaNet
  participant _reset_gdn_states
  participant _reset_gdn_states_kernel
  participant StateBuffers
  Qwen3NextGatedDeltaNet->>_reset_gdn_states: pass state buffers, indices, and initial-state flags
  _reset_gdn_states->>_reset_gdn_states_kernel: launch masked reset
  _reset_gdn_states_kernel->>StateBuffers: zero selected SSM and convolution slots

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.29% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and matches the main change: removing host synchronization from GDN state reset.
Description check	✅ Passed	The description clearly explains the problem, solution, and test coverage, and the checklist is present and checked.
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-22T09:28:27Z

PR_Github #60933 [ run ] completed with state SUCCESS. Commit: f8026d7
/LLM/main/L0_MergeRequest_PR pipeline #49198 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

liji-nv · 2026-07-23T03:07:38Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T03:13:41Z

PR_Github #61176 [ run ] triggered by Bot. Commit: 57f5e21 Link to invocation

tensorrt-cicd · 2026-07-23T11:13:25Z

PR_Github #61176 [ run ] completed with state FAILURE. Commit: 57f5e21
/LLM/main/L0_MergeRequest_PR pipeline #49425 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

liji-nv · 2026-07-24T02:47:59Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-07-24T02:57:05Z

PR_Github #61475 [ run ] triggered by Bot. Commit: 2bd3f1e Link to invocation

tensorrt-cicd · 2026-07-24T07:55:20Z

PR_Github #61475 [ run ] completed with state FAILURE. Commit: 2bd3f1e
/LLM/main/L0_MergeRequest_PR pipeline #49695 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

liji-nv · 2026-07-24T09:59:43Z

/bot run --disable-fail-fast --add-multi-gpu-test

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

PR_Github #61562 [ run ] triggered by Bot. Commit: 2bd3f1e Link to invocation

tensorrt-cicd · 2026-07-24T15:20:54Z

PR_Github #61562 [ run ] completed with state FAILURE. Commit: 2bd3f1e
/LLM/main/L0_MergeRequest_PR pipeline #49773 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

liji-nv · 2026-07-27T05:09:20Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-07-27T05:15:37Z

PR_Github #61810 [ run ] triggered by Bot. Commit: 5699bbc Link to invocation

tensorrt-cicd · 2026-07-27T09:07:17Z

PR_Github #61810 [ run ] completed with state FAILURE. Commit: 5699bbc
/LLM/main/L0_MergeRequest_PR pipeline #50008 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

GDN prefill reset currently selects cache slots with a CUDA boolean mask before assigning zeros. The boolean-indexed selection has a data-dependent output shape, which forces the host to obtain the selected element count and introduces a stream synchronization for every GDN layer. Replace the two indexed assignments with a fixed-shape Triton launch. Each program reads the request's cache index and initialization flag on device, validates the cache slot, and clears both the SSM and convolution state rows in one kernel. Requests with initialized state or invalid sentinel indices remain untouched. Compute the state row offsets explicitly in int64. Recurrent-state views stride across the interleaved SSM and convolution pool, so multiplying an int32 cache index by either stride is not guaranteed to fit in int32 for a large pool. Add a CUDA unit test covering reset slots, preserved initialized slots, untouched unrelated slots, and the negative invalid-index sentinel. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv · 2026-07-28T05:02:10Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-07-28T05:09:26Z

PR_Github #62115 [ run ] triggered by Bot. Commit: 3a00456 Link to invocation

tensorrt-cicd · 2026-07-28T10:31:37Z

PR_Github #62115 [ run ] completed with state FAILURE. Commit: 3a00456
/LLM/main/L0_MergeRequest_PR pipeline #50297 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

liji-nv · 2026-07-28T12:42:29Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-07-28T12:48:54Z

PR_Github #62196 [ run ] triggered by Bot. Commit: 3a00456 Link to invocation

liji-nv · 2026-07-29T08:28:07Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-07-29T08:34:40Z

PR_Github #62474 [ run ] triggered by Bot. Commit: 3a00456 Link to invocation

tensorrt-cicd · 2026-07-29T08:36:48Z

PR_Github #62196 [ run ] completed with state ABORTED. Commit: 3a00456

Link to invocation

tensorrt-cicd · 2026-07-29T10:19:44Z

PR_Github #62474 [ run ] completed with state SUCCESS. Commit: 3a00456
/LLM/main/L0_MergeRequest_PR pipeline #50623 completed with status: 'SUCCESS'

CI Report

Link to invocation

liji-nv requested a review from a team as a code owner July 22, 2026 08:26

liji-nv requested a review from aswinvisva July 22, 2026 08:26

github-actions Bot assigned liji-nv Jul 22, 2026

nv-guomingz approved these changes Jul 22, 2026

View reviewed changes

nv-guomingz changed the title ~~[None][perf] avoid GDN state reset host synchronization~~ [TRTLLM-14551][perf] avoid GDN state reset host synchronization Jul 22, 2026

coderabbitai Bot mentioned this pull request Jul 22, 2026

[None][perf] Precompute GDN cache reset indices once per forward #16748

Closed

1 task

liji-nv force-pushed the liji/avoid-gdn-state-reset-sync branch from f8026d7 to 57f5e21 Compare July 23, 2026 03:07

aswinvisva reviewed Jul 23, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/mamba/gdn_mixer.py Outdated

liji-nv force-pushed the liji/avoid-gdn-state-reset-sync branch from 57f5e21 to 2bd3f1e Compare July 24, 2026 02:39

aswinvisva approved these changes Jul 24, 2026

View reviewed changes

liji-nv force-pushed the liji/avoid-gdn-state-reset-sync branch from 2bd3f1e to 5699bbc Compare July 27, 2026 05:08

liji-nv force-pushed the liji/avoid-gdn-state-reset-sync branch from 5699bbc to 3a00456 Compare July 28, 2026 05:01

liji-nv enabled auto-merge (squash) July 28, 2026 13:29

liji-nv merged commit af64dff into NVIDIA:main Jul 29, 2026
11 checks passed

Uh oh!

Conversation

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

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

liji-nv commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

coderabbitai Bot commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

liji-nv commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

Uh oh!

liji-nv commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

liji-nv commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

liji-nv commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

liji-nv commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

liji-nv commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

liji-nv commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

liji-nv commented Jul 22, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 22, 2026 •

edited

Loading