본문으로 건너뛰기

[SGLang] 외부 스토리지 백엔드: LMCache, 3FS, Mooncake, NIXL

들어가며

GPU 메모리, CPU 메모리를 모두 사용해도 KV Cache 용량이 부족할 수 있다. 특히 멀티노드 환경에서 같은 모델을 서빙하는 여러 인스턴스가 동일한 prefix를 중복 계산하는 문제가 있다.

SGLang은 외부 스토리지 백엔드를 통해 KV Cache를 노드 간 공유하고, 디스크/네트워크 스토리지로 확장한다. python/sglang/srt/mem_cache/storage/backend_factory.py의 팩토리 패턴과 6개 빌트인 백엔드를 분석한다.

구조도

StorageBackendFactory
├── register_backend(name, module_path, class_name)
├── create_backend(backend_name, config, mem_pool_host)
│
│   등록된 빌트인 백엔드:
│   ├── "file"      → HiCacheFile          (로컬 파일)
│   ├── "nixl"      → HiCacheNixl          (NIXL 플러그인)
│   ├── "mooncake"  → MooncakeStore         (Mooncake 분산 KV)
│   ├── "hf3fs"     → HiCacheHF3FS         (3FS 파일시스템)
│   ├── "aibrix"    → AibrixKVCacheStorage  (AIBrix KV 캐시)
│   └── "eic"       → EICStorage            (EIC RDMA)
│
│   동적 백엔드:
│   └── "dynamic"   → 사용자 정의 모듈에서 로드
│
│   별도 캐시 통합:
│   └── LMCRadixCache (RadixCache 서브클래스)

         ┌──────────────────────────┐
         │     HiCacheStorage       │ ← 공통 인터페이스
         │  batch_exists(keys)      │
         │  batch_get(keys, locs)   │
         │  batch_set(keys, locs)   │
         └──────────────────────────┘

StorageBackendFactory: 지연 로딩 팩토리

팩토리는 백엔드 모듈을 등록 시점에 import하지 않는다. 실제로 create_backend()가 호출될 때 비로소 모듈을 로드한다. 이렇게 하면 사용하지 않는 백엔드의 의존성(nixl, mooncake 등)을 설치하지 않아도 된다.

class StorageBackendFactory:
    _registry: Dict[str, Dict[str, Any]] = {}

    @classmethod
    def register_backend(cls, name, module_path, class_name):
        def loader() -> type[HiCacheStorage]:
            return cls._load_backend_class(module_path, class_name, name)
        cls._registry[name] = {
            "loader": loader,
            "module_path": module_path,
            "class_name": class_name,
        }

빌트인 백엔드는 모듈 로드 시점에 등록된다.

StorageBackendFactory.register_backend(
    "nixl",
    "sglang.srt.mem_cache.storage.nixl.hicache_nixl",
    "HiCacheNixl",
)
StorageBackendFactory.register_backend(
    "mooncake",
    "sglang.srt.mem_cache.storage.mooncake_store.mooncake_store",
    "MooncakeStore",
)
StorageBackendFactory.register_backend(
    "hf3fs",
    "sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs",
    "HiCacheHF3FS",
)

동적 백엔드 로딩

"dynamic" 백엔드를 사용하면 사용자가 임의의 모듈에서 스토리지 클래스를 로드할 수 있다.

@classmethod
def _create_dynamic_backend(cls, backend_config, storage_config,
                            mem_pool_host, **kwargs):
    backend_name = backend_config["backend_name"]
    module_path = backend_config["module_path"]
    class_name = backend_config["class_name"]
    backend_class = cls._load_backend_class(module_path, class_name, backend_name)
    return backend_class(storage_config, kwargs)

필수 설정 필드는 backend_name, module_path, class_name 세 가지다.

HiCacheNixl: NVIDIA NIXL 기반

NIXL(NVIDIA Interconnect eXtension Library)은 GPU-GPU, GPU-storage 간 고속 데이터 전송을 제공한다. HiCacheNixl은 NIXL 에이전트를 생성하고 플러그인 기반으로 백엔드를 선택한다.

class HiCacheNixl(HiCacheStorage):
    def __init__(self, storage_config, file_path="/tmp/hicache_storage"):
        nixlconfig = NixlBackendConfig(storage_config.extra_config)
        plugin = nixlconfig.get_specified_plugin()
        self.file_manager = (
            NixlFileManager(file_path)
            if plugin not in NixlBackendSelection.OBJ_PLUGINS
            else None
        )
        agent_config = nixl_agent_config(backends=[])
        self.agent = nixl_agent(f"hicache_nixl_{uuid.uuid4()}", agent_config)

MLA 모델 여부에 따라 suffix가 달라진다. MLA 모델은 TP rank/size를 포함하지 않는다.

MooncakeStore: Mooncake 분산 KV

Mooncake의 분산 KV 스토어를 사용하며, 호스트 메모리를 Mooncake 전용 allocator로 할당한다. 이는 Mooncake의 zero-copy 전송을 위해 필요하다.

