최장 안정 접두사 스케줄러로 DLM 추론 가속

초록

본 논문은 확산 언어 모델(DLM)의 추론 속도를 저해하는 ‘산재 수용’ 방식을 극복하고, 연속적인 안정 접두사를 한 번에 커밋하는 Longest Stable Prefix(LSP) 스케줄러를 제안한다. 단일 전방 패스로 토큰 안정성을 평가하고, 적응형 임계값과 구조적 경계 스냅핑을 통해 왼쪽 정렬된 블록을 원자적으로 고정함으로써 KV 캐시의 연속성을 유지하고 활성 접미사의 길이를 기하급수적으로 감소시켜 최대 3.4배의 가속을 달성한다.

상세 분석

LSP 스케줄러는 기존의 ‘산재 수용’(scattered acceptance) 전략이 초래하는 두 가지 근본적인 비효율성을 직접 해결한다. 첫째, 토큰을 개별적으로 고정하면 고정된 토큰과 아직 미정인 토큰 사이에 다수의 경계가 생겨, 모델이 반복적으로 국소적인 수정을 수행해야 한다. 이는 전체 수렴 과정을 지연시키는 알고리즘적 병목이다. 둘째, 이러한 파편화는 KV 캐시를 작은 조각으로 나누어 메모리 지역성을 파괴하고, 트랜스포머 어텐션 연산이 긴 활성 접미사 전체에 대해 매번 수행되게 만든다. LSP는 ‘최장 안정 접두사’를 한 번에 선택해 원자적으로 커밋함으로써, 고정된 프리픽스가 연속적인 블록으로 성장하도록 만든다. 이때 토큰 안정성은 로그 마진(상위 두 로짓 차이)으로 측정되며, 단일 전방 패스만으로 모든 위치의 마진을 얻는다.

적응형 임계값 탐색은 L′(τ) = 연속적인 마진이 τ 이상인 가장 긴 접두사의 길이로 정의하고, 목표 블록 크기를 현재 활성 접미사 길이 N_k의 αβ 비율(예: 0.250.5) 안에 맞춘다. 이를 위해 마진의 프리픽스 최소값을 미리 계산하고 O(N_k) 시간에 적절한 τ_k를 찾는다. 이렇게 하면 모델이 자신감이 높을 때는 큰 블록을, 불확실할 때는 작은 블록을 선택해 기하급수적인 N_k 감소를 보장한다.

구조적 스냅핑은 후보 블록의 경계를 구두점, 줄바꿈 등 자연어 구조에 맞추어 조정한다. 이는 문법적 일관성을 유지하고, 불필요한 재수정을 방지한다. 또한, LSP는 KV 캐시를 ‘근사 재사용’한다. 고정된 프리픽스는 이후 단계에서 변하지 않는다고 가정하고, 이전 단계의 KV를 그대로 사용한다. 최근 연구에 따르면 인접 단계 간 KV 변화가 미미하므로, 이 근사화는 품질 저하 없이 큰 시스템 효율을 제공한다.

실험에서는 LLaDA‑8B와 Dream‑7B 모델에 LSP를 적용해 수학적 추론, 코드 생성, 다국어(CJK) 작업, 창작 글쓰기 등 다양한 베치에서 평균 2.1~3.4배의 지연 감소를 기록했다. 품질 측면에서는 ROUGE, BLEU, CodeBLEU 등에서 기존 ‘산재 수용’ 기반 스케줄러와 동등하거나 약간 향상된 점수를 보였다. Ablation 연구는 마진 기반 안정성 진단, 적응형 임계값, 구조적 스냅핑 각각이 속도와 품질에 기여함을 확인한다.

요약하면, LSP는 DLM 추론의 핵심 병목인 파편화된 KV 캐시와 불안정한 토큰 경계를 근본적으로 재구성함으로써, 이론적인 병렬성은 유지하면서 실제 하드웨어 효율을 크게 끌어올린다.