Add CUDA graph support for HybridEngine generation (#8271)
## Motivation
Token generation dominates a HybridEngine RLHF iteration, and it is
bound by CPU work rather than by the GPU.
Measured on `facebook/opt-1.3b`, ZeRO-2, single H100, 256-token prompt,
128 new tokens:
| batch | ms / decode step | tok/s |
| ----: | ---------------: | ----: |
| 1 | 5.00 | 200 |
| 4 | 5.10 | 785 |
| 8 | 5.17 | 1549 |
| 16 | 5.58 | 2865 |
| 32 | 5.37 | 5956 |
32× the batch costs 1.07× the latency. The GPU is idle waiting on the
host. A single decode step issues roughly 1300 kernel launches, and of
the 5.22 ms step, 4.52 ms is the model forward while summed kernel time
is only ~2.75 ms.
Replaying a captured CUDA graph replaces those launches with one call.
## Design
Adds an opt-in `hybrid_engine.enable_cuda_graph` flag. Two properties of
the inference kernels shape the design.
**One graph per decode position.** The kernels read the current sequence
length from a host-side counter (`InferenceContext::current_tokens()` in
`csrc/transformer/inference/csrc/pt_binding.cpp`) and pass it to kernels
as a launch parameter. Capture freezes launch parameters, so a single
graph would keep reading and writing one position forever. Host code
*does* run during capture, so capturing one graph per position records
the correct sequence offsets into each.
**All-or-nothing per sequence.** That counter is advanced from host code
(`advance_tokens()`) and is not exposed to Python. Replay runs no host
code, so it leaves the counter behind, and an eager decode step after a
replay would use a stale sequence length and corrupt the KV cache. Eager
and replayed steps therefore must never be mixed within a sequence.
`begin_sequence()` makes the decision once, up front, from the pinned
generation length, before any decode step runs.
Capture is followed immediately by a replay, since capture records work
without executing it; the replay is what actually fills the KV cache for
that position.
## Safety
Graphs are refused, with a warning, wherever captured pointers would not
stay valid:
* **ZeRO stage 3** — parameters are gathered into fresh buffers for each
generate call, and the inference containers hold no persistent weights
(`attn_qkvw is None`). A graph would replay whichever buffers existed at
capture time, which is silently wrong rather than merely slow.
* **`release_inference_cache`** — frees the workspace buffers the graphs
write into.
* **`inference_tp_size > 1`** — untested here.
* **Unpinned generation length** — a sequence that outruns its captured
graphs cannot fall back to eager safely, so graphs engage only when
`min_new_tokens == max_new_tokens`.
Capture failures fall back to eager execution and disable graphs, rather
than failing the training job.
Weight updates were verified explicitly: `reset_params()` writes into
the same inference buffers in place, so the captured pointers stay valid
across optimizer steps. After a step that moved the logits by 11.75,
replay matched a fresh eager forward to within 0.09 and differed from
the pre-step result by the full 11.75 — i.e. graphs track updated
weights rather than replaying stale ones.
## Results
`facebook/opt-1.3b`, ZeRO-2, batch 8, 256-token prompt, 128 new tokens,
single H100, averaged over 4 measured iterations of an RLHF-style loop
(`eval` → `generate` → `train` → forward/backward/step):
| phase | eager | CUDA graph | speedup |
| ----------- | -------- | ---------- | ------- |
| generate | 676.7 ms | 360.8 ms | 1.88× |
| train step | 134.8 ms | 141.0 ms | 0.97× |
| **iteration** | **812.8 ms** | **503.5 ms** | **1.61×** |
* **Generated tokens are unchanged**: 1024/1024 token agreement with the
eager path, all 8 sequences identical end to end.
* **Peak memory**: 27.83 → 27.96 GiB (+132 MiB for 127 captured
positions).
* **One-time capture cost**: the first generation captures the graphs
and takes ~12 s; every later generation replays.
## Tests
`tests/unit/hybrid_engine/test_he_cuda_graph.py`:
* 13 CPU-only tests covering the generation-length gate, the ZeRO-3 /
`release_inference_cache` / `inference_tp_size` rejections, and the
dispatch state machine (unknown length, over-long length, prompt
forwards, and invalidation when the generation length changes).
* One GPU end-to-end test (`seq_inference` marker, `opt-125m`) asserting
that a graphed generation is token-identical to the same generation run
eagerly.
Docs: a Hybrid Engine section added to `docs/_pages/config-json.md`,
covering the config block and the new flag's requirements and
restrictions.
## Scope
Default is off, so nothing changes unless the flag is set; the eager
path measured identically before and after this change (676.7 ms
generate both ways). Only ZeRO-2 with `inference_tp_size=1` was
benchmarked, which is what the guards allow.
Signed-off-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>
