[store] Avoid staging copy for same node tensor put

[zxpdemonio]

Description

Motivation

The Python tensor write path used to copy tensor payloads through a client-side staging buffer before the data reached the Mooncake Store segment.

For local same-node tensor writes, this staging copy is unnecessary when the target Store segment is local to the writer process. It is especially expensive for GPU tensor -> Store CPU segment writes, where the old path first copied the GPU tensor into a client local/staging buffer and then copied that buffer into the Store segment.

This PR removes the extra staging copy from the local tensor write path.

This PR is split from the broader local tensor/offload optimization work. The scope here is intentionally limited to same-node tensor write staging removal.

Scenario

Before this PR:

tensor payload
  -> client local/staging buffer
  -> Store segment

After this PR, when the write can safely use the local same-node path:

tensor payload
  -> Store segment

Changes

• Add a direct multi-buffer write path for tensor payloads.
• Let local same-node tensor writes pass tensor payload buffers directly into the Store write path.
• Extend the transfer task local-copy path to consume caller-provided source buffers for local writes.
• Route put/upsert tensor writes through the shared multi-buffer writer to avoid duplicate tensor staging logic.
• Route single put_tensor and upsert_tensor through the same no-staging multi-buffer path as batch tensor writes.

Performance

Local write staging removal

Environment:

Mode: local same-node tensor put
Protocol: rdma
RDMA device: erdma_0
Sizes: 64 MB, 256 MB, 512 MB
Warmup / iterations: 2 / 7
Baseline: f224dcaf
This PR: c378659e
Benchmark harness: external local script, not included in this PR

The benchmark log confirms MC_STORE_MEMCPY was auto-disabled under RDMA, so the baseline uses the old staging/TransferEngine path. The after case uses the same public tensor APIs and removes the local staging copy when the destination is the local Store segment.

Local CPU tensor -> Store CPU segment put:

Size	Before	After	Improvement
64 MB	5.106 GB/s / 13.142 ms	13.924 GB/s / 4.820 ms	+173%
256 MB	5.160 GB/s / 52.027 ms	13.341 GB/s / 20.121 ms	+159%
512 MB	5.459 GB/s / 98.348 ms	13.340 GB/s / 40.245 ms	+144%

Local GPU tensor -> Store CPU segment put:

Size	Before	After	Improvement
64 MB	2.515 GB/s / 26.689 ms	7.595 GB/s / 8.836 ms	+202%
256 MB	2.586 GB/s / 103.820 ms	7.752 GB/s / 34.626 ms	+200%
512 MB	2.752 GB/s / 195.109 ms	7.776 GB/s / 69.038 ms	+183%

The main performance gain comes from local same-node tensor writes where the extra staging copy is removed.

Pinned Store segment memory local CUDA write

Environment:

Mode: local same-node CUDA tensor write into pinned Store CPU segment
Protocol: tcp
MC_STORE_MEMCPY: 1
Pinned Store segment memory: enabled
MC_STORE_PIN_MEMORY_MAX_BYTES: 4294967296
Store segment size: 4096 MB
Client local buffer size: 1024 MB
Sizes: 64 MB, 256 MB, 512 MB
Warmup / iterations: 3 / 8
Baseline: 6c721eea
This PR: 00230d8f
Benchmark harness: external local one-shot script, not included in this PR

This benchmark uses pinned Store segment memory and local memcpy. It validates that single tensor writes now reach the same no-staging multi-buffer path as batch tensor writes.

Local CUDA tensor -> pinned Store CPU segment put_tensor:

Size	Base GB/s	PR GB/s	Speedup	Base median	PR median
64 MB	3.363	24.923	7.41x	19.955 ms	2.693 ms
256 MB	3.333	25.824	7.75x	80.532 ms	10.395 ms
512 MB	3.345	25.921	7.75x	160.483 ms	20.712 ms

Local CUDA tensor -> pinned Store CPU segment upsert_tensor:

Size	Base GB/s	PR GB/s	Speedup	Base median	PR median
64 MB	3.379	24.702	7.31x	19.858 ms	2.717 ms
256 MB	3.334	25.712	7.71x	80.520 ms	10.440 ms
512 MB	3.346	25.923	7.75x	160.430 ms	20.710 ms

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• 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 commands:

./scripts/code_format.sh --check -b origin/main

cmake --build build-tensor-staging-nocuda-notest -j8 --target store

cmake -S . -B build-tensor-staging-cuda \
  -DUSE_CUDA=ON \
  -DWITH_STORE_RUST=OFF \
  -DWITH_EP=OFF \
  -DBUILD_UNIT_TESTS=OFF \
  -DBUILD_TESTS=OFF \
  -DCMAKE_BUILD_TYPE=Release

