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 2, 02:15 UTC
[Perf] Triton fast path for small CPU→GPU `swap_blocks_batch` in the offloading connector (#42212)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Itay Etelis committed
Jun 3, 2026 at 13:38 UTC
1fa9ea09f61b9b639f59f170eef159d0c7bf163f
3 files changed
+180
-11
tests/v1/kv_offload/cpu/test_swap_blocks_triton.py
new
+31
@@ -0,0 +1,31 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit test for the Triton ``swap_blocks_batch`` fast-path kernel."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.platforms import current_platform
9
+from vllm.v1.kv_offload.cpu.swap_blocks_triton import swap_blocks_batch
10
+
11
+
12
+def _addrs(buffers: list[torch.Tensor]) -> torch.Tensor:
13
+ return torch.tensor([b.data_ptr() for b in buffers], dtype=torch.int64)
14
+
15
+
16
+@pytest.mark.skipif(
17
+ not current_platform.is_cuda(), reason="Triton swap fast path requires CUDA"
18
+)
19
+def test_triton_swap_copies_source_bytes():
20
+ # 8-byte-aligned, sub-threshold sizes covering 8 KiB chunk boundaries and
21
+ # odd tail-mask lengths, with enough descriptors to take the Triton path.
22
+ sizes = [8, 4096, 8192, 8200, 16384, 4088] * 8
23
+ src = [torch.randint(256, (s,), dtype=torch.uint8, device="cuda") for s in sizes]
24
+ dst = [torch.zeros_like(s) for s in src]
25
+ sizes_t = torch.tensor(sizes, dtype=torch.int64)
26
+
27
+ swap_blocks_batch(_addrs(src), _addrs(dst), sizes_t.clone(), bytes_per_chunk=8192)
28
+ torch.accelerator.synchronize()
29
+
30
+ for s, t in zip(src, dst):
31
+ assert torch.equal(t, s) # kernel copied the source bytes verbatim
vllm/v1/kv_offload/cpu/gpu_worker.py
+75
-11
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
import time
5
from collections import deque
6
from dataclasses import dataclass
@@ -10,6 +11,7 @@ from typing_extensions import override
11
12
from vllm import _custom_ops as ops
13
from vllm.logger import init_logger
14
+from vllm.triton_utils import HAS_TRITON, triton
15
from vllm.utils.math_utils import cdiv
16
from vllm.utils.platform_utils import is_pin_memory_available
17
from vllm.v1.kv_offload.base import (
@@ -19,6 +21,10 @@ from vllm.v1.kv_offload.base import (
21
GPULoadStoreSpec,
22
)
23
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
24
+from vllm.v1.kv_offload.cpu.swap_blocks_triton import (
25
+ THRESHOLD_BYTES,
26
+ swap_blocks_batch,
27
+)
28
from vllm.v1.kv_offload.worker.worker import (
29
OffloadingHandler,
30
TransferResult,
@@ -28,6 +34,30 @@ from vllm.v1.kv_offload.worker.worker import (
34
logger = init_logger(__name__)
35
36
37
+def _select_swap_blocks_fn(
38
+ kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
39
+ gpu_to_cpu: bool,
40
+):
41
+ """Resolve the swap_blocks function for a handler at init time."""
42
+ # GPU->CPU is bandwidth-bound; the dedicated copy engine beats Triton.
43
+ if gpu_to_cpu:
44
+ return ops.swap_blocks_batch
45
+ # Fall back to the C++ DMA path on platforms where Triton isn't usable
46
+ # (e.g. ROCm builds without Triton).
47
+ if not HAS_TRITON:
48
+ return ops.swap_blocks_batch
49
+ page_sizes = [r.page_size_bytes for g in kv_cache_groups_data_refs for r in g]
50
+ # Triton wins only on small, 8-byte-aligned payloads.
51
+ if (
52
+ not page_sizes
53
+ or max(page_sizes) >= THRESHOLD_BYTES
54
+ or any(s % 8 for s in page_sizes)
55
+ ):
56
+ return ops.swap_blocks_batch
57
+ chunk = min(triton.next_power_of_2(max(page_sizes)), 8192)
58
+ return functools.partial(swap_blocks_batch, bytes_per_chunk=chunk)
59
+
60
+
61
@dataclass
62
class Transfer:
63
job_id: int
@@ -35,6 +65,9 @@ class Transfer:
65
start_event: torch.Event
66
end_event: torch.Event
67
num_bytes: int
68
+ batch_src: torch.Tensor
69
+ batch_dst: torch.Tensor
70
+ batch_sizes: torch.Tensor
71
72
73
def compute_sub_block_ptrs(
@@ -109,6 +142,17 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None:
142
region.is_pinned = True
143
144
145
+def _new_descriptor_buffers(
146
+ num_copy_ops: int,
147
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
148
+ pin = is_pin_memory_available()
149
+ return (
150
+ torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin),
151
+ torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin),
152
+ torch.empty(num_copy_ops, dtype=torch.int64, pin_memory=pin),
153
+ )
154
+
155
+
156
class SingleDirectionOffloadingHandler(OffloadingHandler):
157
"""
158
SingleDirectionOffloadingHandler handles transfers for a single direction,
@@ -162,6 +206,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
206
)
207
self.gpu_to_cpu: bool = gpu_to_cpu
208
self.kv_cache_groups_data_refs = kv_cache_groups_data_refs
209
+ self._swap_blocks_batch = _select_swap_blocks_fn(
210
+ kv_cache_groups_data_refs, gpu_to_cpu
211
+ )
212
213
# GPU blocks may be smaller
214
# cpu_page_size = gpu_page_size * block_size_factor.
@@ -179,6 +226,8 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
226
self._stream_pool: list[torch.cuda.Stream] = []
227
# list of CUDA events available for re-use
228
self._event_pool: list[torch.Event] = []
229
+ # list of pinned descriptor buffer sets available for re-use
230
+ self._buffer_pool: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
231
232
@override
233
def transfer_async(self, job_id: int, transfer_spec: TransferSpec) -> bool:
@@ -229,9 +278,21 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
278
):
279
num_copy_ops += group_size * len(group_data_refs)
280
232
- all_src = np.empty(num_copy_ops, dtype=np.int64)
233
- all_dst = np.empty(num_copy_ops, dtype=np.int64)
234
- all_sizes = np.empty(num_copy_ops, dtype=np.int64)
281
+ # reuse a pooled buffer set, growing it if this transfer needs more room
282
+ batch_src, batch_dst, batch_sizes = (
283
+ self._buffer_pool.pop()
284
+ if self._buffer_pool
285
+ else _new_descriptor_buffers(num_copy_ops)
286
+ )
287
+ if batch_src.numel() < num_copy_ops:
288
+ batch_src, batch_dst, batch_sizes = _new_descriptor_buffers(num_copy_ops)
289
+
290
+ src = batch_src[:num_copy_ops]
291
+ dst = batch_dst[:num_copy_ops]
292
+ sizes = batch_sizes[:num_copy_ops]
293
+ all_src = src.numpy()
294
+ all_dst = dst.numpy()
295
+ all_sizes = sizes.numpy()
296
297
src_offset = 0
298
dst_offset = 0
@@ -294,10 +355,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
355
assert dst_offset == num_dst_blocks
356
assert op_idx == num_copy_ops
357
297
- batch_src = torch.from_numpy(all_src)
298
- batch_dst = torch.from_numpy(all_dst)
299
- batch_sizes = torch.from_numpy(all_sizes)
300
-
358
stream = self._stream_pool.pop() if self._stream_pool else torch.cuda.Stream()
359
start_event = (
360
self._event_pool.pop()
@@ -328,10 +385,10 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
385
with torch.cuda.stream(stream):
386
start_event.record(stream)
387
if num_copy_ops > 0:
331
- ops.swap_blocks_batch(
332
- batch_src,
333
- batch_dst,
334
- batch_sizes,
388
+ self._swap_blocks_batch(
389
+ src,
390
+ dst,
391
+ sizes,
392
is_src_access_order_any=is_src_access_order_any,
393
)
394
end_event.record(stream)
@@ -344,6 +401,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
401
start_event=start_event,
402
end_event=end_event,
403
num_bytes=num_transfer_bytes,
404
+ batch_src=batch_src,
405
+ batch_dst=batch_dst,
406
+ batch_sizes=batch_sizes,
407
)
408
)
409
@@ -370,6 +430,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
430
self._stream_pool.append(transfer.stream)
431
self._event_pool.append(transfer.end_event)
432
self._event_pool.append(transfer.start_event)
433
+ self._buffer_pool.append(
434
+ (transfer.batch_src, transfer.batch_dst, transfer.batch_sizes)
435
+ )
436
del self._transfer_events[transfer.job_id]
437
return results
438
@@ -388,6 +451,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
451
self._transfer_events.clear()
452
self._stream_pool.clear()
453
self._event_pool.clear()
454
+ self._buffer_pool.clear()
455
self.src_tensors.clear()
456
self.dst_tensors.clear()
457
if self._mmap_region is not None:
vllm/v1/kv_offload/cpu/swap_blocks_triton.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Triton kernel + tuned constants for the ``swap_blocks_batch`` fast path."""
4
+
5
+from __future__ import annotations
6
+
7
+import torch
8
+
9
+from vllm import _custom_ops as ops
10
+from vllm.triton_utils import tl, triton
11
+
12
+# Constants tuned empirically on H100 (PCIe Gen5):
13
+# NUM_SMS - smallest SM slice within 5% of peak bandwidth at the
14
+# 8-32 KB block sizes that matter in practice
15
+# THRESHOLD_BYTES - max payload per descriptor where Triton beats DMA; above
16
+# this the C++ cuMemcpyBatchAsync path takes the lead
17
+# MIN_N - minimum batch size where Triton's per-launch cost is
18
+# amortized; below this DMA wins
19
+NUM_SMS = 12
20
+THRESHOLD_BYTES = 28 * 1024
21
+MIN_N = 16
22
+
23
+
24
+@triton.jit
25
+def _swap_blocks_kernel(
26
+ src_addrs,
27
+ dst_addrs,
28
+ sizes,
29
+ n_jobs, # type: ignore[name-defined]
30
+ BYTES_PER_CHUNK: tl.constexpr, # type: ignore[name-defined]
31
+):
32
+ pid = tl.program_id(0)
33
+ num_progs = tl.num_programs(0)
34
+ WORDS_PER_CHUNK: tl.constexpr = BYTES_PER_CHUNK // 8
35
+ offsets = tl.arange(0, WORDS_PER_CHUNK)
36
+ job = pid
37
+ while job < n_jobs:
38
+ src = tl.load(src_addrs + job).to(tl.pointer_type(tl.int64))
39
+ dst = tl.load(dst_addrs + job).to(tl.pointer_type(tl.int64))
40
+ words = tl.load(sizes + job) // 8
41
+ for start in range(0, words, WORDS_PER_CHUNK):
42
+ idx = start + offsets
43
+ mask = idx < words
44
+ data = tl.load(src + idx, mask=mask, other=0)
45
+ tl.store(dst + idx, data, mask=mask)
46
+ job += num_progs
47
+
48
+
49
+def swap_blocks_batch(
50
+ src_addrs: torch.Tensor,
51
+ dst_addrs: torch.Tensor,
52
+ sizes: torch.Tensor,
53
+ is_src_access_order_any: bool = False,
54
+ *,
55
+ bytes_per_chunk: int,
56
+) -> None:
57
+ """Triton implementation of ``swap_blocks_batch`` for small CPU->GPU batches."""
58
+ n = src_addrs.numel()
59
+ # Too few descriptors to amortize Triton's launch cost.
60
+ if n < MIN_N:
61
+ ops.swap_blocks_batch(
62
+ src_addrs,
63
+ dst_addrs,
64
+ sizes,
65
+ is_src_access_order_any=is_src_access_order_any,
66
+ )
67
+ return
68
+ _swap_blocks_kernel[(min(NUM_SMS, n),)](
69
+ src_addrs.to("cuda", non_blocking=True),
70
+ dst_addrs.to("cuda", non_blocking=True),
71
+ sizes.to("cuda", non_blocking=True),
72
+ n,
73
+ BYTES_PER_CHUNK=bytes_per_chunk,
74
+ )