본문으로 건너뛰기

[sglang] SGLang Ngram 추측 디코딩: 외부 코퍼스 기반 Suffix Automaton 통합으로 성능 최적화

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

들어가며

대규모 언어 모델(LLM)의 추론 속도를 높이는 것은 항상 중요한 과제입니다. SGLang은 이러한 문제를 해결하기 위해 다양한 최적화 기법을 제공하며, 그중 하나가 바로 Ngram 기반의 추측 디코딩(Speculative Decoding)입니다. 추측 디코딩은 작은(draft) 모델이나 Ngram과 같은 예측기를 사용하여 다음 토큰 시퀀스를 미리 생성하고, 이를 큰(target) 모델로 검증하여 전체적인 생성 속도를 향상시키는 기법입니다.

기존 Ngram 추측 디코딩은 주로 현재까지 생성된 토큰 시퀀스(recency)나 빈도(frequency)를 기반으로 Ngram 트라이(Trie)를 구축하여 다음 토큰을 예측했습니다. 하지만 이 방식은 모델이 학습하지 않은, 또는 현재 컨텍스트에 없는 새로운 패턴을 예측하는 데 한계가 있었습니다. 특히, 특정 도메인이나 외부 지식에 대한 예측 성능을 높이기 위해서는 외부 코퍼스(external corpus)의 활용이 필수적입니다.

이번 PR은 이러한 문제를 해결하기 위해 Ngram 추측 디코딩에 외부 코퍼스를 로드하고 Suffix Automaton(SAM)을 구축하여 활용하는 기능을 추가합니다. 이를 통해 Ngram 예측기의 커버리지를 확장하고, 더 풍부하고 정확한 드래프트 토큰 생성을 가능하게 하여 전체적인 추론 성능을 개선하는 것을 목표로 합니다.

코드 분석: 외부 코퍼스 기반 Suffix Automaton 통합

이 PR의 핵심은 Ngram 클래스에 SuffixAutomaton을 통합하고, 외부 코퍼스를 효율적으로 로드하여 SAM을 구축하는 파이프라인을 추가하는 것입니다. 주요 변경사항을 파일별로 살펴보겠습니다.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h

가장 먼저 Ngram 클래스에 SuffixAutomaton 멤버 변수가 추가되었습니다. 이는 Ngram 예측 로직에서 Trie와 함께 SAM을 활용할 수 있도록 하는 기반을 마련합니다.

Before:

class Ngram {
  std::unique_ptr<Trie> trie_;
  Param param_;
  // ... (생략)
};

After:

#include "suffix_automaton.h"
// ... (생략)
class Ngram {
  std::unique_ptr<Trie> trie_;
  std::unique_ptr<SuffixAutomaton> sam_;
  Param param_;
  // ... (생략)
};

sam_는 std::unique_ptr로 선언되어 SuffixAutomaton 객체의 생명 주기를 관리하며, 필요에 따라 동적으로 할당 및 해제될 수 있도록 합니다.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp

이 파일에는 외부 코퍼스 로딩 및 SAM 구축을 위한 새로운 메서드와 batchMatch 로직 변경이 포함됩니다.

1. 외부 코퍼스 로딩 및 SAM 관리 메서드 추가

startExternalCorpusLoad, appendExternalCorpusTokens, finishExternalCorpusLoad, clearExternalCorpus 메서드가 추가되어 외부 코퍼스를 스트리밍 방식으로 처리하고 SAM을 구축할 수 있게 합니다. 특히 appendExternalCorpusTokens는 Pybind가 큰 코퍼스를 메모리에 두 번 로드하는 것을 방지하기 위해 청크 단위로 토큰을 추가하는 역할을 합니다.

After:

void Ngram::startExternalCorpusLoad() {
  std::unique_lock<std::mutex> lock(mutex_);
  sam_ = std::make_unique<SuffixAutomaton>();
}

void Ngram::appendExternalCorpusTokens(const std::vector<int32_t>& tokens) {
  std::unique_lock<std::mutex> lock(mutex_);
  sam_->appendTokens(tokens);
}

void Ngram::finishExternalCorpusLoad() {
  std::unique_lock<std::mutex> lock(mutex_);
  sam_->finalize();
  if (sam_->empty()) {
    sam_.reset();
  }
}

