[onnxruntime] ONNX Runtime MoEPacked GGUF 오프로드 최적화: 3-Stage Prefill 파이프라인과 F16C·Batch-4 Cold Decode

PR 링크: microsoft/onnxruntime#34114 상태: Merged | 변경: +3275 / -523

들어가며

MoE(Mixture-of-Experts) 모델을 제한된 VRAM에서 돌리려면 일부 expert를 CPU RAM에 두고 필요할 때 GPU로 스트리밍하거나 CPU에서 직접 계산하는 hybrid offload가 필요합니다. ONNX Runtime의 com.microsoft::MoEPacked는 GGUF block-packed 가중치(IQ3_XXS, IQ2_S, IQ4_NL 등)를 이 방식으로 처리합니다.

이 PR은 H200에서 nsys --cuda-graph-trace=node로 timeline을 찍어 두 개의 병목을 찾아냈습니다.

1. Prefill: wave 단위 route/tile metadata의 H2D 복사가 compute stream과 Copy Engine에서 head-of-line blocking을 일으켜, 두 stream이 실제로 겹친 시간이 prefill 구간의 6.6%에 불과했습니다. 두 stream이 모두 놀고 있는 시간이 12.5%였습니다.
2. Decode: GPU가 CPU의 cold expert 결과를 기다리는 MoEPackedDoorbellWaitKernel이 cycle 시간의 60.6%(53.40 ms/cycle)를 차지했습니다. 즉 decode 속도를 CPU dot-product 커널이 결정하고 있었습니다.

코드 분석

1. gguf_vec_dot_cpu.h — 소프트웨어 FP16 변환을 F16C 명령어로 교체

GGUF block은 앞부분에 FP16 scale d를 가집니다. 기존 AVX2 루프는 block마다 이 값을 소프트웨어 함수로 FP32로 바꿨습니다. 비트 마스킹, 지수 재조정, subnormal 처리 분기가 들어가는 함수가 SIMD 루프 한가운데서 block마다 호출되던 구조입니다.

새로 추가된 helper는 다음과 같습니다.

__attribute__((target("avx2,fma,f16c"), always_inline)) inline float Fp16ToFloatF16c(uint16_t bits) noexcept {
  return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128(static_cast<int>(bits))));
}

_mm_cvtph_ps(vcvtph2ps)는 FP16→FP32 변환을 명령어 하나로 처리하고 분기가 없습니다. 이 helper로 129개 호출부를 교체했습니다.

Before

const uint8_t* blk = weight_row + i * 66;
const float d = Fp16BitsToFloatCpu(ReadU16Le(blk)) * 0.25f;

After

const uint8_t* blk = weight_row + i * 66;
const float d = Fp16ToFloatF16c(ReadU16Le(blk)) * 0.25f;

IQ2/IQ3 계열은 256개 원소 block마다 scale을 하나 읽고, Q4_0/IQ4_NL 계열은 32개 원소마다 읽습니다. 블록이 작을수록 scale 변환 비용의 비중이 커집니다. 4-row 커널(VecDotGgufRow4xAvx2)에서는 d0..d3 네 개가 연속으로 바뀌어, 분기 없는 변환의 이점이 더 크게 나타납니다.

const float d0 = Fp16ToFloatF16c(ReadU16Le(b0));
const float d1 = Fp16ToFloatF16c(ReadU16Le(b1));
const float d2 = Fp16ToFloatF16c(ReadU16Le(b2));
const float d3 = Fp16ToFloatF16c(ReadU16Le(b3));

같은 파일에 Batch-4 fused AVX2 INT8 커널도 추가됐습니다(AccumIq3Xxs32Q8Batch4Avx2, VecDotGgufRow4xQ8Batch4Avx2 등). 같은 cold expert로 라우팅된 4개 토큰(MTP verification의 S=8 토큰)을 한 번에 처리합니다. 가중치 block을 한 번만 디코드하고 Q8 활성화 4개와 vpmaddubsw로 곱하므로, 비싼 IQ 코드북 디코딩 비용이 4개 route에 나뉩니다. 남는 route는 Batch-2와 single-route 커널이 처리합니다.

2. moe_packed.cc — 3-Stage Prefill 파이프라인

기존에는 wave 하나의 metadata(sorted_routes, tiles, num_tiles)를 단일 버퍼에 두었습니다. 그래서 다음 wave의 H2D 복사가 현재 wave의 compute가 끝날 때까지 기다려야 했습니다. PR은 이 버퍼를 double-buffer로 바꿨습니다.

• d_wave_sorted_routes[2], d_wave_tiles[2], d_wave_num_tiles[2]를 wave_staging.copy_stream에서 복사한 뒤 h2d_done[b]를 record합니다.
• host의 std::memcpy(pinned staging에 expert 가중치를 채우는 작업)는 h2d_done[b]만 기다립니다.
• copy_stream은 device timeline에서 compute_done[b]를 wait합니다.

결과적으로 "host 적재 → H2D 복사 → GPU 연산"이 3단계 파이프라인으로 겹칩니다. host가 GPU compute 완료를 기다리느라 멈추지 않으며, 순서 보장은 device event가 맡습니다. host expert staging도 w_cnt * 3개 task로 병렬화했고, MoEPackedParallelPool 기본 상한은 24에서 40 코어로 올렸습니다.

3. matmul_packed.cu / gguf_quant_cuda.cuh — 작은 N에서 점유율 확보

