Qwen 팀이 GDN 선형 어텐션용 GPU 연산자 FlashQLA를 공개했다

InfoQ 中文10일 전조회 1

알리바바 Qwen 팀의 알고리즘 엔지니어 Zhang Chengruidong이 AICon 2026 선전 행사에서 GDN(Gated Delta Network) 선형 어텐션 전용 GPU 연산자 라이브러리 FlashQLA를 소개했다. 발표에 따르면 FlashQLA는 선형 어텐션 커널이 안고 있던 낮은 병렬도, 중복된 메모리 접근, SFU(특수 기능 유닛) 과부하를 한꺼번에 겨냥한다. 대상 독자는 대규모 모델 시스템 엔지니어와 GPU 커널 개발자, 추론 엔진 아키텍트다.

배경은 컨텍스트 길이다. 시퀀스가 수만에서 수백만 토큰으로 늘면서 softmax 어텐션의 O(N²) 비용이 사실상 상한선이 됐다. Qwen과 Kimi는 선형 어텐션과 softmax 어텐션을 섞는 하이브리드 구조로 방향을 틀었다. Qwen은 GDN 3개 층마다 softmax 어텐션 1개를 넣는 3대1 구성이고, Kimi K3는 KDA 3개에 softmax 1개를 붙인다. 선형 어텐션은 상태를 d×d 행렬에 누적하는 재귀 구조라 계산량이 O(Nd²)로 줄지만, 모든 정보를 그 작은 상태에 압축해야 해서 장거리 회수에 약하다. 그래서 일정 간격으로 softmax 층을 끼워 넣는다.

문제는 이 재귀를 GPU에서 어떻게 빠르게 돌리느냐다. 토큰 단위 재귀는 벡터 내적과 외적에 의존해 텐서 코어를 제대로 쓰지 못한다. 그래서 WY representation으로 재귀 단위를 청크(보통 64)로 키운다. 청크 내 K·Kᵀ를 구해 역행렬을 만들고, 이를 이용해 W와 U를 뽑고, 청크 단위로 상태를 넘기고, 마지막에 출력을 계산하는 4단계다. 1·2·4단계는 청크끼리 독립적이지만 3단계는 앞 청크의 상태에 묶인다. 기존 FLA 저장소 구현은 이 4단계를 각각 별도 연산자로 두고 W, U, S를 HBM에 썼다 다시 읽는다. 추론 프레임워크가 보관하지도 않는 임시값 때문에 대역폭을 그대로 버리는 셈이다.

가장 큰 병목은 병렬도다. 3단계 때문에 어텐션 헤드 하나가 스레드 블록 하나, 즉 SM 하나를 차지한다. Qwen 소형 모델의 선형 어텐션 헤드는 32개 수준인데 H100이나 GB 계열 카드는 SM이 100개를 넘는다. FlashQLA는 시퀀스를 쪼개 병렬화한다. 중간에서 잘라 뒤쪽을 0 상태에서 재귀하면 틀린 상태가 나오는데, 인접 청크의 상태 차이를 다시 차분하면 dK×dK 행렬 하나를 곱하는 관계가 유지된다는 점과 행렬 곱 결합법칙을 이용하면 청크마다 보정 행렬을 따로 계산해 마지막에 한 번만 적용할 수 있다. 모든 청크의 보정 행렬을 저장하면 메모리가 터지므로, CP 랭크마다 마지막 청크 결과만 남기고 상태 의존 재귀를 한 번 돌린 뒤 S를 두 번 계산하는 방식으로 계산을 메모리와 맞바꾼다.

GDN에는 그동안 활용되지 않던 성질이 하나 더 있다. alpha(망각 게이트)가 0과 1 사이 실수이고 K가 L2 정규화돼 있어서 과거 상태의 영향력이 alpha의 거듭제곱으로 줄어든다. 즉 alpha가 항상 1이 아닌 헤드에는 반드시 유효한 슬라이딩 윈도가 존재한다. 발표에 따르면 학습된 모델의 선형 어텐션 헤드 가운데 60~80%가 여기에 해당했다. 이 헤드들은 0번 상태부터 재귀할 필요 없이 윈도 크기만큼만 되돌아가면 되고, 시퀀스 병렬화에서 어차피 계산하는 S*를 그대로 정답으로 쓸 수 있어 보정 행렬 계산을 생략할 수 있다. 어떤 헤드가 해당하는지는 gate 분포를 훑는 가벼운 커널 하나로 판별한다.

연산 자체도 손봤다. 청크 내 gamma 외적 행렬을 계산하려면 C²번의 EXP가 필요해 SFU가 막힌다. 이 행렬의 대각 성분이 1이라는 점을 이용해 역행렬 계산 바깥으로 대각 스케일링을 빼낼 수 있고, 이후 W·U를 구할 때 다시 곱해주면 EXP 연산을 한 곳으로 모을 수 있다. 시퀀스 병렬 보정 행렬 M의 128×128 행렬 연쇄도 같은 기법으로 gamma 항을 정리해 X, Y, Z 형태로 바꾸면 수치 범위가 안정되고 계산량이 줄어든다. 역전파에서는 DA와 DA의 전치를 묶어 K와 한 번만 곱하는 식으로 행렬 곱 횟수를 줄였다.

융합 후에도 이 연산자는 메모리 바운드다. 청크 하나 기준으로 Q, K, V, A, 출력 O의 HBM 트래픽이 18단위인데 H200의 WGMMA 텐서 코어 성능을 환산한 계산량은 30단위 수준이다. L2 적중률이 낮은 워크로드라 2단계 핑퐁 파이프라인으로 충분하다고 판단했다. 대신 레지스터와 공유 메모리가 빠듯하다. 128×128 S 행렬만 해도 워프 그룹당 스레드마다 32비트 레지스터 128개를 채운다. 호퍼에서는 WGMMA를 워프 그룹 단위로 발행해야 해서, TMA로 Q·K·V를 나르고 O를 쓰는 생산자 그룹 하나와 V', O, S를 각각 맡는 소비자 그룹 셋으로 역할을 나눈다.

실무 관점에서 이 발표가 주는 힌트는 선형 어텐션 서빙의 병목이 알고리즘이 아니라 커널 스케줄링과 상태 전달에 있다는 점이다. 다만 아직 진행 중인 작업이 섞여 있다. 카드 내 시퀀스 병렬에서 최적 분할 수를 √N 기반으로 추정하는 모델은 단순화가 거칠고, auto CP 활성화 경계도 정밀 모델 대신 사다리꼴 근사로 두고 있다. FLA 저장소 구현은 현재 카드 간 시퀀스 병렬만 지원한다. Triton으로는 워프 특화나 세밀한 파이프라인 제어가 안 되기 때문에 더 낮은 수준의 도구가 필요했다는 것도 전제다. 원문 발췌에는 성능 벤치마크 수치가 담겨 있지 않아 실제 이득의 크기는 확인할 수 없다.