Add additional KV support for flex_na (#211)

As described in the title.

For unknown reason, flex_na3d with additional KV needs `eps=0.12` rather
than `eps=0.1` to pass the precision check. Will come back to this
later.

---------

Co-authored-by: Ali Hassani <ahassani@nvidia.com>

Author: Leymore

CHANGELOG.md

@@ -10,6 +10,8 @@
     * Now you can use Flex Attention instead of FNA through NATTEN directly.
     * Just import `use_flex_attention()` from `natten`, call it, and enjoy potentially significant
      speedups on newer architectures.
+    * With support for additional KV tokens.
+* Better precision on fused ops with additional KV.
 
 
 ## [0.17.4] - 2025-01-28

requirements-dev.txt

@@ -7,6 +7,6 @@ mypy==1.8.0
 pytest==7.4.4
 click==8.1.7
 rich
-xformers>=0.0.25
+xformers==v0.0.28.post3
 fvcore==0.1.5.post20221221
 twine

src/natten/__init__.py

@@ -30,6 +30,7 @@
     enable_gemm_na,
     enable_tf32,
     enable_tiled_na,
+    force_flex_attention,
     get_memory_usage_preference,
     has_bfloat,
     has_cuda,
@@ -57,6 +58,7 @@
     set_memory_usage_preference,
     use_autotuner,
     use_deterministic_algorithms,
+    use_flex_attention,
     use_fna,
     use_fused_na,
     use_gemm_na,
@@ -81,6 +83,7 @@
     "use_fused_na",
     "is_fused_na_enabled",
     "use_autotuner",
+    "force_flex_attention",
     "disable_autotuner",
     "is_autotuner_enabled",
     "is_autotuner_enabled_for_forward",
@@ -99,6 +102,7 @@
     "use_tf32_in_gemm_na",
     "use_tiled_na",
     "use_gemm_na",
+    "use_flex_attention",
     "is_tf32_in_gemm_na_enabled",
     "is_tiled_na_enabled",
     "is_gemm_na_enabled",
@@ -114,7 +118,6 @@
     "disable_gemm_na",
     "enable_tiled_na",
     "disable_tiled_na",
-    "use_flex_attention",
 ]
 
 __version__ = "0.17.5.dev0"

src/natten/flex.py

@@ -23,12 +23,13 @@
 
 import functools
 import math
-from typing import Optional, Tuple
+from typing import Dict, Optional, Tuple
 
 import torch
 from torch import BoolTensor, IntTensor, Tensor
 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 
