Skip to content

[Bugfix][Store] Open CUDA IPC via exporter PCI bus ID - #5006

Open
huangxiaobai-ydmy wants to merge 3 commits into
kvcache-ai:mainfrom
huangxiaobai-ydmy:fix/issue-4980-cuda-ipc-device-ordinal
Open

huangxiaobai-ydmy wants to merge 3 commits into
kvcache-ai:mainfrom
huangxiaobai-ydmy:fix/issue-4980-cuda-ipc-device-ordinal

Conversation

@huangxiaobai-ydmy

@huangxiaobai-ydmy huangxiaobai-ydmy commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Test plan

  • Real client CUDA_VISIBLE_DEVICES=0,1, dummy CVD=1, buffer on GPU 1, no P2P: put/get succeed (opens on physical GPU 1 via PCI bus ID)
  • Dummy CVD=0 / buffer on GPU 0: still OK
  • Dummy CVD=0,1 (dev 1): still OK (no regression vs opening on exporter GPU)
  • If real client cannot see the exporter PCI bus ID: open fails with INVALID_PARAMS (clear path)
  • CI format + unit/integration jobs green

Dummy-client exports record the exporter's CUDA device ordinal, which is
relative to that process's CUDA_VISIBLE_DEVICES. Opening with
cudaSetDevice(handle.device_id) maps every worker's buffer onto the real
client's GPU 0 when CVD layouts differ. Match TENT NVLinkTransport: open
on the current device with lazy peer access and remember that local
ordinal for close.

Fixes kvcache-ai#4980

Signed-off-by: YDMY_Kirito <202200820046@mail.sdu.edu.cn>

(Rebased onto main after kvcache-ai#5030; packaging commits dropped.)
@github-actions github-actions Bot added run-ci Store size/S 0-99 changed lines (additions + deletions) labels Oct 27, 2026
@0z5a

0z5a commented Oct 27, 2026

Copy link
Copy Markdown
Contributor

Thanks for the quick fix. As the reporter of #4980, I don't think opening on the caller's current device fixes it on CUDA builds, and I think it breaks a layout that works today.

Where the open runs

  • CudaIpcBufferMapping::Open is called from RealClient::batch_write_from_cuda_ipc_dummy_helper and batch_get_into_cuda_ipc_dummy_helper, on the real client's RPC thread.
  • On CUDA builds nothing selects a device for that thread; set_context_if_needed only acts for Ascend.
  • The "current device" is therefore device 0 on a fresh thread. Any SetContext() call on that thread (for example in RuntimeAccelerator::Copy*) leaves its own device selected afterwards, because none of them restore the previous one.
  • TENT's NVLink transport also avoids the peer's ordinal. But the real client serves dummy clients on different GPUs, so it has no single "own" device to open every handle on.

Against the table in #4980 (real client CUDA_VISIBLE_DEVICES=0,1, 2× L40S without P2P):

dummy CVD buffer on GPU before this PR, opens on this PR opens on (fresh RPC thread) expected with this PR
1 1 0, put/get -600 (measured) 0 still -600
0 0 0, OK (measured) 0 OK
0,1 (dev 1) 1 1, OK (measured) 0 -600, a regression

Suggestion (as in #4980)

  • Export a device identity that survives CUDA_VISIBLE_DEVICES: the PCI bus ID (cudaDeviceGetPCIBusId on export, cudaDeviceGetByPCIBusId on open) or the GPU UUID.
  • Open on that device, and fail with a clear error when the real client cannot see it.
  • device_id can stay in the handle for logging and backward compatibility.

AI assistance: this analysis was prepared with Claude; I reviewed it against the code and the measurements in #4980.

Exporter ordinals are relative to CUDA_VISIBLE_DEVICES and cannot be
reused by the real client. Carry cudaDeviceGetPCIBusId in the handle and
resolve with cudaDeviceGetByPCIBusId on open; keep device_id for logging
and open on the caller device only for legacy empty pci_bus_id handles.

(Rebased onto main after kvcache-ai#5030; packaging commits dropped.)
@huangxiaobai-ydmy huangxiaobai-ydmy changed the title [Bugfix][Store] Open CUDA IPC handles on the caller's device [Bugfix][Store] Open CUDA IPC via exporter PCI bus ID Oct 28, 2026
@huangxiaobai-ydmy

Copy link
Copy Markdown
Contributor Author

Addressed @0z5a's analysis of the previous approach:

  • Opening on the RPC thread's "current device" was insufficient (and could regress the CVD=0,1 / buffer-on-1 case).
  • Export now records pci_bus_id via cudaDeviceGetPCIBusId; open resolves it with cudaDeviceGetByPCIBusId and switches to that local ordinal before cudaIpcOpenMemHandle.
  • device_id remains for logging/compat; legacy empty pci_bus_id keeps the TENT-style current-device open.
  • Clear INVALID_PARAMS when the exporter PCI bus ID is not visible to the real client.

(Rebased onto main after kvcache-ai#5030; packaging commits dropped.)
@0z5a

0z5a commented Oct 28, 2026

Copy link
Copy Markdown
Contributor

Thanks for switching to the PCI bus ID. I ran this PR at 8cd8c40c on the same kind of host as #4980: two L40S without P2P (nvidia-smi topo -p2p r shows NS). I used the PR's Linux CI wheel, and compared it against the CI wheel of main at fe4f297a.

Setup:

  • A real mooncake_master, a real client that sees both GPUs, and a dummy client.
  • 1 MiB put_from / get_into from a registered device buffer, compared byte by byte.
dummy CUDA_VISIBLE_DEVICES (device) buffer on main this PR
second GPU only (0) second GPU put/get -600 OK, byte-exact
both (1) second GPU OK OK
first GPU only (0) first GPU OK OK

The failing row from #4980 now passes, and the two control rows are unchanged. LGTM from the hardware side.

AI assistance: the test harness and this summary were prepared with Claude; I reviewed them and ran the tests on the hardware above.

@huangxiaobai-ydmy

Copy link
Copy Markdown
Contributor Author

Thanks @0z5a for re-running the put/get table on the two-GPU no-P2P host and confirming the #4980 case is fixed with no regressions on the control rows — much appreciated.

@huangxiaobai-ydmy
huangxiaobai-ydmy force-pushed the fix/issue-4980-cuda-ipc-device-ordinal branch from 8cd8c40 to cd9a342 Compare October 28, 2026 11:55
@huangxiaobai-ydmy

Copy link
Copy Markdown
Contributor Author

Friendly ping for review when you have a moment — CI is green on the latest head after dropping the unrelated wheel commits.

@ykwd (Store) — reporter @0z5a already re-ran the two-GPU no-P2P put/get table on this PR and LGTM'd from the hardware side.

@huangxiaobai-ydmy

Copy link
Copy Markdown
Contributor Author

Friendly Monday ping for review when you have a moment — CI is green, and @0z5a already confirmed the #4980 no-P2P host case is fixed with no regressions on the control rows.

@ykwd @stmatengss @YiXR @Aionw @XucSh — thanks!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci size/S 0-99 changed lines (additions + deletions) Store Tests

Projects

None yet

2 participants