[Feat] PD/HiSparse: Support hybrid DSA KV pools in PD serving

python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh

@@ -350,6 +350,7 @@ template <
     bool TopKIsBlocks,
     bool RecordMissPlan,
     bool SkipIO,
+    int NUM_TAIL_SLOTS,
     typename SeqLensT,
     typename ReqPoolIndicesT>
 __global__ void load_cache_to_device_buffer_kernel(
@@ -378,6 +379,8 @@ __global__ void load_cache_to_device_buffer_kernel(
     int32_t* __restrict__ miss_count_out,
     int64_t plan_stride) {
   static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA).");
+  static_assert(NUM_TAIL_SLOTS >= 0 && NUM_TAIL_SLOTS < NUM_TOP_K, "Invalid top-k tail slots.");
+  static_assert(!TopKIsBlocks || NUM_TAIL_SLOTS == 0, "Block top-k does not use tail columns.");
   static_assert(SPARSE_BLOCK_SIZE > 0, "SPARSE_BLOCK_SIZE must be positive.");
   // Cache residency and LRU replacement remain token-granular even when the
   // sparse-attention selection arrives as block ids.
@@ -425,7 +428,7 @@ __global__ void load_cache_to_device_buffer_kernel(
         if (token_pos >= 0 && token_pos < seq_len) {
           device_loc = req_device_buffer_locs[token_pos];
         }
-      } else if (i < count && token_pos >= 0) {
+      } else if ((i < count || i >= NUM_TOP_K_TOKENS - NUM_TAIL_SLOTS) && token_pos >= 0 && token_pos < seq_len) {
         device_loc = req_device_buffer_locs[token_pos];
       }
       req_top_k_device_locs[i] = device_loc;
@@ -483,12 +486,10 @@ __global__ void load_cache_to_device_buffer_kernel(
   // Insert top-k tokens into shared-memory hash table.
   for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
     const int32_t token_idx = resolve_selected_token<SPARSE_BLOCK_SIZE, TopKIsBlocks>(req_top_k, i);
-    if constexpr (TopKIsBlocks) {
-      if (token_idx < 0 || token_idx >= seq_len) {
-        s_top_k_tokens[i] = TOKEN_HIT;
-        req_top_k_device_locs[i] = -1;
-        continue;
-      }
+    if (token_idx < 0 || token_idx >= seq_len) {
+      s_top_k_tokens[i] = TOKEN_HIT;
+      req_top_k_device_locs[i] = -1;
+      continue;
     }
     if (token_idx == newest_token) {
       // If topk includes the latest token, bind its canonical occurrence to newest_slot (at HOT_BUFFER_SIZE) and mark
@@ -558,23 +559,13 @@ __global__ void load_cache_to_device_buffer_kernel(
     __syncthreads();
 
     if (warp_id == 0) {
-#ifdef USE_ROCM
-      // ROCm wavefront64: WARP_SIZE (64) > NUM_WARPS (16 at block_size=1024),
-      // so the wide-count form below would let lanes beyond this iteration's
-      // NUM_WARPS-wide window write the accumulator into s_chunk_offset
-      // positions belonging to future iterations, corrupting their reads.
-      // Bound the scan window to NUM_WARPS lanes.
+      // Each iteration produces NUM_WARPS chunks, not a full warp of chunks.
       const int scan_offset = iter * NUM_WARPS + 1;
       const int scan_count = min(scan_offset + NUM_WARPS, NUM_BUFFER_CHUNKS + 1);
       total_hit_count = warp_inclusive_scan(s_chunk_offset, lane_id, scan_offset, scan_count, total_hit_count);
       total_evict_count =
           warp_inclusive_scan(s_evict_chunk_offset, lane_id, scan_offset, scan_count, total_evict_count);
-#else
-      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);
-#endif
+
       if (tid == 0) {
         s_total_hits = total_hit_count;
       }
@@ -633,13 +624,10 @@ __global__ void load_cache_to_device_buffer_kernel(
     __syncthreads();
 
     if (warp_id == 0) {
-#ifdef USE_ROCM
       const int scan_offset = iter * NUM_WARPS + 1;
       const int scan_count = min(scan_offset + NUM_WARPS, NUM_TOKEN_CHUNKS + 1);
       total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, scan_offset, scan_count, total_misses);
-#else
-      total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses);
-#endif
+
       if (tid == 0) {
         s_total_misses = total_misses;
       }
@@ -730,7 +718,8 @@ template <
     int SPARSE_BLOCK_SIZE,
     bool TopKIsBlocks,
     bool RecordMissPlan,
-    bool SkipIO>
+    bool SkipIO,
+    int NUM_TAIL_SLOTS = 0>
 void load_cache_to_device_buffer(
     tvm::ffi::TensorView top_k,
     tvm::ffi::TensorView device_buffer_tokens,
@@ -830,6 +819,7 @@ void load_cache_to_device_buffer(
             TopKIsBlocks,
             RecordMissPlan,
             SkipIO,
+            NUM_TAIL_SLOTS,
             int64_t,
             int64_t>,
         static_cast<const int64_t*>(seq_lens.data_ptr()),
@@ -846,6 +836,7 @@ void load_cache_to_device_buffer(
             TopKIsBlocks,
             RecordMissPlan,
             SkipIO,
+            NUM_TAIL_SLOTS,
             int64_t,
             int32_t>,
         static_cast<const int64_t*>(seq_lens.data_ptr()),
@@ -862,6 +853,7 @@ void load_cache_to_device_buffer(
             TopKIsBlocks,
             RecordMissPlan,
             SkipIO,
+            NUM_TAIL_SLOTS,
             int32_t,
             int64_t>,
         static_cast<const int32_t*>(seq_lens.data_ptr()),
@@ -878,6 +870,7 @@ void load_cache_to_device_buffer(
             TopKIsBlocks,
             RecordMissPlan,
             SkipIO,
+            NUM_TAIL_SLOTS,
             int32_t,
             int32_t>,
         static_cast<const int32_t*>(seq_lens.data_ptr()),

python/sglang/kernels/ops/kvcache/hisparse.py

@@ -178,9 +178,9 @@ def _jit_sparse_module(
     top_k_is_blocks: bool = False,
     record_miss_plan: bool = False,
     skip_io: bool = False,
+    num_tail_slots: int = 0,
 ) -> Module:
-    # record_miss_plan / skip_io are compile-time kernel flags; the
-    # (False, False) production instantiation stays byte-identical.
+    # Compile-time options preserve the ordinary top-k layout by default.
     template_args = make_cpp_args(
         block_size,
         num_top_k,
@@ -191,6 +191,7 @@ def _jit_sparse_module(
         top_k_is_blocks,
         record_miss_plan,
         skip_io,
+        num_tail_slots,
     )
     cache_args = make_cpp_args(
         item_size_bytes,
@@ -203,6 +204,7 @@ def _jit_sparse_module(
         top_k_is_blocks,
         record_miss_plan,
         skip_io,
+        num_tail_slots,
     )
     return load_jit(
         "sparse_cache",
@@ -297,11 +299,26 @@ def _load_cache_to_device_buffer_mla(
     miss_dst: torch.Tensor | None,
     miss_count: torch.Tensor | None,
     skip_io: bool,
+    num_tail_slots: int = 0,
 ) -> None:
     assert hot_buffer_size >= num_top_k, (
         f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})"
     )
 
+    if (
+        top_k_tokens.ndim != 2
+        or top_k_tokens.shape[1] != num_top_k
+        or top_k_device_locs.shape != top_k_tokens.shape
+        or top_k_tokens.dtype != torch.int32
+        or top_k_device_locs.dtype != torch.int32
+        or top_k_tokens.stride(1) != 1
+        or top_k_device_locs.stride(1) != 1
+    ):
+        raise ValueError(
+            "HiSparse top-k must be int32 [batch, num_top_k] with contiguous columns."
+        )
+    if not 0 <= num_tail_slots < num_top_k:
+        raise ValueError("HiSparse tail slots must be in [0, num_top_k).")
     record_miss_plan = miss_src is not None
     module = _jit_sparse_module(
         item_size_bytes,
@@ -312,6 +329,7 @@ def _load_cache_to_device_buffer_mla(
         is_dsv4_layout=is_dsv4_layout,
         record_miss_plan=record_miss_plan,
         skip_io=skip_io,
+        num_tail_slots=num_tail_slots,
     )
 
     empty = torch.empty(0, device=top_k_tokens.device)
@@ -374,9 +392,13 @@ def load_cache_to_device_buffer_mla(
     miss_dst: torch.Tensor | None = None,
     miss_count: torch.Tensor | None = None,
     skip_io: bool = False,
+    num_tail_slots: int = 0,
 ) -> None:
     """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes).
 
+    num_tail_slots counts trailing KPool columns, valid independently of the
+    ordinary top-k prefix for short sequences.
+
     Optional miss_src/miss_dst/miss_count record the miss plan for replay by
     copy_cache_planned_mla; skip_io elides only the KV bytes (timing probe).
     """
@@ -402,6 +424,7 @@ def load_cache_to_device_buffer_mla(
         miss_dst=miss_dst,
         miss_count=miss_count,
         skip_io=skip_io,
+        num_tail_slots=num_tail_slots,
     )
 
 

python/sglang/srt/arg_groups/hisparse_hook.py

@@ -17,7 +17,7 @@
 
 HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE = {
     "bfloat16": {"flashmla_sparse"},
-    "fp8_e4m3": {"flashmla_kv", "flashinfer_sparse_mla"},
+    "fp8_e4m3": {"flashmla_kv", "flashinfer_sparse_mla", "tilelang"},
 }
 HISPARSE_ROCM_DSA_BACKENDS = {"tilelang", "triton", "aiter"}
 HISPARSE_XPU_DSA_BACKENDS_BY_DTYPE = {"bfloat16": {"intel_xpu"}}

python/sglang/srt/arg_groups/overrides.py

@@ -684,15 +684,7 @@ def _dsa_split_backend_resolution(view: Any) -> dict:
             declared["dsa_decode_backend"] = backend
         prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend)
         decode = declared.get("dsa_decode_backend", view.dsa_decode_backend)
-        if (
-            not get_platform().is_hip
-            and kv_cache_dtype == "fp8_e4m3"
-            and "tilelang" in (prefill, decode)
-        ):
-            raise ValueError("CUDA TileLang FP8 KV does not support HiSparse.")
-        # The hisparse allow-list in hisparse_hook is platform- but not
-        # dtype-aware, so an explicitly requested backend still has to clear the
-        # shared backend/kv-cache-dtype rules before this arm returns early.
+        # HiSparse also requires compatible KV layouts on both DSA backends.
         _check_dsa_backend_constraints(
             kv_cache_dtype,
             prefill,

python/sglang/srt/disaggregation/base/conn.py

@@ -12,6 +12,7 @@
 
 if TYPE_CHECKING:
     from sglang.srt.disaggregation.mode import DisaggregationMode
+    from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
 
 
 class StateType(str, enum.Enum):
@@ -66,6 +67,8 @@ class KVArgs:
     host_kv_data_ptrs: Optional[List[int]] = None
     host_kv_data_lens: Optional[List[int]] = None
     host_kv_item_lens: Optional[List[int]] = None
+    # Decode-owned host KV pool exported for same-node NIXL transfers.
+    local_shm_pool: Optional[MLATokenToKVPoolHost] = None
     host_state_data_ptrs: Optional[List[List[int]]] = None
     host_state_data_lens: Optional[List[List[int]]] = None
     host_state_item_lens: Optional[List[List[int]]] = None
@@ -163,6 +166,9 @@ def register_to_bootstrap(self):
     # (ABORT_ACK carrying the sender rank). Without it a hold only ends on timeout.
     supports_deferred_decode_kv_release: bool = False
 
+    def release_host_resources(self) -> None:
+        """Release backend-owned host mappings on graceful shutdown."""
+
     def teardown(self) -> None:
         raise NotImplementedError(
             f"{type(self).__name__} does not support PD role switch teardown"

python/sglang/srt/disaggregation/decode.py

@@ -749,6 +749,11 @@ def _init_kv_manager(self) -> CommonKVManager:
         )
         if self.transfer_backend == TransferBackend.NIXL:
             kv_args.kv_data_mem_kinds = kv_data_mem_kinds
+            if (
+                envs.SGLANG_DISAGGREGATION_NIXL_LOCAL_SHM.get()
+                and self.scheduler.enable_hisparse
+            ):
+                kv_args.local_shm_pool = transfer_kv_pool
         kv_args.page_size = self.token_to_kv_pool.page_size
 
         kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
@@ -3131,6 +3136,13 @@ def _init_staging_handler(self, kv_manager):
         )
         kv_manager._staging_handler = self.staging_handler
 
+    def _requires_host_drain(self, decode_req: DecodeRequest) -> bool:
+        # HiSparse receives KV directly into its own host pool, independently
+        # of the HiCache host-staging path.
+        return self.scheduler.enable_hisparse or (
+            self.enable_host_receive and decode_req.host_staged
+        )
+
     def _release_request(self, decode_req: DecodeRequest) -> None:
         # HiSparse host pages are the transfer destination; free them with the KV.
         if self.scheduler.enable_hisparse:
@@ -3262,7 +3274,7 @@ def pop_transferred_with_outcomes(
                     [decode_req.req],
                     decode_req.req.return_logprob,
                 )
-                requires_host_drain = self._is_host_destination(decode_req)
+                requires_host_drain = self._requires_host_drain(decode_req)
                 deferred_release_enabled = (
                     self.enable_deferred_kv_release
                     and decode_req.kv_receiver.kv_mgr.enable_deferred_decode_kv_release
@@ -3282,6 +3294,10 @@ def pop_transferred_with_outcomes(
                     # already hold destinations; this path always reaches
                     # _do_release, which clears the tracker.
                     decode_req.kv_receiver.ensure_abort_notified(force_arm=True)
+                if self.scheduler.enable_hisparse and deferred_release_enabled:
+                    # No notified peer means metadata was never published; no
+                    # remote write can still reference these host pages.
+                    requires_host_drain = decode_req.kv_receiver.abort_notified
                 failed.append(
                     (
                         i,
@@ -3458,9 +3474,15 @@ def resolve_deferred_releases(self) -> None:
         ready = []
         releases = list(self._deferred_releases.items())
         for decode_req, _, deadline, _, required_acks in releases:
-            if self._is_host_destination(decode_req):
+            if self._requires_host_drain(decode_req):
                 # A timeout cannot prove that a remote write has stopped.
-                is_drained = decode_req.kv_receiver.is_abort_release_safe()
+                is_drained = (
+                    decode_req.kv_receiver.kv_mgr.is_abort_release_safe(
+                        decode_req.req.bootstrap_room, required_acks
+                    )
+                    if self.scheduler.enable_hisparse
+                    else decode_req.kv_receiver.is_abort_release_safe()
+                )
                 is_ready = is_drained
             else:
                 is_drained = decode_req.kv_receiver.kv_mgr.is_abort_release_safe(