feat(rollout): add continuous batching generation prototype

[nathon-lee]

Summary

Add an experimental continuous-batching generation path to DeepSpeed rollout for OPSD/OPD evaluation.

The implementation maintains a bounded number of active request rows, retires requests independently on EOS or when the shared generation budget is reached, compacts surviving StaticCache rows when necessary, and refills released slots with pending requests.

The existing HybridEngineRollout.generate() behavior remains unchanged when continuous batching is not enabled.

What changed

• Add SamplingConfig.continuous_batch_size to enable continuous batching through the existing generate(request, sampling) entry point.
• Preserve the original request-row order in the returned right-padded RolloutBatch.
• Keep the continuous-batching scheduler and cache lifecycle helpers internal to the rollout implementation.
• Extend DeepSpeedStaticCache with per-row write positions and active-row compaction.
• Skip unnecessary cache and attention-mask compaction during steady-state decode.
• Add an identity-compaction fast path while preserving dead-row zeroing for safe row reuse.
• Require Transformers cache-class support for continuous batching while keeping the default generation path available for older models.
• Add deterministic CPU coverage for the modern StaticCache path.
• Update rollout tests and documentation.

Design

Continuous batching is enabled with:

SamplingConfig(continuous_batch_size=N)

When enabled:

1. Up to N request rows remain active at a time.
2. Requests retire independently when they reach EOS or the shared generation budget.
3. Released cache rows are reused for pending requests.
4. Surviving rows are compacted only when required.
5. The final RolloutBatch is restored to the original request-row order.

The scheduler is kept separate from the model backend. The rollout backend owns prompt prefill, decode execution, cache movement, and attention metadata construction.

Scope and limitations

This is an experimental rollout feature intended for OPSD/OPD evaluation.

The current implementation supports:

• One prompt row per request.
• Equal prompt widths.
• Greedy decoding.
• One sample per prompt.
• Transformers models with cache-class support.
• Single-device validation.

The following are intentionally left for follow-up work:

• Variable-width prompt alignment and dead-zone trimming.
• In-place refill.
• Paged or block-based KV cache.
• Prefix/radix cache.
• CUDA Graph support for continuous batching.
• Continuous-batching profiling.
• Shared-prefill integration.
• KV-cache offload.
• Multi-GPU/tensor-parallel validation.

Models without Transformers cache-class support continue to use the default generate() path.

Validation

Focused rollout tests:

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

Focused H200 validation:

52 passed, 2 warnings

Coverage includes:

• Scheduler bookkeeping.
• Modern StaticCache prefill and decode.
• EOS-based request retirement.
• Row compaction and per-row write positions.
• Pending-request refill.
• Identity-compaction fast path.
• Dead-row clearing.
• Default generation behavior.

[delock]

Hi Nathon, from my understanding, your design would assume the decode front is on the same column. So this is why in the draft prompts need to have equal length. And when one row hit EOS, you will fill another prompt and make sure the second token after prefill will align with the decode front.

If this understanding is correct, then I have two suggestions:

1. sort the prompt from longest to shortest, right align prompts in the first batch then prefill in a way that the decoding front would be aligned. And when one row hit EOS and you want to fill another prompt in, there is always enough space to put the (shorter) new prompt kv in a way that the second token after prefill will align with the decode front.

