[flashinfer] NVIDIA Blackwell(SM100) 최적화: FlashInfer의 Cake FMHA 성능 개선 분석

PR 링크: flashinfer-ai/flashinfer#5855 상태: Merged | 변경: +2025 / -1716

들어가며

LLM 추론 가속화 라이브러리인 FlashInfer가 NVIDIA의 최신 아키텍처인 Blackwell(SM100/SM103)을 위한 네 번째 최적화 라운드(Round 4)를 공개했습니다. 이번 업데이트의 핵심은 Cake DCP(Speculative Decoding)와 Packed-MTP(Multi-Token Prediction) 커널의 성능을 쥐어짜는 것입니다.

기존 Round 3 커널에서도 준수한 성능을 보여주었으나, 두 가지 주요 과제가 남아 있었습니다. 첫째는 FP8(E4M3) 캐시를 사용하는 환경에서 Blackwell의 하드웨어 잠재력을 완전히 활용하기 위한 연산 루프의 최적화였고, 둘째는 분산 처리된 청크(Chunk)를 합치는 과정에서 발생하는 비결정성(Non-determinism) 문제였습니다. 이번 PR은 이 두 문제를 해결하며 특정 워크로드에서 최대 18%의 성능 향상을 달성했습니다.

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

1. 벡터화 및 루프 언롤링을 통한 처리량 극대화

가장 눈에 띄는 변화는 kernel_cake_fmha_dcp_spec_bf16_fp8_balanced 커널의 reduction 로직입니다. 기존에는 비교적 작은 단위의 레지스터 배열을 사용해 루프를 돌았으나, 새로운 버전에서는 더 넓은 범위의 벡터화와 공격적인 언롤링(Unrolling)을 적용했습니다.

Before (Round 3):

// 단순한 4개 단위의 누적 및 루프 구조
float acc_f[4];
float out4[4];
#pragma unroll 1
for (int g_r = 0; g_r < n_groups_r; g_r++) {
    acc_f[0] = 0.0f; // ... 초기화
    // ... 중략 ...
    #pragma unroll 8
    for (int c_m = 0; c_m < n_pad_r; c_m++) {
        // 청크별로 하나씩 로드 및 FMA 연산
        float4 _v4 = *reinterpret_cast<const float4*>(partial_o + (o_col_r + c_c * o_stride_r) + 0);
        #pragma unroll
        for (int k = 0; k < 4; k++) {
            acc_f[k] = __fmaf_rn(_vec_load_0[k], b_k, acc_f[k] * a_k);
        }
    }
}

After (Round 4):

// 더 넓은 누적기(16개)와 4개 청크 동시 처리(Batching)
float acc_w[16];
// ... 초기화 ...
#pragma unroll 1
for (int c_b = 0; c_b < n_pad_w; c_b += 4) {
    float m_b[4];
    float l_b[4];
    float o_b[64]; // 16개 요소 * 4개 청크를 한꺼번에 로드
    #pragma unroll
    for (int cb = 0; cb < 4; cb++) {
        // 4개의 청크 데이터를 한 번에 준비하여 메모리 대역폭 활용 극대화
        m_b[cb] = partial_stats[stats_row_r + c_c * stats_stride_r];
        // ... o_b에 데이터 로드 ...
    }
    // 이후 한꺼번에 FMA 연산 수행
}

왜 좋은 최적화인가? Blackwell 아키텍처는 더 높은 메모리 대역폭과 연산 유닛을 가지고 있습니다. 기존처럼 청크를 하나씩 처리하면 명령어 사이의 의존성(Dependency) 때문에 파이프라인이 놀게 됩니다. 변경된 코드는 4개의 청크를 하나의 배치로 묶어 64개의 float 요소를 한꺼번에 다룹니다. 이는 ILP(Instruction Level Parallelism)를 높여 GPU의 연산 유닛 점유율을 극대화합니다.

2. 결정론적(Deterministic) Reduction Fold

MTP(Multi-Token Prediction) 커널에서 중요한 개선 사항은 결과의 재현성입니다. 기존에는 여러 CTA(Cooperative Thread Array)가 공유 메모리에 결과를 쓸 때, 먼저 도착하는 순서대로 연산이 이루어지는 "In-place fold" 방식을 사용했습니다. 부동 소수점 연산은 결합 법칙이 성립하지 않기 때문에, 도착 순서가 바뀌면 결과값에 미세한 차이(1-11 ulp)가 발생했습니다.

개선 내용: 이번 PR에서는 deterministic 두-chunk fold를 도입했습니다. 도착 순서가 아니라 청크의 인덱스(Index)에 따라 연산 순서를 고정했습니다.

• 변경 전: atomicAdd 또는 먼저 도착한 스레드가 중간 결과를 업데이트 (비결정적).
• 변경 후: 각 청크가 자신의 고정된 슬롯을 가지고, 정해진 순서대로 FMA(Fused Multiply-Add) 연산을 수행.

이를 통해 수치적 안정성을 확보했으며, 동일한 입력에 대해 항상 비트 단위로 동일한(Bit-identical) 출력을 보장하게 되었습니다.

성능 수치 분석

이번 최적화의 결과는 놀랍습니다. NVIDIA GB300 및 B200 GPU에서 측정된 지표는 다음과 같습니다.

Engine	GPU	Speedup (Geomean)	Max Speedup
cake_fmha_dcp_spec_bf16_fp8_balanced	GB300	1.075x	1.161x
cake_fmha_dcp_spec_bf16_fp8_d256_balanced	B200	1.062x	1.161x

특히 head_dim=256인 모델의 특정 워크로드(prod_d256_b1_ctx32768)에서는 18.1%의 성능 향상을 보였습니다. 이는 대규모 컨텍스트를 처리하는 최신 LLM 환경에서 매우 유의미한 수치입니다.

리뷰어 피드백과 검증

리뷰 과정에서 yyihuang은 다양한 테스트 스위트(test_dcp_spec_balanced.py, test_cake_fmha.py 등)를 B200 환경에서 실행할 것을 요청했습니다.

결과적으로:

• Correctness: fp64 참조 모델 대비 오차가 Round 3와 동일하거나 더 적음을 확인했습니다.
• Sanitizer: compute-sanitizer를 통해 메모리 오류 및 동기화 오류(synccheck)가 없음을 검증했습니다.

마치며: 시니어 엔지니어의 교훈

이번 FlashInfer의 PR은 최신 하드웨어가 출시되었을 때 소프트웨어 엔지니어가 어떻게 대응해야 하는지를 잘 보여줍니다.

1. 하드웨어 특화 최적화: 단순히 코드를 짜는 것이 아니라, Blackwell의 메모리 계층과 벡터 연산 특성에 맞춰 루프 구조를 완전히 재설계했습니다.
2. 수치적 엄밀성: 성능을 위해 정확도를 희생하지 않았습니다. 오히려 결정론적 연산을 도입하여 디버깅과 신뢰성을 높였습니다.
3. 점진적 개선: Round 1부터 4까지 지속적으로 커널을 다듬으며, 각 단계마다 명확한 벤치마크 수치를 제시하는 데이터 중심의 접근 방식을 취했습니다.

고성능 컴퓨팅(HPC) 분야에서 일하는 엔지니어라면, 이번 PR의 루프 언롤링 기법과 메모리 배치 전략을 반드시 참고해 보시기 바랍니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#floating-point-and-ieee-754
• https://github.com/flashinfer-ai/flashinfer
• https://www.nvidia.com/en-us/data-center/blackwell/

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