[None][fix] helix: compensate position_id for the overlap scheduler

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -5699,6 +5699,7 @@ def append_cross_attention_state(request: LlmRequest,
                     input_ids.append(
                         request.get_tokens(0)[request.context_current_position])
                     past_seen_token_num = request.context_current_position
+                    request_has_previous_tensor = False
                 # The request has no previous tensor:
                 # (1) new_tokens_device is None, which means overlap scheduler is disabled; or
                 # (2) a dummy request; or
@@ -5720,17 +5721,30 @@ def append_cross_attention_state(request: LlmRequest,
                             else:
                                 input_ids.append(request.get_last_tokens(beam))
                     past_seen_token_num = request.max_beam_num_tokens - 1
+                    request_has_previous_tensor = False
                 else:
                     # the request has previous tensor
                     # previous_batch_indices is per-request, not per-beam
                     previous_batch_indices.append(request.py_batch_idx)
                     past_seen_token_num = request.max_beam_num_tokens
+                    request_has_previous_tensor = True
 
                 position_id = past_seen_token_num
                 if _has_cp_helix:
                     # We compute a global position_id because each helix rank has only a subset of
                     # tokens for a sequence.
                     position_id = request.total_input_len_cp + request.py_decoding_iter - 1
+                    if request_has_previous_tensor:
+                        # With the overlap scheduler this batch is prepared
+                        # before the previous iteration's _update_requests has
+                        # advanced py_decoding_iter, so the counter is one
+                        # behind. Compensate exactly like the non-helix path
+                        # above, which uses max_beam_num_tokens *without* the
+                        # -1 in this case. Without this, the position repeats
+                        # once (L, L, L+1, ...) and the new token's K is roped
+                        # at the wrong position before being written to the KV
+                        # cache, corrupting every later step.
+                        position_id += 1
                     if request.py_helix_is_inactive_rank:
                         past_seen_token_num = request.seqlen_this_rank_cp
                     else:

tests/integration/defs/accuracy/test_disaggregated_serving.py

@@ -1780,23 +1780,8 @@ def _test_chunked_prefill_helper(self, *, ctx_pp: int):
     def test_chunked_prefill(self):
         self._test_chunked_prefill_helper(ctx_pp=1)
 
-    @skip_pre_blackwell
-    @pytest.mark.skip_less_device(8)
-    @pytest.mark.parametrize("gen_pp,gen_tp,gen_cp,enable_attention_dp", [
-        (1, 2, 2, False),
-        (1, 2, 2, True),
-    ],
-                             ids=["pp1tp2cp2", "pp1dp2cp2"])
-    @pytest.mark.parametrize("cuda_graph_config", [
-        {
-            "enable_padding": True,
-            "batch_sizes": [1, 2, 4, 8, 16, 32, 64]
-        },
-    ],
-                             ids=["cudagraph:with_padding"])
-    @pytest.mark.parametrize("comms_medium", ["fifo_v2"])
-    def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config,
-                                   gen_pp, gen_tp, gen_cp, enable_attention_dp):
+    def _run_helix_test(self, comms_medium, cuda_graph_config, gen_pp, gen_tp,
+                        gen_cp, enable_attention_dp, disable_overlap_scheduler):
         # Parse comms_medium to get use_nccl_for_alltoall and fifo_version.
         if comms_medium == "nccl":
             use_nccl_for_alltoall = True
@@ -1841,12 +1826,12 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config,
                 "use_nccl_for_alltoall": use_nccl_for_alltoall,
                 "fifo_version": fifo_version,
             },
-            "disable_overlap_scheduler": True,
             "kv_cache_config": kv_cache_config,
             "enable_chunked_prefill": False,
             "cuda_graph_config": cuda_graph_config,
             "cache_transceiver_config": cache_transceiver_config.copy(),
             "enable_attention_dp": enable_attention_dp,
+            "disable_overlap_scheduler": disable_overlap_scheduler,
         }
         disaggregated_server_config = {
             "hostname": "localhost",
@@ -1866,6 +1851,51 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config,
                                       self.MODEL_PATH) as llm:
             run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"])
 
+    @skip_pre_blackwell
+    @pytest.mark.skip_less_device(8)
+    @pytest.mark.parametrize("gen_pp,gen_tp,gen_cp,enable_attention_dp", [
+        (1, 2, 2, False),
+        (1, 2, 2, True),
+    ],
+                             ids=["pp1tp2cp2", "pp1dp2cp2"])
+    @pytest.mark.parametrize("cuda_graph_config", [
+        {
+            "enable_padding": True,
+            "batch_sizes": [1, 2, 4, 8, 16, 32, 64]
+        },
+    ],
+                             ids=["cudagraph:with_padding"])
+    @pytest.mark.parametrize("comms_medium", ["fifo_v2"])
+    def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config,
+                                   gen_pp, gen_tp, gen_cp, enable_attention_dp):
+        self._run_helix_test(comms_medium,
+                             cuda_graph_config,
+                             gen_pp,
+                             gen_tp,
+                             gen_cp,
+                             enable_attention_dp,
+                             disable_overlap_scheduler=True)
+
+    @skip_pre_blackwell
+    @pytest.mark.skip_less_device(8)
+    def test_auto_dtype_with_helix_overlap(self):
+        # Regression test for the helix x overlap-scheduler position_id
+        # off-by-one: with the overlap scheduler enabled, generation batches
+        # are prepared one iteration ahead of py_decoding_iter, and an
+        # uncompensated helix position repeats once and corrupts the KV
+        # cache from the second decode step on. GSM8K fails hard without
+        # the compensation.
+        self._run_helix_test(comms_medium="fifo_v2",
+                             cuda_graph_config={
+                                 "enable_padding": True,
+                                 "batch_sizes": [1, 2, 4, 8, 16, 32, 64]
+                             },
+                             gen_pp=1,
+                             gen_tp=2,
+                             gen_cp=2,
+                             enable_attention_dp=False,
+                             disable_overlap_scheduler=False)
+
     @pytest.mark.skip_less_device(2)
     def test_gen_first(self):
         """Gen-first dense-model smoke test on KVCacheManagerV2 + NIXL python."""

tests/integration/test_lists/test-db/l0_dgx_b200.yml

@@ -164,6 +164,7 @@ l0_dgx_b200:
   tests:
   - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1dp2cp2] TIMEOUT (60)
   - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2] TIMEOUT (60)
+  - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix_overlap TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_mtp] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_bs8_mtp] TIMEOUT (60)