Kimi Delta Attention 연산을 GPU에서 빠르게 처리하는 방법
대규모 언어 모델 추론 시 주의력 연산이 병목이 되는 상황을 겪고 있다면, 이 저장소가 그 시간을 줄여줍니다. FlashKDA는 Kimi Delta Attention(KDA) 알고리즘을 위한 고성능 CUDA 커널 모음입니다. PyTorch 환경에서 기존 라이브러리를 그대로 쓰면서도 내부적으로 이 빠른 커널을 자동으로 호출하도록 설계되었습니다.
v1 설계 배경 공개
이 글은 커널 설계 시 어떤 판단을 내렸는지 설명합니다. 저장소는 CUTLASS 기반의 고성능 KDA 커널을 제공하며, 특히 SM90 이상의 GPU 아키텍처를 대상으로 최적화되었습니다. CUDA 12.9 이상과 PyTorch 2.4 이상 환경에서 동작합니다.
자동 분기되는 설치와 사용
설치는 pip install 명령으로 진행되며, 빌드 시 현재 CUDA 장치를 감지해 해당 아키텍처에 맞게 컴파일합니다. 모든 아키텍처를 명시적으로 컴파일하려면 FLASH_KDA_CUDA_ARCHS=all 환경 변수를 설정해야 합니다. 설치 후 flash-linear-attention 라이브러리의 chunk_kda 함수를 호출하면, FlashKDA가 자동으로 백엔드로 선택됩니다. torch.inference_mode() 컨텍스트 내에서 호출하는 것이 권장되며, FLA_FLASH_KDA=0 환경 변수를 설정하면 기존 Triton 경로로 되돌릴 수 있습니다.
커널 API와 제약 조건
핵심 함수인 flash_kda.fwd는 쿼리, 키, 밸류, 게이트, 베타 등 다양한 텐서를 입력으로 받습니다. 현재는 K와 V의 차원이 128로 고정되어야 합니다. 가변 길이 배치 처리를 위해 cu_seqlens를 전달할 수 있으며, 이 경우 배치 크기 B는 1이어야 합니다. 초기 상태와 최종 상태는 bf16 또는 fp32 텐서로 제공하거나 None으로 설정할 수 있습니다. 상태가 제공될 경우 두 텐서의 데이터 타입이 일치해야 합니다.
정확성 검증과 디버깅
정확성 테스트는 tests/test_fwd.py를 통해 수행되며, PyTorch 참조 구현 및 flash-linear-attention과의 정확 일치 여부를 확인합니다. 성능 벤치마크 결과는 BENCHMARK_H20.md 파일에 기록되어 있습니다. 디버깅 시 logging.basicConfig(level=logging.INFO)를 추가하면 FlashKDA가 선택되었는지 또는 거부된 이유를 로그로 확인할 수 있습니다.
적용 전 점검 포인트
이 저장소는 SM90 이상의 GPU와 CUDA 12.9, PyTorch 2.4를 요구하므로 하드웨어와 소프트웨어 환경을 먼저 점검해야 합니다. K와 V 차원이 128이 아닌 모델에는 적용할 수 없습니다. flash-linear-attention 버전이 0.5.0 이상인지 확인해야 하며, 특정 아키텍처를 타겟팅하는 빌드가 필요한 경우 환경 변수 설정을 잊지 말아야 합니다.
