FuTransformer는 Hopper에서 프리필과 디코드를 분리해 최적화한다

무엇인가

이 논문은 Llama형 Transformer 추론에서 프리필과 단일 토큰 디코드의 병목이 다르다는 점을 출발점으로 삼는다. 프리필은 토큰 여러 개가 가중치를 재사용해 행렬곱과 활성값 이동이 지배적이고, 디코드는 재사용이 적어 가중치와 KV 캐시 읽기가 지배적이다. Tensor Core가 빨라질수록 행렬곱이 아닌 데이터 이동과 변환 커널의 비중이 커지므로, 저자들은 NVIDIA Hopper용 CUDA 구현인 FuTransformer로 CODA의 GEMM-plus-epilogue 방식을 추론에 맞게 확장한다.

어떻게 동작하나

프리필 경로는 영속 FP8 GEMM과 컴파일 시점에 선택되는 epilogue 정책을 쓴다. 128×256×128 CTA 타일, 4단계 TMA 파이프라인, 생산자 워프그룹 1개와 소비자 워프그룹 2개, Hopper FP8 WGMMA와 TMA를 사용한다. QKV epilogue는 K와 V를 offset(request, head, position, channel) 위치의 KV 캐시에 직접 쓰고, 밀집 K/V 텐서와 별도 scatter 커널을 없앤다. 추론 가중치는 고정이므로 RMSNorm gain γ를 양자화 전에 W′ = γW로 가중치에 접고, 행별 정규화 계수 ρ는 다음 투영의 epilogue에서 적용해 정규화된 활성값 텐서를 HBM에 쓰지 않는다. 지연된 per-tensor 스케일링으로 FP8 출력을 생산 epilogue에서 바로 쓰고, 현재 최대값은 다음 호출의 스케일로 계산한다. 패킹된 프리필은 prefix-sum 매핑으로 요청과 토큰 위치를 추적해 패딩 행을 건너뛴다.

무엇과 다른가

디코드 경로는 INT4 투영 가중치와 INT4 KV 캐시 복사본을 사용한다. Hopper에는 INT4 Tensor Core MMA가 없고 FP8 MMA는 디코드의 적은 쿼리 행으로 채우기 어려운 큰 워프그룹 형태를 요구하므로, 커널은 INT4 값을 16비트 MMA 피연산자로 확장한다. 가중치와 키를 MMA 프래그먼트 순서로 저장해 레지스터 안에서 또는 256개 항목 조회표로 직접 역양자화하며, K는 mma.sync.m16n8k16의 B 피연산자 바이트 순서를 따르고 V는 자연 채널 순서로 둔 뒤 공유 메모리에서 전치한다. 그룹 스케일은 누산 뒤에 적용하고, 가중치 니블은 BF16 128.0의 가수 비트에 넣은 뒤 패킹된 BF16 FMA로 136을 빼는 방식으로 처리한다. 캐시 양자화 전에는 블록 대각 Hadamard 변환을 적용하고 쿼리와 어텐션 출력에도 같은 변환을 적용해 qᵀk를 보존하면서 이상치 채널의 오차를 줄인다. 디코드 커널은 M=1용 레지스터 스트리밍과 2≤M≤32용 스키니 행렬, 긴 문맥용 split-KV 감소로 나뉜다.

어떻게 쓰나

실험은 H100 80GB SXM5 한 장, CUDA 13.0, PyTorch 2.14.1+cu130에서 수행됐다. 프리필은 컴파일된 PyTorch FP8보다 최대 1.48배 빨랐고, 1B 형상의 S=512에서만 0.92배로 느렸다. 405B 레이어 형상은 S=2048에서 1286 TFLOP/s, 측정 범위에서 1.20–1.29 PFLOP/s를 기록했다. 8B 형상 S=2048은 804.2μs로 PyTorch 1004.5μs보다 1.25배 빨랐다. CODA와 공유하는 gate/up 투영과 SwiGLU 단계에서는 BF16 기준 8B가 613μs로 CODA보다 1.13배, 컴파일 PyTorch보다 1.18배 빨랐고, 405B는 9442μs로 CODA 대비 1.03배 차이였으나 측정 변동 범위 안이었다. FuTransformer의 FP8 경로는 316μs와 4819μs로, 평가된 BF16 CODA 경로보다 2.02–2.18배 빨랐다.

전제와 한계

