A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 04:45 UTC
[PERF] Speedup of MRoPE prepare inputs (#19939)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@centml.ai>
Vadim Gimpelson committed
Jun 24, 2025 at 10:01 UTC
9a3b88328f7e434cac35b90ee463de6689f9a833
2 files changed
+17
-18
vllm/model_executor/layers/rotary_embedding.py
+9
-9
@@ -26,6 +26,7 @@
26
import math
27
from typing import Any, Optional, Union
28
29
+import numpy as np
30
import torch
31
import torch.nn as nn
32
from transformers import PretrainedConfig
@@ -1458,15 +1459,14 @@ class MRotaryEmbedding(RotaryEmbedding):
1459
]
1460
1461
@staticmethod
1461
- def get_next_input_positions_tensor(
1462
- mrope_position_delta: int,
1463
- context_len: int,
1464
- seq_len: int,
1465
- ) -> torch.Tensor:
1466
- return torch.arange(
1467
- mrope_position_delta + context_len,
1468
- mrope_position_delta + seq_len,
1469
- ).expand(3, -1)
1462
+ def get_next_input_positions_tensor(out: np.ndarray, out_offset: int,
1463
+ mrope_position_delta: int,
1464
+ context_len: int, num_new_tokens: int):
1465
+
1466
+ values = np.arange(mrope_position_delta + context_len,
1467
+ mrope_position_delta + context_len + num_new_tokens,
1468
+ dtype=out.dtype)
1469
+ out[:, out_offset:out_offset + num_new_tokens] = values
1470
1471
@classmethod
1472
def omni_get_updates_use_audio_in_video(
vllm/v1/worker/gpu_model_runner.py
+8
-9
@@ -262,6 +262,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
262
dtype=torch.int64,
263
device="cpu",
264
pin_memory=self.pin_memory)
265
+ self.mrope_positions_np = self.mrope_positions_cpu.numpy()
266
267
# Only relevant for models using ALiBi (e.g, MPT)
268
self.use_alibi = check_use_alibi(model_config)
@@ -889,15 +890,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
890
dst_start = mrope_pos_ptr
891
dst_end = mrope_pos_ptr + completion_part_len
892
892
- self.mrope_positions_cpu[:, dst_start:dst_end] = \
893
- MRotaryEmbedding.get_next_input_positions_tensor(
894
- req.mrope_position_delta,
895
- context_len=num_computed_tokens +
896
- prompt_part_len,
897
- seq_len=num_computed_tokens +
898
- prompt_part_len +
899
- completion_part_len,
900
- )
893
+ MRotaryEmbedding.get_next_input_positions_tensor(
894
+ out=self.mrope_positions_np,
895
+ out_offset=dst_start,
896
+ mrope_position_delta=req.mrope_position_delta,
897
+ context_len=num_computed_tokens + prompt_part_len,
898
+ num_new_tokens=completion_part_len,
899
+ )
900
901
mrope_pos_ptr += completion_part_len
902