diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 13880e0d61..40c3182259 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -184,6 +184,7 @@ list(APPEND transformer_engine_cpp_sources gemm/config.cpp normalization/common.cpp normalization/rtc_dispatch.cpp + cast/mxfp8/specialized/rtc_dispatch.cpp normalization/layernorm/ln_api.cpp normalization/rmsnorm/rmsnorm_api.cpp util/cuda_driver.cpp @@ -554,6 +555,21 @@ make_string_header_from_file(transpose/rtc/transpose.cu string_code_transpose_rtc_transpose_cu) make_string_header_from_file(transpose/rtc/swap_first_dims.cu string_code_transpose_rtc_swap_first_dims_cu) + +# MXFP8 specialized quantize: NVRTC bundled headers + RTC source file +make_string_header_from_file(util/ptx.cuh + string_code_util_ptx_cuh) +make_string_header_from_file(cast/mxfp8/specialized/swizzle.cuh + string_code_cast_mxfp8_specialized_swizzle_cuh) +make_string_header_from_file(cast/mxfp8/specialized/state_counter.cuh + string_code_cast_mxfp8_specialized_state_counter_cuh) +make_string_header_from_file(cast/mxfp8/specialized/quantize_mxfp8.cuh + string_code_cast_mxfp8_specialized_quantize_mxfp8_cuh) +make_string_header_from_file(cast/mxfp8/specialized/rtc/quantize_mxfp8_rowwise.cu + string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_rowwise_cu) +make_string_header_from_file(cast/mxfp8/specialized/rtc/quantize_mxfp8_bidimensional.cu + string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_bidimensional_cu) + make_string_header_from_file(fused_softmax/scaled_masked_softmax.cu string_code_fused_softmax_scaled_masked_softmax_cu) make_string_header_from_file(fused_softmax/scaled_upper_triang_masked_softmax.cu @@ -586,6 +602,7 @@ make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_fwd_kernel.cu string_code_normalization_rmsnorm_rtc_rmsnorm_fwd_kernel_cu) make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu string_code_normalization_rmsnorm_rtc_rmsnorm_bwd_kernel_cu) + target_include_directories(transformer_engine PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/string_headers") @@ -607,6 +624,16 @@ target_compile_definitions(transformer_engine PRIVATE NVTE_BUILD_LEGACY_STATIC_NORM=$) +# Default OFF: the 1x32 rowwise and 32x32 bidimensional MXFP8 specialized +# cast-only kernels are compiled via NVRTC at runtime. Set ON to additionally +# instantiate their static fallbacks for NVTE_DISABLE_NVRTC=1. +option(NVTE_BUILD_LEGACY_STATIC_MXFP8 + "Also compile static specialized MXFP8 cast-only kernels for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_MXFP8=$) + # Compiler options set(nvte_sources_with_fast_math) list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 8c57f112d2..773fed991d 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -19,9 +19,11 @@ #include "../../common.h" #include "../../util/math.h" #include "../../util/ptx.cuh" +#include "../../util/rtc.h" #include "../../utils.cuh" #include "../core/common.cuh" #include "specialized/quantize_mxfp8.cuh" +#include "specialized/rtc_dispatch.cuh" #include "swizzle.cuh" namespace transformer_engine { @@ -680,30 +682,50 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, switch (scaling_type) { case ScalingType::ROWWISE: { using traits = specialized::CastTraits; - auto kernel = specialized::quantize_mxfp8_kernel_cast_only; - - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem)); - - dim3 block(traits::threadLayout::num, traits::warpLayout::N, - traits::warpLayout::M); - dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN, - (rows + traits::blockDimM - 1) / traits::blockDimM); - kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(output->data.dptr), - scales_rowwise_ptr, noop_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + auto *rowwise_input = + reinterpret_cast(input.data.dptr); + auto *rowwise_output = + reinterpret_cast(output->data.dptr); + + // Prefer the NVRTC-compiled kernel; fall back to the static + // instantiation only when it was compiled in and NVRTC is disabled. +#if NVTE_BUILD_LEGACY_STATIC_MXFP8 + const bool use_rtc = rtc::is_enabled(); +#else + constexpr bool use_rtc = true; +#endif + if (use_rtc) { + specialized::launch_rowwise_cast_only_rtc( + rowwise_input, rowwise_output, scales_rowwise_ptr, noop_ptr, + static_cast(rows), static_cast(cols), + static_cast(scale_stride_rowwise), + static_cast(scale_stride_colwise), stream); + } else { +#if NVTE_BUILD_LEGACY_STATIC_MXFP8 + auto kernel = specialized::quantize_mxfp8_kernel_cast_only; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem)); + dim3 block(traits::threadLayout::num, traits::warpLayout::N, + traits::warpLayout::M); + dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN, + (rows + traits::blockDimM - 1) / traits::blockDimM); + kernel<<>>( + rowwise_input, rowwise_output, scales_rowwise_ptr, noop_ptr, rows, cols, + scale_stride_rowwise, scale_stride_colwise); +#else + NVTE_ERROR( + "MXFP8 rowwise specialized cast-only kernel requires NVRTC. Unset " + "NVTE_DISABLE_NVRTC, or rebuild with NVTE_BUILD_LEGACY_STATIC_MXFP8=ON " + "for the static fallback."); +#endif + } break; } case ScalingType::BIDIMENSIONAL: { using traits = specialized::CastTraits; - auto kernel = specialized::quantize_mxfp8_kernel_cast_only; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem)); - // TMA for loading, so that we don't need STS for transposing alignas(64) CUtensorMap tensor_map_input{}; constexpr size_t input_type_bit_size = TypeInfo::size; create_2D_tensor_map(tensor_map_input, input.data, rows, cols, @@ -725,13 +747,38 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, cols, 0, output_type_bit_size, traits::output_swizzle_pattern); - dim3 block(traits::rowThreadLayout::num, traits::numWarps); - dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N, - (rows + traits::blockDIM::M - 1) / traits::blockDIM::M); - kernel<<>>( - tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, - scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, rows, cols, - scale_stride_rowwise, scale_stride_colwise); +#if NVTE_BUILD_LEGACY_STATIC_MXFP8 + const bool use_rtc = rtc::is_enabled(); +#else + constexpr bool use_rtc = true; +#endif + if (use_rtc) { + specialized::launch_bidimensional_cast_only_rtc( + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, + scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, + static_cast(rows), static_cast(cols), + static_cast(scale_stride_rowwise), + static_cast(scale_stride_colwise), stream); + } else { +#if NVTE_BUILD_LEGACY_STATIC_MXFP8 + auto kernel = specialized::quantize_mxfp8_kernel_cast_only; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem)); + dim3 block(traits::rowThreadLayout::num, traits::numWarps); + dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N, + (rows + traits::blockDIM::M - 1) / traits::blockDIM::M); + kernel<<>>( + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, + scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, rows, cols, + scale_stride_rowwise, scale_stride_colwise); +#else + NVTE_ERROR( + "MXFP8 bidimensional specialized cast-only kernel requires NVRTC. Unset " + "NVTE_DISABLE_NVRTC, or rebuild with NVTE_BUILD_LEGACY_STATIC_MXFP8=ON " + "for the static fallback."); +#endif + } break; } diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh index 81e11c40b8..75d2803dd9 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -11,11 +11,35 @@ #ifndef TRANSFORMER_ENGINE_SPECIALIZED_QUANTIZE_MXFP8_CUH_ #define TRANSFORMER_ENGINE_SPECIALIZED_QUANTIZE_MXFP8_CUH_ +#if !defined(__CUDACC_RTC__) #include #include "../../../util/ptx.cuh" #include "state_counter.cuh" #include "swizzle.cuh" +#else +// NVRTC build: these are provided as in-memory headers (see rtc_dispatch), and +// common.h-style host includes are unavailable. +#include "ptx.cuh" +#include "state_counter.cuh" +#include "swizzle.cuh" + +// CUtensorMap / CUtensorMapSwizzle normally come from , which NVRTC +// cannot include. The bidimensional kernels only pass a CUtensorMap by value +// (__grid_constant__) and never introspect it, and ptx.cuh's TMA ops take a +// uint64_t*, so an ABI-compatible opaque struct (128 B, 64-B aligned, matching +// the real CUtensorMap_st) is sufficient. The host builds the real descriptor. +struct alignas(64) CUtensorMap_st { + uint64_t opaque[16]; +}; +using CUtensorMap = CUtensorMap_st; +enum CUtensorMapSwizzle { + CU_TENSOR_MAP_SWIZZLE_NONE = 0, + CU_TENSOR_MAP_SWIZZLE_32B = 1, + CU_TENSOR_MAP_SWIZZLE_64B = 2, + CU_TENSOR_MAP_SWIZZLE_128B = 3, +}; +#endif namespace transformer_engine { namespace dispatch { @@ -68,9 +92,9 @@ __device__ __forceinline__ e8m0_t to_e8m0(IType amax) { constexpr uint16_t max_norm_rcp = _Quantized_Limits::max_norm_rcp; float amax_fp32; - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { ptx::fma_f32_f16(amax_fp32, reinterpret_cast(amax), max_norm_rcp); - } else if constexpr (std::is_same_v) { + } else if constexpr (detail::is_same_v) { ptx::fma_f32_bf16(amax_fp32, reinterpret_cast(amax), max_norm_rcp); } else { amax_fp32 = 0.0f; @@ -78,7 +102,7 @@ __device__ __forceinline__ e8m0_t to_e8m0(IType amax) { } return ptx::float_to_e8m0(amax_fp32); #else - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { return ptx::float_to_e8m0(__fmaf_ieee_rn(amax, Quantized_Limits::max_norm_rcp, 0.0f)); } else { float amax_fp32 = static_cast(amax); @@ -166,14 +190,14 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { }; // 1x32 -template = 0> -__global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__restrict__ input, - typename CastTraits::OType *__restrict__ output, - e8m0_t *__restrict__ scales_rowwise, - const float *noop, int32_t rows, int32_t cols, - int32_t scale_stride_rowwise, - int32_t scale_stride_colwise) { +// Shared device body for the 1x32 rowwise cast-only kernel. Extracted so both +// the statically-instantiated __global__ below and the NVRTC entry point +// (rtc/quantize_mxfp8_rowwise.cu) can reuse it without duplication. +template +__device__ __forceinline__ void quantize_mxfp8_rowwise_cast_only_body( + typename CastTraits::IType *__restrict__ input, typename CastTraits::OType *__restrict__ output, + e8m0_t *__restrict__ scales_rowwise, const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -262,7 +286,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re return; } - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { float thread_amax = 0.f; IType2 *rInput2 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); #pragma unroll @@ -383,7 +407,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re return; } - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { float thread_amax = 0.f; IType2 *rInput2 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); #pragma unroll @@ -499,17 +523,17 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re block_coords.x = blockIdx.x * CastTraits::blockDimN; constexpr int32_t stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems; - using PreferredDataType = std::conditional_t< + using PreferredDataType = detail::conditional_t< stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + detail::conditional_t>>>; int2 end_coords; - end_coords.y = std::min(block_coords.y + CastTraits::blockDimM, rows); - end_coords.x = std::min((block_coords.x + CastTraits::blockDimN) / CastTraits::chunkElems, - scale_stride_rowwise); + end_coords.y = detail::min(block_coords.y + CastTraits::blockDimM, rows); + end_coords.x = detail::min((block_coords.x + CastTraits::blockDimN) / CastTraits::chunkElems, + scale_stride_rowwise); int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::chunkElems); @@ -558,6 +582,24 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +// 1x32 rowwise cast-only kernel (statically-instantiated entry point). +template = 0> +__global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__restrict__ input, + typename CastTraits::OType *__restrict__ output, + e8m0_t *__restrict__ scales_rowwise, + const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, + int32_t scale_stride_colwise) { + quantize_mxfp8_rowwise_cast_only_body(input, output, scales_rowwise, noop, rows, cols, + scale_stride_rowwise, scale_stride_colwise); +} + +// The 32x32 bidimensional (rowwise+colwise) kernels use TMA: CUtensorMap +// grid-constant params and cp_async_bulk_tensor. Under NVRTC, CUtensorMap is +// provided as an opaque struct (above) and the descriptor is built host-side, +// so these are NVRTC-capable (the non-warp-specialized variant is the RTC entry +// point; see rtc/quantize_mxfp8_bidimensional.cu). enum class ColwiseReduceMax : int32_t { Atom = 0, Red = 1, // it's actually the same to Atom @@ -566,9 +608,13 @@ enum class ColwiseReduceMax : int32_t { Num = 4 }; -// 32x32 -template -struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { +// 32x32. Parameterized on the perf-relevant knobs (pipeline stages, per-block +// iteration width, 4x-vs-2x convert) so they can be autotuned at runtime via +// NVRTC. Defaults reproduce the shipped configuration; CastTraits<...,true,true> +// below is exactly this with the defaults, so existing behavior is unchanged. +template +struct BidimTraitsImpl { static constexpr bool isRowwise = true; static constexpr bool isColwise = true; using IType = _IType; @@ -586,7 +632,7 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { using rowWarpDim = Layout; using colWarpDim = Layout; using warpDim = - Layout; + Layout; static constexpr bool _tma_swizzle = true; using warpLayout = Layout<1, 2>; @@ -601,17 +647,17 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { using blockIterDim = Layout; - using iterLayout = Layout<1, 4>; + using iterLayout = Layout<1, _IterN>; using blockDIM = Layout; - static constexpr int32_t numStages = 2; + static constexpr int32_t numStages = _NumStages; using inputUnitType = uint4; static constexpr int32_t rowNumElemsPerUnit = sizeof(inputUnitType) / sizeof(IType); static constexpr int32_t rowNumUnitsPerChunk = rowChunkElems / rowNumElemsPerUnit; // TODO: set condition for float - using inputElemSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>; - using inputUnitSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>; + using inputElemSwz = detail::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>; + using inputUnitSwz = detail::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>; using colIndexSwz = swz::Swizzle<5, 0, 5>; @@ -620,10 +666,10 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { rowChunkElems * sizeof(OType) / sizeof(rowOutputUnitType); static constexpr int32_t rowOutNumElemsPerUnit = sizeof(rowOutputUnitType) / sizeof(OType); - using rowOutputChunkSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>; - using colOutputSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>; + using rowOutputChunkSwz = detail::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>; + using colOutputSwz = detail::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>; - static constexpr bool _use_cvt_4x = true; + static constexpr bool _use_cvt_4x = _UseCvt4x; static constexpr bool _use_warp_specialization = false; static constexpr bool _need_wait_group = iterLayout::num > numStages; static constexpr bool _reuse_input_out_smem = false; @@ -663,21 +709,33 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { _need_smem_for_colwise_reduce ? 32 * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul; static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; - static constexpr size_t smem = _reuse_input_out_smem - ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + - smem_alignment + smem_rowwise_scale + smem_colwise_reduce) - : (smemInput + smemRowwiseOutput + smemColwiseOutput + - smem_alignment + smem_rowwise_scale + smem_colwise_reduce); + static constexpr size_t smem = + _reuse_input_out_smem ? (detail::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + + smem_alignment + smem_rowwise_scale + smem_colwise_reduce) + : (smemInput + smemRowwiseOutput + smemColwiseOutput + smem_alignment + + smem_rowwise_scale + smem_colwise_reduce); }; +// Shipped bidimensional config: BidimTraitsImpl with the default knobs. Keeping +// this as the CastTraits<...,true,true> specialization means all existing callers +// and the static kernel are unchanged. +template +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> + : BidimTraitsImpl<_IType, _OType> {}; + +// Tunable alias for runtime autotuning over (numStages, iterLayout::N, use_cvt_4x). +template +using BidimTunableTraits = BidimTraitsImpl; + __device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) { return (x + align - 1) & ~((align)-1); } // 32x32 template = 0, - std::enable_if_t = 0> + detail::enable_if_t = 0, + detail::enable_if_t = 0> // __launch_bounds__(CastTraits::numThreads) __global__ void quantize_mxfp8_kernel_cast_only( const __grid_constant__ CUtensorMap tensor_map_input, @@ -917,7 +975,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( ptx::mbarrier_arrive_expect_tx(&ldg_consumer[read_state.index()], 0u); } - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { } else { static_assert(CastTraits::_colwise_reduce_max == ColwiseReduceMax::Redux, "Only Redux is implemented"); @@ -1088,18 +1146,18 @@ __global__ void quantize_mxfp8_kernel_cast_only( ptx::numbered_barrier_sync(CastTraits::warpLayout::num * 32, 0u); constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< + using PreferredDataType = detail::conditional_t< stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + detail::conditional_t>>>; int2 end_coords; - end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); + end_coords.y = detail::min(block_coords.y + CastTraits::blockDIM::M, rows); end_coords.x = - std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, - scale_stride_rowwise); + detail::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); @@ -1150,15 +1208,18 @@ __global__ void quantize_mxfp8_kernel_cast_only( #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } -template = 0, - std::enable_if_t = 0> -__global__ void quantize_mxfp8_kernel_cast_only( - const __grid_constant__ CUtensorMap tensor_map_input, - const __grid_constant__ CUtensorMap tensor_map_rowwise_output, - const __grid_constant__ CUtensorMap tensor_map_colwise_output, e8m0_t *scales_rowwise, - e8m0_t *scales_colwise, const float *noop, int32_t rows, int32_t cols, - int32_t scale_stride_rowwise, int32_t scale_stride_colwise) { +// Shared device body for the 32x32 bidimensional (non-warp-specialized) +// cast-only kernel. Extracted so both the statically-instantiated __global__ +// below and the NVRTC entry point (rtc/quantize_mxfp8_bidimensional.cu) reuse +// it. __grid_constant__ is only valid on __global__ params, so the maps are +// passed by const-ref here; __forceinline__ collapses this back into the global +// so the grid-constant param addresses are preserved. +template +__device__ __forceinline__ void quantize_mxfp8_bidimensional_cast_only_body( + const CUtensorMap &tensor_map_input, const CUtensorMap &tensor_map_rowwise_output, + const CUtensorMap &tensor_map_colwise_output, e8m0_t *scales_rowwise, e8m0_t *scales_colwise, + const float *noop, int32_t rows, int32_t cols, int32_t scale_stride_rowwise, + int32_t scale_stride_colwise) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -1362,7 +1423,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( } } - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { if constexpr (CastTraits::_colwise_reduce_max == ColwiseReduceMax::Atom || CastTraits::_colwise_reduce_max == ColwiseReduceMax::Red) { } else if constexpr (CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { @@ -1546,17 +1607,18 @@ __global__ void quantize_mxfp8_kernel_cast_only( if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; - using PreferredDataType = std::conditional_t< + using PreferredDataType = detail::conditional_t< stride_in_smem % 16 == 0, uint4, - std::conditional_t< - stride_in_smem % 8 == 0, uint2, - std::conditional_t>>>; + detail::conditional_t>>>; int2 end_coords; - end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); - end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, - scale_stride_rowwise); + end_coords.y = detail::min(block_coords.y + CastTraits::blockDIM::M, rows); + end_coords.x = + detail::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); int2 valid_coords; valid_coords.y = end_coords.y - block_coords.y; valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); @@ -1607,6 +1669,21 @@ __global__ void quantize_mxfp8_kernel_cast_only( #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +// 32x32 bidimensional cast-only kernel, non-warp-specialized (static entry point). +template = 0, + detail::enable_if_t = 0> +__global__ void quantize_mxfp8_kernel_cast_only( + const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_rowwise_output, + const __grid_constant__ CUtensorMap tensor_map_colwise_output, e8m0_t *scales_rowwise, + e8m0_t *scales_colwise, const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise) { + quantize_mxfp8_bidimensional_cast_only_body( + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, scales_rowwise, + scales_colwise, noop, rows, cols, scale_stride_rowwise, scale_stride_colwise); +} + } // namespace specialized } // namespace quantize_kernel } // namespace mxfp8 diff --git a/transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_bidimensional.cu b/transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_bidimensional.cu new file mode 100644 index 0000000000..3e8c30ec11 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_bidimensional.cu @@ -0,0 +1,38 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC source file for the 32x32 bidimensional (rowwise+colwise) MXFP8 +// cast-only kernel (non-warp-specialized variant). The concrete input/output +// element types are substituted at runtime via the __ITYPE__/__OTYPE__ +// placeholders (see rtc_dispatch). CUtensorMap is provided as an opaque struct +// by the included header; the host builds the real TMA descriptors. + +#include "specialized_quantize_mxfp8.cuh" + +using namespace transformer_engine; +namespace specialized = + transformer_engine::dispatch::mxfp8::quantize_kernel::specialized; // NOLINT(*) + +namespace { +// Substituted at compile time by the host dispatch. Defaults (2, 4, true) +// reproduce the shipped CastTraits<...,true,true> tiling. +using IType = __ITYPE__; +using OType = __OTYPE__; +using BidimTraits = + specialized::BidimTunableTraits; +} // namespace + +// Non-template entry point so the host can request it by a stable name. +__global__ void __launch_bounds__(BidimTraits::numThreads) quantize_mxfp8_bidimensional_rtc_kernel( + const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_rowwise_output, + const __grid_constant__ CUtensorMap tensor_map_colwise_output, e8m0_t *scales_rowwise, + e8m0_t *scales_colwise, const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise) { + specialized::quantize_mxfp8_bidimensional_cast_only_body( + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, scales_rowwise, + scales_colwise, noop, rows, cols, scale_stride_rowwise, scale_stride_colwise); +} diff --git a/transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_rowwise.cu b/transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_rowwise.cu new file mode 100644 index 0000000000..3b9259d76a --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_rowwise.cu @@ -0,0 +1,33 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC source file for the 1x32 rowwise MXFP8 cast-only kernel. The host +// bundles this (and the specialized kernel headers it needs) as in-memory +// strings; the concrete input/output element types are substituted at runtime +// via the __ITYPE__/__OTYPE__ placeholders (see rtc_dispatch). + +#include "specialized_quantize_mxfp8.cuh" + +using namespace transformer_engine; +namespace specialized = + transformer_engine::dispatch::mxfp8::quantize_kernel::specialized; // NOLINT(*) + +namespace { +// Substituted at compile time by the host dispatch. +using IType = __ITYPE__; +using OType = __OTYPE__; +using RowwiseTraits = specialized::CastTraits; +} // namespace + +// Non-template entry point so the host can request it by a stable name. +__global__ void __launch_bounds__(RowwiseTraits::numThreads) + quantize_mxfp8_rowwise_rtc_kernel(IType *__restrict__ input, OType *__restrict__ output, + e8m0_t *__restrict__ scales_rowwise, const float *noop, + int32_t rows, int32_t cols, int32_t scale_stride_rowwise, + int32_t scale_stride_colwise) { + specialized::quantize_mxfp8_rowwise_cast_only_body( + input, output, scales_rowwise, noop, rows, cols, scale_stride_rowwise, scale_stride_colwise); +} diff --git a/transformer_engine/common/cast/mxfp8/specialized/rtc_dispatch.cpp b/transformer_engine/common/cast/mxfp8/specialized/rtc_dispatch.cpp new file mode 100644 index 0000000000..c4242f9d58 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/rtc_dispatch.cpp @@ -0,0 +1,93 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include + +#include "../../../util/rtc.h" +#include "../../../util/string.h" + +// Generated string headers: raw source of the RTC kernel and the device headers +// it needs as in-memory includes. These are large; keeping this in a host-only +// .cpp (never compiled by cicc) avoids blowing up the device compiler in every +// TU that merely launches the kernel. +#include "string_code_cast_mxfp8_specialized_quantize_mxfp8_cuh.h" +#include "string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_bidimensional_cu.h" +#include "string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_rowwise_cu.h" +#include "string_code_cast_mxfp8_specialized_state_counter_cuh.h" +#include "string_code_cast_mxfp8_specialized_swizzle_cuh.h" +#include "string_code_util_ptx_cuh.h" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { +namespace specialized { + +void compile_rowwise_cast_only_rtc(const std::string &kernel_label, const std::string &itype_name, + const std::string &otype_name) { + auto &mgr = rtc::KernelManager::instance(); + if (mgr.is_compiled(kernel_label)) { + return; + } + + std::string code = string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_rowwise_cu; + code = regex_replace(code, "__ITYPE__", itype_name); + code = regex_replace(code, "__OTYPE__", otype_name); + + const std::vector headers = { + {string_code_cast_mxfp8_specialized_quantize_mxfp8_cuh, "specialized_quantize_mxfp8.cuh"}, + {string_code_util_ptx_cuh, "ptx.cuh"}, + {string_code_cast_mxfp8_specialized_state_counter_cuh, "state_counter.cuh"}, + {string_code_cast_mxfp8_specialized_swizzle_cuh, "swizzle.cuh"}, + }; + + // --device-int128: ptx.cuh uses __uint128_t; -default-device: treat the + // unannotated constexpr/inline helpers in ptx.cuh as __device__ under JIT. + const std::vector options = {"--device-int128", "-default-device"}; + constexpr rtc::ArchRequirement arch_requirement{100, rtc::ArchSpecificity::BlackwellSpecific}; + + mgr.compile(kernel_label, "quantize_mxfp8_rowwise_rtc_kernel", code, + "transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_rowwise.cu", + options, headers, arch_requirement); +} + +void compile_bidimensional_cast_only_rtc(const std::string &kernel_label, + const std::string &itype_name, + const std::string &otype_name, int num_stages, int iter_n, + bool use_cvt_4x) { + auto &mgr = rtc::KernelManager::instance(); + if (mgr.is_compiled(kernel_label)) { + return; + } + + std::string code = string_code_cast_mxfp8_specialized_rtc_quantize_mxfp8_bidimensional_cu; + code = regex_replace(code, "__ITYPE__", itype_name); + code = regex_replace(code, "__OTYPE__", otype_name); + code = regex_replace(code, "__NUM_STAGES__", std::to_string(num_stages)); + code = regex_replace(code, "__ITER_N__", std::to_string(iter_n)); + code = regex_replace(code, "__USE_CVT_4X__", use_cvt_4x ? "true" : "false"); + + const std::vector headers = { + {string_code_cast_mxfp8_specialized_quantize_mxfp8_cuh, "specialized_quantize_mxfp8.cuh"}, + {string_code_util_ptx_cuh, "ptx.cuh"}, + {string_code_cast_mxfp8_specialized_state_counter_cuh, "state_counter.cuh"}, + {string_code_cast_mxfp8_specialized_swizzle_cuh, "swizzle.cuh"}, + }; + const std::vector options = {"--device-int128", "-default-device"}; + constexpr rtc::ArchRequirement arch_requirement{100, rtc::ArchSpecificity::BlackwellSpecific}; + + mgr.compile( + kernel_label, "quantize_mxfp8_bidimensional_rtc_kernel", code, + "transformer_engine/common/cast/mxfp8/specialized/rtc/quantize_mxfp8_bidimensional.cu", + options, headers, arch_requirement); +} + +} // namespace specialized +} // namespace quantize_kernel +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine diff --git a/transformer_engine/common/cast/mxfp8/specialized/rtc_dispatch.cuh b/transformer_engine/common/cast/mxfp8/specialized/rtc_dispatch.cuh new file mode 100644 index 0000000000..068a2a712d --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/rtc_dispatch.cuh @@ -0,0 +1,214 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file rtc_dispatch.cuh + * \brief Host-side NVRTC dispatch for the specialized MXFP8 cast-only kernels. + */ + +#ifndef TRANSFORMER_ENGINE_SPECIALIZED_MXFP8_RTC_DISPATCH_CUH_ +#define TRANSFORMER_ENGINE_SPECIALIZED_MXFP8_RTC_DISPATCH_CUH_ + +#if !defined(__CUDACC_RTC__) + +#include + +#include + +#include "../../../util/rtc.h" +// NB: do not include util/string.h here — it pulls in , which is heavy +#include "quantize_mxfp8.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { +namespace specialized { + +// Compile (if not already cached) the rowwise cast-only RTC kernel for the given +// element-type spellings. +void compile_rowwise_cast_only_rtc(const std::string &kernel_label, const std::string &itype_name, + const std::string &otype_name); +void compile_bidimensional_cast_only_rtc(const std::string &kernel_label, + const std::string &itype_name, + const std::string &otype_name, int num_stages, int iter_n, + bool use_cvt_4x); + +template +inline const char *rtc_type_name() { + return detail::type_name(); +} +template <> +inline const char *rtc_type_name() { + return "fp16"; +} +template <> +inline const char *rtc_type_name() { + return "bf16"; +} +template <> +inline const char *rtc_type_name() { + return "fp8e4m3"; +} +template <> +inline const char *rtc_type_name() { + return "fp8e5m2"; +} + +template +inline void launch_rowwise_cast_only_rtc(IType *input, OType *output, e8m0_t *scales_rowwise, + const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise, + cudaStream_t stream) { + using traits = CastTraits; + + const std::string itype_name = rtc_type_name(); + const std::string otype_name = rtc_type_name(); + + const std::string kernel_label = + std::string("quantize_mxfp8_rowwise_cast_only,itype=") + itype_name + ",otype=" + otype_name; + + auto &mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(kernel_label)) { + compile_rowwise_cast_only_rtc(kernel_label, itype_name, otype_name); + } + + if (traits::smem > 0) { + mgr.set_function_attribute(kernel_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + static_cast(traits::smem)); + } + + dim3 block(traits::threadLayout::num, traits::warpLayout::N, traits::warpLayout::M); + dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN, + (rows + traits::blockDimM - 1) / traits::blockDimM); + mgr.launch(kernel_label, grid, block, static_cast(traits::smem), stream, input, + output, scales_rowwise, noop, rows, cols, scale_stride_rowwise, scale_stride_colwise); +} + +// Compile-time config for the bidimensional kernel. +struct BidimConfig { + int32_t num_stages; + int32_t iter_n; + bool use_cvt_4x; +}; + +// Derive the default from the static kernel traits so the static and RTC paths +// cannot drift when the shipped bidimensional configuration changes. +template +constexpr BidimConfig static_bidim_config() { + using traits = CastTraits; + using iter_layout = typename traits::iterLayout; + return {traits::numStages, iter_layout::N, traits::_use_cvt_4x}; +} + +template +inline BidimConfig select_bidim_config(int32_t rows, int32_t cols) { + // No tuned bidimensional overrides are shipped for now. In the future, + // this selector can be expanded with measured shape- or dtype-specific configurations while + // leaving unlisted problems on the static kernel's configuration. + (void)rows; + (void)cols; + return static_bidim_config(); +} + +// Compile+launch one concrete bidimensional config. Geometry/smem come straight +// from BidimTunableTraits, so every config is exactly what the JIT'd kernel uses. +template +inline void launch_bidim_impl(const CUtensorMap &tensor_map_input, + const CUtensorMap &tensor_map_rowwise_output, + const CUtensorMap &tensor_map_colwise_output, e8m0_t *scales_rowwise, + e8m0_t *scales_colwise, const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise, + cudaStream_t stream, const std::string &itype_name, + const std::string &otype_name) { + using traits = BidimTunableTraits; + const std::string kernel_label = std::string("quantize_mxfp8_bidimensional_cast_only,itype=") + + itype_name + ",otype=" + otype_name + + ",ns=" + std::to_string(NS) + ",itn=" + std::to_string(ITN) + + ",cvt=" + (CVT ? "4x" : "2x"); + auto &mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(kernel_label)) { + compile_bidimensional_cast_only_rtc(kernel_label, itype_name, otype_name, NS, ITN, CVT); + } + if (traits::smem > 0) { + mgr.set_function_attribute(kernel_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + static_cast(traits::smem)); + } + dim3 block(traits::rowThreadLayout::num, traits::numWarps); + dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N, + (rows + traits::blockDIM::M - 1) / traits::blockDIM::M); + mgr.launch(kernel_label, grid, block, static_cast(traits::smem), stream, + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, scales_rowwise, + scales_colwise, noop, rows, cols, scale_stride_rowwise, scale_stride_colwise); +} + +// Compile (on first use) and launch the 32x32 bidimensional cast-only kernel via +// NVRTC, selecting the (numStages, iterN, cvt) config for this shape. TMA +// descriptors are built host-side by the caller (config-independent). +template +inline void launch_bidimensional_cast_only_rtc(const CUtensorMap &tensor_map_input, + const CUtensorMap &tensor_map_rowwise_output, + const CUtensorMap &tensor_map_colwise_output, + e8m0_t *scales_rowwise, e8m0_t *scales_colwise, + const float *noop, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, + int32_t scale_stride_colwise, cudaStream_t stream) { + const std::string itype_name = rtc_type_name(); + const std::string otype_name = rtc_type_name(); + const BidimConfig config = select_bidim_config(rows, cols); + constexpr BidimConfig default_config = static_bidim_config(); + + // The static configuration is always supported and is the only configuration + // selected today. + if (config.num_stages == default_config.num_stages && config.iter_n == default_config.iter_n && + config.use_cvt_4x == default_config.use_cvt_4x) { + launch_bidim_impl( + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, scales_rowwise, + scales_colwise, noop, rows, cols, scale_stride_rowwise, scale_stride_colwise, stream, + itype_name, otype_name); + return; + } + +#define NVTE_MXFP8_BIDIM_CASE(NS, ITN, CVT) \ + if (config.num_stages == (NS) && config.iter_n == (ITN) && config.use_cvt_4x == (CVT)) { \ + launch_bidim_impl( \ + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, scales_rowwise, \ + scales_colwise, noop, rows, cols, scale_stride_rowwise, scale_stride_colwise, stream, \ + itype_name, otype_name); \ + return; \ + } + // Keep candidate configurations available for a future tuning PR. The + // selector above does not currently choose any of them. cvt is fixed at 4x. + NVTE_MXFP8_BIDIM_CASE(2, 1, true) + NVTE_MXFP8_BIDIM_CASE(2, 2, true) + NVTE_MXFP8_BIDIM_CASE(2, 4, true) + NVTE_MXFP8_BIDIM_CASE(2, 8, true) + NVTE_MXFP8_BIDIM_CASE(2, 16, true) + NVTE_MXFP8_BIDIM_CASE(3, 1, true) + NVTE_MXFP8_BIDIM_CASE(3, 2, true) + NVTE_MXFP8_BIDIM_CASE(3, 4, true) + NVTE_MXFP8_BIDIM_CASE(3, 8, true) + NVTE_MXFP8_BIDIM_CASE(3, 16, true) + NVTE_MXFP8_BIDIM_CASE(4, 1, true) + NVTE_MXFP8_BIDIM_CASE(4, 2, true) + NVTE_MXFP8_BIDIM_CASE(4, 4, true) + NVTE_MXFP8_BIDIM_CASE(4, 8, true) + NVTE_MXFP8_BIDIM_CASE(4, 16, true) +#undef NVTE_MXFP8_BIDIM_CASE + + NVTE_ERROR("Unsupported MXFP8 bidimensional RTC config: num_stages=", config.num_stages, + ", iter_n=", config.iter_n, ", use_cvt_4x=", config.use_cvt_4x); +} + +} // namespace specialized +} // namespace quantize_kernel +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // !__CUDACC_RTC__ + +#endif // TRANSFORMER_ENGINE_SPECIALIZED_MXFP8_RTC_DISPATCH_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh b/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh index 5e68b3760c..56f7374a93 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh @@ -11,7 +11,11 @@ #ifndef TRANSFORMER_ENGINE_SPECIALIZED_STATE_COUNTER_CUH_ #define TRANSFORMER_ENGINE_SPECIALIZED_STATE_COUNTER_CUH_ +#if !defined(__CUDACC_RTC__) #include +#else +#include "utils.cuh" +#endif namespace transformer_engine { diff --git a/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh b/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh index dc2d650e7c..d5ecc9e4a6 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh @@ -11,8 +11,12 @@ #ifndef TRANSFORMER_ENGINE_SPECIALIZED_SWIZZLE_CUH_ #define TRANSFORMER_ENGINE_SPECIALIZED_SWIZZLE_CUH_ +#if !defined(__CUDACC_RTC__) #include #include +#else +#include "utils.cuh" +#endif namespace transformer_engine { namespace swz { @@ -33,9 +37,9 @@ using constant = C; template __host__ __device__ __forceinline__ constexpr T shiftr(T x) { - if constexpr (std::is_same_v) { + if constexpr (detail::is_same_v) { return x >> s; - } else if constexpr (std::is_same_v) { + } else if constexpr (detail::is_same_v) { if constexpr (s >= 0) { return x >> s; } else { @@ -44,6 +48,7 @@ __host__ __device__ __forceinline__ constexpr T shiftr(T x) { } } +// avoid use of standard math lib to preserve NVRTC compatibility template struct Swizzle { static constexpr int32_t num_bits = BBits; // number of rows @@ -52,13 +57,12 @@ struct Swizzle { static_assert(num_base >= 0, "MBase must be non-negative"); static_assert(num_bits >= 0, "BBits must be non-negative"); - static_assert(abs(num_shft) >= num_bits, "abs(SShift) must be greater than or equal to num_bits"); + static_assert((num_shft < 0 ? -num_shft : num_shft) >= num_bits, + "abs(SShift) must be greater than or equal to num_bits"); using bit_mask = constant; - using yyy_mask = - constant; - using zzz_mask = - constant; + using yyy_mask = constant 0 ? num_shft : 0))>; + using zzz_mask = constant; using msk_shft = constant; static constexpr int32_t swz_code = int32_t(yyy_mask{} | zzz_mask{}); diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 2814aa3490..9f51920a39 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -11,6 +11,7 @@ #ifndef TRANSFORMER_ENGINE_PTX_CUH_ #define TRANSFORMER_ENGINE_PTX_CUH_ +#if !defined(__CUDACC_RTC__) #include #include @@ -22,6 +23,23 @@ #include #include "common/utils.cuh" +#else +// NVRTC build: common.h drags in host-only headers (cuDNN etc.) and cannot be +// compiled by NVRTC. +#ifndef FP4_TYPE_SUPPORTED +#define FP4_TYPE_SUPPORTED (CUDA_VERSION >= 12080) +#endif +#include +#if FP4_TYPE_SUPPORTED +#include +#endif // FP4_TYPE_SUPPORTED +#include "utils.cuh" + +namespace transformer_engine { +using fp16 = half; +using bf16 = nv_bfloat16; +} // namespace transformer_engine +#endif // __CUDACC_RTC__ namespace transformer_engine { @@ -791,7 +809,7 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_round_to_nearest( uint32_t out_8x = 0; constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; if constexpr (is_blackwell) { - if constexpr (std::is_same::value) { + if constexpr (detail::is_same::value) { asm volatile( "{\n" ".reg.f32 zero; \n\t" @@ -822,7 +840,7 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_round_to_nearest( "}" : "=r"(out_8x) : "l"(in03), "l"(in47), "h"(reinterpret_cast(scaling_coefficient))); - } else if constexpr (std::is_same::value) { + } else if constexpr (detail::is_same::value) { asm volatile( "{\n" ".reg.b64 scaling_coeff_2x; \n\t" @@ -883,7 +901,7 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding( uint32_t out_8x = 0; constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; if constexpr (has_rs) { - if constexpr (std::is_same::value) { + if constexpr (detail::is_same::value) { asm volatile( "{\n" ".reg.f32 zero; \n\t" @@ -913,7 +931,7 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding( : "=r"(out_8x) : "l"(in03), "l"(in47), "h"(reinterpret_cast(scaling_coefficient)), "r"(rbits03), "r"(rbits47)); - } else if constexpr (std::is_same::value) { + } else if constexpr (detail::is_same::value) { asm volatile( "{\n" ".reg.b16 v0_bf16, v1_bf16, v2_bf16, v3_bf16, v4_bf16, v5_bf16, v6_bf16, v7_bf16; \n\t" diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 20616f8cb6..1536cb0eef 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -6,6 +6,7 @@ #include "../util/rtc.h" +#include #include #include #include @@ -26,22 +27,80 @@ namespace { #include "string_code_util_math_h.h" #include "string_code_utils_cuh.h" +/*! \brief Compute capabilities that NVRTC supports */ +const std::vector& supported_sm_archs() { + static const std::vector archs_ = [] { + int num_archs = 0; + NVTE_CHECK_NVRTC(nvrtcGetNumSupportedArchs(&num_archs)); + NVTE_CHECK(num_archs > 0, "Could not determine SM archs that NVRTC supports"); + std::vector archs(num_archs); + NVTE_CHECK_NVRTC(nvrtcGetSupportedArchs(archs.data())); + return archs; + }(); + return archs_; +} + /*! \brief Latest compute capability that NVRTC supports * * \return Compute capability as int. Last digit is minor revision, * remaining digits are major revision. */ inline int max_supported_sm_arch() { - static int arch_ = -1; - if (arch_ < 0) { - int num_archs = 0; - NVTE_CHECK_NVRTC(nvrtcGetNumSupportedArchs(&num_archs)); - NVTE_CHECK(num_archs > 0, "Could not determine SM archs that NVRTC supports"); - std::vector archs(num_archs); - NVTE_CHECK_NVRTC(nvrtcGetSupportedArchs(archs.data())); - arch_ = archs.back(); + const auto& archs = supported_sm_archs(); + return *std::max_element(archs.begin(), archs.end()); +} + +/*! \brief Whether NVRTC supports an exact compute capability */ +inline bool is_supported_sm_arch(int sm_arch) { + const auto& archs = supported_sm_archs(); + return std::find(archs.begin(), archs.end(), sm_arch) != archs.end(); +} + +/*! \brief Run-time NVRTC version encoded like CUDA_VERSION */ +inline int nvrtc_version() { + static const int version_ = [] { + int major = 0; + int minor = 0; + NVTE_CHECK_NVRTC(nvrtcVersion(&major, &minor)); + return major * 1000 + minor * 10; + }(); + return version_; +} + +/*! \brief Latest supported architecture in the device's architecture family */ +inline int max_supported_family_sm_arch(int sm_arch) { + const int family = sm_arch / 10; + int target_arch = -1; + for (const int candidate : supported_sm_archs()) { + if (candidate <= sm_arch && candidate / 10 == family) { + target_arch = std::max(target_arch, candidate); + } + } + return target_arch; +} + +/*! \brief Resolve the static build's Blackwell target policy for NVRTC */ +inline ArchSpecificity resolve_arch_specificity(ArchSpecificity specificity, int sm_arch) { + if (specificity != ArchSpecificity::BlackwellSpecific) { + return specificity; + } + + // Keep this matrix synchronized with NVTE_SPECIFIC_ARCHS in CMakeLists.txt. + // The runtime NVRTC version is the JIT equivalent of CUDAToolkit_VERSION. + switch (sm_arch) { + case 100: + case 101: + case 103: + return ArchSpecificity::ArchitectureSpecific; + case 110: + return ArchSpecificity::FamilySpecific; + default: + if (sm_arch / 10 == 12) { + return nvrtc_version() >= 12090 ? ArchSpecificity::FamilySpecific + : ArchSpecificity::ArchitectureSpecific; + } + NVTE_ERROR("No Blackwell RTC target policy is defined for sm_", sm_arch); } - return arch_; } } // namespace @@ -148,7 +207,8 @@ KernelManager& KernelManager::instance() { void KernelManager::compile(const std::string& kernel_label, const std::string& kernel_name, const std::string& code, const std::string& filename, const std::vector& extra_options, - const std::vector
& extra_headers) { + const std::vector
& extra_headers, + ArchRequirement arch_requirement) { const int device_id = cuda::current_device(); const auto key = get_kernel_cache_key(kernel_label, device_id); std::unique_lock lock_guard_(lock_); @@ -158,8 +218,39 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& // Choose whether to compile to PTX or cubin const int sm_arch_ = cuda::sm_arch(device_id); - const int compile_sm_arch = std::min(sm_arch_, max_supported_sm_arch()); - const bool compile_ptx = sm_arch_ != compile_sm_arch; + NVTE_CHECK(sm_arch_ >= arch_requirement.min_sm_arch, "RTC kernel ", kernel_label, " requires sm_", + arch_requirement.min_sm_arch, " or newer, but the current device is sm_", sm_arch_); + + int compile_sm_arch = std::min(sm_arch_, max_supported_sm_arch()); + bool compile_ptx = sm_arch_ != compile_sm_arch; + const char* arch_suffix = ""; + const ArchSpecificity arch_specificity = + resolve_arch_specificity(arch_requirement.specificity, sm_arch_); + if (arch_specificity == ArchSpecificity::ArchitectureSpecific) { + NVTE_CHECK( + is_supported_sm_arch(sm_arch_), "RTC kernel ", kernel_label, + " requires an architecture-specific target for sm_", sm_arch_, + ", but the runtime NVRTC does not support that architecture. Use a newer CUDA toolkit, " + "or disable NVRTC and rebuild with the corresponding legacy static kernel enabled."); + NVTE_CHECK(sm_arch_ >= 90, "Architecture-specific RTC targets require sm_90 or newer"); + compile_sm_arch = sm_arch_; + compile_ptx = false; + arch_suffix = "a"; + } else if (arch_specificity == ArchSpecificity::FamilySpecific) { + const int rtc_version = nvrtc_version(); + NVTE_CHECK(rtc_version >= 12090, "RTC kernel ", kernel_label, + " requires a family-specific target, but NVRTC ", rtc_version / 1000, ".", + (rtc_version % 1000) / 10, " does not support f-suffixed architectures"); + compile_sm_arch = max_supported_family_sm_arch(sm_arch_); + NVTE_CHECK( + compile_sm_arch >= arch_requirement.min_sm_arch, "RTC kernel ", kernel_label, + " requires a family-specific target compatible with sm_", sm_arch_, + ", but the runtime NVRTC does not support that architecture family. Use a newer CUDA " + "toolkit, or disable NVRTC and rebuild with the corresponding legacy static kernel " + "enabled."); + compile_ptx = compile_sm_arch != sm_arch_; + arch_suffix = "f"; + } // Compilation flags std::vector opts = { @@ -168,9 +259,9 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& #endif "--std=c++17"}; if (compile_ptx) { - opts.push_back(concat_strings("--gpu-architecture=compute_", compile_sm_arch)); + opts.push_back(concat_strings("--gpu-architecture=compute_", compile_sm_arch, arch_suffix)); } else { - opts.push_back(concat_strings("--gpu-architecture=sm_", compile_sm_arch)); + opts.push_back(concat_strings("--gpu-architecture=sm_", compile_sm_arch, arch_suffix)); } opts.push_back(concat_strings("-I", cuda::include_directory(true))); opts.insert(opts.end(), extra_options.begin(), extra_options.end()); diff --git a/transformer_engine/common/util/rtc.h b/transformer_engine/common/util/rtc.h index 745566ee60..6c2bc92787 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -40,6 +40,26 @@ struct Header { const char *include_name; }; +/*! \brief Architecture specificity required by an RTC kernel */ +enum class ArchSpecificity { + /*! Compile for a generic architecture, preserving PTX forward compatibility */ + Generic, + /*! Compile for the exact architecture-specific target (e.g. sm_100a) */ + ArchitectureSpecific, + /*! Compile for a family-specific target (e.g. sm_120f) */ + FamilySpecific, + /*! Match the static build's Blackwell a/f target policy */ + BlackwellSpecific, +}; + +/*! \brief Architecture requirements for an RTC kernel */ +struct ArchRequirement { + /*! Minimum compute capability, encoded as major * 10 + minor */ + int min_sm_arch = 0; + /*! Target-specificity policy for the kernel */ + ArchSpecificity specificity = ArchSpecificity::Generic; +}; + /*! \brief Wrapper class for a runtime-compiled CUDA kernel */ class Kernel { public: @@ -145,11 +165,14 @@ class KernelManager { * primarily for debugging * \param[in] extra_options Additional NVRTC compiler options * \param[in] extra_headers Additional in-memory headers available to the program + * \param[in] arch_requirement Minimum and specificity requirements for the + * compilation target */ void compile(const std::string &kernel_label, const std::string &kernel_name, const std::string &code, const std::string &filename, const std::vector &extra_options = {}, - const std::vector
&extra_headers = {}); + const std::vector
&extra_headers = {}, + ArchRequirement arch_requirement = {}); /*! \brief Whether CUDA kernel has been compiled for CUDA device * diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 0b75da622c..9493dd8bb9 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -24,10 +24,21 @@ using uint8_t = unsigned char; using uint16_t = unsigned short int; // NOLINT(*) using uint32_t = unsigned int; using uint64_t = unsigned long long int; // NOLINT(*) +using int8_t = signed char; +using int16_t = short int; // NOLINT(*) +using int32_t = int; // NOLINT(*) +using int64_t = long long int; // NOLINT(*) +using intptr_t = long long int; // NOLINT(*) +using uintptr_t = unsigned long long int; // NOLINT(*) static_assert(sizeof(uint8_t) == 1); static_assert(sizeof(uint16_t) == 2); static_assert(sizeof(uint32_t) == 4); static_assert(sizeof(uint64_t) == 8); +static_assert(sizeof(int8_t) == 1); +static_assert(sizeof(int16_t) == 2); +static_assert(sizeof(int32_t) == 4); +static_assert(sizeof(int64_t) == 8); +static_assert(sizeof(intptr_t) == sizeof(void *)); #endif // Minimal subset of used by RTC kernel headers. Keep these in a @@ -58,6 +69,24 @@ struct conditional { template using conditional_t = typename conditional::type; +template +struct enable_if {}; +template +struct enable_if { + using type = T; +}; +template +using enable_if_t = typename enable_if::type; + +template +__host__ __device__ constexpr T min(const T &a, const T &b) { + return b < a ? b : a; +} +template +__host__ __device__ constexpr T max(const T &a, const T &b) { + return a < b ? b : a; +} + } // namespace detail } // namespace transformer_engine