[chuangz0]

Summary

This PR adds the Mooncake driver for RI-08D's blob backend, which is on main through #20884.

• MooncakeDriver maps [None][feat] Add the blob store driver boundary and an in-memory test driver #20217's private StoreDriver boundary onto the pinned Mooncake 0.3.13 SDK.
• open_mooncake_driver opens the driver from the store's own config file.
• Tests run it over a fake SDK, a real Mooncake pool, and a GPU round trip between two staging lenders.

The factory that RI-02's TransferManager calls is left to a follow-up on top of #20266, whose BackendFactory shape (cleanup dispatcher, FactoryRejected) is still in review. Nothing is wired into the executor yet, so the path stays off until a runtime assembles it.

The initial profile follows the design: KVCM v2, TP1/DP1/PP1/CP1, one pool, host staging from the #19722 lender, and None hints only. Names like RI-08D are items of the Runtime Integration design; M1–M5 are milestones of the Store Backend design. Both documents are linked below.

What changed since the last review

The backend moved to main. #20884 (merged as b351c407bba) copied blob/backend.py, keys.py, store.py, the testing package and their tests verbatim from this PR's head 15c9491ba4a. It then added:

• batch registration (RegistersRegionsTogether);
• a cut of each call at a repeated object key;
• MemoryDriver fixes and new conformance checks;
• a README.

This PR was rebuilt on main 592bd6fd8d9 and now holds only the Mooncake driver, its setup and their tests. The review threads on the backend and the boundary have replies, and that code is unchanged on main in those respects.

New in this PR since that head:

• drain makes no SDK call (thread from @chienchunhung). It used to revoke the put sessions left open, in a blocking RPC outside its budget, which could hold up drain(0.0) and with it the backend's revocation and the lender's hold. close, the owner's last step under main's rule, now does that revocation before it closes the SDK client. A test sets up a gate that would hold a revoke inside the SDK, and checks that drain returns at once without reaching it.
• Who closes the driver. open_mooncake_driver says what main's README says: the caller that opened the driver closes it last. close adds that this happens off the executor loop's thread. The tests close their drivers that way.
• A repeated object key is pinned over Mooncake. The backend cuts the call there, so no operation names a key twice; the driver keys a part's results by store key and relies on that.
• NOT_READY follows what the store reports (thread from @nv-xtf). Main's RowStatus.NOT_READY now covers an object of the key that the store reports without a readable complete replica, with a get's slot untouched, and another put of the same driver, still running, that writes the key; RPC, transfer and integrity errors stay FAULT. README rule 4 now defers to RowStatus for these conditions. So Mooncake's first-read -703 and a put start's -705 with no complete object are NOT_READY under the rule, and the departure note is gone.
• A put of a key an earlier put of this driver left open is DECLINED, with its own reason and counter (rows_declined_left_open_by_own_put), when no complete object stands behind the key (thread from @nv-xtf, agreed). The driver never writes such a key again, and a FAULT here would count the earlier put's outcome again on every later put of the key. Beyond a store's own refusal, RowStatus.DECLINED covers only this case: the driver will not write the key again, as an earlier put of its own may still hold the key or act on a new write of it.
• A key whose first put end went unanswered is left open even when the retry ends the session (second thread from @nv-xtf). Mooncake's master matches a session end or revoke to an object by key and client alone, and an end answered RPC_FAIL or RPC_TIMEOUT may still run later: the SDK closes the timed-out connection, the retry goes out on another, and the master runs the first end whenever it arrives. If the object was then evicted and this driver put the key again, the late end would mark the new, unwritten object complete. This was reproduced on a real 0.3.13 master with a proxy holding the first end. The row stays READY; the key is never put again. A start answered RPC_FAIL or RPC_TIMEOUT still frees its key: leaving it open would bar every key whose start met a master outage from ever being put again, while a late start can only reserve an absent key (its next put is NOT_READY until put_start_discard_timeout_sec) or replace a put of the key unfinished past that timeout. The ledger and its names now say keys rather than sessions (_keys_left_open, put_keys_left_for_close, put_keys_left_to_master).

The previous head is kept as backup/pr20473-before-rebase-onto-20884 in the fork.

Design

• Unified KV Sharing: Runtime Integration, tab "v1: design and roadmap", with RI-08D and RI-09 in M2.
• Unified KV Sharing: Store Backend. This PR covers M3, the Mooncake driver.

