RL 롤아웃의 추론 최적화

앞 글에서는 생성과 학습의 진행을 분리해 준비된 경험부터 학습하는 방법을 보았습니다. 이제 생성기 자체의 비용을 줄여보겠습니다. RL은 같은 질문에서 여러 응답을 만들고, 멀티턴에서 긴 문맥을 반복해서 읽으며, 학습한 가중치로 생성 정책을 계속 갱신합니다. 이 특성이 추론 최적화의 조건이 됩니다.

이번 글에서는 다시 계산할 입력을 줄이기, 메모리 안에서 여러 요청을 함께 진행하기, 한 번의 정책 계산으로 여러 후보 토큰을 검증하기라는 세 방향을 보겠습니다. 각 방법이 줄이는 비용과 추가로 필요한 작업을 구별하고, 마지막에는 정책이 계속 바뀌는 RL에서 캐시와 draft를 어떻게 생각해야 하는지 연결하겠습니다.

같은 prefix의 계산을 재사용하기

모델이 입력 문맥을 처리하는 구간을 prefill, 이후 토큰을 하나씩 이어 생성하는 구간을 decode라고 부릅니다. Attention의 KV 캐시는 이미 처리한 토큰의 key·value를 저장해 이후 계산에 재사용합니다.

GRPO처럼 한 질문에서 여러 응답을 생성할 때는 질문과 시스템 지시문 같은 앞부분이 반복됩니다. 이 공통 prefix의 토큰 ID와 계산 조건이 같다면, 이미 계산한 KV를 공유해 입력을 반복 계산하는 비용을 줄일 수 있습니다.

그림의 두 응답 a와 b는 P를 공유하지만 이후 토큰은 다릅니다. 따라서 a를 생성하며 쌓은 KV를 b의 서로 다른 suffix에 그대로 사용할 수는 없습니다. 또한 문장 중간에 같은 단어가 나왔다는 이유만으로 같은 KV가 되는 것도 아닙니다. 그 토큰 앞에 무엇이 있었는지가 계산에 영향을 줍니다.

멀티턴에서는 P + 응답 a + 도구 관측 + 새 입력으로 문맥이 길어질 수 있습니다. 이 중 실제로 캐시에 남아 있고 완전히 일치하는 앞부분까지 재사용합니다. 새 관측과 새 입력에 대한 계산은 여전히 필요합니다. SGLang의 RadixAttention은 이러한 prefix KV 재사용을 관리하는 대표적인 방식입니다. SGLang 논문

어느 엔진으로 요청을 보낼지도 중요합니다. 그림의 엔진 A는 긴 prefix의 캐시를 가지고 있지만 대기열이 깁니다. 엔진 B는 캐시가 없어 다시 계산해야 하지만 덜 바쁩니다. 캐시만 보고 A로 계속 보내면 재계산은 줄어도 대기가 늘 수 있습니다.

Cache-aware routing은 재사용할 수 있는 입력 계산과 엔진의 부하를 함께 고려해 요청을 어느 추론 엔진에 보낼지 정하는 문제입니다. 그림의 “재계산 비용 + 대기 비용”은 판단을 이해하기 위한 개념적 표현이며, 모든 router가 그 식을 그대로 계산한다는 뜻은 아닙니다. 또한 여기의 router는 요청을 추론 엔진에 보내는 구성요소입니다. 7편의 MoE 전문가 router와는 다른 역할입니다. Miles v0.1 §2.1

빈자리와 KV 용량을 함께 보기

RL 응답은 길이가 서로 다릅니다. 짧은 응답이 끝났을 때 나머지 긴 응답을 기다리는 대신, 새 요청을 받아 함께 진행할 수 있습니다. 이런 방식은 생성 중인 요청들의 배치 구성을 계속 바꾸어 자원을 활용합니다.

하지만 동시 요청 수만 보고 새 작업을 받을 수는 없습니다. 긴 응답은 생성하는 동안 KV 캐시를 계속 쌓고, 새 요청은 입력을 처리하는 prefill도 수행해야 합니다. 실행 자리에 여유가 생기는 것과 KV 메모리에 여유가 생기는 것은 관련이 있지만 같은 조건은 아닙니다.

그림에서 A는 시점 2에 끝나고 D가 들어옵니다. D는 23에 입력을 처리한 뒤 36에 생성합니다. B는 4에, C는 6에 끝납니다. 타임라인은 요청이 진행되는 구간이며, 겹쳐 보이는 모든 구간의 GPU 커널이 동시에 실행된다는 뜻은 아닙니다.

시점 1의 진행 중 KV는 A 2, B 3, C 3을 합해 8입니다. 3.5에는 D 2, B 4, C 4로 10입니다. B가 끝난 뒤 5.5에는 D 3과 C 6으로 9입니다. 요청 수가 줄었더라도 남은 긴 응답의 KV가 커지므로 총 점유가 같은 비율로 줄지는 않습니다.

