[sglang] UniPC 스케줄러에서 GPU 동기화 제거를 통한 성능 최적화 분석

PR 링크: sgl-project/sglang#27440 상태: Merged | 변경: +6 / -6

들어가며

최근 sglang 레포지토리의 Pull Request(PR)에서는 UniPC 스케줄러의 성능 병목 현상을 해결하기 위한 중요한 개선이 이루어졌습니다. UniPC 스케줄러는 Stable Diffusion과 같은 이미지 생성 모델에서 노이즈 제거 과정을 효율적으로 수행하는 데 사용되는 알고리즘입니다. 이 PR의 핵심은 기존 코드에서 불필요하게 발생하던 GPU 동기화(synchronization) 오버헤드를 제거하여 전체적인 처리 속도를 향상시키는 것입니다.

특히, Cosmos3 모델을 H200 GPU에서 실행했을 때 발견된 aten::item / _local_scalar_dense / cudaStreamSynchronize와 같은 함수 호출들은 GPU 연산 결과를 CPU로 가져오거나, GPU 간의 연산 흐름을 맞추기 위해 발생하는 오버헤드였습니다. 이러한 동기화는 GPU의 병렬 처리 능력을 온전히 활용하지 못하게 만드는 주요 원인 중 하나입니다. 본 글에서는 이 PR에서 어떤 코드가 어떻게 변경되었고, 왜 이러한 변경이 성능 향상으로 이어졌는지 상세히 분석해보겠습니다.

코드 분석

이번 PR의 변경 사항은 주로 python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py 파일 내의 UniPC 스케줄러 구현부에 집중되어 있습니다. 특히, multistep_uni_p_bh_update와 multistep_uni_c_bh_update 함수 내에서 스칼라 값들을 텐서로 변환하는 방식이 개선되었습니다.

1. multistep_uni_p_bh_update 함수

이 함수는 UniPC 스케줄러의 예측(predictor) 단계를 처리합니다. 변경 전후의 코드를 비교해보겠습니다.

Before:

-        rks.append(1.0)
-        rks = torch.tensor(rks, device=device)

After:

+        rks.append(torch.ones((), dtype=h.dtype, device=h.device))
+        rks = torch.stack(rks).to(device=device)

Before:

-        b = torch.tensor(b, device=device)

After:

+        b = torch.stack(b).to(device=device)

기존 코드에서는 rks와 b 리스트에 스칼라 값(1.0 또는 계산된 값)을 추가한 뒤, torch.tensor(list, device=device)를 사용하여 이 리스트를 GPU 상의 텐서로 변환했습니다. torch.tensor() 함수는 리스트의 각 요소를 처리할 때, 만약 해당 요소가 CPU에 있다면 GPU로 복사하고, 이미 GPU에 있더라도 내부적으로 스칼라 값을 추출하고 새로운 텐서를 생성하는 과정에서 GPU 커널 실행 및 동기화를 유발할 수 있습니다. 특히 torch.tensor(list_of_gpu_scalars)와 같은 형태는 각 스칼라를 개별적으로 처리하며 암묵적인 CPU 동기화를 일으킬 가능성이 높습니다.

개선된 코드에서는 먼저 torch.ones((), dtype=h.dtype, device=h.device)와 같이 GPU 상에서 직접 1.0 값을 가지는 0차원 텐서를 생성하여 리스트에 추가합니다. 이후 torch.stack(rks).to(device=device)를 사용하여 리스트에 쌓인 GPU 텐서들을 하나로 합칩니다. torch.stack은 이미 GPU에 있는 텐서들을 효율적으로 결합하며, .to(device=device)는 최종 텐서가 올바른 장치에 있는지 확인하는 역할을 합니다. 이 방식은 GPU에 이미 존재하는 스칼라 값들을 CPU로 옮기지 않고 GPU 내에서 바로 텐서로 조립하므로, 불필요한 GPU-CPU 동기화 및 데이터 전송을 피할 수 있습니다.

2. multistep_uni_c_bh_update 함수

이 함수는 UniPC 스케줄러의 교정(corrector) 단계를 처리하며, multistep_uni_p_bh_update와 동일한 방식으로 코드가 수정되었습니다.

Before:

