ReplaySSM: 매 토큰마다 상태를 저장하지 않아 하이브리드 모델을 빠르게 하기

배경과 문제

하이브리드 언어 모델은 어텐션 레이어 사이에 상태-공간 모델(SSM) 레이어를 끼워 넣는다. 어텐션의 KV 캐시는 지나간 토큰의 키와 값을 보관하므로 입력이 길어질수록 커진다. SSM은 과거 정보를 고정 크기 상태 행렬로 요약해 이 증가를 줄이지만 정확한 회상에는 제약이 있다. Nemotron 3과 Qwen3.5 같은 모델은 두 방식을 섞어 효율성과 회상 능력을 함께 추구한다.

SSM 디코딩은 매 스텝 상태를 읽고, 감쇠와 랭크 1 갱신을 적용한 뒤, 출력을 계산하고, 상태를 다시 저장한다. 이 과정에서 상태는 두 가지 일을 한다. 하나는 이번 토큰의 출력을 만들기 위해 읽히는 일이고, 다른 하나는 다음 토큰의 상태가 되기 위해 갱신되는 일이다. 기존 디코딩은 이 둘을 한 번에 처리한다. 그 결과 세 가지 문제가 생긴다. 첫째, 연산량에 비해 메모리 이동이 커서 커널이 HBM 대역폭에 묶인다. 둘째, 투기적 디코딩에서 거절된 토큰을 되돌릴 값싼 방법이 없다. 셋째, 상태를 낮은 정밀도로 저장하면 반올림 오차가 재귀 갱신으로 다시 들어가 누적된다.

기존 접근의 한계

상태의 크기가 고정돼도 매번 옮기는 비용은 작지 않다. SSM 디코딩은 상태 전체에 비교적 적은 연산만 수행하므로, 큰 배치에서 계산 속도보다 HBM이라는 GPU 고대역폭 메모리의 이동 속도에 묶인다. Nemotron-3-Super-120B-A12B-NVFP4를 B300 한 장에서 배치 256으로 측정했을 때, 컨텍스트 약 100K 토큰까지 SSM이 디코드 스텝의 가장 큰 구성 요소였다. 이는 해당 모델·환경에서의 측정이지 모든 하이브리드 모델의 고정 경계는 아니다.

투기적 디코딩은 드래프터가 여러 토큰을 먼저 제안하고 대상 모델이 한꺼번에 검증해 수용한 부분만 남기는 방식이다. 어텐션은 KV 캐시의 위치만 되돌리면 거절된 토큰을 버릴 수 있다. SSM은 그 토큰들이 이미 상태에 섞인 뒤라 복구가 어렵다. 기존 vLLM은 드래프트마다 상태 스냅샷을 저장하므로 쓰기 비용과 메모리 점유가 커져 큰 배치에서 가속 효과가 줄어든다.

상태를 낮은 정밀도로 저장하면 이동량은 줄지만 반올림한 값이 다음 갱신의 입력이 되어 오차가 반복 유입된다. 같은 값을 읽기만 하는 KV 캐시와 다른 점이다. 이 때문에 낮은 정밀도의 가중치를 써도 SSM 상태는 FP32로 두는 모델이 있다. 저자는 단순히 상태를 압축하는 대신, 현재 출력을 만드는 역할과 미래 상태를 갱신하는 역할을 나눈다.

연구방법을 단계별로

첫 단계는 여러 토큰의 재귀 갱신식을 펼치는 것이다. 현재 상태는 과거 체크포인트에 누적 감쇠를 적용한 부분과 최근 입력들이 더한 부분으로 다시 표현할 수 있다. ReplaySSM은 큰 상태 행렬을 매번 저장하는 대신, 체크포인트 하나와 감쇠값·키·값 같은 최근 입력 벡터를 담은 작은 버퍼를 유지한다. 이전 입력을 버리는 요약 변경이 아니라 같은 계산을 다른 순서로 구성하는 접근이다.

대부분의 스텝(non-flush 스텝)에서는 현재 입력을 버퍼에 덧붙이고, 체크포인트와 버퍼로 출력을 계산할 뿐 상태를 저장하지 않는다. 버퍼에 기록이 L개 차면 flush 스텝이 버퍼의 입력을 체크포인트에 접어 넣고, 새 체크포인트를 HBM에 한 번 쓴 뒤 버퍼를 비운다. 출력은 부동소수점 오차 범위 안에서 기존 디코딩과 같다. GDN에는 조정이 하나 필요하다. GDN의 갱신은 이전 상태에 의존하는 보정값 u_t = β_t(v_t - α_t S_{t-1} k_t)를 포함하므로, 원시 입력 v_t 대신 이 u_t를 버퍼에 저장한다. 그러면 GDN도 Mamba-2와 같은 외적 형태로 펼칠 수 있다.

