diff --git a/csrc/transformer/inference/csrc/pt_binding.cpp b/csrc/transformer/inference/csrc/pt_binding.cpp index 19dbe73726f7..988cec016b53 100644 --- a/csrc/transformer/inference/csrc/pt_binding.cpp +++ b/csrc/transformer/inference/csrc/pt_binding.cpp @@ -478,12 +478,13 @@ std::vector ds_softmax_context(at::Tensor& query_key_value, auto output = torch::from_blob(workspace + 4 * buf_size, {bsz, seq_len, hidden_dim}, options); auto query_cont = workspace + 5 * buf_size; + unsigned cache_bsz = InferenceContext::Instance().GetBatchSize(); size_t offset = - 10 * (hidden_dim * bsz * InferenceContext::Instance().GetMaxTokenLength()) + - layer_id * 2 * bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; + 10 * (hidden_dim * cache_bsz * InferenceContext::Instance().GetMaxTokenLength()) + + layer_id * 2 * cache_bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; unsigned all_tokens = soft_len; auto kv_cache = workspace + offset + (hidden_dim / heads) * (is_prompt ? 0 : soft_len - 1); - size_t value_offset = bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; + size_t value_offset = cache_bsz * InferenceContext::Instance().GetMaxTokenLength() * hidden_dim; T* temp_buf = (T*)output.data_ptr() + at::numel(output); launch_bias_add_transform_0213((T*)query_cont, @@ -1943,6 +1944,83 @@ void ds_release_workspace() { InferenceContext::Instance().release_workspace(); bool ds_retake_workspace() { return InferenceContext::Instance().retake_workspace(); } +template +at::ScalarType workspace_scalar_type(); + +template <> +at::ScalarType workspace_scalar_type() +{ + return torch::kFloat32; +} + +template <> +at::ScalarType workspace_scalar_type<__half>() +{ + return torch::kFloat16; +} + +#ifdef BF16_AVAILABLE +template <> +at::ScalarType workspace_scalar_type<__nv_bfloat16>() +{ + return torch::kBFloat16; +} +#endif + +template +std::vector repeat_kv_cache(unsigned source_batch_size, unsigned repeats) +{ + auto& context = InferenceContext::Instance(); + const auto target_batch_size = source_batch_size * repeats; + if (repeats < 1 || source_batch_size < 1 || target_batch_size != context.GetBatchSize()) { + throw std::runtime_error( + "KV cache repeat does not match the allocated workspace batch size"); + } + + const auto num_layers = context.GetNumLayers(); + const auto num_heads = context.GetNumHeads(); + const auto max_tokens = context.GetMaxTokenLength(); + const auto hidden_dim = context.GetHiddenDim(); + const auto head_dim = hidden_dim / num_heads; + const auto current_tokens = context.current_tokens(); + if (current_tokens <= 1) { + throw std::runtime_error("KV cache repeat requires a completed prompt forward"); + } + const auto prompt_tokens = current_tokens - 1; + auto options = at::TensorOptions() + .dtype(workspace_scalar_type()) + .layout(at::kStrided) + .device(at::kCUDA) + .requires_grad(false); + T* workspace = (T*)context.GetWorkSpace(); + const auto cache_offset = 10 * hidden_dim * target_batch_size * max_tokens; + auto cache = torch::from_blob(workspace + cache_offset, + {(long)num_layers, + 2, + (long)target_batch_size, + (long)num_heads, + (long)max_tokens, + (long)head_dim}, + options); + // Backward copies preserve source rows that overlap the expanded destination range. + for (unsigned destination = target_batch_size; destination-- > 0;) { + const auto source = destination / repeats; + if (source == destination) { continue; } + auto destination_cache = cache.select(2, destination).slice(3, 0, prompt_tokens); + auto source_cache = cache.select(2, source).slice(3, 0, prompt_tokens); + destination_cache.copy_(source_cache); + } + + std::vector repeated_cache; + repeated_cache.reserve(num_layers * 2); + for (unsigned layer = 0; layer < num_layers; layer++) { + auto layer_cache = cache.select(0, layer); + repeated_cache.push_back(layer_cache.select(0, 0).slice(2, 0, prompt_tokens)); + repeated_cache.push_back(layer_cache.select(0, 1).slice(2, 0, prompt_tokens)); + } + return repeated_cache; +} + template at::Tensor ds_dequantize(at::Tensor& weight, at::Tensor& qscale, int groups) { @@ -2032,6 +2110,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("allocate_workspace_" #_name, \ &allocate_workspace<_dtype>, \ "DeepSpeed memory allocation for GPT inference with " #_name " (CUDA)"); \ + m.def("repeat_kv_cache_" #_name, \ + &repeat_kv_cache<_dtype>, \ + "Repeat prompt KV cache entries across the inference batch with " #_name " (CUDA)"); \ m.def("dequantize_" #_name, \ &ds_dequantize<_dtype>, \ "DeepSpeed dequantize with " #_name " (CUDA)"); diff --git a/csrc/transformer/inference/includes/inference_context.h b/csrc/transformer/inference/includes/inference_context.h index 378fd4e5e990..5190ff50f516 100644 --- a/csrc/transformer/inference/includes/inference_context.h +++ b/csrc/transformer/inference/includes/inference_context.h @@ -60,6 +60,10 @@ class InferenceContext { { _workSpaceSize = 0; _workspace = 0; + _batch_size = 0; + _num_layers = 0; + _num_heads = 0; + _hidden_dim = 0; cublasStatus_t stat = cublasCreate(&_cublasHandle); if (stat != CUBLAS_STATUS_SUCCESS) { @@ -108,6 +112,10 @@ class InferenceContext { unsigned min_out_tokens) { size_t total_size; + _batch_size = batch_size; + _num_layers = num_layers; + _num_heads = num_heads; + _hidden_dim = hidden_dim; if (!_free_memory_size) { cudaMemGetInfo(&_free_memory_size, &total_size); } // Flash attention requires padded heads and we'll conservatively allocate @@ -181,6 +189,10 @@ class InferenceContext { _attention_unfused_workspace_offset = workSpaceSize - temp_size; } inline size_t GetMaxTokenLength() const { return _max_seq_len; } + inline size_t GetBatchSize() const { return _batch_size; } + inline unsigned GetNumLayers() const { return _num_layers; } + inline unsigned GetNumHeads() const { return _num_heads; } + inline size_t GetHiddenDim() const { return _hidden_dim; } cudaEvent_t GetCompEvent(int id) { return id == 1 ? _comp1_event : _comp2_event; } @@ -275,6 +287,10 @@ class InferenceContext { size_t _free_memory_size; size_t _max_seq_len; + size_t _batch_size; + unsigned _num_layers; + unsigned _num_heads; + size_t _hidden_dim; cudaEvent_t _comp1_event; cudaEvent_t _comp2_event; diff --git a/deepspeed/ops/transformer/inference/op_binding/workspace.py b/deepspeed/ops/transformer/inference/op_binding/workspace.py index 19de7d9576af..0565666559bf 100644 --- a/deepspeed/ops/transformer/inference/op_binding/workspace.py +++ b/deepspeed/ops/transformer/inference/op_binding/workspace.py @@ -158,10 +158,14 @@ def __init__(self, config: DeepSpeedInferenceConfig = None): super(WorkspaceOp, self).__init__(config) if config.dtype == torch.float32: self.allocate_workspace_func = self.inference_module.allocate_workspace_fp32 + repeat_kv_cache_name = "repeat_kv_cache_fp32" elif config.dtype == torch.bfloat16: self.allocate_workspace_func = self.inference_module.allocate_workspace_bf16 + repeat_kv_cache_name = "repeat_kv_cache_bf16" else: self.allocate_workspace_func = self.inference_module.allocate_workspace_fp16 + repeat_kv_cache_name = "repeat_kv_cache_fp16" + self.repeat_kv_cache_func = getattr(self.inference_module, repeat_kv_cache_name, None) self.release_workspace_func = self.inference_module.release_workspace self.retake_workspace_func = self.inference_module.retake_workspace self.reset_cache_func = self.inference_module.reset_cache @@ -176,6 +180,7 @@ def __init__(self, config: DeepSpeedInferenceConfig = None): self.release_workspace_func = self.release_workspace_fallback self.retake_workspace_func = self.retake_workspace_fallback self.reset_cache_func = self.reset_cache_fallback + self.repeat_kv_cache_func = self.repeat_kv_cache_fallback def allocate_workspace(self, *args, **kwargs): self._is_allocated = True @@ -191,6 +196,11 @@ def reset_cache(self): def retake_workspace(self): return self.retake_workspace_func() if self.retake_workspace_func else None + def repeat_kv_cache(self, source_batch_size, repeats): + if self.repeat_kv_cache_func is None: + raise RuntimeError("Shared prefill requires rebuilding the transformer inference extension") + return self.repeat_kv_cache_func(source_batch_size, repeats) + def allocate_workspace_fp32_fallback(self, hidden_dim, num_heads, prompt_length, batch_size, num_layers, mp_size, external_cache, rank, max_out_tokens, min_out_tokens): return self.inference_context.gen_workspace(num_layers, num_heads, batch_size, prompt_length, hidden_dim, @@ -218,5 +228,28 @@ def release_workspace_fallback(self): def retake_workspace_fallback(self): return self.inference_context.retake_workspace() + def repeat_kv_cache_fallback(self, source_batch_size, repeats): + target_batch_size = source_batch_size * repeats + cache_size = self.inference_context.kv_cache_size + if cache_size is None or cache_size[0] != target_batch_size: + raise RuntimeError("KV cache repeat does not match the allocated workspace batch size") + if self.inference_context.kv_cache is None: + raise RuntimeError("KV cache repeat requires a completed prompt forward") + current_tokens = self.inference_context.current_tokens() + if current_tokens <= 1: + raise RuntimeError("KV cache repeat requires a completed prompt forward") + prompt_tokens = current_tokens - 1 + repeated_cache = [] + for key_cache, value_cache in self.inference_context.kv_cache: + # Backward copies preserve source rows that overlap the expanded destination range. + for destination in range(target_batch_size - 1, -1, -1): + source = destination // repeats + if source == destination: + continue + key_cache[destination, :, :prompt_tokens, :].copy_(key_cache[source, :, :prompt_tokens, :]) + value_cache[destination, :, :prompt_tokens, :].copy_(value_cache[source, :, :prompt_tokens, :]) + repeated_cache.extend((key_cache[:, :, :prompt_tokens, :], value_cache[:, :, :prompt_tokens, :])) + return repeated_cache + def is_allocated(self): return self._is_allocated diff --git a/deepspeed/runtime/hybrid_engine.py b/deepspeed/runtime/hybrid_engine.py index 6565b88c0b0f..b5f81a801403 100644 --- a/deepspeed/runtime/hybrid_engine.py +++ b/deepspeed/runtime/hybrid_engine.py @@ -177,6 +177,46 @@ def retake_inference_cache(self): if not retake_success: raise RuntimeError("Unable to retake inference workspace.") + def prepare_shared_prefill(self, source_batch_size, repeats, prompt_length): + """Allocate a target-batch workspace before a shared prompt forward.""" + hybrid_config = self._config.hybrid_engine + if self.Z3_enabled: + raise RuntimeError("Shared prefill does not support ZeRO stage 3") + if hybrid_config.inference_tp_size != 1: + raise RuntimeError("Shared prefill does not support inference tensor parallelism") + if hybrid_config.release_inference_cache: + raise RuntimeError("Shared prefill does not support release_inference_cache") + if hybrid_config.enable_cuda_graph: + raise RuntimeError("Shared prefill does not support CUDA graph capture") + if len(self._inference_containers) == 0: + raise RuntimeError("Shared prefill requires HybridEngine inference containers") + + target_batch_size = source_batch_size * repeats + inference_module = self._inference_containers[0].module + config = inference_module.config + if config.bigscience_bloom: + raise RuntimeError("Shared prefill does not support external KV caches") + inference_module.workspace.allocate_workspace( + config.hidden_size, + config.heads, + prompt_length, + target_batch_size, + len(self._inference_containers), + config.mp_size, + config.bigscience_bloom, + dist.get_rank() if dist.is_initialized() else 0, + config.max_out_tokens, + config.min_out_tokens, + ) + for container in self._inference_containers: + container.module._should_allocate_workspace = False + self._shared_prefill_workspace = inference_module.workspace + + def repeat_shared_prefill_cache(self, source_batch_size, repeats): + """Expand the completed prompt cache for independent response branches.""" + cache_tensors = self._shared_prefill_workspace.repeat_kv_cache(source_batch_size, repeats) + return tuple(zip(cache_tensors[::2], cache_tensors[1::2])) + def generate(self, *inputs, **kwargs): if self._total_batch_size is None: bsz = inputs[0].shape[0] if len(inputs) > 0 else \ diff --git a/deepspeed/runtime/rollout/hybrid_engine_rollout.py b/deepspeed/runtime/rollout/hybrid_engine_rollout.py index 3b51e643f9e5..1249e0f68d43 100644 --- a/deepspeed/runtime/rollout/hybrid_engine_rollout.py +++ b/deepspeed/runtime/rollout/hybrid_engine_rollout.py @@ -26,6 +26,7 @@ class HybridEngineRolloutConfig: """Configuration for HybridEngineRollout.""" use_graph_capture: bool = False enable_profiling: bool = False + use_shared_prefill: bool = False class HybridEngineRollout(RolloutEngine): @@ -42,6 +43,7 @@ def __init__(self, engine, tokenizer, cfg=None): self.tokenizer = tokenizer self.use_graph_capture = getattr(cfg, 'use_graph_capture', False) if cfg else False 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._last_profile = None @torch.no_grad() @@ -77,20 +79,31 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout is_greedy = sampling.temperature <= 0.0 - if self.use_graph_capture and is_greedy: - output_ids = self._generate_graph(prompt_ids, prompt_attn, max_new_tokens, pad_token_id, module, device) - else: - temperature = max(sampling.temperature, 1e-8) - do_sample = not is_greedy - output_ids = module.generate( - prompt_ids, - attention_mask=prompt_attn, - max_new_tokens=max_new_tokens, - do_sample=do_sample, - temperature=temperature if do_sample else 1.0, - top_p=sampling.top_p if do_sample else 1.0, - pad_token_id=pad_token_id, - ) + shared_prefill_handles = [] + if self.use_shared_prefill and n > 1: + if self.use_graph_capture: + raise RuntimeError("Shared prefill does not support CUDA graph capture") + self.engine.prepare_shared_prefill(B, n, prompt_len) + shared_prefill_handles = self._register_shared_prefill_hooks(module, B, n) + try: + if self.use_graph_capture and is_greedy: + output_ids = self._generate_graph(prompt_ids, prompt_attn, max_new_tokens, pad_token_id, module, + device) + else: + temperature = max(sampling.temperature, 1e-8) + do_sample = not is_greedy + output_ids = module.generate( + prompt_ids, + attention_mask=prompt_attn, + max_new_tokens=max_new_tokens, + do_sample=do_sample, + temperature=temperature if do_sample else 1.0, + top_p=sampling.top_p if do_sample else 1.0, + pad_token_id=pad_token_id, + ) + finally: + for handle in shared_prefill_handles: + handle.remove() if self.enable_profiling: accelerator.synchronize() @@ -141,6 +154,43 @@ def get_last_profile(self): """Return the most recent profiling snapshot for this rollout instance.""" return self._last_profile + def _register_shared_prefill_hooks(self, module, batch_size, repeats): + state = {"pending": True, "reduced": False} + + def reduce_prompt_batch(_module, args, kwargs): + input_ids = kwargs.get("input_ids") + if not state["pending"]: + return args, kwargs + if input_ids is None: + raise RuntimeError("Shared prefill requires input_ids as a keyword argument") + expected_batch_size = batch_size * repeats + if input_ids.shape[0] != expected_batch_size: + raise RuntimeError("Shared prefill input batch does not match the expanded rollout batch") + if input_ids.shape[1] <= 1: + raise RuntimeError("Shared prefill requires a prompt with more than one token") + kwargs = dict(kwargs) + kwargs["input_ids"] = input_ids[::repeats] + for name in ("attention_mask", "position_ids", "token_type_ids"): + value = kwargs.get(name) + if isinstance(value, torch.Tensor) and value.shape[0] == expected_batch_size: + kwargs[name] = value[::repeats] + state["reduced"] = True + return args, kwargs + + def expand_prompt_output(_module, _args, _kwargs, output): + if not state["pending"]: + return output + if not state["reduced"]: + raise RuntimeError("Shared prefill did not reduce the prompt batch") + state["pending"] = False + output.past_key_values = self.engine.repeat_shared_prefill_cache(batch_size, repeats) + output.logits = output.logits.repeat_interleave(repeats, dim=0) + return output + + pre_handle = module.register_forward_pre_hook(reduce_prompt_batch, with_kwargs=True) + post_handle = module.register_forward_hook(expand_prompt_output, with_kwargs=True) + return pre_handle, post_handle + # ------------------------------------------------------------------ # Graph capture decode loop (greedy only) # ------------------------------------------------------------------ diff --git a/docs/code-docs/source/inference-engine.rst b/docs/code-docs/source/inference-engine.rst index 23bd9a5db63b..a1fae868d0c6 100644 --- a/docs/code-docs/source/inference-engine.rst +++ b/docs/code-docs/source/inference-engine.rst @@ -43,3 +43,17 @@ batch size, samples per prompt, prompt length, and returned response length. For benchmark matrices, cases execute from the largest effective batch to the smallest because HybridEngine sizes its inference workspace on the first forward. Results remain in the user-requested matrix order. + +Shared Prompt Prefill +--------------------- + +When one prompt branches into multiple response samples, +``HybridEngineRolloutConfig(use_shared_prefill=True)`` computes the prompt +forward once and repeats its KV cache before decoding the independent response +branches. The option is disabled by default. + +Shared prefill currently requires HybridEngine kernel injection, ZeRO stage 0, +inference tensor-parallel size 1, an internal KV cache, and a prompt longer than +one token. It cannot be combined with CUDA graph capture or +``release_inference_cache``. Sampling still happens independently for every +response branch after the shared prompt forward. diff --git a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py index 116f7560fc26..9f34b3e134e4 100644 --- a/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py +++ b/tests/unit/runtime/rollout/test_hybrid_engine_rollout.py @@ -2,16 +2,20 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team -"""CPU-only unit tests for HybridEngineRollout (no GPU needed). +"""Unit tests for HybridEngineRollout. -Tests cover configuration, profiling, generation dispatch, and the pure-tensor sampling helper. +Most tests are CPU-only; the native shared-prefill cache test runs only when CUDA and +the transformer inference extension are available. """ +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest import torch +from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp +from deepspeed.runtime.hybrid_engine import DeepSpeedHybridEngine from deepspeed.runtime.rollout.base import RolloutRequest, SamplingConfig from deepspeed.runtime.rollout.hybrid_engine_rollout import ( HybridEngineRollout, @@ -51,6 +55,7 @@ def test_config_defaults(): cfg = HybridEngineRolloutConfig() assert cfg.use_graph_capture is False assert cfg.enable_profiling is False + assert cfg.use_shared_prefill is False # -- constructor -------------------------------------------------------- @@ -63,6 +68,7 @@ def test_constructor_stores_config(): rollout = HybridEngineRollout(engine, tok, cfg=cfg) assert rollout.use_graph_capture is True assert rollout.enable_profiling is True + assert rollout.use_shared_prefill is False assert rollout.engine is engine assert rollout.tokenizer is tok @@ -71,6 +77,7 @@ def test_constructor_defaults_without_cfg(): rollout = HybridEngineRollout(_make_engine(), _make_tokenizer()) assert rollout.use_graph_capture is False assert rollout.enable_profiling is False + assert rollout.use_shared_prefill is False @patch("deepspeed.runtime.rollout.hybrid_engine_rollout.time.perf_counter") @@ -150,6 +157,207 @@ def test_generate_preserves_zero_pad_token_id(): assert output.attention_mask[:, -1].tolist() == [0, 0] +def test_native_repeat_kv_cache_fp16_reverse_copy(): + """Exercise the native reverse copy with multiple source cache rows.""" + if not torch.cuda.is_available(): #ignore-cuda + pytest.skip("CUDA is required for the native inference kernel") + + from deepspeed.ops.op_builder import InferenceBuilder + + builder = InferenceBuilder() + try: + is_compatible = builder.is_compatible() + except Exception as exc: + pytest.skip(f"Unable to inspect native transformer inference compatibility: {exc}") + if not is_compatible: + pytest.skip("The native transformer inference extension is not compatible") + try: + inference_op = builder.load() + except Exception as exc: + pytest.skip(f"Unable to load the native transformer inference extension: {exc}") + + repeat_kv_cache = getattr(inference_op, "repeat_kv_cache_fp16", None) + if repeat_kv_cache is None: + pytest.skip("The native transformer inference extension lacks repeat_kv_cache_fp16") + + device = torch.device("cuda") + source_batch_size = 2 + repeats = 2 + target_batch_size = source_batch_size * repeats + prompt_length = 2 + hidden_dim = 8 # The FP16 transform kernel processes eight values per thread. + num_heads = 1 + + inference_op.allocate_workspace_fp16( + hidden_dim, + num_heads, + prompt_length, + target_batch_size, + 1, + 1, + False, + 0, + 4, + 1, + ) + try: + query_key_value = torch.zeros((source_batch_size, prompt_length, hidden_dim * 3), + dtype=torch.float16, + device=device) + query_key_value = query_key_value.view(source_batch_size, prompt_length, 3, num_heads, hidden_dim) + query_key_value[0, :, 1, :, :] = 1 + query_key_value[0, :, 2, :, :] = 10 + query_key_value[1, :, 1, :, :] = 3 + query_key_value[1, :, 2, :, :] = 30 + query_key_value = query_key_value.reshape(source_batch_size, prompt_length, hidden_dim * 3) + + empty_mask = torch.empty(1, dtype=torch.float16, device=device) + inference_op.softmax_context_fp16( + query_key_value, + empty_mask, + 0, + False, + False, + num_heads, + 0, + 1.0, + False, + False, + 1, + True, + 0, + 1, + empty_mask, + 1.0, + True, + None, + None, + ) + torch.cuda.synchronize() #ignore-cuda + + repeated_cache = repeat_kv_cache(source_batch_size, repeats) + torch.cuda.synchronize() #ignore-cuda + + assert len(repeated_cache) == 2 + expected_key = torch.empty((target_batch_size, 1, prompt_length, hidden_dim), + dtype=torch.float16, + device=device) + expected_key[:source_batch_size] = 1 + expected_key[source_batch_size:] = 3 + expected_value = expected_key * 10 + assert torch.equal(repeated_cache[0], expected_key) + assert torch.equal(repeated_cache[1], expected_value) + finally: + inference_op.release_workspace() + + +def test_shared_prefill_hooks_reduce_prompt_and_expand_output(): + + class PromptModule(torch.nn.Module): + + def __init__(self): + super().__init__() + self.forward_batch_sizes = [] + + def forward(self, input_ids, attention_mask=None): + self.forward_batch_sizes.append(input_ids.shape[0]) + values = input_ids[:, None, :, None].float() + return SimpleNamespace(logits=input_ids[:, :, None].float(), past_key_values=((values, values), )) + + engine = _make_engine() + engine.repeat_shared_prefill_cache.return_value = ((torch.zeros(4, 1, 2, 1), torch.zeros(4, 1, 2, 1)), ) + module = PromptModule() + rollout = HybridEngineRollout(engine, _make_tokenizer()) + handles = rollout._register_shared_prefill_hooks(module, batch_size=2, repeats=2) + prompt_ids = torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4]]) + + output = module(input_ids=prompt_ids, attention_mask=torch.ones_like(prompt_ids)) + decode_output = module(input_ids=torch.ones(4, 1, dtype=torch.long)) + for handle in handles: + handle.remove() + + assert module.forward_batch_sizes == [2, 4] + assert output.logits.shape[0] == 4 + assert output.past_key_values[0][0].shape[0] == 4 + assert decode_output.logits.shape[0] == 4 + engine.repeat_shared_prefill_cache.assert_called_once_with(2, 2) + + +def test_generate_uses_shared_prefill_for_multiple_samples(): + + class GenerateModule(torch.nn.Module): + + def __init__(self): + super().__init__() + self.forward_batch_sizes = [] + + def forward(self, input_ids, attention_mask=None): + self.forward_batch_sizes.append(input_ids.shape[0]) + values = input_ids[:, None, :, None].float() + return SimpleNamespace(logits=input_ids[:, :, None].float(), past_key_values=((values, values), )) + + def generate(self, input_ids, attention_mask=None, max_new_tokens=1, **_kwargs): + self(input_ids=input_ids, attention_mask=attention_mask) + response = torch.ones(input_ids.shape[0], max_new_tokens, dtype=input_ids.dtype) + return torch.cat((input_ids, response), dim=1) + + engine = _make_engine() + engine.module = GenerateModule() + config = HybridEngineRolloutConfig(use_shared_prefill=True) + rollout = HybridEngineRollout(engine, _make_tokenizer(), config) + + output = rollout.generate(_make_request(), _make_sampling(n_samples_per_prompt=2)) + + assert engine.module.forward_batch_sizes == [2] + assert output.input_ids[:, 3:].shape == (4, 2) + engine.prepare_shared_prefill.assert_called_once_with(2, 2, 3) + engine.repeat_shared_prefill_cache.assert_called_once_with(2, 2) + + +def test_shared_prefill_rejects_graph_capture(): + engine = _make_engine() + config = HybridEngineRolloutConfig(use_graph_capture=True, use_shared_prefill=True) + rollout = HybridEngineRollout(engine, _make_tokenizer(), config) + + with pytest.raises(RuntimeError, match="does not support CUDA graph capture"): + rollout.generate(_make_request(), _make_sampling(n_samples_per_prompt=2)) + + +def test_shared_prefill_fallback_repeats_prompt_cache(): + key_cache = torch.zeros(4, 1, 3, 1) + value_cache = torch.zeros_like(key_cache) + key_cache[:2, :, :2, :] = torch.tensor([[[[1.0], [2.0]]], [[[3.0], [4.0]]]]) + value_cache[:2, :, :2, :] = key_cache[:2, :, :2, :] + 10 + workspace = WorkspaceOp.__new__(WorkspaceOp) + workspace.inference_context = SimpleNamespace( + kv_cache_size=key_cache.shape, + kv_cache=[(key_cache, value_cache)], + current_tokens=lambda: 3, + ) + + repeated_cache = workspace.repeat_kv_cache_fallback(source_batch_size=2, repeats=2) + + expected_key = torch.tensor([1.0, 1.0, 3.0, 3.0]) + assert torch.equal(key_cache[:, 0, 0, 0], expected_key) + assert torch.equal(value_cache[:, 0, 0, 0], expected_key + 10) + assert repeated_cache[0].shape == (4, 1, 2, 1) + assert repeated_cache[1].shape == (4, 1, 2, 1) + + +def test_engine_pairs_shared_prefill_cache_tensors(): + key_cache = torch.zeros(4, 1, 2, 1) + value_cache = torch.ones_like(key_cache) + workspace = MagicMock() + workspace.repeat_kv_cache.return_value = [key_cache, value_cache] + engine = SimpleNamespace(_shared_prefill_workspace=workspace) + + repeated_cache = DeepSpeedHybridEngine.repeat_shared_prefill_cache(engine, 2, 2) + + assert repeated_cache[0][0] is key_cache + assert repeated_cache[0][1] is value_cache + workspace.repeat_kv_cache.assert_called_once_with(2, 2) + + # -- _sample_top_p ------------------------------------------------------