[flashinfer] FlashInfer의 DeepSeek V4 Sparse MLA 최적화 분석

PR 링크: flashinfer-ai/flashinfer#3269 상태: Merged | 변경: +1557 / -103

들어가며

최근 LLM 추론 가속기인 FlashInfer에 DeepSeek V4의 Sparse MLA(Multi-Head Latent Attention)를 지원하기 위한 커널 업데이트가 반영되었습니다. 이번 PR은 특히 BF16 및 FP8 데이터 타입에서의 디코드(decode) 성능을 최적화하고, 가변적인 Top-K 길이와 슬라이딩 윈도우(SWA) KV 캐시를 효율적으로 처리하는 데 초점을 맞추고 있습니다. 본 글에서는 이 최적화가 어떻게 구현되었는지 코드 수준에서 살펴봅니다.

코드 분석

1. csrc/fmhaReduction.cu: 가변 Sparse MLA 지원을 위한 커널 구조 개선

핵심 변경 사항은 fmhaReductionKernel의 유연성 확보입니다. 기존에는 고정된 타일 사이즈와 로직을 사용했으나, 이제는 가변적인 Sparse Top-K 길이를 지원하도록 로직이 확장되었습니다.

Before:

// 기존에는 sparseMla 플래그로 단순 분기 처리
if (sparseMla) {
  seqLenKv = min(seqLenKv, params.mSparseAttnTopK);
}

After:

// 가변적인 Top-K 길이를 지원하기 위해 포인터를 통한 동적 접근
if (supportsVarSparseMlaTopKLens) {
  seqLenKv = params.ptrSparseMlaTopKLens[seqOffsetQ + ctaIdxQ * params.mNumTokensPerCtaQ];
} else if (isTokenSparse) {
  seqLenKv = min(seqLenKv, params.mSparseAttnTopK);
}

이 변경은 각 토큰마다 다른 Sparse 패턴을 가지는 DeepSeek V4의 특성을 반영하기 위함입니다. 또한 SELECT_FMHA_REDUCTION_KERNEL 매크로를 개선하여 TileSizePerCtaQ를 런타임에 선택할 수 있도록 하여, 다양한 시퀀스 길이와 헤드 차원에 대응할 수 있게 되었습니다.

2. 메모리 오프셋 및 인덱싱 최적화

groupsTokensHeadsQ 플래그를 도입하여 토큰과 헤드를 그룹화하여 처리할 때의 메모리 접근 방식을 최적화했습니다.

if (groupsTokensHeadsQ) {
  int32_t tokenIdx{validRowIdx / params.mNumHeadsQPerKv};
  int32_t headIdxInGrp{validRowIdx % params.mNumHeadsQPerKv};
  localHeadIdxO = headIdxInGrp;
  softmaxStatsRowIdx = tokenIdx * params.mNumHeadsQ + headIdxInGrp;
  gmemStoreOffset = int64_t(softmaxStatsRowIdx) * headDimV + headDimIdx;
}

이러한 인덱싱 개선은 GPU의 메모리 병합(Coalesced Access) 효율을 높여, 대규모 모델 추론 시 발생하는 메모리 대역폭 병목을 완화합니다.

왜 이게 좋은가

1. 유연성(Flexibility): 기존의 고정된 Sparse 설정에서 벗어나, ptrSparseMlaTopKLens를 통해 토큰별로 최적화된 어텐션 범위를 지정할 수 있게 되었습니다. 이는 모델의 정확도를 유지하면서도 불필요한 연산을 획기적으로 줄여줍니다.
2. 성능(Performance): TileSizePerCtaQ를 64와 128로 동적으로 선택할 수 있게 함으로써, 하드웨어 리소스 활용도를 극대화했습니다. 테스트 결과 수만 건의 케이스를 통과하며 안정성을 입증했습니다.
3. 확장성: groupsTokensHeadsQ 로직 추가는 향후 더 복잡한 어텐션 패턴이 도입되더라도 커널을 재작성할 필요 없이 대응 가능한 구조를 제공합니다.

리뷰어 피드백 반영

코드 리뷰 과정에서 qsang-nv는 _check_dsv4_sparse_mla_inputs 함수 내의 불필요한 D2H(Device-to-Host) 동기화 문제를 지적했습니다. .item() 호출은 GPU 연산 파이프라인을 멈추게 하므로, 이를 적절한 게이트 조건으로 감싸 성능 저하를 방지하도록 수정되었습니다. 또한, 테스트 커버리지를 높이기 위해 is_varlen=False 케이스를 추가하여 엣지 케이스에서의 안정성을 확보했습니다.

결론

이번 PR은 단순한 기능 추가를 넘어, FlashInfer의 커널이 어떻게 동적이고 복잡한 어텐션 패턴을 효율적으로 처리할 수 있는지 보여주는 좋은 사례입니다. 특히 메모리 접근 패턴의 최적화와 유연한 커널 선택 로직은 고성능 LLM 추론 엔진 개발의 핵심임을 다시 한번 확인시켜 줍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#coalesced-access-to-global-memory

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