[Spec][Ngram] Return token counts in list_external_corpora API (#22471)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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

@@ -115,14 +115,14 @@ void Ngram::clearExternalCorpus() {
   staging_sam_.reset();
 }
 
-std::vector<std::string> Ngram::listExternalCorpora() const {
+std::vector<std::pair<std::string, int64_t>> Ngram::listExternalCorpora() const {
   std::unique_lock<std::mutex> lock(mutex_);
-  std::vector<std::string> ids;
+  std::vector<std::pair<std::string, int64_t>> entries;
-  ids.reserve(sams_.size());
+  entries.reserve(sams_.size());
-  for (const auto& [id, _] : sams_) {
+  for (const auto& [id, sam] : sams_) {
-    ids.push_back(id);
+    entries.emplace_back(id, sam->tokenCount());
   }
-  return ids;
+  return entries;
 }
 
 void Ngram::insertWorker() {

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

@@ -59,7 +59,7 @@ class Ngram {
 
   void clearExternalCorpus();
 
-  std::vector<std::string> listExternalCorpora() const;
+  std::vector<std::pair<std::string, int64_t>> listExternalCorpora() const;
 
   Result batchMatch(const std::vector<std::vector<int32_t>>& tokens);
 

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

@@ -129,11 +129,11 @@ struct NgramCorpusObj : public tvm::ffi::Object {
   }
 
   std::string list_external_corpora() {
-    auto ids = ngram_->listExternalCorpora();
+    auto entries = ngram_->listExternalCorpora();
     std::string result;
-    for (size_t i = 0; i < ids.size(); ++i) {
+    for (size_t i = 0; i < entries.size(); ++i) {
       if (i > 0) result += "\n";
-      result += ids[i];
+      result += entries[i].first + "\t" + std::to_string(entries[i].second);
     }
     return result;
   }

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

@@ -39,6 +39,10 @@ class SuffixAutomaton {
     return !loaded_;
   }
 
+  int64_t tokenCount() const {
+    return pos_;
+  }
+
   Result buildRecency(
       const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const;
 

python/sglang/jit_kernel/ngram_corpus.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Sequence
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 
 import numpy as np
 import torch
@@ -144,10 +144,14 @@ def get_ngram_corpus_cls():
         def remove_corpus(self, corpus_id: str) -> None:
             self.remove_external_corpus(corpus_id)  # type: ignore
 
-        def list_corpora(self) -> List[str]:
+        def list_corpora(self) -> Dict[str, int]:
             result = self.list_external_corpora()  # type: ignore
             if not result:
-                return []
+                return {}
-            return result.split("\n")
+            out: Dict[str, int] = {}
+            for line in result.split("\n"):
+                corpus_id, token_count = line.split("\t", 1)
+                out[corpus_id] = int(token_count)
+            return out
 
     return NgramCorpusFFI

python/sglang/srt/entrypoints/http_server.py

@@ -814,7 +814,7 @@ async def list_external_corpora():
     return ORJSONResponse(
         {
             "success": result.success,
-            "corpus_ids": result.corpus_ids,
+            "corpus_token_counts": result.corpus_token_counts,
             "message": result.message,
         },
         status_code=200 if result.success else HTTPStatus.BAD_REQUEST,

python/sglang/srt/managers/io_struct.py

@@ -1232,7 +1232,7 @@ class ListExternalCorporaReqInput(BaseReq):
 @dataclass
 class ListExternalCorporaReqOutput(BaseReq):
     success: bool
-    corpus_ids: List[str] = field(default_factory=list)
+    corpus_token_counts: Dict[str, int] = field(default_factory=dict)
     message: str = ""
 
 

python/sglang/srt/managers/tokenizer_communicator_mixin.py

@@ -473,10 +473,12 @@ class TokenizerCommunicatorMixin:
             ListExternalCorporaReqInput()
         )
         all_success, all_message = _Communicator.merge_results(results)
-        # Merge corpus IDs from all DP ranks (each rank loads the same set).
+        # Merge corpus token counts from all DP ranks (each rank loads the same set).
-        corpus_ids = results[0].corpus_ids if all_success else []
+        corpus_token_counts = results[0].corpus_token_counts if all_success else {}
         return ListExternalCorporaReqOutput(
-            success=all_success, corpus_ids=corpus_ids, message=all_message
+            success=all_success,
+            corpus_token_counts=corpus_token_counts,
+            message=all_message,
         )
 
     async def flush_cache(

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

@@ -90,7 +90,7 @@ class NgramCorpus:
         old_count = self._corpus_token_counts.pop(corpus_id, 0)
         self._total_loaded_tokens -= old_count
 
-    def list_external_corpora(self) -> List[str]:
+    def list_external_corpora(self) -> Dict[str, int]:
         return self._obj.list_corpora()
 
     def reset(self):

python/sglang/srt/speculative/external_corpus_manager.py

@@ -101,7 +101,10 @@ class ExternalCorpusManager:
         self, recv_req: ListExternalCorporaReqInput
     ) -> ListExternalCorporaReqOutput:
         try:
-            ids = self._worker.list_external_corpora()
+            token_counts = self._worker.list_external_corpora()
-            return ListExternalCorporaReqOutput(success=True, corpus_ids=ids)
+            return ListExternalCorporaReqOutput(
+                success=True,
+                corpus_token_counts=token_counts,
+            )
         except Exception as e:
             return ListExternalCorporaReqOutput(success=False, message=str(e))

python/sglang/srt/speculative/ngram_worker.py

@@ -90,7 +90,7 @@ class NGRAMWorker:
     def remove_external_corpus(self, corpus_id: str) -> None:
         self.ngram_corpus.remove_external_corpus(corpus_id)
 
-    def list_external_corpora(self) -> list[str]:
+    def list_external_corpora(self) -> dict[str, int]:
         return self.ngram_corpus.list_external_corpora()
 
     def _efficient_concat_last_n(self, seq1: List[int], seq2: List[int], n: int):

test/registered/unit/spec/test_ngram_corpus.py

@@ -925,8 +925,10 @@ class TestNgramCorpusMultiSam(CustomTestCase):
             "b", [[10, 20, 30, 40, 50]]
         )
         corpus.commit_external_corpus_load("b", loaded_token_count)
-        ids = corpus.list_external_corpora()
+        token_counts = corpus.list_external_corpora()
-        self.assertEqual(sorted(ids), ["a", "b"])
+        self.assertEqual(sorted(token_counts.keys()), ["a", "b"])
+        self.assertEqual(token_counts["a"], 5)
+        self.assertEqual(token_counts["b"], 5)
 
     def test_remove(self):
         corpus = _make_corpus("BFS", draft_token_num=4, external_sam_budget=3)
@@ -937,12 +939,12 @@ class TestNgramCorpusMultiSam(CustomTestCase):
         )
         corpus.commit_external_corpus_load("b", loaded_token_count)
         corpus.remove_external_corpus("a")
-        self.assertEqual(corpus.list_external_corpora(), ["b"])
+        self.assertEqual(list(corpus.list_external_corpora().keys()), ["b"])
 
     def test_remove_nonexistent_is_noop(self):
         corpus = _make_corpus("BFS", draft_token_num=4, external_sam_budget=3)
         corpus.remove_external_corpus("nonexistent")
-        self.assertEqual(corpus.list_external_corpora(), [])
+        self.assertEqual(corpus.list_external_corpora(), {})
 
     def test_multi_sam_candidates(self):
         corpus = _make_corpus("BFS", draft_token_num=6, external_sam_budget=4)
@@ -983,8 +985,8 @@ class TestNgramCorpusMultiSam(CustomTestCase):
             external_sam_budget=3,
             external_corpus_documents=[[1, 2, 3, 4, 5]],
         )
-        ids = corpus.list_external_corpora()
+        token_counts = corpus.list_external_corpora()
-        self.assertIn("test_corpus", ids)
+        self.assertIn("test_corpus", token_counts)
 
     def test_remove_frees_token_budget(self):
         """Removing a corpus should free its tokens from the total budget."""
@@ -1008,7 +1010,7 @@ class TestNgramCorpusMultiSam(CustomTestCase):
         # Now there's room for a new corpus.
         loaded_token_count = corpus.load_external_corpus_named("c", [[100, 200, 300]])
         corpus.commit_external_corpus_load("c", loaded_token_count)
-        self.assertEqual(sorted(corpus.list_external_corpora()), ["b", "c"])
+        self.assertEqual(sorted(corpus.list_external_corpora().keys()), ["b", "c"])
 
     def test_duplicate_corpus_id_is_rejected(self):
         """Adding a duplicate corpus_id should fail without replacing the original corpus."""
@@ -1024,7 +1026,7 @@ class TestNgramCorpusMultiSam(CustomTestCase):
             corpus.load_external_corpus_named("a", [[10, 20, 30]])
 
         self.assertEqual(corpus.remaining_token_budget, 5)
-        self.assertEqual(corpus.list_external_corpora(), ["a"])
+        self.assertEqual(list(corpus.list_external_corpora().keys()), ["a"])
 
         # The original corpus must still be usable for matching.
         ids, masks = _batch_get(corpus, [[1, 2, 3]])
@@ -1051,7 +1053,7 @@ class TestNgramCorpusMultiSam(CustomTestCase):
         with self.assertRaises(ValueError):
             corpus.load_external_corpus_named("b", [[10, 20, 30, 40, 50, 60]])
 
-        self.assertEqual(corpus.list_external_corpora(), ["a"])
+        self.assertEqual(list(corpus.list_external_corpora().keys()), ["a"])
         self.assertEqual(corpus.remaining_token_budget, 5)
 
         # "a" must still be usable for matching.
@@ -1106,7 +1108,7 @@ class TestMultiSamHttpMock(CustomTestCase):
         )
         tm.list_external_corpora = AsyncMock(
             return_value=ListExternalCorporaReqOutput(
-                success=True, corpus_ids=["a", "b"]
+                success=True, corpus_token_counts={"a": 100, "b": 200}
             )
         )
         set_global_state(mock_state)
@@ -1152,7 +1154,7 @@ class TestMultiSamHttpMock(CustomTestCase):
         self.assertEqual(resp.status_code, 200)
         data = resp.json()
         self.assertTrue(data["success"])
-        self.assertEqual(sorted(data["corpus_ids"]), ["a", "b"])
+        self.assertEqual(data["corpus_token_counts"], {"a": 100, "b": 200})
 
 
 if __name__ == "__main__":