🤖 machine learning

StreamKL: Fast and Memory-Efficient KL Divergence for Boosting Attention Distillation

StreamKL은 쿼리-키 타일을 단 한 번의 패스로 스트리밍함으로써 어텐션 증류(attention distillation)의 이차적인 메모리 및 I/O 병목 현상을 제거하는 새로운 융합 GPU 프리미티브를 도입하여, 단일 GPU에서의 긴 컨텍스트 증류를 가능하게 하기 위해 메모리 사용량을 $O(N_QN_K)$ 에서 $O(1)$ 로 줄이고 상당한 속도 향상을 달성했습니다.

원저자: Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

게시일 2026-06-19

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

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

당신은 작고 빠른 학생(학생 모델)에게 똑똑하고 느린 선생님(선생님 모델)처럼 생각하는 법을 가르치려 한다고 상상해 보십시오. AI의 세계에서, 사람들은 이들이 이야기나 문장의 서로 다른 부분에 어떻게 "주의(attention)"를 기울이는지 비교함으로써 이 작업을 수행합니다. 이 과정을 **어텐션 증류(Attention Distillation)**라고 부릅니다.

이 비교를 수행하기 위해, 컴퓨터는 **KL 발산(KL Divergence)**이라는 특정 숫자를 계산합니다. 이것을 "거리 점수"라고 생각하면 되는데, 학생의 주의가 선생님의 주의와 얼마나 다른지를 알려줍니다. 목표는 이 점수를 최대한 작게 만드는 것입니다.

문제점: "메모리 폭발"

이 논문은 긴 이야기(예: 10만 단어로 된 소설)에 대해 이 비교를 수행하는 것이 현재 컴퓨터 메모리에 악몽과 같다고 설명합니다.

다음은 비유입니다:
당신에게 선생님의 어텐션을 위한 거대한 화이트보드 하나와 학생의 어텐션을 위한 거대한 화이트보드 하나가 있다고 상상해 보십시오. 이들을 비교하려면, 기존 방식은 이 두 보드에 있는 모든 가능한 단어의 쌍을 일일이 적어야 합니다.

만약 64,000개의 단어가 있다면, 당신은 64,000 × 64,000개의 쌍을 적어야 합니다. 이는 무려 40억 개가 넘는 숫자입니다.
이를 수행하려면 컴퓨터의 메인 메모리(HBM)에 들어가지 않을 정도로 거대한 화이트보드가 필요합니다. 이는 마치 도서관의 책들을 신발 상자에 담으려는 것과 같습니다.
컴퓨터가 전체 그림을 한 번에 담을 수 없기 때문에, 이야기를 작은 조각으로 나누고, 처리한 다음, 다시 합쳐야 합니다. 이는 마치 책을 한 글자씩 보고 적은 뒤 다음 글자로 넘어가는 방식으로 책을 읽는 것처럼 느립니다.

해결책: StreamKL ("스트리밍" 접근 방식)

연구자들은 StreamKL이라는 새로운 도구를 만들었습니다. 모든 것을 거대한 화이트보드에 먼저 적는 대신, StreamKL은 마치 컨베이어 벨트처럼 데이터를 실시간으로 흘려보내며 "거리 점수"를 계산하는 영리한 트릭을 사용합니다.

창의적 비유: 공장 조립 라인
두 개의 제품 컨베이어 벨트(선생님의 어텐션과 학생의 어텐션)를 비교하는 공장을 상상해 보십시오.

기존 방식: 라인을 멈추고, 모든 제품을 거대한 창고 바닥(HBM)에 쏟아부은 다음, 모두 측정하고 나서 정리합니다. 이는 창고 전체를 차지하며 매우 느립니다.
StreamKL 방식: 제품이 계속 컨베이어 벨트 위를 움직이도록 둡니다. 각 쌍의 아이템이 센서(GPU 칩)를 통과할 때마다, 즉시 비교하고 차이를 계산한 뒤, 다음 쌍이 도착하기 전에 결과를 아주 작은 주머니(SRAM)에 던져 넣습니다. 당신은 라인을 멈출 필요도 없고, 창고도 필요하지 않습니다. 오직 주머니 하나만 있으면 됩니다.

작동 원리 (마법의 기술)

논문은 이 마법의 두 가지 주요 부분을 설명합니다.

