[TRTLLM-9295][fix] unflake test_overlap_scheduler.py::test_overlap_scheduler_consis…

[ixlmar]

Description

Test case unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[TorchSampler] was failing:

• only when FlashInfer.sampling was enabled

• in B200_PCIe-PyTorch-1 CI stage

• repro on another B200 system, but not on L40S, nor on A10

To root cause the failure, the seq. slots scheduled in each iteration and the resulting sampled tokens
were logged.

The observed flakiness is rooted in the following:

• Overlap scheduling typically speculatively runs an additional generation step, even for requests
  which have already reached their maximum number of output tokens.

• There is hardware dependence in the impact of the above, because e.g. on L40S some requests
  terminate with EOS, whereas they run until max output tokens on B200.

• Although sampling is deterministic as a whole, neither the PyTorch-native, nor the FlashInfer.sampling backends, are batch invariant:

  coin = torch.tensor([0.5, 0.5], device="cuda")
  
  torch.multinomial(
      torch.tile(coin.unsqueeze(0), (6, 1)),
      num_samples=1,
      generator=torch.Generator(device="cuda").manual_seed(42),
  )
  
  # tensor([[0],
  #       [0],
  #       [1],
  #       [0],
  #       [0],
  #       [1]], device='cuda:0')
  
  flashinfer.sampling.sampling_from_probs(
      torch.tile(coin.unsqueeze(0), (9, 1)),
      deterministic=True,
      check_nan=False,
      generator=torch.Generator(device="cuda").manual_seed(42),
  )
  
  # tensor([1, 0, 0, 0, 0, 0, 1, 0, 1], device='cuda:0', dtype=torch.int32)

• Thus, whenever overlap scheduling runs an extra generation step for some request
  before all requests internally ordered after that request have completed, the
  test fails.

• In addition, the extra generation steps run by overlap scheduling can change
  the number of prefill chunks needed (the test comprises three prompts, with
  one being much longer than the other two). This changes the number of samples
  drawn from the torch.Generator (TorchSampler._generator) and thus causes
  a discrepancy in the tokens generated for the affected request.

Whereas the last issue can be worked around by resetting the RNG state after every
sampling iteration, the lack of batch invariance still requires requests to
complete in an ordering consistent with the internal request batch ordering,
making the test brittle (fail on different hardware or upon sampling/RNG implementation
changes).

Reducing the sampling temperature and tweaking max_num_tokens (which affects
the chunked prefill) increases the chances that the test passes but does not remove
the conceptual flakiness.

As long as batch-invariant generation is not available, the test is therefore
modified to use greedy sampling, which is deterministic (given same logits).

In fact, this test was implemented before #7294, when greedy sampling was still
the default (test unchanged by PR, enable_mixed_sampler not used in test, so that temperature was not used).

Test Coverage

PR fixes existing test.

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)

• 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

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

Summary by CodeRabbit

• Tests
  • Updated test input data and configuration parameters to enhance test validation scenarios and improve coverage accuracy.

[ixlmar]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

The pull request modifies a test input JSON file, replacing the final prompt block's complete assertion "The best answer is [A/B/C/D]" with an incomplete string "The best answer is" and changing the temperature field from 1.0 to null.

Changes

Cohort / File(s)	Change Summary
Test data updates tests/unittest/_torch/executor/test_overlap_scheduler_input.json	Modified final prompt assertion to remove explicit answer placeholder; changed temperature field from 1.0 to null

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

• Verify that the incomplete assertion string "The best answer is" is intentional and serves the test's purpose
• Confirm the temperature change to null aligns with test expectations and doesn't unintentionally alter scheduler behavior

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: fixing a flaky test in test_overlap_scheduler.py by addressing consistency issues, with proper JIRA ticket reference and type annotation.
Description check	✅ Passed	The PR description provides comprehensive explanation of the issue, root cause analysis, and solution. It includes detailed test coverage information and a completed PR checklist confirming adherence to guidelines.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

tests/unittest/_torch/executor/test_overlap_scheduler_input.json (1)

5-5: Incomplete prompt ending is intentional for deterministic test.

The truncated final prompt (ending with "The best answer is" instead of a complete assertion) is consistent with the PR objective to use greedy decoding for deterministic output. By requiring the model to generate the answer token, combined with greedy sampling (via temperature: null), the test output becomes deterministic and reproducible across hardware.

One minor suggestion: consider adding an inline comment in the JSON explaining why this prompt is intentionally incomplete.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 9ef7eb7 and 346a13b.

📒 Files selected for processing (1)

• tests/unittest/_torch/executor/test_overlap_scheduler_input.json (1 hunks)

🧰 Additional context used

🧠 Learnings (1)

📓 Common learnings

Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: Check PR Checklist Resolution
• GitHub Check: Pre-commit Check

🔇 Additional comments (1)

tests/unittest/_torch/executor/test_overlap_scheduler_input.json (1)

8-8: Verify that temperature: null triggers deterministic greedy sampling.

Per the PR objectives, the fix replaces non-deterministic sampling (temperature=1.0) with greedy sampling to eliminate batch-invariant variance. Setting temperature to null should default to greedy decoding, but this warrants verification given the criticality to the test fix.

Could you confirm that temperature: null in this test input JSON results in deterministic greedy (temperature=0) sampling behavior in the TensorRT-LLM executor? If the executor uses a different configuration mechanism for greedy mode (e.g., a separate use_greedy flag), that should be set instead.

[tensorrt-cicd]

PR_Github #24479 [ run ] triggered by Bot. Commit: 346a13b

[yongwww]

/bot run

[tensorrt-cicd]

PR_Github #24479 [ run ] completed with state SUCCESS. Commit: 346a13b
/LLM/main/L0_MergeRequest_PR pipeline #18473 completed with status: 'FAILURE'

[ixlmar]

/bot run

[tensorrt-cicd]

PR_Github #24569 [ run ] triggered by Bot. Commit: 346a13b

[tensorrt-cicd]

PR_Github #24569 [ run ] completed with state SUCCESS. Commit: 346a13b
/LLM/main/L0_MergeRequest_PR pipeline #18544 completed with status: 'SUCCESS'