[sglang] SGLang에서 Qwen3-Next FP8 MoE 최적화: H200을 위한 Shared-Expert Fusion

PR 링크: sgl-project/sglang#26204 상태: Merged | 변경: +427 / -28

들어가며

최신 대규모 언어 모델(LLM) 아키텍처, 특히 Mixture-of-Experts(MoE) 모델은 추론 성능을 높이기 위해 Shared-Expert를 활용하는 경우가 많습니다. 하지만 Shared-Expert를 별도의 연산으로 처리하면 메모리 접근과 커널 실행 오버헤드가 발생합니다. 본 PR은 Qwen3-Next 모델을 NVIDIA H200 GPU 환경에서 구동할 때, Shared-Expert를 MoE Expert 슬롯으로 융합(Fusion)하여 성능을 개선하고, FP8 연산을 위한 최적화된 Triton 커널 설정을 추가하는 내용을 담고 있습니다.

코드 분석

1. Triton 커널 설정 추가

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/ 경로에 H200을 위한 새로운 JSON 설정 파일들이 추가되었습니다. 이는 FP8(W8A8) 연산 시 E=513, N=512 구성에 최적화된 블록 크기와 스테이지 설정을 제공합니다.

// E=513,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json
{
    "1": {
        "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128,
        "GROUP_SIZE_M": 32, "num_warps": 4, "num_stages": 4
    },
    // ... (다양한 배치 사이즈별 설정)
}

이 설정은 TMA(Tensor Memory Accelerator)를 활용하여 H200의 메모리 대역폭을 최대한 활용하도록 설계되었습니다.

2. 모델 로직 수정: qwen2_moe.py

Shared-Expert를 MoE Expert로 융합하기 위해 can_fuse_shared_expert 로직이 강화되었고, Qwen2MoeSparseMoeBlock 초기화 시 enable_cuda_shared_expert_fusion 플래그가 추가되었습니다.

# Before
def can_fuse_shared_expert(config, quant_config):
    # ... 기존 로직

# After
def can_fuse_shared_expert(config, quant_config):
    if quant_config is not None:
        exclude_layers = getattr(quant_config, "exclude_layers", None)
        if exclude_layers is None:
            exclude_layers = getattr(quant_config, "ignored_layers", [])
        # ... 레이어 체크 로직 보완

왜 이게 좋은가

이번 최적화의 핵심은 Shared-Expert Fusion입니다. 기존에는 Shared-Expert가 별도의 연산으로 처리되어 GPU 커널 실행 간의 컨텍스트 스위칭 비용이 발생했으나, 이를 MoE Expert 슬롯으로 융합함으로써 단일 커널 내에서 처리가 가능해졌습니다.

성능 지표:

• Chat 워크로드(1000/1000) 기준, SGLang baseline 대비 약 6% 이상의 처리량(tok/s) 향상을 보였습니다.
• p50 TTFT(Time To First Token)가 983ms에서 924ms로 단축되어 응답 속도가 개선되었습니다.

교훈:

1. 커널 융합(Kernel Fusion)의 중요성: 메모리 바운드 연산이 많은 MoE 모델에서 불필요한 커널 호출을 줄이는 것은 성능 향상의 핵심입니다.
2. 하드웨어 특화 설정: H200과 같은 최신 GPU에서는 TMA와 같은 하드웨어 가속 기능을 활용하기 위해 Triton 설정(num_stages, block_size)을 세밀하게 튜닝하는 것이 필수적입니다.

리뷰어 피드백 반영

리뷰 과정에서 Qwen3-Next 모델이 밀집(dense) 레이어를 포함할 수 있다는 점이 지적되었습니다. 기존에는 layers[0]만 확인하여 융합 여부를 결정했으나, 이제는 모델 전체를 스캔하여 Qwen2MoeSparseMoeBlock이 존재하는지 확인한 후 융합 여부를 결정하도록 로직이 개선되었습니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://github.com/sgl-project/sglang

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