diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index 6d63fc5db493..d48f71c43185 100644 --- a/tests/utils_/test_torch_utils.py +++ b/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") diff --git a/vllm/envs.py b/vllm/envs.py index 3f23fb615296..b11944338300 100755 --- a/vllm/envs.py +++ b/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")) diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index af1937754861..4094c7b4ed6f 100644 --- a/vllm/platforms/xpu.py +++ b/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 diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 0a188875b648..f046e5a2e5ac 100644 --- a/vllm/utils/torch_utils.py +++ b/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