[Fix] Merge adjacent KV-row frees so a mid-page split under DCP cannot double-free (#38941)

This commit is contained in:
Kurt Shuster
2026-09-14 21:44:32 +08:00
committed by GitHub
parent 2123aca87e
commit f81fbc749a
3 changed files with 107 additions and 2 deletions
@@ -415,12 +415,14 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
"""Give back ascending, disjoint, half-open row-position ranges
of the ``kv`` record's row; one call keeps a shared page freed once.
"""
from sglang.srt.mem_cache.common import free_kv_row_segments
from sglang.srt.mem_cache.common import coalesce_ranges, free_kv_row_segments
row = self.req_to_token_pool.req_to_token[kv.req_pool_idx]
# Adjacent pieces whose seam falls inside one (DCP-widened) page would
# free that page twice; the allocator rejects that, so merge them first.
free_kv_row_segments(
self.token_to_kv_pool_allocator,
[(row[start:end], start) for start, end in ranges],
[(row[start:end], start) for start, end in coalesce_ranges(ranges)],
swa_evicted_seqlen=kv.swa_evicted_seqlen,
)
+11
View File
@@ -108,6 +108,17 @@ def free_swa_out_of_window_slots(
req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
def coalesce_ranges(ranges: list[tuple[int, int]]) -> list[tuple[int, int]]:
"""Merge adjacent half-open ranges so a split that falls mid-page frees that page once."""
merged: list[tuple[int, int]] = []
for start, end in ranges:
if merged and start == merged[-1][1]:
merged[-1] = (merged[-1][0], end)
else:
merged.append((start, end))
return merged
def free_kv_row_segments(
allocator: BaseTokenToKVPoolAllocator,
segments: list[tuple[torch.Tensor, int]],