[sglang] fused_qknorm_rope 최적화 - interleave RoPE에서 sincosf 중복 제거

2026년 4월 1일수정: 2026년 4월 1일

PR 링크: sgl-project/sglang#21654 상태: Merged | 변경: +208 / -77

들어가며

fused_qknorm_rope는 RMSNorm과 RoPE(Rotary Positional Encoding)를 하나의 CUDA 커널로 융합한 연산이다. interleave(GPT-J) 방식의 RoPE에서 인접한 요소 쌍 (2k, 2k+1)은 동일한 frequency/theta 값을 사용한다. 기존 구현은 모든 요소에 대해 개별적으로 powf와 __sincosf를 호출했지만, 이 PR은 쌍당 한 번만 계산하여 이 비싼 수학 함수 호출을 절반으로 줄인다.

핵심 코드 분석

1. compute_freq 함수 템플릿화

Before:

__device__ inline float
compute_freq_yarn(float base, int rotary_dim, int half_dim, float factor, float low, float high) {
  float freq = powf(base, -2.0f * half_dim / static_cast<float>(rotary_dim));
  if (factor != 1.0f) {
    // YaRN blending 로직
  }
  return freq;
}

After:

template <bool yarn>
__device__ inline float compute_freq(float base, int rotary_dim, int half_dim, float factor, float low, float high) {
  float freq = powf(base, -2.0f * half_dim / static_cast<float>(rotary_dim));
  if constexpr (yarn) {
    // YaRN blending 로직 - 컴파일 타임에 분기 제거
  }
  return freq;
}

런타임 if (factor != 1.0f) 조건을 if constexpr (yarn) 템플릿 파라미터로 변환하여 컴파일 타임에 분기를 제거한다.

2. Interleave RoPE에서 sincosf 중복 제거

Before:

if constexpr (interleave) {
  for (int i = 0; i < numElemsPerThread; i++) {
    elements2[i] = (i % 2 == 0) ? -elements[i + 1] : elements[i - 1];
    // 매 요소마다 개별 freq 계산 + __sincosf 호출
    int half_dim = (offsetThread + i) / 2;
    float freq = compute_freq(...);
    float theta = pos_id * freq;
    __sincosf(theta, &sin_vals[i], &cos_vals[i]);
  }
}

After:

if constexpr (interleave) {
  // 쌍 단위로 처리: (2k, 2k+1)이 같은 freq/theta 공유
  for (int p = 0; p < numElemsPerThread; p += 2) {
    int half_dim = (offsetThread + p) / 2;
    float freq = compute_freq<yarn>(...);
    float theta = pos_id * freq;
    float s, c;
    __sincosf(theta, &s, &c);
    // 동일한 sin/cos를 쌍의 두 요소에 적용
    float e0 = elements[p], e1 = elements[p + 1];
    elements[p]     = e0 * c - e1 * s;
    elements[p + 1] = e0 * s + e1 * c;
  }
}

핵심 최적화: powf + __sincosf 호출이 요소당 1회에서 쌍당 1회로 줄어든다. 이 함수들은 GPU에서 상대적으로 비싼 특수함수(SFU)이므로 호출 횟수 절반 감소는 직접적인 성능 향상으로 이어진다.

3. Production shape 벤치마크 추가

PRODUCTION_SHAPES = [
    ("flux_1024", 4096, 24, 24, 24, 128, 128),
    ("qwen_image_1024", 4096, 32, 32, 32, 128, 128),
    ("zimage_1024", 4096, 30, 30, 30, 128, 128),
]

실제 프로덕션 모델 구성(Flux, Qwen, ZImage)에 대한 벤치마크를 추가하여 최적화 효과를 검증할 수 있다.