디코드에서는 8B와 70B 형상, 배치 1–32에서 INT4 블록이 각 지점의 PyTorch BF16 및 INT4 블록 중 더 빠른 쪽보다 1.3–2.4배 빨랐다. 컴파일·CUDA 그래프 BF16 블록보다는 2.1–3.3배, PyTorch INT4 투영 커널을 쓴 같은 블록보다는 1.3–6.3배 빨랐고, FuTransformer 자체 FP8 경로보다도 31–51% 빨랐다. 완전한 Llama-3.1-8B 디코드는 배치 1에서 스텝당 2.8ms로, 컴파일된 transformers 구현의 8.1ms보다 2.9배 빨랐다. 저자들은 배치 32·문맥 4096에서 PyTorch 스텝 60.8ms 중 어텐션이 89%를 차지하고, 8개 KV 헤드를 32개 쿼리 헤드로 확장하고 cuDNN용으로 캐시를 배치하는 복사 커널이 29ms, cuDNN 어텐션이 23ms, 투영이 5.9ms라고 분석한다. FuTransformer는 8헤드 INT4 캐시를 헤드 확장 없이 프래그먼트 순서로 직접 읽는다.

품질 평가는 Llama-3.1-8B와 WikiText-2 테스트셋의 2048토큰 윈도 141개로 수행됐다. BF16 커널 경로는 PyTorch 퍼플렉시티와 소수점 셋째 자리까지 일치했고, INT4 커널 경로는 같은 양자화 격자로 만든 시뮬레이션 참조와 0.001 퍼플렉시티 이내로 일치했다. INT4 가중치는 round-to-nearest에서 퍼플렉시티를 6.24에서 7.23으로 0.99 올렸고, 같은 형식의 GPTQ 가중치는 6.99로 0.75 올렸다. 회전된 INT4 캐시는 0.20을 더해 GPTQ 완전 INT4 블록이 7.18, BF16 대비 0.94가 됐다. 회전이 없으면 캐시 양자화만으로 0.54가 늘고 완전 블록은 7.76이 되어, 회전이 학습된 모델에서 캐시 품질 비용의 약 70%를 제거한다. 레이어별로는 BF16 커널이 FP32 참조와 0.3–0.4% 차이였고, INT4는 round-to-nearest가 10–13%, GPTQ가 4.5–7% 상대 RMS 오차를 더했다. 31층 이후 누적 오차는 각각 36%와 26%, 11층 근처에서는 42%와 30%였고, 로짓 KL 발산은 0.14와 0.08 nats/token, top-1 일치율은 83.5%와 88%였다. 서빙 워크플로에서는 앞 1024토큰을 BF16 또는 E4M3 프리필로 처리하고 나머지 1023토큰을 디코드할 때 완전 INT4 블록이 퍼플렉시티를 0.53, 9.4% 올렸고, 모든 토큰을 처음부터 디코드하면 0.94, 15.1% 올렸다. FP8 프리필은 0.004, E4M3 디코드 가중치는 0.02만 더했다. 객관식 평가에서는 INT4 가중치가 HellaSwag 2.1점, ARC-Challenge 4.9점을 떨어뜨렸고 INT4 캐시가 각각 0.4점, 2.7점을 더 떨어뜨려 ARC-Challenge 손실이 최대 7.6점이었다. 가변 길이 프리필에서는 B=8, Smax=2048에서 같은 길이 입력의 패킹 오버헤드는 보이지 않았고, 한 요청이 2048, 나머지 일곱이 128일 때 패킹이 5.07배 빨랐다.

개발자에게 이 논문은 Hopper에서 Llama형 모델을 서빙할 때 프리필과 디코드를 별도 커널 경로로 설계하고, 디코드에서 가중치와 KV 캐시 바이트 수를 함께 줄이는 선택이 실제 지연에 큰 영향을 준다는 근거를 제공한다. 다만 저자들이 밝힌 전제와 한계가 분명하다. PyTorch 프리필 비교는 활성값 스케일링 정책이 다르고, INT4 디코드와 PyTorch BF16 기준 비교는 정밀도가 다르며, CODA 비교는 전체 블록이 아니라 가장 가까운 공유 단계다. 성능 범위는 H100 한 장의 단일 레이어와 완전한 Llama-3.1-8B 모델에 한정되고, 다층 스케줄링, 분산 통신, 요청 수락, 완전한 서빙 지연은 포함하지 않는다. 405B 표기는 한 레이어의 차원을 뜻하며 405B 전체 모델을 한 GPU에서 실행한 것이 아니다. 정확도 평가는 한 모델, 퍼플렉시티용 한 데이터셋, 두 개의 제로샷 객관식 과제에 그치고, INT4 디코드는 프로토콜에 따라 퍼플렉시티를 9–15% 올리고 ARC-Challenge 정확도를 최대 7.6점 떨어뜨린다. 지연된 FP8 스케일링은 워밍업 후 안정적인 워크로드에서 시험됐으므로, 실제 배포에는 안전한 스케일 초기화, 포화 모니터링, 요청별 상태가 필요하고 학습된 체크포인트의 과제 정확도를 별도로 검증해야 한다. 구현은 Hopper의 프래그먼트 레이아웃, 비동기 복사, 자원 한도에 맞춰져 있어 다른 GPU 세대에서는 타일 형태, 파이프라인, 패킹 레이아웃을 바꿔야 한다.