본문으로 건너뛰기

[onnxruntime] ONNX Runtime: fpA_intB GEMM 최적화 및 CUDA 그래프 호환성 강화

PR 링크: microsoft/onnxruntime#29622 상태: Merged | 변경: +467 / -116

들어가며

대규모 언어 모델(LLM)의 추론 성능을 최적화하는 핵심 기법 중 하나는 Weight-only Quantization입니다. 그중에서도 fpA_intB 방식(Floating-point Activation x Integer Weight)은 가중치만 저비트(int4, int8)로 양자화하여 메모리 대역폭을 절약하면서도 연산 정밀도를 유지하는 효율적인 방법입니다.

기존 ONNX Runtime(ORT)에서는 이 기능이 복잡한 비트마스크(ORT_FPA_INTB_GEMM) 뒤에 숨겨져 있었고, 빌드 시 기본적으로 비활성화되어 있었습니다. 또한, 최적의 커널을 찾는 프로파일러(Profiler)가 멀티스레드 환경에서 병목을 일으키거나 CUDA Graph 캡처 도중 오작동하는 등의 한계가 있었습니다. 이번 PR은 이러한 제약을 해결하고 fpA_intB 경로를 메인스트림으로 끌어올린 중요한 업데이트입니다.

코드 분석: 핵심 변경 사항

1. 기본 활성화 및 빌드 옵션 변경

가장 먼저 눈에 띄는 변화는 fpA_intB 커널이 CUDA 빌드에서 기본적으로 포함되도록 변경된 점입니다.

Before (cmake/CMakeLists.txt):

option(onnxruntime_USE_FPA_INTB_GEMM "Build FpA IntB gemm cuda kernels" OFF)

After (cmake/CMakeLists.txt):

cmake_dependent_option(onnxruntime_USE_FPA_INTB_GEMM "Build FpA IntB gemm cuda kernels" ON "onnxruntime_USE_CUDA" OFF)

이제 onnxruntime_USE_CUDA가 활성화되면 별도의 플래그 없이도 고성능 양자화 커널을 사용할 수 있습니다.

2. 프로파일러의 락 경합(Lock Contention) 최적화

시니어 엔지니어로서 가장 주목할 만한 부분은 gemm_profiler.h에서 발생하던 전역 락(Global Lock) 문제를 해결한 방식입니다. 기존에는 최적의 커널 설정을 찾는 동안 맵 전체에 쓰기 락을 걸어 다른 스레드의 조회까지 차단했습니다.

Before (onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h):

// getBestConfigOrProfile 내부에서 락을 잡은 채로 무거운 CUDA 연산 수행
std::unique_lock<std::mutex> lock(mMNKProfileMap->mutex);
// ... 커널 프로파일링 (오래 걸림) ...

After (onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h):

// 1. 먼저 읽기 락이나 짧은 검사 후 락 해제
// 2. 락이 없는 상태에서 무거운 CUDA 작업 수행
auto stream = create_local_stream();
auto workspace = allocate_local_workspace();
profileTacticsForProblem(..., workspace, stream);

// 3. 마지막에 결과를 저장할 때만 다시 락을 획득
std::unique_lock<std::mutex> lock(mMNKProfileMap->mutex);
if (mMNKProfileMap->find(id) == mMNKProfileMap->end()) {
    mMNKProfileMap->insert({id, bestConfig});
}

이 변경을 통해 멀티스레드 추론 환경에서 프로파일링으로 인한 스톨(Stall) 현상을 획기적으로 줄였습니다.

3. CUDA Graph 안전성 확보

CUDA Graph는 커널 실행 오버헤드를 줄여주지만, 그래프 캡처 도중에 새로운 커널을 실행하거나 메모리를 할당하는 프로파일링 작업은 허용되지 않습니다. 이번 PR은 이를 명확히 체크합니다.

Code (onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc):

bool is_capturing = false;
// stream이 nullptr(default stream)인 경우도 포함하여 체크
CUDA_CALL_THROW(cudaStreamIsCapturing(stream, &is_capturing));

if (is_capturing) {
    // 캡처 중에는 프로파일링 없이 기존 캐시된 최적 설정만 조회
    return getBestConfig(m, n, k);
} else {
    // 캡처가 아닐 때만 Lazy Profiling 허용
    return getBestConfigOrProfile(m, n, k, stream);
}

리뷰어 tianleiwu의 피드백에 따라 stream != nullptr 가드를 제거하여 기본 스트림에서의 캡처 상태도 정확히 감지하도록 수정되었습니다.

4. 설정 인터페이스 단순화 및 버그 수정

기존의 난해한 비트마스크 대신 직관적인 on/off 문자열과 세션 설정을 도입했습니다. 또한, 환경 변수 로딩 시 발생하던 논리적 오류도 수정되었습니다.

Bug Fix (cuda_runtime_utils.h):

// Before: firstCall이 true로 시작하는데 !firstCall일 때만 로직이 실행됨 (데드 코드)
if (!firstCall) { ... }

// After: 올바른 초기화 로직
if (firstCall) {
#if defined(_WIN32)
    _dupenv_s(&env, &env_len, "CUDA_LAUNCH_BLOCKING");
#else
    char const* env = std::getenv("CUDA_LAUNCH_BLOCKING");
#endif
    result = (env != nullptr && std::string(env) == "1");
    firstCall = false;
}

왜 이게 좋은 개선인가?

  1. 사용성(UX) 향상: 사용자는 이제 ep.cuda.fpa_intb_gemm이라는 명확한 세션 옵션으로 기능을 제어할 수 있습니다. 특히 Prepacked 가중치를 사용할 경우 이 옵션을 자동으로 강제 활성화하여 런타임 에러를 방지합니다.
  2. 성능 안정성: 프로파일러의 락 범위를 최소화함으로써 고부하 서버 환경에서의 동시성 문제를 해결했습니다. 이는 단순한 기능 추가를 넘어 프로덕션 레벨의 안정성을 고려한 설계입니다.
  3. 유연한 튜닝: ep.cuda.fpa_intb_profile_m 옵션을 통해 특정 M(Batch size) 버킷에 대해서만 프로파일링을 수행하도록 설정할 수 있어, 웜업(Warm-up) 시간을 단축할 수 있습니다.

결론

이번 PR은 ONNX Runtime의 CUDA 실행 엔진을 더욱 견고하게 만들었습니다. 특히 LLM 서비스 운영 시 중요한 CUDA Graph 지원과 멀티스레드 성능 최적화가 포함되어 있어, vLLM이나 FasterTransformer와 같은 프레임워크와 대등한 수준의 최적화 도구를 제공하게 되었습니다. 시니어 개발자라면 이러한 '락 범위 최소화'와 '상태 머신(CUDA Graph)에 대한 방어적 프로그래밍' 패턴을 눈여겨볼 필요가 있습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글