[None][perf] prepare_inputs: avoid O(seq_len) get_tokens(0) marshalling on the host by hyukn · Pull Request #16791 · NVIDIA/TensorRT-LLM

hyukn · 2026-07-23T08:27:27Z

Summary

request.get_tokens(0) marshals the whole C++ VecTokens (the entire sequence) into a fresh Python list of boxed ints — O(seq_len) host work that grows linearly with ISL. Three call sites in _prepare_tp_inputs / cuda_graph_runner pay this only to read a length or a small slice, and re-pay it every iteration. In chunked prefill the context loop re-marshals the full prompt for every chunk → O(L²/chunk) per request over the prefill. Normal decode is already immune (it uses get_last_tokens(0), O(1)); the waste is in the prefill/context path and the MTP first-draft paths.

Changes

New nanobind binding LlmRequest.get_tokens_range(beam, begin, end) — returns only the [begin, end) window, so nanobind copies just (end-begin) tokens (O(chunk)) instead of the whole VecTokens. Bounds are clamped.
context loop (_prepare_tp_inputs): use get_tokens_range(0, begin, end) for the current chunk and get_num_tokens(0) (O(1)) for the prompt length, instead of get_tokens(0) + slice + len(...).
first_draft loop (MTP): same — get_num_tokens(0) for the length, get_tokens_range for the last original_max_draft_len+1 tokens.
cuda_graph_runner first-draft branch: len(get_tokens(0)) → get_num_tokens(0).

Output is bit-identical — get_tokens_range returns exactly the same subrange the old get_tokens(0)[begin:end] produced.

Perf — `_prepare_inputs` host time (DSv4-Pro disagg, c3120, matched A/B, nsys, `nvtx_range("_prepare_inputs")`)

CTX / prefill worker (non-overlap, host-bound; primary target; iters 400–500, N=80):

`_prepare_inputs`	p50	p90	success
baseline	27.53 ms	43.29 ms	—
this PR	15.10 ms (−45%)	20.43 ms (−53%)	100% (3000/3000)

GEN / decode worker (overlap on; only the MTP first-draft path hits get_tokens(0), so only the large-L tail moves; iters 6000–6050, N=50):

`_prepare_inputs`	mean	p50	p90
baseline	17.84 ms	13.31 ms	32.78 ms
this PR	13.26 ms (−26%)	11.05 ms (−17%)	14.64 ms (−55%)

Each site drops from O(L) to O(1)/O(chunk); a synthetic microbench (bs128) shows each site flat across ISL after the fix (baseline get_tokens ~18 ms/iter at ISL 50K → ~0.05 ms). On CTX the whole _prepare_inputs shifts down (it sits on the exposed critical path). On GEN the median iter barely moves (normal decode is already O(1) via get_last_tokens) while the tail collapses (−55% p90) — the large-accumulated-length first-draft iters. GPU-timeline analysis on the same GEN capture shows the reduction is exposed, not hidden: GEN GPU bubble 25.7% → 20.7% and per-iteration cycle −6.3% with GPU compute unchanged (host win on the local critical path); pooled E2E SSE/throughput stay flat because GEN decode is not this workload's binding constraint. The binding was rebuilt and the new get_tokens_range verified present on the C++ LlmRequest; the full disagg run served at 100% success (functional parity).

🤖 Generated with Claude Code

…ng on the host `request.get_tokens(0)` marshals the whole C++ VecTokens (entire sequence) into a fresh Python list of boxed ints -- O(seq_len) host work that grows linearly with ISL. Three call sites in `_prepare_tp_inputs` / `cuda_graph_runner` pay this only to read a length or a small slice, and re-pay it every iteration. In chunked prefill the context loop re-marshals the full prompt for every chunk -> O(L^2/chunk) per request over the prefill. Normal decode is already immune (get_last_tokens(0), O(1)); the waste is in the prefill/context and MTP first-draft paths. Changes: - New nanobind binding `LlmRequest.get_tokens_range(beam, begin, end)` that copies only [begin, end) (O(chunk)) instead of the whole VecTokens. - context loop and first_draft loop: use get_tokens_range for the chunk and get_num_tokens(0) (O(1)) for lengths, instead of get_tokens(0) + slice. - cuda_graph_runner first-draft branch: len(get_tokens(0)) -> get_num_tokens(0). Output is bit-identical (get_tokens_range returns the same subrange). Measured on a DSv4-Pro disagg CTX worker (c3120, non-overlap, CTX nsys): _prepare_inputs p50 27.53 ms -> 15.10 ms (-45%), p90 43.29 -> 20.43 ms; 100% request success. Complements NVIDIA#16734 (which removes the overlap device-scalar sync, 268 -> 19.96 ms); together they target the two independent costs in _prepare_inputs. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

