Intel GPU에서의 효율적인 LLM 추론 방법에 대한 연구 소개 (feat. Intel)

읽을거리&정보공유
, , , , ,
1

PyTorchKR​:fire::kr: :thinking::speech_balloon:

  • 대규모 언어 모델(LLM) 사용, 특히 추론을 최적화하기 위해 다양한 연구들이 이뤄지고 있습니다. 이번에 소개하는 논문은 Intel GPU에서의 효율적인 추론 방법에 대한 연구로, 예상하셨겠지만 Intel이 연구하여 공개하였습니다. :sweat_smile: 접근 방식에 대해서는 한 번쯤 살펴보셔도 좋을 것 같아 어떠한 아이디어들이 있는지 간단히 살펴보려고 합니다. :disguised_face:

  • 이 글은 GPT 모델로 정리한 것으로, 잘못된 부분이 있을 수 있으니 글 아래쪽의 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

Intel GPU에서 효율적인 LLM 추론 솔루션 (Efficient LLM inference solution on Intel GPU)

소개

트랜스포머 기반 대규모 언어 모델(LLM)은 여러 분야에서 널리 사용되고 있으며, 실제 애플리케이션에서도 LLM 추론의 효율성이 화두가 되고 있습니다. 그러나 일반적으로 LLM은 대규모 연산을 통해 모델 구조가 복잡하게 설계되고 자기회귀 형태로 추론을 수행하기 때문에 높은 효율을 가진 시스템을 설계하는 것은 어려운 과제입니다.

이를 해결하기 위해, 메모리 효율을 높이고 연산 비용을 낮추기 위해 모델 압축이나 효율적인 배치 크기 관리, 연산 퓨전 등과 같은 주제들이 연구되고 있습니다. 하지만 이러한 방법들은 GPU 특성에 충분히 최적화되지 않았으며, 특히 인텔 GPU의 아키텍처 특성과 맞지 않는 경우가 많았습니다. 이로 인해 인텔 GPU에서 LLM 추론의 효율성이 크게 제한되었습니다.

논문에서는 지연 시간이 짧고 처리량이 높은 효율적인 LLM 추론 솔루션을 제안합니다. 1) 데이터 이동과 요소별 연산을 융합하여 LLM 디코더 계층을 단순화함으로써 다음과 같이 메모리 액세스 빈도를 줄이고 메모리 액세스 빈도를 줄이고 시스템 지연 시간을 줄입니다. 2) 세그먼트 KV 캐시 정책을 제안하여 요청 및 응답 토큰의 키/값을 별도의 물리적 메모리에 보관하여 효과적인 장치 메모리를 관리하여 런타임 배치 크기를 늘리고 시스템 처리량을 개선하는 데 도움이 됩니다. 3) 커스터마이징된 Scaled-Dot-Product-Attention 커널은 퓨전 정책과 일치하도록 설계되었습니다.

Intel GPU에서 주요 LLM 모델을 실행하였을 때, 표준 허깅페이스 구현과 비교하여 최대 7배 낮은 토큰 지연 시간과 27배 높은 처리량을 확인하였습니다.


제안하는 방법

논문에서는 효율적인 추론을 위해 3가지 방법을 제안하고 있습니다. 먼저 LLM 디코더 레이어 단순화는 데이터 이동과 요소별 연산을 퓨전하여 디코더 레이어를 단순화함으로써 메모리 접근 빈도와 연산 복잡성을 줄입니다. 이를 통해 전체 시스템의 지연 시간(latency)이 줄어들며, 메모리 대역폭의 이용이 효율적으로 이루어집니다. 다음으로 세그먼트 KV 캐시 정책은 Key/Value(KV) 메모리를 효과적으로 관리하는 새로운 정책을 도입하는 것입니다. 이 정책은 요청과 응답에 대한 메모리 할당을 최적화하는 것으로, 더 큰 런타임 배치 크기를 지원하며, 시스템의 처리량을 향상시킵니다. 마지막으로 맞춤형 스케일-닷-프로덕트-어텐션 커널은 인텔 GPU의 아키텍처에 맞춤화된 스케일-닷-프로덕트-어텐션 커널을 개발하는 것입니다. 앞에 설명한 두가지 방법과 조화를 이루면서 연산의 효율성을 극대화하며 전체적인 추론 속도를 향상시킵니다.

