[onnxruntime] [CUDA] QMoE MXFP4/NVFP4 가중치 역양자화 성능 최적화: Coalesced Memory Access의 힘

PR 링크: microsoft/onnxruntime#31349 상태: Merged | 변경: +147 / -28

들어가며

최근 대규모 언어 모델(LLM)에서 Mixture-of-Experts(MoE) 아키텍처는 모델의 파라미터 수를 크게 늘리면서도 계산 비용을 효율적으로 유지하는 강력한 방법으로 주목받고 있습니다. 그러나 MoE 모델의 효율적인 운영에는 여러 가지 기술적 도전 과제가 따르며, 그중 하나는 가중치 역양자화(weight dequantization) 과정에서의 성능 병목입니다. 특히 MXFP4/NVFP4와 같은 저정밀도(low-precision) 가중치를 사용하는 경우, 이를 부동 소수점(float)으로 변환하는 과정이 전체 추론 시간의 상당 부분을 차지할 수 있습니다.

오늘 분석할 microsoft/onnxruntime 레포지토리의 PR "[CUDA] Coalesce QMoE MXFP4/NVFP4 weight dequantization"은 바로 이 문제를 해결하기 위한 중요한 최적화입니다. 기존 LaunchQMoEDequantizeFp4Weights 및 LaunchQMoEDequantizeNvfp4Weights 커널은 비효율적인 메모리 접근 패턴으로 인해 HBM(High Bandwidth Memory) 피크 성능의 약 5%만 활용하며, 40-layer NVFP4 MoE 모델의 prefill 과정에서 307ms라는 상당한 시간을 소모했습니다. 이 PR은 메모리 접근 패턴을 최적화하여 이 병목을 제거하고 MoE 모델의 추론 속도를 크게 향상시키는 것을 목표로 합니다.

코드 분석

이 PR은 주로 onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu 파일에 대한 변경사항을 포함하며, 크게 두 가지 핵심 최적화와 몇 가지 부수적인 개선 사항으로 구성됩니다.

1. DecodeFp4E2M1 함수의 분기 없는(Branch-Free) 구현

기존 DecodeFp4E2M1 함수는 kValues 배열을 사용하여 FP4 코드를 float 값으로 변환했습니다. 이는 룩업 테이블(table lookup) 방식으로 구현되었는데, GPU 아키텍처에서 런타임 인덱싱되는 로컬 테이블은 constant-bank load로 컴파일되어 워프(warp) 내에서 주소가 다를 때마다 하드웨어가 이를 다시 재생(replay)해야 하는 비효율성이 있었습니다. 인접한 가중치들이 코드를 공유하는 경우가 드물기 때문에 이러한 비효율성은 더욱 두드러졌습니다.

이 PR에서는 DecodeFp4E2M1을 직접 비트 조작을 통해 분기 없이 구현하여 이러한 오버헤드를 제거했습니다.

Before:

__device__ __forceinline__ float DecodeFp4E2M1(uint8_t code) {
  constexpr float kValues[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f};
  float value = kValues[code & 0x7];
  return (code & 0x8) ? -value : value;
}

After:

__device__ __forceinline__ float DecodeFp4E2M1(uint8_t code) {
  const uint32_t e = (code >> 1) & 0x3u;
  const uint32_t m = code & 0x1u;
  const uint32_t sign = static_cast<uint32_t>(code & 0x8u) << 28;
  const uint32_t normal = ((126u + e) << 23) | (m << 22);
  const uint32_t subnormal = m ? (126u << 23) : 0u;
  return __uint_as_float(sign | (e ? normal : subnormal));
}

개선점:

• 분기 제거: if 문이나 테이블 룩업 대신 비트 연산으로 직접 float 비트 패턴을 조립하여 CPU의 분기 예측 실패와 유사한 GPU의 워프 재실행(replay) 오버헤드를 제거했습니다.
• 레지스터 활용: 모든 연산이 레지스터 내에서 이루어져 메모리 접근을 최소화합니다.
• 안전성 강화 (리뷰 반영): 초기 구현에서 uint32_t를 int로 변환하는 과정에서 발생할 수 있는 구현 정의(implementation-defined) 동작의 위험을 __uint_as_float를 사용하여 uint32_t 비트 패턴을 직접 float으로 캐스팅함으로써 제거했습니다. 이는 리뷰어의 중요한 피드백이 반영된 결과입니다.

