Enhance GroupedLinear with integrating AITER triton kernels #413

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

sudhu2k merged 25 commits into dev from sudhu/aiter_grouped_gemm_integration

Feb 4, 2026

ci/pytorch.sh

-Original file line number
+Diff line change
@@ Expand Up / @@ -70,6 +70,7 @@ run_test_config(){ @@
         run_default_fa 1 attention/test_kv_cache.py
         run_default_fa 1 triton_kernels/test_cast.py
         run_default_fa 1 triton_kernels/test_cast_mxfp8.py
+        run_default_fa 1 triton_kernels/test_grouped_gemm.py
         run_default_fa 1 triton_kernels/test_norm_common.py
         run_default_fa 1 triton_kernels/test_norms.py
         NVTE_TEST_TRITON_AUTOTUNE=1 run_default_fa_lbl "autotune" 3 triton_kernels/test_norms.py
@@ Expand Down @@

setup.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -196,7 +196,10 @@ def setup_requirements() -> Tuple[List[str], List[str]]: @@
             install_requires, test_requires = setup_requirements()
             ext_modules = [setup_common_extension()]
             cmdclass = {"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist}
-            package_data = {"": ["VERSION.txt"]}
+            package_data = {
+                "": ["VERSION.txt"],
+                "transformer_engine.pytorch.triton_kernels.gmm": ["configs/*.json"],
+            }
             include_package_data = True
             extras_require = {"test": test_requires}
@@ Expand Down @@

tests/pytorch/test_numerics.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -10,6 +10,7 @@ @@
     import pytest
     import random
+    from triton_kernels.test_common import get_tolerances
     import torch
     import torch.nn as nn
     from torch.nn import Parameter
@@ Expand Down Expand Up / @@ -2019,6 +2020,7 @@ def _test_grouped_linear_accuracy( @@
     @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
     @pytest.mark.parametrize("bias", all_boolean)
     @pytest.mark.parametrize("delay_wgrad_compute", all_boolean)
+    @pytest.mark.parametrize("use_triton", all_boolean)
     def test_grouped_linear_accuracy(
         dtype,
         num_gemms,
@@ Expand All / @@ -2029,10 +2031,14 @@ def test_grouped_linear_accuracy( @@
         fuse_wgrad_accumulation,
         bias,
         delay_wgrad_compute,
+        use_triton,
         parallel_mode=None,
         use_cutlass=False,
     ):
         fp8 = recipe is not None
+        if not IS_HIP_EXTENSION and use_triton:
+            pytest.skip("Triton grouped gemm is only supported on HIP.")
         if IS_HIP_EXTENSION:
             if dtype not in (torch.float32,) and fuse_wgrad_accumulation and not fp8:
@@ Expand All / @@ -2044,6 +2050,9 @@ def test_grouped_linear_accuracy( @@
         if config.max_seqlen_q % 16 != 0 and fp8:
             pytest.skip("FP8 requires sequence length to be divisible by 16.")
+        if use_triton:
+            os.environ["NVTE_USE_GROUPED_GEMM_TRITON"] = "1"
         with fp8_model_init(enabled=fp8 and fp8_model_params, recipe=recipe):
             grouped_linear = GroupedLinear(
                 num_gemms,
@@ Expand Down Expand Up / @@ -2106,12 +2115,19 @@ def test_grouped_linear_accuracy( @@
             delay_wgrad_compute,
         )
+        if use_triton:
+            os.environ.pop("NVTE_USE_GROUPED_GEMM_TRITON", None)
+        atol, rtol = 0, 0
+        if use_cutlass:
+            atol, rtol = 1e-3, 1e-3
+        if use_triton:
+            atol, rtol = get_tolerances(dtype)
+            if dtype == torch.float32:
+                atol = 2.6e-6
+                rtol = 5e-2
         for o, o_ref in zip(outputs, outputs_ref):
-            if use_cutlass:
-                torch.testing.assert_close(o, o_ref, rtol=1e-3, atol=1e-3)
-            else:
-                # cuBLAS implementation should be bit-wise match
-                torch.testing.assert_close(o, o_ref, rtol=0, atol=0)
+            torch.testing.assert_close(o, o_ref, rtol=rtol, atol=atol)
     @pytest.mark.skipif(
@@ Expand Down Expand Up / @@ -2143,6 +2159,7 @@ def test_grouped_linear_accuracy_cutlass( @@
             fuse_wgrad_accumulation,
             False,
             delay_wgrad_compute,
+            False,
             None,
             use_cutlass=True,
         )
@@ Expand Down Expand Up / @@ -2260,6 +2277,7 @@ def test_grouped_linear_accuracy_single_gemm(recipe): @@
             fuse_wgrad_accumulation=True,
             bias=True,
             delay_wgrad_compute=False,
+            use_triton=False,
         )
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enhance GroupedLinear with integrating AITER triton kernels #413

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!