[CB] Make CB more device agnostic and add support for XPU

@@ -18,7 +18,7 @@ import threading from abc import abstractmethod from collections.abc import Callable, Generator-from contextlib import contextmanager, nullcontext+from contextlib import contextmanager from datetime import timedelta from time import perf_counter from typing import Any@@ -45,7 +45,7 @@ from .offloading_manager import OffloadingManager from .requests import GenerationOutput, RequestState, RequestStatus, logger from .scheduler import SCHEDULER_MAPPING, FIFOScheduler, Scheduler-from .utils import ThreadLocalCounter, WorkloadHints, drain_queue+from .utils import ThreadLocalCounter, WorkloadHints, drain_queue, stream_context   """@@ -403,8 +403,8 @@ def __repr__(self) -> str:     def __del__(self) -> None:         self.inputs_and_outputs = None  # clean up CUDA graphs in priority         gc.collect()-        if torch.cuda.is_available():-            torch.cuda.empty_cache()+        if hasattr(device_module := torch.get_device_module(), "empty_cache"):+            device_module.empty_cache()      def reset(self) -> None:         """Reset the batch processor for a new generation loop."""@@ -600,9 +600,7 @@ def update_batch(self) -> None:                     self.scheduler.add_waiting_request(child_state)              # Actually perform the block copies-            compute_stream = self.inputs_and_outputs.compute_stream-            maybe_stream = torch.cuda.stream(compute_stream) if compute_stream is not None else nullcontext()-            with maybe_stream:+            with stream_context(self.inputs_and_outputs.compute_stream):                 self.cache.perform_cache_copy(fork_src_and_dst)      def has_pending_requests(self) -> bool:@@ -859,8 +857,8 @@ def stop(          # In all cases, a little cleanup is good         gc.collect()-        if torch.cuda.is_available():-            torch.cuda.empty_cache()+        if hasattr(device_module := torch.get_device_module(), "empty_cache"):+            device_module.empty_cache()      def join(self, stop_trigger_time: float, timeout: float | None = None) -> None:         """Wait for the background thread to finish. Wait can be capped using the timeout argument (in seconds)."""

@@ -26,7 +26,7 @@ from ...utils.generic import is_flash_attention_requested from .cache import ATTN_TYPE_TO_ALLOCATOR, group_layers_by_attn_type from .requests import logger-from .utils import WorkloadHints+from .utils import WorkloadHints, get_accelerator_graph_name   FALLBACK_DEFAULTS = {@@ -226,12 +226,12 @@ def decide_use_cuda_graphs(      This function modifies the `use_cuda_graph` attribute of the config in place, to a tuple of booleans.     """-    # If cuda is not available, we cannot use cuda graphs-    if not torch.cuda.is_available():+    # If no graph backend is available on this accelerator, we cannot use graphs.+    if get_accelerator_graph_name() is None:         intended_use_cuda_graph = any(cb_config.cuda_graph_booleans)         if intended_use_cuda_graph:  # throw a warning only if the user intended to use cuda graphs             logger.warning(-                f"{cb_config.use_cuda_graph = } but {torch.cuda.is_available() = }: turning off cuda graphs"+                f"{cb_config.use_cuda_graph = } but no graph backend is available: turning off accelerator graphs."             )         cb_config.use_cuda_graph = (False, False) 