void Ngram::clearExternalCorpus() {
  std::unique_lock<std::mutex> lock(mutex_);
  sam_.reset();
}

startExternalCorpusLoad에서 SuffixAutomaton 객체를 생성하고, appendExternalCorpusTokens로 토큰 청크를 추가한 후, finishExternalCorpusLoad에서 SAM을 최종화(finalize)합니다. 이 과정은 mutex_로 보호되어 스레드 안전성을 확보합니다.

2. batchMatch 로직 변경: Trie와 SAM 후보 병합

batchMatch 함수는 이제 Trie와 SAM에서 생성된 드래프트 토큰 후보를 병합하여 사용합니다. param_.external_sam_budget을 통해 SAM이 생성할 드래프트 토큰의 예산을 설정하고, 나머지 예산은 Trie가 사용하도록 분배합니다.

Before:

  using BuildFn = Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&, MatchState&, size_t) const;
  BuildFn build_fn;
  if (param_.match_type == "BFS") {
    build_fn = &Trie::buildRecency;
  } else if (param_.match_type == "PROB") {
    build_fn = &Trie::buildFrequency;
  } else {
    throw std::runtime_error("Unknown match_type: '" + param_.match_type + "'. Must be 'BFS' or 'PROB'.");
  }

  // ... (생략)

    auto draft_token_num = param_.get_draft_token_num(tokens.size());
    auto res = (trie_.get()->*build_fn)(
        suffix.data(), suffix.size(), suffix.back(), draft_token_num, param_, state, total_lens[i]);
    merged.token.insert(merged.token.end(), res.token.begin(), res.token.end());
    merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end());

After:

  using TrieResultBuildFn =
      Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&, MatchState&, size_t) const;
  using SamResultBuildFn = Result (SuffixAutomaton::*)(const int32_t*, size_t, int32_t, size_t, const Param&) const;
  TrieResultBuildFn trie_result_build_fn;
  SamResultBuildFn sam_result_build_fn;
  if (param_.match_type == "BFS") {
    trie_result_build_fn = &Trie::buildRecency;
    sam_result_build_fn = &SuffixAutomaton::buildRecency;
  } else if (param_.match_type == "PROB") {
    trie_result_build_fn = &Trie::buildFrequency;
    sam_result_build_fn = &SuffixAutomaton::buildFrequency;
  } else {
    throw std::runtime_error("Unknown match_type: '" + param_.match_type + "'. Must be 'BFS' or 'PROB'.");
  }

  // ... (생략)

    const auto total_draft_token_num = param_.get_draft_token_num(tokens.size());
    const auto sam_budget =
        sam_ && !sam_->empty() ? std::min(param_.external_sam_budget, total_draft_token_num) : size_t{0};
    const auto trie_budget = total_draft_token_num - sam_budget;

    if (sam_budget == 0) {
      auto res = (trie_.get()->*trie_result_build_fn)(
          suffix.data(), suffix.size(), suffix.back(), total_draft_token_num, param_, state, total_lens[i]);
      merged.token.insert(merged.token.end(), res.token.begin(), res.token.end());
      merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end());
      continue;
    }

    auto trie_res = (trie_.get()->*trie_result_build_fn)(
        suffix.data(), suffix.size(), suffix.back(), trie_budget, param_, state, total_lens[i]);
    auto sam_res = (sam_.get()->*sam_result_build_fn)(suffix.data(), suffix.size(), suffix.back(), sam_budget, param_);
    auto res = combineRootResults_(suffix.back(), static_cast<int>(total_draft_token_num + 1), trie_res, sam_res);
    merged.token.insert(merged.token.end(), res.token.begin(), res.token.end());
    merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end());

이제 batchMatch는 trie_result_build_fn과 sam_result_build_fn 두 가지 함수 포인터를 사용하여 Trie와 SAM 각각에서 결과를 생성합니다. sam_budget이 0이 아니면, Trie와 SAM에서 각각 trie_budget과 sam_budget만큼의 드래프트 토큰을 생성하고, combineRootResults_ 함수를 통해 이들을 하나의 추측 트리에 병합합니다. 이 병합 과정은 오버헤드가 낮도록 설계되었습니다.

