perf(rollout): select generation and prefill buckets with a tunable cost model

[SuperGoodGame]

Summary

Closes #8785. Related tracking: #8197, #8462.

Performance follow-up to #8497. Aligned CB avoids padding on long-tail prompts, but one prefill Forward per request is expensive for short prompts. A padding-ratio threshold cannot distinguish 128/4 from 512/16, despite identical ratios and different fastest baselines.

This change adds one opt-in automatic policy that compares ordinary no-CB generation with aligned CB using DP-selected prefill buckets. For 1×512 + 31×16, the selected buckets are [1,512] and [31,16], reducing prefill Forwards from 32 to 2.

Auto is disabled by default. Existing manual modes and positional configuration arguments are preserved.

cfg = HybridEngineRolloutConfig(adaptive_prefill=True)
rollout = HybridEngineRollout(engine, tokenizer, cfg)
sampling = SamplingConfig(max_new_tokens=16, temperature=0, continuous_batch_size=32)

Selection and configuration

A length-sorted O(n²)-time / O(n)-space DP minimizes the additive estimated cost of contiguous buckets. It explores one batch, singleton batches and intermediate partitions without running the GPU for candidate evaluation. Each bucket uses ordinary left padding and a 2D attention mask; no varlen kernel is needed. Model geometry is read once per plan.

C(B,L) = fixed + token × B×L + attention × estimated_attention_GFLOPs
         + kv × estimated_KV_MiB
attention_GFLOPs = 4 × B × L² × layers × query_heads × head_dim / 1e9
KV_MiB = 2 × B × L × layers × KV_heads × head_dim × element_bytes / 2^20

Parameter	Default	Meaning
prefill_fixed_cost_ms	25.0	Fixed milliseconds per prefill Forward
prefill_token_cost_ms	0.1	Milliseconds per padded token position; projection/MLP term
prefill_attention_cost_ms	0.01	Milliseconds per estimated attention GFLOP
prefill_kv_cost_ms	0.04	Milliseconds per MiB of KV copied into the CB cache
continuous_decode_cost_ms	10.0	Additional CB milliseconds per remaining Decode step
prefill_max_tokens	65536	Maximum padded token positions per prefill Forward; None removes this limit

All cost coefficients must be finite and non-negative. These are starting estimates, not universal GPU constants; calibrate them for the model, device, precision and attention backend. Head dimension, query/KV heads, layers and parameter element size are read automatically. The DP is optimal for the configured estimates, not guaranteed globally optimal measured latency.

Ordinary generation is eligible when every row fits continuous_batch_size, the physical input fits the prefill/context limits, and neither explicit static capacity nor trimming is requested. Its estimated cost omits CB cache-copy traffic; CB adds continuous_decode_cost_ms × (max_new_tokens - 1). The first output token comes from Prefill.

For CB, every admission/refill replans the newly admitted requests. Effective KV spans are copied into the existing right-aligned cache; surviving rows retain their layout. The policy does not change caller configuration or active-row bounds. Disabling Auto restores manual selection. Model/OOM errors propagate without trying another path.

Correctness

• Apply repetition penalty from the first token using original padded history, even when bucket model inputs are trimmed.
• Preserve generated PAD-valued tokens through the first EOS, excluding structural padding after termination.
• Include route planning and result normalization in end-to-end profiling; count productive output tokens and retain active/configured capacity fields.
• Resolve HF's effective generation configuration before rejecting unsupported processors/output modes, including legacy model-config refresh; preserve explicit generation-config overrides. Auto requires single-process greedy rollout, one sample per prompt and cache-class support; existing CUDA Graph/shared-prefill restrictions remain.

Experiments

Related work: TurboTransformers, §5 (PPoPP 2021) already uses length sorting, profiled batch costs and an O(n²) DP to partition variable-length inference requests. DynaPipe, §4 applies cost-based DP partitioning to training micro-batches. This PR applies that established scheduling approach to HybridEngine rollout, with ordinary-generation comparison, refill planning and existing aligned-cache semantics; it makes no claim of a new DP algorithm.

Single GPU; Qwen/Qwen2.5-3B-Instruct; FP32 / SDPA; PyTorch 2.9.1+cu128; Transformers 4.57.1. Each candidate receives one warmup and five interleaved timed runs, 32 requests and 16 output tokens per request. CB capacity is 32; trimming is disabled, cache capacity unset, TF32 disabled. The default cost coefficients above are unchanged across workloads.

