SwinUNETR: fused attention + channels-last + bf16 for AMD MI300X inference

monai/inferers/utils.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import itertools
+import os
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from typing import Any
 
@@ -258,6 +259,8 @@ def sliding_window_inference(
                 win_condition = condition[s0_idx].to(sw_device)
                 kwargs["condition"] = win_condition
 
+        if torch.version.hip is not None:
+            torch._dynamo.maybe_mark_dynamic(win_data, 0)
         if with_coord:
             seg_prob_out = predictor(win_data, unravel_slice, *args, **kwargs)
         else:

monai/networks/nets/swin_unetr.py

@@ -457,6 +457,7 @@ def __init__(
         qkv_bias: bool = False,
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
+        use_sdpa: bool | None = None,
     ) -> None:
         """
         Args:
@@ -466,6 +467,8 @@ def __init__(
             qkv_bias: add a learnable bias to query, key, value.
             attn_drop: attention dropout rate.
             proj_drop: dropout rate of output.
+            use_sdpa: use fused scaled_dot_product_attention. Defaults to True on ROCm (AMD),
+                False elsewhere. Pass explicitly to override.
         """
 
         super().__init__()
@@ -523,31 +526,53 @@ def __init__(
         self.proj_drop = nn.Dropout(proj_drop)
         trunc_normal_(self.relative_position_bias_table, std=0.02)
         self.softmax = nn.Softmax(dim=-1)
+        if use_sdpa is None:
+            use_sdpa = torch.version.hip is not None and hasattr(F, "scaled_dot_product_attention")
+        self._use_sdpa: bool = use_sdpa
 
     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()
-        attn = attn + relative_position_bias.unsqueeze(0)
+
+        bias = (
+            self.relative_position_bias_table[
+                self.relative_position_index[:n, :n].reshape(-1)  # type: ignore[operator]
+            ]
+            .reshape(n, n, -1)
+            .permute(2, 0, 1)
+            .contiguous()
+            .unsqueeze(0)
+        )  # (1, num_heads, n, n)
+
+        if self._use_sdpa:
+            out = self._attn_sdpa(q, k, v, bias, mask, b, n)
+        else:
+            out = self._attn_explicit(q, k, v, bias, mask, b, n)
+
+        return self.proj_drop(self.proj(out.transpose(1, 2).reshape(b, n, c)))
+
+    def _attn_explicit(self, q, k, v, bias, mask, b, n):
+        attn = q * self.scale @ k.transpose(-2, -1) + bias
         if mask is not None:
             nw = mask.shape[0]
             attn = attn.view(b // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)
             attn = attn.view(-1, self.num_heads, n, n)
-            attn = self.softmax(attn)
-        else:
-            attn = self.softmax(attn)
+        return self.attn_drop(self.softmax(attn)).to(v.dtype) @ v
 
-        attn = self.attn_drop(attn).to(v.dtype)
-        x = (attn @ v).transpose(1, 2).reshape(b, n, c)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
+    def _attn_sdpa(self, q, k, v, bias, mask, b, n):
+        if mask is not None:
+            nw = mask.shape[0]
+            attn_mask = (
+                (bias.unsqueeze(0) + mask.view(1, nw, 1, n, n))
+                .expand(b // nw, nw, self.num_heads, n, n)
+                .reshape(b, self.num_heads, n, n)
+                .to(q.dtype)
+            )
+        else:
+            attn_mask = bias.to(q.dtype)
+        drop_p = self.attn_drop.p if self.training else 0.0
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_p, scale=self.scale)
 
 
 class SwinTransformerBlock(nn.Module):

tests/networks/nets/test_swin_unetr.py

@@ -21,7 +21,7 @@
 
 from monai.apps import download_url
 from monai.networks import eval_mode
-from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR, filter_swinunetr
+from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR, WindowAttention, filter_swinunetr
 from monai.networks.utils import copy_model_state
 from monai.utils import optional_import
 from tests.test_utils import (
@@ -125,6 +125,36 @@ def test_filter_swinunetr(self, input_param, key, value):
                 assert_allclose(dst_dict[key][:8], value, atol=1e-4, rtol=1e-4, type_test=False)
                 self.assertTrue(len(loaded) == 157 and len(not_loaded) == 2)
 
+    @skipUnless(has_einops, "Requires einops")
+    def test_window_attention_sdpa(self):
+        """Test WindowAttention SDPA vs explicit attention numerical equivalence."""
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        dim = 96
+        num_heads = 3
+        window_size = (7, 7, 7)
+
+        # Create two WindowAttention blocks with different attention modes
+        attn_sdpa = WindowAttention(dim=dim, num_heads=num_heads, window_size=window_size, use_sdpa=True).to(device)
+        attn_explicit = WindowAttention(dim=dim, num_heads=num_heads, window_size=window_size, use_sdpa=False).to(
+            device
+        )
+
+        # Share weights to ensure only attention mechanism differs
+        attn_explicit.load_state_dict(attn_sdpa.state_dict())
+        attn_sdpa.eval()
+        attn_explicit.eval()
+
+        # Test input
+        b = 2
+        n = window_size[0] * window_size[1] * window_size[2]  # 343
+        x = torch.randn(b, n, dim).to(device)
+
+        with torch.no_grad():
+            out_sdpa = attn_sdpa(x, mask=None)
+            out_explicit = attn_explicit(x, mask=None)
+
+        assert_allclose(out_sdpa, out_explicit, atol=1e-4)
+
 
 if __name__ == "__main__":
     unittest.main()