[sglang] SGLang Mamba 캐시의 상태 손상 및 슬롯 누수 버그 수정

2026년 7월 23일수정: 2026년 7월 23일

PR 링크: sgl-project/sglang#30986 상태: Merged | 변경: +345 / -9

들어가며

대규모 언어 모델(LLM)의 추론 성능을 최적화하는 과정에서 캐시 관리의 중요성은 아무리 강조해도 지나치지 않습니다. 특히 SGLang과 같이 복잡한 캐시 메커니즘을 사용하는 경우, 예상치 못한 버그는 모델의 정확성 저하 또는 자원 누수로 이어질 수 있습니다.

이번 PR은 SGLang의 Mamba 캐시 구현에서 발생하는 심각한 버그를 해결합니다. 이 버그는 특정 조건 하에서 Mamba 캐시의 load_back 작업이 실패했을 때, 해당 요청의 상태가 손상되거나 메모리 슬롯이 누수되는 문제를 야기했습니다. 이러한 문제는 모델의 출력이 조용히 잘못 생성되거나, 호스트 메모리 압박 시 재현될 수 있어 서비스 안정성에 큰 위협이 되었습니다.

본 글에서는 이 PR이 해결하는 문제의 근본 원인을 분석하고, 코드 변경 사항을 통해 어떻게 버그를 수정했는지, 그리고 이 수정이 왜 좋은 최적화인지 상세히 설명하겠습니다.

코드 분석

이번 PR의 핵심은 Mamba 컴포넌트의 load_back 작업 실패 시 상태 관리 로직을 개선하는 것입니다. 기존에는 load_back 작업이 실패하더라도 할당된 메모리 슬롯이 제대로 해제되지 않아 상태 손상 및 슬롯 누수가 발생했습니다. 이 문제는 TreeComponent의 새로운 훅(hook)을 도입하여 해결되었습니다.

1. `TreeComponent`의 새로운 훅 도입

리뷰 피드백을 반영하여 컴포넌트별 상태 관리가 TreeComponent 내부에 통합되었습니다. 이를 위해 두 개의 새로운 훅이 도입되었습니다:

prepare_load_back: load_back 작업이 실제로 시작되기 전에 호출됩니다. Mamba 컴포넌트의 경우, 이 단계에서 필요한 CoW(Copy-on-Write) 디바이스 슬롯을 미리 할당하고 req.mamba_pool_idx에 기록합니다. 이 슬롯은 PrepareLoadBackResult 객체에 저장됩니다.
finalize_load_back: load_back 작업의 성공 여부를 인자로 받아 호출됩니다. 작업이 실패했을 경우(success가 False일 때), prepare_load_back에서 할당했던 슬롯을 해제하고 req.mamba_pool_idx를 초기화합니다. 이를 통해 실패 시에도 자원이 올바르게 정리되도록 보장합니다.

이 두 훅은 load_back 작업의 모든 종료 경로(성공, 실패, 예외 발생)에서 일관되게 호출되어 상태를 정리합니다.

2. Mamba 컴포넌트의 변경사항

MambaComponent는 TreeComponent의 새로운 훅을 구현하여 Mamba 관련 상태 관리를 담당합니다.

Before:

build_hicache_transfers 함수 내에서 load_back 작업 시 Mamba 슬롯 할당이 직접 이루어졌습니다.

def build_hicache_transfers(
    self,
    node: UnifiedTreeNode,
    phase: CacheTransferPhase,
    *, # MambaComponent specific
    req: Optional[Req] = None,
) -> PrepareTransfersResult:
    ...
    # Per-request mamba CoW (H→D copy into request's device slot)
    cd = node.component_data[ct]
    if req is not None and cd.host_value is not None:
        if req.mamba_pool_idx is None:
            dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
            if dst is None:
                self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
                dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
                assert dst is not None, "Cannot alloc mamba for load_back"
            req.mamba_pool_idx = dst[0]
        transfers.append(
            PoolTransfer(
                name=PoolName.MAMBA,
                host_indices=cd.host_indices,
                device_indices=torch.tensor([req.mamba_pool_idx], dtype=torch.long),
            )
        )
    ...

After:

prepare_load_back에서 슬롯 할당 및 req.mamba_pool_idx 설정이 이루어지고, build_hicache_transfers에서는 이미 할당된 슬롯이 존재함을 가정합니다.

def prepare_load_back(
    self,
    node: UnifiedTreeNode,
    *,
    req: Optional[Req] = None,
) -> PrepareLoadBackResult:
    cd = node.component_data[self.component_type]
    if (
        req is None
        or req.mamba_pool_idx is not None
        or cd.host_value is None
        or cd.value is not None
    ):
        return PrepareLoadBackResult()
    dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
    if dst is None:
        self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
        dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
        assert dst is not None, "Cannot alloc mamba for load_back"
    req.mamba_pool_idx = dst[0]
    return PrepareLoadBackResult(allocated_mamba_slot=dst)

