Spaces:
Running on Zero
Running on Zero
File size: 2,560 Bytes
414b4fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """FlashAttention-2 backend for PaDoc's packed tree visibility pattern."""
from __future__ import annotations
import torch
# Transformers treats custom names containing "flash_attention" as Hub kernels.
ATTN_IMPLEMENTATION = "padoc_tree_fa2"
def tree_flash_attention_forward(
module: torch.nn.Module,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: torch.Tensor | None,
dropout: float = 0.0,
scaling: float | None = None,
*,
tree_q_indices: torch.LongTensor,
tree_kv_indices: torch.LongTensor,
tree_cu_seqlens_q: torch.IntTensor,
tree_cu_seqlens_kv: torch.IntTensor,
tree_max_seqlen_q: int | torch.Tensor,
tree_max_seqlen_kv: int | torch.Tensor,
**kwargs,
) -> tuple[torch.Tensor, None]:
"""Run exact tree attention as a virtual FlashAttention varlen batch."""
del module, attention_mask, kwargs
if tree_cu_seqlens_q.dtype != torch.int32 or tree_cu_seqlens_kv.dtype != torch.int32:
raise TypeError("FlashAttention cumulative sequence lengths must be torch.int32.")
try:
from flash_attn import flash_attn_varlen_func
except ImportError as exc: # pragma: no cover - requires a CUDA environment
raise ImportError(
f"{ATTN_IMPLEMENTATION} requires the optional 'flash-attn' package."
) from exc
batch_size, num_heads, sequence_length, head_dim = query.shape
num_kv_heads = key.shape[1]
query_flat = query.transpose(1, 2).reshape(batch_size * sequence_length, num_heads, head_dim)
key_flat = key.transpose(1, 2).reshape(batch_size * sequence_length, num_kv_heads, head_dim)
value_flat = value.transpose(1, 2).reshape(batch_size * sequence_length, num_kv_heads, head_dim)
output_packed = flash_attn_varlen_func(
query_flat.index_select(0, tree_q_indices),
key_flat.index_select(0, tree_kv_indices),
value_flat.index_select(0, tree_kv_indices),
tree_cu_seqlens_q,
tree_cu_seqlens_kv,
int(tree_max_seqlen_q),
int(tree_max_seqlen_kv),
dropout_p=dropout,
softmax_scale=scaling,
causal=True,
)
output_flat = torch.zeros_like(query_flat)
output_flat.index_copy_(0, tree_q_indices, output_packed)
return output_flat.view(batch_size, sequence_length, num_heads, head_dim), None
def register_tree_flash_attention() -> None:
from transformers import AttentionInterface
AttentionInterface.register(ATTN_IMPLEMENTATION, tree_flash_attention_forward)
register_tree_flash_attention()
|