feat(opsd): share prompt prefill across rollout samples

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

Author: nathon-lee

csrc/transformer/inference/csrc/pt_binding.cpp

@@ -478,12 +478,13 @@ std::vector<at::Tensor> 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>((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 <typename T>
+at::ScalarType workspace_scalar_type();
+
+template <>
+at::ScalarType workspace_scalar_type<float>()
+{
+    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 <typename T>
+std::vector<at::Tensor> 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<T>())
+                       .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<at::Tensor> 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 <typename T>
 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)");

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;

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

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 \

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)
     # ------------------------------------------------------------------

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.