Agentic AI / Generative AI

NVIDIA Dynamo-Triton으로 HSTU 생성형 추천 시스템 배포

Reading Time: 7 minutes

생성형 추천 시스템인 GR은 대규모 개인화를 위한 강력한 접근법으로 떠오르고 있습니다. GR은 추천을 검색, 랭킹, 예측이라는 개별 단계의 집합으로 다루지 않고, 사용자 행동에 대한 시퀀스 모델링으로 재정의합니다. 사용자의 상호작용, 컨텍스트, 후보 아이템, 행동이 카디널리티가 높은 이벤트 스트림의 토큰이 되고, 모델은 이 시퀀스에서 다음에 관련될 아이템을 생성하거나 점수화하도록 학습합니다.

이 접근법은 사용자 이력이 길고, 아이템 카탈로그가 끊임없이 바뀌며, 개인화 품질이 풍부한 순차적 행동을 얼마나 잘 모델링하느냐에 좌우되는 최신 추천 워크로드에 특히 적합합니다. 그러나 서빙에서는 새로운 과제가 생깁니다. 이력이 길고, 임베딩 테이블이 크고, 시퀀스 중심의 모델 아키텍처를 쓰더라도 GR 모델은 낮은 지연 시간으로 추론해야 합니다.

NVIDIA Dynamo-Triton(이전 이름 NVIDIA Triton Inference Server)은 이제 NVIDIA recsys-examples 리포지토리를 통해 엔드투엔드 Hierarchical Sequential Transduction Unit, 즉 HSTU GR 추론 워크플로를 지원합니다. 이 워크플로는 HSTU, PyTorch Ahead-of-Time Inductor 컴파일, FlexKV 기반 KV 캐싱, 네이티브 C++ 검증, NV 임베딩 캐시, Dynamo-Triton 배포를 결합합니다.

그 결과 HSTU 랭킹 모델을 우수한 지연 시간 성능으로 서빙하는 실용적인 경로가 마련되었습니다. NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU에서 동적 배치 크기 8로 측정한 결과, PyTorch AOTI를 적용한 Dynamo-Triton은 GPU KV 캐시 적중률 100% 조건에서 KV 캐싱을 쓰지 않은 동일한 AOTI 구성보다 3계층 HSTU 모델은 최대 4.47배, 8계층 모델은 최대 5.93배 빠른 최상 사례 성능을 보였습니다.

이 글에서는 HSTU 생성형 추천 모델을 PyTorch 개발 단계에서 NVIDIA Dynamo-Triton, PyTorch AOTI, FlexKV 기반의 프로덕션 추론으로 옮기는 방법을 소개합니다. 모델을 내보내고 사전 컴파일하는 방법, 결과물인 배포 아티팩트를 Python과 네이티브 C++에서 검증하는 방법, 별도 런타임용으로 모델을 다시 작성하지 않고 Dynamo-Triton으로 서빙하는 방법을 다룹니다.

또한 GPU 기반 KV 캐싱이 반복 연산을 줄이는 원리를 살펴보고, 지연 시간이 최대 5.93배 낮아진 벤치마크 결과를 통해 이 배포 워크플로의 실질적인 성능 이점을 보여 드립니다.

생성형 추천에 HSTU를 사용하는 이유

HSTU는 카디널리티가 높고 비정상적인 이벤트 스트림을 다루는 GR 워크로드를 위해 고안되었습니다. 기존 추천 시스템에서는 검색과 랭킹을 여러 특화 모델과 피처 파이프라인의 조합으로 구성하는 경우가 많습니다. 반면 GR은 추천을 순차 예측 문제로 모델링하여, 사용자 컨텍스트, 아이템 이력, 행동 이력, 후보 아이템을 하나의 시퀀스 인식 아키텍처 안에서 추론합니다.

NVIDIA HSTU 랭킹 예제에서 모델 입력은 범주형 토큰으로 구성됩니다. 컨텍스트 토큰은 사용자 측 정보, 아이템 토큰은 아이템, 선택적 행동 토큰은 해당 아이템에 대한 사용자의 상호작용을 나타냅니다.

