How Deep Learning Computes

Optimized Attention Benchmark

마지막 수정:

pytorchattentionoptimizationbenchmark

이 카드의 목표는 attention 최적화를 하나의 숫자로 끝내지 않는 것이다.

비교 대상은 세 가지다.

1. manual PyTorch attention
2. PyTorch SDPA
3. optimized attention backend

처음에는 optimized backend를 직접 CUDA로 만들 필요는 없다. PyTorch SDPA의 FlashAttention 계열 backend, torch.compile, 또는 별도 optimized implementation을 비교 대상으로 둘 수 있다.

비교표는 prefill-like와 decode-like를 나눠 작성한다.

workload | B | H | T | Dh | dtype | implementation | latency | memory
prefill  |   |   |   |    |       | manual         |         |
prefill  |   |   |   |    |       | SDPA           |         |
prefill  |   |   |   |    |       | optimized      |         |
decode   |   |   |   |    |       | manual         |         |
decode   |   |   |   |    |       | SDPA           |         |
decode   |   |   |   |    |       | optimized      |         |

이 실험은 나중에 custom inference kernel과 nano-vLLM으로 이어지는 다리 역할을 한다.

확인

  • Attention 최적화 benchmark를 prefill과 decode로 나눠야 하는 이유는 무엇인가?
  • optimized backend가 prefill에서는 빠르지만 decode에서는 덜 유리할 수 있는 이유는 무엇인가?
  • 이 비교가 nano-vLLM의 어떤 부분으로 이어지는가?
먼저 보면 좋은 카드Attention Workload: Prefill vs Decodeattention benchmark를 prefill-like와 decode-like workload로 나누어 측정한다.먼저 보면 좋은 카드FlashAttention IO AnalysisFlashAttention의 핵심 이득이 FLOP 감소가 아니라 HBM read/write 감소라는 점을 이해한다.먼저 보면 좋은 카드torch.compile과 Fused Op 비교eager 실행, torch.compile, fused operation 후보를 비교하며 PyTorch 레벨 최적화의 한계를 확인한다.다음에 보기 좋은 카드nano-vLLM Attention Kernel Swapmanual/SDPA/optimized attention을 nano-vLLM prefill/decode loop에 연결하고 전후 성능을 비교한다.다음에 보기 좋은 카드nano-vLLM Kernel Integration Pointsnano-vLLM serving loop에서 RMSNorm, attention, sampling 같은 kernel 교체 지점을 찾는다.다음에 보기 좋은 카드PyTorch Transformer Benchmark Report 작성하기RMSNorm과 attention profiling, 최적화 결과를 하나의 재현 가능한 실험 리포트로 정리한다.