perf(hisparse): eliminate redundant swap output fill (#32483)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

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

@@ -217,10 +217,18 @@ __global__ void load_cache_to_device_buffer_kernel(
   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;
+  int32_t* req_top_k_device_locs = top_k_device_locs + bid * top_k_device_locs_stride;
+
+  // CUDA graph pads the batch to a captured size. Keep padded output rows
+  // invalid without a separate fill kernel.
+  if (bid >= num_real_reqs[0]) {
+    for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
+      req_top_k_device_locs[i] = -1;
+    }
+    return;
+  }
+
   const int warp_id = tid / WARP_SIZE;
   const int lane_id = tid % WARP_SIZE;
   const BallotMask lanes_before = (BallotMask(1) << lane_id) - BallotMask(1);
@@ -230,7 +238,6 @@ __global__ void load_cache_to_device_buffer_kernel(
 
   // 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;
@@ -241,11 +248,15 @@ __global__ void load_cache_to_device_buffer_kernel(
   // 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) {
+    for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
-      int32_t token_pos = req_top_k_tokens[i];
+      int32_t device_loc = -1;
-      if (token_pos >= 0) {
+      if (i < count) {
-        req_top_k_device_locs[i] = req_device_buffer_locs[token_pos];
+        int32_t token_pos = req_top_k_tokens[i];
+        if (token_pos >= 0) {
+          device_loc = req_device_buffer_locs[token_pos];
+        }
       }
+      req_top_k_device_locs[i] = device_loc;
     }
     return;
   }

python/sglang/srt/managers/hisparse_coordinator.py

@@ -813,7 +813,6 @@ class HiSparseCoordinator:
         num_reqs = req_pool_indices.size(0)
 
         top_k_indices = self.top_k_device_locs_buffer[:num_reqs]
-        top_k_indices.fill_(-1)
 
         swap_in_fn = (
             load_cache_to_device_buffer_dsv4_mla

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

@@ -125,7 +125,6 @@ def _time_kernel(batch_size: int, hot_buffer_size: int, miss_rate: float) -> flo
     def run_once():
         state["device_buffer_tokens"].copy_(state["initial_device_buffer_tokens"])
         state["lru_slots"].copy_(state["initial_lru_slots"])
-        state["top_k_device_locs"].fill_(-1)
         load_cache_to_device_buffer_mla(
             top_k_tokens=state["top_k_tokens"],
             device_buffer_tokens=state["device_buffer_tokens"],

test/registered/kernels/ops/kvcache/test_hisparse.py

@@ -100,6 +100,7 @@ def _run_kernel(
     seq_lens_dtype: torch.dtype = torch.int32,
     req_pool_indices: torch.Tensor | None = None,
     num_real_reqs: int | None = None,
+    output_fill_value: int = -1,
 ) -> torch.Tensor:
     batch_size = top_k_tokens.shape[0]
     if req_pool_indices is None:
@@ -111,7 +112,7 @@ def _run_kernel(
     if num_real_reqs is None:
         num_real_reqs = batch_size
 
-    out = torch.full_like(top_k_tokens, -1)
+    out = torch.full_like(top_k_tokens, output_fill_value)
     load_cache_to_device_buffer_mla(
         top_k_tokens=top_k_tokens,
         device_buffer_tokens=device_buffer_tokens,
@@ -312,6 +313,19 @@ def test_load_cache_to_device_buffer_fast_path(seq_lens_dtype: torch.dtype) -> N
     assert torch.equal(device_buffer.cpu(), device_buffer_before.cpu())
 
 
+def test_load_cache_to_device_buffer_fast_path_overwrites_stale_output() -> None:
+    state = _make_state([[9, 7, 3, 5, 11]], [[0, 1, 2, 3, -1]], [4])
+
+    out = _run_kernel(
+        top_k_tokens=torch.tensor([[1, -1, 0, 0]], dtype=torch.int32, device=DEVICE),
+        seq_len=2,
+        output_fill_value=123456,
+        **state,
+    )
+
+    assert torch.equal(out.cpu(), torch.tensor([[7, -1, -1, -1]], dtype=torch.int32))
+
+
 def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None:
     state = _long_case()
 
@@ -427,6 +441,7 @@ def test_load_cache_to_device_buffer_batched_with_padding() -> None:
         ),
         seq_lens=torch.tensor([8, 3, 8], dtype=torch.int32, device=DEVICE),
         num_real_reqs=2,
+        output_fill_value=123456,
         **state,
     )