2. QMoEDequantizeFp4WeightsVecKernel을 통한 Coalesced Memory Access 최적화

이 PR의 핵심은 새로운 벡터화 커널인 QMoEDequantizeFp4WeightsVecKernel을 도입하여 메모리 접근 패턴을 근본적으로 개선하는 것입니다. 기존 스칼라 커널은 하나의 스레드가 하나의 출력 요소를 담당하고 k 인덱스가 가장 빠르게 변하는 방식으로 매핑되었습니다. 이 방식은 [E, K, N/2] 형태로 저장된 압축된 가중치와 [E, N, K] 형태의 출력 간의 불일치로 인해 비효적입니다. 즉, 연속적인 스레드들이 packed_n만큼 떨어진 압축 바이트를 읽게 되어, 각 스레드가 페이로드의 절반 바이트를 위해 별도의 32바이트 섹터를 접근해야 했습니다. 이는 64-bit div/mod 연산과 함께 HBM 대역폭 활용률을 극히 낮추는 원인이었습니다.

새로운 커널은 출력(output)을 타일링(tiling)하는 방식을 채택하여 이 문제를 해결합니다.

새로운 커널 정의 (주요 부분):

// Tile shape for QMoEDequantizeFp4WeightsVecKernel. kTileN = 64 rows is exactly 32 packed bytes,
// so a block consumes every sector of the packed weights it touches; kTileK / kVecK = 8 threads
// per row make each store 128 contiguous bytes. Block is (8, 64) = 512 threads.
constexpr int kQMoEDequantizeFp4VecK = 8;
constexpr int kQMoEDequantizeFp4TileN = 64;
constexpr int kQMoEDequantizeFp4TileK = 64;

// Coalescing-optimized FP4 / NVFP4 weight dequantization.
// ... (설명 생략)
template <typename T, int kBlockSize, bool kE4M3Scale>
__global__ void QMoEDequantizeFp4WeightsVecKernel(
    const uint8_t* __restrict__ packed_weights,
    const uint8_t* __restrict__ block_scales,
    const float* __restrict__ global_scales,
    T* __restrict__ output,
    int n,
    int k) {
  constexpr int kVecK = kQMoEDequantizeFp4VecK;
  const int row = static_cast<int>(blockIdx.x) * kQMoEDequantizeFp4TileN + static_cast<int>(threadIdx.y);
  if (row >= n) {
    return;
  }
  const int k_base = static_cast<int>(blockIdx.y) * kQMoEDequantizeFp4TileK +
                     static_cast<int>(threadIdx.x) * kVecK;
  const int expert = static_cast<int>(blockIdx.z);

  const int packed_n = n >> 1;
  const int shift = (row & 1) ? 4 : 0;
  const int64_t weight_base = (static_cast<int64_t>(expert) * k + k_base) * packed_n + (row >> 1);

  const int scale_k = k / kBlockSize;
  const uint8_t scale_code = block_scales[(static_cast<int64_t>(expert) * n + row) * scale_k + k_base / kBlockSize];
  const float scale = (kE4M3Scale ? DecodeFloat8E4M3FN(scale_code) : DecodeUE8M0(scale_code)) * global_scales[expert];

  // uint4 storage keeps the staging buffer 16-byte aligned for the vector store below.
  uint4 staged[kVecK * sizeof(T) / sizeof(uint4)];
  T* values = reinterpret_cast<T*>(staged);
#pragma unroll
  for (int j = 0; j < kVecK; ++j) {
    const uint8_t packed = packed_weights[weight_base + static_cast<int64_t>(j) * packed_n];
    values[j] = static_cast<T>(DecodeFp4E2M1(static_cast<uint8_t>((packed >> shift) & 0x0F)) * scale);
  }

  uint4* dst = reinterpret_cast<uint4*>(output + (static_cast<int64_t>(expert) * n + row) * k + k_base);
#pragma unroll
  for (int v = 0; v < static_cast<int>(kVecK * sizeof(T) / sizeof(uint4)); ++v) {
    dst[v] = staged[v];
  }
}

개선점:

