[sglang] SGLang 성능 최적화: D2H 복사 연산의 비동기 오버랩 구현

PR 링크: sgl-project/sglang#29075 상태: Merged | 변경: +61 / -39

들어가며

대규모 언어 모델(LLM) 추론 엔진인 SGLang에서 추론 성능을 저해하던 병목 중 하나는 GPU에서 CPU로 데이터를 전송하는 D2H(Device-to-Host) 복사 과정이었습니다. 기존에는 forward_stream에서 D2H 복사가 동기적으로 수행되어, 다음 추론 단계(Forward step)가 시작되기 전까지 스케줄러 스레드가 블로킹되는 문제가 있었습니다. 이번 PR은 이 D2H 복사를 전용 copy_stream으로 분리하여 다음 추론 단계와 병렬로 실행되도록 개선했습니다.

코드 분석

1. 전용 비동기 복사 프리미티브 도입 (python/sglang/srt/managers/utils.py)

가장 핵심적인 변경은 _async_d2h 함수의 도입입니다. 기존의 단순 .to("cpu", non_blocking=True) 방식은 페이지 가능한 메모리(pageable memory)를 사용하여 호스트 측에서 블로킹을 유발했습니다. 이를 pin_memory를 사용하는 방식으로 변경하고 record_stream을 통해 GPU 메모리 수명 주기를 안전하게 관리하도록 했습니다.

def _async_d2h(t: torch.Tensor) -> torch.Tensor:
    if not t.is_cuda:
        return t.to("cpu", non_blocking=True)
    cpu_t = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
    cpu_t.copy_(t, non_blocking=True)
    t.record_stream(torch.cuda.current_stream(t.device))
    return cpu_t

2. 스케줄러에서의 스트림 오버랩 (python/sglang/srt/managers/scheduler.py)

스케줄러는 이제 forward_stream과 copy_stream을 명시적으로 분리합니다. wait_stream을 사용하여 데이터 의존성을 보장하면서도, 복사 작업이 forward_stream을 가로막지 않도록 설계되었습니다.

self.copy_stream.wait_stream(self.forward_stream)
with self.copy_stream_ctx:
    batch_result.copy_to_cpu(...)

3. 모듈화된 텐서 매핑 (python/sglang/srt/eplb/expert_distribution.py)

기존에는 각 클래스가 직접 복사 로직을 구현했으나, 이제는 map_device_tensors 인터페이스를 통해 _async_d2h를 주입받는 방식으로 구조화되었습니다. 이는 코드 중복을 줄이고 유지보수성을 크게 향상시킵니다.

왜 이게 좋은가

이 최적화의 핵심은 '파이프라이닝(Pipelining)'입니다.

1. 스케줄러 블로킹 해소: D2H 복사가 forward_stream에서 분리됨으로써, GPU는 다음 토큰 생성을 위한 연산을 즉시 시작할 수 있습니다. 이는 추론 지연 시간(Latency)을 직접적으로 줄여줍니다.
2. 메모리 효율성: pin_memory를 사용하여 호스트와 디바이스 간의 데이터 전송 효율을 극대화했습니다.
3. 안정성: record_stream을 통해 비동기 복사가 완료되기 전까지 GPU 텐서가 해제되지 않도록 보장하여, 기존의 잠재적인 메모리 레이스 컨디션을 방지했습니다.

결과적으로 이 변경은 고성능 추론 엔진에서 흔히 발생하는 '데이터 전송 병목'을 효과적으로 해결한 모범적인 사례입니다. 특히 프로파일링을 위해 @torch.profiler.record_function을 추가하여 성능 추적을 용이하게 만든 점도 운영 측면에서 훌륭한 개선입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html
• https://pytorch.org/docs/stable/generated/torch.Tensor.pin_memory.html
• https://pytorch.org/docs/stable/generated/torch.cuda.Stream.html

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