SwinUNETR: optional flash attention (scaled_dot_product_attention) in WindowAttention

monai/networks/nets/swin_unetr.py

@@ -84,6 +84,7 @@ def __init__(
         spatial_dims: int = 3,
         downsample: str | nn.Module = "merging",
         use_v2: bool = False,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -110,6 +111,7 @@ def __init__(
                 user-specified `nn.Module` following the API defined in :py:class:`monai.networks.nets.PatchMerging`.
                 The default is currently `"merging"` (the original version defined in v0.9.0).
             use_v2: using swinunetr_v2, which adds a residual convolution block at the beggining of each swin stage.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
 
         Examples::
 
@@ -170,6 +172,7 @@ def __init__(
             spatial_dims=spatial_dims,
             downsample=look_up_option(downsample, MERGING_MODE) if isinstance(downsample, str) else downsample,
             use_v2=use_v2,
+            use_flash_attention=use_flash_attention,
         )
 
         self.encoder1 = UnetrBasicBlock(
@@ -457,6 +460,7 @@ def __init__(
         qkv_bias: bool = False,
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -466,12 +470,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__
@@ -528,12 +537,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]
@@ -572,6 +595,7 @@ def __init__(
         act_layer: str = "GELU",
         norm_layer: type[LayerNorm] = nn.LayerNorm,
         use_checkpoint: bool = False,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -587,6 +611,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__()
@@ -604,6 +629,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()
@@ -856,6 +882,7 @@ def __init__(
         norm_layer: type[LayerNorm] = nn.LayerNorm,
         downsample: nn.Module | None = None,
         use_checkpoint: bool = False,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -871,6 +898,7 @@ def __init__(
             norm_layer: normalization layer.
             downsample: an optional downsampling layer at the end of the layer.
             use_checkpoint: use gradient checkpointing for reduced memory usage.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -893,6 +921,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)
             ]
@@ -961,6 +990,7 @@ def __init__(
         spatial_dims: int = 3,
         downsample="merging",
         use_v2=False,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -983,6 +1013,7 @@ def __init__(
                 user-specified `nn.Module` following the API defined in :py:class:`monai.networks.nets.PatchMerging`.
                 The default is currently `"merging"` (the original version defined in v0.9.0).
             use_v2: using swinunetr_v2, which adds a residual convolution block at the beginning of each swin stage.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -1025,6 +1056,7 @@ def __init__(
                 norm_layer=norm_layer,
                 downsample=down_sample_mod,
                 use_checkpoint=use_checkpoint,
+                use_flash_attention=use_flash_attention,
             )
             if i_layer == 0:
                 self.layers1.append(layer)

tests/networks/nets/test_swin_unetr.py

@@ -101,6 +101,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)))