[vllm] vLLM, Arm CPU의 BF16 GELU 연산을 LUT 기반 구현으로 8배 가속

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

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께 추론 속도 향상에 대한 요구가 높아지고 있습니다. 특히, 모델의 크기가 커지고 복잡해짐에 따라 각 연산의 효율성이 전체 추론 성능에 미치는 영향이 더욱 중요해지고 있습니다. vLLM은 LLM 추론을 위한 고성능 라이브러리로, 지속적인 최적화 작업을 통해 모델의 성능을 개선하고 있습니다. 이번 PR은 Arm CPU 환경에서 BF16(Brain Floating Point) 데이터 타입의 GELU(Gaussian Error Linear Unit) 연산 속도를 획기적으로 개선하는 내용을 담고 있습니다.

기존 PyTorch의 GELU 구현은 특정 워크로드, 예를 들어 양자화된 Whisper 모델의 추론 시 약 5%의 시간을 차지할 정도로 성능 병목 현상을 일으킬 수 있었습니다. 이는 행렬 곱셈(GEMM)과 같은 주요 연산이 아님에도 불구하고 상당한 비중을 차지하는 것이었습니다. 본 PR은 이러한 문제를 해결하기 위해 Arm CPU에 특화된 Look-Up Table (LUT) 기반의 BF16 GELU 구현을 도입했습니다. 이를 통해 기존 대비 최대 8배의 속도 향상을 달성했으며, 결과적으로 양자화된 Whisper 모델의 추론 속도를 최대 5%까지 단축시키는 효과를 가져왔습니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 최적화가 효과적인지, 그리고 이를 통해 얻을 수 있는 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 분석

이번 PR은 주로 Arm CPU 환경에서의 BF16 GELU 연산 성능 향상에 초점을 맞추고 있으며, 관련 C++ 소스 코드와 테스트 파일, 그리고 빌드 설정 및 Python 인터페이스를 수정했습니다.

1. csrc/cpu/activation_lut_bf16.cpp 및 csrc/cpu/torch_bindings.cpp: 새로운 LUT 기반 GELU 구현

가장 핵심적인 변경은 csrc/cpu/activation_lut_bf16.cpp 파일에 새로운 LUT 기반 GELU 구현이 추가된 것입니다. 이 파일은 Arm CPU의 SIMD(Single Instruction, Multiple Data) 명령어셋(ASIMD)을 활용하여 BF16 데이터 타입의 GELU 연산을 효율적으로 처리합니다.

Before (기존 방식 - PyTorch 기본 구현):

명시적인 LUT 기반 구현이 없었으며, PyTorch의 torch.gelu 함수를 사용했습니다. 이는 일반적으로 부동 소수점 연산을 통해 계산되므로, 특히 BF16과 같이 정밀도가 낮은 데이터 타입에서는 상대적으로 느릴 수 있습니다.

After (LUT 기반 구현):

// csrc/cpu/activation_lut_bf16.cpp

// ... (헤더 및 상수 정의) ...

// GELU의 참조 구현 (PyTorch 기본)
at::Tensor gelu_reference(const at::Tensor& x) {
    return at::gelu(x, "none");
}

// LUT 초기화 함수
void maybe_init_activation_lut_bf16(
    uint16_t* lut, std::once_flag& once,
    at::Tensor (*activation)(const at::Tensor&)) {
  std::call_once(once, [&]() {
    // LUT 입력으로 사용할 float 범위 생성 (0.0 ~ 65504.0)
    auto lut_input = at::empty({static_cast<int64_t>(ActivationLutSize)},
                              at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
    auto* lut_input_ptr = lut_input.data_ptr<float>();
#pragma omp parallel for
    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
      // uint16_t 비트를 float로 변환
      lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
    }

    // 참조 활성화 함수를 사용하여 LUT 출력 계산
    auto lut_output = activation(lut_input);
    const auto* lut_output_ptr = lut_output.data_ptr<float>();
#pragma omp parallel for
    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
      // 계산된 float 값을 반올림하여 uint16_t (BF16)로 저장
      lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
    }
  });
}

// 실제 LUT 적용 함수
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
                         const uint16_t* lut, const char* op_name) {
  // ... (입력/출력 텐서 및 타입 검증) ...

  const auto* src = reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
  auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
  const int64_t n = input.numel();

  CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
#pragma omp parallel for
  for (int64_t i = 0; i < n; ++i) {
    // LUT를 사용하여 입력 값을 직접 조회하여 출력에 저장
    dst[i] = lut[src[i]];
  }
  CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
}

// 오버로드된 함수 (활성화 함수 이름 인자로 받음)
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
                         const std::string& activation) {
  if (activation == "gelu") {
    static std::array<uint16_t, ActivationLutSize> lut{};
    static std::once_flag once;
    maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
    activation_lut_bf16(out, input, lut.data(), "gelu_lut");
    return;
  }

  TORCH_CHECK(false, "Unsupported activation: ", activation);
}

