[Hang-get]

Dev Engineer Review

SpecSampler.update_requests now filters completed requests before conversion. It converts only the host-tensor rows for active requests' sequence slots (new_tokens, new_tokens_lens, next_draft_tokens). The loop reads those rows by compact batch position (row). token_windows[request_index] and draft_lens[req_idx] still use request-position indexing. This matches the diff.

Review points:

• add_token is removed from the import and the loop. Each token now comes from new_tokens_rows[row][i] and req.add_new_token(new_token, beam_idx). The diff does not show the old helper's body. Confirm that add_token only appended the token, so that behavior matches.
• _request_common_handling now takes one request's draft-token list instead of the batch-shaped list. It no longer indexes by py_seq_slot. Any other caller of this method must pass a per-request list. The diff shows only this call site.
• The num_new_tokens <= max_accepted_path_len assertion still runs before token reads. Its behavior is unchanged.
• No public API or configuration changes appear. The import removal affects only the internal pyexecutor.sampler import.

QA Engineer Review

Test files changed:

• Added: tests/unittest/_torch/speculative/hw_agnostic/test_spec_sampler_update_requests.py, with one parameterized test, test_update_requests_reads_each_request_slot_row. It is marked pytest.mark.cpu_only.
• Parameter values: verify_lens is None or [3, 2, 1, 2].

The test covers:

• Non-contiguous sequence slots (11, 0, 7) in a 16-slot store.
• A completed request placed between active requests.
• Accepted-token lists, py_num_accepted_draft_tokens, py_num_draft_tokens_verified, py_rewind_len, py_draft_tokens, and py_decoding_iter.
• Unchanged state for the completed request.

Test-list membership: a git grep of tests/integration/test_lists/ at HEAD for the test filename returned no matches. The test does not appear in test-db/ CI lists or qa/ manual-QA lists. This is a unit test, so these lists do not apply.

Coverage verdict: sufficient for row selection and request updates. It does not cover GPU execution, end-to-end serving, or the stop-criteria paths, because handle_stop_criteria is stubbed to return False.

Current-head CI and test results are not available.

Per-File QA Perspective

• tensorrt_llm/_torch/speculative/spec_sampler_base.py: Verify updates for active requests whose sequence slots are scattered, including batches that contain completed requests. Verify that verification-window checks, accepted-token bounds, and stop handling still behave correctly.
• tests/unittest/_torch/speculative/hw_agnostic/test_spec_sampler_update_requests.py: Adds a CPU-only unit test for row selection and request updates, covering both verification-length paths. It is not in the integration CI or manual-QA lists.

Description

SpecSampler.update_requests runs once per step for every one-model speculative decoding mode. After sampler_event.synchronize() it converted the host copies of the whole sampler store to nested Python lists: new_tokens ([max_draft_len + 1, max_num_sequences, 1]), new_tokens_lens and next_draft_tokens. The store has one row per preallocated sequence slot, but the loop only reads the rows of the batch's requests, so the conversion cost depended on max_num_sequences rather than on the batch.

This change indexes the three host tensors with the slots of the batch's unfinished requests, converts only those rows, and reads them by row position in the validation and update loops. token_windows and draft_lens are still indexed by request position. _request_common_handling now receives the request's own draft row. Device work, the D2H copies and the request updates are unchanged.

Closes #22016.

Benchmark

CPU only, main 37e6e677 against this change (spec_sampler_base.py is identical on the current base 70fe2b2a6). Both versions of update_requests ran on identical host tensors, with request stubs implementing add_new_token, get_num_tokens, get_tokens and finish_by, and the real handle_stop_criteria; native bindings were stubbed for the import. Intel Xeon Platinum 8352Y, CPython 3.12.3, PyTorch 2.15.0+cpu, one pinned core; 20 warmups, then nine rounds of 15-50 timed calls, median of the round medians.

Slots	Active requests	Draft tokens	main	This PR	Speedup
256	1	3	129.6 us	27.1 us	4.8x
256	256	3	651.4 us	584.1 us	1.1x
512	1	3	256.1 us	27.1 us	9.5x
512	16	3	289.1 us	60.8 us	4.8x
2048	1	3	1086.3 us	27.7 us	39.2x
2048	8	3	1108.4 us	45.0 us	24.6x
2048	64	3	1249.0 us	159.0 us	7.9x
2048	2048	3	5748.5 us	4994.5 us	1.2x
2048	8	7	2223.6 us	52.2 us	42.6x

2048 is the default max_batch_size; most DeepSeek-R1-0528 MTP configs under examples/configs/database use 512, including all concurrency-1 configs. A full store is not slower, because each request's tokens become one flat list instead of one single-element list per token.

Test Coverage

• New tests/unittest/_torch/speculative/hw_agnostic/test_spec_sampler_update_requests.py (cpu_only): requests on scattered slots of a 16-slot store, with a finished request between them so that row positions differ from request positions, for uniform and ragged verification windows. It checks the appended tokens, the accepted/verified/rewind counts and the next draft tokens against each request's slot row. 2 passed, on this branch and on main.
• Existing tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py, which drives update_requests through sample_async (verify windows, pending steps, slot reuse): 85 passed.
• Randomized differential check against main, not included in the PR: 20,000 random states (finished requests, ragged and uniform windows, runtime_draft_len above and below the store width, end-id, length and stop-word finishes, invalid windows) gave identical request updates and identical exceptions.
• pre-commit run --files on both files: all hooks passed, including Ruff, ruff-format, codespell, the test-list AST check and Pyrefly (no diagnostics).

