From 14e3ad216a7b9c768557c962aa44917830aa3d00 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Sun, 30 Aug 2026 08:21:16 +0800
Subject: [PATCH 01/12] feat: add continuous batching rollout prototype

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 deepspeed/runtime/rollout/__init__.py         |   5 +
 .../runtime/rollout/continuous_batching.py    | 134 ++++++++++++++++++
 deepspeed/utils/static_cache.py               |  79 +++++++++--
 docs/code-docs/source/inference-engine.rst    |  17 +++
 .../rollout/test_continuous_batching.py       |  68 +++++++++
 tests/unit/utils/test_static_cache.py         |  65 +++++++++
 6 files changed, 360 insertions(+), 8 deletions(-)
 create mode 100644 deepspeed/runtime/rollout/continuous_batching.py
 create mode 100644 tests/unit/runtime/rollout/test_continuous_batching.py
 create mode 100644 tests/unit/utils/test_static_cache.py

diff --git a/deepspeed/runtime/rollout/__init__.py b/deepspeed/runtime/rollout/__init__.py
index 16f6fc595da6..97518d6f9070 100644
--- a/deepspeed/runtime/rollout/__init__.py
+++ b/deepspeed/runtime/rollout/__init__.py
@@ -18,6 +18,8 @@
     SamplingConfig,
 )
 from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout
+from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler,
+                                                           ContinuousBatchUpdate)
 
 __all__ = [
     "HybridEngineRollout",
@@ -27,6 +29,9 @@
     "RolloutRequest",
     "SamplingConfig",
     "build_rollout",
+    "ContinuousBatchRequest",
+    "ContinuousBatchScheduler",
+    "ContinuousBatchUpdate",
 ]
 
 
diff --git a/deepspeed/runtime/rollout/continuous_batching.py b/deepspeed/runtime/rollout/continuous_batching.py
new file mode 100644
index 000000000000..2b0071dd86ed
--- /dev/null
+++ b/deepspeed/runtime/rollout/continuous_batching.py
@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Request scheduling primitives for continuous-batching rollouts.
+
+This module deliberately stops at the scheduler/cache boundary. The model
+backend owns prompt prefill and decode; the scheduler reports which old cache
+rows survive, which requests retire, and which pending requests can be
+admitted into the newly free rows.
+"""
+
+from collections import deque
+from dataclasses import dataclass
+from typing import Hashable
+
+
+@dataclass(frozen=True)
+class ContinuousBatchRequest:
+    """A request waiting for a slot in a continuous decode batch."""
+
+    request_id: Hashable
+    max_new_tokens: int
+
+    def __post_init__(self) -> None:
+        if self.max_new_tokens <= 0:
+            raise ValueError("max_new_tokens must be positive")
+
+
+@dataclass(frozen=True)
+class ContinuousBatchUpdate:
+    """Result of one scheduler transition.
+
+    ``keep_slots`` indexes the previous active batch. The caller should
+    compact the KV cache with these indices, then prefill ``admitted`` into
+    the free rows at the end of the compacted batch.
+    """
+
+    active: tuple[ContinuousBatchRequest, ...]
+    keep_slots: tuple[int, ...]
+    retired: tuple[Hashable, ...]
+    admitted: tuple[ContinuousBatchRequest, ...]
+
+    @property
+    def active_ids(self) -> tuple[Hashable, ...]:
+        return tuple(request.request_id for request in self.active)
+
+    @property
+    def admitted_slots(self) -> tuple[int, ...]:
+        """Rows available for prefilling the newly admitted requests."""
+        start = len(self.active) - len(self.admitted)
+        return tuple(range(start, len(self.active)))
+
+
+class ContinuousBatchScheduler:
+    """FIFO scheduler for bounded, slot-based continuous batching.
+
+    ``schedule`` performs admission/retirement without advancing tokens.
+    ``advance`` represents one decode step for every active request and also
+    retires requests whose token budget has been consumed. A caller may pass
+    explicit finished IDs when the model emits EOS before that budget.
+    """
+
+    def __init__(self, max_batch_size: int):
+        if max_batch_size <= 0:
+            raise ValueError("max_batch_size must be positive")
+        self.max_batch_size = max_batch_size
+        self._pending = deque()
+        self._active = []
+        self._generated = {}
+        self._known_ids = set()
+
+    @property
+    def active(self) -> tuple[ContinuousBatchRequest, ...]:
+        return tuple(request for request, _ in self._active)
+
+    @property
+    def pending(self) -> tuple[ContinuousBatchRequest, ...]:
+        return tuple(self._pending)
+
+    def submit(self, request: ContinuousBatchRequest) -> None:
+        if not isinstance(request, ContinuousBatchRequest):
+            raise TypeError("request must be a ContinuousBatchRequest")
+        if request.request_id in self._known_ids:
+            raise ValueError(f"duplicate request_id: {request.request_id!r}")
+        self._known_ids.add(request.request_id)
+        self._pending.append(request)
+
+    def schedule(self, finished_ids=()) -> ContinuousBatchUpdate:
+        finished_ids = tuple(finished_ids)
+        active_by_id = {request.request_id: slot for slot, (request, _) in enumerate(self._active)}
+        unknown = set(finished_ids) - active_by_id.keys()
+        if unknown:
+            raise ValueError(f"finished request is not active: {next(iter(unknown))!r}")
+
+        finished = set(finished_ids)
+        survivors = [(request, self._generated[request.request_id]) for request, _ in self._active
+                     if request.request_id not in finished]
+        keep_slots = tuple(slot for slot, (request, _) in enumerate(self._active)
+                           if request.request_id not in finished)
+        retired = tuple(request.request_id for request, _ in self._active if request.request_id in finished)
+
+        free_slots = self.max_batch_size - len(survivors)
+        admitted = []
+        for _ in range(free_slots):
+            if not self._pending:
+                break
+            request = self._pending.popleft()
+            admitted.append(request)
+            survivors.append((request, 0))
+            self._generated[request.request_id] = 0
+
+        self._active = survivors
+        for request_id in retired:
+            self._generated.pop(request_id, None)
+            self._known_ids.discard(request_id)
+        return ContinuousBatchUpdate(tuple(request for request, _ in survivors), keep_slots, retired, tuple(admitted))
+
+    def advance(self, finished_ids=()) -> ContinuousBatchUpdate:
+        explicit_finished = set(finished_ids)
+        active_ids = {request.request_id for request, _ in self._active}
+        unknown = explicit_finished - active_ids
+        if unknown:
+            raise ValueError(f"finished request is not active: {next(iter(unknown))!r}")
+        finished = set(explicit_finished)
+        updated = []
+        for request, generated in self._active:
+            generated += 1
+            self._generated[request.request_id] = generated
+            if generated >= request.max_new_tokens:
+                finished.add(request.request_id)
+            updated.append((request, generated))
+        self._active = updated
+        return self.schedule(finished)
diff --git a/deepspeed/utils/static_cache.py b/deepspeed/utils/static_cache.py
index 520bef9c7314..90ee8c80d790 100644
--- a/deepspeed/utils/static_cache.py
+++ b/deepspeed/utils/static_cache.py
@@ -22,8 +22,8 @@
 graph replays read the current value each time.
 
 The caller (HybridEngineRollout) must call ``cache.set_write_position(pos)``
-before each replay, where ``pos`` is a scalar ``torch.long`` tensor on the
-correct device.
+before each replay. ``pos`` may be a scalar (the existing fixed-batch behavior)
+or one position per cache row for continuous batching.
 """
 
 import torch
@@ -49,6 +49,10 @@ def __init__(self, max_cache_len: int):
         self._write_position: torch.Tensor | None = None
 
     def set_write_position(self, pos: torch.Tensor):
+        if not isinstance(pos, torch.Tensor) or pos.dim() not in (0, 1):
+            raise ValueError("write position must be a scalar or a 1-D tensor")
+        if pos.dtype not in (torch.int32, torch.int64):
+            raise ValueError("write position must use an integer dtype")
         self._write_position = pos
 
     def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None:
@@ -86,6 +90,19 @@ def update(
 
         kv_length = key_states.shape[-2]
 
+        if self._write_position is not None and self._write_position.dim() == 1:
+            if self._write_position.numel() < key_states.shape[0]:
+                raise ValueError("per-row write positions must cover the active batch size")
+            if kv_length != 1:
+                raise ValueError("per-row write positions currently support one decode token at a time")
+            cache_position = self._write_position[:key_states.shape[0]].to(self.device)
+            if (cache_position < 0).any() or (cache_position >= self.max_cache_len).any():
+                raise ValueError("per-row write positions must be within the cache bounds")
+            row_indices = torch.arange(key_states.shape[0], device=self.device)
+            self.keys[row_indices, :, cache_position, :] = key_states[:, :, 0, :]
+            self.values[row_indices, :, cache_position, :] = value_states[:, :, 0, :]
+            return self.keys, self.values
+
         if self._write_position is not None:
             cache_position = torch.arange(kv_length, device=self.device) + self._write_position
         else:
@@ -103,7 +120,7 @@ def update(
     def get_mask_sizes(self, query_length: int) -> tuple[int, int]:
         return self.max_cache_len, 0
 
-    def get_seq_length(self) -> int:
+    def get_seq_length(self) -> int | torch.Tensor:
         if not self.is_initialized:
             return 0
         if self._write_position is not None:
@@ -118,10 +135,51 @@ def reset(self) -> None:
             self.keys.zero_()
             self.values.zero_()
 
+    def compact(self, active_indices: torch.Tensor) -> None:
+        """Move active cache rows to the front and update their positions.
+
+        Continuous batching retires requests from arbitrary rows. Copying the
+        survivors in one operation avoids in-place overlap when a later row is
+        moved into an earlier slot, while retaining the tensors' static
+        addresses for CUDA graph users.
+        """
+        if not self.is_initialized:
+            raise RuntimeError("cannot compact an uninitialized cache")
+        if not isinstance(active_indices, torch.Tensor) or active_indices.dim() != 1:
+            raise ValueError("active_indices must be a 1-D tensor")
+        active_indices = active_indices.to(device=self.keys.device, dtype=torch.long)
+        if active_indices.numel() > self.max_batch_size:
+            raise ValueError("active_indices exceeds the cache batch size")
+        if active_indices.numel() and ((active_indices < 0).any() or (active_indices >= self.max_batch_size).any()):
+            raise ValueError("active_indices contains an out-of-range row")
+        if active_indices.unique().numel() != active_indices.numel():
+            raise ValueError("active_indices must not contain duplicates")
+
+        count = active_indices.numel()
+        if count:
+            keys = self.keys.index_select(0, active_indices).clone()
+            values = self.values.index_select(0, active_indices).clone()
+            self.keys[:count].copy_(keys)
+            self.values[:count].copy_(values)
+        if count < self.max_batch_size:
+            self.keys[count:].zero_()
+            self.values[count:].zero_()
+        if self._write_position is not None and self._write_position.dim() == 1:
+            positions = self._write_position
+            if positions.numel() != self.max_batch_size:
+                raise ValueError("per-row write positions must match the cache batch size")
+            position_indices = active_indices.to(positions.device)
+            compacted = positions.index_select(0, position_indices).clone() if count else positions[:0]
+            self._write_position[:count].copy_(compacted)
+            if count < self.max_batch_size:
+                self._write_position[count:].fill_(-1)
+
     def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
         if self.is_initialized:
             self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))
             self.values = self.values.index_select(0, beam_idx.to(self.values.device))
+            if self._write_position is not None and self._write_position.dim() == 1:
+                self._write_position = self._write_position.index_select(0, beam_idx.to(self._write_position.device))
 
 
 class DeepSpeedStaticCache:
@@ -173,15 +231,20 @@ def layers(self):
     def set_write_position(self, pos: torch.Tensor):
         """Set the write position shared by all layers.
 
-        Must be called before each graph replay with the decode step position
-        as a scalar ``torch.long`` tensor on the correct device.  The tensor is
-        stored by reference so subsequent in-place updates (e.g.
-        ``pos.fill_(new_val)``) are immediately visible to all layers.
+        Must be called before each graph replay. A scalar tensor preserves the
+        fixed-batch path; a vector supplies one decode position per cache row.
+        The tensor is stored by reference so subsequent in-place updates are
+        immediately visible to all layers.
         """
         self._write_position = pos
         for layer in self._layers:
             layer.set_write_position(pos)
 
