[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를 기반으로 작성했습니다.