Workload	Effective prompt lengths	Reason
Short, similar	20..13 repeated four times	Normal workload shape from #8497
Long tail	1×512 + 31×16	Long-tail workload from #8497
Short tail	1×20 + 31×1	Counterexample to ratio-only routing
Moderate tail	1×128 + 31×4	Same padding ratio as 512/16, different best legacy path
Spread	16, 32, …, 512	General distribution with intermediate buckets

Inputs repeat non-special token ID 14990. #8497 did not publish complete short-prompt sequences, so that case reconstructs its documented shape. Every timed response matches independent single-request eager generation; each workload produces 512 effective output tokens.

GPU-synchronized wall-clock timing includes the full generate() call, including Auto planning and normalization. Model loading and oracle work are outside timing; profiling is disabled during timed runs. All reported E2E values are arithmetic means of five runs.

Workload	No-CB (ms)	Equal-width CB (ms)	Aligned CB (ms)	Auto (ms)	Auto / best	Auto choice
Short, similar	505.54	673.66	1549.74	501.77	0.9925×	No-CB
1×512 + 31×16	2458.46	2596.38	1682.10	840.06	0.4994×	2 buckets
1×20 + 31×1	501.74	667.57	1331.82	497.57	0.9917×	No-CB
1×128 + 31×4	897.17	1053.26	1468.58	634.59	0.7073×	2 buckets
16, 32, …, 512	2458.13	2595.33	2337.75	1984.86	0.8490×	6 buckets

All five workloads meet the 5% acceptance tolerance. Auto chooses ordinary generation for both short cases; their small timing differences are not claimed as a kernel speedup. Relative to the fastest legacy baseline, mean latency falls by 50.06% for 1×512 + 31×16, 29.27% for 1×128 + 31×4, 15.10% for 16, 32, …, 512. These improvements come from fewer prefill Forwards and reduced padded work.

Throughput and peak allocated memory for every path

Throughput is 512 divided by mean latency. Memory is peak PyTorch allocated memory during timed calls, including model weights.

Workload	Path	Output tok/s	Peak allocated memory (MiB)
Short, similar	No-CB	1012.77	11988.96
Short, similar	Equal-width CB	760.03	12036.75
Short, similar	Aligned CB	330.38	11928.77
Short, similar	Auto	1020.40	11988.96
1×512 + 31×16	No-CB	208.26	15558.69
1×512 + 31×16	Equal-width CB	197.20	16747.20
1×512 + 31×16	Aligned CB	304.38	13278.51
1×512 + 31×16	Auto	609.48	13278.49
1×20 + 31×1	No-CB	1020.46	11994.64
1×20 + 31×1	Equal-width CB	766.96	12033.50
1×20 + 31×1	Aligned CB	384.44	11925.07
1×20 + 31×1	Auto	1029.00	11994.64
1×128 + 31×4	No-CB	570.68	12742.62
1×128 + 31×4	Equal-width CB	486.11	13066.75
1×128 + 31×4	Aligned CB	348.64	12221.97
1×128 + 31×4	Auto	806.82	12221.95
16, 32, …, 512	No-CB	208.29	15558.69
16, 32, …, 512	Equal-width CB	197.28	16747.20
16, 32, …, 512	Aligned CB	219.01	13278.56
16, 32, …, 512	Auto	257.95	13682.33

Separate profiling breakdown

One additional profiling call per candidate, outside acceptance timing. These diagnostics are not averages. Cache management is part of the generation overhead, not an additional independent total.

