[DRAFT] Stage resharding transfers in plan-allocated host blocks - #736
Draft
yinlin09 wants to merge 1 commit into
Draft
[DRAFT] Stage resharding transfers in plan-allocated host blocks#736yinlin09 wants to merge 1 commit into
yinlin09 wants to merge 1 commit into
Conversation
yinlin09
force-pushed
the
yinlin/reshard-demand-staging
branch
from
August 21, 2026 14:15
9b2db5f to
cf44cde
Compare
Resharding resolved every block on the wire by identity into the host mirror: a sender staged each device block into the host block of the same id, and a receiver landed each incoming block at its destination id and uploaded from there. That only works when the host mirror spans the whole device block space, so both sides of a reshard pinned a host copy of their entire device KV pool. Under TPU_RAIDEN_DYNAMIC_HOST_STAGING=1 a plan now allocates host blocks for the device blocks it names -- its source blocks on a sender, its destination blocks on a receiver -- and records the mapping. Chunk resolution consults the mapping, the receiver's upload reads from the mapped blocks, and the blocks return to the pool when the receive completes or the sender plan is unregistered. The wire is unchanged: it still names device blocks. Senders obtain their staging blocks through plan_host_blocks() and copy into those instead of into the mirror. Off by default. Pool-addressed plans keep their existing addressing. **Validation** 8->4 reshard (prefill TP8 -> decode attention TP4 x DP2), Qwen3-1.7B on v7x, 128-token pages, 8k context, 64 x 512-token requests at 64-way concurrency. Flag off vs on, same build and load, two runs per arm. The build includes the multi-sender completion fix for pushed transfers. | | off | on (512-block pool) | |---|---|---| | host staging per rank | whole device pool mirrored: 1025 blocks, 1.8 GB prefill / 3.5 GB decode | 512 blocks, 0.9 GB / 1.8 GB | | pinned host memory, both hosts | 42 GB | 21 GB | | gsm8k, 200 questions (symmetric serve: 0.675, 0.705) | 0.690, 0.670 | 0.685, 0.665 | | greedy probe, 5 sequential + 8 concurrent repeats | repeatable, equal to the symmetric serve | same | | req/s | 12.1, 12.8 | 13.0, 12.4 | | mean TTFT / ITL, ms | 1187 / 28.2, 1225 / 27.0 | 1571 / 23.3, 1208 / 26.8 | | failed transfers | 0 | 0 | - The pool no longer tracks the device pool, so the saving grows with the model: on gpt-oss-120b TP8 the mirror is 56.5 GiB per rank (452 GiB per host); the same 64 in-flight requests fit in about 2.3 GiB per host. - One probe prompt alternates between two sane continuations across its repeats, in both arms.
yinlin09
force-pushed
the
yinlin/reshard-demand-staging
branch
from
August 21, 2026 14:30
cf44cde to
f25d9d9
Compare
yinlin09
force-pushed
the
yinlin/dynamic-host-staging
branch
from
August 21, 2026 18:21
f5a23d7 to
f630926
Compare
yinlin09
force-pushed
the
yinlin/reshard-demand-staging
branch
from
August 21, 2026 18:21
f25d9d9 to
642bb06
Compare
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
1 participant
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.
TL;DR
each plan stages its device blocks in host blocks allocated for the plan.
device pool mirrored → 512-block pool) at equal accuracy and throughput,
zero failed transfers. Expected on gpt-oss-120b TP8 for the same load as
[DRAFT] Allocate host staging per transfer instead of fixed slots #722: 452 GiB → ~2.3 GiB per 8-rank host. Wire protocol unchanged.
TPU_RAIDEN_DYNAMIC_HOST_STAGING=1); off by default.What changes
Resharding resolved every block on the wire by identity into the host mirror:
a sender staged each device block into the host block of the same id, and a
receiver landed each incoming block at its destination id and uploaded from
there. That only works when the host mirror spans the whole device block
space, so both sides of a reshard pinned a host copy of their entire device KV
pool.
RegisteredPlanhost_block_of: device block → host block for the blocks the plan names on this sideGetBlockChunks(plan branch)RegisterActivePlan(receiver)RegisterActivePlan(sender)UnregisterActivePlanplan_host_blocks(uuid, block_ids)Pool-addressed plans keep their existing addressing.
Validation
Build: torch wheel via
ci/build_wheel.sh; the new API is present in allthree ABI extensions.
End to end: 8→4 reshard — prefill TP8 → decode TP8 with attention TP4 × DP2,
dense Qwen3-1.7B on v7x, 128-token pages, context length 8192. Same image,
wheels and load in every arm; the only differences are the flag and the host
pool size. Host staging per rank is read from the engine start-up line
(
max_blocks × num_slotsblocks). Every arm below runs on a build that alsocarries the multi-sender completion fix for pushed transfers ("Complete a pushed block array only after every declared sender lands", separate PR) — without
it the asymmetric path returns wrong output in every configuration, on or off.
Correctness (greedy probe: 3 prompts × 5 sequential repeats + 8 identical
concurrent requests; 200-question few-shot gsm8k):
Throughput (
bench_serving, 64 × 512-token requests, 64-way, req/s · mean TTFT · mean ITL):One probe prompt alternates between two sane continuations across its 5 repeats
in both reshard arms (not in the symmetric reference); same with and without
this PR. gsm8k second runs: 0.670 / 0.000 off, 0.665 / 0.000 on.
21 GB on for this load, and the on-side figure is set by the pool size, not
by the model. On gpt-oss-120b TP8 the mirror is 56.5 GiB per rank (452 GiB
per host); the same 64 in-flight requests need about 2.3 GiB per host
(projected from the page arithmetic in [DRAFT] Allocate host staging per transfer instead of fixed slots #722).
Engine lines (one per instance; identical across ranks):