+from .ops import additional_sdpa, merge_attentions
 from .types import (
     CausalArg1DTypeOrDed,
     CausalArg2DTypeOrDed,
@@ -129,7 +130,10 @@ def flex_na1d(
     kernel_size: Dimension1DTypeOrDed,
     dilation: Dimension1DTypeOrDed = 1,
     is_causal: Optional[CausalArg1DTypeOrDed] = False,
-) -> torch.Tensor:
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    xformers_kwargs: Optional[Dict] = None,
+) -> Tensor:
 
     kernel_size_, dilation_, is_causal_ = check_all_args(
         1, kernel_size, dilation, is_causal
@@ -170,9 +174,30 @@ def flex_na1d(
 
     na_mask = get_na_flex_mask(1, num_tokens_tuple, kernel_size_, dilation_, is_causal_)
     flex_attention_compiled = get_flex_attention_compiled()
-    out_ = flex_attention_compiled(query_, key_, value_, block_mask=na_mask)
+    out_, lse_ = flex_attention_compiled(
+        query_, key_, value_, block_mask=na_mask, return_lse=True
+    )
 
     out = out_.transpose(1, 2)
+    lse = lse_.transpose(1, 2)
+
+    if additional_keys is not None and additional_values is not None:
+        if additional_keys is None or additional_values is None:
+            raise ValueError(
+                "Both `additional_keys` and `additional_values` must be "
+                "either Tensors or NoneTypes."
+            )
+
+        scale = query.shape[-1] ** -0.5
+        additional_output, additional_lse = additional_sdpa(
+            query,
+            additional_keys,
+            additional_values,
+            scale=scale,
+            attn_kwargs=xformers_kwargs,
+        )
+
+        return merge_attentions(out, additional_output, lse, additional_lse)
 
     return out
 
@@ -184,7 +209,10 @@ def flex_na2d(
     kernel_size: Dimension2DTypeOrDed,
     dilation: Dimension2DTypeOrDed = 1,
     is_causal: Optional[CausalArg2DTypeOrDed] = False,
-) -> torch.Tensor:
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    xformers_kwargs: Optional[Dict] = None,
+) -> Tensor:
 
     kernel_size_, dilation_, is_causal_ = check_all_args(
         2, kernel_size, dilation, is_causal
@@ -225,9 +253,30 @@ def flex_na2d(
 
     na_mask = get_na_flex_mask(2, num_tokens_tuple, kernel_size_, dilation_, is_causal_)
     flex_attention_compiled = get_flex_attention_compiled()
-    out_ = flex_attention_compiled(query_, key_, value_, block_mask=na_mask)
+    out_, lse_ = flex_attention_compiled(
+        query_, key_, value_, block_mask=na_mask, return_lse=True
+    )
 
     out = out_.transpose(1, 2).view(batch_size, *num_tokens_tuple, num_heads, head_dim)
+    lse = lse_.transpose(1, 2).view(batch_size, *num_tokens_tuple, num_heads)
+
+    if additional_keys is not None and additional_values is not None:
+        if additional_keys is None or additional_values is None:
+            raise ValueError(
+                "Both `additional_keys` and `additional_values` must be "
+                "either Tensors or NoneTypes."
+            )
+
+        scale = query.shape[-1] ** -0.5
+        additional_output, additional_lse = additional_sdpa(
+            query,
+            additional_keys,
+            additional_values,
+            scale=scale,
+            attn_kwargs=xformers_kwargs,
+        )
+
+        return merge_attentions(out, additional_output, lse, additional_lse)
 
     return out
 
@@ -239,7 +288,10 @@ def flex_na3d(
     kernel_size: Dimension3DTypeOrDed,
     dilation: Dimension3DTypeOrDed = 1,
     is_causal: Optional[CausalArg3DTypeOrDed] = False,
-) -> torch.Tensor:
+    additional_keys: Optional[Tensor] = None,
+    additional_values: Optional[Tensor] = None,
+    xformers_kwargs: Optional[Dict] = None,
+) -> Tensor:
 
     kernel_size_, dilation_, is_causal_ = check_all_args(
         3, kernel_size, dilation, is_causal
@@ -280,8 +332,29 @@ def flex_na3d(
 
     na_mask = get_na_flex_mask(3, num_tokens_tuple, kernel_size_, dilation_, is_causal_)
     flex_attention_compiled = get_flex_attention_compiled()
-    out_ = flex_attention_compiled(query_, key_, value_, block_mask=na_mask)
+    out_, lse_ = flex_attention_compiled(
+        query_, key_, value_, block_mask=na_mask, return_lse=True
+    )
 
     out = out_.transpose(1, 2).view(batch_size, *num_tokens_tuple, num_heads, head_dim)
+    lse = lse_.transpose(1, 2).view(batch_size, *num_tokens_tuple, num_heads)
+
+    if additional_keys is not None and additional_values is not None:
+        if additional_keys is None or additional_values is None:
+            raise ValueError(
+                "Both `additional_keys` and `additional_values` must be "
+                "either Tensors or NoneTypes."
+            )
+
+        scale = query.shape[-1] ** -0.5
+        additional_output, additional_lse = additional_sdpa(
+            query,
+            additional_keys,
+            additional_values,
+            scale=scale,
+            attn_kwargs=xformers_kwargs,
+        )
+
+        return merge_attentions(out, additional_output, lse, additional_lse)
 
     return out

src/natten/functional.py

@@ -1735,10 +1735,6 @@ def na1d(
             raise NotImplementedError(
                 "RPB is not supported in the Flex Attention backend."
             )
-        if additional_keys is not None or additional_values is not None:
-            raise NotImplementedError(
-                "Additional keys/values is not supported in the Flex Attention backend."
-            )
 
         return flex_na1d(
             query,
@@ -1747,6 +1743,9 @@ def na1d(
             kernel_size,
             dilation,
             is_causal,
+            additional_keys=additional_keys,
+            additional_values=additional_values,
+            xformers_kwargs=xformers_kwargs,
         )
 
     tiling_config_forward, tiling_config_backward = autotune_fna(
@@ -1817,10 +1816,6 @@ def na2d(
             raise NotImplementedError(
                 "RPB is not supported in the Flex Attention backend."
             )
-        if additional_keys is not None or additional_values is not None:
-            raise NotImplementedError(
-                "Additional keys/values is not supported in the Flex Attention backend."
-            )
 
         return flex_na2d(
             query,
@@ -1829,6 +1824,9 @@ def na2d(
             kernel_size,
             dilation,
             is_causal,
+            additional_keys=additional_keys,
+            additional_values=additional_values,
+            xformers_kwargs=xformers_kwargs,
         )
 
     tiling_config_forward, tiling_config_backward = autotune_fna(
@@ -1899,10 +1897,6 @@ def na3d(
             raise NotImplementedError(
                 "RPB is not supported in the Flex Attention backend."
             )
-        if additional_keys is not None or additional_values is not None:
-            raise NotImplementedError(
-                "Additional keys/values is not supported in the Flex Attention backend."
-            )
 
         return flex_na3d(
             query,
@@ -1911,6 +1905,9 @@ def na3d(
             kernel_size,
             dilation,
             is_causal,
+            additional_keys=additional_keys,
+            additional_values=additional_values,
+            xformers_kwargs=xformers_kwargs,
         )
 
     tiling_config_forward, tiling_config_backward = autotune_fna(

src/natten/ops.py

@@ -164,17 +164,16 @@ def merge_attentions(
     output_0 = output_fna.reshape(input_shape).to(accum_type)
     output_1 = output_sdpa.reshape(input_shape).to(accum_type)
 
-    sum_of_exps_0 = lse_0.exp().unsqueeze(-1).expand(*input_shape)
-    sum_of_exps_1 = lse_1.exp().unsqueeze(-1).expand(*input_shape)
+    lse_max = torch.maximum(lse_0, lse_1)
+    exp_diff_0 = torch.exp(lse_0 - lse_max).unsqueeze(-1)
+    exp_diff_1 = torch.exp(lse_1 - lse_max).unsqueeze(-1)
 
-    assert sum_of_exps_0.shape == sum_of_exps_1.shape == output_0.shape
-
-    output_0_rescaled = output_0 * sum_of_exps_0
-    output_1_rescaled = output_1 * sum_of_exps_1
+    output_0_rescaled = output_0 * exp_diff_0
+    output_1_rescaled = output_1 * exp_diff_1
 
     assert output_0_rescaled.shape == output_1_rescaled.shape == output_0.shape
 
-    sum_of_exps = sum_of_exps_0 + sum_of_exps_1
+    sum_of_exps = exp_diff_0 + exp_diff_1
 
     output = (output_0_rescaled + output_1_rescaled) / sum_of_exps