minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix SWA cache loc slicing for all attention backends (#29460)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-06-27 20:37:58 -07:00
committed by GitHub
parent 4e6ce8e4ab
commit e6cbc8f5fe
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+9
View File
@@ -1101,6 +1101,15 @@ class KVWriteLoc:
loc: torch.Tensor loc: torch.Tensor
swa_loc: Optional[torch.Tensor] = None swa_loc: Optional[torch.Tensor] = None
def __post_init__(self):
# swa_out_cache_loc is computed once at metadata-init time from the
# full (possibly padded) out_cache_loc. Piecewise CUDA graphs later
# narrow out_cache_loc to real_num_tokens per layer, so swa_loc can
# be longer than loc. Slice to match since both are in the same
# per-token order.
if self.swa_loc is not None and self.swa_loc.shape[0] != self.loc.shape[0]:
self.swa_loc = self.swa_loc[: self.loc.shape[0]]
def unwrap_write_loc(loc_info): def unwrap_write_loc(loc_info):
"""Return ``(loc, swa_loc)`` from a ``KVWriteLoc`` or a bare loc tensor.""" """Return ``(loc, swa_loc)`` from a ``KVWriteLoc`` or a bare loc tensor."""