본문으로 건너뛰기

[onnxruntime] GPU 점유율의 미학: Qwen MTP를 위한 ONNX Runtime NVFP4 GEMV 최적화 분석

PR 링크: microsoft/onnxruntime#32140 상태: Merged | 변경: +130 / -36

들어가며

최신 대규모 언어 모델(LLM)인 Qwen 시리즈, 특히 Multi-Token Prediction(MTP) 구조를 사용하는 모델들은 추론 속도를 높이기 위해 저정밀도 양자화 기술을 적극적으로 활용합니다. 그중에서도 NVFP4(NVIDIA 4-bit Floating Point)는 Hopper 아키텍처(H100, H200)에서 지원하는 강력한 데이터 타입입니다.

하지만 단순히 저정밀도 타입을 사용한다고 해서 항상 최고의 성능이 보장되는 것은 아닙니다. GPU의 연산 자원인 SM(Streaming Multiprocessor)을 얼마나 효율적으로 채우느냐, 즉 'Grid Waves'를 어떻게 관리하느냐가 성능의 핵심입니다. 이번 글에서는 Microsoft의 onnxruntime 레포지토리에 반영된 [CUDA] Tune NVFP4 GEMV tiling for Qwen MTP PR을 통해, 특정 모델 쉐이프(Shape)에서 발생하는 GPU 유휴 자원 문제를 타일링(Tiling) 전략 수정으로 어떻게 해결했는지 분석해 보겠습니다.

문제 진단: 너무 일찍 선택된 'Wide Column Tiles'

기존의 NVFP4 GEMV 디스패치 로직은 단순했습니다. "Wide Column Tiles(ColTiles=4)를 사용했을 때 전체 SM을 한 번이라도 다 채울 수 있는가?"를 기준으로 삼았습니다.

Before (기존 로직)

// onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu

if ((col_tiles + 3) / 4 >= sm_count) {
  // Wide-column: ColTiles=4, 작은 KSplit 선택
  return {std::min(2, windows), 4};
}
// 그 외의 경우: Column-starved 상태로 판단하여 KSplit을 최대화

이 로직의 맹점은 H200(132 SMs) 환경에서 Qwen MTP의 특정 쉐이프(N=17408, K=5120)를 만났을 때 드러납니다.

  • col_tiles = 17408 / 16 = 1088
  • wide_col_blocks = 1088 / 4 = 272
  • 272개 블록은 132개 SM을 가진 H200에서 약 2.06 웨이브(Wave)를 형성합니다.

문제는 2개 웨이브를 채우고 남은 아주 적은 양의 블록(8개)이 마지막 세 번째 웨이브에서 실행될 때, 대다수의 SM이 유휴(Idle) 상태가 된다는 것입니다. 이 경우 차라리 타일 크기를 줄여 블록 수를 4배로 늘리는 것이 GPU 점유율(Occupancy) 측면에서 훨씬 유리합니다.

코드 분석: 3단계 의사결정 트리로의 진화

이번 PR의 핵심은 타일링 로직을 별도의 헤더(matmul_block_scaled_fp4_tiling.h)로 분리하고, 더 정교한 웨이브 기반의 정책을 도입한 것입니다.

After (개선된 로직)

// onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_tiling.h

inline Fp4MmaConfig PickFp4MmaConfig(int n, int k, int sm_count) {
  constexpr int kTargetGridWaves = 4; // 최소 4 웨이브 타겟
  constexpr int kLongReductionGridWaves = 8;
  constexpr int kLongReductionWindows = 64; // K=8192 기준

  const int wide_col_blocks = (col_tiles + 3) / 4;

  // 1단계: Wide Column (ColTiles=4)을 쓰려면 최소 4 웨이브(긴 연산은 8 웨이브)가 보장되어야 함
  if (wide_col_blocks >= kTargetGridWaves * sm_count &&
      (windows < kLongReductionWindows || wide_col_blocks >= kLongReductionGridWaves * sm_count)) {
    return {windows < 2 ? windows : 2, 4};
  }

  // 2단계: Moderate Column (ColTiles=1) 선택, 4 웨이브 보장 시
  if (col_tiles >= kTargetGridWaves * sm_count) {
    const int k_split = windows >= kLongReductionWindows ? 8 : std::min(2, windows);
    return {k_split, 1};
  }

  // 3단계: Column-starved (기존과 동일하게 K-split 최대화)
  // ...
}

주요 변경점

  1. Wave Threshold 도입: 단순히 1 웨이브가 아니라 kTargetGridWaves = 4를 기준으로 삼아, 테일 이펙트(Tail effect)에 의한 성능 저하를 방지했습니다.
  2. Long Reduction 대응: K 값이 큰 경우(K >= 8192), 리덕션(Reduction) 부하가 커지므로 KSplit=8을 사용하여 병렬성을 더 확보하도록 설계했습니다.
  3. 테스트 용이성: CUDA 의존성이 없는 헤더로 로직을 분리하여, 호스트 코드에서 타일링 결정을 유닛 테스트할 수 있게 되었습니다.

성능 최적화의 결과

H200 GPU에서 측정된 성능 수치는 이 최적화가 왜 필요한지를 명확히 보여줍니다.

  • Qwen Gate/Up (N=17408, K=5120): 50.7 us → 46.6 us (1.09x)
  • Qwen Long Reduction (N=17408, K=8192): 61.4 us → 55.0 us (1.12x)

기존 정책이 너무 성급하게 Wide Tile을 선택하여 SM을 놀게 만들었다면, 새로운 정책은 블록 수를 늘려 GPU가 쉴 틈 없이 일하게 만들었습니다.

리뷰어 피드백 및 논의 사항

PR 리뷰 과정에서 유의미한 기술적 논의가 있었습니다.

  1. Environment Variable Overrides: 벤치마킹을 위해 ORT_FP4_GEMV_KSPLIT 등의 환경 변수를 추가했는데, 리뷰어 hariharans29는 ORT_ENFORCE 체크가 매 호출마다 실행되는 점을 지적했습니다. 비록 오버헤드는 작지만, 정적 초기화(Static Initialization) 시점에 한 번만 체크하는 것이 더 깔끔한 패턴이라는 의견이 있었습니다.
  2. Magic Numbers: kTargetGridWaves = 4와 같은 값들은 실험적으로 도출된 값입니다. 리뷰어는 이 값들이 아키텍처(SM 80 vs 90)마다 다를 수 있음을 언급하며, 향후 아키텍처별 파라미터화의 가능성을 열어두었습니다.
  3. CI Slip: 마지막 커밋에서 CI 테스트 통과 개수가 줄어든 현상이 발견되었습니다. 이는 새로운 타일링 테스트 코드가 CUDA 12.8 미만 환경에서 어떻게 동작하는지에 대한 의존성 문제일 가능성이 제기되어, 머지 전 확인이 필요함이 강조되었습니다.

결론: 시니어 엔지니어의 관점

이번 최적화는 "GPU 연산은 단순히 연산량이 아니라, 스케줄링의 효율성 싸움"이라는 진리를 다시 한번 확인시켜 줍니다.

  • 데이터 로컬리티(Weight Locality)를 위해 타일을 크게 가져가는 것이 이론적으로는 좋지만,
  • 병렬성(Parallelism)이 부족하여 SM을 다 채우지 못한다면 과감히 로컬리티를 포기하고 타일을 쪼개야 합니다.

LLM 추론 엔진을 개발하는 엔지니어라면, 자신의 커널이 다양한 쉐이프에서 GPU의 'Waves'를 어떻게 형성하고 있는지 반드시 프로파일링해 보아야 합니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글