Design item	Where
RI-08D: "Assembly creates BlobBackend; TransferManager retains it"	Not in this PR. The factory builds BlobBackend(RegionBinding(regions, driver, hold), wake, config) from main, over the driver open_mooncake_driver returns
RI-08D: the store owns batching, the driver, the pinned SDK, and access-end evidence	The driver owns batching and its threads: a fixed pool of workers cuts parts and decides the order of gets and puts
RI-08D: the store owns SDK capacity	The mooncake section of the store config: workers (default 4), rows_per_call (256), bytes_per_call (128 MiB), each counted per SDK call. Gets go first; once four gets have gone ahead of a waiting put, that put goes next. Puts take at most max(1, workers − 1) threads, so with two workers or more a get never waits behind puts alone
RI-08D: "initial Store accepts only None as its hint"	Mooncake declares no hints, so main's backend refuses a hinted fetch before the driver is called
Store §3.2: resources come back only after access has ended	An operation whose access cannot be proven ended keeps its region from being unregistered and makes drain answer False
Store §3.3: advisory cancel	Cancel ends the rows no worker has taken; the SDK cannot stop a running call
Store §3.5: put is READY only for a complete, readable object, never on "already exists" alone	A session put ends only the rows whose ranged write returned the full size. A key that already exists is READY only if batch_is_exist sees it committed
Store A.2: what a driver must answer	blob/mooncake.py; see "Mooncake mapping" below
Store §5 M3: SDK mapping and real-service tests that reuse the round-trip cases	test_mooncake_real.py; test_blob_backend_round_trip.py, with a real lender on GPU; main's driver_conformance suite run against MooncakeDriver over the fake SDK and over the real pool
Store §6: kv_sharing.backends: [{type: mooncake, config: <path>}], with store parameters in the store's own file	blob/mooncake_setup.py reads a YAML file with blob (key prefix) and mooncake sections; no LLM args are added

Departures

• No trailer. Store §3.5 describes a format of payload plus trailer, but [None][feat] Add the blob store driver boundary and an in-memory test driver #20217 has since left the object format to each driver. A Mooncake object is the payload alone, checked by the SDK's size check and the session protocol (an object is readable only once its session ended).

Mooncake mapping

• put calls the session API, one SDK call per step for a part of up to rows_per_call rows: batch_put_session_start, then batch_put_from_multi_buffer_ranges (one scatter for the part), then batch_put_session_end for the rows written whole. The rest get batch_put_session_revoke, except a row the target may still write into, whose session is left open (see "Put sessions" below).

  • The transfer engine's scatter wait() has no timeout. It returns only once every fragment has physically finished.
  • So once that call returns, its access has ended, even if it answered TRANSFER_FAIL, and a failed row is revoked only after its write has stopped.
  • With batch_put_from, by contrast, a transfer that times out after 60 s is revoked at once while RDMA writes may still be in flight. A late write could then land in space the master has already given to another object.

• get is one batch_get_into per part.

  • A session get makes three SDK calls, and each return must reacquire the GIL.
  • Next to a thread running Python nonstop, a 32 × 1 MiB session get took 16 ms at p99, against 5.6 ms for batch_get_into.

• Put sessions are per key and per client.

  • The driver never starts a session for a key another of its puts is writing now; that row is NOT_READY without asking the SDK. A key repeated within one view lands in two calls; a second put of it that runs alongside the first is NOT_READY, never a fault.
  • A key whose session calls did not all have a known outcome, or that may still hold a session, is kept in one ledger:
    • left to the master, never revoked by the driver, because the target may still write into the object: a scatter that raised or miscounted, or a row the scatter answered TRANSFER_FAIL. The master frees these after put_start_release_timeout_sec;
    • left for close to revoke: a session start that raised or miscounted, an end whose first answer was RPC_FAIL or RPC_TIMEOUT (even when the retry ended the session) or that raised or miscounted, a failed revoke, or a put stopped by a raise. They hold no lender memory, so drain leaves them; one whose revocation fails is freed by the master. A put stopped by a raise is unprovable, so close never comes and the master frees its sessions too.
  • A key left open is never put again, as a late end or revoke would act on the next put of it. A later put of it is READY while the master shows its object complete, and DECLINED otherwise. A start answered RPC_FAIL or RPC_TIMEOUT frees its key (see above).

