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

Support unified memory page-envelope transfers in PD (#39477)

Browse Source 
Co-authored-by: yhzhuang <yhzhuang@fb.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
Co-authored-by: Yonghao Zhuang <yhzhuang@users.noreply.github.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
Yonghao Zhuang
2026-09-18 17:39:50 -07:00
committed by GitHub
co-authored by yhzhuang Lianmin Zheng Yonghao Zhuang Cheng Wan
parent d0730a0e8b
commit 5931fd60ee
27 changed files with 1012 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/fields/memory.py
+2 -4
View File
@@ -88,10 +88,8 @@ class Memory(msgspec.Struct):
"Replace the statically-partitioned hybrid-model pools (full-attn KV + "
"SWA/Mamba state) with one byte buffer split dynamically between "
"sub-pools. Requires the Triton attention / linear-attn / Mamba "
"backends. PD disaggregation is supported over mooncake at equal "
"attention TP with pp=1; not yet compatible with hierarchical / "
"host-tiered KV cache, prefill cuda-graph capture, or speculative "
"decoding other than DSPARK.",
"backends. Supported PD-disaggregation and speculative-decoding "
"configurations are validated at startup.",
] = False
enable_session_radix_cache: A[
bool,
python/sglang/srt/arg_groups/kv_cache_hook.py
+22 -5
View File
@@ -437,11 +437,12 @@ def handle_unified_memory_pool(server_args: Any) -> None:
if not cfg.enable_unified_memory:
return
if cfg.disaggregation_mode != "null":
# Constraints of the whole-envelope transfer; see
# UnifiedMLATokenToKVPool.get_contiguous_buf_infos.
assert cfg.disaggregation_transfer_backend == "mooncake", (
"--enable-unified-memory with PD disaggregation supports only "
"the mooncake transfer backend; got "
# Constraints of the whole-envelope transfer; see the unified MHA and
# MLA pool get_contiguous_buf_infos implementations.
supported_backends = server_args._unified_memory_pd_transfer_backends()
assert cfg.disaggregation_transfer_backend in supported_backends, (
"--enable-unified-memory with PD disaggregation supports only these "
f"transfer backends: {', '.join(sorted(supported_backends))}; got "
f"{cfg.disaggregation_transfer_backend!r}."
)
assert cfg.pp_size == 1, (
@@ -449,6 +450,13 @@ def handle_unified_memory_pool(server_args: Any) -> None:
"pipeline parallelism (whole-envelope transfer has no per-layer "
"entries to subset)."
)
assert not (
cfg.disaggregation_transfer_backend == "mooncake"
and cfg.speculative_algorithm is not None
), (
"--enable-unified-memory with PD disaggregation does not support "
"speculative decoding with the Mooncake transfer backend."
)
assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), (
"--enable-unified-memory with PD disaggregation requires lazy "
"compaction; unset SGLANG_DISABLE_LAZY_COMPACTION."
@@ -459,6 +467,15 @@ def handle_unified_memory_pool(server_args: Any) -> None:
"ships host/C4 rows straight from the allocator, bypassing the "
"virtual->physical translation the unified pool needs."
)
assert cfg.disaggregation_decode_retraction_backup != "host_pool", (
"--enable-unified-memory with PD disaggregation does not support "
"--disaggregation-decode-retraction-backup=host_pool; use "
"cpu_tensor (the automatic default for unified pools)."
)
assert not cfg.disaggregation_decode_enable_offload_kvcache, (
"--enable-unified-memory with PD disaggregation does not yet support "
"--disaggregation-decode-enable-offload-kvcache."
)
assert cfg.speculative_algorithm in (None, "DSPARK"), (
"--enable-unified-memory only supports --speculative-algorithm "
"DSPARK (chain draft); other speculative algorithms are not yet "
python/sglang/srt/disaggregation/decode.py
+127 -87
View File
@@ -455,6 +455,28 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail")
)
def _uses_swa_reservation(self) -> bool:
return (
self._uses_swa_tail_prealloc()
or self.token_to_kv_pool_allocator.prealloc_fits_assumes_reclaim()
)
def _prealloc_reservation_fits(
self,
full_tokens: int,
swa_tokens: int,
*,
full_allocatable_tokens: int,
swa_allocatable_tokens: Optional[int],
) -> bool:
return self.token_to_kv_pool_allocator.prealloc_fits(
self.tree_cache,
full_tokens,
swa_tokens,
full_budget_tokens=full_allocatable_tokens,
swa_budget_tokens=swa_allocatable_tokens,
)
def _release_matched_prefix_lock(self, req: Req) -> None:
if req.swa_prefix_lock_released:
self.tree_cache.dec_lock_ref(req.last_node, req.lock_receipt, skip_swa=True)
@@ -463,24 +485,16 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
self.tree_cache.dec_lock_ref(req.last_node, req.lock_receipt)
def _reclaim_swa_tail_capacity(
self, swa_tail_len: int, req_id: str
self, swa_tail_len: int, req_id: str, *, full_len: int = 0
) -> Optional[str]:
page_size = self.token_to_kv_pool_allocator.page_size
required = ceil_align(swa_tail_len, page_size)
available = self.token_to_kv_pool_allocator.swa_available_size()
if available < required:
self.tree_cache.evict_for_alloc(
EvictParams(swa_num_tokens=required - available)
)
available = self.token_to_kv_pool_allocator.swa_available_size()
if available < required:
return (
f"SWA eviction insufficient: needed={required}, "
f"available={available}, req={req_id}"
)
return None
allocator = self.token_to_kv_pool_allocator
page_size = allocator.page_size
shortfall = allocator.reclaim_for_prealloc(
self.tree_cache,
ceil_align(full_len, page_size),
ceil_align(swa_tail_len, page_size),
)
return None if shortfall is None else f"{shortfall}, req={req_id}"
# SWA caches expose full-attention accounting through full_* accessors.
def _radix_full_evictable(self) -> int:
@@ -581,8 +595,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:])
num_draft_entries = 0
if self.draft_token_to_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are
# always shared with a target model.
# Draft KV shares target virtual ids. Unified target KV is transferred
# with physical ids, so it needs a separate draft index vector.
draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
self.draft_token_to_kv_pool.get_contiguous_buf_infos()
)
@@ -814,30 +828,40 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
return len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool:
# HiSparse admits up to the host-backed logical capacity.
if self.scheduler.enable_hisparse:
capacity = self.scheduler.tp_worker.model_runner.max_token_pool_size
else:
capacity = self.max_total_num_tokens
input_len = self._rebootstrap_prefill_len(req)
if input_len > capacity:
message = f"Request {req.rid} exceeds the maximum number of tokens: {input_len} > {capacity}"
message = None
allocator = self.token_to_kv_pool_allocator
full_required, swa_required = self._prealloc_required_tokens(req)
if not self._uses_swa_tail_prealloc():
swa_required = full_required
ceiling_fits = allocator.prealloc_ceiling_fits(full_required, swa_required)
if ceiling_fits is False:
message = (
f"Request {req.rid} exceeds the unified FULL/SWA KV byte "
f"budget: full={full_required}, swa={swa_required}"
)
elif ceiling_fits is None:
# HiSparse admits up to the host-backed logical capacity.
capacity = (
self.scheduler.tp_worker.model_runner.max_token_pool_size
if self.scheduler.enable_hisparse
else self.max_total_num_tokens
)
input_len = self._rebootstrap_prefill_len(req)
if input_len > capacity:
message = f"Request {req.rid} exceeds the maximum number of tokens: {input_len} > {capacity}"
elif self._uses_swa_tail_prealloc():
_, swa_required = self._prealloc_required_tokens(req)
swa_capacity = self.token_to_kv_pool_allocator.size_swa
if swa_required > swa_capacity:
message = (
f"Request {req.rid} requires too many SWA KV tokens for "
f"decode preallocation: {swa_required} > {swa_capacity}"
)
if message is not None:
logger.error(message)
prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST)
self.scheduler.output_streamer.stream_output([req], req.return_logprob)
return True
if self._uses_swa_tail_prealloc():
_, swa_required = self._prealloc_required_tokens(req)
swa_capacity = self.token_to_kv_pool_allocator.size_swa
if swa_required > swa_capacity:
message = (
f"Request {req.rid} requires too many SWA KV tokens for "
f"decode preallocation: {swa_required} > {swa_capacity}"
)
logger.error(message)
prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST)
self.scheduler.output_streamer.stream_output([req], req.return_logprob)
return True
return False
def extend(self, reqs: List[Req], is_retracted: bool = False) -> None:
@@ -870,8 +894,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
# allocate memory
resumed_reqs = []
indices_to_remove = set()
uses_swa_tail_prealloc = self._uses_swa_tail_prealloc()
if uses_swa_tail_prealloc:
swa_allocatable_tokens = None
if self._uses_swa_reservation():
full_allocatable_tokens, swa_allocatable_tokens = (
self._swa_aware_allocatable_token_budgets(count_retracted=False)
)
@@ -888,17 +912,31 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
break
full_required, swa_required = self._prealloc_required_tokens(req)
if full_required > full_allocatable_tokens:
break
if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens:
if not self._prealloc_reservation_fits(
full_required,
swa_required,
full_allocatable_tokens=full_allocatable_tokens,
swa_allocatable_tokens=swa_allocatable_tokens,
):
break
if self.token_to_kv_pool_allocator.prealloc_fits_assumes_reclaim():
full_len, swa_len = self._prealloc_kv_lens(req)
if (
self._reclaim_swa_tail_capacity(swa_len, req.rid, full_len=full_len)
is not None
):
break
resumed_reqs.append(req)
indices_to_remove.add(i)
req.is_retracted = False
self._pre_alloc(req)
full_allocatable_tokens -= full_required
if uses_swa_tail_prealloc:
full_allocatable_tokens = self._allocatable_token_budgets(
count_retracted=False,
extra_reserved_reqs=len(resumed_reqs),
)
if swa_allocatable_tokens is not None:
swa_allocatable_tokens = self._swa_tail_allocatable_token_budget(
count_retracted=False,
extra_reserved_reqs=len(resumed_reqs),
@@ -1159,8 +1197,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
)
uses_swa_tail_prealloc = self._uses_swa_tail_prealloc()
swa_allocatable_tokens = 0
if uses_swa_tail_prealloc:
swa_allocatable_tokens = None
if self._uses_swa_reservation():
retractable_swa_tokens = sum(
self._swa_retractable_len(r) for r in self.scheduler.running_batch.reqs
)
@@ -1352,27 +1390,18 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
required_alloc_tokens + self.num_reserved_decode_tokens
)
if (
max(
required_tokens_for_request,
origin_input_len
- prefix_len
+ min(
decode_req.req.sampling_params.max_new_tokens,
CLIP_MAX_NEW_TOKEN,
)
- retractable_tokens,
full_required_for_admission = max(
required_tokens_for_request,
origin_input_len
- prefix_len
+ min(
decode_req.req.sampling_params.max_new_tokens,
CLIP_MAX_NEW_TOKEN,
)
> full_allocatable_tokens
):
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
self._release_matched_prefix_lock(decode_req.req)
break
if required_tokens_for_request > full_allocatable_tokens:
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
self._release_matched_prefix_lock(decode_req.req)
break
- retractable_tokens,
)
swa_required_for_admission = 0
swa_len = required_alloc_tokens
if uses_swa_tail_prealloc:
_, swa_required = self._prealloc_required_tokens(decode_req.req)
_, swa_len = self._prealloc_kv_lens(decode_req.req)
@@ -1380,19 +1409,28 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
decode_req.req.sampling_params.max_new_tokens,
CLIP_MAX_NEW_TOKEN,
)
if (
max(
swa_required,
swa_len + max_new_tokens - retractable_swa_tokens,
)
> swa_allocatable_tokens
):
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
self._release_matched_prefix_lock(decode_req.req)
break
swa_required_for_admission = max(
swa_required,
swa_len + max_new_tokens - retractable_swa_tokens,
)
elif swa_allocatable_tokens is not None:
swa_required_for_admission = full_required_for_admission
if not self._prealloc_reservation_fits(
full_required_for_admission,
swa_required_for_admission,
full_allocatable_tokens=full_allocatable_tokens,
swa_allocatable_tokens=swa_allocatable_tokens,
):
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
self._release_matched_prefix_lock(decode_req.req)
break
if swa_allocatable_tokens is not None:
reclaim_error = self._reclaim_swa_tail_capacity(
swa_len, decode_req.req.rid
swa_len,
decode_req.req.rid,
full_len=required_alloc_tokens,
)
if reclaim_error is not None:
if prefix_match is not None and prefix_match.l1_prefix_len > 0:
@@ -1431,7 +1469,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
extra_reserved_reqs=len(preallocated_reqs) + 1,
hicache_reserved_tokens=reserved_restore_tokens,
)
if uses_swa_tail_prealloc:
if swa_allocatable_tokens is not None:
swa_allocatable_tokens = self._swa_tail_allocatable_token_budget(
retractable_tokens=retractable_tokens,
retractable_swa_tokens=retractable_swa_tokens,
@@ -1442,6 +1480,9 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
page_size = self.token_to_kv_pool_allocator.page_size
kv_transfer_page_size = page_size
raw_kv_indices = self.req_to_token_pool.req_to_token[
decode_req.req.kv.req_pool_idx
][total_prefix_len:origin_input_len]
if self.scheduler.enable_hisparse:
# Direct-to-host sends host/C4 rows; keep allocator.page_size
# logical and use the compressed page size only for these indices.
@@ -1453,12 +1494,9 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
kv_indices = dst_kv_indices[: origin_input_len - prefix_len]
else:
# Only send delta indices (beyond prefix) to prefill.
kv_indices = self.req_to_token_pool.req_to_token[
decode_req.req.kv.req_pool_idx
][total_prefix_len:origin_input_len]
kv_indices = (
self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
kv_indices
raw_kv_indices
)
)
@@ -1735,7 +1773,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
# HiSparse pre-alloc only allocates logical indices, so the
# logical pool is the binding constraint for admission control.
available_size = logical_allocator.available_size()
elif self._uses_swa_tail_prealloc():
elif self._uses_swa_reservation():
available_size = self.token_to_kv_pool_allocator.full_available_size()
if get_disagg().disaggregation_decode_enable_radix_cache:
available_size += self._radix_full_evictable()
@@ -1802,8 +1840,10 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
# pool) over-reserves SWA in steady state. Cap by the actual
# remaining headroom up to per-req window cap.
window_size = self.scheduler.sliding_window_size or 0
swa_total = self.token_to_kv_pool_allocator.size_swa
swa_available = self.token_to_kv_pool_allocator.swa_available_size()
allocator = self.token_to_kv_pool_allocator
_, (swa_total, swa_available) = allocator.swa_capacity_and_available(
full_capacity=allocator.size_full, swa_capacity=allocator.size_swa
)
# Per-request SWA ring: cached prefixes still report swa_evictable, but
# evicting them frees no ring space.
swa_evictable = (
python/sglang/srt/disaggregation/mooncake/conn.py
+7 -6
View File
@@ -790,11 +790,18 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
# Published layer IDs give exact pairing; plain-MHA peers publish none
# and keep positional slicing.
has_layer_ids = bool(src_layer_ids or dst_layer_ids)
# Unified SWA publishes one page-envelope region even on an MHA backend.
is_single_region_swa = (
state_type == StateType.SWA
and len(src_data_ptrs) == 1
and len(dst_data_ptrs) == 1
)
if (
self.is_mla_backend
or self.is_hybrid_mla_backend
or force_flat
or has_layer_ids
or is_single_region_swa
):
# Layer IDs map PP-local buffers to global decode entries.
# Registrations without them retain the existing PP mapping.
@@ -1060,12 +1067,6 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
prefill_data_indices=prefill_kv_indices,
dst_data_indices=dst_kv_indices,
executor=executor,
# The unified pool registers ONE region holding every layer's K and
# V inside each page envelope. The MHA branch would half-split that
# single region into K and V halves and compute num_kv_layers = 0,
# transferring nothing at all; the flat branch addresses the region
# as-is. MLA-unified already reaches the flat branch via
# is_mla_backend, so this only adds the MHA-unified peer.
force_flat=get_memory().enable_unified_memory,
src_layer_ids=self.kv_args.kv_layer_ids,
dst_layer_ids=dst_layer_ids,
python/sglang/srt/disaggregation/prefill.py
+6 -5
View File
@@ -271,8 +271,8 @@ class PrefillBootstrapQueue:
)
num_draft_entries = 0
if draft_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are
# always shared with a target model.
# Draft KV shares target virtual ids. Unified target KV is transferred
# with physical ids, so it needs a separate draft index vector.
draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
draft_kv_pool.get_contiguous_buf_infos()
)
@@ -1463,14 +1463,14 @@ class SchedulerDisaggregationPrefillMixin:
for seg_start, seg_end in segments:
is_final_segment = seg_end == end_idx
kv_indices = self.req_to_token_pool.req_to_token[
raw_kv_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, seg_start:seg_end
]
# Unified memory: req_to_token holds VIRTUAL ids; the transfer needs
# physical ones. Per segment, since each is its own gather.
kv_indices = (
self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer(
kv_indices
raw_kv_indices
)
)
page_indices = kv_to_page_indices(kv_indices, page_size)
@@ -1479,9 +1479,10 @@ class SchedulerDisaggregationPrefillMixin:
len(page_indices), segment_is_last
):
continue
send_state_indices = state_indices if segment_is_last else None
req.disagg_kv_sender.send(
page_indices,
state_indices if segment_is_last else None,
send_state_indices,
num_kv_tokens=seg_end - seg_start,
)
req.start_send_idx = end_idx
python/sglang/srt/mem_cache/allocator/base.py
+37
View File
@@ -95,6 +95,43 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
0, min(max_new_tokens, token_capacity - paged_input - self.page_size - 1)
)
def prealloc_fits_assumes_reclaim(self) -> bool:
"""Whether `prealloc_fits` answers about the state reachable AFTER
reclaiming the evictable pages, so admitting on it still owes the
reclaim. False when the answer describes the pool as it stands.
"""
return False
def prealloc_ceiling_fits(self, full_tokens: int, swa_tokens: int) -> bool | None:
"""Whether a demand this size could EVER be preallocated, or None when
this pool has no ceiling of its own and the caller's token capacity is
the only bound.
"""
return None
def prealloc_fits(
self,
tree_cache,
full_tokens: int,
swa_tokens: int,
*,
full_budget_tokens: int,
swa_budget_tokens: int | None = None,
) -> bool:
"""Whether a decode-node preallocation of this size fits.
The budgets are the scheduler's policy: what each side has left once
decode headroom and retraction are reserved. Separate buffers make the
two sides independent, so each is checked against its own budget and
``tree_cache`` is never read -- what it could reclaim is already
inside that budget. A pool that cuts both sides from one buffer
overrides this to price them together, since a per-side token budget
cannot express a shared byte envelope.
"""
return full_tokens <= full_budget_tokens and (
swa_budget_tokens is None or swa_tokens <= swa_budget_tokens
)
def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None:
"""Evict unlocked prefix-cache entries until this allocator can serve
``num_tokens`` or nothing evictable remains.
python/sglang/srt/mem_cache/allocator/hisparse.py
-4
View File
@@ -349,10 +349,6 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def translate_swa_indices_for_transfer(
self, kv_indices: torch.Tensor
) -> torch.Tensor:
# Delegated like the read-path translate above: this composite is not a
# SWA allocator itself, so it inherits neither the default nor an
# override, and the PD payload path calls this on whatever allocator
# the scheduler holds.
return self.logical_attn_allocator.translate_swa_indices_for_transfer(
kv_indices
)
python/sglang/srt/mem_cache/allocator/swa.py
+25 -8
View File
@@ -166,6 +166,30 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens
)
def reclaim_for_prealloc(
self, tree_cache, full_tokens: int, swa_tokens: int
) -> str | None:
"""Free room for a decode-node preallocation; None means it is ready.
Returns a description of the shortfall when it cannot be met, for the
caller to attach to whichever request it was admitting. Separate
buffers make the sliding-window side the only one that needs
reclaiming here, since the full side is priced by the caller's budget.
"""
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
available = self.swa_available_size()
if available < swa_tokens:
tree_cache.evict_for_alloc(
EvictParams(swa_num_tokens=swa_tokens - available)
)
available = self.swa_available_size()
if available < swa_tokens:
return (
f"SWA eviction insufficient: needed={swa_tokens}, available={available}"
)
return None
def swa_capacity_and_available(self, *, full_capacity, swa_capacity):
return (
(full_capacity, self.full_available_size()),
@@ -232,14 +256,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def translate_swa_indices_for_transfer(
self, kv_indices: torch.Tensor
) -> torch.Tensor:
"""Sliding-window token ids as the PD transfer engine addresses them.
The sibling of `translate_kv_indices_for_transfer` for the SWA state
component. On a static pool the sliding-window buffers are indexed by
the same ids the kernels use, so the read-path translate IS the answer.
A virtual-id pool must override: the transfer addresses raw bytes and
needs PHYSICAL ids, not kernel-facing ones.
"""
"""Map full-pool token ids to SWA-buffer token ids for PD transfer."""
return self.translate_loc_from_full_to_swa(kv_indices)
def alloc(self, need_size: int):
python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py
+157 -48
View File
@@ -388,6 +388,8 @@ class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor,
extend_num_tokens: int,
*,
num_swa_pages: Optional[int] = None,
) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
"""Run the full side's paged extend and report which virtual PAGES it
newly took. Returns (virtual TOKEN ids, new virtual PAGE ids), or None
@@ -402,7 +404,10 @@ class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
prefix_lens=prefix_lens_cpu,
)
need_tokens = num_new_pages * self.page_size
if not self.ensure_capacity(need_tokens, need_tokens):
swa_tokens = (
need_tokens if num_swa_pages is None else num_swa_pages * self.page_size
)
if not self.ensure_capacity(need_tokens, swa_tokens):
return None
# Snapshot the virtual PAGES the kernel will consume; clone so swa keeps
@@ -478,14 +483,22 @@ class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
sink and is skipped by `free`'s `swa_v2p_page > 0` mask -- exactly the
out-of-window state the ratchet produces via `free_swa`.
Admission is priced at the FULL side's page count, as plain
`alloc_extend` is: pessimistic when the tail is short, but it reuses
the composite's audited joint capacity path, and the bytes actually
held still follow the tail.
Admission prices FULL's new pages and only the new pages in the SWA
tail. A partial prefix page is already bound and costs no new SWA page.
"""
assert len(prefix_lens_cpu) == 1
assert 0 <= swa_tail_len <= extend_num_tokens
with record_function("UnifiedSWAAlloc.alloc_extend_swa_tail"):
prefix_len = int(prefix_lens_cpu[0])
seq_len = int(seq_lens_cpu[0])
first_new_page = (prefix_len + self.page_size - 1) // self.page_size
first_tail_page = (seq_len - swa_tail_len) // self.page_size
num_swa_pages = (
(seq_len + self.page_size - 1) // self.page_size
- max(first_new_page, first_tail_page)
if swa_tail_len
else 0
)
extended = self._extend_in_virtual_space(
prefix_lens,
prefix_lens_cpu,
@@ -493,6 +506,7 @@ class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator):
seq_lens_cpu,
last_loc,
extend_num_tokens,
num_swa_pages=num_swa_pages,
)
if extended is None:
return None
@@ -819,6 +833,63 @@ class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase):
"""No float in a two-END chain -- nothing can slide."""
return None
def prealloc_fits_assumes_reclaim(self) -> bool:
return True
def prealloc_ceiling_fits(self, full_tokens: int, swa_tokens: int) -> bool | None:
return self.can_reserve(full_tokens, swa_tokens, empty_pool=True)
def reclaim_for_prealloc(
self, tree_cache, full_tokens: int, swa_tokens: int
) -> str | None:
"""Reclaim both sides together: freeing FULL pages can open SWA room
and the reverse, so the shared envelope is the only gate worth
re-checking."""
ready = self.evict_to_free_tokens(
tree_cache, full_tokens, swa_num_tokens=swa_tokens
)
if ready is None:
ready = self.ensure_capacity(full_tokens, swa_tokens)
if ready:
return None
return (
"Unified FULL/SWA byte reclamation insufficient: "
f"needed=({full_tokens}, {swa_tokens})"
)
def prealloc_fits(
self,
tree_cache,
full_tokens: int,
swa_tokens: int,
*,
full_budget_tokens: int,
swa_budget_tokens: int | None = None,
) -> bool:
"""Price both sides against the shared byte envelope.
There is no per-side capacity for the scheduler's budget to be
compared against, so the gap between that budget and what this side
can currently hand out is folded back into the demand; `can_reserve`
then prices the whole ask in bytes. Reachable only for hybrid-SWA
models, so the tree's `full_*` accounting is the full-attention one.
"""
full_evictable_tokens = tree_cache.full_evictable_size()
swa_evictable_tokens = tree_cache.swa_evictable_size()
full_tokens += (
self.full_available_size() + full_evictable_tokens - full_budget_tokens
)
if swa_budget_tokens is not None:
swa_tokens += (
self.swa_available_size() + swa_evictable_tokens - swa_budget_tokens
)
return self.can_reserve(
full_tokens,
swa_tokens,
full_evictable_tokens=full_evictable_tokens,
swa_evictable_tokens=swa_evictable_tokens,
)
def reclaim_plan(
self,
full_tokens: int | float,
@@ -1086,14 +1157,17 @@ class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase):
hi = mid - 1
return lo
def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None:
def evict_to_free_tokens(
self, tree_cache, num_tokens: int, *, swa_num_tokens: Optional[int] = None
) -> bool | None:
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
if tree_cache is None or tree_cache.is_chunk_cache():
return
required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens
reclaim_plan = self.reclaim_plan(
num_tokens,
num_tokens,
required_swa,
full_evictable_tokens=tree_cache.full_evictable_size(),
swa_evictable_tokens=tree_cache.swa_evictable_size(),
)
@@ -1105,7 +1179,7 @@ class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase):
EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim)
)
# A zero-reclaim plan can still depend on compaction before allocation.
return self.ensure_capacity(num_tokens, num_tokens)
return self.ensure_capacity(num_tokens, required_swa)
def verify_byte_accounting(self) -> List[str]:
return (
@@ -1234,69 +1308,104 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase):
if (
full_tokens < 0
or swa_tokens < 0
or full_tokens != swa_tokens
or full_evictable_tokens
or swa_evictable_tokens
or empty_pool
):
return False
return full_tokens <= self.available_size()
return self._fits_page_demand(
math.ceil(full_tokens / self.page_size),
math.ceil(swa_tokens / self.page_size),
)
def prealloc_fits(
self,
tree_cache,
full_tokens: int,
swa_tokens: int,
*,
full_budget_tokens: int,
swa_budget_tokens: int | None = None,
) -> bool:
"""Price the pair on the float chain's grid, then against the budgets.
Each side's `available_size` takes `schedulable_available_size()`,
which credits the peer's drainable holes, so the two are backed by the
same bytes and a pair that fits each side alone can fail together. The
budgets still apply on top: they carry decode headroom this allocator
cannot see.
"""
page_size = self.page_size
if not self._fits_page_demand(
-(-full_tokens // page_size), -(-swa_tokens // page_size)
):
return False
return full_tokens <= full_budget_tokens and (
swa_budget_tokens is None or swa_tokens <= swa_budget_tokens
)
def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool:
if full_tokens < 0 or swa_tokens < 0 or full_tokens != swa_tokens:
if full_tokens < 0 or swa_tokens < 0:
return False
if full_tokens == 0:
if self.can_reserve(full_tokens, swa_tokens):
return True
need_tokens = int(full_tokens)
if need_tokens <= self.available_size():
for allocator in self._flush_targets():
allocator.flush_for_allocation()
if self.can_reserve(full_tokens, swa_tokens):
return True
return _relieve_for_alloc(self, need_tokens)
_float_open_short_side(
self.swa_attn_allocator,
{
self.full_attn_allocator: -(-full_tokens // self.page_size),
self.swa_attn_allocator: -(-swa_tokens // self.page_size),
self.mamba_allocator: 0,
},
)
return self.can_reserve(full_tokens, swa_tokens)
def _compute_available_size(self) -> int:
"""Joint TOKENS for `alloc(N)`: N costs N full pages AND N swa pages, drawn
from DIFFERENT bands -- full extends only into the high band, the float into
either side but only ONE per batch alloc. Feasibility is monotone in N, so
binary search; the order matches the alloc path (full takes the high band).
"""
def _fits_page_demand(self, full_pages: int, swa_pages: int) -> bool:
"""Price FULL first, then SWA in one contiguous band on the float grid."""
fa, sa = self.full_attn_allocator, self.swa_attn_allocator
e_f = fa.entry_bytes_per_page
# full is grow-down: its chain gap IS the high band.
b_high = fa._current_gap_bytes()
h_f = len(fa._free_phys_pages) if fa.lazy_compaction else 0
h_s = sa._hole_pages()
r_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
r_s = sa.num_pages - sa.min_page_index - sa._allocated_pages()
def feasible(n: int) -> bool:
if n > h_f + r_f or n > h_s + r_s:
return False
ext_f = max(0, n - h_f)
if ext_f * e_f > b_high:
return False
ext_s = max(0, n - h_s)
# On the float's page grid, never in raw bytes: a byte budget
# credits a page `take_physical_pages` cannot yield.
full_low_after = fa._byte_low_frontier() - ext_f * e_f
if sa._is_frontier_transparent():
room = sa.pages_in_band(
low_byte=sa._chain_high_frontier_below_bytes(),
high_byte=full_low_after,
)
return ext_s <= room
p_low = sa.pages_in_band(
if full_pages > h_f + r_f or swa_pages > h_s + r_s:
return False
full_bytes = max(0, full_pages - h_f) * fa.entry_bytes_per_page
if full_bytes > fa._current_gap_bytes():
return False
ext_s = max(0, swa_pages - h_s)
full_low_after = fa._byte_low_frontier() - full_bytes
if sa._is_frontier_transparent():
room = sa.pages_in_band(
low_byte=sa._chain_high_frontier_below_bytes(),
high_byte=sa._byte_low_frontier(),
)
p_high = sa.pages_in_band(
low_byte=sa._byte_high_frontier(),
high_byte=full_low_after,
)
return ext_s <= max(p_low, p_high)
return ext_s <= room
p_low = sa.pages_in_band(
low_byte=sa._chain_high_frontier_below_bytes(),
high_byte=sa._byte_low_frontier(),
)
p_high = sa.pages_in_band(
low_byte=sa._byte_high_frontier(),
high_byte=full_low_after,
)
return ext_s <= max(p_low, p_high)
def _compute_available_size(self) -> int:
"""Joint TOKENS for equal FULL/SWA demand, using the same page predicate
as tail allocation. FULL takes the high band before SWA binds its pages.
"""
fa, sa = self.full_attn_allocator, self.swa_attn_allocator
h_f = len(fa._free_phys_pages) if fa.lazy_compaction else 0
h_s = sa._hole_pages()
r_f = fa.num_pages - fa.min_page_index - fa._allocated_pages()
r_s = sa.num_pages - sa.min_page_index - sa._allocated_pages()
lo_n, hi_n = 0, min(h_f + r_f, h_s + r_s)
while lo_n < hi_n:
mid = (lo_n + hi_n + 1) // 2
if feasible(mid):
if self._fits_page_demand(mid, mid):
lo_n = mid
else:
hi_n = mid - 1
python/sglang/srt/mem_cache/allocator/unified_mamba.py
+24
View File
@@ -307,6 +307,30 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
)
return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64))
def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
"""Retraction backup for the FULL + mamba pair.
`Req.offload_kv_cache` hands over `req_to_token` rows, which hold
VIRTUAL ids here; both unified full pools index their host copy by
PHYSICAL ids. The mamba side is already slot-addressed and is
translated by the pool.
"""
return self._kvcache.get_cpu_copy(
self.full_attn_allocator.translate_kv_loc(indices.to(torch.int64)),
mamba_indices=mamba_indices,
req_pool_index=req_pool_index,
)
def load_cpu_copy(
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
return self._kvcache.load_cpu_copy(
kv_cache_cpu,
self.full_attn_allocator.translate_kv_loc(indices.to(torch.int64)),
mamba_indices=mamba_indices,
req_pool_index=req_pool_index,
)
def _move_gate_targets(self):
"""Every member a compaction gate must cover. The mamba end is gated
even where its state is not itself transferred: the gate is about the
python/sglang/srt/mem_cache/allocator/unified_sub_pool.py
+5 -9
View File
@@ -265,7 +265,11 @@ def install_move_gate(
class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
"""Allocator for one sub-pool over a `UnifiedKVPool`."""
"""Allocator for one sub-pool over a `UnifiedKVPool`.
``need_sort`` applies to transfer-facing physical ids, not virtual ids.
Physical free pages are sorted during compaction.
"""
# Capacity-bearing state: any rebind bumps `_capacity_epoch`, invalidating
# the chain's capacity memos (see `_CapacityField`).
@@ -1207,11 +1211,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
if not _relieve_for_alloc(self, need_tokens):
return None
bs = len(prefix_lens)
if self.need_sort and extend_num_tokens // self.page_size + bs + 1 > len(
self.free_virtual_ids
):
self.merge_and_sort_free()
# Snapshot the virtual pages the kernel will consume, to bind them
# to physical pages afterward.
if num_new_pages > 0:
@@ -1274,9 +1273,6 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
if need_tokens > self.available_size():
if not _relieve_for_alloc(self, need_tokens):
return None
if self.need_sort and bs > len(self.free_virtual_ids):
self.merge_and_sort_free()
# Most decode steps reuse the prefix's tail page -> num_new_pages == 0.
if num_new_pages > 0:
new_virtual_pages = self.free_virtual_ids[:num_new_pages].clone()
python/sglang/srt/mem_cache/kv_cache_builder.py
+4 -12
View File
@@ -154,19 +154,11 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str:
if tp_worker.is_hybrid_swa
else None
)
# Host-pool retraction transfers full and sliding-window components
# only, so a model with recurrent state stays on cpu_tensor.
#
# The unified pool is excluded for the same reason hierarchical cache is
# (see `handle_unified_memory_pool`): the host-transfer path indexes the
# device buffers with the ids it is handed, and under the unified pool
# those are VIRTUAL. It also cannot be sized from `kv_cache.size`, which
# is a KERNEL-FACING row count (`num_pages * 2 * layer_num * page_size`)
# rather than a token capacity -- gpt-oss-20b reports 85M "tokens" and
# asks for 418 GB of host memory per component.
# Host-pool retraction does not address unified page envelopes or
# recurrent state, so those configurations stay on cpu_tensor.
supports_host_pool = (
not uses_ssm_state(tp_worker.model_runner.model_config)
and not memory.enable_unified_memory
not memory.enable_unified_memory
and not uses_ssm_state(tp_worker.model_runner.model_config)
and (
isinstance(kv_cache, MHATokenToKVPool)
or (isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0)
python/sglang/srt/mem_cache/kv_cache_configurator.py
+1 -26
View File
@@ -868,32 +868,7 @@ class KVCacheConfigurator:
assert not self.use_mla_backend, (
"unified memory pool does not support MLA-SWA hybrid yet"
)
# Mirror the non-shared path's extra_max_context_len computation.
extra_max_context_len = 4
if get_spec().speculative_num_draft_tokens is not None:
extra_max_context_len += get_spec().speculative_num_draft_tokens
if get_disagg().disaggregation_mode == "decode":
# A decode node hands out request rows to PREALLOCATED transfers on
# top of its running set, so it needs the extra-slot pool (and the
# `pre_alloc_size` the scheduler's invariant checker reads). Mirrors
# `_build_req_to_token_pool`'s decode branch; the mamba composite
# already takes `decode_pre_alloc_size` the same way.
from sglang.srt.disaggregation.decode import DecodeReqToTokenPool
req_to_token_pool = DecodeReqToTokenPool(
size=max_num_reqs,
max_context_len=self.model_config.context_len + extra_max_context_len,
device=self.device,
enable_memory_saver=get_exec().features.enable_memory_saver,
pre_alloc_size=get_disagg().disaggregation_decode_extra_slots,
)
else:
req_to_token_pool = ReqToTokenPool(
size=max_num_reqs,
max_context_len=self.model_config.context_len + extra_max_context_len,
device=self.device,
enable_memory_saver=get_exec().features.enable_memory_saver,
)
req_to_token_pool = self._build_req_to_token_pool(max_num_reqs=max_num_reqs)
head_num = self.model_config.get_num_kv_heads(
get_parallel().attn_tp_size, get_parallel().attn_dcp_size
python/sglang/srt/mem_cache/unified_memory_pool.py
+41 -18
View File
@@ -624,29 +624,29 @@ class UnifiedMHATokenToKVPool(MHATokenToKVPool):
env[tgt_pages] = env[src_pages]
def get_contiguous_buf_infos(self):
"""PD-transfer registration: ONE entry, the raw buffer, addressed as
``raw_ptr + physical_page_id * page_envelope_bytes``.
"""Register the raw buffer as physical page envelopes for PD transfer.
Same whole-envelope contract as `UnifiedMLATokenToKVPool`: the transfer
item is one page across ALL layers and both K and V, because the
per-layer views overlap inside the envelope and index in kernel-facing
ids. A peer must therefore build an identical spec -- enforced on the
wire by `_validate_envelope_kv_layout`.
Full and SWA expose the same allocation with different envelope sizes;
the transfer backend preserves both logical entries while deduplicating
the underlying memory registration.
"""
# The address formula omits the anchor; a nonzero one would mis-address.
assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0
raw = self._unified_buffer._raw
return [raw.data_ptr()], [raw.numel()], [self._page_bytes]
def get_cpu_copy(self, indices, mamba_indices=None):
raise NotImplementedError(
"CPU offloading is unsupported under the unified layout."
)
def _physical_to_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor:
return (indices // self.page_size) * (
self.page_size * self.kernel_page_blocks
) + indices % self.page_size
def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
raise NotImplementedError(
"CPU offloading is unsupported under the unified layout."
)
def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
"""Translate physical host-pool ids for the page-major parent path."""
return super().get_cpu_copy(self._physical_to_kernel_indices(indices))
def load_cpu_copy(
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
super().load_cpu_copy(kv_cache_cpu, self._physical_to_kernel_indices(indices))
def set_kv_buffer_prefix_valid(self, *args, **kwargs):
raise NotImplementedError(
@@ -738,6 +738,22 @@ class UnifiedMLATokenToKVPool(MLATokenToKVPool):
raw = self._unified_buffer._raw
return [raw.data_ptr()], [raw.numel()], [self._page_bytes]
def _physical_to_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor:
"""Physical TOKEN ids -> the kernel-facing ids this class's `kv_buffer`
views are indexed by; the formula is the one in the class docstring."""
return (indices // self.page_size) * (
self.page_size * self.kernel_page_blocks
) + indices % self.page_size
def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
"""Translate physical host-pool ids for the page-major parent path."""
return super().get_cpu_copy(self._physical_to_kernel_indices(indices))
def load_cpu_copy(
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
super().load_cpu_copy(kv_cache_cpu, self._physical_to_kernel_indices(indices))
def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
"""Relocate whole page envelopes.
@@ -1676,7 +1692,7 @@ class UnifiedSWAKVPool(SWAKVPool):
swa_cpu = None
if bool(valid.any().item()):
swa_cpu = self.swa_kv_pool.get_cpu_copy(swa_phys[valid])
return {"full": full_cpu, "swa": swa_cpu}
return {"full": full_cpu, "swa": swa_cpu, "swa_mask": valid.cpu()}
def load_cpu_copy(
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
@@ -1689,7 +1705,14 @@ class UnifiedSWAKVPool(SWAKVPool):
if kv_cache_cpu.get("swa") is not None:
assert self._swa_allocator is not None
swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator)
self.swa_kv_pool.load_cpu_copy(kv_cache_cpu["swa"], swa_phys)
old_swa_mask = kv_cache_cpu["swa_mask"].to(indices.device)
assert old_swa_mask.shape == indices.shape
row_mask = (swa_phys >= 0)[old_swa_mask].cpu()
swa_phys = swa_phys[old_swa_mask][row_mask.to(indices.device)]
if swa_phys.numel() == 0:
return
swa_cpu = self._filter_swa_cpu_copy(kv_cache_cpu["swa"], row_mask)
self.swa_kv_pool.load_cpu_copy(swa_cpu, swa_phys)
class UnifiedSWAPoolBundle(NamedTuple):
python/sglang/srt/server_args.py
+3
View File
@@ -341,6 +341,9 @@ class ServerArgs:
# _handle_page_major_kv_layout); the model-family gate is enforced at pool
# construction in model_runner_kv_cache_mixin._init_pools.
def _unified_memory_pd_transfer_backends(self) -> set[str]:
return {"mooncake"}
@staticmethod
def add_cli_args(parser: argparse.ArgumentParser):
python/sglang/test/separate_buffer_allocator_double.py
+55
View File
@@ -0,0 +1,55 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Allocator double for scheduler tests that admit requests.
`DecodePreallocQueue` asks the allocator to price a preallocation rather than
doing the arithmetic itself, so a bare `MagicMock` returns a truthy `Mock` and
the admission decision under test stops being made anywhere. Binding the real
separate-buffer implementations keeps the arithmetic live while leaving the
per-test stubs (`size_swa`, `swa_available_size`, ...) in charge of the state.
"""
from unittest.mock import MagicMock
from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
# Bound on the double, so each reads the stubs the caller set on it.
_SEPARATE_BUFFER_METHODS = {
"prealloc_fits_assumes_reclaim": BaseTokenToKVPoolAllocator.prealloc_fits_assumes_reclaim,
"prealloc_ceiling_fits": BaseTokenToKVPoolAllocator.prealloc_ceiling_fits,
"prealloc_fits": BaseTokenToKVPoolAllocator.prealloc_fits,
"reclaim_for_prealloc": SWATokenToKVPoolAllocator.reclaim_for_prealloc,
"swa_capacity_and_available": SWATokenToKVPoolAllocator.swa_capacity_and_available,
}
def bind_separate_buffer_capacity(allocator) -> None:
"""Make `allocator` price capacity like a pool whose sides own their own
buffers. Call on any allocator double a `DecodePreallocQueue` will read."""
for name, impl in _SEPARATE_BUFFER_METHODS.items():
setattr(
allocator,
name,
(lambda impl: lambda *args, **kwargs: impl(allocator, *args, **kwargs))(
impl
),
)
def separate_buffer_allocator_double(**attrs) -> MagicMock:
"""A `MagicMock` allocator that prices capacity as separate buffers."""
allocator = MagicMock(**attrs)
bind_separate_buffer_capacity(allocator)
return allocator