Repository navigation
Multi-host disk prefix caching - #29
Merged
Merged
Conversation
A multi-host package refused any host prefix store ("the host prefix
store requires every rank in one process"), so long prompts on ranks
that run as separate processes always prefilled from scratch after
GPU eviction or a restart. On three DGX Sparks serving DeepSeek V4.1
Flash with TP3, a cold 1M-token prompt takes about 10 minutes.
Each host now keeps the records of its own rank under its own
--prefix-cache-disk directory, in the unchanged record format, with a
layout fingerprint bound to the rank and world size. Rank zero's
DiskStore indexes, sizes, evicts and plans every record as before;
its transport (offload::shard::Sharded) runs every store, fetch,
truncate and unlink on rank zero's Staged transport and sends the same
job, with the same file names, GPU ids and state slot, to each peer
over a dedicated TCP connection. Each peer runs it on its own Staged
transport. A job completes only after every rank finished it, so a
fetch is reported once every GPU holds its pages and row and a
checkpoint stays pinned until every rank copied it out. A store some
ranks wrote and another failed is unlinked from the writers before the
failure is reported.
Ranks join their stores after connecting RoCE and before the execution
protocol: every rank must have a disk store or none; peers connect to
a listener on rank zero's rendezvous address with a nonce sent over
the rendezvous connection; and the ranks keep only records every rank
holds with an equal header (rows only where every rank has the row),
closed under parents, deleting or truncating the rest. A memory tier
would keep every rank's bytes in one process, so multi-host packages
support disk-only retention and reject --prefix-cache-memory > 0.
Validation:
- cargo test --features server,runtime-package,roce,test-tools --lib
offload:: : 119 passed, 8 ignored (GPU package tests). New tests
cover the scan agreement and its parent closure, event and file-name
encoding, rank zero gating each job on every rank, a peer failure, a
partly written store removed from the rank that wrote it (both
directions, with a real disk thread), a broken peer connection, and
a peer service mirroring unlinks and truncates into its own
directory until closed.
- Three DGX Sparks (GB10), DeepSeek V4.1 Flash TP3 package exported on
6301710, flare run per rank with --prefix-cache-memory 0
--prefix-cache-disk ... --prefix-cache-disk-limit 200G, 16-key recall
prompts, temperature 0. "Disk" restarts all three ranks first, so
the prefix comes only from each host's NVMe:
| Prompt tokens | Cold | GPU-cached | Disk after restart | Per host |
|---------------|---------|------------|--------------------|----------|
| 65,164 | 21.5 s | 1.8 s | 2.0 s | 78 MB |
| 261,934 | 87.1 s | 2.1 s | 2.4 s | 364 MB* |
| 1,039,629 | 581.9 s | 3.5 s | 4.1 s | 1.37 GB* |
(* cumulative directory size, including earlier records)
Every answer had 16/16 keys, and each length's answer was identical
across cold, cached and disk runs; the 262k and 1M answers match the
earlier TP3 runs and the vLLM TP3 reference byte for byte.
- Disagreement: after deleting one host's 262k root record and
truncating another host's row of a 64k record, startup kept 3 of 6
records on every rank and truncated that row everywhere; the 64k
prompt restored 65,162 tokens and the 262k prompt prefilled cold,
both with unchanged answers.
- Eviction: with a 400M limit, three new prompts evicted older records;
after each request all three hosts held the same files, sizes and
totals.
Limitations: no memory tier across hosts; the partial-store cleanup
and broken-connection paths are unit-tested, not fault-injected on
hardware; compaction (truncate) is mirrored by the same path as unlink
but was not triggered in the hardware runs.
Signed-off-by: Javier López <59578422+javieratomic@users.noreply.github.com>
flare orbit run accepted --prefix-cache-memory, --prefix-cache-disk, --prefix-cache-disk-limit and --prefix-cache-disk-transport through the shared serving options, but ServingOptions::rank_options did not pass them on, so the ranks Orbit started ran without a host prefix store and the flags were silently ignored. Every rank now receives them, as each rank of a multi-host package keeps its own records on its host. Validation: - cargo test --bin flare: 115 passed, 1 ignored; a new test checks that the forwarded arguments reproduce the same load options and that omitted settings forward nothing. - flare orbit run of the DeepSeek V4.1 Flash TP3 package on three DGX Sparks with --executable, --prefix-cache-memory 0, --prefix-cache-disk and --prefix-cache-disk-limit 200G: rank zero restored a 1,039,629-token prompt from the three hosts' records in 3.8 s (16/16 keys, answer unchanged). Signed-off-by: Javier López <59578422+javieratomic@users.noreply.github.com>
A peer rank's record connection could fail while that peer was still copying into the GPU page ids of a job it had already received. Rank zero reported those jobs as failed, the scheduler returned the ids and the state slot to its free lists, and the peer could then overwrite them on its model stream. A peer's record service also released its staging pool before joining its copy and disk threads when a close did not drain. Rank zero's sharded transport now stops instead of reporting when a link fails or a job outlives its deadline. `PrefixStore::poll` returns an error in that state, and the worker stops through its existing failure path, so ids and pins held by unaccounted transfers are never reused or released. A failed link also shuts the socket down so the peer drains. The peer's service closes its transport on every exit path and retains the transport, its threads, and its registered staging together when the drain is not confirmed. Writes are queued to a per-connection writer thread, so the model worker never waits on a peer: a peer that stops reading fills a bounded queue and fails the job. Frames larger than the protocol limit fail before any job starts. A peer now checks every header against its own cache geometry before starting I/O, and the ranks compare their rank-local layout fingerprints at join, so a package whose ranks differ in plane or row size fails startup instead of never matching. Rank zero binds the record listener on a wildcard address and advertises, per peer, the local address of that peer's rendezvous connection; it also requires the connecting address to be that peer's rendezvous address, uses a short handshake deadline per connection, and fails startup when the nonce cannot be read. The disk limit is the smallest limit among the ranks, because a record every host must hold is bounded by the host with the least room. The join phase reports the agreed limit on every rank. A record that fails to restore is no longer offered to later lookups, and a failed cleanup unlink is reported with the store failure instead of discarded. Compatibility: the record format, package format, and CLI are unchanged. A multi-host store still requires disk-only retention with the same settings on every rank. A failed record link now stops serving instead of degrading to local-only copies. When a drain is unconfirmed, the affected process retains its transport and registered staging and exits; that memory is not reclaimed within the process. Validation: helios lib tests 809 passed, 18 ignored and flare bin tests 115 passed, 1 ignored on this branch. The four-Spark run used the same commits built on the private mainline, which also carries the RoCE proxy ABI 10 alignment this package requires. Four DGX Sparks, DeepSeek V4.1 Flash TP4, disk-only store: cold 25,658-token prompt 5.94 s, GPU-cached 0.44 s, disk restore after restart 0.54 s; cold 1,044,058-token prompt 474.5 s, disk restore after restart 2.98 s. Records were byte-identical on all four hosts. With no explicit limit, all four ranks reported the same agreed limit, 46.2 GiB, the smallest of the four derived values. Killing one rank's process during a 1M-token disk restore made the request fail with HTTP 503 and rank zero exit with a tensor-parallel control-socket error; no output was generated from reused ids. The record-link failure path and the unconfirmed-drain path have unit coverage only: the hardware fault injection removed a whole rank, so its control socket failed before the record link could be isolated. The join-time layout comparison is not unit-tested because preparing a store initializes CUDA; its per-header counterpart is.
…st caches A peer built its record transport after it reported its scan, so a rank that could not open its transfer threads, disk thread, or record directory failed after rank zero had already treated it as joined. The failure then surfaced mid-serving as a broken record link. The peer now starts its transport before it sends its scan, so such a failure fails the join. Add `qualify/multihost_cache.py`: it runs one `flare run` process per rank of a multi-host package on one host, each on its own GPU and record directory, and requires a GPU hit without a host hit, a host restore after evicting the GPU copies, and a restore after restarting every rank. After each phase it checks that every rank holds the same record files, because a record one rank lacks is unusable everywhere. The ranks-on-separate-hosts configuration had no scripted qualification. `qualify.host_cache.run_for` now accepts anything bound to a local HTTP port, so the multi-rank cluster reuses the single-host scenario helpers. Validation: helios lib tests 809 passed, 18 ignored and flare bin tests 115 passed, 1 ignored on this branch. The four-Spark run used the same commits built on the private mainline: cold 25,658-token prompt 6.02 s; disk restore after restarting the group 0.50 s; the join phase reported the agreed limit on every rank. Limitations: `qualify/multihost_cache.py` was not run. It needs a host with as many GPUs as the package's rank count, and the test cluster has one GB10 per host for a TP4 package. The transport-startup failure path has no fault injection.
original-el8
force-pushed
the
multihost-prefix-cache
branch
from
October 22, 2026 18:20
768e0a5 to
0fa229d
Compare
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
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.
Closes #23.
Multi-host packages can now use disk prefix caching. Before this change, the disk record directory was treated as local to one rank; with several hosts, records written by one rank were missing on the others, so a restart recovered little or nothing. Now every rank in a multi-host package writes and reads its own record set, and a restart reuses the cache across the whole package.
Attribution
The core of this work is from @javieratomic's branch in #23 (
3290d84,bcfb2b7), carried here unchanged as the first two commits. That includes the record format, the shard transport, and the per-host record layout;crates/flare/src/prefix/shard.rsis 91% his code as submitted (1,512 of 1,658 lines unchanged). The last two commits are follow-up work on failure paths and qualification.What the follow-up commits change
The original path is correct when every peer is healthy. These commits make failure paths explicit instead of silent:
PrefixStore::pollreturnsResult. A failed peer link, or a job that exceedsJOB_TIMEOUT(120 s), poisons the shard set: it reports nothing further and the worker stops. A broken record link is fatal rather than degrading to a partial cache.serveretains its transport and staging buffers viamem::forgetwhen a drain is unconfirmed, so a failed restore cannot free memory another thread is still using.HANDSHAKE_TIMEOUT.qualify/multihost_cache.pyis a new script that exercises the multi-host path end to end.Validation
Tests on this branch: helios 809 passed / 18 ignored; flare binary 115 passed / 1 ignored.
Hardware: four DGX Sparks, DeepSeek V4.1 Flash TP4, disk-only prefix caching, 100 GiB limit.