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 10, 06:30 UTC
[Model Runner V2] Add support for M-RoPE (#32143)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 12, 2026 at 13:37 UTC
0a7dd23754ed5e01303e0eb4e64ace5e70251f46
5 files changed
+203
-7
vllm/v1/worker/gpu/cudagraph_utils.py
+6
-1
@@ -25,10 +25,12 @@ class CudaGraphManager:
25
def __init__(
26
self,
27
vllm_config: VllmConfig,
28
+ uses_mrope: bool,
29
device: torch.device,
30
):
31
self.vllm_config = vllm_config
32
self.scheduler_config = vllm_config.scheduler_config
33
+ self.uses_mrope = uses_mrope
34
self.device = device
35
36
self.max_model_len = vllm_config.model_config.max_model_len
@@ -79,7 +81,10 @@ class CudaGraphManager:
81
) -> None:
82
num_reqs = min(num_tokens, self.max_num_reqs)
83
input_ids = input_buffers.input_ids[:num_tokens]
82
- positions = input_buffers.positions[:num_tokens]
84
+ if not self.uses_mrope:
85
+ positions = input_buffers.positions[:num_tokens]
86
+ else:
87
+ positions = input_buffers.mrope_positions[:, :num_tokens]
88
attn_metadata = prepare_inputs_to_capture(
89
num_reqs,
90
num_tokens,
vllm/v1/worker/gpu/input_batch.py
+21
-2
@@ -31,6 +31,19 @@ class InputBuffers:
31
)
32
self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
33
34
+ # NOTE: `mrope_positions` is implemented with one additional dummy
35
+ # position on purpose to make it non-contiguous so that it can work
36
+ # with torch compile.
37
+ # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
38
+ # NOTE: When M-RoPE is enabled, position ids are 3D regardless of
39
+ # the modality of inputs. For text-only inputs, each dimension has
40
+ # identical position IDs, making M-RoPE functionally equivalent to
41
+ # 1D-RoPE.
42
+ # See page 5 of https://arxiv.org/abs/2409.12191
43
+ self.mrope_positions = torch.zeros(
44
+ (3, max_num_tokens + 1), dtype=torch.int64, device=device
45
+ )
46
+
47
48
@dataclass
49
class InputBatch:
@@ -62,6 +75,8 @@ class InputBatch:
75
input_ids: torch.Tensor
76
# [num_tokens_after_padding]
77
positions: torch.Tensor
78
+ # [3, num_tokens_after_padding]
79
+ mrope_positions: torch.Tensor
80
81
# layer_name -> Metadata
82
attn_metadata: dict[str, Any]
@@ -107,8 +122,11 @@ class InputBatch:
122
input_buffers.query_start_loc[num_reqs + 1 :] = num_tokens
123
query_start_loc = input_buffers.query_start_loc[: num_reqs + 1]
124
110
- input_ids = input_buffers.input_ids[:num_tokens]
111
- positions = input_buffers.positions[:num_tokens]
125
+ input_ids = input_buffers.input_ids[:num_tokens].zero_()
126
+ positions = input_buffers.positions[:num_tokens].zero_()
127
+ input_buffers.mrope_positions.zero_()
128
+ mrope_positions = input_buffers.mrope_positions[:, :num_tokens]
129
+
130
# attn_metadata = defaultdict(lambda: None)
131
logits_indices = query_start_loc[1:] - 1
132
cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
@@ -128,6 +146,7 @@ class InputBatch:
146
seq_lens=seq_lens,
147
input_ids=input_ids,
148
positions=positions,
149
+ mrope_positions=mrope_positions,
150
attn_metadata=None, # type: ignore
151
logits_indices=logits_indices,
152
cu_num_logits=cu_num_logits,
vllm/v1/worker/gpu/mm/__init__.py
vllm/v1/worker/gpu/mm/mrope_utils.py
new
+127
@@ -0,0 +1,127 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.model_executor.models.interfaces import SupportsMRoPE
6
+from vllm.triton_utils import tl, triton
7
+from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
8
+
9
+
10
+class MRopeState:
11
+ def __init__(
12
+ self,
13
+ max_num_reqs: int,
14
+ max_model_len: int,
15
+ device: torch.device,
16
+ ):
17
+ self.max_num_reqs = max_num_reqs
18
+ self.max_model_len = max_model_len
19
+ self.device = device
20
+
21
+ # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
22
+ # wasting a lot of CPU memory.
23
+ self.prefill_mrope_positions = StagedWriteTensor(
24
+ (max_num_reqs, 3 * max_model_len),
25
+ dtype=torch.int32,
26
+ device=device,
27
+ uva_instead_of_gpu=True,
28
+ )
29
+ self.prefill_mrope_delta = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
30
+
31
+ def init_prefill_mrope_positions(
32
+ self,
33
+ req_idx: int,
34
+ mrope_model: SupportsMRoPE,
35
+ prefill_token_ids: list[int],
36
+ mm_features: list,
37
+ ) -> None:
38
+ prefill_mrope_positions, prefill_mrope_delta = (
39
+ mrope_model.get_mrope_input_positions(
40
+ prefill_token_ids,
41
+ mm_features,
42
+ )
43
+ )
44
+ for i in range(3):
45
+ pos = prefill_mrope_positions[i].tolist()
46
+ self.prefill_mrope_positions.stage_write(
47
+ req_idx, i * self.max_model_len, pos
48
+ )
49
+ self.prefill_mrope_delta.np[req_idx] = prefill_mrope_delta
50
+
51
+ def apply_staged_writes(self) -> None:
52
+ self.prefill_mrope_positions.apply_write()
53
+ self.prefill_mrope_delta.copy_to_uva()
54
+
55
+ def prepare_mrope_positions(
56
+ self,
57
+ idx_mapping: torch.Tensor,
58
+ query_start_loc: torch.Tensor,
59
+ prefill_lens: torch.Tensor,
60
+ num_computed_tokens: torch.Tensor,
61
+ mrope_positions: torch.Tensor,
62
+ ) -> None:
63
+ num_reqs = idx_mapping.shape[0]
64
+ _prepare_mrope_positions_kernel[(num_reqs,)](
65
+ mrope_positions,
66
+ mrope_positions.stride(0),
67
+ self.prefill_mrope_positions.gpu,
68
+ self.prefill_mrope_positions.gpu.stride(0),
69
+ self.max_model_len,
70
+ self.prefill_mrope_delta.gpu,
71
+ idx_mapping,
72
+ query_start_loc,
73
+ prefill_lens,
74
+ num_computed_tokens,
75
+ BLOCK_SIZE=1024,
76
+ )
77
+
78
+
79
+@triton.jit
80
+def _prepare_mrope_positions_kernel(
81
+ mrope_positions_ptr,
82
+ mrope_positions_stride,
83
+ prefill_mrope_positions_ptr,
84
+ prefill_mrope_positions_stride0,
85
+ prefill_mrope_positions_stride1,
86
+ prefill_mrope_delta_ptr,
87
+ idx_mapping_ptr,
88
+ query_start_loc_ptr,
89
+ prefill_lens_ptr,
90
+ num_computed_tokens_ptr,
91
+ BLOCK_SIZE: tl.constexpr,
92
+):
93
+ batch_idx = tl.program_id(0)
94
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
95
+
96
+ prefill_len = tl.load(prefill_lens_ptr + req_state_idx)
97
+ num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
98
+ is_prefill = num_computed < prefill_len
99
+
100
+ query_start = tl.load(query_start_loc_ptr + batch_idx)
101
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
102
+ query_len = query_end - query_start
103
+
104
+ mrope_delta = tl.load(prefill_mrope_delta_ptr + req_state_idx)
105
+ for i in range(0, query_len, BLOCK_SIZE):
106
+ block = i + tl.arange(0, BLOCK_SIZE)
107
+ mask = block < query_len
108
+ orig_pos = num_computed + block
109
+
110
+ for j in tl.static_range(3):
111
+ if is_prefill:
112
+ # Read from pre-computed M-RoPE positions.
113
+ pos = tl.load(
114
+ prefill_mrope_positions_ptr
115
+ + req_state_idx * prefill_mrope_positions_stride0
116
+ + j * prefill_mrope_positions_stride1
117
+ + orig_pos,
118
+ mask=mask,
119
+ )
120
+ else:
121
+ # Apply M-RoPE delta.
122
+ pos = orig_pos + mrope_delta
123
+ tl.store(
124
+ mrope_positions_ptr + j * mrope_positions_stride + query_start + block,
125
+ pos,
126
+ mask=mask,
127
+ )
vllm/v1/worker/gpu/model_runner.py
+49
-4
@@ -47,6 +47,7 @@ from vllm.v1.worker.gpu.input_batch import (
47
prepare_pos_seq_lens,
48
prepare_prefill_inputs,
49
)
50
+from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
51
from vllm.v1.worker.gpu.sample.logprob import compute_prompt_logprobs
52
from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
53
from vllm.v1.worker.gpu.sample.output import SamplerOutput
@@ -94,6 +95,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
95
self.max_num_reqs = self.scheduler_config.max_num_seqs
96
self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
97
98
+ # Multimodal
99
+ self.uses_mrope = self.model_config.uses_mrope
100
+ if self.uses_mrope:
101
+ self.mrope_states = MRopeState(
102
+ max_num_reqs=self.max_num_reqs,
103
+ max_model_len=self.max_model_len,
104
+ device=self.device,
105
+ )
106
+
107
self.use_async_scheduling = self.scheduler_config.async_scheduling
108
self.output_copy_stream = torch.cuda.Stream(self.device)
109
self.output_copy_event = torch.cuda.Event()
@@ -132,7 +142,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
142
self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
143
144
# CUDA graphs.
135
- self.cudagraph_manager = CudaGraphManager(self.vllm_config, self.device)
145
+ self.cudagraph_manager = CudaGraphManager(
146
+ self.vllm_config, self.uses_mrope, self.device
147
+ )
148
# Structured outputs worker.
149
self.structured_outputs_worker = StructuredOutputsWorker(
150
max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
@@ -268,6 +280,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
280
dp_size = self.parallel_config.data_parallel_size
281
num_tokens_across_dp = make_num_tokens_across_dp(dp_size, num_tokens)
282
num_sampled_tokens = np.ones(input_batch.num_reqs, dtype=np.int32)
283
+ if not self.uses_mrope:
284
+ positions = input_batch.positions
285
+ else:
286
+ positions = input_batch.mrope_positions
287
with (
288
self.maybe_dummy_run_with_lora(
289
self.lora_config,
@@ -283,7 +299,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
299
):
300
hidden_states = self.model(
301
input_ids=input_batch.input_ids,
286
- positions=input_batch.positions,
302
+ positions=positions,
303
)
304
sample_hidden_states = hidden_states[input_batch.logits_indices]
305
return hidden_states, sample_hidden_states
@@ -393,8 +409,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
409
sampling_params=new_req_data.sampling_params,
410
lora_request=new_req_data.lora_request,
411
)
396
-
412
req_index = self.req_states.req_id_to_index[req_id]
413
+
414
+ # Pre-compute M-RoPE positions for prefill.
415
+ if self.uses_mrope:
416
+ self.mrope_states.init_prefill_mrope_positions(
417
+ req_index,
418
+ self.model, # type: ignore
419
+ new_req_data.prefill_token_ids,
420
+ mm_features=[], # TODO
421
+ )
422
+
423
self.block_tables.append_block_ids(
424
req_index, new_req_data.block_ids, overwrite=True
425
)
@@ -411,6 +436,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
436
437
self.req_states.apply_staged_writes()
438
self.block_tables.apply_staged_writes()
439
+ if self.uses_mrope:
440
+ self.mrope_states.apply_staged_writes()
441
442
def prepare_inputs(
443
self,
@@ -511,6 +538,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
538
)
539
seq_lens = self.input_buffers.seq_lens[:num_reqs]
540
541
+ # Prepare M-RoPE positions.
542
+ if self.uses_mrope:
543
+ self.mrope_states.prepare_mrope_positions(
544
+ idx_mapping,
545
+ query_start_loc,
546
+ self.req_states.prefill_len.gpu,
547
+ self.req_states.num_computed_tokens.gpu,
548
+ self.input_buffers.mrope_positions,
549
+ )
550
+
551
# Some input token ids are directly read from the last sampled tokens
552
# and draft tokens. Also, get the logits indices to sample tokens from.
553
logits_indices = combine_sampled_and_draft_tokens(
@@ -546,6 +583,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
583
584
input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
585
positions = self.input_buffers.positions[:num_tokens_after_padding]
586
+ mrope_positions = self.input_buffers.mrope_positions[
587
+ :, :num_tokens_after_padding
588
+ ]
589
return InputBatch(
590
req_ids=req_ids,
591
num_reqs=num_reqs,
@@ -561,6 +601,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
601
seq_lens=seq_lens,
602
input_ids=input_ids,
603
positions=positions,
604
+ mrope_positions=mrope_positions,
605
attn_metadata=attn_metadata,
606
logits_indices=logits_indices,
607
cu_num_logits=cu_num_logits,
@@ -889,6 +930,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
930
else:
931
# Run PyTorch model in eager mode.
932
# TODO(woosuk): Support piecewise CUDA graph.
933
+ if not self.uses_mrope:
934
+ positions = input_batch.positions
935
+ else:
936
+ positions = input_batch.mrope_positions
937
with set_forward_context(
938
input_batch.attn_metadata,
939
self.vllm_config,
@@ -898,7 +943,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
943
):
944
hidden_states = self.model(
945
input_ids=input_batch.input_ids,
901
- positions=input_batch.positions,
946
+ positions=positions,
947
)
948
949
self.execute_model_state = hidden_states, input_batch, sampling_metadata