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 )