[SuperGoodGame]

November 11 update

Rebased onto master b9100c5, which includes #8390 and #9008; new head 82ac56e.

Uses #8939's shared sampling kwargs, token selector and FP32 repetition-penalty helper. Manual graph and continuous-batching paths retain upstream temperature/top-k/top-p sampling. Adaptive prefill retains its single-sample greedy scope, including fixed-slot graph decode. #9004's pre-built 4-D graph mask is preserved.

• Synchronize ZeRO-3 rollout early stopping #8390 early stopping: ordinary generation and Auto's batched fallback share Synchronize ZeRO-3 rollout early stopping #8390's generation kwargs, so they stop on natural EOS; ZeRO-3 keeps the group-aware synchronized stop. Auto remains single-process.
• Review fixes: the graph path rejects align_decode_fronts=True only when Auto is disabled, because Auto always uses aligned bucketed prefill. The accelerator graph test covers both flag values. The documentation states that Auto takes precedence over align_decode_fronts on the graph path; cache trimming is still rejected.
• Size HybridEngine rollout KV caches from local prefill shapes under AutoTP #9008 lazy KV layers: the bucketed prefill copy initializes uninitialized cache layers through the shared _continuous_cache_layer() helper, like the equal-width path.
• Codex P2 (preparer findings): after a custom _prepare_generation_config() runs, Auto rejects the prepared config if it changes any request override (token budget, EOS, PAD or greedy settings). Otherwise the batched route would honor the change inside generate() while the CB routes keep the request values. A preparer that replaces or edits the attention mask is rejected for the same reason.
• Codex P2 (mask lengths): effective prompt lengths use count_nonzero, so BF16/FP16 masks longer than their exact-integer range no longer round the bucket width.
• Codex P1 (test reliability): the route-planning profile test uses a frozen clock advanced by the model's generation-config preparer, instead of wall-clock bounds or mocked rollout internals.
• The decode cost still uses the configured maximum generation budget. Natural EOS can shorten decode, so the extra-CB estimate can be too large; the documentation states this assumption.

Validation on an 8 × H200 node, PyTorch 2.9.1+cu128, CUDA 12.8, Transformers 4.57.1: 232 passed, 3 skipped, 2 deselected across the runtime rollout, scheduler, interface, static-cache and v1 accelerator suites, including #8390's 4-rank subgroup and 2-GPU ZeRO-3 synchronized-stop tests. Skips require the newer cache query-offset API; deselections are the native shared-prefill kernel and optional pretrained-model test. Pre-commit passed on all changed files. A small natural-EOS check (Qwen2.5-3B-Instruct, four short questions) ended every response at EOS, with Auto matching ordinary generation. This update used AI assistance.

The performance tables below retain their original revision attribution; large-model throughput and peak memory were not remeasured for this rebase.

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. With graph capture, Auto uses eager bucketed prefill and fixed-slot graph decode; cache trimming and shared prefill remain unsupported.

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?

November refresh

Rebased onto b62afa8 and squashed to one commit; the PR is now mergeable. The refresh keeps #8732 graph state, #8839 last-position prefill logits, #8933 accelerator-neutral graph calls and #8852 DeepSpeedStaticCache.get_query_offset(). #8980's explicit accelerator opt-in is respected: a backend without is_graph_capture_supported() falls back to eager before any routing decision.

