[None][fix] Fix logits post processor for beam search in PyTorch backend

[cascade812]

Summary

Problem

The bug occurs when beam search (beam_width > 1) is combined with logits post processors. In _execute_logit_post_processors, the PyTorch backend indexed the logits tensor by request index (logits_tensor[idx]), which assumes one logits row per request. With beam search, each generation request contributes beam_width rows, so the row indexing was wrong: processors were applied to the wrong rows and only to a single beam.

This is why the bart beam2 tests failed while t5 passed — bart registers a _BartForcedTokensLogitsProcessor (to force decoder start/EOS tokens), while t5 has no logits post processor, so t5 never hit this path.

Fix

• Track a running logits_row_offset across context and generation requests, advancing it by beam_width per request (1 for context requests), so each request maps to its correct logits rows.
• Apply the processors to all beam_width rows of a generation request at once via the new _apply_logits_processors helper, passing per-beam token ids (request.get_tokens(beam_idx)), with logits shaped (beam_width, 1, vocab) to match the TRT backend's LogitsProcessor interface.
• Un-waive the two bart-large-cnn beam2 encoder-decoder tests previously skipped under this bug.

[coderabbitai]

📝 Walkthrough

Walkthrough

Refactors logit post-processing in PyTorchModelEngine into a new _apply_logits_processors helper and rewrites _execute_logit_post_processors to use context/generation grouping with explicit row-offset tracking instead of enumerate-based indexing. Also removes two BART end-to-end test waiver entries.

Changes

Logits Post-Processing Refactor

Layer / File(s)	Summary
Post-processor application helper tensorrt_llm/_torch/pyexecutor/model_engine.py	New _apply_logits_processors static-style method slices logits rows, reshapes to [beam_width, 1, vocab], validates processor signature length, invokes processors with request/token args, and writes results back into the tensor.
Row-offset based execution flow tensorrt_llm/_torch/pyexecutor/model_engine.py	_execute_logit_post_processors reworked to process context and generation request groups separately, compute per-request beam widths and token IDs, and track logits_row_offset to align flat logits rows, replacing the prior enumerate-based indexing.

BART Test Waiver Cleanup

Layer / File(s)	Summary
Waiver list update tests/integration/test_lists/waives.txt	Removes two SKIP entries for the bf16-kv-v1 beam2-bart-large-cnn BART encoder/decoder end-to-end test, covering both cuda-graph off and on variants.

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

Sequence Diagram(s)

sequenceDiagram
  participant ModelEngine as PyTorchModelEngine
  participant ScheduledRequests
  participant ApplyHelper as _apply_logits_processors

  ModelEngine->>ScheduledRequests: get context requests
  ModelEngine->>ModelEngine: compute beam_width=1, token_ids via get_tokens(0)
  ModelEngine->>ApplyHelper: apply processors on context logits slice
  ApplyHelper-->>ModelEngine: updated logits_tensor
  ModelEngine->>ScheduledRequests: get generation requests
  ModelEngine->>ModelEngine: compute per-request beam_width, per-beam token_ids
  ModelEngine->>ApplyHelper: apply processors on generation logits slice
  ApplyHelper-->>ModelEngine: updated logits_tensor
  ModelEngine->>ModelEngine: update logits_row_offset

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#15959: Both PRs touch the same CI-only file tests/integration/test_lists/waives.txt, removing/altering waiver entries.

Suggested reviewers: pcastonguay, yuxianq, mzweilz

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title is concise, specific, and matches the main PyTorch beam-search logits post-processor fix.
Description check	✅ Passed	The description explains the bug, fix, and impacted tests, but it omits the template's Test Coverage and PR Checklist sections.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 5870-5873: The logit post processor call in `model_engine.py`
accepts processors with either 4 or 5 parameters, but the invocation currently
always passes 5 positional arguments. Update the `lp(...)` call to match the
validated arity by branching on the processor signature (for example, using the
same `lp_params` check around `lp`) so 4-argument processors are called with 4
args and 5-argument processors receive the extra `None` argument.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2948cce2-6d41-48a7-b0fd-3949364f07ac

📥 Commits

Reviewing files that changed from the base of the PR and between a0c406f and a6cddd7.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #58033 [ run ] triggered by Bot. Commit: a6cddd7 Link to invocation

[tensorrt-cicd]

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

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #58079 [ run ] triggered by Bot. Commit: a6cddd7 Link to invocation

[tensorrt-cicd]

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

[cascade812]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58091 [ run ] triggered by Bot. Commit: a6cddd7 Link to invocation

[tensorrt-cicd]

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

[cascade812]

/bot run

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #58182 [ run ] triggered by Bot. Commit: a6cddd7 Link to invocation

[tensorrt-cicd]

PR_Github #58182 [ run ] completed with state SUCCESS. Commit: a6cddd7
/LLM/main/L0_MergeRequest_PR pipeline #46830 completed with status: 'SUCCESS'

CI Report

Link to invocation