[CI] Remove Flash Attention dependency (#705)

* Update flash-attn version in requirements-test.txt from <=2.2.0 to ==2.5.8

* lint fix

* Remove unused dependencies from requirements-test.txt

* Update import path for padding functions in example MHA forward variable length script

* Refactor code formatting in bert_padding.py for improved readability

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_paged.py

@@ -517,11 +517,21 @@ def main(args):
     output_sparse = sparse_attn.forward(Q, K_cache, V_cache, block_indices, cache_seqlens,
                                         block_table)
 
-    output_ref_fa = ref_program_fa(Q, K_cache, V_cache, cache_seqlens, block_table)
+    is_flash_attn_2_available = False
+    try:
+        import flash_attn  # noqa: F401
+        is_flash_attn_2_available = True
+    except:
+        pass
 
     output_ref_torch = ref_program_torch_paged(Q, K_cache, V_cache, block_indices, cache_seqlens,
                                                block_table, page_block_size, block_N)
 
+    if not is_flash_attn_2_available:
+        print("FlashAttn 2 is not available, skipping FA reference and performance measurement")
+        return
+
+    output_ref_fa = ref_program_fa(Q, K_cache, V_cache, cache_seqlens, block_table)
     # Check correctness
     if sparse_ratio == 0.0:
         max_diff = torch.max(torch.abs(output_sparse - output_ref_fa)).item()

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py

@@ -439,6 +439,17 @@ def main(batch=8,
     out = sparse_kernel(Q, K, V, block_indices, cache_seqlens)
     debug("output", ref, out, atol=1e-3, rtol=1e-3)
 
+    is_flash_attn_2_available = False
+    try:
+        import flash_attn  # noqa: F401
+        is_flash_attn_2_available = True
+    except ImportError:
+        pass
+
+    if not is_flash_attn_2_available:
+        print("FlashAttn 2 is not available, skipping FA reference and performance measurement")
+        return
+
     ## latency reference
     for _ in range(10):
         ref = ref_program_fa(Q, K, V, block_indices, cache_seqlens, max_cache_seqlen,

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py

@@ -419,6 +419,17 @@ def main(batch=8,
     out = model(Q, K, V, block_mask, cache_seqlens)
     debug("output", ref, out, atol=1e-3, rtol=1e-3)
 
+    is_flash_attn_2_available = False
+    try:
+        import flash_attn  # noqa: F401
+        is_flash_attn_2_available = True
+    except ImportError:
+        pass
+
+    if not is_flash_attn_2_available:
+        print("FlashAttn 2 is not available, skipping FA reference and performance measurement")
+        return
+
     ## latency reference
     for _ in range(10):
         ref = ref_program_fa(Q, K, V, block_mask, cache_seqlens, max_cache_seqlen, num_blocks,

examples/blocksparse_attention/example_triton_sparse_gqa_decode_varlen_indice.py

@@ -449,6 +449,17 @@ def main(batch=64,
     print(f"Average time: {avg_time:.6f} seconds")
 
     # Measure performance of reference implementation
+    is_flash_attn_2_available = False
+    try:
+        import flash_attn  # noqa: F401
+        is_flash_attn_2_available = True
+    except ImportError:
+        pass
+
+    if not is_flash_attn_2_available:
+        print("FlashAttn 2 is not available, skipping FA reference and performance measurement")
+        return
+
     start = time.time()
     for _ in range(1000):
         ref_program_fa(Q, K, V, cache_seqlens)

examples/blocksparse_attention/example_triton_sparse_gqa_decode_varlen_mask.py

@@ -429,10 +429,22 @@ def main(batch=64,
     print(f"Average time: {avg_time:.6f} seconds")
     print(f"Average flops: {avg_flops:.2f} GFLOPS")
 
+    is_flash_attn_2_available = False
+    try:
+        import flash_attn  # noqa: F401
+        is_flash_attn_2_available = True
+    except ImportError:
+        pass
+
     # Measure performance of reference implementation
+    if not is_flash_attn_2_available:
+        print("FlashAttn 2 is not available, skipping FA reference and performance measurement")
+        return
+
     start = time.time()
     for _ in range(1000):
         ref_program_fa(Q, K, V, cache_seqlens)
+
     torch.cuda.synchronize()
     end = time.time()
     elapsed_time_ref = end - start

examples/flash_attention/bert_padding.py

+# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
+# ruff: noqa
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+
+class IndexFirstAxis(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input, indices):
+        ctx.save_for_backward(indices)
+        assert input.ndim >= 2
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        second_dim = other_shape.numel()
+        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+        # return input[indices]
+        return torch.gather(
+            rearrange(input, "b ... -> b (...)"), 0,
+            repeat(indices, "z -> z d", d=second_dim)).reshape(-1, *other_shape)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        assert grad_output.ndim >= 2
+        other_shape = grad_output.shape[1:]
+        grad_output = rearrange(grad_output, "b ... -> b (...)")
+        grad_input = torch.zeros(
+            [ctx.first_axis_dim, grad_output.shape[1]],
+            device=grad_output.device,
+            dtype=grad_output.dtype,
+        )
+        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
+        # grad_input[indices] = grad_output
+        grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis = IndexFirstAxis.apply
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, values, indices, first_axis_dim):
+        ctx.save_for_backward(indices)
+        assert indices.ndim == 1
+        assert values.ndim >= 2
+        output = torch.zeros(
+            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype)
+        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
+        output[indices] = values
+        # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (indices,) = ctx.saved_tensors
+        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+        grad_values = grad_output[indices]
+        # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1]))
+        return grad_values, None, None
+
+
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+class IndexFirstAxisResidual(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input, indices):
+        ctx.save_for_backward(indices)
+        assert input.ndim >= 2
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        second_dim = other_shape.numel()
+        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+        output = input[indices]
+        # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last
+        # memory format to channel_first. In other words, input might not be contiguous.
+        # If we don't detach, Pytorch complains about output being a view and is being modified inplace
+        return output, input.detach()
+
+    @staticmethod
+    def backward(ctx, grad_output, grad_residual):
+        (indices,) = ctx.saved_tensors
+        assert grad_output.ndim >= 2
+        other_shape = grad_output.shape[1:]
+        assert grad_residual.shape[1:] == other_shape
+        grad_input = grad_residual
+        # grad_input[indices] += grad_output
+        indices = indices.reshape(indices.shape[0], *((1,) * (grad_output.ndim - 1)))
+        indices = indices.expand_as(grad_output)
+        grad_input.scatter_add_(0, indices, grad_output)
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis_residual = IndexFirstAxisResidual.apply
+
+
+def unpad_input(hidden_states, attention_mask):
+    """
+    Arguments:
+        hidden_states: (batch, seqlen, ...)
+        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+    Return:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
+        max_seqlen_in_batch: int
+    """
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
+    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
+    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
+    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
+    # so we write custom forward and backward to make it a bit faster.
+    return (
+        index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices),
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+def unpad_input_for_concatenated_sequences(hidden_states, attention_mask_in_length):
+    """
+    Supports concatenating short samples in one sequence. The attention_mask_in_length is utilized to mask other short samples. It helps efficient training of variant lengths-based samples (e.g., the supervised fine-tuning task in large language model).
+    The motivation for this function is explained [here](https://github.com/Dao-AILab/flash-attention/issues/432#issuecomment-1668822286).
+    
+    For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is:
+        ```
+        [
+          [2, 3, 0, 0, 0, 0],
+          [3, 2, 0, 0, 0, 0],
+          [6, 0, 0, 0, 0, 0]
+        ]
+        ```
+    , which refers to the 3D-attention mask:
+        ```
+        [
+          [
+            [1, 0, 0, 0, 0, 0],
+            [1, 1, 0, 0, 0, 0],
+            [0, 0, 1, 0, 0, 0],
+            [0, 0, 1, 1, 0, 0],
+            [0, 0, 1, 1, 1, 0],
+            [0, 0, 0, 0, 0, 1]
+          ],
+          [
+            [1, 0, 0, 0, 0, 0],
+            [1, 1, 0, 0, 0, 0],
+            [1, 1, 1, 0, 0, 0],
+            [0, 0, 0, 1, 0, 0],
+            [0, 0, 0, 1, 1, 0],
+            [0, 0, 0, 0, 0, 1]
+          ],
+          [
+            [1, 0, 0, 0, 0, 0],
+            [1, 1, 0, 0, 0, 0],
+            [1, 1, 1, 0, 0, 0],
+            [1, 1, 1, 1, 0, 0],
+            [1, 1, 1, 1, 1, 0],
+            [1, 1, 1, 1, 1, 1]
+          ]
+        ]
+        ```.
+
+    Arguments:
+        hidden_states: (batch, seqlen, ...)
+        attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none.
+    Return:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
+        max_seqlen_in_batch: int
+    """
+    length = attention_mask_in_length.sum(dim=-1)
+    seqlen = attention_mask_in_length.size(-1)
+    attention_mask_2d = torch.arange(
+        seqlen, device=length.device, dtype=length.dtype).expand(len(length),
+                                                                 seqlen) < length.unsqueeze(1)
+    real_indices_idx = torch.nonzero(attention_mask_in_length.flatten(), as_tuple=False).flatten()
+    seqlens_in_batch = attention_mask_in_length.flatten()[real_indices_idx]
+    indices = torch.nonzero(attention_mask_2d.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
+    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
+    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
+    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
+    # so we write custom forward and backward to make it a bit faster.
+    return (
+        index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices),
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+def pad_input(hidden_states, indices, batch, seqlen):
+    """
+    Arguments:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        indices: (total_nnz)
+    Return:
+        hidden_states: (batch, seqlen, ...)
+    """
+    dim = hidden_states.shape[-1]
+    # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype)
+    # output[indices] = hidden_states
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)

examples/flash_attention/example_mha_fwd_varlen.py

@@ -7,7 +7,7 @@ import argparse
 
 import torch
 from einops import rearrange, repeat
-from flash_attn.bert_padding import pad_input, unpad_input
+from bert_padding import pad_input, unpad_input
 
 
 def generate_random_padding_mask(max_seqlen, batch_size, device, mode="random"):
@@ -410,7 +410,19 @@ def main(batch: int = 2, heads: int = 16, seq_len: int = 256, dim: int = 32):
         key_padding_mask,
         causal=causal,
     )
-    import flash_attn
+    torch.testing.assert_close(out, out_ref, rtol=1e-2, atol=1e-2)
+
+    is_flash_attn_2_available = False
+    try:
+        import flash_attn
+        is_flash_attn_2_available = True
+    except:
+        pass
+
+    if not is_flash_attn_2_available:
+        print("FlashAttn 2 is not available, skipping FA reference and performance measurement")
+        return
+
     fla_out_unpad = flash_attn.flash_attn_varlen_func(
         q_unpad,
         k_unpad,
@@ -423,8 +435,8 @@ def main(batch: int = 2, heads: int = 16, seq_len: int = 256, dim: int = 32):
         causal=causal,
     )
     fla_out = output_pad_fn(fla_out_unpad)
-    torch.testing.assert_close(out, out_ref, rtol=1e-2, atol=1e-2)
     torch.testing.assert_close(out, fla_out, rtol=1e-2, atol=1e-2)
+
     print("Assert Equal Passed")
 
 

requirements-test.txt

@@ -26,8 +26,5 @@ tabulate
 wheel
 setuptools
 einops
-attrs
-decorator
-flash-attn<=2.2.0
 scipy
 tornado