prefill과 decode에 서로 다른 양자화를 배정하는 disaggregated quantization

무엇인가

LLM 추론은 프롬프트를 처리하는 prefill과 토큰을 하나씩 생성하는 decode로 나뉜다. 이 논문의 출발점은 두 단계의 병목이 정반대라는 관찰이다. H×H 가중치와 T×H 활성값에 대해 선형 계층은 O(H²+TH)개의 원소를 옮기고 O(TH²)번의 연산을 수행한다. 토큰 수 T가 큰 prefill에서는 가중치 전송 비용이 연산에 묻혀 compute-bound가 되고, 배치 1(T=1)의 decode에서는 가중치 하나가 곱셈-덧셈 한 번에만 쓰여 메모리 대역폭이 병목이 된다. 그래서 기존 양자화는 둘 중 하나만 노린다. 가중치만 압축하는 weight-only 인코딩은 decode 트래픽을 줄이지만 prefill 연산을 가속하지 못하고, INT4나 NVFP4 같은 하드웨어 네이티브 연산 포맷은 prefill을 빠르게 하는 대신 decode 정확도를 깎는다. 대규모 서빙에서는 이미 prefill과 decode를 별도 가속기에 올리는 disaggregated serving이 쓰이지만, 그 분리는 통신과 스케줄링에 관한 것이었고 양자화 자체는 두 단계가 공유했다. 이 논문은 바로 그 지점을 파고든다.

어떻게 동작하나

제안하는 학습 절차는 QADD(Quantization-aware distillation with disaggregation)다. SFT 데이터의 라벨 마스크를 손실 마스킹에만 쓰지 않고 계산 경로 선택에도 사용한다. 사용자 턴(prompt)은 prefill 경로로, 어시스턴트 턴(response)은 decode 경로로 흘려보낸다. 증류 손실은 response 타깃만 감독하지만, 그 그래디언트는 decode가 소비하는 prompt의 key와 value를 거쳐 prefill 가중치까지 도달한다. 결과적으로 두 경로가 하나의 forward-backward pass 안에서 같은 응답 목표를 향해 함께 학습된다. 마스터 가중치를 공유할 수도, 단계별로 분리할 수도 있고, 디코더를 얼려둔 채 prefill만 적응시키는 것도 가능하다. 핵심 실험은 Qwen 3와 Gemma 3에서 Tülu 3 SFT 코퍼스 100M 토큰을 쓰고, 얼린 비양자화 교사 모델에 대한 KL(p_teacher || p_student)를 최소화하는 방식으로 이뤄졌다.

무엇과 다른가

논문은 disaggregated quantization을 세 가지 상호 보완적 스킴으로 정리한다. 첫째는 포맷 분리다. NVFP4는 가중치와 활성값을 모두 양자화해 prefill을 빠르게 하지만, 활성값을 남겨두는 NVFP4A16은 정확도가 높은 대신 prefill 연산 이점을 포기한다. 여기서 활성값 양자화를 decode에서만 끄면 저장량과 prefill 비용은 NVFP4 그대로 유지하면서 메모리 병목인 decode를 2~3% 더 빠르게 하고 decode 중심 태스크 정확도를 올린다. 같은 원리를 2~3비트 LUT 인코딩(LUT2A16, LUT3A16)에도 적용한다. 낮은 비트 가중치 위에서 네이티브 저정밀 연산을 하려면 가중치와 활성값을 NVFP4로 즉석 재양자화(autocast)해야 하는데, 이 autocast를 prefill에만 허용하고 decode는 weight-only LUT로 두는 것이 포맷 분리다. 둘째는 완전 분리다. prefill에 별도의 compute-native NVFP4 가중치를 학습시킨다. 두 가중치 모두 같은 비양자화 모델에서 출발해 QADD로 각자의 포맷 아래 공동 최적화된다. KV 캐시는 원 모델의 layer·head 차원을 유지하므로 attention 구조를 바꿀 필요가 없다. 대가는 prefill 체크포인트를 하나 더 저장해야 한다는 점이다. 셋째는 ODP(Offloaded disaggregated prefill)다. prefill 트랜스포머 블록이 출력을 내면 그 블록의 가중치는 남은 턴 동안 쓰이지 않는다. ODP는 prefill 가중치를 SSD에서 블록 단위로 스트리밍하고, prefill 중에는 쓰이지 않는 decode 가중치 영역을 같은 크기만큼 잘라 블록 버퍼로 재사용한 뒤 생성 전에 복원한다. 버퍼 두 개로 다음 블록 로딩을 연산과 겹치며, 프롬프트가 길수록 로딩 비용이 상대적으로 줄어든다. 여기에 더해 이미 공개된 weight-only 양자화 체크포인트를 그대로 두고 prefill만 학습시키는 prefiller 방식이 있다. 디코더는 역양자화된 가중치로 고정된 채 그래디언트만 prefill 경로로 흘려보내므로, 디코더의 양자화 파이프라인은 블랙박스로 남겨둘 수 있다.

어떻게 쓰나

