[Hybrid Engine] Enable graph capture for HybridEngine on non-CUDA accelerators (#8933)
## 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 #8271).
* The continuous path on Transformers 5.x still needs the
`get_query_offset` Cache-API compatibility work, tracked in #8852; this
PR does not include it.
## Related work
* #8271 added the per-position decode graphs for CUDA; this PR
implements its comment thread's `enable_graph_capture` rename request.
* #8732 added the continuous-rollout graphs for CUDA and routed the
capture-state check through the accelerator abstraction.
* #8852 (open) adds the Transformers 5.x `get_query_offset` support this
PR depends on for newer Transformers versions.
* Fixes #8932 — the accompanying feature request ([REQUEST][Hybrid
Engine] support hybrid engine for non-CUDA), which frames the motivation
and the hard-coding analysis.
---------
Signed-off-by: Zizhao Tang <1601677496@qq.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
DeepSpeed
ab7fd98f - [Hybrid Engine] Enable graph capture for HybridEngine on non-CUDA accelerators (#8933)
