🤖 AI

FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs

이 논문은 WGMMA 연산을 최적화하여 기존 방식 대비 상당한 속도 향상을 달성하는 동시에 높은 수치적 안정성을 유지하면서, NVIDIA H20 GPU에서 Multi-Head Latent Attention 추론을 크게 가속화하기 위해 효율적인 전치 어텐션 파이프라인(Efficient Transpose Attention Pipeline)을 활용하는 새로운 프레임워크인 FlashMLA-ETAP을 소개한다.

원저자: Pengcuo Dege, Qiuming Luo, Rui Mao, Chang Kong

게시일 2026-06-03

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Pengcuo Dege, Qiuming Luo, Rui Mao, Chang Kong

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 아주 긴 책(컨텍스트)을 읽고 단 하나의 짧은 질문(쿼리)에 답하려고 노력하고 있다고 상상해 보세요. 이것은 본질적으로 DeepSeek-R1과 같은 거대 AI 모델이 텍스트를 생성할 때 하는 일과 같습니다. 즉, 다음에 어떤 단어를 말할지 결정하기 위해 지금까지 읽은 모든 내용을 되돌아보는 것입니다.

이 논문은 특정 컴퓨터 칩인 NVIDIA H20을 위해, AI에게 그 책의 페이지를 더 똑똑하게 넘기는 방법을 가르치는 것과 같은 새로운 방법인 FlashMLA-ETAP를 소개합니다.

다음은 일상적인 비유를 사용한 문제와 해결책의 분석입니다:

문제점: "어색한 선반" 이슈

NVIDIA H20 GPU를 다음과 같은 규칙이 있는 도서관에서 일하는 사서라고 생각해 보세요: 선반은 안정성을 위해 최소 64권의 책 너비여야 합니다. 만약 선반에 16권의 책만 놓으려고 한다면, 사서는 선반이 무너지지 않도록 빈 공간을 가짜 책(패딩)으로 채워야 합니다. 이는 시간과 에너지를 낭비합니다.

AI 세계에서 "책"은 어텐션 헤드(서로 다른 것에 집중하는 뇌의 부분들)를 의미합니다. 8개의 H20 GPU가 있는 단일 서버에서 거대한 DeepSeek-R1 모델을 실행할 때, 작업은 분할됩니다. 각 GPU는 결국 16개의 헤드만을 처리하게 됩니다.

16은 요구되는 64보다 작기 때문에, H20 GPU는 하드웨어 규칙을 충족하기 위해 많은 "가짜 작업"(패딩)을 수행해야 합니다. 이는 마치 사서에게 16개의 진짜 상자를 운반하기 위해 64개의 빈 상자를 들고 다니라고 강요하는 것과 같습니다. 이로 인해 AI는 매우 긴 "책"(롱 컨텍스트)을 읽고 있지만 질문은 매우 짧은(한 번에 한 단어씩) 경우, 속도가 느려지고 비효율적이 됩니다.

해결책: 책을 옆으로 돌리기 (ETAP)

저자들은 ETAP(Efficient Transpose Attention Pipeline)이라고 불리는 기술을 만들었습니다. 16개의 헤드를 64 너비의 선반 규칙에 억지로 맞추려고 노력하는 대신, 그들은 문제를 옆으로 돌렸습니다.

짧은 행에 16개의 헤드를 배치하는 대신, 책의 길이(수천 페이지가 될 수 있는)를 주요 차원으로 보는 것입니다. 책이 매우 길기 때문에, 가짜 책 없이도 "64 너비의 선반" 요구 사항을 쉽게 채울 수 있습니다.

차원을 바꿈으로써—대화의 긴 기록을 주요 "너비"로 취급하고 짧은 질문을 "높이"로 취급함으로써—H20 GPU는 가짜 패딩으로 시간을 낭비하지 않고 풀 스피드로 작동할 수 있습니다. 이는 16개의 작은 아이템을 큰 상자에 넣으려고 하는 대신, 여유 공간이 많은 곳에 수직으로 쌓는 법을 깨닫는 것과 같습니다.

결과: 더 빠르고 더 정확하게

논문은 이 "옆으로 돌리는" 접근 방식이 H20 GPU에서 엄청난 차이를 만든다고 주장합니다:

훨씬 빠름: 긴 대화 길이(64,000 단어)에서 FlashMLA-ETAP는 이 특정 모델에 대한 이전의 최고 방법(FlashMLA)보다 2.78배 더 빠릅니다. 또한 FlashAttention-3 및 FlashInfer와 같은 다른 인기 있는 방법들보다 현저히 빠릅니다.
더 정확함: 보통 AI 계산 속도를 높이려고 하면 정밀도(숫자를 너무 공격적으로 반올림하는 것과 같은)를 약간 잃을 수 있습니다. 하지만 이 새로운 방법은 실제로 FlashAttention-3보다 더 정확하며, 훨씬 낮은 오차율(RMSE)을 보입니다. 이는 책을 더 빨리 읽으면서도 동시에 더 잘 이해하는 것과 같습니다.

이것이 중요한 이유

