update logic for IFB

Signed-off-by: Yue Weng <[email protected]>

Author: yweng0828

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1315,7 +1315,7 @@ def previous_seq_slots_device():
 
         num_tokens = len(input_ids)
         num_draft_tokens = len(draft_tokens)
-        num_requests = len(request_ids)
+        len(request_ids)
         total_num_tokens = len(position_ids)
         assert total_num_tokens <= self.max_num_tokens, (
             "total_num_tokens should be less than or equal to max_num_tokens")
@@ -1358,34 +1358,34 @@ def previous_seq_slots_device():
                 previous_pos_indices_host = torch.tensor(previous_pos_indices,
                                                          dtype=torch.int,
                                                          pin_memory=True)
-                new_tokens_len = new_tokens.shape[0]
                 self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_(
                     previous_pos_indices_host, non_blocking=True)
 
-                # The order of request in Batch: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests']
+                # The order of requests in a batch: [context requests, generation requests]
+                # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests']
+                #   1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving.
+                #   2) 'requests that already have previous batch': previous iteration's requests.
+                #   3) 'dummy requests': pad dummy requests for CUDA graph or attention dp.
                 # Therefore, both of self.previous_pos_id_offsets_cuda and self.previous_kv_lens_offsets_cuda are also 3 segments.
-                # 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving.
-                #    Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs.
-                #    self.previous_pos_id_offsets_cuda[0 : new_tokens_len] *= 0
-                #    self.previous_kv_lens_offsets_cuda[0 : num_requests - previous_batch_len - len(extend_dummy_requests)] *= 0
-                #    Already set to '0' during initialization.
-                # 2) 'requests that already have previous batch': enable overlap scheduler.
-                #    Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device.
-                #    self.previous_pos_id_offsets_cuda[new_tokens_len : new_tokens_len + previous_batch_tokens]
-                #    self.previous_kv_lens_offsets_cuda[num_requests - previous_batch_len - len(extend_dummy_requests) : num_requests - len(extend_dummy_requests)]
-                # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp.
-                #    self.previous_pos_id_offsets_cuda[new_tokens_len + previous_batch_tokens : num_requests * (1 + max_draft_len)]
-                #    self.previous_kv_lens_offsets_cuda[num_requests - len(extend_dummy_requests) : num_requests]
-                #    Already set to '0' during initialization.
+                #   For 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving.
+                #       Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs.
+                #       Already set to '0' during initialization.
+                #   For 2) 'requests that already have previous batch': enable overlap scheduler.
+                #       Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device.
+                #   For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp.
+                #       Already set to '0' during initialization.
                 self.previous_pos_id_offsets_cuda[
-                    new_tokens_len:new_tokens_len +
+                    (len(extend_requests) - len(extend_dummy_requests) -
+                     previous_batch_len) * (1 + self.max_draft_len):
+                    (len(extend_requests) - len(extend_dummy_requests) -
+                     previous_batch_len) * (1 + self.max_draft_len) +
                     previous_batch_tokens].copy_(
                         new_tokens_lens_device[self.previous_pos_indices_cuda[
                             0:previous_batch_tokens]],
                         non_blocking=True)
                 self.previous_kv_lens_offsets_cuda[
-                    num_requests - previous_batch_len -
-                    len(extend_dummy_requests):num_requests -
+                    len(extend_requests) - previous_batch_len -
+                    len(extend_dummy_requests):len(extend_requests) -
                     len(extend_dummy_requests)].copy_(
                         kv_len_offsets_device[previous_slots],
                         non_blocking=True)