feat(rollout): add continuous batching generation prototype

[nathon-lee]

Summary

Add a bounded continuous‑batching prototype to DeepSpeed rollout.
The implementation admits requests in FIFO order, retires completed requests,
compacts surviving KV‑cache rows, and prefills pending requests into newly
available slots. It supports requests with different response budgets while
preserving the existing HybridEngineRollout.generate() behavior.

This PR provides the Core API and cache primitives required by the companion
DeepSpeedExamples benchmark.

What changed

• Add ContinuousBatchRequest, ContinuousBatchUpdate, and
  ContinuousBatchScheduler.
• Add HybridEngineRollout.generate_continuous(...).
• Support request retirement and pending‑request refill.
• Add per‑row write positions and active‑row compaction to
  DeepSpeedStaticCache.
• Support both newer Transformers cache objects and legacy tuple‑based KV
  caches.
• Handle Transformers StaticCache constructor variations across versions.
• Add validation for unsupported inputs and invalid scheduler transitions.
• Document the prototype API and its current limitations.

Design

The scheduler is intentionally separated from the model backend. It reports:

• active requests;
• cache rows that survive a retirement event;
• retired request IDs;
• newly admitted requests and their target slots.

The rollout backend owns prompt prefill, decode execution, cache movement, and
attention metadata construction.

The default HybridEngineRollout.generate() path is unchanged.

Scope and limitations

This is an initial prototype intended for evaluation and experimentation.
Current limitations:

• one prompt row per request;
• equal prompt widths;
• greedy decoding only;
• one sample per prompt;
• CUDA Graph capture is not supported;
• paged attention is not implemented;
• the default rollout generation path is unchanged;
• legacy KV‑cache refill currently uses cache selection and concatenation;
• multi‑GPU/tensor‑parallel validation is not included in this PR.

Validation

Focused unit tests

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

[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".