[SGLang] HiRadixCache: 계층적 GPU/CPU/Disk KV 캐시

2026년 4월 10일수정: 2026년 4월 10일

들어가며

GPU 메모리는 비싸고 한정적이다. A100 80GB에서 KV 캐시가 70GB를 차지하면, 동시에 처리할 수 있는 요청 수가 급격히 줄어든다. 그렇다고 캐시를 작게 유지하면 프리픽스 재사용률이 떨어진다. SGLang의 HiRadixCache는 이 딜레마를 GPU-CPU-Disk 3계층 캐시로 해결한다. 자주 사용되는 KV 캐시는 GPU에, 덜 사용되는 것은 CPU 메모리에, 거의 사용되지 않는 것은 디스크에 저장한다.

이 글에서는 python/sglang/srt/mem_cache/hiradix_cache.py를 중심으로 HiRadixCache의 설계를 분석한다.

3계층 캐시 아키텍처

┌─────────────────────────────────────────────────────┐
│  L1: GPU (Hot)                                      │
│  ┌─────┐ ┌─────┐ ┌─────┐                           │
│  │Node │ │Node │ │Node │  ← 활성 요청의 KV 캐시     │
│  │value│ │value│ │value│                             │
│  └──┬──┘ └──┬──┘ └──┬──┘                            │
├─────┼───────┼───────┼───────────────────────────────┤
│  L2: CPU (Warm)          write_through ↑  load_back │
│  ┌──┴──┐ ┌──┴──┐ ┌──┴──┐              ↓            │
│  │Node │ │Node │ │Node │  ← 백업된 KV 캐시          │
│  │host │ │host │ │host │                             │
│  └──┬──┘ └──┬──┘ └──┬──┘                            │
├─────┼───────┼───────┼───────────────────────────────┤
│  L3: Disk/Storage (Cold)   backup ↑  prefetch       │
│  ┌──┴──┐ ┌──┴──┐ ┌──┴──┐          ↓                │
│  │Blob │ │Blob │ │Blob │  ← 영구 저장               │
│  └─────┘ └─────┘ └─────┘                            │
└─────────────────────────────────────────────────────┘

초기화: 캐시 유형별 호스트 풀

HiRadixCache는 GPU KV 캐시 유형에 맞는 CPU 호스트 풀을 생성한다.

class HiRadixCache(RadixCache):
    def __init__(self, params: CacheInitParams, server_args: ServerArgs):
        self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()

        if isinstance(self.kv_cache, MHATokenToKVPool):
            self.token_to_kv_pool_host = MHATokenToKVPoolHost(
                self.kv_cache, server_args.hicache_ratio,
                server_args.hicache_size, self.page_size,
                server_args.hicache_mem_layout,
                allocator_type=server_args.hicache_storage_backend,
            )
        elif isinstance(self.kv_cache, MLATokenToKVPool):
            self.token_to_kv_pool_host = MLATokenToKVPoolHost(...)

MHA, MLA, NSA 세 가지 attention 아키텍처를 모두 지원한다. 각각의 KV 캐시 레이아웃이 다르므로 호스트 풀도 별도로 구현된다.

Write-Through: GPU → CPU 백업

노드의 hit_count가 임계값에 도달하면 CPU로 백업한다.

def _inc_hit_count(self, node: TreeNode, chunked=False):
    if self.cache_controller.write_policy == "write_back" or chunked:
        return
    node.hit_count += 1
    if not node.backuped:
        if node.hit_count >= self.write_through_threshold:
            self.write_backup(node)

write_through_threshold는 쓰기 정책에 따라 달라진다. write_through 모드에서는 1(즉시 백업), 그 외에는 2(두 번째 접근 시 백업)이다.

실제 백업은 write_backup 메서드에서 수행한다.

def write_backup(self, node: TreeNode, write_back=False) -> int:
    # 백업 연속성 보장: 부모가 백업되지 않았으면 건너뛰기
    if not write_back and (
        node.parent != self.root_node and not node.parent.backuped):
        return 0

    host_indices = self.cache_controller.write(
        device_indices=node.value, node_id=node.id)
    if host_indices is None:
        self.evict_host(len(node.value))
        host_indices = self.cache_controller.write(
            device_indices=node.value, node_id=node.id)

    if host_indices is not None:
        node.host_value = host_indices.clone()
        self.ongoing_write_through[node.id] = node
        if not write_back:
            self.inc_lock_ref(node)  # 전송 중 eviction 방지
    return len(host_indices) if host_indices is not None else 0

핵심 설계 원칙은 백업 연속성이다. 부모 노드가 백업되지 않은 상태에서 자식을 백업하면, load back 시 프리픽스 체인이 깨진다. 따라서 루트부터 연속적으로 백업된 노드만 허용한다.

Storage Backend: CPU → Disk 백업

CPU 백업이 완료되면, storage backend로 디스크에 영구 저장할 수 있다.

def write_backup_storage(self, node: TreeNode):
    prefix_keys = (
        node.get_prefix_hash_values(node.parent)
        if self.hicache_storage_pass_prefix_keys
        else None
    )
    operation_id = self.cache_controller.write_storage(
        node.host_value, node.key, node.hash_value, prefix_keys)
    self.ongoing_backup[operation_id] = node
    node.protect_host()  # 디스크 전송 중 CPU eviction 방지