python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp

Python 바인딩 파일에는 NgramCorpusObj 생성자에 새로운 파라미터가 추가되었고, SAM 관련 메서드들이 Python에서 호출될 수 있도록 FFI(Foreign Function Interface)가 추가되었습니다.

Before:

struct NgramCorpusObj : public tvm::ffi::Object {
  // ... (생략)
  NgramCorpusObj(
      int64_t capacity,
      int64_t min_bfs_breadth,
      int64_t max_bfs_breadth,
      int64_t draft_token_num,
      int64_t match_type) {
    // ... (생략)
  }
  // ... (생략)
};

void register_ngram_corpus() {
  namespace refl = tvm::ffi::reflection;
  refl::ObjectDef<NgramCorpusObj>()
      .def(refl::init<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t>(), "__init__")
      .def("async_insert", &NgramCorpusObj::async_insert)
      .def("batch_match", &NgramCorpusObj::batch_match)
      .def("batch_match_stateful", &NgramCorpusObj::batch_match_stateful)
      .def("erase_match_state", &NgramCorpusObj::erase_match_state)
      .def("synchronize", &NgramCorpusObj::synchronize)
      .def("reset", &NgramCorpusObj::reset);
}

After:

struct NgramCorpusObj : public tvm::ffi::Object {
  // ... (생략)
  NgramCorpusObj(
      int64_t capacity,
      int64_t min_bfs_breadth,
      int64_t max_bfs_breadth,
      int64_t draft_token_num,
      int64_t match_type,
      int64_t external_sam_budget,
      int64_t external_corpus_max_tokens) {
    // ... (생략)
    param.external_sam_budget = static_cast<size_t>(external_sam_budget);
    param.external_corpus_max_tokens = static_cast<size_t>(external_corpus_max_tokens);
    // ... (생략)
  }
  // ... (생략)

  void start_external_corpus_load() {
    ngram_->startExternalCorpusLoad();
  }

  void append_external_corpus_tokens(const tvm::ffi::TensorView tokens_tv) {
    auto* data = static_cast<const int32_t*>(tokens_tv.data_ptr());
    int64_t n = tokens_tv.size(0);
    std::vector<int32_t> tokens(data, data + n);
    ngram_->appendExternalCorpusTokens(tokens);
  }

  void finish_external_corpus_load() {
    ngram_->finishExternalCorpusLoad();
  }

  void clear_external_corpus() {
    ngram_->clearExternalCorpus();
  }
  // ... (생략)
};

void register_ngram_corpus() {
  namespace refl = tvm::ffi::reflection;
  refl::ObjectDef<NgramCorpusObj>()
      .def(refl::init<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t, int64_t>(), "__init__")
      .def("async_insert", &NgramCorpusObj::async_insert)
      .def("batch_match", &NgramCorpusObj::batch_match)
      .def("batch_match_stateful", &NgramCorpusObj::batch_match_stateful)
      .def("erase_match_state", &NgramCorpusObj::erase_match_state)
      .def("start_external_corpus_load", &NgramCorpusObj::start_external_corpus_load)
      .def("append_external_corpus_tokens", &NgramCorpusObj::append_external_corpus_tokens)
      .def("finish_external_corpus_load", &NgramCorpusObj::finish_external_corpus_load)
      .def("clear_external_corpus", &NgramCorpusObj::clear_external_corpus)
      .def("synchronize", &NgramCorpusObj::synchronize)
      .def("reset", &NgramCorpusObj::reset);
}

NgramCorpusObj의 생성자는 이제 external_sam_budget과 external_corpus_max_tokens를 인자로 받아 Param 객체에 설정합니다. 또한, start_external_corpus_load, append_external_corpus_tokens, finish_external_corpus_load, clear_external_corpus와 같은 SAM 관리 함수들이 Python에서 직접 호출될 수 있도록 FFI 바인딩이 추가되었습니다. 특히 append_external_corpus_tokens는 tvm::ffi::TensorView를 사용하여 토큰 데이터를 효율적으로 전달받습니다.

python/sglang/jit_kernel/csrc/ngram_corpus/param.h

