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

[Kernel] Fuse KV-cache writes for asymmetric K/V (head_dim != v_head_dim) (#32813)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-30 00:26:10 -07:00
committed by GitHub
parent 2625fdfe6b
commit c192145830
6 changed files with 511 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh
+146 -39
View File
@@ -21,7 +21,9 @@ struct StoreKVCacheParams {
const void* __restrict__ indices;
int64_t stride_k_bytes;
int64_t stride_v_bytes;
int64_t stride_cache_bytes;
// Independent slot strides: head_dim != v_head_dim gives K and V different row widths.
int64_t stride_k_cache_bytes;
int64_t stride_v_cache_bytes;
int64_t stride_indices;
uint32_t batch_size;
int64_t size_limit;
@@ -31,9 +33,37 @@ struct StoreKVCacheParams {
constexpr uint32_t kNumWarps = 4;
constexpr uint32_t kThreadsPerBlock = kNumWarps * device::kWarpThreads;
/**
* \brief How a warp vectorizes one row of kElementBytes: the widest aligned
* vector type it can use, and how many full loop iterations that takes.
* Shared by the interleaved and single-row copies so the two cannot drift.
* kElementBytes == 0 is a valid (empty) plan, so a zero-width tail can be
* queried before being branched away.
*/
template <int64_t kElementBytes>
struct RowVecPlan {
static constexpr int64_t kAlignment = (kElementBytes % (16 * device::kWarpThreads) == 0) ? 16
: kElementBytes % (8 * device::kWarpThreads) == 0 ? 8
: kElementBytes % (4 * device::kWarpThreads) == 0 ? 4
: kElementBytes % 4 == 0 ? 4
: 0;
static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes");
using vec_t = device::AlignedStorage<uint32_t, kAlignment / 4>;
static constexpr int64_t kLoopBytes = sizeof(vec_t) * device::kWarpThreads;
static constexpr int64_t kLoopCount = kElementBytes / kLoopBytes;
static constexpr int64_t kElementCount = kElementBytes / sizeof(vec_t);
static constexpr bool kHasEpilogue = kLoopCount * kLoopBytes < kElementBytes;
};
/**
* \brief Use a single warp to copy key and value data from source to destination.
* Each thread in the warp copies a portion of the data in a coalesced manner.
* Both loads are issued before either store: the two rows live in different
* tensors, and the params' __restrict__ does not survive into the kernel body,
* so the compiler cannot prove k_dst and v_src disjoint and will not sink the
* V load past the K store on its own.
* \tparam kElementBytes The size of each key/value element in bytes.
* \param k_src Pointer to the source key data.
* \param v_src Pointer to the source value data.
@@ -47,17 +77,9 @@ SGL_DEVICE void copy_kv_warp(
void* __restrict__ k_dst,
void* __restrict__ v_dst) {
using namespace device;
constexpr int64_t kAlignment = (kElementBytes % (16 * kWarpThreads) == 0) ? 16
: kElementBytes % (8 * kWarpThreads) == 0 ? 8
: kElementBytes % (4 * kWarpThreads) == 0 ? 4
: kElementBytes % 4 == 0 ? 4
: 0;
static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes");
using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
constexpr auto kLoopCount = kElementBytes / kLoopBytes;
using plan_t = RowVecPlan<kElementBytes>;
using vec_t = typename plan_t::vec_t;
constexpr auto kLoopCount = plan_t::kLoopCount;
const auto gmem = tile::Memory<vec_t>::warp();
@@ -70,8 +92,8 @@ SGL_DEVICE void copy_kv_warp(
}
// handle the epilogue if any
if constexpr (kLoopCount * kLoopBytes < kElementBytes) {
if (gmem.in_bound(kElementBytes / sizeof(vec_t), kLoopCount)) {
if constexpr (plan_t::kHasEpilogue) {
if (gmem.in_bound(plan_t::kElementCount, kLoopCount)) {
const auto k = gmem.load(k_src, kLoopCount);
const auto v = gmem.load(v_src, kLoopCount);
gmem.store(k_dst, k, kLoopCount);
@@ -80,24 +102,100 @@ SGL_DEVICE void copy_kv_warp(
}
}
/**
* \brief Use a single warp to copy one row from source to destination.
* Serves the width by which asymmetric K/V rows differ, which has no counterpart
* row to interleave with.
* \tparam kElementBytes The size of the row in bytes.
* \param src Pointer to the source data.
* \param dst Pointer to the destination data.
*/
template <int64_t kElementBytes>
SGL_DEVICE void copy_row_warp(const void* __restrict__ src, void* __restrict__ dst) {
using namespace device;
using plan_t = RowVecPlan<kElementBytes>;
using vec_t = typename plan_t::vec_t;
constexpr auto kLoopCount = plan_t::kLoopCount;
const auto gmem = tile::Memory<vec_t>::warp();
#pragma unroll kLoopCount
for (int64_t i = 0; i < kLoopCount; ++i) {
gmem.store(dst, gmem.load(src, i), i);
}
// handle the epilogue if any
if constexpr (plan_t::kHasEpilogue) {
if (gmem.in_bound(plan_t::kElementCount, kLoopCount)) {
gmem.store(dst, gmem.load(src, kLoopCount), kLoopCount);
}
}
}
/**
* \brief Copy a K row of kKBytes and a V row of kVBytes with one warp.
* The overlapping prefix goes through the interleaved copy; only the width by
* which the rows differ is left as a serial tail. Equal widths degenerate to a
* single interleaved copy with no tail.
*/
template <int64_t kKBytes, int64_t kVBytes>
SGL_DEVICE void copy_kv_rows_warp(
const void* __restrict__ k_src,
const void* __restrict__ v_src,
void* __restrict__ k_dst,
void* __restrict__ v_dst) {
using namespace device;
constexpr auto kCommon = kKBytes < kVBytes ? kKBytes : kVBytes;
constexpr auto kTail = (kKBytes < kVBytes ? kVBytes : kKBytes) - kCommon;
// The interleaved copy indexes BOTH rows with kCommon's vector width, so that
// width must divide each row's split offset -- the narrower row's alignment
// does not imply the wider one's (e.g. 512 picks 16B, but 516 is not 16B
// aligned). The tail's own width must likewise divide its kCommon start.
// Whatever these gates admit is alignment-safe for the strides too, since a
// stride is a whole multiple of its split size.
constexpr auto kTailOrCommon = kTail == 0 ? kCommon : kTail;
constexpr auto kCommonAlign = RowVecPlan<kCommon>::kAlignment;
constexpr auto kTailAlign = RowVecPlan<kTailOrCommon>::kAlignment;
constexpr bool kCanInterleave =
kKBytes % kCommonAlign == 0 && kVBytes % kCommonAlign == 0 && kCommon % kTailAlign == 0;
if constexpr (kCanInterleave) {
copy_kv_warp<kCommon>(k_src, v_src, k_dst, v_dst);
if constexpr (kTail > 0) {
if constexpr (kKBytes > kVBytes) {
copy_row_warp<kTail>(pointer::offset(k_src, kCommon), pointer::offset(k_dst, kCommon));
} else {
copy_row_warp<kTail>(pointer::offset(v_src, kCommon), pointer::offset(v_dst, kCommon));
}
}
} else {
copy_row_warp<kKBytes>(k_src, k_dst);
copy_row_warp<kVBytes>(v_src, v_dst);
}
}
/**
* \brief Kernel to store key-value pairs into the KV cache.
* Each element is split into multiple parts to allow parallel memory copy.
* \tparam kElementBytes The size of each key/value element in bytes.
* \tparam kKElementBytes The size of each key element in bytes.
* \tparam kVElementBytes The size of each value element in bytes. Differs from
* kKElementBytes for asymmetric KV (head_dim != v_head_dim).
* \tparam kSplit The number of warps that handle each element.
* \tparam kUsePDL Whether to use PDL feature.
* \tparam T The data type of the indices (`int32_t` or `int64_t`).
*/
template <int64_t kElementBytes, int kSplit, bool kUsePDL, typename T>
template <int64_t kKElementBytes, int64_t kVElementBytes, int kSplit, bool kUsePDL, typename T>
__global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) {
using namespace device;
constexpr auto kSplitSize = kElementBytes / kSplit;
constexpr auto kKSplitSize = kKElementBytes / kSplit;
constexpr auto kVSplitSize = kVElementBytes / kSplit;
const uint32_t warp_id = blockIdx.x * kNumWarps + threadIdx.x / kWarpThreads;
const uint32_t item_id = warp_id / kSplit;
const uint32_t split_id = warp_id % kSplit;
const auto& [
k_input, v_input, k_cache, v_cache, indices, // ptr
stride_k, stride_v, stride_cache, stride_indices, batch_size, // size
stride_k, stride_v, stride_k_cache, stride_v_cache, stride_indices, batch_size, // size
size_limit, reserved_skip_index // bounds and reserved sink
] = params;
if (item_id >= batch_size) return;
@@ -109,36 +207,37 @@ __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params)
// A stale/OOB slot id would cause an illegal memory access in the store below;
// fail fast at the culprit instead. always-on (kvcache JIT compiles without NDEBUG).
assert(index >= 0 && index < size_limit);
const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kSplitSize);
const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kSplitSize);
const auto k_dst = pointer::offset(k_cache, index * stride_cache, split_id * kSplitSize);
const auto v_dst = pointer::offset(v_cache, index * stride_cache, split_id * kSplitSize);
const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kKSplitSize);
const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kVSplitSize);
const auto k_dst = pointer::offset(k_cache, index * stride_k_cache, split_id * kKSplitSize);
const auto v_dst = pointer::offset(v_cache, index * stride_v_cache, split_id * kVSplitSize);
if (index != reserved_skip_index) {
copy_kv_warp<kSplitSize>(k_src, v_src, k_dst, v_dst);
copy_kv_rows_warp<kKSplitSize, kVSplitSize>(k_src, v_src, k_dst, v_dst);
}
PDLTriggerSecondary<kUsePDL>();
}
template <int64_t kElementBytes, bool kUsePDL>
template <int64_t kKElementBytes, int64_t kVElementBytes, bool kUsePDL>
struct StoreKVCacheKernel {
static_assert(kElementBytes > 0 && kElementBytes % 4 == 0);
static_assert(kKElementBytes > 0 && kKElementBytes % 4 == 0);
static_assert(kVElementBytes > 0 && kVElementBytes % 4 == 0);
template <int kSplit, typename T>
static constexpr auto store_kernel = store_kvcache<kElementBytes, kSplit, kUsePDL, T>;
static constexpr auto store_kernel = store_kvcache<kKElementBytes, kVElementBytes, kSplit, kUsePDL, T>;
template <typename T>
static auto get_kernel(const int num_split) {
using namespace host;
// only apply split optimization when element size is aligned
if constexpr (kElementBytes % (4 * 128) == 0) {
// only apply split optimization when both element sizes are aligned
if constexpr (kKElementBytes % (4 * 128) == 0 && kVElementBytes % (4 * 128) == 0) {
if (num_split == 4) return store_kernel<4, T>;
}
if constexpr (kElementBytes % (2 * 128) == 0) {
if constexpr (kKElementBytes % (2 * 128) == 0 && kVElementBytes % (2 * 128) == 0) {
if (num_split == 2) return store_kernel<2, T>;
}
if (num_split == 1) return store_kernel<1, T>;
Panic("Unsupported num_split {} for element size {}", num_split, kElementBytes);
Panic("Unsupported num_split {} for element sizes k={} v={}", num_split, kKElementBytes, kVElementBytes);
}
static void
@@ -152,31 +251,37 @@ struct StoreKVCacheKernel {
const int64_t reserved_skip_index) {
using namespace host;
auto B = SymbolicSize{"batch_size"};
auto D = SymbolicSize{"element_size"};
auto DK = SymbolicSize{"k_element_size"};
auto DV = SymbolicSize{"v_element_size"};
auto KS = SymbolicSize{"k_stride"};
auto VS = SymbolicSize{"v_stride"};
auto S = SymbolicSize{"cache_stride"};
auto SK = SymbolicSize{"k_cache_stride"};
auto SV = SymbolicSize{"v_cache_stride"};
auto I = SymbolicSize{"indices_stride"};
auto dtype = SymbolicDType{};
auto device = SymbolicDevice{};
auto indice_dtype = SymbolicDType{};
device.set_options<kDLCUDA, kDLROCM>();
TensorMatcher({B, D}) //
TensorMatcher({B, DK}) //
.with_strides({KS, 1})
.with_dtype(dtype)
.with_device(device)
.verify(k);
TensorMatcher({B, D}) //
TensorMatcher({B, DV}) //
.with_strides({VS, 1})
.with_dtype(dtype)
.with_device(device)
.verify(v);
TensorMatcher({-1, D}) //
.with_strides({S, 1})
TensorMatcher({-1, DK}) //
.with_strides({SK, 1})
.with_dtype(dtype)
.with_device(device)
.verify(k_cache);
TensorMatcher({-1, DV}) //
.with_strides({SV, 1})
.with_dtype(dtype)
.with_device(device)
.verify(k_cache)
.verify(v_cache);
TensorMatcher({B}) //
.with_strides({I})
@@ -186,7 +291,8 @@ struct StoreKVCacheKernel {
const int64_t dtype_size = dtype_bytes(dtype.unwrap());
const uint32_t num_elements = static_cast<uint32_t>(B.unwrap());
RuntimeCheck(kElementBytes == dtype_size * D.unwrap());
RuntimeCheck(kKElementBytes == dtype_size * DK.unwrap());
RuntimeCheck(kVElementBytes == dtype_size * DV.unwrap());
const auto params = StoreKVCacheParams{
.k = k.data_ptr(),
@@ -196,7 +302,8 @@ struct StoreKVCacheKernel {
.indices = indices.data_ptr(),
.stride_k_bytes = KS.unwrap() * dtype_size,
.stride_v_bytes = VS.unwrap() * dtype_size,
.stride_cache_bytes = S.unwrap() * dtype_size,
.stride_k_cache_bytes = SK.unwrap() * dtype_size,
.stride_v_cache_bytes = SV.unwrap() * dtype_size,
.stride_indices = I.unwrap(),
.batch_size = static_cast<uint32_t>(B.unwrap()),
.size_limit = size_limit,
python/sglang/kernels/ops/kvcache/kvcache.py
+27 -16
View File
@@ -18,8 +18,8 @@ if TYPE_CHECKING:
@cache_once
def _jit_kvcache_module(row_bytes: int) -> Module:
args = make_cpp_args(row_bytes, is_arch_support_pdl())
def _jit_kvcache_module(k_row_bytes: int, v_row_bytes: int) -> Module:
args = make_cpp_args(k_row_bytes, v_row_bytes, is_arch_support_pdl())
return load_jit(
"kvcache",
*args,
@@ -29,20 +29,25 @@ def _jit_kvcache_module(row_bytes: int) -> Module:
@cache_once
def can_use_store_cache(size: int) -> bool:
def can_use_store_cache(k_row_bytes: int, v_row_bytes: int = 0) -> bool:
"""Whether the JIT store_cache kernel can serve these row widths.
v_row_bytes=0 means symmetric, i.e. it defaults to k_row_bytes."""
logger = logging.getLogger(__name__)
if size % 4 != 0:
logger.warning(
f"Unsupported row_bytes={size} for JIT KV-Cache kernel:"
" must be multiple of 4"
)
return False
v_row_bytes = v_row_bytes or k_row_bytes
for name, size in (("k_row_bytes", k_row_bytes), ("v_row_bytes", v_row_bytes)):
if size % 4 != 0:
logger.warning(
f"Unsupported {name}={size} for JIT KV-Cache kernel:"
" must be multiple of 4"
)
return False
try:
_jit_kvcache_module(size)
_jit_kvcache_module(k_row_bytes, v_row_bytes)
return True
except Exception as e:
logger.warning(
f"Failed to load JIT KV-Cache kernel " f"with row_bytes={size}: {e}"
f"Failed to load JIT KV-Cache kernel with "
f"k_row_bytes={k_row_bytes} v_row_bytes={v_row_bytes}: {e}"
)
return False
@@ -56,6 +61,7 @@ def store_cache(
indices: torch.Tensor,
*,
row_bytes: int = 0,
v_row_bytes: int = 0,
num_split: int = 0, # can be tuned for performance
size_limit: int = 0,
reserved_skip_index: int = 0,
@@ -64,10 +70,13 @@ def store_cache(
Args:
k (torch.Tensor): Key tensor of shape (batch_size, H * D).
v (torch.Tensor): Value tensor of shape (batch_size, H * D).
v (torch.Tensor): Value tensor of shape (batch_size, H * Dv).
k_cache (torch.Tensor): Key cache tensor of shape (num_pages, H * D).
v_cache (torch.Tensor): Value cache tensor of shape (num_pages, H * D).
v_cache (torch.Tensor): Value cache tensor of shape (num_pages, H * Dv).
indices (torch.Tensor): Indices tensor of shape (batch_size,).
row_bytes (int): Key row width in bytes. Inferred from k when 0.
v_row_bytes (int): Value row width in bytes; differs from row_bytes for
asymmetric KV (head_dim != v_head_dim). Inferred from v when 0.
size_limit (int): Valid slot bound (cache row count = real slots + the
reserved padding slot); an index outside [0, size_limit) fails fast
(device assert) instead of an illegal memory access. Defaults to the
@@ -77,11 +86,13 @@ def store_cache(
pass -1 to disable skipping.
"""
row_bytes = row_bytes or k.shape[-1] * k.element_size()
module = _jit_kvcache_module(row_bytes)
v_row_bytes = v_row_bytes or v.shape[-1] * v.element_size()
module = _jit_kvcache_module(row_bytes, v_row_bytes)
if num_split <= 0:
if row_bytes % 2048 == 0:
# A split must divide BOTH rows, so require the alignment on each.
if row_bytes % 2048 == 0 and v_row_bytes % 2048 == 0:
num_split = 4
elif row_bytes % 1024 == 0:
elif row_bytes % 1024 == 0 and v_row_bytes % 1024 == 0:
num_split = 2
else:
num_split = 1
python/sglang/srt/mem_cache/memory_pool.py
+21 -7
View File
@@ -149,21 +149,26 @@ def _set_kv_buffer_impl(
device_module: Any,
size_limit: int,
alt_stream: Optional[torch.cuda.Stream] = None,
same_kv_dim: bool = True,
v_row_dim: Optional[int] = None, # head_num * v_head_dim; defaults to row_dim
) -> None:
v_row_dim = row_dim if v_row_dim is None else v_row_dim
row_bytes = row_dim * store_dtype.itemsize
if (_is_cuda or _is_hip) and same_kv_dim and can_use_store_cache(row_bytes):
v_row_bytes = v_row_dim * store_dtype.itemsize
if (_is_cuda or _is_hip) and can_use_store_cache(row_bytes, v_row_bytes):
return store_cache(
k.view(-1, row_dim),
v.view(-1, row_dim),
v.view(-1, v_row_dim),
k_cache.view(-1, row_dim),
v_cache.view(-1, row_dim),
v_cache.view(-1, v_row_dim),
indices,
row_bytes=row_bytes,
v_row_bytes=v_row_bytes,
size_limit=size_limit,
)
if _is_cpu and _cpu_has_amx_support:
# store_cache_cpu takes a single row_dim for both K and V, so it only serves
# equal-width rows; asymmetric KV falls through to the naive path below.
if _is_cpu and _cpu_has_amx_support and v_row_dim == row_dim:
return torch.ops.sgl_kernel.store_cache_cpu(
k,
v,
@@ -1819,7 +1824,7 @@ class MHATokenToKVPool(KVCache):
# for store_cache JIT kernel
self.row_dim = self.head_num * self.head_dim
self.same_kv_dim = self.head_dim == self.v_head_dim
self.v_row_dim = self.head_num * self.v_head_dim
def _init_kv_copy_and_warmup(self):
# Zero-layer pool (e.g. all-SWA model's full sub-pool) has no buffers.
@@ -2386,7 +2391,7 @@ class MHATokenToKVPool(KVCache):
# dummy tokens write there); valid index range is [0, size + page_size).
size_limit=self.size + self.page_size,
alt_stream=self.alt_stream,
same_kv_dim=self.same_kv_dim,
v_row_dim=self.v_row_dim,
)
def _quantized_scales(self, global_layer_id: int, k_scale, v_scale):
@@ -2712,6 +2717,15 @@ class MHATokenToKVPool(KVCache):
)
return
# The tiled kernel takes one ROW_BYTES for both tensors, so an asymmetric V
# row would be written at K's width and bleed into the next slot. Only this
# path needs the gate; the non-CUDA branch above handles both widths.
if self.v_row_dim != self.row_dim:
raise NotImplementedError(
"prefix-valid commit requires equal-width K/V rows, got "
f"head_dim={self.head_dim} v_head_dim={self.v_head_dim}."
)
_set_kv_buffer_prefix_valid_impl(
cache_k,
cache_v,