SpecFold: Folding Multi-Branch Redundancy for Faster Speculative Decoding in Diffusion Language Models

Diffusion large language models (DLLMs) generate text through iterative block denoising, and multi-branch speculative decoding accelerates this process by verifying a main branch together with multiple draft branches in a single forward pass. While prior DLLM acceleration methods primarily exploit temporal redundancy across denoising steps, we identify a complementary redundancy axis within each speculative verification step: multi-branch computational redundancy. During speculative verification, draft branches inherit most tokens from their parents while unmasking a small set of additional positions, causing large portions of hidden states to remain highly similar across branches. We propose SpecFold, an algorithm-system co-design that exploits this multi-branch redundancy to reduce the cost of multi-branch speculative verification. Algorithmically, SpecFold performs token-level residual gating and selectively reuses parent computation through folded attention and FFN while preserving residual hidden states. Systemically, a Triton kernel implementation translates this fine-grained reuse into end-to-end throughput gains through efficient sparse multi-branch execution. SpecFold is orthogonal to temporal caching and compatible with existing DLLM speculation strategies. Across two DLLM families, five models, and five standard benchmarks, SpecFold achieves up to 1.64x throughput over Spiffy and up to 1.99x over vanilla decoding, while maintaining comparable task performance.

한 줄 요약

SpecFold는 다중 분기 추측 디코딩에서 계산 중복성을 활용해 디퓨전 언어 모델의 처리 속도를 최대 1.99배 향상시키는 알고리즘-시스템 협력 설계이다.

핵심 기여도

다중 분기 추측 디코딩에서 발생하는 **계산 중복성**(multi-branch computational redundancy)을 새로운 효율화 기회로 제시.
**SpecFold** 알고리즘은 토큰 수준 잔여 게이트와 **folded attention**, **FFN 재사용**을 통해 분기 간 계산 중복을 제거.
**Triton 커널**을 통해 **fine-grained reuse**를 효율적인 GPU 실행으로 변환.
**Spiffy 대비 최대 1.64×**, **Vanilla decoding 대비 최대 1.99×** 처리 속도 향상.

핵심 아이디어

기존 디퓨전 언어 모델(DLLM)의 디코딩 가속화 연구는 주로 **시간적 중복**(temporal redundancy)을 활용했으나, 본 연구는 **분기 간 계산 중복**(multi-branch redundancy)이라는 새로운 축을 제시한다. 추측 디코딩 과정에서, **draft branches는 부모 분기에서 대부분의 토큰을 상속**받고, **추가로 언마스크된 위치만 다름**으로 인해, **은닉 상태가 분기 간 매우 유사**하다는 점을 관찰했다.

이를 기반으로, **SpecFold**는 토큰 수준에서 **잔여 게이트**(residual gate)를 적용해 부모 계산을 선택적으로 재사용하고, **folded attention**과 **FFN 재사용**을 통해 중복 계산을 제거한다. 특히, **bidirectional attention** 구조를 고려해 부모의 attention state를 재사용하면서 자식의 query, key, value로 보정하는 방식을 제안한다.

기술적 접근법

**Algorithm**:
**Token-level residual gating**을 통해 유사한 토큰 위치는 부모 계산을 재사용 (folding).
**Folded attention**: 부모의 attention state를 재사용, 자식의 query, key, value로 보정.
**Folded FFN**: FFN 출력 재사용.
**Residual stream 유지**: 분기 간 다양성 유지.
**System**:
**Triton 커널**로 **branch 간 재사용 위치를 compaction**하고, **attention correction** 수행.
**Redundant branch tensor 생성 방지**.
**Compatibility**:
**Spiffy**와 같은 기존 추측 전략과 호환.
**Temporal caching**과 독립적.

주요 결과

**Spiffy 대비 최대 1.64×**, **Vanilla decoding 대비 최대 1.99× 처리 속도 향상**.
**Fast-dLLM-v2**, **Nemotron-Labs-Diff** 등 2개 DLLM 패밀리, 5개 모델, 5개 벤치마크에서 평가.
**72% 이상의 draft–layer–position–step 튜플에서 부모-자식 은닉 상태 유사도 높음**.

의의 및 한계

**의의**:
기존 DLLM 디코딩 가속화 방법과 **정교한 중복 제거 기법**(fine-grained reuse)을 결합.
**Triton 커널 기반의 GPU 최적화**로 실제 처리 속도 향상 달성.
**기존 추측 전략과 호환**되며, **draft 생성 정책 변경 없이 적용 가능**.
**한계**:
**추측 분기 간의 유사도가 낮은 경우**(예: 높은 발산 분기) 성능 개선 효과가 제한적일 수 있음.
**모든 DLLM 아키텍처에 일반화 가능**한지에 대한 추가 검증 필요.

실용적 활용

**대규모 DLLM 기반의 실시간 텍스트 생성** 시스템에서 처리 지연 감소에 활용 가능.
**고성능 GPU 환경**에서 Triton 커널 기반의 효율적 실행으로, **클라우드 인프라 및 엔드-유저 애플리케이션**에 적용 가능.
**기존 추측 디코딩 전략**(예: Spiffy)과 병행하여 **추가 가속화**를 도모할 수 있음.