From 4092e1e65f420b9fade989e7a690c9dd85b2b6a7 Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Fri, 19 Jun 2026 22:01:48 +0000 Subject: [PATCH] Relax MXFP8 GEMM K constraint from multiple-of-128 to multiple-of-64 The ROCm MXFP8 GEMM guard required K % 128 == 0, but hipBLASLt has MXFP8 (VEC32_UE8M0, scaleA/B=3) solutions for every K % 64 == 0. Verified on MI355X (gfx950) via hipblaslt-bench: K=64/128/192/256/384 succeed; K=32/96/160/224/288 return "no solution". The 128 bound was over-strict and rejected supported shapes such as K=192. M/N multiple-of-16 requirements are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- transformer_engine/common/gemm/rocm_gemm.cu | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index 20d05e3db..028482002 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -1737,12 +1737,16 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK((is_transb ? B0 : B1) == k, "GEMM inputs have incompatible dimensions (A is ", A0, "x", A1, ", B is ", B0, "x", B1, ")"); - // Check that K is compatible with the MXFP8 scale layout, and M/N are multiples of 16 + // Check that K is compatible with the MXFP8 scale layout, and M/N are multiples of 16. + // On non-gfx1250 archs K only needs to be a multiple of 64: hipBLASLt has MXFP8 + // solutions for every K % 64 == 0 (verified on gfx950 via hipblaslt-bench, scaleA/B=3 + // VEC32_UE8M0: K=64/128/192/256/384 succeed; K=32/96/160/224/288 have no solution). + // The prior 128 bound was over-strict and rejected supported shapes (e.g. K=192). if (inputA->scaling_mode == NVTE_MXFP8_1D_SCALING || inputB->scaling_mode == NVTE_MXFP8_1D_SCALING) { const bool is_gfx1250 = cuda::sm_arch() == 125; - // TODO: Also use 32 for gfx950 once hipBLASLt (and TE) support MXFP8 GEMM with + // TODO: Also use 32 for gfx950 once hipBLASLt (and TE) support MXFP8 GEMM with // swizzled scales on that architecture. - const int required_k_multiple = is_gfx1250 ? 32 : 128; + const int required_k_multiple = is_gfx1250 ? 32 : 64; NVTE_CHECK(inputBias->data.dptr == nullptr, "MXFP8 GEMM does not yet support bias."); NVTE_CHECK((k % required_k_multiple) == 0, "GEMM K dimension must be multiple of ", required_k_multiple,