Support efficient flash attention for packed sequences using flash-attn-2.0

pyproject.toml

@@ -40,6 +40,10 @@ examples = [
     "tqdm",
 ]
 
+flash-pack-seq = [
+  "flash-attn>=2.0",
+]
+
 test = [
     "pytest",
 ]

tests/test_x_transformers.py

@@ -11,6 +11,9 @@
     Encoder,
     Decoder,
     LinearNoBias,
+    AttentionLayers,
+    Attention,
+    Attend
 )
 
 from x_transformers.neo_mlp import (
@@ -1569,6 +1572,137 @@ def test_seq_start_pos_parity():
 
     assert torch.allclose(parallel_logits[is_not_masked], seq_logits[is_not_masked], atol = 1e-5)
 
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or \
+    torch.cuda.get_device_capability()[0] < 8 or \
+    __import__('importlib').util.find_spec('flash_attn') is None,
+    reason="CUDA compute capability must be >= 8 and flash_attn must be installed"
+)
+@param('exp', (
+    dict(causal=True, same_partition=True, pos_enc='rotary_pos_emb'),
+    dict(causal=False, same_partition=True, pos_enc='rotary_pos_emb'),
+    dict(causal=False, same_partition=False, pos_enc='rotary_pos_emb'),
+    dict(causal=True, same_partition=True, pos_enc='rotary_xpos'))
+)
+def test_flash_pack_seq(exp):
+    seq_len = 1024
+    dim = 256
+    n_part = 4
+    n_layers = 4
+    causal = exp['causal']
+    same_partition = exp.get('same_partition', False)
+    atl_kwargs = {exp['pos_enc']: True}
+    mem_len = 128 if not same_partition else seq_len
+    x = torch.randn((seq_len, dim)).cuda().half()
+    mem = torch.randn((mem_len, dim)).cuda().half()
+
+    pad_val = 99.0 # float('-inf')
+    def partition(x, num_parts):
+        total = x.shape[0]
+        split_points = sorted(torch.randint(1, total, (num_parts - 1,)).tolist())
+        splits = torch.tensor_split(x, split_points, dim=0)
+        import numpy as np
+        attn_cu_lengths = torch.tensor([0] + np.cumsum([split.shape[0] for split in splits]).tolist()).int().cuda()
+        return splits, attn_cu_lengths
+    with torch.no_grad():
+        with torch.autocast("cuda", dtype=torch.float16):        
+            if same_partition and causal:
+                splits, attn_cu_lengths = partition(x, n_part)
+                # Split mem using the cutting points from attn_cu_lengths
+                split_points = attn_cu_lengths[1:-1].tolist()
+                splits_mem = torch.tensor_split(mem, split_points, dim=0)
+                attn_cu_lengths_context = attn_cu_lengths
+                padded_batch = torch.nn.utils.rnn.pad_sequence(splits, batch_first=True, padding_value=pad_val)
+                padded_mem = torch.nn.utils.rnn.pad_sequence(splits_mem, batch_first=True, padding_value=pad_val)
+
+                mask = (padded_batch != pad_val).any(dim=-1)
+                context_mask = (padded_mem != pad_val).any(dim=-1)
+                max_len = padded_batch.shape[1]
+                attn_mask = mask.unsqueeze(1) & torch.tril(torch.ones((max_len, max_len), dtype=torch.bool)).cuda().unsqueeze(0)
+                attn_mask = attn_mask.unsqueeze(1)
+            else:
+                splits, attn_cu_lengths = partition(x, n_part)
+                splits_mem, attn_cu_lengths_context = partition(mem, n_part)
+                padded_batch = torch.nn.utils.rnn.pad_sequence(splits, batch_first=True, padding_value=pad_val)
+                padded_mem = torch.nn.utils.rnn.pad_sequence(splits_mem, batch_first=True, padding_value=pad_val)
+
+                mask = (padded_batch != pad_val).any(dim=-1)
+                context_mask = (padded_mem != pad_val).any(dim=-1)
+                attn_mask = mask.unsqueeze(2) & context_mask.unsqueeze(1)
+                attn_mask = attn_mask.unsqueeze(1) 
+
+
+            # Standard padding            
+            reset_exp_det()
+            atd = Attend(flash = False, flash_pack_seq = False,causal=causal).cuda().eval()
+            o_atd = atd(q=padded_batch[:,None], k=padded_mem[:,None], v=padded_mem[:,None], mask=attn_mask)
+            o_atd = o_atd[0][:,0]
+            o_atd = torch.cat([o[~(m == pad_val).all(-1)] for o, m in zip(o_atd, padded_batch)], dim=0)
+
+            att=Attention(dim=dim,flash=False,causal=causal).cuda().eval()
+            o_att = att(
+                x=padded_batch,
+                context=padded_mem,
+                attn_mask = attn_mask
+            )
+            o_att = torch.cat([o[~(m == pad_val).all(-1)] for o, m in zip(o_att, padded_batch)], dim=0)
+
+            atl = AttentionLayers(dim=dim, depth=n_layers, cross_attend=True, causal=causal, attn_flash=True, attn_flash_pack_seq=False, **atl_kwargs).cuda().eval()
+            o_atl = atl(
+                x=padded_batch,
+                context=padded_mem,
+                context_mask = context_mask,
+                mask = mask,
+            )
+            o_atl = torch.cat([o[~(m == pad_val).all(-1)] for o, m in zip(o_atl, padded_batch)], dim=0)
+
+
+            # Block masking
+            reset_exp_det()
+            atd_block = Attend(flash = True, flash_pack_seq = True, causal=causal).cuda().eval()        
+
+            flash_pack_seq_kwargs = dict(
+                cu_seqlens_q=attn_cu_lengths,
+                max_seqlen_q = attn_cu_lengths.diff().max().item(),
+                cu_seqlens_k=attn_cu_lengths,
+                max_seqlen_k = attn_cu_lengths.diff().max().item()
+            )
+            flash_pack_seq_kwargs_context = dict(
+                cu_seqlens_q=attn_cu_lengths,
+                max_seqlen_q = attn_cu_lengths.diff().max().item(),
+                cu_seqlens_k=attn_cu_lengths_context,
+                max_seqlen_k = attn_cu_lengths_context.diff().max().item()
+            )            
+            o_atd_block = atd_block(x[None,None], mem[None,None], mem[None,None], flash_pack_seq_kwargs=flash_pack_seq_kwargs_context)
+            o_atd_block = o_atd_block[0][0,0]
+
+
+            att_block=Attention(dim=dim,flash=True,flash_pack_seq=True, causal=causal).cuda().eval()
+            o_att_block = att_block(
+                x = x.unsqueeze(0),
+                context = mem.unsqueeze(0),
+                flash_pack_seq_kwargs=flash_pack_seq_kwargs_context
+            )[0]
+
+            atl_block = AttentionLayers(dim=dim, depth=n_layers, cross_attend=True, causal=causal, attn_flash=True, attn_flash_pack_seq=True, **atl_kwargs).cuda().eval()
+            o_atl_block = atl_block(
+                x = x.unsqueeze(0),
+                context = mem.unsqueeze(0),
+                flash_pack_seq_kwargs=flash_pack_seq_kwargs,
+                flash_pack_seq_context_kwargs=flash_pack_seq_kwargs_context,
+            )[0]
+            torch.testing.assert_close(o_atd, o_atd_block , atol=5e-3, rtol=5e-3)
+            torch.testing.assert_close(o_att, o_att_block , atol=5e-3, rtol=5e-3)
+            torch.testing.assert_close(o_atl, o_atl_block , atol=5e-3, rtol=5e-3)
+
+
+def reset_exp_det():
+    # pass
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
 @param('pos_emb_type', ('rotary', 'polar'))
 def test_pos_emb_parity(pos_emb_type):
     pos_emb_kwargs = {f'{pos_emb_type}_pos_emb': True}

