deepseek-ai/FlashMLA

FlashMLA – DeepSeek 모델을 위한 고성능 어텐션 커널

개요 – FlashMLA는 DeepSeek의 대규모 언어 모델(최근 DeepSeek-V4.1-Flash 모델 등)에서 사용되는 고도로 최적화된 어텐션 연산자를 구현한 C++/CUDA(및 Huawei Ascend) 라이브러리입니다. 주요 기능은 다음과 같습니다.

• 희소 어텐션 커널(토큰 수준의 "DeepSeek Sparse Attention"): prefill 단계(모델이 초기 프롬프트를 처리할 때)와 decoding 단계(한 번에 한 토큰씩 생성할 때) 모두에 대응합니다. 이 커널들은 FP8 또는 FP4 양자화 키-값 캐시를 사용하여 정확도를 유지하면서 메모리 대역폭을 획기적으로 줄입니다.
• 밀집 멀티 헤드 어텐션(MHA) 커널: 표준(희소하지 않은) prefill 및 역방향 패스용.
• 퓨전 커널: Q-norm, RoPE(회전 위치 인코딩), 핵심 어텐션 계산, 출력 측의 공액 RoPE, 그리고 최종적으로 FP8로 캐스팅하는 과정을 단일 실행으로 병합합니다. 이를 통해 수많은 작은 커널로 인한 오버헤드를 제거합니다.
• NVIDIA GPU (SM100/SM103, CUDA 13.1+) 및 Huawei Ascend 950 NPU(CANN 9.2+ 및 torch_npu 경유) 지원.

중요성 – 대규모 언어 모델의 추론 시간 대부분은 어텐션 연산에 소요됩니다. FlashMLA는 저정밀도(FP8/FP4)로 어텐션을 재작성하고 희소성을 활용함으로써 최신 하드웨어에서 >1,000 TFlops를 달성하여 장치의 이론적 피크 성능에 근접합니다. 이는 DeepSeek-V4.1과 같은 모델을 서비스할 때 더 빠른 응답 시간과 더 낮은 비용을 의미합니다.

사용 방법

1. 설치 – 리포지토리를 클론하고 서브모듈을 가져온 뒤 pip install -v . --no-build-isolation을 실행합니다. 빌드 스크립트는 CUDA 환경인지 Ascend 환경인지 자동으로 감지하여 적절한 확장 기능을 컴파일합니다.
2. Prefill (희소) – 쿼리 텐서, KV 캐시, 그리고 각 쿼리 토큰이 참조해야 할 KV 토큰을 커널에 알려주는 3-D indices 텐서를 사용하여 flash_mla_sparse_fwd(q, kv, indices, …)를 호출합니다.
3. Decoding (희소) – 먼저 get_mla_metadata()를 통해 스케줄러 메타데이터를 얻은 다음, 생성 단계마다 flash_mla_with_kvcache(...)를 반복 호출합니다. 이 함수는 README에 설명된 FP8/FP4 KV 캐시 레이아웃을 사용합니다.
4. 밀집 MHA – 널리 사용되는 flash_attn 패키지와 호환되는 flash_attn_varlen_func / flash_attn_varlen_qkvpacked_func API를 사용하되, 추가적인 head_dim_qk 인수가 필요하며 반환 값에 출력과 log-sum-exp(lse)가 포함된다는 점에 유의하십시오.
5. 퓨전 커널 – Q-projection 및 Wv-projection 가중치를 퓨전 커널에 필요한 레이아웃으로 정렬한 후(도우미 함수는 fused_norm_rope_attn_rope_cast.permute_*에 있음), prefill(...) 또는 decode(...)를 호출하여 후속 FP8-einsum을 위한 FP8 출력을 얻습니다.

성능 하이라이트(리포지토리 보고서 기준)

• 희소 prefill: NVIDIA B200 GPU (CUDA 13.3)에서 최대 1350 TFlops. Ascend 950 NPU에서 410 TFlops(하드웨어 피크의 약 95%).
• 희소 decoding: B200에서 최대 1024 TFlops, Ascend 950에서 360 TFlops(피크의 약 83%).
• 퓨전 norm-RoPE-attn-RoPE-cast: B200에서 최대 1460 TFlops(prefill) 및 950 TFlops(decoding)로, 비퓨전 버전 대비 10~15% 향상.
• 밀집 MHA: B200에서 순방향 최대 1460 TFlops, 역방향 최대 1000 TFlops.

기술적 세부 사항

• KV 캐시 형식 – FP8(float8_e4m3) 또는 FP4(e2m1) 양자화 캐시만 지원합니다. 레이아웃은 원시 양자화 값과 토큰별 스케일 팩터를 저장하여 실시간으로 빠른 역양자화를 가능하게 합니다.
• Indices 텐서 – 각 쿼리 토큰이 참조해야 할 KV "페이지"를 인코딩하는 3-D 정수 텐서입니다. 유효하지 않은 항목은 -1로 표시되어 커널에서 무시됩니다.
• 레지스터 스필 안전성 – CUDA 확장 빌드 후 빌드 시스템은 레지스터 스필을 확인합니다. FLASH_MLA_SKIP_REG_SPILL_CHECK=1을 설정하지 않으면 빌드가 실패합니다.
• 호환성 참고 – 2026년 9월 릴리스에서는 NVIDIA Hopper GPU 및 이전 DeepSeek 모델(V3/V3.2/V4.0)에 대한 지원이 중단되었으며 KV 캐시 형식이 변경되었습니다. 이전 구성이 필요한 사용자는 README에 링크된 레거시 커밋을 체크아웃해야 합니다.

추가 정보 – 리포지토리에는 알고리즘과 저수준 최적화를 설명하는 상세한 markdown 보고서(docs/…)와 하드웨어 성능을 검증하기 위한 테스트 스크립트(tests/*.py)가 포함되어 있습니다.

---

FlashMLA는 최첨단 NVIDIA GPU와 Huawei Ascend NPU를 모두 타겟팅하는 DeepSeek의 LLM 추론 파이프라인에서 성능이 중요한 빌딩 블록입니다.