루프 언어모델의 깊이 적응 추론을 루프 스텝 사이 배칭으로 실현한 첫 종단간 구현
Depth-adaptive Inference of Looped Language Models via Continuous Depth Batching
무엇인가
루프 언어모델(looped LM)은 공유 레이어 블록, 즉 순환 코어를 여러 번 반복 적용해 토큰 표현을 정제한다. 이 구조의 최대 장점은 깊이 적응 추론이다. 쉬운 토큰은 코어를 몇 번만 돌고 일찍 빠져나가며, 어려운 토큰은 더 깊이 돈다. 문제는 토큰마다 루프 횟수가 다르면 하나의 균일한 순전파로 묶을 수 없다는 점이다. vLLM 같은 표준 연속 배칭(continuous batching)은 토큰 경계에서만 스케줄링하고 순전파 전체를 하나의 스케줄링 단위로 취급하므로, 코어를 빠져나간 토큰은 코다를 돌려야 하는데 남은 토큰은 다음 루프를 돌아야 하는 상황을 처리하지 못한다. 이 논문은 깊이 적응 추론의 실용적 가치가 결국 배칭 효율에 달려 있다고 보고, 루프 스텝 사이에 새 배치를 구성하는 연속 깊이 배칭(CDB)의 첫 종단간 구현을 제시한다.
어떻게 동작하나
CDB는 요청의 경로를 네 개의 큐로 나눈다. 프리필을 기다리는 비활성 요청, 프렐류드를 기다리는 디코드 토큰, 다음 루프 스텝을 기다리는 순환 항목, 코다를 기다리는 탈출 토큰이다. 이 네 단계는 각각 다른 배치 크기와 빈도로 따로 배칭·실행될 수 있고, 스케줄러는 이들을 임의 순서로 인터리빙한다. 실행 모드는 두 가지다. 노리필 모드는 프렐류드, 모든 토큰이 탈출할 때까지의 순환 스텝, 코다 순으로 고정 실행하고 탈출 토큰을 순환 배치에서 제거해 이후 스텝의 배치 크기를 줄인다. 리필 모드는 단계 순서를 동적으로 정해 비워진 슬롯을 새 토큰으로 채워 배치 크기를 유지한다. 우선순위 규칙은 프리필이나 코다 직후에는 프렐류드를 돌리고, 나머지 중에서는 다음 토큰을 예측해 리필 재료를 만드는 코다를 최우선으로, 그다음 프리필, 마지막으로 순환 코어를 둔다. 코다 레이어가 비싼 모델에서 탈출 한 건마다 코다를 돌리면 이득이 비용을 넘기므로, 최소 코다 배치 크기 K를 두어 탈출이 K개 이상 쌓였을 때만 코다에 우선권을 준다. 예를 들어 64개 배치에서 토큰 하나가 탈출했을 때 K=1이면 즉시 코다·프렐류드를 돌려 다음 순환 배치를 64로 채우고, K=16이면 63개로 먼저 계속 돌다가 30개가 쌓인 뒤 함께 처리한다.
무엇과 다른가
깊이가 다른 토큰이 섞이면 KV 상태에도 구멍이 생긴다. 깊이 인덱스 KV 캐시는 공유 레이어마다 루프 스텝별 슬롯을 두어 토큰당 r_max개 슬롯을 쓰는데, 일찍 탈출한 토큰은 뒤쪽 슬롯을 비워 두고 더 깊이 도는 후속 토큰은 참조할 상태가 없다. 논문은 두 가지 해법을 검토한다. last-exited 캐시는 탈출 시점에 마지막 KV 상태를 건너뛴 슬롯에 복사해 정확도를 유지하고 재계산을 피하지만 슬롯 수와 복사 오버헤드가 남는다. shared 캐시는 순환 레이어마다 토큰당 슬롯 하나만 두어 순환 코어 KV 메모리를 r_max분의 1로 줄이고, 각 루프 스텝이 슬롯을 덮어써 후속 토큰이 항상 가장 최근 상태를 참조하게 한다. 대신 정확도 유지를 위한 별도 학습이 필요할 수 있다. 스케줄링 지연도 해결해야 한다. 탈출 결정은 루프 스텝 r이 끝난 뒤에야 나오는데 다음 배치 구성은 그 전에 끝나야 하므로 GPU가 논다. 논문은 룩어헤드 게이트를 제안해 스텝 r의 탈출 신호가 스텝 r+1이 아니라 r+2의 참여 여부를 결정하게 함으로써 CPU에 한 스텝의 여유를 준다. EOS 토큰을 한 스텝 늦게 잡는 관행과 달리 계산을 낭비하지 않으며, 정확도 저하 없이 탈출 결정을 한 스텝 앞당길 수 있다고 보고한다. 이와 함께 정적 사전할당 텐서, CUDA 그래프, 2의 거듭제곱 패딩으로 GPU 유휴 시간을 거의 0에 가깝게 줄인다.
어떻게 쓰나
논문은 이론적 상한도 함께 유도한다. 토큰 하나의 순환 코어 1회 FLOP을 F_r, 루프와 무관한 프렐류드·코다 FLOP을 F_0라 하면, 고정 깊이 디코딩은 F_0 + r_max·F_r, 적응 깊이는 F_0 + d̄·F_r이고, 디코드 속도 향상 상한은 (F_0 + r_max F_r)/(F_0 + d̄ F_r) ≤ r_max/d̄다. 프렐류드·코다에 트랜스포머 레이어가 많을수록 F_0가 커져 이득이 줄고, r_max가 크면 스킵한 루프 하나가 절약하는 비중이 작아진다. 프리필은 깊이 적응의 영향을 받지 않으므로, 프리필이 차지하는 시간 비율 α를 사후 측정해 종단간 상한 (α + (1-α)/S_dec)^-1을 제시한다. 순환 스텝 지연은 t(N_c, N_b) = δ + γN_b + αN_c로 모델링된다. 작은 배치에서는 가중치 재로딩 비용 δ가 지배해 지연이 거의 일정한 메모리 바운드 구간이고, 큰 배치에서는 연산과 KV 스트리밍이 지배해 배치 크기에 거의 선형으로 늘어난다. 두 구간의 경계가 포화 배치 크기 B*이며, 컨텍스트가 길수록 B*가 작아진다. 따라서 리필은 메모리 바운드 구간에서 가장 효과적이다.
전제와 한계
실험은 Ouro 1.4B(0-24-0, r_max=4, 완전 루프 구조로 경계 단계에 토큰 임베딩과 LM 헤드만 존재)와 Huginn 3.5B(2-4-2, r_max=16, 트랜스포머 레이어의 50%가 코어 밖)에서 수행했다. 워크로드는 Alpaca(짧은 지시), ShareGPT(채팅), ArXiv-Summarization(긴 프리필 중심) 세 가지이며, 단일 80GB H100에서 모든 토큰에 r_max 루프를 도는 표준 토큰 수준 CB와 비교했다. 오프라인 처리량에서 Ouro는 리필 모드가 CB 대비 1.32~1.53배를 내며 추정 상한 S_e2e의 96~99%를 달성했다. Huginn은 노리필이 가장 좋아 1.30~1.37배, S_e2e의 83~96%를 기록했다. 완전 깊이에서는 탈출이 없어 CDB가 CB를 앞설 수 없는데, 그럼에도 단계 분할 순전파와 추가 스케줄러 상호작용 비용을 포함해 CB 처리량의 95~99%를 유지했다. 리필의 이득은 배치가 B* 근처나 그 이하인 메모리 바운드 구간에서 가장 컸고, 배치가 커져 연산 바운드가 되면 줄어들었다. 프리필은 GPU 시간의 3.7~36%를 차지했고, 긴 컨텍스트 ArXiv에서 비중이 가장 커 S_e2e가 S_dec보다 크게 낮아졌다. 온라인 서빙에서는 포아송 도착을 10분간 재생했을 때 CDB가 정규화 지연이 급증하기 전까지 더 높은 요청률을 견뎠다. Ouro에서는 리필이 가장 큰 용량 이득을 냈고 노리필은 그 절반 정도였으며, Huginn에서는 리필과 노리필이 같은 서빙 용량에 도달했다. 최소 코다 배치 크기 K를 바꾼 실험에서 Huginn은 K=W/4가, 경계 단계가 싼 Ouro는 K=1이 최적이었다. Huginn의 경계 단계에서 트랜스포머 레이어를 제거해 0-4-0으로 만들면 K=1이 거의 최적이 되고 상한과의 격차도 줄었다.
실무 관점에서 이 논문이 주는 메시지는 깊이 적응 추론의 이득이 모델 구조와 서빙 시스템에 함께 달려 있다는 것이다. 완전 루프 구조처럼 경계 단계가 가벼운 모델은 리필 모드가 잘 맞고, 프렐류드·코다에 트랜스포머 레이어가 많은 모델은 코다 배칭 효율과 리필 지연 사이의 균형을 잡는 K 값 조정이 필요하다. 또한 탈출이 잦아도 프리필 비중이 크면 종단간 이득이 줄어들므로, 실제 워크로드의 프리필 비율과 컨텍스트 길이를 먼저 확인해야 한다. KV 캐시 레이아웃(깊이 인덱스 대 공유)은 메모리 사용량과 정확도, 구현 복잡도를 동시에 좌우하므로 모델 학습 단계에서부터 결정해 두는 편이 좋다. vLLM·SGLang 같은 기존 배칭 시스템 위에서 그대로 돌릴 수 없다는 점, 즉 루프 스텝 사이 스케줄링과 비동기 배치 준비가 별도로 필요하다는 점도 실무적으로 중요하다.
저자들이 밝힌 한계는 분명하다. 두 개의 루프 아키텍처를 단일 H100에서만 평가했고, mixture-of-experts 같은 향후 구성 요소가 순환·경계 단계의 비용과 스케줄링 요구를 바꿀 수 있다고 인정한다. 멀티 GPU 실행, 프리엠션, 프리픽스 캐싱, 추측 디코딩 같은 프로덕션 서빙 기능은 실험에 포함되지 않아, 이런 기법이 프리필과 디코드의 균형을 바꾸면 CDB의 상대적 이득도 달라질 수 있다. 룩어헤드 게이트는 최소 탈출 깊이를 r_min=2로 제약하는데, Ouro와 Huginn에서는 실제 탈출 패턴과 어긋나지 않아 손해가 없었지만 한 번의 루프만으로 유효한 출력을 내는 모델에서는 이득이 줄 수 있다. 또한 더 공격적인 탈출 임계값과 누락된 KV 상태는 정확도를 떨어뜨릴 수 있어, 원하는 정확도에서 충분한 조기 탈출이 나오도록 모델이 학습되어 있어야 CDB가 쓸 계산 절약을 확보할 수 있다고 강조한다.