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 8, 15:15 UTC
[ModelRunner V2][BugFix] Fix `max_query_len` calculation (#34167)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 9, 2026 at 13:47 UTC
e7e52781ff636bf772301c9282a7601c73b8b905
4 files changed
+6
-1
vllm/v1/worker/gpu/attn_utils.py
+1
-1
@@ -149,13 +149,13 @@ def build_attn_metadata(
149
num_tokens: int,
150
query_start_loc_gpu: torch.Tensor,
151
query_start_loc_cpu: torch.Tensor,
152
+ max_query_len: int,
153
seq_lens: torch.Tensor,
154
max_seq_len: int,
155
block_tables: Sequence[torch.Tensor],
156
slot_mappings: torch.Tensor,
157
kv_cache_config: KVCacheConfig,
158
) -> dict[str, Any]:
158
- max_query_len = int(query_start_loc_cpu.max())
159
seq_lens = seq_lens[:num_reqs]
160
161
attn_metadata: dict[str, Any] = {}
vllm/v1/worker/gpu/cudagraph_utils.py
+1
@@ -267,6 +267,7 @@ def prepare_inputs_to_capture(
267
num_tokens=num_tokens,
268
query_start_loc_gpu=query_start_loc,
269
query_start_loc_cpu=query_start_loc_cpu,
270
+ max_query_len=num_tokens_per_req,
271
seq_lens=input_buffers.seq_lens,
272
max_seq_len=max_model_len,
273
block_tables=input_block_tables,
vllm/v1/worker/gpu/model_runner.py
+3
@@ -274,6 +274,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
274
num_tokens=input_batch.num_tokens,
275
query_start_loc_gpu=input_batch.query_start_loc,
276
query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np),
277
+ max_query_len=input_batch.num_scheduled_tokens.max().item(),
278
seq_lens=input_batch.seq_lens,
279
max_seq_len=self.max_model_len,
280
block_tables=block_tables,
@@ -561,6 +562,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
562
query_start_loc_np = query_start_loc_np[: num_reqs + 1]
563
query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
564
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
565
+ max_query_len = num_scheduled_tokens.max().item()
566
567
# Get prefill tokens.
568
prepare_prefill_inputs(
@@ -624,6 +626,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
626
num_tokens=num_tokens,
627
query_start_loc_gpu=query_start_loc,
628
query_start_loc_cpu=query_start_loc_cpu,
629
+ max_query_len=max_query_len,
630
seq_lens=self.input_buffers.seq_lens,
631
max_seq_len=self.max_model_len,
632
block_tables=block_tables,
vllm/v1/worker/gpu/spec_decode/eagle.py
+1
@@ -301,6 +301,7 @@ class EagleSpeculator:
301
num_tokens=num_reqs,
302
query_start_loc_gpu=query_start_loc,
303
query_start_loc_cpu=query_start_loc_cpu,
304
+ max_query_len=1,
305
seq_lens=self.input_buffers.seq_lens[:num_reqs],
306
max_seq_len=self.max_model_len,
307
block_tables=block_tables,