Workload	Path	Prefill (ms)	Decode (ms)	Cache management (ms)	Prefill / Decode Forwards
Short, similar	No-CB	91.01	401.30	N/A	1 / 15
Short, similar	Equal-width CB	91.99	499.53	26.80	1 / 15
Short, similar	Aligned CB	955.52	500.30	35.25	32 / 15
Short, similar	Auto	90.74	396.90	N/A	1 / 15
1×512 + 31×16	No-CB	1918.22	531.89	N/A	1 / 15
1×512 + 31×16	Equal-width CB	1919.76	595.14	27.09	1 / 15
1×512 + 31×16	Aligned CB	981.04	594.52	35.02	32 / 15
1×512 + 31×16	Auto	153.66	592.66	31.43	2 / 15
1×20 + 31×1	No-CB	91.06	402.45	N/A	1 / 15
1×20 + 31×1	Equal-width CB	91.84	494.84	26.88	1 / 15
1×20 + 31×1	Aligned CB	743.30	491.93	34.48	32 / 15
1×20 + 31×1	Auto	90.62	394.78	N/A	1 / 15
1×128 + 31×4	No-CB	490.08	399.13	N/A	1 / 15
1×128 + 31×4	Equal-width CB	490.85	482.08	27.09	1 / 15
1×128 + 31×4	Aligned CB	889.01	483.21	34.83	32 / 15
1×128 + 31×4	Auto	63.80	485.00	31.43	2 / 15
16, 32, …, 512	No-CB	1918.03	531.64	N/A	1 / 15
16, 32, …, 512	Equal-width CB	1918.95	594.68	27.32	1 / 15
16, 32, …, 512	Aligned CB	1648.16	593.25	35.91	32 / 15
16, 32, …, 512	Auto	1301.59	593.31	32.49	6 / 15

The 128/4 and 512/16 workloads reduce prefill Forwards from 32 on aligned CB to 2 on Auto. Spread uses 6 buckets. All CB paths in these performance runs have zero trims and zero moved bytes; ordinary generation has no static-cache statistics. This makes no new trimming-speedup claim.

Validation and limits

A separate microbenchmark measures planner overhead, without model Prefill: for 128/256 requests of 65,536 tokens, CPU DP takes 0.48/1.82 ms on average with the default per-bucket token limit. It still checks all 8,256/32,896 candidate segments; over-limit segments skip cost evaluation. For 32 requests, DP takes 0.089 ms at length512 and 0.040 ms at length65,536 (0.078 ms with no token limit). GPU mask validation/length extraction for 32×65,536 takes 0.83 ms separately. Measurements use ten repeated timing samples after warmup. These are not 64k model Prefill timings or the complete Auto-path overhead.

The original validation passed 105 related rollout/cache tests, including independent exhaustive partition comparisons for DP, device-backed tiny-Qwen tests for grouped prefill and refill, both cache APIs, effective generation-config/legacy refresh, repetition history, EOS/PAD masks, capacity bounds and external-wall-time checks for profiling. Existing tests cover retirement, compaction and trimming. Independent mixed-bucket checks also exercised live-survivor refill and trimming on/off against padded per-request HF generation.

One pre-existing native shared-prefill extension test is excluded because the host has no CUDA compiler; all added tests ran. Changed-file pre-commit and git diff --check pass. Repeated local review ended with no actionable findings.

python -m pytest -q \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \
  tests/unit/runtime/rollout/test_continuous_batching.py \
  tests/unit/runtime/rollout/test_rollout_interface.py \
  tests/unit/utils/test_static_cache.py \
  -k 'not native_repeat_kv_cache_fp16_reverse_copy'

The prefill token limit does not bound allocated decode KV memory. A single prompt over that limit raises an explicit error; there is no chunked prefill. 64k prompts require a model context limit covering prompt plus output budget, sufficient KV capacity and temporary memory. The measured checkpoint is configured for 32k. Planner/position-limit boundaries are tested; real 64k inference performance is not claimed.

Performance evidence uses unsharded HF models behind HybridEngineRollout; the additional calibrated and real-question checks are listed below. Other model/device configurations, FP16, native kernel injection, ZeRO/TP training integration and production arrival/EOS distributions remain unmeasured. Controlled performance runs do not exercise staggered EOS or trimming; those have correctness coverage.

No new dependency, kernel, online learner or trimming redesign. Related foundation: #8368; profiling: #8494.

Reviewer follow-up

Commit 65a9115 propagates the resolved generation configuration through continuous Prefill, Decode and refill, and omits output_logits from legacy manual generation. 27 targeted tests pass, including six new cases covering both Auto routes, both cache APIs and manual equal-width/aligned compatibility. The new tests reproduced four failures before the fix. The compatibility test uses a strict generation wrapper that rejects the unsupported keyword.

Calibration and evaluation tooling moved to DeepSpeedExamples #1018, with its fitting regression test and usage documentation. Merge the Examples companion before this PR so the documentation link is available. The Examples companion includes MHA-head fallback and validates a source fingerprint covering rollout, planning and static-cache code. A seeded tiny-model GPU smoke check completed calibration and all four saved-config sweep workloads with matching output IDs on every path. Two independent local reviews returned no findings.

