Repository navigation
Mlx: the shim drives mlx-lm 0.32's caches, whose state carries everything, beside 0.31's - #18
Open
IsuraManchanayake wants to merge 18 commits into
Open
IsuraManchanayake wants to merge 18 commits into
IsuraManchanayake wants to merge 18 commits into
Conversation
baseRT 0.3.1 defines BASERT_ABI_VERSION and exports baseRT_abi_version(): two libraries with the same number drive the same way, whatever their version numbers say. The build reads the macro from the vendored header (now 0.3.1's), the loader resolves the symbol when the library has it, and the numbers decide: equal serves, unequal is refused by number. A library that states none (0.3.0 and older), or a superfluid built against a header without the macro, is judged by major.minor as before. The symbol and configuration-struct checks stay. Library carries the number it found.
…s; superfluid compat; exit 4 A probe may state the contract number it checked the runtime against (Probe.abi: baseRT's engine ABI), the worker's check report carries it, and `superfluid runtimes --json` shows it as `abi` on a ready runtime. The basert row also carries `requires`, the header version and ABI number this build was built for, whether or not an engine was found, so a host can compare numbers itself. `superfluid compat [--json]` prints the same for a release. `superfluid serve` exits 4 when the runtime for the model is not installed or cannot run here (its engine library missing or of another ABI), where a flag it got wrong exits 2, so a host can tell the two apart. The catalog finds superfluid-workerd and the shipped adapters beside the real executable, links resolved: a superfluid reached through a symlink (a package manager's bin/, a host's shim) kept running the engine in-process because the link's directory held neither.
…at.json The page names the contract a host builds on: how to start superfluid (the real path, --runtime, --basert-lib and BASERT_KERNEL_LIB, --offline, one --sessions per server, SUPERFLUID_HOME), how to ask for the compatibility verdict before serving, how to choose a release from the compat.json each release now carries beside its tarball, and the exit codes. The CLI reference, the runtimes page and the changelog follow.
…he unnumbered line; the engine install chooses by ABI number The vendored header is baseRT 0.4.0's, so the fallback for a library that states no ABI number can no longer be "the header's minor": that would refuse 0.3.0, which carries ABI 1's functions and layouts without stating it. The line that does so is named instead, 0.3.x for ABI 1, and a library without a number serves only if it is of that line. No later ABI has such libraries. A library of 0.4.0 or later is asked for the model configuration by size (baseRT_get_config_into), into a zeroed struct of these bindings' own size, so the struct may grow or shrink between them without a change of ABI number; an older library is read by value with the size check as before. Every read of the configuration goes through config_of(). `superfluid runtime install basert` chooses by the same rule the loader applies: a release's ABI number, read from the basert-engine-abi-<N>.txt asset baseRT publishes beside its bundles, or a pin's, with the 0.3.x line accepted without one; another engine installs only with --untested. `superfluid compat` and the runtimes report name the unnumbered line too.
The pack step wrote dist/<name>.compat.json, but the artifact the release job downloads carried only the tarball and its digest, so the sidecar the docs promise beside the tarball never reached the release.
`superfluid compat` and the release's compat.json described basert alone. Now an adapter declares its part through Runtime::compat, its worker answers `compat --json`, and the daemon merges the answers of every adapter it ships: for basert the header version, the ABI number, the line accepted without one and the engines it installs on its own; for llama.cpp the upstream builds it was built and tested against; for MLX the versions it installs and the Python they run on. The llama.cpp and MLX adapters are not linked into the daemon, so asking their workers is what keeps the daemon free of them.
…without credentials The report inserted the linked basert adapter by name and then asked the shipped workers, so a source build with --features llamacpp and no standalone worker left llamacpp out. Every linked adapter now answers in process first. The shipped worker's compat is run through exec::without_secrets, as check, plan, install and serving are, so a token in the caller's environment does not reach it.
… install here `superfluid runtime versions <id> [--json]` lists the versions an adapter installs on its own, newest first: which install on this machine and as what, and which are installed. `superfluid runtimes` gains an INSTALLABLE column, the newest such version here, and `installable` in its JSON. The adapter kit answers through a new `Runtime::versions` method, by default each tested version `compat` names planned for the host, and a worker command `versions`; both from the recipe alone, since the table asks on every run. The basert adapter lists its pins only, with the short ABI verdict for one this build does not drive, because its `plan` consults the release listing when a pin fails.
…efore the linked adapter The listing went to the linked adapter whenever this build links one, so the worker override that runtime execution and runtime install honour was ignored, and the INSTALLABLE column could name versions the selected installer does not provide. A named worker is asked first.
…a.cpp's parameters An adapter declares the options its runtime takes on serve (Runtime::options), each with the shape of its value; the daemon checks a value against the declaration before a worker starts and forwards it as --<name> <value>, the way --pull-<name> already works. superfluid runtime options <id> lists them and superfluid runtimes --json carries them. llama.cpp gets its model and context parameters as --llamacpp-<name>: n-gpu-layers, main-gpu, split-mode, tensor-split, load-mode, lazy-mode, repack, threads, threads-batch, batch-size, ubatch-size, flash-attn, cache-type-k, cache-type-v, kv-offload, op-offload, swa-full, the RoPE and YaRN settings, and lib. MLX gets venv, prefill-chunk and cache-limit-mb. A knob every runtime could take stays a flag of serve.
…ced by its cache types, and the options reach every worker start Four findings of the review: - The state identity folds in the RoPE, YaRN and cache-type settings, so a store written under one configuration is refused under another instead of continuing from keys that mean something else. With none of them set the key is the weights alone, as before. - A KV cell is priced by the cache types chosen, by ggml's row sizes, in the descriptor and in the sizing before a load; f16 remains the default. - The options are set for the process as given, before the runtime is required, and ride in the worker's prefix from the first probe on, so a library or environment option counts for availability and sizing. When the probe fails and the worker answers without the options, the option at fault is named as before. - In-process workers get their arguments through the worker's own parser, so --venv lands in its field and the MLX environment option is honoured in a linked daemon and at a fleet node.
…on's own tokenizer sees the llama.cpp library Two more findings of the review: the worker's pull now takes the runtime's options apart from the pull's own arguments, and both pull paths pass them, so a pull by id runs in the environment the check passed and the serve will use. The llama.cpp tokenizer library the daemon loads in process finds libllama as the worker does, by SUPERFLUID_LLAMA_LIB, which the daemon sets from --llamacpp-lib before the first load.
…hing, beside 0.31's mlx-lm 0.32 dropped meta_state: a cache's state now carries its offset and settings beside its arrays, a grown cache's state holds the whole buffer rather than the live tokens, from_state takes the state alone, and a CacheList's state names its members. The shim told a cache by meta_state, took every state leaf for an array and cut nothing, so a copy failed, a layer of several caches could not be exported, and nothing freed the room past the live tokens. The shim now reads a cache either way: a live view cuts a grown cache's keys and values to its offset, copies touch arrays only, exports carry the plain values in the header, and imports build the cache as its mlx-lm wants. mlx 0.32.4 with mlx-lm 0.32.0 joins the tested versions, pinned to the macOS 14 wheels. Found by the nightly runtime job on a fork: #10
…s state holds a growth step _measure_kv_bytes read a cache's state after one token, which from mlx-lm 0.32 is the whole buffer the first growth step made: 256 tokens' worth, so a runtime priced a token 256 times over, the pool was sized by it, and _slack_tokens counted 256 times too few, which left a cut sequence uncompacted. The measurement now reads the live state, as the compaction accounting does. The test builds a runtime the usual way over a one-layer model whose token is 16 bytes: before the fix it measured 4096 and a cut from 8192 to 128 tokens kept every slot.
2 of 3 tasks
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #15: its diff includes #10, #11, #14 and #15 until they merge; this PR's own commit is the last one. I will rebase it onto
rc-0.1.3as they land.What changed
mlx-lm 0.32.0, which pip resolves with mlx 0.32.4, dropped
meta_state: a cache'sstatenow carries its offset and settings beside its arrays, a grown cache'sstateholds the whole buffer rather than the live tokens,from_statetakes the state alone, and aCacheList's state names its members. The shim told a cache apart bymeta_state, took every state leaf for an array, and relied onstatecovering the live tokens, so a copy failed on a non-array leaf, a layer of several caches could not be exported, and nothing freed the room past the live tokens.The change is all in the Python shim, which now reads a cache either way:
_has_metatells the two mlx-lm generations apart, per class or instance._live_stateis whatstateused to be: for the caches whose buffers grow in steps (KVCache,QuantizedKVCache) the keys and values sliced to the offset, lazily; a ring buffer or a cache of another shape as it is; aCacheListmember by member.macosx_15ormacosx_26wheel of mlx, which would not install on an older Mac. The recording script gets the same rule in CI: a nightly job tests each runtime's newest upstream release and records the result #13.Why
The fork's nightly job found the pair failing the adapter's tests: IsuraManchanayake#10. Without this the shim works only with mlx-lm 0.31 and the recipe's tested list stops at mlx 0.32.2 with mlx-lm 0.31.3. Reviewed in the fork as IsuraManchanayake#12.
How it was tested
cargo clippy --all-targets -- -D warnings: clean.cargo test(CI at this head on macos-14, ubuntu-24.04 and ubuntu-24.04-arm).Against both pairs locally, mlx 0.32.4 + mlx-lm 0.32.0 and the tested 0.32.2 + 0.31.3, each in its own environment: the adapter's 24 unit tests, the 15 conformance tests and the 11 real-model tests on Qwen3-0.6B-4bit pass on both.
Not run: a full
superfluid servein the 0.32 environment, and a state exported under 0.31 restored under 0.32 (the shapes match by construction, but it was not exercised).The review on the fork (mlx: the shim drives mlx-lm 0.32's caches, whose state carries everything, beside 0.31's IsuraManchanayake/superfluid#12) found the construction-time measurement reading a 0.32 cache's whole growth step, 256 times a token; fixed in f6d5922 with a test that builds a runtime the usual way, and the runs above repeated on both pairs.
cargo clippy --all-targets -- -D warningscargo testChanges a wire protocol, the WAL format or the tick ABI (says so above, version bumped)