본문으로 건너뛰기

[vllm] [vLLM] ROCm 환경에서 4바이트 스칼라 할당이 유발하는 성능 병목 해결하기

PR 링크: vllm-project/vllm#56170 상태: Merged | 변경: +4 / -1

들어가며

고성능 LLM 서빙 엔진인 vLLM에서는 아주 작은 연산 하나도 전체 처리량(Throughput)과 지연 시간(Latency)에 큰 영향을 미칠 수 있습니다. 특히 GPU 연산이 진행되는 동안 CPU(Host)가 멈추지 않고 다음 작업을 계속해서 큐에 쌓아두는 '비동기 실행'은 성능 최적화의 핵심입니다.

최근 vLLM의 MiniMax M3 모델 구현체에서 흥미로운 성능 개선 PR이 올라왔습니다. ROCm(AMD GPU) 환경에서 단 4바이트의 정수 값을 0으로 초기화하는 코드가 호스트의 실행을 블로킹(Blocking)하여 수백 밀리초의 지연을 발생시키고 있었고, 이를 아주 간단한 트릭으로 해결한 사례입니다. 왜 x[0] = 0이라는 평범한 코드가 문제가 되었는지, 그리고 어떻게 해결했는지 살펴보겠습니다.

문제의 배경: Scalar Upload와 Host-Device 동기화

MiniMax M3 모델은 Sparse Attention을 사용하며, 이 과정에서 KV 캐시의 누적 길이를 관리하는 prefill_cu_seqlens_k 텐서를 생성합니다. 이 텐서의 첫 번째 요소는 항상 0이어야 합니다.

기존 코드는 다음과 같았습니다.

# Before
prefill_cu_seqlens_k = torch.empty(
    num_prefills + 1, dtype=torch.int32, device=seq_lens.device
)
prefill_cu_seqlens_k[0] = 0

여기서 prefill_cu_seqlens_k[0] = 0은 Python의 스칼라 값 0을 GPU 메모리에 있는 텐서의 첫 번째 위치에 할당하는 작업입니다.

ROCm 환경에서 PyTorch는 이 작업을 수행할 때 내부적으로 hipMemcpyWithStream API를 호출합니다. 문제는 이 API의 동작 방식입니다. 만약 해당 스트림(Compute Stream)에 이전에 요청된 연산들이 아직 실행 중이라면, 이 API는 이전 작업들이 완료될 때까지 호스트(CPU)를 대기시킨 후 4바이트를 복사하고 제어권을 반환합니다. 즉, 아주 작은 데이터를 복사하기 위해 CPU가 GPU의 작업 완료를 기다리는 동기화(Synchronization) 현상이 발생하는 것입니다.

코드 분석: 비동기 연산으로의 전환

이 문제를 해결하기 위해 작성된 PR의 변경 사항은 매우 직관적입니다.

vllm/models/minimax_m3/common/sparse_attention.py 수정 사항

             prefill_cu_seqlens_k = torch.empty(
                 num_prefills + 1, dtype=torch.int32, device=seq_lens.device
             )
-            prefill_cu_seqlens_k[0] = 0
+            if current_platform.is_rocm():
+                prefill_cu_seqlens_k[:1].zero_()
+            else:
+                prefill_cu_seqlens_k[0] = 0
             torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:])

왜 zero_()가 더 나은가?

  1. Before (prefill_cu_seqlens_k[0] = 0):

    • CPU에 있는 값 0을 GPU로 보내는 'H2D(Host-to-Device) Copy' 작업입니다.
    • ROCm에서는 이 과정에서 hipMemcpyWithStream이 호출되며, 스트림에 쌓인 이전 작업이 끝날 때까지 CPU를 멈추게 합니다(Blocking).

  2. After (prefill_cu_seqlens_k[:1].zero_()):

    • 텐서의 첫 번째 요소를 슬라이싱([:1])한 뒤, GPU 내부에서 해당 영역을 0으로 채우는 zero_() 커널을 실행합니다.
    • 이는 'Device-only' 연산입니다. CPU는 GPU에게 "이 영역을 0으로 채워라"라는 명령만 던지고 즉시 다음 코드로 넘어갑니다(Non-blocking).
    • 결과적으로 스트림의 순서는 유지되면서 호스트의 대기 시간은 사라집니다.

왜 이게 좋은 최적화인가?

1. 극적인 성능 개선 수치

PR 작성자가 공유한 마이크로 벤치마크 결과에 따르면, 이전 작업이 큐에 쌓여 있는 상태에서 호스트의 명령 제출(Submission) 시간이 다음과 같이 변했습니다.

  • 기존 (Scalar Assignment): 8.360 ms
  • 개선 (In-place zero_): 0.013 ms

무려 640배 이상의 속도 향상이 있었습니다. 실제 TP4(Tensor Parallelism 4) 환경의 추적 결과에서는 각 Rank당 약 700~800ms의 호스트 시간이 절약되는 것이 확인되었습니다. 이는 모델의 전체 추론 지연 시간을 줄이는 데 직접적인 기여를 합니다.

2. 하드웨어 특성에 맞는 최적화

이 문제는 모든 플랫폼에서 동일하게 발생하는 것은 아닙니다. CUDA 환경에서는 PyTorch의 스칼라 할당이 좀 더 효율적으로 처리되거나 비동기적으로 동작할 수 있지만, ROCm의 특정 구현 경로에서는 블로킹이 발생했습니다. current_platform.is_rocm() 분기 처리를 통해 타 플랫폼의 안정성을 해치지 않으면서 ROCm 환경의 병목만 정확히 타격했습니다.

3. 일반적인 교훈: "GPU 텐서에 직접 스칼라를 대입하지 마라"

이 PR은 딥러닝 엔지니어들에게 중요한 교훈을 줍니다. 루프 내부나 빈번하게 호출되는 코드 경로(Hot path)에서 GPU 텐서의 개별 요소에 Python 스칼라 값을 대입하는 것은 잠재적인 성능 폭탄이 될 수 있습니다. 대신 다음과 같은 방법을 지향해야 합니다.

  • tensor.fill_(value)
  • tensor.zero_()
  • 미리 준비된 GPU 텐서 간의 복사

결론

이번 최적화는 코드 단 한 줄의 변경으로 수백 밀리초의 성능 이득을 얻은 훌륭한 사례입니다. 대규모 분산 학습이나 서빙 환경에서는 아주 작은 동기화 지연이 전체 클러스터의 효율을 떨어뜨릴 수 있습니다. 프로파일링 도구(Trace)를 통해 hipMemcpyWithStream과 같은 API가 호스트를 붙잡고 있는지 면밀히 관찰하는 습관이 왜 중요한지 다시 한번 일깨워줍니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글