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

Co-authored-by: Tingwei Huang <huangtingwei9988@gmail.com>

docs/docs/advanced_features/hisparse_guide.mdx

@@ -118,6 +118,12 @@ Pass as a JSON string via `--hisparse-config`:
 
 Example: `--hisparse-config='{"top_k": 2048, "device_buffer_size": 6144, "host_to_device_ratio": 10, "swap_in_block_size": 960}'`
 
+### Shared-index prefetch (automatic)
+
+When a model reuses one anchor layer's top-k selection across a run of subsequent "skip" layers (DSA `index_topk_freq` / `index_topk_pattern`; native in GLM-5.2 as IndexShare), the working set of every skip layer is known the moment the anchor's index is computed. HiSparse exploits this automatically: the anchor's swap-in kernel records its miss plan (which host slots go to which device-buffer slots), and each skip layer replays that plan with a copy-only kernel issued ahead on a side stream, so the skip layers' host→device IO overlaps the intervening layers' compute instead of sitting on the decode critical path. The replay kernel uses a small fixed grid to keep its SM footprint low while overlapped.
+
+The prefetch is enabled automatically for eligible models (no pipeline parallelism, no speculative decoding) and can be turned off for A/B comparison with `SGLANG_DISABLE_HISPARSE_PREFETCH=1`.
+
 ## Deployment
 
 HiSparse currently requires **PD disaggregation mode** and is enabled only on the **decode instance**.

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

@@ -174,6 +174,55 @@ __device__ __forceinline__ int popc_mask(BallotMask mask) {
 #endif
 }
 
+// Copy one missed item host->device with one warp. Shared by the fused swap-in
+// kernel and copy_cache_planned_kernel so the layout dispatch cannot drift.
+template <bool IsMLA, bool IsDsv4Layout>
+__device__ __forceinline__ void copy_miss_item(
+    int32_t lane_id,
+    const void* __restrict__ host_cache_k,
+    const void* __restrict__ host_cache_v,
+    void* __restrict__ device_buffer_k,
+    void* __restrict__ device_buffer_v,
+    int64_t src_loc,
+    int64_t dst_loc,
+    int64_t item_size_bytes) {
+  static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA).");
+  if constexpr (IsDsv4Layout) {
+#ifdef USE_ROCM
+    // ROCm path: host cache and device buffer both use the page-padded C4
+    // layout (same as the write path and the CUDA branch). We can't reuse
+    // device::hisparse::transfer_item here because its warp logic is hardcoded
+    // to a 32-lane warp; on wavefront64 we use transfer_dsv4_item_warp, which
+    // moves the value and the scale in one warp-width-agnostic copy.
+    using namespace device::hisparse;
+    const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(device_buffer_k, static_cast<int32_t>(dst_loc));
+    const auto [src_value_ptr, src_scale_ptr] =
+        get_pointer_paged(const_cast<void*>(host_cache_k), static_cast<int32_t>(src_loc));
+    transfer_dsv4_item_warp(lane_id, src_value_ptr, src_scale_ptr, dst_value_ptr, dst_scale_ptr);
+#else
+    // CUDA path: page-padded device layout + page-padded host layout, K-only.
+    // The host cache is pinned DRAM but uses the same row layout as the GPU C4
+    // cache, so use the page-padded address calculation for both ends.
+    device::hisparse::transfer_item(
+        /*dst_cache=*/device_buffer_k,
+        /*src_cache=*/const_cast<void*>(host_cache_k),
+        /*dst_index=*/static_cast<int32_t>(dst_loc),
+        /*src_index=*/static_cast<int32_t>(src_loc));
+#endif
+  } else {
+    // Generic path: device + host both linear, stride = item_size_bytes.
+    const auto src_k = static_cast<const char*>(host_cache_k) + src_loc * item_size_bytes;
+    auto dst_k = static_cast<char*>(device_buffer_k) + dst_loc * item_size_bytes;
+    transfer_item_warp(lane_id, src_k, dst_k, item_size_bytes);
+
+    if constexpr (!IsMLA) {
+      const auto src_v = static_cast<const char*>(host_cache_v) + src_loc * item_size_bytes;
+      auto dst_v = static_cast<char*>(device_buffer_v) + dst_loc * item_size_bytes;
+      transfer_item_warp(lane_id, src_v, dst_v, item_size_bytes);
+    }
+  }
+}
+
 template <int BLOCK_SIZE>
 __global__ __launch_bounds__(BLOCK_SIZE, 1) void transfer_cache_dsv4_mla_kernel(
     void** src_caches,
@@ -268,12 +317,20 @@ struct SmemLayout {
 // IsDsv4Layout selects the miss-copy addressing:
 //   false -> generic byte-stride: device + host both linear, stride = item_size_bytes
 //   true  -> DSv4 page-padded device + page-padded host (kvcacheio.cuh constants)
+//
+// RecordMissPlan records this step's miss plan (miss_src/dst = host/device loc
+// per miss, miss_count per request) for shared-index skip layers to replay via
+// copy_cache_planned_kernel. SkipIO elides only the KV byte movement (timing
+// probe; output is garbage). Both are compile-time flags so the production
+// (false, false) instantiation stays byte-identical.
 template <
     int BLOCK_SIZE,
     int NUM_TOP_K,
     int HOT_BUFFER_SIZE,
     bool IsMLA,
     bool IsDsv4Layout,
+    bool RecordMissPlan,
+    bool SkipIO,
     typename SeqLensT,
     typename ReqPoolIndicesT>
 __global__ void load_cache_to_device_buffer_kernel(
@@ -296,7 +353,11 @@ __global__ void load_cache_to_device_buffer_kernel(
     int64_t top_k_tokens_stride,
     int64_t top_k_device_locs_stride,
     int64_t page_size,
-    int64_t item_size_bytes) {
+    int64_t item_size_bytes,
+    int64_t* __restrict__ miss_src_out,
+    int32_t* __restrict__ miss_dst_out,
+    int32_t* __restrict__ miss_count_out,
+    int64_t plan_stride) {
   static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA).");
   // todo hisparse: support page wise sparsity
   constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
@@ -345,6 +406,12 @@ __global__ void load_cache_to_device_buffer_kernel(
       }
       req_top_k_device_locs[i] = device_loc;
     }
+    // Short sequences load nothing from host: an empty miss plan for this request.
+    if constexpr (RecordMissPlan) {
+      if (tid == 0) {
+        miss_count_out[bid] = 0;
+      }
+    }
     return;
   }
 
@@ -554,11 +621,22 @@ __global__ void load_cache_to_device_buffer_kernel(
       s_top_k_tokens[miss_offset] = my_token;
       req_top_k_device_locs[my_token_idx] = req_device_buffer_locs[evict_slot];
       req_device_buffer_tokens[evict_slot] = my_token;
+      // Record the plan where the eviction is decided so it cannot disagree
+      // with the copy phase; locs are layer-independent (lockstep buffers).
+      if constexpr (RecordMissPlan) {
+        miss_src_out[bid * plan_stride + miss_offset] = req_host_cache_locs[my_token];
+        miss_dst_out[bid * plan_stride + miss_offset] = req_device_buffer_locs[evict_slot];
+      }
     }
   }
   __syncthreads();
 
   total_misses = NUM_TOP_K - s_total_hits - s_newest_hit;
