[flashinfer] FlashInfer: Rubin(SM107) 아키텍처를 위한 네이티브 확률적 반올림(Stochastic Rounding) 최적화

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

PR 링크: flashinfer-ai/flashinfer#5267 상태: Merged | 변경: +53 / -28

들어가며

최근 고성능 추론 엔진인 FlashInfer는 Mamba 모델의 상태 업데이트(State Update) 성능을 극대화하기 위해 하드웨어 가속 확률적 반올림(Stochastic Rounding, SR) 기능을 지원해 왔습니다. 하지만 기존 구현에서는 Blackwell 아키텍처의 특정 변형(SM100a, SM103a)만을 타겟으로 하여, 최신 Rubin(SM107a) 아키텍처가 하드웨어 가속 대신 느린 소프트웨어 에뮬레이션 경로를 타는 문제가 있었습니다. 본 PR은 SM107a를 지원 목록에 추가하여 하드웨어 네이티브 명령어를 활용함으로써 성능 병목을 해결했습니다.

코드 분석

1. `flashinfer/utils.py`: 아키텍처 지원 테이블 중앙화

기존에는 하드웨어 지원 여부를 하드코딩된 조건문으로 관리했으나, 이를 CVT_RS_SUPPORTED_ARCHES 딕셔너리로 통합하여 유지보수성을 높였습니다.

# Before
# (기존에는 하드코딩된 조건문으로 관리)

# After
CVT_RS_SUPPORTED_ARCHES = {
    (10, 0): "sm_100a",
    (10, 3): "sm_103a",
    (10, 7): "sm_107a",
}

def is_cvt_rs_supported(device: torch.device = None) -> bool:
    # ... 구현 생략 ...
    return get_compute_capability(device) in CVT_RS_SUPPORTED_ARCHES

2. `include/flashinfer/mamba/conversion.cuh`: CUDA 아키텍처 가드 업데이트

PTX 명령어 cvt.rs를 호출하기 위한 컴파일 타임 가드를 업데이트하여 SM107a 환경에서도 하드웨어 명령어가 생성되도록 수정했습니다.

// Before
// (SM100a, SM103a만 고려된 가드)

// After
// sm_100a (B200), sm_103a (B300), and sm_107a (Rubin) 지원
#if FLASHINFER_MAMBA_HAS_CVT_RS
__device__ __forceinline__ __half cvt_rs_f16_f32(float x, uint32_t rand13) {
  return __nvvm_cvt_rs_f16_f32(x, rand13);
}
#endif

왜 이게 좋은가

이번 최적화는 소프트웨어 에뮬레이션(약 12개 명령어 소요)을 하드웨어 네이티브 명령어(cvt.rs)로 대체함으로써, 연산 효율을 극적으로 개선했습니다. 벤치마크 결과에 따르면, 배치 사이즈 256 기준 약 36.41%의 지연 시간(Latency) 감소 효과를 보였습니다.

성능 지표 (Batch 256 기준)

Before: 502.880 µs
After: 319.760 µs
결과: 1.573배 속도 향상

교훈

하드웨어 가속의 중요성: 범용적인 소프트웨어 에뮬레이션은 편리하지만, 고성능 추론 엔진에서는 특정 아키텍처의 전용 명령어를 활용하는 것이 필수적입니다.
설정의 중앙화: 리뷰어의 피드백처럼 지원 아키텍처 목록을 단일 소스(Single Source of Truth)로 관리하면, 향후 새로운 아키텍처(예: SM108 등)가 출시되었을 때 코드 수정 범위를 최소화할 수 있습니다.