[sglang] [AMD ROCm] GLM-5.x Prefill 성능을 66% 끌어올린 Top-K 커널 최적화 분석

PR 링크: sgl-project/sglang#37889 상태: Merged | 변경: +477 / -8

들어가며

대규모 언어 모델(LLM)의 추론 속도를 결정짓는 핵심 요소 중 하나는 효율적인 커널 설계입니다. 특히 GLM-5.x와 같은 최신 모델을 AMD ROCm 환경에서 구동할 때, Prefill 단계의 Top-K 연산은 성능 병목의 원인이 되곤 했습니다. 기존 SGLang의 GLM-5.x 구현에서는 Decode 단계에만 최적화된 'v2' Top-K 커널을 적용했고, Prefill 단계는 여전히 레거시 커널(fused transform)에 의존하고 있었습니다.

이로 인해 발생하는 문제는 명확했습니다. Prefill 단계에서 사용하는 DSA(DeepSeek-V4/V5 스타일)의 'Packed' 레이아웃은 여러 요청이 하나의 거대한 스코어 버퍼를 공유하며, 각 행(row)이 임의의 컬럼 오프셋에서 시작됩니다. 기존의 Paged 커널은 각 행이 0번 컬럼에서 시작한다고 가정했기 때문에 Prefill에 바로 적용할 수 없었습니다.

이번 PR은 이러한 주소 지정(Addressing) 문제를 해결한 topk_packed_kernel을 도입하여, AMD GPU에서 Prefill Top-K 런칭 시간을 약 66% 단축하고 전체 토큰 처리량(Throughput)을 4.9% 향상시킨 최적화 사례입니다.

코드 분석: Packed 레이아웃을 위한 새로운 커널 설계

1. 주소 지정 방식의 변화: TopKPackedParams 도입

가장 큰 변화는 스코어 버퍼 내에서 각 행의 시작 위치를 유연하게 다룰 수 있게 된 점입니다. 기존 Paged 방식과 이번에 도입된 Packed 방식의 차이를 코드로 살펴보겠습니다.

Before (Paged Layout - Decode 중심): 기존 방식은 각 행이 독립적인 페이지 테이블 행을 소유한다고 가정했습니다.

// topk_v2.cuh (기존 논리적 구조)
// score window of row i: scores[i, :seq_lens[i]]
// page-table row: page_tables[i]

After (Packed Layout - Prefill 최적화): 새로운 TopKPackedParams는 row_starts와 row_to_batch를 통해 공유 버퍼 내의 임의 위치를 참조합니다.

// topk_v2.cuh
struct TopKPackedParams {
  float* __restrict__ scores;
  const int32_t* __restrict__ seq_lens;      // 행별 윈도우 길이
  const int32_t* __restrict__ row_starts;    // 행별 스코어 컬럼 오프셋
  const int32_t* __restrict__ row_to_batch;  // 행별 페이지 테이블 매핑
  // ... 생략 ...

  SGL_DEVICE PageTransform get_transform(uint32_t bx) const {
    // row_to_batch가 없으면 identity 매핑, 있으면 해당 배치 인덱스 참조
    const auto table_row = row_to_batch == nullptr ? bx : static_cast<uint32_t>(row_to_batch[bx]);
    return {page_table + static_cast<int64_t>(table_row) * page_table_stride, page_bits, nullptr};
  }
};

이 구조 덕분에 327K와 같은 거대한 컨텍스트에서도 스코어 버퍼를 복사하거나 행별로 거대한 페이지 테이블을 새로 만들 필요가 없어졌습니다. 이는 수백 MB의 메모리 오버헤드를 방지합니다.

2. Vectorized Load를 위한 정렬 및 마스킹

성능을 위해 16-byte vectorized load를 사용하려면 메모리 접근이 4-float 경계에 정렬되어야 합니다. 하지만 row_starts[i]는 임의의 값을 가질 수 있습니다. 이를 해결하기 위해 커널 내부에서 시작 지점을 내림(round down)하고 앞부분을 마스킹하는 기법을 사용했습니다.

