diff --git a/python/sglang/test/send_one.py b/python/sglang/test/send_one.py index 260019abc..1f307b2b4 100644 --- a/python/sglang/test/send_one.py +++ b/python/sglang/test/send_one.py @@ -11,6 +11,7 @@ python3 -m sglang.test.send_one --stop "<|separator|>" "<|eos|>" --max-new-token import argparse import dataclasses import json +import random from typing import Optional import requests @@ -25,6 +26,8 @@ class BenchArgs: port: int = 30000 batch_size: int = 1 different_prompts: bool = False + random_input_len: Optional[int] = None + random_input_vocab_size: int = 32768 seed: Optional[int] = None temperature: float = 0.0 max_new_tokens: int = 512 @@ -54,6 +57,21 @@ class BenchArgs: action="store_true", default=BenchArgs.different_prompts, ) + parser.add_argument( + "--random-input-len", + type=int, + default=BenchArgs.random_input_len, + help="Generate a random prompt of exactly this many tokens (random token IDs). " + "Each request in the batch gets unique random IDs, avoiding radix cache hits. " + "Useful for profiling to ensure the full prefill is captured.", + ) + parser.add_argument( + "--random-input-vocab-size", + type=int, + default=BenchArgs.random_input_vocab_size, + help="Vocab size for --random-input-len. Token IDs are sampled from " + "[0, vocab_size). Default: 32768.", + ) parser.add_argument("--seed", type=int, default=BenchArgs.seed) parser.add_argument("--temperature", type=float, default=BenchArgs.temperature) parser.add_argument( @@ -91,7 +109,35 @@ def send_one_prompt(args: BenchArgs): base_url = f"http://{args.host}:{args.port}" # Construct the input + if args.random_input_len is not None: + # Generate random input ids within the vocab size + n = args.random_input_len + v = args.random_input_vocab_size + if args.batch_size == 1: + input_ids = random.choices(range(v), k=n) + else: + if args.different_prompts: + input_ids = [ + random.choices(range(v), k=n) for _ in range(args.batch_size) + ] + else: + input_ids = [random.choices(range(v), k=n)] * args.batch_size + else: + # Use the user inputs + input_ids = None + if args.batch_size == 1: + prompt = args.prompt + else: + if args.different_prompts: + prompt = [ + f"Test case {i+1}: " + args.prompt for i in range(args.batch_size) + ] + else: + prompt = [args.prompt] * args.batch_size + + # If need image if args.image: + assert args.batch_size == 1 and not args.random_input_len args.prompt = ( "Human: Describe this image in a very short sentence.\n\nAssistant:" ) @@ -110,9 +156,9 @@ def send_one_prompt(args: BenchArgs): else: image_data = None - prompt = args.prompt - + # If need json output if args.json: + assert args.batch_size == 1 and not args.random_input_len prompt = ( "Human: What is the capital of France and how is that city like. " "Give me 3 trivial information about that city. " @@ -122,14 +168,8 @@ def send_one_prompt(args: BenchArgs): else: json_schema = None - if args.batch_size > 1: - if not args.different_prompts: - prompt = [prompt] * args.batch_size - else: - prompt = [f"Test case {i+1}: " + prompt for i in range(args.batch_size)] - json_data = { - "text": prompt, + **({"input_ids": input_ids} if input_ids is not None else {"text": prompt}), "image_data": image_data, "sampling_params": { "sampling_seed": args.seed,