[https://nvbugs/6276842][fix] Replace the two-case warmup distribution with a single shape: one real request…

[tensorrt-cicd]

Summary

• Root cause: Encoder CUDA graph captured the FMHA kernel for "all real, full bucket" seq_lens while runtime fed "real prompts + 1-token dummies" — different regime, FMHA grid/tokensInfo seqEnd-extension behavior diverged on Ada (SM89), drifting logits beyond 1e-3.
• Fix: Replace the two-case warmup distribution with a single shape: one real request capped at min(max_seq_len, num_tokens - (batch_size - 1)) plus (batch_size - 1) 1-token dummies — directly mirroring how pad_batch shapes the runtime input. Drop the now-passing test from waives.txt.
• 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/6276842

Summary by CodeRabbit

• Refactor

  • Improved encoder initialization with streamlined warmup input generation. Sequence length distribution now handles edge cases and complex configurations more robustly and efficiently.

• Tests

  • Re-enabled previously skipped encoder test case to expand test coverage and strengthen validation reliability.

[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: beabcceb-7eb2-4124-8163-43cb2189a5cc

📥 Commits

Reviewing files that changed from the base of the PR and between 6254f3a and 5829292.

📒 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

---

📝 Walkthrough

Walkthrough

The PR simplifies the encoder warmup input construction logic in the PyExecutor model engine by replacing a two-case token distribution strategy with a unified "runtime-mirror" approach that caps the first sequence length and appends single-token dummy entries. The previously-waived decoder encode test is removed from the waives list, indicating the logic fix resolves the test failure.

Changes

Encoder Warmup Simplification

Layer / File(s)	Summary
Encoder warmup sequence length construction and test waiver removal tensorrt_llm/_torch/pyexecutor/model_engine.py, tests/integration/test_lists/waives.txt	The _create_encoder_warmup_inputs function uses simplified length distribution: first length is min(max_seq_len, num_tokens - (batch_size - 1)), followed by batch_size - 1 single-token entries, with stricter validation for infeasible combinations. The tinyllama decoder encode cuda graph test waiver is removed from the skip list.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#14839: Modifies the same waives.txt file to add skip entries for related test cases.
• NVIDIA/TensorRT-LLM#14854: Also modifies waives.txt with test case waiver adjustments.
• NVIDIA/TensorRT-LLM#15056: Updates waives.txt to add new SKIP entries in the same test configuration.

Suggested reviewers

• greg-kwasniewski1
• tcherckez-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: replacing a two-case warmup distribution with a single shape that mirrors runtime behavior, directly aligned with the primary fix implemented in the changeset.
Description check	✅ Passed	The description includes root cause analysis, the fix explanation, test verification, and bug links. However, it lacks explicit test coverage listing and some PR checklist items.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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