[onnxruntime] ONNX Runtime WebGPU: AMD GPU를 위한 Subgroup Matrix MatMulNBits 최적화

PR 링크: microsoft/onnxruntime#32796 상태: Merged | 변경: +295 / -146

들어가며

최근 microsoft/onnxruntime 레포지토리에 흥미로운 WebGPU 최적화 PR이 병합되었습니다. 기존의 MatMulNBits 연산은 주로 wave32 설정에 의존하고 있었으나, 이는 AMD D3D12 어댑터의 잠재력을 충분히 활용하지 못하는 한계가 있었습니다. 본 PR은 AMD 환경에서 wave64를 활용하고, 워크로드의 M, N, K 차원에 따라 타일 크기와 K-blocking 전략을 동적으로 조정함으로써 성능을 극대화하는 것을 목표로 합니다.

코드 분석

1. subgroup_matrix_matmul_nbits.cc: 동적 타일링 및 Wave64 도입

가장 핵심적인 변화는 ApplySubgroupMatrixMatMulNBits 함수 내에서 AMD 어댑터에 대해 wave64를 선택하고, 워크로드 크기에 따라 타일링 전략을 분기하는 로직입니다.

// Before: 고정된 128x128 타일 사용
tile_m = 128;
tile_n = 128;
workgroup_size = 128;

// After: 워크로드에 따른 동적 타일링
if (config.subgroupSize == 64) {
  if (M <= 64) {
    tile_m = 64;
    tile_n = N == 64 ? 64 : 128;
    tile_size_k = 64;
  } else {
    const bool use_four_waves = M > 384 && row_tiles_128 % 2 == 0;
    tile_m = use_four_waves ? 256 : 128;
    tile_n = 64;
    tile_size_k = M < 256 ? 32 : (use_four_waves || M <= 384 ? 128 : 64);
  }
}

이 변경은 단순히 wave64로 전환하는 것을 넘어, M이 작을 때 불필요한 패딩을 방지하고, 큰 M에 대해서는 4개의 웨이브를 사용하여 연산 효율을 높입니다. 또한 tile_size_k를 가변적으로 설정하여 메모리 배리어(barrier)를 줄이는 최적화를 수행했습니다.

2. subgroup_matrix_matmul_nbits.h: 프로그램 캐시 키 확장

타일 크기가 동적으로 변함에 따라, 이를 프로그램 캐시 키에 포함시켜 올바른 셰이더가 생성되도록 수정되었습니다.

// After
SubgroupMatrixMatMulNBitsProgram(..., uint32_t tile_m, uint32_t tile_n, uint32_t tile_k, ...)
    : Program{config.subgroupSize == 64 ? "SubgroupMatrixMatMulNBitsWave64" : "SubgroupMatrixMatMulNBits"},
      tile_m_(tile_m), tile_n_(tile_n), tile_k_(tile_k) {}

왜 이게 좋은가

이번 최적화의 핵심은 '워크로드 인지형 타일링(Workload-aware Tiling)'입니다.

1. 성능 향상: 512-token prefill 기준 약 11%의 성능 향상을 보였으며, 전반적인 prefill 작업에서 약 9%의 이득을 얻었습니다.
2. 리소스 효율성: N=64인 경우 단일 웨이브를 사용하도록 최적화하여 워크그룹 스토리지 사용량을 19,456바이트에서 9,728바이트로 절반으로 줄였습니다.
3. 일반적 교훈: GPU 커널 최적화 시 고정된 타일 사이즈를 고집하기보다, 하드웨어의 특성(Wave64 vs Wave32)과 입력 데이터의 형상(M, N, K)에 맞춰 런타임에 최적의 파라미터를 선택하는 것이 성능의 핵심임을 보여줍니다.

리뷰 과정에서 qjia7님은 N=64일 때 불필요한 패딩을 제거하는 로직을 제안했고, 이는 실제 메모리 대역폭 절감으로 이어졌습니다. 또한, 다양한 에지 케이스(K가 타일 사이즈로 나누어떨어지지 않는 경우 등)에 대한 테스트 케이스 보강을 통해 안정성을 확보했습니다.

결론

이번 PR은 WebGPU 환경에서 하드웨어 가속을 최대한 끌어내기 위한 세밀한 튜닝의 좋은 예시입니다. 특히 AMD 하드웨어의 특성을 고려한 wave64 전략은 향후 다른 연산 최적화에도 중요한 벤치마크가 될 것입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://www.w3.org/TR/WGSL/

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