[sglang] SGLang MoE Shared Expert 최적화: 3개 커널을 1개로 융합하여 GPU 오버헤드 제거

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

PR 링크: sgl-project/sglang#28658 상태: Merged | 변경: +50 / -11

들어가며

대규모 언어 모델(LLM)에서 Mixture-of-Experts (MoE) 아키텍처는 모델의 파라미터 수를 크게 늘리면서도 추론 비용을 효율적으로 유지할 수 있게 해주는 강력한 방법입니다. SGLang 프로젝트는 이러한 MoE 모델의 효율적인 서빙을 목표로 하며, 지속적인 성능 최적화를 진행하고 있습니다. 이번에 살펴볼 GitHub PR은 특히 AMD AITER 백엔드에서 Qwen2MoeSparseMoeBlock의 shared-expert fusion 경로에서 발생하는 비효율성을 해결하는 중요한 최적화입니다.

기존에는 MoE 레이어당 디코딩 단계에서 shared expert 라우팅 가중치를 준비하기 위해 세 가지 별개의 GPU 커널이 순차적으로 실행되었습니다:

sigmoid_kernel_cuda: _get_shared_expert_weights에서 w = F.sigmoid(shared_logits)를 계산하는 데 약 5.5 us 소요.
bfloat16tofloat32_copy_kernel_cuda: shared_weights.to(topk_weights.dtype)를 통해 bf16을 fp32로 캐스팅하는 데 약 4.5 us 소요.
_fused_append_shared_experts_with_weights_kernel: 실제 expert를 append하는 데 약 4.2 us 소요.

이 세 커널은 모두 작은 텐서에 대한 element-wise 연산으로, GPU 커널 런치 오버헤드와 메모리 대역폭에 의해 성능이 제한되는 경향이 있었습니다. 특히 sigmoid와 bf16->fp32 캐스팅은 append 커널에 데이터를 공급하기 위한 '접착제(glue)' 역할에 불과했습니다. 이 PR의 목표는 이 두 가지 접착제 연산을 append 커널 내부로 융합하여, 3개의 커널을 1개로 줄이고 두 번의 전역 메모리 왕복(global round-trip)을 제거하는 것입니다.

코드 분석

이 최적화는 주로 두 파일, fused_moe_triton_kernels.py와 qwen2_moe.py에서 이루어졌습니다.

`python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py`

이 파일은 Triton 커널을 정의하고 호출하는 로직을 포함합니다. 핵심 변경사항은 _fused_append_shared_experts_with_weights_kernel 커널 자체와 이를 호출하는 래퍼 함수 fused_append_shared_experts_with_weights에 있습니다.

`_fused_append_shared_experts_with_weights_kernel`

Before:

def _fused_append_shared_experts_with_weights_kernel(
    out_ids_ptr,
    out_weights_ptr,
    N_BASE,
    K: tl.constexpr,
    S: tl.constexpr,
    BLOCK_K: tl.constexpr,
    BLOCK_S: tl.constexpr,
):
    # ... (생략)
    shared_ids = tl.cast(N_BASE + offs_s, ids.dtype)
    shared_ws = tl.load(shared_weights_ptr + pid * S + offs_s, mask=mask_s)

    tl.store(out_ids_ptr + out_row_ptr + K + offs_s, shared_ids, mask=mask_s)
    tl.store(out_weights_ptr + out_row_ptr + K + offs_s, shared_ws, mask=mask_s)

After:

def _fused_append_shared_experts_with_weights_kernel(
    out_ids_ptr,
    out_weights_ptr,
    N_BASE,
    scale,
    K: tl.constexpr,
    S: tl.constexpr,
    BLOCK_K: tl.constexpr,
    BLOCK_S: tl.constexpr,
    APPLY_SIGMOID: tl.constexpr,
):
    # ... (생략)
    shared_ids = tl.cast(N_BASE + offs_s, ids.dtype)
    shared_ws = tl.load(shared_weights_ptr + pid * S + offs_s, mask=mask_s)
    if APPLY_SIGMOID:
        # Fuse sigmoid(shared_gate) + dtype upcast (+ optional 1/ep_size scale)
        # in-register so the raw bf16 logits stream straight into the fp32
        # output, eliminating the standalone sigmoid and bf16->fp32 copy kernels.
        shared_ws = tl.sigmoid(shared_ws.to(tl.float32)) * scale

    tl.store(out_ids_ptr + out_row_ptr + K + offs_s, shared_ids, mask=mask_s)
    tl.store(out_weights_ptr + out_row_ptr + K + offs_s, shared_ws, mask=mask_s)

