Repository navigation
Expose no flat block array span when host layers are bounded - #1494
Draft
jcc-google wants to merge 1 commit into
Draft
jcc-google wants to merge 1 commit into
jcc-google wants to merge 1 commit into
Conversation
GetBlockArrayHostPointer returned nullptr for bounded host layers only on the legacy layer path. For explicit pools it still returned the pool's own storage buffer, although with bounded host layers a pool's bytes sit in whichever host buffer its transfer leased (and storages past num_host_layers() have no buffer of their own). ValidateChunks would then reject chunks in any other leased buffer. Hoist the bounded-host-layers check above the explicit-pools branch so both paths report no flat span, and ValidateChunks defers to the per-chunk addressing. Signed-off-by: Jincheng Chen <chenjincheng@google.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
With bounded host layers (
num_staging_layers < num_layers),GetBlockArrayHostPointerreturnsnullptron the legacy layer path but, for explicit pools (RegisterPools), still returns the pool's own storage buffer. A pool's bytes then sit in whichever host buffer its transfer leased, and storages pastnum_host_layers()have no buffer of their own. SoValidateChunkswould reject pool-reshard chunks in any other leased buffer (Chunk out of bounds).Split out of #1489, which fixes the full-host-mirror case (the regression hit by sglang-torchtpu PP->TP). No current user runs bounded host layers with explicit pools, so this is a latent issue in the configuration #1391 targets.
Fix
Hoist the bounded-host-layers check to the top of
GetBlockArrayHostPointer, so both the layer and pool paths report no flat span andValidateChunksdefers to the per-chunk addressing.Open question
This keeps the legacy behavior of skipping the flat check. A stricter alternative is to validate each chunk against the uuid's leased buffer. Happy to switch if preferred.
Tests
PoolStagingTest.BoundedHostLayersExposeNoFlatPoolSpan(3 explicit pools, 2 host layers: every pool reports no flat span).