FLAT은 구글 TPU에서 선형 어텐션, Mamba-2, Gated Linear Attention(GLA)을 학습하기 위한 Pallas 커널과 선형 어텐션 앞단의 짧은 인과 합성곱 커널을 공개한 프로젝트다. 제작진은 단일 TPU v6e 칩에서 측정한 결과, 설정에 따라 NVIDIA H100에서 실행한 FLA보다 순전파와 역전파를 합쳐 최대 3.08배 빠르며, XLA로 컴파일한 같은 청크 알고리즘보다 최대 11.4배 빠르다고 보고했다. 다만 측정은 모두 단일 v6e 칩에서 이뤄졌고, DeltaNet·Gated DeltaNet·Kimi Delta Attention 커널은 아직 공개하지 않았다.
선형 어텐션은 문맥 길이에 따라 커지는 KV 캐시 대신 고정 크기의 순환 상태를 사용한다. 따라서 토큰 생성의 시간과 메모리 비용이 일정하며, 학습 비용은 문맥 길이 T에 대해 O(T)라고 설명한다. Kimi K3는 Kimi Delta Attention을 사용하고, Qwen 3.5는 Gated DeltaNet 층과 전체 어텐션을 교차 배치한다. GPU에는 FLA와 Mamba의 CUDA 커널이 있지만, 프로젝트 측은 TPU에서 이러한 층을 학습하는 오픈소스 커널을 알지 못했다고 밝혔다. TPU 추론용 커널은 기존에 연구돼 왔다.
TPU용 커널 설계와 공개 범위
제작진은 TPU v6e의 행렬 연산 장치(MXU)와 벡터 처리 장치(VPU)의 성능 차이, 정적 명령어 스케줄링, 헤드별 행렬 곱의 작은 크기가 구현의 주요 장애였다고 설명한다. 자체 측정에서 VPU는 약 7 TFLOP/s였고 MXU의 정격 성능은 918 TFLOP/s였다. 지수 계산, 마스킹, 누적합 등 선형 어텐션의 벡터 연산은 VPU를 사용한다. TPU 컴파일러는 실행 전에 작업을 고정 폭 명령어 묶음으로 배치하므로, 메모리를 기다리지 않는 커널의 실행 시간은 발행한 묶음 수에 좌우된다고 한다.
FLAT은 여러 헤드를 한 MXU 행렬 곱에 묶어 행렬 장치 활용도를 높이고, 감쇠 게이트의 누적합을 1로 채운 하삼각 행렬과의 곱으로 바꿔 MXU에서 처리한다. 또한 v6e 칩의 단일 코어 특성을 활용해 순전파의 상태 전달과 출력 계산을 한 커널로 합치고, 순환 상태를 전체 시퀀스 동안 온칩 VMEM에 유지한다. GPU용 FLA는 다수의 스트리밍 멀티프로세서(SM)를 바쁘게 유지하기 위해 청크 상태를 HBM에 기록했다가 다시 읽는다.
공개 구현은 accelerated-model-architectures 저장소에 있다. Linear Attention 순전파는 청크 단위 순환 계산을 하나의 Pallas 프로그램으로 결합한다. 현재 구현은 배치와 시퀀스를 타일링하고, 시퀀스 축은 이전 청크 상태에 의존하므로 순차 처리한다. 기본 청크 크기는 256이며, 헤드 상태는 fp32로 VMEM에 둔다. Mamba-2는 감쇠값의 청크 내 누적합을 로그 공간에서 계산한다. 한 헤드의 스칼라 게이트에서는 별도 행렬 곱 방식이 6마이크로초, XLA의 jnp.cumsum은 86마이크로초였다고 보고했다. GLA는 채널별 게이트를 사용하며, 최적화 경로에서는 누적합을 하삼각 행렬 곱으로 바꿔 MXU에서 계산한다.
역전파는 순전파를 다시 실행해 각 청크 시작 상태를 재계산하고 HBM에 저장하는 상태 전달 커널, 청크를 역순으로 돌며 입력과 게이트 기울기를 계산하는 기울기 커널로 구성된다. 게이트 기울기의 누적합은 MXU의 삼각 행렬 곱으로 처리한다. 서로 상쇄되는 fp32 항을 bf16 한 번의 연산으로 변환하면 강한 감쇠 조건에서 오차가 2%였기 때문에, 커널은 이를 두 개의 bf16 부분으로 나눠 두 차례 계산한다.
짧은 깊이별 인과 합성곱도 별도 커널로 제공한다. 일반적인 선형 어텐션 변형은 순환 계산 전에 q, k, v에 이 합성곱을 적용한다. 프로젝트는 커널을 순환 계산에 융합하지 않았으며, 융합하면 HBM 왕복을 줄일 수 있다고 향후 과제로 제시했다.
벤치마크와 조건
모든 측정은 단일 TPU v6e 칩의 장치 시간이며, 별도 언급이 없으면 입력과 출력은 bf16, 누적과 감소 연산은 fp32다. 주요 비교 설정은 배치 크기 B=8, 시퀀스 길이 S=4096, 헤드 수 N=16, K=V=128이다. 메모리 대역폭 기준으로 jnp.copy에서 1.46 TB/s를 측정했고, 이는 v6e HBM 정격 1.638 TB/s의 89%다. MXU 정격은 bf16 기준 918 TFLOP/s이며, 별도 커널로 측정한 VPU 최고 성능은 7.091 TFLOP/s다.
이 설정에서 순전파는 Linear Attention 384.9마이크로초(1.42 TB/s), Mamba-2 438.0마이크로초(1.25 TB/s), GLA 549.8마이크로초(1.48 TB/s)였다. 청크 단위 일반 JAX 구현보다 각각 6.4배, 6.1배, 9.3배 빠르다. 한 위치씩 계산하는 기준 구현이나 jax.lax.scan은 33~100밀리초로 FLAT보다 87~260배 느렸다. 제작진은 이러한 구현이 작은 행렬-벡터 연산을 반복하고, XLA 구현에서는 중간 상태가 HBM을 오간다고 설명했다.
상태 전달 커널은 Linear Attention 1.48 TB/s, Mamba-2 1.31 TB/s, GLA 1.53 TB/s를 기록했다. 청크 JAX 구현보다 각각 3.7배, 5.8배, 8.2배 빠르다. 기울기 커널은 각각 1.42, 1.00, 1.37 TB/s로, 청크 JAX보다 7.4배, 8.1배, 11.4배 빠르다. 다만 한 위치씩 계산하는 기준 구현의 역전파는 45분 안에 컴파일되지 않았고, jax.lax.scan 역전파는 4096개 단계의 상태를 유지하다 HBM을 초과했다. 이 조건에서는 두 구현의 수치를 제시하지 않았다.
FLA의 Triton 구현은 H100 80GB에서 같은 형태로 측정했다. H100의 jnp.copy에 해당하는 메모리 기준은 2.98 TB/s로 v6e의 2.04배였다. 그럼에도 FLAT은 순전파에서 Linear Attention, Mamba-2, GLA가 각각 1.22배, 1.14배, 2.94배 빨랐고, 역전파에서는 1.76배, 1.64배, 3.12배 빨랐다. 순전파와 역전파를 합친 전체 비교에서는 각각 1.59배, 1.50배, 3.08배였다. 비교는 FLA의 청크 크기 64와 chunk_simple_gla 또는 chunk_gla 구현을 사용했다.
합성곱 벤치마크는 B=8, 커널 크기 K=4, 시퀀스 길이 4096, 은닉 크기 2048, 편향과 SiLU 활성화를 사용했다. 순전파는 FLAT 239.5마이크로초, XLA 495.1마이크로초로 2.07배 차이가 났다. 역전파는 FLAT 431.1마이크로초, XLA 2,806마이크로초로 6.5배 차이가 났다. 양방향 합산 시간은 671마이크로초와 3,301마이크로초로 4.9배 차이다. XLA는 가중치 기울기를 시퀀스 전체에 대한 추가 합성곱으로 계산해 3.36GB를 이동했으며, FLAT의 해당 역전파 이동량 407MB의 8.2배였다고 한다.
전체 층에서의 결과와 한계
전체 학습 단계 비교는 FLA 표준 GLA 층과 비슷한 구성을 사용했다. 합성곱은 6,144개 채널에 4개 탭을 적용하며, 출력에는 헤드별 RMSNorm과 swish 게이트가 있다. 선형 어텐션 커널만 적용했을 때 순환 계산은 1.3~2.0밀리초로 줄었고, 전체 단계는 Linear Attention, Mamba-2, GLA에서 각각 1.58배, 1.66배, 2.14배 빨라졌다. 합성곱 커널까지 적용하면 전체 단계는 10.2, 10.9, 12.5밀리초가 되어 커널을 사용하지 않은 경우보다 각각 2.12배, 2.18배, 2.75배 빨랐다. 이 비교의 커널 미사용 기준은 청크 JAX 구현이다.
같은 크기의 인과 소프트맥스 어텐션 층은 TPU용 Pallas 커널인 Splash Attention을 사용해 한 단계에 15.9밀리초가 걸렸다. 모든 커널을 적용한 선형 어텐션 층은 이보다 1.27~1.56배 빨랐다. 다만 선형 어텐션 층에는 소프트맥스 층에 없는 합성곱, 감쇠 게이트, 출력 게이트가 포함된다. 시퀀스 길이 4096에서의 결과이며, 제작진은 소프트맥스 어텐션 비용은 시퀀스 길이에 따라 제곱으로, 이 구현의 비용은 선형으로 증가하므로 더 긴 시퀀스에서는 격차가 커진다고 설명한다.
이 결과를 다른 TPU 세대에 일반화할 수는 없다. 측정은 단일 v6e 칩에서만 했으며, 다른 세대와 긴 문맥을 위한 칩 간 시퀀스 분할은 아직 시험하지 않았다. 또한 현재 공개 커널은 DeltaNet, Gated DeltaNet, Kimi Delta Attention을 지원하지 않는다. 이 변형은 상태 갱신이 단순한 외적의 합이 아니어서 별도 커널이 필요하며, 제작진은 지원을 곧 추가할 계획이라고 밝혔다. 합성곱과 순환 계산의 융합, 출력 정규화와 게이트 전용 커널, Amazon Trainium 지원도 향후 과제로 들었다.
관련 자료: FLAT 저장소, Rigel, FLA, TPU v6e 사양.
댓글 (0)
로그인하면 이 기사에 내 생각을 남길 수 있어요