[DRAFT] Allocate host staging per transfer instead of fixed slots by yinlin09 · Pull Request #722 · google/tpu-sync

yinlin09 · 2026-08-19T22:22:04Z

TL;DR

Host staging for disaggregated KV transfer is allocated per transfer
(TPU_RAIDEN_DYNAMIC_HOST_STAGING=1) instead of in fixed worst-case slots.
8× less host memory for the same load, measured (36 → 4.5 GiB per host),
zero failed transfers, throughput unchanged. At long context lengths fixed
slots stop fitting in the machine at all (table below).
An exhausted pool now fails loudly with a reason instead of reporting a
transfer that never happened.
Next, in [DRAFT] Stage resharding transfers in plan-allocated host blocks #736 (same allocator): the resharding path, which today mirrors
the whole device KV pool on the host — 452 GiB → ~2.3 GiB per host expected
for the same load (validated end to end on a small model in [DRAFT] Stage resharding transfers in plan-allocated host blocks #736).

Why it matters — host memory per 8-rank host, gpt-oss-120b TP8

scenario	fixed slots (today)	on demand (this PR)
8k context, 64-way concurrency — measured	36 GiB (64 slots, else requests fail)	4.5 GiB, zero failures, same throughput
128k context, same traffic — projected	576 GiB (64 slots); 36 GiB buys only 4 transfers in flight	~14 GiB
1M context, mostly 8k requests — projected	4.6 TiB (64 slots); 288 GiB buys 4 in flight	~100 GiB for 64 short + 1 long in flight
Resharding (asymmetric TP), 8k context, same load — #736 (measured on Qwen3-1.7B: host mirror → half the device pool at equal accuracy and throughput; gpt-oss figures projected)	452 GiB (full host mirror of the device KV pool, both sides)	~2.3 GiB (in-flight pages only)

Fixed slots are sized for the longest request the server accepts × the
concurrency wanted; on-demand staging is sized for the bytes actually in
flight. Resharding today goes further and mirrors the entire device KV pool on
the host so that block ids can be resolved by identity; the same allocator plus
a per-plan id map removes that requirement — that is #736, stacked on this PR.

What changes

commit	change
Stage disaggregation transfers on demand instead of in fixed slots	Both sides (`StartRead`, `StartPushInternal`) allocate exactly the pages they stage from the existing host block manager, locked for the transfer's lifetime, and return them on completion/timeout/error. Copy plans already address host blocks explicitly: nothing changes downstream or on the wire. Off by default; reshard/pool-plan paths untouched.
Fail loudly when transfer staging is exhausted	Producer waits for staging to free (bounded by the transfer deadline) instead of marking the send done without transferring; a request that still cannot be seated, or can never fit, is reported as a failure. Both sides log request, pages asked, pages free.

Validation


Test	End-to-end disaggregated-prefill serving (prefill/decode split as in vLLM's P/D disaggregation), SGLang on TPU with raiden as the KV transport
Hardware / model	1 prefill host + 1 decode host, v7x, TP8 each; gpt-oss-120b (hybrid: two engines per request)
KV / context	64-token pages, 0.56 MiB per page per rank; context length 8192
Load	96 short prompts at 32-way burst; 64 × 512-tok; 32 × 4096-tok; 3 × 600 × 2048-tok paced; 200-question gsm8k
Per-request path	each decode rank pulls the request's pages from its prefill peer via `StartRead` / pull-serve
Metric	raiden `failed_recving` per consumer rank ("failed reads"); the serving layer retries, so client success rates hide failures
Build	jax + torch wheels via `ci/build_wheel.sh`; each commit builds standalone; the new path is present in all 3 torch ABI extensions

Results

claim	flag off	flag on
Failed reads / timeouts, all phases	0 / 0	0 / 0
gsm8k accuracy	0.870	0.850 (pre-change baseline 0.850)
Pool needed for this load, zero failures	8192 pages = 4.50 GiB/rank = 36 GiB/host	1024 pages = 0.56 GiB/rank = 4.5 GiB/host
Same load at a pool fixed slots can't serve	4 slots → 1112 failed reads across 139 requests	—
Throughput at the ⅛ pool (512-tok / 4096-tok / burst, req/s)	9.7 / 4.9 / 21.2 (full pool)	9.1 / 4.8 / 20.6

Throughput, same 8192-page pool, flag off vs on (req/s · mean/p99 TTFT ms)

phase	flag off	flag on
96 short prompts, 32-way burst	21.2 · 644/800	22.0 · 722/1230
64 × 512-tok	9.7 · 1793/2540	8.6 · 1688/2735
32 × 4096-tok	4.9 · 3012/5590	4.6 · 3216/5903
600 × 2048-tok, paced (×3)	3.86–3.87 · 79–204	3.86–3.87 · 90–217

Run-to-run variation, measured with two flag-off runs: the three unpaced
phases (the 32-way burst and the 64 × 512-token and 32 × 4096-token runs,
each a few seconds long, requests sent as fast as the server accepts them)
differ from each other by up to 12% in req/s; the paced phase (600 requests
offered at a fixed 4 req/s over about 2.5 minutes) is identical across runs,
as expected while the server keeps up with the offered rate.

Exhaustion behaviour: offered load above pool capacity → reads wait
(bounded by the transfer deadline); a read that can never fit fails with a
logged reason (request, blocks asked, blocks free). Before: producer reported
success without transferring, consumer timed out.

Found during validation, fixed here: the first build leaked staging on
every successful read (event-driven completion paths still released by slot
index). The added error line made it diagnosable in one run.

Follow-ups (not in this PR)

Resharding: replace the identity-addressed host mirror with per-plan host
staging on both sides (452 GiB → ~2.3 GiB per host projected for this load)
— [DRAFT] Stage resharding transfers in plan-allocated host blocks #736.
Consumer-side wait-with-deadline on a transiently full pool (today the
consumer fails fast while the producer waits).
Longer mixed-size soak for fragmentation before making the flag the default
(the paced phase ran 1800 transfers without drift).

A producer that could not seat a pull-serve marked the send done without transferring, so the consumer learned of the failure only by waiting out its transfer timeout; a consumer that could not seat a read failed it silently. The producer now waits for staging to free, bounded by the transfer deadline, and reports a transfer failure when a request still cannot be seated or can never fit. Both sides log the request, the pages asked, and the pages free, so an exhausted pool is attributed where it happens. Measured on the setup of the previous commit: at 4 fixed slots the load exhausted staging 1112 times across 139 requests, each visible only as a consumer-side failure. With staging on demand none occurred. Memory and performance results for on-demand staging itself are in the commit that introduces it, "Stage disaggregation transfers on demand instead of in fixed slots" (github.com//pull/722).

Disaggregation staging today either pre-carves fixed host slots sized for the maximum context length, over-provisioned when most in-flight requests are small, or mirrors the device KV pool on the host, wasteful when the blocks actually transferred are a small fraction of that pool. This change addresses the slot path on both ends. Under TPU_RAIDEN_DYNAMIC_HOST_STAGING=1 an incoming read and an outgoing pull-serve each allocate exactly the pages they stage and return them on completion, so a pool seats transfers by their size rather than by a fixed count. The copy plans already address host blocks explicitly, so nothing downstream changes. Off by default; the pool-plan receive path and the behaviour on an exhausted pool are unchanged. **Validation** gpt-oss-120b, prefill TP8 -> decode TP8 on v7x, 64-token pages, 8k context, 512-4096-token requests at 64-way concurrency, same build and load in every arm. Flag off carves the pool into fixed 128-page slots; flag on allocates pages per transfer from the same pool. | | off, 8192-page pool (64 slots) | on, 8192-page pool | on, 1024-page pool | |---|---|---|---| | pinned host staging per rank / per 8-rank host | 4.5 GiB / 36 GiB | 4.5 GiB / 36 GiB | 0.56 GiB / 4.5 GiB | | failed transfers | 0 | 0 | 0 | | gsm8k, 200 questions, before / after the load | 0.870 / 0.845 | 0.850 / 0.880 | 0.845 / 0.850 | | req/s: 96 x 128-tok burst / 64 x 512-tok / 32 x 4096-tok | 21.2 / 9.7 / 4.9 | 22.0 / 8.6 / 4.6 | 20.6 / 9.1 / 4.8 | | req/s: 600 x 2048-tok offered at 4 req/s, three runs | 3.86-3.87 | 3.86-3.87 | - | - Two flag-off runs differ by up to 12% req/s on the unpaced phases; the differences in the table are within that spread. - As fixed 128-page slots, 1024 pages seat 8 transfers at a time; the load runs 64-way. - At a 1M-token context a fixed slot is 16K pages while an 8k-token request uses 128: 64 slots would pin 4.6 TiB per host, 288 GiB would seat 4 transfers, and per-transfer pages serve 64 short requests plus one 1M request from about 100 GiB.

A producer that could not seat a pull-serve marked the send done without transferring, so the consumer learned of the failure only by waiting out its transfer timeout; a consumer that could not seat a read failed it silently. The producer now waits for staging to free, bounded by the transfer deadline, and reports a transfer failure when a request still cannot be seated. A request larger than the staging can ever seat -- a fixed slot, or the whole per-transfer pool -- fails at once. The misses while waiting are not logged; the one failure message on each side carries the request, the pages asked, and the pages free, so an exhausted pool is attributed where it happens. Measured on the setup of the previous commit: at 4 fixed slots the load exhausted staging 1112 times across 139 requests, each visible only as a consumer-side failure. With staging on demand none occurred. Memory and performance results for on-demand staging itself are in the commit that introduces it, "Stage disaggregation transfers on demand instead of in fixed slots" (github.com//pull/722).

A producer that could not seat a pull-serve marked the send done without transferring, so the consumer learned of the failure only by waiting out its transfer timeout; a consumer that could not seat a read failed it silently. The producer now waits for staging to free, bounded by the transfer deadline, and reports a transfer failure when a request still cannot be seated. A request larger than the staging can ever seat -- a fixed slot, or the whole per-transfer pool -- fails at once. The misses while waiting are not logged; the one failure message on each side carries the request, the pages asked, and the pages free, so an exhausted pool is attributed where it happens. Measured on the setup of the previous commit: at 4 fixed slots the load exhausted staging 1112 times across 139 requests, each visible only as a consumer-side failure. With staging on demand none occurred. Memory and performance results for on-demand staging itself are in the commit that introduces it, "Stage disaggregation transfers on demand instead of in fixed slots" (github.com//pull/722). During shutdown a producer waiting for staging stops instead of sleeping out its deadline, and the manager's destructor waits for every pull-serve worker before tearing down the state they read.

A producer that could not seat a pull-serve marked the send done without transferring, so the consumer learned of the failure only by waiting out its transfer timeout; a consumer that could not seat a read failed it silently. The producer now waits for staging to free, bounded by the transfer deadline, and reports a transfer failure when a request still cannot be seated. A request larger than the staging can ever seat -- a fixed slot, or the whole per-transfer pool -- fails at once. The misses while waiting are not logged; the one failure message on each side carries the request, the pages asked, and the pages free, so an exhausted pool is attributed where it happens. Measured on the setup of the previous commit: at 4 fixed slots the load exhausted staging 1112 times across 139 requests, each visible only as a consumer-side failure. With staging on demand none occurred. Memory and performance results for on-demand staging itself are in the commit that introduces it, "Stage disaggregation transfers on demand instead of in fixed slots" (github.com//pull/722). During shutdown a producer waiting for staging stops instead of sleeping out its deadline, and the manager's destructor waits for every pull-serve worker before tearing down the state they read. A producer entry that reaches its deadline without being pulled is reported as failed, never as sent, and the wait for staging shares that same deadline instead of starting a later one. Shutdown wakes workers parked waiting for a request that will never arrive, so tearing the manager down cannot hang on them.

A producer that could not seat a pull-serve marked the send done without transferring, so the consumer learned of the failure only by waiting out its transfer timeout; a consumer that could not seat a read failed it silently. The producer now waits for staging to free, bounded by the transfer deadline, and reports a transfer failure when a request still cannot be seated. A request larger than the staging can ever seat -- a fixed slot, or the whole per-transfer pool -- fails at once. The misses while waiting are not logged; the one failure message on each side carries the request, the pages asked, and the pages free, so an exhausted pool is attributed where it happens. Measured on the setup of the previous commit: at 4 fixed slots the load exhausted staging 1112 times across 139 requests, each visible only as a consumer-side failure. With staging on demand none occurred. Memory and performance results for on-demand staging itself are in the commit that introduces it, "Stage disaggregation transfers on demand instead of in fixed slots" (github.com//pull/722). During shutdown a producer waiting for staging stops instead of sleeping out its deadline, and the manager's destructor waits for every pull-serve worker before tearing down the state they read. A producer entry that reaches its deadline without being pulled is reported as failed, never as sent, and the wait for staging shares that same deadline instead of starting a later one. Shutdown wakes workers parked waiting for a request that will never arrive, so tearing the manager down cannot hang on them. A copy or network failure while serving a pull is the transfer's failure: it is reported as such instead of as a completed send, and no longer escapes the worker as an exception.

Previously, handling staging buffer exhaustion had two flaws: 1. Producer false success: When serving a pull request without free staging buffers, the producer marked the send as completed without transferring any data. The consumer received nothing and hung until its transfer deadline timed out. 2. Consumer silent failure: When initiating a read without free staging buffers, the consumer failed without logging an error. This change improves reliability, observability, and shutdown safety: - Producer retry and failure propagation: The producer now waits for staging buffers to free up, bounded by the request deadline. If the request exceeds total pool capacity or times out, it explicitly records a transfer failure. - Diagnostic logging: Both endpoints now log the request ID, requested page count, and available capacity on allocation failure, pinpointing which node ran out of memory. - Clean shutdown: Tracks active pull workers so the destructor waits for in-flight transfers to finish before tearing down resources. Validation: - With 4 fixed slots, the workload encountered 1,112 buffer exhaustion events across 139 requests (all previously surfacing only as consumer timeouts). - With on-demand dynamic staging enabled, 0 exhaustion events occurred.

yinlin09 force-pushed the yinlin/dynamic-host-staging branch 6 times, most recently from cd04936 to e53b335 Compare August 21, 2026 05:54

yinlin09 force-pushed the yinlin/dynamic-host-staging branch from e53b335 to d1fcf2c Compare August 21, 2026 05:56

yinlin09 force-pushed the yinlin/dynamic-host-staging branch from d1fcf2c to f5a23d7 Compare August 21, 2026 06:03

yinlin09 mentioned this pull request Aug 21, 2026

[DRAFT] Stage resharding transfers in plan-allocated host blocks #736

Draft

yinlin09 force-pushed the yinlin/dynamic-host-staging branch from f5a23d7 to f630926 Compare August 21, 2026 18:21

yinlin09 force-pushed the yinlin/dynamic-host-staging branch from f630926 to 6217af1 Compare August 23, 2026 17:37

yinlin09 force-pushed the yinlin/dynamic-host-staging branch from 6217af1 to db7a7aa Compare August 23, 2026 20:38

yinlin09 force-pushed the yinlin/dynamic-host-staging branch from db7a7aa to 7074f0d Compare August 23, 2026 20:56

yinlin09 force-pushed the yinlin/dynamic-host-staging branch 3 times, most recently from 8d88d6b to 89ef10d Compare August 24, 2026 07:42

yinlin09 force-pushed the yinlin/dynamic-host-staging branch from 89ef10d to b47cf1c Compare August 24, 2026 07:51

Navigation Menu

Uh oh!

[DRAFT] Allocate host staging per transfer instead of fixed slots - #722

[DRAFT] Allocate host staging per transfer instead of fixed slots#722
yinlin09 wants to merge 2 commits into
mainfrom
yinlin/dynamic-host-staging

yinlin09 commented Aug 19, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yinlin09 commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TL;DR

Why it matters — host memory per 8-rank host, gpt-oss-120b TP8

What changes

Validation

Follow-ups (not in this PR)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yinlin09 commented Aug 19, 2026 •

edited

Loading