[flashinfer] FlashInfer Mamba2 SSD 커널 최적화: Rubin(sm_107a) 아키텍처를 위한 고성능 전략

PR 링크: flashinfer-ai/flashinfer#7869 상태: Merged | 변경: +28373 / -103

들어가며

FlashInfer의 cake_ssd_combined는 Mamba2 모델의 SSD(State Space Duality) 연산을 위한 핵심 커널 모음입니다. 이번 업데이트는 NVIDIA의 차세대 Rubin(sm_107a) 아키텍처를 타겟으로, 다양한 워크로드 패턴(batched, varlen, chunk-parallel 등)에 최적화된 실행 경로를 추가하고, 데이터 로드 효율을 극대화하는 것을 목표로 합니다. 특히, 메모리 대역폭 활용도를 높이기 위한 Z-prefetch 도입과 CTA(Cooperative Thread Array) 클러스터 활용을 통해 실제 추론 성능을 향상시켰습니다.

코드 분석

1. Z-prefetch를 통한 데이터 로드 최적화

이번 PR에서는 exactfused 계열 커널에 Z-prefetch 기능을 추가했습니다. 이는 현재 청크(chunk)를 처리하는 동안 다음 청크의 Z 컬럼을 미리 로드하여 파이프라인 지연 시간을 줄이는 전략입니다.

// Before (기존 코드의 일부)
reinterpret_cast<int4*>(out_native + (out_row_base + dim_base))[0] = reinterpret_cast<int4*>(group_values_pva)[0];

// After (변수명 정리 및 최적화)
reinterpret_cast<int4*>(out_native + (out_row_base + dim_base))[0] = reinterpret_cast<int4*>(group_values_pv)[0];

변수명 group_values_pva를 group_values_pv로 간결하게 수정하고, 루프 내에서 효율적인 데이터 처리를 수행하도록 개선되었습니다. 이와 함께 exactfused_{bf16,f32}_batched_zpf와 같은 새로운 커널이 추가되어, Z 차원이 2개 청크 이상인 경우에 대해 선제적인 메모리 로드를 수행합니다.

2. CTA-Cluster 및 Wrap-around Plan 도입

Rubin 아키텍처의 특성을 활용하기 위해 1x1x1 클러스터 속성을 가진 'Twin' 커널들을 도입했습니다. 또한, exactpre_{bf16,f32}_varlen_wrap 커널을 통해 워크로드가 SM 개수보다 많을 때, CTA 간의 작업을 효율적으로 분배하는 Wrap-around 계획을 적용했습니다.

왜 이게 좋은가

이번 최적화의 핵심은 '워크로드 적응형 커널 선택'입니다. 단순히 하나의 범용 커널을 사용하는 대신, 시퀀스 길이, 헤드 개수, 데이터 타입에 따라 최적화된 커널을 동적으로 선택하도록 라우팅 로직을 개선했습니다.

• 성능 수치: 93개의 테스트 케이스 중 30개 케이스에서 유의미한 성능 향상이 관찰되었습니다. 특히 perf_varlen_32x32 케이스에서는 기존 112.59 µs에서 100.04 µs로 약 12.5%의 속도 향상을 기록했습니다. 전체적인 기하평균 속도 향상은 약 0.73%입니다.
• 교훈: GPU 커널 최적화에서 '데이터 재사용성'과 '파이프라인 병렬성'은 핵심입니다. Z-prefetch와 같이 다음 연산에 필요한 데이터를 미리 로드하는 기법은 메모리 바운드(Memory-bound) 연산이 많은 LLM 커널에서 매우 효과적입니다.

결론

이번 FlashInfer 업데이트는 Rubin 아키텍처의 하드웨어 기능을 최대한 활용하여 Mamba2 SSD 연산의 병목을 성공적으로 해결했습니다. 특히 정교한 커널 라우팅과 파이프라인 최적화는 향후 다른 아키텍처로의 확장성 측면에서도 좋은 레퍼런스가 될 것입니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-data-copy
• https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#cooperative-thread-arrays

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