[sglang] SGLang NGRAM 성능 최적화: 호스트 기반 트리 링크 유도로 GPU 병목 제거하기

PR 링크: sgl-project/sglang#32380 상태: Merged | 변경: +58 / -13

들어가며

LLM 추론 가속을 위한 Speculative Decoding, 특히 NGRAM 기반 방식에서는 draft tree를 구성하고 이를 검증하는 과정이 필수적입니다. 기존 SGLang의 NGRAM 구현체는 draft tree의 구조 정보를 GPU 커널에서 계산한 뒤, 이를 다시 호스트(CPU)로 읽어오는 .cpu() 호출을 수행했습니다. 이 과정에서 GPU는 호스트의 응답을 기다리며 유휴 상태(idle)에 빠지게 되며, 이는 전체 추론 파이프라인의 병목으로 작용했습니다. 본 PR은 이 트리 링크 계산을 호스트에서 직접 수행하도록 변경하여, 불필요한 GPU-CPU 동기화 지점을 제거하고 성능을 극대화했습니다.

코드 분석

1. 호스트 기반 트리 링크 계산 (ngram_worker.py)

기존에는 GPU 커널의 결과물을 기다려야 했으나, 이제는 호스트에서 mask 정보를 이용해 직접 next_token과 next_sibling을 유도합니다.

Before (기존 방식):

# GPU 커널 결과를 기다리는 blocking readback
retrieve_next_token_cpu = verify_input.retrieve_next_token.cpu()
retrieve_next_sibling_cpu = verify_input.retrieve_next_sibling.cpu()

After (개선된 방식):

def _derive_tree_links(mask, bs, draft_token_num):
    # 호스트에서 직접 트리 구조를 계산하여 GPU readback 제거
    tree = mask.reshape(bs, draft_token_num, draft_token_num)
    # ... (중략: 부모 노드 및 형제 노드 관계 유도 로직)
    return torch.from_numpy(next_token), torch.from_numpy(next_sibling)

2. 비동기 처리 및 파이프라인 최적화

데이터 전송 방식을 non_blocking으로 변경하고, 연산 순서를 조정하여 GPU의 forward pass와 CPU의 연산이 겹치도록(overlap) 설계했습니다.

# vocab-mask H2D 전송을 non_blocking으로 변경
vocab_mask = vocab_mask.to(
    verify_input.retrieve_next_token.device, non_blocking=True
)

왜 이게 좋은가

이번 최적화의 핵심은 'GPU-CPU 간의 불필요한 동기화 제거'와 '연산의 병렬화'입니다.

1. 성능 수치: Qwen2.5-Coder-7B 모델 기준, 처리량(throughput)이 기존 670.3 tok/s에서 725.7 tok/s로 약 8.3% 향상되었습니다. per-step 지연 시간은 11.89ms에서 10.98ms로 단축되었습니다.
2. GPU 활용도: GPU가 데이터를 기다리는 시간을 줄임으로써 fwd_occupancy가 84.15%에서 91.29%로 크게 개선되었습니다.
3. 교훈: LLM 추론 엔진 설계 시, 호스트가 이미 알고 있는 정보(draft tree 구조 등)를 굳이 GPU로 넘겨 계산한 뒤 다시 가져오는 것은 전형적인 성능 저하 패턴입니다. 데이터의 흐름을 분석하여 호스트에서 처리 가능한 연산을 분리해내는 것만으로도 상당한 오버헤드를 줄일 수 있습니다.

결론

이번 PR은 단순히 코드를 옮기는 것을 넘어, GPU 아키텍처의 특성을 고려한 비동기 파이프라인 설계의 중요성을 잘 보여줍니다. 특히 non_blocking 전송과 호스트 연산의 병렬화를 통해 추론 엔진의 실질적인 처리량을 크게 끌어올렸습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.to.html

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