Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Make IPEX dyn atten care only about the sdpa memory usage

Browse Source
This commit is contained in:
Disty0
2025-01-26 04:45:05 +03:00
parent bb07dd7a8f
commit a21bc9d9cd
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/intel/ipex/attention.py
+4 -4
View File
@@ -7,7 +7,7 @@ from functools import cache, wraps
# ARC GPUs can't allocate more than 4GB to a single block so we slice the attetion layers
sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 6))
sdpa_slice_trigger_rate = float(os.environ.get('IPEX_SDPA_SLICE_TRIGGER_RATE', 4))
attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4))
# Find something divisible with the input_tokens
@@ -38,7 +38,7 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz
_, _, key_len, _ = key_shape
_, _, _, head_dim = value_shape
slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2
slice_batch_size = attn_heads * math.sqrt(query_len * key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024
split_batch_size = batch_size
split_head_size = attn_heads
@@ -53,12 +53,12 @@ def find_sdpa_slice_sizes(query_shape, key_shape, value_shape, query_element_siz
split_batch_size = find_split_size(split_batch_size, slice_batch_size)
if split_batch_size * slice_batch_size > attention_slice_rate:
slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * head_dim * query_element_size / 1024 / 1024 / 2
slice_head_size = split_batch_size * math.sqrt(query_len * key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024
do_head_split = True
split_head_size = find_split_size(split_head_size, slice_head_size)
if split_batch_size * slice_batch_size > attention_slice_rate:
slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * head_dim * query_element_size / 1024 / 1024 / 2
slice_query_size = split_batch_size * attn_heads * math.sqrt(key_len) * math.sqrt(head_dim) * 2 * query_element_size / 1024 / 1024
do_query_split = True
split_query_size = find_query_size(split_query_size, slice_query_size)