class MooncakeHostTensorAllocator(HostTensorAllocator):
    def allocate(self, dims, dtype, device="cpu") -> torch.Tensor:
        from mooncake.store import MooncakeHostMemAllocator
        self.allocator = MooncakeHostMemAllocator()
        size = 1
        for d in dims:
            size *= d
        size *= torch.tensor([], dtype=self.dtype).element_size()
        ptr_int = self.allocator.alloc(size)
        c_array = (ctypes.c_byte * size).from_address(ptr_int)
        tensor = torch.frombuffer(c_array, dtype=torch.uint8, count=size)
        return tensor.view(dims)

일반 torch.empty()가 아닌 Mooncake 전용 allocator를 사용하는 이유는 RDMA 등록이 필요한 메모리 영역을 할당하기 위해서다.

HiCacheHF3FS: 3FS 파일시스템

3FS(High-Performance File System for AI)는 메타데이터 서버와 데이터 노드로 구성된 분산 파일시스템이다. 페이지 단위로 KV를 저장하며, usrbio(user-space block I/O) 인터페이스로 고속 I/O를 수행한다.

class Hf3fsMetadataInterface(ABC):
    @abstractmethod
    def reserve_and_allocate_page_indices(
        self, rank, keys: List[Tuple[str, str]]
    ) -> List[Tuple[bool, int]]:
        """키별 페이지 인덱스를 할당하고, 기존 존재 여부를 반환"""
        pass

    @abstractmethod
    def confirm_write(self, rank, written_keys, pages_to_release):
        """쓰기 완료를 확인하고, 불필요한 페이지를 반환"""
        pass

    @abstractmethod
    def get_page_indices(self, rank, keys) -> List[Optional[int]]:
        """키에 대한 페이지 인덱스를 조회"""
        pass

AibrixKVCacheStorage: AIBrix

AIBrix KVCache 라이브러리를 사용하며, 블록 단위 해시 기반으로 KV를 관리한다.

class AibrixKVCacheStorage(HiCacheStorage):
    def __init__(self, storage_config, mem_pool):
        kv_cache = mem_pool.device_pool
        self.block_spec = KVCacheBlockSpec(
            block_ntokens=self.page_size,
            block_dtype=self.kv_cache_dtype,
            block_layout=KVCacheBlockLayout(KVCacheBlockLayout.NCLD),
            tensor_spec=KVCacheTensorSpec(
                heads=self.kv_head_ids,
                layers=self.layer_ids,
                head_size=kv_cache.head_dim,
            ),
        )
        config = KVCacheConfig(block_spec=self.block_spec, model_spec=ModelSpec(102400))
        self.kv_cache_manager = BaseKVCacheManager(config)

EICStorage: GPU Direct RDMA

EIC는 GPU Direct RDMA를 지원하는 고성능 KV 스토리지다. GPU NIC affinity를 설정하여 GPU-NIC 간 최적의 경로로 데이터를 전송한다.

GPUNicAffinity = {
    "cuda:0": "eth1", "cuda:1": "eth1",
    "cuda:2": "eth2", "cuda:3": "eth2",
    "cuda:4": "eth3", "cuda:5": "eth3",
    "cuda:6": "eth4", "cuda:7": "eth4",
}

LMCRadixCache: RadixCache 서브클래스

LMCache는 다른 백엔드와 달리 HiCacheStorage 인터페이스를 사용하지 않는다. 대신 RadixCache를 서브클래스로 확장하여 match_prefixcache_finished_req를 직접 오버라이드한다. 레이어 단위 비동기 전송을 위한 전용 CUDA stream을 관리한다.

class LMCRadixCache(RadixCache):
    """RadixCache + LMCache IO.
    Adds: LMCache connector setup, two CUDA streams for async load/store,
    layer-wise transfer executor wiring to KV cache."""

백엔드 비교표

백엔드 전송 방식 멀티노드 공유 GPU Direct 특징
file 로컬 파일 I/O X X 가장 단순, 디버깅용
NIXL 플러그인 기반 O O 다양한 NIXL 플러그인 선택 가능
Mooncake 분산 KV 스토어 O X Zero-copy 호스트 메모리, 글로벌 메타데이터
3FS 분산 파일시스템 O X usrbio 고속 I/O, 페이지 단위
AIBrix 블록 KV 관리자 O X 블록 해시 기반, 자동 메트릭
EIC RDMA O O GPU NIC affinity 최적화
LMCache 레이어별 비동기 O X RadixCache 서브클래스, 독립 통합

설계 근거: 팩토리 + 지연 로딩

백엔드마다 의존하는 외부 라이브러리가 다르다(nixl, mooncake, eic, lmcache 등). 모든 백엔드를 import하면 하나의 라이브러리만 없어도 SGLang 자체가 시작되지 않는다. 지연 로딩 팩토리는 실제로 사용하는 백엔드만 import하므로, 최소 의존성으로 동작한다.

또한 "dynamic" 백엔드를 통해 사용자가 HiCacheStorage를 상속한 커스텀 백엔드를 SGLang 코드 수정 없이 플러그인할 수 있다.

관련 포스트

  • Hybrid Cache Controller: GPU/CPU 하이브리드 캐시 관리
  • 캐시 Eviction 정책: LRU, LFU, FIFO 비교 분석
  • Mamba Radix Cache: SSM 모델을 위한 상태 캐싱

참고

댓글

관련 포스트

SGLang 의 다른글