From 8ffd1861c1c2d05c13b2ca31512a1e088efcc3c3 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 2 Jul 2026 22:19:14 +0100 Subject: [PATCH 1/2] Add optional flash attention to SwinUNETR WindowAttention Add a use_flash_attention flag (default False) threaded from SwinUNETR through SwinTransformer, BasicLayer and SwinTransformerBlock to WindowAttention. When enabled and autograd is disabled (and not scripting), attention is computed with torch.nn.functional.scaled_dot_product_attention, folding the relative position bias and the shifted-window mask into a single additive attn_mask cast to the query dtype. The fused kernel avoids materializing the score matrix. Output matches the default path; training, scripting and the default path are unchanged. This mirrors the flash-attention option already in MONAI's SelfAttention, CrossAttention and CABlock. Signed-off-by: Soumya Snigdha Kundu --- monai/networks/nets/swin_unetr.py | 36 +++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/monai/networks/nets/swin_unetr.py b/monai/networks/nets/swin_unetr.py index 0db2d50d26a..7b47da1436b 100644 --- a/monai/networks/nets/swin_unetr.py +++ b/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) From 1713e893ea9eb17567b12f86f8815719205c2465 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 2 Jul 2026 22:19:14 +0100 Subject: [PATCH 2/2] Add SwinUNETR flash attention equivalence test Assert the SwinViT encoder features with use_flash_attention=True match the default path, in float64 for a tight bit-level comparison. Signed-off-by: Soumya Snigdha Kundu --- tests/networks/nets/test_swin_unetr.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/networks/nets/test_swin_unetr.py b/tests/networks/nets/test_swin_unetr.py index ba94aab4f90..7a184bb5320 100644 --- a/tests/networks/nets/test_swin_unetr.py +++ b/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)))