[Sparse & HICache]: Enables hierarchical sparse KV cache management and scheduling for DeepSeek V32.

python/sglang/srt/disaggregation/decode.py

@@ -49,9 +49,15 @@
 from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.managers.schedule_batch import FINISH_ABORT, RequestStage, ScheduleBatch
 from sglang.srt.managers.utils import GenerationBatchResult
-from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator import (
+    BaseTokenToKVPoolAllocator,
+    is_enable_hierarchical_nsa,
+)
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
-from sglang.srt.mem_cache.common import release_kv_cache
+from sglang.srt.mem_cache.common import (
+    release_kv_cache,
+    truncate_kv_cache_after_prefill,
+)
 from sglang.srt.mem_cache.memory_pool import (
     HybridLinearKVPool,
     HybridReqToTokenPool,
@@ -60,6 +66,7 @@
     ReqToTokenPool,
     SWAKVPool,
 )
+from sglang.srt.mem_cache.sparsity import get_sparse_coordinator
 from sglang.srt.tracing.trace import trace_event_batch, trace_slice_end
 from sglang.srt.utils import get_int_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
@@ -134,6 +141,40 @@ def clear(self):
         self.free_slots = list(range(self.size + self.pre_alloc_size))
 
 
+class NSADecodeReqToTokenPool(DecodeReqToTokenPool):
+    """NSA DecodeReqToTokenPool: separate mapping for KV cache and nsa indexer_k"""
+
+    def __init__(
+        self,
+        size: int,
+        max_context_len: int,
+        device: str,
+        enable_memory_saver: bool,
+        pre_alloc_size: int,
+    ):
+        super().__init__(
+            size, max_context_len, device, enable_memory_saver, pre_alloc_size
+        )
+
+        memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+        )
+        with memory_saver_adapter.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
+            self.req_to_nsa_index_k = torch.zeros(
+                (size + pre_alloc_size, max_context_len),
+                dtype=torch.int32,
+                device=device,
+            )
+
+    def write_index_token(self, indices, values):
+        """Write indexer_k mapping"""
+        self.req_to_nsa_index_k[indices] = values
+
+    def clear(self):
+        super().clear()
+        self.req_to_nsa_index_k.zero_()
+
+
 class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
 
     def __init__(
@@ -510,9 +551,14 @@ def pop_preallocated(self) -> List[DecodeRequest]:
                 state_indices = kv_to_page_indices(state_indices, page_size)
             elif isinstance(self.token_to_kv_pool, NSATokenToKVPool):
                 seq_len = len(decode_req.req.origin_input_ids)
-                kv_indices_full = self.req_to_token_pool.req_to_token[
-                    decode_req.req.req_pool_idx, :seq_len
-                ]
+                if isinstance(self.req_to_token_pool, NSADecodeReqToTokenPool):
+                    kv_indices_full = self.req_to_token_pool.req_to_nsa_index_k[
+                        decode_req.req.req_pool_idx, :seq_len
+                    ]
+                else:
+                    kv_indices_full = self.req_to_token_pool.req_to_token[
+                        decode_req.req.req_pool_idx, :seq_len
+                    ]
                 state_indices = kv_indices_full.cpu().numpy()
                 state_indices = kv_to_page_indices(state_indices, page_size)
             else:
@@ -624,10 +670,10 @@ def _pre_alloc(self, req: Req) -> torch.Tensor:
         req.kv_allocated_len = fill_len
         req.kv_committed_len = fill_len
         if self.token_to_kv_pool_allocator.page_size == 1:
-            kv_loc = self.token_to_kv_pool_allocator.alloc(fill_len)
+            alloc_result = self.token_to_kv_pool_allocator.alloc(fill_len)
         else:
             device = self.token_to_kv_pool_allocator.device
-            kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
+            alloc_result = self.token_to_kv_pool_allocator.alloc_extend(
                 prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
                 prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
                 seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
@@ -637,11 +683,25 @@ def _pre_alloc(self, req: Req) -> torch.Tensor:
             )
 
         assert (
-            kv_loc is not None
+            alloc_result is not None
         ), "KV cache is full! There is a bug in memory estimation."
 
+        if is_enable_hierarchical_nsa(self.token_to_kv_pool_allocator):
+            kv_loc, index_k_loc = alloc_result
+        else:
+            kv_loc = alloc_result
+            index_k_loc = None
+
+        # Write KV indices to req_to_token
         self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
 
+        # Write index_k indices for NSA
+        if index_k_loc is not None:
+            self.req_to_token_pool.write_index_token(
+                (req.req_pool_idx, slice(0, len(index_k_loc))),
+                index_k_loc.to(torch.int32),
+            )
+
         # populate metadata
         req.fill_ids = req.origin_input_ids + req.output_ids
         req.extend_input_len = len(req.origin_input_ids)
@@ -959,4 +1019,15 @@ def process_decode_queue(self: Scheduler):
             alloc_reqs = (
                 self.disagg_decode_transfer_queue.pop_transferred()
             )  # the requests which kv has arrived
+
+            # NSA: Register, Offload and Truncate after KV transfer completes
+            sparse_coordinator = get_sparse_coordinator()
+            if sparse_coordinator is not None:
+                for req in alloc_reqs:
+                    sparse_coordinator.on_request_begin(req)
+                    sparse_coordinator.on_request_prefill_end(req)
+                    truncate_kv_cache_after_prefill(
+                        req, self.req_to_token_pool, self.tree_cache
+                    )
+
             self.waiting_queue.extend(alloc_reqs)

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -11,6 +11,7 @@
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.radix_attention import AttentionType
+from sglang.srt.mem_cache.sparsity import DeepSeekNSAAlgorithm, get_sparse_coordinator
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpecInput
@@ -362,6 +363,12 @@ def __init__(
             1 if model_runner.server_args.enable_deterministic_inference else 0
         )
 
+        # Sparse attention coordinator
+        self.sparse_coordinator = get_sparse_coordinator()
+        if self.sparse_coordinator is not None:
+            if isinstance(self.sparse_coordinator.algorithm, DeepSeekNSAAlgorithm):
+                self.sparse_coordinator = None
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Initialize forward metadata hence all layers in the forward pass can reuse it."""
         metadata = FlashAttentionMetadata()
@@ -961,6 +968,13 @@ def forward_extend(
                 else:
                     o = result
 
+        if self.sparse_coordinator is not None:
+            self.sparse_coordinator.attention_end(
+                output=o,
+                layer=layer,
+                forward_batch=forward_batch,
+            )
+
         return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
     def forward_decode(
@@ -999,6 +1013,18 @@ def forward_decode(
 
         # Use precomputed metadata across all layers
         metadata = self.forward_metadata
+
+        # Apply sparse attention: modify metadata based on query
+        if self.sparse_coordinator is not None:
+            self.sparse_coordinator.attention_begin(
+                query=q,
+                key=k,
+                value=v,
+                layer=layer,
+                forward_batch=forward_batch,
+                attn_metadata=metadata,
+            )
+
         local_attn_metadata = getattr(metadata, "local_attn_metadata", None)
         use_local_attn = (
             self.attention_chunk_size is not None
@@ -1231,6 +1257,13 @@ def forward_decode(
             else:
                 o = result
 
+        if self.sparse_coordinator is not None:
+            self.sparse_coordinator.attention_end(
+                output=o,
+                layer=layer,
+                forward_batch=forward_batch,
+            )
+
         return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -440,11 +440,10 @@ def _forward_cuda_k_only(
         key = self._get_k_bf16(x, positions, enable_dual_stream)
         k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt)
 
-        if not forward_batch.out_cache_loc.is_contiguous():
-            forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous()
+        index_loc = self._get_index_cache_loc(forward_batch)
         forward_batch.token_to_kv_pool.set_index_k_scale_buffer(
             layer_id=layer_id,
-            loc=forward_batch.out_cache_loc,
+            loc=index_loc,
             index_k=k_fp8,
             index_k_scale=k_scale,
         )
@@ -621,11 +620,10 @@ def forward_cuda(
         # k_buffer: (num_total_tokens + page_size, head_dim) fp8_e4m3fn
         # k_scale: (seq_len, head_dim // block_size = 1) fp8_e4m3fn
         # k_scale_cache: (num_total_tokens + page_size, head_dim // block_size = 1) fp8_e4m3fn
-        if not forward_batch.out_cache_loc.is_contiguous():
-            forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous()
+        index_loc = self._get_index_cache_loc(forward_batch)
         forward_batch.token_to_kv_pool.set_index_k_scale_buffer(
             layer_id=layer_id,
-            loc=forward_batch.out_cache_loc,
+            loc=index_loc,
             index_k=k_fp8,
             index_k_scale=k_scale,
         )
@@ -667,6 +665,23 @@ def forward_cuda(
             )
         return topk_result
 
+    def _get_index_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
+        pool = forward_batch.req_to_token_pool
+
+        if (
+            forward_batch.forward_mode.is_decode()
+            and hasattr(pool, "req_to_nsa_index_k")
+        ):
+            index_loc = pool.req_to_nsa_index_k[
+                forward_batch.req_pool_indices, forward_batch.seq_lens - 1
+            ].to(torch.int64)
+        else:
+            index_loc = forward_batch.out_cache_loc
+
+        if not index_loc.is_contiguous():
+            index_loc = index_loc.contiguous()
+        return index_loc
+
     def forward_npu(
         self,
         x: torch.Tensor,