본문으로 건너뛰기

[vllm] vLLM 멀티모달 스케줄러 오버헤드 최적화: Python List 캐싱으로 27% 성능 향상

PR 링크: vllm-project/vllm#40143 상태: Merged | 변경: +None / -None

들어가며

vLLM은 대규모 언어 모델(LLM) 추론을 위한 고성능 서빙 프레임워크로, 특히 멀티모달(Multimodal) 워크로드 처리 시 효율적인 스케줄링이 중요합니다. 멀티모달 모델은 텍스트뿐만 아니라 이미지, 오디오 등 다양한 형태의 입력을 처리하며, 이러한 입력들은 임베딩(embedding) 형태로 변환되어 모델에 전달됩니다. 이 과정에서 get_num_embeds와 같은 함수가 반복적으로 호출되는데, 이 함수의 비효율적인 구현은 전체 시스템의 성능 저하로 이어질 수 있습니다.

이번 PR은 vLLM의 멀티모달 스케줄러에서 발생하는 오버헤드를 줄여 전체 처리량을 크게 향상시킨 최적화에 대한 내용입니다. 특히 get_num_embeds 함수 호출 시 발생하는 불필요한 CPU Torch Tensor 연산 및 Python 객체 변환 비용을 제거하여, 스케줄링이 GPU 작업과 완전히 겹치지 못하고 유휴 시간(idle bubbles)을 발생시키는 문제를 해결합니다.

코드 분석: 무엇이 어떻게 개선되었나?

핵심 변경사항은 vllm/multimodal/inputs.py 파일의 PlaceholderRange 클래스 내 embeds_cumsum 속성과 get_num_embeds, get_embeds_indices_in_range 메서드에 있습니다.

vllm/multimodal/inputs.py

1. embeds_cumsum 속성 캐싱 방식 변경

이전에는 embeds_cumsumtorch.Tensor를 반환하도록 되어 있었습니다. torch.Tensor는 강력한 기능을 제공하지만, Python 스칼라로 변환하거나 접근할 때 추가적인 오버헤드가 발생할 수 있습니다. 특히 이 값이 자주 접근될 경우, Torch 객체의 생성 및 소멸, 그리고 Python 스칼라로의 변환 과정에서 CPU 비용이 누적됩니다.

Before:

class PlaceholderRange:
    # ...
    @cached_property
    def embeds_cumsum(self) -> torch.Tensor | None:
        return None if self.is_embed is None else self.is_embed.cumsum(dim=0)

After:

class PlaceholderRange:
    # ...
    @cached_property
    def embeds_cumsum(self) -> list[int] | None:
        # python list so python indexing avoids torch C++ overhead/conversions/deallocs
        return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist()

개선점: embeds_cumsum의 결과가 torch.Tensor 대신 list[int]로 캐싱되도록 변경되었습니다. torch.Tensor.tolist() 메서드를 사용하여 Torch Tensor의 내용을 Python 리스트로 변환합니다. 이 변경으로 인해 embeds_cumsum에 접근할 때마다 Torch C++ 오버헤드, Torch 객체와 Python 객체 간의 변환, 그리고 불필요한 Torch 객체 할당/해제 비용을 피할 수 있게 됩니다. Python 리스트는 Python 인터프리터 내에서 훨씬 가볍게 처리될 수 있습니다.

2. get_num_embeds 메서드 최적화

get_num_embedsembeds_cumsum의 마지막 요소를 가져와 임베딩의 총 개수를 반환하는 함수입니다. 이전에는 torch.Tensor의 마지막 요소를 가져와 int()로 명시적 형 변환을 수행했습니다. 이 과정에서 torch.Tensor 객체에 대한 접근과 Python 스칼라로의 변환 비용이 발생했습니다.

Before:

    def get_num_embeds(self) -> int:
        if self.embeds_cumsum is None:
            return self.length

        return int(self.embeds_cumsum[-1])

After:

    def get_num_embeds(self) -> int:
        if self.embeds_cumsum is None:
            return self.length

        return self.embeds_cumsum[-1] if self.embeds_cumsum else 0

개선점: embeds_cumsum이 이제 Python 리스트이므로, int() 형 변환 없이 바로 self.embeds_cumsum[-1]로 마지막 요소에 접근할 수 있습니다. 이는 Python 리스트 인덱싱이 Torch Tensor 인덱싱보다 훨씬 가볍기 때문에 성능상 이점을 제공합니다. 또한, self.embeds_cumsum이 비어 있을 경우를 대비하여 if self.embeds_cumsum else 0와 같은 방어 로직이 추가되었습니다.

3. get_embeds_indices_in_range 메서드 최적화

이 메서드 또한 embeds_cumsum의 요소를 사용하여 임베딩 인덱스 범위를 계산합니다. 이전에는 각 접근마다 int() 형 변환을 수행했습니다.

Before:

    def get_embeds_indices_in_range(
        self, start_idx: int, end_idx: int
    ) -> tuple[int, int]:
        # ...
        if self.embeds_cumsum is None:
            return start_idx, end_idx

        embeds_start_idx = (
            int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0
        )
        embeds_end_idx = int(self.embeds_cumsum[end_idx - 1])

        return embeds_start_idx, embeds_end_idx

After:

    def get_embeds_indices_in_range(
        self, start_idx: int, end_idx: int
    ) -> tuple[int, int]:
        # ...
        if self.embeds_cumsum is None:
            return start_idx, end_idx

        embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0
        embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0

        return embeds_start_idx, embeds_end_idx

