Repository: unslothai/unsloth
File: unsloth/models/llama.py (Python)

Code before:
            else:
                Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
        pass
        A = xformers_attention(Q, K, V, attn_bias = causal_mask)
        A = A.view(bsz, q_len, n_heads, head_dim)

    elif HAS_FLASH_ATTENTION and attention_mask is None:
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        A = flash_attn_func(Q, K, V, causal = True)
    else:
        # when qlen==vlen and attn_mask is None, we should use causal attention
        Q_len = Q.shape[-2]
        K_len = K.shape[-2]
        if attention_mask is None and Q_len == K_len:
            is_causal = True
        else:
            is_causal = False
        # Grouped query attention
        if SDPA_HAS_GQA:
            # Needs (batch_size, n_heads, seq_len, head_dim)
            # is_casual and attention_mask must not be both set!
            A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False, enable_gqa = n_groups != 1)
            # Go back to (batch_size, seq_len, n_heads, head_dim)
            A = A.transpose(1, 2)#.contiguous()
        else:
            if n_groups != 1:
                K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
                V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
                K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
                V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
            pass
            # Must be contiguous or else results are False!
            # https://github.com/pytorch/pytorch/issues/112577
            Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
            # Needs (batch_size, n_heads, seq_len, head_dim)
            # is_casual and attention_mask must not be both set!
            A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
            # Go back to (batch_size, seq_len, n_heads, head_dim)
            A = A.transpose(1, 2).contiguous()
        pass
    pass
    attn_output = A.reshape(bsz, q_len, n_heads*head_dim)
    attn_output = self.apply_o(self, attn_output)
    attn_weights = None
    return attn_output, attn_weights, past_key_value
pass


# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590

Review comment (suggestion):
May also need to change the usage of the `is_causal` argument in `scaled_dot_product_attention` function instead of `is_causal=False`, such as here: https://github.com/unslothai/unsloth/pull/2868/files#diff-a45b72bb533eda979990bd79cde5fe9c9fde424779a4f1fc1195b75853d93b45R532

Code after:
            K = K  .view(bsz, kv_seq_len, n_kv_heads,        1, head_dim)
            V = V  .view(bsz, kv_seq_len, n_kv_heads,        1, head_dim)
            K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
            V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
            if hidden_states.requires_grad:
                K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
                V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
            else:
                Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
        pass
        A = xformers_attention(Q, K, V, attn_bias = causal_mask)
        A = A.view(bsz, q_len, n_heads, head_dim)

    elif HAS_FLASH_ATTENTION and attention_mask is None:
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        A = flash_attn_func(Q, K, V, causal = True)
    else:
        # when qlen==vlen and attn_mask is None, we should use causal attention
        Q_len = Q.shape[-2]
        K_len = K.shape[-2]
        if attention_mask is None and Q_len == K_len:
            is_causal = True
        else:
            is_causal = False
        # Grouped query attention
        if SDPA_HAS_GQA:
            # Needs (batch_size, n_heads, seq_len, head_dim)
            # is_casual and attention_mask must not be both set!
            A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal, enable_gqa = n_groups != 1)
            # Go back to (batch_size, seq_len, n_heads, head_dim)
            A = A.transpose(1, 2)#.contiguous()
        else:
            if n_groups != 1:
                K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
                V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
                K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
                V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
            pass
            # Must be contiguous or else results are False!
            # https://github.com/pytorch/pytorch/issues/112577
            Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
            # Needs (batch_size, n_heads, seq_len, head_dim)
            # is_casual and attention_mask must not be both set!
            A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = is_causal)
            # Go back to (batch_size, seq_len, n_heads, head_dim)
            A = A.transpose(1, 2).contiguous()
        pass
    pass
    attn_output = A.reshape(bsz, q_len, n_heads*head_dim)