Support mrope_position_delta cache

This commit is contained in:
shadowxz109
2026-03-06 08:50:53 +08:00
committed by GitHub
parent 9ebffef1ef
commit 261be85ecc
2 changed files with 6 additions and 4 deletions

View File

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

View File

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