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

[amukkara]

@coderabbitai summary

Description

TTFT improvement

Qwen3.5-35B-A3B-FP8, 1 H100 400W, BS=1, ISL=10K / OSL=1

Config	Before	After	Δ
piecewise = off	268.7 ms	263.0 ms	2.1%
piecewise = on	282.7 ms	278.0 ms	1.7%

Test Coverage

Existing Qwen3.5 accuracy tests in tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B

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.

[coderabbitai]

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: 59cfbdc8-6f43-45a6-a329-865fedddee10

📥 Commits

Reviewing files that changed from the base of the PR and between 1fbd240 and 127257b.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/modules/mamba/gdn_mixer.py
• tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py

---

📝 Walkthrough

Walkthrough

Prefill metadata now computes device reset indices for sequences without initial states. Qwen3NextGatedDeltaNet.forward_core uses those indices to clear corresponding SSM and convolution cache slots.

Changes

Prefill state reset

Layer / File(s)	Summary
Prepare prefill reset indices tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py	Mamba2Metadata stages positions without initial states and derives reset_state_indices during prefill preparation.
Apply reset indices to caches tensorrt_llm/_torch/modules/mamba/gdn_mixer.py	forward_core zeros SSM and convolution state slots using reset_state_indices when available.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#16716: Updates the same mixer reset path with a similar index-based cache reset mechanism.

Suggested reviewers: kaiyux

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately reflects the performance-focused GDN cache-reset precomputation change.
Description check	✅ Passed	The description includes the required summary, performance impact, test coverage, and completed checklist items.
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.}

[amukkara]

Closing since #16716 fixes the same issue.