[None][feat] Add low-latency host task dispatch mode for guided decoding by achartier · Pull Request #15863 · NVIDIA/TensorRT-LLM

achartier · 2026-07-02T02:39:13Z

Summary by CodeRabbit

New Features
- Added an optional low-latency host dispatch setting for faster callback handling in supported environments.
- Exposed a new configuration option so it can be enabled from the public API.
Bug Fixes
- Older CUDA setups now fall back gracefully to the standard dispatch path when low-latency mode isn’t supported.
Documentation
- Updated API help text to describe the new option and its compatibility requirements.

Description

Adds enable_low_latency_host_dispatch (default False) to TorchLlmArgs.
When enabled, CUDA host callbacks are dispatched via cudaLaunchHostFunc_v2
with cudaHostTaskSpinWait instead of the default blocking mode, reducing
guided-decoder callback latency by keeping the CPU thread spinning until the
GPU event fires rather than sleeping.

This matters most when Eagle3 speculative decoding is combined with guided
decoding: in that configuration the guided-decoder @hostfunc callbacks are
captured inside the CUDA graph (called from _execute_guided_decoder_if_present
during the target model forward pass), so the spin-wait policy applies to every
decode step and the per-step wake-up latency directly affects inter-token delay.

Requires CUDA 13.2+ (cudaHostTaskSyncMode enum, cudaHostTaskSpinWait = 0x1).
Falls back silently to cudaLaunchHostFunc on older drivers with a log warning.

Files changed:

cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp — add bool lowLatency param; use cudaLaunchHostFunc_v2 behind CUDART_VERSION >= 13020 guard
tensorrt_llm/_torch/hostfunc.py — module-level _low_latency_dispatch flag + set_low_latency_dispatch() setter
tensorrt_llm/_torch/pyexecutor/_util.py — call setter from create_py_executor_instance
tensorrt_llm/llmapi/llm_args.py — TorchLlmArgs field (status="prototype")

Test Coverage

Unit tests (no GPU required beyond existing fixture):

tests/unittest/bindings/test_hostfunc.py — test_low_latency_dispatch: verifies the module-level flag is toggled and hostfuncs fire correctly in both dispatch modes
tests/unittest/llmapi/test_llm_args.py — TestEnableLowLatencyHostDispatch: default value, explicit True/False, YAML round-trip, invalid-type rejection

API stability:

tests/unittest/api_stability/references/llm.yaml — field entry added (required to keep test_api_status green)

Integration test:

tests/integration/defs/accuracy/test_llm_api_pytorch.py — test_guided_decoding_with_eagle3_low_latency_dispatch: exercises the flag with Eagle3 one-model + xgrammar/llguidance + CUDA graphs on Hopper+, validating end-to-end correctness via JsonModeEval

PR Checklist

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.

coderabbitai · 2026-07-02T02:46:28Z

📝 Walkthrough

Walkthrough

This PR adds a low-latency dispatch mode for CUDA host function launches. The C++ binding gains a lowLatency parameter using cudaLaunchHostFunc_v2 with spin-wait on CUDA 13.2+, wired through a Python setter, executor initialization, and a new TorchLlmArgs.enable_low_latency_host_dispatch config field, with accompanying tests.

Changes

Low-latency host dispatch

Layer / File(s)	Summary
CUDA host-function launch binding `cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp`	`launchHostFunc` gains a `lowLatency` parameter, using `cudaLaunchHostFunc_v2` with `cudaHostTaskSpinWait` on CUDA 13.2+ or falling back to `cudaLaunchHostFunc`; docstring updated.
Python dispatch toggle `tensorrt_llm/_torch/hostfunc.py`	New `_low_latency_dispatch` flag and `set_low_latency_dispatch` setter; `launch_hostfunc` forwards the flag to the binding call.
Executor wiring `tensorrt_llm/_torch/pyexecutor/_util.py`	`create_py_executor_instance` imports and invokes `set_low_latency_dispatch` using `llm_args.enable_low_latency_host_dispatch`.
LLM args config field `tensorrt_llm/llmapi/llm_args.py`, `tests/unittest/api_stability/references/llm.yaml`, `tests/unittest/llmapi/test_llm_args.py`	Adds `enable_low_latency_host_dispatch` bool field (prototype status) to `TorchLlmArgs`, API stability reference, and validation tests.
Tests `tests/unittest/bindings/test_hostfunc.py`, `tests/integration/defs/accuracy/test_llm_api_pytorch.py`	New unit test toggling low-latency dispatch and verifying callback output; new Eagle3 guided-decoding accuracy test enabling the flag.

