Prefilling-dLLM: Diffusion Language Models의 긴 문맥 추론을 위한 예측적 프리필링 (Predictive

Diffusion large language models (dLLMs)는 매 디노이징 (denoising) 단계마다 전체 접두사 (prefix)를 다시 인코딩하며, 이는 문맥 길이 (context length)에 따라 제곱으로 증가하는 재계산을 유발하여 긴 문맥 (long-context) 시나리오에서 비용이 매우 높아지게 만듭니다. 우리는 dLLM을 위한 별도의 학습이 필요 없는 프리필-디코드 분리 (prefill-decode disaggregation) 프레임워크인 Prefilling-dLLM을 제안합니다. 이 프레임워크는 접두사를 N개의 청크 (chunks)로 분할하고, 이들의 KV 표현 (KV representations)을 한 번만 캐싱하며, 디코딩을 위해 청크 내 토큰 희소성 (intra-chunk token sparsity)을 활용하여 가장 관련성이 높은 상위 K개의 청크를 선택합니다. 이를 통해 희소 프리필링 (sparse prefilling)이 밀집 어텐션 (dense attention)보다 뛰어난 성능을 보일 수 있음을 입증하였으며, 단계별 복잡도를 전체 시퀀스 길이에 대한 제곱에서 디코드 길이 (decode length)에 대한 제곱으로 감소시켰습니다. LongBench 및 InfiniteBench에서 Prefilling-dLLM은 dLLM 가속 방법들 중 최첨단 (state-of-the-art) 품질을 달성하였으며, 비연속적으로 캐싱된 청크 KV에 대해 디코딩을 병렬화하는 어텐션 커널 (attention kernel)을 통해 8K32K 문맥에서 9.128.0배의 속도 향상을 구현했습니다. 나아가, 각 청크 앞에 추가된 문장 시작 (beginning-of-sequence) 토큰들이 주기적인 어텐션 앵커 (attention anchors) 역할을 하여 'lost-in-the-middle' 현상을 제거함을 보여줍니다. 코드는 https://github.com/menik1126/Prefilling-dLLM 에서 확인할 수 있습니다.