개선점: get_num_embeds와 마찬가지로, embeds_cumsum이 Python 리스트가 됨에 따라 불필요한 int() 형 변환을 제거했습니다. 이는 코드 가독성을 높이고 런타임 오버헤드를 줄이는 효과를 가져옵니다. 또한 embeds_end_idx 계산 시 end_idx > 0 조건이 추가되어 잠재적인 인덱스 에러를 방지합니다.

tests/multimodal/test_inputs.py

테스트 코드도 변경된 embeds_cumsum의 반환 타입에 맞춰 수정되었습니다. torch.Tensor 비교 대신 Python 리스트 비교로 변경되었습니다.

Before:

    # ...
    assert torch.equal(pr.embeds_cumsum, expected)
    # ...

After:

    # ...
    assert pr.embeds_cumsum == expected
    # ...

개선점: 변경된 구현에 맞춰 테스트 코드를 업데이트하여, 새로운 list[int] 반환 타입에 대한 정확성을 보장합니다. 이는 리팩토링 시 테스트 코드의 중요성을 보여줍니다.

왜 이게 좋은 최적화인가?

이 최적화는 멀티모달 워크로드에서 스케줄러의 CPU 오버헤드를 줄이는 데 초점을 맞춥니다. _try_schedule_encoder_inputs 함수는 배치 내 모든 요청을 반복하며, 각 멀티모달 피처에 대해 get_num_embeds를 호출합니다. 이 함수는 embeds_cumsum 속성에 의존하는데, 이 속성이 torch.Tensor로 캐시되어 있었기 때문에 다음과 같은 문제가 발생했습니다:

  1. Torch Tensor 오버헤드: torch.Tensor는 GPU 연산에 최적화되어 있지만, CPU에서 작은 스칼라 값을 자주 접근하고 Python 스칼라로 변환하는 과정은 비효율적입니다. Torch 객체의 생성, 소멸, 그리고 Python 객체와의 상호작용은 상당한 CPU 시간을 소모합니다.
  2. 불필요한 형 변환: int(self.embeds_cumsum[-1])와 같이 매번 명시적으로 형 변환을 수행하는 것도 작은 오버헤드를 발생시킵니다.

이러한 오버헤드는 특히 멀티모달 아이템이 많거나 임베딩 크기가 큰 워크로드에서 심화되어, 스케줄링이 GPU 작업과 완전히 겹치지 못하고 'idle bubbles'을 생성하게 됩니다. 즉, GPU가 놀고 있는 시간이 발생하여 전체 처리량이 저하되는 것입니다.

성능 수치

PR에 포함된 벤치마크 결과는 이 최적화의 효과를 명확히 보여줍니다. Gemma-4-E4B 모델에 128개의 요청과 32개의 이미지(1024x1536)를 사용하는 멀티모달 워크로드에서 다음과 같은 개선이 있었습니다.

Metric baseline This PR Δ
Duration (s) 413.16 325.41 −21.2 %
Request throughput (req/s) 0.31 0.39 +26.9 %
Output tok/s 2 478.5 3 146.8 +27.0 %
Total tok/s 5 093.9 6 467.5 +27.0 %
Peak output tok/s 6 158 9 088 +47.6 %
Mean TTFT (s) 137.9 124.3 −9.8 %
Mean TPOT (ms) 33.81 24.58 −27.3 %

주요 개선 사항:

  • Duration (총 소요 시간): 약 21.2% 감소
  • Request throughput (요청 처리량): 약 26.9% 증가
  • Output tok/s (출력 토큰 처리량): 약 27.0% 증가
  • Mean TPOT (평균 토큰 당 처리 시간): 약 27.3% 감소

이러한 수치들은 embeds_cumsum을 Python 리스트로 캐싱하고 관련 함수에서 불필요한 Torch Tensor 연산을 제거함으로써 스케줄러의 CPU 오버헤드가 크게 줄어들었음을 입증합니다. 스케줄링 오버헤드가 감소하면서 GPU가 더 효율적으로 활용되고, 결과적으로 전체 시스템의 처리량이 향상되었습니다.

일반적 교훈

이 최적화는 다음과 같은 중요한 교훈을 제공합니다:

  1. 적절한 자료 구조 선택: 데이터의 사용 패턴에 따라 가장 효율적인 자료 구조를 선택하는 것이 중요합니다. GPU 연산이 필요 없는 단순한 스칼라 값의 집합이라면 torch.Tensor보다 Python list가 훨씬 효율적일 수 있습니다.
  2. CPU/GPU 경계 최적화: 딥러닝 시스템에서 CPU와 GPU 간의 데이터 전송 및 객체 변환은 종종 성능 병목의 원인이 됩니다. 이 경계를 넘나드는 오버헤드를 최소화하는 것이 중요합니다.
  3. 프로파일링의 중요성: PR 설명에 포함된 프로파일링 이미지는 최적화 전후의 CPU 사용 패턴 변화를 시각적으로 보여줍니다. 이처럼 실제 워크로드에서 병목 지점을 정확히 식별하는 것이 효과적인 최적화의 첫걸음입니다.
  4. 캐싱의 현명한 활용: cached_property는 계산 비용이 높은 속성을 한 번만 계산하고 재사용하는 데 유용하지만, 캐싱되는 객체의 특성(예: torch.Tensor vs list) 또한 성능에 큰 영향을 미칠 수 있습니다.

이 PR은 vLLM과 같은 고성능 추론 시스템에서 작은 코드 변경이 어떻게 전체 시스템 성능에 큰 영향을 미칠 수 있는지 보여주는 좋은 예시입니다. 특히 멀티모달과 같이 복잡한 워크로드에서는 이러한 세밀한 최적화가 사용자 경험과 서비스 비용에 직접적인 영향을 미칩니다.

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

댓글

관련 포스트

PR Analysis 의 다른글