25 lines
645 B
Python
25 lines
645 B
Python
#sage_attn.py
|
|
import os
|
|
from typing import Optional
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from sageattention import sageattn
|
|
import math
|
|
|
|
__all__ = ['SageAttention', 'sage_attention']
|
|
|
|
|
|
def enable_sage_attention():
|
|
"""
|
|
Enable SageAttention by replacing PyTorch's scaled_dot_product_attention
|
|
with sageattn from the SageAttention library.
|
|
"""
|
|
F.scaled_dot_product_attention = sageattn
|
|
return True
|
|
|
|
def disable_sage_attention():
|
|
"""
|
|
Restore PyTorch's original scaled_dot_product_attention function.
|
|
"""
|
|
F.scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
|
|
return True |