LLM 추론 속도, 3배 끌어올린 혁신?

현재 대규모 언어 모델(LLM) 추론 가속 분야는 순차적인 ‘자가회귀적 추측 디코딩’ 방식이 주를 이루고 있습니다.

경량화된 ‘드래프터’ 모델이 토큰을 예측하고, 이를 메인 모델이 검증하는 방식입니다.

하지만 이 방식은 K개의 후보 토큰을 생성하기 위해 K번의 순차적인 연산이 필요하다는 근본적인 병목 현상을 안고 있습니다.

즉, 각 토큰 예측이 완료되어야 다음 단계로 넘어갈 수 있어 추측 단계에서의 속도 향상 잠재력을 제한하는 요인이 됩니다.

이러한 한계를 극복하기 위해 연구자들은 기존의 토큰 단위 예측 방식을 넘어, ‘블록 확산(block diffusion)’이라는 새로운 패러다임을 제시하고 있습니다.

이는 단 한 번의 순방향 연산으로 여러 개의 후보 토큰 블록을 생성하는 혁신적인 접근 방식입니다.

UCSD 연구팀, DFlash로 TPU 추론 혁신

Google Cloud는 AI 하드웨어의 경계를 넓히는 외부 연구를 적극 지원하고 있습니다.

이번에 주목할 만한 오픈소스 성과는 캘리포니아 대학교 샌디에이고(UCSD) 연구팀의 DFlash 구현입니다.

DFlash는 UCSD Z Lab에서 개발한 뛰어난 확산 기반 추측 디코딩 기술로, Google TPU 환경에 성공적으로 통합되었습니다.

특히 페이징 어텐션(paged attention) 및 사전 로딩/디코딩 분리 서빙(prefill/decode disaggregated serving)의 공동 개발자인 Hao Zhang 교수가 이끈 연구팀은 이 혁신적인 아키텍처를 오픈소스 vLLM TPU 추론 생태계에 직접 통합했습니다.

그 결과, 연구팀은 TPU v5p 환경에서 평균 3.13배의 초당 토큰 생성 속도 증가를 달성했으며, 복잡한 수학 문제의 경우 최대 6배에 달하는 성능 향상을 기록했습니다.

또한, TPU v5p에서의 DFlash와 EAGLE-3 비교 서빙 테스트에서 DFlash는 2.29배의 종단 간(end-to-end) 서빙 속도 향상을 보이며, EAGLE-3의 1.30배 성능 향상을 크게 앞질렀습니다.

기존 추론 방식의 근본적 한계 극복

표준 LLM 추론은 텍스트를 생성할 때마다 모델이 전체 순방향 연산(forward pass)을 수행하는 ‘자가회귀적(autoregressive)’ 방식을 따릅니다.

이는 TPU와 같은 AI 가속기의 막대한 병렬 연산 능력을 제대로 활용하지 못하게 하며, 특히 배치 크기가 작을 때 성능 저하를 유발합니다.

추측 디코딩은 이러한 비효율성을 완화하기 위해, 더 작고 효율적인 ‘드래프트(draft)’ 모델이나 메커니즘을 사용하여 여러 개의 미래 토큰을 동시에 예측합니다.

이후 메인 ‘타겟(target)’ 모델이 이 드래프트 토큰들을 단일 병렬 순방향 연산으로 검증합니다.

만약 드래프트 토큰이 정확하다면, 시스템은 단 한 번의 단계로 여러 토큰을 수용할 수 있어 지연 시간을 획기적으로 줄입니다.

그러나 기존 추측 디코딩 방식의 잠재력은 드래프트 모델 자체의 한계에 부딪히는 경우가 많습니다.

대부분의 기존 방법은 순차적으로 후보 토큰을 생성하는 자가회귀적 드래프트 메커니즘에 의존합니다.

즉, 타겟 모델의 검증은 병렬적으로 이루어지더라도, 드래프트 단계는 여전히 O(K)의 순차적 단계에 의해 병목 현상을 겪습니다.

결과적으로, 토큰을 ‘추측’하는 데 소비되는 시간이 검증을 통해 절약되는 시간을 잠식하여 실제적인 속도 향상 잠재력을 제한합니다.

Google TPU 환경에서의 확산 기반 드래프팅

확산 LLM(dLLM)은 이러한 순차적 프로세스를 블록 확산 메커니즘으로 대체함으로써 게임의 판도를 근본적으로 바꾸고 있습니다.

dLLM은 다음 단어를 예측하는 대신, 전체 토큰 ‘블록’을 한 번에 ‘그려냅니다(paint)’.

DFlash는 이러한 dLLM 기반 드래프팅 방법론 중 하나입니다.

타겟 모델에서 추출된 숨겨진 특징(hidden features)을 활용함으로써, DFlash는 단 한 번의 순방향 연산으로 전체 드래프트 토큰 블록을 생성할 수 있습니다.

이는 O(K) 복잡도에서 O(1) 복잡도로의 전환을 의미하며, 드래프팅 지연 시간을 거의 무시할 수 있는 수준으로 줄여줍니다.

이는 TPU의 고대역폭 행렬 곱셈 장치(MXU)에 완벽하게 부합하는 아키텍처적 이점입니다.

