[TRTLLM-14198][refactor] Make FlashInfer a hard dependency for the Torch sampler by zhaoyangwang-nvidia · Pull Request #16160 · NVIDIA/TensorRT-LLM

zhaoyangwang-nvidia · 2026-07-09T02:14:37Z

…sampler and spec

Convert the Torch sampler and one-model speculative decoding from "FlashInfer-optional with torch fallback" to a hard FlashInfer dependency. The failure is raised at sampler / speculative-worker construction time (resolve_sampling_backend, SpecWorkerBase) via a new require_flashinfer helper and a sentinel-based bind_flashinfer_sampling, so sampling_utils stays importable for torch-only consumers (AutoDeploy demollm, drafting greedy helpers).

Add require_flashinfer() in flashinfer_utils.py; IS_FLASHINFER_AVAILABLE semantics are unchanged (its ~46 legitimate fallback consumers are not touched).
ops/flashinfer.py: drop the import guard, bind flashinfer.sampling at construction via a sentinel; compiled bodies do no import work.
Remove the SimpleGroupedStrategySampler torch backend and the disable_flashinfer_sampling LlmArg + its plumbing; update the api_stability and usage golden files and the sampling/telemetry docs.
compute_probs_from_logits becomes FlashInfer-only (drop the CUDA C++ op and CPU vanilla branches).
speculative/interface.py: drop the availability guards, keep the flashinfer>=0.6.4 version gate and its version-gated branches.
Delete now-dead fallback kernels (forward_native_sampling, _apply_top_k_top_p, _random_sample, vanilla.compute_probs_from_logits_op); keep beam-search/greedy/rejection/_Fusions and the torch sample() path used by rejection-sampling draft probs and demollm.
Prune torch-vs-flashinfer parametrized sampler tests.

Summary by CodeRabbit

New Features
- PyTorch sampling now consistently uses the FlashInfer-backed backend for supported workflows.
Bug Fixes
- Removed an unsupported sampling toggle from user-facing configuration and APIs.
- Sampling-related behavior and logprob handling were aligned with the new default backend.
Documentation
- Updated developer and feature docs to reflect the simplified sampling configuration and current backend behavior.
Tests
- Updated sampling tests and API stability references to match the revised configuration surface.

Description

Make FlashInfer required for the PyTorch-backend Torch sampler, and delete
the now-dead torch fallback paths and their tests. When flashinfer is missing,
TorchSampler.__init__ now raises a clear ImportError instead of silently
running a slower/divergent torch path.

IS_FLASHINFER_AVAILABLE semantics are unchanged — its ~46 non-sampler
fallback consumers are not touched.

What changed

TorchSampler.__init__ hard-fails (ImportError) when flashinfer is missing,
before entering the CUDA-graph-captured sampling loop.
Flatten the sampler backend: remove SimpleGroupedStrategySampler and the
GroupedStrategySampler ABC; FlashInferGroupedStrategySampler is the only
implementation. The one-line factory is inlined and ops/interface.py deleted.
Remove the disable_flashinfer_sampling LlmArg and its plumbing (updates
api_stability / usage golden files and sampling & telemetry docs).
compute_probs_from_logits becomes FlashInfer-only; delete dead fallback
kernels (forward_native_sampling, _apply_top_k_top_p, _random_sample,
vanilla.compute_probs_from_logits_op).
Prune torch-vs-flashinfer parametrized sampler tests.

Intentionally kept (not fallbacks)

Torch sample() for two-model speculative rejection sampling and AutoDeploy
demollm; beam search (torch-only by design); _Fusions log-prob machinery;
and the flashinfer>=0.6.4 version gate in one-model spec decoding.

Test Coverage

Updated sampler unit tests (removed use_flashinfer / disable_flashinfer_sampling
parametrizations); updated api_stability references and usage golden manifest.

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.

coderabbitai · 2026-07-09T08:01:53Z

📝 Walkthrough

Walkthrough

This PR removes the disable_flashinfer_sampling configuration option throughout the codebase (public API, sampler construction, golden manifest, API stability reference, docs) and makes the Torch sampler hard-depend on FlashInfer, removing the prior fallback grouped-sampling backend and native probability/sampling implementations. Tests are updated accordingly.

Changes

FlashInfer-only sampling refactor

