D2F: 확산 언어 모델이 KV 캐시를 얻고 자기회귀 모델보다 2.5배 빨라지는 방법 (arXiv:2508.09192)

요약

확산 언어 모델(diffusion language model, dLLM)은 토큰 여러 개를 한 번에 디노이즈할 수 있다는 이론적 장점을 가지고 있다. 그러나 실제 추론 속도는 LLaMA3 같은 자기회귀(AR) 모델보다 훨씬 느렸다. 원인은 양방향 어텐션이었다.

D2F(Discrete Diffusion Forcing, arXiv:2508.09192)는 이 병목을 구조적으로 제거한다. 블록 내 양방향 어텐션과 블록 간 단방향 어텐션을 결합한 블록 단위 인과 어텐션으로 KV 캐시를 가능하게 하고, 선행 블록 완료 없이 후속 블록을 병렬 예측하는 인터블록 병렬 디코딩을 추가한다.

핵심 결과:

바닐라 dLLM(LLaDA, Dream) 대비 50배 이상 속도 향상
LLaMA3-8B / Qwen2.5 대비 2.5배 빠른 추론 (GSM8K 벤치마크)
품질 손실 없음 (수학·코딩 벤치마크 동등 성능)
12시간 학습 (8× A100-SXM4-40GB, Bespoke-Stratos-17k 데이터셋)

배경: 확산 언어 모델이 느렸던 이유

양방향 어텐션이 KV 캐시를 막는 구조

자기회귀 모델은 왼쪽 → 오른쪽 단방향으로 토큰을 생성한다. 각 디코딩 스텝에서 이미 처리한 토큰의 K·V 행렬을 재사용할 수 있다(KV 캐시). GPU 메모리 대역폭 비용이 선형이다.

dLLM은 다르다. 노이즈로 마스킹된 출력 전체를 반복해서 디노이즈한다. 각 스텝에서 모든 위치가 모든 위치를 참조한다(양방향 어텐션). 어떤 포지션 하나가 커밋되면 나머지 포지션의 K·V 값이 달라지므로 캐시를 재사용할 수 없다. 디노이징을 T번 반복하면 어텐션 비용이 O(N² × T)다.

바닐라 dLLM 디노이징 루프:

  [MASK MASK MASK MASK MASK]  ← 초기 상태
  ↓ 디노이징 스텝 1
  [MASK  the  MASK MASK  cat]
  ↓ 디노이징 스텝 2 (전체 재계산 필요)
  [ the   cat  sat  on  MASK]
  ↓ 디노이징 스텝 3 (전체 재계산)
  [ the   cat  sat  on   mat]

  각 스텝마다 모든 위치를 full attention으로 재계산 → KV 캐시 불가

기존 근사 캐싱의 한계

Fast-dLLM, dKV-Cache 같은 이전 연구는 K·V를 주기적으로 갱신하는 근사 캐싱으로 속도를 높였다. 그러나 근사화가 들어가면 품질 저하가 발생한다. 수학 문제나 코딩 과제처럼 정확성이 중요한 태스크에서는 허용하기 어렵다.

D2F는 다른 접근을 택했다. 어텐션 구조 자체를 바꿔서 근사 없이 KV 캐시를 사용할 수 있도록 만든다.

D2F 핵심 메커니즘

블록 단위 인과 어텐션

D2F는 출력 시퀀스를 고정 크기의 블록(block)으로 나눈다. 어텐션 규칙이 두 레벨로 다르다.

블록 내(intra-block): 같은 블록 안의 토큰들은 서로 양방향으로 참조 → 풍부한 지역 컨텍스트 유지
블록 간(inter-block): 블록은 이전 블록(완료된 블록)만 참조 → 단방향 인과 관계

이 설계 덕분에 완료된 블록의 K·V 행렬은 변하지 않는다. 표준 KV 캐시를 그대로 적용할 수 있다.

D2F 블록 단위 인과 어텐션 구조와 KV 캐시 재사용

인터블록 병렬 디코딩

D2F의 두 번째 혁신은 선행 블록 완료 없이 후속 블록을 예측할 수 있다는 점이다. 이론적으로 dLLM이 KV 캐시 없이 병렬 생성을 지원했던 특성을 블록 단위로 유지한다. 이전 블록이 디노이징을 마치는 동안, 다음 블록 후보 토큰을 병렬로 계산할 수 있다.