이 코드는 다음과 같은 단계를 거칩니다:

1. LUT 초기화 (maybe_init_activation_lut_bf16): ActivationLutSize (65536) 크기의 LUT를 생성합니다. 이 LUT는 BF16의 모든 가능한 값(0부터 65535까지)을 float로 변환한 후, gelu_reference 함수(PyTorch의 기본 GELU 구현)를 적용하여 계산된 float 결과를 반올림하여 uint16_t (BF16) 형태로 저장합니다. 이 초기화는 std::call_once를 사용하여 단 한 번만 수행됩니다.
2. LUT 적용 (activation_lut_bf16): 실제 추론 시에는 입력 텐서의 각 BF16 값을 그대로 LUT의 인덱스로 사용하여 미리 계산된 값을 조회합니다. 이 과정은 매우 빠르며, Arm CPU의 SIMD 명령어를 활용하여 병렬로 처리됩니다.

csrc/cpu/torch_bindings.cpp 파일에서는 이 새로운 C++ 함수를 PyTorch C++ API를 통해 노출시키고, Arm CPU 환경에서만 컴파일되도록 #if (defined(__aarch64__) && !defined(__APPLE__)) 조건문을 사용했습니다.

2. vllm/model_executor/layers/activation.py: Python 인터페이스 통합

vllm/model_executor/layers/activation.py 파일의 GELU 클래스는 새로운 C++ 연산을 사용하도록 수정되었습니다.

Before: GELU 클래스는 기본적으로 PyTorch의 F.gelu를 사용했습니다.

After: GELU 클래스의 __init__ 메서드에서 현재 플랫폼이 Arm CPU이고 torch.ops._C.activation_lut_bf16 연산이 사용 가능하다면, 해당 연산을 self.op으로 설정합니다. forward_cpu 메서드에서는 이 self.op을 사용하여 BF16 입력에 대해 LUT 기반 GELU를 호출합니다. 만약 Arm CPU가 아니거나 해당 연산이 없다면, 기존의 forward_native (PyTorch 기본 구현)를 사용합니다.

# vllm/model_executor/layers/activation.py

# ... (기존 코드) ...

@CustomOp.register("gelu")
class GELU(CustomOp):
    # --8<-- [start:gelu]

    def __init__(self):
        super().__init__()
        # Arm CPU이고 activation_lut_bf16 연산이 사용 가능하면 해당 연산 사용
        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM and hasattr(
            torch.ops._C, "activation_lut_bf16"
        ):
            self.op = torch.ops._C.activation_lut_bf16
        else:
            self.op = None

    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        return F.gelu(x, approximate="none")

    def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
        # BF16 타입이고 contiguous한 경우에만 LUT 연산 사용
        if self.op and x.dtype == torch.bfloat16 and x.is_contiguous():
            out = torch.empty_like(x)
            self.op(out, x, "gelu")
            return out
        return self.forward_native(x)

    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
        return self.forward_native(x)

# ... (기존 코드) ...

3. tests/kernels/core/test_cpu_activation.py: 새로운 테스트 추가

tests/kernels/core/test_cpu_activation.py 파일에 새로운 테스트 케이스가 추가되었습니다. 이 테스트는 새로 구현된 activation_lut_bf16 연산을 포함한 다양한 CPU 활성화 함수들을 검증합니다. 특히, GELU 클래스와 torch.ops._C.activation_lut_bf16 연산을 사용하는 테스트는 Arm CPU 환경에서만 실행되도록 pytest.mark.skipif 데코레이터를 사용하여 명시했습니다.

# tests/kernels/core/test_cpu_activation.py

# ... (기존 코드) ...

@pytest.mark.parametrize(
    ("activation_cls", "fn", "op_args"),
    [
        # ... (다른 활성화 함수 테스트) ...
        pytest.param(
            GELU,
            getattr(torch.ops._C, "activation_lut_bf16", None),
            ("gelu",), # op_args
            marks=pytest.mark.skipif(
                current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
                reason="activation_lut_bf16 is only built on Arm CPU",
            ),
        ),
    ],
)
# ... (테스트 함수 본문) ...

4. .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh: CI 설정 변경

Arm CPU 하드웨어 CI 환경에서 새로운 활성화 함수 테스트가 실행되도록 .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh 스크립트가 수정되었습니다. pytest -x -v -s tests/kernels/core/test_cpu_activation.py 라인이 추가되어 해당 테스트 파일이 CI 파이프라인에 포함되도록 했습니다.

5. cmake/cpu_extension.cmake: 빌드 시스템 통합

Arm CPU 확장을 빌드할 때 csrc/cpu/activation_lut_bf16.cpp 파일을 포함하도록 cmake/cpu_extension.cmake 파일이 수정되었습니다. 이는 Arm CPU 환경에서 vLLM을 빌드할 때 해당 C++ 소스 코드가 컴파일되고 라이브러리에 포함되도록 보장합니다.

