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 new file mode 100644 index 000000000000..b4c780287f5e --- /dev/null +++ b/deepspeed/runtime/rollout/continuous_batching.py @@ -0,0 +1,133 @@ +# 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 + + +@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. + + 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, 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 = {} + 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 >= self.max_new_tokens: + finished.add(request.request_id) + updated.append((request, generated)) + self._active = updated + return self.schedule(finished) diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index de491a92c23e..120a8264900b 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -10,15 +10,20 @@ 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 (experimental)**: an opt-in bounded greedy path + selected through ``SamplingConfig.continuous_batch_size``. """ import time +from copy import copy from dataclasses import dataclass +from inspect import signature import torch 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 class _ForwardProfiler: @@ -96,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 @@ -236,6 +244,319 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout return rollout_batch + @torch.no_grad() + def _generate_continuous(self, request, sampling, max_batch_size): + """Generate independent greedy requests in an experimental continuous batch. + + 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( + 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: + 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, + [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): + 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 + + device = requests[0].prompt_ids.device + model_dtype = next(module.parameters()).dtype + + 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] = [] + + 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 + trim_threshold = max(1, 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: + 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_() + + 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) + attention_mask.zero_() + cache_position = prompt_len + + admitted_tokens = self._continuous_prefill( + module, + StaticCache, + cache, + update, + request_by_id, + attention_mask, + write_positions, + 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 = self._call_model( + 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_batch(original_request, responses) + + @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 + + def _validate_continuous_inputs(self, requests, sampling, max_batch_size): + if not requests: + raise ValueError("continuous batching requires at least one request") + 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 + 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") + + @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: + 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 = self._create_static_cache(static_cache_type, module.config, len(admitted_ids), prompt_len, + device, model_dtype) + prefill_output = self._call_model( + 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[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 _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.""" + 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": cache_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) + + 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.full((request.prompt_ids.shape[0], ), + 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 @@ -316,13 +637,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( diff --git a/deepspeed/utils/static_cache.py b/deepspeed/utils/static_cache.py index b5e0ce838b58..3c35be5db539 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,20 @@ 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, :] + 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 else: @@ -103,7 +121,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 +136,39 @@ def reset(self) -> None: self.keys.zero_() self.values.zero_() + 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() + 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_() + + 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)) 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: @@ -175,15 +222,49 @@ 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.""" + 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_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 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, @@ -209,10 +290,16 @@ 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() + 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 416f5a556dae..b978dcaf1f89 100644 --- a/docs/code-docs/source/inference-engine.rst +++ b/docs/code-docs/source/inference-engine.rst @@ -73,3 +73,28 @@ 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 (experimental) +----------------------------------- + +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 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 +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 +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..034625a56877 --- /dev/null +++ b/tests/unit/runtime/rollout/test_continuous_batching.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import pytest + +from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler) + + +def _request(request_id): + return ContinuousBatchRequest(request_id) + + +def test_scheduler_admits_fifo_and_respects_capacity(): + scheduler = ContinuousBatchScheduler(max_batch_size=2, max_new_tokens=3) + 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, max_new_tokens=3) + 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, max_new_tokens=1) + scheduler.submit(_request("a")) + scheduler.submit(_request("b")) + scheduler.schedule() + + update = scheduler.advance() + 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, max_new_tokens=3) + with pytest.raises(ValueError, match="max_new_tokens"): + ContinuousBatchScheduler(max_batch_size=1, max_new_tokens=0) + + scheduler = ContinuousBatchScheduler(max_batch_size=1, max_new_tokens=3) + 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/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index aeb550b26d78..c6b63738adcc 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 @@ -81,6 +80,191 @@ 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]]), + ) + + 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(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(request, SamplingConfig(max_new_tokens=2, temperature=0.5, continuous_batch_size=1)) + + +def test_static_cache_constructor_supports_max_batch_keyword(): + + 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(MaxBatchStaticCache, config, 2, 8, "cpu", torch.float32) + + assert cache.max_batch_size == 2 + 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_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(request, SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1)) + + +def test_continuous_generation_rejects_legacy_cache_model(): + + class LegacyModel(torch.nn.Module): + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(1)) + self.config = SimpleNamespace(max_position_embeddings=32) + + 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)) + 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]]), + 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 == (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) + 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): @@ -419,6 +603,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() @@ -443,6 +628,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) @@ -499,7 +685,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 --------------------------------- diff --git a/tests/unit/utils/test_static_cache.py b/tests/unit/utils/test_static_cache.py new file mode 100644 index 000000000000..4b54474b338c --- /dev/null +++ b/tests/unit/utils/test_static_cache.py @@ -0,0 +1,100 @@ +# 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 cache.get_seq_length().item() == 4 + 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_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)) + 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