변경사항은 scale과 APPLY_SIGMOID라는 두 개의 새로운 인자가 추가된 것입니다. 특히 APPLY_SIGMOID가 True일 경우, shared_ws (shared expert의 gate logits)를 레지스터에서 tl.float32로 캐스팅한 후 tl.sigmoid를 적용하고 scale을 곱하는 연산이 추가되었습니다. 이로써 기존에 별개의 커널로 실행되던 sigmoid 활성화 함수와 bf16->fp32 캐스팅이 이 커널 내부로 융합됩니다.

`fused_append_shared_experts_with_weights`

Before:

def fused_append_shared_experts_with_weights(
    topk_ids, topk_weights, shared_weights, num_fused_shared_experts, N=None
):
    # ... (생략)
    shared_weights_2d = shared_weights.to(topk_weights.dtype)
    # ... (생략)
    _fused_append_shared_experts_with_weights_kernel[
        grid
    ](
        out_ids,
        out_weights,
        N_BASE=N,
        K=k,
        S=s,
        BLOCK_K=block_k,
        BLOCK_S=block_s,
        num_warps=1,
    )
    return out_ids, out_weights

After:

def fused_append_shared_experts_with_weights(
    topk_ids,
    topk_weights,
    shared_weights,
    num_fused_shared_experts,
    N=None,
    apply_sigmoid=False,
    scale=1.0,
):
    # ... (생략)
    # When fusing sigmoid in-kernel, keep the raw logits dtype (the kernel emits
    # fp32 directly); otherwise match the output weight dtype as before.
    shared_weights_2d = (
        shared_weights if apply_sigmoid else shared_weights.to(topk_weights.dtype)
    )
    # ... (생략)
    _fused_append_shared_experts_with_weights_kernel[
        grid
    ](
        out_ids,
        out_weights,
        N_BASE=N,
        scale=scale,
        K=k,
        S=s,
        BLOCK_K=block_k,
        BLOCK_S=block_s,
        APPLY_SIGMOID=apply_sigmoid,
        num_warps=1,
    )
    return out_ids, out_weights

래퍼 함수 fused_append_shared_experts_with_weights에는 apply_sigmoid와 scale 인자가 추가되었습니다. apply_sigmoid가 True일 경우, shared_weights는 raw logits로 간주되어 shared_weights.to(topk_weights.dtype) 캐스팅을 건너뛰고, scale과 APPLY_SIGMOID 플래그가 Triton 커널로 전달됩니다. 이는 호스트 측에서 불필요한 데이터 타입 변환을 제거하고, 커널 내부에서 모든 연산을 처리하도록 합니다.

`python/sglang/srt/models/qwen2_moe.py`

이 파일은 Qwen2 MoE 모델의 로직을 담당하며, shared expert 가중치를 계산하고 append하는 부분을 수정합니다.

`_get_shared_expert_weights`

Before:

def _get_shared_expert_weights(
    self, hidden_states: torch.Tensor
) -> Optional[torch.Tensor]:
    # ... (생략)
    shared_logits = shared_out[0] if isinstance(shared_out, tuple) else shared_out
    w = F.sigmoid(shared_logits)
    # ... (생략)
    if moe_ep_size > 1 and not is_deepep_class_backend():
        w = w / float(moe_ep_size)
    return w

After:

def _get_shared_expert_weights(
    self, hidden_states: torch.Tensor
) -> Optional[Tuple[torch.Tensor, float]]:
    # ... (생략)
    shared_logits = shared_out[0] if isinstance(shared_out, tuple) else shared_out
    scale = 1.0
    moe_ep_size = get_parallel().moe_ep_size
    if moe_ep_size > 1 and not is_deepep_class_backend():
        scale = 1.0 / float(moe_ep_size)
    # Only AITER fuses sigmoid + cast in-kernel; on CUDA keep the legacy
    # eager activation so the NVIDIA path behavior is unchanged.
    if not _use_aiter:
        return F.sigmoid(shared_logits) * scale, 1.0
    return shared_logits, scale

_get_shared_expert_weights 함수는 이제 _use_aiter 플래그에 따라 다른 동작을 합니다. AMD AITER 백엔드를 사용하는 경우 (즉, _use_aiter가 True인 경우), F.sigmoid를 미리 적용하지 않고 raw gate logits와 1/ep_size 스케일 팩터를 튜플 형태로 반환합니다. 이는 sigmoid와 스케일링 연산이 Triton 커널 내부에서 처리될 수 있도록 하기 위함입니다. NVIDIA CUDA 경로의 경우, 기존 로직을 유지하여 F.sigmoid(shared_logits) * scale을 미리 계산하여 반환합니다.

