QwenLM/FlashQLA

무엇인가

FlashQLA는 Qwen 스타일 대규모 언어 모델에서 사용되는 Gated‑Delta‑Rule (GDN) 선형 어텐션 알고리즘을 구현하는 고성능 커널 라이브러리입니다. TileLang DSL 위에 작성되었으며 Python 패키지(flash‑qla)로 배포됩니다. 이 라이브러리는 편리한 고수준 함수(chunk_gated_delta_rule)와 PyTorch 코드에서 호출할 수 있는 저수준 순전파/역전파 커널을 모두 제공합니다.

왜 중요한가

선형 어텐션 메커니즘은 기존 소프트맥스 어텐션의 2차 비용을 시퀀스 길이에 따라 선형적으로 확장되는 계산으로 대체하여 매우 긴 컨텍스트를 가능하게 합니다. FlashQLA는 참조 Flash Linear Attention (FLA) Triton 커널과 비교하여 특히 NVIDIA Hopper/Blackwell GPU에서 순전파 속도를 2‑3배, 역전파 속도를 약 2배 향상시킵니다. 더 빠른 커널은 Qwen 스타일 모델에 의존하는 추론 중심 또는 엣지 사이드 에이전트의 사전 학습 비용을 절감하고 추론 지연 시간을 낮춥니다.

핵심 기능

기능	설명
게이트 기반 자동 카드 내 컨텍스트 병렬화	GDN 게이트의 지수 감쇠가 작업을 여러 CUDA 워프/SM으로 분할할 수 있게 하는 시점을 감지하여 단일 GPU 카드 내에서 자동으로 병렬화를 활성화합니다 (TP, 긴 시퀀스, 소형 헤드 설정).
하드웨어 친화적 대수 재구성	GDN 순전파/역전파 방정식을 재작성하여 수치 정밀도를 유지하면서 Tensor‑Core, CUDA‑Core 및 SFU 사용량을 줄입니다.
TileLang 융합 워프 전용 커널	TileLang DSL을 사용하여 데이터 이동, Tensor‑Core 행렬 곱셈 및 CUDA‑Core 연산의 균형을 맞추는 단일 모놀리식 커널 대신 몇 개의 융합 커널을 생성합니다.
이중 CP (카드 간/카드 내 2단계 컨텍스트 병렬화) (v0.1.3)	GPU 간(카드 간) 및 GPU 내(카드 내)에 두 번째 병렬화 수준을 추가합니다.
SM120 역전파 경로 (v0.1.3)	역전파 커널에서 최신 NVIDIA SM120 아키텍처 지원을 추가합니다.
flash‑linear‑attention용 플러그 앤 플레이 백엔드 (v0.1.2)	표준 FLA API를 통해 FLA GDN 커널의 드롭인 대체품으로 사용할 수 있습니다.

시작하는 방법

# 사전 빌드된 휠 설치 (CUDA 12.8+, PyTorch 2.8+ 필요)
pip install flash-qla

또는 소스에서 빌드:

git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .

간단한 사용 예시

import torch
from flash_qla import chunk_gated_delta_rule

# q, k, v, g, beta는 README에 설명된 형태의 텐서입니다
out, final_state = chunk_gated_delta_rule(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=scale,
    initial_state=init_state,
    output_final_state=True,
    cu_seqlens=cu_seqlens,
)

사용자 정의 학습 루프의 경우 별도의 순전파 및 역전파 함수(*_fwd / *_bwd)를 호출하고 반환된 캐시 객체를 역전파 패스에 다시 전달할 수 있습니다.

사용 시기

GDN 선형 어텐션을 사용하는 Qwen 스타일 모델을 학습하거나 미세 조정하는 경우.
워크로드가 최신 NVIDIA GPU(H100, H200, Blackwell)에서 실행되고 긴 시퀀스 사전 학습 또는 저지연 추론을 위해 최대 속도가 필요한 경우.
이미 PyTorch를 사용하고 있으며 기존 flash-linear-attention API와 호환되는 드롭인 커널을 원하는 경우.

제한 사항 / 미해결 지점

SM 90‑SM 121 및 CUDA 12.8+를 갖춘 NVIDIA GPU만 지원합니다. CPU 또는 AMD는 지원하지 않습니다.
라이브러리는 GDN Chunked Prefill 변형에 중점을 둡니다. 다른 선형 어텐션 공식은 다루지 않습니다.
역전파 속도 향상은 이중 CP 구현에 의존합니다. 오래된 GPU에서는 이점이 적을 수 있습니다.
TileLang 런타임이 필요합니다. 사용자는 TileLang 생성 커널 컴파일에 익숙해야 합니다.

커뮤니티 및 유지 관리

프로젝트는 QwenLM 팀이 유지 관리합니다(인용 목록 참조). 최근 릴리스(v0.1.1‑v0.1.3)는 활발한 개발을 보여줍니다.
이슈와 풀 리퀘스트는 GitHub 저장소에 제출할 수 있습니다. README는 개발자를 위한 벤치마크 스크립트 및 프로파일링 유틸리티 링크를 제공합니다.

위의 모든 정보는 저장소의 README에서 직접 가져온 것이며 외부 추측은 추가되지 않았습니다.