[YAMY1234]

Motivation

With --enable-decoder-swa-bounded-replay, the layers after the last KV-source layer (21–39 on DeepSeek-V4.1-Flash) run only on each request's last 128 extend rows. That removes about 40% of a prefill step's GPU work. It does not make the step faster. An eager V4.1 prefill forward issues about 2,000 kernel launches per 8K chunk, so it is paced by the host, not the GPU. The trim removes GPU time but no launches.

python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V4.1-Flash --trust-remote-code \
  --tp 4 --ep-size 4 --max-running-requests 64 --chunked-prefill-size 8192 \
  --enable-decoder-swa-bounded-replay
python3 prefill_repro.py cold 8192 8     # 8 concurrent cold 8K prompts: input tok/s
python3 prefill_repro.py hit 65536 1     # 64K cached prefix + 1,024 new tokens: TTFT

prefill_repro.py (stdlib only)

# usage: python3 prefill_repro.py cold|hit TOKENS CONCURRENCY   (server at $URL, default :30000)
import json, os, random, statistics, sys, time, urllib.request
from concurrent.futures import ThreadPoolExecutor

URL = os.environ.get("URL", "http://127.0.0.1:30000") + "/generate"
MODE, N, CONC = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
rng = random.Random(N)
prefix = [rng.randrange(1000, 120000) for _ in range(N)]


def ttft(ids):
    body = {"input_ids": ids, "stream": True,
            "sampling_params": {"temperature": 0, "max_new_tokens": 1, "ignore_eos": True}}
    req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"})
    t0 = time.perf_counter()
    with urllib.request.urlopen(req) as r:
        while not r.readline().startswith(b"data:"):
            pass
        first = time.perf_counter() - t0
        r.read()
    return first


def wave():
    if MODE == "cold":
        prompts = [[rng.randrange(1000, 120000) for _ in range(N)] for _ in range(CONC)]
    else:
        prompts = [prefix + [rng.randrange(1000, 120000) for _ in range(1024)] for _ in range(CONC)]
    t0 = time.perf_counter()
    with ThreadPoolExecutor(CONC) as ex:
        times = list(ex.map(ttft, prompts))
    return times, sum(map(len, prompts)) / (time.perf_counter() - t0)


if MODE == "hit":
    ttft(prefix)  # put the prefix in the cache
for _ in range(2):
    wave()
runs = [wave() for _ in range(5)]
times = [t for ts, _ in runs for t in ts]
print(f"{MODE} {N} x{CONC}: TTFT p50 {1e3 * statistics.median(times):.0f} ms, "
      f"input {statistics.median(r for _, r in runs):,.0f} tok/s")

On main, cold prefill with the flag runs at about 49K input tok/s, the same as without it.

Modifications

• New python/sglang/srt/models/deepseek_v4_replay_graphs.py: DecoderReplayGraphs keeps one breakable CUDA graph per tail-row bucket (multiples of 128) and replays the trimmed late layers from it on eager prefill steps.
  • The SWA KV store, the attention and the low-ratio indexer/compressor sources stay eager breaks that read the step's live tail metadata. Everything else reads static input buffers that the bank owns; a run copies the live rows in and the outputs out.
  • A bucket is captured on first use, on every TP rank in the same step, and only after a full-chunk step has run, so row-sized workspaces that grow lazily cannot move under an earlier capture.
• DeepseekV4Model sends the tail to the bank when its row count fits a bucket. Steps with DSpark hidden-state capture, context parallelism, a prefill-graph replay, or on ROCm keep the eager path.
• On by default under the decoder flag: SGLANG_DSV4_DECODER_REPLAY_GRAPH_MAX_ROWS=2048 (16 graphs); 0 keeps the late layers eager. SGLANG_DSV4_DECODER_REPLAY_GRAPH_DEBUG=1 replays segment by segment with a sync and names a faulting segment.
• No change without --enable-decoder-swa-bounded-replay.
• Unit tests: test/registered/unit/model_executor/test_decoder_replay_graphs.py (buckets, state round trip, replay refill).

Memory: buckets are captured on first use into their own graph pool; after a benchmark run that captured 3 buckets it held 1.15 GiB per GPU. KV capacity is unchanged (62,784,256 tokens on main and with this PR).

Accuracy Tests

DeepSeek-V4.1-Flash, 4×GB300, TP4/EP4, both bounded-replay flags, no tuning flags. Cold prompts, then full prefix hits:

	main	this PR
GSM8K-200, cold / hit	196 / 197	195 / 195
needles at 32K (20) and 64K (10): cold / first hit / second hit	30/30, 30/30, 30/30	30/30, 30/30, 30/30

Speed Tests and Profiling

DeepSeek-V4.1-Flash on 4×GB300, TP4/EP4, --enable-decoder-swa-bounded-replay --chunked-prefill-size 8192 --max-running-requests 64 --mem-fraction-static 0.8, no other tuning flags. Main and this PR on the same node, two rounds in reversed order; each value is the average of both rounds. Cold points send random token-id prompts with one output token; hit points cache the prefix first and add 1,024 new tokens.

point	main	this PR	change
cold 8K, 8 concurrent, input tok/s	49,975	69,929 ¹	+40%
cold 32K, 8 concurrent, input tok/s	49,161	74,750	+52%
64K prefix + 1K new, TTFT p50	176.7 ms	121.5 ms	−31%
64K prefix + 1K new, 16 concurrent, TTFT p50	624.0 ms	398.2 ms	−36%
7K prefix + 1K new, TTFT p50	174.0 ms	115.7 ms	−33%
decode step, 32 concurrent, 8K prompt	7.51 ms	7.52 ms	noise

¹ 74.8K and 65.1K in the two rounds. Layers 0–20 still run eagerly, so the step is host-paced and sensitive to CPU noise (#43535 removes that).

KV capacity is unchanged (62,784,256 tokens). Buckets are captured on first use; in these runs the bank captured 3 buckets, 1.15 GiB per GPU.

## Checklist

- [x] Format your code according to the code style guidelines.
- [x] Add unit tests.
- [ ] Update documentation as needed (no user-facing flag; the env switches are documented in `environ.py`).
- [x] Provide accuracy and speed benchmark results.

<!-- pr-states:start -->
---
### CI States

Latest PR Test (Base): <!-- slot:pr-test:start -->:x: [Run #38288834408](https://github.com/sgl-project/sglang/actions/runs/38288834408)<!-- slot:pr-test:end -->
Latest PR Test (Extra): <!-- slot:pr-test-extra:start -->:x: [Run #38288834146](https://github.com/sgl-project/sglang/actions/runs/38288834146)<!-- slot:pr-test-extra:end -->
Latest PR Test (AMD ROCm 10): <!-- slot:pr-test-amd:start -->:x: [Run #38288834383](https://github.com/sgl-project/sglang/actions/runs/38288834383)<!-- slot:pr-test-amd:end -->
<!-- pr-states:end -->

[Oasis-Git]

Closed due to #43476 merged