How Deep Learning Computes

nano-vLLM Attention Kernel Swap

마지막 수정:

inferencenano-vllmattentionkernel

Attention kernel swap은 nano-vLLM 최적화의 중심 실험이다.

비교 대상은 세 가지다.

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

반드시 prefill과 decode를 나눠 측정한다.

prefill:
prompt 전체 attention, TTFT에 영향

decode:
KV cache read 중심, TPOT에 영향

이 실험은 vLLM의 PagedAttention을 이해하기 위한 전 단계다. 작은 엔진에서 attention 교체 지점을 경험해야 production vLLM backend를 읽을 때 책임 경계가 보인다.

확인

  • Prefill attention 최적화는 주로 어떤 metric에 영향을 주는가?
  • Decode attention 최적화는 주로 어떤 metric에 영향을 주는가?
  • SDPA가 빠르더라도 PagedAttention 문제가 따로 남는 이유는 무엇인가?
먼저 보면 좋은 카드nano-vLLM Kernel Integration Pointsnano-vLLM serving loop에서 RMSNorm, attention, sampling 같은 kernel 교체 지점을 찾는다.먼저 보면 좋은 카드Attention Workload: Prefill vs Decodeattention benchmark를 prefill-like와 decode-like workload로 나누어 측정한다.먼저 보면 좋은 카드Optimized Attention Benchmarkmanual attention, PyTorch SDPA, optimized attention backend를 prefill/decode workload에서 비교한다.다음에 보기 좋은 카드nano-vLLM Prefill/Decode Benchmark Reportkernel swap 전후의 TTFT, TPOT, throughput, memory를 prefill/decode workload별로 정리한다.같은 경로의 이전 카드nano-vLLM RMSNorm Kernel SwapPyTorch RMSNorm을 custom/optimized RMSNorm kernel로 교체하고 model output과 serving metric을 비교한다.같은 경로의 다음 카드nano-vLLM Sampling Kernel Candidatelogits에서 token을 고르는 sampling 경로가 언제 kernel 최적화 후보가 되는지 판단한다.