[Spec][Ngram] 6/N: Load an external corpus and construct a Suffix Automaton

[kpham-sgl]

Motivation

Part of Ngram refactoring series #21052
Following #21243
We want to support loading an external corpus and constructing a Suffix Automaton for suffix matching during drafting.

Modifications

• Add optional external-corpus support to NGRAM speculative decoding with --speculative-ngram-external-corpus-path, --speculative-ngram-external-sam-budget, and speculative_ngram_external_corpus_max_tokens
• Stream corpus chunks -> tokenize -> SAM.extend() (and pipeline them) to avoid Pybind materialize the big corpus on memory 2x
• Use speculative_ngram_external_corpus_max_tokens as a proxy to avoid CPU OOM during SAM construction (SAM memory overhead is ~ 2 * external_corpus_tokens_count)
• Use the SAM to generate external-corpus continuations in both Recency and Frequency modes, while reserving a fixed portion of the draft budget for SAM-backed candidates.
• Merge trie and SAM candidates into the same speculative tree with low overhead

Accuracy Tests

Verify with these tests locally

pytest test/registered/unit/server_args/test_server_args.py TestNgramExternalSamArgs
pytest test/registered/spec/utils/test_ngram_corpus.py
pytest test/registered/spec/test_ngram_speculative_decoding.py

Benchmarking and Profiling

This change affects server start time. Here is a quick script to benchmark it against the external corpus size
https://gist.github.com/kpham-sgl/00c547a6fb7ab431ad2a09e42e371bd2

python benchmark_external_corpus_overhead.py 
Trials per config: 3
Corpus sizes (tokens): [1000000, 5000000, 10000000]
Tokenizer: meta-llama/Llama-3.1-8B-Instruct

Running baseline (no external corpus)...
Running benchmark: ~1,000,000 tokens...
Running benchmark: ~5,000,000 tokens...
Running benchmark: ~10,000,000 tokens...

================================================================================
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

We can look into optimizing this in later PRs

[TODO] Run any speculative decoding here with external corpus dataset to see how this help speedup Ngram speculative decoding

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[hnyls2002]

/rerun-test test_ngram_corpus.py test_server_args.py test_ngram_speculative_decoding.py

[hnyls2002]

/tag-and-rerun-ci

[github-actions]

✅ ubuntu-latest (2 tests): View workflow run

cd test/ && python3 registered/unit/spec/test_ngram_corpus.py
cd test/ && python3 registered/unit/server_args/test_server_args.py

✅ 1-gpu-h100 (1 test): View workflow run

cd test/ && python3 registered/spec/test_ngram_speculative_decoding.py