선형 어텐션 반복 상태를 6비트로 압축해 서빙 메모리를 68.7% 줄인다
STEPQuant: When and Where Errors Matter in Delta-Rule Recurrent State Quantization
무엇인가
선형 어텐션은 시퀀스 길이에 따라 커지는 KV 캐시 대신 고정 크기 반복 상태에 과거 토큰을 요약한다. Qwen3.8-27B와 Kimi-Linear-48B-A3B-Instruct 같은 하이브리드 모델은 게이트 델타 규칙 기반 반복 메모리(GDN, KDA)를 표준 어텐션과 섞어 쓴다. 문제는 상태 크기가 컨텍스트 길이와 무관하게 고정이어도, 동시 요청마다 별도의 영속 상태가 필요하고 서빙 시스템이 캐싱·스케줄링용 슬롯을 추가로 예약한다는 점이다. 논문에 따르면 공식 SGLang 배포에서 Qwen의 FP32 상태 풀은 동시 요청 70건에서 BF16 가중치 메모리를 초과한다. 그런데 반복 상태에 균일 양자화를 그대로 적용하면 4비트와 6비트에서 평균 정확도가 급락하고 8비트에서도 격차가 남는다.
어떻게 동작하나
논문은 이 실패를 시간과 공간 두 축으로 분석한다. 시간적으로는 매 디코딩 스텝마다 상태가 갱신되고 다시 양자화되면서 오차가 누적된다. 논문의 Proposition 1은 누적 오차를 E_t = A_t·E_{t-1} + ε_t (A_t = (I − β_t k_t k_t^T)·D_t)로 정식화하고, ‖k_t‖₂ ≤ 1, 0 ≤ β_t ≤ 1, 0 ⪯ D_t ⪯ I 조건에서 ‖A_t‖₂ ≤ ‖D_t‖₂ ≤ 1임을 보인다. 리텐션 게이트 D_t가 이전 오차를 감쇠시키고, 델타 업데이트가 현재 키 방향 성분을 한 번 더 줄이지만 키에 직교하는 성분은 그대로 남는다. 리텐션이 1에 가까우면 오차가 수십 스텝을 버틴다. 실측에서도 Qwen3.8-27B 전체 반복 헤드를 INT6로 양자화했을 때 게이트 반감기가 긴 헤드일수록 누적 상태 오차가 컸고, Spearman 상관계수는 약 0.80이었다. 공간적으로는 같은 크기의 오차라도 어느 키 로우에 있느냐에 따라 리드아웃 영향이 달라진다. 리드아웃 오차는 Δy_t = Σ_i (A_t^T q_t)_i · E_{t-1,i,:}^T로 분해되고, 로우 임팩트 점수 ω_i = E[g_{t,i}²] (g_t = A_t^T q_t)가 그 가중치다. 여기에 더해 상태 행렬은 키 로우와 밸류 컬럼 양쪽 축에서 큰 크기 편차를 보인다. 최대 대 중앙 RMS 비율이 키 로우 축 10.3배, 밸류 컬럼 축 19.4배였고, 샘플링한 상태의 98.6%에서 두 축 모두 3배를 넘었다.
무엇과 다른가
STEPQuant는 이 두 관찰을 각각 담당하는 두 구성 요소로 이뤄진다. 첫째, Lifetime-aware Bit Allocation은 반복 상태 유닛( Qwen에서는 헤드 전체, KDA에서는 키 로우)마다 b비트 재구성 왜곡 d_u(b)와 평균 로그 리텐션 ℓ_u를 추정하고, 수명 가중치 L_u = Σ_{j=0}^{H-1} exp(2jℓ_u)를 곱한 목적 Σ_u L_u·d_u(b_u)를 평균 비트 예산 제약 Σ_u n_u b_u ≤ b̄ Σ_u n_u 아래에서 최소화한다. 즉 오차가 크고 오래 사는 유닛에 더 많은 비트를 준다. 그래도 남는 최고 위험 유닛은 FP16 피벗으로 보존한다. 둘째, Key-Row-Aware Dual-Axis Fitting은 상태를 X̂_ij = r_i·c_j·z_ij로 표현한다. 로우 스케일은 r_i = m_i^{1/2}·w_i^{−1/2}로, m_i는 해당 로우의 평균 절대값(크기가 큰 로우에 넓은 범위를 주기 위함), w_i는 임팩트 점수 ω_i를 γ=0.25로 완화해 헤드 내 기하평균이 1이 되도록 정규화한 항(영향이 큰 로우에 더 촘촘한 분해능을 주기 위함)이다. 컬럼 스케일 c_j는 임팩트 가중 재구성 오차 Σ_{i,j} w_i²(X_ij − r_i c_j z_ij)²를 최소화하도록 피팅해, 영향이 큰 로우의 오차에 더 큰 페널티를 준다. 구현에서는 SGLang의 반복 상태 풀에 패킹된 상태 커널로 통합해, 타일 단위 상태 복원·델타 업데이트·현재 리드아웃을 한 커널로 융합하고 스케일 피팅과 패킹 라이트백은 별도 CUDA 스트림에서 모델 연산과 겹쳐 실행한다.
어떻게 쓰나
실험은 Qwen3.8-27B(GDN)와 Kimi-Linear-48B-A3B-Instruct(KDA)를 BF16 및 4비트 AWQ 가중치로, SGLang과 NVIDIA A800 4장에서 돌렸다. 캘리브레이션은 WikiText-2 학습 세그먼트 32개(각 2048 토큰)로 비트 할당·FP16 피벗·로우 임팩트 점수를 정한다. 긴 생성 추론 7개(LiveCodeBench v6, EvalPlus, AIME 2026, MATH-500, HMMT February 2026, GPQA Diamond, IFBench)와 짧은 생성 6개(MMLU, ARC-C, OpenBookQA, HellaSwag, WinoGrande, LAMBADA)에서 평가했다. 6비트 STEPQuant의 7과제 평균 정확도는 Qwen 80.59%, Kimi 61.47%로 FP32 상태 기준선에 근접했고, 균일 INT6는 각각 45.04%, 45.70%에 그쳤다. 짧은 생성에서는 4비트 STEPQuant가 FP32 대비 Qwen 0.15점, Kimi 0.25점 차이에 불과했지만 균일 INT4는 크게 무너져 유효한 답을 내지 못하는 경우도 있었다. 4비트 AWQ 가중치와 조합한 실전 설정에서도 6비트 STEPQuant는 7과제 평균 79.27%(Qwen), 58.62%(Kimi)로 각 FP32 상태 기준선보다 0.05점, 0.33점만 낮았다.
전제와 한계
구성 요소 분해 실험도 인상적이다. 4비트에서 공간 피팅만 적용한 경우 Qwen 73.95%로, Mamba용으로 설계된 Q-Mamba의 이중축 상태 양자화(DSQ) 7.64%를 크게 앞섰다. FP16 피벗은 Qwen 헤드의 1.39%만 보호하고도 4비트 3과제 평균을 6.70점, 6비트 AIME 정확도를 14.16점 끌어올렸다. 두 요소를 합치면 4비트 STEPQuant가 84.72%로 공간만(73.95%)·시간만(12.87%)을 모두 앞서고, 6비트에서는 84.51%로 개별 요소 80.19%, 75.63%와 FP32의 84.61%에 맞먹는다. 또한 균일 양자화는 과잉 사고를 유발해 Kimi가 4비트에서 AIME 평균 63.40K, HMMT 64.22K 토큰을 생성하면서도 정확도는 거의 0이었던 반면, STEPQuant는 FP32 상태 모델과 비슷한 출력 길이를 유지했다. 서빙 효율은 배치 512, Qwen W4 기준 총 메모리 419.73GiB에서 131.18GiB로 68.7% 감소했고, 반복 상태 메모리는 80.1% 줄어 5.03배 압축, 상태 업데이트 시간은 65.6% 줄어 2.91배 빨라졌다. Kimi는 5.08배 압축과 총 메모리 53.7% 감소를 기록했다.
개발자 관점에서 이 논문이 유효한 지점은 명확하다. 선형 어텐션이나 하이브리드 어텐션 모델을 높은 동시성으로 서빙하면서 상태 풀이 메모리 병목이 된 경우, 가중치 양자화만으로는 해결되지 않는 부분을 건드린다. 특히 상태를 무조건 균일 비트로 내리면 정확도가 붕괴한다는 점, 그리고 그 원인이 시간적 누적과 공간적 편차 두 가지라는 진단은 다른 반복 구조 모델을 다룰 때도 그대로 참고할 만하다. 코드는 공개돼 있고 SGLang 통합 커널까지 포함한다. 다만 도입 전에 확인할 것은 캘리브레이션 데이터 의존성이다. 비트 할당과 피벗, 로우 임팩트 점수는 오프라인 캘리브레이션에서 한 번 정해져 모든 요청에 고정되므로, 서비스 도메인이 WikiText-2와 크게 다르면 재캘리브레이션이 필요할 수 있다.
저자들이 밝힌 전제와 한계도 분명하다. 6비트라는 수치는 "명목상(nominal)" 예산이며 FP16 피벗이 섞이므로 실제 상태당 비트는 이보다 높다. 평가는 두 모델과 지정된 벤치마크에 한정되고, 4비트에서는 Qwen과 달리 Kimi에서 FP32 대비 격차가 남는다. 관련 연구로 소개된 동시 작업 DAMP는 감쇠 기반 지속성으로 일부 채널을 FP16으로 보호하고 나머지를 INT8로 양자화해 상태값당 9.9비트에서 정확도를 유지했다고 보고하는데, 이 논문은 그보다 낮은 6비트가 가능함을 보이는 것을 기여로 내세운다.