Support Kimi K3 Tokenizer

docs/CLI_QUICK_REFERENCE.md

@@ -96,6 +96,8 @@ Flag names shown as `--full.dotted.path --alias`. Both forms work.
 - `--model-params.max-new-tokens --max-output-tokens` - Max output tokens (default: 1024)
 - `--model-params.osl-distribution.min --min-output-tokens` - Min output tokens (default: 1)
 - `--model-params.streaming --streaming` - Streaming mode: auto/on/off (default: auto)
+- `--model-params.tokenizer-name --tokenizer` - HF repo ID or local tokenizer path used for client-side ISL/OSL/TPOT; overrides the served model name
+- `--model-params.enable-token-metrics / --model-params.no-enable-token-metrics` - Enable or disable client-side ISL/OSL/TPOT (default: enabled). Disabling skips tokenizer discovery and loading; non-token metrics are still collected.
 - `--runtime.min-duration-ms --duration` - Min duration: ms default, or with suffix (600s, 10m) (default: 600000)
 - `--runtime.n-samples-to-issue --num-samples` - Explicit sample count override
 - `--client.num-workers --workers` - HTTP workers (-1=auto, default: -1)
@@ -118,6 +120,30 @@ Flag names shown as `--full.dotted.path --alias`. Both forms work.
 
 **All other schema fields** are accessible via dotted paths (e.g., `--model-params.temperature`, `--model-params.top-k`, `--runtime.scheduler-random-seed`). Run `--help` to see the full list.
 
+### Tokenizer selection and Kimi K3
+
+By default, the benchmark uses `model_params.name` to find the tokenizer. Set
+`tokenizer_name` when the endpoint exposes a served alias or when the tokenizer
+is stored at a different path on the benchmark host:
+
+```yaml
+model_params:
+  name: "kimi-k3" # Name sent to the endpoint.
+  tokenizer_name: "moonshotai/Kimi-K3" # HF repo or benchmark-host path.
+```
+
+Kimi K3's custom Transformers tokenizer is supported even though Transformers
+labels it as a slow tokenizer: its BPE core uses the optimized Rust `tiktoken`
+backend. This framework capability does not by itself add Kimi K3 to an MLPerf
+ruleset's list of official submission models.
+
+If token metrics are not needed, avoid tokenizer discovery and loading with:
+
+```yaml
+model_params:
+  enable_token_metrics: false
+```
+
 ## Environment Variables
 
 **In YAML files** — use `${VAR}` or `${VAR:-default}` syntax:

docs/async_utils/services/metrics_aggregator/DESIGN.md

@@ -51,16 +51,21 @@ counted in `pending`. `flush_remaining` never raises.
 ### Sharded batch encoding (`BatchTokenizer`)
 
 The drain fans the whole buffer out across worker **processes**, one pinned
-per `CORES_PER_WORKER` (8) core block. Each worker runs the raw `tokenizers`
-backend's `encode_batch_fast` (Rust, rayon); a single BPE rayon pool
-saturates ~8 cores, so disjoint pinned blocks are how the whole machine is
-used. Workers are spawn-context, warmed in parallel at construction (bounded
-— a hung load is a startup error), and ignore SIGINT.
+per `CORES_PER_WORKER` (8) core block. Hugging Face Fast tokenizers run the raw
+`tokenizers` backend's `encode_batch_fast` (Rust, rayon). Custom Transformers
+wrappers backed by `tiktoken.Encoding`, including Kimi K3, run their native
+`encode` method through a length-counting adapter. Keeping the wrapper in the
+path preserves its special-token and long-input handling; process sharding
+provides parallelism when the wrapper has no batch API. A single BPE rayon pool
+saturates ~8 cores, so disjoint pinned blocks are how the whole machine is used.
+Workers are spawn-context, warmed in parallel at construction (bounded — a hung
+load is a startup error), and ignore SIGINT.
 
 The shard pool has no knob: it auto-sizes to one shard per 8-core block of