2. periodically looking for 'dead zone'(columns before the earliest active row's span start) in the beginning and trim them by shift the rest tokens to the left. The trim would always be safe because newly added prompt will be shorter than active prompts in all rows, so there is always space to add new prompt after trim. By keeping trimming, we don't need to maintain a very large buffer for continuous batching and make it GPU memory size friendly (the length we needed for buffer is max(prompt_len + max_new_tokens)).

These suggestions assume we are not in serving scenario, where samples arrive in random order and serve in FIFO manner. Because it is rollout, we don't necessarily honor generation order in the batch (no TTFT and TPOT SLA), this is the condition we should exploit.

[nathon-lee]

Thanks for the detailed suggestions, @delock. Your understanding is correct. I agree that sorting and right-aligning the prompts, together with trimming dead zones, could improve decode-front alignment and reduce the required buffer size. I’ll investigate these ideas in a follow-up and keep this PR focused on the baseline prototype.

[delock]

Thanks for the detailed suggestions, @delock. Your understanding is correct. I agree that sorting and right-aligning the prompts, together with trimming dead zones, could improve decode-front alignment and reduce the required buffer size. I’ll investigate these ideas in a follow-up and keep this PR focused on the baseline prototype.

I saw your PR is still in draft state. Let me know when you finished your change and ready for review. Thanks!

[nathon-lee]

Thanks for the detailed suggestions, @delock. I ran the focused tests for this PR in the H200 environment. All 35 tests passed, including the native CUDA repeat_kv_cache test.

PR #8368 is now ready for review.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 67b246ef19

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[nathon-lee]

Sorry for the delayed response. I was traveling for work recently. Thanks for your patience.

Thanks, @delock. I agree that RolloutRequest already represents a batch, so we can use one SamplingConfig for the whole batch and avoid per-sample max_new_tokens.
The different max_new_tokens values in the prototype were only used to exercise deterministic retirement and refill behavior. They are not required by OPD/OPSD, since rows can still retire independently when they reach EOS.
I propose the following interface:

• Keep generate(request, sampling) as the single public entry point.
• Add continuous_batch_size: Optional[int] = None to SamplingConfig.
• When it is None, preserve the current generation behavior.
• When it is set to a value >= 1, use continuous batching with at most that many active rows from RolloutRequest.
• Return a RolloutBatch in the original request-row order.
• Keep ContinuousBatchRequest, ContinuousBatchScheduler, and ContinuousBatchUpdate internal and remove them from the public exports.
• Move the current generate_continuous() implementation behind generate() as an internal helper.

The continuous path will retain the current prototype limitations. If this matches your expectation, I’ll update the implementation, tests, and documentation accordingly.

[delock]

Hi Nathon, the interface looks good. I'll start review your implementation.

Sorry for the delayed response. I was traveling for work recently. Thanks for your patience.

Thanks, @delock. I agree that RolloutRequest already represents a batch, so we can use one SamplingConfig for the whole batch and avoid per-sample max_new_tokens. The different max_new_tokens values in the prototype were only used to exercise deterministic retirement and refill behavior. They are not required by OPD/OPSD, since rows can still retire independently when they reach EOS. I propose the following interface:

• Keep generate(request, sampling) as the single public entry point.
• Add continuous_batch_size: Optional[int] = None to SamplingConfig.
• When it is None, preserve the current generation behavior.
• When it is set to a value >= 1, use continuous batching with at most that many active rows from RolloutRequest.
• Return a RolloutBatch in the original request-row order.
• Keep ContinuousBatchRequest, ContinuousBatchScheduler, and ContinuousBatchUpdate internal and remove them from the public exports.
• Move the current generate_continuous() implementation behind generate() as an internal helper.

The continuous path will retain the current prototype limitations. If this matches your expectation, I’ll update the implementation, tests, and documentation accordingly.

[delock]

Thanks @nathon-lee — the unified interface looks right, and the rework faithfully implements what we settled in the interface thread: single generate() entry point, SamplingConfig.continuous_batch_size: Optional[int] = None, a single right-padded RolloutBatch in original row order, and the scheduler types moved out of the public exports. I verified the core bookkeeping (row compaction, per-row write positions, prefill copies, logical positions) end-to-end on CPU and found no regressions; the full unit suite passes locally as well.

Remaining items before merge — 2 blocking, 5 minor. Details inline. Two known gaps at the generate() dispatch boundary (use_shared_prefill + CB is silently ignored; enable_profiling + CB produces no profile) are near-harmless today since CB requires n_samples_per_prompt=1 — can you open open follow-up issues covering both? They won't hold this PR.

[B2] No test coverage for the modern (StaticCache) path — _supports_cache_class = True never appears in the test suite. The only end-to-end test exercises the legacy path, while real models run the modern one, so our flagship path currently has no CI coverage. This is testable deterministically on CPU: a fake cache-class model whose next token depends on the row's cache contents (e.g. logits derived from the row's cached-value sum) makes any write-position / compaction / prefill-copy error change the outputs. Happy to share a reference implementation.

[B3] cache.compact runs unconditionally every decode step — steady-state steps pass an identity permutation and still pay a full-KV index_select + clone + copy (GB-scale per step at production sizes), and tail retirements leave the prefix already correct. Suggested fix, three layers:

