Flash-dLLM, I/O 인식 KV 캐시로 확산형 LLM 추론 가속

확산형 대규모 언어 모델(dLLM)은 자기회귀 모델처럼 토큰을 하나씩 생성하는 대신 여러 위치를 병렬로 갱신할 수 있습니다. 하지만 이론적인 병렬성만으로 실제 추론 속도가 보장되지는 않습니다. 추론 과정에서 KV 캐시를 반복적으로 읽고 쓰면 연산량보다 GPU 메모리 I/O가 더 큰 병목이 될 수 있기 때문입니다.

Flash-dLLM은 KV 캐싱과 병렬 디코딩을 별개의 최적화 문제가 아니라 하나의 시스템 문제로 다룹니다. 논문은 기존 가속 기법이 캐시 재사용과 병렬 검증을 각각 연구하는 경우가 많아, 두 기능을 함께 사용할 때 발생하는 데이터 이동 비용을 충분히 고려하지 못했다고 설명합니다. 제안 프레임워크는 추가 학습 없이 적용되며 두 가지 핵심 구성으로 이뤄집니다.

핵심 내용

• Flash-Cache: QKV 프로젝션, RoPE 처리, 캐시 기록을 하나의 Triton 커널로 융합합니다. 중간 결과의 불필요한 메모리 이동을 줄이고, 배치마다 질의 길이가 다를 때는 블록 단위 스케줄링으로 처리합니다.
• 선택적 캐시 갱신: 매 단계 전체 캐시를 다시 쓰는 대신 새로 디코딩된 토큰과 주의도가 높은 일부 기존 디코딩 토큰을 중심으로 갱신합니다. 캐시의 유용한 정보를 유지하면서 쓰기 트래픽을 줄이려는 방식입니다.
• Flash-Verify: 별도의 보조 모델을 사용하지 않고 dLLM 자체가 초안 생성기와 검증기 역할을 모두 맡습니다. 두 가지 관점을 사용하는 인과 어텐션 마스크로 후보 생성과 검증을 통합하며, 제공된 자료에서는 단계별 수락 토큰 수가 대략 두 배가 된다고 설명합니다.

LLaDA-1.5 실험 결과로는 초당 148211토큰의 속도가 제시됐습니다. 캐시를 사용하지 않는 탐욕적 디코딩보다 22.3148.2배 빠르고, GSM8K와 HumanEval에서는 Elastic-Cache보다 각각 5.1배와 11.0배 빠른 것으로 보고됐습니다. 또한 Fast-dLLM보다 GPU 메모리 사용량이 약 48% 적으며, 배치 크기 32까지 확장할 수 있었다고 합니다.

다만 이 수치는 자료에 제시된 실험 조건의 결과이며 모든 환경에서 동일하게 재현된다는 의미는 아닙니다. 실제 성능은 모델 크기, 시퀀스 길이, GPU 종류, 배치 구성, 커널 구현에 따라 달라질 수 있습니다. 그럼에도 이 연구는 dLLM 서빙에서 계산 병렬성만 높이는 것보다 메모리 시스템과 캐시 접근을 함께 설계하는 일이 중요하다는 점을 보여줍니다.

Flash-dLLM의 의미는 캐시 최적화와 자기 검증형 디코딩을 단순히 병렬로 추가한 것이 아니라 하나의 추론 경로에 결합했다는 데 있습니다. 긴 시퀀스나 큰 배치를 처리할 때 불필요한 KV 캐시 이동을 줄이면, 추가적인 계산 병렬화보다 더 직접적인 효과를 낼 수 있습니다. 향후에는 더 다양한 모델과 하드웨어에서 보고된 속도 및 메모리 절감 효과를 검증해야 합니다.

출처: Hugging Face Daily Papers