tensorrt_llm/_torch/speculative/interface.py (1)

1-1: Add NVIDIA Apache-2.0 header.

Per coding guidelines, prepend the 2025 NVIDIA Apache-2.0 copyright header to all source files.

+# 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.

tests/unittest/_torch/speculative/test_eagle3.py (2)

1-1: Add NVIDIA Apache-2.0 header.

Apply the standard 2025 NVIDIA Apache-2.0 header at the top.

+# 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.

---

29-39: Param matrix expansion looks good; add IDs and guard against degenerate accept-rate division.

• Add ids to parametrize for easier triage.
• Some combinations can yield zero drafted tokens (e.g., early stop); guard accept-rate division.

Apply:

 @pytest.mark.parametrize(
-    "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill",
-    [
+    "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill",
+    [
       ...
     ])

And later near accept-rate computation:

-    accept_rate = num_accepted / num_drafted
+    accept_rate = (num_accepted / num_drafted) if num_drafted > 0 else 0.0

Optionally:

-@pytest.mark.parametrize(..., [...], [...], ...)
+@pytest.mark.parametrize(
+    "...",
+    [
+        pytest.param(True, "TRTLLM", True,  False, True,  False, id="cg_on_trtllm_overlap_off_one_model"),
+        ...
+    ]
+)

tensorrt_llm/_torch/speculative/model_drafter.py (6)

1-1: Add NVIDIA Apache-2.0 header.

Insert the 2025 NVIDIA Apache-2.0 header at the file top.

+# 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.

---

429-433: Docstring arg names do not match function signature.

The docstring mentions draft_sample_state/iteration, but the function takes draft_tensors, draft_position, and draft_length.

-        Args:
-            target_inputs: The target input tensors to update
-            draft_sample_state: The draft sample state containing new tokens
-            iteration: The current iteration index
+        Args:
+            target_inputs: The target input tensors to update.
+            draft_tensors: Draft tokens tensor for this update step.
+            draft_position: Starting draft position to write into.
+            draft_length: Number of draft tokens to write.
+            num_draft_reqs: Number of draft requests in the batch.

---

368-421: Make SampleStateTensorsMTP fields non-None to avoid downstream None-handling.

SampleStateTensorsMTP defines tensor fields; returning None for new_tokens_lens/next_draft_tokens risks attr errors downstream. Initialize to ones/zeros with expected shapes even when there are no generation requests.

-        new_tokens_lens = None
-        next_draft_tokens = None
+        new_tokens_lens = torch.ones(scheduled_batch.batch_size,
+                                     device=device)
+        next_draft_tokens = torch.zeros(scheduled_batch.batch_size,
+                                        self.max_draft_tokens,
+                                        device=device)
 ...
-            else:
-                # Create new tensors with the correct device
-                # We already updated the target state, so the new_tokens_lens should be all ones.
-                new_tokens_lens = torch.ones(scheduled_batch.batch_size,
-                                             device=device)
-                next_draft_tokens = torch.zeros(scheduled_batch.batch_size,
-                                                self.max_draft_tokens,
-                                                device=device)
+            else:
+                # Shapes already initialized above; just fill per-request entries as needed.
                 num_accepted_tokens = request.py_num_accepted_draft_tokens

---

383-399: Guard against missing py_num_accepted_draft_tokens on context requests.

On some paths, context requests may not have this attribute set. Default to 0 to avoid AttributeError.

-                num_accepted_tokens = request.py_num_accepted_draft_tokens
+                num_accepted_tokens = getattr(request, "py_num_accepted_draft_tokens", 0)

---

667-673: guided_decoder.execute signature inconsistency across paths.

Here you call execute(draft_logits, target_logits, d2t=...), while in prepare_draft_tokens it’s execute(logits, d2t=...). Align usage to the expected API to avoid silent misbehavior when guided decoding is enabled.

Would you like me to adjust both call sites to a single signature after confirming the current API?

---

634-639: Return-type docstring mismatch.

The function returns three values (target_inputs, previous_draft_state, draft_batch), but the docstring lists two.

-        Returns:
-            Tuple[Optional[SampleStateTensorsMTP], Optional[SampleState]]:
-                - Updated target inputs or None
-                - Draft sample state or None
+        Returns:
+            Tuple[Optional[SampleStateTensorsMTP], Optional[Any], Optional[ScheduledRequests]]:
+                - Updated target inputs or None
+                - Final draft sample state or None
+                - Draft batch used for the iteration or None

tensorrt_llm/_torch/pyexecutor/model_engine.py (3)

1304-1307: Improve error message for assertion

The assertion should provide a more informative error message explaining when SampleStateTensorsMTP is expected.

-            assert self.enable_spec_decode and not self.is_draft_model
+            assert self.enable_spec_decode and not self.is_draft_model, \
+                "SampleStateTensorsMTP should only be used with speculative decoding enabled on target model"

---

1473-1486: Simplify conditional logic for cached tokens calculation

The nested conditions can be simplified for better readability. Also, the comment at line 1474 could be clearer about when this path is taken.

-                if self.spec_config.spec_dec_mode.has_draft_model():
-                    # In the overlap scheduler workflow, if having draft model, we already updated the previous batch before launching the target model,
-                    # so we only need to add the runtime_draft_len to the past_seen_token_num.
-                    num_cached_tokens_per_seq.append(past_seen_token_num +
-                                                     self.runtime_draft_len)
-                else:
-                    num_cached_tokens_per_seq.append(past_seen_token_num +
-                                                     self.runtime_draft_len + 1)
+                # In overlap scheduler with draft model, the previous batch was already updated,
+                # so we only add runtime_draft_len. Without draft model, we add +1 for the current token.
+                offset = self.runtime_draft_len if self.spec_config.spec_dec_mode.has_draft_model() else self.runtime_draft_len + 1
+                num_cached_tokens_per_seq.append(past_seen_token_num + offset)

---

1667-1674: Consider extracting magic number to a named field

The field num_extended_ctx_requests is initialized here and used in _preprocess_inputs. Consider documenting this field more clearly in the AttentionMetadata class.

Would you like me to help add proper documentation for the num_extended_ctx_requests field in the AttentionMetadata class to clarify its purpose and usage?

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

282-285: Improve error message clarity

The error message could be more specific about which execution modes are incompatible with drafting.

-                raise NotImplementedError(
-                    "Drafting is not supported for selected executor loop. "
-                    "Please disable disagg/pipeline parallelism scheduler.")
+                raise NotImplementedError(
+                    "Drafting is not supported with pipeline parallelism (_executor_loop_pp). "
+                    "Please disable pipeline parallelism to use drafting.")

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.

Learnt from: pcastonguay
PR: NVIDIA/TensorRT-LLM#7455
File: tensorrt_llm/_torch/pyexecutor/py_executor.py:728-731
Timestamp: 2025-09-02T13:43:22.657Z
Learning: The user pcastonguay prefers creating dedicated handler classes to encapsulate complex subsystem logic rather than spreading it across the main class. For disaggregated pipeline parallel termination, they suggest creating a `_disagg_pp_termination_handler` with a `cleanup()` method instead of manually waiting on termination handles during shutdown.