diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py index de53108d1d..7a8fccede0 100644 --- a/monai/inferers/utils.py +++ b/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: diff --git a/monai/networks/nets/swin_unetr.py b/monai/networks/nets/swin_unetr.py index 0db2d50d26..dd8d9325ec 100644 --- a/monai/networks/nets/swin_unetr.py +++ b/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): diff --git a/tests/networks/nets/test_swin_unetr.py b/tests/networks/nets/test_swin_unetr.py index ba94aab4f9..14ba674c74 100644 --- a/tests/networks/nets/test_swin_unetr.py +++ b/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()