SkipKV: Selective Skipping of KV Generation and Storage for Efficient Inference with Large Reasoning Models
Requests for name changes in the electronic proceedings will be accepted with no questions asked. However name changes may cause bibliographic tracking issues. Authors are asked to consider this carefully and discuss it with their co-authors prior to reque
proceedings.mlsys.org
처음보는 컨퍼런스긴 한데 비슷한 방향성을 가진 것 같아서 ...
불필요하게 반복되는 문장은 KV cache에서 제거하고, 불필요한 추론 문장이 생성되기 전에는 activation steering으로 억제하여 정확도를 유지하며 KV 메모리와 생성 길이를 동시에 줄이는 추론 최적화 방법을 고안해냄
정확도는 유지하며 토큰 길이와 KV 메모리 감소를 보인다.
기존 방법들은 일부 숫자나 단어만 제거하여 추론 흐름을 깨드려 추론 길이도 길어지지만 이 방법론은 추론 길이도 짧게 유지한다.
또한 여기선 Multi-batch를 신경쓰며 배치가 커질 때 padding token의 영향을 받게 되는데 계산을 불안정하게 만들어 품질이 떨어지게 됨
거의 같은 내용을 표현하는 두 문자이 있다면 이전 문장의 KV를 우선 제거하고, 더 최근 문장은 유지하는 식으로 진행함
여기선 따로 Sentence BERT를 쓰면 비용이 크니까 모델의 마지막 layer hidden state를 직접 사용하여 문장마다 mean pooling한 뒤 Cos 유사도를 통해 구합니다. 거기서 유사도가 높으면 중복으로 두고, 이전 문장을 지우게 됩니다.
그리고 필요 없는 wait, alternatively 등등 ... 의 hidden state와 reasoning token의 hidden state를 통해 steering도 진행합니다.
| 연구 문제 | Large Reasoning Model은 긴 Chain-of-Thought를 생성하면서 KV cache가 토큰 수에 비례해 증가한다. 이로 인해 GPU 메모리 사용량이 커지고, 처리 가능한 batch size와 decoding throughput이 제한된다. |
| 기존 방법의 한계 | H2O, R-KV와 같은 기존 KV eviction 방법은 개별 토큰의 attention score나 중복도를 기준으로 제거한다. 이 과정에서 수식, 숫자, 최종 답변의 일부가 파편적으로 삭제되어 추론 흐름이 깨지고, 모델이 이전 내용을 반복적으로 재검증하면서 오히려 더 긴 출력을 생성한다. 또한 multi-batch 환경에서는 padding token이 고정 KV budget을 차지해 정확도가 크게 하락한다. |
| 핵심 관찰 | ① 기존 KV eviction은 batch size가 증가할수록 정확도가 감소한다. ② KV cache를 줄였음에도 FullKV보다 생성 토큰 수가 증가할 수 있다. ③ 오답 reasoning에는 정답 reasoning보다 의미적으로 유사한 문장과 Wait, Alternatively, again 같은 non-execution thought가 더 많이 나타난다. |
| 핵심 아이디어 | 토큰 단위가 아니라 문장·reasoning segment 단위의 의미 중복성을 우선적으로 고려한다. 이미 생성된 중복 문장은 KV cache에서 제거하고, 불필요한 문장이 생성되려는 경우 activation steering으로 생성을 억제한다. |
| 방법 1: Sentence-level KV Storage Skipping | 생성 문장을 newline 및 punctuation 기준으로 분할하고, 각 문장의 마지막 layer hidden state 평균을 문장 representation으로 사용한다. 문장 간 cosine similarity가 임계값 τ보다 높으면 이전 문장을 중복으로 판단해 해당 문장의 KV를 우선 제거한다. 별도의 sentence encoder를 실행하지 않아 추가 연산을 줄인다. |
| 최종 Eviction Score | SkipKV는 attention 기반 token importance (I), token redundancy (R), sentence similarity λ_{i,j}를 결합한다. 중 복 문장에 속한 토큰에는 sentence similarity를 추가 감점하여 일반 토큰보다 먼저 제거되도록 한다. 즉, 실제 KV 크기는 token budget으로 맞추되 eviction 우선순위는 문장 의미 중복성이 지배한다. |
| Cache Range Monitoring | 반복적인 eviction 이후에는 원래 generation sequence의 문장 위치와 실제 KV cache 내부 위치가 달라진다. SkipKV는 generation space와 cache space 사이의 문장 범위를 지속적으로 매핑하여 sentence score가 정확한 KV 영역에 적용되도록 한다. |
| 방법 2: Adaptive Steering | MATH 학습 데이터 1,000개에서 execution thought와 non-execution thought의 평균 hidden-state 차이를 steering vector로 구성한다. 생성 중 non-execution thought가 누적될수록 steering strength를 α_t=α_0+γ N_o로 증가시켜 모델을 실제 계산과 문제 해결 중심의 reasoning 방향으로 유도한다. |
| 방법 3: Batch Grouping | 입력을 prefill length 순으로 정렬한 뒤 길이가 비슷한 sample끼리 batch를 구성한다. 이를 통해 padding token을 줄이고, 설정된 KV budget 중 실제 reasoning token이 사용할 수 있는 valid KV budget을 증가시킨다. |
| 학습 여부 | 모델 파라미터를 업데이트하지 않는 training-free inference method이다. 다만 adaptive steering에 사용할 steering vector는 사전에 데이터로부터 계산하며, 모델별 steering layer와 strength 설정이 필요하다. |
| 평가 모델 | DeepSeek-R1-Distill-Qwen-7B, DeepSeek-R1-Distill-Qwen-14B, DeepSeek-R1-Distill-Llama-8B |
| 평가 데이터셋 | 수학 reasoning: AIME-24, MATH-500, GSM8K / 코드 reasoning: LiveCodeBench |
| 비교 방법 | FullKV, H2O, R-KV, SEAL |
| 정확도 결과 | 동일하거나 유사한 KV compression budget에서 기존 방법보다 최대 26.7% 높은 정확도를 기록했다. AIME-24의 R1-Qwen-14B에서는 FullKV와 동일한 정확도를 유지하면서 약 6.7배 적은 KV memory를 사용했다. 일부 조건에서는 FullKV보다 적은 메모리로 더 높은 정확도를 달성했다. |
| 생성 길이 결과 | 기존 token-level eviction은 FullKV보다 더 긴 출력을 생성하는 경우가 많았지만, SkipKV는 FullKV 대비 최대 약 28%, R-KV 대비 모델에 따라 최대 32∼48% 적은 토큰을 생성했다. 논문 전체 기준으로 SoTA 대비 최대 약 1.6배 짧은 생성 길이를 달성했다. |
| Throughput 결과 | GSM8K의 single A100-40GB 환경에서 FullKV 대비 최대 9.6배 높은 throughput을 달성했다. 동일 batch size에서는 더 짧은 생성 길이 덕분에 R-KV보다 최대 1.7배 높은 throughput을 보였다. FullKV와 SEAL이 OOM이 발생한 큰 batch에서도 고정 KV budget을 통해 추론이 가능했다. |
| Ablation 결과 | Sentence Scoring은 추론 coherence와 정확도를 개선했고, Adaptive Steering은 생성 길이를 크게 줄였으며, Batch Grouping은 multi-batch 정확도를 추가로 향상했다. 세 요소를 모두 적용했을 때 R-KV 대비 정확도 최대 +20%p, 생성 길이 최대 −30%를 기록했다. |
| 주요 기여 | ① KV eviction의 단위를 token에서 의미적 문장 구조로 확장했다. ② 이미 생성된 KV의 제거와 불필요한 KV의 생성 억제를 동시에 수행한다. ③ padding으로 인한 multi-batch 성능 저하를 분석하고 batch grouping으로 해결한다. ④ vLLM의 paged KV cache 및 continuous batching에서도 효과를 확인했다. |
| 한계 | Batch grouping은 prompt 길이 분포가 지나치게 불규칙한 환경에서 효과가 제한될 수 있다. 문장 분할 delimiter와 non-execution keyword가 수작업으로 정의되어 있어 언어·도메인별 조정이 필요하다. 평가도 주로 DeepSeek-R1 Distill 계열과 수학·코드 reasoning에 집중되어 있어 agent, 장문 추론, 다국어 환경에 대한 추가 검증이 필요하다. |