PyTorch 2.14 SDPA, rank-3 입력도 fused CUDA 경로를 타는 변화 | DAKER 커뮤니티
어텐션 코드를 다루다 보면 입력 텐서를 늘 [B, H, L, D] 형태로만 생각하기 쉽습니다. 하지만 실제 실험이나 추론 코드에서는 이미 펼쳐 둔 [H, L, D] 형태나 배치 축이 없는 텐서를 그대로 넘기는 경우도 자주 생깁니다. 이때 겉으로는 같은 연산처럼 보여도, 내부에서는 조용히 math 폴백으로 빠져 성능만 떨어지는 일이 있었습니다.
PyTorch 2.14에서는 이 지점이 달라집니다. SDPA가 rank-3 입력도 fused CUDA 백엔드로 보낼 수 있게 되면서, reshape로 우회하던 부담이 줄고 배치 없는 입력에서도 기대한 경로를 타는지 확인하기가 한결 쉬워졌습니다.
PyTorch 2.14의 SDPA는 rank-3 입력도 fused CUDA 백엔드로 보냅니다.
무엇이 바뀌었나
기존에는 어텐션 입력이 rank-4가 아닐 때, 특히 이미 flatten한 [H, L, D]나 배치 없는 텐서를 그대로 넘기면 math 경로로 떨어지는 경우가 문제였습니다. 코드 자체는 정상 동작해도 성능 저하 원인을 바로 알아차리기 어렵다는 점이 함정이었습니다.
이번 변화의 핵심은 텐서의 ndim이 3이더라도 fused CUDA 경로를 탈 수 있다는 점입니다. 따라서 배치 차원이 없거나 이미 평탄해진 입력을 다룰 때도, 불필요한 형태 변환 없이 같은 연산 경로를 기대해 볼 수 있습니다.
재현해 볼 실험 조건
이 변화는 간단한 비교 실험으로 확인할 수 있습니다. 동일한 내용을 rank-4와 rank-3 두 형태로 준비한 뒤, 둘 다 scaled_dot_product_attention에 넣고 백엔드, 커널 이름, 지연 시간을 비교하면 됩니다. 이어서 reshape만으로 fused를 억지로 맞추던 기존 우회를 제거한 뒤에도 결과가 유지되는지 보면 변화가 더 분명해집니다.
특히 먼저 볼 것은 프로파일에 math 커널이 남는지 여부입니다. 배치 차원이 없거나 이미 평탄해진 입력이 여전히 math 경로로 가는지, 아니면 fused CUDA로 처리되는지를 확인하면 핵심을 빠르게 짚을 수 있습니다.
rank-3인데도 fused가 도는지가 이번 변경의 핵심입니다.
실무에서 바로 볼 포인트
이번 변화가 모든 경우를 자동으로 해결하는 것은 아닙니다. 헤드 차원이나 시퀀스 길이가 fused 조건을 벗어나면 여전히 다른 경로로 갈 수 있으므로, 형상 가드는 유지하는 것이 좋습니다.
또한 학습 루프와 추론 루프의 텐서 layout이 다르면 rank만 같아도 성능 차이가 날 수 있습니다. 따라서 두 경로를 함께 재현해 보는 편이 안전합니다. 실험 기록이나 제출 노트에는 PyTorch 버전과 SDPA 입력 rank를 함께 적어 두면 리뷰도 더 빨라집니다.
관련 DAKER 학습
참고 자료
https://daker.ai/public/learning/materials/pytorch-advanced-2-embedding-self-attention-mini-i
https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
실제로는 rank-3 입력에서 어떤 프로파일 차이를 보셨는지 함께 이야기해 주셔도 좋겠습니다.