[onnxruntime] [CUDA] NVFP4 QMoE GEMV 최적화: ALU 바운드 커널의 한계를 넘어서는 방법

PR 링크: microsoft/onnxruntime#31159 상태: Merged | 변경: +728 / -66

들어가며\n\n최근 LLM 추론 최적화에서 NVFP4(E2M1)와 같은 저정밀도 데이터 타입은 메모리 대역폭을 절약하는 핵심 요소로 자리 잡았습니다. 하지만 데이터가 압축될수록 이를 다시 부동소수점으로 변환하는 'Dequantization(역양자화)' 과정에서 발생하는 연산 오버헤드가 새로운 병목 지점이 되곤 합니다.\n\nmicrosoft/onnxruntime 레포지토리에 제출된 이번 PR은 Qwen3.6 모델의 NVFP4 QMoE(Quantized Mixture of Experts) 디코딩 성능을 대폭 개선했습니다. 특히 이 커널이 메모리 대역폭(DRAM)이 아닌 ALU 파이프라인 바운드(ALU-pipeline bound) 상태라는 점에 착안하여, 명령어 수를 줄이는 데 집중한 최적화 사례입니다. ncu 프로파일링 결과 ALU 점유율 78.9% 대비 DRAM 점유율은 7.3%에 불과했던 이 커널을 어떻게 개선했는지 살펴보겠습니다.\n\n## 핵심 변경 사항 분석\n\n### 1. Packed E2M1 Dequantize (prmt 활용)\n\n기존에는 NVFP4 코드를 하나씩 읽어 개별적으로 디코딩했습니다. 이는 분기(branch)가 없더라도 마스킹, 시프트, 오어 연산 등 8개의 코드를 처리하는 데 약 53개의 ALU 명령어가 소모되는 비효율적인 구조였습니다.\n\nBefore (Scalar Decode):\ncpp\n// 기존의 개별 디코딩 방식 (의사 코드)\nfor (int i = 0; i < 8; ++i) {\n uint8_t code = (word >> (i * 4)) & 0xF;\n output[i] = decode_single_fp4(code);\n}\n\n\nAfter (Packed Quad Decode):\ncpp\n// prmt.b32 명령어를 활용한 4개 동시 디코딩\n// details.h 내 Fp4I2FConverter::decode_quad\nstatic __device__ __forceinline__ void decode_quad(uint32_t w, T* dst) {\n uint32_t mag = w & 0x77777777; // 하위 3비트(Magnitude) 유지\n uint32_t sgn = (w >> 3) & 0x11111111; // 부호 비트 추출\n // prmt를 사용하여 4개의 바이트를 한 번에 선택 및 변환\n uint32_t res_low = __byte_perm(mag_table_low, sgn_low, selector);\n // ... 중략 ...\n}\n\n\n새로운 decode_quad 방식은 prmt.b32 명령어를 사용하여 32비트 워드(8개 코드)를 한 번에 처리합니다. prmt는 바이트 단위 선택 기능을 제공하므로, 4개의 Magnitude 룩업을 단 하나의 명령어로 수행할 수 있습니다. 이를 통해 SASS(Assembly) 코드량이 약 30% 감소했으며, 실제 FC1 커널 속도가 33.2µs에서 26.2µs로 개선되었습니다.\n\n### 2. 메모리 및 ALU 트래픽 최적화\n\nColumnMajor 레이아웃에서 Scale 값들이 연속적으로 배치되어 있음에도 불구하고 기존에는 스칼라 형태로 접근했습니다. 이를 벡터화된 접근으로 변경하여 불필요한 메모리 섹터 접근을 줄였습니다.\n\nBefore:\ncpp\n// 각 컬럼마다 개별적으로 Scale 로드\nfor (int i = 0; i < CtaN; ++i) {\n scales[i] = scale_ptr[col + i];\n}\n\n\nAfter:\ncpp\n// ScalesAccess를 통한 벡터화된 로드\nusing ScalesAccess = ScalesAccessT<TypeA, Interleave, CtaN>;\nif constexpr (ScalesAccess::kVectorized) {\n *reinterpret_cast<typename ScalesAccess::TVec*>(vec_scale) = \n *reinterpret_cast<const typename ScalesAccess::TVec*>(scale_ptr);\n}\n\n\n이 변경으로 인해 moe_gemv_kernel의 실행 시간이 0.700ms에서 0.494ms로 약 30% 단축되었습니다. DRAM 사용량이 적더라도 불필요한 로드 명령어를 줄이는 것이 ALU 바운드 커널에서 얼마나 중요한지 보여줍니다.\n\n### 3. MTP(Multi-Token Prediction)를 위한 확장성 확보\n\nMTP 검증 단계에서는 한 번에 여러 토큰을 처리하므로 expanded_num_rows가 증가합니다. 기존에는 이 값이 8로 제한되어 있어, MTP 환경에서는 최적화된 GEMV 경로를 타지 못하고 느린 Fallback 경로(Dequantize + CUTLASS)를 사용해야 했습니다.\n\nBefore/After:\ndiff\n- static constexpr int kMaxProfiledExpandedRows = 8;\n+ static constexpr int kMaxProfiledExpandedRows = 64;\n\n\n이 단순한 상숫값 변경으로 MTP 검증 성능이 2.4 tok/s에서 30~55 tok/s로 최대 23배 향상되었습니다.\n\n### 4. 분석적 기본 타이링 (Analytic Default Tiling)\n\n기존에는 최적의 스레드 배치를 위해 Autotuner에 의존했으나, 이는 추론 스트림을 동기화(Sync)시키는 오버헤드가 있었습니다. 이번 PR에서는 셰이프(Shape)를 분석하여 최적의 스레드 수를 결정하는 휴리스틱을 도입했습니다.\n\ncpp\n// moe_gemv_fp4.cu 내 Fp4MoeGemvDefaultConfig\nif (k < 1024) { \n config.threads = 64; // Idle 스레드 방지\n} else if (grid_size >= 16 * multiProcessorCount) {\n config.threads = 64; // 에필로그 비용 절감\n} else {\n config.threads = 128;\n}\n\n\n## 왜 이게 좋은가\n\n1. 명령어 수준 병렬성(ILP) 극대화: ALU 바운드 커널에서 prmt와 같은 특수 명령어를 활용해 명령어 수를 줄인 것은 교과서적인 최적화입니다.\n2. MTP 실용성 확보: 이론적인 성능 개선에 그치지 않고, 실제 MTP 검증 시나리오에서 발생하는 병목을 정확히 짚어내어 수십 배의 성능 향상을 이끌어냈습니다.\n3. 동기화 오버헤드 제거: Autotuner 없이도 최적에 가까운 Tiling을 선택함으로써 CUDA Graph 캡처 환경에서도 안정적인 성능을 보장합니다.\n\n## 리뷰어 피드백 분석: 수치적 정밀도와 정렬 문제\n\n리뷰 과정에서 중요한 기술적 논의가 있었습니다.\n\n* 부동소수점 비결합성(Non-associativity): 스레드 수를 128에서 64로 변경하거나 누적 방식을 바꾸면 부동소수점 연산 순서가 달라집니다. 리뷰어 titaiwangms는 이로 인해 결과값의 하위 비트가 달라질 수 있음을 지적했습니다. 이는 비트 단위 동일성(Bit-exactness)이 중요한 환경에서 주의해야 할 점입니다.\n* 메모리 정렬(Alignment): tile_w나 vec_scale을 uint32_t나 float4로 캐스팅하여 접근할 때, 적절한 정렬(alignas)이 없으면 미정의 동작(Undefined Behavior)이 발생할 수 있습니다. tianleiwu는 이를 해결하기 위해 alignas(alignof(uint32_t)) 등을 명시적으로 추가했습니다.\n\n## 결론\n\n이번 PR은 하드웨어의 특성(ALU Bound)을 정확히 이해하고, PTX 수준의 명령어 최적화와 아키텍처 기반의 휴리스틱을 결합하여 실질적인 성능 이득을 얻은 훌륭한 사례입니다. 특히 prmt를 이용한 비트 트릭은 고성능 커널 개발자라면 반드시 참고할 만한 기법입니다.

참고 자료

• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-prmt
• https://en.cppreference.com/w/cpp/language/alignas
• https://onnxruntime.ai/docs/performance/model-optimizations/moe.html

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