[onnxruntime] MoEPacked 멀티토큰 검증 최적화: CPU 콜드 라우트 Batch-2 AVX2와 GPU FC2 디코드 재설계

PR 링크: microsoft/onnxruntime#34067 상태: Merged | 변경: +4376 / -80

들어가며

Speculative decoding은 draft 모델이 제안한 여러 토큰을 target 모델이 한 번에 검증(verify)하는 방식입니다. 이때 target 모델의 MoE 레이어는 M = 4 또는 M = 8 토큰을 동시에 처리합니다. com.microsoft::MoEPacked는 GGUF 저비트 양자화(IQ2_XXS, IQ3_XXS 등) 전문가(expert) 가중치를 직접 다루는 연산자이고, VRAM이 부족하면 일부 expert를 CPU에 두는 hybrid offload를 지원합니다.

문제는 기존 경로가 사실상 M = 1 decode에 맞춰져 있었다는 점입니다.

• CPU cold route: 같은 cold expert로 여러 토큰이 라우팅되어도 가중치 블록을 토큰마다 다시 unpack했습니다. 레이어마다 std::vector를 두 번 할당하는 오버헤드도 있었습니다.
• GPU FC2: 라우트 수가 적은 decode/verify 상황에서도 10.75 KB shared memory staging과 __syncthreads()를 거쳤고, CTA당 출력 채널이 128개라 SM을 충분히 채우지 못했습니다.

이 PR은 두 경로를 모두 손봐서 T_verify(M=8)을 100% GPU 기준 -33.2%, 50/50 Hybrid 기준 -32.5% 줄였습니다.

참고: 제공된 diff는 새로 추가된 gguf_vec_dot_cpu.h(3154줄)의 앞부분에서 잘려 있습니다. 아래 인용 코드는 실제 diff에 있는 부분이고, Batch-2와 CUDA 커널 변경은 PR 설명과 리뷰 내용을 근거로 설명합니다.

코드 분석

1. gguf_vec_dot_cpu.h — 부호 처리를 분기 없는 XOR로

IQ2 계열 포맷은 8개 grid 값과 8비트 sign mask를 쌍으로 저장합니다. 이 PR은 sign 비트를 float sign bit XOR 마스크로 바꾸는 테이블을 constexpr로 컴파일 타임에 생성합니다.

struct SignXorMaskTable {
  alignas(32) uint32_t mask[256][8];
  constexpr SignXorMaskTable() : mask{} {
    for (int s = 0; s < 256; ++s) {
      for (int j = 0; j < 8; ++j) {
        mask[s][j] = ((s >> j) & 1) ? 0x80000000u : 0u;
      }
    }
  }
};
inline constexpr SignXorMaskTable kSignXorMask256{};

이 테이블 덕분에 decode 함수에는 분기가 하나도 없습니다.

inline __m256 DecodeGrid8SignedAvx2(uint64_t grid8, uint8_t signs) {
  const __m256i v_i32 = _mm256_cvtepu8_epi32(_mm_cvtsi64_si128(static_cast<int64_t>(grid8)));
  const __m256 v_f32 = _mm256_cvtepi32_ps(v_i32);
  const __m256 v_sign = _mm256_load_ps(reinterpret_cast<const float*>(kSignXorMask256.mask[signs]));
  return _mm256_xor_ps(v_f32, v_sign);
}

흐름은 다음과 같습니다.

1. 8바이트 grid를 cvtepu8_epi32로 8개 lane에 펼칩니다.
2. cvtepi32_ps로 float로 변환합니다.
3. 32바이트 정렬된 마스크를 _mm256_load_ps로 읽어 XOR합니다.

alignas(32)가 aligned load를 보장하고, 곱셈이나 블렌드 없이 sign bit만 뒤집으므로 결과는 정확히 -x가 됩니다.

2. Single-route 내적: 재사용되지 않는 decode

기준이 되는 single-route 커널(VecDotGgufRowAvx2의 IQ2_XXS 분기)은 다음과 같습니다.

for (int ib32 = 0; ib32 < 8; ++ib32) {
  const uint8_t* aux8 = qs + 8 * ib32;
  const uint32_t aux32 = ReadU32Le(aux8 + 4);
  const float db = d * (0.5f + static_cast<float>(aux32 >> 28));
  const __m256 w0 = DecodeGrid8SignedAvx2(iq2xxs_grid[aux8[0]], ksigns_iq2xs[(aux32 >> 0) & 127u]);
  ...
  __m256 sub0 = _mm256_mul_ps(w0, _mm256_loadu_ps(xp + 0));
  __m256 sub1 = _mm256_mul_ps(w1, _mm256_loadu_ps(xp + 8));
  sub0 = _mm256_fmadd_ps(w2, _mm256_loadu_ps(xp + 16), sub0);
  sub1 = _mm256_fmadd_ps(w3, _mm256_loadu_ps(xp + 24), sub1);
  total_acc0 = _mm256_fmadd_ps(_mm256_add_ps(sub0, sub1), _mm256_set1_ps(db), total_acc0);
}

여기서도 sub0/sub1 두 accumulator로 FMA 의존 체인을 끊는 ILP 기법이 쓰였습니다. 하지만 w0..w3를 만드는 과정, 즉 LUT gather, 부호 XOR, int→float 변환은 활성값 한 행에 대해서만 사용됩니다.

3. Batch-2 fused dot product (VecDotGgufRow4xQ8Batch2Avx2, VecDotGgufGateUpRow2xQ8Batch2Cpu)