x_transformers/attend.py

@@ -198,7 +198,8 @@ def __init__(
             enable_flash = True,
             enable_math = True,
             enable_mem_efficient = True
-        )
+        ),
+        flash_pack_seq = False,          # efficient flash attention packed sequence masking for variable length sequences.
     ):
         super().__init__()
         self.scale = scale
@@ -299,13 +300,22 @@ def __init__(
         # flash attention
 
         self.flash = flash
-
+        self.flash_pack_seq = flash_pack_seq
+
         torch_version = version.parse(torch.__version__)
         assert not (flash and torch_version < version.parse('2.0.0')), 'in order to use flash attention, you must be using pytorch 2.0 or above'
 
-        # torch 2.3 uses new backend and context manager
-
-        if self.flash:
+        if self.flash:            
+            if self.flash_pack_seq:
+                try:
+                    from flash_attn import flash_attn_varlen_func
+                    self.flash_attn_varlen_func = flash_attn_varlen_func
+                except ImportError:
+                    raise ImportError("block masking with Flash Attention requires the flash-attn package. Please install it with `pip install flash-attn`.")
+                major, minor = torch.cuda.get_device_capability()
+                assert major >= 8, f"block masking with Flash Attention requires SM80+ (Ampere or newer) GPUs, but your GPU has SM{major}{minor}."
+
+            # torch 2.3 uses new backend and context manager
             if torch_version >= version.parse('2.3'):
                 from torch.nn.attention import SDPBackend
 
@@ -326,7 +336,8 @@ def flash_attn(
         self,
         q, k, v,
         mask = None,
-        attn_bias = None
+        attn_bias = None,
+        flash_pack_seq_kwargs = None, # https://github.com/Dao-AILab/flash-attention/blob/v2.8.3/flash_attn/flash_attn_interface.py#L1370
     ):
         batch, heads, q_len, _, k_len, is_cuda, device = *q.shape, k.shape[-2], q.is_cuda, q.device
 
@@ -421,14 +432,32 @@ def flash_attn(
             mask = attn_bias
 
         # pytorch 2.0 flash attn: q, k, v, mask, dropout, causal, softmax_scale
-
-        with self.sdp_context_manager():
-            out = F.scaled_dot_product_attention(
-                q, k, v,
-                attn_mask = mask,
-                dropout_p = self.dropout if self.training else 0., 
-                is_causal = causal
+        if self.flash_pack_seq:
+            assert exists(flash_pack_seq_kwargs), "flash_pack_seq_kwargs must be provided when self.flash_pack_seq is True"
+            cu_seqlens_k = flash_pack_seq_kwargs.get('cu_seqlens_k', None)
+            cu_seqlens_q = flash_pack_seq_kwargs.get('cu_seqlens_q', None)
+            assert q.shape[0] == 1 and k.shape[0] == 1 and v.shape[0] == 1, f"batch size must be 1 for block masking. Shape was q={q.shape}, k={k.shape}, v={v.shape}"
+            assert not exists(mask) and exists(cu_seqlens_q) and exists(cu_seqlens_k), "mask cannot be passed with cu_seqlens for block masking"
+            assert cu_seqlens_q.shape == cu_seqlens_k.shape and cu_seqlens_q.ndim == 1 and not cu_seqlens_q.is_floating_point() and not cu_seqlens_k.is_floating_point() and (cu_seqlens_q.diff() > 0).all() and (cu_seqlens_k.diff() > 0).all(), "cu_seqlens_q/k should be same-length 1D cumulative sequence lengths for block masking"
+            assert not causal or (cu_seqlens_q == cu_seqlens_k).all(), "causal attention with different cu_seqlens for q and k not supported"
+            # efficient packed sequences flash attentino masking
+            att = self.flash_attn_varlen_func(
+                q = rearrange(q, '1 h t d ->t h d'),
+                k = rearrange(k, '1 h t d ->t h d'),
+                v = rearrange(v, '1 h t d ->t h d'),
+                causal=causal,
+                dropout_p=self.dropout if self.training else 0.,
+                **flash_pack_seq_kwargs
             )
+            out = rearrange(att, 't h d -> 1 h t d')
+        else:
+            with self.sdp_context_manager():
+                out = F.scaled_dot_product_attention(
+                    q, k, v,
+                    attn_mask = mask,
+                    dropout_p = self.dropout if self.training else 0., 
+                    is_causal = causal
+                )
 
         # for a row that is entirely masked out, should zero out the output of that row token
 
@@ -442,7 +471,8 @@ def forward(
         q, k, v,
         mask = None,
         attn_bias = None,
-        prev_attn = None
+        prev_attn = None,
+        flash_pack_seq_kwargs = None,
     ):
         """
         einstein notation
@@ -488,7 +518,7 @@ def forward(
 
         if self.flash:
             assert not exists(prev_attn), 'residual attention not compatible with flash attention'
-            return self.flash_attn(q, k, v, mask = mask, attn_bias = attn_bias)
+            return self.flash_attn(q, k, v, mask = mask, attn_bias = attn_bias, flash_pack_seq_kwargs = flash_pack_seq_kwargs)
 
         kv_einsum_eq = 'b j d' if k.ndim == 3 else 'b h j d'
 

x_transformers/x_transformers.py

@@ -1513,7 +1513,8 @@ def __init__(
             enable_flash = True,
             enable_math = True,
             enable_mem_efficient = True
-        )
+        ),
+        flash_pack_seq = False,
     ):
         super().__init__()
         dim_kv = default(dim_context, dim)
@@ -1695,7 +1696,8 @@ def __init__(
             logit_softclamp_value = logit_softclamp_value,
             cope = cope,
             onnxable = onnxable,
-            sdp_kwargs = attend_sdp_kwargs
+            sdp_kwargs = attend_sdp_kwargs,
+            flash_pack_seq=flash_pack_seq,
         )
 
         # head scaling
@@ -1833,6 +1835,7 @@ def forward(
         additional_key_values: tuple[Tensor, Tensor] | None = None,
         additional_key_value_mask = None,
         kv_input_residual = None,
+        flash_pack_seq_kwargs = None,
     ):
         b, n, h, kv_h, head_scale, num_mem_kv, device, has_context, qkv_receive_diff_residuals, is_multi_latent_attn = x.shape[0], x.shape[1], self.heads, self.kv_heads, self.head_scale, self.num_mem_kv, x.device, exists(context), self.qkv_receive_diff_residuals, self.use_latent_kv
 
@@ -2099,7 +2102,8 @@ def forward(
             q, k, v,
             mask = final_attn_mask,
             attn_bias = attn_bias,
-            prev_attn = prev_attn
+            prev_attn = prev_attn,
+            flash_pack_seq_kwargs = flash_pack_seq_kwargs,
         )
 
         # laser
@@ -2306,7 +2310,7 @@ def __init__(
 
         dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
         data_dependent_alibi = attn_kwargs.get('data_dependent_alibi', False)
-
+        flash_pack_seq = attn_kwargs.get('flash_pack_seq', False)
         assert len(kwargs) == 0, f'unrecognized kwargs passed in {kwargs.keys()}'
 
         self.dim = dim
@@ -2353,6 +2357,7 @@ def __init__(
 
         assert at_most_one_of(rotary_pos_emb, polar_pos_emb), f'either rotary positional embedding or polar positional embedding can be turned on'
         assert not (rotary_xpos and not causal), 'rotary xpos is not compatible with bidirectional attention'
+        assert not flash_pack_seq or rotary_pos_emb, 'block masking only tested for rotary positional embeddings'
         self.rotary_pos_emb = RotaryEmbedding(rotary_emb_dim, use_xpos = rotary_xpos, scale_base = rotary_xpos_scale_base, interpolation_factor = rotary_interpolation_factor, base_rescale_factor = rotary_base_rescale_factor) if rotary_pos_emb else None
 
         # polar positional embedding (PoPE) - https://arxiv.org/abs/2509.10534
@@ -2733,10 +2738,14 @@ def forward(
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
         layers_execute_order: tuple[int, ...] | None = None,
         self_attn_kv_residuals: Tensor | None = None,
-        cross_attn_kv_residuals: Tensor | None = None
+        cross_attn_kv_residuals: Tensor | None = None,
+        flash_pack_seq_kwargs = None,
+        flash_pack_seq_context_kwargs = None,
     ):
         assert not (self.cross_attend ^ exists(context)), 'context must be passed in if cross_attend is set to True'
         assert not (exists(condition) ^ self.need_condition), 'condition needs to be passed in if using adaptive layernorm or vice versa'
+        assert not (exists(flash_pack_seq_kwargs) and (exists(attn_mask) or exists(mask))), 'attn_mask or mask cannot be used with flash block masking'
+        assert not (exists(flash_pack_seq_context_kwargs) and (exists(context_mask))), 'context_mask cannot be used with flash block masking'
 
         # handle seq pos offset if not passed in from wrapper
         # default to 0, but if cache is detected, set appropriate for the relative positional embeddings
@@ -3016,9 +3025,9 @@ def forward(
             # forward depending on layer type
 
             if layer_type == 'a':
-                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, polar_pos_emb = polar_pos_emb, additional_key_values = next(iter_self_attn_kv, None), additional_key_value_mask = additional_kv_mask, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, kv_input_residual = next(self_attn_kv_residuals_iter, None), value_residual = maybe_self_attn_value_residual, return_intermediates = True)
+                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, polar_pos_emb = polar_pos_emb, additional_key_values = next(iter_self_attn_kv, None), additional_key_value_mask = additional_kv_mask, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, kv_input_residual = next(self_attn_kv_residuals_iter, None), value_residual = maybe_self_attn_value_residual, flash_pack_seq_kwargs = flash_pack_seq_kwargs, return_intermediates = True)
             elif layer_type == 'c':
-                out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), kv_input_residual = next(cross_attn_kv_residuals_iter, None), value_residual = maybe_cross_attn_value_residual, **cross_attn_rotary_pos_emb, return_intermediates = True)
+                out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), kv_input_residual = next(cross_attn_kv_residuals_iter, None), value_residual = maybe_cross_attn_value_residual, **cross_attn_rotary_pos_emb, flash_pack_seq_kwargs = flash_pack_seq_context_kwargs, return_intermediates = True)
             elif layer_type == 'f':
                 out = block(x, deep_embed = next(deep_embeds_iter, None))