MPS에서 F.linear에 [B, 1, K] 디코드 활성화를 넣으면 시퀀스길이 1 빠른 경로와 GEMV로 느린 폴백을 피합니다 | DAKER 커뮤니티

시퀀스길이 1 선형
설명용 생성 이미지입니다.

한 줄 요약
PyTorch 2.14의 MPS 경로에서 torch.nn.functional.linear에 시퀀스 길이 1인 활성화 [B, 1, K]를 넣으면, 디코드용 빠른 경로와 GEMV 커널로 이어집니다. 예전처럼 느린 폴백에 떨어지는 경우를 줄입니다.

오토리그레시브 디코드는 한 토큰씩 [B, 1, K] 활성화를 선형층에 넣습니다. Apple Silicon MPS에서는 이 모양이 빠른 경로를 벗어나 bf16·fp16에서 크게 느려질 수 있었습니다. 오늘은 MPS 디바이스에서 작은 선형층에 [B, 1, K]와 [B, T, K](T가 1보다 큼)를 각각 넣어 지연을 비교하십시오. 참가자 노트에는 배치, 은닉 크기, dtype, 시퀀스 길이, 측정 시간을 적습니다. “체감으로 빨라진 것 같다”보다 “같은 조건에서 시퀀스 길이 1이 폴백을 피하는지”가 재현에 도움이 됩니다.

재현해 볼 실험 조건

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

MPS 네이티브 선형대수(SVD·QR 등)와 이번 글의 선형층 디코드 경로는 다른 주제입니다. 이번 실습은 F.linear의 시퀀스 길이 1 빠른 경로만 봅니다. 실험이 끝나면 사용한 PyTorch 버전·Mac 칩·dtype·모양을 한 줄로 고정해 두면, 다음 참가자가 같은 조건을 바로 따라올 수 있습니다.

실무에서 바로 볼 포인트

관련 DAKER 학습

설명용 생성 이미지입니다. 동작은 로컬 MPS에서 F.linear 지연으로 확인하십시오.