diff --git a/python/sglang/jit_kernel/benchmark/bench_hicache.py b/python/sglang/jit_kernel/benchmark/bench_hicache.py index 0218b562a..8929ce706 100644 --- a/python/sglang/jit_kernel/benchmark/bench_hicache.py +++ b/python/sglang/jit_kernel/benchmark/bench_hicache.py @@ -16,6 +16,8 @@ capture doesn't support CPU-GPU memory transfers. """ import itertools +import os +from dataclasses import dataclass from typing import Tuple import torch @@ -23,17 +25,59 @@ import triton import triton.testing from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer -from sglang.jit_kernel.benchmark.utils import ( - DEFAULT_DTYPE, - DEFAULT_QUANTILES, - get_benchmark_range, -) +from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range from sglang.jit_kernel.hicache import ( can_use_hicache_jit_kernel, transfer_hicache_all_layer, transfer_hicache_one_layer, ) +# NOTE: Adjustable hyperparameters for better benchmark stability + +# NOTE: torch impl is too slow in benchmark +DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1" +PAGE_SIZE = 1 +ENABLE_SORT = True +GPU_CACHE_SIZE = 256 * 1024 # 256K tokens on GPU +HOST_CACHE_SIZE = 512 * 1024 # 512K tokens on CPU +NUM_LAYERS = 8 + + +@dataclass(frozen=True) +class HiCacheCache: + k_cache_cuda: torch.Tensor + v_cache_cuda: torch.Tensor + k_cache_host: torch.Tensor + v_cache_host: torch.Tensor + + def get_slice(self, num_layers: int, element_size: int) -> "HiCacheCache": + def slice_cuda(t: torch.Tensor) -> torch.Tensor: + needed_cuda = num_layers * GPU_CACHE_SIZE + return t.view(-1, element_size)[:needed_cuda].unflatten(0, (num_layers, -1)) + + def slice_host(t: torch.Tensor) -> torch.Tensor: + needed_host = num_layers * HOST_CACHE_SIZE + return t.view(-1, element_size)[:needed_host].unflatten(0, (num_layers, -1)) + + return HiCacheCache( + k_cache_cuda=slice_cuda(self.k_cache_cuda), + v_cache_cuda=slice_cuda(self.v_cache_cuda), + k_cache_host=slice_host(self.k_cache_host), + v_cache_host=slice_host(self.v_cache_host), + ) + + +def gen_indices( + size: int, max_size: int, *, page_size: int = PAGE_SIZE +) -> torch.Tensor: + def align(x: int) -> int: + return (x + page_size - 1) // page_size + + assert size <= max_size and max_size % page_size == 0 + indices = torch.randperm(align(max_size))[: align(size)] + offsets = torch.arange(page_size) + return (indices[:, None] * page_size + offsets).flatten().cuda()[:size] + def sglang_aot_transfer_one( k_cache_dst: torch.Tensor, @@ -138,34 +182,10 @@ def pytorch_transfer( v_cache_dst[indices_dst_on_dst] = v_cache_src[indices_src_on_src].to(dst_device) -alt_stream = torch.cuda.Stream() - - -def torch_streams_transfer( - k_cache_dst: torch.Tensor, - v_cache_dst: torch.Tensor, - indices_dst_on_dst: torch.Tensor, - k_cache_src: torch.Tensor, - v_cache_src: torch.Tensor, - indices_src_on_src: torch.Tensor, -) -> None: - """PyTorch 2 Stream baseline.""" - dst_device = k_cache_dst.device - current_stream = torch.cuda.current_stream() - alt_stream.wait_stream(current_stream) - k_cache_dst[indices_dst_on_dst] = k_cache_src[indices_src_on_src].to(dst_device) - with torch.cuda.stream(alt_stream): - v_cache_dst[indices_dst_on_dst] = v_cache_src[indices_src_on_src].to(dst_device) - current_stream.wait_stream(alt_stream) - - # Benchmark configuration -GPU_CACHE_SIZE = 32 * 1024 # 32K tokens on GPU -HOST_CACHE_SIZE = 128 * 1024 # 128K tokens on CPU -NUM_LAYERS = 8 BS_RANGE = get_benchmark_range( - full_range=[2**n for n in range(0, 15)], + full_range=[2**n for n in range(0, 16)], ci_range=[16], ) ELEMENT_SIZE_RANGE = get_benchmark_range( @@ -173,9 +193,9 @@ ELEMENT_SIZE_RANGE = get_benchmark_range( ci_range=[1024], ) -LINE_VALS = ["aot", "jit", "pytorch", "torch_streams"] -LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch", "PyTorch 2 Stream"] -STYLES = [("orange", "-"), ("blue", "--"), ("red", ":"), ("green", "-.")] +LINE_VALS = ["aot", "jit", "pytorch"] +LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"] +STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")] CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE)) @@ -202,76 +222,78 @@ def benchmark_one_layer_h2d( element_size: int, batch_size: int, provider: str ) -> Tuple[float, float, float]: """One Layer: Host (CPU) -> Device (GPU).""" - k_cache_src = torch.randn( - (HOST_CACHE_SIZE, element_size), - dtype=DEFAULT_DTYPE, - device="cpu", - pin_memory=True, - ) - v_cache_src = torch.randn( - (HOST_CACHE_SIZE, element_size), - dtype=DEFAULT_DTYPE, - device="cpu", - pin_memory=True, - ) - k_cache_dst = torch.randn( - (GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" - ) - v_cache_dst = torch.randn( - (GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" - ) + global cache + cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size) + k_cache_src = cache_local.k_cache_host + v_cache_src = cache_local.v_cache_host + k_cache_dst = cache_local.k_cache_cuda + v_cache_dst = cache_local.v_cache_cuda + # to avoid fluctutation, we set the seed as const + torch.manual_seed(batch_size * 65536 + element_size) + indices_src_gpu = gen_indices(batch_size, HOST_CACHE_SIZE) + indices_dst_gpu = gen_indices(batch_size, GPU_CACHE_SIZE) - indices_src_gpu = torch.randperm(HOST_CACHE_SIZE, device="cuda")[:batch_size] - indices_dst_gpu = torch.randperm(GPU_CACHE_SIZE, device="cuda")[:batch_size] + # sort by host indices to improve host access performance + if ENABLE_SORT: + indices_src_gpu, mapping = indices_src_gpu.sort() + indices_dst_gpu = indices_dst_gpu[mapping] indices_src_cpu = indices_src_gpu.cpu() torch.cuda.synchronize() element_bytes = element_size * k_cache_src.element_size() FN_MAP = { - "aot": lambda: sglang_aot_transfer_one( - k_cache_dst, - v_cache_dst, - indices_dst_gpu, - k_cache_src, - v_cache_src, - indices_src_gpu, - element_bytes, - ), - "jit": lambda: sglang_jit_transfer_one( - k_cache_dst, - v_cache_dst, - indices_dst_gpu, - k_cache_src, - v_cache_src, - indices_src_gpu, - element_size, - ), - "pytorch": lambda: pytorch_transfer( - k_cache_dst, - v_cache_dst, - indices_dst_gpu, - k_cache_src, - v_cache_src, - indices_src_cpu, - ), - "torch_streams": lambda: torch_streams_transfer( - k_cache_dst, - v_cache_dst, - indices_dst_gpu, - k_cache_src, - v_cache_src, - indices_src_cpu, - ), + "aot": lambda: [ + sglang_aot_transfer_one( + k_cache_dst[i], + v_cache_dst[i], + indices_dst_gpu, + k_cache_src[i], + v_cache_src[i], + indices_src_gpu, + element_bytes, + ) + for i in range(NUM_LAYERS) + ], + "jit": lambda: [ + sglang_jit_transfer_one( + k_cache_dst[i], + v_cache_dst[i], + indices_dst_gpu, + k_cache_src[i], + v_cache_src[i], + indices_src_gpu, + element_size, + ) + for i in range(NUM_LAYERS) + ], + "pytorch": lambda: [ + pytorch_transfer( + k_cache_dst[i], + v_cache_dst[i], + indices_dst_gpu, + k_cache_src[i], + v_cache_src[i], + indices_src_cpu, + ) + for i in range(NUM_LAYERS) + ], } if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes): return (float("nan"), float("nan"), float("nan")) - ms, min_ms, max_ms = triton.testing.do_bench( - FN_MAP[provider], quantiles=DEFAULT_QUANTILES + if DISABLE_TORCH and provider in ["pytorch"]: + return (float("nan"), float("nan"), float("nan")) + + ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore + FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25 + ) + return ( + 1000 * ms / NUM_LAYERS, + 1000 * max_ms / NUM_LAYERS, + 1000 * min_ms / NUM_LAYERS, ) - return 1000 * ms, 1000 * max_ms, 1000 * min_ms # ============================================================================= @@ -305,27 +327,21 @@ def benchmark_all_layer_d2h( element_size: int, batch_size: int, provider: str ) -> Tuple[float, float, float]: """All Layer: Device (GPU) -> Host (CPU).""" - k_caches_src = torch.randn( - (NUM_LAYERS, GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" - ) - v_caches_src = torch.randn( - (NUM_LAYERS, GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" - ) - k_caches_dst = torch.randn( - (NUM_LAYERS, HOST_CACHE_SIZE, element_size), - dtype=DEFAULT_DTYPE, - device="cpu", - pin_memory=True, - ) - v_caches_dst = torch.randn( - (NUM_LAYERS, HOST_CACHE_SIZE, element_size), - dtype=DEFAULT_DTYPE, - device="cpu", - pin_memory=True, - ) + global cache + cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size) + k_caches_src = cache_local.k_cache_cuda + v_caches_src = cache_local.v_cache_cuda + k_caches_dst = cache_local.k_cache_host + v_caches_dst = cache_local.v_cache_host + # to avoid fluctutation, we set the seed as const + torch.manual_seed(batch_size * 65536 + element_size) - indices_src_gpu = torch.randperm(GPU_CACHE_SIZE, device="cuda")[:batch_size] - indices_dst_gpu = torch.randperm(HOST_CACHE_SIZE, device="cuda")[:batch_size] + indices_src_gpu = gen_indices(batch_size, GPU_CACHE_SIZE) + indices_dst_gpu = gen_indices(batch_size, HOST_CACHE_SIZE) + # sort by host indices to improve host access performance + if ENABLE_SORT: + indices_dst_gpu, mapping = indices_dst_gpu.sort() + indices_src_gpu = indices_src_gpu[mapping] indices_dst_cpu = indices_dst_gpu.cpu() torch.cuda.synchronize() @@ -368,24 +384,16 @@ def benchmark_all_layer_d2h( ) for i in range(NUM_LAYERS) ], - "torch_streams": lambda: [ - torch_streams_transfer( - k_caches_dst[i], - v_caches_dst[i], - indices_dst_cpu, - k_caches_src[i], - v_caches_src[i], - indices_src_gpu, - ) - for i in range(NUM_LAYERS) - ], } if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes): return (float("nan"), float("nan"), float("nan")) - ms, min_ms, max_ms = triton.testing.do_bench( - FN_MAP[provider], quantiles=DEFAULT_QUANTILES + if DISABLE_TORCH and provider in ["pytorch"]: + return (float("nan"), float("nan"), float("nan")) + + ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore + FN_MAP[provider], quantiles=DEFAULT_QUANTILES, warmup=5, rep=25 ) return ( 1000 * ms / NUM_LAYERS, @@ -395,6 +403,17 @@ def benchmark_all_layer_d2h( if __name__ == "__main__": + MAX_SIZE = max(ELEMENT_SIZE_RANGE) + DEVICE_SHAPE = (NUM_LAYERS * GPU_CACHE_SIZE, MAX_SIZE) + HOST_SHAPE = (NUM_LAYERS * HOST_CACHE_SIZE, MAX_SIZE) + + cache = HiCacheCache( + k_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"), + v_cache_cuda=torch.empty(DEVICE_SHAPE, dtype=torch.bfloat16, device="cuda"), + k_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True), + v_cache_host=torch.empty(HOST_SHAPE, dtype=torch.bfloat16, pin_memory=True), + ) + print("=" * 60) print("One Layer: Host -> Device (CPU -> GPU)") print("=" * 60) diff --git a/python/sglang/jit_kernel/csrc/hicache.cuh b/python/sglang/jit_kernel/csrc/hicache.cuh index 555282ee4..04f093a02 100644 --- a/python/sglang/jit_kernel/csrc/hicache.cuh +++ b/python/sglang/jit_kernel/csrc/hicache.cuh @@ -2,25 +2,19 @@ #include #include +#include #include #include -#include -#include #include #include -namespace device::warp { - -template -struct device_vec { - T data[N]; -}; +namespace device { namespace details { -template +template inline constexpr auto get_mem_package() { if constexpr (kUnit == 16) { return uint4{}; @@ -33,90 +27,95 @@ inline constexpr auto get_mem_package() { } } -template -using mem_package_t = decltype(get_mem_package()); +template +using PackageType = decltype(get_mem_package()); -__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { +SGL_DEVICE uint1 load_nc(const uint1* __restrict__ src) { uint32_t tmp; - asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); + asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; } -__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { +SGL_DEVICE uint2 load_nc(const uint2* __restrict__ src) { uint32_t tmp0, tmp1; - asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); + asm volatile("ld.global.L1::no_allocate.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 { +SGL_DEVICE uint4 load_nc(const uint4* __restrict__ src) { 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)); + asm volatile("ld.global.L1::no_allocate.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) { +SGL_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)); + asm volatile("st.global.L1::no_allocate.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); } -__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { +SGL_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)); + asm volatile("st.global.L1::no_allocate.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); } -__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { +SGL_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)); + asm volatile( + "st.global.L1::no_allocate.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); } } // namespace details -template -__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"); +template +SGL_DEVICE auto load_vec(const void* __restrict__ src) { + static_assert(kBytes % 128 == 0, "kBytes must be multiple of 128 bytes"); + static_assert(128 % kNumThreads == 0, "kNumThreads must divide 128 bytes"); + constexpr uint32_t kLoopCount = kBytes / 128; + using Package = details::PackageType<128 / kNumThreads>; + using Storage = AlignedStorage; const auto src_packed = static_cast(src); - const auto lane_id = threadIdx.x % kThreads; - device_vec vec; + const auto lane_id = threadIdx.x % kNumThreads; + Storage 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); + for (uint32_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kNumThreads + lane_id; + vec.data[i] = details::load_nc(&src_packed[j]); } return vec; } -template -__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>); +template +SGL_DEVICE void store_vec(void* __restrict__ dst, const Storage& vec) { + using Package = std::decay_t; + constexpr uint32_t kBytesPerLoop = sizeof(Package) * kNumThreads; + constexpr uint32_t kLoopCount = kBytes / kBytesPerLoop; + static_assert(kBytes % kBytesPerLoop == 0, "Invalid Storage configuration"); const auto dst_packed = static_cast(dst); - const auto lane_id = threadIdx.x % kThreads; + const auto lane_id = threadIdx.x % kNumThreads; #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]); + for (uint32_t i = 0; i < kLoopCount; ++i) { + const auto j = i * kNumThreads + lane_id; + details::store_nc(&dst_packed[j], vec.data[i]); } } -} // namespace device::warp +} // namespace device namespace { +#define SGL_HICACHE_KERNEL __global__ __launch_bounds__(kBlockSize, 1) + struct HicacheKernelParams { void* __restrict__ k_cache_dst; void* __restrict__ v_cache_dst; @@ -124,118 +123,89 @@ struct HicacheKernelParams { 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 + int64_t kv_cache_src_stride; + int64_t kv_cache_dst_stride; + uint32_t length; + uint32_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 +template +SGL_HICACHE_KERNEL void hicache_transfer_per_layer(const __grid_constant__ HicacheKernelParams params) { using namespace device; - static_assert(kNumThreads % kWarpThreads == 0); + static_assert(kBlockSize % kWarpThreads == 0); static_assert(kWarpThreads % kUnroll == 0); - constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; - constexpr auto kWarpsPerBlock = kNumThreads / kWarpThreads; - constexpr auto kWorkers = kWarpsPerBlock * kBlockQuota; + constexpr uint32_t kNumThreads = kWarpThreads / kUnroll; + constexpr uint32_t kWorkersPerBlock = kBlockSize / kNumThreads; + constexpr uint32_t kNumWorkers = kWorkersPerBlock * 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 + kv_cache_src_stride, kv_cache_dst_stride, length, _ // 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 uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads; + for (uint32_t i = work_id; i < length; i += kNumWorkers) { 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); + const auto vec_k = load_vec(src_k); + const auto vec_v = load_vec(src_v); + store_vec(dst_k, vec_k); + 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 +template +SGL_HICACHE_KERNEL void hicache_transfer_all_layer(const __grid_constant__ HicacheKernelParams params) { using namespace device; - using src_ptr_t = std::add_pointer_t; - using dst_ptr_t = std::add_pointer_t; + using src_ptr_t = const void*; + using dst_ptr_t = void*; - static_assert(kNumThreads % kWarpThreads == 0); - constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; - constexpr auto kWarpsPerBlock = static_cast(kNumThreads) / kWarpThreads; - constexpr auto kWorkers = kWarpsPerBlock * kBlockQuota; + static_assert(kBlockSize % kWarpThreads == 0); + static_assert(kWarpThreads % kUnroll == 0); + + constexpr uint32_t kNumThreads = kWarpThreads / kUnroll; + constexpr uint32_t kWorkersPerBlock = kBlockSize / kNumThreads; + constexpr uint32_t kNumWorkers = kWorkersPerBlock * 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 + kv_cache_src_stride, kv_cache_dst_stride, length, 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 uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads; + for (uint32_t i = work_id; i < length; i += kNumWorkers) { 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]; + for (uint32_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); + const auto vec_k = load_vec(src_k); + const auto vec_v = load_vec(src_v); + store_vec(dst_k, vec_k); + 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> +template struct HiCacheKernel { template - static constexpr auto _kernel_one = - hicache_transfer_per_layer; + static constexpr auto kernel_one = hicache_transfer_per_layer; template - static constexpr auto _kernel_all = - hicache_transfer_all_layer; + static constexpr auto kernel_all = hicache_transfer_all_layer; static void run_one( const tvm::ffi::TensorView k_cache_dst, @@ -283,13 +253,13 @@ struct HiCacheKernel { 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 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); + constexpr auto kWorkersPerBlock = kBlockSize / (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, @@ -298,12 +268,12 @@ struct HiCacheKernel { .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, + .length = length, }; - const auto kernel = use_int32 ? _kernel_one : _kernel_one; - LaunchKernel(num_blocks, kNumThreads, device)(kernel, params); + const auto kernel = use_int32 ? kernel_one : kernel_one; + LaunchKernel(num_blocks, kBlockSize, device)(kernel, params); } static void run_all( @@ -313,8 +283,8 @@ struct HiCacheKernel { 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) { + const int64_t kv_src_stride_bytes, + const int64_t kv_dst_stride_bytes) { using namespace host; auto N = SymbolicSize{"num_layers"}; @@ -342,11 +312,11 @@ struct HiCacheKernel { 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 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); + constexpr auto kWorkersPerBlock = kBlockSize / (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, @@ -355,14 +325,16 @@ struct HiCacheKernel { .k_cache_src = k_cache_src_ptr, .v_cache_src = v_cache_src_ptr, .indices_src = indices_src_ptr, + .kv_cache_src_stride = kv_src_stride_bytes, + .kv_cache_dst_stride = kv_dst_stride_bytes, .length = length, - .kv_cache_src_stride = kv_src_stride, - .kv_cache_dst_stride = kv_dst_stride, - .num_layers = static_cast(N.unwrap()), + .num_layers = static_cast(N.unwrap()), }; - const auto kernel = use_int32 ? _kernel_all : _kernel_all; - LaunchKernel(num_blocks, kNumThreads, device)(kernel, params); + const auto kernel = use_int32 ? kernel_all : kernel_all; + LaunchKernel(num_blocks, kBlockSize, device)(kernel, params); } }; +#undef SGL_HICACHE_KERNEL + } // namespace diff --git a/python/sglang/jit_kernel/hicache.py b/python/sglang/jit_kernel/hicache.py index 828b790b2..aa78d5cb4 100644 --- a/python/sglang/jit_kernel/hicache.py +++ b/python/sglang/jit_kernel/hicache.py @@ -14,13 +14,11 @@ DEFAULT_BLOCK_QUOTA = 2 @cache_once 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, + 1024, # num_threads, can be tuned for performance ) return load_jit( "hicache", @@ -39,6 +37,10 @@ def can_use_hicache_jit_kernel( unroll: int | None = None, # can be tuned for performance block_quota: int | None = None, # can be tuned for less interference ) -> bool: + logger = logging.getLogger(__name__) + if element_size % 128 != 0: + logger.warning(f"Unsupported {element_size = } for JIT HiCache kernel") + return False try: unroll = unroll or _default_unroll(element_size) block_quota = block_quota or DEFAULT_BLOCK_QUOTA @@ -49,7 +51,6 @@ def can_use_hicache_jit_kernel( ) return True except Exception as e: - logger = logging.getLogger(__name__) logger.warning(f"Failed to load JIT HiCache kernel: {e}") return False