[vllm] vLLM, Cohere 임베딩 바이너리 압축 성능 4배 개선: NumPy를 활용한 최적화 분석

2026년 8월 14일수정: 2026년 8월 14일

PR 링크: vllm-project/vllm#52277 상태: Merged | 변경: +57 / -24

들어가며

vLLM은 대규모 언어 모델을 효율적으로 서빙하기 위한 오픈소스 라이브러리입니다. 특히 텍스트 임베딩 생성 시 다양한 포맷을 지원하는데, Cohere API의 /v2/embed 엔드포인트에서 제공하는 binary 및 ubinary 타입은 임베딩 벡터를 효율적으로 압축하여 전송합니다. 이전 버전의 vLLM에서는 이 바이너리 임베딩을 생성하는 과정에서 파이썬의 중첩된 루프를 사용하여 성능 병목이 발생했습니다. 본 기술 블로그 글에서는 이 문제를 해결하고 성능을 획기적으로 개선한 GitHub PR(#2753)의 코드 변경 사항을 분석하고, 왜 이러한 변경이 성능 향상에 기여했는지 심층적으로 살펴보겠습니다.

코드 변경 분석

이번 PR의 핵심은 vllm/entrypoints/pooling/embed/protocol.py 파일 내 _pack_binary_embeddings 함수의 구현을 최적화하는 것입니다. 기존의 파이썬 루프 기반 구현을 NumPy의 packbits 함수를 활용하도록 변경하여 성능을 개선했습니다.

`vllm/entrypoints/pooling/embed/protocol.py`

기존 코드 (Before)

기존 _pack_binary_embeddings 함수는 입력된 float 임베딩 벡터를 순회하며 각 요소의 부호에 따라 비트를 설정하고, 8개 비트씩 모아 하나의 바이트로 만드는 작업을 파이썬 리스트와 루프를 사용하여 수행했습니다.

def _pack_binary_embeddings(
    float_embeddings: Sequence[Sequence[float]],
    signed: bool,
) -> list[list[int]]:
    """Bit-pack float embeddings: positive -> 1, negative -> 0.

    Each bit is shifted left by ``7 - idx%8``, and every 8 bits are packed
    into one byte.
    """
    result: list[list[int]] = []
    for embedding in float_embeddings:
        dim = len(embedding)
        if dim % 8 != 0:
            raise ValueError(
                "Embedding dimension must be a multiple of 8 for binary "
                f"embedding types, but got {dim}."
            )
        packed_len = dim // 8
        packed: list[int] = []
        byte_val = 0
        for idx, value in enumerate(embedding):
            bit = 1 if value >= 0 else 0
            byte_val += bit << (7 - idx % 8)
            if (idx + 1) % 8 == 0:
                if signed:
                    byte_val -= _UNSIGNED_TO_SIGNED_DIFF
                packed.append(byte_val)
                byte_val = 0
        assert len(packed) == packed_len
        result.append(packed)
    return result

이 코드는 각 임베딩 벡터에 대해 반복하고, 각 벡터의 차원 내에서도 다시 반복하며 비트를 계산하고 누적하는 구조입니다. 특히 byte_val += bit << (7 - idx % 8)와 같은 연산은 파이썬 레벨에서 수행되므로, 대량의 데이터를 처리할 때 상당한 오버헤드를 발생시킵니다.

변경 후 코드 (After)

NumPy를 도입하여 코드가 훨씬 간결하고 효율적으로 변경되었습니다. 먼저 입력 float_embeddings를 NumPy 배열로 변환하고, np.packbits 함수를 사용하여 비트 압축을 수행합니다. array >= 0 연산은 배열의 모든 요소에 대해 브로드캐스팅되어 부호에 따른 불리언 배열을 생성하고, np.packbits는 이 불리언 배열을 입력받아 각 행별로 8비트씩 묶어 바이트로 압축합니다. signed 플래그에 따라 후처리도 NumPy 연산으로 처리됩니다.

import numpy as np

def _pack_binary_embeddings(
    float_embeddings: Sequence[Sequence[float]],
    signed: bool,
) -> list[list[int]]:
    """Bit-pack float embeddings: positive -> 1, negative -> 0.

    Bits are packed MSB-first, eight per byte.
    """
    if not float_embeddings:
        return []

    array = np.asarray(float_embeddings, dtype=np.float64)
    if array.ndim != 2:
        raise ValueError(
            f"Expected a 2D batch of embeddings, but got {array.ndim}D input."
        )

    dim = array.shape[1]
    if dim % 8 != 0:
        raise ValueError(
            "Embedding dimension must be a multiple of 8 for binary "
            f"embedding types, but got {dim}."
        )

    packed = np.packbits(array >= 0, axis=-1)
    if signed:
        packed = packed.astype(np.int16) - _UNSIGNED_TO_SIGNED_DIFF
    return packed.tolist()

주요 변경 사항은 다음과 같습니다:

NumPy 배열 변환: np.asarray(float_embeddings, dtype=np.float64)를 통해 입력 데이터를 NumPy 배열로 변환합니다. 이는 후속 연산의 효율성을 높입니다.
차원 검증: array.ndim != 2 검사를 통해 입력이 2차원 배열인지 확인합니다. 기존 코드에서는 암묵적으로 처리되었을 수 있는 부분을 명시적으로 처리하여 오류를 방지합니다.
np.packbits 활용: np.packbits(array >= 0, axis=-1)가 핵심입니다. 이 함수는 NumPy 배열의 각 행(axis=-1)에 대해 불리언 값을 받아 8비트씩 묶어 압축합니다. array >= 0은 각 요소가 0 이상인지 여부를 나타내는 불리언 배열을 생성하며, 이는 packbits의 입력으로 사용됩니다.
부호 처리: signed 옵션에 따라 packed.astype(np.int16) - _UNSIGNED_TO_SIGNED_DIFF 연산을 통해 부호 있는 정수 형태로 변환합니다. 이 역시 NumPy 연산으로 처리됩니다.
결과 반환: 최종 결과를 파이썬 리스트 형태로 변환하여 반환합니다 (packed.tolist()).

`tests/entrypoints/pooling/embed/test_protocol.py`

기존 코드의 테스트 케이스 외에, 변경된 로직을 검증하기 위한 새로운 테스트 케이스들이 추가되었습니다. 특히 부동 소수점 표현의 미묘한 차이(예: -0.0, denormalized numbers)가 올바르게 처리되는지, 빈 입력이나 잘못된 차원의 입력에 대한 예외 처리가 잘 동작하는지 등을 검증하는 테스트가 추가되었습니다.

test_negative_zero_treated_as_positive: -0.0이 0.0과 동일하게 처리되어 1로 비트가 설정되는지 확인합니다.
test_negative_denormal_treated_as_negative: 매우 작은 음수 값(denormalized number)이 부동 소수점 정밀도에 따라 음수로 올바르게 처리되는지 확인합니다.
test_empty_input: 빈 입력에 대해 빈 리스트를 반환하는지 확인합니다.
test_non_2d_input_raises: 2D가 아닌 입력에 대해 ValueError가 발생하는지 확인합니다.

이러한 테스트 케이스 추가는 변경된 로직의 정확성과 견고성을 보장하는 데 중요한 역할을 합니다.

왜 이게 좋은가?

성능 향상

PR 설명에 따르면, 이 변경으로 인해 M1 Mac 환경에서 약 4배의 성능 향상을 달성했습니다. 이는 순수 파이썬 루프 대비 NumPy의 벡터화된 연산이 제공하는 엄청난 성능 이점을 보여줍니다. NumPy는 내부적으로 C나 Fortran으로 구현된 최적화된 라이브러리를 사용하므로, 파이썬 인터프리터의 오버헤드 없이 대량의 데이터를 효율적으로 처리할 수 있습니다.

벡터화(Vectorization): array >= 0 연산과 np.packbits 함수는 배열 전체에 대해 한 번에 연산을 수행합니다. 이는 파이썬에서 각 요소를 개별적으로 처리하는 것보다 훨씬 빠릅니다.
메모리 효율성: NumPy 배열은 파이썬 리스트보다 메모리 사용이 효율적이며, 데이터 접근 속도도 빠릅니다.
최적화된 구현: np.packbits와 같은 NumPy 함수는 저수준에서 최적화되어 있어, 복잡한 비트 조작 연산을 매우 빠르게 수행합니다.

일반적인 교훈

파이썬 루프의 한계 인지: 수치 계산이나 데이터 집계 작업에서 파이썬의 기본 루프는 성능 병목이 되기 쉽습니다. 특히 반복 횟수가 많거나 데이터 크기가 클 경우 더욱 그렇습니다.
NumPy의 적극적인 활용: 데이터 과학 및 머신러닝 분야에서 NumPy는 필수적인 라이브러리입니다. 배열 기반 연산이 필요한 경우, NumPy의 벡터화된 함수들을 적극적으로 활용하여 성능을 크게 향상시킬 수 있습니다.
정확한 부동 소수점 처리: 임베딩 벡터의 부호 비트나 값의 크기에 따라 압축 방식이 달라질 수 있습니다. 특히 -0.0이나 denormalized number와 같은 엣지 케이스를 올바르게 처리하는 것이 중요하며, 이를 위해 테스트 커버리지를 충분히 확보해야 합니다.
명시적인 오류 처리: 입력 데이터의 형태나 크기에 대한 검증을 명확하게 하여, 예상치 못한 입력으로 인한 잠재적인 오류를 방지하고 사용자에게 유용한 피드백을 제공해야 합니다.

리뷰 요약

제공된 리뷰 댓글은 /ci run으로, 이는 CI(Continuous Integration) 파이프라인을 실행하라는 요청입니다. 특별한 기술적 논의나 피드백은 없었으며, PR 제출자가 변경 사항을 적용한 후 CI 시스템이 이를 자동으로 테스트하고 검증하는 일반적인 절차의 일부로 보입니다. 이는 PR이 성공적으로 머지되기 위한 필수 단계입니다.

결론

이번 vLLM의 PR은 _pack_binary_embeddings 함수의 구현을 파이썬 루프에서 NumPy의 np.packbits로 전환함으로써, Cohere 바이너리 임베딩 생성 성능을 4배나 향상시킨 훌륭한 예시입니다. 이는 파이썬의 성능 한계를 극복하고 효율적인 라이브러리를 구축하기 위해 NumPy와 같은 최적화된 라이브러리를 적극적으로 활용하는 것의 중요성을 다시 한번 강조합니다. 또한, 엣지 케이스를 고려한 철저한 테스트의 중요성도 보여줍니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [sglang] [MoE] SwiGLU 퓨전: Triton 커널 최적화로 메모리 대역폭 한계 돌파하기
현재글 : [vllm] vLLM, Cohere 임베딩 바이너리 압축 성능 4배 개선: NumPy를 활용한 최적화 분석
다음글 없음

댓글

관련 포스트

PR Analysis 의 다른글