From a66e7978a8020141b64bf19c4d562f7435f74ed2 Mon Sep 17 00:00:00 2001 From: Alex Magro Date: Fri, 26 Jun 2026 18:30:43 +0000 Subject: [PATCH 1/3] hipblaslt Fallback hotfix --- transformer_engine/common/gemm/rocm_gemm.cu | 27 +++++++++++---------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index 625fdef69..30bd8383f 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -1985,7 +1985,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #ifdef USE_HIPKITTENS_GEMM - bool use_hipkittens = false; + bool use_hipkittens = false, hipkittens_gemm_complete = false; if (is_mxfp8) { bool is_gfx950 = (cuda::sm_arch() == 95); bool force_hipblaslt = false; @@ -2001,18 +2001,19 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, if (use_hipkittens) { auto param = CanonicalizeGemmInput(*inputA, transa, *inputB, transb, m, n, k); - kittens_mxfp8_gemm(param.A, param.B, outputD->data.dptr, - param.A_scale_inv, param.B_scale_inv, - m, n, k, is_transa, is_transb, - static_cast(param.Atype), - static_cast(param.Btype), - inputBias->data.dptr, - static_cast(inputBias->data.dtype), - outputPreGelu->data.dptr, - static_cast(outputD->data.dtype), - static_cast(outputPreGelu->data.dtype), - workspace, workspaceSize, s); - } else { + hipkittens_gemm_complete = kittens_mxfp8_gemm(param.A, param.B, outputD->data.dptr, + param.A_scale_inv, param.B_scale_inv, + m, n, k, is_transa, is_transb, + static_cast(param.Atype), + static_cast(param.Btype), + inputBias->data.dptr, + static_cast(inputBias->data.dtype), + outputPreGelu->data.dptr, + static_cast(outputD->data.dtype), + static_cast(outputPreGelu->data.dtype), + workspace, workspaceSize, s); + } + if (!use_hipkittens || !hipkittens_gemm_complete) { #endif // FIXME(https://amd-hub.atlassian.net/browse/ROCM-26110): Remove this workaround once hipBLASLt supports NN/NT // layouts for MXFP8 on gfx1250. From db368903c9da8c3698e70656a4021e3a16f1a352 Mon Sep 17 00:00:00 2001 From: Alex Magro Date: Mon, 29 Jun 2026 13:52:14 +0000 Subject: [PATCH 2/3] Remove hipkittens_gemm_complete var --- transformer_engine/common/gemm/rocm_gemm.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index fc2e78ae5..ab3341d50 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -1985,7 +1985,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #ifdef USE_HIPKITTENS_GEMM - bool use_hipkittens = false, hipkittens_gemm_complete = false; + bool use_hipkittens = false; if (is_mxfp8) { bool is_gfx950 = (cuda::sm_arch() == 95); bool force_hipblaslt = false; From c84068c00f10b260b8279fcc1b382bf4cc477b16 Mon Sep 17 00:00:00 2001 From: Alex Magro Date: Mon, 29 Jun 2026 18:37:18 +0000 Subject: [PATCH 3/3] Fix JAX workspace and bias check --- transformer_engine/common/gemm/rocm_gemm.cu | 4 ++++ transformer_engine/jax/cpp_extensions/gemm.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index ab3341d50..7e11c8054 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -2012,6 +2012,10 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, workspace, workspaceSize, gemm_stream); } if (!use_hipkittens) { + if (is_mxfp8) { + NVTE_CHECK(inputBias->data.dptr == nullptr, + "hipBLASLt MXFP8 GEMM does not support bias"); + } #endif // FIXME(https://amd-hub.atlassian.net/browse/ROCM-26110): Remove this workaround once hipBLASLt supports NN/NT // layouts for MXFP8 on gfx1250. diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 7864c0438..f8177952a 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -653,7 +653,7 @@ def _dims_are_consecutive(dims): n = reduce(operator.mul, rhs_non_contracting_shape) k = lhs_contracting_size layout = ("T" if lhs_is_transposed else "N") + ("T" if rhs_is_transposed else "N") - workspace_size = _hipkittens_workspace_bytes(m, n, k, layout) + workspace_size = max(_hipkittens_workspace_bytes(m, n, k, layout), get_cublas_workspace_size_bytes()) else: # Declare cuBLAS workspace workspace_size = get_cublas_workspace_size_bytes()