HSTU 전처리 경로는 임베딩을 조회하고, 행동 토큰이 있으면 아이템 임베딩과 행동 임베딩을 교차 배치합니다. 이어서 컨텍스트 정보를 덧붙이고 위치 인코딩을 적용합니다. 그다음 HSTU 블록이 시퀀스를 처리하고, 예측 헤드가 멀티태스크 랭킹 출력을 생성합니다.

이 구조는 최신성, 순서, 반복되는 상호작용 패턴이 중요한 추천 시스템에 적합합니다. 다만 요청마다 긴 이력 시퀀스를 반복해서 처리하면 추론 비용이 커질 수 있습니다. 프로덕션 시스템은 HSTU 모델링의 이점을 유지하면서 서빙 중의 중복 연산을 줄여야 합니다.

대규모 순차 추천 모델의 서빙이 어려운 이유

대규모 순차 추천 모델의 서빙은 작은 밀집 랭킹 모델의 서빙과 다릅니다. 서빙 스택은 가변 길이(jagged) 시퀀스 입력, 대규모 범주형 임베딩 상태, 긴 이력을 처리해야 합니다. 같은 사용자가 새로운 정보를 조금만 추가한 채 반복해서 요청하는 패턴도 감당해야 합니다. 요청마다 사용자 이력 전체의 키-값 상태를 다시 계산하면 연산이 낭비되고 지연 시간이 늘어납니다.

이때 중요해지는 것이 KV 캐싱입니다. KV 캐시는 이전 시퀀스 연산에서 얻은 재사용 가능한 키-값 데이터를 저장하여, 모델이 캐시된 이력 구간을 다시 계산하지 않게 합니다. 추천 추론에서는 사용자의 장기 이력이 대부분 그대로이고 새로운 후보 아이템이나 최근 행동만 들어올 때 특히 유용합니다.

NVIDIA HSTU 추론 워크플로에는 GPU 메모리와 호스트 스토리지를 KV 데이터 캐시로 활용하는 KVCacheManager가 있습니다. GPU 캐시는 페이지 단위 KV 데이터 테이블로 구성되며 조회, 할당, 추가, 퇴출을 지원합니다. GPU 캐시 공간이 부족하면 LRU 방식 정책에 따라 오래된 사용자를 퇴출할 수 있습니다. 호스트 측 스토리지는 캐시된 KV 데이터를 담는 또 하나의 계층이며, 이 워크플로에는 KV 캐시 런타임용 FlexKV 기반 백엔드가 포함됩니다.

HSTU 어텐션 커널은 페이지 단위 캐시에서 KV 데이터를 읽을 수 있습니다. 내보낸 추론 경로에는 조회, 할당, 온보딩, 추가, 오프로딩을 위한 캐시 인식 커스텀 연산이 포함되어 있습니다. 덕분에 서빙 경로는 모델의 시퀀스 의미론을 유지하면서 중복 연산을 줄일 수 있습니다.

네이티브 추론을 위한 PyTorch AOTI

PyTorch AOTI(Ahead-of-Time Inductor) 워크플로는 PyTorch 모델에서 시작하여 torch.export와 PyTorch AOTI로 모델을 내보냅니다. AOTI는 모델을 사전에 컴파일하여 네이티브 C++ 런타임이 로드할 수 있는 패키지로 만듭니다. 이렇게 하면 Python 런타임 오버헤드가 줄고, Dynamo-Triton의 PyTorch AOTI 백엔드에 적합한 배포용 아티팩트를 얻을 수 있습니다.

내보낸 모델 패키지에는 AOTI 모델 아카이브와 함께 메타데이터, 임베딩 테이블 파일이 들어 있습니다. NVIDIA 예제의 임베딩 구현은 DynamicEmb 추론 임베딩 테이블과 NV 임베딩 캐시를 결합합니다. NV 임베딩 캐시는 자주 쓰이는 임베딩만 GPU 메모리에 두고 전체 테이블은 CPU 메모리에 유지하여 GPU 메모리 사용량을 줄입니다. 내보내기 경로는 컴파일된 .pt2 아카이브와 함께 레이어 메타데이터와 임베딩 테이블 데이터를 기록하므로, 임베딩 테이블을 불필요하게 복제하지 않고도 모델을 로드할 수 있습니다.

