엔비디아, Dynamo-Triton으로 HSTU 생성형 추천 시스템 추론 가속화

주요 변화와 방식

엔비디아는 Dynamo-Triton(이전 NVIDIA Triton Inference Server)을 통해 계층적 순차 변환 단위(HSTU) 생성형 추천 시스템(GR)의 종단간 추론 워크플로우를 지원한다. 이 워크플로우는 엔비디아 recsys-examples 저장소를 통해 제공된다. HSTU, PyTorch AOTI(Ahead-of-Time Inductor) 컴파일, FlexKV 기반 KV 캐싱, 네이티브 C++ 검증, NV 임베딩 캐시, 그리고 Dynamo-Triton 배포를 통합한다.

기술적 배경과 맥락

생성형 추천 시스템(GR)은 대규모 개인화에 대한 새로운 접근 방식으로, 추천을 사용자 행동에 대한 시퀀스 모델링으로 재구성한다. 이는 사용자 이력이 길고, 아이템 카탈로그가 지속적으로 변경되며, 개인화 품질이 풍부한 순차적 행동 모델링에 의존하는 현대 추천 워크로드에 특히 적합하다. 그러나 긴 이력, 대규모 임베딩 테이블, 시퀀스 중심 모델 아키텍처로 인해 저지연 추론이 어렵다는 과제를 안고 있다.

HSTU는 고카디널리티(high-cardinality) 및 비정상(nonstationary) 이벤트 스트림을 처리하는 GR 워크로드에 도입되었다. 기존 추천 시스템은 검색 및 랭킹을 별도의 모델과 특징 파이프라인으로 구성하는 경우가 많지만, GR은 추천을 순차적 예측 문제로 모델링하여 사용자 컨텍스트, 아이템 이력, 액션 이력, 후보 아이템을 하나의 시퀀스 인식 아키텍처에서 추론할 수 있도록 한다. 하지만 각 요청마다 긴 이력 시퀀스를 반복적으로 처리해야 하므로 추론 비용이 증가할 수 있다.

대규모 순차 추천 시스템을 서빙하는 것은 작은 밀집 랭킹 모델과 다르다. 서빙 스택은 불규칙한 시퀀스 입력, 대규모 범주형 임베딩 상태, 긴 이력, 그리고 동일한 사용자가 적은 양의 새로운 정보만으로 반복적으로 돌아오는 요청 패턴을 처리해야 한다. 매 요청마다 사용자 이력에 대한 전체 키-값 상태를 재계산하는 것은 작업 낭비와 지연 시간 증가로 이어진다. 이를 해결하기 위해 KV 캐싱은 이전 시퀀스 계산에서 재사용 가능한 키-값 데이터를 저장하여 사용자 이력의 캐시된 부분을 재계산하지 않도록 한다.

구체적인 동작 및 수치

HSTU 모델의 입력은 범주형 토큰으로 구성된다. 컨텍스트 토큰은 사용자 측 정보를, 아이템 토큰은 아이템을, 선택적 액션 토큰은 해당 아이템에 대한 사용자 상호 작용을 나타낸다. HSTU 전처리 경로는 임베딩을 검색하고, 액션 토큰이 있을 경우 아이템 및 액션 임베딩을 인터리빙하며, 컨텍스트 정보를 추가하고 위치 인코딩을 적용한다. 이후 HSTU 블록이 시퀀스를 처리하고 예측 헤드가 다중 작업 랭킹 출력을 생성한다.

엔비디아 HSTU 추론 워크플로우에는 KVCacheManager가 포함되어 GPU 메모리와 호스트 스토리지를 KV 데이터 캐시로 활용한다. GPU 캐시는 페이지 기반 KV 데이터 테이블로 구성되며, 조회, 할당, 추가, 제거 기능을 지원한다. GPU 캐시 공간이 부족할 경우 LRU(Least Recently Used) 방식의 정책에 따라 오래된 사용자 데이터가 제거될 수 있다. 호스트 측 스토리지는 캐시된 KV 데이터의 또 다른 계층을 제공하며, 워크플로우에는 FlexKV 기반 백엔드가 KV 캐시 런타임을 지원한다. HSTU 어텐션 커널은 페이지 기반 캐시에서 KV 데이터를 소비할 수 있으며, 내보내진 추론 경로는 조회, 할당, 온보딩, 추가, 오프로딩을 위한 캐시 인식 사용자 정의 연산을 포함한다.

