From aa064f66329696bef6a6aa1705b24b365a73785a Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Thu, 17 Jul 2025 01:39:07 -0700 Subject: [PATCH 1/5] fix previous_pos_id_offsets_cuda previous_kv_lens_offsets_cuda update Signed-off-by: Yue Weng <25103990+yweng0828@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 41 +++++++++++++------ tensorrt_llm/_torch/pyexecutor/py_executor.py | 4 -- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9f9d3ea184dd..754f53a3373d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1340,6 +1340,10 @@ def previous_seq_slots_device(): self.draft_tokens_cuda[:len(draft_tokens)].copy_(draft_tokens, non_blocking=True) if next_draft_tokens_device is not None: + # Initialize these two values to zeros + self.previous_pos_id_offsets_cuda *= 0 + self.previous_kv_lens_offsets_cuda *= 0 + if previous_batch_len > 0: previous_slots = previous_seq_slots_device() # previous input ids @@ -1362,26 +1366,37 @@ 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'] + # 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. self.previous_pos_id_offsets_cuda[ - 0:previous_batch_tokens].copy_( + new_tokens_len:new_tokens_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[0:previous_batch_len].copy_( - kv_len_offsets_device[previous_slots], non_blocking=True) - # for the requests that do not have previous batch, set the previous_pos_id_offsets and - # previous_kv_lens_offsets to zeros to skip the value changes in _preprocess_inputs - self.previous_pos_id_offsets_cuda[ - previous_batch_tokens:num_requests * - (1 + self.max_draft_len)] *= 0 self.previous_kv_lens_offsets_cuda[ - previous_batch_len:num_requests] *= 0 - else: - # change the data to zeros to skip the value changes in _preprocess_inputs - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 + num_requests - previous_batch_len - + len(extend_dummy_requests):num_requests - + len(extend_dummy_requests)].copy_( + kv_len_offsets_device[previous_slots], + non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() max_draft_len = max(draft_lens) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6303be150d27..f983a9797ad8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1022,10 +1022,6 @@ def _executor_loop_overlap(self): ) if self.kv_cache_transceiver: - # For generation requests which have completed KV cache transfer - self._prepare_disagg_gen_transmission_complete( - scheduled_batch) - # Return the first token to the client self._handle_first_token_response(scheduled_batch) From f1c156063f1d4278701769aec4b48afb6fe67df3 Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Fri, 18 Jul 2025 17:39:17 +0000 Subject: [PATCH 2/5] update logic for IFB Signed-off-by: Yue Weng <25103990+yweng0828@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 754f53a3373d..a4effab01fcd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1323,7 +1323,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") @@ -1366,34 +1366,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) From 16f22084706aa9217fed19e0ac5a1ee08f0b180b Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Mon, 21 Jul 2025 06:29:48 +0000 Subject: [PATCH 3/5] refactor Signed-off-by: Yue Weng <25103990+yweng0828@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index a4effab01fcd..d07af6ae11bd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1382,21 +1382,23 @@ def previous_seq_slots_device(): # 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. + + num_extend_reqeust_wo_dummy = len(extend_requests) - len( + extend_dummy_requests) self.previous_pos_id_offsets_cuda[ - (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_( + (num_extend_reqeust_wo_dummy - previous_batch_len) * + (1 + self.max_draft_len):num_extend_reqeust_wo_dummy * + (1 + self.max_draft_len)].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ 0:previous_batch_tokens]], non_blocking=True) + self.previous_kv_lens_offsets_cuda[ - len(extend_requests) - previous_batch_len - - len(extend_dummy_requests):len(extend_requests) - - len(extend_dummy_requests)].copy_( + num_extend_reqeust_wo_dummy - + previous_batch_len:num_extend_reqeust_wo_dummy].copy_( kv_len_offsets_device[previous_slots], non_blocking=True) + elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() max_draft_len = max(draft_lens) From 0fb3f006f4ae94e24797dc0e499e9884669c8c41 Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Tue, 22 Jul 2025 02:25:11 +0000 Subject: [PATCH 4/5] Remove unused code Signed-off-by: Yue Weng <25103990+yweng0828@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d07af6ae11bd..0cbc67114ec8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1323,7 +1323,6 @@ def previous_seq_slots_device(): num_tokens = len(input_ids) num_draft_tokens = len(draft_tokens) - 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") From 2d1e020c854b51cd6e01c8a6a79f699c3a049c80 Mon Sep 17 00:00:00 2001 From: Yue Weng <25103990+yweng0828@users.noreply.github.com> Date: Tue, 22 Jul 2025 10:58:48 +0000 Subject: [PATCH 5/5] unwaive tests Signed-off-by: Yue Weng <25103990+yweng0828@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 3e0b9c62eda5..7e9267006338 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -371,8 +371,6 @@ perf/test_perf.py::test_perf[bert_large-bench-float16-maxbs:32-input_len:128+512 perf/test_perf.py::test_perf[roberta_base-bench-float16-maxbs:32-input_len:128+512] SKIP (https://nvbugspro.nvidia.com/bug/5295411) disaggregated/test_disaggregated.py::test_disaggregated_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/5328160) stress_test/stress_test.py::test_run_stress_test[llama-v3-8b-instruct-hf_tp1-stress_time_300s_timeout_450s-MAX_UTILIZATION-pytorch-stress-test] SKIP (https://nvbugs/5328495) -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=True] SKIP (https://nvbugs/5322354) -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/5322354) full:B200/examples/test_gemma.py::test_llm_gemma_1gpu_summary_vswa[gemma-3-1b-it-other-bfloat16-8] SKIP (https://nvbugs/5292737) full:B200/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype SKIP (https://nvbugs/5295470) examples/test_mistral.py::test_llm_mistral_v1_1gpu[mistral-7b-v0.1-float16-max_attention_window_size_4096-summarization_long] SKIP (https://nvbugs/5324976)