dLLM 추론의 메모리 I/O 병목을 겨냥한 융합 KV 캐시와 자기 검증 병렬 디코딩

무엇인가

이 논문이 다루는 문제는 확산 언어모델(dLLM)의 추론이 실전에서 너무 느리다는 것이다. dLLM은 텍스트 생성을 반복적 디노이징으로 재정식화해 좌에서 우로만 토큰을 뽑는 자기회귀 모델보다 유연한 생성 스케줄과 병렬성을 제공하지만, KV 캐시와 어텐션 커널, 추측 디코딩 주변에서 수년간 최적화를 축적한 자기회귀 LLM 시스템에 비해 효율이 뒤처진다. 저자들이 지목하는 핵심 원인은 두 가지다. 첫째, dLLM은 디노이징 반복마다 시퀀스를 다시 방문하고 갱신되는 토큰 집합이 매번 바뀌기 때문에 KV 캐시의 읽기·쓰기·갱신이 빈번하게 발생한다. 순진한 KV 캐싱은 부동소수점 연산량을 줄여주지만 GPU 메모리 접근 오버헤드를 크게 늘리고, 실제로는 이 캐시 읽기/쓰기 중복이 런타임을 지배해 캐시 재사용으로 얻을 수 있는 실질적 속도 향상을 제한한다. 둘째, dLLM 디코딩에는 강한 토큰 수준 희소성이 있다. 각 디노이징 스텝이 전체 시퀀스를 처리하더라도 현재 예측 분포에 유의미한 영향을 주는 토큰은 소수뿐이다. 저자들은 LLaDA-1.5에서 중간 레이어(5층에서 20층)의 경우 어텐션을 가장 많이 받는 상위 32개 토큰이 전체 어텐션 가중치의 약 50%를 차지한다는 관찰을 제시하며, 모든 캐시 상태를 균등하게 다루는 것이 메모리 대역폭과 연산을 낭비한다고 주장한다. 기존 가속 기법들이 KV 캐싱과 병렬 디코딩을 각각 따로 연구해, 캐시 재사용과 병렬 토큰 검증을 동시에 적용할 때 생기는 I/O 병목을 놓쳤다는 것이 이 논문의 출발점이다.

어떻게 동작하나

제안 시스템 Flash-dLLM은 학습이 필요 없는 추론 가속 프레임워크로, 크게 Flash-Cache와 Flash-Verify 두 축으로 구성된다. Flash-Cache의 첫 요소는 I/O 인지 융합 커널이다. 기존 구현은 각 트랜스포머 레이어에서 QKV 프로젝션, 회전 위치 임베딩(RoPE), 캐시 쓰기, 어텐션을 각각 별도의 CUDA 커널로 실행하며, 각 커널이 출력을 GPU HBM에 쓴 뒤 다음 커널이 다시 읽는다. 이때 레이어당 HBM 트래픽은 대략 4×O(Q·d_model) + O(N·d)에 이른다. 비어텐션 연산들은 산술 강도가 낮아 캐시 갱신 경로가 메모리 바운드가 된다. Flash Attention에서 착안한 융합 커널은 QKV 프로젝션과 RoPE를 SRAM 안에서 수행하고 그 결과인 key와 value를 KV 캐시에 직접 기록해 중간 key/value 텐서의 materialization을 제거한다. 그 결과 HBM 트래픽과 메모리 사용량이 줄고, RTX 3090에서 1.37× 속도 향상을 얻었다. 두 번째 요소는 스케줄드 플래시 어텐션이다. KV 캐싱은 작은 윈도우만 계산하는 캐싱 단계와 전체 시퀀스를 재계산하는 업데이트 단계를 번갈아 수행하므로 같은 배치 안의 샘플들이 매 반복마다 전혀 다른 양의 계산을 요구할 수 있다. 저자들은 배치를 여러 시퀀스 블록으로 나누고 블록 테이블로 쿼리 블록과 대응하는 KV 블록을 정렬해 스케줄링함으로써 길이 편차를 완화하고, 블록을 유연하게 추가·제거해 언제 캐시하고 언제 갱신할지에 대한 적응적 결정을 가능하게 한다. 세 번째 요소는 선택적 캐시 갱신이다. 각 디코딩 스텝에서 고정 크기 쿼리를 슬라이딩 언마스킹 윈도우 β_m과 추적 예산 β_t로 구성하는데, 마스크된 쿼리가 이전에 디코딩된 위치 i에 실제로 얼마나 어텐션을 주는지를 레이어와 헤드에 걸쳐 평균한 점수 a_i^t를 계산하고 그 상위 k개를 추적 집합 T^(t+1)로 삼는다. 새로 디코딩된 토큰은 순위 계산 전에 자동 포함되고, 나머지 디코딩 위치는 쿼리로 참여하지 않고 캐시에서 서빙되어 스텝당 연산량이 β_t + β_m으로 제한된다.

