[TRTLLM-15017][chore] Unify one-model speculative decoding samplers into SpecSampler

[zhaoyangwang-nvidia]

Description

One-model speculative decoding had five sampler classes: SpecSamplerBase plus MTPSampler, Eagle3OneModelSampler, DraftTargetOneModelSampler and SASampler. None of them overrode sample_async, update_requests, validate_request or Store — they differed only in how they sized four buffers, via three hooks (_get_max_tokens, _get_max_new_tokens, _get_draft_tokens_storage_size).

Those sizes all derive from TorchSampler.Args alone. args.max_total_draft_tokens is set to spec_config.tokens_per_gen_step - 1, which equals the draft length each mode computed for itself — including PARD, where it is 2K-1 rather than the tree size. The per-mode subclasses therefore carried no information args did not already have.

Accepted-token buffer (third commit)

new_tokens holds only what a step accepts -- the golden token the target always
accepts, plus the drafts that matched -- so it is sized to max_draft_len + 1
rather than the wire width max_total_draft_tokens + 1. Modes that carry more
draft slots than they can accept shrink; the six linear modes are unchanged
because max_total_draft_tokens == max_draft_len there:

mode	before	after
Eagle3 dynamic tree (K=6, topK=10)	61	7
MTP dynamic tree (K=1, topK=4)	5	2
PARD (K=6, T=2K-1)	12	7
Eagle3 static tree (K=2, T=6)	7	7 (unchanged, see below)
linear modes (K=6, T=6)	7	7

Three of those over-allocations predate this series. Only Eagle3 dynamic tree was
tight, via the Eagle3OneModelSampler override the first commit dropped with the
subclass, so this generalizes that bound instead of reinstating a special case.

Static tree workaround. The deprecated eagle_choices static tree keeps the
wire width. Its one-model drafter never builds the tree: _forward_draft_loop is
a linear loop over runtime_draft_len, which for a non-linear tree is
max_total_draft_tokens, so a step can accept up to max_total_draft_tokens + 1
tokens while max_draft_len only describes a tree depth nothing constructs.
Tree-aware acceptance exists solely in the two-model TorchSampler path. Both
that path and eagle_choices are slated for removal in release 1.4, and this
workaround goes with them. get_spec_decoder carries the override so the special
case sits at the one place that knows the config.

update_requests asserts the bound, so a mode that ever exceeds it fails with a
clear message instead of an opaque IndexError from add_token indexing a plain
host-side list.

Test Coverage

• tests/unittest/_torch/speculative/ — 251 passed, 25 skipped, covering MTP, PARD, SA, DraftTarget, DFlash, DSpark, dynamic draft length, and tree sampling/verification. Compared against a pre-change baseline: no test regressed, and the new assertion never fired.
• Two DeepSeek NVFP4 eagle3 tests (test_deepseek_eagle3, test_deepseek_mla_eagle3) fail identically with and without these changes on an sm90-only build, since CUTLASS FP4 GEMM requires sm100. Confirmed by re-running them on the pre-change tree.
• tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py — 5 passed, exercising the one external consumer of the renamed sampler.

Baseline showed one failure, test_advanced_sampling_mode.py::test_no_topk_matches_full[0.9], which is pre-existing and flaky (3 failures in 8 consecutive runs on H200) with no call path to the changed code.

PR Checklist

• Commit message follows the [JIRA][type] description format
• Commits are signed off (DCO)
• Pre-commit hooks pass
• Tests added/updated as needed — the first two commits are behavior-preserving and covered by existing tests; the third changes buffer sizing and is covered by the existing suites plus the benchmark runs listed above

🤖 Generated with Claude Code

Dev Engineer Review

• Consolidated one-model speculative decoding into SpecSampler.
• Derived buffer sizes from TorchSampler.Args.
• Sized the accepted-token buffer to max_draft_len + 1.
• Added validation for requests that exceed the accepted-token buffer.
• Updated get_spec_decoder to use TorchSampler for two-model modes and SpecSampler for one-model modes.
• Removed obsolete sampler classes, aliases, sizing hooks, and dummy-token hooks.
• Updated imports, comments, docstrings, and the AutoDeploy call site.
• No configuration or test-list files changed.
• No additional correctness, performance, error-handling, or API consistency issues were identified from the provided changes.

QA Engineer Review

• Modified test file: tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py.
• Updated the SSM replay smoke test docstring.
• No test functions were added, modified, or removed.
• Test-list coverage is not applicable because the change only updates documentation.
• Reported speculative decoding and AutoDeploy tests passed. Two DeepSeek NVFP4 Eagle3 failures remain unchanged and are hardware-specific.
• Verdict: sufficient.

[coderabbitai]

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

Changes

Speculative sampler unification

Layer / File(s)	Summary
Unified SpecSampler implementation tensorrt_llm/_torch/speculative/spec_sampler_base.py	SpecSampler derives buffer dimensions from sampler arguments, validates accepted paths, adapts runtime widths, and supplies context draft tokens.
Mode-specific sampler removal tensorrt_llm/_torch/speculative/{draft_target,eagle3,mtp,sa_worker}.py	Mode-specific sampler classes and related dependencies are removed.
Decoder selection and exports tensorrt_llm/_torch/speculative/{__init__,utils}.py, tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py	One-model modes use SpecSampler, while two-model Eagle3 and MTP-Eagle modes use TorchSampler. The package exports SpecSampler.
Sampler reference updates tensorrt_llm/_torch/{auto_deploy/models/custom/modeling_eagle.py,pyexecutor/sampler/top_p_decay.py}, tensorrt_llm/_torch/speculative/{drafter,interface}.py, tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py	Comments, documentation, and the smoke-test description now reference SpecSampler.

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

Suggested reviewers: allisonlim-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.29% 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
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.
Description check	✅ Passed	The description explains the consolidation, buffer-sizing changes, test coverage, known failures, and checklist status.
Title check	✅ Passed	The title clearly and concisely describes the unification of one-model speculative decoding samplers into SpecSampler.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.