PyTorch 2.14 MPS FlexAttention에 KV 배치 브로드캐스트가 들어갔습니다. 쿼리 배치가 달라도 같은 KV를 공유할 수 있어, 페이지드 어텐션으로 가기 위한 전제 조건 하나가 열립니다. 시퀀스 길이 1 F.linear 디코드 수정과는 다른 축입니다.

자세한 내용과 근거는 DAKER 원문에서 확인하실 수 있습니다.

핵심 요약

• 맥에서 어텐션 배치만 키우고 KV는 그대로 두고 싶을 때가 있습니다. 예전에는 쿼리 배치를 늘릴 때마다 KV를 복제해 맞추는 경우가 많았습니다. 오늘은 쿼리 배치와 KV 배치가 1:1로 묶여 있는지부터 보십시오. 공유 KV가 되면 불필요한 KV 복사를 줄일 여지가 생깁니다. 배치를 키운 뒤 메모리가 갑자기 불어났다면, 복

• MPS에서 FlexAttention 입력을 준비하되, 쿼리 배치만 늘리고 KV는 공유 형태로 둡니다.

• 동일 연산을 KV를 배치마다 복제한 경우와 브로드캐스트한 경우로 나눕니다.

원문 보기: https://daker.ai/community/pytorch-2-14-mps-flexattention-shared-kv-queries

DAKER: https://daker.ai/community/pytorch-2-14-mps-flexattention-shared-kv-queries