[None][fix] Make FlashInfer sampling op wrappers opaque to Dynamo by qiaoxj07 · Pull Request #16732 · NVIDIA/TensorRT-LLM

qiaoxj07 · 2026-07-22T12:46:39Z

Description

The @torch.compiled spec-decode sampling functions in sampling_utils.py (compute_probs_from_logits, sampling_batch_spec_dec_one_model, sampling_batch_spec_dec_one_model_for_rejection) inline the plain-Python FlashInfer op wrappers in sampler/ops/flashinfer.py, and nothing below those wrappers is opaque to Dynamo:

FlashInfer's sampling kernels are reached through get_sampling_module(), a functools.cache-wrapped lazy JIT bootstrap.
FlashInfer's own register_custom_op/register_fake_op are deliberate no-ops (flashinfer/utils.py, citing torch.library.custom_op overhead), so the kernel entry points are plain closures over an FFI module.

Dynamo ignores functools.cache/lru_cache wrappers and traces the wrapped body directly (torch/_dynamo/variables/functions.py, WrapperUserFunctionVariable). Tracing the JIT bootstrap hits untraceable builtins (posix.stat, datetime.now, _thread.allocate_lock, sys._getframe, plus nested lru_cached helpers), each emitting a warn-once UserWarning and a permanent graph break. Because the cache wrapper is bypassed inside the compiled artifact, the break-site calls (file stat, module re-load) re-execute eagerly on every invocation of the compiled sampler, defeating the bootstrap's cache fast path. In any MTP/spec-decode serving log this shows up as a burst of Dynamo does not know how to trace the builtin ... / functools.lru_cache-wrapped function warnings per rank.

Fix: decorate every wrapper in sampler/ops/flashinfer.py with @torch.compiler.disable. Each wrapper becomes a single clean graph break, the warnings disappear, and FlashInfer's functools.cache fast path is restored. The kernels already executed eagerly after the old graph breaks, so numerics and CUDA-graph behavior are unchanged. Registering the wrappers as torch.library.custom_op instead is not viable: the merged dual-mode randomness signature (torch.Generator + Union[int, torch.Tensor] seed/offset, intentional as of #16365) uses parameter types unsupported by the custom-op schema. This follows the existing @torch.compiler.disable pattern used for the cute_dsl fmha wrappers.

Test Coverage

A/B on GB300 (torch 2.12 / flashinfer 0.6.15), calling sampling_batch_spec_dec_one_model and sampling_batch_spec_dec_one_model_for_rejection under torch.compile with fixed seed/offset tensors — identical script, stock image vs. this patch bind-mounted:

Metric	Baseline	Patched
`Dynamo does not know how to trace the builtin` warnings	4	0
`functools.lru_cache-wrapped` warnings	6	0
Sampled tokens (both entry points)	—	bit-identical to baseline
Filtered-probs checksum	8.000000142	8.000000142
Steady-state host time per call	280.5 us	195.1 us
First-call compile time	62.3 s + 2.1 s	45.4 s + 0.2 s

Existing tests covering these paths: tests/unittest/_torch/sampler/ and the speculative one-model sampling 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.

Dev Engineer Review

Expanded tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py module docstring to document the rationale: preventing TorchDynamo from tracing FlashInfer’s lazy JIT bootstrap/cache initialization during Python wrapper execution.
Added a private signature-preserving _compiler_disable helper that applies torch.compiler.disable(fn) while preserving the wrapped callable type via cast.
Decorated the module-level FlashInfer sampling/probability wrapper ops with @_compiler_disable (without changing their function signatures or internal FlashInfer call logic):
- top_k_top_p_sampling_from_logits_op
- sampling_from_probs_op
- top_k_sampling_from_probs_op
- top_p_sampling_from_probs_op
- softmax_op
- top_k_mask_logits_op
- top_p_renorm_probs_op
- compute_probs_from_logits_op
No other exported/public declarations or operator logic were changed; this is intended to reduce Dynamo-related warnings while keeping numerics and CUDA-graph behavior stable.

QA Engineer Review

No test changes.

qiaoxj07 · 2026-07-22T12:47:06Z

/bot run

coderabbitai · 2026-07-22T12:53:05Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

FlashInfer sampling and probability-computation wrappers are decorated with @torch.compiler.disable, and the module documentation explains that this prevents Dynamo from tracing FlashInfer bootstrap and caching behavior.

Changes

FlashInfer compiler control

Layer / File(s)	Summary
Disable tracing for FlashInfer wrappers `tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py`	Adds a signature-preserving compiler-disable helper and applies it to FlashInfer sampling, probability transformation, and probability computation ops without changing their signatures or internal calls.

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

