99b669f8b9
Prefill CP only needs the local hidden shard for DeepSeek NextN draft extend. The change adds a draft shared-KV path that captures target hidden locally, feeds only the CP-local slice into the draft model, and keeps draft KV writes/transfers on the same shared logical-to-physical page mapping as target KV.\n\nDebug logs are gated behind SGLANG_CP_DRAFT_SHARED_KV_DEBUG and cover scheduler pool selection, KV manager buffer registration, local physical writes, prefill sender filtering, transfer pages, and decode commit metadata so ETE runs can prove draft KV is sharded rather than full-concatenated on a prefill rank.\n\nConstraint: Prefill runs CP while decode remains DP, so prefill must avoid full hidden/KV materialization but decode still receives full logical KV pages.\nRejected: Keep draft extend on full hidden state | preserves correctness but wastes prefill memory and defeats CP shared-KV intent.\nRejected: Transfer draft KV with a separate mapping | target and draft pools share req_to_token logical indices, so duplicating mapping adds risk without benefit.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the debug logs until ETE evidence confirms draft MLA/index writes and transfer pages are CP-sharded on all ranks.\nTested: Remote compileall for changed CP draft, transfer, scheduler, NSA index, MLA write, and EAGLE files.\nNot-tested: Full GLM-5 EAGLE ETE with SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1 after this logging addition; local pytest intentionally not run.
740 lines
25 KiB
Python
740 lines
25 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
from sglang.srt.mem_cache.base_prefix_cache import (
|
|
BasePrefixCache,
|
|
EvictParams,
|
|
EvictResult,
|
|
)
|
|
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
|
|
build_in_seq_page_compute_owners,
|
|
get_in_seq_page_compute_owner_unavailable_reason,
|
|
)
|
|
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
|
|
from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
|
|
from sglang.srt.server_args import get_global_server_args
|
|
from sglang.srt.utils import support_triton
|
|
from sglang.srt.utils.common import ceil_align
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
|
|
|
|
# Needs 2 + 1 slots for mamba request with prefix cache. 2 for ping pong cache, 1 for running mamba state.
|
|
MAMBA_STATE_PER_REQ_PREFIX_CACHE = 3
|
|
MAMBA_STATE_PER_REQ_NO_CACHE = 1
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _log_cp_shared_kv_alloc_fallback(
|
|
reason: str,
|
|
message: str,
|
|
*args,
|
|
) -> None:
|
|
logger.info(
|
|
"CP shared KV compute-owner allocation fallback (%s): " + message,
|
|
reason,
|
|
*args,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def write_req_to_token_pool_triton(
|
|
req_to_token_ptr, # [max_batch, max_context_len]
|
|
req_pool_indices,
|
|
prefix_tensors,
|
|
pre_lens,
|
|
seq_lens,
|
|
extend_lens,
|
|
out_cache_loc,
|
|
req_to_token_ptr_stride: tl.constexpr,
|
|
):
|
|
BLOCK_SIZE: tl.constexpr = 512
|
|
pid = tl.program_id(0)
|
|
|
|
req_pool_index = tl.load(req_pool_indices + pid)
|
|
pre_len = tl.load(pre_lens + pid)
|
|
seq_len = tl.load(seq_lens + pid)
|
|
prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64))
|
|
|
|
# write prefix
|
|
num_loop = tl.cdiv(pre_len, BLOCK_SIZE)
|
|
for i in range(num_loop):
|
|
offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
|
|
mask = offset < pre_len
|
|
value = tl.load(prefix_tensor + offset, mask=mask)
|
|
tl.store(
|
|
req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset,
|
|
value,
|
|
mask=mask,
|
|
)
|
|
|
|
# NOTE: This can be slow for large bs
|
|
cumsum_start = tl.cast(0, tl.int64)
|
|
for i in range(pid):
|
|
cumsum_start += tl.load(extend_lens + i)
|
|
|
|
num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE)
|
|
for i in range(num_loop):
|
|
offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
|
|
mask = offset < (seq_len - pre_len)
|
|
value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask)
|
|
tl.store(
|
|
req_to_token_ptr
|
|
+ req_pool_index * req_to_token_ptr_stride
|
|
+ offset
|
|
+ pre_len,
|
|
value,
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
def write_cache_indices(
|
|
out_cache_loc: torch.Tensor,
|
|
req_pool_indices_tensor: torch.Tensor,
|
|
req_pool_indices_cpu: torch.Tensor,
|
|
prefix_lens_tensor: torch.Tensor,
|
|
prefix_lens_cpu: torch.Tensor,
|
|
seq_lens_tensor: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
extend_lens_tensor: torch.Tensor,
|
|
extend_lens_cpu: torch.Tensor,
|
|
prefix_tensors: list[torch.Tensor],
|
|
req_to_token_pool: ReqToTokenPool,
|
|
):
|
|
if support_triton(get_global_server_args().attention_backend):
|
|
prefix_pointers = torch.tensor(
|
|
[t.data_ptr() for t in prefix_tensors],
|
|
device=req_to_token_pool.device,
|
|
dtype=torch.uint64,
|
|
)
|
|
# TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start)
|
|
write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)](
|
|
req_to_token_pool.req_to_token,
|
|
req_pool_indices_tensor,
|
|
prefix_pointers,
|
|
prefix_lens_tensor,
|
|
seq_lens_tensor,
|
|
extend_lens_tensor,
|
|
out_cache_loc,
|
|
req_to_token_pool.req_to_token.shape[1],
|
|
)
|
|
else:
|
|
pt = 0
|
|
for i in range(req_pool_indices_cpu.shape[0]):
|
|
req_idx = req_pool_indices_cpu[i].item()
|
|
prefix_len = prefix_lens_cpu[i].item()
|
|
seq_len = seq_lens_cpu[i].item()
|
|
extend_len = extend_lens_cpu[i].item()
|
|
|
|
req_to_token_pool.write(
|
|
(req_idx, slice(0, prefix_len)),
|
|
prefix_tensors[i],
|
|
)
|
|
req_to_token_pool.write(
|
|
(req_idx, slice(prefix_len, seq_len)),
|
|
out_cache_loc[pt : pt + extend_len],
|
|
)
|
|
pt += extend_len
|
|
|
|
|
|
def get_last_loc(
|
|
req_to_token: torch.Tensor,
|
|
req_pool_indices_tensor: torch.Tensor,
|
|
prefix_lens_tensor: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
if (
|
|
get_global_server_args().attention_backend != "ascend"
|
|
and get_global_server_args().attention_backend != "torch_native"
|
|
):
|
|
impl = get_last_loc_triton
|
|
else:
|
|
impl = get_last_loc_torch
|
|
|
|
return impl(req_to_token, req_pool_indices_tensor, prefix_lens_tensor)
|
|
|
|
|
|
def get_last_loc_torch(
|
|
req_to_token: torch.Tensor,
|
|
req_pool_indices_tensor: torch.Tensor,
|
|
prefix_lens_tensor: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return torch.where(
|
|
prefix_lens_tensor > 0,
|
|
req_to_token[req_pool_indices_tensor, prefix_lens_tensor - 1],
|
|
torch.full_like(prefix_lens_tensor, -1),
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def get_last_loc_kernel(
|
|
req_to_token,
|
|
req_pool_indices_tensor,
|
|
prefix_lens_tensor,
|
|
result,
|
|
num_tokens,
|
|
req_to_token_stride,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(0)
|
|
offset = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE
|
|
mask = offset < num_tokens
|
|
|
|
prefix_lens = tl.load(prefix_lens_tensor + offset, mask=mask, other=0)
|
|
req_pool_indices = tl.load(req_pool_indices_tensor + offset, mask=mask, other=0)
|
|
|
|
token_mask = prefix_lens > 0
|
|
token_index = req_pool_indices * req_to_token_stride + (prefix_lens - 1)
|
|
tokens = tl.load(req_to_token + token_index, mask=token_mask, other=-1)
|
|
|
|
tl.store(result + offset, tokens, mask=mask)
|
|
|
|
|
|
def get_last_loc_triton(
|
|
req_to_token: torch.Tensor,
|
|
req_pool_indices_tensor: torch.Tensor,
|
|
prefix_lens_tensor: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
BLOCK_SIZE = 256
|
|
num_tokens = prefix_lens_tensor.shape[0]
|
|
result = torch.empty_like(prefix_lens_tensor)
|
|
grid = (triton.cdiv(num_tokens, BLOCK_SIZE),)
|
|
|
|
get_last_loc_kernel[grid](
|
|
req_to_token,
|
|
req_pool_indices_tensor,
|
|
prefix_lens_tensor,
|
|
result,
|
|
num_tokens,
|
|
req_to_token.stride(0),
|
|
BLOCK_SIZE,
|
|
)
|
|
return result
|
|
|
|
|
|
def alloc_token_slots(
|
|
tree_cache: BasePrefixCache,
|
|
num_tokens: int,
|
|
backup_state: bool = False,
|
|
):
|
|
allocator = tree_cache.token_to_kv_pool_allocator
|
|
evict_result = evict_from_tree_cache(tree_cache, num_tokens)
|
|
|
|
state = None
|
|
if backup_state:
|
|
state = allocator.backup_state()
|
|
|
|
out_cache_loc = allocator.alloc(num_tokens)
|
|
|
|
if out_cache_loc is None:
|
|
error_msg = (
|
|
f"Out of memory. Try to lower your batch size.\n"
|
|
f"Try to allocate {num_tokens} tokens.\n"
|
|
f"{available_and_evictable_str(tree_cache)}"
|
|
f"{_evict_result_str(evict_result)}\n"
|
|
)
|
|
logger.error(error_msg)
|
|
if tree_cache is not None:
|
|
tree_cache.pretty_print()
|
|
raise RuntimeError(error_msg)
|
|
|
|
return (out_cache_loc, state) if backup_state else out_cache_loc
|
|
|
|
|
|
def _evict_result_str(evict_result: EvictResult | None) -> str:
|
|
if evict_result is None:
|
|
return "evict_result=None"
|
|
return (
|
|
"evict_result=("
|
|
f"num_tokens_evicted={evict_result.num_tokens_evicted}, "
|
|
f"swa_num_tokens_evicted={evict_result.swa_num_tokens_evicted}, "
|
|
f"mamba_num_evicted={evict_result.mamba_num_evicted})"
|
|
)
|
|
|
|
|
|
def evict_from_tree_cache(
|
|
tree_cache: BasePrefixCache | None, num_tokens: int
|
|
) -> EvictResult:
|
|
if tree_cache is None:
|
|
return EvictResult()
|
|
|
|
if tree_cache.is_chunk_cache():
|
|
return EvictResult()
|
|
|
|
allocator = tree_cache.token_to_kv_pool_allocator
|
|
|
|
if isinstance(allocator, SWATokenToKVPoolAllocator):
|
|
# Hybrid allocator
|
|
full_available_size = allocator.full_available_size()
|
|
swa_available_size = allocator.swa_available_size()
|
|
|
|
if full_available_size < num_tokens or swa_available_size < num_tokens:
|
|
full_num_tokens = max(0, num_tokens - full_available_size)
|
|
swa_num_tokens = max(0, num_tokens - swa_available_size)
|
|
return tree_cache.evict(
|
|
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
|
|
)
|
|
return EvictResult()
|
|
else:
|
|
# Standard allocator
|
|
available = allocator.available_size()
|
|
if available < num_tokens:
|
|
logger.info(
|
|
"[MemCache-evict] evict_from_tree_cache: available=%d < num_tokens=%d deficit=%d, triggering eviction",
|
|
available,
|
|
num_tokens,
|
|
num_tokens - available,
|
|
)
|
|
return tree_cache.evict(EvictParams(num_tokens=num_tokens))
|
|
return EvictResult()
|
|
|
|
|
|
def _evict_for_compute_owner_lanes(
|
|
*,
|
|
tree_cache: BasePrefixCache | None,
|
|
allocator,
|
|
page_compute_owners: list[int],
|
|
) -> None:
|
|
if tree_cache is None or tree_cache.is_chunk_cache():
|
|
return
|
|
|
|
compute_owner_lane_stats = getattr(allocator, "compute_owner_lane_stats", None)
|
|
if compute_owner_lane_stats is None:
|
|
return
|
|
|
|
max_attempts = max(2, min(8, int(getattr(allocator, "cp_size", 1))))
|
|
for attempt in range(max_attempts):
|
|
_required, _available, deficits = compute_owner_lane_stats(page_compute_owners)
|
|
deficit_pages = sum(deficits)
|
|
if deficit_pages <= 0:
|
|
return
|
|
|
|
try:
|
|
evictable_size = tree_cache.evictable_size()
|
|
except Exception:
|
|
evictable_size = allocator.page_size
|
|
if isinstance(evictable_size, tuple):
|
|
evictable_size = evictable_size[0]
|
|
if evictable_size <= 0:
|
|
logger.info(
|
|
"[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size=%d <= 0, giving up",
|
|
evictable_size,
|
|
)
|
|
return
|
|
|
|
evict_tokens = max(
|
|
allocator.page_size,
|
|
deficit_pages * allocator.page_size * int(getattr(allocator, "cp_size", 1)),
|
|
)
|
|
before_available = allocator.available_size()
|
|
logger.info(
|
|
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d: deficit_pages=%d evict_tokens=%d before_available=%d evictable_size=%d",
|
|
attempt,
|
|
deficit_pages,
|
|
evict_tokens,
|
|
before_available,
|
|
evictable_size,
|
|
)
|
|
evict_result = tree_cache.evict(EvictParams(num_tokens=evict_tokens))
|
|
after_available = allocator.available_size()
|
|
evicted_tokens = getattr(evict_result, "num_tokens_evicted", 0)
|
|
logger.info(
|
|
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result: evicted=%d after_available=%d",
|
|
attempt,
|
|
evicted_tokens,
|
|
after_available,
|
|
)
|
|
if after_available <= before_available and evicted_tokens <= 0:
|
|
return
|
|
|
|
|
|
def alloc_paged_token_slots_extend(
|
|
tree_cache: BasePrefixCache,
|
|
prefix_lens: torch.Tensor,
|
|
prefix_lens_cpu: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
last_loc: torch.Tensor,
|
|
extend_num_tokens: int,
|
|
backup_state: bool = False,
|
|
):
|
|
# Over estimate the number of tokens: assume each request needs a new page.
|
|
allocator = tree_cache.token_to_kv_pool_allocator
|
|
num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
|
|
evict_result = evict_from_tree_cache(tree_cache, num_tokens)
|
|
|
|
# logger.info(
|
|
# "[MemCache-alloc] alloc_paged_token_slots_extend: extend_num_tokens=%d batch_size=%d num_tokens=%d page_size=%d "
|
|
# "available_size=%d evicted=%d",
|
|
# extend_num_tokens,
|
|
# len(seq_lens_cpu),
|
|
# num_tokens,
|
|
# allocator.page_size,
|
|
# allocator.available_size(),
|
|
# getattr(evict_result, "num_tokens_evicted", 0),
|
|
# )
|
|
|
|
alloc_extend_compute_owner = getattr(
|
|
allocator, "alloc_extend_compute_owner", None
|
|
)
|
|
page_compute_owners = None
|
|
compute_owner_unavailable_reason = None
|
|
if alloc_extend_compute_owner is not None and len(prefix_lens_cpu) == 1:
|
|
try:
|
|
server_args = get_global_server_args()
|
|
except ValueError:
|
|
server_args = None
|
|
if (
|
|
server_args is not None
|
|
and server_args.enable_nsa_prefill_cp_shared_kv
|
|
and server_args.enable_nsa_prefill_context_parallel
|
|
and server_args.nsa_prefill_cp_mode == "in-seq-split"
|
|
):
|
|
extend_len = int(seq_lens_cpu[0].item() - prefix_lens_cpu[0].item())
|
|
page_compute_owners = build_in_seq_page_compute_owners(
|
|
extend_len=extend_len,
|
|
extend_prefix_len=int(prefix_lens_cpu[0].item()),
|
|
page_size=int(allocator.page_size),
|
|
cp_size=int(allocator.cp_size),
|
|
)
|
|
if page_compute_owners is None:
|
|
compute_owner_unavailable_reason = (
|
|
get_in_seq_page_compute_owner_unavailable_reason(
|
|
extend_len=extend_len,
|
|
extend_prefix_len=int(prefix_lens_cpu[0].item()),
|
|
page_size=int(allocator.page_size),
|
|
cp_size=int(allocator.cp_size),
|
|
)
|
|
or "unknown"
|
|
)
|
|
else:
|
|
compute_owner_unavailable_reason = "server_args_not_enabled"
|
|
elif alloc_extend_compute_owner is not None:
|
|
compute_owner_unavailable_reason = (
|
|
"multi_batch" if len(prefix_lens_cpu) != 1 else "unknown"
|
|
)
|
|
|
|
state = None
|
|
if backup_state:
|
|
state = allocator.backup_state()
|
|
|
|
if page_compute_owners is not None:
|
|
out_cache_loc = alloc_extend_compute_owner(
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
extend_num_tokens,
|
|
page_compute_owners,
|
|
)
|
|
if out_cache_loc is None:
|
|
_evict_for_compute_owner_lanes(
|
|
tree_cache=tree_cache,
|
|
allocator=allocator,
|
|
page_compute_owners=page_compute_owners,
|
|
)
|
|
if backup_state:
|
|
state = allocator.backup_state()
|
|
out_cache_loc = alloc_extend_compute_owner(
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
extend_num_tokens,
|
|
page_compute_owners,
|
|
)
|
|
if out_cache_loc is None:
|
|
required = available = deficits = None
|
|
compute_owner_lane_stats = getattr(
|
|
allocator, "compute_owner_lane_stats", None
|
|
)
|
|
if compute_owner_lane_stats is not None:
|
|
required, available, deficits = compute_owner_lane_stats(
|
|
page_compute_owners
|
|
)
|
|
_log_cp_shared_kv_alloc_fallback(
|
|
"owner_lane_exhausted",
|
|
"failed to allocate pages from compute-owner lanes; "
|
|
"falling back to legacy page allocation. extend_num_tokens=%s page_size=%s "
|
|
"required_by_owner=%s available_by_owner=%s deficit_by_owner=%s",
|
|
extend_num_tokens,
|
|
allocator.page_size,
|
|
required,
|
|
available,
|
|
deficits,
|
|
)
|
|
out_cache_loc = allocator.alloc_extend(
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
extend_num_tokens,
|
|
)
|
|
else:
|
|
if alloc_extend_compute_owner is not None:
|
|
_log_cp_shared_kv_alloc_fallback(
|
|
compute_owner_unavailable_reason or "compute_owner_not_available",
|
|
"page-aligned compute-owner page assignment is unavailable; "
|
|
"falling back to legacy page allocation. batch_size=%s extend_num_tokens=%s page_size=%s reason=%s",
|
|
len(prefix_lens_cpu),
|
|
extend_num_tokens,
|
|
allocator.page_size,
|
|
compute_owner_unavailable_reason or "compute_owner_not_available",
|
|
)
|
|
out_cache_loc = allocator.alloc_extend(
|
|
prefix_lens,
|
|
prefix_lens_cpu,
|
|
seq_lens,
|
|
seq_lens_cpu,
|
|
last_loc,
|
|
extend_num_tokens,
|
|
)
|
|
|
|
if out_cache_loc is None:
|
|
error_msg = (
|
|
f"Prefill out of memory. Try to lower your batch size.\n"
|
|
f"Try to allocate {extend_num_tokens} tokens.\n"
|
|
f"{available_and_evictable_str(tree_cache)}"
|
|
f"{_evict_result_str(evict_result)}\n"
|
|
)
|
|
logger.error(error_msg)
|
|
if tree_cache is not None:
|
|
tree_cache.pretty_print()
|
|
raise RuntimeError(error_msg)
|
|
|
|
return (out_cache_loc, state) if backup_state else out_cache_loc
|
|
|
|
|
|
def alloc_req_slots(
|
|
req_to_token_pool: ReqToTokenPool,
|
|
reqs: list[Req],
|
|
tree_cache: BasePrefixCache | None,
|
|
) -> list[int]:
|
|
"""Allocate request slots from the pool."""
|
|
num_reqs = len(reqs)
|
|
if isinstance(req_to_token_pool, HybridReqToTokenPool):
|
|
mamba_available_size = req_to_token_pool.mamba_pool.available_size()
|
|
factor = (
|
|
MAMBA_STATE_PER_REQ_PREFIX_CACHE
|
|
if tree_cache.supports_mamba()
|
|
else MAMBA_STATE_PER_REQ_NO_CACHE
|
|
)
|
|
mamba_state_needed = num_reqs * factor
|
|
if mamba_available_size < mamba_state_needed:
|
|
if tree_cache is not None and tree_cache.supports_mamba():
|
|
mamba_num = max(0, mamba_state_needed - mamba_available_size)
|
|
tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
|
|
req_pool_indices = req_to_token_pool.alloc(reqs)
|
|
|
|
if req_pool_indices is None:
|
|
raise RuntimeError(
|
|
"alloc_req_slots runs out of memory. "
|
|
"Please set a smaller number for `--max-running-requests`. "
|
|
f"{req_to_token_pool.available_size()=}, "
|
|
f"{num_reqs=}, "
|
|
)
|
|
return req_pool_indices
|
|
|
|
|
|
def alloc_for_extend(
|
|
batch: ScheduleBatch,
|
|
) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
|
|
"""
|
|
Allocate KV cache for extend batch and write to req_to_token_pool.
|
|
|
|
Returns:
|
|
out_cache_loc: allocated cache locations
|
|
req_pool_indices_device: request pool indices at a device tensor
|
|
req_pool_indices: request pool indices as list
|
|
"""
|
|
# free out-of-window swa tokens
|
|
batch.maybe_evict_swa()
|
|
|
|
prefix_tensors = [r.prefix_indices for r in batch.reqs]
|
|
|
|
# Create tensors for allocation
|
|
prefix_lens_cpu = torch.tensor(batch.prefix_lens, dtype=torch.int64)
|
|
extend_lens_cpu = torch.tensor(batch.extend_lens, dtype=torch.int64)
|
|
prefix_lens_device = prefix_lens_cpu.to(batch.device, non_blocking=True)
|
|
extend_lens_device = extend_lens_cpu.to(batch.device, non_blocking=True)
|
|
|
|
# Allocate req slots
|
|
req_pool_indices = alloc_req_slots(
|
|
batch.req_to_token_pool, batch.reqs, batch.tree_cache
|
|
)
|
|
req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64)
|
|
req_pool_indices_device = req_pool_indices_cpu.to(batch.device, non_blocking=True)
|
|
|
|
# Allocate KV cache (throws exception on failure)
|
|
if batch.tree_cache.page_size == 1:
|
|
out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens)
|
|
else:
|
|
# Paged allocation - build last_loc
|
|
last_loc = [
|
|
(t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device))
|
|
for t in prefix_tensors
|
|
]
|
|
out_cache_loc = alloc_paged_token_slots_extend(
|
|
tree_cache=batch.tree_cache,
|
|
prefix_lens=prefix_lens_device,
|
|
prefix_lens_cpu=prefix_lens_cpu,
|
|
seq_lens=batch.seq_lens,
|
|
seq_lens_cpu=batch.seq_lens_cpu,
|
|
last_loc=torch.cat(last_loc),
|
|
extend_num_tokens=batch.extend_num_tokens,
|
|
)
|
|
|
|
# Write to req_to_token_pool
|
|
write_cache_indices(
|
|
out_cache_loc,
|
|
req_pool_indices_device,
|
|
req_pool_indices_cpu,
|
|
prefix_lens_device,
|
|
prefix_lens_cpu,
|
|
batch.seq_lens,
|
|
batch.seq_lens_cpu,
|
|
extend_lens_device,
|
|
extend_lens_cpu,
|
|
prefix_tensors,
|
|
batch.req_to_token_pool,
|
|
)
|
|
|
|
return out_cache_loc, req_pool_indices_device, req_pool_indices
|
|
|
|
|
|
def alloc_paged_token_slots_decode(
|
|
tree_cache: BasePrefixCache,
|
|
seq_lens: torch.Tensor,
|
|
seq_lens_cpu: torch.Tensor,
|
|
last_loc: torch.Tensor,
|
|
token_per_req: int = 1,
|
|
) -> torch.Tensor:
|
|
"""Allocate paged KV cache for decode batch."""
|
|
allocator = tree_cache.token_to_kv_pool_allocator
|
|
# Over estimate the number of tokens: assume each request needs a new page.
|
|
num_tokens = len(seq_lens) * allocator.page_size
|
|
evict_from_tree_cache(tree_cache, num_tokens)
|
|
|
|
out_cache_loc = allocator.alloc_decode(seq_lens, seq_lens_cpu, last_loc)
|
|
|
|
if out_cache_loc is None:
|
|
error_msg = (
|
|
f"Decode out of memory. Try to lower your batch size.\n"
|
|
f"Try to allocate {len(seq_lens) * token_per_req} tokens.\n"
|
|
f"{available_and_evictable_str(tree_cache)}"
|
|
)
|
|
logger.error(error_msg)
|
|
if tree_cache is not None:
|
|
tree_cache.pretty_print()
|
|
raise RuntimeError(error_msg)
|
|
|
|
return out_cache_loc
|
|
|
|
|
|
def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
|
|
"""
|
|
Allocate KV cache for decode batch and write to req_to_token_pool.
|
|
|
|
Returns:
|
|
out_cache_loc: allocated cache locations
|
|
"""
|
|
|
|
batch.maybe_evict_swa()
|
|
|
|
bs = batch.seq_lens.shape[0]
|
|
|
|
if batch.tree_cache.page_size == 1:
|
|
# Non-paged allocation
|
|
out_cache_loc = alloc_token_slots(batch.tree_cache, bs * token_per_req)
|
|
else:
|
|
# Paged allocation
|
|
last_loc = batch.req_to_token_pool.req_to_token[
|
|
batch.req_pool_indices, batch.seq_lens - 1
|
|
]
|
|
seq_lens_next = batch.seq_lens + token_per_req
|
|
out_cache_loc = alloc_paged_token_slots_decode(
|
|
tree_cache=batch.tree_cache,
|
|
seq_lens=seq_lens_next,
|
|
seq_lens_cpu=batch.seq_lens_cpu + token_per_req,
|
|
last_loc=last_loc,
|
|
token_per_req=token_per_req,
|
|
)
|
|
|
|
# Write to req_to_token_pool
|
|
if batch.model_config.is_encoder_decoder:
|
|
locs = batch.encoder_lens + batch.seq_lens
|
|
else:
|
|
locs = batch.seq_lens.clone()
|
|
|
|
batch.req_to_token_pool.write(
|
|
(batch.req_pool_indices, locs), out_cache_loc.to(torch.int32)
|
|
)
|
|
|
|
return out_cache_loc
|
|
|
|
|
|
def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
|
|
# MambaRadixCache may alloc mamba state before alloc KV cache
|
|
if req.req_pool_idx is None:
|
|
assert (
|
|
tree_cache.supports_mamba()
|
|
), "Only MambaRadixCache allow freeing before alloc"
|
|
# TODO (csy, hanming): clean up this early allocation logic
|
|
if req.mamba_pool_idx is not None:
|
|
tree_cache.req_to_token_pool.mamba_pool.free(
|
|
req.mamba_pool_idx.unsqueeze(-1)
|
|
)
|
|
req.mamba_pool_idx = None
|
|
return
|
|
|
|
tree_cache.cache_finished_req(req, is_insert=is_insert)
|
|
|
|
# FIXME: SessionAwareCache.cache_finished_req sets req_pool_idx = None to
|
|
# transfer KV ownership to the SessionSlot, so we skip the remaining
|
|
# cleanup (overalloc free + pool slot free). This means over-allocated
|
|
# tokens from speculative decoding are NOT freed between turns.
|
|
if req.req_pool_idx is None:
|
|
return
|
|
|
|
start_p, end_p = req.pop_overallocated_kv_cache()
|
|
|
|
global_server_args = get_global_server_args()
|
|
page_size = global_server_args.page_size
|
|
spec_algo = global_server_args.speculative_algorithm
|
|
|
|
if spec_algo is None:
|
|
assert (
|
|
start_p == end_p
|
|
), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}"
|
|
|
|
if page_size > 1:
|
|
start_p = ceil_align(start_p, page_size)
|
|
|
|
if start_p < end_p:
|
|
indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][
|
|
start_p:end_p
|
|
]
|
|
tree_cache.token_to_kv_pool_allocator.free(indices_to_free)
|
|
# If the prefix cache doesn't manage mamba states, we must free them here.
|
|
if isinstance(tree_cache.req_to_token_pool, HybridReqToTokenPool) and (
|
|
not tree_cache.supports_mamba()
|
|
):
|
|
assert (
|
|
req.mamba_pool_idx is not None
|
|
), "mamba state is freed while the tree cache does not manage mamba states"
|
|
tree_cache.req_to_token_pool.free_mamba_cache(req)
|
|
tree_cache.req_to_token_pool.free(req)
|
|
|
|
|
|
def available_and_evictable_str(tree_cache: BasePrefixCache) -> str:
|
|
return tree_cache.available_and_evictable_str()
|