MTP verification(M=2..8)의 GDN projection은 N이 작아서 CTA 수가 SM 132개를 채우기에 부족했습니다. N <= 2560이면 MatMulPackedBf16TcWarpKSplit<4, 1>로 dispatch하도록 바꿔 CTA를 1,280개에서 2,560개로 늘렸고, 호출당 지연은 131.1us에서 71.6us로 줄었습니다(-45.4%). ReadU32Dev는 2-byte aligned block(Q8_0, Q6_K 등)에서 uint16x2 __ldg 로드를 일관되게 사용하도록 고쳐 warp divergence를 없앴습니다.

리뷰에서 나온 중요한 논의

리뷰 과정에서 성능보다 정확성과 안전성 문제가 더 많이 지적됐습니다.

• F16C 런타임 체크 누락 (Copilot): 커널은 F16C를 요구하는데 CpuSupportsAvx2Fma()는 AVX2/FMA만 확인했습니다. F16C를 가린 VM에서는 illegal instruction이 발생할 수 있었습니다. 수정 후 CPUIDInfo::HasF16C()와 __builtin_cpu_supports("f16c")를 함께 검사합니다. intrinsic을 새로 쓸 때는 dispatch 조건도 같이 갱신해야 한다는 점을 보여주는 사례입니다.
• Q8 활성화 양자화의 정확도 (apsonawane): 처음에는 shape와 dtype 조건을 만족하면 Q8을 자동으로 켰습니다. 리뷰어가 BF16 입력 [1, 1/512, ...]이면 Q8 block에서 q[1]=0이 되어 출력이 0.73046875에서 0으로 무너지는 반례를 제시했습니다. 결국 Q8은 ep.cuda.moe_packed_cpu_q8_dot=1로 명시적으로 opt-in한 경우에만 쓰이도록 바뀌었습니다. 손실이 있는 최적화는 휴리스틱으로 켜지 말아야 한다는 원칙이 지켜진 셈입니다.
• VRAM 예산 회계: 더블버퍼 metadata bank가 device_working_bytes 계산에서 빠져 있었습니다. 수정 후 WaveRouteDeviceBytes를 예산에 포함하고, cold streaming이 필요 없으면 버퍼를 할당하지 않습니다.
• Neumaier 보상합과 inf: v - t에서 inf - inf = NaN이 생기는 문제는 isfinite() 가드로 막았습니다.
• 스레드 풀 수명: prefill 전용 세션에서는 Shutdown()이 호출되지 않아 최대 39개 worker가 남고 DLL이 pin되는 문제가 있었습니다. 이 세션도 live_states_로 추적하도록 고쳤습니다.
• 테스트 oracle 문제: 참조 구현은 (src / max_abs) * 127을, 프로덕션은 src * (127 / max_abs)를 계산하고 있었습니다. float32에서 두 식은 반올림 경계에서 결과가 달라지고, 이 차이가 실제 CI 실패 원인이었습니다. 참조 구현이 QuantizeBlockQ8Scalar를 그대로 쓰도록 바꿨습니다.

왜 이게 좋은가

수치 (H200, 30 GiB VRAM + 50 GiB RAM, baseline 대비)

지표	Before	After
8K prefill 2-stream overlap	6.6%	24.9% (3.77x)
두 stream 모두 idle	12.5%	4.4%
DoorbellWait	53.40 ms/cyc	26.91 (exact) / 9.66 (tail-sub)
2K TTFT	1.061s	0.786s (+35.0% tok/s)
8K decode	38.87 t/s	69.30 t/s (+78.3%)
CPU ms/step	61.58	32.80

일반적 교훈

1. 먼저 프로파일링한다. "CPU가 느리다"는 짐작 대신 nsys로 stream overlap 비율과 idle bubble을 수치로 확인한 뒤 손을 댔습니다.
2. 공유 버퍼는 직렬화 지점이 된다. metadata 버퍼 하나 때문에 파이프라인 전체가 멈췄습니다. 더블버퍼링을 하고 host 동기화와 device 동기화를 분리하는 것이 copy/compute overlap의 기본입니다.
3. hot loop 안의 작은 변환도 비용이다. scale 하나를 변환하는 소프트웨어 함수도 block마다 호출되면 커널 시간을 크게 차지합니다. 전용 하드웨어 명령어(F16C)가 있으면 쓰고, 반드시 런타임 CPUID로 확인합니다.
4. 비싼 디코딩은 여러 입력에 나눈다. Batch-4 커널처럼 가중치 디코드 한 번을 여러 route에 재사용하면 연산 강도가 올라갑니다.
5. 손실이 있는 가속은 opt-in으로 둔다. 벤치마크가 좋아 보여도 기존 모델의 수치를 조용히 바꾸면 안 됩니다. 리뷰어가 구체적인 반례를 만들어 이를 증명했습니다.

참고 자료

• _mm_cvtph_ps — Intel Intrinsics Guide — F16C FP16→FP32 변환 intrinsic
• CUDA Runtime API: Stream Management — cudaStreamWaitEvent 등 stream 간 동기화
• CUDA Runtime API: Event Management — h2d_done/compute_done 이벤트 기반 동기화
• Nsight Systems User Guide — PR에서 사용한 timeline 프로파일러
• GCC x86 Built-in Functions (__builtin_cpu_supports) — F16C 런타임 검사에 사용

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