[nathon-lee]

Background

Follow-up to #8732 and its
review suggestions.
Training updates invalidate the continuous rollout graph, so steady-state latency alone does not capture the cost of alternating training and generation.

Implementation status

The initial benchmark records CUDA-synchronized latency and throughput for eager / graph rollouts around one real SGD update. Each measured call's tokens, attention masks, and response offsets are checked against native Hugging Face generate() using the same weights.

Existing GPU measurements

Recorded on an RTX A4500 (20 GiB), PyTorch 2.9.1+cu128 / CUDA 12.8, Transformers 5.20.1, FP16, using a randomly initialized tiny Llama. Steady-state values are medians of 10 calls; first-call values are single samples.

Phase	Eager (ms)	CUDA Graph (ms)
First call	240.182	61.933
Steady state	228.014	45.424
First call after SGD update	226.304	58.800
Steady state after update	234.919	45.281

Graph first-call overhead was 16.509 ms before SGD and 13.519 ms afterward; the training update took 55.849 ms. These are synthetic observations, not production-model performance claims. The overhead values are end-to-end differences, not isolated capture timings.

Reproduce on the benchmark checkout with the Issue 4 fix:

python -m benchmarks.rollout.bench_cuda_graph_first_call \
    --requests 8 --batch-size 4 --prompt-len 32 --max-new-tokens 32 \
    --iterations 10 --output rollout-first-call.json

CUDA / model libraries were initialized without warming up the rollout graph; each phase measured eager before graph. The original combined working tree reported 60 passed, 1 skipped and passing pre-commit checks. The skip required an optional pretrained model. This benchmark was not rerun during the separate Issue 4 validation.

Remaining work and acceptance criteria

• Separate cache / static-buffer allocation, warmup, and capture timings from total rollout latency, without double-counting.
• Report per-phase capture counts, allocated / reserved and peak GPU memory, latency, and generated-token throughput.
• Measure repeated training / rollout cycles with real optimizer updates; preserve the four phases and generation-parity checks.
• Keep synchronized timing, fixed inputs, reproducible commands, and complete model / attention / hardware / software / workload metadata. Clearly label synthetic versus pretrained-model results; none of the latter are available yet.

Scope and dependencies

Depends on #8732 and, on Transformers 5.20.1, Issue 4. Keep the runtime fix separate. Measure existing behavior without weakening parameter-version checks. Benchmark code and usage belong in benchmarks/rollout/bench_cuda_graph_first_call.py and benchmarks/README.md, not the Issue 4 branch.