HiSparse for Sparse Attention (#20343)

python/sglang/jit_kernel/csrc/hisparse.cuh

@@ -0,0 +1,390 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cuda_runtime.h>
+#include <stdexcept>
+#include <stdint.h>
+#include <string>
+
+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<const uint64_t*>(src_addr);
+  uint64_t* __restrict__ dst = static_cast<uint64_t*>(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 <int BLOCK_SIZE, int NUM_TOP_K, int HOT_BUFFER_SIZE, bool IsMLA>
+__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<int>(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<int16_t>(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<int32_t>(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<int64_t>(req_device_buffer_locs[evict_slot]);
+
+    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, int NUM_TOP_K, int HOT_BUFFER_SIZE, bool IsMLA>
+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<BLOCK_SIZE, NUM_TOP_K, HOT_BUFFER_SIZE, IsMLA>,
+      static_cast<const int32_t*>(top_k_tokens.data_ptr()),
+      static_cast<int32_t*>(device_buffer_tokens.data_ptr()),
+      static_cast<const int64_t*>(host_cache_locs.data_ptr()),
+      static_cast<const int32_t*>(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<int32_t*>(top_k_device_locs.data_ptr()),
+      static_cast<const int64_t*>(req_pool_indices.data_ptr()),
+      static_cast<const int32_t*>(seq_lens.data_ptr()),
+      static_cast<int16_t*>(lru_slots.data_ptr()),
+      static_cast<const int32_t*>(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

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,
+    )

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):

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

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(

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:

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()

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()

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,

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
+        )

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

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)
 
 

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",
 ]

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

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(

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,

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
 

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,

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,

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"