Add setup for Helion to compile on MTIA with basic test #1169
+54
−8
Conversation
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
Summary: Enable Helion with MTIA for the first time. This is very basic setup to just make showcase it working. Following diffs will try to enable more modes (e.g. eager) and more tests. I used default mtia machine arch setup for now + set triton_mtia backend + compile it through JIT flow. Highlights/doubts: - Remove RefEager mode from the test, seems to not be used? - For some reason i have to call mtia_triton_launcher.init() twice (?), otherwise it result with inaccuracy.. i did not debug it fully - There is a check `if TYPE_CHECKING:` in __init__.py, not sure how does this work for functions, can i import triton? will CI be able to check this for me? - I disabled one test when run on mtia. Test complains about different type interfered from Triton-MTIA compiler (error: IntList but got GenericList) - Hardcoded grid size for mtia architectures, I will have to cover this part separately, sadly this is not fully covered in torch by Triton-MTIA, when doing Triton-MTIA update i created hardcoded stub, since it is not being used. v2 We cannot expose arch names to opensource, hide them then in triton_mtia. CPU added support for their backend, use their _get_triton_backend instead of calling triton.runtime.driver.active.get_current_target directly v3 Aligned further to changes made when diff was not landed and i removed custom implementation for handling torch device. Removed calling is_available from torch/mtia/__init__.py Differential Revision: D76375133
meta-cla
bot
added
the
CLA Signed
jansel
requested changes
Nov 26, 2025
| f"PyTorch version {min_version} or higher required", | ||
| ) | ||
|
|
||
|
|
Contributor
There was a problem hiding this comment.
I'd expect this to cause some lints, though for some reason CI isn't running. Can you do ./lint.sh install && ./lint.sh?
|
|
||
| init_mtia_device() | ||
| # Ignore disk cache. Kernels will still keep an in-memory cache. | ||
| os.environ.setdefault("TRITON_ALWAYS_COMPILE", "1") |
Contributor
There was a problem hiding this comment.
Is this required? Seems a bit odd to be running this code on every kernel launch and not cleaning it up.
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.
Summary:
Enable Helion with MTIA for the first time. This is very basic setup to just make showcase it working.
Following diffs will try to enable more modes (e.g. eager) and more tests. I used default mtia machine arch setup for now + set triton_mtia backend + compile it through JIT flow.
Highlights/doubts:
if TYPE_CHECKING:in init.py, not sure how does this work for functions, can i import triton? will CI be able to check this for me?v2
We cannot expose arch names to opensource, hide them then in triton_mtia.
CPU added support for their backend, use their _get_triton_backend instead of calling triton.runtime.driver.active.get_current_target directly
v3
Aligned further to changes made when diff was not landed and i removed custom implementation for handling torch device. Removed calling is_available from torch/mtia/init.py
Differential Revision: D76375133