• Results:

  Mooncake answer	Row result
  get: object not found (-704)	MISS
  get, first read: no readable complete replica (-703), the slot untouched; put: a key another put of this driver is writing now, or an object that exists with no readable complete replica (-705 and batch_is_exist 0)	NOT_READY
  put: a key an earlier put of this driver left open, with no readable complete object (batch_is_exist 0)	DECLINED
  put start: pool full (-200)	DECLINED
  put start: already exists (-705) and batch_is_exist 1	READY
  Lease expired (-707)	Read once more. The reread is READY only if it reads the whole object again; anything else is FAULT, never MISS or NOT_READY, because the first read already wrote the slot
  put start -600 (the driver holds no session of the key, so it is something else, such as a cachelib master's 16 MiB − 16 B object cap), a wrong size, or any other code	FAULT

• Access-end evidence:

  • A put part's access ends when its ranged write returns an answer for every row, whatever the answers are, TRANSFER_FAIL included.
  • A get part's access ends when batch_get_into answers every row and no answer is TRANSFER_FAIL.
  • In every other case the operation turns access_unprovable: a TRANSFER_FAIL on a get, an exception, or a wrong answer count. It keeps its region from being unregistered and makes drain answer False. The backend keeps the attempt's slots for good.

• Capabilities:

  • get, put and cancel. Cancel ends the rows no worker has taken; the SDK cannot stop a running call.
  • absent_is_miss: the master answers OBJECT_NOT_FOUND for a key it does not hold, before any transfer.
  • Host memory only (device_memory=False). Main's backend refuses a row in any other region with InvalidSubmission, and a lender with no such region with ConstructionRejected.
  • Rows of one span (multi_span=False): batch_get_into reads each object into one buffer.
  • No probe, and no row limit of its own (max_rows_per_op=None), since the driver cuts parts itself.
  • Regions are registered one at a time; the driver does not implement RegistersRegionsTogether.

• Shutdown: drain cancels the rows no worker took, ends the worker threads and waits for the running parts, and makes no SDK call. close, once a drain returned True and every region is unregistered, revokes the sessions left for it and then closes the SDK client, even if the revocation raised. The revocation waits on the master, and closing holds the GIL through the unmount, so the owner calls close off the executor loop's thread.

Scope

Commit	Production	Tests
961ee1277ee Add the Mooncake store driver	mooncake.py	a fake SDK, mapping, session and scheduling tests, and the conformance suite over the fake SDK
ec7bd8d3742 Open a Mooncake store driver from a store config	mooncake_setup.py	configuration, setup and layering tests
ec265c74665 Run the blob backend over a real Mooncake pool		opt-in tests (TRTLLM_MOONCAKE_REAL=1) against mooncake_master and a segment provider, plus the conformance suite on the real store
4d311477bb9 Round-trip KV blocks between two staging lenders through the blob backend		GPU test with real KV cache managers, over the in-memory store and a real Mooncake pool
840590063e7 Revoke the Mooncake sessions left open in close, so drain makes no SDK call	drain, close	a revoke held inside the SDK; owners close their drivers
321fc530e8d Pin a repeated object key over Mooncake and say where NOT_READY departs	comments	a view naming one key twice
83496c94f31 Close the Mooncake client even when its revocation raises, and pin the order	close	the fake SDK answers nothing once closed, as 0.3.13 does, so the revoke-before-close order is pinned
aaeeb4cdede Widen NOT_READY to what the store reports, and decline a put of a key left open	main's RowStatus wording (NOT_READY, DECLINED, the class and probe docstrings), README rule 4, the left-open reason in mooncake.py	the left-open session tests expect DECLINED and its counter
a9f6a74c3cb Leave a Mooncake key open when its first put end went unanswered	_end_sessions (the first end answer decides the key), the ledger renamed to keys, main's RowStatus.DECLINED wording	the retried end leaves its key open; a timed-out start frees its key; a key left open, to the master or for close, is never put again

The first two commits add the Mooncake modules to blob/__init__.py; the first three list them in main's blob/README.md.

Production code: +1497 / −12 lines, of which main's store.py is +15 / −11, plus +12 / −3 in the README. Tests: +3177 / −2 lines. Each commit passes the CPU tests present at it.

Verification

• CPU, at a9f6a74c3cb, over main's backend:
  • Blob tests: 647 passed, 170 skipped. This count includes main's own blob tests; the skips are the real-pool tests and conformance checks that a driver's capabilities make inapplicable. These ran with the package's heavy imports stubbed, not from a wheel; CI runs them in the full environment through l0_cpu.
  • The Mooncake files alone: 195 passed, 62 skipped.
  • 16 workers, 30 repeats: 19410 passed, 5100 skipped, no failures.
  • The left-open change was checked against 17 mutants of its logic: all but one fail a test. The survivor names the first rather than the final end answer in a fault's text.
  • The new drain/close tests were checked against mutants: revoking in drain again, revoking after the client is closed, and closing without the finally each fail them.
• Real Mooncake over TCP, at a9f6a74c3cb: 28 passed, 20 skipped. The conformance suite on the real store now also runs main's extent check; the three batch-registration checks are skipped, as Mooncake registers one region at a time.
• Conformance over the fake SDK: runs every check the driver's capabilities and the fake's scenario hooks allow (blocked get and put, a key another writer holds, a full store, a crash in put and get, a failing read, an unprovable get, a driver that leaves its operations unstarted).
• Not re-run since the rebase:
  • Mutation testing. About 1,000 mutants ran over the production modules at the previous head. Every surviving mutant that changed a row result, an access end or when memory is released was given a test that kills it.
  • The GB200 RDMA runs. They passed on f2d736b023d.
• GPU round trip (kv_cache/sharing, in l0_a10), written against main's backend: passed in CI's A10-PyTorch-2 stage (L0 pipeline 66862), on the same diff before the rebase onto 592bd6fd8d9. It moves each row on its own: the in-memory driver takes one row per operation on a thread of its own, and Mooncake cuts each operation into one-row parts for two workers.

Objects of the wrong size. Both cases are checked.

• If the object is longer than the row, Mooncake refuses it with -600 before any transfer (real_client.cpp:5134-5141).
• If it is shorter, Mooncake writes it into the head of the slot and returns its size (real_client.cpp:5169-5180). The driver reports FAULT, its access ends normally, and nothing outside the row is touched.

Performance

Driver-owned concurrency against the previous design, on CPU. These numbers were measured at 15c9491ba4a, before the rebase, against f2d736b023d, and have not been re-run. The backend part of what they measure is now main's; main has since moved the repeated-key check into the call cut, and nothing else on the measured path.

• Setup, for both versions:
  • a fixed submission rate;
  • a fake SDK whose data calls sleep 1 ms with the GIL released;
  • each version's Mooncake defaults: 4 workers, 256 rows and 128 MiB per SDK call.
• The executor loop runs about 1 ms of Python per iteration. It is measured two ways:
  • torch-like: it releases the GIL 4 times per iteration, as torch and CUDA calls do;
  • pure Python: it never releases the GIL, the worst case.
• In each cell the previous version's number comes first.

Mean submission to settle, torch-like executor:

	1 × 256 rows	4 × 256 rows	1 × 4096 rows
fetch	2.87 → 2.38 ms	6.56 → 3.83 ms	16.17 → 9.36 ms
all submissions	3.05 → 2.62 ms	6.42 → 4.85 ms	16.81 → 11.20 ms

The executor's own cost, pure-Python executor:

	1 × 256 rows	4 × 256 rows	1 × 4096 rows
Submit call (fetch/publish) p99	4.8–5.7 → 0.27 ms	11.4–11.7 → 0.45 ms	8.6–9.2 → 1.4–1.5 ms
Executor time lost to the GIL per submission	0.83–1.04 → 0.65–0.80 ms	2.9–3.0 → 2.0–2.1 ms	12.3–12.8 → 5.2–5.9 ms
Executor stalls over 1 ms	2–5 → 0	33–35 → 1–2	129–133 → 6–8

Worker CPU per submission is less than half of what it was. Settle with the pure-Python executor stays as before for 1 × 256, improves for 1 × 4096 (p50 56.7–57.0 → 52.3–52.4 ms), and is slower for 4 × 256 (p50 20.3–20.6 → 24.8–25.0 ms); see "Known limitations".

Throughput on two GB200 nodes over RDMA, measured on f2d736b023d, not yet on this branch.

• Setup: segments on the other node; 4 workers; 128 MiB per SDK call.
• NIC use: one process uses the two NICs of the NUMA node where its staging memory lives; four ranks use all four NICs.
• "Busy" means another thread runs Python nonstop.
• The microbenchmark will follow in a separate PR.

setup	main thread	put GB/s	get GB/s
one process, 4 MiB rows	idle	89–91	94
one process, 4 MiB rows	busy	15–19	31–43
one process, 1 MiB rows	busy	12.5–14	24–27
four TP ranks on one node, NUMA-bound, 1 MiB rows, node total	idle	160	179
same	busy	52	95

A row is one KV block (32 tokens) across all layers of a layer group, holding this rank's heads. With TP the rows shrink, so realistic rows are 0.25–2.5 MiB: Llama-3.1-8B at TP4 1 MiB, Llama-3.3-70B at TP4 2.5 MiB, DeepSeek-V3 (MLA) 2.1 MiB.

Please confirm

• @chienchunhung:
  • A driver call that cannot prove its access ended keeps its attempt's slots for good, and the binding ends in the profile's containment. RI-02 implements only executor-fatal, so today a get that answers TRANSFER_FAIL, raises, or miscounts stops the executor.
    • A fetch reaches this state on a TRANSFER_FAIL answer, an exception, or a wrong answer count from batch_get_into. SDK 0.3.13 answers TRANSFER_FAIL alike for a failed slice, a status-query error, a partly submitted batch and its 60 s timeout. In the last three cases the transfer may still be running, and the answer does not say which case it is (transfer_task.cpp:746-805, 813-905, 1239-1258).
    • A publish reaches it only on an exception or a wrong answer count from the ranged write, because that write returns only after every fragment has finished.
    • The alternative is a session get, whose access ends even on TRANSFER_FAIL, at two more SDK calls per part. This PR keeps batch_get_into and asks for a containment profile that is not executor-fatal for a store fetch, since the request can recompute instead.
  • A driver call (one put or get) has no time bound; the lease holder's deadlines bound how long a request waits, and shutdown labels name what is still held.
  • Mooncake declares no row limit, so main's driver_call_limit(region) returns None. One driver call is then a layer group of a submission, cut only where it names an object key twice, and a holder bounds what one unprovable call holds by the size of its leases.
• @nv-xtf: a key whose first put end went unanswered is now left open even when the retry ends the session (a9f6a74c3cb); a start answered RPC_FAIL or RPC_TIMEOUT still frees its key. See the thread for why.

Known limitations

• The data path does its per-part work under the GIL. The SDK holds the GIL while converting a call's lists, and each worker runs some Python per part. Next to a busy executor, store throughput drops sharply (see the GB200 table). A GIL-free data path (a batched Mooncake API with completion notification, or a C++ worker) could bring busy throughput close to idle; whether it is worth building should wait for the end-to-end TTFT measurement (Store M4). Passing numpy arrays to the SDK instead of lists does not help: the 0.3.13 binding takes list[int], so it converts an array element by element, measured 30 times slower than tolist().
• An executor that never releases the GIL delays the first transfers. Such an executor leaves the store's workers waiting for CPython's 5 ms forced switch, so submission to settle can take one switch interval longer than with the previous design (4 × 256 rows: p50 20.5 → 24.8 ms). Yielding after a submission fixes this case only and slows the 4 × 256 fetch next to a torch-like executor, so nothing yields.
• Fetch latency next to a sustained publish load. Gets go first, but fetches and publishes still share the GIL and the transfer engine.
• Deployment.
  • Keep each rank on its GPU's NUMA node. TRT-LLM's NUMA-aware worker affinity (on by default, TLLM_NUMA_AWARE_WORKER_AFFINITY) does this.
  • One process tops out near 100 GB/s; four ranks on their own NUMA nodes reach 160–180 GB/s together.
• No probe. The driver declares probe=False. batch_is_exist could back one later, but in 0.3.13 it answers 0 both for an absent key and for one still being written, takes no timeout, and renews the object's read lease on a hit.
• Lender regions name no memory kind yet. Until the lender's MemoryRegion adds memory_kind, whoever builds the backend hands RegionBinding host-memory SlotRegions; the GPU round trip does this.
• Other.
  • Session puts carry no checksum.
  • A key left open is never put again by this driver, so a key whose first end went unanswered is not refilled from this process once the store evicts it; other processes still can. The root fix is in Mooncake: an end or revoke that names the attempt it ends, as [Store] Fence stale put finalization with ReplicaID kvcache-ai/Mooncake#3180 did for RFC [RFC]: Replica-ID Validation for PutEnd / PutRevoke and GetReplicaList kvcache-ai/Mooncake#3030 (closed unmerged). Once a release has it, the rule can go.
  • The SDK is pinned to 0.3.13, and the access-end reasoning relies on that version's transfer engine.
  • MC_STORE_LOCAL_HOT_CACHE_SIZE other than empty or 0, TCP's legacy framing and the newer transfer engine are refused when the driver is opened.
  • The last worker to exit calls on_change for the rows still queued after it stopped counting. A drain returns True before that call only once every worker died of a BaseException that left no operation unprovable, which only an on_change raising against the protocol does; under BlobBackend the owner cannot close the driver before that call is recorded.
  • The store config is read into frozen dataclasses with strict type checks rather than Pydantic models; it does not become LLM args.
  • The real-Mooncake tests are opt-in and do not run in CI.

Notes

The base is main 592bd6fd8d9, which contains #20117, #20217 and #20884.

• Contribution checklist reviewed: tests are included; there are no new dependencies and no new public LLM configuration fields; the new path is off unless a runtime assembles it.

🤖 Generated with Claude Code

Dev Engineer Review

MooncakeDriver adds a host-memory implementation of the StoreDriver boundary for Mooncake SDK 0.3.13. It supports batched transfers, session management, access tracking, cancellation, and shutdown. open_mooncake_driver validates configuration and opens the driver. The PR does not add a runtime factory or wire the driver into the executor.

drain does not call the SDK. close revokes eligible remaining sessions before closing the client. The driver retains access slots when it cannot prove an operation has ended. This can prevent a successful drain and block executor progress. The supported setup profile is limited; verify its configuration checks and assumptions against deployment needs.

QA Engineer Review

The changes add CPU-only fake-SDK driver and setup suites, opt-in real-Mooncake integration tests, and CUDA-gated blob-backend round trips. Driver tests cover status mapping, transfer limits and scheduling, cancellation, drain and close behavior, session cleanup, repeated keys, and unprovable access. Setup tests cover configuration, SDK selection, environment checks, networking, cleanup, and import boundaries. Real-pool tests cover declines, object-size handling, concurrent transfers, and duplicate publishes. Backend tests cover full and partial fetches.

Current test results are unavailable. The supplied author-reported results apply to earlier revisions and do not establish the current status. The changed tests were not found in the scanned CI or manual-QA lists. The Mooncake entries found in those lists refer to other tests. Coverage verdict: needs follow-up.

Per-File QA Perspective

• tensorrt_llm/_torch/disaggregation/blob/README.md: Documents the driver, setup module, test requirements, and result classifications. Verify the documented verdict conditions against RowStatus behavior.
• tensorrt_llm/_torch/disaggregation/blob/__init__.py: Clarifies driver and setup responsibilities and lazy SDK imports. Verify that package imports do not load mooncake.store prematurely.
• tensorrt_llm/_torch/disaggregation/blob/mooncake.py: Adds driver transfer, scheduling, access-tracking, cancellation, and shutdown paths. Verify SDK result mapping and session cleanup, especially when access completion is unprovable.
• tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py: Adds configuration parsing, SDK and environment validation, master checks, and driver setup. Verify accepted settings and cleanup after failures.
• tests/test_common/mooncake_utils.py: Adds a subprocess-backed Mooncake segment provider with readiness and cleanup handling. No matching entry was found in the scanned CI or manual-QA lists.
• tests/unittest/_torch/disaggregation/blob/fake_mooncake.py: Adds a scripted fake SDK for driver tests. Verify that fake session and buffer behavior represents the SDK cases under test. No matching entry was found in the scanned lists.
• tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py: Covers driver conformance, transfer scheduling and limits, result mapping, access tracking, and lifecycle errors. No matching entry was found in the scanned lists.
• tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py: Adds opt-in tests for real-pool transfers, concurrency, and duplicate publishes. No matching entry was found in the scanned lists.
• tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py: Covers configuration, SDK setup, networking, failure cleanup, and import boundaries. No matching entry was found in the scanned lists.
• tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py: Adds CUDA-gated full and partial fetch tests. No matching entry was found in the scanned lists.
• tensorrt_llm/_torch/disaggregation/blob/store.py: Clarifies NOT_READY and DECLINED semantics. Verify that callers handle concurrent or unresolved puts as documented.

[Shixiaowei02]

This registers every region the lender exposes and treats each one as host memory. Nothing checks that, so a region backed by device memory (a pool lent in place) would be registered with the transfer engine as host memory. Please register only host regions that have slots, keep their indices, refuse construction when none qualifies, and add a test that a device region is never registered.

[chuangz0]

Done. RegionBinding.served_region_indices keeps only regions with slots in memory the driver can access: HOST, or DEVICE when the driver declares device_memory (Mooncake does not). With none, register() releases the hold and raises ConstructionRejected. Tests: test_region_binding_registers_the_regions_with_slots_the_driver_can_access, test_a_binding_with_no_region_to_register_is_rejected_and_lets_the_hold_go.

[Shixiaowei02]

A row whose region_index names a region this backend did not register is a row it does not recognize. Please reject it here with InvalidSubmission instead of passing the driver an address in memory it never registered.

[chuangz0]

Done: _checked_group_rows raises InvalidSubmission for a region not in served_region_indices (test_rows_in_a_region_the_backend_did_not_register_are_rejected).

[Shixiaowei02]

Rejecting every hint closes this class to drivers that declare hints. Could a fetch pass the hint through to driver.get(rows, hint) when the driver declares hints, and reject it otherwise? A publish has nowhere to put a hint, so rejecting a non-None destination there is right.

[chuangz0]

Done: a fetch's hint goes to every get of a driver that declares hints; a fetch with a hint over any other driver raises InvalidSubmission before the driver is called, and a publish still takes only None (test_a_fetch_hint_reaches_every_get_of_a_driver_that_reads_hints, test_invalid_submissions_are_rejected_before_the_driver[hint]).

[Shixiaowei02]

After a driver breaks its protocol (a wrong result count here, or a raise other than DriverRejected at line 549), the backend keeps accepting work. The holder then only sees attempts that never quiesce, and stops using the binding only when its retention limit trips. Please refuse every later submission with SubmissionRejected, with a fault text that names the violation.

[chuangz0]

Done: a raise other than DriverRejected from put/get, a put/get that returns no handle, a query that raises, and a wrong result count are each recorded as a protocol violation. The call is held as unprovable (unless its access end was already recorded), and every later submission gets SubmissionRejected naming the violation. A submission being handed over when the violation is recorded hands no further call to the driver: the next call fails as a whole, the rest end CANCELLED, and none of them holds memory; calls already in the driver run on. Tests: test_any_other_raise_from_the_driver_holds_the_call_and_refuses_later_work, test_a_poll_with_the_wrong_row_count_faults_the_call_and_keeps_it_held, test_the_calls_after_a_protocol_violation_never_reach_the_driver, test_a_violation_told_during_the_hand_over_stops_the_rest_of_it, test_calls_cut_off_by_another_attempts_violation_end_their_access, test_a_protocol_violation_leaves_calls_already_in_the_driver_alone.

[Shixiaowei02]

When the only work left is unprovable, the progress thread can never finish, so every shutdown waits out its whole budget. Could shutdown return as soon as no worker is inside the driver and every open batch is unprovable?

[chuangz0]

Done: shutdown returns as soon as every attempt not yet quiesced holds a call whose access can never be proven to end, since from then on neither its label nor a region's can clear (_shutdown_answer_is_final). That also covers an attempt whose other calls are still running in the driver. Tests: test_shutdown_returns_once_every_open_attempt_has_an_unprovable_call, test_shutdown_returns_once_a_call_in_the_driver_turns_unprovable.

[Shixiaowei02]

The code mapping and the access-end rules were read from 0.3.13. On another SDK they may not hold, and an early return from a put would let a revoke free space that a transfer is still writing. Please refuse other versions, with a test-only override if needed.

[chuangz0]

Done: open_sdk refuses a mooncake.store not shipped at 0.3.13 by one of the pinned distributions (_require_pinned_builds, which also refuses when any installed build that lists the module's file is not one of those distributions at 0.3.13, as a second install over the same path can leave). No override was needed: the tests fake the installed distributions (test_only_pinned_sdk_builds_are_opened).

[Shixiaowei02]

This holds the first scatter for at most 5 s, well under the master's put discard timeout (30 s by default). The case that matters is a first put stuck past that timeout: the second put of the same key from this client must still be refused, and batch_is_exist must stay 0.

[chuangz0]

Added test_real_put_stuck_past_the_discard_timeout_keeps_its_keys: the master runs with --put_start_discard_timeout_sec=2 and the first scatter is held past it. The second put of the same keys from this client is skipped without asking the SDK, and batch_is_exist stays 0 until the first put ends.

[Shixiaowei02]

Please check that mooncake_master comes from the same pinned distribution as the client, and record both versions in the run log. Beyond the current cases, the real-pool run should also cover a TRANSFER_FAIL get (access stays unproven), a lease-expired reread, a master kill, freeze and restart, a store-side freeze during a put, and process exit while a call is blocked.

[chuangz0]

Version part done: the master must be shipped by one of the pinned distributions and report 0.3.13, every installed pinned client build must be 0.3.13, and both versions are printed and recorded as a test-suite property. Of the extra cases, the real pool covers a put stuck past the discard timeout. A TRANSFER_FAIL get and a lease-expired reread are covered only over the fake SDK (test_mooncake_driver.py). Master kill/freeze/restart, a store-side freeze during a put, and process exit during a blocked call are not in this PR.

[Shixiaowei02]

This strict xfail, and the one at test_mooncake_real.py:228, depend on the suite's "destination unchanged" assertion for short objects. batch_get_into writes a shorter object into the head of the slot, so that assertion does not fit this store; both xfails should go when the suite drops it.

[chuangz0]

Done: for a smaller object the suite now checks only that the row is a FAULT and that the canaries outside the slot are intact (test_an_object_smaller_than_its_row_is_a_fault); both xfails are gone.

[Shixiaowei02]

Loading the in-memory driver by file path ties this conftest to the test-tree layout. Once the driver suite and the in-memory driver are importable from the package, please import them instead.

[chuangz0]

Done: the in-memory driver and the driver suite moved to tensorrt_llm/_torch/disaggregation/testing/. The conftest's path loading is gone, and the round-trip test imports them from the package.

[Shixiaowei02]

This rejects every non-None hint, but blob/store.py in this PR says the backend passes a hint only to a driver that declares hints and rejects it before dispatch for the others (Capabilities.hints, StoreDriver.get). As written, a driver that fetches across instances by hint can't be used under this backend. Could a fetch's hint go to every get of a driver that declares hints (the same hint for each get when a submission is split into several), keeping the InvalidSubmission rejection for drivers that don't, with a test for both cases? Publish can keep rejecting non-None hints, since put takes none.

[chuangz0]

Done, same change as the thread above: when a fetch is split, each get receives the same hint object; a fetch with a hint over a driver without hints raises InvalidSubmission before the driver is called, and publish keeps rejecting a non-None destination. Both cases are tested in test_a_fetch_hint_reaches_every_get_of_a_driver_that_reads_hints and test_invalid_submissions_are_rejected_before_the_driver[hint].

[chienchunhung]

If a put-end RPC fails, this revoke call can make drain(0.0) wait for the SDK. The driver contract requires drain to return within its time budget. While it waits, the backend retains the region registration and lender hold, even after all data access has ended.

Maybe we should keep session cleanup within the remaining drain budget, and add a test with a blocked revoke call to verify that drain(0.0) returns promptly.

[chuangz0]

Fixed in 282e8493d01 and a53e895d184: drain makes no SDK call any more, so drain(0.0) returns at once whatever the master does. The sessions left open hold no lender memory, so nothing needs them revoked before the regions and the hold go. close now tries the revocation once and then closes the SDK client, even if the revocation raised. Under main's rule close is the owner's last step, after shutdown returned no labels and release_hold(), and it already waits on the master for the unmount. A session whose revocation fails is freed by the master after put_start_release_timeout_sec. The ledger's FOR_DRAIN is now FOR_CLOSE. test_drain_returns_at_once_while_a_revocation_would_block holds a revoke inside the SDK and checks that drain returns without calling it and that close revokes afterwards. test_sessions_left_open_are_revoked_only_in_close covers a revocation that fails.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/TensorRT-LLM
• Branch: main
• File: .coderabbit.yaml
• Commit: 592bd6f
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: ac04114d-54c5-45e7-aea8-6b91591ba366

Commits

Reviewing files that changed from the base of the PR and between aaeeb4c and a9f6a74.

Files selected for processing: 3

• tensorrt_llm/_torch/disaggregation/blob/mooncake.py
• tensorrt_llm/_torch/disaggregation/blob/store.py
• tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

The change adds Mooncake configuration and SDK setup, a worker-based blob driver, and tests using a fake SDK, a real Mooncake pool, and the KV-cache blob backend. It also updates the documented RowStatus conditions.

Changes

Mooncake Blob Store

Layer / File(s)	Summary
Configuration and SDK setup tensorrt_llm/_torch/disaggregation/blob/mooncake_setup.py, tensorrt_llm/_torch/disaggregation/blob/__init__.py, tests/unittest/_torch/disaggregation/blob/test_mooncake_setup.py, tensorrt_llm/_torch/disaggregation/blob/README.md, tensorrt_llm/_torch/disaggregation/blob/store.py	Adds configuration parsing and validation, SDK build and environment checks, master availability checks, address and RDMA-device selection, and driver setup. Updates row-status documentation and Mooncake test instructions.
Driver lifecycle and scheduling tensorrt_llm/_torch/disaggregation/blob/mooncake.py	Adds host-memory registration, row admission, worker scheduling, operation polling and cancellation, draining, and closing.
Put and get transfer handling tensorrt_llm/_torch/disaggregation/blob/mooncake.py	Adds put-session and scatter handling, get reads and lease-expired rereads, per-row result mapping, and access-safety tracking.
Fake SDK and driver tests tests/unittest/_torch/disaggregation/blob/fake_mooncake.py, tests/unittest/_torch/disaggregation/blob/test_mooncake_driver.py	Adds a scripted in-memory SDK fake and CPU-only tests for driver transfers, scheduling, session handling, lifecycle, and backend round trips.
Real-pool and backend integration tests tests/test_common/mooncake_utils.py, tests/unittest/_torch/disaggregation/blob/test_mooncake_real.py, tests/unittest/_torch/executor/kv_cache/sharing/test_blob_backend_round_trip.py	Adds an opt-in Mooncake segment provider and real-pool tests, plus full and partial KV-cache blob backend fetch tests.

Priority: Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant BlobBackend
  participant MooncakeDriver
  participant Worker
  participant MooncakeStore
  BlobBackend->>MooncakeDriver: submit put or get rows
  MooncakeDriver->>Worker: queue a bounded part
  Worker->>MooncakeStore: issue batch SDK calls
  MooncakeStore-->>Worker: return per-key answers
  Worker-->>MooncakeDriver: record row results and access state
  BlobBackend->>MooncakeDriver: poll operation
  MooncakeDriver-->>BlobBackend: return row results

Loading

Merge Risk: Minimal · up to a9f6a

This change adds a Mooncake blob driver that is not yet wired into the executor, so it has no effect on runtime behavior. Prior review concerns about session cleanup and row-status mapping appear to have been addressed. No unresolved merge-blocking issue was identified.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 35.44% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely identifies the addition of the Mooncake store driver for the blob backend.
Description check		The description explains the purpose, scope, design, testing, verification results, limitations, and checklist status. It provides the information required by the template, although it uses expanded s…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description explains the purpose, scope, design, testing, verification results, limitations, and checklist status. It provides the information required by the template, although it uses expanded section headings.

Full details: Docstring Coverage

Explanation

Docstring coverage is 35.44% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 158 functions across 10 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[chuangz0]

/bot run

[tensorrt-cicd]

PR_Github #81167 [ run ] triggered by Bot. Commit: 36dce6d Link to invocation

[tensorrt-cicd]

PR_Github #81167 [ run ] completed with state SUCCESS. Commit: 36dce6d
/LLM/main/L0_MergeRequest_PR pipeline #66862 completed with status: 'FAILURE'

CI Report

⚠️ Resource Maintenance — Test Coverage Reduced:

Resource maintenance in LLM/main/L0_Test-x86_64-Single-GPU #9587: L40S-PyTorch-2 is temporarily paused due to Kubernetes agent provisioning and node issues.
Affected stages:
L40S-PyTorch-2

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chuangz0]

/bot run

[tensorrt-cicd]

PR_Github #81317 [ run ] triggered by Bot. Commit: aaeeb4c Link to invocation

[tensorrt-cicd]

PR_Github #81317 [ run ] completed with state FAILURE. Commit: aaeeb4c
/LLM/main/L0_MergeRequest_PR pipeline #67001 completed with status: 'FAILURE'

CI Report

⚠️ Resource Maintenance — Test Coverage Reduced:

Resource maintenance in LLM/main/L0_Test-x86_64-Single-GPU #10112: L40S-PyTorch-2 is temporarily paused due to Kubernetes agent provisioning and node issues.
Affected stages:
L40S-PyTorch-2

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[nv-xtf]

Does a successful put-end retry guarantee that the original timed-out RPC can no longer execute? Otherwise, could a late end prematurely commit a new object after eviction and reuse of the same key by this client? If so, should we retain the no-reuse guard even after retry success?