Add new apex module to jit load system#294
Merged
amd-sriram merged 16 commits intomasterfrom Feb 18, 2026
Merged
Conversation
… builder and parameter to supply builder module name
…odule with builder provided
…odule with builder provided
…me from the builder, update missing entries in the table in readme for adding new module in jit
|
I am not sure of the request here, |
…and fields in the initial builder code created by script
Collaborator
Author
|
! cherry-pick --onto release/1.10.0 |
amd-sriram
added a commit
that referenced
this pull request
Feb 19, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
amd-sriram
added a commit
that referenced
this pull request
Feb 19, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
amd-sriram
added a commit
that referenced
this pull request
Mar 2, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
amd-sriram
added a commit
that referenced
this pull request
Mar 2, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
amd-sriram
added a commit
that referenced
this pull request
Mar 2, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
amd-sriram
added a commit
that referenced
this pull request
Mar 2, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
This was referenced Mar 2, 2026
jithunnair-amd
pushed a commit
to ROCm/pytorch
that referenced
this pull request
Mar 3, 2026
Commit Messages: - Update release notes for clarity and conciseness (ROCm/apex#311) Removed redundant information and clarified release notes. - Add new apex module to jit load system (ROCm/apex#294) (ROCm/apex#302) Co-authored-by: sriram <sriram.kumar@silo.ai>
jithunnair-amd
pushed a commit
to ROCm/pytorch
that referenced
this pull request
Mar 3, 2026
Commit Messages: - Update release notes (ROCm/apex#312) - Add new apex module to jit load system (ROCm/apex#294) (ROCm/apex#307) Co-authored-by: sriram <sriram.kumar@silo.ai>
jithunnair-amd
pushed a commit
to ROCm/pytorch
that referenced
this pull request
Mar 3, 2026
Commit Messages: - release notes for 1.8.0 (ROCm/apex#313) Updated release notes to reflect changes for version 1.8.0 including unit test fixes and new features. - Add new apex module to jit load system (ROCm/apex#294) (ROCm/apex#308) Co-authored-by: sriram <sriram.kumar@silo.ai>
amd-sriram
added a commit
that referenced
this pull request
Mar 4, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
amd-sriram
added a commit
that referenced
this pull request
Mar 4, 2026
* add code to add loader module for jit module * fix errors to create jit module adder - use correct file name to save code to * fix errors to create jit module adder - use correct class name of the builder and parameter to supply builder module name * fix errors to create jit module loader * add description about jit module script to add jit loader for a jit module with builder provided * add description about jit module script to add jit loader for a jit module with builder provided * add attributes and methods to override when creating a jit module builder * add extra new lines * update jit module to take the builder file name and extract module name from the builder, update missing entries in the table in readme for adding new module in jit * refine the description about module to jit * add description about jit * add description about jit * add code to create a builder based on user inputs * change the example from fused_dense to swiglu * allow user to skip sources list * change description of cxx and nvcc flags, add description of methods and fields in the initial builder code created by script
jithunnair-amd
pushed a commit
to ROCm/pytorch
that referenced
this pull request
Mar 5, 2026
Commit Messages: - Create custom python operators for MixedFusedLayerNorm and MixedFusedRMSNorm. (ROCm/apex#304) - Add new apex module to jit load system (ROCm/apex#294) - add details of fused_conv_bias_relu in table of modules and fix error of maximum depth reached (ROCm/apex#297) - Port fused_conv_bias_relu to ROCm (ROCm/apex#295) Co-authored-by: sriram <sriram.kumar@silo.ai>
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
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.
Motivation
Integrate a new apex extension to jit loader system. Users create a apex module with C++/CUDA source code. The user would like to easily add it to JIT. The script helps to create initial builder and loader modules to help the module integrate into JIT system.
Technical Details
Run the scripts/jit_module.py script.
The script scripts/jit_module.py creates a loader module in compatibility folder and builder module in op_builder.
More details in the readme file.
Test Plan
Test Result
Apex builds and the unit test passes.
Submission Checklist