[onnxruntime] ONNX Runtime의 CPU int4 가중치 프리패킹 최적화: 병렬 처리 효율성 개선

PR 링크: microsoft/onnxruntime#31690 상태: Merged | 변경: +363 / -186

들어가며

대규모 언어 모델(LLM)을 로컬 환경에서 실행할 때, 모델 로딩 시간은 사용자 경험에 큰 영향을 미칩니다. 특히 Qwen3.6-27B와 같은 대형 모델을 int4로 양자화하여 사용할 경우, 가중치 프리패킹(Prepacking) 과정에서 병목 현상이 발생하곤 합니다. Microsoft의 onnxruntime 레포지토리에 올라온 이번 PR은 CPU 환경에서 int4 가중치 프리패킹의 병렬화 전략을 개선하여, 세션 생성 시간을 기존 대비 약 25% 단축하는 성과를 거두었습니다.

코드 분석

이번 최적화의 핵심은 기존의 고정된 청크(Chunk) 크기 방식에서 벗어나, 스레드 풀(Thread Pool)의 가용성과 작업의 특성을 고려한 동적 병렬화로 전환한 것입니다.

1. onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx_common.h

기존에는 MLAS_PACK_BLKS_PER_CHUNK를 64로 고정하여 사용했습니다. 이는 대규모 행렬 연산에는 적합하지만, N(출력 차원)이 작은 경우 스레드 풀의 자원을 충분히 활용하지 못하고 유휴 상태(Idle)로 남는 문제가 있었습니다.

Before:

const size_t Iterations = N * SubBlkCountK; // one iteration per sub block
// ...
MlasTrySimpleParallel(ThreadPool, Iterations, [&](ptrdiff_t tid) {
    const size_t n = tid / SubBlkCountK;
    const size_t k_subblk = tid % SubBlkCountK;
    // ...
});

After:

const size_t ChunkSubBlks = MlasQNBitPackChunkSize(SubBlkCountK, N, ThreadPool);
const size_t ChunkCount = MlasDivRoundup(SubBlkCountK, ChunkSubBlks);
const size_t Iterations = N * ChunkCount;

MlasTrySimpleParallel(ThreadPool, Iterations, [&](ptrdiff_t tid) {
    const size_t n = tid / ChunkCount;
    const size_t chunk_idx = tid % ChunkCount;
    const size_t k_subblk_start = chunk_idx * ChunkSubBlks;
    // ... 루프를 통해 청크 단위로 처리
});

새롭게 추가된 MlasQNBitPackChunkSize 함수는 MaxThreads와 N을 비교하여, 작업 단위가 너무 작아 스레드 오버헤드가 발생하는 것을 방지하면서도, 동시에 너무 적은 작업으로 인해 코어가 놀지 않도록 청크 크기를 동적으로 조정합니다.

왜 이게 좋은가

1. 스레드 활용도 극대화: N이 작은 경우에도 ChunkCount를 조절하여 스레드 풀의 모든 코어를 효율적으로 사용합니다. 이는 특히 프로젝션 행렬이 좁은 모델에서 큰 성능 향상을 가져옵니다.
2. 스케줄링 오버헤드 감소: 너무 작은 단위로 작업을 쪼개면 스레드 간 컨텍스트 스위칭과 디스패치 오버헤드가 실제 연산 시간을 압도하게 됩니다. MLAS_PACK_MIN_BLKS_PER_CHUNK를 도입하여 최소 작업 단위를 보장함으로써 이를 해결했습니다.
3. 성능 수치: Intel Xeon Platinum 8480C 환경에서 테스트한 결과, 기존 약 44초가 소요되던 로딩 시간이 약 33초로 줄어들어 약 25%의 성능 향상을 확인했습니다.

교훈

병렬화는 단순히 '많이 쪼개는 것'이 아니라 '적절한 크기로 쪼개는 것'이 핵심입니다. 특히 라이브러리 레벨의 코드에서는 다양한 하드웨어 환경과 다양한 모델 구조(N의 크기)를 모두 고려해야 하므로, 이번 PR처럼 하드웨어의 스레드 가용성을 런타임에 확인하여 작업 단위를 결정하는 방식이 매우 효과적입니다.

리뷰어 피드백 반영

리뷰 과정에서 SubBlkCountK가 0일 경우 발생할 수 있는 0으로 나누기(Divide-by-zero) 오류에 대한 지적이 있었으며, 이를 방지하기 위한 가드 코드(if (N == 0 || SubBlkCountK == 0) return;)가 추가되었습니다. 또한, 경계 조건(Boundary conditions)에 대한 테스트 케이스 보강이 필요하다는 피드백을 통해 코드의 안정성을 높였습니다.

참고 자료

• https://onnxruntime.ai/docs/api/python/api_summary.html
• https://github.com/microsoft/onnxruntime

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