Suggested reviewers: asfiyab-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 62.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change and follows the required [None][fix] format.
Description check	✅ Passed	The PR description includes Description, Test Coverage, and PR Checklist, and it clearly explains the issue and fix.
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.}

tensorrt-cicd · 2026-07-22T12:53:34Z

PR_Github #60977 [ run ] triggered by Bot. Commit: 7ccf02b Link to invocation

tensorrt-cicd · 2026-07-22T13:23:36Z

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

qiaoxj07 · 2026-07-22T16:48:40Z

/bot run

tensorrt-cicd · 2026-07-22T16:57:43Z

PR_Github #61034 [ run ] triggered by Bot. Commit: 71cab53 Link to invocation

tensorrt-cicd · 2026-07-22T18:42:24Z

PR_Github #61034 [ run ] completed with state FAILURE. Commit: 71cab53
/LLM/main/L0_MergeRequest_PR pipeline #49289 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

qiaoxj07 · 2026-07-22T18:53:56Z

/bot run

tensorrt-cicd · 2026-07-22T19:00:21Z

PR_Github #61066 [ run ] triggered by Bot. Commit: 80516ad Link to invocation

tensorrt-cicd · 2026-07-22T19:13:31Z

PR_Github #61066 [ run ] completed with state FAILURE. Commit: 80516ad
/LLM/main/L0_MergeRequest_PR pipeline #49321 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

qiaoxj07 · 2026-07-22T19:16:56Z

/bot run

tensorrt-cicd · 2026-07-22T19:23:18Z

PR_Github #61072 [ run ] triggered by Bot. Commit: 4c47731 Link to invocation

tensorrt-cicd · 2026-07-22T21:00:36Z

PR_Github #61072 [ run ] completed with state FAILURE. Commit: 4c47731
/LLM/main/L0_MergeRequest_PR pipeline #49327 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

qiaoxj07 · 2026-07-22T21:04:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-22T21:15:01Z

PR_Github #61093 [ run ] triggered by Bot. Commit: 4c47731 Link to invocation

tensorrt-cicd · 2026-07-23T00:24:29Z

PR_Github #61093 [ run ] completed with state SUCCESS. Commit: 4c47731
/LLM/main/L0_MergeRequest_PR pipeline #49346 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

qiaoxj07 · 2026-07-23T00:34:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T00:41:09Z

PR_Github #61133 [ run ] triggered by Bot. Commit: cc2555c Link to invocation

The torch.compile'd spec-decode sampling functions (sampling_utils.py) inline these plain-Python wrappers, and nothing below them is opaque to Dynamo: flashinfer's kernels are reached via get_sampling_module(), a functools.cache'd lazy JIT bootstrap whose custom-op registration is a no-op in flashinfer. Dynamo bypasses the functools.cache wrapper and traces the bootstrap body, so every untraceable builtin in it (posix.stat, datetime.now, lock allocation, lru_cache'd get_cuda_version) emits a warn-once and becomes a permanent graph break whose call re-executes eagerly on every invocation of the compiled sampler, defeating the bootstrap's cache fast path. Decorate all wrappers with torch.compiler.disable: each becomes a single clean graph break, the Dynamo warnings disappear, and the functools.cache fast path is restored. Kernels already executed eagerly after the old graph breaks, so numerics and CUDA-graph behavior are unchanged. The decorator is applied through a typed shim (_compiler_disable) because torch.compiler.disable ships untyped: used directly it fails mypy strict (untyped-decorator on every wrapper) and cascades no-any-return errors into sampling_utils.py. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

qiaoxj07 · 2026-07-23T06:45:30Z

/bot kill

qiaoxj07 · 2026-07-23T06:45:38Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T06:51:50Z

PR_Github #61220 [ kill ] triggered by Bot. Commit: de8c7c8 Link to invocation

tensorrt-cicd · 2026-07-23T06:53:38Z

PR_Github #61133 [ run ] completed with state ABORTED. Commit: cc2555c

Link to invocation

tensorrt-cicd · 2026-07-23T06:53:43Z

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

Link to invocation

tensorrt-cicd · 2026-07-23T06:59:51Z

PR_Github #61221 [ run ] triggered by Bot. Commit: de8c7c8 Link to invocation

tensorrt-cicd · 2026-07-23T11:05:14Z

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

qiaoxj07 · 2026-07-23T11:20:55Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-23T11:26:53Z

PR_Github #61283 [ run ] triggered by Bot. Commit: de8c7c8 Link to invocation

tensorrt-cicd · 2026-07-23T13:28:30Z

PR_Github #61283 [ run ] completed with state SUCCESS. Commit: de8c7c8
/LLM/main/L0_MergeRequest_PR pipeline #49516 completed with status: 'SUCCESS'