순전파 (점수 계산):
연구자들은 컴퓨터가 점수를 점진적으로 업데이트할 수 있도록 하는 새로운 수학적 공식을 발명했습니다. 데이터를 스트리밍하면서, 전체 목록 대신 몇 개의 숫자(예: 실행 중인 최댓값과 합계)만을 사용하여 계속해서 집계를 유지합니다. 이는 메모리 부족 없이 어떤 길이의 이야기도 처리할 수 있게 해줍니다.
역전파 (실수를 통한 학습):
컴퓨터가 점수를 통해 학생으로부터 배워야 할 때, 보통 데이터를 다시 살펴봐야 합니다. 기존 방식은 전체 거대한 데이터 목록을 저장했다가 나중에 다시 봅니다. StreamKL은 더 똑똑합니다. 목록을 버리는 대신 몇 가지 "비밀 키"(LSE 값이라고 불림)를 기억합니다. 나중에 다시 돌아와야 할 때, 이 키들을 사용하여 필요한 특정 데이터 조각을 그 자리에서 바로 재구성하고, 교훈을 계산한 뒤, 다시 잊어버립니다. 이는 케이크를 통째로 구워 냉장고에 보관하는 대신, 케이크의 레시피를 기억해 두었다가 필요할 때마다 한 조각씩 구워 맛을 보는 것과 같습니다.

결과: 속도와 공간

논문은 강력한 NVIDIA GPU(H200 및 A100)를 사용하여 매우 긴 문맥(최대 512,000 단어)에 대해 테스트를 진행했습니다.

메모리 절감: StreamKL은 추가로 필요한 메모리를 "이차적(quadratic)"인 수준(테라바이트 단위로 폭발하는 수준)에서 "상수적(constant)"인 수준(매우 작게 유지되는 수준)으로 줄였습니다. 64k 문맥에서 512GB의 메모리가 필요했던 것을 거의 없는 수준으로 줄였습니다. 이를 통해 단일 GPU가 이전에는 슈퍼컴퓨터가 필요했거나 불가능했던 작업들을 처리할 수 있게 되었습니다.
속도: 데이터를 주고받는 과정이 방대하지 않기 때문에 믿을 수 없을 정도로 빠릅니다.
- 일부 테스트에서, 점수를 계산하는 데 있어 표준 방식보다 43배 더 빨랐습니다.
- 학습 단계에서는 14배 더 빨랐습니다.

요약

StreamKL은 AI 모델이 주의를 기울이는 법을 가르치는 새로운 방법입니다. 이 방식은 컴퓨터가 전체 비교 목록을 작성하지 않도록 함으로써, 긴 텍text를 다룰 때 발생하는 "메모리 부족" 문제를 해결합니다. 대신, 데이터를 작고 효율적인 파이프라인을 통해 흘려보내어 결과를 즉시 계산합니다. 이를 통해 이전에는 너무 커서 단일 컴퓨터로는 다룰 수 없었던 작업들을 일반 컴퓨터에서도 훈련하고 실행할 수 있게 되었습니다.

기술 요약: StreamKL

문제 정의

어텐션 증류(Attention distillation)는 하나의 어텐션 분포를 다른 분포와 일치시키기 위해 쿨백-라이블러(Kullback-Leibler, KL) 발산을 최소화하는 기술로, 지식 증류(knowledge distillation), 모델 프루닝(pruning), 지속 학습(continual learning), 그리고 희소 어텐션(sparse-attention) 대규모 언어 모델(LLM) 학습에 있어 매우 중요합니다. 그러나 기존의 구현 방식들은 긴 컨텍스트 길이( $N_Q, N_K$ )를 처리할 때 심각한 확장성 병목 현상에 직면합니다.

현재의 접근 방식들은 두 어텐션 분포( $P_1, P_2 \in \mathbb{R}^{N_Q \times N_K}$ )를 요소별 KL 감소(element-wise KL reduction)를 계산하기 전에 고대역폭 메모리(HBM)에 먼저 "실체화(materialize)"합니다. 이는 다음과 같은 결과를 초래합니다:

이차적 메모리 점유(Quadratic Memory Footprint): 두 개의 $N_Q \times N_K$ 행렬을 저장하는 데 $O(N_Q N_K)$ 메모리가 필요합니다. 32개의 헤드와 BF16 정밀도를 사용하는 64K 토큰 컨텍스트의 경우, 이 과정에서만 약 512 GB가 소모되어 단일 NVIDIA H200 GPU의 용량인 141 GB를 초과하게 됩니다.
이차적 IO 비용(Quadratic IO Costs): 이러한 분포를 HBM에 쓰고 다시 읽어오는 과정이 실행 시간을 지배하며, 컨텍스트 길이에 따라 이차적으로 증가합니다.
청킹(Chunking)의 비효율성: 청크 단위 처리는 메모리를 위해 레이턴시(latency)를 희생하는 일반적인 해결책이지만, 작은 청크는 GPU 병렬성을 저하시키고 큰 청크는 메모리 부족(OOM) 오류를 유발합니다.

