Make changes suggested in PR.

Signed-off-by: Daniel Campora <[email protected]>

Author: dcampora

tensorrt_llm/_torch/pyexecutor/handle_logits.py

@@ -16,8 +16,6 @@ def __call__(
         context_requests: List[LlmRequest],
         generation_requests: List[LlmRequest],
         logits: torch.Tensor,
-        num_context_logits_prefix_sum: List[int],
-        max_num_sequences: int,
         beam_width: int,
     ):
         """Handles context and generation logits for a batch of requests.
@@ -26,10 +24,14 @@ def __call__(
             context_requests: List of context requests to process
             generation_requests: List of generation requests to process
             logits: Input logits tensor
-            num_context_logits_prefix_sum: Prefix sum of context logits for each request
-            max_num_sequences: Maximum number of sequences to process
             beam_width: Beam width for the generation requests
         """
+        num_context_logits_prefix_sum = [0]
+        prefix_sum = 0
+        for request in context_requests:
+            prefix_sum += request.context_chunk_size if request.py_return_context_logits else 1
+            num_context_logits_prefix_sum.append(prefix_sum)
+
         # Copy logits into decoderBuffers.logits
         for batch_index, llm_req in enumerate(context_requests):
             logits_begin = num_context_logits_prefix_sum[batch_index]

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -176,7 +176,6 @@ def __init__(self,
         self.guided_decoder = guided_decoder
         self.dist = dist
         self.disable_overlap_scheduler = disable_overlap_scheduler
-        self.max_num_sequences = max_num_sequences
 
         # enqueue and _fetch_new_requests used data
         self.active = True
@@ -1489,16 +1488,9 @@ def _sample_async(self, scheduled_batch,
     def _handle_logits(self, scheduled_batch, batch_outputs):
         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)
-
             HandleLogits()(
                 scheduled_batch.context_requests,
                 scheduled_batch.generation_requests, batch_outputs["logits"],
-                num_context_logits_prefix_sum, self.max_num_sequences,
                 self.sampler.beam_width(scheduled_batch.all_requests()))
 
     @nvtx_range("_setup_sampler_step")

tests/unittest/_torch/test_return_logits.py

@@ -12,7 +12,7 @@
 global_kvcache_config = KvCacheConfig(max_tokens=10000)
 
 
-# @force_ampere  # Save H100 resource
+@force_ampere  # Save H100 resource
 @pytest.mark.parametrize("return_log_probs", [False, True])
 @pytest.mark.parametrize("gather_generation_logits", [False, True])
 @pytest.mark.parametrize("gather_context_logits", [False, True])