[AIMIGRAPHX-1209] optimize kernel 2 for non kv cache flash decoding and update tests

[bdevorem]

Motivation

Rewrites the flash decoding kernel 2 recombination step in find_flash_decoding to use the exp-normalize form:

O = sum(O' * exp(LSE - max)) / sum(exp(LSE - max))

instead of normalizing weights first, then scaling and summing partial outputs. The result is mathematically equivalent but produces IR that fuses more cleanly downstream (e.g. with rewrite_broadcast in a follow-up PR).

This is a standalone slice extracted from #4546.

Technical Details

After kernel 1, flash decoding returns partial outputs O' and log-sum-exp values LSE per group. Kernel 2 combines them into the final attention output.

The previous IR did:

scale = exp(LSE - max) / sum(exp(LSE - max))
broadcast scale to O' shape
mul(O', scale) → reduce_sum

The new form avoids the intermediate normalized scale broadcast and instead computes numerator and denominator separately:

broadcast exp(LSE - max) to O' shape
numerator = sum(O' * weights)
denominator = sum(weights)
O = numerator / denominator

A follow-up PR will introduce a rewrite_broadcast pass that further optimizes this down the line by rewriting multibroadcast → convert into convert → multibroadcast and multibroadcast → reduce into reduce → multibroadcast, so convert and reduction run on the smaller tensor before broadcasting to the full O' shape.

Opus was used to make this PR description more concise.

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

• • Added: New functionality.
• • Changed: Changes to existing functionality.
• • Removed: Functionality or support that has been removed. (Compared to a previous release)
• • Optimized: Component performance that has been optimized or improved.
• • Resolved Issues: Known issues from a previous version that have been resolved.
• • Not Applicable: This PR is not to be included in the changelog.

Follow the LLVM AI Tool Use Policy for contributions using AI.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #5090      +/-   ##
===========================================
- Coverage    92.90%   92.90%   -0.00%     
===========================================
  Files          603      603              
  Lines        32528    32525       -3     
===========================================
- Hits         30218    30215       -3     
  Misses        2310     2310              

Files with missing lines	Coverage Δ
src/fuse_attention.cpp	98.13% <100.00%> (-0.01%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[gh-app-migraphx-bot-pr-write]

Test	Batch	New Rate (241080)	Old Rate (0c63af)*	Diff	Status
torchvision-resnet50	64	3,278.72	3,266.06	0.39%	✅
torchvision-resnet50_fp16	64	7,613.08	7,546.56	0.88%	✅
torchvision-densenet121	32	2,501.06	2,482.18	0.76%	✅
torchvision-densenet121_fp16	32	5,009.68	4,976.34	0.67%	✅
torchvision-inceptionv3	32	2,058.58	2,057.16	0.07%	✅
torchvision-inceptionv3_fp16	32	4,418.52	4,415.87	0.06%	✅
cadene-inceptionv4	16	821.51	818.77	0.33%	✅
cadene-resnext64x4	16	783.40	783.71	-0.04%	✅
slim-mobilenet	64	8,394.00	8,356.38	0.45%	✅
slim-nasnetalarge	64	230.18	229.42	0.33%	✅
slim-resnet50v2	64	3,201.34	3,185.54	0.50%	✅
bert-mrpc-onnx	8	1,167.75	1,167.46	0.02%	✅
bert-mrpc-tf	1	480.46	486.41	-1.22%	✅
pytorch-examples-wlang-gru	1	483.25	480.59	0.55%	✅
pytorch-examples-wlang-lstm	1	388.07	390.19	-0.54%	✅
torchvision-resnet50_1	1	1,028.27	1,037.04	-0.85%	✅
cadene-dpn92_1	1	440.25	475.84	-7.48%	🔴
cadene-resnext101_1	1	364.62	365.49	-0.24%	✅
onnx-taau-downsample	1	848.21	843.97	0.50%	✅
dlrm-criteoterabyte	1	32.59	32.42	0.52%	✅
dlrm-criteoterabyte_fp16	1	52.86	52.07	1.52%	✅
agentmodel	1	8,060.62	9,645.99	-16.44%	🔴
unet_fp16	2	59.16	58.92	0.41%	✅
resnet50v1_fp16	1	1,354.42	1,412.17	-4.09%	✅
resnet50v1_int8	1	1,656.51	1,602.24	3.39%	✅
bert_base_cased_fp16	64	1,103.60	1,098.36	0.48%	✅
bert_large_uncased_fp16	32	345.66	344.75	0.26%	✅
bert_large_fp16	1	205.44	204.70	0.36%	✅
distilgpt2_fp16	16	2,101.34	2,090.27	0.53%	✅
yolov5s	1	561.53	560.47	0.19%	✅
tinyllama	1	45.79	45.79	0.02%	✅
vicuna-fastchat	1	43.84	44.12	-0.64%	✅
whisper-tiny-encoder	1	419.99	412.35	1.85%	✅
whisper-tiny-decoder	1	416.16	413.51	0.64%	✅
llama2_7b	1	20.99	20.84	0.74%	✅
qwen1.5-7b	1	23.59	23.52	0.29%	✅
phi3-3.8b	1	26.87	26.73	0.50%	✅
llama3-8b	1	21.78	21.66	0.54%	✅
whisper-large-encoder	1	10.22	10.18	0.33%	✅
whisper-large-decoder	1	105.89	106.43	-0.51%	✅
mistral-7b	1	23.77	23.62	0.63%	✅
FLUX.1-schnell	1	720.17	756.55	-4.81%	✅

Regressions detected 🔴

* No develop baseline was found for this PR's branch point; compared against the latest available develop run instead.

[gh-app-migraphx-bot-pr-write]

Test	Status	Result
bert-mrpc-onnx	✅	PASSED: MIGraphX meets tolerance
bert-mrpc-tf	✅	PASSED: MIGraphX meets tolerance
pytorch-examples-wlang-gru	✅	PASSED: MIGraphX meets tolerance
pytorch-examples-wlang-lstm	✅	PASSED: MIGraphX meets tolerance
dlrm-criteoterabyte	✅	PASSED: MIGraphX meets tolerance
agentmodel	✅	PASSED: MIGraphX meets tolerance
unet	✅	PASSED: MIGraphX meets tolerance
resnet50v1	✅	PASSED: MIGraphX meets tolerance
bert_base_cased_fp16	✅	PASSED: MIGraphX meets tolerance
bert_large_uncased_fp16	🔴	FAILED: MIGraphX is not within tolerance - check verbose output
bert_large	✅	PASSED: MIGraphX meets tolerance
yolov5s	✅	PASSED: MIGraphX meets tolerance
tinyllama	✅	PASSED: MIGraphX meets tolerance
vicuna-fastchat	✅	PASSED: MIGraphX meets tolerance
whisper-tiny-encoder	✅	PASSED: MIGraphX meets tolerance
whisper-tiny-decoder	✅	PASSED: MIGraphX meets tolerance
distilgpt2_fp16	🔴	FAILED: MIGraphX is not within tolerance - check verbose output
llama2_7b	✅	PASSED: MIGraphX meets tolerance
qwen1.5-7b	✅	PASSED: MIGraphX meets tolerance
phi3-3.8b	✅	PASSED: MIGraphX meets tolerance
llama3-8b	✅	PASSED: MIGraphX meets tolerance
whisper-large-encoder	✅	PASSED: MIGraphX meets tolerance
whisper-large-decoder	✅	PASSED: MIGraphX meets tolerance
mistral-7b	✅	PASSED: MIGraphX meets tolerance
FLUX.1-schnell	✅	PASSED: MIGraphX meets tolerance

[Copilot]

Pull request overview

Optimizes the flash-decoding “kernel 2” recombination logic inside find_flash_decoding (in fuse_attention) by switching to an exp-normalize formulation that produces IR more amenable to downstream fusion/rewrites, and updates the corresponding attention fusion tests and changelog entry.

Changes:

• Rewrites flash-decoding recombination to compute numerator/denominator via unnormalized weights exp(LSE - max) and divide.
• Updates multiple flash-decoding test expectations to match the new kernel-2 IR structure.
• Adds a changelog note for the optimization.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
src/fuse_attention.cpp	Reworks flash-decoding kernel-2 recombination IR to exp-normalize form for better downstream fusion.
test/fuse_attention.cpp	Updates expected graphs/IR construction in flash-decoding tests to reflect the new recombination form.
CHANGELOG.md	Documents the flash-decoding recombination optimization.

Comments suppressed due to low confidence (1)

src/fuse_attention.cpp:776

• The shape-specific comment "[B, G, M] -> [B, G, M, D]" is misleading here because the actual LSE/partial-output ranks vary (e.g., LSE often has a trailing singleton dim and 4D attention adds a head dim). Consider making this comment rank-agnostic and describing the intent (broadcast to match O' for elementwise mul).

        // broadcast weights to match O' shape
        // [B, G, M] -> [B, G, M, D]