From b51f9bbee765337f960f7d219982a7dd32e880d8 Mon Sep 17 00:00:00 2001 From: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Date: Thu, 20 Nov 2025 16:03:32 +0800 Subject: [PATCH] [Feature] Introduce JIT Kernel in sglang (with hicache JIT kernel) (#13453) --- benchmark/hicache/perf.py | 248 +++++++++ python/sglang/jit_kernel/.clang-format | 19 + python/sglang/jit_kernel/csrc/hicache.cuh | 264 ++++++++++ python/sglang/jit_kernel/hicache.py | 113 ++++ .../jit_kernel/include/sgl_kernel/tensor.h | 488 ++++++++++++++++++ .../jit_kernel/include/sgl_kernel/utils.cuh | 101 ++++ .../jit_kernel/include/sgl_kernel/utils.h | 88 ++++ .../jit_kernel/include/sgl_kernel/warp.cuh | 145 ++++++ python/sglang/jit_kernel/utils.py | 103 ++++ .../sglang/srt/mem_cache/memory_pool_host.py | 40 +- 10 files changed, 1589 insertions(+), 20 deletions(-) create mode 100644 benchmark/hicache/perf.py create mode 100644 python/sglang/jit_kernel/.clang-format create mode 100644 python/sglang/jit_kernel/csrc/hicache.cuh create mode 100644 python/sglang/jit_kernel/hicache.py create mode 100644 python/sglang/jit_kernel/include/sgl_kernel/tensor.h create mode 100644 python/sglang/jit_kernel/include/sgl_kernel/utils.cuh create mode 100644 python/sglang/jit_kernel/include/sgl_kernel/utils.h create mode 100644 python/sglang/jit_kernel/include/sgl_kernel/warp.cuh create mode 100644 python/sglang/jit_kernel/utils.py diff --git a/benchmark/hicache/perf.py b/benchmark/hicache/perf.py new file mode 100644 index 000000000..2349af4b1 --- /dev/null +++ b/benchmark/hicache/perf.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +from typing import Any, Callable, NamedTuple + +import torch + + +def jit_hicache_impl( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src: torch.Tensor, + item_bytes: int, + block_quota: int, +) -> None: + from sglang.jit_kernel.hicache import transfer_hicache_one_layer + + _ = item_bytes + + transfer_hicache_one_layer( + k_cache_dst=k_cache_dst, + v_cache_dst=v_cache_dst, + indices_dst=indices_dst, + k_cache_src=k_cache_src, + v_cache_src=v_cache_src, + indices_src=indices_src, + block_quota=block_quota, + ) + + +def ref_hicache_impl( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src: torch.Tensor, + item_bytes: int, + block_quota: int, +) -> None: + from sgl_kernel import transfer_kv_per_layer + + transfer_kv_per_layer( + src_k=k_cache_src, + src_v=v_cache_src, + dst_k=k_cache_dst, + dst_v=v_cache_dst, + src_indices=indices_src, + dst_indices=indices_dst, + item_size=item_bytes, + block_quota=block_quota, + ) + + +class HicacheBenchArgs(NamedTuple): + cache_item_size: int + dtype: torch.dtype + block_quota: int + + +def perf(f: Callable[[], Any], loop: int = 100) -> float: + tic = torch.cuda.Event(enable_timing=True) + toc = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + # warm up + f() + torch.cuda._sleep(10**8) + tic.record() + for _ in range(loop): + f() + toc.record() + toc.synchronize() + return tic.elapsed_time(toc) / loop + + +@torch.inference_mode() +def test_hicache_kernel(args: HicacheBenchArgs) -> None: + CACHE_ITEM_SIZE, DTYPE, BLOCK_QUOTA = args + + CUDA_CACHE_SIZE = 1024 * 1024 + HOST_CACHE_SIZE = CUDA_CACHE_SIZE * 2 + + cuda_cache = torch.randn( + (2, CUDA_CACHE_SIZE, CACHE_ITEM_SIZE), + dtype=DTYPE, + device="cuda", + ) + host_cache = torch.empty( + (2, HOST_CACHE_SIZE, CACHE_ITEM_SIZE), + dtype=DTYPE, + device="cpu", + pin_memory=True, + ) + + ITEM_BYTES = cuda_cache.element_size() * CACHE_ITEM_SIZE + + def _gen_indices(size: int, bs: int) -> torch.Tensor: + assert bs <= size + result = ( + (torch.randperm(size, dtype=torch.int64, device="cuda")[:bs]).sort().values + ) + if not (torch.all(result >= 0) and torch.all(result < size)): + where = (result < 0) | (result >= size) + place = where.nonzero(as_tuple=False) + print("Invalid indices at positions:", place) + print("Invalid indices values:", result[place]) + raise ValueError("Generated invalid indices") + return result + + def _calc_tput(dur: float) -> float: + return (MEM / (1024**3)) / (dur / 1000) # GB/s + + def _gain_str(aot_dur: float, jit_dur: float) -> str: + gain = 100 * (aot_dur / jit_dur - 1) + if gain >= 0: + return f"+{gain:>6.2f}%" + else: + return f"-{-gain:>6.2f}%" + + print(f"{CACHE_ITEM_SIZE = }, {DTYPE = }, {BLOCK_QUOTA = }") + + def _fast_test_correctness(bs: int): + src_indices = _gen_indices(CUDA_CACHE_SIZE, bs) + dst_indices = _gen_indices(HOST_CACHE_SIZE, bs) + host_cache_cuda = torch.randn_like(host_cache, device="cuda") + host_cache.copy_(host_cache_cuda, non_blocking=True) + + # copy from cuda to host + jit_hicache_impl( + k_cache_dst=host_cache[0], + v_cache_dst=host_cache[1], + indices_dst=dst_indices, + k_cache_src=cuda_cache[0], + v_cache_src=cuda_cache[1], + indices_src=src_indices, + item_bytes=ITEM_BYTES, + block_quota=BLOCK_QUOTA, + ) + dst_indices = dst_indices.cpu() + assert torch.all( + host_cache[0][dst_indices].cuda() == cuda_cache[0][src_indices] + ) + + BS_RANGE = [2**n for n in range(8, 18)] + for bs in BS_RANGE: + _fast_test_correctness(bs) + + print("Correctness passed! Start HiCache kernel performance test...") + print("=" * 70) + + for bs in BS_RANGE: + indices_dst = _gen_indices(CUDA_CACHE_SIZE, bs) + indices_src = _gen_indices(HOST_CACHE_SIZE, bs) + MEM = 2 * bs * ITEM_BYTES + + def _run_kernel_h2d(impl): + return impl( + k_cache_dst=cuda_cache[0], + v_cache_dst=cuda_cache[1], + indices_dst=indices_dst, + k_cache_src=host_cache[0], + v_cache_src=host_cache[1], + indices_src=indices_src, + item_bytes=ITEM_BYTES, + block_quota=BLOCK_QUOTA, + ) + + our_h2d_dur = perf(lambda: _run_kernel_h2d(jit_hicache_impl)) + ref_h2d_dur = perf(lambda: _run_kernel_h2d(ref_hicache_impl)) + print( + f"{bs = :6d}, H->D", + f"| aot {_calc_tput(ref_h2d_dur):<6.2f} GB/s", + f"| jit {_calc_tput(our_h2d_dur):<6.2f} GB/s", + f"| {_gain_str(ref_h2d_dur, our_h2d_dur)}", + ) + + print("=" * 70) + + for bs in BS_RANGE: + indices_dst = _gen_indices(HOST_CACHE_SIZE, bs) + indices_src = _gen_indices(CUDA_CACHE_SIZE, bs) + MEM = 2 * bs * ITEM_BYTES + + def _run_kernel_d2h(impl): + return impl( + k_cache_dst=host_cache[0], + v_cache_dst=host_cache[1], + indices_dst=indices_dst, + k_cache_src=cuda_cache[0], + v_cache_src=cuda_cache[1], + indices_src=indices_src, + item_bytes=ITEM_BYTES, + block_quota=BLOCK_QUOTA, + ) + + our_d2h_dur = perf(lambda: _run_kernel_d2h(jit_hicache_impl)) + ref_d2h_dur = perf(lambda: _run_kernel_d2h(ref_hicache_impl)) + print( + f"{bs = :6d}, D->H", + f"| aot {_calc_tput(ref_d2h_dur):<6.2f} GB/s", + f"| jit {_calc_tput(our_d2h_dur):<6.2f} GB/s", + f"| {_gain_str(ref_d2h_dur, our_d2h_dur)}", + ) + + print("=" * 70) + + +def main() -> None: + torch.cuda.set_device(0) + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + + tic = torch.cuda.Event(enable_timing=True) + toc = torch.cuda.Event(enable_timing=True) + + BUF_SIZE = 1024 * 1024 * 1024 + cuda_mem = torch.empty(BUF_SIZE, dtype=torch.uint8, device="cuda") + host_mem = torch.empty(BUF_SIZE, dtype=torch.uint8, device="cpu", pin_memory=True) + + # test peak bandwidth + tic.record() + cuda_mem.copy_(host_mem, non_blocking=True) + toc.record() + toc.synchronize() + dur = tic.elapsed_time(toc) + print(f"Peak H->D Bandwidth: {(BUF_SIZE / (1024**3)) / (dur / 1000):.2f} GB/s") + + tic.record() + host_mem.copy_(cuda_mem, non_blocking=True) + toc.record() + toc.synchronize() + dur = tic.elapsed_time(toc) + print(f"Peak D->H Bandwidth: {(BUF_SIZE / (1024**3)) / (dur / 1000):.2f} GB/s") + + for block_quota in [1, 2, 3, 4]: + for cache_item_size in [128, 256, 512, 1024]: + args = HicacheBenchArgs( + cache_item_size=cache_item_size, + dtype=torch.float16, + block_quota=block_quota, + ) + test_hicache_kernel(args) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/jit_kernel/.clang-format b/python/sglang/jit_kernel/.clang-format new file mode 100644 index 000000000..75fe1387c --- /dev/null +++ b/python/sglang/jit_kernel/.clang-format @@ -0,0 +1,19 @@ +BasedOnStyle: Google +IndentWidth: 2 +ColumnLimit: 120 +AllowShortFunctionsOnASingleLine: Empty +DerivePointerAlignment: false +PointerAlignment: Left +NamespaceIndentation: None +SortIncludes: true +AllowShortLoopsOnASingleLine: false +BinPackParameters: false # Prevents packing parameters in declarations +BinPackArguments: false # Prevents packing arguments in function calls +AlignAfterOpenBracket: AlwaysBreak # Forces a break after the opening parenthesis +AlignOperands: Align # Aligns arguments vertically +PenaltyBreakBeforeFirstCallParameter: 1 # Encourages breaking before the first argument +PenaltyReturnTypeOnItsOwnLine: 100 # Keeps return type with function name + +IncludeCategories: + - Regex: '^$' + Priority: 0 diff --git a/python/sglang/jit_kernel/csrc/hicache.cuh b/python/sglang/jit_kernel/csrc/hicache.cuh new file mode 100644 index 000000000..e52ecbd3a --- /dev/null +++ b/python/sglang/jit_kernel/csrc/hicache.cuh @@ -0,0 +1,264 @@ +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace { + +struct HicacheKernelParams { + void* __restrict__ k_cache_dst; + void* __restrict__ v_cache_dst; + const void* __restrict__ indices_dst; + void* __restrict__ k_cache_src; + void* __restrict__ v_cache_src; + const void* __restrict__ indices_src; + std::size_t length; + std::size_t kv_cache_src_stride; + std::size_t kv_cache_dst_stride; + std::size_t num_layers = 0; // only used in all_layer transfer +}; + +template < + std::integral T, + std::size_t kElementSize, + std::size_t kUnroll, + std::size_t kBlockQuota, + std::size_t kNumThreads, + std::size_t kMaxOccupancy> +__global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void hicache_transfer_per_layer( + const __grid_constant__ HicacheKernelParams params) { + // each warp acts as a worker + using namespace device; + static_assert(kNumThreads % kWarpThreads == 0); + static_assert(kWarpThreads % kUnroll == 0); + + constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; + constexpr auto kWarpsPerBlock = kNumThreads / kWarpThreads; + constexpr auto kWorkers = kWarpsPerBlock * kBlockQuota; + + const auto& [ + k_cache_dst, v_cache_dst, indices_dst, // dst + k_cache_src, v_cache_src, indices_src, // src + length, kv_cache_src_stride, kv_cache_dst_stride, _ // metadata + ] = params; + const auto warp_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads; + + // force to transfer 128 bytes per iteration + // since the PCIe transaction size is 128 bytes aligned + constexpr auto kGranularity = 128 / kWarpThreads; + + for (auto i = warp_id; i < length; i += kWorkers) { + const auto pos_src = static_cast(indices_src)[i]; + const auto pos_dst = static_cast(indices_dst)[i]; + const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); + const auto dst_k = pointer::offset(k_cache_dst, pos_dst * kv_cache_dst_stride); + const auto src_v = pointer::offset(v_cache_src, pos_src * kv_cache_src_stride); + const auto dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride); + const auto vec_k = warp::load_vec(src_k); + const auto vec_v = warp::load_vec(src_v); + warp::store_vec(dst_k, vec_k); + warp::store_vec(dst_v, vec_v); + } +} + +template < + std::integral T, + std::size_t kElementSize, + std::size_t kUnroll, + std::size_t kBlockQuota, + std::size_t kNumThreads, + std::size_t kMaxOccupancy> +__global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void hicache_transfer_all_layer( + const __grid_constant__ HicacheKernelParams params) { + // each warp acts as a worker + using namespace device; + using src_ptr_t = std::add_pointer_t; + using dst_ptr_t = std::add_pointer_t; + + static_assert(kNumThreads % kWarpThreads == 0); + constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; + constexpr auto kWarpsPerBlock = static_cast(kNumThreads) / kWarpThreads; + constexpr auto kWorkers = kWarpsPerBlock * kBlockQuota; + + const auto& [ + k_ptr_dst, v_ptr_dst, indices_dst, // dst + k_ptr_src, v_ptr_src, indices_src, // src + length, kv_cache_src_stride, kv_cache_dst_stride, num_layers // metadata + ] = params; + const auto warp_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads; + + // force to transfer 128 bytes per iteration + // since the PCIe transaction size is 128 bytes aligned + constexpr auto kGranularity = 128 / kWarpThreads; + + for (auto i = warp_id; i < length; i += kWorkers) { + const auto pos_src = static_cast(indices_src)[i]; + const auto pos_dst = static_cast(indices_dst)[i]; + for (std::size_t layer = 0; layer < num_layers; ++layer) { + const auto k_cache_src = static_cast(k_ptr_src)[layer]; + const auto v_cache_src = static_cast(v_ptr_src)[layer]; + const auto k_cache_dst = static_cast(k_ptr_dst)[layer]; + const auto v_cache_dst = static_cast(v_ptr_dst)[layer]; + const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); + const auto dst_k = pointer::offset(k_cache_dst, pos_dst * kv_cache_dst_stride); + const auto src_v = pointer::offset(v_cache_src, pos_src * kv_cache_src_stride); + const auto dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride); + const auto vec_k = warp::load_vec(src_k); + const auto vec_v = warp::load_vec(src_v); + warp::store_vec(dst_k, vec_k); + warp::store_vec(dst_v, vec_v); + } + } +} + +template < + std::size_t kElementSize, + std::size_t kUnroll, + std::size_t kBlockQuota, + std::size_t kNumThreads, + std::size_t kMaxOccupancy> +struct HiCacheKernel { + template + static constexpr auto _kernel_one = + hicache_transfer_per_layer; + template + static constexpr auto _kernel_all = + hicache_transfer_all_layer; + + static void run_one( + const tvm::ffi::TensorView k_cache_dst, + const tvm::ffi::TensorView v_cache_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView k_cache_src, + const tvm::ffi::TensorView v_cache_src, + const tvm::ffi::TensorView indices_src) { + using namespace host; + + auto D = SymbolicSize{"D"}; // cache dimension + auto N = SymbolicSize{"N"}; // src kv stride + auto M = SymbolicSize{"M"}; // dst kv stride + auto L = SymbolicSize{"L"}; // indices length + auto cache_dtype = SymbolicDType{}; + auto indices_dtype = SymbolicDType{}; + auto indices_device = SymbolicDevice{}; + + TensorMatcher({-1, D}) // + .with_strides({N, 1}) + .with_dtype(cache_dtype) + .with_device() + .verify(k_cache_src) + .verify(v_cache_src); + TensorMatcher({-1, D}) // + .with_strides({M, 1}) + .with_dtype(cache_dtype) + .with_device() + .verify(k_cache_dst) + .verify(v_cache_dst); + TensorMatcher({L}) // + .with_dtype(indices_dtype) + .with_device(indices_device) + .verify(indices_src) + .verify(indices_dst); + + // verify dimension match + const auto dtype_size = dtype_bytes(cache_dtype.unwrap()); + const auto element_bytes = D.unwrap() * dtype_size; + RuntimeCheck(kElementSize == element_bytes, "HicacheKernel: cache dimension mismatch."); + + const auto k_cache_dst_ptr = k_cache_dst.data_ptr(); + const auto v_cache_dst_ptr = v_cache_dst.data_ptr(); + const auto k_cache_src_ptr = k_cache_src.data_ptr(); + const auto v_cache_src_ptr = v_cache_src.data_ptr(); + const auto indices_dst_ptr = indices_dst.data_ptr(); + const auto indices_src_ptr = indices_src.data_ptr(); + const auto length = static_cast(L.unwrap()); + const auto kv_cache_src_stride = static_cast(N.unwrap()) * dtype_size; + const auto kv_cache_dst_stride = static_cast(M.unwrap()) * dtype_size; + const auto use_int32 = indices_dtype.unwrap().bits == 32; + const auto device = indices_device.unwrap(); + + constexpr auto kWorkersPerBlock = kNumThreads / (device::kWarpThreads / kUnroll); + const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); + const auto params = HicacheKernelParams{ + .k_cache_dst = k_cache_dst_ptr, + .v_cache_dst = v_cache_dst_ptr, + .indices_dst = indices_dst_ptr, + .k_cache_src = k_cache_src_ptr, + .v_cache_src = v_cache_src_ptr, + .indices_src = indices_src_ptr, + .length = length, + .kv_cache_src_stride = kv_cache_src_stride, + .kv_cache_dst_stride = kv_cache_dst_stride, + }; + const auto kernel = use_int32 ? _kernel_one : _kernel_one; + LaunchKernel(num_blocks, kNumThreads, device)(kernel, params); + } + + static void run_all( + const tvm::ffi::TensorView k_ptr_dst, + const tvm::ffi::TensorView v_ptr_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView k_ptr_src, + const tvm::ffi::TensorView v_ptr_src, + const tvm::ffi::TensorView indices_src, + const std::size_t kv_src_stride, + const std::size_t kv_dst_stride) { + using namespace host; + + auto N = SymbolicSize{"N"}; // num layers + auto L = SymbolicSize{"L"}; // indices length + auto dtype_ = SymbolicDType{}; + auto device_ = SymbolicDevice{}; + + TensorMatcher({N}) // + .with_dtype() + .with_device(device_) + .verify(k_ptr_src) + .verify(v_ptr_src) + .verify(k_ptr_dst) + .verify(v_ptr_dst); + TensorMatcher({L}) // + .with_dtype(dtype_) + .with_device(device_) + .verify(indices_src) + .verify(indices_dst); + + // verify dimension match + const auto k_cache_dst_ptr = k_ptr_dst.data_ptr(); + const auto v_cache_dst_ptr = v_ptr_dst.data_ptr(); + const auto k_cache_src_ptr = k_ptr_src.data_ptr(); + const auto v_cache_src_ptr = v_ptr_src.data_ptr(); + const auto indices_dst_ptr = indices_dst.data_ptr(); + const auto indices_src_ptr = indices_src.data_ptr(); + const auto length = static_cast(L.unwrap()); + const auto use_int32 = dtype_.unwrap().bits == 32; + const auto device = device_.unwrap(); + + constexpr auto kWorkersPerBlock = kNumThreads / (device::kWarpThreads / kUnroll); + const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); + const auto params = HicacheKernelParams{ + .k_cache_dst = k_cache_dst_ptr, + .v_cache_dst = v_cache_dst_ptr, + .indices_dst = indices_dst_ptr, + .k_cache_src = k_cache_src_ptr, + .v_cache_src = v_cache_src_ptr, + .indices_src = indices_src_ptr, + .length = length, + .kv_cache_src_stride = kv_src_stride, + .kv_cache_dst_stride = kv_dst_stride, + .num_layers = static_cast(N.unwrap()), + }; + const auto kernel = use_int32 ? _kernel_all : _kernel_all; + LaunchKernel(num_blocks, kNumThreads, device)(kernel, params); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/hicache.py b/python/sglang/jit_kernel/hicache.py new file mode 100644 index 000000000..5612809bf --- /dev/null +++ b/python/sglang/jit_kernel/hicache.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING + +from sglang.jit_kernel.utils import load_jit, make_cpp_args + +if TYPE_CHECKING: + import torch + from tvm_ffi.module import Module + +DEFAULT_BLOCK_QUOTA = 2 + + +@lru_cache(maxsize=None) +def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> Module: + num_threads, occupancy = 1024, 1 + args = make_cpp_args( + element_size, + unroll, + block_quota, + num_threads, + occupancy, + ) + return load_jit( + "hicache", + *args, + cuda_files=["hicache.cuh"], + cuda_wrappers=[("launch_one", f"HiCacheKernel<{args}>::run_one")], + ) + + +def _default_unroll(element_size: int) -> int: + if element_size <= 512: + return 4 + + if element_size <= 1024: + return 2 + + # fallback: no unroll + return 1 + + +def transfer_hicache_one_layer( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src: torch.Tensor, + *, + element_dim: int | None = None, + unroll: int | None = None, # can be tuned for performance + block_quota: int | None = None, # can be tuned for less interference +) -> None: + element_dim = element_dim or k_cache_dst.size(-1) + k_cache_src = k_cache_src.view(-1, element_dim) + v_cache_src = v_cache_src.view(-1, element_dim) + k_cache_dst = k_cache_dst.view(-1, element_dim) + v_cache_dst = v_cache_dst.view(-1, element_dim) + element_size = element_dim * k_cache_dst.element_size() + block_quota = block_quota or DEFAULT_BLOCK_QUOTA + unroll = unroll or _default_unroll(element_size) + module = _jit_hicache_module( + element_size=element_size, + unroll=unroll, + block_quota=block_quota, + ) + module.launch_one( + k_cache_dst, + v_cache_dst, + indices_dst, + k_cache_src, + v_cache_src, + indices_src, + ) + + +def transfer_hicache_all_layer( + k_ptr_dst: torch.Tensor, + v_ptr_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_ptr_src: torch.Tensor, + v_ptr_src: torch.Tensor, + indices_src: torch.Tensor, + kv_cache_src_stride_bytes: int, + kv_cache_dst_stride_bytes: int, + *, + element_size: int | None = None, + unroll: int | None = None, # can be tuned for performance + block_quota: int | None = None, # can be tuned for less interference +) -> None: + if element_size is None: # assume both contiguous + assert kv_cache_dst_stride_bytes == kv_cache_src_stride_bytes + element_size = kv_cache_dst_stride_bytes + + block_quota = block_quota or DEFAULT_BLOCK_QUOTA + unroll = unroll or _default_unroll(element_size) + module = _jit_hicache_module( + element_size=element_size, + unroll=unroll, + block_quota=block_quota, + ) + module.launch_all( + k_ptr_dst, + v_ptr_dst, + indices_dst, + k_ptr_src, + v_ptr_src, + indices_src, + kv_cache_src_stride_bytes, + kv_cache_dst_stride_bytes, + ) diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h new file mode 100644 index 000000000..5df90dc53 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h @@ -0,0 +1,488 @@ +#pragma once +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace host { + +namespace stdr = std::ranges; +namespace stdv = std::views; + +namespace details { + +struct SizeRef; +struct DTypeRef; +struct DeviceRef; + +template +struct dtype_trait {}; + +template +struct dtype_trait { + inline static constexpr auto value = DLDataType{ + .code = std::is_signed_v ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt, + .bits = static_cast(sizeof(T) * 8), + .lanes = 1}; +}; + +template +struct dtype_trait { + inline static constexpr auto value = + DLDataType{.code = DLDataTypeCode::kDLFloat, .bits = static_cast(sizeof(T) * 8), .lanes = 1}; +}; + +inline constexpr auto kAnyDeviceID = -1; +inline constexpr auto kAnySize = static_cast(-1); +inline constexpr auto kNullSize = static_cast(0); +inline constexpr auto kNullDType = static_cast(18u); +inline constexpr auto kNullDevice = static_cast(-1); + +template +inline constexpr auto kDTypeList = std::array{dtype_trait::value...}; + +template +inline constexpr auto kDeviceList = std::array{ + DLDevice{.device_type = static_cast(Codes), .device_id = kAnyDeviceID}...}; + +template +struct PrintAbleSpan { + explicit PrintAbleSpan(std::span data) : data(data) {} + std::span data; +}; + +// define DLDataType comparison and printing in root namespace +template +inline constexpr auto kDeviceStringMap = [] { + constexpr auto map = std::array{ + std::pair{DLDeviceType::kDLCPU, "cpu"}, + std::pair{DLDeviceType::kDLCUDA, "cuda"}, + std::pair{DLDeviceType::kDLCUDAHost, "cuda_host"}, + std::pair{DLDeviceType::kDLOpenCL, "opencl"}, + std::pair{DLDeviceType::kDLVulkan, "vulkan"}, + std::pair{DLDeviceType::kDLMetal, "metal"}, + std::pair{DLDeviceType::kDLVPI, "vpi"}, + std::pair{DLDeviceType::kDLROCM, "rocm"}, + std::pair{DLDeviceType::kDLROCMHost, "rocm_host"}, + std::pair{DLDeviceType::kDLExtDev, "ext_dev"}, + std::pair{DLDeviceType::kDLCUDAManaged, "cuda_managed"}, + std::pair{DLDeviceType::kDLOneAPI, "oneapi"}, + std::pair{DLDeviceType::kDLWebGPU, "webgpu"}, + std::pair{DLDeviceType::kDLHexagon, "hexagon"}, + std::pair{DLDeviceType::kDLMAIA, "maia"}, + std::pair{DLDeviceType::kDLTrn, "trn"}, + }; + constexpr auto max_type = stdr::max(map | stdv::keys); + auto result = std::array{}; + for (const auto& [code, name] : map) { + result[static_cast(code)] = name; + } + return result; +}(); + +struct PrintableDevice { + DLDevice device; +}; + +inline auto& operator<<(std::ostream& os, DLDevice device) { + const auto& mapping = kDeviceStringMap<>; + const auto entry = static_cast(device.device_type); + host::RuntimeCheck(entry < mapping.size()); + const auto name = mapping[entry]; + host::RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type)); + os << name; + if (device.device_id != kAnyDeviceID) os << "[" << device.device_id << "]"; + return os; +} + +inline auto& operator<<(std::ostream& os, PrintableDevice pd) { + return os << pd.device; +} + +template +inline auto& operator<<(std::ostream& os, PrintAbleSpan span) { + os << "["; + for (const auto i : stdv::iota(std::size_t{0}, span.data.size())) { + if (i > 0) { + os << ", "; + } + os << span.data[i]; + } + os << "]"; + return os; +} + +} // namespace details + +struct SymbolicSize { + public: + SymbolicSize(std::string_view annotation = {}) : m_value(details::kNullSize), m_annotation(annotation) {} + + auto get_name() const -> std::string_view { + return m_annotation; + } + auto set_value(int64_t value) -> void { + host::RuntimeCheck(!this->has_value(), "Size value already set"); + m_value = value; + } + auto has_value() const -> bool { + return m_value != 0; + } + auto get_value() const -> std::optional { + return this->has_value() ? std::optional{m_value} : std::nullopt; + } + auto unwrap() const -> int64_t { + host::RuntimeCheck(this->has_value(), "Size value is not set"); + return m_value; + } + + SymbolicSize(const SymbolicSize&) = delete; + SymbolicSize& operator=(const SymbolicSize&) = delete; + + auto verify(int64_t dim) -> void { + if (this->has_value()) { + host::RuntimeCheck(m_value == dim, "Size mismatch: expected ", m_value, " but got ", dim); + } else { + this->set_value(dim); + } + } + + private: + std::int64_t m_value; + std::string_view m_annotation; +}; + +inline auto operator==(DLDevice lhs, DLDevice rhs) -> bool { + return lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id; +} + +struct SymbolicDType { + public: + SymbolicDType() : m_value({details::kNullDType, 0, 0}) {} + + auto set_value(DLDataType value) -> void { + host::RuntimeCheck(!this->has_value(), "Dtype value already set"); + host::RuntimeCheck( + m_check(value), "Dtype value [", value, "] not in the allowed options: ", details::PrintAbleSpan{m_options}); + m_value = value; + } + auto has_value() const -> bool { + return m_value.code != details::kNullDType; + } + auto get_value() const -> std::optional { + return this->has_value() ? std::optional{m_value} : std::nullopt; + } + auto unwrap() const -> DLDataType { + host::RuntimeCheck(this->has_value(), "Dtype value is not set"); + return m_value; + } + + auto set_options(std::span options) -> void { + m_options = options; + } + template + auto set_options() -> void { + m_options = details::kDTypeList; + } + + auto verify(DLDataType dtype) -> void { + if (this->has_value()) { + host::RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype); + } else { + this->set_value(dtype); + } + } + + private: + auto m_check(DLDataType value) const -> bool { + return stdr::empty(m_options) || (stdr::find(m_options, value) != stdr::end(m_options)); + } + + std::span m_options; + DLDataType m_value; +}; + +struct SymbolicDevice { + public: + SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {} + + auto set_value(DLDevice value) -> void { + host::RuntimeCheck(!this->has_value(), "Device value already set"); + host::RuntimeCheck( + m_check(value), + "Device value [", + details::PrintableDevice{value}, + "] not in the allowed options: ", + details::PrintAbleSpan{m_options}); + m_value = value; + } + auto has_value() const -> bool { + return m_value.device_type != details::kNullDevice; + } + auto get_value() const -> std::optional { + return this->has_value() ? std::optional{m_value} : std::nullopt; + } + auto unwrap() const -> DLDevice { + host::RuntimeCheck(this->has_value(), "Device value is not set"); + return m_value; + } + + auto set_options(std::span options) -> void { + m_options = options; + } + template + auto set_options() -> void { + m_options = details::kDeviceList; + } + + auto verify(DLDevice device) -> void { + if (this->has_value()) { + host::RuntimeCheck( + m_value == device, + "Device mismatch: expected ", + details::PrintableDevice{m_value}, + " but got ", + details::PrintableDevice{device}); + } else { + this->set_value(device); + } + } + + private: + auto m_check(DLDevice value) const -> bool { + return stdr::empty(m_options) || (stdr::any_of(m_options, [value](const DLDevice& opt) { + // device type must exactly match + if (opt.device_type != value.device_type) return false; + // device id can be wildcarded + return opt.device_id == details::kAnyDeviceID || opt.device_id == value.device_id; + })); + } + + std::span m_options; + DLDevice m_value; +}; + +namespace details { + +template +struct BaseRef { + public: + BaseRef(const BaseRef&) = delete; + BaseRef& operator=(const BaseRef&) = delete; + + auto operator->() const -> T* { + return m_ref; + } + auto operator*() const -> T& { + return *m_ref; + } + auto rebind(T& other) -> void { + m_ref = &other; + } + + explicit BaseRef() : m_ref(&m_cache), m_cache() {} + BaseRef(T& size) : m_ref(&size), m_cache() {} + + private: + T* m_ref; + T m_cache; +}; + +struct SizeRef : BaseRef { + using BaseRef::BaseRef; + SizeRef(int64_t value) { + if (value != kAnySize) { + (**this).set_value(value); + } else { + // otherwise, we can match any size + } + } + + auto value_or_name(std::size_t dim) const -> std::string { + if (const auto value = (**this).get_value()) { + return std::to_string(*value); + } else { + const auto annotation = (**this).get_name(); + if (annotation.empty()) { + return "dim#" + std::to_string(dim); + } else { + return static_cast(annotation); + } + } + } +}; + +struct DTypeRef : BaseRef { + using BaseRef::BaseRef; + DTypeRef(DLDataType options) { + (**this).set_value(options); + } + DTypeRef(std::initializer_list options) { + (**this).set_options(options); + } + DTypeRef(std::span options) { + (**this).set_options(options); + } +}; + +struct DeviceRef : BaseRef { + using BaseRef::BaseRef; + DeviceRef(DLDevice options) { + (**this).set_value(options); + } + DeviceRef(std::initializer_list options) { + (**this).set_options(options); + } + DeviceRef(std::span options) { + (**this).set_options(options); + } +}; + +} // namespace details + +struct TensorMatcher { + private: + using SizeRef = details::SizeRef; + using DTypeRef = details::DTypeRef; + using DeviceRef = details::DeviceRef; + using Loc_t = std::source_location; + + public: + TensorMatcher(const TensorMatcher&) = delete; + TensorMatcher& operator=(const TensorMatcher&) = delete; + + explicit TensorMatcher(std::initializer_list shape) : m_shape(shape), m_strides(), m_dtype() {} + + auto with_strides(std::initializer_list strides) && -> TensorMatcher&& { + // no partial update allowed + host::RuntimeCheck(m_strides.size() == 0, "Strides already specified"); + host::RuntimeCheck(m_shape.size() == strides.size(), "Strides size must match shape size"); + m_strides = strides; + return std::move(*this); + } + + template + auto with_dtype(DTypeRef&& dtype) && -> TensorMatcher&& { + m_init_dtype(); + m_dtype.rebind(*dtype); + return std::move(*this); + } + + template + auto with_dtype() && -> TensorMatcher&& { + static_assert(sizeof...(Ts) > 0, "At least one dtype option must be specified"); + m_init_dtype(); + m_dtype->set_options(); + return std::move(*this); + } + + template + auto with_device(DeviceRef&& device) && -> TensorMatcher&& { + m_init_device(); + m_device.rebind(*device); + return std::move(*this); + } + + template + auto with_device() && -> TensorMatcher&& { + static_assert(sizeof...(Codes) > 0, "At least one device option must be specified"); + m_init_device(); + m_device->set_options(); + return std::move(*this); + } + + // once we start verification, we cannot modify anymore + auto verify(tvm::ffi::TensorView view, Loc_t loc = Loc_t::current()) const&& -> const TensorMatcher&& { + try { + this->m_verify_impl(view); + } catch (PanicError& e) { + auto oss = std::ostringstream{}; + oss << "Tensor match failed for " << this->debug_str() << " at " << loc.file_name() << ":" << loc.line() + << "\n- Root cause: " << e.detail(); + throw PanicError(std::move(oss).str()); + } + return std::move(*this); + } + + auto debug_str() const -> std::string { + auto oss = std::ostringstream{}; + oss << "Tensor<"; + std::size_t dim = 0; + for (const auto& size_ref : m_shape) { + if (dim > 0) { + oss << ", "; + } + oss << size_ref.value_or_name(dim++); + } + oss << ">"; + if (m_strides.size() > 0) { + oss << " [strides=<"; + dim = 0; + for (const auto& stride_ref : m_strides) { + if (dim > 0) { + oss << ", "; + } + oss << stride_ref.value_or_name(dim++); + } + oss << ">]"; + } + return std::move(oss).str(); + } + + private: + auto m_verify_impl(tvm::ffi::TensorView view) const -> void { + const auto dim = static_cast(view.dim()); + host::RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim); + for (const auto i : stdv::iota(std::size_t{0}, dim)) { + m_shape[i]->verify(view.size(i)); + } + if (this->m_has_strides()) { + for (const auto i : stdv::iota(std::size_t{0}, dim)) { + m_strides[i]->verify(view.stride(i)); + } + } else { + host::RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected"); + } + // since we may double verify, we will force to check + m_dtype->verify(view.dtype()); + m_device->verify(view.device()); + } + + auto m_init_dtype() -> void { + host::RuntimeCheck(!m_has_dtype, "DType already specified"); + m_has_dtype = true; + } + auto m_init_device() -> void { + host::RuntimeCheck(!m_has_device, "Device already specified"); + m_has_device = true; + } + auto m_has_strides() const -> bool { + return !m_strides.empty(); + } + + std::span m_shape; + std::span m_strides; + DTypeRef m_dtype; + DeviceRef m_device; + bool m_has_dtype = false; + bool m_has_device = false; +}; + +} // namespace host diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh new file mode 100644 index 000000000..cf03d8c07 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -0,0 +1,101 @@ +#pragma once + +#include + +#include +#include + +#include +#include +#include +#include + +namespace device { + +inline constexpr auto kWarpThreads = 32u; + +namespace pointer { + +// we only allow void * pointer arithmetic for safety + +template +__always_inline __device__ auto offset(T* ptr, U... offset) -> void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +template +__always_inline __device__ auto offset(const T* ptr, U... offset) -> const void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +} // namespace pointer + +} // namespace device + +namespace host { + +inline auto +RuntimeDeviceCheck(::cudaError_t error, std::source_location location = std::source_location::current()) -> void { + if (error != ::cudaSuccess) { + [[unlikely]]; + ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); + } +} + +inline auto RuntimeCudaCheck(std::source_location location = std::source_location::current()) -> void { + return RuntimeDeviceCheck(::cudaGetLastError(), location); +} + +template +inline void set_smem_once(std::size_t smem_size) { + static const auto last_smem_size = [&] { + RuntimeDeviceCheck(::cudaFuncSetAttribute(F, ::cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + return smem_size; + }(); + RuntimeCheck( + smem_size <= last_smem_size, + "Dynamic shared memory size exceeds the previously set maximum size: ", + last_smem_size, + " bytes"); +} + +struct LaunchKernel { + public: + explicit LaunchKernel( + dim3 grid_dim, dim3 block_dim, DLDevice device, std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)) {} + + explicit LaunchKernel( + dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t dynamic_shared_mem_bytes = 0) noexcept + : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)) {} + + static auto resolve_device(DLDevice device) -> cudaStream_t { + return static_cast(::TVMFFIEnvGetStream(device.device_type, device.device_id)); + } + + LaunchKernel(const LaunchKernel&) = delete; + LaunchKernel& operator=(const LaunchKernel&) = delete; + + template + auto operator()(T&& kernel, Args&&... args) const -> void { + host::RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward(args)...)); + } + + private: + static auto + s_make_config(dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t smem) -> cudaLaunchConfig_t { + auto config = ::cudaLaunchConfig_t{}; + config.gridDim = grid_dim; + config.blockDim = block_dim; + config.dynamicSmemBytes = smem; + config.stream = stream; + config.numAttrs = 0; + return config; + } + cudaLaunchConfig_t m_config; + /// TODO: We can add a queue to store the attributes if needed in the future. +}; + +} // namespace host diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/python/sglang/jit_kernel/include/sgl_kernel/utils.h new file mode 100644 index 000000000..fd9723df6 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.h @@ -0,0 +1,88 @@ +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace host { + +struct PanicError : public std::runtime_error { + public: + // copy and move constructors + explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {} + auto detail() const -> std::string_view { + const auto sv = std::string_view{m_message}; + const auto pos = sv.find(": "); + return pos == std::string_view::npos ? sv : sv.substr(pos + 2); + } + + private: + std::string m_message; +}; + +template +[[noreturn]] +inline auto panic(std::source_location location, Args&&... args) -> void { + std::ostringstream os; + os << "Runtime check failed at " << location.file_name() << ":" << location.line(); + if constexpr (sizeof...(args) > 0) { + os << ": "; + (os << ... << std::forward(args)); + } else { + os << " in " << location.function_name(); + } + throw PanicError(std::move(os).str()); +} + +template +struct RuntimeCheck { + using Loc_t = std::source_location; + template + explicit RuntimeCheck(Cond&& condition, Args&&... args, Loc_t location = Loc_t::current()) { + if (!condition) { + [[unlikely]]; + ::host::panic(location, std::forward(args)...); + } + } +}; + +template +explicit RuntimeCheck(Cond&&, Args&&...) -> RuntimeCheck; + +template +inline constexpr auto div_ceil(T a, U b) { + return (a + b - 1) / b; +} + +template +inline constexpr auto div_ceil(T a, U b) { + return (a + b - 1) / b; +} + +inline auto dtype_bytes(DLDataType dtype) -> std::size_t { + return static_cast(dtype.bits / 8); +} + +namespace pointer { + +// we only allow void * pointer arithmetic for safety + +template +inline auto offset(T* ptr, U... offset) -> void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +template +inline auto offset(const T* ptr, U... offset) -> const void* { + static_assert(std::is_same_v, "Pointer arithmetic is only allowed for void* pointers"); + return static_cast(ptr) + (... + offset); +} + +} // namespace pointer + +} // namespace host diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh new file mode 100644 index 000000000..904531f30 --- /dev/null +++ b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh @@ -0,0 +1,145 @@ +#pragma once +#include + +#include +#include +#include + +namespace device::warp { + +namespace details { + +template +inline constexpr auto get_mem_package() { + if constexpr (kUnit == 16) { + return uint4{}; + } else if constexpr (kUnit == 8) { + return uint2{}; + } else if constexpr (kUnit == 4) { + return uint1{}; + } else { + static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4, "Unsupported memory package size"); + } +} + +inline constexpr auto default_unit_size(std::size_t x) -> std::size_t { + if (x % (16 * kWarpThreads) == 0) return 16; + if (x % (8 * kWarpThreads) == 0) return 8; + if (x % (4 * kWarpThreads) == 0) return 4; + return 0; // trigger static assert in _get_mem_package +} + +template +using mem_package_t = decltype(get_mem_package()); + +template +struct storage_vec { + T data[N]; +}; + +__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { + uint32_t tmp; + asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); + return uint1{tmp}; +} + +__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { + uint32_t tmp0, tmp1; + asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); + return uint2{tmp0, tmp1}; +} + +__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { + uint32_t tmp0, tmp1, tmp2, tmp3; + asm volatile("ld.global.cs.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); + return uint4{tmp0, tmp1, tmp2, tmp3}; +} + +__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { + uint32_t tmp = value.x; + asm volatile("st.global.cs.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +} + +__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { + uint32_t tmp0 = value.x; + uint32_t tmp1 = value.y; + asm volatile("st.global.cs.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +} + +__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { + uint32_t tmp0 = value.x; + uint32_t tmp1 = value.y; + uint32_t tmp2 = value.z; + uint32_t tmp3 = value.w; + asm volatile("st.global.cs.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +} + +} // namespace details + +template < + std::size_t kBytes, + std::size_t kUnit = details::default_unit_size(kBytes), + std::size_t kThreads = ::device::kWarpThreads> +__always_inline __device__ void copy(void* __restrict__ dst, const void* __restrict__ src) { + using Package = details::mem_package_t; + constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; + constexpr auto kLoopCount = kBytes / kBytesPerLoop; + static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + + const auto dst_packed = static_cast(dst); + const auto src_packed = static_cast(src); + const auto lane_id = threadIdx.x % kThreads; + +#pragma unroll kLoopCount + for (std::size_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kThreads + lane_id; + dst_packed[j] = src_packed[j]; + } +} + +template < + std::size_t kBytes, + std::size_t kUnit = details::default_unit_size(kBytes), + std::size_t kThreads = ::device::kWarpThreads> +__always_inline __device__ auto load_vec(const void* __restrict__ src) { + using Package = details::mem_package_t; + constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; + constexpr auto kLoopCount = kBytes / kBytesPerLoop; + static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + + const auto src_packed = static_cast(src); + const auto lane_id = threadIdx.x % kThreads; + details::storage_vec vec; + +#pragma unroll kLoopCount + for (std::size_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kThreads + lane_id; + vec.data[i] = details::load_nc(src_packed + j); + } + + return vec; +} + +template < + std::size_t kBytes, + std::size_t kUnit = details::default_unit_size(kBytes), + std::size_t kThreads = ::device::kWarpThreads, + typename Tp> +__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) { + using Package = details::mem_package_t; + constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; + constexpr auto kLoopCount = kBytes / kBytesPerLoop; + static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); + static_assert(std::is_same_v>); + + const auto dst_packed = static_cast(dst); + const auto lane_id = threadIdx.x % kThreads; + +#pragma unroll kLoopCount + for (std::size_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kThreads + lane_id; + details::store_nc(dst_packed + j, vec.data[i]); + } +} + +} // namespace device::warp diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py new file mode 100644 index 000000000..6462cf41c --- /dev/null +++ b/python/sglang/jit_kernel/utils.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import pathlib +from functools import lru_cache +from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union + +if TYPE_CHECKING: + from tvm_ffi import Module + + +def _make_wrapper(tup: Tuple[str, str]) -> str: + export_name, kernel_name = tup + return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));" + + +@lru_cache() +def _resolve_kernel_path() -> pathlib.Path: + cur_dir = pathlib.Path(__file__).parent.resolve() + + # first, try this directory structure + def _environment_install(): + candidate = cur_dir.resolve() + if (candidate / "include").exists() and (candidate / "csrc").exists(): + return candidate + return None + + def _package_install(): + # TODO: support find path by package + return None + + path = _environment_install() or _package_install() + if path is None: + raise RuntimeError("Cannot find sgl-kernel/jit path") + return path + + +KERNEL_PATH = _resolve_kernel_path() +DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] +DEFAULT_CFLAGS = ["-std=c++20", "-O3"] +DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_LDFLAGS = [] +CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] + + +class CPPArgList(list[str]): + def __str__(self) -> str: + return ", ".join(self) + + +def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList: + def _convert(arg: CPP_TEMPLATE_TYPE) -> str: + if isinstance(arg, bool): + return "true" if arg else "false" + if isinstance(arg, (int, float)): + return str(arg) + raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}") + + return CPPArgList(_convert(arg) for arg in args) + + +def load_jit( + *args: str, + cpp_files: List[str] | None = None, + cuda_files: List[str] | None = None, + cpp_wrappers: List[Tuple[str, str]] | None = None, + cuda_wrappers: List[Tuple[str, str]] | None = None, + extra_cflags: List[str] | None = None, + extra_cuda_cflags: List[str] | None = None, + extra_ldflags: List[str] | None = None, + extra_include_paths: List[str] | None = None, + build_directory: str | None = None, +) -> Module: + from tvm_ffi.cpp import load_inline + + cpp_files = cpp_files or [] + cuda_files = cuda_files or [] + cpp_wrappers = cpp_wrappers or [] + cuda_wrappers = cuda_wrappers or [] + extra_cflags = extra_cflags or [] + extra_cuda_cflags = extra_cuda_cflags or [] + extra_ldflags = extra_ldflags or [] + extra_include_paths = extra_include_paths or [] + + # include cpp files + cpp_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cpp_files] + cpp_sources = [f'#include "{path}"' for path in cpp_paths] + cpp_sources += [_make_wrapper(tup) for tup in cpp_wrappers] + + # include cuda files + cuda_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cuda_files] + cuda_sources = [f'#include "{path}"' for path in cuda_paths] + cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] + + return load_inline( + "sgl_kernel_jit_" + "_".join(str(arg) for arg in args), + cpp_sources=cpp_sources, + cuda_sources=cuda_sources, + extra_cflags=DEFAULT_CFLAGS + extra_cflags, + extra_cuda_cflags=DEFAULT_CUDA_CFLAGS + extra_cuda_cflags, + extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, + build_directory=build_directory, + ) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index f8c456c19..affb23a79 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -7,6 +7,7 @@ from typing import Optional import psutil import torch +from sglang.jit_kernel.hicache import transfer_hicache_one_layer from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool, MLATokenToKVPool from sglang.srt.utils import is_npu, is_xpu @@ -14,14 +15,12 @@ _is_npu = is_npu() _is_xpu = is_xpu() if not (_is_npu or _is_xpu): from sgl_kernel.kvcacheio import ( - transfer_kv_all_layer, transfer_kv_all_layer_direct_lf_pf, transfer_kv_all_layer_lf_pf, transfer_kv_all_layer_lf_ph, transfer_kv_all_layer_mla, transfer_kv_all_layer_mla_lf_pf, transfer_kv_direct, - transfer_kv_per_layer, transfer_kv_per_layer_direct_pf_lf, transfer_kv_per_layer_mla, transfer_kv_per_layer_mla_pf_lf, @@ -282,14 +281,14 @@ class MHATokenToKVPoolHost(HostKVCache): ): if io_backend == "kernel": if self.layout == "layer_first": - transfer_kv_per_layer( - src_k=self.k_buffer[layer_id], - dst_k=device_pool.k_buffer[layer_id], - src_v=self.v_buffer[layer_id], - dst_v=device_pool.v_buffer[layer_id], - src_indices=host_indices, - dst_indices=device_indices, - item_size=self.token_stride_size, + transfer_hicache_one_layer( + k_cache_dst=device_pool.k_buffer[layer_id], + v_cache_dst=device_pool.v_buffer[layer_id], + k_cache_src=self.k_buffer[layer_id], + v_cache_src=self.v_buffer[layer_id], + indices_dst=device_indices, + indices_src=host_indices, + element_dim=self.head_num * self.head_dim, ) elif self.layout == "page_first": transfer_kv_per_layer_pf_lf( @@ -369,16 +368,17 @@ class MHATokenToKVPoolHost(HostKVCache): ): if io_backend == "kernel": if self.layout == "layer_first": - transfer_kv_all_layer( - src_k_layers=device_pool.k_data_ptrs, - dst_k_layers=self.k_data_ptrs, - src_v_layers=device_pool.v_data_ptrs, - dst_v_layers=self.v_data_ptrs, - src_indices=device_indices, - dst_indices=host_indices, - item_size=self.token_stride_size, - num_layers=self.layer_num, - ) + element_dim = self.head_num * self.head_dim + for layer_id in range(self.layer_num): + transfer_hicache_one_layer( + k_cache_dst=self.k_buffer[layer_id], + v_cache_dst=self.v_buffer[layer_id], + k_cache_src=device_pool.k_buffer[layer_id], + v_cache_src=device_pool.v_buffer[layer_id], + indices_dst=host_indices, + indices_src=device_indices, + element_dim=element_dim, + ) elif self.layout == "page_first": transfer_kv_all_layer_lf_pf( src_k_layers=device_pool.k_data_ptrs,