모델 구조 단순화 (Model Structure Simplification)

대규모 언어 모델(LLM)은 다수의 디코더 레이어(Decoder Layer)를 포함하는 복잡한 구조로 구성되어 있습니다. 이러한 디코더 레이어는 컨텍스트 정보를 포착하기 위해 방대한 연산을 수행하며, 각 디코더 레이어에는 멀티-헤드 어텐션(MHA, Multi-Head Attention)과 피드-포워드(Feed-Forward, FF) 레이어로 구성되어 있습니다.

llama2 모델에서 최적화된 멀티헤드 어텐션 모듈

이 중 멀티-헤드 어텐션(MHA) 모듈의 데이터 이동 연산(예: Transpose, Cat, Index Select)은 메모리 접근과 커널 실행에 대한 오버헤드를 유발하고 있어, MHA 구조를 단순화하여 이러한 연산들을 줄였습니다​​. 예를 들어, RMSNorm, RoPE, SDPA 모듈 내의 다중 연산들은 단일 커널로 통합됩니다. 이러한 모듈들에서 사용되는 세 개의 Linear 연산(쿼리, 키, 값 생성을 위한)은 하나로 결합됩니다. 또한, 요소별 연산(element-wise operations)은 이전의 Linear 연산과 함께 통합됩니다​​.

llama2 모델 추론 시의 최적화된 연산 흐름
llama2 모델 추론 시의 최적화된 연산 흐름1038×747 88 KB

이러한 최적화를 통해, Llama2 모델의 각 디코더 레이어에서 수행되는 연산의 수가 원래 수백 개에서 단 아홉 개로 줄어듭니다. 이 최적화는 데이터 이동 연산과 요소별 연산을 모두 제거하며, 여러 연산을 단일 커널로 통합합니다​​. 디코딩 단계에서 시퀀스/배치 레이아웃 변환 작업을 도입하였음에도 불구하고, 이러한 작업은 각 시간 단계에서 첫 번째 및 마지막 디코더 레이어 전에만 한 번 발생하기 때문에 전체 추론 과정에 비해 거의 부담이 되지 않습니다​​.

세그먼트 KV 캐시 (Segment KV Cache)

대규모 언어 모델(LLM) 추론은 큰 매개변수 크기와 키/값(KV, Key-Value) 캐시 정책 때문에 많은 메모리를 소비합니다. 이는 배치 크기를 제한하고, 결과적으로 시스템 처리량에 영향을 미칩니다​. 특히, 기존의 표준 KV 캐시 구현에서는 프롬프트와 응답의 키/값을 함께 연결하여 연속적인 KV 캐시를 형성합니다. 이는 프롬프트 키/값이 BW(Beam Width, 빔 크기)만큼 확장되어야 하므로, 장치 메모리를 낭비하게 됩니다​​. 특히, 디코딩 단계에서 새로운 더 큰 KV 캐시 버퍼가 할당되면서 이전의 작은 KV 캐시 버퍼를 재사용하지 못하게 되어, 많은 메모리 파편화가 발생합니다​​.

Segment KV Cache는 이를 해결하기 위해 제안한 방식으로, 프롬프트와 응답의 KV값을 별도의 버퍼에 보관하고, 주기적으로 캐시의 파편을 비우는 방식입니다. Prefill 단계에서는 프롬프트의 KV값이 [BatchSize, N_{prompt}, H, D] 형태로 각 디코더 레이어마다 생성되어 보관되고, Decoding 단계에서는 응답 KV 캐시가 [N_{step}, BatchSize \times BeamWidth, H, D] 형태로 미리 할당되어, 각 단계(Step)에서 응답 KV 값이 해당하는 위치에 저장됩니다. 또한, 이전에는 N_{step} 을 일반적으로 모델별 최대 길이(예. Llama2의 경우 4096)로 정해두었지만, 이렇게 하는 경우 실행 중 응답이 더 짧을 경우 메모리 낭비가 발생합니다. 이를 방지하기 위해 N_{Step} 값을 동적으로 16씩 증가시키는 방법을 사용합니다. (응답 길이가 16 보다 커지는 경우는 논문 7page, Eq(3) 위쪽 문단을 참고해주세요)

