A high-throughput and memory-efficient inference and serving engine for LLMs
[PERF] Fuse multi-group block table staged writes (#44944)
Signed-off-by: jesse <szxfml@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Song Zhixin committed
Jun 17, 2026 at 01:53 UTC
8e27a9c215560ad700036ad4b4d32676d61ef4b1
3 files changed
+242
-17
tests/v1/worker/test_gpu_block_table.py
new
+132
@@ -0,0 +1,132 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.v1.worker.gpu.block_table import BlockTables
9
+
10
+pytestmark = pytest.mark.skipif(
11
+ not current_platform.is_cuda(),
12
+ reason="requires CUDA",
13
+)
14
+
15
+
16
+def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch):
17
+ device = torch.device("cuda")
18
+ block_tables = BlockTables(
19
+ block_sizes=[16, 32, 8],
20
+ max_num_reqs=4,
21
+ max_num_batched_tokens=64,
22
+ max_num_blocks_per_group=[8, 8, 8],
23
+ device=device,
24
+ kernel_block_sizes=[16, 16, 8],
25
+ )
26
+
27
+ def fail_if_apply_write_called():
28
+ pytest.fail("multi-group writes should use the fused apply kernel")
29
+
30
+ for block_table in block_tables.block_tables:
31
+ monkeypatch.setattr(block_table, "apply_write", fail_if_apply_write_called)
32
+
33
+ block_tables.append_block_ids(
34
+ req_index=0,
35
+ new_block_ids=([1, 2], [10, 11], []),
36
+ overwrite=True,
37
+ )
38
+ block_tables.append_block_ids(
39
+ req_index=1,
40
+ new_block_ids=([3], [12], [5, 6]),
41
+ overwrite=True,
42
+ )
43
+ block_tables.apply_staged_writes()
44
+ torch.accelerator.synchronize()
45
+
46
+ assert torch.equal(
47
+ block_tables.block_tables[0].gpu[0, :2],
48
+ torch.tensor([1, 2], dtype=torch.int32, device=device),
49
+ )
50
+ # Group 1 has blocks_per_kv_block == 2, so each KV block expands to two
51
+ # kernel block IDs.
52
+ assert torch.equal(
53
+ block_tables.block_tables[1].gpu[0, :4],
54
+ torch.tensor([20, 21, 22, 23], dtype=torch.int32, device=device),
55
+ )
56
+ assert torch.equal(
57
+ block_tables.block_tables[0].gpu[1, :1],
58
+ torch.tensor([3], dtype=torch.int32, device=device),
59
+ )
60
+ assert torch.equal(
61
+ block_tables.block_tables[1].gpu[1, :2],
62
+ torch.tensor([24, 25], dtype=torch.int32, device=device),
63
+ )
64
+ assert torch.equal(
65
+ block_tables.block_tables[2].gpu[1, :2],
66
+ torch.tensor([5, 6], dtype=torch.int32, device=device),
67
+ )
68
+ assert block_tables.num_blocks.np[0, 0] == 2
69
+ assert block_tables.num_blocks.np[1, 0] == 4
70
+ assert block_tables.num_blocks.np[2, 0] == 0
71
+ assert block_tables.num_blocks.np[0, 1] == 1
72
+ assert block_tables.num_blocks.np[1, 1] == 2
73
+ assert block_tables.num_blocks.np[2, 1] == 2
74
+ assert torch.equal(
75
+ block_tables.num_blocks.gpu[:, :2],
76
+ torch.tensor([[2, 1], [4, 2], [0, 2]], dtype=torch.int32, device=device),
77
+ )
78
+
79
+ for block_table in block_tables.block_tables:
80
+ assert not block_table._staged_write_indices
81
+ assert not block_table._staged_write_starts
82
+ assert not block_table._staged_write_contents
83
+ assert not block_table._staged_write_cu_lens
84
+
85
+ block_tables.append_block_ids(
86
+ req_index=0,
87
+ new_block_ids=([7], [13], [8]),
88
+ overwrite=False,
89
+ )
90
+ block_tables.apply_staged_writes()
91
+ torch.accelerator.synchronize()
92
+
93
+ assert torch.equal(
94
+ block_tables.block_tables[0].gpu[0, :3],
95
+ torch.tensor([1, 2, 7], dtype=torch.int32, device=device),
96
+ )
97
+ assert torch.equal(
98
+ block_tables.block_tables[1].gpu[0, :6],
99
+ torch.tensor([20, 21, 22, 23, 26, 27], dtype=torch.int32, device=device),
100
+ )
101
+ assert torch.equal(
102
+ block_tables.block_tables[2].gpu[0, :1],
103
+ torch.tensor([8], dtype=torch.int32, device=device),
104
+ )
105
+ assert block_tables.num_blocks.np[0, 0] == 3
106
+ assert block_tables.num_blocks.np[1, 0] == 6
107
+ assert block_tables.num_blocks.np[2, 0] == 1
108
+
109
+
110
+def test_block_tables_apply_staged_writes_single_group():
111
+ device = torch.device("cuda")
112
+ block_tables = BlockTables(
113
+ block_sizes=[16],
114
+ max_num_reqs=2,
115
+ max_num_batched_tokens=16,
116
+ max_num_blocks_per_group=[4],
117
+ device=device,
118
+ kernel_block_sizes=[16],
119
+ )
120
+
121
+ block_tables.append_block_ids(
122
+ req_index=0,
123
+ new_block_ids=([1, 2],),
124
+ overwrite=True,
125
+ )
126
+ block_tables.apply_staged_writes()
127
+ torch.accelerator.synchronize()
128
+
129
+ assert torch.equal(
130
+ block_tables.block_tables[0].gpu[0, :2],
131
+ torch.tensor([1, 2], dtype=torch.int32, device=device),
132
+ )
vllm/v1/worker/gpu/block_table.py
+21
-12
@@ -6,7 +6,12 @@ import torch
6
7
from vllm.triton_utils import tl, triton
8
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
9
-from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
9
+from vllm.v1.worker.gpu.buffer_utils import (
10
+ FusedStagedWriter,
11
+ StagedWriteTensor,
12
+ UvaBackedTensor,
13
+ _load_ptr,
14
+)
15
16
17
class BlockTables:
@@ -52,6 +57,12 @@ class BlockTables:
57
(self.num_kv_cache_groups, self.max_num_reqs),
58
dtype=torch.int32,
59
)
60
+ self.fused_writer: FusedStagedWriter | None = None
61
+ if self.num_kv_cache_groups > 1:
62
+ # Only the multi-group path uses the fused writer.
63
+ self.fused_writer = FusedStagedWriter(
64
+ self.device, self.num_kv_cache_groups * self.max_num_reqs
65
+ )
66
67
# Block tables used for model's forward pass.
68
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
@@ -109,10 +120,15 @@ class BlockTables:
120
self.num_blocks.np[i, req_index] = start + len(block_ids)
121
122
def apply_staged_writes(self) -> None:
112
- # TODO(woosuk): This can be inefficient since it launches one kernel per
113
- # block table. Implement a kernel to handle all block tables at once.
114
- for block_table in self.block_tables:
115
- block_table.apply_write()
123
+ if self.num_kv_cache_groups == 1:
124
+ # Single group: write directly, skipping the per-write group lookup.
125
+ self.block_tables[0].apply_write()
126
+ else:
127
+ # Multiple groups: apply all block tables with one fused kernel.
128
+ assert self.fused_writer is not None
129
+ self.fused_writer.apply(
130
+ self.block_tables, self.block_table_ptrs, self.block_table_strides
131
+ )
132
self.num_blocks.copy_to_uva()
133
134
def gather_block_tables(
@@ -283,10 +299,3 @@ def _compute_slot_mappings_kernel(
299
slot_ids = tl.where(is_local, slot_ids, PAD_ID)
300
301
tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
286
-
287
-
288
-@triton.jit
289
-def _load_ptr(ptr_to_ptr, elem_dtype):
290
- ptr = tl.load(ptr_to_ptr)
291
- ptr = tl.cast(ptr, tl.pointer_type(elem_dtype))
292
- return tl.multiple_of(ptr, 16)
vllm/v1/worker/gpu/buffer_utils.py
+89
-5
@@ -194,7 +194,9 @@ class StagedWriteTensor:
194
starts_uva,
195
write_contents,
196
cu_lens_uva,
197
+ None,
198
BLOCK_SIZE=1024,
199
+ MULTI_GROUP=False,
200
)
201
# Clear the staged writes
202
self.clear_staged_writes()
@@ -206,15 +208,81 @@ class StagedWriteTensor:
208
self._staged_write_cu_lens.clear()
209
210
211
+class FusedStagedWriter:
212
+ """Applies the staged writes of several `StagedWriteTensor`s at once."""
213
+
214
+ def __init__(
215
+ self, device: torch.device, max_writes: int, max_concurrency: int | None = None
216
+ ):
217
+ new_pool = partial(
218
+ UvaBufferPool, dtype=torch.int32, max_concurrency=max_concurrency
219
+ )
220
+ self.group_ids = new_pool(max_writes)
221
+ self.indices = new_pool(max_writes)
222
+ self.starts = new_pool(max_writes)
223
+ self.cu_lens = new_pool(max_writes)
224
+ self.device = device
225
+
226
+ def apply(
227
+ self,
228
+ tensors: Sequence[StagedWriteTensor],
229
+ output_ptrs: torch.Tensor,
230
+ output_strides: torch.Tensor,
231
+ ) -> None:
232
+ """Apply and clear the staged writes of `tensors` with one kernel."""
233
+ group_ids: list[int] = []
234
+ indices: list[int] = []
235
+ starts: list[int] = []
236
+ contents: list[int | float] = []
237
+ cu_lens: list[int] = []
238
+
239
+ for group_id, t in enumerate(tensors):
240
+ n = len(t._staged_write_indices)
241
+ if n == 0:
242
+ continue
243
+
244
+ group_ids.extend([group_id] * n)
245
+ indices.extend(t._staged_write_indices)
246
+ starts.extend(t._staged_write_starts)
247
+ content_base = len(contents)
248
+ contents.extend(t._staged_write_contents)
249
+ cu_lens.extend(content_base + cu_len for cu_len in t._staged_write_cu_lens)
250
+
251
+ if not group_ids:
252
+ return
253
+
254
+ group_ids_uva = self.group_ids.copy_to_uva(group_ids)
255
+ indices_uva = self.indices.copy_to_uva(indices)
256
+ starts_uva = self.starts.copy_to_uva(starts)
257
+ cu_lens_uva = self.cu_lens.copy_to_uva(cu_lens)
258
+ contents_gpu = async_tensor_h2d(contents, torch.int32, self.device)
259
+
260
+ _apply_write_kernel[(len(group_ids),)](
261
+ output_ptrs,
262
+ output_strides,
263
+ indices_uva,
264
+ starts_uva,
265
+ contents_gpu,
266
+ cu_lens_uva,
267
+ group_ids_uva,
268
+ BLOCK_SIZE=1024,
269
+ MULTI_GROUP=True,
270
+ )
271
+ for t in tensors:
272
+ t.clear_staged_writes()
273
+
274
+
275
@triton.jit
276
def _apply_write_kernel(
211
- output_ptr,
212
- output_stride,
277
+ output_ptr, # MULTI_GROUP: ptr-to-ptrs [num_groups]; else: data ptr
278
+ output_stride, # MULTI_GROUP: ptr-to-strides [num_groups]; else: row stride
279
write_indices_ptr,
280
write_starts_ptr,
281
write_contents_ptr,
282
write_cu_lens_ptr,
283
+ write_group_ids_ptr, # [num_writes], used only when MULTI_GROUP
284
BLOCK_SIZE: tl.constexpr,
285
+ MULTI_GROUP: tl.constexpr,
286
):
287
pid = tl.program_id(0)
288
row_idx = tl.load(write_indices_ptr + pid)
@@ -224,10 +292,26 @@ def _apply_write_kernel(
292
cu_end = tl.load(write_cu_lens_ptr + pid)
293
content_len = cu_end - cu_start
294
295
+ if MULTI_GROUP:
296
+ # Each write targets a different output tensor (KV cache group);
297
+ # resolve its base pointer and row stride per write.
298
+ group_id = tl.load(write_group_ids_ptr + pid)
299
+ row_ptr = _load_ptr(output_ptr + group_id, tl.int32)
300
+ row_stride = tl.load(output_stride + group_id)
301
+ else:
302
+ row_ptr = output_ptr
303
+ row_stride = output_stride
304
+ row_ptr += row_idx * row_stride + start_idx
305
+
306
for i in range(0, content_len, BLOCK_SIZE):
307
block = i + tl.arange(0, BLOCK_SIZE)
308
mask = block < content_len
309
content = tl.load(write_contents_ptr + cu_start + block, mask=mask)
231
- tl.store(
232
- output_ptr + row_idx * output_stride + start_idx + block, content, mask=mask
233
- )
310
+ tl.store(row_ptr + block, content, mask=mask)
311
+
312
+
313
+@triton.jit
314
+def _load_ptr(ptr_to_ptr, elem_dtype):
315
+ ptr = tl.load(ptr_to_ptr)
316
+ ptr = tl.cast(ptr, tl.pointer_type(elem_dtype))
317
+ return tl.multiple_of(ptr, 16)