[onnxruntime] 소비자 GPU에서 QMoE 오프로드 최적화: Windows WDDM 대응, Tail-Expert 가지치기, AVX2 INT2/INT3 커널 개선

PR 링크: microsoft/onnxruntime#33916 상태: Merged | 변경: +810 / -160

들어가며

수백억 파라미터 규모의 Mixture-of-Experts(MoE) 모델을 24 GB VRAM의 소비자 GPU에서 돌리려면, 모든 expert를 GPU에 올릴 수 없습니다. ONNX Runtime의 QMoE CPU/GPU offload는 자주 쓰이는 expert는 GPU에 두고, 나머지 "cold expert"는 CPU 메모리(혹은 파일 매핑)에 둔 채 CPU에서 직접 계산하거나 필요할 때 PCIe로 올려 보내는 구조입니다.

이 PR은 RTX 4090(24 GB) + i9-13900(30 GiB RAM) 환경에서 드러난 문제를 한꺼번에 다룹니다.

• Windows 파일 매핑 경로 버그: source.file_path == nullptr이지만 이미 매핑된 메모리가 들어오면 FileExpertStore 초기화가 실패했습니다.
• Pinned memory 예산 회계 오류: promotion용 bounce buffer(16 MiB)와 decode용 headroom(24 MiB)이 섞여 계산되어 decode 단계가 pinned memory를 못 얻는 starvation이 생겼습니다.
• 확률이 낮은 tail expert에 대한 낭비: top-k 라우팅에서 확률이 아주 낮은 expert도 CPU cold path로 계산되어 decode 지연의 주범이 되었습니다.
• CPU AVX2 INT2/INT3 GEMV 커널의 비효율: activation load가 column 루프 안에서 반복되었습니다.

코드 분석

1. qmoe_shared_layout.cc — Windows 매핑 메모리 수용

기존에는 파일 경로가 없고 작은 resident 복사도 불가능하면 곧바로 return false였습니다. 이제는 VirtualQuery로 해당 주소가 실제로 commit된 매핑 메모리인지 확인한 뒤 그대로 사용합니다.

// Before
} else if (length <= (4ULL << 20) && read_opt_enabled) {
  bank->resident.assign(...);
} else {
  return false;
}

// After
} else if (source.data != nullptr) {
  MEMORY_BASIC_INFORMATION mbi{};
  if (VirtualQuery(source.data, &mbi, sizeof(mbi)) == 0 || mbi.State != MEM_COMMIT ||
      (mbi.Type != MEM_MAPPED && !read_opt_enabled)) {
    return false;
  }
  bank->mapped_ptr = static_cast<const uint8_t*>(source.data);
  bank->is_file_backed = (mbi.Type == MEM_MAPPED);
}

핵심은 복사 없이 OS 파일 매핑을 그대로 쓰는 것입니다. 수 GiB 규모의 expert 가중치를 private heap에 다시 commit하면 30 GiB RAM 머신에서는 곧바로 페이징이 발생합니다. MEM_MAPPED면 OS page cache가 관리하므로 메모리 압박 시 깨끗하게 버릴 수 있고, is_file_backed 플래그로 이후 VirtualUnlock 기반 page eviction도 가능해집니다.

2. 세션 옵션 — Tail-Expert 확률 임계값 가지치기

새 세션 옵션 두 개가 추가되었습니다.

static const char* const kOrtSessionOptionsMoeOffloadMinExpertWeight = "session.moe.offload.min_expert_weight";
// - "1" (default): retain GPU-resident (`slot >= 0`) tail experts during decode at zero wall-time cost and prune only
//   CPU-cold (`slot < 0`) tail experts whose routing probability is below `min_expert_weight`.
static const char* const kOrtSessionOptionsMoeOffloadSkipTailColdOnly = "session.moe.offload.skip_tail_cold_only";

