Repository navigation
fix(offload): key multimodal prompts by media, one LMCache engine id per server, MP anchors on the chunk grid, keep a hybrid's last chunk - #3027
Conversation
…plit servers Four model-independent LMCache fixes, moved out of the Qwen3.8 offload PR: - tier_token_ids folds Sequence.cache_seed into the first ids of every LMCache lookup, save and load key the chunked schedulers and the DSV4 scheduler build, so prompts with the same text and different images no longer share KV keys; DSV4's SLOT boundary hashes start from the same seed, as BlockManager's chain does. - A hybrid's full-prompt hit on a chunk-aligned prompt keeps its last chunk when a joint boundary lies in it (KimiK3OffloadScheduler and its Qwen subclass); _loadable_hit always takes the sequence. - Under lmcache_mp with state images the prompt-end anchor is floored onto the chunk grid, the only place MP saves an image, and a prefill chunk ends on such an anchor. - lmcache_engine_id adds data_parallel_base_port, so two servers on one host no longer bind the same lookup socket. Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
|
Opened at How each point from that review maps here:
Self-check with the repo's
Validation on
Pending: the four GPU checks in the description:
Please don't merge before those are in. CI is green (Black, Ruff, non-GPU unit tests). |
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
some changes needed..
kimi-k3 + lmcache_offload: same chunk-aligned prompt twice (offload_tokens == whole prompt) and a gsm8k pass1/pass2 round? deepseek-v4.1: two images same text, B must not hit A? qwen3-next/3.5 + lmcache_mp: gsm8k pass1/pass2 vs #2636 963d7fa? two servers on one host: lmcache sockets do not cross?
ci green.
With images now saved on the chunk grid, lmcache_mp restarts restore state for most hits. A worker polls a retrieve only between forwards, and the first forward after a cold start can spend ~30 s compiling kernels. By the time the retrieve is seen terminal and its restore enqueued, the 30 s transfer deadline, counted from the retrieve's submission, has already passed, so the first pending restore stopped the engine (Qwen3.5 and Qwen3-Next GSM8K pass 2 after an ATOM restart). The retrieve is terminal at that point, so the deadline now bounds only the restore copy, counted from when it is enqueued. A restore whose fence raises keeps the submission clock, and a copy that never completes still stops the engine. Co-authored-by: Cursor <cursoragent@cursor.com>
The DSV4 scheduler's lookup and save keys had tests; its load did not, so building the load request from raw token ids passed every test. The new test fails under that change. Co-authored-by: Cursor <cursoragent@cursor.com>
…ed one A native-state worker restores two loads at a time, in order, between forwards. On a restarted server, a burst of 64 restores plus two slow first forwards (29 s and 21 s under host load) kept the last loads queued past the scheduler's 40 s bound with nothing in flight, and the engine stopped. A finished load now restarts the bound of the loads seen after it; a load whose completion is lost still stops the engine one bound after the loads ahead of it finish. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Pushed The GPU runs for valar-cop's review found two real problems in 1. Restore deadline counted from the retrieve (
2. Scheduler bound counted from dispatch, including the queue (
Self-check with the repo's
Validation on
Pending: Kimi-K3 + The description is updated to this head. |
|
@valar-cop, results for the four checks in your review. MI355X, 1. Kimi-K3 + 2. DeepSeek-V4.1-Flash, two images with the same text: B does not hit A (
3. Qwen3-Next / Qwen3.5 +
Accuracy matches This round also found two real problems, both at the start of pass 2 after the restart, where every request now restores:
The causes and deterministic A/Bs (old head stops, new head passes) are in this comment. Qwen3-Next has no row on the head. Since ~06:15 UTC on 11-05 it hits a HIP illegal memory access in pass 1 in this environment, including on main 4. Two servers on one host: the sockets don't cross (
cc @valarLip |
|
@ZLkanyo009 already reviewing this commit, results in a few minutes.. say |
| again. A load is restarted only by loads seen before it, so one whose | ||
| completion is lost still stops the engine one bound after the last of | ||
| those finishes.""" | ||
| loads = {operation for _, operation in self._active_load_operations.values()} |
There was a problem hiding this comment.
loads is a set, so position is hash order not dispatch order.. iterate self._active_load_operations.values() instead? one meta can dispatch several loads and the worker restores in that order
There was a problem hiding this comment.
@valar-cop Right, thanks. A restart's first step dispatches the whole burst at once, so hash order could leave a load queued behind a finished one on its old bound, which is the failure this method is for. Fixed in 753a552: positions now come from _active_load_operations.values(), which build_connector_meta fills in meta.add_request order. The set is kept only for membership. New test test_loads_dispatched_together_restart_in_dispatch_order (8 loads in one step) fails on 23d9ee2 under PYTHONHASHSEED 0-5 and passes with the fix.
cc @valarLip
There was a problem hiding this comment.
looked at it, yes: _active_load_operations is filled in _decide_load_after_alloc right before meta.add_request (mp/native_state_scheduler.py:565, chunked_scheduler.py:812), so the numbering is in dispatch order now, also within one step. resolved from my side.
The restart numbered the loads it first saw in one pass by iterating a set, so loads dispatched in the same step got hash order. A load dispatched after a finished one could then hold a lower number and keep its old bound. Iterate `_active_load_operations`, which is in dispatch order, the order the worker restores in. Co-authored-by: Cursor <cursoragent@cursor.com>
|
ci triage for 45123b6:
|
|
Pushed valar-cop's r4477625065 is a real bug in change 6.
Validation on
Also: the description's "Not changed" no longer links #3029. That issue was closed; the plugin-path note stays in the description. Pending: Kimi-K3 results (server loading); CI running. The description is updated to this head. |
main's #2912 made `CoreManager._init_shared_state` read `keep_finished_outputs` and `asyncio_mode`, which the hand-built namespace lacked. `atom_config_double` derives its fields from `dataclasses.fields(Config)`, so a field main adds next arrives with its default. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Pushed Pre Checkin was red on this PR's own test (run 41987910858):
Validation on
Pending: CI on this head. Kimi-K3 + |
|
Review 5779752875, item 1: Kimi-K3 + Setup: one MI355X node, TP8,
Before each tier load, 107 random 8192-token prompts (876k tokens, 1.3× the HBM pool) push the test prompt out of HBM. Chunk-aligned prompt, sent twice:
GSM8K pass1 / pass2 (questions 0-199, 5-shot). After pass 1, the same 1.3× flush pushes pass 1 out of HBM, so pass 2's hits come from the tier.
State loads: 111 attempted, 111 completed, 0 failed. No tracebacks in the server log. Prefix caching:
|
|
@ZLkanyo009 already looked at 45123b6, nothing new since my last review. say |
- mp scheduler: the load restart keeps one insertion-ordered dict of live loads instead of a dict of positions, a counter and a min(); same rule, same tests (dispatch-order mutation still fails under PYTHONHASHSEED 0-3). - BlockManager._joint_kv_boundary takes the tier hit as a required argument; the optional one gave the method two modes. Its two other callers pass `seq.offload_joint.kv_prefix_tokens`, the value it read. - tier_token_ids: `cache_seed` is never None on a Sequence (it is computed when the multimodal data has none); drop that check. - Comments: drop a measured figure (~30 s compile) and a default (8192-token ladder) that go stale; keep the constraints. - Drop the test that counted `_walk_prefix` calls through a monkeypatched private method: it tested the memo's mechanics, not a result. Co-authored-by: Cursor <cursoragent@cursor.com>
…nch, name the replay decision - align_checkpoint_anchor: its one caller already validates the grid against the hash block size, so the ValueError and its test go. - Kimi _withholds_last_prompt_token: a native seq always has a bound block manager; the plugin SeqView returns before the check. - _replay_decision returns a _ReplayDecision instead of a 3-tuple. - _record_checkpoint_end: cached_tokens -> resume_tokens. - Shorter restart docstring; drop a measured figure from a test docstring. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Update What changed
Size vs. merge-base Validation
Not changed
Pending: Pre Checkin on this head (CI running). After that, #2636 merges this head. cc @valarLip |
Split out of #2636 at valarLip's request (review, the
offload/config.py:463comment): these offload fixes are not Qwen3.8-specific, so they get their own PR and their own validation. #2636 now targets this branch and keeps only the Qwen3.8 changes. Merge order: this PR → #2636 → #2668.What
offload/config.py,lmcache_engine_id). The id names the ZMQ socket under a fixed host directory, so two servers on one host both bounddp0, and one server's scheduler could follow the other server's hits. The id is nowatom-offload-r{data_parallel_base_port}-dp{rank}. If the port is unresolved (0), the DP rank alone is used, as on main. The id is not part of a chunk's key, so stored chunks keep their keys._offload_common.tier_token_ids). Image placeholder ids are the same for every picture. Before this change, two prompts with the same text and different images had identical tier keys, and B loaded A's KV. Every lookup, save and load request now keys its ids throughtier_token_ids(seq, end), which foldsSequence.cache_seedinto the first 4 ids. This coversChunkedOffloadSchedulerBase(in-processlmcache_offloadandlmcache_mp), the native-state lookup, and the DSV4 PAGE+SLOT scheduler. DSV4's SLOT boundary hash chain (_chained_prefix_hashes(seed=)) now starts fromcache_seedasBlockManager's chain does; before, it started from −1. Text prompts key exactly as before.lmcache_mpstate images on its chunk grid (BlockManager.align_checkpoint_anchor,align_midstep_chunk;mp/native_state_scheduler.py). MP keeps a checkpoint only on its own chunk grid, so the prompt-end anchor is floored onto that grid when the grid still lies past where the request resumes. A prefill chunk that crosses an on-grid anchor ends at the anchor. Off-grid anchors are not cut for. The duplicate_record_checkpoint_endcall after_commit_joint_boundaryis removed._withholds_last_prompt_token,_loadable_hit). A full-prompt hit gives back its last token and is then floored to a chunk. For a chunk-aligned prompt, that floor drops a whole chunk. The generic default in_offload_commonkeeps that behavior. The Kimi-K3 scheduler, and therefore its Qwen GDN subclasses, does not give back when the joint boundary lies past the withheld chunk: the joint boundary already leaves a token to compute._loadable_hit(hit, num_prompt, seq)now has one signature, and all three callers passseq: chunked, DSV4 and mooncake_store.lmcache_mprestore is timed from when it is enqueued (mp/native_state_worker.py,_begin_restore). Found by the GPU runs below. The transfer deadline (30 s since fix(runtime): propagate MP failures and return unhealthy status #2837) was counted from the retrieve's submission. The worker polls only between forwards, so one long forward, such as the first one after a restart, which compiles the restore kernels, could use up the deadline before the restore copy was even enqueued. The engine then stopped withLMCacheTransferUnprovable. Main rarely reaches this path, because MP keeps few state images off its chunk grid; with change 3, every restarted request restores. While the retrieve is pending, the deadline still counts from its submission. Once the copy is enqueued, it bounds only the copy.lmcache_mpload queued behind a finished one starts its scheduler bound again (mp/scheduler.py,_restart_loads_queued_behind_a_finished_one). Also found by the GPU runs. The scheduler stops the engine when a load is not terminal 40 s (30 + 10) after it first sees it. A native-state worker restores two loads at a time, in the order it received them, and only between forwards. With change 3, 64 requests arriving at once after a restart need 64 restores. On a busy host (a 29 s first forward, then a 21 s step), the last 12 were still queued when the bound ran out, with nothing in flight:LoadOperationId(req_id=62) is still not terminal after 51s. Now, when a load finishes, every load dispatched after it starts its bound again. Dispatch order comes from_active_load_operations, whichbuild_connector_metafills inmeta.add_requestorder, so it holds within one step too. A load is restarted only by loads dispatched before it, so if its completion is lost, the engine still stops one bound after the last of those finishes. The defaults (30 s, 10 s, two slots) are unchanged. Main has the same bound but rarely restores.Each change has its tests (
test_lmcache_offload_connector.py,test_state_checkpoint.py,test_gdn_midstep_checkpoints.py,test_lmcache_mp_*), and the docs are updated (offload/README.md,docs/environment_variables.md).Not changed
atom/plugin/vllm/kv_transfer/) keeps main's behavior for changes 1 and 2.SeqViewhas nocache_seed, so plugin media keys are raw ids, as on main. vLLM'sParallelConfighas nodata_parallel_base_port, so the plugin engine id staysatom-offload-dp{rank}. Neither is a regression. For the same reasontier_token_idsreadsgetattr(seq, "cache_seed", -1)(main's tests' fake seqs lack it too); declaringcache_seedonSeqViewis left for a follow-up.Validation
Head
1e6fe588d. The branch started from main69d7931b0;d513f0ca9merges main30d2ad88f. Each GPU number is labeled with the head it ran on.f1662fe4cis change 5,4162fcde5adds only a CPU test,23d9ee232is change 6, and753a55282makes change 6 number loads in dispatch order (valar-cop's r4477625065).45123b668builds the two-servers test's configs withconftest.atom_config_double(), so the test keeps up with fields main adds toConfig(main's #2912 addedkeep_finished_outputs).00dd30e79and1e6fe588dtidy this PR's own code without changing behavior.CPU
review-pr'saffected_tests.py, CPU arm:4162fcde51975 passed / 12 skipped;23d9ee2321988 passed / 12 skipped;753a552821989 passed / 12 skipped;45123b668against main30d2ad88f1989 passed / 12 skipped;00dd30e791988 / 12;1e6fe588d1987 / 12 (each drops one test); nothing failing on head.run_unit_tests.sh) on45123b668: 11234 passed, 1164 skipped, 3 failed; on00dd30e79: 11233 passed, the same 3 failed. The 3 failures aretest_graph_marker×2 andtest_sparse_indexer_side_effect_op, which fail the same way on main (inductor g++).test_lmcache_mp_native_worker.py, GPU);test_lmcache_mp_native_scheduler.py);GPU, MI355X,
rocm/atom-dev:nightly_202610221638:GPU unit tests (the offload, checkpoint and
lmcache_mpfiles):2caf12f1e1249 passed / 9 skipped / 0 failed;4162fcde51252 / 9 / 0;23d9ee2321254 / 9 / 0. The new tests account for the increases.affected_tests.py's GPU-only files on1e6fe588d: 325 passed / 9 skipped. The skips needATOM_DSV41_REFERENCEor a Mooncake Store service.Two servers on one host (
2caf12f1e). Qwen3.5-35B-A3B-FP8, in-processlmcache_offload, one server per GPU, interval 256. Prompt X was stored on S1 and prompt Y on S2. Then X was sent to S2 and Y to S1. Then 60 fillers per server pushed both out of HBM, and X was sent to S1 and Y to S2 again.engine_atom-offload-r35849-dp0_…andengine_atom-offload-r34969-dp0_….offload_tokens0, same answers as when stored.offload_tokens,state_tier1, same answers.69d7931b0, both servers bindengine_atom-offload-dp0_…. The crossing did not show in that run, because each scheduler connected before the other server rebound the path; the shared path is the bug.DeepSeek-V4.1-Flash, two images with the same text (
4162fcde5). TP4,lmcache_mp(V4.1 accepts onlylmcache_mpor Mooncake), thinking off. A (red circle, "CAT") and B (blue square, "DOG") each have 559 prompt tokens with the same ids.offload_tokens0,cached_tokens0, and it answered "a square, it is blue … DOG".offload_tokens512 from its own entries, and each answered with its own picture (red circle "CAT", blue square "DOG").69d7931b0, same steps: B gotoffload_tokens512 from A's entries, and answered "a circle, it is red, and the word written on it is "CAT"".Qwen3.5 / Qwen3-Next +
lmcache_mpGSM8K. Questions 0-199. Pass 1 stores. ATOM is then restarted with LMCache kept, and pass 2 loads from LMCache.69d7931b02caf12f1e2caf12f1ef1662fe4cf1662fe4c4162fcde523d9ee23200dd30e79963d7fa0d([Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636, before the split)Main's low hit rate is the anchor bug change 3 fixes: MP keeps a state image only on its chunk grid, and main's anchors are off it.
Qwen3-Next has no
23d9ee232row. Since about 06:15 UTC on 11-05, Qwen3-Next stops in pass 1 with a HIP illegal memory access (recv_async_output,copy_done.synchronize()) in this environment. That includes main69d7931b0without any offload: 4 runs on two GPUs, including runs with a freshTMPDIRand inductor cache. The Qwen3-Next rows above ran before that.A/B for changes 5 and 6. Each natural failure needs a slow host, so the A/B shortens
lmcache.mp.transfer_deadline_son the pass-2 server only:2caf12f1e, pass 2 stopped at once withload:3:5 is still not terminal after 9s(native_state_worker.py:502). Onf1662fe4c, it scored 0.97 / 0.965 with all 200 retrieved.4162fcde5, pass 2 stopped at once withLoadOperationId(req_id=62, generation=62) is still not terminal after 18s, the same message as the natural failure; LMCache had retrieved only the first 64. On23d9ee232, it scored 0.975 / 0.975, with all 200 hitting (88.5%).Kimi-K3 +
lmcache_offload, TP8 on one node,recipes/Kimi-K3.mdflags but with prefix caching on (details: comment). Before each tier load, a 1.3× HBM flush;/debug/cache_statsshows the load came from the tier.30d2ad88f753a5528245123b6681920 is 2048 minus one 128-token block: the joint KV + state boundary is the last block boundary before the end. With the recipe's
--no-enable_prefix_caching, Kimi-K3 never hits the tier, on main and here alike: state checkpoints need prefix caching (block_manager.py:364). The run with the recipe's flags on753a55282had GSM8K 0.99 / 0.995 and 0 tier tokens.The
963d7fa0drow did not have the Kimi-K3 give-back, the DSV4 keys or seed,anchor % grid == 0, or changes 5 and 6. It is context only. Earlier, on963d7fa0d(comment): MiniMax-M3.1 with two images and the same text, on both in-processlmcache_offloadandlmcache_mp. B did not hit A in either.Notes
type(port) is intcheck is removed per valarLip's review. valar-cop had asked for it in r4439120010. The port is now read asint(... or 0)and used only when > 0.align_midstep_chunkdocstring says so.@valarLip