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

1-1: Prepend NVIDIA copyright header (current year).

Per coding guidelines, all source files must start with the NVIDIA copyright header.

Apply at the top of the file.

---

373-381: Bug: request.seq_slot should be request.py_seq_slot.

LlmRequest elsewhere uses py_seq_slot; request.seq_slot likely does not exist and will raise AttributeError during speculative acceptance/writeback.

Apply this diff:

-                new_tokens[i, request.seq_slot, self.BEAM] = new_token
+                new_tokens[i, request.py_seq_slot, self.BEAM] = new_token
@@
-                new_tokens[num_accepted, request.seq_slot,
+                new_tokens[num_accepted, request.py_seq_slot,
                            self.BEAM] = new_token

Also applies to: 391-399

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

1-1: Missing NVIDIA copyright header.

Per guidelines, prepend the NVIDIA copyright header.

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

1-1: Insert NVIDIA copyright header.

This file is missing the required header.

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

311-325: handle_logprobs builds correct per-token logprobs; minor robustness tweaks.

Logic is sound for TorchSampler (beam_width=1). Consider guarding against count > available length and documenting semantics to match request.py_result expectations.

Apply this diff to tighten safety and clarity:

-    def handle_logprobs(self, request: LlmRequest, state: SampleState, *, beam: int, count: int):
+    def handle_logprobs(self, request: LlmRequest, state: SampleState, *, beam: int, count: int):
+        """Append per-step token logprobs for the latest `count` steps (TorchSampler assumes beam=0)."""
         current_slice = slice(0, count), request.py_seq_slot, beam
         if request.py_return_log_probs:
             assert state.host.log_probs is not None
-            log_probs = state.host.log_probs[request.py_seq_slot][beam][:count]
+            lp = state.host.log_probs[request.py_seq_slot][beam]
+            assert lp.numel() >= count, f"Requested {count} logprobs, only {lp.numel()} available"
+            log_probs = lp[:count]
             current_tokens = state.host.new_tokens[current_slice]
 
             token_log_probs = [{
                 int(token): Logprob(logprob=logprob, rank=1)
             } for token, logprob in zip(current_tokens, log_probs.tolist())]
             assert beam == 0, "The following call relies on beam_width to be 1 - hence the list with a single element"
             request.py_result.append_log_probs([token_log_probs])

tests/unittest/_torch/test_return_logits.py (1)

1-1: Add NVIDIA copyright header.

All test sources should also prepend the current-year NVIDIA copyright.

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

684-699: Parameter ‘sampler’ is now unused; simplify signature and document behavior.

You’ve removed TorchSampler-specific gating, so sampler is unused. Rename to “_sampler” to avoid linter warnings and clarify backward compatibility; keep the assignment solely based on overlap scheduler as you do now.

Apply this diff:

-    def set_exclude_last_generation_logits(self,
-                                           disable_overlap_scheduler: bool,
-                                           sampler: Sampler) -> None:
+    def set_exclude_last_generation_logits(self,
+                                           disable_overlap_scheduler: bool,
+                                           _sampler: Sampler) -> None:
         # When overlap scheduler is enabled then when starting to handle a new prompt,
@@
-        # getter is required.
-        self.should_exclude_last_generation_logits = not disable_overlap_scheduler
+        # getter is required.
+        self.should_exclude_last_generation_logits = not disable_overlap_scheduler

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

1485-1500: Document _handle_logits and avoid duplication of prefix-sum logic across components.

Function works, but:

• Add a short docstring for maintainability.
• Consider centralizing the “num_context_logits_prefix_sum” computation to avoid duplicated logic (also computed in TRTLLMSampler) and ensure future changes don’t diverge.

Apply this diff to add a docstring now:

 @nvtx_range("_handle_logits")
 def _handle_logits(self, scheduled_batch, batch_outputs):
+    """Handle context/generation logits outside Sampler.
+
+    Builds per-context prefix sums and delegates tensor marshalling to HandleLogits,
+    using the batch logits and beam width derived from the Sampler.
+    No-op if no request asks for logits.
+    """
     if any(r.py_return_context_logits or r.py_return_generation_logits
            for r in scheduled_batch.all_requests()):
         num_context_logits_prefix_sum = [0]
         prefix_sum = 0
         for request in scheduled_batch.context_requests:
             prefix_sum += request.context_chunk_size if request.py_return_context_logits else 1
             num_context_logits_prefix_sum.append(prefix_sum)

If you want, I can factor this prefix-sum logic into a small helper and update both call sites.

---

752-756: Avoid dynamic attributes on dataclasses for logits.

Assigning sample_state.host.logits at runtime works in Python but breaks type-checkers and contracts implied by SampleStateTensors. Consider introducing a dedicated optional field on the host tensor type used in PP, or a lightweight wrapper, to keep the interface explicit.

I can propose a minimal type-safe wrapper to carry transient logits for PP without reintroducing logits into all Sampler paths.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro