Skip to content

[Store] Skip hot-cache writeback for device-resident read destinations - #4764

Merged
ykwd merged 1 commit into
kvcache-ai:mainfrom
he-yufeng:fix/hot-cache-device-dest-skip
Oct 20, 2026
Merged

ykwd merged 1 commit into
kvcache-ai:mainfrom
he-yufeng:fix/hot-cache-device-dest-skip

Conversation

@he-yufeng

@he-yufeng he-yufeng commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

What

Fix the SIGSEGV reported in #4761: with the local hot cache enabled
(MC_STORE_LOCAL_HOT_CACHE_SIZE), batch_get_into /
batch_get_into_multi_buffers crash the reading process when the destination
buffer is CUDA memory.

Root cause

After a successful remote read, Client::ProcessSlicesAsync submits every
destination slice to LocalHotCacheHandler::SubmitPutTask, which copies the
data into the cache block with a host std::memcpy(block->addr, slice.ptr, slice.size) (local_hot_cache.cpp). When the destination is a device pointer,
that host memcpy dereferences a device address and the process dies
(__memcpy_avx512_unaligned in the reporter's stack).

Single get_into into CUDA and batch reads with the hot cache disabled never
reach this path, which matches the reporter's matrix of what works and what
crashes.

Fix

In ProcessSlicesAsync, skip the writeback for slices whose destination is
device memory, using the existing accelerator registry probe
(RuntimeAccelerators().FindDeviceForPointer) — the same discriminator the
checksum path (ComputeObjectChecksumForSlices) and the DFS staging paths
already use. Device-resident read results are simply not promoted into the
host hot cache; the read itself is unaffected. On builds without any
accelerator backend the registry is empty and behavior is unchanged.

Caching device-destined reads by staging D2H into the cache block is a
possible follow-up, but silently dropping the writeback is the safe fix: it
restores the pre-hot-cache behavior for exactly the configurations that crash
today.

Test

New regression test
LocalHotCacheTest.BatchGetSkipsHotCacheWritebackForDeviceDestination:

  • registers a fake AcceleratorDevice (same pattern as
    runtime_accelerator_test.cpp) that classifies one static buffer as device
    memory, so the guard is exercised on CPU-only runners;
  • runs a real two-client TCP transfer (writer owns the segment, reader owns
    none) so the writeback path in ProcessSlicesAsync actually executes;
  • asserts the device-destined key is not in the hot cache while the
    host-destined control key from the same batch is (waits for the async
    publish, so the skip cannot pass vacuously).

Verified on Linux (ubuntu:22.04 container, aarch64): the new test fails with
the guard removed (the fake-device key gets cached) and passes with it.

Fixes #4761

// Identify TE transfer slices (non-local) and submit async put tasks
for (size_t i = 0; i < slices.size(); ++i) {
if (runtime_accelerator.FindDeviceForPointer(slices[i].ptr) !=
nullptr) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: maybe a logging here to show how many slice we skipped.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Counted and logged per batch now (device_dest_slices -> one INFO line with key and skip count). 40/40 in the Linux container.

@he-yufeng

Copy link
Copy Markdown
Collaborator Author

CI note: the one CTest failure is dfs_sync_client_test, and it fails identically on main (run 39242136165, same "1 tests failed out of 258", same test #180), so it is pre-existing and not caused by this PR. The new client_local_hot_cache_test including the added BatchGetSkipsHotCacheWritebackForDeviceDestination case passes (17.1s in this run). Re-triggered once already and the failure reproduced on the same test, which also matches the main-branch behavior.

@he-yufeng
he-yufeng force-pushed the fix/hot-cache-device-dest-skip branch from c46d268 to c0c5ccb Compare October 19, 2026 18:15
@he-yufeng

Copy link
Copy Markdown
Collaborator Author

Done — skipped slices are now counted and logged once per batch (Skipped hot-cache writeback for N device-resident slice(s), key=...). Rebuilt and re-ran client_local_hot_cache_test in the Linux container: 40/40 pass.

@he-yufeng

Copy link
Copy Markdown
Collaborator Author

Status note: the CTest failure blocking the CI Gate here is the DfsOnlyBatchFailureOnlyRevokesFailedKey main breakage from #4637; the fix is now #4781 (I closed my duplicate #4785 in its favor after verifying both in a container). Once either lands, a rebase here should turn the gate green.

Fixes kvcache-ai#4761.

With the local hot cache enabled, batch_get_into into CUDA memory crashed
the reader with SIGSEGV: ProcessSlicesAsync handed every destination slice
to LocalHotCacheHandler::SubmitPutTask, which copies into the cache block
with a host std::memcpy. For a device pointer that dereference faults.

Probe each slice with the accelerator registry (the same
FindDeviceForPointer discriminator the checksum and DFS staging paths
already use) and skip the writeback for device-resident destinations. The
read path is unaffected; the local-memcpy worker pool is already
device-aware, so cache hits into device memory keep working when the entry
was cached by a host read.

Add a regression test that registers a fake accelerator claiming a static
buffer as device memory and drives a real two-client TCP batch read: the
device-destined key stays out of the hot cache while the host-destined
control key in the same batch is written back.
@he-yufeng
he-yufeng force-pushed the fix/hot-cache-device-dest-skip branch from c0c5ccb to a3b0562 Compare October 20, 2026 03:43
@he-yufeng

Copy link
Copy Markdown
Collaborator Author

Rebased onto current main (5cf142e), which now carries #4781's DfsOnlyBatchFailureOnlyRevokesFailedKey fix — that was the failing gate here. Diff unchanged. Re-verified in the Linux container after the rebase: client_local_hot_cache_test builds clean and all 40 tests pass.

@ykwd
ykwd merged commit 67df8c4 into kvcache-ai:main Oct 20, 2026
27 checks passed
zxpdemonio added a commit to zxpdemonio/Mooncake that referenced this pull request Oct 21, 2026
…20261021

* upstream/main: (507 commits)
  [Reshard] Plan Store KV reuse across TP, PP and host layouts (kvcache-ai#4762)
  [Store] Walk tenants through shard-backed cursors and build the removes on them (kvcache-ai#4796)
  [Store-RS][Bugfix] Match TENT Request Layout to the C ABI (kvcache-ai#4732)
  [TransferEngine] Recover auto GID after data-path retry exhaustion (kvcache-ai#4290)
  [Store] Add the namespace layer over the object route (kvcache-ai#4721)
  [Store] Fix silent loss of offload keys deferred to the ungrouped pool (kvcache-ai#3006) (kvcache-ai#3009)
  [Store] Remove redundant alignment headroom from DFS shard allocation (kvcache-ai#4753)
  [CI/Build] Gate ROCm hardware tests and bound nightly Docker builds (kvcache-ai#4750)
  [Bugfix][Store] Fix double-counting of segment OpLog metrics in OpLogApplier (kvcache-ai#4771)
  [Store] Skip hot-cache writeback for device-resident read destinations (kvcache-ai#4764)
  [Reshard] Add bounded in-flight execution for weight transfers (kvcache-ai#4628)
  [Store] Avoid racing with detached hot-cache fills during key scans (kvcache-ai#4334)
  [Store] HA: drain the master on SIGTERM/SIGINT before exit (part 1 of kvcache-ai#4710) (kvcache-ai#4745)
  TENT: report a notification the receiver could not take (kvcache-ai#4754)
  [Store] Reach object metadata through a per-key accessor contract (kvcache-ai#4717)
  [Store][Bugfix] Fix DFS sync client test retry failure (kvcache-ai#4781)
  [Wheel] Keep integer dtype for typed_ragged rows when a row is an empty list (kvcache-ai#4611)
  [Doc] Sync README news onto the docs homepage (kvcache-ai#4767)
  [CI] remove stale Torch version binding after PG JIT (kvcache-ai#4760)
  [TransferEngine] Add buffer-protocol and page-gathered batch transfer APIs (kvcache-ai#4763)
  ...
zxpdemonio added a commit to zxpdemonio/Mooncake that referenced this pull request Oct 21, 2026
…-ranges-20261021

* upstream/main: (394 commits)
  [Reshard] Plan Store KV reuse across TP, PP and host layouts (kvcache-ai#4762)
  [Store] Walk tenants through shard-backed cursors and build the removes on them (kvcache-ai#4796)
  [Store-RS][Bugfix] Match TENT Request Layout to the C ABI (kvcache-ai#4732)
  [TransferEngine] Recover auto GID after data-path retry exhaustion (kvcache-ai#4290)
  [Store] Add the namespace layer over the object route (kvcache-ai#4721)
  [Store] Fix silent loss of offload keys deferred to the ungrouped pool (kvcache-ai#3006) (kvcache-ai#3009)
  [Store] Remove redundant alignment headroom from DFS shard allocation (kvcache-ai#4753)
  [CI/Build] Gate ROCm hardware tests and bound nightly Docker builds (kvcache-ai#4750)
  [Bugfix][Store] Fix double-counting of segment OpLog metrics in OpLogApplier (kvcache-ai#4771)
  [Store] Skip hot-cache writeback for device-resident read destinations (kvcache-ai#4764)
  [Reshard] Add bounded in-flight execution for weight transfers (kvcache-ai#4628)
  [Store] Avoid racing with detached hot-cache fills during key scans (kvcache-ai#4334)
  [Store] HA: drain the master on SIGTERM/SIGINT before exit (part 1 of kvcache-ai#4710) (kvcache-ai#4745)
  TENT: report a notification the receiver could not take (kvcache-ai#4754)
  [Store] Reach object metadata through a per-key accessor contract (kvcache-ai#4717)
  [Store][Bugfix] Fix DFS sync client test retry failure (kvcache-ai#4781)
  [Wheel] Keep integer dtype for typed_ragged rows when a row is an empty list (kvcache-ai#4611)
  [Doc] Sync README news onto the docs homepage (kvcache-ai#4767)
  [CI] remove stale Torch version binding after PG JIT (kvcache-ai#4760)
  [TransferEngine] Add buffer-protocol and page-gathered batch transfer APIs (kvcache-ai#4763)
  ...
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: [Store] batch_get_into into GPU memory crashes (SIGSEGV) when the local hot cache is enabled

3 participants