SwinUNETR: optional flash attention (scaled_dot_product_attention) in WindowAttention (#8977)

Fixes #8973 .

### Description

SwinUNETR's WindowAttention builds the full (nWindows*heads, N, N) score
matrix by hand before softmax. This adds an opt-in use_flash_attention
flag (default False, so existing behaviour is unchanged) that routes
attention through torch.nn.functional.scaled_dot_product_attention,
folding the relative position bias and, for shifted windows, the
attention mask into one additive attn_mask cast to the query dtype. This
mirrors the flash-attention option already in MONAI's SelfAttention,
CrossAttention and CABlock.

Measured on the SwinUNETR encoder (SwinViT) forward, inference, single
GPU, best-of-5. Float32 output matches the default path to within 3e-6
and is bit-exact in float64, verified across 2D and 3D, batch sizes 1 to
4, and non-cubic inputs.

| ROI | dtype | default | flash | speedup |
|---|---|---|---|---|
| 96^3 | fp32 | 12.59 ms | 8.34 ms | 1.51x |
| 96^3 | bf16 | 10.71 ms | 5.46 ms | 1.96x |
| 128^3 | fp32 | 34.27 ms | 21.74 ms | 1.58x |
| 128^3 | bf16 | 29.45 ms | 14.12 ms | 2.09x |
| 160^3 | fp32 | 59.05 ms | 37.66 ms | 1.57x |
| 160^3 | bf16 | 51.25 ms | 25.33 ms | 2.02x |

The flag is threaded through SwinTransformer, BasicLayer and
SwinTransformerBlock to WindowAttention, exactly as use_v2 and
use_checkpoint are. No parameters or buffers change, so pretrained
weights load unchanged. The flash path is used only when autograd is
disabled and the module is not scripted, so training and TorchScript
keep the original path byte-for-byte; this is a deliberate choice to
leave training numerics untouched.

### Types of changes
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [x] New tests added to cover the changes.
- [x] In-line docstrings updated.

---------

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
Signed-off-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>

Author: aymuos15

monai/networks/nets/swin_unetr.py

@@ -98,6 +98,7 @@ def __init__(
         hyena_omega_0: float = 10.0,
         hyena_l_cache: int = 32,
         hyena_short_conv_fft_chunks: int = 0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -149,6 +150,7 @@ def __init__(
             hyena_omega_0: SIREN frequency. Default 10.0 (stable).
             hyena_l_cache: SIREN coordinate-grid cache size per spatial dim.
             hyena_short_conv_fft_chunks: channel chunk size for the FFT short conv (0 = no chunking).
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
 
         Examples::
 
@@ -224,6 +226,7 @@ def __init__(
             hyena_omega_0=hyena_omega_0,
             hyena_l_cache=hyena_l_cache,
             hyena_short_conv_fft_chunks=hyena_short_conv_fft_chunks,
+            use_flash_attention=use_flash_attention,
         )
 
         self.encoder1 = UnetrBasicBlock(
@@ -513,6 +516,7 @@ def __init__(
         qkv_bias: bool = False,
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -522,12 +526,17 @@ def __init__(
             qkv_bias: add a learnable bias to query, key, value.
             attn_drop: attention dropout rate.
             proj_drop: dropout rate of output.
+            use_flash_attention: if True, use ``torch.nn.functional.scaled_dot_product_attention`` for the
+                windowed attention. Equivalent to the default path but faster at inference; only used when
+                autograd is disabled (e.g. under ``torch.no_grad()`` or ``torch.inference_mode()``, not
+                ``eval()`` alone) and the module is not scripted.
         """
 
         super().__init__()
         self.dim = dim
         self.window_size = window_size
         self.num_heads = num_heads
+        self.use_flash_attention = use_flash_attention
         head_dim = dim // num_heads
         self.scale = head_dim**-0.5
         mesh_args = torch.meshgrid.__kwdefaults__
@@ -584,12 +593,26 @@ def forward(self, x, mask):
         b, n, c = x.shape
         qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
-        q = q * self.scale
-        attn = q @ k.transpose(-2, -1)
         relative_position_bias = self.relative_position_bias_table[
             self.relative_position_index.clone()[:n, :n].reshape(-1)  # type: ignore[operator]
         ].reshape(n, n, -1)
         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        if self.use_flash_attention and not torch.jit.is_scripting() and not torch.is_grad_enabled():
+            # additive bias combines the relative position bias and, for shifted windows, the attention mask
+            if mask is not None:
+                nw = mask.shape[0]
+                bias = relative_position_bias.view(1, 1, self.num_heads, n, n) + mask.reshape(1, nw, 1, n, n)
+                bias = bias.expand(b // nw, nw, self.num_heads, n, n).reshape(b, self.num_heads, n, n)
+            else:
+                bias = relative_position_bias.unsqueeze(0)
+            x = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=bias.to(q.dtype), dropout_p=0.0, scale=self.scale
+            )
+            x = x.transpose(1, 2).reshape(b, n, c)
+            return self.proj_drop(self.proj(x))
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
         attn = attn + relative_position_bias.unsqueeze(0)
         if mask is not None:
             nw = mask.shape[0]
@@ -628,6 +651,7 @@ def __init__(
         act_layer: str = "GELU",
         norm_layer: type[LayerNorm] = nn.LayerNorm,
         use_checkpoint: bool = False,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -643,6 +667,7 @@ def __init__(
             act_layer: activation layer.
             norm_layer: normalization layer.
             use_checkpoint: use gradient checkpointing for reduced memory usage.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -660,6 +685,7 @@ def __init__(
             qkv_bias=qkv_bias,
             attn_drop=attn_drop,
             proj_drop=drop,
+            use_flash_attention=use_flash_attention,
         )
 
         self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
@@ -924,6 +950,7 @@ def __init__(
         hyena_omega_0: float = 10.0,
         hyena_l_cache: int = 32,
         hyena_short_conv_fft_chunks: int = 0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -946,6 +973,7 @@ def __init__(
                 hyena_use_chunked_fft, hyena_use_fft_short_conv, hyena_omega_0, hyena_l_cache,
                 hyena_short_conv_fft_chunks: forwarded to :class:`HyenaTransformerBlock`. See its
                 docstring for semantics.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -996,6 +1024,7 @@ def __init__(
                         drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
                         norm_layer=norm_layer,
                         use_checkpoint=use_checkpoint,
+                        use_flash_attention=use_flash_attention,
                     )
                     for i in range(depth)
                 ]
@@ -1079,6 +1108,7 @@ def __init__(
         hyena_omega_0: float = 10.0,
         hyena_l_cache: int = 32,
         hyena_short_conv_fft_chunks: int = 0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -1112,6 +1142,7 @@ def __init__(
                 hyena_use_chunked_fft, hyena_use_fft_short_conv, hyena_omega_0, hyena_l_cache,
                 hyena_short_conv_fft_chunks: HyenaND configuration. See
                 :class:`monai.networks.blocks.HyenaTransformerBlock` for semantics.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -1193,6 +1224,7 @@ def __init__(
                 hyena_omega_0=hyena_omega_0,
                 hyena_l_cache=hyena_l_cache,
                 hyena_short_conv_fft_chunks=hyena_short_conv_fft_chunks,
+                use_flash_attention=use_flash_attention,
             )
             if i_layer == 0:
                 self.layers1.append(layer)

tests/networks/nets/test_swin_unetr.py

@@ -111,6 +111,19 @@ def test_invalid_input_shape(self):
         with self.assertRaises(ValueError):
             net_2d(torch.randn(1, 1, 48, 33))  # 33 is not divisible by 32
 
+    @skipUnless(has_einops, "Requires einops")
+    def test_flash_attention(self):
+        input_param = {"in_channels": 1, "out_channels": 2, "feature_size": 12, "spatial_dims": 3}
+        net_ref = SwinUNETR(use_flash_attention=False, **input_param).double()
+        net_flash = SwinUNETR(use_flash_attention=True, **input_param).double()
+        net_flash.load_state_dict(net_ref.state_dict())
+        x = torch.randn(1, 1, 64, 64, 64, dtype=torch.float64)
+        with eval_mode(net_ref, net_flash):
+            ref = net_ref.swinViT(x, net_ref.normalize)
+            out = net_flash.swinViT(x, net_flash.normalize)
+        for a, b in zip(ref, out, strict=True):
+            assert_allclose(a, b, atol=1e-6, rtol=1e-6, type_test=False)
+
     def test_patch_merging(self):
         dim = 10
         t = PatchMerging(dim)(torch.zeros((1, 21, 20, 20, dim)))