diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 12afdfd74..23d6dfc35 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -362,6 +362,7 @@ class MultimodalInputs: # QWen2-VL related mrope_positions: Optional[torch.Tensor] = None mrope_position_delta: Optional[torch.Tensor] = None + mrope_position_delta_repeated_cache: Optional[torch.Tensor] = None @staticmethod def from_dict(obj: dict): diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a713f399b..aaeb4fa7f 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -654,10 +654,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): seq_len: int, ) -> torch.Tensor: # doing below compute on cpu to avoid frequent small kernels - mrope_position_deltas = mm_input.mrope_position_delta.flatten() - mrope_positions = ( - (mrope_position_deltas + seq_len - 1).unsqueeze(0).repeat(3, 1) - ) + if mm_input.mrope_position_delta_repeated_cache is None: + mm_input.mrope_position_delta_repeated_cache = ( + (mm_input.mrope_position_delta - 1).flatten().unsqueeze(0).repeat(3, 1) + ) + mrope_positions = mm_input.mrope_position_delta_repeated_cache + seq_len return mrope_positions def _compute_mrope_positions(