From a309f1f8f4cb73dffa58010a1c37f569bb535b48 Mon Sep 17 00:00:00 2001 From: JoyFuture <35593546+JoyFuture@users.noreply.github.com> Date: Sat, 9 May 2026 18:22:12 +0800 Subject: [PATCH] fix(cuda_graph): zero out_cache_loc_swa on pad and use int32 (hybrid-SWA accuracy fix) (#24743) --- python/sglang/srt/model_executor/cuda_graph_runner.py | 8 +++++++- .../srt/model_executor/piecewise_cuda_graph_runner.py | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 00ca12fff..b04ee47cd 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -182,7 +182,7 @@ class DecodeInputBuffers(ForwardInputBuffers): seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32) out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype) out_cache_loc_swa = ( - torch.zeros((max_num_token,), dtype=torch.int64) + torch.zeros((max_num_token,), dtype=torch.int32) if is_hybrid_swa else None ) @@ -289,6 +289,12 @@ class DecodeInputBuffers(ForwardInputBuffers): if bs != raw_bs: self.seq_lens.fill_(seq_len_fill_value) self.out_cache_loc.zero_() + # Padded SWA indices left over from a previous replay would point + # into real SWA slots, so set_kv_buffer on padded tokens would + # corrupt active requests' KV. Zero the whole buffer so padded + # positions map to the sentinel slot (matches piecewise runner). + if self.out_cache_loc_swa is not None: + self.out_cache_loc_swa.zero_() if self.mamba_track_indices is not None: self.mamba_track_indices.zero_() if self.mamba_track_mask is not None: diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index c89fa5958..da547df15 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -238,7 +238,7 @@ class PiecewiseCudaGraphRunner: (self.max_num_tokens,), dtype=self._cache_loc_dtype() ) out_cache_loc_swa = ( - torch.zeros((self.max_num_tokens,), dtype=torch.int64) + torch.zeros((self.max_num_tokens,), dtype=torch.int32) if model_runner.is_hybrid_swa else None )