AWS Trainium 기반 추론적 디코딩을 통한 생성 속도 3배 향상

주장아마존 웹 서비스(AWS)의 트레이니엄(Trainium) 환경에서 추론적 디코딩(Speculative Decoding)을 적용하면 토큰 생성 속도를 최대 3배까지 높일 수 있습니다. 이는 토큰 생성 단계가 많은 작업에서 하드웨어 활용도를 극대화하고 운영 비용을 절감하는 핵심 전략입니다.

팩트추론적 디코딩은 작은 초안 모델(Draft Model)이 여러 토큰을 미리 제안하고, 타겟 모델이 이를 한 번의 연산으로 검증하는 방식을 취합니다. 이 과정은 순차적인 디코딩 단계를 줄여 메모리 대역폭 병목 현상을 완화합니다.

주장해당 기술은 키-값(KV) 캐시 메모리 접근 횟수를 감소시켜 하드웨어 효율을 높입니다. 표준 디코딩 방식은 토큰 하나를 생성할 때마다 메모리에서 전체 캐시를 읽어야 하므로 메모리 대역폭에 크게 의존합니다.

팩트AWS 뉴런엑스 분산 추론(NeuronX Distributed Inference, NxDI) 라이브러리는 바닐라, 퓨즈드, 이글(EAGLE), 메두사(Medusa) 등 네 가지 추론적 디코딩 모드를 지원합니다. 이번 벤치마크에서는 최적의 성능을 확보하고자 퓨즈드 추론 방식을 채택했습니다.

팩트벤치마크는 트레이니엄2(trn2.48xlarge) 인스턴스와 아마존 엘라스틱 쿠버네티스 서비스(Amazon EKS) 환경에서 수행했습니다. 동일한 부하 조건에서 엘엘엠퍼프(llmperf) 도구를 사용하여 표준 디코딩과 추론적 디코딩의 성능을 비교했습니다.

팩트큐웬3-32B(Qwen3-32B) 모델을 타겟으로 큐웬3-1.7B(Qwen3-1.7B)를 초안 모델로 활용했을 때 7개의 추론 토큰을 설정하는 것이 가장 균형 잡힌 성능을 보였습니다. 0.6B 모델은 속도는 빠르지만 토큰 수락률이 낮아 효율성이 떨어졌습니다.

교차검증초안 모델과 타겟 모델은 동일한 토크나이저와 어휘 사전을 공유해야 합니다. 서로 다른 아키텍처를 사용할 경우 토큰 일치율이 낮아져 성능 향상 효과가 반감될 수 있습니다.

교차검증추론 토큰 수(num_speculative_tokens) 설정값에 따라 결과가 달라집니다. 너무 낮게 설정하면 속도 향상이 제한되고, 너무 높게 설정하면 조기 거부 확률이 높아져 검증 비용이 증가하는 부작용이 발생합니다.

교차검증성능 향상 폭은 프롬프트의 구조에 따라 차이를 보입니다. 반복적인 코드나 숫자 시퀀스처럼 구조화된 프롬프트에서 더 높은 성능 개선을 확인했습니다.

출처아마존 웹 서비스의 공식 머신러닝 블로그를 통해 해당 기술의 벤치마크 결과와 추론적 디코딩 적용 방식을 교차 검증했습니다.