Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs

Flash-Verify: 캐시 기반 초안-검증 병렬 디코딩

기존 신뢰도 기반 디코딩은 $c_i \ge \epsilon$ 인 토큰만 수락하고 나머지는 버린다 (근거: §2.3). 불확실한 과제에서는 스텝당 수락량이 천장이 된다.

Flash-Verify는 드래프트 패스에서 $Q_t = T_t \cup M_t^{\beta_m}$ 에 대해 전체 KV 캐시를 조회하고 마스크드 위치를 신뢰도 순으로 정렬해 확신 집합 $D_t$ 와 탐색 집합 $S_t$ 로 나눈다 (근거: Alg. 2). 검증 패스는 세 그룹으로 새 쿼리를 만든다. 조정된 추적 집합 $T_v$ , $S_t$ 에 초안 예측을 채운 드래프트 뷰, 같은 위치에 [MASK] 를 채운 마스크 뷰이다. 두 뷰는 위치 임베딩을 공유하지만 인과적 마스크로 서로 볼 수 없다 (근거: §2.3, Fig. 2).

수락 규칙은 단순하다:

$\hat{x}_i$ 는 드래프트 예측, $\tilde{x}_i$ 와 $\tilde{c}_i$ 는 마스크 뷰 예측과 신뢰도다. 디코딩 순서대로 첫 불일치까지만 순차 수락한다 (근거: §2.3). 추가 비용은 전체 시퀀스가 아니라 $2\beta_m$ 에 비례하며 외부 모델이 필요 없다.

이론적 보장도 있다 (근거: Appx. C). $p_i$ 를 마스크 뷰 검증 조건부 분포, $\pi_i$ 를 커밋된 점질량이라 하면 수락된 토큰마다 다음이 성립한다:

블록 $A = \{i_1,\dots,i_m\}$ 에 대해서는 $\text{TV}(\pi_A, p_A) = 1 - \prod_k \tilde{c}_{i_k} \le 1 - \gamma^m \le m(1-\gamma)$ 다. 검증 마스크가 인과적이므로 $p_A$ 는 주변 분포의 곱이 아니라 모델의 진짜 조인트이며 검증 토큰은 병렬 언마스킹 의존성 오차를 내지 않는다 (근거: Appx. C, Remark 1). $\epsilon$ 을 낮추는 것과 검증은 동등하지 않다. $\epsilon$ 은 모두 마스크된 상태의 초안 주변 분포를 가르고 $\gamma$ 는 실제로 커밋하는 분포를 가르기 때문이다 (근거: Appx. C, Remark 2).

작동 원리: 구체적인 예시로 살펴보기

대학원생을 위한 미니 예시를 만들자. 어휘가 {A, B, [MASK]} 뿐이고 생성 길이 $N=4$ , $\beta_m=2$ , $\beta_t=1$ , $\epsilon=0.9$ , $\gamma=0.8$ 이라 하자 (근거: §2.2-§2.3의 변수 정의 차용).

입력 $x_0 =$ [BOS] [MASK] [MASK] [MASK] [MASK] , $D_{<1} = \{1\}$ , $M_1 = \{2,3,4,5\}$ 라 두자. 여기서 $Q$ 는 쿼리 위치, $K$ 와 $V$ 는 캐시, $D_t$ 는 새로 디코딩된 위치, $M_t^{\beta_m}$ 은 앞쪽 $\beta_m$ 개 마스크, $T_t$ 는 추적 집합, $S_t$ 는 검증 대상이다.

스텝 1: 드래프트. $M_1^{\beta_m}=\{2,3\}$ 이고 $T_1=\emptyset$ 이므로 $Q_1=\{2,3\}$ 이다. 융합 커널이 $x[Q_1]$ 만 읽어 $q,k,v$ 를 SRAM에서 만들고 캐시에 직접 쓴다 (근거: Alg. 1). 전체 캐시에 대해 어텐션한 뒤 모델이 $q_2=[0.92,0.05,0.03]$ , $q_3=[0.70,0.20,0.10]$ 을 냈다고 하자. 신뢰도는 $c_2=0.92$ , $c_3=0.70$ 이다. $\epsilon=0.9$ 이므로 $D_1=\{2: \hat{x}_2=A\}$ , $S_1=\{3: \hat{x}_3=A\}$ 이다.