• 출력 타일링: 블록(block)이 64개의 연속적인 행(row)과 64개의 연속적인 k를 커버하도록 설계되었습니다. threadIdx.x는 k 그룹을, threadIdx.y는 행을 선택합니다.
• Coalesced Store: 각 스레드는 8개의 값을 16바이트 uint4로 출력합니다. 같은 행의 8개 스레드(lane)가 모여 하나의 128바이트 연속 저장(contiguous store)을 수행합니다. 이는 하나의 워프(warp)가 32개의 개별 요청 대신 4개의 행에 대해 4개의 요청만 발행하도록 하여 메모리 쓰기 효율을 극대화합니다.
• Coalesced Read: kTileN = 64 행은 정확히 32개의 압축 바이트에 해당하므로, 블록은 자신이 접근하는 모든 압축 섹터를 완전히 소비합니다. 이는 packed_weights 읽기 효율을 높입니다.
• 메모리 요청 감소: PR 설명에 따르면, 이전 시도에서 단순히 스레드당 처리하는 k의 개수를 늘리는(16개, 32개) 방식은 Max Bandwidth의 95%를 달성했지만, DRAM 처리량은 26%에 불과했습니다. 이는 메모리 *요청(requests)*에 의해 병목이 발생했기 때문입니다. 새로운 타일링 방식은 메모리 요청당 처리하는 바이트 수를 늘려 이 문제를 해결합니다.
• 효율적인 인덱스 계산: 인덱스 분해(index decomposition)가 순수한 그리드 산술(pure grid arithmetic)로 이루어져 정수 나눗셈(integer division)이 필요 없어 계산 오버헤드가 줄어듭니다.
• 스케일 값 재사용: 블록 스케일 바이트와 expert별 글로벌 스케일은 kVecK 값당 한 번만 읽히므로 메모리 접근 횟수가 줄어듭니다.

3. 새로운 커널 호출 로직 및 폴백(Fallback)

LaunchQMoEDequantizeFp4WeightsImpl 및 LaunchQMoEDequantizeNvfp4WeightsImpl 함수는 QMoEDequantizeFp4VecApplies 헬퍼 함수를 통해 새로운 벡터화 커널을 조건부로 호출하도록 변경되었습니다. 이는 특정 조건(예: n이 짝수, k가 kQMoEDequantizeFp4TileK의 배수 등)을 만족할 때만 최적화된 경로를 사용하고, 그렇지 않은 경우에는 기존 스칼라 커널을 폴백으로 사용하도록 합니다.

Before (LaunchQMoEDequantizeFp4WeightsImpl):

void LaunchQMoEDequantizeFp4WeightsImpl(
    const uint8_t* packed_weights,
    const uint8_t* block_scales,
    const float* global_scales,
    T* output,
    int num_experts,
    int n,
    int k,
    cudaStream_t stream) {
  int64_t total = static_cast<int64_t>(num_experts) * n * k;
  constexpr int block = 256;
  int grid = onnxruntime::narrow<int>((total + block - 1) / block);
  QMoEDequantizeFp4WeightsKernel<<<grid, block, 0, stream>>>(
      packed_weights, block_scales, global_scales, output, num_experts, n, k);
}

After (LaunchQMoEDequantizeFp4WeightsImpl):

void LaunchQMoEDequantizeFp4WeightsImpl(
    const uint8_t* packed_weights,
    const uint8_t* block_scales,
    const float* global_scales,
    T* output,
    int num_experts,
    int n,
    int k,
    cudaStream_t stream) {
  constexpr int block = 256;
  if (QMoEDequantizeFp4VecApplies<32>(num_experts, n, k)) {
    const dim3 tile_block(kQMoEDequantizeFp4TileK / kQMoEDequantizeFp4VecK, kQMoEDequantizeFp4TileN);
    const dim3 tile_grid((n + kQMoEDequantizeFp4TileN - 1) / kQMoEDequantizeFp4TileN,
                         k / kQMoEDequantizeFp4TileK, num_experts);
    QMoEDequantizeFp4WeightsVecKernel<T, 32, false><<<tile_grid, tile_block, 0, stream>>>(
        packed_weights, block_scales, global_scales, output, n, k);
    CUDA_CALL_THROW(cudaGetLastError());
    return;
  }
  int64_t total = static_cast<int64_t>(num_experts) * n * k;
  int grid = onnxruntime::narrow<int>((total + block - 1) / block);
  QMoEDequantizeFp4WeightsKernel<<<grid, block, 0, stream>>>(
      packed_weights, block_scales, global_scales, output, num_experts, n, k);
  CUDA_CALL_THROW(cudaGetLastError());
}

개선점:

