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 12, 12:15 UTC
[Model Runner V2] Minor cleanup for build_attn_metadata (#29576)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 26, 2025 at 20:10 UTC
ee80aee1cab6f0b6893cf54c8aa2f2c23512ec82
3 files changed
+13
-8
vllm/v1/worker/gpu/attn_utils.py
+3
-5
@@ -18,7 +18,6 @@ from vllm.v1.kv_cache_interface import (
18
KVCacheConfig,
19
KVCacheSpec,
20
)
21
-from vllm.v1.utils import CpuGpuBuffer
21
from vllm.v1.worker.utils import bind_kv_cache
22
23
@@ -145,7 +144,8 @@ def build_attn_metadata(
144
attn_metadata_builders: list[AttentionMetadataBuilder],
145
num_reqs: int,
146
num_tokens: int,
148
- query_start_loc: CpuGpuBuffer,
147
+ query_start_loc_gpu: torch.Tensor,
148
+ query_start_loc_cpu: torch.Tensor,
149
seq_lens: torch.Tensor,
150
seq_lens_np: np.ndarray,
151
num_computed_tokens_cpu: torch.Tensor | None,
@@ -153,9 +153,7 @@ def build_attn_metadata(
153
slot_mappings: torch.Tensor,
154
kv_cache_config: KVCacheConfig,
155
) -> dict[str, Any]:
156
- query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
157
- query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
158
- max_query_len = int(query_start_loc.np[: num_reqs + 1].max())
156
+ max_query_len = int(query_start_loc_cpu.max())
157
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())
vllm/v1/worker/gpu/cudagraph_utils.py
+2
-1
@@ -120,7 +120,8 @@ class CudaGraphManager:
120
attn_metadata_builders=attn_metadata_builders,
121
num_reqs=batch_size,
122
num_tokens=batch_size,
123
- query_start_loc=input_buffers.query_start_loc,
123
+ query_start_loc_gpu=input_buffers.query_start_loc.gpu[: batch_size + 1],
124
+ query_start_loc_cpu=input_buffers.query_start_loc.cpu[: batch_size + 1],
125
seq_lens=input_buffers.seq_lens,
126
seq_lens_np=np.full(batch_size, self.max_model_len, dtype=np.int32),
127
num_computed_tokens_cpu=None, # FIXME
vllm/v1/worker/gpu/model_runner.py
+8
-2
@@ -226,11 +226,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
226
num_computed_tokens = torch.zeros(
227
input_batch.num_reqs, dtype=torch.int32, device=self.device
228
)
229
+ query_start_loc = self.input_buffers.query_start_loc
230
+ query_start_loc_gpu = query_start_loc.gpu[: input_batch.num_reqs + 1]
231
+ query_start_loc_cpu = query_start_loc.cpu[: input_batch.num_reqs + 1]
232
attn_metadata = build_attn_metadata(
233
attn_metadata_builders=self.attn_metadata_builders,
234
num_reqs=input_batch.num_reqs,
235
num_tokens=input_batch.num_tokens,
233
- query_start_loc=self.input_buffers.query_start_loc,
236
+ query_start_loc_gpu=query_start_loc_gpu,
237
+ query_start_loc_cpu=query_start_loc_cpu,
238
seq_lens=self.input_buffers.seq_lens,
239
seq_lens_np=input_batch.seq_lens_np,
240
num_computed_tokens_cpu=num_computed_tokens,
@@ -515,6 +519,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
519
self.input_buffers.query_start_loc.np[num_reqs + 1 :] = num_tokens
520
self.input_buffers.query_start_loc.copy_to_gpu()
521
query_start_loc_gpu = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
522
+ query_start_loc_cpu = self.input_buffers.query_start_loc.cpu[: num_reqs + 1]
523
query_start_loc_np = self.input_buffers.query_start_loc.np[: num_reqs + 1]
524
525
# Copy prefill tokens from CPU to GPU.
@@ -572,7 +577,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
577
attn_metadata_builders=self.attn_metadata_builders,
578
num_reqs=num_reqs,
579
num_tokens=num_tokens,
575
- query_start_loc=self.input_buffers.query_start_loc,
580
+ query_start_loc_gpu=query_start_loc_gpu,
581
+ query_start_loc_cpu=query_start_loc_cpu,
582
seq_lens=self.input_buffers.seq_lens,
583
seq_lens_np=seq_lens_np,
584
num_computed_tokens_cpu=num_computed_tokens,