[jcc-google]

Summary

Pool reshard sessions (ReshardReceiveSession and ReshardSendSession) stop leasing StagingLayerManager host-buffer slots when the manager already has a full host mirror (num_host_layers() == num_layers(), i.e. no num_staging_layers). Transfers then stage and validate in place against each pool's own host storage, matching the existing layer-path sessions while keeping bounded host-layer staging unchanged.

Commits (review one at a time)

1. 1f08788 Address pool reshard receives in place on a full host mirror. ReshardReceiveSession stops leasing StagingLayerManager slots when every storage has its own host buffer. Fixes the sgl-project/sglang-torchtpu PP->TP failure (regression from Support bounded host layer staging for pool reshard receives #1391).
2. 6d869e3 Stage pool reshard sends in place on a full host mirror. The symmetric fix for ReshardSendSession (regression from Support bounded host layer staging for pool reshard sends #1392).

The bounded-host-layers ValidateChunks gap on explicit pools is split out to #1494; it is independent of this full-host-mirror regression.

Problem

Pool reshard transfers into a manager without num_staging_layers (every storage has its own host buffer) fail on the receiver with:

OUT_OF_RANGE: Chunk out of bounds. Chunk ptr: ..., Block array base: ..., Block array size: S

The chunk offset from the block array base is an exact multiple of the host buffer size (1x, 2x, 88x observed). The failed push finishes the ReshardReceiveSession (Incoming push failed), so subsequent pushes of the same transfer fail with Receive session ... is draining, No transfer chunks found for block N, and explicit-destination push ... has no registered plan on this pool-mode worker, while the sender sees Connection reset by peer.

Observed in sgl-project/sglang-torchtpu PP8 prefill -> TP8 decode (SGLANG_RAIDEN_PIPELINE_TRANSFER=1, Muse Spark). tpu-sync-torch 0.0.1.dev20261029081917 and earlier pass; 0.0.1.dev20261029204745 and later fail.

Root cause

• Support bounded host layer staging for pool reshard receives #1391 (951b49b, receive) and Support bounded host layer staging for pool reshard sends #1392 (3d98f15, send) route pool-reshard transfers through StagingLayerManager unconditionally, leasing a host-buffer slot for every pool even on a full host mirror (num_host_layers() == num_layers()).
• On explicit pools, addressing and D2H/H2D then use GetHostPointer(slot) + base_offset_bytes, whereas transport ValidateChunks checks incoming chunks against GetBlockArrayHostPointer(pool) (GetHostPointer(pool.storage_index) + base_offset_bytes).
• Because pushes from pipeline stages arrive in arbitrary order, slot usually differs from pool.storage_index, placing the chunk in a different host buffer and failing bounds validation.
• Layer-path TransferReceiveSession and TransferSendSession already size their slot manager to 0 and return kNotFound ("no bounded host staging") on a full host mirror; the reshard sessions missed that guard.

Fix

• In ReshardReceiveSession and ReshardSendSession, size staging_layer_manager_ to HostLayersBounded() ? base->num_host_layers() : 0 and return kNotFound from AcquireStagingSlot (and LookupStagingSlot on send) when !HostLayersBounded().
• Callers treat kNotFound as host_layer_idx = std::nullopt, so GetBlockChunks, D2hPoolBlocks, and H2dPoolBlocks address pool.storage_index in place.
• Bounded host-layer staging (num_host_layers() < num_layers()) from Support bounded host layer staging for pool reshard receives #1391 and Support bounded host layer staging for pool reshard sends #1392 is unchanged.

Tests

• Unit tests added / updated:
  • ReshardReceiveSessionTest.FullHostMirrorAddressesPoolsInPlace: pool 1 (storage 1) arrives first; verifies its chunks lie inside the span accepted by ValidateChunks and H2D reads each pool's own storage.
  • ReshardSendSessionTest.FullHostMirrorStagesPoolsInPlace: loopback push where pool order != storage order; passes transport ValidateChunks and delivers each pool's bytes in place.
  • ReshardSendSessionTest.StagingSlotReleasedAfterLastPeerPush: configured with bounded host layers (num_staging_layers), where slot leasing applies; kv_cache_manager_with_transfer_pool_reshard_test helper accepts kNotFound on a full mirror.
• Bazel verification (Linux ml-build container, clang-18):
  • //tpu_sync/kv_cache:reshard_receive_session_test, reshard_send_session_test, pool_staging_test, and kv_cache_manager_with_transfer_pool_reshard_test all pass at both 1f08788 and 6d869e3 individually.
  • With the production fix reverted and only the new tests kept, both new tests fail as expected (Chunk out of bounds on the send loopback test; {pool 1 -> buffer 0, pool 0 -> buffer 1} mismatch on the receive test).
  • Multi-chip TPU MPMD Bazel targets were not run directly inside Bazel; TPU hardware coverage was run end-to-end via the built tpu-sync-torch wheel below.

End-to-end validation

Setup: sgl-project/sglang-torchtpu Muse Spark (816B NVFP4 + DFlash2), PP8 prefill -> TP8 decode with SGLANG_RAIDEN_PIPELINE_TRANSFER=1 (pool reshard), tpu7x 2x2x1 per side (main 7e580f5936 + sgl-project/sglang-torchtpu#1950). Load: random ISL 1024 / OSL 256, 32 prompts, concurrency 4.

tpu-sync-torch	Result
0.0.1.dev20261028081934	PASS
0.0.1.dev20261029081917 (before 951b49b)	PASS
0.0.1.dev20261029204745 (9f5270b, sgl-project/sglang-torchtpu#1869 pin; rolled back in sgl-project/sglang-torchtpu#1970)	FAIL: Chunk out of bounds on decode
0.0.1.dev20261030081907	FAIL: same
9f5270b + this fix (backport)	PASS: 32/32 complete, 8192/8192 output tokens (retokenized match), greedy probe correct, DFlash accept length 6.35, 0 Raiden errors (Chunk out of bounds / no registered plan / Incoming push failed / Connection reset)
This PR (50b50d3 + fix)	PASS: 32/32 complete, 8192/8192 output tokens, greedy probe correct, DFlash accept length 6.34, 0 Raiden errors

Additional checks with the patched wheel:

• sgl-project/sglang-torchtpu PP8 P/D nightly suite (test_tpu_muse_spark_pd_singlehost_pp8_dflash.py): server config, KV offload reload, SWA passkey (cold & warm), GSM8K 0.975 (floor 0.88), DFlash accept length 4.72 (floor 4.0) — all PASS.
• TP8 -> TP8 P/D with bounded host-layer staging (num_staging_layers = [3, 5] of [21, 46] layers): 32/32 complete, accept length 6.01, 0 Raiden errors — PASS.

Receive-only backport onto 9f5270b (which predates send-side #1392): https://github.com/jcc-google/tpu-sync/tree/jcc/pool-reshard-staging-fix-on-9f5270b

[jcc-google]

Superseded by the internal CL, now landed on main as 47e9f20 ("Stage explicit-pool reshard transfers in place on a full host mirror"), rebased on top of the UUID-resolved validation fix (9dfd755c / d0f471b0). Closing.