monkeypatch.attention.sage_attn

monkeypatch.attention.sage_attn

Monkeypatch for SageAttention for use with transformers.

https://github.com/thu-ml/SageAttention/

Functions

Name	Description
patch_sageattn	Validate SageAttention is available. Registration in the attention/mask
sage_attention_forward	Forward pass for SageAttention compatible with transformers attention interfaces.

patch_sageattn

monkeypatch.attention.sage_attn.patch_sageattn()

Validate SageAttention is available. Registration in the attention/mask function registries is handled by register_sage_attn() in init.py.

sage_attention_forward

monkeypatch.attention.sage_attn.sage_attention_forward(
    module,
    query,
    key,
    value,
    attention_mask=None,
    dropout=0.0,
    scaling=None,
    is_causal=None,
    **kwargs,
)

Forward pass for SageAttention compatible with transformers attention interfaces.

https://github.com/thu-ml/SageAttention/