From 7cb9828038e9ec51d0215b1093fc92a90980cff1 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Fri, 9 Oct 2026 12:27:22 +0800 Subject: [PATCH 1/7] perf(rollout): select generation and prefill buckets with tunable costs Signed-off-by: SuperGoodGame <985236470@qq.com> --- .../runtime/rollout/continuous_batching.py | 35 +++ .../runtime/rollout/hybrid_engine_rollout.py | 271 ++++++++++++---- docs/code-docs/source/inference-engine.rst | 62 ++++ .../rollout/test_continuous_batching.py | 46 ++- .../rollout/test_hybrid_engine_rollout.py | 296 +++++++++++++++++- 5 files changed, 642 insertions(+), 68 deletions(-) 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..a62ceaa7e124 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,6 +123,18 @@ 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 @@ -273,6 +295,7 @@ 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]) @@ -280,14 +303,33 @@ def _generate_continuous(self, request, sampling, max_batch_size): self._validate_continuous_inputs(requests, 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) - } + 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 + # HF can refresh untouched defaults from legacy model settings during generate(). + prepare_config = getattr(module, "_prepare_generation_config", None) + if callable(prepare_config): + generation_config, _ = prepare_config(None) + defaults = GenerationConfig() + 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 value != getattr(defaults, name, value): + raise ValueError(f"adaptive prefill does not support non-default generation setting: {name}") + if align_decode_fronts: + effective_lengths = request.prompt_attention_mask.sum(dim=1).tolist() + prompt_lengths = dict(enumerate(map(int, effective_lengths))) + else: + prompt_lengths = dict.fromkeys(range(len(requests)), 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,15 +340,64 @@ 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] * 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 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 (len(requests) <= 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) + _, bucket_cost = plan_prefill_buckets(list(prompt_lengths.values()), prefill_cost, + self.prefill_max_tokens) + batched_cost = prefill_cost(len(requests), full_width, include_kv=False) + continuous_cost = bucket_cost + self.continuous_decode_cost_ms * (sampling.max_new_tokens - 1) + use_batched = batched_cost <= continuous_cost + if use_batched: + # A bounded rollout with no pending rows can use the existing batch path. + output = self.generate(request, replace(sampling, continuous_batch_size=None)) + pad_token_id = self.tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = self.tokenizer.eos_token_id + _, response_attention = self._pad_after_eos(output.input_ids, + request.prompt_ids.shape[1], + self.tokenizer.eos_token_id, + pad_token_id, + mask_padding=False) + output.attention_mask[:, request.prompt_ids.shape[1]:] = response_attention + 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"] = len(requests) + 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 @@ -318,7 +409,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 +470,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 +484,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 +513,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 +537,10 @@ def _generate_continuous(self, request, sampling, max_batch_size): prompt_lengths, model_dtype, device, - self.align_decode_fronts, + align_decode_fronts, profile, ) - 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 +551,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) @@ -524,6 +615,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 @@ -655,13 +748,27 @@ 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, @@ -670,23 +777,47 @@ def _continuous_prefill(self, prompt_attention, model_dtype, device, - ) + prompt_requests=prompt_requests if self.adaptive_prefill else None) 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 +882,14 @@ 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): 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,22 +904,36 @@ 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=self.adaptive_prefill, ) - 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) + return next_tokens, prefill_cache @staticmethod def _continuous_next_tokens(logits, request_ids, request_by_id, responses, module): @@ -943,10 +1095,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 +1115,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/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst index 6f6e5ea591ca..41b08728d15a 100644 --- a/docs/code-docs/source/inference-engine.rst +++ b/docs/code-docs/source/inference-engine.rst @@ -125,6 +125,68 @@ 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. + +``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..1b111e9ae38d 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]]), @@ -1032,3 +1042,273 @@ 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("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("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", [True, None]) +@pytest.mark.parametrize("bucketed", [True, False]) +def test_adaptive_generation_preserves_padded_repetition_history(cache_support, bucketed): + 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=True, + 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)) + + +@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 From 359b3d5002f8af11e1565ebf59c67b93b1af77cd Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sat, 10 Oct 2026 16:45:10 +0800 Subject: [PATCH 2/7] fix(rollout): reduce adaptive routing overhead and add cost calibration Signed-off-by: SuperGoodGame <985236470@qq.com> --- benchmarks/rollout_prefill.py | 305 ++++++++++++++++++ .../runtime/rollout/hybrid_engine_rollout.py | 99 +++--- docs/code-docs/source/inference-engine.rst | 32 ++ .../rollout/test_continuous_batching.py | 9 + .../rollout/test_hybrid_engine_rollout.py | 40 +++ 5 files changed, 445 insertions(+), 40 deletions(-) create mode 100644 benchmarks/rollout_prefill.py diff --git a/benchmarks/rollout_prefill.py b/benchmarks/rollout_prefill.py new file mode 100644 index 000000000000..b463529b1751 --- /dev/null +++ b/benchmarks/rollout_prefill.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Calibrate prefill costs separately from a small, reproducible route sweep.""" + +import argparse +from dataclasses import asdict +import hashlib +from itertools import combinations +import json +from pathlib import Path +import statistics +import time +from types import SimpleNamespace + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config, Qwen2ForCausalLM + +from deepspeed.runtime.rollout.base import RolloutRequest, SamplingConfig +from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout, HybridEngineRolloutConfig +import deepspeed.runtime.rollout.hybrid_engine_rollout as rollout_backend +from deepspeed.accelerator import get_accelerator + + +def nonnegative_fit(matrix, targets): + """Solve three-parameter NNLS by checking its eight possible active sets.""" + matrix = np.asarray(matrix, dtype=float) + targets = np.asarray(targets, dtype=float) + scale = np.maximum(np.abs(matrix).max(axis=0), 1e-12) + scaled = matrix / scale + best = np.zeros(matrix.shape[1]) + error = float(targets @ targets) + for count in range(1, matrix.shape[1] + 1): + for columns in combinations(range(matrix.shape[1]), count): + values = np.linalg.lstsq(scaled[:, columns], targets, rcond=None)[0] + if np.any(values < 0): + continue + candidate = np.zeros(matrix.shape[1]) + candidate[list(columns)] = values + residual = scaled @ candidate - targets + if residual @ residual < error: + best, error = candidate, float(residual @ residual) + return (best / scale).tolist() + + +def make_request(lengths, token, pad, device): + width = max(lengths) + ids = torch.full((len(lengths), width), pad, dtype=torch.long, device=device) + mask = torch.zeros_like(ids) + for row, length in enumerate(lengths): + ids[row, -length:] = token + mask[row, -length:] = 1 + return RolloutRequest(ids, mask) + + +def measure(rollout, request, steps, continuous, repeats): + sampling = SamplingConfig(max_new_tokens=steps, + temperature=0, + continuous_batch_size=request.prompt_ids.shape[0] if continuous else None) + samples = [] + for iteration in range(-1, repeats): + get_accelerator().synchronize() + start = time.perf_counter() + output = rollout.generate(request, sampling) + get_accelerator().synchronize() + if iteration >= 0: + samples.append((time.perf_counter() - start) * 1000) + return samples, output + + +def calibrate(engine, tokenizer, token, pad, device, repeats): + model = engine.module + config = model.config.get_text_config() + head_dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads + attention_scale = 4 * config.num_hidden_layers * config.num_attention_heads * head_dim / 1e9 + kv_scale = 2 * config.num_hidden_layers * config.num_key_value_heads * head_dim * next( + model.parameters()).element_size() / 2**20 + matrix, times, copy_rates, records = [], [], [], [] + # These calibration shapes differ from every held-out sweep workload below. + shapes = [(1, 32), (1, 256), (4, 64), (4, 384), (8, 768), (16, 96)] + for batch_size, width in shapes: + request = make_request([width] * batch_size, token, pad, device) + cfg = HybridEngineRolloutConfig(adaptive_prefill=True, + enable_profiling=True, + continuous_cache_capacity=width + 1, + prefill_fixed_cost_ms=1e9, + prefill_token_cost_ms=0, + prefill_attention_cost_ms=0, + prefill_kv_cost_ms=0) + rollout = HybridEngineRollout(engine, tokenizer, cfg) + profiles = [] + for iteration in range(-1, repeats): + rollout.generate(request, SamplingConfig(max_new_tokens=1, temperature=0, + continuous_batch_size=batch_size)) + if iteration >= 0: + profiles.append(rollout.get_last_profile()) + prefill = statistics.mean(p['prefill_forward_ms'] for p in profiles) + copying = statistics.mean(p['cache_management_overhead_ms'] for p in profiles) + assert all(p['num_prefill_forwards'] == 1 for p in profiles) + matrix.append([1, batch_size * width, attention_scale * batch_size * width**2]) + times.append(prefill) + copy_rates.append(copying / (kv_scale * batch_size * width)) + records.append({ + 'batch_size': batch_size, + 'width': width, + 'prefill_ms': prefill, + 'cache_management_ms': copying + }) + print('CALIBRATE', batch_size, width, round(prefill, 3), 'ms', flush=True) + fixed, linear, attention = nonnegative_fit(matrix, times) + # Differences across decode budgets cancel most prefill and allocation costs. + request = make_request([80] * 16, token, pad, device) + differences = {} + for steps in (1, 8): + cb_cfg = HybridEngineRolloutConfig(adaptive_prefill=True, + continuous_cache_capacity=80 + steps, + prefill_fixed_cost_ms=1e9, + prefill_token_cost_ms=0, + prefill_attention_cost_ms=0, + prefill_kv_cost_ms=0) + cb = HybridEngineRollout(engine, tokenizer, cb_cfg) + ordinary = HybridEngineRollout(engine, tokenizer) + cb_times, _ = measure(cb, request, steps, True, repeats) + ordinary_times, _ = measure(ordinary, request, steps, False, repeats) + differences[steps] = statistics.mean(cb_times) - statistics.mean(ordinary_times) + costs = { + 'prefill_fixed_cost_ms': fixed, + 'prefill_token_cost_ms': linear, + 'prefill_attention_cost_ms': attention, + 'prefill_kv_cost_ms': statistics.median(copy_rates), + 'continuous_decode_cost_ms': max(0, (differences[8] - differences[1]) / 7) + } + return costs, { + 'prefill_shapes': records, + 'decode_cb_minus_no_cb_ms': differences, + 'note': 'KV rate includes cache management overhead; coefficients are configuration-specific estimates.' + } + + +def sweep(engine, tokenizer, token, pad, device, steps, repeats, costs): + workloads = { + 'short128': [16] * 128, + 'moderate_tail': [128] + [4] * 31, + 'spread': [16 * (i + 1) for i in range(32)], + 'long_tail': [512] + [16] * 31 + } + report = {} + modes = ['no_cb', 'equal_width', 'aligned', 'auto'] + for name, lengths in workloads.items(): + request = make_request(lengths, token, pad, device) + cfgs = { + mode: + HybridEngineRolloutConfig(adaptive_prefill=mode == 'auto', + align_decode_fronts=mode == 'aligned', + **(costs if mode == 'auto' else {})) + for mode in modes + } + rollouts = {mode: HybridEngineRollout(engine, tokenizer, cfg) for mode, cfg in cfgs.items()} + samples = {mode: [] for mode in modes} + outputs = {} + for iteration in range(-1, repeats): + shift = iteration % len(modes) + for mode in modes[shift:] + modes[:shift]: + rollout = rollouts[mode] + sampling = SamplingConfig(max_new_tokens=steps, + temperature=0, + continuous_batch_size=None if mode == 'no_cb' else len(lengths)) + get_accelerator().synchronize() + start = time.perf_counter() + output = rollout.generate(request, sampling) + get_accelerator().synchronize() + if iteration >= 0: + samples[mode].append((time.perf_counter() - start) * 1000) + outputs[mode] = output.input_ids[:, max(lengths):].cpu().tolist() + means = {mode: statistics.mean(values) for mode, values in samples.items()} + best = min(modes[:-1], key=means.get) + stats = rollouts['auto'].get_last_continuous_stats() + report[name] = { + 'lengths': lengths, + 'max_new_tokens': steps, + 'mean_ms': means, + 'raw_ms': samples, + 'best_baseline': best, + 'auto_vs_best': means['auto'] / means[best], + 'auto_strategy': 'batched' if stats is None else 'bucketed', + 'output_ids_match': all(value == outputs['no_cb'] for value in outputs.values()), + 'output_match_by_path': { + mode: value == outputs['no_cb'] + for mode, value in outputs.items() + }, + 'outputs': outputs + } + rollouts['auto'].enable_profiling = True + rollouts['auto'].generate( + request, SamplingConfig(max_new_tokens=steps, temperature=0, continuous_batch_size=len(lengths))) + report[name]['auto_profile'] = rollouts['auto'].get_last_profile() + print('SWEEP', + name, + means, + 'auto/best', + round(report[name]['auto_vs_best'], 4), + report[name]['auto_strategy'], + flush=True) + return report + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--model', default='tiny-qwen2') + parser.add_argument('--dtype', choices=['float32', 'bfloat16'], default='float32') + parser.add_argument('--new-tokens', type=int, default=2) + parser.add_argument('--repeats', type=int, default=5) + parser.add_argument('--calibrate', + type=Path, + help='Write coefficients using separate calibration shapes, then exit.') + parser.add_argument('--cost-config', type=Path) + parser.add_argument('--output', type=Path, default=Path('prefill-sweep.json')) + args = parser.parse_args() + if args.repeats <= 0 or args.new_tokens <= 0: + parser.error('repeats and new-tokens must be positive') + torch.manual_seed(8497) + torch.set_num_threads(4) + device = get_accelerator().device_name() + dtype = getattr(torch, args.dtype) + if args.model == 'tiny-qwen2': + config = Qwen2Config(vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=4096, + pad_token_id=0, + eos_token_id=None, + attn_implementation='sdpa') + model = Qwen2ForCausalLM(config).to(device=device, dtype=dtype).eval() + token, pad = 3, 0 + else: + native_tokenizer = AutoTokenizer.from_pretrained(args.model) + token = native_tokenizer.encode('hello', add_special_tokens=False)[0] + pad = native_tokenizer.pad_token_id + if pad is None: + pad = native_tokenizer.eos_token_id + model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype, + attn_implementation='sdpa').to(device).eval() + tokenizer = SimpleNamespace(pad_token_id=pad, eos_token_id=None) + engine = SimpleNamespace(module=model) + source = Path(rollout_backend.__file__) + metadata = { + 'model': args.model, + 'model_config': json.loads(json.dumps(model.config.to_dict())), + 'dtype': args.dtype, + 'gpu': torch.cuda.get_device_name(0), #ignore-cuda + 'sm_count': torch.cuda.get_device_properties(0).multi_processor_count, #ignore-cuda + 'torch': torch.__version__, + 'transformers': __import__('transformers').__version__, + 'backend': 'sdpa', + 'seed': 8497, + 'cpu_threads': 4, + 'allow_tf32': torch.backends.cuda.matmul.allow_tf32, + 'source_sha256': hashlib.sha256(source.read_bytes()).hexdigest(), + 'scope': + 'Controlled EOS-disabled generation; includes routing and cache work; excludes loading and tokenization.' + } + with torch.inference_mode(): + if args.calibrate: + costs, measurements = calibrate(engine, tokenizer, token, pad, device, args.repeats) + args.calibrate.parent.mkdir(parents=True, exist_ok=True) + args.calibrate.write_text( + json.dumps({ + 'metadata': metadata, + 'costs': costs, + 'calibration': measurements + }, indent=2) + '\n') + print('COSTS', json.dumps(costs), flush=True) + return + calibration = json.loads(args.cost_config.read_text()) if args.cost_config else {} + if 'metadata' in calibration: + for key in ('model', 'model_config', 'dtype', 'gpu', 'sm_count', 'backend', 'torch', 'transformers', + 'allow_tf32'): + if calibration['metadata'][key] != metadata[key]: + parser.error(f'calibration differs in {key}; recalibrate for this configuration') + costs = calibration.get('costs', {}) + # A supplied JSON uses the same validated public configuration fields as an application. + cfg = HybridEngineRolloutConfig(adaptive_prefill=True, **costs) + report = { + 'metadata': metadata, + 'auto_config': asdict(cfg), + 'repeats': args.repeats, + 'workloads': sweep(engine, tokenizer, token, pad, device, args.new_tokens, args.repeats, costs) + } + worst = max(report['workloads'].items(), key=lambda pair: pair[1]['auto_vs_best']) + report['worst_case'] = { + 'workload': worst[0], + 'auto_vs_best': worst[1]['auto_vs_best'], + 'regression_percent': (worst[1]['auto_vs_best'] - 1) * 100 + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + '\n') + print('WORST', report['worst_case'], flush=True) + + +if __name__ == '__main__': + main() diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index a62ceaa7e124..3d451fa98892 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -139,12 +139,15 @@ def __init__(self, engine, tokenizer, cfg=None): 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 @@ -210,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: @@ -230,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 @@ -297,24 +307,34 @@ def _generate_continuous(self, request, sampling, max_batch_size): """ 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 + 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 + 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): - generation_config, _ = prepare_config(None) - defaults = GenerationConfig() + 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" @@ -326,10 +346,9 @@ def _generate_continuous(self, request, sampling, max_batch_size): if value != getattr(defaults, name, value): raise ValueError(f"adaptive prefill does not support non-default generation setting: {name}") if align_decode_fronts: - effective_lengths = request.prompt_attention_mask.sum(dim=1).tolist() prompt_lengths = dict(enumerate(map(int, effective_lengths))) else: - prompt_lengths = dict.fromkeys(range(len(requests)), request.prompt_ids.shape[1]) + 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: @@ -342,7 +361,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): else: 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] * len(requests), max_batch_size)) + 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: @@ -356,26 +375,29 @@ def _generate_continuous(self, request, sampling, max_batch_size): 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 (len(requests) <= max_batch_size and full_tokens_fit and full_positions_fit + 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) - _, bucket_cost = plan_prefill_buckets(list(prompt_lengths.values()), prefill_cost, - self.prefill_max_tokens) - batched_cost = prefill_cost(len(requests), full_width, include_kv=False) - continuous_cost = bucket_cost + self.continuous_decode_cost_ms * (sampling.max_new_tokens - 1) - use_batched = batched_cost <= continuous_cost + 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(request, replace(sampling, continuous_batch_size=None)) - pad_token_id = self.tokenizer.pad_token_id - if pad_token_id is None: - pad_token_id = self.tokenizer.eos_token_id - _, response_attention = self._pad_after_eos(output.input_ids, - request.prompt_ids.shape[1], - self.tokenizer.eos_token_id, - pad_token_id, - mask_padding=False) - output.attention_mask[:, request.prompt_ids.shape[1]:] = response_attention + 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 @@ -389,7 +411,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): 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"] = len(requests) + profile["active_batch_size"] = batch_size profile["continuous_batch_size"] = max_batch_size return output @@ -401,6 +423,9 @@ def _generate_continuous(self, request, sampling, max_batch_size): 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 @@ -661,8 +686,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") @@ -671,23 +696,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): diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst index 41b08728d15a..99a64c7b905e 100644 --- a/docs/code-docs/source/inference-engine.rst +++ b/docs/code-docs/source/inference-engine.rst @@ -159,6 +159,38 @@ 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. + +``benchmarks/rollout_prefill.py`` provides separate calibration and evaluation: + +.. code-block:: bash + + python benchmarks/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ + --repeats 3 --calibrate costs.json + python benchmarks/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``. + +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. diff --git a/tests/unit/runtime/rollout/test_continuous_batching.py b/tests/unit/runtime/rollout/test_continuous_batching.py index 2a8b0e2edaea..3f4c4729496f 100644 --- a/tests/unit/runtime/rollout/test_continuous_batching.py +++ b/tests/unit/runtime/rollout/test_continuous_batching.py @@ -122,3 +122,12 @@ def test_prefill_bucket_capacity_supports_long_lengths_without_token_expansion() 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) + + +def test_prefill_calibration_fit_recovers_nonnegative_costs(): + from benchmarks.rollout_prefill import nonnegative_fit + + matrix = [[1, 1, 1], [1, 2, 4], [1, 4, 16], [1, 8, 64]] + assert nonnegative_fit(matrix, [5.5, 10, 22, 58]) == pytest.approx([2, 3, 0.5]) + # Decreasing noisy measurements must not produce negative execution costs. + assert nonnegative_fit(matrix, [4, 3, 2, 1]) == pytest.approx([2.5, 0, 0]) diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 1b111e9ae38d..cbf2e1a038f7 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -1289,6 +1289,46 @@ def test_dp_prefill_rejects_single_prompt_above_forward_limit(): 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("force_cb", [False, True]) def test_adaptive_profile_includes_route_planning_time(force_cb): from deepspeed.accelerator import get_accelerator From 65a9115cdbee1a071d13236615067f404ba65b9e Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 00:45:41 +0800 Subject: [PATCH 3/7] fix(rollout): preserve resolved penalties and manual compatibility Signed-off-by: SuperGoodGame <985236470@qq.com> --- benchmarks/rollout_prefill.py | 305 ------------------ .../runtime/rollout/hybrid_engine_rollout.py | 30 +- docs/code-docs/source/inference-engine.rst | 13 +- .../rollout/test_continuous_batching.py | 9 - .../rollout/test_hybrid_engine_rollout.py | 53 +++ 5 files changed, 85 insertions(+), 325 deletions(-) delete mode 100644 benchmarks/rollout_prefill.py diff --git a/benchmarks/rollout_prefill.py b/benchmarks/rollout_prefill.py deleted file mode 100644 index b463529b1751..000000000000 --- a/benchmarks/rollout_prefill.py +++ /dev/null @@ -1,305 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -"""Calibrate prefill costs separately from a small, reproducible route sweep.""" - -import argparse -from dataclasses import asdict -import hashlib -from itertools import combinations -import json -from pathlib import Path -import statistics -import time -from types import SimpleNamespace - -import numpy as np -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config, Qwen2ForCausalLM - -from deepspeed.runtime.rollout.base import RolloutRequest, SamplingConfig -from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout, HybridEngineRolloutConfig -import deepspeed.runtime.rollout.hybrid_engine_rollout as rollout_backend -from deepspeed.accelerator import get_accelerator - - -def nonnegative_fit(matrix, targets): - """Solve three-parameter NNLS by checking its eight possible active sets.""" - matrix = np.asarray(matrix, dtype=float) - targets = np.asarray(targets, dtype=float) - scale = np.maximum(np.abs(matrix).max(axis=0), 1e-12) - scaled = matrix / scale - best = np.zeros(matrix.shape[1]) - error = float(targets @ targets) - for count in range(1, matrix.shape[1] + 1): - for columns in combinations(range(matrix.shape[1]), count): - values = np.linalg.lstsq(scaled[:, columns], targets, rcond=None)[0] - if np.any(values < 0): - continue - candidate = np.zeros(matrix.shape[1]) - candidate[list(columns)] = values - residual = scaled @ candidate - targets - if residual @ residual < error: - best, error = candidate, float(residual @ residual) - return (best / scale).tolist() - - -def make_request(lengths, token, pad, device): - width = max(lengths) - ids = torch.full((len(lengths), width), pad, dtype=torch.long, device=device) - mask = torch.zeros_like(ids) - for row, length in enumerate(lengths): - ids[row, -length:] = token - mask[row, -length:] = 1 - return RolloutRequest(ids, mask) - - -def measure(rollout, request, steps, continuous, repeats): - sampling = SamplingConfig(max_new_tokens=steps, - temperature=0, - continuous_batch_size=request.prompt_ids.shape[0] if continuous else None) - samples = [] - for iteration in range(-1, repeats): - get_accelerator().synchronize() - start = time.perf_counter() - output = rollout.generate(request, sampling) - get_accelerator().synchronize() - if iteration >= 0: - samples.append((time.perf_counter() - start) * 1000) - return samples, output - - -def calibrate(engine, tokenizer, token, pad, device, repeats): - model = engine.module - config = model.config.get_text_config() - head_dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads - attention_scale = 4 * config.num_hidden_layers * config.num_attention_heads * head_dim / 1e9 - kv_scale = 2 * config.num_hidden_layers * config.num_key_value_heads * head_dim * next( - model.parameters()).element_size() / 2**20 - matrix, times, copy_rates, records = [], [], [], [] - # These calibration shapes differ from every held-out sweep workload below. - shapes = [(1, 32), (1, 256), (4, 64), (4, 384), (8, 768), (16, 96)] - for batch_size, width in shapes: - request = make_request([width] * batch_size, token, pad, device) - cfg = HybridEngineRolloutConfig(adaptive_prefill=True, - enable_profiling=True, - continuous_cache_capacity=width + 1, - prefill_fixed_cost_ms=1e9, - prefill_token_cost_ms=0, - prefill_attention_cost_ms=0, - prefill_kv_cost_ms=0) - rollout = HybridEngineRollout(engine, tokenizer, cfg) - profiles = [] - for iteration in range(-1, repeats): - rollout.generate(request, SamplingConfig(max_new_tokens=1, temperature=0, - continuous_batch_size=batch_size)) - if iteration >= 0: - profiles.append(rollout.get_last_profile()) - prefill = statistics.mean(p['prefill_forward_ms'] for p in profiles) - copying = statistics.mean(p['cache_management_overhead_ms'] for p in profiles) - assert all(p['num_prefill_forwards'] == 1 for p in profiles) - matrix.append([1, batch_size * width, attention_scale * batch_size * width**2]) - times.append(prefill) - copy_rates.append(copying / (kv_scale * batch_size * width)) - records.append({ - 'batch_size': batch_size, - 'width': width, - 'prefill_ms': prefill, - 'cache_management_ms': copying - }) - print('CALIBRATE', batch_size, width, round(prefill, 3), 'ms', flush=True) - fixed, linear, attention = nonnegative_fit(matrix, times) - # Differences across decode budgets cancel most prefill and allocation costs. - request = make_request([80] * 16, token, pad, device) - differences = {} - for steps in (1, 8): - cb_cfg = HybridEngineRolloutConfig(adaptive_prefill=True, - continuous_cache_capacity=80 + steps, - prefill_fixed_cost_ms=1e9, - prefill_token_cost_ms=0, - prefill_attention_cost_ms=0, - prefill_kv_cost_ms=0) - cb = HybridEngineRollout(engine, tokenizer, cb_cfg) - ordinary = HybridEngineRollout(engine, tokenizer) - cb_times, _ = measure(cb, request, steps, True, repeats) - ordinary_times, _ = measure(ordinary, request, steps, False, repeats) - differences[steps] = statistics.mean(cb_times) - statistics.mean(ordinary_times) - costs = { - 'prefill_fixed_cost_ms': fixed, - 'prefill_token_cost_ms': linear, - 'prefill_attention_cost_ms': attention, - 'prefill_kv_cost_ms': statistics.median(copy_rates), - 'continuous_decode_cost_ms': max(0, (differences[8] - differences[1]) / 7) - } - return costs, { - 'prefill_shapes': records, - 'decode_cb_minus_no_cb_ms': differences, - 'note': 'KV rate includes cache management overhead; coefficients are configuration-specific estimates.' - } - - -def sweep(engine, tokenizer, token, pad, device, steps, repeats, costs): - workloads = { - 'short128': [16] * 128, - 'moderate_tail': [128] + [4] * 31, - 'spread': [16 * (i + 1) for i in range(32)], - 'long_tail': [512] + [16] * 31 - } - report = {} - modes = ['no_cb', 'equal_width', 'aligned', 'auto'] - for name, lengths in workloads.items(): - request = make_request(lengths, token, pad, device) - cfgs = { - mode: - HybridEngineRolloutConfig(adaptive_prefill=mode == 'auto', - align_decode_fronts=mode == 'aligned', - **(costs if mode == 'auto' else {})) - for mode in modes - } - rollouts = {mode: HybridEngineRollout(engine, tokenizer, cfg) for mode, cfg in cfgs.items()} - samples = {mode: [] for mode in modes} - outputs = {} - for iteration in range(-1, repeats): - shift = iteration % len(modes) - for mode in modes[shift:] + modes[:shift]: - rollout = rollouts[mode] - sampling = SamplingConfig(max_new_tokens=steps, - temperature=0, - continuous_batch_size=None if mode == 'no_cb' else len(lengths)) - get_accelerator().synchronize() - start = time.perf_counter() - output = rollout.generate(request, sampling) - get_accelerator().synchronize() - if iteration >= 0: - samples[mode].append((time.perf_counter() - start) * 1000) - outputs[mode] = output.input_ids[:, max(lengths):].cpu().tolist() - means = {mode: statistics.mean(values) for mode, values in samples.items()} - best = min(modes[:-1], key=means.get) - stats = rollouts['auto'].get_last_continuous_stats() - report[name] = { - 'lengths': lengths, - 'max_new_tokens': steps, - 'mean_ms': means, - 'raw_ms': samples, - 'best_baseline': best, - 'auto_vs_best': means['auto'] / means[best], - 'auto_strategy': 'batched' if stats is None else 'bucketed', - 'output_ids_match': all(value == outputs['no_cb'] for value in outputs.values()), - 'output_match_by_path': { - mode: value == outputs['no_cb'] - for mode, value in outputs.items() - }, - 'outputs': outputs - } - rollouts['auto'].enable_profiling = True - rollouts['auto'].generate( - request, SamplingConfig(max_new_tokens=steps, temperature=0, continuous_batch_size=len(lengths))) - report[name]['auto_profile'] = rollouts['auto'].get_last_profile() - print('SWEEP', - name, - means, - 'auto/best', - round(report[name]['auto_vs_best'], 4), - report[name]['auto_strategy'], - flush=True) - return report - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--model', default='tiny-qwen2') - parser.add_argument('--dtype', choices=['float32', 'bfloat16'], default='float32') - parser.add_argument('--new-tokens', type=int, default=2) - parser.add_argument('--repeats', type=int, default=5) - parser.add_argument('--calibrate', - type=Path, - help='Write coefficients using separate calibration shapes, then exit.') - parser.add_argument('--cost-config', type=Path) - parser.add_argument('--output', type=Path, default=Path('prefill-sweep.json')) - args = parser.parse_args() - if args.repeats <= 0 or args.new_tokens <= 0: - parser.error('repeats and new-tokens must be positive') - torch.manual_seed(8497) - torch.set_num_threads(4) - device = get_accelerator().device_name() - dtype = getattr(torch, args.dtype) - if args.model == 'tiny-qwen2': - config = Qwen2Config(vocab_size=128, - hidden_size=64, - intermediate_size=128, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=4096, - pad_token_id=0, - eos_token_id=None, - attn_implementation='sdpa') - model = Qwen2ForCausalLM(config).to(device=device, dtype=dtype).eval() - token, pad = 3, 0 - else: - native_tokenizer = AutoTokenizer.from_pretrained(args.model) - token = native_tokenizer.encode('hello', add_special_tokens=False)[0] - pad = native_tokenizer.pad_token_id - if pad is None: - pad = native_tokenizer.eos_token_id - model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype, - attn_implementation='sdpa').to(device).eval() - tokenizer = SimpleNamespace(pad_token_id=pad, eos_token_id=None) - engine = SimpleNamespace(module=model) - source = Path(rollout_backend.__file__) - metadata = { - 'model': args.model, - 'model_config': json.loads(json.dumps(model.config.to_dict())), - 'dtype': args.dtype, - 'gpu': torch.cuda.get_device_name(0), #ignore-cuda - 'sm_count': torch.cuda.get_device_properties(0).multi_processor_count, #ignore-cuda - 'torch': torch.__version__, - 'transformers': __import__('transformers').__version__, - 'backend': 'sdpa', - 'seed': 8497, - 'cpu_threads': 4, - 'allow_tf32': torch.backends.cuda.matmul.allow_tf32, - 'source_sha256': hashlib.sha256(source.read_bytes()).hexdigest(), - 'scope': - 'Controlled EOS-disabled generation; includes routing and cache work; excludes loading and tokenization.' - } - with torch.inference_mode(): - if args.calibrate: - costs, measurements = calibrate(engine, tokenizer, token, pad, device, args.repeats) - args.calibrate.parent.mkdir(parents=True, exist_ok=True) - args.calibrate.write_text( - json.dumps({ - 'metadata': metadata, - 'costs': costs, - 'calibration': measurements - }, indent=2) + '\n') - print('COSTS', json.dumps(costs), flush=True) - return - calibration = json.loads(args.cost_config.read_text()) if args.cost_config else {} - if 'metadata' in calibration: - for key in ('model', 'model_config', 'dtype', 'gpu', 'sm_count', 'backend', 'torch', 'transformers', - 'allow_tf32'): - if calibration['metadata'][key] != metadata[key]: - parser.error(f'calibration differs in {key}; recalibrate for this configuration') - costs = calibration.get('costs', {}) - # A supplied JSON uses the same validated public configuration fields as an application. - cfg = HybridEngineRolloutConfig(adaptive_prefill=True, **costs) - report = { - 'metadata': metadata, - 'auto_config': asdict(cfg), - 'repeats': args.repeats, - 'workloads': sweep(engine, tokenizer, token, pad, device, args.new_tokens, args.repeats, costs) - } - worst = max(report['workloads'].items(), key=lambda pair: pair[1]['auto_vs_best']) - report['worst_case'] = { - 'workload': worst[0], - 'auto_vs_best': worst[1]['auto_vs_best'], - 'regression_percent': (worst[1]['auto_vs_best'] - 1) * 100 - } - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(report, indent=2) + '\n') - print('WORST', report['worst_case'], flush=True) - - -if __name__ == '__main__': - main() diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index 3d451fa98892..2fe068254c21 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -564,6 +564,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): device, align_decode_fronts, profile, + generation_config=prepared_generation_config, ) if align_decode_fronts: for admitted, target_row in zip(update.admitted, update.admitted_slots): @@ -600,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))) @@ -730,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 {} @@ -748,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) @@ -796,7 +800,8 @@ def _continuous_prefill(self, prompt_attention, model_dtype, device, - prompt_requests=prompt_requests if self.adaptive_prefill else None) + prompt_requests=prompt_requests if self.adaptive_prefill else None, + generation_config=generation_config) self._profile_end(profile, "prefill_forward_ms", prefill_start, count="num_prefill_forwards") cache_copy_start = self._profile_start(profile) for source_row, row in enumerate(bucket): @@ -908,7 +913,8 @@ def _continuous_prefill_cache(self, prompt_attention, model_dtype, device, - prompt_requests=None): + 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 @@ -923,7 +929,9 @@ def _continuous_prefill_cache(self, eos_token_id=None, pad_token_id=pad_token_id, return_dict_in_generate=True, - output_logits=self.adaptive_prefill, + **({ + "output_logits": True + } if self.adaptive_prefill else {}), ) prefill_cache = prefill_output.past_key_values if not self.adaptive_prefill: @@ -950,13 +958,19 @@ def _continuous_prefill_cache(self, 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) + 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: diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst index 99a64c7b905e..807302bdc773 100644 --- a/docs/code-docs/source/inference-engine.rst +++ b/docs/code-docs/source/inference-engine.rst @@ -165,13 +165,15 @@ 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. -``benchmarks/rollout_prefill.py`` provides separate calibration and evaluation: +The `DeepSpeedExamples adaptive prefill benchmark +`_ +provides separate calibration and evaluation. From the DeepSpeedExamples checkout: .. code-block:: bash - python benchmarks/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ + python benchmarks/opsd/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ --repeats 3 --calibrate costs.json - python benchmarks/rollout_prefill.py --model Qwen/Qwen3-32B --dtype bfloat16 \ + 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. @@ -181,6 +183,11 @@ 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 diff --git a/tests/unit/runtime/rollout/test_continuous_batching.py b/tests/unit/runtime/rollout/test_continuous_batching.py index 3f4c4729496f..2a8b0e2edaea 100644 --- a/tests/unit/runtime/rollout/test_continuous_batching.py +++ b/tests/unit/runtime/rollout/test_continuous_batching.py @@ -122,12 +122,3 @@ def test_prefill_bucket_capacity_supports_long_lengths_without_token_expansion() 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) - - -def test_prefill_calibration_fit_recovers_nonnegative_costs(): - from benchmarks.rollout_prefill import nonnegative_fit - - matrix = [[1, 1, 1], [1, 2, 4], [1, 4, 16], [1, 8, 64]] - assert nonnegative_fit(matrix, [5.5, 10, 22, 58]) == pytest.approx([2, 3, 0.5]) - # Decreasing noisy measurements must not produce negative execution costs. - assert nonnegative_fit(matrix, [4, 3, 2, 1]) == pytest.approx([2.5, 0, 0]) diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index cbf2e1a038f7..130ccbc59350 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -1329,6 +1329,59 @@ def prepare(*args, **kwargs): SamplingConfig(max_new_tokens=2, temperature=0, continuous_batch_size=1)) +@pytest.mark.parametrize("cache_support", [True, None]) +@pytest.mark.parametrize("capacity", [1, 2]) +def test_adaptive_generation_uses_resolved_penalty_on_both_routes(cache_support, capacity): + from transformers import GenerationConfig + + model = _make_small_qwen() + 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 From ecacbce12c3f6749e4d80f8718a285618e6fd7e9 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 01:15:31 +0800 Subject: [PATCH 4/7] fix(rollout): expose static cache capability for MHA decoders Signed-off-by: SuperGoodGame <985236470@qq.com> --- deepspeed/utils/static_cache.py | 2 ++ .../rollout/test_hybrid_engine_rollout.py | 28 +++++++++++++++++++ 2 files changed, 30 insertions(+) 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/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 130ccbc59350..2c307c21de50 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -1102,6 +1102,34 @@ def test_adaptive_generation_matches_single_request_eager(lengths, capacity, str 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")), From ad3b247899aaece066c8b71fa47c15b5fca0620f Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 01:28:01 +0800 Subject: [PATCH 5/7] fix(rollout): apply repetition penalty in FP32 like HF Signed-off-by: SuperGoodGame <985236470@qq.com> --- .../runtime/rollout/hybrid_engine_rollout.py | 1 + .../rollout/test_hybrid_engine_rollout.py | 19 +++++++++++++++++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index 2fe068254c21..aeda93699a6f 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -976,6 +976,7 @@ def _continuous_next_tokens(logits, request_ids, request_by_id, responses, modul 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) diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 2c307c21de50..3886c4e406e8 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -580,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): @@ -1359,10 +1373,11 @@ def prepare(*args, **kwargs): @pytest.mark.parametrize("cache_support", [True, None]) @pytest.mark.parametrize("capacity", [1, 2]) -def test_adaptive_generation_uses_resolved_penalty_on_both_routes(cache_support, capacity): +@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() + 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) From d4def3db036f721a375d8582cc762d04262712b6 Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 01:46:29 +0800 Subject: [PATCH 6/7] fix(rollout): reject unknown generation switches Signed-off-by: SuperGoodGame <985236470@qq.com> --- .../runtime/rollout/hybrid_engine_rollout.py | 2 +- .../rollout/test_hybrid_engine_rollout.py | 31 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index aeda93699a6f..7533700413c3 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -343,7 +343,7 @@ def _generate_continuous(self, request, sampling, max_batch_size): for name, value in vars(generation_config).items(): if name.startswith("_") or name in supported_settings: continue - if value != getattr(defaults, name, value): + 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))) diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 3886c4e406e8..98b328988ef5 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -1190,6 +1190,37 @@ def test_adaptive_generation_rejects_unsupported_generation_settings(setting, va 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): From a0588b7dea72f2d23e8c90ca1daf51675633b4db Mon Sep 17 00:00:00 2001 From: SuperGoodGame <985236470@qq.com> Date: Sun, 11 Oct 2026 02:09:44 +0800 Subject: [PATCH 7/7] fix(rollout): retain padded history in manual aligned prefill Signed-off-by: SuperGoodGame <985236470@qq.com> --- .../runtime/rollout/hybrid_engine_rollout.py | 17 ++++++++--------- .../rollout/test_hybrid_engine_rollout.py | 10 ++++++---- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index 7533700413c3..56ce43035f80 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -793,15 +793,14 @@ def _continuous_prefill(self, 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, - prompt_requests=prompt_requests if self.adaptive_prefill else None, - generation_config=generation_config) + 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) for source_row, row in enumerate(bucket): diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 98b328988ef5..61b9e0d096f7 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -1279,9 +1279,10 @@ def test_adaptive_generation_applies_repetition_penalty_from_first_token(capacit assert torch.equal(result.input_ids[:, 5:], expected.unsqueeze(0).expand(2, -1)) -@pytest.mark.parametrize("cache_support", [True, None]) -@pytest.mark.parametrize("bucketed", [True, False]) -def test_adaptive_generation_preserves_padded_repetition_history(cache_support, bucketed): +@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 @@ -1303,7 +1304,8 @@ def test_adaptive_generation_preserves_padded_repetition_history(cache_support, ]) rollout = HybridEngineRollout( SimpleNamespace(module=model), SimpleNamespace(pad_token_id=11, eos_token_id=11), - HybridEngineRolloutConfig(adaptive_prefill=True, + 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,