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 13, 10:45 UTC
[V0 deprecation] Clean up num_prefill_tokens logic for V0 (#28203)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Canlin Guo committed
Nov 12, 2025 at 02:57 UTC
de120bc94f2e51633824093c626423ec8e7cb3a9
1 file changed
+3
-15
vllm/forward_context.py
+3
-15
@@ -5,7 +5,7 @@ import time
5
from collections import defaultdict
6
from contextlib import contextmanager
7
from dataclasses import dataclass
8
-from typing import TYPE_CHECKING, Any, NamedTuple, Union
8
+from typing import TYPE_CHECKING, Any, NamedTuple
9
10
import torch
11
@@ -185,18 +185,13 @@ class ForwardContext:
185
# copy from vllm_config.compilation_config.static_forward_context
186
no_compile_layers: dict[str, Any]
187
"""
188
- Type AttentionMetadata for v0,
188
Type Dict[str, AttentionMetadata] for v1, map from layer_name of each
189
attention layer to its attention metadata
190
Type List[Dict[str, AttentionMetadata]] for DBO. List of size two, one
191
for each microbatch.
192
Set dynamically for each forward pass
193
"""
195
- attn_metadata: Union[
196
- "AttentionMetadata",
197
- dict[str, "AttentionMetadata"],
198
- list[dict[str, "AttentionMetadata"]],
199
- ]
194
+ attn_metadata: dict[str, "AttentionMetadata"] | list[dict[str, "AttentionMetadata"]]
195
# TODO: remove after making all virtual_engines share the same kv cache
196
virtual_engine: int # set dynamically for each forward pass
197
# set dynamically for each forward pass
@@ -324,14 +319,7 @@ def set_forward_context(
319
finally:
320
global last_logging_time, batchsize_logging_interval
321
if need_to_track_batchsize:
327
- if hasattr(attn_metadata, "num_prefill_tokens"):
328
- # for v0 attention backends
329
- batchsize = (
330
- attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens
331
- )
332
- else:
333
- # for v1 attention backends
334
- batchsize = num_tokens
322
+ batchsize = num_tokens
323
# we use synchronous scheduling right now,
324
# adding a sync point here should not affect
325
# scheduling of the next batch