Repository navigation
Conversation
c537c01 to
f1b1689
Compare
| revoked and the hold is released. | ||
| RuntimeError: A registration may be live; the hold is kept. | ||
| """ | ||
| for index, region in enumerate(self.regions): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done. RegionBinding.served_regions 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.
| f"{where}: addresses must be int64 [{rows}], got {addresses.dtype} " | ||
| f"{addresses.shape}" | ||
| ) | ||
| if not 0 <= group.region_index < len(self._regions): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done: _group_rows raises InvalidSubmission for a region not in served_regions (test_rows_in_a_region_the_backend_did_not_register_are_rejected).
| # Submission, on the executor thread. | ||
|
|
||
| def _submit(self, view: View, hint: Hint, is_fetch: bool) -> BlobAttempt: | ||
| if hint is not None: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done: a fetch's hint goes to every get of a driver that declares hints; any other driver gets InvalidSubmission, 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]).
| key=(self._name, "observe", type(error)), | ||
| ) | ||
| return None, False, False | ||
| if count != len(batch): # outside the driver protocol: trust neither answer again |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done: a driver raise other than DriverRejected, or a wrong result count, is recorded as a protocol violation. Every later submission gets SubmissionRejected naming it, and batches other attempts still have queued are dropped (_stop_after_protocol_violation; test_any_other_raise_from_the_driver_holds_the_batch_and_refuses_later_work, test_a_poll_with_the_wrong_row_count_faults_the_batch_and_keeps_it_held).
| deadline = time.monotonic() + max(0.0, timeout) | ||
| if self._close_and_force_outcomes(): | ||
| self._call_wake() | ||
| self._progress.join(max(0.0, deadline - time.monotonic())) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
Done: shutdown returns as soon as every attempt still open holds an unprovable batch, since from then on neither its label nor a region's can clear (_shutdown_answer_is_final). That also covers a worker still inside the driver for such an attempt. Tests: test_shutdown_returns_once_every_open_attempt_is_pinned_by_an_unprovable_call, test_shutdown_returns_once_a_call_in_the_driver_turns_unprovable.
| return f"cannot import mooncake.store ({error}). {hint}" | ||
|
|
||
|
|
||
| def _warn_unless_pinned_version() -> None: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done: open_sdk refuses a mooncake.store not shipped at 0.3.13 by one of the pinned distributions (_require_pinned_version). No override was needed: the tests fake the installed distributions (test_only_pinned_sdk_builds_are_opened).
| return self.store.batch_put_from_multi_buffer_ranges(*args) | ||
|
|
||
|
|
||
| def test_real_concurrent_publishes_of_the_same_keys_from_one_client(clients, request) -> None: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| SLOT_BYTES = 64 << 10 | ||
|
|
||
|
|
||
| def _require_mooncake() -> None: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Version part done: the master must come from the same pinned build as the client, 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.
| def make_driver(self) -> StoreDriver: | ||
| return MooncakeDriver(FakeMooncakeStore()) | ||
|
|
||
| @pytest.mark.xfail(strict=True, reason=SHORT_OBJECT_LANDS) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
|
|
||
|
|
||
| @pytest.fixture(scope="session") | ||
| def blob_memory_store(): |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| # Submission, on the executor thread. | ||
|
|
||
| def _submit(self, view: View, hint: Hint, is_fetch: bool) -> BlobAttempt: | ||
| if hint is not None: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Done, same change as the thread above: when a fetch is split, each get receives the same hint object; drivers without hints get InvalidSubmission, 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].
4e583ed to
7f29395
Compare
The store driver boundary had no backend on top of it. BlobBackend implements the public Backend and Attempt contracts over any StoreDriver, and RegionBinding owns the lender's region hold and the registrations made under it. - Submission, in the backend's call sequence, validates the whole view with numpy, takes references to its arrays, cuts batches by rows, bytes and the driver's limit, and queues them; it never calls the driver. batch_size(region) gives the rows and bytes of one driver call, so a holder can size its submissions to bound what one call that never proves its access end keeps held. - BlobConfig's defaults assume nothing about the driver: one worker, and no limit of the backend's own on a call, so a group goes in one call unless the driver declares max_rows_per_op (batch_size then answers None). How many calls to overlap and how large to make them depends on the store, so a store's values come from its store config. - A fetch's hint goes unchanged to every get of a driver that declares hints; other drivers, and every publish, refuse a non-None hint with InvalidSubmission before dispatch. - BlobConfig.workers threads call put/get, which the driver protocol lets run concurrently; fetches go first, and once four fetch batches were taken while a publish batch waited, that publish batch goes next, so publishes do not starve. A batch counts as started before it enters the driver, which is what keeps a cancelled or faulted attempt from quiescing while the SDK still touches its slots. - One progress thread is the only observer of operations a worker could not settle; each batch records its results and its access end once. An operation the driver reports access_unprovable is counted (stats batches_unprovable) and polled again only until its results are recorded; its access stays open whatever the driver later says, so its attempt never quiesces and shutdown keeps its label. - DriverRejected from put/get is the one raise that certifies nothing escaped: those rows fault and their access counts as ended. Any other raise keeps them held. NOT_READY and DECLINED rows are undelivered, without a fault and without stopping the attempt. - A fault stops a fetch's later batches, since a prefix reuse has no use for the rows past it; a publish's rows stand alone, so it stops only when a call failed as a whole (a raise, a refusal, a wrong row count). After a fault, the outcome waits for the batches already started, so with several workers a later batch that fails first does not drop the earlier rows a prefix reuse needs. - Shutdown forces outcomes, wakes, and waits within its budget for the progress thread, or only until every attempt still open is pinned by a call that can never prove its access ended, after which its answer cannot change; it never calls the driver. The progress thread drains and revokes registrations once workers are gone and every attempt has quiesced, retrying failed revocations. - A driver call has no upper bound (Mooncake waits for every transfer to finish), so the progress thread logs one warning for any call past 60 s and stats() reports the oldest call's age. - A poll answered with the wrong result count breaks the driver protocol, so that batch becomes unprovable too, unless its access end was already recorded: counted, no longer polled, its slots held. That count, or a raise other than DriverRejected, is a protocol violation: the backend refuses every later submission with SubmissionRejected naming it and stops the attempts that still have batches queued, with the violation as their fault; batches already in the driver run on. - RegionBinding rolls back a registration failure that escaped nothing (DriverRejected, or RegistrationError without a handle: it releases the hold and raises ConstructionRejected) and keeps the hold on any failure that may have left a registration live. - RegionBinding registers only the regions with slots in memory the driver can access, HOST or, once the driver declares device_memory, DEVICE, each under its region's memory kind. Rows in any other region are refused with InvalidSubmission at submission, never skipped, and a binding with no region to register is refused with ConstructionRejected after releasing the hold. - Delivery masks are bool. An attempt that concludes normally hands its final masks to Outcome read-only, so Outcome keeps them without a copy; a forced outcome is copied, since started batches still apply. keys.key_stem checks a prefix once and gives the stem every key under it starts with, so the key format stays defined in keys.py alone while the backend builds each key as stem + hex(name); a batch's names are hex-encoded in one pass, since at real row sizes (0.25-2.5 MiB) a backend prepares hundreds of thousands of rows a second under the GIL. cpu_fixtures.py and test_blob_backend.py add the CPU behaviour cases over the in-memory driver. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
MooncakeDriver implements the StoreDriver protocol over a set-up MooncakeDistributedStore from SDK 0.3.13, zero copy from and into the lender's registered host slots: a put starts a put session per key, writes every row in one scatter (batch_put_from_multi_buffer_ranges), ends the sessions of rows written whole and revokes the rest; a get is one batch_get_into. Get stays on batch_get_into because get sessions lost on GB200: every SDK call waits for the GIL on return, so next to a busy Python thread the three session calls took 10-12 ms longer at p99 than one batch_get_into, and they read at 69-97% of its bandwidth with one or two workers. Put sessions came within 6% of batch_put_from or beat it at four workers, and they revoke a failed row only after the scatter finished, so a late write cannot land in space the master already gave another object (batch_put_from frees it on a timeout or a partial post). - Access end differs by call. The put scatter waits, without a timeout, until every fragment it posted is physically finished, so its return ends access whatever it answers. batch_get_into can answer TRANSFER_FAIL for a future that gave up while its transfer still runs, so a get ends access only when no row answers TRANSFER_FAIL. A raise or a wrong answer count leaves access unproven. The driver's put and get return only after every SDK call they make has returned, so an operation that did not end its access then never will: access_unprovable is true, it keeps its regions from being unregistered (DriverRejected) and makes drain answer False at once. Start, end, revoke and existence calls touch metadata only. The module docstring cites the 0.3.13 sources each rule was read from. - Put sessions are per key and per client, so the driver ends or revokes only sessions this put started, and ends only rows written whole, since ending marks the object complete. It never starts a session for a key it already holds: a key another of its puts is writing is NOT_READY at once, and a key whose session it left open is READY when batch_is_exist shows its object complete, whoever wrote it, and NOT_READY otherwise. - A session is left open when it cannot be ended or revoked, when its end went unanswered (RPC_FAIL or RPC_TIMEOUT, even if a retry is then refused, or a raise), since the master matches a late end on key and client only and would complete the next allocation of the key, and when a session start or a whole put raised or miscounted. drain revokes those once no call runs and none is unprovable. A row whose scatter answered TRANSFER_FAIL is left open too, but drain never revokes it: over TCP the initiator may give up while the target still applies data into the object's space, and a revoke frees that space at once, while the master frees an unfinished put only after its release timeout. - Codes become row results here and nowhere else: OBJECT_NOT_FOUND is a MISS; REPLICA_IS_NOT_READY on a get, a key this driver holds and a key another client holds uncommitted (OBJECT_ALREADY_EXISTS, batch_is_exist 0) are NOT_READY; a session start's INVALID_PARAMS, which can then only be something else, such as a cachelib master's 16 MiB - 16 B object cap, is a FAULT; a full pool (NO_AVAILABLE_HANDLE) is DECLINED; a key that already exists is READY once batch_is_exist confirms it; a row whose read lease expired is read once more, and that reread is READY or a FAULT, never a MISS or NOT_READY, since the first read already wrote the slot; an object of the wrong size and the rest are FAULTs with codes. An object shorter than its row lands in the slot's head before batch_get_into answers its size, so that row is a FAULT whose slot is not untouched; a longer one is refused without a transfer. - Admission and answer mapping take the common case batch-wise (every row in one registration; every row written or read whole) and fall back to row by row otherwise: at 0.25-1 MiB rows the driver handles hundreds of thousands of rows a second, all under the GIL. - Admission refusals (draining or closed, no rows, a repeated key, a row of more than one span or over 2 GiB, a span outside every live registration, a hint) raise DriverRejected before anything is recorded; after that put/get cannot raise, and any later fault becomes row faults with access left open. Rows must stay under 2 GiB because the put scatter answers in 32 bits. put/get run concurrently with each other and with drain, which waits for the calls in progress; shared state sits behind one small lock never held across an SDK call. - Capabilities: absent_is_miss, since the master answers OBJECT_NOT_FOUND for a key it does not hold, while failures keep their own codes and stay FAULTs; host memory only, so register refuses a DEVICE region before the SDK is asked; rows of one span, since batch_get_into reads each object into one buffer (a multi-span get would need batch_get_into_multi_buffers). - close() closes the SDK client once a drain returned True and every region is unregistered. Production never calls it: the SDK's close holds the GIL through RPCs, while the lender's hold is given back on the manager's thread; process exit releases the client. fake_mooncake.py implements only the SDK calls the driver may use, with 0.3.13's semantics over real bytes; test_mooncake_driver.py pins the mapping, the put session bookkeeping, the access-end rules, drain, close and a backend round trip over it, and runs NVIDIA#20217's driver_conformance checks over the fake SDK with every scenario hook the fake can stage: a get of one key held at the SDK call, a key another client is writing, a full pool, a raise inside a put, a failing read and a get whose access end cannot be proven. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
A deployment describes its store in a config file of its own, while the backend needs a MooncakeDriver over a set-up SDK client and a BlobConfig. mooncake_setup turns the one into the other. load_store_config reads the store config, a YAML file or the mapping parsed from one: a blob section, laid over MOONCAKE_BLOB_DEFAULTS (two workers, 64 rows and 128 MiB a call, with the reasons for the bytes cap beside them) for BlobConfig, and a mooncake section, read into MooncakeConfig. Unknown sections and keys are refused with their section, and so are values of the wrong type (YAML 1.1 reads 1.28e+8 as a float); BlobConfig refuses those of the blob section. An unreadable file or a YAML error is a ValueError too, and so is a master address that is not one host:port, such as an HA master list. open_sdk refuses TCP's legacy framing, which reports a write before it lands, the local hot cache, which can serve a torn read as a whole object, and MC_USE_TENT and MC_USE_TEV1, which select a transfer engine the driver's rules were not read for; logs the Mooncake environment before setup; imports mooncake.store only when called, with an install and NIXL-shadowing hint, and refuses it unless every installed distribution that lists the imported module is a 0.3.13 build with the RDMA and TCP transports; waits for the master; resolves the RDMA devices and the local address with the connector's rules (inlined, since the connector pulls in torch); and closes a client whose setup failed. A client in service is never closed: the driver's close() holds the GIL through unmount RPCs. open_mooncake_driver reads both sections, the blob one with the caller's reader (BlobConfig for the blob backend), before it opens the SDK, which waits for the master and mounts a segment, and returns the driver with what the reader returned. A process may open one driver per lender, each with an SDK client of its own. The factory that builds the backend for one lender from these pieces follows on top of RI-02, whose BackendFactory, OwnerCleanupFn and FactoryRejected it has to fit. test_mooncake_setup.py covers config errors, setup arguments, device detection, the refused TCP framing, hot cache, newer engine and SDK builds, the missing SDK and the half-open client, and the blob package's import layering, on a fake SDK module. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
The fake SDK shows what the driver asks for; only a real master shows whether the SDK answers that way. test_mooncake_real.py runs NVIDIA#20217's driver_conformance checks over the real pool, one fresh client per check, with the two scenarios a real master can stage: a key another client is still writing, which both directions find not ready, and a get held inside the SDK, which delays neither another get nor drain. It runs the backend between real clients for each answer the fake SDK models beyond them: a full pool that declines rows, a get that lands a shorter object's bytes and refuses a longer one, and a second publish of keys whose put sessions the same client still holds, also once the first put is stuck past the master's put_start_discard_timeout_sec (set to 2 s through a wrapper around mooncake_master): its keys stay skipped and never read as existing until it ends; and four workers round trip 2000 distinct rows byte for byte. Clients are closed only after every backend over them shut down cleanly, so a failing test cannot close one under a running call. The tested clients give the pool no segment and keep no local buffer, so every object sits in another process's segment and nothing can fall back to the SDK's copying path. running_segment_provider, in test_common/mooncake_utils.py, gives a pool its memory from that separate process. The file runs only with TRTLLM_MOONCAKE_REAL=1, and fails unless mooncake_master ships with an installed SDK build and both report 0.3.13; it prints both versions and records them as a test-suite property. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
… the blob backend The CPU and real-pool tests drive the backend with host slots of their own; none shows that what the staging lender lends is what the backend moves. test_blob_backend_round_trip.py does it end to end on a GPU: one real KV cache manager publishes a request's blocks through its lender's staging slots, another fetches them into a request of its own, and a test lease holder drives the lender as the runtime will (lend, wait for the view, submit it as it stands, wait for quiescence, mark what arrived, release). - Pages come back byte for byte and resume_bounds covers the prompt's three whole blocks, over the in-memory store and, with TRTLLM_MOONCAKE_REAL=1, a real Mooncake pool. - A fetch whose later blocks fault marks only the delivered prefix, so resume_bounds stops at the first gap. The in-memory store and driver come from NVIDIA#20217's testing package. The test builds each backend itself, as a factory will: it opens the driver, takes the lender's hold, registers the regions with RegionBinding, starts BlobBackend, and gives the hold back once the backend shut down. The lender's MemoryRegion names no memory kind, which NVIDIA#20117's SlotRegion requires, so the test registers host-memory SlotRegions made from the lender's regions, as a runtime must until the lender's regions name their memory kind. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…ch rows Measured on two GB200 nodes over RDMA, with the main thread idle or next to a thread running Python bytecode nonstop (busy, standing in for a busy executor), 4 MiB rows, four workers unless noted: - Publishes now take at most max(1, workers - 1) workers, so a fetch never waits behind publish calls on every worker. Busy, with 128 MiB batches, a 16 MiB fetch next to two publish views in flight rose at p99 by 12.9 ms over no load instead of 22.3 ms without the cap (where a fetch without load already takes 16 ms at p99 waiting for the GIL). Idle the gain is less sure: 2.7 vs 6.3 ms in the one job that ran both, but another job without the cap also saw 2.7 ms. The cost is busy publish bandwidth, 14.6 vs 19.3 GB/s (-24%) in the same job, since only three workers publish; idle it stays at 90 GB/s. - MOONCAKE_BLOB_DEFAULTS workers 2 -> 4. One put call alone does not fill the links, and with the cap two workers would leave publishes one: next to fetches they moved 48.7 GB/s with two workers and 80.4 GB/s with four. Before the cap, four workers already put 88.5 vs 84.6 GB/s idle and 19.5 vs 11.6 GB/s busy. - MOONCAKE_BLOB_DEFAULTS rows_per_batch 64 -> 256. A rank's row is one KV block of a layer group, 0.25-2.5 MiB for common models at TP1-8 (bf16, 32 tokens per block), so 64 rows kept calls of 0.25-1 MiB rows at 16-64 MiB, under the byte cap. With 256 (median of three passes), 0.25 MiB rows put 52 vs 41 GB/s idle and 6.1 vs 2.3 GB/s busy; 0.5 MiB rows put 79 vs 69 idle and 11.5 vs 4.6 busy; busy throughput at 0.25-1 MiB rows rises 1.4-3.2x across jobs. The costs: 0.25 MiB rows get 8% less idle (85.4 vs 92.7 GB/s); fetch latency under publish load rises at p99 by +2.1 vs +1.0 ms (0.25 MiB, idle), +20 vs +13 ms (0.25 MiB, busy) and +40 vs +17 ms (0.5 MiB, busy); and a 1 ms sleeper wakes 0.74-1.34 ms late at p99 with 0.25 MiB rows (two jobs), about our 1 ms target and over it in one job, while 0.5 MiB rows and up stay under it at 0.53-0.8 ms. 512 rows would make it 1.3-2.0 ms. Rows as small as 0.25 MiB come only from small models with FP8 KV at TP8. - bytes_per_batch stays 128 MiB. Busy, each call's return waits for the GIL, so 128 MiB batches move 3.0-3.7x (put) and 1.7-2.3x (get) what 32 MiB batches do; across jobs busy 128 MiB batches put 15-19 and get 31-43 GB/s. Idle the bandwidth is the same (put 90.6 vs 93.4, get 94.3 vs 95.3 GB/s). They cost fetch latency under publish load: +2.4 vs +0.8 ms at p99 idle, +12.6 vs +1.3 ms busy. Our fetch latency target, at most twice the no-load p99, is not met idle (4.5x, 0.68 -> 3.09 ms): a fetch still shares the GIL and the transfer engine with the publishes. Busy it is 1.8x. The interpreter's switch interval matters more than either knob: at 0.5 ms instead of 5 ms, busy puts move 68-74 GB/s and gets 84-90 GB/s at 128 MiB, while the busy thread keeps 87-91% of the loops per second it runs alone (99% at 5 ms). test_blob_backend.py pins the cap. test_mooncake_real.py's concurrent publishes of the same keys need two publishes in the store at once, so they now run three workers. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
7f29395 to
f2d736b
Compare
This branch has not been deployed
3 participants
Summary
This PR adds the store backend for RI-08D, with Mooncake as its first store.
BlobBackendserves [TRTLLM-16186][feat] Define backend contract for KV cache transfers #20117'sBackend/Attemptcontract on top of [None][feat] Add the blob store driver boundary and an in-memory test driver #20217's privateStoreDriverboundary.MooncakeDrivermaps that boundary onto the pinned Mooncake 0.3.13 SDK.RegionBindingregisters the lender's regions with the driver under the lender's hold, andopen_mooncake_driveropens the store from its own config file. The factory that RI-02's TransferManager calls is left to a follow-up on top of [None][feat] Add shared KV transfer ownership and progress (RI-02) #20268, whoseBackendFactoryshape (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
Nonehints 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.Design
BlobBackend(RegionBinding(regions, driver, hold), wake, config)is what a factory builds; the RI-02 factory adapter follows on top of #20268RegionBinding(blob/backend.py) registers only under a hold the caller took, and releases it on the manager's thread only after every region is confirmed unregistered, or at once when a failed registration was fully rolled backBlobConfig.workers(default 4) bounds concurrent driver calls. Publishes take at mostmax(1, workers − 1)of them, and fetches go first, with one publish batch after every four.Nonehint raisesInvalidSubmissionbefore dispatch, and the driver leavesDriverCapabilities.hintsat its default, Falsetest_blob_backend.py::test_shutdown_never_calls_the_driver_and_unregisters_in_the_background; every CPU endpoint also checks that the hold is released on the manager's threadmemory_kindit can access, and rejects rows it does not recognize withInvalidSubmission, never skipping themRegionBinding.registeredinblob/backend.py;BlobBackendchecks each group's region before queueingOutcomeits final masks read-only, so they are not copiedBlobAttempt.poll()andquiesced()report separate facts, and each is set once under the backend lockwake()runs after the state is published, outside the lock.shutdown(timeout)never calls the driver; drain and unregister run on the progress thread.batch_is_existsees it committed.blob/mooncake.py; see "Mooncake mapping" belowtest_mooncake_real.py;test_blob_backend_round_trip.py, which uses a real lender on GPU; #20217'sdriver_conformancesuite run againstMooncakeDriverkv_sharing.backends: [{type: mooncake, config: <path>}], with store parameters in the store's own fileblob/mooncake_setup.pyreads a YAML file withblobandmooncakesections; no LLM args are addedWhere the implementation departs from the design text:
cancel.Mooncake mapping
put calls the session API:
batch_put_session_start, thenbatch_put_from_multi_buffer_ranges(one scatter for the whole batch), thenbatch_put_session_end. Rows that did not land getbatch_put_session_revoke.wait()has no timeout. It returns only once every fragment has physically finished.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 uses
batch_get_into, one SDK call per batch.batch_get_into.Results:
Access-end evidence:
batch_get_intoanswers every row and no answer is TRANSFER_FAIL.access_unprovableturns true: a TRANSFER_FAIL on a get, an exception, or a wrong answer count. The backend then counts the call, stops polling it and keeps its attempt's slots. It treats a driver'spollthat answers the wrong number of rows the same way, unless that batch's access end was already recorded.Capabilities:
getandput.absent_is_miss: the master answers OBJECT_NOT_FOUND for a key it does not hold, before any transfer.device_memory=False). The binding registers only the regions whosememory_kindthe driver can access, which is HOST for Mooncake, and passes each region's kind on. A row in any other region raisesInvalidSubmission, and a lender with no such region is refused withConstructionRejected.multi_span=False):batch_get_intoreads each object into one buffer; a get of several spans would needbatch_get_into_multi_buffers.Scope
backend.py,keys.key_stem, and CPU behavior tests over #20217'sMemoryDrivermooncake.py, a fake SDK, mapping tests, and the conformance suite run over the fake SDK with its scenario hooksmooncake_setup.py, configuration and layering testsTRTLLM_MOONCAKE_REAL=1) againstmooncake_masterand a segment provider, plus the conformance suite on the real storemax(1, workers − 1)workers; defaultsworkers=4,bytes_per_batch=128 MiB,rows_per_batch=256Production code: +1680 / −6 lines. Tests: +3259 lines.
Verification
At
0e78604f12, on main083447cd22's prebuilt wheel with this branch's Python files on top (main after083447cd22, up to the PR's baseee7c59aa78, changes no C++):test_backend_contract.py57 passed;kv_cache/sharing531 passed and 16 skipped (15 need a Hopper or Blackwell GPU and pass on GB200 below; 1 needs a real pool). pre-commit, run in the development container, is clean over the six commits._torch/disaggregation/blob,_torch/disaggregation/test_backend_contract.pyand_torch/executor/kv_cache/sharing. The 13tests/unittest/disaggregated/files that [TRTLLM-16186][feat] Define backend contract for KV cache transfers #20117 touches give 489 passed, 1 skipped and 2 xfailed (marked xfail in the source) on the same image.tensorrt_llm._torch.disaggregation.testing.driver_conformance) runs againstMooncakeDrivertwice.Objects of the wrong size. Both checks pass.
real_client.cpp:5134-5141).real_client.cpp:5169-5180). The driver reports FAULT, its access ends normally, and nothing outside the row is touched.Performance. The microbenchmark behind these numbers will follow in a separate PR.
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. Examples:
Please confirm
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).batch_get_session_start+ ranges + end), whose access ends even on TRANSFER_FAIL. It costs two more SDK calls per batch: next to a busy Python thread, a 32 × 1 MiB get took 16 ms at p99 against 5.6 ms. This PR keepsbatch_get_intoand asks for a containment profile that is not executor-fatal for a store fetch, since the request can recompute instead.RegionBindingon the manager's thread before the executor loop starts, because the lender's hold bookkeeping has no lock.RegionBinding.register()raisesConstructionRejectedwhen nothing stays registered and the hold is released; the RI-02 factory adapter maps that toFactoryRejected.keys.key_stemis a small addition to [None][feat] Add the blob store driver boundary and an in-memory test driver #20217'skeys.py.blob/store.py.Known limitations
sys.setswitchinterval(0.0005), busy throughput rose 2–7x, and the busy thread lost 9–13% of its loop rate while transfers ran. This is an observation; the PR does not change the switch interval.TLLM_NUMA_AWARE_WORKER_AFFINITY) already does this; a rank whose affinity is turned off or set elsewhere can end up on another node's two NICs.probe=False: the Store Backend design leaves a public probe out of this version, and the [TRTLLM-16186][feat] Define backend contract for KV cache transfers #20117 contract has nothing that would call it.batch_is_existcould back a probe later, with two gaps against [None][feat] Add the blob store driver boundary and an in-memory test driver #20217's probe contract. In 0.3.13 it answers 0 both for an absent key and for one still being written (the master'sExistKeyneeds a readable replica), so it cannot report NOT_READY. It also takes no timeout, so the driver would need its own thread to answer within the budget.MemoryRegion(pyexecutor/kv_cache/sharing/_types.py) has nomemory_kind, which [TRTLLM-16186][feat] Define backend contract for KV cache transfers #20117'sSlotRegionnow requires.RegionBindingSlotRegions withmemory_kind=HOSTfor a staging lender. The GPU round trip does this withHostRegion.RegionBindingreads every region's kind when it is constructed, so a region without one fails there, before anything is registered. The caller still owns the hold it passed in and must release it.MC_STORE_LOCAL_HOT_CACHE_SIZEother than empty or 0 is refused when the driver is opened (open_sdkraisesValueError). The 0.3.13 SDK copies a get's slot into its hot cache before it checks the read lease and keeps the copy after the master evicts the object, so a torn LEASE_EXPIRED read could later be served as a whole object.close(): it holds the GIL and sends RPCs, while backend teardown runs on the manager's thread.blobandmooncakesections of its YAML) is read into frozen dataclasses with strict type checks rather than Pydantic models; it does not become LLM args. Converting them to Pydantic is a small change if reviewers prefer it.Notes
Stack: main
ee7c59aa78, which already contains #20217 (032ea6d414) → backend contract #20117 (8ac96cf4ed, Xiaowei's three commits, cherry-picked) → this PR.Review only the last six commits: Mooncake-only diff.
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