def finalize_load_back(
    self, req: Optional[Req], prep: PrepareLoadBackResult, success: bool
) -> None:
    # A called-off load-back returns the slot prepare allocated and clears req (the H->D copy never ran).
    if not success and prep.allocated_mamba_slot is not None:
        self.cache.req_to_token_pool.mamba_allocator.free(prep.allocated_mamba_slot)
        req.mamba_pool_idx = None

def build_hicache_transfers(
    self,
    node: UnifiedTreeNode,
    phase: CacheTransferPhase,
    *,
    req: Optional[Req] = None,
) -> PrepareTransfersResult:
    ...
    # Per-request mamba CoW: H→D copy into the request's device slot allocated by prepare_load_back.
    cd = node.component_data[ct]
    if req is not None and cd.host_value is not None:
        assert req.mamba_pool_idx is not None
        transfers.append(
            PoolTransfer(
                name=PoolName.MAMBA,
                host_indices=cd.host_indices,
                device_indices=torch.tensor([req.mamba_pool_idx], dtype=torch.long),
            )
        )
    ...

또한, finalize_load_back 함수는 load_back 작업이 실패했을 때 할당된 슬롯을 해제하고 req.mamba_pool_idx를 None으로 설정합니다. 이는 실패 시에도 요청이 잘못된 상태를 참조하는 것을 방지합니다.

3. `UnifiedRadixCache.load_back`의 변경사항

UnifiedRadixCache의 load_back 함수는 이제 prepare_load_back과 finalize_load_back 훅을 호출하도록 수정되었습니다.

Before:

load_back 함수는 직접 슬롯 할당 및 전송 빌드 로직을 포함했습니다.

def load_back(
    self,
    best_match_node: UnifiedTreeNode,
    mem_quota: Optional[int] = None,
    req: Optional[Req] = None,
) -> bool:
    ...
    kv_tokens = len(kv_xfer.host_indices)
    ...
    # Build aux transfers, keyed per component.
    comp_xfers: dict[ComponentType, list] = {}
    for comp in self._components_tuple:
        if comp.component_type == ComponentType.KV:
            continue
        comp_xfers[comp.component_type] = comp.build_hicache_transfers(
            best_match_node,
            CacheTransferPhase.LOAD_BACK,
            req=req,
        )
    ...
    # If load_back fails, the slot stays allocated until request end (a bounded leak).
    # The corruption is the severe half.
    return self._load_back_transfers(
        best_match_node=best_match_node,
        mem_quota=mem_quota,
        req=req,
        kv_xfer=kv_xfer,
        result=result,
        ancestor_lock_params=ancestor_lock_params,
        host_anchor_params=host_anchor_params,
    )

After:

load_back 함수는 try...finally 블록을 사용하여 prepare_load_back을 호출하고, _load_back_transfers 함수 내부에서 실제 로직을 수행한 후, finally 블록에서 finalize_load_back을 호출하여 상태를 정리합니다.

def load_back(
    self,
    best_match_node: UnifiedTreeNode,
    mem_quota: Optional[int] = None,
    req: Optional[Req] = None,
) -> bool:
    # Let each component pre-allocate per-request state for the load-back;
    # the finally below lets components recover it unless the load succeeds.
    preps: dict[ComponentType, PrepareLoadBackResult] = {
        comp.component_type: comp.prepare_load_back(best_match_node, req=req)
        for comp in self._components_tuple
    }
    success = False
    try:
        success = self._load_back_transfers(
            best_match_node=best_match_node,
            mem_quota=mem_quota,
            req=req,
            kv_xfer=kv_xfer,
            result=result,
            ancestor_lock_params=ancestor_lock_params,
            host_anchor_params=host_anchor_params,
        )
        return success
    finally:
        for comp in self._components_tuple:
            comp.finalize_load_back(req, preps[comp.component_type], success)

def _load_back_transfers(
    self,
    *,
    best_match_node: UnifiedTreeNode,
    mem_quota: Optional[int],
    req,
    kv_xfer: PoolTransfer,
    result: IncLockRefResult,
    ancestor_lock_params: Optional[DecLockRefParams],
    host_anchor_params: Optional[DecLockRefParams],
) -> bool:
    kv_tokens = len(kv_xfer.host_indices)
    # Build aux transfers, keyed per component.
    comp_xfers: dict[ComponentType, list] = {}
    for comp in self._components_tuple:
        if comp.component_type == ComponentType.KV:
            continue
        comp_xfers[comp.component_type] = comp.build_hicache_transfers(
            best_match_node,
            CacheTransferPhase.LOAD_BACK,
            req=req,
        )
    ...

