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:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user