[Hybrid Engine] Enable graph capture for HybridEngine on non-CUDA accelerators

accelerator/npu_accelerator.py

@@ -247,6 +247,9 @@ def is_triton_supported(self):
         return False
 
     # Graph operations
+    def is_current_stream_capturing(self) -> bool:
+        return torch.npu.is_current_stream_capturing()
+
     def create_graph(self):
         return torch.npu.NPUGraph()
 

deepspeed/runtime/config.py

@@ -363,7 +363,9 @@ class HybridEngineConfig(DeepSpeedConfigModel):
     release_inference_cache: bool = False
     pin_parameters: bool = True
     tp_gather_partition_size: int = 8
-    enable_cuda_graph: bool = False
+    enable_graph_capture: bool = False
+    # Deprecated alias for enable_graph_capture; the value migrates with a warning.
+    enable_cuda_graph: bool = Field(False, json_schema_extra={"deprecated": True, "new_param": "enable_graph_capture"})
 
 
 def get_hybrid_engine_config(param_dict):

deepspeed/runtime/hybrid_engine.py

@@ -21,7 +21,7 @@
 from deepspeed.utils import logger
 from deepspeed.module_inject.layers import LinearLayer, Normalize, EmbeddingLayer, OPTEmbedding
 from ..ops.transformer.inference.op_binding.workspace import WorkspaceOp
-from .hybrid_engine_graph import (DecodeGraphCache, decode_steps_from_generate_kwargs, validate_cuda_graph_support)
+from .hybrid_engine_graph import (DecodeGraphCache, decode_steps_from_generate_kwargs, validate_graph_capture_support)
 
 
 class DeepSpeedHybridEngine(DeepSpeedEngine):
@@ -61,10 +61,10 @@ def __init__(self, args, model, **kwargs):
 
         self._orig_module_forward = None
         self._decode_graphs = None
-        if self._config.hybrid_engine.enable_cuda_graph and len(self._inference_containers) > 0:
-            unsupported = validate_cuda_graph_support(self._config.hybrid_engine, self._config.zero_config.stage)
+        if self._config.hybrid_engine.enable_graph_capture and len(self._inference_containers) > 0:
+            unsupported = validate_graph_capture_support(self._config.hybrid_engine, self._config.zero_config.stage)
             if unsupported is not None:
