diff --git a/docs/docs/advanced_features/hisparse_guide.mdx b/docs/docs/advanced_features/hisparse_guide.mdx index f3dc2c321..323ba28d9 100644 --- a/docs/docs/advanced_features/hisparse_guide.mdx +++ b/docs/docs/advanced_features/hisparse_guide.mdx @@ -118,6 +118,12 @@ Pass as a JSON string via `--hisparse-config`: Example: `--hisparse-config='{"top_k": 2048, "device_buffer_size": 6144, "host_to_device_ratio": 10, "swap_in_block_size": 960}'` +### Shared-index prefetch (automatic) + +When a model reuses one anchor layer's top-k selection across a run of subsequent "skip" layers (DSA `index_topk_freq` / `index_topk_pattern`; native in GLM-5.2 as IndexShare), the working set of every skip layer is known the moment the anchor's index is computed. HiSparse exploits this automatically: the anchor's swap-in kernel records its miss plan (which host slots go to which device-buffer slots), and each skip layer replays that plan with a copy-only kernel issued ahead on a side stream, so the skip layers' host→device IO overlaps the intervening layers' compute instead of sitting on the decode critical path. The replay kernel uses a small fixed grid to keep its SM footprint low while overlapped. + +The prefetch is enabled automatically for eligible models (no pipeline parallelism, no speculative decoding) and can be turned off for A/B comparison with `SGLANG_DISABLE_HISPARSE_PREFETCH=1`. + ## Deployment HiSparse currently requires **PD disaggregation mode** and is enabled only on the **decode instance**. diff --git a/python/sglang/kernels/jit/csrc/hisparse.cuh b/python/sglang/kernels/jit/csrc/hisparse.cuh index 815378ae7..33bad3ceb 100644 --- a/python/sglang/kernels/jit/csrc/hisparse.cuh +++ b/python/sglang/kernels/jit/csrc/hisparse.cuh @@ -174,6 +174,55 @@ __device__ __forceinline__ int popc_mask(BallotMask mask) { #endif } +// Copy one missed item host->device with one warp. Shared by the fused swap-in +// kernel and copy_cache_planned_kernel so the layout dispatch cannot drift. +template +__device__ __forceinline__ void copy_miss_item( + int32_t lane_id, + const void* __restrict__ host_cache_k, + const void* __restrict__ host_cache_v, + void* __restrict__ device_buffer_k, + void* __restrict__ device_buffer_v, + int64_t src_loc, + int64_t dst_loc, + int64_t item_size_bytes) { + static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); + if constexpr (IsDsv4Layout) { +#ifdef USE_ROCM + // ROCm path: host cache and device buffer both use the page-padded C4 + // layout (same as the write path and the CUDA branch). We can't reuse + // device::hisparse::transfer_item here because its warp logic is hardcoded + // to a 32-lane warp; on wavefront64 we use transfer_dsv4_item_warp, which + // moves the value and the scale in one warp-width-agnostic copy. + using namespace device::hisparse; + const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(device_buffer_k, static_cast(dst_loc)); + const auto [src_value_ptr, src_scale_ptr] = + get_pointer_paged(const_cast(host_cache_k), static_cast(src_loc)); + transfer_dsv4_item_warp(lane_id, src_value_ptr, src_scale_ptr, dst_value_ptr, dst_scale_ptr); +#else + // CUDA path: page-padded device layout + page-padded host layout, K-only. + // The host cache is pinned DRAM but uses the same row layout as the GPU C4 + // cache, so use the page-padded address calculation for both ends. + device::hisparse::transfer_item( + /*dst_cache=*/device_buffer_k, + /*src_cache=*/const_cast(host_cache_k), + /*dst_index=*/static_cast(dst_loc), + /*src_index=*/static_cast(src_loc)); +#endif + } else { + // Generic path: device + host both linear, stride = item_size_bytes. + const auto src_k = static_cast(host_cache_k) + src_loc * item_size_bytes; + auto dst_k = static_cast(device_buffer_k) + dst_loc * item_size_bytes; + transfer_item_warp(lane_id, src_k, dst_k, item_size_bytes); + + if constexpr (!IsMLA) { + const auto src_v = static_cast(host_cache_v) + src_loc * item_size_bytes; + auto dst_v = static_cast(device_buffer_v) + dst_loc * item_size_bytes; + transfer_item_warp(lane_id, src_v, dst_v, item_size_bytes); + } + } +} + template __global__ __launch_bounds__(BLOCK_SIZE, 1) void transfer_cache_dsv4_mla_kernel( void** src_caches, @@ -268,12 +317,20 @@ struct SmemLayout { // IsDsv4Layout selects the miss-copy addressing: // false -> generic byte-stride: device + host both linear, stride = item_size_bytes // true -> DSv4 page-padded device + page-padded host (kvcacheio.cuh constants) +// +// RecordMissPlan records this step's miss plan (miss_src/dst = host/device loc +// per miss, miss_count per request) for shared-index skip layers to replay via +// copy_cache_planned_kernel. SkipIO elides only the KV byte movement (timing +// probe; output is garbage). Both are compile-time flags so the production +// (false, false) instantiation stays byte-identical. template < int BLOCK_SIZE, int NUM_TOP_K, int HOT_BUFFER_SIZE, bool IsMLA, bool IsDsv4Layout, + bool RecordMissPlan, + bool SkipIO, typename SeqLensT, typename ReqPoolIndicesT> __global__ void load_cache_to_device_buffer_kernel( @@ -296,7 +353,11 @@ __global__ void load_cache_to_device_buffer_kernel( int64_t top_k_tokens_stride, int64_t top_k_device_locs_stride, int64_t page_size, - int64_t item_size_bytes) { + int64_t item_size_bytes, + int64_t* __restrict__ miss_src_out, + int32_t* __restrict__ miss_dst_out, + int32_t* __restrict__ miss_count_out, + int64_t plan_stride) { static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); // todo hisparse: support page wise sparsity constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; @@ -345,6 +406,12 @@ __global__ void load_cache_to_device_buffer_kernel( } req_top_k_device_locs[i] = device_loc; } + // Short sequences load nothing from host: an empty miss plan for this request. + if constexpr (RecordMissPlan) { + if (tid == 0) { + miss_count_out[bid] = 0; + } + } return; } @@ -554,11 +621,22 @@ __global__ void load_cache_to_device_buffer_kernel( 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; + // Record the plan where the eviction is decided so it cannot disagree + // with the copy phase; locs are layer-independent (lockstep buffers). + if constexpr (RecordMissPlan) { + miss_src_out[bid * plan_stride + miss_offset] = req_host_cache_locs[my_token]; + miss_dst_out[bid * plan_stride + miss_offset] = req_device_buffer_locs[evict_slot]; + } } } __syncthreads(); total_misses = NUM_TOP_K - s_total_hits - s_newest_hit; + if constexpr (RecordMissPlan) { + if (tid == 0) { + miss_count_out[bid] = total_misses; + } + } // Write back LRU order: evictables at front (LRU), hits at back (MRU). { const int total_evictable = HOT_BUFFER_SIZE - s_total_hits; @@ -596,51 +674,28 @@ __global__ void load_cache_to_device_buffer_kernel( } // each warp copies one miss directly, can be separated into a new kernel if parallelism is a concern - 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]; + if constexpr (!SkipIO) { + 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]); + 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) { -#ifdef USE_ROCM - // ROCm path: host cache and device buffer both use the page-padded C4 - // layout (same as the write path and the CUDA branch). We can't reuse - // device::hisparse::transfer_item here because its warp logic is hardcoded - // to a 32-lane warp; on wavefront64 we use transfer_dsv4_item_warp, which - // moves the value and the scale in one warp-width-agnostic copy. - using namespace device::hisparse; - const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(device_buffer_k, static_cast(dst_loc)); - const auto [src_value_ptr, src_scale_ptr] = - get_pointer_paged(const_cast(host_cache_k), static_cast(src_loc)); - transfer_dsv4_item_warp(lane_id, src_value_ptr, src_scale_ptr, dst_value_ptr, dst_scale_ptr); -#else - // CUDA path: page-padded device layout + page-padded host layout, K-only. - // The host cache is pinned DRAM but uses the same row layout as the GPU C4 - // cache, so use the page-padded address calculation for both ends. - device::hisparse::transfer_item( - /*dst_cache=*/device_buffer_k, - /*src_cache=*/const_cast(host_cache_k), - /*dst_index=*/static_cast(dst_loc), - /*src_index=*/static_cast(src_loc)); -#endif - } else { - // Generic path: device + host both linear, stride = item_size_bytes. - const auto src_k = static_cast(host_cache_k) + src_loc * item_size_bytes; - auto dst_k = static_cast(device_buffer_k) + dst_loc * item_size_bytes; - transfer_item_warp(lane_id, src_k, dst_k, item_size_bytes); - - if constexpr (!IsMLA) { - const auto src_v = static_cast(host_cache_v) + src_loc * item_size_bytes; - auto dst_v = static_cast(device_buffer_v) + dst_loc * item_size_bytes; - transfer_item_warp(lane_id, src_v, dst_v, item_size_bytes); - } + copy_miss_item( + lane_id, host_cache_k, host_cache_v, device_buffer_k, device_buffer_v, src_loc, dst_loc, item_size_bytes); } } } -template +template < + int BLOCK_SIZE, + int NUM_TOP_K, + int HOT_BUFFER_SIZE, + bool IsMLA, + bool IsDsv4Layout, + bool RecordMissPlan, + bool SkipIO> void load_cache_to_device_buffer( tvm::ffi::TensorView top_k_tokens, tvm::ffi::TensorView device_buffer_tokens, @@ -656,11 +711,22 @@ void load_cache_to_device_buffer( tvm::ffi::TensorView lru_slots, tvm::ffi::TensorView num_real_reqs, int64_t page_size, - int64_t item_size_bytes) { + int64_t item_size_bytes, + tvm::ffi::TensorView miss_src_out, + tvm::ffi::TensorView miss_dst_out, + tvm::ffi::TensorView miss_count_out) { using namespace host; const int64_t bs = top_k_tokens.shape()[0]; const int64_t host_stride = host_cache_locs.shape()[1]; + // Miss-plan side outputs; 0-dim sentinels when RecordMissPlan is false. + 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 (RecordMissPlan && miss_dst_out.strides()[0] != plan_stride) { + throw std::runtime_error("load_cache_to_device_buffer: miss_src/miss_dst row strides differ"); + } const int64_t buffer_stride_0 = device_buffer_tokens.strides()[0]; const int64_t lru_slot_stride_0 = lru_slots.strides()[0]; const int64_t top_k_tokens_stride = top_k_tokens.strides()[0]; @@ -697,7 +763,11 @@ void load_cache_to_device_buffer( top_k_tokens_stride, top_k_device_locs_stride, page_size, - item_size_bytes); + item_size_bytes, + miss_src_ptr, + miss_dst_ptr, + miss_count_ptr, + plan_stride); }; const auto seq_dtype = seq_lens.dtype(); @@ -713,6 +783,8 @@ void load_cache_to_device_buffer( HOT_BUFFER_SIZE, IsMLA, IsDsv4Layout, + RecordMissPlan, + SkipIO, int64_t, int64_t>, static_cast(seq_lens.data_ptr()), @@ -725,6 +797,8 @@ void load_cache_to_device_buffer( HOT_BUFFER_SIZE, IsMLA, IsDsv4Layout, + RecordMissPlan, + SkipIO, int64_t, int32_t>, static_cast(seq_lens.data_ptr()), @@ -737,6 +811,8 @@ void load_cache_to_device_buffer( HOT_BUFFER_SIZE, IsMLA, IsDsv4Layout, + RecordMissPlan, + SkipIO, int32_t, int64_t>, static_cast(seq_lens.data_ptr()), @@ -749,6 +825,8 @@ void load_cache_to_device_buffer( HOT_BUFFER_SIZE, IsMLA, IsDsv4Layout, + RecordMissPlan, + SkipIO, int32_t, int32_t>, static_cast(seq_lens.data_ptr()), @@ -756,4 +834,93 @@ void load_cache_to_device_buffer( } } +// Copy-only swap-in for shared-index skip layers: replays the anchor's recorded +// miss plan (no hit detection / LRU; the anchor's slot table stays valid). The +// small fixed grid (num_blocks) keeps the SM footprint low while overlapping +// compute on a side stream. SkipIO is the same probe as in the fused kernel. +template +__global__ __launch_bounds__(BLOCK_SIZE, 1) void copy_cache_planned_kernel( + const int64_t* __restrict__ miss_src_locs, + const int32_t* __restrict__ miss_dst_locs, + const int32_t* __restrict__ miss_counts, + const int32_t* __restrict__ num_real_reqs, + const void* __restrict__ host_cache_k, + const void* __restrict__ host_cache_v, + void* __restrict__ device_buffer_k, + void* __restrict__ device_buffer_v, + int64_t plan_stride, + int64_t item_size_bytes) { + constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; + const int lane_id = threadIdx.x % WARP_SIZE; + const int warp_global = blockIdx.x * NUM_WARPS + threadIdx.x / WARP_SIZE; + const int total_warps = gridDim.x * NUM_WARPS; + const int real = num_real_reqs[0]; + + // Warp-sized windows amortize the miss_counts loads; warps then round-robin + // the flattened (request, miss) space so a large sparse batch spreads over + // all warps (183us -> 29us at bs=100 with 2 misses/req on H200) while one + // request's miss burst still uses every warp. + int start = 0; // flat index of the current request's first miss + for (int base = 0; base < real; base += WARP_SIZE) { + const int r_lane = base + lane_id; + const int cnt_lane = (r_lane < real) ? miss_counts[r_lane] : 0; + const int window = (real - base < WARP_SIZE) ? (real - base) : WARP_SIZE; + for (int j = 0; j < window; ++j) { + const int cnt = __shfl_sync(FULL_WARP_MASK, cnt_lane, j); + if (cnt == 0) continue; + int m0 = (warp_global - start) % total_warps; + if (m0 < 0) m0 += total_warps; + const int64_t r = base + j; + const int64_t* src_row = miss_src_locs + r * plan_stride; + const int32_t* dst_row = miss_dst_locs + r * plan_stride; + for (int m = m0; m < cnt; m += total_warps) { + // Timing probe: the plan is still walked; only the bytes stay put. + if constexpr (SkipIO) continue; + copy_miss_item( + lane_id, + host_cache_k, + host_cache_v, + device_buffer_k, + device_buffer_v, + src_row[m], + static_cast(dst_row[m]), + item_size_bytes); + } + start += cnt; + } + } +} + +template +void copy_cache_planned( + tvm::ffi::TensorView miss_src_locs, + tvm::ffi::TensorView miss_dst_locs, + tvm::ffi::TensorView miss_counts, + tvm::ffi::TensorView num_real_reqs, + tvm::ffi::TensorView host_cache_k, + tvm::ffi::TensorView host_cache_v, + tvm::ffi::TensorView device_buffer_k, + tvm::ffi::TensorView device_buffer_v, + int64_t num_blocks, + int64_t item_size_bytes) { + using namespace host; + const int64_t plan_stride = miss_src_locs.strides()[0]; + if (miss_dst_locs.strides()[0] != plan_stride) { + throw std::runtime_error("copy_cache_planned: miss_src/miss_dst row strides differ"); + } + const auto device = LaunchKernel::resolve_device(miss_src_locs.device()); + LaunchKernel(num_blocks, BLOCK_SIZE, device)( + copy_cache_planned_kernel, + static_cast(miss_src_locs.data_ptr()), + static_cast(miss_dst_locs.data_ptr()), + static_cast(miss_counts.data_ptr()), + static_cast(num_real_reqs.data_ptr()), + host_cache_k.data_ptr(), + (IsMLA || host_cache_v.ndim() == 0) ? (const void*)nullptr : host_cache_v.data_ptr(), + device_buffer_k.data_ptr(), + (IsMLA || device_buffer_v.ndim() == 0) ? (void*)nullptr : device_buffer_v.data_ptr(), + plan_stride, + item_size_bytes); +} + } // namespace sglang diff --git a/python/sglang/kernels/ops/kvcache/hisparse.py b/python/sglang/kernels/ops/kvcache/hisparse.py index 0abdc438b..fc0305d2c 100644 --- a/python/sglang/kernels/ops/kvcache/hisparse.py +++ b/python/sglang/kernels/ops/kvcache/hisparse.py @@ -19,12 +19,29 @@ def _jit_sparse_module( hot_buffer_size: int, is_mla: bool = False, is_dsv4_layout: bool = False, + record_miss_plan: bool = False, + skip_io: bool = False, ) -> Module: + # record_miss_plan / skip_io are compile-time kernel flags; the + # (False, False) production instantiation stays byte-identical. template_args = make_cpp_args( - block_size, num_top_k, hot_buffer_size, is_mla, is_dsv4_layout + block_size, + num_top_k, + hot_buffer_size, + is_mla, + is_dsv4_layout, + record_miss_plan, + skip_io, ) cache_args = make_cpp_args( - item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla, is_dsv4_layout + item_size_bytes, + block_size, + num_top_k, + hot_buffer_size, + is_mla, + is_dsv4_layout, + record_miss_plan, + skip_io, ) return load_jit( "sparse_cache", @@ -39,6 +56,30 @@ def _jit_sparse_module( ) +@functools.cache +def _jit_copy_planned_module( + block_size: int, + is_mla: bool, + is_dsv4_layout: bool, + skip_io: bool, +) -> Module: + template_args = make_cpp_args(block_size, is_mla, is_dsv4_layout, skip_io) + return load_jit( + "sparse_copy_planned", + block_size, + is_mla, + is_dsv4_layout, + skip_io, + cuda_files=["hisparse.cuh"], + cuda_wrappers=[ + ( + "copy_cache_planned", + f"copy_cache_planned<{template_args}>", + ) + ], + ) + + @functools.cache def _jit_dsv4_transfer_module(block_size: int) -> Module: template_args = make_cpp_args(block_size) @@ -91,11 +132,16 @@ def _load_cache_to_device_buffer_mla( page_size: int, block_size: int, num_real_reqs: torch.Tensor | None, + miss_src: torch.Tensor | None, + miss_dst: torch.Tensor | None, + miss_count: torch.Tensor | None, + skip_io: bool, ) -> None: assert ( hot_buffer_size >= num_top_k ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + record_miss_plan = miss_src is not None module = _jit_sparse_module( item_size_bytes, block_size, @@ -103,6 +149,8 @@ def _load_cache_to_device_buffer_mla( hot_buffer_size, is_mla=True, is_dsv4_layout=is_dsv4_layout, + record_miss_plan=record_miss_plan, + skip_io=skip_io, ) empty = torch.empty(0) @@ -112,6 +160,16 @@ def _load_cache_to_device_buffer_mla( [top_k_tokens.size(0)], dtype=torch.int32, device=top_k_tokens.device ) + if record_miss_plan: + assert miss_dst is not None and miss_count is not None + assert miss_src.dtype == torch.int64 and miss_dst.dtype == torch.int32 + assert miss_count.dtype == torch.int32 + # The kernel indexes both plan rows with one stride. + assert miss_src.stride(0) == miss_dst.stride(0) + else: + # Unused sentinels; the RecordMissPlan=false instantiation never reads them. + miss_src = miss_dst = miss_count = empty + module.load_cache_to_device_buffer( top_k_tokens, device_buffer_tokens, @@ -128,6 +186,9 @@ def _load_cache_to_device_buffer_mla( num_real_reqs, page_size, item_size_bytes, + miss_src, + miss_dst, + miss_count, ) @@ -148,8 +209,16 @@ def load_cache_to_device_buffer_mla( page_size: int = 1, block_size: int = 256, 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, + skip_io: bool = False, ) -> None: - """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes).""" + """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes). + + Optional miss_src/miss_dst/miss_count record the miss plan for replay by + copy_cache_planned_mla; skip_io elides only the KV bytes (timing probe). + """ _load_cache_to_device_buffer_mla( is_dsv4_layout=False, top_k_tokens=top_k_tokens, @@ -168,6 +237,47 @@ def load_cache_to_device_buffer_mla( page_size=page_size, block_size=block_size, num_real_reqs=num_real_reqs, + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + skip_io=skip_io, + ) + + +def copy_cache_planned_mla( + *, + miss_src: torch.Tensor, + miss_dst: torch.Tensor, + miss_count: torch.Tensor, + num_real_reqs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + item_size_bytes: int, + num_blocks: int = 4, + block_size: int = 1024, + is_dsv4_layout: bool = False, + skip_io: bool = False, +) -> None: + """Replay a recorded miss plan (host_cache -> device_buffer) for a skip layer. + + IO-only, no planning; the small fixed grid keeps the SM footprint low while + overlapped on a side stream. The anchor's slot table stays valid (lockstep). + """ + assert miss_src.dtype == torch.int64 and miss_dst.dtype == torch.int32 + assert miss_count.dtype == torch.int32 + module = _jit_copy_planned_module(block_size, True, is_dsv4_layout, skip_io) + empty = torch.empty(0) + module.copy_cache_planned( + miss_src, + miss_dst, + miss_count, + num_real_reqs, + host_cache, + empty, + device_buffer, + empty, + num_blocks, + item_size_bytes, ) @@ -188,6 +298,10 @@ def load_cache_to_device_buffer_dsv4_mla( page_size: int = 1, block_size: int = 256, 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, + skip_io: bool = False, ) -> None: """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout.""" _load_cache_to_device_buffer_mla( @@ -208,4 +322,8 @@ def load_cache_to_device_buffer_dsv4_mla( page_size=page_size, block_size=block_size, num_real_reqs=num_real_reqs, + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + skip_io=skip_io, ) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f5e7eaf8b..bf3c1eab3 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -846,6 +846,14 @@ class Envs: # Triton two_dot variant, 1.16-1.38x faster across GLM/DS shapes). SGLANG_OPT_Q8KV8_QPREP_VARIANT = EnvStr("auto") + # HiSparse + # Kill-switch for the shared-index (IndexShare) swap-in prefetch + # (auto-enabled for GLM-5.2-style DSA); set True to A/B synchronous swap-in. + SGLANG_DISABLE_HISPARSE_PREFETCH = EnvBool(False) + # Timing probe: run the swap-in fully but skip the host->device KV bytes, + # measuring the "IO is free" floor. GARBAGE OUTPUT -- benchmarking only. + SGLANG_DEBUG_HISPARSE_SKIP_IO = EnvBool(False) + # TRT-LLM-gen fused MoE (SiTU) via sglang JIT: path to an unpacked SiTU # cubin pool (cubins + flat ABI headers + overlay/; distributed as a # single downloadable archive). Needs the public flashinfer package diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 48bf5db43..430bd0692 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -1,14 +1,17 @@ # to be combined with the sparse coordinator class and sparse algorithm family import logging -from typing import List, NamedTuple, Union +from typing import Dict, List, NamedTuple, Optional, Tuple, Union import torch from sglang.kernels.ops.kvcache.hisparse import ( + copy_cache_planned_mla, load_cache_to_device_buffer_dsv4_mla, load_cache_to_device_buffer_mla, ) +from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa +from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, @@ -42,6 +45,69 @@ class HiSparseTokenStats(NamedTuple): host_token_usage: float +def resolve_shared_index_layers( + *, + hf_text_config, + pp_size: int, + is_speculative: bool, +) -> Optional[List[bool]]: + """Per-layer "reuses the previous layer's top-k index" pattern, or None. + + Mirrors DeepseekV2AttentionMLA's skip_topk derivation (index_topk_pattern / + index_topk_freq / cli_factor); None when the model has no sharing or the + prefetch cannot run (PP, speculative decoding, kill-switch). + """ + if not is_deepseek_dsa(hf_text_config): + return None + num_layers = hf_text_config.num_hidden_layers + cli_factor = getattr(hf_text_config, "cli_factor", 1) or 1 + if cli_factor > 1: + pattern = [i % cli_factor != 0 for i in range(num_layers)] + else: + pattern = [dsa_layer_skips_topk(hf_text_config, i) for i in range(num_layers)] + if not any(pattern): + return None + if pp_size != 1 or is_speculative: + logger.warning( + "HiSparse shared-index prefetch is unsupported under pipeline " + "parallelism / speculative decoding; falling back to synchronous " + "swap-in." + ) + return None + if envs.SGLANG_DISABLE_HISPARSE_PREFETCH.get(): + logger.info( + "HiSparse shared-index prefetch disabled via " + "SGLANG_DISABLE_HISPARSE_PREFETCH; using synchronous swap-in." + ) + return None + return pattern + + +def _build_prefetch_groups( + is_shared_index_layer: List[bool], +) -> Tuple[Dict[int, List[int]], List[int]]: + """Group consecutive shared-index (skip) layers under their anchor layer. + + Returns (groups, slot): anchor layer_id -> ordered skip layers, and each + skip layer's position in its group (indexes the per-slot prefetch events). + """ + groups: Dict[int, List[int]] = {} + slot = [0] * len(is_shared_index_layer) + anchor = None + for i, is_shared in enumerate(is_shared_index_layer): + if not is_shared: + anchor = i # compute layer; anchors the skip layers after it + continue + assert anchor is not None, ( + f"shared-index (skip) layer {i} has no preceding compute layer; " + "the model's index-topk pattern is invalid" + ) + group = groups.setdefault(anchor, []) + slot[i] = len(group) + group.append(i) + return groups, slot + + class HiSparseCoordinator: def __init__( self, @@ -56,6 +122,7 @@ class HiSparseCoordinator: tp_group, host_to_device_ratio: int = 2, swap_in_block_size: int = 960, + shared_index_layers: Optional[List[bool]] = None, ): self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator @@ -63,6 +130,9 @@ class HiSparseCoordinator: self.device_buffer_size = device_buffer_size self.device = device self.swap_in_block_size = swap_in_block_size + # Timing probe: skip the host->device KV bytes to measure the "IO is + # free" floor. Produces garbage output; benchmarking only. + self.skip_io = envs.SGLANG_DEBUG_HISPARSE_SKIP_IO.get() self.compress_ratio = self.token_to_kv_pool_allocator.compress_ratio self.is_dsv4_hisparse = isinstance( @@ -186,6 +256,65 @@ class HiSparseCoordinator: # staging already backed up all prefill tokens. Cleared after one step. self._skip_first_backup = [False] * max_num_req_slots + self._init_shared_index_prefetch( + shared_index_layers=shared_index_layers, + layer_num=layer_num, + max_num_req_slots=max_num_req_slots, + ) + + def _init_shared_index_prefetch( + self, + shared_index_layers: Optional[List[bool]], + layer_num: int, + max_num_req_slots: int, + ) -> None: + """Set up the plan-then-IO prefetch for shared-index (IndexShare) models: + the anchor's kernel records its miss plan and skip layers replay it on + `prefetch_stream`, overlapping their IO with the intervening compute.""" + if shared_index_layers is not None and len(shared_index_layers) != layer_num: + # Attention-layer count differs from num_hidden_layers (e.g. Longcat + # doubles it): pattern would be misindexed, fall back to synchronous. + logger.warning( + "HiSparse shared-index prefetch disabled: pattern length %d != " + "KV pool layer_num %d; using synchronous swap-in.", + len(shared_index_layers), + layer_num, + ) + shared_index_layers = None + self._is_shared_index_layer = list(shared_index_layers or [False] * layer_num) + self.enable_prefetch = any(self._is_shared_index_layer) + self._prefetch_groups, self._prefetch_slot = _build_prefetch_groups( + self._is_shared_index_layer + ) + if not self.enable_prefetch: + return + + # Small fixed grid for the copy-only kernel: low SM footprint so the + # copies overlap compute with little contention. + self._prefetch_copy_blocks = 4 + max_group_size = max(len(g) for g in self._prefetch_groups.values()) + self.prefetch_stream = device_module.Stream() + self._prefetch_events = [device_module.Event() for _ in range(max_group_size)] + # Plan recorded by the current anchor, replayed by its skip layers. One + # buffer set suffices: the last skip layer's event wait orders the next + # anchor's writes after this group's copies. + self._miss_src = torch.zeros( + (max_num_req_slots, self.top_k), dtype=torch.int64, device=self.device + ) + self._miss_dst = torch.zeros( + (max_num_req_slots, self.top_k), dtype=torch.int32, device=self.device + ) + self._miss_count = torch.zeros( + (max_num_req_slots,), dtype=torch.int32, device=self.device + ) + logger.info( + "HiSparse: shared-index prefetch (plan-then-IO) enabled; %d anchor " + "group(s), %d skip layer(s) of %d total.", + len(self._prefetch_groups), + sum(self._is_shared_index_layer), + layer_num, + ) + def set_decode_producer_stream(self, stream) -> None: self.decode_producer_stream = stream @@ -194,6 +323,9 @@ class HiSparseCoordinator: # See HostKVCache.destroy for why the explicit unregister matters. self.write_staging_stream.synchronize() self.decode_backup_stream.synchronize() + if self.enable_prefetch: + # Skip-layer copies read the pinned host pool on the prefetch stream. + self.prefetch_stream.synchronize() self.mem_pool_host.destroy() def get_token_stats(self) -> HiSparseTokenStats: @@ -802,16 +934,20 @@ class HiSparseCoordinator: self.lru_slots[:, req.req_pool_idx, :].copy_(self._lru_init) self._skip_first_backup[req.req_pool_idx] = False - def swap_in_selected_pages( + def _run_swap_in_kernel( self, req_pool_indices: torch.Tensor, compressed_seq_lens: torch.Tensor, top_k_result: torch.Tensor, layer_id: int, + record_plan: bool = False, ) -> torch.Tensor: - """Swap selected top-k tokens into device memory and return their indices.""" - num_reqs = req_pool_indices.size(0) + """Run the full plan+IO swap-in kernel for one layer; return its slot table. + record_plan (set on the anchor of a shared-index group) also records the + miss plan into self._miss_{src,dst,count} for the skip layers to replay. + """ + num_reqs = req_pool_indices.size(0) top_k_indices = self.top_k_device_locs_buffer[:num_reqs] swap_in_fn = ( @@ -819,6 +955,15 @@ class HiSparseCoordinator: if self.is_dsv4_hisparse else load_cache_to_device_buffer_mla ) + plan = ( + dict( + miss_src=self._miss_src[:num_reqs], + miss_dst=self._miss_dst[:num_reqs], + miss_count=self._miss_count[:num_reqs], + ) + if record_plan + else {} + ) swap_in_fn( top_k_tokens=top_k_result, device_buffer_tokens=self.req_device_buffer_tokens[layer_id], @@ -836,5 +981,70 @@ class HiSparseCoordinator: page_size=1, block_size=self.swap_in_block_size, num_real_reqs=self.num_real_reqs, + skip_io=self.skip_io, + **plan, ) return top_k_indices + + def _run_copy_only_kernel(self, num_reqs: int, skip_layer: int) -> None: + """Replay the anchor's recorded miss plan into a skip layer's buffers + (IO-only; the anchor's slot table stays valid -- lockstep layout).""" + copy_cache_planned_mla( + miss_src=self._miss_src[:num_reqs], + miss_dst=self._miss_dst[:num_reqs], + miss_count=self._miss_count[:num_reqs], + num_real_reqs=self.num_real_reqs, + host_cache=self.mem_pool_host.kv_buffer[skip_layer], + device_buffer=self.mem_pool_device.kv_buffer[skip_layer], + item_size_bytes=self.item_size_bytes, + num_blocks=self._prefetch_copy_blocks, + is_dsv4_layout=self.is_dsv4_hisparse, + skip_io=self.skip_io, + ) + + def swap_in_selected_pages( + self, + req_pool_indices: torch.Tensor, + compressed_seq_lens: torch.Tensor, + top_k_result: torch.Tensor, + layer_id: int, + ) -> torch.Tensor: + """Swap selected top-k tokens into device memory and return their indices. + + With prefetch enabled, anchors swap in synchronously (recording the miss + plan) and prefetch their skip layers' copies; skip layers just wait. + """ + if not self.enable_prefetch: + return self._run_swap_in_kernel( + req_pool_indices, compressed_seq_lens, top_k_result, layer_id + ) + + num_reqs = req_pool_indices.size(0) + if self._is_shared_index_layer[layer_id]: + # Skip layer: wait for its prefetched copy; the anchor's slot table + # applies (shared index + lockstep buffers). + slot = self._prefetch_slot[layer_id] + self._prefetch_events[slot].wait(device_module.current_stream()) + return self.top_k_device_locs_buffer[:num_reqs] + + # Anchor: swap in synchronously (recording the plan), then prefetch the + # skip layers' copies on the side stream. + group = self._prefetch_groups.get(layer_id) + anchor_locs = self._run_swap_in_kernel( + req_pool_indices, + compressed_seq_lens, + top_k_result, + layer_id, + record_plan=group is not None, + ) + if group: + # Fork: the prefetch stream must observe the anchor's plan (produced + # on the current stream) before replaying it. + self.prefetch_stream.wait_stream(device_module.current_stream()) + with device_module.stream(self.prefetch_stream): + for skip_layer in group: + self._run_copy_only_kernel(num_reqs, skip_layer) + self._prefetch_events[self._prefetch_slot[skip_layer]].record( + self.prefetch_stream + ) + return anchor_locs diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index fc51e70de..84bd42ffe 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -847,7 +847,10 @@ class ModelRunner: def maybe_init_hisparse_coordinator(self): if not self.enable_hisparse: return - from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator + from sglang.srt.managers.hisparse_coordinator import ( + HiSparseCoordinator, + resolve_shared_index_layers, + ) from sglang.srt.mem_cache.sparsity import parse_hisparse_config hisparse_cfg = parse_hisparse_config(self.server_args) @@ -867,6 +870,11 @@ class ModelRunner: ), host_to_device_ratio=hisparse_cfg.host_to_device_ratio, swap_in_block_size=hisparse_cfg.swap_in_block_size, + shared_index_layers=resolve_shared_index_layers( + hf_text_config=self.model_config.hf_text_config, + pp_size=self.ps.pp_size, + is_speculative=self.spec_algorithm.is_speculative(), + ), ) def post_capture_resize_kv_pool(self): diff --git a/test/manual/kernels/test_hisparse_prefetch.py b/test/manual/kernels/test_hisparse_prefetch.py new file mode 100644 index 000000000..12aad76aa --- /dev/null +++ b/test/manual/kernels/test_hisparse_prefetch.py @@ -0,0 +1,394 @@ +"""Extended tests for the HiSparse shared-index (plan-then-IO) prefetch. + +Local-only (not registered to CI): these cover the CUDA-graph capture/replay +pattern, the DSv4 page-padded layout, and the SGLANG_DEBUG_HISPARSE_SKIP_IO +probe, each of which JIT-compiles extra kernel instantiations. The cheap +plan-replay correctness guards run in CI via +test/registered/kernels/ops/kvcache/test_hisparse.py, which this file imports +its fixtures from. + +Run: python3 test/manual/kernels/test_hisparse_prefetch.py +""" + +import sys +from pathlib import Path + +import pytest +import torch + +from sglang.kernels.ops.kvcache.hisparse import ( + copy_cache_planned_mla, + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, +) +from sglang.srt.utils import is_hip + +sys.path.insert( + 0, + str( + Path(__file__).resolve().parents[2] + / "registered" + / "kernels" + / "ops" + / "kvcache" + ), +) +from test_hisparse import ( # noqa: F401 (GPU/platform guard applies here too); noqa: E402 + DEVICE, + DEVICE_CACHE_SIZE, + DSV4_ITEM_BYTES, + DSV4_PAGE_BYTES, + DSV4_PAGE_SIZE, + DTYPE, + HOST_CACHE_SIZE, + HOT_BUFFER_SIZE, + ITEM_SIZE_BYTES, + KV_DIM, + _host_cache, + _long_case, + _make_plan, + _make_state, + _run_kernel, + _write_dsv4_token, + pytestmark, +) + + +def test_plan_then_io_dsv4_matches_sync_swap_in() -> None: + """DSv4 layout: replaying the recorded plan lands the page-padded value+scale + bytes exactly where the fused swap-in copy puts them.""" + num_pages = 2 + state = _long_case() + plan_state = _long_case() + + def _dsv4_caches(): + host = torch.zeros( + (num_pages, DSV4_PAGE_BYTES), + dtype=torch.uint8, + device="cpu", + pin_memory=True, + ) + for token in range(HOST_CACHE_SIZE): + _write_dsv4_token(host, token, seed=token + 1) + dev = torch.full( + (num_pages, DSV4_PAGE_BYTES), 0xFF, dtype=torch.uint8, device=DEVICE + ) + return host, dev + + common = dict( + top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE), + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE), + 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), + ) + + # Reference: fused swap-in copies host token 6 into evict slot 0 (loc 9). + ref_host, ref_dev = _dsv4_caches() + out = torch.full((1, 1), -1, dtype=torch.int32, device=DEVICE) + load_cache_to_device_buffer_dsv4_mla( + device_buffer_tokens=state["device_buffer_tokens"], + host_cache_locs=state["host_cache_locs"], + device_buffer_locs=state["device_buffer_locs"], + host_cache=ref_host, + device_buffer=ref_dev, + top_k_device_locs=out, + lru_slots=state["lru_slots"], + **common, + ) + + # Anchor: same swap-in on a twin state, recording the plan. + miss_src, miss_dst, miss_count = _make_plan(1, 1) + anchor_host, anchor_dev = _dsv4_caches() + anchor_out = torch.full((1, 1), -1, dtype=torch.int32, device=DEVICE) + load_cache_to_device_buffer_dsv4_mla( + device_buffer_tokens=plan_state["device_buffer_tokens"], + host_cache_locs=plan_state["host_cache_locs"], + device_buffer_locs=plan_state["device_buffer_locs"], + host_cache=anchor_host, + device_buffer=anchor_dev, + top_k_device_locs=anchor_out, + lru_slots=plan_state["lru_slots"], + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + **common, + ) + + # Skip layer: replay the plan into a fresh buffer; must match the reference. + replay_host, replay_dev = _dsv4_caches() + copy_cache_planned_mla( + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + num_real_reqs=common["num_real_reqs"], + host_cache=replay_host, + device_buffer=replay_dev, + item_size_bytes=DSV4_ITEM_BYTES, + num_blocks=4, + is_dsv4_layout=True, + ) + torch.cuda.synchronize() + + assert torch.equal(miss_count.cpu(), torch.tensor([1], dtype=torch.int32)) + assert torch.equal(anchor_out.cpu(), out.cpu()) + assert torch.equal(replay_dev.cpu(), ref_dev.cpu()) + + +def test_skip_io_probe_plans_without_moving_bytes() -> None: + """skip_io still runs all planning (slot table, LRU, miss plan) but must + leave the device buffer untouched; replaying the plan then repairs it.""" + locs = [[9, 7, 3, 5, 11]] + toks = [[1, 4, 2, 5, -1]] + top_k = torch.tensor([[6, 4]], dtype=torch.int32, device=DEVICE) + nr, K = top_k.shape + + ref = _make_state(locs, toks, [7]) + ref_out = _run_kernel(top_k_tokens=top_k, seq_len=8, **ref) + + probe = _make_state(locs, toks, [7]) + probe_buffer_before = probe["device_buffer"].clone() + miss_src, miss_dst, miss_count = _make_plan(nr, K) + probe_out = _run_kernel( + top_k_tokens=top_k, + seq_len=8, + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + skip_io=True, + **probe, + ) + + # All planning outputs match the real run; only the bytes stayed put. + assert torch.equal(probe_out.cpu(), ref_out.cpu()) + assert torch.equal(probe["lru_slots"].cpu(), ref["lru_slots"].cpu()) + assert torch.equal(probe["device_buffer"].cpu(), probe_buffer_before.cpu()) + assert not torch.equal(probe["device_buffer"].cpu(), ref["device_buffer"].cpu()) + + copy_cache_planned_mla( + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + num_real_reqs=torch.tensor([nr], dtype=torch.int32, device=DEVICE), + host_cache=probe["host_cache"], + device_buffer=probe["device_buffer"], + item_size_bytes=ITEM_SIZE_BYTES, + num_blocks=4, + ) + torch.cuda.synchronize() + assert torch.equal(probe["device_buffer"].cpu(), ref["device_buffer"].cpu()) + + +_PIO_LAYERS = 4 # one anchor (layer 0) + three skip layers (GLM group of freq 4) +_PIO_REQS = 2 +_PIO_SEQ = 10 # > HOT_BUFFER_SIZE -> long path; newest token = 9 +_PIO_DBL = [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]] # per-req [4 LRU slots + newest] +_PIO_STEPS = [ + [[4, 5, 0, 9], [6, 7, 1, 9]], + [[4, 10, 2, 9], [8, 5, 3, 9]], + [[11, 4, 5, 9], [6, 12, 7, 9]], + [[0, 1, 2, 9], [3, 4, 5, 9]], +] + + +def _pio_fresh(host_cache, dbl): + """Fresh per-layer buffers/tokens/lru, all layers initialized identically.""" + buffers, dbt, lru = [], [], [] + init_tokens = [[0, 1, 2, 3, -1], [0, 1, 2, 3, -1]] + for _ in range(_PIO_LAYERS): + db = torch.full((DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE) + for rid in range(_PIO_REQS): + for slot, tok in enumerate(init_tokens[rid][:HOT_BUFFER_SIZE]): + db[dbl[rid, slot]].copy_(host_cache[tok].to(DEVICE)) + db[dbl[rid, HOT_BUFFER_SIZE]].copy_(host_cache[_PIO_SEQ - 1].to(DEVICE)) + buffers.append(db) + dbt.append(torch.tensor(init_tokens, dtype=torch.int32, device=DEVICE)) + lru.append( + torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(_PIO_REQS, 1) + .contiguous() + ) + torch.cuda.synchronize() + return buffers, dbt, lru + + +def _pio_swap_in(topk, dbt, lru, dbl, hcl, host, buffer, out, seq_lens, nrr, rpi, plan): + miss_src, miss_dst, miss_count = plan if plan else (None, None, None) + load_cache_to_device_buffer_mla( + top_k_tokens=topk, + device_buffer_tokens=dbt, + host_cache_locs=hcl, + device_buffer_locs=dbl, + host_cache=host, + device_buffer=buffer, + top_k_device_locs=out, + req_pool_indices=rpi, + seq_lens=seq_lens, + lru_slots=lru, + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=topk.shape[1], + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=1, + block_size=256, + num_real_reqs=nrr, + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + ) + + +def _pio_prefetch_step( + topk, buffers, dbt, lru, dbl, hcl, host, seq_lens, nrr, rpi, out, plan, side, events +): + """Anchor records the plan; skip layers replay it copy-only on a side stream.""" + miss_src, miss_dst, miss_count = plan + _pio_swap_in( + topk, dbt[0], lru[0], dbl, hcl, host, buffers[0], out, seq_lens, nrr, rpi, plan + ) + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for layer in range(1, _PIO_LAYERS): + copy_cache_planned_mla( + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + num_real_reqs=nrr, + host_cache=host, + device_buffer=buffers[layer], + item_size_bytes=ITEM_SIZE_BYTES, + num_blocks=4, + ) + events[layer].record(side) + for layer in range(1, _PIO_LAYERS): + events[layer].wait(torch.cuda.current_stream()) + + +def _pio_sync_step(topk, buffers, dbt, lru, dbl, hcl, host, seq_lens, nrr, rpi): + """Reference: run the full swap-in independently on every layer.""" + outs = [] + for layer in range(_PIO_LAYERS): + out = torch.full_like(topk, -1) + _pio_swap_in( + topk, + dbt[layer], + lru[layer], + dbl, + hcl, + host, + buffers[layer], + out, + seq_lens, + nrr, + rpi, + None, + ) + outs.append(out.clone()) + return outs + + +@pytest.mark.skipif(is_hip(), reason="CUDA graph capture test is CUDA-only.") +def test_plan_then_io_cuda_graph_replay() -> None: + """The plan-then-IO prefetch pattern captures into a CUDA graph and replays + bit-identically to the eager synchronous swap-in across multiple steps.""" + host = _host_cache() + dbl = torch.tensor(_PIO_DBL, dtype=torch.int32, device=DEVICE) + hcl = ( + torch.arange(HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE) + .view(1, -1) + .repeat(_PIO_REQS, 1) + .contiguous() + ) + seq_lens = torch.full((_PIO_REQS,), _PIO_SEQ, dtype=torch.int32, device=DEVICE) + nrr = torch.tensor([_PIO_REQS], dtype=torch.int32, device=DEVICE) + rpi = torch.arange(_PIO_REQS, dtype=torch.int64, device=DEVICE) + K = len(_PIO_STEPS[0][0]) + steps = [torch.tensor(s, dtype=torch.int32, device=DEVICE) for s in _PIO_STEPS] + + # Reference: full synchronous swap-in on every layer, snapshotted per step. + ref_buf, ref_dbt, ref_lru = _pio_fresh(host, dbl) + ref_slots, ref_snap = [], [] + for topk in steps: + ref_slots.append( + _pio_sync_step( + topk, ref_buf, ref_dbt, ref_lru, dbl, hcl, host, seq_lens, nrr, rpi + ) + ) + torch.cuda.synchronize() + ref_snap.append([b.clone() for b in ref_buf]) + torch.cuda.synchronize() + + # Graph-captured prefetch replayed step by step against a fixed topk buffer. + buf, dbt, lru = _pio_fresh(host, dbl) + topk_buf = torch.zeros((_PIO_REQS, K), dtype=torch.int32, device=DEVICE) + out = torch.full((_PIO_REQS, K), -1, dtype=torch.int32, device=DEVICE) + plan = _make_plan(_PIO_REQS, K) + side = torch.cuda.Stream() + events = [torch.cuda.Event() for _ in range(_PIO_LAYERS)] + + warm = torch.cuda.Stream() + warm.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warm): + topk_buf.copy_(steps[0]) + _pio_prefetch_step( + topk_buf, + buf, + dbt, + lru, + dbl, + hcl, + host, + seq_lens, + nrr, + rpi, + out, + plan, + side, + events, + ) + torch.cuda.current_stream().wait_stream(warm) + torch.cuda.synchronize() + + # Reset state mutated by warmup so capture starts from a clean identical state. + buf, dbt, lru = _pio_fresh(host, dbl) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _pio_prefetch_step( + topk_buf, + buf, + dbt, + lru, + dbl, + hcl, + host, + seq_lens, + nrr, + rpi, + out, + plan, + side, + events, + ) + torch.cuda.synchronize() + + for s, topk in enumerate(steps): + topk_buf.copy_(topk) + graph.replay() + torch.cuda.synchronize() + # Anchor slot table matches the synchronous layer-0 result. + assert torch.equal( + out.cpu(), ref_slots[s][0].cpu() + ), f"slots differ at step {s}" + # Every layer's device buffer stays bit-identical to synchronous swap-in. + for layer in range(_PIO_LAYERS): + assert torch.equal( + buf[layer].cpu(), ref_snap[s][layer].cpu() + ), f"buffer differs at step {s}, layer {layer}" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/kernels/benchmark/kvcache/bench_hisparse.py b/test/registered/kernels/benchmark/kvcache/bench_hisparse.py index 069ed2f88..17a768dcc 100644 --- a/test/registered/kernels/benchmark/kvcache/bench_hisparse.py +++ b/test/registered/kernels/benchmark/kvcache/bench_hisparse.py @@ -6,13 +6,16 @@ import triton import triton.testing from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE -from sglang.kernels.ops.kvcache.hisparse import load_cache_to_device_buffer_mla +from sglang.kernels.ops.kvcache.hisparse import ( + copy_cache_planned_mla, + load_cache_to_device_buffer_mla, +) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci( - est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" + est_time=18, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" ) -register_amd_ci(est_time=12, stage="jit-kernel-benchmark", runner_config="amd") +register_amd_ci(est_time=18, stage="jit-kernel-benchmark", runner_config="amd") DEVICE = DEFAULT_DEVICE DTYPE = DEFAULT_DTYPE @@ -159,6 +162,66 @@ def _time_kernel(batch_size: int, hot_buffer_size: int, miss_rate: float) -> flo return start.elapsed_time(end) * 1000.0 / ROUNDS +def _time_planned_copy( + batch_size: int, hot_buffer_size: int, miss_rate: float +) -> float: + """Time the copy-only replay used by shared-index skip layers: one anchor + swap-in records a real plan, each timed round replays it (num_blocks=4).""" + state = _build_inputs(batch_size, hot_buffer_size, miss_rate) + miss_src = torch.zeros((batch_size, TOP_K), dtype=torch.int64, device=DEVICE) + miss_dst = torch.zeros((batch_size, TOP_K), dtype=torch.int32, device=DEVICE) + miss_count = torch.zeros((batch_size,), dtype=torch.int32, device=DEVICE) + load_cache_to_device_buffer_mla( + top_k_tokens=state["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=state["top_k_device_locs"], + req_pool_indices=state["req_pool_indices"], + seq_lens=state["seq_lens"], + lru_slots=state["lru_slots"], + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=TOP_K, + hot_buffer_size=hot_buffer_size, + block_size=1024, + num_real_reqs=state["num_real_reqs"], + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + ) + torch.cuda.synchronize() + skip_layer_buffer = torch.empty_like(state["device_buffer"]) + + def run_once(): + copy_cache_planned_mla( + miss_src=miss_src, + miss_dst=miss_dst, + miss_count=miss_count, + num_real_reqs=state["num_real_reqs"], + host_cache=state["host_cache"], + device_buffer=skip_layer_buffer, + item_size_bytes=ITEM_SIZE_BYTES, + num_blocks=4, + ) + + run_once() + torch.cuda.synchronize() + for _ in range(WARMUP_ROUNDS): + run_once() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(ROUNDS): + run_once() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / ROUNDS + + @triton.testing.perf_report( triton.testing.Benchmark( x_names=["batch_size", "hot_buffer_size", "miss_rate", "miss_tokens_cnt"], @@ -188,5 +251,35 @@ def benchmark_latency( return avg_us, avg_us, avg_us +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "hot_buffer_size", "miss_rate", "miss_tokens_cnt"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="hisparse-planned-copy-latency", + args={}, + ) +) +def benchmark_planned_copy_latency( + batch_size: int, + hot_buffer_size: int, + miss_rate: float, + miss_tokens_cnt: int, + provider: str, +) -> Tuple[float, float, float]: + assert provider == "jit" + batch_size = int(batch_size) + hot_buffer_size = int(hot_buffer_size) + miss_rate = float(miss_rate) + assert miss_tokens_cnt == batch_size * _miss_tokens_per_req(miss_rate) + avg_us = _time_planned_copy(batch_size, hot_buffer_size, miss_rate) + return avg_us, avg_us, avg_us + + if __name__ == "__main__": benchmark_latency.run(print_data=True) + benchmark_planned_copy_latency.run(print_data=True)