Skip to content

Fix for MHA in attn refactor #4152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jul 23, 2025
Merged

Fix for MHA in attn refactor #4152

merged 5 commits into from
Jul 23, 2025

Conversation

shivadbhavsar
Copy link
Contributor

Update the algorithm in fuse_attention for finding the attention subgraph.

  • Current implementation was fusing unwanted instructions into the attention subgraph (eg. pointwise instructions that have multiple outputs)
  • Implement find_instructions_between routine that only gives instructions directly connected to both start and end node
  • Any additional fusions are handled already by the fuse_mlir pass

@shivadbhavsar shivadbhavsar self-assigned this Jul 18, 2025
@shivadbhavsar shivadbhavsar requested a review from causten as a code owner July 18, 2025 19:14
@shivadbhavsar shivadbhavsar added the bugfix Fixes a bug found in the code. label Jul 18, 2025
@shivadbhavsar shivadbhavsar linked an issue Jul 18, 2025 that may be closed by this pull request
@shivadbhavsar shivadbhavsar requested a review from CharlieL7 July 21, 2025 15:36
Copy link

codecov bot commented Jul 21, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Additional details and impacted files
@@             Coverage Diff             @@
##           develop    #4152      +/-   ##
===========================================
+ Coverage    92.23%   92.24%   +0.01%     
===========================================
  Files          548      548              
  Lines        25187    25185       -2     
===========================================
  Hits         23230    23230              
+ Misses        1957     1955       -2     
Files with missing lines Coverage Δ
src/fuse_attention.cpp 98.18% <100.00%> (+2.87%) ⬆️
src/include/migraphx/instruction.hpp 100.00% <ø> (ø)
src/instruction.cpp 89.05% <100.00%> (+0.66%) ⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Collaborator

@kahmed10 kahmed10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm, please fix the cppcheck issue

@shivadbhavsar shivadbhavsar requested a review from CharlieL7 July 22, 2025 21:35
@migraphx-bot
Copy link
Collaborator

Test Batch Rate new
00e4b4
Rate old
408b98
Diff Compare
torchvision-resnet50 64 3,249.36 3,249.84 -0.01%
torchvision-resnet50_fp16 64 6,936.29 6,927.96 0.12%
torchvision-densenet121 32 2,451.47 2,450.72 0.03%
torchvision-densenet121_fp16 32 4,181.88 4,199.41 -0.42%
torchvision-inceptionv3 32 1,635.94 1,638.39 -0.15%
torchvision-inceptionv3_fp16 32 2,757.36 2,755.49 0.07%
cadene-inceptionv4 16 771.60 772.10 -0.06%
cadene-resnext64x4 16 814.55 818.89 -0.53%
slim-mobilenet 64 7,470.48 7,467.98 0.03%
slim-nasnetalarge 64 211.21 211.20 0.00%
slim-resnet50v2 64 3,339.32 3,342.30 -0.09%
bert-mrpc-onnx 8 1,147.88 1,148.63 -0.07%
bert-mrpc-tf 1 457.25 457.81 -0.12%
pytorch-examples-wlang-gru 1 343.88 342.34 0.45%
pytorch-examples-wlang-lstm 1 477.06 466.93 2.17%
torchvision-resnet50_1 1 789.45 793.10 -0.46%
cadene-dpn92_1 1 414.51 413.32 0.29%
cadene-resnext101_1 1 387.75 392.78 -1.28%
onnx-taau-downsample 1 395.73 396.46 -0.18%
dlrm-criteoterabyte 1 33.77 33.77 -0.01%
dlrm-criteoterabyte_fp16 1 51.18 51.15 0.05%
agentmodel 1 10,157.81 10,131.48 0.26%
unet_fp16 2 60.67 60.68 -0.01%
resnet50v1_fp16 1 1,031.78 1,037.89 -0.59%
resnet50v1_int8 1 1,059.20 1,046.81 1.18%
bert_base_cased_fp16 64 1,170.02 1,169.82 0.02%
bert_large_uncased_fp16 32 361.32 361.33 -0.00%
bert_large_fp16 1 202.22 203.50 -0.63%
distilgpt2_fp16 16 2,240.61 2,239.25 0.06%
yolov5s 1 540.90 540.44 0.09%
tinyllama 1 43.97 44.01 -0.09%
vicuna-fastchat 1 45.38 45.32 0.13%
whisper-tiny-encoder 1 419.14 419.11 0.01%
whisper-tiny-decoder 1 411.22 402.75 2.10%
llama2_7b 1 19.18 19.18 -0.00%
qwen1.5-7b 1 23.68 23.65 0.12%
phi3-3.8b 1 26.89 26.86 0.11%
mask-rcnn 1 12.82 12.82 0.00%
llama3-8b 1 21.85 21.84 0.07%
whisper-large-encoder 1 10.22 10.22 -0.00%
whisper-large-decoder 1 103.53 103.52 0.02%
mistral-7b 1 23.87 23.83 0.16%
FLUX.1-schnell 1 766.31 767.72 -0.18%
nan nan nan nan nan%

This build is not recommended to merge 🔴

