Skip to content

Fix prompt-cache round-trip support for ArraysCache, MambaCache, and CacheList#155

Open
ronaldmannak wants to merge 6 commits intoml-explore:mainfrom
PicoMLX:kvcache-restore
Open

Fix prompt-cache round-trip support for ArraysCache, MambaCache, and CacheList#155
ronaldmannak wants to merge 6 commits intoml-explore:mainfrom
PicoMLX:kvcache-restore

Conversation

@ronaldmannak
Copy link
Copy Markdown
Contributor

Proposed changes

This change keeps the existing public savePromptCache / loadPromptCache API intact and extends cache reconstruction so composite and array-backed caches can be restored correctly from persisted prompt caches.

What changed

  • ArraysCache now serializes enough metadata in metaState to preserve:
    • total slot count
    • sparse slot positions
    • leftPadding
  • MambaCache now round-trips through the same array-cache restore path while preserving its concrete type
  • CacheList now serializes child cache structure in metaState and reconstructs children recursively on load
  • prompt-cache save/load continues to use the existing safetensors metadata format and Python-compatible cache class names

Tests

Added round-trip coverage for:

  • sparse ArraysCache
  • ArraysCache with leftPadding
  • MambaCache type preservation
  • CacheList with KV-only children
  • CacheList with hybrid children (MambaCache + KVCacheSimple)

Notes

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

return result
}
set {
// Handled by restoreFromMetaState
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this OK that it is a NOP? I don't see any calls to set in the code, but other KVCache types do allow set. Should this be a fatalError() instead -- is it a bug to call the setter on CacheList?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@davidkoski The set is required to conform to the KVCache protocol. The same NOP pattern exists in CacheList.metaState (line 1220) as well, btw. I agree a silent NOP could bite us in the future and I've replaced both with assertionFailure() so they trigger in debug builds but are stripped in release. Happy to switch to fatalError() if you'd prefer a hard crash in all builds.

Copy link
Copy Markdown
Collaborator

@davidkoski davidkoski left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks good, thank you. See what you think of my question -- not a blocker for merge.

@davidkoski
Copy link
Copy Markdown
Collaborator

Has a few conflicts from #158 now

@ronaldmannak
Copy link
Copy Markdown
Contributor Author

@davidkoski done. should be good to go

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants