[flashinfer] Blackwell(SM100a) 성능 극대화: FlashInfer의 CTA-Pair 커널 최적화 분석

2026년 10월 19일수정: 2026년 10월 19일

PR 링크: flashinfer-ai/flashinfer#6586 상태: Merged | 변경: +804 / -46

들어가며

LLM(Large Language Model) 추론 성능을 결정짓는 핵심 연산 중 하나는 분산 환경에서의 통신과 연산을 결합한 all_gather_matmul입니다. FlashInfer 프로젝트의 Cake 백엔드는 이러한 연산을 가속화하기 위한 고성능 커널을 제공합니다.

이번 분석에서 다룰 PR은 NVIDIA의 최신 아키텍처인 Blackwell(SM100a, B200)에서 특정 워크로드(Shape)에 대해 CTA-pair 메인 커널을 활성화하여 성능을 최적화한 사례입니다. 기존에는 SM100a에서 모든 호출이 단일 CTA(Single-CTA) 커널로 실행되었으나, 이번 변경을 통해 9개의 주요 Shape에서 2개의 CTA가 협력하는 'Pair' 구조를 사용하여 하드웨어 활용도를 극대화했습니다.

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

1. JIT 라우팅 로직의 확장 (`flashinfer/jit/cake_all_gather_matmul.py`)

가장 먼저, 어떤 상황에서 최적화된 CTA-pair 커널을 사용할지 결정하는 라우팅 테이블이 업데이트되었습니다. 이전에는 SM100a에 대한 CG2_REGIONS(CTA-pair 적용 영역)가 정의되지 않았거나 비어 있었습니다.

Before (Conceptual):

# SM100a에 대한 CTA-pair 영역이 정의되지 않음
CG2_REGIONS = {
    "sm_103a": [...], # SM103a(B300)만 지원
}

After:

# SM100a를 위한 9개의 측정된 Shape(Region Points) 추가
CG2_REGIONS["sm_100a"] = [
    # [world_size, dtype, b_layout, min_padded_rows, max_padded_rows, min_cols, max_cols]
    [2, "bfloat16", "n-major", 2048, 2048, 2048, 2048], # M=2048
    [2, "bfloat16", "n-major", 4096, 4096, 2048, 2048], # M=4096
    [4, "bfloat16", "k-major", 1024, 1024, 2560, 2560], # TP4 QKV
    # ... 총 9개의 포인트 추가
]
CG2_CLASSES["sm_100a"] = ["ws2_bf16_n_major", "ws4_bf16_k_major", "ws8_bf16_k_major"]

이 변경을 통해 특정 world_size, dtype, M, N 조건이 맞을 경우 JIT 컴파일러는 일반 커널 대신 _cg2 접미사가 붙은 CTA-pair 커널을 선택하게 됩니다.

2. 커널 선언 및 TMA Descriptor 인코딩 (`.cu` 파일)

실제 커널 실행 파일에서는 새로운 CTA-pair 커널(_cg2)에 대한 선언과 Blackwell의 핵심 기술인 TMA(Tensor Memory Accelerator)를 활용하기 위한 Descriptor 생성 로직이 추가되었습니다.

Before:

// 단일 CTA 커널만 선언됨
extern "C" __global__ void kernel_cake_all_gather_matmul_c27adeb3c14305afd206(
    // ... arguments
);

After:

// 새로운 CTA-pair 커널(_06c1...) 선언 추가
extern "C" __global__ void kernel_cake_all_gather_matmul_06c1f7cb4aa36c5918bd(
    const __grid_constant__ CUtensorMap A_local, 
    const __grid_constant__ CUtensorMap A_scratch,
    const __grid_constant__ CUtensorMap B, 
    // ... 나머지 인자 동일
);

