대량 언어 모델의 효율적인 배치 추론을 위한 BASS: Batched Attention-optimized Speculative Sampling

Core Concepts

BASS는 대량 언어 모델의 추론 속도와 GPU 활용도를 크게 향상시키는 새로운 배치 추론 기법입니다.

Abstract

이 논문은 대량 언어 모델(LLM)의 추론 속도와 GPU 활용도를 높이기 위한 BASS(Batched Attention-optimized Speculative Sampling) 기법을 소개합니다.
주요 내용은 다음과 같습니다:

기존 단일 시퀀스 기반 투기적 디코딩(speculative decoding)을 확장하여 다중 시퀀스를 동시에 처리할 수 있도록 하였습니다. 이를 통해 GPU 활용도를 크게 높일 수 있습니다.

가변 길이 시퀀스를 효율적으로 처리하기 위해 BASS-PAD와 BASS-SPLIT이라는 두 가지 커널 구현 방식을 제안하였습니다.

각 시퀀스의 최적 초안 길이를 동적으로 조절하는 휴리스틱 알고리즘을 개발하였습니다.

OPT, CodeGen-Mono, 7.8B 코드 모델 등 다양한 모델에 대한 실험 결과, BASS가 기존 방식 대비 최대 3.23배 빠른 추론 속도와 2.94배 높은 GPU 활용도를 달성했음을 보여줍니다.

시간 예산 내에서 BASS는 단일 시퀀스 투기적 디코딩 대비 훨씬 높은 정확도(Pass@First 43%, Pass@All 61%)를 달성할 수 있습니다.

Stats

7.8B 코드 모델에서 BASS는 배치 크기 4일 때 첫 번째 완료 시퀀스의 토큰 당 지연 시간이 4.1ms로, 정규 디코딩 대비 3.23배 빠릅니다.
7.8B 코드 모델에서 BASS는 배치 크기 8일 때 모든 시퀀스의 평균 토큰 당 지연 시간이 5.8ms로, 정규 디코딩 대비 2.15배 빠릅니다.
7.8B 코드 모델에서 BASS는 배치 크기 8일 때 GPU 활용도가 최대 15.8%에 달하며, 이는 정규 디코딩의 3배 이상, 단일 시퀀스 투기적 디코딩의 10배 이상입니다.

Quotes

"BASS는 대량 언어 모델의 추론 속도와 GPU 활용도를 크게 향상시키는 새로운 배치 추론 기법입니다."
"BASS는 배치 크기 8일 때 GPU 활용도가 최대 15.8%에 달하며, 이는 정규 디코딩의 3배 이상, 단일 시퀀스 투기적 디코딩의 10배 이상입니다."
"BASS는 시간 예산 내에서 단일 시퀀스 투기적 디코딩 대비 훨씬 높은 정확도(Pass@First 43%, Pass@All 61%)를 달성할 수 있습니다."

Key Insights Distilled From

BASS: Batched Attention-optimized Speculative Sampling

by Haifeng Qian... at arxiv.org 04-25-2024

https://arxiv.org/pdf/2404.15778.pdf

Deeper Inquiries

대량 언어 모델의 추론 속도와 GPU 활용도를 더욱 향상시킬 수 있는 방법은 무엇일까요?

BASS(Batched Attention-optimized Speculative Sampling) 기법은 대량 언어 모델의 추론 속도와 GPU 활용도를 향상시키는 방법 중 하나입니다. 이 기법은 여러 시퀀스를 동시에 처리하면서 각 시퀀스의 특정 시점에서의 draft token을 동적으로 조정하여 GPU 활용도를 최적화합니다. 또한 BASS-PAD와 BASS-SPLIT과 같은 커스텀 CUDA 커널을 사용하여 attention 계산을 최적화하고, 동적 draft 길이 조정 알고리즘을 통해 효율적인 추론을 가능하게 합니다. 이러한 방법을 통해 GPU 활용도를 높이고 추론 속도를 향상시킬 수 있습니다.

BASS 기법의 단점은 무엇이며, 이를 극복할 수 있는 방법은 무엇일까요?

BASS 기법의 주요 단점은 attention 계산에서 발생하는 ragged tensor 문제와 각 시퀀스의 draft token 수를 동적으로 조정하는 어려움입니다. 이로 인해 효율적인 배치 처리가 어려워지고, 일부 시퀀스의 길이가 다른 경우에 대응하기 어려울 수 있습니다. 이러한 단점을 극복하기 위해 BASS-PAD와 BASS-SPLIT과 같은 커스텀 커널을 사용하여 ragged tensor 문제를 해결하고, 동적 draft 길이 조정 알고리즘을 통해 각 시퀀스에 맞는 최적의 draft 길이를 설정할 수 있습니다.

BASS 기법을 다른 분야의 모델 추론에 적용할 수 있을까요? 그 경우 어떤 추가적인 고려사항이 필요할까요?

BASS 기법은 대량 언어 모델뿐만 아니라 다른 분야의 모델 추론에도 적용할 수 있습니다. 다른 분야의 모델에서도 여러 시퀀스를 동시에 처리하면서 GPU 활용도를 최적화하고 추론 속도를 향상시킬 수 있습니다. 추가적인 고려사항으로는 각 모델의 특성에 맞게 커스텀화된 커널 및 알고리즘을 적용해야 하며, 모델의 특정 요구사항에 맞게 BASS 기법을 조정해야 합니다. 또한 다른 분야의 모델에서도 ragged tensor 문제와 동적 draft 길이 조정 알고리즘을 고려하여 최적의 성능을 달성해야 합니다.