diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 4dbcc9e7f..a0136faaa 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -746,6 +746,21 @@ else() # USE_ROCM list(APPEND transformer_engine_LINKER_LIBS hiprtc) find_package(hipblaslt) list(APPEND transformer_engine_LINKER_LIBS roc::hipblaslt) + # Device scalar alpha (ROCm/rocm-libraries#13482) is not in a hipBLASLt release yet; detect the + # enum rather than a version, since dev builds can share a version number with releases. + include(CheckCXXSourceCompiles) + include(CMakePushCheckState) + cmake_push_check_state(RESET) + set(CMAKE_REQUIRED_LIBRARIES roc::hipblaslt hip::host) + unset(NVTE_HIPBLASLT_DEVICE_SCALAR_ALPHA CACHE) + check_cxx_source_compiles(" + #include + int main() { return HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_SCALAR_BETA_HOST_EXT; }" + NVTE_HIPBLASLT_DEVICE_SCALAR_ALPHA) + cmake_pop_check_state() + if(NVTE_HIPBLASLT_DEVICE_SCALAR_ALPHA) + target_compile_definitions(transformer_engine PRIVATE NVTE_HIPBLASLT_DEVICE_SCALAR_ALPHA) + endif() if(USE_FUSED_ATTN_AOTRITON) target_compile_definitions(transformer_engine PUBLIC USE_FUSED_ATTN_AOTRITON) list(APPEND transformer_engine_LINKER_LIBS aotriton) diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index f98efd5dd..f5866fc9d 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -396,7 +397,7 @@ static void launch_dequant_fp4_to_bf16( } // Compute per-row alpha vector on device for NVFP4 GEMM: -// alpha_out[i] = alpha_in * amax_A * amax_B / (fp4_max^2 * fp8_max^2) for i in [0, m) +// alpha_out[i] = alpha_in * amax_A * amax_B / (fp4_max^2 * fp8_max_A * fp8_max_B) for i in [0, m) // Used with HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST, which // expects a device vector of length m for alpha, while beta stays on the host. __global__ void compute_fp4_alpha_vector_kernel(float alpha_in, const float* __restrict__ amax_A, @@ -415,7 +416,8 @@ __global__ void compute_fp4_alpha_vector_kernel(float alpha_in, const float* __r // [dim x 8] row-major with block stride 8 (one 128-element K-tile == 8 blocks of 16), dim rows tall // (NOT padded to 128). Iterating over the OUTPUT writes every byte once -- valid scale or 0 for the // padded blocks of the last K-tile -- so no separate zeroing memset is needed. Both operands are -// processed in one grid; thread 0 also writes c0 = alpha * amax_A * amax_B / (fp4_max^2 fp8_max^2). +// processed in one grid; thread 0 also writes +// c0 = alpha * amax_A * amax_B / (fp4_max^2 * fp8_max_A * fp8_max_B). __global__ void nvfp4_prep_kernel(const uint8_t* __restrict__ inA, uint8_t* __restrict__ outA, int dimA, int in_strideA, const uint8_t* __restrict__ inB, uint8_t* __restrict__ outB, int dimB, int in_strideB, int k_scale, @@ -604,32 +606,29 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla return ret; } -// Prepare a native (gfx1250) NVFP4 hipBLASLt GEMM: swizzle the block scales and fold the per-tensor -// amax correction into a HOST scalar alpha. +// Prepare a native (gfx1250) NVFP4 hipBLASLt GEMM: swizzle the block scales and compute the +// per-tensor amax correction on device. // -// hipBLASLt applies the FP8 (E4M3) block scales but has no NVFP4 solution that accepts a device -// alpha, so the per-tensor (global) correction cannot ride on a device alpha the way the CUDA path -// does. A single fused kernel swizzles both operands' scales and forms -// c0 = alpha * amax_A * amax_B / (fp4_max^2 * fp8_max^2) -// on device; only that scalar is copied back to the host. The GEMM then computes c0*(A*B) + beta*C -// in its fp32 accumulator with a single round to the output dtype (matching the reference addmm). -// No post-GEMM scaling or C copy is needed. -// -// FIXME(AIHPBLAS-3571): the host-scalar-alpha detour (device->host c0 copy + stream sync) only -// exists because hipBLASLt has no NVFP4 device-alpha solution. Once -// https://amd-hub.atlassian.net/browse/AIHPBLAS-3571 is fixed, pass c0 as a device scalar alpha and -// drop the c0 readback / hipStreamSynchronize below. +// hipBLASLt applies the FP8 (E4M3) block scales. A single fused kernel swizzles both operands' +// scales and forms +// c0 = alpha * amax_A * amax_B / (fp4_max^2 * fp8_max_A * fp8_max_B) +// in device memory. c0 is then used as the GEMM alpha, so the GEMM computes c0*(A*B) + beta*C in +// its fp32 accumulator with a single round to the output dtype (matching the reference addmm). +// No post-GEMM scaling or C copy is needed. hipblaslt_gemm() passes c0 as a device scalar alpha +// when hipBLASLt has a matching solution and otherwise reads it back to the host. static void setup_nvfp4_gemm_native( const transformer_engine::Tensor& inputA, cublasOperation_t transa, const transformer_engine::Tensor& inputB, cublasOperation_t transb, int m, int n, int k, float alpha_in, - void* workspace, size_t& workspaceSize, float* alpha_out, + void* workspace, size_t& workspaceSize, const float** c0_dev_out, const void** a_scale_swz_out, const void** b_scale_swz_out, hipStream_t stream) { const float fp4_max = 6.0f; - const float fp8_max = te_fp8_fnuz() ? 240.0f : 448.0f; - const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max * fp8_max); + // 4over6 tensors carry an E4M3 bound of 256, not 448. + const float fp8_max_A = static_cast(inputA.nvfp4_e4m3_max); + const float fp8_max_B = static_cast(inputB.nvfp4_e4m3_max); + const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max_A * fp8_max_B); const float* amax_A = (transa == CUBLAS_OP_T) ? reinterpret_cast(inputA.amax.dptr) @@ -672,15 +671,11 @@ static void setup_nvfp4_gemm_native( float* c0_dev = reinterpret_cast(base + res_a + res_b); // One fused launch swizzles both operands' scales (padding written inline, no memset) and forms - // c0 on device; a single 4-byte copy then brings c0 to the host for the GEMM's host scalar alpha. + // c0 on device. launch_nvfp4_prep(a_sinv.dptr, swz_a, m, static_cast(a_sinv.shape[1]), b_sinv.dptr, swz_b, n, static_cast(b_sinv.shape[1]), k_scale, num_ktiles, amax_A, amax_B, alpha_in, factor_inv, c0_dev, stream); - float host_c0 = 0.0f; - // FIXME(AIHPBLAS-3571): drop this c0 readback + sync once hipBLASLt supports NVFP4 device alpha. - NVTE_CHECK_CUDA(hipMemcpyAsync(&host_c0, c0_dev, sizeof(float), hipMemcpyDeviceToHost, stream)); - NVTE_CHECK_CUDA(hipStreamSynchronize(stream)); - *alpha_out = host_c0; + *c0_dev_out = c0_dev; *a_scale_swz_out = swz_a; *b_scale_swz_out = swz_b; } @@ -702,8 +697,8 @@ static void dequant_fp4_gemm_inputs( const float fp4_max = 6.0f; // Read the bound per operand as upstream does: 4over6 tensors carry 256, not 448. - const float fp8_max_A = te_fp8_fnuz() ? 240.0f : static_cast(inputA.nvfp4_e4m3_max); - const float fp8_max_B = te_fp8_fnuz() ? 240.0f : static_cast(inputB.nvfp4_e4m3_max); + const float fp8_max_A = static_cast(inputA.nvfp4_e4m3_max); + const float fp8_max_B = static_cast(inputB.nvfp4_e4m3_max); const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max_A * fp8_max_B); const float* amax_A = (transa == CUBLAS_OP_T) @@ -1447,6 +1442,61 @@ static void init_hipblaslt_handles(hipblasLtHandle_t* hipblaslt_handles) { } } +// Select device scalar alpha if hipBLASLt has a kernel for this GEMM config (not all NVFP4 kernels +// support it yet), probed once per config. Returns false and leaves host pointer mode otherwise. +static bool use_device_scalar_alpha(hipblasLtHandle_t handle, hipblasLtMatmulDesc_t operationDesc, + hipblasLtMatrixLayout_t Adesc, hipblasLtMatrixLayout_t Bdesc, + hipblasLtMatrixLayout_t Ddesc, size_t workspaceSize, + const GemmAlgoCache::Key& cfg) { +#ifdef NVTE_HIPBLASLT_DEVICE_SCALAR_ALPHA + static std::mutex mt; + static std::map supported; + + int32_t pointer_mode = HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_SCALAR_BETA_HOST_EXT; + NVTE_CHECK_HIPBLASLT(hipblasLtMatmulDescSetAttribute( + operationDesc, HIPBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode))); + + std::lock_guard lock(mt); + auto it = supported.find(cfg); + if (it == supported.end()) { + hipblasLtMatmulPreference_t preference = nullptr; + NVTE_CHECK_HIPBLASLT(hipblasLtMatmulPreferenceCreate(&preference)); + NVTE_CHECK_HIPBLASLT(hipblasLtMatmulPreferenceSetAttribute( + preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize))); + hipblasLtMatmulHeuristicResult_t result; + int count = 0; + const hipblasStatus_t status = hipblasLtMatmulAlgoGetHeuristic( + handle, operationDesc, Adesc, Bdesc, Ddesc, Ddesc, preference, 1, &result, &count); + NVTE_CHECK_HIPBLASLT(hipblasLtMatmulPreferenceDestroy(preference)); + it = supported.emplace(cfg, status == HIPBLAS_STATUS_SUCCESS && count > 0).first; + if (!it->second) { + // Shapes vary with MoE routing but kernel availability does not, so warn once per shape-free config. + static std::set warned; + GemmAlgoCache::Key family = cfg; + family.m = family.n = family.k = family.lda = family.ldb = family.ldd = 0; + if (warned.insert(family).second) { + // std::string because concat_strings would print a string_view char by char. + const std::string d_type(typeNameMapper.getName(cfg.d_type)); + const std::string epilogue(cfg.epilogue == HIPBLASLT_EPILOGUE_DEFAULT + ? "none" : epilogueNameMapper.getName(cfg.epilogue)); + NVTE_WARN("hipBLASLt has no device scalar alpha kernel for NVFP4 GEMM d_type=", d_type, + " epilogue=", epilogue, " (first seen at m=", cfg.m, " n=", cfg.n, " k=", cfg.k, + "); falling back to host alpha, which synchronizes the stream on every call."); + } + } + } + + if (!it->second) { + pointer_mode = HIPBLASLT_POINTER_MODE_HOST; + NVTE_CHECK_HIPBLASLT(hipblasLtMatmulDescSetAttribute( + operationDesc, HIPBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode))); + } + return it->second; +#else + return false; +#endif +} + void hipblaslt_gemm(const Tensor *inputA, const Tensor *inputB, @@ -1479,11 +1529,11 @@ void hipblaslt_gemm(const Tensor *inputA, // On every other architecture hipBLASLt cannot consume the block scales, so we fall // back to dequantizing FP4 -> BF16 (block scales only) and running a standard BF16 GEMM. // - // The per-tensor amax correction (alpha * amax_A * amax_B / (fp4_max^2 * fp8_max^2)) is applied - // differently per path because hipBLASLt has no NVFP4 solution that accepts a device alpha: the - // native path folds it into a HOST scalar alpha (setup_nvfp4_gemm_native) so the GEMM computes - // c0*(A*B) + beta*C directly, while the BF16 dequant fallback folds it into a device alpha vector - // (HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST). + // The per-tensor amax correction c0 = alpha * amax_A * amax_B / (fp4_max^2 * fp8_max_A * fp8_max_B) + // is folded into the GEMM alpha. The native path passes c0 as a device scalar alpha + // (HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_SCALAR_BETA_HOST_EXT) when hipBLASLt has a solution for it + // and otherwise reads c0 back to a host scalar alpha; the BF16 dequant fallback uses a device + // alpha vector (HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST). // // Only NVFP4 uses this dequant fallback. MXFP4 (also an FP4 data dtype) has a native // hipBLASLt path, so gate on the scaling mode rather than the FP4 data dtype (which @@ -1502,17 +1552,16 @@ void hipblaslt_gemm(const Tensor *inputA, const void* alpha_ptr = static_cast(&alpha); const void* beta_ptr = static_cast(&beta); - // Native NVFP4 (gfx1250) folds the per-tensor amax correction into a host scalar alpha (c0) so - // hipBLASLt computes c0*(A*B) + beta*C in one fp32-accumulated round; beta stays the caller's. + // Native NVFP4 (gfx1250) uses c0 as alpha so hipBLASLt computes c0*(A*B) + beta*C in one + // fp32-accumulated round; beta stays the caller's. alpha_ptr is set once the pointer mode is known. float nvfp4_c0 = 1.0f; + const float* nvfp4_c0_dev = nullptr; const void* nvfp4_scaleA_swz = nullptr; const void* nvfp4_scaleB_swz = nullptr; if (use_nvfp4_native) { setup_nvfp4_gemm_native(*inputA, transa, *inputB, transb, m, n, k, alpha, - workspace, workspaceSize, &nvfp4_c0, + workspace, workspaceSize, &nvfp4_c0_dev, &nvfp4_scaleA_swz, &nvfp4_scaleB_swz, stream); - alpha_ptr = static_cast(&nvfp4_c0); - beta_ptr = static_cast(&beta); } else if (use_nvfp4) { dequant_fp4_gemm_inputs(param, *inputA, transa, *inputB, transb, m, n, k, alpha, workspace, workspaceSize, @@ -1714,7 +1763,7 @@ void hipblaslt_gemm(const Tensor *inputA, // Native NVFP4 GEMM (gfx1250): the FP4 (E2M1) operands carry FP8 (E4M3) scales, one per // 16-element block along the contraction (innermost) dimension. hipBLASLt applies these block // scales (K-tiled, see setup_nvfp4_gemm_native); the per-tensor amax correction rides on the - // host scalar alpha (c0) so it is folded into the GEMM itself. + // alpha (c0) so it is folded into the GEMM itself. scaling_mode = HIPBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; NVTE_CHECK_HIPBLASLT( hipblasLtMatmulDescSetAttribute(operationDesc, HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER, @@ -1776,8 +1825,8 @@ void hipblaslt_gemm(const Tensor *inputA, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); - // Only the BF16 dequant fallback uses the device alpha vector; the native NVFP4 path applies - // its per-tensor correction after the GEMM, so it keeps the default host pointer mode. + // Only the BF16 dequant fallback uses the device alpha vector; the native NVFP4 path selects its + // pointer mode below. if (use_nvfp4 && !use_nvfp4_native) { int32_t pointer_mode = HIPBLASLT_POINTER_MODE_ALPHA_DEVICE_VECTOR_BETA_HOST; NVTE_CHECK_HIPBLASLT(hipblasLtMatmulDescSetAttribute( @@ -1790,6 +1839,18 @@ void hipblaslt_gemm(const Tensor *inputA, ((use_fp8 || use_mxfp4) && gelu) ? aux_type : (hipDataType)-1, m, n, k, param.lda, param.ldb, ldd, param.transA, param.transB, scaling_mode, epilogue, use_nvfp4 && !use_nvfp4_native); + + if (use_nvfp4_native) { + if (use_device_scalar_alpha(handle, operationDesc, Adesc, Bdesc, Ddesc, workspaceSize, gemm_cfg)) { + alpha_ptr = static_cast(nvfp4_c0_dev); + } else { + NVTE_CHECK_CUDA(hipMemcpyAsync(&nvfp4_c0, nvfp4_c0_dev, sizeof(float), hipMemcpyDeviceToHost, + stream)); + NVTE_CHECK_CUDA(hipStreamSynchronize(stream)); + alpha_ptr = static_cast(&nvfp4_c0); + } + } + GemmAlgoCache::Algo cached_algo; if (algoCache.find(gemm_cfg, workspaceSize, cached_algo) == 0 || !cached_algo.algo.has_value()) {