[Bugfix][TENT] Fix silent TPU data corruption for transfers larger than one staging chunk#2815
Conversation
…an one staging chunk Caught while testing the TPU staging path on a real TPU VM (v5p-8, libtpu 0.0.32, PJRT C API 0.83) for the first time. Until now the feature had only ever run against the mock adapter. ProxyManager stages a transfer in chunk_size (4 MiB) pieces and passes `token + chunk_offset` to the platform for every chunk after the first (proxy_manager.cpp:76). The adapter ABI only ever specified base-address classification, so `isDevicePtr(token + off)` returned false, TENT classified TPU HBM as host memory, and TpuPlatform::copy fell through to CpuPlatform::copy -- a plain memcpy of the device token. On real PJRT that memcpy does not crash: PJRT_Buffer_UnsafePointer returns an address that is host-readable but does not hold the buffer's data. The staging copy therefore produced garbage and reported COMPLETED. Chunk 0 was correct and every subsequent chunk was silently wrong, so any transfer over 4 MiB was corrupted without an error anywhere. The mock could not catch this because its "device" pointers were ordinary host memory, so the accidental memcpy happened to produce the right bytes. Fixes: - tpu_pjrt_abi.h / tpu_pjrt_shim.h: state that classification and copy entrypoints must resolve interior addresses to the registered buffer whose range contains them, that a copy may not run past a buffer's end, and that the token must never be dereferenced. - tpu_transport.cpp: require exactly one TPU-device side per staging hop and fail loudly otherwise, so a non-conforming or absent adapter can no longer degrade into a silent memcpy. - mock_tpu_pjrt_adapter.cpp: model the two properties that matter -- tokens are opaque (a poisoned read-only mapping, data lives in shadow storage) and addresses may be interior. A copy that bypasses the adapter now yields 0xDD instead of accidentally passing. - common.cmake: -DUSE_TPU=ON silently compiled zero TPU code, because all TPU sources live under tent/ which is gated on USE_TENT. This is why the bug was invisible to every build. Now a hard error. Tests: two new regression tests in tent_tpu_pjrt_shim_test and a new tent_tpu_transport_test (6 cases) that drives the staging hop the way ProxyManager does. Verified they fail against the pre-fix code (LocalStageCopiesFromInteriorDeviceOffset reports COMPLETED while delivering 0xDD) and pass after. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support and validation for interior pointers (offsets within registered buffers) during TPU transfers. It updates the TPU/PJRT ABI, shim, and transport layer to correctly resolve and copy to/from interior addresses, preventing silent corruption during chunked staging transfers. It also enhances the mock TPU adapter to poison tokens and simulate real PJRT behavior, and adds comprehensive unit tests for both the shim and the transport layer. The feedback identifies a potential resource leak in mock_tpu_pjrt_register_device where mmap'ed memory could be leaked if the subsequent allocation of the shadow buffer throws a std::bad_alloc exception, and suggests allocating the shadow buffer first to ensure exception safety.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| void *mock_tpu_pjrt_register_device(size_t size, int index) { | ||
| void *token = mmap(nullptr, size, PROT_READ | PROT_WRITE, | ||
| MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); | ||
| if (token == MAP_FAILED) return nullptr; | ||
| std::memset(token, kPoison, size); | ||
| mprotect(token, size, PROT_READ); | ||
|
|
||
| std::lock_guard<std::mutex> lock(g_mutex); | ||
| g_device_registry[addr] = index; | ||
| g_device_registry.push_back(Buffer{reinterpret_cast<uintptr_t>(token), | ||
| new unsigned char[size](), size, index}); | ||
| return token; | ||
| } |
There was a problem hiding this comment.
In mock_tpu_pjrt_register_device, if new unsigned char[size]() throws a std::bad_alloc exception, the previously allocated token via mmap will be leaked because munmap is never called. Allocating the shadow buffer before calling mmap ensures exception safety and prevents resource leaks.
void *mock_tpu_pjrt_register_device(size_t size, int index) {
unsigned char *shadow = new unsigned char[size]();
void *token = mmap(nullptr, size, PROT_READ | PROT_WRITE,
MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
if (token == MAP_FAILED) {
delete[] shadow;
return nullptr;
}
std::memset(token, kPoison, size);
mprotect(token, size, PROT_READ);
std::lock_guard<std::mutex> lock(g_mutex);
g_device_registry.push_back(Buffer{reinterpret_cast<uintptr_t>(token),
shadow, size, index});
return token;
}
staryxchen
left a comment
staryxchen
left a comment
There was a problem hiding this comment.
Thanks for the deep investigation and the thorough write-up. The ABI contract clarification, the transport-level guard in TpuTransport::startTransfer, the mock adapter rewrite with poisoned tokens, and the new test coverage all look excellent — approving those.
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
4 participants
Description
Fixes a silent data-corruption bug in the TENT TPU staging path added by #2733 (RFC #2662).
Found by running the TPU path on a real TPU VM for the first time (
v5p-8,libtpu 0.0.32, PJRT C API 0.83). Everything before this had only ever run against the mock adapter.The bug
ProxyManagerstages a transfer inchunk_size(4 MiB) pieces and passestoken + chunk_offsetto the platform for every chunk after the first:The adapter ABI only ever specified base-address classification. So for
offset > 0:mc_tpu_pjrt_is_device_ptr(token + offset)returns 0,TpuPlatform::copyconcludes neither side is device memory,CpuPlatform::copy— a plainmemcpyof the device token.On real PJRT that
memcpydoes not crash.PJRT_Buffer_UnsafePointerreturns an address that is host-readable but does not hold the buffer's data. The staging copy produced garbage and reportedCOMPLETED.Net effect: chunk 0 is correct, every later chunk is silently wrong. Any TPU transfer larger than 4 MiB was corrupted, with no error raised anywhere — which is essentially every real KV-cache transfer.
Why CI never caught it
Two independent reasons, both fixed here:
-DUSE_TPU=ONcompiled zero TPU code. All TPU sources live undertent/, which is gated onUSE_TENT. The build configured and succeeded while building nothing.memcpyproduced the right bytes. The mock structurally could not reproduce the failure.Changes
tpu_pjrt_abi.h/tpu_pjrt_shim.h— specify that classification and copy entrypoints must resolve an interior address to the registered buffer whose range contains it; that a copy may not run past a buffer's end; and that the token must never be dereferenced.tpu_transport.cpp— require exactly one TPU-device side per staging hop, and fail loudly otherwise. Defense in depth: a non-conforming or absent adapter can no longer degrade into a silentmemcpy.mock_tpu_pjrt_adapter.cpp— model the two properties that actually matter on hardware: tokens are opaque (a poisoned read-only mapping; data lives in shadow storage) and addresses may be interior. A copy that bypasses the adapter now yields0xDDinstead of accidentally passing.common.cmake—USE_TPU=ONwithoutUSE_TENT=ONis now a hard error instead of a silent no-op.Module
mooncake-transfer-engine)mooncake-store)mooncake-ep)mooncake-pg)mooncake-integration)mooncake-p2p-store)mooncake-wheel)mooncake-common)mooncake-rl)Type of Change
Note: the
common.cmakeguard is technically build-breaking for anyone passing-DUSE_TPU=ONalone — but such a build produced no TPU code, so nothing can regress.How Has This Been Tested?
Test commands:
cmake -S . -B build -G Ninja -DUSE_TENT=ON -DUSE_TPU=ON \ -DWITH_STORE=OFF -DWITH_STORE_RUST=OFF -DBUILD_BENCHMARK=OFF -DUSE_RDMA=ON cmake --build build -j ctest --test-dir build/mooncake-transfer-engine/tentTest results:
100% tests passed, 0 tests failed out of 22tent_tpu_transport_test(6 cases) drivesTpuTransport's staging hop exactly asProxyManagerissues it — local stage at an interior device offset, the mirrored remote stage, the READ direction, and the three failure modes (no device side, two device sides, non-local target).ClassifiesInteriorDevicePointersandCopyFromInteriorOffsetMovesTheRightBytesintent_tpu_pjrt_shim_test.Verified the tests actually catch the bug. Reverting the fixes (exact-match registry + no transport guard) makes
LocalStageCopiesFromInteriorDeviceOffsetfail exactly as the hardware behaves: statusCOMPLETED, staging buffer filled with0xDD. Re-applying the fixes turns it green.Hardware findings behind this were reproduced with a standalone PJRT probe (
GetPjrtApifromlibtpu.so→BufferFromHostBuffer→UnsafePointer→memcpy), confirming the token is readable and holds unrelated bytes.Not covered: the cross-node hop. This VM is single-host (gVNIC, no RDMA), so the device↔host staging hop is exercised, not host↔remote-host.
Checklist
./scripts/code_format.shpre-commit run --all-filesand all hooks pass — see note belowOn
pre-commit: every hook passes on the touched files exceptcmake-format, which rewrites large amounts of unrelated code intent/tests/CMakeLists.txtandcommon.cmake. PerAGENTS.md("do not include those unrelated edits in the PR") I kept the diff minimal and left those rewrites out; happy to letpre-commit.ciautofix, or to land them separately.Related work
[TransferEngine] Add TPU integration scaffolding and guards), which targets the classic transfer engine (mooncake-transfer-engine/src,include/gpu_vendor/tpu.h). That PR introduces its ownUSE_TPUoption. If it lands, theUSE_TPU ⇒ USE_TENTguard here needs revisiting, sinceUSE_TPUwould then also gate classic-TE code. Flagging for whoever merges second.Separately, hardware benchmarking turned up an architectural issue — the PJRT sub-range copy entrypoints that
Platform::copyrequires run ~20× slower than the whole-buffer ones (0.8 GB/s vs 18.5 GB/s D2H), capping the staging design near ~1 GB/s per chip. That is a design discussion, not a bugfix, so I'll raise it on #2662 rather than here.AI Assistance Disclosure
Claude Code was used to run the TPU-VM investigation (building the PJRT probes that identified the token/interior-pointer behaviour), to implement the fixes, and to draft the tests. I have reviewed every changed line and can defend the change end to end.