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, 07:45 UTC
[Model Runner V2] Skip building deprecated fields in attn metadata (#32132)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 11, 2026 at 14:31 UTC
19504ac07fda211744bd67e62c03ab6b32c92ab1
5 files changed
+7
-39
vllm/v1/worker/gpu/attn_utils.py
+1
-7
@@ -3,7 +3,6 @@
3
from collections.abc import Sequence
4
from typing import Any, cast
5
6
-import numpy as np
6
import torch
7
8
from vllm.config import VllmConfig, get_layers_from_vllm_config
@@ -147,16 +146,13 @@ def build_attn_metadata(
146
query_start_loc_gpu: torch.Tensor,
147
query_start_loc_cpu: torch.Tensor,
148
seq_lens: torch.Tensor,
150
- seq_lens_np: np.ndarray,
151
- num_computed_tokens_cpu: torch.Tensor | None,
149
+ max_seq_len: int,
150
block_tables: Sequence[torch.Tensor],
151
slot_mappings: torch.Tensor,
152
kv_cache_config: KVCacheConfig,
153
) -> dict[str, Any]:
154
max_query_len = int(query_start_loc_cpu.max())
155
seq_lens = seq_lens[:num_reqs]
158
- seq_lens_cpu = torch.from_numpy(seq_lens_np)
159
- max_seq_len = int(seq_lens_np.max())
156
157
attn_metadata: dict[str, Any] = {}
158
kv_cache_groups = kv_cache_config.kv_cache_groups
@@ -168,9 +164,7 @@ def build_attn_metadata(
164
query_start_loc=query_start_loc_gpu,
165
query_start_loc_cpu=query_start_loc_cpu,
166
seq_lens=seq_lens,
171
- _seq_lens_cpu=seq_lens_cpu,
167
max_seq_len=max_seq_len,
173
- _num_computed_tokens_cpu=num_computed_tokens_cpu,
168
num_reqs=num_reqs,
169
num_actual_tokens=num_tokens,
170
max_query_len=max_query_len,
vllm/v1/worker/gpu/cudagraph_utils.py
+3
-6
@@ -232,11 +232,9 @@ def prepare_inputs_to_capture(
232
query_start_loc.np[: num_reqs + 1] = np.arange(num_reqs + 1) * num_tokens_per_req
233
query_start_loc.np[num_reqs:] = num_tokens
234
query_start_loc.copy_to_gpu()
235
- seq_lens_np = np.full(num_reqs, max_model_len, dtype=np.int32)
235
+
236
# HACK(woosuk): For faster warmup, we set seq_lens (GPU) to num_tokens
237
- # rather than max_model_len. This introduces a discrepancy between
238
- # seq_lens (on GPU) and seq_lens_np (on CPU), which may cause issues for
239
- # certain attention backends.
237
+ # rather than max_model_len.
238
input_buffers.seq_lens[:num_reqs] = num_tokens
239
input_buffers.seq_lens[num_reqs:] = 0
240
@@ -250,8 +248,7 @@ def prepare_inputs_to_capture(
248
query_start_loc_gpu=query_start_loc.gpu[: num_reqs + 1],
249
query_start_loc_cpu=query_start_loc.cpu[: num_reqs + 1],
250
seq_lens=input_buffers.seq_lens,
253
- seq_lens_np=seq_lens_np,
254
- num_computed_tokens_cpu=None, # FIXME
251
+ max_seq_len=max_model_len,
252
block_tables=input_block_tables,
253
slot_mappings=slot_mappings,
254
kv_cache_config=kv_cache_config,
vllm/v1/worker/gpu/input_batch.py
-4
@@ -70,7 +70,6 @@ class InputBatch:
70
query_start_loc_np: np.ndarray
71
# [num_reqs]
72
seq_lens: torch.Tensor
73
- seq_lens_np: np.ndarray
73
74
# [num_tokens_after_padding]
75
input_ids: torch.Tensor
@@ -109,8 +108,6 @@ class InputBatch:
108
query_start_loc_np = input_buffers.query_start_loc.np[: num_reqs + 1]
109
query_start_loc = input_buffers.query_start_loc.copy_to_gpu()[: num_reqs + 1]
110
# seq_len equals to query_len
112
- seq_lens_np = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
113
- seq_lens_np[-1] += num_tokens % num_reqs
111
input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs
112
input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs
113
input_buffers.seq_lens[num_reqs:] = 0
@@ -133,7 +130,6 @@ class InputBatch:
130
query_start_loc=query_start_loc,
131
query_start_loc_np=query_start_loc_np,
132
seq_lens=seq_lens,
136
- seq_lens_np=seq_lens_np,
133
input_ids=input_ids,
134
positions=positions,
135
attn_metadata=None, # type: ignore
vllm/v1/worker/gpu/model_runner.py
+2
-18
@@ -228,9 +228,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
228
slot_mappings = self.block_tables.get_dummy_slot_mappings(
229
input_batch.num_tokens
230
)
231
- num_computed_tokens = torch.zeros(
232
- input_batch.num_reqs, dtype=torch.int32, device=self.device
233
- )
231
query_start_loc = self.input_buffers.query_start_loc
232
query_start_loc_gpu = query_start_loc.gpu[: input_batch.num_reqs + 1]
233
query_start_loc_cpu = query_start_loc.cpu[: input_batch.num_reqs + 1]
@@ -241,8 +238,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
238
query_start_loc_gpu=query_start_loc_gpu,
239
query_start_loc_cpu=query_start_loc_cpu,
240
seq_lens=self.input_buffers.seq_lens,
244
- seq_lens_np=input_batch.seq_lens_np,
245
- num_computed_tokens_cpu=num_computed_tokens,
241
+ max_seq_len=self.max_model_len,
242
block_tables=block_tables,
243
slot_mappings=slot_mappings,
244
kv_cache_config=self.kv_cache_config,
@@ -522,16 +518,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
518
query_start_loc_gpu, self.input_buffers.positions[:num_tokens]
519
)
520
525
- # Get num_computed_tokens.
526
- # HACK(woosuk): Here, we use num_computed_tokens on GPU instead of
527
- # num_computed_tokens_cpu. This works for most cases.
528
- num_computed_tokens = self.req_states.num_computed_tokens[idx_mapping]
529
- # HACK(woosuk): Only GPU has the exact seq_lens because at this point
530
- # CPU does not know how many draft tokens are accepted/rejected in the
531
- # previous step. Therefore, we use max_model_len to be safe.
532
- # NOTE(woosuk): This only works for FA3 backend.
533
- seq_lens_np = np.full(num_reqs, self.max_model_len, dtype=np.int32)
534
-
521
# Layer name -> attention metadata.
522
attn_metadata = build_attn_metadata(
523
attn_metadata_builders=self.attn_metadata_builders,
@@ -540,8 +526,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
526
query_start_loc_gpu=query_start_loc_gpu,
527
query_start_loc_cpu=query_start_loc_cpu,
528
seq_lens=self.input_buffers.seq_lens,
543
- seq_lens_np=seq_lens_np,
544
- num_computed_tokens_cpu=num_computed_tokens,
529
+ max_seq_len=self.max_model_len,
530
block_tables=block_tables,
531
slot_mappings=slot_mappings,
532
kv_cache_config=self.kv_cache_config,
@@ -561,7 +546,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
546
query_start_loc=query_start_loc_gpu,
547
query_start_loc_np=query_start_loc_np,
548
seq_lens=seq_lens,
564
- seq_lens_np=seq_lens_np,
549
input_ids=input_ids,
550
positions=positions,
551
attn_metadata=attn_metadata,
vllm/v1/worker/gpu/spec_decode/eagle.py
+1
-4
@@ -288,8 +288,6 @@ class EagleSpeculator:
288
# Run eager mode.
289
query_start_loc.np[: num_reqs + 1] = np.arange(num_reqs + 1)
290
query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
291
- # HACK(woosuk)
292
- seq_lens_np = np.full(num_reqs, self.max_model_len, dtype=np.int32)
291
block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
292
293
# FIXME(woosuk): This is UNSAFE!!
@@ -300,8 +298,7 @@ class EagleSpeculator:
298
query_start_loc_gpu=query_start_loc_gpu,
299
query_start_loc_cpu=query_start_loc_cpu,
300
seq_lens=self.input_buffers.seq_lens[:num_reqs],
303
- seq_lens_np=seq_lens_np,
304
- num_computed_tokens_cpu=None, # FIXME
301
+ max_seq_len=self.max_model_len,
302
block_tables=block_tables,
303
slot_mappings=slot_mappings,
304
kv_cache_config=self.kv_cache_config,