[sglang] Intel XPU를 위한 고성능 그래프 캡처 및 리플레이 지원 구현

PR 링크: sgl-project/sglang#29053 상태: Merged | 변경: +886 / -43

들어가며

LLM 추론 서비스에서 호스트 측의 커널 실행 오버헤드는 전체 지연 시간(latency)을 결정짓는 중요한 요소입니다. 기존의 CUDA Graph는 이러한 오버헤드를 획기적으로 줄여주는 기술로 널리 사용되어 왔습니다. 이번 PR은 Intel XPU 환경에서도 이와 유사한 그래프 캡처 및 리플레이 기능을 구현하여, Decode 단계의 Full-graph와 Prefill 단계의 tc_piecewise 최적화를 가능하게 했습니다.

코드 분석

1. XPUGraphRunner 및 백엔드 구현

새로운 XPUGraphRunner 클래스를 도입하여 기존 DecodeCudaGraphRunner의 로직을 XPU 환경에 맞게 확장했습니다. ModelRunner와 통합되어 Decode 단계에서 torch.xpu.XPUGraph를 활용한 그래프 캡처를 수행합니다.

# python/sglang/srt/compilation/xpu_piecewise_backend.py
class XPUPiecewiseBackend(CUDAPiecewiseBackend):
    # ... (생략)
    xpugraph = torch.xpu.XPUGraph()
    with torch.xpu.graph(xpu_graph=xpugraph, pool=self.graph_pool, stream=stream):
        output = entry.runnable(*args)

2. 분산 처리 최적화 (Dynamo Capture)

Dynamo가 all_reduce나 all_gather를 저수준의 _c10d_functional로 분해하지 않도록, inplace_all_reduce와 같은 불투명(opaque) 커스텀 연산으로 라우팅했습니다. 이는 그래프 캡처 도중 분산 연산이 깨지는 문제를 방지합니다.

3. Weak Ref Tensor 지원

sgl_kernel을 통해 XPU 환경에서도 weak_ref_tensor를 사용할 수 있도록 하여, 그래프 캡처 시 메모리 관리 효율성을 높였습니다.

# python/sglang/srt/compilation/weak_ref_tensor.py
if is_cuda() or is_hip() or is_musa() or is_xpu():
    from sgl_kernel import weak_ref_tensor

왜 이게 좋은가

이번 최적화의 핵심은 호스트 오버헤드 감소입니다. 반복적인 커널 호출 대신 미리 캡처된 그래프를 리플레이함으로써, 특히 배치 크기가 작거나 추론 속도가 빠른 모델에서 유의미한 성능 향상을 기대할 수 있습니다.

• 일반적 교훈: 하드웨어 가속기(XPU, GPU 등)에서 그래프 캡처는 정적인 계산 그래프를 가진 모델에 매우 효과적입니다. 특히 torch.compile과 결합할 때, 컴파일된 커널을 그래프 내부에 캡처함으로써 런타임 오버헤드를 최소화할 수 있습니다.
• 리뷰어 피드백: 리뷰 과정에서 sgl-kernel-xpu와의 의존성 문제가 논의되었으며, weak_ref_tensor의 적절한 등록을 통해 안정성을 확보했습니다.

결론

Intel XPU 사용자는 이제 --cuda-graph-backend-prefill tc_piecewise 옵션을 통해 Prefill 단계의 성능을 최적화할 수 있게 되었습니다. 이는 SGLang이 다양한 하드웨어 백엔드에서 일관된 고성능 추론을 제공하기 위한 중요한 이정표입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.xpu.XPUGraph.html
• https://docs.sglang.ai/developer_guide/contribution_guide.html

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