여기서는 모두 용량 12 안에 있습니다. 실제로는 최대 요청 수 제한을 만족해도 필요한 KV를 확보할 수 없다면 요청을 기다리게 하거나 다른 메모리 관리 전략을 적용해야 합니다. PagedAttention은 KV를 블록 단위로 관리해 메모리 낭비를 줄이는 접근입니다. 그렇다고 물리적인 KV 용량 제한이 사라지는 것은 아닙니다. 실제 메모리 관리에는 그림에서 제외한 보존 prefix 캐시와 기타 예약 공간도 들어갑니다. PagedAttention 논문

새 요청을 많이 받으면 prefill과 기존 요청의 decode가 같은 실행 자원을 두고 경쟁할 수 있습니다. 입력을 나누어 처리하거나 두 단계의 배분을 조절하는 이유입니다. 별도 GPU로 prefill과 decode를 분리하는 방식도 있지만, 이 그림의 필수 구성은 아닙니다. RL에서는 최종 토큰 하나의 지연뿐 아니라 그룹이 학습 준비를 마치는 시간까지 연결해서 봐야 합니다.

모델 구조도 이 선택에 영향을 줍니다. KV를 압축하는 구조는 저장량을 바꾸고, sparse attention은 attention 계산량과 접근 패턴을 바꿀 수 있습니다. 다만 sparse attention을 사용한다는 사실만으로 보관하는 KV 전체가 같은 비율로 줄어든다고 가정할 수는 없습니다. 현재 모델의 캐시 구조를 기준으로 자원 예산을 잡아야 합니다.

여러 후보를 한 번에 검증하기

자기회귀 생성에서는 다음 토큰을 확정해야 그 뒤의 토큰을 선택할 수 있습니다. Speculative decoding은 빠른 제안자와 목표 정책의 검증을 결합해 이 과정을 효율적으로 수행합니다.

Draft는 먼저 여러 후보 토큰을 제안합니다. Target은 우리가 실제로 사용하려는 생성 정책이며, 제안된 경로의 여러 위치를 함께 계산해 후보를 검증합니다. 여기서 검증은 정답이나 보상을 매기는 일이 아닙니다. 해당 후보를 target의 생성 절차에서 채택할 수 있는지 판단하는 계산입니다.

먼저 후보 제안과 검증 흐름을 읽고, 후보를 만드는 방법인 MTP는 다음 절에서 살펴보겠습니다.

그림에서는 draft가 x₁, x₂, x₃, x₄를 제안합니다. Target 검증에서 앞의 두 개는 채택되지만 x₃는 거절됩니다. 그러면 x₄도 그대로 가져갈 수 없습니다. x₄는 거절된 x₃까지 포함한 문맥에서 제안된 토큰이기 때문입니다. 거절 지점에서 정해진 보정 분포로 z를 뽑고, 확정된 x₁, x₂, z 뒤에서 다시 시작합니다.

확률적으로 샘플링하는 경우 이 검증은 서로 가장 높은 확률을 주는 토큰이 같은지만 비교하는 검사와 다릅니다. 고전적인 speculative sampling은 해당 문맥에서 target과 draft가 후보에 부여한 확률을 사용해 채택 확률을 계산하고, 거절하면 두 분포의 차이를 이용한 보정 분포에서 토큰을 뽑습니다. 여기의 확률은 temperature·top-p 등 실제 샘플링 규칙을 반영한 분포를 기준으로 합니다. 이 절차가 유효하게 수행될 때 target 분포를 보존하면서 여러 토큰을 진행할 수 있습니다. Speculative Decoding 원논문

따라서 draft가 target을 완벽하게 따라 할 필요는 없습니다. 다만 제안 품질이 나쁘면 자주 거절되어 절약하는 작업이 줄어듭니다. 분포를 올바르게 유지하는가와 실제로 빨라지는가는 다른 질문입니다. 총 연산량은 오히려 늘 수 있습니다. 여러 위치를 함께 검증하면서 순차적인 target 호출과 반복적인 가중치·KV 읽기를 줄이는 것이 가속의 근거입니다. Draft 생성, target 검증, 거절 뒤 처리, 추가 메모리 비용까지 포함해야 속도를 판단할 수 있습니다.

MTP는 후보를 만드는 한 방법이다

Draft를 만드는 방식은 여러 가지입니다. 별도의 작은 모델을 사용할 수도 있고, target의 내부 특징을 이용하는 EAGLE 계열의 모듈이나 모델에 결합된 MTP를 사용할 수도 있습니다.