동작은 단순합니다. SoftmaxTopK 이후 rank ≥ 1인 expert 중 확률이 min_expert_weight 미만이면 topk_scales[route] = 0.0f, topk_indices[route] = -2로 마킹합니다. top-1은 항상 유지합니다. skip_tail_cold_only가 영리한 부분인데, GPU에 이미 상주하는 expert는 계산 비용이 거의 0이므로 남겨 두고, 비싼 CPU cold expert만 잘라냅니다. 정확도 손실은 최소화하고 wall-time의 병목만 제거하는 설계입니다.

또한 cold_count가 0이면 input_dev -> host_input_out PCIe 복사 자체를 건너뛰고, merge/finalize 커널은 weight가 0인 route의 초기화되지 않은 contribution row를 읽지 않도록 수정되었습니다.

3. qmoe_shared_layout_int3_avx2.cc — unpack과 dequant 융합

기존 코드는 정수 코드를 float로 변환했다가, lookup이 필요한 경우 다시 int로 되돌려 permutevar를 수행했습니다.

// Before
inline __m256 Unpack8Int3WeightsAvx2(...) {
  return _mm256_cvtepi32_ps(
      _mm256_sub_epi32(_mm256_and_si256(_mm256_srlv_epi32(...), mask7), four));
}
template <bool DeferredScale>
inline __m256 DequantInt3Avx2(__m256 codes, __m256 table, __m256i four) {
  if constexpr (DeferredScale) return codes;
  return _mm256_permutevar8x32_ps(table, _mm256_add_epi32(_mm256_cvtps_epi32(codes), four));
}

// After
template <bool DeferredScale>
inline __m256 UnpackDequant8Int3WeightsAvx2(..., const __m256& table) noexcept {
  const __m256i ucodes = _mm256_and_si256(_mm256_srlv_epi32(_mm256_set1_epi32(static_cast<int>(u24)), shifts), mask7);
  if constexpr (DeferredScale) {
    return _mm256_cvtepi32_ps(_mm256_sub_epi32(ucodes, four));
  } else {
    return _mm256_permutevar8x32_ps(table, ucodes);
  }
}

lookup 경로에서 sub → cvtepi32_ps → cvtps_epi32 → add라는 무의미한 왕복이 사라졌습니다. 0~7의 unsigned code를 바로 8-lane table 인덱스로 쓰면 되기 때문입니다. AVX-512 버전도 _mm512_permutexvar_ps로 동일하게 정리되었습니다.

4. 진짜 8-column Int3Group8Avx2

기존 8-column 함수는 사실 4-column 함수를 두 번 호출하는 래퍼였습니다. 즉 activation a_natural을 두 번 읽었습니다.

// Before
Int3Group4Avx2<...>(w_base, ..., n0, k, a_natural, first);
Int3Group4Avx2<...>(w_base + 4 * k * 3 / 8, ..., n0 + 4, k, a_natural, second);

// After (발췌)
for (int64_t kb = 0; kb < blocks; ++kb) {
  __m256 a0[kChunks32], a1[kChunks32], a2[kChunks32], a3[kChunks32];
  for (int ch = 0; ch < kChunks32; ++ch) {
    a0[ch] = _mm256_loadu_ps(a_natural + kb * BlockSize + ch * 32 + 0);
    ...
  }
  auto step_col = [&](int c) QMOE_AVX2_LAMBDA_INLINE { ... };
  step_col(0); ... step_col(7);
}

블록당 activation을 한 번 로드해 8개 column이 공유하고, 4블록마다 다음 weight를 _mm_prefetch합니다. 이론상 메모리 트래픽이 줄어드는 전형적인 register blocking입니다.

리뷰에서 드러난 트레이드오프

이 PR의 리뷰는 매우 구체적이었고, 배울 점이 많습니다.

