FA-3/4의 최적화 기술 중 일부가 RTX GPU로 전이될 수 있는지 궁금했습니다. vLLM/SGLang의 어텐션(attention)은 소비자용 카드에서는 FA-2로 폴백(fallback)되므로(FA-3 및 FA-4는 데이터센터 전용), 성능상 남겨진 이점이 있는지 알고 싶어 어텐션 커널(attention kernels)을 처음부터 다시 빌드했습니다. 해당 커널은 FA-2와 대등한 수준에 도달했습니다 (RTX 5090에서 batch=1, heads=8, seq_len=4096, head_dim=64 기준 206us). 하지만 불행히도 FA-3/4의 최적화 기술들은 소비자용 카드에서 적용이 불가능하거나 도움이 되지 않습니다. FA-2가 한계치(ceiling)인 것으로 보입니다.

요약하자면: FA-3의 주요 동력은 더 빠른 텐서 코어 명령어(WGMMA)이지만, 이는 RTX GPU에서는 사용할 수 없습니다. TMA (tensor memory accelerator)는 sm_120 (RTX 5XXX)에서 사용 가능합니다. 이론상으로는 도움이 되지만(LSU 감소), 전송(transport)이 병목(bottleneck)이 아니기 때문에 최종 수치는 거의 변하지 않습니다. 워프 전문화(Warp specialization) 또한 사용 가능합니다. 그러나 이는 주로 스케줄링 최적화, 즉 파이프라인 버블(pipeline bubbles)을 제거하고 텐서 코어를 더 잘 활용하도록 돕는 기술이지만, 비동기 텐서 코어 명령어(asynchronous tensor core instructions) 없이는 효과가 없습니다. 결과는 부정적이었습니다: 213us vs 206us. FA-4에서는 B200의 텐서 코어가 매우 빨라 전체 파이프라인이 SFU(special functions unit)에 의해 제한(bound)되기 때문에 FMA 명령어를 사용하여 exp를 시뮬레이션하기도 했습니다. RTX 5090은 텐서 코어 제한(tensor-core bound) 상태이므로 이 최적화 역시 의미가 없습니다. 실제로 소프트맥스(softmax)를 위해 expf 대신 더 빠른 exp2f를 사용하는 전통적인 최적화조차 수치를 변화시키지 못했습니다.

더 깊은 파이프라인(deeper pipeline)이나 레지스터 핑퐁(register ping-pong)과 같이 소비자용 실리콘에서 잠재적으로 작동할 수 있는 몇 가지 다른 최적화들도 시도해 보았으나, 성과가 없었습니다. 전체 파이프라인이 텐서 파이프 제한(tensor-pipe-bound) 상태이기 때문에, FA-2가 한계라고 믿으며, 모든 의미 있는 동력은 더 빠르고 낮은 정밀도의 텐서 코어를 활용하기 위해 어느 정도의 정확도를 희생해야 할 것입니다. 참고로 이것은 프리필(prefill) 및 연산 제한(compute-bound) 영역을 지배하는 일반적인 어텐션(regular attention)에 대한 탐구입니다.

대규모 KV 캐시(KV cache)를 대상으로 하는 디코딩(Decoding)은 다른 양상을 보이며, 이는 메모리 대역폭 제한(memory-bound)의 문제로, 위에서 언급한 그 어떤 것보다 split-KV/Flash-Decoding이 더 중요하게 작용합니다. 전체 기사: https://riftstack.ai/research/learning-flashattention-the-hard-way-part-2 Github: github.com/cloudrift-ai/emmy /u/NoVibeCoding 님이 r/LocalLLaMA에 제출함 [link] [comments]