[flashinfer] FlashInfer, Muse-Spark-Open을 위한 융합 정규화 및 TP/EP MoE 최적화로 성능 극대화

2026년 10월 16일수정: 2026년 10월 16일

PR 링크: flashinfer-ai/flashinfer#6544 상태: Merged | 변경: +9123 / -93

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기와 복잡성을 증가시키면서, 동시에 추론 성능에 대한 요구를 더욱 높이고 있습니다. 특히 Mixture-of-Experts (MoE) 아키텍처는 수십억 개의 파라미터를 효율적으로 활용할 수 있게 해주지만, 그만큼 복잡한 연산과 통신 오버헤드를 동반합니다. NVIDIA의 Muse-Spark-Open 모델과 같은 최신 아키텍처는 이러한 MoE 모델의 성능을 더욱 끌어올리기 위해 다양한 최적화 기법을 요구합니다.

이번 PR은 FlashInfer 라이브러리에 Muse-Spark-Open 모델을 위한 새로운 커널들을 도입하여, 특히 정규화(Normalization)와 MoE 관련 연산을 융합(fusion)하고 최적화하는 데 중점을 둡니다. 이를 통해 중간 텐서 생성을 줄이고, 커널 실행 횟수를 최소화하며, 데이터 이동을 최적화하여 전반적인 추론 성능을 크게 향상시키는 것을 목표로 합니다.

이 글에서는 해당 PR에서 어떤 코드 변경이 이루어졌고, 왜 이러한 변경이 성능 향상으로 이어지는지, 그리고 이 최적화가 가지는 일반적인 교훈은 무엇인지 심층적으로 분석해보겠습니다.

코드 분석

이번 PR은 주로 flashinfer/norm 및 flashinfer/comm 모듈에 새로운 커널들을 추가하고 기존 커널들을 개선하는 데 집중합니다. 특히, 다양한 정규화 기법과 MoE 관련 통신 및 연산을 융합하는 데 초점을 맞추고 있습니다.

1. 융합 정규화 커널 (`flashinfer/norm`)

MoE 모델에서는 각 전문가(expert)의 출력을 처리하기 위해 다양한 정규화 기법이 사용됩니다. 이 PR은 이러한 정규화 연산들을 다른 연산들과 융합하여 GPU의 연산 효율성을 높입니다.

1.1. Gated Residual RMSNorm Fusion

기존에는 RMSNorm, 게이팅(gating), 잔차(residual) 결합이 별도의 연산으로 수행되었을 수 있습니다. 이 PR은 이를 하나의 커널로 융합하여 중간 텐서 생성을 줄이고 커널 실행 오버헤드를 감소시킵니다. 특히 FP32 잔차 스트림을 유지하면서 융합을 수행하는 것이 특징입니다.

--- a/docs/api/norm.rst
+++ b/docs/api/norm.rst
@@ -25,6 +25,23 @@ Kernels for normalization layers.
     fused_dit_gate_residual_layernorm_scale_shift
     fused_dit_gate_residual_layernorm_gamma_beta
 
+Fixed-Width Fused Norms
+-----------------------
+
+Fused gated-residual and per-head gated RMSNorm kernels for SM100, SM103 and
+SM107. They support only the widths in their docstrings (hidden size 8192, head
+dim 64); check the matching ``supports_*`` helper and fall back otherwise.
+
+.. autosummary::
+    :toctree: ../generated
+
+    gated_residual_rmsnorm
+    supports_gated_residual_rmsnorm
+    gated_residual_next_rmsnorm
+    supports_gated_residual_next_rmsnorm
+    rmsnorm_sigmoid_gate
+    supports_rmsnorm_sigmoid_gate
+
 Training RMSNorm (Cake)
 -----------------------

새로운 gated_residual_rmsnorm, gated_residual_next_rmsnorm, rmsnorm_sigmoid_gate 함수들이 추가되었으며, 이는 특정 하드웨어 아키텍처(SM100, SM103, SM107)에 최적화되어 있습니다. 이 함수들은 고정된 너비(hidden size 8192, head dim 64)를 지원하며, supports_* 함수를 통해 해당 커널 사용 가능 여부를 확인할 수 있습니다.

1.2. RMSNorm + Sigmoid Gate Fusion

이 최적화는 각 헤드(head)별 RMSNorm과 시그모이드 게이팅을 융합합니다. 런타임 토큰 커널 재사용, 풀-타일 전문화(full-tile specialization), PDL(Programmable Data Layout) 지원을 통해 효율성을 높입니다.

2. 융합 통신 및 MoE 커널 (`flashinfer/comm`)

