SwinUNETR: fused attention + channels-last + bf16 for AMD MI300X inference#1
Open
nilapate wants to merge 3 commits into
Open
SwinUNETR: fused attention + channels-last + bf16 for AMD MI300X inference#1nilapate wants to merge 3 commits into
nilapate wants to merge 3 commits into
Conversation
…bf16 inferer.py (SlidingWindowInferer.__call__): - channels_last_3d conversion on 5-D inputs on ROCm (torch.version.hip is not None) - bf16 autocast override (bundle evaluator defaults to fp16; MI300X wedge uses bf16) applied only when autocast is already active (fp32 golden runs are untouched) utils.py (sliding_window_inference): - maybe_mark_dynamic(win_data, 0) before the compiled predictor on ROCm bounds torch.compile recompiles to ~2 over a full multi-image run swin_unetr.py (WindowAttention): - _forward_sdpa: fused scaled_dot_product_attention path (AOTriton EFFICIENT backend on ROCm; (b,H,n,n) intermediate never materializes in HBM) - Gated on PYTORCH_MIOPEN_SUGGEST_NHWC env var (refactored in next commit) Signed-off-by: Patel, Nilaykumar K <nilapate@amd.com>
Refactor WindowAttention to eliminate ~30 lines of duplicated scaffolding
(QKV projection, RPB lookup, output proj) shared between the two attention
paths. Only the 5-line attention kernel differs.
swin_unetr.py:
- Unified forward() with shared QKV + RPB + output-proj scaffolding
- _attn_explicit / _attn_sdpa: narrow kernel-only helpers
- Gate moved from os.environ.get("PYTORCH_MIOPEN_SUGGEST_NHWC") to
torch.version.hip is not None, checked once in __init__ as self._use_sdpa
(trace-time constant — no graph break under torch.compile)
- Drop unnecessary .clone() on relative_position_index (registered buffer,
never written after __init__)
- Add use_sdpa: bool | None = None constructor param for explicit override
- Remove now-unused import os
inferers/utils.py:
- maybe_mark_dynamic gate: PYTORCH_MIOPEN_SUGGEST_NHWC env var replaced with
torch.version.hip is not None (authoritative ROCm discriminant, no env-var
in the compiled path)
Signed-off-by: Patel, Nilaykumar K <nilapate@amd.com>
- Add test_window_attention_sdpa() to validate numerical equivalence - Follows same pattern as SABlock, CABlock, CrossAttentionBlock tests - Validates WindowAttention SDPA vs explicit attention with atol=1e-4 - Covers 3D medical imaging windows (7x7x7) Signed-off-by: Patel, Nilaykumar K <nilapate@amd.com>
soumitra-chatterjee-tech
approved these changes
Jul 3, 2026
| from __future__ import annotations | ||
|
|
||
| import itertools | ||
| import os |
There was a problem hiding this comment.
The changes don't suggest this is needed - can you double check?
|
|
||
| bias = ( | ||
| self.relative_position_bias_table[ | ||
| self.relative_position_index[:n, :n].reshape(-1) # type: ignore[operator] |
There was a problem hiding this comment.
Did you drop the .clone() intentionally? Wouldn't it cause side effects elsewhere?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe the change
Adds AMD MI300X inference optimizations to SwinUNETR.
Description
Adds ROCm-specific optimizations to SwinUNETR inference: fused scaled_dot_product_attention in WindowAttention, channels_last layout and bf16 autocast in SlidingWindowInferer, and maybe_mark_dynamic to bound torch.compile recompiles to ~2 per run. All changes gate on torch.version.hip is not None — strict no-op on CUDA and CPU builds.
PR Checklist
Commit quality
Change type and validation
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmltested in thedocs/folder.