PyTorch 2.14 MPS FlexAttention에서 KV를 공유해 쿼리 배치를 늘릴 때 볼 점 | DAKER 커뮤니티
맥에서 어텐션 배치를 키우려 할 때 가장 먼저 부딪히는 문제 중 하나는 메모리입니다. 쿼리 배치만 늘리고 싶은데 KV까지 함께 복제되면, 기대보다 메모리 사용량이 빠르게 커질 수 있습니다.
PyTorch 2.14의 MPS FlexAttention에 추가된 KV 배치 브로드캐스트는 바로 이 지점을 다시 보게 합니다. 같은 KV를 여러 배치가 함께 쓸 수 있다면, 쿼리 배치를 늘리면서도 불필요한 복사를 줄일 수 있기 때문입니다.
PyTorch 2.14의 MPS FlexAttention에 KV 배치 브로드캐스트가 추가됐습니다.
브로드캐스트는 같은 데이터를 복제하지 않고 여러 배치가 함께 쓰는 방식입니다. 이제 쿼리 배치가 달라도 같은 KV를 공유할 수 있습니다. 이는 페이지드 어텐션으로 가기 위한 중요한 조건 중 하나입니다. 이 내용은 시퀀스 길이 1에서의 F.linear 디코드 수정과는 다른 주제입니다.
왜 KV 공유 여부를 먼저 봐야 하는가
어텐션 배치만 늘리고 KV는 그대로 유지하고 싶을 때가 있습니다. 예전에는 쿼리 배치를 늘릴 때마다 KV도 함께 복제해 크기를 맞추는 경우가 많았습니다. 이럴 때는 먼저 쿼리 배치와 KV 배치가 꼭 1:1로 묶여 있는지 확인하면 됩니다.
같은 KV를 공유할 수 있다면, 불필요한 KV 복사를 줄일 수 있습니다. 배치를 키운 뒤 메모리가 갑자기 늘었다면, KV를 복제한 것인지 공유한 것인지부터 구분해 보는 것이 좋습니다.
직접 확인해 볼 실험 조건
MPS에서 FlexAttention 입력을 준비하되, 쿼리 배치만 늘리고 KV는 공유 형태로 둡니다. 그리고 같은 연산을 두 경우로 나눠 비교하면 됩니다. 하나는 KV를 배치마다 복제한 경우이고, 다른 하나는 브로드캐스트로 공유한 경우입니다.
이때 프리필 구간의 지연 시간과 메모리를 확인하고, 보조 통계인 log-sum-exp와 max 경로가 깨지지 않는지도 함께 보면 됩니다.
같은 KV를 공유할 수 있다면, 쿼리 배치를 늘리면서도 불필요한 KV 복사를 줄일 수 있습니다.
PyTorch 2.14에서 함께 봐야 할 점
PyTorch 2.14는 KV 배치 브로드캐스트뿐 아니라 보조 텐서와 동적 크기 캡처도 보강합니다. 동적 크기 캡처는 입력 크기가 바뀌어도 실행 경로를 더 유연하게 다루는 기능입니다.
공식 정리에는 특정 macOS·커널 조합에서 프리필 가속이 보고된 바 있습니다. 다만 그 수치는 저자의 벤치마크와 하드웨어 조건을 바탕으로 한 결과이므로, 기기마다 다시 측정해야 합니다. 숫자를 그대로 믿기보다 같은 입력으로 로컬 환경에서 한 번 더 재보는 편이 안전합니다.
실무에서 먼저 정리할 포인트
우선 디코드 선형층 이슈와 혼동하지 않는 것이 좋습니다. 그쪽은 시퀀스 길이 1에서의 F.linear 경로이고, 이번 글은 FlexAttention에서 KV를 공유하는 이야기입니다.
페이지드 어텐션을 목표로 한다면, 먼저 KV 공유가 실제로 되는지 단위 테스트로 확인해 두는 것이 좋습니다. 또한 배포 노트에는 PyTorch·macOS 버전을 함께 적는 편이 좋습니다. 커널 전제가 바뀌면 체감 성능도 달라질 수 있습니다.
어텐션 학습 자료와 배포 성능 체크를 함께 보면, 배치 설계와 측정 포인트를 한 흐름으로 정리할 수 있습니다.
함께 보면 좋은 DAKER 학습 자료
참고 자료
https://daker.ai/public/learning/materials/pytorch-advanced-2-embedding-self-attention-mini-i
https://daker.ai/public/learning/materials/pytorch-advanced-deployment-torchscript-onnx-perfo
https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
실제 환경에서 쿼리 배치와 KV 공유를 분리해 측정해 본 경험이 있다면, 어떤 지점에서 차이가 가장 크게 느껴졌는지 궁금합니다.