Repository navigation
feat(remote_cxl): read hints, so the memory node can prepare KV ranges before the load reads them - #130
Draft
youngrok-XCENA wants to merge 23 commits into
Draft
feat(remote_cxl): read hints, so the memory node can prepare KV ranges before the load reads them#130youngrok-XCENA wants to merge 23 commits into
youngrok-XCENA wants to merge 23 commits into
Conversation
A remote CXL agent can load one out-of-tree region extension, registered
under the maru.region_extensions entry-point group and named with
--region-extension. It serves memory that is slow until a device has
brought a page in (a CXL-SSD with a DRAM cache in front of its flash):
- validate_access may refuse a region before it is mapped, mapped hears
it once the NIC has it (raising keeps the region out of use), and
retire is asked before a region goes back; while it says no, the
region stays here unregistered and the sweep asks again.
- hint takes read_hint / read_hint_end messages (op "extension") from
any authenticated connection, once the agent has checked the epoch, the
action and that every range lies in a region it maps and some
connection holds, under a matching handle. It runs under the agent's
lock and only queues.
- The hello advertises region_extension {"version": 1, "actions": [...]}
and RegionMapClient keeps the hint actions it advertises; anything else
means no hints.
An agent that fails to load its extension exits.
… bytes early A serving engine knows which stored keys a request will load well before its worker starts the load. Read hints carry that to the memory side, where a region extension can bring the bytes into a faster tier. Maru's KV data path stays as it is: a hint is advice that never gates a READ or WRITE, and one that is late or lost costs only the time it would have saved. - read_hints (maru_read_hints in vLLM, remote_cxl only) turns them on. MaruHandler.hint_read / end_read_hint queue a hint and its end; one thread per backend probes maru-server for where the hits are (a probe with locations, which takes no lease) and sends read_hint and read_hint_end to every agent whose hello advertises them, in messages of at most 4096 ranges. At most 1024 hints wait; ends are never dropped. - A client that loads with read_hints on hints each load as its READs start, under a name of its own, and ends that hint after the load (for engines that name no request to their storage). - The vLLM scheduler connector hints a request when it first finds its asynchronous load, or on arrival with maru_read_hint_on_arrival; only the scheduler gets read_hints, so vLLM loads are not hinted twice.
…region A region its owner allocated in this session holds no KV until the owner writes to it, and the owner writes only after a successful map reply. Until the first one, its maps carry fresh=true, retries included; after it (an agent restart, a remap) they do not. The agent takes fresh only with the region's own handle and only on the map that maps the region there first, and passes it to RegionExtension.mapped(fresh=) before the reply leaves, so a device can drop the region's stale contents before any WRITE reaches it. If mapped raises, the map fails and the region stays out of use until a later map succeeds. An agent that does not know fresh ignores it, as it ignores any map field it does not read.
A design page on read hints: who owns what between Maru, the remote CXL agent and a region extension, and one request's hint, load and end. The configuration reference lists read_hints, maru_read_hints, maru_read_hint_on_arrival and the agent's --region-extension options.
Some memory behind an agent is a write-back cache that loses what it held on a reset without a reboot of its node. Maru already treats a changed agent boot as a reboot: owners give their regions back and readers load misses. A region extension may now tell the memory's generation (optional memory_generation(), local state, called outside the agent's lock), and the agent advertises "<node boot id>:<generation>" as its boot, the boot id alone without the hook. The agent reads the generation at start (a failure stops it) and on every hello, map and sweep. When it changes, every mapping is deregistered from the NIC, so RDMA through an old registration fails, and dropped as a release does; from then on hello and map replies carry the new boot. Map replies now say the boot, and a worker takes it over its connection's. A generation the extension cannot tell later keeps the old one.
…loses A hint's sender may close its connection without ending its hints. The agent now calls the extension's optional connection_closed(conn) for every connection that closes, after it dropped the connection's holds and outside its lock, so the extension can end what that connection's hints started instead of waiting for its own expiry. An exception is logged and ignored.
…an unknown memory generation From the review of the region extension hooks: - hint: the agent checks a hint under its lock and hands it to the extension after, so no other request waits for it. A hint about a region that went in between is advice about nothing. - A connection may keep at most --max-hint-requests (4096) hint requests open; read_hint opens one, read_hint_end or the connection's end closes it, and a hint the extension fails leaves none open. More are refused with an error that says so. - The memory generation is read one caller at a time, on a dedicated thread every --memory-generation-poll-s (100 ms) as well as on hello, map and sweep, so a caller that read earlier never sets an older value and a reset is acted on within the extension's detection time plus that interval. - A generation the extension does not tell for --memory-generation-grace-s (2 s) is unknown: every mapping is deregistered and dropped, the boot reads <boot id>:unknown and maps are refused until it answers. The boot that follows carries the count of unknown periods, so nothing mapped before one is trusted after it. - A map whose boot changed while it ran (a reset, or an unknown generation) fails and its registration goes, instead of answering with the new boot for memory the extension saw before the reset. - extension.py states the latency contract (retire, hint and memory_generation answer from local state within microseconds) and that fresh is the owner's word, to be used only on the region's own bytes.
From the review of the worker-side read hints: - The hint thread keeps at most 4096 hinted requests open and ends the oldest past that, so a request whose end never comes cannot fill the agent's per-connection limit. The vLLM scheduler bounds the requests it keeps hint state for (65536) the same way. - The arrival hint wraps vLLM's EngineCore.preprocess_add_request only if it takes (self, request), and passes on only a result whose first item is a request with an id; otherwise hints on arrival are off with a warning. - The data path forgets which own regions had a map reply once they go back. - The hint thread's lookups share the MaruServer RPC client with the scheduler's: a test runs both at once, and a comment says why that is safe and that a hint probe refreshes the keys' last use as a lookup does.
… ~<n> After the n-th period in which the memory generation was unknown, the agent advertises <boot id>:<generation>~<n>, a new boot even when the extension tells the same generation as before: a reset while it could not tell would not show in the generation. The test now covers two such periods with an unchanged generation.
1 of 4 tasks
…'s pages Without them, a store's RDMA WRITEs into memory with a slow tier behind a cache make the device fill each written page from the slow tier first, page by page: on a CXL-SSD stores ran several times slower than with the pages prepared, and landed late enough for the next turn's lookups to miss. - write_hints (vLLM: maru_write_hints, which follows maru_read_hints unless set) makes a store, once it has its pages and before it posts the WRITEs, queue a write_hint on the hint thread: the pages, in write order, under a token for the store. Fire and forget; no end. With vLLM only the worker sends them (it stores; the scheduler hints loads). - The agent takes write_hint with the read hints' checks; it keeps no request open for it. A region extension says which hint actions it takes (optional hint_actions; read_hint and read_hint_end without it), the hello advertises those, Maru sends only those and the agent refuses others. A write hint gives no leave to change the pages. - The hint thread class is MemoryHints now that it sends both kinds.
The memory side keeps a hinted request's bytes accounted until the hint ends. Ending it at request_finished kept that for the whole decode, so the memory side's cap filled with requests already loaded. The scheduler now ends the hint when vLLM reports the request's load done (finished_recving), with request_finished as the fallback; a hint is ended once.
The memory side brings hinted requests' bytes in in hint order, which is not the order their loads run in, so loads read chunks still being brought in while others' wait ready. A worker now tells the memory side when the load of a request starts. - read_promotes (vLLM: on in the worker with maru_read_hints and maru_async_load) makes MaruHandler.promote_read(request_id) queue a read_promote on the hint thread: no ranges, no end, fire and forget. The vLLM worker calls it once per remote_cxl load, before its first batch is read. - The worker did not send the request's hint (its scheduler did, on another connection), so it sends the promote to every agent of the pool whose extension takes it; the extension matches it to a hint by request id alone. - The agent takes read_promote like read_hint_end and keeps no request open for it; it advertises and accepts it only if the region extension lists it in hint_actions.
A queued write hint reaches the memory side about when the store's WRITEs do, so most WRITEs still landed on pages the device had to fill from its slow tier first. Stores run off the request path (vLLM's write-behind thread), so they can afford a short wait. - When an agent's region extension takes write_wait, a store sends its pages for that agent in a write_wait (the store's token, the ranges in write order, timeout_ms) on the store's own connection to the agent and waits for the answer, then posts the WRITEs whatever it is. The waits of one store share write_wait_ms (MaruConfig, default 50, 0 never waits; vLLM maru_write_wait_ms); an agent that does not answer costs the reply timeout on top. Other agents still get a queued write hint. An agent that cannot be reached is not asked for 5 s. - The agent checks a write_wait like a write hint, plus timeout_ms (1 to 10000), keeps no request open for it, and calls the extension outside its lock on the requesting connection's own thread; the extension's answer (a map, or None) is the reply. Only that connection waits. It is advertised and accepted only if the extension lists it in hint_actions.
Prefetching on the memory side runs barely faster than RDMA READ, so a load's READs catch up with it and land on ranges still being brought in while later ones are ready. A load can now ask which are ready. - read_wait_ms (MaruConfig, default 0 = off; vLLM maru_read_wait_ms, worker only): when the agent's region extension takes read_wait, a load reads each batch per peer in steps. It sends read_wait with the ranges it still has to READ, in READ order, and timeout_ms; the reply's ready_prefix says how many leading ranges are ready (answered as soon as one is, or 0 after timeout_ms); the load READs that many, at least one. A wait that fails READs the rest at once; without read_wait the batch is one READ as before. Transfer errors propagate as from one READ, so the slot handling is unchanged. - The agent takes read_wait like write_wait: checked ranges, timeout_ms from 1 to 10000, the extension called outside its lock on the connection's thread, its answer the reply, no request kept open, and only if the extension lists it in hint_actions. - Waits move to waits.py (MemoryWaits), with one connection per agent and thread, so a load's read wait never queues behind a store's write wait and neither holds up hints or region maps.
A load waited up to read_wait_ms at every step, so memory that never got ready made a batch of n ranges take n waits. A batch now has one deadline, read_wait_ms from its first wait: each wait asks with what is left of it, and after a wait that times out (0 ready) or fails, or once the deadline passed, the rest of that peer's batch is READ at once.
…nvironment The agent service gets its settings from an environment file, but the region extension name and config were command-line only, so a package that provides an extension could not turn it on for the installed service. Both now default from MARU_REMOTE_CXL_AGENT_REGION_EXTENSION and MARU_REMOTE_CXL_AGENT_REGION_EXTENSION_CONFIG, which a drop-in on maru-remote-cxl-agent.service can set.
A read wait READ only the leading ranges the memory side had ready and asked again for the rest. It kept the READ order, so it moved the wait from the device to the client and added a round trip per step; the device already holds an early READ until its range is in.
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
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.
🤔 Background & Motivation (Why)
With this PR, Maru tells the memory node which KV ranges a request will read before the load starts (a read hint), so an extension on the memory node can prepare them, for example bring them into a front cache. The KV data path is unchanged.
Replaces #124, which gated every RDMA transfer from the GPU node; gating added a round trip per READ, and a transparent device cache does not need it.
Measured with an out-of-tree package for one CXL-SSD (48 conversations × 6 turns, paired against the same remote path with the device's DRAM holding everything): reuse-turn TTFT 1.0038× [0.993, 1.015] over 5 pairs, KV load +1.24 ms, against 1.0041× and +2.04 ms for the earlier per-READ gating prototype. Two findings shaped the hooks below: writes into pages the device has not prepared are filled page by page from the SSD (store bandwidth fell from 10.3 to 1.3 GB/s), so a store names its pages ahead (a write hint) and may wait briefly for them. In serving, turning off only that wait changed neither TTFT nor load time measurably (0.990× [0.957, 1.024], 3 pairs, write hints still on); it bounds the case where preparation lags behind stores. An early READ needs nothing: the device holds it until its range is in.
🏗️ Design Changes
Grey boxes are Maru; the orange box is an installed package, selected by name; the cylinder is the memory device.
flowchart LR subgraph gpu["GPU node"] sch["Maru (vLLM scheduler)"] wrk["Maru (vLLM worker)"] end subgraph mem["memory node"] ag["remote CXL agent<br/>(checks every hinted range)"] ext["region extension<br/>(maru.region_extensions)"] dev[("memory with a front cache<br/>e.g. CXL-SSD")] end sch -->|"read hint, read hint end"| ag wrk -.->|"read hint at load start<br/>(clients without early hints)"| ag ag --> ext ext -.->|"prepare ranges"| dev wrk ==>|"RDMA READ / WRITE (unchanged)"| dev classDef maru fill:#e8e8e8,stroke:#555,color:#111 classDef pkg fill:#fbe3cc,stroke:#c06000,color:#111 class sch,wrk,ag maru class ext pkgmaru://) can hint each load just before its READs. Off unlessmaru_read_hintsis set; hints go only to agents that advertise support, so without an extension nothing changes.maru_write_wait_ms(default 50) for them to be prepared, then writes anyway. Stores run on the write-behind thread, so this adds nothing to TTFT.fresh. The extension can then reset that region's old contents before the map completes, so the reset finishes before the owner's first WRITE can reach the region.📝 Implementation Details
maru_handler/storage/remote_cxl/hints.py(background thread, 1024 queued hints, ends never dropped, at most 4096 ranges per message, at most 4096 open requests),client.py(hint_read,end_read_hint, load-start hints),maru_server/replica_directory.py(probewithlocations: true). vLLM:maru_read_hints, opt-inmaru_read_hint_on_arrival(wrapsEngineCore.preprocess_add_request; installed only if its signature matches).--region-extensionand--region-extension-configalso come fromMARU_REMOTE_CXL_AGENT_REGION_EXTENSIONandMARU_REMOTE_CXL_AGENT_REGION_EXTENSION_CONFIG, so a package that provides an extension can turn it on for the installed agent service with a drop-in.maru_remote_cxl_agent/extension.py(validate_access,mapped(fresh=),hint,retire,close, optionalmemory_generation,connection_closed; extension calls must answer from local state without blocking on the device).service.py: hint checks under the lock, extension call outside it; per-connection limit--max-hint-requests(default 4096); generation watcher--memory-generation-poll-s(default 0.1 s), fail-closed after--memory-generation-grace-s(default 2 s) (boot<boot_id>:unknown, maps refused, then<boot_id>:<gen>~<n>). hello advertisesregion_extension: {"version": 1, "actions": [...]}; map replies carryboot.region.pyandlifecycle.pyare unchanged;datapath.pyadds onlyfresh. Older agents ignorefreshand advertise no hint support, so they get no hints.docs/source/design_doc/read_hints.md,docs/source/api_reference/config.md.✅ Tests
🔗 Related Issues (optional)
📦 Release Note (for auto-generation / write in English)
NEW
maru_read_hints,maru_write_hints, opt-in), store waits (maru_write_wait_ms), load promotes, and a region extension point on the remote CXL agent (maru.region_extensions) with fresh-region maps, memory generations and connection-close notices.CHANGED
FIXED
IMPORTANT NOTES