Skip to content

Replace fatalError with throwing KVCacheError in recoverable KV cache paths#451

Open
JuanColilla wants to merge 2 commits into
ml-explore:mainfrom
JuanColilla:upstream/kvcache-throwing-errors
Open

Replace fatalError with throwing KVCacheError in recoverable KV cache paths#451
JuanColilla wants to merge 2 commits into
ml-explore:mainfrom
JuanColilla:upstream/kvcache-throwing-errors

Conversation

@JuanColilla

Copy link
Copy Markdown

Why

Several KV cache paths crashed the whole process with fatalError on
conditions that are recoverable, not programmer bugs:

  • KVCacheSimple.toQuantized / RotatingKVCache.toQuantized crash when the
    requested group size doesn't evenly divide the model's key/value head
    dimensions, or (for RotatingKVCache) unconditionally, since rotating-cache
    quantization isn't implemented yet.
  • restoreCacheFromMetaState (used by loadPromptCache) had no validation
    of the array/metaState counts it reads out of a saved prompt-cache file,
    so a truncated or corrupted file crashed the process instead of failing
    the load.

A library shouldn't fatalError on attacker- or disk-corruption-controlled
input, or on a caller passing a group size that happens not to fit a given
model's head dimension — both are inputs the caller can reasonably want to
handle (skip quantization, reject a bad file, log and continue).

What

  • KVCacheSimple.toQuantized and RotatingKVCache.toQuantized now throws KVCacheError instead of calling fatalError, matching the existing
    throwing surface of loadPromptCache/savePromptCache.
  • restoreCacheFromMetaState now validates array/metaState counts for
    KVCacheSimple, RotatingKVCache, QuantizedKVCache, and ChunkedKVCache
    before constructing the cache, throwing KVCacheError with a message
    identifying which field was short.
  • maybeQuantizeKVCache's public, non-throwing signature is unchanged: it
    now falls back to leaving a cache unquantized (try?) rather than
    propagating, since silently keeping the unquantized cache was already its
    behavior for other ineligibility reasons.

How tested

  • Tests/MLXLMTests/KVCacheTests.swift: new regression tests constructing
    hand-crafted corrupt prompt-cache files (wrong array count for
    KVCacheSimple/RotatingKVCache, wrong metaState count for
    QuantizedKVCache/ChunkedKVCache) and asserting loadPromptCache throws
    KVCacheError instead of crashing; a test asserting
    RotatingKVCache.toQuantized() throws; a test asserting
    KVCacheSimple.toQuantized(groupSize:bits:) throws on a head dimension
    that divides none of the supported group sizes (32/64/128).
  • swift build compiles clean across every target.
  • This branch was rebased on top of main's current state, including the
    TurboQuant KV-cache-compression path merged after this change was first
    written; its one call site to the now-throwing toQuantized is updated in
    a second, clearly separated commit so the diff shows exactly what had to
    adapt and why.

KVCacheSimple.toQuantized, RotatingKVCache.toQuantized, and the prompt
cache deserialization path (restoreCacheFromMetaState) crashed the
process on corrupt cache files or incompatible group-size/head-dim
combinations. These are now throwing KVCacheError instead, matching
the existing throwing surface of loadPromptCache/savePromptCache.

maybeQuantizeKVCache falls back to leaving the cache unquantized on
failure rather than propagating, so its existing non-throwing public
signature is unchanged.

(cherry picked from commit 3440b9f)
…antized()

maybeTurboQuantizeKVCache's boundary-layer protection path called
KVCacheSimple.toQuantized() before it became throwing. The call site
already guards that groupSize resolves against the layer's head
dimensions, so the throw is not expected here in practice; use try? and
skip the layer (matching this function's existing skip-on-ineligibility
pattern) rather than force-unwrap, so a future edge case degrades to
leaving that layer at fp16 instead of crashing.
@aleroot

aleroot commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

@JuanColilla please see my PR: #453
just removing fatal errors is not the right approach...
For quantization, simply changing the leaf conversion methods from fatalError to throws moves the failure to callers rather than preventing invalid configurations.

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.

2 participants