[sglang] HiCache 메모리 누수 수정: host indices clone으로 참조 해제 보장

PR 링크: sgl-project/sglang#21624 상태: Merged | 변경: +3 / -3

들어가며

PyTorch에서 텐서를 다른 변수에 할당하면 데이터가 복사되지 않고 동일한 메모리를 공유합니다. 이는 성능상 이점이지만, 원본 텐서가 더 큰 버퍼의 일부(예: GPU 메모리 풀의 슬라이스)일 때 그 큰 버퍼가 해제되지 못하는 메모리 누수의 원인이 됩니다. 이번 PR은 SGLang의 HiCache에서 바로 이 문제를 수정합니다.

핵심 코드 분석

host_indices 참조를 clone으로 교체

Before (hiradix_cache.py):

if host_indices is not None:
    node.host_value = host_indices
    assert len(node.host_value) > 0

After:

if host_indices is not None:
    node.host_value = host_indices.clone()
    assert len(node.host_value) > 0

동일한 패턴이 hi_mamba_radix_cache.py에서도 두 곳 수정되었습니다:

# KV cache 쪽
node.host_value = host_indices.clone()

# Mamba state cache 쪽
node.mamba_host_value = host_indices.clone()

왜 이게 좋은가

1. 메모리 누수 방지: host_indices가 GPU-to-Host 전송 버퍼의 슬라이스라면, clone() 없이는 전체 전송 버퍼가 GC되지 못합니다. clone()은 필요한 데이터만 독립된 메모리에 복사하여 원본 버퍼의 해제를 허용합니다.
2. 데이터 무결성: 원본 host_indices가 이후 다른 전송에 재사용될 경우, clone 없이는 노드의 값이 의도치 않게 변경될 수 있습니다.

정리

3줄만 변경된 간결한 수정이지만, LLM 추론 서버처럼 장시간 운영되는 시스템에서는 이런 작은 메모리 누수가 OOM(Out of Memory)으로 이어질 수 있습니다. PyTorch 텐서의 참조 의미론(reference semantics)을 정확히 이해하는 것이 중요합니다.

참고 자료

• sgl-project/sglang#21624
• python/sglang/srt/mem_cache/hiradix_cache.py

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