[논문리뷰] RidgeFill: Training-Free Late-Layer Cache Fill for Faster Prefill

링크: 논문 PDF로 바로 열기

저자: Dylan Johnston, Robert P. Smith

1. Key Terms & Definitions (핵심 용어 및 정의)

• Split / Tail / Prefix: Deployment마다 split layer S와 tail 길이 T를 정합니다. Layer 0~S−1은 모든 row(prompt token 위치)에서 exact하게 실행됩니다. 마지막 T개 token(tail)과 generated token은 모든 layer를 exact하게 실행합니다. 나머지 N−T개 token이 prefix입니다.
• Ridge Map (Fill Map): 각 late layer ℓ마다 따로 두는 linear map x̂_ℓ = A_ℓ·x_S + b_ℓ입니다. Split layer의 hidden state로 해당 layer의 input을 예측합니다. Calibration text를 한 번 통과시킨 뒤 ridge regression으로 closed form에서 fitting하며, training step은 없습니다.
• Filled Row vs. Exact Row: Filled row는 late layer에서 예측된 input으로 K/V projection만 계산하고 MLP를 포함한 나머지 연산을 건너뜁니다. Gated DeltaNet layer에서는 recurrent state update만 계산합니다. Exact row는 selector가 고른 prefix row로, 모든 layer를 정상적으로 실행합니다.
• Identity Fill: Split state x_S를 late layer의 input으로 그대로 쓰는 baseline입니다. POP는 이 방식에서 split을 L−⌊L/3⌋에, tail을 1 token으로 둔 operating point입니다.
• Novel / Key Token ΔNLL: Average loss가 fill error를 가리기 때문에 도입한 지표입니다. Novel token은 앞선 prompt에 같은 8-token run이 없어 그대로 copy할 수 없는 token입니다. Key token은 context를 마지막 4,096 token으로 자르면 exact loss가 2 nats 넘게 오르는 token입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

이 논문은 long-context prompt의 prefill latency를 줄이는 것을 목표로 합니다. 이를 위해 trained model 수정 없이 late layer 연산의 대부분을 건너뛰는 training-free 방법을 제안합니다. 예를 들어 Qwen3.8-27B에서 120K token을 prefill하면 A100에서 128초가 걸리고, 사용자는 그동안 아무 출력도 받지 못합니다.

기존 접근에는 각각 한계가 있습니다.

• YOCO와 DeepSeek-V4.1-Flash는 middle layer state의 projection으로 upper layer cache를 채우는 구조를 쓰지만, 모델 전체를 그 구조로 pretraining해야 합니다.
• SwiftKV(약 680M token distillation)와 KVA projector는 retrofit 방식이지만 별도 training이 필요합니다.
• POP 같은 identity fill은 training이 없는 대신 split이 얕아지면 retrieval이 크게 무너집니다.

저자들은 approximation theory 기반 correction(Jensen, Gaussian closure, Nyquist interpolation)도 시도했지만 효과가 없었습니다. Fill error가 row 간에 white noise에 가까워 보간할 구조가 없었기 때문입니다. 그래서 핵심 과제는 error를 평균적으로 줄이는 것이 아니라 어떤 row를 exact하게 남길지 고르는 것이라고 봅니다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

RidgeFill의 설계 논리는 speculative decoding을 반대 방향으로 적용한 것입니다. Prefill은 compute-bound이므로 남는 memory bandwidth로 per-layer map을 읽고, 그 대가로 late layer의 arithmetic을 건너뜁니다.

방법 구성:

• Late layer input 예측: Ridge map은 late layer의 K/V가 아니라 input을 예측합니다. 따라서 각 layer의 원래 projection과 recurrent state update가 그대로 실행됩니다.
• Value Rescale: Tail의 처음 c=512 row에 대해 exact value norm과 predicted value norm의 비율을 KV head별로 구하고, filled value에 곱합니다. 이 scale은 request마다 새로 측정합니다.
• Terminal-State Correction: Gated DeltaNet hybrid model에서 filled row들의 error는 shared state에 누적됩니다. 이 correction은 그 상태를 exact 쪽으로 되돌립니다.
• Exact row selector: Selector는 세 가지입니다. Reference study는 마지막 64 token에서 attention을 가장 많이 받는 10% row를 고릅니다. 이는 attention weight a_ij에 비례하는 error bound에 근거합니다. Flash-Next는 protected-token class를, 27B는 late-layer attention을 직접 재는 probe를 씁니다.
• Speedup 상한: Compute ceiling은 r = S/L + (1−S/L)·T/N이고 speedup은 최대 1/r입니다.

