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, 07:15 UTC
[Bugfix][TPU] Fix tpu model runner testcase failure (#18810)
Signed-off-by: Carol Zheng <cazheng@google.com>
Carol Zheng committed
May 30, 2025 at 10:04 UTC
fba02e3bd1775f28cb18390a9485460a9aabdeec
2 files changed
+50
-16
tests/v1/tpu/worker/test_tpu_model_runner.py
+26
-5
@@ -81,7 +81,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
81
mm_hashes=[],
82
mm_positions=[],
83
sampling_params=SamplingParams(),
84
- block_ids=[0],
84
+ block_ids=[[0]], # block_ids should be list[list[int]]
85
num_computed_tokens=0,
86
lora_request=None,
87
))
@@ -112,14 +112,35 @@ def _is_req_added(model_runner, req_id: str) -> bool:
112
113
114
def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
115
+ """Check if the request state block IDs match the block table.
116
+
117
+ This function handles both legacy BlockTable and new MultiGroupBlockTable
118
+ structures for backward compatibility.
119
+ """
120
+
121
req_index = model_runner.input_batch.req_id_to_index[req_id]
116
- block_table = model_runner.input_batch.block_table
122
+ multi_group_block_table = model_runner.input_batch.block_table
123
req_state = model_runner.requests[req_id]
118
- if block_table.num_blocks_per_row[req_index] != len(req_state.block_ids):
124
+
125
+ # Access the first block table from MultiGroupBlockTable
126
+ # This is safe since we currently only use single KV cache groups
127
+ block_table = multi_group_block_table[0]
128
+
129
+ # req_state.block_ids is now list[list[int]] for MultiGroupBlockTable
130
+ # Extract the first group's block IDs
131
+ if isinstance(req_state.block_ids[0], list):
132
+ # New format: list[list[int]] - extract first group
133
+ req_block_ids = req_state.block_ids[0]
134
+ else:
135
+ # Legacy format: list[int] - use directly
136
+ req_block_ids = req_state.block_ids
137
+
138
+ if block_table.num_blocks_per_row[req_index] != len(req_block_ids):
139
return False
140
+
141
num_blocks = block_table.num_blocks_per_row[req_index]
121
- return (block_table.block_table_np[req_index, :num_blocks] ==
122
- req_state.block_ids).all()
142
+ block_table_values = block_table.block_table_np[req_index, :num_blocks]
143
+ return (block_table_values == req_block_ids).all()
144
145
146
def test_update_states_new_request(model_runner):
vllm/v1/worker/tpu_model_runner.py
+24
-11
@@ -175,11 +175,21 @@ class TPUModelRunner(LoRAModelRunnerMixin):
175
self.kv_caches: list[torch.Tensor] = []
176
# req_id -> (input_id -> encoder_output)
177
self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}
178
- # self.input_batch: InputBatch # Persistent batch.
178
179
# Request states.
180
self.requests: dict[str, CachedRequestState] = {}
181
182
+ # Initialize input batch early to avoid AttributeError in _update_states
183
+ self.input_batch = InputBatch(
184
+ max_num_reqs=self.max_num_reqs,
185
+ max_model_len=self.max_model_len,
186
+ max_num_batched_tokens=self.max_num_tokens,
187
+ device=self.device,
188
+ pin_memory=self.pin_memory,
189
+ vocab_size=self.model_config.get_vocab_size(),
190
+ block_size=self.block_size,
191
+ )
192
+
193
# Cached torch/numpy tensor
194
# The pytorch tensor and numpy array share the same buffer.
195
# Sometimes the numpy op is faster so we create both.
@@ -1286,16 +1296,19 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1296
"Hybrid models with more than one KV cache type are not "
1297
"supported yet.")
1298
1289
- self.input_batch = InputBatch(
1290
- max_num_reqs=self.max_num_reqs,
1291
- max_model_len=self.max_model_len,
1292
- max_num_batched_tokens=self.max_num_tokens,
1293
- device=self.device,
1294
- pin_memory=self.pin_memory,
1295
- vocab_size=self.model_config.get_vocab_size(),
1296
- block_size=kv_cache_config.kv_cache_groups[0].kv_cache_spec.
1297
- block_size,
1298
- )
1299
+ if kv_cache_config.kv_cache_groups[
1300
+ 0].kv_cache_spec.block_size != self.block_size:
1301
+ self.input_batch = InputBatch(
1302
+ max_num_reqs=self.max_num_reqs,
1303
+ max_model_len=self.max_model_len,
1304
+ max_num_batched_tokens=self.max_num_tokens,
1305
+ device=self.device,
1306
+ pin_memory=self.pin_memory,
1307
+ vocab_size=self.model_config.get_vocab_size(),
1308
+ block_size=kv_cache_config.kv_cache_groups[0].kv_cache_spec.
1309
+ block_size,
1310
+ )
1311
+ # Verify dtype compatibility between block_table_cpu and input_batch
1312
assert self.block_table_cpu.dtype == self.input_batch.block_table[
1313
0].get_cpu_tensor().dtype
1314