대부분의 AI 최적화는 초고가형 하이엔드 칩(H100과 같은)을 위해 설계되었습니다. H20은 "중급형" 칩으로, 성능은 좋지만 가장 강력하지는 않습니다. 이 논문은 적절한 소프트웨어 기술(ETAP)을 사용하면 특정 작업에 대해 중급형 칩의 성능을 훨씬 더 좋게 만들 수 있음을 보여줍니다. 이는 마치 기어를 바꾸는 방식을 변경함으로써 일반 세단이 특정 종류의 도로에서 스포츠카처럼 효율적으로 달리게 만드는 것과 같습니다.

요약하자면, FlashMLA-ETAP는 NVIDIA H20 GPU가 긴 AI 대화를 읽는 방식을 재구성하도록 알려주는 소프트웨어 업데이트이며, 이를 통해 낭비되는 노력을 제거하고 AI가 더 빠르고 정확하게 응답하도록 만듭니다.

기술 요약: FlashMLA-ETAP

문제 정의

DeepSeek-R1 671B와 같은 대규모 모델을 단일 인스턴스의 미드 티어(mid-tier) GPU 서버(예: NVIDIA H20)에 배포하는 것은 심각한 추론 병목 현상에 직면해 있습니다. H100/H800과 같은 고성능 GPU는 Tensor Memory Accelerator(TMA) 및 높은 연산 처리량과 같은 고급 기능을 활용하는 반면, NVIDIA H20은 148 TFLOPS라는 완만한 FP16 연산 능력으로 인해 제약을 받습니다.

핵심적인 비효율성은 다중 헤드 잠재 어텐션(Multi-Head Latent Attention, MLA) 추론의 디코딩 단계에서 발생합니다. 이 시나리오에서는 쿼리(Query) 시퀀스 길이는 일반적으로 짧은 반면(1~2 토큰), 키-값(Key-Value, KV) 컨텍스트 길이는 매우 깁니다(최대 64K 토큰). FlashMLA, FlashAttention-3, FlashInfer와 같이 고성능 하드웨어에 최적화된 기존 프레임워크들은 다음과 같은 이유로 H20에서 심각한 저활용 문제를 겪습니다:

WGMMA 패딩 오버헤드: Hopper 아키텍처의 WarpGroup Matrix-Multiply-Accumulate(WGMMA) 명령은 효율적인 실행을 위해 M-차원이 최소 64여야 합니다. 단일 인스턴스 배포 시, 128개의 모델 헤드는 8개의 GPU에 분산되어 각 GPU당 16개의 헤드만 할당됩니다. 이 per-GPU 헤드 수(M=16)는 WGMMA 최소 요구치를 충족하지 못하므로, 불필요한 패딩을 강제하여 연산 활용도를 25% 미만으로 떨어뜨립니다.
차원 불일치: 전통적인 어텐션 계산은 짧은 쿼리 길이를 M-차원에 정렬시키는데, 이는 패딩 문제를 악화시킵니다.
메모리 대역폭 압박: 긴 KV 컨텍스트는 메모리 액세스 지연 시간을 증가시키며, 이는 풍부한 자원을 갖춘 고성능 GPU에 비해 미드 티어 GPU에서 특히 치명적입니다.

방법론: 효율적인 전치 어텐션 파이프라인 (Efficient Transpose Attention Pipeline, ETAP)

이러한 제약을 해결하기 위해 저자들은 **Efficient Transpose Attention Pipeline (ETAP)**을 중심으로 하는 FlashMLA-ETAP 프레임워크를 제안합니다. 이 새로운 계산 모드는 H20의 하드웨어 제약에 맞게 표준 어텐션 메커니즘을 재구성합니다.

핵심 메커니즘

표준적인 어텐션 흐름( $S = Q \cdot K^T$ , $P = \text{softmax}(S)$ , $O = P \cdot V$ ) 대신, ETAP은 전략적인 전치(transposition)를 도입합니다:

전치 계산: 파이프라인은 $S^T = K \cdot Q^T$ 를 계산한 다음, $P^T = \text{softmax}(S^T)$ 를 수행하고, 최종적으로 $O = (V^T \cdot P^T)^T$ 를 계산합니다.
차원 재정렬: 연산을 전치함으로써, 자연스럽게 큰 값(예: 16K–64K)을 가지는 KV 컨텍스트 길이가 WGMMA 연산의 M-차원이 됩니다. 짧은 쿼리 길이는 N-차원이 됩니다.
패딩 제거: KV 컨텍스 길이가 크기 때문에, WGMMA 최소 차원 요구사항(≥64)을 자연스럽게 충족하여 쿼리 차원에서의 불필요한 패딩을 제거합니다. 이는 불필요한 연산의 비율을 크게 줄여줍니다.

구현 세부 사항