투기적 디코딩에서 ReplaySSM은 검증 전에 드래프트 토큰 T개의 입력을 모두 버퍼에 덧붙이고, 검증이 끝나면 수용된 개수만큼 버퍼 포인터를 앞으로 옮긴다. KV 캐시처럼 포인터 뒤의 기록은 무시되고 다음에 덧붙일 때 덮어써진다. 드래프트마다 저장할 상태 스냅샷도, 복원할 상태도 없다. flush는 이미 확정된 기록만 체크포인트에 접어 넣으므로 수용 길이를 기다리지 않고 실행할 수 있다. 구현은 링 버퍼(용량 L+T)를 써서 확정된 앞부분을 버릴 때 데이터를 복사하지 않고 시작 위치만 옮긴다.

출력 전용 평가(output-only evaluation)는 중간 상태 행렬을 만들지 않고 체크포인트와 입력 기록으로 출력을 직접 계산한다. 예를 들어 벡터 v와 k로 큰 행렬을 만든 뒤 q를 곱하는 대신, 먼저 k와 q의 내적으로 숫자 하나를 구해 v에 곱할 수 있다. 여러 드래프트에도 이를 적용하면 체크포인트를 공유하고 작은 행렬곱으로 검증을 처리할 수 있다. GDN의 드래프트 간 의존성은 작은 삼각 행렬 풀이로 처리한다.

저정밀도 상태를 위해 ReplaySSM은 두 가지 정밀도를 분리한다. 재귀용으로는 고정밀도 체크포인트를, non-flush 읽기용으로는 저정밀도 복사본을 둔다. flush는 고정밀도 체크포인트에서 시작해 버퍼의 입력을 반영한 뒤, 결과를 두 형식으로 각각 변환해 저장한다. 저정밀도 복사본은 재귀로 다시 읽히지 않으므로 그 반올림 오차는 현재 구간의 출력에만 영향을 준다. GDN에서는 non-flush 스텝에서 계산한 보정값 u가 저정밀도 복사본에서 나온 것이라 오차를 품고 있으므로, 원시 입력 v도 함께 캐시하고 flush 때 고정밀도 상태로 보정값을 다시 계산한다.

결과와 숫자의 의미

성능 실험은 Mamba-2를 쓰는 Nemotron 3 모델 두 개와 GDN을 쓰는 Qwen3.5 모델 두 개를 대상으로 했다. 각 모델은 H100 또는 B300 한 장에서 텐서 병렬 없이 실행했고, 가장 빠른 실행 설정으로 튜닝한 vLLM을 기준선으로 썼다. 일반 성능 비교는 FP32 SSM 상태·버퍼 길이 16이며, 이후의 저정밀도 정확도 평가와 구별해야 한다.

배치 256에서 SSM 커널은 1.46~1.80배, 전체 디코드 스텝은 1.20~1.47배 빨라졌다. 커널 수치에는 버퍼를 체크포인트에 합치는 flush 비용도 포함된다. 어텐션과 다른 계산은 그대로이므로 전체 이득은 커널 이득보다 작다. 배치 1에서는 전체적으로 1~4% 느렸다. 요청이 적으면 원래 커널부터 메모리 대역폭을 충분히 쓰지 않아, 상태 이동을 줄인 효과보다 추가 계산 부담이 클 수 있다.

투기적 디코딩은 120B급 두 모델의 NVFP4 가중치·FP32 상태, 모델당 B300 한 장, GSM8K 입력에서 요청당 256토큰 생성, 검증 길이 4와 버퍼 16이라는 조건이다. 배치 512의 처리량은 기존 vLLM 투기적 디코딩보다 Nemotron-3-Super에서 1.85배, Qwen3.5-122B에서 2.10배였다. 표준 디코딩과 비교하면 각각 2.16배와 1.90배다. 검증 커널 개선과 스냅샷 메모리 절감이 기여했으며 수용 길이를 바꿔 얻은 이득은 아니다.

저정밀도 상태의 KL 발산은 PG-19 책 네 권에서 같은 정답 토큰을 입력하며 최대 128K 토큰까지 출력 분포를 비교했다. Qwen3.5-4B에서 읽기와 재귀에 NVFP4를 함께 쓰는 경우보다, 읽기는 NVFP4·재귀는 FP32로 나누면 128K 지점의 KL이 0.136에서 0.0038로 약 36배 작아졌다. 이는 FP32 출력 분포와의 차이가 줄었다는 뜻이지 정답률이 36배 올랐다는 뜻이 아니다. FP16과 확률적 반올림을 쓴 재귀도 이 평가에서 FP32에 가까웠다.

다운스트림 정확도는 MMLU-Pro, GPQA-Diamond, Math-L5에서 측정했다. 공유 NVFP4+SR 체크포인트는 GPQA-Diamond에서 4개 모델 평균 13.05점, MMLU-Pro에서 5.88점, Math-L5에서 10.92점을 잃었다. NVFP4 읽기와 FP16+SR 재귀를 분리하면 손실이 각각 2.36점, 0.15점, 0.10점으로 줄고, FP8+SR 재귀를 쓰면 1.77점, 0.40점, 0.38점으로 줄어든다. 모델별로 보면 가장 작은 Nemotron-3-Nano-4B는 분리 구성에서도 GPQA-Diamond에서 5.6~6.7점을 잃는다. 나머지 세 모델은 분리 구성에서 FP32 대비 약 2점 이내를 유지한다. GPQA-Diamond는 칸마다 594개 생성에서 문항 단위 표준 오차가 2.5~3점이라, 큰 차이만 구분된다.

