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,
test/registered/kernels/benchmark/kvcache/bench_store_cache.py
+28
View File
@@ -72,5 +72,33 @@ def benchmark(batch_size: int, item_size: int, impl: str):
)
# Asymmetric K/V (head_dim != v_head_dim). The item_size sweep above drives both
# rows from one value, so it never reaches the split-prefix-plus-tail path.
# 192/128 and 384/256 both go live in a single MiMoV2 TP=4 deployment, whose
# layers carry either 1 or 2 kv heads per rank; the reversed and wide pairs cover
# a V-side tail and a num_split > 1 shape.
ASYM_ITEM_SIZES = [(192, 128), (384, 256), (128, 192), (1024, 512)]
@marker.parametrize("k_item,v_item", ASYM_ITEM_SIZES, [(192, 128), (1024, 512)])
@marker.parametrize("batch_size", [2**n for n in range(0, 15)], [16])
@marker.benchmark("impl", ["jit", "torch_compile", "torch_streams"])
def benchmark_asymmetric(batch_size: int, k_item: int, v_item: int, impl: str):
torch.manual_seed(42)
k = create_random(batch_size, k_item)
k_cache = create_empty(CACHE_SIZE, k_item)
v = create_random(batch_size, v_item)
v_cache = create_empty(CACHE_SIZE, v_item)
indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size]
return marker.do_bench(
FN_MAP[impl],
input_args=(k, v, k_cache, v_cache, indices),
graph_clone_args=(0, 1, 4), # not need to clone cache, which is large
memory_args=(k, v, indices), # k_cache / v_cache excluded
memory_output=(k, v), # inplace write, size = k + v
)
if __name__ == "__main__":
benchmark.run()
benchmark_asymmetric.run()
test/registered/kernels/ops/kvcache/test_store_cache.py
+87
View File
@@ -174,11 +174,98 @@ def test_store_cache_num_split(
assert torch.all(v_cache[indices] == v)
# Asymmetric K/V (head_dim != v_head_dim): different row widths AND cache strides.
# MiMoV2 is 192/128. Both orderings, since nothing may assume K is the wider one.
ASYM_DIM_PAIRS = get_ci_test_range(
[(192, 128), (128, 192), (1024, 512), (512, 1024), (96, 64), (2048, 1024)],
[(192, 128), (512, 1024)],
)
# The kernel is a byte copier specialized on (k_row_bytes, v_row_bytes) -- no dtype
# in its template args -- so equal-itemsize dtypes share one instantiation. bf16 and
# fp32 are the two distinct itemsizes; fp16 would just re-run the bf16 one.
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
@pytest.mark.parametrize("k_dim,v_dim", ASYM_DIM_PAIRS)
def test_store_cache_asymmetric(k_dim: int, v_dim: int, dtype: torch.dtype) -> None:
batch_size = 128
k = torch.randn((batch_size, k_dim), dtype=dtype, device=DEVICE)
v = torch.randn((batch_size, v_dim), dtype=dtype, device=DEVICE)
k_cache = torch.randn((SMALL_CACHE, k_dim), dtype=dtype, device=DEVICE)
v_cache = torch.randn((SMALL_CACHE, v_dim), dtype=dtype, device=DEVICE)
k_before, v_before = k_cache.clone(), v_cache.clone()
indices = torch.randperm(SMALL_CACHE - 1, device=DEVICE)[:batch_size] + 1
store_cache(k, v, k_cache, v_cache, indices)
assert torch.all(k_cache[indices] == k)
assert torch.all(v_cache[indices] == v)
# Applying K's stride to V (or vice versa) would corrupt neighbouring slots,
# which the target-slot assertions above cannot see.
untouched = torch.ones(SMALL_CACHE, dtype=torch.bool, device=DEVICE)
untouched[indices] = False
assert torch.all(k_cache[untouched] == k_before[untouched])
assert torch.all(v_cache[untouched] == v_before[untouched])
def _valid_asym_num_splits(k_dim: int, v_dim: int, dtype: torch.dtype) -> list:
"""num_split values valid for BOTH rows; a split must divide each of them."""
k_bytes, v_bytes = k_dim * dtype.itemsize, v_dim * dtype.itemsize
splits = [1]
if k_bytes % (2 * 128) == 0 and v_bytes % (2 * 128) == 0:
splits.append(2)
if k_bytes % (4 * 128) == 0 and v_bytes % (4 * 128) == 0:
splits.append(4)
return splits
def _default_num_split(k_dim: int, v_dim: int, dtype: torch.dtype) -> int:
"""Mirrors the heuristic in store_cache(); the default is already exercised
by test_store_cache_asymmetric, which does not pass num_split."""
k_bytes, v_bytes = k_dim * dtype.itemsize, v_dim * dtype.itemsize
if k_bytes % 2048 == 0 and v_bytes % 2048 == 0:
return 4
if k_bytes % 1024 == 0 and v_bytes % 1024 == 0:
return 2
return 1
# Only splits the default heuristic would NOT pick: the split gate is two-sided
# (K and V must both align), so the off-default branches are what needs pinning.
_ASYM_NUM_SPLIT_CASES = [
(_k, _v, _ns)
for _k, _v in ASYM_DIM_PAIRS
for _ns in _valid_asym_num_splits(_k, _v, DTYPE)
if _ns != _default_num_split(_k, _v, DTYPE)
]
@pytest.mark.parametrize("k_dim,v_dim,num_split", _ASYM_NUM_SPLIT_CASES)
def test_store_cache_asymmetric_num_split(
k_dim: int, v_dim: int, num_split: int
) -> None:
batch_size = 128
k = torch.randn((batch_size, k_dim), dtype=DTYPE, device=DEVICE)
v = torch.randn((batch_size, v_dim), dtype=DTYPE, device=DEVICE)
k_cache = torch.randn((SMALL_CACHE, k_dim), dtype=DTYPE, device=DEVICE)
v_cache = torch.randn((SMALL_CACHE, v_dim), dtype=DTYPE, device=DEVICE)
indices = torch.randperm(SMALL_CACHE - 1, device=DEVICE)[:batch_size] + 1
store_cache(k, v, k_cache, v_cache, indices, num_split=num_split)
assert torch.all(k_cache[indices] == k)
assert torch.all(v_cache[indices] == v)
def test_can_use_store_cache() -> None:
assert can_use_store_cache(128)
assert can_use_store_cache(256)
assert can_use_store_cache(1024)
assert can_use_store_cache(2048)
# asymmetric widths, and the documented default (v falls back to k)
assert can_use_store_cache(384, 256)
assert can_use_store_cache(256, 384)
assert can_use_store_cache(1024, 0) == can_use_store_cache(1024)
if __name__ == "__main__":
test/registered/unit/mem_cache/test_asymmetric_mha_pool.py
+202
View File
@@ -0,0 +1,202 @@
"""Device-side tests for MHATokenToKVPool with asymmetric KV (head_dim != v_head_dim).
Covers the wiring the kernel-level tests cannot see: that the pool derives
``v_row_dim`` from ``v_head_dim`` and threads it into the fused store_cache kernel.
A mis-wired ``v_row_dim`` still writes the right bytes into the right K slots, so
the untouched-slot assertions are what pin the V width and stride down.
Skipped on CPU -- the fused path is CUDA/HIP only.
python -m pytest test/registered/unit/mem_cache/test_asymmetric_mha_pool.py -v
"""
import unittest
from types import SimpleNamespace
import torch
from sglang.kernels.ops.kvcache.kvcache import can_use_store_cache
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
from sglang.test.ci.ci_register import register_cuda_ci
_HAS_CUDA = torch.cuda.is_available()
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
DTYPE = torch.bfloat16
HEAD_NUM = 2
POOL_SIZE = 63 # buffers get POOL_SIZE + page_size rows
NUM_WRITES = 16
# (head_dim, v_head_dim). Both orderings, since nothing may assume K is wider.
# The last pair is wide enough for the split heuristic to pick num_split=2.
ASYM_DIM_PAIRS = [(192, 128), (128, 192), (512, 256)]
def _build_pool(head_dim: int, v_head_dim: int) -> MHATokenToKVPool:
return MHATokenToKVPool(
size=POOL_SIZE,
page_size=1,
dtype=DTYPE,
head_num=HEAD_NUM,
head_dim=head_dim,
v_head_dim=v_head_dim,
layer_num=1,
device="cuda",
enable_memory_saver=False,
enable_alt_stream=False,
)
@unittest.skipUnless(_HAS_CUDA, "fused store_cache path requires CUDA")
class TestAsymmetricMHAPoolRowDims(unittest.TestCase):
def test_v_row_dim_tracks_v_head_dim(self):
for head_dim, v_head_dim in ASYM_DIM_PAIRS:
with self.subTest(head_dim=head_dim, v_head_dim=v_head_dim):
pool = _build_pool(head_dim, v_head_dim)
self.assertEqual(pool.row_dim, HEAD_NUM * head_dim)
self.assertEqual(pool.v_row_dim, HEAD_NUM * v_head_dim)
def test_v_row_dim_defaults_to_row_dim_when_symmetric(self):
pool = _build_pool(128, 128)
self.assertEqual(pool.v_row_dim, pool.row_dim)
def test_swa_dims_override_row_dims(self):
# A hybrid sliding-window model builds a second pool through the swa_*
# parameters, which override head_num/head_dim/v_head_dim wholesale. Both
# of MiMoV2's pools are asymmetric, so v_row_dim has to follow
# swa_v_head_dim rather than the full pool's v_head_dim.
pool = MHATokenToKVPool(
size=POOL_SIZE,
page_size=1,
dtype=DTYPE,
head_num=HEAD_NUM,
head_dim=512,
v_head_dim=256,
swa_head_num=1,
swa_head_dim=192,
swa_v_head_dim=128,
layer_num=1,
device="cuda",
enable_memory_saver=False,
enable_alt_stream=False,
)
self.assertEqual(pool.row_dim, 1 * 192)
self.assertEqual(pool.v_row_dim, 1 * 128)
def test_swa_v_head_dim_falls_back_to_v_head_dim(self):
# swa_v_head_dim omitted: head_dim comes from swa_head_dim but v_head_dim
# does not, so the two are read from different sources. Pinned because a
# pool built this way is asymmetric in a way neither config states.
pool = MHATokenToKVPool(
size=POOL_SIZE,
page_size=1,
dtype=DTYPE,
head_num=HEAD_NUM,
head_dim=512,
v_head_dim=256,
swa_head_num=1,
swa_head_dim=192,
layer_num=1,
device="cuda",
enable_memory_saver=False,
enable_alt_stream=False,
)
self.assertEqual(pool.row_dim, 1 * 192)
self.assertEqual(pool.v_row_dim, 1 * 256)
@unittest.skipUnless(_HAS_CUDA, "fused store_cache path requires CUDA")
class TestAsymmetricMHAPoolSetKVBuffer(unittest.TestCase):
"""set_kv_buffer round-trip through the fused kernel, per dim pair."""
def _run_roundtrip(self, head_dim: int, v_head_dim: int):
pool = _build_pool(head_dim, v_head_dim)
k_buf, v_buf = pool.k_buffer[0], pool.v_buffer[0]
self.assertEqual(tuple(k_buf.shape[1:]), (HEAD_NUM, head_dim))
self.assertEqual(tuple(v_buf.shape[1:]), (HEAD_NUM, v_head_dim))
itemsize = pool.store_dtype.itemsize
self.assertTrue(
can_use_store_cache(pool.row_dim * itemsize, pool.v_row_dim * itemsize),
"fused store_cache unavailable; the naive fallback is also correct, so "
"this test would pass without covering anything",
)
# Seed every slot so an over-wide V write shows up on a slot never targeted.
k_buf.copy_(torch.randn_like(k_buf))
v_buf.copy_(torch.randn_like(v_buf))
k_before, v_before = k_buf.clone(), v_buf.clone()
# Slot 0 is the reserved padding slot store_cache skips; target [1, num_slots).
num_slots = k_buf.shape[0]
loc = torch.randperm(num_slots - 1, device="cuda")[:NUM_WRITES] + 1
cache_k = torch.randn(
(NUM_WRITES, HEAD_NUM, head_dim), dtype=DTYPE, device="cuda"
)
cache_v = torch.randn(
(NUM_WRITES, HEAD_NUM, v_head_dim), dtype=DTYPE, device="cuda"
)
pool.set_kv_buffer(SimpleNamespace(layer_id=0), loc, cache_k, cache_v)
self.assertTrue(torch.equal(k_buf[loc], cache_k), "K target slots")
self.assertTrue(torch.equal(v_buf[loc], cache_v), "V target slots")
untouched = torch.ones(num_slots, dtype=torch.bool, device="cuda")
untouched[loc] = False
self.assertTrue(
torch.equal(k_buf[untouched], k_before[untouched]),
"K bled outside its target slots",
)
self.assertTrue(
torch.equal(v_buf[untouched], v_before[untouched]),
"V bled outside its target slots (wrong row width or stride)",
)
def test_asymmetric_roundtrip(self):
for head_dim, v_head_dim in ASYM_DIM_PAIRS:
with self.subTest(head_dim=head_dim, v_head_dim=v_head_dim):
self._run_roundtrip(head_dim, v_head_dim)
def test_symmetric_roundtrip_unchanged(self):
self._run_roundtrip(128, 128)
@unittest.skipUnless(_HAS_CUDA, "prefix-valid tiled kernel requires CUDA")
class TestAsymmetricPrefixValidGuard(unittest.TestCase):
"""set_kv_buffer_prefix_valid's tiled kernel takes one row width for both
tensors, so it must refuse asymmetric KV rather than truncate V."""
def _call_prefix_valid(self, pool, head_dim, v_head_dim):
rows = 2
loc_2d = torch.tensor([[1, 2]], dtype=torch.int64, device="cuda")
commit_lens = torch.tensor([rows], dtype=torch.int32, device="cuda")
cache_k = torch.randn((rows, HEAD_NUM, head_dim), dtype=DTYPE, device="cuda")
cache_v = torch.randn((rows, HEAD_NUM, v_head_dim), dtype=DTYPE, device="cuda")
pool.set_kv_buffer_prefix_valid(
SimpleNamespace(layer_id=0, k_scale=None, v_scale=None),
loc_2d,
commit_lens,
cache_k,
cache_v,
)
def test_rejects_asymmetric(self):
for head_dim, v_head_dim in ASYM_DIM_PAIRS:
with self.subTest(head_dim=head_dim, v_head_dim=v_head_dim):
pool = _build_pool(head_dim, v_head_dim)
with self.assertRaises(NotImplementedError):
self._call_prefix_valid(pool, head_dim, v_head_dim)
def test_accepts_symmetric(self):
# The guard must not tighten the equal-width path it already served.
pool = _build_pool(128, 128)
self._call_prefix_valid(pool, 128, 128)
expected = torch.arange(1, 3, device="cuda")
self.assertTrue(torch.any(pool.k_buffer[0][expected] != 0))
self.assertTrue(torch.any(pool.v_buffer[0][expected] != 0))
if __name__ == "__main__":
unittest.main()