무엇과 다른가

Flash-Verify는 dLLM 자체가 drafter이자 verifier가 되는 자기 검증 병렬 디코딩이다. 기존 confidence-aware 디코딩은 확신도 c^i = max_x p_θ(x^i | x_t)가 임계값 ε을 넘는 토큰만 언마스킹하고 나머지는 버리는데, 모델이 불확실한 과제에서는 스텝당 통과하는 토큰이 적어 처리량 상한이 생긴다. Flash-Verify는 이때 버려지는 정답 예측을 회수한다. 각 디노이징 스텝에서 표준 draft 패스가 전체 KV 캐시에 대해 쿼리 집합을 실행하고, 마스크된 위치들을 확신도로 정렬해 ε 이상인 확신 집합 D^t(즉시 수락)와 ε 미만인 탐색 집합 S^t(검증 후보)로 나눈다. 이어지는 verify 패스는 세 그룹으로 새 쿼리를 만든다. 이전에 디코딩된 토큰과 D^t로 이루어진 조정된 추적 집합 T_v, 탐색 위치에 draft 예측을 채운 draft view, 같은 위치에 [MASK]를 채운 mask view다. 두 view는 위치 임베딩을 공유하지만 융합 Triton 커널 내부에 로드된 causal 어텐션 마스크로 격리되어, 추적된 컨텍스트는 draft view를 참조할 수 없고 같은 위치의 draft view와 mask view도 서로를 참조하지 못하므로 공유 컨텍스트로부터 독립적인 예측을 낸다. 탐색 토큰은 두 view의 예측이 일치하고 mask view의 확신도가 임계값 γ 이상일 때 수락된다: accept(i) = I[x̂^i = x̃^i] · I[c̃^i ≥ γ]. 추측 디코딩 관례에 따라 디코딩 순서를 따라 순차적으로 수락하고 첫 불일치 이후의 토큰은 모두 거부한다. verify 패스는 같은 융합 커널과 미리 할당된 KV 캐시를 재사용하고 쿼리 토큰과 어텐션 마스크만 바뀌므로 추가 비용은 전체 시퀀스가 아니라 2β_m에 비례한다. 별도의 자기회귀 검증 모델이나 여러 번의 독립적 순전파, 추가 학습이 필요 없다는 점이 기존 dLLM용 draft-and-verify 기법과의 차이다.

어떻게 쓰나