+  if constexpr (RecordMissPlan) {
+    if (tid == 0) {
+      miss_count_out[bid] = total_misses;
+    }
+  }
   // Write back LRU order: evictables at front (LRU), hits at back (MRU).
   {
     const int total_evictable = HOT_BUFFER_SIZE - s_total_hits;
@@ -596,6 +674,7 @@ __global__ void load_cache_to_device_buffer_kernel(
   }
 
   // each warp copies one miss directly, can be separated into a new kernel if parallelism is a concern
+  if constexpr (!SkipIO) {
     for (int miss_idx = warp_id; miss_idx < total_misses; miss_idx += NUM_WARPS) {
       const int32_t miss_token = s_top_k_tokens[miss_idx];
       const int16_t evict_slot = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - miss_idx];
@@ -603,44 +682,20 @@ __global__ void load_cache_to_device_buffer_kernel(
       const int64_t src_loc = req_host_cache_locs[miss_token];
       const int64_t dst_loc = static_cast<int64_t>(req_device_buffer_locs[evict_slot]);
 
-    if constexpr (IsDsv4Layout) {
+      copy_miss_item<IsMLA, IsDsv4Layout>(
-#ifdef USE_ROCM
+          lane_id, host_cache_k, host_cache_v, device_buffer_k, device_buffer_v, src_loc, dst_loc, item_size_bytes);
-      // ROCm path: host cache and device buffer both use the page-padded C4
-      // layout (same as the write path and the CUDA branch). We can't reuse
-      // device::hisparse::transfer_item here because its warp logic is hardcoded
-      // to a 32-lane warp; on wavefront64 we use transfer_dsv4_item_warp, which
-      // moves the value and the scale in one warp-width-agnostic copy.
-      using namespace device::hisparse;
-      const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(device_buffer_k, static_cast<int32_t>(dst_loc));
-      const auto [src_value_ptr, src_scale_ptr] =
-          get_pointer_paged(const_cast<void*>(host_cache_k), static_cast<int32_t>(src_loc));
-      transfer_dsv4_item_warp(lane_id, src_value_ptr, src_scale_ptr, dst_value_ptr, dst_scale_ptr);
-#else
-      // CUDA path: page-padded device layout + page-padded host layout, K-only.
-      // The host cache is pinned DRAM but uses the same row layout as the GPU C4
-      // cache, so use the page-padded address calculation for both ends.
-      device::hisparse::transfer_item(
-          /*dst_cache=*/device_buffer_k,
-          /*src_cache=*/const_cast<void*>(host_cache_k),
-          /*dst_index=*/static_cast<int32_t>(dst_loc),
-          /*src_index=*/static_cast<int32_t>(src_loc));
-#endif
-    } else {
-      // Generic path: device + host both linear, stride = item_size_bytes.
-      const auto src_k = static_cast<const char*>(host_cache_k) + src_loc * item_size_bytes;
-      auto dst_k = static_cast<char*>(device_buffer_k) + dst_loc * item_size_bytes;
-      transfer_item_warp(lane_id, src_k, dst_k, item_size_bytes);
-
-      if constexpr (!IsMLA) {
-        const auto src_v = static_cast<const char*>(host_cache_v) + src_loc * item_size_bytes;
-        auto dst_v = static_cast<char*>(device_buffer_v) + dst_loc * item_size_bytes;
-        transfer_item_warp(lane_id, src_v, dst_v, item_size_bytes);
-      }
     }
   }
 }
 
-template <int BLOCK_SIZE, int NUM_TOP_K, int HOT_BUFFER_SIZE, bool IsMLA, bool IsDsv4Layout>
+template <
+    int BLOCK_SIZE,
+    int NUM_TOP_K,
+    int HOT_BUFFER_SIZE,
+    bool IsMLA,
+    bool IsDsv4Layout,
+    bool RecordMissPlan,
+    bool SkipIO>
 void load_cache_to_device_buffer(
     tvm::ffi::TensorView top_k_tokens,
     tvm::ffi::TensorView device_buffer_tokens,
@@ -656,11 +711,22 @@ void load_cache_to_device_buffer(
     tvm::ffi::TensorView lru_slots,
     tvm::ffi::TensorView num_real_reqs,
     int64_t page_size,
-    int64_t item_size_bytes) {
+    int64_t item_size_bytes,
+    tvm::ffi::TensorView miss_src_out,
+    tvm::ffi::TensorView miss_dst_out,
+    tvm::ffi::TensorView miss_count_out) {
   using namespace host;
 
   const int64_t bs = top_k_tokens.shape()[0];
   const int64_t host_stride = host_cache_locs.shape()[1];
+  // Miss-plan side outputs; 0-dim sentinels when RecordMissPlan is false.
+  int64_t* const miss_src_ptr = RecordMissPlan ? static_cast<int64_t*>(miss_src_out.data_ptr()) : nullptr;
+  int32_t* const miss_dst_ptr = RecordMissPlan ? static_cast<int32_t*>(miss_dst_out.data_ptr()) : nullptr;
+  int32_t* const miss_count_ptr = RecordMissPlan ? static_cast<int32_t*>(miss_count_out.data_ptr()) : nullptr;
+  const int64_t plan_stride = RecordMissPlan ? miss_src_out.strides()[0] : 0;
+  if (RecordMissPlan && miss_dst_out.strides()[0] != plan_stride) {
+    throw std::runtime_error("load_cache_to_device_buffer: miss_src/miss_dst row strides differ");
+  }
   const int64_t buffer_stride_0 = device_buffer_tokens.strides()[0];
   const int64_t lru_slot_stride_0 = lru_slots.strides()[0];
   const int64_t top_k_tokens_stride = top_k_tokens.strides()[0];
@@ -697,7 +763,11 @@ void load_cache_to_device_buffer(
         top_k_tokens_stride,
         top_k_device_locs_stride,
         page_size,
-        item_size_bytes);
+        item_size_bytes,
+        miss_src_ptr,
+        miss_dst_ptr,
+        miss_count_ptr,
+        plan_stride);
   };
 
   const auto seq_dtype = seq_lens.dtype();
@@ -713,6 +783,8 @@ void load_cache_to_device_buffer(
             HOT_BUFFER_SIZE,
             IsMLA,
             IsDsv4Layout,
+            RecordMissPlan,
+            SkipIO,
             int64_t,
             int64_t>,
         static_cast<const int64_t*>(seq_lens.data_ptr()),
@@ -725,6 +797,8 @@ void load_cache_to_device_buffer(
             HOT_BUFFER_SIZE,
             IsMLA,
             IsDsv4Layout,
+            RecordMissPlan,
+            SkipIO,
             int64_t,
             int32_t>,
         static_cast<const int64_t*>(seq_lens.data_ptr()),
@@ -737,6 +811,8 @@ void load_cache_to_device_buffer(
             HOT_BUFFER_SIZE,
             IsMLA,
             IsDsv4Layout,
+            RecordMissPlan,
+            SkipIO,
             int32_t,
             int64_t>,
         static_cast<const int32_t*>(seq_lens.data_ptr()),
@@ -749,6 +825,8 @@ void load_cache_to_device_buffer(
             HOT_BUFFER_SIZE,
             IsMLA,
             IsDsv4Layout,
+            RecordMissPlan,
+            SkipIO,
             int32_t,
             int32_t>,
         static_cast<const int32_t*>(seq_lens.data_ptr()),
@@ -756,4 +834,93 @@ void load_cache_to_device_buffer(
   }
 }
 
+// Copy-only swap-in for shared-index skip layers: replays the anchor's recorded
+// miss plan (no hit detection / LRU; the anchor's slot table stays valid). The
+// small fixed grid (num_blocks) keeps the SM footprint low while overlapping
+// compute on a side stream. SkipIO is the same probe as in the fused kernel.
+template <int BLOCK_SIZE, bool IsMLA, bool IsDsv4Layout, bool SkipIO>
+__global__ __launch_bounds__(BLOCK_SIZE, 1) void copy_cache_planned_kernel(
+    const int64_t* __restrict__ miss_src_locs,
+    const int32_t* __restrict__ miss_dst_locs,
+    const int32_t* __restrict__ miss_counts,
+    const int32_t* __restrict__ num_real_reqs,
+    const void* __restrict__ host_cache_k,
+    const void* __restrict__ host_cache_v,
+    void* __restrict__ device_buffer_k,
+    void* __restrict__ device_buffer_v,
+    int64_t plan_stride,
+    int64_t item_size_bytes) {
+  constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
+  const int lane_id = threadIdx.x % WARP_SIZE;
+  const int warp_global = blockIdx.x * NUM_WARPS + threadIdx.x / WARP_SIZE;
+  const int total_warps = gridDim.x * NUM_WARPS;
+  const int real = num_real_reqs[0];
+
+  // Warp-sized windows amortize the miss_counts loads; warps then round-robin
+  // the flattened (request, miss) space so a large sparse batch spreads over
+  // all warps (183us -> 29us at bs=100 with 2 misses/req on H200) while one
+  // request's miss burst still uses every warp.
+  int start = 0;  // flat index of the current request's first miss
+  for (int base = 0; base < real; base += WARP_SIZE) {
+    const int r_lane = base + lane_id;
+    const int cnt_lane = (r_lane < real) ? miss_counts[r_lane] : 0;
+    const int window = (real - base < WARP_SIZE) ? (real - base) : WARP_SIZE;
+    for (int j = 0; j < window; ++j) {
+      const int cnt = __shfl_sync(FULL_WARP_MASK, cnt_lane, j);
+      if (cnt == 0) continue;
+      int m0 = (warp_global - start) % total_warps;
+      if (m0 < 0) m0 += total_warps;
+      const int64_t r = base + j;
+      const int64_t* src_row = miss_src_locs + r * plan_stride;
+      const int32_t* dst_row = miss_dst_locs + r * plan_stride;
+      for (int m = m0; m < cnt; m += total_warps) {
+        // Timing probe: the plan is still walked; only the bytes stay put.
+        if constexpr (SkipIO) continue;
+        copy_miss_item<IsMLA, IsDsv4Layout>(
+            lane_id,
+            host_cache_k,
+            host_cache_v,
+            device_buffer_k,
+            device_buffer_v,
+            src_row[m],
+            static_cast<int64_t>(dst_row[m]),
+            item_size_bytes);
+      }
+      start += cnt;
+    }
+  }
+}
+
+template <int BLOCK_SIZE, bool IsMLA, bool IsDsv4Layout, bool SkipIO>
+void copy_cache_planned(
+    tvm::ffi::TensorView miss_src_locs,
+    tvm::ffi::TensorView miss_dst_locs,
+    tvm::ffi::TensorView miss_counts,
+    tvm::ffi::TensorView num_real_reqs,
+    tvm::ffi::TensorView host_cache_k,
+    tvm::ffi::TensorView host_cache_v,
+    tvm::ffi::TensorView device_buffer_k,
+    tvm::ffi::TensorView device_buffer_v,
+    int64_t num_blocks,
+    int64_t item_size_bytes) {
+  using namespace host;
+  const int64_t plan_stride = miss_src_locs.strides()[0];
+  if (miss_dst_locs.strides()[0] != plan_stride) {
+    throw std::runtime_error("copy_cache_planned: miss_src/miss_dst row strides differ");
+  }
+  const auto device = LaunchKernel::resolve_device(miss_src_locs.device());
+  LaunchKernel(num_blocks, BLOCK_SIZE, device)(
+      copy_cache_planned_kernel<BLOCK_SIZE, IsMLA, IsDsv4Layout, SkipIO>,
+      static_cast<const int64_t*>(miss_src_locs.data_ptr()),
+      static_cast<const int32_t*>(miss_dst_locs.data_ptr()),
+      static_cast<const int32_t*>(miss_counts.data_ptr()),
+      static_cast<const int32_t*>(num_real_reqs.data_ptr()),
+      host_cache_k.data_ptr(),
+      (IsMLA || host_cache_v.ndim() == 0) ? (const void*)nullptr : host_cache_v.data_ptr(),
+      device_buffer_k.data_ptr(),
+      (IsMLA || device_buffer_v.ndim() == 0) ? (void*)nullptr : device_buffer_v.data_ptr(),
+      plan_stride,
+      item_size_bytes);
+}
+
 }  // namespace sglang

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

@@ -19,12 +19,29 @@ def _jit_sparse_module(
     hot_buffer_size: int,
     is_mla: bool = False,
     is_dsv4_layout: bool = False,
+    record_miss_plan: bool = False,
+    skip_io: bool = False,
 ) -> Module:
+    # record_miss_plan / skip_io are compile-time kernel flags; the
+    # (False, False) production instantiation stays byte-identical.
     template_args = make_cpp_args(
-        block_size, num_top_k, hot_buffer_size, is_mla, is_dsv4_layout
+        block_size,
+        num_top_k,
+        hot_buffer_size,
+        is_mla,
+        is_dsv4_layout,
+        record_miss_plan,
+        skip_io,
     )
     cache_args = make_cpp_args(
-        item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla, is_dsv4_layout
+        item_size_bytes,
+        block_size,
+        num_top_k,
+        hot_buffer_size,
+        is_mla,
+        is_dsv4_layout,
+        record_miss_plan,
+        skip_io,
     )
     return load_jit(
         "sparse_cache",
@@ -39,6 +56,30 @@ def _jit_sparse_module(
     )
 
 
+@functools.cache
+def _jit_copy_planned_module(
+    block_size: int,
+    is_mla: bool,
+    is_dsv4_layout: bool,
+    skip_io: bool,
+) -> Module:
+    template_args = make_cpp_args(block_size, is_mla, is_dsv4_layout, skip_io)
+    return load_jit(
+        "sparse_copy_planned",
+        block_size,
+        is_mla,
+        is_dsv4_layout,
+        skip_io,
+        cuda_files=["hisparse.cuh"],
+        cuda_wrappers=[
+            (
+                "copy_cache_planned",
+                f"copy_cache_planned<{template_args}>",
+            )
+        ],
+    )
+
+
 @functools.cache
 def _jit_dsv4_transfer_module(block_size: int) -> Module:
     template_args = make_cpp_args(block_size)
@@ -91,11 +132,16 @@ def _load_cache_to_device_buffer_mla(
     page_size: int,
     block_size: int,
     num_real_reqs: torch.Tensor | None,
+    miss_src: torch.Tensor | None,
+    miss_dst: torch.Tensor | None,
+    miss_count: torch.Tensor | None,
+    skip_io: bool,
 ) -> None:
     assert (
         hot_buffer_size >= num_top_k
     ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})"
 
+    record_miss_plan = miss_src is not None
     module = _jit_sparse_module(
         item_size_bytes,
         block_size,
@@ -103,6 +149,8 @@ def _load_cache_to_device_buffer_mla(
         hot_buffer_size,
         is_mla=True,
         is_dsv4_layout=is_dsv4_layout,
+        record_miss_plan=record_miss_plan,
+        skip_io=skip_io,
     )
 
     empty = torch.empty(0)
@@ -112,6 +160,16 @@ def _load_cache_to_device_buffer_mla(
             [top_k_tokens.size(0)], dtype=torch.int32, device=top_k_tokens.device
         )
 
+    if record_miss_plan:
+        assert miss_dst is not None and miss_count is not None
+        assert miss_src.dtype == torch.int64 and miss_dst.dtype == torch.int32
+        assert miss_count.dtype == torch.int32
+        # The kernel indexes both plan rows with one stride.
+        assert miss_src.stride(0) == miss_dst.stride(0)
+    else:
+        # Unused sentinels; the RecordMissPlan=false instantiation never reads them.
+        miss_src = miss_dst = miss_count = empty
+
     module.load_cache_to_device_buffer(
         top_k_tokens,
         device_buffer_tokens,
@@ -128,6 +186,9 @@ def _load_cache_to_device_buffer_mla(
         num_real_reqs,
         page_size,
         item_size_bytes,
+        miss_src,
+        miss_dst,
+        miss_count,
     )
 
 
@@ -148,8 +209,16 @@ def load_cache_to_device_buffer_mla(
     page_size: int = 1,
     block_size: int = 256,
     num_real_reqs: torch.Tensor | None = None,
+    miss_src: torch.Tensor | None = None,
+    miss_dst: torch.Tensor | None = None,
+    miss_count: torch.Tensor | None = None,
+    skip_io: bool = False,
 ) -> None:
-    """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes)."""
+    """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes).
+
+    Optional miss_src/miss_dst/miss_count record the miss plan for replay by
+    copy_cache_planned_mla; skip_io elides only the KV bytes (timing probe).
+    """
     _load_cache_to_device_buffer_mla(
         is_dsv4_layout=False,
         top_k_tokens=top_k_tokens,
@@ -168,6 +237,47 @@ def load_cache_to_device_buffer_mla(
         page_size=page_size,
         block_size=block_size,
         num_real_reqs=num_real_reqs,
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+        skip_io=skip_io,
+    )
+
+
+def copy_cache_planned_mla(
+    *,
+    miss_src: torch.Tensor,
+    miss_dst: torch.Tensor,
+    miss_count: torch.Tensor,
+    num_real_reqs: torch.Tensor,
+    host_cache: torch.Tensor,
+    device_buffer: torch.Tensor,
+    item_size_bytes: int,
+    num_blocks: int = 4,
+    block_size: int = 1024,
+    is_dsv4_layout: bool = False,
+    skip_io: bool = False,
+) -> None:
+    """Replay a recorded miss plan (host_cache -> device_buffer) for a skip layer.
+
+    IO-only, no planning; the small fixed grid keeps the SM footprint low while
+    overlapped on a side stream. The anchor's slot table stays valid (lockstep).
+    """
+    assert miss_src.dtype == torch.int64 and miss_dst.dtype == torch.int32
+    assert miss_count.dtype == torch.int32
+    module = _jit_copy_planned_module(block_size, True, is_dsv4_layout, skip_io)
+    empty = torch.empty(0)
+    module.copy_cache_planned(
+        miss_src,
+        miss_dst,
+        miss_count,
+        num_real_reqs,
+        host_cache,
+        empty,
+        device_buffer,
+        empty,
+        num_blocks,
+        item_size_bytes,
     )
 
 
@@ -188,6 +298,10 @@ def load_cache_to_device_buffer_dsv4_mla(
     page_size: int = 1,
     block_size: int = 256,
     num_real_reqs: torch.Tensor | None = None,
+    miss_src: torch.Tensor | None = None,
+    miss_dst: torch.Tensor | None = None,
+    miss_count: torch.Tensor | None = None,
+    skip_io: bool = False,
 ) -> None:
     """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout."""
     _load_cache_to_device_buffer_mla(
@@ -208,4 +322,8 @@ def load_cache_to_device_buffer_dsv4_mla(
         page_size=page_size,
         block_size=block_size,
         num_real_reqs=num_real_reqs,
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+        skip_io=skip_io,
     )

python/sglang/srt/environ.py

@@ -846,6 +846,14 @@ class Envs:
     # Triton two_dot variant, 1.16-1.38x faster across GLM/DS shapes).
     SGLANG_OPT_Q8KV8_QPREP_VARIANT = EnvStr("auto")
 
+    # HiSparse
+    # Kill-switch for the shared-index (IndexShare) swap-in prefetch
+    # (auto-enabled for GLM-5.2-style DSA); set True to A/B synchronous swap-in.
+    SGLANG_DISABLE_HISPARSE_PREFETCH = EnvBool(False)
+    # Timing probe: run the swap-in fully but skip the host->device KV bytes,
+    # measuring the "IO is free" floor. GARBAGE OUTPUT -- benchmarking only.
+    SGLANG_DEBUG_HISPARSE_SKIP_IO = EnvBool(False)
+
     # TRT-LLM-gen fused MoE (SiTU) via sglang JIT: path to an unpacked SiTU
     # cubin pool (cubins + flat ABI headers + overlay/; distributed as a
     # single downloadable archive). Needs the public flashinfer package

python/sglang/srt/managers/hisparse_coordinator.py

@@ -1,14 +1,17 @@
 # to be combined with the sparse coordinator class and sparse algorithm family
 
 import logging
-from typing import List, NamedTuple, Union
+from typing import Dict, List, NamedTuple, Optional, Tuple, Union
 
 import torch
 
 from sglang.kernels.ops.kvcache.hisparse import (
+    copy_cache_planned_mla,
     load_cache_to_device_buffer_dsv4_mla,
     load_cache_to_device_buffer_mla,
 )
+from sglang.srt.configs.model_config import dsa_layer_skips_topk, is_deepseek_dsa
+from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.allocator.hisparse import (
     DeepSeekV4HiSparseTokenToKVPoolAllocator,
@@ -42,6 +45,69 @@ class HiSparseTokenStats(NamedTuple):
     host_token_usage: float
 
 
+def resolve_shared_index_layers(
+    *,
+    hf_text_config,
+    pp_size: int,
+    is_speculative: bool,
+) -> Optional[List[bool]]:
+    """Per-layer "reuses the previous layer's top-k index" pattern, or None.
+
+    Mirrors DeepseekV2AttentionMLA's skip_topk derivation (index_topk_pattern /
+    index_topk_freq / cli_factor); None when the model has no sharing or the
+    prefetch cannot run (PP, speculative decoding, kill-switch).
+    """
+    if not is_deepseek_dsa(hf_text_config):
+        return None
+    num_layers = hf_text_config.num_hidden_layers
+    cli_factor = getattr(hf_text_config, "cli_factor", 1) or 1
+    if cli_factor > 1:
+        pattern = [i % cli_factor != 0 for i in range(num_layers)]
+    else:
+        pattern = [dsa_layer_skips_topk(hf_text_config, i) for i in range(num_layers)]
+    if not any(pattern):
+        return None
+    if pp_size != 1 or is_speculative:
+        logger.warning(
+            "HiSparse shared-index prefetch is unsupported under pipeline "
+            "parallelism / speculative decoding; falling back to synchronous "
+            "swap-in."
+        )
+        return None
+    if envs.SGLANG_DISABLE_HISPARSE_PREFETCH.get():
+        logger.info(
+            "HiSparse shared-index prefetch disabled via "
+            "SGLANG_DISABLE_HISPARSE_PREFETCH; using synchronous swap-in."
+        )
+        return None
+    return pattern
+
+
+def _build_prefetch_groups(
+    is_shared_index_layer: List[bool],
+) -> Tuple[Dict[int, List[int]], List[int]]:
+    """Group consecutive shared-index (skip) layers under their anchor layer.
+
+    Returns (groups, slot): anchor layer_id -> ordered skip layers, and each
+    skip layer's position in its group (indexes the per-slot prefetch events).
+    """
+    groups: Dict[int, List[int]] = {}
+    slot = [0] * len(is_shared_index_layer)
+    anchor = None
+    for i, is_shared in enumerate(is_shared_index_layer):
+        if not is_shared:
+            anchor = i  # compute layer; anchors the skip layers after it
+            continue
+        assert anchor is not None, (
+            f"shared-index (skip) layer {i} has no preceding compute layer; "
+            "the model's index-topk pattern is invalid"
+        )
+        group = groups.setdefault(anchor, [])
+        slot[i] = len(group)
+        group.append(i)
+    return groups, slot
+
+
 class HiSparseCoordinator:
     def __init__(
         self,
@@ -56,6 +122,7 @@ class HiSparseCoordinator:
         tp_group,
         host_to_device_ratio: int = 2,
         swap_in_block_size: int = 960,
+        shared_index_layers: Optional[List[bool]] = None,
     ):
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
@@ -63,6 +130,9 @@ class HiSparseCoordinator:
         self.device_buffer_size = device_buffer_size
         self.device = device
         self.swap_in_block_size = swap_in_block_size
+        # Timing probe: skip the host->device KV bytes to measure the "IO is
+        # free" floor. Produces garbage output; benchmarking only.
+        self.skip_io = envs.SGLANG_DEBUG_HISPARSE_SKIP_IO.get()
         self.compress_ratio = self.token_to_kv_pool_allocator.compress_ratio
 
         self.is_dsv4_hisparse = isinstance(
@@ -186,6 +256,65 @@ class HiSparseCoordinator:
         # staging already backed up all prefill tokens.  Cleared after one step.
         self._skip_first_backup = [False] * max_num_req_slots
 
+        self._init_shared_index_prefetch(
+            shared_index_layers=shared_index_layers,
+            layer_num=layer_num,
+            max_num_req_slots=max_num_req_slots,
+        )
+
+    def _init_shared_index_prefetch(
+        self,
+        shared_index_layers: Optional[List[bool]],
+        layer_num: int,
+        max_num_req_slots: int,
+    ) -> None:
+        """Set up the plan-then-IO prefetch for shared-index (IndexShare) models:
+        the anchor's kernel records its miss plan and skip layers replay it on
+        `prefetch_stream`, overlapping their IO with the intervening compute."""
+        if shared_index_layers is not None and len(shared_index_layers) != layer_num:
+            # Attention-layer count differs from num_hidden_layers (e.g. Longcat
+            # doubles it): pattern would be misindexed, fall back to synchronous.
+            logger.warning(
+                "HiSparse shared-index prefetch disabled: pattern length %d != "
+                "KV pool layer_num %d; using synchronous swap-in.",
+                len(shared_index_layers),
+                layer_num,
+            )
+            shared_index_layers = None
+        self._is_shared_index_layer = list(shared_index_layers or [False] * layer_num)
+        self.enable_prefetch = any(self._is_shared_index_layer)
+        self._prefetch_groups, self._prefetch_slot = _build_prefetch_groups(
+            self._is_shared_index_layer
+        )
+        if not self.enable_prefetch:
+            return
+
+        # Small fixed grid for the copy-only kernel: low SM footprint so the
+        # copies overlap compute with little contention.
+        self._prefetch_copy_blocks = 4
+        max_group_size = max(len(g) for g in self._prefetch_groups.values())
+        self.prefetch_stream = device_module.Stream()
+        self._prefetch_events = [device_module.Event() for _ in range(max_group_size)]
+        # Plan recorded by the current anchor, replayed by its skip layers. One
+        # buffer set suffices: the last skip layer's event wait orders the next
+        # anchor's writes after this group's copies.
+        self._miss_src = torch.zeros(
+            (max_num_req_slots, self.top_k), dtype=torch.int64, device=self.device
+        )
+        self._miss_dst = torch.zeros(
+            (max_num_req_slots, self.top_k), dtype=torch.int32, device=self.device
+        )
+        self._miss_count = torch.zeros(
+            (max_num_req_slots,), dtype=torch.int32, device=self.device
+        )
+        logger.info(
+            "HiSparse: shared-index prefetch (plan-then-IO) enabled; %d anchor "
+            "group(s), %d skip layer(s) of %d total.",
+            len(self._prefetch_groups),
+            sum(self._is_shared_index_layer),
+            layer_num,
+        )
+
     def set_decode_producer_stream(self, stream) -> None:
         self.decode_producer_stream = stream
 
@@ -194,6 +323,9 @@ class HiSparseCoordinator:
         # See HostKVCache.destroy for why the explicit unregister matters.
         self.write_staging_stream.synchronize()
         self.decode_backup_stream.synchronize()
+        if self.enable_prefetch:
+            # Skip-layer copies read the pinned host pool on the prefetch stream.
+            self.prefetch_stream.synchronize()
         self.mem_pool_host.destroy()
 
     def get_token_stats(self) -> HiSparseTokenStats:
@@ -802,16 +934,20 @@ class HiSparseCoordinator:
         self.lru_slots[:, req.req_pool_idx, :].copy_(self._lru_init)
         self._skip_first_backup[req.req_pool_idx] = False
 
-    def swap_in_selected_pages(
+    def _run_swap_in_kernel(
         self,
         req_pool_indices: torch.Tensor,
         compressed_seq_lens: torch.Tensor,
         top_k_result: torch.Tensor,
         layer_id: int,
+        record_plan: bool = False,
     ) -> torch.Tensor:
-        """Swap selected top-k tokens into device memory and return their indices."""
+        """Run the full plan+IO swap-in kernel for one layer; return its slot table.
-        num_reqs = req_pool_indices.size(0)
 
+        record_plan (set on the anchor of a shared-index group) also records the
+        miss plan into self._miss_{src,dst,count} for the skip layers to replay.
+        """
+        num_reqs = req_pool_indices.size(0)
         top_k_indices = self.top_k_device_locs_buffer[:num_reqs]
 
         swap_in_fn = (
@@ -819,6 +955,15 @@ class HiSparseCoordinator:
             if self.is_dsv4_hisparse
             else load_cache_to_device_buffer_mla
         )
+        plan = (
+            dict(
+                miss_src=self._miss_src[:num_reqs],
+                miss_dst=self._miss_dst[:num_reqs],
+                miss_count=self._miss_count[:num_reqs],
+            )
+            if record_plan
+            else {}
+        )
         swap_in_fn(
             top_k_tokens=top_k_result,
             device_buffer_tokens=self.req_device_buffer_tokens[layer_id],
@@ -836,5 +981,70 @@ class HiSparseCoordinator:
             page_size=1,
             block_size=self.swap_in_block_size,
             num_real_reqs=self.num_real_reqs,
+            skip_io=self.skip_io,
+            **plan,
         )
         return top_k_indices
+
+    def _run_copy_only_kernel(self, num_reqs: int, skip_layer: int) -> None:
+        """Replay the anchor's recorded miss plan into a skip layer's buffers
+        (IO-only; the anchor's slot table stays valid -- lockstep layout)."""
+        copy_cache_planned_mla(
+            miss_src=self._miss_src[:num_reqs],
+            miss_dst=self._miss_dst[:num_reqs],
+            miss_count=self._miss_count[:num_reqs],
+            num_real_reqs=self.num_real_reqs,
+            host_cache=self.mem_pool_host.kv_buffer[skip_layer],
+            device_buffer=self.mem_pool_device.kv_buffer[skip_layer],
+            item_size_bytes=self.item_size_bytes,
+            num_blocks=self._prefetch_copy_blocks,
+            is_dsv4_layout=self.is_dsv4_hisparse,
+            skip_io=self.skip_io,
+        )
+
+    def swap_in_selected_pages(
+        self,
+        req_pool_indices: torch.Tensor,
+        compressed_seq_lens: torch.Tensor,
+        top_k_result: torch.Tensor,
+        layer_id: int,
+    ) -> torch.Tensor:
+        """Swap selected top-k tokens into device memory and return their indices.
+
+        With prefetch enabled, anchors swap in synchronously (recording the miss
+        plan) and prefetch their skip layers' copies; skip layers just wait.
+        """
+        if not self.enable_prefetch:
+            return self._run_swap_in_kernel(
+                req_pool_indices, compressed_seq_lens, top_k_result, layer_id
+            )
+
+        num_reqs = req_pool_indices.size(0)
+        if self._is_shared_index_layer[layer_id]:
+            # Skip layer: wait for its prefetched copy; the anchor's slot table
+            # applies (shared index + lockstep buffers).
+            slot = self._prefetch_slot[layer_id]
+            self._prefetch_events[slot].wait(device_module.current_stream())
+            return self.top_k_device_locs_buffer[:num_reqs]
+
+        # Anchor: swap in synchronously (recording the plan), then prefetch the
+        # skip layers' copies on the side stream.
+        group = self._prefetch_groups.get(layer_id)
+        anchor_locs = self._run_swap_in_kernel(
+            req_pool_indices,
+            compressed_seq_lens,
+            top_k_result,
+            layer_id,
+            record_plan=group is not None,
+        )
+        if group:
+            # Fork: the prefetch stream must observe the anchor's plan (produced
+            # on the current stream) before replaying it.
+            self.prefetch_stream.wait_stream(device_module.current_stream())
+            with device_module.stream(self.prefetch_stream):
+                for skip_layer in group:
+                    self._run_copy_only_kernel(num_reqs, skip_layer)
+                    self._prefetch_events[self._prefetch_slot[skip_layer]].record(
+                        self.prefetch_stream
+                    )
+        return anchor_locs

python/sglang/srt/model_executor/model_runner.py

@@ -847,7 +847,10 @@ class ModelRunner:
     def maybe_init_hisparse_coordinator(self):
         if not self.enable_hisparse:
             return
-        from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
+        from sglang.srt.managers.hisparse_coordinator import (
+            HiSparseCoordinator,
+            resolve_shared_index_layers,
+        )
         from sglang.srt.mem_cache.sparsity import parse_hisparse_config
 
         hisparse_cfg = parse_hisparse_config(self.server_args)
@@ -867,6 +870,11 @@ class ModelRunner:
             ),
             host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
             swap_in_block_size=hisparse_cfg.swap_in_block_size,
+            shared_index_layers=resolve_shared_index_layers(
+                hf_text_config=self.model_config.hf_text_config,
+                pp_size=self.ps.pp_size,
+                is_speculative=self.spec_algorithm.is_speculative(),
+            ),
         )
 
     def post_capture_resize_kv_pool(self):

test/manual/kernels/test_hisparse_prefetch.py

@@ -0,0 +1,394 @@
+"""Extended tests for the HiSparse shared-index (plan-then-IO) prefetch.
+
+Local-only (not registered to CI): these cover the CUDA-graph capture/replay
+pattern, the DSv4 page-padded layout, and the SGLANG_DEBUG_HISPARSE_SKIP_IO
+probe, each of which JIT-compiles extra kernel instantiations. The cheap
+plan-replay correctness guards run in CI via
+test/registered/kernels/ops/kvcache/test_hisparse.py, which this file imports
+its fixtures from.
+
+Run: python3 test/manual/kernels/test_hisparse_prefetch.py
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+from sglang.kernels.ops.kvcache.hisparse import (
+    copy_cache_planned_mla,
+    load_cache_to_device_buffer_dsv4_mla,
+    load_cache_to_device_buffer_mla,
+)
+from sglang.srt.utils import is_hip
+
+sys.path.insert(
+    0,
+    str(
+        Path(__file__).resolve().parents[2]
+        / "registered"
+        / "kernels"
+        / "ops"
+        / "kvcache"
+    ),
+)
+from test_hisparse import (  # noqa: F401  (GPU/platform guard applies here too); noqa: E402
+    DEVICE,
+    DEVICE_CACHE_SIZE,
+    DSV4_ITEM_BYTES,
+    DSV4_PAGE_BYTES,
+    DSV4_PAGE_SIZE,
+    DTYPE,
+    HOST_CACHE_SIZE,
+    HOT_BUFFER_SIZE,
+    ITEM_SIZE_BYTES,
+    KV_DIM,
+    _host_cache,
+    _long_case,
+    _make_plan,
+    _make_state,
+    _run_kernel,
+    _write_dsv4_token,
+    pytestmark,
+)
+
+
+def test_plan_then_io_dsv4_matches_sync_swap_in() -> None:
+    """DSv4 layout: replaying the recorded plan lands the page-padded value+scale
+    bytes exactly where the fused swap-in copy puts them."""
+    num_pages = 2
+    state = _long_case()
+    plan_state = _long_case()
+
+    def _dsv4_caches():
+        host = torch.zeros(
+            (num_pages, DSV4_PAGE_BYTES),
+            dtype=torch.uint8,
+            device="cpu",
+            pin_memory=True,
+        )
+        for token in range(HOST_CACHE_SIZE):
+            _write_dsv4_token(host, token, seed=token + 1)
+        dev = torch.full(
+            (num_pages, DSV4_PAGE_BYTES), 0xFF, dtype=torch.uint8, device=DEVICE
+        )
+        return host, dev
+
+    common = dict(
+        top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE),
+        req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
+        seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE),
+        item_size_bytes=DSV4_ITEM_BYTES,
+        num_top_k=1,
+        hot_buffer_size=HOT_BUFFER_SIZE,
+        page_size=DSV4_PAGE_SIZE,
+        block_size=256,
+        num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
+    )
+
+    # Reference: fused swap-in copies host token 6 into evict slot 0 (loc 9).
+    ref_host, ref_dev = _dsv4_caches()
+    out = torch.full((1, 1), -1, dtype=torch.int32, device=DEVICE)
+    load_cache_to_device_buffer_dsv4_mla(
+        device_buffer_tokens=state["device_buffer_tokens"],
+        host_cache_locs=state["host_cache_locs"],
+        device_buffer_locs=state["device_buffer_locs"],
+        host_cache=ref_host,
+        device_buffer=ref_dev,
+        top_k_device_locs=out,
+        lru_slots=state["lru_slots"],
+        **common,
+    )
+
+    # Anchor: same swap-in on a twin state, recording the plan.
+    miss_src, miss_dst, miss_count = _make_plan(1, 1)
+    anchor_host, anchor_dev = _dsv4_caches()
+    anchor_out = torch.full((1, 1), -1, dtype=torch.int32, device=DEVICE)
+    load_cache_to_device_buffer_dsv4_mla(
+        device_buffer_tokens=plan_state["device_buffer_tokens"],
+        host_cache_locs=plan_state["host_cache_locs"],
+        device_buffer_locs=plan_state["device_buffer_locs"],
+        host_cache=anchor_host,
+        device_buffer=anchor_dev,
+        top_k_device_locs=anchor_out,
+        lru_slots=plan_state["lru_slots"],
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+        **common,
+    )
+
+    # Skip layer: replay the plan into a fresh buffer; must match the reference.
+    replay_host, replay_dev = _dsv4_caches()
+    copy_cache_planned_mla(
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+        num_real_reqs=common["num_real_reqs"],
+        host_cache=replay_host,
+        device_buffer=replay_dev,
+        item_size_bytes=DSV4_ITEM_BYTES,
+        num_blocks=4,
+        is_dsv4_layout=True,
+    )
+    torch.cuda.synchronize()
+
+    assert torch.equal(miss_count.cpu(), torch.tensor([1], dtype=torch.int32))
+    assert torch.equal(anchor_out.cpu(), out.cpu())
+    assert torch.equal(replay_dev.cpu(), ref_dev.cpu())
+
+
+def test_skip_io_probe_plans_without_moving_bytes() -> None:
+    """skip_io still runs all planning (slot table, LRU, miss plan) but must
+    leave the device buffer untouched; replaying the plan then repairs it."""
+    locs = [[9, 7, 3, 5, 11]]
+    toks = [[1, 4, 2, 5, -1]]
+    top_k = torch.tensor([[6, 4]], dtype=torch.int32, device=DEVICE)
+    nr, K = top_k.shape
+
+    ref = _make_state(locs, toks, [7])
+    ref_out = _run_kernel(top_k_tokens=top_k, seq_len=8, **ref)
+
+    probe = _make_state(locs, toks, [7])
+    probe_buffer_before = probe["device_buffer"].clone()
+    miss_src, miss_dst, miss_count = _make_plan(nr, K)
+    probe_out = _run_kernel(
+        top_k_tokens=top_k,
+        seq_len=8,
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+        skip_io=True,
+        **probe,
+    )
+
+    # All planning outputs match the real run; only the bytes stayed put.
+    assert torch.equal(probe_out.cpu(), ref_out.cpu())
+    assert torch.equal(probe["lru_slots"].cpu(), ref["lru_slots"].cpu())
+    assert torch.equal(probe["device_buffer"].cpu(), probe_buffer_before.cpu())
+    assert not torch.equal(probe["device_buffer"].cpu(), ref["device_buffer"].cpu())
+
+    copy_cache_planned_mla(
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+        num_real_reqs=torch.tensor([nr], dtype=torch.int32, device=DEVICE),
+        host_cache=probe["host_cache"],
+        device_buffer=probe["device_buffer"],
+        item_size_bytes=ITEM_SIZE_BYTES,
+        num_blocks=4,
+    )
+    torch.cuda.synchronize()
+    assert torch.equal(probe["device_buffer"].cpu(), ref["device_buffer"].cpu())
+
+
+_PIO_LAYERS = 4  # one anchor (layer 0) + three skip layers (GLM group of freq 4)
+_PIO_REQS = 2
+_PIO_SEQ = 10  # > HOT_BUFFER_SIZE -> long path; newest token = 9
+_PIO_DBL = [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]  # per-req [4 LRU slots + newest]
+_PIO_STEPS = [
+    [[4, 5, 0, 9], [6, 7, 1, 9]],
+    [[4, 10, 2, 9], [8, 5, 3, 9]],
+    [[11, 4, 5, 9], [6, 12, 7, 9]],
+    [[0, 1, 2, 9], [3, 4, 5, 9]],
+]
+
+
+def _pio_fresh(host_cache, dbl):
+    """Fresh per-layer buffers/tokens/lru, all layers initialized identically."""
+    buffers, dbt, lru = [], [], []
+    init_tokens = [[0, 1, 2, 3, -1], [0, 1, 2, 3, -1]]
+    for _ in range(_PIO_LAYERS):
+        db = torch.full((DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE)
+        for rid in range(_PIO_REQS):
+            for slot, tok in enumerate(init_tokens[rid][:HOT_BUFFER_SIZE]):
+                db[dbl[rid, slot]].copy_(host_cache[tok].to(DEVICE))
+            db[dbl[rid, HOT_BUFFER_SIZE]].copy_(host_cache[_PIO_SEQ - 1].to(DEVICE))
+        buffers.append(db)
+        dbt.append(torch.tensor(init_tokens, dtype=torch.int32, device=DEVICE))
+        lru.append(
+            torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE)
+            .view(1, -1)
+            .repeat(_PIO_REQS, 1)
+            .contiguous()
+        )
+    torch.cuda.synchronize()
+    return buffers, dbt, lru
+
+
+def _pio_swap_in(topk, dbt, lru, dbl, hcl, host, buffer, out, seq_lens, nrr, rpi, plan):
+    miss_src, miss_dst, miss_count = plan if plan else (None, None, None)
+    load_cache_to_device_buffer_mla(
+        top_k_tokens=topk,
+        device_buffer_tokens=dbt,
+        host_cache_locs=hcl,
+        device_buffer_locs=dbl,
+        host_cache=host,
+        device_buffer=buffer,
+        top_k_device_locs=out,
+        req_pool_indices=rpi,
+        seq_lens=seq_lens,
+        lru_slots=lru,
+        item_size_bytes=ITEM_SIZE_BYTES,
+        num_top_k=topk.shape[1],
+        hot_buffer_size=HOT_BUFFER_SIZE,
+        page_size=1,
+        block_size=256,
+        num_real_reqs=nrr,
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+    )
+
+
+def _pio_prefetch_step(
+    topk, buffers, dbt, lru, dbl, hcl, host, seq_lens, nrr, rpi, out, plan, side, events
+):
+    """Anchor records the plan; skip layers replay it copy-only on a side stream."""
+    miss_src, miss_dst, miss_count = plan
+    _pio_swap_in(
+        topk, dbt[0], lru[0], dbl, hcl, host, buffers[0], out, seq_lens, nrr, rpi, plan
+    )
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        for layer in range(1, _PIO_LAYERS):
+            copy_cache_planned_mla(
+                miss_src=miss_src,
+                miss_dst=miss_dst,
+                miss_count=miss_count,
+                num_real_reqs=nrr,
+                host_cache=host,
+                device_buffer=buffers[layer],
+                item_size_bytes=ITEM_SIZE_BYTES,
+                num_blocks=4,
+            )
+            events[layer].record(side)
+    for layer in range(1, _PIO_LAYERS):
+        events[layer].wait(torch.cuda.current_stream())
+
+
+def _pio_sync_step(topk, buffers, dbt, lru, dbl, hcl, host, seq_lens, nrr, rpi):
+    """Reference: run the full swap-in independently on every layer."""
+    outs = []
+    for layer in range(_PIO_LAYERS):
+        out = torch.full_like(topk, -1)
+        _pio_swap_in(
+            topk,
+            dbt[layer],
+            lru[layer],
+            dbl,
+            hcl,
+            host,
+            buffers[layer],
+            out,
+            seq_lens,
+            nrr,
+            rpi,
+            None,
+        )
+        outs.append(out.clone())
+    return outs
+
+
+@pytest.mark.skipif(is_hip(), reason="CUDA graph capture test is CUDA-only.")
+def test_plan_then_io_cuda_graph_replay() -> None:
+    """The plan-then-IO prefetch pattern captures into a CUDA graph and replays
+    bit-identically to the eager synchronous swap-in across multiple steps."""
+    host = _host_cache()
+    dbl = torch.tensor(_PIO_DBL, dtype=torch.int32, device=DEVICE)
+    hcl = (
+        torch.arange(HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE)
+        .view(1, -1)
+        .repeat(_PIO_REQS, 1)
+        .contiguous()
+    )
+    seq_lens = torch.full((_PIO_REQS,), _PIO_SEQ, dtype=torch.int32, device=DEVICE)
+    nrr = torch.tensor([_PIO_REQS], dtype=torch.int32, device=DEVICE)
+    rpi = torch.arange(_PIO_REQS, dtype=torch.int64, device=DEVICE)
+    K = len(_PIO_STEPS[0][0])
+    steps = [torch.tensor(s, dtype=torch.int32, device=DEVICE) for s in _PIO_STEPS]
+
+    # Reference: full synchronous swap-in on every layer, snapshotted per step.
+    ref_buf, ref_dbt, ref_lru = _pio_fresh(host, dbl)
+    ref_slots, ref_snap = [], []
+    for topk in steps:
+        ref_slots.append(
+            _pio_sync_step(
+                topk, ref_buf, ref_dbt, ref_lru, dbl, hcl, host, seq_lens, nrr, rpi
+            )
+        )
+        torch.cuda.synchronize()
+        ref_snap.append([b.clone() for b in ref_buf])
+    torch.cuda.synchronize()
+
+    # Graph-captured prefetch replayed step by step against a fixed topk buffer.
+    buf, dbt, lru = _pio_fresh(host, dbl)
+    topk_buf = torch.zeros((_PIO_REQS, K), dtype=torch.int32, device=DEVICE)
+    out = torch.full((_PIO_REQS, K), -1, dtype=torch.int32, device=DEVICE)
+    plan = _make_plan(_PIO_REQS, K)
+    side = torch.cuda.Stream()
+    events = [torch.cuda.Event() for _ in range(_PIO_LAYERS)]
+
+    warm = torch.cuda.Stream()
+    warm.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(warm):
+        topk_buf.copy_(steps[0])
+        _pio_prefetch_step(
+            topk_buf,
+            buf,
+            dbt,
+            lru,
+            dbl,
+            hcl,
+            host,
+            seq_lens,
+            nrr,
+            rpi,
+            out,
+            plan,
+            side,
+            events,
+        )
+    torch.cuda.current_stream().wait_stream(warm)
+    torch.cuda.synchronize()
+
+    # Reset state mutated by warmup so capture starts from a clean identical state.
+    buf, dbt, lru = _pio_fresh(host, dbl)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        _pio_prefetch_step(
+            topk_buf,
+            buf,
+            dbt,
+            lru,
+            dbl,
+            hcl,
+            host,
+            seq_lens,
+            nrr,
+            rpi,
+            out,
+            plan,
+            side,
+            events,
+        )
+    torch.cuda.synchronize()
+
+    for s, topk in enumerate(steps):
+        topk_buf.copy_(topk)
+        graph.replay()
+        torch.cuda.synchronize()
+        # Anchor slot table matches the synchronous layer-0 result.
+        assert torch.equal(
+            out.cpu(), ref_slots[s][0].cpu()
+        ), f"slots differ at step {s}"
+        # Every layer's device buffer stays bit-identical to synchronous swap-in.
+        for layer in range(_PIO_LAYERS):
+            assert torch.equal(
+                buf[layer].cpu(), ref_snap[s][layer].cpu()
+            ), f"buffer differs at step {s}, layer {layer}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/benchmark/kvcache/bench_hisparse.py

@@ -6,13 +6,16 @@ import triton
 import triton.testing
 
 from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
-from sglang.kernels.ops.kvcache.hisparse import load_cache_to_device_buffer_mla
+from sglang.kernels.ops.kvcache.hisparse import (
+    copy_cache_planned_mla,
+    load_cache_to_device_buffer_mla,
+)
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(
-    est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+    est_time=18, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
-register_amd_ci(est_time=12, stage="jit-kernel-benchmark", runner_config="amd")
+register_amd_ci(est_time=18, stage="jit-kernel-benchmark", runner_config="amd")
 
 DEVICE = DEFAULT_DEVICE
 DTYPE = DEFAULT_DTYPE
@@ -159,6 +162,66 @@ def _time_kernel(batch_size: int, hot_buffer_size: int, miss_rate: float) -> flo
     return start.elapsed_time(end) * 1000.0 / ROUNDS
 
 
+def _time_planned_copy(
+    batch_size: int, hot_buffer_size: int, miss_rate: float
+) -> float:
+    """Time the copy-only replay used by shared-index skip layers: one anchor
+    swap-in records a real plan, each timed round replays it (num_blocks=4)."""
+    state = _build_inputs(batch_size, hot_buffer_size, miss_rate)
+    miss_src = torch.zeros((batch_size, TOP_K), dtype=torch.int64, device=DEVICE)
+    miss_dst = torch.zeros((batch_size, TOP_K), dtype=torch.int32, device=DEVICE)
+    miss_count = torch.zeros((batch_size,), dtype=torch.int32, device=DEVICE)
+    load_cache_to_device_buffer_mla(
+        top_k_tokens=state["top_k_tokens"],
+        device_buffer_tokens=state["device_buffer_tokens"],
+        host_cache_locs=state["host_cache_locs"],
+        device_buffer_locs=state["device_buffer_locs"],
+        host_cache=state["host_cache"],
+        device_buffer=state["device_buffer"],
+        top_k_device_locs=state["top_k_device_locs"],
+        req_pool_indices=state["req_pool_indices"],
+        seq_lens=state["seq_lens"],
+        lru_slots=state["lru_slots"],
+        item_size_bytes=ITEM_SIZE_BYTES,
+        num_top_k=TOP_K,
+        hot_buffer_size=hot_buffer_size,
+        block_size=1024,
+        num_real_reqs=state["num_real_reqs"],
+        miss_src=miss_src,
+        miss_dst=miss_dst,
+        miss_count=miss_count,
+    )
+    torch.cuda.synchronize()
+    skip_layer_buffer = torch.empty_like(state["device_buffer"])
+
+    def run_once():
+        copy_cache_planned_mla(
+            miss_src=miss_src,
+            miss_dst=miss_dst,
+            miss_count=miss_count,
+            num_real_reqs=state["num_real_reqs"],
+            host_cache=state["host_cache"],
+            device_buffer=skip_layer_buffer,
+            item_size_bytes=ITEM_SIZE_BYTES,
+            num_blocks=4,
+        )
+
+    run_once()
+    torch.cuda.synchronize()
+    for _ in range(WARMUP_ROUNDS):
+        run_once()
+    torch.cuda.synchronize()
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(ROUNDS):
+        run_once()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000.0 / ROUNDS
+
+
 @triton.testing.perf_report(
     triton.testing.Benchmark(
         x_names=["batch_size", "hot_buffer_size", "miss_rate", "miss_tokens_cnt"],
@@ -188,5 +251,35 @@ def benchmark_latency(
     return avg_us, avg_us, avg_us
 
 
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "hot_buffer_size", "miss_rate", "miss_tokens_cnt"],
+        x_vals=CONFIGS,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=STYLES,
+        ylabel="us",
+        plot_name="hisparse-planned-copy-latency",
+        args={},
+    )
+)
+def benchmark_planned_copy_latency(
+    batch_size: int,
+    hot_buffer_size: int,
+    miss_rate: float,
+    miss_tokens_cnt: int,
+    provider: str,
+) -> Tuple[float, float, float]:
+    assert provider == "jit"
+    batch_size = int(batch_size)
+    hot_buffer_size = int(hot_buffer_size)
+    miss_rate = float(miss_rate)
+    assert miss_tokens_cnt == batch_size * _miss_tokens_per_req(miss_rate)
+    avg_us = _time_planned_copy(batch_size, hot_buffer_size, miss_rate)
+    return avg_us, avg_us, avg_us
+
+
 if __name__ == "__main__":
     benchmark_latency.run(print_data=True)
+    benchmark_planned_copy_latency.run(print_data=True)