[SGLang] Mamba 캐시 누수 수정: adder 실패 시 pool index 회수

PR 링크: sgl-project/sglang#21404 상태: Merged | 변경: +7 / -0

들어가며

Mamba 기반 모델(Jamba, Qwen3.5 등)은 SSM state를 위한 전용 메모리 풀(mamba_pool)을 사용한다. 스케줄러가 새 요청을 prefill 배치에 추가할 때 mamba_pool_idx가 할당되는데, 배치가 가득 차거나 메모리 부족으로 요청 추가가 실패하면 이미 할당된 pool index가 해제되지 않는 버그가 있었다.

핵심 코드 분석

Before:

# adder가 요청 추가를 시도하고 실패하면
# break로 루프를 빠져나가지만, 이미 할당된
# req.mamba_pool_idx는 해제되지 않음
break

After:

# 요청이 실제로 can_run_list에 추가되었는지 확인
added = len(adder.can_run_list) > 0 and req is adder.can_run_list[-1]
if not added and req.mamba_pool_idx is not None:
    self.tree_cache.req_to_token_pool.mamba_pool.free(
        req.mamba_pool_idx.unsqueeze(-1)
    )
    req.mamba_pool_idx = None
break

왜 이게 좋은가

1. Silent memory leak 수정: 요청이 실패할 때마다 SSM state 슬롯이 하나씩 누수되어, 장시간 운영 시 max-mamba-cache-size에 도달하면 모든 Mamba 요청이 거부된다.
2. 최소 변경: 7줄 추가만으로 문제를 정확히 해결한다. adder.can_run_list의 마지막 요소가 현재 요청인지 확인하는 단순한 로직이다.
3. 방어적 프로그래밍: mamba_pool_idx is not None 체크로 Mamba가 아닌 모델에서는 아무 영향이 없다.

정리

리소스 할당과 해제의 비대칭은 서버 애플리케이션에서 가장 흔한 메모리 누수 패턴이다. 이 PR은 "할당했으면 실패 경로에서 반드시 해제"라는 원칙을 7줄로 구현한 깔끔한 수정이다.

참고 자료

• PR #21404

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석과 해석에서 오류가 있을 수 있으니, 정확한 내용은 원본 PR을 참고해주세요.