[sglang] Intel GPU 가속을 위한 SGLang MoE 커널 최적화: GPT-OSS bf16 지원 분석

PR 링크: sgl-project/sglang#22417 상태: Merged | 변경: +None / -None

들어가며

대규모 언어 모델(LLM)의 추론 성능을 높이기 위한 핵심 기술 중 하나는 **Kernel Fusion(커널 퓨전)**입니다. 특히 Mixture of Experts(MoE) 구조를 가진 모델은 수많은 전문가(Expert) 레이어를 거쳐야 하므로, 개별 연산을 하나로 묶어 메모리 대역폭 낭비를 줄이는 것이 필수적입니다.

최근 SGLang 프로젝트에 반영된 PR에서는 Intel GPU(XPU) 환경에서 GPT-OSS와 같은 BF16 기반 MoE 모델의 성능을 최적화하기 위해 fused_experts 커널 경로를 활성화하고, 정밀한 연산 제어를 위한 파라미터를 추가했습니다. 이번 글에서는 이 변경사항이 왜 중요한지, 그리고 어떤 기술적 이점을 제공하는지 분석해 보겠습니다.

코드 분석: 정밀한 GEMM 제어를 위한 파라미터 주입

이번 업데이트의 핵심은 python/sglang/srt/layers/quantization/unquant.py 파일 내 forward_xpu 함수에 있습니다. 기존에는 Intel GPU 전용 MoE 커널을 호출할 때 가중치와 활성화 함수 정보만 전달했으나, 이제는 GEMM(General Matrix Multiply) 연산의 스케일링과 클램핑을 위한 파라미터가 추가되었습니다.

Before

# 기존 코드: 기본적인 가중치와 활성화 함수만 전달
output = sgl_kernel_xpu.fused_experts(
    hidden_states, 
    w1, 
    w2, 
    topk_weights, 
    topk_ids,
    b1=getattr(layer, "w13_weight_bias", None),
    b2=getattr(layer, "w2_weight_bias", None),
    activation=moe_runner_config.activation,
)

After

# 변경된 코드: gemm1_alpha 및 gemm1_limit 파라미터 추가
output = sgl_kernel_xpu.fused_experts(
    hidden_states, 
    w1, 
    w2, 
    topk_weights, 
    topk_ids,
    b1=getattr(layer, "w13_weight_bias", None),
    b2=getattr(layer, "w2_weight_bias", None),
    activation=moe_runner_config.activation,
    gemm1_alpha=moe_runner_config.gemm1_alpha,        # 추가
    gemm1_limit=moe_runner_config.gemm1_clamp_limit,  # 추가
)

왜 이게 좋은 최적화인가?

1. 하드웨어 가속기(XPU) 최적화 경로 활성화

GPT-OSS 모델은 BF16 정밀도를 사용합니다. 기존에는 Intel GPU 환경에서 이러한 모델들이 일반적인(Generic) 연산 경로를 탔을 가능성이 높으나, 이번 PR을 통해 sgl-kernel-xpu의 전용 fused_experts 커널을 사용할 수 있게 되었습니다. 이는 커널 실행 오버헤드를 줄이고 연산 유닛의 활용도를 극대화합니다.

2. 수치 안정성 및 정확도 확보 (gemm1_alpha, gemm1_limit)

MoE 레이어의 첫 번째 GEMM 연산 결과에 적용되는 alpha(스케일링 인자)와 limit(클램핑 임계값)을 커널 내부로 직접 전달합니다.

• gemm1_alpha: 활성화 함수 적용 전 값의 범위를 조정하여 수치적 안정성을 높입니다.
• gemm1_limit: 특정 범위를 벗어나는 값을 제한(Clamp)하여 FP16/BF16 연산 시 발생할 수 있는 오버플로우나 언더플로우를 방지합니다.

이러한 파라미터를 커널 외부가 아닌 Fused Kernel 내부에서 처리함으로써, 별도의 클램핑 커널을 실행할 필요가 없어 메모리 읽기/쓰기 횟수가 줄어듭니다.

성능 및 검증 결과

PR 작성자가 공유한 테스트 결과에 따르면, Intel B60 GPU 4장을 사용한 TP(Tensor Parallelism) 4 환경에서 GSM8K 벤치마크 결과가 **Nvidia A100과 유사한 수준(Exact Match 0.8711)**으로 나타났습니다. 이는 최적화된 커널이 단순히 빠를 뿐만 아니라, 모델의 추론 정확도를 완벽하게 유지하고 있음을 증명합니다.

결론

이번 변경사항은 Intel GPU 환경에서 SGLang의 경쟁력을 높이는 중요한 단계입니다. 특히 fused_experts 커널에 세부적인 연산 제어 파라미터를 노출함으로써, 다양한 MoE 아키텍처 모델들이 하드웨어 특화 가속을 받을 수 있는 기반을 마련했습니다. 시니어 엔지니어 관점에서 볼 때, 이러한 파라미터화된 커널 설계는 코드의 재사용성을 높이고 하드웨어 추상화 계층을 견고하게 만드는 훌륭한 접근 방식입니다.

참고 자료

• https://github.com/sgl-project/sgl-kernel
• https://pytorch.org/docs/stable/notes/cuda.html#bfloat16-support

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