[Disagg][StagingBuffer][2/2] Support radix cache by YAMY1234 · Pull Request #30545 · sgl-project/sglang

YAMY1234 · 2026-07-08T16:14:39Z

Part 2 of 2 of "Staging buffer + RadixCache support". Its first commit is [1/2] #31217 (the staging-buffer robustness groundwork — the identical commit); this PR adds RadixCache support in the commits on top. #31217 should land first; review the RadixCache commits here.

Motivation

The heterogeneous-TP staging path identifies chunks positionally (chunk_idx = start_page // full_chunk_pages) against a uniform prefetched ring-allocation grid, so any send whose boundaries do not land on that grid silently breaks the protocol. A prefill radix-cache hit inflates the first send past one grid slot (and skews every later chunk index), causing ring-allocation overruns, decode-prefix scatter-offset corruption, and a blind tail scatter from chunk_infos[-1]. As a result the staging buffer could not be combined with radix cache.

Modifications

Builds on the robustness groundwork in #31217.

Prefill: floor non-last sends to the staging grid (base = decode_prefix_len) and split every send into one TransferKVChunk per grid slot; re-enable the cached-prefix early-send under staging, capped at an at-rest prefix snapshot (prefix_indices grows under overlap scheduling to cover launched-but-unresolved chunks).
Decode: scatter is fully arrival-driven (the last chunk now also sends CHUNK_READY); a room completes only after all-ranks Success with every allocation scattered and every scatter event finished, and stalls past the disaggregation waiting timeout surface as KVPoll.Failed; scatter offsets account for the decode-side cached prefix; writer fan-in accounting moves to the handler so room teardown can purge partial counts.
Tests: a prefill-TP4 → decode-TP2 e2e (test_disaggregation_different_tp.py) with chunked-prefill-size 256 + decode radix, so the shared prefix spans several grid slots and both sides reuse the cached prefix — the configuration that was broken before this change.

Accuracy Tests

Qwen3.5-397B-A17B-FP8, DEP4 prefill → TP4 decode, staging buffer + decode radix cache both enabled — the combination this PR fixes (a radix-cache hit under staging previously corrupted decode output / wedged the ring).

Test	Before	After
gsm8k (1319)	—	0.9802
128k needle-in-haystack (16 needles/haystack, ~127.5k-token shared cached prefix)	2 / 128	128 / 128

gsm8k's few-shot prefix fits within one grid slot, so it mainly confirms staging adds no regression; the needle probe uses a genuine ~127.5k-token shared cached prefix and directly exercises the radix-hit grid-split and decode-prefix scatter-offset paths.

Speed Tests and Profiling

Topology-level Pareto comparison on 8×GB300 (Qwen3.5-397B-A17B-NVFP4, AgentX 256k-context workload, one 4-GPU prefill node + one 4-GPU decode node). The heterogeneous run enables the staging buffer.

Prefill → Decode	Concurrency	Output TPS/user (p50) ↑	Total TPS/GPU ↑
TP4 → TP4	128	256.38	27,329
DEP4 → TP4 (staging)	192	281.02	27,286
Relative	—	+9.61%	−0.16%

At essentially matched cost efficiency (~27.3k total tokens/s/GPU), the DEP4-prefill → TP4-decode point has 9.61% higher output TPS/user in this comparison.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ✅ Run #30984792496
Latest PR Test (Extra): ✅ Run #31090378169

gemini-code-assist · 2026-07-08T16:14:45Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

Builds on the staging lifecycle hardening in the previous commit. The heterogeneous-TP staging path identifies chunks positionally (chunk_idx = start_page // full_chunk_pages) against a uniform prefetched ring-allocation grid, so any send whose boundaries do not land on that grid silently breaks the protocol. A prefill radix-cache hit inflates the first send past one grid slot (and skews every later chunk index), causing ring-allocation overruns, decode-prefix scatter offset corruption, and a blind tail scatter from chunk_infos[-1]. Prefill: floor non-last sends to the staging grid (base = decode_prefix_len) and split every send into one TransferKVChunk per grid slot; re-enable the cached-prefix early-send under staging, capped at an at-rest prefix snapshot (prefix_indices grows under overlap scheduling to cover launched-but-unresolved chunks). Decode: scatter is fully arrival-driven (the last chunk now also sends CHUNK_READY); a room completes only after all-ranks Success with every allocation scattered and every scatter event finished, and stalls past the disaggregation waiting timeout surface as KVPoll.Failed; scatter offsets account for the decode-side cached prefix; writer fan-in accounting moves to the handler so room teardown can purge partial counts. Adds test_staging_radix_grid.py covering the grid math, arrival-driven completion, fan-in accounting, release_room, and the decode-prefix page alignment guard.

The existing hetero-TP staging e2e uses the default chunked-prefill-size, so the shared gsm8k few-shot prefix fits in a single grid slot and never exercises the radix-hit grid-split or decode-prefix scatter-offset paths. Add a prefill-TP4->decode-TP2 class with chunked-prefill-size 256 and decode radix enabled, so the shared prefix spans several grid slots and both sides reuse cached prefix -- the configuration that was broken before this change.

Per .claude/rules/no-getattr-defensive.md, declare fields at construction and access them directly instead of defensive getattr/hasattr: - TransferKVChunk.staging_counted (bool = False) replaces the dynamically set _staging_counted flag. - StagingAllocator._scatter_stream (None) replaces the lazily monkey-patched attribute; _scatter_region uses "is None" and release_room accesses it directly. - Receivers always initialize chunk_staging_infos = [] so readers access it directly (the receiver-None checks are kept -- those are real). - Direct access for the always-set manager fields enable_staging / _staging_ctx and for receiver.bootstrap_infos.

The staging teardown-ordering fix only dropped per-room state on Success, so a Failed last chunk left transfer_infos / req_to_decode_prefix_len behind (cleaned up later only via the sender failure path). Restore the original behavior of tearing down on the last chunk regardless of Success/Failed -- but keep the outstanding-chunk guard so a deferred earlier chunk is never stranded, and keep a non-last Failed chunk's room (more chunks may follow). Applied symmetrically to the nixl and mooncake workers. Also seed _staging_outstanding on the transfer-worker test mock, which builds its manager via object.__new__ and so must mirror __init__.

ShangmingCai · 2026-08-03T08:56:05Z

+        if elapsed > self.completion_timeout:
+            logger.error(
+                "[STAGING] room=%s not complete %.0fs after all-ranks Success "
+                "(a scatter never arrived); failing the request.",
+                room,
+                elapsed,
+            )
+            decode_req._staging_failed = True


Should this also respect SGLANG_DISAGGREGATION_WAITING_TIMEOUT?

Yes, self.completion_timeout is initialized from SGLANG_DISAGGREGATION_WAITING_TIMEOUT in __init__, so this check already follows that env var.

ShangmingCai

Others look good. But the logic is becoming more and more complicated, we should do some cleanup in the near future.

register_decode_req raised (killing the prefill scheduler) when a decode request's cache_protected_len was not page-aligned. Mark the request failed via the existing _staging_failed flag instead -- the staging poll path surfaces it as KVPoll.Failed -- so one misconfigured decode endpoint cannot take down the whole prefill scheduler.

YAMY1234 · 2026-08-05T04:37:28Z

/rerun-test test/registered/disaggregation/test_disaggregation_different_tp.py

github-actions · 2026-08-05T04:37:55Z

Results for /rerun-test test/registered/disaggregation/test_disaggregation_different_tp.py:

🚀 8-gpu-h20 (1 test): ❌ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_different_tp.py

YAMY1234 · 2026-08-05T15:13:43Z

/rerun-test test/registered/disaggregation/test_disaggregation_different_tp.py

github-actions · 2026-08-05T15:14:08Z

Results for /rerun-test test/registered/disaggregation/test_disaggregation_different_tp.py:

🚀 8-gpu-h20 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_different_tp.py

YAMY1234 · 2026-08-06T07:38:14Z

/rerun-failed-ci

Co-authored-by: Shangming Cai <csmthu@gmail.com>

A prefill send could cross a staging slot boundary when a radix-cache hit or scheduler slicing left the send cursor off the chunked-prefill grid, so a single chunk exceeded one staging slot and failed geometry validation on decode (pages > slot capacity). Anchor the grid at the decode prefix length recorded at bootstrap, hold back partial slots on non-final sends, and split each send into one sub-send per staging slot. Only the final sub-send carries state indices and is marked last. Prefetch and the sender share staging_grid_tokens(). Backport of the prefill grid segmentation from upstream sgl-project#30545. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

YAMY1234 changed the title ~~[PD] Fix staging buffer incompatibility with radix cache~~ [PD] Support staging buffer with radix cache and radix-prefix pre-transfer Jul 8, 2026

YAMY1234 force-pushed the fix/staging-radix-grid-align branch from 0d3192c to ce4170f Compare July 14, 2026 17:54

YAMY1234 changed the title ~~[PD] Support staging buffer with radix cache and radix-prefix pre-transfer~~ [PD][2/2] Support staging buffer with radix cache Jul 14, 2026

YAMY1234 force-pushed the fix/staging-radix-grid-align branch from ce4170f to 9037dcf Compare July 14, 2026 22:29

YAMY1234 changed the title ~~[PD][2/2] Support staging buffer with radix cache~~ [Disagg][StagingBuffer][2/2] Support radix cache Jul 14, 2026

YAMY1234 force-pushed the fix/staging-radix-grid-align branch 4 times, most recently from e3527fd to 3e0c9cc Compare July 17, 2026 06:47

YAMY1234 commented Jul 17, 2026

View reviewed changes

Comment thread python/sglang/srt/disaggregation/prefill.py Outdated

YAMY1234 commented Jul 17, 2026

View reviewed changes

Comment thread python/sglang/srt/disaggregation/common/staging_handler.py Outdated

YAMY1234 force-pushed the fix/staging-radix-grid-align branch 6 times, most recently from ef2a87f to cdbab30 Compare July 18, 2026 05:49

YAMY1234 marked this pull request as ready for review July 18, 2026 07:36

YAMY1234 requested review from ByronHsu, HaiShaw, ShangmingCai, Ying1123, hnyls2002, merrymercy, sogalin and xiezhq-hermann as code owners July 18, 2026 07:36

ShangmingCai reviewed Jul 18, 2026

View reviewed changes

Comment thread python/sglang/srt/disaggregation/mooncake/conn.py Outdated

YAMY1234 force-pushed the fix/staging-radix-grid-align branch 2 times, most recently from 56250c0 to 462b1f7 Compare July 20, 2026 16:55

YAMY1234 added 4 commits July 31, 2026 11:13

YAMY1234 force-pushed the fix/staging-radix-grid-align branch from 0fef2a0 to ad0747a Compare July 31, 2026 18:40

ShangmingCai reviewed Aug 3, 2026

View reviewed changes

Comment thread python/sglang/srt/disaggregation/common/staging_handler.py Outdated

ShangmingCai reviewed Aug 3, 2026

View reviewed changes

ShangmingCai approved these changes Aug 3, 2026

View reviewed changes

YAMY1234 force-pushed the fix/staging-radix-grid-align branch from a3126da to 02644c3 Compare August 4, 2026 03:57

Merge branch 'main' into fix/staging-radix-grid-align

a56c894

Merge branch 'main' into fix/staging-radix-grid-align

086592c

ShangmingCai added the bypass-fastfail label Aug 6, 2026

ShangmingCai merged commit 05c7ebf into sgl-project:main Aug 6, 2026
366 of 410 checks passed

guptaishaan mentioned this pull request Aug 7, 2026

[disagg][mooncake] Keep decode status thread alive on bad messages #32909

Closed

amd-bot mentioned this pull request Aug 8, 2026

[CI Monitor] Daily Report - 2026-08-08 bingxche/sglang-ci-bot#160

Open

saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026

[Disagg][StagingBuffer][2/2] Support radix cache (sgl-project#30545)

c395088

Co-authored-by: Shangming Cai <csmthu@gmail.com>

YAMY1234 deleted the fix/staging-radix-grid-align branch August 25, 2026 16:16

hanming-lu mentioned this pull request Aug 27, 2026

[Kimi-K3] Deferred decode-side KV release (#35049 + #35360) on the zmq-mitigation base #36610

Draft

Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026

[Disagg][StagingBuffer][2/2] Support radix cache (sgl-project#30545)

cf1b0cf

Co-authored-by: Shangming Cai <csmthu@gmail.com>

kpham-sgl mentioned this pull request Sep 19, 2026

[PD] Bound cached-prefix DCP transfers by pack capacity #40376

Merged

4 tasks done

Navigation Menu

Repository navigation

[Disagg][StagingBuffer][2/2] Support radix cache - #30545

[Disagg][StagingBuffer][2/2] Support radix cache#30545
ShangmingCai merged 8 commits into
sgl-project:mainfrom
YAMY1234:fix/staging-radix-grid-align

YAMY1234 commented Jul 8, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jul 8, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ShangmingCai Aug 3, 2026

Uh oh!

YAMY1234 Aug 4, 2026

Uh oh!

ShangmingCai left a comment

Uh oh!

YAMY1234 commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026 •

edited

Loading

Uh oh!

YAMY1234 commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026 •

edited

Loading

Uh oh!

YAMY1234 commented Aug 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

YAMY1234 commented Jul 8, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

gemini-code-assist Bot commented Jul 8, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ShangmingCai Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

YAMY1234 Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai left a comment

Choose a reason for hiding this comment

Uh oh!

YAMY1234 commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

YAMY1234 commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

YAMY1234 commented Aug 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

YAMY1234 commented Jul 8, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Aug 5, 2026 •

edited

Loading

github-actions Bot commented Aug 5, 2026 •

edited

Loading