The calibrated sweep from the preceding routing-overhead fix is separate from the original 16-token experiments above. It uses two EOS-disabled output tokens, SDPA and one GPU, with independent uniform calibration shapes and one unchanged coefficient set per configuration:

Configuration	128×16	1×128 + 31×4	16..512	1×512 + 31×16	Worst regression
Tiny Qwen2 / FP32	1.0220×	0.9685×	0.9539×	0.9647×	+2.205%
Qwen2.5-3B / FP32	1.0005×	0.2345×	0.7808×	0.2265×	+0.054%
Qwen3-32B / BF16	1.0027×	0.3881×	0.7189×	0.1811×	+0.273%

Values are Auto E2E / fastest manual baseline E2E. The two larger models select no-CB for 128×16 and bucketed CB for the other three cases; the tiny model selects no-CB throughout. Auto matches ordinary output IDs in all twelve cases. Target-model means use three interleaved repetitions after warmup; the tiny test uses ten. Full sweep details and calibration limits.

A separate Qwen3-32B real-question check selected no-CB on both workloads: MATH-500 averaged 31.452 s / 31.393 s for ordinary/Auto over three runs; ARC-AGI-2 measured 237.560 s / 237.888 s in one paired run. All eight responses matched. These checks cover complete generation calls, excluding tokenization, loading, rewards and training.

A real tiny-OPT calibration additionally exposed a missing HF cache capability flag. Commit ecacbce adds is_compileable=True to the static-cache container, matching its existing static layers. Six targeted GPU checks pass, including new OPT manual/Auto comparisons with padded prompts and refill against the HF oracle, plus existing Qwen bucket/cache-path checks. The three Examples regression tests cover fitting, complete OPT calibration and saved-source compatibility. Both local reviews found no issues in this follow-up.

Final correctness and performance acceptance

Commit d4def3d includes the final fixes: non-default repetition penalties are calculated in FP32 like HF, and unknown public generation switches are rejected before routing. The default-penalty path retains its early return. Five focused BF16 checks and six unknown-setting CPU cases pass. The Examples CLI now requires complete metadata for saved calibration files. Both independent correctness and performance reviews returned no findings after these fixes.

A new calibration and sweep were run against measured commit d4def3d. Qwen3-32B, BF16, SDPA, one GPU; two EOS-disabled output tokens; one warmup plus three interleaved timed runs per path. Calibration and held-out evaluation use separate shapes. The calibration file, sweep and that measured revision's rollout/planner/static-cache sources have the same combined SHA256: f9d79a43d2e6c0e3fef164ef11e81a26c643820beb1a0d3f1a86f41a46dbdd0c.

Workload	No-CB ms	Equal-width CB ms	Aligned CB ms	Auto ms	Auto / best	Auto path
128×16	269.83	439.91	5717.89	270.03	1.000743×	No-CB
1×128 + 31×4	496.15	540.76	1501.57	191.96	0.386886×	2 buckets
16,32,…,512	1962.87	2009.18	1864.55	1336.57	0.716828×	6 buckets
1×512 + 31×16	1962.56	2005.33	1575.08	288.31	0.183045×	2 buckets

All four pass the 5% acceptance tolerance; worst measured regression is +0.07428%. Auto output IDs match ordinary generation in every case, and both routes are exercised. Means and ratios were independently recomputed from raw samples. The legacy aligned spread baseline retains its previously recorded second-token difference on one row; Auto matches no-CB there. Fixed output length prevents this numerical-sensitive baseline difference from changing the amount of timed work.

The fitted coefficients are fixed=28.791679 ms, token=0.093362064 ms, attention=0.010090454 ms/GFLOP, KV=0.056288540 ms/MiB and additional decode=17.144541 ms/step. The earlier tables remain historical measurements on their original versions. This final check does not establish long-context, natural-EOS, 256-request E2E or dynamic-arrival performance. Auto remains disabled by default.

Commit a0588b7 then fixes manual aligned prefill to retain the original padded repetition history. Its focused GPU regression reproduced the mismatch and now passes. Auto already passed this same mapping, and the measured default-penalty paths do not read it; independent correctness and performance delta reviews returned no findings. The table remains attributed to d4def3d, with no timings or hashes relabeled as a new-HEAD measurement. The Examples benchmark explicitly requires Transformers >=4.56; the library's existing compatibility scope is unchanged.

