Add generated-shared-prefix dataset in bench_one_batch

[Qiaolin-Yu]

Motivation

Also refactor to reuse the dataset from bench_serving

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @Qiaolin-Yu, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the benchmarking capabilities of the bench_one_batch_server_internal.py script by introducing a new dataset type: "generated-shared-prefix". This dataset is designed to simulate scenarios where multiple requests share a common prompt prefix, which is particularly useful for evaluating the efficiency of prefix caching mechanisms in LLM serving systems. The changes include new configuration options and a dedicated request generation function, enabling more targeted and realistic performance testing.

Highlights

• New Dataset Type: Added a new "generated-shared-prefix" dataset option to bench_one_batch_server_internal.py for specialized benchmarking scenarios.
• New Configuration Arguments: Introduced new command-line arguments (--gsp-num-groups, --gsp-system-prompt-len, --gsp-question-len, --gsp-output-len) to configure the parameters of the generated shared prefix dataset.
• Shared Prefix Request Generation: Implemented a new function, sample_generated_shared_prefix_requests, to create synthetic requests where multiple prompts share a common system prompt prefix but have unique questions.
• Benchmarking Integration: Integrated the new dataset type and its configuration into the run_one_case and run_benchmark_internal functions, enabling the use of this dataset for performance testing.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/sglang/test/bench_one_batch_server_internal.py
  • Imported the gen_prompt utility function.
  • Extended the BenchArgs class with new fields (gsp_num_groups, gsp_system_prompt_len, gsp_question_len, gsp_output_len) to configure the generated shared prefix dataset.
  • Updated the add_cli_args function to expose these new configuration options as command-line arguments.
  • Added the sample_generated_shared_prefix_requests function to generate requests with shared system prompts and unique questions.
  • Modified the run_one_case function to conditionally use the new sample_generated_shared_prefix_requests function when the "generated-shared-prefix" dataset is selected, and adjusted how mmmu dataset requests are processed.
  • Updated the run_benchmark_internal function to pass the new gsp_kwargs to run_one_case during both warmup and actual benchmark runs.

Activity

• No human activity (comments, reviews, etc.) has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new generated-shared-prefix dataset option to the bench_one_batch benchmark, which is a valuable addition for testing performance scenarios with shared prefixes. The implementation is well-structured. I've identified one correctness issue regarding special token handling during tokenization and a couple of opportunities to improve the efficiency of the test script by using batch tokenization. My review includes specific suggestions to address these points.

[hnyls2002]

/tag-and-rerun-ci