이 워크플로는 동일한 내보내기 아티팩트를 여러 방식으로 검증합니다. Python 내보내기 스크립트가 패키지와 리플레이 텐서를 생성하고, 네이티브 C++ 실행 파일이 내보낸 모델을 로드하여 리플레이하면서 정확성과 성능을 검증합니다. Dynamo-Triton 배포도 같은 AOTI 패키지와 리플레이 경로를 사용하므로, 개발 단계의 검증과 프로덕션 서빙이 일치하도록 유지할 수 있습니다.

Dynamo-Triton 배포 경로

Dynamo-Triton은 내보낸 HSTU 모델을 위한 프로덕션 서빙 계층을 제공합니다. AOTI 배포에는 platform: "torch_aoti"를 지정한 Dynamo-Triton PyTorch 백엔드를 사용합니다. 이를 통해 Dynamo-Triton은 사전 컴파일된 PyTorch 모델 패키지를 로드하여 서빙할 수 있습니다.

전체 워크플로는 다음 다섯 단계로 이루어집니다.

  • 필요한 커스텀 연산자와 런타임 라이브러리 빌드
  • PyTorch AOTI로 HSTU 랭킹 모델 내보내기
  • FlexKV 기반 KV 캐시 서비스 시작
  • 네이티브 C++ 리플레이로 내보낸 아티팩트 검증
  • Dynamo-Triton으로 내보낸 KV 캐시 AOTI 모델 서빙

추천 시스템 서빙에는 빠른 모델 커널 이상의 요소가 필요하기 때문에 이 접근법이 중요합니다. Dynamo-Triton은 모델 리포지토리 관리, 요청 처리, 백엔드 통합, 메트릭, 배포 구조를 제공합니다. AOTI는 오버헤드가 낮은 컴파일된 모델 아티팩트를 제공하고, NV 임베딩 캐시는 임베딩 테이블에서 자주 쓰이는 부분만 GPU 메모리에 두어 GPU 메모리 요구량을 낮춥니다. FlexKV는 어텐션 블록을 캐싱하여 긴 사용자 이력의 재연산을 줄입니다. 이들이 함께 현실적인 생성형 추천 추론을 겨냥한 서빙 스택을 이룹니다.

HSTU 서빙 지연 시간 벤치마크

이 벤치마크는 Dynamo-Triton 백엔드, 모델 크기, 배치 크기, KV 캐시 상태에 따라 HSTU 서빙 지연 시간을 비교합니다.

목표는 프로덕션 HSTU 서빙 스택의 성능 이점을 정량화하는 것입니다. Dynamo-Triton의 PyTorch AOTI 백엔드를 Python 백엔드와 비교하고, GPU KV 캐싱으로 지연 시간이 추가로 얼마나 줄어드는지 측정하며, 이러한 이점이 모델 깊이와 배치 크기에 따라 어떻게 달라지는지 평가합니다. 이를 통해 캐시를 쓰지 않는 Python 기반 추론에서 Dynamo-Triton을 이용한 컴파일 및 캐시 인식 HSTU 배포로 전환할 때 기대할 수 있는 성능을 개발자에게 보여 줍니다.

recsys-examples의 벤치마크 결과는 단일 GPU에서 KuaiRand-1K 랭킹 구성으로 측정했습니다. 모델 구조는 3계층과 8계층 HSTU 변형, 히든 크기 512, 어텐션 헤드 4개, BF16 모델 가중치, BF16 KV 캐시, 최대 이력 시퀀스 길이 총 8,192(아이템과 행동 쌍 4,096개)의 이력 스트림, 최대 후보 시퀀스 길이 100, 컨텍스트 피처 6개입니다. 정렬 전 유효 시퀀스 길이는 8,298개 토큰이며, 최대 정렬 시퀀스 길이 8,320으로 내보냈습니다.

벤치마크 프로토콜은 논리 요청당 지연 시간을 보고합니다. Dynamo-Triton AOTI 벤치마크에서는 Dynamo-Triton 호출 한 번에 논리 배치 하나가 들어가며, 논리 요청당 지연 시간은 엔드투엔드 패스 시간을 Dynamo-Triton 호출 횟수와 논리 배치 크기의 곱으로 나누어 계산합니다. 데이터셋 로딩, 검증, 리배칭, 사용자 ID 생성, 서버 시작, 워밍업, 워밍업 이후 대기 시간은 측정 시간에서 제외했습니다.

