[sglang] SGLang VLM 최적화: CUDA IPC Staging 오버헤드 제거를 통한 성능 향상

PR 링크: sgl-project/sglang#26096 상태: Merged | 변경: +127 / -79

들어가며

SGLang은 고성능 LLM 및 VLM 서빙을 위한 프레임워크입니다. 특히 VLM(Vision-Language Model) 처리 과정에서 멀티모달 데이터는 여러 프로세스 간의 효율적인 전달이 필수적입니다. 기존 SGLang 구현에서는 MultimodalProcessor 출력 단계에서 CUDA IPC를 통한 텐서 전송을 위해 매번 불필요한 스테이징(staging) 작업을 수행하고 있었습니다. 본 PR은 이러한 중복 작업을 제거하고, pad/hash 계산을 최적화하여 TTFT(Time To First Token)를 개선하는 데 목적이 있습니다.

코드 분석

1. python/sglang/srt/multimodal/processors/base_processor.py

기존에는 process_and_combine_mm_data 내부에서 텐서를 CudaIpcTensorTransportProxy로 변환하는 로직이 중복되어 있었습니다. 이를 _wrap_tensor_for_cuda_ipc라는 헬퍼 함수로 캡슐화하여 코드 가독성과 유지보수성을 높였습니다.

# Before: 중복된 로직이 반복됨
if isinstance(item.feature, torch.Tensor) and item.feature.is_cuda:
    # ... (IPC 변환 로직)

# After: 헬퍼 함수 사용
def _wrap_tensor_for_cuda_ipc(self, tensor: torch.Tensor):
    # ... (IPC 변환 로직 캡슐화)
    return CudaIpcTensorTransportProxy(...)

또한, item.set_pad_value()를 미리 호출하여 스케줄러 측에서 반복적인 해싱 작업을 방지했습니다.

2. python/sglang/srt/managers/schedule_batch.py

MultimodalDataItem의 reconstruct 메서드를 개선하여, 타겟 디바이스에서 한 번만 프록시를 재구성하도록 변경했습니다. 기존에는 reconstruct 호출 시 디바이스를 매번 확인했으나, 이제는 target_device를 명시적으로 전달받아 오버헤드를 줄였습니다.

# Before
def reconstruct(self):
    reconstruct_device = torch.cuda.current_device()
    # ...

# After
def reconstruct(self, target_device: int):
    # ...
    self.feature = self.feature.reconstruct_on_target_device(target_device)

왜 이게 좋은가

이번 최적화의 핵심은 '불필요한 데이터 복사 및 재계산 방지'입니다.

1. IPC 스테이징 최적화: pad와 hash를 미리 계산하여 스케줄러의 각 TP(Tensor Parallel) 랭크에서 동일한 작업을 반복하지 않도록 했습니다.
2. 메모리 효율성: reconstruct 로직을 최적화하여 CUDA IPC 프록시가 타겟 디바이스에서 효율적으로 물질화(materialize)되도록 개선했습니다.

성능 지표: Qwen3.5-397B 모델 기준, OCRBench 테스트에서 TTFT가 약 1.6% 개선되었으며, E2E 지연 시간 또한 소폭 감소했습니다. 텍스트 전용 워크로드에서는 회귀(regression)가 발생하지 않음을 확인했습니다.

교훈: 분산 환경에서 텐서를 전송할 때는 IPC 오버헤드가 성능의 병목이 될 수 있습니다. 특히 VLM과 같이 큰 멀티모달 데이터를 다룰 때는 '한 번 계산하고 여러 번 재사용(Compute once, reuse many)'하는 전략이 매우 중요합니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.cuda.current_device.html
• https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__IPC.html

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