JVP Flash Attention

Description

Flash Attention Triton kernel with support for second-order derivatives, such as Jacobian-Vector Products (JVPs) and Hessian-Vector Products (HVPs)

Installation

Using pip, one can install jvp_flash_attention as follows.

# Install package
pip install jvp_flash_attention

# [OPTIONAL, for development] Install package and pre-commit hooks
pip install -e .
pre-commit install

Usage

Once installed, one can use jvp_flash_attention in place of PyTorch's scaled_dot_product_attention as follows.

import torch.nn.functional as F

from torch.nn.attention import SDPBackend, sdpa_kernel
from jvp_flash_attention.jvp_attention import JVPAttn, attention as jvp_attention

with sdpa_kernel(SDPBackend.MATH):
  # Regular (quadratic) attention
  x = F.scaled_dot_product_attention(
      q,
      k,
      v,
      attn_mask=attn_mask,
      dropout_p=attn_dropout_p if self.training else 0.0,
  )

# JVP flash attention
x = jvp_attention(
    q,
    k,
    v,
    attn_mask=attn_mask,
    # dropout_p=attn_dropout_p if self.training else 0.0,  # NOTE: Attention dropout is currently unsupported
)

Anecdotally, one can also swap out F.scaled_dot_product_attention with jvp_attention even for pretrained models with minimal impact on numerical accuracy.

Note: If calling torch.func.jvp manually in your model's forward pass like pred, df = torch.func.jvp(*(lambda x_jvp: model(x_jvp), (x,), (gt,))), make sure to use JVP Flash Attention in your model as model = lambda q, k, v: JVPAttn.fwd_dual(q, k, v) instead of as model = lambda q, k, v: jvp_attention(q, k, v) to ensure each input's tangent vectors are computed prior to running PyTorch's autograd engine. Models that rely on torch.autograd.grad to compute higher-order derivatives in their forward pass (e.g., energy-based models) should not require this change.

Contributions or enhancements are welcome!

Results

Loss matching

Model training with either F.scaled_dot_product_attention or JVPAttn.fwd_dual produces the same loss trajectory.

Speed matching

Model training with either F.scaled_dot_product_attention or JVPAttn.fwd_dual achieves the same iteration speed.

Note: The following results can be reproduced (for float32 precision) by running python tests/test_jvp_attention.py --dtype float32.

Time scaling

jvp_attention outscales the speed of (SDPBackend.MATH-based) F.scaled_dot_product_attention when calculating second-order derivatives.

Memory scaling

jvp_attention improves the memory usage of (SDPBackend.MATH-based) F.scaled_dot_product_attention when calculating second-order derivatives.

Tests

If you want to run all the unit tests verifying the correctness of the JVP Flash Attention Triton kernel, run the following command(s).

python tests/test_jvp_attention.py --dtype {float16,bfloat16,float32}

In principle, the kernel should support ROCm systems as well, though it has not yet been tested on them. macOS is currently unsupported except using a CPU-only backend.

Full results for float16:

==============================================================================================================
BENCHMARK SUMMARY
==============================================================================================================
Seq Len    Causal   Mask       Method     Time (ms)    Mem (MB)     TFLOP/s      Max Error    Grad Check
--------------------------------------------------------------------------------------------------------------
32         False    additive   sdpa       0.821        3.09           0.0 TFLOP/s baseline     N/A
32         False    additive   jvp_attn   0.723        1.08           0.0 TFLOP/s 1.83e+01     ✗

32         False    boolean    sdpa       0.961        3.14           0.0 TFLOP/s baseline     N/A
32         False    boolean    jvp_attn   0.504        1.03           0.0 TFLOP/s 3.91e-03     ✓

32         False    none       sdpa       0.576        3.09           0.0 TFLOP/s baseline     N/A
32         False    none       jvp_attn   0.447        1.03           0.0 TFLOP/s 1.95e-03     ✓

32         True     none       sdpa       0.934        3.10           0.0 TFLOP/s baseline     N/A
32         True     none       jvp_attn   0.458        1.03           0.0 TFLOP/s 3.91e-03     ✓

