PyTorch 2.14에서 달라진 SDPA: rank-3 입력도 이제 융합 커널을 탑니다 | DAKER 커뮤니티

Scaled Dot-Product Attention 성능을 점검할 때, 코드가 크게 다르지 않은데도 예상보다 느리게 나오는 경우가 있습니다. 특히 배치 차원이 없거나 이미 펼쳐진 rank-3 텐서를 넣었을 때라면, 연산 자체보다 커널 선택이 원인일 수 있습니다.

PyTorch 2.14는 이 지점을 바로잡았습니다. 이전에는 rank-3 입력이 조용히 빠른 fused CUDA 백엔드를 건너뛰고 math 경로로 떨어질 수 있었지만, 이제는 rank-3 입력도 융합 백엔드로 디스패치됩니다. 작은 reshape 차이가 성능 차이로 이어지던 함정을 이해해 두면, 벤치마크와 실서비스 모두에서 원인 파악이 훨씬 쉬워집니다.

rank-3 입력에서 생기던 문제

공식 PyTorch 2.14 Release Blog의 Core Features 섹션에는 「SDPA Fused Backends for Rank-3 Inputs」가 포함되어 있습니다. 여기서는 PR #192271과 함께, 배치 차원이 없거나 펼쳐진 입력이 빠른 fused 커널을 우회하던 문제가 언급됩니다.

missing or flattened batch dimensions silently bypassed the fast fused kernels

실무에서는 입력을 항상 [B, H, L, D] 형태로 유지하지 않는 경우가 적지 않습니다. 배치를 접어서 [H, L, D]처럼 넘기거나, 동등한 rank-3 텐서로 처리하는 흐름도 흔합니다. 문제는 같은 연산량처럼 보여도 이 형태 차이 때문에 fused SDPA 커널 대신 math 또는 attention 폴백이 선택될 수 있었다는 점입니다.

그래서 모델 구조를 바꾸지 않았는데도, reshape 방식만 달리해 다시 측정하면 속도가 크게 달라지는 일이 생겼습니다. 이 경우 원인은 모델이 아니라 커널 선택에 있다고 보는 것이 맞습니다.

재현할 때 확인할 지점

이 문제는 비교적 단순한 조건에서 드러납니다. 입력을 [B, H, L, D]가 아니라 배치를 접은 [H, L, D] 또는 동등한 rank-3 형태로 넘기고, 프로파일에서 fused SDPA 커널 대신 math 경로가 보이는지 확인하면 됩니다.

같은 연산량인데도 프로파일 결과가 다르게 나타난다면, 입력 rank가 백엔드 선택에 영향을 준 상황일 가능성이 큽니다. 이때는 코드 전체를 의심하기보다, 텐서 형태와 디스패치 경로를 먼저 보는 것이 좋습니다.

PyTorch 2.14에서 달라진 점

PyTorch 2.14 이후에는 rank-3 입력도 융합 CUDA 백엔드로 디스패치됩니다. 즉, 호출부에서 배치 차원을 다시 끼워 넣지 않아도 fused 경로를 탈 수 있게 된 것입니다.

PyTorch 2.14는 rank-3 입력도 fused SDPA 경로로 보냅니다.

이 변화는 코드 가독성이나 텐서 정리 방식 때문에 rank-3 입력을 쓰던 경우에 특히 의미가 있습니다. 예전에는 같은 계산을 하더라도 입력 모양 때문에 성능상 불이익을 받을 수 있었지만, 2.14에서는 그 차이가 줄어듭니다.

벤치마크할 때 함께 기록하면 좋은 항목

이 이슈를 정확히 비교하려면 PyTorch 버전과 입력 조건을 함께 남기는 것이 중요합니다. 특히 2.13 이하와 2.14를 같은 입력으로 비교해야 원인 분리가 됩니다.

함께 기록할 항목으로는 PyTorch 버전, 입력 rank, dtype, 헤드 차원이 있습니다. 같은 모델이라도 이 조건에 따라 실제 선택되는 커널이 달라질 수 있기 때문입니다.

같이 볼 만한 PyTorch 2.14 성능 변화

이번 릴리스에는 SDPA 외에도 성능과 관련된 항목이 있습니다. 예를 들어 Inductor simple_overlap이 기본으로 켜져 분산 학습에서 통신과 연산의 겹침이 기본값이 됩니다.

또한 NVGEMM이 Inductor GEMM 후보에 포함되어 epilogue fusion과 저정밀 경로가 확장됩니다. 이 부분은 max_autotune과 max_autotune_gemm_backends 설정이 필요합니다.

참고 자료

PyTorch 2.14 Release Blog

실제로 rank-3 입력 때문에 SDPA 성능 차이를 겪은 적이 있는지, 있다면 어떤 형태의 입력에서 확인됐는지 궁금합니다.