-                logger.warning(f"HybridEngine: running without CUDA graphs. {unsupported}.")
+                logger.warning(f"HybridEngine: running without graph capture. {unsupported}.")
             else:
                 self._orig_module_forward = self.module.forward
                 self._decode_graphs = DecodeGraphCache(self._orig_module_forward,
@@ -203,8 +203,8 @@ def prepare_shared_prefill(self, source_batch_size, repeats, prompt_length):
             raise RuntimeError("Shared prefill does not support inference tensor parallelism")
         if hybrid_config.release_inference_cache:
             raise RuntimeError("Shared prefill does not support release_inference_cache")
-        if hybrid_config.enable_cuda_graph:
-            raise RuntimeError("Shared prefill does not support CUDA graph capture")
+        if hybrid_config.enable_graph_capture:
+            raise RuntimeError("Shared prefill does not support graph capture")
         if len(self._inference_containers) == 0:
             raise RuntimeError("Shared prefill requires HybridEngine inference containers")
 

deepspeed/runtime/hybrid_engine_graph.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # DeepSpeed Team
-"""CUDA graph capture for HybridEngine token generation.
+"""Graph capture for HybridEngine token generation.
 
 Generation in the HybridEngine is bound by CPU work rather than by the GPU: a
 single decode step issues on the order of a thousand kernel launches, and the
 kernels finish well before the host can queue the next step. Capturing the decode
-forward into a CUDA graph replaces all of those launches with one replay.
+forward into a graph replaces all of those launches with one replay.
 
 Two properties of the inference kernels shape this design.
 
@@ -37,7 +37,7 @@ def _tensor_kwargs(kwargs):
 
 
 class DecodeGraphCache:
-    """Dispatches generation forwards to a per-position CUDA graph.
+    """Dispatches generation forwards to a per-position graph.
 
     Args:
         forward: The original ``forward``, used both for capture and as the
@@ -97,7 +97,7 @@ def begin_sequence(self, num_decode_steps):
         # A different generation length needs its own set of graphs, because each
         # graph has both its sequence offset and its attention-mask width baked in.
         if self._captured_length is not None and self._captured_length != num_decode_steps:
-            logger.info(f"HybridEngine CUDA graph: generation length changed "
+            logger.info(f"HybridEngine graph capture: generation length changed "
                         f"{self._captured_length} -> {num_decode_steps}, recapturing.")
             self.invalidate()
 
@@ -141,7 +141,7 @@ def _fall_back(self, reason, args, kwargs):
         Only safe before any replay has happened in the current sequence, which
         is why this is reachable only from the capture path.
         """
-        logger.warning(f"HybridEngine CUDA graph disabled: {reason}")
+        logger.warning(f"HybridEngine graph capture disabled: {reason}")
         self.invalidate()
         self._disabled = True
         self._use_graphs = False
@@ -161,9 +161,9 @@ def __call__(self, *args, **kwargs):
             # begin_sequence() bounds the length, so this means the caller ran
             # longer than it declared. Replays have already happened, so eager
             # execution would read a stale sequence counter.
-            raise RuntimeError(f"HybridEngine CUDA graph: generation exceeded the declared length "
+            raise RuntimeError(f"HybridEngine graph capture: generation exceeded the declared length "
                                f"({self._max_positions} decode steps). Set hybrid_engine.max_out_tokens "
-                               f"to cover the longest generation, or disable hybrid_engine.enable_cuda_graph.")
+                               f"to cover the longest generation, or disable hybrid_engine.enable_graph_capture.")
 
         if position in self._graphs:
             if self._shapes_match(position, kwargs):
@@ -195,23 +195,23 @@ def decode_steps_from_generate_kwargs(kwargs):
     return int(max_new) - 1
 
 
-def validate_cuda_graph_support(config, zero_stage):
-    """Return the reason CUDA graphs cannot be used, or ``None`` if they can."""
-    if not get_accelerator().is_available() or get_accelerator().device_name() != "cuda":
-        return "CUDA graphs require the CUDA accelerator"
+def validate_graph_capture_support(config, zero_stage):
+    """Return the reason graph capture cannot be used, or ``None`` if it can."""
+    if not get_accelerator().is_available() or get_accelerator().create_graph() is None:
+        return "graph capture is not supported on the current accelerator"
 
     if zero_stage == 3:
         # Under ZeRO-3 the inference containers hold no persistent weights; the
         # parameters are gathered into fresh buffers for each generate call. A
         # graph would keep replaying whichever buffers existed at capture time,
         # which is silently wrong rather than merely slow.
-        return "CUDA graphs are not supported with ZeRO stage 3"
+        return "graph capture is not supported with ZeRO stage 3"
 
     if config.release_inference_cache:
         # Releasing the workspace frees the buffers the graphs write into.
-        return "CUDA graphs are not supported with release_inference_cache"
+        return "graph capture is not supported with release_inference_cache"
 
     if config.inference_tp_size > 1:
-        return "CUDA graphs are not supported with inference_tp_size > 1"
+        return "graph capture is not supported with inference_tp_size > 1"
 
     return None

deepspeed/runtime/rollout/hybrid_engine_rollout.py

@@ -7,9 +7,9 @@
   1. **model.generate()** (default): delegates to HuggingFace generate.
      Supports sampling (temperature, top_p) and greedy.
   2. **graph capture + DeepSpeedStaticCache**: only for greedy (temperature=0).
-     Pre-allocates a StaticCache, captures the decode forward pass with a
-     CUDA graph, and replays it for each decode step.  Eliminates kernel
-     launch overhead.
+     Pre-allocates a StaticCache, captures the decode forward pass into a
+     graph, and replays it for each decode step.  Eliminates kernel launch
+     overhead.
   3. **continuous batching (experimental)**: an opt-in bounded greedy path
      selected through ``SamplingConfig.continuous_batch_size``.
 """
@@ -181,7 +181,7 @@ def generate(self, request: RolloutRequest, sampling: SamplingConfig) -> Rollout
         try:
             if self.use_shared_prefill and n > 1:
                 if self.use_graph_capture:
-                    raise RuntimeError("Shared prefill does not support CUDA graph capture")
+                    raise RuntimeError("Shared prefill does not support graph capture")
                 self.engine.prepare_shared_prefill(B, n, prompt_len)
                 shared_prefill_handles = self._register_shared_prefill_hooks(module, B, n)
             if self.use_graph_capture and is_greedy:
@@ -545,7 +545,7 @@ def _generate_continuous(self, request, sampling, max_batch_size):
 
     @torch.no_grad()
     def _generate_continuous_graph(self, original_request, requests, sampling, max_batch_size):
-        """Run continuous greedy decoding with a fixed-capacity CUDA graph.
+        """Run continuous greedy decoding with a fixed-capacity graph.
 
         Scheduler slots stay attached to physical cache rows for their whole
         lifetime. Retired rows are masked until a new prompt is prefetched into
@@ -559,10 +559,10 @@ def _generate_continuous_graph(self, original_request, requests, sampling, max_b
             raise ValueError("continuous batching requires a model with cache-class support; use the default "
                              "generate() path or upgrade transformers")
         if getattr(module.config, "_attn_implementation", None) == "flash_attention_2":
-            raise ValueError("continuous CUDA graph capture does not support flash_attention_2")
+            raise ValueError("continuous graph capture does not support flash_attention_2")
         device = requests[0].prompt_ids.device
-        if device.type != "cuda":
-            raise ValueError("continuous CUDA graph capture requires CUDA rollout inputs")
+        if get_accelerator().create_graph() is None:
+            raise ValueError("continuous graph capture is not supported on the current accelerator")
 
         from transformers import StaticCache
         from deepspeed.utils.static_cache import DeepSpeedStaticCache
@@ -735,9 +735,9 @@ def _generate_continuous_graph(self, original_request, requests, sampling, max_b
 
     def _validate_continuous_graph_options(self, prompt_len, sampling):
         if self.align_decode_fronts:
-            raise ValueError("continuous CUDA graph capture does not support aligned decode fronts")
+            raise ValueError("continuous graph capture does not support aligned decode fronts")
         if self.enable_cache_trimming:
-            raise ValueError("continuous CUDA graph capture does not support cache trimming")
+            raise ValueError("continuous graph capture does not support cache trimming")
         if (self.continuous_cache_capacity is not None
                 and self.continuous_cache_capacity < prompt_len + sampling.max_new_tokens):
             raise ValueError("continuous_cache_capacity must fit the prompt and all generated tokens")
@@ -1304,7 +1304,7 @@ def _pad_after_eos(output_ids, response_start, eos_token_id, pad_token_id):
     # ------------------------------------------------------------------
 
     def _generate_graph(self, prompt_ids, prompt_attn, max_new_tokens, pad_token_id, module, device):
-        """Greedy decode with DeepSpeedStaticCache + CUDA graph capture."""
+        """Greedy decode with DeepSpeedStaticCache + graph capture."""
         from transformers import StaticCache
         from deepspeed.utils.static_cache import DeepSpeedStaticCache
 

docs/_pages/config-json.md

@@ -870,17 +870,17 @@ The Hybrid Engine (`DeepSpeedHybridEngine`) switches a model between training mo
     "release_inference_cache": false,
     "pin_parameters": true,
     "tp_gather_partition_size": 8,
-    "enable_cuda_graph": false
+    "enable_graph_capture": false
   }
 ```
 
-***enable_cuda_graph***: [boolean]
+***enable_graph_capture***: [boolean]
 
 | Description                                                                                                                                                                                                                                                                                                                                        | Default |
 | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
-| Capture token generation into CUDA graphs. Generation issues on the order of a thousand kernel launches per token and is bound by CPU launch overhead rather than by the GPU, so replaying a captured graph removes most of the per-token cost. One graph is captured per decode position, and generated tokens are unchanged. | `false` |
+| Capture token generation into device graphs. Generation issues on the order of a thousand kernel launches per token and is bound by CPU launch overhead rather than by the accelerator, so replaying a captured graph removes most of the per-token cost. One graph is captured per decode position, and generated tokens are unchanged. The previous key `enable_cuda_graph` is deprecated: its value migrates to this field with a warning, so existing configs keep working. | `false` |
 
-`enable_cuda_graph` requires a pinned generation length (`min_new_tokens` equal to `max_new_tokens`) and `max_out_tokens` large enough to cover the longest generation. It is ignored, with a warning, when any of the following apply, since captured graphs would not stay valid:
+`enable_graph_capture` requires a pinned generation length (`min_new_tokens` equal to `max_new_tokens`) and `max_out_tokens` large enough to cover the longest generation. It is ignored, with a warning, when any of the following apply, since captured graphs would not stay valid:
 
 * ZeRO stage 3, where parameters are gathered into fresh buffers for each generation
 * `release_inference_cache: true`, which frees the buffers the graphs write into

op_builder/npu/inference.py

@@ -219,6 +219,12 @@ def _softmax_context(query_key_value,
         is_promt = seq_len > 1
         if not InferenceContext.kv_caches:
             InferenceContext.kv_caches = [[None, None, 0] for _ in range(num_layers)]
+        # layer_id is derived from DeepSpeedSelfAttention.num_layers, a process-wide
+        # counter that only ever grows (ds_attention.py:24-25): a second engine built
+        # in the same process sees layer_ids beyond the first engine's num_layers.
+        # Grow the cache on demand instead of trusting num_layers to cover them.
+        while len(InferenceContext.kv_caches) <= layer_id:
+            InferenceContext.kv_caches.append([None, None, 0])
         if is_promt:
             InferenceContext.reset_tokens(seq_len)
             # Keep the previous buffer when it fits, so fixed-length workloads (e.g.
@@ -278,19 +284,22 @@ def _softmax_context(query_key_value,
         layer_scale = max(1, layer_id) if len(alibi.size()) > 1 else 1.0
         alpha = norm_factor * norm_factor / layer_scale
 
-        output = torch_npu.npu_fusion_attention(q,
-                                                k_buf[:, :kv_len],
-                                                v_buf[:, :kv_len],
-                                                heads,
-                                                "BSH",
-                                                pse=None,
-                                                padding_mask=None,
-                                                atten_mask=attn_mask.bool(),
-                                                scale=alpha,
-                                                pre_tockens=65536,
-                                                next_tockens=65536,
-                                                keep_prob=1,
-                                                inner_precise=0)[0]
+        # v3 is the graph-mode version of npu_fusion_attention: v1/v2 launch work on
+        # an internal secondary stream that CANN capture rejects (107025), v3 does not.
+        # (torch_npu documents aclgraph support for BNSD only, but BSH captures fine.)
+        output = torch_npu.npu_fusion_attention_v3(q,
+                                                   k_buf[:, :kv_len],
+                                                   v_buf[:, :kv_len],
+                                                   heads,
+                                                   "BSH",
+                                                   pse=None,
+                                                   padding_mask=None,
+                                                   atten_mask=attn_mask.bool(),
+                                                   scale=alpha,
+                                                   pre_tockens=65536,
+                                                   next_tockens=65536,
+                                                   keep_prob=1,
+                                                   inner_precise=0)[0]
 
         # Align with the CUDA csrc (pt_binding.cpp:548): increment the token count after the
         # last layer's forward, so the next decode token's seq_offset (soft_len-1) lands on