주요 결과:

• Qwen3.8-27B (llama.cpp, A100 40GB)
  • TTFT가 16K120K token에서 **1.251.31×** 빨라졌습니다(maps resident 기준). 예를 들어 120K에서는 128.10초 → 100.01초입니다.
  • Map을 host memory에서 streaming하면 extra VRAM이 약 1.3 GiB 줄고, 16K에서 0.5초가 추가됩니다.
  • RULER는 16K에서 −1.19점, 32K에서 −1.42점 변했습니다.
  • LongBench v2는 38.57 → 37.18(−1.39점)입니다. 95% CI는 0을 포함하지만, 수 점 수준의 손실 가능성을 배제하지는 못합니다.
• Qwen3.8-Flash-Next (vLLM)
  • Exact 대비 TTFT가 9K32K에서 **1.501.60×** 빨라졌습니다. 이론 상한은 각각 1.64×, 1.88×입니다.
  • RULER needle task는 99.8(exact 100.0)입니다.
  • Novel-token perplexity는 +2.65% [−0.02%, +4.93%] 높고, quality mode에서는 +1.60%입니다.
• 8~9B reference study (Llama-3.1-8B, Qwen3.5-9B, Qwen3-8B)
  • 같은 split과 tail 조건에서 identity fill 대비 filled-region RULER가 두 dense model에서 +49.1점, +49.6점 높았습니다.
  • Released SwiftKV checkpoint와 비교하면 novel-token loss 차이가 0.03 nats 이내입니다(SwiftKV − RidgeFill = +0.010 nats).
  • 다만 같은 4,096-token tail을 주면 SwiftKV의 trained projection이 더 좋습니다(novel token −0.024, key token −0.224 nats).
• Error 위치 (splice test): Llama-3.1-8B에서는 error가 key에 있고, 그중 절반이 key의 scale 문제입니다. Key rescale을 적용하면 novel-token 추가 loss가 +1.006 → +0.502 nats로 줄어듭니다. Qwen3.5-9B에서는 value와 recurrent state가 error를 담고 있습니다.

4. Conclusion & Impact (결론 및 시사점)

RidgeFill은 training 없이 closed-form ridge map만으로 기존 모델에 YOCO 스타일의 late-layer cache fill을 붙입니다. 이를 통해 실제 serving engine(llama.cpp, vLLM)에서 **1.25~1.60×**의 TTFT 개선을 얻으면서 benchmark 손실을 작게 유지했습니다.

다만 한계도 분명합니다.

• QK-norm이 없는 모델(Llama-3.1-8B): 얕은 split에서 identity fill보다 loss가 나쁩니다. 64K 이상 문서에서는 fill을 하지 않는 경우보다도 나쁘며, 128K source code 문서에서는 ΔNLL이 +8.19 nats까지 올라갑니다. 저자들은 이런 모델에 half depth 이상의 split이나 attention-selected row를 권장합니다.
• Key token에서의 잔여 손실: Needle score가 exact와 같아도 key-token loss는 0.4 nats 이상 남을 수 있습니다.

방법론 측면의 기여도 큽니다. 이 논문은 평균 loss만으로는 approximation 방법의 error가 드러나지 않음을 보였고, 그 대안으로 filled-region retrieval, novel/key-token loss, splice test라는 평가 체계를 제시했습니다. 산업적으로는 재학습 없이 기존 checkpoint에 바로 붙일 수 있는 long-context prefill 가속 기법이라는 의미가 있습니다. 동시에 needle 지표만으로 품질을 판단하면 위험하다는 점도 보여줍니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.