통합: ETAP은 KV 캐시 크기를 최소화하기 위해 저차원 결합 압축(low-rank joint compression)을 사용하는 FlashMLA에 통합되었습니다. 구현은 프로듀서-컨슈머 동기화를 사용하는 협동 스레드 배열(Cooperative Thread Array, CTA) 내에서 작동합니다.
파이프라인 구조: 알고리즘은 두 개의 워프그룹(warpgroup)을 사용합니다:
- 컨슈머 워프그룹(Consumer Warpgroup): 전치된 어텐션 점수( $K \cdot Q^T$ )를 계산하고 소프트맥스 통계량 및 출력 세그먼트를 업데이트합니다.
- 프로듀서 워프그룹(Producer Warpgroup): 고대역폭 메모리(HBM)에서 공유 메모리(SMEM)로의 비동기 데이터 로딩을 처리하고 두 번째 출력 세그먼트를 계산합니다.
일반화 가능성: 저자들은 ETAP의 전치 전략이 FlashAttention-3 및 FlashInfer와 같은 다른 프레임워크에 원활하게 통합될 수 있음을 보여주는 이론적 분석을 제공합니다. 이 통합은 행렬 곱셈 및 소프트맥스 단계를 전치된 흐름을 채택하도록 수정해야 하지만, 해당 프레임워크들의 핵심적인 I/O 인식 타일링(I/O-aware tiling) 및 커널 퓨전(kernel fusion) 이점은 그대로 유지합니다.

주요 기여

새로운 계산 모드: 컴퓨팅 자원이 제한된 미드 티어 GPU, 특히 H20을 위해 전치를 활용하여 WGMMA 패딩 오버헤드를 완화하는 ETAP을 도입했습니다.
FlashMLA-ETAP 프레임워크: H20을 위해 설계되었으며 다른 어텐션 프레임워크에도 쉽게 적응할 수 있도록 제작된 실용적인 FlashMLA 구현체입니다.
이론적 분석: ETAP의 전치 전략이 특정 프레임워크에 국한된 것이 아니라, 디코딩 시나리오(짧은 쿼리, 긴 컨텍스트)의 보편적인 과제를 해결한다는 것을 입증하는 상세한 분석을 제공합니다.

실험 결과

실험은 NVIDIA H20 GPU(96GB HBM3, 148 TFLOPS FP16)에서 DeepSeek-R1 모델을 사용하여 배치 크기 16 및 32, 시퀀스 길이 512에서 64K 범위에 대해 수행되었습니다.

처리량 개선:
- 시퀀스 길이 64K (배치 크기 16)에서, FlashMLA-ETAP은 89 TFLOPS/s를 달성하여 베이스라인인 FlashMLA(32 TFLOPS/s) 대비 2.78배의 속도 향상을 보였습니다.
- 동일한 구성에서 FlashAttention-3(17 TFLOPS/s) 및 FlashInfer(18 TFLOPS/s)와 비교했을 때, FlashMLA-ETAP은 각각 5.24배 및 4.94배의 성능 향상을 나타냈습니다.
- 성능 이득은 시퀀스 길이가 길어질수록 증가했으며, 이는 KV 컨텍스트가 성장함에 따라 패딩 오버헤드를 줄이는 데 있어 ETAP의 효과를 입증합니다.
수치적 안정성:
- FlashMLA-ETAP은 배정밀도(double-precision) 참조 모델 대비 $1.25 \times 10^{-5}$ 의 제곱평균제곱근 오차(RMSE)를 달축성했습니다.
- 이는 FlashAttention-3( $1.9 \times 10^{-4}$ )와 비교하여 오차가 15.2배 감소한 것으로, 전치 및 저차원 압축과의 통합이 높은 수치적 정확도를 유지함을 확인시켜 줍니다.

의의 및 주장

본 논문은 FlashMLA-ETAP이 자원이 제한된 환경에서의 추론에 대한 중요한 격차를 해소하며, H20과 같은 미드 티어 GPU를 위한 확장 가능한 솔루션을 제공한다고 주장합니다. 어텐션 계산을 하드웨어 특유의 제약(WGMMA 차원)에 맞게 재지향함으로써, 이 프레임워크는 이전 프레임워크들이 심각한 저활용을 겪었던 단일 인스턴스 서버에서 대규모 모델(예: DeepSeek-R1)의 효율적인 배포를 가능하게 합니다.

저자들은 이 연구를 효율적인 추론을 민주화하기 위한 단계로 규정하며, 고성능 GPU 클러스터에 접근할 수 없는 환경에서도 고급 모델의 폭넓은 채택을 가능하게 한다고 강조합니다. 또한, 이 설계가 일반화 가능하다는 점을 들어, H20을 넘어 더 넓은 범위의 GPU 아키텍처에서 하드웨어 인식 최적화를 위한 길을 열어줄 수 있다고 설명합니다.

저자가 언급한 한계점:
평가는 현재 H20 GPU, 자기회귀 디코딩(순방향 패스당 하나의 토큰), 그리고 특정 구성(16개 헤드, 576 헤드 차원, 최대 64K 시퀀스)으로 제한되어 있습니다. 저자들은 다양한 GPU, 추론 패러다임 또는 더 큰 모델에 대한 일반성을 검증하기 위해서는 추가적인 조사가 필요함을 인정합니다. 향-후 과제로 FlashAttention-3 및 FlashInfer로의 ETAP 통합을 경험적으로 검증하고, 다른 Hopper 기반 미드 티어 GPU를 위해 최적화하는 것을 제시했습니다.