PyTorch AOTI 워크플로우는 `torch.export`와 PyTorch AOTI를 사용하여 PyTorch 모델을 내보낸다. AOTI는 모델을 미리 컴파일하여 네이티브 C++ 런타임이 로드할 수 있는 패키지를 생성한다. 이 패키지에는 AOTI 모델 아카이브와 메타데이터, 임베딩 테이블 파일이 포함된다. 엔비디아 예시에서 임베딩 구현은 DynamicEmb 추론 임베딩 테이블과 NV 임베딩 캐시를 결합한다. NV 임베딩 캐시는 인기 있는 임베딩만 GPU 메모리에 저장하고 전체 테이블은 CPU 메모리에 유지하여 GPU 메모리 사용량을 줄인다.

Dynamo-Triton 배포 경로는 내보내진 HSTU 모델을 위한 프로덕션 서빙 계층을 제공한다. AOTI 배포는 `platform: "torch_aoti"`를 사용하여 Dynamo-Triton PyTorch 백엔드를 활용한다. 전체 워크플로우는 필요한 사용자 정의 연산자 및 런타임 라이브러리 빌드, PyTorch AOTI를 사용한 HSTU 랭킹 모델 내보내기, FlexKV 기반 KV 캐시 서비스 시작, 네이티브 C++ 리플레이를 통한 내보내진 아티팩트 검증, 그리고 Dynamo-Triton을 통한 내보내진 KV 캐시 AOTI 모델 서빙의 다섯 단계로 구성된다.

개발자 및 사용자 영향

이 워크플로우는 개발자에게 HSTU 랭킹 모델을 낮은 지연 시간으로 서빙할 수 있는 실용적인 경로를 제공한다. 개발자는 모델을 별도의 런타임용으로 재작성할 필요 없이 PyTorch 개발 환경에서 프로덕션 추론으로 전환할 수 있다. Dynamo-Triton은 모델 저장소 관리, 요청 처리, 백엔드 통합, 메트릭 및 배포 구조를 제공하며, AOTI는 오버헤드가 낮은 컴파일된 모델 아티팩트를 제공한다. NV 임베딩 캐시는 GPU 메모리 요구 사항을 낮추고, FlexKV는 긴 사용자 이력에 대한 어텐션 블록 캐싱을 통해 재계산을 줄여준다.

추천 시스템은 엄격한 지연 시간 예산 내에서 작동하므로, 추가적인 랭킹 지연 시간은 페이지 로드 시간, 피드 응답성 및 광고 서빙 마감 시간에 영향을 미칠 수 있다. 이 HSTU 서빙 워크플로우는 상호 보완적인 기술들을 결합하여 이러한 과제를 해결한다. 특히 연속적인 요청이 사용자 상호 작용 이력의 변경되지 않은 접두사를 공유할 때 이 접근 방식은 매우 유용하다. 모델은 해당 시퀀스 부분에 대한 어텐션을 재계산하는 대신 캐시된 키-값 상태를 재사용하고 새로 추가된 토큰에 필요한 부분만 계산한다. 잠재적 절감 효과는 시퀀스가 길어지고 모델이 깊어질수록 증가한다.

적용 조건 및 한계

이 워크플로우는 NVIDIA/recsys-examples GitHub 저장소에서 재현 및 확장할 수 있다. 이를 위해서는 필요한 사용자 정의 연산자 및 런타임 라이브러리 빌드, KuaiRand-1K 데이터 준비, 체크포인트 훈련이 선행되어야 한다. 또한, Dynamo-Triton 런타임 이미지를 패키징하고 Dynamo-Triton 서버를 통해 요청을 리플레이하는 과정이 필요하다.