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, 01:30 UTC
[Model Runner V2] Move mrope_positions buffer to MRopeState (#32532)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 17, 2026 at 20:09 UTC
4147910f1e893ba69aa86a210c73e02ae8a0dfde
4 files changed
+49
-34
vllm/v1/worker/gpu/cudagraph_utils.py
+7
-4
@@ -75,16 +75,17 @@ class CudaGraphManager:
75
num_tokens: int,
76
model: nn.Module,
77
input_buffers: InputBuffers,
78
+ mrope_positions: torch.Tensor | None,
79
block_tables: BlockTables,
80
attn_metadata_builders: list[AttentionMetadataBuilder],
81
kv_cache_config: KVCacheConfig,
82
) -> None:
83
num_reqs = min(num_tokens, self.max_num_reqs)
84
input_ids = input_buffers.input_ids[: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]
85
+ positions = input_buffers.positions[:num_tokens]
86
+ if self.uses_mrope:
87
+ assert mrope_positions is not None
88
+ positions = mrope_positions[:, :num_tokens]
89
attn_metadata = prepare_inputs_to_capture(
90
num_reqs,
91
num_tokens,
@@ -136,6 +137,7 @@ class CudaGraphManager:
137
self,
138
model: nn.Module,
139
input_buffers: InputBuffers,
140
+ mrope_positions: torch.Tensor | None,
141
block_tables: BlockTables,
142
attn_metadata_builders: list[AttentionMetadataBuilder],
143
kv_cache_config: KVCacheConfig,
@@ -146,6 +148,7 @@ class CudaGraphManager:
148
self.capture_graph,
149
model=model,
150
input_buffers=input_buffers,
151
+ mrope_positions=mrope_positions,
152
block_tables=block_tables,
153
attn_metadata_builders=attn_metadata_builders,
154
kv_cache_config=kv_cache_config,
vllm/v1/worker/gpu/input_batch.py
+2
-17
@@ -31,19 +31,6 @@ 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
-
34
35
@dataclass
36
class InputBatch:
@@ -76,7 +63,7 @@ class InputBatch:
63
# [num_tokens_after_padding]
64
positions: torch.Tensor
65
# [3, num_tokens_after_padding]
79
- mrope_positions: torch.Tensor
66
+ mrope_positions: torch.Tensor | None
67
68
# layer_name -> Metadata
69
attn_metadata: dict[str, Any]
@@ -124,8 +111,6 @@ class InputBatch:
111
112
input_ids = input_buffers.input_ids[:num_tokens].zero_()
113
positions = input_buffers.positions[:num_tokens].zero_()
127
- input_buffers.mrope_positions.zero_()
128
- mrope_positions = input_buffers.mrope_positions[:, :num_tokens]
114
115
# attn_metadata = defaultdict(lambda: None)
116
logits_indices = query_start_loc[1:] - 1
@@ -146,7 +131,7 @@ class InputBatch:
131
seq_lens=seq_lens,
132
input_ids=input_ids,
133
positions=positions,
149
- mrope_positions=mrope_positions,
134
+ mrope_positions=None,
135
attn_metadata=None, # type: ignore
136
logits_indices=logits_indices,
137
cu_num_logits=cu_num_logits,
vllm/v1/worker/gpu/mm/mrope_utils.py
+17
-3
@@ -11,10 +11,12 @@ 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
@@ -28,6 +30,19 @@ class MRopeState:
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,
@@ -58,12 +73,11 @@ class MRopeState:
73
query_start_loc: torch.Tensor,
74
prefill_lens: torch.Tensor,
75
num_computed_tokens: torch.Tensor,
61
- mrope_positions: torch.Tensor,
76
) -> None:
77
num_reqs = idx_mapping.shape[0]
78
_prepare_mrope_positions_kernel[(num_reqs,)](
65
- mrope_positions,
66
- mrope_positions.stride(0),
79
+ self.mrope_positions,
80
+ self.mrope_positions.stride(0),
81
self.prefill_mrope_positions.gpu,
82
self.prefill_mrope_positions.gpu.stride(0),
83
self.max_model_len,
vllm/v1/worker/gpu/model_runner.py
+23
-10
@@ -99,6 +99,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
99
if self.uses_mrope:
100
self.mrope_states = MRopeState(
101
max_num_reqs=self.max_num_reqs,
102
+ max_num_tokens=self.max_num_tokens,
103
max_model_len=self.max_model_len,
104
device=self.device,
105
)
@@ -284,15 +285,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
285
input_buffers=self.input_buffers,
286
device=self.device,
287
)
288
+ if self.uses_mrope:
289
+ input_batch.mrope_positions = self.mrope_states.mrope_positions[
290
+ :, :num_tokens
291
+ ]
292
if not skip_attn:
293
self.prepare_dummy_attn_metadata(input_batch)
294
295
dp_size = self.parallel_config.data_parallel_size
296
num_tokens_across_dp = make_num_tokens_across_dp(dp_size, num_tokens)
297
num_sampled_tokens = np.ones(input_batch.num_reqs, dtype=np.int32)
293
- if not self.uses_mrope:
294
- positions = input_batch.positions
295
- else:
298
+ positions = input_batch.positions
299
+ if self.uses_mrope:
300
positions = input_batch.mrope_positions
301
with (
302
self.maybe_dummy_run_with_lora(
@@ -371,9 +375,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
375
start_free_gpu_memory = torch.cuda.mem_get_info()[0]
376
377
with self.maybe_setup_dummy_loras(self.lora_config):
378
+ mrope_positions = None
379
+ if self.uses_mrope:
380
+ mrope_positions = self.mrope_states.mrope_positions
381
self.cudagraph_manager.capture(
382
model=self.model,
383
input_buffers=self.input_buffers,
384
+ mrope_positions=mrope_positions,
385
block_tables=self.block_tables,
386
attn_metadata_builders=self.attn_metadata_builders,
387
kv_cache_config=self.kv_cache_config,
@@ -566,7 +574,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
574
query_start_loc,
575
self.req_states.prefill_len.gpu,
576
self.req_states.num_computed_tokens.gpu,
569
- self.input_buffers.mrope_positions,
577
)
578
579
# Some input token ids are directly read from the last sampled tokens
@@ -604,9 +611,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
611
612
input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
613
positions = self.input_buffers.positions[:num_tokens_after_padding]
607
- mrope_positions = self.input_buffers.mrope_positions[
608
- :, :num_tokens_after_padding
609
- ]
614
+ mrope_positions = None
615
+ if self.uses_mrope:
616
+ mrope_positions = self.mrope_states.mrope_positions[
617
+ :, :num_tokens_after_padding
618
+ ]
619
return InputBatch(
620
req_ids=req_ids,
621
num_reqs=num_reqs,
@@ -936,6 +945,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
945
input_buffers=self.input_buffers,
946
device=self.device,
947
)
948
+ if self.uses_mrope:
949
+ input_batch.mrope_positions = self.mrope_states.mrope_positions[
950
+ :, :num_tokens_after_padding
951
+ ]
952
self.prepare_dummy_attn_metadata(input_batch)
953
954
# Run model.
@@ -949,9 +962,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
962
else:
963
# Run PyTorch model in eager mode.
964
# TODO(woosuk): Support piecewise CUDA graph.
952
- if not self.uses_mrope:
953
- positions = input_batch.positions
954
- else:
965
+ positions = input_batch.positions
966
+ if self.uses_mrope:
967
+ assert input_batch.mrope_positions is not None
968
positions = input_batch.mrope_positions
969
with set_forward_context(
970
input_batch.attn_metadata,