MoE 모델의 핵심은 여러 전문가를 효율적으로 선택하고 그 결과를 종합하는 것입니다. 이는 종종 All-Gather와 같은 통신 연산과 결합되며, 이 PR은 이러한 연산들을 융합하여 성능을 개선합니다.

2.1. TP MoE Finalize Fusion

Tensor Parallelism (TP) 환경에서 MoE 모델의 최종 출력을 계산하는 과정은 복잡합니다. 이 PR은 TP reduction, 전문가별 RMSNorm, 그리고 가중치 합산(weighted accumulation)을 GEMM2 연산의 결과로부터 직접 융합하는 새로운 커널을 도입합니다. 이는 중간 결과를 저장하고 다시 로드하는 과정을 줄여줍니다. 특히, 토큰 소유자 경로(token-owner paths)를 사용하여 중간 및 큰 토큰 개수에 대해 최적화합니다.

--- a/docs/api/comm.rst
+++ b/docs/api/comm.rst
@@ -679,6 +679,47 @@ the MNNVL workspace accepted by ``decode_cp_a2a_alltoall``.
     decode_cp_a2a_lse_reduce_create_workspace
     decode_cp_a2a_lse_reduce
 
+TP MoE Finalize + RMSNorm
+-------------------------
+
+Tensor-parallel MoE finalize fused with the all-reduce and per-expert RMSNorm,
+over an NVLink multicast workspace (hidden size 4096, top-k 8). Initialize the
+process-wide workspace collectively with ``init_tp_moe_finalize_rmsnorm`` before
+the first call.
+
+.. currentmodule:: flashinfer.comm
+
+.. autosummary::
+    :toctree: ../generated
+
+    init_tp_moe_finalize_rmsnorm
+    supports_tp_moe_finalize_rmsnorm
+    tp_moe_finalize_rmsnorm
+
+MNNVL CuTe DSL Fused Collectives
+--------------------------------
+
+.. currentmodule:: flashinfer.comm.mnnvl_cutedsl_moe_finalize
+
+.. autosummary::
+    :toctree: ../generated
+
+    MNNVLCuteDSLMoEFinalizeAllReduceWorkspace
+    supports_moe_finalize_rmsnorm_allreduce
+    moe_finalize_rmsnorm_allreduce
+
+.. currentmodule:: flashinfer.comm.mnnvl_cutedsl_ag
+
+.. autosummary::
+    :toctree: ../generated
+
+    MNNVLCuteDSLAllGatherDualNormWorkspace
+    supports_all_gather_gated_residual_next_rmsnorm
+    all_gather_gated_residual_next_rmsnorm
+    MNNVLCuteDSLAllGatherNormQuantWorkspace
+    supports_all_gather_rmsnorm_fp4_quantize
+    all_gather_rmsnorm_fp4_quantize
+
 Mixed Communication
 -------------------

새로운 tp_moe_finalize_rmsnorm 함수와 관련 초기화 함수 init_tp_moe_finalize_rmsnorm이 추가되었습니다. 이는 TP 환경에서 MoE 최종 연산을 최적화하며, 특히 All-Reduce와 RMSNorm을 융합합니다. 또한, MNNVL CuTe DSL을 활용한 all_gather_gated_residual_next_rmsnorm 및 all_gather_rmsnorm_fp4_quantize와 같은 융합 커널들도 도입되었습니다.

2.2. All-gather + Gated Residual Dual Norm Fusion

이 최적화는 컬럼 All-Gather 연산을 게이트 잔차 이중 정규화(gated residual dual normalization)와 융합합니다. 이는 통신과 계산을 동시에 수행하여 효율성을 높입니다.

2.3. All-gather + RMSNorm + NVFP4 Quantization Fusion

All-Gather, RMSNorm, 그리고 NVFP4 양자화(quantization)를 하나의 커널로 융합합니다. 이는 데이터 이동과 정규화, 그리고 양자화 과정을 효율적으로 처리합니다.

2.4. MoE Finalize + RMSNorm + All-reduce Fusion

MoE 최종 연산과 RMSNorm, All-reduce 연산을 융합합니다. 비-로컬 전문가(nonlocal experts)의 경우 정규화를 건너뛰는 최적화도 포함됩니다.

3. 커널 지원 및 검증

새로 추가된 커널들은 SM100, SM103, SM107 아키텍처를 지원합니다. 또한, PR에는 이러한 새로운 기능들을 검증하기 위한 테스트 케이스 추가 및 기존 테스트의 수정이 포함되어 있습니다. 리뷰어들의 피드백을 통해 pre-commit 및 테스트 관련 이슈들이 해결되었음을 확인할 수 있습니다.

