From dcdde0bf7f89ba8dd9eda0cbe558e039f0795da7 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 00:44:42 +0800 Subject: [PATCH 1/4] benchmarks: move adaptive prefill calibration to Examples Signed-off-by: SuperGoodGame <985236470@qq.com> --- benchmarks/opsd/README.md | 51 ++++ benchmarks/opsd/rollout_prefill.py | 305 ++++++++++++++++++++++++ benchmarks/opsd/test_rollout_prefill.py | 21 ++ 3 files changed, 377 insertions(+) create mode 100644 benchmarks/opsd/rollout_prefill.py create mode 100644 benchmarks/opsd/test_rollout_prefill.py diff --git a/benchmarks/opsd/README.md b/benchmarks/opsd/README.md index e58be3800..0fdb8a557 100644 --- a/benchmarks/opsd/README.md +++ b/benchmarks/opsd/README.md @@ -51,3 +51,54 @@ torchrun --nproc_per_node=1 \ Use `--help` for the complete argument list. Model download, GPU memory, and the fused inference kernels supported by the selected model can limit which matrix shapes run successfully. + +## Adaptive prefill calibration + +`rollout_prefill.py` calibrates the opt-in adaptive prefill policy and compares +it with ordinary generation, equal-width CB and aligned CB. It requires a +DeepSpeed checkout containing [DeepSpeed PR #8786](https://github.com/deepspeedai/DeepSpeed/pull/8786), +installed in the active Python environment, plus PyTorch, Transformers and NumPy. +The scope is one GPU, greedy decoding and SDPA, with an unsharded HuggingFace +model behind the rollout interface; it does not measure a full training step. + +From the DeepSpeedExamples root: + +```sh +python benchmarks/opsd/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ + --repeats 3 --calibrate costs.json +python benchmarks/opsd/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ + --cost-config costs.json --repeats 3 --output sweep.json +``` + +Use `--model tiny-qwen2 --dtype float32` for a seeded smoke model. Calibration +uses six uniform shapes, distinct from the four mixed-length evaluation cases, +to fit non-negative fixed/token/attention costs and estimate KV and decode costs. +The built-in calibration grid reaches 768 input tokens. Extend the `shapes` +list in `calibrate()` to cover the deployment's actual lengths and feasible batch +sizes before using the estimates for longer contexts. Coefficients are specific +to the model, GPU, precision and attention backend; the sweep rejects mismatched +configuration metadata. + +Calibration runs during deployment preparation; load its saved coefficients +once when constructing the rollout. It does not profile or refit on each request: + +```python +import json +from pathlib import Path +from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRolloutConfig + +costs = json.loads(Path("costs.json").read_text())["costs"] +cfg = HybridEngineRolloutConfig(adaptive_prefill=True, **costs) +``` + +The sweep reports raw timings, selected routes, output-ID agreement, +`auto_vs_best` and the worst regression. Its default two-token, EOS-disabled +generation isolates routing overhead; it is not a production rollout trace. +Use `--new-tokens` for other generation budgets and validate held-out workloads +representative of the deployment. Timing excludes loading and tokenization. + +The fitting regression test needs no GPU: + +```sh +python -m unittest discover -s benchmarks/opsd -p test_rollout_prefill.py +``` diff --git a/benchmarks/opsd/rollout_prefill.py b/benchmarks/opsd/rollout_prefill.py new file mode 100644 index 000000000..b463529b1 --- /dev/null +++ b/benchmarks/opsd/rollout_prefill.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Calibrate prefill costs separately from a small, reproducible route sweep.""" + +import argparse +from dataclasses import asdict +import hashlib +from itertools import combinations +import json +from pathlib import Path +import statistics +import time +from types import SimpleNamespace + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config, Qwen2ForCausalLM + +from deepspeed.runtime.rollout.base import RolloutRequest, SamplingConfig +from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout, HybridEngineRolloutConfig +import deepspeed.runtime.rollout.hybrid_engine_rollout as rollout_backend +from deepspeed.accelerator import get_accelerator + + +def nonnegative_fit(matrix, targets): + """Solve three-parameter NNLS by checking its eight possible active sets.""" + matrix = np.asarray(matrix, dtype=float) + targets = np.asarray(targets, dtype=float) + scale = np.maximum(np.abs(matrix).max(axis=0), 1e-12) + scaled = matrix / scale + best = np.zeros(matrix.shape[1]) + error = float(targets @ targets) + for count in range(1, matrix.shape[1] + 1): + for columns in combinations(range(matrix.shape[1]), count): + values = np.linalg.lstsq(scaled[:, columns], targets, rcond=None)[0] + if np.any(values < 0): + continue + candidate = np.zeros(matrix.shape[1]) + candidate[list(columns)] = values + residual = scaled @ candidate - targets + if residual @ residual < error: + best, error = candidate, float(residual @ residual) + return (best / scale).tolist() + + +def make_request(lengths, token, pad, device): + width = max(lengths) + ids = torch.full((len(lengths), width), pad, dtype=torch.long, device=device) + mask = torch.zeros_like(ids) + for row, length in enumerate(lengths): + ids[row, -length:] = token + mask[row, -length:] = 1 + return RolloutRequest(ids, mask) + + +def measure(rollout, request, steps, continuous, repeats): + sampling = SamplingConfig(max_new_tokens=steps, + temperature=0, + continuous_batch_size=request.prompt_ids.shape[0] if continuous else None) + samples = [] + for iteration in range(-1, repeats): + get_accelerator().synchronize() + start = time.perf_counter() + output = rollout.generate(request, sampling) + get_accelerator().synchronize() + if iteration >= 0: + samples.append((time.perf_counter() - start) * 1000) + return samples, output + + +def calibrate(engine, tokenizer, token, pad, device, repeats): + model = engine.module + config = model.config.get_text_config() + head_dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads + attention_scale = 4 * config.num_hidden_layers * config.num_attention_heads * head_dim / 1e9 + kv_scale = 2 * config.num_hidden_layers * config.num_key_value_heads * head_dim * next( + model.parameters()).element_size() / 2**20 + matrix, times, copy_rates, records = [], [], [], [] + # These calibration shapes differ from every held-out sweep workload below. + shapes = [(1, 32), (1, 256), (4, 64), (4, 384), (8, 768), (16, 96)] + for batch_size, width in shapes: + request = make_request([width] * batch_size, token, pad, device) + cfg = HybridEngineRolloutConfig(adaptive_prefill=True, + enable_profiling=True, + continuous_cache_capacity=width + 1, + prefill_fixed_cost_ms=1e9, + prefill_token_cost_ms=0, + prefill_attention_cost_ms=0, + prefill_kv_cost_ms=0) + rollout = HybridEngineRollout(engine, tokenizer, cfg) + profiles = [] + for iteration in range(-1, repeats): + rollout.generate(request, SamplingConfig(max_new_tokens=1, temperature=0, + continuous_batch_size=batch_size)) + if iteration >= 0: + profiles.append(rollout.get_last_profile()) + prefill = statistics.mean(p['prefill_forward_ms'] for p in profiles) + copying = statistics.mean(p['cache_management_overhead_ms'] for p in profiles) + assert all(p['num_prefill_forwards'] == 1 for p in profiles) + matrix.append([1, batch_size * width, attention_scale * batch_size * width**2]) + times.append(prefill) + copy_rates.append(copying / (kv_scale * batch_size * width)) + records.append({ + 'batch_size': batch_size, + 'width': width, + 'prefill_ms': prefill, + 'cache_management_ms': copying + }) + print('CALIBRATE', batch_size, width, round(prefill, 3), 'ms', flush=True) + fixed, linear, attention = nonnegative_fit(matrix, times) + # Differences across decode budgets cancel most prefill and allocation costs. + request = make_request([80] * 16, token, pad, device) + differences = {} + for steps in (1, 8): + cb_cfg = HybridEngineRolloutConfig(adaptive_prefill=True, + continuous_cache_capacity=80 + steps, + prefill_fixed_cost_ms=1e9, + prefill_token_cost_ms=0, + prefill_attention_cost_ms=0, + prefill_kv_cost_ms=0) + cb = HybridEngineRollout(engine, tokenizer, cb_cfg) + ordinary = HybridEngineRollout(engine, tokenizer) + cb_times, _ = measure(cb, request, steps, True, repeats) + ordinary_times, _ = measure(ordinary, request, steps, False, repeats) + differences[steps] = statistics.mean(cb_times) - statistics.mean(ordinary_times) + costs = { + 'prefill_fixed_cost_ms': fixed, + 'prefill_token_cost_ms': linear, + 'prefill_attention_cost_ms': attention, + 'prefill_kv_cost_ms': statistics.median(copy_rates), + 'continuous_decode_cost_ms': max(0, (differences[8] - differences[1]) / 7) + } + return costs, { + 'prefill_shapes': records, + 'decode_cb_minus_no_cb_ms': differences, + 'note': 'KV rate includes cache management overhead; coefficients are configuration-specific estimates.' + } + + +def sweep(engine, tokenizer, token, pad, device, steps, repeats, costs): + workloads = { + 'short128': [16] * 128, + 'moderate_tail': [128] + [4] * 31, + 'spread': [16 * (i + 1) for i in range(32)], + 'long_tail': [512] + [16] * 31 + } + report = {} + modes = ['no_cb', 'equal_width', 'aligned', 'auto'] + for name, lengths in workloads.items(): + request = make_request(lengths, token, pad, device) + cfgs = { + mode: + HybridEngineRolloutConfig(adaptive_prefill=mode == 'auto', + align_decode_fronts=mode == 'aligned', + **(costs if mode == 'auto' else {})) + for mode in modes + } + rollouts = {mode: HybridEngineRollout(engine, tokenizer, cfg) for mode, cfg in cfgs.items()} + samples = {mode: [] for mode in modes} + outputs = {} + for iteration in range(-1, repeats): + shift = iteration % len(modes) + for mode in modes[shift:] + modes[:shift]: + rollout = rollouts[mode] + sampling = SamplingConfig(max_new_tokens=steps, + temperature=0, + continuous_batch_size=None if mode == 'no_cb' else len(lengths)) + get_accelerator().synchronize() + start = time.perf_counter() + output = rollout.generate(request, sampling) + get_accelerator().synchronize() + if iteration >= 0: + samples[mode].append((time.perf_counter() - start) * 1000) + outputs[mode] = output.input_ids[:, max(lengths):].cpu().tolist() + means = {mode: statistics.mean(values) for mode, values in samples.items()} + best = min(modes[:-1], key=means.get) + stats = rollouts['auto'].get_last_continuous_stats() + report[name] = { + 'lengths': lengths, + 'max_new_tokens': steps, + 'mean_ms': means, + 'raw_ms': samples, + 'best_baseline': best, + 'auto_vs_best': means['auto'] / means[best], + 'auto_strategy': 'batched' if stats is None else 'bucketed', + 'output_ids_match': all(value == outputs['no_cb'] for value in outputs.values()), + 'output_match_by_path': { + mode: value == outputs['no_cb'] + for mode, value in outputs.items() + }, + 'outputs': outputs + } + rollouts['auto'].enable_profiling = True + rollouts['auto'].generate( + request, SamplingConfig(max_new_tokens=steps, temperature=0, continuous_batch_size=len(lengths))) + report[name]['auto_profile'] = rollouts['auto'].get_last_profile() + print('SWEEP', + name, + means, + 'auto/best', + round(report[name]['auto_vs_best'], 4), + report[name]['auto_strategy'], + flush=True) + return report + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--model', default='tiny-qwen2') + parser.add_argument('--dtype', choices=['float32', 'bfloat16'], default='float32') + parser.add_argument('--new-tokens', type=int, default=2) + parser.add_argument('--repeats', type=int, default=5) + parser.add_argument('--calibrate', + type=Path, + help='Write coefficients using separate calibration shapes, then exit.') + parser.add_argument('--cost-config', type=Path) + parser.add_argument('--output', type=Path, default=Path('prefill-sweep.json')) + args = parser.parse_args() + if args.repeats <= 0 or args.new_tokens <= 0: + parser.error('repeats and new-tokens must be positive') + torch.manual_seed(8497) + torch.set_num_threads(4) + device = get_accelerator().device_name() + dtype = getattr(torch, args.dtype) + if args.model == 'tiny-qwen2': + config = Qwen2Config(vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=4096, + pad_token_id=0, + eos_token_id=None, + attn_implementation='sdpa') + model = Qwen2ForCausalLM(config).to(device=device, dtype=dtype).eval() + token, pad = 3, 0 + else: + native_tokenizer = AutoTokenizer.from_pretrained(args.model) + token = native_tokenizer.encode('hello', add_special_tokens=False)[0] + pad = native_tokenizer.pad_token_id + if pad is None: + pad = native_tokenizer.eos_token_id + model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype, + attn_implementation='sdpa').to(device).eval() + tokenizer = SimpleNamespace(pad_token_id=pad, eos_token_id=None) + engine = SimpleNamespace(module=model) + source = Path(rollout_backend.__file__) + metadata = { + 'model': args.model, + 'model_config': json.loads(json.dumps(model.config.to_dict())), + 'dtype': args.dtype, + 'gpu': torch.cuda.get_device_name(0), #ignore-cuda + 'sm_count': torch.cuda.get_device_properties(0).multi_processor_count, #ignore-cuda + 'torch': torch.__version__, + 'transformers': __import__('transformers').__version__, + 'backend': 'sdpa', + 'seed': 8497, + 'cpu_threads': 4, + 'allow_tf32': torch.backends.cuda.matmul.allow_tf32, + 'source_sha256': hashlib.sha256(source.read_bytes()).hexdigest(), + 'scope': + 'Controlled EOS-disabled generation; includes routing and cache work; excludes loading and tokenization.' + } + with torch.inference_mode(): + if args.calibrate: + costs, measurements = calibrate(engine, tokenizer, token, pad, device, args.repeats) + args.calibrate.parent.mkdir(parents=True, exist_ok=True) + args.calibrate.write_text( + json.dumps({ + 'metadata': metadata, + 'costs': costs, + 'calibration': measurements + }, indent=2) + '\n') + print('COSTS', json.dumps(costs), flush=True) + return + calibration = json.loads(args.cost_config.read_text()) if args.cost_config else {} + if 'metadata' in calibration: + for key in ('model', 'model_config', 'dtype', 'gpu', 'sm_count', 'backend', 'torch', 'transformers', + 'allow_tf32'): + if calibration['metadata'][key] != metadata[key]: + parser.error(f'calibration differs in {key}; recalibrate for this configuration') + costs = calibration.get('costs', {}) + # A supplied JSON uses the same validated public configuration fields as an application. + cfg = HybridEngineRolloutConfig(adaptive_prefill=True, **costs) + report = { + 'metadata': metadata, + 'auto_config': asdict(cfg), + 'repeats': args.repeats, + 'workloads': sweep(engine, tokenizer, token, pad, device, args.new_tokens, args.repeats, costs) + } + worst = max(report['workloads'].items(), key=lambda pair: pair[1]['auto_vs_best']) + report['worst_case'] = { + 'workload': worst[0], + 'auto_vs_best': worst[1]['auto_vs_best'], + 'regression_percent': (worst[1]['auto_vs_best'] - 1) * 100 + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + '\n') + print('WORST', report['worst_case'], flush=True) + + +if __name__ == '__main__': + main() diff --git a/benchmarks/opsd/test_rollout_prefill.py b/benchmarks/opsd/test_rollout_prefill.py new file mode 100644 index 000000000..4a186cc39 --- /dev/null +++ b/benchmarks/opsd/test_rollout_prefill.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +import unittest + +import numpy as np + +from rollout_prefill import nonnegative_fit + + +class CalibrationFitTest(unittest.TestCase): + + def test_recovers_nonnegative_costs(self): + matrix = [[1, 1, 1], [1, 2, 4], [1, 4, 16], [1, 8, 64]] + np.testing.assert_allclose(nonnegative_fit(matrix, [5.5, 10, 22, 58]), [2, 3, 0.5]) + # Noisy measurements must not produce negative execution costs. + np.testing.assert_allclose(nonnegative_fit(matrix, [4, 3, 2, 1]), [2.5, 0, 0], atol=1e-12) + + +if __name__ == "__main__": + unittest.main() From 79b2140eaff9b5c316799e18a613165a8dbfa775 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 01:15:34 +0800 Subject: [PATCH 2/4] fix(benchmarks): support MHA calibration and verify backend sources Signed-off-by: SuperGoodGame <985236470@qq.com> --- benchmarks/opsd/README.md | 5 +- benchmarks/opsd/rollout_prefill.py | 14 +++-- benchmarks/opsd/test_rollout_prefill.py | 83 ++++++++++++++++++++++++- 3 files changed, 94 insertions(+), 8 deletions(-) diff --git a/benchmarks/opsd/README.md b/benchmarks/opsd/README.md index 0fdb8a557..feb79ec77 100644 --- a/benchmarks/opsd/README.md +++ b/benchmarks/opsd/README.md @@ -77,7 +77,8 @@ The built-in calibration grid reaches 768 input tokens. Extend the `shapes` list in `calibrate()` to cover the deployment's actual lengths and feasible batch sizes before using the estimates for longer contexts. Coefficients are specific to the model, GPU, precision and attention backend; the sweep rejects mismatched -configuration metadata. +configuration metadata, including a source hash covering rollout, bucket planning +and static-cache implementations. Recalibrate when those implementations change. Calibration runs during deployment preparation; load its saved coefficients once when constructing the rollout. It does not profile or refit on each request: @@ -97,7 +98,7 @@ generation isolates routing overhead; it is not a production rollout trace. Use `--new-tokens` for other generation budgets and validate held-out workloads representative of the deployment. Timing excludes loading and tokenization. -The fitting regression test needs no GPU: +Run the fitting, MHA calibration and saved-configuration regression tests: ```sh python -m unittest discover -s benchmarks/opsd -p test_rollout_prefill.py diff --git a/benchmarks/opsd/rollout_prefill.py b/benchmarks/opsd/rollout_prefill.py index b463529b1..a5c17e5d6 100644 --- a/benchmarks/opsd/rollout_prefill.py +++ b/benchmarks/opsd/rollout_prefill.py @@ -19,7 +19,9 @@ from deepspeed.runtime.rollout.base import RolloutRequest, SamplingConfig from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout, HybridEngineRolloutConfig +import deepspeed.runtime.rollout.continuous_batching as batching_backend import deepspeed.runtime.rollout.hybrid_engine_rollout as rollout_backend +import deepspeed.utils.static_cache as cache_backend from deepspeed.accelerator import get_accelerator @@ -74,7 +76,8 @@ def calibrate(engine, tokenizer, token, pad, device, repeats): config = model.config.get_text_config() head_dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads attention_scale = 4 * config.num_hidden_layers * config.num_attention_heads * head_dim / 1e9 - kv_scale = 2 * config.num_hidden_layers * config.num_key_value_heads * head_dim * next( + kv_heads = getattr(config, 'num_key_value_heads', None) or config.num_attention_heads + kv_scale = 2 * config.num_hidden_layers * kv_heads * head_dim * next( model.parameters()).element_size() / 2**20 matrix, times, copy_rates, records = [], [], [], [] # These calibration shapes differ from every held-out sweep workload below. @@ -246,7 +249,8 @@ def main(): attn_implementation='sdpa').to(device).eval() tokenizer = SimpleNamespace(pad_token_id=pad, eos_token_id=None) engine = SimpleNamespace(module=model) - source = Path(rollout_backend.__file__) + sources = (rollout_backend.__file__, batching_backend.__file__, cache_backend.__file__) + source_hash = hashlib.sha256(b'\0'.join(Path(source).read_bytes() for source in sources)).hexdigest() metadata = { 'model': args.model, 'model_config': json.loads(json.dumps(model.config.to_dict())), @@ -259,7 +263,7 @@ def main(): 'seed': 8497, 'cpu_threads': 4, 'allow_tf32': torch.backends.cuda.matmul.allow_tf32, - 'source_sha256': hashlib.sha256(source.read_bytes()).hexdigest(), + 'source_sha256': source_hash, 'scope': 'Controlled EOS-disabled generation; includes routing and cache work; excludes loading and tokenization.' } @@ -278,8 +282,8 @@ def main(): calibration = json.loads(args.cost_config.read_text()) if args.cost_config else {} if 'metadata' in calibration: for key in ('model', 'model_config', 'dtype', 'gpu', 'sm_count', 'backend', 'torch', 'transformers', - 'allow_tf32'): - if calibration['metadata'][key] != metadata[key]: + 'allow_tf32', 'source_sha256'): + if calibration['metadata'].get(key) != metadata[key]: parser.error(f'calibration differs in {key}; recalibrate for this configuration') costs = calibration.get('costs', {}) # A supplied JSON uses the same validated public configuration fields as an application. diff --git a/benchmarks/opsd/test_rollout_prefill.py b/benchmarks/opsd/test_rollout_prefill.py index 4a186cc39..165157da4 100644 --- a/benchmarks/opsd/test_rollout_prefill.py +++ b/benchmarks/opsd/test_rollout_prefill.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team +from contextlib import redirect_stdout, redirect_stderr +import io +import json +from pathlib import Path +import tempfile +from types import SimpleNamespace import unittest +from unittest.mock import patch import numpy as np +import torch -from rollout_prefill import nonnegative_fit +from rollout_prefill import calibrate, main, nonnegative_fit class CalibrationFitTest(unittest.TestCase): @@ -16,6 +24,79 @@ def test_recovers_nonnegative_costs(self): # Noisy measurements must not produce negative execution costs. np.testing.assert_allclose(nonnegative_fit(matrix, [4, 3, 2, 1]), [2.5, 0, 0], atol=1e-12) + def test_calibrates_multihead_model_without_kv_head_setting(self): + from deepspeed.accelerator import get_accelerator + from transformers import OPTConfig, OPTForCausalLM + + config = OPTConfig(vocab_size=32, + hidden_size=32, + ffn_dim=64, + num_hidden_layers=2, + num_attention_heads=4, + max_position_embeddings=1024, + pad_token_id=0, + eos_token_id=None, + attn_implementation="sdpa") + self.assertFalse(hasattr(config, "num_key_value_heads")) + device = get_accelerator().device_name() + model = OPTForCausalLM(config).to(device).eval() + threads = torch.get_num_threads() + try: + torch.set_num_threads(4) + with torch.inference_mode(), redirect_stdout(io.StringIO()): + costs, measurements = calibrate(SimpleNamespace(module=model), + SimpleNamespace(pad_token_id=0, eos_token_id=None), + 3, + 0, + device, + repeats=1) + finally: + torch.set_num_threads(threads) + self.assertEqual(len(measurements["prefill_shapes"]), 6) + self.assertTrue(all(np.isfinite(value) and value >= 0 for value in costs.values())) + + def test_cli_rejects_stale_or_missing_source_hash(self): + with tempfile.TemporaryDirectory() as directory: + costs_path = Path(directory) / "costs.json" + output_path = Path(directory) / "sweep.json" + # Isolate file compatibility from measurements; the model/config remain real. + with patch("rollout_prefill.get_accelerator", return_value=SimpleNamespace(device_name=lambda: "cpu")), \ + patch("torch.cuda.get_device_name", return_value="test-device"), \ + patch("torch.cuda.get_device_properties", return_value=SimpleNamespace(multi_processor_count=1)), \ + patch("rollout_prefill.calibrate", return_value=({}, {})), \ + patch("rollout_prefill.sweep", return_value={"smoke": {"auto_vs_best": 1.0}}) as sweep, \ + redirect_stdout(io.StringIO()): + with patch("sys.argv", ["rollout_prefill.py", "--calibrate", str(costs_path)]): + main() + saved = json.loads(costs_path.read_text()) + args = ["rollout_prefill.py", "--cost-config", str(costs_path), "--output", str(output_path)] + with patch("sys.argv", args): + main() + self.assertEqual(sweep.call_count, 1) + for source_hash in ("0" * 64, None): + with self.subTest(source_hash=source_hash): + calibration = json.loads(json.dumps(saved)) + if source_hash is None: + calibration["metadata"].pop("source_sha256") + else: + calibration["metadata"]["source_sha256"] = source_hash + costs_path.write_text(json.dumps(calibration)) + stderr = io.StringIO() + with patch("sys.argv", args), redirect_stderr(stderr), self.assertRaises(SystemExit) as error: + main() + self.assertEqual(error.exception.code, 2) + self.assertIn("source_sha256", stderr.getvalue()) + self.assertEqual(sweep.call_count, 1) + costs_path.write_text(json.dumps(saved)) + changed_cache = Path(directory) / "changed_cache.py" + changed_cache.write_text("# A different static-cache implementation\n") + with patch("rollout_prefill.cache_backend.__file__", str(changed_cache)), \ + patch("sys.argv", args), redirect_stderr(io.StringIO()), \ + self.assertRaises(SystemExit) as error: + main() + self.assertEqual(error.exception.code, 2) + self.assertEqual(sweep.call_count, 1) + if __name__ == "__main__": unittest.main() From beb3401dceabfa9342c4fce313b004fc8fdb7032 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 01:46:33 +0800 Subject: [PATCH 3/4] fix(benchmarks): require metadata for saved calibration Signed-off-by: SuperGoodGame <985236470@qq.com> --- benchmarks/opsd/README.md | 2 ++ benchmarks/opsd/rollout_prefill.py | 2 ++ benchmarks/opsd/test_rollout_prefill.py | 5 +++++ 3 files changed, 9 insertions(+) diff --git a/benchmarks/opsd/README.md b/benchmarks/opsd/README.md index feb79ec77..7ff1bfe75 100644 --- a/benchmarks/opsd/README.md +++ b/benchmarks/opsd/README.md @@ -79,6 +79,8 @@ sizes before using the estimates for longer contexts. Coefficients are specific to the model, GPU, precision and attention backend; the sweep rejects mismatched configuration metadata, including a source hash covering rollout, bucket planning and static-cache implementations. Recalibrate when those implementations change. +`--cost-config` requires a complete calibration metadata object. Applications +may still pass hand-tuned coefficients directly to `HybridEngineRolloutConfig`. Calibration runs during deployment preparation; load its saved coefficients once when constructing the rollout. It does not profile or refit on each request: diff --git a/benchmarks/opsd/rollout_prefill.py b/benchmarks/opsd/rollout_prefill.py index a5c17e5d6..c7d2cd791 100644 --- a/benchmarks/opsd/rollout_prefill.py +++ b/benchmarks/opsd/rollout_prefill.py @@ -280,6 +280,8 @@ def main(): print('COSTS', json.dumps(costs), flush=True) return calibration = json.loads(args.cost_config.read_text()) if args.cost_config else {} + if args.cost_config and not isinstance(calibration.get('metadata'), dict): + parser.error('cost-config requires calibration metadata; rerun --calibrate') if 'metadata' in calibration: for key in ('model', 'model_config', 'dtype', 'gpu', 'sm_count', 'backend', 'torch', 'transformers', 'allow_tf32', 'source_sha256'): diff --git a/benchmarks/opsd/test_rollout_prefill.py b/benchmarks/opsd/test_rollout_prefill.py index 165157da4..63246ff31 100644 --- a/benchmarks/opsd/test_rollout_prefill.py +++ b/benchmarks/opsd/test_rollout_prefill.py @@ -87,6 +87,11 @@ def test_cli_rejects_stale_or_missing_source_hash(self): self.assertEqual(error.exception.code, 2) self.assertIn("source_sha256", stderr.getvalue()) self.assertEqual(sweep.call_count, 1) + costs_path.write_text(json.dumps({"costs": saved["costs"]})) + with patch("sys.argv", args), redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as error: + main() + self.assertEqual(error.exception.code, 2) + self.assertEqual(sweep.call_count, 1) costs_path.write_text(json.dumps(saved)) changed_cache = Path(directory) / "changed_cache.py" changed_cache.write_text("# A different static-cache implementation\n") From 3fd685fc557a81e03f8c76ea9cbbf89e4b047696 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 02:09:47 +0800 Subject: [PATCH 4/4] docs(benchmarks): specify minimum Transformers version Signed-off-by: SuperGoodGame <985236470@qq.com> --- benchmarks/opsd/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/opsd/README.md b/benchmarks/opsd/README.md index 7ff1bfe75..b6083b926 100644 --- a/benchmarks/opsd/README.md +++ b/benchmarks/opsd/README.md @@ -57,7 +57,7 @@ matrix shapes run successfully. `rollout_prefill.py` calibrates the opt-in adaptive prefill policy and compares it with ordinary generation, equal-width CB and aligned CB. It requires a DeepSpeed checkout containing [DeepSpeed PR #8786](https://github.com/deepspeedai/DeepSpeed/pull/8786), -installed in the active Python environment, plus PyTorch, Transformers and NumPy. +installed in the active Python environment, plus PyTorch, Transformers >=4.56 and NumPy. The scope is one GPU, greedy decoding and SDPA, with an unsharded HuggingFace model behind the rollout interface; it does not measure a full training step.