RT-Lynx: Diffusion Transformers에서 GEMM Sparsity를 Activation으로 옮겨 실현하는 방법

Diffusion Transformer(DiT) 계열 모델의 추론 비용은 높은 편이다. 가중치 중심의 sparsification은 성능 저하를 야기하는 반면, Activation sparsity는 본질적으로 더 희소한 패턴을 만들고 품질 손실을 크게 줄인다. RT-Lynx는 Activation sparsity를 online으로 적용하고 Norm Compensation 및 LoRA 보정으로 품질 저하를 억제하며, CUDA 기반의 엔드투엔드 파이프라인으로 실사용 속도 향상을 달성한다.

왜 중요한가

핵심 기여

Activation 중심의 sparsification으로 DiT의 성능 저하 최소화

Activation sparsity가 Weight sparsity보다 DiT의 이미지 생성 품질 저하를 훨씬 덜 유발한다는 경험적 발견을 바탕으로, 2:4 패턴의 활성화 희소화를 제안하고 Norm Compensation으로 노름 손실을 보정한다.

Norm Compensation과 LoRA를 결합한 종합 보정

S(X)로 희소화된 활성화의 노름을 원래 활성화의 노름과 일치시키고, 잔차 정보를 LoRA 분기로 보정하여 고주파 디테일 회복 및 품질 유지에 기여한다.

online Sparse GEMM과 Sparsified Tensor Core의 엔드-투-엔드 파이프라인

Pattern 결정, Top-K 선택, 압축 등을 하나의 CUDA 실행 경로로 결합해 SpTC compatible 레이아웃에 2:4 활성화를 직접 생성하고, sparse GEMM과 dense LoRA를 결합하는 엔드-투-엔드 실행 파이프라인을 구현했다.

여러 DiT 모델 및 벤치마크에서의 실증과 호환성

Qwen-Image, Flux.1-dev, Z-Image 등 다양한 DiT 아키텍처에서 일관된 속도 향상과 품질 유지가 확인되었으며, 다른 가속 기법(distillation, quantization, cache, sparse attention)과의 결합에 문제없이 작동한다.

핵심 아이디어 이해하기

출발점: DiT의 가중치 분포는 quasi-Gaussian하고 비교적 밀집된 반면, 토큰별 활성화는 각 토큰 내에서 소수의 채널만 활발히 작용한다. 이로 인해 activation sparsity가 더 자연스러운 제약으로 작동하며, 2:4와 같은 semi-structured 패턴에서도 출력 오차가 가중치 sparsity보다 작다.
해결 원리: Activation sparsity를 온라인으로 적용하고, norm 보정(NS)과 LoRA 보정으로 잔차를 보완한다. S(X)로 활성화를 2:4 구조로 축소하되, s를 통해 원래 활성화의 L2 노름을 재현한다. 추가로 LoRA 분기(LA, LB)를 도입해 고주파 디테일의 손실을 보정한다.
달라지는 점: 이로써 DiT의 추론 속도는 엔드-투-엔드에서 밀도 기반의 가중치 pruning보다 큰 이득을 얻으며, 평균적으로 선형 계층에서 최대 1.55×의 속도향상을 실현하고, 1.88× Sparse GEMM 가속을 달성한다. 또한 SpTC와의 결합으로 하드웨어 친화적 구현이 가능하다.

방법론