PR에 따르면 cold-route reuse 비율 C/U(cold route 수 / unique cold expert 수)는 M=4에서 1.391.48x, M=8에서 1.862.07x로 측정되었습니다. 즉 검증 배치에서는 여러 토큰이 같은 expert를 자주 공유합니다.

그래서 ColdExpertGroup의 grp.count >= 2인 경우, 가중치 sub-block을 레지스터에서 한 번만 unpack하고 두 route의 Q8 활성값에 동시에 곱합니다. 개념을 그림으로 나타내면 다음과 같습니다(실제 코드가 아닌 설명용 의사코드).

// 개념도: w = decode(block) 1회 → 두 route에 재사용
w = DecodeSubBlock(blk);
acc_a = fma(w, q8_a, acc_a);
acc_b = fma(w, q8_b, acc_b);

두 route의 계산 순서는 single-route와 동일하게 유지되므로, 결과는 0 ULP 차이, 즉 bit-exact입니다.

4. Zero-heap 스택 버퍼와 병렬 Q8 양자화

ExecuteMoEPackedColdRoutesCpuWithScratch는 num_cold_routes <= 128이면 sorted_c_stack[128]과 groups_stack[128]을 씁니다. 그 결과 48-layer verify 한 번에 발생하던 heap 할당/해제 96회가 사라집니다.

활성값 양자화(BF16→FP32→QuantizeRowQ8Cpu)는 num_unique_toks >= 2일 때 토큰 단위로 병렬화됩니다. Q8 블록 구조는 내적에 필요한 bsum을 미리 계산해 둡니다.

struct alignas(16) GgufQ8Block {
  float d;
  int32_t bsum;
  int8_t qs[32];
};

5. moe_packed.cu — MoEPackedFc2Bf16TcKernel<1>

kMaxSubTiles == 1(total_routes <= 128)일 때 GPU FC2 경로가 세 가지로 바뀝니다.

• 10.75 KB s_act_u4 staging 버퍼, 정수 나눗셈 루프, K-loop __syncthreads()를 제거하고, fc1_out을 128-bit __ldg로 직접 읽습니다. 활성값이 작으므로 read-only cache 경로로 충분합니다.
• IQ 계열 grid LUT를 shared memory(s_lut)에 올리고 DecodeGgufSubBlock8SmemLutTpl<kDownType>로 decode합니다.
• CTA 구성을 8 warp(128 채널)에서 4 warp(64 채널)로 줄여 grid 크기를 2배로 늘립니다. decode처럼 작업량이 작을 때 SM 점유율을 높이기 위해서입니다.

리뷰에서 나온 중요한 논의

• "production caller가 없다": Copilot은 Q8 cold-route helper가 실제 hybrid 경로(RunColdDecodeStep)에서 호출되지 않는다고 지적했습니다. 그대로였다면 최적화는 벤치마크에서만 존재했을 것입니다. 작성자는 ExecuteMoEPackedColdRoutesCpuWithScratch를 RunColdDecodeStep에 연결하고 hybrid EP parity 테스트로 검증했습니다.
• 4-warp grid의 경계 조건: hidden size가 64의 배수가 아니면 마지막 CTA에 비활성 warp가 생깁니다. 이 경우를 검증하기 위해 hidden_size=96 테스트 케이스가 추가됐습니다.
• bit-exact 주장 검증: K=2560의 multi-superblock 환경에서 모든 Batch-2 포맷과 mixed gate/up 조합을 single-route 호출 2회와 비교하는 테스트가 추가됐습니다.
• 빌드 이식성: void를 반환하는 함수를 if 조건으로 쓴 오류가 있었습니다. 또 MSVC /arch:AVX2 환경에서는 __attribute__((target)) 기반 함수가 선언되지 않습니다. 둘 다 ORT_GGUF_CPU_HAS_AVX2_TARGET과 CpuSupportsAvx2Fma() 가드로 수정됐습니다.

왜 이게 좋은가

지표 (UD-IQ2_XS)	Before	After
100% GPU T_verify(M=8)	50.32–51.68 ms	32.90–35.27 ms (-33.2%)
50/50 Hybrid T_verify(M=8)	56.08–80.69 ms	38.69–53.59 ms (-32.5%)

End-to-end로는 fixed M=8에서 최대 197.4 tok/s(3.93x)를 기록했고, greedy token parity는 64/64를 유지했습니다.

일반적인 교훈

1. 배치 내 재사용을 찾아라: 저비트 양자화 커널에서는 비용의 대부분이 unpack/decode입니다. 같은 가중치를 공유하는 입력이 2개 이상이면 decode를 한 번만 하는 것이 가장 효과적인 최적화입니다.
2. 작은 M에는 작은 CTA: decode 상황에서는 타일을 키우는 것보다 CTA 수를 늘려 SM을 채우는 편이 낫습니다.
3. 필요 없는 staging은 걷어내라: 데이터가 작으면 shared memory와 barrier가 오히려 오버헤드가 됩니다.
4. 최적화는 실제 호출 경로에 연결되어야 의미가 있다: 리뷰가 잡아낸 "caller 없음" 문제가 그 예입니다.

참고 자료

• Intel Intrinsics Guide — _mm256_fmadd_ps, _mm256_xor_ps, _mm256_cvtepu8_epi32 등 PR에서 사용한 AVX2/FMA intrinsic 공식 문서
• CUDA C++ Programming Guide — __ldg, shared memory, __syncthreads 관련 공식 문서
• ONNX Runtime Contrib Operators — com.microsoft 도메인 contrib 연산자 문서

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