[Fix] Merge adjacent KV-row frees so a mid-page split under DCP cannot double-free (#38941)
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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]],
|
||||
|
||||
Reference in New Issue
Block a user