• 유연한 최적화 적용: 모든 경우에 최적화된 커널을 강제하지 않고, 특정 조건(예: n이 짝수, k가 타일 크기의 배수 등)을 만족할 때만 벡터화 커널을 사용합니다. 이는 기존의 다양한 모델 형태에 대한 호환성을 유지하면서도 성능 이점을 취할 수 있게 합니다.
• 오류 처리 강화: cudaGetLastError() 호출을 추가하여 커널 실행 후 발생할 수 있는 CUDA 런타임 오류를 즉시 감지하고 처리할 수 있도록 했습니다.

4. DecodeFloat8E4M3FN 함수의 위치 변경

DecodeFloat8E4M3FN 함수가 QMoEDequantizeFp4WeightsVecKernel에서 사용될 수 있도록 파일 내에서 더 앞쪽으로 이동되었습니다. 이는 동작 변경 없이 코드 구조를 개선한 것입니다.

왜 이게 좋은 최적화인가?

이 PR은 CUDA 프로그래밍에서 성능 최적화의 핵심 원칙들을 모범적으로 보여줍니다.

1. 메모리 접근 패턴 최적화 (Coalesced Memory Access): GPU 컴퓨팅에서 가장 중요한 성능 요소 중 하나는 전역 메모리(global memory) 접근 패턴입니다. 이 PR은 스레드 블록과 그리드 매핑을 재설계하여 메모리 읽기/쓰기가 Coalesced Access 패턴을 따르도록 했습니다. 이는 여러 스레드가 인접한 메모리 위치에 동시에 접근하여 하나의 큰 트랜잭션으로 처리되게 함으로써, HBM 대역폭 활용률을 극대화하고 메모리 요청 수를 줄입니다. 기존 5% HBM 피크 활용률에서 크게 개선되어, MoE 모델의 prefill 단계에서 가중치 역양자화가 더 이상 병목이 되지 않도록 합니다.

2. 분기 없는 코드 (Branch-Free Code): DecodeFp4E2M1 함수의 개선은 GPU의 SIMT(Single Instruction, Multiple Thread) 아키텍처에서 분기문이 성능에 미치는 악영향을 회피하는 좋은 예시입니다. 분기문은 워프 내 스레드들이 서로 다른 경로를 실행하게 하여 워프 발산(warp divergence)을 유발하고, 이는 유휴 스레드를 발생시켜 효율을 저하시킵니다. 비트 조작을 통한 분기 없는 구현은 이러한 문제를 해결하고 예측 가능한 실행 경로를 보장합니다.

3. 프로파일링 기반 최적화: PR 설명에서 Nsight Compute를 사용하여 이전 시도(단순 벡터 폭 확장)가 메모리 요청에 의해 병목이 발생했음을 확인했다는 점은 매우 중요합니다. 이는 막연한 추측이 아닌 실제 프로파일링 데이터를 기반으로 문제의 근본 원인을 파악하고, 그에 맞는 해결책을 제시했음을 의미합니다. 최적화는 항상 측정(measurement)에서 시작되어야 한다는 교훈을 줍니다.

4. 유연한 적용: 최적화된 커널이 모든 입력 형태에 대해 강제되지 않고, QMoEDequantizeFp4VecApplies 함수를 통해 특정 조건을 만족할 때만 적용되도록 한 점은 견고한 소프트웨어 설계의 좋은 예시입니다. 이는 광범위한 사용 사례에 대한 호환성을 유지하면서도 특정 고성능 시나리오에서 이점을 취할 수 있게 합니다.

마치며

이 PR은 ONNX Runtime에서 MoE 모델의 MXFP4/NVFP4 가중치 역양자화 성능을 획기적으로 개선하는 중요한 기여를 합니다. 특히 CUDA 커널 최적화의 핵심인 Coalesced Memory Access 패턴의 중요성을 명확히 보여주며, 분기 없는 코드 구현과 프로파일링 기반의 문제 해결 접근 방식은 모든 GPU 개발자에게 좋은 교훈이 될 것입니다. 이러한 최적화를 통해 MoE 모델은 더욱 빠르고 효율적으로 동작할 수 있게 되어, 대규모 AI 모델의 실제 배포 및 활용에 큰 도움이 될 것입니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#device-memory-accesses
• https://docs.nvidia.com/cuda/cuda-math-api/index.html#device-intrinsics
• https://docs.nvidia.com/cuda/cuda-math-api/index.html#ldexpf
• https://docs.nvidia.com/cuda/cuda-math-api/index.html#exp2f

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