[sglang] SGLang LTX-2 모델을 위한 Breakable CUDA Graph 최적화 분석

PR 링크: sgl-project/sglang#33885 상태: Merged | 변경: +102 / -4

들어가며

최근 sgl-project/sglang 레포지토리에 LTX-2 모델을 위한 Breakable CUDA Graph(BCG) 지원이 추가되었습니다. LTX-2와 같은 고성능 비디오 생성 모델은 H200과 같은 최신 GPU 환경에서도 커널 런칭 오버헤드(launch-bound)로 인해 GPU 활용도가 저하되는 문제가 있었습니다. 본 PR은 LTX-2 파이프라인이 BCG를 활용할 수 있도록 6가지 핵심 제약 사항을 해결하여, H200 기준 E2E 추론 시간을 10.75초에서 6.90초로 약 1.56배 단축했습니다.

코드 분석

1. BCG 허용 목록 및 Warmup 로직 수정 (server_args.py, warmup_request_builder.py)

기존에는 LTX-2가 BCG 허용 목록에 없어 기능이 비활성화되어 있었습니다. 또한, Warmup 시 프레임 수를 제한하거나 불필요한 이미지를 삽입하여 그래프 캡처 시그니처가 실제 서빙 환경과 달라지는 문제가 있었습니다.

# warmup_request_builder.py
if getattr(server_args, "enable_breakable_cuda_graph", False) is True:
    # BCG는 정확한 latent shape에서만 동작하므로 Warmup 시 전체 프레임 수를 사용해야 함
    return num_frames

2. 모델 호출 라우팅 및 H2D 복사 최적화 (denoising.py)

LTX-2는 모델 kwargs를 로컬에서 생성하여 predict_noise 훅을 우회하고 있었습니다. 이를 BCG 러너를 거치도록 수정하고, 그래프 캡처 중 금지된 호스트-디바이스(H2D) 복사를 제거했습니다.

# denoising.py
# Before: torch.tensor(list, device=cuda) -> Illegal in graph capture
# After: Build coords outside the captured region
if video_coords is None:
    video_coords = step.current_model.rope.prepare_video_coords(...)

3. 세그먼트 제한 상향 (runner.py)

LTX-2는 복잡한 구조로 인해 기본 세그먼트 제한(128)을 초과했습니다. 이를 512로 상향하여 캡처 성공률을 높였습니다.

# runner.py
self.max_segments = max(0, _env_int("SGLANG_DIFFUSION_BCG_MAX_SEGMENTS", 512))

왜 이게 좋은가

이 최적화는 CPU 커널 런칭 오버헤드를 제거하는 데 핵심이 있습니다. 프로파일링 결과, CPU 커널 런칭 API 호출이 22,020회에서 2,472회로 대폭 감소했으며, GPU 활용도는 29%에서 95%로 비약적으로 상승했습니다.

교훈:

1. Signature Matching: CUDA Graph는 입력 시그니처가 완벽히 일치해야 합니다. Warmup 시 실제 서빙 환경과 동일한 파라미터(프레임 수, 조건부 입력 등)를 구성하는 것이 필수적입니다.
2. Capture-Illegal Operations: torch.tensor와 같은 동적 H2D 복사는 그래프 캡처 중 런타임 에러를 유발하거나 캡처 실패의 원인이 됩니다. 이러한 연산은 캡처 영역 밖으로 분리해야 합니다.
3. Compute-bound vs Launch-bound: BCG는 연산량이 적고 런칭 오버헤드가 큰 모델에서 효과적입니다. 고해상도(1080p)와 같이 이미 Compute-bound인 환경에서는 torch.compile을 통한 커널 퓨전이 더 유리할 수 있습니다.

결론

이번 PR은 단순히 기능을 활성화하는 것을 넘어, 모델의 내부 구조를 BCG에 맞게 재설계함으로써 실질적인 성능 향상을 이끌어냈습니다. 특히 시그니처 미스 발생 시 로그를 남기는 진단 도구를 추가하여 디버깅 편의성을 높인 점이 인상적입니다.

참고 자료

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

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