@migraphx-bot
Copy link
Collaborator


     ✅ bert-mrpc-onnx: PASSED: MIGraphX meets tolerance

❌bert-mrpc-tf: ERROR - check error outputerror: unknown warning option '-Wnrvo' [-Werror,-Wunknown-warning-option]

error: unknown warning option '-Wnrvo' [-Werror,-Wunknown-warning-option]

2025-07-23 07:18:22.729960: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1753273108.051094 183252 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 62973 MB memory: -> device: 0, name: AMD Instinct MI250X/MI250, pci bus id: 0000:b3:00.0
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1753273108.928082 183252 mlir_graph_optimization_pass.cc:401] MLIR V1 optimization pass is not enabled
2025-07-23 07:18:38.784485: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-07-23 07:18:38.784891: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-07-23 07:18:38.784944: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-07-23 07:18:38.784990: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-07-23 07:18:38.785030: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-07-23 07:18:38.785079: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-07-23 07:18:38.785131: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-07-23 07:18:38.785284: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
2025-07-23 07:18:38.786262: E tensorflow/compiler/mlir/tools/kernel_gen/tf_framework_c_interface.cc:228] INTERNAL: Generating device code failed.
2025-07-23 07:18:38.787385: W tensorflow/core/framework/op_kernel.cc:1829] UNKNOWN: JIT compilation failed.
2025-07-23 07:18:38.787405: I tensorflow/core/framework/local_rendezvous.cc:405] Local rendezvous is aborting with status: UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
2025-07-23 07:18:38.787417: I tensorflow/core/framework/local_rendezvous.cc:405] Local rendezvous is aborting with status: UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
[[import/loss/output/_21]]
2025-07-23 07:18:38.787459: I tensorflow/core/framework/local_rendezvous.cc:424] Local rendezvous recv item cancelled. Key hash: 11217777527359497193
Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1407, in _do_call
return fn(*args)
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1390, in _run_fn
return self._call_tf_sessionrun(options, feed_dict, fetch_list,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1483, in _call_tf_sessionrun
return tf_session.TF_SessionRun_wrapper(self._session, options, feed_dict,
tensorflow.python.framework.errors_impl.UnknownError: 2 root error(s) found.
(0) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
[[import/loss/output/_21]]
(1) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
0 successful operations.
0 derived errors ignored.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 359, in
main()
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 335, in main
y_out = sess.run(y, feed_dict=tf_dict)
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 977, in run
result = self._run(None, fetches, feed_dict, options_ptr,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1220, in _run
results = self._do_run(handle, final_targets, final_fetches,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1400, in _do_run
return self._do_call(_run_fn, feeds, fetches, targets, options,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1426, in _do_call
raise type(e)(node_def, op, message) # pylint: disable=no-value-for-parameter
tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:

Detected at node 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' defined at (most recent call last):
Node: 'import/bert/embeddings/LayerNorm/moments/SquaredDifference'
Detected at node 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' defined at (most recent call last):
Node: 'import/bert/embeddings/LayerNorm/moments/SquaredDifference'
2 root error(s) found.
(0) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
[[import/loss/output/_21]]
(1) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
0 successful operations.
0 derived errors ignored.

Original stack trace for 'import/bert/embeddings/LayerNorm/moments/SquaredDifference':


     ✅ 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: FAILED: MIGraphX is not within tolerance - check verbose output


     ✅ 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: PASSED: MIGraphX meets tolerance

     ✅ llama2_7b: PASSED: MIGraphX meets tolerance

     ✅ qwen1.5-7b: PASSED: MIGraphX meets tolerance

     ✅ phi3-3.8b: PASSED: MIGraphX meets tolerance

🔴mask-rcnn: FAILED: MIGraphX is not within tolerance - check verbose output


     ✅ llama3-8b: PASSED: MIGraphX meets tolerance

     ✅ whisper-large-decoder: PASSED: MIGraphX meets tolerance

     ✅ mistral-7b: PASSED: MIGraphX meets tolerance

     ✅ FLUX.1-schnell: PASSED: MIGraphX meets tolerance

@causten causten merged commit 3281eaa into develop Jul 23, 2025
40 checks passed
@causten causten deleted the mha_fix branch July 23, 2025 15:18
@@ -578,9 +578,9 @@ bool reaches(instruction_ref start, instruction_ref end, const_module_ref m)
{
if(start == end)
return true;
if(not m->has_instruction(start) or not m->has_instruction(end))
if(not m->has_instruction(start) or not m->has_instruction(end) or
std::distance(m->begin(), start) > std::distance(m->begin(), end))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should not be added here, as this traverses the entire program which can be really slow. The algorithm for find_instructions_between needs to be updated to check for this since it can check a smaller subset of the program.

causten pushed a commit that referenced this pull request Jul 26, 2025
Update the algorithm in fuse_attention for finding the attention subgraph.

Current implementation was fusing unwanted instructions into the attention subgraph (eg. pointwise instructions that have multiple outputs)
Implement find_instructions_between routine that only gives instructions directly connected to both start and end node
Any additional fusions are handled already by the fuse_mlir pass
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
backport bugfix Fixes a bug found in the code.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Attention refactor bug with MHA
7 participants