[VenusTZZ]

Summary

The HybridEngine graph-capture paths (per-position decode graphs from #8271 and the continuous-rollout graphs from #8732) were gated on CUDA in several places: a device-name check, an input device-type check, and hard-coded "CUDA graph" naming. This PR replaces them with accelerator capability hooks so any backend that implements graph capture (create_graph() returning a graph, capture_to_graph, replay_graph, and is_current_stream_capturing) takes the same code path as CUDA.

Two accelerator-specific fixes ride along to make the NPU path actually work:

• NPU inference attention switches to the graph-mode flash attention variant (npu_fusion_attention_v3): the v1/v2 variants launch work on an internal secondary stream that CANN capture rejects at capture_end (error 107025).
• The NPU inference shim grows its KV-cache table on demand: layer_id is derived from DeepSpeedSelfAttention.num_layers, a process-wide counter that only ever grows, so any process that builds a second engine would otherwise index out of bounds.

Design

Capability gates instead of name gates. validate_cuda_graph_support becomes validate_graph_capture_support and rejects backends with get_accelerator().create_graph() is None instead of comparing device_name() with "cuda". The continuous-rollout path uses the same capability check. Its capture guard (is_current_stream_capturing, used by DeepSpeedStaticCache to skip host-side bounds checks while recording) already routes through the accelerator abstraction since #8732 (CUDA implementation, safe False default); this PR adds the NPU implementation.

Naming. Internal names, docstrings, and error messages drop the "CUDA" prefix (validate_graph_capture_support, "continuous graph capture", "per-position graph"). The config key is renamed to hybrid_engine.enable_graph_capture so the flag reads as a backend capability rather than a CUDA-specific feature, following the rename request in #8271's thread. The old hybrid_engine.enable_cuda_graph key still works as a deprecated alias: its value migrates to the new field with a warning (the existing DeepSpeed deprecated-field machinery), so existing config files keep working.

NPU flash attention. NPUInference._softmax_context calls torch_npu.npu_fusion_attention_v3, the graph-mode variant. The v1/v2 variants rejected capture with "capture model contains a stream that was not joined to the original stream" (107025); after the failed capture the fallback eager path then hit a secondary failure (561103) because host-side counters had already advanced. v3 captures cleanly in the BSH layout (the docs only promise BNSD), and both failure modes disappear together.

Safety

The existing refusals are preserved and remain device-agnostic: ZeRO stage 3, release_inference_cache, inference_tp_size > 1, and unpinned generation lengths all refuse graphs with a warning and fall back to eager, as do capture failures at runtime. The continuous-rollout path keeps its explicit rejections of aligned decode fronts, cache trimming, and flash_attention_2 implementations.

Results

Old container path (#8271 methodology)

facebook/opt-1.3b local weights, ZeRO-2, FP16, single Ascend NPU, 256-token prompt, 128 new tokens:

batch	ms / decode step	tok/s
1	7.16	140
4	10.86	368
8	12.27	652
16	15.77	1015
32	22.37	1431

Throughput scales 10.2× over a 32× batch increase. Batch 8, averaged over 4 measured iterations of an RLHF-style loop (eval → generate → train → forward/backward/step):

phase	eager	graph	speedup
generate	5436.0 ms	1578.4 ms	3.44×
train step	797.7 ms	797.7 ms	1.00×
iteration	6233.7 ms	2376.1 ms	2.62×

• Generated tokens are unchanged: 1024/1024 token agreement with the eager path (8 sequences × 128 tokens, same input).
• Peak memory: 8.94 → 9.04 GiB (+0.10 GiB for 127 captured positions).
• One-time capture cost: 8.9 s for the first generation; every later generation replays.

Continuous rollouts (#8732 methodology)

gpt2 local weights, FP16, batch 1, 32-token prompt, 32 new tokens, greedy, continuous batching on a single Ascend NPU, two warmup calls and five steady-state measurements:

Mode	Median latency	Throughput
Eager	1060.44 ms	30 tokens/s
Graph reuse	177.59 ms	180 tokens/s

That is a 5.97× steady-state speedup, with the first call (including capture) at 395.5 ms reported separately. Two different prompts (seeds 1234/1235) produced token-identical eager and graph outputs (32/32 tokens each) while reusing one capture.

Tests

tests/unit/hybrid_engine/test_he_graph.py (renamed from test_he_cuda_graph.py):

• 16 CPU-only tests covering the generation-length gate, the ZeRO-3 / release_inference_cache / inference_tp_size rejections, the dispatch state machine, capture-failure fallback, and invalidation on length change.
• One end-to-end test (seq_inference marker) asserting a graphed generation is token-identical to the same generation run eagerly.

Continuous-rollout suite, run on NPU with Transformers 4.57.6:

pytest -q \
  tests/unit/runtime/rollout/test_continuous_batching.py \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py \
  tests/unit/v1/test_hybrid_engine_rollout_v1.py
# 55 passed, 3 skipped

NPU end-to-end validation ran locally on a single Ascend NPU (PyTorch 2.10 + torch_npu 2.10.0.post7; Transformers 4.57.6): both paths capture and replay correctly with token-identical outputs, and the eager path is unchanged. The three HybridEngine scripts in DeepSpeedExamples (TODO: fill in the script names) were also run end-to-end on NPU with graph capture enabled.

Scope

The graph flag remains opt-in and off by default; nothing changes for CUDA, whose eager path measured identically before and after. Only ZeRO-2 with inference_tp_size=1 was benchmarked on the container path, matching the upstream guards. Known limitations carried over or found while measuring:

• The NPU train step was measured over 2 iterations: ZeRO-2 on NPU retains ~4.9 GiB of live device memory per train step (not freed by del/GC), so a 4-step loop exhausts the 32 GB card. This is unrelated to graph capture and tracked separately.
• Generate and train were measured in separate processes (iteration = sum of means; the eval↔train switch overhead is ~0.07% per Add CUDA graph support for HybridEngine generation #8271).
• The continuous path on Transformers 5.x still needs the get_query_offset Cache-API compatibility work, tracked in Support Transformers 5.20 cache query offsets in DeepSpeedStaticCache #8852; this PR does not include it.

Related work

• Add CUDA graph support for HybridEngine generation #8271 added the per-position decode graphs for CUDA; this PR implements its comment thread's enable_graph_capture rename request.
• Enable CUDA Graph for continuous rollouts #8732 added the continuous-rollout graphs for CUDA and routed the capture-state check through the accelerator abstraction.
• Support Transformers 5.20 cache query offsets in DeepSpeedStaticCache #8852 (open) adds the Transformers 5.x get_query_offset support this PR depends on for newer Transformers versions.
• Fixes [REQUEST][Hybrid Engine] Enable graph capture for HybridEngine on non-CUDA accelerators #8932 — the accompanying feature request ([REQUEST][Hybrid Engine] support hybrid engine for non-CUDA), which frames the motivation and the hard-coding analysis.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-29T06:28:23.785292Z We couldn’t complete this review. Comment @codex review to try again.	454611a	Manual request

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

Codex Review: Something went wrong. Try again later by commenting “@codex review”.

An unknown error occurred

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Codex Review: Something went wrong. Try again later by commenting “@codex review”.

An unknown error occurred

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".