[amarbarosys]

Motivation

Closes #255.

In inference, nn.TransformerEncoderLayer takes a fast path that dynamo records as one aten._transformer_encoder_layer_fwd node, and code calling nn.MultiheadAttention's fast path records aten._native_multi_head_attention (#255). Neither has a converter or a torch decomposition, so every encoder layer stays in torch and the model is split into separate MIGraphX partitions.

On torch 2.13 the #255 example (vit_b_16) no longer emits the fused MHA op (it reaches the backend as SDPA with the fast path on or off); the encoder layer is where it still appears. The suggestion on #255, torch.backends.mha.set_fastpath_enabled(False), acts at dynamo trace time, so inside the backend it is too late.

Technical Details

• dynamo/decompositions.py: decompositions for both ops (packed qkv projection, heads, SDPA, or an explicit softmax when attention weights are requested, output projection; layer norms, residuals and FFN for the encoder layer).
• backends.py: passed to the AOT retrace in both backends (aot_export_joint_simple and aot_autograd take decompositions since torch 2.1, the oldest torch the dynamo backend loads on).
• The graph that reaches the converters then holds linear, layer_norm and scaled_dot_product_attention, which already lower.
• No breaking changes: no new option or API; graphs without these two ops are unchanged.

Test Plan

• tests/dynamo/test_decompositions_dynamo.py (CPU): each decomposition against the real aten op over mask types 0/1/2, bool and float masks, need_weights, average, norm_first, relu/gelu.
• tests/dynamo/test_fused_attention_dynamo.py (GPU): encoder layer and the direct MHA op in fp16, use_aot on and off, compared with eager, asserting nothing is left to torch. Run on master as the control.
• End to end: torch.compile of vit_b_16 and a 6-layer nn.TransformerEncoder, fp16, both backends, vs eager.
• Full tests/dynamo; latency of 6 encoder layers vs eager.
  Setup: RX 7900 XTX (gfx1100), ROCm 7.2, torch 2.13, MIGraphX develop.

Test Result

• Decompositions: 42 passed.
• End to end: all 4 cases match eager on this branch (max abs diff 4.9e-3 vit_b_16, 9.8e-3 encoder); on master vit_b_16 passes (already SDPA), the encoder stays in torch with the default backend, and with use_aot fails to compile (AttributeError: 'GraphModule' object has no attribute '_frozen_param0').
• Fused attention: 20 passed, 4 strict xfails; the same 20 fail on master. The xfails are the encoder layer with a padding mask on the non-AOT path: F._canonical_mask's in-place masked_fill_ is rejected by aot_export_joint_simple, also on master (separate issue).
• Full tests/ (fx + dynamo): 906 passed, 3 failed, 14 skipped, 6 xfailed (run with pytest-xdist -n 8). 2 failures are test_quant_LLM (transformers DynamicCache), same on master, fixed in Fix 0-d new_zeros/new_ones in remove_const_ops; run test_quant_LLM without the KV cache #296. 1 is an encoder case that read ~/.cache/migraphx/problem.json while another worker was writing it; run alone it passes.

Performance

Median ms per forward, 6x TransformerEncoderLayer (E512, 8 heads, FFN 2048), fp16, RX 7900 XTX, default MIGraphX settings. On master the encoder layers stay in torch, so master runs at about eager speed.

• Batch 8, length 128, 50 runs: eager 1.300, master 1.338, this PR 1.047 (1.24x vs eager).
• Batch 1, one run per point:

length	eager / master	this PR
512	1.07	0.93
1024	1.65	1.37
1500	2.23	1.83
2048	3.19	3.57
4096	7.00	13.01
8192	21.87	46.10
15000	51.29	145.75

From 2k tokens this PR is slower than eager on gfx1100. That is MIGraphX's attention on gfx11, not the decomposition: MIGraphX fuses the decomposed pattern, but its MLIR attention is on by default only on gfx94/gfx95, so gfx11 materialises the full score matrix (tracked in ROCm/AMDMIGraphX#5396).

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.
• Code builds and tests pass locally
• New/changed functionality has test coverage
• Commit messages are clear and use imperative mood
• Documentation updated if applicable (no user-facing option or env var added)