Review invitation

cc @nathon-lee @delock — this follows your continuous-batching, profiling, and aligned-cache work in #8368, #8494, and #8497. Could you review the automatic route selection, bucket/refill cache semantics, and cost defaults?

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-10T18:19:47.126293Z	a0588b7	New commits

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7cb9828038

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[delock]

Hi @SuperGoodGame thanks for your PR. I have a question about estimation of prefill ratio in rollout e2e. What is typical rollout decode length and prompt length for RL rollout? I'm trying to understand the scope where this optimization will be most effective. @nathon-lee @LiRunGuo if you have any inputs.

[SuperGoodGame]

Thanks @delock. Prefill's share of rollout E2E depends on the checkpoint and task, including the prompt and response distributions. There is no single task-independent “typical” RL length.

Motivation: prompt lengths vary within a batch. Too few buckets waste padded computation; too many repeatedly pay per-forward/kernel-launch overhead. Cost-based DP balances these costs. The same calibrated policy adapts across successive batches, compares eligible no-CB with CB, and replans new admissions within CB. Coefficients are adjustable for the model, GPU, precision, and backend.

Additional measurements: same Qwen2.5-3B-Instruct FP32/SDPA setup, unchanged coefficients, 1×512 + 31×16 prompts. EOS was disabled to isolate output-length effects:

Output tokens/request	No-CB Prefill / E2E	Fastest manual baseline	Auto mean	Auto choice
128	29.2%	Aligned CB: 6.450 s	5.666 s	Bucketed CB
512	8.25%	No-CB: 23.270 s	23.067 s	No-CB

E2E values average three interleaved runs after warmup; fractions use a separate profiling call. Outputs matched independent per-request HF generation.

An EOS-enabled example with eight math/coding prompts (43–56 input tokens, 512-token cap) produced 2–512 output tokens: six ended at EOS, two hit the cap. Prefill was only 0.57% of no-CB E2E; Auto selected no-CB and matched it (12.836 s vs 12.987 s). These are controlled examples, not representative RL traces.

The largest gains are expected when padding waste and Prefill's time share are substantial. When Decode dominates, Auto can retain no-CB. Future EOS length is unknown, so selection uses the configured generation budget and estimated costs. The feature remains opt-in, single-process, greedy, and single-sample.

Here E2E covers the complete generation call, excluding tokenization, loading, reward computation, and training updates.

[nathon-lee]

Thanks for the detailed explanation and thorough benchmarking.

My main concern is keeping the default OPSD path simple and maintainable. I'd like to take some more time to evaluate this approach before making a decision.

I'll take a closer look and share my thoughts later.

[SuperGoodGame]

Thanks @nathon-lee for sharing your concern. I agree that keeping the default conservative makes sense. Auto is already disabled by default (adaptive_prefill=False), so this PR provides an opt-in optimization.

The goal of Auto is to retain baseline performance when bucketing offers little benefit, and capture the gains when it does. It evaluates each batch and selects ordinary batch generation when that is estimated to be the better choice.

Related community practice also exists: verl’s batching utilities implement token-budgeted dynamic micro-batching with sequence-length-based workload estimates. This provides a practical precedent for workload-aware batch composition, which this PR applies to rollout prefill as an opt-in policy.

I also ran a small real-question A/B check with Qwen3-32B, BF16, on a single GPU, using greedy decoding and aligned natural-EOS stopping behavior:

Workload	Ordinary no-CB	Auto
MATH-500 — mean of 3 runs	31.452 s	31.393 s
ARC-AGI-2 — single paired run	237.560 s	237.888 s

The timing differences were below 0.2%, and all eight responses were token-identical. Auto selected ordinary batch generation in both cases, as intended. These results complement the prefill-heavy improvements reported in the PR by checking baseline alignment on a larger model with real questions.

Would you consider this default-off, opt-in scope, with baseline alignment as the acceptance target, suitable for the PR? If further validation would help, could you suggest the specific workloads or baselines most relevant to your target deployments?

[nathon-lee]

Hi, @SuperGoodGame Thanks for the additional Qwen3-32B results. It is useful to see that Auto falls back to No-CB and remains close to the baseline on real questions. I agree that default-off meaningfully limits compatibility risk.

