Google TPU, DFlash로 추론 3배 가속

Google 개발자 블로그가 UCSD 연구팀의 DFlash를 소개했다. DFlash는 autoregressive draft 대신 block diffusion으로 후보 토큰 블록을 한 번의 forward pass에 생성해 speculative decoding의 serial bottleneck을 줄인다.

vLLM TPU inference에 통합한 뒤 TPU v5p에서 평균 3.13배의 tokens/sec 향상을 기록했고, 복잡한 수학 태스크에서는 최대 약 6배까지 빨라졌다. Llama-3.1-8B 기준 head-to-head에서는 DFlash가 2.29배의 end-to-end speedup을 내며 EAGLE-3의 1.30배를 앞섰다.

구현은 TPU/JAX 환경에 맞춰 재설계됐다.

• target model은 paged KV cache, draft model은 static on-device JAX arrays를 쓰는 dual-cache 구조를 적용
• context buffer는 power-of-2 padding으로 host-TPU 전송 오버헤드를 줄임
• accepted token count와 proposer state를 동기화해 diffusion draft와 target의 불일치를 제거

블로그는 TPU급 하드웨어에서 16토큰과 1024토큰 검증 비용이 거의 같다는 점을 들어, 병목이 verification보다 draft quality에 있다고 정리했다. 블록 크기 확대보다 acceptance probability 개선이 더 큰 성능 레버라는 해석도 제시했다.