[PD]feat(bench): add --fake-prefill flag for decode-only stress testing (#22973)

This commit is contained in:
ybyang
2026-04-16 13:57:55 -07:00
committed by GitHub
parent 29f56cb230
commit 41258f874d
3 changed files with 66 additions and 0 deletions
+35
View File
@@ -341,6 +341,41 @@ python3 -m sglang.bench_serving \
--random-output-len 256 --random-output-len 256
``` ```
10) Fake decode stress testing (PD disaggregation, decode-only):
When benchmarking pure decode performance in a PD disaggregation setup, you can bypass the prefill node entirely by using `--fake-prefill`. This requires the decode server to be started with `--disaggregation-transfer-backend fake`:
```bash
# Step 1: Start a decode-only server with fake transfer backend
python -m sglang.launch_server \
--model-path meta-llama/Llama-3.1-8B-Instruct \
--disaggregation-mode decode \
--disaggregation-transfer-backend fake \
--port 30001
# Step 2: Run bench_serving with --fake-prefill
python3 -m sglang.bench_serving \
--backend sglang \
--host 127.0.0.1 --port 30001 \
--model meta-llama/Llama-3.1-8B-Instruct \
--dataset-name random \
--num-prompts 500 \
--random-input-len 1024 --random-output-len 256 \
--fake-prefill
```
Similarly, `bench_one_batch_server` also supports `--fake-prefill`:
```bash
python3 -m sglang.bench_one_batch_server \
--base-url http://127.0.0.1:30001 \
--model-path meta-llama/Llama-3.1-8B-Instruct \
--batch-size 32 --input-len 1024 --output-len 256 \
--fake-prefill
```
The `--fake-prefill` flag automatically injects special sentinel values into each request, telling the decode server to skip real KV transfer and generate fake KV data locally.
### Troubleshooting ### Troubleshooting
- All requests failed: verify `--backend`, server URL/port, `--model`, and authentication. Check warmup errors printed by the script. - All requests failed: verify `--backend`, server URL/port, `--model`, and authentication. Check warmup errors printed by the script.
+14
View File
@@ -44,6 +44,7 @@ from sglang.benchmark.utils import (
remove_prefix, remove_prefix,
set_ulimit, set_ulimit,
) )
from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST
from sglang.srt.utils.network import NetworkAddress from sglang.srt.utils.network import NetworkAddress
_ROUTING_KEY_HEADER = "X-SMG-Routing-Key" _ROUTING_KEY_HEADER = "X-SMG-Routing-Key"
@@ -1709,6 +1710,11 @@ def run_benchmark(args_: argparse.Namespace):
if args.extra_request_body: if args.extra_request_body:
extra_request_body = json.loads(args.extra_request_body) extra_request_body = json.loads(args.extra_request_body)
# Inject bootstrap fields for fake decode benchmarking
if getattr(args, "fake_prefill", False):
extra_request_body["bootstrap_host"] = FAKE_BOOTSTRAP_HOST
extra_request_body["bootstrap_room"] = 0
if args.tokenize_prompt: if args.tokenize_prompt:
assert ( assert (
args.backend == "sglang" args.backend == "sglang"
@@ -2339,6 +2345,14 @@ if __name__ == "__main__":
], ],
help="Underlying workload for the mooncake dataset.", help="Underlying workload for the mooncake dataset.",
) )
parser.add_argument(
"--fake-prefill",
action="store_true",
default=False,
help="Enable fake prefill mode for decode-only benchmarking. "
"Use with a decode server running --disaggregation-transfer-backend fake "
"to benchmark pure decode performance without a real prefill node.",
)
parser.add_argument( parser.add_argument(
"--tag", type=str, default=None, help="The tag to be dumped to output." "--tag", type=str, default=None, help="The tag to be dumped to output."
) )
@@ -19,6 +19,7 @@ from transformers import AutoProcessor, PreTrainedTokenizer
from sglang.benchmark.datasets import get_dataset from sglang.benchmark.datasets import get_dataset
from sglang.benchmark.utils import get_processor, get_tokenizer from sglang.benchmark.utils import get_processor, get_tokenizer
from sglang.profiler import run_profile from sglang.profiler import run_profile
from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST
from sglang.srt.entrypoints.http_server import launch_server from sglang.srt.entrypoints.http_server import launch_server
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import is_blackwell, kill_process_tree from sglang.srt.utils import is_blackwell, kill_process_tree
@@ -113,6 +114,7 @@ class BenchArgs:
seed: int = 42 seed: int = 42
cache_hit_rate: float = 0.0 cache_hit_rate: float = 0.0
backend: str = "sglang" backend: str = "sglang"
fake_prefill: bool = False
server_args_for_metrics: Optional[List[str]] = None server_args_for_metrics: Optional[List[str]] = None
@staticmethod @staticmethod
@@ -242,6 +244,14 @@ class BenchArgs:
choices=["sglang", "vllm"], choices=["sglang", "vllm"],
help="Backend server type (sglang or vllm).", help="Backend server type (sglang or vllm).",
) )
parser.add_argument(
"--fake-prefill",
action="store_true",
default=BenchArgs.fake_prefill,
help="Enable fake prefill mode for decode-only benchmarking. "
"Use with a decode server running --disaggregation-transfer-backend fake "
"to benchmark pure decode performance without a real prefill node.",
)
parser.add_argument( parser.add_argument(
"--server-args-for-metrics", "--server-args-for-metrics",
type=str, type=str,
@@ -426,6 +436,7 @@ def run_one_case(
gsp_system_prompt_len: int = BenchArgs.gsp_system_prompt_len, gsp_system_prompt_len: int = BenchArgs.gsp_system_prompt_len,
gsp_question_len: int = BenchArgs.gsp_question_len, gsp_question_len: int = BenchArgs.gsp_question_len,
gsp_output_len: int = BenchArgs.gsp_output_len, gsp_output_len: int = BenchArgs.gsp_output_len,
fake_prefill: bool = False,
): ):
if backend == "vllm": if backend == "vllm":
# You need to have export VLLM_SERVER_DEV_MODE=1 in your environment to use this endpoint. # You need to have export VLLM_SERVER_DEV_MODE=1 in your environment to use this endpoint.
@@ -520,6 +531,9 @@ def run_one_case(
payload["input_ids"] = input_ids payload["input_ids"] = input_ids
if image_data is not None: if image_data is not None:
payload["image_data"] = image_data payload["image_data"] = image_data
if fake_prefill:
payload["bootstrap_host"] = FAKE_BOOTSTRAP_HOST
payload["bootstrap_room"] = 0
gen_url = url + "/generate" gen_url = url + "/generate"
# Warm up cache if cache_hit_rate > 0.0 # Warm up cache if cache_hit_rate > 0.0
@@ -864,6 +878,7 @@ def run_benchmark_internal(
parallel_batch=bench_args.parallel_batch, parallel_batch=bench_args.parallel_batch,
backend=bench_args.backend, backend=bench_args.backend,
model_name=model_name, model_name=model_name,
fake_prefill=bench_args.fake_prefill,
**gsp_kwargs, **gsp_kwargs,
) )
print("=" * 8 + " Warmup End " + "=" * 8 + "\n") print("=" * 8 + " Warmup End " + "=" * 8 + "\n")
@@ -900,6 +915,7 @@ def run_benchmark_internal(
cache_hit_rate=bench_args.cache_hit_rate, cache_hit_rate=bench_args.cache_hit_rate,
backend=bench_args.backend, backend=bench_args.backend,
model_name=model_name, model_name=model_name,
fake_prefill=bench_args.fake_prefill,
**gsp_kwargs, **gsp_kwargs,
) )
) )
@@ -945,6 +961,7 @@ def run_benchmark_internal(
profile_output_dir=bench_args.profile_output_dir, profile_output_dir=bench_args.profile_output_dir,
backend=bench_args.backend, backend=bench_args.backend,
model_name=model_name, model_name=model_name,
fake_prefill=bench_args.fake_prefill,
**gsp_kwargs, **gsp_kwargs,
) )
) )