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 9, 14:45 UTC
[Model Runner V2] Fix slot_mapping after #25954 (#33046)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jan 25, 2026 at 18:29 UTC
edf927bc9f8dbb88b4ae1f37c8d4cea8d88b0c78
5 files changed
+40
-23
vllm/v1/worker/gpu/cudagraph_utils.py
+7
-7
@@ -101,9 +101,6 @@ class CudaGraphManager:
101
kv_cache_config,
102
)
103
num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
104
- slot_mappings_by_layer = build_slot_mappings_by_layer(
105
- slot_mappings, kv_cache_config
106
- )
104
105
# Warm up.
106
with set_forward_context(
@@ -112,7 +109,7 @@ class CudaGraphManager:
109
num_tokens=num_tokens,
110
cudagraph_runtime_mode=CUDAGraphMode.NONE,
111
num_tokens_across_dp=num_tokens_across_dp,
115
- slot_mapping=slot_mappings_by_layer,
112
+ slot_mapping=slot_mappings,
113
):
114
hidden_states = model(
115
input_ids=input_ids,
@@ -132,7 +129,7 @@ class CudaGraphManager:
129
num_tokens=num_tokens,
130
cudagraph_runtime_mode=CUDAGraphMode.NONE,
131
num_tokens_across_dp=num_tokens_across_dp,
135
- slot_mapping=slot_mappings_by_layer,
132
+ slot_mapping=slot_mappings,
133
),
134
torch.cuda.graph(graph, self.pool),
135
):
@@ -252,7 +249,7 @@ def prepare_inputs_to_capture(
249
attn_metadata_builders: list[AttentionMetadataBuilder],
250
max_model_len: int,
251
kv_cache_config: KVCacheConfig,
255
-) -> tuple[dict[str, Any], torch.Tensor]:
252
+) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
253
num_tokens_per_req = num_tokens // num_reqs
254
255
query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) * num_tokens_per_req
@@ -269,6 +266,9 @@ def prepare_inputs_to_capture(
266
267
input_block_tables = [x[:num_reqs] for x in block_tables.input_block_tables]
268
slot_mappings = block_tables.slot_mappings[:, :num_tokens]
269
+ slot_mappings_by_layer = build_slot_mappings_by_layer(
270
+ slot_mappings, kv_cache_config
271
+ )
272
273
attn_metadata = build_attn_metadata(
274
attn_metadata_builders=attn_metadata_builders,
@@ -282,4 +282,4 @@ def prepare_inputs_to_capture(
282
slot_mappings=slot_mappings,
283
kv_cache_config=kv_cache_config,
284
)
285
- return attn_metadata, slot_mappings
285
+ return attn_metadata, slot_mappings_by_layer
vllm/v1/worker/gpu/input_batch.py
+3
@@ -66,6 +66,8 @@ class InputBatch:
66
67
# layer_name -> Metadata
68
attn_metadata: dict[str, Any]
69
+ # layer_name -> slot_mapping
70
+ slot_mappings: dict[str, torch.Tensor]
71
72
# [total_num_logits]
73
logits_indices: torch.Tensor
@@ -133,6 +135,7 @@ class InputBatch:
135
mrope_positions=None,
136
inputs_embeds=None,
137
attn_metadata=None, # type: ignore
138
+ slot_mappings=None, # type: ignore
139
logits_indices=logits_indices,
140
cu_num_logits=cu_num_logits,
141
cu_num_logits_np=cu_num_logits_np,
vllm/v1/worker/gpu/model_runner.py
+11
-9
@@ -269,6 +269,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
269
slot_mappings = self.block_tables.get_dummy_slot_mappings(
270
input_batch.num_tokens
271
)
272
+ slot_mappings_by_layer = build_slot_mappings_by_layer(
273
+ slot_mappings, self.kv_cache_config
274
+ )
275
attn_metadata = build_attn_metadata(
276
attn_metadata_builders=self.attn_metadata_builders,
277
num_reqs=input_batch.num_reqs,
@@ -282,6 +285,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
285
kv_cache_config=self.kv_cache_config,
286
)
287
input_batch.attn_metadata = attn_metadata
288
+ input_batch.slot_mappings = slot_mappings_by_layer
289
290
@torch.inference_mode()
291
def _dummy_run(
@@ -345,6 +349,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
349
self.speculator.run_model(
350
self.max_num_tokens,
351
attn_metadata=None,
352
+ slot_mappings=None,
353
num_tokens_across_dp=num_tokens_across_dp,
354
)
355
torch.cuda.synchronize()
@@ -615,6 +620,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
620
query_start_loc,
621
self.input_buffers.positions[:num_tokens],
622
)
623
+ # Layer name -> slot mapping.
624
+ slot_mappings_by_layer = build_slot_mappings_by_layer(
625
+ slot_mappings, self.kv_cache_config
626
+ )
627
628
# Layer name -> attention metadata.
629
attn_metadata = build_attn_metadata(
@@ -655,6 +664,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
664
mrope_positions=mrope_positions,
665
inputs_embeds=None,
666
attn_metadata=attn_metadata,
667
+ slot_mappings=slot_mappings_by_layer,
668
logits_indices=logits_indices,
669
cu_num_logits=cu_num_logits,
670
cu_num_logits_np=cu_num_logits_np,
@@ -882,14 +892,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
892
if self.uses_mrope:
893
assert input_batch.mrope_positions is not None
894
positions = input_batch.mrope_positions
885
- slot_mappings = self.block_tables.compute_slot_mappings(
886
- input_batch.idx_mapping,
887
- input_batch.query_start_loc,
888
- input_batch.positions[: input_batch.num_tokens],
889
- )
890
- slot_mappings_by_layer = build_slot_mappings_by_layer(
891
- slot_mappings, self.kv_cache_config
892
- )
895
with set_forward_context(
896
input_batch.attn_metadata,
897
self.vllm_config,
@@ -897,7 +899,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
899
# TODO(woosuk): Support piecewise CUDA graph.
900
cudagraph_runtime_mode=CUDAGraphMode.NONE,
901
num_tokens_across_dp=num_tokens_across_dp,
900
- slot_mapping=slot_mappings_by_layer,
902
+ slot_mapping=input_batch.slot_mappings,
903
):
904
self.kv_connector.pre_forward(scheduler_output)
905
hidden_states = self.model(
vllm/v1/worker/gpu/spec_decode/eagle.py
+16
-4
@@ -13,7 +13,10 @@ from vllm.model_executor.model_loader import get_model
13
from vllm.triton_utils import tl, triton
14
from vllm.v1.attention.backend import AttentionMetadataBuilder
15
from vllm.v1.kv_cache_interface import KVCacheConfig
16
-from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
16
+from vllm.v1.worker.gpu.attn_utils import (
17
+ build_attn_metadata,
18
+ build_slot_mappings_by_layer,
19
+)
20
from vllm.v1.worker.gpu.block_table import BlockTables
21
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
22
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
@@ -108,7 +111,8 @@ class EagleSpeculator:
111
def run_model(
112
self,
113
num_tokens: int,
111
- attn_metadata: dict[str, Any],
114
+ attn_metadata: dict[str, Any] | None,
115
+ slot_mappings: dict[str, torch.Tensor] | None,
116
num_tokens_across_dp: torch.Tensor | None,
117
) -> tuple[torch.Tensor, torch.Tensor]:
118
with set_forward_context(
@@ -117,6 +121,7 @@ class EagleSpeculator:
121
num_tokens=num_tokens,
122
cudagraph_runtime_mode=CUDAGraphMode.NONE,
123
num_tokens_across_dp=num_tokens_across_dp,
124
+ slot_mapping=slot_mappings,
125
):
126
ret_hidden_states = self.model(
127
input_ids=self.input_buffers.input_ids[:num_tokens],
@@ -134,6 +139,7 @@ class EagleSpeculator:
139
self,
140
num_reqs: int,
141
attn_metadata: dict[str, Any],
142
+ slot_mappings: dict[str, torch.Tensor],
143
num_tokens_across_dp: torch.Tensor | None,
144
) -> None:
145
pos = self.input_buffers.positions[:num_reqs]
@@ -142,7 +148,7 @@ class EagleSpeculator:
148
for step in range(1, self.num_speculative_steps):
149
# Run the eagle model.
150
last_hidden_states, hidden_states = self.run_model(
145
- num_reqs, attn_metadata, num_tokens_across_dp
151
+ num_reqs, attn_metadata, slot_mappings, num_tokens_across_dp
152
)
153
logits = self.model.compute_logits(last_hidden_states)
154
@@ -235,6 +241,7 @@ class EagleSpeculator:
241
last_hidden_states, hidden_states = self.run_model(
242
num_tokens,
243
input_batch.attn_metadata,
244
+ input_batch.slot_mappings,
245
num_tokens_across_dp=None, # FIXME
246
)
247
sample_hidden_states = last_hidden_states[last_token_indices]
@@ -311,7 +318,12 @@ class EagleSpeculator:
318
slot_mappings=slot_mappings,
319
kv_cache_config=self.kv_cache_config,
320
)
314
- self.generate_draft(num_reqs, attn_metadata, num_tokens_across_dp=None) # FIXME
321
+ slot_mappings_by_layer = build_slot_mappings_by_layer(
322
+ slot_mappings, self.kv_cache_config
323
+ )
324
+ self.generate_draft(
325
+ num_reqs, attn_metadata, slot_mappings_by_layer, num_tokens_across_dp=None
326
+ ) # FIXME
327
return self.draft_tokens[:num_reqs]
328
329
vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py
+3
-3
@@ -69,7 +69,7 @@ class EagleCudaGraphManager:
69
kv_cache_config: KVCacheConfig,
70
) -> None:
71
num_reqs = min(num_tokens, self.max_num_reqs)
72
- attn_metadata = prepare_inputs_to_capture(
72
+ attn_metadata, slot_mappings = prepare_inputs_to_capture(
73
num_reqs,
74
num_tokens,
75
input_buffers,
@@ -81,13 +81,13 @@ class EagleCudaGraphManager:
81
num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
82
83
# Warm up.
84
- generate_fn(num_tokens, attn_metadata, num_tokens_across_dp)
84
+ generate_fn(num_tokens, attn_metadata, slot_mappings, num_tokens_across_dp)
85
86
# Capture the graph.
87
assert num_tokens not in self.graphs
88
graph = torch.cuda.CUDAGraph()
89
with torch.cuda.graph(graph, self.pool):
90
- generate_fn(num_tokens, attn_metadata, num_tokens_across_dp)
90
+ generate_fn(num_tokens, attn_metadata, slot_mappings, num_tokens_across_dp)
91
self.graphs[num_tokens] = graph
92
93
@torch.inference_mode()