본문으로 건너뛰기

[onnxruntime] CUDA GEMV 최적화: M=8 Narrow Tile 확장을 통한 GPU 점유율 극대화 (A100/H200 대응)

PR 링크: microsoft/onnxruntime#33600 상태: Merged | 변경: +1107 / -59

들어가며

LLM(대규모 언어 모델) 추론 성능을 높이기 위한 기술 중 하나인 Speculative Decoding(추측 디코딩)은 작은 모델(Drafter)이 생성한 여러 토큰을 큰 모델(Target)이 한 번에 검증하는 방식을 취합니다. 이 과정에서 배치 사이즈($M$)가 8 내외인 GEMV(General Matrix-Vector Multiplication) 연산이 빈번하게 발생합니다.

최근 Microsoft ONNX Runtime 레포지토리에 반영된 PR에서는 이 $M=8$ 구간의 GEMV 성능을 극대화하기 위한 정교한 최적화가 포함되었습니다. 핵심 문제는 'Wave Under-utilization'입니다. GPU의 SM(Streaming Multiprocessor) 개수에 비해 연산해야 할 타일(Thread Block)의 수가 너무 적으면 GPU 자원을 충분히 활용하지 못하게 됩니다. 이번 글에서는 타일 크기($CtaN$)를 절반으로 줄여(Narrow Tile) 블록 수를 늘림으로써 어떻게 성능을 개선했는지 코드 수준에서 분석합니다.

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

1. 타일 선택 로직의 분리와 정교화 (PickM8GemvCtaN 도입)

기존에는 $M=8$인 경우에도 고정된 타일 크기를 사용하거나 단순한 $N$ 값의 임계치(Cutoff)만으로 타일 크기를 결정했습니다. 이번 변경에서는 PickM8GemvCtaN이라는 전용 선택기를 도입하여 디바이스 특성(A100, H200 등)에 맞는 최적의 타일을 선택하도록 개선되었습니다.

Before (기존 dispatcher.h 일부 추정):

// 기존에는 특정 조건에서 CtaNLargeM / 2를 하드코딩하거나 단순 비교함
if (M <= 8 && N < 5136) {
    cta_n = 2; // Narrow tile
} else {
    cta_n = 4; // Default tile
}

After (fpA_intB_gemv.h):

// 정교해진 디바이스별 타일 선택 로직
inline int PickM8GemvCtaN(int N, int K, int block_size, int device_sm, int device_sm_count, bool is_scale_only) {
  // A100 (SM 8.0, 108 SMs) 전용 윈도우
  if (device_sm == 80 && device_sm_count == 108) {
    if (is_scale_only && (block_size == 32 || block_size == 64 || block_size == 128)) {
      // INT4 전용 측정 기반 윈도우
      if (is_int4 && ((N > 3456 && N <= 4320) || (K <= 3072 && N > 1728 && N <= 2560))) return 2;
      // INT8 전용 측정 기반 윈도우
      if (is_int8 && (N > 864 && N <= 1304)) return 2;
    }
    return 4; // 기본값
  }
  // 기타 디바이스는 기존 Cutoff 유지
  return (N / (CtaNLargeM * kInterleave) <= 320) ? 2 : 4;
}

2. Host-side 오버헤드 제거: cudaGetDevice 호출 최적화

디바이스의 SM 정보를 알기 위해 매번 커널 런칭 시점에 CUDA API를 호출하는 것은 CPU 오버헤드를 발생시킵니다. 이번 PR에서는 Params 구조체에 device_sm 정보를 미리 포함시켜 Dispatcher 내부에서의 런타임 호출을 제거했습니다.

Before (dispatcher.h):

// 매 런칭마다 디바이스 속성을 쿼리함 (오버헤드 발생)
int device;
cudaGetDevice(&device);
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device);

After (dispatcher.h):

// Params에 미리 저장된 값을 사용
template <typename T>
void dispatch_gemv(const Params& params) {
    // ...
    int cta_n = PickM8GemvCtaN(params.N, params.K, params.block_size, 
                               params.device_sm, params.device_sm_count, 
                               params.is_scale_only);
    // ...
}

이 변경을 통해 A100 기준 CPU Enqueue 시간이 약 0.1µs에서 0.02µs 수준으로 80% 가량 감소했습니다.

3. 테스트 커버리지 강화

리뷰어들의 피드백에 따라 $M=8$ 케이스에 대한 엄격한 회귀 테스트가 추가되었습니다. 특히 $M=7, 9$와 같은 인접 구간에서 Fallback이 정상 작동하는지, 그리고 $M=8$ 결과가 $M=4$를 두 번 실행한 결과와 비트 단위로 일치하는지 검증합니다.

Test Code (fpA_intB_gemm_kernel_test.cc):

// M=8 Narrow Tile 검증 테스트
TEST(FpAIntBGemvTest, M8NarrowTile_INT8_BF16) {
    // N=1280 (Narrow 적용 구간)과 N=1536 (Fallback 구간)을 모두 테스트
    for (int n : {1280, 1536}) {
        auto params = CreateGemvParams(8, n, 1024);
        RunAndVerify(params);
        // CtaN이 예상대로 선택되었는지 로그 확인
        EXPECT_EQ(GetLaunchedCtaN(), (n == 1280 ? 2 : 4));
    }
}

왜 이게 좋은 최적화인가?

1. Occupancy(점유율)와 Wave의 이해

GPU는 수많은 SM으로 구성되어 있습니다. 예를 들어 A100은 108개의 SM을 가집니다. 만약 $N=1280$인 연산에서 $CtaN=4$를 사용하면 블록 수가 적어 일부 SM이 놀게 됩니다(Under-utilization). $CtaN=2$로 줄이면 블록 수가 두 배로 늘어나 더 많은 SM에 작업을 분배할 수 있고, 이는 곧 지연 시간(Latency) 단축으로 이어집니다.

2. 성능 수치 기반의 의사결정

PR 설명에 포함된 A100 측정 결과에 따르면:

  • INT4 (N=4320): Narrow Tile 사용 시 17.52µs, Default 사용 시 18.71µs (약 6.4% 개선).
  • H200 (BF16/INT8): Drafter 모델의 o_proj 레이어 지연 시간이 14.50µs에서 13.83µs로 단축되었습니다.

3. 정교한 Heuristics

무조건 타일을 작게 만드는 것이 정답은 아닙니다. 타일이 작아지면 블록 간 중복 작업이 늘어날 수 있습니다. 이번 PR은 A100에서 $N$이 4336을 넘어서는 순간 Narrow Tile이 오히려 28% 이상 느려진다는 점을 발견하고, 이를 기반으로 'Measured Window'를 설정했습니다. 이는 단순한 이론적 접근이 아닌 철저한 벤치마크 기반의 최적화임을 보여줍니다.

결론 및 교훈

  1. Micro-optimization의 가치: LLM 추론처럼 수조 번 반복되는 연산에서는 1µs 미만의 CPU 오버헤드 제거와 5%의 커널 성능 향상이 거대한 비용 절감으로 이어집니다.
  2. 디바이스 특성 고려: SM 개수와 레지스터 파일 크기가 다른 디바이스(A100 vs H200)에 대해 동일한 전략을 고수하지 않고 개별 튜닝을 적용한 점이 인상적입니다.
  3. 테스트의 중요성: 성능 최적화 과정에서 결과값이 변하지 않았음을 보장하기 위해 비트 단위 일치(Bit-identical) 테스트를 수행하는 것은 엔지니어링의 기본이자 핵심입니다.

이 PR은 고성능 컴퓨팅 라이브러리가 지향해야 할 정교한 최적화의 정석을 보여주고 있습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글