diff --git a/transformer_engine/common/common.cu b/transformer_engine/common/common.cu index 6f2bc7d124..ff1bdd93f2 100644 --- a/transformer_engine/common/common.cu +++ b/transformer_engine/common/common.cu @@ -28,7 +28,6 @@ __global__ void __launch_bounds__(1) } // namespace -#ifndef __HIP_PLATFORM_AMD__ cudaDataType_t get_cuda_dtype(const transformer_engine::DType t) { using namespace transformer_engine; switch (t) { @@ -50,7 +49,6 @@ cudaDataType_t get_cuda_dtype(const transformer_engine::DType t) { NVTE_ERROR("Invalid type"); } } -#endif void update_tensor_scale_inv(Tensor *t, cudaStream_t stream) { if (is_fp8_dtype(t->data.dtype) && is_tensor_scaling(t->scaling_mode)) { diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 576953739f..05bc284d6e 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -25,16 +25,19 @@ #include "../util/vectorized_pointwise.h" #include "./config.h" -#ifndef __HIP_PLATFORM_AMD__ - -namespace { - -inline void CreateCublasHandle(cublasLtHandle_t *handle) { - NVTE_CHECK_CUBLAS(cublasLtCreate(handle)); -} - -} // namespace +#ifdef __HIP_PLATFORM_AMD__ +#include +#include +// Grouped GEMM via hipblaslt_ext::GroupedGemm requires hipBLASLt 1.3.0+. +#if (HIPBLASLT_VERSION_MAJOR > 1) || \ + (HIPBLASLT_VERSION_MAJOR == 1 && HIPBLASLT_VERSION_MINOR >= 3) +#define TE_HIPBLASLT_GROUPED_GEMM_SUPPORTED 1 +#else +#define TE_HIPBLASLT_GROUPED_GEMM_SUPPORTED 0 +#endif +#endif +#ifndef __HIP_PLATFORM_AMD__ // MXFP8 support for grouped GEMM requires cuBLAS 13.3+ #define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300 @@ -49,8 +52,24 @@ inline void CreateCublasHandle(cublasLtHandle_t *handle) { // BF16 support for grouped GEMM requires cuBLAS 13.3+ #define CUBLAS_GROUPED_GEMM_VERSION 130300 +#endif // !__HIP_PLATFORM_AMD__ + +#if (defined(__HIP_PLATFORM_AMD__) && TE_HIPBLASLT_GROUPED_GEMM_SUPPORTED) || \ + (defined(CUBLAS_VERSION) && CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION) + +namespace { -#if CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION +#ifndef __HIP_PLATFORM_AMD__ +inline void CreateCublasHandle(cublasLtHandle_t *handle) { + NVTE_CHECK_CUBLAS(cublasLtCreate(handle)); +} +#else +inline void CreateHipblasLtHandle(hipblasLtHandle_t *handle) { + NVTE_CHECK_HIPBLASLT(hipblasLtCreate(handle)); +} +#endif + +} // namespace namespace { @@ -169,6 +188,10 @@ struct GroupedGemmSetupWorkspace { int *d_cols = nullptr; // N (last dim) - also used for C // NVFP4: per-group computed alpha values (alpha * amax_A * amax_B * factor_inv) float *nvfp4_computed_alpha = nullptr; +#ifdef __HIP_PLATFORM_AMD__ + // Device buffer for hipBLASLt grouped GEMM UserArguments (one entry per group). + hipblaslt_ext::UserArguments *user_args = nullptr; +#endif // End-of-layout offset in bytes (unaligned). required_setup_size rounds this up. size_t total_bytes = 0; @@ -219,6 +242,11 @@ struct GroupedGemmSetupWorkspace { place(ws.d_rows, int_size); place(ws.d_cols, int_size); place(ws.nvfp4_computed_alpha, float_size); +#ifdef __HIP_PLATFORM_AMD__ + // hipBLASLt requires the UserArguments buffer to be 16-byte aligned. + align_ptr(); + place(ws.user_args, num_tensors * sizeof(hipblaslt_ext::UserArguments)); +#endif ws.total_bytes = offset; return ws; @@ -232,7 +260,13 @@ struct GroupedGemmSetupWorkspace { } }; -inline bool grouped_gemm_supports_per_group_alpha_beta(int sm) { return sm >= 100 && sm <= 110; } +inline bool grouped_gemm_supports_per_group_alpha_beta(int sm) { +#ifdef __HIP_PLATFORM_AMD__ + return sm == 94 || sm == 95 || sm == 125; +#else + return sm >= 100 && sm <= 110; +#endif +} inline size_t validate_grouped_gemm_inputs( size_t num_tensors, std::initializer_list inputs, @@ -337,6 +371,10 @@ inline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) { inline void check_grouped_gemm_requirements(const char *api_name) { const int current_device = transformer_engine::cuda::current_device(); const int sm = transformer_engine::cuda::sm_arch(current_device); +#ifdef __HIP_PLATFORM_AMD__ + NVTE_CHECK(sm == 94 || sm == 95 || sm == 125, api_name, + " requires gfx942, gfx950, or gfx1250, but device compute capability is ", sm, "."); +#else const int cublas_ver = transformer_engine::cuda::cublas_version(); #if CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_HOPPER_VERSION NVTE_CHECK(sm >= 90 && sm <= 110, api_name, @@ -353,6 +391,7 @@ inline void check_grouped_gemm_requirements(const char *api_name) { NVTE_CHECK(cublas_ver >= CUBLAS_GROUPED_GEMM_VERSION, api_name, " requires cuBLAS 13.3+, but run-time cuBLAS version is ", cublas_ver); #endif +#endif } inline transformer_engine::GroupedMatmulConfig parse_grouped_gemm_config( @@ -385,13 +424,16 @@ inline void validate_nvfp4_grouped_gemm_support(const GroupedOperandSelection &A const bool nvfp4 = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) || transformer_engine::is_nvfp_scaling(B_sel.scaling_mode); if (!nvfp4) return; - +#ifdef __HIP_PLATFORM_AMD__ + NVTE_CHECK(false, "Grouped GEMM: NVFP4 is not supported on ROCm."); +#else NVTE_CHECK(transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) && transformer_engine::is_nvfp_scaling(B_sel.scaling_mode), "Grouped GEMM: A and B must both use NVFP4 scaling or both not."); NVTE_CHECK(use_per_group_alpha_beta, "Grouped GEMM: NVFP4 requires per-group alpha/beta support because each group " "has its own amax-derived global scale."); +#endif } // FP8 block scaling grouped GEMM is only supported on Hopper (SM90). @@ -431,10 +473,14 @@ inline void validate_grouped_gemm_scaling_modes(NVTEScalingMode a_mode, NVTEScal ": incompatible A/B scaling modes."); if (transformer_engine::is_fp8_block_scaling(a_mode) || transformer_engine::is_fp8_block_scaling(b_mode)) { +#ifdef __HIP_PLATFORM_AMD__ + NVTE_CHECK(false, api_name, ": FP8 block scaling grouped GEMM is not supported on ROCm."); +#else NVTE_CHECK(sm >= 90 && sm < 100, api_name, ": FP8 block scaling grouped GEMM is only supported on Hopper (SM90-SM99), " "not SM", sm, "."); +#endif } } @@ -814,6 +860,8 @@ inline void *validate_and_get_workspace_ptr(transformer_engine::Tensor *ws, size return ws->data.dptr; } +#ifndef __HIP_PLATFORM_AMD__ + inline void init_matrix_layouts( cublasLtMatrixLayoutOpaque_t &descA, cublasLtMatrixLayoutOpaque_t &descB, cublasLtMatrixLayoutOpaque_t &descC, cublasLtMatrixLayoutOpaque_t &descD, @@ -1003,6 +1051,8 @@ inline cublasLtMatmulAlgo_t select_grouped_gemm_algo(cublasLtHandle_t handle, return heuristicResult.algo; } +#endif // !__HIP_PLATFORM_AMD__ + struct GroupedGemmWorkspace { GroupedGemmSetupWorkspace setup_workspace; void *cublas_workspace_ptr = nullptr; @@ -1027,12 +1077,231 @@ inline GroupedGemmWorkspace setup_grouped_gemm_workspace(transformer_engine::Ten return {std::move(setup_workspace), cublas_workspace_ptr, num_tensors}; } +#ifdef __HIP_PLATFORM_AMD__ + +__device__ __forceinline__ void store_user_arg_scalar(int8_t *dst, float value) { + *reinterpret_cast(dst) = value; +} + +// Fill hipBLASLt UserArguments on device from the setup-kernel workspace. Skips empty groups +// (d_rows/d_cols <= 0) using the same compaction order as the host setProblem path. +__global__ void populate_hipblaslt_user_args_kernel(hipblaslt_ext::UserArguments *user_args, + void *const *A_ptrs, void *const *B_ptrs, + void *const *D_ptrs, const int *a_rows, + const int *a_cols, const int *b_rows, + const int *b_cols, const int *d_rows, + const int *d_cols, float *const *alpha_ptrs, + float *const *beta_ptrs, size_t num_tensors, + bool trans_A, bool trans_B) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + size_t out_idx = 0; + for (size_t i = 0; i < num_tensors; ++i) { + const int m = d_rows[i]; + const int n = d_cols[i]; + if (m <= 0 || n <= 0) { + continue; + } + + hipblaslt_ext::UserArguments ua{}; + ua.m = static_cast(m); + ua.n = static_cast(n); + ua.batch = 1; + const int kA = trans_A ? a_rows[i] : a_cols[i]; + ua.k = static_cast(kA); + + ua.a = A_ptrs[i]; + ua.b = B_ptrs[i]; + ua.c = D_ptrs[i]; + ua.d = D_ptrs[i]; + ua.strideA1 = static_cast(a_rows[i]); + ua.strideB1 = static_cast(b_rows[i]); + ua.strideC1 = static_cast(m); + ua.strideD1 = static_cast(m); + + store_user_arg_scalar(ua.alpha, *alpha_ptrs[i]); + store_user_arg_scalar(ua.beta, *beta_ptrs[i]); + + user_args[out_idx++] = ua; + } +} + +inline void rocm_execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspace, + const GroupedOperandSelection &A_sel, + const GroupedOperandSelection &B_sel, + transformer_engine::DType d_dtype, size_t num_tensors, + const GroupedGemmConfig &config, + transformer_engine::Tensor *wspace_cublas, + cudaStream_t stream) { + using namespace transformer_engine; + using namespace hipblaslt_ext; + NVTE_CHECK(num_tensors <= static_cast(kMaxGroups), + "ROCm grouped GEMM supports up to ", kMaxGroups, " groups."); + + NVTE_CHECK(A_sel.scaling_mode == NVTE_DELAYED_TENSOR_SCALING && + B_sel.scaling_mode == NVTE_DELAYED_TENSOR_SCALING, + "ROCm V2 grouped GEMM via hipBLASLt currently supports unscaled FP16 only."); + NVTE_CHECK(is_fp16_dtype(A_sel.dtype) && is_fp16_dtype(B_sel.dtype) && is_fp16_dtype(d_dtype), + "ROCm V2 grouped GEMM via hipBLASLt supports FP16 only, got A=", + to_string(A_sel.dtype), ", B=", to_string(B_sel.dtype), ", D=", to_string(d_dtype), + "."); + + const size_t n = num_tensors; + std::vector h_A(n), h_B(n), h_D(n); + std::vector h_alpha_ptrs(n), h_beta_ptrs(n); + std::vector h_a_rows(n), h_a_cols(n), h_b_rows(n), h_b_cols(n), h_d_rows(n), h_d_cols(n); + + auto dtoh = [&](const void *src, void *dst, size_t bytes) { + NVTE_CHECK_CUDA(cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToHost, stream)); + }; + dtoh(setup_workspace.A_ptrs, h_A.data(), n * sizeof(void *)); + dtoh(setup_workspace.B_ptrs, h_B.data(), n * sizeof(void *)); + dtoh(setup_workspace.D_ptrs, h_D.data(), n * sizeof(void *)); + dtoh(setup_workspace.alpha_ptrs, h_alpha_ptrs.data(), n * sizeof(float *)); + dtoh(setup_workspace.beta_ptrs, h_beta_ptrs.data(), n * sizeof(float *)); + dtoh(setup_workspace.a_rows, h_a_rows.data(), n * sizeof(int)); + dtoh(setup_workspace.a_cols, h_a_cols.data(), n * sizeof(int)); + dtoh(setup_workspace.b_rows, h_b_rows.data(), n * sizeof(int)); + dtoh(setup_workspace.b_cols, h_b_cols.data(), n * sizeof(int)); + dtoh(setup_workspace.d_rows, h_d_rows.data(), n * sizeof(int)); + dtoh(setup_workspace.d_cols, h_d_cols.data(), n * sizeof(int)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector active; + active.reserve(n); + for (size_t i = 0; i < n; ++i) { + if (h_d_rows[i] > 0 && h_d_cols[i] > 0) { + active.push_back(i); + } + } + if (active.empty()) { + return; + } + + const hipDataType a_type = get_cuda_dtype(A_sel.dtype); + const hipDataType b_type = get_cuda_dtype(B_sel.dtype); + const hipDataType d_type = get_cuda_dtype(d_dtype); + const hipblasOperation_t op_A = A_sel.trans ? HIPBLAS_OP_T : HIPBLAS_OP_N; + const hipblasOperation_t op_B = B_sel.trans ? HIPBLAS_OP_T : HIPBLAS_OP_N; + + using HipblasHandleManager = + transformer_engine::detail::HandleManager; + hipblasLtHandle_t handle = HipblasHandleManager::Instance().GetHandle(); + + std::vector heuristic_results; + NVTE_CHECK_HIPBLASLT(hipblaslt_ext::getAllAlgos( + handle, GemmType::HIPBLASLT_GROUPED_GEMM, op_A, op_B, a_type, b_type, d_type, d_type, + HIPBLAS_COMPUTE_32F, heuristic_results)); + NVTE_CHECK(!heuristic_results.empty(), + "No hipBLASLt grouped GEMM algorithms were found for the requested problem type."); + + GroupedGemm grouped_gemm(handle, op_A, op_B, a_type, b_type, d_type, d_type, HIPBLAS_COMPUTE_32F); + size_t max_workspace_bytes = kGroupedGemmCublasWorkspaceSize; + if (wspace_cublas != nullptr) { + max_workspace_bytes = + get_buffer_size_bytes(wspace_cublas->data.numel(), wspace_cublas->data.dtype); + } + grouped_gemm.setMaxWorkspaceBytes(max_workspace_bytes); + + const size_t g = active.size(); + std::vector m(g), n_dim(g), k(g), batch_count(g, 1); + std::vector lda(g), ldb(g), ldc(g), ldd(g); + std::vector strideA(g, 0), strideB(g, 0), strideC(g, 0), strideD(g, 0); + std::vector h_alpha_vals(g), h_beta_vals(g); + std::vector epilogue(g); + std::vector inputs(g); + GemmProblemType problem_type(op_A, op_B, a_type, b_type, d_type, d_type, HIPBLAS_COMPUTE_32F); + + for (size_t j = 0; j < g; ++j) { + const size_t i = active[j]; + m[j] = h_d_rows[i]; + n_dim[j] = h_d_cols[i]; + const int64_t kA = A_sel.trans ? h_a_rows[i] : h_a_cols[i]; + const int64_t kB = B_sel.trans ? h_b_cols[i] : h_b_rows[i]; + NVTE_CHECK(kA == kB, "Grouped GEMM K mismatch for group ", i, ": ", kA, " vs ", kB); + k[j] = kA; + lda[j] = h_a_rows[i]; + ldb[j] = h_b_rows[i]; + ldc[j] = h_d_rows[i]; + ldd[j] = h_d_rows[i]; + NVTE_CHECK_CUDA(cudaMemcpy(&h_alpha_vals[j], h_alpha_ptrs[i], sizeof(float), + cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA( + cudaMemcpy(&h_beta_vals[j], h_beta_ptrs[i], sizeof(float), cudaMemcpyDeviceToHost)); + epilogue[j].setMode(HIPBLASLT_EPILOGUE_DEFAULT); + inputs[j].setA(h_A[i]); + inputs[j].setB(h_B[i]); + inputs[j].setC(h_D[i]); + inputs[j].setD(h_D[i]); + inputs[j].setAlpha(&h_alpha_vals[j]); + inputs[j].setBeta(&h_beta_vals[j]); + } + + NVTE_CHECK_HIPBLASLT(grouped_gemm.setProblem(m, n_dim, k, batch_count, lda, ldb, ldc, ldd, strideA, + strideB, strideC, strideD, epilogue, inputs, + problem_type)); + + size_t chosen_algo = heuristic_results.size(); + size_t chosen_workspace = 0; + for (size_t algo_idx = 0; algo_idx < heuristic_results.size(); ++algo_idx) { + hipblasLtMatmulAlgo_t algo = heuristic_results[algo_idx].algo; + size_t workspace_required = 0; + if (grouped_gemm.isAlgoSupported(algo, workspace_required) != HIPBLAS_STATUS_SUCCESS) { + continue; + } + if (workspace_required <= max_workspace_bytes) { + chosen_algo = algo_idx; + chosen_workspace = workspace_required; + break; + } + } + if (chosen_algo >= heuristic_results.size()) { + NVTE_CHECK_HIPBLASLT(grouped_gemm.setProblem(m, n_dim, k, batch_count, epilogue, inputs)); + for (size_t algo_idx = 0; algo_idx < heuristic_results.size(); ++algo_idx) { + hipblasLtMatmulAlgo_t algo = heuristic_results[algo_idx].algo; + size_t workspace_required = 0; + if (grouped_gemm.isAlgoSupported(algo, workspace_required) != HIPBLAS_STATUS_SUCCESS) { + continue; + } + if (workspace_required <= max_workspace_bytes) { + chosen_algo = algo_idx; + chosen_workspace = workspace_required; + break; + } + } + } + NVTE_CHECK(chosen_algo < heuristic_results.size(), + "No suitable hipBLASLt grouped GEMM algorithm was found."); + (void)chosen_workspace; + + NVTE_CHECK(setup_workspace.user_args != nullptr, + "ROCm grouped GEMM setup workspace is missing the UserArguments buffer."); + populate_hipblaslt_user_args_kernel<<<1, 1, 0, stream>>>( + setup_workspace.user_args, setup_workspace.A_ptrs, setup_workspace.B_ptrs, + setup_workspace.D_ptrs, setup_workspace.a_rows, setup_workspace.a_cols, setup_workspace.b_rows, + setup_workspace.b_cols, setup_workspace.d_rows, setup_workspace.d_cols, + setup_workspace.alpha_ptrs, setup_workspace.beta_ptrs, n, A_sel.trans, B_sel.trans); + + void *workspace_ptr = wspace_cublas != nullptr ? wspace_cublas->data.dptr : nullptr; + NVTE_CHECK_HIPBLASLT( + grouped_gemm.initialize(heuristic_results[chosen_algo].algo, workspace_ptr, true, stream)); + NVTE_CHECK_HIPBLASLT(grouped_gemm.run(setup_workspace.user_args, stream)); + (void)config; +} +#endif // __HIP_PLATFORM_AMD__ + inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspace, const GroupedOperandSelection &A_sel, const GroupedOperandSelection &B_sel, transformer_engine::DType d_dtype, size_t num_tensors, - const GroupedGemmConfig &config, void *cublas_workspace_ptr, - cudaStream_t stream) { + const GroupedGemmConfig &config, + transformer_engine::Tensor *wspace_cublas, cudaStream_t stream) { +#ifdef __HIP_PLATFORM_AMD__ + rocm_execute_grouped_gemm(setup_workspace, A_sel, B_sel, d_dtype, num_tensors, config, + wspace_cublas, stream); +#else + void *cublas_workspace_ptr = wspace_cublas != nullptr ? wspace_cublas->data.dptr : nullptr; using cublasHandleManager = transformer_engine::detail::HandleManager; cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); @@ -1069,8 +1338,6 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac cublasLtMatmulAlgo_t algo = select_grouped_gemm_algo( handle, matmulDesc, descA, descB, descC, descD, config.avg_m, config.avg_n, config.avg_k); - // Hopper uses a single scalar alpha/beta for the whole grouped GEMM; - // Blackwell+ uses per-matrix alpha/beta arrays. void *alpha_arg = config.use_per_group_alpha_beta ? static_cast(setup_workspace.alpha_ptrs) : config.alpha_dptr; @@ -1081,6 +1348,7 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac setup_workspace.B_ptrs, &descB, beta_arg, setup_workspace.C_ptrs, &descC, setup_workspace.D_ptrs, &descD, &algo, cublas_workspace_ptr, kGroupedGemmCublasWorkspaceSize, stream)); +#endif } // Device helper: compute the element offset for tensor `idx` given shape metadata. @@ -1234,7 +1502,11 @@ __global__ void grouped_bias_add_kernel(char *__restrict__ d_base, : static_cast(d_meta.uniform_first); } for (int offset = 16; offset > 0; offset >>= 1) { +#ifdef __HIP_PLATFORM_AMD__ + local_sum += __shfl_down_sync(0xffffffffULL, local_sum, offset); +#else local_sum += __shfl_down_sync(0xffffffff, local_sum, offset); +#endif } if (tid == 0) s_valid_rows = local_sum; } @@ -1643,7 +1915,7 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, - gemm_config, workspace.cublas_workspace_ptr, stream); + gemm_config, wspace_cublas, stream); } void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, @@ -1793,7 +2065,7 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num gemm_config.avg_k = config_.avg_k.value_or(transa ? avg_first_dim : avg_last_dim); gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, - gemm_config, workspace.cublas_workspace_ptr, stream); + gemm_config, wspace_cublas, stream); } void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, @@ -1881,7 +2153,7 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, d_dtype, num_tensors, gemm_config, - workspace.cublas_workspace_ptr, stream); + wspace_cublas, stream); } namespace { @@ -1999,15 +2271,16 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream); } -#else // CUBLAS_VERSION < CUBLAS_GROUPED_GEMM_VERSION +#elif defined(__HIP_PLATFORM_AMD__) void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb, const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha, const NVTETensor beta, NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.3+, but compile-time cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); + NVTE_ERROR("nvte_grouped_gemm requires hipBLASLt 1.3.0+, but compile-time hipBLASLt version is ", + HIPBLASLT_VERSION_MAJOR, ".", HIPBLASLT_VERSION_MINOR, ".", + HIPBLASLT_VERSION_PATCH, "."); } void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, @@ -2017,9 +2290,9 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { NVTE_ERROR( - "nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.3+, but compile-time " - "cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); + "nvte_grouped_gemm_with_discrete_inputA requires hipBLASLt 1.3.0+, but compile-time " + "hipBLASLt version is ", + HIPBLASLT_VERSION_MAJOR, ".", HIPBLASLT_VERSION_MINOR, ".", HIPBLASLT_VERSION_PATCH, "."); } void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, @@ -2030,43 +2303,43 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { NVTE_ERROR( - "nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.3+, but compile-time " - "cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); + "nvte_grouped_gemm_with_discrete_out requires hipBLASLt 1.3.0+, but compile-time " + "hipBLASLt version is ", + HIPBLASLT_VERSION_MAJOR, ".", HIPBLASLT_VERSION_MINOR, ".", HIPBLASLT_VERSION_PATCH, "."); } void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); + NVTE_ERROR("nvte_grouped_bias_add requires hipBLASLt 1.3.0+, but compile-time hipBLASLt version is ", + HIPBLASLT_VERSION_MAJOR, ".", HIPBLASLT_VERSION_MINOR, ".", + HIPBLASLT_VERSION_PATCH, "."); } void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, const NVTETensor scale, cudaStream_t stream) { NVTE_ERROR( - "nvte_grouped_scaled_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version " - "is ", - CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); + "nvte_grouped_scaled_bias_add requires hipBLASLt 1.3.0+, but compile-time hipBLASLt " + "version is ", + HIPBLASLT_VERSION_MAJOR, ".", HIPBLASLT_VERSION_MINOR, ".", HIPBLASLT_VERSION_PATCH, "."); } size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) { NVTE_ERROR( - "nvte_get_grouped_gemm_setup_workspace_size requires cuBLAS 13.3+, but compile-time cuBLAS " - "version is ", - CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); + "nvte_get_grouped_gemm_setup_workspace_size requires hipBLASLt 1.3.0+, but compile-time " + "hipBLASLt version is ", + HIPBLASLT_VERSION_MAJOR, ".", HIPBLASLT_VERSION_MINOR, ".", HIPBLASLT_VERSION_PATCH, "."); return 0; } -#endif // CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION - -#else //__HIP_PLATFORM_AMD__ +#else // !__HIP_PLATFORM_AMD__ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb, const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha, const NVTETensor beta, NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_gemm is not supported on ROCm yet"); + NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.3+, but compile-time cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, @@ -2075,7 +2348,10 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num const NVTETensor alpha, const NVTETensor beta, NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_gemm_with_discrete_inputA is not supported on ROCm yet"); + NVTE_ERROR( + "nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.3+, but compile-time " + "cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, @@ -2085,20 +2361,36 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, const NVTETensor alpha, const NVTETensor beta, NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_gemm_with_discrete_out is not supported on ROCm yet"); + NVTE_ERROR( + "nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.3+, but compile-time " + "cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_bias_add is not supported on ROCm yet"); + NVTE_ERROR("nvte_grouped_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); +} + +void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, + const NVTETensor scale, cudaStream_t stream) { + NVTE_ERROR( + "nvte_grouped_scaled_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version " + "is ", + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) { - NVTE_ERROR("nvte_get_grouped_gemm_setup_workspace_size is not supported on ROCm yet"); + NVTE_ERROR( + "nvte_get_grouped_gemm_setup_workspace_size requires cuBLAS 13.3+, but compile-time cuBLAS " + "version is ", + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); return 0; } -#endif // __HIP_PLATFORM_AMD__ +#endif // (HIP && TE_HIPBLASLT_GROUPED_GEMM_SUPPORTED) || CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION + namespace { __global__ void convert_int32_to_int64_kernel(const int32_t *src, int64_t *dst, size_t n) { diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 64f0289682..3e13e42114 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -27,6 +27,7 @@ get_num_compute_streams, JAXX_Collective_Op, get_device_compute_capability, + get_grouped_gemm_setup_workspace_size, ) if not is_hip_extension(): from transformer_engine_jax import ( @@ -34,7 +35,6 @@ initialize_cgemm_communicator, is_collective_gemm_with_cublasmp, get_cgemm_num_max_streams, - get_grouped_gemm_setup_workspace_size, ) from .base import BasePrimitive, register_primitive @@ -90,13 +90,11 @@ # compiled against cuBLAS < 13.2, in which case the cuda-graphable path is unavailable. _v2_grouped_gemm_available_reason = "" try: - if is_hip_extension(): - _v2_grouped_gemm_available = False - else: - get_grouped_gemm_setup_workspace_size(1) - _v2_grouped_gemm_available = True + get_grouped_gemm_setup_workspace_size(1) + _v2_grouped_gemm_available = True except RuntimeError as e: - if "cublas" in str(e).lower(): + err = str(e).lower() + if "cublas" in err or "rocm" in err or "not supported" in err: _v2_grouped_gemm_available = False _v2_grouped_gemm_available_reason = str(e) else: @@ -2138,9 +2136,18 @@ def _is_v2_grouped_gemm_supported( ), ) - # nvte_grouped_gemm (the v2 kernel) requires SM100+ (Blackwell or newer). - # Fall back to the v1 path on SM90 (Hopper) and older architectures. - if get_min_device_compute_capability() < 100: + # nvte_grouped_gemm requires Blackwell (SM100+) on CUDA, or gfx942/950/1250 on ROCm. + if is_hip_extension(): + if get_min_device_compute_capability() not in (94, 95, 125): + return ( + False, + ( + "The TE V2 grouped GEMM on ROCm requires gfx942, gfx950, or gfx1250 but" + f" current min device compute capability is" + f" {get_min_device_compute_capability()}." + ), + ) + elif get_min_device_compute_capability() < 100: return ( False, ( @@ -2152,10 +2159,14 @@ def _is_v2_grouped_gemm_supported( if has_bias: return False, "Grouped GEMM with bias is not supported in the TE V2 grouped GEMM kernel." - if scaling_mode == ScalingMode.NO_SCALING and dtype == jnp.bfloat16: - return True, "" + if scaling_mode == ScalingMode.NO_SCALING: + if is_hip_extension(): + if dtype == jnp.float16: + return True, "" + elif dtype == jnp.bfloat16: + return True, "" - if scaling_mode == ScalingMode.MXFP8_1D_SCALING: + if scaling_mode == ScalingMode.MXFP8_1D_SCALING and not is_hip_extension(): # V2 MXFP8 requires that the total first dimension of both operands (up to # axis_boundary) is divisible by 128, matching the quantize V2 kernel requirement. # Individual group sizes must also be 128-aligned (dynamic constraint). @@ -2215,8 +2226,9 @@ def _is_v2_grouped_gemm_supported( return ( False, ( - "The TE V2 grouped GEMM currently only supports non-quantized BF16 and MXFP8 with 1D" - " block scaling, but NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is enabled and the input" + "The TE V2 grouped GEMM currently only supports non-quantized BF16 (CUDA), FP16" + " (ROCm hipBLASLt), and MXFP8 with 1D block scaling (CUDA only), but" + " NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is enabled and the input" f" parameters do not meet these requirements (scaling_mode= {scaling_mode}," f" dtype={dtype}, has_bias={has_bias}, lhs_shape={lhs_shape}, rhs_shape={rhs_shape}," f" lhs_axis_boundary={lhs_axis_boundary}, rhs_axis_boundary={rhs_axis_boundary})." @@ -2238,8 +2250,9 @@ def is_v2_grouped_gemm_supported( Returns: A tuple of (is_supported: bool, reason: str) where is_supported indicates whether the V2 grouped GEMM can be used, and reason provides an explanation if it is not supported. """ - # Use the V2 path for plain BF16 non-quantized inputs and MXFP8; fall back to - # the legacy nvte_multi_tensor_gemm path for all other cases (tensor-scaled FP8, etc.). + # Use the V2 path for plain BF16 non-quantized inputs and MXFP8(CUDA only), + # ROCm: FP16 unscaled via hipBLASLt only; fall back to the legacy + # nvte_multi_tensor_gemm path for all other cases (tensor-scaled FP8, etc.). # Bias can be supported in a kernel or in pure-JAX in the future. enforce_v2_gmm = _should_enforce_v2_grouped_gemm() diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 725ac26e80..99aef4a8e9 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -178,12 +178,12 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("get_topk_workspace_sizes", &GetTopkWorkspaceSizes); m.def("nvte_get_qkv_format", &nvte_get_qkv_format); m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported); + m.def("get_grouped_gemm_setup_workspace_size", &nvte_get_grouped_gemm_setup_workspace_size); #ifndef USE_ROCM m.def("nvte_built_with_cublasmp", &::nvte_built_with_cublasmp); m.def("initialize_cgemm_communicator", &InitializeCgemmCommunicator); m.def("is_collective_gemm_with_cublasmp", &IsCollectiveGemmWithCublasmp); m.def("get_cgemm_num_max_streams", &GetCgemmNumMaxStreams); - m.def("get_grouped_gemm_setup_workspace_size", &nvte_get_grouped_gemm_setup_workspace_size); #ifdef NVTE_WITH_NCCL_EP m.def("set_ep_bootstrap_params", &SetEpBootstrapParams, pybind11::arg("unique_id_bytes"), pybind11::arg("ep_size"), pybind11::arg("rank_within_group"), pybind11::arg("num_experts"),