Skip to content

Cache integer/decimal token sets in ConstrainedJSONGenerator (like the string set already is) - #220

Open
oscarbc96 wants to merge 1 commit into
huggingface:mainfrom
oscarbc96:cache-numeric-token-sets
Open

Cache integer/decimal token sets in ConstrainedJSONGenerator (like the string set already is)#220
oscarbc96 wants to merge 1 commit into
huggingface:mainfrom
oscarbc96:cache-numeric-token-sets

Conversation

@oscarbc96

Copy link
Copy Markdown

Problem

ConstrainedJSONGenerator.init unconditionally calls buildValidIntegerTokens
and buildValidDecimalTokens on every construction (Sources/AnyLanguageModel/Shared/StructuredGeneration.swift:110-111).
Each does a full 0 ..< backend.vocabSize scan and, per token, calls
backend.isSpecialToken(token) and backend.tokenText(token). On the MLX
backend, isSpecialToken calls tokenizer.decode(tokens:skipSpecialTokens:)
twice and tokenText calls it once more — three tokenizer decode calls per
vocab entry.

On a 151,936-token vocabulary (Qwen3-1.7B), that's ~456K decode calls per
scan, ~911K for the pair, on every single ConstrainedJSONGenerator
construction
— regardless of whether the schema being constrained contains
any numeric field at all (the scan isn't gated on schema content).

buildValidStringTokens already solves this exact problem: it's memoized in
StringTokenCache, keyed by tokenizer identity (vocabSize/eosToken/endTokens/sampleTexts).
The integer and decimal scans had no such cache.

Fix

Mirror the existing StringTokenCache pattern for the integer and decimal
scans: two new sibling Locked<[Key: Set<Int>]> dictionaries, with a
guard-and-store block around each build function. No other behavior change.

Measured (Qwen3-1.7B-4bit/MLX, Apple Silicon, M-series)

  • Each scan: 1.8–1.9s cold → 44µs/17µs on cache hit (five to six orders of
    magnitude).
  • End-to-end constrained generation (one call = one ConstrainedJSONGenerator
    construction): 9.9–12.3s → 6.2–8.6s, a ~30–40% reduction in the warm decode
    floor.
  • Output verified unchanged: cached token sets are byte-identical to their
    uncached computation, so this cannot itself introduce nondeterminism.

Testing

git apply --check verified clean against 0.8.0 (163f385). Manual
instrumentation (temporary ContinuousClock prints, not part of this diff)
confirmed cache hits on the second call within the same process. No existing
test coverage exercises ConstrainedJSONGenerator timing directly, so this
is a pure perf fix with no behavior change to verify against.

… are

ConstrainedJSONGenerator.init unconditionally called
buildValidIntegerTokens/buildValidDecimalTokens on every construction, each
scanning the full vocabulary (0..<backend.vocabSize) and, per token, calling
isSpecialToken (two tokenizer decodes) and tokenText (one more) - three
decode calls per vocab entry. On a 151,936-token vocabulary (Qwen3-1.7B) that
is roughly 456K decode calls per scan, ~911K for the pair, and it ran again
on every single ConstrainedJSONGenerator construction regardless of whether
the schema being constrained contained any numeric field at all.

buildValidStringTokens already memoizes its result in StringTokenCache,
keyed by tokenizer identity (vocabSize/eosToken/endTokens/sampleTexts). This
mirrors that exact pattern for the integer and decimal scans: two new sibling
dictionaries on StringTokenCache, guard-and-store around each build function,
no other behavior change.

Measured on Qwen3-1.7B-4bit/MLX, Apple Silicon: each scan drops from
1.8-1.9s (cold) to 44us/17us (cache hit) - five to six orders of magnitude.
End-to-end constrained generation drops from 9.9-12.3s to 6.2-8.6s per call,
a 30-40% reduction in warm decode time, with byte-identical cached token sets
(output verified unchanged across runs).
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