However, these two cases both select No-CB, so they do not yet validate the more important part of the policy: whether the fixed-cost model reliably selects bucketed CB when it predicts a benefit. I ran a small independent smoke test with the default coefficients on an RTX A4500 using a randomly initialized small Qwen2 model. This is not representative of the intended deployment, but it found Auto selecting bucketed CB and regressing by 18% for 1x128 + 31x4 and by 27% for a 16..512 length spread. A 128-request short-prompt case also selected ordinary batching but was 1.82x slower because planning dominated the very short generation.

This does not change my view that the opt-in scope is reasonable. At the same time, No-CB selections alone provide limited evidence about how well the cost model is calibrated when Auto selects bucketed CB.

To build confidence in that part of the policy, it would be helpful to document the coefficients as configuration-specific and provide a reproducible sweep that includes workloads where Auto selects both routes. Reporting Auto E2E / best baseline E2E, including the worst-case regression, on the target model/GPU/precision configurations would make the trade-offs easier to evaluate.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 359b3d5002

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[SuperGoodGame]

Thanks @nathon-lee for the independent smoke test. I agree that the default coefficients are starting estimates, not portable deployment settings. Users enabling Auto should calibrate for their model, GPU, precision and attention backend, and for the prompt-length/batch-size range they actually use. Auto remains disabled by default.

Recommended calibration: choose a small set of feasible (batch size, prompt length) pairs covering the deployment's short, typical and long prompts, together with small and typical batches. For example, a workload spanning 128–8192 input tokens should include representative points around 128, 512, 2048 and 8192, subject to memory and model context limits. Fit once, save the coefficients, and validate on separate mixed-length batches with representative generation lengths. Compare Auto against all available manual paths and report the worst regression before enabling it.

Why exact DP is practical here: for per-group batches around 128–256 requests, there are only 8,256–32,896 candidate segments. The O(n²) DP can therefore compute the exact optimum for the calibrated cost model over length-sorted contiguous partitions, rather than requiring heuristic search. Here n is the number of requests, not the number of tokens. A small n makes exact solving practical; prediction of measured GPU latency still depends on calibration.

Calibration is deployment setup work: run the supplied script during deployment preparation or warmup, save the coefficients in costs.json, and have the application load them once when constructing HybridEngineRolloutConfig. Calibration and fitting do not run inside generation requests. Runtime reuses the saved coefficients and performs only batch-length extraction and route planning. The same configuration can be reused across tasks whose lengths and batch sizes are covered by calibration; changing tasks alone does not require recalibration.

The current script's built-in calibration grid reaches only 768 input tokens. Those measurements do not establish long-context accuracy; longer deployments need an expanded grid. At present this means changing the shapes list in calibrate(); there is no automatic workload-range discovery.

There were also two distinct causes of the reported regressions:

• Wrong route selection: fixed estimates transferred poorly to the tiny model. Configuration-specific calibration keeps that model on ordinary generation in the measured sweep.
• Unnecessary routing overhead: per-request mask validation added 128 device scalar reads, and DP ran even when ordinary batching could be selected directly. Commit 359b3d5 batches validation/length extraction, skips DP using a cost lower bound or the exact uniform-length estimate, delays per-row request construction, and reuses generation configuration and result normalization. The 128-short-request case now has the same seven scalar reads as ordinary generation and no DP call.

The separate calibrated sweep reports Auto E2E / fastest manual baseline E2E (ordinary, equal-width CB or aligned CB):

Configuration	128×16	1×128 + 31×4	16..512	1×512 + 31×16	Worst regression
Tiny Qwen2 / FP32	1.0220×	0.9685×	0.9539×	0.9647×	+2.205%
Qwen2.5-3B / FP32	1.0005×	0.2345×	0.7808×	0.2265×	+0.054%
Qwen3-32B / BF16	1.0027×	0.3881×	0.7189×	0.1811×	+0.273%

All used SDPA on a single GPU. Each configuration uses one coefficient set fitted on separate uniform shapes. The two larger models choose ordinary generation for 128×16 and bucketed CB for the other three workloads; the tiny model chooses ordinary throughout. Auto's output IDs match ordinary generation in all twelve cases. These are controlled, EOS-disabled, two-token measurements exposing routing overhead, not long-context or production rollout validation. Results average three interleaved runs after warmup, or ten for the tiny model.

Reproduce the built-in short-context calibration and sweep with:

python benchmarks/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 --repeats 3 --calibrate costs.json
python benchmarks/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 --cost-config costs.json --repeats 3 --output sweep.json