UCSD 연구팀은 DFlash를 vLLM TPU 추론 프레임워크에 통합하면서, Google Cloud 엔지니어들의 아키텍처 지침을 받아 오버헤드를 최소화하여 메모리 대역폭과 행렬 곱셈 장치가 최대한 활용되도록 했습니다.

DFlash 프로포저와 검증 파이프라인을 TPU 아키텍처에 효율적으로 매핑함으로써, 드래프팅 단계의 오버헤드를 최소화하는 동시에 타겟 모델의 병렬 검증 처리량을 극대화할 수 있었습니다.

TPU/JAX 환경으로의 DFlash 포팅 과제와 해결

DFlash를 기존의 GPU/PyTorch 구현에서 Google TPU/JAX AI 스택 생태계로 포팅하는 과정은 단순한 코드 변환 이상의 재설계가 필요했습니다.

TPU의 고유한 아키텍처 강점에 맞춰 시스템을 재구성하는 작업이었습니다.

UCSD 팀은 다음과 같은 세 가지 주요 기술적 난관을 극복했습니다.

1. 어텐션을 위한 ‘듀얼 캐시’ 솔루션

PyTorch 환경에서는 DFlash가 간단하고 동적인 KV 관리를 사용합니다.

하지만 tpu-inference를 통한 고성능 TPU 서빙은 페이징 어텐션과 Pallas 커널을 사용하며, 이는 메모리를 고정 크기 페이지로 분할하여 효율성을 극대화하는 시스템입니다.

DFlash의 비인과적(non-causal) 블록 확산 방식은 이러한 표준 페이징 어텐션과 기본적으로 호환되지 않습니다.

이를 해결하기 위해 연구팀은 듀얼 캐시 아키텍처를 설계했습니다.

타겟 모델은 계속해서 페이징 KV 캐시를 사용하여 대규모 서빙에 필요한 고성능 Pallas 커널의 이점을 누립니다.

드래프트 모델은 정적 온디바이스 JAX 배열을 사용하는 특화된 경로를 활용하여, 원래 DFlash 설계를 성공적으로 복제하면서 TPU 네이티브 성능을 유지했습니다.

2. 지능적인 컨텍스트 관리

DFlash의 독특한 특징은 드래프트 모델이 ‘타겟 조건부(target-conditioned)’라는 점입니다.

즉, 타겟 모델의 중간 추론 단계를 관찰하며 지속적으로 최신 정보를 유지합니다.

이러한 ‘숨겨진 상태’는 시간이 지남에 따라 증가하는 컨텍스트 버퍼에 저장됩니다.

CPU 호스트와 TPU 가속기 간의 통신을 가능한 한 빠르게 유지하기 위해, 연구팀은 ‘2의 거듭제곱(power-of-2) 패딩 전략’을 구현했습니다.

이는 새롭게 투영된 특징이 버퍼에 추가될 때 최적화된 청크 단위로 전송되도록 보장합니다.

드래프트 모델이 이미 ‘소비’한 컨텍스트의 양을 정확하게 추적함으로써, 중복 처리나 데이터 손실을 방지하여 병렬 드래프팅의 높은 정확도를 유지할 수 있었습니다.

3. TPU 추론에서의 메타데이터 간극 해소

표준 드래프트 방식과 달리, DFlash는 반복 간에 지속적인 상태(컨텍스트 버퍼, KV 캐시 위치, RoPE 오프셋 포함)를 유지하며 병렬 블록 예측을 수행하는 상태 저장(stateful) 방식을 사용합니다.

TPU 최적화 vLLM 파이프라인에서 프로포저로 전달되는 메타데이터에는 현재 검증 중인 드래프트 토큰이 포함되었습니다.

이는 대부분의 모델에 표준적인 처리 방식이지만, 확산 기반 아키텍처의 경우 ‘시퀀스 길이 인플레이션(sequence length inflation)’이라는 문제가 발생했습니다.

이는 내부 드래프트 상태가 타겟 모델의 실제 상태와 어긋나는 불일치 현상입니다.

프로포저를 재설계하여 실제 수용된 토큰 수와 엄격하게 동기화함으로써, 연구팀은 두 모델 간의 완벽한 정렬을 복원했습니다.

이 조정 덕분에 블록 확산 로직이 TPU 하드웨어에서 완전한 수학적 정밀도로 작동할 수 있게 되었고, 최종 결과에서 나타나는 극적인 속도 향상을 이끌어낼 수 있었습니다.

TPU 서빙의 미래를 벤치마킹하다

TPU v5p에서의 DFlash와 EAGLE-3의 성능 비교는 DFlash가 단순히 이론적인 우위를 넘어 실제 환경에서 체감 가능한 성능 향상을 제공함을 입증했습니다.

DFlash는 기존의 자가회귀적 추측 디코딩 방식의 한계를 명확히 보여주었으며, 블록 확산 방식이 LLM 추론 속도를 획기적으로 개선할 수 있는 잠재력을 지녔음을 시사합니다.

이러한 발전은 Google TPU의 활용성을 더욱 높이고, LLM 기반 애플리케이션의 성능을 한 단계 끌어올릴 것으로 기대됩니다.

앞으로 이러한 혁신적인 기술이 LLM 생태계 전반에 어떻게 확산될지 주목해야 할 것입니다.

출처: https://developers.googleblog.com/supercharging-llm-inference-on-google-tpus-achieving-3x-speedups-with-diffusion-style-speculative-decoding/