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 15, 12:15 UTC
[V1] Remove num_input_tokens from attn_metadata (#17193)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Apr 30, 2025 at 00:28 UTC
24e6ad3f16d59005cdfc4de6c7bdeb4359b5d21c
6 files changed
+14
-21
vllm/forward_context.py
+7
-9
@@ -74,15 +74,13 @@ def set_forward_context(attn_metadata: Any,
74
if vllm_config.parallel_config.data_parallel_size > 1:
75
dp_size = vllm_config.parallel_config.data_parallel_size
76
dp_rank = vllm_config.parallel_config.data_parallel_rank
77
- if attn_metadata is not None:
78
- if hasattr(attn_metadata, "num_prefill_tokens"):
79
- # for v0 attention backends
80
- batchsize = attn_metadata.num_prefill_tokens + \
81
- attn_metadata.num_decode_tokens
82
- else:
83
- # for v1 attention backends
84
- batchsize = attn_metadata.num_input_tokens
77
+ if attn_metadata is not None and hasattr(attn_metadata,
78
+ "num_prefill_tokens"):
79
+ # for v0 attention backends
80
+ batchsize = attn_metadata.num_prefill_tokens + \
81
+ attn_metadata.num_decode_tokens
82
else:
83
+ # for v1 attention backends or no attn_metadata
84
batchsize = num_tokens
85
num_tokens_across_dp = [0] * dp_size
86
num_tokens_across_dp[dp_rank] = batchsize
@@ -124,7 +122,7 @@ def set_forward_context(attn_metadata: Any,
122
attn_metadata.num_decode_tokens
123
else:
124
# for v1 attention backends
127
- batchsize = attn_metadata.num_input_tokens
125
+ batchsize = num_tokens
126
# we use synchronous scheduling right now,
127
# adding a sync point here should not affect
128
# scheduling of the next batch
vllm/v1/attention/backends/flash_attn.py
-3
@@ -94,9 +94,6 @@ class FlashAttentionMetadata:
94
scheduler_metadata: Optional[torch.Tensor] = None
95
prefix_scheduler_metadata: Optional[torch.Tensor] = None
96
97
- # For logging.
98
- num_input_tokens: int = 0 # Number of tokens including padding.
99
-
97
# for local attention
98
@dataclass
99
class LocalAttentionMetadata:
vllm/v1/attention/backends/flashinfer.py
-3
@@ -183,9 +183,6 @@ class FlashInferMetadata:
183
decode_wrapper: Optional[BatchDecodeWithPagedKVCacheWrapper] = None
184
cascade_wrapper: Optional[MultiLevelCascadeAttentionWrapper] = None
185
186
- # For logging.
187
- num_input_tokens: int = 0 # Number of tokens including padding.
188
-
186
@property
187
def query_start_loc(self):
188
# The GPUModelRunner expects to be able to access this property.
vllm/v1/attention/backends/mla/common.py
-3
@@ -312,9 +312,6 @@ class MLACommonMetadata(Generic[D]):
312
num_decode_tokens: int
313
num_prefills: int
314
315
- # For logging.
316
- num_input_tokens: int = 0 # Number of tokens including padding.
317
-
315
# The dimension of the attention heads
316
head_dim: Optional[int] = None
317
vllm/v1/worker/gpu_model_runner.py
+3
-2
@@ -1036,7 +1036,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1036
num_input_tokens = round_up(num_scheduled_tokens, tp_size)
1037
else:
1038
num_input_tokens = num_scheduled_tokens
1039
- attn_metadata.num_input_tokens = num_input_tokens
1039
1040
# _prepare_inputs may reorder the batch, so we must gather multi
1041
# modal outputs after that to ensure the correct order
@@ -1088,7 +1087,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1087
1088
# Run the decoder.
1089
# Use persistent buffers for CUDA graphs.
1091
- with set_forward_context(attn_metadata, self.vllm_config):
1090
+ with set_forward_context(attn_metadata,
1091
+ self.vllm_config,
1092
+ num_tokens=num_input_tokens):
1093
output = self.model(
1094
input_ids=input_ids,
1095
positions=positions,
vllm/v1/worker/tpu_model_runner.py
+4
-1
@@ -769,7 +769,10 @@ class TPUModelRunner:
769
xm.mark_step()
770
num_reqs = self.input_batch.num_reqs
771
# Run the decoder
772
- with set_forward_context(attn_metadata, self.vllm_config):
772
+ with set_forward_context(
773
+ attn_metadata,
774
+ self.vllm_config,
775
+ num_tokens=scheduler_output.total_num_scheduled_tokens):
776
hidden_states = self.model(
777
input_ids=input_ids,
778
positions=self.position_ids,