FlashQLA(Flash Qwen Linear Attention)는 Qwen 팀이 공개한 고성능 선형 어텐션 커널 라이브러리다. TileLang 위에 구축되어 GDN(Gated Delta Network) Chunked Prefill의 순전파·역전파에 합리적인 연산자 퓨전과 성능 최적화를 적용한다. NVIDIA Hopper 아키텍처에서 FLA Triton 커널 대비 순전파 2~3배, 역전파 2배 속도 향상을 달성한다.
GitHub: QwenLM/FlashQLA | MIT 라이선스
배경
Qwen3.5 / Qwen3.6 계열 모델은 선형 어텐션 기반 GDN 아키텍처를 활용한다. 기존 FLA(Flash Linear Attention) Triton 커널은 기능적으로는 충분하지만, 프리트레이닝이나 에지 추론 환경에서 GPU 활용률이 낮은 문제가 있었다. FlashQLA는 이 간극을 메우기 위해 설계됐다.
주요 특징
1. 게이트 기반 자동 내부 컨텍스트 병렬화
GDN 게이트의 지수 감쇠(exponential decay) 특성을 활용해 텐서 병렬(TP), 긴 시퀀스, 소헤드 수 설정에서 카드 내 자동 컨텍스트 병렬(CP)을 활성화한다. GPU SM 활용률이 향상된다.
2. 하드웨어 친화적 대수 재공식화
수치 정밀도를 유지하면서 Tensor Core·CUDA Core·SFU 오버헤드를 효과적으로 줄이도록 GDN Chunked Prefill의 순전파·역전파 흐름을 재공식화했다.
3. TileLang 퓨전 워프 특화 커널
단계별 독립 커널 분해도, 전체 연산 흐름의 단일 커널 퓨전도 아닌 중간 방식을 채택한다. CP·역전파 요구사항을 고려해 핵심 퓨전 커널을 TileLang으로 구현하고, 데이터 이동·Tensor Core 연산·CUDA Core 연산을 중첩(overlap)하도록 워프그룹 특화를 수동 구현했다.
요구사항
- SM90 이상 (NVIDIA Hopper 이상)
- CUDA 12.8 이상
- PyTorch 2.8 이상
설치
git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .사용법
고수준 API
from flash_qla import chunk_gated_delta_rule
o, final_state = chunk_gated_delta_rule(
q=q, # [B, T, H_q, K]
k=k, # [B, T, H_q, K]
v=v, # [B, T, H_v, V]
g=g, # [B, T, H_v]
beta=beta, # [B, T, H_v]
scale=scale,
initial_state=initial_state,
output_final_state=True,
cu_seqlens=cu_seqlens, # 가변 길이 시퀀스 지원
)저수준 API
순전파·역전파를 분리 호출:
from flash_qla import chunk_gated_delta_rule_fwd, chunk_gated_delta_rule_bwd
# 순전파
g, A, o, h, final_state = chunk_gated_delta_rule_fwd(
q, k, v, g, beta, scale=scale, initial_state=h0, cu_seqlens=cu_seqlens
)벤치마크
Qwen3.5 / Qwen3.6 계열의 헤드 설정(h_k,v ∈ {64, 48, 32, 24, 16, 8}, TP1~TP8)에서 FLA Triton 및 FlashInfer baseline 대비 측정 (FLA 0.5.0, Triton 3.5.1, FlashInfer 0.6.9, TileLang 0.1.8):
- 순전파: 다양한 배치 길이에서 단일 커널 지연 측정
- 역전파: 총 토큰 수 대비 단일 업데이트 스텝 지연 측정
H200 기준 상세 결과: benchmark_results_H200.txt
적용 대상
- Qwen3.5 / Qwen3.6 계열 모델을 프리트레이닝 또는 파인튜닝하는 연구자·엔지니어
- GDN 기반 선형 어텐션 모델의 에지 추론 성능을 개선하려는 개발자
- NVIDIA Hopper 이상 GPU 환경에서 커스텀 어텐션 커널을 교체하려는 팀
라이선스
MIT