From 82c2efc4b076875c3fb83558639d61a2f282e254 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 6 Aug 2026 02:31:12 +0000 Subject: [PATCH] [Benchmark] Add DFlash vs baseline accuracy/perf sweep scripts Add five standalone benchmark drivers under benchmark/dflash/ that launch servers across (attention_backend, tp_size) configs and run a workload for each (concurrency, num_questions) setting, reporting DFlash throughput and acceptance length against a baseline run. - bench_dflash_gsm8k_sweep.py GSM8K - bench_dflash_math500_sweep.py MATH500, scored by \boxed{...} string equivalence (conservative lower bound) - bench_dflash_humaneval_sweep.py HumanEval pass@1, sandboxed subprocess - bench_dflash_mbpp_sweep.py MBPP pass@1, sandboxed subprocess - bench_dflash_mtbench_sweep.py MT-Bench, performance only (no judge) For the code benchmarks, test execution runs after the timed generation region so it never contaminates throughput or accept-length numbers. Launch environment matches launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh and is applied via setdefault before importing torch/sglang, so explicit caller exports still win. These are benchmark scripts, not CI tests -- they are long-running by design. --- benchmark/dflash/bench_dflash_gsm8k_sweep.py | 867 ++++++++++++++ .../dflash/bench_dflash_humaneval_sweep.py | 1051 +++++++++++++++++ .../dflash/bench_dflash_math500_sweep.py | 1002 ++++++++++++++++ benchmark/dflash/bench_dflash_mbpp_sweep.py | 1036 ++++++++++++++++ .../dflash/bench_dflash_mtbench_sweep.py | 897 ++++++++++++++ 5 files changed, 4853 insertions(+) create mode 100644 benchmark/dflash/bench_dflash_gsm8k_sweep.py create mode 100644 benchmark/dflash/bench_dflash_humaneval_sweep.py create mode 100644 benchmark/dflash/bench_dflash_math500_sweep.py create mode 100644 benchmark/dflash/bench_dflash_mbpp_sweep.py create mode 100644 benchmark/dflash/bench_dflash_mtbench_sweep.py diff --git a/benchmark/dflash/bench_dflash_gsm8k_sweep.py b/benchmark/dflash/bench_dflash_gsm8k_sweep.py new file mode 100644 index 000000000000..306821ca3f81 --- /dev/null +++ b/benchmark/dflash/bench_dflash_gsm8k_sweep.py @@ -0,0 +1,867 @@ +"""DFLASH vs baseline GSM8K sweep. + +This is a *benchmark script* (not a CI test): it can take a long time because it +launches servers for multiple (attention_backend, tp_size) configs and runs a +GSM8K workload for each (concurrency, num_questions) setting. + +Example usage: + ./venv/bin/python benchmark/dflash/bench_dflash_gsm8k_sweep.py + ./venv/bin/python benchmark/dflash/bench_dflash_gsm8k_sweep.py --skip-baseline --concurrencies 32 --tp-sizes 8 +""" + +from __future__ import annotations + +import argparse +import ast +import os +import re +import statistics +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Optional + +# Environment matching launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh. +# Applied via setdefault (caller's explicit exports win) BEFORE importing torch +# and sglang: e.g. TVM_FFI_DISABLE_TORCH_C_DLPACK avoids a libtorch_cuda.so load +# failure on this ROCm build during the sglang import chain. +_LAUNCH_ENV = { + "SGLANG_DISABLE_CUDNN_CHECK": "1", + "SGLANG_USE_CUDA_IPC_TRANSPORT": "1", + "SGLANG_VLM_CACHE_SIZE_MB": "8192", + "SGLANG_USE_AITER": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_SHUFFLE": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_FP8HIPB": "1", + "USE_AITER_COMM": "1", + "AITER_MOE_SMALL_BATCH": "1", + "SGLANG_USE_AITER_NEW_CA": "false", + "SGLANG_USE_IPC_POOL_HANDLE_CACHE": "1", + "AITER_MOE_PADDING_SIZE": "192", + "HIP_GDN_SORT_IDX_BS": "32768", + "TVM_FFI_DISABLE_TORCH_C_DLPACK": "1", +} + +for _k, _v in _LAUNCH_ENV.items(): + os.environ.setdefault(_k, _v) + +import requests # noqa: E402 +import torch # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 + +from sglang.srt.utils import get_device_sm, kill_process_tree # noqa: E402 +from sglang.test.test_utils import ( # noqa: E402 + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + find_available_port, + popen_launch_server, +) +from sglang.utils import download_and_cache_file, read_jsonl # noqa: E402 + +INVALID = -9999999 + + +def _parse_int_csv(value: str) -> list[int]: + return [int(x) for x in value.split(",") if x.strip()] + + +def _filter_attention_backends(backends: list[str], *, device_sm: int) -> list[str]: + if not (80 <= device_sm <= 90): + backends = [b for b in backends if b != "fa3"] + if device_sm < 100: + backends = [b for b in backends if b not in ("fa4", "trtllm_mha")] + return backends or ["flashinfer"] + + +def _get_answer_value(answer_str: str) -> int: + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def _maybe_download_gsm8k(data_path: str) -> str: + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + if os.path.isfile(data_path): + return data_path + return download_and_cache_file(url) + + +def _flush_cache(base_url: str) -> None: + resp = requests.get(base_url + "/flush_cache", timeout=60) + resp.raise_for_status() + + +def _send_generate( + base_url: str, + text: str | list[str], + *, + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + timeout_s: int, +) -> list[dict]: + if isinstance(text, list) and not text: + return [] + sampling_params: dict = { + "temperature": float(temperature), + "top_p": float(top_p), + "top_k": int(top_k), + "max_new_tokens": int(max_new_tokens), + } + resp = requests.post( + base_url + "/generate", + json={ + "text": text, + "sampling_params": sampling_params, + }, + timeout=int(timeout_s), + ) + resp.raise_for_status() + out = resp.json() + if isinstance(text, list): + if not isinstance(out, list): + raise RuntimeError( + "Expected a list response for batched /generate, but got " + f"type={type(out).__name__}." + ) + if len(out) != len(text): + raise RuntimeError( + "Batched /generate output length mismatch: " + f"got {len(out)} outputs for {len(text)} prompts." + ) + return out + + if isinstance(out, list): + raise RuntimeError( + "Expected an object response for single /generate, but got " + f"type={type(out).__name__}." + ) + return [out] + + +@dataclass(frozen=True) +class BenchMetrics: + latency_s: float + output_tokens: int + output_toks_per_s: float + accuracy: Optional[float] + invalid_rate: Optional[float] + spec_accept_length: Optional[float] + spec_accept_length_weighted: Optional[float] + spec_verify_ct_sum: int + + +def _run_gsm8k_requests( + base_url: str, + *, + prompts: list[str], + labels: Optional[list[int]], + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + concurrency: int, + batch_requests: bool, + timeout_s: int, + expect_dflash: bool, +) -> BenchMetrics: + if labels is not None and len(labels) != len(prompts): + raise ValueError("labels length must match prompts length") + + # Drop the first batch from metrics to exclude one-time JIT/cuda-graph overhead + # that often happens immediately after /flush_cache for large batch sizes. + bs = max(int(concurrency), 1) + if len(prompts) > bs: + warmup_prompts = prompts[:bs] + if batch_requests: + _send_generate( + base_url, + warmup_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = [ + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for prompt in warmup_prompts + ] + for fut in as_completed(futures): + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate warmup request." + ) + + prompts = prompts[bs:] + labels = labels[bs:] if labels is not None else None + + start = time.perf_counter() + total_tokens = 0 + spec_verify_ct_sum = 0 + spec_accept_lengths: list[float] = [] + correct = 0 + invalid = 0 + + def _handle_output(out: dict, label: Optional[int]) -> None: + nonlocal total_tokens, spec_verify_ct_sum, correct, invalid + meta = out.get("meta_info", {}) or {} + total_tokens += int(meta.get("completion_tokens", 0)) + spec_verify_ct_sum += int(meta.get("spec_verify_ct", 0)) + if "spec_accept_length" in meta: + try: + spec_accept_lengths.append(float(meta["spec_accept_length"])) + except (TypeError, ValueError): + pass + + if label is not None: + pred = _get_answer_value(out.get("text", "")) + if pred == INVALID: + invalid += 1 + if pred == label: + correct += 1 + + if batch_requests: + bs = max(int(concurrency), 1) + for start_idx in range(0, len(prompts), bs): + chunk_prompts = prompts[start_idx : start_idx + bs] + chunk_labels = ( + labels[start_idx : start_idx + bs] if labels is not None else None + ) + outs = _send_generate( + base_url, + chunk_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + if chunk_labels is None: + for out in outs: + _handle_output(out, None) + else: + for out, label in zip(outs, chunk_labels): + _handle_output(out, label) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = { + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ): i + for i, prompt in enumerate(prompts) + } + for fut in as_completed(futures): + i = futures[fut] + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate request." + ) + label = None if labels is None else labels[i] + _handle_output(outs[0], label) + + latency = time.perf_counter() - start + toks_per_s = total_tokens / max(latency, 1e-6) + + if expect_dflash and spec_verify_ct_sum <= 0: + raise RuntimeError( + "DFLASH sanity check failed: did not observe any `spec_verify_ct` in responses " + "(DFLASH may not have been enabled)." + ) + + spec_accept_length = ( + float(statistics.mean(spec_accept_lengths)) if spec_accept_lengths else None + ) + # Token-weighted global accept length: total completion tokens / total verify + # steps (each request weighted by its length, unlike the per-request mean above). + spec_accept_length_weighted = ( + float(total_tokens) / float(spec_verify_ct_sum) + if spec_verify_ct_sum > 0 + else None + ) + + if labels is None: + acc = None + invalid_rate = None + else: + acc = correct / max(len(prompts), 1) + invalid_rate = invalid / max(len(prompts), 1) + + return BenchMetrics( + latency_s=float(latency), + output_tokens=int(total_tokens), + output_toks_per_s=float(toks_per_s), + accuracy=acc, + invalid_rate=invalid_rate, + spec_accept_length=spec_accept_length, + spec_accept_length_weighted=spec_accept_length_weighted, + spec_verify_ct_sum=int(spec_verify_ct_sum), + ) + + +def _format_table( + *, + tp_sizes: list[int], + concurrencies: list[int], + values: dict[tuple[int, int], Optional[float]], + float_fmt: str, +) -> str: + header = ["tp\\conc"] + [str(c) for c in concurrencies] + rows: list[list[str]] = [header] + for tp in tp_sizes: + row = [str(tp)] + for c in concurrencies: + v = values.get((tp, c), None) + row.append("N/A" if v is None else format(v, float_fmt)) + rows.append(row) + + col_widths = [ + max(len(row[col_idx]) for row in rows) for col_idx in range(len(rows[0])) + ] + + lines: list[str] = [] + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(rows[0]))) + lines.append(" ".join("-" * w for w in col_widths)) + for row in rows[1:]: + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(row))) + return "\n".join(lines) + + +def _build_common_server_args( + args: argparse.Namespace, *, backend: str, tp: int +) -> list[str]: + # Mirror launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh (baseline flags). + common_server_args: list[str] = [ + "--tp-size", + str(tp), + "--reasoning-parser", + "qwen3", + "--tool-call-parser", + "qwen3_coder", + "--enable-multimodal", + "--trust-remote-code", + "--chunked-prefill-size", + "65536", + "--mem-fraction-static", + str(args.mem_fraction_static if args.mem_fraction_static is not None else 0.9), + "--max-prefill-tokens", + "65536", + "--max-running-requests", + str(args.max_running_requests), + "--attention-backend", + backend, + "--mm-attention-backend", + "aiter_attn", + "--mamba-scheduler-strategy", + str(args.mamba_scheduler_strategy), + "--disable-custom-all-reduce", + "--kv-cache-dtype", + "fp8_e4m3", + "--page-size", + str(int(args.page_size) if args.page_size is not None else 64), + ] + if args.disable_radix_cache: + common_server_args.append("--disable-radix-cache") + return common_server_args + + +def _build_mode_runs( + args: argparse.Namespace, common_server_args: list[str] +) -> list[tuple[str, str, list[str], bool]]: + mode_runs: list[tuple[str, str, list[str], bool]] = [] + if not args.skip_baseline: + mode_runs.append(("baseline", "baseline", common_server_args, False)) + mode_runs.append( + ( + "dflash", + "DFLASH", + [ + *common_server_args, + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + args.draft_model, + "--speculative-num-draft-tokens", + str(int(args.speculative_num_draft_tokens)), + *( + [ + "--speculative-dflash-draft-window-size", + str(int(args.speculative_dflash_draft_window_size)), + ] + if args.speculative_dflash_draft_window_size is not None + else [] + ), + *( + [ + "--speculative-draft-attention-backend", + args.speculative_draft_attention_backend, + ] + if args.speculative_draft_attention_backend + else [] + ), + ], + True, + ) + ) + return mode_runs + + +def _collect_metric( + *, + results: dict[tuple[str, int, int, str], BenchMetrics], + backend: str, + tp_sizes: list[int], + concurrencies: list[int], + mode: str, + field: str, +) -> dict[tuple[int, int], Optional[float]]: + out: dict[tuple[int, int], Optional[float]] = {} + for tp in tp_sizes: + for conc in concurrencies: + metrics = results.get((backend, tp, conc, mode), None) + out[(tp, conc)] = None if metrics is None else getattr(metrics, field) + return out + + +def _compute_speedup( + baseline: dict[tuple[int, int], Optional[float]], + dflash: dict[tuple[int, int], Optional[float]], +) -> dict[tuple[int, int], Optional[float]]: + return { + key: None if (b is None or d is None or b <= 0) else (d / b) + for key, b in baseline.items() + for d in [dflash.get(key, None)] + } + + +def _print_kv_lines(items: list[tuple[str, object]]) -> None: + for key, value in items: + print(f"{key}={value}") + + +def _run_mode_for_backend_tp( + *, + mode_label: str, + model_path: str, + base_url: str, + server_args: list[str], + expect_dflash: bool, + prompts: list[str], + labels: list[int], + concurrencies: list[int], + num_questions_by_conc: dict[int, int], + args: argparse.Namespace, +) -> dict[int, BenchMetrics]: + print(f"\n=== {mode_label} ===") + server_start_timeout_s = int(max(DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, args.timeout_s)) + proc = popen_launch_server( + model_path, + base_url, + timeout=server_start_timeout_s, + other_args=server_args, + ) + try: + _send_generate( + base_url, + "Hello", + max_new_tokens=8, + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + timeout_s=min(int(args.timeout_s), 300), + ) + + metrics_by_conc: dict[int, BenchMetrics] = {} + for conc in concurrencies: + n = num_questions_by_conc[conc] + _flush_cache(base_url) + print( + f"[warmup] run 1 warmup batch (size={conc}) after /flush_cache; excluded from metrics." + ) + metrics = _run_gsm8k_requests( + base_url, + prompts=prompts[: n + conc], + labels=labels[: n + conc], + max_new_tokens=int(args.max_new_tokens), + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + concurrency=int(conc), + batch_requests=bool(args.batch_requests), + timeout_s=int(args.timeout_s), + expect_dflash=expect_dflash, + ) + metrics_by_conc[conc] = metrics + line = ( + f"[{mode_label}] conc={conc:>2} n={n:<4} " + f"toks/s={metrics.output_toks_per_s:,.2f} " + f"latency={metrics.latency_s:.1f}s " + f"acc={metrics.accuracy:.3f} invalid={metrics.invalid_rate:.3f}" + ) + if expect_dflash: + accept_len = ( + "N/A" + if metrics.spec_accept_length is None + else f"{metrics.spec_accept_length:.3f}" + ) + accept_len_w = ( + "N/A" + if metrics.spec_accept_length_weighted is None + else f"{metrics.spec_accept_length_weighted:.3f}" + ) + line += ( + f" accept_len={accept_len} " + f"accept_len_weighted={accept_len_w} " + f"spec_verify_ct_sum={metrics.spec_verify_ct_sum}" + ) + print(line) + return metrics_by_conc + finally: + kill_process_tree(proc.pid) + try: + proc.wait(timeout=30) + except Exception: + pass + + +def _print_summary( + *, + args: argparse.Namespace, + attention_backends: list[str], + tp_sizes: list[int], + concurrencies: list[int], + device_sm: int, + results: dict[tuple[str, int, int, str], BenchMetrics], +) -> None: + print("\n=== DFLASH GSM8K Sweep Summary ===") + _print_kv_lines( + [ + ("target_model", args.target_model), + ("draft_model", args.draft_model), + ("max_new_tokens", args.max_new_tokens), + ( + "sampling", + f"temperature:{args.temperature}, top_p:{args.top_p}, top_k:{args.top_k}", + ), + ("attention_backends", ",".join(attention_backends)), + ( + "speculative_draft_attention_backend", + args.speculative_draft_attention_backend, + ), + ( + "speculative_dflash_draft_window_size", + args.speculative_dflash_draft_window_size, + ), + ("tp_sizes", ",".join(str(x) for x in tp_sizes)), + ("concurrencies", ",".join(str(x) for x in concurrencies)), + ( + "questions_per_concurrency_base", + args.questions_per_concurrency_base, + ), + ("device_sm", device_sm), + ("skip_baseline", bool(args.skip_baseline)), + ] + ) + + section_fields = [ + ("Baseline output tok/s", "baseline", "output_toks_per_s", ",.2f"), + ("Baseline accuracy", "baseline", "accuracy", ".3f"), + ("DFLASH output tok/s", "dflash", "output_toks_per_s", ",.2f"), + ("DFLASH accuracy", "dflash", "accuracy", ".3f"), + ( + "DFLASH acceptance length (mean spec_accept_length)", + "dflash", + "spec_accept_length", + ".3f", + ), + ( + "DFLASH acceptance length (token-weighted: total_tokens/verify_ct)", + "dflash", + "spec_accept_length_weighted", + ".3f", + ), + ] + + for backend in attention_backends: + print(f"\n=== Backend: {backend} ===") + metrics_map = { + (mode, field): _collect_metric( + results=results, + backend=backend, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + mode=mode, + field=field, + ) + for _, mode, field, _ in section_fields + } + sections: list[tuple[str, dict[tuple[int, int], Optional[float]], str]] = [ + (title, metrics_map[(mode, field)], fmt) + for title, mode, field, fmt in section_fields + ] + sections.insert( + 4, + ( + "Speedup (DFLASH / baseline)", + _compute_speedup( + metrics_map[("baseline", "output_toks_per_s")], + metrics_map[("dflash", "output_toks_per_s")], + ), + ".3f", + ), + ) + + for title, values, fmt in sections: + print(f"\n{title}") + print( + _format_table( + tp_sizes=tp_sizes, + concurrencies=concurrencies, + values=values, + float_fmt=fmt, + ) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", default="test.jsonl") + parser.add_argument( + "--target-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-PTPC-FP8/" + ) + parser.add_argument( + "--draft-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-DFlash/" + ) + parser.add_argument( + "--skip-baseline", + action="store_true", + help="Skip running the baseline (target-only) sweep; only run DFLASH and report N/A for baseline/speedup.", + ) + parser.add_argument( + "--batch-requests", + action="store_true", + help="Send prompts as server-side batched /generate requests (batch size = concurrency) instead of client-side concurrent requests.", + ) + parser.add_argument("--max-new-tokens", type=int, default=2048) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=1) + parser.add_argument( + "--timeout-s", + type=int, + default=3600, + help=( + "Timeout in seconds for benchmarked /generate calls and server startup " + "health checks." + ), + ) + parser.add_argument( + "--mem-fraction-static", + type=float, + default=None, + help="Optional server --mem-fraction-static override. If unset, use the server auto heuristic.", + ) + parser.add_argument("--disable-radix-cache", action="store_true") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument( + "--speculative-num-draft-tokens", + type=int, + default=16, + help="DFLASH verify window length (matches the 397B launch script).", + ) + parser.add_argument( + "--page-size", + type=int, + default=None, + help="Optional server --page-size override for both baseline and DFLASH runs.", + ) + parser.add_argument("--max-running-requests", type=int, default=32) + parser.add_argument( + "--mamba-scheduler-strategy", + default="extra_buffer", + help=( + "Server --mamba-scheduler-strategy value to pass through to benchmark " + "runs, e.g. `no_buffer` or `extra_buffer`." + ), + ) + parser.add_argument("--tp-sizes", default="8") + parser.add_argument("--concurrencies", default="1") + parser.add_argument( + "--questions-per-concurrency-base", + type=int, + default=128, + help="num_questions = base * concurrency (default matches the sweep plan).", + ) + parser.add_argument( + "--max-questions-per-config", + type=int, + default=1024, + help="Cap num_questions per (tp, concurrency) run (default: 1024).", + ) + parser.add_argument("--attention-backends", default="aiter") + parser.add_argument( + "--speculative-draft-attention-backend", + default="triton", + help="Optional server --speculative-draft-attention-backend override for DFLASH runs.", + ) + parser.add_argument( + "--speculative-dflash-draft-window-size", + type=int, + default=None, + help="Optional server --speculative-dflash-draft-window-size override for DFLASH runs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this sweep.") + if args.temperature < 0.0: + raise RuntimeError(f"--temperature must be >= 0, got {args.temperature}.") + if not (0.0 < args.top_p <= 1.0): + raise RuntimeError(f"--top-p must be in (0, 1], got {args.top_p}.") + if args.top_k == 0 or args.top_k < -1: + raise RuntimeError(f"--top-k must be -1 (all vocab) or >= 1, got {args.top_k}.") + if args.timeout_s <= 0: + raise RuntimeError(f"--timeout-s must be > 0, got {args.timeout_s}.") + + visible_gpus = int(torch.cuda.device_count()) + tp_sizes = _parse_int_csv(args.tp_sizes) + tp_sizes = [tp for tp in tp_sizes if tp >= 1 and tp <= visible_gpus] + if not tp_sizes: + raise RuntimeError( + f"No tp sizes are runnable with visible_gpus={visible_gpus}. " + "Set CUDA_VISIBLE_DEVICES accordingly." + ) + + concurrencies = _parse_int_csv(args.concurrencies) + concurrencies = [c for c in concurrencies if c >= 1] + if not concurrencies: + raise RuntimeError("No concurrencies specified.") + + num_questions_by_conc = { + c: min( + int(args.questions_per_concurrency_base) * int(c), + int(args.max_questions_per_config), + ) + for c in concurrencies + } + max_questions = max(num_questions_by_conc.values()) + + attention_backends = [ + s.strip() for s in args.attention_backends.split(",") if s.strip() + ] + device_sm = get_device_sm() + attention_backends = _filter_attention_backends( + attention_backends, device_sm=device_sm + ) + + data_path = _maybe_download_gsm8k(args.data_path) + lines = list(read_jsonl(data_path)) + if len(lines) < max_questions: + raise RuntimeError( + f"GSM8K file only has {len(lines)} lines, but need {max_questions}." + ) + + tokenizer = AutoTokenizer.from_pretrained(args.target_model) + + prompts: list[str] = [] + labels: list[int] = [] + for i in range(max_questions): + user_content = ( + lines[i]["question"] + + "\nPlease reason step by step, and put your final answer within \\boxed{}." + ) + prompts.append( + tokenizer.apply_chat_template( + [{"role": "user", "content": user_content}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + ) + labels.append(_get_answer_value(lines[i]["answer"])) + if not all(label != INVALID for label in labels): + raise RuntimeError("Invalid labels in GSM8K data.") + + # Results indexed by (backend, tp, concurrency, mode). + results: dict[tuple[str, int, int, str], BenchMetrics] = {} + # Baseline metrics are backend-agnostic in this sweep; run once per TP and reuse. + baseline_cache_by_tp: dict[int, dict[int, BenchMetrics]] = {} + + for backend_idx, backend in enumerate(attention_backends): + for tp in tp_sizes: + port_base = find_available_port(20000) + common_server_args = _build_common_server_args(args, backend=backend, tp=tp) + mode_runs = _build_mode_runs(args, common_server_args) + + for idx, ( + mode_key, + mode_name, + mode_server_args, + expect_dflash, + ) in enumerate(mode_runs): + if ( + mode_key == "baseline" + and not args.skip_baseline + and backend_idx > 0 + and tp in baseline_cache_by_tp + ): + mode_metrics = baseline_cache_by_tp[tp] + else: + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"backend={backend} tp={tp} ({mode_name})", + model_path=args.target_model, + base_url=f"http://127.0.0.1:{find_available_port(port_base + idx)}", + server_args=mode_server_args, + expect_dflash=expect_dflash, + prompts=prompts, + labels=labels, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + ) + if mode_key == "baseline" and not args.skip_baseline: + baseline_cache_by_tp[tp] = mode_metrics + + for conc, metrics in mode_metrics.items(): + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/dflash/bench_dflash_humaneval_sweep.py b/benchmark/dflash/bench_dflash_humaneval_sweep.py new file mode 100644 index 000000000000..7713bdce65cc --- /dev/null +++ b/benchmark/dflash/bench_dflash_humaneval_sweep.py @@ -0,0 +1,1051 @@ +"""DFLASH vs baseline HumanEval sweep. + +This is a *benchmark script* (not a CI test): it can take a long time because it +launches servers for multiple (attention_backend, tp_size) configs and runs a +HumanEval workload for each (concurrency, num_questions) setting. + +Accuracy is pass@1: the model completion is combined with the problem's unit +test and executed in a sandboxed subprocess. Code execution happens AFTER the +timed generation region so it never contaminates throughput / accept-length. + +Security note: this EXECUTES model-generated code in subprocesses. Run only in a +disposable/trusted environment. + +Example usage: + ./venv/bin/python benchmark/dflash/bench_dflash_humaneval_sweep.py + ./venv/bin/python benchmark/dflash/bench_dflash_humaneval_sweep.py --skip-baseline --concurrencies 32 --tp-sizes 8 +""" + +from __future__ import annotations + +import argparse +import gzip +import multiprocessing as mp +import os +import re +import shutil +import statistics +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Optional + +# Environment matching launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh. +# Applied via setdefault (caller's explicit exports win) BEFORE importing torch +# and sglang: e.g. TVM_FFI_DISABLE_TORCH_C_DLPACK avoids a libtorch_cuda.so load +# failure on this ROCm build during the sglang import chain. +_LAUNCH_ENV = { + "SGLANG_DISABLE_CUDNN_CHECK": "1", + "SGLANG_USE_CUDA_IPC_TRANSPORT": "1", + "SGLANG_VLM_CACHE_SIZE_MB": "8192", + "SGLANG_USE_AITER": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_SHUFFLE": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_FP8HIPB": "1", + "USE_AITER_COMM": "1", + "AITER_MOE_SMALL_BATCH": "1", + "SGLANG_USE_AITER_NEW_CA": "false", + "SGLANG_USE_IPC_POOL_HANDLE_CACHE": "1", + "AITER_MOE_PADDING_SIZE": "192", + "HIP_GDN_SORT_IDX_BS": "32768", + "TVM_FFI_DISABLE_TORCH_C_DLPACK": "1", +} + +for _k, _v in _LAUNCH_ENV.items(): + os.environ.setdefault(_k, _v) + +import requests # noqa: E402 +import torch # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 + +from sglang.srt.utils import get_device_sm, kill_process_tree # noqa: E402 +from sglang.test.test_utils import ( # noqa: E402 + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + find_available_port, + popen_launch_server, +) +from sglang.utils import download_and_cache_file, read_jsonl # noqa: E402 + + +def _parse_int_csv(value: str) -> list[int]: + return [int(x) for x in value.split(",") if x.strip()] + + +def _filter_attention_backends(backends: list[str], *, device_sm: int) -> list[str]: + if not (80 <= device_sm <= 90): + backends = [b for b in backends if b != "fa3"] + if device_sm < 100: + backends = [b for b in backends if b not in ("fa4", "trtllm_mha")] + return backends or ["flashinfer"] + + +_CODE_BLOCK_RE = re.compile(r"```(?:python)?\s*\n?(.*?)```", re.DOTALL | re.IGNORECASE) + + +def _extract_code(text: str, entry_point: str) -> str: + """Pull a code body from a model response. + + Prefers the last fenced ```python block; falls back to the raw text. + """ + blocks = _CODE_BLOCK_RE.findall(text) + if blocks: + # Prefer a block that actually defines the target function. + for block in reversed(blocks): + if f"def {entry_point}" in block: + return block + return blocks[-1] + return text + + +def _build_humaneval_program( + *, problem_prompt: str, completion_text: str, test: str, entry_point: str +) -> str: + code = _extract_code(completion_text, entry_point) + # If the model regenerated the signature, use its code as the full module; + # otherwise treat its output as a body to append to the provided stub. + if f"def {entry_point}" in code: + program_body = code + else: + program_body = problem_prompt + "\n" + code + return program_body + "\n" + test + f"\n\ncheck({entry_point})\n" + + +def _exec_worker(program: str, queue) -> None: # pragma: no cover - child process + # Best-effort resource limits; the parent also enforces a wall-clock timeout. + try: + import resource + + resource.setrlimit(resource.RLIMIT_CPU, (10, 10)) + except Exception: + pass + import builtins # noqa: F401 + + g: dict = {"__name__": "__humaneval__"} + try: + exec(program, g) + queue.put(True) + except Exception: + queue.put(False) + + +def _check_program(program: str, timeout_s: float) -> bool: + """Execute `program` in a subprocess; True iff it runs without raising.""" + ctx = mp.get_context("fork") + queue = ctx.Queue() + proc = ctx.Process(target=_exec_worker, args=(program, queue)) + proc.start() + proc.join(timeout_s) + if proc.is_alive(): + proc.terminate() + proc.join(1) + if proc.is_alive(): + proc.kill() + proc.join() + return False + try: + return bool(queue.get_nowait()) + except Exception: + return False + + +def _maybe_download_humaneval(data_path: str) -> str: + url = "https://raw.githubusercontent.com/openai/human-eval/master/data/HumanEval.jsonl.gz" + if os.path.isfile(data_path): + return data_path + cached = download_and_cache_file(url) + if not cached.endswith(".gz"): + return cached + # The upstream file is gzipped but read_jsonl() opens the path as utf-8 text, + # so unpack it once next to the cached download. + unpacked = cached[: -len(".gz")] + if not os.path.isfile(unpacked): + tmp = unpacked + ".tmp" + with gzip.open(cached, "rb") as src, open(tmp, "wb") as dst: + shutil.copyfileobj(src, dst) + os.replace(tmp, unpacked) + return unpacked + + +def _flush_cache(base_url: str) -> None: + resp = requests.get(base_url + "/flush_cache", timeout=60) + resp.raise_for_status() + + +def _send_generate( + base_url: str, + text: str | list[str], + *, + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + timeout_s: int, +) -> list[dict]: + if isinstance(text, list) and not text: + return [] + sampling_params: dict = { + "temperature": float(temperature), + "top_p": float(top_p), + "top_k": int(top_k), + "max_new_tokens": int(max_new_tokens), + } + resp = requests.post( + base_url + "/generate", + json={ + "text": text, + "sampling_params": sampling_params, + }, + timeout=int(timeout_s), + ) + resp.raise_for_status() + out = resp.json() + if isinstance(text, list): + if not isinstance(out, list): + raise RuntimeError( + "Expected a list response for batched /generate, but got " + f"type={type(out).__name__}." + ) + if len(out) != len(text): + raise RuntimeError( + "Batched /generate output length mismatch: " + f"got {len(out)} outputs for {len(text)} prompts." + ) + return out + + if isinstance(out, list): + raise RuntimeError( + "Expected an object response for single /generate, but got " + f"type={type(out).__name__}." + ) + return [out] + + +@dataclass(frozen=True) +class BenchMetrics: + latency_s: float + output_tokens: int + output_toks_per_s: float + accuracy: Optional[float] + invalid_rate: Optional[float] + spec_accept_length: Optional[float] + spec_accept_length_weighted: Optional[float] + spec_verify_ct_sum: int + + +def _run_requests( + base_url: str, + *, + prompts: list[str], + tasks: Optional[list[dict]], + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + concurrency: int, + batch_requests: bool, + timeout_s: int, + expect_dflash: bool, + exec_timeout_s: float, +) -> BenchMetrics: + """`tasks[i]` carries the fields needed to score prompt `i` (or None to skip).""" + if tasks is not None and len(tasks) != len(prompts): + raise ValueError("tasks length must match prompts length") + + # Drop the first batch from metrics to exclude one-time JIT/cuda-graph overhead + # that often happens immediately after /flush_cache for large batch sizes. + bs = max(int(concurrency), 1) + if len(prompts) > bs: + warmup_prompts = prompts[:bs] + if batch_requests: + _send_generate( + base_url, + warmup_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = [ + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for prompt in warmup_prompts + ] + for fut in as_completed(futures): + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate warmup request." + ) + + prompts = prompts[bs:] + tasks = tasks[bs:] if tasks is not None else None + + start = time.perf_counter() + total_tokens = 0 + spec_verify_ct_sum = 0 + spec_accept_lengths: list[float] = [] + # Collect completions during the timed region; score AFTER timing stops. + completions: list[Optional[str]] = [None] * len(prompts) + + def _handle_output(out: dict, index: int) -> None: + nonlocal total_tokens, spec_verify_ct_sum + meta = out.get("meta_info", {}) or {} + total_tokens += int(meta.get("completion_tokens", 0)) + spec_verify_ct_sum += int(meta.get("spec_verify_ct", 0)) + if "spec_accept_length" in meta: + try: + spec_accept_lengths.append(float(meta["spec_accept_length"])) + except (TypeError, ValueError): + pass + completions[index] = out.get("text", "") + + if batch_requests: + bs = max(int(concurrency), 1) + for start_idx in range(0, len(prompts), bs): + chunk_prompts = prompts[start_idx : start_idx + bs] + outs = _send_generate( + base_url, + chunk_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for offset, out in enumerate(outs): + _handle_output(out, start_idx + offset) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = { + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ): i + for i, prompt in enumerate(prompts) + } + for fut in as_completed(futures): + i = futures[fut] + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate request." + ) + _handle_output(outs[0], i) + + latency = time.perf_counter() - start + toks_per_s = total_tokens / max(latency, 1e-6) + + if expect_dflash and spec_verify_ct_sum <= 0: + raise RuntimeError( + "DFLASH sanity check failed: did not observe any `spec_verify_ct` in responses " + "(DFLASH may not have been enabled)." + ) + + spec_accept_length = ( + float(statistics.mean(spec_accept_lengths)) if spec_accept_lengths else None + ) + # Token-weighted global accept length: total completion tokens / total verify + # steps (each request weighted by its length, unlike the per-request mean above). + spec_accept_length_weighted = ( + float(total_tokens) / float(spec_verify_ct_sum) + if spec_verify_ct_sum > 0 + else None + ) + + # --- Post-timing pass@1 scoring via sandboxed code execution --- + if tasks is None: + acc = None + invalid_rate = None + else: + correct = 0 + invalid = 0 + for i, task in enumerate(tasks): + comp = completions[i] + if not comp: + invalid += 1 + continue + program = _build_humaneval_program( + problem_prompt=task["prompt"], + completion_text=comp, + test=task["test"], + entry_point=task["entry_point"], + ) + if f"def {task['entry_point']}" not in program: + invalid += 1 + continue + if _check_program(program, exec_timeout_s): + correct += 1 + acc = correct / max(len(tasks), 1) + invalid_rate = invalid / max(len(tasks), 1) + + return BenchMetrics( + latency_s=float(latency), + output_tokens=int(total_tokens), + output_toks_per_s=float(toks_per_s), + accuracy=acc, + invalid_rate=invalid_rate, + spec_accept_length=spec_accept_length, + spec_accept_length_weighted=spec_accept_length_weighted, + spec_verify_ct_sum=int(spec_verify_ct_sum), + ) + + +def _format_table( + *, + tp_sizes: list[int], + concurrencies: list[int], + values: dict[tuple[int, int], Optional[float]], + float_fmt: str, +) -> str: + header = ["tp\\conc"] + [str(c) for c in concurrencies] + rows: list[list[str]] = [header] + for tp in tp_sizes: + row = [str(tp)] + for c in concurrencies: + v = values.get((tp, c), None) + row.append("N/A" if v is None else format(v, float_fmt)) + rows.append(row) + + col_widths = [ + max(len(row[col_idx]) for row in rows) for col_idx in range(len(rows[0])) + ] + + lines: list[str] = [] + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(rows[0]))) + lines.append(" ".join("-" * w for w in col_widths)) + for row in rows[1:]: + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(row))) + return "\n".join(lines) + + +def _build_common_server_args( + args: argparse.Namespace, *, backend: str, tp: int +) -> list[str]: + # Mirror launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh (baseline flags). + common_server_args: list[str] = [ + "--tp-size", + str(tp), + "--reasoning-parser", + "qwen3", + "--tool-call-parser", + "qwen3_coder", + "--enable-multimodal", + "--trust-remote-code", + "--chunked-prefill-size", + "65536", + "--mem-fraction-static", + str(args.mem_fraction_static if args.mem_fraction_static is not None else 0.9), + "--max-prefill-tokens", + "65536", + "--max-running-requests", + str(args.max_running_requests), + "--attention-backend", + backend, + "--mm-attention-backend", + "aiter_attn", + "--mamba-scheduler-strategy", + str(args.mamba_scheduler_strategy), + "--disable-custom-all-reduce", + "--kv-cache-dtype", + "fp8_e4m3", + "--page-size", + str(int(args.page_size) if args.page_size is not None else 64), + ] + if args.disable_radix_cache: + common_server_args.append("--disable-radix-cache") + return common_server_args + + +def _build_mode_runs( + args: argparse.Namespace, common_server_args: list[str] +) -> list[tuple[str, str, list[str], bool]]: + mode_runs: list[tuple[str, str, list[str], bool]] = [] + if not args.skip_baseline: + mode_runs.append(("baseline", "baseline", common_server_args, False)) + mode_runs.append( + ( + "dflash", + "DFLASH", + [ + *common_server_args, + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + args.draft_model, + "--speculative-num-draft-tokens", + str(int(args.speculative_num_draft_tokens)), + *( + [ + "--speculative-dflash-draft-window-size", + str(int(args.speculative_dflash_draft_window_size)), + ] + if args.speculative_dflash_draft_window_size is not None + else [] + ), + *( + [ + "--speculative-draft-attention-backend", + args.speculative_draft_attention_backend, + ] + if args.speculative_draft_attention_backend + else [] + ), + ], + True, + ) + ) + return mode_runs + + +def _collect_metric( + *, + results: dict[tuple[str, int, int, str], BenchMetrics], + backend: str, + tp_sizes: list[int], + concurrencies: list[int], + mode: str, + field: str, +) -> dict[tuple[int, int], Optional[float]]: + out: dict[tuple[int, int], Optional[float]] = {} + for tp in tp_sizes: + for conc in concurrencies: + metrics = results.get((backend, tp, conc, mode), None) + out[(tp, conc)] = None if metrics is None else getattr(metrics, field) + return out + + +def _compute_speedup( + baseline: dict[tuple[int, int], Optional[float]], + dflash: dict[tuple[int, int], Optional[float]], +) -> dict[tuple[int, int], Optional[float]]: + return { + key: None if (b is None or d is None or b <= 0) else (d / b) + for key, b in baseline.items() + for d in [dflash.get(key, None)] + } + + +def _print_kv_lines(items: list[tuple[str, object]]) -> None: + for key, value in items: + print(f"{key}={value}") + + +def _run_mode_for_backend_tp( + *, + mode_label: str, + model_path: str, + base_url: str, + server_args: list[str], + expect_dflash: bool, + prompts: list[str], + tasks: list[dict], + concurrencies: list[int], + num_questions_by_conc: dict[int, int], + args: argparse.Namespace, + launch_server: bool = True, +) -> dict[int, BenchMetrics]: + print(f"\n=== {mode_label} ===") + proc = None + if launch_server: + server_start_timeout_s = int( + max(DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, args.timeout_s) + ) + proc = popen_launch_server( + model_path, + base_url, + timeout=server_start_timeout_s, + other_args=server_args, + ) + else: + print(f"[external] reusing already-running server at {base_url}.") + try: + _send_generate( + base_url, + "Hello", + max_new_tokens=8, + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + timeout_s=min(int(args.timeout_s), 300), + ) + + metrics_by_conc: dict[int, BenchMetrics] = {} + for conc in concurrencies: + n = num_questions_by_conc[conc] + _flush_cache(base_url) + print( + f"[warmup] run 1 warmup batch (size={conc}) after /flush_cache; excluded from metrics." + ) + metrics = _run_requests( + base_url, + prompts=prompts[: n + conc], + tasks=tasks[: n + conc], + max_new_tokens=int(args.max_new_tokens), + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + concurrency=int(conc), + batch_requests=bool(args.batch_requests), + timeout_s=int(args.timeout_s), + expect_dflash=expect_dflash, + exec_timeout_s=float(args.exec_timeout_s), + ) + metrics_by_conc[conc] = metrics + line = ( + f"[{mode_label}] conc={conc:>2} n={n:<4} " + f"toks/s={metrics.output_toks_per_s:,.2f} " + f"latency={metrics.latency_s:.1f}s " + f"pass@1={metrics.accuracy:.3f} invalid={metrics.invalid_rate:.3f}" + ) + if expect_dflash or metrics.spec_verify_ct_sum > 0: + accept_len = ( + "N/A" + if metrics.spec_accept_length is None + else f"{metrics.spec_accept_length:.3f}" + ) + accept_len_w = ( + "N/A" + if metrics.spec_accept_length_weighted is None + else f"{metrics.spec_accept_length_weighted:.3f}" + ) + line += ( + f" accept_len={accept_len} " + f"accept_len_weighted={accept_len_w} " + f"spec_verify_ct_sum={metrics.spec_verify_ct_sum}" + ) + print(line) + return metrics_by_conc + finally: + if proc is not None: + kill_process_tree(proc.pid) + try: + proc.wait(timeout=30) + except Exception: + pass + + +def _print_summary( + *, + args: argparse.Namespace, + attention_backends: list[str], + tp_sizes: list[int], + concurrencies: list[int], + device_sm: int, + results: dict[tuple[str, int, int, str], BenchMetrics], +) -> None: + print("\n=== DFLASH HumanEval Sweep Summary ===") + if args.base_url: + _print_kv_lines( + [ + ("base_url", args.base_url), + ("external_mode", args.external_mode), + ( + "note", + "server-side flags below were NOT applied; the running server's own " + "config was measured (tp/backend are labels only)", + ), + ] + ) + _print_kv_lines( + [ + ("target_model", args.target_model), + ("draft_model", args.draft_model), + ("max_new_tokens", args.max_new_tokens), + ( + "sampling", + f"temperature:{args.temperature}, top_p:{args.top_p}, top_k:{args.top_k}", + ), + ("attention_backends", ",".join(attention_backends)), + ( + "speculative_draft_attention_backend", + args.speculative_draft_attention_backend, + ), + ( + "speculative_dflash_draft_window_size", + args.speculative_dflash_draft_window_size, + ), + ("tp_sizes", ",".join(str(x) for x in tp_sizes)), + ("concurrencies", ",".join(str(x) for x in concurrencies)), + ( + "questions_per_concurrency_base", + args.questions_per_concurrency_base, + ), + ("device_sm", device_sm), + ("skip_baseline", bool(args.skip_baseline)), + ] + ) + + section_fields = [ + ("Baseline output tok/s", "baseline", "output_toks_per_s", ",.2f"), + ("Baseline pass@1", "baseline", "accuracy", ".3f"), + ("DFLASH output tok/s", "dflash", "output_toks_per_s", ",.2f"), + ("DFLASH pass@1", "dflash", "accuracy", ".3f"), + ( + "DFLASH acceptance length (mean spec_accept_length)", + "dflash", + "spec_accept_length", + ".3f", + ), + ( + "DFLASH acceptance length (token-weighted: total_tokens/verify_ct)", + "dflash", + "spec_accept_length_weighted", + ".3f", + ), + ] + + for backend in attention_backends: + print(f"\n=== Backend: {backend} ===") + metrics_map = { + (mode, field): _collect_metric( + results=results, + backend=backend, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + mode=mode, + field=field, + ) + for _, mode, field, _ in section_fields + } + sections: list[tuple[str, dict[tuple[int, int], Optional[float]], str]] = [ + (title, metrics_map[(mode, field)], fmt) + for title, mode, field, fmt in section_fields + ] + sections.insert( + 4, + ( + "Speedup (DFLASH / baseline)", + _compute_speedup( + metrics_map[("baseline", "output_toks_per_s")], + metrics_map[("dflash", "output_toks_per_s")], + ), + ".3f", + ), + ) + + for title, values, fmt in sections: + print(f"\n{title}") + print( + _format_table( + tp_sizes=tp_sizes, + concurrencies=concurrencies, + values=values, + float_fmt=fmt, + ) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", default="HumanEval.jsonl") + parser.add_argument( + "--target-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-PTPC-FP8/" + ) + parser.add_argument( + "--draft-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-DFlash/" + ) + parser.add_argument( + "--base-url", + default=None, + help=( + "Benchmark an already-running server at this URL (e.g. http://127.0.0.1:9080) " + "instead of launching one server per config. All server-side flags " + "(--tp-size, --attention-backends, DFLASH flags, ...) are then ignored: the " + "running server's own configuration is what gets measured, and the first " + "--tp-sizes / --attention-backends value is used only as a summary label." + ), + ) + parser.add_argument( + "--external-mode", + default="auto", + choices=["auto", "baseline", "dflash"], + help=( + "Which summary row a --base-url run fills in. `auto` classifies by whether " + "the server reports spec_verify_ct; `dflash` additionally enforces the " + "DFLASH sanity check." + ), + ) + parser.add_argument( + "--skip-baseline", + action="store_true", + help="Skip running the baseline (target-only) sweep; only run DFLASH and report N/A for baseline/speedup.", + ) + parser.add_argument( + "--batch-requests", + action="store_true", + help="Send prompts as server-side batched /generate requests (batch size = concurrency) instead of client-side concurrent requests.", + ) + parser.add_argument("--max-new-tokens", type=int, default=1024) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=1) + parser.add_argument( + "--timeout-s", + type=int, + default=3600, + help=( + "Timeout in seconds for benchmarked /generate calls and server startup " + "health checks." + ), + ) + parser.add_argument( + "--exec-timeout-s", + type=float, + default=10.0, + help="Per-problem wall-clock timeout for pass@1 code execution.", + ) + parser.add_argument( + "--mem-fraction-static", + type=float, + default=None, + help="Optional server --mem-fraction-static override. If unset, use the server auto heuristic.", + ) + parser.add_argument("--disable-radix-cache", action="store_true") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument( + "--speculative-num-draft-tokens", + type=int, + default=16, + help="DFLASH verify window length (matches the 397B launch script).", + ) + parser.add_argument( + "--page-size", + type=int, + default=None, + help="Optional server --page-size override for both baseline and DFLASH runs.", + ) + parser.add_argument("--max-running-requests", type=int, default=32) + parser.add_argument( + "--mamba-scheduler-strategy", + default="extra_buffer", + help=( + "Server --mamba-scheduler-strategy value to pass through to benchmark " + "runs, e.g. `no_buffer` or `extra_buffer`." + ), + ) + parser.add_argument("--tp-sizes", default="8") + parser.add_argument("--concurrencies", default="1") + parser.add_argument( + "--questions-per-concurrency-base", + type=int, + default=128, + help="num_questions = base * concurrency (default matches the sweep plan).", + ) + parser.add_argument( + "--max-questions-per-config", + type=int, + default=164, + help="Cap num_questions per (tp, concurrency) run (HumanEval has 164 items).", + ) + parser.add_argument("--attention-backends", default="aiter") + parser.add_argument( + "--speculative-draft-attention-backend", + default="triton", + help="Optional server --speculative-draft-attention-backend override for DFLASH runs.", + ) + parser.add_argument( + "--speculative-dflash-draft-window-size", + type=int, + default=None, + help="Optional server --speculative-dflash-draft-window-size override for DFLASH runs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + # With --base-url the server runs elsewhere (or was launched by hand), so this + # process is a pure HTTP client and needs no local GPU. + external = bool(args.base_url) + if not external and not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this sweep.") + if args.temperature < 0.0: + raise RuntimeError(f"--temperature must be >= 0, got {args.temperature}.") + if not (0.0 < args.top_p <= 1.0): + raise RuntimeError(f"--top-p must be in (0, 1], got {args.top_p}.") + if args.top_k == 0 or args.top_k < -1: + raise RuntimeError(f"--top-k must be -1 (all vocab) or >= 1, got {args.top_k}.") + if args.timeout_s <= 0: + raise RuntimeError(f"--timeout-s must be > 0, got {args.timeout_s}.") + + tp_sizes = _parse_int_csv(args.tp_sizes) + if external: + # Labels only: nothing is launched, so there is no GPU count to validate against. + tp_sizes = tp_sizes[:1] or [0] + else: + visible_gpus = int(torch.cuda.device_count()) + tp_sizes = [tp for tp in tp_sizes if tp >= 1 and tp <= visible_gpus] + if not tp_sizes: + raise RuntimeError( + f"No tp sizes are runnable with visible_gpus={visible_gpus}. " + "Set CUDA_VISIBLE_DEVICES accordingly." + ) + + concurrencies = _parse_int_csv(args.concurrencies) + concurrencies = [c for c in concurrencies if c >= 1] + if not concurrencies: + raise RuntimeError("No concurrencies specified.") + + num_questions_by_conc = { + c: min( + int(args.questions_per_concurrency_base) * int(c), + int(args.max_questions_per_config), + ) + for c in concurrencies + } + max_questions = max(num_questions_by_conc.values()) + + attention_backends = [ + s.strip() for s in args.attention_backends.split(",") if s.strip() + ] + if external: + # Label only; the running server already picked its backend. + device_sm = 0 + attention_backends = attention_backends[:1] or ["external"] + else: + device_sm = get_device_sm() + attention_backends = _filter_attention_backends( + attention_backends, device_sm=device_sm + ) + + data_path = _maybe_download_humaneval(args.data_path) + lines = list(read_jsonl(data_path)) + if len(lines) < max_questions: + raise RuntimeError( + f"HumanEval file only has {len(lines)} lines, but need {max_questions}." + ) + + tokenizer = AutoTokenizer.from_pretrained(args.target_model) + + prompts: list[str] = [] + tasks: list[dict] = [] + for i in range(max_questions): + row = lines[i] + user_content = ( + "Complete the following Python function. Return the complete function " + "definition inside a single ```python code block, and do not add " + "example usage or explanations.\n\n```python\n" + + row["prompt"] + + "\n```" + ) + prompts.append( + tokenizer.apply_chat_template( + [{"role": "user", "content": user_content}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + ) + tasks.append( + { + "prompt": row["prompt"], + "test": row["test"], + "entry_point": row["entry_point"], + } + ) + + # Results indexed by (backend, tp, concurrency, mode). + results: dict[tuple[str, int, int, str], BenchMetrics] = {} + # Baseline metrics are backend-agnostic in this sweep; run once per TP and reuse. + baseline_cache_by_tp: dict[int, dict[int, BenchMetrics]] = {} + + if external: + base_url = args.base_url.rstrip("/") + backend = attention_backends[0] + tp = tp_sizes[0] + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"external server {base_url} (mode={args.external_mode})", + model_path=args.target_model, + base_url=base_url, + server_args=[], + expect_dflash=(args.external_mode == "dflash"), + prompts=prompts, + tasks=tasks, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + launch_server=False, + ) + for conc, metrics in mode_metrics.items(): + mode_key = args.external_mode + if mode_key == "auto": + mode_key = "dflash" if metrics.spec_verify_ct_sum > 0 else "baseline" + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + return + + for backend_idx, backend in enumerate(attention_backends): + for tp in tp_sizes: + port_base = find_available_port(20000) + common_server_args = _build_common_server_args(args, backend=backend, tp=tp) + mode_runs = _build_mode_runs(args, common_server_args) + + for idx, ( + mode_key, + mode_name, + mode_server_args, + expect_dflash, + ) in enumerate(mode_runs): + if ( + mode_key == "baseline" + and not args.skip_baseline + and backend_idx > 0 + and tp in baseline_cache_by_tp + ): + mode_metrics = baseline_cache_by_tp[tp] + else: + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"backend={backend} tp={tp} ({mode_name})", + model_path=args.target_model, + base_url=f"http://127.0.0.1:{find_available_port(port_base + idx)}", + server_args=mode_server_args, + expect_dflash=expect_dflash, + prompts=prompts, + tasks=tasks, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + ) + if mode_key == "baseline" and not args.skip_baseline: + baseline_cache_by_tp[tp] = mode_metrics + + for conc, metrics in mode_metrics.items(): + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/dflash/bench_dflash_math500_sweep.py b/benchmark/dflash/bench_dflash_math500_sweep.py new file mode 100644 index 000000000000..df257dc5ddf1 --- /dev/null +++ b/benchmark/dflash/bench_dflash_math500_sweep.py @@ -0,0 +1,1002 @@ +"""DFLASH vs baseline MATH500 sweep. + +This is a *benchmark script* (not a CI test): it can take a long time because it +launches servers for multiple (attention_backend, tp_size) configs and runs a +MATH500 workload for each (concurrency, num_questions) setting. + +Accuracy is scored by extracting the final \\boxed{...} answer from the model +output and comparing (after light normalization) against the dataset's gold +answer. This is a string-equivalence check, NOT a symbolic-math or LLM-judge +comparison, so accuracy here is a conservative lower bound. + +Example usage: + ./venv/bin/python benchmark/dflash/bench_dflash_math500_sweep.py + ./venv/bin/python benchmark/dflash/bench_dflash_math500_sweep.py --skip-baseline --concurrencies 32 --tp-sizes 8 +""" + +from __future__ import annotations + +import argparse +import os +import re +import statistics +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Optional + +# Environment matching launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh. +# Applied via setdefault (caller's explicit exports win) BEFORE importing torch +# and sglang: e.g. TVM_FFI_DISABLE_TORCH_C_DLPACK avoids a libtorch_cuda.so load +# failure on this ROCm build during the sglang import chain. +_LAUNCH_ENV = { + "SGLANG_DISABLE_CUDNN_CHECK": "1", + "SGLANG_USE_CUDA_IPC_TRANSPORT": "1", + "SGLANG_VLM_CACHE_SIZE_MB": "8192", + "SGLANG_USE_AITER": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_SHUFFLE": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_FP8HIPB": "1", + "USE_AITER_COMM": "1", + "AITER_MOE_SMALL_BATCH": "1", + "SGLANG_USE_AITER_NEW_CA": "false", + "SGLANG_USE_IPC_POOL_HANDLE_CACHE": "1", + "AITER_MOE_PADDING_SIZE": "192", + "HIP_GDN_SORT_IDX_BS": "32768", + "TVM_FFI_DISABLE_TORCH_C_DLPACK": "1", +} + +for _k, _v in _LAUNCH_ENV.items(): + os.environ.setdefault(_k, _v) + +import requests # noqa: E402 +import torch # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 + +from sglang.srt.utils import get_device_sm, kill_process_tree # noqa: E402 +from sglang.test.test_utils import ( # noqa: E402 + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + find_available_port, + popen_launch_server, +) +from sglang.utils import download_and_cache_file, read_jsonl # noqa: E402 + +INVALID = "\x00__INVALID__\x00" + + +def _parse_int_csv(value: str) -> list[int]: + return [int(x) for x in value.split(",") if x.strip()] + + +def _filter_attention_backends(backends: list[str], *, device_sm: int) -> list[str]: + if not (80 <= device_sm <= 90): + backends = [b for b in backends if b != "fa3"] + if device_sm < 100: + backends = [b for b in backends if b not in ("fa4", "trtllm_mha")] + return backends or ["flashinfer"] + + +def _extract_last_boxed(text: str) -> Optional[str]: + """Return the content of the last \\boxed{...}, handling nested braces.""" + idx = text.rfind("\\boxed") + if idx < 0: + return None + # Move to the first '{' after \boxed. + brace = text.find("{", idx) + if brace < 0: + return None + depth = 0 + for j in range(brace, len(text)): + c = text[j] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + if depth == 0: + return text[brace + 1 : j] + return None + + +def _normalize_answer(ans: str) -> str: + """Light normalization so string comparison is less brittle.""" + if ans is None: + return INVALID + s = ans.strip() + # Strip a wrapping \boxed{...} if the gold answer includes it. + boxed = _extract_last_boxed(s) + if boxed is not None: + s = boxed.strip() + # Remove common LaTeX text wrappers and spacing. + s = s.replace("\\left", "").replace("\\right", "") + s = s.replace("\\!", "").replace("\\,", "").replace("\\;", "").replace("\\ ", "") + s = s.replace("$", "").replace("\\$", "") + s = s.replace("\\%", "").replace("%", "") + s = s.replace("^{\\circ}", "").replace("^\\circ", "") + s = s.replace("\\text{", "").replace("\\mathrm{", "").replace("\\mbox{", "") + s = s.replace(" ", "") + s = s.rstrip(".") + # Drop a trailing lone '}' left over from removed wrappers. + while s.endswith("}") and s.count("{") < s.count("}"): + s = s[:-1] + if s == "": + return INVALID + return s + + +def _score_math(pred_text: str, gold: str) -> tuple[bool, bool]: + """Return (is_correct, is_invalid) for a MATH500 prediction.""" + boxed = _extract_last_boxed(pred_text) + if boxed is None: + return (False, True) + pred = _normalize_answer(boxed) + if pred == INVALID: + return (False, True) + return (pred == _normalize_answer(gold), False) + + +def _maybe_download_math500(data_path: str) -> str: + url = "https://huggingface.co/datasets/HuggingFaceH4/MATH-500/resolve/main/test.jsonl" + if os.path.isfile(data_path): + return data_path + return download_and_cache_file(url) + + +def _flush_cache(base_url: str) -> None: + resp = requests.get(base_url + "/flush_cache", timeout=60) + resp.raise_for_status() + + +def _send_generate( + base_url: str, + text: str | list[str], + *, + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + timeout_s: int, +) -> list[dict]: + if isinstance(text, list) and not text: + return [] + sampling_params: dict = { + "temperature": float(temperature), + "top_p": float(top_p), + "top_k": int(top_k), + "max_new_tokens": int(max_new_tokens), + } + resp = requests.post( + base_url + "/generate", + json={ + "text": text, + "sampling_params": sampling_params, + }, + timeout=int(timeout_s), + ) + resp.raise_for_status() + out = resp.json() + if isinstance(text, list): + if not isinstance(out, list): + raise RuntimeError( + "Expected a list response for batched /generate, but got " + f"type={type(out).__name__}." + ) + if len(out) != len(text): + raise RuntimeError( + "Batched /generate output length mismatch: " + f"got {len(out)} outputs for {len(text)} prompts." + ) + return out + + if isinstance(out, list): + raise RuntimeError( + "Expected an object response for single /generate, but got " + f"type={type(out).__name__}." + ) + return [out] + + +@dataclass(frozen=True) +class BenchMetrics: + latency_s: float + output_tokens: int + output_toks_per_s: float + accuracy: Optional[float] + invalid_rate: Optional[float] + spec_accept_length: Optional[float] + spec_accept_length_weighted: Optional[float] + spec_verify_ct_sum: int + + +def _run_requests( + base_url: str, + *, + prompts: list[str], + labels: Optional[list[str]], + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + concurrency: int, + batch_requests: bool, + timeout_s: int, + expect_dflash: bool, +) -> BenchMetrics: + if labels is not None and len(labels) != len(prompts): + raise ValueError("labels length must match prompts length") + + # Drop the first batch from metrics to exclude one-time JIT/cuda-graph overhead + # that often happens immediately after /flush_cache for large batch sizes. + bs = max(int(concurrency), 1) + if len(prompts) > bs: + warmup_prompts = prompts[:bs] + if batch_requests: + _send_generate( + base_url, + warmup_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = [ + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for prompt in warmup_prompts + ] + for fut in as_completed(futures): + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate warmup request." + ) + + prompts = prompts[bs:] + labels = labels[bs:] if labels is not None else None + + start = time.perf_counter() + total_tokens = 0 + spec_verify_ct_sum = 0 + spec_accept_lengths: list[float] = [] + correct = 0 + invalid = 0 + + def _handle_output(out: dict, label: Optional[str]) -> None: + nonlocal total_tokens, spec_verify_ct_sum, correct, invalid + meta = out.get("meta_info", {}) or {} + total_tokens += int(meta.get("completion_tokens", 0)) + spec_verify_ct_sum += int(meta.get("spec_verify_ct", 0)) + if "spec_accept_length" in meta: + try: + spec_accept_lengths.append(float(meta["spec_accept_length"])) + except (TypeError, ValueError): + pass + + if label is not None: + is_correct, is_invalid = _score_math(out.get("text", ""), label) + if is_invalid: + invalid += 1 + if is_correct: + correct += 1 + + if batch_requests: + bs = max(int(concurrency), 1) + for start_idx in range(0, len(prompts), bs): + chunk_prompts = prompts[start_idx : start_idx + bs] + chunk_labels = ( + labels[start_idx : start_idx + bs] if labels is not None else None + ) + outs = _send_generate( + base_url, + chunk_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + if chunk_labels is None: + for out in outs: + _handle_output(out, None) + else: + for out, label in zip(outs, chunk_labels): + _handle_output(out, label) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = { + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ): i + for i, prompt in enumerate(prompts) + } + for fut in as_completed(futures): + i = futures[fut] + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate request." + ) + label = None if labels is None else labels[i] + _handle_output(outs[0], label) + + latency = time.perf_counter() - start + toks_per_s = total_tokens / max(latency, 1e-6) + + if expect_dflash and spec_verify_ct_sum <= 0: + raise RuntimeError( + "DFLASH sanity check failed: did not observe any `spec_verify_ct` in responses " + "(DFLASH may not have been enabled)." + ) + + spec_accept_length = ( + float(statistics.mean(spec_accept_lengths)) if spec_accept_lengths else None + ) + # Token-weighted global accept length: total completion tokens / total verify + # steps (each request weighted by its length, unlike the per-request mean above). + spec_accept_length_weighted = ( + float(total_tokens) / float(spec_verify_ct_sum) + if spec_verify_ct_sum > 0 + else None + ) + + if labels is None: + acc = None + invalid_rate = None + else: + acc = correct / max(len(prompts), 1) + invalid_rate = invalid / max(len(prompts), 1) + + return BenchMetrics( + latency_s=float(latency), + output_tokens=int(total_tokens), + output_toks_per_s=float(toks_per_s), + accuracy=acc, + invalid_rate=invalid_rate, + spec_accept_length=spec_accept_length, + spec_accept_length_weighted=spec_accept_length_weighted, + spec_verify_ct_sum=int(spec_verify_ct_sum), + ) + + +def _format_table( + *, + tp_sizes: list[int], + concurrencies: list[int], + values: dict[tuple[int, int], Optional[float]], + float_fmt: str, +) -> str: + header = ["tp\\conc"] + [str(c) for c in concurrencies] + rows: list[list[str]] = [header] + for tp in tp_sizes: + row = [str(tp)] + for c in concurrencies: + v = values.get((tp, c), None) + row.append("N/A" if v is None else format(v, float_fmt)) + rows.append(row) + + col_widths = [ + max(len(row[col_idx]) for row in rows) for col_idx in range(len(rows[0])) + ] + + lines: list[str] = [] + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(rows[0]))) + lines.append(" ".join("-" * w for w in col_widths)) + for row in rows[1:]: + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(row))) + return "\n".join(lines) + + +def _build_common_server_args( + args: argparse.Namespace, *, backend: str, tp: int +) -> list[str]: + # Mirror launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh (baseline flags). + common_server_args: list[str] = [ + "--tp-size", + str(tp), + "--reasoning-parser", + "qwen3", + "--tool-call-parser", + "qwen3_coder", + "--enable-multimodal", + "--trust-remote-code", + "--chunked-prefill-size", + "65536", + "--mem-fraction-static", + str(args.mem_fraction_static if args.mem_fraction_static is not None else 0.9), + "--max-prefill-tokens", + "65536", + "--max-running-requests", + str(args.max_running_requests), + "--attention-backend", + backend, + "--mm-attention-backend", + "aiter_attn", + "--mamba-scheduler-strategy", + str(args.mamba_scheduler_strategy), + "--disable-custom-all-reduce", + "--kv-cache-dtype", + "fp8_e4m3", + "--page-size", + str(int(args.page_size) if args.page_size is not None else 64), + ] + if args.disable_radix_cache: + common_server_args.append("--disable-radix-cache") + return common_server_args + + +def _build_mode_runs( + args: argparse.Namespace, common_server_args: list[str] +) -> list[tuple[str, str, list[str], bool]]: + mode_runs: list[tuple[str, str, list[str], bool]] = [] + if not args.skip_baseline: + mode_runs.append(("baseline", "baseline", common_server_args, False)) + mode_runs.append( + ( + "dflash", + "DFLASH", + [ + *common_server_args, + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + args.draft_model, + "--speculative-num-draft-tokens", + str(int(args.speculative_num_draft_tokens)), + *( + [ + "--speculative-dflash-draft-window-size", + str(int(args.speculative_dflash_draft_window_size)), + ] + if args.speculative_dflash_draft_window_size is not None + else [] + ), + *( + [ + "--speculative-draft-attention-backend", + args.speculative_draft_attention_backend, + ] + if args.speculative_draft_attention_backend + else [] + ), + ], + True, + ) + ) + return mode_runs + + +def _collect_metric( + *, + results: dict[tuple[str, int, int, str], BenchMetrics], + backend: str, + tp_sizes: list[int], + concurrencies: list[int], + mode: str, + field: str, +) -> dict[tuple[int, int], Optional[float]]: + out: dict[tuple[int, int], Optional[float]] = {} + for tp in tp_sizes: + for conc in concurrencies: + metrics = results.get((backend, tp, conc, mode), None) + out[(tp, conc)] = None if metrics is None else getattr(metrics, field) + return out + + +def _compute_speedup( + baseline: dict[tuple[int, int], Optional[float]], + dflash: dict[tuple[int, int], Optional[float]], +) -> dict[tuple[int, int], Optional[float]]: + return { + key: None if (b is None or d is None or b <= 0) else (d / b) + for key, b in baseline.items() + for d in [dflash.get(key, None)] + } + + +def _print_kv_lines(items: list[tuple[str, object]]) -> None: + for key, value in items: + print(f"{key}={value}") + + +def _run_mode_for_backend_tp( + *, + mode_label: str, + model_path: str, + base_url: str, + server_args: list[str], + expect_dflash: bool, + prompts: list[str], + labels: list[str], + concurrencies: list[int], + num_questions_by_conc: dict[int, int], + args: argparse.Namespace, + launch_server: bool = True, +) -> dict[int, BenchMetrics]: + print(f"\n=== {mode_label} ===") + proc = None + if launch_server: + server_start_timeout_s = int( + max(DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, args.timeout_s) + ) + proc = popen_launch_server( + model_path, + base_url, + timeout=server_start_timeout_s, + other_args=server_args, + ) + else: + print(f"[external] reusing already-running server at {base_url}.") + try: + _send_generate( + base_url, + "Hello", + max_new_tokens=8, + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + timeout_s=min(int(args.timeout_s), 300), + ) + + metrics_by_conc: dict[int, BenchMetrics] = {} + for conc in concurrencies: + n = num_questions_by_conc[conc] + _flush_cache(base_url) + print( + f"[warmup] run 1 warmup batch (size={conc}) after /flush_cache; excluded from metrics." + ) + metrics = _run_requests( + base_url, + prompts=prompts[: n + conc], + labels=labels[: n + conc], + max_new_tokens=int(args.max_new_tokens), + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + concurrency=int(conc), + batch_requests=bool(args.batch_requests), + timeout_s=int(args.timeout_s), + expect_dflash=expect_dflash, + ) + metrics_by_conc[conc] = metrics + line = ( + f"[{mode_label}] conc={conc:>2} n={n:<4} " + f"toks/s={metrics.output_toks_per_s:,.2f} " + f"latency={metrics.latency_s:.1f}s " + f"acc={metrics.accuracy:.3f} invalid={metrics.invalid_rate:.3f}" + ) + if expect_dflash or metrics.spec_verify_ct_sum > 0: + accept_len = ( + "N/A" + if metrics.spec_accept_length is None + else f"{metrics.spec_accept_length:.3f}" + ) + accept_len_w = ( + "N/A" + if metrics.spec_accept_length_weighted is None + else f"{metrics.spec_accept_length_weighted:.3f}" + ) + line += ( + f" accept_len={accept_len} " + f"accept_len_weighted={accept_len_w} " + f"spec_verify_ct_sum={metrics.spec_verify_ct_sum}" + ) + print(line) + return metrics_by_conc + finally: + if proc is not None: + kill_process_tree(proc.pid) + try: + proc.wait(timeout=30) + except Exception: + pass + + +def _print_summary( + *, + args: argparse.Namespace, + attention_backends: list[str], + tp_sizes: list[int], + concurrencies: list[int], + device_sm: int, + results: dict[tuple[str, int, int, str], BenchMetrics], +) -> None: + print("\n=== DFLASH MATH500 Sweep Summary ===") + if args.base_url: + _print_kv_lines( + [ + ("base_url", args.base_url), + ("external_mode", args.external_mode), + ( + "note", + "server-side flags below were NOT applied; the running server's own " + "config was measured (tp/backend are labels only)", + ), + ] + ) + _print_kv_lines( + [ + ("target_model", args.target_model), + ("draft_model", args.draft_model), + ("max_new_tokens", args.max_new_tokens), + ( + "sampling", + f"temperature:{args.temperature}, top_p:{args.top_p}, top_k:{args.top_k}", + ), + ("attention_backends", ",".join(attention_backends)), + ( + "speculative_draft_attention_backend", + args.speculative_draft_attention_backend, + ), + ( + "speculative_dflash_draft_window_size", + args.speculative_dflash_draft_window_size, + ), + ("tp_sizes", ",".join(str(x) for x in tp_sizes)), + ("concurrencies", ",".join(str(x) for x in concurrencies)), + ( + "questions_per_concurrency_base", + args.questions_per_concurrency_base, + ), + ("device_sm", device_sm), + ("skip_baseline", bool(args.skip_baseline)), + ] + ) + + section_fields = [ + ("Baseline output tok/s", "baseline", "output_toks_per_s", ",.2f"), + ("Baseline accuracy", "baseline", "accuracy", ".3f"), + ("DFLASH output tok/s", "dflash", "output_toks_per_s", ",.2f"), + ("DFLASH accuracy", "dflash", "accuracy", ".3f"), + ( + "DFLASH acceptance length (mean spec_accept_length)", + "dflash", + "spec_accept_length", + ".3f", + ), + ( + "DFLASH acceptance length (token-weighted: total_tokens/verify_ct)", + "dflash", + "spec_accept_length_weighted", + ".3f", + ), + ] + + for backend in attention_backends: + print(f"\n=== Backend: {backend} ===") + metrics_map = { + (mode, field): _collect_metric( + results=results, + backend=backend, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + mode=mode, + field=field, + ) + for _, mode, field, _ in section_fields + } + sections: list[tuple[str, dict[tuple[int, int], Optional[float]], str]] = [ + (title, metrics_map[(mode, field)], fmt) + for title, mode, field, fmt in section_fields + ] + sections.insert( + 4, + ( + "Speedup (DFLASH / baseline)", + _compute_speedup( + metrics_map[("baseline", "output_toks_per_s")], + metrics_map[("dflash", "output_toks_per_s")], + ), + ".3f", + ), + ) + + for title, values, fmt in sections: + print(f"\n{title}") + print( + _format_table( + tp_sizes=tp_sizes, + concurrencies=concurrencies, + values=values, + float_fmt=fmt, + ) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", default="math500_test.jsonl") + parser.add_argument( + "--target-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-PTPC-FP8/" + ) + parser.add_argument( + "--draft-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-DFlash/" + ) + parser.add_argument( + "--base-url", + default=None, + help=( + "Benchmark an already-running server at this URL (e.g. http://127.0.0.1:9080) " + "instead of launching one server per config. All server-side flags " + "(--tp-size, --attention-backends, DFLASH flags, ...) are then ignored: the " + "running server's own configuration is what gets measured, and the first " + "--tp-sizes / --attention-backends value is used only as a summary label." + ), + ) + parser.add_argument( + "--external-mode", + default="auto", + choices=["auto", "baseline", "dflash"], + help=( + "Which summary row a --base-url run fills in. `auto` classifies by whether " + "the server reports spec_verify_ct; `dflash` additionally enforces the " + "DFLASH sanity check." + ), + ) + parser.add_argument( + "--skip-baseline", + action="store_true", + help="Skip running the baseline (target-only) sweep; only run DFLASH and report N/A for baseline/speedup.", + ) + parser.add_argument( + "--batch-requests", + action="store_true", + help="Send prompts as server-side batched /generate requests (batch size = concurrency) instead of client-side concurrent requests.", + ) + parser.add_argument("--max-new-tokens", type=int, default=2048) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=1) + parser.add_argument( + "--timeout-s", + type=int, + default=3600, + help=( + "Timeout in seconds for benchmarked /generate calls and server startup " + "health checks." + ), + ) + parser.add_argument( + "--mem-fraction-static", + type=float, + default=None, + help="Optional server --mem-fraction-static override. If unset, use the server auto heuristic.", + ) + parser.add_argument("--disable-radix-cache", action="store_true") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument( + "--speculative-num-draft-tokens", + type=int, + default=16, + help="DFLASH verify window length (matches the 397B launch script).", + ) + parser.add_argument( + "--page-size", + type=int, + default=None, + help="Optional server --page-size override for both baseline and DFLASH runs.", + ) + parser.add_argument("--max-running-requests", type=int, default=32) + parser.add_argument( + "--mamba-scheduler-strategy", + default="extra_buffer", + help=( + "Server --mamba-scheduler-strategy value to pass through to benchmark " + "runs, e.g. `no_buffer` or `extra_buffer`." + ), + ) + parser.add_argument("--tp-sizes", default="8") + parser.add_argument("--concurrencies", default="1") + parser.add_argument( + "--questions-per-concurrency-base", + type=int, + default=128, + help="num_questions = base * concurrency (default matches the sweep plan).", + ) + parser.add_argument( + "--max-questions-per-config", + type=int, + default=500, + help="Cap num_questions per (tp, concurrency) run (MATH500 has 500 items).", + ) + parser.add_argument("--attention-backends", default="aiter") + parser.add_argument( + "--speculative-draft-attention-backend", + default="triton", + help="Optional server --speculative-draft-attention-backend override for DFLASH runs.", + ) + parser.add_argument( + "--speculative-dflash-draft-window-size", + type=int, + default=None, + help="Optional server --speculative-dflash-draft-window-size override for DFLASH runs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + # With --base-url the server runs elsewhere (or was launched by hand), so this + # process is a pure HTTP client and needs no local GPU. + external = bool(args.base_url) + if not external and not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this sweep.") + if args.temperature < 0.0: + raise RuntimeError(f"--temperature must be >= 0, got {args.temperature}.") + if not (0.0 < args.top_p <= 1.0): + raise RuntimeError(f"--top-p must be in (0, 1], got {args.top_p}.") + if args.top_k == 0 or args.top_k < -1: + raise RuntimeError(f"--top-k must be -1 (all vocab) or >= 1, got {args.top_k}.") + if args.timeout_s <= 0: + raise RuntimeError(f"--timeout-s must be > 0, got {args.timeout_s}.") + + tp_sizes = _parse_int_csv(args.tp_sizes) + if external: + # Labels only: nothing is launched, so there is no GPU count to validate against. + tp_sizes = tp_sizes[:1] or [0] + else: + visible_gpus = int(torch.cuda.device_count()) + tp_sizes = [tp for tp in tp_sizes if tp >= 1 and tp <= visible_gpus] + if not tp_sizes: + raise RuntimeError( + f"No tp sizes are runnable with visible_gpus={visible_gpus}. " + "Set CUDA_VISIBLE_DEVICES accordingly." + ) + + concurrencies = _parse_int_csv(args.concurrencies) + concurrencies = [c for c in concurrencies if c >= 1] + if not concurrencies: + raise RuntimeError("No concurrencies specified.") + + num_questions_by_conc = { + c: min( + int(args.questions_per_concurrency_base) * int(c), + int(args.max_questions_per_config), + ) + for c in concurrencies + } + max_questions = max(num_questions_by_conc.values()) + + attention_backends = [ + s.strip() for s in args.attention_backends.split(",") if s.strip() + ] + if external: + # Label only; the running server already picked its backend. + device_sm = 0 + attention_backends = attention_backends[:1] or ["external"] + else: + device_sm = get_device_sm() + attention_backends = _filter_attention_backends( + attention_backends, device_sm=device_sm + ) + + data_path = _maybe_download_math500(args.data_path) + lines = list(read_jsonl(data_path)) + if len(lines) < max_questions: + raise RuntimeError( + f"MATH500 file only has {len(lines)} lines, but need {max_questions}." + ) + + tokenizer = AutoTokenizer.from_pretrained(args.target_model) + + prompts: list[str] = [] + labels: list[str] = [] + for i in range(max_questions): + user_content = ( + lines[i]["problem"] + + "\nPlease reason step by step, and put your final answer within \\boxed{}." + ) + prompts.append( + tokenizer.apply_chat_template( + [{"role": "user", "content": user_content}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + ) + labels.append(str(lines[i]["answer"])) + + # Results indexed by (backend, tp, concurrency, mode). + results: dict[tuple[str, int, int, str], BenchMetrics] = {} + # Baseline metrics are backend-agnostic in this sweep; run once per TP and reuse. + baseline_cache_by_tp: dict[int, dict[int, BenchMetrics]] = {} + + if external: + base_url = args.base_url.rstrip("/") + backend = attention_backends[0] + tp = tp_sizes[0] + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"external server {base_url} (mode={args.external_mode})", + model_path=args.target_model, + base_url=base_url, + server_args=[], + expect_dflash=(args.external_mode == "dflash"), + prompts=prompts, + labels=labels, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + launch_server=False, + ) + for conc, metrics in mode_metrics.items(): + mode_key = args.external_mode + if mode_key == "auto": + mode_key = "dflash" if metrics.spec_verify_ct_sum > 0 else "baseline" + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + return + + for backend_idx, backend in enumerate(attention_backends): + for tp in tp_sizes: + port_base = find_available_port(20000) + common_server_args = _build_common_server_args(args, backend=backend, tp=tp) + mode_runs = _build_mode_runs(args, common_server_args) + + for idx, ( + mode_key, + mode_name, + mode_server_args, + expect_dflash, + ) in enumerate(mode_runs): + if ( + mode_key == "baseline" + and not args.skip_baseline + and backend_idx > 0 + and tp in baseline_cache_by_tp + ): + mode_metrics = baseline_cache_by_tp[tp] + else: + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"backend={backend} tp={tp} ({mode_name})", + model_path=args.target_model, + base_url=f"http://127.0.0.1:{find_available_port(port_base + idx)}", + server_args=mode_server_args, + expect_dflash=expect_dflash, + prompts=prompts, + labels=labels, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + ) + if mode_key == "baseline" and not args.skip_baseline: + baseline_cache_by_tp[tp] = mode_metrics + + for conc, metrics in mode_metrics.items(): + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/dflash/bench_dflash_mbpp_sweep.py b/benchmark/dflash/bench_dflash_mbpp_sweep.py new file mode 100644 index 000000000000..7cc7ccdabe7b --- /dev/null +++ b/benchmark/dflash/bench_dflash_mbpp_sweep.py @@ -0,0 +1,1036 @@ +"""DFLASH vs baseline MBPP sweep. + +This is a *benchmark script* (not a CI test): it can take a long time because it +launches servers for multiple (attention_backend, tp_size) configs and runs an +MBPP workload for each (concurrency, num_questions) setting. + +Accuracy is pass@1: the model completion is combined with the problem's +`test_setup_code` and `test_list` assertions and executed in a sandboxed +subprocess. Code execution happens AFTER the timed generation region so it never +contaminates throughput / accept-length. + +Security note: this EXECUTES model-generated code in subprocesses. Run only in a +disposable/trusted environment. + +Example usage: + ./venv/bin/python benchmark/dflash/bench_dflash_mbpp_sweep.py + ./venv/bin/python benchmark/dflash/bench_dflash_mbpp_sweep.py --skip-baseline --concurrencies 32 --tp-sizes 8 +""" + +from __future__ import annotations + +import argparse +import multiprocessing as mp +import os +import re +import statistics +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Optional + +# Environment matching launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh. +# Applied via setdefault (caller's explicit exports win) BEFORE importing torch +# and sglang: e.g. TVM_FFI_DISABLE_TORCH_C_DLPACK avoids a libtorch_cuda.so load +# failure on this ROCm build during the sglang import chain. +_LAUNCH_ENV = { + "SGLANG_DISABLE_CUDNN_CHECK": "1", + "SGLANG_USE_CUDA_IPC_TRANSPORT": "1", + "SGLANG_VLM_CACHE_SIZE_MB": "8192", + "SGLANG_USE_AITER": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_SHUFFLE": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_FP8HIPB": "1", + "USE_AITER_COMM": "1", + "AITER_MOE_SMALL_BATCH": "1", + "SGLANG_USE_AITER_NEW_CA": "false", + "SGLANG_USE_IPC_POOL_HANDLE_CACHE": "1", + "AITER_MOE_PADDING_SIZE": "192", + "HIP_GDN_SORT_IDX_BS": "32768", + "TVM_FFI_DISABLE_TORCH_C_DLPACK": "1", +} + +for _k, _v in _LAUNCH_ENV.items(): + os.environ.setdefault(_k, _v) + +import requests # noqa: E402 +import torch # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 + +from sglang.srt.utils import get_device_sm, kill_process_tree # noqa: E402 +from sglang.test.test_utils import ( # noqa: E402 + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + find_available_port, + popen_launch_server, +) +from sglang.utils import download_and_cache_file, read_jsonl # noqa: E402 + + +def _parse_int_csv(value: str) -> list[int]: + return [int(x) for x in value.split(",") if x.strip()] + + +def _filter_attention_backends(backends: list[str], *, device_sm: int) -> list[str]: + if not (80 <= device_sm <= 90): + backends = [b for b in backends if b != "fa3"] + if device_sm < 100: + backends = [b for b in backends if b not in ("fa4", "trtllm_mha")] + return backends or ["flashinfer"] + + +_CODE_BLOCK_RE = re.compile(r"```(?:python)?\s*\n?(.*?)```", re.DOTALL | re.IGNORECASE) + + +def _extract_code(text: str) -> str: + """Pull a code body from a model response. + + Prefers the last fenced ```python block; falls back to the raw text. + """ + blocks = _CODE_BLOCK_RE.findall(text) + if blocks: + # Prefer a block that defines a function. + for block in reversed(blocks): + if "def " in block: + return block + return blocks[-1] + return text + + +def _build_mbpp_program( + *, completion_text: str, test_setup_code: str, test_list: list[str] +) -> str: + code = _extract_code(completion_text) + parts = [code] + if test_setup_code: + parts.append(test_setup_code) + parts.extend(test_list) + return "\n".join(parts) + "\n" + + +def _exec_worker(program: str, queue) -> None: # pragma: no cover - child process + # Best-effort resource limits; the parent also enforces a wall-clock timeout. + try: + import resource + + resource.setrlimit(resource.RLIMIT_CPU, (10, 10)) + except Exception: + pass + + g: dict = {"__name__": "__mbpp__"} + try: + exec(program, g) + queue.put(True) + except Exception: + queue.put(False) + + +def _check_program(program: str, timeout_s: float) -> bool: + """Execute `program` in a subprocess; True iff it runs without raising.""" + ctx = mp.get_context("fork") + queue = ctx.Queue() + proc = ctx.Process(target=_exec_worker, args=(program, queue)) + proc.start() + proc.join(timeout_s) + if proc.is_alive(): + proc.terminate() + proc.join(1) + if proc.is_alive(): + proc.kill() + proc.join() + return False + try: + return bool(queue.get_nowait()) + except Exception: + return False + + +def _maybe_download_mbpp(data_path: str) -> str: + url = "https://raw.githubusercontent.com/google-research/google-research/master/mbpp/mbpp.jsonl" + if os.path.isfile(data_path): + return data_path + return download_and_cache_file(url) + + +def _flush_cache(base_url: str) -> None: + resp = requests.get(base_url + "/flush_cache", timeout=60) + resp.raise_for_status() + + +def _send_generate( + base_url: str, + text: str | list[str], + *, + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + timeout_s: int, +) -> list[dict]: + if isinstance(text, list) and not text: + return [] + sampling_params: dict = { + "temperature": float(temperature), + "top_p": float(top_p), + "top_k": int(top_k), + "max_new_tokens": int(max_new_tokens), + } + resp = requests.post( + base_url + "/generate", + json={ + "text": text, + "sampling_params": sampling_params, + }, + timeout=int(timeout_s), + ) + resp.raise_for_status() + out = resp.json() + if isinstance(text, list): + if not isinstance(out, list): + raise RuntimeError( + "Expected a list response for batched /generate, but got " + f"type={type(out).__name__}." + ) + if len(out) != len(text): + raise RuntimeError( + "Batched /generate output length mismatch: " + f"got {len(out)} outputs for {len(text)} prompts." + ) + return out + + if isinstance(out, list): + raise RuntimeError( + "Expected an object response for single /generate, but got " + f"type={type(out).__name__}." + ) + return [out] + + +@dataclass(frozen=True) +class BenchMetrics: + latency_s: float + output_tokens: int + output_toks_per_s: float + accuracy: Optional[float] + invalid_rate: Optional[float] + spec_accept_length: Optional[float] + spec_accept_length_weighted: Optional[float] + spec_verify_ct_sum: int + + +def _run_requests( + base_url: str, + *, + prompts: list[str], + tasks: Optional[list[dict]], + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + concurrency: int, + batch_requests: bool, + timeout_s: int, + expect_dflash: bool, + exec_timeout_s: float, +) -> BenchMetrics: + """`tasks[i]` carries the fields needed to score prompt `i` (or None to skip).""" + if tasks is not None and len(tasks) != len(prompts): + raise ValueError("tasks length must match prompts length") + + # Drop the first batch from metrics to exclude one-time JIT/cuda-graph overhead + # that often happens immediately after /flush_cache for large batch sizes. + bs = max(int(concurrency), 1) + if len(prompts) > bs: + warmup_prompts = prompts[:bs] + if batch_requests: + _send_generate( + base_url, + warmup_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = [ + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for prompt in warmup_prompts + ] + for fut in as_completed(futures): + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate warmup request." + ) + + prompts = prompts[bs:] + tasks = tasks[bs:] if tasks is not None else None + + start = time.perf_counter() + total_tokens = 0 + spec_verify_ct_sum = 0 + spec_accept_lengths: list[float] = [] + # Collect completions during the timed region; score AFTER timing stops. + completions: list[Optional[str]] = [None] * len(prompts) + + def _handle_output(out: dict, index: int) -> None: + nonlocal total_tokens, spec_verify_ct_sum + meta = out.get("meta_info", {}) or {} + total_tokens += int(meta.get("completion_tokens", 0)) + spec_verify_ct_sum += int(meta.get("spec_verify_ct", 0)) + if "spec_accept_length" in meta: + try: + spec_accept_lengths.append(float(meta["spec_accept_length"])) + except (TypeError, ValueError): + pass + completions[index] = out.get("text", "") + + if batch_requests: + bs = max(int(concurrency), 1) + for start_idx in range(0, len(prompts), bs): + chunk_prompts = prompts[start_idx : start_idx + bs] + outs = _send_generate( + base_url, + chunk_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for offset, out in enumerate(outs): + _handle_output(out, start_idx + offset) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = { + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ): i + for i, prompt in enumerate(prompts) + } + for fut in as_completed(futures): + i = futures[fut] + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate request." + ) + _handle_output(outs[0], i) + + latency = time.perf_counter() - start + toks_per_s = total_tokens / max(latency, 1e-6) + + if expect_dflash and spec_verify_ct_sum <= 0: + raise RuntimeError( + "DFLASH sanity check failed: did not observe any `spec_verify_ct` in responses " + "(DFLASH may not have been enabled)." + ) + + spec_accept_length = ( + float(statistics.mean(spec_accept_lengths)) if spec_accept_lengths else None + ) + # Token-weighted global accept length: total completion tokens / total verify + # steps (each request weighted by its length, unlike the per-request mean above). + spec_accept_length_weighted = ( + float(total_tokens) / float(spec_verify_ct_sum) + if spec_verify_ct_sum > 0 + else None + ) + + # --- Post-timing pass@1 scoring via sandboxed code execution --- + if tasks is None: + acc = None + invalid_rate = None + else: + correct = 0 + invalid = 0 + for i, task in enumerate(tasks): + comp = completions[i] + if not comp: + invalid += 1 + continue + program = _build_mbpp_program( + completion_text=comp, + test_setup_code=task.get("test_setup_code", ""), + test_list=task["test_list"], + ) + if "def " not in program: + invalid += 1 + continue + if _check_program(program, exec_timeout_s): + correct += 1 + acc = correct / max(len(tasks), 1) + invalid_rate = invalid / max(len(tasks), 1) + + return BenchMetrics( + latency_s=float(latency), + output_tokens=int(total_tokens), + output_toks_per_s=float(toks_per_s), + accuracy=acc, + invalid_rate=invalid_rate, + spec_accept_length=spec_accept_length, + spec_accept_length_weighted=spec_accept_length_weighted, + spec_verify_ct_sum=int(spec_verify_ct_sum), + ) + + +def _format_table( + *, + tp_sizes: list[int], + concurrencies: list[int], + values: dict[tuple[int, int], Optional[float]], + float_fmt: str, +) -> str: + header = ["tp\\conc"] + [str(c) for c in concurrencies] + rows: list[list[str]] = [header] + for tp in tp_sizes: + row = [str(tp)] + for c in concurrencies: + v = values.get((tp, c), None) + row.append("N/A" if v is None else format(v, float_fmt)) + rows.append(row) + + col_widths = [ + max(len(row[col_idx]) for row in rows) for col_idx in range(len(rows[0])) + ] + + lines: list[str] = [] + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(rows[0]))) + lines.append(" ".join("-" * w for w in col_widths)) + for row in rows[1:]: + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(row))) + return "\n".join(lines) + + +def _build_common_server_args( + args: argparse.Namespace, *, backend: str, tp: int +) -> list[str]: + # Mirror launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh (baseline flags). + common_server_args: list[str] = [ + "--tp-size", + str(tp), + "--reasoning-parser", + "qwen3", + "--tool-call-parser", + "qwen3_coder", + "--enable-multimodal", + "--trust-remote-code", + "--chunked-prefill-size", + "65536", + "--mem-fraction-static", + str(args.mem_fraction_static if args.mem_fraction_static is not None else 0.9), + "--max-prefill-tokens", + "65536", + "--max-running-requests", + str(args.max_running_requests), + "--attention-backend", + backend, + "--mm-attention-backend", + "aiter_attn", + "--mamba-scheduler-strategy", + str(args.mamba_scheduler_strategy), + "--disable-custom-all-reduce", + "--kv-cache-dtype", + "fp8_e4m3", + "--page-size", + str(int(args.page_size) if args.page_size is not None else 64), + ] + if args.disable_radix_cache: + common_server_args.append("--disable-radix-cache") + return common_server_args + + +def _build_mode_runs( + args: argparse.Namespace, common_server_args: list[str] +) -> list[tuple[str, str, list[str], bool]]: + mode_runs: list[tuple[str, str, list[str], bool]] = [] + if not args.skip_baseline: + mode_runs.append(("baseline", "baseline", common_server_args, False)) + mode_runs.append( + ( + "dflash", + "DFLASH", + [ + *common_server_args, + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + args.draft_model, + "--speculative-num-draft-tokens", + str(int(args.speculative_num_draft_tokens)), + *( + [ + "--speculative-dflash-draft-window-size", + str(int(args.speculative_dflash_draft_window_size)), + ] + if args.speculative_dflash_draft_window_size is not None + else [] + ), + *( + [ + "--speculative-draft-attention-backend", + args.speculative_draft_attention_backend, + ] + if args.speculative_draft_attention_backend + else [] + ), + ], + True, + ) + ) + return mode_runs + + +def _collect_metric( + *, + results: dict[tuple[str, int, int, str], BenchMetrics], + backend: str, + tp_sizes: list[int], + concurrencies: list[int], + mode: str, + field: str, +) -> dict[tuple[int, int], Optional[float]]: + out: dict[tuple[int, int], Optional[float]] = {} + for tp in tp_sizes: + for conc in concurrencies: + metrics = results.get((backend, tp, conc, mode), None) + out[(tp, conc)] = None if metrics is None else getattr(metrics, field) + return out + + +def _compute_speedup( + baseline: dict[tuple[int, int], Optional[float]], + dflash: dict[tuple[int, int], Optional[float]], +) -> dict[tuple[int, int], Optional[float]]: + return { + key: None if (b is None or d is None or b <= 0) else (d / b) + for key, b in baseline.items() + for d in [dflash.get(key, None)] + } + + +def _print_kv_lines(items: list[tuple[str, object]]) -> None: + for key, value in items: + print(f"{key}={value}") + + +def _run_mode_for_backend_tp( + *, + mode_label: str, + model_path: str, + base_url: str, + server_args: list[str], + expect_dflash: bool, + prompts: list[str], + tasks: list[dict], + concurrencies: list[int], + num_questions_by_conc: dict[int, int], + args: argparse.Namespace, + launch_server: bool = True, +) -> dict[int, BenchMetrics]: + print(f"\n=== {mode_label} ===") + proc = None + if launch_server: + server_start_timeout_s = int( + max(DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, args.timeout_s) + ) + proc = popen_launch_server( + model_path, + base_url, + timeout=server_start_timeout_s, + other_args=server_args, + ) + else: + print(f"[external] reusing already-running server at {base_url}.") + try: + _send_generate( + base_url, + "Hello", + max_new_tokens=8, + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + timeout_s=min(int(args.timeout_s), 300), + ) + + metrics_by_conc: dict[int, BenchMetrics] = {} + for conc in concurrencies: + n = num_questions_by_conc[conc] + _flush_cache(base_url) + print( + f"[warmup] run 1 warmup batch (size={conc}) after /flush_cache; excluded from metrics." + ) + metrics = _run_requests( + base_url, + prompts=prompts[: n + conc], + tasks=tasks[: n + conc], + max_new_tokens=int(args.max_new_tokens), + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + concurrency=int(conc), + batch_requests=bool(args.batch_requests), + timeout_s=int(args.timeout_s), + expect_dflash=expect_dflash, + exec_timeout_s=float(args.exec_timeout_s), + ) + metrics_by_conc[conc] = metrics + line = ( + f"[{mode_label}] conc={conc:>2} n={n:<4} " + f"toks/s={metrics.output_toks_per_s:,.2f} " + f"latency={metrics.latency_s:.1f}s " + f"pass@1={metrics.accuracy:.3f} invalid={metrics.invalid_rate:.3f}" + ) + if expect_dflash or metrics.spec_verify_ct_sum > 0: + accept_len = ( + "N/A" + if metrics.spec_accept_length is None + else f"{metrics.spec_accept_length:.3f}" + ) + accept_len_w = ( + "N/A" + if metrics.spec_accept_length_weighted is None + else f"{metrics.spec_accept_length_weighted:.3f}" + ) + line += ( + f" accept_len={accept_len} " + f"accept_len_weighted={accept_len_w} " + f"spec_verify_ct_sum={metrics.spec_verify_ct_sum}" + ) + print(line) + return metrics_by_conc + finally: + if proc is not None: + kill_process_tree(proc.pid) + try: + proc.wait(timeout=30) + except Exception: + pass + + +def _print_summary( + *, + args: argparse.Namespace, + attention_backends: list[str], + tp_sizes: list[int], + concurrencies: list[int], + device_sm: int, + results: dict[tuple[str, int, int, str], BenchMetrics], +) -> None: + print("\n=== DFLASH MBPP Sweep Summary ===") + if args.base_url: + _print_kv_lines( + [ + ("base_url", args.base_url), + ("external_mode", args.external_mode), + ( + "note", + "server-side flags below were NOT applied; the running server's own " + "config was measured (tp/backend are labels only)", + ), + ] + ) + _print_kv_lines( + [ + ("target_model", args.target_model), + ("draft_model", args.draft_model), + ("max_new_tokens", args.max_new_tokens), + ( + "sampling", + f"temperature:{args.temperature}, top_p:{args.top_p}, top_k:{args.top_k}", + ), + ("attention_backends", ",".join(attention_backends)), + ( + "speculative_draft_attention_backend", + args.speculative_draft_attention_backend, + ), + ( + "speculative_dflash_draft_window_size", + args.speculative_dflash_draft_window_size, + ), + ("tp_sizes", ",".join(str(x) for x in tp_sizes)), + ("concurrencies", ",".join(str(x) for x in concurrencies)), + ( + "questions_per_concurrency_base", + args.questions_per_concurrency_base, + ), + ("device_sm", device_sm), + ("skip_baseline", bool(args.skip_baseline)), + ] + ) + + section_fields = [ + ("Baseline output tok/s", "baseline", "output_toks_per_s", ",.2f"), + ("Baseline pass@1", "baseline", "accuracy", ".3f"), + ("DFLASH output tok/s", "dflash", "output_toks_per_s", ",.2f"), + ("DFLASH pass@1", "dflash", "accuracy", ".3f"), + ( + "DFLASH acceptance length (mean spec_accept_length)", + "dflash", + "spec_accept_length", + ".3f", + ), + ( + "DFLASH acceptance length (token-weighted: total_tokens/verify_ct)", + "dflash", + "spec_accept_length_weighted", + ".3f", + ), + ] + + for backend in attention_backends: + print(f"\n=== Backend: {backend} ===") + metrics_map = { + (mode, field): _collect_metric( + results=results, + backend=backend, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + mode=mode, + field=field, + ) + for _, mode, field, _ in section_fields + } + sections: list[tuple[str, dict[tuple[int, int], Optional[float]], str]] = [ + (title, metrics_map[(mode, field)], fmt) + for title, mode, field, fmt in section_fields + ] + sections.insert( + 4, + ( + "Speedup (DFLASH / baseline)", + _compute_speedup( + metrics_map[("baseline", "output_toks_per_s")], + metrics_map[("dflash", "output_toks_per_s")], + ), + ".3f", + ), + ) + + for title, values, fmt in sections: + print(f"\n{title}") + print( + _format_table( + tp_sizes=tp_sizes, + concurrencies=concurrencies, + values=values, + float_fmt=fmt, + ) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", default="mbpp.jsonl") + parser.add_argument( + "--target-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-PTPC-FP8/" + ) + parser.add_argument( + "--draft-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-DFlash/" + ) + parser.add_argument( + "--base-url", + default=None, + help=( + "Benchmark an already-running server at this URL (e.g. http://127.0.0.1:9080) " + "instead of launching one server per config. All server-side flags " + "(--tp-size, --attention-backends, DFLASH flags, ...) are then ignored: the " + "running server's own configuration is what gets measured, and the first " + "--tp-sizes / --attention-backends value is used only as a summary label." + ), + ) + parser.add_argument( + "--external-mode", + default="auto", + choices=["auto", "baseline", "dflash"], + help=( + "Which summary row a --base-url run fills in. `auto` classifies by whether " + "the server reports spec_verify_ct; `dflash` additionally enforces the " + "DFLASH sanity check." + ), + ) + parser.add_argument( + "--skip-baseline", + action="store_true", + help="Skip running the baseline (target-only) sweep; only run DFLASH and report N/A for baseline/speedup.", + ) + parser.add_argument( + "--batch-requests", + action="store_true", + help="Send prompts as server-side batched /generate requests (batch size = concurrency) instead of client-side concurrent requests.", + ) + parser.add_argument("--max-new-tokens", type=int, default=1024) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=1) + parser.add_argument( + "--timeout-s", + type=int, + default=3600, + help=( + "Timeout in seconds for benchmarked /generate calls and server startup " + "health checks." + ), + ) + parser.add_argument( + "--exec-timeout-s", + type=float, + default=10.0, + help="Per-problem wall-clock timeout for pass@1 code execution.", + ) + parser.add_argument( + "--mem-fraction-static", + type=float, + default=None, + help="Optional server --mem-fraction-static override. If unset, use the server auto heuristic.", + ) + parser.add_argument("--disable-radix-cache", action="store_true") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument( + "--speculative-num-draft-tokens", + type=int, + default=16, + help="DFLASH verify window length (matches the 397B launch script).", + ) + parser.add_argument( + "--page-size", + type=int, + default=None, + help="Optional server --page-size override for both baseline and DFLASH runs.", + ) + parser.add_argument("--max-running-requests", type=int, default=32) + parser.add_argument( + "--mamba-scheduler-strategy", + default="extra_buffer", + help=( + "Server --mamba-scheduler-strategy value to pass through to benchmark " + "runs, e.g. `no_buffer` or `extra_buffer`." + ), + ) + parser.add_argument("--tp-sizes", default="8") + parser.add_argument("--concurrencies", default="1") + parser.add_argument( + "--questions-per-concurrency-base", + type=int, + default=128, + help="num_questions = base * concurrency (default matches the sweep plan).", + ) + parser.add_argument( + "--max-questions-per-config", + type=int, + default=500, + help="Cap num_questions per (tp, concurrency) run (MBPP full set is ~974).", + ) + parser.add_argument("--attention-backends", default="aiter") + parser.add_argument( + "--speculative-draft-attention-backend", + default="triton", + help="Optional server --speculative-draft-attention-backend override for DFLASH runs.", + ) + parser.add_argument( + "--speculative-dflash-draft-window-size", + type=int, + default=None, + help="Optional server --speculative-dflash-draft-window-size override for DFLASH runs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + # With --base-url the server runs elsewhere (or was launched by hand), so this + # process is a pure HTTP client and needs no local GPU. + external = bool(args.base_url) + if not external and not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this sweep.") + if args.temperature < 0.0: + raise RuntimeError(f"--temperature must be >= 0, got {args.temperature}.") + if not (0.0 < args.top_p <= 1.0): + raise RuntimeError(f"--top-p must be in (0, 1], got {args.top_p}.") + if args.top_k == 0 or args.top_k < -1: + raise RuntimeError(f"--top-k must be -1 (all vocab) or >= 1, got {args.top_k}.") + if args.timeout_s <= 0: + raise RuntimeError(f"--timeout-s must be > 0, got {args.timeout_s}.") + + tp_sizes = _parse_int_csv(args.tp_sizes) + if external: + # Labels only: nothing is launched, so there is no GPU count to validate against. + tp_sizes = tp_sizes[:1] or [0] + else: + visible_gpus = int(torch.cuda.device_count()) + tp_sizes = [tp for tp in tp_sizes if tp >= 1 and tp <= visible_gpus] + if not tp_sizes: + raise RuntimeError( + f"No tp sizes are runnable with visible_gpus={visible_gpus}. " + "Set CUDA_VISIBLE_DEVICES accordingly." + ) + + concurrencies = _parse_int_csv(args.concurrencies) + concurrencies = [c for c in concurrencies if c >= 1] + if not concurrencies: + raise RuntimeError("No concurrencies specified.") + + num_questions_by_conc = { + c: min( + int(args.questions_per_concurrency_base) * int(c), + int(args.max_questions_per_config), + ) + for c in concurrencies + } + max_questions = max(num_questions_by_conc.values()) + + attention_backends = [ + s.strip() for s in args.attention_backends.split(",") if s.strip() + ] + if external: + # Label only; the running server already picked its backend. + device_sm = 0 + attention_backends = attention_backends[:1] or ["external"] + else: + device_sm = get_device_sm() + attention_backends = _filter_attention_backends( + attention_backends, device_sm=device_sm + ) + + data_path = _maybe_download_mbpp(args.data_path) + lines = list(read_jsonl(data_path)) + if len(lines) < max_questions: + raise RuntimeError( + f"MBPP file only has {len(lines)} lines, but need {max_questions}." + ) + + tokenizer = AutoTokenizer.from_pretrained(args.target_model) + + prompts: list[str] = [] + tasks: list[dict] = [] + for i in range(max_questions): + row = lines[i] + test_list = list(row.get("test_list", [])) + example_test = test_list[0] if test_list else "" + user_content = ( + "You are an expert Python programmer. Write a Python function that " + "solves the following task. Return only the function inside a single " + "```python code block, with no explanations.\n\n" + f"Task: {row['text']}\n" + f"Your code must pass this test: {example_test}" + ) + prompts.append( + tokenizer.apply_chat_template( + [{"role": "user", "content": user_content}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + ) + tasks.append( + { + "test_setup_code": row.get("test_setup_code", ""), + "test_list": test_list, + } + ) + + # Results indexed by (backend, tp, concurrency, mode). + results: dict[tuple[str, int, int, str], BenchMetrics] = {} + # Baseline metrics are backend-agnostic in this sweep; run once per TP and reuse. + baseline_cache_by_tp: dict[int, dict[int, BenchMetrics]] = {} + + if external: + base_url = args.base_url.rstrip("/") + backend = attention_backends[0] + tp = tp_sizes[0] + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"external server {base_url} (mode={args.external_mode})", + model_path=args.target_model, + base_url=base_url, + server_args=[], + expect_dflash=(args.external_mode == "dflash"), + prompts=prompts, + tasks=tasks, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + launch_server=False, + ) + for conc, metrics in mode_metrics.items(): + mode_key = args.external_mode + if mode_key == "auto": + mode_key = "dflash" if metrics.spec_verify_ct_sum > 0 else "baseline" + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + return + + for backend_idx, backend in enumerate(attention_backends): + for tp in tp_sizes: + port_base = find_available_port(20000) + common_server_args = _build_common_server_args(args, backend=backend, tp=tp) + mode_runs = _build_mode_runs(args, common_server_args) + + for idx, ( + mode_key, + mode_name, + mode_server_args, + expect_dflash, + ) in enumerate(mode_runs): + if ( + mode_key == "baseline" + and not args.skip_baseline + and backend_idx > 0 + and tp in baseline_cache_by_tp + ): + mode_metrics = baseline_cache_by_tp[tp] + else: + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"backend={backend} tp={tp} ({mode_name})", + model_path=args.target_model, + base_url=f"http://127.0.0.1:{find_available_port(port_base + idx)}", + server_args=mode_server_args, + expect_dflash=expect_dflash, + prompts=prompts, + tasks=tasks, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + ) + if mode_key == "baseline" and not args.skip_baseline: + baseline_cache_by_tp[tp] = mode_metrics + + for conc, metrics in mode_metrics.items(): + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/dflash/bench_dflash_mtbench_sweep.py b/benchmark/dflash/bench_dflash_mtbench_sweep.py new file mode 100644 index 000000000000..7aacbb5ef66c --- /dev/null +++ b/benchmark/dflash/bench_dflash_mtbench_sweep.py @@ -0,0 +1,897 @@ +"""DFLASH vs baseline MT-Bench sweep (performance only). + +This is a *benchmark script* (not a CI test): it can take a long time because it +launches servers for multiple (attention_backend, tp_size) configs and runs an +MT-Bench workload for each (concurrency, num_questions) setting. + +MT-Bench is an open-ended, LLM-judged benchmark. This script does NOT run a judge: +it measures throughput and DFLASH acceptance length only (no quality score). Each +MT-Bench item's first-turn prompt is used as a single-turn generation request. + +Example usage: + ./venv/bin/python benchmark/dflash/bench_dflash_mtbench_sweep.py + ./venv/bin/python benchmark/dflash/bench_dflash_mtbench_sweep.py --skip-baseline --concurrencies 32 --tp-sizes 8 +""" + +from __future__ import annotations + +import argparse +import os +import statistics +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Optional + +# Environment matching launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh. +# Applied via setdefault (caller's explicit exports win) BEFORE importing torch +# and sglang: e.g. TVM_FFI_DISABLE_TORCH_C_DLPACK avoids a libtorch_cuda.so load +# failure on this ROCm build during the sglang import chain. +_LAUNCH_ENV = { + "SGLANG_DISABLE_CUDNN_CHECK": "1", + "SGLANG_USE_CUDA_IPC_TRANSPORT": "1", + "SGLANG_VLM_CACHE_SIZE_MB": "8192", + "SGLANG_USE_AITER": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_SHUFFLE": "1", + "SGLANG_ROCM_USE_AITER_LINEAR_FP8HIPB": "1", + "USE_AITER_COMM": "1", + "AITER_MOE_SMALL_BATCH": "1", + "SGLANG_USE_AITER_NEW_CA": "false", + "SGLANG_USE_IPC_POOL_HANDLE_CACHE": "1", + "AITER_MOE_PADDING_SIZE": "192", + "HIP_GDN_SORT_IDX_BS": "32768", + "TVM_FFI_DISABLE_TORCH_C_DLPACK": "1", +} + +for _k, _v in _LAUNCH_ENV.items(): + os.environ.setdefault(_k, _v) + +import requests # noqa: E402 +import torch # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 + +from sglang.srt.utils import get_device_sm, kill_process_tree # noqa: E402 +from sglang.test.test_utils import ( # noqa: E402 + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + find_available_port, + popen_launch_server, +) +from sglang.utils import download_and_cache_file, read_jsonl # noqa: E402 + + +def _parse_int_csv(value: str) -> list[int]: + return [int(x) for x in value.split(",") if x.strip()] + + +def _filter_attention_backends(backends: list[str], *, device_sm: int) -> list[str]: + if not (80 <= device_sm <= 90): + backends = [b for b in backends if b != "fa3"] + if device_sm < 100: + backends = [b for b in backends if b not in ("fa4", "trtllm_mha")] + return backends or ["flashinfer"] + + +def _maybe_download_mtbench(data_path: str) -> str: + url = "https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl" + if os.path.isfile(data_path): + return data_path + return download_and_cache_file(url) + + +def _flush_cache(base_url: str) -> None: + resp = requests.get(base_url + "/flush_cache", timeout=60) + resp.raise_for_status() + + +def _send_generate( + base_url: str, + text: str | list[str], + *, + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + timeout_s: int, +) -> list[dict]: + if isinstance(text, list) and not text: + return [] + sampling_params: dict = { + "temperature": float(temperature), + "top_p": float(top_p), + "top_k": int(top_k), + "max_new_tokens": int(max_new_tokens), + } + resp = requests.post( + base_url + "/generate", + json={ + "text": text, + "sampling_params": sampling_params, + }, + timeout=int(timeout_s), + ) + resp.raise_for_status() + out = resp.json() + if isinstance(text, list): + if not isinstance(out, list): + raise RuntimeError( + "Expected a list response for batched /generate, but got " + f"type={type(out).__name__}." + ) + if len(out) != len(text): + raise RuntimeError( + "Batched /generate output length mismatch: " + f"got {len(out)} outputs for {len(text)} prompts." + ) + return out + + if isinstance(out, list): + raise RuntimeError( + "Expected an object response for single /generate, but got " + f"type={type(out).__name__}." + ) + return [out] + + +@dataclass(frozen=True) +class BenchMetrics: + latency_s: float + output_tokens: int + output_toks_per_s: float + spec_accept_length: Optional[float] + spec_accept_length_weighted: Optional[float] + spec_verify_ct_sum: int + + +def _run_requests( + base_url: str, + *, + prompts: list[str], + max_new_tokens: int, + temperature: float, + top_p: float, + top_k: int, + concurrency: int, + batch_requests: bool, + timeout_s: int, + expect_dflash: bool, +) -> BenchMetrics: + # Drop the first batch from metrics to exclude one-time JIT/cuda-graph overhead + # that often happens immediately after /flush_cache for large batch sizes. + bs = max(int(concurrency), 1) + if len(prompts) > bs: + warmup_prompts = prompts[:bs] + if batch_requests: + _send_generate( + base_url, + warmup_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = [ + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for prompt in warmup_prompts + ] + for fut in as_completed(futures): + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate warmup request." + ) + + prompts = prompts[bs:] + + start = time.perf_counter() + total_tokens = 0 + spec_verify_ct_sum = 0 + spec_accept_lengths: list[float] = [] + + def _handle_output(out: dict) -> None: + nonlocal total_tokens, spec_verify_ct_sum + meta = out.get("meta_info", {}) or {} + total_tokens += int(meta.get("completion_tokens", 0)) + spec_verify_ct_sum += int(meta.get("spec_verify_ct", 0)) + if "spec_accept_length" in meta: + try: + spec_accept_lengths.append(float(meta["spec_accept_length"])) + except (TypeError, ValueError): + pass + + if batch_requests: + bs = max(int(concurrency), 1) + for start_idx in range(0, len(prompts), bs): + chunk_prompts = prompts[start_idx : start_idx + bs] + outs = _send_generate( + base_url, + chunk_prompts, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for out in outs: + _handle_output(out) + else: + with ThreadPoolExecutor(max_workers=int(concurrency)) as pool: + futures = [ + pool.submit( + _send_generate, + base_url=base_url, + text=prompt, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + timeout_s=timeout_s, + ) + for prompt in prompts + ] + for fut in as_completed(futures): + outs = fut.result() + if len(outs) != 1: + raise RuntimeError( + "Expected exactly one output for single /generate request." + ) + _handle_output(outs[0]) + + latency = time.perf_counter() - start + toks_per_s = total_tokens / max(latency, 1e-6) + + if expect_dflash and spec_verify_ct_sum <= 0: + raise RuntimeError( + "DFLASH sanity check failed: did not observe any `spec_verify_ct` in responses " + "(DFLASH may not have been enabled)." + ) + + spec_accept_length = ( + float(statistics.mean(spec_accept_lengths)) if spec_accept_lengths else None + ) + # Token-weighted global accept length: total completion tokens / total verify + # steps (each request weighted by its length, unlike the per-request mean above). + spec_accept_length_weighted = ( + float(total_tokens) / float(spec_verify_ct_sum) + if spec_verify_ct_sum > 0 + else None + ) + + return BenchMetrics( + latency_s=float(latency), + output_tokens=int(total_tokens), + output_toks_per_s=float(toks_per_s), + spec_accept_length=spec_accept_length, + spec_accept_length_weighted=spec_accept_length_weighted, + spec_verify_ct_sum=int(spec_verify_ct_sum), + ) + + +def _format_table( + *, + tp_sizes: list[int], + concurrencies: list[int], + values: dict[tuple[int, int], Optional[float]], + float_fmt: str, +) -> str: + header = ["tp\\conc"] + [str(c) for c in concurrencies] + rows: list[list[str]] = [header] + for tp in tp_sizes: + row = [str(tp)] + for c in concurrencies: + v = values.get((tp, c), None) + row.append("N/A" if v is None else format(v, float_fmt)) + rows.append(row) + + col_widths = [ + max(len(row[col_idx]) for row in rows) for col_idx in range(len(rows[0])) + ] + + lines: list[str] = [] + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(rows[0]))) + lines.append(" ".join("-" * w for w in col_widths)) + for row in rows[1:]: + lines.append(" ".join(cell.rjust(col_widths[i]) for i, cell in enumerate(row))) + return "\n".join(lines) + + +def _build_common_server_args( + args: argparse.Namespace, *, backend: str, tp: int +) -> list[str]: + # Mirror launch_qwen3.5-397B-fp8_tp8_prefix_cache_DFlash.sh (baseline flags). + common_server_args: list[str] = [ + "--tp-size", + str(tp), + "--reasoning-parser", + "qwen3", + "--tool-call-parser", + "qwen3_coder", + "--enable-multimodal", + "--trust-remote-code", + "--chunked-prefill-size", + "65536", + "--mem-fraction-static", + str(args.mem_fraction_static if args.mem_fraction_static is not None else 0.9), + "--max-prefill-tokens", + "65536", + "--max-running-requests", + str(args.max_running_requests), + "--attention-backend", + backend, + "--mm-attention-backend", + "aiter_attn", + "--mamba-scheduler-strategy", + str(args.mamba_scheduler_strategy), + "--disable-custom-all-reduce", + "--kv-cache-dtype", + "fp8_e4m3", + "--page-size", + str(int(args.page_size) if args.page_size is not None else 64), + ] + if args.disable_radix_cache: + common_server_args.append("--disable-radix-cache") + return common_server_args + + +def _build_mode_runs( + args: argparse.Namespace, common_server_args: list[str] +) -> list[tuple[str, str, list[str], bool]]: + mode_runs: list[tuple[str, str, list[str], bool]] = [] + if not args.skip_baseline: + mode_runs.append(("baseline", "baseline", common_server_args, False)) + mode_runs.append( + ( + "dflash", + "DFLASH", + [ + *common_server_args, + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + args.draft_model, + "--speculative-num-draft-tokens", + str(int(args.speculative_num_draft_tokens)), + *( + [ + "--speculative-dflash-draft-window-size", + str(int(args.speculative_dflash_draft_window_size)), + ] + if args.speculative_dflash_draft_window_size is not None + else [] + ), + *( + [ + "--speculative-draft-attention-backend", + args.speculative_draft_attention_backend, + ] + if args.speculative_draft_attention_backend + else [] + ), + ], + True, + ) + ) + return mode_runs + + +def _collect_metric( + *, + results: dict[tuple[str, int, int, str], BenchMetrics], + backend: str, + tp_sizes: list[int], + concurrencies: list[int], + mode: str, + field: str, +) -> dict[tuple[int, int], Optional[float]]: + out: dict[tuple[int, int], Optional[float]] = {} + for tp in tp_sizes: + for conc in concurrencies: + metrics = results.get((backend, tp, conc, mode), None) + out[(tp, conc)] = None if metrics is None else getattr(metrics, field) + return out + + +def _compute_speedup( + baseline: dict[tuple[int, int], Optional[float]], + dflash: dict[tuple[int, int], Optional[float]], +) -> dict[tuple[int, int], Optional[float]]: + return { + key: None if (b is None or d is None or b <= 0) else (d / b) + for key, b in baseline.items() + for d in [dflash.get(key, None)] + } + + +def _print_kv_lines(items: list[tuple[str, object]]) -> None: + for key, value in items: + print(f"{key}={value}") + + +def _run_mode_for_backend_tp( + *, + mode_label: str, + model_path: str, + base_url: str, + server_args: list[str], + expect_dflash: bool, + prompts: list[str], + concurrencies: list[int], + num_questions_by_conc: dict[int, int], + args: argparse.Namespace, + launch_server: bool = True, +) -> dict[int, BenchMetrics]: + print(f"\n=== {mode_label} ===") + proc = None + if launch_server: + server_start_timeout_s = int( + max(DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, args.timeout_s) + ) + proc = popen_launch_server( + model_path, + base_url, + timeout=server_start_timeout_s, + other_args=server_args, + ) + else: + print(f"[external] reusing already-running server at {base_url}.") + try: + _send_generate( + base_url, + "Hello", + max_new_tokens=8, + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + timeout_s=min(int(args.timeout_s), 300), + ) + + metrics_by_conc: dict[int, BenchMetrics] = {} + for conc in concurrencies: + n = num_questions_by_conc[conc] + _flush_cache(base_url) + print( + f"[warmup] run 1 warmup batch (size={conc}) after /flush_cache; excluded from metrics." + ) + metrics = _run_requests( + base_url, + prompts=prompts[: n + conc], + max_new_tokens=int(args.max_new_tokens), + temperature=float(args.temperature), + top_p=float(args.top_p), + top_k=int(args.top_k), + concurrency=int(conc), + batch_requests=bool(args.batch_requests), + timeout_s=int(args.timeout_s), + expect_dflash=expect_dflash, + ) + metrics_by_conc[conc] = metrics + line = ( + f"[{mode_label}] conc={conc:>2} n={n:<4} " + f"toks/s={metrics.output_toks_per_s:,.2f} " + f"latency={metrics.latency_s:.1f}s " + f"out_tokens={metrics.output_tokens}" + ) + if expect_dflash or metrics.spec_verify_ct_sum > 0: + accept_len = ( + "N/A" + if metrics.spec_accept_length is None + else f"{metrics.spec_accept_length:.3f}" + ) + accept_len_w = ( + "N/A" + if metrics.spec_accept_length_weighted is None + else f"{metrics.spec_accept_length_weighted:.3f}" + ) + line += ( + f" accept_len={accept_len} " + f"accept_len_weighted={accept_len_w} " + f"spec_verify_ct_sum={metrics.spec_verify_ct_sum}" + ) + print(line) + return metrics_by_conc + finally: + if proc is not None: + kill_process_tree(proc.pid) + try: + proc.wait(timeout=30) + except Exception: + pass + + +def _print_summary( + *, + args: argparse.Namespace, + attention_backends: list[str], + tp_sizes: list[int], + concurrencies: list[int], + device_sm: int, + results: dict[tuple[str, int, int, str], BenchMetrics], +) -> None: + print("\n=== DFLASH MT-Bench Sweep Summary (performance only; no judge) ===") + if args.base_url: + _print_kv_lines( + [ + ("base_url", args.base_url), + ("external_mode", args.external_mode), + ( + "note", + "server-side flags below were NOT applied; the running server's own " + "config was measured (tp/backend are labels only)", + ), + ] + ) + _print_kv_lines( + [ + ("target_model", args.target_model), + ("draft_model", args.draft_model), + ("max_new_tokens", args.max_new_tokens), + ( + "sampling", + f"temperature:{args.temperature}, top_p:{args.top_p}, top_k:{args.top_k}", + ), + ("attention_backends", ",".join(attention_backends)), + ( + "speculative_draft_attention_backend", + args.speculative_draft_attention_backend, + ), + ( + "speculative_dflash_draft_window_size", + args.speculative_dflash_draft_window_size, + ), + ("tp_sizes", ",".join(str(x) for x in tp_sizes)), + ("concurrencies", ",".join(str(x) for x in concurrencies)), + ( + "questions_per_concurrency_base", + args.questions_per_concurrency_base, + ), + ("device_sm", device_sm), + ("skip_baseline", bool(args.skip_baseline)), + ] + ) + + section_fields = [ + ("Baseline output tok/s", "baseline", "output_toks_per_s", ",.2f"), + ("DFLASH output tok/s", "dflash", "output_toks_per_s", ",.2f"), + ( + "DFLASH acceptance length (mean spec_accept_length)", + "dflash", + "spec_accept_length", + ".3f", + ), + ( + "DFLASH acceptance length (token-weighted: total_tokens/verify_ct)", + "dflash", + "spec_accept_length_weighted", + ".3f", + ), + ] + + for backend in attention_backends: + print(f"\n=== Backend: {backend} ===") + metrics_map = { + (mode, field): _collect_metric( + results=results, + backend=backend, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + mode=mode, + field=field, + ) + for _, mode, field, _ in section_fields + } + sections: list[tuple[str, dict[tuple[int, int], Optional[float]], str]] = [ + (title, metrics_map[(mode, field)], fmt) + for title, mode, field, fmt in section_fields + ] + sections.insert( + 2, + ( + "Speedup (DFLASH / baseline)", + _compute_speedup( + metrics_map[("baseline", "output_toks_per_s")], + metrics_map[("dflash", "output_toks_per_s")], + ), + ".3f", + ), + ) + + for title, values, fmt in sections: + print(f"\n{title}") + print( + _format_table( + tp_sizes=tp_sizes, + concurrencies=concurrencies, + values=values, + float_fmt=fmt, + ) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", default="mt_bench_question.jsonl") + parser.add_argument( + "--target-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-PTPC-FP8/" + ) + parser.add_argument( + "--draft-model", default="/shared/xiaomi/Qwen3.5-397B-A17B-DFlash/" + ) + parser.add_argument( + "--base-url", + default=None, + help=( + "Benchmark an already-running server at this URL (e.g. http://127.0.0.1:9080) " + "instead of launching one server per config. All server-side flags " + "(--tp-size, --attention-backends, DFLASH flags, ...) are then ignored: the " + "running server's own configuration is what gets measured, and the first " + "--tp-sizes / --attention-backends value is used only as a summary label." + ), + ) + parser.add_argument( + "--external-mode", + default="auto", + choices=["auto", "baseline", "dflash"], + help=( + "Which summary row a --base-url run fills in. `auto` classifies by whether " + "the server reports spec_verify_ct; `dflash` additionally enforces the " + "DFLASH sanity check." + ), + ) + parser.add_argument( + "--skip-baseline", + action="store_true", + help="Skip running the baseline (target-only) sweep; only run DFLASH and report N/A for baseline/speedup.", + ) + parser.add_argument( + "--batch-requests", + action="store_true", + help="Send prompts as server-side batched /generate requests (batch size = concurrency) instead of client-side concurrent requests.", + ) + parser.add_argument("--max-new-tokens", type=int, default=1024) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=1) + parser.add_argument( + "--timeout-s", + type=int, + default=3600, + help=( + "Timeout in seconds for benchmarked /generate calls and server startup " + "health checks." + ), + ) + parser.add_argument( + "--mem-fraction-static", + type=float, + default=None, + help="Optional server --mem-fraction-static override. If unset, use the server auto heuristic.", + ) + parser.add_argument("--disable-radix-cache", action="store_true") + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument( + "--speculative-num-draft-tokens", + type=int, + default=16, + help="DFLASH verify window length (matches the 397B launch script).", + ) + parser.add_argument( + "--page-size", + type=int, + default=None, + help="Optional server --page-size override for both baseline and DFLASH runs.", + ) + parser.add_argument("--max-running-requests", type=int, default=32) + parser.add_argument( + "--mamba-scheduler-strategy", + default="extra_buffer", + help=( + "Server --mamba-scheduler-strategy value to pass through to benchmark " + "runs, e.g. `no_buffer` or `extra_buffer`." + ), + ) + parser.add_argument("--tp-sizes", default="8") + parser.add_argument("--concurrencies", default="1") + parser.add_argument( + "--questions-per-concurrency-base", + type=int, + default=80, + help="num_questions = base * concurrency (default matches the sweep plan).", + ) + parser.add_argument( + "--max-questions-per-config", + type=int, + default=80, + help="Cap num_questions per (tp, concurrency) run (MT-Bench has 80 items).", + ) + parser.add_argument("--attention-backends", default="aiter") + parser.add_argument( + "--speculative-draft-attention-backend", + default="triton", + help="Optional server --speculative-draft-attention-backend override for DFLASH runs.", + ) + parser.add_argument( + "--speculative-dflash-draft-window-size", + type=int, + default=None, + help="Optional server --speculative-dflash-draft-window-size override for DFLASH runs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + # With --base-url the server runs elsewhere (or was launched by hand), so this + # process is a pure HTTP client and needs no local GPU. + external = bool(args.base_url) + if not external and not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this sweep.") + if args.temperature < 0.0: + raise RuntimeError(f"--temperature must be >= 0, got {args.temperature}.") + if not (0.0 < args.top_p <= 1.0): + raise RuntimeError(f"--top-p must be in (0, 1], got {args.top_p}.") + if args.top_k == 0 or args.top_k < -1: + raise RuntimeError(f"--top-k must be -1 (all vocab) or >= 1, got {args.top_k}.") + if args.timeout_s <= 0: + raise RuntimeError(f"--timeout-s must be > 0, got {args.timeout_s}.") + + tp_sizes = _parse_int_csv(args.tp_sizes) + if external: + # Labels only: nothing is launched, so there is no GPU count to validate against. + tp_sizes = tp_sizes[:1] or [0] + else: + visible_gpus = int(torch.cuda.device_count()) + tp_sizes = [tp for tp in tp_sizes if tp >= 1 and tp <= visible_gpus] + if not tp_sizes: + raise RuntimeError( + f"No tp sizes are runnable with visible_gpus={visible_gpus}. " + "Set CUDA_VISIBLE_DEVICES accordingly." + ) + + concurrencies = _parse_int_csv(args.concurrencies) + concurrencies = [c for c in concurrencies if c >= 1] + if not concurrencies: + raise RuntimeError("No concurrencies specified.") + + num_questions_by_conc = { + c: min( + int(args.questions_per_concurrency_base) * int(c), + int(args.max_questions_per_config), + ) + for c in concurrencies + } + max_questions = max(num_questions_by_conc.values()) + + attention_backends = [ + s.strip() for s in args.attention_backends.split(",") if s.strip() + ] + if external: + # Label only; the running server already picked its backend. + device_sm = 0 + attention_backends = attention_backends[:1] or ["external"] + else: + device_sm = get_device_sm() + attention_backends = _filter_attention_backends( + attention_backends, device_sm=device_sm + ) + + data_path = _maybe_download_mtbench(args.data_path) + lines = list(read_jsonl(data_path)) + if len(lines) < max_questions: + raise RuntimeError( + f"MT-Bench file only has {len(lines)} lines, but need {max_questions}." + ) + + tokenizer = AutoTokenizer.from_pretrained(args.target_model) + + prompts: list[str] = [] + for i in range(max_questions): + # MT-Bench items store a list of turns; use the first-turn prompt. + turns = lines[i]["turns"] + user_content = turns[0] + prompts.append( + tokenizer.apply_chat_template( + [{"role": "user", "content": user_content}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + ) + + # Results indexed by (backend, tp, concurrency, mode). + results: dict[tuple[str, int, int, str], BenchMetrics] = {} + # Baseline metrics are backend-agnostic in this sweep; run once per TP and reuse. + baseline_cache_by_tp: dict[int, dict[int, BenchMetrics]] = {} + + if external: + base_url = args.base_url.rstrip("/") + backend = attention_backends[0] + tp = tp_sizes[0] + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"external server {base_url} (mode={args.external_mode})", + model_path=args.target_model, + base_url=base_url, + server_args=[], + expect_dflash=(args.external_mode == "dflash"), + prompts=prompts, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + launch_server=False, + ) + for conc, metrics in mode_metrics.items(): + mode_key = args.external_mode + if mode_key == "auto": + mode_key = "dflash" if metrics.spec_verify_ct_sum > 0 else "baseline" + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + return + + for backend_idx, backend in enumerate(attention_backends): + for tp in tp_sizes: + port_base = find_available_port(20000) + common_server_args = _build_common_server_args(args, backend=backend, tp=tp) + mode_runs = _build_mode_runs(args, common_server_args) + + for idx, ( + mode_key, + mode_name, + mode_server_args, + expect_dflash, + ) in enumerate(mode_runs): + if ( + mode_key == "baseline" + and not args.skip_baseline + and backend_idx > 0 + and tp in baseline_cache_by_tp + ): + mode_metrics = baseline_cache_by_tp[tp] + else: + mode_metrics = _run_mode_for_backend_tp( + mode_label=f"backend={backend} tp={tp} ({mode_name})", + model_path=args.target_model, + base_url=f"http://127.0.0.1:{find_available_port(port_base + idx)}", + server_args=mode_server_args, + expect_dflash=expect_dflash, + prompts=prompts, + concurrencies=concurrencies, + num_questions_by_conc=num_questions_by_conc, + args=args, + ) + if mode_key == "baseline" and not args.skip_baseline: + baseline_cache_by_tp[tp] = mode_metrics + + for conc, metrics in mode_metrics.items(): + results[(backend, tp, conc, mode_key)] = metrics + + _print_summary( + args=args, + attention_backends=attention_backends, + tp_sizes=tp_sizes, + concurrencies=concurrencies, + device_sm=device_sm, + results=results, + ) + + +if __name__ == "__main__": + main()