왜 이게 좋은가?

1. 압도적인 성능 향상

이 PR의 가장 큰 장점은 Arm CPU 환경에서 BF16 GELU 연산의 성능을 최대 8배까지 향상시켰다는 점입니다. 이는 모델의 전체 추론 속도에 직접적인 영향을 미칩니다. 특히, PR 설명에 언급된 것처럼 양자화된 Whisper 모델의 경우 전체 추론 시간의 약 5%를 차지하던 GELU 연산이 빨라짐으로써 모델의 최종 추론 속도가 최대 5%까지 개선되었습니다.

이러한 성능 향상은 다음과 같은 이유로 가능했습니다:

• LUT (Look-Up Table) 활용: 복잡한 부동 소수점 계산 대신 미리 계산된 값을 테이블에서 직접 조회하는 방식은 연산량을 획기적으로 줄입니다. 특히, BF16 데이터 타입은 2^16 = 65536개의 고유한 값을 가지므로, 이 값들에 대한 연산 결과를 미리 계산해두면 매우 효율적입니다.
• Arm CPU 특화 최적화: Arm CPU의 SIMD 명령어셋(ASIMD)을 활용하여 병렬 처리를 극대화했습니다. LUT에서 값을 조회하는 간단한 연산이라도 여러 데이터를 동시에 처리할 수 있다면 성능은 크게 향상됩니다.
• 데이터 타입 (BF16) 최적화: BF16은 FP32보다 표현 범위는 유사하지만 정밀도가 낮아 연산량이 적습니다. LUT 기반 구현은 이러한 BF16의 특성을 잘 활용하여 성능을 극대화합니다.

2. 코드 중복 최소화 및 유지보수성

리뷰 과정에서 oneDNN의 LUT 구현을 재사용할지에 대한 논의가 있었습니다. 하지만 oneDNN -> PyTorch -> vLLM으로 이어지는 의존성 전파 지연, 복잡한 리팩토링 필요성, 그리고 vLLM 내에서 oneDNN의 일부 기능만 직접 사용하는 현재 상황 등을 고려했을 때, vLLM 자체적으로 약 50줄 내외의 간단한 LUT 구현을 추가하는 것이 단기적으로는 더 효율적이고 유지보수하기 용이하다는 결론이 내려졌습니다. 또한, 이 LUT 기반 구현은 향후 PyTorch에 fast activation 기능이 통합되면 제거될 예정이므로, 일시적인 최적화로 볼 수 있습니다.

3. 명확한 테스트 및 검증

새로운 테스트 파일(tests/kernels/core/test_cpu_activation.py)이 추가되어, 새로 구현된 LUT 기반 GELU 연산이 기존 PyTorch 구현과 동일한 결과를 내는지, 그리고 Arm CPU 환경에서 올바르게 동작하는지를 검증합니다. 이는 코드의 안정성을 높이고 향후 발생할 수 있는 회귀(regression)를 방지하는 데 중요한 역할을 합니다.

4. 일반적인 교훈

• 하드웨어 특화 최적화의 중요성: 특정 하드웨어 아키텍처(여기서는 Arm CPU)의 명령어셋과 데이터 타입(BF16)의 특성을 깊이 이해하고 이를 활용하면 상당한 성능 향상을 이끌어낼 수 있습니다.
• LUT의 유용성: 복잡한 수학 함수나 연산의 경우, 특히 정밀도 요구사항이 높지 않거나 입력 값의 범위가 제한적인 경우, LUT를 활용하는 것은 매우 효과적인 성능 최적화 기법이 될 수 있습니다.
• 라이브러리 통합 전략: 외부 라이브러리(oneDNN)의 기능을 활용하는 것과 자체 구현을 추가하는 것 사이의 trade-off를 신중하게 고려해야 합니다. 의존성, 개발 속도, 유지보수성, 그리고 장기적인 로드맵을 종합적으로 판단해야 합니다.
• 테스트의 중요성: 새로운 최적화 기법을 도입할 때는 반드시 기존 기능과의 호환성 및 정확성을 검증하는 철저한 테스트가 수반되어야 합니다.

결론

vLLM의 이번 PR은 Arm CPU 환경에서 BF16 GELU 연산의 성능을 LUT 기반 구현을 통해 획기적으로 개선했습니다. 이는 LLM 추론 속도 향상에 기여하며, 특히 특정 모델에서는 실질적인 성능 개선 효과를 보여줍니다. 하드웨어 특화 최적화와 LUT 활용의 좋은 사례이며, 라이브러리 개발 시 고려해야 할 다양한 측면을 보여주는 흥미로운 변경이었습니다.

참고 자료

• https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml
• https://github.com/oneapi-src/oneDNN/pull/4827

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