[sglang] SGLang FLUX 3 Action 성능 최적화: CUDA Graphs 도입을 통한 Host Launch 오버헤드 제거

PR 링크: sgl-project/sglang#42171 상태: Merged | 변경: +371 / -62

들어가며

최신 멀티모달 모델인 FLUX 3 Action을 실행할 때, GPU 연산 시간보다 Host에서 커널을 실행(Launch)하기 위해 소요되는 시간이 더 큰 비중을 차지하는 경우가 많습니다. 특히 VAE 프레임 인코딩 단계에서는 약 2,200개의 커널이 실행되는데, 이때 GPU 작업 시간은 9.3ms인 반면 Host Launch 시간은 18ms에 달합니다. 본 PR은 이러한 Host-bound 문제를 해결하기 위해 CUDA Graphs를 도입하여, 반복적인 커널 실행 시퀀스를 그래프로 캡처함으로써 오버헤드를 획기적으로 줄였습니다.

코드 분석

1. CUDA Graphs 도입 (flux3_action.py)

Flux3ActionPipelineConfig에 enable_cuda_graph 플래그를 추가하고, cuda_graph_enabled 메서드를 통해 단일 GPU 환경에서만 그래프를 활성화하도록 설계했습니다.

# Before
# (기존에는 매 단계마다 eager 모드로 커널을 개별 실행)

# After
class Flux3ActionPipelineConfig(PipelineConfig):
    enable_cuda_graph: bool = True
    def cuda_graph_enabled(self, server_args: Any, component: str) -> bool:
        return self.enable_cuda_graph and server_args.num_gpus == 1 and not server_args.should_start_component_on_cpu(component)

2. 그래프 캡처 및 재실행 (VLATensorGraphRunner)

Flux3ActionObservationEncodingStage와 Flux3ActionDenoisingStage에 VLATensorGraphRunner를 통합했습니다. 이를 통해 관측 인코딩(Observation encoding)과 디노이징 단계의 커널 시퀀스를 캡처합니다.

# 관측 인코딩 단계에서의 그래프 캡처 적용 예시
self._graph = VLATensorGraphRunner(
    "FLUX 3 observation",
    enabled=use_cuda_graph,
    max_entries=config.cuda_graph_max_entries,
)

3. Bit-Exact Fusion 커널의 안전한 캡처 (flux3.py)

기존에는 그래프 캡처 중에는 최적화된 fused 커널(LN+modulate, SwiGLU 등)이 비활성화되었으나, 이제는 _can_verify()를 통해 그래프 캡처 중에도 안전하게 검증된 커널을 사용할 수 있도록 로직을 개선했습니다.

# Before
if _LN_MODULATE.disabled or not _eager_fast_path_allowed(x):
    return (1 + scale) * norm(x) + shift

# After
if not _LN_MODULATE.is_verified(sig) and not _can_verify():
    return (1 + scale) * norm(x) + shift

왜 이게 좋은가

이번 최적화는 특히 Host Launch 오버헤드가 큰 환경에서 강력한 성능 향상을 보여줍니다. RTX 5090 환경에서 벤치마크 결과, sd-fp8r 변형 모델의 경우 Eager 모드 대비 11.6%의 지연 시간(Latency) 감소를 기록했습니다.

• 성능 수치: 관측 인코딩 그래프는 모든 변형에서 9~15ms의 지연 시간을 절감하며, 디노이징 그래프는 커널 호출 횟수에 따라 최대 19ms까지 성능을 개선합니다.
• 메모리 오버헤드: 그래프 메모리 풀 사용으로 인해 약 210~310 MiB의 추가 메모리가 사용되지만, 이는 전체 파이프라인 성능을 고려할 때 충분히 감수할 만한 트레이드오프입니다.
• 교훈: GPU 연산이 가벼운 모델이나 단계에서는 커널 Launch 오버헤드가 전체 성능을 결정짓는 핵심 요소가 됩니다. torch.cuda.graphs를 활용하여 반복적인 커널 시퀀스를 캡처하는 것은 고성능 추론 엔진 설계의 필수적인 패턴입니다.

결론

이번 PR은 단순히 커널을 최적화하는 것을 넘어, Host-GPU 간의 통신 병목을 제거하는 구조적 접근을 취했습니다. 특히 fused quantization과 함께 사용했을 때 시너지가 극대화되어, 총 29.6%의 성능 향상을 이끌어낸 점이 인상적입니다.

참고 자료

• 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를 기반으로 작성했습니다.