fix(cuda_graph): zero out_cache_loc_swa on pad and use int32 (hybrid-SWA accuracy fix) (#24743)

This commit is contained in:
JoyFuture
2026-05-09 18:22:12 +08:00
committed by GitHub
parent ba625d5290
commit a309f1f8f4
2 changed files with 8 additions and 2 deletions
@@ -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:
@@ -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
)