SDPA rank-3 입력이 융합 커널을 건너뛰던 함정 — PyTorch 2.14에서 수정됐습니다 | DAKER 커뮤니티
한 줄 요약
배치 차원이 없거나 이미 펼쳐진 rank-3 텐서를 Scaled Dot-Product Attention에 넣으면, 예전에는 융합 CUDA 백엔드 대신 math 경로로 떨어질 수 있었습니다. PyTorch 2.14는 rank-3 입력도 융합 백엔드로 보냅니다.
근거는 공식 2.14 릴리스 블로그 Core Features의 「SDPA Fused Backends for Rank-3 Inputs」입니다. PR #192271이 언급되어 있습니다.
재현하기 쉬운 실험 조건
- 입력을
[B, H, L, D]가 아니라 배치를 접은[H, L, D]또는 동등한 rank-3로 넘깁니다. - 같은 연산량인데 프로파일에 fused SDPA 커널 대신 math/attention 폴백이 보입니다.
- 코드를 reshape만 바꿔 다시 재면 속도가 크게 달라집니다. 모델 구조가 아니라 커널 선택이 원인입니다.
블로그 표현 그대로, 「missing or flattened batch dimensions silently bypassed the fast fused kernels」가 실무에서 자주 나오는 함정입니다.
2.14 이후 기대 결과
- rank-3 입력도 융합 CUDA 백엔드로 디스패치됩니다.
- 호출부가 배치 차원을 다시 끼우지 않아도 fused 경로를 탈 수 있습니다.
- 벤치할 때는 PyTorch 버전, 입력 rank, dtype, 헤드 차원을 함께 기록하십시오. 2.13 이하와 2.14를 같은 입력으로 비교해야 원인 분리가 됩니다.
같이 보면 좋은 2.14 성능 항목
- Inductor
simple_overlap이 기본 켜져 분산 학습에서 통신·연산 겹침이 기본값이 됩니다. - NVGEMM이 Inductor GEMM 후보에 들어와 epilogue fusion·저정밀 경로를 확장합니다.
max_autotune과max_autotune_gemm_backends설정이 필요합니다.