From 2229d1b620131e7f6ec5c19eb5a20e96dd2d7a22 Mon Sep 17 00:00:00 2001 From: beiyonder <86228410+beiyonder@users.noreply.github.com> Date: Thu, 21 May 2026 23:48:40 +0530 Subject: [PATCH] Add SDPA fallback for FlashAttention --- models/flash_attention_prefixlm_v2.py | 112 +++++++++++++++++++++++++- models/layers.py | 3 +- 2 files changed, 112 insertions(+), 3 deletions(-) diff --git a/models/flash_attention_prefixlm_v2.py b/models/flash_attention_prefixlm_v2.py index b7856ee..9ab8830 100644 --- a/models/flash_attention_prefixlm_v2.py +++ b/models/flash_attention_prefixlm_v2.py @@ -1,10 +1,19 @@ from typing import Tuple, Optional import torch +import torch.nn.functional as F from torch import Tensor import numpy as np -from flash_attn_interface import _flash_attn_backward, maybe_contiguous +try: + from flash_attn_interface import _flash_attn_backward, maybe_contiguous + FLASH_ATTN_AVAILABLE = True +except (ImportError, OSError): + _flash_attn_backward = None + FLASH_ATTN_AVAILABLE = False + + def maybe_contiguous(x): + return x.contiguous() if x is not None and not x.is_contiguous() else x def compute_aux_seq_tensors_scalars(prefix_lens: np.ndarray, causal_lens: np.ndarray, batch_max_tokens: int): @@ -85,6 +94,103 @@ def _custom_flash_attn_forward( return out, softmax_lse +def _torch_sdpa(q: Tensor, k: Tensor, v: Tensor, attn_mask: Tensor) -> Tensor: + # q/k/v: [seq, heads, head_dim]; SDPA expects [batch, heads, seq, head_dim]. + return F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), + attn_mask=attn_mask.unsqueeze(0).unsqueeze(0), + dropout_p=0.0, + is_causal=False, + ).squeeze(0).transpose(0, 1) + + +def _prefixlm_mask(seq_len: int, prefix_len: int, is_causal: bool, device: torch.device) -> Tensor: + query_pos = torch.arange(seq_len, device=device).unsqueeze(1) + key_pos = torch.arange(seq_len, device=device).unsqueeze(0) + if is_causal: + return key_pos <= query_pos + return torch.where(query_pos < prefix_len, key_pos < prefix_len, key_pos <= query_pos) + + +def _torch_flash_attn_varlen_prefixlm(q: Tensor, + k: Tensor, + v: Tensor, + is_causal: bool, + prefix_lens: Tensor, + causal_lens: Tensor, + cu_seqlens: Tensor, + total_seqlen: Tensor, + numseqs: Tensor, + max_seqlen_prefix: Tensor, + max_seqlen_causal: Tensor, + max_seqlen_all: Tensor) -> Tensor: + del max_seqlen_prefix, max_seqlen_causal, max_seqlen_all + out = torch.zeros_like(q) + numseqs_int = int(numseqs.item()) + total_seqlen_int = int(total_seqlen.item()) + + cu = cu_seqlens[:numseqs_int + 1].to(device="cpu", dtype=torch.int64) + prefixes = prefix_lens[:numseqs_int].to(device="cpu", dtype=torch.int64) + causals = causal_lens[:numseqs_int].to(device="cpu", dtype=torch.int64) + + for i in range(numseqs_int): + start = int(cu[i].item()) + seq_len = int((prefixes[i] + causals[i]).item()) + prefix_len = int(prefixes[i].item()) + if seq_len == 0: + continue + + end = start + seq_len + mask = _prefixlm_mask(seq_len, prefix_len, is_causal, q.device) + out[start:end] = _torch_sdpa(q[start:end], k[start:end], v[start:end], mask) + + if total_seqlen_int < out.shape[0]: + out[total_seqlen_int:] = 0 + return out + + +def flash_attn_with_kvcache(q: Tensor, + k: Tensor, + v: Tensor, + k_cache: Tensor, + v_cache: Tensor, + cache_seqlens: Tensor | int, + causal: bool, + **kwargs) -> Tensor: + if FLASH_ATTN_AVAILABLE: + from flash_attn_interface import flash_attn_with_kvcache as _flash_attn_with_kvcache + return _flash_attn_with_kvcache(q=q, k=k, v=v, k_cache=k_cache, v_cache=v_cache, + cache_seqlens=cache_seqlens, causal=causal, **kwargs) + + del kwargs + out = torch.empty_like(q) + batch_size, query_len = q.shape[:2] + if isinstance(cache_seqlens, int): + lengths = torch.full((batch_size,), cache_seqlens, device="cpu", dtype=torch.int64) + else: + lengths = cache_seqlens.detach().to(device="cpu", dtype=torch.int64) + + for b in range(batch_size): + start = int(lengths[b].item()) + end = start + query_len + k_cache[b, start:end] = k[b] + v_cache[b, start:end] = v[b] + + keys = k_cache[b, :end] + values = v_cache[b, :end] + if causal: + query_pos = torch.arange(start, end, device=q.device).unsqueeze(1) + key_pos = torch.arange(end, device=q.device).unsqueeze(0) + mask = key_pos <= query_pos + else: + mask = torch.ones((query_len, end), device=q.device, dtype=torch.bool) + out[b] = _torch_sdpa(q[b], keys, values, mask) + + return out + + @torch.library.custom_op("flash_attn::flash_attn_varlen_prefixlm_compileop", mutates_args=()) def flash_attn_varlen_prefixlm_compileop( q: Tensor, @@ -276,6 +382,10 @@ def flash_attn_varlen_prefixlm(q: Tensor, prefix_lens: Tensor, causal_lens: Tensor, cu_seqlens: Tensor, # CPU tensors (scalars) total_seqlen: Tensor, numseqs: Tensor, max_seqlen_prefix: Tensor, max_seqlen_causal: Tensor, max_seqlen_all: Tensor): + if not FLASH_ATTN_AVAILABLE: + return _torch_flash_attn_varlen_prefixlm(q, k, v, is_causal, + prefix_lens, causal_lens, cu_seqlens, + total_seqlen, numseqs, max_seqlen_prefix, max_seqlen_causal, max_seqlen_all) # Apply function return FlashAttnVarlenPrefixLM.apply(q, k, v, is_causal, prefix_lens, causal_lens, cu_seqlens, diff --git a/models/layers.py b/models/layers.py index 2159964..610d527 100644 --- a/models/layers.py +++ b/models/layers.py @@ -7,8 +7,7 @@ from einops import rearrange from models.common import trunc_normal_init_, unwrap_tensor -from models.flash_attention_prefixlm_v2 import flash_attn_varlen_prefixlm -from flash_attn_interface import flash_attn_with_kvcache +from models.flash_attention_prefixlm_v2 import flash_attn_varlen_prefixlm, flash_attn_with_kvcache Carry = dict[str, Any]