[XPU] Gate FlashAttention-in-graph on oneAPI 2026.0+ runtime support

tests/utils_/test_torch_utils.py

@@ -13,6 +13,7 @@
     is_quantized_kv_cache,
     set_torch_threads_for_runtime,
     startup_omp_num_threads,
+    supports_xpu_fa_in_graph,
 )
 
 
@@ -116,6 +117,25 @@ def child_thread_func():
             pytest.fail("Child thread failed to exit properly")
 
 
+@pytest.mark.parametrize(
+    ("graph_supported", "xpu_ver", "expected"),
+    [
+        (True, "20260000", True),  # oneAPI 2026.0 -> FA capturable
+        (True, "20260100", True),  # newer 2026.x
+        (True, "20250302", False),  # oneAPI 2025.3 -> scratch not capturable
+        (True, None, False),  # non-XPU torch build
+        (True, "not-a-number", False),  # unparsable -> fail closed
+        (False, "20260000", False),  # torch too old for any XPU graph
+    ],
+)
+def test_supports_xpu_fa_in_graph(monkeypatch, graph_supported, xpu_ver, expected):
+    monkeypatch.setattr(
+        "vllm.utils.torch_utils.supports_xpu_graph", lambda: graph_supported
+    )
+    monkeypatch.setattr(torch.version, "xpu", xpu_ver, raising=False)
+    assert supports_xpu_fa_in_graph() is expected
+
+
 def test_current_stream_multithread():
     if not torch.cuda.is_available():
         pytest.skip("CUDA not available")

vllm/envs.py

@@ -296,6 +296,11 @@
     VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True
     VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32
     VLLM_XPU_ENABLE_XPU_GRAPH: bool = False
+    # When XPU graphs are enabled, clamp full graph modes to PIECEWISE so
+    # FlashAttention SYCL kernels with work_group_scratch_memory stay outside
+    # the captured graph. Set to 0 to allow FULL / FULL_AND_PIECEWISE (e.g.
+    # TRITON_ATTN experiments).
+    VLLM_XPU_GRAPH_FORCE_PIECEWISE: bool = True
     VLLM_XPU_USE_SAMPLER_KERNEL: bool = True
     VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
     VLLM_GPU_NIC_PCIE_MAPPING: str = ""
@@ -2040,6 +2045,12 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_XPU_ENABLE_XPU_GRAPH": lambda: bool(
         int(os.getenv("VLLM_XPU_ENABLE_XPU_GRAPH", "0"))
     ),
+    # When XPU graphs are on, force cudagraph_mode down to PIECEWISE if the
+    # configured mode still has full graphs. Default on (safe canary). Set
+    # VLLM_XPU_GRAPH_FORCE_PIECEWISE=0 to keep FULL / FULL_AND_PIECEWISE.
+    "VLLM_XPU_GRAPH_FORCE_PIECEWISE": lambda: bool(
+        int(os.getenv("VLLM_XPU_GRAPH_FORCE_PIECEWISE", "1"))
+    ),
     # whether use xpu specific sample kernel
     "VLLM_XPU_USE_SAMPLER_KERNEL": lambda: bool(
         int(os.getenv("VLLM_XPU_USE_SAMPLER_KERNEL", "1"))

vllm/platforms/xpu.py

@@ -285,7 +285,10 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
             compilation_config.compile_sizes = []
 
         # lazy import to avoid circular import
-        from vllm.utils.torch_utils import supports_xpu_graph
+        from vllm.utils.torch_utils import (
+            supports_xpu_fa_in_graph,
+            supports_xpu_graph,
+        )
 
         if not supports_xpu_graph():
             compilation_config.cudagraph_mode = CUDAGraphMode.NONE
@@ -300,10 +303,48 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 "please set VLLM_XPU_ENABLE_XPU_GRAPH=1 to enable it."
             )
         else:
+            fa_in_graph_ok = supports_xpu_fa_in_graph()
             logger.warning_once(
                 "XPU Graph support is experimental and currently only supports "
                 "single-GPU execution."
             )
+            mode = compilation_config.cudagraph_mode
+            wants_full = mode is not None and mode.has_full_cudagraphs()
+            if wants_full and envs.VLLM_XPU_GRAPH_FORCE_PIECEWISE:
+                # Default-on safety clamp (feature 01): keep FlashAttention
+                # outside the captured graph.
+                logger.warning_once(
+                    "VLLM_XPU_GRAPH_FORCE_PIECEWISE=1: overriding "
+                    "cudagraph_mode from %s to PIECEWISE so FlashAttention "
+                    "stays outside the XPU Graph. Set "
+                    "VLLM_XPU_GRAPH_FORCE_PIECEWISE=0 to capture FlashAttention "
+                    "in a full graph (needs oneAPI 2026.0+).",
+                    mode.name,
+                )
+                compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
+            elif wants_full and not fa_in_graph_ok:
+                # Fail-closed: operator opted out of the clamp but the runtime
+                # cannot capture FA scratch kernels into a SYCL Graph. Capturing
+                # would raise the work_group_scratch_memory RuntimeError, so
+                # re-clamp instead of crashing at model warmup.
+                logger.warning_once(
+                    "VLLM_XPU_GRAPH_FORCE_PIECEWISE=0 requested full graph "
+                    "mode %s, but this runtime cannot capture FlashAttention "
+                    "into a SYCL Graph (needs oneAPI 2026.0+; "
+                    "torch.version.xpu=%s). Falling back to PIECEWISE to avoid "
+                    "the work_group_scratch_memory SYCL Graph error. Upgrade "
+                    "the base image or use TRITON_ATTN for FULL mode.",
+                    mode.name,
+                    getattr(torch.version, "xpu", None),
+                )
+                compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
+            elif wants_full:
+                logger.info_once(
+                    "FlashAttention-in-graph enabled: capturing full XPU "
+                    "Graph (mode %s) including FlashAttention on oneAPI 2026.0+ "
+                    "runtime.",
+                    mode.name,
+                )
 
         # Disable fusion passes not yet supported on XPU.
         from vllm.config.compilation import CompilationMode

vllm/utils/torch_utils.py

@@ -1011,6 +1011,25 @@ def supports_xpu_graph() -> bool:
     return is_torch_equal_or_newer("2.11.0.dev")
 
 
+def supports_xpu_fa_in_graph() -> bool:
+    """Whether FlashAttention SYCL kernels can be captured into an XPU Graph.
+
+    FA kernels use ``sycl_ext_oneapi_work_group_scratch_memory``, which only
+    became capturable by the SYCL Graph extension on oneAPI 2026.0+ runtimes
+    (``torch.version.xpu >= 20260000``). On older runtimes capturing FA into a
+    full graph raises the ``work_group_scratch_memory ... not yet available for
+    use with the SYCL Graph extension`` RuntimeError, so full graph modes must
+    keep attention outside the capture (PIECEWISE).
+    """
+    if not supports_xpu_graph():
+        return False
+    xpu_ver = getattr(torch.version, "xpu", None)
+    try:
+        return xpu_ver is not None and int(xpu_ver) >= 20260000
+    except (TypeError, ValueError):
+        return False
+
+
 # create a library to hold the custom op
 vllm_lib = Library("vllm", "FRAGMENT")  # noqa