[flashinfer] Blackwell 아키텍처를 위한 MLA 디코드 커널 최적화: FlashInfer의 성능 개선 분석

PR 링크: flashinfer-ai/flashinfer#8211 상태: Merged | 변경: +532 / -321

들어가며

LLM 추론 성능의 핵심 중 하나는 Attention 메커니즘의 효율성입니다. 특히 DeepSeek 모델 등에서 사용되는 MLA(Multi-head Latent Attention)는 메모리 대역폭을 절약하면서도 높은 성능을 내기 위해 설계되었습니다. 하지만 새로운 하드웨어 아키텍처가 등장할 때마다, 해당 하드웨어의 특성(SM 개수, 메모리 계층 구조, 명령어 세트 등)에 맞춘 세밀한 튜닝이 필수적입니다.

이번 글에서는 flashinfer-ai/flashinfer 레포지토리에 반영된 Blackwell 아키텍처(sm_100a, sm_103a) 전용 MLA 디코드 커널(cake_trtllm_mla_blackwell)의 최적화 사례를 분석합니다. 이 PR은 특히 BF16 value-quarter split 프로그램과 GB300(152-SM)을 위한 CLC(Continuous Latency Compensation) 프로그램의 성능을 개선하는 데 초점을 맞추고 있습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

이번 최적화의 핵심은 하드웨어 리소스를 더 촘촘하게 활용하기 위한 소프트웨어 파이프라이닝(Software Pipelining)의 강화와 공유 메모리(Shared Memory, SMEM) 레이아웃의 재설계입니다.

1. 파이프라인 스테이지 확장 (Latency Hiding)

가장 눈에 띄는 변화는 KV(Key-Value) 데이터를 처리하는 파이프라인 스테이지의 개수 증가입니다.

Before:

#define NUM_KV_PIPE_STAGES 15

After:

#define NUM_KV_PIPE_STAGES 17

설명: Blackwell 아키텍처는 이전 세대보다 더 많은 SM(Streaming Multiprocessor)과 향상된 메모리 서브시스템을 갖추고 있습니다. NUM_KV_PIPE_STAGES를 15에서 17로 늘린 것은, 메모리로부터 데이터를 가져오는 동안 연산 유닛이 쉬지 않도록 더 많은 데이터를 미리 버퍼링(Prefetching)하겠다는 의도입니다. 이는 특히 대규모 배치 사이즈나 긴 시퀀스 길이를 처리할 때 메모리 지연 시간(Latency)을 효과적으로 숨기는(Hiding) 역할을 합니다.

2. 공유 메모리(SMEM) 레이아웃 및 오프셋 조정

파이프라인 스테이지가 늘어남에 따라, 각 데이터를 저장할 SMEM의 위치도 재조정되었습니다.

Before:

#define SMEM_SMEM_P_OFF 196608
// ...
__nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 196608);
const int smem_p_addr = smem + 196608;

After:

#define SMEM_SMEM_P_OFF 212992
// ...
__nv_bfloat16* smem_p = reinterpret_cast<__nv_bfloat16*>(smem_raw + 212992);
const int smem_p_addr = smem + 212992;

설명: SMEM_SMEM_P_OFF(Partial P matrix의 오프셋)가 뒤로 밀려났습니다. 이는 앞선 KV 파이프라인 스테이지들이 차지하는 공간이 늘어났기 때문입니다. Blackwell은 SM당 가용한 SMEM 용량이 크기 때문에, 이러한 공격적인 스테이지 확장이 가능합니다. 정확한 주소 계산은 tcgen05(Blackwell의 Tensor Core 명령어)와 같은 하드웨어 가속 기능을 사용할 때 정렬(Alignment) 문제를 피하기 위해 매우 중요합니다.

3. Mbarrier 초기화 로직 고도화

비동기 데이터 전송을 관리하는 mbarrier의 개수와 초기화 카운트도 변경되었습니다.

Before:

// Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 56 barriers)
// ...
if (lane < 24) {
    mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32);
}

After:

// Mbarrier init (18 pipeline groups, 0 ordered-sequence groups, 58 barriers)
// ...
if (lane < 26) {
    mbarrier_init(smem + 230688 + lane * 8, _mbarrier_init_count_0_32);
}

설명: 배리어의 개수가 56개에서 58개로 늘어났고, 이를 초기화하는 스레드(lane)의 범위도 확장되었습니다. 이는 추가된 파이프라인 스테이지 간의 동기화를 정밀하게 제어하기 위함입니다. 특히 Blackwell의 비동기 복사(Async Copy) 기능을 최대한 활용하기 위해서는 각 스테이지의 완료 상태를 추적하는 배리어 관리가 성능의 병목이 되지 않아야 합니다.

왜 이게 좋은 최적화인가?

성능 수치로 증명된 결과

PR 설명에 포함된 벤치마크 결과는 이 변경사항의 가치를 명확히 보여줍니다.

• NVIDIA B200 (sm_100a): 측정된 10개의 Shape 중 8개에서 성능 향상이 관찰되었습니다. 최대 1.0299x의 스피드업을 기록했습니다.
• NVIDIA GB300 (sm_103a): 22개의 Shape 중 11개에서 성능이 향상되었으며, 특히 대규모 배치(batch=512, kv_len=4096) 상황에서 1.0299x의 유의미한 개선을 보였습니다.

일반적인 교훈

1. 하드웨어 특화 튜닝의 중요성: 범용적인 커널보다 특정 아키텍처(Blackwell)의 SM 수와 메모리 특성에 맞춘 파이프라인 깊이 조절이 실제 성능 향상으로 이어집니다.
2. 리소스 트레이드오프: 공유 메모리를 더 많이 사용하는 대신 파이프라인 스테이지를 늘려 Latency를 숨기는 전략은 최신 고성능 GPU에서 매우 유효합니다.
3. 정밀한 동기화 제어: mbarrier와 같은 저수준 동기화 프리미티브를 최적화된 파이프라인에 맞춰 재설계하는 것이 복잡한 커널의 안정성과 성능을 동시에 잡는 길입니다.

결론

이번 PR은 Blackwell GPU의 잠재력을 끌어내기 위해 커널의 내부 파이프라인 구조를 과감하게 수정한 사례입니다. 단순한 코드 정리를 넘어, 하드웨어의 한계치까지 리소스를 할당하고 동기화 로직을 재설계함으로써 실제 추론 환경에서 중요한 성능 이득을 얻어냈습니다. FlashInfer와 같은 고성능 라이브러리가 어떻게 최신 GPU에 빠르게 대응하는지 보여주는 좋은 예시라고 할 수 있습니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-barrier
• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-asynchronous-copy

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