64         False    additive   sdpa       0.860        6.75           0.0 TFLOP/s baseline     N/A
64         False    additive   jvp_attn   0.847        2.26           0.1 TFLOP/s 2.23e+00     ✗

64         False    boolean    sdpa       0.908        6.94           0.0 TFLOP/s baseline     N/A
64         False    boolean    jvp_attn   0.521        2.07           0.1 TFLOP/s 3.91e-03     ✓

64         False    none       sdpa       0.542        6.75           0.0 TFLOP/s baseline     N/A
64         False    none       jvp_attn   0.414        2.07           0.1 TFLOP/s 1.95e-03     ✓

64         True     none       sdpa       0.888        6.77           0.0 TFLOP/s baseline     N/A
64         True     none       jvp_attn   0.437        2.07           0.1 TFLOP/s 2.20e-03     ✓

128        False    additive   sdpa       0.834        16.51          0.1 TFLOP/s baseline     N/A
128        False    additive   jvp_attn   0.750        4.89           0.3 TFLOP/s 3.91e-03     ✓

128        False    boolean    sdpa       0.840        17.26          0.1 TFLOP/s baseline     N/A
128        False    boolean    jvp_attn   0.520        4.14           0.4 TFLOP/s 3.91e-03     ✓

128        False    none       sdpa       0.610        16.51          0.2 TFLOP/s baseline     N/A
128        False    none       jvp_attn   0.459        4.14           0.4 TFLOP/s 9.77e-04     ✓

128        True     none       sdpa       1.053        16.57          0.0 TFLOP/s baseline     N/A
128        True     none       jvp_attn   0.438        4.14           0.2 TFLOP/s 2.44e-03     ✓

256        False    additive   sdpa       0.829        47.77          0.5 TFLOP/s baseline     N/A
256        False    additive   jvp_attn   0.738        12.02          1.1 TFLOP/s 3.91e-03     ✓

256        False    boolean    sdpa       0.872        50.77          0.5 TFLOP/s baseline     N/A
256        False    boolean    jvp_attn   0.482        8.27           1.7 TFLOP/s 3.91e-03     ✓

256        False    none       sdpa       0.812        47.27          0.5 TFLOP/s baseline     N/A
256        False    none       jvp_attn   0.460        8.27           1.8 TFLOP/s 9.77e-04     ✓

256        True     none       sdpa       0.964        47.52          0.2 TFLOP/s baseline     N/A
256        True     none       jvp_attn   0.436        8.27           0.9 TFLOP/s 3.91e-03     ✓

512        False    additive   sdpa       1.416        153.55         1.2 TFLOP/s baseline     N/A
512        False    additive   jvp_attn   0.715        30.55          4.6 TFLOP/s 1.95e-03     ✓

512        False    boolean    sdpa       1.441        165.05         1.1 TFLOP/s baseline     N/A
512        False    boolean    jvp_attn   0.500        16.55          6.6 TFLOP/s 1.95e-03     ✓

512        False    none       sdpa       1.374        153.05         1.2 TFLOP/s baseline     N/A
512        False    none       jvp_attn   0.407        16.55          8.1 TFLOP/s 4.88e-04     ✓

512        True     none       sdpa       1.402        154.05         0.6 TFLOP/s baseline     N/A
512        True     none       jvp_attn   0.460        16.55          3.6 TFLOP/s 2.93e-03     ✓

1024       False    additive   sdpa       4.963        546.84         1.3 TFLOP/s baseline     N/A
1024       False    additive   jvp_attn   1.183        96.84         11.1 TFLOP/s 1.95e-03     ✓

1024       False    boolean    sdpa       4.991        594.84         1.3 TFLOP/s baseline     N/A
1024       False    boolean    jvp_attn   0.622        33.84         21.1 TFLOP/s 1.95e-03     ✓

1024       False    none       sdpa       4.227        546.84         1.6 TFLOP/s baseline     N/A
1024       False    none       jvp_attn   0.420        33.84         31.3 TFLOP/s 4.88e-04     ✓

1024       True     none       sdpa       4.861        550.84         0.7 TFLOP/s baseline     N/A
1024       True     none       jvp_attn   0.469        33.84         14.0 TFLOP/s 3.91e-03     ✓

