chore: import upstream snapshot with attribution

python/sglang/srt/speculative/cpp_ngram/.clang-format

@@ -0,0 +1,15 @@
+BasedOnStyle: Google
+IndentWidth: 2
+ColumnLimit: 120
+AllowShortFunctionsOnASingleLine: Empty
+DerivePointerAlignment: false
+PointerAlignment: Left
+NamespaceIndentation: None
+SortIncludes: true
+AllowShortLoopsOnASingleLine: false
+BinPackParameters: false                 # Prevents packing parameters in declarations
+BinPackArguments: false                  # Prevents packing arguments in function calls
+AlignAfterOpenBracket: AlwaysBreak       # Forces a break after the opening parenthesis
+AlignOperands: Align                     # Aligns arguments vertically
+PenaltyBreakBeforeFirstCallParameter: 1  # Encourages breaking before the first argument
+PenaltyReturnTypeOnItsOwnLine: 100       # Keeps return type with function name

python/sglang/srt/speculative/cpp_ngram/external_corpus.py

@@ -0,0 +1,62 @@
+import json
+from collections.abc import Iterator
+from pathlib import Path
+
+# Must match SuffixAutomaton::kSeparatorToken in suffix_automaton.h.
+SEPARATOR_TOKEN = -(2**31)
+
+# Default chunk size for streaming tokenized documents into the SAM.
+DEFAULT_CHUNK_SIZE = 4096
+
+
+def iter_external_corpus_chunks(
+    path: str, tokenizer, max_tokens: int, chunk_size: int = DEFAULT_CHUNK_SIZE
+) -> Iterator[list[int]]:
+    """Chunk documents and yield fixed-size token chunks from a JSONL corpus file."""
+    corpus_path = Path(path)
+    if not corpus_path.is_file():
+        raise ValueError(f"External ngram corpus path does not exist: {path}")
+    if tokenizer is None:
+        raise ValueError("A tokenizer is required to load an external ngram corpus.")
+    if max_tokens <= 0:
+        raise ValueError("External ngram corpus max tokens must be positive.")
+
+    total_tokens = 0
+    has_previous_doc = False
+    with corpus_path.open("r", encoding="utf-8") as f:
+        for line_no, line in enumerate(f, start=1):
+            if not line.strip():
+                continue
+
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError as e:
+                raise ValueError(
+                    f"Invalid JSON in external ngram corpus at line {line_no}: {e.msg}"
+                ) from e
+
+            if not isinstance(record, str):
+                raise ValueError(
+                    "Invalid external ngram corpus record at line "
+                    f"{line_no}: expected a JSON string."
+                )
+
+            token_ids = list(tokenizer.encode(record, add_special_tokens=False))
+            if not token_ids:
+                continue
+
+            separator_cost = 1 if has_previous_doc else 0
+            next_total_tokens = total_tokens + separator_cost + len(token_ids)
+            if next_total_tokens > max_tokens:
+                raise ValueError(
+                    "External ngram corpus exceeds the configured token limit "
+                    f"({max_tokens}) at line {line_no} after loading "
+                    f"{total_tokens} tokens."
+                )
+            total_tokens = next_total_tokens
+
+            if has_previous_doc:
+                token_ids = [SEPARATOR_TOKEN] + token_ids
+            for i in range(0, len(token_ids), chunk_size):
+                yield token_ids[i : i + chunk_size]
+            has_previous_doc = True

python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py

