[sglang] SGLang, JoyEcho 모델에 Breakable CUDA Graph 적용하여 추론 속도 43% 향상

PR 링크: sgl-project/sglang#38110 상태: Merged | 변경: +10 / -5

들어가며

최근 SGLang 프로젝트에서는 jdopensource/JoyAI-Echo 모델의 추론 성능을 획기적으로 개선하는 Pull Request(PR)가 병합되었습니다. 이 PR의 핵심은 Breakable CUDA Graph (BCG) 기술을 JoyEcho 모델에 적용하여 기존의 심각한 launch/dispatch 병목 현상을 해결하는 것입니다. JoyEcho는 비디오 및 오디오를 생성하는 Diffusion 모델로, 단일 GPU 환경에서 추론 시 GPU 활용률이 매우 낮아 성능 개선의 여지가 많았습니다. 본 글에서는 이 PR이 어떻게 JoyEcho 모델의 추론 속도를 크게 향상시켰는지, 코드 변경 사항을 중심으로 상세히 분석하고 그 의미를 짚어보겠습니다.

코드 변경 분석

이번 PR의 변경 사항은 크게 두 부분으로 나눌 수 있습니다. 첫째, JoyEcho 모델의 추론 로직을 BCG가 지원하는 방식으로 수정하는 것이고, 둘째, BCG 지원 목록에 JoyEcho 모델을 명시적으로 추가하는 것입니다.

1. joy_echo/denoising.py - 추론 로직 수정

가장 핵심적인 변경은 JoyEcho 모델의 Denoising 단계에서 모델의 forward pass를 BCG가 관리하는 함수를 통해 호출하도록 수정한 것입니다. 기존 코드에서는 step.current_model(**model_kwargs)를 직접 호출하여 eager 모드로 실행했습니다. 하지만 BCG를 적용하기 위해서는 이 호출을 BCG가 캡처한 그래프를 재사용하는 _ltx2_call_current_model 함수를 통해 이루어져야 합니다.

Before:

        with self._ltx2_model_forward_context(ctx, step):
            model_video, model_audio = step.current_model(**model_kwargs)

After:

        with self._ltx2_model_forward_context(ctx, step):
            model_video, model_audio = self._ltx2_call_current_model(
                ctx, step, model_kwargs
            )

이 변경은 실제 모델의 연산 자체를 바꾸는 것이 아니라, 연산의 실행 방식을 BCG의 그래프 캡처 및 재플레이 메커니즘에 통합시키는 역할을 합니다. 즉, 모델의 forward pass가 매번 새로 컴파일되고 실행되는 대신, 미리 캡처된 CUDA 그래프를 재사용함으로써 launch overhead를 제거하는 것입니다.

2. server_args.py - BCG 지원 목록 업데이트

BCG는 특정 파이프라인 설정 및 모델 ID에 대해서만 활성화됩니다. 따라서 JoyEcho 모델을 BCG와 함께 사용하려면, 해당 모델과 파이프라인 설정을 BCG 지원 목록에 명시적으로 등록해야 합니다.

BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS 업데이트:

기존에는 jdopensource/joyai-echo와 joyai-echo가 이 목록에 없었습니다. PR에서는 이 두 모델 ID를 추가하여 BCG가 해당 모델에 대해 활성화될 수 있도록 합니다.

BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS 업데이트:

마찬가지로 JoyEchoPipelineConfig가 지원 목록에 없었기 때문에, 이를 추가하여 JoyEcho 파이프라인 설정이 BCG를 지원하도록 합니다.

Unsupported Model Warning 메시지 업데이트:

BCG가 지원되지 않는 모델에 대해 출력되는 경고 메시지도 업데이트되었습니다. 기존 메시지에는 JoyEcho가 포함되어 있지 않았으나, PR 이후에는 JoyEcho가 지원 목록에 포함되었음을 명확히 알 수 있도록 수정되었습니다.

Before (Warning Message Snippet):

            

## 참고 자료
- https://pytorch.org/docs/stable/generated/torch.compile.html

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