[sglang] SGLang Ngram Speculative Decoding 최적화: MatchState 증분 업데이트 성능 개선

PR 링크: sgl-project/sglang#22180 상태: Merged | 변경: +None / -None

들어가며

SGLang의 Speculative Decoding 모듈에서 Ngram 기반의 매칭 로직은 추론 속도를 높이는 핵심 요소입니다. 최근 PR에서는 MatchState의 증분 업데이트(incremental advance) 과정에서 발생하는 불필요한 오버헤드를 제거하여 성능을 최적화했습니다. 특히, 매 토큰마다 발생하는 힙 할당을 제거하고, 불필요한 상태 조회 로직을 간소화함으로써 CPU 기반의 추론 효율성을 크게 개선했습니다.

코드 분석

1. trie.cpp: 힙 할당 제거 및 로직 간소화

가장 큰 성능 개선은 advanceMatchState_ 함수 내에서 이루어졌습니다. 기존 코드에서는 매 토큰마다 std::vector를 새로 생성하여 힙 할당이 발생했습니다.

Before:

for (size_t i = 0; i < len; ++i) {
  const auto next_depth = std::min(state.anchors.size() + 1, param_.max_trie_depth);
  std::vector<NodeRef> next(next_depth); // 매 루프마다 힙 할당 발생
  // ...
}

After:

std::vector<NodeRef> next;
next.reserve(param_.max_trie_depth); // 루프 밖에서 미리 할당

for (size_t i = 0; i < len; ++i) {
  const auto next_depth = std::min(state.anchors.size() + 1, param_.max_trie_depth);
  next.assign(next_depth, {}); // 기존 버퍼 재사용
  // ...
}

또한, root_에 직접 접근하도록 변경하여 resolve 함수를 통한 불필요한 검증 단계를 생략했습니다. 이는 root 노드가 절대 퇴거(eviction)되지 않는다는 도메인 지식을 활용한 최적화입니다.

2. 테스트 환경 개선

test_ngram_corpus.py를 spec/utils/에서 unit/spec/으로 이동시키고, register_cuda_ci 대신 register_cpu_ci를 사용하도록 변경했습니다. 이는 Ngram 코퍼스가 순수 C++ 로직임을 고려할 때, 불필요한 GPU 자원 점유를 막고 CI 파이프라인의 효율성을 높이는 적절한 조치입니다.

왜 이게 좋은가

이번 최적화의 핵심은 **'메모리 할당 최소화'**와 **'불필요한 추상화 제거'**입니다.

• 성능 수치: max_trie_depth=18 환경에서 증분 업데이트 방식은 기존 재구축(rebuild) 방식 대비 약 1.40배의 속도 향상을 보였습니다 (32.4 us/step vs 45.3 us/step).
• 교훈:
  1. Hot path에서의 할당 제거: 루프 내부의 std::vector 생성은 작은 크기라도 빈번하게 호출될 경우 성능 저하의 주범이 됩니다. reserve와 assign을 통한 버퍼 재사용은 필수적인 패턴입니다.
  2. 불변성 활용: root_ 노드와 같이 라이프사이클이 명확한 객체는 복잡한 조회 함수 대신 직접 접근하여 오버헤드를 줄일 수 있습니다.
  3. 테스트 격리: 순수 CPU 로직은 GPU CI에서 분리하여 전체적인 CI 속도를 개선하는 것이 운영 측면에서 유리합니다.

이번 PR은 단순히 속도만 높인 것이 아니라, TrieNode의 version 관리와 reset() 시의 epoch 기반 무효화 로직에 대한 명확한 주석을 추가하여 코드의 유지보수성까지 고려한 우수한 변경 사례입니다.

참고 자료

• https://en.cppreference.com/w/cpp/container/vector/reserve
• https://en.cppreference.com/w/cpp/container/vector/assign

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