전체 접근 방식: activation sparsity를 이용하여 DiT의 선형 계층에서 2:4 패턴을 적용하고, Norm Compensation으로 활성화 노름 손실을 보정하며 LoRA 분기를 통해 잔차 정보를 보완한다. Selective Layer Skipping으로 단일 스트림 경로의 주요 계층에서 성능 저하를 줄인다.
핵심 메커니즘: Top-K 기반의 2:4 활성화 선택을 수행하고, X˜를 얻은 후 S(X) = s·X˜로 보정한다. Ys = S(X)W⊤, Yr = X(LA LB)⊤, Y = Ys + Yr로 출력한다. 활성화 sparsification은 온라인으로 수행되며, 2:4 패턴에 맞춰 SpTC 형식으로 커널에 직접 전달된다.
구현 및 학습 세부: LoRA 랭크 R은 64로 고정하고, Loss = ||X·W⊤ − (S(X)·W⊤ + X·(LA LB)⊤)||^2를 최소화하도록 LA, LB만 학습한다. Distillation 기반의 지도교사-학생 방식으로 학습하며, 2k 스텝 내에 수렴하는 것을 확인했다. 단일 스트림 경로에서는 Out/OUp 등 특정 레이어를 건너뛰는 전략이 도입되었다.
CUDA 커널 최적화: 온라인 희소화 파이프라인을 하나의 실행 경로로 합성하고, Sparse GEMM 커널은 streamK 스타일의 블록 병렬 파이프라인으로 대역폭-지연 구조를 활용한다. LoRA 분기와의 연동은 온칸에서 Yr를 생성하고 Ys에 누적한다.

주요 결과

메인 벤치마크: MJHQ에서의 FID(↓)은 Ours가 21.25, IR(↑) 1.304, C.SCR(↑) 27.33, C.IQA(↑) 0.9263으로 Sparse Weight 및 다른 SOTA 대비 우수하다. sDCI에서 FID(↓) 25.78, IR(↑) 1.226, C.SCR(↑) 26.61, C.IQA(↑) 0.9366과 비슷하거나 우수하다.
Ablation: SA-NC-LoRA 구성에서 FID 21.25까지 회복되며, 싱글 스트림 모델의 경우 Layer Skipping으로 추가 개선 가능하다.
속도/효율성: Dense 대비 End2End 속도는 약 1.2×의 이점을 달성하며, Linear 계층에서 1.43×~1.55×의 속도 향상을 보인다. Qwen-Image의 예에서 원래 0.75s의 생성 시간이 0.62s로 감소했다. Sparse-GEMM 백엔드의 대조에서 RT-Lynx 커널은 최대 1.88×의 속도업과 Sparse-Cost를 낮춘 것으로 보고된다.
타 기법과의 호환성: distillation, W8A8 quantization, TeaCache, SpargeAttn 등과의 조합에서도 품질 저하 없이 작동한다.

기술 상세

전체 아키텍처: DiT 기반 Transformer 블록의 QKV 투영 및 Up/Down 매핑에 activation sparsity를 적용한다. 핵심 수식은 S(X) = s·X˜, s = ŝ/ (∥X∥2)로 정의되며, Y = S(X)W⊤ + X(LA LB)⊤로 출력한다. LoRA 랭크 R은 64로 고정되며 LA, LB만 학습한다. Norm Compensation으로 활성화의 노름 손실을 보정하고, Selective Layer Skipping으로 단일 스트림 경로의 특정 레이어를 건너뛰는 전략을 도입한다. CUDA 커널은 online sparsification과 sparse Tensor Core를 하나의 경로로 융합하여 2:4 패턴의 활성화를 직접 레지스터에 구성하고, Ys와 Yr를 온칸에서 누적한다. 실험은 Qwen-Image, Flux.1-dev, Z-Image에서 수행되었고 핵심 지표로 FID, Image Reward(IR), CLIP-IQA(C.SCR, C.IQA)를 사용한다.

실무 활용

RT-Lynx는 activation sparsity를 활용한 엔드-투-엔드 가속 프레임워크로, Diffusion Transformer 계열 모델의 추론 비용과 응답 시간을 감소시키면서도 품질을 크게 유지한다. 다른 가속 기법과도 상호 보완적으로 작동한다.

대형 diffusion 기반 이미지 생성 서비스에서 실시간 응답 시간 단축
저전력 또는 모바일/엣지 디바이스에서 DiT 추론 가속화
다중 모델 파이프라인에서 activation sparsity를 이용한 일괄 처리 속도 향상
스케일링된 DiT 계열 모델의 배포 시 추가적인 하드웨어 가속과의 조합

코드 공개 여부: 미확인

키워드

Activation SparsityN:M sparsityRT-LynxDiTSparse Tensor CoreNorm CompensationLoRA adaptationend-to-end inference

용어 해설

