tensorrt_llm/_torch/pyexecutor/llm_request.py (2)

303-304: Plumb and document the new flag across creators; consider explicit docstring.

You added is_first_draft but it’s not documented and isn’t set in executor_request_to_llm_request (defaults False). If external creators should be able to set it, add a short param docstring and, if applicable, thread it from ExecutorRequest; otherwise, clarify it’s internal-only.

 class LlmRequest(...):
-    def __init__(..., is_first_draft: Optional[bool] = False, **kwargs):
+    def __init__(..., is_first_draft: Optional[bool] = False, **kwargs):
+        """...
+        Args:
+            ...
+            is_first_draft: Internal flag used by Eagle3 first-draft flow.
+        """

---

1-1: Missing NVIDIA copyright header.

Per repo guidelines, prepend the 2025 NVIDIA copyright header to this file.

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

1-1: Missing NVIDIA copyright header.

Add the header as required.

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

1-1: Missing NVIDIA copyright header.

Add the header as required.

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (5)

56-69: Graph key derivation: handle non-Eagle3 managers and keep bool out of semantics unless needed.

Logic is fine. Minor: add a space after comma, and prefer returning a 2-tuple unless the boolean truly changes shapes to reduce graph proliferation.

-    def get_graph_key(self, batch_size,spec_resource_manager: Optional[BaseResourceManager] = None):
+    def get_graph_key(
+        self, batch_size: int, spec_resource_manager: Optional[BaseResourceManager] = None
+    ):

---

93-108: maybe_get_cuda_graph: update docstring and fix long lines.

Docstring mentions 3-tuple but you return 4. Also address E501 by breaking long returns.

-        Returns a tuple containing:
-        - A boolean indicating if a graph can be used.
-        - The attn_metadata for the graph, if applicable.
-        - The spec_metadata for the graph, if applicable.
+        Returns:
+            (can_use: bool, attn_metadata, spec_metadata, key)
@@
-            return False, None, None,None
+            return False, None, None, None
@@
-                return False, None, None,None
+                return False, None, None, None
@@
-            return False, None, None,None
+            return False, None, None, None

---

147-150: needs_capture(key): accessor looks good; keep signature consistent.

Minor nit: remove trailing blank line.

---

183-188: Avoid print in library code; use logger.

Replace print with logger.debug to integrate with existing logging.

-        print(f"capture graph {key}")
+        from tensorrt_llm.logger import logger
+        logger.debug("capture graph %s", key)

---

1-1: Missing NVIDIA copyright header.

Add the header as required.

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

1-1: Missing NVIDIA copyright header.

Add the header as required.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

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

78-89: Fix None-safety and unreachable return; clarify EAGLE3 logic.

• spec_resource_manager may be None and BaseResourceManager doesn’t declare is_first_draft; use getattr with a default.
• The second return is dead code.
• Make the parameter Optional and return a bool explicitly.

-    def attention_need_spec_dec_mode(self,
-                                     spec_resource_manager: BaseResourceManager,
-                                     is_draft_model: bool):
+    def attention_need_spec_dec_mode(
+            self,
+            spec_resource_manager: Optional[BaseResourceManager],
+            is_draft_model: bool) -> bool:
@@
-        return self.is_eagle3_one_model() or (
-            self.is_eagle3() and spec_resource_manager.is_first_draft) or (
-                self.is_eagle3() and not is_draft_model)
-        return self.is_eagle3_one_model() or (
-            self.is_eagle3() and spec_resource_manager.is_first_draft)
+        if self.is_eagle3_one_model():
+            return True
+        if not self.is_eagle3():
+            return False
+        is_first = getattr(spec_resource_manager, "is_first_draft", False)
+        return is_first or not is_draft_model

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

294-296: Don’t mutate shared spec_config; also guard None.

This can NPE when spec_config is None and leaks mutations across engines.

