[GH-ISSUE #4943] [RFC] Optional GPU read staging for fragmented Store destinations #1257

New issue

Closed

opened 2026-10-29 21:24:29 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented 2026-10-29 21:24:29 +00:00

Owner

Copy link

Originally created by @wjinxu on GitHub (Oct 26, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4943

Motivation

Loading cached KV through batch_get_into_multi_buffers can turn each CPU-resident Store object into many small transfers to GPU pages. The proposal keeps the existing key layout and API but receives each key into a contiguous GPU staging buffer, then scatters successful reads to the caller's pages with one CUDA kernel per buffer fill.

This is an opt-in performance experiment, not a claim that Store lookup or all production TTFT is transfer-bound.

Proposed scope

• One reusable buffer per Store client, disabled by default. MC_STORE_GPU_STAGING=1 enables it; MC_STORE_GPU_STAGING_BUFFER_BYTES sets the payload budget (default 256 MiB).
• Configuration is sampled during setup. GPU allocation and registration are lazy, after the destination GPU is known.
• Keep lookup, replica selection, key format, transfer protocol selection and per-key result handling in their existing paths.
• Require complete memory objects, exact byte coverage, registered nonoverlapping device destinations on one GPU, and hot cache disabled. CXL objects keep the existing path.
• A whole key must fit the buffer. Split a larger call at key boundaries, wait for the receive and scatter before reusing the single buffer, and scatter only successful keys.
• Unsupported requests or a busy buffer fall back before any staged transfer is submitted. Transfer/scatter failure returns an error; it must not silently replay partially written data through direct reads.
• The synchronous API returns destination ownership only after completion. Callers must quiesce reads before close or unregister, as with the existing API.
• Keep prefer_alloc_in_same_node unchanged in both paths. Its current read implementation groups submissions by source endpoint, so benchmark it explicitly rather than assuming that grouping and staging are equivalent.

Why a CUDA kernel

This adds an nvcc build requirement to CUDA-enabled Store builds. Other Mooncake modules already contain CUDA sources, but Store previously did not require this extra compilation step. Non-CUDA Store builds keep a no-op staging implementation.

We tried a .cpp implementation using cudaMemcpyBatchAsync. With 256 MiB of device data, measured over 20 post-warmup samples in rotating order:

Fragment size	Descriptor upload + custom scatter	C++ batched CUDA copies	Individual CUDA copies
8 KiB	0.163 ms	72.149 ms	79.619 ms
256 KiB	0.159 ms	2.546 ms	2.513 ms

These are complete host-observed submit-and-wait times on the test system, not universal CUDA performance figures. The batch-copy timing includes constructing pointer arrays. It was therefore rejected as the default implementation for this fragmented workload. The custom kernel keeps a small source file and launches once per buffer fill. No runtime compilation or embedded PTX is proposed.

Controlled Store experiment

H200 reader host, separate CPU Store owner, CUDA 13, RDMA; same-host means separate owner and reader processes using RDMA loopback. Each reader process uses a distinct GPU, Store instances, staging buffer and disjoint private keys. GPU memory registrations and initialization are warmed up before timed samples.

We ran 1, 2 and 4 reader processes, two key layouts, same-host and cross-host owners, three fresh-process rounds, eight timed samples per mode per round. Processes synchronize before each measured API call. All four modes run against the same data with rotating order. Every read checks returned sizes and exact byte equality; keys are removed after each group. There are 36 groups, 2688 timed reads, plus 336 warmup reads.

The table reports the median of the three per-round API latency medians, in milliseconds. It is not an end-to-end serving throughput measurement.

Topology	Keys × key size	Readers	Direct	Grouped direct	Staged	Grouped staged
cross	96 × 12 MiB	1	129.25	362.86	58.29	63.40
cross	96 × 12 MiB	2	132.54	356.52	58.61	64.44
cross	96 × 12 MiB	4	125.97	347.31	57.88	62.19
cross	1536 × 384 KiB	1	76.05	116.85	30.87	35.44
cross	1536 × 384 KiB	2	97.94	125.13	41.10	36.01
cross	1536 × 384 KiB	4	95.86	120.74	33.63	35.53
local	96 × 12 MiB	1	114.80	566.79	61.76	66.84
local	96 × 12 MiB	2	126.36	612.69	59.69	64.86
local	96 × 12 MiB	4	124.96	602.65	59.02	63.08
local	1536 × 384 KiB	1	68.78	170.93	30.00	35.24
local	1536 × 384 KiB	2	83.04	181.20	30.76	35.89
local	1536 × 384 KiB	4	91.07	188.76	43.49	35.65

Grouped direct means prefer_alloc_in_same_node=true. On this system it was slower than ordinary direct reads; this result does not by itself establish the reason or justify changing that existing default. Staging gains survive comparison against the faster of the two direct baselines.

Model validation

Earlier real Qwen3-Coder-30B-A3B-Instruct TP=2 tests used the same native Store build, 1K/8K/24K inputs, five repeated external-cache reads per length, greedy decoding with natural EOS, and a separate CPU owner. Resetting the local prefix cache intentionally forces external reads; this is a controlled hit-path experiment, not evidence of a production hit rate.

Input tokens	Direct / staged TTFT, round 1	Direct / staged TTFT, reverse-order round 2
1024	66.13 / 64.00 ms	65.64 / 63.55 ms
8192	105.84 / 90.30 ms	107.30 / 89.53 ms
24576	209.77 / 149.24 ms	199.30 / 157.18 ms

We additionally ran two vLLM TP=2 instances concurrently against the same cross-host owner for three fresh-instance rounds, swapping GPU assignments in round 2. Each mode completed 21 timed requests per round, including six mixed-length concurrent requests. All outputs matched across modes and ended naturally; the staged instance recorded native staging reads. Direct / staged median TTFT:

Input	Paired round 1	Paired round 2 (GPU assignment swapped)	Paired round 3
1024	63.60 / 62.33 ms	73.96 / 63.78 ms	77.96 / 66.19 ms
8192	96.45 / 86.87 ms	107.79 / 89.29 ms	108.36 / 90.90 ms
24576	172.44 / 145.08 ms	181.01 / 148.85 ms	204.75 / 155.76 ms

The serving instances ran while CPU-only native compilation was in progress, so these are controlled comparisons under shared-host activity, not isolated production capacity measurements. These model results use a saved native binary from before the type-layout-only cleanup; the final code was separately rebuilt and passed the focused regression suite.

Correctness and current draft limitations

After the final type-layout cleanup, CUDA and non-CUDA syntax checks and the CUDA native build passed. Four focused unit tests and both seven-check integration runs passed. The CUDA scatter test checks reordered, unaligned and odd-size fragments plus untouched guard bytes. Existing direct-read regression tests pass. Cross-host RDMA and same-host TCP integration checks cover exact-key buffer capacity, too-small-buffer fallback, changed addresses, irregular fragments, duplicate fanout, missing keys and concurrent calls. Traces confirm the native scatter and four buffer fills in a seven-key boundary case. The concurrent-call check is not a deterministic proof of busy-buffer fallback.

The prototype currently depends on a separate transfer-timeout draining change (f8a6a60 in the development branch). Timeout/exception ownership, failed initialization cleanup, and unregister failure behavior require explicit review before this is merge-ready. Successful-path benchmarks do not establish fault-path safety. Hot-cache admission is excluded because it can retain slices after the synchronous read.

Related work and questions

• #4507 / #4502 pack fragmented CPU source ranges on the owner for get_into_ranges. This proposal addresses fragmented GPU destinations for existing full-object multi-buffer reads; destination unpacking is outside that RFC's scope.
• #4117 and #4242 discuss adaptive scatter/gather at the Transfer Engine layer. This proposal initially keeps policy and buffer ownership in Store and uses the existing transport API.
• #4799 groups multi-buffer transfers by target and improves Ascend local copy. We explicitly measure the existing grouped read path; this proposal adds a contiguous GPU destination plus scatter rather than only grouping submissions.

Feedback requested: is Store the right ownership boundary, is the extra nvcc requirement acceptable for CUDA Store builds

Before submitting a new issue...

• Make sure you already searched for relevant issues and read the documentation

Originally created by @wjinxu on GitHub (Oct 26, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4943 ## Motivation Loading cached KV through `batch_get_into_multi_buffers` can turn each CPU-resident Store object into many small transfers to GPU pages. The proposal keeps the existing key layout and API but receives each key into a contiguous GPU staging buffer, then scatters successful reads to the caller's pages with one CUDA kernel per buffer fill. This is an opt-in performance experiment, not a claim that Store lookup or all production TTFT is transfer-bound. ## Proposed scope - One reusable buffer per Store client, disabled by default. `MC_STORE_GPU_STAGING=1` enables it; `MC_STORE_GPU_STAGING_BUFFER_BYTES` sets the payload budget (default 256 MiB). - Configuration is sampled during setup. GPU allocation and registration are lazy, after the destination GPU is known. - Keep lookup, replica selection, key format, transfer protocol selection and per-key result handling in their existing paths. - Require complete memory objects, exact byte coverage, registered nonoverlapping device destinations on one GPU, and hot cache disabled. CXL objects keep the existing path. - A whole key must fit the buffer. Split a larger call at key boundaries, wait for the receive and scatter before reusing the single buffer, and scatter only successful keys. - Unsupported requests or a busy buffer fall back before any staged transfer is submitted. Transfer/scatter failure returns an error; it must not silently replay partially written data through direct reads. - The synchronous API returns destination ownership only after completion. Callers must quiesce reads before close or unregister, as with the existing API. - Keep `prefer_alloc_in_same_node` unchanged in both paths. Its current read implementation groups submissions by source endpoint, so benchmark it explicitly rather than assuming that grouping and staging are equivalent. ## Why a CUDA kernel This adds an nvcc build requirement to CUDA-enabled Store builds. Other Mooncake modules already contain CUDA sources, but Store previously did not require this extra compilation step. Non-CUDA Store builds keep a no-op staging implementation. We tried a `.cpp` implementation using `cudaMemcpyBatchAsync`. With 256 MiB of device data, measured over 20 post-warmup samples in rotating order: | Fragment size | Descriptor upload + custom scatter | C++ batched CUDA copies | Individual CUDA copies | |---|---:|---:|---:| | 8 KiB | 0.163 ms | 72.149 ms | 79.619 ms | | 256 KiB | 0.159 ms | 2.546 ms | 2.513 ms | These are complete host-observed submit-and-wait times on the test system, not universal CUDA performance figures. The batch-copy timing includes constructing pointer arrays. It was therefore rejected as the default implementation for this fragmented workload. The custom kernel keeps a small source file and launches once per buffer fill. No runtime compilation or embedded PTX is proposed. ## Controlled Store experiment H200 reader host, separate CPU Store owner, CUDA 13, RDMA; same-host means separate owner and reader processes using RDMA loopback. Each reader process uses a distinct GPU, Store instances, staging buffer and disjoint private keys. GPU memory registrations and initialization are warmed up before timed samples. We ran 1, 2 and 4 reader processes, two key layouts, same-host and cross-host owners, three fresh-process rounds, eight timed samples per mode per round. Processes synchronize before each measured API call. All four modes run against the same data with rotating order. Every read checks returned sizes and exact byte equality; keys are removed after each group. There are 36 groups, 2688 timed reads, plus 336 warmup reads. The table reports the median of the three per-round API latency medians, in milliseconds. It is not an end-to-end serving throughput measurement. | Topology | Keys × key size | Readers | Direct | Grouped direct | Staged | Grouped staged | |---|---|---:|---:|---:|---:|---:| | cross | 96 × 12 MiB | 1 | 129.25 | 362.86 | 58.29 | 63.40 | | cross | 96 × 12 MiB | 2 | 132.54 | 356.52 | 58.61 | 64.44 | | cross | 96 × 12 MiB | 4 | 125.97 | 347.31 | 57.88 | 62.19 | | cross | 1536 × 384 KiB | 1 | 76.05 | 116.85 | 30.87 | 35.44 | | cross | 1536 × 384 KiB | 2 | 97.94 | 125.13 | 41.10 | 36.01 | | cross | 1536 × 384 KiB | 4 | 95.86 | 120.74 | 33.63 | 35.53 | | local | 96 × 12 MiB | 1 | 114.80 | 566.79 | 61.76 | 66.84 | | local | 96 × 12 MiB | 2 | 126.36 | 612.69 | 59.69 | 64.86 | | local | 96 × 12 MiB | 4 | 124.96 | 602.65 | 59.02 | 63.08 | | local | 1536 × 384 KiB | 1 | 68.78 | 170.93 | 30.00 | 35.24 | | local | 1536 × 384 KiB | 2 | 83.04 | 181.20 | 30.76 | 35.89 | | local | 1536 × 384 KiB | 4 | 91.07 | 188.76 | 43.49 | 35.65 | Grouped direct means `prefer_alloc_in_same_node=true`. On this system it was slower than ordinary direct reads; this result does not by itself establish the reason or justify changing that existing default. Staging gains survive comparison against the faster of the two direct baselines. ## Model validation Earlier real Qwen3-Coder-30B-A3B-Instruct TP=2 tests used the same native Store build, 1K/8K/24K inputs, five repeated external-cache reads per length, greedy decoding with natural EOS, and a separate CPU owner. Resetting the local prefix cache intentionally forces external reads; this is a controlled hit-path experiment, not evidence of a production hit rate. | Input tokens | Direct / staged TTFT, round 1 | Direct / staged TTFT, reverse-order round 2 | |---:|---:|---:| | 1024 | 66.13 / 64.00 ms | 65.64 / 63.55 ms | | 8192 | 105.84 / 90.30 ms | 107.30 / 89.53 ms | | 24576 | 209.77 / 149.24 ms | 199.30 / 157.18 ms | We additionally ran two vLLM TP=2 instances concurrently against the same cross-host owner for three fresh-instance rounds, swapping GPU assignments in round 2. Each mode completed 21 timed requests per round, including six mixed-length concurrent requests. All outputs matched across modes and ended naturally; the staged instance recorded native staging reads. Direct / staged median TTFT: | Input | Paired round 1 | Paired round 2 (GPU assignment swapped) | Paired round 3 | |---:|---:|---:|---:| | 1024 | 63.60 / 62.33 ms | 73.96 / 63.78 ms | 77.96 / 66.19 ms | | 8192 | 96.45 / 86.87 ms | 107.79 / 89.29 ms | 108.36 / 90.90 ms | | 24576 | 172.44 / 145.08 ms | 181.01 / 148.85 ms | 204.75 / 155.76 ms | The serving instances ran while CPU-only native compilation was in progress, so these are controlled comparisons under shared-host activity, not isolated production capacity measurements. These model results use a saved native binary from before the type-layout-only cleanup; the final code was separately rebuilt and passed the focused regression suite. ## Correctness and current draft limitations After the final type-layout cleanup, CUDA and non-CUDA syntax checks and the CUDA native build passed. Four focused unit tests and both seven-check integration runs passed. The CUDA scatter test checks reordered, unaligned and odd-size fragments plus untouched guard bytes. Existing direct-read regression tests pass. Cross-host RDMA and same-host TCP integration checks cover exact-key buffer capacity, too-small-buffer fallback, changed addresses, irregular fragments, duplicate fanout, missing keys and concurrent calls. Traces confirm the native scatter and four buffer fills in a seven-key boundary case. The concurrent-call check is not a deterministic proof of busy-buffer fallback. The prototype currently depends on a separate transfer-timeout draining change (`f8a6a60` in the development branch). Timeout/exception ownership, failed initialization cleanup, and unregister failure behavior require explicit review before this is merge-ready. Successful-path benchmarks do not establish fault-path safety. Hot-cache admission is excluded because it can retain slices after the synchronous read. ## Related work and questions - #4507 / #4502 pack fragmented CPU source ranges on the owner for `get_into_ranges`. This proposal addresses fragmented GPU destinations for existing full-object multi-buffer reads; destination unpacking is outside that RFC's scope. - #4117 and #4242 discuss adaptive scatter/gather at the Transfer Engine layer. This proposal initially keeps policy and buffer ownership in Store and uses the existing transport API. - #4799 groups multi-buffer transfers by target and improves Ascend local copy. We explicitly measure the existing grouped read path; this proposal adds a contiguous GPU destination plus scatter rather than only grouping submissions. Feedback requested: is Store the right ownership boundary, is the extra nvcc requirement acceptable for CUDA Store builds ### Before submitting a new issue... - [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

bsakaguchi 2026-10-29 21:24:29 +00:00

• closed this issue
• added the
  RFC
  label

bsakaguchi commented 2026-10-29 21:24:30 +00:00

Author

Owner

Copy link

@wjinxu commented on GitHub (Oct 26, 2026):

I've put together a small PR https://github.com/kvcache-ai/Mooncake/pull/4944 and run some basic tests. Feedback is welcome.

<!-- gh-comment-id:6413107658 --> @wjinxu commented on GitHub (Oct 26, 2026): I've put together a small PR https://github.com/kvcache-ai/Mooncake/pull/4944 and run some basic tests. Feedback is welcome.

bsakaguchi commented 2026-10-29 21:24:30 +00:00

Author

Owner

Copy link

@github-actions[bot] commented on GitHub (Oct 26, 2026):

Thanks for opening this issue, @wjinxu!

Field	Value
Issue	#4943
GitHub user ID	105263726
Reporter	@wjinxu

A maintainer will triage this when possible. To help us respond faster, please include:

• Mooncake version or commit SHA
• Environment (OS, CUDA/driver, RDMA stack if relevant)
• Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

<!-- gh-comment-id:6413919456 --> @github-actions[bot] commented on GitHub (Oct 26, 2026): Thanks for opening this issue, @wjinxu! | Field | Value | |-------|-------| | **Issue** | #4943 | | **GitHub user ID** | `105263726` | | **Reporter** | @wjinxu | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi referenced this issue 2026-10-29 21:32:52 +00:00

[PR #1257] [CLOSED] [Feature] Modify nvlink_allocator to enable nvlink intraNode transport in SGlang #2112

Sign in to join this conversation.

No Branch/Tag specified

Branches Tags

main

te-coro

staryxchen-tent-capability-research

feature/musa-ep-jit

zxpdemonio/fix-tent-bootstrap-admission

add-store-rs-tone-tests

scatter-rdma-completion-pr2

P2P-Mooncake-Store

release/v0.3.14

feat/tent-rdma-adaptive-adapter

fix/ep-local-counter-gpu-scope

codex/kvcs-oplog-eviction

give-me-a-star

cruz/mooncake-kvcs-lowlevel-pr1-20261008

ylt-handshake-limits

ci/rocm-e2e-tent

sima/docs-pypi-downloads

fix-ci/hicache-mooncake-metadata-loopback

sima/cuda-driver-loading

sima/ci-performance-gate

fix/tent-rdma-unknown-topology-rail-pairing

ci/tent-e2e-dual-job

mc-patch-1

codex/p2p-merge-prerequisites

sima/store-hca-numa-placement

codex/p2p-merge-cli

codex/p2p-merge-ipv6

copilot/fix-ctests-job-failure

czh/pg-device-collectives-pr6

dependabot/cargo/mooncake-store-rs/pyo3-0.29.0

czh/pg-device-collectives-pr5

codex/segment-pool-master-cutover

sima/nightly-ci-recovery

ci/tent-params

stmatengss-patch-1

dependabot/cargo/mooncake-store-rs/rand-0.8.8

dependabot/cargo/mooncake-store-rs/opentelemetry_sdk-0.32.1

dependabot/cargo/mooncake-store-rs/xxhash-rust-0.8.16

codex/migrate-store-rest-service

fix/tone-wheel-arch-selection

czh/pg-device-collectives-pr4

czh/pg-device-collectives-pr3

ketong/fix_moe_tc

copilot/fix-test-wheel-ubuntu-job

fix/ci-wheel-torch-pg-pin

codex/migrate-structured-object-storage

codex/migrate-allocator-python

codex/migrate-vllm-python

codex/migrate-shared-segment

ci/fix_sgl_eval_install

codex/segment-pool-nof-placement-cleanup

feat/pip-native-release

copilot/bugfix-ha-standby-promotion

copilot/mooncake-store-transport-settings

copilot/fix-tensor-support-issues

copilot/optimize-rl-direction

copilot/mooncake-storeagentic-workflow

copilot/fix-client-reconnect-timeout

copilot/mooncake-error-debugging

copilot/mooncake-ipv6-support

feat/nccl-host-python-api

dependabot/go_modules/mooncake-p2p-store/src/p2pstore/google.golang.org/grpc-1.82.1

copilot/migrate-and-move-operations

copilot/migrate-and-move-operations-again

refactor/centralize-env-vars

idhanani/mooncake-retain-groups

copilot/codeownerreviewer

copilot/give-feedback-on-issue-2123

copilot/create-issue-rfc

copilot/break-cuda-monopoly-logic

copilot/compare-cudamemcpy-cu12-cu13

copilot/report-ci-error-issue

copilot/research-commit-history-analysis

copilot/explain-mooncake-store-service

copilot/analyze-project-functionality

copilot/sub-pr-1680

copilot/review-pr-1680

copilot/sub-pr-1597

copilot/fix-github-actions-job-failure

copilot/move-test-sglang-integration-job

v0.3.14

v0.3.14-pre-4

v0.3.14-pre-3

v0.3.14-pre-2

v0.3.14-pre

v0.3.14-rc1

v0.3.13.post1

v0.3.13-rc2

v0.3.13-rc1

v0.3.13

v0.3.12.post1

v0.3.12

v0.3.11-rc1

v0.3.11.post1

v0.3.11

v0.3.10.post2

v0.3.10.post1

v0.3.10

v0.3.9

v0.3.8.post1

v0.3.8

v0.3.7.post2

v0.3.7.post1

v0.3.7

v0.3.6.post1

v0.3.6

v0.3.5

v0.3.4.post2

v0.3.4.post1

v0.3.4

v0.3.3.post2

v0.3.3.post1

v0.3.3

v0.3.2.post1

v0.3.2

v0.3.1

v0.3.0

v0.3.0-beta

v0.3.0-alpha

v0.2.0

Labels

Clear labels   

ascend

  

auto-closed

  

bug

  

bug

  

CI/Build

  

Common

  

documentation

  

enhancement

  

feature request

  

good first issue

  

help wanted

  

high priority

  

MetaX

  

nightly-failure

  

pull-request

Mirrored from GitHub Pull Request

  

PyTorch Backend

  

question

  

RFC

  

Roadmap

  

stale

  

Store

  

Tests

  

Transfer Engine

No labels

ascend

auto-closed

bug

bug

CI/Build

Common

documentation

enhancement

feature request

good first issue

help wanted

high priority

MetaX

nightly-failure

pull-request

PyTorch Backend

question

RFC

Roadmap

stale

Store

Tests

Transfer Engine

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

1 participant

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/Mooncake#1257

No description provided.