Google, DFlash 추측적 디코딩으로 TPU 속도 향상 3배 강조

Google Cloud는 UCSD 연구진이 구축하고 vLLM TPU 추론 스택에 통합한 TPU v5p에서의 블록디퓨전 사전 디코딩 오픈소스 구현을 소개했다. 팀은 단일 순전파로 초안 토큰 블록을 생성하는 DFlash 접근 방식이 벤치마크 전반에서 평균 처리량을 3.13배 향상시켰고, 일부 수학 작업 부하에서는 거의 6배의 속도 향상을 달성했다고 밝혔다. 같은 TPU 하드웨어와 Llama3.18B 대상 모델에서 EAGLE3와 1대1로 비교했을 때, DFlash는 EAGLE3의 1.30배에 비해 전체 엔드투엔드 서비스 성능이 2.29배 향상되었다. 또한 이 글은 어텐션 캐싱, 컨텍스트 처리, 메타데이터 정렬의 변경을 포함해 이 방법을 TPU와 JAX에 맞게 조정하는 데 필요한 엔지니어링 작업도 설명했다. Google은 이 구현이 vLLM tpuinference 저장소에 제출되었으며, 향후 TPU 기반 LLM 서비스 연구의 방향을 정하는 데 도움이 될 수 있다고 말했다.