Estimated code review effort: 2 (Simple) | ~15 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Executor
  participant hostfuncPy as hostfunc.py
  participant Bindings as bindings.launch_hostfunc
  participant CUDA as CUDA Runtime

  Executor->>hostfuncPy: set_low_latency_dispatch(enabled)
  hostfuncPy->>hostfuncPy: update _low_latency_dispatch
  hostfuncPy->>Bindings: launch_hostfunc(stream, freeUserData, low_latency, callback)
  Bindings->>CUDA: cudaLaunchHostFunc_v2(cudaHostTaskSpinWait) or cudaLaunchHostFunc
  CUDA-->>Bindings: launch result

Related issues: None specified.

Related PRs: None specified.

Suggested labels: low-latency, hostfunc, prototype

Suggested reviewers: None specified.

Poem

A rabbit spins instead of sleeps,
Low-latency, the fast dispatch it keeps.
Through C++ bindings, flags now flow,
To executor and args, the toggle does go.
Tests confirm the switch works right—
Hop, spin, launch, take flight! 🐇⚡

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the PR’s main change and follows the repository’s required ticket/type/summary format.
Description check	✅ Passed	The description follows the template with Description, Test Coverage, and PR Checklist, and includes the key implementation and test details.
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.

✨ 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 `@cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp`:
- Around line 84-100: The fallback in `launchHostFunc` logs a warning every time
`lowLatency` is true on CUDA <13.2, which can spam the hot path. Update the
`launchHostFunc` logic in `hostfunc.cpp` to emit the fallback warning only once
(for example via a static once-flag or equivalent), while keeping the fallback
`cudaLaunchHostFunc` call unchanged. Also reconcile the
`enable_low_latency_host_dispatch` documentation in `llm_args.py` so it matches
the intended once-only or silent fallback behavior.

🪄 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: 338754ea-936b-4c6a-9dd1-b2587064eb6b

📥 Commits

Reviewing files that changed from the base of the PR and between 281acfd and f76068d.

📒 Files selected for processing (8)

cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp
tensorrt_llm/_torch/hostfunc.py
tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/llmapi/llm_args.py
tests/integration/defs/accuracy/test_llm_api_pytorch.py
tests/unittest/api_stability/references/llm.yaml
tests/unittest/bindings/test_hostfunc.py
tests/unittest/llmapi/test_llm_args.py

achartier · 2026-07-02T02:52:24Z

/bot run

tensorrt-cicd · 2026-07-02T03:02:39Z

PR_Github #57078 [ run ] triggered by Bot. Commit: 00a3bf1 Link to invocation

tensorrt-cicd · 2026-07-02T04:09:14Z

PR_Github #57078 [ run ] completed with state FAILURE. Commit: 00a3bf1
/LLM/main/L0_MergeRequest_PR pipeline #45869 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

achartier · 2026-07-02T14:36:30Z

/bot run

tensorrt-cicd · 2026-07-02T14:43:06Z

PR_Github #57197 [ run ] triggered by Bot. Commit: 00a3bf1 Link to invocation

tensorrt-cicd · 2026-07-02T16:30:12Z

PR_Github #57197 [ run ] completed with state SUCCESS. Commit: 00a3bf1
/LLM/main/L0_MergeRequest_PR pipeline #45970 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

achartier · 2026-07-02T16:32:06Z

/bot run

tensorrt-cicd · 2026-07-02T16:38:06Z

PR_Github #57220 [ run ] triggered by Bot. Commit: 00a3bf1 Link to invocation

tensorrt-cicd · 2026-07-02T19:17:46Z

PR_Github #57220 [ run ] completed with state SUCCESS. Commit: 00a3bf1
/LLM/main/L0_MergeRequest_PR pipeline #45991 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

achartier · 2026-07-02T22:21:04Z

/bot run

tensorrt-cicd · 2026-07-02T22:28:40Z

PR_Github #57265 [ run ] triggered by Bot. Commit: 00a3bf1 Link to invocation

tensorrt-cicd · 2026-07-02T22:29:27Z

PR_Github #57265 [ run ] completed with state FAILURE. Commit: 00a3bf1

Link to invocation

achartier · 2026-07-02T22:35:15Z

/bot run

tensorrt-cicd · 2026-07-02T22:40:54Z

PR_Github #57269 [ run ] triggered by Bot. Commit: 00a3bf1 Link to invocation

tensorrt-cicd · 2026-07-02T23:05:19Z

PR_Github #57269 [ run ] completed with state FAILURE. Commit: 00a3bf1
/LLM/main/L0_MergeRequest_PR pipeline #46030 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

yiqingy0 · 2026-07-03T03:47:11Z

/bot run

tensorrt-cicd · 2026-07-03T03:57:24Z

PR_Github #57339 [ run ] triggered by Bot. Commit: 00a3bf1 Link to invocation

