diff --git a/deepspeed/runtime/rollout/continuous_batching.py b/deepspeed/runtime/rollout/continuous_batching.py index 6c85d8fff46a..69b063c2a6dc 100644 --- a/deepspeed/runtime/rollout/continuous_batching.py +++ b/deepspeed/runtime/rollout/continuous_batching.py @@ -131,3 +131,38 @@ def advance(self, finished_ids=()) -> ContinuousBatchUpdate: updated.append((request, generated)) self._active = updated return self.schedule(finished) + + +def plan_prefill_buckets(lengths, cost, max_tokens=None): + """Minimize additive predicted cost over contiguous length-sorted buckets. + + ``cost(batch_size, padded_width)`` belongs to the model backend. Returned + indices refer to the caller's original request order. Complexity is O(n²) + in request count, independent of prompt token count. + """ + if any(length <= 0 for length in lengths): + raise ValueError("prefill lengths must be positive") + if max_tokens is not None and max_tokens <= 0: + raise ValueError("prefill token limit must be positive") + order = sorted(range(len(lengths)), key=lambda i: -lengths[i]) + best = [0.0] + [float("inf")] * len(order) + previous = [0] * (len(order) + 1) + for end in range(1, len(order) + 1): + for start in range(end): + count = end - start + width = lengths[order[start]] + if max_tokens is not None and count * width > max_tokens: + continue + candidate = best[start] + cost(count, width) + if candidate < best[end]: + best[end] = candidate + previous[end] = start + if best[-1] == float("inf"): + raise ValueError("a prompt exceeds the prefill token limit; increase the configured limit") + buckets = [] + end = len(order) + while end: + start = previous[end] + buckets.append(tuple(order[start:end])) + end = start + return tuple(reversed(buckets)), best[-1] diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index f249cb84cd32..56ce43035f80 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -14,17 +14,20 @@ selected through ``SamplingConfig.continuous_batch_size``. """ +import math import time from copy import copy -from dataclasses import dataclass +from dataclasses import dataclass, replace from inspect import signature from typing import Optional import torch from deepspeed.accelerator import get_accelerator +import deepspeed.comm as dist from deepspeed.runtime.rollout.base import RolloutBatch, RolloutEngine, RolloutRequest, SamplingConfig -from deepspeed.runtime.rollout.continuous_batching import ContinuousBatchRequest, ContinuousBatchScheduler +from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler, + plan_prefill_buckets) class _ForwardProfiler: @@ -95,6 +98,13 @@ class HybridEngineRolloutConfig: align_decode_fronts: bool = False enable_cache_trimming: bool = False continuous_cache_capacity: Optional[int] = None + adaptive_prefill: bool = False + prefill_fixed_cost_ms: float = 25.0 + prefill_token_cost_ms: float = 0.1 + prefill_attention_cost_ms: float = 0.01 + prefill_kv_cost_ms: float = 0.04 + continuous_decode_cost_ms: float = 10.0 + prefill_max_tokens: Optional[int] = 65536 class HybridEngineRollout(RolloutEngine): @@ -113,16 +123,31 @@ def __init__(self, engine, tokenizer, cfg=None): self.enable_profiling = getattr(cfg, 'enable_profiling', False) if cfg else False self.use_shared_prefill = getattr(cfg, 'use_shared_prefill', False) if cfg else False self.align_decode_fronts = getattr(cfg, 'align_decode_fronts', False) if cfg else False + self.adaptive_prefill = getattr(cfg, 'adaptive_prefill', False) if cfg else False + for name, default in (("prefill_fixed_cost_ms", 25.0), ("prefill_token_cost_ms", 0.1), + ("prefill_attention_cost_ms", 0.01), ("prefill_kv_cost_ms", + 0.04), ("continuous_decode_cost_ms", 10.0)): + value = getattr(cfg, name, default) if cfg else default + if not math.isfinite(value) or value < 0: + raise ValueError(f"{name} must be finite and non-negative") + setattr(self, name, value) + self.prefill_max_tokens = getattr(cfg, 'prefill_max_tokens', 65536) if cfg else 65536 + if self.prefill_max_tokens is not None: + if not isinstance(self.prefill_max_tokens, int) or self.prefill_max_tokens <= 0: + raise ValueError("prefill_max_tokens must be a positive integer or None") self.enable_cache_trimming = getattr(cfg, 'enable_cache_trimming', False) if cfg else False self.continuous_cache_capacity = getattr(cfg, 'continuous_cache_capacity', None) if cfg else None self._last_profile = None self._last_continuous_stats = None + self._generation_defaults = 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) + return self._generate_batched(request, sampling) + def _generate_batched(self, request, sampling, generation_config=None, mask_padding=True): device = request.prompt_ids.device B = request.prompt_ids.shape[0] n = sampling.n_samples_per_prompt @@ -188,6 +213,12 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout } if do_sample: generate_kwargs["top_k"] = max(sampling.top_k, 0) + if generation_config is not None: + generate_kwargs["generation_config"] = generation_config + prepare_config = getattr(module, "_prepare_generation_config", None) + if callable(prepare_config) and "use_model_defaults" in signature(prepare_config).parameters: + # Auto already resolved and validated all model defaults. + generate_kwargs["use_model_defaults"] = False output_ids = module.generate(prompt_ids, **generate_kwargs) finally: for handle in shared_prefill_handles: @@ -208,6 +239,7 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout response_start=prompt_len, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=pad_token_id, + mask_padding=mask_padding, ) # Build attention mask: pad positions (both left padding from prompt @@ -273,21 +305,50 @@ def _generate_continuous(self, request, sampling, max_batch_size): greedy decoding. Completed rows retire immediately and pending prompts prefill into the released rows before the next decode step. """ + profile = self._start_continuous_profile() if self.adaptive_prefill and self.enable_profiling else None 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) + batch_size = request.prompt_ids.shape[0] + effective_lengths = self._validate_continuous_inputs(request, sampling, max_batch_size) module = self.engine.module - profile = self._start_continuous_profile() if self.enable_profiling else None - profile_accelerator = profile["accelerator"] if profile is not None else None - profile_start = profile["start"] if profile is not None else None - prompt_lengths = { - request_id: (int(request.prompt_attention_mask.sum().item()) - if self.align_decode_fronts else request.prompt_ids.shape[1]) - for request_id, request in enumerate(requests) - } + prepared_generation_config = None + if self.adaptive_prefill and dist.is_initialized() and dist.get_world_size() > 1: + raise ValueError("adaptive prefill currently supports single-process rollouts only") + align_decode_fronts = self.adaptive_prefill or self.align_decode_fronts + if self.adaptive_prefill: + generation_config = getattr(module, "generation_config", None) + if generation_config is not None: + from transformers import GenerationConfig, GenerationMixin + # HF can refresh untouched defaults from legacy model settings during generate(). + prepare_config = getattr(module, "_prepare_generation_config", None) + if callable(prepare_config): + native_prepare = getattr(prepare_config, "__func__", + None) is GenerationMixin._prepare_generation_config + needs_prepare = not native_prepare or getattr(generation_config, "cache_implementation", + None) == "hybrid" + if getattr(generation_config, "_from_model_config", False): + legacy_settings = getattr(module.config, "_get_non_default_generation_parameters", None) + needs_prepare = needs_prepare or not callable(legacy_settings) or bool(legacy_settings()) + if needs_prepare: + generation_config, _ = prepare_config(None) + prepared_generation_config = generation_config + if self._generation_defaults is None: + self._generation_defaults = GenerationConfig() + defaults = self._generation_defaults + supported_settings = { + "max_length", "max_new_tokens", "do_sample", "temperature", "top_p", "top_k", "bos_token_id", + "eos_token_id", "pad_token_id", "repetition_penalty", "transformers_version" + } + # CB implements repetition penalty; other HF processors can change tokens across routes. + for name, value in vars(generation_config).items(): + if name.startswith("_") or name in supported_settings: + continue + if not hasattr(defaults, name) or value != getattr(defaults, name): + raise ValueError(f"adaptive prefill does not support non-default generation setting: {name}") + if align_decode_fronts: + prompt_lengths = dict(enumerate(map(int, effective_lengths))) + else: + prompt_lengths = dict.fromkeys(range(batch_size), request.prompt_ids.shape[1]) prompt_len = max(prompt_lengths.values()) max_positions = getattr(module.config, "max_position_embeddings", None) if max_positions is not None: @@ -298,18 +359,73 @@ def _generate_continuous(self, request, sampling, max_batch_size): raise ValueError("continuous_cache_capacity must be at least the padded prompt length") max_cache_len = self.continuous_cache_capacity else: - max_cache_len = (prompt_len + sampling.max_new_tokens if self.align_decode_fronts else - self._estimate_continuous_cache_len(prompt_len, [sampling.max_new_tokens] * - len(requests), max_batch_size)) + max_cache_len = (prompt_len + + sampling.max_new_tokens if align_decode_fronts else self._estimate_continuous_cache_len( + prompt_len, [sampling.max_new_tokens] * batch_size, 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 getattr(module, "_supports_cache_class", None) is False: raise ValueError("continuous batching requires a model with cache-class support; use the default " "generate() path or upgrade transformers") + use_batched = False + if self.adaptive_prefill: + if self.prefill_max_tokens is not None and prompt_len > self.prefill_max_tokens: + raise ValueError("a prompt exceeds the prefill token limit; increase the configured limit") + full_width = request.prompt_ids.shape[1] + full_tokens_fit = self.prefill_max_tokens is None or request.prompt_ids.numel() <= self.prefill_max_tokens + full_positions_fit = max_positions is None or full_width + sampling.max_new_tokens <= max_positions + if (batch_size <= max_batch_size and full_tokens_fit and full_positions_fit + and self.continuous_cache_capacity is None and not self.enable_cache_trimming): + prefill_cost = self._prefill_cost_function(module) + batched_cost = prefill_cost(batch_size, full_width, include_kv=False) + decode_cost = self.continuous_decode_cost_ms * (sampling.max_new_tokens - 1) + # Every partition pays at least one Forward plus each request's unpadded work. + # If ordinary generation already beats that bound, there is no reason to run DP. + if min(prompt_lengths.values()) == prompt_len: + lower_bound = prefill_cost(batch_size, prompt_len) + else: + lower_bound = self.prefill_fixed_cost_ms + sum( + prefill_cost(1, length) - self.prefill_fixed_cost_ms for length in prompt_lengths.values()) + use_batched = batched_cost <= lower_bound + decode_cost + if not use_batched: + _, bucket_cost = plan_prefill_buckets(list(prompt_lengths.values()), prefill_cost, + self.prefill_max_tokens) + use_batched = batched_cost <= bucket_cost + decode_cost + if use_batched: + # A bounded rollout with no pending rows can use the existing batch path. + output = self._generate_batched(request, + replace(sampling, continuous_batch_size=None), + generation_config=prepared_generation_config, + mask_padding=False) + self._last_continuous_stats = None + if self.enable_profiling: + total_ms = (self._profile_start(profile) - profile["start"]) * 1000 + profile = self._last_profile + routing_ms = total_ms - profile["total_ms"] + profile["generation_ms"] += routing_ms + profile["generation_overhead_ms"] += routing_ms + profile["total_ms"] = total_ms + profile["generation_strategy"] = "batched" + generated_tokens = int(output.attention_mask[:, request.prompt_ids.shape[1]:].sum().item()) + profile["num_generated_tokens"] = generated_tokens + total_ms = profile["total_ms"] + profile["tokens_per_second"] = generated_tokens * 1000 / total_ms if total_ms > 0 else 0.0 + profile["active_batch_size"] = batch_size + profile["continuous_batch_size"] = max_batch_size + return output + + if self.enable_profiling and profile is None: + profile = self._start_continuous_profile() + profile_accelerator = profile["accelerator"] if profile is not None else None + profile_start = profile["start"] if profile is not None else None + from transformers import StaticCache from deepspeed.utils.static_cache import DeepSpeedStaticCache + requests = tuple( + RolloutRequest(request.prompt_ids[index:index + 1], request.prompt_attention_mask[index:index + 1]) + for index in range(batch_size)) device = requests[0].prompt_ids.device model_dtype = next(module.parameters()).dtype @@ -318,7 +434,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): request_by_id = {} responses = {} request_order = (sorted(range(len(requests)), key=lambda request_id: -prompt_lengths[request_id]) - if self.align_decode_fronts else range(len(requests))) + if align_decode_fronts else range(len(requests))) for request_id in request_order: request = requests[request_id] scheduler.submit(ContinuousBatchRequest(request_id)) @@ -379,13 +495,13 @@ 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 self.align_decode_fronts: + if align_decode_fronts: span_starts[:survivor_count] = [span_starts[index] for index in update.keep_slots] span_starts[survivor_count:] = [0] * (max_batch_size - survivor_count) dead_prefix = 0 if self.enable_cache_trimming: - if self.align_decode_fronts: + if align_decode_fronts: dead_prefix = min(span_starts[:survivor_count]) else: trim_threshold = max(1, prompt_len) @@ -393,7 +509,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): if dead_prefix < trim_threshold and cache_position < max_cache_len - trim_threshold: dead_prefix = 0 elif cache_position >= max_cache_len: - if self.align_decode_fronts: + if align_decode_fronts: dead_prefix = min(span_starts[:survivor_count]) else: dead_prefix = self._continuous_dead_prefix(attention_mask, survivor_count) @@ -422,7 +538,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): if profile_accelerator is not None: profile_accelerator.synchronize() stats["trim_latency_ms"] += (time.perf_counter() - trim_start) * 1000.0 - if self.align_decode_fronts: + if align_decode_fronts: span_starts[:survivor_count] = [start - dead_prefix for start in span_starts[:survivor_count]] else: cache.reset() @@ -446,10 +562,11 @@ def _generate_continuous(self, request, sampling, max_batch_size): prompt_lengths, model_dtype, device, - self.align_decode_fronts, + align_decode_fronts, profile, + generation_config=prepared_generation_config, ) - if self.align_decode_fronts: + if align_decode_fronts: for admitted, target_row in zip(update.admitted, update.admitted_slots): span_starts[target_row] = cache_position - prompt_lengths[admitted.request_id] @@ -460,7 +577,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): if cache_position >= max_cache_len: raise self._continuous_cache_exhaustion_error(max_cache_len) write_positions[:survivor_count].fill_(cache_position) - if self.align_decode_fronts: + if align_decode_fronts: position_ids = torch.tensor([logical_positions[request_id] for request_id in survivor_ids], dtype=torch.long, device=device).unsqueeze(1) @@ -484,6 +601,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): request_by_id, responses, module, + generation_config=prepared_generation_config, ) stats["decode_steps"] += 1 decoded_tokens = dict(zip(survivor_ids, decoded.split(1, dim=0))) @@ -524,6 +642,8 @@ def _generate_continuous(self, request, sampling, max_batch_size): if profile is not None: self._finish_continuous_profile(profile, original_request, responses, max_batch_size, prompt_len, generation_end, post_processing_end) + self._last_profile["generation_strategy"] = ("bucketed" if self.adaptive_prefill else + "aligned" if align_decode_fronts else "equal_width") return output @staticmethod @@ -568,8 +688,8 @@ def _estimate_continuous_cache_len(prompt_len, max_new_tokens, max_batch_size): return max_cache_len - def _validate_continuous_inputs(self, requests, sampling, max_batch_size): - if not requests: + def _validate_continuous_inputs(self, request, sampling, max_batch_size): + if request.prompt_ids.shape[0] == 0: raise ValueError("continuous batching requires at least one request") if max_batch_size <= 0: raise ValueError("max_batch_size must be positive") @@ -578,23 +698,17 @@ def _validate_continuous_inputs(self, requests, sampling, max_batch_size): if self.use_shared_prefill: raise ValueError("continuous batching does not support shared prompt prefill") - 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 padded prompt widths") - if request.prompt_ids.device != device: - raise ValueError("continuous batching requests must use the same device") - if not request.prompt_attention_mask.any(): - raise ValueError("continuous batching requires at least one prompt token per request") + # Rollout masks contain 0/1; validate and read all lengths in one device-to-host transfer. + lengths = request.prompt_attention_mask.sum(dim=1).tolist() + if any(length <= 0 for length in lengths): + raise ValueError("continuous batching requires at least one prompt token per request") + return lengths @staticmethod def _continuous_dead_prefix(attention_mask, active_count): @@ -618,7 +732,8 @@ def _continuous_prefill(self, model_dtype, device, align_decode_fronts=False, - profile=None): + profile=None, + generation_config=None): if not update.admitted: return {} @@ -636,6 +751,7 @@ def _continuous_prefill(self, prompt_attention, model_dtype, device, + generation_config=generation_config, ) self._profile_end(profile, "prefill_forward_ms", prefill_start, count="num_prefill_forwards") cache_copy_start = self._profile_start(profile) @@ -655,38 +771,76 @@ def _continuous_prefill(self, return dict(zip(admitted_ids, prefill_tokens.split(1, dim=0))) admitted_tokens = {} - for admitted, target_row in zip(update.admitted, update.admitted_slots): - request_id = admitted.request_id - request = request_by_id[request_id] - prompt_len = prompt_lengths[request_id] - valid = request.prompt_attention_mask[0].bool() - prompt_ids = request.prompt_ids[:, valid] - prompt_attention = torch.ones_like(prompt_ids) + admitted = list(zip(update.admitted, update.admitted_slots)) + if self.adaptive_prefill: + lengths = [prompt_lengths[item.request_id] for item, _ in admitted] + buckets, _ = plan_prefill_buckets(lengths, self._prefill_cost_function(module), self.prefill_max_tokens) + else: + buckets = tuple((row, ) for row in range(len(admitted))) + pad_token_id = self.tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = self.tokenizer.eos_token_id + for bucket in buckets: + width = max(prompt_lengths[admitted[row][0].request_id] for row in bucket) + prompt_ids = torch.full((len(bucket), width), pad_token_id, dtype=torch.long, device=device) + prompt_attention = torch.zeros_like(prompt_ids) + prompt_requests = {} + for source_row, row in enumerate(bucket): + item, _ = admitted[row] + request = request_by_id[item.request_id] + length = prompt_lengths[item.request_id] + prompt_ids[source_row, -length:] = request.prompt_ids[0, request.prompt_attention_mask[0].bool()] + prompt_attention[source_row, -length:] = 1 + prompt_requests[source_row] = request prefill_start = self._profile_start(profile) - prefill_tokens, prefill_cache = self._continuous_prefill_cache( - module, - static_cache_type, - prompt_ids, - prompt_attention, - model_dtype, - device, - ) + prefill_tokens, prefill_cache = self._continuous_prefill_cache(module, + static_cache_type, + prompt_ids, + prompt_attention, + model_dtype, + device, + prompt_requests=prompt_requests, + generation_config=generation_config) self._profile_end(profile, "prefill_forward_ms", prefill_start, count="num_prefill_forwards") cache_copy_start = self._profile_start(profile) - admitted_tokens[request_id] = prefill_tokens - cache_start = cache_position - prompt_len - if cache_start < 0: - raise ValueError("continuous batching prompt does not fit behind the decode front") - for layer_idx in range(self._cache_layer_count(prefill_cache)): - prefill_keys, prefill_values = self._cache_layer_tensors(prefill_cache, layer_idx) - target_layer = cache.layers[layer_idx] - target_layer.keys[target_row, :, cache_start:cache_position].copy_(prefill_keys[0]) - target_layer.values[target_row, :, cache_start:cache_position].copy_(prefill_values[0]) - attention_mask[target_row, cache_start:cache_position].fill_(1) - write_positions[target_row] = cache_position + for source_row, row in enumerate(bucket): + item, target_row = admitted[row] + length = prompt_lengths[item.request_id] + cache_start = cache_position - length + if cache_start < 0: + raise ValueError("continuous batching prompt does not fit behind the decode front") + admitted_tokens[item.request_id] = prefill_tokens[source_row:source_row + 1] + for layer_idx in range(self._cache_layer_count(prefill_cache)): + keys, values = self._cache_layer_tensors(prefill_cache, layer_idx) + target_layer = cache.layers[layer_idx] + target_layer.keys[target_row, :, cache_start:cache_position].copy_(keys[source_row, :, -length:]) + target_layer.values[target_row, :, cache_start:cache_position].copy_(values[source_row, :, + -length:]) + attention_mask[target_row, cache_start:cache_position].fill_(1) + write_positions[target_row] = cache_position self._profile_end(profile, "cache_management_overhead_ms", cache_copy_start) return admitted_tokens + def _prefill_cost_function(self, module): + """Read model geometry once; DP candidates use only scalar arithmetic.""" + config = getattr(module.config, "text_config", module.config) + layers = getattr(config, "num_hidden_layers", 1) + heads = getattr(config, "num_attention_heads", 1) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + head_dim = getattr(config, "hidden_size", heads) // heads + kv_heads = getattr(config, "num_key_value_heads", None) or heads + attention_cost = self.prefill_attention_cost_ms * 4 * layers * heads * head_dim / 1e9 + element_size = next(module.parameters()).element_size() + kv_cost = self.prefill_kv_cost_ms * 2 * layers * kv_heads * head_dim * element_size / (1024 * 1024) + + def cost(count, width, include_kv=True): + tokens = count * width + return self.prefill_fixed_cost_ms + tokens * (self.prefill_token_cost_ms + attention_cost * width + + (kv_cost if include_kv else 0)) + + return cost + def _start_continuous_profile(self): accelerator = get_accelerator() accelerator.synchronize() @@ -751,7 +905,15 @@ def _finish_continuous_profile(self, profile, request, responses, max_batch_size "continuous_batch_size": max_batch_size, } - def _continuous_prefill_cache(self, module, static_cache_type, prompt_ids, prompt_attention, model_dtype, device): + def _continuous_prefill_cache(self, + module, + static_cache_type, + prompt_ids, + prompt_attention, + model_dtype, + device, + prompt_requests=None, + generation_config=None): position_ids = self._prefill_position_ids(prompt_attention) if getattr(module, "_supports_cache_class", None) is None: pad_token_id = self.tokenizer.pad_token_id @@ -766,31 +928,54 @@ def _continuous_prefill_cache(self, module, static_cache_type, prompt_ids, promp eos_token_id=None, pad_token_id=pad_token_id, return_dict_in_generate=True, + **({ + "output_logits": True + } if self.adaptive_prefill else {}), ) - return prefill_output.sequences[:, -1:], prefill_output.past_key_values - - prompt_len = prompt_ids.shape[1] - prefill_cache = self._create_static_cache(static_cache_type, module.config, prompt_ids.shape[0], 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), - position_ids=position_ids, - ) - return prefill_output.logits[:, -1, :].argmax(dim=-1, keepdim=True), prefill_cache + prefill_cache = prefill_output.past_key_values + if not self.adaptive_prefill: + return prefill_output.sequences[:, -1:], prefill_cache + logits = prefill_output.logits[-1] + else: + prompt_len = prompt_ids.shape[1] + prefill_cache = self._create_static_cache(static_cache_type, module.config, prompt_ids.shape[0], + 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), + position_ids=position_ids, + ) + logits = prefill_output.logits[:, -1, :] + request_ids = range(prompt_ids.shape[0]) + if prompt_requests is None: + prompt_requests = { + row: RolloutRequest(prompt_ids[row:row + 1], prompt_attention[row:row + 1]) + for row in request_ids + } + # Trimming model input must not trim HF's repetition-penalty history. + next_tokens = self._continuous_next_tokens(logits, + request_ids, + prompt_requests, + dict.fromkeys(request_ids, ()), + module, + generation_config=generation_config) + return next_tokens, prefill_cache @staticmethod - def _continuous_next_tokens(logits, request_ids, request_by_id, responses, module): - repetition_penalty = getattr(getattr(module, "generation_config", None), "repetition_penalty", 1.0) + def _continuous_next_tokens(logits, request_ids, request_by_id, responses, module, generation_config=None): + if generation_config is None: + generation_config = getattr(module, "generation_config", None) + repetition_penalty = getattr(generation_config, "repetition_penalty", 1.0) if repetition_penalty is None: repetition_penalty = 1.0 if repetition_penalty == 1.0: return logits.argmax(dim=-1, keepdim=True) + logits = logits.float() next_tokens = [] for row, request_id in enumerate(request_ids): input_ids = torch.cat((request_by_id[request_id].prompt_ids, *responses[request_id]), dim=1) @@ -943,10 +1128,11 @@ def expand_prompt_output(_module, _args, _kwargs, output): return pre_handle, post_handle @staticmethod - def _pad_after_eos(output_ids, response_start, eos_token_id, pad_token_id): + def _pad_after_eos(output_ids, response_start, eos_token_id, pad_token_id, mask_padding=True): """Retain the first response EOS and pad every subsequent position.""" response_ids = output_ids[:, response_start:] - response_attn = (response_ids != pad_token_id) + response_attn = (response_ids != pad_token_id) if mask_padding else torch.ones_like(response_ids, + dtype=torch.bool) if eos_token_id is None or response_ids.shape[1] == 0: return output_ids, response_attn.long() @@ -962,7 +1148,7 @@ def _pad_after_eos(output_ids, response_start, eos_token_id, pad_token_id): output_ids = output_ids.clone() output_ids[:, response_start:].masked_fill_(after_first_eos, pad_token_id) # EOS is a valid generated token even when pad_token_id == eos_token_id. - response_attn = ((response_ids != pad_token_id) | first_eos) & ~after_first_eos + response_attn = (response_attn | first_eos) & ~after_first_eos return output_ids, response_attn.long() # ------------------------------------------------------------------ diff --git a/deepspeed/utils/static_cache.py b/deepspeed/utils/static_cache.py index 3c35be5db539..e1a28346b7db 100644 --- a/deepspeed/utils/static_cache.py +++ b/deepspeed/utils/static_cache.py @@ -193,6 +193,8 @@ class DeepSpeedStaticCache: Dtype for eager initialization. """ + is_compileable = True + def __init__( self, config, diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst index 6f6e5ea591ca..807302bdc773 100644 --- a/docs/code-docs/source/inference-engine.rst +++ b/docs/code-docs/source/inference-engine.rst @@ -125,6 +125,107 @@ from that pre-sorted pending queue. Cache trimming remains disabled unless span with ``continuous_cache_capacity``; if the span is exhausted, the rollout raises an error that names both remedies. +Set ``HybridEngineRolloutConfig(adaptive_prefill=True)`` to choose ordinary +batched generation or aligned CB with cost-based prefill buckets. Auto is +opt-in and takes precedence over ``align_decode_fronts`` without modifying the +caller's configuration. A length-sorted O(n²) dynamic program minimizes the +additive estimated cost of contiguous buckets, including one large bucket and +one request per bucket as candidates. Each bucket uses ordinary left padding +and a two-dimensional attention mask; no packed/varlen kernel is required. +The CB cache layout remains right-aligned for the entire call. Newly admitted +requests are replanned together on every refill, without reordering survivors. + +The estimated bucket cost is a fixed Forward term plus linear padded-token +work, quadratic Attention work, and KV traffic. Attention work uses the model's +layer count, query heads and ``head_dim``; KV traffic uses KV heads, ``head_dim`` +and parameter element size. Projection/MLP work is represented by the calibrated +linear coefficient rather than inferred from head dimension alone. + +The following coefficients are starting estimates, not portable guarantees: + +* ``prefill_fixed_cost_ms=25.0``: fixed cost per prefill Forward. +* ``prefill_token_cost_ms=0.1``: milliseconds per padded token position. +* ``prefill_attention_cost_ms=0.01``: milliseconds per estimated GFLOP of + Attention work (``4 * B * L² * layers * query_heads * head_dim``). +* ``prefill_kv_cost_ms=0.04``: milliseconds per MiB of estimated KV traffic. +* ``continuous_decode_cost_ms=10.0``: extra CB milliseconds per Decode step + (generation budget minus the first token returned by Prefill), used when + comparing with ordinary generation. + +All costs must be finite and non-negative. Calibrate them for the model, +device, precision and attention backend. DP is optimal for these estimates; +it cannot guarantee globally minimal measured GPU latency. Ordinary generation +is eligible only when all requests fit the active-row and prefill limits, the +physical input width fits the model position limit, and neither explicit static +capacity nor trimming is requested. Otherwise Auto remains on the CB path. + +Auto validates masks and reads effective lengths in one batch transfer. When +ordinary generation beats a lower bound on every CB partition, it skips DP; +uniform prompt lengths also admit an exact single-bucket estimate. The fallback +reuses the resolved generation configuration and normalizes the result once. +This keeps planning overhead small for short generations and large request counts. + +The `DeepSpeedExamples adaptive prefill benchmark +`_ +provides separate calibration and evaluation. From the DeepSpeedExamples checkout: + +.. code-block:: bash + + python benchmarks/opsd/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ + --repeats 3 --calibrate costs.json + python benchmarks/opsd/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ + --cost-config costs.json --repeats 3 --output sweep.json + +Use ``--model tiny-qwen2 --dtype float32`` for a seeded small-model smoke test. +Calibration uses six uniform shapes distinct from the evaluation workloads, +fits non-negative fixed/token/Attention terms, measures KV management rates, +and estimates extra CB Decode cost from one-token and eight-token calls. +The saved JSON records model, GPU, precision, backend and library versions; +the sweep rejects a calibration from a different configuration. Applications +can pass its ``costs`` dictionary to ``HybridEngineRolloutConfig``. +Calibration runs during deployment preparation; applications load the saved +coefficients once at rollout initialization, without profiling or refitting +inside generation requests. The built-in grid reaches 768 input tokens; extend +it to cover the deployment's lengths and feasible batch sizes before relying +on these estimates for longer contexts. + +The four evaluation workloads cover 128 short requests, a 128/4 tail, a +16..512 spread and a 512/16 tail. Timing includes routing and cache work, with +one warmup and interleaved repetitions of all four paths. The report includes +``auto_vs_best`` and the worst-case regression, plus selected routes and output +ID agreement. Generation is deliberately limited to two tokens with EOS +disabled to expose short-generation overhead; this is a controlled performance +test rather than a natural-EOS rollout trace. Use ``--new-tokens`` to evaluate +other budgets and recalibrate when the deployment configuration changes. + +``prefill_max_tokens=65536`` limits padded token positions in each Forward; +set it to ``None`` to remove this planning limit. A single prompt above the +configured limit raises an error; this does not implement chunked prefill. +Long prompts also need a model position limit covering the prompt and generation +budget, sufficient static KV capacity, and enough temporary Forward memory. +The token limit does not bound allocated Decode KV memory; reduce the configured +active-row capacity or cache capacity when needed. + +Automatic selection currently requires single-process greedy rollout with one +sample per prompt and a supported cache-class model. Other non-default generation +settings beyond repetition penalty, token IDs and the supported length/sampling +settings are rejected; for example, beam search, no-repeat n-grams, forced tokens +and dictionary outputs cannot be reproduced by the continuous path. Auto resolves +HF's effective generation configuration before validation, including legacy model +settings, while retaining explicit generation-configuration overrides. Repetition +penalty retains the original padded prompt history even when bucket model input +is trimmed. + +With profiling, ``get_last_profile()["generation_strategy"]`` reports +``"batched"`` or ``"bucketed"`` for Auto; ``num_prefill_forwards`` counts actual +bucket Forwards on the CB path. End-to-end time includes initial route planning +and result normalization. Ordinary generation has no static-cache snapshot, +so ``get_last_continuous_stats()`` returns ``None`` for that call. Across Auto +paths, productive-token counts and throughput exclude structural padding and +retain EOS plus PAD-valued tokens generated before termination; active/configured +row-capacity fields retain their continuous-call meaning. Manual alignment and +equal-width modes retain their existing profile strategy labels. + The most recent cache statistics are available from ``rollout.get_last_continuous_stats()``. They include ``cache_capacity``, ``peak_cache_length``, ``cache_memory_bytes``, ``trim_count``, diff --git a/tests/unit/runtime/rollout/test_continuous_batching.py b/tests/unit/runtime/rollout/test_continuous_batching.py index 2204112b09c2..2a8b0e2edaea 100644 --- a/tests/unit/runtime/rollout/test_continuous_batching.py +++ b/tests/unit/runtime/rollout/test_continuous_batching.py @@ -2,9 +2,12 @@ # DeepSpeed Team +from itertools import product + import pytest -from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler) +from deepspeed.runtime.rollout.continuous_batching import (ContinuousBatchRequest, ContinuousBatchScheduler, + plan_prefill_buckets) def _request(request_id): @@ -78,3 +81,44 @@ def test_scheduler_rejects_invalid_transitions(): scheduler.submit(_request("a")) with pytest.raises(ValueError, match="not active"): scheduler.schedule(finished_ids=("missing", )) + + +@pytest.mark.parametrize("lengths,max_tokens", [([4, 1, 16, 4], None), ([2, 8, 3, 7, 1], 16)]) +def test_prefill_buckets_match_exhaustive_partition_cost(lengths, max_tokens): + + def cost(count, width): + return 5.0 + count * width + 0.01 * count * width * width + + order = sorted(range(len(lengths)), key=lambda i: -lengths[i]) + reference = float("inf") + for cuts in product([False, True], repeat=len(lengths) - 1): + starts = [0] + [i + 1 for i, split in enumerate(cuts) if split] + ends = starts[1:] + [len(lengths)] + groups = [order[a:b] for a, b in zip(starts, ends)] + if max_tokens is not None and any(len(g) * max(lengths[i] for i in g) > max_tokens for g in groups): + continue + reference = min(reference, sum(cost(len(g), max(lengths[i] for i in g)) for g in groups)) + + buckets, predicted = plan_prefill_buckets(lengths, cost, max_tokens) + assert sorted(i for bucket in buckets for i in bucket) == list(range(len(lengths))) + assert predicted == pytest.approx(reference) + assert predicted == pytest.approx(sum(cost(len(g), max(lengths[i] for i in g)) for g in buckets)) + + +@pytest.mark.parametrize("fixed_cost,expected_count", [(25.0, 2), (5000.0, 1)]) +def test_prefill_bucket_count_responds_to_forward_cost(fixed_cost, expected_count): + lengths = [16] * 31 + [512] + buckets, _ = plan_prefill_buckets(lengths, lambda count, width: fixed_cost + 0.1 * count * width) + assert len(buckets) == expected_count + + +def test_prefill_bucket_capacity_supports_long_lengths_without_token_expansion(): + lengths = [65536, 16, 16] + buckets, _ = plan_prefill_buckets(lengths, lambda count, width: 1 + count * width, 65536) + assert sorted(i for bucket in buckets for i in bucket) == [0, 1, 2] + assert all(len(g) * max(lengths[i] for i in g) <= 65536 for g in buckets) + with pytest.raises(ValueError, match="prefill.*limit"): + plan_prefill_buckets([65537], lambda count, width: count * width, 65536) + assert plan_prefill_buckets([], lambda count, width: count * width) == ((), 0.0) + with pytest.raises(ValueError, match="positive"): + plan_prefill_buckets([0], lambda count, width: count * width) diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 1c5be4984a64..61b9e0d096f7 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -7,6 +7,7 @@ the transformer inference extension are available. """ +import time from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -57,6 +58,8 @@ def test_config_defaults(): assert cfg.enable_profiling is False assert cfg.use_shared_prefill is False assert cfg.align_decode_fronts is False + assert cfg.adaptive_prefill is False + assert cfg.prefill_max_tokens == 65536 assert cfg.enable_cache_trimming is False assert cfg.continuous_cache_capacity is None @@ -93,8 +96,10 @@ def test_constructor_defaults_without_cfg(): assert rollout.continuous_cache_capacity is None -def test_continuous_generation_rejects_unsupported_inputs(): - rollout = HybridEngineRollout(_make_engine(), _make_tokenizer()) +@pytest.mark.parametrize("adaptive_prefill", [False, True]) +def test_continuous_generation_rejects_unsupported_inputs(adaptive_prefill): + rollout = HybridEngineRollout(_make_engine(), _make_tokenizer(), + HybridEngineRolloutConfig(adaptive_prefill=adaptive_prefill)) request = RolloutRequest( prompt_ids=torch.tensor([[0, 1, 2]]), prompt_attention_mask=torch.tensor([[0, 1, 1]]), @@ -131,7 +136,9 @@ def test_continuous_cache_span_does_not_sum_independent_requests(): assert cache_len < 64 + 64 * 100 -def test_continuous_generation_validates_each_request_length(): +@pytest.mark.parametrize("adaptive_prefill", [False, True]) +@pytest.mark.parametrize("length,max_positions", [(3, 4), (65536, 32768)]) +def test_continuous_generation_validates_each_request_length(adaptive_prefill, length, max_positions): class LimitedModel(torch.nn.Module): @@ -140,11 +147,13 @@ class LimitedModel(torch.nn.Module): def __init__(self): super().__init__() self.weight = torch.nn.Parameter(torch.zeros(1)) - self.config = SimpleNamespace(max_position_embeddings=4) + self.config = SimpleNamespace(max_position_embeddings=max_positions) 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)) + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=adaptive_prefill)) + prompt = torch.ones((1, length), dtype=torch.long) + request = RolloutRequest(prompt, torch.ones_like(prompt)) with pytest.raises(ValueError, match="request exceeds"): rollout.generate(request, SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1)) @@ -410,7 +419,8 @@ def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=Tru assert rollout.get_last_continuous_stats()["trim_count"] == 1 -def test_continuous_generation_trims_cache_after_staggered_eos(): +@pytest.mark.parametrize("adaptive_prefill", [False, True]) +def test_continuous_generation_trims_cache_after_staggered_eos(adaptive_prefill): class CacheConfig(SimpleNamespace): @@ -453,7 +463,7 @@ def forward(self, input_ids, attention_mask, past_key_values=None, use_cache=Tru rollout = HybridEngineRollout( SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=2), - cfg=HybridEngineRolloutConfig(enable_cache_trimming=True), + cfg=HybridEngineRolloutConfig(enable_cache_trimming=True, adaptive_prefill=adaptive_prefill), ) request = RolloutRequest( torch.tensor([[1, 2, 3], [1, 2, 4], [1, 2, 5], [1, 2, 6], [1, 2, 7], [1, 2, 8]]), @@ -570,6 +580,20 @@ def test_continuous_generation_treats_none_repetition_penalty_as_one(): assert next_tokens.tolist() == [[0]] +def test_continuous_repetition_penalty_matches_hf_in_bfloat16(): + from transformers import RepetitionPenaltyLogitsProcessor + + logits = torch.tensor([[1.0, 0.333984375]], dtype=torch.bfloat16) + prompt = torch.tensor([[0]]) + request = RolloutRequest(prompt, torch.ones_like(prompt)) + model = SimpleNamespace(generation_config=SimpleNamespace(repetition_penalty=3.0)) + expected = RepetitionPenaltyLogitsProcessor(3.0)(prompt, logits.float()).argmax(dim=-1, keepdim=True) + # Pin the BF16 rounding bug: dividing in BF16 ties the scores and incorrectly selects token 0. + actual = HybridEngineRollout._continuous_next_tokens(logits, (0, ), {0: request}, {0: ()}, model) + assert expected.tolist() == [[1]] + assert torch.equal(actual, expected) + + @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): @@ -1032,3 +1056,428 @@ def test_static_cache_falls_back_to_derived_head_dim(): cache = DeepSpeedStaticCache(_cache_config(), batch_size=1, max_cache_len=8, device="cpu", dtype=torch.float32) assert tuple(cache.layers[0].keys.shape) == (1, 2, 8, 16) + + +def _make_small_qwen(): + from transformers import Qwen2Config, Qwen2ForCausalLM + from deepspeed.accelerator import get_accelerator + + torch.manual_seed(8497) + device = get_accelerator().device_name() + config = Qwen2Config(vocab_size=32, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + pad_token_id=0, + eos_token_id=None) + return Qwen2ForCausalLM(config).to(device).eval() + + +@pytest.mark.parametrize("lengths,capacity,strategy", [ + ([4, 4, 3, 3], 4, "batched"), + ([4, 4, 3, 3], 2, "bucketed"), + ([16, 1, 1, 1], 4, "batched"), + ([4, 2, 1, 1], 4, "batched"), + ([1], 1, "batched"), +]) +def test_adaptive_generation_matches_single_request_eager(lengths, capacity, strategy): + model = _make_small_qwen() + device = next(model.parameters()).device + tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=None) + width = max(lengths) + prompt_ids = torch.zeros((len(lengths), width), dtype=torch.long, device=device) + mask = torch.zeros_like(prompt_ids, dtype=torch.float32) + expected = [] + with torch.no_grad(): + for row, length in enumerate(lengths): + prompt = torch.arange(1, length + 1, device=device).unsqueeze(0) + prompt_ids[row, -length:] = prompt + mask[row, -length:] = 1 + expected.append(model.generate(prompt, max_new_tokens=3, do_sample=False)[0, length:]) + rollout = HybridEngineRollout( + SimpleNamespace(module=model), tokenizer, + HybridEngineRolloutConfig(adaptive_prefill=True, enable_profiling=True, align_decode_fronts=True)) + sampling = SamplingConfig(max_new_tokens=3, temperature=0, continuous_batch_size=capacity) + result = rollout.generate(RolloutRequest(prompt_ids, mask), sampling) + + assert torch.equal(result.input_ids[:, width:], torch.stack(expected)) + assert torch.equal(result.input_ids[:, :width], prompt_ids) + assert torch.equal(result.attention_mask[:, :width], mask) + assert result.response_start_idx.tolist() == [width] * len(lengths) + assert rollout.get_last_profile()["generation_strategy"] == strategy + assert sampling.continuous_batch_size == capacity + assert rollout.align_decode_fronts is True + if strategy == "batched": + assert rollout.get_last_continuous_stats() is None + else: + assert rollout.get_last_profile()["active_batch_size"] <= capacity + + +@pytest.mark.parametrize("adaptive_prefill", [False, True]) +def test_continuous_opt_matches_eager_with_static_cache(adaptive_prefill): + from deepspeed.accelerator import get_accelerator + from transformers import OPTConfig, OPTForCausalLM + + torch.manual_seed(8497) + config = OPTConfig(vocab_size=32, + hidden_size=32, + ffn_dim=64, + num_hidden_layers=2, + num_attention_heads=4, + max_position_embeddings=32, + pad_token_id=0, + eos_token_id=None, + attn_implementation="sdpa") + model = OPTForCausalLM(config).to(get_accelerator().device_name()).eval() + prompt = torch.tensor([[1, 4, 5], [0, 4, 5]], device=next(model.parameters()).device) + mask = prompt.ne(0).long() + with torch.no_grad(): + expected = model.generate(prompt, attention_mask=mask, max_new_tokens=3, do_sample=False) + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=adaptive_prefill)) + result = rollout.generate(RolloutRequest(prompt, mask), + SamplingConfig(max_new_tokens=3, temperature=0, continuous_batch_size=1)) + assert torch.equal(result.input_ids, expected) + assert result.attention_mask[:, :3].tolist() == mask.tolist() + + +@pytest.mark.parametrize("name,value", [ + ("prefill_fixed_cost_ms", -1.0), + ("prefill_token_cost_ms", float("nan")), + ("prefill_attention_cost_ms", float("inf")), + ("prefill_max_tokens", 0), +]) +def test_adaptive_prefill_rejects_invalid_cost_settings(name, value): + with pytest.raises(ValueError, match=name): + HybridEngineRollout(_make_engine(), _make_tokenizer(), + HybridEngineRolloutConfig(adaptive_prefill=True, **{name: value})) + + +@patch("deepspeed.comm.get_world_size", return_value=2) +@patch("deepspeed.comm.is_initialized", return_value=True) +def test_adaptive_prefill_rejects_unsynchronized_distributed_selection(_initialized, _world_size): + rollout = HybridEngineRollout(_make_engine(), _make_tokenizer(), HybridEngineRolloutConfig(adaptive_prefill=True)) + with pytest.raises(ValueError, match="single-process"): + rollout.generate(_make_request(), SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=2)) + + +def test_config_preserves_legacy_positional_arguments(): + cfg = HybridEngineRolloutConfig(False, False, False, False, True, 16) + assert cfg.enable_cache_trimming is True + assert cfg.continuous_cache_capacity == 16 + assert cfg.adaptive_prefill is False + assert cfg.prefill_max_tokens == 65536 + + +@pytest.mark.parametrize("setting,value", [ + ("no_repeat_ngram_size", 1), + ("num_beams", 2), + ("return_dict_in_generate", True), +]) +@pytest.mark.parametrize("config_source", ["generation", "model"]) +def test_adaptive_generation_rejects_unsupported_generation_settings(setting, value, config_source): + model = _make_small_qwen() + config = model.generation_config if config_source == "generation" else model.config + setattr(config, setting, value) + prompt = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]], device=next(model.parameters()).device) + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=True)) + with pytest.raises(ValueError, match=f"adaptive prefill.*{setting}"): + rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=4, temperature=0, continuous_batch_size=2)) + + +@pytest.mark.parametrize("value", [True, False, None]) +@pytest.mark.parametrize("config_source", ["generation", "prepared"]) +def test_adaptive_generation_rejects_unknown_public_settings(value, config_source): + from transformers import Qwen2Config, Qwen2ForCausalLM + + model = Qwen2ForCausalLM( + Qwen2Config(vocab_size=32, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2)).eval() + native_prepare = model._prepare_generation_config + if config_source == "generation": + model.generation_config.custom_generation_switch = value + + def prepare(*args, **kwargs): + config, model_kwargs = native_prepare(*args, **kwargs) + if config_source == "prepared": + config.custom_generation_switch = value + return config, model_kwargs + + prompt = torch.tensor([[1, 2, 3]]) + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=True)) + with patch.object(model, "_prepare_generation_config", side_effect=prepare): + with pytest.raises(ValueError, match="custom_generation_switch"): + rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=1, temperature=0, continuous_batch_size=1)) + + +@pytest.mark.parametrize("capacity", [1, 2]) +@pytest.mark.parametrize("stop_at_first", [True, False]) +def test_adaptive_profile_counts_only_productive_tokens_after_eos(capacity, stop_at_first): + model = _make_small_qwen() + device = next(model.parameters()).device + prompt_tokens = [1, 2, 3, 4] if stop_at_first else [1, 5] + width = len(prompt_tokens) + prompt = torch.tensor([prompt_tokens, prompt_tokens], device=device) + with torch.no_grad(): + if stop_at_first: + eos = model.generate(prompt[:1], max_new_tokens=1, do_sample=False)[0, -1].item() + else: + # A zero output head deterministically emits PAD 0 before EOS 2. + model.lm_head.weight.zero_() + eos = 2 + tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=eos) + rollout = HybridEngineRollout(SimpleNamespace(module=model), tokenizer, + HybridEngineRolloutConfig(adaptive_prefill=True, enable_profiling=True)) + result = rollout.generate( + RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=4 if stop_at_first else 1, temperature=0, continuous_batch_size=capacity)) + profile = rollout.get_last_profile() + assert result.attention_mask[:, width:].sum().item() == 2 + if not stop_at_first: + assert result.input_ids[:, width:].tolist() == [[0], [0]] + assert profile["num_generated_tokens"] == 2 + assert profile["tokens_per_second"] == pytest.approx(2 * 1000 / profile["total_ms"]) + assert profile["active_batch_size"] == min(capacity, 2) + assert profile["continuous_batch_size"] == capacity + + +@pytest.mark.parametrize("capacity", [1, 2]) +@pytest.mark.parametrize("config_source", ["generation", "model"]) +def test_adaptive_generation_applies_repetition_penalty_from_first_token(capacity, config_source): + from transformers import GenerationConfig + + model = _make_small_qwen() + model._supports_cache_class = True + if config_source == "generation": + # An explicitly changed generation config takes precedence over legacy model settings. + model.config.num_beams = 2 + model.generation_config.repetition_penalty = 3.0 + expected_config = model.generation_config + else: + model.config.repetition_penalty = 3.0 + expected_config = GenerationConfig.from_model_config(model.config) + prompt = torch.tensor([[11, 1, 2, 3, 4], [11, 1, 2, 3, 4]], device=next(model.parameters()).device) + with torch.no_grad(): + # Passing the oracle config explicitly leaves the target's stale default untouched. + expected = model.generate(prompt[:1], generation_config=expected_config, max_new_tokens=2, do_sample=False)[0, + 5:] + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=True)) + result = rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=capacity)) + assert torch.equal(result.input_ids[:, 5:], expected.unsqueeze(0).expand(2, -1)) + + +@pytest.mark.parametrize("cache_support,bucketed,adaptive_prefill", [(True, True, True), (True, False, True), + (None, True, True), (None, False, True), + (True, True, False)]) +def test_adaptive_generation_preserves_padded_repetition_history(cache_support, bucketed, adaptive_prefill): + model = _make_small_qwen() + model._supports_cache_class = cache_support + model.generation_config.repetition_penalty = 3.0 + device = next(model.parameters()).device + prompt = torch.full((8, 16), 11, dtype=torch.long, device=device) + mask = torch.zeros_like(prompt) + prompt[0] = torch.arange(1, 17, device=device) + mask[0] = 1 + prompt[1:, -4:] = torch.arange(1, 5, device=device) + mask[1:, -4:] = 1 + with torch.no_grad(): + expected = torch.cat([ + model.generate(prompt[row:row + 1], + attention_mask=mask[row:row + 1], + max_new_tokens=1, + do_sample=False, + eos_token_id=None, + pad_token_id=11)[:, -1:] for row in range(8) + ]) + rollout = HybridEngineRollout( + SimpleNamespace(module=model), SimpleNamespace(pad_token_id=11, eos_token_id=11), + HybridEngineRolloutConfig(adaptive_prefill=adaptive_prefill, + align_decode_fronts=not adaptive_prefill, + prefill_fixed_cost_ms=0.01 if bucketed else 1e6, + prefill_token_cost_ms=1.0, + prefill_attention_cost_ms=0.0, + prefill_kv_cost_ms=0.0, + continuous_decode_cost_ms=0.0)) + result = rollout.generate(RolloutRequest(prompt, mask), + SamplingConfig(max_new_tokens=1, temperature=0, continuous_batch_size=8)) + assert torch.equal(result.input_ids[:, 16:], expected) + assert result.attention_mask[:, 16:].tolist() == [[1]] * 8 + + +@pytest.mark.parametrize("cache_support", [True, None]) +@pytest.mark.parametrize("capacity,prefill_forwards", [(4, 2), (2, 3)]) +def test_dp_prefill_groups_requests_and_matches_eager(cache_support, capacity, prefill_forwards): + model = _make_small_qwen() + model._supports_cache_class = cache_support + device = next(model.parameters()).device + lengths = [1, 16, 2, 3] + prompt = torch.zeros((4, 16), dtype=torch.long, device=device) + mask = torch.zeros_like(prompt) + expected = [] + with torch.no_grad(): + for row, length in enumerate(lengths): + prompt[row, -length:] = torch.arange(1, length + 1, device=device) + mask[row, -length:] = 1 + expected.append( + model.generate(prompt[row:row + 1], + attention_mask=mask[row:row + 1], + max_new_tokens=3, + do_sample=False)[0, 16:]) + cfg = HybridEngineRolloutConfig(adaptive_prefill=True, + enable_profiling=True, + prefill_fixed_cost_ms=2.0, + prefill_token_cost_ms=1.0, + prefill_attention_cost_ms=0.0, + prefill_kv_cost_ms=0.0, + continuous_decode_cost_ms=0.0) + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + cfg) + result = rollout.generate(RolloutRequest(prompt, mask), + SamplingConfig(max_new_tokens=3, temperature=0, continuous_batch_size=capacity)) + assert torch.equal(result.input_ids[:, 16:], torch.stack(expected)) + assert torch.equal(result.attention_mask[:, :16], mask) + assert result.response_start_idx.tolist() == [16] * 4 + assert rollout.get_last_profile()["generation_strategy"] == "bucketed" + assert rollout.get_last_profile()["num_prefill_forwards"] == prefill_forwards + + +def test_dp_prefill_rejects_single_prompt_above_forward_limit(): + model = _make_small_qwen() + prompt = torch.ones((1, 17), dtype=torch.long, device=next(model.parameters()).device) + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=True, prefill_max_tokens=16)) + with pytest.raises(ValueError, match="prefill.*limit"): + rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=1, temperature=0, continuous_batch_size=1)) + + +def test_adaptive_short_batch_does_not_read_one_gpu_scalar_per_request(): + model = _make_small_qwen() + prompt = torch.ones((128, 4), dtype=torch.long, device=next(model.parameters()).device) + request = RolloutRequest(prompt, torch.ones_like(prompt)) + tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=None) + scalar_reads = [] + outputs = [] + for adaptive in (False, True): + rollout = HybridEngineRollout(SimpleNamespace(module=model), tokenizer, + HybridEngineRolloutConfig(adaptive_prefill=adaptive)) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as profile: + output = rollout.generate( + request, + SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=128 if adaptive else None)) + outputs.append(output.input_ids) + # Pin the reported regression: per-row mask.any() caused 128 device-to-host scalar reads. + scalar_reads.append( + sum(event.count for event in profile.key_averages() if event.key == "aten::_local_scalar_dense")) + assert torch.equal(*outputs) + assert scalar_reads[1] <= scalar_reads[0] + 4 + + +def test_adaptive_generation_resolves_custom_generation_config(): + model = _make_small_qwen() + native_prepare = model._prepare_generation_config + + def prepare(*args, **kwargs): + config, model_kwargs = native_prepare(*args, **kwargs) + config.min_length = 5 + return config, model_kwargs + + prompt = torch.ones((1, 4), dtype=torch.long, device=next(model.parameters()).device) + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=True)) + with patch.object(model, "_prepare_generation_config", side_effect=prepare): + with pytest.raises(ValueError, match="min_length"): + rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1)) + + +@pytest.mark.parametrize("cache_support", [True, None]) +@pytest.mark.parametrize("capacity", [1, 2]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_adaptive_generation_uses_resolved_penalty_on_both_routes(cache_support, capacity, dtype): + from transformers import GenerationConfig + + model = _make_small_qwen().to(dtype=dtype) + model._supports_cache_class = cache_support + native_prepare = model._prepare_generation_config + prompt = torch.tensor([[11, 1, 2, 3, 4], [11, 1, 2, 3, 4]], device=next(model.parameters()).device) + expected_config = GenerationConfig.from_model_config(model.config) + expected_config.repetition_penalty = 3.0 + with torch.no_grad(): + expected = model.generate(prompt, generation_config=expected_config, max_new_tokens=3, do_sample=False) + unpenalized = model.generate(prompt, max_new_tokens=3, do_sample=False) + assert not torch.equal(expected[:, 5:], unpenalized[:, 5:]) + + def prepare(*args, **kwargs): + config, model_kwargs = native_prepare(*args, **kwargs) + config.repetition_penalty = 3.0 + return config, model_kwargs + + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=True)) + with patch.object(model, "_prepare_generation_config", side_effect=prepare): + result = rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=3, temperature=0, continuous_batch_size=capacity)) + assert torch.equal(result.input_ids, expected) + assert model.generation_config.repetition_penalty == 1.0 + assert (rollout.get_last_continuous_stats() is None) == (capacity == 2) + + +@pytest.mark.parametrize("align_decode_fronts", [False, True]) +def test_manual_continuous_generation_omits_unsupported_output_logits(align_decode_fronts): + model = _make_small_qwen() + model._supports_cache_class = None + native_generate = model.generate + prompt = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]], device=next(model.parameters()).device) + with torch.no_grad(): + expected = native_generate(prompt, max_new_tokens=2, do_sample=False) + + def legacy_generate(*args, **kwargs): + if "output_logits" in kwargs: + raise ValueError("output_logits is not supported by this generation API") + return native_generate(*args, **kwargs) + + rollout = HybridEngineRollout(SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(align_decode_fronts=align_decode_fronts)) + with patch.object(model, "generate", side_effect=legacy_generate): + result = rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1)) + assert torch.equal(result.input_ids, expected) + + +@pytest.mark.parametrize("force_cb", [False, True]) +def test_adaptive_profile_includes_route_planning_time(force_cb): + from deepspeed.accelerator import get_accelerator + + model = _make_small_qwen() + device = next(model.parameters()).device + prompt = torch.ones((1024, 4), dtype=torch.long, device=device) + rollout = HybridEngineRollout( + SimpleNamespace(module=model), SimpleNamespace(pad_token_id=0, eos_token_id=None), + HybridEngineRolloutConfig(adaptive_prefill=True, enable_profiling=True, enable_cache_trimming=force_cb)) + get_accelerator().synchronize() + start = time.perf_counter() + result = rollout.generate(RolloutRequest(prompt, torch.ones_like(prompt)), + SamplingConfig(max_new_tokens=1, temperature=0, continuous_batch_size=1024)) + get_accelerator().synchronize() + wall_ms = (time.perf_counter() - start) * 1000 + profile = rollout.get_last_profile() + # The former profile excluded initial DP planning and over-reported throughput. + assert profile["total_ms"] >= wall_ms * 0.8 - 5 + assert profile["total_ms"] <= wall_ms + 5 + assert profile["num_generated_tokens"] == result.attention_mask[:, 4:].sum().item() + if force_cb: + assert rollout.get_last_continuous_stats()["end_to_end_ms"] >= wall_ms * 0.8 - 5