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, 21:00 UTC
[v1] Cleanup the BlockTable in InputBatch (#13977)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Mar 1, 2025 at 03:03 UTC
e7bd944e08c97b43b9f7a9b837c4e18e30f33c7e
5 files changed
+25
-17
tests/v1/worker/test_gpu_model_runner.py
+14
@@ -89,6 +89,17 @@ def _is_sampling_metadata_changed(model_runner,
89
sampling_metadata_before)
90
91
92
+def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
93
+ req_index = model_runner.input_batch.req_id_to_index[req_id]
94
+ block_table = model_runner.input_batch.block_table
95
+ req_state = model_runner.requests[req_id]
96
+ if block_table.num_blocks_per_row[req_index] != len(req_state.block_ids):
97
+ return False
98
+ num_blocks = block_table.num_blocks_per_row[req_index]
99
+ return (block_table.block_table_np[req_index, :num_blocks] ==
100
+ req_state.block_ids).all()
101
+
102
+
103
def test_update_states_new_request(model_runner):
104
req_id = "req_0"
105
@@ -100,6 +111,7 @@ def test_update_states_new_request(model_runner):
111
assert _is_sampling_metadata_changed(model_runner, metadata_before)
112
assert _is_req_added(model_runner, req_id)
113
assert _is_req_scheduled(model_runner, req_id)
114
+ assert _is_req_state_block_table_match(model_runner, req_id)
115
116
117
def test_update_states_request_finished(model_runner):
@@ -185,6 +197,7 @@ def test_update_states_request_resumed(model_runner):
197
assert _is_sampling_metadata_changed(model_runner, metadata_before)
198
assert _is_req_added(model_runner, req_id)
199
assert _is_req_scheduled(model_runner, req_id)
200
+ assert _is_req_state_block_table_match(model_runner, req_id)
201
202
203
def test_update_states_no_changes(model_runner):
@@ -215,6 +228,7 @@ def test_update_states_no_changes(model_runner):
228
assert not _is_sampling_metadata_changed(model_runner, metadata_before)
229
assert _is_req_added(model_runner, req_id)
230
assert _is_req_scheduled(model_runner, req_id)
231
+ assert _is_req_state_block_table_match(model_runner, req_id)
232
233
234
def test_update_states_request_unscheduled(model_runner):
vllm/v1/worker/block_table.py
+6
-7
@@ -15,13 +15,11 @@ class BlockTable:
15
def __init__(
16
self,
17
max_num_reqs: int,
18
- max_model_len: int,
18
max_num_blocks_per_req: int,
19
pin_memory: bool,
20
device: torch.device,
21
):
22
self.max_num_reqs = max_num_reqs
24
- self.max_model_len = max_model_len
23
self.max_num_blocks_per_req = max_num_blocks_per_req
24
self.pin_memory = pin_memory
25
self.device = device
@@ -42,18 +40,19 @@ class BlockTable:
40
41
def append_row(
42
self,
45
- row_idx: int,
46
- start: int,
43
block_ids: List[int],
44
+ row_idx: int,
45
) -> None:
46
if not block_ids:
47
return
48
num_blocks = len(block_ids)
49
+ start = self.num_blocks_per_row[row_idx]
50
+ self.num_blocks_per_row[row_idx] += num_blocks
51
self.block_table_np[row_idx, start:start + num_blocks] = block_ids
53
- self.num_blocks_per_row[row_idx] = start + num_blocks
52
55
- def add_row(self, row_idx: int, block_ids: List[int]) -> None:
56
- self.append_row(row_idx, 0, block_ids)
53
+ def add_row(self, block_ids: List[int], row_idx: int) -> None:
54
+ self.num_blocks_per_row[row_idx] = 0
55
+ self.append_row(block_ids, row_idx)
56
57
def move_row(self, src: int, tgt: int) -> None:
58
num_blocks = self.num_blocks_per_row[src]
vllm/v1/worker/gpu_input_batch.py
+1
-2
@@ -92,7 +92,6 @@ class InputBatch:
92
# Block table.
93
self.block_table = BlockTable(
94
max_num_reqs=max_num_reqs,
95
- max_model_len=max_model_len,
95
max_num_blocks_per_req=max_num_blocks_per_req,
96
pin_memory=pin_memory,
97
device=device,
@@ -249,7 +248,7 @@ class InputBatch:
248
self.num_tokens_no_spec[req_index] = request.num_tokens
249
250
self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
252
- self.block_table.add_row(req_index, request.block_ids)
251
+ self.block_table.add_row(request.block_ids, req_index)
252
253
sampling_params = request.sampling_params
254
if sampling_params.sampling_type == SamplingType.GREEDY:
vllm/v1/worker/gpu_model_runner.py
+2
-4
@@ -399,10 +399,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
399
# Update the persistent batch.
400
self.input_batch.num_computed_tokens_cpu[req_index] = (
401
num_computed_tokens)
402
- start_index = (len(req_state.block_ids) -
403
- len(req_data.new_block_ids))
404
- self.input_batch.block_table.append_row(req_index, start_index,
405
- req_data.new_block_ids)
402
+ self.input_batch.block_table.append_row(req_data.new_block_ids,
403
+ req_index)
404
# Add new_token_ids to token_ids_cpu.
405
start_token_index = num_computed_tokens
406
end_token_index = num_computed_tokens + len(req_data.new_token_ids)
vllm/v1/worker/tpu_model_runner.py
+2
-4
@@ -247,10 +247,8 @@ class TPUModelRunner:
247
# Update the persistent batch.
248
self.input_batch.num_computed_tokens_cpu[req_index] = (
249
req_data.num_computed_tokens)
250
- start_index = len(req_state.block_ids) - len(
251
- req_data.new_block_ids)
252
- self.input_batch.block_table.append_row(req_index, start_index,
253
- req_data.new_block_ids)
250
+ self.input_batch.block_table.append_row(req_data.new_block_ids,
251
+ req_index)
252
253
# Add the new or resumed requests to the persistent batch.
254
# The smaller empty indices are filled first.