The CUDA < 13.2 fallback warning was emitted on every launchHostFunc call, producing unbounded log spam on guided-decoding workloads. Guard with std::once_flag so it fires at most once per process lifetime. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>

…patch docs The field description said the CUDA < 13.2 fallback is silent, but the implementation logs a one-time warning. Update the docs to match. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>

achartier · 2026-07-05T23:52:33Z

/bot run

tensorrt-cicd · 2026-07-05T23:57:52Z

PR_Github #57638 [ run ] triggered by Bot. Commit: abc6479 Link to invocation

tensorrt-cicd · 2026-07-06T02:30:30Z

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

…lm_args create_py_executor_instance can receive a SimpleNamespace (e.g. in unit tests) instead of a real TorchLlmArgs. Use getattr with a False default so callers that don't set the field still work. Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>

achartier · 2026-07-06T03:03:47Z

/bot run

tensorrt-cicd · 2026-07-06T03:10:17Z

PR_Github #57673 [ run ] triggered by Bot. Commit: 9e90c00 Link to invocation

tensorrt-cicd · 2026-07-06T05:45:23Z

PR_Github #57673 [ run ] completed with state FAILURE. Commit: 9e90c00
/LLM/main/L0_MergeRequest_PR pipeline #46391 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

achartier · 2026-07-06T16:04:56Z

/bot run

tensorrt-cicd · 2026-07-06T16:10:42Z

PR_Github #57784 [ run ] triggered by Bot. Commit: 9e90c00 Link to invocation

tensorrt-cicd · 2026-07-06T19:55:55Z

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

achartier · 2026-07-06T20:51:43Z

/bot run

tensorrt-cicd · 2026-07-06T20:59:03Z

PR_Github #57821 [ run ] triggered by Bot. Commit: 9e90c00 Link to invocation

tensorrt-cicd · 2026-07-06T21:56:37Z

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

achartier · 2026-07-06T21:59:45Z

/bot run

tensorrt-cicd · 2026-07-06T22:07:03Z

PR_Github #57834 [ run ] triggered by Bot. Commit: 9e90c00 Link to invocation

tensorrt-cicd · 2026-07-07T00:06:35Z

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

achartier · 2026-07-07T01:10:44Z

/bot run

tensorrt-cicd · 2026-07-07T01:16:11Z

PR_Github #57873 [ run ] triggered by Bot. Commit: 9e90c00 Link to invocation

tensorrt-cicd · 2026-07-07T04:41:27Z

PR_Github #57873 [ run ] completed with state SUCCESS. Commit: 9e90c00
/LLM/main/L0_MergeRequest_PR pipeline #46569 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

syuoni

LGTM, thanks!

zhaoyangwang-nvidia

LGTM

…ing (NVIDIA#15863) Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>

achartier requested review from a team as code owners July 2, 2026 02:39

achartier requested review from HuiGao-NV, suyoggupta and zhenhuaw-me July 2, 2026 02:39

github-actions Bot assigned achartier Jul 2, 2026

achartier added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 2, 2026

coderabbitai Bot reviewed Jul 2, 2026

View reviewed changes

Comment thread cpp/tensorrt_llm/nanobind/runtime/hostfunc.cpp

xinhe-nv approved these changes Jul 3, 2026

View reviewed changes

achartier added 2 commits July 5, 2026 16:52

achartier force-pushed the feat/low-latency-host-dispatch branch from 00a3bf1 to abc6479 Compare July 5, 2026 23:52

achartier requested a review from syuoni July 7, 2026 19:49

suyoggupta removed their request for review July 7, 2026 20:12

syuoni approved these changes Jul 8, 2026

View reviewed changes

zhaoyangwang-nvidia approved these changes Jul 8, 2026

View reviewed changes

venkywonka approved these changes Jul 8, 2026

View reviewed changes

achartier merged commit 62b3cc7 into NVIDIA:main Jul 8, 2026
7 checks passed

achartier deleted the feat/low-latency-host-dispatch branch July 8, 2026 14:46

BrianLi23 pushed a commit to BrianLi23/TensorRT-LLM that referenced this pull request Jul 9, 2026

[None][feat] Add low-latency host task dispatch mode for guided decod…

269f5d3

…ing (NVIDIA#15863) Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>

Conversation

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

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 2, 2026

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

achartier commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

achartier commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

achartier commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

achartier commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

achartier commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

yiqingy0 commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

achartier commented Jul 5, 2026

Uh oh!

tensorrt-cicd commented Jul 5, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

achartier commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

achartier commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

achartier commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

achartier commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

achartier commented Jul 2, 2026 •

edited by coderabbitai Bot

Loading