+    def compact(self, active_indices: torch.Tensor) -> None:
+        """Compact active rows after requests retire from a batch."""
+        for layer in self._layers:
+            layer.compact(active_indices)
+
     def update(
         self,
         key_states: torch.Tensor,
@@ -207,7 +270,7 @@ def early_initialization(
             fake_v = torch.zeros((batch_size, num_heads, 0, head_dim), dtype=dtype, device=device)
             layer.lazy_initialization(fake_k, fake_v)
 
-    def get_seq_length(self, layer_idx: int = 0) -> int:
+    def get_seq_length(self, layer_idx: int = 0) -> int | torch.Tensor:
         if layer_idx >= len(self._layers):
             return 0
         return self._layers[layer_idx].get_seq_length()
diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst
index a1fae868d0c6..8c3902067253 100644
--- a/docs/code-docs/source/inference-engine.rst
+++ b/docs/code-docs/source/inference-engine.rst
@@ -57,3 +57,20 @@ inference tensor-parallel size 1, an internal KV cache, and a prompt longer than
 one token. It cannot be combined with CUDA graph capture or
 ``release_inference_cache``. Sampling still happens independently for every
 response branch after the shared prompt forward.
+
+Continuous-batching prototype
+-----------------------------
+
+``deepspeed.runtime.rollout.continuous_batching`` provides the scheduling and
+slot-lifecycle primitive needed to build a continuous decode batch. Requests
+are admitted in FIFO order up to a configured capacity. When a request retires,
+the update identifies the surviving cache rows to compact and the pending
+requests that can be prefetched into the newly free rows. The model backend is
+responsible for applying that update, running prompt prefill, and constructing
+the attention metadata for the active rows.
+
+The prototype intentionally does not implement paged attention or change the
+default ``HybridEngineRollout.generate`` path. ``DeepSpeedStaticCache`` accepts
+one write position per row and can compact active rows while preserving its
+static tensor addresses. This mirrors the scheduler/cache separation used by
+systems such as vLLM and SGLang without copying their backend-specific kernels.
diff --git a/tests/unit/runtime/rollout/test_continuous_batching.py b/tests/unit/runtime/rollout/test_continuous_batching.py
new file mode 100644
index 000000000000..04e7f4c775a9
--- /dev/null
+++ b/tests/unit/runtime/rollout/test_continuous_batching.py
@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler)
+
+
+def _request(request_id, max_new_tokens=3):
+    return ContinuousBatchRequest(request_id, max_new_tokens)
+
+
+def test_scheduler_admits_fifo_and_respects_capacity():
+    scheduler = ContinuousBatchScheduler(max_batch_size=2)
+    scheduler.submit(_request("a"))
+    scheduler.submit(_request("b"))
+    scheduler.submit(_request("c"))
+
+    update = scheduler.schedule()
+    assert update.active_ids == ("a", "b")
+    assert update.keep_slots == ()
+    assert update.admitted == (_request("a"), _request("b"))
+    assert update.admitted_slots == (0, 1)
+    assert scheduler.pending == (_request("c"), )
+
+
+def test_scheduler_compacts_survivors_and_admits_pending_request():
+    scheduler = ContinuousBatchScheduler(max_batch_size=2)
+    scheduler.submit(_request("a"))
+    scheduler.submit(_request("b"))
+    scheduler.submit(_request("c"))
+    scheduler.schedule()
+
+    update = scheduler.schedule(finished_ids=("a", ))
+    assert update.keep_slots == (1, )
+    assert update.retired == ("a", )
+    assert update.admitted == (_request("c"), )
+    assert update.admitted_slots == (1, )
+    assert update.active_ids == ("b", "c")
+
+
+def test_scheduler_advance_retires_by_budget():
+    scheduler = ContinuousBatchScheduler(max_batch_size=2)
+    scheduler.submit(_request("a", max_new_tokens=1))
+    scheduler.submit(_request("b", max_new_tokens=3))
+    scheduler.schedule()
+
+    update = scheduler.advance()
+    assert update.retired == ("a", )
+    assert update.keep_slots == (1, )
+    assert update.active_ids == ("b", )
+    assert scheduler.active[0].request_id == "b"
+
+
+def test_scheduler_rejects_invalid_transitions():
+    with pytest.raises(ValueError, match="max_batch_size"):
+        ContinuousBatchScheduler(max_batch_size=0)
+    with pytest.raises(ValueError, match="max_new_tokens"):
+        _request("bad", max_new_tokens=0)
+
+    scheduler = ContinuousBatchScheduler(max_batch_size=1)
+    scheduler.submit(_request("a"))
+    with pytest.raises(ValueError, match="duplicate"):
+        scheduler.submit(_request("a"))
+    with pytest.raises(ValueError, match="not active"):
+        scheduler.schedule(finished_ids=("missing", ))
diff --git a/tests/unit/utils/test_static_cache.py b/tests/unit/utils/test_static_cache.py
new file mode 100644
index 000000000000..4240b9b17229
--- /dev/null
+++ b/tests/unit/utils/test_static_cache.py
@@ -0,0 +1,65 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.utils.static_cache import DeepSpeedStaticCache, DeepSpeedStaticLayer
+
+
+def test_static_layer_supports_per_row_decode_positions():
+    layer = DeepSpeedStaticLayer(max_cache_len=4)
+    keys = torch.zeros((2, 1, 1, 2))
+    values = torch.zeros_like(keys)
+    layer.lazy_initialization(keys, values)
+    layer.set_write_position(torch.tensor([1, 3], dtype=torch.long))
+
+    keys[:, :, 0, :] = torch.tensor([[[2.0, 2.0]], [[4.0, 4.0]]])
+    values.copy_(keys * 10)
+    layer.update(keys, values)
+
+    assert layer.keys[0, 0, 1].tolist() == [2.0, 2.0]
+    assert layer.keys[1, 0, 3].tolist() == [4.0, 4.0]
+    assert torch.equal(layer.get_seq_length(), torch.tensor([2, 4]))
+
+
+def test_static_cache_compact_preserves_rows_and_positions():
+    config = type("Config", (), {"num_hidden_layers": 1, "num_attention_heads": 1, "hidden_size": 2})()
+    cache = DeepSpeedStaticCache(config=config,
+                                 batch_size=3,
+                                 max_cache_len=4,
+                                 device=torch.device("cpu"),
+                                 dtype=torch.float32)
+    cache.set_write_position(torch.tensor([1, 2, 3], dtype=torch.long))
+    layer = cache.layers[0]
+    layer.keys[:, 0, 0, :] = torch.tensor([[10.0, 10.0], [20.0, 20.0], [30.0, 30.0]])
+    layer.values.copy_(layer.keys)
+
+    cache.compact(torch.tensor([2, 0], dtype=torch.long))
+
+    assert layer.keys[:2, 0, 0, 0].tolist() == [30.0, 10.0]
+    assert torch.equal(cache.get_seq_length(), torch.tensor([4, 2, 0]))
+    assert layer.keys[2].abs().sum().item() == 0
+
+
+def test_static_layer_rejects_mismatched_per_row_positions():
+    layer = DeepSpeedStaticLayer(max_cache_len=4)
+    keys = torch.zeros((2, 1, 1, 2))
+    layer.lazy_initialization(keys, keys)
+    layer.set_write_position(torch.tensor([1], dtype=torch.long))
+    with pytest.raises(ValueError, match="cover the active batch size"):
+        layer.update(keys, keys)
+
+
+def test_static_layer_accepts_active_prefix_of_cache_rows():
+    layer = DeepSpeedStaticLayer(max_cache_len=4)
+    keys = torch.zeros((3, 1, 1, 2))
+    layer.lazy_initialization(keys, keys)
+    layer.set_write_position(torch.tensor([1, 2, -1], dtype=torch.long))
+    active_keys = torch.ones((2, 1, 1, 2))
+
+    layer.update(active_keys, active_keys)
+
+    assert layer.keys[:2, 0, 1:3].sum().item() == 4

From f80610a8b9bf83aa3875386bd16d9b49c122969d Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Sun, 30 Aug 2026 08:47:34 +0800
Subject: [PATCH 02/12] feat: integrate continuous batching rollout prototype

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py  | 192 ++++++++++++++++++
 deepspeed/utils/static_cache.py               |  55 ++++-
 docs/code-docs/source/inference-engine.rst    |  16 +-
 .../rollout/test_hybrid_engine_rollout.py     |  19 ++
 tests/unit/utils/test_static_cache.py         |   2 +-
 5 files changed, 268 insertions(+), 16 deletions(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 0db4965f5077..9b1067404873 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -10,6 +10,8 @@
      Pre-allocates a StaticCache, captures the decode forward pass with a
      CUDA graph, and replays it for each decode step.  Eliminates kernel
      launch overhead.
+  3. **generate_continuous()**: a bounded greedy prototype that refills
+     retired cache rows with pending prompts.
 """
 
 import time
@@ -19,6 +21,7 @@
 
 from deepspeed.accelerator import get_accelerator
 from deepspeed.runtime.rollout.base import RolloutBatch, RolloutEngine, RolloutRequest, SamplingConfig
+from deepspeed.runtime.rollout.continuous_batching import ContinuousBatchRequest, ContinuousBatchScheduler
 
 
 @dataclass
@@ -167,6 +170,195 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout
 
         return rollout_batch
 
+    @torch.no_grad()
+    def generate_continuous(self, requests, sampling_configs, max_batch_size):
+        """Generate independent greedy requests in a continuously refilled batch.
+
+        This first integration targets the OPSD prototype: every request has a
+        single prompt row and one greedy response. Requests may use different
+        response budgets. Completed rows retire immediately and pending prompts
+        prefill into the released rows before the next decode step.
+        """
+        requests = tuple(requests)
+        sampling_configs = tuple(sampling_configs)
+        self._validate_continuous_inputs(requests, sampling_configs, max_batch_size)
+
+        from transformers import StaticCache
+        from deepspeed.utils.static_cache import DeepSpeedStaticCache
+
+        module = self.engine.module
+        device = requests[0].prompt_ids.device
+        prompt_len = requests[0].prompt_ids.shape[1]
+        model_dtype = next(module.parameters()).dtype
+        max_cache_len = prompt_len + sum(config.max_new_tokens for config in sampling_configs)
+        max_positions = getattr(module.config, "max_position_embeddings", max_cache_len)
+        if max_cache_len > max_positions:
+            raise ValueError("continuous batching cache exceeds the model maximum position embeddings")
+
+        scheduler = ContinuousBatchScheduler(max_batch_size)
+        request_by_id = {}
+        responses = {}
+        for request_id, (request, config) in enumerate(zip(requests, sampling_configs)):
+            scheduler.submit(ContinuousBatchRequest(request_id, config.max_new_tokens))
+            request_by_id[request_id] = request
+            responses[request_id] = []
+
+        cache = DeepSpeedStaticCache(
+            module.config,
+            batch_size=max_batch_size,
+            max_cache_len=max_cache_len,
+            device=device,
+            dtype=model_dtype,
+        )
+        write_positions = torch.full((max_batch_size, ), -1, dtype=torch.long, device=device)
+        cache.set_write_position(write_positions)
+        attention_mask = torch.zeros((max_batch_size, max_cache_len), dtype=torch.long, device=device)
+        next_tokens = {}
+        cache_position = prompt_len
+        update = scheduler.schedule()
+
+        while update.active:
+            keep_slots = torch.tensor(update.keep_slots, dtype=torch.long, device=device)
+            survivor_count = keep_slots.numel()
+            if survivor_count:
+                cache.compact(keep_slots)
+                survivor_attention = attention_mask.index_select(0, keep_slots).clone()
+                attention_mask.zero_()
+                attention_mask[:survivor_count].copy_(survivor_attention)
+            else:
+                cache.reset()
+                write_positions.fill_(-1)
+                attention_mask.zero_()
+                cache_position = prompt_len
+
+            admitted_tokens = self._continuous_prefill(
+                module,
+                StaticCache,
+                cache,
+                update,
+                request_by_id,
+                attention_mask,
+                cache_position,
+                prompt_len,
+                model_dtype,
+                device,
+            )
+
+            decoded_tokens = {}
+            if survivor_count:
+                survivor_ids = update.active_ids[:survivor_count]
+                decode_input = torch.cat([next_tokens[request_id] for request_id in survivor_ids], dim=0)
+                write_positions[:survivor_count].fill_(cache_position)
+                position_ids = attention_mask[:survivor_count, :cache_position].sum(dim=1, keepdim=True)
+                attention_mask[:survivor_count, cache_position] = 1
+                output = module(
+                    decode_input,
+                    attention_mask=attention_mask[:survivor_count],
+                    past_key_values=cache,
+                    use_cache=True,
+                    cache_position=torch.tensor([cache_position], dtype=torch.long, device=device),
+                    position_ids=position_ids,
+                )
+                decoded = output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
+                decoded_tokens = dict(zip(survivor_ids, decoded.split(1, dim=0)))
+
+            next_tokens = decoded_tokens | admitted_tokens
+            finished_ids = []
+            for request_id in update.active_ids:
+                token = next_tokens[request_id]
+                responses[request_id].append(token)
+                if self._is_eos(token):
+                    finished_ids.append(request_id)
+
+            update = scheduler.advance(finished_ids)
+            if survivor_count:
+                cache_position += 1
+
+        return [
+            self._build_continuous_output(request, responses[request_id])
+            for request_id, request in enumerate(requests)
+        ]
+
+    def _validate_continuous_inputs(self, requests, sampling_configs, max_batch_size):
+        if not requests:
+            raise ValueError("continuous batching requires at least one request")
+        if len(requests) != len(sampling_configs):
+            raise ValueError("requests and sampling_configs must have the same length")
+        if max_batch_size <= 0:
+            raise ValueError("max_batch_size must be positive")
+        if self.use_graph_capture:
+            raise ValueError("continuous batching does not yet support CUDA graph capture")
+
+        prompt_len = requests[0].prompt_ids.shape[1]
+        device = requests[0].prompt_ids.device
+        for request, config in zip(requests, sampling_configs):
+            if request.prompt_ids.shape[0] != 1:
+                raise ValueError("continuous batching requires one prompt row per request")
+            if request.prompt_ids.shape[1] != prompt_len:
+                raise ValueError("continuous batching currently requires equal prompt widths")
+            if request.prompt_ids.device != device:
+                raise ValueError("continuous batching requests must use the same device")
+            if config.temperature > 0:
+                raise ValueError("continuous batching currently supports greedy decoding only")
+            if config.n_samples_per_prompt != 1:
+                raise ValueError("continuous batching currently supports one sample per prompt")
+
+    def _continuous_prefill(self, module, static_cache_type, cache, update, request_by_id, attention_mask,
+                            cache_position, prompt_len, model_dtype, device):
+        if not update.admitted:
+            return {}
+
+        admitted_ids = tuple(request.request_id for request in update.admitted)
+        prompt_ids = torch.cat([request_by_id[request_id].prompt_ids for request_id in admitted_ids], dim=0)
+        prompt_attention = torch.cat([request_by_id[request_id].prompt_attention_mask for request_id in admitted_ids],
+                                     dim=0)
+        prefill_cache = static_cache_type(
+            config=module.config,
+            batch_size=len(admitted_ids),
+            max_cache_len=prompt_len,
+            device=device,
+            dtype=model_dtype,
+        )
+        prefill_output = module(
+            prompt_ids,
+            attention_mask=prompt_attention,
+            past_key_values=prefill_cache,
+            use_cache=True,
+            cache_position=torch.arange(prompt_len, device=device),
+        )
+        prefill_tokens = prefill_output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
+        cache_start = cache_position - prompt_len
+        for layer_idx, prefill_layer in enumerate(prefill_cache.layers):
+            target_layer = cache.layers[layer_idx]
+            for source_row, target_row in enumerate(update.admitted_slots):
+                target_layer.keys[target_row, :, cache_start:cache_position].copy_(prefill_layer.keys[source_row])
+                target_layer.values[target_row, :, cache_start:cache_position].copy_(prefill_layer.values[source_row])
+        for source_row, target_row in enumerate(update.admitted_slots):
+            attention_mask[target_row, cache_start:cache_position].copy_(prompt_attention[source_row])
+            write_positions = cache._write_position
+            write_positions[target_row] = cache_position
+        return dict(zip(admitted_ids, prefill_tokens.split(1, dim=0)))
+
+    def _is_eos(self, token):
+        eos_token_id = self.tokenizer.eos_token_id
+        if eos_token_id is None:
+            return False
+        eos_ids = torch.as_tensor(eos_token_id, dtype=token.dtype, device=token.device).flatten()
+        return bool((token == eos_ids).any().item())
+
+    @staticmethod
+    def _build_continuous_output(request, response_tokens):
+        response_ids = torch.cat(response_tokens, dim=1)
+        input_ids = torch.cat((request.prompt_ids, response_ids), dim=1)
+        response_attention = torch.ones_like(response_ids)
+        attention_mask = torch.cat((request.prompt_attention_mask, response_attention), dim=1)
+        response_start = request.prompt_ids.shape[1]
+        return RolloutBatch(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            response_start_idx=torch.tensor([response_start], dtype=torch.long, device=input_ids.device),
+        )
+
     def get_last_profile(self):
         """Return the most recent profiling snapshot for this rollout instance."""
         return self._last_profile
diff --git a/deepspeed/utils/static_cache.py b/deepspeed/utils/static_cache.py
index 90ee8c80d790..e60b42f1f80f 100644
--- a/deepspeed/utils/static_cache.py
+++ b/deepspeed/utils/static_cache.py
@@ -101,7 +101,8 @@ def update(
             row_indices = torch.arange(key_states.shape[0], device=self.device)
             self.keys[row_indices, :, cache_position, :] = key_states[:, :, 0, :]
             self.values[row_indices, :, cache_position, :] = value_states[:, :, 0, :]
-            return self.keys, self.values
+            active_batch_size = key_states.shape[0]
+            return self.keys[:active_batch_size], self.values[:active_batch_size]
 
         if self._write_position is not None:
             cache_position = torch.arange(kv_length, device=self.device) + self._write_position
@@ -156,14 +157,7 @@ def compact(self, active_indices: torch.Tensor) -> None:
             raise ValueError("active_indices must not contain duplicates")
 
         count = active_indices.numel()
-        if count:
-            keys = self.keys.index_select(0, active_indices).clone()
-            values = self.values.index_select(0, active_indices).clone()
-            self.keys[:count].copy_(keys)
-            self.values[:count].copy_(values)
-        if count < self.max_batch_size:
-            self.keys[count:].zero_()
-            self.values[count:].zero_()
+        self._compact_rows(active_indices, count)
         if self._write_position is not None and self._write_position.dim() == 1:
             positions = self._write_position
             if positions.numel() != self.max_batch_size:
@@ -174,6 +168,19 @@ def compact(self, active_indices: torch.Tensor) -> None:
             if count < self.max_batch_size:
                 self._write_position[count:].fill_(-1)
 
+    def _compact_rows(self, active_indices: torch.Tensor, count: int | None = None) -> None:
+        """Compact only the row tensors; used by the multi-layer cache."""
+        if count is None:
+            count = active_indices.numel()
+        if count:
+            keys = self.keys.index_select(0, active_indices).clone()
+            values = self.values.index_select(0, active_indices).clone()
+            self.keys[:count].copy_(keys)
+            self.values[:count].copy_(values)
+        if count < self.max_batch_size:
+            self.keys[count:].zero_()
+            self.values[count:].zero_()
+
     def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
         if self.is_initialized:
             self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))
@@ -242,8 +249,28 @@ def set_write_position(self, pos: torch.Tensor):
 
     def compact(self, active_indices: torch.Tensor) -> None:
         """Compact active rows after requests retire from a batch."""
+        if not isinstance(active_indices, torch.Tensor) or active_indices.dim() != 1:
+            raise ValueError("active_indices must be a 1-D tensor")
+        max_batch_size = self._layers[0].max_batch_size if self._layers else 0
+        active_indices = active_indices.to(device=self._layers[0].keys.device, dtype=torch.long)
+        if active_indices.numel() > max_batch_size:
+            raise ValueError("active_indices exceeds the cache batch size")
+        if active_indices.numel() and ((active_indices < 0).any() or (active_indices >= max_batch_size).any()):
+            raise ValueError("active_indices contains an out-of-range row")
+        if active_indices.unique().numel() != active_indices.numel():
+            raise ValueError("active_indices must not contain duplicates")
         for layer in self._layers:
-            layer.compact(active_indices)
+            layer._compact_rows(active_indices)
+        if self._write_position is not None and self._write_position.dim() == 1:
+            positions = self._write_position
+            if positions.numel() != max_batch_size:
+                raise ValueError("per-row write positions must match the cache batch size")
+            position_indices = active_indices.to(positions.device)
+            count = active_indices.numel()
+            compacted = positions.index_select(0, position_indices).clone() if count else positions[:0]
+            positions[:count].copy_(compacted)
+            if count < positions.numel():
+                positions[count:].fill_(-1)
 
     def update(
         self,
@@ -273,7 +300,13 @@ def early_initialization(
     def get_seq_length(self, layer_idx: int = 0) -> int | torch.Tensor:
         if layer_idx >= len(self._layers):
             return 0
-        return self._layers[layer_idx].get_seq_length()
+        length = self._layers[layer_idx].get_seq_length()
+        if isinstance(length, torch.Tensor) and length.dim() == 1:
+            # Transformer decoder implementations use one scalar past length
+            # to build the causal mask. Continuous batching keeps rows aligned
+            # to a shared physical position, so the largest row is sufficient.
+            return length.max()
+        return length
 
     def get_max_cache_shape(self, layer_idx: int = 0) -> int:
         if layer_idx >= len(self._layers):
diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst
index 8c3902067253..613d8cc0397c 100644
--- a/docs/code-docs/source/inference-engine.rst
+++ b/docs/code-docs/source/inference-engine.rst
@@ -70,7 +70,15 @@ responsible for applying that update, running prompt prefill, and constructing
 the attention metadata for the active rows.
 
 The prototype intentionally does not implement paged attention or change the
-default ``HybridEngineRollout.generate`` path. ``DeepSpeedStaticCache`` accepts
-one write position per row and can compact active rows while preserving its
-static tensor addresses. This mirrors the scheduler/cache separation used by
-systems such as vLLM and SGLang without copying their backend-specific kernels.
+default ``HybridEngineRollout.generate`` path. For a first end-to-end trial,
+``HybridEngineRollout.generate_continuous`` accepts one request per prompt row
+and a matching list of greedy ``SamplingConfig`` objects. It dynamically
+prefills admitted prompts and decodes surviving rows until every request has
+finished. CUDA Graph capture, sampling, multiple samples per prompt, and
+different prompt widths are intentionally rejected until the scheduling
+semantics are validated on real workloads.
+
+``DeepSpeedStaticCache`` accepts one write position per row and can compact
+active rows while preserving its static tensor addresses. This mirrors the
+scheduler/cache separation used by systems such as vLLM and SGLang without
+copying their backend-specific kernels.
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index 475c73f5854d..3b132c41abee 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -80,6 +80,25 @@ def test_constructor_defaults_without_cfg():
     assert rollout.use_shared_prefill is False
 
 
+def test_continuous_generation_rejects_unsupported_inputs():
+    rollout = HybridEngineRollout(_make_engine(), _make_tokenizer())
+    request = RolloutRequest(
+        prompt_ids=torch.tensor([[0, 1, 2]]),
+        prompt_attention_mask=torch.tensor([[0, 1, 1]]),
+    )
+
+    with pytest.raises(ValueError, match="at least one request"):
+        rollout.generate_continuous([], [], max_batch_size=1)
+    with pytest.raises(ValueError, match="same length"):
+        rollout.generate_continuous([request], [], max_batch_size=1)
+    with pytest.raises(ValueError, match="greedy"):
+        rollout.generate_continuous(
+            [request],
+            [SamplingConfig(max_new_tokens=2, temperature=0.5)],
+            max_batch_size=1,
+        )
+
+
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.time.perf_counter")
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.get_accelerator")
 def test_generate_records_profile_when_enabled(mock_get_accelerator, mock_perf_counter):
diff --git a/tests/unit/utils/test_static_cache.py b/tests/unit/utils/test_static_cache.py
index 4240b9b17229..5fed81203236 100644
--- a/tests/unit/utils/test_static_cache.py
+++ b/tests/unit/utils/test_static_cache.py
@@ -40,7 +40,7 @@ def test_static_cache_compact_preserves_rows_and_positions():
     cache.compact(torch.tensor([2, 0], dtype=torch.long))
 
     assert layer.keys[:2, 0, 0, 0].tolist() == [30.0, 10.0]
-    assert torch.equal(cache.get_seq_length(), torch.tensor([4, 2, 0]))
+    assert cache.get_seq_length().item() == 4
     assert layer.keys[2].abs().sum().item() == 0
 
 

From bb5f6607569bd1cef178d77b3b9991ad9d0cf3a8 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Sun, 30 Aug 2026 08:50:49 +0800
Subject: [PATCH 03/12] fix: support StaticCache batch API variants

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py  | 34 +++++++++++--------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 9b1067404873..22443e11ffe8 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -16,6 +16,7 @@
 
 import time
 from dataclasses import dataclass
+from inspect import signature
 
 import torch
 
@@ -312,13 +313,8 @@ def _continuous_prefill(self, module, static_cache_type, cache, update, request_
         prompt_ids = torch.cat([request_by_id[request_id].prompt_ids for request_id in admitted_ids], dim=0)
         prompt_attention = torch.cat([request_by_id[request_id].prompt_attention_mask for request_id in admitted_ids],
                                      dim=0)
-        prefill_cache = static_cache_type(
-            config=module.config,
-            batch_size=len(admitted_ids),
-            max_cache_len=prompt_len,
-            device=device,
-            dtype=model_dtype,
-        )
+        prefill_cache = self._create_static_cache(static_cache_type, module.config, len(admitted_ids), prompt_len,
+                                                  device, model_dtype)
         prefill_output = module(
             prompt_ids,
             attention_mask=prompt_attention,
@@ -346,6 +342,22 @@ def _is_eos(self, token):
         eos_ids = torch.as_tensor(eos_token_id, dtype=token.dtype, device=token.device).flatten()
         return bool((token == eos_ids).any().item())
 
+    @staticmethod
+    def _create_static_cache(static_cache_type, config, batch_size, max_cache_len, device, dtype):
+        """Construct StaticCache across Transformers' batch-size API variants."""
+        common_kwargs = {
+            "config": config,
+            "max_cache_len": max_cache_len,
+            "device": device,
+            "dtype": dtype,
+        }
+        parameters = signature(static_cache_type).parameters
+        if "batch_size" in parameters:
+            common_kwargs["batch_size"] = batch_size
+        elif "max_batch_size" in parameters:
+            common_kwargs["max_batch_size"] = batch_size
+        return static_cache_type(**common_kwargs)
+
     @staticmethod
     def _build_continuous_output(request, response_tokens):
         response_ids = torch.cat(response_tokens, dim=1)
@@ -439,13 +451,7 @@ def _generate_graph(self, prompt_ids, prompt_attn, max_new_tokens, pad_token_id,
         model_dtype = next(module.parameters()).dtype
 
         # --- Prefill with HF StaticCache (correct attention semantics) ---
-        prefill_cache = StaticCache(
-            config=module.config,
-            batch_size=batch_size,
-            max_cache_len=max_len,
-            device=device,
-            dtype=model_dtype,
-        )
+        prefill_cache = self._create_static_cache(StaticCache, module.config, batch_size, max_len, device, model_dtype)
         prefill_attn = torch.ones(batch_size, prompt_len, dtype=torch.long, device=device)
         prefill_attn[:, :prompt_len] = prompt_attn
         prefill_out = module(

From 8233436dbc3cb8a5ee86a725363786e9f68c6126 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Sun, 30 Aug 2026 08:55:00 +0800
Subject: [PATCH 04/12] fix: support legacy Transformers StaticCache configs

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py      |  7 ++++++-
 .../runtime/rollout/test_hybrid_engine_rollout.py | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 22443e11ffe8..7ad66d38d807 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -15,6 +15,7 @@
 """
 
 import time
+from copy import copy
 from dataclasses import dataclass
 from inspect import signature
 
@@ -345,8 +346,12 @@ def _is_eos(self, token):
     @staticmethod
     def _create_static_cache(static_cache_type, config, batch_size, max_cache_len, device, dtype):
         """Construct StaticCache across Transformers' batch-size API variants."""
+        cache_config = config
+        if not hasattr(config, "num_key_value_heads") or config.num_key_value_heads is None:
+            cache_config = copy(config)
+            cache_config.num_key_value_heads = config.num_attention_heads
         common_kwargs = {
-            "config": config,
+            "config": cache_config,
             "max_cache_len": max_cache_len,
             "device": device,
             "dtype": dtype,
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index 3b132c41abee..a10e0b969ee7 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -99,6 +99,21 @@ def test_continuous_generation_rejects_unsupported_inputs():
         )
 
 
+def test_static_cache_constructor_supports_legacy_batch_keyword():
+
+    class LegacyStaticCache:
+
+        def __init__(self, config, max_batch_size, max_cache_len, device, dtype):
+            self.config = config
+            self.max_batch_size = max_batch_size
+
+    config = SimpleNamespace(num_attention_heads=4)
+    cache = HybridEngineRollout._create_static_cache(LegacyStaticCache, config, 2, 8, "cpu", torch.float32)
+
+    assert cache.max_batch_size == 2
+    assert cache.config.num_key_value_heads == 4
+
+
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.time.perf_counter")
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.get_accelerator")
 def test_generate_records_profile_when_enabled(mock_get_accelerator, mock_perf_counter):

From a434080db8fd4b880b42216a6dd2ecec5db8b514 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Sun, 30 Aug 2026 09:24:28 +0800
Subject: [PATCH 05/12] fix: support legacy model forward cache arguments

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py      | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 7ad66d38d807..7c4ef3065759 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -253,7 +253,8 @@ def generate_continuous(self, requests, sampling_configs, max_batch_size):
                 write_positions[:survivor_count].fill_(cache_position)
                 position_ids = attention_mask[:survivor_count, :cache_position].sum(dim=1, keepdim=True)
                 attention_mask[:survivor_count, cache_position] = 1
-                output = module(
+                output = self._call_model(
+                    module,
                     decode_input,
                     attention_mask=attention_mask[:survivor_count],
                     past_key_values=cache,
@@ -316,7 +317,8 @@ def _continuous_prefill(self, module, static_cache_type, cache, update, request_
                                      dim=0)
         prefill_cache = self._create_static_cache(static_cache_type, module.config, len(admitted_ids), prompt_len,
                                                   device, model_dtype)
-        prefill_output = module(
+        prefill_output = self._call_model(
+            module,
             prompt_ids,
             attention_mask=prompt_attention,
             past_key_values=prefill_cache,
@@ -343,6 +345,15 @@ def _is_eos(self, token):
         eos_ids = torch.as_tensor(eos_token_id, dtype=token.dtype, device=token.device).flatten()
         return bool((token == eos_ids).any().item())
 
+    @staticmethod
+    def _call_model(module, input_ids, **kwargs):
+        """Call models with only the cache arguments their HF version accepts."""
+        parameters = signature(module.forward).parameters
+        accepts_kwargs = any(parameter.kind == parameter.VAR_KEYWORD for parameter in parameters.values())
+        if not accepts_kwargs:
+            kwargs = {name: value for name, value in kwargs.items() if name in parameters}
+        return module(input_ids, **kwargs)
+
     @staticmethod
     def _create_static_cache(static_cache_type, config, batch_size, max_cache_len, device, dtype):
         """Construct StaticCache across Transformers' batch-size API variants."""

From 67b246ef19d0bdfb21d50673ba0b8ce9b50d354d Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Sun, 30 Aug 2026 09:30:20 +0800
Subject: [PATCH 06/12] fix: support legacy KV cache in continuous batching

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py  | 120 +++++++++++++++++-
 .../rollout/test_hybrid_engine_rollout.py     |  38 ++++++
 2 files changed, 153 insertions(+), 5 deletions(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 7c4ef3065759..e1e441545271 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -185,17 +185,20 @@ def generate_continuous(self, requests, sampling_configs, max_batch_size):
         sampling_configs = tuple(sampling_configs)
         self._validate_continuous_inputs(requests, sampling_configs, max_batch_size)
 
-        from transformers import StaticCache
-        from deepspeed.utils.static_cache import DeepSpeedStaticCache
-
         module = self.engine.module
-        device = requests[0].prompt_ids.device
         prompt_len = requests[0].prompt_ids.shape[1]
-        model_dtype = next(module.parameters()).dtype
         max_cache_len = prompt_len + sum(config.max_new_tokens for config in sampling_configs)
         max_positions = getattr(module.config, "max_position_embeddings", max_cache_len)
         if max_cache_len > max_positions:
             raise ValueError("continuous batching cache exceeds the model maximum position embeddings")
+        if not getattr(module, "_supports_cache_class", False):
+            return self._generate_continuous_legacy(requests, sampling_configs, max_batch_size)
+
+        from transformers import StaticCache
+        from deepspeed.utils.static_cache import DeepSpeedStaticCache
+
+        device = requests[0].prompt_ids.device
+        model_dtype = next(module.parameters()).dtype
 
         scheduler = ContinuousBatchScheduler(max_batch_size)
         request_by_id = {}
@@ -282,6 +285,113 @@ def generate_continuous(self, requests, sampling_configs, max_batch_size):
             for request_id, request in enumerate(requests)
         ]
 
+    def _generate_continuous_legacy(self, requests, sampling_configs, max_batch_size):
+        """Continuous decode for Transformers models that return legacy KV tuples."""
+        module = self.engine.module
+        device = requests[0].prompt_ids.device
+        prompt_len = requests[0].prompt_ids.shape[1]
+        scheduler, request_by_id, responses = self._create_continuous_scheduler(requests, sampling_configs,
+                                                                                max_batch_size)
+        next_tokens = {}
+        attention_mask = None
+        past_key_values = None
+        update = scheduler.schedule()
+
+        while update.active:
+            survivor_count = len(update.keep_slots)
+            survivor_ids = update.active_ids[:survivor_count]
+            survivor_cache = None
+            survivor_attention = None
+            decoded_tokens = {}
+            if survivor_count:
+                keep_slots = torch.tensor(update.keep_slots, dtype=torch.long, device=device)
+                survivor_cache = self._select_legacy_cache_rows(past_key_values, keep_slots)
+                survivor_attention = attention_mask.index_select(0, keep_slots)
+                decode_input = torch.cat([next_tokens[request_id] for request_id in survivor_ids], dim=0)
+                survivor_attention = torch.cat(
+                    (survivor_attention, torch.ones((survivor_count, 1), dtype=torch.long, device=device)), dim=1)
+                decode_output = self._call_model(
+                    module,
+                    decode_input,
+                    attention_mask=survivor_attention,
+                    past_key_values=survivor_cache,
+                    use_cache=True,
+                )
+                survivor_cache = decode_output.past_key_values
+                decoded = decode_output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
+                decoded_tokens = dict(zip(survivor_ids, decoded.split(1, dim=0)))
+
+            admitted_cache, admitted_attention, admitted_tokens = self._legacy_prefill(module, update, request_by_id)
+            past_key_values, attention_mask = self._merge_legacy_cache(
+                survivor_cache,
+                survivor_attention,
+                admitted_cache,
+                admitted_attention,
+                prompt_len,
+            )
+            next_tokens = decoded_tokens | admitted_tokens
+            finished_ids = []
+            for request_id in update.active_ids:
+                token = next_tokens[request_id]
+                responses[request_id].append(token)
+                if self._is_eos(token):
+                    finished_ids.append(request_id)
+            update = scheduler.advance(finished_ids)
+
+        return [
+            self._build_continuous_output(request, responses[request_id])
+            for request_id, request in enumerate(requests)
+        ]
+
+    @staticmethod
+    def _create_continuous_scheduler(requests, sampling_configs, max_batch_size):
+        scheduler = ContinuousBatchScheduler(max_batch_size)
+        request_by_id = {}
+        responses = {}
+        for request_id, (request, config) in enumerate(zip(requests, sampling_configs)):
+            scheduler.submit(ContinuousBatchRequest(request_id, config.max_new_tokens))
+            request_by_id[request_id] = request
+            responses[request_id] = []
+        return scheduler, request_by_id, responses
+
+    def _legacy_prefill(self, module, update, request_by_id):
+        if not update.admitted:
+            return None, None, {}
+        admitted_ids = tuple(request.request_id for request in update.admitted)
+        prompt_ids = torch.cat([request_by_id[request_id].prompt_ids for request_id in admitted_ids], dim=0)
+        prompt_attention = torch.cat(
+            [request_by_id[request_id].prompt_attention_mask for request_id in admitted_ids], dim=0)
+        output = self._call_model(module, prompt_ids, attention_mask=prompt_attention, use_cache=True)
+        tokens = output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
+        return output.past_key_values, prompt_attention, dict(zip(admitted_ids, tokens.split(1, dim=0)))
+
+    @staticmethod
+    def _select_legacy_cache_rows(past_key_values, rows):
+        return tuple((keys.index_select(0, rows), values.index_select(0, rows)) for keys, values in past_key_values)
+
+    @staticmethod
+    def _merge_legacy_cache(survivor_cache, survivor_attention, admitted_cache, admitted_attention, prompt_len):
+        if survivor_cache is None:
+            return admitted_cache, admitted_attention
+        if admitted_cache is None:
+            return survivor_cache, survivor_attention
+
+        cache_len = survivor_cache[0][0].shape[2]
+        left_padding = cache_len - prompt_len
+        padded_admitted = tuple(
+            (torch.nn.functional.pad(keys, (0, 0, left_padding, 0)),
+             torch.nn.functional.pad(values, (0, 0, left_padding, 0))) for keys, values in admitted_cache)
+        merged_cache = tuple((torch.cat((survivor_keys, admitted_keys), dim=0),
+                              torch.cat((survivor_values, admitted_values), dim=0))
+                             for (survivor_keys, survivor_values), (admitted_keys,
+                                                                    admitted_values) in zip(survivor_cache,
+                                                                                            padded_admitted))
+        admitted_padding = torch.zeros((admitted_attention.shape[0], left_padding),
+                                       dtype=admitted_attention.dtype,
+                                       device=admitted_attention.device)
+        padded_attention = torch.cat((admitted_padding, admitted_attention), dim=1)
+        return merged_cache, torch.cat((survivor_attention, padded_attention), dim=0)
+
     def _validate_continuous_inputs(self, requests, sampling_configs, max_batch_size):
         if not requests:
             raise ValueError("continuous batching requires at least one request")
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index a10e0b969ee7..90fb4f7883d8 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -114,6 +114,44 @@ def __init__(self, config, max_batch_size, max_cache_len, device, dtype):
     assert cache.config.num_key_value_heads == 4
 
 
+def test_continuous_generation_refills_legacy_cache_batch():
+    class LegacyCacheModel(torch.nn.Module):
+        _supports_cache_class = False
+
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(1))
+            self.calls = []
+            self.config = SimpleNamespace(max_position_embeddings=32)
+
+        def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=True):
+            past_length = 0 if past_key_values is None else past_key_values[0][0].shape[2]
+            total_length = past_length + input_ids.shape[1]
+            assert attention_mask.shape == (input_ids.shape[0], total_length)
+            self.calls.append((input_ids.shape[0], input_ids.shape[1], past_length))
+            keys = torch.zeros((input_ids.shape[0], 1, total_length, 1))
+            logits = torch.zeros((input_ids.shape[0], input_ids.shape[1], 16))
+            logits[..., 7] = 1
+            return SimpleNamespace(logits=logits, past_key_values=((keys, keys.clone()), ))
+
+    model = LegacyCacheModel()
+    rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None))
+    requests = [
+        RolloutRequest(torch.tensor([[1, 2, token]]), torch.ones((1, 3), dtype=torch.long))
+        for token in (3, 4, 5)
+    ]
+    configs = [
+        SamplingConfig(max_new_tokens=1, temperature=0),
+        SamplingConfig(max_new_tokens=3, temperature=0),
+        SamplingConfig(max_new_tokens=2, temperature=0),
+    ]
+
+    outputs = rollout.generate_continuous(requests, configs, max_batch_size=2)
+
+    assert [output.input_ids.shape[1] - 3 for output in outputs] == [1, 3, 2]
+    assert model.calls[:3] == [(2, 3, 0), (1, 1, 3), (1, 3, 0)]
+
+
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.time.perf_counter")
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.get_accelerator")
 def test_generate_records_profile_when_enabled(mock_get_accelerator, mock_perf_counter):

From a9bc1d213daf7153147d1aa1c59c4bc9828baf17 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Wed, 2 Sep 2026 20:56:31 +0800
Subject: [PATCH 07/12] fix: handle continuous batching cache spans

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py  | 91 +++++++++++++++----
 .../rollout/test_hybrid_engine_rollout.py     | 41 ++++++++-
 2 files changed, 114 insertions(+), 18 deletions(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 8bbaae751ca6..5ffb1714702b 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -256,9 +256,18 @@ def generate_continuous(self, requests, sampling_configs, max_batch_size):
 
         module = self.engine.module
         prompt_len = requests[0].prompt_ids.shape[1]
-        max_cache_len = prompt_len + sum(config.max_new_tokens for config in sampling_configs)
-        max_positions = getattr(module.config, "max_position_embeddings", max_cache_len)
-        if max_cache_len > max_positions:
+        max_positions = getattr(module.config, "max_position_embeddings", None)
+        if max_positions is not None:
+            for config in sampling_configs:
+                logical_length = prompt_len + config.max_new_tokens
+                if logical_length > max_positions:
+                    raise ValueError("continuous batching request exceeds the model maximum position embeddings")
+        max_cache_len = self._estimate_continuous_cache_len(
+            prompt_len,
+            [config.max_new_tokens for config in sampling_configs],
+            max_batch_size,
+        )
+        if max_positions is not None and max_cache_len > max_positions:
             raise ValueError("continuous batching cache exceeds the model maximum position embeddings")
         if not getattr(module, "_supports_cache_class", False):
             return self._generate_continuous_legacy(requests, sampling_configs, max_batch_size)
@@ -374,7 +383,7 @@ def _generate_continuous_legacy(self, requests, sampling_configs, max_batch_size
             decoded_tokens = {}
             if survivor_count:
                 keep_slots = torch.tensor(update.keep_slots, dtype=torch.long, device=device)
-                survivor_cache = self._select_legacy_cache_rows(past_key_values, keep_slots)
+                survivor_cache = self._select_legacy_cache_rows(past_key_values, keep_slots, attention_mask.shape[0])
                 survivor_attention = attention_mask.index_select(0, keep_slots)
                 decode_input = torch.cat([next_tokens[request_id] for request_id in survivor_ids], dim=0)
                 survivor_attention = torch.cat(
@@ -428,15 +437,66 @@ def _legacy_prefill(self, module, update, request_by_id):
             return None, None, {}
         admitted_ids = tuple(request.request_id for request in update.admitted)
         prompt_ids = torch.cat([request_by_id[request_id].prompt_ids for request_id in admitted_ids], dim=0)
-        prompt_attention = torch.cat(
-            [request_by_id[request_id].prompt_attention_mask for request_id in admitted_ids], dim=0)
+        prompt_attention = torch.cat([request_by_id[request_id].prompt_attention_mask for request_id in admitted_ids],
+                                     dim=0)
         output = self._call_model(module, prompt_ids, attention_mask=prompt_attention, use_cache=True)
         tokens = output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
         return output.past_key_values, prompt_attention, dict(zip(admitted_ids, tokens.split(1, dim=0)))
 
     @staticmethod
-    def _select_legacy_cache_rows(past_key_values, rows):
-        return tuple((keys.index_select(0, rows), values.index_select(0, rows)) for keys, values in past_key_values)
+    def _select_legacy_cache_rows(past_key_values, rows, logical_batch_size):
+        if logical_batch_size <= 0:
+            raise ValueError("logical batch size must be positive")
+        selected_cache = []
+        for keys, values in past_key_values:
+            if keys.shape[0] != values.shape[0]:
+                raise ValueError("legacy KV cache key/value batch dimensions must match")
+            cache_batch_size = keys.shape[0]
+            if cache_batch_size == logical_batch_size:
+                cache_rows = rows
+            elif cache_batch_size > logical_batch_size and cache_batch_size % logical_batch_size == 0:
+                heads_per_request = cache_batch_size // logical_batch_size
+                head_offsets = torch.arange(heads_per_request, device=rows.device)
+                cache_rows = (rows[:, None] * heads_per_request + head_offsets).reshape(-1)
+            else:
+                raise ValueError("cannot identify the logical batch dimension in the legacy KV cache")
+            selected_cache.append((keys.index_select(0, cache_rows), values.index_select(0, cache_rows)))
+        return tuple(selected_cache)
+
+    @staticmethod
+    def _estimate_continuous_cache_len(prompt_len, max_new_tokens, max_batch_size):
+        """Estimate the largest cache span before the scheduler becomes empty."""
+        pending = list(max_new_tokens)
+        active = []
+        max_cache_len = prompt_len
+        busy_span = 0
+
+        while active or pending:
+            if not active:
+                active = pending[:max_batch_size]
+                pending = pending[max_batch_size:]
+                busy_span = 0
+
+            busy_span += 1
+            max_cache_len = max(max_cache_len, prompt_len + busy_span)
+            survivors = [remaining - 1 for remaining in active if remaining > 1]
+            free_slots = max_batch_size - len(survivors)
+            admitted = pending[:free_slots]
+            pending = pending[free_slots:]
+
+            if survivors:
+                active = survivors + admitted
+                continue
+
+            if admitted:
+                active = admitted
+                # All previous rows retired, so the implementation resets the cache
+                # before prefilling the newly admitted requests.
+                busy_span = 0
+            else:
+                active = []
+
+        return max_cache_len
 
     @staticmethod
     def _merge_legacy_cache(survivor_cache, survivor_attention, admitted_cache, admitted_attention, prompt_len):
@@ -447,14 +507,13 @@ def _merge_legacy_cache(survivor_cache, survivor_attention, admitted_cache, admi
 
         cache_len = survivor_cache[0][0].shape[2]
         left_padding = cache_len - prompt_len
-        padded_admitted = tuple(
-            (torch.nn.functional.pad(keys, (0, 0, left_padding, 0)),
-             torch.nn.functional.pad(values, (0, 0, left_padding, 0))) for keys, values in admitted_cache)
-        merged_cache = tuple((torch.cat((survivor_keys, admitted_keys), dim=0),
-                              torch.cat((survivor_values, admitted_values), dim=0))
-                             for (survivor_keys, survivor_values), (admitted_keys,
-                                                                    admitted_values) in zip(survivor_cache,
-                                                                                            padded_admitted))
+        padded_admitted = tuple((torch.nn.functional.pad(keys, (0, 0, left_padding, 0)),
+                                 torch.nn.functional.pad(values, (0, 0, left_padding, 0)))
+                                for keys, values in admitted_cache)
+        merged_cache = tuple(
+            (torch.cat((survivor_keys, admitted_keys), dim=0), torch.cat((survivor_values, admitted_values), dim=0))
+            for (survivor_keys, survivor_values), (admitted_keys,
+                                                   admitted_values) in zip(survivor_cache, padded_admitted))
         admitted_padding = torch.zeros((admitted_attention.shape[0], left_padding),
                                        dtype=admitted_attention.dtype,
                                        device=admitted_attention.device)
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index 061ba4d07edc..6c42f1c12413 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -114,7 +114,45 @@ def __init__(self, config, max_batch_size, max_cache_len, device, dtype):
     assert cache.config.num_key_value_heads == 4
 
 
+def test_continuous_cache_span_does_not_sum_independent_requests():
+    cache_len = HybridEngineRollout._estimate_continuous_cache_len(64, [64] * 100, 100)
+
+    assert cache_len == 128
+    assert cache_len < 64 + 64 * 100
+
+
+def test_select_legacy_cache_rows_expands_flattened_heads():
+    keys = torch.arange(2 * 4 * 3).reshape(2 * 4, 1, 3)
+    values = keys + 100
+    rows = torch.tensor([1])
+
+    selected = HybridEngineRollout._select_legacy_cache_rows(((keys, values), ), rows, logical_batch_size=2)
+
+    assert torch.equal(selected[0][0], keys[4:8])
+    assert torch.equal(selected[0][1], values[4:8])
+
+
+def test_continuous_generation_validates_each_request_length():
+
+    class LimitedModel(torch.nn.Module):
+
+        _supports_cache_class = False
+
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(1))
+            self.config = SimpleNamespace(max_position_embeddings=4)
+
+    model = LimitedModel()
+    rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(eos_token_id=None))
+    request = RolloutRequest(torch.tensor([[1, 2, 3]]), torch.ones((1, 3), dtype=torch.long))
+
+    with pytest.raises(ValueError, match="request exceeds"):
+        rollout.generate_continuous([request], [SamplingConfig(max_new_tokens=2, temperature=0)], 1)
+
+
 def test_continuous_generation_refills_legacy_cache_batch():
+
     class LegacyCacheModel(torch.nn.Module):
         _supports_cache_class = False
 
@@ -137,8 +175,7 @@ def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=Tru
     model = LegacyCacheModel()
     rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None))
     requests = [
-        RolloutRequest(torch.tensor([[1, 2, token]]), torch.ones((1, 3), dtype=torch.long))
-        for token in (3, 4, 5)
+        RolloutRequest(torch.tensor([[1, 2, token]]), torch.ones((1, 3), dtype=torch.long)) for token in (3, 4, 5)
     ]
     configs = [
         SamplingConfig(max_new_tokens=1, temperature=0),

From 9510db68ea9b611fc5905bce5eec36c966f729b7 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Mon, 7 Sep 2026 16:02:35 +0800
Subject: [PATCH 08/12] feat(rollout): unify continuous batching generation
 interface

Route continuous batching through generate() with a shared SamplingConfig and batch-level max_new_tokens budget. Return an ordered RolloutBatch, hide scheduler details from public exports, and update the rollout tests and documentation.

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 deepspeed/runtime/rollout/__init__.py         |   5 -
 deepspeed/runtime/rollout/base.py             |   2 +
 .../runtime/rollout/continuous_batching.py    |  14 +-
 .../runtime/rollout/hybrid_engine_rollout.py  | 122 +++++++++++-------
 docs/code-docs/source/inference-engine.rst    |  24 ++--
 .../rollout/test_continuous_batching.py       |  27 ++--
 .../rollout/test_hybrid_engine_rollout.py     |  52 ++++----
 .../runtime/rollout/test_rollout_interface.py |   8 ++
 8 files changed, 142 insertions(+), 112 deletions(-)

diff --git a/deepspeed/runtime/rollout/__init__.py b/deepspeed/runtime/rollout/__init__.py
index 97518d6f9070..16f6fc595da6 100644
--- a/deepspeed/runtime/rollout/__init__.py
+++ b/deepspeed/runtime/rollout/__init__.py
@@ -18,8 +18,6 @@
     SamplingConfig,
 )
 from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout
-from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler,
-                                                           ContinuousBatchUpdate)
 
 __all__ = [
     "HybridEngineRollout",
@@ -29,9 +27,6 @@
     "RolloutRequest",
     "SamplingConfig",
     "build_rollout",
-    "ContinuousBatchRequest",
-    "ContinuousBatchScheduler",
-    "ContinuousBatchUpdate",
 ]
 
 
diff --git a/deepspeed/runtime/rollout/base.py b/deepspeed/runtime/rollout/base.py
index abff6c6ccb12..94aba68e4a30 100644
--- a/deepspeed/runtime/rollout/base.py
+++ b/deepspeed/runtime/rollout/base.py
@@ -10,6 +10,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from typing import Optional
 
 import torch
 
@@ -32,6 +33,7 @@ class SamplingConfig:
     top_p: float = 1.0
     top_k: int = -1
     n_samples_per_prompt: int = 1
+    continuous_batch_size: Optional[int] = None
 
 
 @dataclass
diff --git a/deepspeed/runtime/rollout/continuous_batching.py b/deepspeed/runtime/rollout/continuous_batching.py
index 2b0071dd86ed..36c30ff64848 100644
--- a/deepspeed/runtime/rollout/continuous_batching.py
+++ b/deepspeed/runtime/rollout/continuous_batching.py
@@ -20,11 +20,6 @@ class ContinuousBatchRequest:
     """A request waiting for a slot in a continuous decode batch."""
 
     request_id: Hashable
-    max_new_tokens: int
-
-    def __post_init__(self) -> None:
-        if self.max_new_tokens <= 0:
-            raise ValueError("max_new_tokens must be positive")
 
 
 @dataclass(frozen=True)
@@ -55,16 +50,21 @@ def admitted_slots(self) -> tuple[int, ...]:
 class ContinuousBatchScheduler:
     """FIFO scheduler for bounded, slot-based continuous batching.
 
+    All submitted requests share the scheduler's ``max_new_tokens`` budget;
+    requests may still retire earlier when the model emits EOS.
     ``schedule`` performs admission/retirement without advancing tokens.
     ``advance`` represents one decode step for every active request and also
     retires requests whose token budget has been consumed. A caller may pass
     explicit finished IDs when the model emits EOS before that budget.
     """
 
-    def __init__(self, max_batch_size: int):
+    def __init__(self, max_batch_size: int, max_new_tokens: int):
         if max_batch_size <= 0:
             raise ValueError("max_batch_size must be positive")
+        if max_new_tokens <= 0:
+            raise ValueError("max_new_tokens must be positive")
         self.max_batch_size = max_batch_size
+        self.max_new_tokens = max_new_tokens
         self._pending = deque()
         self._active = []
         self._generated = {}
@@ -127,7 +127,7 @@ def advance(self, finished_ids=()) -> ContinuousBatchUpdate:
         for request, generated in self._active:
             generated += 1
             self._generated[request.request_id] = generated
-            if generated >= request.max_new_tokens:
+            if generated >= self.max_new_tokens:
                 finished.add(request.request_id)
             updated.append((request, generated))
         self._active = updated
diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 5ffb1714702b..1a76e1b54042 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -10,8 +10,8 @@
      Pre-allocates a StaticCache, captures the decode forward pass with a
      CUDA graph, and replays it for each decode step.  Eliminates kernel
      launch overhead.
-  3. **generate_continuous()**: a bounded greedy prototype that refills
-     retired cache rows with pending prompts.
+  3. **continuous batching**: an opt-in bounded greedy path selected through
+     ``SamplingConfig.continuous_batch_size``.
 """
 
 import time
@@ -101,6 +101,9 @@ def __init__(self, engine, tokenizer, cfg=None):
 
     @torch.no_grad()
     def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> RolloutBatch:
+        if sampling.continuous_batch_size is not None:
+            return self._generate_continuous(request, sampling, sampling.continuous_batch_size)
+
         device = request.prompt_ids.device
         B = request.prompt_ids.shape[0]
         n = sampling.n_samples_per_prompt
@@ -242,35 +245,36 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout
         return rollout_batch
 
     @torch.no_grad()
-    def generate_continuous(self, requests, sampling_configs, max_batch_size):
+    def _generate_continuous(self, request, sampling, max_batch_size):
         """Generate independent greedy requests in a continuously refilled batch.
 
         This first integration targets the OPSD prototype: every request has a
-        single prompt row and one greedy response. Requests may use different
-        response budgets. Completed rows retire immediately and pending prompts
-        prefill into the released rows before the next decode step.
+        single prompt row and one greedy response. Completed rows retire
+        immediately and pending prompts prefill into the released rows before
+        the next decode step.
         """
-        requests = tuple(requests)
-        sampling_configs = tuple(sampling_configs)
-        self._validate_continuous_inputs(requests, sampling_configs, max_batch_size)
+        original_request = request
+        requests = tuple(
+            RolloutRequest(request.prompt_ids[index:index + 1], request.prompt_attention_mask[index:index + 1])
+            for index in range(request.prompt_ids.shape[0]))
+        self._validate_continuous_inputs(requests, sampling, max_batch_size)
 
         module = self.engine.module
         prompt_len = requests[0].prompt_ids.shape[1]
         max_positions = getattr(module.config, "max_position_embeddings", None)
         if max_positions is not None:
-            for config in sampling_configs:
-                logical_length = prompt_len + config.max_new_tokens
-                if logical_length > max_positions:
-                    raise ValueError("continuous batching request exceeds the model maximum position embeddings")
+            logical_length = prompt_len + sampling.max_new_tokens
+            if logical_length > max_positions:
+                raise ValueError("continuous batching request exceeds the model maximum position embeddings")
         max_cache_len = self._estimate_continuous_cache_len(
             prompt_len,
-            [config.max_new_tokens for config in sampling_configs],
+            [sampling.max_new_tokens] * len(requests),
             max_batch_size,
         )
         if max_positions is not None and max_cache_len > max_positions:
             raise ValueError("continuous batching cache exceeds the model maximum position embeddings")
         if not getattr(module, "_supports_cache_class", False):
-            return self._generate_continuous_legacy(requests, sampling_configs, max_batch_size)
+            return self._generate_continuous_legacy(requests, sampling, max_batch_size)
 
         from transformers import StaticCache
         from deepspeed.utils.static_cache import DeepSpeedStaticCache
@@ -278,11 +282,11 @@ def generate_continuous(self, requests, sampling_configs, max_batch_size):
         device = requests[0].prompt_ids.device
         model_dtype = next(module.parameters()).dtype
 
-        scheduler = ContinuousBatchScheduler(max_batch_size)
+        scheduler = ContinuousBatchScheduler(max_batch_size, sampling.max_new_tokens)
         request_by_id = {}
         responses = {}
-        for request_id, (request, config) in enumerate(zip(requests, sampling_configs)):
-            scheduler.submit(ContinuousBatchRequest(request_id, config.max_new_tokens))
+        for request_id, request in enumerate(requests):
+            scheduler.submit(ContinuousBatchRequest(request_id))
             request_by_id[request_id] = request
             responses[request_id] = []
 
@@ -358,18 +362,14 @@ def generate_continuous(self, requests, sampling_configs, max_batch_size):
             if survivor_count:
                 cache_position += 1
 
-        return [
-            self._build_continuous_output(request, responses[request_id])
-            for request_id, request in enumerate(requests)
-        ]
+        return self._build_continuous_batch(original_request, responses)
 
-    def _generate_continuous_legacy(self, requests, sampling_configs, max_batch_size):
+    def _generate_continuous_legacy(self, requests, sampling, max_batch_size):
         """Continuous decode for Transformers models that return legacy KV tuples."""
         module = self.engine.module
         device = requests[0].prompt_ids.device
         prompt_len = requests[0].prompt_ids.shape[1]
-        scheduler, request_by_id, responses = self._create_continuous_scheduler(requests, sampling_configs,
-                                                                                max_batch_size)
+        scheduler, request_by_id, responses = self._create_continuous_scheduler(requests, sampling, max_batch_size)
         next_tokens = {}
         attention_mask = None
         past_key_values = None
@@ -416,18 +416,19 @@ def _generate_continuous_legacy(self, requests, sampling_configs, max_batch_size
                     finished_ids.append(request_id)
             update = scheduler.advance(finished_ids)
 
-        return [
-            self._build_continuous_output(request, responses[request_id])
-            for request_id, request in enumerate(requests)
-        ]
+        request = RolloutRequest(
+            torch.cat([request.prompt_ids for request in requests], dim=0),
+            torch.cat([request.prompt_attention_mask for request in requests], dim=0),
+        )
+        return self._build_continuous_batch(request, responses)
 
     @staticmethod
-    def _create_continuous_scheduler(requests, sampling_configs, max_batch_size):
-        scheduler = ContinuousBatchScheduler(max_batch_size)
+    def _create_continuous_scheduler(requests, sampling, max_batch_size):
+        scheduler = ContinuousBatchScheduler(max_batch_size, sampling.max_new_tokens)
         request_by_id = {}
         responses = {}
-        for request_id, (request, config) in enumerate(zip(requests, sampling_configs)):
-            scheduler.submit(ContinuousBatchRequest(request_id, config.max_new_tokens))
+        for request_id, request in enumerate(requests):
+            scheduler.submit(ContinuousBatchRequest(request_id))
             request_by_id[request_id] = request
             responses[request_id] = []
         return scheduler, request_by_id, responses
@@ -520,11 +521,9 @@ def _merge_legacy_cache(survivor_cache, survivor_attention, admitted_cache, admi
         padded_attention = torch.cat((admitted_padding, admitted_attention), dim=1)
         return merged_cache, torch.cat((survivor_attention, padded_attention), dim=0)
 
-    def _validate_continuous_inputs(self, requests, sampling_configs, max_batch_size):
+    def _validate_continuous_inputs(self, requests, sampling, max_batch_size):
         if not requests:
             raise ValueError("continuous batching requires at least one request")
-        if len(requests) != len(sampling_configs):
-            raise ValueError("requests and sampling_configs must have the same length")
         if max_batch_size <= 0:
             raise ValueError("max_batch_size must be positive")
         if self.use_graph_capture:
@@ -532,17 +531,19 @@ def _validate_continuous_inputs(self, requests, sampling_configs, max_batch_size
 
         prompt_len = requests[0].prompt_ids.shape[1]
         device = requests[0].prompt_ids.device
-        for request, config in zip(requests, sampling_configs):
+        if sampling.max_new_tokens <= 0:
+            raise ValueError("max_new_tokens must be positive")
+        if sampling.temperature > 0:
+            raise ValueError("continuous batching currently supports greedy decoding only")
+        if sampling.n_samples_per_prompt != 1:
+            raise ValueError("continuous batching currently supports one sample per prompt")
+        for request in requests:
             if request.prompt_ids.shape[0] != 1:
                 raise ValueError("continuous batching requires one prompt row per request")
             if request.prompt_ids.shape[1] != prompt_len:
                 raise ValueError("continuous batching currently requires equal prompt widths")
             if request.prompt_ids.device != device:
                 raise ValueError("continuous batching requests must use the same device")
-            if config.temperature > 0:
-                raise ValueError("continuous batching currently supports greedy decoding only")
-            if config.n_samples_per_prompt != 1:
-                raise ValueError("continuous batching currently supports one sample per prompt")
 
     def _continuous_prefill(self, module, static_cache_type, cache, update, request_by_id, attention_mask,
                             cache_position, prompt_len, model_dtype, device):
@@ -612,17 +613,42 @@ def _create_static_cache(static_cache_type, config, batch_size, max_cache_len, d
             common_kwargs["max_batch_size"] = batch_size
         return static_cache_type(**common_kwargs)
 
-    @staticmethod
-    def _build_continuous_output(request, response_tokens):
-        response_ids = torch.cat(response_tokens, dim=1)
-        input_ids = torch.cat((request.prompt_ids, response_ids), dim=1)
-        response_attention = torch.ones_like(response_ids)
-        attention_mask = torch.cat((request.prompt_attention_mask, response_attention), dim=1)
+    def _build_continuous_batch(self, request, responses):
+        response_ids = [torch.cat(responses[index], dim=1) for index in range(request.prompt_ids.shape[0])]
+        max_response_len = max(response.shape[1] for response in response_ids)
+        pad_token_id = self.tokenizer.pad_token_id
+        if pad_token_id is None:
+            pad_token_id = self.tokenizer.eos_token_id
+        if pad_token_id is None:
+            raise ValueError("The tokenizer must define pad_token_id or eos_token_id")
+        input_rows = []
+        attention_rows = []
+        for index, response in enumerate(response_ids):
+            padding = max_response_len - response.shape[1]
+            if padding:
+                response_padding = torch.full((1, padding), pad_token_id, dtype=response.dtype, device=response.device)
+                attention_padding = torch.zeros((1, padding), dtype=request.prompt_attention_mask.dtype,
+                                                device=response.device)
+                response = torch.cat((response, response_padding), dim=1)
+            else:
+                attention_padding = torch.empty((1, 0),
+                                                dtype=request.prompt_attention_mask.dtype,
+                                                device=response.device)
+            input_rows.append(torch.cat((request.prompt_ids[index:index + 1], response), dim=1))
+            response_attention = torch.ones((1, response_ids[index].shape[1]),
+                                             dtype=request.prompt_attention_mask.dtype,
+                                             device=response.device)
+            attention_rows.append(torch.cat(
+                (request.prompt_attention_mask[index:index + 1], response_attention, attention_padding), dim=1))
+
+        input_ids = torch.cat(input_rows, dim=0)
+        attention_mask = torch.cat(attention_rows, dim=0)
         response_start = request.prompt_ids.shape[1]
         return RolloutBatch(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            response_start_idx=torch.tensor([response_start], dtype=torch.long, device=input_ids.device),
+            response_start_idx=torch.full((request.prompt_ids.shape[0], ), response_start, dtype=torch.long,
+                                          device=input_ids.device),
         )
 
     def get_last_profile(self):
diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst
index 62eca0928bbb..4efacc6db822 100644
--- a/docs/code-docs/source/inference-engine.rst
+++ b/docs/code-docs/source/inference-engine.rst
@@ -77,22 +77,18 @@ response branch after the shared prompt forward.
 Continuous-batching prototype
 -----------------------------
 
-``deepspeed.runtime.rollout.continuous_batching`` provides the scheduling and
-slot-lifecycle primitive needed to build a continuous decode batch. Requests
-are admitted in FIFO order up to a configured capacity. When a request retires,
-the update identifies the surviving cache rows to compact and the pending
-requests that can be prefetched into the newly free rows. The model backend is
-responsible for applying that update, running prompt prefill, and constructing
-the attention metadata for the active rows.
+Continuous batching is enabled through ``SamplingConfig.continuous_batch_size``
+on the regular ``HybridEngineRollout.generate(request, sampling)`` entry point.
+When unset, generation keeps its existing behavior. When set to a positive
+value, at most that many prompt rows are active at once; completed rows retire
+and pending rows are prefetched into the released slots. The returned
+``RolloutBatch`` remains in the original ``RolloutRequest`` row order.
 
 The prototype intentionally does not implement paged attention or change the
-default ``HybridEngineRollout.generate`` path. For a first end-to-end trial,
-``HybridEngineRollout.generate_continuous`` accepts one request per prompt row
-and a matching list of greedy ``SamplingConfig`` objects. It dynamically
-prefills admitted prompts and decodes surviving rows until every request has
-finished. CUDA Graph capture, sampling, multiple samples per prompt, and
-different prompt widths are intentionally rejected until the scheduling
-semantics are validated on real workloads.
+default generation semantics. It currently requires one prompt width for all
+rows, greedy decoding, and one sample per prompt. CUDA Graph capture and
+multiple prompt widths are rejected until the scheduling semantics are
+validated on real workloads.
 
 ``DeepSpeedStaticCache`` accepts one write position per row and can compact
 active rows while preserving its static tensor addresses. This mirrors the
diff --git a/tests/unit/runtime/rollout/test_continuous_batching.py b/tests/unit/runtime/rollout/test_continuous_batching.py
index 04e7f4c775a9..4f713c426ab0 100644
--- a/tests/unit/runtime/rollout/test_continuous_batching.py
+++ b/tests/unit/runtime/rollout/test_continuous_batching.py
@@ -8,12 +8,12 @@
 from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler)
 
 
-def _request(request_id, max_new_tokens=3):
-    return ContinuousBatchRequest(request_id, max_new_tokens)
+def _request(request_id):
+    return ContinuousBatchRequest(request_id)
 
 
 def test_scheduler_admits_fifo_and_respects_capacity():
-    scheduler = ContinuousBatchScheduler(max_batch_size=2)
+    scheduler = ContinuousBatchScheduler(max_batch_size=2, max_new_tokens=3)
     scheduler.submit(_request("a"))
     scheduler.submit(_request("b"))
     scheduler.submit(_request("c"))
@@ -27,7 +27,7 @@ def test_scheduler_admits_fifo_and_respects_capacity():
 
 
 def test_scheduler_compacts_survivors_and_admits_pending_request():
-    scheduler = ContinuousBatchScheduler(max_batch_size=2)
+    scheduler = ContinuousBatchScheduler(max_batch_size=2, max_new_tokens=3)
     scheduler.submit(_request("a"))
     scheduler.submit(_request("b"))
     scheduler.submit(_request("c"))
@@ -42,25 +42,24 @@ def test_scheduler_compacts_survivors_and_admits_pending_request():
 
 
 def test_scheduler_advance_retires_by_budget():
-    scheduler = ContinuousBatchScheduler(max_batch_size=2)
-    scheduler.submit(_request("a", max_new_tokens=1))
-    scheduler.submit(_request("b", max_new_tokens=3))
+    scheduler = ContinuousBatchScheduler(max_batch_size=2, max_new_tokens=1)
+    scheduler.submit(_request("a"))
+    scheduler.submit(_request("b"))
     scheduler.schedule()
 
     update = scheduler.advance()
-    assert update.retired == ("a", )
-    assert update.keep_slots == (1, )
-    assert update.active_ids == ("b", )
-    assert scheduler.active[0].request_id == "b"
+    assert update.retired == ("a", "b")
+    assert update.keep_slots == ()
+    assert update.active_ids == ()
 
 
 def test_scheduler_rejects_invalid_transitions():
     with pytest.raises(ValueError, match="max_batch_size"):
-        ContinuousBatchScheduler(max_batch_size=0)
+        ContinuousBatchScheduler(max_batch_size=0, max_new_tokens=3)
     with pytest.raises(ValueError, match="max_new_tokens"):
-        _request("bad", max_new_tokens=0)
+        ContinuousBatchScheduler(max_batch_size=1, max_new_tokens=0)
 
-    scheduler = ContinuousBatchScheduler(max_batch_size=1)
+    scheduler = ContinuousBatchScheduler(max_batch_size=1, max_new_tokens=3)
     scheduler.submit(_request("a"))
     with pytest.raises(ValueError, match="duplicate"):
         scheduler.submit(_request("a"))
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index aebc4c29d457..3fb83dcd3ad6 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -88,16 +88,13 @@ def test_continuous_generation_rejects_unsupported_inputs():
         prompt_attention_mask=torch.tensor([[0, 1, 1]]),
     )
 
+    empty_request = RolloutRequest(torch.empty((0, 3), dtype=torch.long), torch.empty((0, 3), dtype=torch.long))
     with pytest.raises(ValueError, match="at least one request"):
-        rollout.generate_continuous([], [], max_batch_size=1)
-    with pytest.raises(ValueError, match="same length"):
-        rollout.generate_continuous([request], [], max_batch_size=1)
+        rollout.generate(empty_request, SamplingConfig(max_new_tokens=2, continuous_batch_size=1))
+    with pytest.raises(ValueError, match="positive"):
+        rollout.generate(request, SamplingConfig(max_new_tokens=2, continuous_batch_size=0))
     with pytest.raises(ValueError, match="greedy"):
-        rollout.generate_continuous(
-            [request],
-            [SamplingConfig(max_new_tokens=2, temperature=0.5)],
-            max_batch_size=1,
-        )
+        rollout.generate(request, SamplingConfig(max_new_tokens=2, temperature=0.5, continuous_batch_size=1))
 
 
 def test_static_cache_constructor_supports_legacy_batch_keyword():
@@ -149,7 +146,7 @@ def __init__(self):
     request = RolloutRequest(torch.tensor([[1, 2, 3]]), torch.ones((1, 3), dtype=torch.long))
 
     with pytest.raises(ValueError, match="request exceeds"):
-        rollout.generate_continuous([request], [SamplingConfig(max_new_tokens=2, temperature=0)], 1)
+        rollout.generate(request, SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1))
 
 
 def test_continuous_generation_refills_legacy_cache_batch():
@@ -170,24 +167,25 @@ def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=Tru
             self.calls.append((input_ids.shape[0], input_ids.shape[1], past_length))
             keys = torch.zeros((input_ids.shape[0], 1, total_length, 1))
             logits = torch.zeros((input_ids.shape[0], input_ids.shape[1], 16))
-            logits[..., 7] = 1
+            for row, token in enumerate(input_ids[:, -1].tolist()):
+                logits[row, :, 2 if token == 3 else 7] = 1
             return SimpleNamespace(logits=logits, past_key_values=((keys, keys.clone()), ))
 
     model = LegacyCacheModel()
-    rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None))
-    requests = [
-        RolloutRequest(torch.tensor([[1, 2, token]]), torch.ones((1, 3), dtype=torch.long)) for token in (3, 4, 5)
-    ]
-    configs = [
-        SamplingConfig(max_new_tokens=1, temperature=0),
-        SamplingConfig(max_new_tokens=3, temperature=0),
-        SamplingConfig(max_new_tokens=2, temperature=0),
-    ]
-
-    outputs = rollout.generate_continuous(requests, configs, max_batch_size=2)
+    rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=2))
+    request = RolloutRequest(
+        torch.tensor([[1, 2, 3], [1, 2, 4], [1, 2, 5]]),
+        torch.ones((3, 3), dtype=torch.long),
+    )
+    output = rollout.generate(request, SamplingConfig(max_new_tokens=3, temperature=0, continuous_batch_size=2))
 
-    assert [output.input_ids.shape[1] - 3 for output in outputs] == [1, 3, 2]
-    assert model.calls[:3] == [(2, 3, 0), (1, 1, 3), (1, 3, 0)]
+    assert output.input_ids.shape == (3, 6)
+    assert output.input_ids[:, :3].tolist() == request.prompt_ids.tolist()
+    assert output.input_ids[:, 3:].tolist() == [[2, 0, 0], [7, 7, 7], [7, 7, 7]]
+    assert output.attention_mask[:, 3:].tolist() == [[1, 0, 0], [1, 1, 1], [1, 1, 1]]
+    assert output.response_start_idx.tolist() == [3, 3, 3]
+    assert model.calls[0] == (2, 3, 0)
+    assert (1, 3, 0) in model.calls
 
 
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.time.perf_counter")
@@ -528,6 +526,7 @@ def test_generate_calls_graph_capture_when_enabled(mock_get_accelerator, mock_pe
     sampling.temperature = 0
     sampling.n_samples_per_prompt = 1
     sampling.max_new_tokens = 3
+    sampling.continuous_batch_size = None
 
     rollout.generate(req, sampling)
     rollout._generate_graph.assert_called_once()
@@ -552,6 +551,7 @@ def test_generate_keeps_ranks_in_lockstep_and_pads_after_eos():
     sampling.n_samples_per_prompt = 1
     sampling.max_new_tokens = 4
     sampling.top_p = 1.0
+    sampling.continuous_batch_size = None
 
     result = rollout.generate(req, sampling)
 
@@ -608,7 +608,11 @@ def test_generate_accepts_zero_pad_token_id():
     req = MagicMock()
     req.prompt_ids = torch.tensor([[10, 11]])
     req.prompt_attention_mask = torch.ones(1, 2, dtype=torch.long)
-    sampling = MagicMock(temperature=0, n_samples_per_prompt=1, max_new_tokens=2, top_p=1.0)
+    sampling = MagicMock(temperature=0,
+                        n_samples_per_prompt=1,
+                        max_new_tokens=2,
+                        top_p=1.0,
+                        continuous_batch_size=None)
 
     rollout.generate(req, sampling)
 
diff --git a/tests/unit/runtime/rollout/test_rollout_interface.py b/tests/unit/runtime/rollout/test_rollout_interface.py
index bb45267ef5ac..7a94c01cc40b 100644
--- a/tests/unit/runtime/rollout/test_rollout_interface.py
+++ b/tests/unit/runtime/rollout/test_rollout_interface.py
@@ -11,6 +11,7 @@
 import pytest
 import torch
 
+import deepspeed.runtime.rollout as rollout_module
 from deepspeed.runtime.rollout import (
     RolloutBatch,
     RolloutEngine,
@@ -60,6 +61,13 @@ def test_sampling_config_defaults():
     assert cfg.top_p == 1.0
     assert cfg.top_k == -1
     assert cfg.n_samples_per_prompt == 1
+    assert cfg.continuous_batch_size is None
+
+
+def test_continuous_batching_details_are_not_public_exports():
+    assert "ContinuousBatchRequest" not in rollout_module.__all__
+    assert "ContinuousBatchScheduler" not in rollout_module.__all__
+    assert "ContinuousBatchUpdate" not in rollout_module.__all__
 
 
 # --- interface conformance via FakeRollout ---------------------------------

From 2987844f28a4765be7d1eee162776c139fe4d525 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Tue, 8 Sep 2026 15:49:47 +0800
Subject: [PATCH 09/12] chore(rollout): align license headers

Remove the corporate copyright line from the new rollout module and tests to match the existing SPDX and DeepSpeed Team header format.

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 deepspeed/runtime/rollout/continuous_batching.py         | 1 -
 tests/unit/runtime/rollout/test_continuous_batching.py   | 1 -
 tests/unit/runtime/rollout/test_hybrid_engine_rollout.py | 1 -
 3 files changed, 3 deletions(-)

diff --git a/deepspeed/runtime/rollout/continuous_batching.py b/deepspeed/runtime/rollout/continuous_batching.py
index 36c30ff64848..b4c780287f5e 100644
--- a/deepspeed/runtime/rollout/continuous_batching.py
+++ b/deepspeed/runtime/rollout/continuous_batching.py
@@ -1,4 +1,3 @@
-# Copyright (c) Microsoft Corporation.
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
diff --git a/tests/unit/runtime/rollout/test_continuous_batching.py b/tests/unit/runtime/rollout/test_continuous_batching.py
index 4f713c426ab0..034625a56877 100644
--- a/tests/unit/runtime/rollout/test_continuous_batching.py
+++ b/tests/unit/runtime/rollout/test_continuous_batching.py
@@ -1,4 +1,3 @@
-# Copyright (c) Microsoft Corporation.
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index 3fb83dcd3ad6..f76db27e934d 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -1,4 +1,3 @@
-# Copyright (c) Microsoft Corporation.
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team

From 765a22b5fced2cfc37af5ed6418aeef44800a832 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Tue, 8 Sep 2026 21:20:39 +0800
Subject: [PATCH 10/12] fix(rollout): address continuous batching review
 feedback

Add CPU coverage for the modern StaticCache continuous-batching path and avoid unnecessary cache and attention-mask compaction during steady-state decode. Remove the unsupported legacy continuous-batching fallback, clean up unused cache APIs, and update the experimental documentation.

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py  | 153 ++----------------
 deepspeed/utils/static_cache.py               |  36 +----
 docs/code-docs/source/inference-engine.rst    |  14 +-
 .../rollout/test_hybrid_engine_rollout.py     |  78 +++++----
 tests/unit/utils/test_static_cache.py         |  20 ++-
 5 files changed, 95 insertions(+), 206 deletions(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 1a76e1b54042..3483f808e555 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -10,8 +10,8 @@
      Pre-allocates a StaticCache, captures the decode forward pass with a
      CUDA graph, and replays it for each decode step.  Eliminates kernel
      launch overhead.
-  3. **continuous batching**: an opt-in bounded greedy path selected through
-     ``SamplingConfig.continuous_batch_size``.
+  3. **continuous batching (experimental)**: an opt-in bounded greedy path
+     selected through ``SamplingConfig.continuous_batch_size``.
 """
 
 import time
@@ -246,12 +246,11 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout
 
     @torch.no_grad()
     def _generate_continuous(self, request, sampling, max_batch_size):
-        """Generate independent greedy requests in a continuously refilled batch.
+        """Generate independent greedy requests in an experimental continuous batch.
 
-        This first integration targets the OPSD prototype: every request has a
-        single prompt row and one greedy response. Completed rows retire
-        immediately and pending prompts prefill into the released rows before
-        the next decode step.
+        Continuous batching currently requires one prompt row per request and
+        greedy decoding. Completed rows retire immediately and pending prompts
+        prefill into the released rows before the next decode step.
         """
         original_request = request
         requests = tuple(
@@ -274,7 +273,8 @@ def _generate_continuous(self, request, sampling, max_batch_size):
         if max_positions is not None and max_cache_len > max_positions:
             raise ValueError("continuous batching cache exceeds the model maximum position embeddings")
         if not getattr(module, "_supports_cache_class", False):
-            return self._generate_continuous_legacy(requests, sampling, max_batch_size)
+            raise ValueError("continuous batching requires a model with cache-class support; use the default "
+                             "generate() path or upgrade transformers")
 
         from transformers import StaticCache
         from deepspeed.utils.static_cache import DeepSpeedStaticCache
@@ -308,10 +308,13 @@ def _generate_continuous(self, request, sampling, max_batch_size):
             keep_slots = torch.tensor(update.keep_slots, dtype=torch.long, device=device)
             survivor_count = keep_slots.numel()
             if survivor_count:
-                cache.compact(keep_slots)
-                survivor_attention = attention_mask.index_select(0, keep_slots).clone()
-                attention_mask.zero_()
-                attention_mask[:survivor_count].copy_(survivor_attention)
+                if update.retired:
+                    cache.compact(keep_slots)
+                    identity = torch.arange(survivor_count, dtype=torch.long, device=device)
+                    if not torch.equal(keep_slots, identity):
+                        survivor_attention = attention_mask.index_select(0, keep_slots).clone()
+                        attention_mask[:survivor_count].copy_(survivor_attention)
+                    attention_mask[survivor_count:].zero_()
             else:
                 cache.reset()
                 write_positions.fill_(-1)
@@ -325,6 +328,7 @@ def _generate_continuous(self, request, sampling, max_batch_size):
                 update,
                 request_by_id,
                 attention_mask,
+                write_positions,
                 cache_position,
                 prompt_len,
                 model_dtype,
@@ -364,106 +368,6 @@ def _generate_continuous(self, request, sampling, max_batch_size):
 
         return self._build_continuous_batch(original_request, responses)
 
-    def _generate_continuous_legacy(self, requests, sampling, max_batch_size):
-        """Continuous decode for Transformers models that return legacy KV tuples."""
-        module = self.engine.module
-        device = requests[0].prompt_ids.device
-        prompt_len = requests[0].prompt_ids.shape[1]
-        scheduler, request_by_id, responses = self._create_continuous_scheduler(requests, sampling, max_batch_size)
-        next_tokens = {}
-        attention_mask = None
-        past_key_values = None
-        update = scheduler.schedule()
-
-        while update.active:
-            survivor_count = len(update.keep_slots)
-            survivor_ids = update.active_ids[:survivor_count]
-            survivor_cache = None
-            survivor_attention = None
-            decoded_tokens = {}
-            if survivor_count:
-                keep_slots = torch.tensor(update.keep_slots, dtype=torch.long, device=device)
-                survivor_cache = self._select_legacy_cache_rows(past_key_values, keep_slots, attention_mask.shape[0])
-                survivor_attention = attention_mask.index_select(0, keep_slots)
-                decode_input = torch.cat([next_tokens[request_id] for request_id in survivor_ids], dim=0)
-                survivor_attention = torch.cat(
-                    (survivor_attention, torch.ones((survivor_count, 1), dtype=torch.long, device=device)), dim=1)
-                decode_output = self._call_model(
-                    module,
-                    decode_input,
-                    attention_mask=survivor_attention,
-                    past_key_values=survivor_cache,
-                    use_cache=True,
-                )
-                survivor_cache = decode_output.past_key_values
-                decoded = decode_output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
-                decoded_tokens = dict(zip(survivor_ids, decoded.split(1, dim=0)))
-
-            admitted_cache, admitted_attention, admitted_tokens = self._legacy_prefill(module, update, request_by_id)
-            past_key_values, attention_mask = self._merge_legacy_cache(
-                survivor_cache,
-                survivor_attention,
-                admitted_cache,
-                admitted_attention,
-                prompt_len,
-            )
-            next_tokens = decoded_tokens | admitted_tokens
-            finished_ids = []
-            for request_id in update.active_ids:
-                token = next_tokens[request_id]
-                responses[request_id].append(token)
-                if self._is_eos(token):
-                    finished_ids.append(request_id)
-            update = scheduler.advance(finished_ids)
-
-        request = RolloutRequest(
-            torch.cat([request.prompt_ids for request in requests], dim=0),
-            torch.cat([request.prompt_attention_mask for request in requests], dim=0),
-        )
-        return self._build_continuous_batch(request, responses)
-
-    @staticmethod
-    def _create_continuous_scheduler(requests, sampling, max_batch_size):
-        scheduler = ContinuousBatchScheduler(max_batch_size, sampling.max_new_tokens)
-        request_by_id = {}
-        responses = {}
-        for request_id, request in enumerate(requests):
-            scheduler.submit(ContinuousBatchRequest(request_id))
-            request_by_id[request_id] = request
-            responses[request_id] = []
-        return scheduler, request_by_id, responses
-
-    def _legacy_prefill(self, module, update, request_by_id):
-        if not update.admitted:
-            return None, None, {}
-        admitted_ids = tuple(request.request_id for request in update.admitted)
-        prompt_ids = torch.cat([request_by_id[request_id].prompt_ids for request_id in admitted_ids], dim=0)
-        prompt_attention = torch.cat([request_by_id[request_id].prompt_attention_mask for request_id in admitted_ids],
-                                     dim=0)
-        output = self._call_model(module, prompt_ids, attention_mask=prompt_attention, use_cache=True)
-        tokens = output.logits[:, -1, :].argmax(dim=-1, keepdim=True)
-        return output.past_key_values, prompt_attention, dict(zip(admitted_ids, tokens.split(1, dim=0)))
-
-    @staticmethod
-    def _select_legacy_cache_rows(past_key_values, rows, logical_batch_size):
-        if logical_batch_size <= 0:
-            raise ValueError("logical batch size must be positive")
-        selected_cache = []
-        for keys, values in past_key_values:
-            if keys.shape[0] != values.shape[0]:
-                raise ValueError("legacy KV cache key/value batch dimensions must match")
-            cache_batch_size = keys.shape[0]
-            if cache_batch_size == logical_batch_size:
-                cache_rows = rows
-            elif cache_batch_size > logical_batch_size and cache_batch_size % logical_batch_size == 0:
-                heads_per_request = cache_batch_size // logical_batch_size
-                head_offsets = torch.arange(heads_per_request, device=rows.device)
-                cache_rows = (rows[:, None] * heads_per_request + head_offsets).reshape(-1)
-            else:
-                raise ValueError("cannot identify the logical batch dimension in the legacy KV cache")
-            selected_cache.append((keys.index_select(0, cache_rows), values.index_select(0, cache_rows)))
-        return tuple(selected_cache)
-
     @staticmethod
     def _estimate_continuous_cache_len(prompt_len, max_new_tokens, max_batch_size):
         """Estimate the largest cache span before the scheduler becomes empty."""
@@ -499,28 +403,6 @@ def _estimate_continuous_cache_len(prompt_len, max_new_tokens, max_batch_size):
 
         return max_cache_len
 
-    @staticmethod
-    def _merge_legacy_cache(survivor_cache, survivor_attention, admitted_cache, admitted_attention, prompt_len):
-        if survivor_cache is None:
-            return admitted_cache, admitted_attention
-        if admitted_cache is None:
-            return survivor_cache, survivor_attention
-
-        cache_len = survivor_cache[0][0].shape[2]
-        left_padding = cache_len - prompt_len
-        padded_admitted = tuple((torch.nn.functional.pad(keys, (0, 0, left_padding, 0)),
-                                 torch.nn.functional.pad(values, (0, 0, left_padding, 0)))
-                                for keys, values in admitted_cache)
-        merged_cache = tuple(
-            (torch.cat((survivor_keys, admitted_keys), dim=0), torch.cat((survivor_values, admitted_values), dim=0))
-            for (survivor_keys, survivor_values), (admitted_keys,
-                                                   admitted_values) in zip(survivor_cache, padded_admitted))
-        admitted_padding = torch.zeros((admitted_attention.shape[0], left_padding),
-                                       dtype=admitted_attention.dtype,
-                                       device=admitted_attention.device)
-        padded_attention = torch.cat((admitted_padding, admitted_attention), dim=1)
-        return merged_cache, torch.cat((survivor_attention, padded_attention), dim=0)
-
     def _validate_continuous_inputs(self, requests, sampling, max_batch_size):
         if not requests:
             raise ValueError("continuous batching requires at least one request")
@@ -546,7 +428,7 @@ def _validate_continuous_inputs(self, requests, sampling, max_batch_size):
                 raise ValueError("continuous batching requests must use the same device")
 
     def _continuous_prefill(self, module, static_cache_type, cache, update, request_by_id, attention_mask,
-                            cache_position, prompt_len, model_dtype, device):
+                            write_positions, cache_position, prompt_len, model_dtype, device):
         if not update.admitted:
             return {}
 
@@ -573,7 +455,6 @@ def _continuous_prefill(self, module, static_cache_type, cache, update, request_
                 target_layer.values[target_row, :, cache_start:cache_position].copy_(prefill_layer.values[source_row])
         for source_row, target_row in enumerate(update.admitted_slots):
             attention_mask[target_row, cache_start:cache_position].copy_(prompt_attention[source_row])
-            write_positions = cache._write_position
             write_positions[target_row] = cache_position
         return dict(zip(admitted_ids, prefill_tokens.split(1, dim=0)))
 
diff --git a/deepspeed/utils/static_cache.py b/deepspeed/utils/static_cache.py
index 083ca8345e02..6825087653e0 100644
--- a/deepspeed/utils/static_cache.py
+++ b/deepspeed/utils/static_cache.py
@@ -136,47 +136,17 @@ def reset(self) -> None:
             self.keys.zero_()
             self.values.zero_()
 
-    def compact(self, active_indices: torch.Tensor) -> None:
-        """Move active cache rows to the front and update their positions.
-
-        Continuous batching retires requests from arbitrary rows. Copying the
-        survivors in one operation avoids in-place overlap when a later row is
-        moved into an earlier slot, while retaining the tensors' static
-        addresses for CUDA graph users.
-        """
-        if not self.is_initialized:
-            raise RuntimeError("cannot compact an uninitialized cache")
-        if not isinstance(active_indices, torch.Tensor) or active_indices.dim() != 1:
-            raise ValueError("active_indices must be a 1-D tensor")
-        active_indices = active_indices.to(device=self.keys.device, dtype=torch.long)
-        if active_indices.numel() > self.max_batch_size:
-            raise ValueError("active_indices exceeds the cache batch size")
-        if active_indices.numel() and ((active_indices < 0).any() or (active_indices >= self.max_batch_size).any()):
-            raise ValueError("active_indices contains an out-of-range row")
-        if active_indices.unique().numel() != active_indices.numel():
-            raise ValueError("active_indices must not contain duplicates")
-
-        count = active_indices.numel()
-        self._compact_rows(active_indices, count)
-        if self._write_position is not None and self._write_position.dim() == 1:
-            positions = self._write_position
-            if positions.numel() != self.max_batch_size:
-                raise ValueError("per-row write positions must match the cache batch size")
-            position_indices = active_indices.to(positions.device)
-            compacted = positions.index_select(0, position_indices).clone() if count else positions[:0]
-            self._write_position[:count].copy_(compacted)
-            if count < self.max_batch_size:
-                self._write_position[count:].fill_(-1)
-
     def _compact_rows(self, active_indices: torch.Tensor, count: int | None = None) -> None:
         """Compact only the row tensors; used by the multi-layer cache."""
         if count is None:
             count = active_indices.numel()
-        if count:
+        identity = torch.arange(count, device=active_indices.device, dtype=active_indices.dtype)
+        if count and not torch.equal(active_indices, identity):
             keys = self.keys.index_select(0, active_indices).clone()
             values = self.values.index_select(0, active_indices).clone()
             self.keys[:count].copy_(keys)
             self.values[:count].copy_(values)
+        # Clear retired rows so their stale KV cannot be observed after refill.
         if count < self.max_batch_size:
             self.keys[count:].zero_()
             self.values[count:].zero_()
diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst
index 4efacc6db822..86b610d85dac 100644
--- a/docs/code-docs/source/inference-engine.rst
+++ b/docs/code-docs/source/inference-engine.rst
@@ -74,8 +74,8 @@ one token. It cannot be combined with CUDA graph capture or
 ``release_inference_cache``. Sampling still happens independently for every
 response branch after the shared prompt forward.
 
-Continuous-batching prototype
------------------------------
+Continuous batching (experimental)
+-----------------------------------
 
 Continuous batching is enabled through ``SamplingConfig.continuous_batch_size``
 on the regular ``HybridEngineRollout.generate(request, sampling)`` entry point.
@@ -84,11 +84,13 @@ value, at most that many prompt rows are active at once; completed rows retire
 and pending rows are prefetched into the released slots. The returned
 ``RolloutBatch`` remains in the original ``RolloutRequest`` row order.
 
-The prototype intentionally does not implement paged attention or change the
+The experimental path intentionally does not implement paged attention or change the
 default generation semantics. It currently requires one prompt width for all
-rows, greedy decoding, and one sample per prompt. CUDA Graph capture and
-multiple prompt widths are rejected until the scheduling semantics are
-validated on real workloads.
+rows, a model with cache-class support, greedy decoding, and one sample per
+prompt. CUDA Graph capture and multiple prompt widths are rejected until the
+scheduling semantics are validated on real workloads. Models without
+cache-class support should use the default ``generate()`` path or upgrade
+Transformers.
 
 ``DeepSpeedStaticCache`` accepts one write position per row and can compact
 active rows while preserving its static tensor addresses. This mirrors the
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index f76db27e934d..36776b355f35 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -96,16 +96,16 @@ def test_continuous_generation_rejects_unsupported_inputs():
         rollout.generate(request, SamplingConfig(max_new_tokens=2, temperature=0.5, continuous_batch_size=1))
 
 
-def test_static_cache_constructor_supports_legacy_batch_keyword():
+def test_static_cache_constructor_supports_max_batch_keyword():
 
-    class LegacyStaticCache:
+    class MaxBatchStaticCache:
 
         def __init__(self, config, max_batch_size, max_cache_len, device, dtype):
             self.config = config
             self.max_batch_size = max_batch_size
 
     config = SimpleNamespace(num_attention_heads=4)
-    cache = HybridEngineRollout._create_static_cache(LegacyStaticCache, config, 2, 8, "cpu", torch.float32)
+    cache = HybridEngineRollout._create_static_cache(MaxBatchStaticCache, config, 2, 8, "cpu", torch.float32)
 
     assert cache.max_batch_size == 2
     assert cache.config.num_key_value_heads == 4
@@ -118,17 +118,6 @@ def test_continuous_cache_span_does_not_sum_independent_requests():
     assert cache_len < 64 + 64 * 100
 
 
-def test_select_legacy_cache_rows_expands_flattened_heads():
-    keys = torch.arange(2 * 4 * 3).reshape(2 * 4, 1, 3)
-    values = keys + 100
-    rows = torch.tensor([1])
-
-    selected = HybridEngineRollout._select_legacy_cache_rows(((keys, values), ), rows, logical_batch_size=2)
-
-    assert torch.equal(selected[0][0], keys[4:8])
-    assert torch.equal(selected[0][1], values[4:8])
-
-
 def test_continuous_generation_validates_each_request_length():
 
     class LimitedModel(torch.nn.Module):
@@ -148,29 +137,58 @@ def __init__(self):
         rollout.generate(request, SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1))
 
 
-def test_continuous_generation_refills_legacy_cache_batch():
+def test_continuous_generation_rejects_legacy_cache_model():
 
-    class LegacyCacheModel(torch.nn.Module):
-        _supports_cache_class = False
+    class LegacyModel(torch.nn.Module):
 
         def __init__(self):
             super().__init__()
             self.weight = torch.nn.Parameter(torch.zeros(1))
-            self.calls = []
             self.config = SimpleNamespace(max_position_embeddings=32)
 
-        def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=True):
-            past_length = 0 if past_key_values is None else past_key_values[0][0].shape[2]
-            total_length = past_length + input_ids.shape[1]
-            assert attention_mask.shape == (input_ids.shape[0], total_length)
-            self.calls.append((input_ids.shape[0], input_ids.shape[1], past_length))
-            keys = torch.zeros((input_ids.shape[0], 1, total_length, 1))
+    model = LegacyModel()
+    rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(eos_token_id=None))
+    request = RolloutRequest(torch.tensor([[1, 2, 3]]), torch.ones((1, 3), dtype=torch.long))
+
+    with pytest.raises(ValueError, match="cache-class support"):
+        rollout.generate(request, SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1))
+
+
+def test_continuous_generation_covers_modern_static_cache_path():
+
+    class CacheClassModel(torch.nn.Module):
+        _supports_cache_class = True
+
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(1))
+            self.calls = []
+
+            class CacheConfig(SimpleNamespace):
+
+                def get_text_config(self, **_kwargs):
+                    return self
+
+            self.config = CacheConfig(
+                max_position_embeddings=32,
+                num_hidden_layers=1,
+                num_attention_heads=1,
+                num_key_value_heads=1,
+                hidden_size=1,
+                head_dim=1,
+            )
+
+        def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=True, **kwargs):
+            key_states = input_ids[:, None, :, None].to(dtype=torch.float32)
+            _, cache_values = past_key_values.update(key_states, key_states, layer_idx=0, **kwargs)
+            cache_sums = cache_values[:, 0].sum(dim=(1, 2))
+            next_tokens = torch.where(cache_sums == 6, 2, 7).long()
+            self.calls.append((input_ids.shape[0], input_ids.shape[1]))
             logits = torch.zeros((input_ids.shape[0], input_ids.shape[1], 16))
-            for row, token in enumerate(input_ids[:, -1].tolist()):
-                logits[row, :, 2 if token == 3 else 7] = 1
-            return SimpleNamespace(logits=logits, past_key_values=((keys, keys.clone()), ))
+            logits.scatter_(2, next_tokens[:, None, None].expand(-1, input_ids.shape[1], 1), 1)
+            return SimpleNamespace(logits=logits, past_key_values=past_key_values)
 
-    model = LegacyCacheModel()
+    model = CacheClassModel()
     rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=2))
     request = RolloutRequest(
         torch.tensor([[1, 2, 3], [1, 2, 4], [1, 2, 5]]),
@@ -183,8 +201,8 @@ def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=Tru
     assert output.input_ids[:, 3:].tolist() == [[2, 0, 0], [7, 7, 7], [7, 7, 7]]
     assert output.attention_mask[:, 3:].tolist() == [[1, 0, 0], [1, 1, 1], [1, 1, 1]]
     assert output.response_start_idx.tolist() == [3, 3, 3]
-    assert model.calls[0] == (2, 3, 0)
-    assert (1, 3, 0) in model.calls
+    assert model.calls[0] == (2, 3)
+    assert (1, 3) in model.calls
 
 
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.time.perf_counter")
diff --git a/tests/unit/utils/test_static_cache.py b/tests/unit/utils/test_static_cache.py
index 5fed81203236..addd9718606a 100644
--- a/tests/unit/utils/test_static_cache.py
+++ b/tests/unit/utils/test_static_cache.py
@@ -1,4 +1,3 @@
-# Copyright (c) Microsoft Corporation.
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
@@ -44,6 +43,25 @@ def test_static_cache_compact_preserves_rows_and_positions():
     assert layer.keys[2].abs().sum().item() == 0
 
 
+def test_static_cache_compact_identity_keeps_active_rows_and_clears_tail():
+    config = type("Config", (), {"num_hidden_layers": 1, "num_attention_heads": 1, "hidden_size": 2})()
+    cache = DeepSpeedStaticCache(config=config,
+                                 batch_size=3,
+                                 max_cache_len=4,
+                                 device=torch.device("cpu"),
+                                 dtype=torch.float32)
+    cache.set_write_position(torch.tensor([1, 2, 3], dtype=torch.long))
+    layer = cache.layers[0]
+    layer.keys[:, 0, 0, :] = torch.tensor([[10.0, 10.0], [20.0, 20.0], [30.0, 30.0]])
+    layer.values.copy_(layer.keys)
+
+    cache.compact(torch.tensor([0, 1], dtype=torch.long))
+
+    assert layer.keys[:2, 0, 0, 0].tolist() == [10.0, 20.0]
+    assert layer.get_seq_length().tolist() == [2, 3, 0]
+    assert layer.keys[2].abs().sum().item() == 0
+
+
 def test_static_layer_rejects_mismatched_per_row_positions():
     layer = DeepSpeedStaticLayer(max_cache_len=4)
     keys = torch.zeros((2, 1, 1, 2))

From 3b0713e4d0bd2bf12811cbdeacac470d337bcb8e Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Wed, 9 Sep 2026 22:25:40 +0800
Subject: [PATCH 11/12] fix(rollout): reclaim cache space after staggered EOS

Trim unused cache prefixes during continuous batching so staggered EOS and request refill cannot exhaust the allocated StaticCache span. Update the cache position, write positions, and attention mask after trimming, and add CPU coverage for staggered EOS refill and cache shifting.

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py  | 21 +++++++
 deepspeed/utils/static_cache.py               | 21 +++++++
 docs/code-docs/source/inference-engine.rst    |  2 +
 .../rollout/test_hybrid_engine_rollout.py     | 60 +++++++++++++++++++
 tests/unit/utils/test_static_cache.py         | 17 ++++++
 5 files changed, 121 insertions(+)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 3483f808e555..0acbfaf1ce8e 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -302,6 +302,7 @@ def _generate_continuous(self, request, sampling, max_batch_size):
         attention_mask = torch.zeros((max_batch_size, max_cache_len), dtype=torch.long, device=device)
         next_tokens = {}
         cache_position = prompt_len
+        trim_threshold = max(1, prompt_len)
         update = scheduler.schedule()
 
         while update.active:
@@ -315,6 +316,17 @@ def _generate_continuous(self, request, sampling, max_batch_size):
                         survivor_attention = attention_mask.index_select(0, keep_slots).clone()
                         attention_mask[:survivor_count].copy_(survivor_attention)
                     attention_mask[survivor_count:].zero_()
+
+                if update.retired or cache_position >= max_cache_len - trim_threshold:
+                    dead_prefix = self._continuous_dead_prefix(attention_mask, survivor_count)
+                    if dead_prefix >= trim_threshold or cache_position >= max_cache_len:
+                        if dead_prefix == 0:
+                            raise ValueError("continuous batching cache exhausted before active requests retired")
+                        cache.trim_left(dead_prefix)
+                        attention_mask[:, :-dead_prefix].copy_(attention_mask[:, dead_prefix:].clone())
+                        attention_mask[:, -dead_prefix:].zero_()
+                        write_positions[:survivor_count].sub_(dead_prefix)
+                        cache_position -= dead_prefix
             else:
                 cache.reset()
                 write_positions.fill_(-1)
@@ -427,6 +439,15 @@ def _validate_continuous_inputs(self, requests, sampling, max_batch_size):
             if request.prompt_ids.device != device:
                 raise ValueError("continuous batching requests must use the same device")
 
+    @staticmethod
+    def _continuous_dead_prefix(attention_mask, active_count):
+        if active_count == 0:
+            return 0
+        occupied = attention_mask[:active_count].any(dim=0)
+        if not occupied.any():
+            return 0
+        return int(occupied.to(dtype=torch.int32).argmax().item())
+
     def _continuous_prefill(self, module, static_cache_type, cache, update, request_by_id, attention_mask,
                             write_positions, cache_position, prompt_len, model_dtype, device):
         if not update.admitted:
diff --git a/deepspeed/utils/static_cache.py b/deepspeed/utils/static_cache.py
index 6825087653e0..3c35be5db539 100644
--- a/deepspeed/utils/static_cache.py
+++ b/deepspeed/utils/static_cache.py
@@ -151,6 +151,18 @@ def _compact_rows(self, active_indices: torch.Tensor, count: int | None = None)
             self.keys[count:].zero_()
             self.values[count:].zero_()
 
+    def _trim_left(self, count: int) -> None:
+        """Shift the cache contents left while preserving tensor addresses."""
+        if count <= 0:
+            return
+        if count >= self.max_cache_len:
+            self.reset()
+            return
+        self.keys[:, :, :-count].copy_(self.keys[:, :, count:].clone())
+        self.values[:, :, :-count].copy_(self.values[:, :, count:].clone())
+        self.keys[:, :, -count:].zero_()
+        self.values[:, :, -count:].zero_()
+
     def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
         if self.is_initialized:
             self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))
@@ -244,6 +256,15 @@ def compact(self, active_indices: torch.Tensor) -> None:
             if count < positions.numel():
                 positions[count:].fill_(-1)
 
+    def trim_left(self, count: int) -> None:
+        """Shift all cache layers left to reclaim an unused prefix span."""
+        if not isinstance(count, int) or count < 0:
+            raise ValueError("trim count must be a non-negative integer")
+        if count == 0:
+            return
+        for layer in self._layers:
+            layer._trim_left(count)
+
     def update(
         self,
         key_states: torch.Tensor,
diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst
index 86b610d85dac..b978dcaf1f89 100644
--- a/docs/code-docs/source/inference-engine.rst
+++ b/docs/code-docs/source/inference-engine.rst
@@ -83,6 +83,8 @@ When unset, generation keeps its existing behavior. When set to a positive
 value, at most that many prompt rows are active at once; completed rows retire
 and pending rows are prefetched into the released slots. The returned
 ``RolloutBatch`` remains in the original ``RolloutRequest`` row order.
+The experimental path periodically trims unused cache columns from the left
+to keep long-running staggered-EOS workloads within the allocated cache span.
 
 The experimental path intentionally does not implement paged attention or change the
 default generation semantics. It currently requires one prompt width for all
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index 36776b355f35..9fb5a52f9fb7 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -205,6 +205,66 @@ def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=Tru
     assert (1, 3) in model.calls
 
 
+def test_continuous_generation_trims_cache_after_staggered_eos():
+
+    class CacheConfig(SimpleNamespace):
+
+        def get_text_config(self, **_kwargs):
+            return self
+
+    class CacheClassModel(torch.nn.Module):
+        _supports_cache_class = True
+
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(1))
+            self.config = CacheConfig(
+                max_position_embeddings=32,
+                num_hidden_layers=1,
+                num_attention_heads=1,
+                num_key_value_heads=1,
+                hidden_size=1,
+                head_dim=1,
+            )
+
+        def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=True, **kwargs):
+            states = input_ids[:, None, :, None].to(dtype=torch.float32)
+            _, values = past_key_values.update(states, states, layer_idx=0, **kwargs)
+            cache_sums = values[:, 0].sum(dim=(1, 2))
+            eos_rows = (cache_sums == 6) | (cache_sums == 8) | (cache_sums == 10)
+            next_tokens = torch.where(eos_rows, 2, 7).long()
+            logits = torch.zeros((input_ids.shape[0], input_ids.shape[1], 16))
+            logits.scatter_(2, next_tokens[:, None, None].expand(-1, input_ids.shape[1], 1), 1)
+            return SimpleNamespace(logits=logits, past_key_values=past_key_values)
+
+    model = CacheClassModel()
+    rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=2))
+    request = RolloutRequest(
+        torch.tensor([[1, 2, 3], [1, 2, 4], [1, 2, 5], [1, 2, 6], [1, 2, 7], [1, 2, 8]]),
+        torch.ones((6, 3), dtype=torch.long),
+    )
+
+    output = rollout.generate(request, SamplingConfig(max_new_tokens=4, temperature=0, continuous_batch_size=2))
+
+    assert output.input_ids.shape == (6, 7)
+    assert output.input_ids[:, 3:].tolist() == [
+        [2, 0, 0, 0],
+        [7, 7, 7, 7],
+        [2, 0, 0, 0],
+        [7, 7, 7, 7],
+        [2, 0, 0, 0],
+        [7, 7, 7, 7],
+    ]
+    assert output.attention_mask[:, 3:].tolist() == [
+        [1, 0, 0, 0],
+        [1, 1, 1, 1],
+        [1, 0, 0, 0],
+        [1, 1, 1, 1],
+        [1, 0, 0, 0],
+        [1, 1, 1, 1],
+    ]
+
+
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.time.perf_counter")
 @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.get_accelerator")
 def test_generate_records_profile_when_enabled(mock_get_accelerator, mock_perf_counter):
diff --git a/tests/unit/utils/test_static_cache.py b/tests/unit/utils/test_static_cache.py
index addd9718606a..4b54474b338c 100644
--- a/tests/unit/utils/test_static_cache.py
+++ b/tests/unit/utils/test_static_cache.py
@@ -62,6 +62,23 @@ def test_static_cache_compact_identity_keeps_active_rows_and_clears_tail():
     assert layer.keys[2].abs().sum().item() == 0
 
 
+def test_static_cache_trim_left_shifts_values_and_clears_tail():
+    config = type("Config", (), {"num_hidden_layers": 1, "num_attention_heads": 1, "hidden_size": 1})()
+    cache = DeepSpeedStaticCache(config=config,
+                                 batch_size=1,
+                                 max_cache_len=4,
+                                 device=torch.device("cpu"),
+                                 dtype=torch.float32)
+    layer = cache.layers[0]
+    layer.keys[0, 0, :, 0] = torch.tensor([0.0, 1.0, 2.0, 3.0])
+    layer.values.copy_(layer.keys)
+
+    cache.trim_left(1)
+
+    assert layer.keys[0, 0, :, 0].tolist() == [1.0, 2.0, 3.0, 0.0]
+    assert layer.values[0, 0, :, 0].tolist() == [1.0, 2.0, 3.0, 0.0]
+
+
 def test_static_layer_rejects_mismatched_per_row_positions():
     layer = DeepSpeedStaticLayer(max_cache_len=4)
     keys = torch.zeros((2, 1, 1, 2))

From 5ffa07218e7678de7b22bc71a6d2db19279220a0 Mon Sep 17 00:00:00 2001
From: nathon-lee <leejianwoo@gmail.com>
Date: Thu, 10 Sep 2026 16:18:24 +0000
Subject: [PATCH 12/12] style: fix formatting via pre-commit

Signed-off-by: nathon-lee <leejianwoo@gmail.com>
---
 .../runtime/rollout/hybrid_engine_rollout.py     | 16 ++++++++++------
 .../rollout/test_hybrid_engine_rollout.py        |  8 ++++----
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
index 0acbfaf1ce8e..120a8264900b 100644
--- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py
+++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py
@@ -529,7 +529,8 @@ def _build_continuous_batch(self, request, responses):
             padding = max_response_len - response.shape[1]
             if padding:
                 response_padding = torch.full((1, padding), pad_token_id, dtype=response.dtype, device=response.device)
-                attention_padding = torch.zeros((1, padding), dtype=request.prompt_attention_mask.dtype,
+                attention_padding = torch.zeros((1, padding),
+                                                dtype=request.prompt_attention_mask.dtype,
                                                 device=response.device)
                 response = torch.cat((response, response_padding), dim=1)
             else:
@@ -538,10 +539,11 @@ def _build_continuous_batch(self, request, responses):
                                                 device=response.device)
             input_rows.append(torch.cat((request.prompt_ids[index:index + 1], response), dim=1))
             response_attention = torch.ones((1, response_ids[index].shape[1]),
-                                             dtype=request.prompt_attention_mask.dtype,
-                                             device=response.device)
-            attention_rows.append(torch.cat(
-                (request.prompt_attention_mask[index:index + 1], response_attention, attention_padding), dim=1))
+                                            dtype=request.prompt_attention_mask.dtype,
+                                            device=response.device)
+            attention_rows.append(
+                torch.cat((request.prompt_attention_mask[index:index + 1], response_attention, attention_padding),
+                          dim=1))
 
         input_ids = torch.cat(input_rows, dim=0)
         attention_mask = torch.cat(attention_rows, dim=0)
@@ -549,7 +551,9 @@ def _build_continuous_batch(self, request, responses):
         return RolloutBatch(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            response_start_idx=torch.full((request.prompt_ids.shape[0], ), response_start, dtype=torch.long,
+            response_start_idx=torch.full((request.prompt_ids.shape[0], ),
+                                          response_start,
+                                          dtype=torch.long,
                                           device=input_ids.device),
         )
 
diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
index 9fb5a52f9fb7..c6b63738adcc 100644
--- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
+++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py
@@ -686,10 +686,10 @@ def test_generate_accepts_zero_pad_token_id():
     req.prompt_ids = torch.tensor([[10, 11]])
     req.prompt_attention_mask = torch.ones(1, 2, dtype=torch.long)
     sampling = MagicMock(temperature=0,
-                        n_samples_per_prompt=1,
-                        max_new_tokens=2,
-                        top_p=1.0,
-                        continuous_batch_size=None)
+                         n_samples_per_prompt=1,
+                         max_new_tokens=2,
+                         top_p=1.0,
+                         continuous_batch_size=None)
 
     rollout.generate(req, sampling)