본문으로 건너뛰기

[flashinfer] NVIDIA Blackwell(SM103a)을 위한 극한의 커널 퓨전: MiniMax-H3 BF16 Pre-attention 최적화 분석

PR 링크: flashinfer-ai/flashinfer#4690 상태: Merged | 변경: +2532 / -0

들어가며

최신 대규모 언어 모델(LLM)과 확산 모델(Diffusion Models)의 추론 성능을 결정짓는 핵심 요소 중 하나는 커널 퓨전(Kernel Fusion)입니다. 개별적인 연산(RMSNorm, Linear, RoPE 등)을 각각 실행하면 매번 GPU 메모리(VRAM)와 연산 유닛 사이에서 데이터를 주고받아야 하는 메모리 대역폭 병목(Memory Wall) 현상이 발생합니다.

이번에 분석할 FlashInfer의 PR은 NVIDIA의 차세대 아키텍처인 Blackwell(SM103a)을 타겟으로, MiniMax-H3 워크로드에서 사용되는 Pre-attention 연산들을 하나의 거대한 커널로 통합한 사례입니다. 특히 CUDA 12.9/13.0에서 도입된 Tensor Map과 TMA(Tensor Memory Accelerator) 기능을 적극 활용하여 성능을 극대화했습니다.

코드 분석: Segmented Baseline vs Fused Kernel

1. 기존 방식 (Segmented Baseline)

기존에는 PyTorch의 표준 연산들을 순차적으로 호출했습니다. 아래의 _segmented_baseline 함수를 보면 데이터가 여러 단계를 거치며 메모리에 반복적으로 써지고 읽히는 것을 알 수 있습니다.

# Before: Segmented Baseline (benchmarks/bench_minimax_h3_bf16_pre_attention.py)
def _segmented_baseline(case):
    # 1. RMSNorm
    norm = F.rms_norm(case["x"], (HIDDEN,), case["x_norm_weight"], eps=EPS).to(torch.bfloat16)
    # 2. Indexed AdaLN (Scale & Shift)
    index = case["adaln_index"].long()
    scale = case["adaln_scale"].index_select(0, index)
    shift = case["adaln_shift"].index_select(0, index)
    adaln = torch.addcmul(shift, norm, (scale + 1.0).to(torch.bfloat16)).to(torch.bfloat16)
    # 3. QKV Projection (Linear)
    qkv = F.linear(adaln, case["qkv_weight"]).to(torch.bfloat16)
    # 4. Per-head RMSNorm & RoPE & Packing
    # ... (중략) ...
    case["baseline_out"].copy_(packed_view)
    return case["baseline_out"]

이 방식의 문제점은 norm, adaln, qkv 등 중간 결과물들이 모두 VRAM에 저장되어야 한다는 점입니다. 이는 연산 속도보다 메모리 전송 속도가 느린 현대 GPU 환경에서 큰 오버헤드가 됩니다.

2. 개선된 방식 (Fused Kernel)

이번 PR에서 도입된 minimax_h3_bf16_pre_attention 커널은 이 모든 과정을 단 한 번의 커널 실행으로 끝냅니다.

# After: Fused Kernel Call
def _run_candidate(case):
    return minimax_h3_bf16_pre_attention(
        case["x"], case["x_norm_weight"],
        case["adaln_scale"], case["adaln_shift"], case["adaln_index"],
        case["qkv_weight"], case["q_norm_weight"], case["k_norm_weight"],
        case["rope_cos_sin"],
        ulysses_degree=case["ulysses_degree"],
        out=case["out"], eps=case["eps"],
    )

내부적으로는 Blackwell 아키텍처의 하드웨어 가속기인 TMA를 사용하기 위해 CUtensorMap 구조체를 정의하고 활용합니다.

/* csrc/cake_minimax_h3_bf16_pre_attention_sm103a.cu */
struct __align__(128) FlashInferTensorMap {
  uint64_t opaque[16];
};
// CUDA의 CUtensorMap과 ABI 호환성을 유지하며 Blackwell의 TMA 기능을 사용
static_assert(sizeof(FlashInferTensorMap) == 128, "tensor-map ABI size mismatch");

왜 이게 좋은 최적화인가?

1. 하드웨어 특화 가속 (Blackwell TMA)

Blackwell 아키텍처(SM103a)의 핵심은 TMA(Tensor Memory Accelerator)입니다. 기존의 cp.async보다 더 진보된 이 기능은 다차원 텐서 레이아웃을 하드웨어 레벨에서 이해하고 데이터를 비동기적으로 로드합니다. 이번 PR은 CUtensorMap을 통해 이 기능을 직접 제어하여 데이터 로딩 오버헤드를 최소화했습니다.

2. 압도적인 성능 향상

벤치마크 결과에 따르면, 모든 생산용 쉐이프(Production-center shapes)에서 평균 1.189x의 속도 향상을 보였습니다. 특히 특정 설정(P=8, M=4824)에서는 지연 시간이 3.007ms에서 2.079ms로 줄어들어 약 1.44배(44%)의 성능 향상을 기록했습니다.

3. 메모리 효율성 및 안전성

리뷰 과정에서 논의된 것처럼, AdaLN 인덱스가 범위를 벗어나는 경우(-1, 9 등)에 대한 예외 처리를 디바이스 코드 내에 포함하여 안정성을 높였습니다. 또한, 중간 텐서 생성을 억제함으로써 VRAM 사용량을 획기적으로 줄였습니다.

리뷰어 피드백 분석

리뷰어 yyihuang은 TMA와 mbarrier의 동기화 순서에 대해 중요한 지적을 남겼습니다. NVIDIA 공식 가이드에 따르면 cp_async_bulk_tensor를 호출한 후 mbarrier_arrive_expect_tx를 호출하는 것이 올바른 순서이며, 이 커널은 해당 표준을 정확히 따르고 있음을 확인했습니다.

또한, 초기 구현에서 발생할 수 있었던 Tensor Map의 생명주기(Lifetime) 문제나 잘못된 인덱스 참조 문제를 테스트 케이스 추가를 통해 해결했습니다. 특히 CUDA Graph 리플레이 시 이전 출력값이 남아있어 테스트를 잘못 통과하는 경우를 방지하기 위해, 리플레이 전 출력 버퍼를 명시적으로 초기화하는 로직을 추가한 점이 인상적입니다.

결론

이번 FlashInfer의 업데이트는 단순히 코드를 합치는 수준을 넘어, 특정 하드웨어(Blackwell)의 기능을 극한으로 끌어올린 최적화의 정석을 보여줍니다. LLM 서빙 성능을 한 단계 더 끌어올리기 위해서는 이처럼 하드웨어 아키텍처에 깊게 밀착된 커널 개발이 필수적임을 다시 한번 확인시켜 줍니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글