Superseded by 8a9641b: Auto now combines eager bucketed prefill with fixed-slot graph decode (see the November 11 update). SegKI (#8602) is not part of this PR.

Validation, 8×H200-class node with one GPU in use (NVIDIA H20Z, PyTorch 2.9.1+cu128, CUDA 12.8, Transformers 4.57.1):

Check	Result
rollout/cache suite	159 passed, 3 skipped, 1 deselected
Skips	Three Transformers 5.20-only query-offset cases added by #8852
Deselected	Native shared-prefill kernel test; the test host has no CUDA compiler
Changed-file pre-commit	Passed

Graph replay, resolved-config propagation on both routes (FP32/BF16), DP bucket/refill oracles, Auto/Graph rejection (since replaced by joint routing) and manual graph dispatch all pass. No new performance numbers are claimed: the tables above were measured before #8839/#8732, so refreshed baselines need a separate calibration run if wanted.

AI assistance was used for the rebase, conflict resolution, test execution and drafting of this description; the submitting human reviewed the diff and the results above.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

Codex Review Summary: 🔄 Running

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

Review	Status	Commit	Review trigger
📝 Code Review	🔄 Running since 2026-11-14T15:38:56.463279Z	0b71e5c	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.

[SuperGoodGame]

Refreshed onto current master (b62afa8) and force-pushed; the PR is now MERGEABLE and squashed to a single commit d9a2cc3.

Changes absorbed from upstream since the last revision:

• Enable CUDA Graph for continuous rollouts #8732 graph state, [OPSD] Request last-position logits only in HybridEngine rollout prefill #8839 last-position prefill logits and [Hybrid Engine] Enable graph capture for HybridEngine on non-CUDA accelerators #8933 accelerator-neutral graph calls are retained.
• Support Transformers 5.20 cache query offsets in DeepSpeedStaticCache #8852 DeepSpeedStaticCache.get_query_offset() is preserved alongside the adaptive-prefill work.
• [Hybrid Engine] Require an explicit accelerator opt-in for graph capture #8980's explicit accelerator opt-in is respected: a backend without is_graph_capture_supported() falls back to eager before routing.

Scope is unchanged. Auto and graph capture remain mutually exclusive because the graph uses fixed physical slots while Auto requires aligned, bucketed prefill. Joint Auto/Graph routing is deferred until the graph path and SegKI settle; SegKI (#8602) is not part of this PR.

Validation on one GPU of an 8×H20Z node (PyTorch 2.9.1+cu128, CUDA 12.8, Transformers 4.57.1):

Check	Result
Rollout/cache suite	159 passed, 3 skipped, 1 deselected
Skips	Three Transformers 5.20-only query-offset tests from #8852
Deselected	Native shared-prefill kernel test (no CUDA compiler on the test host)
Changed-file pre-commit	Passed

Covered: real graph replay, resolved generation-config propagation on both routes (FP32/BF16), DP bucket/refill oracles, Auto/Graph rejection, manual graph dispatch.

The performance tables are unchanged and still attributed to their original revisions; I have not relabeled them as new-HEAD measurements. A refreshed calibration run can follow if the new baselines are wanted for review.

[delock]

refill. Joint Auto/Graph routing is deferred until the graph path and SegKI settle; SegKI (#8602) is not part of this PR.

Validation on one GPU of an 8×H2

Hi @SuperGoodGame thank you for keeping an eye on this. The idea of this PR looks good. Currently it support eager mode only. Does it indicates CB should decouple prefill and decode (so prefill can run in eager and decode can run in graph capture)?

[SuperGoodGame]

Thanks @delock. Yes, prefill and decode can be decoupled. I’ve pushed the implementation in 8a9641b. Eager prefill groups newly admitted prompts by length and writes their KV entries into the slots assigned by the graph scheduler. Surviving requests keep their slots as decode replays the captured graph.

Both columns below use CUDA Graph decode on a single H200, with BF16, SDPA, and 32 requests. Values are mean end-to-end generation latency over three interleaved runs after warmup. Output lengths were fixed with EOS disabled for timing.

Qwen2.5-3B-Instruct

Prompt lengths	Output tokens/request	Equal-width prefill	Bucketed prefill	Latency change
32×16	16	196.24 ms	199.96 ms	+1.90%
1×512 + 31×16	16	497.03 ms	334.98 ms	−32.60%
16, 32, …, 512	16	494.01 ms	482.64 ms	−2.30%
32×16	128	1504.71 ms	1507.89 ms	+0.21%
1×512 + 31×16	128	2520.23 ms	2352.57 ms	−6.65%

Qwen3-32B

Prompt lengths	Output tokens/request	Equal-width prefill	Bucketed prefill	Latency change
32×16	16	648.53 ms	648.85 ms	+0.05%
1×512 + 31×16	16	3179.83 ms	1466.65 ms	−53.88%
16, 32, …, 512	16	3183.78 ms	2518.05 ms	−20.91%
32×16	128	6052.51 ms	6052.99 ms	+0.01%
1×512 + 31×16	128	13602.31 ms	11897.50 ms	−12.53%

Generated token IDs matched across paths in 9 of 10 cases. In the last 32B row, the 31 identical short prompts first diverged at output token 107. Equal-width prefill produced a BF16 tie between tokens 198 and 220 (both 32.75). With two buckets, their scores were 32.50 and 32.75. Forcing a single bucket reproduced the equal-width output exactly.

A separate GPU regression test covered early EOS and slot refill. Each response matched Hugging Face token for token, and two rollout calls reused one graph capture.

With both options enabled, the caller selects graph decode and the adaptive policy chooses the prefill buckets. Does that match the separation you had in mind?

[delock]

Hi @SuperGoodGame , the separation of prefill and decode sounds good. It allows your PR compatible with graph capture in decode. Now it has conflict with merged #8939 which introduced greedy sampling. Can you take a look at the conflict? Also #8390 introduced early stopping that may affect cost model.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 32cd9619c3

ℹ️ 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: f2ae71be4b

ℹ️ 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: 5d134f5abb

ℹ️ 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: 7ce82f24e2

ℹ️ 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]

Hi @delock, the previous CI run failed due to a Transformers version compatibility issue. This is now fixed, and the targeted tests pass locally with PyTorch 2.10 CPU and Transformers 5.23 (147 passed). I've also rebased onto the latest master; the current head is edfbf88.

Could you please approve the CI workflows again for the latest push? They are currently marked action_required. Thanks for your help!

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: edfbf883ff

ℹ️ 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]

Skip the BF16 cases on unsupported accelerators

On MPS running macOS 13, MPS_Accelerator.is_bf16_supported() returns false, but _make_small_qwen() places the model on MPS and this case unconditionally converts it to BF16. The test therefore raises an unsupported-dtype error before exercising rollout; the new BF16-mask test above has the same problem. Gate both cases on the selected accelerator's BF16 capability, following the explicit unavailable-hardware skip requirement in the checklist mandated by AGENTS.md.

AGENTS.md reference: AGENTS.md:L40-L40

Useful? React with 👍 / 👎.

[delock]

Hi @SuperGoodGame can you address the latest codex comments (MPS and prepare called twice)? I'd prefer a small change size fix for these two comments, thanks!

[SuperGoodGame]

@codex review

Fixed both findings in 810fb39:

• BF16 tests now skip before allocating tensors when the selected accelerator does not support BF16.
• Auto's batched fallback lets generate() prepare the original configuration, rather than preparing an already-transformed copy again. The existing custom-preparer test now multiplies the repetition penalty and checks the effective settings on both routes. Before the fix, four batched cases observed 3 → 9; all pass after the fix.

Validation: PyTorch 2.10 CPU / Transformers 5.23: 147 passed, 1 skipped, 4 deselected; Transformers 4.57 CPU targeted tests: 64 passed; single H200 targeted tests: 10 passed. Pre-commit passed. macOS/MPS was not available for hardware validation.

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Another round soon, please!

Reviewed commit: 810fb39cb9

ℹ️ 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 @delock for the approval. The failed cpu-torch-latest-v1-multirank / unit tests run is not caused by this PR:

• Why the full suite ran: the test selector diffs against pull_request.base.sha (a75ab89, this branch's previous base), while the checked-out merge commit sits on 84481c3. Fused CPU Adam step #8841's csrc/adam/* changes were therefore counted as part of this PR, which selects the full v1 suite (MODE: all). The previous run (edfbf88) selected only tests/unit/v1/test_hybrid_engine_rollout_v1.py.
• Why it failed: this run landed on an AWS runner whose CPU supports AVX512-FP16. The 10 failures are the fp16 cases of TestZeroAutoCast::test / test_safe_modules_conf, which the usual Azure runners skip because torch.ops.mkldnn._is_mkldnn_fp16_supported() is False there, and TestMuonMatchesZeroStage0 (relative error ~1.5e-3 against the 1e-5 bound). I reproduced the same 10 failures on plain master, both a75ab89 and 84481c3, without this PR, on a Xeon Platinum 8558 (AVX512-FP16/AMX).

I rebased onto current master (84481c3); the new head is 3e651b9 and the PR diff is unchanged. With this base the selector again picks only the rollout test file (locally: 1 passed, 1 skipped on CPU; 193 rollout/cache tests pass on GPU).

If this looks good, please feel free to merge. If you would like to rerun CI first, please let me know and I will rebase onto the latest master right before it runs, so the selection stays scoped to this PR. Note that the merge queue runs the full v1 suite, so it can hit the same failures if it lands on an AVX512-FP16 runner.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3e651b9c37

ℹ️ 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]

Thanks @delock for the approval. The failed cpu-torch-latest-v1-multirank / unit tests run is not caused by this PR:

• Why the full suite ran: the test selector diffs against pull_request.base.sha (a75ab89, this branch's previous base), while the checked-out merge commit sits on 84481c3. Fused CPU Adam step #8841's csrc/adam/* changes were therefore counted as part of this PR, which selects the full v1 suite (MODE: all). The previous run (edfbf88) selected only tests/unit/v1/test_hybrid_engine_rollout_v1.py.
• Why it failed: this run landed on an AWS runner whose CPU supports AVX512-FP16. The 10 failures are the fp16 cases of TestZeroAutoCast::test / test_safe_modules_conf, which the usual Azure runners skip because torch.ops.mkldnn._is_mkldnn_fp16_supported() is False there, and TestMuonMatchesZeroStage0 (relative error ~1.5e-3 against the 1e-5 bound). I reproduced the same 10 failures on plain master, both a75ab89 and 84481c3, without this PR, on a Xeon Platinum 8558 (AVX512-FP16/AMX).

I rebased onto current master (84481c3); the new head is 3e651b9 and the PR diff is unchanged. With this base the selector again picks only the rollout test file (locally: 1 passed, 1 skipped on CPU; 193 rollout/cache tests pass on GPU).

If this looks good, please feel free to merge. If you would like to rerun CI first, please let me know and I will rebase onto the latest master right before it runs, so the selection stays scoped to this PR. Note that the merge queue runs the full v1 suite, so it can hit the same failures if it lands on an AVX512-FP16 runner.

Hi @SuperGoodGame I'll take a look tomorrow. Also please use merge instead of rebase, this will help track what had been changed from last approval. Thanks!

[delock]

Hi @SuperGoodGame can you help resolve the conflicts? Thanks!