[sglang] SGLang 성능 최적화: Speculative Decoding의 H2D 병목 해결 및 코드 중복 제거

PR 링크: sgl-project/sglang#28500 상태: Merged | 변경: +23 / -30

들어가며

LLM 추론 엔진인 SGLang에서 Speculative Decoding(특히 Eagle 모델)을 사용할 때, 매 디코딩 단계마다 발생하는 불필요한 동기화 문제가 성능 저하의 원인이 되고 있었습니다. 구체적으로 EagleDraftInput.prepare_for_decode 과정에서 torch.tensor(..., device=...)를 호출하며 발생하는 Host-to-Device(H2D) 복사가 메인 추론 스트림을 블로킹(Blocking)하고 있었습니다. 본 PR은 이 문제를 해결하기 위해 H2D 복사를 비동기(Non-blocking) 방식으로 전환하고, 기존의 중복된 로직을 ScheduleBatch로 통합하여 효율성을 극대화했습니다.

코드 분석

1. python/sglang/srt/managers/schedule_batch.py: 로직 통합 및 비동기 H2D 구현

가장 핵심적인 변경은 cumulate_penalty_output_tokens라는 공통 메서드를 신설한 것입니다. 기존에는 prepare_for_decode 내부에 하드코딩되어 있던 로직을 분리하여 재사용성을 높였습니다.

Before:

latest_output_ids = torch.tensor(last_tokens, dtype=torch.int64).to(
    self.device, non_blocking=True
)

After:

def cumulate_penalty_output_tokens(self):
    # ... (생략) ...
    latest_output_ids = torch.tensor(
        last_tokens,
        dtype=torch.int64,
        pin_memory=is_pin_memory_available(self.device),
    ).to(self.device, non_blocking=True)
    self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(latest_output_ids)

pin_memory를 명시적으로 사용함으로써 H2D 복사가 비동기적으로 처리되도록 보장했습니다. 이는 GPU가 연산을 수행하는 동안 CPU가 데이터를 준비하여 전송하는 파이프라이닝을 가능하게 합니다.

2. python/sglang/srt/speculative/eagle_info_v2.py: Spec-decode 경로 최적화

Speculative Decoding 경로에서도 동일한 동기화 문제가 발생하고 있었습니다. 기존 코드는 device=batch.device를 사용하여 즉각적인 동기화를 유발했습니다.

Before:

output_ids = torch.tensor(
    [...], 
    dtype=torch.int64,
    device=batch.device,
)
batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(output_ids)

After:

if batch.sampling_info.penalizer_orchestrator.is_required:
    batch.cumulate_penalty_output_tokens()

이제 batch.cumulate_penalty_output_tokens()를 호출함으로써, 앞서 최적화된 비동기 H2D 로직을 그대로 활용하게 되었습니다.

왜 이게 좋은가

1. 파이프라인 스톨(Stall) 제거: 기존의 동기식 H2D 복사는 GPU의 Forward 연산이 끝날 때까지 CPU가 대기하거나, 반대로 GPU가 CPU의 데이터 전송을 기다리게 만드는 '동기화 병목'을 유발했습니다. non_blocking=True와 pin_memory 조합은 이 병목을 제거하여 추론 지연 시간(Latency)을 줄입니다.
2. 코드 유지보수성 향상: 동일한 로직이 여러 곳에 분산되어 있던 것을 ScheduleBatch 클래스로 중앙 집중화했습니다. 이는 향후 페널티 로직이 변경되더라도 한 곳만 수정하면 되므로 버그 발생 가능성을 낮춥니다.
3. 일반적 교훈: 고성능 추론 엔진을 설계할 때, 매 스텝마다 발생하는 호스트-디바이스 간 데이터 전송은 항상 주의 깊게 살펴야 합니다. 특히 torch.tensor 생성 시 device 인자를 직접 전달하는 것은 동기화를 유발할 수 있으므로, pin_memory와 non_blocking 옵션을 활용한 비동기 전송 패턴을 적극적으로 도입해야 합니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.tensor.html
• https://pytorch.org/docs/stable/notes/cuda.html#asynchronous-execution

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