[TRTLLM-7384][feat] enable rejection sampling for CDL #7731
Conversation
📝 WalkthroughWalkthroughAdds d2t attribute to LlmRequest. Updates sampler to realign draft probabilities using d2t and enforce single-strategy batching when mixed sampler is disabled. ChainDrafter.forward now returns both draft tokens and logits. ModelDrafter wires logits through static draft loop and attaches d2t to requests with debug prints. Test adds greedy/non-greedy parameterization. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant ModelDrafter
participant Drafter as ChainDrafter
participant Sampler
participant Target as TargetModel
Note over ModelDrafter: Prepare draft batch<br/>attach req.d2t
Client->>ModelDrafter: request drafts
ModelDrafter->>Drafter: forward(batch)
Drafter-->>ModelDrafter: {new_draft_tokens, draft_tokens_logits}
Note over ModelDrafter: Static draft loop<br/>accumulate tokens + logits<br/>set py_draft_tokens/py_draft_logits
ModelDrafter->>Target: run target step
Target-->>Sampler: target_probs
Note over Sampler: If d2t present,<br/>realign draft_probs to target_probs
ModelDrafter->>Sampler: draft_tokens + draft_probs (from logits)
Sampler-->>ModelDrafter: accept/reject decisions
ModelDrafter-->>Client: generated tokens
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
tensorrt_llm/_torch/pyexecutor/llm_request.py (1)
1-1: Add NVIDIA Apache-2.0 header (2025).Repository guidelines require the NVIDIA Apache-2.0 copyright header at the top of all .py files. Please add it here.
Example header:
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Licensed under the Apache License, Version 2.0 (the "License"); # You may not use this file except in compliance with the License. # You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License.tensorrt_llm/_torch/pyexecutor/sampler.py (1)
1-1: Add NVIDIA Apache-2.0 header (2025).Please add the required header to this file as well.
tensorrt_llm/_torch/speculative/drafting_loops.py (1)
1-1: Add NVIDIA Apache-2.0 header (2025).Please add the required header at the top.
tensorrt_llm/_torch/speculative/model_drafter.py (1)
1-1: Add NVIDIA Apache-2.0 header (2025).Please add the required header to this file.
tests/unittest/_torch/speculative/test_draft_target.py (1)
1-1: Add NVIDIA Apache-2.0 header (2025).Please add the required header at the top of the test file as well.
🧹 Nitpick comments (6)
tensorrt_llm/_torch/pyexecutor/llm_request.py (1)
358-359: d2t attribute: type, lifecycle, and child requests.
- Add a brief doc on what d2t is (shape=vocab, dtype=long, device semantics).
- Consider copying d2t in create_child_request if children will also need alignment; otherwise explicitly reset to None there to avoid accidental reuse.
Do you want me to wire a safe copy/reset in create_child_request?
tensorrt_llm/_torch/pyexecutor/sampler.py (3)
538-539: Use py_seq_slot consistently.Other paths use request.py_seq_slot. Mixing with request.seq_slot risks divergence. Use the python-side slot for consistency.
Apply:
- new_tokens[i, request.seq_slot, BEAM_0] = new_token + new_tokens[i, request.py_seq_slot, BEAM_0] = new_token
944-946: Single-strategy enforcement LGTM; add assertion message context.The assert is correct when mixed sampler is disabled. Consider including the set(strategies) in the message to aid debugging.
516-517: Small style: iterable unpacking over list concat.Replace list(dims) + [x] with (*dims, x) as in the patch above.
tensorrt_llm/_torch/speculative/drafting_loops.py (1)
134-137: Return-type change: update docstring and callers.ChainDrafter.forward now returns a dict with tokens and logits. Add a short docstring comment here documenting shapes to prevent misuse. Callers appear updated (ModelDrafter), so LGTM.
tests/unittest/_torch/speculative/test_draft_target.py (1)
17-23: Parametrize greedy vs non-greedy: good coverage; seed to avoid flakes.Non-greedy (top-p) paths can be nondeterministic. Consider fixing a seed in SamplingParams (if supported) or via torch.manual_seed/torch.cuda.manual_seed to minimize test flakes.
Would you like me to add a deterministic seed for this test?
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/llm_request.py(1 hunks)tensorrt_llm/_torch/pyexecutor/sampler.py(2 hunks)tensorrt_llm/_torch/speculative/drafting_loops.py(2 hunks)tensorrt_llm/_torch/speculative/model_drafter.py(4 hunks)tests/unittest/_torch/speculative/test_draft_target.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Use only spaces, no tabs; indent with 4 spaces.
Files:
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/_torch/speculative/drafting_loops.pytensorrt_llm/_torch/pyexecutor/sampler.pytests/unittest/_torch/speculative/test_draft_target.py
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.
Files:
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/_torch/speculative/drafting_loops.pytensorrt_llm/_torch/pyexecutor/sampler.pytests/unittest/_torch/speculative/test_draft_target.py
**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).
Files:
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/_torch/speculative/drafting_loops.pytensorrt_llm/_torch/pyexecutor/sampler.pytests/unittest/_torch/speculative/test_draft_target.py
🧠 Learnings (2)
📚 Learning: 2025-08-28T10:25:22.370Z
Learnt from: ixlmar
PR: NVIDIA/TensorRT-LLM#7294
File: tensorrt_llm/_torch/pyexecutor/sampler.py:887-891
Timestamp: 2025-08-28T10:25:22.370Z
Learning: In tensorrt_llm/_torch/pyexecutor/sampler.py, the draft_probs and target_probs tensors have shapes [1, steps] not [steps, vocab_size] as might be expected, making the .squeeze(0) operations appropriate for removing the batch dimension of size 1.
Applied to files:
tensorrt_llm/_torch/pyexecutor/sampler.py
📚 Learning: 2025-08-28T10:22:02.288Z
Learnt from: ixlmar
PR: NVIDIA/TensorRT-LLM#7294
File: tensorrt_llm/_torch/pyexecutor/sampler.py:1191-1197
Timestamp: 2025-08-28T10:22:02.288Z
Learning: In tensorrt_llm/_torch/pyexecutor/sampler.py, the object identity comparison `softmax_req_indices is not group_req_indices_cuda` on line ~1191 is intentional and used as an optimization to determine whether to reuse an existing indexer or create a new one, based on which code path was taken during tensor assignment.
Applied to files:
tensorrt_llm/_torch/pyexecutor/sampler.py
🧬 Code graph analysis (2)
tensorrt_llm/_torch/speculative/model_drafter.py (3)
tensorrt_llm/_torch/pyexecutor/scheduler.py (1)
all_requests(38-39)cpp/include/tensorrt_llm/batch_manager/llmRequest.h (1)
LlmRequestState(45-206)tensorrt_llm/_torch/pyexecutor/llm_request.py (2)
append(81-100)append(127-144)
tensorrt_llm/_torch/speculative/drafting_loops.py (2)
tensorrt_llm/_torch/pyexecutor/llm_request.py (2)
append(81-100)append(127-144)tensorrt_llm/_torch/speculative/eagle3.py (1)
Eagle3SpecMetadata(85-190)
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/sampler.py
516-517: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (2)
tensorrt_llm/_torch/speculative/drafting_loops.py (1)
108-110: Per-step logits collection LGTM.Capturing per-step logits matches downstream sampler needs.
Also applies to: 126-128
tests/unittest/_torch/speculative/test_draft_target.py (1)
58-61: Conditional SamplingParams LGTM.The construction matches the new parameter. No issues.
kris1025
force-pushed
the
mr-enable-rejection-sampling-for-cdl
branch
7 times, most recently
from
5740036 to
be3cf64
Compare
kris1025
changed the title
|
@coderabbitai help |
ChatThere are 3 ways to chat with CodeRabbit:
CodeRabbit commands
Other keywords and placeholders
Status, support, documentation and community
|
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
kris1025
force-pushed
the
mr-enable-rejection-sampling-for-cdl
branch
2 times, most recently
from
c036408 to
f6d0d5d
Compare
kris1025
force-pushed
the
mr-enable-rejection-sampling-for-cdl
branch
2 times, most recently
from
057904c to
f9afe36
Compare
|
/bot run |
|
PR_Github #19323 [ run ] triggered by Bot |
|
PR_Github #19323 [ run ] completed with state |
|
/bot run --reuse-test |
kris1025
force-pushed
the
mr-enable-rejection-sampling-for-cdl
branch
from
09259bf to
66b3973
Compare
|
/bot run |
|
PR_Github #20867 [ run ] triggered by Bot |
|
/bot kill |
kris1025
force-pushed
the
mr-enable-rejection-sampling-for-cdl
branch
from
66b3973 to
54340b2
Compare
|
PR_Github #20872 [ kill ] triggered by Bot |
|
PR_Github #20867 [ run ] completed with state |
|
PR_Github #20872 [ kill ] completed with state |
|
/bot run |
|
PR_Github #20876 [ run ] triggered by Bot |
|
PR_Github #20876 [ run ] completed with state |
kris1025
force-pushed
the
mr-enable-rejection-sampling-for-cdl
branch
from
54340b2 to
c42a62d
Compare
|
/bot run |
|
PR_Github #20891 [ run ] triggered by Bot |
|
PR_Github #20891 [ run ] completed with state |
|
/bot run |
|
PR_Github #21004 [ run ] triggered by Bot |
|
PR_Github #21004 [ run ] completed with state |
kris1025
force-pushed
the
mr-enable-rejection-sampling-for-cdl
branch
from
0bd452f to
c6e2a32
Compare
|
/bot run |
|
PR_Github #21026 [ run ] triggered by Bot |
|
PR_Github #21026 [ run ] completed with state |
4 participants
Summary by CodeRabbit
New Features
Tests
Description
Test Coverage
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)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
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
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.