`_append_shared_to_topk_output`

Before:

def _append_shared_to_topk_output(
    self,
    topk_output: StandardTopKOutput,
    hidden_states: torch.Tensor,
) -> StandardTopKOutput:
    # ... (생략)
    shared_weights = self._get_shared_expert_weights(hidden_states)
    # ... (생략)
    fused_topk_ids, fused_topk_weights = fused_append_shared_experts_with_weights(
        topk_output.topk_ids,
        topk_output.topk_weights,
        shared_weights,
        self.num_fused_shared_experts,
        N=self.num_experts,
    )
    return StandardTopKOutput(
        topk_weights=fused_topk_weights,
    )

After:

def _append_shared_to_topk_output(
    self,
    topk_output: StandardTopKOutput,
    hidden_states: torch.Tensor,
) -> StandardTopKOutput:
    # ... (생략)
    shared = self._get_shared_expert_weights(hidden_states)
    if shared is None:
        return topk_output
    shared_weights, shared_scale = shared

    # ... (생략)
    # AITER returns raw logits + scale for in-kernel sigmoid fusion; CUDA
    # returns pre-activated weights (scale already folded in) → no fusion.
    fused_topk_ids, fused_topk_weights = fused_append_shared_experts_with_weights(
        topk_output.topk_ids,
        topk_output.topk_weights,
        shared_weights,
        self.num_fused_shared_experts,
        N=self.num_experts,
        apply_sigmoid=_use_aiter,
        scale=shared_scale,
    )
    return StandardTopKOutput(
        topk_weights=fused_topk_weights,
    )

_append_shared_to_topk_output 함수는 _get_shared_expert_weights에서 반환된 튜플을 shared_weights와 shared_scale로 언팩합니다. 그리고 fused_append_shared_experts_with_weights를 호출할 때 apply_sigmoid=_use_aiter와 scale=shared_scale 인자를 전달하여, AMD AITER 백엔드에서만 커널 융합이 활성화되도록 합니다.

왜 이게 좋은가

성능 최적화

이 PR의 가장 큰 장점은 GPU 커널 런치 오버헤드를 크게 줄여 성능을 개선한다는 점입니다. 기존에는 MoE 레이어당 3개의 개별 커널이 실행되었지만, 이제는 1개의 커널로 모든 작업을 처리합니다. 이는 특히 작은 텐서에 대한 연산에서 두드러지는 효과를 발휘합니다.

torch profiler를 사용한 벤치마킹 결과는 다음과 같습니다 (디코딩 단계):

Kernel	Before	After	Notes
`sigmoid_kernel_cuda`	~5.5 us / layer	0 occurrences	제거됨
`bfloat16tofloat32_copy_kernel_cuda`	~4.5 us / layer	제거됨
`_fused_append_shared_experts_with_weights_kernel`	~4.2 us	780× @ 3.79 us	sigmoid + cast를 레지스터에서 처리 (오버헤드 거의 없음)

프로파일링 결과에서 보듯이, sigmoid_kernel_cuda와 bfloat16tofloat32_copy_kernel_cuda는 완전히 사라졌습니다. _fused_append_shared_experts_with_weights_kernel은 이제 sigmoid와 bf16->fp32 캐스팅을 레지스터 내에서 처리함에도 불구하고, 실행 시간이 오히려 약간 감소했습니다 (4.2 us → 3.79 us). 이는 커널 런치 오버헤드와 전역 메모리 접근을 줄인 효과가 매우 크다는 것을 보여줍니다.

단일 호출당 약 10 us (5.5 + 4.5)의 시간을 절약하며, 이는 기존 append 커널에 거의 추가 비용 없이 통합되었습니다. 이러한 최적화는 'decode glue' 부분에서 발생하며, 전체 추론 시간의 0.1% 미만을 차지하는 부분이지만, 수많은 MoE 레이어와 디코딩 단계에서 누적되면 상당한 이득을 가져올 수 있습니다.

정확도 및 수치적 안정성

GSM8K (200 질문, 병렬 2000, greedy) 테스트 결과, 평균 차이가 0.013으로, 실행 간 표준 편차(σ ≈ 0.045) 내에 있어 회귀가 없음을 확인했습니다. 이는 최적화가 모델의 정확도에 부정적인 영향을 미치지 않았음을 의미합니다.

오프라인 커널 등가성 테스트에서는 라우팅된 가중치(routed weights)가 비트 단위로 동일하고, expert ID도 정확했습니다. shared-expert 컬럼은 약 4e-4 정도의 차이를 보였는데, 이는 융합된 경로가 sigmoid 연산을 fp32로 계산하는 반면, 기존 경로는 bf16으로 계산했기 때문입니다. 즉, 융합된 경로는 수치적으로 더 정밀하다고 볼 수 있습니다.