MTP(Multi-Token Prediction)는 여러 미래 토큰을 예측하도록 하는 구조·학습 방식을 가리킵니다. 이를 추론의 후보 제안에 사용할 수 있습니다. 따라서 “speculative decoding을 쓰는가”와 “MTP를 쓰는가”를 완전히 독립적인 가속 기법 두 개로 세기보다, 검증 구조 안에서 어떤 제안 방식을 사용하는가로 보는 편이 관계를 이해하기 쉽습니다.

MTP 모듈이 있다는 사실만으로 모든 추론 경로가 정확한 확률적 검증을 수행한다고 결론 내릴 수는 없습니다. 어떤 후보 구조와 검증 절차를 사용하는지 구현을 확인해야 합니다. 또한 모델이 MTP를 제공한다고 RL 중 target과 draft를 함께 학습하는 기능까지 공개되어 있다는 뜻은 아닙니다.

정책 갱신 뒤 캐시와 draft를 확인하기

일반적인 고정 모델 서비스와 달리 RL은 target 정책을 계속 갱신합니다. 이 때문에 캐시와 draft의 수명도 모델 버전과 연결됩니다.

먼저 이전 가중치에서 만든 KV를 새 가중치에서도 정확히 사용할 수 있다고 가정해서는 안 됩니다. 같은 토큰이라도 가중치가 바뀌면 key·value가 달라질 수 있습니다. Weight sync 이후 이전 캐시를 무효화하거나 다시 계산하는 등 구현의 처리 규칙을 확인해야 합니다. 반면 temperature처럼 logits 뒤의 선택에만 적용되는 설정은, 그 자체로 앞선 KV 계산을 바꾸는 설정과는 구별해야 합니다.

Draft에는 두 문제가 있습니다. 하나는 새 target이 선호하는 토큰과 제안이 얼마나 잘 맞는가입니다. Target이 v10에서 v11로 바뀌면 오래된 draft의 채택률이 낮아질 수 있습니다. 다른 하나는 target의 특징이나 KV를 사용하던 내부 상태가 새 버전에서도 유효한가입니다. 단순히 draft 가중치를 유지할 수 있다는 것과 내부 상태를 그대로 재사용할 수 있다는 것은 다릅니다.

유효한 검증·보정 절차가 유지된다면 draft가 덜 잘 맞는 것은 주로 속도에 영향을 줍니다. 다만 한 검증 반복 안에서 target 버전이 바뀌거나, 검증에 쓰는 확률과 상태가 어긋나면 분포 보존의 조건 자체를 깨뜨릴 수 있습니다. 새 target 적용, 진행 중 검증, 내부 상태 정리를 일관되게 연결해야 합니다.

NeMo RL의 연구에서는 과제의 rollout으로 EAGLE-3 draft를 사전 학습하고, 주실험의 RL 진행 중에는 draft를 고정했습니다. Qwen3-8B-Base를 32개 GB200에서 동기 GRPO로 학습한 실측에서 생성 시간은 100.0초에서 56.6초로, 전체 step은 151.2초에서 107.5초로 줄었습니다. 생성의 약 1.77배 가속과 전체 step의 약 1.41배 가속이 다른 이유를 보여줍니다. 이는 해당 모델·과제·장비의 측정이며 모든 RL 구성의 기대치가 아닙니다. NeMo RL 연구 §3.1–3.2, Tables 1–2

같은 연구에서 draft를 RL 중 추가 학습하는 방법은 별도 실험입니다. 과제에 잘 맞게 초기화한 경우 rollout 가속 배율은 고정 1.77배에서 추가 학습 1.78배로 거의 같았고, 일반 대화 데이터로 초기화한 경우에는 1.51배에서 1.63배로 개선됐습니다. 따라서 target 갱신마다 draft도 반드시 학습해야 한다기보다, 현재 채택 효율과 갱신 비용을 보고 결정할 문제입니다. 공식 구현은 정책의 forward에서 얻은 특징과 분포를 재사용하되 gradient를 분리해 EAGLE-3를 학습하는 경로를 제공합니다. 연구 §3.3, NeMo RL EAGLE-3 가이드

최적화 결과를 볼 때는 제안 길이, 실제 채택 길이, 생성 시간, 전체 RL 시간을 차례로 구별할 수 있습니다. 평균 채택 길이가 늘어도 draft 비용이 함께 커지면 생성은 빨라지지 않을 수 있습니다. 생성이 빨라져도 학습이나 가중치 전달이 병목이 되면 전체 이득은 제한됩니다.

다음 글에서는 하나의 모델 호출보다 긴 에이전트 프로그램의 수명을 보겠습니다. 도구가 실행되는 동안 어떤 문맥을 메모리에 남기고, 어떤 작업을 잠시 내렸다가 다시 올릴지 결정하는 scheduling으로 연결하겠습니다.