diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 05de30dd3..e8f3eb611 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -835,3 +835,45 @@ def hash_feature(f): reconstruct_t = f.reconstruct_on_target_device(torch.cuda.current_device()) return tensor_hash([reconstruct_t]) return data_hash(f) + + +def extend_mrope_positions_for_retracted_request( + mrope_positions: torch.Tensor, output_ids_len: int +) -> torch.Tensor: + """ + Extend mrope_positions for retracted requests by appending positions for output_ids. + + When a request is retracted and has multimodal inputs with mrope_positions, + we need to extend the positions to cover the output_ids that were already generated. + For pure text tokens, all three dimensions use the same incremental sequence. + + Args: + mrope_positions: The original mrope positions tensor, shape (3, origin_input_ids_len) + output_ids_len: The number of output tokens to generate positions for + + Returns: + Extended mrope_positions tensor with shape (3, origin_input_ids_len + output_ids_len) + """ + if output_ids_len <= 0: + return mrope_positions + + # Get the last position value corresponding to origin_input_ids + # mrope_positions shape: (3, origin_input_ids_len) + last_position = mrope_positions[:, -1] # shape: (3,) + + # Generate pure text mrope positions for output_ids + # All three dimensions for pure text are the same incremental sequence + start_pos = last_position[0] + 1 # Start from last position + 1 + output_positions = ( + torch.arange( + start_pos, + start_pos + output_ids_len, + dtype=torch.int64, + device=mrope_positions.device, + ) + .unsqueeze(0) + .expand(3, -1) + ) # shape: (3, output_ids_len) + + # Concatenate to the original mrope_positions + return torch.cat([mrope_positions, output_positions], dim=1) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index aaa0a00b2..00257318b 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -810,6 +810,22 @@ class Req: match_result.host_hit_length, ) self.cache_protected_len = len(self.prefix_indices) + + if ( + self.is_retracted + and self.multimodal_inputs is not None + and self.multimodal_inputs.mrope_positions is not None + ): + from sglang.srt.managers.mm_utils import ( + extend_mrope_positions_for_retracted_request, + ) + + self.multimodal_inputs.mrope_positions = ( + extend_mrope_positions_for_retracted_request( + self.multimodal_inputs.mrope_positions, len(self.output_ids) + ) + ) + self.extend_input_len = len(self.fill_ids) - len(self.prefix_indices) # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313