표준 어텐션 계산은 FlashAttention을 통해 온라인 융합 커널(online, fused kernels)로 유사한 문제를 해결했습니다. 하지만 어텐션 증류에 이를 적용하는 것은 쉽지 않은데, 왜냐하면 KL 발산은 가중치가 적용된 로짓 차이(weighted logit-difference) 항을 통해 두 개의 독립적인 소프트맥스 분포를 결합하기 때문에, 단일 패스(single-pass) 온라인 공식화가 수학적으로 복잡하기 때문입니다.

방법론: StreamKL

StreamKL은 중간 분포의 $O(N_Q N_K)$ 실체화를 제거하는, 어텐션 KL 발산을 위한 최초의 융합형 단일 패스(one-pass) GPU 프리미티브를 도입합니다.

1. 온라인 공식화 (순전파, Forward Pass)

StreamKL은 결합된 두 분포의 KL 감소를 위한 새로운 온라인 업데이트 스킴을 도출합니다. 저자들은 행 단위 KL 발산을 재구성함으로써, 목적 함수를 다섯 가지의 행별 실행 스칼라(per-row running scalars)로 표현합니다:

$m_1, m_2$ : 두 로짓 분포에 대한 실행 최대값(running maximums).
$l_1, l_2$ : 지수 함수의 비정규화된 합(unnormalized sums of exponentials, 온라인 소프트맥스와 유사).
$acc$: 가중치가 적용된 로짓 차이를 위한 누산기(accumulator).

재구성된 목적 함수는 다음과 같습니다:
$L = \frac{acc}{l_1} + \text{LSE}_2 - \text{LSE}_1$
여기서 $\text{LSE}_t = m_t + \log l_t$ 입니다.

이 공식화를 통해 커널은 쿼리-키 타일(query-key tiles)을 온칩 SRAM으로 스트리밍하며, 이 다섯 가지 스칼라를 점진적으로 업데이트할 수 있습니다. 결정적으로, 누산기 $acc $는 전체 행렬이 아닌$ B_Q$(블록 크기) 크기의 벡터이며, 이는 표준 FlashAttention에 비해 산술적 리스케일링 오버헤드를 줄여줍니다. 또한, 이 알고리즘은 인과적 마스킹(causal masking)을 처리하기 위해 인과적 경계(causal frontier) 너머의 키 타일을 완전히 건너뜁니다.

2. 재계산 전략 (역전파, Backward Pass)

역전파 과정에서 이차적인 $P_1$ 및 $P_2$ 행렬을 저장하는 것을 피하기 위해, StreamKL은 재계산 전략을 채택합니다:

저장되는 상태: 최종 KL 값과 로그-합-지수(Log-Sum-Exp, LSE) 벡터( $\text{LSE}_1, \text{LSE}_2$ )만이 HBM에 $O(N_Q)$ 로 저장됩니다.
재생성: 역전파 중에, 저장된 LSE와 원본 $Q, K$ 행렬을 사용하여 어텐션 확률을 타일 단위로 재생성합니다 ( $P_i = e^{S_i - \text{LSE}}$ ).
그래디언트 계산: 커널은 다음 두 가지 설정에 대해 그래디언트를 계산합니다:
- 설정 1: $P_1$ 은 고정하고 $P_2$ 를 최적화 (표준 증류).
- 설정 2: $P_2$ 는 고정하고 $P_1$ 을 최적화.
커널 융합: 저자들은 각 $QK$ 타일 쌍을 정확히 한 번씩 방문하는 융합된 역전파 커널을 설계하여 HBM 트래픽을 줄였습니다. $dQ $(K를 따름)를 누적하는 것과$ dK $(Q를 따름)를 누적하는 것 사이의 충돌을 관리하기 위해, 커널은 버퍼가 L2 캐시에 들어갈 때는$ dQ$를 위한 원자적 연산(atomic operations)을 사용하고, 경합이 너무 높아지면 별도의 커널로 전환합니다.

3. GPU 커널 최적화

