Skip to content
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

Attention fusion (part 2) #2013

Merged
merged 45 commits into from
Jan 24, 2025
Merged

Attention fusion (part 2) #2013

merged 45 commits into from
Jan 24, 2025

Conversation

gramalingam
Copy link
Collaborator

@gramalingam gramalingam commented Jan 16, 2025

Continuation of attention fusion.

  • Adds a version of GroupQueryAttention
  • Adds support in Cos-Sin cache fusion for constant-folded position-ids
  • Restructure MHA fusion into a class-based rewrite rule

Also restructure the folder structure.

  • Eventually eliminate folders called "onnxruntime" and "transfomers", which hinder importing the original packages with those names. For now moving just the relevant new files. (Will restructure older files later.)
  • ORT-specific fusions go into the ort_fusions folder.

@gramalingam gramalingam marked this pull request as draft January 16, 2025 21:48
onnxscript/rewriter/ort_fusions/__init__.py Fixed Show fixed Hide fixed
onnxscript/rewriter/ort_fusions/__init__.py Fixed Show fixed Hide fixed
onnxscript/rewriter/ort_fusions/_core.py Fixed Show fixed Hide fixed
onnxscript/rewriter/ort_fusions/_core.py Fixed Show fixed Hide fixed
onnxscript/rewriter/ort_fusions/_core.py Fixed Show fixed Hide fixed
onnxscript/rewriter/ort_fusions/_core.py Fixed Show fixed Hide fixed
onnxscript/rewriter/ort_fusions/_core.py Fixed Show fixed Hide fixed
Copy link

codecov bot commented Jan 16, 2025

❌ 65 Tests Failed:

Tests completed Failed Passed Skipped
11931 65 11866 2459
View the top 3 failed tests by shortest run time
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0263_test_concat_1d_axis_0
Stack Traces | 0.003s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.10.11\x64\lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_concat_1d_axis_0'

The above exception was the direct cause of the following exception:
.nox\test\lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_concat_1d_axis_0' (e=No module named 'tests.onnx_backend_test_code.test_concat_1d_axis_0') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_concat_1d_axis_0.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_concat_1d_axis_0.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_concat_1d_axis_0(value0: FLOAT[2], value1: FLOAT[2]) -> (FLOAT[4]):
E       output = opset13.Concat(value0, value1, axis=0)
E       return output
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1101_test_shape_end_negative_1
Stack Traces | 0.003s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.10.11\x64\lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_shape_end_negative_1'

The above exception was the direct cause of the following exception:
.nox\test\lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_shape_end_negative_1' (e=No module named 'tests.onnx_backend_test_code.test_shape_end_negative_1') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_end_negative_1.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_end_negative_1.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset21
E   
E   @script()
E   def bck_test_shape_end_negative_1(x: FLOAT[3,4,5]) -> (INT64[2]):
E       y = opset21.Shape(x, end=-1)
E       return y
onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1257_test_triu_one_row
Stack Traces | 0.003s run time
onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.10.11\x64\lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_triu_one_row'

The above exception was the direct cause of the following exception:
.nox\test\lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_triu_one_row' (e=No module named 'tests.onnx_backend_test_code.test_triu_one_row') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_triu_one_row.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_triu_one_row.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import INT64
E   from onnxscript.onnx_opset import opset14
E   
E   @script()
E   def bck_test_triu_one_row(x: INT64[3,1,5], k: INT64) -> (INT64[3,1,5]):
E       y = opset14.Trilu(x, k)
E       return y

To view more test analytics, go to the Test Analytics Dashboard
📢 Thoughts on this report? Let us know!

@gramalingam gramalingam changed the title [WIP] Attention fusion (part 2) Attention fusion (part 2) Jan 18, 2025
@gramalingam gramalingam marked this pull request as ready for review January 18, 2025 01:15
Comment on lines +140 to +143
# if bindings["B"] * bindings["H"] != bindings["B*H"]:
# return False
# if bindings["H"] * bindings["d_h"] != bindings["H*d_h"]:
# return False

Check notice

Code scanning / CodeQL

Commented-out code Note

This comment appears to contain commented-out code.
@gramalingam gramalingam enabled auto-merge (squash) January 24, 2025 17:04
@gramalingam gramalingam merged commit d44853e into main Jan 24, 2025
22 of 29 checks passed
@gramalingam gramalingam deleted the rama/fuse3 branch January 24, 2025 17:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
Development

Successfully merging this pull request may close these issues.

2 participants