Dynamo-Triton 백엔드 비교와 배치 크기별 결과에는 NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU를 사용했습니다.

Dynamo-Triton 백엔드 비교

Dynamo-Triton 배치 크기 2에서는 KV 캐시 적중이 없어도 PyTorch AOTI가 Dynamo-Triton Python 백엔드보다 지연 시간을 개선합니다. 20 GB GPU KV 캐시에서 적중이 발생하면 지연 시간 개선 폭이 훨씬 커집니다.

이 결과에서는 서로 다른 두 가지 이득이 드러납니다. 첫째, AOTI는 Python 백엔드보다 서빙 오버헤드를 줄입니다. 둘째, KV 캐시 적중은 캐시된 시퀀스 상태를 재사용하여 모델이 처리할 연산량을 줄입니다. 재연산을 피하는 효과가 더 큰 8계층 모델에서 캐시의 이점이 특히 뚜렷합니다.

AOTI와 KV 캐시의 배치 크기별 확장성

배치 크기별 PyTorch AOTI 백엔드 결과를 보면, 논리 배치 크기가 커질수록 KV 캐싱의 효과도 커집니다.

배치 크기 8에서 3계층 HSTU 모델은 GPU KV 캐시 적중 시 논리 요청당 지연 시간 0.423 ms에 도달합니다. 8계층 HSTU 모델은 0.678 ms에 도달합니다. 긴 시퀀스 랭킹 추론에서 우수한 결과이며, 컴파일된 모델 실행과 캐시 인식 서빙을 결합하는 방식의 가치를 보여 줍니다.

HSTU GR 추론 가속의 이점

추천 시스템은 엄격한 지연 시간 예산 안에서 동작합니다. 랭킹 지연 시간이 늘어나면 페이지 로딩 시간, 피드 반응성, 광고 서빙 마감 시간에 영향을 줄 수 있습니다. 그런데 시퀀스를 인식하고 개인화 수준이 높아지는 모델일수록 더 많은 추론 연산이 필요할 수 있으며, 사용자 이력이 길어질수록 특히 그렇습니다.

HSTU 서빙 워크플로는 서로 보완하는 기술을 결합하여 이 문제를 해결합니다. HSTU가 생성형 추천 아키텍처를 제공하고, PyTorch AOTInductor가 사전 컴파일된 배포 아티팩트를 생성합니다. FlexKV 기반 KV 캐싱은 이전에 계산한 어텐션 상태를 재사용하게 하고, NVIDIA Dynamo-Triton은 프로덕션 서빙 환경을 제공합니다.

이 접근법은 연속된 요청에서 사용자 상호작용 이력의 앞부분이 변하지 않을 때 특히 유용합니다. 해당 구간의 어텐션을 다시 계산하는 대신 캐시된 키-값 상태를 재사용하고, 새로 추가된 토큰에 필요한 부분만 계산하면 됩니다. 시퀀스가 길고 모델이 깊을수록 절감 효과가 커집니다. 그렇지 않으면 여러 HSTU 계층에 걸친 반복 연산이 상당한 지연 시간을 더하기 때문입니다.

HSTU GR 추론 가속 시작하기

이 워크플로는 NVIDIA/recsys-examples GitHub 리포지토리에서 재현하고 확장할 수 있습니다. HSTU 개요 문서는 컨텍스트 토큰, 아이템 토큰, 행동 토큰, 임베딩 테이블, HSTU 블록, 예측 헤드를 포함한 GR 모델 구조를 소개합니다.

AOTI 추론 가이드는 필요한 이미지와 라이브러리 빌드, KuaiRand-1K 데이터 준비, 체크포인트 학습, KV 캐시 AOTI 모델 내보내기, C++ 리플레이를 통한 검증, Dynamo-Triton 런타임 이미지 패키징, Dynamo-Triton 서버를 통한 요청 리플레이까지 단계별로 안내합니다.

더 알아보려면 다음 관련 자료를 참고하십시오.

감사의 말

이 글은 여러 NVIDIA 팀이 협력한 결과물입니다. J, Runchu Zhao, Yulu Liu, Lin Hu, Zhuofan Li, Jacob Subag, Tomer Bar-On의 기여에 감사드립니다.

Discuss (0)

Tags