| 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 | ) |