chore: import upstream snapshot with attribution
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,647 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
BasePrefixCache,
|
||||
DecLockRefParams,
|
||||
DecLockRefResult,
|
||||
EvictParams,
|
||||
EvictResult,
|
||||
IncLockRefResult,
|
||||
InitLoadBackParams,
|
||||
MatchPrefixParams,
|
||||
MatchResult,
|
||||
)
|
||||
from sglang.srt.utils.common import ceil_align
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class _VirtualNode:
|
||||
"""Sentinel node for streaming session requests.
|
||||
|
||||
Passed to inc_lock_ref / dec_lock_ref so the cache can distinguish
|
||||
streaming-session locks (no-op) from real radix-tree locks (forwarded).
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionSlot:
|
||||
"""Holds KV state between streaming session turns."""
|
||||
|
||||
virtual_node: _VirtualNode = field(default_factory=_VirtualNode)
|
||||
|
||||
# KV pool state (None means no KV is currently held by this slot)
|
||||
req_pool_idx: Optional[int] = None
|
||||
kv_committed_len: int = 0
|
||||
kv_allocated_len: int = 0
|
||||
|
||||
# First req's radix tree node (for dec_lock_ref on session close)
|
||||
last_node: Any = None
|
||||
cache_protected_len: int = 0
|
||||
swa_uuid_for_lock: Optional[str] = None
|
||||
|
||||
# SWA state
|
||||
swa_evicted_seqlen: int = 0
|
||||
|
||||
# Mamba states
|
||||
mamba_pool_idx: Any = None
|
||||
mamba_ping_pong_track_buffer: Any = None
|
||||
mamba_next_track_idx: Any = None
|
||||
mamba_last_track_seqlen: Any = None
|
||||
mamba_branching_seqlen: Any = None
|
||||
|
||||
@property
|
||||
def is_holding_kv(self) -> bool:
|
||||
"""Whether this slot currently holds KV pool resources."""
|
||||
return self.req_pool_idx is not None
|
||||
|
||||
def save_from_req(self, req: Req, is_first: bool):
|
||||
"""Save KV state from a finishing request into this slot."""
|
||||
self.req_pool_idx = req.req_pool_idx
|
||||
self.kv_committed_len = req.kv_committed_len
|
||||
self.kv_allocated_len = req.kv_allocated_len
|
||||
self.swa_evicted_seqlen = req.swa_evicted_seqlen
|
||||
|
||||
if is_first:
|
||||
self.last_node = req.last_node
|
||||
self.cache_protected_len = req.cache_protected_len
|
||||
self.swa_uuid_for_lock = req.swa_uuid_for_lock
|
||||
|
||||
self.mamba_pool_idx = req.mamba_pool_idx
|
||||
self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer
|
||||
self.mamba_next_track_idx = req.mamba_next_track_idx
|
||||
self.mamba_last_track_seqlen = req.mamba_last_track_seqlen
|
||||
self.mamba_branching_seqlen = req.mamba_branching_seqlen
|
||||
|
||||
# Ownership has transferred to the slot. Null *all* of the req's
|
||||
# references so any later alloc()/free path that inspects the req
|
||||
# (e.g. the alloc-skip check on `req.mamba_ping_pong_track_buffer
|
||||
# is None`, or the retract cleanup) sees no dangling pointers
|
||||
# into slot-owned tensors. Without this the alloc path can decide
|
||||
# the req still has a ping-pong buffer and skip alloc, causing
|
||||
# the slot's tensor to be reused by a new req and leaked when
|
||||
# the slot is later freed.
|
||||
req.req_pool_idx = None
|
||||
req.mamba_pool_idx = None
|
||||
req.mamba_ping_pong_track_buffer = None
|
||||
req.mamba_next_track_idx = None
|
||||
req.mamba_last_track_seqlen = None
|
||||
req.mamba_branching_seqlen = None
|
||||
|
||||
def restore_to_req(self, req: Req):
|
||||
"""Restore KV state from this slot into an incoming request."""
|
||||
req.req_pool_idx = self.req_pool_idx
|
||||
req.kv_committed_len = self.kv_committed_len
|
||||
req.kv_allocated_len = self.kv_allocated_len
|
||||
req.swa_evicted_seqlen = self.swa_evicted_seqlen
|
||||
req.swa_uuid_for_lock = self.swa_uuid_for_lock
|
||||
|
||||
req.mamba_pool_idx = self.mamba_pool_idx
|
||||
req.mamba_ping_pong_track_buffer = self.mamba_ping_pong_track_buffer
|
||||
req.mamba_next_track_idx = self.mamba_next_track_idx
|
||||
req.mamba_last_track_seqlen = self.mamba_last_track_seqlen
|
||||
req.mamba_branching_seqlen = self.mamba_branching_seqlen
|
||||
|
||||
# NOTE: req_pool_idx and mamba_pool_idx are intentionally NOT cleared
|
||||
# from the slot. During chunked prefill, a request may be rejected by
|
||||
# the scheduler (e.g. budget exhausted) and retried in the next cycle.
|
||||
# Each retry calls match_prefix -> restore_to_req again, so the slot
|
||||
# must remain intact for idempotent restoration.
|
||||
|
||||
|
||||
def _is_streaming(req: Optional[Req]) -> bool:
|
||||
return req is not None and req.session is not None and req.session.streaming
|
||||
|
||||
|
||||
class StreamingSession(BasePrefixCache):
|
||||
"""Adds streaming-session KV save/restore on top of any BasePrefixCache.
|
||||
|
||||
Works both as an external wrapper (``StreamingSession(RadixCache(...))``)
|
||||
and in embedded composition (``StreamingSession(inner=self)``). For the
|
||||
embedded case, the composing cache must pre-check dispatch conditions
|
||||
(``_is_streaming`` / ``find_active_slot`` / ``has_slot``) so the internal
|
||||
fall-through to ``self.inner.xxx`` never fires -- otherwise it recurses.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: BasePrefixCache):
|
||||
self.inner = inner
|
||||
self.slots: Dict[str, SessionSlot] = {}
|
||||
|
||||
# -- Forward PrefixCacheTrait properties to inner cache --
|
||||
|
||||
@property
|
||||
def req_to_token_pool(self):
|
||||
return self.inner.req_to_token_pool
|
||||
|
||||
@req_to_token_pool.setter
|
||||
def req_to_token_pool(self, value):
|
||||
self.inner.req_to_token_pool = value
|
||||
|
||||
@property
|
||||
def token_to_kv_pool_allocator(self):
|
||||
return self.inner.token_to_kv_pool_allocator
|
||||
|
||||
@token_to_kv_pool_allocator.setter
|
||||
def token_to_kv_pool_allocator(self, value):
|
||||
self.inner.token_to_kv_pool_allocator = value
|
||||
|
||||
@property
|
||||
def page_size(self):
|
||||
return self.inner.page_size
|
||||
|
||||
@page_size.setter
|
||||
def page_size(self, value):
|
||||
self.inner.page_size = value
|
||||
|
||||
@property
|
||||
def disable(self):
|
||||
return self.inner.disable
|
||||
|
||||
@disable.setter
|
||||
def disable(self, value):
|
||||
self.inner.disable = value
|
||||
|
||||
@property
|
||||
def metrics_collector(self):
|
||||
return self.inner.metrics_collector
|
||||
|
||||
@metrics_collector.setter
|
||||
def metrics_collector(self, value):
|
||||
self.inner.metrics_collector = value
|
||||
|
||||
# -- Condition helpers (used by embedded-mode callers for pre-dispatch) --
|
||||
|
||||
def has_slot(self, session_id: str) -> bool:
|
||||
return session_id in self.slots
|
||||
|
||||
def any_holding_kv(self) -> bool:
|
||||
return any(s.is_holding_kv for s in self.slots.values())
|
||||
|
||||
# -- Try-handle entries for composition (see class docstring) --
|
||||
|
||||
def try_inc_lock_ref(self, node: Any) -> Optional[IncLockRefResult]:
|
||||
"""No-op lock if ``node`` is a session-internal sentinel; returns
|
||||
None to tell the caller to run its raw tree lock path."""
|
||||
if isinstance(node, _VirtualNode):
|
||||
return IncLockRefResult()
|
||||
return None
|
||||
|
||||
def try_dec_lock_ref(
|
||||
self, node: Any, params: Optional[DecLockRefParams] = None
|
||||
) -> Optional[DecLockRefResult]:
|
||||
if isinstance(node, _VirtualNode):
|
||||
return DecLockRefResult()
|
||||
return None
|
||||
|
||||
def find_active_slot(self, req: Req) -> Optional[SessionSlot]:
|
||||
"""Returns an active slot for this req, or None.
|
||||
|
||||
Side effect: if req is pre-aborted (to_finish set, e.g. input too
|
||||
long), detach it from the session so cache_finished_req treats it
|
||||
as a normal req. The slot stays intact for the next request.
|
||||
"""
|
||||
if not _is_streaming(req):
|
||||
return None
|
||||
slot = self.slots.get(req.session.session_id)
|
||||
if slot is None or slot.req_pool_idx is None:
|
||||
return None
|
||||
if req.to_finish is not None:
|
||||
req.session.abort_req()
|
||||
req.session = None
|
||||
return None
|
||||
return slot
|
||||
|
||||
# -- BasePrefixCache abstract methods --
|
||||
|
||||
def reset(self):
|
||||
self.slots.clear()
|
||||
self.inner.reset()
|
||||
|
||||
# -- Streaming entries: contract with embedded composers (e.g.
|
||||
# UnifiedRadixCache) is a uniform "try_handle_*" pattern. Each method
|
||||
# executes the streaming body if applicable and signals whether the
|
||||
# caller still needs to run its raw path.
|
||||
|
||||
def try_match_prefix(self, params: MatchPrefixParams) -> Optional[MatchResult]:
|
||||
"""Returns a MatchResult iff the request hits an active session slot;
|
||||
otherwise None (caller falls back to its raw match)."""
|
||||
slot = self.find_active_slot(params.req)
|
||||
if slot is None:
|
||||
return None
|
||||
|
||||
req = params.req
|
||||
slot.restore_to_req(req)
|
||||
|
||||
# token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve
|
||||
# already applied). min handles retract retry where committed_len
|
||||
# can exceed len(token_ids) by 1.
|
||||
prefix_len = min(req.kv_committed_len, len(params.key))
|
||||
|
||||
# Streaming sessions are append-only (session_controller rollback
|
||||
# ensures req_nodes always points to the last successful req).
|
||||
assert prefix_len >= slot.cache_protected_len, (
|
||||
f"streaming session prefix shrank: {prefix_len=} < "
|
||||
f"{slot.cache_protected_len=}"
|
||||
)
|
||||
|
||||
# Free orphaned tail: alloc_for_extend will overwrite
|
||||
# req_to_token[prefix_len:] with new indices. The range
|
||||
# [prefix_len, kv_allocated_len) has stale indices from the
|
||||
# previous turn's decode (e.g. alloc-commit gap on retract,
|
||||
# or speculative draft tokens).
|
||||
self._free_tail(slot, req, prefix_len)
|
||||
|
||||
device_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, :prefix_len
|
||||
].to(dtype=torch.int64)
|
||||
|
||||
return MatchResult(
|
||||
device_indices=device_indices,
|
||||
last_device_node=slot.virtual_node,
|
||||
last_host_node=slot.virtual_node,
|
||||
best_match_node=slot.virtual_node,
|
||||
cache_protected_len=slot.cache_protected_len,
|
||||
)
|
||||
|
||||
def try_cache_finished_req(
|
||||
self, req: Req, is_insert: bool = True, **kwargs
|
||||
) -> bool:
|
||||
"""Handles a streaming-session finish (save slot / mid-abort nuke).
|
||||
Returns True if handled; False means caller runs its raw path."""
|
||||
if not _is_streaming(req):
|
||||
return False
|
||||
|
||||
from sglang.srt.managers.schedule_batch import FINISH_ABORT
|
||||
|
||||
session_id = req.session.session_id
|
||||
slot = self.slots.get(session_id)
|
||||
is_first = slot is None
|
||||
|
||||
# Mid-processing abort only. Pre-aborted reqs have session=None
|
||||
# (set in find_active_slot) and never reach here.
|
||||
# Nuke all KV via release_session, delete slot. Token IDs stay
|
||||
# in req_nodes (finish_req was never called -> last successful
|
||||
# req). Next request re-prefills from scratch.
|
||||
if isinstance(req.finished_reason, FINISH_ABORT):
|
||||
if slot is None:
|
||||
# First-request mid-processing abort: create ephemeral
|
||||
# slot from req state so release_session handles cleanup.
|
||||
# Include last_node/cache_protected_len from the req so
|
||||
# release_session calls dec_lock_ref on the tree lock.
|
||||
# Also carry the mamba refs over so _free_slot_mamba can
|
||||
# return the (possibly extra_buffer ping-pong) slots to
|
||||
# the mamba pool; otherwise the abort orphans them.
|
||||
slot = SessionSlot(
|
||||
req_pool_idx=req.req_pool_idx,
|
||||
kv_allocated_len=req.kv_allocated_len,
|
||||
last_node=req.last_node,
|
||||
cache_protected_len=req.cache_protected_len,
|
||||
swa_uuid_for_lock=req.swa_uuid_for_lock,
|
||||
mamba_pool_idx=req.mamba_pool_idx,
|
||||
mamba_ping_pong_track_buffer=req.mamba_ping_pong_track_buffer,
|
||||
)
|
||||
self.slots[session_id] = slot
|
||||
# Slot now owns the mamba state — drop the req's refs so
|
||||
# the abort fall-through doesn't double-free.
|
||||
req.mamba_pool_idx = None
|
||||
req.mamba_ping_pong_track_buffer = None
|
||||
slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len)
|
||||
self.release_session(session_id)
|
||||
req.req_pool_idx = None
|
||||
req.session.abort_req()
|
||||
self._mark_kv_freed(req)
|
||||
return True
|
||||
|
||||
if is_first:
|
||||
slot = SessionSlot()
|
||||
self.slots[session_id] = slot
|
||||
|
||||
finished_len = (
|
||||
req.finished_len if req.finished_len is not None else len(req.output_ids)
|
||||
)
|
||||
target = len(req.origin_input_ids) + finished_len
|
||||
self._trim_overshoot(req, finished_len)
|
||||
|
||||
slot.save_from_req(req, is_first=is_first)
|
||||
# Inherit the authoritative finished length on the slot, not the lagging
|
||||
# req clock (under overlap + honest committed the clock lags the in-flight
|
||||
# verify by ~1, which would short-change inheritance). Clamp to allocated
|
||||
# to keep committed <= allocated for prepare_for_decode.
|
||||
slot.kv_committed_len = min(target, slot.kv_allocated_len)
|
||||
|
||||
# Update req_nodes to this successfully finished request.
|
||||
req.session.finish_req(req)
|
||||
|
||||
self._mark_kv_freed(req)
|
||||
return True
|
||||
|
||||
def try_cache_unfinished_req(
|
||||
self, req: Req, chunked: bool = False, **kwargs
|
||||
) -> bool:
|
||||
"""Handles a streaming-session mid-flight cache op:
|
||||
- chunked prefill: snapshot current KV as prefix, skip radix
|
||||
- subsequent turn: skip radix (slot already holds KV)
|
||||
Returns False for first-turn non-chunked (caller must run raw radix
|
||||
insert to set up the initial tree lock)."""
|
||||
if not _is_streaming(req):
|
||||
return False
|
||||
if chunked:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
req.req_pool_idx, : req.extend_range.end
|
||||
]
|
||||
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
|
||||
return True
|
||||
if req.session.session_id in self.slots:
|
||||
return True
|
||||
return False
|
||||
|
||||
# -- BasePrefixCache abstract methods: thin adapters over try_handle_* --
|
||||
|
||||
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
||||
result = self.try_match_prefix(params)
|
||||
if result is not None:
|
||||
return result
|
||||
return self.inner.match_prefix(params)
|
||||
|
||||
def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
|
||||
if self.try_cache_finished_req(req, is_insert=is_insert, **kwargs):
|
||||
return
|
||||
self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs)
|
||||
|
||||
def cache_unfinished_req(self, req: Req, **kwargs):
|
||||
if self.try_cache_unfinished_req(req, **kwargs):
|
||||
return
|
||||
self.inner.cache_unfinished_req(req, **kwargs)
|
||||
|
||||
def evict(self, params: EvictParams) -> EvictResult:
|
||||
return self.inner.evict(params)
|
||||
|
||||
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
|
||||
result = self.try_inc_lock_ref(node)
|
||||
if result is not None:
|
||||
return result
|
||||
return self.inner.inc_lock_ref(node)
|
||||
|
||||
def dec_lock_ref(
|
||||
self, node: Any, params: Optional[DecLockRefParams] = None
|
||||
) -> DecLockRefResult:
|
||||
result = self.try_dec_lock_ref(node, params)
|
||||
if result is not None:
|
||||
return result
|
||||
return self.inner.dec_lock_ref(node, params)
|
||||
|
||||
# -- Session lifecycle --
|
||||
|
||||
def release_session(self, session_id: str) -> None:
|
||||
slot = self.slots.pop(session_id, None)
|
||||
if slot is None:
|
||||
return
|
||||
protected_len = slot.cache_protected_len
|
||||
lock_node = slot.last_node
|
||||
tokens_freed = (
|
||||
max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0
|
||||
)
|
||||
logger.info(
|
||||
"Session KV released: %s (%d tokens freed)", session_id, tokens_freed
|
||||
)
|
||||
|
||||
if lock_node is not None:
|
||||
if slot.swa_uuid_for_lock is not None:
|
||||
self.inner.dec_lock_ref(
|
||||
lock_node,
|
||||
DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock),
|
||||
)
|
||||
else:
|
||||
self.inner.dec_lock_ref(lock_node)
|
||||
|
||||
if slot.is_holding_kv:
|
||||
start = protected_len
|
||||
end = slot.kv_allocated_len
|
||||
if start < end:
|
||||
kv_indices = self.req_to_token_pool.req_to_token[
|
||||
slot.req_pool_idx, start:end
|
||||
]
|
||||
self.token_to_kv_pool_allocator.free(kv_indices)
|
||||
self.req_to_token_pool.free_slots.append(slot.req_pool_idx)
|
||||
|
||||
self._free_slot_mamba(slot)
|
||||
|
||||
def release_radix_session(self, session_id: str) -> None:
|
||||
self.inner.release_radix_session(session_id)
|
||||
|
||||
def session_held_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
|
||||
"""Total KV tokens held by session slots, not tracked by the tree.
|
||||
|
||||
Excludes slots whose KV is currently owned by an owning request --
|
||||
those tokens are counted via uncached_size in the busy mem check.
|
||||
A slot's pool_idx being in active_pool_idxs indicates a req owns it.
|
||||
"""
|
||||
total = 0
|
||||
for slot in self.slots.values():
|
||||
in_batch = (
|
||||
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
|
||||
)
|
||||
if slot.is_holding_kv and not in_batch:
|
||||
allocated = ceil_align(slot.kv_allocated_len, self.page_size)
|
||||
total += allocated - slot.cache_protected_len
|
||||
return total
|
||||
|
||||
def session_held_full_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
|
||||
"""An alias to align the naming style of SWA"""
|
||||
return self.session_held_tokens(active_pool_idxs)
|
||||
|
||||
def session_held_swa_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
|
||||
"""Total SWA tokens held by session slots, not tracked by the tree."""
|
||||
total = 0
|
||||
for slot in self.slots.values():
|
||||
in_batch = (
|
||||
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
|
||||
)
|
||||
if slot.is_holding_kv and not in_batch:
|
||||
allocated = ceil_align(slot.kv_allocated_len, self.page_size)
|
||||
total += allocated - max(
|
||||
slot.cache_protected_len, slot.swa_evicted_seqlen
|
||||
)
|
||||
return total
|
||||
|
||||
def session_held_req_count(self, active_pool_idxs: Optional[set] = None) -> int:
|
||||
"""Number of req pool slots held by session slots."""
|
||||
|
||||
def _owned(s):
|
||||
in_batch = (
|
||||
active_pool_idxs is not None and s.req_pool_idx in active_pool_idxs
|
||||
)
|
||||
return s.is_holding_kv and not in_batch
|
||||
|
||||
return sum(_owned(s) for s in self.slots.values())
|
||||
|
||||
def session_held_mamba_slots(self, active_pool_idxs: Optional[set] = None) -> int:
|
||||
"""Total mamba_pool entries held by session slots (mamba_pool_idx +
|
||||
mamba_ping_pong_track_buffer). Excludes slots whose owning req is
|
||||
currently in the batch -- those slots are counted via the normal
|
||||
alloc/free paths (same convention as the sibling ``session_held_*``
|
||||
accessors).
|
||||
"""
|
||||
total = 0
|
||||
for slot in self.slots.values():
|
||||
in_batch = (
|
||||
active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs
|
||||
)
|
||||
if in_batch:
|
||||
continue
|
||||
if slot.mamba_pool_idx is not None:
|
||||
total += slot.mamba_pool_idx.numel()
|
||||
if slot.mamba_ping_pong_track_buffer is not None:
|
||||
total += slot.mamba_ping_pong_track_buffer.numel()
|
||||
return total
|
||||
|
||||
def _free_slot_mamba(self, slot: SessionSlot) -> None:
|
||||
"""Return a session slot's mamba pool state to the allocator."""
|
||||
mamba_allocator = getattr(self.req_to_token_pool, "mamba_allocator", None)
|
||||
if mamba_allocator is None:
|
||||
return
|
||||
if slot.mamba_pool_idx is not None:
|
||||
mamba_allocator.free(slot.mamba_pool_idx.unsqueeze(0))
|
||||
slot.mamba_pool_idx = None
|
||||
if slot.mamba_ping_pong_track_buffer is not None:
|
||||
mamba_allocator.free(slot.mamba_ping_pong_track_buffer)
|
||||
slot.mamba_ping_pong_track_buffer = None
|
||||
|
||||
# -- Internal helpers (streaming body bits) --
|
||||
|
||||
def _free_tail(self, slot: SessionSlot, req: Req, prefix_len: int) -> None:
|
||||
"""match_prefix path: free orphaned KV in [prefix_len, kv_allocated_len)
|
||||
before alloc_for_extend overwrites it. The gap appears when spec
|
||||
decoding pushes allocated above committed, or when retract retry's
|
||||
logit-reserve pulls prefix_len below committed.
|
||||
"""
|
||||
self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len)
|
||||
slot.kv_allocated_len = prefix_len
|
||||
slot.kv_committed_len = min(slot.kv_committed_len, prefix_len)
|
||||
slot.swa_evicted_seqlen = min(slot.swa_evicted_seqlen, prefix_len)
|
||||
req.kv_allocated_len = prefix_len
|
||||
req.kv_committed_len = min(req.kv_committed_len, prefix_len)
|
||||
req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, prefix_len)
|
||||
|
||||
def _trim_overshoot(self, req: Req, finished_len: int) -> None:
|
||||
"""Trim slot KV to finished_len boundary. Spec v2 may overshoot
|
||||
max_new_tokens (verify round commits M+1 at a time); next turn's
|
||||
input is output_ids[:finished_len], so positions past that must
|
||||
be released to avoid token/KV mismatch.
|
||||
"""
|
||||
target = len(req.origin_input_ids) + finished_len
|
||||
self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len)
|
||||
req.kv_allocated_len = min(req.kv_allocated_len, target)
|
||||
req.kv_committed_len = min(req.kv_committed_len, target)
|
||||
req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, target)
|
||||
req.output_ids = req.output_ids[:finished_len]
|
||||
|
||||
def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None:
|
||||
"""Free req_to_token[pool_idx, ceil_align(target):end). Page-aligned
|
||||
because PagedTokenToKVPoolAllocator.free returns whole pages
|
||||
(free_index // page_size), so partial-page free would corrupt pages
|
||||
still holding committed tokens. The range [target, ceil_align(target))
|
||||
stays attached until release_session frees the whole page.
|
||||
"""
|
||||
if end <= target:
|
||||
return
|
||||
start = target
|
||||
if self.page_size > 1:
|
||||
start = ceil_align(start, self.page_size)
|
||||
if start < end:
|
||||
tail = self.req_to_token_pool.req_to_token[pool_idx, start:end]
|
||||
self.token_to_kv_pool_allocator.free(tail)
|
||||
|
||||
@staticmethod
|
||||
def _mark_kv_freed(req: Req) -> None:
|
||||
"""Set bookkeeping flags so busy check skips this finished req."""
|
||||
if not req.kv_committed_freed:
|
||||
req.pop_committed_kv_cache()
|
||||
if not req.kv_overallocated_freed:
|
||||
req.pop_overallocated_kv_cache()
|
||||
|
||||
# -- Pass-through methods --
|
||||
|
||||
def evictable_size(self):
|
||||
return self.inner.evictable_size()
|
||||
|
||||
def full_evictable_size(self):
|
||||
return self.inner.full_evictable_size()
|
||||
|
||||
def swa_evictable_size(self):
|
||||
return self.inner.swa_evictable_size()
|
||||
|
||||
def protected_size(self):
|
||||
return self.inner.protected_size()
|
||||
|
||||
def full_protected_size(self):
|
||||
return self.inner.full_protected_size()
|
||||
|
||||
def swa_protected_size(self):
|
||||
return self.inner.swa_protected_size()
|
||||
|
||||
def total_size(self):
|
||||
return self.inner.total_size()
|
||||
|
||||
def pretty_print(self):
|
||||
return self.inner.pretty_print()
|
||||
|
||||
def init_load_back(self, params: InitLoadBackParams):
|
||||
return self.inner.init_load_back(params)
|
||||
|
||||
def ready_to_load_host_cache(self):
|
||||
return self.inner.ready_to_load_host_cache()
|
||||
|
||||
def flush_write_through_acks(self) -> None:
|
||||
return self.inner.flush_write_through_acks()
|
||||
|
||||
def check_hicache_events(self):
|
||||
return self.inner.check_hicache_events()
|
||||
|
||||
def take_events(self):
|
||||
return self.inner.take_events()
|
||||
|
||||
def supports_swa(self):
|
||||
return self.inner.supports_swa()
|
||||
|
||||
def supports_mamba(self):
|
||||
return self.inner.supports_mamba()
|
||||
|
||||
def supports_streaming_session(self) -> bool:
|
||||
return True
|
||||
|
||||
def is_chunk_cache(self):
|
||||
return self.inner.is_chunk_cache()
|
||||
|
||||
def is_tree_cache(self):
|
||||
return self.inner.is_tree_cache()
|
||||
|
||||
def available_and_evictable_str(self):
|
||||
return self.inner.available_and_evictable_str()
|
||||
|
||||
def init_metrics_collector(self):
|
||||
return self.inner.init_metrics_collector()
|
||||
|
||||
def sanity_check(self):
|
||||
# Skip inner sanity check when sessions hold tree locks, because
|
||||
# the check asserts all nodes are unlocked during idle.
|
||||
if self.any_holding_kv():
|
||||
return
|
||||
self.inner.sanity_check()
|
||||
|
||||
# Forward attribute access for cache-specific methods (e.g.
|
||||
# sliding_window_size, all_values_flatten, etc.)
|
||||
def __getattr__(self, name):
|
||||
return getattr(self.inner, name)
|
||||
Reference in New Issue
Block a user