Repository navigation
feat: add the remote_cxl storage backend (CXL memory of memory nodes over RDMA) - #107
Conversation
…for remote workers A region mapper runs on every memory node of a pool, shared by every app. Handlers ask it to map a region (map), to drop their hold (unmap) and, as the region's owner, to deregister it before giving it back (release). It checks the handle with the resource manager (GET_ACCESS), maps the region from the node's DAX device, registers it with the node's NIC through NIXL and replies with the address, the key and NIXL metadata of that region only. Release deregisters the region whoever else holds it and refuses later maps of that handle, so an RDMA WRITE still in flight from the old owner fails at the NIC instead of landing in the region's next owner. This holds only when the NIC checks keys and the deregistration reaches it: the region mapper runs UCX on rc/dc only, with the registration cache off, and keeps a released address range unmapped (no access) for 600 s, since the NIC hands out the same key when the same range is registered again. maru_remote holds the Handler side: the region map client and the NIXL transport. pip install 'maru[rdma]' installs NIXL; the daemon is maru-region-mapper.
An RDMA worker opens a storage session as a cpu worker does. maru-server allocates its regions from the resource manager, as request_alloc does for a cxl Handler, records the keys the worker stores in them and gives one more region when the worker's regions are full (expand). Other sessions of the same namespace and KV layout find those keys and read them under a lease. A worker gives a region back (return) only after the region's region mapper released it; maru-server then stops finding its keys, ends its leases and frees the region at the resource manager. Victims and remove act only on keys of the caller's own regions. maru-server takes the pool's region mappers with --region-mapper and hands them to workers in the handshake, as it hands out the RM address; it never calls them. It advertises rdma_storage_v1 when it has CXL and region mappers.
MaruHandler(storage_backend="rdma") keeps KV in the CXL memory of memory nodes and moves it with one-sided RDMA through a host staging buffer registered with the worker's NIC. It gets its regions from maru-server's storage session, asks each region's region mapper for the address and key (map), and reads and writes the region directly; the bytes pass through no Maru process. Before it gives a region back, the worker waits for its WRITEs to that region (bounded) and asks the region mapper to release it, then returns only the regions whose release was confirmed. A region skips the release only when the worker never mapped it and the region mapper holds no mapping for it: no map reply, the latest map refused or failed, or the reply came on a connection that has since been reopened. A region mapper that does not answer is retried later; regions that cannot be released stay out of use. Several MaruServer URLs (comma separated) give one pool each. A request stays in one pool, picked by rendezvous hashing of its first chunk key, and a pool whose MaruServer cannot be reached is retried in the background while its prefixes miss.
The rdma backend uses the storage session contract of the cpu backend (synchronous chunkwise loads and stores), with maru_pool_size as the size of one region and the settings maru_rdma_device, maru_rdma_staging_size and maru_rdma_timeout_s. The staging buffer defaults to two max_model_len requests (at least 1 GiB), and an explicit size that cannot hold one such request is rejected at start. A configuration error stops the worker instead of being retried. Every storage-session call of a request carries its first chunk key as the route key, so a request is stored and found in one pool.
examples/vllm/rdma runs the region mapper, maru-server and vLLM engines with storage_backend "rdma", on one host (rdma_example.sh) or across nodes (the step-by-step scripts). run_simple_query.sh and run_benchmark.sh send a new prompt to one engine and the same prompt to the other, and report cached tokens, TTFT, pool hits and whether both answers match. The vLLM integration guide and the examples index describe the deployment.
… rdma The rdma staging buffer is registered with CUDA, so the connector now lets the rdma backend use the copy kernel and maru_async_load / maru_async_store: - A load returns its staging slots only after its GPU copy has finished. If the copy fails, the slots go back once the queued copies have finished, and stay taken if that cannot be waited for. - With asynchronous loads or stores, stores leave half of the staging slots for loads (rdma_load_reserve); a store that finds no slot skips the chunk for the step and logs at most once a minute. - An rdma request is loaded in batches that fit the free load slots: each batch is read under a lease, copied and released before the next one. A failed batch fails the request from that batch on, and an asynchronous load reports the request done only after its last batch. - The loader thread copies on its own stream without waiting for the compute stream, as the cxl asynchronous load does, and waits up to 1 s for a slot another load holds. - A partial chunk takes the per-layer copy; the kernel needs a whole slab. - Synchronous rdma loads with vLLM async scheduling on are rejected at start: a failed synchronous load is reported after vLLM has already scheduled the request's next step. The scheduler now puts the external token range on asynchronous loads, so they fill only the blocks vLLM allocated for them, and does not offer the same hit again to a request whose asynchronous load failed (a preempted request gets it again). Asynchronous stores pass the request's route key, so with several pools they allocate in the pool of their request. cpu, mixed and cxl paths are unchanged.
MARU_ASYNC=1 on the launcher turns on maru_async_load and maru_async_store. The example README and the vLLM integration guide describe the copy kernel over the staging buffer and when vLLM async scheduling must be off.
An engine gives its regions back when vLLM shuts it down. When the stop signal reaches the engine process itself, as Ctrl-C does for every process of the terminal's foreground group, or a kill of the whole process tree, the engine can end before its regions go back, and maru-server keeps them until it stops. On the testbed, a SIGTERM or SIGINT sent to an engine's process group left its regions allocated, while the same signal sent to the API server alone returned them. The launcher now runs vLLM in its own session and passes Ctrl-C and TERM on to the API server only. rdma_example.sh stops each engine by signaling its launcher and waits up to 30 s before killing what is left. The README says how to stop an engine.
…gnal Under set -e, the wait interrupted by a trapped signal returned non-zero and ended the launcher before vLLM had exited, so a caller that stopped maru-server next could do so before the engine gave its regions back. The launcher now keeps waiting until vLLM exits, and also passes HUP on.
a637f40 to
3adb7cf
Compare
…he RM The region mapper is a platform service like the resource manager, one per memory node, but nothing installed it: whoever deployed the rdma backend had to start it by hand. install-maru-resource-manager (run by install.sh) now also installs maru-region-mapper.service when NIXL is installed. The unit is wanted by maru-resource-manager.service and part of it, so starting, stopping or restarting the resource manager does the same to the region mapper; the resource manager's own unit is unchanged. Settings come from /etc/maru/region-mapper.env, written once and kept across reinstalls: the resource manager address (127.0.0.1:9850), the bind address and port (0.0.0.0:9860, so workers of other nodes reach it) and the RDMA NIC (UCX default). The region mapper's options now take their defaults from these environment variables, and --rm-address defaults to the local resource manager. install-maru-region-mapper installs or removes only the service; --no-region-mapper leaves it out of the RM install.
The README now has memory nodes run the region mapper as the service that starts with the resource manager, with region_mapper.sh for hosts without it. rdma_example.sh uses the service when it runs on the host instead of starting a second region mapper.
It compares a computed answer with one that reuses stored KV, which vLLM itself can make differ; a difference is not a load failure.
… guide Worker nodes (MaruKVConnector, MaruHandler with RdmaStorageClient and RdmaDataPath), the metadata node (maru-server) and the memory nodes of a shared pool (resource manager, region mapper service), with the RPCs, the NIC registration and the RDMA path between them.
…de setup The vLLM guide shows the cxl and rdma backends taking the same steps, with only the mapping step moved to the memory node. The guide and the RDMA example now say how a memory node of a shared pool that runs no resource manager installs the region mapper service, and that the resource manager then listens on a network address as in the multi-node configuration.
UD has no one-sided RDMA, so UCX would carry RDMA over it in software, where a released region's key is not checked by the NIC. UCX_TLS may now name rc and dc transports and self; UD passes only with ":aux", which UCX uses to set up rc connections. "ib" is refused because it enables UD for RDMA too.
region_mapper.sh cannot take the service's node lock under /run/maru, so it checks for maru-region-mapper.service and refuses to start beside it. The component figure marks parts of a new box as new, names RegionMapService, and numbers the shared pool's memory nodes 1 and 2; the concept figure names the address and NIC key the region mapper returns.
On a memory node of a shared pool that runs no resource manager, the service wanted by maru-resource-manager.service never started at boot. install-maru-region-mapper --standalone installs a unit wanted by multi-user.target instead, so that node's region mapper starts at boot on its own. The default unit is unchanged.
…nimum Loads go in batches of the free staging slots, so the minimum size the connector checks comes from stores, which put every chunk a step completes into staging before writing them.
…etrying A standalone service runs on a node without a resource manager, but it was given the settings file that points at a local resource manager and says the service starts with it. It now gets settings with an empty MARU_RM_ADDRESS to fill in, and the region mapper refuses to start with an empty resource manager address. The unit stops retrying after 5 failed starts in a minute, so a node without an RDMA NIC does not restart it forever.
The backend keeps KV in the CXL memory of memory nodes; RDMA is only how a worker without CXL reaches it. Backend names say where the KV lives (cpu, cxl, mixed), so the backend is now maru_storage_backend="remote_cxl", with its package maru_handler/storage/remote_cxl, RemoteCxlStorageClient, the remote_cxl_storage_v1 handshake capability, examples/vllm/remote_cxl and the matching tests and docs. Names of the RDMA transport stay: RdmaDataPath, the maru_rdma_device, maru_rdma_staging_size and maru_rdma_timeout_s settings, and the maru[rdma] extra that installs NIXL.
Adds the return rule, the shared pool's region-to-node choice, several pools and the staging slots as figures, and names RemoteCxlStorageClient in the component figure.
youngrok-XCENA
changed the title
Only synchronous loads need --no-async-scheduling. With MARU_ASYNC=1 the launcher now leaves vLLM's default on; two engines on another node loaded each other's KV with it on, with answers equal to vLLM's prefix cache.
systemctl enable adds the links of the unit's current [Install] section and keeps any others, so installing the standalone form and then the RM-bound form left the service wanted by multi-user.target as well. The installer now disables the service before enabling it.
Stores leave half of the free slots to loads; which slots they take is not fixed, so the figure labels the two shares by what they hold.
…e, keyed region mappers Three changes close the open design tasks of the remote_cxl backend. Keys outlive their writer (task 15). As MaruServer keeps a cxl Handler's regions while keys refer to them, it now takes the regions of a remote_cxl worker that left (drained close or session expiry) over: regions with keys stay readable, the others go back. A RegionKeeper thread in maru-server holds a kept region's mapping at its region mapper, maps it again after a region mapper restart, gives it back if the memory node rebooted or the mapper refuses the handle, and gives back the kept region read longest ago when the resource manager refuses a region request (one at a time). Every region still goes to the resource manager only after its region mapper released it. The owner no longer releases at close; it settles its WRITEs, returns never-mapped regions and closes drained. It still releases when MaruServer restarted. Owners report each region's boot id in heartbeats and the drained close. Readers get a read handle (task 20). acquire gives other sessions' regions as (region id, offset, length, read token), the token an HMAC of the resource manager's auth token. The region mapper maps a read handle only for a region it has mapped, and release needs the region's handle, which only the owner and MaruServer hold. Who may map is bound again (task 22). Every connection to a region mapper proves the region mapper key (HMAC over per-connection nonces, both ways) before it may map, unmap or release. maru-region-mapper, maru-server (--region-mapper-key-file) and workers (maru_region_mapper_key_file) read the key file; the installer creates it. With MARU_REGION_MAPPER_HOST empty (the installed default) the service listens on its RDMA NIC's address.
…nfig docs The architecture overview gains a section on CXL memory of other nodes over RDMA: the same layers and resource flow, the region mapper on memory nodes, and how the data path, capability-based access and region lifetime carry over. The MaruHandler design gets the remote_cxl backend (components, store, retrieve, eviction, giving regions back, staging slots) and the MaruServer design the remote CXL regions (sessions, read handles, the regions of a worker that left and the region keeper), with the figure of a kept region. The installation guide and the configuration reference list the region mapper service, its key and the remote_cxl settings. The failed-deregistration storage test now checks the keeper's path: a region whose release fails at its region mapper does not reach the RM.
…s away A worker that cannot reach one of its regions removes the region's keys, so probes stop reporting hits that cannot be read. MaruServer now does the same for a region it keeps after its owner left: if the region's mapper does not answer for 30 seconds, the region is given back (its keys stop being found) and is released and freed once the mapper answers again. Before, a crashed owner on a node that went down left keys that kept failing to read.
The agent dropped an unhold the resource manager answered with any error, and queued no unhold after any failed hold, but the resource manager keeps the new state in memory when only saving it failed. Only a refusal (HoldRefusedError: a bad token, a region not allocated, a bad name, or an RM without holds) now means nothing is held; other failures keep the unhold queued for the sweep. The component figure labels the agent's resource manager requests HOLD / UNHOLD, and the hold figure carries labels only.
…manager holds The component figure and the MaruServer design said RegionKeeper holds a region at the remote CXL agent, the word now used for the agent's holds at the resource manager.
|
@jooho-XCENA Thanks for the reproductions; both are real. Fixed in 7aee9c1 (resource manager), 88a31fb and 55a6a2a (agent) by gating range reuse on confirmed deregistration, as you suggested.
Why both of your paths were open: this PR picks a region's agent by hashing the region id, so the same range can come back through another agent, and the agent-local fence never sees it. In your local- One part is older than this PR. The resource manager also reclaims a maru-server's regions while cxl Handlers still map them: a cxl Handler keeps taking pages of its region locally (handler.py#L521) and writing through its mapping, which stays until it closes. That needs the reclaim to wait for every mapper, cxl Handlers included, so I opened #115 for it. |
hyunyul-XCENA
left a comment
hyunyul-XCENA
left a comment
There was a problem hiding this comment.
The remote_cxl backend reads solidly — region hand-back ordering, isolated transfers and lease handling all hold up under the paths I traced; what I found left are three small operational issues: a store-skip warning that is rate-limited on one path but not on the default write-behind one, a pool log line that prints the whole comma-separated server URL instead of the failing pool's, and a sweep that stops retrying stuck deregistrations while the resource manager is unreachable.
| if self._remote_cxl_mode and isinstance( | ||
| e, MemoryError | StorageUnavailableError | ||
| ): | ||
| self._note_store_skip(base_key, e) # rate-limited |
There was a problem hiding this comment.
🟡 minor — The new rate limiting does not take effect on the write-behind path, which is the default for remote_cxl (maru_async_store defaults to true). _note_store_skip suppresses its own warning for 60 s, but line 4721 then calls _complete_write_behind_keys([base_key], [False]), and that method logs "Maru write-behind store failed: %s" at warning level for every key it is given. So a staging buffer that is full while loads hold the slots — which the docstring calls expected — still produces one warning line per chunk, which is exactly what the change set out to avoid. The two synchronous store paths do not have this problem: a chunk skipped there never reaches batch_store, so _note_store_skip is the only log. Either pass the skipped key through a variant of _complete_write_behind_keys that does not log, or let _note_store_skip record the key and have _complete_write_behind_keys stay quiet for keys it already reported.
There was a problem hiding this comment.
Confirmed. Fixed in 208ab5e: the write-behind path passes reported=True to _complete_write_behind_keys for a chunk whose failure it has already logged (rate-limited by _note_store_skip, or as a prepare error), which then publishes the outcome without a second line. test_a_full_staging_buffer_skips_the_async_store_quietly now also checks that no per-chunk line is logged; it fails on b553c14.
| except Exception as exc: | ||
| failures.append(exc) | ||
| logger.warning( | ||
| "Remote CXL pool %s is unavailable: %s", c.config.server_url, exc |
There was a problem hiding this comment.
🟡 minor — All pools share one MaruConfig, so c.config.server_url is the whole comma-separated maru_server_url string, not the URL of the pool that failed. With several pools this log line — the one an operator reads to find out which MaruServer is down — names every pool instead of the unreachable one. The same applies to the "is available" line in _retry_loop (L131-133). PoolRouter already keeps the per-pool identity in self.urls, paired one-to-one with self.clients, so iterating with zip(self.clients, self.urls) (or enumerate) and logging self.urls[i] fixes both sites.
| conn: Connection id. | ||
| """ | ||
| with self._lock: | ||
| for rid, m in list(self._maps.items()): |
There was a problem hiding this comment.
🟡 minor — sweep retries the deregistration of holderless live mappings before any resource-manager query (the first block, under the lock), but a mapping that is already out of use (m.error set) and still registered is only retried in the second loop, after _allocated(). When the resource manager cannot be reached, the first _allocated() call returns out of sweep entirely, so those mappings never get their retry — even though retrying them needs no resource-manager query at all, exactly as for the holderless ones. The NIC registration then stays alive for as long as the resource manager is away. Widening the first block's condition to also cover m.error is not None (if m.error is not None or not m.holders:) would give both kinds the same lock-held retry and leave the second loop to the resource-manager check only.
There was a problem hiding this comment.
Confirmed. Fixed in 208ab5e: the first, lock-held block of the sweep now retries the deregistration of every mapping out of use (failed, or with no holder left) before any resource manager query, and the resource manager loop checks only live mappings. test_the_sweep_deregisters_a_failed_mapping_while_the_rm_is_away (a live mapping is checked first and the resource manager does not answer) fails on b553c14.
When the holds file was renamed but syncing its directory failed, HOLD returned the error and kept the hold, but a retry of the same HOLD (or a retried FREE of the draining region) succeeded without saving again, so the agent could register a region whose hold a restart might lose. Such a retry now syncs the directory again and succeeds only once that works.
- The agent opens a device whose UUID it cannot find here only if the RM's path carries that UUID in its header, as MaruShmClient does. - Expired remote_cxl sessions hold nothing and no longer count toward the session limit; the oldest of their records go beyond 4096. - maru-server refuses for good a region larger than the RM's largest pool, and a worker whose handler keeps failing says so at ERROR every 12 tries. - Workers and maru-server read the key copy in ~/.maru by default; the docs install workers with --no-rm. - A region release that raises at the RM is retried like an unconfirmed one. - A write-behind chunk skipped for a full staging buffer is not logged again when it completes. - PoolRouter logs each pool's own URL. - The agent's sweep retries the deregistration of failed mappings before it asks the RM, so it runs while the RM is away.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Lines 3540 to 3542 in 208ab5e
Release CPU leases after deferred copies complete
When maru_async_load is enabled with the cpu or mixed backend, infos contains CpuReadLease objects, but completing the CUDA event only removes _deferred_refs; it never calls handler.release_retrieved(infos). CpuReadLease has no destructor that invokes release(), so every successful deferred load leaves its server lease active, preventing eviction of those keys and eventually reaching the 4096-lease session limit, after which subsequent loads are refused. Release the retained leases when the event completes, including the failure and shutdown paths.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # This map's hold supersedes an unhold of the same handle still | ||
| # waiting for the RM (another handle's stays queued). | ||
| with self._lock: | ||
| if self._unheld.get(h.region_id) == h: | ||
| del self._unheld[h.region_id] | ||
| return self._rm.hold_access(h, self.node_id, self.epoch) |
There was a problem hiding this comment.
When an earlier unhold RPC failed before reaching the resource manager, _unheld is the only record that the old hold still needs removal. If the owner then frees the region, it becomes draining, and a delayed map of the same handle reaches this code: the pending unhold is deleted before hold_access(), which is then definitively refused because draining regions cannot be held. _establish() deliberately does not requeue an unhold for HoldRefusedError, so the old hold remains forever and the resource manager can never finish freeing that region until the agent restarts or an administrator intervenes. Remove the pending entry only after the replacement hold is confirmed.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed. Fixed in a8b8436: _hold now drops a queued unhold of the same handle only after hold_access has succeeded. A refused hold (the region drains, so it cannot be held again) leaves the unhold queued, and the sweep sends it, after which the resource manager frees the range. test_a_refused_hold_keeps_the_unhold_still_queued fails on 208ab5e.
…succeeds A map of a region whose unhold had not reached the resource manager dropped the queued unhold before holding again. If the region had been freed meanwhile, it was draining, the new hold was refused, and the old hold stayed until the agent restarted, so the range was never freed. The queued unhold now goes only once the new hold is recorded.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a8b8436ca8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self._give_back_rebooted(region) | ||
| return | ||
| first = not region.boot | ||
| region.boot = access.boot |
There was a problem hiding this comment.
When /proc/sys/kernel/random/boot_id is unavailable, the agent deliberately reports an empty boot ID, and this assignment leaves region.boot empty permanently. After that memory node actually reboots and the client reconnects, _adopt() cannot distinguish the new boot, so it makes the region reachable again while retaining keys whose completed RDMA writes may have been lost with the old node's CPU caches; subsequent cache hits can therefore inject stale or corrupt KV. Fail closed when the agent supplies no boot identity rather than continuing to use the region across reconnections.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed. Fixed in 72fb137 by failing closed on both sides: the agent refuses to start when it cannot read the node's boot id, and a worker (and MaruServer's region keeper) treats an agent whose hello carries no boot as unreachable, so its regions are set aside and their keys go after 30 s, instead of being used across a reboot it cannot see. test_main_refuses_a_node_without_a_boot_id and test_a_worker_refuses_an_agent_without_a_boot_id fail on a8b8436.
…s boot An agent that could not read /proc/sys/kernel/random/boot_id reported an empty boot, and a worker then could not tell a reboot of the memory node, which loses writes its CPU caches held, from a restart of the agent. The agent now refuses to start without a boot id, and a worker treats an agent that reports none as unreachable.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 72fb137bf7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| with self._lock: | ||
| self._isolated.append( | ||
| IsolatedTransfer( | ||
| exc.pending, "READ", exc.pending.peer, None, [], list(slots) |
There was a problem hiding this comment.
When one read batch spans regions on multiple remote CXL agents, the pending transfer for a timed-out peer can touch only that peer's pairs, but this records every slot in the batch. If that transfer remains unconfirmed, the reaper retains slots belonging to completed or not-yet-issued peers indefinitely, potentially exhausting the entire staging buffer and turning subsequent cache hits into misses. Track the slots alongside each by_peer group and isolate only the timed-out peer's subset.
Useful? React with 👍 / 👎.
hyunyul-XCENA
left a comment
hyunyul-XCENA
left a comment
There was a problem hiding this comment.
The remote_cxl backend is well structured and its safety ordering (release-before-return, RM holds, boot checks) holds up under review; I found no new correctness defect beyond the open comments, only a few small issues in the connector's load pacing, the deferred-load bookkeeping, and two installer/config edges.
| if b > 0: | ||
| capacity = handler.retrieve_capacity() | ||
| if capacity == 0: | ||
| capacity = self._leased_batch_size(handler, n, True) |
There was a problem hiding this comment.
🟡 minor — _leased_batch_size documents that it waits for slots only off the engine thread ("on the engine thread it fails at once rather than hold up the step"), and _load_leased honours that for the first batch (stream is not None). This mid-request refill, however, passes offthread=True unconditionally, so a synchronous load (maru_async_load=false) can sleep up to _LEASED_SLOT_WAIT_S (1 s), in 10 ms increments, inside start_load_kv — i.e. inside the forward pass — every time a chunk boundary finds retrieve_capacity() == 0. Either pass stream is not None here, or drop the 1 s wait on the engine thread and let the batch shrink/fail. I kept this at minor because the window where capacity reaches 0 on a synchronous-load deployment is narrow (stores and loads share the engine thread there, and rdma_load_reserve is 0), so I could not construct a reproducer — but the call contradicts the method's own contract either way.
| start = self._load_starts.pop(request.request_id, 0) | ||
| self._load_ranges[request.request_id] = (start, start + num_external_tokens) | ||
| if self._deferred_loading: | ||
| self._async_load_issued.add(request.request_id) |
There was a problem hiding this comment.
🟡 minor — _async_load_issued is added for every deferred load, not only for loads that failed, and it is only cleared on request_finished or on the preempted/finished sweep in build_connector_meta. Any other path that makes vLLM call get_num_new_matched_tokens again for the same request id while it is still live — e.g. a request that is unparked after a successful load but cannot be scheduled and goes back to WAITING — then loses its (valid) hit and logs req %s returned after a failed load, which is untrue for that request. The record is a "load was issued" flag being read as a "load failed" flag. Consider setting it only when _fail_deferred_load actually reports the request, or clearing it when get_finished_loading reports the load as successful, so the misleading warning and the lost hit are limited to real failures.
|
|
||
| # Key that workers and maru-server prove on every connection. Copy this file | ||
| # to them, readable only by the users that run them. | ||
| MARU_REMOTE_CXL_AGENT_KEY_FILE=/etc/maru/remote-cxl-agent.key |
There was a problem hiding this comment.
⚪ nit — install() takes key_file as a parameter and writes the key there, but DEFAULT_ENV and STANDALONE_ENV hardcode MARU_REMOTE_CXL_AGENT_KEY_FILE=/etc/maru/remote-cxl-agent.key. With a non-default key_file the generated env file points the service at a key that was never created, and the agent exits at startup with a KeyFileError. Interpolating key_file into the template would keep the two in step. Today only the tests pass a custom path, so the impact is limited to future callers.
| self._config.server_url, | ||
| timeout_ms=self._config.timeout_ms, | ||
| ) | ||
| self._rpc = self._make_rpc(self._config.server_urls[0]) |
There was a problem hiding this comment.
⚪ nit — self._config.server_urls[0] raises IndexError when server_url is empty or contains only commas/whitespace, because server_urls filters out blank entries. The previous code passed server_url straight to RpcClient and failed later with a connection error naming the URL. MaruConfig.__post_init__ does not validate that server_url is non-empty, so a blank value now surfaces as a bare IndexError from the constructor. A check in __post_init__ (or server_urls or [self.server_url]) would keep the error message actionable.
| return out | ||
|
|
||
|
|
||
| def _rdma_staging_size(extra: dict[str, Any], page_size: int) -> int | None: |
There was a problem hiding this comment.
Is a request-sized store buffer needed on remote_cxl?
The store path defers chunked-prefill steps and gathers the whole prefix into staging at the prompt-completing step (L4662–4669, from #66). On cxl that costs nothing, since the store target is the CXL page itself. On remote_cxl it means pinned host memory proportional to max_model_len: 16 GiB on the store side for Llama-3.1-8B at 128k.
Could remote_cxl stream stores instead, writing the chunks each step completes? Completed chunks don't change and keys are prefix-chained, so partial stores are valid, and request_finished already holds the GPU blocks until batch_store finishes, so it needs no extra GPU memory. Staging would then be sized by bandwidth × latency rather than max_model_len. Since #66 coalesced to avoid contending with prefill, an A/B on TTFT/TPOT would show whether that still holds here; cxl could stay as is.
There was a problem hiding this comment.
Agreed that the request-sized store buffer is expensive. #107 already introduces a large backend/lifetime change, so I will take streamed per-step stores and bounded store staging to a separate follow-up PR rather than expand this one further.
That PR should compare TTFT/TPOT and peak registered host memory against the current coalesced store, and retain publish-after-WRITE and GPU-block lifetime guarantees. #107 keeps the current request-sized budget; this is a scope split, not a claim that the memory cost has been optimized.
| def load_capacity(self) -> int: | ||
| """Slots the next load batch may take: the free slots, at most the | ||
| load reserve when there is one.""" | ||
| with self._free_lock: | ||
| free = len(self._free) | ||
| return min(free, self.load_reserve) if self.load_reserve else free |
There was a problem hiding this comment.
Could you share the reasoning for the fixed 50:50 load/store split?
Loads are capped at half the slots even when no store is running, so the default has to hold two max_model_len requests. But the two sides need very different amounts: loads are bounded by the pipeline depth (2 × _LOAD_PIPELINE_CHUNKS = 16 slots, 512 MiB for Llama-3.1-8B) and hold slots briefly, while stores hold slots until the WRITE completes.
Would separate load and store pools (two free lists over one registered mapping) be easier to manage? Each could be sized for its own need (32 GiB → ~16.5 GiB with no other change), neither side could starve the other, and the load_reserve logic would go away.
There was a problem hiding this comment.
Agreed; I will address independently sized load/store budgets in the same follow-up PR as bounded/streamed stores, keeping that redesign out of the already large #107. The current implementation uses one free list with a store reserve threshold and a load-capacity cap, rather than two physical buffer halves, but the default budget is still larger than the load pipeline needs.
The follow-up should size loads by the actual pipeline/concurrency requirement, size stores by their retained work, and test simultaneous load/store progress, starvation and slot reuse as well as memory usage and TTFT/TPOT. #107 retains its existing budget policy.
| # Only a confirmed terminal state permits releasing the handle. | ||
| # A post/query exception may leave a live RMA behind even when | ||
| # release_xfer_handle would return success. | ||
| released = state == "ERR" and self.release_handle(op, peer, handle) |
There was a problem hiding this comment.
Do the staging slots of a failed transfer ever come back?
released requires state == "ERR", but NativeNixlAgent._state() raises instead of returning "ERR", so after an RDMA error the transfer is never released and PendingTransfer.poll() keeps failing. Its staging slots then stay quarantined for the life of the worker, and repeated errors gradually shrink the buffer until stores are skipped and loads have no room.
Is there a recovery path I'm missing? If not, could the slots be reclaimed at a definite fence point, e.g. after the endpoint to that peer is closed?
There was a problem hiding this comment.
Confirmed: with the native binding, a persistent transfer-status exception can keep the pending handle and its staging slots retained. 9674da9 only returns unrelated peers' slots; it does not make the failed peer's own transfer recover automatically.
I will handle safe in-process reclamation in a separate recovery PR, so #107 does not grow into an endpoint/fencing redesign as well. That PR needs a proven DMA-quiescence boundary and persistent-native-error/late-DMA tests before reclaiming memory; a timeout or an endpoint-close call alone is not sufficient evidence.
For #107, the current boundary is fail-closed retention: an affected worker may lose cache capacity and require a worker-process restart. Retrying or closing the Handler can leave the registered buffers retained; restarting only the memory-node agent is not an established recovery guarantee. The operator guidance and restart-recovery validation need to make this limit explicit before merge. I am deferring automatic reclamation, not claiming this path is fixed or resolving the question by moving it; I will leave this thread open for agreement on that boundary and the follow-up scope.
hyunyul-XCENA
left a comment
hyunyul-XCENA
left a comment
There was a problem hiding this comment.
The remote_cxl backend is coherent end to end and the release/hold ordering that keeps late RDMA out of a reused region holds up under reading; I only found three small issues (an RPC issued under the client lock, one inconsistent offthread flag, and a world-writable lock path in the example).
seohui-XCENA
left a comment
seohui-XCENA
left a comment
There was a problem hiding this comment.
Rechecked my three review comments against 9674da9. The concerns are addressed:
- Refusals: an oversized region produces a permanent, actionable error; repeated transient refusals log at ERROR every 12 failures with guidance about engine IDs and pool capacity. Expired remote_cxl sessions no longer consume the active-session quota.
- Pool sizing: the existing cxl path also uses maru_pool_size for the initial allocation and, by default, subsequent growth. Keeping the shared name is reasonable, and the updated documentation explains it consistently. The server now rejects a region larger than the largest eligible pool.
- Worker installation and keys: the installation guide and worker example explicitly use ./install.sh --no-rm. Worker and maru-server key defaults now point to ~/.maru/remote-cxl-agent.key, separate from the agent's /etc/maru key. Keeping agent installation coupled to the RM is acceptable with these changes.
Validation on this commit: 285 targeted directory/agent/installer/storage unit tests and 29 vLLM remote_cxl configuration tests passed (314 total). Current CI checks are also green. Approving based on this follow-up review.
🤔 Background & Motivation (Why)
Maru shares KV through CXL shared memory, so only nodes attached to the CXL device share it today; a GPU node without CXL recomputes a prefix that an engine on another node has already computed.
This PR adds the remote_cxl backend: GPU nodes without CXL store and load KV in the CXL memory of memory nodes with one-sided RDMA. The metadata server manages keys and region ownership; the memory-node agent maps and registers regions, and the NICs move their payloads. The README and docs now explain both directly attached and RDMA-accessible CXL pools with the figure above. Earlier two-node TTFT measurements from
0b71898are preserved below as historical results; they are not measurements of the current implementation.🏗️ Design Changes
Peer-scoped READ isolation follow-up (
9674da9). When one peer times out, has an uncertain transfer state, or is interrupted after posting, retain only that peer's destination slots. Return completed and unposted peers' slots immediately; keep the uncertain transfer protected until its completion and handle release are confirmed. This fixes the over-isolation reported in the review without changing the staging allocation or payload path.Validation: 15 new fault-injection regressions cover three peer positions, two interleaved slots per peer, timeouts, uncertain state queries, interrupts before/after posting, safe reuse during a late READ, and exactly-once slot return. Local unit tests: 1,681 passed, 30 skipped; targeted integration: 151 passed, 3 skipped, 6 deselected (
TestAllocFreeexcluded). Ruff, changed-file formatting and diff checks pass.mypy .stops at the existing duplicaterun_benchmarkexample module names. These are local/fake-based regressions; the hardware and TTFT evidence below remains tied to the pre-follow-up source11/1aca0cd, rather than a device remeasurement of this commit.Architecture. GPU worker nodes, which have no CXL, run vLLM with MaruHandler. A metadata node runs one maru-server per app per pool. The memory nodes' CXL memory forms the RDMA memory pool: one memory node runs the resource manager, and every memory node runs a remote CXL agent. The agent is needed because one-sided RDMA reaches only memory that a process on the memory node has mapped and registered with its NIC, and a worker cannot map CXL memory it is not attached to. The agent maps a region R, registers it with the NIC and returns R's address and remote key; the worker's RDMA then moves the bytes NIC to NIC, through no process on the memory node. Regions flow as on the cxl path (resource manager → maru-server → Handler); only the mapping moves from the Handler to the agent.
Components. Each worker keeps requests within their selected pool and shares one staging/data path across its pool clients. The metadata node tracks keys, leases and region ownership, including regions retained after a worker leaves. Each memory-node agent checks access with the resource manager, maps and registers regions, and releases those registrations before their ranges can be reused.
Progress and CPU use. Worker transfer progress continues independently of the GPU's execution. When no work is ready, it waits for NIC events with a periodic fallback. The memory-node agent drives connection setup for a bounded period and performs periodic housekeeping. This avoids idle busy-poll loops while keeping completion and buffer ownership separate from event-wait timeouts.
Who sets what.
Rules that keep KV correct.
Figures: concept, deployment, the rules
Other structures considered
📝 Implementation Details
Components. In the worker,
PoolRouterpicks a request's pool when the app uses several pools, and that pool'sRemoteCxlStorageClientruns the storage session (regions, pages, keys, eviction); oneRdmaDataPath, shared by all pools, holds the staging buffer, NIXL and the connections to the agents. In maru-server,ReplicaDirectorykeeps the storage sessions, regions, keys and leases and gets regions throughAllocationManager, andRegionKeeperholds the regions of workers that left. In each agent,RegionMapServertakes the connections and checks the agent key, andRegionMapServiceholds the region at the resource manager (HOLD, which checks the handle;UNHOLDonce it is deregistered), maps it withDaxRegionMapperand registers it with NIXL.Progress and CPU use. A worker uses NIXL's public native API with progress enabled and a 1 ms event-wait fallback. The C++ thread progresses and arms UCX, then waits on its event descriptors; a ready event can wake it before the timeout. The posting thread also progresses UCX when it queries an active transfer. The memory-node agent keeps background progress OFF, drives connection setup for a bounded window, and drains notifications on map/hello/sweep. Neither side needs an idle busy-poll loop.
Transfer ownership and synchronization. Post/query exceptions and timeouts retain pending handles. Metadata loading checks the affected pending transfers under the same transport lock that creates new transfers, and drops the lock while waiting for them to finish. The client retains connection probes through late-transfer cleanup, rejects new read leases after close starts, and distinguishes retained incomplete resources from successful startup.
Files
maru_remote_cxl_agent/(new, called the region mapper in earlier revisions of this PR: the service, its TCP server with the key proof, DAX mapping and NIXL registration after the resource manager'sHOLD, UCX checks, installer),maru_resource_manager/(region holds:HOLD,UNHOLD,RESET_HOLDS, theholdsstate file,maru-rm-ctl drop-holds),maru_shm/(get_access(),hold_access(),unhold(),reset_holds()).maru_server/replica_directory.py(remote_cxl storage sessions, cross-session leases with read handles, eviction victims, taking over the regions of workers that left),maru_server/region_keeper.py(new),maru_server/server.py(--remote-cxl-agent,--remote-cxl-agent-key-file; usage of remote_cxl regions),maru_tools/usage.py.maru_remote/(remote CXL agent client, key files, NIXL transport and a minimal public-native worker adapter),maru_handler/storage/remote_cxl/(RemoteCxlStorageClient,RdmaDataPath, region lifecycle,PoolRouter),maru_handler/storage/staging.py,maru_handler/handler.py,maru_common/config.py,maru_vllm/connector.py(remote_cxl settings, batched and asynchronous loads and stores on the staging buffer).examples/vllm/remote_cxl/(its sharing check compares the answer from the pool with vLLM's own prefix cache, since vLLM itself can answer differently when it reuses a prefix),docs/(vLLM integration, architecture, MaruHandler and MaruServer designs, installation, system setup, configuration),README.mdanddocs/index.md(the figure above, a title and opening that define both kinds of memory pool, with the zero-copy, fan-out and energy claims scoped to nodes attached to CXL, and a vLLM section in the README).Settings
maru_storage_backend="remote_cxl",maru_pool_size(the CXL memory an engine gets at a time, as oncxl; one region),maru_engine_id,maru_cache_namespace,maru_remote_cxl_agent_key_file, and optionalmaru_rdma_device,maru_rdma_staging_sizeandmaru_rdma_timeout_s; a comma-separatedmaru_server_urlgives several pools.maru_async_loadandmaru_async_storedefault to true on this backend; setting them to false runs loads and stores on the engine thread. Engine requirements are those of thecpubackend, except that asynchronous loads and stores are accepted; with vLLM's asynchronous scheduling on, asynchronous loads are required.storage_backend="remote_cxl"with the matchingrdma_*andremote_cxl_agent_key_filefields;alloc,batch_existsandbatch_retrievetake an optionalroute_key, which other backends ignore.--remote-cxl-agent host:port(repeated for a shared pool) and--remote-cxl-agent-key-file.install.shalso installs the remote CXL agent service and its key;install-maru-remote-cxl-agent --standaloneserves a memory node without a resource manager. The native transport requiresnixl>=1.2.0,<2, matching CUDA-wheel versions and the required public APIs. For the validated vLLM 0.22.1 / Torch 2.11 / CUDA 13 serving stack, use the example constraints selecting NIXL 1.2.0.cpu,mixed,remote_cxl): the model path and the node-local config fields (_name_or_path,name_or_path,transformers_version,_commit_hash) no longer enter it, so engines that keep the same weights at different paths share KV; the storage backend enters it, so engines of different backends never share a namespace.cpuandmixedcopies do not survive a restart, so nothing stored is lost.Deployment commands (memory node 10.0.0.2 with NIC
mlx5_1, metadata node 10.0.0.3, workers' NICmlx5_0; the hardware runs below put maru-server on the memory node)Avoid duplicate stores after a hit. After every READ and GPU copy of a request succeeds, the connector skips gathering and allocating full chunks that this request already loaded from the external pool. Failed loads and partially loaded boundary chunks remain eligible for storage. The record ends on finish, preemption or abort; it does not become a global assertion that a key still exists. A subsequent eviction can cost that request one cache refill, but does not invalidate the bytes already on its GPU. Other backends and chunked-prefill coalescing retain their existing policy.
Points for review.
UCX_TLS=rc,dc,cuda_copywith the registration cache off;cuda_copysupplies local memory-type support and is accepted only alongside RC/DC.HOLD,UNHOLDandRESET_HOLDSand keeps holds inholdsin its state directory. The agent and the resource manager are upgraded together: an agent refuses to map against a resource manager without holds. A memory node gone for good keeps its held ranges untilmaru-rm-ctl drop-holds NODE --confirm. The resource manager still reclaims regions that cxl Handlers map without waiting for them; that is older than this PR ([Bug] resource manager: a reclaimed region's range is handed out while cxl Handlers still map it #115).✅ Tests
The tested execution snapshot is source11, compared with public
72fb137againstmain@9a56e39. Final commit1aca0cdc61cf3ee4ccf9115e1653b573e92e6f0dhas the same executable runtime AST as source11; subsequent runtime-file changes are docstrings, with documentation and serving constraints added separately. Source11 adds a startup/partial-cleanup ready-state guard to source10. The PR remains a draft. Final-commit CI status: GitHub Actions Ruff, Python 3.12/3.13/3.14 unit tests, Resource Manager and aggregate unit-test checks passed on this head; Jenkins build 79 skipped tests under draft policy, and its external PR-head context remains PENDING. Jenkins build 79 for this head wasNOT_BUILTbecause draft validation was skipped, with zero downstream builds; the external PR-headPENDINGcontext is not a running build. Completed local validation is separate from approval to merge.A subsequent test-only change removes a race in the notification-drain assertion. The executable runtime and the source11 device/A/B measurements are unchanged.
TestAllocFreecases that launch their own resource manager were explicitly excluded; this is not a claim that those six passed.mypy --explicit-package-bases .checked 172 source files and retained the public head's 321 existing errors, with zero new or removed diagnostics under normalized multiset comparison. Defaultmypy .aborts on duplicate module discovery in this checkout; neither command is reported as a clean type-check pass. Sphinx builds with the same 20 existing warnings and zero new warnings.Final matched A/B: source11 versus public 72fb137. Three alternating pairs on the first host used the same prompt SHA, token counts and READ chunk counts; all generated outputs matched. The table is the mean of three run means. Excluding each run's first 6.5K request leaves 57 included requests and 468 RPC-inclusive retrieval batches per version.
Both serving engines and the memory agent ran on one host, with one H200 NVL GPU per engine, CXL DRAM and physical-RoCE NIC loopback. The stack was Llama-3.1-8B-Instruct, vLLM 0.22.1, Torch 2.11.0 and NIXL 1.2.0, with
--enforce-eagerand the default asynchronous load/store path. Each length used five sequential synthetic document prompts asking to repeat the last word, withmax_tokens=16,temperature=0and a 0.5-second store-to-load wait. The three output paths were initial compute-and-store, external load and local prefix reuse; the table compares external-load TTFT between implementations.The 32K second pair was slower: 265.10 → 267.50 ms (+0.9053%). Across the included subset, retrieval-batch maxima were 12.35 / 12.22 ms (public/source11), with no batch over 20 ms in either arm. The independent repeats are three run pairs, not 57 independent trials; this is not a formal noninferiority test. Lifetime fixes, duplicate-store removal and event progress are combined, so the result does not isolate the event timeout's causal effect.
A separate five-second, single-worker-factory observation of public72 versus source10, without peers or registered payload buffers, measured 100.92% → 1.43% CPU, relative to one core. It is not whole-vLLM CPU or a power measurement. Source11's concurrent-decode busy gaps remained 257.01–372.50 ms while two fresh prefills also ran; those gaps cannot be attributed to RDMA alone.
The separate shared-CXL platform uses Qwen3-8B. Earlier source8/9 validation and source9's fixed-route comparison are retained separately: 16K/32K means increased 2.81% / 4.85% versus public, while reader re-stores fell from 950 chunks per run (2,850 over three runs) to zero. Those results do not test the latest native worker. The subsequent outage of one node in that platform has no established cause.
Evidence IDs:
final11-ab-analysis.json(120 request rows, 42 evidence-file hashes),runs/val_fix11/status.tsv,runs/regress_final11.out,runs/native_twohost_fix11_01/sealed-manifest.jsonandmypy-final11-delta.json, in campaign20261021_pr107_upstream_validation. Exact check commands and scope are in the test record. Hardware results do not establish every vendor error's DMA quiescence, real host-reboot cache visibility, or performance at arbitrary pool/worker counts.Earlier completed source10 A/B — retained under its original snapshot
Source10's three-pair comparison measured mean cache-load TTFT changes of −3.46%, −3.57%, −2.71% and −2.82% at the same four target lengths. Its included-subset retrieval maxima were 12.01 / 12.36 ms (public/source10), with no batch over 20 ms. These are
final10-ab-analysis.jsonresults, not source11 measurements.Earlier unit and resource-manager regression record
The original public
72fb137record with CUDA hidden was 1,563 passed, 33 skipped, with 383 new tests in ten files at that revision. The threaded files passed 15 repeated runs on0b71898and ten on3316416. These historical counts and repeats are preserved separately from the final source11 suite.On
208ab5e, 97 ctest tests passed, including 11 hold tests, together with eightmaru-rm-ctltests and a Python client against the resource manager's fake-DAX test daemon. This is the earlier test record, not a newly repeated source11 C++ suite.Historical hardware checks and TTFT — 0b71898, not the current source
These original results used one CXL memory node and a separate two-GPU worker over RoCE. They are retained for history and are not the current source11 A/B result. The 338.8 ms figure below is the earlier rounded 339 ms claim.
At
0b71898, computing and storing a prompt took at most 0.6% longer than recomputing it with asynchronous stores. This historical store-overhead result has not been relabeled as a source11 measurement.TTFT in ms (median of 5 prompts; recompute is vLLM without the connector; "async" is asynchronous loads with vLLM's asynchronous scheduling on):
At that historical revision, covered by unit tests only: a memory node reboot, the layout check, lease-release retries and late WRITEs that end in an error.
🔗 Related Issues (optional)
🌿 Related PRs (optional)
🌿 Related Branches (optional)
📦 Release Note (for auto-generation / write in English)
NEW
CHANGED
FIXED
IMPORTANT NOTES