Grouped MXFP8 GEMMs with HipKittens

README.rst

@@ -208,6 +208,22 @@ If you want to check that only previously tuned algorithms are used by your appl
   diff algo_tune.csv algo_tune_check.csv
 
 
+Single-Launch Grouped GEMM Backends on ROCm
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+For Mixture of Experts workloads, ROCm TE supports single-launch grouped GEMMs. Two backends are available:
+
+* **HipKittens** -- MXFP8 grouped GEMM only, requires expert dimensions to be 256-aligned. Default when the grouped GEMM path is enabled on gfx950.
+* **CK** (Composable Kernel) -- Grouped GEMM used as a fallback when HipKittens constraints are not met, or when explicitly selected.
+
+These backends are not enabled by default. The following environment variables control backend selection:
+
+* ``NVTE_USE_CUTLASS_GROUPED_GEMM=1`` -- Enable the single-launch grouped GEMM path. On gfx950, defaults to HipKittens with CK fallback.
+* ``NVTE_USE_HIPKITTENS_GROUPED_GEMM=1`` -- Explicitly enable HipKittens grouped GEMM. Takes priority over CK.
+* ``NVTE_USE_CK_GROUPED_GEMM=1`` -- Explicitly select CK, bypassing HipKittens.
+
+When none are set, TE uses multi-stream dispatch (one hipBLASLt GEMM per expert).
+
+
 Fused Attention Backends on ROCm
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Currently ROCm TE supports two backends, AOTriton and CK, for fused attention.

tests/pytorch/test_numerics.py

@@ -2249,6 +2249,7 @@ def test_grouped_linear_accuracy(
 )
 @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean)
 @pytest.mark.parametrize("delay_wgrad_compute", all_boolean)
+@pytest.mark.parametrize("grouped_gemm_backend", ["hipkittens", "ck"] if IS_HIP_EXTENSION else ["cutlass"], ids=str)
 def test_grouped_linear_accuracy_cutlass(
     dtype,
     num_gemms,
@@ -2258,8 +2259,21 @@ def test_grouped_linear_accuracy_cutlass(
     fp8_model_params,
     fuse_wgrad_accumulation,
     delay_wgrad_compute,
+    grouped_gemm_backend,
 ):
+    if (IS_HIP_EXTENSION and grouped_gemm_backend == "ck"
+            and recipe is not None and recipe.mxfp8()
+            and get_device_compute_capability() != (12, 5)):
+        pytest.skip("CK MXFP8 grouped GEMM only supported on gfx1250.")
+
     os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1"
+    if IS_HIP_EXTENSION:
+        os.environ.pop("NVTE_USE_HIPKITTENS_GROUPED_GEMM", None)
+        os.environ.pop("NVTE_USE_CK_GROUPED_GEMM", None)
+        if grouped_gemm_backend == "hipkittens":
+            os.environ["NVTE_USE_HIPKITTENS_GROUPED_GEMM"] = "1"
+        elif grouped_gemm_backend == "ck":
+            os.environ["NVTE_USE_CK_GROUPED_GEMM"] = "1"
     test_grouped_linear_accuracy(
         dtype,
         num_gemms,
@@ -2275,6 +2289,9 @@ def test_grouped_linear_accuracy_cutlass(
         use_cutlass=True,
     )
     os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None)
+    if IS_HIP_EXTENSION:
+        os.environ.pop("NVTE_USE_HIPKITTENS_GROUPED_GEMM", None)
+        os.environ.pop("NVTE_USE_CK_GROUPED_GEMM", None)
 
 
 @pytest.mark.parametrize("dtype", param_types, ids=str)

transformer_engine/common/gemm/cublaslt_gemm.cu

@@ -33,6 +33,9 @@
 #include "./cutlass_grouped_gemm.cuh"
 #else
 #include "ck_grouped_gemm/ck_grouped_gemm.h"
+#ifdef USE_HIPKITTENS_GEMM
+#include "kittens/mxfp8_gemm.h"
+#endif
 #endif
 
 #ifndef __HIP_PLATFORM_AMD__
@@ -1097,6 +1100,17 @@ void nvte_cublas_handle_init() { auto _ = cublasHandleManager::Instance().GetHan
 }  //  namespace transformer_engine
 #endif // __HIP_PLATFORM_AMD__
 
+#ifdef USE_HIPKITTENS_GEMM
+namespace transformer_engine {
+bool try_kittens_grouped_mxfp8_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,
+    int num_gemms, bool transa, bool transb, NVTETensor *workspace,
+    bool accumulate, cudaStream_t stream);
+bool try_kittens_grouped_mxfp8_wgrad(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,
+    int num_gemms, bool transa, bool transb, NVTETensor *workspace,
+    bool accumulate, cudaStream_t stream);
+}
+#endif
+
 void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,
                             const NVTETensor *bias, NVTETensor *pre_gelu_out, const int num_gemms,
                             bool transa, bool transb, bool grad, NVTETensor *workspace,
@@ -1111,7 +1125,9 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor
   const int current_device = transformer_engine::cuda::current_device();
   const bool is_hopper = (transformer_engine::cuda::sm_arch(current_device) == 90);
 #endif
-  const bool use_cutlass = transformer_engine::getenv<bool>("NVTE_USE_CUTLASS_GROUPED_GEMM", false);
+  const bool use_cutlass = transformer_engine::getenv<bool>("NVTE_USE_CUTLASS_GROUPED_GEMM", false)
+                        || transformer_engine::getenv<bool>("NVTE_USE_HIPKITTENS_GROUPED_GEMM", false)
+                        || transformer_engine::getenv<bool>("NVTE_USE_CK_GROUPED_GEMM", false);
   const bool warn_fallback =
       transformer_engine::getenv<bool>("NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK", false);
 
@@ -1201,6 +1217,16 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor
 
     bool handled_by_ck = false;
     if (transformer_engine::is_mxfp8_scaling(inputA->scaling_mode)) {
+#ifdef USE_HIPKITTENS_GEMM
+      if (transformer_engine::try_kittens_grouped_mxfp8_gemm(A, B, D, num_gemms, transa, transb,
+                                         workspace, accumulate, stream)) {
+        return;
+      }
+      if (transformer_engine::try_kittens_grouped_mxfp8_wgrad(A, B, D, num_gemms, transa, transb,
+                                         workspace, accumulate, stream)) {
+        return;
+      }
+#endif
       handled_by_ck = ck_tile_mx_grouped_gemm(
           A, B, D, num_gemms, transa, transb, workspace, accumulate, stream);
     } else {