값 토큰 재배열과 E4M3 직접 인코딩으로 비디오 DiT 어텐션을 가속하는 VC-Attention

무엇인가

이 논문은 비디오 Diffusion Transformer(DiT)의 자기어텐션을 저비트로 실행할 때 생기는 두 문제를 다룬다. 하나는 정확도다. 저비트 양자화는 하드웨어 블록마다 하나의 스케일을 공유하는데, 블록 안의 소수 이상치가 스케일을 결정해 나머지 값들이 좁은 표현 범위에 갇힌다. 기존 연구는 쿼리와 키를 스무딩했지만, 값(value)의 이상치는 고정된 채널이나 시공간 구조를 따르지 않아 출력 오차의 주원인으로 남는다. Wan2.2에서는 키 스무딩과 회전을 적용한 뒤에도 식 O - O_q = (P - P_q)V + P_q(V - V_q)의 두 번째 항이 출력 오차의 82%를 차지한다. 다른 하나는 속도다. 저비트 Tensor Core는 QK와 PV 두 행렬곱만 가속하므로, 그 사이의 고정밀 softmax, 특히 FP32 지수 연산과 FP32에서 E4M3로의 변환이 데이터센터 GPU에서 가장 긴 파이프라인 단계가 된다. 5초 720p Wan2.2-14B는 약 70K 토큰을 만들고 RTX 5090에서 어텐션이 생성 시간의 64% 이상을 차지한다. MiniMax-H3 1344×768에서는 SageAttention2가 기준 영상에서 벗어나면서도 B200에서 BF16 FlashAttention-4의 0.29배 속도에 그친다.

어떻게 동작하나

제안 방법은 학습이 필요 없는 VC-Attention이며, V-Smooth와 ExpCast-FP8 두 요소로 구성된다. V-Smooth는 각 배치와 헤드에서 값 토큰들에 대해 가벼운 온라인 k-means를 수행해 레이블 z_t를 얻고, π = argsort(z)로 키와 값을 함께 재배열한다. 키를 재배열하면 확률 행렬 P의 열이 값 행렬 V의 행과 같은 순서로 바뀌므로 P'V' = PV가 되어 비인과 자기어텐션의 출력은 변하지 않는다. 그다음 하드웨어 블록 크기 B_v = 128행마다 블록 평균 μ_j를 빼고 잔차 R_j = V'_j - 1 μ_j^T만 양자화한다. 평균은 온라인 softmax가 이미 유지하는 행 합 r_ij를 이용해 A_i ← α_i A_i + P~_{q,ij} R_{q,j} + r_ij μ_j^T로 복원되므로 V를 다시 읽는 패스나 추가 버퍼가 필요 없다. 블록 평균이 제거하는 에너지 비율은 Wan2.2 100개 헤드 평균으로 원래 순서에서 8%, DeltaQuant식 고정 큐브에서 12%, 정렬 후 36%다. 평균 자체는 블록당 16비트 벡터 하나로 값 원소당 0.125비트에 불과하다. 그룹화는 모든 디노이징 스텝이 아니라 처음 25% 스텝에서만 수행하고, 그룹화가 들어간 스텝은 어텐션 시간을 30% 늘리지만 전체 스텝 평균으로는 3~4% 비용으로 상각된다.

무엇과 다른가

