[onnxruntime] ONNX Runtime의 FP4/FP8 GEMV 커널 최적화: Tensor Core와 M-Tiling을 통한 성능 극대화

2026년 8월 1일수정: 2026년 8월 1일

PR 링크: microsoft/onnxruntime#31155 상태: Merged | 변경: +1911 / -97

들어가며

최근 LLM 추론에서 Speculative Decoding이나 MTP(Multi-Token Prediction)와 같은 기법들이 도입되면서, 기존의 M=1 (단일 토큰 생성) 환경에 최적화된 GEMV 커널들이 M > 1 환경에서 성능 저하를 겪는 문제가 발생했습니다. 특히 FP8/FP4와 같은 저정밀도 연산에서 cuBLAS 대비 성능이 역전되는 현상이 관찰되었습니다. 본 PR은 MatMulBlockQuantizedFp4Weight 및 MatMulBlockQuantizedFp8Weight 연산의 Fast Path에 5가지 핵심 최적화를 적용하여 이러한 문제를 해결합니다.

코드 분석

1. FP8: `RowsPerWarp > 1` 시 FP32 Widening Hoisting

M > 1인 경우, 기존 커널은 각 행마다 가중치를 FP32로 확장했습니다. 이를 루프 밖으로 끌어올려(hoisting) 중복 연산을 제거했습니다.

// Before: 루프 내부에서 매번 확장
// After: RowsPerWarp > 1일 때 B를 한 번만 확장하여 재사용
// Instructions per 16 weight bytes: 8 + 48*R -> 8*C + 16*R + 16*C + 16*R*C

2. FP4: `prmt` Quad Decode

prmt.b32 명령어를 사용하여 32비트 워드에서 8개의 E2M1 코드를 한 번에 디코딩합니다.

// Before: 요소별 디코딩
// After: prmt.b32를 이용한 4개 코드 동시 디코딩
// SASS instruction count: 352 -> 288 (-18%)

3. FP4: Grid-Gated Row Tiling

RowsPerBlock 행을 한 번에 처리하여 가중치 로드와 디코딩 비용을 상쇄합니다.

// Grid-gated: N이 작을 때는 tiling이 오히려 손해이므로 조건부 적용
if (col_blocks >= sm_count) { /* Tiling 적용 */ }

4. Tensor Core 활용 (`mma.m16n8k16`)

가장 큰 성능 향상을 가져온 부분으로, 기존 FMA 기반 연산을 mma.m16n8k16 명령어로 대체했습니다.

// 가중치를 A 슬롯, 활성화 함수를 B 슬롯에 배치하여 Transpose 없이 연산
// mma.m16n8k16을 통해 16개 출력 컬럼을 한 번에 처리

왜 이게 좋은가

이번 최적화는 단순히 연산 효율을 높이는 것을 넘어, 메모리 대역폭 활용도를 극대화했습니다. 특히 FP4 Tensor Core 최적화는 lm_head 연산에서 기존 대비 4.96배의 속도 향상을 기록했습니다.

핵심 교훈:

데이터 재사용성: M > 1 환경에서는 가중치 로드 비용을 여러 행에 걸쳐 상쇄(amortize)하는 전략이 필수적입니다.
ISA 활용: prmt와 같은 하드웨어 특화 명령어를 활용하면 디코딩 오버헤드를 획기적으로 줄일 수 있습니다.
Tensor Core의 유연성: 가중치와 활성화 함수의 레이아웃을 mma 명령어의 요구사항에 맞게 정렬하면, 별도의 전치(transpose) 비용 없이 고성능 연산이 가능합니다.

리뷰 피드백 반영

리뷰 과정에서 ORT_FP4_GEMV_MMA와 같은 환경 변수가 int로 파싱되어 false 값을 무시하는 버그가 발견되어 bool 타입으로 수정되었습니다. 또한, FP8 Tensor Core 테스트의 수치적 정확도 검증을 위해 허용 오차(tolerance)를 엄격하게 조정하여 테스트의 신뢰성을 확보했습니다.