@@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License.-from contextlib import nullcontext from functools import partial from itertools import repeat from typing import TypedDict@@ -26,7 +25,15 @@ from .cache_allocators import FULL_ATTENTION, SLIDING_ATTENTION from .cb_logits_processors import ContinuousBatchingLogitsProcessorList from .requests import TMP_TOKEN_ID, FutureRequestState, logger-from .utils import CudaGraphBuffer, aligned_divide, attn_mask_is_needed, build_attention_mask, pad_to_pow2+from .utils import (+    CudaGraphBuffer,+    aligned_divide,+    attn_mask_is_needed,+    build_attention_mask,+    create_device_stream,+    pad_to_pow2,+    stream_context,+)   class PagedAttentionArgs(TypedDict):@@ -114,10 +121,11 @@ def __init__(         self.requests_in_batch: list[FutureRequestState] = []         self.req_id_to_new_token_position: dict[str, int] = {}  # only used for async API         self.graphs: CudaGraphBuffer = CudaGraphBuffer()-        # Setup static tensors and compute stream+        # Setup static tensors         self._setup_static_tensors(logit_processor=logit_processor)         self._reset_static_tensors(full_reset=True)-        self.compute_stream = torch.cuda.Stream(device=self.device) if device.type == "cuda" else None+        # If the device is an accelerator that supports streams, also create a compute stream+        self.compute_stream = create_device_stream(device) if device.type != "cpu" else None      def _setup_static_tensors(self, logit_processor: ContinuousBatchingLogitsProcessorList) -> None:         """Allocates static tensors for generation inputs and outputs. This is called only once at init time, to avoid@@ -227,8 +235,7 @@ def _transfer_inputs(         other.max_length_q = self.max_length_q         other.max_length_k = dict(self.max_length_k)         # Transfer static tensors-        maybe_stream = torch.cuda.stream(stream) if stream is not None else nullcontext()-        with maybe_stream:+        with stream_context(stream):             other._bulk_input_tensor.copy_(self._bulk_input_tensor, non_blocking=non_blocking)  # fast bulk transfer             # Only transfer block_table for decode-only batches (when it's actually used)             if self.use_block_table:@@ -596,24 +603,21 @@ def __init__(             model_dtype=model_dtype,             logit_processor=logit_processor,         )-        # Create events only on CUDA devices-        self.h2d_over = torch.cuda.Event() if torch.cuda.is_available() else None-        self.compute_over = torch.cuda.Event() if torch.cuda.is_available() else None-        self.d2h_over = torch.cuda.Event() if torch.cuda.is_available() else None+        self.h2d_over = torch.Event(device, enable_timing=False)+        self.compute_over = torch.Event(device, enable_timing=False)+        self.d2h_over = torch.Event(device, enable_timing=False)      def reset(self) -> None:         self.host_io.reset()         self.device_io.reset()         for event in [self.h2d_over, self.compute_over, self.d2h_over]:-            if event is not None:-                event.synchronize()+            event.synchronize()      def transfer_inputs_h2d(self, stream: torch.cuda.Stream) -> None:         self.host_io._transfer_inputs(self.device_io, stream=stream, non_blocking=True)      def transfer_outputs_d2h(self, stream: torch.cuda.Stream | None) -> None:-        maybe_stream = torch.cuda.stream(stream) if stream is not None else nullcontext()-        with maybe_stream:+        with stream_context(stream):             self.host_io.output_ids.copy_(self.device_io.output_ids, non_blocking=True)  @@ -663,6 +667,8 @@ class ContinuousBatchingAsyncIOs:     Proper ordering of steps is ensured through the use of CUDA events and streams.     """ +    _supported_device_types = ("cuda", "xpu")+     def __init__(         self,         cache: PagedAttentionCache,@@ -672,9 +678,13 @@ def __init__(         model_dtype: torch.dtype,         logit_processor: ContinuousBatchingLogitsProcessorList,     ) -> None:-        # Async batching needs streams to function, so check is CUDA is available-        if not torch.cuda.is_available():-            raise RuntimeError(f"Async batching requires CUDA, but {torch.cuda.is_available() = }")+        # Check device module is compatible with async batching+        if device.type not in self._supported_device_types:+            raise RuntimeError(+                f"Async batching requires a device type in {self._supported_device_types} but got {device.type = }"+            )+        if not torch.get_device_module(device).is_available():+            raise RuntimeError("Async batching requires an available device.")         # IO pairs used to avoid race conditions         self.current_pair = 0         self.io_pairs = [@@ -689,9 +699,11 @@ def __init__(             for _ in range(2)         ]         # CUDA streams-        self.h2d_stream = torch.cuda.Stream(device=device)-        self.d2h_stream = torch.cuda.Stream(device=device)-        self.compute_stream = torch.cuda.Stream(device=device)+        self.h2d_stream: torch.cuda.Stream = create_device_stream(device)+        self.d2h_stream: torch.cuda.Stream = create_device_stream(device)+        self.compute_stream: torch.cuda.Stream = create_device_stream(device)+        if self.h2d_stream is None or self.d2h_stream is None or self.compute_stream is None:+            raise RuntimeError(f"Async batching requires stream to function. Stream creation failed with {device = }.")         # Set all unused compute streams to None         self.io_pairs[0].host_io.compute_stream = None         self.io_pairs[0].device_io.compute_stream = None

@@ -13,7 +13,6 @@ # limitations under the License. import time from collections.abc import Callable-from contextlib import nullcontext  import torch from torch import nn@@ -23,7 +22,16 @@ from .cb_logits_processors import ContinuousBatchingLogitsProcessorList from .input_outputs import ContinuousBatchingAsyncIOs, ContinuousBatchingIOs from .requests import RequestState, RequestStatus, logger-from .utils import create_warmup_future_states, get_cuda_pools, mem_pool_ctx, pad_to_interval, pad_to_pow2+from .utils import (+    accelerator_graph_capture_context,+    create_accelerator_graph,+    create_warmup_future_states,+    get_memory_pools,+    memory_pool_context,+    pad_to_interval,+    pad_to_pow2,+    stream_context,+)   def infer_max_single_request_tokens(cache: PagedAttentionCache) -> int:@@ -72,7 +80,7 @@ def __init__(          # Set up the graph pool. This allows all graphs to share the same memory pool, greatly saving memory.         if self.use_cuda_graph_varlen or self.use_cuda_graph_decode:-            self.mem_pool, self.graph_pool_id = get_cuda_pools()+            self.mem_pool, self.graph_pool_id = get_memory_pools()         else:             self.mem_pool, self.graph_pool_id = None, None @@ -133,16 +141,15 @@ def compute_batch(self, model: nn.Module, batch_data: dict) -> None:          # If we are not using CUDA graphs, we perform the generation step and return         if not use_cuda_graph:-            maybe_stream = torch.cuda.stream(compute_stream) if compute_stream is not None else nullcontext()-            with maybe_stream:+            with stream_context(compute_stream):                 forward_fn(model, batch_data, carry_over_ids, prev_output_ids, output_ids) -        # Otherwise, we either create or replay the graph (CUDA is available in this path)+        # Otherwise, we either create or replay the graph (an accelerator is available in this path)         else:             graph = self.inputs_and_outputs.get_graph()             # Case: the graph already exists, so we replay it             if graph is not None:-                with torch.cuda.stream(compute_stream):+                with stream_context(compute_stream):                     graph.replay()             # Otherwise, the graph does not exist, so we create it             else:@@ -162,13 +169,11 @@ def _get_forward_fn(self, use_block_table: bool) -> tuple[Callable, bool]:     def _capture_graph(self, forward_fn: Callable, compute_stream: torch.cuda.Stream, *args) -> None:         """Helper function to capture and store a graph for a given forward function."""         # Warmup (ensures the right result is computed before capturing the graph)-        with torch.cuda.stream(compute_stream), mem_pool_ctx(self.mem_pool):+        with stream_context(compute_stream), memory_pool_context(self.mem_pool):             forward_fn(*args)         # Capture using a thread-local capture mode to avoid capturing GPU operations from outside the model forward-        graph = torch.cuda.CUDAGraph()-        with torch.cuda.graph(-            graph, stream=compute_stream, pool=self.graph_pool_id, capture_error_mode="thread_local"-        ):+        graph = create_accelerator_graph()+        with accelerator_graph_capture_context(graph, stream=compute_stream, pool=self.graph_pool_id):             forward_fn(*args)         # Store         self.inputs_and_outputs.set_graph(graph)

@@ -24,7 +24,6 @@ """  import logging-from contextlib import nullcontext from math import ceil  import torch@@ -35,6 +34,7 @@ from .distributed import DistributedHelper from .requests import FutureRequestState, RequestState, RequestStatus, logger from .scheduler import Scheduler+from .utils import stream_context   def contiguous_runs(indices: list[int]) -> list[tuple[int, int, int]]:@@ -151,7 +151,7 @@ def _compute_num_cpu_sectors(self, cpu_offload_space_gib: float | None, safety_t      def _stream_ctx(self):         """Returns a context manager that runs enclosed ops on the compute stream, or a no-op when none is set."""-        return torch.cuda.stream(self._compute_stream) if self._compute_stream is not None else nullcontext()+        return stream_context(self._compute_stream)      def offload_requests(self) -> bool:         """Evict enough active requests that, at the next batch, every remaining starved request can allocate the

@@ -13,7 +13,8 @@ # limitations under the License. import queue import threading-from contextlib import contextmanager+import warnings+from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from math import ceil, log2 from typing import Any@@ -24,6 +25,10 @@ from .requests import FutureRequestState, RequestState, RequestStatus  +# NOTE: AMD GPUs are supported under the "cuda" module+DEVICE_TYPE_TO_GRAPH_NAME = {"cuda": "CUDAGraph", "xpu": "XPUGraph"}++ class CudaGraphBuffer:     """A dict for CUDA graphs with a special __del__ method to make sure the graphs are properly reset.""" @@ -204,20 +209,6 @@ def drain_queue(request_queue: queue.Queue) -> list[RequestState]:     return new_states  -def get_cuda_pools() -> tuple:-    """Returns a tuple of (mem_pool, graph_pool_id) for CUDA graphs."""-    mem_pool = torch.cuda.MemPool()-    graph_pool_id = mem_pool.id-    return mem_pool, graph_pool_id---@contextmanager-def mem_pool_ctx(mem_pool):-    """A context manager to use a CUDA mem pool."""-    with torch.cuda.use_mem_pool(mem_pool):-        yield-- def find_num_key_value_heads(config: PreTrainedConfig) -> int:     """Finds the number of key-value heads for the given config."""     # If the model supports GQA, we leverage it by using the num_key_value_heads attribute@@ -251,3 +242,107 @@ def exact_div(a: int, b: int) -> int:     if remainder:         raise ValueError(f"Division of {a} by {b} is not exact: {remainder = } != 0")     return quotient+++# ---------------------------------------------- DEVICE AGNOSTIC UTILS ----------------------------------------------- #+++def get_available_accelerator_module() -> Any | None:+    """Returns the accelerator module if there is one and it is available, None otherwise."""+    device_module = torch.get_device_module()  # the device is never specified, so that we always get the same module+    if device_module.__name__.endswith("cpu") or not device_module.is_available():+        return None+    return device_module+++def create_device_stream(device: torch.device) -> torch.cuda.Stream | None:+    """If the given device is available and supports streams, returns a new stream. Otherwise, returns None."""+    device_module = torch.get_device_module()+    if device_module.is_available() and hasattr(device_module, "Stream"):+        return device_module.Stream(device=device)+    return None+++def stream_context(stream: torch.cuda.Stream | None) -> torch.cuda.StreamContext:+    """If the stream is not None, returns a context manager to use the stream. Otherwise, returns a null context.+    This function assumes that the current device supports streams."""+    if stream is None:+        return nullcontext()+    return torch.get_device_module().stream(stream)+++def get_accelerator_graph_name() -> str | None:+    """If an accelerator is available and has cuda graph-like objects, returns their name. Otherwise, returns None."""+    # Stop if no accelerator is available+    accelerator_module = get_available_accelerator_module()+    if accelerator_module is None:+        return None+    # Stop if the accelerator does not support graphs no matter the version (eg. MPS)+    accelerator_type = accelerator_module.__name__.split(".")[-1]+    graph_name = DEVICE_TYPE_TO_GRAPH_NAME.get(accelerator_type)+    if graph_name is None:+        return None+    # Stop if the accelerator does not support the graph for this version (eg. XPU before 2.11)+    return graph_name if hasattr(accelerator_module, graph_name) else None+++def create_accelerator_graph() -> torch.cuda.CUDAGraph:+    """If an accelerator is available and supports graphs, returns a new graph. Otherwise, raises an error."""+    graph_name = get_accelerator_graph_name()+    if graph_name is None:+        raise RuntimeError(+            f"The current device is unavailable or does not support graphs: {torch.get_device_module().__name__ = }."+        )+    graph_cls = getattr(torch.get_device_module(), graph_name)+    return graph_cls()+++def accelerator_graph_capture_context(+    graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream | None, pool: int+) -> torch.cuda.graph:+    """Returns a context manager used to capture the graph. This function assumes that the current device supports+    graphs."""+    device_module = torch.get_device_module()+    capture_kwargs = {"stream": stream, "pool": pool}+    if device_module.__name__.endswith("cuda"):+        capture_kwargs["capture_error_mode"] = "thread_local"+    return device_module.graph(graph, **capture_kwargs)+++def get_memory_pools() -> tuple[Any, int]:+    """If an accelerator is available and supports graphs, returns a tuple of (mem_pool, graph_pool_id) for the+    accelerator graphs. Otherwise, raises an error."""+    if get_accelerator_graph_name() is None:+        raise RuntimeError(+            "No graph-capture backend is available for the current accelerator: memory pools cannot be used."+        )+    mem_pool = torch.get_device_module().MemPool()+    graph_pool_id = mem_pool.id+    return mem_pool, graph_pool_id+++def memory_pool_context(mem_pool: Any) -> AbstractContextManager:+    """A context manager to use a specific memory pool for allocations. This function assumes that the current device+    supports memory pools."""+    return torch.get_device_module().use_mem_pool(mem_pool)+++# ------------------------------------------------- DEPRECATED UTILS ------------------------------------------------- #+++def get_cuda_pools() -> tuple:+    warnings.warn(+        "Deprecated. Use get_memory_pools() instead. Will be removed in transformers 5.23",+        FutureWarning,+        stacklevel=2,+    )+    return get_memory_pools()+++def mem_pool_ctx(mem_pool: Any):+    warnings.warn(+        "Deprecated. Use memory_pool_context() instead. Will be removed in transformers 5.23",+        FutureWarning,+        stacklevel=2,+    )+    return memory_pool_context(mem_pool)

@@ -64,6 +64,7 @@     RequestStatus,     get_device_and_memory_breakdown, )+from transformers.generation.continuous_batching.utils import DEVICE_TYPE_TO_GRAPH_NAME from transformers.integrations.eager_paged import eager_paged_attention_forward from transformers.testing_utils import (     Expectations,@@ -1226,9 +1227,11 @@ def _test_continuous_batching_parity(          # Skip the test if Flash Attention 2 or 3 is required but not available.         flexible_flash_skip(self, attn_implementation)-        # Skip the test if cuda graph is on but the device is not CUDA-        if continuous_batching_config.use_cuda_graph and torch_device != "cuda":-            self.skipTest("CUDA graph is only supported on CUDA devices. Skipping test.")+        # Skip the test if cuda graph is on but the device does not support graph capture+        device_type = torch_device.type if isinstance(torch_device, torch.device) else torch_device+        if any(continuous_batching_config.cuda_graph_booleans) and device_type not in DEVICE_TYPE_TO_GRAPH_NAME:+            supported = list(DEVICE_TYPE_TO_GRAPH_NAME.keys())+            self.skipTest(f"CUDA graph is only supported on {supported}, but {device_type = }. Skipping test.")          # If the config turns on compile, change the generation config to use the default mode instead of         # max-autotune-no-cudagraphs which can change the kernels between generate_batch and generate@@ -1281,7 +1284,9 @@ def _test_continuous_batching_parity(         model.generation_config.max_new_tokens = max_new_tokens         model.generation_config.do_sample = False -        model.generation_config.use_cuda_graph = continuous_batching_config.use_cuda_graph+        model.generation_config.use_cuda_graph = (+            any(continuous_batching_config.cuda_graph_booleans) and device_type in DEVICE_TYPE_TO_GRAPH_NAME+        )         model.generation_config.compile_config = continuous_batching_config.varlen_compile_config         # Create a static cache if compile_config is set, because regular generate requires a compileable cache         past_key_values = None