Motivation

In PD disaggregation mode, KV cache transfer happens after full prefill computation completes. For long prompts (≥1K tokens), this creates a significant TTFT bottleneck — the decode side must wait for all layers to be computed and then transferred sequentially.

This PR implements layer-pipelined KV transfer: instead of computing all layers then transferring all KV at once, we split layers into groups and transfer each group incrementally. Transfer of group N overlaps with GPU compute of group N+1, significantly reducing TTFT.

Related: #19931 (same direction, different approach)

Key Results

Benchmark environment: SGLang v0.4.10.post2, torch 2.7.1+cu126, Qwen2.5-72B-Instruct, 2×8 H20 GPUs (TP=4 each), 4×400G IB (RDMA), PD + Mooncake backend.
The pipelined KV transfer logic is algorithmically identical between v0.4.10.post2 and this PR — differences are limited to upstream API adaptation (environ.py registration, PP-aware pointers, EAGLE/staging guards). See "Code equivalence" section below.

TTFT (ms) — Prompt Length Sweep (C=32, output=256)

TTFT p95 (ms)

Throughput (output tok/s)

Multi-turn Dialogue (16 sessions × 10 turns)

Extreme Stress (C=64, prompt=4096, output=1024)

Design

The feature is controlled by three environment variables (registered in `environ.py`), disabled by default:

• `SGLANG_PIPELINED_KV_TRANSFER=true` — enable the feature
• `SGLANG_PIPELINE_GROUP_SIZE=10` — layers per group (override adaptive default)
• `SGLANG_PIPELINE_MIN_TOKENS=3072` — threshold; short prompts use normal path

How it works

1. `_get_pipeline_group_size(batch)` — per-batch decision: returns adaptive group_size (>0) or 0 to skip pipeline. A universal guard ensures models without `forward_split_prefill` safely fallback to the normal path. Short prompts also fall back with zero overhead.

2. `run_batch_pipelined(batch, group_size)` in `Scheduler` — splits forward into layer groups using `model_runner.forward_split_prefill()`, enqueues per-layer KV transfer via `send_layer()` after each group. CUDA events synchronize GPU→transfer ordering. Pre-computes state indices for hybrid models via `_prepare_pipelined_state_indices()`.

3. `process_batch_result_pipelined_prefill()` — result handler that dispatches to `run_batch_pipelined` instead of `run_batch`, then follows the same downstream logic (including EAGLE spec_info propagation, staging sync, A2A MoE finalization).

4. `MooncakeKVManager.send_kvcache_layer()` — single-layer RDMA transfer supporting both MHA and MLA architectures via `get_mha_kv_ptrs_with_pp` / `get_mla_kv_ptrs_with_pp`.

5. `TpModelWorker.forward_batch_generation_split_{init,layer,sample}()` — three-phase split forward: init attention backend → run N layers per call → sample after last group.

Call chain

```
event_loop_normal_disagg_prefill
→ _get_pipeline_group_size(batch)
→ >0: run_batch_pipelined → split_init → [split_layer + send_layer] × N → split_sample
→ 0: run_batch (unchanged)
```

Adaptive group_size (E1)

Instead of a fixed `SGLANG_PIPELINE_GROUP_SIZE`, group_size is automatically tuned based on prompt length to keep pipeline iterations in [6, 10]:

• Short prompts (<4K tokens): 10 iterations (maximize overlap)
• Medium prompts (4K-8K): 8 iterations (sweet spot)
• Long prompts (>8K): 6 iterations (reduce loop overhead)

User can still override via `SGLANG_PIPELINE_GROUP_SIZE` env var.

Different TP support (E2)

`send_kvcache_layer()` supports MHA head slicing when prefill TP ≠ decode TP, using vectorized numpy addressing (same math as `send_kvcache_slice`). MLA is TP-invariant and needs no slicing.

Mamba/SWA/NSA state support (E4)

Hybrid models (Jamba, FalconH1, DeepSeek-R1 with SWA) are fully supported. `_prepare_pipelined_state_indices()` pre-computes state indices before the layer loop, then passes them through `send_layer(state_indices=...)` on the last layer to trigger `maybe_send_extra()`. This covers:

• HybridLinearKVPool (Mamba SSM): `req_index_to_mamba_index_mapping`
• SWAKVPool (Sliding Window): windowed page indices via `translate_loc_from_full_to_swa`
• NSATokenToKVPool (Native Sparse Attention): full sequence page indices

No decode-side changes needed — decode already waits for all data (KV + state) before starting.

Universal guard + FalconH1 support (E7)

A universal `hasattr(model, "forward_split_prefill")` guard replaces the previous multimodal-only guard. This ensures:

• Models without `forward_split_prefill` (e.g. DeepSeek-V2/V3) safely fallback — no crash
• Models with `forward_split_prefill` (LLaMA, Qwen, Gemma, FalconH1, etc.) use pipelined path

FalconH1 (Mamba hybrid) now has `forward_split_prefill`, enabling layer-pipelined transfer. Each layer's attention produces KV cache (transferred per-layer via pipeline), while SSM state is sent once at the end via `maybe_send_extra()` (SSM state is fixed-size, independent of sequence length — no benefit from per-layer pipelining).

MTP/EAGLE compatibility (E8)

Reviewed and confirmed that `process_batch_result_pipelined_prefill` correctly propagates EAGLE `spec_info` (`topk_p`, `topk_index`, `hidden_states`) to requests — identical to the normal path. MTP decode-side rollback is purely a decode-phase operation with no interaction with prefill-time pipelined transfer. Also aligned `copy_done.synchronize()`, `routed_experts_output.finalize()`, and `maybe_cache_unfinished_req` with the normal result handler.

Zero regression guarantee

When `SGLANG_PIPELINED_KV_TRANSFER=false` (default):

• `_get_pipeline_group_size()` returns `0` on the first line
• `run_batch_pipelined`, `process_batch_result_pipelined_prefill`, and all `split_*` methods are never called
• `TransferKVChunk.layer_id` defaults to `None`, so `transfer_worker` always takes the existing path
• `add_transfer_request` new parameters have default values — existing callers unaffected

Code equivalence (v0.4.10.post2 → this PR)

Benchmarks were collected on v0.4.10.post2. This PR ports the same logic to upstream main with these adaptations:

• `os.environ.get()` → `envs.XXX.get()` (central environ.py registry)
• `send_kvcache_layer` uses PP-aware `get_mha_kv_ptrs_with_pp` / `get_mla_kv_ptrs_with_pp` (upstream helpers, equivalent at PP=1)
• `process_batch_result_pipelined_prefill` adapted for upstream's EAGLE spec info, staging buffer, and `report_prefill_stats` APIs
• Core algorithm (grouped layer forward loop, CUDA event sync, per-layer RDMA enqueue) is identical

Checklist

Modified Files

Future Work

• MLA `forward_split_prefill` — DeepSeek-V2/V3 has complex interactions with TBO, NSA context parallel, and A2A MoE. Deferred as separate PR by someone familiar with DeepSeek internals.
• PP support — Cross-PP-stage pipeline coordination (long-term)