From 13f4f010d8ea7fbf628cd3b5313a73cac6c0285e Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Sun, 22 Mar 2026 23:09:31 -0700 Subject: [PATCH] HiSparse for Sparse Attention (#20343) --- python/sglang/jit_kernel/csrc/hisparse.cuh | 390 ++++++++++++ python/sglang/jit_kernel/hisparse.py | 88 +++ .../srt/layers/attention/nsa_backend.py | 26 +- .../srt/managers/hisparse_coordinator.py | 596 ++++++++++++++++++ python/sglang/srt/managers/schedule_batch.py | 15 + python/sglang/srt/managers/scheduler.py | 108 +++- .../scheduler_output_processor_mixin.py | 4 + .../scheduler_runtime_checker_mixin.py | 3 + python/sglang/srt/managers/tp_worker.py | 3 + .../srt/mem_cache/hisparse_memory_pool.py | 341 ++++++++++ python/sglang/srt/mem_cache/memory_pool.py | 5 +- .../sglang/srt/mem_cache/memory_pool_host.py | 2 +- .../sglang/srt/mem_cache/sparsity/__init__.py | 2 + .../sparsity/core/sparse_coordinator.py | 11 +- .../sglang/srt/mem_cache/sparsity/factory.py | 54 +- .../srt/model_executor/cuda_graph_runner.py | 9 + .../srt/model_executor/forward_batch_info.py | 4 + .../sglang/srt/model_executor/model_runner.py | 28 + .../model_runner_kv_cache_mixin.py | 37 +- python/sglang/srt/server_args.py | 23 +- 20 files changed, 1691 insertions(+), 58 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/hisparse.cuh create mode 100644 python/sglang/jit_kernel/hisparse.py create mode 100644 python/sglang/srt/managers/hisparse_coordinator.py create mode 100644 python/sglang/srt/mem_cache/hisparse_memory_pool.py diff --git a/python/sglang/jit_kernel/csrc/hisparse.cuh b/python/sglang/jit_kernel/csrc/hisparse.cuh new file mode 100644 index 000000000..3cf12178f --- /dev/null +++ b/python/sglang/jit_kernel/csrc/hisparse.cuh @@ -0,0 +1,390 @@ +#include +#include + +#include + +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int WARP_SIZE = 32; +constexpr int32_t TOKEN_HIT = 0xFFFFFFFF; +constexpr int32_t HASH_EMPTY = -1; + +// Knuth multiplicative hash for open-addressing table of size hash_size. +__device__ __forceinline__ int hash_slot(int32_t key, int hash_size) { + return ((uint32_t)key * 2654435761u) % (uint32_t)hash_size; +} + +__device__ __forceinline__ void +transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_t item_size_bytes) { + const uint64_t* __restrict__ src = static_cast(src_addr); + uint64_t* __restrict__ dst = static_cast(dst_addr); + const int total_chunks = item_size_bytes / sizeof(uint64_t); + +#pragma unroll + for (int j = lane_id; j < total_chunks; j += WARP_SIZE) { + uint64_t tmp; + asm volatile("ld.global.nc.b64 %0,[%1];" : "=l"(tmp) : "l"(src + j) : "memory"); + asm volatile("st.global.cg.b64 [%0],%1;" ::"l"(dst + j), "l"(tmp) : "memory"); + } +} + +__device__ __forceinline__ int warp_inclusive_scan(int* s_data, int lane_id, int offset, int count, int accumulator) { + int idx = lane_id + offset; + int val = (idx < count) ? s_data[idx] : 0; + +#pragma unroll + for (int i = 1; i < 32; i *= 2) { + int n = __shfl_up_sync(0xffffffff, val, i); + if (lane_id >= i) val += n; + } + val += accumulator; + if (idx < count) { + s_data[idx] = val; + } + accumulator = __shfl_sync(0xffffffff, val, 31); + return accumulator; +} + +// Each block processes one request +// req_pool_indices are int64_t (pool indices can be large), seq_lens are int32_t +// Layout: [HOT_BUFFER_SIZE slots for LRU] + [page_size slots for newest token] +// newest_slot is at HOT_BUFFER_SIZE (first position of extra page) +template +__global__ void load_cache_to_device_buffer_kernel( + const int32_t* __restrict__ top_k_tokens, + int32_t* __restrict__ device_buffer_tokens, + const int64_t* __restrict__ host_cache_locs, + const int32_t* __restrict__ device_buffer_locs, + const void* __restrict__ host_cache_k, + const void* __restrict__ host_cache_v, + void* __restrict__ device_buffer_k, + void* __restrict__ device_buffer_v, + int32_t* __restrict__ top_k_device_locs, + const int64_t* __restrict__ req_pool_indices, + const int32_t* __restrict__ seq_lens, + int16_t* __restrict__ lru_slots, + const int32_t* __restrict__ num_real_reqs, + int64_t buffer_stride_0, + int64_t host_stride, + int64_t lru_slot_stride_0, + int64_t top_k_tokens_stride, + int64_t top_k_device_locs_stride, + int64_t page_size, + int64_t item_size_bytes) { + // todo hisparse: support page wise sparsity + constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; + constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K + WARP_SIZE - 1) / WARP_SIZE; + constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + WARP_SIZE - 1) / WARP_SIZE; + + const int bid = blockIdx.x; + // Early exit for padded blocks (CUDA graph pads batch to a captured size) + if (bid >= num_real_reqs[0]) return; + + const int tid = threadIdx.x; + const int warp_id = tid / WARP_SIZE; + const int lane_id = tid % WARP_SIZE; + const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1; + + const int64_t rid = req_pool_indices[bid]; + const int64_t seq_len = seq_lens[bid]; + + // Calculate offsets for this request + const int32_t* req_top_k_tokens = top_k_tokens + bid * top_k_tokens_stride; + int32_t* req_top_k_device_locs = top_k_device_locs + bid * top_k_device_locs_stride; + + const int64_t buffer_offset = rid * buffer_stride_0; + int32_t* req_device_buffer_tokens = device_buffer_tokens + buffer_offset; + const int32_t* req_device_buffer_locs = device_buffer_locs + buffer_offset; + const int64_t* req_host_cache_locs = host_cache_locs + rid * host_stride; + int16_t* req_lru_slots = lru_slots + rid * lru_slot_stride_0; + + // Fast path: short sequences have all tokens in the device buffer in order. + if (seq_len <= HOT_BUFFER_SIZE) { + const int count = (seq_len < NUM_TOP_K) ? static_cast(seq_len) : NUM_TOP_K; + for (int i = tid; i < count; i += BLOCK_SIZE) { + int32_t token_pos = req_top_k_tokens[i]; + if (token_pos >= 0) { + req_top_k_device_locs[i] = req_device_buffer_locs[token_pos]; + } + } + return; + } + + // Top-k token positions; reused as miss-token scratch in the copy phase + __shared__ int32_t s_top_k_tokens[NUM_TOP_K]; + // Prefix-sum offsets for hit counting and miss counting + __shared__ int32_t s_chunk_offset[NUM_BUFFER_CHUNKS + 1]; + // Prefix-sum offsets for evictable counting + __shared__ int32_t s_evict_chunk_offset[NUM_BUFFER_CHUNKS + 1]; + // Compacted slot ordering: [hits fwd→ ... ←evictables bwd] + __shared__ int16_t s_lru_slots_out[HOT_BUFFER_SIZE]; + // Open-addressing hash table: top-k token_id → top-k index + constexpr int HASH_SIZE = NUM_TOP_K * 2; + __shared__ int32_t s_hash_keys[HASH_SIZE]; + __shared__ int16_t s_hash_vals[HASH_SIZE]; + + __shared__ int32_t s_total_hits; + __shared__ int32_t s_newest_hit; + + // Initialize shared memory: counters, hash table, prefix-sum offsets. + if (tid == 0) { + s_total_hits = 0; + s_newest_hit = 0; + } + for (int i = tid; i < HASH_SIZE; i += BLOCK_SIZE) { + s_hash_keys[i] = HASH_EMPTY; + } + for (int i = tid; i < NUM_BUFFER_CHUNKS + 1; i += BLOCK_SIZE) { + s_chunk_offset[i] = 0; + s_evict_chunk_offset[i] = 0; + } + __syncthreads(); + + const int newest_slot = HOT_BUFFER_SIZE; + const int32_t newest_token = seq_len - 1; + + // Insert top-k tokens into shared-memory hash table. + for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) { + int32_t token_idx = req_top_k_tokens[i]; + if (token_idx == newest_token) { + // If topk includes the latest token, bind its canonical occurrence to newest_slot (at HOT_BUFFER_SIZE) and mark + // it as a hit. newest_slot is at the first position of the extra page, excluded from LRU tracking. + s_top_k_tokens[i] = TOKEN_HIT; + req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot]; + s_newest_hit = 1; + } else { + int slot = hash_slot(token_idx, HASH_SIZE); + while (true) { + int32_t old = atomicCAS(&s_hash_keys[slot], HASH_EMPTY, token_idx); + if (old == HASH_EMPTY || old == token_idx) { + s_hash_vals[slot] = static_cast(i); + break; + } + slot = (slot + 1) % HASH_SIZE; + } + s_top_k_tokens[i] = token_idx; + } + } + __syncthreads(); + + constexpr int ITERATIONS_PER_WARP_BUFFER = (NUM_BUFFER_CHUNKS + NUM_WARPS - 1) / NUM_WARPS; + int total_hit_count = 0; + int total_evict_count = 0; + for (int iter = 0; iter < ITERATIONS_PER_WARP_BUFFER; iter++) { + int chunk_idx = warp_id + iter * NUM_WARPS; + bool has_valid_chunk = chunk_idx < NUM_BUFFER_CHUNKS; + + const int slot_idx = chunk_idx * WARP_SIZE + lane_id; + const bool has_valid_slot = has_valid_chunk && (slot_idx < HOT_BUFFER_SIZE); + const int16_t buf_slot = has_valid_slot ? req_lru_slots[slot_idx] : -1; + int32_t my_buffer_token = (buf_slot >= 0) ? req_device_buffer_tokens[buf_slot] : -1; + int my_found_top_k_idx = -1; + if (my_buffer_token >= 0) { + int h = hash_slot(my_buffer_token, HASH_SIZE); + while (true) { + int32_t k = s_hash_keys[h]; + if (k == my_buffer_token) { + my_found_top_k_idx = static_cast(s_hash_vals[h]); + break; + } + if (k == HASH_EMPTY) break; + h = (h + 1) % HASH_SIZE; + } + } + bool is_hit = my_found_top_k_idx >= 0; + bool is_evictable = has_valid_slot && !is_hit; + + // Record hits + if (is_hit) { + s_top_k_tokens[my_found_top_k_idx] = TOKEN_HIT; + req_top_k_device_locs[my_found_top_k_idx] = req_device_buffer_locs[buf_slot]; + } + + int local_hit_offset = 0; + int local_evict_offset = 0; + if (has_valid_chunk) { + const unsigned int hit_mask = __ballot_sync(0xFFFFFFFF, is_hit); + const unsigned int evict_mask = __ballot_sync(0xFFFFFFFF, is_evictable); + local_hit_offset = __popc(hit_mask & lanes_before); + local_evict_offset = __popc(evict_mask & lanes_before); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = __popc(hit_mask); + s_evict_chunk_offset[chunk_idx + 1] = __popc(evict_mask); + } + } + __syncthreads(); + + if (warp_id == 0) { + total_hit_count = + warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_hit_count); + total_evict_count = + warp_inclusive_scan(s_evict_chunk_offset, lane_id, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_evict_count); + if (tid == 0) { + s_total_hits = total_hit_count; + } + } + __syncthreads(); + + // Hits grow forward from index 0 + if (is_hit) { + int hit_offset = s_chunk_offset[chunk_idx] + local_hit_offset; + s_lru_slots_out[hit_offset] = buf_slot; + } + // Evictables grow backward from HOT_BUFFER_SIZE - 1 + if (is_evictable) { + int evict_offset = s_evict_chunk_offset[chunk_idx] + local_evict_offset; + s_lru_slots_out[HOT_BUFFER_SIZE - 1 - evict_offset] = buf_slot; + } + } + __syncthreads(); + + // Write back LRU order: evictables at front (LRU), hits at back (MRU). + { + const int total_evictable = HOT_BUFFER_SIZE - s_total_hits; + for (int i = tid; i < HOT_BUFFER_SIZE; i += BLOCK_SIZE) { + if (i < total_evictable) { + // Evictables: source at backward end, dest at LRU front + req_lru_slots[i] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i]; + } else { + // Hits: source at forward end, dest at MRU back + req_lru_slots[i] = s_lru_slots_out[i - total_evictable]; + } + } + } + + // Reset offsets for the miss counting phase (only NUM_TOKEN_CHUNKS + 1 entries needed). + for (int i = tid; i < NUM_TOKEN_CHUNKS + 1; i += BLOCK_SIZE) { + s_chunk_offset[i] = 0; + } + __syncthreads(); + + // Third pass to identify misses and their evictable slots + int total_misses = 0; + constexpr int ITERATIONS_PER_WARP_TOKEN = (NUM_TOKEN_CHUNKS + NUM_WARPS - 1) / NUM_WARPS; + for (int iter = 0; iter < ITERATIONS_PER_WARP_TOKEN; iter++) { + int chunk_idx = warp_id + iter * NUM_WARPS; + bool has_valid_chunk = chunk_idx < NUM_TOKEN_CHUNKS; + + const int chunk_token_start = chunk_idx * WARP_SIZE; + const int my_token_idx = chunk_token_start + lane_id; + const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K); + + int32_t my_token = 0; + bool is_miss = false; + int local_miss_offset = 0; + + if (has_valid_token) { + is_miss = s_top_k_tokens[my_token_idx] != TOKEN_HIT; + if (is_miss) { + my_token = s_top_k_tokens[my_token_idx]; + } + } + + if (has_valid_chunk) { + const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss); + local_miss_offset = __popc(miss_mask & lanes_before); + const int warp_miss_count = __popc(miss_mask); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = warp_miss_count; + } + } + __syncthreads(); + + if (warp_id == 0) { + total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses); + } + __syncthreads(); + + if (is_miss) { + int miss_offset = s_chunk_offset[chunk_idx] + local_miss_offset; + int16_t evict_slot = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - miss_offset]; + // Reuse s_top_k_tokens as miss scratch: miss_offset < my_token_idx always + // holds (hits are skipped), so compacted writes never overrun pending reads. + s_top_k_tokens[miss_offset] = my_token; + req_top_k_device_locs[my_token_idx] = req_device_buffer_locs[evict_slot]; + req_device_buffer_tokens[evict_slot] = my_token; + } + } + __syncthreads(); + + total_misses = NUM_TOP_K - s_total_hits - s_newest_hit; + // 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]; + + 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 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 +void load_cache_to_device_buffer( + 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 host_cache_v, + tvm::ffi::TensorView device_buffer_k, + tvm::ffi::TensorView device_buffer_v, + tvm::ffi::TensorView top_k_device_locs, + tvm::ffi::TensorView req_pool_indices, + tvm::ffi::TensorView seq_lens, + tvm::ffi::TensorView lru_slots, + tvm::ffi::TensorView num_real_reqs, + int64_t page_size, + int64_t item_size_bytes) { + using namespace host; + + const int64_t bs = top_k_tokens.shape()[0]; + const int64_t host_stride = host_cache_locs.shape()[1]; + 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]; + const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0]; + const auto device = LaunchKernel::resolve_device(top_k_tokens.device()); + + LaunchKernel(bs, BLOCK_SIZE, device)( + load_cache_to_device_buffer_kernel, + 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(), + (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(), + static_cast(top_k_device_locs.data_ptr()), + static_cast(req_pool_indices.data_ptr()), + static_cast(seq_lens.data_ptr()), + static_cast(lru_slots.data_ptr()), + static_cast(num_real_reqs.data_ptr()), + buffer_stride_0, + host_stride, + lru_slot_stride_0, + top_k_tokens_stride, + top_k_device_locs_stride, + page_size, + item_size_bytes); +} + +} // namespace diff --git a/python/sglang/jit_kernel/hisparse.py b/python/sglang/jit_kernel/hisparse.py new file mode 100644 index 000000000..25db57b0b --- /dev/null +++ b/python/sglang/jit_kernel/hisparse.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import load_jit, make_cpp_args + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@functools.cache +def _jit_sparse_module( + item_size_bytes: int, + block_size: int, + num_top_k: int, + hot_buffer_size: int, + is_mla: bool = False, +) -> Module: + template_args = make_cpp_args(block_size, num_top_k, hot_buffer_size, is_mla) + cache_args = make_cpp_args( + item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla + ) + return load_jit( + "sparse_cache", + *cache_args, + cuda_files=["hisparse.cuh"], + cuda_wrappers=[ + ( + "load_cache_to_device_buffer", + f"load_cache_to_device_buffer<{template_args}>", + ) + ], + ) + + +def load_cache_to_device_buffer_mla( + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int = 1, + block_size: int = 256, + num_real_reqs: torch.Tensor | None = None, +) -> None: + assert ( + hot_buffer_size >= num_top_k + ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" + + module = _jit_sparse_module( + item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla=True + ) + + empty = torch.empty(0) + + if num_real_reqs is None: + num_real_reqs = torch.tensor( + [top_k_tokens.size(0)], dtype=torch.int32, device=top_k_tokens.device + ) + + module.load_cache_to_device_buffer( + top_k_tokens, + device_buffer_tokens, + host_cache_locs, + device_buffer_locs, + host_cache, + empty, + device_buffer, + empty, + top_k_device_locs, + req_pool_indices, + seq_lens, + lru_slots, + num_real_reqs, + page_size, + item_size_bytes, + ) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 84e93c0e1..ad0d5840c 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -177,6 +177,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata): attn_metadata: NSAMetadata topk_transform_method: TopkTransformMethod paged_mqa_schedule_metadata: Optional[torch.Tensor] = None + force_unfused_topk: bool = False def get_seqlens_int32(self) -> torch.Tensor: return self.attn_metadata.cache_seqlens_int32 @@ -246,7 +247,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata): else: page_table_size_1 = self.attn_metadata.page_table_1 - if not envs.SGLANG_NSA_FUSE_TOPK.get(): + if not envs.SGLANG_NSA_FUSE_TOPK.get() or self.force_unfused_topk: return fast_topk_v2(logits, seq_lens_topk, topk, row_starts=ks) elif self.topk_transform_method == TopkTransformMethod.PAGED: # NOTE(dark): if fused, we return a transformed page table directly @@ -1367,6 +1368,14 @@ class NativeSparseAttnBackend( page_size=1, ) + # todo hisparse: to cover more backends + if forward_batch.hisparse_coordinator is not None: + page_table_1 = ( + forward_batch.token_to_kv_pool.translate_loc_to_hisparse_device( + page_table_1 + ) + ) + if nsa_impl == "tilelang": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) @@ -1512,7 +1521,14 @@ class NativeSparseAttnBackend( if topk_indices is not None: topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) - if envs.SGLANG_NSA_FUSE_TOPK.get(): + if forward_batch.hisparse_coordinator is not None: + page_table_1 = forward_batch.hisparse_coordinator.swap_in_selected_pages( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + topk_indices, + layer.layer_id, + ) + elif envs.SGLANG_NSA_FUSE_TOPK.get(): page_table_1 = topk_indices else: page_table_1 = transform_index_page_table_decode( @@ -2055,6 +2071,7 @@ class NativeSparseAttnBackend( and sum_seq_lens <= forward_batch.get_max_chunk_capacity() # Fits in chunk and (not is_nsa_enable_prefill_cp()) # CP not enabled + and (forward_batch.hisparse_coordinator is None) ) else: self.use_mha = False # Decode/verify always use MLA @@ -2096,10 +2113,15 @@ class NativeSparseAttnBackend( def get_indexer_metadata( self, layer_id: int, forward_batch: ForwardBatch ) -> NSAIndexerMetadata: + force_unfused = ( + forward_batch.hisparse_coordinator is not None + and forward_batch.forward_mode.is_decode_or_idle() + ) return NSAIndexerMetadata( attn_metadata=self.forward_metadata, topk_transform_method=self.get_topk_transform_method(), paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata, + force_unfused_topk=force_unfused, ) def _compute_flashmla_metadata(self, cache_seqlens: torch.Tensor, seq_len_q: int): diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py new file mode 100644 index 000000000..92ef22f40 --- /dev/null +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -0,0 +1,596 @@ +# to be combined with the sparse coordinator class and sparse algorithm family + +import logging +from typing import List, NamedTuple + +import torch + +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.mem_cache.hisparse_memory_pool import ( + HiSparseNSATokenToKVPool, + HiSparseTokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost +from sglang.srt.utils import get_device_module + +device_module = get_device_module() + +from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + +logger = logging.getLogger(__name__) + + +class HiSparseAct(NamedTuple): + start_event: device_module.Event + finish_event: device_module.Event + req: Req + + +class HiSparseCoordinator: + def __init__( + self, + req_to_token_pool: ReqToTokenPool, + token_to_kv_pool_allocator: HiSparseTokenToKVPoolAllocator, + top_k: int, + device_buffer_size: int, + device: str, + tp_group: torch.distributed.ProcessGroup, + host_to_device_ratio: int = 2, + ): + self.req_to_token_pool = req_to_token_pool + self.token_to_kv_pool_allocator = token_to_kv_pool_allocator + self.top_k = top_k + self.device_buffer_size = device_buffer_size + self.device = device + + self.mem_pool_device: HiSparseNSATokenToKVPool = ( + self.token_to_kv_pool_allocator.get_kvcache() + ) + self.mem_pool_host = MLATokenToKVPoolHost( + device_pool=self.mem_pool_device, + host_to_device_ratio=host_to_device_ratio, + host_size=0, + page_size=1, # for simplicity, we set page size to 1 to enable backup one token at a time + layout="layer_first", + override_kv_cache_dim=self.mem_pool_device.kv_cache_dim, + ) + + max_num_reqs = req_to_token_pool.size + max_context_len = req_to_token_pool.max_context_len + + # to have an extra page for new tokens + self.padded_buffer_size = ( + self.device_buffer_size + self.mem_pool_device.page_size + ) + + self.req_to_device_buffer = torch.zeros( + (max_num_reqs, self.padded_buffer_size), dtype=torch.int64, device=device + ) + self.req_device_buffer_size = torch.zeros( + max_num_reqs, dtype=torch.int64, device="cpu" + ) + self.req_to_host_pool = torch.full( + (max_num_reqs, max_context_len), + -1, + dtype=torch.int64, + device=device, + ) + + self.write_staging_stream = device_module.Stream() + self.ack_staging_queue: List[HiSparseAct] = [] + self.decode_producer_stream = None + + self.tp_group = tp_group + self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group) + + # initialize data structures for swap-in kernel + layer_num = self.mem_pool_device.layer_num + self.req_device_buffer_tokens = torch.full( + (layer_num, max_num_reqs, self.padded_buffer_size), + -1, + dtype=torch.int32, + device=device, + ) + self.req_device_buffer_token_locs = torch.full( + (layer_num, max_num_reqs, self.padded_buffer_size), + -1, + dtype=torch.int32, + device=device, + ) + self._lru_init = torch.arange( + self.device_buffer_size, dtype=torch.int16, device=device + ) + self.lru_slots = ( + self._lru_init.view(1, 1, -1) + .repeat(layer_num, max_num_reqs, 1) + .contiguous() + ) + + # Pre-allocated output buffer for swap_in_selected_pages (CUDA-graph safe) + self.top_k_device_locs_buffer = torch.full( + (max_num_reqs, self.top_k), -1, dtype=torch.int32, device=device + ) + # Scalar tensor: number of real (non-padded) requests in the batch. + # Updated before each graph replay so padded blocks early-return. + self.num_real_reqs = torch.zeros(1, dtype=torch.int32, device=device) + + # CPU flag: True means "skip backup on the next decode step" because + # staging already backed up all prefill tokens. Cleared after one step. + self._skip_first_backup = [False] * max_num_reqs + + def set_decode_producer_stream(self, stream) -> None: + self.decode_producer_stream = stream + + def admit_request_into_staging(self, req: Req) -> None: + req.staging = True + logical_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : len(req.fill_ids) + ] + device_indices = self.mem_pool_device._translate_loc_to_hisparse_device( + logical_indices + ) + + prefill_len = len(device_indices) + host_indices = self.mem_pool_host.alloc(prefill_len) + if host_indices is None: + logger.error( + "HiSparse: host mem pool alloc failed for %d tokens (req %s)", + prefill_len, + req.rid, + ) + raise RuntimeError( + f"HiSparse host mem pool alloc failed for {prefill_len} tokens" + ) + host_indices = host_indices.to(device=self.device) + self.req_to_host_pool[req.req_pool_idx, :prefill_len] = host_indices + + start_event = device_module.Event() + finish_event = device_module.Event() + start_event.record() + with device_module.stream(self.write_staging_stream): + start_event.wait(self.write_staging_stream) + self.mem_pool_host.backup_from_device_all_layer( + self.mem_pool_device, host_indices, device_indices, io_backend="kernel" + ) + finish_event.record() + if host_indices.is_cuda: + host_indices.record_stream(self.write_staging_stream) + if device_indices.is_cuda: + device_indices.record_stream(self.write_staging_stream) + + self.ack_staging_queue.append(HiSparseAct(start_event, finish_event, req)) + + def alloc_device_buffer(self, req: Req) -> None: + allocated_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : req.kv_allocated_len + ] + page_size = self.mem_pool_device.page_size + # Allocate only enough for current tokens (page-aligned). + # When prefill already fills device_buffer_size, include the reserved page. + alloc_size = min( + ((req.kv_allocated_len + page_size - 1) // page_size) * page_size, + self.device_buffer_size, + ) + if alloc_size == self.device_buffer_size: + alloc_size = self.padded_buffer_size + buffer_indices = self.token_to_kv_pool_allocator.alloc_device_buffer( + allocated_indices, + alloc_size, + ) + if buffer_indices is None: + logger.error( + "HiSparse: alloc_device_buffer failed for req %s " + "(kv_allocated_len=%d, alloc_size=%d)", + req.rid, + req.kv_allocated_len, + alloc_size, + ) + raise RuntimeError("HiSparse alloc_device_buffer returned None") + + self.req_to_device_buffer[req.req_pool_idx, :alloc_size] = buffer_indices + self.req_device_buffer_size[req.req_pool_idx] = alloc_size + + self.req_device_buffer_tokens[ + :, req.req_pool_idx, : self.device_buffer_size + ] = torch.arange(self.device_buffer_size, device=self.device) + self.req_device_buffer_token_locs[:, req.req_pool_idx, :alloc_size] = ( + buffer_indices[:alloc_size] + ) + + def has_ongoing_staging(self) -> bool: + return len(self.ack_staging_queue) > 0 + + def collect_ready_reqs(self) -> List[Req]: + ready_reqs = [] + if len(self.ack_staging_queue) == 0: + return ready_reqs + + finish_count = 0 + for _, finish_event, _ in self.ack_staging_queue: + if not finish_event.query(): + break + finish_count += 1 + queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu") + if self.tp_world_size > 1: + # synchronize TP workers to make sure the same update to scheduler + torch.distributed.all_reduce( + queue_size, + op=torch.distributed.ReduceOp.MIN, + group=self.tp_group, + ) + finish_count = int(queue_size.item()) + while finish_count > 0: + _, _, req = self.ack_staging_queue.pop(0) + # prepare device buffer and update req + self.alloc_device_buffer(req) + req.staging = False + self._skip_first_backup[req.req_pool_idx] = True + finish_count -= 1 + ready_reqs.append(req) + return ready_reqs + + def _grow_device_buffers( + self, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens_cpu: torch.Tensor, + req_pool_indices_cpu: torch.Tensor, + ) -> torch.Tensor: + """Grow device buffers for requests whose sequence length exceeds current capacity.""" + current_caps = self.req_device_buffer_size[req_pool_indices_cpu] + short_reqs_cpu = seq_lens_cpu <= self.device_buffer_size + needs_grow_cpu = short_reqs_cpu & (seq_lens_cpu > current_caps) + + if torch.any(needs_grow_cpu): + page_size = self.mem_pool_device.page_size + grow_indices = torch.where(needs_grow_cpu)[0] + + # Compute all grow sizes on CPU, then do a single bulk allocation + req_idxs = [] + old_caps = [] + new_caps = [] + grow_sizes = [] + total_grow = 0 + for i in grow_indices.tolist(): + req_idx = int(req_pool_indices_cpu[i]) + current_cap = int(current_caps[i]) + seq_len = int(seq_lens_cpu[i]) + + new_cap = min( + ((seq_len + page_size - 1) // page_size) * page_size, + self.device_buffer_size, + ) + if new_cap == self.device_buffer_size: + new_cap = self.padded_buffer_size + grow_size = new_cap - current_cap + if grow_size <= 0: + continue + req_idxs.append(req_idx) + old_caps.append(current_cap) + new_caps.append(new_cap) + grow_sizes.append(grow_size) + total_grow += grow_size + + if total_grow > 0: + all_new_indices = ( + self.token_to_kv_pool_allocator.hisparse_attn_allocator.alloc( + total_grow + ) + ) + if all_new_indices is None: + logger.error( + "HiSparse: _grow_device_buffers bulk alloc failed " + "(total_grow=%d)", + total_grow, + ) + raise RuntimeError( + f"HiSparse _grow_device_buffers failed (total_grow={total_grow})" + ) + + offset = 0 + for req_idx, current_cap, new_cap, grow_size in zip( + req_idxs, old_caps, new_caps, grow_sizes + ): + chunk = all_new_indices[offset : offset + grow_size] + offset += grow_size + self.req_to_device_buffer[req_idx, current_cap:new_cap] = chunk + self.req_device_buffer_token_locs[ + :, req_idx, current_cap:new_cap + ] = chunk + self.req_device_buffer_size[req_idx] = new_cap + + reserved_positions = (seq_lens - 1).clamp(max=self.device_buffer_size) + return self.req_to_device_buffer[req_pool_indices, reserved_positions] + + def map_last_loc_to_buffer( + self, + seq_lens: torch.Tensor, + out_cache_loc: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens_cpu: torch.Tensor, + ) -> None: + req_pool_indices_cpu = req_pool_indices.cpu() + + self._eager_backup_previous_token( + seq_lens, req_pool_indices, seq_lens_cpu, req_pool_indices_cpu + ) + # Grow device buffers if needed and resolve the latest-token slot. + reserved_buffer_loc = self._grow_device_buffers( + seq_lens, req_pool_indices, seq_lens_cpu, req_pool_indices_cpu + ) + + self.req_device_buffer_token_locs[ + :, req_pool_indices, self.device_buffer_size + ] = reserved_buffer_loc.to(torch.int32) + + # todo, clear the prior mapping as well + self.mem_pool_device.full_to_hisparse_device_index_mapping[out_cache_loc] = ( + reserved_buffer_loc + ) + + def _eager_backup_previous_token( + self, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens_cpu: torch.Tensor, + req_pool_indices_cpu: torch.Tensor, + ) -> None: + """Back up the previous decode token to host memory. + + Every decode step, the token written in the *previous* step must be + backed up to host so the swap-in kernel can later recover it. + + The only exception is the first decode step right after staging: all + prefill tokens were already backed up during staging, so there is nothing new to save yet. + """ + if self.decode_producer_stream is not None: + device_module.current_stream().wait_stream(self.decode_producer_stream) + + # Build the list of batch positions that need a host backup. + # Skip the first decode step after staging (prefill already backed up). + backup_indices = [] + for i in range(len(seq_lens_cpu)): + req_idx = int(req_pool_indices_cpu[i]) + if self._skip_first_backup[req_idx]: + self._skip_first_backup[req_idx] = False + continue + backup_indices.append(i) + + if not backup_indices: + return + + backup_indices_gpu = torch.tensor( + backup_indices, dtype=torch.int64, device=self.device + ) + # The previous token's position and its device buffer slot: + # - short seq: slot = seq_len - 2 (within the regular buffer) + # - long seq: slot = device_buffer_size (the reserved slot) + actual_token_pos = seq_lens[backup_indices_gpu] - 2 + buffer_slot = actual_token_pos.clamp(max=self.device_buffer_size) + + backup_req_indices = req_pool_indices[backup_indices_gpu] + device_locs = self.req_to_device_buffer[backup_req_indices, buffer_slot] + + host_locs = self.mem_pool_host.alloc(len(device_locs)) + if host_locs is None: + logger.error( + "HiSparse: host mem pool alloc failed for %d decode backup tokens", + len(device_locs), + ) + raise RuntimeError( + f"HiSparse host mem pool alloc failed for {len(device_locs)} decode backup tokens" + ) + host_locs = host_locs.to(device=self.device) + self.req_to_host_pool[backup_req_indices, actual_token_pos] = host_locs + + self.mem_pool_host.backup_from_device_all_layer( + self.mem_pool_device, + host_locs, + device_locs.contiguous(), + io_backend="kernel", + ) + + def get_front_topk_tokens( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + ) -> torch.Tensor: + top_k_indices = self.req_to_device_buffer[req_pool_indices, : self.top_k].to( + torch.int32 + ) + topk_col_indices = torch.arange(self.top_k, device=self.device).unsqueeze(0) + # Mask out positions beyond each request's seq_len + mask = topk_col_indices >= seq_lens.unsqueeze(1) + top_k_indices[mask] = -1 + return top_k_indices + + def naive_load_topk( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + top_k_tokens: torch.Tensor, + layer_id: int, + ) -> torch.Tensor: + """Load top-k selected tokens into device memory and return their device indices. + + This is a naive per-request loop implementation for debugging/validation. + Production code uses swap_in_selected_pages (JIT CUDA kernel) instead. + + Args: + req_pool_indices: Pool indices for each request. Shape: (num_reqs,) + seq_lens: Sequence lengths for each request. Shape: (num_reqs,) + top_k_tokens: Selected token positions per request. Shape: (num_reqs, top_k) + layer_id: The layer to load KV cache for. + + Returns: + Device KV cache indices for the selected tokens. Shape: (num_reqs, top_k) + """ + num_reqs = req_pool_indices.size(0) + top_k_indices = torch.full( + (num_reqs, self.top_k), -1, dtype=torch.int32, device=self.device + ) + + for i in range(num_reqs): + seq_len = int(seq_lens[i].item()) + top_n = min(seq_len, self.top_k) + if top_n == 0: + continue + + req_idx = int(req_pool_indices[i].item()) + selected_tokens = top_k_tokens[i, :top_n].to(dtype=torch.int64) + + assert torch.all( + selected_tokens >= 0 + ), f"Req {req_idx}: selected tokens contain negative positions" + assert torch.all(selected_tokens < seq_len), ( + f"Req {req_idx}: selected tokens {selected_tokens.tolist()} " + f"out of range for seq_len={seq_len}" + ) + + if seq_len <= self.device_buffer_size: + device_indices = self.req_to_device_buffer[req_idx, selected_tokens] + else: + device_indices = torch.empty( + top_n, dtype=torch.int64, device=self.device + ) + + is_latest_token = selected_tokens == (seq_len - 1) + needs_host_load = ~is_latest_token + + device_indices[is_latest_token] = self.req_to_device_buffer[ + req_idx, self.device_buffer_size + ] + + num_to_load = int(needs_host_load.sum().item()) + if num_to_load > 0: + tokens_to_load = selected_tokens[needs_host_load] + host_locs = self.req_to_host_pool[req_idx, tokens_to_load] + + invalid_mask = host_locs < 0 + if torch.any(invalid_mask): + bad_positions = tokens_to_load[invalid_mask].tolist() + raise AssertionError( + f"Req {req_idx} (seq_len={seq_len}, layer={layer_id}): " + f"missing host backup at token positions {bad_positions}" + ) + + buffer_locs = self.req_to_device_buffer[req_idx, :num_to_load] + device_indices[needs_host_load] = buffer_locs + + self.mem_pool_host.load_to_device_per_layer( + self.mem_pool_device, + host_locs, + buffer_locs, + layer_id, + io_backend="kernel", + ) + + top_k_indices[i, :top_n] = device_indices.to(torch.int32) + + return top_k_indices + + def abort_staging_request(self, req: Req) -> None: + """Remove a request from the staging queue and free its host resources. + + Must be called when aborting a request that has been admitted into staging + but has not yet completed (i.e. req.staging is True). + """ + # Remove from staging queue + self.ack_staging_queue = [ + act for act in self.ack_staging_queue if act.req is not req + ] + # Wait for any in-flight staging DMA to complete before freeing + self.write_staging_stream.synchronize() + + # Free host memory that was allocated during admit_request_into_staging + host_indices = self.req_to_host_pool[req.req_pool_idx, : req.kv_allocated_len] + host_indices = host_indices[host_indices >= 0] + if host_indices.numel() > 0: + self.mem_pool_host.free(host_indices) + self.req_to_host_pool[req.req_pool_idx, :] = -1 + self._skip_first_backup[req.req_pool_idx] = False + req.staging = False + + def retract_req(self, req: Req) -> None: + if req.staging: + self.abort_staging_request(req) + else: + self.request_finished(req) + + def request_finished(self, req: Req): + # release resources only after the execution of a potential overlapped batch + if self.decode_producer_stream is not None: + device_module.current_stream().wait_stream(self.decode_producer_stream) + + # release memory — only free actually-allocated buffer indices + current_cap = int(self.req_device_buffer_size[req.req_pool_idx]) + buffer_indices = self.req_to_device_buffer[req.req_pool_idx, :current_cap] + self.token_to_kv_pool_allocator.free_hisparse_indices(buffer_indices) + + allocated_locs = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : req.kv_allocated_len + ] + self.token_to_kv_pool_allocator.full_to_hisparse_device_index_mapping[ + allocated_locs + ] = 0 + + host_indices = self.req_to_host_pool[req.req_pool_idx, : req.kv_allocated_len] + host_indices = host_indices[host_indices >= 0] + if host_indices.numel() > 0: + self.mem_pool_host.free(host_indices) + # clear req info + self.req_device_buffer_tokens[:, req.req_pool_idx, :] = -1 + self.req_device_buffer_token_locs[:, req.req_pool_idx, :] = -1 + self.req_to_device_buffer[req.req_pool_idx, :] = 0 + self.req_device_buffer_size[req.req_pool_idx] = 0 + self.req_to_host_pool[req.req_pool_idx, :] = -1 + 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( + self, + req_pool_indices: torch.Tensor, + 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.""" + # The CUDA kernel expects req_pool_indices as int64 and seq_lens as int32. + if req_pool_indices.dtype != torch.int64: + raise ValueError( + f"req_pool_indices dtype {req_pool_indices.dtype} is not int64 as expected" + ) + if seq_lens.dtype != torch.int32: + raise ValueError( + f"seq_lens dtype {seq_lens.dtype} is not int32 as expected" + ) + if top_k_result.dtype != torch.int32: + raise ValueError( + f"top_k_result dtype {top_k_result.dtype} is not int32 as expected" + ) + + num_reqs = req_pool_indices.size(0) + top_k_indices = self.top_k_device_locs_buffer[:num_reqs] + top_k_indices.fill_(-1) + # todo, adjustable for performance + block_size = 1024 + load_cache_to_device_buffer_mla( + top_k_tokens=top_k_result, + device_buffer_tokens=self.req_device_buffer_tokens[layer_id], + host_cache_locs=self.req_to_host_pool, + device_buffer_locs=self.req_device_buffer_token_locs[layer_id], + host_cache=self.mem_pool_host.kv_buffer[layer_id], + device_buffer=self.mem_pool_device.kv_buffer[layer_id], + top_k_device_locs=top_k_indices, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + lru_slots=self.lru_slots[layer_id], + item_size_bytes=self.mem_pool_host.token_stride_size, + num_top_k=self.top_k, + hot_buffer_size=self.device_buffer_size, + page_size=1, + block_size=block_size, + num_real_reqs=self.num_real_reqs, + ) + return top_k_indices diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 063b8d7c3..670cc1c85 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -94,6 +94,7 @@ if TYPE_CHECKING: from typing import Any, Dict from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator from sglang.srt.managers.session_controller import Session from sglang.srt.observability.scheduler_metrics_mixin import PrefillStats from sglang.srt.speculative.eagle_info import EagleDraftInput @@ -799,6 +800,9 @@ class Req(ReqDllmMixin): # For diffusion LLM self.init_diffusion_llm(dllm_config) + # For hisparse + self.staging = False + @property def seqlen(self) -> int: """Get the current sequence length of the request.""" @@ -1343,6 +1347,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): dp_cooperation_info: Optional[DPCooperationInfo] = None prefill_stats: Optional[PrefillStats] = None + # HiSparse + hisparse_coordinator: Optional[HiSparseCoordinator] = None + @classmethod def init_new( cls, @@ -2065,6 +2072,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.orig_seq_lens.add_(1) self.seq_lens_sum += bs + if self.hisparse_coordinator is not None: + self.hisparse_coordinator.map_last_loc_to_buffer( + self.seq_lens, + self.out_cache_loc, + self.req_pool_indices, + self.seq_lens_cpu, + ) + if get_global_server_args().enable_mamba_extra_buffer(): if len(self.reqs) == 0: self.mamba_track_indices = torch.empty( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 42e841fa4..f31303e8b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -78,6 +78,7 @@ from sglang.srt.layers.moe import initialize_moe_config from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config from sglang.srt.lora.lora_overlap_loader import LoRAOverlapLoader +from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator from sglang.srt.managers.io_struct import ( AbortReq, ActiveRanksOutput, @@ -193,6 +194,7 @@ from sglang.srt.observability.scheduler_metrics_mixin import ( ) from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info from sglang.srt.parser.reasoning_parser import ReasoningParser +from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( @@ -334,6 +336,8 @@ class Scheduler( self.enable_hierarchical_cache = server_args.enable_hierarchical_cache self.enable_hicache_storage = server_args.hicache_storage_backend is not None self.max_recv_per_poll = envs.SGLANG_SCHEDULER_MAX_RECV_PER_POLL.get() + self.enable_hisparse = server_args.enable_hisparse + self.hisparse_coordinator: Optional[HiSparseCoordinator] = None # Distributed rank info self.attn_tp_rank, self.attn_tp_size, self.attn_dp_rank = ( @@ -767,9 +771,13 @@ class Scheduler( self.tree_cache = RadixCache(params) if server_args.enable_streaming_session: - self.tree_cache = SessionAwareCache(self.tree_cache) + if self.enable_hisparse: + # Coordinator was created inside ModelRunner.initialize() before CUDA graph capture + self.hisparse_coordinator = self.tp_worker.model_runner.hisparse_coordinator + self.hisparse_coordinator.set_decode_producer_stream(self.forward_stream) + if ( server_args.disaggregation_mode == "decode" and server_args.disaggregation_decode_enable_offload_kvcache @@ -2034,6 +2042,45 @@ class Scheduler( def stash_chunked_request(self, req: Req): self.tree_cache.cache_unfinished_req(req, chunked=True) + def _build_hisparse_decode_batch(self, reqs): + """Build a ScheduleBatch for hisparse requests transitioning from staging to decode.""" + device = self.device + + batch = ScheduleBatch.init_new( + reqs=reqs, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + tree_cache=self.tree_cache, + model_config=self.model_config, + enable_overlap=self.enable_overlap, + spec_algorithm=self.spec_algorithm, + ) + + batch.req_pool_indices = torch.tensor( + [r.req_pool_idx for r in reqs], dtype=torch.int64, device=device + ) + seq_lens = [len(r.origin_input_ids) + len(r.output_ids) - 1 for r in reqs] + batch.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device=device) + batch.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64) + batch.orig_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + batch.seq_lens_sum = sum(seq_lens) + # output_ids = last generated token, used as input_ids by prepare_for_decode + batch.output_ids = torch.tensor( + [r.output_ids[-1] for r in reqs], dtype=torch.int64, device=device + ) + + # Set logprob fields if any request needs them + if batch.return_logprob: + batch.top_logprobs_nums = [r.top_logprobs_num for r in reqs] + batch.token_ids_logprobs = [list(r.origin_input_ids) for r in reqs] + + # Build sampling info from scratch for these requests + batch.sampling_info = SamplingBatchInfo.from_schedule_batch( + batch, self.model_config.vocab_size + ) + # todo hisparse, maybe other info to contain for the new batch + return batch + def get_next_batch_to_run(self) -> Optional[ScheduleBatch]: self._abort_on_waiting_timeout() self._abort_on_running_timeout() @@ -2054,30 +2101,40 @@ class Scheduler( chunked_req_to_exclude.add(self.chunked_req) self.stash_chunked_request(self.chunked_req) - if self.last_batch and self.last_batch.forward_mode.is_extend(): - if self.last_batch.chunked_req is not None: - # In the context pipeline parallelism, after the last chunk, the current microbatch still track outdated chunked_req. - # We need to discard it. - chunked_req_to_exclude.add(self.last_batch.chunked_req) - - if self.dllm_config is not None and self.last_batch.reqs: - chunked_req_to_exclude.update(self.last_batch.reqs) - - # Filter batch - last_bs = self.last_batch.batch_size() - self.last_batch.filter_batch( - chunked_req_to_exclude=list(chunked_req_to_exclude) - ) - if self.last_batch.batch_size() < last_bs: - self.running_batch.batch_is_full = False - - # Merge the new batch into the running batch. - if not self.last_batch.is_empty(): + if self.enable_hisparse: + ready_reqs = self.hisparse_coordinator.collect_ready_reqs() + if len(ready_reqs) > 0: + new_batch = self._build_hisparse_decode_batch(ready_reqs) if self.running_batch.is_empty(): - self.running_batch = self.last_batch + self.running_batch = new_batch else: - # Merge running_batch with prefill batch - self.running_batch.merge_batch(self.last_batch) + self.running_batch.merge_batch(new_batch) + self.running_batch.hisparse_coordinator = self.hisparse_coordinator + else: + if self.last_batch and self.last_batch.forward_mode.is_extend(): + if self.last_batch.chunked_req is not None: + # In the context pipeline parallelism, after the last chunk, the current microbatch still track outdated chunked_req. + # We need to discard it. + chunked_req_to_exclude.add(self.last_batch.chunked_req) + + if self.dllm_config is not None and self.last_batch.reqs: + chunked_req_to_exclude.update(self.last_batch.reqs) + + # Filter batch + last_bs = self.last_batch.batch_size() + self.last_batch.filter_batch( + chunked_req_to_exclude=list(chunked_req_to_exclude) + ) + if self.last_batch.batch_size() < last_bs: + self.running_batch.batch_is_full = False + + # Merge the new batch into the running batch. + if not self.last_batch.is_empty(): + if self.running_batch.is_empty(): + self.running_batch = self.last_batch + else: + # Merge running_batch with prefill batch + self.running_batch.merge_batch(self.last_batch) # For prefill-only batch, filter out finished requests since they # won't go through the decode step. This keeps running_batch accurate @@ -2444,6 +2501,8 @@ class Scheduler( for req in retracted_reqs: self._add_request_to_queue(req, is_retracted=True) + if self.enable_hisparse: + self.hisparse_coordinator.retract_req(req) else: self.new_token_ratio = max( self.new_token_ratio - self.new_token_ratio_decay, @@ -2980,6 +3039,7 @@ class Scheduler( return RpcReqOutput(success, "" if not exec else str(exec)) def abort_request(self, recv_req: AbortReq): + # todo hisparse, release resources for abort requests in hisparse coordinator # Delete requests in the waiting queue to_del = [] for i, req in enumerate(self.waiting_queue): @@ -2998,6 +3058,8 @@ class Scheduler( self.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req) # For disaggregation decode mode, the request in the waiting queue has KV cache allocated. if self.disaggregation_mode == DisaggregationMode.DECODE: + if self.enable_hisparse: + self.hisparse_coordinator.request_finished(req) release_kv_cache(req, self.tree_cache) # For disaggregation prefill mode, free the metadata buffer index if self.disaggregation_mode == DisaggregationMode.PREFILL: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index f4b2b9978..a5b4fb9f3 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -176,6 +176,8 @@ class SchedulerOutputProcessorMixin: req.time_stats.set_completion_time() elif not batch.decoding_reqs or req not in batch.decoding_reqs: self.tree_cache.cache_unfinished_req(req) + if self.enable_hisparse: + self.hisparse_coordinator.admit_request_into_staging(req) self.maybe_collect_customized_info(i, req, logits_output) @@ -448,6 +450,8 @@ class SchedulerOutputProcessorMixin: if not self.decode_offload_manager.offload_kv_cache(req): self.decode_offload_manager.finalize_release_on_finish(req) else: + if self.enable_hisparse: + self.hisparse_coordinator.request_finished(req) release_kv_cache(req, self.tree_cache) req.time_stats.set_completion_time() diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index dc8b7597e..8d01f7792 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -371,6 +371,9 @@ class SchedulerRuntimeCheckerMixin: queue_size += len(self.decode_offload_manager.ongoing_offload) if queue_size: return + elif self.enable_hisparse: + if self.hisparse_coordinator.has_ongoing_staging(): + return self.check_memory() self.check_tree_cache() diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 8f5500639..46a87aeb8 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -407,6 +407,9 @@ class TpModelWorker(BaseTpWorker): if self.hicache_layer_transfer_counter is not None: self.hicache_layer_transfer_counter.set_consumer(consumer_index) + def register_hisparse_coordinator(self, coordinator): + self.model_runner.hisparse_coordinator = coordinator + def get_worker_info(self): return ( self.max_total_num_tokens, diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py new file mode 100644 index 000000000..d1ada2f17 --- /dev/null +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -0,0 +1,341 @@ +# mapping on device memory, host memory and memory allocator + +import weakref +from typing import Optional + +import torch +from sgl_kernel.kvcacheio import transfer_kv_all_layer_mla + +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.mem_cache.allocator import ( + BaseTokenToKVPoolAllocator, + PagedTokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool + + +class HiSparseNSATokenToKVPool(NSATokenToKVPool): + def __init__( + self, + size: int, + page_size: int, + kv_lora_rank: int, + dtype: torch.dtype, + qk_rope_head_dim: int, + layer_num: int, + device: str, + index_head_dim: int, + enable_memory_saver: bool, + kv_cache_dim: int, + start_layer: Optional[int] = None, + end_layer: Optional[int] = None, + host_to_device_ratio: int = 2, + ): + super().__init__( + size=size, + page_size=page_size, + kv_lora_rank=kv_lora_rank, + dtype=dtype, + qk_rope_head_dim=qk_rope_head_dim, + layer_num=layer_num, + device=device, + index_head_dim=index_head_dim, + enable_memory_saver=enable_memory_saver, + kv_cache_dim=kv_cache_dim, + start_layer=start_layer, + end_layer=end_layer, + index_buf_size=size * host_to_device_ratio, + ) + self.bytes_per_token = self.kv_cache_dim * self.dtype.itemsize + + def register_mapping(self, full_to_hisparse_device_index_mapping: torch.Tensor): + self.full_to_hisparse_device_index_mapping = ( + full_to_hisparse_device_index_mapping + ) + + def translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor): + return self.full_to_hisparse_device_index_mapping[compressed_indices].to( + torch.int32 + ) + + def _translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor): + return self.full_to_hisparse_device_index_mapping[compressed_indices] + + def set_kv_buffer( + self, + layer: RadixAttention, + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + ): + loc = self.translate_loc_to_hisparse_device(loc) + super().set_kv_buffer(layer, loc, cache_k, cache_v) + + def set_mla_kv_buffer( + self, + layer: RadixAttention, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + ): + loc = self.translate_loc_to_hisparse_device(loc) + super().set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope) + + def get_mla_kv_buffer( + self, + layer: RadixAttention, + loc: torch.Tensor, + dst_dtype: Optional[torch.dtype] = None, + ): + loc = self.translate_loc_to_hisparse_device(loc) + return super().get_mla_kv_buffer(layer, loc, dst_dtype) + + def transfer_values_on_device(self, dst_indices, src_indices): + transfer_kv_all_layer_mla( + src_layers=self.data_ptrs, + dst_layers=self.data_ptrs, + src_indices=src_indices, + dst_indices=dst_indices, + item_size=self.bytes_per_token, + num_layers=self.layer_num, + ) + + def get_cpu_copy(self, indices): + raise NotImplementedError("HiSparseDevicePool does not support get_cpu_copy") + + def load_cpu_copy(self, kv_cache_cpu, indices): + raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy") + + +class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): + def __init__( + self, + size: int, + page_size: int, + dtype: torch.dtype, + device: torch.device, + kvcache: NSATokenToKVPool, + need_sort: bool, + host_to_device_ratio: int = 2, + ): + self._kvcache = kvcache + self._size_full = size * host_to_device_ratio + self._size_hisparse = size + self.dtype = dtype + self.device = device + self.page_size = page_size + self.need_sort = need_sort + + self.logical_attn_allocator = PagedTokenToKVPoolAllocator( + self._size_full, + self.page_size, + self.dtype, + self.device, + kvcache, + need_sort, + ) + + self.hisparse_attn_allocator = PagedTokenToKVPoolAllocator( + self._size_hisparse, + self.page_size, + self.dtype, + self.device, + kvcache, + need_sort, + ) + + self.full_to_hisparse_device_index_mapping = torch.cat( + [ + torch.zeros( + self._size_full + self.page_size, + dtype=torch.int64, + device=self.device, + ), + torch.tensor([-1], dtype=torch.int64, device=self.device), + ] + ) + + self.free_pages = None + self.release_pages = None + self.is_not_in_free_group = True + self.free_group = [] + self.clear() + + self._kvcache.register_mapping( + weakref.proxy(self.full_to_hisparse_device_index_mapping) + ) + + @property + def size_full(self) -> int: + return self._size_full + + def available_size(self) -> int: + return min( + self.logical_attn_allocator.available_size(), + self.hisparse_attn_allocator.available_size(), + ) + + def alloc(self, need_size: int): + raise NotImplementedError( + "Page size = 1 is not supported in HiSparse allocator" + ) + + def alloc_device_buffer(self, allocated_indices, need_size: int): + assert need_size % self.page_size == 0 + # clear original reference and isolate the buffer from outside addressing, allocate new buffer if needed + hisparse_indices = self.full_to_hisparse_device_index_mapping[allocated_indices] + self.full_to_hisparse_device_index_mapping[allocated_indices] = 0 + if len(hisparse_indices) >= need_size: + buffer_indices = hisparse_indices[:need_size] + self.free_hisparse_indices(hisparse_indices[need_size:]) + else: + # page alignment, claiming the residual space for an incomplete page + page_residual_length = len(hisparse_indices) % self.page_size + if page_residual_length != 0: + hisparse_indices = torch.cat( + [ + hisparse_indices, + torch.arange( + hisparse_indices[-1] + 1, + hisparse_indices[-1] + + self.page_size + - page_residual_length + + 1, + device=self.device, + ), + ] + ) + extra_indices = self.hisparse_attn_allocator.alloc( + need_size - len(hisparse_indices) + ) + assert ( + extra_indices is not None + ), "Hisparse allocation failed in alloc_device_buffer" + buffer_indices = torch.cat([hisparse_indices, extra_indices]) + return buffer_indices + + def free_hisparse_indices(self, buffer_indices: torch.Tensor): + # disable free group mechanism for device buffer free + self.hisparse_attn_allocator.is_not_in_free_group = True + self.hisparse_attn_allocator.free(buffer_indices[buffer_indices > 0]) + + def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): + hisparse_last_locs = self._kvcache._translate_loc_to_hisparse_device(last_locs) + return hisparse_last_locs + + def alloc_extend( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, # last_loc for full layers + extend_num_tokens: int, + ): + assert self.page_size > 1 + num_tokens = extend_num_tokens + len(seq_lens) * self.page_size + + if num_tokens > self.available_size(): + return None + + logical_indices = self.logical_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + ) + assert logical_indices is not None, "Logical allocation failed in alloc_extend" + + hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc) + hisparse_indices = self.hisparse_attn_allocator.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + hisparse_last_loc, + len(logical_indices), + ) + assert ( + hisparse_indices is not None + ), "Hisparse allocation failed in alloc_extend" + + self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices + + return logical_indices + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, # last_loc for full layers + ): + logical_indices = self.logical_attn_allocator.alloc_decode( + seq_lens, seq_lens_cpu, last_loc + ) + + return logical_indices + + def alloc_decode_debug( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, # last_loc for full layers + ): + logical_indices = self.logical_attn_allocator.alloc_decode( + seq_lens, seq_lens_cpu, last_loc + ) + + hisparse_last_loc = self.get_last_loc_hisparse_device(last_loc) + hisparse_indices = self.hisparse_attn_allocator.alloc_decode( + seq_lens, + seq_lens_cpu, + hisparse_last_loc, + ) + + if logical_indices is None or hisparse_indices is None: + return None + + self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices + + return logical_indices + + def free_hisparse(self, free_indices: torch.Tensor): + hisparse_indices = self._kvcache._translate_loc_to_hisparse_device(free_indices) + hisparse_indices = hisparse_indices[hisparse_indices > 0] + self.free_hisparse_indices(hisparse_indices) + self.full_to_hisparse_device_index_mapping[free_indices] = 0 + + def clear(self): + self.logical_attn_allocator.clear() + self.hisparse_attn_allocator.clear() + + # Note: the last item is -1, we don't clear it, see the comment in __init__ + self.full_to_hisparse_device_index_mapping[:-1].fill_(0) + self.is_not_in_free_group = True + self.free_group = [] + + def free_group_begin(self): + return + + def free_group_end(self): + return + + def free(self, free_index: torch.Tensor): + if free_index.numel() == 0: + return + + if self.is_not_in_free_group: + self.logical_attn_allocator.free(free_index) + self.free_hisparse(free_index) + else: + self.free_group.append(free_index) + assert ( + self.logical_attn_allocator.available_size() + <= self.logical_attn_allocator.size + ) + assert ( + self.hisparse_attn_allocator.available_size() + <= self.hisparse_attn_allocator.size + ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 140d1b514..51530e865 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1779,6 +1779,7 @@ class NSATokenToKVPool(MLATokenToKVPool): kv_cache_dim: int, start_layer: Optional[int] = None, end_layer: Optional[int] = None, + index_buf_size: Optional[int] = None, ): override_dim = ( @@ -1802,6 +1803,8 @@ class NSATokenToKVPool(MLATokenToKVPool): # self.index_k_dtype = torch.float8_e4m3fn # self.index_k_scale_dtype = torch.float32 self.index_head_dim = index_head_dim + if index_buf_size is None: + index_buf_size = size # num head == 1 and head dim == 128 for index_k in NSA assert index_head_dim == 128 @@ -1823,7 +1826,7 @@ class NSATokenToKVPool(MLATokenToKVPool): # * buf[i, :page_size * head_dim] for fp8 data # * buf[i, page_size * head_dim:].view(float32) for scale ( - (size + page_size + 1) // self.page_size, + (index_buf_size + page_size + 1) // self.page_size, self.page_size * ( index_head_dim + index_head_dim // self.quant_block_size * 4 diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d98824420..ea220bef0 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -268,7 +268,7 @@ class HostKVCache(abc.ABC): @synchronized def free(self, indices: torch.Tensor) -> int: - self.free_slots = torch.cat([self.free_slots, indices]) + self.free_slots = torch.cat([self.free_slots, indices.cpu()]) return len(indices) diff --git a/python/sglang/srt/mem_cache/sparsity/__init__.py b/python/sglang/srt/mem_cache/sparsity/__init__.py index 66e9ee899..e226ab5b9 100644 --- a/python/sglang/srt/mem_cache/sparsity/__init__.py +++ b/python/sglang/srt/mem_cache/sparsity/__init__.py @@ -9,6 +9,7 @@ from sglang.srt.mem_cache.sparsity.core import SparseConfig, SparseCoordinator from sglang.srt.mem_cache.sparsity.factory import ( create_sparse_coordinator, get_sparse_coordinator, + parse_hisparse_config, register_sparse_coordinator, ) @@ -23,5 +24,6 @@ __all__ = [ "SparseCoordinator", "create_sparse_coordinator", "get_sparse_coordinator", + "parse_hisparse_config", "register_sparse_coordinator", ] diff --git a/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py b/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py index f3c37dc2c..cf7c1d06c 100644 --- a/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py +++ b/python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py @@ -55,10 +55,13 @@ class RequestTrackers: class SparseConfig: """Configuration for sparse attention.""" - backend: str - algorithm: str - page_size: int = 64 - min_sparse_prompt_len: int = 2048 + top_k: int = 2048 + device_buffer_size: int = 4096 + host_to_device_ratio: int = 2 + algorithm: Optional[str] = None + backend: Optional[str] = None + page_size: Optional[int] = None + min_sparse_prompt_len: Optional[int] = None sparse_extra_config: dict = field( default_factory=dict ) # Algorithm-specific config, parsed by each algorithm diff --git a/python/sglang/srt/mem_cache/sparsity/factory.py b/python/sglang/srt/mem_cache/sparsity/factory.py index 39d0f4682..c8b29f041 100644 --- a/python/sglang/srt/mem_cache/sparsity/factory.py +++ b/python/sglang/srt/mem_cache/sparsity/factory.py @@ -59,33 +59,51 @@ def _create_backend_adaptor( def _parse_sparse_config(server_args) -> SparseConfig: - """Parse hierarchical sparse config""" - # Parse extra config if provided - extra_config_str = server_args.hierarchical_sparse_attention_extra_config + """Parse hierarchical sparse config from JSON string. + + Required fields with defaults: top_k (2048), device_buffer_size (2*top_k), + host_to_device_ratio (2). + Optional fields (default None): algorithm, backend, min_sparse_prompt_len, + page_size. All remaining fields go to sparse_extra_config. + """ + extra_config_str = server_args.hisparse_config if extra_config_str is not None: try: extra_config = json.loads(extra_config_str) - - # Extract algorithm and backend - algorithm = extra_config.pop("algorithm", "quest") - backend = extra_config.pop("backend", "flashattention") - min_sparse_prompt_len = extra_config.pop("min_sparse_prompt_len", 2048) - - # Everything else goes to algorithm_extra_config - sparse_extra_config = extra_config except json.JSONDecodeError as e: - logger.warning( - f"Failed to parse hierarchical_sparse_attention_extra_config: {e}" - ) + raise ValueError(f"Failed to parse hisparse_config: {e}") from e + else: + extra_config = {} - config = SparseConfig( + top_k = extra_config.pop("top_k", 2048) + device_buffer_size = extra_config.pop("device_buffer_size", 2 * top_k) + host_to_device_ratio = extra_config.pop("host_to_device_ratio", 2) + + if device_buffer_size < top_k: + raise ValueError( + f"device_buffer_size ({device_buffer_size}) must be no smaller than top_k ({top_k})" + ) + + algorithm = extra_config.pop("algorithm", None) + backend = extra_config.pop("backend", None) + min_sparse_prompt_len = extra_config.pop("min_sparse_prompt_len", None) + page_size = extra_config.pop("page_size", None) + + return SparseConfig( + top_k=top_k, + device_buffer_size=device_buffer_size, + host_to_device_ratio=host_to_device_ratio, algorithm=algorithm, backend=backend, - page_size=server_args.page_size, + page_size=page_size, min_sparse_prompt_len=min_sparse_prompt_len, - sparse_extra_config=sparse_extra_config, + sparse_extra_config=extra_config, ) - return config + + +def parse_hisparse_config(server_args) -> SparseConfig: + """Parse hisparse config from server_args, returning defaults if no config provided.""" + return _parse_sparse_config(server_args) def create_sparse_coordinator( diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 05e2d5887..ea8c93963 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -955,6 +955,12 @@ class CudaGraphRunner: global_forward_mode=self.capture_forward_mode, lora_ids=lora_ids, ) + + # HiSparse: set coordinator so the hisparse code path is captured into the graph + forward_batch.hisparse_coordinator = self.model_runner.hisparse_coordinator + if forward_batch.hisparse_coordinator is not None: + forward_batch.hisparse_coordinator.num_real_reqs.fill_(bs) + if buffers.ngram_embedding_info is not None: forward_batch.ngram_embedding_info = buffers.ngram_embedding_info.slice(bs) @@ -1121,6 +1127,9 @@ class CudaGraphRunner: self.raw_num_token = raw_num_token self.bs = bs + if self.model_runner.hisparse_coordinator is not None: + self.model_runner.hisparse_coordinator.num_real_reqs.fill_(raw_bs) + def replay( self, forward_batch: ForwardBatch, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 54211238f..bec834965 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -68,6 +68,7 @@ from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator from sglang.srt.managers.schedule_batch import ModelWorkerBatch, MultimodalInputs from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool from sglang.srt.model_executor.model_runner import ModelRunner @@ -422,6 +423,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # For hidden states before normal return_hidden_states_before_norm: bool = False + # For hisparse + hisparse_coordinator: Optional[HiSparseCoordinator] = None + # For ngram embedding ngram_embedding_info: Optional[NgramEmbeddingInfo] = None diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 93a5ccce0..fc9afafac 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -345,6 +345,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.forward_pass_id = 0 self.init_new_workspace = False self.draft_model_idx = draft_model_idx + self.enable_hisparse = server_args.enable_hisparse self.remote_instance_transfer_engine = None self.remote_instance_transfer_engine_session_id = "" @@ -418,6 +419,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM: deep_gemm_wrapper.update_deep_gemm_config(gpu_id, server_args) + # For hisparse (must be set before initialize() so CUDA graph capture can see it) + self.hisparse_coordinator = None + # Initialize the model runner self.initialize(pre_model_load_memory) self.check_quantized_moe_compatibility() @@ -611,6 +615,26 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Init ngram embedding token table self.maybe_init_ngram_embedding() + # Init hisparse coordinator (must happen before CUDA graph capture) + if self.enable_hisparse: + from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator + from sglang.srt.mem_cache.sparsity import parse_hisparse_config + + hisparse_cfg = parse_hisparse_config(self.server_args) + self.hisparse_coordinator = HiSparseCoordinator( + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + top_k=hisparse_cfg.top_k, + device_buffer_size=hisparse_cfg.device_buffer_size, + device=self.device, + tp_group=( + self.attention_tp_group.cpu_group + if self.server_args.enable_dp_attention + else self.tp_group.cpu_group + ), + host_to_device_ratio=hisparse_cfg.host_to_device_ratio, + ) + # Init routed experts capturer self.init_routed_experts_capturer() @@ -2686,6 +2710,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): if forward_batch.out_cache_loc_swa is not None: self.token_to_kv_pool.set_swa_loc(forward_batch.out_cache_loc_swa) + forward_batch.hisparse_coordinator = self.hisparse_coordinator + if self.hisparse_coordinator is not None: + self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size) + if forward_batch.forward_mode.is_decode(): ret = self.forward_decode( forward_batch, diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 588da8034..a021c6561 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -13,6 +13,10 @@ from sglang.srt.mem_cache.allocator import ( PagedTokenToKVPoolAllocator, TokenToKVPoolAllocator, ) +from sglang.srt.mem_cache.hisparse_memory_pool import ( + HiSparseNSATokenToKVPool, + HiSparseTokenToKVPoolAllocator, +) from sglang.srt.mem_cache.memory_pool import ( DoubleSparseTokenToKVPool, HybridLinearKVPool, @@ -481,8 +485,8 @@ class ModelRunnerKVCacheMixin: end_layer=self.end_layer, ) elif self.use_mla_backend and is_nsa_model: - self.token_to_kv_pool = NSATokenToKVPool( - self.max_total_num_tokens, + nsa_pool_kwargs = dict( + size=self.max_total_num_tokens, page_size=self.page_size, dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, @@ -495,6 +499,16 @@ class ModelRunnerKVCacheMixin: end_layer=self.end_layer, index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config), ) + if self.enable_hisparse: + from sglang.srt.mem_cache.sparsity import parse_hisparse_config + + hisparse_cfg = parse_hisparse_config(self.server_args) + nsa_pool_kwargs["host_to_device_ratio"] = ( + hisparse_cfg.host_to_device_ratio + ) + self.token_to_kv_pool = HiSparseNSATokenToKVPool(**nsa_pool_kwargs) + else: + self.token_to_kv_pool = NSATokenToKVPool(**nsa_pool_kwargs) elif self.use_mla_backend and not self.mambaish_config: assert not is_nsa_model if is_float4_e2m1fn_x2(self.kv_cache_dtype): @@ -669,7 +683,24 @@ class ModelRunnerKVCacheMixin: need_sort=need_sort, ) else: - if self.page_size == 1: + if self.enable_hisparse: + from sglang.srt.mem_cache.sparsity import ( + parse_hisparse_config, + ) + + hisparse_cfg = parse_hisparse_config(self.server_args) + self.token_to_kv_pool_allocator = ( + HiSparseTokenToKVPoolAllocator( + self.max_total_num_tokens, + page_size=self.page_size, + dtype=self.kv_cache_dtype, + device=self.device, + kvcache=self.token_to_kv_pool, + need_sort=need_sort, + host_to_device_ratio=hisparse_cfg.host_to_device_ratio, + ) + ) + elif self.page_size == 1: self.token_to_kv_pool_allocator = TokenToKVPoolAllocator( self.max_total_num_tokens, dtype=self.kv_cache_dtype, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7a775028c..3a8d7fbf8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -562,7 +562,8 @@ class ServerArgs: hicache_storage_backend_extra_config: Optional[str] = None # Hierarchical sparse attention - hierarchical_sparse_attention_extra_config: Optional[str] = None + enable_hisparse: bool = False + hisparse_config: Optional[str] = None # LMCache enable_lmcache: bool = False @@ -5073,13 +5074,17 @@ class ServerArgs: # Hierarchical sparse attention parser.add_argument( - "--hierarchical-sparse-attention-extra-config", + "--enable-hisparse", + action="store_true", + help="Enable hierarchical sparse attention", + ) + + parser.add_argument( + "--hisparse-config", type=str, - default=ServerArgs.hierarchical_sparse_attention_extra_config, + default=ServerArgs.hisparse_config, help="A dictionary in JSON string format for hierarchical sparse attention configuration. " - "Required fields: algorithm (str), backend (str). " - "All other fields are algorithm-specific and passed to the algorithm constructor. " - 'Example: \'{"algorithm": "quest", "backend": "flashattention", "sparsity_ratio": 0.7, "min_sparse_prompt_len": 2048}\'', + 'Example: \'{"top_k": 2048, "device_buffer_size": 4096}\'', ) # LMCache @@ -6052,6 +6057,12 @@ class ServerArgs: "Please set --chunked-prefill-size -1 when using --multi-item-scoring-delimiter." ) + # Check hisparse + if self.enable_hisparse: + assert ( + self.disable_radix_cache + ), "Hierarchical sparse attention currently requires --disable-radix-cache." + assert ( self.schedule_conservativeness >= 0 ), "schedule_conservativeness must be non-negative"