Enhance GroupedLinear with integrating AITER triton kernels #413
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
added 2 commits
January 15, 2026 16:23
…for JSON configs - Added support for using Triton kernels in GroupedLinear, allowing for optimized performance based on environment variables. - Updated the setup.py to include JSON configuration files for Triton kernels in the package data. - Added a new test case for grouped GEMM functionality in the CI pipeline. - Refactored the handling of input tensors and gradients to accommodate the new Triton kernel logic.
sudhu2k
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
added 4 commits
January 15, 2026 19:18
Collaborator
|
Update copyright date of modified files |
ipanfilo
reviewed
Jan 17, 2026
| run_default_fa 1 triton_kernels/test_cast_mxfp8.py | ||
| run_default_fa 1 triton_kernels/test_norm_common.py | ||
| run_default_fa 1 triton_kernels/test_norms.py | ||
| run_default_fa 1 triton_kernels/test_grouped_gemm.py |
Collaborator
There was a problem hiding this comment.
Please move it two lines higher, alphabetical sort helps to find tests
| delay_wgrad_compute, | ||
| parallel_mode=None, | ||
| ): | ||
| os.environ["NVTE_USE_GROUPED_GEMM_TRITON"] = "1" |
Collaborator
There was a problem hiding this comment.
This env won't be cleared if the test is skipped of failed
| else: | ||
| inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) | ||
|
|
||
| if not use_grouped_gemm_triton: |
| group_sizes_list=kwargs.get("m_splits_list", []), | ||
| ) | ||
|
|
||
| grad_biases = [None] * len(m_splits) if bias is None else bias |
Collaborator
There was a problem hiding this comment.
m_splits.shape[0] or len(m_splits_list)?
| package_data = {"": ["VERSION.txt"]} | ||
| package_data = { | ||
| "": ["VERSION.txt"], | ||
| "transformer_engine.pytorch.triton_kernels.gmm": ["configs/*.json"], |
Collaborator
There was a problem hiding this comment.
They should be part of pytorch extension installation not TE core
| _ = general_grouped_gemm( | ||
| general_grouped_gemm_func = general_grouped_gemm_triton if use_grouped_gemm_triton else general_grouped_gemm | ||
| # Prepare m_splits for each backend | ||
| m_splits_for_kernel = m_splits |
Collaborator
There was a problem hiding this comment.
It may be more straightforward to keep m_splits as-is and add mandatory parameter m_splits_tensor or m_splits_for_kernel to general_grouped_gemm_triton(), instead of swapping them here.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Description
This PR enhances the GroupedLinear module with Triton kernel support for grouped GEMM operations, providing optimized performances. The implementation includes a complete Triton-based grouped matrix multiplication (GMM) backend that can be enabled via environment variables, along with pre-tuned configurations for optimal performance.
Benchmark results:
https://github.com/ROCm/frameworks-internal/issues/13792#issuecomment-3739558113
https://github.com/ROCm/frameworks-internal/issues/13792#issuecomment-3746418683
Type of change
Changes
NVTE_USE_GROUPED_GEMM_TRITON)m_splits_tensorparameter to keep tensor data on GPU and avoid redundant CPU-GPU data transfers for improved performancetransformer_engine/pytorch/triton_kernels/gmm/from AITER including:gmm_common.py: Common utilities and helper functionsgmm_kernels.py: Core Triton kernel implementations for grouped GEMM operationsgmm_wrapper.py: High-level wrapper functions from AITERpid_preprocessing.py: Process ID preprocessing for efficient kernel schedulinggfx942-GMM.json: pre-tuned configs for gfx942 archgfx950-GMM.json: pre-tuned configs for gfx950 archgrouped_linear.pyto support Triton kernel path with proper tensor handling.triton_kernels/common.pytests/pytorch/triton_kernels/test_grouped_gemm.py(516 lines)ci/pytorch.shto include grouped GEMM tests in the CI pipelineChecklist: