본문으로 건너뛰기

[flashinfer] Blackwell 시대를 위한 Mamba-2 SSD 최적화: FlashInfer의 네 번째 성능 개선 라운드 분석

PR 링크: flashinfer-ai/flashinfer#6291 상태: Merged | 변경: +3528 / -1776

들어가며

최근 언어 모델 아키텍처의 흐름에서 Mamba-2와 같은 SSD(Selective State Space) 모델은 Transformer의 선형적 어텐션 한계를 극복할 대안으로 주목받고 있습니다. 특히 대규모 시퀀스를 처리할 때 prefill 단계의 효율성은 전체 추론 성능을 결정짓는 핵심 요소입니다.

이번 FlashInfer PR은 NVIDIA의 차세대 아키텍처인 Blackwell(SM100, SM103)을 타겟으로 cake_ssd_combined 커널의 네 번째 성능 최적화 라운드를 다룹니다. 이 PR의 핵심 목표는 비트 단위의 정확성(Bitwise Identity)을 유지하면서도, 데이터 이동(Data Movement)과 스케줄링 최적화를 통해 하드웨어 한계치(Roofline)에 도달하는 것입니다. 시니어 엔지니어의 관점에서 이번 변경사항이 왜 뛰어난 최적화인지 코드 수준에서 분석해 보겠습니다.

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

1. 정밀도 확장 및 Shared Memory 레이아웃 최적화

가장 먼저 눈에 띄는 변화는 delta 피연산자의 처리 방식입니다. 이전 버전에서는 f16으로 Shared Memory(SMEM)에 스테이징한 후 연산 시점에 변환했지만, 이번 PR에서는 아예 f32로 확장하여 저장합니다.

Before:

// csrc/cake_mamba_ssd_combined/generated/device/..._39deb388f5.cu
#define SMEM_SMEM_DELTA_ALL_STAGE_BYTES 512
#define SMEM_TOTAL 231936

__half* smem_delta_all = reinterpret_cast<__half*>(smem_raw + SMEM_SMEM_DELTA_ALL_OFF);

// 루프 내부에서 매번 변환 발생
float _cvt_f32_2 = __half2float(smem_delta_all[stage1_3 * 128 + (unsigned int)col]);
b_scale *= _cvt_f32_2;

After:

// csrc/cake_mamba_ssd_combined/generated/device/..._df1813f0fc.cu
#define SMEM_SMEM_DELTA_ALL_STAGE_BYTES 1024 // 용량 2배 증가
#define SMEM_TOTAL 232448 // 전체 SMEM 사용량 증가

float* smem_delta_all = reinterpret_cast<float*>(smem_raw + SMEM_SMEM_DELTA_ALL_OFF);

// 루프 내부: 변환 없이 직접 연산
b_scale *= smem_delta_all[stage1_3 * 128 + (unsigned int)col];

왜 좋은 최적화인가? SMEM 사용량이 512바이트 늘어났지만, Blackwell 아키텍처의 SMEM 한도(227KB) 내에서 충분히 수용 가능한 수준입니다. 대신 Hot-path인 스캔 루프 내부에서 __half2float 변환 명령어를 제거함으로써 명령어 발행(Issue) 효율을 높였습니다. 이는 연산 집약적인 커널에서 지연 시간(Latency)을 줄이는 결정적인 역할을 합니다.

2. 16바이트 벡터화 로드 (Vectorized SMEM Access)

기존에는 SMEM에서 데이터를 읽을 때 스칼라 방식으로 접근했으나, 변경 후에는 ld.shared.v4.b32 인라인 어셈블리를 사용하여 16바이트씩 묶어서 읽어옵니다.

Before:

// 루프 내에서 개별 접근
float _exp2_2 = approx_exp2((row_cumsum - smem_cumsum_all[stage3_4 * 128 + (unsigned int)col_2]) * 1.4426950408889634f);

After:

// 16바이트(float 4개)씩 벡터화하여 로드
#pragma unroll
for (int quad = 0; quad < 8; quad++) {
    asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];"
        : "=r"(*reinterpret_cast<uint32_t*>(&cumsum_values[4 * quad])), ...
        : "r"(smem_cumsum_all_addr + (unsigned int)((chunk_elem + 4 * quad) * 4)));
}
// 이후 로컬 레지스터(cumsum_values)에서 사용
float _exp2_2 = approx_exp2((row_cumsum - cumsum_values[local_col]) * 1.4426950408889634f);

왜 좋은 최적화인가? CUDA 커널에서 SMEM 대역폭을 최대로 활용하려면 접근 횟수를 줄이고 한 번에 큰 단위를 읽는 것이 유리합니다. v4.b32를 사용하면 뱅크 충돌(Bank Conflict) 위험을 관리하면서도 처리량을 극대화할 수 있습니다. 특히 Blackwell의 TMEM(Tensor Memory) 연산과 병행될 때, 이러한 메모리 파이프라인 최적화는 연산 유닛이 쉬지 않고 돌아가게 만듭니다.

3. Host-side: CUDA Graph Capture 지원

호스트 측면에서는 CakeSSDCombined.prepare API를 도입하여 CUDA Graph 캡처가 가능하도록 개선되었습니다. 짧은 시퀀스(sub-chunk)의 경우 기존에는 호스트에서 패딩 복사(pad copies)를 수행했으나, 이제는 TMA(Tensor Memory Accelerator)의 Out-of-bounds 기능을 활용하여 커널 수를 6개에서 2개로 대폭 줄였습니다.

왜 이게 좋은가: 성능 수치와 교훈

이 PR의 결과는 놀랍습니다.

  1. 압도적인 성능 향상: B200(SM100) 기준, 주요 벤치마크 행(saturated rows)에서 이전 대비 1.63x ~ 1.64x의 속도 향상을 기록했습니다. Triton 구현체와 비교하면 최대 10배 이상 빠릅니다.
  2. HBM Bound 근접: 최적화된 커널은 HBM 대역폭 한계치의 56~73%에 도달했습니다. 이는 소프트웨어 수준에서 하드웨어 성능을 거의 한계까지 끌어올렸음을 의미합니다.
  3. 지연 시간 단축: 커널 간 간격(Kernel-to-kernel gap)을 0.3µs에서 0.2µs로 줄여, 작은 배치 사이즈에서의 응답성도 개선했습니다.

시니어 엔지니어의 인사이트

이 최적화 사례에서 배울 수 있는 일반적인 교훈은 다음과 같습니다.

  • Precision vs Throughput: 때로는 메모리를 조금 더 쓰더라도(f16 -> f32) 연산 파이프라인의 병목(변환 오버헤드)을 제거하는 것이 전체 성능에 이득입니다.
  • Hardware-Specific Features: Blackwell의 TMA와 TMEM 특성을 깊이 이해하고, 이를 위해 데이터 레이아웃을 재설계하는 것이 범용적인 코드를 작성하는 것보다 훨씬 강력한 결과를 냅니다.
  • Bitwise Identity: 성능을 올리면서도 결과값이 이전과 완전히 동일함을 보장하는 엄격한 검증 과정이 수반되어야 실제 프로덕션 환경에서 신뢰받는 라이브러리가 될 수 있습니다.

FlashInfer의 이번 업데이트는 Blackwell 아키텍처를 사용하는 AI 인프라에서 Mamba-2 모델의 추론 효율을 한 단계 격상시킨 훌륭한 사례입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글