From fbc995d3e08cf394b7418eb6d8583a89a7cf12f9 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 11 Jun 2026 21:04:16 -0500 Subject: [PATCH] Account for reasoning token stripping Signed-off-by: Teodor-Dumitru Ene --- .../dynamic_text_gen_server/endpoints/chat_completions.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index 460acf39e9b..0e471064293 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -343,12 +343,14 @@ def _replace_prefix_tokens( from the previous generation (rather than the ones from the chat template application).""" # Strip the EOS from the previous turn token ids if it exists - if previous_turn_token_ids[-1] == eos_token_id: + if previous_turn_token_ids and previous_turn_token_ids[-1] == eos_token_id: previous_turn_token_ids = previous_turn_token_ids[:-1] # Find the last EOS token id in the previous turn token ids last_eos_token_id_index = len(retokeenized_previous_turn_token_ids) - 1 - for i in reversed(range(len(retokeenized_previous_turn_token_ids))): + # Note that the current conversation stat may be shorter than the previous conversation state. + scan_len = min(len(retokeenized_previous_turn_token_ids), len(current_turn_token_ids)) + for i in reversed(range(scan_len)): if current_turn_token_ids[i] == eos_token_id: last_eos_token_id_index = i break