FlexAttention Part II — 추론·디코딩

FlexAttention Part II — 추론·디코딩

FlexAttention의 후속 블로그는 추론(decoding) 을 위한 전용 백엔드를 소개해요. GQA와 PageAttention을 지원하고, prefill과 decode를 구분해 최적화하는 방법을 다룹니다.

추론의 특수성

  • 디코딩은 토큰을 하나씩 생성하며 KV 캐시를 재활용해요. 매 반복 전체 어텐션을 다시 계산하지 않아요.
  • FlexAttention은 prefill-phase 커널(학습용)과 FlexDecoding 백엔드(추론용)를 자동으로 골라서 실행해요.

예시

# prefill: query shape [B, H, 8000, D]
flex_attention(q_prefill, k_cache, v_cache, ...)
# decoding: query shape [B, H, 1, D]
flex_attention(q_last_token, k_cache, v_cache, ...)  # FlexDecoding 커널
  • GQA(그룹드 쿼리 어텐션)와 PageAttention(고정 크기 페이지 KV 캐시)을 지원해요.
  • score_mod로 오프셋을 바꿔가며도 재컴파일 없이 다음 토큰을 처리할 수 있는 점이 장점이에요.

더 알아보기