Param 구조체에 external_sam_budget과 external_corpus_max_tokens 필드가 추가되어, 외부 코퍼스 기반 SAM의 동작을 제어할 수 있게 되었습니다.

Before:

struct Param {
  // ... (생략)
  size_t draft_token_num;
  std::string match_type;
  // ... (생략)
};

After:

struct Param {
  // ... (생략)
  size_t draft_token_num;
  size_t external_sam_budget = 0;
  size_t external_corpus_max_tokens = 10000000;
  std::string match_type;
  // ... (생략)
};

external_sam_budget은 SAM이 생성할 드래프트 토큰의 최대 개수를, external_corpus_max_tokens는 SAM 구축에 사용할 외부 코퍼스의 최대 토큰 수를 지정합니다. 이는 CPU OOM(Out Of Memory)을 방지하기 위한 중요한 제어 메커니즘입니다.

python/sglang/jit_kernel/csrc/ngram_corpus/result.cpp

combineRootResults_ 함수가 추가되어 Trie와 SAM에서 생성된 Result 객체를 효율적으로 병합합니다. 이 함수는 두 결과의 토큰과 마스크 정보를 결합하여 최종 추측 트리를 생성합니다.

After:

std::vector<std::vector<int32_t>> extractLeafPaths_(const Result& result) {
  // ... (생략)
}

Result buildResultFromLeafPaths_(int last_token, int draft_token_num, const std::vector<std::vector<int32_t>>& paths) {
  // ... (생략)
}

Result combineRootResults_(int last_token, int draft_token_num, const Result& res1, const Result& res2) {
  auto paths1 = extractLeafPaths_(res1);
  auto paths2 = extractLeafPaths_(res2);
  paths1.insert(paths1.end(), std::make_move_iterator(paths2.begin()), std::make_move_iterator(paths2.end()));
  return buildResultFromLeafPaths_(last_token, draft_token_num, paths1);
}

extractLeafPaths_는 주어진 Result 객체에서 리프 노드까지의 경로를 추출하고, buildResultFromLeafPaths_는 추출된 경로들을 기반으로 새로운 Result 객체를 구축합니다. combineRootResults_는 이 두 헬퍼 함수를 사용하여 두 Result를 병합합니다. 이 방식은 두 개의 개별적인 예측 트리를 효율적으로 하나의 통합된 추측 트리로 변환합니다.

왜 이게 좋은 최적화/개선인가?

이 PR은 SGLang의 Ngram 추측 디코딩 기능을 크게 향상시키는 여러 가지 중요한 최적화 및 개선 사항을 포함합니다.

  1. 예측 커버리지 확장: 기존 Ngram Trie는 주로 현재 컨텍스트 내의 Ngram 패턴에 의존했습니다. Suffix Automaton은 외부 코퍼스를 통해 더 넓은 범위의 텍스트 패턴을 학습하고 예측에 활용할 수 있게 합니다. 이는 모델이 생성해야 할 내용이 현재 컨텍스트를 넘어서는 경우, 특히 특정 도메인 지식이 필요한 경우에 드래프트 토큰의 정확도를 크게 높일 수 있습니다.

  2. CPU OOM 방지 및 효율적인 코퍼스 로딩: 대규모 외부 코퍼스를 로드할 때 발생할 수 있는 CPU 메모리 부족(OOM) 문제를 speculative_ngram_external_corpus_max_tokens 파라미터로 제어합니다. 또한, appendExternalCorpusTokens를 통해 코퍼스를 청크 단위로 스트리밍하고 토큰화하여 Pybind가 전체 코퍼스를 메모리에 두 번 로드하는 비효율성을 방지합니다. 이는 대용량 데이터 처리 시 메모리 사용량을 최적화하는 데 필수적입니다.

  3. 유연한 드래프트 토큰 예산 관리: external_sam_budget 파라미터를 통해 Trie와 SAM 간에 드래프트 토큰 생성 예산을 동적으로 분배할 수 있습니다. 이는 특정 시나리오에서 Trie 또는 SAM 중 어느 것이 더 효과적인지 실험하고 최적의 조합을 찾을 수 있는 유연성을 제공합니다.

  4. 낮은 오버헤드의 후보 병합: Trie와 SAM에서 생성된 드래프트 토큰 후보들을 combineRootResults_ 함수를 통해 효율적으로 병합합니다. PR 설명에 따르면 이 병합 과정은 "low overhead"를 목표로 설계되었으며, 이는 두 예측기의 장점을 활용하면서도 추가적인 성능 저하를 최소화합니다.

