perf(opsd): benchmark shared prefill with CUDA Graph decode

benchmarks/opsd/README.md

@@ -2,18 +2,20 @@
 
 `benchmark_hybrid_engine_rollout.py` measures rollout-level performance for an
 OPSD workload backed by DeepSpeed HybridEngine. It runs a matrix of synthetic,
-exact-length prompts and reports prompt expansion, generation, post-processing,
-total latency, generated-token throughput, and peak accelerator memory. Each
-case includes raw iteration profiles, mean and p50 summaries, and p95 summaries
-for latency metrics.
+exact-length prompts and reports prompt expansion, generation, prefill/decode
+forward timings, generation overhead, post-processing, total latency,
+generated-token throughput, and peak accelerator memory. Each case includes
+raw iteration profiles, mean/p50/p95 summaries for latency metrics, and a
+consistent `num_decode_forwards` count.
 
 This benchmark depends on the rollout profiling API introduced by
 DeepSpeed PR #8295:
 
 https://github.com/deepspeedai/DeepSpeed/pull/8295
 
 Use a DeepSpeed checkout that contains that API and place it first on
-`PYTHONPATH`. The current validation scope is one process, one GPU, and ZeRO-0.
+`PYTHONPATH`. The current validation scope is one process, one GPU, and ZeRO-0
+with `inference_tp_size=1`.
 This is a HybridEngine rollout benchmark, not a complete OPSD training-step
 benchmark; it does not measure teacher inference, loss computation, backward,
 or optimizer work.
@@ -27,18 +29,89 @@ control unreported warmup calls and recorded calls. Pass
 `--release-inference-cache` to release the inference cache after generation;
 otherwise the benchmark retains it. `--temperature`, `--top-p`, `--seed`, and
 `--output` control sampling, reproducibility, and the JSON output path.
+Pass `--use-shared-prefill` to enable shared prompt prefill. Shared-prefill
+currently supports only ZeRO-0, `inference_tp_size=1`, and the internal KV
+cache; it cannot be combined with `--release-inference-cache`. The selected
+mode is recorded as `use_shared_prefill` in the top-level JSON result.
+Pass `--use-graph-capture` to benchmark CUDA Graph decode. CUDA Graph capture
+can be combined with `--use-shared-prefill`, but supports only greedy,
+fixed-length generation (`--temperature 0`). The first generation performs
+graph capture, so use `--warmup 1` or greater before recording measurements.
 
 The largest effective batch (`batch_size * samples_per_prompt`) executes first
 so HybridEngine initializes a sufficiently large inference workspace. Results
 in the output JSON retain the matrix order requested on the command line.
 
