[PD]feat(bench): add --fake-prefill flag for decode-only stress testing (#22973)
This commit is contained in:
@@ -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.
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user