[sglang] SGLang에 KDA FlashInfer 백엔드 도입: Blackwell 아키텍처에서의 효율적인 추론

PR 링크: sgl-project/sglang#30113 상태: Merged | 변경: +1077 / -39

들어가며

최근 LLM 추론 가속을 위해 KDA(Kimi Delta Attention)와 같은 선형 어텐션(Linear Attention) 메커니즘이 주목받고 있습니다. 하지만 기존 SGLang의 KDA 구현체는 Speculative Decoding(MTP)을 위한 target_verify 경로가 부재하여, Kimi-Linear와 같은 모델에서 추론 가속 기술을 활용하는 데 제약이 있었습니다. 본 PR은 FlashInfer의 recurrent_kda 커널을 활용하여 Blackwell(SM100) 아키텍처에서 KDA의 decode 및 target_verify를 지원함으로써 이 문제를 해결합니다.

코드 분석

1. kernels/kda_flashinfer.py (신규)

FlashInfer의 고성능 커널을 SGLang의 LinearAttnKernelBase에 래핑합니다. recurrent_kda는 단일 토큰 디코드와 fused spec-decode를 모두 처리할 수 있습니다.

# FlashInferKDAKernel 구현 예시
class FlashInferKDAKernel(LinearAttnKernelBase):
    def decode(self, q, k, v, a, b, A_log, dt_bias, ssm_states, ...):
        # recurrent_kda 호출을 통한 인플레이스 업데이트
        return flashinfer.kda_decode.recurrent_kda(..., use_gate_in_kernel=True)

2. kda_backend.py

KDAKernelDispatcher에 FlashInfer 분기를 추가하여 target_verify 경로를 활성화했습니다. 특히 MTP(Multi-Token Prediction)를 위해 forward_extend 경로를 연결하여, 기존 Triton 구현체에서 발생하던 체크포인트 롤백 문제를 해결했습니다.

# MTP verify 경로 추가
if decode_backend.is_flashinfer():
    self.verify_kernel = FlashInferKDAKernel()

3. mem_cache/memory_pool.py

KDA의 특수한 메모리 레이아웃 요구사항을 반영했습니다. KDA는 transpose 연산이 포함되므로, 기존 GDN의 deduplicated conv-window 레이아웃을 사용하면 데이터가 오염됩니다. 이를 위해 is_kda 가드를 추가하여 dense 레이아웃을 유지하도록 수정했습니다.

왜 이게 좋은가

1. 정확성(Correctness) 확보: 기존 Triton 기반 target_verify는 체크포인트 롤백 로직의 결함으로 인해 end-to-end 추론 시 결과가 퇴화하는 문제가 있었습니다. FlashInfer 백엔드는 시퀀셜 디코드와 동일한 상태 업데이트를 보장하여 gsm8k 벤치마크에서 0.895의 정확도를 달성했습니다.
2. 성능 최적화: SM100 아키텍처에서 recurrent_kda 커널은 B=256 배치 환경에서 기존 Triton 커널 대비 약 1.35배의 속도 향상을 보여줍니다.
3. 일반적 교훈: 특수 연산(KDA)을 위한 커널 최적화 시, 단순히 연산 속도뿐만 아니라 메모리 레이아웃(dense vs deduplicated)이 전체 파이프라인의 정합성에 미치는 영향을 고려해야 합니다. 또한, Speculative Decoding 도입 시 target_verify의 상태 체크포인트가 시퀀셜 디코드와 정확히 일치하는지 검증하는 것이 필수적입니다.

리뷰 피드백 반영

리뷰어들의 의견을 반영하여 speculative_topk > 1인 경우 명시적인 에러를 발생시키도록 KDAAttnBackend에 가드 코드를 추가했습니다. 이는 지원되지 않는 설정으로 인한 런타임 오류를 방지합니다.

참고 자료

• FlashInfer KDA Documentation
• SGLang GitHub Repository

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.