Count structured chat tokens accurately - #441
Conversation
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #441 +/- ##
=======================================
Coverage ? 81.90%
=======================================
Files ? 148
Lines ? 19700
Branches ? 0
=======================================
Hits ? 16136
Misses ? 3564
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| # parts) remain unavailable for ISL reporting. | ||
| if token_ids is not None: | ||
| prompt_data = PromptData(token_ids=tuple(token_ids)) | ||
| elif isinstance(data.get("messages"), list | tuple): |
There was a problem hiding this comment.
I see there's a priority order of message prep here - Should message, prompt, token_id be mutually exclusive? E.g. throwing warning/error if more than one is present in the parquet?
There were frequent confusions about which field is used in the parquet in the loadgen - and would be good to call this priority / mutual exclusiveness out somewhere
| if isinstance(ev_rec.data, PromptData) and ev_rec.data.token_ids is not None: | ||
| self.registry.record(self.metric_name, len(ev_rec.data.token_ids)) | ||
| return | ||
| if isinstance(ev_rec.data, PromptData) and ev_rec.data.messages is not None: |
There was a problem hiding this comment.
Seems missing comment (which inherits from the prepopulated token ids comment above)
| return self.tokenizer.encode(text, add_special_tokens=add_special_tokens) | ||
|
|
||
| def encode_lengths(self, texts: list[str]) -> list[int]: | ||
| return [len(self.encode(text, add_special_tokens=False)) for text in texts] |
There was a problem hiding this comment.
Should we use batch tokenizer for this path?
| if isinstance(ev_rec.data, PromptData) and ev_rec.data.token_ids is not None: | ||
| self.registry.record(self.metric_name, len(ev_rec.data.token_ids)) | ||
| return | ||
| if isinstance(ev_rec.data, PromptData) and ev_rec.data.messages is not None: | ||
| if self._queue is not None: | ||
| self._queue.enqueue_prompt( | ||
| (ev_rec.data.messages, ev_rec.data.tools), | ||
| self._make_recorder(ev_rec, pre_change), | ||
| ) | ||
| return |
There was a problem hiding this comment.
Can this be combined so the common if isinstance(ev_rec.data, PromptData) is evaluated one making it more legible.
The messages and self._queue is not None and then be combined.
| if isinstance(ev_rec.data, TextModelOutput) and ( | ||
| ev_rec.data.reasoning or ev_rec.data.tool_calls | ||
| ): |
There was a problem hiding this comment.
Did this change the TPS we are observing as we didn't account for the reasoning output before?
|
|
||
|
|
||
| def _backend_from_tokenizer(tokenizer: Any) -> Any | None: | ||
| """Return the tokenizer's supported length-counting path.""" |
There was a problem hiding this comment.
Is this docstring accurate?
| else: | ||
| prompt_text = data.get("prompt") | ||
| prompt_data = PromptData( | ||
| text=prompt_text if isinstance(prompt_text, str) else None | ||
| ) |
There was a problem hiding this comment.
Recommend hardening this to elif isinstance(data.get("prompt"), str) and adding a warning in the else clause to catch multimodal inputs or any other cases - otherwise these may get silently ignored.
| ) | ||
| self._prefix_len = len(tok.tokenize(prefix)) | ||
| self._prefix_len = len(prefix) | ||
| with_assistant = cast( |
There was a problem hiding this comment.
rename to with_assistant_tokens to help disambiguate text data from tokens.
| try: | ||
| prefix = cast( | ||
| str, | ||
| list[int], | ||
| tok.apply_chat_template( | ||
| [_PREFIX_USER_MSG], tokenize=False, add_generation_prompt=False | ||
| [_PREFIX_USER_MSG], tokenize=True, add_generation_prompt=False | ||
| ), | ||
| ) |
There was a problem hiding this comment.
apply_chat_template returns a dictionary by default if tokenize is set to True. This is only trigerred when tokenize is True. Please update call sites. Also, can you add tests that would catch this behavior - fail with current change and pass with return_dict set to False for tokenize set to True.
arekay-nv
left a comment
arekay-nv
left a comment
There was a problem hiding this comment.
Consolidated review
This change targets accurate structured-chat token counting, but as written the core counts are incorrect on the pinned transformers==5.5.0, and two further issues persist even after the primary fix. All findings below are verified against the pinned tokenizer (live repro on Qwen/Qwen2.5-0.5B-Instruct / Qwen3-8B).
Must fix
tokenize=Truereturns aBatchEncoding;len()counts keys, not tokens (token_metrics.py309/318/465/513). ISL → ~2, structured OSL/TPOT → 0. Fix:return_dict=Falseon everytokenize=Truesite.- A unit-test fake masks the above — it returns a bare list for
tokenize=True, so the suite stays green (test_token_metrics.py).
Persists after the fix above
- Reasoning silently dropped from OSL/TPOT — input templates ignore
reasoning_content(metrics_table.py360/404). - Empty/multimodal
messagesdropped from ISL —messages: []→IndexError; multimodalcontenterrors;_extract_prompt_textis now dead (session.py272).
Should fix
_token_count_promptlacks tool-call normalization, a try/fallback, andchat_template_kwargsparity that_token_count_message/the adapter already have (token_metrics.py499).- Custom-tiktoken backend unreachable —
tiktokenundeclared (token_metrics.py138). - Prompt lane drains serially under the flush lock, never sharded — high-QPS runs can exceed
--drain-timeout(token_metrics.py775).
Checked and fine: msgspec wire round-trip + gc=False safety, _inflight accounting across the three flush phases, and add_generation_prompt=True for ISL.
Suggested tests are inline on the relevant lines; each is written to fail on the current code and pass after the described fix. Real-tokenizer tests use the cached Qwen/Qwen2.5-0.5B-Instruct and can be pytest.skip-guarded when uncached.
| encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] | ||
| list(messages), **kwargs | ||
| ) | ||
| return len(encoded) |
There was a problem hiding this comment.
[critical] len(apply_chat_template(..., tokenize=True)) counts BatchEncoding keys, not tokens.
On the pinned transformers==5.5.0, apply_chat_template(..., tokenize=True) defaults to return_dict=True and returns a BatchEncoding, so len(encoded) is the key count (~2), not the token count. Live repro on Qwen/Qwen2.5-0.5B-Instruct: a 44‑token prompt → ISL 2. The same pattern corrupts the baseline (_prefix_len/_baseline become 2/0) and _token_count_message (line 465), so structured OSL/TPOT collapse to 0 for every reasoning/tool-call sample.
Fix: pass return_dict=False (or index ["input_ids"]) on every tokenize=True site — here plus lines 309, 318, 465.
Suggested test (fails now, passes after the fix):
def test_structured_prompt_isl_is_token_count_not_batchencoding_keys():
bt = BatchTokenizer("Qwen/Qwen2.5-0.5B-Instruct",
live_workers=1, cores_per_worker=0, n_workers=0)
try:
msgs = ({"role": "user", "content": "Write a detailed essay about the fall of Rome."},)
assert bt._token_count_prompt(msgs, None) > 10 # returns 2 today
finally:
bt.close()|
|
||
| def apply_chat_template( | ||
| self, messages, tokenize=False, add_generation_prompt=False | ||
| self, messages, tools=None, tokenize=False, add_generation_prompt=False |
There was a problem hiding this comment.
[high] This fake hides the BatchEncoding defect, so the suite stays green while production is wrong.
apply_chat_template(..., tokenize=True) here returns a bare list, so len() equals the token count — but real transformers returns a BatchEncoding whose len() is the key count (~2). That is why the new chat-template tests pass while the real path returns 2.
Make the fake return a mapping for tokenize=True unless return_dict=False, so the suite fails until the call sites are fixed:
class _BE(dict): # BatchEncoding stand-in: len() == number of keys
pass
def apply_chat_template(self, messages, tools=None, tokenize=False,
add_generation_prompt=False, return_dict=True):
ids = list(range(len(...))) # existing token logic
if tokenize and return_dict:
be = _BE(); be["input_ids"] = ids; be["attention_mask"] = [1]*len(ids)
return be
return ids if tokenize else rendered| def _extract_message(self, ev_rec, row, pre_change): | ||
| if isinstance(ev_rec.data, TextModelOutput) and ev_rec.data.tool_calls: | ||
| if isinstance(ev_rec.data, TextModelOutput) and ( | ||
| ev_rec.data.reasoning or ev_rec.data.tool_calls |
There was a problem hiding this comment.
[high] Reasoning is silently dropped from OSL/TPOT (independent of the return_dict fix).
Routing reasoning-only outputs through the input chat template renders {"role": "assistant", "reasoning_content": ...}, which reference templates (Qwen, DeepSeek-R1, …) ignore — and it does not raise, so the whitespace fallback never fires. Verified on Qwen2.5: _token_count_message("hi", "<long reasoning>", None) == _token_count_message("hi", None, None). Previously the plain-text path counted reasoning via str(TextModelOutput), so this under-reports OSL/TPOT (and inflates TPS) for exactly the reasoning models this change targets. Same applies to TpotTrigger (line 404).
Consider counting reasoning as generated text rather than through the input-side template.
Suggested test:
def test_reasoning_tokens_are_counted():
bt = BatchTokenizer("Qwen/Qwen2.5-0.5B-Instruct",
live_workers=1, cores_per_worker=0, n_workers=0)
try:
assert bt._token_count_message("hi", "a long private chain of reasoning", None) \
> bt._token_count_message("hi", None, None)
finally:
bt.close()| # parts) remain unavailable for ISL reporting. | ||
| if token_ids is not None: | ||
| prompt_data = PromptData(token_ids=tuple(token_ids)) | ||
| elif isinstance(data.get("messages"), list | tuple): |
There was a problem hiding this comment.
[medium] Empty and multimodal messages are silently dropped from ISL.
isinstance([], list | tuple) is True, so a dataset messages: [] builds PromptData(messages=()); _token_count_prompt then calls apply_chat_template([], add_generation_prompt=True), which raises IndexError (verified) — and with no fallback that drops the sample's ISL and taints the drain. This branch is also a regression for multimodal chat: the previous code extracted text via _extract_prompt_text(...) (now dead code), whereas list-shaped content sent here makes text-only templates error.
Guard with truthiness so empty falls through to the text branch, and sanitize/fall back for non-text content:
elif data.get("messages"):
...Suggested test:
def test_empty_messages_fall_through_not_indexerror():
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
with pytest.raises(IndexError): # documents current behavior; guard removes it
tok.apply_chat_template([], tokenize=True, add_generation_prompt=True, return_dict=False)| self._thread, self._token_count_message, content, reasoning, tool_calls | ||
| ) | ||
|
|
||
| def _token_count_prompt( |
There was a problem hiding this comment.
[medium] _token_count_prompt lacks the safeguards _token_count_message already has.
Three gaps versus the message path:
- No
_normalize_tool_calls_for_templateon assistanttool_callsembedded in prompt history. Agentic multi-turn history carries wire-JSONargumentsstrings; Hermes/Qwen templates that iterateargumentsas a mapping then diverge or raise (verified 50 vs 49 tokens on Qwen2.5), while OSL handles them. - No try/except fallback — a template failure propagates and taints the drain instead of degrading to an approximate count.
chat_template_kwargsignored — the adapter forwards these to the server (e.g. Qwen3enable_thinking: 19 vs 15 tokens), but they are not threaded throughPromptData, so ISL diverges from the served prompt.
Mirror the message path: normalize embedded tool calls, wrap in the same try/fallback, and carry chat_template_kwargs through PromptData.
| model = getattr(tokenizer, "model", None) | ||
| model_type = type(model) | ||
| if ( | ||
| model_type.__module__ == "tiktoken.core" |
There was a problem hiding this comment.
[medium] The custom-tiktoken backend is unreachable — tiktoken is not a declared dependency.
For a Kimi-style tokenizer, AutoTokenizer.from_pretrained(...) imports tiktoken during load and raises ImportError before _backend_from_tokenizer is ever reached (tiktoken is not in pyproject.toml; not importable in the current env). As-is the advertised backend can never activate in a clean install. Declare tiktoken in the appropriate dependency group (pinned ==) or document the requirement, and add a test that exercises a real tiktoken-backed tokenizer rather than a fake.
| self._record(on_count, count) | ||
| for i, ((messages, tools), on_count) in enumerate(prompt_items): | ||
| try: | ||
| count = await self._tokenizer.token_count_prompt_async( |
There was a problem hiding this comment.
[medium] The prompt lane never shards and drains serially under the flush lock.
token_count_prompt_async runs on the small in-process thread pool (like the message path), and _flush awaits each prompt item one-by-one while holding self._lock. On main this in-process lane was justified as "rare (tool calls)"; this change makes it the path for every structured-chat ISL sample. At high QPS the end-of-run drain can exceed --drain-timeout, leaving ISL permanently uncounted (n_pending_tasks > 0) and blocking concurrent live flushes. Consider batching prompt render+encode through the sharded path like count_texts_async.
Suggested test (peak-concurrency probe with a stub tokenizer):
async def test_prompt_drain_is_not_serialized_per_item():
peak = active = 0
class T:
async def count_texts_async(self, *a, **k): return []
async def token_count_message_async(self, *a, **k): return 0
async def token_count_prompt_async(self, m, t, loop):
nonlocal active, peak
active += 1; peak = max(peak, active)
await asyncio.sleep(0.01); active -= 1; return 3
q = TokenBatchQueue(T(), asyncio.get_running_loop())
for _ in range(8):
q.enqueue_prompt((({"role":"user","content":"x"},), None), lambda c: None)
await q.drain_all()
assert peak > 1
nv-alicheng
left a comment
nv-alicheng
left a comment
There was a problem hiding this comment.
Review Council — Multi-AI Code Review
Reviewed by: Claude + Code-Quality | Depth: standard
codex unavailable in this environment. This PR already has a thorough human review; this pass independently corroborates the existing critical/high findings and adds only net-new items (see summary).
| if isinstance(ev_rec.data, PromptData) and ev_rec.data.token_ids is not None: | ||
| self.registry.record(self.metric_name, len(ev_rec.data.token_ids)) | ||
| return | ||
| if isinstance(ev_rec.data, PromptData) and ev_rec.data.messages is not None: |
There was a problem hiding this comment.
[Claude] medium (testing): The new ISL structured-prompt lane has no queue-integration test. Grep over tests/ finds zero references to IslTrigger or enqueue_prompt. Untested: (a) IslTrigger.fire messages routing here (322-328), including the silent drop when self._queue is None; (b) TokenBatchQueue.enqueue_prompt and the whole prompt lane in _flush (token_metrics.py:773-787) — the _record/_inflight decrement, drain-vs-live requeue, and the CancelledError handler that requeues all prompt_items (token_metrics.py:765). The only prompt coverage is test_token_count_prompt_preserves_messages_tools_and_generation_prompt, which calls token_count_prompt_async directly and never touches the queue or trigger; the metrics_table message-routing test covers OSL/TPOT only, not ISL. This is a hot-path lane with nontrivial cancellation/requeue logic — add direct queue + trigger tests.
|
|
||
| text: str | None = None | ||
| token_ids: tuple[int, ...] | None = None | ||
| messages: tuple[dict[str, Any], ...] | None = None |
There was a problem hiding this comment.
[Both] low (code-quality): PromptData is gc=False (types.py:283) but carries two mutable-container fields, messages: tuple[dict[str, Any], ...] and tools: tuple[dict[str, Any], ...], holding arbitrary nested chat/tool dicts. Every sibling struct in this file documents the gc-audit convention: TextModelOutput carries an explicit AT-RISK (gc=False): Has mutable container field 'tool_calls'... note, StreamChunk/ErrorData justify gc=False with "metadata contains only scalar key-value pairs." PromptData's docstring has no such note even though its dicts are the least scalar of any. The choice is practically safe (JSON-native dataset dicts don't cycle back; encoded synchronously and short-lived) — so this is a convention/documentation gap, not a runtime defect — but it breaks the file's own documented gc-safety process. Add the AT-RISK annotation. Separately, dict[str, Any] for messages/tools is loosely typed where a TypedDict (role/content/optional reasoning_content/tool_calls) would pin the keys the template code reaches for.
Review Council — Multi-AI Code ReviewReviewed by: Claude + Code-Quality | Depth: standard This PR already carries a thorough human review (arekay-nv, nvzhihanj). This council pass independently corroborates those findings and adds net-new ones — it does not re-file what's already commented. ✅ Corroborated (already commented — not re-filed)
🟡 Net-new — Should Fix (medium)
🔵 Net-new — Consider (low)
Dropped as duplicates of existing threads: IslTrigger double-
|
5 participants
What changed
messagesand top-level tool definitions for ISL instead of flattening them into plain text.apply_chat_template(..., tokenize=True)so template-inserted and special tokens follow the official tokenizer semantics.prompttokenization is the complete chat input used for ISL, whilemessagetokenization is structured assistant output used for OSL/TPOT.Why
The previous ISL path flattened multi-turn messages, losing roles, prior reasoning, tool-call history, tool results, tool definitions, and chat-template tokens. The existing output path also omitted structured reasoning unless tool calls were present. In addition, directly using Kimi's underlying tiktoken model would bypass model-specific preprocessing implemented by its official tokenizer wrapper.
These changes preserve the structured request and response fields and delegate their serialization to the model tokenizer.
Validation
8 passed, 160 deselectedin focused tests covering structured prompt round trips, ISL preservation, structured reasoning/tool OSL routing, Kimi wrapper selection, and process-shard support.The constant ISL difference comes from SGLang adding
function.strict: falseand top-leveldefer_loading: nullto each of the 62 benchmark tool definitions before applying Kimi's chat template. These unset JSON fields add 496 model-input tokens; SGLang excludes a 3-token generation stub from reported prompt usage, producing the observed 493-token per-turn difference. Excluding unset fields makes the server and client input token IDs match exactly.git diff --cached --check