Split-K 변형: 쿼리 수가 적은 워크로드(예: $N_Q=1$ 인 자기회귀 디코딩)를 위해, StreamKL은 스레드 블록 간에 키 차원을 분할하여 스트리밍 멀티프로세서(SM)를 포화시키는 Split-K 변형을 제공합니다.
TMA 가속: Hopper GPU에서 구현체는 비동기 벌크 메모리 복사를 위해 텐서 메모리 액셀러레이터(TMA)를 활용하여 메모리 레이턴시를 숨깁니다.
수치적 안정성: 설정 2에서의 로그 비율 계산은 수치적 불안정성( $-\infty$ 또는 NaN)을 피하기 위해 직접적인 확률 로그 대신 로짓 차이와 저장된 LSE를 사용하여 수행됩니다.

주요 기여

병목 현상 식별: 본 논문은 어텐션 KL 발산이 긴 컨텍스트로 확장되는 것을 방해하는 이차적 메모리 및 IO 병목 현상을 식별했습니다.
온라인 재구성: 최초의 온라인 재구성을 도입하여, 이차적인 중간 결과물 없이도 완전히 융합된 순전파 프리미티브와 재계산 기반의 역전파를 가능하게 했습니다.
효율적인 구현: 저자들은 순전파와 역전파 모두를 위해 Split-K, TMA, 적응형 융합(adaptive fusion)과 같은 전용 최적화가 포함된 효율적인 GPU 커널을 설계하고 구현했습니다.
확장성: StreamKL은 어텐션 증류의 추가 HBM 점유를 $O(N_Q N_K)$ 에서 $O(1)$ 로 줄여, 단일 GPU에서 긴 컨텍스트 증류를 가능하게 합니다.

실험 결과

NVIDIA A100 및 H200 GPU에서 4K부터 512K까지의 컨텍스트 길이에 대해 평가되었습니다:

메모리 효율성: StreamKL은 컨텍스트 길이에 관계없이 평탄한(flat) 메모리 점유를 유지합니다. 반면, 베이스라인(PyTorch, torch.compile, FLA)은 이차적으로 증가하며 16K–64K 구간에서 OOM 오류가 발생합니다. 64K에서 StreamKL은 torch.compile/FLA보다 16,000배 이상, PyTorch보다 135,000배 이상 적은 HBM을 사용합니다.
순전파 레이턴시:
- StreamKL은 비인과적(non-causal) 설정에서 PyTorch 대비 13–18배, torch.compile/FLA 대비 2.4–3.7배 빠른 속도를 달성합니다.
- 인과적 마스킹(causal masking) 환경에서는, StreamKL이 마스킹된 계산을 완전히 건너뛰는 반면 베이스라인은 전체 행렬을 실체화하기 때문에, PyTorch 대비 26–43배, 베이스라인 대비 4.2–7.1배 더 빠른 속도를 보입니다.
역전파 레이턴시:
- 설정 1에서 StreamKL은 비인과적 설정 시 PyTorch보다 2.6–4.0배, 인과적 설정 시 5.3–9.9배 빠릅니다.
- 설정 2에서 StreamKL은 비인과적 설정 시 PyTorch보다 3.8–6.5배, 인과적 설정 시 9.1–14.0배 빠른 속도를 달성합니다.
- torch.compile과 비교했을 때, 특히 마스킹된 타일을 건너뜀으로써 재계산 비용이 절반으로 줄어드는 인과적 설정에서 StreamKL은 경쟁력이 있거나 더 빠릅니다.
절제 연구(Ablation): Split-K 변형은 $N_Q$ 가 작은 워크로드에서 최대 8.4배의 속도 향상을 제공합니다. 융합된 역전파 커널은 $N_Q$ 가 작을 때 별도의 커널보다 최대 2.0배 더 빠르며, 이는 적응형 선택 전략의 정당성을 입증합니다.

의의

본 논문은 StreamKL이 단일 GPU에서 64K 이상의 컨텍스트 길이를 가진 어텐션 증류 워크로드를 지속할 수 있는 최초의 솔루션이라고 주장합니다. 이차적인 어텐션 분포 실체화를 제거함으로써, StreamKL은 현대 LLM 학습(예: DeepSeek V3.2 및 GLM-5와 같은 희소 어텐션 모델)에서 어텐션 증류의 확장을 가로막는 주요 시스템 병목 현상을 제거합니다. 이는 융합된 온라인 프리미티브가 표준 어텐션을 넘어 KL 발산과 같은 복잡한 결합 분포로 성공적으로 확장될 수 있음을 보여줍니다.