효율적인 저비트 MXFP 추론을 위한 대각-타일 혼합 정밀도 어텐션

Diagonal-Tiled Mixed-Precision Attention for Efficient Low-Bit MXFP Inference

트랜스포머 기반의 거대 언어 모델(LLM)은 다양한 실제 작업에서 뛰어난 성능을 보여주었지만, 어텐션의 이중 제곱 복잡성과 고정밀 연산의 메모리 대역폭 제한으로 인해 추론 비용이 여전히 매우 높습니다. 본 연구에서는 차세대 GPU 아키텍처의 컴퓨팅 능력을 활용하여, 마이크로 스케일링 부동 소수점(MXFP) 데이터 형식을 사용하는 저비트 혼합 정밀도 어텐션 커널을 제시합니다. 제안하는 대각-타일 혼합 정밀도 어텐션(DMA)은 타일 레벨에서 두 가지 유형의 저비트 연산을 적용하며, Triton을 사용하여 구현된 정교한 퓨즈드 커널입니다. DMA는 하드웨어 레벨의 병렬성과 메모리 효율성을 활용하여 모델 성능을 저하시키지 않으면서 빠르고 효율적인 추론을 가능하게 합니다. NVIDIA B200 GPU에서 수행한 광범위한 실험 결과, 제안하는 커널은 미미한 성능 저하로 생성 품질을 유지하면서 동시에 커널 퓨전을 통해 상당한 속도 향상을 달성합니다. 저희 코드는 https://github.com/yifu-ding/MP-Sparse-Attn 에서 확인할 수 있습니다.

Transformer-based large language models (LLMs) have demonstrated remarkable performance across a wide range of real-world tasks, but their inference cost remains prohibitively high due to the quadratic complexity of attention and the memory bandwidth limitations of high-precision operations. In this work, we present a low-bit mixed-precision attention kernel using the microscaling floating-point (MXFP) data format, utilizing the computing capability on next-generation GPU architectures. Our Diagonal-Tiled Mixed-Precision Attention (DMA) incorporates two kinds of low-bit computation at the tiling-level, and is a delicate fused kernel implemented using Triton, exploiting hardware-level parallelism and memory efficiency to enable fast and efficient inference without compromising model performance. Extensive empirical evaluations on NVIDIA B200 GPUs show that our kernel maintains generation quality with negligible degradation, and meanwhile achieves significant speedup by kernel fusion. We release our code at https://github.com/yifu-ding/MP-Sparse-Attn.