CI Report

Link to invocation

…nnahz/dep-1083-port-flashinfer-stable-va-lifecycle-for-native-all-reduce * 'main' of https://github.com/NVIDIA/TensorRT-LLM: (83 commits) [None][feat] Bind SourceIdentity to checkpoint artifacts (NVIDIA#16159) [None][infra] Waive 4 failed cases for main in pre-merge 49550 (NVIDIA#16798) [None][fix] Make FlashInfer sampling op wrappers opaque to Dynamo (NVIDIA#16732) [None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e (NVIDIA#16420) [None][feat] Default GLM-5 to the Python KV-cache transceiver (NVIDIA#16524) [None][chore] Add NVTX ranges to per-iteration ADP sync points in PyExecutor (NVIDIA#16422) [https://nvbugs/6426850][test] Unwaive Qwen3.5 397B NVFP4 ADP4 TRTLLM test (NVIDIA#16348) [https://nvbugs/6445456][fix] Restore inplace ops for functionalization v2 (NVIDIA#16410) [None][infra] Waive 1 failed cases for main in pre-merge 49229 (NVIDIA#16786) [None][fix] Load DeepSeek V4 mixed-precision NVFP4 checkpoints (NVIDIA#16433) [None][feat] ADP conversation router: configurable least-queued placement for new conversations (NVIDIA#16294) [None][infra] Waive 1 failed cases for main in pre-merge 49424 (NVIDIA#16780) [None][infra] Waive 1 failed cases for main in pre-merge 49424 (NVIDIA#16781) [TRTLLMINF-188][infra] Require approval for PerfSanity wildcard runs (NVIDIA#16777) [TRTLLM-13969][feat] Support MiniMax M3 for Disaggregated Serving (NVIDIA#16017) [NVIDIA#11932][fix] Filter CUTLASS MoE GEMM tile configs by device shared memory on SM121 (NVIDIA#12704) [None][chore] Remove attention backend test waivers (NVIDIA#16723) [TRTLLM-14540][perf] Skip fp32 state round-trip in FlashInfer GDN pre… (NVIDIA#16703) [TRTLLM-13694][feat] Add IBDB recipe provenance and refresh configs (NVIDIA#16254) [None][infra] Preview/bump/main (NVIDIA#16758) ... Signed-off-by: Hannah Zhang <hannahz@nvidia.com>

…IDIA#16732) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

qiaoxj07 requested a review from a team as a code owner July 22, 2026 12:46

qiaoxj07 requested review from JunyiXu-nv and asfiyab-nvidia July 22, 2026 12:46

github-actions Bot assigned qiaoxj07 Jul 22, 2026

qiaoxj07 force-pushed the fix/flashinfer-sampling-dynamo-opaque branch from 7ccf02b to 71cab53 Compare July 22, 2026 16:48

qiaoxj07 force-pushed the fix/flashinfer-sampling-dynamo-opaque branch from 71cab53 to 80516ad Compare July 22, 2026 18:53

qiaoxj07 force-pushed the fix/flashinfer-sampling-dynamo-opaque branch from 80516ad to 4c47731 Compare July 22, 2026 19:16

qiaoxj07 force-pushed the fix/flashinfer-sampling-dynamo-opaque branch from 4c47731 to cc2555c Compare July 23, 2026 00:34

QiJune requested review from lori-ren and zhaoyangwang-nvidia July 23, 2026 03:22

zhaoyangwang-nvidia approved these changes Jul 23, 2026

View reviewed changes

qiaoxj07 force-pushed the fix/flashinfer-sampling-dynamo-opaque branch from cc2555c to de8c7c8 Compare July 23, 2026 06:45

qiaoxj07 merged commit 77774d9 into NVIDIA:main Jul 23, 2026
7 checks passed

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

[None][fix] Make FlashInfer sampling op wrappers opaque to Dynamo (NV…

8e35afa

…IDIA#16732) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

Uh oh!

Conversation

qiaoxj07 commented Jul 22, 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

Dev Engineer Review

QA Engineer Review

Uh oh!

qiaoxj07 commented Jul 22, 2026

Uh oh!

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

Uh oh!

Reviews paused

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

qiaoxj07 commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

qiaoxj07 commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

qiaoxj07 commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

qiaoxj07 commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

qiaoxj07 commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

qiaoxj07 commented Jul 23, 2026

Uh oh!

qiaoxj07 commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

qiaoxj07 commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

tensorrt-cicd commented Jul 23, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

qiaoxj07 commented Jul 22, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 22, 2026 •

edited

Loading