-        self.original_max_draft_len = spec_config.max_draft_len
-        if spec_config is not None and is_draft_model:
-            spec_config.max_draft_len = 0
-        self.spec_config = spec_config
+        self.original_max_draft_len = spec_config.max_draft_len if spec_config else 0
+        self.spec_config = copy.deepcopy(spec_config) if spec_config else None
+        if self.spec_config is not None and is_draft_model:
+            self.spec_config.max_draft_len = 0

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

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

119-126: Fix padding: handle Tensor/list, avoid negative pad, and don’t mutate upstream token container

input_tokens.extend(...) will fail when input_tokens is a Tensor and may compute a negative pad. Convert to a list, clamp pad to ≥0, and pass the copy to the request.

Apply this diff:

-        #pad input_tokens to max_draft_tokens
-        input_tokens.extend(
-            0 for _ in range(self.max_draft_tokens - num_accepted_tokens))
-        new_request = self._create_draft_request(request, input_tokens)
+        # Ensure list and pad to max_draft_tokens
+        if isinstance(input_tokens, torch.Tensor):
+            tokens = input_tokens.tolist()
+        else:
+            tokens = list(input_tokens)
+        pad = max(0, self.max_draft_tokens - num_accepted_tokens)
+        tokens.extend([0] * pad)
+        new_request = self._create_draft_request(request, tokens)
         new_request.state = LlmRequestState.GENERATION_IN_PROGRESS
-        new_request.py_num_accepted_draft_tokens = request.py_num_accepted_draft_tokens
+        new_request.py_num_accepted_draft_tokens = request.py_num_accepted_draft_tokens
         new_request.py_is_first_draft = True

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (2)

16-18: Add guarded import for LlmRequest used in annotations

Ensures the forward-ref resolves for type checkers without causing import cycles.

 if TYPE_CHECKING:
     from .model_engine import PyTorchModelEngine
+    from .llm_request import LlmRequest

---

54-55: Fix undefined type in annotation (Ruff F821): use LlmRequest in TYPE_CHECKING

Optional["Request"] is undefined at runtime and breaks lint/CI.

Apply this diff here:

-        self.padding_dummy_request: Optional["Request"] = None
+        self.padding_dummy_request: Optional["LlmRequest"] = None

And add the guarded import (see separate comment on Lines 16-18).

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

115-118: Docstring/state mismatch for accepted-tokens path

Docstring says “chunked context request,” but state is set to GENERATION_IN_PROGRESS. Confirm intent and align doc/comment accordingly.

---

1-1: Add NVIDIA copyright header

Per repo guidelines, prepend the header.

+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (5)

27-30: Update class docstring to reflect new 2- or 3-tuple graph keys

Doc still says “(batch size, draft_len)”; keys can also be (batch_size, draft_len, is_first_draft).

-    multiple graphs, keyed by (batch size, draft_len).
+    multiple graphs, keyed by (batch_size, draft_len) or
+    (batch_size, draft_len, is_first_draft) for Eagle3 first-draft flow.

---

101-112: Update maybe_get_cuda_graph docstring to include returned key

The function now returns 4 values (including the key).

-        Returns a tuple containing:
+        Returns a tuple containing:
         - A boolean indicating if a graph can be used.
         - The attn_metadata for the graph, if applicable.
         - The spec_metadata for the graph, if applicable.
+        - The graph key used for lookup/capture/replay.

---

19-21: Correct misleading comment about dummy request ID

(1 << 64) - 1 isn’t a prime; it’s a sentinel. Update the comment.

-# A large prime number used for dummy request IDs to avoid collisions
+# Sentinel 64-bit ID used for CUDA-graph dummy requests to avoid collisions
 CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1

---

198-198: Use logger instead of print

Prefer repo logger to avoid stdout noise.

-        print(f"capture graph {key}")
+        logger.debug(f"capture graph {key}")

Add the import near the other imports:

+from tensorrt_llm.logger import logger

---

1-1: Add NVIDIA copyright header

Per repo guidelines, prepend the header.

+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.