Segment KV Cache 정책을 사용하는 경우 메모리 소비량은 다음과 같이 계산할 수 있습니다.
cache_{segment} = BatchSize \times (N_{prompt} + BeamWidth \times Ceil(\frac{N_{response}} {step}) \times cache_{token}

이렇게 Segment KV Cache 정책을 적용하여 프롬프트 KV을 다양한 응답 토큰과 공유하게 하여 메모리 낭비를 줄일 수 있으며, N_{step} 값의 동적 할당으로 인해 메모리 파편화도 감소시키고, 장치 메모리 재사용 확률을 높입니다​​. 이러한 메모리 관리는 대규모 언어 모델 추론에서 처리량을 향상시키는 데 기여합니다.

주요 LLM에서 표준 KV 캐시와 세그먼트 KV 캐시 사용 시의 메모리 소비 비교
주요 LLM에서 표준 KV 캐시와 세그먼트 KV 캐시 사용 시의 메모리 소비 비교1098×287 54.1 KB

맞춤형 SDPA 커널 (Customized SDPA Kernel)

맞춤형 SDPA(Scaled-Dot-Product-Attention) 커널은 FlashAttention 및 FasterTransformer에서 영감을 받아 SDPA 모듈 내의 모든 계산 단계와 인덱스 선택(Index Select) 연산을 하나의 커널로 통합합니다.

맞춤형 SDPA 커널 구형 예시
맞춤형 SDPA 커널 구형 예시1051×617 103 KB

디코딩 단계에서 SDPA 커널의 입력 텐서는 쿼리, 프롬프트 키/값, 응답 키/값의 형태로 구성됩니다. 이 때 각 입력 텐서의 크기는 다음과 같습니다.

  • Shape of Query: [1, BatchSize \times BeamWidth, H, D]
  • Shape of Prompt KV: [BatchSize, N_{prompt}, H, D]
  • Shape of Response KV: [N_{response}, BatchSize \times BeamWidth, H, D]

이 때, 결과로 출력되는 텐서는 [1, BatchSize \times BeamWidth, H, D] 의 크기를 가집니다.

맞춤형 SDPA 커널에서는 BS(Batch Size)와 H(Head 수)가 GPU의 다른 작업 그룹(블록)에서 병렬 처리됩니다. 단일 GPU 작업 그룹에서는 쿼리 및 프롬프트 키/값, 응답 키/값의 입력 텐서가 각각 [1, BeamWidth, D], [N_{prompt}, D], [N_{response}, BeamWidth, D] 의 형태로 처리됩니다​​. Prefill 단계에서의 맞춤형 SDPA 커널은 디코딩 단계의 특별한 경우로, 쿼리와 프롬프트 키/값의 입력 텐서만을 사용하여 어텐션 컨텍스트를 계산합니다​​.

결론

(성능 개선에 대한 내용은 논문의 4. Experiments 부분(9p~12p)을 참고해주세요. )

대규모 언어 모델은 추론 시 메모리를 많이 사용하기 때문에, 메모리 접근 빈도를 줄이면 효율성을 높일 수 있습니다. 논문에서는 메모리 접근 빈도를 낮추기 위해 데이터 이동과 요소별 연산을 융햡하여 디코더 부분을 단순화하는데 중점을 두었습니다. 또한, 프롬프트가 서로 다른 응답 간에 공유되도록 하고, 정기적으로 캐시를 비우는 세그먼트 KV 캐시를 사용하여 장치의 메모리를 절약하고 처리량을 개선합니다. 마지막으로 앞에서 언급한 융합과 Segment KV 캐시 정책을 통합하여 효율적인 맞춤형 SDPA 커널을 설계하였습니다.

더 읽어보기

논문



:fire:파이토치 한국 사용자 모임:kr:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

:gift: 아래:arrow_lower_right:쪽에 좋아요:heart:를 눌러주시면 뉴스 발행에 힘이 됩니다~ :star_struck:

4개의 좋아요