ExpCast-FP8은 softmax 단계를 줄이기 위해 로그 영역 점수에서 E4M3 확률 코드를 직접 만든다. E4M3 바이트는 지수 필드 e와 가수 필드 m으로 v = 2^{e-7}(1+m/8)을 나타내므로, 정수로 읽은 8e+m은 8 log2 v + 56 + ε(m)에 해당한다. 온라인 softmax가 이미 유지하는 u = (s - m_i') log2 e ≤ 0에 대해 행 최댓값이 256이 되도록 2^8을 곱하면, 코드는 c(u) = clip_[0,120](Round(8(u+8)+56+β)), β = -0.35로 계산된다. 여기서 8과 56은 E4M3 인코딩에서 바로 나오고 β는 남는 항 ε(m)을 최소최대 기준으로 중앙화한 값 -0.3443을 쓴 것으로, 모델별 피팅이 아니다. 이렇게 하면 FP32 지수 연산과 FP32에서 E4M3로의 캐스트를 하나의 융합 곱셈-덧셈과 정수 변환으로 대체한다. 원소별 오차는 최대 한 코드, 상대오차 7.5%까지 가능하지만 정규화된 행에서는 가수 비트에만 의존해 대부분 상쇄된다. Proposition 3.1은 정상 E4M3 범위 u_k ≥ -14인 행에서 총변동거리 TV(p, p_hat) < 3.64%, ||o_hat - o||_2 ≤ 3.64% diam_2(V')를 보장하고, 언더플로 꼬리의 정규화 질량을 τ라 하면 우변은 0.0364 + τ가 된다. Wan2.2 100개 헤드의 204.8K 어텐션 행에서 총변동은 평균 1.6%, 최대 17.8%였고 언더플로가 없는 행은 1.4% 아래였다. 같은 행에서 FP32 지수 후 E4M3 캐스트는 평균 1.1%였다.

어떻게 쓰나

실험은 Wan2.2-T2V-A14B 720p, LongCat-Video 480p, HunyuanVideo-1.5 720p, MiniMax-H3 1344×768 네 모델에서 수행했다. 각 모델마다 MovieGen Bench 프롬프트로 100개 비디오를 생성하고 프롬프트와 시드를 모든 방법에 공유했다. 충실도는 같은 모델의 BF16 출력 대비 PSNR, SSIM, LPIPS로 재고, VBench의 subject consistency와 imaging quality도 함께 본다. 속도는 BF16 FlashAttention-4를 기준으로 정규화한 어텐션 속도 향상과 클립 하나의 end-to-end 생성 시간으로 측정한다. 8비트에서는 SageAttention2, 스무딩 없는 8비트 FlashAttention-4, Attn-QAT와 비교하고, 4비트에서는 SageAttention3와 비교한다. 결과적으로 V-Smooth는 두 정밀도 모두에서 가장 충실한 저비트 어텐션이다. 8비트에서 V-Smooth는 모든 모델의 모든 열에서 앞서며 SageAttention2 대비 PSNR을 Wan2.2에서 2.3dB, HunyuanVideo-1.5에서 2.8dB 높이고 LPIPS를 13~29% 줄였다. ExpCast-FP8을 융합하면 이 PSNR 이득 중 0.7~2.1dB를 속도와 맞바꾸지만, 결합 커널도 네 모델 모두에서 SageAttention2를 앞선다. 4비트에서는 ExpCast-FP8 없이 V-Smooth만으로 SageAttention3 대비 Wan2.2에서 2.9dB, LongCat-Video에서 3.6dB PSNR을 높이고 LPIPS를 최대 41% 줄였다. VBench 두 지표에서는 모든 학습 불필요 방법이 BF16 모델과 0.01 이내에 머물러 방법 간 차이를 구분하지 못했다. QK에 Hadamard 회전을 추가해도 PSNR 변화가 SageAttention3에서 0.1dB, 8비트 FlashAttention-4에서 0.3dB 이내여서 남은 오차는 QK가 아니라 값 연산에 있음이 확인된다. Attn-QAT는 학습 없이 실행하면 모든 모델에서 SageAttention2보다 3.4~6.7dB 낮은 PSNR을 보였고 VBench 점수를 흔드는 유일한 방법이었다.

전제와 한계