protect_host()로 노드의 host_ref_counter를 증가시켜, 디스크 전송이 완료될 때까지 CPU 메모리에서 해제되지 않도록 보호한다.

Writing Check: 비동기 전송 완료 확인

GPU-CPU 전송은 비동기로 수행된다. 스케줄러가 주기적으로 완료를 확인한다.

def writing_check(self, write_back=False):
    if len(self.ongoing_write_through) == 0:
        return

    finish_count = 0
    for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
        if not finish_event.query():
            break
        finish_count += 1

    # TP 워커 간 동기화: 모든 워커가 완료한 수만큼만 처리
    queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
    if self.tp_world_size > 1:
        torch.distributed.all_reduce(
            queue_size, op=torch.distributed.ReduceOp.MIN,
            group=self.tp_group)

    finish_count = int(queue_size.item())
    while finish_count > 0:
        _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0)
        finish_event.synchronize()
        for ack_id in ack_list:
            backuped_node = self.ongoing_write_through.pop(ack_id)
            self.dec_lock_ref(backuped_node)
            if self.enable_storage:
                self.write_backup_storage(backuped_node)
        finish_count -= 1

Tensor Parallel 환경에서 핵심적인 동기화 로직이 있다. ReduceOp.MIN으로 모든 TP 워커가 완료한 최소 수만큼만 처리한다. 이는 워커 간 Radix Tree 상태를 일관되게 유지하기 위한 것이다.

동적 Storage Backend 관리

런타임에 storage backend를 동적으로 연결/해제할 수 있다.

def attach_storage_backend(self, storage_backend, ...):
    if self.enable_storage:
        current_backend = self.cache_controller.storage_backend_type
        if current_backend == storage_backend:
            # 같은 백엔드: 정책만 업데이트
            if hicache_storage_prefetch_policy is not None:
                self.prefetch_stop_policy = hicache_storage_prefetch_policy
            return True, "policies updated."
        return False, "Detach first."

    self.cache_controller.attach_storage_backend(
        storage_backend=storage_backend, ...)
    return True, "Attached successfully."

def detach_storage_backend(self):
    self._drain_storage_control_queues_local()
    self.cache_controller.detach_storage_backend()
    self._force_release_pending_storage_ops()
    self.enable_storage = False
    return True, "Detached successfully."

detach 시 세 단계의 안전장치가 작동한다. 먼저 대기 중인 제어 큐를 비우고, 컨트롤러의 스토리지 스레드를 중지하고, 남은 보류 중인 작업을 강제 해제한다.

Prefetch 정책

디스크에서 KV 캐시를 미리 가져오는 세 가지 정책을 지원한다.

# server_args.hicache_storage_prefetch_policy
"best_effort"    # 비동기, 안 오면 말고
"wait_complete"  # 동기, 완료될 때까지 대기
"timeout"        # 타임아웃 내에서 대기

타임아웃은 토큰 수에 비례하여 계산된다.

self.prefetch_timeout_base = prefetch_timeout_base        # 기본 1초
self.prefetch_timeout_per_page = (
    self.page_size / 1024 * prefetch_timeout_per_ki_token  # 1024토큰당 0.25초
)

Host Eviction 관리

CPU 메모리도 무한하지 않으므로 eviction이 필요하다.

def _update_host_leaf_status(self, node: TreeNode):
    # GPU와 별도로 CPU 리프 노드 상태를 추적
    ...

GPU의 evictable_leaves와 별도로 evictable_host_leaves를 관리하여 CPU 메모리의 eviction을 독립적으로 처리한다.

설계 근거

Write-through vs Write-back: Write-through(hit_count >= 1)는 데이터 손실 위험이 적지만 대역폭을 많이 사용한다. Write-back은 eviction 시점에만 백업하므로 대역폭 효율이 높지만, GPU eviction이 CPU 전송 완료를 기다려야 하는 지연이 발생할 수 있다. SGLang은 두 정책을 모두 지원하며, hit_count 기반의 write_through_selective까지 제공한다.

백업 연속성 제약: 부모 없이 자식만 백업하면 load back 시 프리픽스 체인을 재구성할 수 없다. 이 제약은 load back의 정확성을 보장하면서도, 인기 있는 프리픽스가 자연스럽게 루트부터 순서대로 백업되는 효과가 있다.

TP 동기화에 MIN 사용: 각 GPU가 서로 다른 속도로 CPU 전송을 완료할 수 있다. 가장 느린 GPU의 완료 수에 맞추면, 모든 GPU의 Radix Tree 상태가 항상 동일하게 유지된다. 이는 collective 통신의 정확성을 보장한다.

참고

SGLang 의 다른글

이전글 [SGLang] Allocator: 토큰-KV 풀 할당 전략의 설계
현재글 : [SGLang] HiRadixCache: 계층적 GPU/CPU/Disk KV 캐시
다음글 [SGLang] Sliding Window Attention 캐시: SWA 최적화 설계