NVIDIA, SSM 상태 쓰기 최소화하는 ReplaySSM 공개

NVIDIA와 프린스턴대 Dao AI Lab이 ReplaySSM 기법을 공개했습니다. 이 기법은 하이브리드 모델의 SSM 레이어에서 매 스텝마다 HBM에 상태를 저장하는 병목을 해결하기 위해, 상태를 체크포인트와 최근 입력 버퍼로 분리해 표현합니다.

핵심 메커니즘

• 비플러시(Non-flush) 스텝: 버퍼에 입력을 추가하고 출력만 계산하며 상태 저장을 생략합니다.
• 플러시(Flush) 스텝: 버퍼 내용을 체크포인트에 반영해 HBM에 한 번만 저장합니다.
• 투기적 디코딩 롤백: 링 버퍼 포인터 이동으로 처리하여 상태 스냅샷 저장 오버헤드를 제거합니다.

성능 개선 효과

vLLM 튜닝 기준선 대비 FP32 상태, 길이 16 조건에서 다음과 같은 성능 향상을 달성했습니다.

• 표준 디코딩: 배치 256에서 SSM 커널 1.46~1.80배, 디코드 스텝 전체 1.20~1.47배 가속.
• 투기적 디코딩: 배치 512에서 Nemotron-3-Super 1.85배, Qwen3.5-122B 2.10배 처리량 향상.
• 검증 지연: vLLM의 T=8 지연이 4.66배인 반면, ReplaySSM은 0.89~1.21배 수준으로 유지.

양자화 및 정확도

읽기용 저정밀도(NVFP4)와 재귀용 고정밀도(FP32) 상태를 분리하여 오차 누적을 방지했습니다. NVFP4 상태 사용 시 GPQA-Diamond 정확도 손실을 1.77~2.36점으로 줄였습니다. Mamba-2 non-flush 커널은 NVFP4 읽기 시 FP32 대비 3.8배 빠릅니다.

구현 및 제한사항

• vLLM v0.31.0: Mamba-2 표준 디코딩(--use-replayssm)만 지원.
• ReplaySSM GitHub: GDN 경로 및 저정밀도 커널은 연구용 포크에만 존재.
• 적용 조건: 배치가 클수록 이득이 커지며, 배치 1 로컬 추론에서는 오히려 느려집니다.