엔비디아·프린스턴대, 하이브리드 모델 디코딩 속도 개선 기술 발표

하이브리드 모델 디코딩 병목

엔비디아와 미국 프린스턴 대학교 다오 인공지능 연구실(Dao AI Lab) 연구진이 하이브리드 언어 모델의 디코딩 속도를 개선하는 기술인 리플레이에스에스엠(ReplaySSM)을 공개했다. 어텐션(Attention) 계층과 상태 공간 모델(SSM) 계층을 교차 배치하는 하이브리드 아키텍처는 긴 문맥에서도 일정한 상태 크기를 유지하도록 설계되었으나, 실제 서빙 환경에서는 오히려 상태 공간 모델 계층이 메모리 대역폭 병목을 유발하는 문제가 있었다.

기존 디코딩 방식은 매 토큰 생성 단계마다 고대역폭 메모리에서 전체 상태를 읽어와 갱신한 뒤 다시 기록하는 과정을 거친다. 연산량에 비해 상태 크기가 크고 상태 공간 모델 계층 수가 어텐션 계층보다 훨씬 많기 때문에, 동시 처리 배치 크기가 큰 환경에서는 상태 메모리 입출력이 전체 지연 시간의 대부분을 차지했다. 또한 이전 상태로 되돌리기 어려운 순환 구조 특성상 투기적 디코딩(Speculative Decoding) 적용이 까다롭고, 양자화 오차가 다음 상태로 누적되는 한계가 존재했다.

입력 버퍼링 기반 상태 분리

연구진은 매 단계마다 거대한 전체 상태를 갱신해 저장하는 대신 주기적인 체크포인트와 최근 입력값 버퍼를 조합해 상태를 표현하는 방식을 고안했다. 윈도 크기에 도달하기 전 일반 단계에서는 새로 들어온 입력값을 버퍼에 추가하고, 기존 체크포인트와 버퍼 내용만을 참조해 현재 토큰의 출력을 계산한다. 이 과정에서는 전체 상태를 메모리에 다시 쓰지 않는다.

버퍼가 가득 차는 시점에는 축적된 입력값을 한 번에 체크포인트에 반영하여 고대역폭 메모리에 기록하고 버퍼를 비우는 플러시(Flush) 과정을 수행한다. 이로써 출력 연산과 순환 상태 갱신이 분리되며, 지정된 윈도 단위마다 메모리 쓰기 작업이 단 한 번으로 줄어들어 지연 시간의 원인이 되던 메모리 트래픽이 절반 수준으로 감소한다.

투기적 디코딩 및 양자화 최적화

투기적 디코딩 검증 과정에서는 제안된 초안 토큰들의 입력값을 버퍼에 일괄 추가한 뒤 채택된 토큰 수만큼 버퍼 포인터를 이동시키는 방식으로 롤백을 단순화했다. 매 초안 토큰마다 상태 스냅샷을 저장하고 복원하던 기존 추론 엔진의 메모리 낭비를 제거함으로써 고배치 환경에서도 동시 처리량을 유지할 수 있게 되었다. 또한 초안 검증 시 중간 상태 행렬을 만들지 않고 스칼라 축약과 행렬곱으로 출력을 직접 구하는 연산 방식을 적용해 검증 지연을 대폭 단축했다.

양자화 적용 시 발생하는 오차 누적 문제는 출력 계산용 상태와 순환 갱신용 상태의 정밀도를 분리해 해결했다. 순환 갱신에 쓰이는 체크포인트는 고정밀도 형식을 유지하고, 비플러시 단계의 출력 연산에만 저정밀도 복사본을 읽도록 구성했다. 저정밀도 상태에서 발생한 반올림 오차가 다음 순환 루프로 유입되지 않아 초장문 생성 시에도 벤치마크 정확도 저하가 최소화되었다.

성능 평가 및 적용 제약

엔비디아 H100 및 블랙웰 울트라 B300 환경에서 40억 매개변수 모델부터 1200억 매개변수급 전문가 혼합(MoE) 모델까지 테스트한 결과, 배치 크기 256 환경에서 커널 실행 속도는 1.46배에서 1.80배, 디코딩 단계 속도는 1.20배에서 1.47배 빨라졌다. 투기적 서빙 처리량 역시 기존 프레임워크 대비 1.85배에서 2.10배 향상되는 결과를 보였다.

단일 요청을 처리하는 배치 크기 1 환경에서는 메모리 대역폭 외의 다른 지연 요소가 지배적이어서 속도 향상이 나타나지 않는 제약이 있다. 성능 측정은 윈도 크기를 16으로 설정하고 부동소수점 형식을 적용한 상태에서 진행되었으며, 현재 해당 구현체는 오픈소스 추론 프레임워크 및 일부 상용 모델 서비스에 통합되어 제공되고 있다.