실험은 단일 NVIDIA A100 80GB에서 LLaDA-1.5를 대상으로 GSM8K, MATH, HumanEval, MBPP 네 벤치마크에서 수행했고, 융합 KV 캐시 커널은 Triton 2.0으로 구현했다. 기본 하이퍼파라미터는 확신도 임계값 ε = 0.9, 검증 임계값 γ = 0.8, 블록 크기 β = 16, 추적 예산 β_t = 80, 슬라이딩 윈도우 β_m = 64, 생성 길이 512다. 비교 대상은 KV 캐싱 없는 표준 dLLM 추론(No Cache, greedy와 confidence-aware 양쪽), prefix caching과 confidence-aware 디코딩을 쓰는 Fast-dLLM, 어텐션 패턴 기반으로 캐시를 적응적으로 재사용하는 Elastic-Cache다. Flash-Cache는 greedy 디코딩에서 8.0×에서 58.5×의 속도 향상을 냈고, confidence-aware 디코딩과 결합하면 17.0×에서 102.2×로 범위가 넓어져 캐시 가속이 병렬 디코딩과 함께 쓸 때도 유효함을 보였다. Flash-Verify와 Flash-Cache를 결합한 구성은 여덟 개 설정 모두에서 최고 처리량을 기록해 148.0에서 210.6 tokens/s, 속도 향상 22.3×에서 148.2×를 달성했고, 두 번째로 빠른 구성인 confidence-aware Flash-Cache 대비 처리량을 약 23.5%에서 45.0% 끌어올렸다. 생성 길이가 길어질수록 이득이 커져 256토큰에서 512토큰으로 갈 때 GSM8K는 29.1×에서 81.0×, MATH는 22.3×에서 42.0×, HumanEval은 29.9×에서 58.0×, MBPP는 61.7×에서 148.2×로 증가했다. 초록 기준으로는 이전 최강 베이스라인인 Elastic-Cache 대비 GSM8K에서 5.1×, HumanEval에서 11.0× 속도 향상을 보고한다.

전제와 한계

정확도와 메모리 관련 수치도 구체적으로 제시된다. GSM8K-512에서 Flash-Verify와 Flash-Cache 결합 구성이 최고 정확도 83.02%와 최고 처리량 210.6 tokens/s를 동시에 달성했다. 다만 가장 빠른 구성이 항상 가장 정확한 것은 아니라는 점을 저자들도 명시한다. Flash-Verify를 캐싱 없이 쓴 구성이 GSM8K-256과 MATH-512에서 가장 정확했고, confidence-aware Flash-Cache는 HumanEval-512와 MBPP 두 설정에서 앞섰다. 수학 추론 과제에서는 결합 방법이 최고 정확도 대비 1.78%p 이내에 머물지만, 256토큰 코드 생성에서는 HumanEval 3.66%p, MBPP 3.60%p까지 격차가 벌어진다. 별도 비교에서 기존 베이스라인들은 정확도 79.32%에서 82.79%, 처리량 14.9에서 42.8 tokens/s를 기록했는데, Elastic-Cache가 베이스라인 중 최고 정확도 82.79%를, FreeDave가 최고 처리량 42.8 tokens/s를 냈다. Flash-Cache에 confidence-aware 디코딩을 붙이면 82.87%와 149.4 tokens/s, 여기에 Flash-Verify까지 결합하면 83.02%와 210.6 tokens/s가 된다. 확장성 측면에서는 GSM8K-512 1-shot에서 Flash-dLLM이 배치 크기 32까지 GPU 메모리 고갈 없이 거의 선형으로 처리량이 늘어난 반면 Fast-dLLM은 배치 24에서 OOM이 발생했다. 배치 16에서 Flash-dLLM은 약 26GB를 쓰고 Fast-dLLM은 50GB를 써서 약 48%의 메모리 절감을 보였는데, 저자들은 이를 평탄하게 미리 할당된 캐시 레이아웃이 동적 할당과 기존 4차원 KV 캐시 구현의 패딩 오버헤드를 피한 결과로 설명한다. GSM8K-512 5-shot에서는 모든 구성이 배치 1에서 32까지 단조롭게 확장하며 Flash-Verify 결합 구성이 배치 32에서 199.8 tokens/s에 도달했다.