3. 테스트 케이스 추가

이 버그를 재현하고 수정되었음을 검증하기 위해 세 가지 새로운 테스트 케이스가 추가되었습니다:

test_load_back_abort_frees_unpublished_mamba_slot: mem_quota 부족으로 load_back이 중단되었을 때, 할당되었지만 게시되지 않은 Mamba 슬롯이 올바르게 해제되는지 확인합니다.
test_load_back_load_failure_frees_unpublished_mamba_slot: cache_controller.load() 호출 실패 시에도 슬롯이 해제되는지 확인합니다.
test_load_back_abort_keeps_preexisting_mamba_slot: load_back 호출 전에 이미 요청이 슬롯을 소유하고 있었을 때, 중단되더라도 해당 슬롯이 유지되는지 확인합니다.

이 테스트들은 수정 전에는 실패했지만, 수정 후에는 모두 통과하며 버그가 해결되었음을 입증합니다. 전체 유닛 테스트 스위트(test_unified_radix_cache_unittest.py)는 850개 통과 / 702개 건너뜀으로 녹색 상태를 유지했습니다.

왜 이게 좋은가

이 PR은 다음과 같은 이유로 훌륭한 최적화 및 버그 수정입니다.

정확성 보장: 가장 중요한 개선점은 모델의 정확성입니다. 이전에는 load_back 실패 시 초기화되지 않았거나 이전 상태가 남아있는 메모리 슬롯을 읽어 조용히 잘못된 출력을 생성할 수 있었습니다. 이 수정으로 인해 이러한 상태 손상 경로가 제거되어 모델의 예측 정확성이 보장됩니다.
메모리 누수 방지: 실패한 load_back 작업 후에도 Mamba 슬롯이 해제되지 않고 요청이 끝날 때까지 남아있는 '바운드된 누수(bounded leak)'가 있었습니다. 이 PR은 실패 시 슬롯을 즉시 해제함으로써 이러한 메모리 누수를 방지하고 시스템 자원을 효율적으로 관리합니다.
코드 구조 개선: 컴포넌트별 상태 관리 로직을 TreeComponent의 훅으로 분리함으로써 코드의 응집성과 재사용성이 향상되었습니다. 이는 향후 유사한 캐시 관리 로직을 추가하거나 수정할 때 유지보수성을 높이는 데 기여합니다.
견고한 오류 처리: try...finally 블록과 새로운 훅 메커니즘을 통해 load_back 작업의 모든 종료 경로(성공, 실패, 예외)에서 일관된 상태 정리 로직을 적용하여 오류 처리의 견고성을 높였습니다.

성능 영향: 이 PR은 주로 오류 처리 및 상태 관리 로직을 재구성한 것이므로 직접적인 성능 향상(속도 증가)은 기대하기 어렵습니다. 오히려 prepare_load_back과 finalize_load_back 훅 호출 오버헤드가 약간 발생할 수 있습니다. 하지만 이는 정확성 및 안정성 확보라는 더 큰 이득에 비하면 미미한 수준입니다. PR 설명에도 'N/A — no-op hooks for non-mamba components; allocation work unchanged, only relocated.'라고 명시되어 있습니다.

일반적인 교훈: 복잡한 시스템에서 자원 할당 및 해제는 매우 중요합니다. 특히 비동기 작업이나 실패 가능성이 있는 작업에서는 할당된 자원이 작업 실패 시에도 반드시 회수될 수 있도록 명확한 콜백 또는 예외 처리 메커니즘을 설계해야 합니다. 또한, 컴포넌트별 로직을 중앙 집중식으로 관리하되, 각 컴포넌트가 필요한 부분을 확장할 수 있도록 인터페이스(훅)를 제공하는 것은 좋은 설계 패턴입니다.

결론

이번 PR은 SGLang의 Mamba 캐시에서 발생하던 중요한 버그를 성공적으로 해결했습니다. load_back 작업 실패 시의 상태 손상 및 메모리 누수 문제를 TreeComponent의 새로운 훅을 통해 체계적으로 관리함으로써, 모델의 정확성을 높이고 시스템의 안정성을 크게 향상시켰습니다. 이는 LLM 추론 시스템의 신뢰성을 확보하는 데 중요한 기여를 했습니다.

참고 자료

https://pytorch.org/docs/stable/generated/torch.compile.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

PR Analysis 의 다른글

이전글 [vllm] vLLM Inkling 모델: 인플레이스 연산으로 GPU 메모리 최적화 및 OOM 방지
현재글 : [sglang] SGLang Mamba 캐시의 상태 손상 및 슬롯 누수 버그 수정
다음글 [flashinfer] FlashInfer SM103 FP4 GEMM 최적화: Store256 및 Fused Epilogue 도입