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 6, 18:15 UTC
[Model Runner V2] Add Support for XD-RoPE (#36817)
Signed-off-by: Santino Ramos <elsantinoramos@gmail.com>
Santino Ramos committed
Mar 14, 2026 at 09:27 UTC
3ed46f374b17d98ca6f098e74cb7c5fd4146179c
5 files changed
+224
-165
vllm/v1/worker/gpu/cudagraph_utils.py
+3
@@ -320,6 +320,9 @@ class ModelCudaGraphManager(CudaGraphManager):
320
model_inputs = {
321
"input_ids": input_buffers.input_ids[:num_tokens],
322
"positions": input_buffers.positions[:num_tokens],
323
+ # TODO: Pass intermediate_tensors for PP CUDA graph
324
+ # support (https://github.com/vllm-project/vllm/pull/35162).
325
+ "intermediate_tensors": None,
326
**model_state.prepare_dummy_inputs(num_reqs, num_tokens),
327
}
328
model_output = model(**model_inputs)
vllm/v1/worker/gpu/mm/mrope_utils.py
deleted
-136
@@ -1,136 +0,0 @@
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_num_tokens: int,
15
- max_model_len: int,
16
- device: torch.device,
17
- ):
18
- self.max_num_reqs = max_num_reqs
19
- self.max_num_tokens = max_num_tokens
20
- self.max_model_len = max_model_len
21
- self.device = device
22
-
23
- # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
24
- # wasting a lot of CPU memory.
25
- self.prefill_mrope_positions = StagedWriteTensor(
26
- (max_num_reqs * 3, max_model_len),
27
- dtype=torch.int32,
28
- device=device,
29
- uva_instead_of_gpu=True,
30
- )
31
- self.prefill_mrope_delta = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
32
-
33
- # NOTE: `mrope_positions` is implemented with one additional dummy
34
- # position on purpose to make it non-contiguous so that it can work
35
- # with torch compile.
36
- # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
37
- # NOTE: When M-RoPE is enabled, position ids are 3D regardless of
38
- # the modality of inputs. For text-only inputs, each dimension has
39
- # identical position IDs, making M-RoPE functionally equivalent to
40
- # 1D-RoPE.
41
- # See page 5 of https://arxiv.org/abs/2409.12191
42
- self.mrope_positions = torch.zeros(
43
- (3, max_num_tokens + 1), dtype=torch.int64, device=device
44
- )
45
-
46
- def init_prefill_mrope_positions(
47
- self,
48
- req_idx: int,
49
- mrope_model: SupportsMRoPE,
50
- prefill_token_ids: list[int],
51
- mm_features: list,
52
- ) -> None:
53
- prefill_mrope_positions, prefill_mrope_delta = (
54
- mrope_model.get_mrope_input_positions(prefill_token_ids, mm_features)
55
- )
56
- for i in range(3):
57
- pos = prefill_mrope_positions[i].tolist()
58
- self.prefill_mrope_positions.stage_write(3 * req_idx + i, 0, pos)
59
- self.prefill_mrope_delta.np[req_idx] = prefill_mrope_delta
60
-
61
- def apply_staged_writes(self) -> None:
62
- self.prefill_mrope_positions.apply_write()
63
- self.prefill_mrope_delta.copy_to_uva()
64
-
65
- def prepare_mrope_positions(
66
- self,
67
- idx_mapping: torch.Tensor,
68
- query_start_loc: torch.Tensor,
69
- prefill_lens: torch.Tensor,
70
- num_computed_tokens: torch.Tensor,
71
- ) -> None:
72
- num_reqs = idx_mapping.shape[0]
73
- _prepare_mrope_positions_kernel[(num_reqs,)](
74
- self.mrope_positions,
75
- self.mrope_positions.stride(0),
76
- self.prefill_mrope_positions.gpu,
77
- 3 * self.max_model_len,
78
- self.max_model_len,
79
- self.prefill_mrope_delta.gpu,
80
- idx_mapping,
81
- query_start_loc,
82
- prefill_lens,
83
- num_computed_tokens,
84
- BLOCK_SIZE=1024,
85
- )
86
-
87
-
88
-@triton.jit
89
-def _prepare_mrope_positions_kernel(
90
- mrope_positions_ptr,
91
- mrope_positions_stride,
92
- prefill_mrope_positions_ptr,
93
- prefill_mrope_positions_stride0,
94
- prefill_mrope_positions_stride1,
95
- prefill_mrope_delta_ptr,
96
- idx_mapping_ptr,
97
- query_start_loc_ptr,
98
- prefill_lens_ptr,
99
- num_computed_tokens_ptr,
100
- BLOCK_SIZE: tl.constexpr,
101
-):
102
- batch_idx = tl.program_id(0)
103
- req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
104
-
105
- prefill_len = tl.load(prefill_lens_ptr + req_state_idx)
106
- num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
107
- is_prefill = num_computed < prefill_len
108
-
109
- query_start = tl.load(query_start_loc_ptr + batch_idx)
110
- query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
111
- query_len = query_end - query_start
112
-
113
- mrope_delta = tl.load(prefill_mrope_delta_ptr + req_state_idx)
114
- for i in range(0, query_len, BLOCK_SIZE):
115
- block = i + tl.arange(0, BLOCK_SIZE)
116
- mask = block < query_len
117
- orig_pos = num_computed + block
118
-
119
- for j in tl.static_range(3):
120
- if is_prefill:
121
- # Read from pre-computed M-RoPE positions.
122
- pos = tl.load(
123
- prefill_mrope_positions_ptr
124
- + req_state_idx * prefill_mrope_positions_stride0
125
- + j * prefill_mrope_positions_stride1
126
- + orig_pos,
127
- mask=mask,
128
- )
129
- else:
130
- # Apply M-RoPE delta.
131
- pos = orig_pos + mrope_delta
132
- tl.store(
133
- mrope_positions_ptr + j * mrope_positions_stride + query_start + block,
134
- pos,
135
- mask=mask,
136
- )
vllm/v1/worker/gpu/mm/rope.py
new
+197
@@ -0,0 +1,197 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import cast
4
+
5
+import torch
6
+import torch.nn as nn
7
+
8
+from vllm.config import ModelConfig
9
+from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsXDRoPE
10
+from vllm.triton_utils import tl, triton
11
+from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
12
+
13
+
14
+class RopeState:
15
+ """Unified state for multi-dimensional RoPE variants (M-RoPE, XD-RoPE).
16
+
17
+ M-RoPE: 3 dims, uses position delta for decode.
18
+ XD-RoPE: 3 or 4 dims, delta is 0 (decode uses orig_pos for all dims).
19
+
20
+ NOTE: `positions` is implemented with one additional dummy position on
21
+ purpose to make it non-contiguous so that it can work with torch compile.
22
+ See detailed explanation in
23
+ https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
24
+
25
+ NOTE: When M-RoPE is enabled, position ids are 3D regardless of the
26
+ modality of inputs. For text-only inputs, each dimension has identical
27
+ position IDs, making M-RoPE functionally equivalent to 1D-RoPE.
28
+ See page 5 of https://arxiv.org/abs/2409.12191
29
+ """
30
+
31
+ def __init__(
32
+ self,
33
+ num_dims: int,
34
+ has_delta: bool,
35
+ max_num_reqs: int,
36
+ max_num_tokens: int,
37
+ max_model_len: int,
38
+ device: torch.device,
39
+ ):
40
+ self.num_dims = num_dims
41
+ self.has_delta = has_delta
42
+ self.max_num_reqs = max_num_reqs
43
+ self.max_num_tokens = max_num_tokens
44
+ self.max_model_len = max_model_len
45
+ self.device = device
46
+
47
+ # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
48
+ # wasting a lot of CPU memory.
49
+ self.prefill_positions = StagedWriteTensor(
50
+ (max_num_reqs * num_dims, max_model_len),
51
+ dtype=torch.int32,
52
+ device=device,
53
+ uva_instead_of_gpu=True,
54
+ )
55
+ self.positions = torch.zeros(
56
+ (num_dims, max_num_tokens + 1), dtype=torch.int64, device=device
57
+ )
58
+
59
+ # Delta is non-zero for M-RoPE, always 0 for XD-RoPE.
60
+ self.prefill_delta = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
61
+
62
+ def init_prefill_positions(
63
+ self,
64
+ req_idx: int,
65
+ model: nn.Module,
66
+ prefill_token_ids: list[int],
67
+ mm_features: list,
68
+ ) -> None:
69
+ if self.has_delta:
70
+ mrope_model = cast(SupportsMRoPE, model)
71
+ prefill_positions, delta = mrope_model.get_mrope_input_positions(
72
+ prefill_token_ids, mm_features
73
+ )
74
+ self.prefill_delta.np[req_idx] = delta
75
+ else:
76
+ xdrope_model = cast(SupportsXDRoPE, model)
77
+ prefill_positions = xdrope_model.get_xdrope_input_positions(
78
+ prefill_token_ids, mm_features
79
+ )
80
+
81
+ for i in range(self.num_dims):
82
+ pos = prefill_positions[i].tolist()
83
+ self.prefill_positions.stage_write(self.num_dims * req_idx + i, 0, pos)
84
+
85
+ def apply_staged_writes(self) -> None:
86
+ self.prefill_positions.apply_write()
87
+ if self.has_delta:
88
+ self.prefill_delta.copy_to_uva()
89
+
90
+ def get_positions(self, num_tokens: int) -> torch.Tensor:
91
+ return self.positions[:, :num_tokens]
92
+
93
+ def prepare_positions(
94
+ self,
95
+ idx_mapping: torch.Tensor,
96
+ query_start_loc: torch.Tensor,
97
+ prefill_lens: torch.Tensor,
98
+ num_computed_tokens: torch.Tensor,
99
+ ) -> None:
100
+ num_reqs = idx_mapping.shape[0]
101
+ _prepare_rope_positions_kernel[(num_reqs,)](
102
+ self.positions,
103
+ self.positions.stride(0),
104
+ self.prefill_positions.gpu,
105
+ self.num_dims * self.max_model_len,
106
+ self.max_model_len,
107
+ self.prefill_delta.gpu,
108
+ idx_mapping,
109
+ query_start_loc,
110
+ prefill_lens,
111
+ num_computed_tokens,
112
+ BLOCK_SIZE=1024,
113
+ NUM_DIMS=self.num_dims,
114
+ )
115
+
116
+
117
+def get_rope_state(
118
+ model_config: ModelConfig,
119
+ model: nn.Module,
120
+ max_num_reqs: int,
121
+ max_num_tokens: int,
122
+ max_model_len: int,
123
+ device: torch.device,
124
+) -> RopeState | None:
125
+ """Create a RopeState if the model uses multi-dimensional RoPE."""
126
+ if model_config.uses_mrope:
127
+ assert isinstance(model, SupportsMRoPE)
128
+ return RopeState(
129
+ num_dims=3,
130
+ has_delta=True,
131
+ max_num_reqs=max_num_reqs,
132
+ max_num_tokens=max_num_tokens,
133
+ max_model_len=max_model_len,
134
+ device=device,
135
+ )
136
+ if model_config.uses_xdrope_dim > 0:
137
+ assert isinstance(model, SupportsXDRoPE)
138
+ return RopeState(
139
+ num_dims=model_config.uses_xdrope_dim,
140
+ has_delta=False,
141
+ max_num_reqs=max_num_reqs,
142
+ max_num_tokens=max_num_tokens,
143
+ max_model_len=max_model_len,
144
+ device=device,
145
+ )
146
+ return None
147
+
148
+
149
+@triton.jit
150
+def _prepare_rope_positions_kernel(
151
+ positions_ptr,
152
+ positions_stride,
153
+ prefill_positions_ptr,
154
+ prefill_positions_stride0,
155
+ prefill_positions_stride1,
156
+ prefill_delta_ptr,
157
+ idx_mapping_ptr,
158
+ query_start_loc_ptr,
159
+ prefill_lens_ptr,
160
+ num_computed_tokens_ptr,
161
+ BLOCK_SIZE: tl.constexpr,
162
+ NUM_DIMS: tl.constexpr,
163
+):
164
+ batch_idx = tl.program_id(0)
165
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
166
+
167
+ prefill_len = tl.load(prefill_lens_ptr + req_state_idx)
168
+ num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
169
+ is_prefill = num_computed < prefill_len
170
+
171
+ query_start = tl.load(query_start_loc_ptr + batch_idx)
172
+ query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
173
+ query_len = query_end - query_start
174
+
175
+ delta = tl.load(prefill_delta_ptr + req_state_idx)
176
+
177
+ for i in range(0, query_len, BLOCK_SIZE):
178
+ block = i + tl.arange(0, BLOCK_SIZE)
179
+ mask = block < query_len
180
+ orig_pos = num_computed + block
181
+
182
+ for j in tl.static_range(NUM_DIMS):
183
+ if is_prefill:
184
+ pos = tl.load(
185
+ prefill_positions_ptr
186
+ + req_state_idx * prefill_positions_stride0
187
+ + j * prefill_positions_stride1
188
+ + orig_pos,
189
+ mask=mask,
190
+ )
191
+ else:
192
+ pos = orig_pos + delta
193
+ tl.store(
194
+ positions_ptr + j * positions_stride + query_start + block,
195
+ pos,
196
+ mask=mask,
197
+ )
vllm/v1/worker/gpu/model_runner.py
+2
-1
@@ -992,6 +992,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
992
"input_ids": input_batch.input_ids,
993
"positions": input_batch.positions,
994
"inputs_embeds": inputs_embeds,
995
+ "intermediate_tensors": intermediate_tensors,
996
# NOTE: Values returned by `prepare_inputs` will override the default
997
# values above.
998
**self.model_state.prepare_inputs(input_batch, self.req_states),
@@ -1000,7 +1001,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1001
# Update for non-first PP ranks.
1002
model_inputs["input_ids"] = None
1003
model_inputs["inputs_embeds"] = None
1003
- model_inputs["intermediate_tensors"] = intermediate_tensors
1004
+ assert intermediate_tensors is not None
1005
1006
# Run model.
1007
if batch_desc.cg_mode == CUDAGraphMode.FULL:
vllm/v1/worker/gpu/model_states/default.py
+22
-28
@@ -13,7 +13,7 @@ from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
13
from vllm.v1.worker.gpu.input_batch import InputBatch
14
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
15
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
16
-from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
16
+from vllm.v1.worker.gpu.mm.rope import get_rope_state
17
from vllm.v1.worker.gpu.model_states.interface import ModelState
18
from vllm.v1.worker.gpu.states import RequestState
19
from vllm.v1.worker.utils import AttentionGroup
@@ -52,29 +52,28 @@ class DefaultModelState(ModelState):
52
device=self.device,
53
)
54
55
- self.uses_mrope = self.model_config.uses_mrope
56
- if self.uses_mrope:
57
- self.mrope_state = MRopeState(
58
- max_num_reqs=self.max_num_reqs,
59
- max_num_tokens=self.max_num_tokens,
60
- max_model_len=self.max_model_len,
61
- device=self.device,
62
- )
55
+ self.rope_state = get_rope_state(
56
+ self.model_config,
57
+ model,
58
+ max_num_reqs=self.max_num_reqs,
59
+ max_num_tokens=self.max_num_tokens,
60
+ max_model_len=self.max_model_len,
61
+ device=self.device,
62
+ )
63
64
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
65
- if self.uses_mrope:
66
- # Pre-compute M-RoPE positions for prefill.
65
+ if self.rope_state is not None:
66
assert new_req_data.prefill_token_ids is not None
68
- self.mrope_state.init_prefill_mrope_positions(
67
+ self.rope_state.init_prefill_positions(
68
req_index,
70
- self.model, # type: ignore
69
+ self.model,
70
new_req_data.prefill_token_ids,
71
mm_features=new_req_data.mm_features,
72
)
73
74
def apply_staged_writes(self) -> None:
76
- if self.uses_mrope:
77
- self.mrope_state.apply_staged_writes()
75
+ if self.rope_state is not None:
76
+ self.rope_state.apply_staged_writes()
77
78
def get_mm_embeddings(
79
self,
@@ -109,31 +108,26 @@ class DefaultModelState(ModelState):
108
109
def prepare_inputs(
110
self, input_batch: InputBatch, req_states: RequestState
112
- ) -> dict[str, Any]:
113
- if not self.uses_mrope:
114
- # Common case (1D positions).
115
- return {}
111
+ ) -> dict[str, torch.Tensor | None]:
112
+ if self.rope_state is None:
113
+ return {} # Common case (1D positions).
114
117
- # Prepare M-RoPE positions.
118
- self.mrope_state.prepare_mrope_positions(
115
+ self.rope_state.prepare_positions(
116
input_batch.idx_mapping,
117
input_batch.query_start_loc,
118
req_states.prefill_len.gpu,
119
req_states.num_computed_tokens.gpu,
120
)
124
- mrope_positions = self.mrope_state.mrope_positions[
125
- :, : input_batch.num_tokens_after_padding
126
- ]
127
- return {"positions": mrope_positions}
121
+ positions = self.rope_state.get_positions(input_batch.num_tokens_after_padding)
122
+ return {"positions": positions}
123
124
def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
125
model_inputs = {}
126
if self.supports_mm_inputs:
127
inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens]
128
model_inputs["inputs_embeds"] = inputs_embeds
134
- if self.uses_mrope:
135
- mrope_positions = self.mrope_state.mrope_positions[:, :num_tokens]
136
- model_inputs["positions"] = mrope_positions
129
+ if self.rope_state is not None:
130
+ model_inputs["positions"] = self.rope_state.get_positions(num_tokens)
131
return model_inputs
132
133
def prepare_attn(