[onnxruntime] ONNX Runtime CUDA QMoE 및 FP8 최적화: 성능 향상 및 메모리 효율성 증대

PR 링크: microsoft/onnxruntime#33094 상태: Merged | 변경: +1367 / -109

들어가며

최근 대규모 언어 모델(LLM) 분야에서는 Mixture-of-Experts (MoE) 아키텍처가 뛰어난 성능과 효율성으로 주목받고 있습니다. 특히 ONNX Runtime은 다양한 하드웨어 가속기 지원을 통해 이러한 모델의 배포를 용이하게 합니다. 이번 PR은 ONNX Runtime의 CUDA Execution Provider에서 MoE 모델, 특히 양자화된 MoE(QMoE) 모델의 성능과 메모리 사용량을 개선하기 위한 중요한 최적화를 포함하고 있습니다. 주요 내용은 오프라인 NVFP4 가중치 레이아웃 지원과 기본 활성화된 융합 블록 스케일 FP8 실행입니다. 이 PR은 대규모 MoE 모델에서 발생하는 메모리 및 초기화 비용을 절감하고, 추론 시 계산 효율성을 극대화하는 것을 목표로 합니다.

코드 분석

이번 PR은 주로 docs/ 디렉토리의 문서 파일과 내부 CUDA 커널 구현에 영향을 미치는 변경사항을 포함합니다. 핵심 변경 사항은 다음과 같습니다.

1. 문서 업데이트 (docs/ContribOperators.md, docs/contrib_ops/cuda/moe_qmoe.md)

가장 눈에 띄는 변경은 QMoE 연산자에 대한 문서 업데이트입니다. 새로운 기능과 옵션이 추가됨에 따라 관련 설명이 상세하게 보강되었습니다.

**1.1. 블록 스케일 FP8 실행 설명 추가 (docs/ContribOperators.md)

--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -6121,6 +6121,10 @@
         Partial edge blocks are allowed. No zero points or activation scales are accepted.
         Without a positive block_size, FP8 uses the legacy per-expert fc*_global_scale inputs instead.
         Block-scaled FP8 does not use global scales. Activations retain the input type (weight-only quantization).
+        On CUDA SM80-or-newer GPUs, SiLU/SwiGLU block-scaled FP8 uses fused weight-only GEMV/GEMM by default.
+        These kernels decode and scale FP8 tiles on chip without materializing full FP16/BF16 weight buffers.
+        Set ORT_ENABLE_FP8_FUSED=0 before session creation to select the compact dense-dequant fallback.
+        Other activation types and unsupported routing sizes retain the fallback.
         The WebGPU block-FP8 kernel supports block_size=128 with float32 scales and rejects projections
         that exceed 32-bit shader addressing or the device's per-dimension dispatch limit.
 

설명: SM80 이상 GPU에서 SiLU/SwiGLU 블록 스케일 FP8 워크로드가 기본적으로 융합된 가중치 전용 GEMV/GEMM을 사용하도록 변경되었습니다. 이는 FP8 타일을 온칩에서 디코딩하고 스케일링하여 전체 FP16/BF16 가중치 버퍼를 생성하지 않으므로 메모리 효율성을 크게 향상시킵니다. ORT_ENABLE_FP8_FUSED=0 환경 변수를 통해 이전의 콤팩트한 dense-dequant fallback으로 전환할 수 있습니다.

**1.2. weights_prepacked 속성 설명 확장 (docs/ContribOperators.md, docs/contrib_ops/cuda/moe_qmoe.md)

--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -6182,7 +6186,7 @@
 <dt><tt>use_sparse_mixer</tt> : int</dt>
 <dd>Whether to use sparse mixer</dd>
 <dt><tt>weights_prepacked</tt> : int</dt>
-<dd>Only meaningful when quant_type='int'. Tri-state control over the layout of the int4/int8 fc1/fc2 weight initializers. The concrete prepacked layouts selected by -1 and 1 are determined by the execution provider. 0: the initializers are raw, un-prepacked [E, N, K/pack] tensors as produced by quantize_matmul_{4,8}bits. Defaults to -1.</dd>
+<dd>Expert-weight layout selector with quantization-specific meanings. For quant_type='int', -1 and 1 select execution-provider-specific prepacked int4/int8 fc1/fc2 layouts. 0 selects raw [E,N,K/pack] tensors as produced by quantize_matmul_{4,8}bits. For quant_type='nvfp4' on CUDA, -1 (default) and 0 select legacy N-packed [E,K,N/2] bytes; 1 selects K-packed row-major [E,N,K/2] bytes while retaining logical [E,K,N/2] dimensions. NVFP4 block scales and global scales retain their original layouts and values. Ignored for other quantization types. Supported values are -1, 0, and 1; defaults to -1.</dd>
 </dl>
 
 #### Inputs (6 - 21)

