[SGLang] Session-Aware Cache: 사용자별 KV 캐시 파티셔닝
들어가며
멀티턴 대화에서 매 턴마다 이전 대화의 KV Cache를 다시 계산하는 것은 낭비다. 하지만 일반적인 Radix Cache는 요청이 끝나면 KV 리소스를 반환하고, 다음 턴에서 prefix 매칭으로 재활용을 시도한다. 동시 요청이 많으면 그 사이에 evict되어 재활용에 실패할 수 있다.
SGLang의 SessionAwareCache는 세션 단위로 KV 리소스를 점유하는 Decorator 패턴의 캐시다. 스트리밍 세션의 KV를 반환하지 않고 세션 슬롯에 보관해 두었다가, 같은 세션의 다음 턴에서 즉시 재활용한다. python/sglang/srt/mem_cache/session_aware_cache.py를 분석한다.
구조도
SessionAwareCache (Decorator)
├── inner: BasePrefixCache ← 실제 캐시 (RadixCache 등)
├── slots: Dict[str, SessionSlot] ← 세션별 KV 상태 보관
│
│ SessionSlot
│ ├── virtual_node: _VirtualNode ← 락 구분용 센티넬
│ ├── req_pool_idx: int ← KV pool 인덱스
│ ├── kv_committed_len: int ← 커밋된 KV 길이
│ ├── kv_allocated_len: int ← 할당된 KV 길이
│ ├── last_node: TreeNode ← Radix tree 노드 참조
│ ├── mamba_pool_idx ← Mamba 상태 (하이브리드)
│ └── swa_evicted_seqlen: int ← SWA 상태
│
│ 요청 흐름:
│ Turn 1 완료 → save_from_req() → SessionSlot에 저장
│ Turn 2 시작 → restore_to_req() → SessionSlot에서 복원
│ 세션 종료 → release_session() → 리소스 해제
SessionSlot: 턴 간 상태 보존
SessionSlot은 한 스트리밍 세션의 KV 상태를 턴 사이에 보관하는 데이터 클래스다. KV Cache 인덱스, Mamba 상태, SWA(Sliding Window Attention) 상태를 모두 관리한다.
@dataclass
class SessionSlot:
virtual_node: _VirtualNode = field(default_factory=_VirtualNode)
req_pool_idx: Optional[int] = None
kv_committed_len: int = 0
kv_allocated_len: int = 0
last_node: Any = None
cache_protected_len: int = 0
# Mamba 상태
mamba_pool_idx: Any = None
mamba_ping_pong_track_buffer: Any = None
mamba_next_track_idx: Any = None
@property
def is_holding_kv(self) -> bool:
return self.req_pool_idx is not None
save_from_req()는 요청이 끝나면 KV 상태를 슬롯에 저장하고, 요청 객체에서 리소스 참조를 제거한다. 이렇게 하면 스케줄러가 요청 리소스를 해제하더라도 KV는 슬롯에 남는다.
def save_from_req(self, req: Req, is_first: bool):
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.mamba_pool_idx = req.mamba_pool_idx
# 요청에서 참조 제거 (이중 해제 방지)
req.req_pool_idx = None
req.mamba_pool_idx = None
restore_to_req()는 다음 턴의 요청에 저장된 KV 상태를 복원한다. 주의할 점은 슬롯에서 참조를 제거하지 않는다는 것이다.
def restore_to_req(self, req: Req):
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.mamba_pool_idx = self.mamba_pool_idx
# NOTE: req_pool_idx와 mamba_pool_idx는 슬롯에서 제거하지 않음.
# chunked prefill에서 스케줄러가 요청을 reject하고 재시도할 수 있으므로
# 슬롯은 멱등 복원을 위해 유지되어야 한다.
Decorator 패턴: 비침습적 확장
SessionAwareCache는 BasePrefixCache를 감싸는 Decorator다. 비스트리밍 요청은 inner 캐시로 그대로 전달하고, 스트리밍 요청만 SessionSlot 로직을 적용한다.
class SessionAwareCache(BasePrefixCache):
def __init__(self, inner: BasePrefixCache):
self.inner = inner
self.slots: Dict[str, SessionSlot] = {}
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
req = params.req
if not _is_streaming(req):
return self.inner.match_prefix(params)
session_id = req.session.session_id
slot = self.slots.get(session_id)
if slot is None or slot.req_pool_idx is None:
return self.inner.match_prefix(params)
slot.restore_to_req(req)
prefix_len = min(req.kv_committed_len,
max(len(params.key.token_ids) - 1, 0))
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,
)
스트리밍 여부는 단순한 헬퍼 함수로 판별한다.
def _is_streaming(req: Optional[Req]) -> bool:
return req is not None and req.session is not None and req.session.streaming
_VirtualNode: 락 구분 센티넬
_VirtualNode는 실제 Radix Tree 노드가 아닌 센티넬 객체다. inc_lock_ref와 dec_lock_ref에서 이 객체를 인식하면 no-op으로 처리한다.
class _VirtualNode:
"""Sentinel node for streaming session requests."""
pass
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
if isinstance(node, _VirtualNode):
return IncLockRefResult() # no-op
return self.inner.inc_lock_ref(node)
세션이 KV를 점유하는 동안에는 Radix Tree의 락을 사용하지 않는다. 세션 슬롯 자체가 리소스 소유권을 나타내기 때문이다.
세션 생명주기 관리
세션이 종료되면 release_session()이 모든 리소스를 해제한다.
def release_session(self, session_id: str):
slot = self.slots.pop(session_id, None)
if slot is None:
return
# Radix tree 락 해제
if slot.last_node is not None:
if slot.swa_uuid_for_lock is not None:
self.inner.dec_lock_ref(slot.last_node,
DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock))
else:
self.inner.dec_lock_ref(slot.last_node)
# KV pool 리소스 해제
if slot.is_holding_kv:
start = slot.cache_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)
세션이 보유한 토큰 수를 추적하는 메서드들도 제공한다. 스케줄러가 메모리 예산을 계산할 때 세션이 점유한 리소스를 파악해야 하기 때문이다.
def session_held_tokens(self) -> int:
total = 0
for slot in self.slots.values():
if slot.is_holding_kv:
allocated = ceil_align(slot.kv_allocated_len, self.page_size)
total += allocated - slot.cache_protected_len
return total
sanity_check 예외 처리
세션이 KV를 보유하는 동안에는 inner 캐시의 sanity check를 건너뛴다. 일반적으로 idle 상태에서는 모든 노드가 unlocked여야 하지만, 세션이 tree 락을 유지하고 있으면 이 가정이 깨지기 때문이다.
def sanity_check(self):
if any(s.is_holding_kv for s in self.slots.values()):
return # 세션이 락을 유지하는 동안 skip
self.inner.sanity_check()
설계 근거: 왜 Decorator인가
SessionAwareCache가 RadixCache를 상속하지 않고 감싸는 이유는 조합 가능성 때문이다. RadixCache, MambaRadixCache, 어떤 BasePrefixCache 구현체든 inner로 넣을 수 있다. 세션 로직과 캐시 로직이 완전히 분리되므로, 각각 독립적으로 테스트하고 발전시킬 수 있다.
RadixCache → SessionAwareCache(RadixCache)
MambaRadixCache → SessionAwareCache(MambaRadixCache)
LMCRadixCache → SessionAwareCache(LMCRadixCache)
관련 포스트
- 캐시 Eviction 정책: LRU, LFU, FIFO 비교 분석
- Mamba Radix Cache: SSM 모델을 위한 상태 캐싱
- Hybrid Cache Controller: GPU/CPU 하이브리드 캐시 관리
참고
관련 포스트
SGLang 의 다른글
- 이전글 [SGLang] Hybrid Cache Controller: GPU/CPU 하이브리드 캐시 관리
- 현재글 : [SGLang] Session-Aware Cache: 사용자별 KV 캐시 파티셔닝
- 다음글 [SGLang] 외부 스토리지 백엔드: LMCache, 3FS, Mooncake, NIXL
댓글