AI 추론 최적화
장문맥 AI의 병목, 어텐션과 GPU를 함께 설계해야 합니다
NVIDIA가 공개한 어텐션 공동설계 4가지 지침을 공식 기술 블로그 기준으로 정리했습니다.
이 글에서 다룰 내용
장문맥 병목, 그룹 크기, 헤드 차원, KV 상태, GPU 병렬화
장문맥 추론의 진짜 병목은 무엇인가
이 글에서 다룰 내용
장문맥의 계산 병목, 어텐션 설계 4가지 원칙, GPU 공동설계, AI 에이전트 적용 효과
장문맥과 AI 에이전트 워크로드가 늘수록 전체 추론 시간에서 어텐션이 차지하는 비중이 커집니다. NVIDIA가 공개한 DeepSeek-R1 프리필 분석에서는 문맥 길이가 4K에서 128K로 늘 때 어텐션 비중이 18%에서 85%로 증가했습니다.
NVIDIA는 프리필과 디코드를 서로 다른 문제로 구분합니다. 프리필은 전체 입력을 병렬 처리해 주로 연산 병목이 되고, 디코드는 토큰을 하나씩 생성하면서 KV 캐시를 읽기 때문에 주로 HBM 메모리 대역폭 병목이 됩니다.
따라서 장문맥 추론은 최대 컨텍스트 길이만 늘리는 문제가 아닙니다. 모델의 어텐션 구조와 GPU가 실제로 실행하는 방식을 함께 설계하는 공동설계가 필요합니다.
원칙 1: 디코드에 맞춰 그룹 크기를 높인다
첫 번째 지침은 쿼리 헤드가 공유하는 KV 헤드의 비율인 그룹 크기를 디코드 기준으로 선택하고 높이는 것입니다. NVIDIA 분석에서는 프리필이 그룹 크기에 상대적으로 민감하지 않지만, 디코드는 작은 행렬 연산과 KV 캐시 읽기의 영향을 크게 받습니다.
그룹 크기를 높이면 디코드에서 산술 집약도를 끌어올릴 수 있습니다. 다만 실제 선택은 모델 정확도와 GPU 커널 측정 결과를 함께 검증해야 하며, NVIDIA 발표만으로 모든 모델의 최적값이 같다고 볼 수는 없습니다.
원칙 2: 헤드 차원은 128 또는 256에 맞춘다
두 번째 지침은 헤드 차원을 128 또는 256으로 설계하는 것입니다. NVIDIA는 이 크기가 GPU 타일과 128바이트 전송에 맞고, 텐서 메모리 예산 안에서 효율을 높이는 선택이라고 설명합니다.
프리필에서는 더 큰 헤드가 소프트맥스 연산을 가리는 데도 도움이 될 수 있습니다. 이는 NVIDIA GPU와 해당 측정 조건에서 제시된 공동설계 지침이므로, 다른 하드웨어나 커널에서는 별도 벤치마크가 필요합니다.
원칙 3: 유효 KV 상태를 줄인다
세 번째 지침은 긴 문맥에서 읽어야 할 유효 KV 상태를 줄이는 것입니다. NVIDIA는 KV 캐시 압축, 희소 어텐션, 슬라이딩 윈도 어텐션, 하이브리드 모델을 구체적인 선택지로 제시합니다.
이는 문맥을 무조건 짧게 자르라는 뜻이 아닙니다. 실제 서비스에서는 장거리 정보 보존과 처리량을 함께 측정해, 속도를 얻는 대신 필요한 근거를 잃지 않는지 확인해야 합니다.
원칙 4: 병렬화 규모를 KV 헤드 수에 맞춘다
네 번째 지침은 텐서 병렬화 규모를 KV 헤드 수 이하로 유지하는 것입니다. NVIDIA는 각 GPU가 완전한 KV 헤드를 갖도록 텐서 병렬화 수를 KV 헤드 수보다 크게 잡지 않는 체크리스트를 제시합니다.
KV 헤드가 한두 개로 적은 모델은 텐서 병렬화 확장 한계에 빨리 도달합니다. NVIDIA는 이런 경우 어텐션에는 ADP 또는 KVP를, MoE 피드포워드 네트워크에는 EP를 사용하는 Wide EP와 Helix Parallelism을 대안으로 설명합니다.
공동설계가 AI 에이전트에 미치는 영향
AI 에이전트와 멀티턴 앱은 긴 프리픽스 캐시를 반복해서 참조할 수 있습니다. NVIDIA는 새 입력이 짧더라도 긴 프리픽스를 읽는 프리필은 디코드와 비슷한 메모리 병목을 보일 수 있다고 설명합니다.
또한 FlashAttention은 전체 어텐션 행렬을 메모리에 만들지 않고 Q·K·V 타일을 HBM에서 온칩 SRAM으로 가져와 연산합니다. NVIDIA는 이 구조에서 행렬 곱은 Tensor Core, 소프트맥스 지수 연산은 특수 함수 유닛이 담당한다고 설명합니다.
실무에서는 최대 토큰 수 하나보다 첫 토큰 지연, 토큰 생성 지연, 처리량, KV 캐시 메모리, GPU 간 병렬화 효율을 함께 봐야 합니다. 이 네 지침은 같은 하드웨어에서 GPU 활용률과 상호작용성을 높이기 위한 NVIDIA의 체크리스트입니다.
결론: 더 긴 문맥보다 더 영리한 어텐션
NVIDIA 발표의 핵심은 장문맥 모델을 만든 뒤 GPU에 맞추는 것이 아니라, 그룹 크기·헤드 차원·KV 상태·병렬화 전략을 모델 설계 단계부터 함께 결정하라는 것입니다.
네 가지 지침은 범용 정답이 아니라 NVIDIA GPU에서 밀집 어텐션의 처리량과 상호작용성을 높이기 위한 출발점입니다. 모델 품질과 실제 워크로드를 함께 측정해야 최종 설계를 확정할 수 있습니다.
한 줄 요약: 장문맥 추론의 경쟁력은 컨텍스트 창 크기보다 어텐션 구조와 GPU 실행을 함께 최적화하는 데서 나옵니다.