먼저 논문은 단계별 양자화 민감도가 다르다는 것을 수치로 보인다. 각 단계를 NVFP4로 따로 양자화해 보면, decode 중심 벤치마크에서는 decode만 양자화했을 때의 정확도 손실이 prefill만 양자화했을 때의 2~4배이고 Gemma3-1B에서는 최대 7배다. 반대로 prefill 중심 태스크에서는 8개 모델 중 7개에서 prefill만 양자화한 쪽의 손실이 decode만 양자화한 쪽의 1.1~4.1배다. 본 실험은 Qwen 3(0.6B, 1.7B, 4B, 8B)와 Gemma 3(1B, 4B, 12B)에서, decode 중심 평가로 GSM8K·MATH-500·MMLU-Pro(Qwen 3는 reasoning on/off 두 모드), prefill 중심 평가로 RULER(4K·8K·16K·32K)를 사용한다. 정확도는 vLLM과 NIXL로 실제 disaggregated serving을 돌려 측정했고, 지연은 DGX Spark에서 batch-one 기준으로 쟀다. 포맷 분리는 decode 중심 태스크에서 비분리 스킴 대비 Qwen 3와 Gemma 3 각각 NVFP4 +1.9/+3.1점, LUT3 +4.5/+3.7점, LUT2 +2.5/+1.8점을 얻었고, prefill 중심 태스크에서는 효과가 1.3점 미만이었다. 속도는 prefill에서 BF16 대비 Qwen3-8B 1.49배, Gemma3-12B 1.67배다. 완전 분리는 decode 중심에서 LUT3 +6.3/+5.2점, LUT2 +10.7/+7.4점, prefill 중심에서 LUT3 +4.1/+4.3점, LUT2 +5.3/+10.5점을 얻는다. 2비트에서는 LUT2A16 weight-only 서빙보다 4.5~12.5점 높으면서 prefill 연산은 더 빠르다. ODP는 16K 이상 문맥에서 상주 NVFP4 prefill 지연을 Qwen 3에서 5% 미만, Gemma 3에서 8% 미만만 늘리고, 16K에서 스트리밍 스택이 BF16보다 Qwen3-8B 1.47배, Gemma3-12B 1.58배 빠르다.

전제와 한계

논문은 27B 규모로도 검증한다. Qwen3.8-27B에 대해 Unsloth가 공개한 GGUF 체크포인트 8종(IQ2_XXS 같은 벡터 양자화 포함)을 그대로 두고 NVFP4 prefiller만 학습시켰다. BF16 모델의 코드·수학 추론 트레이스로 95M 토큰을 QADD 학습한 뒤 MMLU-Pro와 MMMU-Pro를 reasoning을 켠 상태로 평가했다. 1비트 디코더에서 NVFP4 prefiller는 MMLU-Pro +32.5점, MMMU-Pro +35.3점으로 weight-only 정확도를 두 배 이상으로 끌어올렸다. 2비트에서는 +7.4점, +6.3점으로 이득이 줄고, 3비트에서는 오히려 약간의 성능 저하가 나타난다. 학습 코퍼스에 멀티모달 예시가 하나도 없는데도 시각 추론으로 이득이 전이됐다는 점을 저자들은 짚는다. 같은 정밀도의 RTN prefill보다 학습된 prefiller가 크게 앞선다는 결과도 부록에 있다. 속도는 llama.cpp 확장에서 8K 문맥 기준 TTFT가 12.27초에서 6.90초로 1.78배 빨라졌고, 4K~32K 구간에서 1.38~1.78배다. 마지막으로 재학습 없는 PTQ만으로 포맷 분리를 최대 2.8T 파라미터의 텍스트·멀티모달 모델 8종(Qwen 3.8, Gemma 4, Muse Glimmer, Nemotron 3, Kimi-K3)에 적용해, 4비트에서 13개 모델-벤치마크 조합 중 11개에서 점추정치가 개선되고 6개가 α=0.05에서 유의하며 유의한 악화는 없었다고 보고한다.

실무 관점에서 이 논문이 주는 판단 기준은 명확하다. 단일 사용자 로컬 서빙처럼 디바이스 메모리가 빠듯하고 원래 모델의 응답성을 유지해야 하는 상황이라면, 포맷 분리는 저장량과 prefill 비용을 그대로 두고 decode 정확도만 올리는 플러그인 교체에 가깝다. 이미 2~3비트 weight-only 체크포인트를 쓰고 있고 prefill이 느리다면 완전 분리나 prefiller가 후보가 된다. 특히 prefiller는 공개된 GGUF 체크포인트를 다시 학습하지 않고 그대로 쓸 수 있다는 점이 실용적이다. 다만 ODP는 dense 모델 로컬 배포를 전제로 한 도구이며 MoE에는 그대로 옮겨지지 않는다. 활성 파라미터 비율이 커질수록 로딩 비용이 연산보다 훨씬 커지기 때문이다. 또한 짧은 프롬프트에서는 SSD 로딩에 발목이 잡혀 TTFT가 오히려 나빠질 수 있고, 3비트 이상 디코더에서는 prefiller 이득이 사라지거나 역전되므로 자신의 디코더 비트폭과 프롬프트 길이 분포를 먼저 확인해야 한다.

저자들이 명시한 한계도 분명하다. 평가는 decode 중심 추론과 단일 턴 prefill 중심 태스크에 한정되며, 배치가 큰 서빙 성능이나 멀티턴·에이전트 동작은 다루지 않는다. 멀티턴에서는 캐시된 어시스턴트 토큰이 decode가 만든 KV 엔트리를 유지하는데, 같은 토큰 이력을 prefill로 다시 캐싱하면 다른 표현이 나올 수 있고 이 캐시 정책 의존성에 대한 견고성은 검증되지 않았다. ODP는 짧은 시퀀스의 TTFT가 중요할 때와 MoE 모델에는 덜 유용하다. 지연 측정은 DGX Spark의 batch-one과 llama.cpp 환경에 국한된다. 저자들은 재현용 코드베이스와 ODP를 추가한 llama.cpp 포크, 학습된 Qwen3.8-27B prefiller를 GitHub와 Hugging Face Hub에 공개했다고 밝힌다.