namespace cg2 {
// 3D TMA descriptor 인코딩 로직 추가
inline CUtensorMap EncodeTma_A_local(const TensorView& t) {
    // ... 중략 ...
    CUresult r = cuTensorMapEncodeTiled(&tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3,
                                      const_cast<void*>(tensor_base), global_dim, global_strides,
                                      box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE,
                                      CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE,
                                      CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
    return tm;
}
} // namespace cg2

여기서 cuTensorMapEncodeTiled는 Blackwell 하드웨어가 메모리에서 데이터를 효율적으로 가져올 수 있도록 텐서 레이아웃을 설명하는 객체를 생성합니다. CU_TENSOR_MAP_SWIZZLE_128B와 같은 설정을 통해 L2 캐시 활용도를 높이고 뱅크 컨플릭트를 최소화합니다.

3. 공통 헤더의 헬퍼 모듈 업데이트 (`cake_all_gather_matmul_device_common.cuh`)

컴파일러 헬퍼 라이브러리의 구성이 변경되었습니다. 이는 코드 생성기(Cake Codegen)가 Blackwell 아키텍처에 최적화된 모듈을 더 많이 포함하도록 업데이트되었음을 의미합니다.

Before/After 비교:

// Before: 43 modules
// After: 46 modules
// Generated by cake.codegen — do not edit.
// Shared device preamble of the sm_100a, sm_103a generated kernels

왜 이게 좋은 최적화인가?

1. 실질적인 성능 향상 (Measured Speedups)

PR 설명에 포함된 벤치마크 결과에 따르면, 최적화가 적용된 9개 Shape 중 7개에서 4.1% ~ 10.1%의 유의미한 성능 향상이 관찰되었습니다.

TP4 gate/up (M=512, N=14336): 약 10.1% 빨라짐.
World size 2 (M=4096, N=2048): 약 8.8% 빨라짐.

이러한 수치는 이미 고도로 최적화된 커널에서 추가로 얻어낸 결과라는 점에서 매우 가치가 큽니다.

2. CTA Clustering 활용

Blackwell 아키텍처는 여러 개의 CTA(Thread Block)를 하나의 클러스터로 묶어 하드웨어 자원을 공유하거나 동기화할 수 있는 기능을 제공합니다. CTA-pair 커널은 이러한 특성을 활용하여, 단일 CTA가 처리하기에는 너무 크거나 통신 오버헤드가 발생하는 작업을 두 개의 CTA가 나누어 처리함으로써 처리량(Throughput)을 높입니다.

3. 정밀한 라우팅 전략

모든 Shape에 대해 무조건 Pair 커널을 적용하는 것이 아니라, 실제 측정을 통해 성능 이득이 확실한 'Region Points'를 정의했습니다.

Padded Rows 기반 선택: M 값이 정확히 일치하지 않더라도 128의 배수로 패딩된 값이 포인트와 일치하면 최적화된 커널을 사용하도록 설계하여 범용성을 확보했습니다.
Bitwise Identity 보장: 성능은 올리되 결과값은 기존 레퍼런스 커널과 비트 단위로 동일(Bitwise Identical)함을 검증하여 안정성을 확보했습니다.

결론 및 교훈

이번 PR은 최신 GPU 하드웨어의 특성을 깊이 이해하고, 이를 소프트웨어 스택(JIT 컴파일러 및 커널 라이브러리)에 어떻게 녹여내야 하는지를 잘 보여줍니다.

데이터 기반 최적화: 단순히 이론적인 예측이 아니라, 실제 B200 노드에서 AB 테스트를 통해 검증된 Shape만 라우팅 테이블에 추가했습니다.
하드웨어 가속기 활용: TMA와 같은 최신 하드웨어 기능을 적극적으로 사용하여 메모리 대역폭을 효율적으로 사용했습니다.
점진적 개선: SM103a에서 먼저 검증된 로직을 SM100a로 확장 적용하며 코드 재사용성을 높였습니다.

시니어 엔지니어로서 배울 점은, "최적화는 공짜가 아니며, 하드웨어의 변화에 맞춰 소프트웨어의 라우팅 전략도 세밀하게 진화해야 한다"는 사실입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer, Kimi-K3 FP8 프로젝션 커널 공유 및 코드 최적화를 통한 성능 향상
현재글 : [flashinfer] Blackwell(SM100a) 성능 극대화: FlashInfer의 CTA-Pair 커널 최적화 분석
다음글 [flashinfer] FlashInfer SM90 아키텍처 성능 최적화: QkvNormRopeExchangeOverlap의 Shape Gate 개선