2048       False    additive   sdpa       18.773       2052.19        1.4 TFLOP/s baseline     N/A
2048       False    additive   jvp_attn   3.379        336.19        15.6 TFLOP/s 1.95e-03     ✓

2048       False    boolean    sdpa       18.815       2244.19        1.4 TFLOP/s baseline     N/A
2048       False    boolean    jvp_attn   1.674        66.19         31.4 TFLOP/s 1.95e-03     ✓

2048       False    none       sdpa       16.156       2052.19        1.6 TFLOP/s baseline     N/A
2048       False    none       jvp_attn   1.186        66.19         44.3 TFLOP/s 4.88e-04     ✓

2048       True     none       sdpa       18.587       2068.19        0.7 TFLOP/s baseline     N/A
2048       True     none       jvp_attn   0.720        66.19         36.5 TFLOP/s 1.95e-03     ✓


================================================================================
MASK TYPE PERFORMANCE COMPARISON
================================================================================
Seq Len    Causal   Method     No Mask         Boolean Mask    Additive Mask
--------------------------------------------------------------------------------
32         False    jvp_attn   0.45 ms         0.50 ms (1.13x) 0.72 ms (1.62x)
32         True     jvp_attn   0.46 ms         N/A             N/A
64         False    jvp_attn   0.41 ms         0.52 ms (1.26x) 0.85 ms (2.05x)
64         True     jvp_attn   0.44 ms         N/A             N/A
128        False    jvp_attn   0.46 ms         0.52 ms (1.13x) 0.75 ms (1.63x)
128        True     jvp_attn   0.44 ms         N/A             N/A
256        False    jvp_attn   0.46 ms         0.48 ms (1.05x) 0.74 ms (1.60x)
256        True     jvp_attn   0.44 ms         N/A             N/A
512        False    jvp_attn   0.41 ms         0.50 ms (1.23x) 0.72 ms (1.76x)
512        True     jvp_attn   0.46 ms         N/A             N/A
1024       False    jvp_attn   0.42 ms         0.62 ms (1.48x) 1.18 ms (2.82x)
1024       True     jvp_attn   0.47 ms         N/A             N/A
2048       False    jvp_attn   1.19 ms         1.67 ms (1.41x) 3.38 ms (2.85x)
2048       True     jvp_attn   0.72 ms         N/A             N/A

============================================================
STATISTICS
============================================================
Average speedup: 4.50x
Min speedup: 1.02x
Max speedup: 25.82x

Accuracy: 26/28 tests passed
⚠️  Some accuracy checks failed

Failed configurations:
  - Seq=32, Causal=False, Mask=additive
  - Seq=64, Causal=False, Mask=additive

Full results for bfloat16:

```

BENCHMARK SUMMARY

Seq Len Causal Mask Method Time (ms) Mem (MB) TFLOP/s Max Error Grad Check

32 False additive sdpa 0.864 3.09 0.0 TFLOP/s baseline N/A 32 False additive jvp_attn 0.773 1.08 0.0 TFLOP/s 1.84e+01 ✗

32 False boolean sdpa 0.949 3.14 0.0 TFLOP/s baseline N/A 32 False boolean jvp_attn 0.569 1.03 0.0 TFLOP/s 3.12e-02 ✓

32 False none sdpa 0.662 3.09 0.0 TFLOP/s baseline N/A 32 False none jvp_attn 0.447 1.03 0.0 TFLOP/s 1.56e-02 ✓

32 True none sdpa 0.945 3.10 0.0 TFLOP/s baseline N/A 32 True none jvp_attn 0.469 1.03 0.0 TFLOP/s 3.12e-02 ✓

64 False additive sdpa 0.923 6.75 0.0 TFLOP/s baseline N/A 64 False additive jvp_attn 1.149 2.26 0.0 TFLOP/s 2.23e+00 ✗

64 False boolean sdpa 0.910 6.94 0.0 TFLOP/s baseline N/A 64 False boolean jvp_attn 0.518 2.07 0.1 TFLOP/s 3.12e-02 ✓

64 False none sdpa 0.554 6.75 0.0 TFLOP/s bas