[Store] Fix current-stream readiness for DummyClient CUDA IPC tensor writes

[mo-ke-ke]

Description

The DummyClient CUDA IPC tensor-write fast path can expose a tensor allocation to the receiver before work queued on the caller's PyTorch current stream has completed. For example, after current.wait_stream(producer), the dependency exists on current, but the existing path exports the memory handle and starts the RPC without observing that dependency. The receiver can therefore consume stale or partially written data.

This change:

• captures each CUDA tensor's device and PyTorch current stream while holding the GIL;
• exports all CUDA IPC memory handles first, preserving the existing whole-batch staging fallback for non-exportable and zero-sized payloads;
• validates the captured device against the exported allocation;
• deduplicates (device_id, stream_handle) pairs and synchronizes only streams associated with requests that will be sent;
• completes synchronization before invoking the existing RPC; and
• returns an error without issuing the RPC if stream validation or synchronization fails.

Callers remain responsible for joining work from other producer streams into the current stream, for example with current.wait_stream(producer).

This is deliberately narrow: CPU tensors and the staging fallback are unchanged, and the patch adds no device-wide synchronization, wire-format field, RPC endpoint, serialized ABI change, or capability negotiation.

Related work: #3234 introduced the DummyClient CUDA IPC tensor path fixed here. #1946 targets registered Store buffers and the local-replica GPU/CPU transfer path, while #3159 targets CUDA IPC upsert. Neither covers PyTorch current-stream readiness for the existing put_tensor, batch, and TP write path. This PR does not modify either path or any wire/RPC schema.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Test environment: NVIDIA B300, CUDA 12.8, PyTorch 2.11.0+cu128.

Test commands:

./scripts/code_format.sh --check --changed-lines --base origin/main
python -m pytest -q \
  mooncake-wheel/tests/test_dummy_client.py::TestDistributedObjectStoreSingleStore::test_dummy_cuda_ipc_put_tensor_waits_for_current_stream \
  mooncake-wheel/tests/test_dummy_client.py::TestDistributedObjectStoreSingleStore::test_dummy_cuda_ipc_batch_put_tensor_waits_for_current_stream

Correctness results:

Asynchronous case	Pre-fix main (a6b4db4)	This PR
Single tensor	0/20 correct	20/20 correct
Batch tensor	0/20 correct	20/20 correct

A separate synthetic matrix using a non-default current stream covered async/control and single/batch cases at 20 trials each: 80/80 correct with zero payload mismatches.

Fresh CUDA and no-CUDA configurations built successfully, and the Python extension imported successfully from both builds.

Ready-tensor performance:

• At 4 KiB, differences were noise-dominated; the worst observed P95 increase was 3.1%.
• At 16 MiB, no performance regression was observed.

Pre-commit:

Every changed-file hook except ruff-format passed. ruff check reports zero findings on both origin/main and this branch. ruff format --check has the same whole-file failure on both revisions, so that hook was skipped to avoid an unrelated full-file reformat.

Test results:

• Unit tests pass
• Integration tests pass
• Manual testing done

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit run --all-files and all hooks pass
• I have updated the documentation (not applicable: no public API or configuration change)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue (not applicable: the change is below 500 LOC)

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specify below)

OpenAI Codex assisted with source inspection, CUDA stream-ordering analysis, test design, implementation, validation, and review. The submitter remains responsible for understanding and defending the change.

[mo-ke-ke]

@stmatengss @zxpdemonio @ykwd, when you have a chance, could one of you please take a look at this narrow follow-up to #3234? It fixes a reproduced PyTorch current-stream readiness race in the DummyClient CUDA IPC tensor-write path without changing the wire/RPC schema; the pre-fix case reproduced at 0/20 versus 20/20 here, and CI is green.

[zxpdemonio]

I reviewed this change with emphasis on API compatibility, failure atomicity, CUDA/non-CUDA portability, batch fallback, and data exposure. I found no blocking issues.

Why our earlier testing missed this race: our performance scripts explicitly called torch.cuda.synchronize() after tensor initialization and again before every put_tensor. That made every source tensor host-ready before CUDA IPC export, so those runs measured ready-tensor throughput but did not exercise PyTorch current-stream ordering. The earlier dummy tensor regression also used CPU tensors. This PR closes that gap by asserting that the producer event is still pending at the call boundary after current.wait_stream(producer), for both single and batch writes.

I also validated commit 530e50e6 on an NVIDIA A10 / CUDA 12.8 test host:

• CUDA store, mooncake_client, and mooncake_master targets built successfully.
• The two new stream-readiness tests passed: 2/2.
• Existing dummy tensor operations and mixed put/put_tensor concurrency tests passed: 2/2.
• Changed-line formatting and git diff --check passed.
• The upstream CI matrix is green, including non-CUDA and ROCm builds.

The public Python signatures and RPC/wire schema remain unchanged. CPU, zero-sized, mixed CPU/CUDA, and non-exportable batches retain the existing whole-batch staging fallback. Stream validation/synchronization completes before RPC submission, so a failure cannot publish a partial Store write; the new error logs also do not expose keys, pointers, or raw exception text.

Residual, non-blocking risks: the committed regression relies on private torch.cuda._sleep and will skip where it is unavailable; forced CUDA stream-failure injection is not covered.

[zxpdemonio]

Follow-up after reproducing the test command exactly: the production fix still looks correct, but I found two issues in the new regression test that should be closed before relying on it as a gate.

Validation on NVIDIA A10 / CUDA 12.8 at 530e50e6:

• CUDA Store/client/master build: passed
• changed-line format and git diff --check: passed
• exact new-test invocation: 2 skipped, not passed
• same two tests with PyTorch preloaded before Mooncake: 2 passed

Two test changes are needed:

1. At test_dummy_client.py:10, catch only a top-level ModuleNotFoundError for torch and re-raise other import failures. Otherwise a broken CUDA/PyTorch environment is silently reported as a skip.
2. At test_dummy_client.py:232, assert ready_at_return. The contract under test is that put_tensor/batch_put_tensor do not return before the captured current stream is ready. Payload verification alone can pass accidentally if RPC/server latency outlasts the synthetic sleep, even when client-side synchronization is removed.

I found no additional blocker in the production implementation: current-stream capture happens under the GIL, device/stream validation and synchronization happen before RPC submission, partial writes are not published on synchronization failure, and the public API/wire schema are unchanged.

[mo-ke-ke]

@zxpdemonio Fixed both test-gate issues in e47e341: only a top-level missing torch is treated as unavailable, while dependency and binary import failures now propagate; the regression also fails unless the captured current stream is ready when put_tensor or batch_put_tensor returns.

I reran the exact two-test invocation on B300 without preloading PyTorch: 2 passed, 4 subtests passed, 0 skipped. The updated head's CI gate is currently tripped by an unrelated repo-wide typos finding on current main at mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp:260 (mis-set).

[mo-ke-ke]

@XucSh @YiXR @stmatengss @ykwd The updated head is now green on the current main (24 successful, 1 skipped, 0 failed), and the two test-gate follow-ups from the existing review are addressed in e47e341. When one of you has availability, could you provide the required Store CODEOWNERS review? It is the only remaining merge gate.

[mo-ke-ke]

@ykwd Gentle follow-up on this Store CUDA IPC readiness fix. The two test-gate issues raised during review were addressed in e47e341, and the current head remains conflict-free with main. When convenient, could you please review the Store-side changes needed to clear the CODEOWNER gate? Thanks!