[Perf] Triton fast path for small CPU→GPU swap_blocks_batch in the offloading connector

vllm/v1/kv_offload/cpu/gpu_worker.py

@@ -7,7 +7,6 @@
 import numpy as np
 import torch
 
-from vllm import _custom_ops as ops
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.utils.platform_utils import is_pin_memory_available
@@ -18,6 +17,7 @@
     GPULoadStoreSpec,
 )
 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
+from vllm.v1.kv_offload.cpu.triton_swap import swap_blocks_batch
 from vllm.v1.kv_offload.worker.worker import (
     OffloadingHandler,
     TransferResult,
@@ -316,7 +316,7 @@ def transfer_async(self, job_id: int, transfer_spec: TransferSpec) -> bool:
         with torch.cuda.stream(stream):
             start_event.record(stream)
             if num_copy_ops > 0:
-                ops.swap_blocks_batch(batch_src, batch_dst, batch_sizes)
+                swap_blocks_batch(batch_src, batch_dst, batch_sizes)
             end_event.record(stream)
 
         self._transfer_events[job_id] = end_event

vllm/v1/kv_offload/cpu/triton_swap.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Triton fast path for ``swap_blocks_batch`` on small uniform batches."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.triton_utils import tl, triton
+
+_NUM_SMS = 12
+_THRESHOLD_BYTES = 28 * 1024
+
+
+@triton.jit
+def _kernel(
+    src_addrs,
+    dst_addrs,
+    n_jobs,  # type: ignore[name-defined]
+    bytes_per_job,  # type: ignore[name-defined]
+    BYTES_PER_CHUNK: tl.constexpr,  # type: ignore[name-defined]
+):
+    pid = tl.program_id(0)
+    num_progs = tl.num_programs(0)
+    WORDS_PER_CHUNK: tl.constexpr = BYTES_PER_CHUNK // 8
+    words = bytes_per_job // 8
+    offsets = tl.arange(0, WORDS_PER_CHUNK)
+    job = pid
+    while job < n_jobs:
+        src = tl.load(src_addrs + job).to(tl.pointer_type(tl.int64))
+        dst = tl.load(dst_addrs + job).to(tl.pointer_type(tl.int64))
+        for start in range(0, words, WORDS_PER_CHUNK):
+            idx = start + offsets
+            mask = idx < words
+            data = tl.load(src + idx, mask=mask, other=0)
+            tl.store(dst + idx, data, mask=mask)
+        job += num_progs
+
+
+def swap_blocks_batch(
+    src_addrs: torch.Tensor,
+    dst_addrs: torch.Tensor,
+    sizes: torch.Tensor,
+) -> None:
+    """Drop-in replacement for ``ops.swap_blocks_batch`` with Triton fast path."""
+    n = src_addrs.numel()
+    if n == 0:
+        return
+    bpj = int(sizes[0].item())
+    if bpj >= _THRESHOLD_BYTES or bpj % 8 != 0 or not bool((sizes == bpj).all()):
+        ops.swap_blocks_batch(src_addrs, dst_addrs, sizes)
+        return
+    chunk = min(triton.next_power_of_2(bpj), 8192)
+    _kernel[(min(_NUM_SMS, n),)](
+        src_addrs.to("cuda", non_blocking=True),
+        dst_addrs.to("cuda", non_blocking=True),
+        n,
+        bpj,
+        BYTES_PER_CHUNK=chunk,
+    )