[nvbug/5393888][nvbug/5393042] Always use py_seq_slot (#6147)

Signed-off-by: Netanel Haber <[email protected]>

Author: netanel-haber

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1152,7 +1152,7 @@ def _prepare_tp_inputs(
             if multimodal_params.has_content():
                 multimodal_params_list.append(multimodal_params)
 
-            request.py_batch_idx = request.seq_slot
+            request.py_batch_idx = request.py_seq_slot
 
         num_ctx_requests = len(scheduled_requests.context_requests)
         num_ctx_tokens = len(input_ids)
@@ -1234,11 +1234,11 @@ def _prepare_tp_inputs(
                 num_cached_tokens_per_seq.append(past_seen_token_num)
                 request_ids.append(request.py_request_id)
                 # update batch index
-                request.py_batch_idx = request.seq_slot
+                request.py_batch_idx = request.py_seq_slot
             else:
                 # update batch index
                 previous_batch_idx = request.py_batch_idx
-                request.py_batch_idx = request.seq_slot
+                request.py_batch_idx = request.py_seq_slot
                 # inputs
                 # overlap scheduler can only support the speculative decoding
                 # methods with a fixed number of draft tokens
@@ -1292,8 +1292,8 @@ def _prepare_tp_inputs(
                 gather_ids.append(len(position_ids) - 1)
 
             request_ids.append(request.py_request_id)
-            gen_request_seq_slots.append(request.seq_slot)
-            request.py_batch_idx = request.seq_slot
+            gen_request_seq_slots.append(request.py_seq_slot)
+            request.py_batch_idx = request.py_seq_slot
 
         previous_batch_len = len(previous_batch_indices)
 

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -194,7 +194,7 @@ def add_token(request: LlmRequest,
               *,
               beam: int,
               step: int = 0) -> int:
-    seq_slot = request.seq_slot
+    seq_slot = request.py_seq_slot
     assert seq_slot is not None
     new_token = int(new_tokens[step, seq_slot, beam])
     request.add_new_token(new_token, beam)
@@ -285,14 +285,14 @@ def _handle_stop_criteria(self, request: LlmRequest,
 
     def handle_logits(self, request: LlmRequest, state: SampleState, *,
                       beam: int, count: int):
-        current_slice = slice(0, count), request.seq_slot, beam
+        current_slice = slice(0, count), request.py_seq_slot, beam
         if request.py_return_generation_logits:
             assert state.host.logits is not None
             current_logits = state.host.logits[current_slice]
             request.py_result.append_generation_logits(current_logits)
         if request.py_return_log_probs:
             assert state.host.log_probs is not None
-            log_probs = state.host.log_probs[request.seq_slot][beam][:count]
+            log_probs = state.host.log_probs[request.py_seq_slot][beam][:count]
             current_tokens = state.host.new_tokens[current_slice]
 
             token_log_probs = [{
@@ -406,7 +406,7 @@ def _process_requests(self,
         no_draft_tokens = len(requests) == sum_steps
         fast_path = not self.enable_mixed_sampler and no_draft_tokens and gen_logits_host is None and log_probs_host is None
 
-        seq_slots = torch.as_tensor([r.seq_slot for r in requests])
+        seq_slots = torch.as_tensor([r.py_seq_slot for r in requests])
         seq_slots = seq_slots.to(device="cuda", non_blocking=True)
 
         if fast_path:
@@ -616,9 +616,9 @@ def _update_cache_indirection_buffer(self,
         # Copy cache indirection output to input
         for request in scheduled_requests.generation_requests:
             self.store["decoder_state"].cache_indirection_input[
-                request.seq_slot].copy_(
+                request.py_seq_slot].copy_(
                     self.store["decoder_state"].cache_indirection_output[
-                        request.seq_slot],
+                        request.py_seq_slot],
                     non_blocking=True)
 
     @torch.inference_mode()
@@ -881,7 +881,7 @@ def update_requests_multiple_beams_or_drafting(self,
 
     def _finalize_request(self, request: LlmRequest, streaming: bool):
         """ Finalizes the request. This is necessary for beam search. """
-        seq_slot = request.seq_slot
+        seq_slot = request.py_seq_slot
         event = self.algs.decoder.finalize(self.store["decoder_state"],
                                            seq_slot, request.sampling_config,
                                            streaming)
@@ -893,7 +893,7 @@ def _post_process_request(self, request: LlmRequest,
         request: LlmRequest which shall be post processed
         finalize_event: CudaEvent to wait for the finalize step to finish
         """
-        seq_slot = request.seq_slot
+        seq_slot = request.py_seq_slot
         beam_width = request.sampling_config.beam_width
         # synchronize on the finalize event before continuing the post processing.
         finalize_event.synchronize()

tensorrt_llm/_torch/speculative/mtp.py

@@ -232,7 +232,7 @@ def _request_common_handling(self, request: LlmRequest,
         assert not request.py_return_context_logits, "return_context_logits not implemented for MTPSampler"
         assert not request.py_return_generation_logits, "return_generation_logits not implemented for MTPSampler"
         assert not request.py_return_log_probs, "return_log_probs not implemented for MTPSampler"
-        request.py_draft_tokens = next_draft_tokens[request.seq_slot]
+        request.py_draft_tokens = next_draft_tokens[request.py_seq_slot]
         request.py_decoding_iter += 1
 
     def update_requests(self, state: SampleStateMTP) -> None:
@@ -253,7 +253,7 @@ def update_requests(self, state: SampleStateMTP) -> None:
         for req in state.scheduled_requests.generation_requests:
             if req.state == LlmRequestState.GENERATION_COMPLETE:
                 continue
-            num_new_tokens = new_tokens_lens[req.seq_slot]
+            num_new_tokens = new_tokens_lens[req.py_seq_slot]
             for i in range(num_new_tokens):
                 new_token = add_token(req, new_tokens, beam=beam_idx, step=i)
                 if self._handle_stop_criteria(req, new_token):
@@ -269,7 +269,7 @@ def sample_async(self, scheduled_requests: ScheduledRequests,
         # next_new_tokens_device: input tokens for the next iteration, device tensor, shape: batch_size, nextn + 1
 
         requests = scheduled_requests.all_requests()
-        slots = torch.as_tensor([r.seq_slot for r in requests])
+        slots = torch.as_tensor([r.py_seq_slot for r in requests])
         slots = slots.to(device="cuda", non_blocking=True)
 
         o_new_tokens = outputs['new_tokens'][:len(requests)]