From afeaeccfa2f1f2fe9e245fabef2b2e58c34d13f3 Mon Sep 17 00:00:00 2001 From: DAI0818 <53027952+daii-0818@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:09:44 +0800 Subject: [PATCH] perf(hisparse): eliminate redundant swap output fill (#32483) Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- python/sglang/kernels/jit/csrc/hisparse.cuh | 27 +++++++++++++------ .../srt/managers/hisparse_coordinator.py | 1 - .../benchmark/kvcache/bench_hisparse.py | 1 - .../kernels/ops/kvcache/test_hisparse.py | 17 +++++++++++- 4 files changed, 35 insertions(+), 11 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/hisparse.cuh b/python/sglang/kernels/jit/csrc/hisparse.cuh index 50477dc44..d176deb18 100644 --- a/python/sglang/kernels/jit/csrc/hisparse.cuh +++ b/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(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]; + for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) { + int32_t device_loc = -1; + if (i < count) { + 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; } diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 3120ce0a3..48bf5db43 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/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 diff --git a/test/registered/kernels/benchmark/kvcache/bench_hisparse.py b/test/registered/kernels/benchmark/kvcache/bench_hisparse.py index 051dcba21..069ed2f88 100644 --- a/test/registered/kernels/benchmark/kvcache/bench_hisparse.py +++ b/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"], diff --git a/test/registered/kernels/ops/kvcache/test_hisparse.py b/test/registered/kernels/ops/kvcache/test_hisparse.py index 7f3104cce..8824099e4 100644 --- a/test/registered/kernels/ops/kvcache/test_hisparse.py +++ b/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, )