효율 수치는 하드웨어별로 분명히 제시된다. Wan2.2에서 B200은 BF16 FlashAttention-4 대비 어텐션 1.59배, SageAttention2 대비 6.02배 빠르다. H200은 BF16 FlashAttention-4 대비 1.46배, SageAttention2 대비 1.16배다. 클립 하나의 end-to-end 생성은 B200에서 1.19배, H200에서 1.13배 빨라진다. 워크스테이션 Blackwell의 4비트 구성에서는 V-Smooth가 RTX PRO 6000에서 어텐션 2.27배, RTX 5090에서 3.58배 빠르며, end-to-end로는 각각 1.36배, 1.70배다. 두 4비트 커널의 비용이 같으므로 이들은 Table 2의 충실도로만 구분된다. B300에서는 INT8 QK 행렬곱이 없어 QK와 PV를 모두 FP8로 유지하는데, BF16 FlashAttention-4 대비 1.47배 빠르고 같은 정밀도의 순진한 FP8 커널의 1.31배를 앞선다. MiniMax-H3에서도 18.4dB로 순진한 커널의 17.1dB보다 높다. 전처리 융합 실험에서는 양자화기와 Hadamard 회전 융합이 1.59배, 두 스무딩 평균과 순열 gather가 1.43배, rotary embedding이 1.94배를 더하고, 다섯 패스와 그룹화를 손으로 작성한 커널이 1.97배를 더해 전체 8.74배가 된다. 그룹화는 그때 어텐션 호출당 전체 체인 비용 4.8ms의 29%를 차지한다. 그룹화 위치에 대한 절제 실험에서는 처음 25% 스텝에 그룹화하는 것이 모든 스텝 그룹화보다 PSNR 0.5dB를 포기하면서 SSIM과 LPIPS는 더 좋았고, 같은 양을 균등하게 분산하면 2.6dB 뒤졌다. 이 창은 클립 하나에서 16초를 아낀다(351.0초 대 367.0초).

개발자 입장에서 이 논문은 영상 DiT 추론 최적화에서 재학습 없이 커널 수준으로 적용할 수 있는 선택지를 제시한다. 구현은 CuTe/CUDA로 FlashAttention-4/SageAttention 커널을 제자리 수정하는 방식이라, 기존 어텐션 파이프라인에 V-Smooth와 ExpCast-FP8을 끼워 넣는 형태로 이해할 수 있다. 8비트 데이터센터 구성은 B200과 H200에서 두 메커니즘을 모두 쓰고, 4비트 워크스테이션 구성은 RTX PRO 6000과 RTX 5090에서 V-Smooth만 쓴다. sparse attention이나 양자화된 선형 계층과도 조합할 수 있고 어텐션 패턴과 샘플링 스케줄은 바꾸지 않는다. 다만 정확도와 속도의 맞바꿈이 있다. ExpCast-FP8은 PSNR을 0.7~2.1dB 내주는 대신 속도를 얻으므로, 충실도가 최우선이면 V-Smooth 단독 구성이나 ExpCast를 끈 8비트 구성을 검토해야 한다. 또한 그룹화는 처음 25% 디노이징 스텝에만 넣는 스케줄과 순열 재사용을 전제로 하므로, 실제 파이프라인에서 디노이징 스텝 수와 스케줄을 맞춰 비용을 확인해야 한다.

저자들이 밝힌 한계와 전제도 분명하다. 4비트 데이터센터 구성은 평가하지 않았는데, on-the-fly NVFP4 P가 16원소마다 스케일을 계산해 softmax 임계 경로에 올라가고 ExpCast-FP8이 이를 제거하지 못하기 때문이다. ExpCast-FP8의 보장은 u_k ≥ -14인 정상 E4M3 범위를 전제로 하며, 언더플로 꼬리가 있으면 오차 한계에 τ가 더해진다. 실제 측정에서도 최대 총변동 17.8%인 행이 있었다. V-Smooth는 그룹화 스텝에서 어텐션 시간을 30% 늘리고, 전체 스텝 평균으로 3~4% 비용이 남는다. 처음 25% 스텝만 그룹화하면 모든 스텝 그룹화보다 PSNR 0.5dB를 포기한다. 클러스터를 정확히 B_v 토큰으로 균형 맞추면 오차를 8.5% 더 줄일 수 있지만 그룹화 비용이 2.8배, 배포 형상에서 어텐션 시간의 85%가 들어 plain k-means를 운용점으로 택했다. 값 이상치가 입력마다 달라지는 문제를 온라인 클러스터링으로 다루는 만큼, 모델별 재학습은 필요 없지만 입력 분포와 스텝 스케줄에 따른 그룹화 비용과 정확도 변동은 배포 전에 확인해야 한다.