스텝 1: 검증. 검증 쿼리는 [추적, D_1=A, S_1 드래프트=A, S_1 마스크=[MASK]] 형태다. 위치 3의 드래프트 뷰와 마스크 뷰가 서로를 보지 못하도록 마스크를 씌운다. 마스크 뷰가 $p_3=[0.85,0.10,0.05]$ 를 내면 $\tilde{x}_3=A$ , $\tilde{c}_3=0.85$ 다. $\hat{x}_3=\tilde{x}_3$ 이고 $0.85 \ge \gamma$ 이므로 위치 3도 수락된다. 기존 방식이었다면 0.70 때문에 버려졌을 토큰이다.

  flowchart LR
  Q["Q = T + Mβm<br/>예: {2,3}"] --> F["FusedForward<br/>Flash-Cache"]
  F --> D["D: c ≥ ε<br/>즉시 수락"]
  F --> S["S: c < ε<br/>검증 후보"]
  S --> V["Verify: draft뷰 vs mask뷰<br/>인과 마스크"]
  V --> A["합의 + c̃ ≥ γ<br/>첫 불일치까지 수락"]

스텝 1: 캐시 선택. $a_i^1$ 을 계산해 이전 디코딩 토큰 중 마스크드 쿼리 $\{2,3\}$ 으로부터 가장 주목받은 1개를 고른다. 새로 풀린 $\{2,3\}$ 은 자동 포함 후 랭킹하므로 $T_2$ 는 크기가 1로 유지된다. 다음 쿼리는 $Q_2 = M_2^{\beta_m} \cup T_2$ 다.

위 그림이 왜 이 예시가 통하는지 보여준다. 왼쪽 패널처럼 중간층은 소수 토큰에 집중하고 오른쪽 패널처럼 신뢰도를 올리면 스텝당 수락량이 늘어난다 (근거: Fig. 1b-c).

두 번째 그림은 드래프트 정렬부터 $\epsilon$ 과 $\gamma$ 임계값 판정까지의 실제 검증 파이프라인에 대응한다 (근거: Fig. 2).

비밀 병기 1개를 뜯어보면

퓨즈드 커널과 자체 검증 중 하나만 고른다면 자체 검증을 고른다. 커널이 상수를 깎는다면 검증은 기울기를 바꾸기 때문이다 (근거: §3.3).

메커니즘은 명확하다. 신뢰도 기반은 주변 분포 $q_i$ 에서 $\epsilon$ 을 넘지 못하면 정답이어도 버린다 (근거: Appx. C). Flash-Verify는 앞선 초안을 조건으로 넣은 $p_{i_k}$ 에서 다시 묻는다. 문맥 증거가 모이면 $\tilde{c}$ 가 올라가 $r=0.48$ 상관관계가 처리량으로 전환된다 (근거: Fig. 1c, Fig. 6). $\gamma=0.60$ 에서 7.2 tokens/step 를 풀어도 신뢰도 기반 5.6 tokens/step 와 처리량은 비슷하면서 정확도는 3.5%p 높다. 토큰/반복이 커질수록 정확도 격차는 줄고 처리량 우위는 최대 1.33× 로 벌어진다 (근거: §3.3, Fig. 6).

성능 검증: 주요 결과

실험 설정

단일 NVIDIA A100 80GB GPU, Triton 2.0 퓨즈드 커널, LLaDA-1.5 평가가 기본이다 (근거: §3.1). 벤치마크는 GSM8K 5-shot, MATH 4-shot, HumanEval 0-shot, MBPP 3-shot이며 생성 길이 256 tokens 와 512 tokens 를 본다 (근거: §3.1, Appx. D). 기본 하이퍼파라미터는 $\epsilon=0.9$ , $\gamma=0.8$ , 블록 크기 $\beta=16$ , $\beta_t=80$ , $\beta_m=64$ 다 (근거: §3.1). 처리량은 lm-eval-harness 기준 종료 토큰까지의 평균 tokens/s 다 (근거: Appx. D). 토큰화와 사전학습 코퍼스는 LLaDA-1.5 기반 모델을 그대로 쓰며 추가 학습이나 파인튜닝은 없다. 파인튜닝 전략은 해당 없음이다 (근거: §3.1, Abstract).

위 그림은 배치 확장 실험을 요약한다. 배치 크기 32까지 선형에 가깝게 오르고 Fast-dLLM은 24에서 OOM이 나며 메모리도 전 구간에서 가장 낮다 (근거: Fig. 4).

핵심 성능 지표