1. skip entirely when update.retired is empty;
2. inside _compact_rows, skip the survivor copy when active_indices == arange(count), but keep the zeroing of dead rows — row reuse depends on it (a stale attention-mask row would make the next occupant attend to the previous request's KV);
3. apply the same treatment to the attention-mask compaction at the call site.

In-place refill (admitted row = retired row, no survivor movement on 1:1 churn) can be a follow-up.

Scope of CB with legacy model

I believe the legacy CB path should be removed. For these reasons:

• OPSD rollouts target modern reasoning models, all of which support cache classes;
• the fallback rebuilds the full cache every step — running a throughput feature on the slowest possible cache management undermines its purpose;
• net deletion of ~120 lines plus their dedicated tests;
  @PKUWZP @minjiazhang for this opinion

Concretely: raise a clear ValueError for models without _supports_cache_class ("continuous batching requires a model with cache-class support; use the default generate() path or upgrade transformers"). The goal is to keep unnecessary complexity out. If we decide to keep it, one small fix applies (legacy path should reuse the incoming RolloutRequest instead of re-catting rows).

[delock]

This is our only CB end-to-end test and it exercises the legacy path (_supports_cache_class = False). The modern path — the one every real model takes — has no coverage. Could you add a CPU e2e in the same style, with a fake cache-class model that calls past_key_values.update(...) per layer and derives logits from the row's cached values? That pins down the KV bookkeeping (write positions, compaction, prefill copies) deterministically.

[delock]

This runs on every decode step. With no retirement keep_slots is the identity permutation, so we pay a full-KV copy for nothing; with tail-only retirements the prefix is already correct and only the dead rows need clearing. Suggested guard structure:

if update.retired:
    cache.compact(keep_slots)   # + identity fast path inside _compact_rows (keep the zeroing!)

Important detail: the identity fast path must skip only the survivor copy — the zeroing of rows [count:] is what makes retired rows safe to reuse (a stale mask row would corrupt the next occupant). The attention-mask compaction below deserves the same treatment.

[delock]

DeepSpeedStaticLayer.compact has no callers in production or tests — DeepSpeedStaticCache.compact (via layer._compact_rows) is the only entry point. Please remove it; this also eliminates the duplicated validation block.

[delock]

The rollout already owns this tensor (created in _generate_continuous, registered via set_write_position), so reaching into the cache's private attribute is avoidable — please thread it as a parameter like attention_mask already is.

[nathon-lee]

Thanks, @delock, for the thorough review and CPU validation. The B2/B3 issues and cleanup suggestions are clear. I’ll add modern StaticCache CPU E2E coverage, optimize compaction while preserving dead-row zeroing, remove the unused/private implementation details, and update the experimental wording. I’ll also open follow-up issues for the shared-prefill and profiling dispatch behavior. For the legacy CB path, I agree with the rationale and will confirm with PKUWZP and minjiazhang before removing it.

[delock]

Here cache reset only happens when all rows retire at the same time step. This will likely happen when all rows reaches max new tokens. However if some row retire with EOS (this is why continuous batching is needed), then reset may never happen and we will eventually hit the end of the cache with error.

Trim and shift the whole cache to the left could address this issue. Check periodically if at least N columns in the front of cache does not contain survivors cache, these columns can be removed by moving kv cache elements to the left, and free up the right part of the cache for further decoding. Update decode position accordingly. Pick a proper N could make balance between trim overhead and extra memory space needed.

[nathon-lee]

Thanks, @delock. You’re right that the current reset only happens when all active rows retire at the same step. With staggered EOS and continuous refill, the cache position can keep increasing until it exceeds the allocated cache length.

I’ll add periodic dead-prefix trimming for the modern StaticCache path, update the attention mask, write positions, and decode position after the shift, and add a deterministic CPU regression with staggered EOS/refill that exceeds the original cache span. I’ll keep the trimming threshold internal and avoid changing the public API.

[delock]

@nathon-lee Can you fix formatting error? Thanks!

[nathon-lee]

@nathon-lee Can you fix formatting error? Thanks!

OFC, I’ll fix the formatting and push an update shortly.