hyukn · 2026-07-23T14:16:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T14:23:17Z

PR_Github #61304 [ run ] triggered by Bot. Commit: 1916433 Link to invocation

coderabbitai · 2026-07-23T14:23:24Z

Walkthrough

Adds a bounded get_tokens_range request binding and updates executor code to use direct token counts or limited token windows for chunked prefill, first-draft requests, and CUDA graph sequence-length calculations.

Changes

Token Range Retrieval

Layer / File(s)	Summary
Token range binding `cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp`	Adds `get_tokens_range(beam, begin, end)`, which clamps bounds and returns a half-open token slice.
Model engine token windows `tensorrt_llm/_torch/pyexecutor/model_engine.py`	Uses bounded token retrieval for context chunks and first-draft suffixes, and obtains prompt length from the request token count.
CUDA graph sequence sizing `tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`	Uses `get_num_tokens(0)` for first-draft sequence-length calculation.

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

Suggested reviewers: cascade812

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description explains the change well, but it omits the required Description, Test Coverage, and PR Checklist sections from the template.	Add the missing template sections, especially concrete test coverage and a completed PR checklist, and align the headings with the repository template.

✅ Passed checks (4 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 specific and accurately describes the main performance change in the PR.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

hyukn · 2026-07-24T00:32:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-24T00:39:53Z

PR_Github #61437 [ run ] triggered by Bot. Commit: 1916433 Link to invocation

tensorrt-cicd · 2026-07-24T00:40:00Z

PR_Github #61304 [ run ] completed with state ABORTED. Commit: 1916433

Link to invocation

tensorrt-cicd · 2026-07-24T04:01:12Z

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

hyukn · 2026-07-24T09:54:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-24T10:00:29Z

PR_Github #61560 [ run ] triggered by Bot. Commit: 1916433 Link to invocation

tensorrt-cicd · 2026-07-24T13:51:38Z

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

hyukn · 2026-07-24T16:43:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-24T16:49:43Z

PR_Github #61604 [ run ] triggered by Bot. Commit: 1916433 Link to invocation

tensorrt-cicd · 2026-07-24T19:32:12Z

PR_Github #61604 [ run ] completed with state SUCCESS. Commit: 1916433
/LLM/main/L0_MergeRequest_PR pipeline #49813 completed with status: 'SUCCESS'

CI Report

Link to invocation

…ng on the host (NVIDIA#16791) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> (cherry picked from commit b8ff548)

…nnahz/dep-1083-port-flashinfer-stable-va-lifecycle-for-native-all-reduce * 'main' of https://github.com/NVIDIA/TensorRT-LLM: (54 commits) [NVIDIA#15673][fix] Enable CUDA core fast path for SM89/SM120/SM121 (NVIDIA#12705) [None][test] Adjust timeout cases in QA perf test (NVIDIA#16894) [https://nvbugs/6157892][fix] Mistral format refactor (NVIDIA#15123) [None][feat] Add kimi_k2/glm_5 grouped routing and fused router to bench_moe (NVIDIA#16830) [https://nvbugs/6501376][fix] Test-only fix — drop the `if hidden_size % 2 != 0: with pytest.raises(...)`… (NVIDIA#16844) [TRTLLM-13642][feat] Add perf sanity tests for Llama-3.1-8B and Gemma-3-1B and verify cache transceiver V2 support (NVIDIA#16355) [https://nvbugs/6433376][fix] Update the Dense test to mirror the MoE sibling — assert `bfloat16` under… (NVIDIA#16203) [None][fix] Resolve NVFP4 mixed-precision base layers for the DSpark draft (NVIDIA#16831) [https://nvbugs/6479324][test] Remove waiver for fixed qwen3_5_4b_fp8_stress disaggregated stress test (NVIDIA#16878) [https://nvbugs/6507109][infra] Split slow DGX B300 attention unit tests (NVIDIA#16838) [None][infra] Waive 21 failed cases for main in post-merge 2862 (NVIDIA#16882) [None][perf] prepare_inputs: avoid O(seq_len) get_tokens(0) marshalling on the host (NVIDIA#16791) [None][perf] Optimize Blackwell fused MHC half-MMA kernel (NVIDIA#16799) [None][infra] Auto-update test durations from OpenSearch (last 7 days) [None][perf] Skip DeepGEMM clean_logits in DSA indexer prefill on custom top-k path (NVIDIA#16789) [None][feat] Support DeepSeek-V4 in layer_wise_benchmarks (NVIDIA#16774) [https://nvbugs/6465993][fix] use attention cache dtype for disaggregated transfer (NVIDIA#16505) [https://nvbugs/6463822][fix] Fix LTX2 CUDA graph test leak issue (NVIDIA#16775) [https://nvbugs/5948435][chore] Unwaive DeepSeekV3Lite test_nvfp4_4gpus CUTLASS ep4 fp8kv on RTXPro6000D (NVIDIA#16621) [TRTLLM-14417][fix] Exclude ADP/cuda-graph dummy requests from speculative-decode acceptance stats (NVIDIA#16571) ... Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

…nnahz/dep-1082-shared-mnnvl-moe-lifecycle * 'main' of https://github.com/NVIDIA/TensorRT-LLM: (142 commits) [NVIDIA#15673][fix] Enable CUDA core fast path for SM89/SM120/SM121 (NVIDIA#12705) [None][test] Adjust timeout cases in QA perf test (NVIDIA#16894) [https://nvbugs/6157892][fix] Mistral format refactor (NVIDIA#15123) [None][feat] Add kimi_k2/glm_5 grouped routing and fused router to bench_moe (NVIDIA#16830) [https://nvbugs/6501376][fix] Test-only fix — drop the `if hidden_size % 2 != 0: with pytest.raises(...)`… (NVIDIA#16844) [TRTLLM-13642][feat] Add perf sanity tests for Llama-3.1-8B and Gemma-3-1B and verify cache transceiver V2 support (NVIDIA#16355) [https://nvbugs/6433376][fix] Update the Dense test to mirror the MoE sibling — assert `bfloat16` under… (NVIDIA#16203) [None][fix] Resolve NVFP4 mixed-precision base layers for the DSpark draft (NVIDIA#16831) [https://nvbugs/6479324][test] Remove waiver for fixed qwen3_5_4b_fp8_stress disaggregated stress test (NVIDIA#16878) [https://nvbugs/6507109][infra] Split slow DGX B300 attention unit tests (NVIDIA#16838) [None][infra] Waive 21 failed cases for main in post-merge 2862 (NVIDIA#16882) [None][perf] prepare_inputs: avoid O(seq_len) get_tokens(0) marshalling on the host (NVIDIA#16791) [None][perf] Optimize Blackwell fused MHC half-MMA kernel (NVIDIA#16799) [None][infra] Auto-update test durations from OpenSearch (last 7 days) [None][perf] Skip DeepGEMM clean_logits in DSA indexer prefill on custom top-k path (NVIDIA#16789) [None][feat] Support DeepSeek-V4 in layer_wise_benchmarks (NVIDIA#16774) [https://nvbugs/6465993][fix] use attention cache dtype for disaggregated transfer (NVIDIA#16505) [https://nvbugs/6463822][fix] Fix LTX2 CUDA graph test leak issue (NVIDIA#16775) [https://nvbugs/5948435][chore] Unwaive DeepSeekV3Lite test_nvfp4_4gpus CUTLASS ep4 fp8kv on RTXPro6000D (NVIDIA#16621) [TRTLLM-14417][fix] Exclude ADP/cuda-graph dummy requests from speculative-decode acceptance stats (NVIDIA#16571) ... Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

Cherry-pick only the 3 production files of NVIDIA#16791 (1916433) onto the af4aff4 container base so the A/B isolates PR#16791 from main drift. Co-Authored-By: Claude <noreply@anthropic.com>

Apply ONLY the 3-file diff of NVIDIA#16791 (1916433) onto af4aff4 via git apply (not whole-file checkout, which pulled unrelated bindings drift like expect_snapshot_points whose C++ methods don't exist in af4aff4's header). Isolates PR#16791 from main drift; +47/-8 lines, get_tokens_range only. Co-Authored-By: Claude <noreply@anthropic.com>

github-actions Bot assigned hyukn Jul 23, 2026

hyukn force-pushed the yukunh/prepare-inputs-gettokens-oL branch from 460954a to 1916433 Compare July 23, 2026 14:16

hyukn marked this pull request as ready for review July 23, 2026 14:16

hyukn requested review from a team as code owners July 23, 2026 14:16

hyukn requested review from Tabrizian, YihuiLu512, chienchunhung and lori-ren July 23, 2026 14:16

Tabrizian approved these changes Jul 27, 2026

View reviewed changes

hyukn merged commit b8ff548 into NVIDIA:main Jul 27, 2026
19 checks passed

Uh oh!

Conversation

hyukn commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Perf — _prepare_inputs host time (DSv4-Pro disagg, c3120, matched A/B, nsys, nvtx_range("_prepare_inputs"))

Uh oh!

hyukn commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

coderabbitai Bot commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

hyukn commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

hyukn commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

hyukn commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hyukn commented Jul 23, 2026 •

edited

Loading

Perf — `_prepare_inputs` host time (DSv4-Pro disagg, c3120, matched A/B, nsys, `nvtx_range("_prepare_inputs")`)

coderabbitai Bot commented Jul 23, 2026 •

edited

Loading