[flashinfer] NVIDIA Blackwell의 잠재력을 극한으로: MiniMax-H3 NVFP4 양자화 및 GEMM 최적화 분석

2026년 9월 26일수정: 2026년 9월 26일

PR 링크: flashinfer-ai/flashinfer#5562 상태: Merged | 변경: +154 / -716

들어가며

최근 LLM 및 확산 모델(Diffusion Models)의 추론 성능을 극대화하기 위해 FP8을 넘어 FP4(NVFP4)와 같은 저정밀도 데이터 타입 도입이 가속화되고 있습니다. 특히 NVIDIA의 최신 아키텍처인 Blackwell(SM100, SM103)은 이러한 저정밀도 연산을 하드웨어 수준에서 가속하는 tcgen05 명령어를 도입했습니다.

이번 글에서는 FlashInfer 라이브러리에 반영된 MiniMax-H3 모델의 출력 투영(Output Projection) 및 Gated Residual 연산 최적화 PR을 분석합니다. 이 PR은 단순한 기능 구현을 넘어, Blackwell 아키텍처의 특성을 활용해 NVFP4 양자화 패스에서 HBM 대역폭의 80%를 달성하고, 전체 연산 속도를 기존 대비 최대 4배까지 끌어올린 고도의 최적화 기법을 담고 있습니다.

코드 분석: 핵심 최적화 포인트

1. NVFP4 양자화 패스의 병렬성 강화 (Row Splitting)

기존 구현에서는 Warp 하나가 하나의 행(Row)을 담당하여 양자화를 수행했습니다. 하지만 Blackwell과 같은 고성능 GPU에서는 단일 Warp의 메모리 요청만으로는 HBM의 거대한 대역폭을 충분히 채우기 어렵습니다. 이번 PR에서는 ROW_SPLIT 개념을 도입하여 하나의 행을 여러 파트가 나누어 처리하도록 변경했습니다.

Before:

// 단일 Warp가 하나의 행을 전담
int row = bid * ROWS_PER_CTA + warp;
if (row < M) {
    // ... 행 전체에 대한 루프 수행
}

After:

// ROW_SPLIT을 통해 하나의 행을 여러 파트(Warp)가 분할 처리
#define ROW_SPLIT 1 // 실제 런타임 설정에 따라 확장 가능 구조

int slot = bid * ROWS_PER_CTA + warp;
int row = slot / ROW_SPLIT;
int part = slot - row * ROW_SPLIT;

if (row < M) {
    int vec0 = part * (VECS_PER_LANE / ROW_SPLIT);
    #pragma unroll
    for (int batch = 0; batch < VECS_PER_LANE / ROW_SPLIT / LOAD_BATCH; batch++) {
        // 14개의 독립적인 16-byte 로드를 동시에 수행 (In-flight)
        // ...
    }
}

이 변경을 통해 Lane당 14개의 독립적인 16-byte 로드를 동시에 유지(In-flight)할 수 있게 되었습니다. 이는 메모리 지연 시간(Latency)을 효과적으로 숨기고 HBM 대역폭 활용률을 75-80%까지 끌어올리는 결과로 이어졌습니다.

2. TMEM Accumulator Drain 최적화 (32 -> 64 Columns)

Blackwell 아키텍처의 핵심 중 하나는 Tensor Memory(TMEM)입니다. GEMM 연산 결과가 저장되는 Accumulator를 TMEM에서 메인 메모리로 옮기는(Drain) 과정은 전체 성능에 큰 영향을 미칩니다. 기존에는 32개 컬럼 단위로 데이터를 읽어왔으나, 이를 64개 단위로 확장하여 명령어 오버헤드를 줄였습니다.

Before (32-column load):

// tcgen05.ld.sync.aligned.32x32b.x32.b32 사용
asm volatile(
    "tcgen05.ld.sync.aligned.32x32b.x32.b32"
    " {%0, %1, ..., %31}, [%32];"
    : "=f"(_tmem_load_0[0]), ... 
    : "r"(lane_addr));

After (64-column load):

// x64 변형을 사용하여 한 번에 더 많은 데이터를 로드
asm volatile(
    "tcgen05.ld.sync.aligned.32x32b.x64.b32"
    " {%0, %1, ..., %63}, [%64];"
    : "=f"(_tmem_load_0[0]), ... // 64개의 레지스터에 매핑
    : "r"(lane_addr));

이 최적화는 B200에서 약 0.4~1.0%의 추가 성능 향상을 가져왔습니다. 수치상으로는 작아 보이지만, 이미 극도로 최적화된 커널에서 Epilogue의 병목을 제거했다는 점에서 의미가 큽니다.

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer Kimi-K3 Fused MoE Router 최적화: Warp-per-row 전략 도입
현재글 : [flashinfer] NVIDIA Blackwell의 잠재력을 극한으로: MiniMax-H3 NVFP4 양자화 및 GEMM 최적화 분석
다음글 [flashinfer] FlashInfer NVFP4 QKV GEMM 최적화: SM103a Epilogue 통합 및 CUDA 런처 개선

[flashinfer] NVIDIA Blackwell의 잠재력을 극한으로: MiniMax-H3 NVFP4 양자화 및 GEMM 최적화 분석

들어가며

코드 분석: 핵심 최적화 포인트

1. NVFP4 양자화 패스의 병렬성 강화 (Row Splitting)

2. TMEM Accumulator Drain 최적화 (32 -> 64 Columns)

3. GEMM Raster Group 크기 조정 (16 -> 8)

왜 이게 좋은가?

1. Roofline 모델에 근접한 성능

2. 지연 시간 숨기기(Latency Hiding)의 정석

3. 하드웨어 특화 명령어 활용

마치며

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글