[None][fix] Keep GDN decode recurrence stable in deterministic mixed batches

[shikicloud]

Dev Engineer Review

When FORCE_DETERMINISTIC=1 and a non-verification batch contains decode tokens, forward_extend now runs prefill and decode through separate paths. Decode uses _recurrent_update; prefill receives only prefill state indices, sequence boundaries, and output. This aims to keep decode recurrence results stable when prefill requests share the batch.

The change also caches the prefill query-start view for repeated metadata preparation. Other batches retain the combined path. No changed public declarations or unrelated files are reported. Test execution results and current review findings were not supplied.

QA Engineer Review

The new CUDA-only test covers exact output and state equality between mixed execution and separately executed prefill and decode. It varies prefill backend, sequence lengths, state dtype, and output allocation. It also checks inactive state slots, output-buffer reuse, and CUDA graph replay. FlashInfer cases skip when the backend is unavailable. The test name was not found in the integration test-db/ or qa/ lists. Coverage verdict: sufficient for the targeted mixed-batch regression; execution results were not provided.

Per-File QA Perspective

• tensorrt_llm/_torch/modules/mamba/gdn_mixer.py: Verify deterministic mixed prefill/decode dispatch, decode state updates, output-buffer slicing, and graph-capture metadata reuse. Other batches should continue to use the combined path.
• tests/unittest/_torch/modules/mamba/test_gdn_mixed_decode.py: Adds test_mixed_decode_matches_decode_only, parameterized across two prefill backends, five prefill layouts, two state dtypes, and two output modes. It is not listed in the searched integration test-db/ or qa/ lists.

Description

Decode requests in mixed batches. This prevents batch composition from changing decode outputs and final state for identical inputs.
This issue was found during prefill/decode consistency testing. No associated NVBug is recorded. Related to #19501.

Test Coverage

• Covers standalone versus mixed decode, FlashInfer/Triton prefill, FP32/BF16 state, irregular layouts, and CUDA graphs.
• Initial regressions: 16 failed before, 16 passed after.
• Expanded suite: 40 passed, including with the other GDN fixes removed.

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: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 2bd1b63d-0c25-4322-9656-29b9bae1a72b

📥 Commits

Reviewing files that changed from the base of the PR and between eef4a51 and 99b615e.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/modules/mamba/gdn_mixer.py
• tests/unittest/_torch/modules/mamba/test_gdn_mixed_decode.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The GDN mixer now separates prefill and decode processing for deterministic mixed batches. A CUDA regression test compares mixed execution with separate prefill and decode execution.

Changes

GDN mixed prefill and decode

Layer / File(s)	Summary
Prepare prefill boundaries and recurrent decode tensorrt_llm/_torch/modules/mamba/gdn_mixer.py	The mixer caches a view of prefill query boundaries and routes ordinary decode updates through _recurrent_update.
Split deterministic mixed-batch execution tensorrt_llm/_torch/modules/mamba/gdn_mixer.py, tests/unittest/_torch/modules/mamba/test_gdn_mixed_decode.py	For deterministic, non-verification batches with decode tokens, prefill uses chunk_gated_delta_rule and decode uses _recurrent_update. CUDA tests compare outputs and state, check inactive slots, and validate preallocated output and CUDA graph replay.

Priority: ➖ Normal

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

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant GDN as gdn_mixer
  participant Chunk as chunk_gated_delta_rule
  participant Recurrent as _recurrent_update
  participant Caller
  GDN->>Chunk: Prefill tensors, state indices, and sequence boundaries
  Chunk-->>GDN: Prefill output
  GDN->>Recurrent: Decode tensors, state indices, and adjusted boundaries
  Recurrent-->>GDN: Decode output
  GDN-->>Caller: Return supplied output buffer or concatenated outputs

Loading

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to 99b61

No actionable merge-blocking risk is identified for the deterministic mixed-batch change; it is ready to merge after normal checks.

Pre-merge checks · 4 · 1

Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 10 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Title check		The title follows the required [None][fix] format and clearly describes the main change: stabilizing GDN decode recurrence in deterministic mixed batches.
Description check		The description includes the issue, solution, related issue, test coverage, reported results, and the required checklist. It is sufficiently complete and directly aligned with the pull request changes…

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[shikicloud]

/bot run

[tensorrt-cicd]

PR_Github #77607 [ run ] triggered by Bot. Commit: 99b615e Link to invocation

[tensorrt-cicd]

PR_Github #77607 [ run ] completed with state SUCCESS. Commit: 99b615e
/LLM/main/L0_MergeRequest_PR pipeline #63980 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