• GSM8K-512: Flash-Cache + Flash-Verify가 83.02% 정확도와 210.6 tokens/s 로 정확도와 처리량 모두 1위다. 탐욕적 무캐시 2.6 tokens/s 대비 81.0× 다 (근거: Tab. 1).
• MATH-512: 35.98% , 210.1 tokens/s , 42.0× 다. 최고 정확도는 무캐시 Flash-Verify의 37.76% 다 (근거: Tab. 1).
• HumanEval-512: 40.24% , 185.6 tokens/s , 58.0× 다. 최고 정확도는 신뢰도 기반 Flash-Cache의 42.07% 다 (근거: Tab. 1).
• MBPP-512: 39.00% , 148.2 tokens/s , 148.2× 다. 탐욕적 Flash-Cache가 39.80% 로 정확도는 앞선다 (근거: Tab. 1).
• 확장 비교: dKV-Cache 14.9 tokens/s , Fast-dLLM 36.8 tokens/s , FreeDave 42.8 tokens/s , Elastic-Cache 41.7 tokens/s 대비 Flash-Cache + Flash-Verify는 210.6 tokens/s 다 (근거: Tab. 2).
• 메모리: 배치 크기 16에서 약 26 GB 대 Fast-dLLM 약 50 GB 로 약 48% 절감이다. 평탄한 사전 할당 캐시가 패딩 오버헤드를 없앤 덕분이다 (근거: §3.3, Fig. 4b).
• 마스크드 윈도 $\beta_m$ : 16에서 96으로 키우면 처리량은 줄지만 Flash-Verify가 신뢰도 기반 대비 1.4× 에서 1.5× 로 우위를 유지하고 정확도 격차는 0.2%p 에서 1.2%p 다 (근거: §3.3, Fig. 7).
• 생성 길이와 프리필: 256 tokens 에서 두 방식 모두 처리량 정점을 찍고 1024 tokens 로 가면 약 10.5% 에서 12.8% 하락한다. 1-shot에서 8-shot으로 가면 Flash-Verify 우위는 1.50× 에서 1.35× 로 줄지만 절대값은 42.6 tokens/s 에서 78.6 tokens/s 앞선다 (근거: Appx. F, Tab. 5).

비판적 비교

가장 강한 비교 지점은 동일 하드웨어에서 재실행한 Elastic-Cache와 FreeDave 대비 결과다 (근거: §3.1, Tab. 2). Elastic-Cache는 베이스라인 중 최고 정확도 82.79% , FreeDave는 최고 처리량 42.8 tokens/s 였지만 Flash-Cache + 신뢰도 기반이 82.87% 와 149.4 tokens/s 로 둘 다 넘고 Flash-Verify 결합이 83.02% 와 210.6 tokens/s 로 한 번 더 뛴다.

반대로 코드 256 tokens 설정에서는 처리량 1위가 정확도 1위가 아니다. HumanEval-256에서 39.63% 대 최고 43.29% 로 3.66%p, MBPP-256에서 38.20% 대 최고 41.80% 로 3.60%p 격차가 난다 (근거: §3.2). 수학 추론에서는 1.78%p 이내지만 코드 단문 생성에서는 적극적 병렬 커밋이 엄밀한 단일 토큰 정제를 대체하지 못하는 구간이 있다. 신뢰도 기반이 약 83.4% 로 피크 정확도는 조금 더 높지만 처리량은 약 131 tokens/s 로 떨어진다 (근거: §3.3, Fig. 5b). 즉 파레토 프론티어는 Flash-Verify 쪽이 더 바깥에 있지만 극단적 정확도一点을 노린다면 보수적 디코딩이 남는다.

시스템 관점 지표

주 지표는 지연 시간이 아니라 처리량 tokens/s 와 피크 메모리 GB 다 (근거: §3.3). 배치 크기 1에서 32로 가면 탐욕적 방법은 18.5 tokens/s 에서 55.0 tokens/s 로 2.97× , 신뢰도 기반은 51.7 tokens/s 에서 139.5 tokens/s 로 2.70× , Flash-Verify는 56.0 tokens/s 에서 199.8 tokens/s 로 3.57× 가 된다 (근거: Appx. G, Tab. 6). 학습 FLOPs나 $/1M tokens 비용은 보고하지 않았으며 에너지 지표도 없다. 재현 체크리스트로는 코드 공개, Triton 2.0 명시, A100 80GB 명시, 5개 시드 표준편차 보고는 있으나 드라이버 버전과 정확한 프롬프트 전문은 부록 평가 스크립트에 의존한다 (근거: §3.1, Appx. D-E).

우리의 관점: 강점, 한계, 그리고 이 연구가 중요한 이유