cmake --build build-tensor-staging-cuda -j8 --target store mooncake_master

Manual validation was run with local benchmark/smoke harnesses that are not included in this PR:

• local CPU tensor -> Store CPU segment put
• local GPU tensor -> Store CPU segment put
• CPU tensor write smoke for put_tensor, upsert_tensor, batch_put_tensor, and batch_upsert_tensor
• CUDA duplicate put_tensor smoke to confirm the existing successful return behavior for an existing key
• CUDA tensor local write benchmark with pinned Store segment memory and local memcpy enabled

Test results:

• Format check passes
• Build passes without CUDA enabled
• Build passes with CUDA enabled
• Manual CPU tensor write smoke passes
• Manual CUDA duplicate put_tensor smoke passes
• Manual performance validation done

Manual test results:

• Local CPU tensor put improved from 5.1-5.5 GB/s to 13.3-13.9 GB/s.
• Local GPU tensor put improved from 2.5-2.8 GB/s to 7.6-7.8 GB/s.
• With pinned Store segment memory, local CUDA put_tensor improved from 3.33-3.36 GB/s to 24.92-25.92 GB/s.
• With pinned Store segment memory, local CUDA upsert_tensor improved from 3.33-3.38 GB/s to 24.70-25.92 GB/s.

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 (if applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (Codex was used to help with implementation, benchmarking, and review)

[zxpdemonio]

Verdict

Useful same-process tensor write optimization with solid measured gains, but one behavioral bug should be fixed before merge, and the auto-enable of prefer_alloc_in_same_node needs clearer fallback semantics.

Blockers

1. MakeDefaultLocalTensorWriteConfig cannot honor an explicit prefer_alloc_in_same_node=False.
   For an otherwise-default ReplicateConfig, the helper always flips the flag to true, even when the caller set it to false. There is currently no clean opt-out short of setting a dummy preferred segment / replica policy. That makes the new no-staging path mandatory for common tensor writes and contradicts the PR claim that unsupported paths keep the staging fallback.

2. Default tensor put/upsert semantics change without a remote/staging fallback.
   Default put_tensor / upsert_tensor / batch variants now force same-node allocation + direct buffers. If the local segment cannot allocate, the write fails hard (NO_AVAILABLE_HANDLE / same-process submit failure) instead of falling back to the old staging + TransferEngine path. Please either:

   • keep auto-prefer but add an automatic staging fallback when local alloc/submit fails, or
   • make auto-prefer opt-in (or respect explicit False) and document the capacity/behavior change prominently.

Non-blocking

3. append_tensor_write_buffers uses sizeof(TensorMetadata) while the staging path uses metadata.header.data_offset. Today builders set them equal, but the no-staging path should use data_offset for layout consistency and future-proofing.

4. batch_upsert_from_multi_buffers has no DummyClient implementation (base default returns INVALID_PARAMS). Tensor APIs already reject dummy clients, so this is mostly API asymmetry vs batch_put_from_multi_buffers, but worth closing if upsert multi-buffer is now a public surface.

5. Test coverage is thin relative to the behavior change.

   • Added test_batch_pub_tensor_same_node only covers explicit prefer_alloc_in_same_node=True.
   • Missing coverage for: default config auto-prefer, explicit False opt-out, upsert no-staging, and CUDA tensor correctness under the new default path (beyond the small get equality assert).

6. require_same_process intentionally bypasses MC_STORE_MEMCPY=0. That matches the benchmark story, but please call it out in docs so operators who disabled memcpy are not surprised.

What looks good

• Constraining no-staging to the current client segment via normalize_same_process_tensor_write_config.
• Forcing memcpy through submit_batch(..., require_same_process=true) with local-endpoint checks.
• Enabling upsert on the prefer-same-node path and fixing FinalizeBatchUpsert to use allocated-replica transfer success instead of pending_transfers (which are only tracked on merged ops).
• Sharing put/upsert through batch_*_from_multi_buffers / BatchWriteFromMultiBuffers.
• Accelerator-aware memcpy worker already handles GPU to CPU for the direct path.

CI

Checks were still pending/skipped at review time; please confirm format + the listed tensor tests remain green on the latest head.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 17.87440% with 170 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...oncake-integration/store/store_py_parallel_write.h	0.00%	74 Missing ⚠️
mooncake-integration/store/store_py.cpp	0.00%	40 Missing ⚠️
mooncake-store/src/real_client.cpp	48.05%	40 Missing ⚠️
mooncake-store/src/dummy_client.cpp	0.00%	15 Missing ⚠️
mooncake-store/src/real_client_main.cpp	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[yokinoshitayoki]

@zxpdemonio It seems that this PR has conflicts that must be resolved.

[zxpdemonio]

@zxpdemonio It seems that this PR has conflicts that must be resolved.

@yokinoshitayoki The rebase conflict has been resolved.

The conflict was caused by the CUDA IPC support that recently landed upstream. After rebasing on top of it, we wired the dummy tensor write path to use the CUDA IPC handoff as well.

With this change, same-node dummy-client GPU tensor writes can avoid the dummy/client local staging buffer: the dummy client exports the CUDA tensor via CUDA IPC, and the real client maps that payload and writes it directly into the local Store segment. I also added CUDA coverage for dummy upsert_tensor and verified the path with client metrics showing batch_upsert_from_cuda_ipc.

[yokinoshitayoki]

Should we take a look at the failed CI tests?

[zxpdemonio]

Thanks for the reminder. I checked the failed jobs. They were caused by EPEL mirror metadata download errors (404s and timeouts), not by the PR code. I reran the failed jobs, and the full workflow is now green: https://github.com/kvcache-ai/Mooncake/actions/runs/31068240815. No source changes were needed.

[zxpdemonio]

Update summary for the latest head (a85ec52):

• Kept the default tensor API allocation semantics unchanged. The no-staging write path is opt-in through prefer_alloc_in_same_node; allocation now reuses the existing host-aware local-first allocator and retains its ordered remote fallback.
• Removed the client-side segment override. After allocation, direct copy is used only when the selected Store segment is eligible for the same-process local-memcpy path.
• Routed single and batch put_tensor / upsert_tensor writes through the shared multi-buffer implementation, avoiding duplicated staging logic.
• Completed DummyClient CUDA IPC support for tensor upsert, so dummy and real clients expose the same put/upsert behavior.
• Added coverage for host-aware placement and CUDA tensor write paths, and updated the API documentation to match the final behavior.

The current Linux workflow, including format, builds, wheel tests, and CI Gate, is green: https://github.com/kvcache-ai/Mooncake/actions/runs/31068240815

[yokinoshitayoki]

@zxpdemonio Seems this branch has conflicts that must be resolved, could you take a look?

[zxpdemonio]

Supplemental DummyClient tensor-path validation

The existing performance tables in the PR description use a real client. This is an additional end-to-end measurement for the tensor + DummyClient path.

Environment:

• PR head: a85ec526
• DummyClient process -> standalone real client on the same node
• CUDA IPC source handoff
• pinned Store segment: 4 GiB
• MC_STORE_MEMCPY=1
• protocol: TCP
• client local buffer: 64 MiB
• tensor sizes: 64/256/512 MiB
• warmup/iterations: 3/8
• ready CUDA tensors are explicitly synchronized before timing; asynchronous current-stream readiness is tracked separately by [Store] Fix current-stream readiness for DummyClient CUDA IPC tensor writes #3303
• external one-shot harness, not included in this PR

API	Size	Throughput	Median latency	Readback
put_tensor(cuda)	64 MiB	22.730 GB/s	2.952 ms	passed
put_tensor(cuda)	256 MiB	25.267 GB/s	10.624 ms	passed
put_tensor(cuda)	512 MiB	25.735 GB/s	20.862 ms	passed
upsert_tensor(cuda)	64 MiB	22.757 GB/s	2.949 ms	passed
upsert_tensor(cuda)	256 MiB	25.292 GB/s	10.614 ms	passed
upsert_tensor(cuda)	512 MiB	25.741 GB/s	20.857 ms	passed

Each write was followed outside the timed region by a DummyClient CUDA readback and payload comparison. The real-client log also confirmed that the 4 GiB Store segment was successfully registered as pinned memory. The 256 and 512 MiB tensors both exceed the 64 MiB client local buffer and still complete, confirming that the payload does not depend on a full-tensor client staging buffer.

[zxpdemonio]

Rebased this PR onto the latest main, including the changes from #3197, and reran the same-node dummy-client CUDA tensor write benchmark.

Configuration: MC_STORE_MEMCPY=1, 4 GiB pinned Store segment, 64 MiB client local buffer, CUDA device source. Every write was read back and verified.

Size	put_tensor	upsert_tensor
64 MiB	22.63 GB/s	22.83 GB/s
256 MiB	25.30 GB/s	25.32 GB/s
512 MiB	25.79 GB/s	25.80 GB/s

The 256 MiB and 512 MiB tensors are larger than the 64 MiB client local buffer and still sustain about 25 GB/s, confirming that the same-node CUDA IPC/local-memcpy path remains no-staging after the rebase.