@@ -0,0 +1,200 @@
+# -*- coding: utf-8 -*-
+
+import logging
+from collections.abc import Iterable, Sequence
+from typing import Dict, List, Tuple
+
+import numpy as np
+
+from sglang.jit_kernel.ngram_corpus import get_ngram_corpus_cls
+
+logger = logging.getLogger(__name__)
+
+
+class NgramCorpus:
+    def __init__(
+        self,
+        max_trie_depth=18,
+        min_bfs_breadth=1,
+        max_bfs_breadth=8,
+        draft_token_num=8,
+        match_type="BFS",
+        capacity=1000000,
+        external_sam_budget=0,
+        external_corpus_max_tokens=10000000,
+    ) -> None:
+        cls = get_ngram_corpus_cls()
+        self._obj = cls(
+            capacity=capacity,
+            max_trie_depth=max_trie_depth,
+            min_bfs_breadth=min_bfs_breadth,
+            max_bfs_breadth=max_bfs_breadth,
+            draft_token_num=draft_token_num,
+            match_type=match_type,
+            external_sam_budget=external_sam_budget,
+            external_corpus_max_tokens=external_corpus_max_tokens,
+        )
+        self.draft_token_num = draft_token_num
+        self.external_corpus_max_tokens = external_corpus_max_tokens
+        self._req_id_to_state_id: Dict[str, int] = {}
+        self._next_state_id: int = 0
+        self._corpus_token_counts: Dict[str, int] = {}
+        self._total_loaded_tokens: int = 0
+
+    def _get_state_id(self, req_id: str) -> int:
+        sid = self._req_id_to_state_id.get(req_id)
+        if sid is None:
+            sid = self._next_state_id
+            self._next_state_id += 1
+            self._req_id_to_state_id[req_id] = sid
+        return sid
+
+    def batch_put(self, batch_tokens: List[List[int]]):
+        self._obj.insert(batch_tokens)
+
+    def synchronize(self):
+        self._obj.synchronize()  # type: ignore
+
+    @property
+    def remaining_token_budget(self) -> int:
+        return self.external_corpus_max_tokens - self._total_loaded_tokens
+
+    def load_external_corpus_named(
+        self, corpus_id: str, chunks: Iterable[Sequence[int]]
+    ) -> int:
+        if corpus_id in self._corpus_token_counts:
+            raise ValueError(
+                f"External corpus '{corpus_id}' already exists. Remove it before "
+                f"adding a new corpus with the same id."
+            )
+        # Note(kpham-sgl): remaining_token_budget is stale (e.g if there are removes
+        # during the load), which makes the budget more conservative than it should be.
+        # This is acceptable because otherwise load_external_corpus_named would need to check the budget after each chunk,
+        # which would be inefficient.
+        _, loaded_token_count = self._obj.load_external_corpus_named(
+            corpus_id, chunks, self.remaining_token_budget
+        )
+        return loaded_token_count
+
+    # Commit corpus bookkeeping after successful load. Call only at background thread join.
+    # (or after synchronous load_external_corpus_named returns)
+    def commit_external_corpus_load(
+        self, corpus_id: str, loaded_token_count: int
+    ) -> None:
+        self._corpus_token_counts[corpus_id] = loaded_token_count
+        self._total_loaded_tokens += loaded_token_count
+
+    def remove_external_corpus(self, corpus_id: str) -> None:
+        self._obj.remove_corpus(corpus_id)
+        old_count = self._corpus_token_counts.pop(corpus_id, 0)
+        self._total_loaded_tokens -= old_count
+
+    def list_external_corpora(self) -> Dict[str, int]:
+        return self._obj.list_corpora()
+
+    def reset(self):
+        self._obj.reset()  # type: ignore
+        self._req_id_to_state_id.clear()
+        self._next_state_id = 0
+
+    def batch_get(
+        self,
+        req_ids: List[str],
+        batch_tokens: List[List[int]],
+        total_lens: List[int],
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        state_ids = [self._get_state_id(rid) for rid in req_ids]
+        return self._obj.match_stateful(state_ids, batch_tokens, total_lens)
+
+    def erase_match_state(self, req_ids: List[str]):
+        state_ids = []
+        for rid in req_ids:
+            sid = self._req_id_to_state_id.pop(rid, None)
+            if sid is not None:
+                state_ids.append(sid)
+        if state_ids:
+            self._obj.erase_states(state_ids)
+
+    def leaf_paths_from_mask(
+        self, tokens: List[int], tree_mask: List[List[int]]
+    ) -> List[List[int]]:
+        """
+        Find all leaf paths according to the binary tree_mask (i.e., paths that are not prefixes of any other path).
+
+        Args:
+            mask   : List[List[int]]   # nxn binary matrix
+            tokens : List[int]         # token list corresponding to columns
+
+        Returns:
+            List[List[int]]            # token lists of only the leaf paths, preserving their order of appearance
+        """
+
+        row_sets = [
+            (i, {idx for idx, v in enumerate(row) if v == 1})
+            for i, row in enumerate(tree_mask)
+        ]
+        leaf_sets = []
+        leaf_rows = []
+
+        for i, cur_set in reversed(row_sets):
+            if any(cur_set <= kept for kept in leaf_sets):
+                continue
+            leaf_sets.append(cur_set)
+            leaf_rows.append(i)
+
+        leaf_rows.reverse()
+        result = []
+        for r in leaf_rows:
+            path = [tokens[col] for col in range(len(tokens)) if tree_mask[r][col] == 1]
+            result.append(path)
+
+        return result
+
+    def debug_result(
+        self, decoding_ids: np.ndarray, decoding_masks: np.ndarray, tokenizer=None
+    ):
+        decoding_ids = decoding_ids.reshape(-1, self.draft_token_num)
+        decoding_masks = decoding_masks.reshape(
+            -1, self.draft_token_num, self.draft_token_num
+        )
+        logger.info(f"\n{decoding_ids=}\n{decoding_masks=}")
+        for i in range(decoding_ids.shape[0]):
+            leaf_paths = self.leaf_paths_from_mask(
+                decoding_ids[i].tolist(), decoding_masks[i].tolist()
+            )
+            if tokenizer is None:
+                logger.info(f"draft path {i}: {leaf_paths}")
+            else:
+                logger.info(f"result {i}:")
+                for leaf_path in leaf_paths:
+                    logger.info(
+                        f"draft path {i}: {leaf_path} -> {tokenizer.decode(leaf_path, ensure_ascii=False)}"
+                    )
+
+
+# main function
+if __name__ == "__main__":
+    format = f"%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s"
+    logging.basicConfig(
+        level=logging.DEBUG,
+        format=format,
+        datefmt="%Y-%m-%d %H:%M:%S",
+        force=True,
+    )
+
+    token_ids = [
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+        [1, 2, 3, 44, 55, 66, 77, 88, 99, 100],
+    ]
+    corpus = NgramCorpus(max_trie_depth=12, draft_token_num=8)
+    corpus.batch_put(token_ids)
+
+    corpus.synchronize()
+    queries = [[1, 2, 3], [3, 44], [3, 6, 999]]
+    decoding_ids, decoding_masks = corpus.batch_get(
+        req_ids=[f"query-{i}" for i in range(len(queries))],
+        batch_tokens=queries,
+        total_lens=[len(q) for q in queries],
+    )
+
+    corpus.debug_result(decoding_ids, decoding_masks)