-the allowed CPU universe. There is no fallback — no fast Rust backend, or a
-failed/over-budget warmup, is a startup error, because an in-process slow
-path cannot keep up and would surface much later as an incomplete drain.
+the allowed CPU universe. There is no fallback — no supported optimized
+backend (Hugging Face Fast or `tiktoken`), or a failed/over-budget warmup, is a
+startup error, because an in-process slow path cannot keep up and would surface
+much later as an incomplete drain.
 Platforms without an affinity API (macOS) shard unpinned; each worker caps
 its rayon pool to the block size instead.
 
@@ -71,7 +76,9 @@ SIGTERM-kills the aggregator child — which the aggregator's SIGTERM handler
 turns into a best-effort `INTERRUPTED` final snapshot.
 
 Chat-template items (tool calls) run on the in-process thread lane —
-`apply_chat_template` is Python/Jinja; sharding buys nothing.
+`apply_chat_template` is Python/Jinja; sharding buys nothing. They use the
+tokenizer's native `tokenize=True` path so model-specific segment policies,
+including Kimi K3's XTML special-token boundaries, are preserved.
 
 ### CPU affinity: tokenize is post-run
 

examples/10_Agentic_Inference/README.md

@@ -68,7 +68,7 @@ sglang serve \
   --mem-fraction-static 0.95
 ```
 
-`--model-path` is the checkpoint loaded by the server. It can be a local path visible to the server container or a Hugging Face model ID, depending on your SGLang environment. `--served-model-name` is the OpenAI model name exposed to clients; set `model_params.name` in the YAML to the same value.
+`--model-path` is the checkpoint loaded by the server. It can be a local path visible to the server container or a Hugging Face model ID, depending on your SGLang environment. `--served-model-name` is the OpenAI model name exposed to clients; set `model_params.name` in the YAML to the same value. If that served name is not also a loadable tokenizer on the benchmark host, set `model_params.tokenizer_name` to the tokenizer's Hugging Face repo or benchmark-host path so client-side ISL/OSL/TPOT metrics are collected.
 
 ## Client YAML
 

examples/10_Agentic_Inference/kimi_agentic_benchmark.yaml

@@ -3,7 +3,11 @@ version: "1.0"
 type: "online"
 
 model_params:
-  name: "/model"
+  name: "kimi-k2.6"
+  # The served name above need not be a loadable tokenizer. Point this at the
+  # Kimi tokenizer's HF repo or a local path visible on the benchmark host so
+  # client-side ISL/OSL/TPOT metrics are collected.
+  tokenizer_name: "/path/to/Kimi-K2.6"
   temperature: 1.0
   top_p: 0.95
   max_new_tokens: 8192

pyproject.toml

@@ -80,6 +80,9 @@ dependencies = [
     "rich==14.3.3",
     # Needed for tokenization and OSL reporting
     "transformers==5.5.0",
+    # Required by custom model tokenizers such as Kimi K3. Although exposed as
+    # a Transformers "slow" tokenizer, K3 delegates BPE to this Rust backend.
+    "tiktoken==0.13.0",
     # Required by transformers' apply_chat_template
     "jinja2==3.1.6",
     "numpy>=1.26.4",

src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py

@@ -21,14 +21,16 @@
 per-sample text. The sharded pool is the drain-phase accelerator and is
 auto-sized (one shard per core block); live mid-run flushes run on a small
 in-process thread pool (``--tokenizer-workers``, default 2) owned by the
-queue's live loop. A tokenizer without a fast (Rust) backend is a startup
-error, never a silent slow path. Platforms without CPU affinity (e.g. macOS)
-shard unpinned at full speed; only cache/NUMA locality is lost.
+queue's live loop. A tokenizer without a supported optimized backend (a Hugging
+Face Fast ``tokenizers`` backend or a Rust ``tiktoken.Encoding`` core) is a
+startup error, never a silent slow path. Platforms without CPU affinity (e.g.
+macOS) shard unpinned at full speed; only cache/NUMA locality is lost.
 """
 
 from __future__ import annotations
 
 import asyncio
+import inspect
 import json
 import logging
 import multiprocessing
@@ -38,7 +40,7 @@
 from collections.abc import Callable
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from itertools import chain
-from typing import TYPE_CHECKING, Any, Protocol, cast
+from typing import TYPE_CHECKING, Any, Protocol
 
 import msgspec
 from inference_endpoint.endpoint_client.cpu_affinity import (
@@ -122,14 +124,64 @@ def load_reference_tokenizer(tokenizer_name: str) -> Any:
     return AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
 
 
+class _TikTokenBackend:
+    """Length-counting adapter for a Transformers wrapper around tiktoken.
+
+    Hugging Face labels custom tokenizers such as Kimi K3's
+    ``TikTokenTokenizer`` as "slow" because they inherit
+    ``PreTrainedTokenizer`` and do not expose ``backend_tokenizer``. Their BPE
+    core is nevertheless the Rust implementation in ``tiktoken.Encoding``.
+
+    Keep the model-provided wrapper in the loop instead of calling
+    ``Encoding.encode_batch`` directly. Kimi K3's ``encode`` method preserves
+    its special-token policy and splits very long or pathological strings to
+    avoid tiktoken's input-size and long-whitespace-run failure modes. The
+    outer Endpoints process sharding supplies parallelism across texts.
+    """
+
+    def __init__(self, tokenizer: Any) -> None:
+        self._tokenizer = tokenizer
+        # Kimi K3 exposes this explicit switch so raw user/model text cannot
+        # reinterpret a literal ``<|...|>`` substring as an XTML control token.
+        # Other tiktoken-backed wrappers may not expose it.
+        try:
+            self._supports_allow_special = (
+                "allow_special_tokens" in inspect.signature(tokenizer.encode).parameters
+            )
+        except (TypeError, ValueError):
+            self._supports_allow_special = False
+
+    def count_texts(self, texts: list[str]) -> list[int]:
+        if self._supports_allow_special:
+            return [
+                len(self._tokenizer.encode(text, allow_special_tokens=False))
+                for text in texts
+            ]
+        return [len(self._tokenizer.encode(text)) for text in texts]
+
+
+def _backend_from_tokenizer(tokenizer: Any) -> Any | None:
+    """Return a supported optimized length-counting backend."""
+    backend = getattr(tokenizer, "backend_tokenizer", None)
+    if backend is not None:
+        return backend
+
+    model = getattr(tokenizer, "model", None)
+    model_module = type(model).__module__.partition(".")[0]
+    if model_module == "tiktoken" and callable(getattr(tokenizer, "encode", None)):
+        return _TikTokenBackend(tokenizer)
+    return None
+
+
 def load_reference_backend(tokenizer_name: str) -> Any | None:
-    """Raw tokenizers backend (fast Rust path) for length counting.
+    """Optimized backend for reference-tokenizer length counting.
 
-    Counting through the backend avoids the transformers "sequence longer than
-    model_max_length" warning the Python wrapper emits, so no ``model_max_length``
-    override is needed. ``None`` if the tokenizer has no fast backend.
+    Usually this is the raw Hugging Face ``tokenizers`` backend. A custom
+    Transformers tokenizer backed by ``tiktoken.Encoding`` receives an adapter
+    that preserves its model-specific encode behavior. ``None`` means there is
+    no supported accelerated backend.
     """
-    return getattr(load_reference_tokenizer(tokenizer_name), "backend_tokenizer", None)
+    return _backend_from_tokenizer(load_reference_tokenizer(tokenizer_name))
 
 
 def _init_worker(tokenizer_name: str, core_set: list[int]) -> None:
@@ -160,15 +212,34 @@ def _init_worker(tokenizer_name: str, core_set: list[int]) -> None:
     global _WORKER_BACKEND
     _WORKER_BACKEND = load_reference_backend(tokenizer_name)
     if _WORKER_BACKEND is not None:
-        _WORKER_BACKEND.encode("warmup", add_special_tokens=False)
+        encode_lengths(_WORKER_BACKEND, ["warmup"])
 
 
 def encode_lengths(backend: Any, texts: list[str]) -> list[int]:
-    """Per-text token counts via the raw tokenizers backend, one rayon call."""
+    """Per-text counts through an optimized HF Fast or tiktoken backend."""
+    count_texts = getattr(backend, "count_texts", None)
+    if count_texts is not None:
+        return count_texts(texts)
     encode_batch = getattr(backend, "encode_batch_fast", None) or backend.encode_batch
     return [len(e.ids) for e in encode_batch(texts, add_special_tokens=False)]
 
 
+def _chat_template_token_count(tokenizer: Any, messages: list[dict[str, Any]]) -> int:
+    """Count a rendered conversation through the tokenizer's native path.
+
+    In particular, Kimi K3 assigns ``allow_special`` per XTML segment. Rendering
+    to one string and subsequently calling ``tokenize`` loses those boundaries
+    and can reinterpret literal special-token text from users or tools as
+    structure. ``tokenize=True`` preserves the model tokenizer's exact policy.
+    """
+    encoded = tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=False
+    )
+    if isinstance(encoded, dict):
+        encoded = encoded["input_ids"]
+    return len(encoded)
+
+
 def _worker_encode_lengths(texts: list[str]) -> list[int]:
     """Per-text token counts for a shard, in one rayon-parallel call."""
     backend = _WORKER_BACKEND
@@ -248,6 +319,7 @@ def __init__(
             max_workers=max(1, live_workers), thread_name_prefix="tok-thread"
         )
         self._load_tokenizer()  # also computes the chat-template baseline
+        self._backend = _backend_from_tokenizer(self._tokenizer)
         # Process shards for the batched text path. Empty only when
         # in-process mode was explicitly requested (n_workers=0 or
         # cores_per_worker<=0; ctor overrides used primarily by tests —
@@ -263,22 +335,11 @@ def _load_tokenizer(self) -> None:
         # Baseline = tokens from a [user, empty-assistant] pair minus the [user]
         # prefix alone, so the assistant frame is subtracted from message counts.
         try:
-            prefix = cast(
-                str,
-                tok.apply_chat_template(
-                    [_PREFIX_USER_MSG], tokenize=False, add_generation_prompt=False
-                ),
+            self._prefix_len = _chat_template_token_count(tok, [_PREFIX_USER_MSG])
+            with_assistant = _chat_template_token_count(
+                tok, [_PREFIX_USER_MSG, {"role": "assistant", "content": ""}]
             )
-            self._prefix_len = len(tok.tokenize(prefix))
-            with_assistant = cast(
-                str,
-                tok.apply_chat_template(
-                    [_PREFIX_USER_MSG, {"role": "assistant", "content": ""}],
-                    tokenize=False,
-                    add_generation_prompt=False,
-                ),
-            )
-            self._baseline = len(tok.tokenize(with_assistant)) - self._prefix_len
+            self._baseline = with_assistant - self._prefix_len
         except Exception:
             self._prefix_len = 0
             self._baseline = 0
@@ -296,18 +357,19 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None:
         process's affinity mask (or the online CPU count when the platform
         has no affinity API — shards then run unpinned), always at least one;
         an explicit count is clamped to that capacity. An environment that
-        cannot shard — no fast Rust backend, a warmup that fails or exceeds
-        its budget — raises instead of silently degrading to a slow path
-        that cannot keep up with completions.
+        cannot shard — no optimized backend, a warmup that fails or exceeds
+        its budget — raises instead of silently degrading to a slow path that
+        cannot keep up with completions.
         """
         if cores_per_worker <= 0 or n_workers == 0:
             logger.info("BatchTokenizer: in-process tokenization (explicit)")
             return
-        if getattr(self._tokenizer, "backend_tokenizer", None) is None:
+        if self._backend is None:
             raise RuntimeError(
-                f"tokenizer {self._tokenizer_name!r} has no fast (Rust) "
-                "backend; token metrics require one to keep up with "
-                "completions. Use a fast tokenizer, or disable token metrics."
+                f"tokenizer {self._tokenizer_name!r} has no supported optimized "
+                "backend (Hugging Face Fast or tiktoken); token metrics require "
+                "one to keep up with completions. Use a supported tokenizer, "
+                "add a backend adapter, or disable token metrics."
             )
         # The full allowed CPU universe (cgroup-clamped) drives the shard block
         # math. cgroup_clamped_cpus owns the probe-and-restore of this process's
@@ -364,11 +426,9 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None:
     # -- batched text path --------------------------------------------------
 
     def _encode_lengths_inproc(self, texts: list[str]) -> list[int]:
-        tok = self._tokenizer
-        backend = getattr(tok, "backend_tokenizer", None)
-        if backend is not None:
-            return encode_lengths(backend, texts)
-        return [len(tok.tokenize(t)) for t in texts]  # type: ignore[union-attr]
+        if self._backend is not None:
+            return encode_lengths(self._backend, texts)
+        return [len(self._tokenizer.tokenize(t)) for t in texts]  # type: ignore[union-attr]
 
     async def count_texts_async(
         self,
@@ -424,10 +484,9 @@ def _token_count_message(
         if tool_calls:
             msg["tool_calls"] = _normalize_tool_calls_for_template(tool_calls)
         try:
-            rendered = tok.apply_chat_template(  # type: ignore[union-attr]
-                [_PREFIX_USER_MSG, msg], tokenize=False, add_generation_prompt=False
+            full = _chat_template_token_count(  # type: ignore[arg-type]
+                tok, [_PREFIX_USER_MSG, msg]
             )
-            full = len(tok.tokenize(rendered))  # type: ignore[union-attr]
             return max(0, full - self._prefix_len - self._baseline)
         except Exception as exc:
             key = f"{self._tokenizer_name}:{type(exc).__name__}"

src/inference_endpoint/commands/benchmark/accuracy.py

@@ -189,7 +189,7 @@ def _load_osl_backend(has_accuracy: bool, tokenizer_name: str | None) -> Any | N
     """Load the reference tokenizer backend for accuracy OSL, or None to disable.
 
     Loaded only when a real accuracy dataset exists; a load failure or a tokenizer
-    with no fast (Rust) backend disables OSL rather than failing scoring.
+    with no supported optimized backend disables OSL rather than failing scoring.
     """
     if not (has_accuracy and tokenizer_name is not None):
         return None
@@ -202,15 +202,16 @@ def _load_osl_backend(has_accuracy: bool, tokenizer_name: str | None) -> Any | N
             e,
         )
         return None
-    # A tokenizer with no fast (Rust) backend disables OSL rather than falling
-    # back to a slow Python-tokenizer count: the perf side
-    # (token_metrics._setup_shards) requires a fast backend too and raises without
-    # one, so OSL stays fast-only and consistent on both sides. Warn so the skip is
-    # visible instead of silently dropping the block.
+    # A tokenizer with no supported optimized backend disables OSL rather than
+    # falling back to a slow Python-tokenizer count: the perf side
+    # (token_metrics._setup_shards) has the same requirement and raises without
+    # one, so OSL stays consistent on both sides. Warn so the skip is visible
+    # instead of silently dropping the block.
     if osl_backend is None:
         logger.warning(
-            "Accuracy OSL disabled: tokenizer %r has no fast (Rust) backend "
-            "(token counting requires one, as on the perf side)",
+            "Accuracy OSL disabled: tokenizer %r has no supported optimized "
+            "backend (Hugging Face Fast or tiktoken; token counting requires "
+            "one, as on the perf side)",
             tokenizer_name,
         )
     return osl_backend

src/inference_endpoint/commands/benchmark/execute.py

@@ -488,7 +488,10 @@ def setup_benchmark(
     model_name = config.model_params.name
     tokenizer_override = config.model_params.tokenizer_name
     tokenizer_name: str | None
-    if tokenizer_override:
+    if not config.model_params.enable_token_metrics:
+        logger.info("Client-side token metrics disabled by configuration")
+        tokenizer_name = None
+    elif tokenizer_override:
         if not _check_tokenizer_exists(tokenizer_override):
             raise SetupError(
                 f"Tokenizer override '{tokenizer_override}' could not be verified. "