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 14, 05:30 UTC
[Core] Use `CpuGpuBuffer` for block table tensors (#24795)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Sep 16, 2025 at 19:18 UTC
eeb135eb87277b388e1a2c97ff9c13fcd94fc6ae
6 files changed
+53
-63
tests/v1/tpu/worker/test_tpu_model_runner.py
+1
-1
@@ -125,7 +125,7 @@ def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
125
return False
126
127
num_blocks = block_table.num_blocks_per_row[req_index]
128
- block_table_values = block_table.block_table_np[req_index, :num_blocks]
128
+ block_table_values = block_table.block_table.np[req_index, :num_blocks]
129
return (block_table_values == req_block_ids).all()
130
131
tests/v1/worker/test_gpu_input_batch.py
+4
-1
@@ -15,6 +15,7 @@ from vllm.utils import is_pin_memory_available, make_tensor_with_pad
15
from vllm.v1.pool.metadata import PoolingMetadata
16
from vllm.v1.sample.logits_processor import LogitsProcessors
17
from vllm.v1.sample.metadata import SamplingMetadata
18
+from vllm.v1.utils import CpuGpuBuffer
19
from vllm.v1.worker.block_table import BlockTable, MultiGroupBlockTable
20
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
21
@@ -45,7 +46,7 @@ def _compare_objs(obj1,
46
47
is_same = False
48
if isinstance(a, torch.Tensor):
48
- if (a.numel() == 0 or b.numel() == 0):
49
+ if a.numel() == 0 or b.numel() == 0:
50
is_same = (a.numel() == 0 and b.numel() == 0)
51
elif torch.allclose(a, b):
52
is_same = True
@@ -61,6 +62,8 @@ def _compare_objs(obj1,
62
is_same = True # if we make it here must be same
63
elif a == b:
64
is_same = True
65
+ elif isinstance(a, CpuGpuBuffer):
66
+ is_same = np.allclose(a.np, b.np) and torch.allclose(a.gpu, b.gpu)
67
assert is_same, f"Attribute {attr_name} is different"\
68
f" in {obj1} and {obj2}: {a} != {b}"
69
tests/v1/worker/test_gpu_model_runner.py
+1
-1
@@ -165,7 +165,7 @@ def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
165
req_state.block_ids[0]):
166
return False
167
num_blocks = block_table.num_blocks_per_row[req_index]
168
- return (block_table.block_table_np[req_index, :num_blocks] ==
168
+ return (block_table.block_table.np[req_index, :num_blocks] ==
169
req_state.block_ids[0]).all()
170
171
vllm/v1/worker/block_table.py
+32
-43
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Union
4
5
import numpy as np
6
import torch
@@ -7,6 +8,7 @@ import torch
8
from vllm.distributed import get_dcp_group
9
from vllm.logger import init_logger
10
from vllm.utils import cdiv
11
+from vllm.v1.utils import CpuGpuBuffer
12
13
logger = init_logger(__name__)
14
@@ -29,28 +31,13 @@ class BlockTable:
31
self.pin_memory = pin_memory
32
self.device = device
33
32
- self.block_table = torch.zeros(
33
- (max_num_reqs, max_num_blocks_per_req),
34
- device=self.device,
35
- dtype=torch.int32,
36
- )
37
- self.block_table_cpu = torch.zeros(
38
- (max_num_reqs, max_num_blocks_per_req),
39
- device="cpu",
40
- dtype=torch.int32,
41
- pin_memory=pin_memory,
42
- )
43
- self.block_table_np = self.block_table_cpu.numpy()
34
+ self.block_table = self._make_buffer(max_num_reqs,
35
+ max_num_blocks_per_req,
36
+ dtype=torch.int32)
37
self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)
38
46
- self.slot_mapping_cpu = torch.zeros(self.max_num_batched_tokens,
47
- dtype=torch.int64,
48
- device="cpu",
49
- pin_memory=self.pin_memory)
50
- self.slot_mapping_np = self.slot_mapping_cpu.numpy()
51
- self.slot_mapping = torch.zeros(self.max_num_batched_tokens,
52
- dtype=torch.int64,
53
- device=self.device)
39
+ self.slot_mapping = self._make_buffer(self.max_num_batched_tokens,
40
+ dtype=torch.int64)
41
try:
42
self.dcp_world_size = get_dcp_group().world_size
43
self.dcp_rank = get_dcp_group().rank_in_group
@@ -69,7 +56,7 @@ class BlockTable:
56
num_blocks = len(block_ids)
57
start = self.num_blocks_per_row[row_idx]
58
self.num_blocks_per_row[row_idx] += num_blocks
72
- self.block_table_np[row_idx, start:start + num_blocks] = block_ids
59
+ self.block_table.np[row_idx, start:start + num_blocks] = block_ids
60
61
def add_row(self, block_ids: list[int], row_idx: int) -> None:
62
self.num_blocks_per_row[row_idx] = 0
@@ -77,17 +64,14 @@ class BlockTable:
64
65
def move_row(self, src: int, tgt: int) -> None:
66
num_blocks = self.num_blocks_per_row[src]
80
- self.block_table_np[tgt, :num_blocks] = self.block_table_np[
81
- src, :num_blocks]
67
+ block_table_np = self.block_table.np
68
+ block_table_np[tgt, :num_blocks] = block_table_np[src, :num_blocks]
69
self.num_blocks_per_row[tgt] = num_blocks
70
71
def swap_row(self, src: int, tgt: int) -> None:
85
- num_blocks_src = self.num_blocks_per_row[src]
86
- num_blocks_tgt = self.num_blocks_per_row[tgt]
87
- self.num_blocks_per_row[src] = num_blocks_tgt
88
- self.num_blocks_per_row[tgt] = num_blocks_src
89
-
90
- self.block_table_np[[src, tgt]] = self.block_table_np[[tgt, src]]
72
+ src_tgt, tgt_src = [src, tgt], [tgt, src]
73
+ self.num_blocks_per_row[src_tgt] = self.num_blocks_per_row[tgt_src]
74
+ self.block_table.np[src_tgt] = self.block_table.np[tgt_src]
75
76
def compute_slot_mapping(self, req_indices: np.ndarray,
77
positions: np.ndarray) -> None:
@@ -107,7 +91,7 @@ class BlockTable:
91
virtual_block_size = self.block_size * self.dcp_world_size
92
block_table_indices = (req_indices * self.max_num_blocks_per_req +
93
positions // virtual_block_size)
110
- block_numbers = self.block_table_np.ravel()[block_table_indices]
94
+ block_numbers = self.block_table.np.ravel()[block_table_indices]
95
# Use virtual_block_size for mask calculation, which marks local
96
# tokens.
97
virtual_block_offsets = positions % virtual_block_size
@@ -117,40 +101,45 @@ class BlockTable:
101
# Calculate slot_mapping
102
slot_mapping = block_numbers * self.block_size + block_offsets
103
# Write final slots, use -1 for not-local
120
- self.slot_mapping_np[:req_indices.shape[0]] = np.where(
104
+ self.slot_mapping.np[:req_indices.shape[0]] = np.where(
105
mask, slot_mapping, -1)
106
else:
107
block_table_indices = (req_indices * self.max_num_blocks_per_req +
108
positions // self.block_size)
125
- block_numbers = self.block_table_np.ravel()[block_table_indices]
109
+ block_numbers = self.block_table.np.ravel()[block_table_indices]
110
block_offsets = positions % self.block_size
111
np.add(block_numbers * self.block_size,
112
block_offsets,
129
- out=self.slot_mapping_np[:req_indices.shape[0]])
113
+ out=self.slot_mapping.np[:req_indices.shape[0]])
114
115
def commit_block_table(self, num_reqs: int) -> None:
132
- self.block_table[:num_reqs].copy_(self.block_table_cpu[:num_reqs],
133
- non_blocking=True)
116
+ self.block_table.copy_to_gpu(num_reqs)
117
118
def commit_slot_mapping(self, num_tokens: int) -> None:
136
- self.slot_mapping[:num_tokens].copy_(
137
- self.slot_mapping_cpu[:num_tokens], non_blocking=True)
119
+ self.slot_mapping.copy_to_gpu(num_tokens)
120
121
def clear(self) -> None:
140
- self.block_table.fill_(0)
141
- self.block_table_cpu.fill_(0)
122
+ self.block_table.gpu.fill_(0)
123
+ self.block_table.cpu.fill_(0)
124
143
- def get_device_tensor(self) -> torch.Tensor:
125
+ def get_device_tensor(self, num_reqs: int) -> torch.Tensor:
126
"""Returns the device tensor of the block table."""
145
- return self.block_table
127
+ return self.block_table.gpu[:num_reqs]
128
129
def get_cpu_tensor(self) -> torch.Tensor:
130
"""Returns the CPU tensor of the block table."""
149
- return self.block_table_cpu
131
+ return self.block_table.cpu
132
133
def get_numpy_array(self) -> np.ndarray:
134
"""Returns the numpy array of the block table."""
153
- return self.block_table_np
135
+ return self.block_table.np
136
+
137
+ def _make_buffer(self, *size: Union[int, torch.SymInt],
138
+ dtype: torch.dtype) -> CpuGpuBuffer:
139
+ return CpuGpuBuffer(*size,
140
+ dtype=dtype,
141
+ device=self.device,
142
+ pin_memory=self.pin_memory)
143
144
145
class MultiGroupBlockTable:
vllm/v1/worker/cpu_model_runner.py
+4
-4
@@ -89,7 +89,7 @@ class CPUModelRunner(GPUModelRunner):
89
assert isinstance(device_tensor, torch.Tensor)
90
setattr(obj, device_attr_name, cpu_tensor)
91
92
- for k, v in vars(self).items():
92
+ for v in vars(self).values():
93
if isinstance(v, CpuGpuBuffer):
94
v.gpu = v.cpu
95
@@ -98,9 +98,9 @@ class CPUModelRunner(GPUModelRunner):
98
replace_tensor(self.input_batch, k, k[:-11])
99
100
for block_table in self.input_batch.block_table.block_tables:
101
- for k, v in vars(block_table).items():
102
- if k.endswith("_cpu") and isinstance(v, torch.Tensor):
103
- replace_tensor(block_table, k, k[:-4])
101
+ for v in vars(block_table).values():
102
+ if isinstance(v, CpuGpuBuffer):
103
+ v.gpu = v.cpu
104
105
def load_model(self, eep_scale_up: bool = False) -> None:
106
logger.info("Starting to load model %s...", self.model_config.model)
vllm/v1/worker/gpu_model_runner.py
+11
-13
@@ -427,9 +427,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
427
*size: Union[int, torch.SymInt],
428
dtype: torch.dtype,
429
numpy: bool = True) -> CpuGpuBuffer:
430
- # Bfloat16 torch tensors cannot be directly cast to a numpy array, so
431
- # if a bfloat16 buffer is needed without a corresponding numpy array,
432
- # don't bother instantiating the numpy array.
430
return CpuGpuBuffer(*size,
431
dtype=dtype,
432
device=self.device,
@@ -1062,13 +1059,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1059
num_common_prefix_blocks = 0
1060
else:
1061
blk_table = self.input_batch.block_table[kv_cache_group_id]
1065
- blk_table_tensor = blk_table.get_device_tensor()[:num_reqs]
1066
- slot_mapping = blk_table.slot_mapping[:
1067
- total_num_scheduled_tokens]
1062
+ blk_table_tensor = blk_table.get_device_tensor(num_reqs)
1063
+ slot_mapping = blk_table.slot_mapping.gpu[:
1064
+ total_num_scheduled_tokens]
1065
1066
# Fill unused with -1. Needed for reshape_and_cache in full cuda
1067
# graph mode.
1071
- blk_table.slot_mapping[total_num_scheduled_tokens:].fill_(-1)
1068
+ blk_table.slot_mapping.gpu[total_num_scheduled_tokens:].fill_(
1069
+ -1)
1070
num_common_prefix_blocks = (
1071
scheduler_output.
1072
num_common_prefix_blocks[kv_cache_group_id])
@@ -2903,10 +2901,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2901
num_actual_tokens=num_tokens,
2902
max_query_len=max_query_len,
2903
max_seq_len=self.max_model_len,
2906
- block_table_tensor=self.input_batch.block_table[
2907
- kv_cache_group_id].get_device_tensor()[:num_reqs],
2908
- slot_mapping=self.input_batch.
2909
- block_table[kv_cache_group_id].slot_mapping[:num_tokens],
2904
+ block_table_tensor=self.input_batch.
2905
+ block_table[kv_cache_group_id].get_device_tensor(num_reqs),
2906
+ slot_mapping=self.input_batch.block_table[
2907
+ kv_cache_group_id].slot_mapping.gpu[:num_tokens],
2908
causal=True)
2909
for attn_group in self.attn_groups[kv_cache_group_id]:
2910
if ubatch_slices is not None:
@@ -3265,8 +3263,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3263
cudagraph_runtime_mode=cudagraph_runtime_mode,
3264
uniform_decode=False)
3265
3268
- # Capture full cudagraph for uniform decode batches if we have
3269
- # dont already have full mixed prefill-decode cudagraphs
3266
+ # Capture full cudagraph for uniform decode batches if we
3267
+ # don't already have full mixed prefill-decode cudagraphs.
3268
if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and \
3269
cudagraph_mode.separate_routine():
3270
max_num_tokens = self.scheduler_config.max_num_seqs * \