[vllm] vLLM의 차세대 CPU 가속: Intel Diamond Rapids를 위한 AMX-FP8 어텐션 구현

PR 링크: vllm-project/vllm#49410 상태: Merged | 변경: +1211 / -61

들어가며

vLLM은 GPU뿐만 아니라 CPU 환경에서도 고성능 추론을 지원하기 위해 지속적으로 노력하고 있습니다. 기존의 AMX(Advanced Matrix Extensions) 구현은 FP8 KV 캐시를 사용하더라도 연산 전 BF16으로 디퀀트(Dequantize)하는 과정이 필수적이었습니다. 이로 인해 메모리 대역폭과 연산 효율 면에서 병목이 발생했습니다. 본 PR은 Intel의 차세대 아키텍처인 Diamond Rapids(DMR)를 타겟으로, 하드웨어 수준에서 FP8 연산을 직접 수행하는 AMX_FP8 ISA 경로를 추가하여 이러한 오버헤드를 제거했습니다.

코드 분석

1. 하드웨어 가속 경로 추가 (csrc/cpu/cpu_attn_amx_fp8.hpp)

새로운 AMX_FP8 구현은 하드웨어 인트린직을 직접 사용하여 연산 효율을 극대화합니다.

// Before: BF16으로 변환 후 연산
// After: _tile_dpfp8ps를 통한 네이티브 FP8 MMA 연산
// QK phase: native FP8×FP8 MMA via _tile_dpfp8ps
// PV phase: native FP8 MMA with 64-byte VNNI4 tiles

2. 컴파일러 및 빌드 환경 최적화 (cmake/cpu_extension.cmake)

AMX-FP8 명령어셋을 사용하기 위해 GCC 15와 최신 binutils가 필요합니다. 또한, oneDNN에서 발생하는 GCC 15의 최적화 버그를 우회하기 위한 플래그를 추가했습니다.

# GCC 15 이상에서 발생하는 oneDNN ICE 우회
if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
    CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15)
    list(APPEND CXX_COMPILE_FLAGS "-fno-tree-pre")
endif()

3. 런타임 ISA 감지 (csrc/cpu/cpu_attn.cpp)

컴파일 타임에 AMX-FP8이 활성화되어 있더라도, 실제 실행되는 CPU가 이를 지원하는지 확인하는 런타임 체크가 추가되었습니다.

static bool runtime_has_amx_fp8() {
  unsigned int eax = 0, ebx = 0, ecx = 0, edx = 0;
  if (__get_cpuid_count(7, 1, &eax, &ebx, &ecx, &edx)) {
    return (eax >> 21) & 1u;
  }
  return false;
}

왜 이게 좋은가

이번 최적화의 핵심은 '데이터 이동 및 변환 비용의 최소화'입니다. 기존 BF16 AMX 방식 대비 다음과 같은 성능 향상을 보였습니다.

• Decode 성능: KV128 기준 1.14x 향상
• Prefill 성능: Q128/KV128 기준 0.77x (기존 0.30x 대비 대폭 개선)

일반적 교훈:

1. 하드웨어 가속 활용: 범용 연산보다 특정 데이터 타입(FP8)을 위한 전용 하드웨어 유닛(AMX-FP8)을 활용하는 것이 추론 지연 시간을 획기적으로 줄입니다.
2. 런타임 체크의 중요성: 최신 ISA를 타겟팅할 때는 컴파일 타임 지원 여부뿐만 아니라, 실제 실행 환경의 CPUID를 통한 런타임 검증이 필수적입니다.
3. 컴파일러 버그 대응: 최신 툴체인(GCC 15) 도입 시 발생할 수 있는 컴파일러 최적화 이슈(-fno-tree-pre)를 사전에 파악하고 대응하는 것이 안정적인 빌드 유지에 중요합니다.

리뷰 과정에서 bigPYJ1151은 코드의 범용성을 위해 하드코딩된 정렬 값(64)을 제안했고, zhejiangxiaomai는 이를 반영하여 구현을 간결하게 개선했습니다. 이는 성능뿐만 아니라 유지보수성 측면에서도 훌륭한 개선입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://www.intel.com/content/www/us/en/developer/articles/technical/advanced-matrix-extensions-amx.html

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