From 74d7fc65c6539dc575af691485db43df2116f362 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 4 Feb 2024 02:03:23 +1100 Subject: [PATCH] SDPA --- unsloth/models/llama.py | 1 + unsloth/models/mistral.py | 1 + 2 files changed, 2 insertions(+) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8671d4e022..f24ecafbe3 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -309,6 +309,7 @@ def LlamaAttention_fast_forward( V = V.transpose(1, 2) A = flash_attn_func(Q, K, V, causal = True) else: + print("0", end = "") # Grouped query attention if n_groups != 1: K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index e0246f82b3..4cfeb4ad98 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -137,6 +137,7 @@ def MistralAttention_fast_forward( window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw) A = flash_attn_func(Q, K, V, causal = True, window_size = window) else: + print("0", end = "") # Grouped query attention # if n_groups != 1: K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)