[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;
-  ids.reserve(sams_.size());
-  for (const auto& [id, _] : sams_) {
-    ids.push_back(id);
+  std::vector<std::pair<std::string, int64_t>> entries;
+  entries.reserve(sams_.size());
+  for (const auto& [id, sam] : sams_) {
+    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 result.split("\n")
+                return {}
+            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