이 두 기능이 결합되면:

완료 블록 KV 캐시 → 이전 문맥 재계산 비용 제거
인터블록 병렬 디코딩 → 현재 블록과 다음 블록 추론을 겹쳐 실행

학습: 비대칭 증류

D2F 모델을 처음부터 학습하는 대신, 이미 학습된 양방향 dLLM(교사)에서 D2F 학생 모델로 증류한다.

교사 vs 학생의 역할

역할	어텐션 범위	목적
교사 (LLaDA-8B 원본)	전체 양방향	정확한 조건부 분포 계산
학생 (D2F)	블록 단위 인과	교사 예측 근사

교사는 전체 시퀀스를 바라보며 각 마스크 토큰의 최적 예측을 만든다. 학생은 이전 블록까지만 보면서 교사의 출력을 재현하도록 학습된다.

학습 효율

논문은 Bespoke-Stratos-17k 데이터셋으로 단 12시간 파인튜닝한 것만으로 목표 성능에 도달했다고 보고한다. 하드웨어는 NVIDIA A100-SXM4-40GB 8장이다. 모델 아키텍처 변경 없이 어텐션 마스크와 목적함수만 바꿔 기존 체크포인트를 재사용할 수 있다.

성능 결과

GSM8K 수학 벤치마크

모델	토큰/초	LLaMA3 대비
LLaDA-8B (바닐라)	~2	0.04×
Dream-7B (바닐라)	~2	0.04×
LLaMA3-8B	~50	1× (기준)
D2F (LLaDA 기반)	~125	2.5×

품질: D2F와 LLaMA3의 GSM8K 정확도는 통계적으로 동등하다.

바닐라 dLLM 대비

D2F는 바닐라 dLLM 대비 50배 이상 빠르다. 블록 크기 B를 크게 할수록 토큰 병렬성이 늘어나지만, 블록 내 어텐션 비용도 올라간다. 실험에서는 B = 16이 속도·품질 균형점으로 보고됐다.

운영 관점: dLLM 서빙을 고려하는 엔지니어에게

dLLM 도입 시 어디가 바뀌나

D2F를 채택하면 서빙 스택에서 두 가지가 달라진다.

1. KV 캐시 관리 방식

완료된 블록의 K·V는 영구적으로 캐시된다. PagedAttention 구현체(vLLM, SGLang)를 그대로 사용하되, 블록 완료 이벤트를 캐시 영구화 신호로 처리하는 수정이 필요하다. 현재 구현은 표준 HuggingFace 추론 루프 기반이므로 프로덕션 서빙 시스템 통합은 추가 작업이 필요하다(Open question).

2. 배치 스케줄링

인터블록 병렬 디코딩은 요청 간 블록 완료 시점이 달라 기존 연속 배칭(continuous batching)과 맞지 않는 경우가 생긴다. Sangam(arXiv:2607.04206)이 제안한 deficit token-budget scheduler 같은 dLLM 특화 스케줄링이 필요하다.

언제 D2F가 유리한가

수학·코딩처럼 긴 추론 체인이 필요한 태스크
짧은 TTFT보다 처리량 우선 워크로드
dLLM 고유의 병렬 토큰 커밋 특성을 살려야 하는 경우

AR 모델 대비 절대 우위는 아직 한정적이다. D2F가 2.5배 빠르다는 결과는 배치 크기 1 기준이다. 대형 배치에서 AR 모델도 KV 캐시 이점이 크기 때문에 실서비스 비교는 워크로드에 따라 다를 수 있다.

오픈 퀘스천

D2F가 긴 컨텍스트(100K+ 토큰)에서 AR과 비교해 얼마나 유리한가?
블록 크기 B를 동적으로 조정하는 adaptive D2F 가능성
PagedAttention, FlashAttention-4와의 통합 경로

요약

D2F는 dLLM의 양방향 어텐션 병목을 블록 단위 인과 어텐션으로 해결해 KV 캐시를 활성화하고, 50× 이상의 속도 향상으로 AR 모델보다 빠른 추론을 최초로 달성했다. 12시간 증류로 구현 비용이 낮고, 기존 dLLM 체크포인트를 재사용할 수 있다는 점도 실용적이다. dLLM 서빙 생태계(Sangam, Fast-dLLM 등)와 통합되면 AR과 dLLM의 서빙 효율 격차가 빠르게 줄어들 것으로 예상된다.