Layer / File(s)	Summary
Remove disable_flashinfer_sampling from public config and docs `tensorrt_llm/llmapi/llm_args.py`, `tensorrt_llm/_torch/pyexecutor/_util.py`, `tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`, `tensorrt_llm/usage/llm_args_golden_manifest.json`, `tests/unittest/api_stability/references/llm.yaml`, `docs/source/developer-guide/telemetry.md`, `docs/source/features/sampling.md`	Removes the `disable_flashinfer_sampling` field/parameter from `TorchLlmArgs`, sampler construction helpers, executor creation call, golden manifest, API stability reference, and documentation.
TorchSampler hard-depends on FlashInfer `tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`, `tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py`	`TorchSampler.__init__` raises `ImportError` if FlashInfer is unavailable and directly assigns `FlashInferGroupedStrategySampler`; the `disable_flashinfer_sampling` field is removed and docstrings/comments updated.
Remove fallback grouped-sampler and native probability implementations `tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py`, `tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py`	Removes `GroupedStrategySampler`/`SimpleGroupedStrategySampler`, makes `FlashInferGroupedStrategySampler` standalone, replaces conditional FlashInfer/native dispatch in `compute_probs_from_logits` and spec-decoding sampling with unconditional FlashInfer calls, and removes native top-k/top-p sampling functions.
Test updates for FlashInfer-only sampling `tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py`, `tests/unittest/_torch/sampler/test_beam_search.py`, `tests/unittest/_torch/sampler/test_logits_logprobs.py`, `tests/unittest/_torch/sampler/test_torch_sampler.py`	Removes `disable_flashinfer_sampling`/`use_flashinfer` parameterizations and updates helper functions and assertions to reflect the FlashInfer-only sampling backend.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#8581: Adds/threads a disable_flash_infer_sampling flag through create_torch_sampler_args/instantiate_sampler/create_py_executor, directly conflicting with this PR's removal of that same plumbing.

