diff --git a/python/sglang/kernels/jit/csrc/hisparse.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh similarity index 100% rename from python/sglang/kernels/jit/csrc/hisparse.cuh rename to python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse_spec.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse_spec.cuh new file mode 100644 index 000000000..b84fa6b2a --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse_spec.cuh @@ -0,0 +1,1113 @@ +#pragma once + +// Multi-step speculative HiSparse cache-management kernels. +// DSv4 cache transfer remains in hisparse.cuh. + +#include // TensorMatcher and symbolic tensor validation +#include // RuntimeCheck and host utilities + +#include // LaunchKernel and PDL helpers + +#include +#include + +#include + +namespace sglang { + +#ifdef USE_ROCM +constexpr int WARP_SIZE = 64; +using BallotMask = uint64_t; +constexpr BallotMask FULL_WARP_MASK = 0xFFFFFFFFFFFFFFFFull; +#else +constexpr int WARP_SIZE = 32; +using BallotMask = unsigned int; +constexpr BallotMask FULL_WARP_MASK = 0xFFFFFFFFu; +#endif +constexpr int64_t HASH_DELETED = -2; +constexpr int32_t COMPACT_HASH_BITS = 13; +constexpr int32_t COMPACT_HASH_MASK = (1 << COMPACT_HASH_BITS) - 1; +constexpr int32_t COMPACT_APPROX_CLAIM_FLAG = int32_t{1} << 29; +constexpr int32_t COMPACT_APPROX_ADMIT_FLAG = int32_t{1} << 30; +constexpr int32_t APPROX_ADMISSION_STEPS = 1; +constexpr int32_t CLOCK_VICTIM_SAMPLES = 6; +constexpr uint32_t HASH_DEGRADED_FLAG = uint32_t{1} << 31; +constexpr uint32_t SCRATCH_EPOCH_MASK = HASH_DEGRADED_FLAG - 1; + +struct SpecCacheState { + int64_t* __restrict__ hash_primary; + int64_t* __restrict__ hash_secondary; + int32_t* __restrict__ ring_state; + int32_t* __restrict__ ref_epochs; + int64_t hash_stride; + int64_t hash_size; + int64_t ref_epoch_stride; +}; + +struct SpecMissWorkspace { + int32_t* __restrict__ locs; + int32_t* __restrict__ metadata; + int32_t* __restrict__ counters; + int64_t loc_stride; + int64_t metadata_stride; + int64_t counter_capacity; +}; + +constexpr int32_t ceil_log2_constexpr(int32_t value) { + int32_t bits = 0; + int32_t capacity = 1; + while (capacity < value) { + capacity <<= 1; + ++bits; + } + return bits; +} + +template +struct PackedRingState { + static constexpr int32_t CURSOR_BITS = ceil_log2_constexpr(HOT_BUFFER_SIZE); + static constexpr int32_t EPOCH_BITS = 32 - CURSOR_BITS < 23 ? 32 - CURSOR_BITS : 23; + static constexpr uint32_t CURSOR_MASK = (uint32_t{1} << CURSOR_BITS) - 1; + static constexpr uint32_t EPOCH_MASK = (uint32_t{1} << EPOCH_BITS) - 1; + + static_assert(HOT_BUFFER_SIZE > 1, "speculative hot buffer must contain at least two slots."); + static_assert(CURSOR_BITS <= 15, "speculative CLOCK cursor requires hot_buffer_size <= 32768."); + + __device__ static int32_t next_epoch(int32_t packed_state) { + int32_t epoch = static_cast((static_cast(packed_state) >> CURSOR_BITS) + 1) & EPOCH_MASK; + return epoch == 0 ? 1 : epoch; + } + + __device__ static int32_t cursor(int32_t packed_state) { + return static_cast(static_cast(packed_state) & CURSOR_MASK) % HOT_BUFFER_SIZE; + } + + __device__ static int32_t pack(int32_t epoch, int32_t cursor) { + return static_cast( + (static_cast(epoch) << CURSOR_BITS) | (static_cast(cursor) & CURSOR_MASK)); + } +}; + +template +struct PackedRingEntry { + static constexpr int64_t TOKEN_CAPACITY = int64_t{1} << 31; + + // Keep the token in the low word so the common lookup path can compare it + // without a 64-bit shift; the slot is decoded only after a token match. + __device__ static int64_t pack(int32_t token, int32_t slot) { + return static_cast( + (static_cast(static_cast(slot)) << 32) | static_cast(token)); + } + + __device__ static int32_t token(int64_t packed) { + return static_cast(static_cast(packed)); + } + + __device__ static int32_t slot(int64_t packed) { + return static_cast(static_cast(packed) >> 32); + } +}; + +__device__ __forceinline__ int64_t atomic_cas_i64(int64_t* address, int64_t expected, int64_t desired) { + return static_cast(atomicCAS( + reinterpret_cast(address), + static_cast(expected), + static_cast(desired))); +} + +__device__ __forceinline__ int64_t atomic_exch_i64(int64_t* address, int64_t value) { + return static_cast( + atomicExch(reinterpret_cast(address), static_cast(value))); +} + +// Knuth multiplicative hash for open-addressing table of size hash_size. +__device__ __forceinline__ int hash_slot(int32_t key, int hash_size) { + const uint32_t size = static_cast(hash_size); + const uint32_t hash = static_cast(key) * 2654435761u; + return static_cast((size & (size - 1)) == 0 ? hash & (size - 1) : hash % size); +} + +__device__ __forceinline__ int next_hash_slot(int slot, int hash_size) { + const uint32_t size = static_cast(hash_size); + const uint32_t next = static_cast(slot + 1); + return static_cast((size & (size - 1)) == 0 ? next & (size - 1) : next % size); +} + +__device__ __forceinline__ int ring_hash_slot(int32_t key, int64_t hash_size) { + return static_cast((static_cast(key) * 2654435761u) & static_cast(hash_size - 1)); +} + +__device__ __forceinline__ int ring_hash_slot_secondary(int32_t key, int64_t hash_size) { + uint32_t hash = static_cast(key); + hash ^= hash >> 16; + hash *= 0x7FEB352Du; + hash ^= hash >> 15; + hash *= 0x846CA68Bu; + hash ^= hash >> 16; + return static_cast(hash & static_cast(hash_size - 1)); +} + +__device__ __forceinline__ bool mark_cache_epoch(int32_t* __restrict__ cache_ref, int32_t cache_epoch) { + return atomicExch(cache_ref, cache_epoch) != cache_epoch; +} + +__device__ __forceinline__ int32_t scratch_union_lookup( + const unsigned long long* __restrict__ table, + const unsigned long long* __restrict__ indices, + int32_t table_size, + uint32_t scratch_epoch, + int32_t token) { + int32_t hash_pos = hash_slot(token, table_size); + for (int32_t attempt = 0; attempt < table_size; ++attempt) { + const auto packed = table[hash_pos]; + if (static_cast(packed >> 32) != scratch_epoch) { + return -1; + } + if (static_cast(static_cast(packed)) == token) { + const auto index_entry = indices[hash_pos]; + return static_cast(index_entry >> 32) == scratch_epoch + ? static_cast(static_cast(index_entry)) + : -1; + } + hash_pos = next_hash_slot(hash_pos, table_size); + } + return -1; +} + +template +__device__ __forceinline__ int32_t ring_hash_lookup( + const int64_t* __restrict__ keys, + const int64_t* __restrict__ vals, + int64_t hash_size, + int32_t token, + const int32_t* __restrict__ req_device_buffer_tokens, + bool hash_degraded) { + if (token < 0) { + return -1; + } + using Entry = PackedRingEntry; + const int64_t primary = keys[ring_hash_slot(token, hash_size)]; + if (primary >= 0 && Entry::token(primary) == token) { + const int32_t slot = Entry::slot(primary); + if (slot < HOT_BUFFER_SIZE && req_device_buffer_tokens[slot] == token) { + return slot; + } + } + const int64_t secondary = vals[ring_hash_slot_secondary(token, hash_size)]; + if (secondary >= 0 && Entry::token(secondary) == token) { + const int32_t slot = Entry::slot(secondary); + if (slot < HOT_BUFFER_SIZE && req_device_buffer_tokens[slot] == token) { + return slot; + } + } + if (hash_degraded) { + for (int32_t slot = 0; slot < HOT_BUFFER_SIZE; ++slot) { + if (req_device_buffer_tokens[slot] == token) { + return slot; + } + } + } + return -1; +} + +template +__device__ __forceinline__ int32_t hot_cache_lookup( + const int64_t* __restrict__ keys, + const int64_t* __restrict__ vals, + int64_t hash_size, + int32_t token, + const int32_t* __restrict__ req_device_buffer_tokens, + bool hash_degraded) { + if (token >= 0 && token < HOT_BUFFER_SIZE && req_device_buffer_tokens[token] == token) { + return token; + } + return ring_hash_lookup(keys, vals, hash_size, token, req_device_buffer_tokens, hash_degraded); +} + +template +__device__ __forceinline__ int32_t ring_hash_insert_atomic( + int64_t* __restrict__ keys, int64_t* __restrict__ vals, int64_t hash_size, int32_t token, int32_t buf_slot) { + if (token < 0) { + return -1; + } + using Entry = PackedRingEntry; + const int64_t packed = Entry::pack(token, buf_slot); + const int32_t primary_slot = ring_hash_slot(token, hash_size); + const int32_t secondary_slot = ring_hash_slot_secondary(token, hash_size); + int64_t old = keys[primary_slot]; + if (old >= 0 && Entry::token(old) == token) { + atomic_exch_i64(keys + primary_slot, packed); + return primary_slot; + } + if (old < 0 && atomic_cas_i64(keys + primary_slot, old, packed) == old) { + return primary_slot; + } + old = vals[secondary_slot]; + if (old >= 0 && Entry::token(old) == token) { + atomic_exch_i64(vals + secondary_slot, packed); + return secondary_slot; + } + if (old < 0 && atomic_cas_i64(vals + secondary_slot, old, packed) == old) { + return secondary_slot; + } + + int64_t current = packed; + bool use_secondary = false; + for (int32_t kick = 0; kick < 64; ++kick) { + const int32_t current_token = Entry::token(current); + int64_t* table = use_secondary ? vals : keys; + const int32_t slot = + use_secondary ? ring_hash_slot_secondary(current_token, hash_size) : ring_hash_slot(current_token, hash_size); + const int64_t displaced = atomic_exch_i64(table + slot, current); + if (displaced < 0 || Entry::token(displaced) == current_token) { + return slot; + } + current = displaced; + use_secondary = !use_secondary; + } + return -1; +} + +template +__device__ __forceinline__ void ring_hash_erase_atomic( + int64_t* __restrict__ keys, int64_t* __restrict__ vals, int64_t hash_size, int32_t token, int32_t buf_slot) { + if (token < 0) { + return; + } + using Entry = PackedRingEntry; + const int64_t expected = Entry::pack(token, buf_slot); + const int32_t primary_slot = ring_hash_slot(token, hash_size); + if (atomic_cas_i64(keys + primary_slot, expected, HASH_DELETED) == expected) { + return; + } + const int32_t secondary_slot = ring_hash_slot_secondary(token, hash_size); + atomic_cas_i64(vals + secondary_slot, expected, HASH_DELETED); +} + +#ifdef USE_ROCM +template +__device__ __forceinline__ void +transfer_item_warp(int32_t lane_id, const void* __restrict__ src_addr, void* __restrict__ dst_addr) { + const auto src = static_cast(src_addr); + auto dst = static_cast(dst_addr); + + constexpr int64_t word_count = ITEM_SIZE_BYTES / static_cast(sizeof(uint64_t)); + const auto src_words = reinterpret_cast(src); + auto dst_words = reinterpret_cast(dst); + for (int64_t i = lane_id; i < word_count; i += WARP_SIZE) { + dst_words[i] = src_words[i]; + } + + constexpr int64_t tail_start = word_count * static_cast(sizeof(uint64_t)); + for (int64_t i = tail_start + lane_id; i < ITEM_SIZE_BYTES; i += WARP_SIZE) { + dst[i] = src[i]; + } +} + +#else +template +__device__ __forceinline__ void transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr) { + // Issue the 512B body and 64B edge loads before either store so both host + // reads can remain in flight. Rows alternate between 0B and 64B offsets in + // a 128B transaction, so place the body on the aligned side of each row. + if constexpr (ITEM_SIZE_BYTES == 576) { + const auto src = static_cast(src_addr); + auto dst = static_cast(dst_addr); + const bool edge_first = (reinterpret_cast(src_addr) & 127u) == 64u; + const int32_t body_offset = edge_first ? 64 : 0; + const int32_t edge_offset = edge_first ? 0 : 512; + uint64_t body_lo, body_hi; + uint64_t edge_lo, edge_hi; + const auto body_src = reinterpret_cast(src + body_offset + lane_id * 16); + auto body_dst = reinterpret_cast(dst + body_offset + lane_id * 16); + asm volatile("ld.global.nc.v2.b64 {%0,%1},[%2];" : "=l"(body_lo), "=l"(body_hi) : "l"(body_src) : "memory"); + if (lane_id < 4) { + const auto edge_src = reinterpret_cast(src + edge_offset + lane_id * 16); + asm volatile("ld.global.nc.v2.b64 {%0,%1},[%2];" : "=l"(edge_lo), "=l"(edge_hi) : "l"(edge_src) : "memory"); + } + asm volatile("st.global.cg.v2.b64 [%0],{%1,%2};" ::"l"(body_dst), "l"(body_lo), "l"(body_hi) : "memory"); + if (lane_id < 4) { + auto edge_dst = reinterpret_cast(dst + edge_offset + lane_id * 16); + asm volatile("st.global.cg.v2.b64 [%0],{%1,%2};" ::"l"(edge_dst), "l"(edge_lo), "l"(edge_hi) : "memory"); + } + return; + } + + // 128-bit bulk transfer via paired 64-bit loads (avoids alignment issues with uint4) + constexpr int total_pairs = ITEM_SIZE_BYTES / 16; // number of 16-byte chunks + { + const uint64_t* __restrict__ src = static_cast(src_addr); + uint64_t* __restrict__ dst = static_cast(dst_addr); + for (int j = lane_id; j < total_pairs; j += WARP_SIZE) { + uint64_t lo, hi; + const uint64_t* s = src + j * 2; + asm volatile("ld.global.nc.v2.b64 {%0,%1},[%2];" : "=l"(lo), "=l"(hi) : "l"(s) : "memory"); + uint64_t* d = dst + j * 2; + asm volatile("st.global.cg.v2.b64 [%0],{%1,%2};" ::"l"(d), "l"(lo), "l"(hi) : "memory"); + } + } + + // Tail: 64-bit for remaining 8-byte chunk (if item_size not multiple of 16) + constexpr int tail_8B = (ITEM_SIZE_BYTES - total_pairs * 16) / 8; + if (tail_8B > 0 && lane_id < tail_8B) { + const uint64_t* __restrict__ src8 = + reinterpret_cast(static_cast(src_addr) + total_pairs * 16); + uint64_t* __restrict__ dst8 = reinterpret_cast(static_cast(dst_addr) + total_pairs * 16); + uint64_t tmp; + asm volatile("ld.global.nc.b64 %0,[%1];" : "=l"(tmp) : "l"(src8 + lane_id) : "memory"); + asm volatile("st.global.cg.b64 [%0],%1;" ::"l"(dst8 + lane_id), "l"(tmp) : "memory"); + } +} + +#endif + +__device__ __forceinline__ int first_set_lane(BallotMask mask) { +#ifdef USE_ROCM + return __ffsll(mask) - 1; +#else + return __ffs(mask) - 1; +#endif +} + +template +__device__ __forceinline__ bool try_get_extra_page_device_loc( + int32_t token_idx, + int64_t seq_len, + const int32_t* __restrict__ req_device_buffer_tokens, + const int32_t* __restrict__ req_device_buffer_locs, + int64_t page_size, + int32_t* __restrict__ out_loc) { + int64_t slot = -1; + if (static_cast(token_idx) >= seq_len - page_size) { + if (page_size == 4) { + const int4 page_tokens = *reinterpret_cast(req_device_buffer_tokens + HOT_BUFFER_SIZE); + if (page_tokens.x == token_idx) { + slot = HOT_BUFFER_SIZE; + } else if (page_tokens.y == token_idx) { + slot = HOT_BUFFER_SIZE + 1; + } else if (page_tokens.z == token_idx) { + slot = HOT_BUFFER_SIZE + 2; + } else if (page_tokens.w == token_idx) { + slot = HOT_BUFFER_SIZE + 3; + } + } else { + for (int64_t candidate_slot = HOT_BUFFER_SIZE; candidate_slot < HOT_BUFFER_SIZE + page_size; candidate_slot++) { + if (req_device_buffer_tokens[candidate_slot] == token_idx) { + slot = candidate_slot; + break; + } + } + } + } + + if (slot < HOT_BUFFER_SIZE || slot >= HOT_BUFFER_SIZE + page_size) { + return false; + } + const int32_t loc = req_device_buffer_locs[slot]; + if (loc < 0) { + return false; + } + *out_loc = loc; + return true; +} + +// Flatten all speculative steps. Each lane resolves one occurrence; the warp +// cooperatively copies only lanes that won a unique-miss claim. +template +__global__ void load_cache_to_device_buffer_spec_gather_kernel( + const int32_t* __restrict__ top_k_tokens, + int32_t* __restrict__ device_buffer_tokens, + const int64_t* __restrict__ host_cache_locs, + int32_t* __restrict__ device_buffer_locs, + const void* __restrict__ host_cache_k, + void* __restrict__ device_buffer_k, + int32_t* __restrict__ top_k_device_locs, + const int64_t* __restrict__ req_pool_indices, + const int32_t* __restrict__ seq_lens, + SpecCacheState cache_state, + SpecMissWorkspace miss_workspace, + const int32_t* __restrict__ num_real_reqs, + int64_t* __restrict__ miss_src_out, + int32_t* __restrict__ miss_dst_out, + int32_t* __restrict__ miss_count_out, + int64_t buffer_stride_0, + int64_t host_stride, + int64_t top_k_tokens_stride, + int64_t top_k_device_locs_stride, + int64_t plan_stride, + int64_t page_size) { + const int bid = blockIdx.x; + const int tid = threadIdx.x; + constexpr int64_t total_occurrences = NUM_STEPS * NUM_TOP_K; + int32_t* req_top_k_device_locs = top_k_device_locs + bid * top_k_device_locs_stride; + if (bid >= num_real_reqs[0]) { + if constexpr (RecordMissPlan) { + if (blockIdx.y == 0 && tid == 0) { + miss_count_out[bid] = 0; + } + } + for (int64_t i = tid; i < total_occurrences; i += BLOCK_SIZE) + req_top_k_device_locs[i] = 0; + return; + } + + const int warp_id = tid / WARP_SIZE; + const int lane_id = tid % WARP_SIZE; + constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; + const int64_t total_warps = static_cast(gridDim.y) * NUM_WARPS; + const int64_t global_warp = static_cast(blockIdx.y) * NUM_WARPS + warp_id; + const int64_t occ = static_cast(lane_id) * total_warps + global_warp; + const int64_t rid = req_pool_indices[bid]; + const int64_t buffer_offset = rid * buffer_stride_0; + + int32_t* req_device_buffer_tokens = device_buffer_tokens + buffer_offset; + int32_t* req_device_buffer_locs = device_buffer_locs + buffer_offset; + const int64_t* req_host_cache_locs = host_cache_locs + rid * host_stride; + int64_t* req_ring_hash_keys = cache_state.hash_primary + rid * cache_state.hash_stride; + int64_t* req_ring_hash_vals = cache_state.hash_secondary + rid * cache_state.hash_stride; + int32_t* req_cache_ref_bits = cache_state.ref_epochs + rid * cache_state.ref_epoch_stride; + int32_t* req_scratch_locs = miss_workspace.locs + rid * miss_workspace.loc_stride; + int32_t* req_scratch_tokens = miss_workspace.metadata + rid * miss_workspace.metadata_stride; + auto* req_scratch_table = reinterpret_cast(req_scratch_tokens); + auto* req_scratch_indices = req_scratch_table + total_occurrences; + int32_t* req_compact_hash_positions = req_scratch_tokens + 4 * total_occurrences; + int32_t* req_work_count = miss_workspace.counters + miss_workspace.counter_capacity + rid; + int32_t* req_union_hit_count = miss_workspace.counters + 2 * miss_workspace.counter_capacity + rid; + int32_t* req_scratch_generation = miss_workspace.counters + 3 * miss_workspace.counter_capacity + rid; + using RingState = PackedRingState; + const int32_t cache_epoch = RingState::next_epoch(cache_state.ring_state[rid]); + const uint32_t packed_scratch_generation = static_cast(*req_scratch_generation); + const bool hash_degraded = (packed_scratch_generation & HASH_DEGRADED_FLAG) != 0; + uint32_t next_scratch_epoch = (packed_scratch_generation & SCRATCH_EPOCH_MASK) + 1; + next_scratch_epoch = next_scratch_epoch > SCRATCH_EPOCH_MASK ? 1 : next_scratch_epoch; + const int32_t scratch_epoch = static_cast(next_scratch_epoch); + const int32_t* req_top_k_tokens = top_k_tokens + bid * top_k_tokens_stride; + + int32_t token = -1; + int32_t loc = -1; + int32_t cache_slot = -1; + bool copy_owner = false; + int32_t miss_idx = -1; + int64_t src_loc = -1; + bool needs_cache_lookup = false; + if (occ < total_occurrences) { + const int64_t step = occ / NUM_TOP_K; + token = req_top_k_tokens[occ]; + const int64_t seq_len = static_cast(seq_lens[bid * NUM_STEPS + step]); + if (token >= 0 && token < seq_len) { + int32_t direct_loc = -1; + if (try_get_extra_page_device_loc( + token, seq_len, req_device_buffer_tokens, req_device_buffer_locs, page_size, &direct_loc)) { + loc = direct_loc; + } else { + needs_cache_lookup = true; + } + } + } + + // The flattened occurrence mapping places the same top-k position from + // different speculative steps in one warp. Stable hits only need one Ring lookup. + int32_t lookup_owner_lane = lane_id; + bool reuse_owner_loc = false; + if (total_warps <= NUM_TOP_K && NUM_TOP_K % total_warps == 0) { + const int32_t lanes_per_step = static_cast(NUM_TOP_K / total_warps); + if (lanes_per_step > 0 && NUM_STEPS * lanes_per_step <= WARP_SIZE) { + lookup_owner_lane = lane_id % lanes_per_step; + const int32_t owner_token = __shfl_sync(FULL_WARP_MASK, token, lookup_owner_lane); + const int32_t owner_needs_lookup = __shfl_sync(FULL_WARP_MASK, needs_cache_lookup ? 1 : 0, lookup_owner_lane); + reuse_owner_loc = + needs_cache_lookup && lane_id != lookup_owner_lane && owner_needs_lookup != 0 && owner_token == token; + } + } + + if (needs_cache_lookup && !reuse_owner_loc) { + cache_slot = hot_cache_lookup( + req_ring_hash_keys, req_ring_hash_vals, cache_state.hash_size, token, req_device_buffer_tokens, hash_degraded); + if (cache_slot >= 0) { + loc = req_device_buffer_locs[cache_slot]; + } else { + src_loc = req_host_cache_locs[token]; + if (src_loc >= 0) { + int32_t hash_pos = hash_slot(token, static_cast(total_occurrences)); + const auto epoch_bits = static_cast(static_cast(scratch_epoch)) << 32; + const auto token_bits = static_cast(token); + const auto packed = epoch_bits | token_bits; + for (int64_t attempt = 0; attempt < total_occurrences; ++attempt) { + auto old = atomicCAS(req_scratch_table + hash_pos, 0ull, 0ull); + if (static_cast(old >> 32) != static_cast(scratch_epoch)) { + const auto claimed = atomicCAS(req_scratch_table + hash_pos, old, packed); + if (claimed == old) { + const int32_t unique_idx = atomicAdd(req_work_count, 1); + if constexpr (RecordMissPlan) { + miss_idx = unique_idx; + } + req_compact_hash_positions[unique_idx] = hash_pos; + __threadfence(); + atomicExch(req_scratch_indices + hash_pos, epoch_bits | static_cast(unique_idx)); + if (unique_idx < miss_workspace.loc_stride) { + loc = req_scratch_locs[unique_idx]; + copy_owner = loc >= 0; + } + break; + } + continue; + } + if (static_cast(old) == token_bits) { + auto index_entry = atomicCAS(req_scratch_indices + hash_pos, 0ull, 0ull); + while (static_cast(index_entry >> 32) != static_cast(scratch_epoch)) { + index_entry = atomicCAS(req_scratch_indices + hash_pos, 0ull, 0ull); + } + const int32_t unique_idx = static_cast(static_cast(index_entry)); + if (unique_idx < miss_workspace.loc_stride) { + loc = req_scratch_locs[unique_idx]; + } + break; + } + hash_pos = next_hash_slot(hash_pos, static_cast(total_occurrences)); + } + } + } + } + + const int32_t owner_loc = __shfl_sync(FULL_WARP_MASK, loc, lookup_owner_lane); + const int32_t owner_cache_slot = __shfl_sync(FULL_WARP_MASK, cache_slot, lookup_owner_lane); + if (reuse_owner_loc) { + loc = owner_loc; + cache_slot = owner_cache_slot; + } + if (cache_slot >= 0) { + if (mark_cache_epoch(req_cache_ref_bits + cache_slot, cache_epoch)) { + atomicAdd(req_union_hit_count, 1); + } + } + if (occ < total_occurrences) { + req_top_k_device_locs[occ] = loc; + } + + BallotMask copy_mask = __ballot_sync(FULL_WARP_MASK, copy_owner); + while (copy_mask != 0) { + const int owner_lane = first_set_lane(copy_mask); + const int32_t copy_token = __shfl_sync(FULL_WARP_MASK, token, owner_lane); + const int32_t copy_loc = __shfl_sync(FULL_WARP_MASK, loc, owner_lane); + const int64_t copy_src_loc = __shfl_sync(FULL_WARP_MASK, src_loc, owner_lane); + if (copy_token >= 0 && copy_loc >= 0 && copy_src_loc >= 0) { + if constexpr (RecordMissPlan) { + const int32_t copy_miss_idx = __shfl_sync(FULL_WARP_MASK, miss_idx, owner_lane); + if (lane_id == owner_lane) { + miss_src_out[bid * plan_stride + copy_miss_idx] = copy_src_loc; + miss_dst_out[bid * plan_stride + copy_miss_idx] = copy_loc; + } + } + const auto src_k = static_cast(host_cache_k) + copy_src_loc * ITEM_SIZE_BYTES; + auto dst_k = static_cast(device_buffer_k) + static_cast(copy_loc) * ITEM_SIZE_BYTES; + transfer_item_warp(lane_id, src_k, dst_k); + } + copy_mask &= ~(static_cast(1) << owner_lane); + } +} + +template +__global__ void load_cache_to_device_buffer_spec_commit_kernel( + const int32_t* __restrict__ top_k_tokens, + int32_t* __restrict__ top_k_device_locs, + int32_t* __restrict__ device_buffer_tokens, + int32_t* __restrict__ device_buffer_locs, + const int64_t* __restrict__ host_cache_locs, + const void* __restrict__ host_cache_k, + void* __restrict__ device_buffer_k, + const int64_t* __restrict__ req_pool_indices, + SpecCacheState cache_state, + SpecMissWorkspace miss_workspace, + const int32_t* __restrict__ num_real_reqs, + int64_t* __restrict__ miss_src_out, + int32_t* __restrict__ miss_dst_out, + int32_t* __restrict__ miss_count_out, + int64_t top_k_tokens_stride, + int64_t top_k_device_locs_stride, + int64_t buffer_stride_0, + int64_t host_stride, + int64_t plan_stride) { + device::PDLWaitPrimary(); + const int bid = blockIdx.x; + if (bid >= num_real_reqs[0]) return; + + const int lane_id = threadIdx.x % WARP_SIZE; + const int64_t rid = req_pool_indices[bid]; + constexpr int64_t total_occurrences = NUM_STEPS * NUM_TOP_K; + using RingState = PackedRingState; + const int32_t packed_ring_state = cache_state.ring_state[rid]; + const int32_t cache_epoch = RingState::next_epoch(packed_ring_state); + const int32_t clock_cursor = RingState::cursor(packed_ring_state); + const int64_t buffer_offset = rid * buffer_stride_0; + int32_t* req_device_buffer_tokens = device_buffer_tokens + buffer_offset; + int32_t* req_device_buffer_locs = device_buffer_locs + buffer_offset; + int64_t* req_ring_hash_keys = cache_state.hash_primary + rid * cache_state.hash_stride; + int64_t* req_ring_hash_vals = cache_state.hash_secondary + rid * cache_state.hash_stride; + int32_t* req_cache_ref_bits = cache_state.ref_epochs + rid * cache_state.ref_epoch_stride; + int32_t* req_scratch_locs = miss_workspace.locs + rid * miss_workspace.loc_stride; + int32_t* req_scratch_tokens = miss_workspace.metadata + rid * miss_workspace.metadata_stride; + const auto* req_scratch_table = reinterpret_cast(req_scratch_tokens); + auto* req_scratch_indices = reinterpret_cast(req_scratch_tokens) + total_occurrences; + int32_t* req_compact_hash_positions = req_scratch_tokens + 4 * total_occurrences; + const int64_t* req_host_cache_locs = host_cache_locs + rid * host_stride; + 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; + int32_t* req_work_count = miss_workspace.counters + miss_workspace.counter_capacity + rid; + int32_t* req_union_hit_count = miss_workspace.counters + 2 * miss_workspace.counter_capacity + rid; + int32_t* req_scratch_generation = miss_workspace.counters + 3 * miss_workspace.counter_capacity + rid; + const uint32_t packed_scratch_generation = static_cast(*req_scratch_generation); + const bool hash_degraded = (packed_scratch_generation & HASH_DEGRADED_FLAG) != 0; + uint32_t next_scratch_epoch = (packed_scratch_generation & SCRATCH_EPOCH_MASK) + 1; + next_scratch_epoch = next_scratch_epoch > SCRATCH_EPOCH_MASK ? 1 : next_scratch_epoch; + const int32_t scratch_epoch = static_cast(next_scratch_epoch); + const int32_t miss_count = *req_work_count; + const int32_t union_hit_count = *req_union_hit_count; + const bool scratch_overflow = miss_count > miss_workspace.loc_stride; + const int64_t compact_iterations = (miss_count + blockDim.x - 1) / blockDim.x; + const bool lock_free_single_pass = compact_iterations == 1; + + // The union fast path can preserve every token needed by this speculative group. + // Its victim selection below partitions the ring by miss ordinal, avoiding + // a full hot-cache scan while keeping all current union hits protected. + __shared__ int32_t s_use_union_clock; + __shared__ int32_t s_approx_admission_budget; + __shared__ int32_t s_hash_degraded; + if (threadIdx.x == 0) { + s_use_union_clock = miss_count + union_hit_count <= HOT_BUFFER_SIZE; + const int32_t mandatory_direct_misses = + scratch_overflow ? miss_count - static_cast(miss_workspace.loc_stride) : 0; + s_approx_admission_budget = + s_use_union_clock ? 0 + : (scratch_overflow ? max(0, HOT_BUFFER_SIZE - union_hit_count - mandatory_direct_misses) + : HOT_BUFFER_SIZE); + s_hash_degraded = hash_degraded; + } + __syncthreads(); + if (!s_use_union_clock) { + // The complete union is already available from hot + scratch for this + // attention call. Only the cache update is approximate: keep the most + // recent speculative working set and avoid replaying four LRU passes. + if (!scratch_overflow) { + for (int32_t slot = threadIdx.x; slot < HOT_BUFFER_SIZE; slot += blockDim.x) { + req_cache_ref_bits[slot] = 0; + } + } + __syncthreads(); + if (threadIdx.x == 0) { + *req_union_hit_count = 0; + } + __syncthreads(); + constexpr int64_t admission_step = NUM_STEPS > APPROX_ADMISSION_STEPS ? NUM_STEPS - APPROX_ADMISSION_STEPS : 0; + const int64_t admission_start = admission_step * NUM_TOP_K; + for (int64_t occ = admission_start + threadIdx.x; occ < total_occurrences; occ += blockDim.x) { + const int32_t token = req_top_k_tokens[occ]; + if (token < 0) { + continue; + } + const int32_t slot = hot_cache_lookup( + req_ring_hash_keys, + req_ring_hash_vals, + cache_state.hash_size, + token, + req_device_buffer_tokens, + s_hash_degraded != 0); + if (slot >= 0) { + mark_cache_epoch(req_cache_ref_bits + slot, cache_epoch); + continue; + } + const int32_t unique_idx = scratch_union_lookup( + req_scratch_table, + req_scratch_indices, + static_cast(total_occurrences), + static_cast(scratch_epoch), + token); + if (unique_idx >= 0 && unique_idx < miss_workspace.loc_stride) { + const int32_t old_metadata = atomicOr(req_compact_hash_positions + unique_idx, COMPACT_APPROX_CLAIM_FLAG); + if ((old_metadata & COMPACT_APPROX_CLAIM_FLAG) == 0) { + const int32_t admission_ordinal = atomicAdd(req_union_hit_count, 1); + if (admission_ordinal < s_approx_admission_budget) { + atomicOr(req_compact_hash_positions + unique_idx, COMPACT_APPROX_ADMIT_FLAG); + } + } + } + } + __syncthreads(); + } + + for (int64_t iteration = 0; iteration < compact_iterations; ++iteration) { + const int64_t compact_idx = threadIdx.x + iteration * blockDim.x; + bool copy_owner = false; + int64_t copy_src_loc = -1; + int32_t copy_dst_loc = -1; + int64_t hash_pos = -1; + int32_t token = -1; + int32_t unique_idx = -1; + int32_t victim = -1; + bool direct_overflow = false; + bool rotate_compact = false; + + if (compact_idx < miss_count) { + const int32_t compact_entry = req_compact_hash_positions[compact_idx]; + hash_pos = compact_entry & COMPACT_HASH_MASK; + const auto packed = req_scratch_table[hash_pos]; + const auto index_entry = req_scratch_indices[hash_pos]; + if (static_cast(packed >> 32) == static_cast(scratch_epoch) && + static_cast(index_entry >> 32) == static_cast(scratch_epoch)) { + token = static_cast(static_cast(packed)); + unique_idx = static_cast(static_cast(index_entry)); + direct_overflow = scratch_overflow && unique_idx >= miss_workspace.loc_stride; + const bool approximate_admission = !s_use_union_clock && (compact_entry & COMPACT_APPROX_ADMIT_FLAG) != 0; + rotate_compact = unique_idx < miss_workspace.loc_stride && (s_use_union_clock || approximate_admission); + + if ((direct_overflow || rotate_compact) && s_use_union_clock) { + uint32_t oldest_age = 0; + const int32_t partition_size = (HOT_BUFFER_SIZE - 1 - compact_idx) / miss_count + 1; + const int32_t sample_count = min(partition_size, CLOCK_VICTIM_SAMPLES); + for (int32_t sample = 0; sample < sample_count; ++sample) { + const int32_t partition_offset = sample * partition_size / sample_count; + const int32_t linear = compact_idx + partition_offset * miss_count; + const int32_t candidate = (clock_cursor + linear) % HOT_BUFFER_SIZE; + const int32_t ref_epoch = req_cache_ref_bits[candidate]; + if (ref_epoch == cache_epoch) { + continue; + } + const uint32_t age = + (static_cast(cache_epoch) - static_cast(ref_epoch)) & RingState::EPOCH_MASK; + if (victim < 0 || age > oldest_age) { + victim = candidate; + oldest_age = age; + } + } + if (victim >= 0) { + req_cache_ref_bits[victim] = cache_epoch; + } + } else if (direct_overflow || rotate_compact) { + const int32_t scratch_capacity = static_cast(miss_workspace.loc_stride); + const int32_t victim_count = direct_overflow ? miss_count - scratch_capacity : miss_count; + const int32_t victim_ordinal = direct_overflow ? unique_idx - scratch_capacity : unique_idx; + for (int32_t linear = victim_ordinal; linear < HOT_BUFFER_SIZE; linear += victim_count) { + const int32_t candidate = static_cast( + (static_cast(linear) * 2654435761u + static_cast(cache_epoch)) % HOT_BUFFER_SIZE); + const int32_t observed = req_cache_ref_bits[candidate]; + if (observed == cache_epoch) { + continue; + } + if (lock_free_single_pass) { + req_cache_ref_bits[candidate] = cache_epoch; + victim = candidate; + break; + } + if (atomicCAS(req_cache_ref_bits + candidate, observed, cache_epoch) == observed) { + victim = candidate; + break; + } + } + } + } + } + + // Ring partitions are disjoint. Mark their direct choices before the rare + // fallback probes globally for a partition that had no available slot. + __syncthreads(); + + if ((direct_overflow || rotate_compact) && victim < 0) { + const uint32_t start = + (static_cast(clock_cursor) + static_cast(compact_idx) * 2654435761u) % HOT_BUFFER_SIZE; + for (int32_t attempt = 0; attempt < HOT_BUFFER_SIZE; ++attempt) { + const int32_t candidate = + static_cast((start + static_cast(attempt) * 2654435761u) % HOT_BUFFER_SIZE); + const int32_t observed = req_cache_ref_bits[candidate]; + if (observed == cache_epoch) { + continue; + } + if (atomicCAS(req_cache_ref_bits + candidate, observed, cache_epoch) == observed) { + victim = candidate; + break; + } + } + } + + if (direct_overflow || rotate_compact) { + if (victim >= 0) { + const int32_t old_loc = req_device_buffer_locs[victim]; + const int32_t new_loc = direct_overflow ? old_loc : req_scratch_locs[unique_idx]; + if (new_loc >= 0 && old_loc >= 0) { + const int32_t old_token = req_device_buffer_tokens[victim]; + req_device_buffer_tokens[victim] = token; + req_cache_ref_bits[victim] = cache_epoch; + if (rotate_compact) { + req_device_buffer_locs[victim] = new_loc; + req_scratch_locs[unique_idx] = old_loc; + } else { + const auto epoch_bits = static_cast(static_cast(scratch_epoch)) << 32; + atomicExch(req_scratch_indices + hash_pos, epoch_bits | static_cast(old_loc)); + copy_src_loc = req_host_cache_locs[token]; + copy_dst_loc = old_loc; + copy_owner = copy_src_loc >= 0; + } + ring_hash_erase_atomic( + req_ring_hash_keys, req_ring_hash_vals, cache_state.hash_size, old_token, victim); + const int32_t inserted_slot = ring_hash_insert_atomic( + req_ring_hash_keys, req_ring_hash_vals, cache_state.hash_size, token, victim); + if (inserted_slot < 0) { + atomicExch(&s_hash_degraded, 1); + } + } + } else if (direct_overflow) { + const auto epoch_bits = static_cast(static_cast(scratch_epoch)) << 32; + atomicExch(req_scratch_indices + hash_pos, epoch_bits | UINT32_MAX); + } + } + + BallotMask copy_mask = __ballot_sync(FULL_WARP_MASK, copy_owner); + while (copy_mask != 0) { + const int owner_lane = first_set_lane(copy_mask); + const int64_t src_loc = __shfl_sync(FULL_WARP_MASK, copy_src_loc, owner_lane); + const int32_t dst_loc = __shfl_sync(FULL_WARP_MASK, copy_dst_loc, owner_lane); + if (src_loc >= 0 && dst_loc >= 0) { + if constexpr (RecordMissPlan) { + const int32_t miss_idx = __shfl_sync(FULL_WARP_MASK, unique_idx, owner_lane); + if (lane_id == owner_lane) { + miss_src_out[bid * plan_stride + miss_idx] = src_loc; + miss_dst_out[bid * plan_stride + miss_idx] = dst_loc; + } + } + const auto src_k = static_cast(host_cache_k) + src_loc * ITEM_SIZE_BYTES; + auto dst_k = static_cast(device_buffer_k) + static_cast(dst_loc) * ITEM_SIZE_BYTES; + transfer_item_warp(lane_id, src_k, dst_k); + } + copy_mask &= ~(static_cast(1) << owner_lane); + } + } + + __syncthreads(); + if (scratch_overflow) { + for (int64_t occ = threadIdx.x; occ < total_occurrences; occ += blockDim.x) { + if (req_top_k_device_locs[occ] >= 0) { + continue; + } + const int32_t token = req_top_k_tokens[occ]; + if (token < 0) { + continue; + } + + int32_t hash_pos = hash_slot(token, static_cast(total_occurrences)); + for (int64_t attempt = 0; attempt < total_occurrences; ++attempt) { + const auto packed = req_scratch_table[hash_pos]; + if (static_cast(packed >> 32) != static_cast(scratch_epoch)) { + break; + } + if (static_cast(packed) == static_cast(token)) { + const auto index_entry = req_scratch_indices[hash_pos]; + if (static_cast(index_entry >> 32) == static_cast(scratch_epoch)) { + const uint32_t resolved_loc = static_cast(index_entry); + if (resolved_loc != UINT32_MAX) { + req_top_k_device_locs[occ] = static_cast(resolved_loc); + } + } + break; + } + hash_pos = next_hash_slot(hash_pos, static_cast(total_occurrences)); + } + } + } + + __syncthreads(); + if (threadIdx.x == 0) { + miss_workspace.counters[rid] = miss_count; + if constexpr (RecordMissPlan) { + miss_count_out[bid] = miss_count; + } + *req_work_count = 0; + *req_union_hit_count = 0; + *req_scratch_generation = + static_cast(static_cast(scratch_epoch) | (s_hash_degraded != 0 ? HASH_DEGRADED_FLAG : 0)); + const int32_t next_cursor = + miss_count > 0 ? (clock_cursor + min(miss_count, HOT_BUFFER_SIZE)) % HOT_BUFFER_SIZE : clock_cursor; + cache_state.ring_state[rid] = RingState::pack(cache_epoch, next_cursor); + } +} + +template < + int BLOCK_SIZE, + int NUM_TOP_K, + int HOT_BUFFER_SIZE, + int ITEM_SIZE_BYTES, + int NUM_STEPS, + bool RecordMissPlan, + bool USE_PDL> +void load_cache_to_device_buffer_spec( + tvm::ffi::TensorView top_k_tokens, + tvm::ffi::TensorView device_buffer_tokens, + tvm::ffi::TensorView host_cache_locs, + tvm::ffi::TensorView device_buffer_locs, + tvm::ffi::TensorView host_cache_k, + tvm::ffi::TensorView device_buffer_k, + tvm::ffi::TensorView top_k_device_locs, + tvm::ffi::TensorView req_pool_indices, + tvm::ffi::TensorView seq_lens, + tvm::ffi::TensorView cache_index, + tvm::ffi::TensorView cache_policy, + tvm::ffi::TensorView scratch_locs, + tvm::ffi::TensorView scratch_state, + tvm::ffi::TensorView num_real_reqs, + int64_t page_size, + tvm::ffi::TensorView miss_src_out, + tvm::ffi::TensorView miss_dst_out, + tvm::ffi::TensorView miss_count_out) { + using namespace host; + + static_assert(NUM_STEPS > 1 && NUM_STEPS <= 4, "HiSparse speculative swap requires 2-4 steps."); + static_assert(NUM_TOP_K >= 1024, "HiSparse speculative swap requires top_k >= 1024."); + static_assert(NUM_STEPS * NUM_TOP_K <= 8192, "HiSparse speculative swap supports at most 8192 occurrences."); + + const int64_t bs = top_k_tokens.shape()[0]; + constexpr int64_t total_occurrences = NUM_STEPS * NUM_TOP_K; + RuntimeCheck(top_k_tokens.ndim() == 3, "speculative top_k_tokens must have shape [batch, steps, top_k]."); + RuntimeCheck(top_k_device_locs.ndim() == 3, "speculative output must have shape [batch, steps, top_k]."); + RuntimeCheck(top_k_tokens.shape()[1] == NUM_STEPS, "top_k_tokens step dimension mismatch."); + RuntimeCheck(top_k_tokens.shape()[2] == NUM_TOP_K, "top_k_tokens top-k dimension mismatch."); + RuntimeCheck( + cache_index.ndim() == 3 && cache_index.shape()[1] == 2, + "speculative cache_index must have shape [num_requests, 2, hash_size]."); + const int64_t ring_hash_size = cache_index.shape()[2]; + RuntimeCheck( + ring_hash_size > 0 && (ring_hash_size & (ring_hash_size - 1)) == 0, "ring hash capacity must be a power of two."); + RuntimeCheck( + scratch_locs.ndim() == 2 && scratch_locs.shape()[1] > 0, + "speculative scratch_locs must have shape [num_requests, scratch_capacity]."); + const int64_t num_request_slots = scratch_locs.shape()[0]; + RuntimeCheck( + cache_index.shape()[0] >= num_request_slots, + "speculative cache_index request capacity is smaller than scratch_locs."); + RuntimeCheck( + cache_policy.ndim() == 2 && cache_policy.shape()[0] >= num_request_slots + 1 && + cache_policy.shape()[1] >= HOT_BUFFER_SIZE, + "speculative cache_policy must contain one CLOCK control row and one reference-epoch row per request."); + RuntimeCheck( + scratch_state.ndim() == 2 && scratch_state.shape()[0] >= num_request_slots + 1 && + scratch_state.shape()[1] >= 4 * num_request_slots && scratch_state.shape()[1] >= 5 * total_occurrences, + "speculative scratch_state must contain one counter row and one miss-metadata row per request."); + RuntimeCheck(scratch_state.strides()[0] % 2 == 0, "speculative scratch metadata stride must be 64-bit aligned."); + + const int64_t host_stride = host_cache_locs.shape()[1]; + RuntimeCheck( + host_stride <= PackedRingEntry::TOKEN_CAPACITY, + "speculative packed ring metadata supports sequence lengths up to ", + PackedRingEntry::TOKEN_CAPACITY, + ", got ", + host_stride); + const int64_t buffer_stride_0 = device_buffer_tokens.strides()[0]; + const int64_t top_k_tokens_stride = top_k_tokens.strides()[0]; + const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0]; + const int64_t cache_index_stride_0 = cache_index.strides()[0]; + const int64_t cache_index_stride_1 = cache_index.strides()[1]; + const int64_t cache_policy_stride_0 = cache_policy.strides()[0]; + const int64_t scratch_stride_0 = scratch_locs.strides()[0]; + const int64_t scratch_state_stride_0 = scratch_state.strides()[0]; + int64_t* const miss_src_ptr = RecordMissPlan ? static_cast(miss_src_out.data_ptr()) : nullptr; + int32_t* const miss_dst_ptr = RecordMissPlan ? static_cast(miss_dst_out.data_ptr()) : nullptr; + int32_t* const miss_count_ptr = RecordMissPlan ? static_cast(miss_count_out.data_ptr()) : nullptr; + const int64_t plan_stride = RecordMissPlan ? miss_src_out.strides()[0] : 0; + if constexpr (RecordMissPlan) { + RuntimeCheck( + miss_src_out.ndim() == 2 && miss_src_out.shape()[0] >= bs && miss_src_out.shape()[1] >= total_occurrences, + "speculative miss_src must have shape [batch, >= steps * top_k]."); + RuntimeCheck( + miss_dst_out.ndim() == 2 && miss_dst_out.shape()[0] >= bs && miss_dst_out.shape()[1] >= total_occurrences, + "speculative miss_dst must have shape [batch, >= steps * top_k]."); + RuntimeCheck( + miss_count_out.ndim() == 1 && miss_count_out.shape()[0] >= bs, + "speculative miss_count must have shape [batch]."); + RuntimeCheck(miss_dst_out.strides()[0] == plan_stride, "speculative miss_src/miss_dst row strides differ."); + } + + // Preserve the device kernels' independent restrict-qualified views while + // exposing only four packed tensors through FFI. Row 0 is the control + // plane; request-owned rows begin at row 1. + auto* cache_index_ptr = static_cast(cache_index.data_ptr()); + auto* cache_policy_ptr = static_cast(cache_policy.data_ptr()); + auto* scratch_state_ptr = static_cast(scratch_state.data_ptr()); + const SpecCacheState cache_state{ + cache_index_ptr, + cache_index_ptr + cache_index_stride_1, + cache_policy_ptr, + cache_policy_ptr + cache_policy_stride_0, + cache_index_stride_0, + ring_hash_size, + cache_policy_stride_0}; + const SpecMissWorkspace miss_workspace{ + static_cast(scratch_locs.data_ptr()), + scratch_state_ptr + scratch_state_stride_0, + scratch_state_ptr, + scratch_stride_0, + scratch_state_stride_0, + num_request_slots}; + auto cuda_device = SymbolicDevice{}; + cuda_device.set_options(); + TensorMatcher({bs}).with_dtype().with_device(cuda_device).verify(req_pool_indices); + TensorMatcher({bs * NUM_STEPS}).with_dtype().with_device(cuda_device).verify(seq_lens); + const auto device = cuda_device.unwrap(); + + const uint32_t tiles = static_cast((total_occurrences + BLOCK_SIZE - 1) / BLOCK_SIZE); + LaunchKernel(dim3(static_cast(bs), tiles), BLOCK_SIZE, device)( + load_cache_to_device_buffer_spec_gather_kernel< + BLOCK_SIZE, + NUM_TOP_K, + HOT_BUFFER_SIZE, + ITEM_SIZE_BYTES, + NUM_STEPS, + RecordMissPlan>, + static_cast(top_k_tokens.data_ptr()), + static_cast(device_buffer_tokens.data_ptr()), + static_cast(host_cache_locs.data_ptr()), + static_cast(device_buffer_locs.data_ptr()), + host_cache_k.data_ptr(), + device_buffer_k.data_ptr(), + static_cast(top_k_device_locs.data_ptr()), + static_cast(req_pool_indices.data_ptr()), + static_cast(seq_lens.data_ptr()), + cache_state, + miss_workspace, + static_cast(num_real_reqs.data_ptr()), + miss_src_ptr, + miss_dst_ptr, + miss_count_ptr, + buffer_stride_0, + host_stride, + top_k_tokens_stride, + top_k_device_locs_stride, + plan_stride, + page_size); + + LaunchKernel(dim3(static_cast(bs)), 512, device) + .enable_pdl(USE_PDL)( + load_cache_to_device_buffer_spec_commit_kernel< + NUM_TOP_K, + HOT_BUFFER_SIZE, + ITEM_SIZE_BYTES, + NUM_STEPS, + RecordMissPlan, + USE_PDL>, + static_cast(top_k_tokens.data_ptr()), + static_cast(top_k_device_locs.data_ptr()), + static_cast(device_buffer_tokens.data_ptr()), + static_cast(device_buffer_locs.data_ptr()), + static_cast(host_cache_locs.data_ptr()), + host_cache_k.data_ptr(), + device_buffer_k.data_ptr(), + static_cast(req_pool_indices.data_ptr()), + cache_state, + miss_workspace, + static_cast(num_real_reqs.data_ptr()), + miss_src_ptr, + miss_dst_ptr, + miss_count_ptr, + top_k_tokens_stride, + top_k_device_locs_stride, + buffer_stride_0, + host_stride, + plan_stride); +} + +} // namespace sglang diff --git a/python/sglang/kernels/ops/kvcache/hisparse.py b/python/sglang/kernels/ops/kvcache/hisparse.py index fc0305d2c..cfebdcbf3 100644 --- a/python/sglang/kernels/ops/kvcache/hisparse.py +++ b/python/sglang/kernels/ops/kvcache/hisparse.py @@ -1,16 +1,171 @@ from __future__ import annotations import functools -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple import torch -from sglang.kernels.jit.utils import load_jit, make_cpp_args +from sglang.kernels.jit.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) if TYPE_CHECKING: from tvm_ffi.module import Module +_GATHER_BLOCK_SIZE = 64 + + +class HiSparseSpecState(NamedTuple): + """Persistent cache state and reusable miss workspace for speculative swap. + + ``cache_index`` stores the two int64 hash banks as + ``[num_requests, 2, hash_size]``. ``cache_policy`` uses a control-plane row + for the packed CLOCK states followed by one reference-epoch row per + request: ``[1 + num_requests, hot_buffer_size]``. + + ``scratch_locs`` and ``scratch_state`` hold reusable miss locations, + counters, and metadata shared by all layers. + """ + + cache_index: torch.Tensor + cache_policy: torch.Tensor + scratch_locs: torch.Tensor + scratch_state: torch.Tensor + + +@cache_once +def _jit_spec_module( + item_size_bytes: int, + block_size: int, + num_top_k: int, + hot_buffer_size: int, + num_steps: int, + record_miss_plan: bool, +) -> Module: + template_args = make_cpp_args( + block_size, + num_top_k, + hot_buffer_size, + item_size_bytes, + num_steps, + record_miss_plan, + is_arch_support_pdl(), + ) + return load_jit( + "hisparse_spec", + *template_args, + cuda_files=["kvcacheio/hisparse_spec.cuh"], + cuda_wrappers=[ + ( + "load_cache_to_device_buffer_spec", + f"load_cache_to_device_buffer_spec<{template_args}>", + ) + ], + ) + + +def load_cache_to_device_buffer_spec_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, + state: HiSparseSpecState, + num_real_reqs: torch.Tensor, + miss_src: torch.Tensor | None = None, + miss_dst: torch.Tensor | None = None, + miss_count: torch.Tensor | None = None, +) -> None: + """Resolve all speculative steps and swap unique misses in one launch pair. + + Optional miss-plan outputs use the same protocol as the single-step HiSparse + kernel, so shared-index layers can replay only the Host-to-GPU copies with + ``copy_cache_planned_mla``. + """ + _, num_steps, num_top_k = top_k_tokens.shape + if not 2 <= num_steps <= 4: + raise ValueError( + f"HiSparse speculative swap requires 2-4 steps, got {num_steps}." + ) + hot_buffer_size = state.cache_policy.size(1) + page_size = device_buffer_tokens.size(1) - hot_buffer_size + item_size_bytes = host_cache.stride(0) * host_cache.element_size() + record_miss_plan = miss_src is not None + if record_miss_plan: + if miss_dst is None or miss_count is None: + raise ValueError( + "miss_src, miss_dst, and miss_count must be provided together." + ) + if miss_src.dtype != torch.int64 or miss_dst.dtype != torch.int32: + raise ValueError("miss_src must be int64 and miss_dst must be int32.") + if miss_count.dtype != torch.int32: + raise ValueError("miss_count must be int32.") + plan_capacity = num_steps * num_top_k + batch_size = top_k_tokens.size(0) + if ( + miss_src.ndim != 2 + or miss_dst.ndim != 2 + or miss_src.size(0) < batch_size + or miss_dst.size(0) < batch_size + or miss_src.size(1) < plan_capacity + or miss_dst.size(1) < plan_capacity + ): + raise ValueError( + "speculative miss_src/miss_dst must have shape " + f"[batch, >= steps * top_k] (capacity {plan_capacity})." + ) + if miss_count.ndim != 1 or miss_count.numel() < batch_size: + raise ValueError("speculative miss_count must have shape [batch].") + if miss_src.stride(0) != miss_dst.stride(0): + raise ValueError("miss_src/miss_dst row strides must match.") + else: + if miss_dst is not None or miss_count is not None: + raise ValueError( + "miss_src, miss_dst, and miss_count must be provided together." + ) + empty = torch.empty(0) + miss_src = miss_dst = miss_count = empty + + module = _jit_spec_module( + item_size_bytes, + _GATHER_BLOCK_SIZE, + num_top_k, + hot_buffer_size, + num_steps, + record_miss_plan, + ) + + module.load_cache_to_device_buffer_spec( + 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, + state.cache_index, + state.cache_policy, + state.scratch_locs, + state.scratch_state, + num_real_reqs, + page_size, + miss_src, + miss_dst, + miss_count, + ) + + @functools.cache def _jit_sparse_module( item_size_bytes: int, @@ -46,7 +201,7 @@ def _jit_sparse_module( return load_jit( "sparse_cache", *cache_args, - cuda_files=["hisparse.cuh"], + cuda_files=["kvcacheio/hisparse.cuh"], cuda_wrappers=[ ( "load_cache_to_device_buffer", @@ -70,7 +225,7 @@ def _jit_copy_planned_module( is_mla, is_dsv4_layout, skip_io, - cuda_files=["hisparse.cuh"], + cuda_files=["kvcacheio/hisparse.cuh"], cuda_wrappers=[ ( "copy_cache_planned", @@ -86,7 +241,7 @@ def _jit_dsv4_transfer_module(block_size: int) -> Module: return load_jit( "sparse_cache_dsv4_transfer", block_size, - cuda_files=["hisparse.cuh"], + cuda_files=["kvcacheio/hisparse.cuh"], cuda_wrappers=[ ( "transfer_cache_dsv4_mla", diff --git a/test/registered/jit/benchmark/bench_hisparse_spec.py b/test/registered/jit/benchmark/bench_hisparse_spec.py new file mode 100644 index 000000000..941a222cf --- /dev/null +++ b/test/registered/jit/benchmark/bench_hisparse_spec.py @@ -0,0 +1,359 @@ +from __future__ import annotations + +from typing import NamedTuple + +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.kvcache.hisparse import ( + HiSparseSpecState, + copy_cache_planned_mla, + load_cache_to_device_buffer_mla, + load_cache_to_device_buffer_spec_mla, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=45, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +DEVICE = "cuda" +NUM_STEPS = 4 +NUM_INDEX_LAYERS = 4 +NUM_SHARED_LAYERS = 3 +TOP_K = 2048 +HOT_BUFFER_SIZE = 4096 +PAGE_SIZE = 64 +ITEM_WORDS = 72 +ITEM_SIZE_BYTES = ITEM_WORDS * 8 +MISS_COUNT_PER_STEP = 196 +UNIQUE_MISS_COUNT = 782 +# CUDA Graph benchmarking replays 100 swap calls per graph. Keep every replay +# on a fresh miss range while staying inside the native GLM-5.2 context length. +MISS_ADVANCE = 800 +SEQ_LEN = 1_048_576 +TOKEN_SCALE = 1_000_003 + + +class _BenchmarkState(NamedTuple): + 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_buffers: torch.Tensor + req_pool_indices: torch.Tensor + seq_lens: torch.Tensor + num_real_reqs: torch.Tensor + lru_slots: torch.Tensor + spec_out: torch.Tensor + lru_out: torch.Tensor + miss_src: torch.Tensor + miss_dst: torch.Tensor + miss_count: torch.Tensor + swap_states: tuple[HiSparseSpecState, ...] + + +def _make_top_k_tokens(batch_size: int) -> torch.Tensor: + hit_count = TOP_K - MISS_COUNT_PER_STEP + steps = [] + next_miss = HOT_BUFFER_SIZE + duplicate_tokens = [] + for step in range(NUM_STEPS): + hits = torch.roll( + torch.arange(HOT_BUFFER_SIZE, dtype=torch.int32, device=DEVICE), + step * 137, + )[:hit_count] + if step < 2: + misses = torch.arange( + next_miss, + next_miss + MISS_COUNT_PER_STEP, + dtype=torch.int32, + device=DEVICE, + ) + duplicate_tokens.append(misses[step]) + next_miss += MISS_COUNT_PER_STEP + else: + unique_misses = torch.arange( + next_miss, + next_miss + MISS_COUNT_PER_STEP - 1, + dtype=torch.int32, + device=DEVICE, + ) + misses = torch.cat((duplicate_tokens[step - 2].view(1), unique_misses)) + next_miss += MISS_COUNT_PER_STEP - 1 + steps.append(torch.cat((hits, misses))) + top_k_tokens = torch.stack(steps).unsqueeze(0) + top_k_tokens = top_k_tokens.repeat(batch_size, 1, 1).contiguous() + assert torch.unique(top_k_tokens[0, :, -MISS_COUNT_PER_STEP:]).numel() == ( + UNIQUE_MISS_COUNT + ) + return top_k_tokens + + +def _make_cache_index(batch_size: int) -> torch.Tensor: + hash_size = 1 << (2 * HOT_BUFFER_SIZE - 1).bit_length() + cache_index = torch.full( + (batch_size, 2, hash_size), -1, dtype=torch.int64, device=DEVICE + ) + tokens = torch.arange(HOT_BUFFER_SIZE, dtype=torch.int64, device=DEVICE) + hash_slots = ((tokens * 2654435761) & (hash_size - 1)).to(torch.long) + cache_index[:, 0, hash_slots] = (tokens << 32) | tokens + return cache_index + + +def _build_state(batch_size: int) -> _BenchmarkState: + buffer_size = HOT_BUFFER_SIZE + PAGE_SIZE + scratch_size = HOT_BUFFER_SIZE + physical_tokens_per_req = buffer_size + scratch_size + + top_k_tokens = _make_top_k_tokens(batch_size) + device_buffer_tokens = torch.full( + (batch_size, buffer_size), -1, dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens[:, :HOT_BUFFER_SIZE] = torch.arange( + HOT_BUFFER_SIZE, dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens = ( + device_buffer_tokens.unsqueeze(0).repeat(NUM_INDEX_LAYERS, 1, 1).contiguous() + ) + request_bases = ( + torch.arange(batch_size, dtype=torch.int32, device=DEVICE).view(-1, 1) + * physical_tokens_per_req + ) + device_buffer_locs = ( + request_bases + + torch.arange(buffer_size, dtype=torch.int32, device=DEVICE).view(1, -1) + ).contiguous() + scratch_locs = ( + request_bases + + buffer_size + + torch.arange(scratch_size, dtype=torch.int32, device=DEVICE).view(1, -1) + ).contiguous() + + host_cache_locs = torch.arange(SEQ_LEN, dtype=torch.int64, device=DEVICE) + host_cache_locs = host_cache_locs.view(1, -1).repeat(batch_size, 1).contiguous() + host_cache = torch.empty((SEQ_LEN, ITEM_WORDS), dtype=torch.int64, pin_memory=True) + host_cache.copy_( + torch.arange(SEQ_LEN, dtype=torch.int64).view(-1, 1) * TOKEN_SCALE + + torch.arange(ITEM_WORDS, dtype=torch.int64).view(1, -1) + ) + device_buffers = torch.empty( + (NUM_INDEX_LAYERS, batch_size * physical_tokens_per_req, ITEM_WORDS), + dtype=torch.int64, + device=DEVICE, + ) + hot_locs = device_buffer_locs[:, :HOT_BUFFER_SIZE].to(torch.long) + hot_values = host_cache[:HOT_BUFFER_SIZE].to(DEVICE) + for device_buffer in device_buffers: + device_buffer[hot_locs] = hot_values + + total_occurrences = NUM_STEPS * TOP_K + swap_states = [] + for _ in range(NUM_INDEX_LAYERS): + scratch_state = torch.full( + (batch_size + 1, max(4 * batch_size, 5 * total_occurrences)), + -1, + dtype=torch.int32, + device=DEVICE, + ) + scratch_state[0].zero_() + swap_states.append( + HiSparseSpecState( + cache_index=_make_cache_index(batch_size), + cache_policy=torch.zeros( + (batch_size + 1, HOT_BUFFER_SIZE), + dtype=torch.int32, + device=DEVICE, + ), + scratch_locs=scratch_locs, + scratch_state=scratch_state, + ) + ) + return _BenchmarkState( + 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_buffers=device_buffers, + req_pool_indices=torch.arange(batch_size, dtype=torch.int64, device=DEVICE), + seq_lens=torch.full( + (batch_size * NUM_STEPS,), SEQ_LEN, dtype=torch.int32, device=DEVICE + ), + num_real_reqs=torch.tensor([batch_size], dtype=torch.int32, device=DEVICE), + lru_slots=torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1), + spec_out=torch.full( + (NUM_INDEX_LAYERS, *top_k_tokens.shape), + -1, + dtype=top_k_tokens.dtype, + device=DEVICE, + ), + lru_out=torch.full_like(top_k_tokens, -1), + miss_src=torch.full( + (batch_size, total_occurrences), -1, dtype=torch.int64, device=DEVICE + ), + miss_dst=torch.full( + (batch_size, total_occurrences), -1, dtype=torch.int32, device=DEVICE + ), + miss_count=torch.zeros(batch_size, dtype=torch.int32, device=DEVICE), + swap_states=tuple(swap_states), + ) + + +def _run_spec_layer( + state: _BenchmarkState, layer_idx: int, *, record_plan: bool = False +) -> None: + load_cache_to_device_buffer_spec_mla( + top_k_tokens=state.top_k_tokens, + device_buffer_tokens=state.device_buffer_tokens[layer_idx], + host_cache_locs=state.host_cache_locs, + device_buffer_locs=state.device_buffer_locs, + host_cache=state.host_cache, + device_buffer=state.device_buffers[layer_idx], + top_k_device_locs=state.spec_out[layer_idx], + req_pool_indices=state.req_pool_indices, + seq_lens=state.seq_lens, + state=state.swap_states[layer_idx], + num_real_reqs=state.num_real_reqs, + miss_src=state.miss_src if record_plan else None, + miss_dst=state.miss_dst if record_plan else None, + miss_count=state.miss_count if record_plan else None, + ) + + +def _run_four_full_layers(state: _BenchmarkState) -> None: + for layer_idx in range(NUM_INDEX_LAYERS): + _run_spec_layer(state, layer_idx) + + +def _run_planned_copies(state: _BenchmarkState) -> None: + for layer_idx in range(1, NUM_SHARED_LAYERS + 1): + copy_cache_planned_mla( + miss_src=state.miss_src, + miss_dst=state.miss_dst, + miss_count=state.miss_count, + num_real_reqs=state.num_real_reqs, + host_cache=state.host_cache, + device_buffer=state.device_buffers[layer_idx], + item_size_bytes=ITEM_SIZE_BYTES, + num_blocks=8, + ) + + +def _run_lru_step(state: _BenchmarkState, step: int) -> None: + batch_size = state.top_k_tokens.size(0) + load_cache_to_device_buffer_mla( + top_k_tokens=state.top_k_tokens[:, step], + device_buffer_tokens=state.device_buffer_tokens[0], + host_cache_locs=state.host_cache_locs, + device_buffer_locs=state.device_buffer_locs, + host_cache=state.host_cache, + device_buffer=state.device_buffers[0], + top_k_device_locs=state.lru_out[:, step], + req_pool_indices=state.req_pool_indices, + # All benchmark steps use the same logical length. Reuse a contiguous + # slice so the measured path does not allocate a temporary tensor. + seq_lens=state.seq_lens[:batch_size], + lru_slots=state.lru_slots, + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=TOP_K, + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=PAGE_SIZE, + block_size=1024, + num_real_reqs=state.num_real_reqs, + ) + + +def _assert_current_result(state: _BenchmarkState, impl: str) -> None: + torch.cuda.synchronize() + expected = state.top_k_tokens.to(torch.int64) * TOKEN_SCALE + if impl == "spec_4_full": + for layer_idx in range(NUM_INDEX_LAYERS): + actual = state.device_buffers[layer_idx, :, 0][ + state.spec_out[layer_idx].to(torch.long) + ] + torch.testing.assert_close(actual, expected) + else: + out = state.lru_out if impl == "lru_1_full" else state.spec_out[0] + actual = state.device_buffers[0, :, 0][out.to(torch.long)] + torch.testing.assert_close(actual, expected) + if impl == "spec_1_full_3_shared": + for bid in range(state.top_k_tokens.size(0)): + count = int(state.miss_count[bid].item()) + src = state.miss_src[bid, :count].to(torch.long) + dst = state.miss_dst[bid, :count].to(torch.long) + expected = state.host_cache[src.cpu(), 0].to(DEVICE) + for layer_idx in range(1, NUM_SHARED_LAYERS + 1): + torch.testing.assert_close( + state.device_buffers[layer_idx, :, 0][dst], + expected, + ) + + +def _check_initial_result(state: _BenchmarkState, impl: str) -> None: + if impl == "lru_1_full": + for step in range(NUM_STEPS): + _run_lru_step(state, step) + elif impl == "spec_1_full": + _run_spec_layer(state, 0) + elif impl == "spec_4_full": + _run_four_full_layers(state) + else: + _run_spec_layer(state, 0, record_plan=True) + _run_planned_copies(state) + _assert_current_result(state, impl) + + +@marker.parametrize("batch_size", [1, 2, 4, 8, 16, 32, 64], [1]) +@marker.benchmark( + "impl", + ["lru_1_full", "spec_1_full", "spec_4_full", "spec_1_full_3_shared"], +) +def benchmark(batch_size: int, impl: str): + state = _build_state(batch_size) + _check_initial_result(state, impl) + + if impl == "lru_1_full": + + def run() -> None: + state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE) + for step in range(NUM_STEPS): + _run_lru_step(state, step) + + elif impl == "spec_1_full": + + def run() -> None: + state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE) + _run_spec_layer(state, 0) + + elif impl == "spec_4_full": + + def run() -> None: + state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE) + _run_four_full_layers(state) + + else: + + def run() -> None: + state.top_k_tokens[..., -MISS_COUNT_PER_STEP:].add_(MISS_ADVANCE) + _run_spec_layer(state, 0, record_plan=True) + _run_planned_copies(state) + + result = marker.do_bench( + run, + use_cuda_graph=True, + warmup_iters=5, + replay_iters=40, + disable_log_bandwidth=True, + memory_args=None, + memory_output=None, + ) + _assert_current_result(state, impl) + return result + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/jit/test_hisparse_spec.py b/test/registered/jit/test_hisparse_spec.py new file mode 100644 index 000000000..7f46073a2 --- /dev/null +++ b/test/registered/jit/test_hisparse_spec.py @@ -0,0 +1,550 @@ +from __future__ import annotations + +from typing import NamedTuple + +import pytest +import torch + +from sglang.kernels.ops.kvcache.hisparse import ( + HiSparseSpecState, + copy_cache_planned_mla, + load_cache_to_device_buffer_spec_mla, +) +from sglang.srt.utils import is_npu, is_xpu +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or is_npu() or is_xpu(), + reason="HiSparse speculative swap tests require a CUDA GPU.", +) + +DEVICE = "cuda" +TOKEN_SCALE = 1_000_003 + + +class _SwapState(NamedTuple): + device_buffer_tokens: torch.Tensor + device_buffer_locs: torch.Tensor + host_cache_locs: torch.Tensor + host_cache: torch.Tensor + device_buffer: torch.Tensor + swap_state: HiSparseSpecState + + +def _make_cache_index(num_reqs: int, hot_buffer_size: int) -> torch.Tensor: + hash_size = 1 << (2 * hot_buffer_size - 1).bit_length() + cache_index = torch.full( + (num_reqs, 2, hash_size), -1, dtype=torch.int64, device=DEVICE + ) + tokens = torch.arange(hot_buffer_size, dtype=torch.int64, device=DEVICE) + hash_slots = ((tokens * 2654435761) & (hash_size - 1)).to(torch.long) + packed_entries = (tokens << 32) | tokens + cache_index[:, 0, hash_slots] = packed_entries + return cache_index + + +def _make_state( + *, + num_reqs: int, + hot_buffer_size: int, + page_size: int, + scratch_size: int, + seq_len: int, + item_words: int, + metadata_occurrences: int, +) -> _SwapState: + buffer_size = hot_buffer_size + page_size + device_buffer_tokens = torch.full( + (num_reqs, buffer_size), -1, dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens[:, :hot_buffer_size] = torch.arange( + hot_buffer_size, dtype=torch.int32, device=DEVICE + ) + + physical_tokens_per_req = buffer_size + scratch_size + request_bases = ( + torch.arange(num_reqs, dtype=torch.int32, device=DEVICE).view(-1, 1) + * physical_tokens_per_req + ) + device_buffer_locs = ( + request_bases + + torch.arange(buffer_size, dtype=torch.int32, device=DEVICE).view(1, -1) + ).contiguous() + scratch_locs = ( + request_bases + + buffer_size + + torch.arange(scratch_size, dtype=torch.int32, device=DEVICE).view(1, -1) + ).contiguous() + + host_cache_locs = torch.arange(seq_len, dtype=torch.int64, device=DEVICE) + host_cache_locs = host_cache_locs.view(1, -1).repeat(num_reqs, 1).contiguous() + host_cache = torch.empty((seq_len, item_words), dtype=torch.int64, pin_memory=True) + host_cache.copy_( + torch.arange(seq_len, dtype=torch.int64).view(-1, 1) * TOKEN_SCALE + + torch.arange(item_words, dtype=torch.int64).view(1, -1) + ) + + device_buffer = torch.full( + (num_reqs * physical_tokens_per_req, item_words), + -1, + dtype=torch.int64, + device=DEVICE, + ) + hot_locs = device_buffer_locs[:, :hot_buffer_size].to(torch.long) + device_buffer[hot_locs] = host_cache[:hot_buffer_size].to(DEVICE) + + scratch_state = torch.full( + (num_reqs + 1, max(4 * num_reqs, 5 * metadata_occurrences)), + -1, + dtype=torch.int32, + device=DEVICE, + ) + scratch_state[0].zero_() + swap_state = HiSparseSpecState( + cache_index=_make_cache_index(num_reqs, hot_buffer_size), + cache_policy=torch.zeros( + (num_reqs + 1, hot_buffer_size), + dtype=torch.int32, + device=DEVICE, + ), + scratch_locs=scratch_locs, + scratch_state=scratch_state, + ) + return _SwapState( + device_buffer_tokens=device_buffer_tokens, + device_buffer_locs=device_buffer_locs, + host_cache_locs=host_cache_locs, + host_cache=host_cache, + device_buffer=device_buffer, + swap_state=swap_state, + ) + + +def _run_swap( + *, + top_k_tokens: torch.Tensor, + seq_lens: torch.Tensor, + state: _SwapState, + out: torch.Tensor | None = None, + req_pool_indices: torch.Tensor | None = None, + num_real_reqs: torch.Tensor | None = None, + miss_src: torch.Tensor | None = None, + miss_dst: torch.Tensor | None = None, + miss_count: torch.Tensor | None = None, +) -> torch.Tensor: + if out is None: + out = torch.full_like(top_k_tokens, -1) + else: + out.fill_(-1) + num_reqs = top_k_tokens.size(0) + if req_pool_indices is None: + req_pool_indices = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE) + if num_real_reqs is None: + num_real_reqs = torch.tensor([num_reqs], dtype=torch.int32, device=DEVICE) + load_cache_to_device_buffer_spec_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=state.device_buffer_tokens, + host_cache_locs=state.host_cache_locs, + device_buffer_locs=state.device_buffer_locs, + host_cache=state.host_cache, + device_buffer=state.device_buffer, + top_k_device_locs=out, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + state=state.swap_state, + num_real_reqs=num_real_reqs, + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + ) + return out + + +def _assert_output_matches_tokens( + state: _SwapState, out: torch.Tensor, tokens: torch.Tensor +) -> None: + actual = state.device_buffer[out.to(torch.long)] + expected = tokens.to(torch.int64).unsqueeze(-1) * TOKEN_SCALE + torch.arange( + state.device_buffer.size(-1), dtype=torch.int64, device=DEVICE + ) + torch.testing.assert_close(actual, expected) + + +class TestHiSparseSpec(CustomTestCase): + def test_deduplicates_repeated_misses_and_copies_full_items(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k, item_words = 4, 2048, 72 + total_occurrences = num_steps * top_k + state = _make_state( + num_reqs=1, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=hot_size, + seq_len=16384, + item_words=item_words, + metadata_occurrences=total_occurrences, + ) + + miss_count = 196 + hits = torch.arange(top_k - miss_count, dtype=torch.int32, device=DEVICE) + shared_misses = hot_size + torch.arange( + miss_count, dtype=torch.int32, device=DEVICE + ) + step = torch.cat((hits, shared_misses)) + top_k_tokens = step.view(1, 1, -1).repeat(1, num_steps, 1) + seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE) + + out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state) + torch.cuda.synchronize() + + _assert_output_matches_tokens(state, out, top_k_tokens) + self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), miss_count) + repeated_miss_locs = out[0, :, -miss_count:] + self.assertTrue(torch.all(repeated_miss_locs == repeated_miss_locs[0]).item()) + + def test_copies_782_cross_step_unique_misses(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k, item_words = 4, 2048, 72 + total_occurrences = num_steps * top_k + state = _make_state( + num_reqs=1, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=hot_size, + seq_len=16384, + item_words=item_words, + metadata_occurrences=total_occurrences, + ) + + steps = [] + next_miss = hot_size + for step_idx, miss_count in enumerate((196, 196, 195, 195)): + hits = torch.roll( + torch.arange(hot_size, dtype=torch.int32, device=DEVICE), + step_idx * 137, + )[: top_k - miss_count] + misses = torch.arange( + next_miss, + next_miss + miss_count, + dtype=torch.int32, + device=DEVICE, + ) + next_miss += miss_count + steps.append(torch.cat((hits, misses))) + top_k_tokens = torch.stack(steps).unsqueeze(0).contiguous() + seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE) + out = _run_swap( + top_k_tokens=top_k_tokens, + seq_lens=seq_lens, + state=state, + ) + torch.cuda.synchronize() + + _assert_output_matches_tokens(state, out, top_k_tokens) + self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 782) + + def test_records_union_plan_for_shared_layer_io(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k, item_words = 4, 2048, 72 + total_occurrences = num_steps * top_k + state = _make_state( + num_reqs=1, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=hot_size, + seq_len=16384, + item_words=item_words, + metadata_occurrences=total_occurrences, + ) + + steps = [] + next_miss = hot_size + for step_idx, step_miss_count in enumerate((196, 196, 195, 195)): + hits = torch.roll( + torch.arange(hot_size, dtype=torch.int32, device=DEVICE), + step_idx * 137, + )[: top_k - step_miss_count] + misses = torch.arange( + next_miss, + next_miss + step_miss_count, + dtype=torch.int32, + device=DEVICE, + ) + next_miss += step_miss_count + steps.append(torch.cat((hits, misses))) + top_k_tokens = torch.stack(steps).unsqueeze(0).contiguous() + seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE) + miss_src = torch.full( + (1, total_occurrences), -1, dtype=torch.int64, device=DEVICE + ) + miss_dst = torch.full( + (1, total_occurrences), -1, dtype=torch.int32, device=DEVICE + ) + miss_count = torch.full((1,), -1, dtype=torch.int32, device=DEVICE) + + _run_swap( + top_k_tokens=top_k_tokens, + seq_lens=seq_lens, + state=state, + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + ) + shared_layer_buffer = torch.full_like(state.device_buffer, -1) + copy_cache_planned_mla( + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + num_real_reqs=torch.ones(1, dtype=torch.int32, device=DEVICE), + host_cache=state.host_cache, + device_buffer=shared_layer_buffer, + item_size_bytes=state.host_cache.stride(0) + * state.host_cache.element_size(), + ) + torch.cuda.synchronize() + + self.assertEqual(int(miss_count.item()), 782) + count = int(miss_count.item()) + src = miss_src[0, :count].to(torch.long) + dst = miss_dst[0, :count].to(torch.long) + torch.testing.assert_close( + shared_layer_buffer[dst], state.host_cache[src.cpu()].to(DEVICE) + ) + torch.testing.assert_close(shared_layer_buffer[dst], state.device_buffer[dst]) + + def test_padded_request_clears_stale_plan_count(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k = 4, 2048 + total_occurrences = num_steps * top_k + state = _make_state( + num_reqs=2, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=hot_size, + seq_len=8192, + item_words=1, + metadata_occurrences=total_occurrences, + ) + top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).view( + 1, 1, -1 + ) + top_k_tokens = top_k_tokens.repeat(2, num_steps, 1).contiguous() + seq_lens = torch.full((2 * num_steps,), 8192, dtype=torch.int32, device=DEVICE) + miss_src = torch.full( + (2, total_occurrences), -1, dtype=torch.int64, device=DEVICE + ) + miss_dst = torch.full( + (2, total_occurrences), -1, dtype=torch.int32, device=DEVICE + ) + miss_count = torch.full((2,), 123, dtype=torch.int32, device=DEVICE) + + _run_swap( + top_k_tokens=top_k_tokens, + seq_lens=seq_lens, + state=state, + num_real_reqs=torch.ones(1, dtype=torch.int32, device=DEVICE), + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + ) + torch.cuda.synchronize() + + self.assertEqual(int(miss_count[1].item()), 0) + + def test_resolves_all_speculative_extra_page_slots_without_host_io(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k = 4, 2048 + total_occurrences = num_steps * top_k + seq_len = 8192 + state = _make_state( + num_reqs=1, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=hot_size, + seq_len=seq_len, + item_words=72, + metadata_occurrences=total_occurrences, + ) + + draft_tokens = torch.arange( + seq_len - num_steps, seq_len, dtype=torch.int32, device=DEVICE + ) + extra_offsets = torch.tensor([0, 7, 31, 63], device=DEVICE) + extra_locs = state.device_buffer_locs[0, hot_size + extra_offsets].to( + torch.long + ) + state.device_buffer_tokens[0, hot_size + extra_offsets] = draft_tokens + state.device_buffer[extra_locs] = state.host_cache[ + draft_tokens.to(device="cpu", dtype=torch.long) + ].to(DEVICE) + state.host_cache_locs[0, draft_tokens.to(torch.long)] = -1 + + hits = torch.arange(top_k - 1, dtype=torch.int32, device=DEVICE) + top_k_tokens = torch.stack( + [torch.cat((hits, draft_tokens[step : step + 1])) for step in range(4)] + ).unsqueeze(0) + seq_lens = draft_tokens + 1 + + out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state) + torch.cuda.synchronize() + + _assert_output_matches_tokens(state, out, top_k_tokens) + torch.testing.assert_close(out[0, :, -1].to(torch.long), extra_locs) + self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 0) + + def test_full_union_overflow_preserves_all_8192_outputs(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k = 4, 2048 + total_occurrences = num_steps * top_k + state = _make_state( + num_reqs=1, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=total_occurrences - hot_size, + seq_len=16384, + item_words=72, + metadata_occurrences=total_occurrences, + ) + top_k_tokens = ( + hot_size + torch.arange(total_occurrences, dtype=torch.int32, device=DEVICE) + ).view(1, num_steps, top_k) + seq_lens = torch.full((num_steps,), 16384, dtype=torch.int32, device=DEVICE) + + miss_src = torch.full( + (1, total_occurrences), -1, dtype=torch.int64, device=DEVICE + ) + miss_dst = torch.full( + (1, total_occurrences), -1, dtype=torch.int32, device=DEVICE + ) + miss_count = torch.full((1,), -1, dtype=torch.int32, device=DEVICE) + out = _run_swap( + top_k_tokens=top_k_tokens, + seq_lens=seq_lens, + state=state, + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + ) + torch.cuda.synchronize() + + _assert_output_matches_tokens(state, out, top_k_tokens) + self.assertEqual(torch.unique(out).numel(), total_occurrences) + self.assertEqual( + int(state.swap_state.scratch_state[0, 0].item()), total_occurrences + ) + self.assertEqual(int(miss_count.item()), total_occurrences) + self.assertTrue(miss_src.ge(0).all().item()) + self.assertTrue(miss_dst.ge(0).all().item()) + + def test_packed_ring_supports_glm52_native_context_length(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k = 4, 2048 + total_occurrences = num_steps * top_k + seq_len = 1_048_648 + state = _make_state( + num_reqs=1, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=hot_size, + seq_len=seq_len, + item_words=1, + metadata_occurrences=total_occurrences, + ) + top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).view( + 1, 1, -1 + ) + top_k_tokens = top_k_tokens.repeat(1, num_steps, 1) + high_token = seq_len - 1 + top_k_tokens[:, :, -1] = high_token + seq_lens = torch.full((num_steps,), seq_len, dtype=torch.int32, device=DEVICE) + + out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state) + torch.cuda.synchronize() + + _assert_output_matches_tokens(state, out, top_k_tokens) + self.assertTrue(out.ge(0).all().item()) + + # The first call admits the high token into the packed hash. The + # second call must resolve it as a hot hit rather than truncating the + # packed int64 entry and repeating Host-to-GPU IO. + out = _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state) + torch.cuda.synchronize() + _assert_output_matches_tokens(state, out, top_k_tokens) + self.assertEqual(int(state.swap_state.scratch_state[0, 0].item()), 0) + + def test_cuda_graph_replay_preserves_valid_locations(self) -> None: + hot_size, page_size = 4096, 64 + num_steps, top_k = 4, 2048 + total_occurrences = num_steps * top_k + state = _make_state( + num_reqs=1, + hot_buffer_size=hot_size, + page_size=page_size, + scratch_size=hot_size, + seq_len=65536, + item_words=72, + metadata_occurrences=total_occurrences, + ) + top_k_tokens = torch.arange(top_k, dtype=torch.int32, device=DEVICE).repeat( + num_steps, 1 + ) + for step, miss_count in enumerate((164, 102, 61, 20)): + top_k_tokens[step, -miss_count:] = torch.arange( + 8192 + step * top_k, + 8192 + step * top_k + miss_count, + dtype=torch.int32, + device=DEVICE, + ) + top_k_tokens = top_k_tokens.unsqueeze(0).contiguous() + seq_lens = torch.tensor( + [65533, 65534, 65535, 65536], dtype=torch.int32, device=DEVICE + ) + + _run_swap(top_k_tokens=top_k_tokens, seq_lens=seq_lens, state=state) + torch.cuda.synchronize() + graph_out = torch.full_like(top_k_tokens, -1) + req_pool_indices = torch.arange(1, dtype=torch.int64, device=DEVICE) + num_real_reqs = torch.tensor([1], dtype=torch.int32, device=DEVICE) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _run_swap( + top_k_tokens=top_k_tokens, + seq_lens=seq_lens, + state=state, + out=graph_out, + req_pool_indices=req_pool_indices, + num_real_reqs=num_real_reqs, + ) + + for _ in range(4): + graph.replay() + torch.cuda.synchronize() + + _assert_output_matches_tokens(state, graph_out, top_k_tokens) + self.assertTrue(graph_out.ge(0).all().item()) + + def test_rejects_invalid_step_shape_before_compilation(self) -> None: + state = _make_state( + num_reqs=1, + hot_buffer_size=4096, + page_size=64, + scratch_size=4096, + seq_len=8192, + item_words=1, + metadata_occurrences=8192, + ) + with self.assertRaisesRegex(ValueError, "2-4 steps"): + _run_swap( + top_k_tokens=torch.zeros( + (1, 1, 2048), dtype=torch.int32, device=DEVICE + ), + seq_lens=torch.tensor([8192], dtype=torch.int32, device=DEVICE), + state=state, + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v", "-s"]))