From 66cc11a461bfbef50d20365f925e81e967341ce3 Mon Sep 17 00:00:00 2001 From: Amrutha M Date: Thu, 23 Jul 2026 20:51:36 -0700 Subject: [PATCH 1/6] Add hisparse swap-in kernels Register JIT swap-in kernels (load_cache_to_device_buffer_mla, load_cache_to_device_buffer_dsv4_mla, transfer_cache_dsv4_mla) with supporting headers, Python wrapper, and tests. Add -std=c++20 to SYCL_KERNEL_OPTIONS so ATen/torch headers included from device-compiled runners compile cleanly. --- cmake/BuildFlags.cmake | 4 + .../jit_kernel/hisparse/c4_layout.hpp | 76 +++ .../hisparse/load_cache_to_device_buffer.hpp | 618 ++++++++++++++++++ .../hisparse/transfer_cache_dsv4_mla.hpp | 145 ++++ python/sgl_kernel/jit/__init__.py | 8 + python/sgl_kernel/jit/hisparse.py | 329 ++++++++++ tests/test_hisparse_jit.py | 498 ++++++++++++++ 7 files changed, 1678 insertions(+) create mode 100644 include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp create mode 100644 include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp create mode 100644 include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp create mode 100644 python/sgl_kernel/jit/hisparse.py create mode 100755 tests/test_hisparse_jit.py diff --git a/cmake/BuildFlags.cmake b/cmake/BuildFlags.cmake index 1ffc9b341..acbf5d45e 100644 --- a/cmake/BuildFlags.cmake +++ b/cmake/BuildFlags.cmake @@ -74,6 +74,10 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -ftemplate-backtrace-limit=0) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -fno-sycl-unnamed-lambda) + # ATen/torch headers included in device-compiled FMHA runners require C++20. + # -sycl-std sets the SYCL spec version, not the C++ language standard, so the + # C++ standard must be set explicitly to match SYCL_HOST_FLAGS (-std=c++20). + set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -std=c++20) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -sycl-std=2020) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -fhonor-nans) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -fhonor-infinities) diff --git a/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp b/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp new file mode 100644 index 000000000..7507b32db --- /dev/null +++ b/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp @@ -0,0 +1,76 @@ +/** + * HiSparse C4 paged-cache layout helpers (SYCL / Intel XPU). + * + * Ports device::hisparse::{get_pointer_paged, transfer_item} from the CUDA + * kernel (sglang jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh). + * + * Paged C4 cache layout (per page of kPageSize tokens): + * VALUE 0, VALUE 1, ..., VALUE 63, (kValueBytes each) + * SCALE 0, SCALE 1, ..., SCALE 63, (kScaleBytes each) + * [padding to align the page to a 576-byte boundary] + * + * FlashMLA requires each page to be aligned to 576 bytes. + */ + +#pragma once + +#include +#include + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +// C4 paged layout constants (must match kvcacheio.cuh exactly). +inline constexpr int64_t kPageSize = 64; +inline constexpr int64_t kPageBits = 6; // log2(kPageSize) +inline constexpr int64_t kValueBytes = 576; +inline constexpr int64_t kScaleBytes = 8; +inline constexpr int64_t kItemBytes = kValueBytes + kScaleBytes; + +// div_ceil(kItemBytes * kPageSize, 576) * 576 -> page byte stride. +inline constexpr int64_t kPageBytes = ((kItemBytes * kPageSize + 576 - 1) / 576) * 576; +inline constexpr int64_t kScaleOffset = kValueBytes * kPageSize; + +// int64-word counts for the strided copy loops. +inline constexpr int kValueWords = static_cast(kValueBytes / 8); // 72 +inline constexpr int kScaleWords = static_cast(kScaleBytes / 8); // 1 + +static_assert(kValueBytes % 8 == 0, "kValueBytes must be a multiple of 8"); +static_assert(kScaleBytes % 8 == 0, "kScaleBytes must be a multiple of 8"); +static_assert((int64_t(1) << kPageBits) == kPageSize, "kPageBits must equal log2(kPageSize)"); + +struct PointerInfo { + int64_t* value_ptr; + int64_t* scale_ptr; +}; + +// Resolve the value/scale int64 pointers for a single token slot in a paged +// C4 cache. `index` is the logical token slot; the layout is page-padded. +inline PointerInfo get_pointer_paged(void* cache, int32_t index) { + const int32_t page_num = index >> kPageBits; + const int32_t page_offset = index & (kPageSize - 1); + char* base = static_cast(cache) + static_cast(page_num) * kPageBytes; + char* value_ptr = base + static_cast(page_offset) * kValueBytes; + char* scale_ptr = base + kScaleOffset + static_cast(page_offset) * kScaleBytes; + return {reinterpret_cast(value_ptr), reinterpret_cast(scale_ptr)}; +} + +// Copy one C4 item (value + scale) between page-padded caches, cooperatively +// across a sub-group. `lane_id`/`sg_size` are the sub-group local id and width; +// the strided loops make this correct for any Intel SIMD width (8/16/32). +inline void transfer_item(int lane_id, int sg_size, void* dst_cache, void* src_cache, int32_t dst_index, int32_t src_index) { + const PointerInfo dst = get_pointer_paged(dst_cache, dst_index); + const PointerInfo src = get_pointer_paged(src_cache, src_index); + + for (int j = lane_id; j < kValueWords; j += sg_size) { + dst.value_ptr[j] = src.value_ptr[j]; + } + for (int j = lane_id; j < kScaleWords; j += sg_size) { + dst.scale_ptr[j] = src.scale_ptr[j]; + } +} + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp b/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp new file mode 100644 index 000000000..611c859e7 --- /dev/null +++ b/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp @@ -0,0 +1,618 @@ +/** + * HiSparse: load_cache_to_device_buffer SYCL kernel (Intel XPU). + * + * Ports load_cache_to_device_buffer_kernel<...> from the CUDA source + * (sglang jit_kernel/csrc/hisparse.cuh). One work-group processes one request: + * it hashes the request's top-k token positions, classifies the current hot + * device buffer slots into hits / evictables, streams the missing tokens in + * from the host cache into evicted slots, and rewrites the per-request LRU + * order (evictables at the front, hits at the back). + * + * CUDA -> SYCL mapping: + * - warp (32 lanes) -> sub-group pinned to kWarpSize (32) + * - __ballot + popc(&before) -> exclusive_scan_over_group (local prefix) + * - popc(mask) -> reduce_over_group (sub-group total) + * - __shfl_up / __shfl -> inclusive_scan_over_group / group_broadcast + * - atomicCAS (shared) -> atomic_ref<..., local_space> + * - extern __shared__ -> local_accessor + * - __syncthreads() -> item.barrier(local_space) + * + * The sub-group is pinned to 32 lanes so the slot<->lane mapping + * (slot_idx = chunk * 32 + lane) and the shared-memory layout match the CUDA + * kernel bit-for-bit, giving identical eviction ordering and outputs. + */ + +#pragma once + +#include +#include + +#include "c4_layout.hpp" + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +// Fixed logical warp width (matches the CUDA kernel this ports). +static constexpr int kWarpSize = 32; + +static constexpr int32_t kTokenHit = static_cast(0xFFFFFFFF); // -1 sentinel "already resident" +static constexpr int32_t kHashEmpty = -1; + +// Knuth multiplicative hash for the open-addressing table of size hash_size. +inline int hash_slot(int32_t key, int hash_size) { + return static_cast((static_cast(key) * 2654435761u) % static_cast(hash_size)); +} + +// Cooperative linear (non-paged) item copy across a sub-group. Used by the +// generic (non-DSv4) miss-copy path where device + host are both linear with +// stride == item_size_bytes. +inline void transfer_item_linear(int lane_id, int sg_size, const void* src, void* dst, int64_t item_size_bytes) { + const int64_t nwords = item_size_bytes / 8; + const int64_t* s = static_cast(src); + int64_t* d = static_cast(dst); + for (int64_t j = lane_id; j < nwords; j += sg_size) { + d[j] = s[j]; + } + const int64_t tail_start = nwords * 8; + const char* sc = static_cast(src) + tail_start; + char* dc = static_cast(dst) + tail_start; + for (int64_t j = lane_id; j < item_size_bytes - tail_start; j += sg_size) { + dc[j] = sc[j]; + } +} + +// Shared-memory size calculation (mirrors the CUDA SmemLayout). +// Layout: int32_t region (4-byte aligned) followed by int16_t region. +template +struct SmemLayout { + static constexpr int HASH_SIZE = NUM_TOP_K * 2; + static constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + kWarpSize - 1) / kWarpSize; + // int32_t region: top_k_tokens + chunk_offset + evict_chunk_offset + hash_keys + {total_hits, newest_hit} + static constexpr int TOTAL_INT32 = NUM_TOP_K + (NUM_BUFFER_CHUNKS + 1) + (NUM_BUFFER_CHUNKS + 1) + HASH_SIZE + 2; + // int16_t region: lru_slots_out + hash_vals + static constexpr int TOTAL_INT16 = HOT_BUFFER_SIZE + HASH_SIZE; + static constexpr size_t BYTES = TOTAL_INT32 * sizeof(int32_t) + TOTAL_INT16 * sizeof(int16_t); +}; + +// Local (shared) memory atomic CAS returning the previous value, matching +// CUDA atomicCAS(addr, compare, val) semantics. +inline int32_t atomic_cas_local(int32_t* addr, int32_t compare, int32_t val) { + ::sycl::atomic_ref< + int32_t, + ::sycl::memory_order::relaxed, + ::sycl::memory_scope::work_group, + ::sycl::access::address_space::local_space> + ref(*addr); + int32_t expected = compare; + ref.compare_exchange_strong(expected, val); + // On success `expected` is unchanged (== compare); on failure it holds the + // current value. Either way this equals CUDA atomicCAS's return value. + return expected; +} + +// Single-sub-group inclusive prefix scan over a shared array window +// [offset, offset+kWarpSize), threading a running accumulator. Mirrors the CUDA +// warp_inclusive_scan (which used __shfl_up_sync / __shfl_sync). +inline int warp_inclusive_scan( + const ::sycl::sub_group& sg, int32_t* s_data, int lane_id, int sg_size, int offset, int count, int accumulator) { + const int idx = lane_id + offset; + int val = (idx < count) ? s_data[idx] : 0; + val = ::sycl::inclusive_scan_over_group(sg, val, ::sycl::plus()); + val += accumulator; + if (idx < count) { + s_data[idx] = val; + } + accumulator = ::sycl::group_broadcast(sg, val, sg_size - 1); + return accumulator; +} + +template < + int BLOCK_SIZE, + int NUM_TOP_K, + int HOT_BUFFER_SIZE, + bool IsMLA, + bool IsDsv4Layout, + typename SeqLensT, + typename ReqPoolIndicesT> +class LoadCacheToDeviceBufferKernel { + public: + static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); + + using Layout = SmemLayout; + static constexpr int NUM_WARPS = BLOCK_SIZE / kWarpSize; + static constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K + kWarpSize - 1) / kWarpSize; + static constexpr int NUM_BUFFER_CHUNKS = Layout::NUM_BUFFER_CHUNKS; + static constexpr int HASH_SIZE = Layout::HASH_SIZE; + + LoadCacheToDeviceBufferKernel( + const int32_t* top_k_tokens, + int32_t* device_buffer_tokens, + const int64_t* host_cache_locs, + const int32_t* device_buffer_locs, + const void* host_cache_k, + const void* host_cache_v, + void* device_buffer_k, + void* device_buffer_v, + int32_t* top_k_device_locs, + const ReqPoolIndicesT* req_pool_indices, + const SeqLensT* seq_lens, + int16_t* lru_slots, + const int32_t* num_real_reqs, + int64_t buffer_stride_0, + int64_t host_stride, + int64_t lru_slot_stride_0, + int64_t top_k_tokens_stride, + int64_t top_k_device_locs_stride, + int64_t page_size, + int64_t item_size_bytes, + ::sycl::local_accessor smem) + : top_k_tokens_(top_k_tokens), + device_buffer_tokens_(device_buffer_tokens), + host_cache_locs_(host_cache_locs), + device_buffer_locs_(device_buffer_locs), + host_cache_k_(host_cache_k), + host_cache_v_(host_cache_v), + device_buffer_k_(device_buffer_k), + device_buffer_v_(device_buffer_v), + top_k_device_locs_(top_k_device_locs), + req_pool_indices_(req_pool_indices), + seq_lens_(seq_lens), + lru_slots_(lru_slots), + num_real_reqs_(num_real_reqs), + buffer_stride_0_(buffer_stride_0), + host_stride_(host_stride), + lru_slot_stride_0_(lru_slot_stride_0), + top_k_tokens_stride_(top_k_tokens_stride), + top_k_device_locs_stride_(top_k_device_locs_stride), + page_size_(page_size), + item_size_bytes_(item_size_bytes), + smem_(smem) {} + + [[sycl::reqd_sub_group_size(kWarpSize)]] void operator()(::sycl::nd_item<1> item) const { + const int bid = static_cast(item.get_group(0)); + // Early exit for padded blocks (CUDA graph pads batch to a captured size). + if (bid >= num_real_reqs_[0]) return; + + const ::sycl::sub_group sg = item.get_sub_group(); + const int tid = static_cast(item.get_local_id(0)); + const int warp_id = static_cast(sg.get_group_linear_id()); + const int lane_id = static_cast(sg.get_local_linear_id()); + const int sg_size = static_cast(sg.get_max_local_range()[0]); + + const int64_t rid = static_cast(req_pool_indices_[bid]); + const int64_t seq_len = static_cast(seq_lens_[bid]); + + // Per-request base offsets. + const int32_t* req_top_k_tokens = top_k_tokens_ + bid * top_k_tokens_stride_; + int32_t* req_top_k_device_locs = top_k_device_locs_ + bid * top_k_device_locs_stride_; + + const int64_t buffer_offset = rid * buffer_stride_0_; + int32_t* req_device_buffer_tokens = device_buffer_tokens_ + buffer_offset; + const int32_t* req_device_buffer_locs = device_buffer_locs_ + buffer_offset; + const int64_t* req_host_cache_locs = host_cache_locs_ + rid * host_stride_; + int16_t* req_lru_slots = lru_slots_ + rid * lru_slot_stride_0_; + + // Fast path: short sequences have all tokens resident in device-buffer order. + if (seq_len <= HOT_BUFFER_SIZE) { + const int count = (seq_len < NUM_TOP_K) ? static_cast(seq_len) : NUM_TOP_K; + for (int i = tid; i < count; i += BLOCK_SIZE) { + int32_t token_pos = req_top_k_tokens[i]; + if (token_pos >= 0) { + req_top_k_device_locs[i] = req_device_buffer_locs[token_pos]; + } + } + return; + } + + // Carve up the shared-memory scratch: int32 region first, then int16. + // SYCL local memory is allocated max-aligned, so the int32 reinterpret is safe. + char* smem_raw = &smem_[0]; + int32_t* smem_i32 = reinterpret_cast(smem_raw); + int32_t* s_top_k_tokens = smem_i32; // NUM_TOP_K + int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K; // NUM_BUFFER_CHUNKS + 1 + int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // NUM_BUFFER_CHUNKS + 1 + int32_t* s_hash_keys = s_evict_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // HASH_SIZE + int32_t* s_total_hits_ptr = s_hash_keys + HASH_SIZE; // 1 + int32_t* s_newest_hit_ptr = s_hash_keys + HASH_SIZE + 1; // 1 + + int16_t* smem_i16 = reinterpret_cast(smem_i32 + Layout::TOTAL_INT32); + int16_t* s_lru_slots_out = smem_i16; // HOT_BUFFER_SIZE + int16_t* s_hash_vals = s_lru_slots_out + HOT_BUFFER_SIZE; // HASH_SIZE + + // Initialize counters, hash table, and prefix-sum offsets. + if (tid == 0) { + *s_total_hits_ptr = 0; + *s_newest_hit_ptr = 0; + } + for (int i = tid; i < HASH_SIZE; i += BLOCK_SIZE) { + s_hash_keys[i] = kHashEmpty; + } + for (int i = tid; i < NUM_BUFFER_CHUNKS + 1; i += BLOCK_SIZE) { + s_chunk_offset[i] = 0; + s_evict_chunk_offset[i] = 0; + } + item.barrier(::sycl::access::fence_space::local_space); + + const int newest_slot = HOT_BUFFER_SIZE; + const int32_t newest_token = static_cast(seq_len - 1); + + // Insert top-k token positions into the shared-memory hash table. + for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) { + int32_t token_idx = req_top_k_tokens[i]; + if (token_idx == newest_token) { + // The latest token lives at newest_slot (first slot of the extra page), + // excluded from LRU tracking. Bind and mark it as a hit. + s_top_k_tokens[i] = kTokenHit; + req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot]; + *s_newest_hit_ptr = 1; + } else { + int slot = hash_slot(token_idx, HASH_SIZE); + while (true) { + int32_t old = atomic_cas_local(&s_hash_keys[slot], kHashEmpty, token_idx); + if (old == kHashEmpty || old == token_idx) { + s_hash_vals[slot] = static_cast(i); + break; + } + slot = (slot + 1) % HASH_SIZE; + } + s_top_k_tokens[i] = token_idx; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Pass over hot-buffer slots: classify hits vs evictables and compact them. + constexpr int ITERATIONS_PER_WARP_BUFFER = (NUM_BUFFER_CHUNKS + NUM_WARPS - 1) / NUM_WARPS; + int total_hit_count = 0; + int total_evict_count = 0; + for (int iter = 0; iter < ITERATIONS_PER_WARP_BUFFER; iter++) { + const int chunk_idx = warp_id + iter * NUM_WARPS; + const bool has_valid_chunk = chunk_idx < NUM_BUFFER_CHUNKS; + + const int slot_idx = chunk_idx * kWarpSize + lane_id; + const bool has_valid_slot = has_valid_chunk && (slot_idx < HOT_BUFFER_SIZE); + const int16_t buf_slot = has_valid_slot ? req_lru_slots[slot_idx] : static_cast(-1); + int32_t my_buffer_token = (buf_slot >= 0) ? req_device_buffer_tokens[buf_slot] : -1; + int my_found_top_k_idx = -1; + if (my_buffer_token >= 0) { + int h = hash_slot(my_buffer_token, HASH_SIZE); + while (true) { + int32_t k = s_hash_keys[h]; + if (k == my_buffer_token) { + my_found_top_k_idx = static_cast(s_hash_vals[h]); + break; + } + if (k == kHashEmpty) break; + h = (h + 1) % HASH_SIZE; + } + } + const bool is_hit = my_found_top_k_idx >= 0; + const bool is_evictable = has_valid_slot && !is_hit; + + // Record hits: bind the top-k index to this resident slot's device loc. + if (is_hit) { + s_top_k_tokens[my_found_top_k_idx] = kTokenHit; + req_top_k_device_locs[my_found_top_k_idx] = req_device_buffer_locs[buf_slot]; + } + + int local_hit_offset = 0; + int local_evict_offset = 0; + if (has_valid_chunk) { + local_hit_offset = ::sycl::exclusive_scan_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); + local_evict_offset = ::sycl::exclusive_scan_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); + const int warp_hits = ::sycl::reduce_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); + const int warp_evicts = ::sycl::reduce_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = warp_hits; + s_evict_chunk_offset[chunk_idx + 1] = warp_evicts; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + if (warp_id == 0) { + total_hit_count = + warp_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_hit_count); + total_evict_count = warp_inclusive_scan( + sg, s_evict_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_evict_count); + if (tid == 0) { + *s_total_hits_ptr = total_hit_count; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Hits grow forward from index 0. + if (is_hit) { + int hit_offset = s_chunk_offset[chunk_idx] + local_hit_offset; + s_lru_slots_out[hit_offset] = buf_slot; + } + // Evictables grow backward from HOT_BUFFER_SIZE - 1. + if (is_evictable) { + int evict_offset = s_evict_chunk_offset[chunk_idx] + local_evict_offset; + s_lru_slots_out[HOT_BUFFER_SIZE - 1 - evict_offset] = buf_slot; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Reset offsets for the miss-counting phase (NUM_TOKEN_CHUNKS + 1 entries). + for (int i = tid; i < NUM_TOKEN_CHUNKS + 1; i += BLOCK_SIZE) { + s_chunk_offset[i] = 0; + } + item.barrier(::sycl::access::fence_space::local_space); + + // Pass over top-k tokens: identify misses and assign them evictable slots. + int total_misses = 0; + constexpr int ITERATIONS_PER_WARP_TOKEN = (NUM_TOKEN_CHUNKS + NUM_WARPS - 1) / NUM_WARPS; + for (int iter = 0; iter < ITERATIONS_PER_WARP_TOKEN; iter++) { + const int chunk_idx = warp_id + iter * NUM_WARPS; + const bool has_valid_chunk = chunk_idx < NUM_TOKEN_CHUNKS; + + const int chunk_token_start = chunk_idx * kWarpSize; + const int my_token_idx = chunk_token_start + lane_id; + const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K); + + int32_t my_token = 0; + bool is_miss = false; + int local_miss_offset = 0; + + if (has_valid_token) { + is_miss = s_top_k_tokens[my_token_idx] != kTokenHit; + if (is_miss) { + my_token = s_top_k_tokens[my_token_idx]; + } + } + + if (has_valid_chunk) { + local_miss_offset = ::sycl::exclusive_scan_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); + const int warp_miss_count = ::sycl::reduce_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = warp_miss_count; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + if (warp_id == 0) { + total_misses = + warp_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses); + } + item.barrier(::sycl::access::fence_space::local_space); + + if (is_miss) { + int miss_offset = s_chunk_offset[chunk_idx] + local_miss_offset; + int16_t evict_slot = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - miss_offset]; + // Reuse s_top_k_tokens as miss scratch: miss_offset < my_token_idx always + // holds (hits are skipped), so compacted writes never overrun pending reads. + s_top_k_tokens[miss_offset] = my_token; + req_top_k_device_locs[my_token_idx] = req_device_buffer_locs[evict_slot]; + req_device_buffer_tokens[evict_slot] = my_token; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + total_misses = NUM_TOP_K - *s_total_hits_ptr - *s_newest_hit_ptr; + // Rewrite LRU order: misses then remaining evictables at the front (LRU), + // hits at the back (MRU). + { + const int total_evictable = HOT_BUFFER_SIZE - *s_total_hits_ptr; + for (int i = tid; i < HOT_BUFFER_SIZE; i += BLOCK_SIZE) { + if (i < total_misses) { + req_lru_slots[total_evictable - total_misses + i] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i]; + } else if (i < total_evictable) { + req_lru_slots[i - total_misses] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i]; + } else { + req_lru_slots[i] = s_lru_slots_out[i - total_evictable]; + } + } + } + + // Each sub-group copies one miss directly from host cache to device buffer. + for (int miss_idx = warp_id; miss_idx < total_misses; miss_idx += NUM_WARPS) { + const int32_t miss_token = s_top_k_tokens[miss_idx]; + const int16_t evict_slot = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - miss_idx]; + + const int64_t src_loc = req_host_cache_locs[miss_token]; + const int64_t dst_loc = static_cast(req_device_buffer_locs[evict_slot]); + + if constexpr (IsDsv4Layout) { + // Page-padded C4 device layout + page-padded host layout, K-only. + transfer_item( + lane_id, + sg_size, + device_buffer_k_, + const_cast(host_cache_k_), + static_cast(dst_loc), + static_cast(src_loc)); + } else { + // Generic path: device + host both linear, stride == item_size_bytes. + const char* src_k = static_cast(host_cache_k_) + src_loc * item_size_bytes_; + char* dst_k = static_cast(device_buffer_k_) + dst_loc * item_size_bytes_; + transfer_item_linear(lane_id, sg_size, src_k, dst_k, item_size_bytes_); + + if constexpr (!IsMLA) { + const char* src_v = static_cast(host_cache_v_) + src_loc * item_size_bytes_; + char* dst_v = static_cast(device_buffer_v_) + dst_loc * item_size_bytes_; + transfer_item_linear(lane_id, sg_size, src_v, dst_v, item_size_bytes_); + } + } + } + } + + private: + const int32_t* top_k_tokens_; + int32_t* device_buffer_tokens_; + const int64_t* host_cache_locs_; + const int32_t* device_buffer_locs_; + const void* host_cache_k_; + const void* host_cache_v_; + void* device_buffer_k_; + void* device_buffer_v_; + int32_t* top_k_device_locs_; + const ReqPoolIndicesT* req_pool_indices_; + const SeqLensT* seq_lens_; + int16_t* lru_slots_; + const int32_t* num_real_reqs_; + int64_t buffer_stride_0_; + int64_t host_stride_; + int64_t lru_slot_stride_0_; + int64_t top_k_tokens_stride_; + int64_t top_k_device_locs_stride_; + int64_t page_size_; + int64_t item_size_bytes_; + ::sycl::local_accessor smem_; +}; + +template < + int BLOCK_SIZE, + int NUM_TOP_K, + int HOT_BUFFER_SIZE, + bool IsMLA, + bool IsDsv4Layout, + typename SeqLensT, + typename ReqPoolIndicesT> +void load_cache_to_device_buffer_launcher( + ::sycl::queue& queue, + const void* top_k_tokens, + void* device_buffer_tokens, + const void* host_cache_locs, + const void* device_buffer_locs, + const void* host_cache_k, + const void* host_cache_v, + void* device_buffer_k, + void* device_buffer_v, + void* top_k_device_locs, + const void* req_pool_indices, + const void* seq_lens, + void* lru_slots, + const void* num_real_reqs, + int64_t batch_size, + int64_t buffer_stride_0, + int64_t host_stride, + int64_t lru_slot_stride_0, + int64_t top_k_tokens_stride, + int64_t top_k_device_locs_stride, + int64_t page_size, + int64_t item_size_bytes) { + if (batch_size == 0) { + return; + } + using Kernel = LoadCacheToDeviceBufferKernel< + BLOCK_SIZE, + NUM_TOP_K, + HOT_BUFFER_SIZE, + IsMLA, + IsDsv4Layout, + SeqLensT, + ReqPoolIndicesT>; + constexpr size_t smem_bytes = SmemLayout::BYTES; + + queue.submit([&](::sycl::handler& cgh) { + ::sycl::local_accessor smem(::sycl::range<1>(smem_bytes), cgh); + cgh.parallel_for( + ::sycl::nd_range<1>( + ::sycl::range<1>(static_cast(batch_size) * BLOCK_SIZE), ::sycl::range<1>(BLOCK_SIZE)), + Kernel( + static_cast(top_k_tokens), + static_cast(device_buffer_tokens), + static_cast(host_cache_locs), + static_cast(device_buffer_locs), + host_cache_k, + host_cache_v, + device_buffer_k, + device_buffer_v, + static_cast(top_k_device_locs), + static_cast(req_pool_indices), + static_cast(seq_lens), + static_cast(lru_slots), + static_cast(num_real_reqs), + buffer_stride_0, + host_stride, + lru_slot_stride_0, + top_k_tokens_stride, + top_k_device_locs_stride, + page_size, + item_size_bytes, + smem)); + }); +} + +// ============================================================================ +// C API for Python (ctypes) binding +// ============================================================================ +// +// The compile-time template config (block size, top-k, hot-buffer size, MLA / +// DSv4 flags) is fixed per module via -D macros, mirroring the CUDA JIT that +// bakes the same values into template arguments. The seq_lens / req_pool_indices +// dtype combination (i32/i64) is selected at call time by picking the matching +// exported symbol. + +#define _DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) \ + extern "C" void load_cache_to_device_buffer_##SEQ_SUFFIX##_##RPI_SUFFIX( \ + void* queue_ptr, \ + const void* top_k_tokens, \ + void* device_buffer_tokens, \ + const void* host_cache_locs, \ + const void* device_buffer_locs, \ + const void* host_cache_k, \ + const void* host_cache_v, \ + void* device_buffer_k, \ + void* device_buffer_v, \ + void* top_k_device_locs, \ + const void* req_pool_indices, \ + const void* seq_lens, \ + void* lru_slots, \ + const void* num_real_reqs, \ + int64_t batch_size, \ + int64_t buffer_stride_0, \ + int64_t host_stride, \ + int64_t lru_slot_stride_0, \ + int64_t top_k_tokens_stride, \ + int64_t top_k_device_locs_stride, \ + int64_t page_size, \ + int64_t item_size_bytes) { \ + auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ + load_cache_to_device_buffer_launcher< \ + SGL_HISPARSE_BLOCK_SIZE, \ + SGL_HISPARSE_NUM_TOP_K, \ + SGL_HISPARSE_HOT_BUFFER_SIZE, \ + (SGL_HISPARSE_IS_MLA != 0), \ + (SGL_HISPARSE_IS_DSV4 != 0), \ + SEQ_T, \ + RPI_T>( \ + queue, \ + top_k_tokens, \ + device_buffer_tokens, \ + host_cache_locs, \ + device_buffer_locs, \ + host_cache_k, \ + host_cache_v, \ + device_buffer_k, \ + device_buffer_v, \ + top_k_device_locs, \ + req_pool_indices, \ + seq_lens, \ + lru_slots, \ + num_real_reqs, \ + batch_size, \ + buffer_stride_0, \ + host_stride, \ + lru_slot_stride_0, \ + top_k_tokens_stride, \ + top_k_device_locs_stride, \ + page_size, \ + item_size_bytes); \ + } +#define DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) _DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) + +#if defined(SGL_HISPARSE_BLOCK_SIZE) && defined(SGL_HISPARSE_NUM_TOP_K) && defined(SGL_HISPARSE_HOT_BUFFER_SIZE) && \ + defined(SGL_HISPARSE_IS_MLA) && defined(SGL_HISPARSE_IS_DSV4) +DEFINE_LOAD_CACHE(i64, int64_t, i64, int64_t) +DEFINE_LOAD_CACHE(i64, int64_t, i32, int32_t) +DEFINE_LOAD_CACHE(i32, int32_t, i64, int64_t) +DEFINE_LOAD_CACHE(i32, int32_t, i32, int32_t) +#endif + +#undef DEFINE_LOAD_CACHE +#undef _DEFINE_LOAD_CACHE + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp new file mode 100644 index 000000000..09c60cd02 --- /dev/null +++ b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp @@ -0,0 +1,145 @@ +/** + * HiSparse: transfer_cache_dsv4_mla SYCL kernel (Intel XPU). + * + * Ports transfer_cache_dsv4_mla_kernel from the CUDA source + * (sglang jit_kernel/csrc/hisparse.cuh). Bulk-copies DSv4-MLA C4 tokens between + * two sets of page-padded C4 buffers, one set per model layer. + * + * Mapping to the CUDA original: + * - CUDA "warp" (32 lanes) -> SYCL sub-group (pinned to kSubGroupSize). + * - One sub-group copies one item, iterating over all layers. + * - Grid-stride loop over items across all sub-groups. + * + * src_caches / dst_caches are device arrays of `num_layers` raw cache base + * pointers (uint64_t values), one per layer. + */ + +#pragma once + +#include +#include + +#include "c4_layout.hpp" + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +// Sub-group width used for the cooperative item copy. Intel GPUs support 16/32; +// 32 mirrors the CUDA warp the kernel was written against. The strided copy in +// transfer_item is correct for any width, so this only affects granularity. +static constexpr int kSubGroupSize = 32; + +template +class TransferCacheDsv4MlaKernel { + public: + static constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + + TransferCacheDsv4MlaKernel( + void** src_caches, + void** dst_caches, + const int64_t* src_indices, + const int64_t* dst_indices, + uint32_t num_items, + uint32_t num_layers, + uint32_t total_sub_groups) + : src_caches_(src_caches), + dst_caches_(dst_caches), + src_indices_(src_indices), + dst_indices_(dst_indices), + num_items_(num_items), + num_layers_(num_layers), + total_sub_groups_(total_sub_groups) {} + + [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { + const ::sycl::sub_group sg = item.get_sub_group(); + const int lane_id = static_cast(sg.get_local_linear_id()); + const int sg_size = static_cast(sg.get_max_local_range()[0]); + + // Global sub-group index: group * subgroups_per_group + local subgroup index. + const uint32_t global_sg = + static_cast(item.get_group(0)) * kNumSubGroups + static_cast(sg.get_group_linear_id()); + + for (uint32_t i = global_sg; i < num_items_; i += total_sub_groups_) { + const int32_t src_index = static_cast(src_indices_[i]); + const int32_t dst_index = static_cast(dst_indices_[i]); + for (uint32_t layer_id = 0; layer_id < num_layers_; ++layer_id) { + transfer_item(lane_id, sg_size, dst_caches_[layer_id], src_caches_[layer_id], dst_index, src_index); + } + } + } + + private: + void** src_caches_; + void** dst_caches_; + const int64_t* src_indices_; + const int64_t* dst_indices_; + uint32_t num_items_; + uint32_t num_layers_; + uint32_t total_sub_groups_; +}; + +template +void transfer_cache_dsv4_mla_launcher( + ::sycl::queue& queue, + void** src_caches, + void** dst_caches, + const int64_t* src_indices, + const int64_t* dst_indices, + uint32_t num_items, + uint32_t num_layers) { + if (num_items == 0) { + return; + } + constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + const uint32_t num_groups = (num_items + kNumSubGroups - 1) / kNumSubGroups; + const uint32_t total_sub_groups = num_groups * kNumSubGroups; + + queue.submit([&](::sycl::handler& cgh) { + cgh.parallel_for( + ::sycl::nd_range<1>( + ::sycl::range<1>(static_cast(num_groups) * BLOCK_SIZE), ::sycl::range<1>(BLOCK_SIZE)), + TransferCacheDsv4MlaKernel( + src_caches, dst_caches, src_indices, dst_indices, num_items, num_layers, total_sub_groups)); + }); +} + +// ============================================================================ +// C API for Python (ctypes) binding +// ============================================================================ + +#define _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) \ + extern "C" void transfer_cache_dsv4_mla_##BLOCK_SIZE( \ + void* queue_ptr, \ + void* src_caches, \ + void* dst_caches, \ + const void* src_indices, \ + const void* dst_indices, \ + uint32_t num_items, \ + uint32_t num_layers) { \ + auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ + transfer_cache_dsv4_mla_launcher( \ + queue, \ + static_cast(src_caches), \ + static_cast(dst_caches), \ + static_cast(src_indices), \ + static_cast(dst_indices), \ + num_items, \ + num_layers); \ + } +#define DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) + +#ifdef SGL_HISPARSE_BLOCK_SIZE +DEFINE_TRANSFER_CACHE_DSV4_MLA(SGL_HISPARSE_BLOCK_SIZE) +#else +DEFINE_TRANSFER_CACHE_DSV4_MLA(256) +DEFINE_TRANSFER_CACHE_DSV4_MLA(512) +DEFINE_TRANSFER_CACHE_DSV4_MLA(1024) +#endif + +#undef DEFINE_TRANSFER_CACHE_DSV4_MLA +#undef _DEFINE_TRANSFER_CACHE_DSV4_MLA + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/python/sgl_kernel/jit/__init__.py b/python/sgl_kernel/jit/__init__.py index 7f4d14817..e2b562331 100644 --- a/python/sgl_kernel/jit/__init__.py +++ b/python/sgl_kernel/jit/__init__.py @@ -26,6 +26,11 @@ def is_xpu() -> bool: is_icpx_available, load_jit_sycl, ) + from .hisparse import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, + ) from .moe_align_block_size import moe_align_block_size from .norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm, rmsnorm from .per_tensor_quant_fp8 import per_tensor_quant_fp8 @@ -54,6 +59,9 @@ def is_xpu() -> bool: "apply_rope_inplace", "apply_rope_inplace_with_kvcache", "timestep_embedding", + "transfer_cache_dsv4_mla", + "load_cache_to_device_buffer_mla", + "load_cache_to_device_buffer_dsv4_mla", ] else: # Non-XPU environment - provide stubs diff --git a/python/sgl_kernel/jit/hisparse.py b/python/sgl_kernel/jit/hisparse.py new file mode 100644 index 000000000..d739c39a9 --- /dev/null +++ b/python/sgl_kernel/jit/hisparse.py @@ -0,0 +1,329 @@ +""" +XPU/SYCL HiSparse KV-offload swap-in kernel wrappers. + +Provides JIT-compiled SYCL ports of the CUDA HiSparse kernels used for +hierarchical sparse attention (DeepSeek DSA / V4). Two kernels are exposed: + +- ``transfer_cache_dsv4_mla``: bulk-copy DSv4 C4 tokens between page-padded C4 + buffers, one set of buffers per model layer (evict / backup path). +- ``load_cache_to_device_buffer_mla`` / ``..._dsv4_mla``: per-request swap-in of + the current top-k tokens into a small hot device buffer, maintaining an LRU + ordering and streaming misses in from the host cache. + +These mirror the API of ``sglang.jit_kernel.hisparse`` (the CUDA path). +""" + +from __future__ import annotations + +import ctypes + +import torch + +from .compiler import load_jit_sycl +from .utils import cache_once + +# Block sizes for which the transfer kernel is pre-instantiated in the header's +# default (non-macro) path. Other sizes are compiled on demand via -D. +_SUPPORTED_TRANSFER_BLOCK_SIZES = (256, 512, 1024) + + +# --------------------------------------------------------------------------- +# transfer_cache_dsv4_mla +# --------------------------------------------------------------------------- + + +@cache_once +def _jit_transfer_cache_dsv4_mla_module(block_size: int): + """Compile/load the DSv4 C4 transfer module for a given block size.""" + if block_size % 32 != 0: + raise ValueError(f"block_size must be a multiple of 32, got {block_size}") + return load_jit_sycl( + "hisparse_transfer_cache_dsv4_mla", + str(block_size), + sycl_files=["hisparse/transfer_cache_dsv4_mla.hpp"], + extra_sycl_cflags=[f"-DSGL_HISPARSE_BLOCK_SIZE={block_size}"], + ) + + +_TRANSFER_ARGTYPES = [ + ctypes.c_void_p, # queue + ctypes.c_void_p, # src_caches (void**) + ctypes.c_void_p, # dst_caches (void**) + ctypes.c_void_p, # src_indices (const int64_t*) + ctypes.c_void_p, # dst_indices (const int64_t*) + ctypes.c_uint32, # num_items + ctypes.c_uint32, # num_layers +] + + +def transfer_cache_dsv4_mla( + src_ptrs: torch.Tensor, + dst_ptrs: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + block_size: int = 1024, +) -> None: + """Transfer DSv4 C4 tokens between page-padded C4 buffers. + + Args: + src_ptrs: 1-D uint64 tensor of per-layer source cache base pointers. + dst_ptrs: 1-D uint64 tensor of per-layer destination cache base pointers. + src_indices: 1-D int64 tensor of source token slot indices. + dst_indices: 1-D int64 tensor of destination token slot indices. + block_size: SYCL work-group size (multiple of 32). + """ + assert src_ptrs.dtype == torch.uint64 and dst_ptrs.dtype == torch.uint64 + assert src_indices.dtype == torch.int64 and dst_indices.dtype == torch.int64 + assert src_ptrs.numel() == dst_ptrs.numel() + assert src_indices.numel() == dst_indices.numel() + + num_items = src_indices.numel() + if num_items == 0: + return + num_layers = src_ptrs.numel() + + module = _jit_transfer_cache_dsv4_mla_module(block_size) + func = module.get_function( + f"transfer_cache_dsv4_mla_{block_size}", _TRANSFER_ARGTYPES + ) + queue = torch.xpu.current_stream().sycl_queue + func( + queue, + src_ptrs.data_ptr(), + dst_ptrs.data_ptr(), + src_indices.data_ptr(), + dst_indices.data_ptr(), + num_items, + num_layers, + ) + + +# --------------------------------------------------------------------------- +# load_cache_to_device_buffer +# --------------------------------------------------------------------------- + +_LOAD_CACHE_ARGTYPES = [ + ctypes.c_void_p, # queue + ctypes.c_void_p, # top_k_tokens (const int32_t*) + ctypes.c_void_p, # device_buffer_tokens (int32_t*) + ctypes.c_void_p, # host_cache_locs (const int64_t*) + ctypes.c_void_p, # device_buffer_locs (const int32_t*) + ctypes.c_void_p, # host_cache_k + ctypes.c_void_p, # host_cache_v + ctypes.c_void_p, # device_buffer_k + ctypes.c_void_p, # device_buffer_v + ctypes.c_void_p, # top_k_device_locs (int32_t*) + ctypes.c_void_p, # req_pool_indices + ctypes.c_void_p, # seq_lens + ctypes.c_void_p, # lru_slots (int16_t*) + ctypes.c_void_p, # num_real_reqs (const int32_t*) + ctypes.c_int64, # batch_size + ctypes.c_int64, # buffer_stride_0 + ctypes.c_int64, # host_stride + ctypes.c_int64, # lru_slot_stride_0 + ctypes.c_int64, # top_k_tokens_stride + ctypes.c_int64, # top_k_device_locs_stride + ctypes.c_int64, # page_size + ctypes.c_int64, # item_size_bytes +] + + +@cache_once +def _jit_load_cache_module( + block_size: int, + num_top_k: int, + hot_buffer_size: int, + is_mla: bool, + is_dsv4_layout: bool, +): + """Compile/load the swap-in module for a fixed template configuration.""" + if block_size % 32 != 0: + raise ValueError(f"block_size must be a multiple of 32, got {block_size}") + if hot_buffer_size < num_top_k: + raise ValueError( + f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + ) + return load_jit_sycl( + "hisparse_load_cache_to_device_buffer", + str(block_size), + str(num_top_k), + str(hot_buffer_size), + "mla" if is_mla else "gqa", + "dsv4" if is_dsv4_layout else "linear", + sycl_files=["hisparse/load_cache_to_device_buffer.hpp"], + extra_sycl_cflags=[ + f"-DSGL_HISPARSE_BLOCK_SIZE={block_size}", + f"-DSGL_HISPARSE_NUM_TOP_K={num_top_k}", + f"-DSGL_HISPARSE_HOT_BUFFER_SIZE={hot_buffer_size}", + f"-DSGL_HISPARSE_IS_MLA={1 if is_mla else 0}", + f"-DSGL_HISPARSE_IS_DSV4={1 if is_dsv4_layout else 0}", + ], + ) + + +def _dtype_suffix(t: torch.Tensor) -> str: + return "i64" if t.dtype == torch.int64 else "i32" + + +def _load_cache_to_device_buffer_mla( + *, + is_dsv4_layout: bool, + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int, + block_size: int, + num_real_reqs: torch.Tensor | None, +) -> None: + assert ( + hot_buffer_size >= num_top_k + ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + + module = _jit_load_cache_module( + block_size, + num_top_k, + hot_buffer_size, + True, # is_mla + is_dsv4_layout, + ) + + if num_real_reqs is None: + num_real_reqs = torch.tensor( + [top_k_tokens.size(0)], dtype=torch.int32, device=top_k_tokens.device + ) + + batch_size = top_k_tokens.size(0) + host_stride = host_cache_locs.size(1) + buffer_stride_0 = device_buffer_tokens.stride(0) + lru_slot_stride_0 = lru_slots.stride(0) + top_k_tokens_stride = top_k_tokens.stride(0) + top_k_device_locs_stride = top_k_device_locs.stride(0) + + func_name = ( + f"load_cache_to_device_buffer_" + f"{_dtype_suffix(seq_lens)}_{_dtype_suffix(req_pool_indices)}" + ) + func = module.get_function(func_name, _LOAD_CACHE_ARGTYPES) + queue = torch.xpu.current_stream().sycl_queue + + func( + queue, + top_k_tokens.data_ptr(), + device_buffer_tokens.data_ptr(), + host_cache_locs.data_ptr(), + device_buffer_locs.data_ptr(), + host_cache.data_ptr(), + 0, # host_cache_v (MLA: unused) + device_buffer.data_ptr(), + 0, # device_buffer_v (MLA: unused) + top_k_device_locs.data_ptr(), + req_pool_indices.data_ptr(), + seq_lens.data_ptr(), + lru_slots.data_ptr(), + num_real_reqs.data_ptr(), + batch_size, + buffer_stride_0, + host_stride, + lru_slot_stride_0, + top_k_tokens_stride, + top_k_device_locs_stride, + page_size, + item_size_bytes, + ) + + +def load_cache_to_device_buffer_mla( + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int = 1, + block_size: int = 256, + num_real_reqs: torch.Tensor | None = None, +) -> None: + """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes).""" + _load_cache_to_device_buffer_mla( + is_dsv4_layout=False, + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=top_k_device_locs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + lru_slots=lru_slots, + item_size_bytes=item_size_bytes, + num_top_k=num_top_k, + hot_buffer_size=hot_buffer_size, + page_size=page_size, + block_size=block_size, + num_real_reqs=num_real_reqs, + ) + + +def load_cache_to_device_buffer_dsv4_mla( + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int = 1, + block_size: int = 256, + num_real_reqs: torch.Tensor | None = None, +) -> None: + """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout.""" + _load_cache_to_device_buffer_mla( + is_dsv4_layout=True, + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=top_k_device_locs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + lru_slots=lru_slots, + item_size_bytes=item_size_bytes, + num_top_k=num_top_k, + hot_buffer_size=hot_buffer_size, + page_size=page_size, + block_size=block_size, + num_real_reqs=num_real_reqs, + ) + + +__all__ = [ + "transfer_cache_dsv4_mla", + "load_cache_to_device_buffer_mla", + "load_cache_to_device_buffer_dsv4_mla", +] diff --git a/tests/test_hisparse_jit.py b/tests/test_hisparse_jit.py new file mode 100755 index 000000000..96b6b86a7 --- /dev/null +++ b/tests/test_hisparse_jit.py @@ -0,0 +1,498 @@ +""" +Accuracy tests for the XPU/SYCL HiSparse swap-in kernels. + +Ported from the CUDA oracle (sglang test/registered/jit/test_hisparse.py). The +SYCL kernels pin the logical warp to a 32-lane sub-group, so the slot<->lane +mapping and eviction ordering match the CUDA kernel bit-for-bit; the expected +values below are therefore identical to the CUDA reference. + +Guarded failure modes (derived-property + bug-regression): + - LRU hit/evict compaction and MRU/LRU write-back ordering. + - Miss classification, evict-slot reuse, and host->device miss copy. + - Fast-path (seq_len <= hot_buffer) short-circuit leaves state untouched. + - CUDA-graph padding (num_real_reqs) leaves padded request rows untouched. + - DSv4 page-padded C4 addressing on both the transfer and swap-in paths. +""" + +import pytest +import torch + +HAS_XPU = hasattr(torch, "xpu") and torch.xpu.is_available() + +try: + from sgl_kernel.jit import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, + ) + + HAS_SGL_JIT = True +except ImportError: + HAS_SGL_JIT = False + +pytestmark = [ + pytest.mark.skipif(not HAS_XPU, reason="Requires XPU device"), + pytest.mark.skipif(not HAS_SGL_JIT, reason="Requires sgl_kernel JIT HiSparse"), +] + +DEVICE = "xpu" +DTYPE = torch.float32 +KV_DIM = 8 +HOT_BUFFER_SIZE = 4 +PADDED_BUFFER_SIZE = HOT_BUFFER_SIZE + 1 +HOST_CACHE_SIZE = 16 +DEVICE_CACHE_SIZE = 16 +ITEM_SIZE_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size() +DSV4_PAGE_SIZE = 64 +DSV4_VALUE_BYTES = 576 +DSV4_SCALE_BYTES = 8 +DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES +DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576 +DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE + + +def _pinned(shape, dtype): + """Host tensor pinned for the current XPU device.""" + return torch.empty(shape, dtype=dtype, device="cpu").pin_memory() + + +def _host_cache() -> torch.Tensor: + host_cache = _pinned((HOST_CACHE_SIZE, 1, KV_DIM), DTYPE) + host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache)) + return host_cache + + +def _dsv4_token_pattern(seed: int) -> tuple[torch.Tensor, torch.Tensor]: + value = ( + (torch.arange(DSV4_VALUE_BYTES, dtype=torch.int16) + seed) + .remainder(256) + .to(torch.uint8) + ) + scale = ( + (torch.arange(DSV4_SCALE_BYTES, dtype=torch.int16) + seed + 17) + .remainder(256) + .to(torch.uint8) + ) + return value, scale + + +def _write_dsv4_token(cache: torch.Tensor, loc: int, seed: int) -> None: + page = loc // DSV4_PAGE_SIZE + offset = loc % DSV4_PAGE_SIZE + value, scale = _dsv4_token_pattern(seed) + cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES].copy_( + value.to(cache.device) + ) + scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES + cache[page, scale_start : scale_start + DSV4_SCALE_BYTES].copy_( + scale.to(cache.device) + ) + + +def _read_dsv4_token(cache: torch.Tensor, loc: int) -> torch.Tensor: + page = loc // DSV4_PAGE_SIZE + offset = loc % DSV4_PAGE_SIZE + value = cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES] + scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES + scale = cache[page, scale_start : scale_start + DSV4_SCALE_BYTES] + return torch.cat([value, scale]) + + +def _dsv4_ptrs(cache: torch.Tensor) -> torch.Tensor: + return torch.tensor([cache.data_ptr()], dtype=torch.uint64, device=DEVICE) + + +def _run_kernel( + *, + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + lru_slots: torch.Tensor, + seq_len: int | None = None, + seq_lens: torch.Tensor | None = None, + seq_lens_dtype: torch.dtype = torch.int32, + req_pool_indices: torch.Tensor | None = None, + num_real_reqs: int | None = None, +) -> torch.Tensor: + batch_size = top_k_tokens.shape[0] + if req_pool_indices is None: + req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device=DEVICE) + if seq_lens is None: + seq_lens = torch.full( + (batch_size,), seq_len, dtype=seq_lens_dtype, device=DEVICE + ) + if num_real_reqs is None: + num_real_reqs = batch_size + + out = torch.full_like(top_k_tokens, -1) + load_cache_to_device_buffer_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + lru_slots=lru_slots, + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=top_k_tokens.shape[1], + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=1, + block_size=256, + num_real_reqs=torch.tensor([num_real_reqs], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + return out + + +def _make_state( + device_buffer_locs_rows: list[list[int]], + device_buffer_tokens_rows: list[list[int]], + newest_tokens: list[int], +): + host_cache = _host_cache() + device_buffer = torch.full( + (DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE + ) + device_buffer_locs = torch.tensor( + device_buffer_locs_rows, dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens = torch.tensor( + device_buffer_tokens_rows, dtype=torch.int32, device=DEVICE + ) + lru_slots = ( + torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(device_buffer_locs.shape[0], 1) + ) + host_cache_locs = ( + torch.arange(HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE) + .view(1, -1) + .repeat(device_buffer_locs.shape[0], 1) + ) + + # Slots 0..3 participate in LRU; slot 4 is the reserved newest slot. + for rid, newest_token in enumerate(newest_tokens): + for slot, token in enumerate(device_buffer_tokens_rows[rid][:HOT_BUFFER_SIZE]): + if token >= 0: + device_buffer[device_buffer_locs[rid, slot]].copy_( + host_cache[token].to(DEVICE, non_blocking=True) + ) + device_buffer[device_buffer_locs[rid, HOT_BUFFER_SIZE]].copy_( + host_cache[newest_token].to(DEVICE, non_blocking=True) + ) + torch.xpu.synchronize() + + return { + "host_cache": host_cache, + "device_buffer": device_buffer, + "device_buffer_locs": device_buffer_locs, + "device_buffer_tokens": device_buffer_tokens, + "lru_slots": lru_slots, + "host_cache_locs": host_cache_locs, + } + + +def test_transfer_cache_dsv4_mla_copies_paged_token() -> None: + src_cache = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) + dst_cache = _pinned((2, DSV4_PAGE_BYTES), torch.uint8) + dst_cache.zero_() + src_loc = DSV4_PAGE_SIZE + 6 + dst_loc = DSV4_PAGE_SIZE + 1 + _write_dsv4_token(src_cache, src_loc, seed=41) + + transfer_cache_dsv4_mla( + src_ptrs=_dsv4_ptrs(src_cache), + dst_ptrs=_dsv4_ptrs(dst_cache), + src_indices=torch.tensor([src_loc], dtype=torch.int64, device=DEVICE), + dst_indices=torch.tensor([dst_loc], dtype=torch.int64, device=DEVICE), + ) + torch.xpu.synchronize() + + assert torch.equal( + _read_dsv4_token(dst_cache, dst_loc).to(DEVICE), + _read_dsv4_token(src_cache, src_loc), + ) + + +def test_dsv4_swap_in_reads_paged_host_layout() -> None: + host_cache = _pinned((2, DSV4_PAGE_BYTES), torch.uint8) + host_cache.zero_() + device_buffer = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) + host_loc = DSV4_PAGE_SIZE + 1 + swap_loc = DSV4_PAGE_SIZE + 12 + _write_dsv4_token(host_cache, host_loc, seed=41) + + top_k_tokens = torch.tensor([[3]], dtype=torch.int32, device=DEVICE) + device_buffer_tokens = torch.full( + (1, PADDED_BUFFER_SIZE), -1, dtype=torch.int32, device=DEVICE + ) + host_cache_locs = torch.zeros((1, 8), dtype=torch.int64, device=DEVICE) + host_cache_locs[0, 3] = host_loc + device_buffer_locs = torch.tensor( + [[swap_loc, swap_loc + 1, swap_loc + 2, swap_loc + 3, swap_loc + 4]], + dtype=torch.int32, + device=DEVICE, + ) + lru_slots = torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE).view( + 1, -1 + ) + out = torch.full_like(top_k_tokens, -1) + + load_cache_to_device_buffer_dsv4_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE), + lru_slots=lru_slots, + item_size_bytes=DSV4_ITEM_BYTES, + num_top_k=1, + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=1, + block_size=256, + num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + + assert out.item() == swap_loc + assert torch.equal( + _read_dsv4_token(device_buffer, swap_loc), + _read_dsv4_token(host_cache, host_loc).to(DEVICE), + ) + + +def _long_case(): + # One-request baseline used by the stateful cases below: + # req 0 LRU slots : [0, 1, 2, 3] + # req 0 cached tokens : slot0->1, slot1->4, slot2->2, slot3->5 + # req 0 physical locs : slot0->9, slot1->7, slot2->3, slot3->5 + # req 0 newest slot : slot4/newest -> token 7 at physical loc 11 + return _make_state([[9, 7, 3, 5, 11]], [[1, 4, 2, 5, -1]], [7]) + + +@pytest.mark.parametrize("seq_lens_dtype", [torch.int32, torch.int64]) +def test_load_cache_to_device_buffer_fast_path(seq_lens_dtype: torch.dtype) -> None: + host_cache = _host_cache() + device_buffer = torch.arange( + DEVICE_CACHE_SIZE * KV_DIM, dtype=DTYPE, device=DEVICE + ).view(DEVICE_CACHE_SIZE, 1, KV_DIM) + device_buffer_before = device_buffer.clone() + device_buffer_locs = torch.tensor( + [[13, 9, 5, 1, 15]], dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens = torch.tensor( + [[10, 11, 12, 13, -1]], dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens_before = device_buffer_tokens.clone() + lru_slots = torch.tensor([[0, 1, 2, 3]], dtype=torch.int16, device=DEVICE) + lru_slots_before = lru_slots.clone() + + # seq_len <= HOT_BUFFER_SIZE should skip host loads and LRU mutations, + # so top_k_tokens acts like direct indexing into device_buffer_locs. + out = _run_kernel( + top_k_tokens=torch.tensor([[2, 0, 1]], dtype=torch.int32, device=DEVICE), + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=torch.arange( + HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE + ).view(1, -1), + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + lru_slots=lru_slots, + seq_len=3, + seq_lens_dtype=seq_lens_dtype, + ) + + assert torch.equal(out.cpu(), torch.tensor([[5, 13, 9]], dtype=torch.int32)) + assert torch.equal(device_buffer_tokens.cpu(), device_buffer_tokens_before.cpu()) + assert torch.equal(lru_slots.cpu(), lru_slots_before.cpu()) + assert torch.equal(device_buffer.cpu(), device_buffer_before.cpu()) + + +def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None: + state = _long_case() + + # Query [4, 2, 7]: + # 4 hits slot1 -> loc 7 + # 2 hits slot2 -> loc 3 + # 7 is the newest token -> reserved newest loc 11 + # + # Hits move to the MRU tail, so [0, 1, 2, 3] becomes [0, 3, 1, 2]. + out = _run_kernel( + top_k_tokens=torch.tensor([[4, 2, 7]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[7, 3, 11]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[1, 4, 2, 5, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[0, 3, 1, 2]], dtype=torch.int16) + ) + + +def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None: + state = _long_case() + + # Step 1: touch tokens [4, 2], so LRU becomes [0, 3, 1, 2]. + # Step 2: query token 6, which is a miss. + # The kernel should reuse the new LRU head slot0, whose physical loc is 9. + _run_kernel( + top_k_tokens=torch.tensor([[4, 2]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + out = _run_kernel( + top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[6, 4, 2, 5, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[3, 1, 2, 0]], dtype=torch.int16) + ) + assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6]) + + +def test_load_cache_to_device_buffer_multiple_misses_copy_all_slots() -> None: + state = _make_state( + [[9, 7, 3, 5, 11]], + [[0, 1, 2, 3, -1]], + [8], + ) + + out = _run_kernel( + top_k_tokens=torch.tensor([[4, 5, 6, 7]], dtype=torch.int32, device=DEVICE), + seq_len=9, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[9, 7, 3, 5]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[4, 5, 6, 7, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[0, 1, 2, 3]], dtype=torch.int16) + ) + for token, loc in zip([4, 5, 6, 7], [9, 7, 3, 5]): + assert torch.equal( + state["device_buffer"][loc].cpu(), state["host_cache"][token] + ) + + +def test_load_cache_to_device_buffer_batched_with_padding() -> None: + state = _make_state( + [ + [9, 7, 3, 5, 11], + [12, 10, 8, 6, 14], + [15, 4, 2, 1, 13], + ], + [ + [1, 4, 2, 5, -1], + [0, 1, 2, 3, -1], + [9, 8, 7, 6, -1], + ], + [7, 4, 5], + ) + padded_tokens_before = state["device_buffer_tokens"][2].clone() + padded_lru_before = state["lru_slots"][2].clone() + + # req 0: long path; req 1: fast path; req 2: padded block (must be ignored). + out = _run_kernel( + top_k_tokens=torch.tensor( + [[4, 6, 7], [2, 1, 0], [9, 8, 7]], dtype=torch.int32, device=DEVICE + ), + seq_lens=torch.tensor([8, 3, 8], dtype=torch.int32, device=DEVICE), + num_real_reqs=2, + **state, + ) + + assert torch.equal( + out.cpu(), + torch.tensor([[7, 9, 11], [8, 10, 12], [-1, -1, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["device_buffer_tokens"][:2].cpu(), + torch.tensor([[6, 4, 2, 5, -1], [0, 1, 2, 3, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"][:2].cpu(), + torch.tensor([[2, 3, 0, 1], [0, 1, 2, 3]], dtype=torch.int16), + ) + assert torch.equal( + state["device_buffer_tokens"][2].cpu(), padded_tokens_before.cpu() + ) + assert torch.equal(state["lru_slots"][2].cpu(), padded_lru_before.cpu()) + assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6]) + + +def test_load_cache_to_device_buffer_dsv4_mla_miss_copy_layout() -> None: + # Both the host cache and the device buffer use the page-padded C4 layout. + # The miss copy must read the host source with paged addressing + # (get_pointer_paged), not a linear per-item stride. + num_pages = (HOST_CACHE_SIZE + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE + + state = _long_case() + host_cache = _pinned((num_pages, DSV4_PAGE_BYTES), torch.uint8) + host_cache.zero_() + for token in range(HOST_CACHE_SIZE): + _write_dsv4_token(host_cache, token, seed=token + 1) + + device_buffer = torch.full( + (num_pages, DSV4_PAGE_BYTES), + 0xFF, + dtype=torch.uint8, + device=DEVICE, + ) + out = torch.full((1, 1), -1, dtype=torch.int32, device=DEVICE) + + # Token 6 is a miss in _long_case(), so it should be copied into evict slot 0, + # whose physical device loc is 9. + load_cache_to_device_buffer_dsv4_mla( + top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE), + device_buffer_tokens=state["device_buffer_tokens"], + host_cache_locs=state["host_cache_locs"], + device_buffer_locs=state["device_buffer_locs"], + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE), + lru_slots=state["lru_slots"], + item_size_bytes=DSV4_ITEM_BYTES, + num_top_k=1, + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=DSV4_PAGE_SIZE, + block_size=256, + num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + + assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32)) + assert torch.equal( + _read_dsv4_token(device_buffer, 9).cpu(), + _read_dsv4_token(host_cache, 6), + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From 8de37a36db196934fa48b2c9c5a68707e24c1188 Mon Sep 17 00:00:00 2001 From: Amrutha M Date: Tue, 4 Aug 2026 21:29:35 -0700 Subject: [PATCH 2/6] Address review comments; move to kvcache/ namespace --- .../hisparse/load_cache_to_device_buffer.hpp | 23 +++++++---- .../hisparse/transfer_cache_dsv4_mla.hpp | 3 ++ python/sgl_kernel/jit/__init__.py | 3 +- python/sgl_kernel/jit/kvcache/__init__.py | 26 ++++++++++++ .../sgl_kernel/jit/{ => kvcache}/hisparse.py | 41 +++++++++++++++---- 5 files changed, 80 insertions(+), 16 deletions(-) create mode 100644 python/sgl_kernel/jit/kvcache/__init__.py rename python/sgl_kernel/jit/{ => kvcache}/hisparse.py (87%) diff --git a/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp b/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp index 611c859e7..d8ffb19eb 100644 --- a/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp +++ b/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp @@ -73,6 +73,10 @@ struct SmemLayout { // int16_t region: lru_slots_out + hash_vals static constexpr int TOTAL_INT16 = HOT_BUFFER_SIZE + HASH_SIZE; static constexpr size_t BYTES = TOTAL_INT32 * sizeof(int32_t) + TOTAL_INT16 * sizeof(int16_t); + // Round int16 region up to whole int32 slots so a single int32_t local_accessor + // holds both regions and stays 4-byte aligned throughout. + static constexpr int TOTAL_INT32_SLOTS = + TOTAL_INT32 + (TOTAL_INT16 * sizeof(int16_t) + sizeof(int32_t) - 1) / sizeof(int32_t); }; // Local (shared) memory atomic CAS returning the previous value, matching @@ -118,6 +122,9 @@ template < class LoadCacheToDeviceBufferKernel { public: static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); + static_assert( + BLOCK_SIZE % kWarpSize == 0, + "BLOCK_SIZE must be a multiple of the warp size (32)."); using Layout = SmemLayout; static constexpr int NUM_WARPS = BLOCK_SIZE / kWarpSize; @@ -146,7 +153,7 @@ class LoadCacheToDeviceBufferKernel { int64_t top_k_device_locs_stride, int64_t page_size, int64_t item_size_bytes, - ::sycl::local_accessor smem) + ::sycl::local_accessor smem) : top_k_tokens_(top_k_tokens), device_buffer_tokens_(device_buffer_tokens), host_cache_locs_(host_cache_locs), @@ -205,10 +212,10 @@ class LoadCacheToDeviceBufferKernel { return; } - // Carve up the shared-memory scratch: int32 region first, then int16. - // SYCL local memory is allocated max-aligned, so the int32 reinterpret is safe. - char* smem_raw = &smem_[0]; - int32_t* smem_i32 = reinterpret_cast(smem_raw); + // Scratch is allocated as int32_t (see kernel launch), giving 4-byte + // alignment for the int32 region. int16 region follows immediately after + // TOTAL_INT32 int32 slots, so its base is also 4-byte aligned. + int32_t* smem_i32 = &smem_[0]; int32_t* s_top_k_tokens = smem_i32; // NUM_TOP_K int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K; // NUM_BUFFER_CHUNKS + 1 int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // NUM_BUFFER_CHUNKS + 1 @@ -457,7 +464,7 @@ class LoadCacheToDeviceBufferKernel { int64_t top_k_device_locs_stride_; int64_t page_size_; int64_t item_size_bytes_; - ::sycl::local_accessor smem_; + ::sycl::local_accessor smem_; }; template < @@ -502,10 +509,10 @@ void load_cache_to_device_buffer_launcher( IsDsv4Layout, SeqLensT, ReqPoolIndicesT>; - constexpr size_t smem_bytes = SmemLayout::BYTES; + constexpr size_t smem_slots = SmemLayout::TOTAL_INT32_SLOTS; queue.submit([&](::sycl::handler& cgh) { - ::sycl::local_accessor smem(::sycl::range<1>(smem_bytes), cgh); + ::sycl::local_accessor smem(::sycl::range<1>(smem_slots), cgh); cgh.parallel_for( ::sycl::nd_range<1>( ::sycl::range<1>(static_cast(batch_size) * BLOCK_SIZE), ::sycl::range<1>(BLOCK_SIZE)), diff --git a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp index 09c60cd02..708f7bd3e 100644 --- a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp +++ b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp @@ -33,6 +33,9 @@ static constexpr int kSubGroupSize = 32; template class TransferCacheDsv4MlaKernel { public: + static_assert( + BLOCK_SIZE % kSubGroupSize == 0, + "BLOCK_SIZE must be a multiple of the sub-group size (32)."); static constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; TransferCacheDsv4MlaKernel( diff --git a/python/sgl_kernel/jit/__init__.py b/python/sgl_kernel/jit/__init__.py index e2b562331..a3c3116b2 100644 --- a/python/sgl_kernel/jit/__init__.py +++ b/python/sgl_kernel/jit/__init__.py @@ -26,7 +26,8 @@ def is_xpu() -> bool: is_icpx_available, load_jit_sycl, ) - from .hisparse import ( + # hisparse lives under kvcache/ to mirror sglang's CUDA-side layout + from .kvcache.hisparse import ( load_cache_to_device_buffer_dsv4_mla, load_cache_to_device_buffer_mla, transfer_cache_dsv4_mla, diff --git a/python/sgl_kernel/jit/kvcache/__init__.py b/python/sgl_kernel/jit/kvcache/__init__.py new file mode 100644 index 000000000..b2518eb5f --- /dev/null +++ b/python/sgl_kernel/jit/kvcache/__init__.py @@ -0,0 +1,26 @@ +"""KV-cache-family JIT kernels for Intel XPU. + +Groups SYCL JIT wrappers whose CUDA-side counterparts live under +``sglang.kernels.ops.kvcache.*``. Mirrors that layout so consumers can write +the same import shape across backends: + + from sglang.kernels.ops.kvcache.hisparse import load_cache_to_device_buffer_mla # CUDA + from sgl_kernel.jit.kvcache.hisparse import load_cache_to_device_buffer_mla # XPU + +Modules exported here: + +- ``hisparse``: DSA/DSv4 hisparse swap-in + evict/backup kernels + (``load_cache_to_device_buffer_{mla,dsv4_mla}``, ``transfer_cache_dsv4_mla``). +""" + +from .hisparse import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, +) + +__all__ = [ + "load_cache_to_device_buffer_dsv4_mla", + "load_cache_to_device_buffer_mla", + "transfer_cache_dsv4_mla", +] diff --git a/python/sgl_kernel/jit/hisparse.py b/python/sgl_kernel/jit/kvcache/hisparse.py similarity index 87% rename from python/sgl_kernel/jit/hisparse.py rename to python/sgl_kernel/jit/kvcache/hisparse.py index d739c39a9..f8267c50e 100644 --- a/python/sgl_kernel/jit/hisparse.py +++ b/python/sgl_kernel/jit/kvcache/hisparse.py @@ -19,12 +19,8 @@ import torch -from .compiler import load_jit_sycl -from .utils import cache_once - -# Block sizes for which the transfer kernel is pre-instantiated in the header's -# default (non-macro) path. Other sizes are compiled on demand via -D. -_SUPPORTED_TRANSFER_BLOCK_SIZES = (256, 512, 1024) +from sgl_kernel.jit.compiler import load_jit_sycl +from sgl_kernel.jit.utils import cache_once # --------------------------------------------------------------------------- @@ -162,7 +158,13 @@ def _jit_load_cache_module( def _dtype_suffix(t: torch.Tensor) -> str: - return "i64" if t.dtype == torch.int64 else "i32" + if t.dtype == torch.int64: + return "i64" + if t.dtype == torch.int32: + return "i32" + raise TypeError( + f"seq_lens / req_pool_indices must be int32 or int64, got {t.dtype}" + ) def _load_cache_to_device_buffer_mla( @@ -189,6 +191,31 @@ def _load_cache_to_device_buffer_mla( hot_buffer_size >= num_top_k ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + # Kernel gets only stride(0); rows must be contiguous. device_buffer_tokens + # and device_buffer_locs share stride(0) (reused inside the kernel). + _expected_dtypes = { + "top_k_tokens": (top_k_tokens, torch.int32), + "device_buffer_tokens": (device_buffer_tokens, torch.int32), + "host_cache_locs": (host_cache_locs, torch.int64), + "device_buffer_locs": (device_buffer_locs, torch.int32), + "top_k_device_locs": (top_k_device_locs, torch.int32), + "lru_slots": (lru_slots, torch.int16), + } + for name, (t, expected) in _expected_dtypes.items(): + if t.dtype != expected: + raise TypeError(f"{name} must be {expected}, got {t.dtype}") + if t.device.type != "xpu": + raise ValueError(f"{name} must be on XPU, got {t.device}") + if t.dim() >= 2 and t.stride(-1) != 1: + raise ValueError( + f"{name} must be row-contiguous (stride(-1)==1), got stride={t.stride()}" + ) + if device_buffer_tokens.stride(0) != device_buffer_locs.stride(0): + raise ValueError( + "device_buffer_tokens and device_buffer_locs must share stride(0), " + f"got {device_buffer_tokens.stride(0)} vs {device_buffer_locs.stride(0)}" + ) + module = _jit_load_cache_module( block_size, num_top_k, From eec821a055e9b4f5f111a8af27493133acfa7e74 Mon Sep 17 00:00:00 2001 From: Amrutha M Date: Mon, 10 Aug 2026 20:42:04 -0700 Subject: [PATCH 3/6] Added benchmark file --- benchmark/bench_jit_hisparse.py | 607 ++++++++++++++++++++++++++++++++ 1 file changed, 607 insertions(+) create mode 100644 benchmark/bench_jit_hisparse.py diff --git a/benchmark/bench_jit_hisparse.py b/benchmark/bench_jit_hisparse.py new file mode 100644 index 000000000..2724766be --- /dev/null +++ b/benchmark/bench_jit_hisparse.py @@ -0,0 +1,607 @@ +"""Benchmark the XPU/SYCL HiSparse KV-offload swap-in kernels vs torch eager. + +Covers the three JIT entry points added alongside ``tests/test_hisparse_jit.py``: + +- ``load_cache_to_device_buffer_mla`` — linear host/device layout +- ``load_cache_to_device_buffer_dsv4_mla`` — page-padded C4 layout +- ``transfer_cache_dsv4_mla`` — bulk evict / backup copy + +Each is compared against a pure-PyTorch (``torch`` provider) implementation of +the same operation, so the reported speedup is kernel vs eager on the same +device and the same inputs. The eager versions are written the way a user +naturally would: vectorized advanced indexing, one ``index`` / ``index_copy_`` +per layer or per request batch, no fusion across layers. + +Two swap-in regimes are reported separately, because both the kernel and the +eager reference mutate ``device_buffer_tokens`` / ``lru_slots`` in place: + +- **hit** — every top-k token is already resident. Idempotent across reps, so + ``triton.testing.do_bench`` measures it directly. This is the steady state. +- **miss** — every top-k token must stream in from the host cache. Only the + *first* call on a given state actually misses, so reps are timed one at a + time with XPU events and the state is rebuilt outside the timed window. + +The kernels are templated on ``(block_size, num_top_k, hot_buffer_size, +is_mla, is_dsv4_layout)`` and each distinct tuple triggers its own ``icpx`` +compile, so the sweep varies ``batch_size`` (a runtime argument) and keeps the +template configuration list short. Expect a one-off JIT compile pause on the +first run; later runs hit the ``~/.cache/sgl_kernel/jit_sycl`` ``.so`` cache. + +Usage (with oneAPI on PATH so JIT compilation can find ``icpx``):: + + source /opt/intel/oneapi/2025.3/oneapi-vars.sh + ZE_AFFINITY_MASK=0 python benchmark/bench_jit_hisparse.py + +Pin the run to a *single* device. These kernels are per-rank, and exposing +several devices to one process halves the achieved memory bandwidth on the +device actually used (measured on Arc Pro B60: 826 GB/s with one device +visible vs 377 GB/s with four). That is a runtime/driver effect, not a +property of these kernels -- a bare ``torch.Tensor.copy_`` shows the same 2x +drop -- but it makes multi-device numbers understate the kernel by ~2x. +""" + +import itertools + +import pandas as pd +import torch +import triton + +try: + from sgl_kernel.jit import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, + ) + + HAS_SGL_JIT = True +except ImportError: + HAS_SGL_JIT = False + print("Warning: sgl_kernel JIT HiSparse not available") + +DEVICE = "xpu" + +# Linear-layout MLA item: 512 kv-lora + 64 rope, bf16 (matches DSA MLA cache). +KV_DIM = 576 +DTYPE = torch.bfloat16 +LINEAR_ITEM_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size() + +# Page-padded C4 layout constants (must match c4_layout.hpp). +DSV4_PAGE_SIZE = 64 +DSV4_VALUE_BYTES = 576 +DSV4_SCALE_BYTES = 8 +DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES +DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576 +DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE + +BLOCK_SIZE = 256 + +all_results = [] + + +def _pinned(shape, dtype): + return torch.empty(shape, dtype=dtype, device="cpu").pin_memory() + + +def _release(*objs): + """Drop device buffers and return the memory to the driver. + + The largest transfer config allocates ``num_layers`` src+dst pairs + (~2.3 GB at 32768 items x 61 layers). Holding several configs' worth live + at once pushes a 24 GB card into allocator thrash and depresses the + measured bandwidth by ~2x, so each state is freed once it has been timed. + """ + for o in objs: + if isinstance(o, list): + o.clear() + torch.xpu.synchronize() + torch.xpu.empty_cache() + + +def _dsv4_views(cache): + """Expose a page-padded C4 cache as (value, scale) views, no copy. + + A page is [VALUE 0..63][SCALE 0..63][pad to 576B]. Plain slicing + ``view`` + would fail (the slice is not contiguous, row stride is kPageBytes), so use + ``as_strided`` to build genuine views with the layout's own strides. + """ + pages = cache.shape[0] + value = cache.as_strided( + (pages, DSV4_PAGE_SIZE, DSV4_VALUE_BYTES), + (DSV4_PAGE_BYTES, DSV4_VALUE_BYTES, 1), + ) + scale = cache.as_strided( + (pages, DSV4_PAGE_SIZE, DSV4_SCALE_BYTES), + (DSV4_PAGE_BYTES, DSV4_SCALE_BYTES, 1), + storage_offset=DSV4_SCALE_OFFSET, + ) + return value, scale + + +def _page_split(index): + """Logical token slot -> (page number, offset in page).""" + return index // DSV4_PAGE_SIZE, index % DSV4_PAGE_SIZE + + +# --------------------------------------------------------------------------- +# swap-in state +# --------------------------------------------------------------------------- + + +def _make_swapin_state(batch_size, num_top_k, hot_buffer_size, is_dsv4, regime): + """Build inputs for one swap-in configuration. + + ``regime="hit"`` seeds ``device_buffer_tokens`` with exactly the requested + top-k tokens (no host traffic). ``regime="miss"`` seeds it with a disjoint + token range so every top-k lookup misses and must be streamed in. + """ + slots_per_req = hot_buffer_size + 1 # +1 reserved newest slot + num_device_items = batch_size * slots_per_req + # Host cache holds resident + non-resident tokens for every request. + num_host_items = batch_size * (num_top_k + hot_buffer_size) + 1 + + if is_dsv4: + host_pages = (num_host_items + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE + dev_pages = (num_device_items + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE + host_cache = _pinned((host_pages, DSV4_PAGE_BYTES), torch.uint8) + host_cache.fill_(7) + device_buffer = torch.zeros( + (dev_pages, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE + ) + item_size_bytes = DSV4_ITEM_BYTES + else: + host_cache = _pinned((num_host_items, 1, KV_DIM), DTYPE) + host_cache.fill_(1.0) + device_buffer = torch.zeros( + (num_device_items, 1, KV_DIM), dtype=DTYPE, device=DEVICE + ) + item_size_bytes = LINEAR_ITEM_BYTES + + # Each request owns a contiguous, disjoint span of device slots. + device_buffer_locs = ( + torch.arange(num_device_items, dtype=torch.int32, device=DEVICE) + .view(batch_size, slots_per_req) + .contiguous() + ) + host_cache_locs = ( + torch.arange(num_host_items, dtype=torch.int64, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1) + .contiguous() + ) + + # Request r asks for tokens [base_r, base_r + num_top_k). + base = torch.arange(batch_size, dtype=torch.int32, device=DEVICE).view(-1, 1) * ( + num_top_k + hot_buffer_size + ) + top_k_tokens = ( + base + torch.arange(num_top_k, dtype=torch.int32, device=DEVICE).view(1, -1) + ).contiguous() + + device_buffer_tokens = torch.full( + (batch_size, slots_per_req), -1, dtype=torch.int32, device=DEVICE + ) + if regime == "hit": + # Resident set == requested set, so no host traffic is needed. + device_buffer_tokens[:, :num_top_k] = top_k_tokens + else: + # Resident set is disjoint from the requested set -> every slot misses. + device_buffer_tokens[:, :hot_buffer_size] = ( + base + + num_top_k + + torch.arange(hot_buffer_size, dtype=torch.int32, device=DEVICE).view(1, -1) + ) + device_buffer_tokens[:, hot_buffer_size] = top_k_tokens[:, -1] + + lru_slots = ( + torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1) + .contiguous() + ) + req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device=DEVICE) + # seq_len must exceed hot_buffer_size or the kernel takes the fast path. + seq_lens = torch.full( + (batch_size,), num_top_k + hot_buffer_size, dtype=torch.int32, device=DEVICE + ) + top_k_device_locs = torch.full_like(top_k_tokens, -1) + num_real_reqs = torch.tensor([batch_size], dtype=torch.int32, device=DEVICE) + torch.xpu.synchronize() + + return { + "top_k_tokens": top_k_tokens, + "device_buffer_tokens": device_buffer_tokens, + "host_cache_locs": host_cache_locs, + "device_buffer_locs": device_buffer_locs, + "host_cache": host_cache, + "device_buffer": device_buffer, + "top_k_device_locs": top_k_device_locs, + "req_pool_indices": req_pool_indices, + "seq_lens": seq_lens, + "lru_slots": lru_slots, + "item_size_bytes": item_size_bytes, + "num_top_k": num_top_k, + "hot_buffer_size": hot_buffer_size, + "page_size": DSV4_PAGE_SIZE if is_dsv4 else 1, + "block_size": BLOCK_SIZE, + "num_real_reqs": num_real_reqs, + } + + +def _sglang_swapin(state, is_dsv4): + fn = ( + load_cache_to_device_buffer_dsv4_mla + if is_dsv4 + else load_cache_to_device_buffer_mla + ) + fn(**state) + + +# --------------------------------------------------------------------------- +# torch eager reference: swap-in +# --------------------------------------------------------------------------- + + +def _torch_swapin(state, is_dsv4): + """Pure-PyTorch equivalent of ``load_cache_to_device_buffer_*_mla``. + + Same observable effect as the kernel: resolve which top-k tokens are already + resident, assign the misses to the least-recently-used evictable slots, + stream those items host->device, and refresh the LRU order. + + This is a *timing* reference. It reproduces the resident-token set and the + LRU ordering, but the kernel's exact miss-to-slot assignment is an internal + detail, so slot ids may differ; accuracy is covered by tests/. + + The dominant cost here is structural, not a missing optimization: indexing a + pinned CPU tensor needs the gather indices on the host, which forces a + device->host sync per call. The kernel does the whole thing on device. + """ + top_k = state["top_k_tokens"] # [B, K] int32 + dbt = state["device_buffer_tokens"] # [B, S] int32 + dbl = state["device_buffer_locs"] # [B, S] int32 + hcl = state["host_cache_locs"] # [B, H] int64 + host_cache = state["host_cache"] + device_buffer = state["device_buffer"] + lru = state["lru_slots"] # [B, hot] int16 + hot = state["hot_buffer_size"] + + # ---- classify each requested token as hit or miss ------------------- + eq = top_k.unsqueeze(2) == dbt.unsqueeze(1) # [B, K, S] + is_hit = eq.any(dim=2) # [B, K] + hit_slot = eq.to(torch.uint8).argmax(dim=2) # [B, K] + slot_is_hit = eq.any(dim=1) # [B, S] + + # ---- pick evict slots in LRU order ---------------------------------- + lru_long = lru.to(torch.int64) + evictable = ~slot_is_hit.gather(1, lru_long) # [B, hot] in LRU order + # Stable sort brings evictable slots first while preserving LRU order. + order = torch.argsort(~evictable, dim=1, stable=True) + evict_slots = lru_long.gather(1, order) # [B, hot] + + # Rank each miss, and match it to the evict slot of the same rank. + miss = ~is_hit + miss_rank = miss.cumsum(dim=1) - 1 # [B, K] + evict_pick = evict_slots.gather(1, miss_rank.clamp_(min=0)) + assigned = torch.where(miss, evict_pick, hit_slot) # [B, K] + + # ---- stream the misses in ------------------------------------------- + b_idx, k_idx = miss.nonzero(as_tuple=True) + if b_idx.numel(): + token = top_k[b_idx, k_idx].to(torch.int64) + host_loc = hcl[b_idx, token] + dev_loc = dbl[b_idx, assigned[b_idx, k_idx]].to(torch.int64) + # Pinned host cache must be indexed with CPU indices -> D2H sync. + host_loc_cpu = host_loc.cpu() + if is_dsv4: + h_val, h_scale = _dsv4_views(host_cache) + d_val, d_scale = _dsv4_views(device_buffer) + hp, ho = _page_split(host_loc_cpu) + dp, do = _page_split(dev_loc) + d_val[dp, do] = h_val[hp, ho].to(DEVICE, non_blocking=True) + d_scale[dp, do] = h_scale[hp, ho].to(DEVICE, non_blocking=True) + else: + staged = host_cache[host_loc_cpu].to(DEVICE, non_blocking=True) + device_buffer[dev_loc] = staged + + # ---- refresh residency + LRU order ------------------------------ + dbt[b_idx, assigned[b_idx, k_idx]] = top_k[b_idx, k_idx] + + state["top_k_device_locs"].copy_(dbl.gather(1, assigned)) + # Touched slots become most-recently-used: untouched keep LRU order first. + touched = torch.zeros_like(slot_is_hit) + touched.scatter_(1, assigned, True) + key = touched.gather(1, lru_long).to(torch.uint8) + lru.copy_(lru_long.gather(1, torch.argsort(key, dim=1, stable=True)).to(lru.dtype)) + + +def _swapin_call(provider, state, is_dsv4): + if provider == "sglang": + _sglang_swapin(state, is_dsv4) + else: + _torch_swapin(state, is_dsv4) + + +def _time_miss_regime(provider, batch_size, num_top_k, hot_buffer_size, is_dsv4, reps=20): + """Time the cold (all-miss) path, rebuilding state outside the timed window. + + ``do_bench`` cannot be used here: the first call makes every token + resident, so reps 2..n would measure the hit path instead. + """ + # Warm up JIT compilation and the module cache before timing anything. + _swapin_call( + provider, + _make_swapin_state(batch_size, num_top_k, hot_buffer_size, is_dsv4, "miss"), + is_dsv4, + ) + torch.xpu.synchronize() + + samples = [] + for _ in range(reps): + state = _make_swapin_state( + batch_size, num_top_k, hot_buffer_size, is_dsv4, "miss" + ) + start, end = torch.xpu.Event(enable_timing=True), torch.xpu.Event( + enable_timing=True + ) + start.record() + _swapin_call(provider, state, is_dsv4) + end.record() + torch.xpu.synchronize() + samples.append(start.elapsed_time(end)) + state.clear() # free before building the next rep's state + _release() + samples.sort() + return samples[len(samples) // 2] # median ms + + +# (num_top_k, hot_buffer_size) — each pair is one extra JIT compile per layout. +TEMPLATE_CONFIGS = [(64, 64), (256, 256)] +BATCH_SIZES = [1, 8, 32, 128] + +swapin_configs = [ + (b, k, h) for b, (k, h) in itertools.product(BATCH_SIZES, TEMPLATE_CONFIGS) +] + +SWAPIN_PROVIDERS = [ + f"{p}-{layout}-{regime}" + for layout in ("linear", "dsv4") + for regime in ("hit", "miss") + for p in ("sglang", "torch") +] + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "num_top_k", "hot_buffer_size"], + x_vals=swapin_configs, + line_arg="provider", + line_vals=SWAPIN_PROVIDERS, + line_names=[p.replace("-", " ") for p in SWAPIN_PROVIDERS], + styles=[ + ("blue", "-"), + ("blue", "--"), + ("cyan", "-"), + ("cyan", "--"), + ("green", "-"), + ("green", "--"), + ("orange", "-"), + ("orange", "--"), + ], + ylabel="us", + plot_name="hisparse-load-cache-to-device-buffer-performance", + args={}, + ) +) +def benchmark_swapin(batch_size, num_top_k, hot_buffer_size, provider): + impl, layout, regime = provider.split("-") + is_dsv4 = layout == "dsv4" + + if regime == "miss": + ms = _time_miss_regime( + impl, batch_size, num_top_k, hot_buffer_size, is_dsv4 + ) + min_ms = max_ms = ms + else: + state = _make_swapin_state( + batch_size, num_top_k, hot_buffer_size, is_dsv4, "hit" + ) + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: _swapin_call(impl, state, is_dsv4), quantiles=[0.5, 0.2, 0.8] + ) + state.clear() + _release() + + item_bytes = DSV4_ITEM_BYTES if is_dsv4 else LINEAR_ITEM_BYTES + # A miss reads one item from host and writes one to the device buffer. + moved = 0 if regime == "hit" else batch_size * num_top_k * item_bytes * 2 + all_results.append( + { + "kernel": "load_cache_to_device_buffer", + "provider": impl, + "case": f"{layout}-{regime}", + "batch_size": batch_size, + "num_top_k": num_top_k, + "hot_buffer_size": hot_buffer_size, + "time_us": 1000 * ms, + "GB_s": (moved / (ms * 1e-3) / 1e9) if moved and ms > 0 else float("nan"), + } + ) + return 1000 * ms, 1000 * min_ms, 1000 * max_ms + + +# --------------------------------------------------------------------------- +# transfer_cache_dsv4_mla +# --------------------------------------------------------------------------- + + +def _make_transfer_state(num_items, num_layers, block_size): + pages = (num_items + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE + srcs = [ + torch.full((pages, DSV4_PAGE_BYTES), 3, dtype=torch.uint8, device=DEVICE) + for _ in range(num_layers) + ] + dsts = [ + torch.zeros((pages, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) + for _ in range(num_layers) + ] + src_ptrs = torch.tensor( + [t.data_ptr() for t in srcs], dtype=torch.uint64, device=DEVICE + ) + dst_ptrs = torch.tensor( + [t.data_ptr() for t in dsts], dtype=torch.uint64, device=DEVICE + ) + idx = torch.arange(num_items, dtype=torch.int64, device=DEVICE) + torch.xpu.synchronize() + # Keep srcs/dsts alive: src_ptrs only holds raw addresses. + return srcs, dsts, src_ptrs, dst_ptrs, idx, block_size + + +def _torch_transfer(srcs, dsts, src_indices, dst_indices): + """Pure-PyTorch equivalent of ``transfer_cache_dsv4_mla``. + + The kernel walks all layers inside one launch; eager has to issue an + indexed copy per layer, which is the cost this comparison isolates. + """ + sp, so = _page_split(src_indices) + dp, do = _page_split(dst_indices) + for src, dst in zip(srcs, dsts): + s_val, s_scale = _dsv4_views(src) + d_val, d_scale = _dsv4_views(dst) + d_val[dp, do] = s_val[sp, so] + d_scale[dp, do] = s_scale[sp, so] + + +transfer_configs = list(itertools.product([64, 512, 4096, 32768], [1, 8, 61])) + +TRANSFER_PROVIDERS = ["sglang-bs256", "sglang-bs512", "sglang-bs1024", "torch"] + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_items", "num_layers"], + x_vals=transfer_configs, + line_arg="provider", + line_vals=TRANSFER_PROVIDERS, + line_names=["sglang block=256", "sglang block=512", "sglang block=1024", "torch"], + styles=[("blue", "-"), ("green", "-"), ("red", "-"), ("orange", "--")], + ylabel="us", + plot_name="hisparse-transfer-cache-dsv4-mla-performance", + args={}, + ) +) +def benchmark_transfer(num_items, num_layers, provider): + is_torch = provider == "torch" + block_size = 1024 if is_torch else int(provider.removeprefix("sglang-bs")) + srcs, dsts, src_ptrs, dst_ptrs, idx, bs = _make_transfer_state( + num_items, num_layers, block_size + ) + + if is_torch: + fn = lambda: _torch_transfer(srcs, dsts, idx, idx) + else: + fn = lambda: transfer_cache_dsv4_mla( + src_ptrs=src_ptrs, + dst_ptrs=dst_ptrs, + src_indices=idx, + dst_indices=idx, + block_size=bs, + ) + ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8]) + _release(srcs, dsts) + + moved = num_items * num_layers * DSV4_ITEM_BYTES * 2 # read + write + all_results.append( + { + "kernel": "transfer_cache_dsv4_mla", + "provider": "torch" if is_torch else "sglang", + "case": "transfer" if is_torch else f"transfer-bs{bs}", + "num_items": num_items, + "num_layers": num_layers, + "time_us": 1000 * ms, + "GB_s": moved / (ms * 1e-3) / 1e9 if ms > 0 else float("nan"), + } + ) + return 1000 * ms, 1000 * min_ms, 1000 * max_ms + + +# --------------------------------------------------------------------------- +# speedup analysis +# --------------------------------------------------------------------------- + + +def _report_speedup(df, index_cols, case_label, title): + """Print the torch-vs-sglang speedup summary for one kernel.""" + pivot = df.pivot_table(index=index_cols, columns="provider", values="time_us") + if "torch" not in pivot.columns or "sglang" not in pivot.columns: + return + pivot = pivot.dropna(subset=["torch", "sglang"]) + if pivot.empty: + return + + pivot["speedup"] = pivot["torch"] / pivot["sglang"] + print("\n" + "=" * 80) + print(f"Speedup Analysis (torch vs sglang) — {title}") + print("=" * 80) + print(f"\nOverall average speedup: {pivot['speedup'].mean():.2f}x") + print(f"Overall max speedup: {pivot['speedup'].max():.2f}x") + print(f"Overall min speedup: {pivot['speedup'].min():.2f}x") + + print(f"\nSpeedup by {case_label}:") + levels = pivot.index.get_level_values(case_label) + for value in dict.fromkeys(levels): + sp = pivot.loc[levels == value, "speedup"] + if not sp.empty: + print( + f" {str(value):>14s}: avg={sp.mean():.2f}x " + f"max={sp.max():.2f}x min={sp.min():.2f}x" + ) + + +if __name__ == "__main__": + if not HAS_SGL_JIT: + print("ERROR: sgl_kernel JIT HiSparse kernels unavailable.") + raise SystemExit(1) + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + print("ERROR: no XPU device available.") + raise SystemExit(1) + + print("HiSparse swap-in kernels (JIT SYCL) vs torch eager") + print("First run compiles each template configuration with icpx; be patient.") + print("=" * 80) + benchmark_swapin.run(print_data=True) + + print("\n" + "=" * 80) + print("transfer_cache_dsv4_mla (evict / backup path)") + print("=" * 80) + benchmark_transfer.run(print_data=True) + + df = pd.DataFrame(all_results) + df["time_us"] = df["time_us"].round(2) + df["GB_s"] = df["GB_s"].round(2) + + print("\n" + "=" * 80) + print("Raw Results") + print("=" * 80) + print(df.to_markdown(index=False)) + + swapin = df[df["kernel"] == "load_cache_to_device_buffer"] + _report_speedup( + swapin, + ["batch_size", "num_top_k", "hot_buffer_size", "case"], + "case", + "load_cache_to_device_buffer", + ) + + # Compare eager against the default block size only (block is within noise). + transfer = df[ + (df["kernel"] == "transfer_cache_dsv4_mla") + & (df["case"].isin(["transfer", "transfer-bs1024"])) + ] + _report_speedup( + transfer, + ["num_items", "num_layers"], + "num_layers", + "transfer_cache_dsv4_mla (block=1024)", + ) + + print("\nBenchmark finished!") From 7d7a4a41998e751813eaad86c2952c81576d9b1c Mon Sep 17 00:00:00 2001 From: Amrutha M Date: Mon, 10 Aug 2026 21:02:48 -0700 Subject: [PATCH 4/6] [HiSparse]Fix lint issues --- benchmark/bench_jit_hisparse.py | 19 ++- .../jit_kernel/hisparse/c4_layout.hpp | 3 +- .../hisparse/load_cache_to_device_buffer.hpp | 136 ++++++++-------- .../hisparse/transfer_cache_dsv4_mla.hpp | 148 ------------------ python/sgl_kernel/jit/__init__.py | 1 + python/sgl_kernel/jit/kvcache/hisparse.py | 2 - 6 files changed, 83 insertions(+), 226 deletions(-) delete mode 100644 include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp diff --git a/benchmark/bench_jit_hisparse.py b/benchmark/bench_jit_hisparse.py index 2724766be..10478e8de 100644 --- a/benchmark/bench_jit_hisparse.py +++ b/benchmark/bench_jit_hisparse.py @@ -188,7 +188,9 @@ def _make_swapin_state(batch_size, num_top_k, hot_buffer_size, is_dsv4, regime): device_buffer_tokens[:, :hot_buffer_size] = ( base + num_top_k - + torch.arange(hot_buffer_size, dtype=torch.int32, device=DEVICE).view(1, -1) + + torch.arange(hot_buffer_size, dtype=torch.int32, device=DEVICE).view( + 1, -1 + ) ) device_buffer_tokens[:, hot_buffer_size] = top_k_tokens[:, -1] @@ -321,7 +323,9 @@ def _swapin_call(provider, state, is_dsv4): _torch_swapin(state, is_dsv4) -def _time_miss_regime(provider, batch_size, num_top_k, hot_buffer_size, is_dsv4, reps=20): +def _time_miss_regime( + provider, batch_size, num_top_k, hot_buffer_size, is_dsv4, reps=20 +): """Time the cold (all-miss) path, rebuilding state outside the timed window. ``do_bench`` cannot be used here: the first call makes every token @@ -397,9 +401,7 @@ def benchmark_swapin(batch_size, num_top_k, hot_buffer_size, provider): is_dsv4 = layout == "dsv4" if regime == "miss": - ms = _time_miss_regime( - impl, batch_size, num_top_k, hot_buffer_size, is_dsv4 - ) + ms = _time_miss_regime(impl, batch_size, num_top_k, hot_buffer_size, is_dsv4) min_ms = max_ms = ms else: state = _make_swapin_state( @@ -482,7 +484,12 @@ def _torch_transfer(srcs, dsts, src_indices, dst_indices): x_vals=transfer_configs, line_arg="provider", line_vals=TRANSFER_PROVIDERS, - line_names=["sglang block=256", "sglang block=512", "sglang block=1024", "torch"], + line_names=[ + "sglang block=256", + "sglang block=512", + "sglang block=1024", + "torch", + ], styles=[("blue", "-"), ("green", "-"), ("red", "-"), ("orange", "--")], ylabel="us", plot_name="hisparse-transfer-cache-dsv4-mla-performance", diff --git a/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp b/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp index 7507b32db..1791fa944 100644 --- a/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp +++ b/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp @@ -59,7 +59,8 @@ inline PointerInfo get_pointer_paged(void* cache, int32_t index) { // Copy one C4 item (value + scale) between page-padded caches, cooperatively // across a sub-group. `lane_id`/`sg_size` are the sub-group local id and width; // the strided loops make this correct for any Intel SIMD width (8/16/32). -inline void transfer_item(int lane_id, int sg_size, void* dst_cache, void* src_cache, int32_t dst_index, int32_t src_index) { +inline void +transfer_item(int lane_id, int sg_size, void* dst_cache, void* src_cache, int32_t dst_index, int32_t src_index) { const PointerInfo dst = get_pointer_paged(dst_cache, dst_index); const PointerInfo src = get_pointer_paged(src_cache, src_index); diff --git a/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp b/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp index d8ffb19eb..37a923652 100644 --- a/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp +++ b/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp @@ -122,9 +122,7 @@ template < class LoadCacheToDeviceBufferKernel { public: static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); - static_assert( - BLOCK_SIZE % kWarpSize == 0, - "BLOCK_SIZE must be a multiple of the warp size (32)."); + static_assert(BLOCK_SIZE % kWarpSize == 0, "BLOCK_SIZE must be a multiple of the warp size (32)."); using Layout = SmemLayout; static constexpr int NUM_WARPS = BLOCK_SIZE / kWarpSize; @@ -217,15 +215,15 @@ class LoadCacheToDeviceBufferKernel { // TOTAL_INT32 int32 slots, so its base is also 4-byte aligned. int32_t* smem_i32 = &smem_[0]; int32_t* s_top_k_tokens = smem_i32; // NUM_TOP_K - int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K; // NUM_BUFFER_CHUNKS + 1 - int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // NUM_BUFFER_CHUNKS + 1 - int32_t* s_hash_keys = s_evict_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // HASH_SIZE - int32_t* s_total_hits_ptr = s_hash_keys + HASH_SIZE; // 1 - int32_t* s_newest_hit_ptr = s_hash_keys + HASH_SIZE + 1; // 1 + int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K; // NUM_BUFFER_CHUNKS + 1 + int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // NUM_BUFFER_CHUNKS + 1 + int32_t* s_hash_keys = s_evict_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // HASH_SIZE + int32_t* s_total_hits_ptr = s_hash_keys + HASH_SIZE; // 1 + int32_t* s_newest_hit_ptr = s_hash_keys + HASH_SIZE + 1; // 1 int16_t* smem_i16 = reinterpret_cast(smem_i32 + Layout::TOTAL_INT32); - int16_t* s_lru_slots_out = smem_i16; // HOT_BUFFER_SIZE - int16_t* s_hash_vals = s_lru_slots_out + HOT_BUFFER_SIZE; // HASH_SIZE + int16_t* s_lru_slots_out = smem_i16; // HOT_BUFFER_SIZE + int16_t* s_hash_vals = s_lru_slots_out + HOT_BUFFER_SIZE; // HASH_SIZE // Initialize counters, hash table, and prefix-sum offsets. if (tid == 0) { @@ -317,8 +315,8 @@ class LoadCacheToDeviceBufferKernel { item.barrier(::sycl::access::fence_space::local_space); if (warp_id == 0) { - total_hit_count = - warp_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_hit_count); + total_hit_count = warp_inclusive_scan( + sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_hit_count); total_evict_count = warp_inclusive_scan( sg, s_evict_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_evict_count); if (tid == 0) { @@ -378,8 +376,8 @@ class LoadCacheToDeviceBufferKernel { item.barrier(::sycl::access::fence_space::local_space); if (warp_id == 0) { - total_misses = - warp_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses); + total_misses = warp_inclusive_scan( + sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses); } item.barrier(::sycl::access::fence_space::local_space); @@ -551,61 +549,61 @@ void load_cache_to_device_buffer_launcher( // dtype combination (i32/i64) is selected at call time by picking the matching // exported symbol. -#define _DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) \ - extern "C" void load_cache_to_device_buffer_##SEQ_SUFFIX##_##RPI_SUFFIX( \ - void* queue_ptr, \ - const void* top_k_tokens, \ - void* device_buffer_tokens, \ - const void* host_cache_locs, \ - const void* device_buffer_locs, \ - const void* host_cache_k, \ - const void* host_cache_v, \ - void* device_buffer_k, \ - void* device_buffer_v, \ - void* top_k_device_locs, \ - const void* req_pool_indices, \ - const void* seq_lens, \ - void* lru_slots, \ - const void* num_real_reqs, \ - int64_t batch_size, \ - int64_t buffer_stride_0, \ - int64_t host_stride, \ - int64_t lru_slot_stride_0, \ - int64_t top_k_tokens_stride, \ - int64_t top_k_device_locs_stride, \ - int64_t page_size, \ - int64_t item_size_bytes) { \ - auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ - load_cache_to_device_buffer_launcher< \ - SGL_HISPARSE_BLOCK_SIZE, \ - SGL_HISPARSE_NUM_TOP_K, \ - SGL_HISPARSE_HOT_BUFFER_SIZE, \ - (SGL_HISPARSE_IS_MLA != 0), \ - (SGL_HISPARSE_IS_DSV4 != 0), \ - SEQ_T, \ - RPI_T>( \ - queue, \ - top_k_tokens, \ - device_buffer_tokens, \ - host_cache_locs, \ - device_buffer_locs, \ - host_cache_k, \ - host_cache_v, \ - device_buffer_k, \ - device_buffer_v, \ - top_k_device_locs, \ - req_pool_indices, \ - seq_lens, \ - lru_slots, \ - num_real_reqs, \ - batch_size, \ - buffer_stride_0, \ - host_stride, \ - lru_slot_stride_0, \ - top_k_tokens_stride, \ - top_k_device_locs_stride, \ - page_size, \ - item_size_bytes); \ +#define _DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) \ + extern "C" void load_cache_to_device_buffer_##SEQ_SUFFIX##_##RPI_SUFFIX( \ + void* queue_ptr, \ + const void* top_k_tokens, \ + void* device_buffer_tokens, \ + const void* host_cache_locs, \ + const void* device_buffer_locs, \ + const void* host_cache_k, \ + const void* host_cache_v, \ + void* device_buffer_k, \ + void* device_buffer_v, \ + void* top_k_device_locs, \ + const void* req_pool_indices, \ + const void* seq_lens, \ + void* lru_slots, \ + const void* num_real_reqs, \ + int64_t batch_size, \ + int64_t buffer_stride_0, \ + int64_t host_stride, \ + int64_t lru_slot_stride_0, \ + int64_t top_k_tokens_stride, \ + int64_t top_k_device_locs_stride, \ + int64_t page_size, \ + int64_t item_size_bytes) { \ + auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ + load_cache_to_device_buffer_launcher< \ + SGL_HISPARSE_BLOCK_SIZE, \ + SGL_HISPARSE_NUM_TOP_K, \ + SGL_HISPARSE_HOT_BUFFER_SIZE, \ + (SGL_HISPARSE_IS_MLA != 0), \ + (SGL_HISPARSE_IS_DSV4 != 0), \ + SEQ_T, \ + RPI_T>( \ + queue, \ + top_k_tokens, \ + device_buffer_tokens, \ + host_cache_locs, \ + device_buffer_locs, \ + host_cache_k, \ + host_cache_v, \ + device_buffer_k, \ + device_buffer_v, \ + top_k_device_locs, \ + req_pool_indices, \ + seq_lens, \ + lru_slots, \ + num_real_reqs, \ + batch_size, \ + buffer_stride_0, \ + host_stride, \ + lru_slot_stride_0, \ + top_k_tokens_stride, \ + top_k_device_locs_stride, \ + page_size, \ + item_size_bytes); \ } #define DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) _DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) diff --git a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp deleted file mode 100644 index 708f7bd3e..000000000 --- a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp +++ /dev/null @@ -1,148 +0,0 @@ -/** - * HiSparse: transfer_cache_dsv4_mla SYCL kernel (Intel XPU). - * - * Ports transfer_cache_dsv4_mla_kernel from the CUDA source - * (sglang jit_kernel/csrc/hisparse.cuh). Bulk-copies DSv4-MLA C4 tokens between - * two sets of page-padded C4 buffers, one set per model layer. - * - * Mapping to the CUDA original: - * - CUDA "warp" (32 lanes) -> SYCL sub-group (pinned to kSubGroupSize). - * - One sub-group copies one item, iterating over all layers. - * - Grid-stride loop over items across all sub-groups. - * - * src_caches / dst_caches are device arrays of `num_layers` raw cache base - * pointers (uint64_t values), one per layer. - */ - -#pragma once - -#include -#include - -#include "c4_layout.hpp" - -namespace sgl { -namespace sycl_kernel { -namespace hisparse { - -// Sub-group width used for the cooperative item copy. Intel GPUs support 16/32; -// 32 mirrors the CUDA warp the kernel was written against. The strided copy in -// transfer_item is correct for any width, so this only affects granularity. -static constexpr int kSubGroupSize = 32; - -template -class TransferCacheDsv4MlaKernel { - public: - static_assert( - BLOCK_SIZE % kSubGroupSize == 0, - "BLOCK_SIZE must be a multiple of the sub-group size (32)."); - static constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; - - TransferCacheDsv4MlaKernel( - void** src_caches, - void** dst_caches, - const int64_t* src_indices, - const int64_t* dst_indices, - uint32_t num_items, - uint32_t num_layers, - uint32_t total_sub_groups) - : src_caches_(src_caches), - dst_caches_(dst_caches), - src_indices_(src_indices), - dst_indices_(dst_indices), - num_items_(num_items), - num_layers_(num_layers), - total_sub_groups_(total_sub_groups) {} - - [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { - const ::sycl::sub_group sg = item.get_sub_group(); - const int lane_id = static_cast(sg.get_local_linear_id()); - const int sg_size = static_cast(sg.get_max_local_range()[0]); - - // Global sub-group index: group * subgroups_per_group + local subgroup index. - const uint32_t global_sg = - static_cast(item.get_group(0)) * kNumSubGroups + static_cast(sg.get_group_linear_id()); - - for (uint32_t i = global_sg; i < num_items_; i += total_sub_groups_) { - const int32_t src_index = static_cast(src_indices_[i]); - const int32_t dst_index = static_cast(dst_indices_[i]); - for (uint32_t layer_id = 0; layer_id < num_layers_; ++layer_id) { - transfer_item(lane_id, sg_size, dst_caches_[layer_id], src_caches_[layer_id], dst_index, src_index); - } - } - } - - private: - void** src_caches_; - void** dst_caches_; - const int64_t* src_indices_; - const int64_t* dst_indices_; - uint32_t num_items_; - uint32_t num_layers_; - uint32_t total_sub_groups_; -}; - -template -void transfer_cache_dsv4_mla_launcher( - ::sycl::queue& queue, - void** src_caches, - void** dst_caches, - const int64_t* src_indices, - const int64_t* dst_indices, - uint32_t num_items, - uint32_t num_layers) { - if (num_items == 0) { - return; - } - constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; - const uint32_t num_groups = (num_items + kNumSubGroups - 1) / kNumSubGroups; - const uint32_t total_sub_groups = num_groups * kNumSubGroups; - - queue.submit([&](::sycl::handler& cgh) { - cgh.parallel_for( - ::sycl::nd_range<1>( - ::sycl::range<1>(static_cast(num_groups) * BLOCK_SIZE), ::sycl::range<1>(BLOCK_SIZE)), - TransferCacheDsv4MlaKernel( - src_caches, dst_caches, src_indices, dst_indices, num_items, num_layers, total_sub_groups)); - }); -} - -// ============================================================================ -// C API for Python (ctypes) binding -// ============================================================================ - -#define _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) \ - extern "C" void transfer_cache_dsv4_mla_##BLOCK_SIZE( \ - void* queue_ptr, \ - void* src_caches, \ - void* dst_caches, \ - const void* src_indices, \ - const void* dst_indices, \ - uint32_t num_items, \ - uint32_t num_layers) { \ - auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ - transfer_cache_dsv4_mla_launcher( \ - queue, \ - static_cast(src_caches), \ - static_cast(dst_caches), \ - static_cast(src_indices), \ - static_cast(dst_indices), \ - num_items, \ - num_layers); \ - } -#define DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) - -#ifdef SGL_HISPARSE_BLOCK_SIZE -DEFINE_TRANSFER_CACHE_DSV4_MLA(SGL_HISPARSE_BLOCK_SIZE) -#else -DEFINE_TRANSFER_CACHE_DSV4_MLA(256) -DEFINE_TRANSFER_CACHE_DSV4_MLA(512) -DEFINE_TRANSFER_CACHE_DSV4_MLA(1024) -#endif - -#undef DEFINE_TRANSFER_CACHE_DSV4_MLA -#undef _DEFINE_TRANSFER_CACHE_DSV4_MLA - -} // namespace hisparse -} // namespace sycl_kernel -} // namespace sgl diff --git a/python/sgl_kernel/jit/__init__.py b/python/sgl_kernel/jit/__init__.py index a3c3116b2..58cd8f688 100644 --- a/python/sgl_kernel/jit/__init__.py +++ b/python/sgl_kernel/jit/__init__.py @@ -26,6 +26,7 @@ def is_xpu() -> bool: is_icpx_available, load_jit_sycl, ) + # hisparse lives under kvcache/ to mirror sglang's CUDA-side layout from .kvcache.hisparse import ( load_cache_to_device_buffer_dsv4_mla, diff --git a/python/sgl_kernel/jit/kvcache/hisparse.py b/python/sgl_kernel/jit/kvcache/hisparse.py index f8267c50e..839b5c96b 100644 --- a/python/sgl_kernel/jit/kvcache/hisparse.py +++ b/python/sgl_kernel/jit/kvcache/hisparse.py @@ -18,11 +18,9 @@ import ctypes import torch - from sgl_kernel.jit.compiler import load_jit_sycl from sgl_kernel.jit.utils import cache_once - # --------------------------------------------------------------------------- # transfer_cache_dsv4_mla # --------------------------------------------------------------------------- From acb7e7303b634629cbd9dcf141583cf89dfd515e Mon Sep 17 00:00:00 2001 From: Amrutha M Date: Mon, 17 Aug 2026 19:45:34 -0700 Subject: [PATCH 5/6] [HiSparse] Restore transfer_cache_dsv4_mla.hpp dropped by lint commit --- .../hisparse/transfer_cache_dsv4_mla.hpp | 146 ++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp diff --git a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp new file mode 100644 index 000000000..9f8e0ad16 --- /dev/null +++ b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp @@ -0,0 +1,146 @@ +/** + * HiSparse: transfer_cache_dsv4_mla SYCL kernel (Intel XPU). + * + * Ports transfer_cache_dsv4_mla_kernel from the CUDA source + * (sglang jit_kernel/csrc/hisparse.cuh). Bulk-copies DSv4-MLA C4 tokens between + * two sets of page-padded C4 buffers, one set per model layer. + * + * Mapping to the CUDA original: + * - CUDA "warp" (32 lanes) -> SYCL sub-group (pinned to kSubGroupSize). + * - One sub-group copies one item, iterating over all layers. + * - Grid-stride loop over items across all sub-groups. + * + * src_caches / dst_caches are device arrays of `num_layers` raw cache base + * pointers (uint64_t values), one per layer. + */ + +#pragma once + +#include +#include + +#include "c4_layout.hpp" + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +// Sub-group width used for the cooperative item copy. Intel GPUs support 16/32; +// 32 mirrors the CUDA warp the kernel was written against. The strided copy in +// transfer_item is correct for any width, so this only affects granularity. +static constexpr int kSubGroupSize = 32; + +template +class TransferCacheDsv4MlaKernel { + public: + static_assert(BLOCK_SIZE % kSubGroupSize == 0, "BLOCK_SIZE must be a multiple of the sub-group size (32)."); + static constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + + TransferCacheDsv4MlaKernel( + void** src_caches, + void** dst_caches, + const int64_t* src_indices, + const int64_t* dst_indices, + uint32_t num_items, + uint32_t num_layers, + uint32_t total_sub_groups) + : src_caches_(src_caches), + dst_caches_(dst_caches), + src_indices_(src_indices), + dst_indices_(dst_indices), + num_items_(num_items), + num_layers_(num_layers), + total_sub_groups_(total_sub_groups) {} + + [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { + const ::sycl::sub_group sg = item.get_sub_group(); + const int lane_id = static_cast(sg.get_local_linear_id()); + const int sg_size = static_cast(sg.get_max_local_range()[0]); + + // Global sub-group index: group * subgroups_per_group + local subgroup index. + const uint32_t global_sg = + static_cast(item.get_group(0)) * kNumSubGroups + static_cast(sg.get_group_linear_id()); + + for (uint32_t i = global_sg; i < num_items_; i += total_sub_groups_) { + const int32_t src_index = static_cast(src_indices_[i]); + const int32_t dst_index = static_cast(dst_indices_[i]); + for (uint32_t layer_id = 0; layer_id < num_layers_; ++layer_id) { + transfer_item(lane_id, sg_size, dst_caches_[layer_id], src_caches_[layer_id], dst_index, src_index); + } + } + } + + private: + void** src_caches_; + void** dst_caches_; + const int64_t* src_indices_; + const int64_t* dst_indices_; + uint32_t num_items_; + uint32_t num_layers_; + uint32_t total_sub_groups_; +}; + +template +void transfer_cache_dsv4_mla_launcher( + ::sycl::queue& queue, + void** src_caches, + void** dst_caches, + const int64_t* src_indices, + const int64_t* dst_indices, + uint32_t num_items, + uint32_t num_layers) { + if (num_items == 0) { + return; + } + constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + const uint32_t num_groups = (num_items + kNumSubGroups - 1) / kNumSubGroups; + const uint32_t total_sub_groups = num_groups * kNumSubGroups; + + queue.submit([&](::sycl::handler& cgh) { + cgh.parallel_for( + ::sycl::nd_range<1>( + ::sycl::range<1>(static_cast(num_groups) * BLOCK_SIZE), ::sycl::range<1>(BLOCK_SIZE)), + TransferCacheDsv4MlaKernel( + src_caches, dst_caches, src_indices, dst_indices, num_items, num_layers, total_sub_groups)); + }); +} + +// ============================================================================ +// C API for Python (ctypes) binding +// ============================================================================ + +#define _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) \ + extern "C" void transfer_cache_dsv4_mla_##BLOCK_SIZE( \ + void* queue_ptr, \ + void* src_caches, \ + void* dst_caches, \ + const void* src_indices, \ + const void* dst_indices, \ + uint32_t num_items, \ + uint32_t num_layers) { \ + auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ + transfer_cache_dsv4_mla_launcher( \ + queue, \ + static_cast(src_caches), \ + static_cast(dst_caches), \ + static_cast(src_indices), \ + static_cast(dst_indices), \ + num_items, \ + num_layers); \ + } +#define DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) + +#ifdef SGL_HISPARSE_BLOCK_SIZE +DEFINE_TRANSFER_CACHE_DSV4_MLA(SGL_HISPARSE_BLOCK_SIZE) +#else +DEFINE_TRANSFER_CACHE_DSV4_MLA(256) +DEFINE_TRANSFER_CACHE_DSV4_MLA(512) +DEFINE_TRANSFER_CACHE_DSV4_MLA(1024) +#endif + +#undef DEFINE_TRANSFER_CACHE_DSV4_MLA +#undef _DEFINE_TRANSFER_CACHE_DSV4_MLA + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl From 445bc76e3c6c89fd5951c341c2b00fa329d92598 Mon Sep 17 00:00:00 2001 From: Amrutha M Date: Wed, 19 Aug 2026 19:38:55 -0700 Subject: [PATCH 6/6] Convert HiSparse swap-in kernels to AOT --- benchmark/bench_hisparse.py | 809 ++++++++++++++++++ benchmark/bench_jit_hisparse.py | 614 ------------- .../{jit_kernel => }/hisparse/c4_layout.hpp | 10 +- .../hisparse/load_cache_to_device_buffer.hpp | 449 ++++++++++ .../hisparse/transfer_cache_dsv4_mla.hpp | 58 ++ .../hisparse/load_cache_to_device_buffer.hpp | 623 -------------- .../hisparse/transfer_cache_dsv4_mla.hpp | 146 ---- include/sgl_kernel_ops.h | 29 + python/sgl_kernel/__init__.py | 5 + python/sgl_kernel/hisparse.py | 194 +++++ python/sgl_kernel/jit/__init__.py | 10 - python/sgl_kernel/jit/kvcache/__init__.py | 26 - python/sgl_kernel/jit/kvcache/hisparse.py | 354 -------- src/sycl/HiSparse.cpp | 314 +++++++ src/torch_extension_sycl.cc | 17 + tests/run_suite.py | 1 + ...{test_hisparse_jit.py => test_hisparse.py} | 152 +++- 17 files changed, 2009 insertions(+), 1802 deletions(-) create mode 100644 benchmark/bench_hisparse.py delete mode 100644 benchmark/bench_jit_hisparse.py rename include/sgl_kernel/{jit_kernel => }/hisparse/c4_layout.hpp (90%) create mode 100644 include/sgl_kernel/hisparse/load_cache_to_device_buffer.hpp create mode 100644 include/sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp delete mode 100644 include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp delete mode 100644 include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp create mode 100644 python/sgl_kernel/hisparse.py delete mode 100644 python/sgl_kernel/jit/kvcache/__init__.py delete mode 100644 python/sgl_kernel/jit/kvcache/hisparse.py create mode 100644 src/sycl/HiSparse.cpp rename tests/{test_hisparse_jit.py => test_hisparse.py} (74%) diff --git a/benchmark/bench_hisparse.py b/benchmark/bench_hisparse.py new file mode 100644 index 000000000..d0b855e10 --- /dev/null +++ b/benchmark/bench_hisparse.py @@ -0,0 +1,809 @@ +import argparse +import os +from itertools import product + +import pandas as pd +import torch +import triton +from sgl_kernel import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, +) + +DEVICE = "xpu" +SEED = 42 + +# Linear-layout MLA item: 512 kv-lora + 64 rope, bf16 (matches the DSA MLA cache). +DTYPE = torch.bfloat16 +ELEM = DTYPE.itemsize +KV_DIM = 576 +LINEAR_ITEM_BYTES = KV_DIM * ELEM + +# Page-padded C4 layout constants (must match c4_layout.hpp). +DSV4_PAGE_SIZE = 64 +DSV4_VALUE_BYTES = 576 +DSV4_SCALE_BYTES = 8 +DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES +DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576 +DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE + +VRAM_BUDGET_FRACTION = 0.5 + +QUANTILES = [0.5, 0.25, 0.75] # median, fastest, slowest +RESULT_DIR = "bench_bmg_hisparse_res" + +all_results = [] + + +def _div_up(a, b): + return (a + b - 1) // b + + +def _release(*objs): + """Free device buffers; holding several configs live at once halves bandwidth.""" + for o in objs: + if isinstance(o, (list, dict)): + o.clear() + torch.xpu.synchronize() + torch.xpu.empty_cache() + + +def _gbps(nbytes, ms): + if not nbytes or ms <= 0: + return float("nan") + return nbytes * 1e-9 / (ms * 1e-3) + + +def _dsv4_views(cache): + """Expose a page-padded C4 cache as (value, scale) views, no copy. + + Slicing + ``view`` would fail (the slice is not contiguous, row stride is + DSV4_PAGE_BYTES), so build genuine views with the layout's own strides. + """ + pages = cache.shape[0] + value = cache.as_strided( + (pages, DSV4_PAGE_SIZE, DSV4_VALUE_BYTES), + (DSV4_PAGE_BYTES, DSV4_VALUE_BYTES, 1), + ) + scale = cache.as_strided( + (pages, DSV4_PAGE_SIZE, DSV4_SCALE_BYTES), + (DSV4_PAGE_BYTES, DSV4_SCALE_BYTES, 1), + storage_offset=DSV4_SCALE_OFFSET, + ) + return value, scale + + +def _page_split(index): + return index // DSV4_PAGE_SIZE, index % DSV4_PAGE_SIZE + + +def _bench(fn, warmup=10): + for _ in range(warmup): + fn() + torch.xpu.synchronize() + return triton.testing.do_bench(fn, quantiles=QUANTILES) + + +def _bench_cold(make_state, call, reps=20): + warm = make_state() + call(warm) # warm the kernel and the allocator + warm.clear() + torch.xpu.synchronize() + + samples = [] + for _ in range(reps): + state = make_state() + start = torch.xpu.Event(enable_timing=True) + end = torch.xpu.Event(enable_timing=True) + start.record() + call(state) + end.record() + torch.xpu.synchronize() + samples.append(start.elapsed_time(end)) + state.clear() + + samples.sort() + last = len(samples) - 1 + return tuple(samples[min(int(q * len(samples)), last)] for q in QUANTILES) + + +# --------------------------------------------------------------------------- +# swap-in state +# --------------------------------------------------------------------------- + +# Every miss rep reuses the same host cache; only the *device* state has to be +# rebuilt. torch's caching host allocator never returns pinned blocks to the OS, +# so allocating one per rep starved the transfer suite of its 2.3 GB. +_HOST_CACHES = {} + + +def _host_cache(batch_size, num_top_k, hot_buffer_size, is_dsv4): + key = (batch_size, num_top_k, hot_buffer_size, is_dsv4) + cached = _HOST_CACHES.get(key) + if cached is not None: + return cached + + num_host_items = batch_size * (num_top_k + hot_buffer_size) + 1 + if is_dsv4: + shape = (_div_up(num_host_items, DSV4_PAGE_SIZE), DSV4_PAGE_BYTES) + cache = torch.empty(shape, dtype=torch.uint8, device="cpu").pin_memory() + cache.fill_(7) + else: + shape = (num_host_items, 1, KV_DIM) + cache = torch.empty(shape, dtype=DTYPE, device="cpu").pin_memory() + cache.fill_(1.0) + _HOST_CACHES[key] = cache + return cache + + +def release_host_caches(): + _release(_HOST_CACHES) + + +def _make_swapin_state(batch_size, num_top_k, hot_buffer_size, is_dsv4, regime, block): + slots_per_req = hot_buffer_size + 1 # +1 reserved newest slot + num_device_items = batch_size * slots_per_req + num_host_items = batch_size * (num_top_k + hot_buffer_size) + 1 + host_cache = _host_cache(batch_size, num_top_k, hot_buffer_size, is_dsv4) + + if is_dsv4: + device_buffer = torch.zeros( + (_div_up(num_device_items, DSV4_PAGE_SIZE), DSV4_PAGE_BYTES), + dtype=torch.uint8, + device=DEVICE, + ) + item_size_bytes = DSV4_ITEM_BYTES + else: + device_buffer = torch.zeros( + (num_device_items, 1, KV_DIM), dtype=DTYPE, device=DEVICE + ) + item_size_bytes = LINEAR_ITEM_BYTES + + # Each request owns a contiguous, disjoint span of device slots. + device_buffer_locs = ( + torch.arange(num_device_items, dtype=torch.int32, device=DEVICE) + .view(batch_size, slots_per_req) + .contiguous() + ) + host_cache_locs = ( + torch.arange(num_host_items, dtype=torch.int64, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1) + .contiguous() + ) + + # Request r asks for tokens [base_r, base_r + num_top_k). + base = torch.arange(batch_size, dtype=torch.int32, device=DEVICE).view(-1, 1) * ( + num_top_k + hot_buffer_size + ) + top_k_tokens = ( + base + torch.arange(num_top_k, dtype=torch.int32, device=DEVICE).view(1, -1) + ).contiguous() + + device_buffer_tokens = torch.full( + (batch_size, slots_per_req), -1, dtype=torch.int32, device=DEVICE + ) + if regime == "hit": + device_buffer_tokens[:, :num_top_k] = top_k_tokens + else: + device_buffer_tokens[:, :hot_buffer_size] = ( + base + + num_top_k + + torch.arange(hot_buffer_size, dtype=torch.int32, device=DEVICE).view( + 1, -1 + ) + ) + device_buffer_tokens[:, hot_buffer_size] = top_k_tokens[:, -1] + + lru_slots = ( + torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1) + .contiguous() + ) + torch.xpu.synchronize() + + return { + "top_k_tokens": top_k_tokens, + "device_buffer_tokens": device_buffer_tokens, + "host_cache_locs": host_cache_locs, + "device_buffer_locs": device_buffer_locs, + "host_cache": host_cache, + "device_buffer": device_buffer, + "top_k_device_locs": torch.full_like(top_k_tokens, -1), + "req_pool_indices": torch.arange(batch_size, dtype=torch.int64, device=DEVICE), + # seq_len must exceed hot_buffer_size or the kernel takes the fast path. + "seq_lens": torch.full( + (batch_size,), num_top_k + hot_buffer_size, dtype=torch.int32, device=DEVICE + ), + "lru_slots": lru_slots, + "item_size_bytes": item_size_bytes, + "num_top_k": num_top_k, + "hot_buffer_size": hot_buffer_size, + "page_size": DSV4_PAGE_SIZE if is_dsv4 else 1, + "block_size": block, + "num_real_reqs": torch.tensor([batch_size], dtype=torch.int32, device=DEVICE), + } + + +def _swapin(state, is_dsv4): + fn = ( + load_cache_to_device_buffer_dsv4_mla + if is_dsv4 + else load_cache_to_device_buffer_mla + ) + fn(**state) + + +def _torch_swapin(state, is_dsv4): + """Pure-PyTorch equivalent of ``load_cache_to_device_buffer_*_mla``. + + A *timing* reference only -- the kernel's exact miss-to-slot assignment is + internal, so slot ids may differ; accuracy is covered by tests/. The dominant + cost is structural: indexing a pinned CPU tensor needs the gather indices on + the host, forcing a device->host sync per call. The kernel stays on device. + """ + top_k = state["top_k_tokens"] # [B, K] int32 + dbt = state["device_buffer_tokens"] # [B, S] int32 + dbl = state["device_buffer_locs"] # [B, S] int32 + hcl = state["host_cache_locs"] # [B, H] int64 + host_cache = state["host_cache"] + device_buffer = state["device_buffer"] + lru = state["lru_slots"] # [B, hot] int16 + + # Classify each requested token as hit or miss. + eq = top_k.unsqueeze(2) == dbt.unsqueeze(1) # [B, K, S] + is_hit = eq.any(dim=2) # [B, K] + hit_slot = eq.to(torch.uint8).argmax(dim=2) # [B, K] + slot_is_hit = eq.any(dim=1) # [B, S] + + # Pick evict slots in LRU order; the stable sort brings evictable slots + # first while preserving that order. + lru_long = lru.to(torch.int64) + evictable = ~slot_is_hit.gather(1, lru_long) # [B, hot] in LRU order + order = torch.argsort(~evictable, dim=1, stable=True) + evict_slots = lru_long.gather(1, order) # [B, hot] + + # Rank each miss, and match it to the evict slot of the same rank. + miss = ~is_hit + miss_rank = miss.cumsum(dim=1) - 1 # [B, K] + evict_pick = evict_slots.gather(1, miss_rank.clamp_(min=0)) + assigned = torch.where(miss, evict_pick, hit_slot) # [B, K] + + # Stream the misses in. + b_idx, k_idx = miss.nonzero(as_tuple=True) + if b_idx.numel(): + token = top_k[b_idx, k_idx].to(torch.int64) + host_loc = hcl[b_idx, token] + dev_loc = dbl[b_idx, assigned[b_idx, k_idx]].to(torch.int64) + # Pinned host cache must be indexed with CPU indices -> D2H sync. + host_loc_cpu = host_loc.cpu() + if is_dsv4: + h_val, h_scale = _dsv4_views(host_cache) + d_val, d_scale = _dsv4_views(device_buffer) + hp, ho = _page_split(host_loc_cpu) + dp, do = _page_split(dev_loc) + d_val[dp, do] = h_val[hp, ho].to(DEVICE, non_blocking=True) + d_scale[dp, do] = h_scale[hp, ho].to(DEVICE, non_blocking=True) + else: + staged = host_cache[host_loc_cpu].to(DEVICE, non_blocking=True) + device_buffer[dev_loc] = staged + + dbt[b_idx, assigned[b_idx, k_idx]] = top_k[b_idx, k_idx] + + state["top_k_device_locs"].copy_(dbl.gather(1, assigned)) + # Touched slots become most-recently-used; untouched keep LRU order first. + touched = torch.zeros_like(slot_is_hit) + touched.scatter_(1, assigned, True) + key = touched.gather(1, lru_long).to(torch.uint8) + lru.copy_(lru_long.gather(1, torch.argsort(key, dim=1, stable=True)).to(lru.dtype)) + + +# --------------------------------------------------------------------------- +# transfer state +# --------------------------------------------------------------------------- + +# One live state at a time. triton iterates providers innermost, so a shape's +# block sizes arrive back to back and reuse cuts allocation work 3x. +_TRANSFER_STATE = {"key": None, "value": None} + + +def _transfer_state(num_items, num_layers): + key = (num_items, num_layers) + if _TRANSFER_STATE["key"] == key: + return _TRANSFER_STATE["value"] + + release_transfer_state() + pages = _div_up(num_items, DSV4_PAGE_SIZE) + srcs = [ + torch.full((pages, DSV4_PAGE_BYTES), 3, dtype=torch.uint8, device=DEVICE) + for _ in range(num_layers) + ] + dsts = [ + torch.zeros((pages, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) + for _ in range(num_layers) + ] + value = ( + srcs, + dsts, + # Raw addresses only -- srcs/dsts must stay alive alongside them. + torch.tensor([t.data_ptr() for t in srcs], dtype=torch.uint64, device=DEVICE), + torch.tensor([t.data_ptr() for t in dsts], dtype=torch.uint64, device=DEVICE), + torch.arange(num_items, dtype=torch.int64, device=DEVICE), + ) + torch.xpu.synchronize() + _TRANSFER_STATE.update(key=key, value=value) + return value + + +def _torch_transfer(srcs, dsts, src_indices, dst_indices): + """Pure-PyTorch equivalent of ``transfer_cache_dsv4_mla``. + + The kernel walks all layers in one launch; eager needs an indexed copy per + layer, which is the cost this comparison isolates. + """ + sp, so = _page_split(src_indices) + dp, do = _page_split(dst_indices) + for src, dst in zip(srcs, dsts): + s_val, s_scale = _dsv4_views(src) + d_val, d_scale = _dsv4_views(dst) + d_val[dp, do] = s_val[sp, so] + d_scale[dp, do] = s_scale[sp, so] + + +def release_transfer_state(): + value = _TRANSFER_STATE["value"] + if value is not None: + _release(value[0], value[1]) + _TRANSFER_STATE.update(key=None, value=None) + + +def _vram_budget(): + return int(torch.xpu.get_device_properties(0).total_memory * VRAM_BUDGET_FRACTION) + + +def _transfer_footprint(num_items, num_layers): + return 2 * num_layers * _div_up(num_items, DSV4_PAGE_SIZE) * DSV4_PAGE_BYTES + + +# --------------------------------------------------------------------------- +# swap-in benchmark +# --------------------------------------------------------------------------- + + +def swapin_providers(with_torch): + impls = ("sglang", "torch") if with_torch else ("sglang",) + return [ + f"{impl}-{layout}-{regime}" + for layout in ("linear", "dsv4") + for regime in ("hit", "miss") + for impl in impls + ] + + +_STYLES = [ + ("blue", "-"), + ("blue", "--"), + ("green", "-"), + ("green", "--"), + ("red", "-"), + ("red", "--"), + ("orange", "-"), + ("orange", "--"), +] + + +def benchmark_swapin( + batch_size, num_top_k, hot_buffer_size, provider, block_size, reps +): + impl, layout, regime = provider.split("-") + print( + f"benchmark load_cache_to_device_buffer {provider} batch_size={batch_size} " + f"num_top_k={num_top_k} hot_buffer_size={hot_buffer_size} " + f"block_size={block_size}" + ) + torch.xpu.manual_seed_all(SEED) + + is_dsv4 = layout == "dsv4" + item_bytes = DSV4_ITEM_BYTES if is_dsv4 else LINEAR_ITEM_BYTES + run = _torch_swapin if impl == "torch" else _swapin + + if regime == "miss": + # A miss reads one item from the host and writes one to the device buffer. + nbytes = batch_size * num_top_k * item_bytes * 2 + ms, fast_ms, slow_ms = _bench_cold( + lambda: _make_swapin_state( + batch_size, num_top_k, hot_buffer_size, is_dsv4, "miss", block_size + ), + lambda st: run(st, is_dsv4), + reps=reps, + ) + else: + # Hits move no bytes: index resolution plus LRU bookkeeping only. + nbytes = 0 + state = _make_swapin_state( + batch_size, num_top_k, hot_buffer_size, is_dsv4, "hit", block_size + ) + ms, fast_ms, slow_ms = _bench(lambda: run(state, is_dsv4)) + state.clear() + _release() + + all_results.append( + { + "op": f"load_cache_to_device_buffer [{regime}]", + "impl": impl, + "layout": layout, + "batch_size": batch_size, + "num_top_k": num_top_k, + "hot_buffer_size": hot_buffer_size, + "block_size": block_size, + "us (median)": round(ms * 1000, 2), + "GB/s (median)": round(_gbps(nbytes, ms), 2), + "GB/s (min)": round(_gbps(nbytes, slow_ms), 2), + "GB/s (max)": round(_gbps(nbytes, fast_ms), 2), + } + ) + return ms * 1000, fast_ms * 1000, slow_ms * 1000 + + +def swapin_mark(configs, with_torch=False): + providers = swapin_providers(with_torch) + return triton.testing.Mark( + benchmark_swapin, + triton.testing.Benchmark( + x_names=["batch_size", "num_top_k", "hot_buffer_size"], + x_vals=configs, + line_arg="provider", + line_vals=providers, + line_names=[p.replace("-", " ") for p in providers], + styles=_STYLES[: len(providers)], + ylabel="us", + plot_name="hisparse-load-cache-to-device-buffer", + args={}, + ), + ) + + +# --------------------------------------------------------------------------- +# transfer benchmark +# --------------------------------------------------------------------------- + + +def benchmark_transfer(num_items, num_layers, provider): + is_torch = provider == "torch" + block = 0 if is_torch else int(provider) + print( + f"benchmark transfer_cache_dsv4_mla {provider} num_items={num_items} " + f"num_layers={num_layers}" + ) + torch.xpu.manual_seed_all(SEED) + + footprint = _transfer_footprint(num_items, num_layers) + budget = _vram_budget() + if footprint > budget: + # Say what was skipped: a silent drop reads as "covered" in the table. + print( + f" SKIPPED: needs {footprint / 1024**3:.1f} GiB of VRAM, " + f"budget is {budget / 1024**3:.1f} GiB" + ) + nan = float("nan") + return nan, nan, nan + + srcs, dsts, src_ptrs, dst_ptrs, idx = _transfer_state(num_items, num_layers) + nbytes = num_items * num_layers * DSV4_ITEM_BYTES * 2 # read + write + + if is_torch: + fn = lambda: _torch_transfer(srcs, dsts, idx, idx) # noqa: E731 + else: + fn = lambda: transfer_cache_dsv4_mla( # noqa: E731 + src_ptrs=src_ptrs, + dst_ptrs=dst_ptrs, + src_indices=idx, + dst_indices=idx, + block_size=block, + ) + ms, fast_ms, slow_ms = _bench(fn) + + all_results.append( + { + "op": "transfer_cache_dsv4_mla", + "impl": "torch" if is_torch else "sglang", + "num_items": num_items, + "num_layers": num_layers, + "block_size": block, + "us (median)": round(ms * 1000, 2), + "GB/s (median)": round(_gbps(nbytes, ms), 2), + "GB/s (min)": round(_gbps(nbytes, slow_ms), 2), + "GB/s (max)": round(_gbps(nbytes, fast_ms), 2), + } + ) + return _gbps(nbytes, ms), _gbps(nbytes, slow_ms), _gbps(nbytes, fast_ms) + + +def transfer_mark(configs, block_sizes, with_torch=False): + providers = [str(b) for b in block_sizes] + (["torch"] if with_torch else []) + return triton.testing.Mark( + benchmark_transfer, + triton.testing.Benchmark( + x_names=["num_items", "num_layers"], + x_vals=configs, + line_arg="provider", + line_vals=providers, + line_names=[ + "torch eager" if p == "torch" else f"block={p}" for p in providers + ], + styles=_STYLES[: len(providers)], + ylabel="GB/s", + plot_name="hisparse-transfer-cache-dsv4-mla", + args={}, + ), + ) + + +# --------------------------------------------------------------------------- +# speedup analysis +# --------------------------------------------------------------------------- + + +def speedup_analysis(df, transfer_block_size): + if "impl" not in df.columns or "torch" not in set(df["impl"]): + return + + print("\n" + "=" * 80) + print("SPEEDUP ANALYSIS (torch eager / sglang, higher = kernel is faster)") + print("=" * 80) + + for op, sub in df.groupby("op", sort=False): + sgl = sub[sub["impl"] == "sglang"] + ref = sub[sub["impl"] == "torch"] + if sgl.empty or ref.empty: + continue + if op == "transfer_cache_dsv4_mla": + sgl = sgl[sgl["block_size"] == transfer_block_size] + group, label = "num_layers", f"block={transfer_block_size}" + else: + group, label = "layout", "all layouts" + + keys = [ + c + for c in ("layout", "batch_size", "num_top_k", "num_items", "num_layers") + if c in sub.columns + ] + merged = sgl.merge(ref, on=keys, how="inner", suffixes=("_sgl", "_torch")) + if merged.empty: + continue + merged["speedup"] = merged["us (median)_torch"] / merged["us (median)_sgl"] + + print(f"\n### {op} ({label})\n") + print( + f" overall: avg={merged['speedup'].mean():.2f}x " + f"max={merged['speedup'].max():.2f}x " + f"min={merged['speedup'].min():.2f}x" + ) + print(f"\n by {group}:") + for value, rows in merged.groupby(group, sort=True): + print( + f" {value!s:>8}: avg={rows['speedup'].mean():6.2f}x " + f"max={rows['speedup'].max():6.2f}x " + f"min={rows['speedup'].min():6.2f}x" + ) + + +def check_nonzero(): + probe = torch.tensor([[0, 1, 0], [1, 0, 1]], dtype=torch.bool, device=DEVICE) + rows, cols = probe.nonzero(as_tuple=True) + ok = rows.tolist() == [0, 1, 1] and cols.tolist() == [1, 0, 2] + if not ok: + print( + "WARNING: torch.nonzero is wrong on this stack " + f"(got rows={rows.tolist()} cols={cols.tolist()}, " + "expected rows=[0, 1, 1] cols=[1, 0, 2]). " + "The torch-eager baseline is unreliable; speedups will be overstated." + ) + return ok + + +# --------------------------------------------------------------------------- +# regression tracking +# --------------------------------------------------------------------------- + +# Columns identifying one benchmarked configuration, for joining two runs. +_KEY_COLUMNS = [ + "op", + "impl", + "layout", + "batch_size", + "num_top_k", + "hot_buffer_size", + "num_items", + "num_layers", + "block_size", +] + + +def compare_results(df): + os.makedirs(RESULT_DIR, exist_ok=True) + previous_csv = os.path.join(RESULT_DIR, "previous.csv") + current_csv = os.path.join(RESULT_DIR, "current.csv") + + df.to_csv(current_csv, index=False) + print(f"\nCurrent results saved to: {current_csv}") + + if not os.path.exists(previous_csv): + print(f"No {previous_csv} found, nothing to compare against.") + print("Tip: copy current.csv to previous.csv to set a baseline.") + return + + try: + prev = pd.read_csv(previous_csv) + except Exception as e: # noqa: BLE001 -- a stale CSV must not kill the run + print(f"Error loading {previous_csv}: {e}") + return + print(f"Loaded previous results from: {previous_csv}") + + keys = [c for c in _KEY_COLUMNS if c in df.columns and c in prev.columns] + merged = df.merge(prev, on=keys, how="inner", suffixes=("", "_prev")) + if merged.empty: + print("No configurations in common with the previous run.") + return + + merged["delta %"] = ( + (merged["us (median)"] - merged["us (median)_prev"]) + / merged["us (median)_prev"] + * 100 + ).round(1) + report = merged[keys + ["us (median)_prev", "us (median)", "delta %"]].rename( + columns={"us (median)_prev": "previous us", "us (median)": "current us"} + ) + + print("\n" + "=" * 80) + print("REGRESSION vs previous.csv (positive delta = slower than before)") + print("=" * 80 + "\n") + print(report.dropna(axis=1, how="all").to_markdown(index=False)) + print( + f"\nWorst: {report['delta %'].max():+.1f}% " + f"Best: {report['delta %'].min():+.1f}% " + f"Mean: {report['delta %'].mean():+.1f}%" + ) + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + + +def parse_args(): + p = argparse.ArgumentParser(description="HiSparse KV-cache benchmark for XPU") + p.add_argument( + "--suite", + nargs="+", + choices=["swapin", "transfer"], + default=["swapin", "transfer"], + help="Which ops to benchmark (default: both)", + ) + p.add_argument( + "--batch-sizes", + type=int, + nargs="+", + default=[1, 8, 32, 128], + metavar="B", + help="Swap-in batch sizes (default: 1 8 32 128)", + ) + p.add_argument( + "--top-k", + type=int, + nargs="+", + default=[64, 256], + metavar="K", + help="Swap-in num_top_k values; hot_buffer_size is set equal to each " + "(default: 64 256)", + ) + p.add_argument( + "--swapin-block-size", + type=int, + default=256, + help="block_size for the swap-in op (default: 256)", + ) + p.add_argument( + "--reps", + type=int, + default=20, + help="Reps for the cold (miss) swap-in regime (default: 20)", + ) + p.add_argument( + "--num-items", + type=int, + nargs="+", + default=[64, 512, 4096, 32768], + metavar="N", + help="Transfer item counts (default: 64 512 4096 32768)", + ) + p.add_argument( + "--num-layers", + type=int, + nargs="+", + default=[1, 8, 61], + metavar="L", + help="Transfer layer counts (default: 1 8 61)", + ) + p.add_argument( + "--block-sizes", + type=int, + nargs="+", + default=[256, 512, 1024], + choices=[256, 512, 1024], + metavar="B", + help="block_size values for the transfer op (default: 256 512 1024)", + ) + p.add_argument( + "--with-torch", + action="store_true", + help="Add a torch-eager provider and print a speedup analysis. Off by " + "default: the eager path allocates heavily and its allocator traffic " + "widens the kernels' own measured spread", + ) + p.add_argument( + "--save-csv", + action="store_true", + help="Save results to CSV and diff against the previous run " + "(default: only print results)", + ) + return p.parse_args() + + +def main(): + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + print("ERROR: no XPU device available.") + raise SystemExit(1) + + args = parse_args() + props = torch.xpu.get_device_properties(0) + print(f"Device : {torch.xpu.get_device_name(0)}") + print(f"VRAM : {props.total_memory / 1024**3:.1f} GiB") + print(f"dtype : {DTYPE}") + print(f"suites : {args.suite}") + print(f"torch : {'yes' if args.with_torch else 'no (kernels only)'}\n") + + if args.with_torch: + check_nonzero() + + all_results.clear() + + if "swapin" in args.suite: + configs = [(b, k, k) for b, k in product(args.batch_sizes, args.top_k)] + swapin_mark(configs, args.with_torch).run( + print_data=True, + show_plots=False, + save_path=None, + block_size=args.swapin_block_size, + reps=args.reps, + ) + release_host_caches() + + if "transfer" in args.suite: + configs = list(product(args.num_items, args.num_layers)) + transfer_mark(configs, args.block_sizes, args.with_torch).run( + print_data=True, show_plots=False, save_path=None + ) + release_transfer_state() + + if not all_results: + print("No results collected.") + return + + print("\nBenchmark finished!") + df = pd.DataFrame(all_results) + for op, sub in df.groupby("op", sort=False): + sub = sub.dropna(axis=1, how="all").reset_index(drop=True) + print(f"\n### {op}\n") + print(sub.drop(columns=["op"]).to_markdown(index=False)) + + speedup_analysis(df, max(args.block_sizes)) + + if args.save_csv: + compare_results(df) + + +if __name__ == "__main__": + main() diff --git a/benchmark/bench_jit_hisparse.py b/benchmark/bench_jit_hisparse.py deleted file mode 100644 index 10478e8de..000000000 --- a/benchmark/bench_jit_hisparse.py +++ /dev/null @@ -1,614 +0,0 @@ -"""Benchmark the XPU/SYCL HiSparse KV-offload swap-in kernels vs torch eager. - -Covers the three JIT entry points added alongside ``tests/test_hisparse_jit.py``: - -- ``load_cache_to_device_buffer_mla`` — linear host/device layout -- ``load_cache_to_device_buffer_dsv4_mla`` — page-padded C4 layout -- ``transfer_cache_dsv4_mla`` — bulk evict / backup copy - -Each is compared against a pure-PyTorch (``torch`` provider) implementation of -the same operation, so the reported speedup is kernel vs eager on the same -device and the same inputs. The eager versions are written the way a user -naturally would: vectorized advanced indexing, one ``index`` / ``index_copy_`` -per layer or per request batch, no fusion across layers. - -Two swap-in regimes are reported separately, because both the kernel and the -eager reference mutate ``device_buffer_tokens`` / ``lru_slots`` in place: - -- **hit** — every top-k token is already resident. Idempotent across reps, so - ``triton.testing.do_bench`` measures it directly. This is the steady state. -- **miss** — every top-k token must stream in from the host cache. Only the - *first* call on a given state actually misses, so reps are timed one at a - time with XPU events and the state is rebuilt outside the timed window. - -The kernels are templated on ``(block_size, num_top_k, hot_buffer_size, -is_mla, is_dsv4_layout)`` and each distinct tuple triggers its own ``icpx`` -compile, so the sweep varies ``batch_size`` (a runtime argument) and keeps the -template configuration list short. Expect a one-off JIT compile pause on the -first run; later runs hit the ``~/.cache/sgl_kernel/jit_sycl`` ``.so`` cache. - -Usage (with oneAPI on PATH so JIT compilation can find ``icpx``):: - - source /opt/intel/oneapi/2025.3/oneapi-vars.sh - ZE_AFFINITY_MASK=0 python benchmark/bench_jit_hisparse.py - -Pin the run to a *single* device. These kernels are per-rank, and exposing -several devices to one process halves the achieved memory bandwidth on the -device actually used (measured on Arc Pro B60: 826 GB/s with one device -visible vs 377 GB/s with four). That is a runtime/driver effect, not a -property of these kernels -- a bare ``torch.Tensor.copy_`` shows the same 2x -drop -- but it makes multi-device numbers understate the kernel by ~2x. -""" - -import itertools - -import pandas as pd -import torch -import triton - -try: - from sgl_kernel.jit import ( - load_cache_to_device_buffer_dsv4_mla, - load_cache_to_device_buffer_mla, - transfer_cache_dsv4_mla, - ) - - HAS_SGL_JIT = True -except ImportError: - HAS_SGL_JIT = False - print("Warning: sgl_kernel JIT HiSparse not available") - -DEVICE = "xpu" - -# Linear-layout MLA item: 512 kv-lora + 64 rope, bf16 (matches DSA MLA cache). -KV_DIM = 576 -DTYPE = torch.bfloat16 -LINEAR_ITEM_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size() - -# Page-padded C4 layout constants (must match c4_layout.hpp). -DSV4_PAGE_SIZE = 64 -DSV4_VALUE_BYTES = 576 -DSV4_SCALE_BYTES = 8 -DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES -DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576 -DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE - -BLOCK_SIZE = 256 - -all_results = [] - - -def _pinned(shape, dtype): - return torch.empty(shape, dtype=dtype, device="cpu").pin_memory() - - -def _release(*objs): - """Drop device buffers and return the memory to the driver. - - The largest transfer config allocates ``num_layers`` src+dst pairs - (~2.3 GB at 32768 items x 61 layers). Holding several configs' worth live - at once pushes a 24 GB card into allocator thrash and depresses the - measured bandwidth by ~2x, so each state is freed once it has been timed. - """ - for o in objs: - if isinstance(o, list): - o.clear() - torch.xpu.synchronize() - torch.xpu.empty_cache() - - -def _dsv4_views(cache): - """Expose a page-padded C4 cache as (value, scale) views, no copy. - - A page is [VALUE 0..63][SCALE 0..63][pad to 576B]. Plain slicing + ``view`` - would fail (the slice is not contiguous, row stride is kPageBytes), so use - ``as_strided`` to build genuine views with the layout's own strides. - """ - pages = cache.shape[0] - value = cache.as_strided( - (pages, DSV4_PAGE_SIZE, DSV4_VALUE_BYTES), - (DSV4_PAGE_BYTES, DSV4_VALUE_BYTES, 1), - ) - scale = cache.as_strided( - (pages, DSV4_PAGE_SIZE, DSV4_SCALE_BYTES), - (DSV4_PAGE_BYTES, DSV4_SCALE_BYTES, 1), - storage_offset=DSV4_SCALE_OFFSET, - ) - return value, scale - - -def _page_split(index): - """Logical token slot -> (page number, offset in page).""" - return index // DSV4_PAGE_SIZE, index % DSV4_PAGE_SIZE - - -# --------------------------------------------------------------------------- -# swap-in state -# --------------------------------------------------------------------------- - - -def _make_swapin_state(batch_size, num_top_k, hot_buffer_size, is_dsv4, regime): - """Build inputs for one swap-in configuration. - - ``regime="hit"`` seeds ``device_buffer_tokens`` with exactly the requested - top-k tokens (no host traffic). ``regime="miss"`` seeds it with a disjoint - token range so every top-k lookup misses and must be streamed in. - """ - slots_per_req = hot_buffer_size + 1 # +1 reserved newest slot - num_device_items = batch_size * slots_per_req - # Host cache holds resident + non-resident tokens for every request. - num_host_items = batch_size * (num_top_k + hot_buffer_size) + 1 - - if is_dsv4: - host_pages = (num_host_items + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE - dev_pages = (num_device_items + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE - host_cache = _pinned((host_pages, DSV4_PAGE_BYTES), torch.uint8) - host_cache.fill_(7) - device_buffer = torch.zeros( - (dev_pages, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE - ) - item_size_bytes = DSV4_ITEM_BYTES - else: - host_cache = _pinned((num_host_items, 1, KV_DIM), DTYPE) - host_cache.fill_(1.0) - device_buffer = torch.zeros( - (num_device_items, 1, KV_DIM), dtype=DTYPE, device=DEVICE - ) - item_size_bytes = LINEAR_ITEM_BYTES - - # Each request owns a contiguous, disjoint span of device slots. - device_buffer_locs = ( - torch.arange(num_device_items, dtype=torch.int32, device=DEVICE) - .view(batch_size, slots_per_req) - .contiguous() - ) - host_cache_locs = ( - torch.arange(num_host_items, dtype=torch.int64, device=DEVICE) - .view(1, -1) - .repeat(batch_size, 1) - .contiguous() - ) - - # Request r asks for tokens [base_r, base_r + num_top_k). - base = torch.arange(batch_size, dtype=torch.int32, device=DEVICE).view(-1, 1) * ( - num_top_k + hot_buffer_size - ) - top_k_tokens = ( - base + torch.arange(num_top_k, dtype=torch.int32, device=DEVICE).view(1, -1) - ).contiguous() - - device_buffer_tokens = torch.full( - (batch_size, slots_per_req), -1, dtype=torch.int32, device=DEVICE - ) - if regime == "hit": - # Resident set == requested set, so no host traffic is needed. - device_buffer_tokens[:, :num_top_k] = top_k_tokens - else: - # Resident set is disjoint from the requested set -> every slot misses. - device_buffer_tokens[:, :hot_buffer_size] = ( - base - + num_top_k - + torch.arange(hot_buffer_size, dtype=torch.int32, device=DEVICE).view( - 1, -1 - ) - ) - device_buffer_tokens[:, hot_buffer_size] = top_k_tokens[:, -1] - - lru_slots = ( - torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE) - .view(1, -1) - .repeat(batch_size, 1) - .contiguous() - ) - req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device=DEVICE) - # seq_len must exceed hot_buffer_size or the kernel takes the fast path. - seq_lens = torch.full( - (batch_size,), num_top_k + hot_buffer_size, dtype=torch.int32, device=DEVICE - ) - top_k_device_locs = torch.full_like(top_k_tokens, -1) - num_real_reqs = torch.tensor([batch_size], dtype=torch.int32, device=DEVICE) - torch.xpu.synchronize() - - return { - "top_k_tokens": top_k_tokens, - "device_buffer_tokens": device_buffer_tokens, - "host_cache_locs": host_cache_locs, - "device_buffer_locs": device_buffer_locs, - "host_cache": host_cache, - "device_buffer": device_buffer, - "top_k_device_locs": top_k_device_locs, - "req_pool_indices": req_pool_indices, - "seq_lens": seq_lens, - "lru_slots": lru_slots, - "item_size_bytes": item_size_bytes, - "num_top_k": num_top_k, - "hot_buffer_size": hot_buffer_size, - "page_size": DSV4_PAGE_SIZE if is_dsv4 else 1, - "block_size": BLOCK_SIZE, - "num_real_reqs": num_real_reqs, - } - - -def _sglang_swapin(state, is_dsv4): - fn = ( - load_cache_to_device_buffer_dsv4_mla - if is_dsv4 - else load_cache_to_device_buffer_mla - ) - fn(**state) - - -# --------------------------------------------------------------------------- -# torch eager reference: swap-in -# --------------------------------------------------------------------------- - - -def _torch_swapin(state, is_dsv4): - """Pure-PyTorch equivalent of ``load_cache_to_device_buffer_*_mla``. - - Same observable effect as the kernel: resolve which top-k tokens are already - resident, assign the misses to the least-recently-used evictable slots, - stream those items host->device, and refresh the LRU order. - - This is a *timing* reference. It reproduces the resident-token set and the - LRU ordering, but the kernel's exact miss-to-slot assignment is an internal - detail, so slot ids may differ; accuracy is covered by tests/. - - The dominant cost here is structural, not a missing optimization: indexing a - pinned CPU tensor needs the gather indices on the host, which forces a - device->host sync per call. The kernel does the whole thing on device. - """ - top_k = state["top_k_tokens"] # [B, K] int32 - dbt = state["device_buffer_tokens"] # [B, S] int32 - dbl = state["device_buffer_locs"] # [B, S] int32 - hcl = state["host_cache_locs"] # [B, H] int64 - host_cache = state["host_cache"] - device_buffer = state["device_buffer"] - lru = state["lru_slots"] # [B, hot] int16 - hot = state["hot_buffer_size"] - - # ---- classify each requested token as hit or miss ------------------- - eq = top_k.unsqueeze(2) == dbt.unsqueeze(1) # [B, K, S] - is_hit = eq.any(dim=2) # [B, K] - hit_slot = eq.to(torch.uint8).argmax(dim=2) # [B, K] - slot_is_hit = eq.any(dim=1) # [B, S] - - # ---- pick evict slots in LRU order ---------------------------------- - lru_long = lru.to(torch.int64) - evictable = ~slot_is_hit.gather(1, lru_long) # [B, hot] in LRU order - # Stable sort brings evictable slots first while preserving LRU order. - order = torch.argsort(~evictable, dim=1, stable=True) - evict_slots = lru_long.gather(1, order) # [B, hot] - - # Rank each miss, and match it to the evict slot of the same rank. - miss = ~is_hit - miss_rank = miss.cumsum(dim=1) - 1 # [B, K] - evict_pick = evict_slots.gather(1, miss_rank.clamp_(min=0)) - assigned = torch.where(miss, evict_pick, hit_slot) # [B, K] - - # ---- stream the misses in ------------------------------------------- - b_idx, k_idx = miss.nonzero(as_tuple=True) - if b_idx.numel(): - token = top_k[b_idx, k_idx].to(torch.int64) - host_loc = hcl[b_idx, token] - dev_loc = dbl[b_idx, assigned[b_idx, k_idx]].to(torch.int64) - # Pinned host cache must be indexed with CPU indices -> D2H sync. - host_loc_cpu = host_loc.cpu() - if is_dsv4: - h_val, h_scale = _dsv4_views(host_cache) - d_val, d_scale = _dsv4_views(device_buffer) - hp, ho = _page_split(host_loc_cpu) - dp, do = _page_split(dev_loc) - d_val[dp, do] = h_val[hp, ho].to(DEVICE, non_blocking=True) - d_scale[dp, do] = h_scale[hp, ho].to(DEVICE, non_blocking=True) - else: - staged = host_cache[host_loc_cpu].to(DEVICE, non_blocking=True) - device_buffer[dev_loc] = staged - - # ---- refresh residency + LRU order ------------------------------ - dbt[b_idx, assigned[b_idx, k_idx]] = top_k[b_idx, k_idx] - - state["top_k_device_locs"].copy_(dbl.gather(1, assigned)) - # Touched slots become most-recently-used: untouched keep LRU order first. - touched = torch.zeros_like(slot_is_hit) - touched.scatter_(1, assigned, True) - key = touched.gather(1, lru_long).to(torch.uint8) - lru.copy_(lru_long.gather(1, torch.argsort(key, dim=1, stable=True)).to(lru.dtype)) - - -def _swapin_call(provider, state, is_dsv4): - if provider == "sglang": - _sglang_swapin(state, is_dsv4) - else: - _torch_swapin(state, is_dsv4) - - -def _time_miss_regime( - provider, batch_size, num_top_k, hot_buffer_size, is_dsv4, reps=20 -): - """Time the cold (all-miss) path, rebuilding state outside the timed window. - - ``do_bench`` cannot be used here: the first call makes every token - resident, so reps 2..n would measure the hit path instead. - """ - # Warm up JIT compilation and the module cache before timing anything. - _swapin_call( - provider, - _make_swapin_state(batch_size, num_top_k, hot_buffer_size, is_dsv4, "miss"), - is_dsv4, - ) - torch.xpu.synchronize() - - samples = [] - for _ in range(reps): - state = _make_swapin_state( - batch_size, num_top_k, hot_buffer_size, is_dsv4, "miss" - ) - start, end = torch.xpu.Event(enable_timing=True), torch.xpu.Event( - enable_timing=True - ) - start.record() - _swapin_call(provider, state, is_dsv4) - end.record() - torch.xpu.synchronize() - samples.append(start.elapsed_time(end)) - state.clear() # free before building the next rep's state - _release() - samples.sort() - return samples[len(samples) // 2] # median ms - - -# (num_top_k, hot_buffer_size) — each pair is one extra JIT compile per layout. -TEMPLATE_CONFIGS = [(64, 64), (256, 256)] -BATCH_SIZES = [1, 8, 32, 128] - -swapin_configs = [ - (b, k, h) for b, (k, h) in itertools.product(BATCH_SIZES, TEMPLATE_CONFIGS) -] - -SWAPIN_PROVIDERS = [ - f"{p}-{layout}-{regime}" - for layout in ("linear", "dsv4") - for regime in ("hit", "miss") - for p in ("sglang", "torch") -] - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["batch_size", "num_top_k", "hot_buffer_size"], - x_vals=swapin_configs, - line_arg="provider", - line_vals=SWAPIN_PROVIDERS, - line_names=[p.replace("-", " ") for p in SWAPIN_PROVIDERS], - styles=[ - ("blue", "-"), - ("blue", "--"), - ("cyan", "-"), - ("cyan", "--"), - ("green", "-"), - ("green", "--"), - ("orange", "-"), - ("orange", "--"), - ], - ylabel="us", - plot_name="hisparse-load-cache-to-device-buffer-performance", - args={}, - ) -) -def benchmark_swapin(batch_size, num_top_k, hot_buffer_size, provider): - impl, layout, regime = provider.split("-") - is_dsv4 = layout == "dsv4" - - if regime == "miss": - ms = _time_miss_regime(impl, batch_size, num_top_k, hot_buffer_size, is_dsv4) - min_ms = max_ms = ms - else: - state = _make_swapin_state( - batch_size, num_top_k, hot_buffer_size, is_dsv4, "hit" - ) - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: _swapin_call(impl, state, is_dsv4), quantiles=[0.5, 0.2, 0.8] - ) - state.clear() - _release() - - item_bytes = DSV4_ITEM_BYTES if is_dsv4 else LINEAR_ITEM_BYTES - # A miss reads one item from host and writes one to the device buffer. - moved = 0 if regime == "hit" else batch_size * num_top_k * item_bytes * 2 - all_results.append( - { - "kernel": "load_cache_to_device_buffer", - "provider": impl, - "case": f"{layout}-{regime}", - "batch_size": batch_size, - "num_top_k": num_top_k, - "hot_buffer_size": hot_buffer_size, - "time_us": 1000 * ms, - "GB_s": (moved / (ms * 1e-3) / 1e9) if moved and ms > 0 else float("nan"), - } - ) - return 1000 * ms, 1000 * min_ms, 1000 * max_ms - - -# --------------------------------------------------------------------------- -# transfer_cache_dsv4_mla -# --------------------------------------------------------------------------- - - -def _make_transfer_state(num_items, num_layers, block_size): - pages = (num_items + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE - srcs = [ - torch.full((pages, DSV4_PAGE_BYTES), 3, dtype=torch.uint8, device=DEVICE) - for _ in range(num_layers) - ] - dsts = [ - torch.zeros((pages, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) - for _ in range(num_layers) - ] - src_ptrs = torch.tensor( - [t.data_ptr() for t in srcs], dtype=torch.uint64, device=DEVICE - ) - dst_ptrs = torch.tensor( - [t.data_ptr() for t in dsts], dtype=torch.uint64, device=DEVICE - ) - idx = torch.arange(num_items, dtype=torch.int64, device=DEVICE) - torch.xpu.synchronize() - # Keep srcs/dsts alive: src_ptrs only holds raw addresses. - return srcs, dsts, src_ptrs, dst_ptrs, idx, block_size - - -def _torch_transfer(srcs, dsts, src_indices, dst_indices): - """Pure-PyTorch equivalent of ``transfer_cache_dsv4_mla``. - - The kernel walks all layers inside one launch; eager has to issue an - indexed copy per layer, which is the cost this comparison isolates. - """ - sp, so = _page_split(src_indices) - dp, do = _page_split(dst_indices) - for src, dst in zip(srcs, dsts): - s_val, s_scale = _dsv4_views(src) - d_val, d_scale = _dsv4_views(dst) - d_val[dp, do] = s_val[sp, so] - d_scale[dp, do] = s_scale[sp, so] - - -transfer_configs = list(itertools.product([64, 512, 4096, 32768], [1, 8, 61])) - -TRANSFER_PROVIDERS = ["sglang-bs256", "sglang-bs512", "sglang-bs1024", "torch"] - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["num_items", "num_layers"], - x_vals=transfer_configs, - line_arg="provider", - line_vals=TRANSFER_PROVIDERS, - line_names=[ - "sglang block=256", - "sglang block=512", - "sglang block=1024", - "torch", - ], - styles=[("blue", "-"), ("green", "-"), ("red", "-"), ("orange", "--")], - ylabel="us", - plot_name="hisparse-transfer-cache-dsv4-mla-performance", - args={}, - ) -) -def benchmark_transfer(num_items, num_layers, provider): - is_torch = provider == "torch" - block_size = 1024 if is_torch else int(provider.removeprefix("sglang-bs")) - srcs, dsts, src_ptrs, dst_ptrs, idx, bs = _make_transfer_state( - num_items, num_layers, block_size - ) - - if is_torch: - fn = lambda: _torch_transfer(srcs, dsts, idx, idx) - else: - fn = lambda: transfer_cache_dsv4_mla( - src_ptrs=src_ptrs, - dst_ptrs=dst_ptrs, - src_indices=idx, - dst_indices=idx, - block_size=bs, - ) - ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8]) - _release(srcs, dsts) - - moved = num_items * num_layers * DSV4_ITEM_BYTES * 2 # read + write - all_results.append( - { - "kernel": "transfer_cache_dsv4_mla", - "provider": "torch" if is_torch else "sglang", - "case": "transfer" if is_torch else f"transfer-bs{bs}", - "num_items": num_items, - "num_layers": num_layers, - "time_us": 1000 * ms, - "GB_s": moved / (ms * 1e-3) / 1e9 if ms > 0 else float("nan"), - } - ) - return 1000 * ms, 1000 * min_ms, 1000 * max_ms - - -# --------------------------------------------------------------------------- -# speedup analysis -# --------------------------------------------------------------------------- - - -def _report_speedup(df, index_cols, case_label, title): - """Print the torch-vs-sglang speedup summary for one kernel.""" - pivot = df.pivot_table(index=index_cols, columns="provider", values="time_us") - if "torch" not in pivot.columns or "sglang" not in pivot.columns: - return - pivot = pivot.dropna(subset=["torch", "sglang"]) - if pivot.empty: - return - - pivot["speedup"] = pivot["torch"] / pivot["sglang"] - print("\n" + "=" * 80) - print(f"Speedup Analysis (torch vs sglang) — {title}") - print("=" * 80) - print(f"\nOverall average speedup: {pivot['speedup'].mean():.2f}x") - print(f"Overall max speedup: {pivot['speedup'].max():.2f}x") - print(f"Overall min speedup: {pivot['speedup'].min():.2f}x") - - print(f"\nSpeedup by {case_label}:") - levels = pivot.index.get_level_values(case_label) - for value in dict.fromkeys(levels): - sp = pivot.loc[levels == value, "speedup"] - if not sp.empty: - print( - f" {str(value):>14s}: avg={sp.mean():.2f}x " - f"max={sp.max():.2f}x min={sp.min():.2f}x" - ) - - -if __name__ == "__main__": - if not HAS_SGL_JIT: - print("ERROR: sgl_kernel JIT HiSparse kernels unavailable.") - raise SystemExit(1) - if not (hasattr(torch, "xpu") and torch.xpu.is_available()): - print("ERROR: no XPU device available.") - raise SystemExit(1) - - print("HiSparse swap-in kernels (JIT SYCL) vs torch eager") - print("First run compiles each template configuration with icpx; be patient.") - print("=" * 80) - benchmark_swapin.run(print_data=True) - - print("\n" + "=" * 80) - print("transfer_cache_dsv4_mla (evict / backup path)") - print("=" * 80) - benchmark_transfer.run(print_data=True) - - df = pd.DataFrame(all_results) - df["time_us"] = df["time_us"].round(2) - df["GB_s"] = df["GB_s"].round(2) - - print("\n" + "=" * 80) - print("Raw Results") - print("=" * 80) - print(df.to_markdown(index=False)) - - swapin = df[df["kernel"] == "load_cache_to_device_buffer"] - _report_speedup( - swapin, - ["batch_size", "num_top_k", "hot_buffer_size", "case"], - "case", - "load_cache_to_device_buffer", - ) - - # Compare eager against the default block size only (block is within noise). - transfer = df[ - (df["kernel"] == "transfer_cache_dsv4_mla") - & (df["case"].isin(["transfer", "transfer-bs1024"])) - ] - _report_speedup( - transfer, - ["num_items", "num_layers"], - "num_layers", - "transfer_cache_dsv4_mla (block=1024)", - ) - - print("\nBenchmark finished!") diff --git a/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp b/include/sgl_kernel/hisparse/c4_layout.hpp similarity index 90% rename from include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp rename to include/sgl_kernel/hisparse/c4_layout.hpp index 1791fa944..0546605af 100644 --- a/include/sgl_kernel/jit_kernel/hisparse/c4_layout.hpp +++ b/include/sgl_kernel/hisparse/c4_layout.hpp @@ -1,9 +1,5 @@ /** - * HiSparse C4 paged-cache layout helpers (SYCL / Intel XPU). - * - * Ports device::hisparse::{get_pointer_paged, transfer_item} from the CUDA - * kernel (sglang jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh). - * + * HiSparse C4 paged-cache layout helpers. * Paged C4 cache layout (per page of kPageSize tokens): * VALUE 0, VALUE 1, ..., VALUE 63, (kValueBytes each) * SCALE 0, SCALE 1, ..., SCALE 63, (kScaleBytes each) @@ -21,7 +17,9 @@ namespace sgl { namespace sycl_kernel { namespace hisparse { -// C4 paged layout constants (must match kvcacheio.cuh exactly). +inline constexpr int kSubGroupSize = 32; + +// C4 paged layout constants (must match the model-side cache layout exactly). inline constexpr int64_t kPageSize = 64; inline constexpr int64_t kPageBits = 6; // log2(kPageSize) inline constexpr int64_t kValueBytes = 576; diff --git a/include/sgl_kernel/hisparse/load_cache_to_device_buffer.hpp b/include/sgl_kernel/hisparse/load_cache_to_device_buffer.hpp new file mode 100644 index 000000000..aa62e8e22 --- /dev/null +++ b/include/sgl_kernel/hisparse/load_cache_to_device_buffer.hpp @@ -0,0 +1,449 @@ +/** + * HiSparse: load_cache_to_device_buffer SYCL kernel (Intel XPU). + */ + +#pragma once + +#include +#include +#include + +#include "c4_layout.hpp" + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +static constexpr int32_t kTokenHit = static_cast(0xFFFFFFFF); // -1 sentinel "already resident" +static constexpr int32_t kHashEmpty = -1; + +// Knuth multiplicative hash into an open-addressing table of size hash_size. +// hash_mask is hash_size-1 for power-of-two sizes (else 0), avoiding a modulo. +inline int hash_slot(int32_t key, int hash_size, int hash_mask) { + const uint32_t h = static_cast(key) * 2654435761u; + return hash_mask != 0 ? static_cast(h & static_cast(hash_mask)) + : static_cast(h % static_cast(hash_size)); +} + +// Linear-probe step: (slot + 1) % hash_size without the division. +inline int hash_probe_next(int slot, int hash_size) { + const int next = slot + 1; + return next == hash_size ? 0 : next; +} + +// Cooperative linear (non-paged) item copy across a sub-group, for the generic +// miss path where device + host are both linear with stride item_size_bytes. +inline void transfer_item_linear(int lane_id, int sg_size, const void* src, void* dst, int64_t item_size_bytes) { + const int64_t nwords = item_size_bytes / 8; + const int64_t* s = static_cast(src); + int64_t* d = static_cast(dst); + for (int64_t j = lane_id; j < nwords; j += sg_size) { + d[j] = s[j]; + } + const int64_t tail_start = nwords * 8; + const char* sc = static_cast(src) + tail_start; + char* dc = static_cast(dst) + tail_start; + for (int64_t j = lane_id; j < item_size_bytes - tail_start; j += sg_size) { + dc[j] = sc[j]; + } +} + +// Local-memory atomic CAS returning the previous value. +inline int32_t atomic_cas_local(int32_t* addr, int32_t compare, int32_t val) { + ::sycl::atomic_ref< + int32_t, + ::sycl::memory_order::relaxed, + ::sycl::memory_scope::work_group, + ::sycl::access::address_space::local_space> + ref(*addr); + int32_t expected = compare; + ref.compare_exchange_strong(expected, val); + // On success `expected` is unchanged (== compare); on failure it holds the + // current value. + return expected; +} + +// Single-sub-group inclusive prefix scan over the local-memory window +// [offset, offset+sg_size), threading a running accumulator. +inline int sub_group_inclusive_scan( + const ::sycl::sub_group& sg, int32_t* s_data, int lane_id, int sg_size, int offset, int count, int accumulator) { + const int idx = lane_id + offset; + int val = (idx < count) ? s_data[idx] : 0; + val = ::sycl::inclusive_scan_over_group(sg, val, ::sycl::plus()); + val += accumulator; + if (idx < count) { + s_data[idx] = val; + } + accumulator = ::sycl::group_broadcast(sg, val, sg_size - 1); + return accumulator; +} + +// Local-memory layout: an int32_t region followed by an int16_t region starting +// at int32 slot total_int32, so it inherits 4-byte alignment. num_top_k and +// hot_buffer_size are runtime values, so the host computes the layout and hands +// the offsets to the kernel. +struct SmemLayout { + int hash_size; + int hash_mask; // hash_size-1 if a power of two, else 0 + int num_buffer_chunks; + int num_token_chunks; + int total_int32; + int total_int16; + int total_int32_slots; // allocation size of the int32 local_accessor + + static SmemLayout make(int num_top_k, int hot_buffer_size) { + SmemLayout l{}; + l.hash_size = num_top_k * 2; + l.hash_mask = (l.hash_size & (l.hash_size - 1)) == 0 ? l.hash_size - 1 : 0; + l.num_buffer_chunks = (hot_buffer_size + kSubGroupSize - 1) / kSubGroupSize; + l.num_token_chunks = (num_top_k + kSubGroupSize - 1) / kSubGroupSize; + // int32 region: top_k_tokens + chunk_offset + evict_chunk_offset + hash_keys + // + {total_hits, newest_hit} + l.total_int32 = num_top_k + (l.num_buffer_chunks + 1) + (l.num_buffer_chunks + 1) + l.hash_size + 2; + // int16 region: lru_slots_out + hash_vals + l.total_int16 = hot_buffer_size + l.hash_size; + l.total_int32_slots = + l.total_int32 + static_cast((l.total_int16 * sizeof(int16_t) + sizeof(int32_t) - 1) / sizeof(int32_t)); + return l; + } + + size_t bytes() const { + return static_cast(total_int32_slots) * sizeof(int32_t); + } +}; + +// seq_lens and req_pool_indices are int32 or int64 depending on the caller. Each +// is read once per work-group, so a runtime dtype branch beats instantiating the +// kernel for both index types. +inline int64_t load_index(const void* base, bool is_i64, int bid) { + return is_i64 ? static_cast(base)[bid] : static_cast(static_cast(base)[bid]); +} + +// IsMLA / IsDsv4Layout stay compile-time because they select the inner copy loop. +// Aggregate-initialized so the launcher can name the fields: the list is long +// enough that positional arguments are easy to transpose silently. +template +struct LoadCacheToDeviceBufferKernel { + static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); + + const int32_t* top_k_tokens_; + int32_t* device_buffer_tokens_; + const int64_t* host_cache_locs_; + const int32_t* device_buffer_locs_; + const void* host_cache_k_; + const void* host_cache_v_; + void* device_buffer_k_; + void* device_buffer_v_; + int32_t* top_k_device_locs_; + const void* req_pool_indices_; + const void* seq_lens_; + int16_t* lru_slots_; + const int32_t* num_real_reqs_; + bool req_pool_indices_is_i64_; + bool seq_lens_is_i64_; + int64_t buffer_stride_0_; + int64_t host_stride_; + int64_t lru_slot_stride_0_; + int64_t top_k_tokens_stride_; + int64_t top_k_device_locs_stride_; + int64_t item_size_bytes_; + // Runtime shape / layout, precomputed on the host (see SmemLayout). + int block_size_; + int num_sub_groups_; + int num_top_k_; + int hot_buffer_size_; + int hash_size_; + int hash_mask_; + int num_buffer_chunks_; + int num_token_chunks_; + int iters_per_sg_buffer_; + int iters_per_sg_token_; + int total_int32_; + ::sycl::local_accessor smem_; + + [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { + const int bid = static_cast(item.get_group(0)); + const int tid = static_cast(item.get_local_id(0)); + int32_t* req_top_k_device_locs = top_k_device_locs_ + bid * top_k_device_locs_stride_; + + // A graph-captured batch is padded to the captured size. Keep padded output + // rows invalid without a separate fill kernel. + if (bid >= num_real_reqs_[0]) { + for (int i = tid; i < num_top_k_; i += block_size_) { + req_top_k_device_locs[i] = -1; + } + return; + } + + const ::sycl::sub_group sg = item.get_sub_group(); + const int sg_id = static_cast(sg.get_group_linear_id()); + const int lane_id = static_cast(sg.get_local_linear_id()); + const int sg_size = static_cast(sg.get_max_local_range()[0]); + + const int64_t rid = load_index(req_pool_indices_, req_pool_indices_is_i64_, bid); + const int64_t seq_len = load_index(seq_lens_, seq_lens_is_i64_, bid); + + // Per-request base offsets. + const int32_t* req_top_k_tokens = top_k_tokens_ + bid * top_k_tokens_stride_; + + const int64_t buffer_offset = rid * buffer_stride_0_; + int32_t* req_device_buffer_tokens = device_buffer_tokens_ + buffer_offset; + const int32_t* req_device_buffer_locs = device_buffer_locs_ + buffer_offset; + const int64_t* req_host_cache_locs = host_cache_locs_ + rid * host_stride_; + int16_t* req_lru_slots = lru_slots_ + rid * lru_slot_stride_0_; + + // Fast path: short sequences have all tokens resident in device-buffer order. + if (seq_len <= hot_buffer_size_) { + const int count = (seq_len < num_top_k_) ? static_cast(seq_len) : num_top_k_; + for (int i = tid; i < num_top_k_; i += block_size_) { + int32_t device_loc = -1; + if (i < count) { + const int32_t token_pos = req_top_k_tokens[i]; + if (token_pos >= 0) { + device_loc = req_device_buffer_locs[token_pos]; + } + } + req_top_k_device_locs[i] = device_loc; + } + return; + } + + // Scratch is one int32_t accessor; the int16 region starts at slot + // total_int32_, so both regions stay 4-byte aligned. + int32_t* smem_i32 = &smem_[0]; + int32_t* s_top_k_tokens = smem_i32; // num_top_k + int32_t* s_chunk_offset = s_top_k_tokens + num_top_k_; // num_buffer_chunks + 1 + int32_t* s_evict_chunk_offset = s_chunk_offset + (num_buffer_chunks_ + 1); // num_buffer_chunks + 1 + int32_t* s_hash_keys = s_evict_chunk_offset + (num_buffer_chunks_ + 1); // hash_size + int32_t* s_total_hits_ptr = s_hash_keys + hash_size_; // 1 + int32_t* s_newest_hit_ptr = s_hash_keys + hash_size_ + 1; // 1 + + int16_t* smem_i16 = reinterpret_cast(smem_i32 + total_int32_); + int16_t* s_lru_slots_out = smem_i16; // hot_buffer_size + int16_t* s_hash_vals = s_lru_slots_out + hot_buffer_size_; // hash_size + + // Initialize counters, hash table, and prefix-sum offsets. + if (tid == 0) { + *s_total_hits_ptr = 0; + *s_newest_hit_ptr = 0; + } + for (int i = tid; i < hash_size_; i += block_size_) { + s_hash_keys[i] = kHashEmpty; + } + for (int i = tid; i < num_buffer_chunks_ + 1; i += block_size_) { + s_chunk_offset[i] = 0; + s_evict_chunk_offset[i] = 0; + } + item.barrier(::sycl::access::fence_space::local_space); + + const int newest_slot = hot_buffer_size_; + const int32_t newest_token = static_cast(seq_len - 1); + + // Insert top-k token positions into the local-memory hash table. + for (int i = tid; i < num_top_k_; i += block_size_) { + int32_t token_idx = req_top_k_tokens[i]; + if (token_idx == newest_token) { + // The latest token lives at newest_slot, outside LRU tracking: bind it + // directly and mark it a hit. + s_top_k_tokens[i] = kTokenHit; + req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot]; + *s_newest_hit_ptr = 1; + } else { + int slot = hash_slot(token_idx, hash_size_, hash_mask_); + while (true) { + int32_t old = atomic_cas_local(&s_hash_keys[slot], kHashEmpty, token_idx); + if (old == kHashEmpty || old == token_idx) { + s_hash_vals[slot] = static_cast(i); + break; + } + slot = hash_probe_next(slot, hash_size_); + } + s_top_k_tokens[i] = token_idx; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Pass over hot-buffer slots: classify hits vs evictables and compact them. + int total_hit_count = 0; + int total_evict_count = 0; + for (int iter = 0; iter < iters_per_sg_buffer_; iter++) { + const int chunk_idx = sg_id + iter * num_sub_groups_; + const bool has_valid_chunk = chunk_idx < num_buffer_chunks_; + + const int slot_idx = chunk_idx * kSubGroupSize + lane_id; + const bool has_valid_slot = has_valid_chunk && (slot_idx < hot_buffer_size_); + const int16_t buf_slot = has_valid_slot ? req_lru_slots[slot_idx] : static_cast(-1); + int32_t my_buffer_token = (buf_slot >= 0) ? req_device_buffer_tokens[buf_slot] : -1; + int my_found_top_k_idx = -1; + if (my_buffer_token >= 0) { + int h = hash_slot(my_buffer_token, hash_size_, hash_mask_); + while (true) { + int32_t k = s_hash_keys[h]; + if (k == my_buffer_token) { + my_found_top_k_idx = static_cast(s_hash_vals[h]); + break; + } + if (k == kHashEmpty) break; + h = hash_probe_next(h, hash_size_); + } + } + const bool is_hit = my_found_top_k_idx >= 0; + const bool is_evictable = has_valid_slot && !is_hit; + + // Record hits: bind the top-k index to this resident slot's device loc. + if (is_hit) { + s_top_k_tokens[my_found_top_k_idx] = kTokenHit; + req_top_k_device_locs[my_found_top_k_idx] = req_device_buffer_locs[buf_slot]; + } + + int local_hit_offset = 0; + int local_evict_offset = 0; + if (has_valid_chunk) { + local_hit_offset = ::sycl::exclusive_scan_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); + local_evict_offset = ::sycl::exclusive_scan_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); + const int sg_hits = ::sycl::reduce_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); + const int sg_evicts = ::sycl::reduce_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = sg_hits; + s_evict_chunk_offset[chunk_idx + 1] = sg_evicts; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + if (sg_id == 0) { + // Bound the scan window to num_sub_groups lanes: only that many entries + // were written this iteration, and letting the remaining lanes join with + // the wide count would fold stale values from earlier iterations into the + // accumulator and write it into slots future iterations read. + const int scan_count = ::std::min(chunk_idx + 1 + num_sub_groups_, num_buffer_chunks_ + 1); + total_hit_count = + sub_group_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, scan_count, total_hit_count); + total_evict_count = sub_group_inclusive_scan( + sg, s_evict_chunk_offset, lane_id, sg_size, chunk_idx + 1, scan_count, total_evict_count); + if (tid == 0) { + *s_total_hits_ptr = total_hit_count; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Hits grow forward from index 0. + if (is_hit) { + int hit_offset = s_chunk_offset[chunk_idx] + local_hit_offset; + s_lru_slots_out[hit_offset] = buf_slot; + } + // Evictables grow backward from hot_buffer_size - 1. + if (is_evictable) { + int evict_offset = s_evict_chunk_offset[chunk_idx] + local_evict_offset; + s_lru_slots_out[hot_buffer_size_ - 1 - evict_offset] = buf_slot; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Reset offsets for the miss-counting phase (num_token_chunks + 1 entries). + for (int i = tid; i < num_token_chunks_ + 1; i += block_size_) { + s_chunk_offset[i] = 0; + } + item.barrier(::sycl::access::fence_space::local_space); + + // Pass over top-k tokens: identify misses and assign them evictable slots. + int total_misses = 0; + for (int iter = 0; iter < iters_per_sg_token_; iter++) { + const int chunk_idx = sg_id + iter * num_sub_groups_; + const bool has_valid_chunk = chunk_idx < num_token_chunks_; + + const int chunk_token_start = chunk_idx * kSubGroupSize; + const int my_token_idx = chunk_token_start + lane_id; + const bool has_valid_token = has_valid_chunk && (my_token_idx < num_top_k_); + + int32_t my_token = 0; + bool is_miss = false; + int local_miss_offset = 0; + + if (has_valid_token) { + is_miss = s_top_k_tokens[my_token_idx] != kTokenHit; + if (is_miss) { + my_token = s_top_k_tokens[my_token_idx]; + } + } + + if (has_valid_chunk) { + local_miss_offset = ::sycl::exclusive_scan_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); + const int sg_miss_count = ::sycl::reduce_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = sg_miss_count; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + if (sg_id == 0) { + // Same bounded window as the buffer pass above. + const int scan_count = ::std::min(chunk_idx + 1 + num_sub_groups_, num_token_chunks_ + 1); + total_misses = + sub_group_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, scan_count, total_misses); + } + item.barrier(::sycl::access::fence_space::local_space); + + if (is_miss) { + int miss_offset = s_chunk_offset[chunk_idx] + local_miss_offset; + int16_t evict_slot = s_lru_slots_out[hot_buffer_size_ - 1 - miss_offset]; + // Reuse s_top_k_tokens as miss scratch: miss_offset < my_token_idx always + // holds (hits are skipped), so compacted writes never overrun pending reads. + s_top_k_tokens[miss_offset] = my_token; + req_top_k_device_locs[my_token_idx] = req_device_buffer_locs[evict_slot]; + req_device_buffer_tokens[evict_slot] = my_token; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + total_misses = num_top_k_ - *s_total_hits_ptr - *s_newest_hit_ptr; + // Rewrite LRU order: misses then remaining evictables at the front (LRU), + // hits at the back (MRU). + { + const int total_evictable = hot_buffer_size_ - *s_total_hits_ptr; + for (int i = tid; i < hot_buffer_size_; i += block_size_) { + if (i < total_misses) { + req_lru_slots[total_evictable - total_misses + i] = s_lru_slots_out[hot_buffer_size_ - 1 - i]; + } else if (i < total_evictable) { + req_lru_slots[i - total_misses] = s_lru_slots_out[hot_buffer_size_ - 1 - i]; + } else { + req_lru_slots[i] = s_lru_slots_out[i - total_evictable]; + } + } + } + + // Each sub-group copies one miss directly from host cache to device buffer. + for (int miss_idx = sg_id; miss_idx < total_misses; miss_idx += num_sub_groups_) { + const int32_t miss_token = s_top_k_tokens[miss_idx]; + const int16_t evict_slot = s_lru_slots_out[hot_buffer_size_ - 1 - miss_idx]; + + const int64_t src_loc = req_host_cache_locs[miss_token]; + const int64_t dst_loc = static_cast(req_device_buffer_locs[evict_slot]); + + if constexpr (IsDsv4Layout) { + // Page-padded C4 device layout + page-padded host layout, K-only. + transfer_item( + lane_id, + sg_size, + device_buffer_k_, + const_cast(host_cache_k_), + static_cast(dst_loc), + static_cast(src_loc)); + } else { + // Generic path: device + host both linear, stride == item_size_bytes. + const char* src_k = static_cast(host_cache_k_) + src_loc * item_size_bytes_; + char* dst_k = static_cast(device_buffer_k_) + dst_loc * item_size_bytes_; + transfer_item_linear(lane_id, sg_size, src_k, dst_k, item_size_bytes_); + + if constexpr (!IsMLA) { + const char* src_v = static_cast(host_cache_v_) + src_loc * item_size_bytes_; + char* dst_v = static_cast(device_buffer_v_) + dst_loc * item_size_bytes_; + transfer_item_linear(lane_id, sg_size, src_v, dst_v, item_size_bytes_); + } + } + } + } +}; + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/include/sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp b/include/sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp new file mode 100644 index 000000000..2434689b5 --- /dev/null +++ b/include/sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp @@ -0,0 +1,58 @@ +/** + * HiSparse: transfer_cache_dsv4_mla SYCL kernel. + * + * Bulk-copies DSv4-MLA C4 tokens between two sets of page-padded C4 buffers, one + * set per model layer. One sub-group copies one item across all layers, with a + * global-stride loop over items. + * + * src_caches / dst_caches are device arrays of `num_layers` raw cache base + * pointers (uint64_t values), one per layer. + */ + +#pragma once + +#include +#include + +#include "c4_layout.hpp" + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +// Aggregate-initialized so the launcher can name the fields. +template +struct TransferCacheDsv4MlaKernel { + static_assert(BLOCK_SIZE % kSubGroupSize == 0, "BLOCK_SIZE must be a multiple of the sub-group size (32)."); + static constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + + void** src_caches_; + void** dst_caches_; + const int64_t* src_indices_; + const int64_t* dst_indices_; + uint32_t num_items_; + uint32_t num_layers_; + uint32_t total_sub_groups_; + + [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { + const ::sycl::sub_group sg = item.get_sub_group(); + const int lane_id = static_cast(sg.get_local_linear_id()); + const int sg_size = static_cast(sg.get_max_local_range()[0]); + + // Global sub-group index: group * subgroups_per_group + local subgroup index. + const uint32_t global_sg = + static_cast(item.get_group(0)) * kNumSubGroups + static_cast(sg.get_group_linear_id()); + + for (uint32_t i = global_sg; i < num_items_; i += total_sub_groups_) { + const int32_t src_index = static_cast(src_indices_[i]); + const int32_t dst_index = static_cast(dst_indices_[i]); + for (uint32_t layer_id = 0; layer_id < num_layers_; ++layer_id) { + transfer_item(lane_id, sg_size, dst_caches_[layer_id], src_caches_[layer_id], dst_index, src_index); + } + } + } +}; + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp b/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp deleted file mode 100644 index 37a923652..000000000 --- a/include/sgl_kernel/jit_kernel/hisparse/load_cache_to_device_buffer.hpp +++ /dev/null @@ -1,623 +0,0 @@ -/** - * HiSparse: load_cache_to_device_buffer SYCL kernel (Intel XPU). - * - * Ports load_cache_to_device_buffer_kernel<...> from the CUDA source - * (sglang jit_kernel/csrc/hisparse.cuh). One work-group processes one request: - * it hashes the request's top-k token positions, classifies the current hot - * device buffer slots into hits / evictables, streams the missing tokens in - * from the host cache into evicted slots, and rewrites the per-request LRU - * order (evictables at the front, hits at the back). - * - * CUDA -> SYCL mapping: - * - warp (32 lanes) -> sub-group pinned to kWarpSize (32) - * - __ballot + popc(&before) -> exclusive_scan_over_group (local prefix) - * - popc(mask) -> reduce_over_group (sub-group total) - * - __shfl_up / __shfl -> inclusive_scan_over_group / group_broadcast - * - atomicCAS (shared) -> atomic_ref<..., local_space> - * - extern __shared__ -> local_accessor - * - __syncthreads() -> item.barrier(local_space) - * - * The sub-group is pinned to 32 lanes so the slot<->lane mapping - * (slot_idx = chunk * 32 + lane) and the shared-memory layout match the CUDA - * kernel bit-for-bit, giving identical eviction ordering and outputs. - */ - -#pragma once - -#include -#include - -#include "c4_layout.hpp" - -namespace sgl { -namespace sycl_kernel { -namespace hisparse { - -// Fixed logical warp width (matches the CUDA kernel this ports). -static constexpr int kWarpSize = 32; - -static constexpr int32_t kTokenHit = static_cast(0xFFFFFFFF); // -1 sentinel "already resident" -static constexpr int32_t kHashEmpty = -1; - -// Knuth multiplicative hash for the open-addressing table of size hash_size. -inline int hash_slot(int32_t key, int hash_size) { - return static_cast((static_cast(key) * 2654435761u) % static_cast(hash_size)); -} - -// Cooperative linear (non-paged) item copy across a sub-group. Used by the -// generic (non-DSv4) miss-copy path where device + host are both linear with -// stride == item_size_bytes. -inline void transfer_item_linear(int lane_id, int sg_size, const void* src, void* dst, int64_t item_size_bytes) { - const int64_t nwords = item_size_bytes / 8; - const int64_t* s = static_cast(src); - int64_t* d = static_cast(dst); - for (int64_t j = lane_id; j < nwords; j += sg_size) { - d[j] = s[j]; - } - const int64_t tail_start = nwords * 8; - const char* sc = static_cast(src) + tail_start; - char* dc = static_cast(dst) + tail_start; - for (int64_t j = lane_id; j < item_size_bytes - tail_start; j += sg_size) { - dc[j] = sc[j]; - } -} - -// Shared-memory size calculation (mirrors the CUDA SmemLayout). -// Layout: int32_t region (4-byte aligned) followed by int16_t region. -template -struct SmemLayout { - static constexpr int HASH_SIZE = NUM_TOP_K * 2; - static constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + kWarpSize - 1) / kWarpSize; - // int32_t region: top_k_tokens + chunk_offset + evict_chunk_offset + hash_keys + {total_hits, newest_hit} - static constexpr int TOTAL_INT32 = NUM_TOP_K + (NUM_BUFFER_CHUNKS + 1) + (NUM_BUFFER_CHUNKS + 1) + HASH_SIZE + 2; - // int16_t region: lru_slots_out + hash_vals - static constexpr int TOTAL_INT16 = HOT_BUFFER_SIZE + HASH_SIZE; - static constexpr size_t BYTES = TOTAL_INT32 * sizeof(int32_t) + TOTAL_INT16 * sizeof(int16_t); - // Round int16 region up to whole int32 slots so a single int32_t local_accessor - // holds both regions and stays 4-byte aligned throughout. - static constexpr int TOTAL_INT32_SLOTS = - TOTAL_INT32 + (TOTAL_INT16 * sizeof(int16_t) + sizeof(int32_t) - 1) / sizeof(int32_t); -}; - -// Local (shared) memory atomic CAS returning the previous value, matching -// CUDA atomicCAS(addr, compare, val) semantics. -inline int32_t atomic_cas_local(int32_t* addr, int32_t compare, int32_t val) { - ::sycl::atomic_ref< - int32_t, - ::sycl::memory_order::relaxed, - ::sycl::memory_scope::work_group, - ::sycl::access::address_space::local_space> - ref(*addr); - int32_t expected = compare; - ref.compare_exchange_strong(expected, val); - // On success `expected` is unchanged (== compare); on failure it holds the - // current value. Either way this equals CUDA atomicCAS's return value. - return expected; -} - -// Single-sub-group inclusive prefix scan over a shared array window -// [offset, offset+kWarpSize), threading a running accumulator. Mirrors the CUDA -// warp_inclusive_scan (which used __shfl_up_sync / __shfl_sync). -inline int warp_inclusive_scan( - const ::sycl::sub_group& sg, int32_t* s_data, int lane_id, int sg_size, int offset, int count, int accumulator) { - const int idx = lane_id + offset; - int val = (idx < count) ? s_data[idx] : 0; - val = ::sycl::inclusive_scan_over_group(sg, val, ::sycl::plus()); - val += accumulator; - if (idx < count) { - s_data[idx] = val; - } - accumulator = ::sycl::group_broadcast(sg, val, sg_size - 1); - return accumulator; -} - -template < - int BLOCK_SIZE, - int NUM_TOP_K, - int HOT_BUFFER_SIZE, - bool IsMLA, - bool IsDsv4Layout, - typename SeqLensT, - typename ReqPoolIndicesT> -class LoadCacheToDeviceBufferKernel { - public: - static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); - static_assert(BLOCK_SIZE % kWarpSize == 0, "BLOCK_SIZE must be a multiple of the warp size (32)."); - - using Layout = SmemLayout; - static constexpr int NUM_WARPS = BLOCK_SIZE / kWarpSize; - static constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K + kWarpSize - 1) / kWarpSize; - static constexpr int NUM_BUFFER_CHUNKS = Layout::NUM_BUFFER_CHUNKS; - static constexpr int HASH_SIZE = Layout::HASH_SIZE; - - LoadCacheToDeviceBufferKernel( - const int32_t* top_k_tokens, - int32_t* device_buffer_tokens, - const int64_t* host_cache_locs, - const int32_t* device_buffer_locs, - const void* host_cache_k, - const void* host_cache_v, - void* device_buffer_k, - void* device_buffer_v, - int32_t* top_k_device_locs, - const ReqPoolIndicesT* req_pool_indices, - const SeqLensT* seq_lens, - int16_t* lru_slots, - const int32_t* num_real_reqs, - int64_t buffer_stride_0, - int64_t host_stride, - int64_t lru_slot_stride_0, - int64_t top_k_tokens_stride, - int64_t top_k_device_locs_stride, - int64_t page_size, - int64_t item_size_bytes, - ::sycl::local_accessor smem) - : top_k_tokens_(top_k_tokens), - device_buffer_tokens_(device_buffer_tokens), - host_cache_locs_(host_cache_locs), - device_buffer_locs_(device_buffer_locs), - host_cache_k_(host_cache_k), - host_cache_v_(host_cache_v), - device_buffer_k_(device_buffer_k), - device_buffer_v_(device_buffer_v), - top_k_device_locs_(top_k_device_locs), - req_pool_indices_(req_pool_indices), - seq_lens_(seq_lens), - lru_slots_(lru_slots), - num_real_reqs_(num_real_reqs), - buffer_stride_0_(buffer_stride_0), - host_stride_(host_stride), - lru_slot_stride_0_(lru_slot_stride_0), - top_k_tokens_stride_(top_k_tokens_stride), - top_k_device_locs_stride_(top_k_device_locs_stride), - page_size_(page_size), - item_size_bytes_(item_size_bytes), - smem_(smem) {} - - [[sycl::reqd_sub_group_size(kWarpSize)]] void operator()(::sycl::nd_item<1> item) const { - const int bid = static_cast(item.get_group(0)); - // Early exit for padded blocks (CUDA graph pads batch to a captured size). - if (bid >= num_real_reqs_[0]) return; - - const ::sycl::sub_group sg = item.get_sub_group(); - const int tid = static_cast(item.get_local_id(0)); - const int warp_id = static_cast(sg.get_group_linear_id()); - const int lane_id = static_cast(sg.get_local_linear_id()); - const int sg_size = static_cast(sg.get_max_local_range()[0]); - - const int64_t rid = static_cast(req_pool_indices_[bid]); - const int64_t seq_len = static_cast(seq_lens_[bid]); - - // Per-request base offsets. - const int32_t* req_top_k_tokens = top_k_tokens_ + bid * top_k_tokens_stride_; - int32_t* req_top_k_device_locs = top_k_device_locs_ + bid * top_k_device_locs_stride_; - - const int64_t buffer_offset = rid * buffer_stride_0_; - int32_t* req_device_buffer_tokens = device_buffer_tokens_ + buffer_offset; - const int32_t* req_device_buffer_locs = device_buffer_locs_ + buffer_offset; - const int64_t* req_host_cache_locs = host_cache_locs_ + rid * host_stride_; - int16_t* req_lru_slots = lru_slots_ + rid * lru_slot_stride_0_; - - // Fast path: short sequences have all tokens resident in device-buffer order. - if (seq_len <= HOT_BUFFER_SIZE) { - const int count = (seq_len < NUM_TOP_K) ? static_cast(seq_len) : NUM_TOP_K; - for (int i = tid; i < count; i += BLOCK_SIZE) { - int32_t token_pos = req_top_k_tokens[i]; - if (token_pos >= 0) { - req_top_k_device_locs[i] = req_device_buffer_locs[token_pos]; - } - } - return; - } - - // Scratch is allocated as int32_t (see kernel launch), giving 4-byte - // alignment for the int32 region. int16 region follows immediately after - // TOTAL_INT32 int32 slots, so its base is also 4-byte aligned. - int32_t* smem_i32 = &smem_[0]; - int32_t* s_top_k_tokens = smem_i32; // NUM_TOP_K - int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K; // NUM_BUFFER_CHUNKS + 1 - int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // NUM_BUFFER_CHUNKS + 1 - int32_t* s_hash_keys = s_evict_chunk_offset + (NUM_BUFFER_CHUNKS + 1); // HASH_SIZE - int32_t* s_total_hits_ptr = s_hash_keys + HASH_SIZE; // 1 - int32_t* s_newest_hit_ptr = s_hash_keys + HASH_SIZE + 1; // 1 - - int16_t* smem_i16 = reinterpret_cast(smem_i32 + Layout::TOTAL_INT32); - int16_t* s_lru_slots_out = smem_i16; // HOT_BUFFER_SIZE - int16_t* s_hash_vals = s_lru_slots_out + HOT_BUFFER_SIZE; // HASH_SIZE - - // Initialize counters, hash table, and prefix-sum offsets. - if (tid == 0) { - *s_total_hits_ptr = 0; - *s_newest_hit_ptr = 0; - } - for (int i = tid; i < HASH_SIZE; i += BLOCK_SIZE) { - s_hash_keys[i] = kHashEmpty; - } - for (int i = tid; i < NUM_BUFFER_CHUNKS + 1; i += BLOCK_SIZE) { - s_chunk_offset[i] = 0; - s_evict_chunk_offset[i] = 0; - } - item.barrier(::sycl::access::fence_space::local_space); - - const int newest_slot = HOT_BUFFER_SIZE; - const int32_t newest_token = static_cast(seq_len - 1); - - // Insert top-k token positions into the shared-memory hash table. - for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) { - int32_t token_idx = req_top_k_tokens[i]; - if (token_idx == newest_token) { - // The latest token lives at newest_slot (first slot of the extra page), - // excluded from LRU tracking. Bind and mark it as a hit. - s_top_k_tokens[i] = kTokenHit; - req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot]; - *s_newest_hit_ptr = 1; - } else { - int slot = hash_slot(token_idx, HASH_SIZE); - while (true) { - int32_t old = atomic_cas_local(&s_hash_keys[slot], kHashEmpty, token_idx); - if (old == kHashEmpty || old == token_idx) { - s_hash_vals[slot] = static_cast(i); - break; - } - slot = (slot + 1) % HASH_SIZE; - } - s_top_k_tokens[i] = token_idx; - } - } - item.barrier(::sycl::access::fence_space::local_space); - - // Pass over hot-buffer slots: classify hits vs evictables and compact them. - constexpr int ITERATIONS_PER_WARP_BUFFER = (NUM_BUFFER_CHUNKS + NUM_WARPS - 1) / NUM_WARPS; - int total_hit_count = 0; - int total_evict_count = 0; - for (int iter = 0; iter < ITERATIONS_PER_WARP_BUFFER; iter++) { - const int chunk_idx = warp_id + iter * NUM_WARPS; - const bool has_valid_chunk = chunk_idx < NUM_BUFFER_CHUNKS; - - const int slot_idx = chunk_idx * kWarpSize + lane_id; - const bool has_valid_slot = has_valid_chunk && (slot_idx < HOT_BUFFER_SIZE); - const int16_t buf_slot = has_valid_slot ? req_lru_slots[slot_idx] : static_cast(-1); - int32_t my_buffer_token = (buf_slot >= 0) ? req_device_buffer_tokens[buf_slot] : -1; - int my_found_top_k_idx = -1; - if (my_buffer_token >= 0) { - int h = hash_slot(my_buffer_token, HASH_SIZE); - while (true) { - int32_t k = s_hash_keys[h]; - if (k == my_buffer_token) { - my_found_top_k_idx = static_cast(s_hash_vals[h]); - break; - } - if (k == kHashEmpty) break; - h = (h + 1) % HASH_SIZE; - } - } - const bool is_hit = my_found_top_k_idx >= 0; - const bool is_evictable = has_valid_slot && !is_hit; - - // Record hits: bind the top-k index to this resident slot's device loc. - if (is_hit) { - s_top_k_tokens[my_found_top_k_idx] = kTokenHit; - req_top_k_device_locs[my_found_top_k_idx] = req_device_buffer_locs[buf_slot]; - } - - int local_hit_offset = 0; - int local_evict_offset = 0; - if (has_valid_chunk) { - local_hit_offset = ::sycl::exclusive_scan_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); - local_evict_offset = ::sycl::exclusive_scan_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); - const int warp_hits = ::sycl::reduce_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); - const int warp_evicts = ::sycl::reduce_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); - if (lane_id == 0) { - s_chunk_offset[chunk_idx + 1] = warp_hits; - s_evict_chunk_offset[chunk_idx + 1] = warp_evicts; - } - } - item.barrier(::sycl::access::fence_space::local_space); - - if (warp_id == 0) { - total_hit_count = warp_inclusive_scan( - sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_hit_count); - total_evict_count = warp_inclusive_scan( - sg, s_evict_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_evict_count); - if (tid == 0) { - *s_total_hits_ptr = total_hit_count; - } - } - item.barrier(::sycl::access::fence_space::local_space); - - // Hits grow forward from index 0. - if (is_hit) { - int hit_offset = s_chunk_offset[chunk_idx] + local_hit_offset; - s_lru_slots_out[hit_offset] = buf_slot; - } - // Evictables grow backward from HOT_BUFFER_SIZE - 1. - if (is_evictable) { - int evict_offset = s_evict_chunk_offset[chunk_idx] + local_evict_offset; - s_lru_slots_out[HOT_BUFFER_SIZE - 1 - evict_offset] = buf_slot; - } - } - item.barrier(::sycl::access::fence_space::local_space); - - // Reset offsets for the miss-counting phase (NUM_TOKEN_CHUNKS + 1 entries). - for (int i = tid; i < NUM_TOKEN_CHUNKS + 1; i += BLOCK_SIZE) { - s_chunk_offset[i] = 0; - } - item.barrier(::sycl::access::fence_space::local_space); - - // Pass over top-k tokens: identify misses and assign them evictable slots. - int total_misses = 0; - constexpr int ITERATIONS_PER_WARP_TOKEN = (NUM_TOKEN_CHUNKS + NUM_WARPS - 1) / NUM_WARPS; - for (int iter = 0; iter < ITERATIONS_PER_WARP_TOKEN; iter++) { - const int chunk_idx = warp_id + iter * NUM_WARPS; - const bool has_valid_chunk = chunk_idx < NUM_TOKEN_CHUNKS; - - const int chunk_token_start = chunk_idx * kWarpSize; - const int my_token_idx = chunk_token_start + lane_id; - const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K); - - int32_t my_token = 0; - bool is_miss = false; - int local_miss_offset = 0; - - if (has_valid_token) { - is_miss = s_top_k_tokens[my_token_idx] != kTokenHit; - if (is_miss) { - my_token = s_top_k_tokens[my_token_idx]; - } - } - - if (has_valid_chunk) { - local_miss_offset = ::sycl::exclusive_scan_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); - const int warp_miss_count = ::sycl::reduce_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); - if (lane_id == 0) { - s_chunk_offset[chunk_idx + 1] = warp_miss_count; - } - } - item.barrier(::sycl::access::fence_space::local_space); - - if (warp_id == 0) { - total_misses = warp_inclusive_scan( - sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses); - } - item.barrier(::sycl::access::fence_space::local_space); - - if (is_miss) { - int miss_offset = s_chunk_offset[chunk_idx] + local_miss_offset; - int16_t evict_slot = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - miss_offset]; - // Reuse s_top_k_tokens as miss scratch: miss_offset < my_token_idx always - // holds (hits are skipped), so compacted writes never overrun pending reads. - s_top_k_tokens[miss_offset] = my_token; - req_top_k_device_locs[my_token_idx] = req_device_buffer_locs[evict_slot]; - req_device_buffer_tokens[evict_slot] = my_token; - } - } - item.barrier(::sycl::access::fence_space::local_space); - - total_misses = NUM_TOP_K - *s_total_hits_ptr - *s_newest_hit_ptr; - // Rewrite LRU order: misses then remaining evictables at the front (LRU), - // hits at the back (MRU). - { - const int total_evictable = HOT_BUFFER_SIZE - *s_total_hits_ptr; - for (int i = tid; i < HOT_BUFFER_SIZE; i += BLOCK_SIZE) { - if (i < total_misses) { - req_lru_slots[total_evictable - total_misses + i] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i]; - } else if (i < total_evictable) { - req_lru_slots[i - total_misses] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i]; - } else { - req_lru_slots[i] = s_lru_slots_out[i - total_evictable]; - } - } - } - - // Each sub-group copies one miss directly from host cache to device buffer. - for (int miss_idx = warp_id; miss_idx < total_misses; miss_idx += NUM_WARPS) { - const int32_t miss_token = s_top_k_tokens[miss_idx]; - const int16_t evict_slot = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - miss_idx]; - - const int64_t src_loc = req_host_cache_locs[miss_token]; - const int64_t dst_loc = static_cast(req_device_buffer_locs[evict_slot]); - - if constexpr (IsDsv4Layout) { - // Page-padded C4 device layout + page-padded host layout, K-only. - transfer_item( - lane_id, - sg_size, - device_buffer_k_, - const_cast(host_cache_k_), - static_cast(dst_loc), - static_cast(src_loc)); - } else { - // Generic path: device + host both linear, stride == item_size_bytes. - const char* src_k = static_cast(host_cache_k_) + src_loc * item_size_bytes_; - char* dst_k = static_cast(device_buffer_k_) + dst_loc * item_size_bytes_; - transfer_item_linear(lane_id, sg_size, src_k, dst_k, item_size_bytes_); - - if constexpr (!IsMLA) { - const char* src_v = static_cast(host_cache_v_) + src_loc * item_size_bytes_; - char* dst_v = static_cast(device_buffer_v_) + dst_loc * item_size_bytes_; - transfer_item_linear(lane_id, sg_size, src_v, dst_v, item_size_bytes_); - } - } - } - } - - private: - const int32_t* top_k_tokens_; - int32_t* device_buffer_tokens_; - const int64_t* host_cache_locs_; - const int32_t* device_buffer_locs_; - const void* host_cache_k_; - const void* host_cache_v_; - void* device_buffer_k_; - void* device_buffer_v_; - int32_t* top_k_device_locs_; - const ReqPoolIndicesT* req_pool_indices_; - const SeqLensT* seq_lens_; - int16_t* lru_slots_; - const int32_t* num_real_reqs_; - int64_t buffer_stride_0_; - int64_t host_stride_; - int64_t lru_slot_stride_0_; - int64_t top_k_tokens_stride_; - int64_t top_k_device_locs_stride_; - int64_t page_size_; - int64_t item_size_bytes_; - ::sycl::local_accessor smem_; -}; - -template < - int BLOCK_SIZE, - int NUM_TOP_K, - int HOT_BUFFER_SIZE, - bool IsMLA, - bool IsDsv4Layout, - typename SeqLensT, - typename ReqPoolIndicesT> -void load_cache_to_device_buffer_launcher( - ::sycl::queue& queue, - const void* top_k_tokens, - void* device_buffer_tokens, - const void* host_cache_locs, - const void* device_buffer_locs, - const void* host_cache_k, - const void* host_cache_v, - void* device_buffer_k, - void* device_buffer_v, - void* top_k_device_locs, - const void* req_pool_indices, - const void* seq_lens, - void* lru_slots, - const void* num_real_reqs, - int64_t batch_size, - int64_t buffer_stride_0, - int64_t host_stride, - int64_t lru_slot_stride_0, - int64_t top_k_tokens_stride, - int64_t top_k_device_locs_stride, - int64_t page_size, - int64_t item_size_bytes) { - if (batch_size == 0) { - return; - } - using Kernel = LoadCacheToDeviceBufferKernel< - BLOCK_SIZE, - NUM_TOP_K, - HOT_BUFFER_SIZE, - IsMLA, - IsDsv4Layout, - SeqLensT, - ReqPoolIndicesT>; - constexpr size_t smem_slots = SmemLayout::TOTAL_INT32_SLOTS; - - queue.submit([&](::sycl::handler& cgh) { - ::sycl::local_accessor smem(::sycl::range<1>(smem_slots), cgh); - cgh.parallel_for( - ::sycl::nd_range<1>( - ::sycl::range<1>(static_cast(batch_size) * BLOCK_SIZE), ::sycl::range<1>(BLOCK_SIZE)), - Kernel( - static_cast(top_k_tokens), - static_cast(device_buffer_tokens), - static_cast(host_cache_locs), - static_cast(device_buffer_locs), - host_cache_k, - host_cache_v, - device_buffer_k, - device_buffer_v, - static_cast(top_k_device_locs), - static_cast(req_pool_indices), - static_cast(seq_lens), - static_cast(lru_slots), - static_cast(num_real_reqs), - buffer_stride_0, - host_stride, - lru_slot_stride_0, - top_k_tokens_stride, - top_k_device_locs_stride, - page_size, - item_size_bytes, - smem)); - }); -} - -// ============================================================================ -// C API for Python (ctypes) binding -// ============================================================================ -// -// The compile-time template config (block size, top-k, hot-buffer size, MLA / -// DSv4 flags) is fixed per module via -D macros, mirroring the CUDA JIT that -// bakes the same values into template arguments. The seq_lens / req_pool_indices -// dtype combination (i32/i64) is selected at call time by picking the matching -// exported symbol. - -#define _DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) \ - extern "C" void load_cache_to_device_buffer_##SEQ_SUFFIX##_##RPI_SUFFIX( \ - void* queue_ptr, \ - const void* top_k_tokens, \ - void* device_buffer_tokens, \ - const void* host_cache_locs, \ - const void* device_buffer_locs, \ - const void* host_cache_k, \ - const void* host_cache_v, \ - void* device_buffer_k, \ - void* device_buffer_v, \ - void* top_k_device_locs, \ - const void* req_pool_indices, \ - const void* seq_lens, \ - void* lru_slots, \ - const void* num_real_reqs, \ - int64_t batch_size, \ - int64_t buffer_stride_0, \ - int64_t host_stride, \ - int64_t lru_slot_stride_0, \ - int64_t top_k_tokens_stride, \ - int64_t top_k_device_locs_stride, \ - int64_t page_size, \ - int64_t item_size_bytes) { \ - auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ - load_cache_to_device_buffer_launcher< \ - SGL_HISPARSE_BLOCK_SIZE, \ - SGL_HISPARSE_NUM_TOP_K, \ - SGL_HISPARSE_HOT_BUFFER_SIZE, \ - (SGL_HISPARSE_IS_MLA != 0), \ - (SGL_HISPARSE_IS_DSV4 != 0), \ - SEQ_T, \ - RPI_T>( \ - queue, \ - top_k_tokens, \ - device_buffer_tokens, \ - host_cache_locs, \ - device_buffer_locs, \ - host_cache_k, \ - host_cache_v, \ - device_buffer_k, \ - device_buffer_v, \ - top_k_device_locs, \ - req_pool_indices, \ - seq_lens, \ - lru_slots, \ - num_real_reqs, \ - batch_size, \ - buffer_stride_0, \ - host_stride, \ - lru_slot_stride_0, \ - top_k_tokens_stride, \ - top_k_device_locs_stride, \ - page_size, \ - item_size_bytes); \ - } -#define DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) _DEFINE_LOAD_CACHE(SEQ_SUFFIX, SEQ_T, RPI_SUFFIX, RPI_T) - -#if defined(SGL_HISPARSE_BLOCK_SIZE) && defined(SGL_HISPARSE_NUM_TOP_K) && defined(SGL_HISPARSE_HOT_BUFFER_SIZE) && \ - defined(SGL_HISPARSE_IS_MLA) && defined(SGL_HISPARSE_IS_DSV4) -DEFINE_LOAD_CACHE(i64, int64_t, i64, int64_t) -DEFINE_LOAD_CACHE(i64, int64_t, i32, int32_t) -DEFINE_LOAD_CACHE(i32, int32_t, i64, int64_t) -DEFINE_LOAD_CACHE(i32, int32_t, i32, int32_t) -#endif - -#undef DEFINE_LOAD_CACHE -#undef _DEFINE_LOAD_CACHE - -} // namespace hisparse -} // namespace sycl_kernel -} // namespace sgl diff --git a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp b/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp deleted file mode 100644 index 9f8e0ad16..000000000 --- a/include/sgl_kernel/jit_kernel/hisparse/transfer_cache_dsv4_mla.hpp +++ /dev/null @@ -1,146 +0,0 @@ -/** - * HiSparse: transfer_cache_dsv4_mla SYCL kernel (Intel XPU). - * - * Ports transfer_cache_dsv4_mla_kernel from the CUDA source - * (sglang jit_kernel/csrc/hisparse.cuh). Bulk-copies DSv4-MLA C4 tokens between - * two sets of page-padded C4 buffers, one set per model layer. - * - * Mapping to the CUDA original: - * - CUDA "warp" (32 lanes) -> SYCL sub-group (pinned to kSubGroupSize). - * - One sub-group copies one item, iterating over all layers. - * - Grid-stride loop over items across all sub-groups. - * - * src_caches / dst_caches are device arrays of `num_layers` raw cache base - * pointers (uint64_t values), one per layer. - */ - -#pragma once - -#include -#include - -#include "c4_layout.hpp" - -namespace sgl { -namespace sycl_kernel { -namespace hisparse { - -// Sub-group width used for the cooperative item copy. Intel GPUs support 16/32; -// 32 mirrors the CUDA warp the kernel was written against. The strided copy in -// transfer_item is correct for any width, so this only affects granularity. -static constexpr int kSubGroupSize = 32; - -template -class TransferCacheDsv4MlaKernel { - public: - static_assert(BLOCK_SIZE % kSubGroupSize == 0, "BLOCK_SIZE must be a multiple of the sub-group size (32)."); - static constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; - - TransferCacheDsv4MlaKernel( - void** src_caches, - void** dst_caches, - const int64_t* src_indices, - const int64_t* dst_indices, - uint32_t num_items, - uint32_t num_layers, - uint32_t total_sub_groups) - : src_caches_(src_caches), - dst_caches_(dst_caches), - src_indices_(src_indices), - dst_indices_(dst_indices), - num_items_(num_items), - num_layers_(num_layers), - total_sub_groups_(total_sub_groups) {} - - [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { - const ::sycl::sub_group sg = item.get_sub_group(); - const int lane_id = static_cast(sg.get_local_linear_id()); - const int sg_size = static_cast(sg.get_max_local_range()[0]); - - // Global sub-group index: group * subgroups_per_group + local subgroup index. - const uint32_t global_sg = - static_cast(item.get_group(0)) * kNumSubGroups + static_cast(sg.get_group_linear_id()); - - for (uint32_t i = global_sg; i < num_items_; i += total_sub_groups_) { - const int32_t src_index = static_cast(src_indices_[i]); - const int32_t dst_index = static_cast(dst_indices_[i]); - for (uint32_t layer_id = 0; layer_id < num_layers_; ++layer_id) { - transfer_item(lane_id, sg_size, dst_caches_[layer_id], src_caches_[layer_id], dst_index, src_index); - } - } - } - - private: - void** src_caches_; - void** dst_caches_; - const int64_t* src_indices_; - const int64_t* dst_indices_; - uint32_t num_items_; - uint32_t num_layers_; - uint32_t total_sub_groups_; -}; - -template -void transfer_cache_dsv4_mla_launcher( - ::sycl::queue& queue, - void** src_caches, - void** dst_caches, - const int64_t* src_indices, - const int64_t* dst_indices, - uint32_t num_items, - uint32_t num_layers) { - if (num_items == 0) { - return; - } - constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; - const uint32_t num_groups = (num_items + kNumSubGroups - 1) / kNumSubGroups; - const uint32_t total_sub_groups = num_groups * kNumSubGroups; - - queue.submit([&](::sycl::handler& cgh) { - cgh.parallel_for( - ::sycl::nd_range<1>( - ::sycl::range<1>(static_cast(num_groups) * BLOCK_SIZE), ::sycl::range<1>(BLOCK_SIZE)), - TransferCacheDsv4MlaKernel( - src_caches, dst_caches, src_indices, dst_indices, num_items, num_layers, total_sub_groups)); - }); -} - -// ============================================================================ -// C API for Python (ctypes) binding -// ============================================================================ - -#define _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) \ - extern "C" void transfer_cache_dsv4_mla_##BLOCK_SIZE( \ - void* queue_ptr, \ - void* src_caches, \ - void* dst_caches, \ - const void* src_indices, \ - const void* dst_indices, \ - uint32_t num_items, \ - uint32_t num_layers) { \ - auto& queue = *static_cast<::sycl::queue*>(queue_ptr); \ - transfer_cache_dsv4_mla_launcher( \ - queue, \ - static_cast(src_caches), \ - static_cast(dst_caches), \ - static_cast(src_indices), \ - static_cast(dst_indices), \ - num_items, \ - num_layers); \ - } -#define DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) _DEFINE_TRANSFER_CACHE_DSV4_MLA(BLOCK_SIZE) - -#ifdef SGL_HISPARSE_BLOCK_SIZE -DEFINE_TRANSFER_CACHE_DSV4_MLA(SGL_HISPARSE_BLOCK_SIZE) -#else -DEFINE_TRANSFER_CACHE_DSV4_MLA(256) -DEFINE_TRANSFER_CACHE_DSV4_MLA(512) -DEFINE_TRANSFER_CACHE_DSV4_MLA(1024) -#endif - -#undef DEFINE_TRANSFER_CACHE_DSV4_MLA -#undef _DEFINE_TRANSFER_CACHE_DSV4_MLA - -} // namespace hisparse -} // namespace sycl_kernel -} // namespace sgl diff --git a/include/sgl_kernel_ops.h b/include/sgl_kernel_ops.h index 18000d455..7d3880f66 100644 --- a/include/sgl_kernel_ops.h +++ b/include/sgl_kernel_ops.h @@ -1258,4 +1258,33 @@ void causal_conv1d_update( const std::optional& conv_state_indices_, int64_t pad_slot_id); +/* + * HiSparse hierarchical sparse KV cache (DeepSeek DSA / V4) + */ +void transfer_cache_dsv4_mla( + const at::Tensor& src_ptrs, + const at::Tensor& dst_ptrs, + const at::Tensor& src_indices, + const at::Tensor& dst_indices, + int64_t block_size); + +void load_cache_to_device_buffer_mla( + const at::Tensor& top_k_tokens, + const at::Tensor& device_buffer_tokens, + const at::Tensor& host_cache_locs, + const at::Tensor& device_buffer_locs, + const at::Tensor& host_cache, + const at::Tensor& device_buffer, + const at::Tensor& top_k_device_locs, + const at::Tensor& req_pool_indices, + const at::Tensor& seq_lens, + const at::Tensor& lru_slots, + const std::optional& num_real_reqs, + int64_t item_size_bytes, + int64_t num_top_k, + int64_t hot_buffer_size, + int64_t page_size, + int64_t block_size, + bool is_dsv4_layout); + #pragma GCC visibility pop diff --git a/python/sgl_kernel/__init__.py b/python/sgl_kernel/__init__.py index 2ecbf45ee..03abb0e80 100644 --- a/python/sgl_kernel/__init__.py +++ b/python/sgl_kernel/__init__.py @@ -80,6 +80,11 @@ ) from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda from sgl_kernel.hadamard import hadamard_transform +from sgl_kernel.hisparse import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, +) from sgl_kernel.inkling_attn_prologue import ( compile_inkling_attn_prologue, inkling_attn_prologue_decode, diff --git a/python/sgl_kernel/hisparse.py b/python/sgl_kernel/hisparse.py new file mode 100644 index 000000000..765612ea2 --- /dev/null +++ b/python/sgl_kernel/hisparse.py @@ -0,0 +1,194 @@ +# HiSparse hierarchical sparse KV-cache ops for Intel XPU, used by DeepSeek +# DSA / V4 hierarchical sparse attention. +# Paged C4 layout (per page of 64 tokens): 64 value slots of 576 B, then 64 scale +# slots of 8 B, then padding so each page starts on a 576-byte boundary. + +from typing import Optional + +import torch + +# Work-group size for transfer_cache_dsv4_mla; only 256/512/1024 are compiled. +# All three measure within noise on BMG (Xe2), so this is an escape hatch rather +# than a tunable. +_DEFAULT_TRANSFER_BLOCK_SIZE = 1024 + +# Work-group size for the swap-in kernel: a plain runtime value, so any multiple +# of 32 works. +_DEFAULT_SWAP_IN_BLOCK_SIZE = 256 + + +def transfer_cache_dsv4_mla( + src_ptrs: torch.Tensor, + dst_ptrs: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + block_size: int = _DEFAULT_TRANSFER_BLOCK_SIZE, +) -> None: + """Transfer DSv4 C4 tokens between page-padded C4 buffers, all layers. + + Args: + src_ptrs: 1-D uint64 tensor of per-layer source cache base pointers. + dst_ptrs: 1-D uint64 tensor of per-layer destination cache base pointers. + src_indices: 1-D int64 tensor of source token slot indices. + dst_indices: 1-D int64 tensor of destination token slot indices. + block_size: work-group size; one of 256, 512, 1024. + """ + torch.ops.sgl_kernel.transfer_cache_dsv4_mla.default( + src_ptrs, + dst_ptrs, + src_indices, + dst_indices, + block_size, + ) + + +def _load_cache_to_device_buffer_mla( + is_dsv4_layout: bool, + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int, + block_size: int, + num_real_reqs: Optional[torch.Tensor], +) -> None: + torch.ops.sgl_kernel.load_cache_to_device_buffer_mla.default( + top_k_tokens, + device_buffer_tokens, + host_cache_locs, + device_buffer_locs, + host_cache, + device_buffer, + top_k_device_locs, + req_pool_indices, + seq_lens, + lru_slots, + num_real_reqs, + item_size_bytes, + num_top_k, + hot_buffer_size, + page_size, + block_size, + is_dsv4_layout, + ) + + +def load_cache_to_device_buffer_mla( + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int = 1, + block_size: int = _DEFAULT_SWAP_IN_BLOCK_SIZE, + num_real_reqs: Optional[torch.Tensor] = None, +) -> None: + """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes). + + Swaps each request's top-k tokens into a small hot device buffer, maintaining + the per-request LRU order, streaming misses in from the host cache and writing + every top-k token's device slot to ``top_k_device_locs``. + + Args: + top_k_tokens: (batch, num_top_k) int32 top-k token positions per request. + device_buffer_tokens: (num_reqs, hot_buffer_size + 1) int32 token position + resident in each device-buffer slot; updated in place. + host_cache_locs: (num_reqs, max_seq_len) int64 host cache slot per token. + device_buffer_locs: (num_reqs, hot_buffer_size + 1) int32 device cache slot + per buffer slot. Must share ``stride(0)`` with device_buffer_tokens. + host_cache: host-side KV cache tensor. + device_buffer: device-side hot KV buffer tensor. + top_k_device_locs: (batch, num_top_k) int32 output device slots. + req_pool_indices: (batch,) int32/int64 request-pool row per batch entry. + seq_lens: (batch,) int32/int64 sequence lengths. + lru_slots: (num_reqs, hot_buffer_size) int16 LRU order; updated in place. + item_size_bytes: bytes per KV item (one token, all heads). + num_top_k: top-k count; must be <= hot_buffer_size. + hot_buffer_size: device buffer capacity in tokens, excluding the extra + slot reserved for the newest token. + page_size: accepted for API parity; unused. + block_size: work-group size; a multiple of 32. + num_real_reqs: (1,) int32 count of non-padded requests, for graph-captured + padded batches. Defaults to the full batch. + """ + _load_cache_to_device_buffer_mla( + False, + top_k_tokens, + device_buffer_tokens, + host_cache_locs, + device_buffer_locs, + host_cache, + device_buffer, + top_k_device_locs, + req_pool_indices, + seq_lens, + lru_slots, + item_size_bytes, + num_top_k, + hot_buffer_size, + page_size, + block_size, + num_real_reqs, + ) + + +def load_cache_to_device_buffer_dsv4_mla( + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int = 1, + block_size: int = _DEFAULT_SWAP_IN_BLOCK_SIZE, + num_real_reqs: Optional[torch.Tensor] = None, +) -> None: + """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout. + + Same as :func:`load_cache_to_device_buffer_mla`, except the miss copy walks + the paged C4 layout described at the top of this module, so ``host_cache`` and + ``device_buffer`` must both be page-padded C4 buffers. + """ + _load_cache_to_device_buffer_mla( + True, + top_k_tokens, + device_buffer_tokens, + host_cache_locs, + device_buffer_locs, + host_cache, + device_buffer, + top_k_device_locs, + req_pool_indices, + seq_lens, + lru_slots, + item_size_bytes, + num_top_k, + hot_buffer_size, + page_size, + block_size, + num_real_reqs, + ) diff --git a/python/sgl_kernel/jit/__init__.py b/python/sgl_kernel/jit/__init__.py index 61261e57c..cb8e16a8b 100644 --- a/python/sgl_kernel/jit/__init__.py +++ b/python/sgl_kernel/jit/__init__.py @@ -26,13 +26,6 @@ def is_xpu() -> bool: is_icpx_available, load_jit_sycl, ) - - # hisparse lives under kvcache/ to mirror sglang's CUDA-side layout - from .kvcache.hisparse import ( - load_cache_to_device_buffer_dsv4_mla, - load_cache_to_device_buffer_mla, - transfer_cache_dsv4_mla, - ) from .moe_align_block_size import moe_align_block_size from .moe_fused_gate import can_use_moe_fused_gate, moe_fused_gate from .moe_topk_sigmoid import topk_sigmoid @@ -67,9 +60,6 @@ def is_xpu() -> bool: "apply_rope_inplace", "apply_rope_inplace_with_kvcache", "timestep_embedding", - "transfer_cache_dsv4_mla", - "load_cache_to_device_buffer_mla", - "load_cache_to_device_buffer_dsv4_mla", ] else: # Non-XPU environment - provide stubs diff --git a/python/sgl_kernel/jit/kvcache/__init__.py b/python/sgl_kernel/jit/kvcache/__init__.py deleted file mode 100644 index b2518eb5f..000000000 --- a/python/sgl_kernel/jit/kvcache/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -"""KV-cache-family JIT kernels for Intel XPU. - -Groups SYCL JIT wrappers whose CUDA-side counterparts live under -``sglang.kernels.ops.kvcache.*``. Mirrors that layout so consumers can write -the same import shape across backends: - - from sglang.kernels.ops.kvcache.hisparse import load_cache_to_device_buffer_mla # CUDA - from sgl_kernel.jit.kvcache.hisparse import load_cache_to_device_buffer_mla # XPU - -Modules exported here: - -- ``hisparse``: DSA/DSv4 hisparse swap-in + evict/backup kernels - (``load_cache_to_device_buffer_{mla,dsv4_mla}``, ``transfer_cache_dsv4_mla``). -""" - -from .hisparse import ( - load_cache_to_device_buffer_dsv4_mla, - load_cache_to_device_buffer_mla, - transfer_cache_dsv4_mla, -) - -__all__ = [ - "load_cache_to_device_buffer_dsv4_mla", - "load_cache_to_device_buffer_mla", - "transfer_cache_dsv4_mla", -] diff --git a/python/sgl_kernel/jit/kvcache/hisparse.py b/python/sgl_kernel/jit/kvcache/hisparse.py deleted file mode 100644 index 839b5c96b..000000000 --- a/python/sgl_kernel/jit/kvcache/hisparse.py +++ /dev/null @@ -1,354 +0,0 @@ -""" -XPU/SYCL HiSparse KV-offload swap-in kernel wrappers. - -Provides JIT-compiled SYCL ports of the CUDA HiSparse kernels used for -hierarchical sparse attention (DeepSeek DSA / V4). Two kernels are exposed: - -- ``transfer_cache_dsv4_mla``: bulk-copy DSv4 C4 tokens between page-padded C4 - buffers, one set of buffers per model layer (evict / backup path). -- ``load_cache_to_device_buffer_mla`` / ``..._dsv4_mla``: per-request swap-in of - the current top-k tokens into a small hot device buffer, maintaining an LRU - ordering and streaming misses in from the host cache. - -These mirror the API of ``sglang.jit_kernel.hisparse`` (the CUDA path). -""" - -from __future__ import annotations - -import ctypes - -import torch -from sgl_kernel.jit.compiler import load_jit_sycl -from sgl_kernel.jit.utils import cache_once - -# --------------------------------------------------------------------------- -# transfer_cache_dsv4_mla -# --------------------------------------------------------------------------- - - -@cache_once -def _jit_transfer_cache_dsv4_mla_module(block_size: int): - """Compile/load the DSv4 C4 transfer module for a given block size.""" - if block_size % 32 != 0: - raise ValueError(f"block_size must be a multiple of 32, got {block_size}") - return load_jit_sycl( - "hisparse_transfer_cache_dsv4_mla", - str(block_size), - sycl_files=["hisparse/transfer_cache_dsv4_mla.hpp"], - extra_sycl_cflags=[f"-DSGL_HISPARSE_BLOCK_SIZE={block_size}"], - ) - - -_TRANSFER_ARGTYPES = [ - ctypes.c_void_p, # queue - ctypes.c_void_p, # src_caches (void**) - ctypes.c_void_p, # dst_caches (void**) - ctypes.c_void_p, # src_indices (const int64_t*) - ctypes.c_void_p, # dst_indices (const int64_t*) - ctypes.c_uint32, # num_items - ctypes.c_uint32, # num_layers -] - - -def transfer_cache_dsv4_mla( - src_ptrs: torch.Tensor, - dst_ptrs: torch.Tensor, - src_indices: torch.Tensor, - dst_indices: torch.Tensor, - block_size: int = 1024, -) -> None: - """Transfer DSv4 C4 tokens between page-padded C4 buffers. - - Args: - src_ptrs: 1-D uint64 tensor of per-layer source cache base pointers. - dst_ptrs: 1-D uint64 tensor of per-layer destination cache base pointers. - src_indices: 1-D int64 tensor of source token slot indices. - dst_indices: 1-D int64 tensor of destination token slot indices. - block_size: SYCL work-group size (multiple of 32). - """ - assert src_ptrs.dtype == torch.uint64 and dst_ptrs.dtype == torch.uint64 - assert src_indices.dtype == torch.int64 and dst_indices.dtype == torch.int64 - assert src_ptrs.numel() == dst_ptrs.numel() - assert src_indices.numel() == dst_indices.numel() - - num_items = src_indices.numel() - if num_items == 0: - return - num_layers = src_ptrs.numel() - - module = _jit_transfer_cache_dsv4_mla_module(block_size) - func = module.get_function( - f"transfer_cache_dsv4_mla_{block_size}", _TRANSFER_ARGTYPES - ) - queue = torch.xpu.current_stream().sycl_queue - func( - queue, - src_ptrs.data_ptr(), - dst_ptrs.data_ptr(), - src_indices.data_ptr(), - dst_indices.data_ptr(), - num_items, - num_layers, - ) - - -# --------------------------------------------------------------------------- -# load_cache_to_device_buffer -# --------------------------------------------------------------------------- - -_LOAD_CACHE_ARGTYPES = [ - ctypes.c_void_p, # queue - ctypes.c_void_p, # top_k_tokens (const int32_t*) - ctypes.c_void_p, # device_buffer_tokens (int32_t*) - ctypes.c_void_p, # host_cache_locs (const int64_t*) - ctypes.c_void_p, # device_buffer_locs (const int32_t*) - ctypes.c_void_p, # host_cache_k - ctypes.c_void_p, # host_cache_v - ctypes.c_void_p, # device_buffer_k - ctypes.c_void_p, # device_buffer_v - ctypes.c_void_p, # top_k_device_locs (int32_t*) - ctypes.c_void_p, # req_pool_indices - ctypes.c_void_p, # seq_lens - ctypes.c_void_p, # lru_slots (int16_t*) - ctypes.c_void_p, # num_real_reqs (const int32_t*) - ctypes.c_int64, # batch_size - ctypes.c_int64, # buffer_stride_0 - ctypes.c_int64, # host_stride - ctypes.c_int64, # lru_slot_stride_0 - ctypes.c_int64, # top_k_tokens_stride - ctypes.c_int64, # top_k_device_locs_stride - ctypes.c_int64, # page_size - ctypes.c_int64, # item_size_bytes -] - - -@cache_once -def _jit_load_cache_module( - block_size: int, - num_top_k: int, - hot_buffer_size: int, - is_mla: bool, - is_dsv4_layout: bool, -): - """Compile/load the swap-in module for a fixed template configuration.""" - if block_size % 32 != 0: - raise ValueError(f"block_size must be a multiple of 32, got {block_size}") - if hot_buffer_size < num_top_k: - raise ValueError( - f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" - ) - return load_jit_sycl( - "hisparse_load_cache_to_device_buffer", - str(block_size), - str(num_top_k), - str(hot_buffer_size), - "mla" if is_mla else "gqa", - "dsv4" if is_dsv4_layout else "linear", - sycl_files=["hisparse/load_cache_to_device_buffer.hpp"], - extra_sycl_cflags=[ - f"-DSGL_HISPARSE_BLOCK_SIZE={block_size}", - f"-DSGL_HISPARSE_NUM_TOP_K={num_top_k}", - f"-DSGL_HISPARSE_HOT_BUFFER_SIZE={hot_buffer_size}", - f"-DSGL_HISPARSE_IS_MLA={1 if is_mla else 0}", - f"-DSGL_HISPARSE_IS_DSV4={1 if is_dsv4_layout else 0}", - ], - ) - - -def _dtype_suffix(t: torch.Tensor) -> str: - if t.dtype == torch.int64: - return "i64" - if t.dtype == torch.int32: - return "i32" - raise TypeError( - f"seq_lens / req_pool_indices must be int32 or int64, got {t.dtype}" - ) - - -def _load_cache_to_device_buffer_mla( - *, - is_dsv4_layout: bool, - top_k_tokens: torch.Tensor, - device_buffer_tokens: torch.Tensor, - host_cache_locs: torch.Tensor, - device_buffer_locs: torch.Tensor, - host_cache: torch.Tensor, - device_buffer: torch.Tensor, - top_k_device_locs: torch.Tensor, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - lru_slots: torch.Tensor, - item_size_bytes: int, - num_top_k: int, - hot_buffer_size: int, - page_size: int, - block_size: int, - num_real_reqs: torch.Tensor | None, -) -> None: - assert ( - hot_buffer_size >= num_top_k - ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" - - # Kernel gets only stride(0); rows must be contiguous. device_buffer_tokens - # and device_buffer_locs share stride(0) (reused inside the kernel). - _expected_dtypes = { - "top_k_tokens": (top_k_tokens, torch.int32), - "device_buffer_tokens": (device_buffer_tokens, torch.int32), - "host_cache_locs": (host_cache_locs, torch.int64), - "device_buffer_locs": (device_buffer_locs, torch.int32), - "top_k_device_locs": (top_k_device_locs, torch.int32), - "lru_slots": (lru_slots, torch.int16), - } - for name, (t, expected) in _expected_dtypes.items(): - if t.dtype != expected: - raise TypeError(f"{name} must be {expected}, got {t.dtype}") - if t.device.type != "xpu": - raise ValueError(f"{name} must be on XPU, got {t.device}") - if t.dim() >= 2 and t.stride(-1) != 1: - raise ValueError( - f"{name} must be row-contiguous (stride(-1)==1), got stride={t.stride()}" - ) - if device_buffer_tokens.stride(0) != device_buffer_locs.stride(0): - raise ValueError( - "device_buffer_tokens and device_buffer_locs must share stride(0), " - f"got {device_buffer_tokens.stride(0)} vs {device_buffer_locs.stride(0)}" - ) - - module = _jit_load_cache_module( - block_size, - num_top_k, - hot_buffer_size, - True, # is_mla - is_dsv4_layout, - ) - - if num_real_reqs is None: - num_real_reqs = torch.tensor( - [top_k_tokens.size(0)], dtype=torch.int32, device=top_k_tokens.device - ) - - batch_size = top_k_tokens.size(0) - host_stride = host_cache_locs.size(1) - buffer_stride_0 = device_buffer_tokens.stride(0) - lru_slot_stride_0 = lru_slots.stride(0) - top_k_tokens_stride = top_k_tokens.stride(0) - top_k_device_locs_stride = top_k_device_locs.stride(0) - - func_name = ( - f"load_cache_to_device_buffer_" - f"{_dtype_suffix(seq_lens)}_{_dtype_suffix(req_pool_indices)}" - ) - func = module.get_function(func_name, _LOAD_CACHE_ARGTYPES) - queue = torch.xpu.current_stream().sycl_queue - - func( - queue, - top_k_tokens.data_ptr(), - device_buffer_tokens.data_ptr(), - host_cache_locs.data_ptr(), - device_buffer_locs.data_ptr(), - host_cache.data_ptr(), - 0, # host_cache_v (MLA: unused) - device_buffer.data_ptr(), - 0, # device_buffer_v (MLA: unused) - top_k_device_locs.data_ptr(), - req_pool_indices.data_ptr(), - seq_lens.data_ptr(), - lru_slots.data_ptr(), - num_real_reqs.data_ptr(), - batch_size, - buffer_stride_0, - host_stride, - lru_slot_stride_0, - top_k_tokens_stride, - top_k_device_locs_stride, - page_size, - item_size_bytes, - ) - - -def load_cache_to_device_buffer_mla( - top_k_tokens: torch.Tensor, - device_buffer_tokens: torch.Tensor, - host_cache_locs: torch.Tensor, - device_buffer_locs: torch.Tensor, - host_cache: torch.Tensor, - device_buffer: torch.Tensor, - top_k_device_locs: torch.Tensor, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - lru_slots: torch.Tensor, - item_size_bytes: int, - num_top_k: int, - hot_buffer_size: int, - page_size: int = 1, - block_size: int = 256, - num_real_reqs: torch.Tensor | None = None, -) -> None: - """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes).""" - _load_cache_to_device_buffer_mla( - is_dsv4_layout=False, - top_k_tokens=top_k_tokens, - device_buffer_tokens=device_buffer_tokens, - host_cache_locs=host_cache_locs, - device_buffer_locs=device_buffer_locs, - host_cache=host_cache, - device_buffer=device_buffer, - top_k_device_locs=top_k_device_locs, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens, - lru_slots=lru_slots, - item_size_bytes=item_size_bytes, - num_top_k=num_top_k, - hot_buffer_size=hot_buffer_size, - page_size=page_size, - block_size=block_size, - num_real_reqs=num_real_reqs, - ) - - -def load_cache_to_device_buffer_dsv4_mla( - top_k_tokens: torch.Tensor, - device_buffer_tokens: torch.Tensor, - host_cache_locs: torch.Tensor, - device_buffer_locs: torch.Tensor, - host_cache: torch.Tensor, - device_buffer: torch.Tensor, - top_k_device_locs: torch.Tensor, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - lru_slots: torch.Tensor, - item_size_bytes: int, - num_top_k: int, - hot_buffer_size: int, - page_size: int = 1, - block_size: int = 256, - num_real_reqs: torch.Tensor | None = None, -) -> None: - """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout.""" - _load_cache_to_device_buffer_mla( - is_dsv4_layout=True, - top_k_tokens=top_k_tokens, - device_buffer_tokens=device_buffer_tokens, - host_cache_locs=host_cache_locs, - device_buffer_locs=device_buffer_locs, - host_cache=host_cache, - device_buffer=device_buffer, - top_k_device_locs=top_k_device_locs, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens, - lru_slots=lru_slots, - item_size_bytes=item_size_bytes, - num_top_k=num_top_k, - hot_buffer_size=hot_buffer_size, - page_size=page_size, - block_size=block_size, - num_real_reqs=num_real_reqs, - ) - - -__all__ = [ - "transfer_cache_dsv4_mla", - "load_cache_to_device_buffer_mla", - "load_cache_to_device_buffer_dsv4_mla", -] diff --git a/src/sycl/HiSparse.cpp b/src/sycl/HiSparse.cpp new file mode 100644 index 000000000..e2ee88536 --- /dev/null +++ b/src/sycl/HiSparse.cpp @@ -0,0 +1,314 @@ +#include + +#include + +#include "Utils.h" +#include "comm/General.h" +#include "sgl_kernel/hisparse/load_cache_to_device_buffer.hpp" +#include "sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp" +#include "sgl_kernel_export.h" + +using namespace sgl::sycl_kernel::hisparse; + +namespace { + +// --------------------------------------------------------------------------- +// transfer_cache_dsv4_mla +// --------------------------------------------------------------------------- + +template +void launch_transfer_cache_dsv4_mla( + void** src_caches, + void** dst_caches, + const int64_t* src_indices, + const int64_t* dst_indices, + uint32_t num_items, + uint32_t num_layers) { + constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + const uint32_t num_groups = div_up(num_items, static_cast(kNumSubGroups)); + const uint32_t total_sub_groups = num_groups * kNumSubGroups; + + TransferCacheDsv4MlaKernel kernel{ + .src_caches_ = src_caches, + .dst_caches_ = dst_caches, + .src_indices_ = src_indices, + .dst_indices_ = dst_indices, + .num_items_ = num_items, + .num_layers_ = num_layers, + .total_sub_groups_ = total_sub_groups, + }; + + auto cgf = DPCPP_Q_CGF(cgh) { + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>(static_cast(num_groups) * BLOCK_SIZE), sycl::range<1>(BLOCK_SIZE)), + kernel); + }; + dpcppGetCurrentQueue().submit(cgf); +} + +// Validate a uint64 pointer table the same way KVCacheIO.cpp does. +void check_ptr_table(const at::Tensor& tbl, int64_t num_layers, const char* name) { + TORCH_CHECK(tbl.scalar_type() == at::kUInt64, name, " must be a uint64 pointer table"); + TORCH_CHECK(tbl.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tbl.numel() == num_layers, name, " must have num_layers entries, got ", tbl.numel()); +} + +// --------------------------------------------------------------------------- +// load_cache_to_device_buffer +// --------------------------------------------------------------------------- + +struct LoadCacheArgs { + const int32_t* top_k_tokens; + int32_t* device_buffer_tokens; + const int64_t* host_cache_locs; + const int32_t* device_buffer_locs; + const void* host_cache_k; + const void* host_cache_v; + void* device_buffer_k; + void* device_buffer_v; + int32_t* top_k_device_locs; + const void* req_pool_indices; + const void* seq_lens; + int16_t* lru_slots; + const int32_t* num_real_reqs; + bool req_pool_indices_is_i64; + bool seq_lens_is_i64; + int64_t buffer_stride_0; + int64_t host_stride; + int64_t lru_slot_stride_0; + int64_t top_k_tokens_stride; + int64_t top_k_device_locs_stride; + int64_t item_size_bytes; + int64_t batch_size; + int block_size; + int num_top_k; + int hot_buffer_size; +}; + +template +void launch_load_cache_to_device_buffer(const LoadCacheArgs& a, const SmemLayout& layout) { + const int num_sub_groups = a.block_size / kSubGroupSize; + + auto cgf = DPCPP_Q_CGF(cgh) { + sycl::local_accessor smem(sycl::range<1>(static_cast(layout.total_int32_slots)), cgh); + LoadCacheToDeviceBufferKernel kernel{ + .top_k_tokens_ = a.top_k_tokens, + .device_buffer_tokens_ = a.device_buffer_tokens, + .host_cache_locs_ = a.host_cache_locs, + .device_buffer_locs_ = a.device_buffer_locs, + .host_cache_k_ = a.host_cache_k, + .host_cache_v_ = a.host_cache_v, + .device_buffer_k_ = a.device_buffer_k, + .device_buffer_v_ = a.device_buffer_v, + .top_k_device_locs_ = a.top_k_device_locs, + .req_pool_indices_ = a.req_pool_indices, + .seq_lens_ = a.seq_lens, + .lru_slots_ = a.lru_slots, + .num_real_reqs_ = a.num_real_reqs, + .req_pool_indices_is_i64_ = a.req_pool_indices_is_i64, + .seq_lens_is_i64_ = a.seq_lens_is_i64, + .buffer_stride_0_ = a.buffer_stride_0, + .host_stride_ = a.host_stride, + .lru_slot_stride_0_ = a.lru_slot_stride_0, + .top_k_tokens_stride_ = a.top_k_tokens_stride, + .top_k_device_locs_stride_ = a.top_k_device_locs_stride, + .item_size_bytes_ = a.item_size_bytes, + .block_size_ = a.block_size, + .num_sub_groups_ = num_sub_groups, + .num_top_k_ = a.num_top_k, + .hot_buffer_size_ = a.hot_buffer_size, + .hash_size_ = layout.hash_size, + .hash_mask_ = layout.hash_mask, + .num_buffer_chunks_ = layout.num_buffer_chunks, + .num_token_chunks_ = layout.num_token_chunks, + .iters_per_sg_buffer_ = div_up(layout.num_buffer_chunks, num_sub_groups), + .iters_per_sg_token_ = div_up(layout.num_token_chunks, num_sub_groups), + .total_int32_ = layout.total_int32, + .smem_ = smem, + }; + cgh.parallel_for( + sycl::nd_range<1>( + sycl::range<1>(static_cast(a.batch_size) * a.block_size), + sycl::range<1>(static_cast(a.block_size))), + kernel); + }; + dpcppGetCurrentQueue().submit(cgf); +} + +void check_swap_in_tensor(const at::Tensor& t, at::ScalarType expected, const char* name) { + TORCH_CHECK(t.scalar_type() == expected, name, " must be ", expected, ", got ", t.scalar_type()); + TORCH_CHECK(t.device().is_xpu(), name, " must be on XPU, got ", t.device()); + // The kernel is given only stride(0); every row must be contiguous. + TORCH_CHECK(t.dim() < 2 || t.stride(-1) == 1, name, " must be row-contiguous (stride(-1) == 1)"); +} + +} // namespace + +SGL_KERNEL_EXPORT void transfer_cache_dsv4_mla( + const at::Tensor& src_ptrs, + const at::Tensor& dst_ptrs, + const at::Tensor& src_indices, + const at::Tensor& dst_indices, + int64_t block_size) { + TORCH_CHECK(src_indices.scalar_type() == at::kLong, "src_indices must be int64"); + TORCH_CHECK(dst_indices.scalar_type() == at::kLong, "dst_indices must be int64"); + TORCH_CHECK(src_indices.is_contiguous(), "src_indices must be contiguous"); + TORCH_CHECK(dst_indices.is_contiguous(), "dst_indices must be contiguous"); + TORCH_CHECK( + src_indices.numel() == dst_indices.numel(), + "index count mismatch: ", + src_indices.numel(), + " vs ", + dst_indices.numel()); + + const int64_t num_layers = src_ptrs.numel(); + check_ptr_table(src_ptrs, num_layers, "src_ptrs"); + check_ptr_table(dst_ptrs, num_layers, "dst_ptrs"); + + const int64_t num_items = src_indices.numel(); + if (num_items == 0 || num_layers == 0) return; // nothing to transfer + + auto** src_caches = reinterpret_cast(src_ptrs.data_ptr()); + auto** dst_caches = reinterpret_cast(dst_ptrs.data_ptr()); + const auto* src_idx = src_indices.data_ptr(); + const auto* dst_idx = dst_indices.data_ptr(); + const auto items = static_cast(num_items); + const auto layers = static_cast(num_layers); + + // block_size is a template parameter; 1024 is the default, the rest are + // escape hatches (all three measure within noise on Xe2). + switch (block_size) { + case 256: + launch_transfer_cache_dsv4_mla<256>(src_caches, dst_caches, src_idx, dst_idx, items, layers); + break; + case 512: + launch_transfer_cache_dsv4_mla<512>(src_caches, dst_caches, src_idx, dst_idx, items, layers); + break; + case 1024: + launch_transfer_cache_dsv4_mla<1024>(src_caches, dst_caches, src_idx, dst_idx, items, layers); + break; + default: + TORCH_CHECK(false, "block_size must be one of 256, 512, 1024, got ", block_size); + } +} + +SGL_KERNEL_EXPORT void load_cache_to_device_buffer_mla( + const at::Tensor& top_k_tokens, + const at::Tensor& device_buffer_tokens, + const at::Tensor& host_cache_locs, + const at::Tensor& device_buffer_locs, + const at::Tensor& host_cache, + const at::Tensor& device_buffer, + const at::Tensor& top_k_device_locs, + const at::Tensor& req_pool_indices, + const at::Tensor& seq_lens, + const at::Tensor& lru_slots, + const std::optional& num_real_reqs, + int64_t item_size_bytes, + int64_t num_top_k, + int64_t hot_buffer_size, + int64_t page_size, + int64_t block_size, + bool is_dsv4_layout) { + TORCH_CHECK(num_top_k > 0, "num_top_k must be positive, got ", num_top_k); + TORCH_CHECK( + hot_buffer_size >= num_top_k, "hot_buffer_size (", hot_buffer_size, ") must be >= num_top_k (", num_top_k, ")"); + TORCH_CHECK( + block_size > 0 && block_size % kSubGroupSize == 0, + "block_size must be a positive multiple of ", + kSubGroupSize, + ", got ", + block_size); + TORCH_CHECK(item_size_bytes > 0, "item_size_bytes must be positive, got ", item_size_bytes); + // int16_t slot indices are stored in the LRU array and the hash values. + TORCH_CHECK(hot_buffer_size < 32767, "hot_buffer_size must fit in int16, got ", hot_buffer_size); + + check_swap_in_tensor(top_k_tokens, at::kInt, "top_k_tokens"); + check_swap_in_tensor(device_buffer_tokens, at::kInt, "device_buffer_tokens"); + check_swap_in_tensor(host_cache_locs, at::kLong, "host_cache_locs"); + check_swap_in_tensor(device_buffer_locs, at::kInt, "device_buffer_locs"); + check_swap_in_tensor(top_k_device_locs, at::kInt, "top_k_device_locs"); + check_swap_in_tensor(lru_slots, at::kShort, "lru_slots"); + TORCH_CHECK( + req_pool_indices.scalar_type() == at::kInt || req_pool_indices.scalar_type() == at::kLong, + "req_pool_indices must be int32 or int64, got ", + req_pool_indices.scalar_type()); + TORCH_CHECK( + seq_lens.scalar_type() == at::kInt || seq_lens.scalar_type() == at::kLong, + "seq_lens must be int32 or int64, got ", + seq_lens.scalar_type()); + TORCH_CHECK(req_pool_indices.is_contiguous(), "req_pool_indices must be contiguous"); + TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); + TORCH_CHECK(host_cache_locs.dim() >= 2, "host_cache_locs must be at least 2-D"); + + TORCH_CHECK( + device_buffer_tokens.stride(0) == device_buffer_locs.stride(0), + "device_buffer_tokens and device_buffer_locs must share stride(0), got ", + device_buffer_tokens.stride(0), + " vs ", + device_buffer_locs.stride(0)); + + const int64_t batch_size = top_k_tokens.size(0); + TORCH_CHECK( + req_pool_indices.numel() >= batch_size && seq_lens.numel() >= batch_size, + "req_pool_indices / seq_lens must cover the batch (", + batch_size, + "), got ", + req_pool_indices.numel(), + " / ", + seq_lens.numel()); + if (batch_size == 0) return; + + const SmemLayout layout = SmemLayout::make(static_cast(num_top_k), static_cast(hot_buffer_size)); + const size_t local_mem_size = dpcppGetCurrentQueue().get_device().get_info(); + TORCH_CHECK( + layout.bytes() <= local_mem_size, + "hisparse swap-in needs ", + layout.bytes(), + " bytes of shared local memory for num_top_k=", + num_top_k, + " hot_buffer_size=", + hot_buffer_size, + ", but the device provides only ", + local_mem_size); + + at::Tensor real_reqs = + num_real_reqs.has_value() ? *num_real_reqs : at::full({1}, batch_size, top_k_tokens.options().dtype(at::kInt)); + TORCH_CHECK(real_reqs.scalar_type() == at::kInt, "num_real_reqs must be int32, got ", real_reqs.scalar_type()); + TORCH_CHECK(real_reqs.device().is_xpu(), "num_real_reqs must be on XPU, got ", real_reqs.device()); + + LoadCacheArgs args{ + .top_k_tokens = top_k_tokens.data_ptr(), + .device_buffer_tokens = device_buffer_tokens.data_ptr(), + .host_cache_locs = host_cache_locs.data_ptr(), + .device_buffer_locs = device_buffer_locs.data_ptr(), + .host_cache_k = host_cache.data_ptr(), + .host_cache_v = nullptr, // MLA: K-only + .device_buffer_k = device_buffer.data_ptr(), + .device_buffer_v = nullptr, // MLA: K-only + .top_k_device_locs = top_k_device_locs.data_ptr(), + .req_pool_indices = req_pool_indices.data_ptr(), + .seq_lens = seq_lens.data_ptr(), + .lru_slots = lru_slots.data_ptr(), + .num_real_reqs = real_reqs.data_ptr(), + .req_pool_indices_is_i64 = req_pool_indices.scalar_type() == at::kLong, + .seq_lens_is_i64 = seq_lens.scalar_type() == at::kLong, + .buffer_stride_0 = device_buffer_tokens.stride(0), + .host_stride = host_cache_locs.size(1), + .lru_slot_stride_0 = lru_slots.stride(0), + .top_k_tokens_stride = top_k_tokens.stride(0), + .top_k_device_locs_stride = top_k_device_locs.stride(0), + .item_size_bytes = item_size_bytes, + .batch_size = batch_size, + .block_size = static_cast(block_size), + .num_top_k = static_cast(num_top_k), + .hot_buffer_size = static_cast(hot_buffer_size), + }; + (void)page_size; + + if (is_dsv4_layout) { + launch_load_cache_to_device_buffer(args, layout); + } else { + launch_load_cache_to_device_buffer(args, layout); + } +} diff --git a/src/torch_extension_sycl.cc b/src/torch_extension_sycl.cc index a87d03388..e0908a5a3 100644 --- a/src/torch_extension_sycl.cc +++ b/src/torch_extension_sycl.cc @@ -596,6 +596,23 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "Tensor out_loc, Tensor! kvcache, bool is_decode, int compress_ratio, int page_size, bool use_fp4, " "int preshuffle_size=0, bool use_bf16_store=False) -> ()"); m.impl("fused_norm_rope_store", torch::kXPU, &at::native::xpu::fused_norm_rope_store); + + /* + * HiSparse hierarchical sparse KV cache kernels + */ + m.def( + "transfer_cache_dsv4_mla(Tensor src_ptrs, Tensor(a!) dst_ptrs, " + "Tensor src_indices, Tensor dst_indices, int block_size) -> ()"); + m.impl("transfer_cache_dsv4_mla", torch::kXPU, &transfer_cache_dsv4_mla); + + m.def( + "load_cache_to_device_buffer_mla(Tensor top_k_tokens, Tensor(a!) device_buffer_tokens, " + "Tensor host_cache_locs, Tensor device_buffer_locs, Tensor host_cache, " + "Tensor(b!) device_buffer, Tensor(c!) top_k_device_locs, Tensor req_pool_indices, " + "Tensor seq_lens, Tensor(d!) lru_slots, Tensor? num_real_reqs, int item_size_bytes, " + "int num_top_k, int hot_buffer_size, int page_size, int block_size, " + "bool is_dsv4_layout) -> ()"); + m.impl("load_cache_to_device_buffer_mla", torch::kXPU, &load_cache_to_device_buffer_mla); } REGISTER_EXTENSION(common_ops) diff --git a/tests/run_suite.py b/tests/run_suite.py index 75a515f94..9c0b1ff57 100644 --- a/tests/run_suite.py +++ b/tests/run_suite.py @@ -75,6 +75,7 @@ class TestFile: TestFile("test_sconv_metadata_and_windows.py"), TestFile("test_sconv_update_sconv_cache.py"), TestFile("test_inkling_attn_prologue.py"), + TestFile("test_hisparse.py"), ], # Nightly suite: exercises the wheel installed in the intel/sgl-kernel-xpu-dev # nightly image. Populate with longer-running or full-shape tests that are diff --git a/tests/test_hisparse_jit.py b/tests/test_hisparse.py similarity index 74% rename from tests/test_hisparse_jit.py rename to tests/test_hisparse.py index 96b6b86a7..88b17c70e 100755 --- a/tests/test_hisparse_jit.py +++ b/tests/test_hisparse.py @@ -1,39 +1,31 @@ """ Accuracy tests for the XPU/SYCL HiSparse swap-in kernels. -Ported from the CUDA oracle (sglang test/registered/jit/test_hisparse.py). The -SYCL kernels pin the logical warp to a 32-lane sub-group, so the slot<->lane -mapping and eviction ordering match the CUDA kernel bit-for-bit; the expected -values below are therefore identical to the CUDA reference. +The kernels pin the sub-group to 32 lanes, so the slot<->lane mapping and +eviction ordering are fully determined and the expected values below are exact. Guarded failure modes (derived-property + bug-regression): - LRU hit/evict compaction and MRU/LRU write-back ordering. - Miss classification, evict-slot reuse, and host->device miss copy. - - Fast-path (seq_len <= hot_buffer) short-circuit leaves state untouched. - - CUDA-graph padding (num_real_reqs) leaves padded request rows untouched. + - Multi-iteration chunk scan (hot_buffer_size past one sub-group window). + - Fast-path (seq_len <= hot_buffer) short-circuit leaves cache state untouched + while still writing -1 to every output slot it cannot resolve. + - Graph-capture padding (num_real_reqs) leaves padded requests' cache state + untouched and writes -1 across their output rows. - DSv4 page-padded C4 addressing on both the transfer and swap-in paths. """ import pytest import torch +from sgl_kernel import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, +) -HAS_XPU = hasattr(torch, "xpu") and torch.xpu.is_available() +HAS_XPU = torch.xpu.is_available() -try: - from sgl_kernel.jit import ( - load_cache_to_device_buffer_dsv4_mla, - load_cache_to_device_buffer_mla, - transfer_cache_dsv4_mla, - ) - - HAS_SGL_JIT = True -except ImportError: - HAS_SGL_JIT = False - -pytestmark = [ - pytest.mark.skipif(not HAS_XPU, reason="Requires XPU device"), - pytest.mark.skipif(not HAS_SGL_JIT, reason="Requires sgl_kernel JIT HiSparse"), -] +pytestmark = pytest.mark.skipif(not HAS_XPU, reason="Requires XPU device") DEVICE = "xpu" DTYPE = torch.float32 @@ -116,6 +108,7 @@ def _run_kernel( seq_lens_dtype: torch.dtype = torch.int32, req_pool_indices: torch.Tensor | None = None, num_real_reqs: int | None = None, + output_fill_value: int = -1, ) -> torch.Tensor: batch_size = top_k_tokens.shape[0] if req_pool_indices is None: @@ -127,7 +120,9 @@ def _run_kernel( if num_real_reqs is None: num_real_reqs = batch_size - out = torch.full_like(top_k_tokens, -1) + # Cases that assert the kernel *writes* -1 (rather than merely leaving a slot + # alone) pass output_fill_value, so the default -1 cannot mask a skipped write. + out = torch.full_like(top_k_tokens, output_fill_value) load_cache_to_device_buffer_mla( top_k_tokens=top_k_tokens, device_buffer_tokens=device_buffer_tokens, @@ -319,6 +314,23 @@ def test_load_cache_to_device_buffer_fast_path(seq_lens_dtype: torch.dtype) -> N assert torch.equal(device_buffer.cpu(), device_buffer_before.cpu()) +def test_load_cache_to_device_buffer_fast_path_overwrites_stale_output() -> None: + # The fast path must write every one of the num_top_k output slots, not just + # the first `count = min(seq_len, num_top_k)`: slots past seq_len, and slots + # whose token position is negative, have to be set to -1 rather than left + # holding whatever the caller's buffer contained. + state = _make_state([[9, 7, 3, 5, 11]], [[0, 1, 2, 3, -1]], [4]) + + out = _run_kernel( + top_k_tokens=torch.tensor([[1, -1, 0, 0]], dtype=torch.int32, device=DEVICE), + seq_len=2, + output_fill_value=123456, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[7, -1, -1, -1]], dtype=torch.int32)) + + def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None: state = _long_case() @@ -399,6 +411,98 @@ def test_load_cache_to_device_buffer_multiple_misses_copy_all_slots() -> None: ) +def test_load_cache_to_device_buffer_multi_iteration_scan_compacts_evictables() -> None: + # Regression for the multi-iteration prefix scan over the hot-buffer chunks. + # + # If the classification loop scans s_evict_chunk_offset with the full + # num_buffer_chunks + 1 element count, all 32 lanes participate even though + # only num_sub_groups entries were written this iteration: the lanes past that + # window re-read what an earlier iteration's scan left behind and fold it into + # the accumulator, so every later iteration compacts evictable slots to the + # wrong positions in s_lru_slots_out. + # + # Reproducing it needs a stale region inside the second iteration's window: + # num_buffer_chunks > 2 * num_sub_groups -> hot_buffer_size > 512 at block + # 256. At block_size 512 or 1024 num_sub_groups is >= 16 and it cannot + # trigger. 8192 additionally keeps the corrupted offsets inside + # s_lru_slots_out, so the failure is a deterministic wrong LRU order rather + # than an out-of-bounds local-memory write. + hot_buffer_size = 8192 + num_top_k = 32 + padded_size = hot_buffer_size + 1 + seq_len = hot_buffer_size + 1 # > hot_buffer_size, so the fast path is skipped + + host_cache = _pinned((num_top_k, 1, KV_DIM), DTYPE) + host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache)) + device_buffer = torch.full( + (padded_size, 1, KV_DIM), -1.0, dtype=DTYPE, device=DEVICE + ) + # Identity slot -> device loc mapping keeps the expected values readable. + device_buffer_locs = torch.arange( + padded_size, dtype=torch.int32, device=DEVICE + ).view(1, -1) + # Nothing is resident, so every buffer slot is evictable (the path whose + # chunk counts are non-zero, hence the one that carries the corruption) and + # every query is a miss. + device_buffer_tokens = torch.full( + (1, padded_size), -1, dtype=torch.int32, device=DEVICE + ) + lru_slots = torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE).view( + 1, -1 + ) + host_cache_locs = torch.arange(seq_len, dtype=torch.int64, device=DEVICE).view( + 1, -1 + ) + top_k_tokens = torch.arange(num_top_k, dtype=torch.int32, device=DEVICE).view(1, -1) + out = torch.full_like(top_k_tokens, -1) + + load_cache_to_device_buffer_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE), + lru_slots=lru_slots, + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=num_top_k, + hot_buffer_size=hot_buffer_size, + page_size=1, + block_size=256, + num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + + # Evictables compact backwards from hot_buffer_size - 1 in LRU order, so + # miss j takes buffer slot j, and the LRU list rotates left by num_top_k as + # the just-filled slots move to the MRU tail. + assert torch.equal( + out.cpu(), torch.arange(num_top_k, dtype=torch.int32).view(1, -1) + ) + assert torch.equal( + device_buffer_tokens[0, :num_top_k].cpu(), + torch.arange(num_top_k, dtype=torch.int32), + ) + assert torch.equal( + device_buffer_tokens[0, num_top_k:].cpu(), + torch.full((padded_size - num_top_k,), -1, dtype=torch.int32), + ) + assert torch.equal( + lru_slots[0].cpu(), + torch.cat( + [ + torch.arange(num_top_k, hot_buffer_size, dtype=torch.int16), + torch.arange(num_top_k, dtype=torch.int16), + ] + ), + ) + for token in range(num_top_k): + assert torch.equal(device_buffer[token].cpu(), host_cache[token]) + + def test_load_cache_to_device_buffer_batched_with_padding() -> None: state = _make_state( [ @@ -423,6 +527,8 @@ def test_load_cache_to_device_buffer_batched_with_padding() -> None: ), seq_lens=torch.tensor([8, 3, 8], dtype=torch.int32, device=DEVICE), num_real_reqs=2, + # The padded row must be *written* as -1, not merely left alone. + output_fill_value=123456, **state, )