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 7, 16:30 UTC
[Model Runner V2] Prepare attn metadata in ModelState [2/N] (#35383)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 26, 2026 at 11:47 UTC
3d66502e1bf48d4ca92ca0d54f7c9bba39a8556c
4 files changed
+110
-92
vllm/v1/worker/gpu/input_batch.py
+3
-8
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
-from typing import Any
4
5
import numpy as np
6
import torch
@@ -60,6 +59,8 @@ class InputBatch:
59
query_start_loc_np: np.ndarray
60
# [num_reqs]
61
seq_lens: torch.Tensor
62
+ # [num_reqs]
63
+ dcp_local_seq_lens: torch.Tensor | None
64
65
# [num_tokens_after_padding]
66
input_ids: torch.Tensor
@@ -68,11 +69,6 @@ class InputBatch:
69
# [num_tokens_after_padding, hidden_size]
70
inputs_embeds: torch.Tensor | None
71
71
- # layer_name -> Metadata
72
- attn_metadata: dict[str, Any]
73
- # layer_name -> slot_mapping
74
- slot_mappings: dict[str, torch.Tensor]
75
-
72
# [total_num_logits]
73
logits_indices: torch.Tensor
74
# [num_reqs + 1]
@@ -139,11 +135,10 @@ class InputBatch:
135
query_start_loc=query_start_loc,
136
query_start_loc_np=query_start_loc_np,
137
seq_lens=seq_lens,
138
+ dcp_local_seq_lens=None,
139
input_ids=input_ids,
140
positions=positions,
141
inputs_embeds=None,
145
- attn_metadata=None, # type: ignore
146
- slot_mappings=None, # type: ignore
142
logits_indices=logits_indices,
143
cu_num_logits=cu_num_logits,
144
cu_num_logits_np=cu_num_logits_np,
vllm/v1/worker/gpu/model_runner.py
+72
-82
@@ -46,7 +46,6 @@ from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
46
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
47
from vllm.v1.worker.gpu.async_utils import AsyncOutput
48
from vllm.v1.worker.gpu.attn_utils import (
49
- build_attn_metadata,
49
build_slot_mappings_by_layer,
50
get_kv_cache_spec,
51
init_attn_backend,
@@ -317,31 +316,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
316
)
317
self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)
318
320
- def prepare_dummy_attn_metadata(self, input_batch: InputBatch) -> None:
321
- block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
322
- slot_mappings = self.block_tables.get_dummy_slot_mappings(
323
- input_batch.num_tokens
324
- )
325
- slot_mappings_by_layer = build_slot_mappings_by_layer(
326
- slot_mappings, self.kv_cache_config
327
- )
328
- attn_metadata = build_attn_metadata(
329
- attn_groups=self.attn_groups,
330
- num_reqs=input_batch.num_reqs,
331
- num_tokens=input_batch.num_tokens,
332
- query_start_loc_gpu=input_batch.query_start_loc,
333
- query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np),
334
- max_query_len=input_batch.num_scheduled_tokens.max().item(),
335
- seq_lens=input_batch.seq_lens,
336
- max_seq_len=self.max_model_len,
337
- block_tables=block_tables,
338
- slot_mappings=slot_mappings,
339
- kv_cache_config=self.kv_cache_config,
340
- dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
341
- )
342
- input_batch.attn_metadata = attn_metadata
343
- input_batch.slot_mappings = slot_mappings_by_layer
344
-
319
@torch.inference_mode()
320
def _dummy_run(
321
self, num_tokens: int, *args, skip_attn: bool = True, **kwargs
@@ -384,7 +358,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
358
return None, None
359
360
assert self.execute_model_state is not None
387
- hidden_states, _, input_batch, _ = self.execute_model_state
361
+ input_batch, _, _, _, hidden_states, _, _ = self.execute_model_state
362
self.execute_model_state = None
363
assert hidden_states is not None # Last PP rank always has hidden_states
364
sample_hidden_states = hidden_states[input_batch.logits_indices]
@@ -546,7 +520,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
520
self.encoder_runner.add_request(req_id, new_req_data.mm_features)
521
522
self.model_state.add_request(req_index, new_req_data)
549
-
523
self.block_tables.append_block_ids(
524
req_index, new_req_data.block_ids, overwrite=True
525
)
@@ -624,9 +597,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
597
idx_mapping, total_num_logits, cu_num_logits, max_expand_len
598
)
599
627
- # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
628
- block_tables = self.block_tables.gather_block_tables(idx_mapping)
629
-
600
# Get query_start_loc.
601
query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
602
query_start_loc_np[0] = 0
@@ -635,11 +605,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
605
# Some attention backends like FA3 require query_start_loc to be non-decreasing.
606
query_start_loc_np[num_reqs + 1 :] = num_tokens
607
async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
638
-
608
query_start_loc_np = query_start_loc_np[: num_reqs + 1]
640
- query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
609
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
642
- max_query_len = num_scheduled_tokens.max().item()
610
611
# Get prefill tokens if any.
612
if self.req_states.any_prefills(idx_mapping_np):
@@ -663,6 +630,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
630
)
631
seq_lens = self.input_buffers.seq_lens[:num_reqs]
632
633
+ dcp_local_seq_lens = None
634
if self.use_dcp:
635
# Prepare dcp local seq_lens.
636
prepare_dcp_local_seq_lens(
@@ -673,7 +641,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
641
self.dcp_rank,
642
self.cp_interleave,
643
)
676
- dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs]
644
+ dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs]
645
646
# Some input token ids are directly read from the last sampled tokens
647
# and draft tokens. Also, get the logits indices to sample tokens from.
@@ -689,35 +657,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
657
total_num_logits,
658
)
659
692
- # Compute slot mappings: [num_kv_cache_groups, num_tokens]
693
- slot_mappings = self.block_tables.compute_slot_mappings(
694
- idx_mapping,
695
- query_start_loc,
696
- self.input_buffers.positions[:num_tokens],
697
- )
698
- # Layer name -> slot mapping.
699
- slot_mappings_by_layer = build_slot_mappings_by_layer(
700
- slot_mappings, self.kv_cache_config
701
- )
702
-
703
- # Layer name -> attention metadata.
704
- attn_metadata = build_attn_metadata(
705
- attn_groups=self.attn_groups,
706
- num_reqs=num_reqs,
707
- num_tokens=num_tokens,
708
- query_start_loc_gpu=query_start_loc,
709
- query_start_loc_cpu=query_start_loc_cpu,
710
- max_query_len=max_query_len,
711
- seq_lens=self.input_buffers.seq_lens,
712
- max_seq_len=self.max_model_len,
713
- block_tables=block_tables,
714
- slot_mappings=slot_mappings,
715
- kv_cache_config=self.kv_cache_config,
716
- dcp_local_seq_lens=dcp_local_seq_lens,
717
- )
718
-
719
- input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
720
- positions = self.input_buffers.positions[:num_tokens_after_padding]
660
return InputBatch(
661
req_ids=req_ids,
662
num_reqs=num_reqs,
@@ -732,17 +671,38 @@ class GPUModelRunner(LoRAModelRunnerMixin):
671
query_start_loc=query_start_loc,
672
query_start_loc_np=query_start_loc_np,
673
seq_lens=seq_lens,
735
- input_ids=input_ids,
736
- positions=positions,
674
+ dcp_local_seq_lens=dcp_local_seq_lens,
675
+ input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
676
+ positions=self.input_buffers.positions[:num_tokens_after_padding],
677
inputs_embeds=None,
738
- attn_metadata=attn_metadata,
739
- slot_mappings=slot_mappings_by_layer,
678
logits_indices=logits_indices,
679
cu_num_logits=cu_num_logits,
680
cu_num_logits_np=cu_num_logits_np,
681
has_structured_output_reqs=scheduler_output.has_structured_output_requests,
682
)
683
684
+ def prepare_attn(
685
+ self, input_batch: InputBatch
686
+ ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
687
+ # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
688
+ block_tables = self.block_tables.gather_block_tables(input_batch.idx_mapping)
689
+ # Compute slot mappings: [num_kv_cache_groups, num_tokens]
690
+ slot_mappings = self.block_tables.compute_slot_mappings(
691
+ input_batch.idx_mapping,
692
+ input_batch.query_start_loc,
693
+ input_batch.positions,
694
+ )
695
+ return block_tables, slot_mappings
696
+
697
+ def prepare_dummy_attn(
698
+ self, input_batch: InputBatch
699
+ ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
700
+ block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
701
+ slot_mappings = self.block_tables.get_dummy_slot_mappings(
702
+ input_batch.num_tokens
703
+ )
704
+ return block_tables, slot_mappings
705
+
706
@torch.inference_mode()
707
def get_mm_embeddings(
708
self,
@@ -899,6 +859,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
859
input_batch = self.prepare_inputs(
860
scheduler_output, num_tokens_after_padding
861
)
862
+ block_tables, slot_mappings = self.prepare_attn(input_batch)
863
+
864
if self.lora_config:
865
# Activate LoRA adapters.
866
lora_inputs = self.lora_state.make_lora_inputs(
@@ -929,9 +891,28 @@ class GPUModelRunner(LoRAModelRunnerMixin):
891
device=self.device,
892
)
893
if not skip_attn_for_dummy_run:
932
- self.prepare_dummy_attn_metadata(input_batch)
894
+ block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
895
+ else:
896
+ block_tables = None
897
+ slot_mappings = None
898
# FIXME(woosuk): Fix warmup for LoRA.
899
900
+ attn_metadata = None
901
+ slot_mappings_by_layer = None
902
+ if not (dummy_run and skip_attn_for_dummy_run):
903
+ assert slot_mappings is not None
904
+ slot_mappings_by_layer = build_slot_mappings_by_layer(
905
+ slot_mappings, self.kv_cache_config
906
+ )
907
+ assert block_tables is not None
908
+ attn_metadata = self.model_state.prepare_attn(
909
+ input_batch,
910
+ block_tables,
911
+ slot_mappings,
912
+ self.attn_groups,
913
+ self.kv_cache_config,
914
+ )
915
+
916
model_inputs = {
917
"input_ids": input_batch.input_ids,
918
"positions": input_batch.positions,
@@ -968,13 +949,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
949
)
950
951
with set_forward_context(
971
- input_batch.attn_metadata,
952
+ attn_metadata,
953
self.vllm_config,
954
num_tokens=input_batch.num_tokens_after_padding,
955
cudagraph_runtime_mode=cudagraph_runtime_mode,
956
num_tokens_across_dp=num_tokens_across_dp,
957
batch_descriptor=batch_descriptor,
977
- slot_mapping=input_batch.slot_mappings,
958
+ slot_mapping=slot_mappings_by_layer,
959
):
960
self.kv_connector.pre_forward(scheduler_output)
961
model_output = self.model(**model_inputs)
@@ -985,22 +966,23 @@ class GPUModelRunner(LoRAModelRunnerMixin):
966
aux_hidden_states = None
967
968
kv_connector_output = self.kv_connector.post_forward(scheduler_output)
969
+ self.execute_model_state = (
970
+ input_batch,
971
+ model_inputs,
972
+ attn_metadata,
973
+ slot_mappings_by_layer,
974
+ hidden_states,
975
+ aux_hidden_states,
976
+ kv_connector_output,
977
+ )
978
979
if not self.is_last_pp_rank:
980
# Non-last PP rank: return IntermediateTensors for sending.
981
assert isinstance(hidden_states, IntermediateTensors)
982
hidden_states.kv_connector_output = kv_connector_output
993
- self.execute_model_state = (None, None, input_batch, kv_connector_output)
983
return hidden_states
995
-
984
# Last rank (or no PP): hidden_states is a tensor for sampling.
985
assert isinstance(hidden_states, torch.Tensor)
998
- self.execute_model_state = (
999
- hidden_states,
1000
- aux_hidden_states,
1001
- input_batch,
1002
- kv_connector_output,
1003
- )
986
return None
987
988
@torch.inference_mode()
@@ -1010,9 +992,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
992
if self.execute_model_state is None:
993
# The prior execute_model call must have failed.
994
return None
1013
- hidden_states, aux_hidden_states, input_batch, kv_connector_output = (
1014
- self.execute_model_state
1015
- )
995
+ (
996
+ input_batch,
997
+ model_inputs,
998
+ attn_metadata,
999
+ slot_mappings_by_layer,
1000
+ hidden_states,
1001
+ aux_hidden_states,
1002
+ kv_connector_output,
1003
+ ) = self.execute_model_state
1004
self.execute_model_state = None
1005
1006
if not self.is_last_pp_rank:
@@ -1075,6 +1063,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1063
if self.speculator is not None:
1064
draft_tokens = self.speculator.propose(
1065
input_batch,
1066
+ attn_metadata,
1067
+ slot_mappings_by_layer,
1068
hidden_states,
1069
aux_hidden_states,
1070
num_sampled,
vllm/v1/worker/gpu/model_states.py
+31
@@ -1,13 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Any
4
+
5
import torch
6
import torch.nn as nn
7
8
from vllm.config import VllmConfig
9
from vllm.v1.core.sched.output import NewRequestData
10
+from vllm.v1.kv_cache_interface import KVCacheConfig
11
+from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
12
from vllm.v1.worker.gpu.input_batch import InputBatch
13
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
14
from vllm.v1.worker.gpu.states import RequestState
15
+from vllm.v1.worker.utils import AttentionGroup
16
17
18
class ModelState:
@@ -72,3 +77,29 @@ class ModelState:
77
return {}
78
mrope_positions = self.mrope_state.mrope_positions[:, :num_tokens]
79
return {"positions": mrope_positions}
80
+
81
+ def prepare_attn(
82
+ self,
83
+ input_batch: InputBatch,
84
+ block_tables: tuple[torch.Tensor, ...],
85
+ slot_mappings: torch.Tensor,
86
+ attn_groups: list[list[AttentionGroup]],
87
+ kv_cache_config: KVCacheConfig,
88
+ ) -> dict[str, Any]:
89
+ query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
90
+ max_query_len = input_batch.num_scheduled_tokens.max().item()
91
+ attn_metadata = build_attn_metadata(
92
+ attn_groups=attn_groups,
93
+ num_reqs=input_batch.num_reqs,
94
+ num_tokens=input_batch.num_tokens,
95
+ query_start_loc_gpu=input_batch.query_start_loc,
96
+ query_start_loc_cpu=query_start_loc_cpu,
97
+ max_query_len=max_query_len,
98
+ seq_lens=input_batch.seq_lens,
99
+ max_seq_len=self.max_model_len,
100
+ block_tables=block_tables,
101
+ slot_mappings=slot_mappings,
102
+ kv_cache_config=kv_cache_config,
103
+ dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
104
+ )
105
+ return attn_metadata
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+4
-2
@@ -182,6 +182,8 @@ class EagleSpeculator:
182
def propose(
183
self,
184
input_batch: InputBatch,
185
+ attn_metadata: dict[str, Any],
186
+ slot_mappings: dict[str, torch.Tensor],
187
# [num_tokens, hidden_size]
188
last_hidden_states: torch.Tensor,
189
# num_layers x [num_tokens, hidden_size]
@@ -229,8 +231,8 @@ class EagleSpeculator:
231
# TODO(woosuk): Support CUDA graph for prefill.
232
last_hidden_states, hidden_states = self.run_model(
233
num_tokens,
232
- input_batch.attn_metadata,
233
- input_batch.slot_mappings,
234
+ attn_metadata,
235
+ slot_mappings,
236
num_tokens_across_dp=None, # FIXME
237
)
238
sample_hidden_states = last_hidden_states[last_token_indices]