[https://nvbugs/6434535][fix] Replace boolean-mask indexed writes with a static-shape keep-mask…

[trtllm-agent]

Summary

• Root cause: Boolean-mask indexed write state_indices_p[~has_initial_states_p] on CUDA tensors lowers to torch.nonzero() with a shape-dependent D2H sync; under piecewise CUDA-graph capture this stalls per-rank divergently and HangDetector fires after 300s.
• Fix: Replace boolean-mask indexed writes with a static-shape keep-mask multiplication: gather rows via integer index state_indices_p, multiply by has_initial_states_p broadcast-shaped to the row dims. Rows with has_initial_states=True multiply by 1 (unchanged); False rows multiply by 0 (zeroed). Semantically identical, no dynamic-shape sync.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6434535

Summary by CodeRabbit

• Bug Fixes
  • Improved state handling for prefill requests without cached initial state.
  • Helps ensure cached sequence and convolution states are reset consistently during processing.

[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: 05b070bc-5e55-49d4-9d75-6d306932229d

📥 Commits

Reviewing files that changed from the base of the PR and between b4e7808 and b815c8a.

📒 Files selected for processing (1)

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

---

📝 Walkthrough

Walkthrough

Updates Qwen3NextGatedDeltaNet.forward_core to clear cached SSM and convolution states for prefill requests without initial states using a broadcasted keep mask.

Changes

Gated DeltaNet state handling

Layer / File(s)	Summary
Apply broadcasted keep mask tensorrt_llm/_torch/modules/mamba/gdn_mixer.py	Cached SSM and convolution states now use a read-modify-write mask that preserves entries with initial states and scales missing initial states to zero.

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

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and accurately summarizes the main keep-mask fix.
Description check	✅ Passed	The description clearly explains the root cause, fix, and test coverage, even though the checklist section is omitted.
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.}