[sglang] DeepSeek NextN을 위한 Fused EH Norm 최적화: 커널 융합으로 성능 극대화하기

PR 링크: sgl-project/sglang#29667 상태: Merged | 변경: +386 / -12

들어가며

최신 LLM 추론 엔진인 SGLang에서 DeepSeek 모델의 성능을 극대화하기 위한 중요한 최적화가 진행되었습니다. 특히 DeepSeek NextN 모델 구조에서 사용되는 EH Norm(Embeddings + Hidden states Normalization) 연산은 기존에 개별적인 torch.norm과 torch.cat 연산으로 수행되어 메모리 접근 오버헤드가 컸습니다. 본 PR은 이 과정을 하나의 CUDA 커널로 융합(Fusion)하여 메모리 대역폭 사용을 최적화하고 추론 속도를 획기적으로 개선했습니다.

코드 분석

1. 커널 구현: fused_eh_norm.cuh

핵심은 fused_eh_norm_kernel을 통해 두 개의 정규화 연산을 하나의 GPU 블록 내에서 처리하는 것입니다.

Before (Conceptual):

# 기존 방식: 두 번의 정규화 후 cat 연산으로 메모리 복사 발생
output = torch.cat((self.enorm(hidden_states), self.hnorm(previous_hidden_states)), dim=-1)

After (CUDA Kernel):

// FusedEHNormParams를 통해 데이터를 한 번에 로드하고 정규화 수행
const auto embeds_vec = gmem.load(embeds);
const auto embeds_output_vec = norm::apply_norm_cta<kHidden>(embeds_vec, enorm_weight_vec, params.eps, smem, kNumWarps);
gmem.store(output, embeds_output_vec);

const auto prev_vec = gmem.load(previous_hidden);
const auto prev_output_vec = norm::apply_norm_cta<kHidden>(prev_vec, hnorm_weight_vec, params.eps, smem, kNumWarps);
gmem.store(output + kHidden, prev_output_vec);

이 방식은 중간 결과를 메모리에 쓰지 않고 레지스터와 공유 메모리(Shared Memory)를 적극 활용하여 Global Memory 접근 횟수를 최소화합니다.

2. 모델 적용: deepseek_nextn.py

모델의 forward 패스에서 _is_cuda 환경일 경우 최적화된 커널을 호출하도록 분기 처리했습니다.

if _is_cuda:
    eh_input = fused_eh_norm(hidden_states, previous_hidden_states, ...)
else:
    eh_input = torch.cat((self.enorm(hidden_states), self.hnorm(previous_hidden_states)), dim=-1)

왜 이게 좋은가

성능 향상

벤치마크 결과, 기존 PyTorch 구현 대비 약 10배에서 20배 이상의 처리량(GB/s) 향상을 보였습니다. 특히 토큰 수가 많아질수록 커널 융합의 이점이 극대화되어, 512 토큰 기준 torch 구현이 115us 걸리던 연산이 jit 커널에서는 7us 수준으로 단축되었습니다.

교훈

1. Memory Bound 문제 해결: 정규화와 같은 연산은 연산량보다 메모리 대역폭이 병목인 경우가 많습니다. 커널 융합은 이러한 메모리 I/O를 획기적으로 줄여줍니다.
2. JIT의 활용: sglang의 jit_kernel 프레임워크를 사용하여 특정 하드웨어 아키텍처에 최적화된 코드를 동적으로 생성함으로써 범용성과 성능을 모두 잡았습니다.

리뷰어 피드백

리뷰 과정에서 테스트 등록 방식에 대한 논의가 있었으며, 최종적으로는 CI 안정성을 위해 유지하기로 결정되었습니다. 이는 성능 최적화 코드가 단순한 구현을 넘어 운영 환경에서도 안정적으로 동작해야 함을 시사합니다.

References

• SGLang GitHub
• CUDA C++ Programming Guide

참고 자료

• https://github.com/sgl-project/sglang
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/

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