PyTorch 2.14 MPS에서 F.linear의 [B, 1, K] 디코드 입력이 달라진 점 | DAKER 커뮤니티

오토리그레시브 디코드에서는 선형층에 한 토큰씩 입력하는 경우가 많습니다. 이때 입력 모양이 [B, 1, K]라면, 같은 선형층이라도 디바이스와 커널 경로에 따라 체감 성능이 크게 달라질 수 있습니다.

이번 글은 PyTorch 2.14의 MPS 경로에서 torch.nn.functional.linear에 시퀀스 길이 1인 활성화 [B, 1, K]를 넣었을 때 무엇이 달라졌는지 짚습니다. 핵심은 디코드용 빠른 경로와 GEMV 커널로 이어지면서, 예전처럼 느린 폴백에 떨어지는 경우를 줄인다는 점입니다.

PyTorch 2.14의 MPS 경로에서 F.linear에 [B, 1, K]를 넣으면 디코드용 빠른 경로와 GEMV 커널로 이어져, 느린 폴백을 피하는 데 도움이 됩니다.

왜 시퀀스 길이 1 입력이 중요했나

오토리그레시브 디코드는 한 토큰씩 [B, 1, K] 활성화를 선형층에 넣는 방식으로 동작합니다. Apple Silicon MPS에서는 이 모양이 빠른 경로를 벗어나 bf16·fp16에서 크게 느려질 수 있었습니다.

그래서 단순히 선형층 자체를 보는 것보다, 같은 가중치와 같은 dtype에서 시퀀스 길이 1 입력이 어떤 경로를 타는지 확인하는 것이 더 중요했습니다. 체감상 빨라졌는지를 보는 것보다, 같은 조건에서 시퀀스 길이 1이 폴백을 피하는지 기록하는 편이 재현에 도움이 됩니다.

재현해 볼 실험 조건

MPS 디바이스에서 작은 선형층에 [B, 1, K]와 [B, T, K](T가 1보다 큼)를 각각 넣고 지연을 비교해 보면 됩니다. 참가자 노트에는 배치, 은닉 크기, dtype, 시퀀스 길이, 측정 시간을 적어 두는 것이 좋습니다.

MPS가 있는 Mac에서 작은 F.linear 가중치를 만들고, 입력 dtype을 bf16 또는 fp16으로 맞춥니다.
입력 모양 [B, 1, K]와 [B, 8, K]처럼 시퀀스 길이가 다른 경우를 같은 가중치로 각각 측정합니다.
시퀀스 길이 1에서만 느렸던 이전 경험과 비교해, 이번 버전에서 지연·커널 경로가 어떻게 달라졌는지 한 줄을 남깁니다.

실험할 때 함께 봐야 할 범위

MPS 네이티브 선형대수(SVD·QR 등)와 이번 글의 선형층 디코드 경로는 다른 주제입니다. 이번 실습은 F.linear의 시퀀스 길이 1 빠른 경로만 봅니다.

실험이 끝나면 사용한 PyTorch 버전·Mac 칩·dtype·모양을 한 줄로 정리해 두면, 다음 참가자가 같은 조건을 바로 따라오기 좋습니다.

실무에서 바로 볼 포인트

디코드 루프에서 입력이 실제로 [B, 1, K]인지 먼저 확인합니다.
시퀀스 길이 1과 긴 시퀀스의 지연, 그리고 dtype을 함께 기록해 두는 것이 좋습니다.
SVD·QR 실습과 섞지 않고, 이번 글에서는 선형층 디코드 경로만 따로 보는 편이 좋습니다.

참고 자료

https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/community/mps-native-svd-qr-cholesky-pytorch-apple-silicon
https://daker.ai/community

같은 조건에서 [B, 1, K] 입력의 지연 변화를 직접 확인해 보셨다면 어떤 차이가 있었는지 궁금합니다.