N:M 희소성(N:M sparsity): — 고정된 로컬 N:M 패턴으로 각 계층의 출력에서 N개의 비제로 원소만 남기고 나머지를 0으로 설정하는 semi-structured sparsity 기법. 이를 통해 하드웨어에서의 처리 효율을 높이고 FLOPs를 대폭 감소시킬 수 있다.
Activation sparsity: — 토큰당 활성화되는 채널 수가 매우 작아지는 현상으로, DiT의 FFN에서 활성화가 주로 0에 가까운 값으로 분포하며 비제로 채널 수가 적다. 가중치 희소화보다 출력 오차가 작고 품질 저하가 완화된다.
Sparse Tensor Core: — NVIDIA의 Semi-structured sparsity를 지원하는 전용 실행 경로로, 2:4와 같은 패턴의 연산을 하드웨어 차원에서 가속하고 nnz를 절반으로 줄이되 기억접근 패턴을 규칙적으로 유지한다.
Norm Compensation(Norm Compensation): — 활성화 희소화로 인한 활성화 벡터의 L2 노름 감소를 보정하기 위한 방법으로, 희소화된 활성화를 dense 원래 활성화의 노름과 동일하게 만들기 위해 스케일링을 적용한다.
LoRA 어댑테이션(LoRA Adaptation): — 희소 활성화로 인한 미세한 잔차 정보를 보완하기 위해 LoRA 매개변수 LA, LB를 도입해 저차원 어댑터를 학습시키는 방법. 주 모듈은 고정하고 LoRA 분기만 학습한다.
Diffusion TransformERs(Diffusion Transformers (DiT)): — Diffusion 모델에 Transformer 모듈을 결합한 구조로, 고해상도 이미지 생성 품질과 확장성을 높이는 아키텍처이며 본 연구의 주 타깃 모델이다.

코드 예제

text

S(X) = s · X˜, s = (∥X∥2 / (∥X˜∥2 + ϵ))

activation sparsification에서 비제로 구성의 크기를 원래 활성화의 노름과 동일하게 보정하는 수식.

text

Y = Ys + Yr = S(X) · W⊤ + X · (LALB)⊤

활성화 기반 sparse 경로와 LoRA 보조 경로의 결합으로 최종 출력을 구성하는 수식.

전체 접근 방식: activation sparsity를 이용하여 DiT의 선형 계층에서 2:4 패턴을 적용하고, Norm Compensation으로 활성화 노름 손실을 보정하며 LoRA 분기를 통해 잔차 정보를 보완한다. Selective Layer Skipping으로 단일 스트림 경로의 주요 계층에서 성능 저하를 줄인다.
핵심 메커니즘: Top-K 기반의 2:4 활성화 선택을 수행하고, X˜를 얻은 후 S(X) = s·X˜로 보정한다. Ys = S(X)W⊤, Yr = X(LA LB)⊤, Y = Ys + Yr로 출력한다. 활성화 sparsification은 온라인으로 수행되며, 2:4 패턴에 맞춰 SpTC 형식으로 커널에 직접 전달된다.
구현 및 학습 세부: LoRA 랭크 R은 64로 고정하고, Loss = ||X·W⊤ − (S(X)·W⊤ + X·(LA LB)⊤)||^2를 최소화하도록 LA, LB만 학습한다. Distillation 기반의 지도교사-학생 방식으로 학습하며, 2k 스텝 내에 수렴하는 것을 확인했다. 단일 스트림 경로에서는 Out/OUp 등 특정 레이어를 건너뛰는 전략이 도입되었다.
CUDA 커널 최적화: 온라인 희소화 파이프라인을 하나의 실행 경로로 합성하고, Sparse GEMM 커널은 streamK 스타일의 블록 병렬 파이프라인으로 대역폭-지연 구조를 활용한다. LoRA 분기와의 연동은 온칸에서 Yr를 생성하고 Ys에 누적한다.

RT-Lynx: Diffusion Transformers에서 GEMM Sparsity를 Activation으로 옮겨 실현하는 방법

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

코드 예제

관련 토론

댓글

관련 기사

트랜스포머 성능을 높이는 동적 단기 컨볼루션 기법

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 기사

트랜스포머 성능을 높이는 동적 단기 컨볼루션 기법