DeepSpeed

a9505fd6 - feat(opsd): share prompt prefill across rollout samples (#8296)

feat(opsd): share prompt prefill across rollout samples (#8296) ## 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: ```python from deepspeed.ops.op_builder import InferenceBuilder op = InferenceBuilder().load(verbose=True) assert hasattr(op, "repeat_kv_cache_fp16") ``` Focused unit tests: ```bash 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. --------- Signed-off-by: nathon-lee <leejianwoo@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>