구체적인 예시

Nemotron-3-Super-120B-A12B-NVFP4의 GPQA-Diamond 문항을 FP32 출력 길이 기준 사분위로 나눈 그래프를 보면, 공유 NVFP4 체크포인트는 짧은 두 구간에서는 FP32와 0~2점 차이지만, 긴 두 구간(중앙값 10.7K, 61K 토큰)에서는 20.9점과 28.2점 뒤처진다. NVFP4 읽기와 FP8+SR 재귀 구성은 같은 구간에서 2.0점과 3.4점 차이였다. 이는 오차가 생성 길이에 따라 누적된다는 것을 보여준다.

같은 Nemotron 모델에서 공유 NVFP4 생성의 20.7%는 출력 토큰 상한에 걸려 오답으로 처리됐고 FP32는 5.7%였다. 둘 다 완료한 문항만 보면 손실은 2.8점이다. 따라서 전체 점수 손실에는 답의 내용뿐 아니라 생성이 길어져 완료하지 못한 영향도 섞여 있다. 반면 약 130K 토큰 입력의 RULER-128K에서는 모든 구성이 FP32 대비 0.3점 이내였다. 양자화가 디코딩에만 적용된 이번 실험에서는 긴 입력과 긴 출력을 구별해야 한다.

한계와 오해하기 쉬운 점

이득은 큰 배치에서 SSM의 메모리 이동이 병목인 경우에 집중된다. 배치 1의 로컬 추론에는 거의 도움이 없거나 조금 느리고, 매우 긴 컨텍스트에서 어텐션 비중이 커지면 전체 효과가 줄어든다. 투기적 디코딩도 배치와 검증 길이에 따라 달라지므로 대표 배치의 최고값만 자신의 서비스에 적용할 수 없다.

정확도와 KL은 Triton에서 에뮬레이션한 양자화로 측정했고, 실제 낮은 정밀도 저장의 커널 지연은 별도 CuTe DSL 구현으로 측정했다. 정확도와 속도를 하나의 완성된 서빙 구현에서 동시에 측정한 결과가 아니다. NVFP4 읽기 3.83배는 배치 256의 Mamba-2 non-flush 커널 수치이지 전체 서비스 속도가 아니다. 고정밀도 갱신 경로와 현재 출력에도 오차가 남아 FP32와 비트 단위로 동일하다는 보장은 없다.

공식 코드는 연구용 vLLM 포크다. 현재 README는 upstream에 Mamba-2 표준 디코딩 커널과 FlashInfer의 표준·투기적 경로가 있으며, 나머지 Triton 커널·에뮬레이션 양자화·저정밀도 CuTe는 연구용 포크에 있다고 설명한다. main 브랜치와 출시 버전의 지원 범위가 같다는 뜻은 아니다. 프레임워크 이름만 보고 모든 모델과 정밀도 모드가 지원된다고 판단해서는 안 된다.

6월 Dao AI Lab 글과 10월 논문은 최적 버퍼 크기와 평가 모델·GPU 구성·수치가 다르다. 이번 설명의 성능 수치는 10월 논문의 조건을 따른다. 또한 네 모델의 결과를 모든 SSM이나 일반 Transformer에 확대해서는 안 되며, 전체 배포 비용이나 전력 효율 향상을 직접 측정한 연구도 아니다.

실제로 어떻게 활용할까

검토할 만한 상황은 많은 요청을 동시에 처리하는 하이브리드 모델 서빙이다. 먼저 현재 작업에서 SSM 커널이 차지하는 시간과 메모리 부담을 측정해야 한다. 작은 배치의 대화형 추론인지 큰 배치의 처리량 중심 서비스인지에 따라 선택이 달라진다. 논문의 배치 32·128·256 결과처럼 규모가 바뀌면 이득도 달라진다.

적용할 때는 모델의 SSM 종류, 프레임워크의 정확한 버전과 커널 경로, 투기적 디코딩 지원을 함께 확인해야 한다. 저정밀도 상태 분리는 정확도 손실을 줄이는 연구 결과이며, 자신의 작업에서 긴 출력의 완료율과 정답률까지 검증해야 한다. 읽기 정밀도를 낮추면 항상 같은 정확도로 더 빨라진다는 뜻은 아니다.

운영 사례로 언급되는 GLM-5.3-Flash와 Kimi-K3도 같은 성능 표를 그대로 적용할 대상은 아니다. 공식 README는 GLM의 서빙 적용을 설명하고, Kimi는 드래프트 입력을 캐시하는 동시기 접근으로 인용했다고 구별한다. 연구의 핵심 활용법은 한 번 읽기 위한 복사본과 계속 갱신할 상태를 나눠 생각하는 설계 원리이며, 실제 효과는 각 배포 환경에서 측정해야 한다.