절제 분석은 트레이드오프의 형태를 보여준다. 추적 예산 β_t와 검증 임계값 γ를 바꿔가며 측정한 결과 정확도와 처리량은 대체로 반비례해, 정확도를 약 80.2%에서 83.2%로 올리면 처리량은 278에서 186 tokens/s로 떨어진다. β_t를 키우면 추가 연산을 대가로 정확도가 오르고, γ는 고정 예산 안에서 병렬성을 더 세밀하게 조절하는 손잡이가 된다. 저자들이 제시하는 유리한 운용점은 β_t를 비교적 크게 유지하면서 γ를 낮춰 디코딩 병렬성을 높이는 것이다. Flash-Verify와 confidence-aware 디코딩을 직접 비교하면, 두 방법이 겹치는 정확도 구간 대부분에서 Flash-Verify가 처리량이 훨씬 높다. 약 82.6%에서 82.9% 정확도에서 Flash-Verify는 190에서 210 tokens/s를 내는 반면 confidence-aware는 140에서 160 tokens/s에 그친다. confidence-aware는 약 83.4%의 근소하게 더 높은 최고 정확도를 달성하지만 그때 처리량은 약 131 tokens/s로 떨어진다. 반복당 디코딩 토큰 수를 보면 Flash-Verify는 검증 스텝이 추가되는 비용 때문에 반복당 7.2토큰을 디코딩할 때 confidence-aware가 5.6토큰일 때와 비슷한 처리량을 내지만 정확도는 3.5% 더 높고, 반복당 토큰 수가 늘수록 두 방법의 정확도 격차는 좁아지면서 Flash-Verify의 처리량 우위는 최대 1.33×까지 커진다. 마스크 윈도우 크기 β_m을 키우면 두 방법 모두 중간 크기까지 정확도가 오르다가 처리량은 감소하는데, Flash-Verify는 모든 설정에서 정확도 격차가 0.2에서 1.2%p에 불과하고 처리량 우위는 윈도우가 커질수록 약 1.4×에서 1.5×로 증가한다.

한국 개발자 입장에서 이 논문의 실무적 의미는 명확하다. 학습이 필요 없는 프레임워크이므로 LLaDA 계열 dLLM 위에 그대로 얹을 수 있고, 별도의 drafter 모델이나 검증용 자기회귀 모델을 서빙 스택에 추가하지 않아도 된다는 점이 배포 부담을 낮춘다. 구현은 Triton 2.0 커널이고 코드가 VILA-Lab/Flash-dLLM 저장소에 공개되어 있다. 배치 추론과 긴 시퀀스에서 메모리 절감 폭이 크기 때문에 처리량과 동시 접속 수를 함께 끌어올려야 하는 서빙 환경에서 특히 유효하다. 다만 도입 전에 확인해야 할 것들이 있다. 정확도와 처리량이 명확한 트레이드오프 관계에 있으므로 β_t, γ, ε, β_m을 자신의 워크로드에 맞게 튜닝해야 하고, 코드 생성처럼 생성 길이가 짧은 과제에서는 최고 정확도 구성 대비 3%p대의 정확도 손실이 날 수 있다. 검증 패스가 추가 순전파 비용을 발생시키므로 반복당 수락 토큰 수가 충분히 늘지 않는 설정에서는 이득이 상쇄될 수 있다는 점, 그리고 캐시가 현재 쿼리 위치 Q^t에서는 정확하지만 나머지 위치에서는 근사라는 점도 함께 고려해야 한다.

한계와 전제에 관해 저자들은 별도의 한계 절을 두지 않았고, 원문에 명시된 실패 사례나 윤리적 논의도 제시되지 않았다. 원문에서 확인할 수 있는 전제는 다음과 같다. 모든 실험이 단일 NVIDIA A100 80GB에서 LLaDA-1.5 한 모델로, GSM8K·MATH·HumanEval·MBPP 네 벤치마크에 한정되어 수행되었다. 선택적 캐시 갱신은 근사 캐시로, 현재 쿼리 집합에 속한 위치에서만 캐시가 정확하고 나머지는 근사값이며, 상위 32개 토큰이 어텐션의 약 50%를 차지한다는 관찰은 중간 레이어(5층에서 20층)에 국한된 것이다. 또한 정확도와 처리량의 트레이드오프가 존재해 어떤 설정에서는 Flash-dLLM이 최고 정확도 구성이 아니며, Flash-Verify는 검증 스텝 때문에 반복당 토큰 수가 늘어도 처리량 이득이 상쇄될 수 있다. 저자들이 밝힌 정량적 한계 수치나 실패 조건은 원문에 별도로 제시되지 않았다.