[Disagg] GPU staging buffer with dynamic ring allocator for heterogeneous TP KV transfer

[YAMY1234]

Motivation

Design Docs

When prefill and decode use different TP layouts in disaggregated serving (e.g., TP4 prefill + DEP4 decode), KV cache transfer requires head slice redistribution. The original per-token RDMA approach (send_kvcache_slice) generates ~30,000 small RDMA requests per request, severely limiting transfer throughput at high concurrency.

This PR introduces a GPU staging buffer mechanism that gathers scattered head slices into contiguous memory for bulk RDMA transfer, reducing RDMA request count by ~1000x. It also includes a dynamic ring buffer allocator on the decode side to support variable-length requests without fixed slot exhaustion.

Modifications

Core staging mechanism (sglang/srt/disaggregation/common/staging.py):

• StagingBuffer: Pre-allocated GPU buffer (cuMemCreate-backed for NVLink compatibility)
• StagingAllocator: Ring buffer allocator — 1 large buffer with dynamic offset allocation, FIFO recycling
• gather_kv_head_slices / scatter_kv_head_slices: GPU gather/scatter between KV cache and staging buffer
• gather_all_layers_to_staging: Batched gather for all layers K+V with dedicated CUDA stream

Mooncake transfer integration (sglang/srt/disaggregation/mooncake/conn.py):

• send_kvcache_staged: Prefill-side gather + bulk RDMA to decode staging pool
• Updated bootstrap protocol: 1 base_ptr + total_size (replaces N fixed slot pointers)
• Updated per-request protocol: staging_offset (replaces staging_slot_id)
• Routing logic: staging path when attn_tp mismatch + staging enabled, graceful fallback to slice

Decode-side async scatter (sglang/srt/disaggregation/decode.py):

• Scatter on dedicated scatter_stream, zero synchronize on scheduler main thread
• CUDA Event-based completion tracking across scheduler iterations
• Scatter overlaps with decode forward on different CUDA streams

Configuration (sglang/srt/environ.py):

• SGLANG_DISAGG_STAGING_BUFFER: Enable/disable staging (default: False)
• SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB: Prefill per-worker staging size (default: 64MB)
• SGLANG_DISAGG_STAGING_POOL_SIZE_MB: Decode ring buffer total size (default: 4096MB)

Accuracy Tests

Test with the fix: #20655

Prefill TP=4, decode DEP4:

GPQA:

Non-staging(send_kvcache_slice):

Repeat: 8, mean: 0.859
Scores: ['0.833', '0.894', '0.864', '0.854', '0.854', '0.869', '0.859', '0.843']

Staging:

Repeat: 8, mean: 0.859
Scores: ['0.859', '0.864', '0.879', '0.889', '0.859', '0.848', '0.838', '0.833']

LongbenchV2(128k ISL, test chunk prefill transfer):

Non-staging:

100%|██████████| 503/503 [1:51:35<00:00, 13.31s/it]
Total latency: 6695.624 s
Score: 0.655

Staging:

100%|██████████| 503/503 [19:10<00:00,  2.29s/it]
Total latency: 1150.807 s
Score: 0.652

Benchmarking and Profiling

Setup: Qwen3.5-397B-A17B-FP8 on GB200, 1P1D (1 prefill node + 1 decode node), 4 GPUs per node, ISL=1000, OSL=1000, Mooncake MNNVL transport.

With staging buffer enabled (red), TP4→DEP4 heterogeneous transfer now matches the performance of homogeneous DEP4→DEP4 (green) and TP4→TP4 (blue) across all concurrency levels, a major improvement over the original per-token slice path (purple) which plateaus at ~2.4K tok/s.

Zooming into low-concurrency, without staging (purple, 1p1d-tp4dep4) throughput plateaus around 300 TPS/GPU due to per-toke overhead. With staging (red), throughput scales smoothly to 1700+ TPS/GPU, matching the homogeneous baselines (blue/green) that were previously unreachable.

Concurrency	TP4→DEP4 no staging	TP4→DEP4 + staging	DEP4→DEP4 (homogeneous)	TP4→TP4
1	87	90	90	133
8	556	611	614	777
32	1,510	1,758	1,788	2,002
64	2,200	2,984	3,015	3,226
128	2,466	4,914	4,959	5,253
256	2,344	7,842	8,054	8,269
512	2,383	11,795	12,465	11,649

(Output token throughput, tok/s. Higher is better.)

Key observations:

• Staging provides 2-3x throughput improvement over per-token slice at high concurrency (128+)
• Performance is within ~5% of homogeneous DEP4→DEP4 baseline
• At concurrency 512, staging approaches TP4→TP4 performance

Kernel Profiling

Prefill TP4, decode DEP4(attention tp size=1):

Prefill python stack(under 1k1k, concurrency=128) python stack: 840 us/writer:

Decode python stack(under 1k1k, concurrency=128) python stack: 678 us/scatter submission:

Because decode aggregates the KV from 4 prefill ranks and performs a unified scatter, the scatter kernel latency is roughly 4× that of prefill, which matches our expectations.

_fused_gather_to_staging_kernel
Metric                    1k/128      128k (16k chunk)/128
----------------------------------------------------------
Avg per kernel            7.97 µs     74.0 µs
Median per kernel         6.38 µs     70.5 µs
_fused_scatter_from_staging_kernel

Scatter (Staging Buffer → KV Cache)
Metric                    1k/128      128k (16k chunk)/128
----------------------------------------------------------
Avg per kernel            29.8 µs     385 µs
Median per kernel         28.0 µs     394 µs

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.

[gemini-code-assist]

Warning

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

[ShangmingCai]

Thx for the PR. Another contributor also discussed this idea with me before, the author of #18163. This implementation is cleaner. I think we can work together to get it done while minimizing the risk of OOM.
CC: @Ch3ngY1

[ShangmingCai]

Maybe use a better name with staging as a prefix or suffix to distinguish from the normal use case.

[YAMY1234]

Makes sense, change to use register_staging_room_bootstrap

[ShangmingCai]

Others LGTM. This is the final round of review. Thank you for putting up with me, and most importantly, thank you for contributing this great feature.

[ShangmingCai]

/tag-and-rerun-ci

[ShangmingCai]

LGTM

[Fridge003]

@ShangmingCai The failing test is unrelated to this PR, and it has been fixed on main branch

[Fridge003]

We need to add this environ to document