Fix mrope_positions size when req is retracted (#13700)

Signed-off-by: Kun(llfl) <i@imux.top>
Co-authored-by: Xuchun Shang <xuchun.shang@gmail.com>
This commit is contained in:
kun-llfl
2025-12-02 11:38:20 +08:00
committed by GitHub
parent 427a19b643
commit 106df4eac5
2 changed files with 58 additions and 0 deletions

View File

@@ -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)

View File

@@ -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