1 files changed, 85 insertions, 6 deletions
diff --git a/modules/xpu_specific.py b/modules/xpu_specific.py
index d8da94a0..2971dbc3 100644
--- a/modules/xpu_specific.py
+++ b/modules/xpu_specific.py
@@ -27,11 +27,90 @@ def torch_xpu_gc():
 
 has_xpu = check_for_xpu()
 
+
+# Arc GPU cannot allocate a single block larger than 4GB: https://github.com/intel/compute-runtime/issues/627
+# Here we implement a slicing algorithm to split large batch size into smaller chunks,
+# so that SDPA of each chunk wouldn't require any allocation larger than ARC_SINGLE_ALLOCATION_LIMIT.
+# The heuristic limit (TOTAL_VRAM // 8) is tuned for Intel Arc A770 16G and Arc A750 8G,
+# which is the best trade-off between VRAM usage and performance.
+ARC_SINGLE_ALLOCATION_LIMIT = {}
+orig_sdp_attn_func = torch.nn.functional.scaled_dot_product_attention
+def torch_xpu_scaled_dot_product_attention(
+    query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, *args, **kwargs
+):
+    # cast to same dtype first
+    key = key.to(query.dtype)
+    value = value.to(query.dtype)
+    if attn_mask is not None and attn_mask.dtype != torch.bool:
+        attn_mask = attn_mask.to(query.dtype)
+
+    N = query.shape[:-2]  # Batch size
+    L = query.size(-2)  # Target sequence length
+    E = query.size(-1)  # Embedding dimension of the query and key
+    S = key.size(-2)  # Source sequence length
+    Ev = value.size(-1)  # Embedding dimension of the value
+
+    total_batch_size = torch.numel(torch.empty(N))
+    device_id = query.device.index
+    if device_id not in ARC_SINGLE_ALLOCATION_LIMIT:
+        ARC_SINGLE_ALLOCATION_LIMIT[device_id] = min(torch.xpu.get_device_properties(device_id).total_memory // 8, 4 * 1024 * 1024 * 1024)
+    batch_size_limit = max(1, ARC_SINGLE_ALLOCATION_LIMIT[device_id] // (L * S * query.element_size()))
+
+    if total_batch_size <= batch_size_limit:
+        return orig_sdp_attn_func(
+            query,
+            key,
+            value,
+            attn_mask,
+            dropout_p,
+            is_causal,
+            *args, **kwargs
+        )
+
+    query = torch.reshape(query, (-1, L, E))
+    key = torch.reshape(key, (-1, S, E))
+    value = torch.reshape(value, (-1, S, Ev))
+    if attn_mask is not None:
+        attn_mask = attn_mask.view(-1, L, S)
+    chunk_count = (total_batch_size + batch_size_limit - 1) // batch_size_limit
+    outputs = []
+    for i in range(chunk_count):
+        attn_mask_chunk = (
+            None
+            if attn_mask is None
+            else attn_mask[i * batch_size_limit : (i + 1) * batch_size_limit, :, :]
+        )
+        chunk_output = orig_sdp_attn_func(
+            query[i * batch_size_limit : (i + 1) * batch_size_limit, :, :],
+            key[i * batch_size_limit : (i + 1) * batch_size_limit, :, :],
+            value[i * batch_size_limit : (i + 1) * batch_size_limit, :, :],
+            attn_mask_chunk,
+            dropout_p,
+            is_causal,
+            *args, **kwargs
+        )
+        outputs.append(chunk_output)
+    result = torch.cat(outputs, dim=0)
+    return torch.reshape(result, (*N, L, Ev))
+
+
+def is_xpu_device(device: str | torch.device = None):
+    if device is None:
+        return False
+    if isinstance(device, str):
+        return device.startswith("xpu")
+    return device.type == "xpu"
+
+
 if has_xpu:
-    # W/A for https://github.com/intel/intel-extension-for-pytorch/issues/452: torch.Generator API doesn't support XPU device
-    CondFunc('torch.Generator',
-        lambda orig_func, device=None: torch.xpu.Generator(device),
-        lambda orig_func, device=None: device is not None and device.type == "xpu")
+    try:
+        # torch.Generator supports "xpu" device since 2.1
+        torch.Generator("xpu")
+    except RuntimeError:
+        # W/A for https://github.com/intel/intel-extension-for-pytorch/issues/452: torch.Generator API doesn't support XPU device (for torch < 2.1)
+        CondFunc('torch.Generator',
+            lambda orig_func, device=None: torch.xpu.Generator(device),
+            lambda orig_func, device=None: is_xpu_device(device))
 
     # W/A for some OPs that could not handle different input dtypes
     CondFunc('torch.nn.functional.layer_norm',
@@ -55,5 +134,5 @@ if has_xpu:
         lambda orig_func, tensors, dim=0, out=None: orig_func([t.to(tensors[0].dtype) for t in tensors], dim=dim, out=out),
         lambda orig_func, tensors, dim=0, out=None: not all(t.dtype == tensors[0].dtype for t in tensors))
     CondFunc('torch.nn.functional.scaled_dot_product_attention',
-        lambda orig_func, query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False: orig_func(query, key.to(query.dtype), value.to(query.dtype), attn_mask, dropout_p, is_causal),
-        lambda orig_func, query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False: query.dtype != key.dtype or query.dtype != value.dtype)
+        lambda orig_func, *args, **kwargs: torch_xpu_scaled_dot_product_attention(*args, **kwargs),
+        lambda orig_func, query, *args, **kwargs: query.is_xpu)