[TRTLLM-12982][feat] support multi item scoring in LLM.encode by ixlmar · Pull Request #14693 · NVIDIA/TensorRT-LLM

ixlmar · 2026-05-28T13:32:38Z

Description

This exposes the functionality added in flashinfer-ai/flashinfer#1015 via LLM.encode (#12801). SGLang has been supporting this since sgl-project/sglang#10979.

Note: Commit e50c18a is being integrated via #14634 but is required for tests to pass.

Test Coverage

PR includes new tests.

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.

Summary by CodeRabbit

Release Notes

New Features
- Added multi-item scoring capability to the encode API, enabling evaluation of multiple candidate answers within a single batch request.
- Added batch_indexed_model_output and copy_logits_to_host parameters to the encode method for flexible output formatting and device control.
Tests
- Added comprehensive test coverage for multi-item scoring scenarios, input validation, and API constraints.

ixlmar · 2026-05-28T14:38:30Z

/bot run

tensorrt-cicd · 2026-05-28T14:45:44Z

PR_Github #50809 [ run ] triggered by Bot. Commit: d0824ae Link to invocation

ixlmar · 2026-05-28T15:01:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-28T15:07:02Z

PR_Github #50813 [ run ] triggered by Bot. Commit: d0824ae Link to invocation

tensorrt-cicd · 2026-05-28T15:07:07Z

PR_Github #50809 [ run ] completed with state ABORTED. Commit: d0824ae
/LLM/main/L0_MergeRequest_PR pipeline #40280 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

ixlmar · 2026-05-28T15:26:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-28T15:33:28Z

PR_Github #50819 [ run ] triggered by Bot. Commit: ccc4520 Link to invocation

tensorrt-cicd · 2026-05-28T15:37:05Z

PR_Github #50813 [ run ] completed with state ABORTED. Commit: d0824ae

Link to invocation

tensorrt-cicd · 2026-05-28T23:39:27Z

PR_Github #50819 [ run ] completed with state FAILURE. Commit: ccc4520
/LLM/main/L0_MergeRequest_PR pipeline #40289 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

ixlmar · 2026-05-29T07:50:11Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-29T07:56:34Z

PR_Github #51008 [ run ] triggered by Bot. Commit: dc49554 Link to invocation

ixlmar · 2026-05-29T09:08:57Z

/bot kill

tensorrt-cicd · 2026-05-29T09:15:26Z

PR_Github #51021 [ kill ] triggered by Bot. Commit: dc49554 Link to invocation

tensorrt-cicd · 2026-05-29T09:19:26Z

PR_Github #51008 [ run ] completed with state ABORTED. Commit: dc49554

Link to invocation

tensorrt-cicd · 2026-05-29T09:19:40Z

PR_Github #51021 [ kill ] completed with state SUCCESS. Commit: dc49554
Successfully killed previous jobs for commit dc49554

Link to invocation

ixlmar · 2026-05-29T09:43:34Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-29T09:49:01Z

PR_Github #51029 [ run ] triggered by Bot. Commit: 3b78547 Link to invocation

tensorrt-cicd · 2026-06-15T07:28:12Z

PR_Github #54253 [ run ] triggered by Bot. Commit: 9a8e8ad Link to invocation

tensorrt-cicd · 2026-06-15T09:33:59Z

PR_Github #54253 [ run ] completed with state SUCCESS. Commit: 9a8e8ad
/LLM/main/L0_MergeRequest_PR pipeline #43327 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

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

ixlmar · 2026-06-15T11:51:47Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-15T11:57:28Z

PR_Github #54296 [ run ] triggered by Bot. Commit: 735c68e Link to invocation

ixlmar · 2026-06-15T16:26:22Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-15T16:32:24Z

PR_Github #54326 [ run ] triggered by Bot. Commit: 735c68e Link to invocation

tensorrt-cicd · 2026-06-15T16:36:59Z

PR_Github #54296 [ run ] completed with state ABORTED. Commit: 735c68e

Link to invocation

tensorrt-cicd · 2026-06-15T21:36:13Z

PR_Github #54326 [ run ] completed with state SUCCESS. Commit: 735c68e
/LLM/main/L0_MergeRequest_PR pipeline #43397 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

…#14693) Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>

ixlmar requested a review from Funatiq May 28, 2026 13:32

github-actions Bot assigned ixlmar May 28, 2026

ixlmar changed the title ~~[TRTLLM-12982][feat] multi item scoring~~ [TRTLLM-12982][feat] support multi item scoring in LLM.encode May 28, 2026

ixlmar force-pushed the feat/multi-item-scoring branch 2 times, most recently from 78c8f63 to d0824ae Compare May 28, 2026 14:38

ixlmar force-pushed the feat/multi-item-scoring branch from d0824ae to ccc4520 Compare May 28, 2026 15:26

ixlmar added the api-compatible Accepted LLM API contract change that is backwards-compatible label May 29, 2026

ixlmar marked this pull request as ready for review May 29, 2026 14:30

ixlmar requested review from a team as code owners May 29, 2026 14:30

ixlmar requested a review from venkywonka May 29, 2026 14:30

ixlmar force-pushed the feat/multi-item-scoring branch from 9a8e8ad to a8a17f0 Compare June 15, 2026 09:35

ixlmar added 12 commits June 15, 2026 13:51

feat: support multi-item attention masking in LLM.encode()

7ca5bd9

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

fix: update LLMAPI test refs

823e4e1

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

fix: adjust doc-strings to pass LLMAPI stability test

d6308a4

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

fix: adjust to upstream CUDA graphs changes

d38b876

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

fix: reorder test fixtures

7bde25e

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

fix: reduce logging frequency

dc287da

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

typing: add _PageSizeOnePrefillWrapper

e93366a

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

doc: explain multi-item scoring in doc-string

0370b59

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

doc: FlashInferMultiItemParams

28ed980

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

fix: 1d-logits with batch_indexed_model_output

f94bc23

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

doc: improve comment

534776c

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

fix: do not forward **kwargs

735c68e

Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com>

ixlmar force-pushed the feat/multi-item-scoring branch from a8a17f0 to 735c68e Compare June 15, 2026 11:51

ixlmar enabled auto-merge (squash) June 15, 2026 19:26

ixlmar merged commit 35c9704 into NVIDIA:main Jun 15, 2026
8 checks passed

ixlmar deleted the feat/multi-item-scoring branch June 16, 2026 07:09

coderabbitai Bot mentioned this pull request Jun 16, 2026

[TRTLLM-12982][perf] reuse multi-item scoring position_ids and params #15413

Merged

1 task

ixlmar mentioned this pull request Jun 16, 2026

[TRTLLM-12982][chore] relocate torch_multi_arange #15416

Merged

1 task

tensorrt-cicd mentioned this pull request Jun 18, 2026

[https://nvbugs/6336801][fix] Add the two skip_softmax_threshold_scale_factor_decode/prefill aliases to… #15482

Closed

2 tasks

Uh oh!

Conversation

ixlmar commented May 28, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Summary by CodeRabbit

Release Notes

Uh oh!

ixlmar commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

ixlmar commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

ixlmar commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

ixlmar commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

ixlmar commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

ixlmar commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

ixlmar commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

ixlmar commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ixlmar commented May 28, 2026 •

edited by coderabbitai Bot

Loading