Suggested reviewers: dcampora, Funatiq, nv-guomingz, chzblych, ziyixiong-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 41.67% 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 includes a ticket plus lowercase type.
Description check	✅ Passed	The description includes summary, what changed, test coverage, and checklist items, so it largely matches the template.
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 `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 1888-1917: The `_sample_grouped_strategies` wrapper is accepting
`group_metadata` but not passing it through to `sample_grouped_strategies_orig`,
so the helper silently drops metadata. Update the delegation in
`_sample_grouped_strategies` to forward the `group_metadata` argument alongside
the existing `group_logit_indices`, `generator`, and `return_probs` parameters
so the instrumentation helper preserves real metadata.

🪄 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: ff798a0a-64d7-499d-9610-83539c3ce283

📥 Commits

Reviewing files that changed from the base of the PR and between 8ef7190 and ad0e8b5.

📒 Files selected for processing (16)

docs/source/developer-guide/telemetry.md
docs/source/features/sampling.md
tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py
tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.py
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/usage/llm_args_golden_manifest.json
tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
tests/unittest/_torch/sampler/test_beam_search.py
tests/unittest/_torch/sampler/test_logits_logprobs.py
tests/unittest/_torch/sampler/test_torch_sampler.py
tests/unittest/api_stability/references/llm.yaml

💤 Files with no reviewable changes (9)

tests/unittest/api_stability/references/llm.yaml
tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/usage/llm_args_golden_manifest.json
docs/source/developer-guide/telemetry.md
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py

zhaoyangwang-nvidia · 2026-07-09T08:21:59Z

/bot run --disable-fail-fast

…sampler and spec Convert the Torch sampler and one-model speculative decoding from "FlashInfer-optional with torch fallback" to a hard FlashInfer dependency. The failure is raised at sampler / speculative-worker construction time (TorchSampler.__init__, SpecWorkerBase) with a clear ImportError, so sampling_utils stays importable for torch-only consumers (AutoDeploy demollm, drafting greedy helpers). - IS_FLASHINFER_AVAILABLE semantics are unchanged (its ~46 legitimate fallback consumers are not touched). - ops/flashinfer.py: keep flashinfer.sampling behind an import guard so the module stays importable without flashinfer; ops are only reached after the hard-dependency check at construction time. - TorchSampler.__init__ enforces the hard dependency once (out of the CUDA-graph-captured loop) before selecting the grouped sampler class. - Flatten the single-implementation grouped sampler: drop the GroupedStrategySampler ABC and keep only FlashInferGroupedStrategySampler. - Remove the SimpleGroupedStrategySampler torch backend and the disable_flashinfer_sampling LlmArg + its plumbing; update the api_stability and usage golden files and the sampling/telemetry docs. - compute_probs_from_logits becomes FlashInfer-only (drop the CUDA C++ op and CPU vanilla branches). - speculative/interface.py: drop the availability guards, keep the flashinfer>=0.6.4 version gate and its version-gated branches. - Delete now-dead fallback kernels (forward_native_sampling, _apply_top_k_top_p, _random_sample, vanilla.compute_probs_from_logits_op); keep beam-search/greedy/rejection/_Fusions and the torch sample() path used by rejection-sampling draft probs and demollm. - Prune torch-vs-flashinfer parametrized sampler tests. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

tensorrt-cicd · 2026-07-09T08:28:46Z

PR_Github #58426 [ run ] triggered by Bot. Commit: 5d97c03 Link to invocation

zhaoyangwang-nvidia · 2026-07-09T08:33:04Z

/bot kill

zhaoyangwang-nvidia · 2026-07-09T08:35:26Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-09T08:38:54Z

PR_Github #58427 [ kill ] triggered by Bot. Commit: 5d97c03 Link to invocation

tensorrt-cicd · 2026-07-09T08:41:17Z

PR_Github #58429 [ run ] triggered by Bot. Commit: 5d97c03 Link to invocation

tensorrt-cicd · 2026-07-09T08:41:22Z

PR_Github #58427 [ kill ] completed with state ABORTED. Commit: 5d97c03

Link to invocation

tensorrt-cicd · 2026-07-09T08:41:46Z

PR_Github #58426 [ run ] completed with state ABORTED. Commit: 5d97c03

Link to invocation

…vailable One-model speculative sampling (sampling_batch_spec_dec_one_model) now calls flashinfer unconditionally after the torch fallback was removed. SpecWorkerBase previously only set use_flashinfer=False when flashinfer was missing or <0.6.4, so the worker could initialize successfully and then crash mid-inference on the first non-greedy sampling step. Raise ImportError at SpecWorkerBase.__init__ when flashinfer is unavailable or below 0.6.4, and drop the now-redundant use_flashinfer flag (always True past the check). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

QiJune

LGTM

zhaoyangwang-nvidia · 2026-07-09T09:08:19Z

Hi @sunnyqgg Added some spec-side changes here touching the spec ↔ sampler integration — could you help review? Thanks!

tensorrt-cicd · 2026-07-09T10:30:48Z

PR_Github #58429 [ run ] completed with state SUCCESS. Commit: 5d97c03
/LLM/main/L0_MergeRequest_PR pipeline #47044 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

sampler.py lives in tensorrt_llm/_torch/pyexecutor/sampler/, so the relative import 'from ..flashinfer_utils' resolved to the wrong package and broke 'import tensorrt_llm' during pytest collection. Use the absolute import 'from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE', matching the sibling ops/flashinfer.py. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

zhaoyangwang-nvidia · 2026-07-09T11:01:26Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-09T11:07:15Z

PR_Github #58454 [ run ] triggered by Bot. Commit: 721eb3a Link to invocation

github-actions Bot assigned zhaoyangwang-nvidia Jul 9, 2026

zhaoyangwang-nvidia added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 9, 2026

zhaoyangwang-nvidia force-pushed the remove-non-flashinfer-support branch 4 times, most recently from 78f0aed to 0fe449d Compare July 9, 2026 04:06

zhaoyangwang-nvidia changed the title ~~[TRTLLM-13212][refactor] Make FlashInfer a hard dependency for the Torch sampler and speculative decoding~~ [TRTLLM-14198][refactor] Make FlashInfer a hard dependency for the Torch sampler and speculative decoding Jul 9, 2026

zhaoyangwang-nvidia force-pushed the remove-non-flashinfer-support branch 5 times, most recently from 4093d14 to ad0e8b5 Compare July 9, 2026 07:48

zhaoyangwang-nvidia changed the title ~~[TRTLLM-14198][refactor] Make FlashInfer a hard dependency for the Torch sampler and speculative decoding~~ [TRTLLM-14198][refactor] Make FlashInfer a hard dependency for the Torch sampler Jul 9, 2026

zhaoyangwang-nvidia marked this pull request as ready for review July 9, 2026 07:51

zhaoyangwang-nvidia requested review from a team as code owners July 9, 2026 07:51

zhaoyangwang-nvidia requested review from QiJune and arysef July 9, 2026 07:51

coderabbitai Bot reviewed Jul 9, 2026

View reviewed changes

Comment thread tests/unittest/_torch/sampler/test_torch_sampler.py

zhaoyangwang-nvidia force-pushed the remove-non-flashinfer-support branch from ad0e8b5 to f979046 Compare July 9, 2026 08:05

zhaoyangwang-nvidia force-pushed the remove-non-flashinfer-support branch from f979046 to 5d97c03 Compare July 9, 2026 08:22

QiJune reviewed Jul 9, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py

zhaoyangwang-nvidia requested a review from a team as a code owner July 9, 2026 08:53

zhaoyangwang-nvidia requested a review from sunnyqgg July 9, 2026 08:54

QiJune approved these changes Jul 9, 2026

View reviewed changes

Uh oh!

Conversation

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

Uh oh!

Summary by CodeRabbit

Description

What changed

Intentionally kept (not fallbacks)

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 9, 2026

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zhaoyangwang-nvidia commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

zhaoyangwang-nvidia commented Jul 9, 2026

Uh oh!

zhaoyangwang-nvidia commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

Uh oh!

QiJune left a comment

Choose a reason for hiding this comment

Uh oh!

zhaoyangwang-nvidia commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

zhaoyangwang-nvidia commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zhaoyangwang-nvidia commented Jul 9, 2026 •

edited by coderabbitai Bot

Loading