From c44c815597b8ec7411f3cfdd8a9b0aff79b52df3 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson Date: Sat, 21 Jun 2025 17:37:38 +0000 Subject: [PATCH 1/3] Speedup MRoPE prepare inputs Signed-off-by: Vadim Gimpelson --- vllm/model_executor/layers/rotary_embedding.py | 13 +++++++++++++ vllm/v1/worker/gpu_model_runner.py | 17 ++++++++--------- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding.py b/vllm/model_executor/layers/rotary_embedding.py index 9de2338968a1..96c9df05bc95 100644 --- a/vllm/model_executor/layers/rotary_embedding.py +++ b/vllm/model_executor/layers/rotary_embedding.py @@ -26,6 +26,7 @@ import math from typing import Any, Optional, Union +import numpy as np import torch import torch.nn as nn from transformers import PretrainedConfig @@ -1468,6 +1469,18 @@ def get_next_input_positions_tensor( mrope_position_delta + seq_len, ).expand(3, -1) + # Faster version of `get_next_input_positions_tensor` + @staticmethod + def mrope_assign_next_input_positions(out: np.ndarray, out_offset: int, + mrope_position_delta: int, + context_len: int, + num_new_tokens: int): + + for dim in range(3): + for idx in range(num_new_tokens): + out[dim, out_offset + + idx] = mrope_position_delta + context_len + idx + @classmethod def omni_get_updates_use_audio_in_video( cls, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f96fb64342c9..a6d07a53f581 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -261,6 +261,7 @@ def __init__( dtype=torch.int64, device="cpu", pin_memory=self.pin_memory) + self.mrope_positions_np = self.mrope_positions_cpu.numpy() # Only relevant for models using ALiBi (e.g, MPT) self.use_alibi = check_use_alibi(model_config) @@ -888,15 +889,13 @@ def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"): dst_start = mrope_pos_ptr dst_end = mrope_pos_ptr + completion_part_len - self.mrope_positions_cpu[:, dst_start:dst_end] = \ - MRotaryEmbedding.get_next_input_positions_tensor( - req.mrope_position_delta, - context_len=num_computed_tokens + - prompt_part_len, - seq_len=num_computed_tokens + - prompt_part_len + - completion_part_len, - ) + MRotaryEmbedding.mrope_assign_next_input_positions( + out=self.mrope_positions_np, + out_offset=dst_start, + mrope_position_delta=req.mrope_position_delta, + context_len=num_computed_tokens + prompt_part_len, + num_new_tokens=completion_part_len, + ) mrope_pos_ptr += completion_part_len From 029f1e361638bd286fcd4cfbdcb47bed8cc34958 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson Date: Sat, 21 Jun 2025 18:19:12 +0000 Subject: [PATCH 2/3] fix comment Signed-off-by: Vadim Gimpelson --- vllm/model_executor/layers/rotary_embedding.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding.py b/vllm/model_executor/layers/rotary_embedding.py index 96c9df05bc95..bc4f559aba25 100644 --- a/vllm/model_executor/layers/rotary_embedding.py +++ b/vllm/model_executor/layers/rotary_embedding.py @@ -1476,10 +1476,10 @@ def mrope_assign_next_input_positions(out: np.ndarray, out_offset: int, context_len: int, num_new_tokens: int): - for dim in range(3): - for idx in range(num_new_tokens): - out[dim, out_offset + - idx] = mrope_position_delta + context_len + idx + values = np.arange(mrope_position_delta + context_len, + mrope_position_delta + context_len + num_new_tokens, + dtype=out.dtype) + out[:, out_offset:out_offset + num_new_tokens] = values @classmethod def omni_get_updates_use_audio_in_video( From 8baa18e013f9a951d1881415aa21db931cf9c985 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson Date: Tue, 24 Jun 2025 00:59:35 +0000 Subject: [PATCH 3/3] fix another comment Signed-off-by: Vadim Gimpelson --- .../model_executor/layers/rotary_embedding.py | 19 +++---------------- vllm/v1/worker/gpu_model_runner.py | 2 +- 2 files changed, 4 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding.py b/vllm/model_executor/layers/rotary_embedding.py index bc4f559aba25..b7bb2affc4fa 100644 --- a/vllm/model_executor/layers/rotary_embedding.py +++ b/vllm/model_executor/layers/rotary_embedding.py @@ -1459,22 +1459,9 @@ def get_next_input_positions( ] @staticmethod - def get_next_input_positions_tensor( - mrope_position_delta: int, - context_len: int, - seq_len: int, - ) -> torch.Tensor: - return torch.arange( - mrope_position_delta + context_len, - mrope_position_delta + seq_len, - ).expand(3, -1) - - # Faster version of `get_next_input_positions_tensor` - @staticmethod - def mrope_assign_next_input_positions(out: np.ndarray, out_offset: int, - mrope_position_delta: int, - context_len: int, - num_new_tokens: int): + def get_next_input_positions_tensor(out: np.ndarray, out_offset: int, + mrope_position_delta: int, + context_len: int, num_new_tokens: int): values = np.arange(mrope_position_delta + context_len, mrope_position_delta + context_len + num_new_tokens, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index a6d07a53f581..c3b04a3f2da1 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -889,7 +889,7 @@ def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"): dst_start = mrope_pos_ptr dst_end = mrope_pos_ptr + completion_part_len - MRotaryEmbedding.mrope_assign_next_input_positions( + MRotaryEmbedding.get_next_input_positions_tensor( out=self.mrope_positions_np, out_offset=dst_start, mrope_position_delta=req.mrope_position_delta,