-From the DeepSpeedExamples repository root, run a single-GPU benchmark with:
+## Profile fields
+
+The raw profile for every iteration contains these latency fields in
+milliseconds:
+
+- `prompt_expansion_ms`: time to expand prompts for `n_samples_per_prompt`.
+- `generation_ms`: time spent in generation after prompt expansion.
+- `prefill_forward_ms`: the first top-level model forward (the prefill).
+- `decode_forward_ms`: cumulative time of subsequent top-level model forwards.
+- `generation_overhead_ms`: `generation_ms` minus prefill and decode forward
+  time, covering the remaining generation work.
+- `post_processing_ms`: time to construct the rollout result after generation.
+- `total_ms`: end-to-end time for the profiled rollout.
+
+`num_decode_forwards` is metadata counting the forwards included in
+`decode_forward_ms`; it is retained in each raw profile and summarized once per
+case after verifying that every iteration has the same count. It is not a
+latency statistic. Forward breakdown fields can be `null` on paths such as
+CUDA Graph replay. If all iterations are unavailable, their case summary is
+`null`; a mixture of `null` and numeric values is rejected.
+
+For A/B comparisons, keep model, dtype, seed, workload matrix, warmup, and
+iterations exactly identical between runs; change only the feature under test.
+
+From the DeepSpeedExamples repository root, run the baseline single-GPU OPT-6.7B
+FP16 benchmark with:
+
+```bash
+PYTHONPATH=/workspace/DeepSpeed_woo:/workspace/DeepSpeedExamples \
+torchrun --nproc_per_node=1 \
+  benchmarks/opsd/benchmark_hybrid_engine_rollout.py \
+  --model facebook/opt-6.7b \
+  --dtype fp16 \
+  --batch-sizes 1 \
+  --samples-per-prompt 1 4 \
+  --prompt-lengths 128 512 \
+  --response-lengths 32 128 \
+  --temperature 0 \
+  --warmup 5 \
+  --iterations 20 \
+  --output /workspace/results/opsd_prefill_decode_baseline.json
+```
+
+Run the shared-prefill variant with the same workload and controls:
+
+```bash
+PYTHONPATH=/workspace/DeepSpeed_woo:/workspace/DeepSpeedExamples \
+torchrun --nproc_per_node=1 \
+  benchmarks/opsd/benchmark_hybrid_engine_rollout.py \
+  --model facebook/opt-6.7b \
+  --dtype fp16 \
+  --batch-sizes 1 \
+  --samples-per-prompt 1 4 \
+  --prompt-lengths 128 512 \
+  --response-lengths 32 128 \
+  --temperature 0 \
+  --warmup 5 \
+  --iterations 20 \
+  --use-shared-prefill \
+  --output /workspace/results/opsd_prefill_decode_shared.json
+```
+
+For a quick local smoke test, reduce warmup and iterations:
 
 ```bash
 PYTHONPATH=/workspace/DeepSpeed_woo:/workspace/DeepSpeedExamples \
 torchrun --nproc_per_node=1 \
   benchmarks/opsd/benchmark_hybrid_engine_rollout.py \
   --model facebook/opt-6.7b \
+  --dtype fp16 \
   --batch-sizes 1 \
   --samples-per-prompt 1 4 \
   --prompt-lengths 128 512 \

benchmarks/opsd/benchmark_hybrid_engine_rollout.py

@@ -17,8 +17,16 @@
 from pathlib import Path
 
 
-_LATENCY_FIELDS = ("prompt_expansion_ms", "generation_ms", "post_processing_ms", "total_ms")
-_THROUGHPUT_FIELDS = ("tokens_per_second", )
+_LATENCY_FIELDS = (
+    "prompt_expansion_ms",
+    "generation_ms",
+    "prefill_forward_ms",
+    "decode_forward_ms",
+    "generation_overhead_ms",
+    "post_processing_ms",
+    "total_ms",
+)
+_THROUGHPUT_FIELDS = ("tokens_per_second",)
 
 
 def _percentile(values, percentile):
@@ -31,6 +39,11 @@ def _summarize(profiles):
     summary = {}
     for field in _LATENCY_FIELDS:
         values = [profile[field] for profile in profiles]
+        if all(value is None for value in values):
+            summary[field] = None
+            continue
+        if any(value is None for value in values):
+            raise ValueError(f"Profile field {field} is None for only some iterations")
         summary[field] = {
             "mean": statistics.mean(values),
             "p50": statistics.median(values),
@@ -45,6 +58,13 @@ def _summarize(profiles):
     return summary
 
 
+def _validate_decode_forward_counts(profiles):
+    counts = [profile["num_decode_forwards"] for profile in profiles]
+    if len(set(counts)) != 1:
+        raise ValueError(f"num_decode_forwards differs across iterations: {counts}")
+    return counts[0]
+
+
 def _build_parser():
     parser = argparse.ArgumentParser(description="Benchmark HybridEngine rollout stage profiling")
     parser.add_argument("--model", default="facebook/opt-6.7b", help="HuggingFace model ID or local model path")
@@ -58,6 +78,8 @@ def _build_parser():
     parser.add_argument("--warmup", type=int, default=5)
     parser.add_argument("--iterations", type=int, default=20)
     parser.add_argument("--release-inference-cache", action="store_true")
+    parser.add_argument("--use-shared-prefill", action="store_true")
+    parser.add_argument("--use-graph-capture", action="store_true")
     parser.add_argument("--seed", type=int, default=1234)
     parser.add_argument("--output", default="opsd_rollout_profile.json")
     return parser
@@ -72,6 +94,10 @@ def _validate_args(args):
         raise ValueError("temperature must be non-negative")
     if not 0.0 < args.top_p <= 1.0:
         raise ValueError("top-p must be in the interval (0, 1]")
+    if getattr(args, "use_shared_prefill", False) and args.release_inference_cache:
+        raise ValueError("--use-shared-prefill cannot be combined with --release-inference-cache")
+    if args.use_graph_capture and args.temperature != 0.0:
+        raise ValueError("--use-graph-capture requires --temperature 0.0 for greedy generation")
 
 
 def _load_model_and_tokenizer(model_name, dtype, device):
@@ -106,6 +132,7 @@ def _build_engine(model, args):
         },
         "hybrid_engine": {
             "enabled": True,
+            "enable_cuda_graph": args.use_graph_capture,
             "max_out_tokens": max(args.prompt_lengths) + max(args.response_lengths),
             "release_inference_cache": args.release_inference_cache,
         },
@@ -157,6 +184,7 @@ def _run_case(rollout, model, args, batch_size, samples_per_prompt, prompt_lengt
         "requested_response_length": response_length,
         "returned_response_length": profiles[-1]["response_length"],
         "peak_memory_mb": accelerator.max_memory_allocated() / (1024**2),
+        "num_decode_forwards": _validate_decode_forward_counts(profiles),
         "summary": _summarize(profiles),
         "profiles": profiles,
     }
@@ -181,6 +209,8 @@ def _build_result(args, device, cases):
         "temperature": args.temperature,
         "top_p": args.top_p,
         "release_inference_cache": args.release_inference_cache,
+        "use_shared_prefill": getattr(args, "use_shared_prefill", False),
+        "use_graph_capture": args.use_graph_capture,
         "cases": cases,
     }
 
@@ -206,7 +236,12 @@ def _run(args):
     try:
         model, tokenizer = _load_model_and_tokenizer(args.model, dtype, device)
         engine = _build_engine(model, args)
-        rollout = HybridEngineRollout(engine, tokenizer, HybridEngineRolloutConfig(enable_profiling=True))
+        rollout = HybridEngineRollout(
+            engine,
+            tokenizer,
+            HybridEngineRolloutConfig(enable_profiling=True,
+                                      use_shared_prefill=args.use_shared_prefill),
+        )
 
         case_specs, execution_order = _ordered_case_specs(args)
         cases = [None] * len(case_specs)