강점. 저자들의 주장은 일관된다. 캐시 재사용과 병렬 검증을 따로 최적화하지 않고 같은 융합 커널과 사전 할당 캐시 위에서 묶었다는 점이다 (근거: §2.3). 외부 검증기가 없다는 점은 배포 관점에서 크다. FlashDLM이 외부 자가회귀 검증기를 쓰고 FreeDave가 2회 독립 패스를 쓰는 것과 달리 추가 모델과 학습이 없다 (근거: §4). 이론 부록이 단순한 직관을 TV 거리 보장으로 닫은 것도 드물다 (근거: Appx. C).

명시적 한계. 저자들은 두 개 마스크드 확산 LLM과 수학·코드 구조화 출력 위주로 평가했음을 인정한다 (근거: Appx. A). 연속 공간 확산, 장문 쓰기나 대화 같은 평탄한 신뢰도 분포, 고정된 $\gamma$ 와 $\beta_m$ 은 미검증 영역이다.

잠재적 한계. 첫째 강한 가정이다. 보장은 모델 자신의 순차 분포에 대한 편차이지 데이터 분포에 대한 정확도가 아니다 (근거: Appx. C, Assumption 1). 둘째 $D_t$ 의 병렬 오차는 여전히 신뢰도 기반 정리 $[46, Thm. 1]$ 에 의존한다. 셋째 비용 투명성이 약하다. 훈련 연산량은 해당 없지만 추론 비용과 전력은 없다. 넷째 사회적 영향은 기반 dLLM의 유해·편향 생성을 그대로 상속하며 안전 필터를 바꾸지 않는다 (근거: Appx. B).

그럼에도 중요한 이유는 dLLM 실용화의 병목을 FLOPs가 아니라 I/O와 스케줄링으로 재정의했기 때문이다. 배치 크기 32까지 OOM 없이 처리량을 유지하고 메모리를 절반 가까이 줄인 결과는 서빙 시스템 설계에 직접 쓸 수 있다 (근거: Fig. 4, Tab. 6).

다음 단계는?: 앞으로의 길

저자들은 적응형 $\gamma$ 와 $\beta_m$ , 개방형 생성 검증, 연속 공간 확산으로 확장을 다음 방향으로 제시한다 (근거: Appx. A). 합리적인 다음 단계는 세 가지다.

• 실행 중 신뢰도 통계 기반 스케줄러: 고정 임계값 대신 프리픽스가 길거나 불확실성이 클 때 $\beta_m$ 을 줄이고 $\gamma$ 를 올리는 제어기를 붙인다. 부록의 1-shot에서 8-shot까지 상대 속도 하락을 완화할 수 있다 (근거: Appx. F).
• 긴 문맥과 배치 서빙 결합: 블록 테이블 설계를 PagedAttention류 서빙 스택과 붙여 배치 크기 16 이후 포화 구간을 미는 실험이 필요하다. 현재도 배치 크기 16이 배치 크기 32의 약 93.2% 를 내므로 효율점이 명확하다 (근거: Appx. G).
• 정확도 민감 코드 과제 분리: 256 tokens 코드에서 생긴 3%p대 격차를 닫기 위해 검증 실패율이 높은 위치만 보수적으로 되돌리는 하이브리드 커밋을 시험한다 (근거: §3.2).

논문 원문의 표

arXiv e-print 의 LaTeX 원본에서 기계적으로 옮긴 표입니다. 숫자는 논문의 값이며 모델을 거치지 않았습니다.

표 1. Accuracy and decoding efficiency of LLaDA-1.5 across different benchmarks and decoding configurations. Each cell reports accuracy (top) and throughput with speedup over greedy decoding without caching (bottom; blueblue: tokens/s, orangeorange: speedup). Bold indicates the highest accuracy in each row, while yellow!20yellow shading indicates the highest throughput.

표 2. Comparison of accuracy and decoding throughput across different methods.

표 3. The hyper-parameters of Flash-dLLM under various settings.

표 4. Mean accuracy (%) and throughput over five random seeds. Accuracy is shown on the first line and throughput is shown in blue on the second line. Values are mean $\pm$ standard deviation.

표 5. Ablation of generation and prefill lengths.

표 6. Effect of batch size on throughput. Using LLaDA-1.5 for GSM8K-512 (5-shot). Speedup is measured relative to Flash-Cache with greedy decoding at the same batch size.

표 7. Example LLaDA-1.5 (without cache) — A qualitative comparison of different decoding methods.

표 8. Example Flash-dLLM: A qualitative comparison of different decoding methods, LLaDA-1.5

이 글의 그림은 arXiv:2609.26796 원본에서 가져왔습니다 (CC BY 4.0). 크기와 형식만 바꿨습니다.