벤치마킹 및 성능 수치

PR에 포함된 벤치마킹 결과는 외부 코퍼스 로딩 및 SAM 구축에 필요한 시간과 메모리 오버헤드를 명확히 보여줍니다.

================================================================================
External Corpus Load + SAM Construction Overhead
================================================================================
  corpus_tokens |  load_time_mean(s) |  load_time_std | mem_delta_mean(MB) |  mem_delta_std
-------------------------------------------------------------------------------------------
   0 (baseline) |             0.0828 |         0.0046 |             177.85 |           0.04
      1,000,000 |             2.6590 |         0.0418 |             398.58 |           0.35
      5,000,000 |            14.7094 |         0.1811 |            1845.80 |           0.19
     10,000,000 |            31.1640 |         0.1667 |            3846.36 |           0.40

이 수치는 외부 코퍼스 크기가 증가함에 따라 SAM 구축 시간이 선형적으로 증가하며, 메모리 사용량도 2 * external_corpus_tokens_count에 비례하여 증가함을 보여줍니다. 예를 들어, 1천만 토큰 코퍼스의 경우 약 31초의 로딩 시간과 3.8GB의 추가 메모리가 필요합니다. 이는 서버 시작 시간에 영향을 미치지만, 추측 디코딩의 정확도 향상을 통해 얻을 수 있는 추론 속도 개선 효과와 트레이드오프 관계에 있습니다. PR 작성자는 "We can look into optimizing this in later PRs"라고 언급하며, 초기 로딩 오버헤드에 대한 추가 최적화 가능성을 열어두고 있습니다.

일반적 교훈

이 PR은 다음과 같은 일반적인 소프트웨어 엔지니어링 및 시스템 최적화 교훈을 제공합니다.

  • 하이브리드 예측 전략: 단일 예측기에 의존하기보다, 여러 예측기(여기서는 Trie와 SAM)의 장점을 결합하여 더 강력하고 유연한 시스템을 구축할 수 있습니다.
  • 리소스 관리의 중요성: 대규모 데이터를 다룰 때는 메모리(CPU OOM)와 로딩 시간 같은 리소스 제약 조건을 명확히 인지하고, 이를 관리하기 위한 메커니즘(예: max_tokens 제한, 스트리밍 로딩)을 설계해야 합니다.
  • 성능과 기능의 트레이드오프: 새로운 기능을 추가하거나 성능을 개선할 때, 항상 그에 따른 비용(예: 초기 로딩 시간, 메모리 사용량)을 고려하고, 이들 간의 적절한 트레이드오프 지점을 찾아야 합니다.
  • 모듈화 및 확장성: Ngram 클래스에 SuffixAutomaton을 unique_ptr로 추가하고, FFI를 통해 Python에서 제어할 수 있도록 설계함으로써, 시스템의 모듈성과 확장성을 높였습니다. 이는 향후 추가적인 예측기나 최적화 기법을 통합하기 용이하게 만듭니다.

결론

이 PR은 SGLang의 Ngram 추측 디코딩에 외부 코퍼스 기반 Suffix Automaton을 성공적으로 통합하여, 예측기의 커버리지를 확장하고 더 정확한 드래프트 토큰 생성을 가능하게 했습니다. 이는 LLM의 추론 속도 향상에 기여할 중요한 개선 사항입니다. 초기 로딩 오버헤드가 존재하지만, 이는 향후 최적화를 통해 개선될 여지가 있으며, 전반적인 추론 성능 향상이라는 큰 그림에서 볼 때 가치 있는 투자입니다. 이러한 개선은 SGLang이 다양한 사용 사례와 도메인에 걸쳐 LLM 추론을 더욱 효율적으로 수행할 수 있도록 하는 중요한 발판이 될 것입니다.

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

댓글

관련 포스트

PR Analysis 의 다른글