The saved JSON records the configuration, and the sweep rejects mismatched metadata. 108 related tests and changed-file pre-commit checks pass.

[nathon-lee]

@SuperGoodGame Thanks for adding the calibration workflow, reproducible sweep, and the routing-overhead optimization. This addresses my earlier concern about configuration-specific performance evidence, and the default-off experimental scope now looks much clearer.

Before I can approve, could we please resolve the remaining route-invariance and compatibility issues:

1. Auto resolves a generation config and passes it to the ordinary batched path, but the bucketed path still reads module.generation_config for repetition penalty. If _prepare_generation_config() returns a copied config with a non-default penalty, the selected route can change generated tokens. Please propagate the resolved config (or penalty) through the continuous prefill/decode path and add a regression test.
2. The legacy manual continuous-prefill path now passes output_logits=False even when Auto is disabled. Since the inference requirements support older Transformers versions, this can turn into an unsupported generation/model kwarg and regress existing manual continuous batching. Please pass output_logits only when the adaptive path needs it, with a compatibility regression test.

The calibration benchmark is useful, but it is deployment evaluation tooling rather than DeepSpeed library runtime
code. Could we move it to https://github.com/deepspeedai/DeepSpeedExamples and link to it from the documentation
instead?

With those items addressed, I have no further concerns. @delock, do you have any remaining feedback before we merge?

[SuperGoodGame]

Thanks @nathon-lee. Addressed in 65a9115: the resolved config now reaches continuous Prefill, Decode and refill, and manual generation omits output_logits.

27 targeted tests pass, including regressions for both routes/cache APIs and the older generation interface. Both independent local reviews returned no findings.

The calibration script and fitting test moved to DeepSpeedExamples #1018; the docs now link there. The migrated script is unchanged, and its calibration/saved-config sweep smoke check passes. The Examples companion should merge first so the documentation link is available.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ecacbce12c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[SuperGoodGame]

The final fixes and validation are now in d4def3d. Independent review found a BF16 repetition-penalty rounding discrepancy; it now uses FP32 like HF. Unknown public generation switches are also rejected, with regressions for both stored and prepared configurations. The Examples CLI requires complete saved-calibration metadata.

I recalibrated Qwen3-32B BF16/SDPA on the final implementation, then ran the same four held-out workloads against all three manual baselines. Worst mean regression: +0.07428%; all four Auto outputs match no-CB. Auto chooses no-CB for 128 short requests and 2/6/2 prefill buckets for the mixed cases. These remain controlled two-token, EOS-disabled measurements; the PR description now includes the complete current-version table and source fingerprint. Both independent final reviews returned no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d4def3db03

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a0588b7dea

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Resolve generation config with the request overrides

For a custom/remote model whose _prepare_generation_config derives settings from generation kwargs (for example, enabling no_repeat_ngram_size based on max_new_tokens), this call sees none of the kwargs that _generate_batched later supplies to module.generate. Fresh evidence after the earlier prepared-config propagation fix is that the final code still calls prepare_config(None) without those overrides, so validation can accept the config while the batched route enables the processor during its second preparation and the bucketed direct-forward route never does, making token IDs depend on routing; prepare and validate with the same max-token and sampling kwargs used for generation.

Useful? React with 👍 / 👎.

[delock]

Hi @SuperGoodGame @nathon-lee . I think even when we set it to default off, the complexity is still there. I would suggest we focus on improve decode performance right now (we have segKI and graph capture WIP), then come back to prefill part when the decode part is settle. In that way we can make decode and prefill work more seamlessly. Let's say we revisit it one month from now, that's mid November. At that time we should know how batch segKI and graph capture works with CB and have better image to know how to put prefill optimization in.

[SuperGoodGame]

Thanks @delock. I understand the priority of stabilizing the decode path first.

I’ll keep this PR open for follow-up and revisit the integration around mid-November, once the CB, segKI, and graph capture direction is clearer. The implementation and benchmark results will remain available for that discussion.

Thanks for taking the time to review.

[SuperGoodGame]

I see prefill bucketing and decode optimizations as complementary. Our measurements show meaningful end-to-end generation gains when padding and prefill overhead are substantial.

I understand that the cache and scheduling interfaces need to align. If useful, I’d be happy to discuss how the prefill policy can work with CB, segKI, and graph capture, and help with the integration.