--- a/docs/contrib_ops/cuda/moe_qmoe.md
+++ b/docs/contrib_ops/cuda/moe_qmoe.md
@@ -82,7 +82,7 @@
 | `expert_weight_bits` (QMoE only) | int | 4 | 4 (INT4/MXFP4) or 8 (INT8/FP8). |
 | `block_size` (QMoE only) | int | -1 | Group size for INT4/INT8 group-wise quantization. -1 = per-output-channel. FP4/WFP4AFP8 normalize an omitted value to 32 and require 32 when present; NVFP4 normalizes an omitted value to 16 and requires 16 when present. FP8 uses square blocks for positive values, per-expert globals otherwise. | 
 | `quant_type` (QMoE only) | string | "int" | "int", "fp4", "nvfp4", "fp8", "wfp4afp8". See [§3](#3-quantization-modes). |
- | `weights_prepacked` (QMoE only) | int | -1 | Tri-state, only meaningful when `quant_type="int"`. The prepacked layouts selected by `-1` and `1` are **EP-determined**. `-1` (default): the INT4/INT8 `fc1`/`fc2` initializers are already prepacked in the EP's default layout (e.g. from `pack_weights_for_cuda_mixed_gemm` for the CUDA EP). `1`: already prepacked in an alternate EP-selected layout. `0`: the initializers are raw `[E, N, K/pack]` tensors (as produced by `quantize_matmul_{4,8}bits`) and the kernel runs the CUTLASS layout transform in `PrePack()`. **Note:** the CUDA EP INT4/INT8 MoE GEMM always runs the Ampere (SM80) kernel — even on SM90 — so it consumes the SM80 `fpA_intB` layout on all architectures; `-1` and `1` are therefore equivalent for the CUDA EP today, and `1` is reserved for a possible future Hopper-specific layout. See [§5.1](#51-weights-input-2--5--8). |
+| `weights_prepacked` (QMoE only) | int | -1 | Layout selector for integer quantization and CUDA NVFP4. Integer semantics are unchanged: `-1` and `1` select EP-determined prepacked layouts (both SM80 `fpA_intB` on CUDA today); `0` selects raw `[E, N, K/pack]` weights transformed in `PrePack()`. For CUDA NVFP4, omitted/`-1` and `0` select legacy N-packed bytes; `1` selects offline K-packed row-major bytes. Mode `2` is unsupported for NVFP4. See [§5.1](#51-weights-input-2--5--8). |
+
+**CUDA NVFP4** uses a separate `weights_prepacked` contract:
+
+| `weights_prepacked` | Physical packed-byte storage |
+|---------------------|------------------------------|
+| Omitted / `-1` | Legacy N-packed `[E, K, N/2]` |
+| `0` | Explicit legacy N-packed `[E, K, N/2]` |
+| `1` | Offline K-packed row-major `[E, N, K/2]` |
+
+Logical ONNX dimensions remain `[E, K, N/2]` in every mode. Mode `1` changes
+only the physical weight-byte layout: E2M1 codes, raw E4M3 block scales
+`[E, N, K/16]`, and FP32 global scales `[E]` retain their values. Mode `2` is
unsupported. Integer-quantization semantics are unchanged.
+
+Mode `1` shares the offline payload across eligible native GEMV decode,
+Blackwell native GEMM prefill, and dense dequantization fallback without a
+decode-only weight repack or expanded GEMV scale bank. Blackwell native GEMM
+still derives its TMA-swizzled block scales; the dense fallback reads the
+retained raw scale inputs, not that swizzled buffer. Disabling prepacking with
+`session.disable_prepacking` preserves correct offline-layout dequantization
+and eligible GEMV execution; native prefill requires prepacking.
+
 MXFP4 weights must be packed by `pack_fp4_weights_for_cuda_moe_gemm`. FP8 weights
 are stored as raw e4m3 bytes (no packing).
 

설명: weights_prepacked 속성의 의미가 quant_type='int' 뿐만 아니라 quant_type='nvfp4' (CUDA)에 대해서도 확장되었습니다. 특히 weights_prepacked=1은 오프라인 K-packed row-major 저장 방식을 사용하도록 지정합니다. 이 방식은 기존의 N-packed 방식과 달리 물리적인 바이트 레이아웃을 변경하여 메모리 접근 패턴을 최적화하고, GEMV, GEMM, 그리고 dense fallback 간에 동일한 사전 패킹된 페이로드를 공유할 수 있게 합니다. 이는 런타임 시 가중치 재패킹 및 스케일 버퍼로 인한 추가 메모리 및 초기화 비용을 줄여줍니다.

**1.3. NVFP4 관련 상세 설명 추가 (docs/contrib_ops/cuda/moe_qmoe.md)

--- a/docs/contrib_ops/cuda/moe_qmoe.md
+++ b/docs/contrib_ops/cuda/moe_qmoe.md
@@ -1060,18 +1081,20 @@
 - `is_moe_gemv_fp4_supported` accepts `group_size ∈ {16, 32}`; the dispatch instantiates the
   `GroupSize=16` cases in `dispatch_moe_gemv_group_size` /
   `dispatch_moe_gemv_interleaved_swiglu_group_size` ([moe_gemv_device.cuh](onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh)).
-- Prepacked NVFP4 uses the non-interleaved `ColumnMajor` layout; the opt-in interleaved path
+- Legacy prepacked NVFP4 uses the non-interleaved `ColumnMajor` layout; the opt-in interleaved path
   ([§9.10](#910-interleaved-gemv-layout--dtype-conditional-accumulation)) is MXFP4-only because its
   `kStepK=32` tile is tied to the block-32 scale layout.
-- `QMoECombineNvfp4ScalesForGemv` ([qmoe_kernels.cu](onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu))
+- For legacy prepacked NVFP4, `QMoECombineNvfp4ScalesForGemv` ([qmoe_kernels.cu](onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu))
   decodes the `float8e4m3fn` block scales, folds in the per-expert FP32 global scale, and rewrites
   `[E, n, k/16] → [E, k/16, n]` in the activation dtype (`TypeA`) that the GEMV expects.
  - The decode gate ([moe_quantization.cc](onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc)) fires
-  when `expanded = num_tokens·top_k ∈ (0, 64]`, SwiGLU is fused, and both FC1
+  when `expanded = num_tokens·top_k ∈ (0, 64]` for legacy N-packed storage or
+  `(0, 80]` for offline K-packed storage, SwiGLU is fused, and both FC1
   (`n=2·inter`, `k=hidden`) and FC2 (`n=hidden`, `k=inter`) satisfy `n,k ≥ 512` and group-16 block
   alignment.
    The fused expert-map prologue supports `top_k ∈ {1, 2, 4, 6, 8, 10}` and up to 1022 experts.
-  Qwen Flash (`hidden=2560`, `inter=640`, `top_k=10`) qualifies for one through six tokens.
+  Qwen Flash (`hidden=2560`, `inter=640`, `top_k=10`) qualifies for one through six tokens
+  with legacy storage, or one through eight tokens with offline K-packed storage.
 
 #### Raw-layout memory option
 

설명: 오프라인 K-packed 저장 방식(weights_prepacked=1)을 사용할 때, GEMV 디스패치 조건이 expanded 행의 최대 개수가 64개에서 80개로 증가했습니다. 이는 멀티 토큰 처리 시 더 많은 전문가를 효율적으로 활용할 수 있게 하여, 이전에는 fallback으로 처리되었을 경우를 최적화된 경로로 유도합니다. Qwen Flash 모델의 경우, 8개 토큰까지도 최적화된 경로를 탈 수 있게 됩니다.

2. 내부 CUDA 커널 및 로직 변경 (리뷰 코멘트 기반 추정)

PR 설명과 리뷰 코멘트를 종합해 볼 때, 다음과 같은 내부적인 CUDA 커널 및 로직 변경이 이루어졌을 것으로 추정됩니다.

• NVFP4 GEMM 최적화: MoeGemvFp4RawKPackedKernel과 같은 새로운 커널이 도입되었을 가능성이 높습니다. 이는 K-packed 레이아웃을 활용하여 contiguous vector loads, warp reductions, 그리고 short reductions 최적화를 수행합니다. 리뷰 코멘트에서는 "one contiguous uint2 load per 16-value K scale block, eight K-reduction lanes per output column, warp-shuffle reductions, and no shared-memory weight transpose or block barriers" 와 같은 구체적인 최적화 기법이 언급되었습니다.
• FP8 융합 실행: SM80+ GPU에서 FP8 워크로드를 위한 융합 GEMV/GEMM 커널이 구현되었습니다. 이 커널은 온칩에서 FP8 타일을 디코딩하고 스케일링하며, FP16/BF16 가중치 버퍼를 생성하지 않습니다. 이는 kunal-vaishnavi의 리뷰 코멘트에서 "Implemented fused weight-only block-scaled FP8 QMoE GEMV and tensor-core GEMM"으로 명시되어 있습니다.
• FC1 전문가 재사용 (실험적): ORT_NVFP4_FC1_EXPERT_REUSE=1 환경 변수를 통해 FC1 가중치 및 스케일 재사용 기능이 실험적으로 추가되었습니다. 이는 특정 조건(예: 융합 SwiGLU, 충분한 전문가 행 수) 하에서 두 개의 활성화 행에 대해 가중치를 한 번만 디코딩하여 계산량을 줄입니다. 이 기능은 기본적으로 비활성화되어 있으며, 모델의 전문가 중복성 및 타겟 GPU에서의 성능 측정을 통해 신중하게 활성화해야 합니다.

왜 이게 좋은가?

이번 PR의 최적화는 여러 측면에서 큰 이점을 제공합니다.

1. 성능 향상:

   • FP8 융합 실행: 리뷰 코멘트에서 A100에서의 합성 테스트 결과, 융합 실행은 fallback 대비 2.92배 (1 토큰) 및 1.63배 (17 토큰)의 성능 향상을 보였습니다. 이는 온칩 디코딩 및 스케일링, 그리고 FP16/BF16 가중치 버퍼를 사용하지 않음으로써 달성되었습니다.
   • NVFP4 K-packed GEMM: 새로운 K-packed GEMM 커널은 기존 N-packed 호환성 경로 대비 약 3.3배 ~ 4.8배 더 빠른 속도를 보였습니다. 또한, 특정 조건에서는 6.6배의 결합 커널 속도 향상도 관찰되었습니다. 이는 메모리 접근 패턴 최적화와 효율적인 스레드 활용 덕분입니다.
   • 확장된 GEMV 자격: 80개의 확장된 전문가 행까지 지원하게 되면서, 더 많은 멀티 토큰 시나리오가 fallback 대신 최적화된 GEMV 경로를 타게 되어 전반적인 처리량이 향상됩니다. 7개 토큰 시나리오에서 14.20배, 8개 토큰 시나리오에서 14.60배의 속도 향상이 보고되었습니다.

2. 메모리 효율성 증대:

   • 오프라인 NVFP4 가중치: K-packed 레이아웃은 가중치 페이로드를 공유하여 런타임 재패킹 및 추가 스케일 버퍼 할당을 줄입니다. 이는 특히 대규모 MoE 모델에서 상당한 메모리 절약을 가져옵니다.
   • FP8 융합 실행: FP16/BF16 가중치 버퍼를 생성하지 않음으로써 메모리 사용량을 더욱 줄입니다.

3. 일반적인 교훈:

   • 데이터 레이아웃의 중요성: 가중치 데이터의 물리적 저장 방식(N-packed vs K-packed)이 메모리 접근 패턴과 계산 효율성에 미치는 영향을 보여줍니다. 올바른 레이아웃 선택은 성능에 결정적인 영향을 미칩니다.
   • 하드웨어 특성 활용: SM80 이상 GPU의 기능을 활용하여 FP8 계산을 온칩에서 처리하는 것은 최신 하드웨어의 잠재력을 최대한 끌어내는 좋은 예입니다.
   • 점진적 최적화 및 실험: FC1 전문가 재사용과 같이 실험적인 기능은 기본적으로 비활성화하고, 실제 모델과 타겟 하드웨어에서의 충분한 검증 후에 사용하도록 하는 접근 방식은 안정성과 성능 향상의 균형을 맞추는 데 중요합니다.
   • 문서화의 중요성: 새로운 기능과 옵션(예: weights_prepacked의 NVFP4 의미, FP8 융합 실행 제어)을 명확하게 문서화하는 것은 사용자가 최적의 성능을 끌어낼 수 있도록 돕는 필수적인 과정입니다.

결론

이번 ONNX Runtime CUDA PR은 NVFP4 가중치 레이아웃 최적화와 블록 스케일 FP8 실행의 기본 활성화를 통해 대규모 QMoE 모델의 성능과 메모리 효율성을 크게 향상시켰습니다. 이러한 개선은 LLM 모델의 배포 및 운영 비용 절감에 기여할 것이며, ONNX Runtime이 최신 AI 하드웨어 가속 기능을 효과적으로 활용하는 능력을 보여주는 좋은 사례입니다.

참고 자료

• https://onnxruntime.ai/docs/reference/contrib_ops/cuda/moe_qmoe.html
• https://onnxruntime.ai/docs/reference/contrib_operators.html

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