[sglang] SGLang의 add_constant 커널 최적화: 아키텍처 인지 벡터화(Vectorization) 도입

PR 링크: sgl-project/sglang#24755 상태: Merged | 변경: +130 / -10

들어가며

SGLang의 add_constant 연산은 텐서의 모든 요소에 상수를 더하는 단순한 작업이지만, LLM 추론 엔진의 성능을 극대화하기 위해서는 이러한 기초 연산의 효율성이 매우 중요합니다. 기존 구현은 스칼라 기반의 커널을 사용하고 있어, 데이터 크기가 매우 클 경우 메모리 대역폭을 충분히 활용하지 못하는 한계가 있었습니다. 본 PR은 하드웨어 아키텍처(Pre-Blackwell vs Blackwell)를 고려한 벡터화(Vectorized) 경로를 도입하여 대규모 텐서 연산 성능을 비약적으로 향상시켰습니다.

코드 분석

python/sglang/jit_kernel/csrc/add_constant.cuh

핵심 변경 사항은 add_constant_vectorized_kernel의 도입입니다. 기존 스칼라 커널은 각 스레드가 하나의 int32_t 요소를 처리했으나, 새로운 구현은 device::AlignedVector를 사용하여 여러 요소를 한 번에 로드/저장합니다.

Before (스칼라 방식):

__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < length) {
    dst[idx] = src[idx] + kConstant;
  }
}

After (벡터화 방식):

__global__ void add_constant_vectorized_kernel(int32_t* dst, const int32_t* src, size_t length) {
  using Vector = device::AlignedVector<int32_t, kElementsPerVector>;
  // ... (생략) ...
  if (work_idx < vector_count) {
    auto values = device::load_as<Vector>(src, work_idx);
#pragma unroll
    for (size_t i = 0; i < kElementsPerVector; ++i) {
      values[i] += kConstant;
    }
    device::store_as<Vector>(dst, values, work_idx);
  }
  // ... (tail 처리) ...
}

또한, add_constant 함수 내부에서 텐서의 정렬(Alignment) 상태와 크기를 확인하여, 최적화된 벡터 경로를 탈지 스칼라 경로를 탈지 동적으로 결정하는 로직이 추가되었습니다.

왜 이게 좋은가

1. 메모리 대역폭 활용 극대화: 벡터화된 로드/스토어 명령어를 사용함으로써 GPU의 메모리 인터페이스를 더 효율적으로 점유합니다. H200 환경에서 16M 요소 기준 약 32%~35%의 성능 향상을 기록했습니다.
2. 아키텍처 유연성: device::kMaxVecBytes를 활용하여 하드웨어(Blackwell의 32바이트 vs 이전 세대의 16바이트)에 맞게 벡터 폭을 자동으로 조정합니다.
3. 안전한 폴백(Fallback): 정렬되지 않은 메모리 주소나 작은 크기의 텐서에 대해서는 기존 스칼라 커널을 유지하여 정확성과 범용성을 보장합니다.

일반적 교훈

• 데이터 정렬(Alignment)의 중요성: 벡터화 연산은 메모리 정렬이 필수적입니다. is_aligned_for_vector와 같은 체크 로직을 통해 정렬되지 않은 데이터로 인한 성능 저하나 오류를 방지해야 합니다.
• Tail 처리: 벡터화 시 전체 길이가 벡터 크기의 배수가 아닐 경우, 남은 요소를 처리하는 'Tail' 로직을 반드시 포함해야 합니다.
• Launch-bound 연산: 작은 텐서 연산은 커널 실행 오버헤드(Launch overhead)가 지배적이므로, 무리한 최적화보다는 스칼라 커널을 유지하는 것이 전략적으로 유리합니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#vector-types

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