[Rust][Benchmark] Support HF ShareGPT datasets in multi-turn mode (#51104)
Signed-off-by: xiaguan <751080330@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
JinYan Su committed
Sep 15, 2026 at 07:48 UTC
f0a61bd4386bf0d694cfe848e853f2450462e454
3 files changed
+56
-9
rust/src/bench/README.md
+1
@@ -667,6 +667,7 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n
667
668
- `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`.
669
- `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns.
670
+- `--dataset-name hf` — downloads a ShareGPT-format HuggingFace config, then loads it like `sharegpt`.
671
672
**Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0.
673
rust/src/bench/src/config.rs
+32
-6
@@ -549,13 +549,15 @@ impl BenchConfig {
549
));
550
}
551
552
- // Normalize and validate min/max turns. ShareGPT only consumes max_turns
553
- // (the loader walks all available turns up to the cap), so the
554
- // min/num/max coupling used for synthetic generation does not apply.
555
- if args.dataset_name == DatasetName::ShareGpt {
552
+ // Normalize and validate min/max turns. Conversation datasets only
553
+ // consume max_turns (the loader walks all available turns up to the
554
+ // cap), so the min/num/max coupling used for synthetic generation
555
+ // does not apply.
556
+ if matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) {
557
if args.multi_turn_max_turns == 1 {
558
return Err(BenchError::Config(
558
- "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
559
+ "--multi-turn-max-turns must be at least 2 for ShareGPT-format multi-turn"
560
+ .into(),
561
));
562
}
563
} else {
@@ -722,7 +724,7 @@ impl BenchConfig {
724
multi_turn_min_turns,
725
multi_turn_max_turns,
726
sharegpt_multi_turn_max_turns: if args.multi_turn
725
- && args.dataset_name == DatasetName::ShareGpt
727
+ && matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf)
728
&& args.multi_turn_max_turns != 0
729
{
730
Some(args.multi_turn_max_turns)
@@ -1067,6 +1069,30 @@ mod tests {
1069
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
1070
}
1071
1072
+ #[test]
1073
+ fn test_hf_multi_turn_uses_conversation_turn_cap() {
1074
+ let args = vec![
1075
+ "vllm-bench",
1076
+ "--backend",
1077
+ "openai-chat",
1078
+ "--multi-turn",
1079
+ "--model",
1080
+ "test-model",
1081
+ "--dataset-name",
1082
+ "hf",
1083
+ "--dataset-path",
1084
+ "org/sharegpt-dataset",
1085
+ "--hf-subset",
1086
+ "sharegpt",
1087
+ "--multi-turn-max-turns",
1088
+ "2",
1089
+ ];
1090
+ let args = parse_args(args);
1091
+ let config = BenchConfig::from_args(&args).unwrap();
1092
+
1093
+ assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
1094
+ }
1095
+
1096
#[test]
1097
fn test_sweep_summary_percentiles_default_empty() {
1098
let args = base_multi_turn_args();
rust/src/bench/src/multi_turn.rs
+23
-3
@@ -186,9 +186,29 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
186
));
187
}
188
DatasetName::Hf => {
189
- return Err(BenchError::Config(
190
- "HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
191
- ));
189
+ let tok = tokenizer
190
+ .as_ref()
191
+ .ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?;
192
+ let dataset_id = config.dataset_path.as_deref().ok_or_else(|| {
193
+ BenchError::Config("--dataset-path is required for --dataset-name hf".into())
194
+ })?;
195
+ let (downloaded_path, _config, _split) =
196
+ crate::datasets::hf_dataset::download_hf_dataset(
197
+ dataset_id,
198
+ config.hf_subset.as_deref(),
199
+ config.hf_split.as_deref(),
200
+ config.num_prompts,
201
+ )
202
+ .await?;
203
+ crate::datasets::multi_turn::load_sharegpt_multi_turn(
204
+ tok,
205
+ &downloaded_path,
206
+ config.num_prompts,
207
+ config.hf_output_len,
208
+ config.sharegpt_multi_turn_max_turns,
209
+ config.seed,
210
+ &config.request_id_prefix,
211
+ )?
212
}
213
};
214