MPS FlexAttention KV 배치 브로드캐스트 — PyTorch 2.14에서 공유 KV로 쿼리 배치를 묶습니다 | DAKER 커뮤니티
한 줄 요약
PyTorch 2.14 MPS FlexAttention에 KV 배치 브로드캐스트가 들어갔습니다. 쿼리 배치가 달라도 같은 KV를 공유할 수 있어, 페이지드 어텐션으로 가기 위한 전제 조건 하나가 열립니다. 시퀀스 길이 1 F.linear 디코드 수정과는 다른 축입니다.
맥에서 어텐션 배치만 키우고 KV는 그대로 두고 싶을 때가 있습니다. 예전에는 쿼리 배치를 늘릴 때마다 KV를 복제해 맞추는 경우가 많았습니다. 오늘은 쿼리 배치와 KV 배치가 1:1로 묶여 있는지부터 보십시오. 공유 KV가 되면 불필요한 KV 복사를 줄일 여지가 생깁니다. 배치를 키운 뒤 메모리가 갑자기 불어났다면, 복제 KV인지 공유 KV인지부터 가르십시오.
재현해 볼 실험 조건
- MPS에서 FlexAttention 입력을 준비하되, 쿼리 배치만 늘리고 KV는 공유 형태로 둡니다.
- 동일 연산을 KV를 배치마다 복제한 경우와 브로드캐스트한 경우로 나눕니다.
- 프리필 구간 지연·메모리와, 보조 통계(log-sum-exp·max) 경로가 깨지지 않는지 확인합니다.
2.14는 KV 배치 브로드캐스트와 보조 텐서·동적 크기 캡처를 보강합니다. 공식 정리에는 특정 macOS·커널 조합에서 프리필 가속이 보고된 바 있으나, 수치는 저자 벤치·하드웨어 전제라 기기마다 다시 재야 합니다. 숫자를 글에 박아 두지 말고, 같은 입력으로 로컬에서 한 번 더 재는 편이 안전합니다.
실무에서 바로 볼 포인트
- 디코드 선형층 함정과 혼동하지 마십시오. 그쪽은
F.linear시퀀스 길이 1 경로이고, 이번 글은 FlexAttention KV 공유입니다. - 페이지드 어텐션을 목표로 한다면, KV 공유가 되는지를 먼저 단위 테스트로 잠그십시오.
- 배포 노트에 PyTorch·macOS 버전을 함께 적으십시오. 커널 전제가 바뀌면 체감이 달라집니다.
- 어텐션 학습 자료와 배포 성능 체크를 같이 보면, 배치 설계와 측정 포인트가 한 흐름으로 정리됩니다.
관련 DAKER 학습
설명용 생성 이미지입니다. 배속·메모리는 사용 기기에서 다시 측정하십시오.