SpecFold가 확산 언어모델 추측 디코딩의 분기 중복을 접어 속도를 높인다

무엇인가

확산 대형 언어모델(DLLM)은 토큰 블록을 반복적으로 디노이징하며 텍스트를 생성한다. 다중 분기 추측 디코딩은 메인 분기와 여러 드래프트 분기를 한 번의 배치 forward pass로 검증해 미래의 디노이징 단계를 건너뛰는 기법이다. 논문은 전체 디코딩 비용을 "디노이징 forward pass 횟수 × forward pass당 비용"으로 분해한다. 기존 연구는 앞 항(추측으로 줄이는 횟수)에 집중했지만, 분기 수가 늘면 뒤 항(밀집 검증 비용)이 함께 커진다는 것이 이 논문이 지적하는 새 병목이다.

어떻게 동작하나

저자들은 각 추측 검증 단계 안에 아직 활용되지 않은 중복 축이 있다고 본다. 드래프트 분기는 부모로부터 대부분의 토큰을 물려받고 최대 S_t개 위치만 추가로 언마스킹하므로, 은닉 상태가 분기 간에 거의 같게 유지된다. Fast-dLLM-v2-7B(블록 크기 32, 드래프트 B=8)의 28개 레이어를 프로파일링한 결과, (드래프트, 스텝, 레이어, 위치) 조합의 72%가 부모-자식 상대 잔차 ρ≤0.1, 88%가 ρ≤0.3을 만족했고, 큰 발산은 새로 언마스킹된 위치 근처에 몰렸다. 또한 Nemotron-Labs-Diffusion-3B에서 드래프트를 2개에서 8개로 늘리면 forward pass당 지연이 계속 증가해, 밀집 검증이 모든 분기의 모든 위치를 계산하고 있음이 확인됐다.

무엇과 다른가

알고리즘의 핵심은 토큰 수준 부모-자식 잔차 게이트다. 각 레이어와 블록, 분기 i, 위치 n마다 ρ = ||H_i[n] − H_par(i)[n]|| / ||H_par(i)[n]|| 를 계산하고, 임계값 δ 이하이면 그 위치를 "접힌(folded)" 집합에 넣는다. 접힌 위치는 부모의 QKV와 FFN 출력을 그대로 물려받고, 어텐션은 부모의 스트리밍 어텐션 상태를 재사용하되 자식이 다시 계산한 키·값으로 보정해 해당 분기 자신의 KV 캐시 기준으로 정확하게 만든다. 각 분기는 잔차 스트림을 따로 유지해 분기 간 다양성이 레이어를 넘어 누적되게 하고, 그 누적으로 ρ가 δ를 넘으면 해당 위치는 다시 계산된다. 즉 QKV 투영과 FFN은 재계산 위치 집합에 대해서만 실행된다.

어떻게 쓰나

시스템 구현은 이 세밀한 재사용을 Triton 커널로 옮긴다. 게이트 커널이 잔차를 평가하고 접힌 위치의 출처를 기록하며 분기 전체의 재계산 위치를 압축해, 투영과 FFN이 압축된 배치 위에서 돌아간다. 소스 해석 커널은 접힘 체인을 따라가 접힌 위치가 부모 상태에 접근하게 하고, 어텐션은 FlashAttention 스타일 스트리밍 커널로 처리한다. 공유 프리픽스 KV 캐시는 분기 전체에 제로카피 스트라이드 뷰로 확장되며 그룹 쿼리 어텐션도 지원한다. 이 기법은 시간적 캐싱과 직교하고, 드래프트 생성 정책을 바꾸지 않아 Spiffy 같은 기존 추측 전략 위에 그대로 얹을 수 있다.

전제와 한계

실험은 Fast-dLLM-v2(1.5B, 7B)와 Nemotron-Labs-Diffusion(3B, 8B, 14B) 두 계열, 다섯 모델에서 GSM8K·MATH·HumanEval·MBPP·IFEval 다섯 벤치마크로 수행했고, 단일 NVIDIA H100 80GB에서 돌렸다. 베이스라인은 각 모델의 기본 블록 디코딩(Vanilla)과 Spiffy이며, Spiffy가 만든 드래프트 그래프를 SpecFold가 그대로 재사용한다. Fast-dLLM-v2-1.5B에서 Spiffy 대비 최대 1.64배, Vanilla 대비 최대 1.99배 처리량을, 7B에서 각각 최대 1.22배와 1.31배를 기록했다. Nemotron-Labs-Diffusion 3B·8B·14B에서는 Spiffy 대비 최대 1.34배·1.18배·1.07배, Vanilla 대비 최대 1.48배·1.32배·1.24배였다. 정확도는 Vanilla·Spiffy와 비슷한 수준을 유지했고, Spiffy의 정확도 변동은 부동소수점 잡음이며 직렬 실행 검증으로 구현 정확성을 확인했다고 밝힌다.

절제 실험도 구체적이다. Nemotron-Labs-Diffusion-8B에서 δ를 0.01부터 0.7까지 쓸어보면, δ=0.2까지는 정확도가 Vanilla 수준을 유지하면서 Vanilla 대비 1.3~1.4배 처리량 이득이 나오고, δ=0.3을 넘으면 처리량은 계속 늘지만 정확도가 떨어진다. δ를 0에 가깝게 만들어 모든 드래프트를 재계산하면 결과가 Spiffy에 수렴해 구현이 검증된다. 본 실험에서는 태스크별 튜닝 없이 δ=0.1을 보수적 기본값으로 고정했다. 드래프트 예산을 B=2에서 8까지 늘린 실험에서는 수용률이 B=6 근처에서 포화되며 NFE가 Vanilla 대비 30% 줄었는데, Spiffy는 밀집 검증 비용 때문에 Vanilla TPS의 0.99~1.17배에 머문 반면 SpecFold는 같은 드래프트와 NFE로 최대 1.73배 처리량 이득을 냈다. 또한 토큰 하나만 다른 부모-자식 쌍을 분석해, 발산이 편집된 위치에 지속적으로 남고 확정된 위치에서는 낮게 유지되며 마스킹·최근 언마스킹 위치에 집중되는 "잔차 리플" 구조를 확인했다.

개발자 관점에서 이 논문은 확산 LLM 추론 서버에서 추측 디코딩을 켤 때 분기 수를 늘리는 것만으로는 처리량이 늘지 않는다는 점을 데이터로 보여준다. Spiffy 같은 기존 스킴을 쓰고 있다면 드래프트 그래프 구성은 그대로 두고 검증 단계의 커널만 교체하는 방식으로 이득을 볼 수 있고, δ 하나로 지연-정확도 트레이드오프를 조절할 수 있다는 점이 실무적으로 유용하다. 다만 이득은 모델이 커질수록 줄어든다(14B에서 Spiffy 대비 1.07배).

논문이 제시한 텍스트에는 별도의 한계 절이 없고 윤리 성명만 있다. 대신 저자들이 밝힌 전제는 분명하다. δ는 태스크별 튜닝 없이 0.1로 고정한 보수적 값이며 δ가 0.3을 넘으면 정확도가 희생된다. 드래프트 그래프는 Spiffy 방식으로 모델당 50개 샘플(MBPP·MATH 학습 분할에서 25개씩, 평가 데이터와 분리)로 보정해야 하고, Fast-dLLM-v2는 블록 캐싱을 끈 기본 설정, Nemotron-Labs-Diffusion은 블록 경계마다 프리픽스 KV 캐시를 갱신하는 설정을 전제로 한다. 이득은 Triton 커널 구현에 의존하므로 다른 하드웨어·런타임으로의 이식성은 별도 검증이 필요하다.