[KVConnector] native scale-up KV connector (HIP VMM, kv_connector=native)#1499
Open
carlushuang wants to merge 4 commits into
Open
[KVConnector] native scale-up KV connector (HIP VMM, kv_connector=native)#1499carlushuang wants to merge 4 commits into
carlushuang wants to merge 4 commits into
Conversation
…nsfer
Adds a small VMM (Virtual Memory Management) helper that lets one process
export a GPU buffer as a POSIX fd and another process import it, grant its own
device peer access, and copy directly over the fabric (XGMI). This is the
transport foundation for a single-node (scale-up) prefill/decode KV connector.
Unlike legacy hipIpc handles, VMM shareable handles + explicit hipMemSetAccess
give reliable cross-process peer access at high concurrency and do not require
the source GPU to be in the consumer's visible set (so prefill/decode keep
natural device==rank placement; no reorder or base_gpu_id changes needed).
- atom/kv_transfer/disaggregation/xgmi/{__init__,vmm}.py: VmmBuffer API
(alloc / export_fd / import_fd / tensor). The tiny C++ helper is
JIT-compiled lazily on first use (never at import time), so importing the
module is cheap and safe on CPU-only hosts.
- tests/test_xgmi_vmm_transfer.py: cross-process test (producer on GPU0
exports, consumer on GPU1 imports, peer-copies and verifies). Skips with
fewer than 2 GPUs or no VMM support.
Validated on 8xMI355X (gfx950, ROCm 7.2.4): cross-process peer copy is correct
at high concurrency where legacy hipIpc corrupts/faults.
Refs #1483
Rename the xgmi/ package to native/ (it depends on no third-party transport)
and wire the HIP VMM transport primitive into a registered KVConnector.
- atom/kv_transfer/disaggregation/native/native_connector.py:
NativeConnector + NativeConnectorScheduler. Scheduler is transport-agnostic
(transfer_id<->request_id, ConnectorMetadata). Worker uses a VMM staging
buffer + direct hipMemcpy peer copies over XGMI: the consumer exports a VMM
staging fd over a UNIX side channel; the producer imports it, gathers the
request's KV blocks straight into it, and the consumer scatters into its KV
pool. No RDMA, no IPC-handle churn, no host staging.
- factory.py: register "native".
- native/README.md: when to use it, selection, and the 4P4D launch recipe
(no "protocol" field, natural device==rank placement).
- _vmm_ext.cpp / vmm.py: add vmm_ptr + vmm_copy (raw peer DtoD copy).
Users select it with:
--kv-transfer-config '{"kv_connector":"native","kv_role":"kv_producer"|"kv_consumer",...}'
Scope: v1 wires the VMM primitive (validated cross-process by
tests/test_native_vmm_transfer.py) into the connector interface; requests in a
step are transferred sequentially and the DeepSeek-V4 slot/index-region fast
paths + a concurrent staging pool are follow-ups. End-to-end 4P4D serving
validation is tracked in #1483.
Refs #1483
carlushuang
changed the title
carlushuang
changed the title
The producer must cache reqs_to_save under the request seq.id (the transfer handle the consumer requests, per request_finished), not params["transfer_id"] (which defaults to 0). Symptom before the fix: the first request succeeded (seq.id 0 == default 0) but subsequent requests timed out waiting for the prefill entry, then gathered nothing. Validated end-to-end on 8xMI355X, DeepSeek-V4-Pro, single-node 4P4D (kv_connector=native): correct output, and no meaningful perf regression vs mooncake protocol=hip under identical config (TP=4, util 0.85, max_num_seqs 128). Serving (ISL/OSL=1024) within ~0.5%; prefill-isolation within 1-5% (staging gather/scatter overhead). Both use hipMemcpyPeer over XGMI.
Collaborator
Author
E2E DSV4 4P4D validated — no perf regression vs mooncake-hipBrought up single-node 4P4D on 8×MI355X (gfx950) with Prefill-isolation (ISL=1024/OSL=16), req/s | median TTFT:
Real serving (ISL/OSL=1024), out tok/s | TTFT | TPOT:
No meaningful regression: serving-level within ~0.5% (TPOT identical); prefill-isolation within 1–5% (native's staging gather/scatter overhead — 91 block + 62 slot + compressor-state copies/req vs mooncake's batched transfer). Both use Notes: native ran c16 cleanly (no hang), used disjoint visibility (no reorder trick), and needs no third-party module. A live fix went in for the |
…xgmi-vmm-transfer # Conflicts: # atom/kv_transfer/disaggregation/factory.py
ZhangLirong-amd
approved these changes
Jul 8, 2026
Jasen2201
approved these changes
Jul 8, 2026
| self._prefills_cv.wait_for( | ||
| lambda: transfer_id in self._prefills, timeout=_PREFILL_WAIT_S | ||
| ) | ||
| src_block_ids, src_slot = self._prefills.get(transfer_id, ([], -1)) |
Contributor
There was a problem hiding this comment.
memory leak risk, after finish, self._prefills should pop,
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
3 participants
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.
Native single-node KV connector (HIP VMM)
A fully in-tree prefill/decode KV connector for the single-node (scale-up / XGMI) case, using only the HIP Virtual Memory Management API — no MoRI, no Mooncake. Legacy hipIpc GPU-direct paths are unreliable cross-process here; VMM shareable handles are reliable and don't require the source GPU to be visible to the consumer (so natural
device==rank, no reorder trick). Refs #1483.What it adds
atom/kv_transfer/disaggregation/native/vmm.py(+_vmm_ext.cpp):VmmBuffer(alloc / export_fd / import_fd / tensor) +copy. C++ helper is JIT-compiled lazily (never at import → CPU-CI safe).native_connector.py:NativeConnector+NativeConnectorScheduler. Consumer sends a VMM staging fd over a UNIX socket; producer imports it, gathers the request's KV blocks + SWA slots + compressor state into it (hipMemcpy peer over XGMI); consumer scatters into its KV pool.kv_connector="native"in the factory;native/README.mdhas the 4P4D launch recipe.tests/test_native_vmm_transfer.py: cross-process transfer test (skips with < 2 GPUs / no VMM).Usage
No
protocolfield (always XGMI); prefillHIP_VISIBLE_DEVICES=0,1,2,3/ decode4,5,6,7(no reorder).E2E result — DeepSeek-V4-Pro 4P4D, no perf regression vs mooncake-hip
8×MI355X (gfx950), identical config (TP=4, util 0.85,
--max-num-seqs 128, default chunk):Serving-level within ~0.5% (TPOT identical); prefill-isolation within 1–5% (native's staging gather/scatter overhead). Both use
hipMemcpyPeer/XGMI. Native also ran c16 cleanly (no hang), unit test passes, black+ruff clean.Scope / follow-ups
v1: requests transferred sequentially per step; follow-ups are a concurrent staging pool and batching the per-region copies (would close the small prefill-iso gap).