// topk_v2.cuh 내 topk_packed_kernel
const auto rem = row_start % kVecSize;
if (rem != 0) {
  // 인덱서가 완료될 때까지 대기 (PDL 사용 시)
  device::PDLWaitPrimary<kPDL>();
  if (const auto tx = threadIdx.x; tx < rem) {
    // 정렬을 위해 끌어온 앞쪽 컬럼들을 padding_value로 마스킹
    score[row_start - rem + tx] = impl::padding_value();
  }
}
// 실제 연산 범위 조정
problem.in -= rem;
problem.seq_len = seq_len + rem;
problem.input_start = broadcast(rem);

이 방식은 하드웨어 수준의 정렬 요구사항을 충족하면서도 정확한 Top-K 결과를 보장합니다.

3. Python 백엔드 라우팅

dsa_topk_backend.py에서는 이제 Prefill(Extend) 단계에서 조건이 맞을 경우 새로운 v2 packed 커널을 호출합니다.

# dsa/dsa_topk_backend.py (개념적 변경)
if is_hip_runtime() and can_use_v2_packed:
    # 기존 legacy 커널 대신 v2 packed 커널 호출
    _topk_transform_v2_packed(
        scores, seq_lens, page_tables, out_page_indices, 
        page_size, row_starts=ks, row_to_batch=batch_idx_list
    )

왜 이게 좋은가?

1. 압도적인 커널 성능 향상

프로파일링 결과에 따르면, topk_packed_kernel은 기존 coop_topk_kernel(legacy) 대비 런칭당 실행 시간을 약 64%~67% 감소시켰습니다.

• Baseline: 896.8 µs ~ 1972.3 µs
• This PR: 319.7 µs ~ 649.1 µs

2. 실제 서비스 지표 개선

커널 수준의 개선은 실제 End-to-End 성능으로 이어졌습니다. ISL 70000 / OSL 300 설정에서:

• Token Throughput: +4.9% 향상
• Median TTFT (Time to First Token): -4.8% 감소
• Median TPOT (Time Per Output Token): -3.5% 감소

이는 대규모 컨텍스트를 처리하는 GLM-5.x 모델에서 사용자가 느끼는 체감 속도가 유의미하게 빨라졌음을 의미합니다.

3. 유연한 아키텍처 유지 (Reviewer Feedback 반영)

코드 리뷰 과정에서 EricKing626은 기존 Paged 커널의 추상화를 깨뜨리지 않기 위해 별도의 커널 함수(topk_packed_kernel)로 분리할 것을 제안했습니다. 이를 통해 코드의 복잡성을 관리하면서도 Prefill 특화 최적화를 달성할 수 있었습니다. 또한, 현재는 AMD ROCm 환경에서만 활성화되도록 USE_ROCM 가드를 사용했는데, 이는 CUDA 환경에서는 이미 FlashInfer 등을 통해 유사한 최적화가 제공되고 있기 때문입니다.

마치며

이번 최적화는 단순히 알고리즘을 개선한 것이 아니라, GPU 하드웨어의 메모리 정렬 특성과 모델 특유의 데이터 레이아웃(Packed DSA)을 깊이 이해하고 결합한 결과입니다. 특히 AMD ROCm 생태계에서 SGLang이 NVIDIA GPU에 못지않은 최적화 수준을 갖추어 가고 있다는 점에서 큰 의의가 있습니다.

고성능 커널 개발을 꿈꾸는 엔지니어라면, 이번 사례처럼 "불필요한 메모리 복사 제거"와 "하드웨어 친화적인 데이터 정렬"이 얼마나 큰 성능 차이를 만드는지 주목해 보시기 바랍니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.topk.html
• https://rocm.docs.amd.com/en/latest/
• https://github.com/sgl-project/sglang

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