minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Refactor] Fix test and clean up hicache code (#18555)

Browse Source
This commit is contained in:
DarkSharpness
2026-02-18 14:37:46 +08:00
committed by GitHub
parent 95c44cea29
commit 9d138685c1
3 changed files with 252 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/benchmark/bench_hicache.py
+142 -123
View File
@@ -16,6 +16,8 @@ capture doesn't support CPU-GPU memory transfers.
""" """
import itertools import itertools
import os
from dataclasses import dataclass
from typing import Tuple from typing import Tuple
import torch import torch
@@ -23,17 +25,59 @@ import triton
import triton.testing import triton.testing
from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
from sglang.jit_kernel.benchmark.utils import ( from sglang.jit_kernel.benchmark.utils import DEFAULT_QUANTILES, get_benchmark_range
DEFAULT_DTYPE,
DEFAULT_QUANTILES,
get_benchmark_range,
)
from sglang.jit_kernel.hicache import ( from sglang.jit_kernel.hicache import (
can_use_hicache_jit_kernel, can_use_hicache_jit_kernel,
transfer_hicache_all_layer, transfer_hicache_all_layer,
transfer_hicache_one_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( def sglang_aot_transfer_one(
k_cache_dst: torch.Tensor, 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) 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 # 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( 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], ci_range=[16],
) )
ELEMENT_SIZE_RANGE = get_benchmark_range( ELEMENT_SIZE_RANGE = get_benchmark_range(
@@ -173,9 +193,9 @@ ELEMENT_SIZE_RANGE = get_benchmark_range(
ci_range=[1024], ci_range=[1024],
) )
LINE_VALS = ["aot", "jit", "pytorch", "torch_streams"] LINE_VALS = ["aot", "jit", "pytorch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch", "PyTorch 2 Stream"] LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("red", ":"), ("green", "-.")] STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")]
CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE)) 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 element_size: int, batch_size: int, provider: str
) -> Tuple[float, float, float]: ) -> Tuple[float, float, float]:
"""One Layer: Host (CPU) -> Device (GPU).""" """One Layer: Host (CPU) -> Device (GPU)."""
k_cache_src = torch.randn( global cache
(HOST_CACHE_SIZE, element_size), cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
dtype=DEFAULT_DTYPE, k_cache_src = cache_local.k_cache_host
device="cpu", v_cache_src = cache_local.v_cache_host
pin_memory=True, k_cache_dst = cache_local.k_cache_cuda
) v_cache_dst = cache_local.v_cache_cuda
v_cache_src = torch.randn( # to avoid fluctutation, we set the seed as const
(HOST_CACHE_SIZE, element_size), torch.manual_seed(batch_size * 65536 + element_size)
dtype=DEFAULT_DTYPE, indices_src_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
device="cpu", indices_dst_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
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"
)
indices_src_gpu = torch.randperm(HOST_CACHE_SIZE, device="cuda")[:batch_size] # sort by host indices to improve host access performance
indices_dst_gpu = torch.randperm(GPU_CACHE_SIZE, device="cuda")[:batch_size] 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() indices_src_cpu = indices_src_gpu.cpu()
torch.cuda.synchronize() torch.cuda.synchronize()
element_bytes = element_size * k_cache_src.element_size() element_bytes = element_size * k_cache_src.element_size()
FN_MAP = { FN_MAP = {
"aot": lambda: sglang_aot_transfer_one( "aot": lambda: [
k_cache_dst, sglang_aot_transfer_one(
v_cache_dst, k_cache_dst[i],
indices_dst_gpu, v_cache_dst[i],
k_cache_src, indices_dst_gpu,
v_cache_src, k_cache_src[i],
indices_src_gpu, v_cache_src[i],
element_bytes, indices_src_gpu,
), element_bytes,
"jit": lambda: sglang_jit_transfer_one( )
k_cache_dst, for i in range(NUM_LAYERS)
v_cache_dst, ],
indices_dst_gpu, "jit": lambda: [
k_cache_src, sglang_jit_transfer_one(
v_cache_src, k_cache_dst[i],
indices_src_gpu, v_cache_dst[i],
element_size, indices_dst_gpu,
), k_cache_src[i],
"pytorch": lambda: pytorch_transfer( v_cache_src[i],
k_cache_dst, indices_src_gpu,
v_cache_dst, element_size,
indices_dst_gpu, )
k_cache_src, for i in range(NUM_LAYERS)
v_cache_src, ],
indices_src_cpu, "pytorch": lambda: [
), pytorch_transfer(
"torch_streams": lambda: torch_streams_transfer( k_cache_dst[i],
k_cache_dst, v_cache_dst[i],
v_cache_dst, indices_dst_gpu,
indices_dst_gpu, k_cache_src[i],
k_cache_src, v_cache_src[i],
v_cache_src, indices_src_cpu,
indices_src_cpu, )
), for i in range(NUM_LAYERS)
],
} }
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes): if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan")) return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( if DISABLE_TORCH and provider in ["pytorch"]:
FN_MAP[provider], quantiles=DEFAULT_QUANTILES 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 element_size: int, batch_size: int, provider: str
) -> Tuple[float, float, float]: ) -> Tuple[float, float, float]:
"""All Layer: Device (GPU) -> Host (CPU).""" """All Layer: Device (GPU) -> Host (CPU)."""
k_caches_src = torch.randn( global cache
(NUM_LAYERS, GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" cache_local = cache.get_slice(num_layers=NUM_LAYERS, element_size=element_size)
) k_caches_src = cache_local.k_cache_cuda
v_caches_src = torch.randn( v_caches_src = cache_local.v_cache_cuda
(NUM_LAYERS, GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" k_caches_dst = cache_local.k_cache_host
) v_caches_dst = cache_local.v_cache_host
k_caches_dst = torch.randn( # to avoid fluctutation, we set the seed as const
(NUM_LAYERS, HOST_CACHE_SIZE, element_size), torch.manual_seed(batch_size * 65536 + 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,
)
indices_src_gpu = torch.randperm(GPU_CACHE_SIZE, device="cuda")[:batch_size] indices_src_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
indices_dst_gpu = torch.randperm(HOST_CACHE_SIZE, device="cuda")[:batch_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() indices_dst_cpu = indices_dst_gpu.cpu()
torch.cuda.synchronize() torch.cuda.synchronize()
@@ -368,24 +384,16 @@ def benchmark_all_layer_d2h(
) )
for i in range(NUM_LAYERS) 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): if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan")) return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( if DISABLE_TORCH and provider in ["pytorch"]:
FN_MAP[provider], quantiles=DEFAULT_QUANTILES 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 ( return (
1000 * ms / NUM_LAYERS, 1000 * ms / NUM_LAYERS,
@@ -395,6 +403,17 @@ def benchmark_all_layer_d2h(
if __name__ == "__main__": 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("=" * 60)
print("One Layer: Host -> Device (CPU -> GPU)") print("One Layer: Host -> Device (CPU -> GPU)")
print("=" * 60) print("=" * 60)
python/sglang/jit_kernel/csrc/hicache.cuh
+105 -133
View File
@@ -2,25 +2,19 @@
#include <sgl_kernel/utils.h> #include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh> #include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <dlpack/dlpack.h> #include <dlpack/dlpack.h>
#include <algorithm> #include <algorithm>
#include <concepts>
#include <cstddef>
#include <cstdint> #include <cstdint>
#include <type_traits> #include <type_traits>
namespace device::warp { namespace device {
template <typename T, std::size_t N>
struct device_vec {
T data[N];
};
namespace details { namespace details {
template <std::size_t kUnit> template <int kUnit>
inline constexpr auto get_mem_package() { inline constexpr auto get_mem_package() {
if constexpr (kUnit == 16) { if constexpr (kUnit == 16) {
return uint4{}; return uint4{};
@@ -33,90 +27,95 @@ inline constexpr auto get_mem_package() {
} }
} }
template <std::size_t kBytes, std::size_t kUnit> template <int kUnit>
using mem_package_t = decltype(get_mem_package<kUnit>()); using PackageType = decltype(get_mem_package<kUnit>());
__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { SGL_DEVICE uint1 load_nc(const uint1* __restrict__ src) {
uint32_t tmp; 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}; 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; 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}; 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; 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}; 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; 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 tmp0 = value.x;
uint32_t tmp1 = value.y; 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 tmp0 = value.x;
uint32_t tmp1 = value.y; uint32_t tmp1 = value.y;
uint32_t tmp2 = value.z; uint32_t tmp2 = value.z;
uint32_t tmp3 = value.w; 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 } // namespace details
template <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads> template <int64_t kBytes, uint32_t kNumThreads>
__always_inline __device__ auto load_vec(const void* __restrict__ src) { SGL_DEVICE auto load_vec(const void* __restrict__ src) {
using Package = details::mem_package_t<kBytes, kUnit>; static_assert(kBytes % 128 == 0, "kBytes must be multiple of 128 bytes");
constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; static_assert(128 % kNumThreads == 0, "kNumThreads must divide 128 bytes");
constexpr auto kLoopCount = kBytes / kBytesPerLoop; constexpr uint32_t kLoopCount = kBytes / 128;
static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); using Package = details::PackageType<128 / kNumThreads>;
using Storage = AlignedStorage<Package, kLoopCount>;
const auto src_packed = static_cast<const Package*>(src); const auto src_packed = static_cast<const Package*>(src);
const auto lane_id = threadIdx.x % kThreads; const auto lane_id = threadIdx.x % kNumThreads;
device_vec<Package, kLoopCount> vec; Storage vec;
#pragma unroll kLoopCount #pragma unroll kLoopCount
for (std::size_t i = 0; i < kLoopCount; ++i) { for (uint32_t i = 0; i < kLoopCount; ++i) {
const auto j = i * kThreads + lane_id; const auto j = i * kNumThreads + lane_id;
vec.data[i] = details::load_nc(src_packed + j); vec.data[i] = details::load_nc(&src_packed[j]);
} }
return vec; return vec;
} }
template <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads, typename Tp> template <int64_t kBytes, uint32_t kNumThreads, typename Storage>
__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) { SGL_DEVICE void store_vec(void* __restrict__ dst, const Storage& vec) {
using Package = details::mem_package_t<kBytes, kUnit>; using Package = std::decay_t<decltype(vec.data[0])>;
constexpr auto kBytesPerLoop = sizeof(Package) * kThreads; constexpr uint32_t kBytesPerLoop = sizeof(Package) * kNumThreads;
constexpr auto kLoopCount = kBytes / kBytesPerLoop; constexpr uint32_t kLoopCount = kBytes / kBytesPerLoop;
static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes"); static_assert(kBytes % kBytesPerLoop == 0, "Invalid Storage configuration");
static_assert(std::is_same_v<Tp, device_vec<Package, kLoopCount>>);
const auto dst_packed = static_cast<Package*>(dst); const auto dst_packed = static_cast<Package*>(dst);
const auto lane_id = threadIdx.x % kThreads; const auto lane_id = threadIdx.x % kNumThreads;
#pragma unroll kLoopCount #pragma unroll kLoopCount
for (std::size_t i = 0; i < kLoopCount; ++i) { for (uint32_t i = 0; i < kLoopCount; ++i) {
const auto j = i * kThreads + lane_id; const auto j = i * kNumThreads + lane_id;
details::store_nc(dst_packed + j, vec.data[i]); details::store_nc(&dst_packed[j], vec.data[i]);
} }
} }
} // namespace device::warp } // namespace device
namespace { namespace {
#define SGL_HICACHE_KERNEL __global__ __launch_bounds__(kBlockSize, 1)
struct HicacheKernelParams { struct HicacheKernelParams {
void* __restrict__ k_cache_dst; void* __restrict__ k_cache_dst;
void* __restrict__ v_cache_dst; void* __restrict__ v_cache_dst;
@@ -124,118 +123,89 @@ struct HicacheKernelParams {
void* __restrict__ k_cache_src; void* __restrict__ k_cache_src;
void* __restrict__ v_cache_src; void* __restrict__ v_cache_src;
const void* __restrict__ indices_src; const void* __restrict__ indices_src;
std::size_t length; int64_t kv_cache_src_stride;
std::size_t kv_cache_src_stride; int64_t kv_cache_dst_stride;
std::size_t kv_cache_dst_stride; uint32_t length;
std::size_t num_layers = 0; // only used in all_layer transfer uint32_t num_layers = 0; // only used in all_layer transfer
}; };
template < template <typename T, int64_t kElementSize, uint32_t kUnroll, uint32_t kBlockQuota, uint32_t kBlockSize>
std::integral T, SGL_HICACHE_KERNEL void hicache_transfer_per_layer(const __grid_constant__ HicacheKernelParams params) {
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; using namespace device;
static_assert(kNumThreads % kWarpThreads == 0); static_assert(kBlockSize % kWarpThreads == 0);
static_assert(kWarpThreads % kUnroll == 0); static_assert(kWarpThreads % kUnroll == 0);
constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; constexpr uint32_t kNumThreads = kWarpThreads / kUnroll;
constexpr auto kWarpsPerBlock = kNumThreads / kWarpThreads; constexpr uint32_t kWorkersPerBlock = kBlockSize / kNumThreads;
constexpr auto kWorkers = kWarpsPerBlock * kBlockQuota; constexpr uint32_t kNumWorkers = kWorkersPerBlock * kBlockQuota;
const auto& [ const auto& [
k_cache_dst, v_cache_dst, indices_dst, // dst k_cache_dst, v_cache_dst, indices_dst, // dst
k_cache_src, v_cache_src, indices_src, // src 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; ] = params;
const auto warp_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads;
// force to transfer 128 bytes per iteration const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads;
// since the PCIe transaction size is 128 bytes aligned for (uint32_t i = work_id; i < length; i += kNumWorkers) {
constexpr auto kGranularity = 128 / kWarpThreads;
for (auto i = warp_id; i < length; i += kWorkers) {
const auto pos_src = static_cast<const T*>(indices_src)[i]; const auto pos_src = static_cast<const T*>(indices_src)[i];
const auto pos_dst = static_cast<const T*>(indices_dst)[i]; const auto pos_dst = static_cast<const T*>(indices_dst)[i];
const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); 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 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 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 dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride);
const auto vec_k = warp::load_vec<kElementSize, kGranularity, kWarpThreads>(src_k); const auto vec_k = load_vec<kElementSize, kNumThreads>(src_k);
const auto vec_v = warp::load_vec<kElementSize, kGranularity, kWarpThreads>(src_v); const auto vec_v = load_vec<kElementSize, kNumThreads>(src_v);
warp::store_vec<kElementSize, kGranularity, kWarpThreads>(dst_k, vec_k); store_vec<kElementSize, kNumThreads>(dst_k, vec_k);
warp::store_vec<kElementSize, kGranularity, kWarpThreads>(dst_v, vec_v); store_vec<kElementSize, kNumThreads>(dst_v, vec_v);
} }
} }
template < template <typename T, int64_t kElementSize, uint32_t kUnroll, uint32_t kBlockQuota, uint32_t kBlockSize>
std::integral T, SGL_HICACHE_KERNEL void hicache_transfer_all_layer(const __grid_constant__ HicacheKernelParams params) {
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 namespace device;
using src_ptr_t = std::add_pointer_t<const void* const>; using src_ptr_t = const void*;
using dst_ptr_t = std::add_pointer_t<void* const>; using dst_ptr_t = void*;
static_assert(kNumThreads % kWarpThreads == 0); static_assert(kBlockSize % kWarpThreads == 0);
constexpr auto kWarpThreads = device::kWarpThreads / kUnroll; static_assert(kWarpThreads % kUnroll == 0);
constexpr auto kWarpsPerBlock = static_cast<uint32_t>(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& [ const auto& [
k_ptr_dst, v_ptr_dst, indices_dst, // dst k_ptr_dst, v_ptr_dst, indices_dst, // dst
k_ptr_src, v_ptr_src, indices_src, // src 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; ] = params;
const auto warp_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads;
// force to transfer 128 bytes per iteration const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads;
// since the PCIe transaction size is 128 bytes aligned for (uint32_t i = work_id; i < length; i += kNumWorkers) {
constexpr auto kGranularity = 128 / kWarpThreads;
for (auto i = warp_id; i < length; i += kWorkers) {
const auto pos_src = static_cast<const T*>(indices_src)[i]; const auto pos_src = static_cast<const T*>(indices_src)[i];
const auto pos_dst = static_cast<const T*>(indices_dst)[i]; const auto pos_dst = static_cast<const T*>(indices_dst)[i];
for (std::size_t layer = 0; layer < num_layers; ++layer) { for (uint32_t layer = 0; layer < num_layers; ++layer) {
const auto k_cache_src = static_cast<src_ptr_t>(k_ptr_src)[layer]; const auto k_cache_src = static_cast<const src_ptr_t*>(k_ptr_src)[layer];
const auto v_cache_src = static_cast<src_ptr_t>(v_ptr_src)[layer]; const auto v_cache_src = static_cast<const src_ptr_t*>(v_ptr_src)[layer];
const auto k_cache_dst = static_cast<dst_ptr_t>(k_ptr_dst)[layer]; const auto k_cache_dst = static_cast<const dst_ptr_t*>(k_ptr_dst)[layer];
const auto v_cache_dst = static_cast<dst_ptr_t>(v_ptr_dst)[layer]; const auto v_cache_dst = static_cast<const dst_ptr_t*>(v_ptr_dst)[layer];
const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); 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 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 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 dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride);
const auto vec_k = warp::load_vec<kElementSize, kGranularity, kWarpThreads>(src_k); const auto vec_k = load_vec<kElementSize, kNumThreads>(src_k);
const auto vec_v = warp::load_vec<kElementSize, kGranularity, kWarpThreads>(src_v); const auto vec_v = load_vec<kElementSize, kNumThreads>(src_v);
warp::store_vec<kElementSize, kGranularity, kWarpThreads>(dst_k, vec_k); store_vec<kElementSize, kNumThreads>(dst_k, vec_k);
warp::store_vec<kElementSize, kGranularity, kWarpThreads>(dst_v, vec_v); store_vec<kElementSize, kNumThreads>(dst_v, vec_v);
} }
} }
} }
template < template <int64_t kElementSize, uint32_t kUnroll, uint32_t kBlockQuota, uint32_t kBlockSize>
std::size_t kElementSize,
std::size_t kUnroll,
std::size_t kBlockQuota,
std::size_t kNumThreads,
std::size_t kMaxOccupancy>
struct HiCacheKernel { struct HiCacheKernel {
template <typename T> template <typename T>
static constexpr auto _kernel_one = static constexpr auto kernel_one = hicache_transfer_per_layer<T, kElementSize, kUnroll, kBlockQuota, kBlockSize>;
hicache_transfer_per_layer<T, kElementSize, kUnroll, kBlockQuota, kNumThreads, kMaxOccupancy>;
template <typename T> template <typename T>
static constexpr auto _kernel_all = static constexpr auto kernel_all = hicache_transfer_all_layer<T, kElementSize, kUnroll, kBlockQuota, kBlockSize>;
hicache_transfer_all_layer<T, kElementSize, kUnroll, kBlockQuota, kNumThreads, kMaxOccupancy>;
static void run_one( static void run_one(
const tvm::ffi::TensorView k_cache_dst, 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 v_cache_src_ptr = v_cache_src.data_ptr();
const auto indices_dst_ptr = indices_dst.data_ptr(); const auto indices_dst_ptr = indices_dst.data_ptr();
const auto indices_src_ptr = indices_src.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr();
const auto length = static_cast<std::size_t>(L.unwrap()); const auto length = static_cast<uint32_t>(L.unwrap());
const auto kv_cache_src_stride = static_cast<std::size_t>(N.unwrap()) * dtype_size; const auto kv_cache_src_stride = static_cast<int64_t>(N.unwrap() * dtype_size);
const auto kv_cache_dst_stride = static_cast<std::size_t>(M.unwrap()) * dtype_size; const auto kv_cache_dst_stride = static_cast<int64_t>(M.unwrap() * dtype_size);
const auto use_int32 = indices_dtype.unwrap().bits == 32; const auto use_int32 = indices_dtype.unwrap().bits == 32;
const auto device = indices_device.unwrap(); 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 num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota);
const auto params = HicacheKernelParams{ const auto params = HicacheKernelParams{
.k_cache_dst = k_cache_dst_ptr, .k_cache_dst = k_cache_dst_ptr,
@@ -298,12 +268,12 @@ struct HiCacheKernel {
.k_cache_src = k_cache_src_ptr, .k_cache_src = k_cache_src_ptr,
.v_cache_src = v_cache_src_ptr, .v_cache_src = v_cache_src_ptr,
.indices_src = indices_src_ptr, .indices_src = indices_src_ptr,
.length = length,
.kv_cache_src_stride = kv_cache_src_stride, .kv_cache_src_stride = kv_cache_src_stride,
.kv_cache_dst_stride = kv_cache_dst_stride, .kv_cache_dst_stride = kv_cache_dst_stride,
.length = length,
}; };
const auto kernel = use_int32 ? _kernel_one<int32_t> : _kernel_one<int64_t>; const auto kernel = use_int32 ? kernel_one<int32_t> : kernel_one<int64_t>;
LaunchKernel(num_blocks, kNumThreads, device)(kernel, params); LaunchKernel(num_blocks, kBlockSize, device)(kernel, params);
} }
static void run_all( static void run_all(
@@ -313,8 +283,8 @@ struct HiCacheKernel {
const tvm::ffi::TensorView k_ptr_src, const tvm::ffi::TensorView k_ptr_src,
const tvm::ffi::TensorView v_ptr_src, const tvm::ffi::TensorView v_ptr_src,
const tvm::ffi::TensorView indices_src, const tvm::ffi::TensorView indices_src,
const std::size_t kv_src_stride, const int64_t kv_src_stride_bytes,
const std::size_t kv_dst_stride) { const int64_t kv_dst_stride_bytes) {
using namespace host; using namespace host;
auto N = SymbolicSize{"num_layers"}; auto N = SymbolicSize{"num_layers"};
@@ -342,11 +312,11 @@ struct HiCacheKernel {
const auto v_cache_src_ptr = v_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_dst_ptr = indices_dst.data_ptr();
const auto indices_src_ptr = indices_src.data_ptr(); const auto indices_src_ptr = indices_src.data_ptr();
const auto length = static_cast<std::size_t>(L.unwrap()); const auto length = static_cast<uint32_t>(L.unwrap());
const auto use_int32 = dtype_.unwrap().bits == 32; const auto use_int32 = dtype_.unwrap().bits == 32;
const auto device = device_.unwrap(); 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 num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota);
const auto params = HicacheKernelParams{ const auto params = HicacheKernelParams{
.k_cache_dst = k_cache_dst_ptr, .k_cache_dst = k_cache_dst_ptr,
@@ -355,14 +325,16 @@ struct HiCacheKernel {
.k_cache_src = k_cache_src_ptr, .k_cache_src = k_cache_src_ptr,
.v_cache_src = v_cache_src_ptr, .v_cache_src = v_cache_src_ptr,
.indices_src = indices_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, .length = length,
.kv_cache_src_stride = kv_src_stride, .num_layers = static_cast<uint32_t>(N.unwrap()),
.kv_cache_dst_stride = kv_dst_stride,
.num_layers = static_cast<std::size_t>(N.unwrap()),
}; };
const auto kernel = use_int32 ? _kernel_all<int32_t> : _kernel_all<int64_t>; const auto kernel = use_int32 ? kernel_all<int32_t> : kernel_all<int64_t>;
LaunchKernel(num_blocks, kNumThreads, device)(kernel, params); LaunchKernel(num_blocks, kBlockSize, device)(kernel, params);
} }
}; };
#undef SGL_HICACHE_KERNEL
} // namespace } // namespace
python/sglang/jit_kernel/hicache.py
+5 -4
View File
@@ -14,13 +14,11 @@ DEFAULT_BLOCK_QUOTA = 2
@cache_once @cache_once
def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> Module: def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> Module:
num_threads, occupancy = 1024, 1
args = make_cpp_args( args = make_cpp_args(
element_size, element_size,
unroll, unroll,
block_quota, block_quota,
num_threads, 1024, # num_threads, can be tuned for performance
occupancy,
) )
return load_jit( return load_jit(
"hicache", "hicache",
@@ -39,6 +37,10 @@ def can_use_hicache_jit_kernel(
unroll: int | None = None, # can be tuned for performance unroll: int | None = None, # can be tuned for performance
block_quota: int | None = None, # can be tuned for less interference block_quota: int | None = None, # can be tuned for less interference
) -> bool: ) -> bool:
logger = logging.getLogger(__name__)
if element_size % 128 != 0:
logger.warning(f"Unsupported {element_size = } for JIT HiCache kernel")
return False
try: try:
unroll = unroll or _default_unroll(element_size) unroll = unroll or _default_unroll(element_size)
block_quota = block_quota or DEFAULT_BLOCK_QUOTA block_quota = block_quota or DEFAULT_BLOCK_QUOTA
@@ -49,7 +51,6 @@ def can_use_hicache_jit_kernel(
) )
return True return True
except Exception as e: except Exception as e:
logger = logging.getLogger(__name__)
logger.warning(f"Failed to load JIT HiCache kernel: {e}") logger.warning(f"Failed to load JIT HiCache kernel: {e}")
return False return False