feat(opsd): share prompt prefill across rollout samples

[nathon-lee]

Summary

This PR adds an opt‑in shared‑prefill path for HybridEngineRollout.

When one prompt branches into multiple response samples, the current rollout
path expands the prompt batch before generation and computes the same prompt
prefill once for every response branch.

With shared prefill enabled, HybridEngine now:

1. computes each unique prompt prefill once;
2. expands the resulting KV cache to the response‑sample batch;
3. continues decoding each response branch independently.

For the tested OPT‑6.7B workload with four samples per prompt, this reduced
end‑to‑end rollout latency by 23.77% while preserving identical generated
tokens.

Motivation

For a rollout with:

batch_size = B
samples_per_prompt = N

the default path expands the prompt batch to B * N before the first model
forward.

This repeats the same prompt prefill N times even though all response branches
have identical prompt tokens and prompt KV state.

Profiling showed that this redundant prefill work is a significant part of the
total latency when generating multiple responses per prompt.

Implementation

Shared prompt forward

This PR adds the default‑off configuration option:

HybridEngineRolloutConfig(use_shared_prefill=True)

For samples_per_prompt > 1, the first model forward is reduced from the
expanded B * N prompt batch to the original B unique prompts.

The first‑forward logits are expanded back to B * N before sampling.
Sampling and all subsequent decode forwards remain independent for every
response branch.

Native KV‑cache expansion

The transformer inference extension now records the allocated workspace
metadata required to address the KV cache independently from the current
forward batch:

• allocated batch size
• number of layers
• number of attention heads
• hidden dimension

The prompt KV cache is exposed using its native layout:

[layer, key/value, batch, heads, tokens, head_dim]

The cache is expanded from B to B * N in place.

Copies proceed from the destination batch in reverse order. This preserves
source rows that overlap the expanded destination range and avoids allocating
a full temporary clone of the prompt KV cache.

The returned key and value tensors are zero‑copy views into the HybridEngine
inference workspace.

Fallback implementation

The Python fallback workspace implements the same cache‑expansion ordering and
returned‑cache contract for environments that do not use the native transformer
inference extension.

Current restrictions

Shared prefill is disabled by default and currently requires:

• HybridEngine kernel injection
• ZeRO stage 0
• inference tensor‑parallel size 1
• the internal HybridEngine KV cache
• prompt length greater than one token

It currently cannot be combined with:

• ZeRO stage 3
• inference tensor parallelism
• external KV caches
• CUDA graph capture
• release_inference_cache

Unsupported combinations fail explicitly instead of silently falling back to
an incorrect execution path.

Performance validation

The following A/B measurements were collected before the benchmark was moved
to DeepSpeedExamples. The benchmark itself is not included in this PR.

Environment

• GPU: NVIDIA RTX A4500
• GPU memory: 20,470 MiB
• Driver: 580.159.04
• PyTorch: 2.9.1+cu128
• CUDA runtime: 12.8
• Transformers: 4.40.2
• Model: facebook/opt‑6.7b
• Dtype: FP16
• Batch size: 1
• Samples per prompt: 4
• Prompt length: 512
• Response length: 32
• Warmup iterations: 5
• Measured iterations: 20

Results

Metric	Expanded prefill	Shared prefill	Change
Prefill latency	443.98 ms	132.78 ms	-70.1%
End‑to‑end latency	1327.07 ms	1011.62 ms	-23.77%
Throughput	96.45 tokens/s	126.53 tokens/s	+31.2%
Peak memory	13,144.66 MiB	13,178.66 MiB	+34.0 MiB

For the tested deterministic workload, the baseline and shared‑prefill paths
produced the same response‑token hash:

8811c53689938cf065bd44ad4c9093c876379ab9d93a2486cdcacfbcaf60c5fa

The executable benchmark and its CLI integration will be submitted separately
to DeepSpeedExamples, following the repository‑maintainer guidance.

Validation

The native transformer inference extension was rebuilt and verified to export
the new operation:

from deepspeed.ops.op_builder import InferenceBuilder
op = InferenceBuilder().load(verbose=True)
assert hasattr(op, "repeat_kv_cache_fp16")

Focused unit tests:

pytest -q tests/unit/runtime/rollout/test_hybrid_engine_rollout.py

Result:

18 passed

Tests cover:

• default‑off shared‑prefill configuration
• incompatible CUDA graph configuration
• prompt‑batch reduction
• logits and KV‑cache expansion
• decode forwards retaining the expanded batch
• Python fallback KV‑cache expansion
• native cache tensor pairing
• profiling behavior inherited from the prerequisite PR
• zero‑valued pad‑token handling

The modified files also pass the repository pre‑commit hooks.

Scope

This PR implements prompt sharing only for multiple response samples derived
from the same prompt within one rollout call.

It does not implement:

• the executable OPSD benchmark
• cross‑request prefix caching
• radix‑tree caching
• persistent prefix reuse between rollout calls
• paged KV‑cache allocation
• distributed shared prefill

The executable benchmark is being moved to DeepSpeedExamples as a separate
change.

Related to #8197.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7ce74d1d46

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[delock]

Hi @nathon-lee thanks for your PR. The performance data looks great! Noting picking inference engine v1 means we need to invest into inference kernels for modern models (Qwen for example), which current v1 engine does not support yet. But from the perspective of OPD/OPSD rollout v1 engine should be a good start point. @sfc-gh-truwase @PKUWZP

[delock]

@nathon-lee one comments is we will need to run the test under GPU in order to cover C++ kernels. Your benchmark test validates GPU but batch size =1 does not cover the reverse copy algorithm that is targeted for situation when batch>1. Can you make UT cover C++ kernels as well? Thanks!

[nathon-lee]

Thanks, @delock! I agree this is a gap. The current CPU tests only cover the reverse-copy ordering through the Python fallback, while the existing GPU benchmark uses B=1 and therefore doesn’t exercise the overlapping source rows.

I’ll add a CUDA-only native-kernel test with B=2 and repeats=2 to cover this case. It will verify that the C++ repeat_kv_cache implementation produces [A, A, B, B] and matches the fallback behavior. The test will also skip cleanly when CUDA or the native inference extension is unavailable.

[nathon-lee]

Hi @delock, thanks very much for taking a look!

I’ve added a CUDA-only native-kernel test with B=2 and repeats=2. It exercises the overlapping reverse-copy path and verifies that the resulting cache ordering is [A, A, B, B].

I also confirmed that the focused GPU test passes with the native inference extension: 1 passed, 18 deselected. Thanks again for your review!