Repository navigation
Conversation
Contributor
|
Documentation preview: https://vllm--59276.org.readthedocs.build/en/59276/ |
mergify
Bot
added
documentation
Contributor
Author
|
Multi-host validation (MI355X, AMD Pollara RDMA): three engine hosts and the
Replicas made by readers ( |
Add a shared KV connector and an embedded runtime for single-engine, single-node offloading to MORI-managed host DRAM. Shared connector / base: - Register UMBPStoreConnector and separate scheduler, worker, transfer metadata, runtime contract, and statistics responsibilities. - Plan hash-keyed, cache-group-scoped objects with rank-local identities and logical completeness checks across participating workers. - Integrate external lookup, asynchronous load completion, ranged and partial-prefix transfer plans, and hybrid cache hit-window semantics. - Implement eager stores and free-queue/store-event-driven lazy offload, retaining GPU block ownership until transfer completion. - Track request preemption, cancellation, partial failures, cache resets, eviction events, and transfer metrics through the connector lifecycle. Embedded runtime: - Register GPU KV buffers and use MORI ranged pointer APIs for stores to host DRAM and restores to GPU memory. - Keep DRAM allocation, residency, publication, capacity, and eviction authoritative in MORI rather than maintaining a connector-owned cache. - Configure per-rank or total DRAM capacity, watermarks, NUMA placement, huge pages, prefaulting, shared-memory backing, and bounded execution. - Handle asynchronous lookup and transfer completion, timeouts, retries under DRAM pressure, and per-object errors. - Limit embedded scope to DRAM offloading; standalone and distributed runtimes, cross-engine sharing, and SSD tiers are outside this change. - Keep pipeline parallelism rejected and layer-wise stores disabled until a safe native full-object publication protocol exists. Tests and documentation: - Include shared and embedded unit suites, runtime contracts, MORI GPU range and eviction tests, token-restore integration tests, and a status guide describing configuration, supported scope, and remaining work. - Test coverage does not imply every model or topology case passes. This consolidates the functional snapshot at 264c17e without changing tracked contents. Subsequent simplification and performance work remains a separate commit. History and message changes are verified by exact Git tree equality; no new unit tests, GPU tests, or model evaluations are claimed by this rewrite. AI assistance was used in development and commit-history consolidation. Co-authored-by: Codex <codex@openai.com> Co-authored-by: Cursor Agent <agent@cursor.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>
Simplify the connector and lifecycle handling, and reduce repeated work in external lookup and batched load preparation. MORI remains authoritative for object residency; no potentially stale hit-state cache is introduced. Shared connector / base: - Remove unused identity and capability scaffolding and consolidate transfer state, completion accounting, and materialization fallbacks. - Preserve lazy free-queue/store-event ownership and all-rank completion; finalize stores exactly once even when a no-forward step skips the normal model-runner save hook. - Reuse request-local lookup identities. Apply the unchanged core KVCacheCoordinator to a read-only per-query hit view instead of allocating, indexing, and freeing shadow cache blocks. - Carry whole-block loads in compact columns with shared request and generation metadata. Preserve rank localization, failure-to-block mapping, byte accounting, and fallbacks for special layouts. - Filter null GPU source blocks before eager residency checks and store plan construction. Valid sources still require authoritative MORI lookup. Embedded runtime: - Keep DRAM-specific option validation and rank-capacity resolution in the adapter, and move the CPU byte-copy fake into test utilities. - Reuse registered per-group layouts to construct batched load addresses directly and consume compact load metadata through native completion. - Retain transfer buffer ownership across timeout and cancellation, and preserve safe teardown and scheduler/worker lookup communication. - Do not change native MORI, model code, kernels, or core cache algorithms. Native scoped write admission remains unimplemented. Add precise annotations, avoid conflicting local variable types, use monkeypatch for test method replacement, and test metadata serialization through vLLM Msgpack rather than a direct pickle import. Historical validation of the GPU-tested snapshot before lint cleanup, using MI355X hardware, a pinned ROCm development image, and real weights: - Unit and runtime contracts: 208 passed, 1 skipped. The skipped case requires four visible GPUs for TP2/PP2. - Differential checks: 206 passed, including 1,600 changing-residency windows. This overlaps the unit suite; unsupported DCP+SWA cases are rejection tests, not evidence of support. - Llama-3.2-1B TP1 eager, lazy, synchronous lookup, and TP2: actual external restore; identical 32 decoded tokens, text, and finish reason; chosen-logprob error below 0.02; clean engine exits. - GPT-OSS C8/C32: 240 replay requests with external hits and no failed load. - Same-GPU GPT-OSS eager restore without diagnostic hooks, 20 replays per aggregate: baseline 25.8554 ms, candidate 24.2346 ms, Simple 18.1520 ms. Each candidate replay restores 10,000 tokens, 633 objects, and 248,905,728 bytes. This is not the full serving test plan or a claim of Simple parity. Historical validation after lint cleanup, on the committed source tree: - Combined unit, contract, and differential suites: 252 passed, 5 GPU-required tests skipped; run without GPU device access. - Pre-commit checks, including mypy 3.10 and forbidden imports, passed; the manual mypy 3.12 hook passed. Only Actionlint was skipped because an unrelated Go dependency download failed TLS certificate verification. This change does not modify workflow files. - The preceding GPU and performance results describe the pre-lint snapshot, not a post-lint rerun of those same cases. Known limitation: strict GPT-OSS cold/replay numerical checks still fail the unchanged 0.02 logprob tolerance. Simple and old UMBP controls also fail, but the cause is unresolved and an additional regression is not ruled out. A diagnostic verifies identical stored/restored bytes for all 39 objects (15,335,424 bytes); added HIP synchronization means it does not prove uninstrumented async ordering or strict output equivalence. This preserves the tracked tree and optimization diff of 954d661. History and message changes are verified by exact Git tree equality; the validation above is historical, not a new test run performed for this rewrite. Co-authored-by: Codex <codex@openai.com> Signed-off-by: CharlesWu <Charles.Wu@amd.com>
vLLM copies a restored Mamba state into the request's running slot while it prepares inputs, before a synchronous load in start_load_kv could have landed, so the copy reads stale state. With load_async=false, Kimi-Linear GSM8K (256 questions, conc 64) fell from 89% to 52% on restored prompts, for both layer-wise and bulk loads, while single requests still matched, since a lone load finished in time. Waiting for the load inside start_load_kv does not help because the copy has already happened. The scheduler now loads asynchronously whenever the KV cache has a Mamba group and warns that load_async=false is ignored. With it the same run restores at 89.8% (layer-wise), 88.7% (bulk) and 89.8% (async). Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>
With load_async, a request whose KV is loading is not scheduled until get_finished reports its load, so no forward pass reads the blocks being filled. wait_for_layer_load still waited for those loads: the batched load of each request is kept under the job key None, and _drain_load_jobs selects None jobs for every layer name. The first attention layer of each step therefore blocked until every in-flight asynchronous load had finished, putting other requests' restores on the running batch's critical path. wait_for_layer_load now returns at once for asynchronous loads. Preemption without request ids, which must settle loads before vLLM reuses their blocks, drains the load jobs directly instead of going through it. Signed-off-by: Yichao Zhu <yichao.zhu@amd.com>
The embedded lookup socket path hashed only the key namespace and the rank's (tp, pcp, dcp, pp) coordinates. The namespace identifies the model and layout, not the engine, so with data_parallel_size > 1 on one host, or two engines serving the same model, worker tp0 of each resolved to the same /tmp path. The second one unlinked the first one's socket and bound its own, so one engine's scheduler silently queried the other engine's private pool: hits it could not load, while its own pool went unused. The path now also hashes the engine's data-parallel index and an optional lookup_instance, which separates independent engines on one host. A lookup server that finds its path served by a live listener fails startup and names lookup_instance instead of taking the socket over; a stale socket file left by a dead process is still replaced. Signed-off-by: Yichao Zhu <yichao.zhu@amd.com>
YukioZzz
force-pushed
the
yichaozhu/umbp-distributed
branch
from
September 30, 2026 16:31
6f5319a to
bf6acae
Compare
10 of 18 tasks
Loads of whole blocks already skip plan materialization: the embedded handle computes every block's addresses from the registered per-group layouts with numpy. Stores still materialized one BlockTransferPlan per block on the forward thread, building one KVRange per layer. For a Llama-3.1-8B 8192-token prefill that is 512 plans and 16384 ranges, 8.6 ms of Python in wait_for_save before the step can finish. Stores of whole blocks now take the same path through store_blocks: the worker only localizes the keys, and the handle builds the put arguments from the registered layouts, with the object size of each block taken from its layout. Plans with token ranges, sub-block layouts or a missing layout still fall back to materialized ranges. The MORI call receives identical arguments either way. Signed-off-by: Yichao Zhu <yichao.zhu@amd.com>
Part of vllm-project#57187. Standalone mode stores KV in one umbp_standalone_server per host. Every engine, data-parallel rank, and worker on the host attaches to it over a Unix-domain socket, so a prefix stored by one of them is restored by any other and survives engine restarts. Workers reuse the MORI worker handle from embedded mode, including its batched block loads, with a client configured for the server; they register GPU KV allocations through HIP IPC and serve no lookup socket. The scheduler handle queries the server directly. Runtime failures degrade instead of propagating: lookups report misses, loads mark their blocks for recompute, stores fail without raising, and publish, clear, and close tolerate a server that is gone. An unreachable server still fails engine startup. A stalled server must not turn into corrupted KV or a failed step. A store that outlives timeout_ms stays pending rather than raising from poll(): MORI still reads its source blocks, and the store event keeps them pinned until the server answers. While such a store is blocked, new loads and stores fail immediately with "transfer not attempted" instead of queueing behind it, and shutdown does not wait for it. The adapter validates its own options: it rejects TCP endpoints, total_capacity_bytes, and DRAM sizing options unless auto_start spawns the server, since the server owns its pool. Pipeline parallelism is rejected as in embedded mode. reset_cache clears the shared pool for every attached engine. Signed-off-by: Yichao Zhu <yichao.zhu@amd.com>
Distributed mode lets engines on different hosts reuse each other's KV: a umbp_master indexes objects cluster-wide, and every worker is a pool node that serves its own host-DRAM slice to peers over MORI-IO RDMA. Scheduler lookups ask the master through a client with no pool, no I/O engine and no peer service, so the scheduler process allocates nothing and the master never routes a put to it. Workers join under an identity that includes the pid, because the master refuses to re-register a node id until its previous holder expires, which would otherwise block an engine restart. Each worker binds peer_service_port (and io_engine_port when set) offset by its physical GPU index, which is unique per host across engines and DP ranks. node_address defaults to the vLLM host IP. The adapter rejects options distributed nodes would silently ignore (shared-memory and watermark DRAM knobs) and local-mode options, validates the master address and ports, splits total_capacity_bytes across an engine's ranks, and fails at startup when ranged_scratch_size cannot hold one KV object, since every off-node transfer would then fail. Warnings flag page padding and large 4 KiB-paged pools, which NICs with small translation caches cannot register for RDMA. The DRAM option checks shared with embedded mode move into helpers, and the local modes reject every distributed-only option. Resetting one engine does not clear the shared pool. An unreachable master degrades lookups to misses and loads/stores to failed transfers. MORI's lookup RPC has no deadline, so against a master that accepts but never answers it blocks until the master recovers; lookups run on the scheduler's critical path, so each is bounded by lookup_timeout_ms and reported as all-miss, and while one is still blocked later lookups miss without asking the master. Workers reuse the standalone shared-pool handle, so a store stalled on the master stays pending with its source pinned. Pipeline parallelism is rejected in every mode. Unit tests use a fake master/node cluster; integration tests run a real umbp_master with two RDMA pool nodes on two GPUs, covering byte-exact restore, index cleanup when a node exits, master loss, a restarted master relearning every live node, and a SIGSTOPped master costing one bounded lookup. Signed-off-by: Yichao Zhu <yichao.zhu@amd.com>
A distributed lookup can keep reporting an object that no longer loads:
the master lists a crashed node's objects until it expires that node
(30 s by default). vLLM recovers a failed load by rescheduling the
request, whose next lookup hits the same objects, so the request loaded,
failed and rescheduled every step until the node expired: 691 recoveries
and a 8.9 s TTFT for one 2k-token request in a two-host run with a 9 s
expiry.
Workers now report the rank-local keys of failed loads with their error,
and the scheduler treats those keys as misses for
load_failure_quarantine_ms. It defaults to 30 s in distributed mode and
to 0 (off) otherwise, since local pools forget such objects at once. A
load the runtime refused without attempting it ("transfer not
attempted", e.g. while a store is blocked) says nothing about the object
and is not quarantined. A later store of the same key does not lift the
quarantine, because a pool that still lists the dead copy reports that
store as already done. With it, the same run recovers once and serves
the request with a 0.05 s TTFT.
Signed-off-by: Yichao Zhu <yichao.zhu@amd.com>
YukioZzz
force-pushed
the
yichaozhu/umbp-distributed
branch
from
October 1, 2026 03:02
bf6acae to
e8bd5e3
Compare
This was referenced Oct 9, 2026
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Adds the distributed mode of the UMBP KV connector (RFC #57187): a
umbp_masterindexes KV objects cluster-wide, and every vLLM worker is a pool node that serves its
host-DRAM slice to peers over MORI-IO RDMA. Any engine attached to the master reuses
KV stored by any other, on the same host or another one.
Stacked on the standalone PR (#59275) and #58568, rebased onto #58568's current
head
37832fa99. Please review only the last two commits:[UMBP] Add distributed mode backed by a cluster-wide master- runtime, optionvalidation in the adapter, docs (
docs/features/umbp_distributed.md), unit andreal-RDMA tests.
[UMBP] Treat objects that just failed to load as misses for a while- workersreport failed load keys, and the scheduler treats them as misses for
load_failure_quarantine_ms(30 s in distributed mode, off elsewhere). Without ita crashed node's objects caused 691 load-fail/reschedule cycles for one request
until the master expired the node. Loads refused without being attempted are not
quarantined.
The rebase dropped
[UMBP] Keep a timed-out store's source pinned until MORI returns: #58568's new worker already keeps stores pending across steps under astore event, and the runtime-side stall handling moved into the standalone PR's
shared-pool worker handle, which this mode reuses.
Design
so the scheduler allocates nothing and the master never places objects there.
binds
peer_service_port + physical GPU index.lookup_timeout_ms: MORI's lookup RPC has no deadline, and ahung master would otherwise stall the scheduler.
(
kv_load_failure_policy=recompute).Test Plan
tests/v1/kv_connector/unit/test_umbp_distributed.py.umbp_masterand two RDMA pool nodes on two GPUs:tests/v1/kv_connector/umbp_integration/test_mori_distributed_gpu.py(byte-exactcross-node restore, node exit, master loss / restart / SIGSTOP). Peer ports are
now drawn below the ephemeral range so concurrent suites cannot take them.
Test Result
passed, 1 skipped (needs four GPUs).
99.2% of prompt tokens restored externally on average.
64): cold on host 1 89.23%, RDMA hits from host 2 88.78%, host 3 reading host 2's
replicas after host 1 is stopped 88.93%, a fresh host-1 engine reading host 3's
copies after host 2 is killed 88.86%; zero failed loads.
Known limits
registers only when backed by reserved 2 MiB hugepages
(
dram_use_hugepages=trueplusvm.nr_hugepages); MORI aborts the worker when aregistration fails. The same limit blocks Mooncake's segment registration.
nodes, reconnect backoff.
Draft until the stack below it lands.