SDPA rank-3 입력도 fused CUDA 경로 — PyTorch 2.14에서 배치 없는 텐서가 math 폴백을 피합니다 | DAKER 커뮤니티

한 줄 요약
PyTorch 2.14의 SDPA는 rank-3 입력도 fused CUDA 백엔드로 보냅니다. 배치 축을 빼거나 이미 펼쳐 둔 텐서가 math 폴백으로 조용히 떨어지던 함정이 줄어듭니다.

어텐션 입력을 [B, H, L, D]로만 가정하다 보면, 이미 flatten한 [H, L, D]나 배치 없는 텐서를 그대로 넘기기 쉽습니다. 오늘은 텐서 ndim이 3인지 4인지부터 로그로 남기십시오. rank-3인데도 fused가 도는지가 이번 변경의 핵심입니다.

재현해 볼 실험 조건

동일 내용을 rank-4와 rank-3 두 형태로 준비합니다.
둘 다 scaled_dot_product_attention에 넣고, 백엔드·커널 이름·지연 시간을 비교합니다.
reshape만으로 fused를 억지로 맞추던 기존 우회를 제거한 뒤 결과가 유지되는지 확인합니다.

배치 차원이 없거나 이미 평탄해진 입력이 math 경로로 가면, 코드는 그대로인데 속도만 빠집니다. 2.14는 그 입력을 fused CUDA로 보내 형태 변환 부담을 줄입니다. 참가자는 프로파일에 math 커널이 남는지만 먼저 보면 됩니다.

실무에서 바로 볼 포인트

헤드 차원·시퀀스 길이가 fused 조건을 벗어나면 여전히 다른 경로로 갈 수 있습니다. 형상 가드를 유지하십시오.
학습 루프와 추론 루프의 텐서 layout이 다르면 rank만 같아도 성능이 갈립니다. 두 경로를 같이 재현하십시오.
제출 노트에는 PyTorch 버전과 SDPA 입력 rank를 적어 두면 리뷰가 빨라집니다.

SDPA rank-3 입력도 fused CUDA 경로 — PyTorch 2.14에서 배치 없는 텐서가 math 폴백을 피합니다 | DAKER 커뮤니티

재현해 볼 실험 조건

실무에서 바로 볼 포인트

관련 DAKER 학습