Add CUDA graph support for HybridEngine generation

[PKUWZP]

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.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[delock]

Is it possible to change flag name from enable_cuda_graph to enable_graph_capture? I believe it should be universal across accelerators.

[PKUWZP]

Is it possible to change flag name from enable_cuda_graph to enable_graph_capture? I believe it should be universal across accelerators.

Let me submit another PR to change it. Thanks for flagging it out:)