-        rks.append(1.0)
-        rks = torch.tensor(rks, device=device)

After:

+        rks.append(torch.ones((), dtype=h.dtype, device=h.device))
+        rks = torch.stack(rks).to(device=device)

Before:

-        b = torch.tensor(b, device=device)

After:

+        b = torch.stack(b).to(device=device)

마찬가지로, 스칼라 값들을 torch.tensor()로 변환하는 대신, GPU 상에서 직접 생성된 0차원 텐서를 torch.stack으로 결합함으로써 GPU 동기화 오버헤드를 제거했습니다. 이는 예측 단계와 동일한 원리로 성능 향상에 기여합니다.

왜 이게 좋은가?

이번 PR의 핵심 개선 사항은 GPU 연산 과정에서 발생하는 불필요한 동기화 비용을 제거하는 것입니다. PR 설명에 따르면, H200 GPU에서 Cosmos3 모델을 사용한 벤치마크 결과는 다음과 같습니다:

Metric	Before	After
Total duration	57457.98 ms	56509.18 ms
Denoising stage	53571.69 ms	52634.58 ms
Peak reserved memory	43792 MB	43792 MB

벤치마크 결과, 총 처리 시간(Total duration)은 약 1.55% (57457.98ms -> 56509.18ms), Denoising 스테이지 시간은 약 1.75% (53571.69ms -> 52634.58ms) 단축되었습니다. 메모리 사용량은 동일하게 유지되었습니다.

이러한 성능 향상은 다음과 같은 이유로 매우 의미 있습니다:

1. GPU 활용률 극대화: GPU는 병렬 처리에 특화된 하드웨어입니다. CPU와의 잦은 동기화는 GPU가 다음 연산을 기다리게 만들어 병렬 처리의 이점을 상쇄시킵니다. GPU 상에서 모든 연산을 완료하고 필요한 경우에만 결과를 CPU로 가져오는 것이 GPU 활용률을 높이는 핵심입니다.
2. 암묵적 동기화 제거: torch.tensor(list_of_scalars)와 같은 코드는 개발자가 명시적으로 동기화를 호출하지 않더라도 내부적으로 발생할 수 있습니다. 이러한 암묵적인 동기화는 코드 분석을 어렵게 하고 예상치 못한 성능 저하를 유발할 수 있습니다. 명시적으로 torch.stack과 같은 GPU 네이티브 연산을 사용함으로써 코드의 의도를 명확히 하고 최적의 성능을 달성할 수 있습니다.
3. 일반화 가능한 교훈: 이 PR은 단순히 특정 모델의 스케줄러를 최적화하는 것을 넘어, GPU 컴퓨팅에서 흔히 발생하는 최적화 패턴을 보여줍니다. 즉, GPU 상에서 생성된 데이터를 CPU로 불필요하게 옮기거나, 스칼라 값들을 텐서로 변환할 때 GPU 네이티브 연산을 우선적으로 고려해야 한다는 일반적인 교훈을 제공합니다.

리뷰어인 mickqian 님께서도 tag-and-rerun-ci를 통해 변경 사항의 검증을 요청하셨으며, 이는 코드 변경이 의도한 대로 작동하는지, 그리고 CI/CD 파이프라인에서 정상적으로 처리되는지를 확인하는 중요한 절차입니다. 이러한 검증 과정을 거쳐 안정성이 확보된 최적화는 실제 서비스에 적용될 때 더욱 신뢰를 줄 수 있습니다.

결론

sglang의 UniPC 스케줄러에 적용된 이번 GPU 동기화 제거 최적화는 작지만 매우 효과적인 개선입니다. torch.tensor 대신 torch.stack을 사용하여 GPU 상에서 데이터를 직접 조립함으로써 불필요한 동기화 오버헤드를 제거하고, 결과적으로 이미지 생성 속도를 향상시켰습니다. 이는 GPU 컴퓨팅의 기본 원리를 잘 이해하고 적용한 모범 사례이며, 향후 유사한 성능 병목 현상을 해결하는 데 중요한 참고 자료가 될 것입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.stack.html
• https://pytorch.org/docs/stable/generated/torch.tensor.html
• https://pytorch.org/docs/stable/generated/torch.ones.html

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