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 10, 11:30 UTC
[Model Runner V2] Simplify BlockTables with UVA (#31965)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 8, 2026 at 10:24 UTC
750824324903f6dfc289d633ff5c513f16304f40
2 files changed
+52
-154
vllm/v1/worker/gpu/block_table.py
+46
-122
@@ -6,8 +6,9 @@ import torch
6
7
from vllm.triton_utils import tl, triton
8
from vllm.utils.math_utils import cdiv
9
+from vllm.utils.platform_utils import is_uva_available
10
+from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
11
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
10
-from vllm.v1.utils import CpuGpuBuffer
12
13
14
class BlockTables:
@@ -18,51 +19,53 @@ class BlockTables:
19
max_num_batched_tokens: int,
20
max_model_len: int,
21
device: torch.device,
21
- pin_memory: bool,
22
):
23
self.block_sizes = block_sizes
24
self.max_num_reqs = max_num_reqs
25
self.max_num_batched_tokens = max_num_batched_tokens
26
self.max_model_len = max_model_len
27
self.device = device
28
- self.pin_memory = pin_memory
28
+
29
+ if not is_uva_available():
30
+ raise RuntimeError("UVA is not available")
31
32
self.num_kv_cache_groups = len(self.block_sizes)
33
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
32
- self.block_tables: list[torch.Tensor] = []
34
+ self.block_tables: list[UvaBuffer] = []
35
for i in range(self.num_kv_cache_groups):
36
block_size = self.block_sizes[i]
37
max_num_blocks = cdiv(self.max_model_len, block_size)
36
- block_table = torch.zeros(
38
+ block_table = UvaBuffer(
39
self.max_num_reqs,
40
max_num_blocks,
41
dtype=torch.int32,
40
- device=self.device,
42
)
43
self.block_tables.append(block_table)
43
- self.block_table_ptrs = self._make_ptr_tensor(self.block_tables)
44
-
45
- # Block tables used for model's forward pass.
46
- # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
47
- self.input_block_tables: list[torch.Tensor] = [
48
- torch.zeros_like(block_table) for block_table in self.block_tables
49
- ]
50
- self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
51
-
44
+ self.block_table_ptrs = self._make_ptr_tensor(
45
+ [b.gpu for b in self.block_tables]
46
+ )
47
self.block_table_strides = torch.tensor(
53
- [b.stride(0) for b in self.block_tables],
48
+ [b.gpu.stride(0) for b in self.block_tables],
49
dtype=torch.int64,
50
device=self.device,
51
)
52
+
53
self.block_sizes_tensor = torch.tensor(
54
self.block_sizes, dtype=torch.int32, device=self.device
55
)
60
- self.num_blocks = torch.zeros(
56
+ self.num_blocks = UvaBuffer(
57
self.num_kv_cache_groups,
58
self.max_num_reqs,
59
dtype=torch.int32,
64
- device=self.device,
60
)
61
+
62
+ # Block tables used for model's forward pass.
63
+ # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
64
+ self.input_block_tables: list[torch.Tensor] = [
65
+ torch.zeros_like(b.gpu) for b in self.block_tables
66
+ ]
67
+ self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
68
+
69
self.slot_mappings = torch.zeros(
70
self.num_kv_cache_groups,
71
self.max_num_batched_tokens,
@@ -70,74 +73,36 @@ class BlockTables:
73
device=self.device,
74
)
75
73
- # Misc buffers.
74
- self.req_indices = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
75
- self.overwrite = self._make_buffer(self.max_num_reqs, dtype=torch.bool)
76
- self.cu_num_new_blocks = self._make_buffer(
77
- self.num_kv_cache_groups, self.max_num_reqs + 1, dtype=torch.int32
78
- )
79
-
80
- def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
81
- return CpuGpuBuffer(
82
- *args, dtype=dtype, pin_memory=self.pin_memory, device=self.device
83
- )
84
-
76
def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
77
# NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
78
ptrs_tensor_cpu = torch.tensor(
79
[t.data_ptr() for t in x],
80
dtype=torch.uint64,
81
device="cpu",
91
- pin_memory=self.pin_memory,
82
+ pin_memory=True,
83
)
84
return ptrs_tensor_cpu.to(self.device, non_blocking=True)
85
86
def append_block_ids(
87
self,
97
- # [num_reqs]
98
- req_indices: list[int],
99
- # [num_kv_cache_groups, num_reqs + 1]
100
- cu_num_new_blocks: tuple[list[int], ...],
101
- # [num_kv_cache_groups, num_new_blocks]
88
+ req_index: int,
89
new_block_ids: tuple[list[int], ...],
103
- # [num_reqs]
104
- overwrite: list[bool],
90
+ overwrite: bool,
91
) -> None:
106
- num_reqs = len(req_indices)
107
- self.req_indices.np[:num_reqs] = req_indices
108
- self.overwrite.np[:num_reqs] = overwrite
109
- for i in range(self.num_kv_cache_groups):
110
- self.cu_num_new_blocks.np[i, : num_reqs + 1] = cu_num_new_blocks[i]
111
-
112
- # NOTE(woosuk): Here, we cannot use a fixed-size buffer because there's
113
- # no clear upper bound to the number of new blocks in a single step.
114
- # NOTE(woosuk): The buffer has to be cached, because otherwise we cannot
115
- # guarantee that the buffer is not freed before the copy is completed.
116
- self.new_block_ids_cpu = torch.empty(
117
- self.num_kv_cache_groups,
118
- max(len(x) for x in new_block_ids),
119
- dtype=torch.int32,
120
- device="cpu",
121
- pin_memory=self.pin_memory,
122
- )
123
- new_block_ids_np = self.new_block_ids_cpu.numpy()
92
for i in range(self.num_kv_cache_groups):
125
- new_block_ids_np[i, : len(new_block_ids[i])] = new_block_ids[i]
126
- new_block_ids_gpu = self.new_block_ids_cpu.to(self.device, non_blocking=True)
127
-
128
- _append_block_ids_kernel[(self.num_kv_cache_groups, num_reqs)](
129
- self.req_indices.copy_to_gpu(num_reqs),
130
- self.cu_num_new_blocks.copy_to_gpu(),
131
- self.cu_num_new_blocks.gpu.stride(0),
132
- new_block_ids_gpu,
133
- new_block_ids_gpu.stride(0),
134
- self.overwrite.copy_to_gpu(num_reqs),
135
- self.block_table_strides,
136
- self.block_table_ptrs,
137
- self.num_blocks,
138
- self.num_blocks.stride(0),
139
- BLOCK_SIZE=1024, # type: ignore
140
- )
93
+ block_ids = new_block_ids[i]
94
+ num_new_blocks = len(block_ids)
95
+ if num_new_blocks == 0:
96
+ continue
97
+
98
+ # TODO(woosuk): Too many Numpy invocations. Optimize this.
99
+ start = self.num_blocks.np[i, req_index] if not overwrite else 0
100
+ end = start + num_new_blocks
101
+ if num_new_blocks == 1:
102
+ self.block_tables[i].np[req_index, start] = block_ids[0]
103
+ else:
104
+ self.block_tables[i].np[req_index, start:end] = block_ids
105
+ self.num_blocks.np[i, req_index] = end
106
107
def gather_block_tables(
108
self,
@@ -149,8 +114,8 @@ class BlockTables:
114
self.block_table_ptrs,
115
self.input_block_table_ptrs,
116
self.block_table_strides,
152
- self.num_blocks,
153
- self.num_blocks.stride(0),
117
+ self.num_blocks.gpu,
118
+ self.num_blocks.gpu.stride(0),
119
BLOCK_SIZE=1024, # type: ignore
120
)
121
return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
@@ -186,54 +151,6 @@ class BlockTables:
151
return self.slot_mappings[:, :num_tokens]
152
153
189
-@triton.jit
190
-def _append_block_ids_kernel(
191
- # Inputs
192
- req_indices, # [num_reqs]
193
- cu_num_new_blocks_ptr, # [num_kv_cache_groups, num_reqs + 1]
194
- cu_num_new_blocks_stride,
195
- new_block_ids_ptr, # [num_kv_cache_groups, num_new_blocks]
196
- new_block_ids_stride,
197
- overwrite, # [num_reqs]
198
- block_table_strides, # [num_kv_cache_groups]
199
- # Outputs
200
- block_table_ptrs, # [num_kv_cache_groups]
201
- num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs]
202
- num_blocks_stride,
203
- # Constants
204
- BLOCK_SIZE: tl.constexpr,
205
-):
206
- group_id = tl.program_id(0)
207
- batch_idx = tl.program_id(1)
208
- req_idx = tl.load(req_indices + batch_idx)
209
- do_overwrite = tl.load(overwrite + batch_idx)
210
-
211
- group_new_blocks_ptr = cu_num_new_blocks_ptr + group_id * cu_num_new_blocks_stride
212
- start_idx = tl.load(group_new_blocks_ptr + batch_idx)
213
- end_idx = tl.load(group_new_blocks_ptr + batch_idx + 1)
214
- num_new_blocks = end_idx - start_idx
215
-
216
- group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride
217
- dst_start_idx = tl.load(group_num_blocks_ptr + req_idx) if not do_overwrite else 0
218
- dst_end_idx = dst_start_idx + num_new_blocks
219
- tl.store(group_num_blocks_ptr + req_idx, dst_end_idx)
220
-
221
- # Destination
222
- block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
223
- block_table_stride = tl.load(block_table_strides + group_id)
224
- row_ptr = block_table_ptr + req_idx * block_table_stride
225
-
226
- group_new_block_ids_ptr = new_block_ids_ptr + group_id * new_block_ids_stride
227
- for i in range(0, num_new_blocks, BLOCK_SIZE):
228
- offset = i + tl.arange(0, BLOCK_SIZE)
229
- block_ids = tl.load(
230
- group_new_block_ids_ptr + start_idx + offset, mask=offset < num_new_blocks
231
- )
232
- tl.store(
233
- row_ptr + dst_start_idx + offset, block_ids, mask=offset < num_new_blocks
234
- )
235
-
236
-
154
@triton.jit
155
def _gather_block_tables_kernel(
156
batch_idx_to_req_idx, # [batch_size]
@@ -312,3 +229,10 @@ def _load_ptr(ptr_to_ptr, elem_dtype):
229
ptr = tl.load(ptr_to_ptr)
230
ptr = tl.cast(ptr, tl.pointer_type(elem_dtype))
231
return tl.multiple_of(ptr, 16)
232
+
233
+
234
+class UvaBuffer:
235
+ def __init__(self, *size, dtype: torch.dtype):
236
+ self.cpu = torch.zeros(*size, dtype=dtype, device="cpu", pin_memory=True)
237
+ self.np = self.cpu.numpy()
238
+ self.gpu = get_cuda_view_from_cpu_tensor(self.cpu)
vllm/v1/worker/gpu/model_runner.py
+6
-32
@@ -193,7 +193,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
193
max_num_batched_tokens=self.max_num_tokens,
194
max_model_len=self.max_model_len,
195
device=self.device,
196
- pin_memory=self.pin_memory,
196
)
197
198
self.attn_backends, self.attn_metadata_builders = init_attn_backend(
@@ -382,16 +381,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
381
for req_id in scheduler_output.finished_req_ids:
382
self.req_states.remove_request(req_id)
383
385
- # TODO(woosuk): Change SchedulerOutput.
386
- req_indices: list[int] = []
387
- cu_num_new_blocks = tuple(
388
- [0] for _ in range(self.block_tables.num_kv_cache_groups)
389
- )
390
- new_block_ids: tuple[list[int], ...] = tuple(
391
- [] for _ in range(self.block_tables.num_kv_cache_groups)
392
- )
393
- overwrite: list[bool] = []
394
-
384
# Add new requests.
385
for new_req_data in scheduler_output.scheduled_new_reqs:
386
assert new_req_data.prompt_token_ids is not None
@@ -408,12 +397,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
397
)
398
399
req_index = self.req_states.req_id_to_index[req_id]
411
- req_indices.append(req_index)
412
- for i, block_ids in enumerate(new_req_data.block_ids):
413
- x = cu_num_new_blocks[i][-1]
414
- cu_num_new_blocks[i].append(x + len(block_ids))
415
- new_block_ids[i].extend(block_ids)
416
- overwrite.append(True)
400
+ self.block_tables.append_block_ids(
401
+ req_index, new_req_data.block_ids, overwrite=True
402
+ )
403
if scheduler_output.scheduled_new_reqs:
404
self.req_states.prefill_len.copy_to_gpu()
405
@@ -421,23 +407,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
407
cached_reqs = scheduler_output.scheduled_cached_reqs
408
for i, req_id in enumerate(cached_reqs.req_ids):
409
req_index = self.req_states.req_id_to_index[req_id]
424
-
410
req_new_block_ids = cached_reqs.new_block_ids[i]
411
if req_new_block_ids is not None:
427
- req_indices.append(req_index)
428
- for group_id, block_ids in enumerate(req_new_block_ids):
429
- x = cu_num_new_blocks[group_id][-1]
430
- cu_num_new_blocks[group_id].append(x + len(block_ids))
431
- new_block_ids[group_id].extend(block_ids)
432
- overwrite.append(False)
433
-
434
- if req_indices:
435
- self.block_tables.append_block_ids(
436
- req_indices=req_indices,
437
- cu_num_new_blocks=cu_num_new_blocks,
438
- new_block_ids=new_block_ids,
439
- overwrite=overwrite,
440
- )
412
+ self.block_tables.append_block_ids(
413
+ req_index, req_new_block_ids, overwrite=False
414
+ )
415
416
def prepare_inputs(
417
self,