SIMD 최적화는 측정이 전부다. baijumeswani는 GCC13 Release, EPYC 7V12에서 실제 Int3Fc2Range 호출을 측정해, activation load hoisting 후 2-row INT3 케이스가 약 322 → 400 µs(약 24%) 느려졌다고 보고했습니다. AVX2는 ymm 레지스터가 16개뿐인데, BlockSize=128이면 a0~a3[4]만으로 16개를 소모하고 acc[8]까지 더하면 spill이 불가피합니다. 스택 프레임이 커졌다는 관찰도 이를 뒷받침합니다. INT2 8-column 경로 역시 1-row decode에서 124 → 131 µs(약 6%) 느려졌습니다. "load를 공유하면 빠르다"는 직관이 register pressure 앞에서 깨지는 전형적 사례입니다.

명시적 예산은 상한이어야 한다. 자동 튜닝이 사용자가 지정한 device_expert_bytes나 host cap을 넘어 확장하는 문제가 여러 번 지적되었습니다. 16 GiB RAM/4 GiB cap 시뮬레이션에서 bounded cache가 꺼지며 8 GiB cold 데이터가 admit되거나, 생성자가 admitted cache 크기를 0으로 덮어써 InitializeBoundedHostCache가 아예 생략되는 버그도 있었습니다. 후자는 tianleiwu가 #33935에서 수정하고 64-expert 회귀 테스트를 추가했습니다.

설정 우선순위의 일관성. 세션 옵션으로 min_expert_weight=0을 명시해도 환경변수 ORT_QMOE_MIN_EXPERT_WEIGHT=0.065가 다시 덮어쓰는 문제, 그리고 cold-only 가지치기가 doorbell 경로에만 적용되고 callback 경로에는 빠져 출력이 2.125 vs 0.921875로 달라지는 문제도 지적되었습니다. 정확도에 영향을 주는 옵션일수록 opt-out이 확실해야 합니다.

그 밖에 noexcept 함수 안에서 env 파싱이 throw하면 프로세스가 종료되는 문제, SM89 INT8 Tensor-Core 경로가 테스트에서 SM90만 실행되어 검증되지 않는 문제도 언급되었습니다.

왜 이게 좋은가

PR 설명의 벤치마크(RTX 4090 + 30 GiB RAM, Qwen 계열 모델):

시나리오	기본	tail-skip
2K TTFT	8.92 s	7.55 s (0.08)
8K TTFT	14.42–15.32 s	10.84–12.30 s
MMLU-Pro 100 정확도	86.0%	84.0% (0.065)
MMLU-Pro decode	27.39 tok/s	34.70 tok/s
wall time	3,196 s	2,099 s

정확도 2%p를 내주고 wall time을 약 34% 줄였습니다. 128K 컨텍스트에서도 peak VRAM 21,345 MiB로 24 GB 안에 들어옵니다.

일반적 교훈

1. 가장 비싼 경로만 골라 자른다: tail expert 전체가 아니라 CPU-cold tail만 제거해 정확도/속도 트레이드오프를 최적화했습니다.
2. 복사보다 매핑: OS의 file-backed mapping을 활용하면 private commit 없이 메모리 압박에 유연하게 대응할 수 있습니다. 단, working-set 상한은 별도로 보장해야 합니다.
3. 불필요한 형 변환 왕복 제거: int→float→int 같은 왕복은 SIMD 핫루프에서 순수한 낭비입니다.
4. Register blocking은 레지스터 수와 함께 설계: AVX2의 16개 ymm 한계를 넘으면 hoisting이 오히려 역효과를 냅니다. row 수·block 크기별 dispatch와 실측이 필수입니다.
5. 자동 튜닝은 사용자 명시값을 존중: 명시적 예산은 상한으로, 명시적 0은 진짜 비활성화로 처리해야 합니다.

참고 자료

• VirtualQuery (Win32) — 매핑 메모리 상태 확인에 사용된 API
• MEMORY_BASIC_INFORMATION — MEM_COMMIT/MEM_MAPPED 판별 구조체
• VirtualUnlock — 업로드된 host page eviction에 사용
• PrefetchVirtualMemory — Windows mmap prefetch 기본 활성화
• Intel Intrinsics Guide — _mm256_permutevar8x32_ps, _mm256_srlv_epi32 등 AVX2 intrinsic
• ONNX Runtime Session Options — 세션 옵션 설정 API

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