에이전트형 AI가 확산되면서 LLM 추론의 부하 특성이 달라지고 있다. 긴 프롬프트와 도구 호출 결과, 검색 결과, 여러 턴에 걸친 추론이 누적되면서 하나의 요청이 100만 토큰에 이르기도 한다. 문제는 텍스트를 실제로 생성하는 디코드 단계다. 새 토큰 하나를 만들 때마다 KV 캐시에 쌓인 과거 상태 전부를 참조해야 하므로, 커널은 HBM에서 지난 상태를 반복적으로 읽어 들인다. 그 결과 성능을 좌우하는 지점이 연산 유닛에서 메모리 시스템으로 옮겨간다. AMD가 새 Instinct MI450 GPU를 소개하면서 어텐션 디코드를 사례로 든 이유가 여기에 있다.
MI450은 메모리에 민감한 AI 워크로드를 겨냥해 이전 세대인 MI350 대비 온칩 자원을 늘리고 HBM 용량과 대역폭을 키웠다. 하드웨어 구성 단위 명칭도 정리됐다. 기존 MI 시리즈에서 CU로 부르던 연산 단위는 이제 워크그룹 프로세서(WGP)로 불린다. 하나의 WGP는 각각 32개 레인을 가진 SIMD32 유닛 4개로 구성되며, 각 SIMD32는 벡터 ALU와 행렬 곱셈 누산 유닛(WMMA)이 사용하는 자체 벡터 레지스터(VGPR)를 갖는다. WGP 안의 SIMD32들은 하나의 로컬 데이터 공유 메모리(LDS)를 공유한다.
MI450이 새로 얹은 하드웨어
가장 눈에 띄는 변화는 TDM이라는 전용 유닛이다. TDM은 글로벌 메모리와 LDS 사이에서 구조화된 텐서 데이터를 옮기는 역할을 한다. 커널이 접근할 텐서를 주소와 형태, 스트라이드, 레이아웃으로 기술하면 TDM이 대량 데이터 전송을 비동기로 처리한다. 작은 벡터 로드를 수없이 발행해야 했던 MI350 방식과 달리, 연산과 메모리 이동을 겹치는 파이프라인 커널을 짜기가 쉬워진다. 여기에 더해 워크그룹 클러스터가 도입됐다. 각기 다른 WGP에서 도는 여러 워크그룹이 하드웨어 클러스터 배리어로 협력하고, 같은 데이터를 쓰는 워크그룹끼리는 멀티캐스트 로드로 데이터를 공유한다. 별도 커널을 다시 실행하거나 글로벌 메모리로 동기화하지 않고도 더 넓은 범위의 협업을 표현할 수 있다는 의미다.
커널은 Gluon으로 작성한다. Gluon은 Triton 기반 DSL로, 타일 단위 SPMD 프로그래밍 모델을 유지하되 텐서마다 명시적 레이아웃을 요구한다. 각 원소가 레지스터, 레인, 웨이브, 워크그룹에 어떻게 분배되는지를 직접 기술하게 하므로 Triton보다 낮은 수준의 제어가 가능하다. 생성되는 명령어를 세밀하게 다루려는 커널 전문가에게 유용하지만, 그만큼 레이아웃 설계 부담을 개발자가 진다는 뜻이기도 하다.
어텐션 디코드에서 무엇을 최적화하나
어텐션은 Q와 K를 곱하고 소프트맥스로 P를 얻은 뒤 P와 V를 곱해 출력 O를 만드는 연산이다. 프리필은 Q 토큰 수와 KV 토큰 수가 같아 여러 토큰을 병렬 처리할 수 있지만, 디코드는 Q 토큰이 하나뿐이고 KV 토큰만 컨텍스트 길이를 따라 계속 늘어난다. 각 헤드가 독립적인 표준 멀티헤드 어텐션(MHA)에서는 단일 토큰 디코드가 연산 자원을 크게 낭비한다. 반면 GPT-OSS 같은 최신 모델이 쓰는 MQA에서는 여러 Q 헤드가 하나의 K·V 헤드를 공유하므로, 그 Q 헤드들을 묶어 처리할 수 있다. 커널은 (B, H_kv, 1) 그리드로 실행되고, 각 프로그램은 배치와 KV 헤드 한 쌍을 맡아 Flash Attention 방식으로 KV 타일을 순회한다.
블로그가 제시하는 최적화는 텐서 레이아웃, 데이터 로딩, 파이프라이닝, Split-k 병렬화 네 가지다. 레이아웃에서 가장 먼저 정해야 할 것은 WMMA 피연산자 레이아웃이다. 예를 들어 [16, 16, 128] 명령 형태는 FP8 피연산자 두 개를 소비해 FP32 출력 타일을 만든다. 웨이브는 소프트맥스 축소가 KV 방향으로 일어나므로 웨이브 간 통신을 피하기 위해 그룹화된 Q 방향으로 분배한다. 다만 H_q/H_kv 값이 작은 편이라 웨이브를 너무 많이 쓰지 않아야 하는데, 가령 이 값이 32면 웨이브 2개를 선택한다. 또한 QK 출력과 PV 입력의 레이아웃이 맞지 않으면 핫 루프가 변환 비용을 치르므로, WMMA 출력에 transpose 옵션을 주어 추가 명령 없이 전치하고 K Width를 8로 맞춰 두 연산을 연결한다.
대역폭의 85%, 그리고 남은 조건
데이터 로딩 최적화의 핵심은 캐시 우회다. MI450에는 LDS와 같은 계층에 WGP별 캐시가 있지만 프로그래머에게 직접 보이지 않는다. TDM은 데이터를 LDS로 바로 넣거나 캐시를 거쳐 넣을 수 있는데, 메모리 바운드 워크로드는 용량이 제한된 캐시를 우회해 LDS로 직행하는 편이 낫다. 직행 경로를 쓰려면 전송 요청의 최내차원이 최소 128바이트, 권장 256바이트여야 한다. K·V 타일의 최내차원은 헤드 차원 D인데, FP8 기준으로는 D가 256일 때 최적이다. 헤드 차원은 모델 구조가 정하므로 커널이 바꿀 수 없지만, K와 V를 재구성해 최내차원을 넓힌 뒤 LDS에서 레지스터로 읽을 때 되돌리는 우회로가 제시된다. 이렇게 레이아웃과 데이터 경로를 정리하고 온라인 소프트맥스의 루프 반송 값을 파이프라인으로 겹치면, AMD는 이 Gluon 커널이 MI450 최대 HBM 대역폭의 85%에 도달했다고 밝힌다.
다만 이는 AMD가 자사 블로그에서 공개한 초기 결과이며, 특정 구성을 전제로 한 수치라는 점을 감안할 필요가 있다. 최적화 상당수가 FP8, 특정 헤드 차원, MQA 같은 조건에 맞춰져 있어 다른 모델 구조나 데이터 타입에서 그대로 재현된다고 보기는 어렵다. 무엇보다 Gluon은 레이아웃과 K Width, 전치, 캐시 경로까지 개발자가 직접 지정해야 하는 저수준 도구다. 국내 실무자 관점에서 보면, 차세대 가속기의 실효 성능은 결국 이런 하드웨어 특성을 이해하고 커널 수준에서 메모리 이동을 설계할 수 있는 역량에 달려 있음을 보여주는 사례로 읽는 편이 현실적이다.