KV-streams는 KV 캐시를 압축해 흘려보내 RL을 최대 11배 가속한다.
KV-streams for Efficient Compaction in Agentic Reinforcement Learning
무엇인가
긴 호라이즌 에이전트 태스크에 강화학습을 돌리면 동시에 많은 롤아웃을 생성해야 하고, 각 롤아웃의 KV 캐시는 길이에 비례해 선형으로 커진다. GPU 메모리 예산이 고정되어 있으면 가장 오래 실행되는 롤아웃이 메모리를 독점해 배치 크기와 처리량을 눌러버린다. 컨텍스트 압축(compaction)은 토큰 예산을 넘을 때 오래된 컨텍스트 일부를 요약으로 대체하거나 삭제해 메모리를 상수로 유지하는 대표적 해법이지만, 기존 방식은 압축이 일어날 때마다 남겨둔 컨텍스트를 새 트레이스에 다시 프리필한다. 논문은 이 반복 프리필 비용을 정량화한다. 압축 한 번에 유지되는 엔트리 수를 r, 토큰/턴 예산을 B라 하면 압축 사이에 새로 생성되는 엔트리는 B−r개이고, N개 엔트리 롤아웃의 추가 프리필량은 Δ_re-prefill = N·r/(B−r)로, 유지되는 컨텍스트가 예산에 가까워질수록 발산한다.
어떻게 동작하나
제안 방법인 KV-streams는 이 재프리필을 아예 건너뛴다. 핵심은 압축이 일어나도 KV 캐시를 플러시하지 않고 앞으로 스트리밍하는 것이다. 유지되는 오버랩과 생성된 요약은 위치 인코딩이 리셋된 새 캐시 (k'_i, v'_i)로 다시 계산되지 않고, 원래의 (k_i, v_i)를 그대로 유지한다. 따라서 Δ_re-prefill = 0이 된다. 구현은 두 갈래다. 추론 엔진 쪽에서는 삭제 대상 KV 엔트리를 엔진 내부에서 직접 제거(evict)해 메모리 절감 효과는 그대로 누리면서 하나의 연속된 KV 스트림을 유지한다. 트레이너 쪽에서는 삭제된 구간에 대한 어텐션을 차단하는 커스텀 어텐션 마스크로 같은 제거 패턴을 재현해, 단일 포워드 패스로 압축 효과를 복제한다. 이 방식은 특정 압축 규칙에 종속되지 않는 플러그인이며, 논문은 유지량 r = o + s(오버랩 o + 요약 s)의 스펙트럼을 아우르는 네 가지 전략에 적용한다. 컨텍스트를 거의 남기지 않는 Summary(o=0), 절반을 삭제하는 Markovian Thinker(o=B/2, s=0), 모델이 남길 턴을 직접 고르는 Markovian Pick, 그리고 거의 전체 윈도를 유지해 재프리필 비용이 가장 큰 Sliding Window(o≈B)다.
무엇과 다른가
실험은 텍스트 기반 게임 두 개와 소프트웨어 엔지니어링 태스크에서 이뤄졌다. TextWorld는 60,000개의 합성 게임을 생성해 256개 대표 태스크로 평가하고, 10턴마다 압축을 트리거하며 32k 토큰 예산 안에서 최대 200턴, 500 그래디언트 스텝, 배치 512, 8-GPU H100 노드(추론 4, 트레이너 4)로 학습한다. ALFWorld는 16k 토큰 예산, 200 스텝, 배치 128, A100 4장으로 학습하고 70개 본 게임과 67개 미학습 게임에서 평가한다. 소프트웨어 엔지니어링은 SWE-rebench 1,028개와 ScaleSWE 443개 태스크(베이스 모델이 항상 풀거나 전혀 못 푸는 태스크는 pass@3 기준으로 제거)로 학습해 SWE-bench Verified에서 평가한다. 결과는 TextWorld에서 재프리필 압축 대비 3.7배에서 11.3배 적은 GPU-hours로 동일한 최종 성능에 도달했고, ALFWorld에서는 트레이스가 짧아 격차가 줄지만 1.3배에서 3.0배 처리량 향상을 보였다. Summary와 Sliding Window를 재프리필로 돌리면 풀 컨텍스트 학습보다 GPU-hours를 더 많이 쓰는 반면, KV-streams 변형들은 풀 컨텍스트보다도 벽시계 시간을 크게 줄였다. 소프트웨어 엔지니어링에서는 Sliding Window + KV-streams가 SWE-bench Verified에서 52.7±2.1%로 풀 컨텍스트 51.5±1.2%, 재프리필 Markovian Thinker 53.4±0.2%와 오차 범위 내에 들면서, 약 20시간에 피크에 도달해 65시간 이상 걸린 풀 컨텍스트 대비 3배 빨랐다. Markovian Thinker도 KV-streams에서 약 32시간, 재프리필에서 약 60시간이었다. 시간 분해를 보면 재프리필의 병목은 유지 컨텍스트 r 때문에 반복 재계산이 강제되는 forward/backward 패스이고, KV-streams는 forward/backward와 생성 양쪽을 모두 줄인다. 미학습 텍스트 게임 전이 평가에서도 KV-streams는 풀 컨텍스트와 재프리필 대비 손실이 없었고, 어떤 단일 압축 전략도 모든 게임에서 최고는 아니었다.
어떻게 쓰나
논문의 두 번째 주장은 SFT가 필요 없다는 것이다. 유지된 KV가 이미 컨텍스트에서 사라진 토큰에 어텐션했기 때문에 KV 캐시가 의사 순환(pseudo-recurrent) 상태로 작동할 수 있다는 것은 선행 연구가 보였지만, 그쪽은 그 행동이 나타나게 하려는 첫 단계로 지도 미세조정(SFT)을 요구했다. 저자들은 Qwen3-4B-Instruct-2507로 32k 컨텍스트에서 10k 궤적을 모아 무작위 제거 어텐션 마스크로 SFT를 한 뒤 RL을 하는 조건과, 베이스 모델에서 바로 RL을 하는 조건을 TextWorld에서 비교했다. SFT 계산량을 포함해 세면 RL 단독이 더 계산 효율적이었고 성능 저하도 관찰되지 않았다. Markovian Thinker에서는 SFT가 오히려 학습을 크게 느리게 만들었고, Sliding Window에서는 두 방식이 같은 시간에 피크에 도달했다. 통제된 합성 실험도 이를 뒷받침한다. 모델에 과일을 하나 배정하고 k개 토큰을 디코딩하게 한 뒤, 배정이 담긴 턴을 제거하고 어떤 과일이었는지 묻는 과제에서 k ∈ {16, 32, 64, 128, 256, 512}로 KV 캐시 용량을 늘리면 회상률이 100%까지 올라가 SFT와 맞먹었고, 16과 32 토큰의 최소 예산에서만 불안정했다. 학습된 과일 회상 능력은 배우(actor)라는 새 객체 범주로도 일반화됐고, 후보 목록에 과일만 제시하면서 실제로는 텔레비전을 배정하는, 베이스 모델이 샘플링조차 하지 않을 객체에 대해서도 KV 토큰 예산이 충분하면 RL만으로 100% 회상에 도달했다.
전제와 한계
개발자 관점에서 이 논문이 주는 실무적 의미는 명확하다. 이미 요약 기반이든 슬라이딩 윈도든 어떤 압축 전략을 쓰고 있더라도, 추론 엔진에서 KV를 플러시하는 대신 제거하고 트레이너에 어텐션 마스크를 미러링하는 것만으로 학습 처리량을 크게 올릴 수 있다. 다만 도입 시 확인할 점이 있다. 논문의 vLLM 구성은 16토큰 KV 캐시 블록을 쓰고 완전한 블록만 프리픽스 캐시에 넣기 때문에, 요청이 16의 배수로 끝나지 않으면 꼬리 토큰이 커밋되지 않아 스트림이 끊기고 성능이 크게 떨어졌다. 저자들은 각 완성(completion)을 16의 배수로 패딩해 해결했지만, 이 패딩은 추론 엔진과 트레이너 양쪽에서 프리필되어야 해 전체 토큰 수를 부풀린다. 소프트웨어 엔지니어링 실험에서는 블록 크기 1로 개별 토큰을 커밋하는 SGLang 포크를 써서 패딩이 필요 없었다. 또한 트레이너는 최대 오프폴리시 지연 3의 비동기 RL을 쓰고, 가중치 업데이트 시 KV 캐시를 플러시하지 않도록 prime-RL을 수정했다. 매 가중치 업데이트마다 플러시하면 긴 트레이스 전체를 다시 프리필해야 해 메모리 절감이 무효가 되기 때문이다. 즉 KV-streams의 이득은 비동기 RL 루프와 캐시를 유지하는 가중치 업데이트 구현에 의존한다.
저자들이 밝힌 한계도 분명하다. 계산 자원 제약 때문에 다중 시드 실험은 ALFWorld에서만 수행했고, 소프트웨어 엔지니어링과 TextWorld는 단일 시드로 돌렸다. 일부 실험은 실행에 일주일 이상 걸린다고 언급한다. 소프트웨어 엔지니어링에서는 계산 한계로 네 가지 압축 전략 중 Markovian Thinker와 Sliding Window 두 가지만 풀 컨텍스트와 비교했다. 합성 실험에 대해서는, 첫 번째 과제가 후보 목록에 항상 정답이 포함되어 무작위 샘플링으로도 일부 성공이 나와 부트스트랩이 가능했고, 특정 용례에 과적합될 여지가 있으며, 제거된 컨텍스트와 그렇지 않은 컨텍스트 사이에 확률 지지집합의 겹침이 있었다는 점을 저자들이 직접 지적한다. 또한 KV-streams는 새로운 압축 규칙을 제안하는 것이 아니라 어떤 규칙이든 학습 비용을 낮추는 직교적 기법이므로, 무엇을 남길지 학습하는 기존 방법들과 결합해 쓸 수 있다는 것이 논문의 위치 설정이다.