Local test runs used --noconftest with an in-process stub for the native bindings; this machine has no usable GPU or TensorRT-LLM build. The l0_cpu list already includes unittest/_torch/speculative/hw_agnostic. Not run: GPU tests and end-to-end serving benchmarks.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

AI assistance (Cursor) was used for this change.

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

SpecSampler.update_requests now selects host token data for active requests by sequence slot. It appends accepted tokens and passes each request’s draft-token row to common handling. A CPU-only test checks non-contiguous slots, an interleaved completed request, and both verification-length modes.

Changes

Speculative sampler request updates

Layer / File(s)	Summary
Select rows and update active requests tensorrt_llm/_torch/speculative/spec_sampler_base.py, tests/unittest/_torch/speculative/hw_agnostic/test_spec_sampler_update_requests.py	update_requests gathers host data for active requests by sequence slot, appends accepted tokens directly, and passes each request’s draft-token row to common handling. The CPU-only test checks active and completed requests, request state updates, and both verification-length modes.

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

Change: Refactor

Possibly related PRs

• NVIDIA/TensorRT-LLM#20524: Adds an optional SamplerStoreWriter path in SpecSampler; this PR also changes host-side sampler row selection and request updates.

Suggested reviewers: zhaoyuanh-nvidia

Merge Risk: Minimal · up to da824

This change reduces host-side conversion work in the speculative sampler's request updates. No concrete merge-blocking risk was identified, and a new CPU test covers scattered slots and completed requests.

Pre-merge checks: 4

Check name	Status	Explanation
Title check		The title follows the repository template with an issue reference [#22016] and a [perf] type. It names the main change: SpecSampler.update_requests now converts only the batch's slot rows instead of t…
Description check		The description covers the required areas from the template. It explains the problem and the fix in the Description section. It lists the relevant tests under Test Coverage, and it includes benchmark …
Linked Issues check		[#22016] asks SpecSampler.update_requests to convert only the host-tensor rows for the batch's unfinished requests, and to keep device work, D2H copies, and request updates unchanged. The reviewed dif…
Out of Scope Changes check		The changes in spec_sampler_base.py serve #22016. They touch only the conversion, row selection, and per-request reads in update_requests, plus the _request_common_handling signature that receives the…

Full details: Title check

Explanation

The title follows the repository template with an issue reference [#22016] and a [perf] type. It names the main change: SpecSampler.update_requests now converts only the batch's slot rows instead of the whole sampler store. The title is specific and matches the raw summary and PR objectives.

Full details: Description check

Explanation

The description covers the required areas from the template. It explains the problem and the fix in the Description section. It lists the relevant tests under Test Coverage, and it includes benchmark data and a checklist. It does not use the exact template comment text, and the PR Checklist shows only one box, which the author filled in. These gaps are minor, so the description is mostly complete.

Full details: Linked Issues check

Explanation

[#22016] asks SpecSampler.update_requests to convert only the host-tensor rows for the batch's unfinished requests, and to keep device work, D2H copies, and request updates unchanged. The reviewed diff (merge base efec60e to head da82467) does this in spec_sampler_base.py. After sampler_event.synchronize(), it builds active_requests from state.requests, skipping GENERATION_COMPLETE. It indexes new_tokens[:, active_slots, beam_idx].T, new_tokens_lens[active_slots], and next_draft_tokens[active_slots] before .tolist(). The loops read these by compact row position, while token_windows and state.draft_lens stay indexed by request position. _request_common_handling now receives the request's own draft row. The new CPU-only test covers scattered slots, a finished request between active requests, and uniform and ragged windows. The PR description reports the test passing. The rebase onto main does not change this file's logic. No coding requirement from #22016 remains unmet. The GPU path and end-to-end behavior were not run, so their results are not established by this review.

Full details: Out of Scope Changes check

Explanation

The changes in spec_sampler_base.py serve #22016. They touch only the conversion, row selection, and per-request reads in update_requests, plus the _request_common_handling signature that receives the row directly. The removed add_token helper is replaced by direct reads and req.add_new_token, which is part of the same row-based change. The new test file tests/unittest/_torch/speculative/hw_agnostic/test_spec_sampler_update_requests.py verifies that row-to-request mapping, so it supports the issue. The rebase onto main adds no separate unrelated change in the reviewed evidence. Residual uncertainty: this review did not search other modules for callers of the changed _request_common_handling signature or imports of the removed add_token name. Those callers are not shown in the reviewed evidence.

Finishing Touches

Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/TensorRT-LLM
• Branch: main
• File: .coderabbit.yaml
• Commit: efec60e
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 6caab302-0c85-4d1b-9050-7b640244bd98

Commits

Reviewing files that changed from the base of the PR and between afb753f and da82467.

Files selected for processing: 1

• tensorrt_llm/_torch/speculative/spec_sampler_base.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

[JunyiXu-nv]

Runtime part LGTM. The compact rows match the per-slot reads (new_tokens[step][slot][beam], new_tokens_lens and next_draft_tokens by slot) on random shapes and slot subsets, add_token did only the read plus add_new_token, and _request_common_handling has no other caller. The new cpu_only test runs via l0_cpu.yml's speculative/hw_agnostic entry.

[hanjingtian]

Could you please rebase this branch onto the latest upstream main to include the CI hang fix in #22101 (6edbe383497a)? The current branch still contains the regression that can cause pipelines to hang for an extended period and worsen CI queueing

[Hang-get]

Rebased onto current main (efec60ea5), which includes the CI hang fix from #22101 (6edbe383). This branch is now da82467ca. The patch is unchanged.