[TRTLLM-9687][chore] Update testcase for write_finish_reasons and adjusted override of setup_sampler_step

Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com>

Author: stnie

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -887,6 +887,8 @@ def _create_store(self) -> Store:
                 first_finish_reasons=int_tensor(
                     self.CACHE_INDIRECTION_SHAPE[:-1],
                 ),
+                max_lengths_tensor=int_tensor(self.max_num_sequences),
+                end_ids=int_tensor(self.max_num_sequences),
             )
         else:
             return self.Store(
@@ -1330,15 +1332,16 @@ def _is_new_request(self, request: LlmRequest) -> bool:
             or request.is_disagg_generation_transmission_complete
         )
 
-    def setup_sampler_step(self, requests: ScheduledRequests):
+    @override
+    def setup_sampler_step(self, scheduled_requests: ScheduledRequests):
         """Setup the sampler step for the requests
 
         Args:
             requests: list[LlmRequest]. The requests to setup the sampler step for
         """
         if self._use_beam_search:
-            self._prepare_beam_search(requests.all_requests())
-        for request in requests.all_requests():
+            self._prepare_beam_search(scheduled_requests.all_requests())
+        for request in scheduled_requests.all_requests():
             if self._is_new_request(request):
                 self.store.max_lengths_tensor[request.py_seq_slot].fill_(
                     min(self.max_seq_len, request.orig_prompt_len + request.py_max_new_tokens)

tensorrt_llm/_torch/speculative/mtp.py

@@ -222,7 +222,10 @@ class Store(TorchSampler.Store):
         next_draft_tokens: torch.Tensor
         new_tokens_lens: torch.Tensor
         max_total_draft_tokens: torch.Tensor
-        finish_reasons: None = None  # Necessary to satisfy the interface of TorchSampler.Store
+        # Necessary to satisfy the interface of TorchSampler.Store
+        finish_reasons: None = None
+        end_ids: None = None
+        max_lengths_tensor: None = None
 
         def __post_init__(self):
             pass  # finish_reasons has no size to compare against new_tokens in MTPSampler

tests/unittest/_torch/sampler/test_torch_sampler.py

@@ -691,16 +691,40 @@ def setup(requests: list["RequestCase"]):
         seq_slots = torch.tensor(
             [req.request.py_seq_slot for req in requests], device="cuda", dtype=torch.int64
         )
+        seq_lens = torch.tensor(
+            [req.request.max_beam_num_tokens for req in requests], dtype=torch.int32, device="cuda"
+        )
         new_tokens = torch.tensor(
             [req.new_tokens for req in requests], dtype=torch.int32, device="cuda"
         ).T
         sampler.store.new_tokens[:, seq_slots, BEAM] = new_tokens
+        max_seq_lens = torch.tensor(
+            [
+                min(
+                    sampler.max_seq_len, req.request.orig_prompt_len + req.request.py_max_new_tokens
+                )
+                for req in requests
+            ],
+            dtype=torch.int32,
+            device="cuda",
+        )
+        end_ids = torch.tensor(
+            [
+                req.request.py_end_id if req.request.py_end_id is not None else -1
+                for req in requests
+            ],
+            dtype=torch.int32,
+            device="cuda",
+        )
+        sampler.store.max_lengths_tensor[seq_slots] = max_seq_lens
+        sampler.store.end_ids[seq_slots] = end_ids
 
         def run():
             sampler._write_finish_reasons(
                 [req.request for req in requests],
                 finish_reasons=sampler.store.finish_reasons,
                 new_tokens=sampler.store.new_tokens,
+                seq_lens=seq_lens,
                 seq_slots=seq_slots,
             )