Skip to content

feat(model-host): accurate prompt token counting + pre-flight context-fit gate (WOR-1671)#715

Merged
rickcrawford merged 1 commit into
mainfrom
wor-1671-tokenizer
Jul 19, 2026
Merged

feat(model-host): accurate prompt token counting + pre-flight context-fit gate (WOR-1671)#715
rickcrawford merged 1 commit into
mainfrom
wor-1671-tokenizer

Conversation

@rickcrawford

Copy link
Copy Markdown
Contributor

What

The token-counting (tokenizers) and chat-template (minijinja) primitives shipped a while ago but had zero callers, so budgets and context-fit for locally served models still ran on a chars/4 heuristic, and an over-long prompt only failed opaquely inside the engine. This wires the real tokenizer into the managed-local request path.

Four layers:

  1. Acquisition — the pull best-effort prefetches tokenizer.json + tokenizer_config.json alongside a non-GGUF model (GGUF embeds its tokenizer), mirroring the existing config.json prefetch. A miss is non-fatal.
  2. Cache — a process-wide TokenizerCache parses a model's tokenizer.json once and reuses it per request (they're multi-MB files).
  3. Runtime — a launched deployment captures its tokenizer paths + served context window into a ServingContext, reachable per deployment via a defaulted PreparedDeploymentRuntime method (no blast radius) and a manager accessor.
  4. Dispatch — before forwarding a managed-local chat request, the gateway renders the model's chat_template (matching what the engine tokenizes) and counts the prompt against the model's own tokenizer. A prompt exceeding the served context window is refused pre-flight with a legible context_length_exceeded error instead of a cold path to an engine rejection.

Graceful degradation (three safe tiers)

  • No tokenizer on disk → keep the existing estimate (gate skipped).
  • Tokenizer but no/broken template → count concatenated message content (small, safe undercount).
  • Non-chat body → skip the gate.

Scoping notes

  • Chat-template rendering is used only to make the count accurate. The served engines (vLLM, SGLang, llama-server) apply the template themselves on /v1/chat/completions, so nothing renders on the way to the engine — the render_chat_template primitive is exercised here for counting and remains available for a future raw-prompt engine.
  • The pre-reservation budget path resolves before the managed lane is chosen, so it still uses the estimate; threading the exact count into it is a follow-up. The context-fit gate and the accurate count for logging land now.

Tests

Cache load/reuse/miss; render_prompt_for_count template + fallback + malformed-template paths; PromptTokenFit::fits() boundary (a prompt that exactly fills the window leaves no room to generate). Full gate green: nextest 6984/6984, clippy -D warnings, rustdoc, schema + capability checks.

Offline-verifiable: the counting mechanism is deterministic against a real tokenizer.json; no GPU needed.

…er and gate over-context prompts

The token-counting and chat-template primitives shipped but had no callers,
so budgets and context-fit for locally served models still ran on a
chars/4 heuristic and an over-long prompt only failed opaquely inside the
engine. This wires the real tokenizer into the managed-local request path.

- Acquisition: the pull now best-effort prefetches tokenizer.json and
  tokenizer_config.json alongside a non-GGUF model (GGUF embeds its
  tokenizer), mirroring the config.json prefetch. A miss is non-fatal and
  leaves the request path on its estimate.
- Cache: a process-wide TokenizerCache parses a model's tokenizer.json once
  and reuses it for every request (they are multi-megabyte files).
- Runtime: a launched deployment captures its tokenizer paths and served
  context window into a ServingContext, reachable per deployment through a
  defaulted PreparedDeploymentRuntime method and a manager accessor.
- Dispatch: before forwarding a managed-local chat request, the gateway
  renders the model's chat_template (matching what the engine tokenizes)
  and counts the prompt against the model's own tokenizer. A prompt that
  exceeds the served context window is refused pre-flight with a legible
  context_length_exceeded error instead of a cold path to an engine
  rejection. A model that shipped no tokenizer, a non-chat body, or a
  template that fails to render all fall back safely (concatenated content,
  then the existing estimate).

Chat-template rendering is used only to make the count accurate; the served
engines (vLLM, SGLang, llama-server) apply the template themselves on
/v1/chat/completions, so nothing renders on the way to the engine. Wiring
the accurate count into the pre-reservation budget path (it resolves before
the managed lane) is a follow-up.
@rickcrawford
rickcrawford merged commit 96a33f0 into main Jul 19, 2026
4 checks passed
@rickcrawford
rickcrawford deleted the wor-1671-tokenizer branch July 19, 2026 05:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant