[sglang] SGLang 성능 최적화: torch.cuda.empty_cache() 호출 제어를 통한 가중치 업데이트 병목 해결

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

들어가며

대규모 언어 모델(LLM)을 서빙하는 환경에서 가중치(Weight)를 동적으로 업데이트하는 것은 매우 빈번한 작업입니다. 하지만 SGLang의 기존 구현에서는 가중치 업데이트 후 캐시를 비우는 flush_cache() 과정에서 항상 torch.cuda.empty_cache()를 호출하고 있었습니다. torch.cuda.empty_cache()는 모든 CUDA 스트림을 동기화(Synchronize)하는 무거운 작업으로, 동시 추론 요청이 많은 환경에서 심각한 지연 시간(Latency)을 유발합니다. 본 PR은 이 호출을 선택적으로 변경하여 성능을 최적화합니다.

코드 분석

1. scheduler.py: flush_cache 함수 시그니처 변경

기존에는 flush_cache가 호출될 때마다 무조건적으로 torch.cuda.empty_cache()를 실행했습니다. 이를 선택적 인자로 변경하여 제어권을 확보했습니다.

# Before
def flush_cache(self):
    # ...
    torch.cuda.empty_cache()

# After
def flush_cache(self, empty_cache: bool = True):
    # ...
    if empty_cache:
        torch.cuda.empty_cache()

2. io_struct.py: 요청 입력 데이터 구조 확장

가중치 업데이트 요청 시 torch_empty_cache 플래그를 전달할 수 있도록 각 요청 클래스에 필드를 추가했습니다.

# 예시: UpdateWeightsFromDistributedReqInput
class UpdateWeightsFromDistributedReqInput(BaseReq):
    # ...
    torch_empty_cache: bool = False

3. scheduler_update_weights_mixin.py: 업데이트 경로 수정

가중치 업데이트 로직(from_disk, from_distributed 등)에서 recv_req.torch_empty_cache 값을 전달하도록 수정하여, 필요할 때만 캐시를 비우도록 했습니다.

# Before
flush_cache_success = self.flush_cache()

# After
flush_cache_success = self.flush_cache(
    empty_cache=recv_req.torch_empty_cache
)

왜 이게 좋은가

torch.cuda.empty_cache()는 PyTorch의 캐시 할당자(Caching Allocator)가 보유한 미사용 메모리를 OS에 반환하도록 시도하는 함수입니다. 이 과정에서 GPU의 모든 활성 스트림을 대기시켜야 하므로, 추론 엔진이 실시간으로 요청을 처리 중일 때 호출하면 Stop-the-world 현상과 유사한 지연이 발생합니다.

이번 최적화의 핵심 교훈은 다음과 같습니다:

1. 동기화 지점 최소화: 고성능 추론 엔진에서는 불필요한 GPU 동기화 호출을 제거하는 것만으로도 P99 지연 시간을 크게 개선할 수 있습니다.
2. 설정의 유연성: 모든 상황에서 캐시를 비우는 대신, 사용자가 필요 여부를 결정하게 함으로써 시스템의 처리량(Throughput)을 최적화할 수 있습니다.

리뷰 과정에서 확인된 CI 실패 사례들은 본 변경사항과 무관한 인프라 이슈(AMD/NPU 환경의 타임아웃 등)로 판명되었으며, 핵심 로직인 가중치 업데이트 및 캐시 관리 경로는 안전하게 유지되었습니다.

결론

이번 변경은 SGLang의 가중치 업데이트 메커니즘을 더욱 정교하게 다듬어, 대규모 트래픽 환경에서의 안정성을 높였습니다. 특히 실시간 모델 교체가 잦은 서비스 환경에서 유의미한 성능 향상을 기대할 수 있습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.cuda.empty_cache.html

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