일반적인 교훈

커널 융합(Kernel Fusion)의 중요성: GPU에서 작은 연산들이 연속적으로 실행될 때 발생하는 커널 런치 오버헤드는 무시할 수 없습니다. 서로 의존적인 작은 연산들을 하나의 커널로 융합하는 것은 성능 향상에 매우 효과적인 전략입니다.
Triton의 활용: Triton과 같은 DSL(Domain-Specific Language)은 GPU 커널을 직접 작성하여 세밀한 최적화를 가능하게 합니다. 특히 tl.constexpr을 사용하여 컴파일 타임에 조건부 로직을 결정함으로써, 런타임 오버헤드 없이 유연한 커널을 구현할 수 있습니다.
백엔드별 최적화: _use_aiter와 같은 플래그를 통해 특정 하드웨어(AMD) 및 백엔드에 특화된 최적화를 적용하면서도, 다른 백엔드(NVIDIA CUDA)의 기존 동작을 유지하는 유연한 코드 구조는 다양한 환경을 지원하는 데 필수적입니다.
정확도와 성능의 균형: 성능 최적화 시 정확도 회귀가 없는지, 오히려 수치적 정밀도가 향상되는지 면밀히 검토하는 것이 중요합니다.

리뷰 댓글 분석 및 추가 검증

PR 리뷰 과정에서 amd-bot은 중요한 피드백을 제공했습니다. 가장 핵심적인 내용은 PR CI (Continuous Integration) 테스트가 이 PR의 변경된 코드 경로를 충분히 검증하지 못했다는 점입니다. 특히 AMD AITER 백엔드에서 shared-expert sigmoid 융합이 활성화되는 조건(_use_aiter가 True이고 Qwen2MoeSparseMoeBlock 모델이 shared expert를 사용하는 경우)이 PR CI 테스트 스위트에서 충분히 커버되지 않았다는 지적이 있었습니다.

amd-bot은 nightly-amd 스위트에서 실행되는 Qwen2-57B-A14B-Instruct와 같은 shared-expert 모델에 대한 테스트나, fused_append_shared_experts_with_weights(apply_sigmoid=True)에 대한 단위 테스트 추가를 권고했습니다. 이는 최적화된 코드 경로가 실제 프로덕션 환경에서 예상대로 작동하고 수치적 정확성을 유지하는지 확인하기 위한 필수적인 단계입니다.

이에 대해 yichiche는 새로운 단위 테스트 test/registered/moe/test_fused_append_shared_experts.py를 추가하여 검증을 수행했습니다. 이 테스트는 test_fused_sigmoid_matches_eager를 포함하여 융합된 sigmoid 연산이 eager 모드와 일치하는지 확인하고, test_legacy_path_unchanged를 통해 기존 경로의 불변성을 검증했습니다. 모든 테스트가 PASSED로 확인되면서, CI 커버리지 부족 문제는 해소되었고, 융합된 커널의 정확성이 성공적으로 검증되었습니다.

이러한 리뷰 과정은 단순한 코드 변경을 넘어, 시스템의 복잡성을 이해하고 잠재적인 위험을 식별하며, 충분한 검증을 통해 고품질의 소프트웨어를 제공하는 데 있어 CI와 단위 테스트의 중요성을 다시 한번 강조합니다.

결론

이번 SGLang PR은 MoE shared expert 처리 과정에서 발생하는 불필요한 GPU 커널 런치 오버헤드를 제거하고, 3개의 커널을 1개로 융합하는 성공적인 최적화 사례를 보여주었습니다. Triton 커널 내부에서 sigmoid 활성화 함수와 데이터 타입 캐스팅을 인라인으로 처리함으로써, 약 10 us/layer의 시간을 절약하고 수치적 정밀도를 향상시켰습니다. 이는 SGLang의 MoE 모델 서빙 효율성을 한 단계 끌어올리는 중요한 기여이며, GPU 프로그래밍에서 커널 융합의 강력함을 다시 한번 입증했습니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM, Diffusion-Gemma 샘플러 메모리 최적화: 요청 기반 타일링으로 OOM 문제 해결
현재글 : [sglang] SGLang MoE Shared Expert 최적화: 3개 커널을 1개로 융합하여 GPU 오버헤드 제거
다음글 [flashinfer] FlashInfer의 BF16 GEMM 성능 극대화: CUDA Graph와 Cold L2 Cache 도입

댓글

관련 포스트

PR Analysis 의 다른글