--- a/flashinfer/__init__.py
+++ b/flashinfer/__init__.py
@@ -264,6 +264,18 @@ def __dir__() -> list[str]:
 from .cake_rmsnorm_train import (
     cake_rmsnorm_train_forward as cake_rmsnorm_train_forward,
 )
+from .norm import gated_residual_rmsnorm as gated_residual_rmsnorm
+from .norm import gated_residual_next_rmsnorm as gated_residual_next_rmsnorm
+from .norm import rmsnorm_sigmoid_gate as rmsnorm_sigmoid_gate
+from .norm import (
+    supports_gated_residual_next_rmsnorm as supports_gated_residual_next_rmsnorm,
+)
+from .norm import (
+    supports_gated_residual_rmsnorm as supports_gated_residual_rmsnorm,
+)
+from .norm import (
+    supports_rmsnorm_sigmoid_gate as supports_rmsnorm_sigmoid_gate,
+)
 from . import nvfp4_attention_sm120 as nvfp4_attention_sm120
 from .nvfp4_attention_sm120 import (
     nvfp4_attention_sm120_fwd as nvfp4_attention_sm120_fwd,

__init__.py 파일에서 새로운 융합 정규화 함수들이 임포트되어 사용자가 쉽게 접근할 수 있도록 합니다.

왜 이게 좋은가?

이 PR에서 도입된 최적화들은 여러 측면에서 성능 향상에 기여합니다.

연산 융합 (Kernel Fusion): 여러 개의 작은 커널 호출을 하나의 큰 커널로 통합함으로써, GPU 커널 실행 오버헤드(launch overhead)를 크게 줄일 수 있습니다. 또한, 각 커널 간의 중간 결과 저장을 위한 글로벌 메모리 접근을 최소화하여 데이터 이동(data movement)을 줄입니다. 이는 GPU의 메모리 대역폭 병목 현상을 완화하고 연산 유닛의 활용률을 높이는 데 결정적인 역할을 합니다.
데이터 이동 감소: 융합된 커널은 중간 결과를 레지스터나 공유 메모리(shared memory) 내에서 처리하는 경우가 많습니다. 이는 느린 글로벌 메모리 접근 횟수를 줄여 성능을 크게 향상시킵니다. 특히, All-Gather와 같은 통신 연산과 계산 연산을 융합하는 것은 데이터가 GPU에 머무는 시간을 최적화하는 데 매우 효과적입니다.
특정 아키텍처 최적화: SM100, SM103, SM107과 같은 최신 GPU 아키텍처에 특화된 커널을 제공함으로써, 해당 하드웨어의 기능을 최대한 활용할 수 있습니다. 예를 들어, CuCu DSL과 같은 최신 프로그래밍 모델을 사용하여 하드웨어 기능을 세밀하게 제어하고 최적의 성능을 이끌어냅니다.
MoE 효율성 증대: MoE 모델은 파라미터는 많지만 실제 추론 시에는 일부 전문가만 활성화됩니다. 이 PR에서 도입된 TP MoE Finalize 융합 커널은 이러한 MoE 구조의 특성을 고려하여, 전문가 선택 및 결과 종합 과정을 더욱 효율적으로 만들어 MoE 모델의 추론 속도를 높입니다.

일반적인 교훈:

융합은 성능 향상의 핵심: GPU 컴퓨팅에서 연산 융합은 반복적으로 입증된 성능 향상 기법입니다. 특히 LLM과 같이 연산 집약적인 워크로드에서는 그 효과가 더욱 두드러집니다.
메모리 접근 패턴 최적화: 데이터 이동을 최소화하는 것은 GPU 성능 최적화의 핵심입니다. 통신과 계산을 융합하거나, 중간 결과를 레지스터/공유 메모리에 유지하는 전략은 항상 고려해야 합니다.
하드웨어 특성 활용: 타겟 하드웨어 아키텍처의 기능을 최대한 활용하는 커스텀 커널 개발은 성능을 극대화하는 강력한 방법입니다.
MoE 최적화의 중요성: MoE 모델의 복잡성을 관리하고 효율성을 높이기 위한 전용 최적화 기법은 LLM 분야에서 점점 더 중요해지고 있습니다.

리뷰 댓글에서 언급된 enable_pdl 옵션의 일관성 문제나, tp_moe_finalize_rmsnorm의 워크스페이스 관리 방식에 대한 논의는 코드의 견고성과 사용 편의성을 높이기 위한 중요한 피드백입니다. 이러한 피드백을 통해 라이브러리의 완성도가 더욱 높아질 것입니다.