FlashAttention을 켰는데도 32K 컨텍스트 prefill이 512 토큰보다 수백 배 느린 현상을 처음 마주치면 당황스럽다. "FlashAttention이 메모리 문제를 해결했다고 들었는데"라는 반응이 자연스럽다. 하지만 FlashAttention이 줄이는 것은 HBM 접근량이고, 줄이지 못하는 것은 행렬 내적 횟수다. 이 두 가지가 독립적인 병목이기 때문에, IO 최적화만으로는 긴 컨텍스트 prefill의 속도 문제를 완전히 해소할 수 없다.
O(n²)는 어디서 나오는가
Self-Attention의 비용은 FLOPs와 HBM 접근량, 두 축으로 분리된다.
시퀀스 길이 n, 헤드 차원 d, 헤드 수 h를 기준으로 QKᵀ 행렬곱에 2n²d FLOPs가 필요하다. Softmax 이후 AV 행렬곱에 또 2n²d. 총 4n²dh FLOPs로 n에 대해 O(n²)다.
HBM 접근량은 다른 계산이다. 표준 Attention 구현은 S = QKᵀ 결과를 n×n 행렬째로 HBM에 쓴 다음 softmax를 위해 다시 읽는다. FP16 기준 읽기·쓰기 합산하면 4n²h bytes다. 큰 n에서 Q, K, V 적재량(n에 선형)을 압도하는 dominant term이다.
두 병목을 하나의 지표로 합치면 Arithmetic Intensity(AI)가 된다:
AI = FLOPs / bytes = 2n²d / 4n² = d/2 (FP16, QKᵀ 커널 기준)
d=128이면 AI=64 ops/byte다. A100 SXM의 roofline은 312 TFLOPS ÷ 2000 GB/s = 156 ops/byte. 64 < 156이므로 QKᵀ 단계는 IO-bound다. 표준 Attention 구현이 HBM 읽기·쓰기에 묶여 있다는 뜻이고, 이 병목을 겨냥한 것이 FlashAttention이다.
FlashAttention이 줄이는 것과 줄이지 못하는 것
FlashAttention-2(Dao et al., 2023)는 S 행렬을 HBM에 내보내지 않는다. SRAM 크기에 맞는 블록으로 쪼개 softmax까지 on-chip에서 끝내고 최종 출력 O만 HBM에 저장한다. 결과적으로 HBM 읽기·쓰기가 O(n²)에서 O(n·d)로, n에 대해 선형으로 줄어든다.
FLOPs는 동일하다. QKᵀ를 블록으로 나눠 계산해도 전체 내적 횟수는 바뀌지 않는다.
IO 병목이 해소된 이후 GPU를 막는 것은 연산량 자체다. FlashAttention이 적용된 상태에서의 실효 AI는 올라가는데, 그러면 GPU가 compute-bound에 진입하는 시점이 생긴다:
prefill compute time ≈ 4n²dh / peak_TFLOPS
prefill IO time ≈ n × d × h × 4 bytes / bandwidth (Q,K,V + O 적재·저장)
n이 커질수록 compute time이 n²로 커지는 반면 IO time은 n에 선형이다. 어느 시점을 넘으면 compute time이 지배적이 되고, 그 구간에서는 FlashAttention이 IO를 아무리 줄여도 FLOPs 자체가 병목이 된다.
A100 80GB, Llama-2-7B, 배치 크기=1 조건에서 측정한 수치는 다음과 같다:
| 시퀀스 길이 | FlashAttn OFF | FlashAttn ON | 속도 향상 |
|---|---|---|---|
| 512 | ~28ms | ~18ms | 1.6× |
| 2,048 | ~200ms | ~100ms | 2.0× |
| 4,096 | ~720ms | ~300ms | 2.4× |
| 8,192 | ~2,700ms | ~980ms | 2.8× |
| 16,384 | ~10,800ms | ~3,100ms | 3.5× |
| 32,768 | ~43,000ms | ~11,500ms | 3.7× |
FlashAttention ON 상태에서도 512→32768 구간은 약 640배 느려진다. 이론적 n² 증가(64배)와의 괴리는 16K 이상 구간에서 GPU가 compute-bound로 전환되고, 여기에 KV cache 메모리 관리 오버헤드가 가산된 결과다. 속도 향상 비율(1.6×→3.7×)이 시퀀스가 길수록 커 보이는 것도 이 때문이다 — OFF 상태에서 IO 낭비가 그만큼 심각해지는 것이지, FlashAttention이 긴 시퀀스에서 더 효과적인 게 아니다.
Chunked Prefill이 실제로 하는 일
Sarathi-Serve(Agrawal et al., 2024, USENIX OSDI)가 제안한 Chunked Prefill은 O(n²) 총 연산량을 바꾸지 않는다. 긴 prefill 요청을 청크 크기 c로 나눠 n/c 스텝에 걸쳐 처리하면 스텝당 연산은 O(c·n)으로 작아지지만 총합은 여전히 O(n²)다.
이 방법의 효과는 스케줄링에 있다. 표준 스케줄러에서는 prefill이 끝날 때까지 decode 요청이 기다린다. 32K prefill이 11초 걸리는 동안 decode 큐 전체가 멈춘다. Chunked Prefill은 청크 하나가 끝날 때마다 decode 요청을 배치에 끼워 넣는다. GPU 활용률이 오르고 decode 대기가 줄어드는 구조다.
vLLM에서는 다음과 같이 켠다:
vllm serve meta-llama/Llama-2-7b-hf \
--enable-chunked-prefill \
--max-num-batched-tokens 512
기본 청크 크기는 512 토큰이다. vLLM V1에서는 가능한 경우 기본 활성화된다. 청크가 작을수록 decode 대기가 줄어들어 TPOT(time per output token)가 개선되지만, 청크 경계마다 KV cache 부분 적재와 스케줄러 오버헤드가 쌓여 단일 요청의 TTFT p50은 오를 수 있다. Sarathi-Serve 논문에서 청크 크기 256~512 구간이 TTFT p95와 throughput 사이 최적점으로 나타나며, Mistral-7B 기준 단일 A100에서 vLLM 대비 2.6배 serving capacity 향상을 보고한다.
"TTFT를 낮춘다"는 설명이 Chunked Prefill 관련 문서 곳곳에 있는데, 정확히는 p95 이상의 tail latency를 억제하는 것이다. throughput 향상 레버라는 점을 오해한 채 도입하면 p50 TTFT가 오히려 나빠졌다는 결과를 보게 된다.
Ring Attention은 O(n²)를 나눈다
단일 GPU에서 불가능한 128K, 1M 토큰 컨텍스트를 처리하려면 시퀀스 차원을 여러 GPU에 분산해야 한다. Ring Attention(Liu et al., 2023)은 k개 GPU에 시퀀스를 n/k씩 할당해 각 GPU의 attention 연산이 O((n/k)²)로 줄어들게 한다. GPU 2배면 단일 GPU 연산이 이론상 4분의 1이다.
단, 각 GPU는 자신의 Q에 대해 다른 GPU의 K, V 블록을 참조해야 하므로 cross-device 통신이 불가피하다. Llama-2-7B (d=128, h=32) 기준 n=32K, k=2이면 교환해야 할 KV 블록은 약 512MB다. NVLink(A100 양방향 ~600 GB/s)에서 교환 시간은 ~1.7ms, PCIe Gen4(단방향 32 GB/s)에서는 ~16ms다. 32K prefill 시간이 FlashAttention ON 상태에서 11.5초이므로 NVLink 환경에서는 통신이 연산에 완전히 묻히지만, PCIe 환경에서는 임계 시퀀스 길이가 64K128K 이상으로 올라가야 분산이 단일 GPU보다 유리해진다.
Ring Attention 논문은 통신과 연산을 pipeline overlap으로 숨길 수 있다고 서술하지만, 이 overlap의 실효성은 하드웨어 토폴로지와 배치 크기에 따라 달라진다. GPU를 늘린다고 prefill이 선형으로 빨라지지 않는 이유가 여기 있다.
세 전략을 하나의 프레임으로
| FlashAttention | Chunked Prefill | Ring Attention | |
|---|---|---|---|
| 건드리는 병목 | IO (HBM reads/writes) | 스케줄링 단위 | 단일 GPU 연산량 |
| FLOPs 변화 | 동일 | 동일 | GPU당 1/k² |
| 추가 비용 | 없음 | 청크 경계 오버헤드 | cross-device KV 통신 |
| TTFT 효과 | IO-bound 구간에서 직접 단축 | p95 억제, p50 소폭 상승 | NVLink 환경에서 직접 단축 |
| 주요 조건 | 단일 GPU, 범용 | 혼합 워크로드, tail SLO | 64K+ 컨텍스트, NVLink |
실무 순서는 대체로 이렇다. FlashAttention은 무조건 켠다 — IO-bound 구간에서 명확한 이득이 있고, compute-bound 구간에서도 손해는 없다. Chunked Prefill은 decode와 prefill이 섞이는 혼합 워크로드에서 TPOT 개선과 tail latency 억제를 위해 추가한다. Ring Attention은 단일 GPU 메모리·연산 한계를 벗어난 초장 컨텍스트, 그것도 NVLink 환경이 확보됐을 때 고려한다.
prefill 비용이 n²로 오르는 구조에서 컨텍스트 길이 상한선은 TTFT p95 SLO와 직접 연결된다. 평균 컨텍스트가 4K여도 p99 요청이 32K라면 그 한 요청이 진행되는 동안 decode 큐 전체가 지연을 받는다. Chunked Prefill이 이 충격을 완화하지만, max_model_len을 SLO 예산 기반으로 설정하는 게 선행이다.