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

HiSparse: shared-index (IndexShare) plan-then-IO swap-in prefetch (#34329)

Browse Source 
Co-authored-by: Tingwei Huang <huangtingwei9988@gmail.com>
This commit is contained in:
Zhiqiang Xie
2026-08-11 01:58:28 -07:00
committed by GitHub
co-authored by Tingwei Huang
parent 396722e490
commit 5469faec45
8 changed files with 1057 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/hisparse.cuh
+209 -42
View File
@@ -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 <bool IsMLA, bool IsDsv4Layout>
__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<int32_t>(dst_loc));
const auto [src_value_ptr, src_scale_ptr] =
get_pointer_paged(const_cast<void*>(host_cache_k), static_cast<int32_t>(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<void*>(host_cache_k),
/*dst_index=*/static_cast<int32_t>(dst_loc),
/*src_index=*/static_cast<int32_t>(src_loc));
#endif
} else {
// Generic path: device + host both linear, stride = item_size_bytes.
const auto src_k = static_cast<const char*>(host_cache_k) + src_loc * item_size_bytes;
auto dst_k = static_cast<char*>(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<const char*>(host_cache_v) + src_loc * item_size_bytes;
auto dst_v = static_cast<char*>(device_buffer_v) + dst_loc * item_size_bytes;
transfer_item_warp(lane_id, src_v, dst_v, item_size_bytes);
}
}
}
template <int BLOCK_SIZE>
__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<int64_t>(req_device_buffer_locs[evict_slot]);
const int64_t src_loc = req_host_cache_locs[miss_token];
const int64_t dst_loc = static_cast<int64_t>(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<int32_t>(dst_loc));
const auto [src_value_ptr, src_scale_ptr] =
get_pointer_paged(const_cast<void*>(host_cache_k), static_cast<int32_t>(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<void*>(host_cache_k),
/*dst_index=*/static_cast<int32_t>(dst_loc),
/*src_index=*/static_cast<int32_t>(src_loc));
#endif
} else {
// Generic path: device + host both linear, stride = item_size_bytes.
const auto src_k = static_cast<const char*>(host_cache_k) + src_loc * item_size_bytes;
auto dst_k = static_cast<char*>(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<const char*>(host_cache_v) + src_loc * item_size_bytes;
auto dst_v = static_cast<char*>(device_buffer_v) + dst_loc * item_size_bytes;
transfer_item_warp(lane_id, src_v, dst_v, item_size_bytes);
}
copy_miss_item<IsMLA, IsDsv4Layout>(
lane_id, host_cache_k, host_cache_v, device_buffer_k, device_buffer_v, src_loc, dst_loc, item_size_bytes);
}
}
}
template <int BLOCK_SIZE, int NUM_TOP_K, int HOT_BUFFER_SIZE, bool IsMLA, bool IsDsv4Layout>
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<int64_t*>(miss_src_out.data_ptr()) : nullptr;
int32_t* const miss_dst_ptr = RecordMissPlan ? static_cast<int32_t*>(miss_dst_out.data_ptr()) : nullptr;
int32_t* const miss_count_ptr = RecordMissPlan ? static_cast<int32_t*>(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<const int64_t*>(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<const int64_t*>(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<const int32_t*>(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<const int32_t*>(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 <int BLOCK_SIZE, bool IsMLA, bool IsDsv4Layout, bool SkipIO>
__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<IsMLA, IsDsv4Layout>(
lane_id,
host_cache_k,
host_cache_v,
device_buffer_k,
device_buffer_v,
src_row[m],
static_cast<int64_t>(dst_row[m]),
item_size_bytes);
}
start += cnt;
}
}
}
template <int BLOCK_SIZE, bool IsMLA, bool IsDsv4Layout, bool SkipIO>
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<BLOCK_SIZE, IsMLA, IsDsv4Layout, SkipIO>,
static_cast<const int64_t*>(miss_src_locs.data_ptr()),
static_cast<const int32_t*>(miss_dst_locs.data_ptr()),
static_cast<const int32_t*>(miss_counts.data_ptr()),
static_cast<const int32_t*>(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
python/sglang/kernels/ops/kvcache/hisparse.py
+121 -3
View File
@@ -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,
)
python/sglang/srt/environ.py
+8
View File
@@ -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
python/sglang/srt/managers/hisparse_coordinator.py
+214 -4
View File
@@ -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
python/sglang/srt/model_executor/model_runner.py
+9 -1
View File
@@ -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):