diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 13880e0d61..c811307d75 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -181,6 +181,7 @@ list(APPEND transformer_engine_cpp_sources cudnn_utils.cpp transformer_engine.cpp fused_attn/fused_attn.cpp + cast/nvfp4/rtc_dispatch.cpp gemm/config.cpp normalization/common.cpp normalization/rtc_dispatch.cpp @@ -565,6 +566,20 @@ make_string_header_from_file(utils.cuh make_string_header_from_file(util/math.h string_code_util_math_h) +# NVFP4 4over6 quantize: NVRTC bundled headers + RTC source file +make_string_header_from_file(include/transformer_engine/nvfp4/4over6.h + string_code_transformer_engine_nvfp4_4over6_h) +make_string_header_from_file(util/type_extrema.h + string_code_util_type_extrema_h) +make_string_header_from_file(util/ptx.cuh + string_code_util_ptx_cuh) +make_string_header_from_file(cast/nvfp4/core_nvfp4.cuh + string_code_cast_nvfp4_core_nvfp4_cuh) +make_string_header_from_file(cast/nvfp4/quantize_4over6_kernel.cuh + string_code_cast_nvfp4_quantize_4over6_kernel_cuh) +make_string_header_from_file(cast/nvfp4/rtc/quantize_4over6.cu + string_code_cast_nvfp4_rtc_quantize_4over6_cu) + # Norm NVRTC bundled headers + RTC source files make_string_header_from_file(normalization/kernel_params.h string_code_normalization_kernel_params_h) @@ -589,6 +604,16 @@ make_string_header_from_file(normalization/rmsnorm/rtc/rmsnorm_bwd_kernel.cu target_include_directories(transformer_engine PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/string_headers") +# Default OFF: the NVFP4 4over6 quantize kernel (384-way fanout) is compiled via +# NVRTC at runtime. Set ON to additionally statically instantiate it as the +# fallback selected when NVTE_DISABLE_NVRTC=1. +option(NVTE_BUILD_LEGACY_STATIC_NVFP4 + "Also compile the static NVFP4 4over6 quantize kernel for NVTE_DISABLE_NVRTC fallback" + OFF) +target_compile_definitions(transformer_engine + PRIVATE + NVTE_BUILD_LEGACY_STATIC_NVFP4=$) + option(NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX "Also compile legacy static fused softmax kernels for NVTE_DISABLE_NVRTC fallback" OFF) diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index 3820430d5b..f0e40a6137 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -11,6 +11,7 @@ #ifndef TRANSFORMER_ENGINE_CORE_NVFP4_CUH_ #define TRANSFORMER_ENGINE_CORE_NVFP4_CUH_ +#if !defined(__CUDACC_RTC__) #include #include #include @@ -22,11 +23,28 @@ #include "../../util/math.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" +#else +// NVRTC build: common.h (host-only: cuDNN/cutlass) cannot be parsed by NVRTC. +// utils.cuh, util/math.h and ptx.cuh are injected as in-memory headers by the +// RTC dispatch and already provide the integer typedefs, detail::is_same, and +// the fp8 element types. util/type_extrema.h (also injected) provides the +// transformer_engine-namespace fp4 aliases and detail::TypeExtrema +// specializations that would otherwise come from common.h. +#include "ptx.cuh" +#include "util/math.h" +#include "utils.cuh" +#endif // __CUDACC_RTC__ #if FP4_TYPE_SUPPORTED #include #endif // FP4_TYPE_SUPPORTED +#if defined(__CUDACC_RTC__) +namespace transformer_engine { +#include "util/type_extrema.h" +} // namespace transformer_engine +#endif // __CUDACC_RTC__ + namespace transformer_engine { namespace dispatch { namespace nvfp4 { @@ -94,6 +112,8 @@ __device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const return global_encode_scale; } +#if !defined(__CUDACC_RTC__) +// The RTC 4over6 path rejects stochastic rounding; static transpose kernels use this helper. __device__ __forceinline__ uint32_t get_rbits( transformer_engine::curanddx::detail::philox4x32_native_state &rng, @@ -108,6 +128,7 @@ __device__ __forceinline__ uint32_t get_rbits( const uint32_t rbits = rbits_arr[rnd_idx++]; return rbits; } +#endif // !__CUDACC_RTC__ #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_kernel.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_kernel.cuh new file mode 100644 index 0000000000..770040f6eb --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_kernel.cuh @@ -0,0 +1,620 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_4over6_kernel.cuh + * \brief Dedicated kernels for NVFP4 4over6 quantization. + * + * Four Over Six evaluates two TE-style NVFP4 encodings for every 1x16 + * quantization group. The map-to-6 candidate uses the normal scale. The + * map-to-4 candidate expands the E4M3 block scale by 1.5x so FP4 value 4 + * reaches the same range that FP4 value 6 reaches in the normal encoding. + * The selected candidate is the one with lower configured dequantization + * error; ties select map-to-6. The quantized candidates, dequantized values, + * and errors are kept in registers, matching the structure of the official + * Four Over Six implementation. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_4OVER6_KERNEL_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_4OVER6_KERNEL_CUH_ + +#if !defined(__CUDACC_RTC__) +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../utils.cuh" +#else +// NVRTC build: the host-only headers above are unavailable. +#include +#endif // __CUDACC_RTC__ + +#include "core_nvfp4.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { + +#if FP4_TYPE_SUPPORTED + +namespace quantize_4over6_kernel { + +constexpr int kThreads = 128; +constexpr int kWarpThreads = 32; +constexpr int kGroupSize = 16; +constexpr int kTileRows = 128; +constexpr int kTileCols = 64; +constexpr int kTileColGroups = kTileCols / kGroupSize; +constexpr int kTileRowGroups = kTileRows / kGroupSize; +constexpr int kPipelineStages = 2; +constexpr int kStageRows = kTileRows / kPipelineStages; +constexpr int kStageRowGroups = kStageRows / kGroupSize; +constexpr int kElementsPerHalfGroup = 8; +constexpr int kPackedWordsPerGroup = 2; +static_assert(kTileRows == kPipelineStages * kStageRows); +static_assert(kStageRows % kGroupSize == 0); + +template +struct Config { + static constexpr NVTENVFP44Over6Mode mode = kMode; + static constexpr bool err_use_fast_math = kErrUseFastMath; +}; + +struct Candidate { + uint32_t packed[kPackedWordsPerGroup]; + float err; +}; + +struct CandidatePair { + Candidate map4; + Candidate map6; +}; + +struct ScalePair { + nvfp4_scale_t map4; + nvfp4_scale_t map6; + float inv_map4; + float inv_map6; + float global_encode_scale; +}; + +struct FP16ErrorScalePair { + uint32_t map4; + uint32_t map6; +}; + +template +__device__ __forceinline__ float compute_error_rn(const float diff) { + if constexpr (kMode == kNVTENVFP44Over6MinMSE) { + return __fmul_rn(diff, diff); + } else if constexpr (kMode == kNVTENVFP44Over6MinMAE) { + return fabsf(diff); + } else { + NVTE_DEVICE_ERROR("Unsupported NVFP4 4over6 mode."); + return fabsf(diff); + } +} + +template +__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, + const float global_amax) { + static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); + constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f + constexpr float fp8_max = detail::TypeExtrema::max; // 448.0f + constexpr float expand_to_map4 = 1.5f; + const float S_enc = core::compute_global_encode_scaling_factor_FP4(global_amax); + const float base = block_amax / fp4_max * S_enc; + + ScalePair scales; + scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); + scales.map6 = static_cast(fminf(base, fp8_max)); + + const float S_dec = 1.0f / S_enc; + scales.inv_map4 = + fminf(1.0f / (static_cast(scales.map4) * S_dec), detail::TypeExtrema::max); + scales.inv_map6 = + fminf(1.0f / (static_cast(scales.map6) * S_dec), detail::TypeExtrema::max); + scales.global_encode_scale = S_enc; + return scales; +} + +template +__device__ __forceinline__ float load_input(const IType *ptr, const size_t idx) { + return static_cast(ptr[idx]); +} + +template +__device__ __forceinline__ void load_row_group(const IType *tile, const int row, + const int col_start, float (&x0)[8], float (&x1)[8], + float *amax) { + Vec x0_vec; + Vec x1_vec; + x0_vec.load_from(&tile[row * kTileCols + col_start]); + x1_vec.load_from(&tile[row * kTileCols + col_start + kElementsPerHalfGroup]); + + *amax = 0.0f; +#pragma unroll + for (int i = 0; i < kElementsPerHalfGroup; ++i) { + const float v0 = static_cast(x0_vec.data.elt[i]); + const float v1 = static_cast(x1_vec.data.elt[i]); + x0[i] = v0; + x1[i] = v1; + *amax = fmaxf(*amax, fabsf(v0)); + *amax = fmaxf(*amax, fabsf(v1)); + } +} + +template +__device__ __forceinline__ void load_col_group(const IType *tile, const int row_start, + const int col, float (&x0)[8], float (&x1)[8], + float *amax) { + *amax = 0.0f; +#pragma unroll + for (int i = 0; i < kElementsPerHalfGroup; ++i) { + const float v0 = load_input(tile, (row_start + i) * kTileCols + col); + const float v1 = load_input(tile, (row_start + i + kElementsPerHalfGroup) * kTileCols + col); + x0[i] = v0; + x1[i] = v1; + *amax = fmaxf(*amax, fabsf(v0)); + *amax = fmaxf(*amax, fabsf(v1)); + } +} + +template +__device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_bits, const float x, + const float sf, const float global_amax, + float *err) { + constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f + constexpr float fp8_max = static_cast(E4M3_MAX); + constexpr float err_denom = fp4_max * fp8_max; + const uint16_t half_bits = (dequant_bits >> SHIFT) & 0xFFFF; + const float dequant = __half2float(__ushort_as_half(half_bits)); + const float val = __fdiv_rn(__fmul_rn(__fmul_rn(dequant, sf), global_amax), err_denom); + const float diff = __fsub_rn(val, x); + *err = __fadd_rn(*err, compute_error_rn(diff)); +} + +__device__ __forceinline__ uint8_t fp8_bits(const nvfp4_scale_t sf) { + return *reinterpret_cast(&sf); +} + +__device__ __forceinline__ FP16ErrorScalePair compute_fp16_error_scales(const ScalePair &scales) { + FP16ErrorScalePair result; + const uint32_t packed_scales = static_cast(fp8_bits(scales.map4)) | + (static_cast(fp8_bits(scales.map6)) << 8); + asm volatile( + "{\n" + ".reg .b16 fp8_pair;\n" + ".reg .b16 map4_h, map6_h;\n" + ".reg .b32 scale_h2;\n" + "cvt.u16.u32 fp8_pair, %2;\n" + "cvt.rn.f16x2.e4m3x2 scale_h2, fp8_pair;\n" + "mov.b32 {map4_h, map6_h}, scale_h2;\n" + "mov.b32 %0, {map4_h, map4_h};\n" + "mov.b32 %1, {map6_h, map6_h};\n" + "}" + : "=r"(result.map4), "=r"(result.map6) + : "r"(packed_scales)); + return result; +} + +__device__ __forceinline__ float2 f16x2_scaled_to_float2(const uint32_t q_h2, + const uint32_t scale_h2) { + float2 result; + asm volatile( + "{\n" + ".reg .b16 lo, hi;\n" + ".reg .b32 prod_h2;\n" + "mul.rn.f16x2 prod_h2, %2, %3;\n" + "mov.b32 {lo, hi}, prod_h2;\n" + "cvt.f32.f16 %0, lo;\n" + "cvt.f32.f16 %1, hi;\n" + "}" + : "=f"(result.x), "=f"(result.y) + : "r"(q_h2), "r"(scale_h2)); + return result; +} + +template +__device__ __forceinline__ void accumulate_fp16_scaled_error_pair(const uint32_t q_h2, + const float x0, const float x1, + const uint32_t scale_h2, + const float global_encode_scale, + float *err) { + const float2 candidate = f16x2_scaled_to_float2(q_h2, scale_h2); + const float original0 = __fmul_rn(x0, global_encode_scale); + const float original1 = __fmul_rn(x1, global_encode_scale); + const float diff0 = __fsub_rn(candidate.x, original0); + const float diff1 = __fsub_rn(candidate.y, original1); + *err = __fadd_rn(*err, compute_error_rn(diff0)); + *err = __fadd_rn(*err, compute_error_rn(diff1)); +} + +template +__device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( + const float (&x)[8], const float block_scale_inverse, const nvfp4_scale_t sf, + const uint32_t fp16_error_scale, const float global_amax, const float global_encode_scale, + float *err) { + uint32_t out = 0; + uint32_t out_dequant_1 = 0; + uint32_t out_dequant_2 = 0; + uint32_t out_dequant_3 = 0; + uint32_t out_dequant_4 = 0; + + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + asm volatile( + "{\n" + ".reg .b8 byte0, byte1, byte2, byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %8, %7;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %10, %9;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %12, %11;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "cvt.rn.f16x2.e2m1x2 %1, byte0;\n" + "cvt.rn.f16x2.e2m1x2 %2, byte1;\n" + "cvt.rn.f16x2.e2m1x2 %3, byte2;\n" + "cvt.rn.f16x2.e2m1x2 %4, byte3;\n" + "}" + : "=r"(out), "=r"(out_dequant_1), "=r"(out_dequant_2), "=r"(out_dequant_3), + "=r"(out_dequant_4) + : "f"(__fmul_rn(x[0], block_scale_inverse)), "f"(__fmul_rn(x[1], block_scale_inverse)), + "f"(__fmul_rn(x[2], block_scale_inverse)), "f"(__fmul_rn(x[3], block_scale_inverse)), + "f"(__fmul_rn(x[4], block_scale_inverse)), "f"(__fmul_rn(x[5], block_scale_inverse)), + "f"(__fmul_rn(x[6], block_scale_inverse)), "f"(__fmul_rn(x[7], block_scale_inverse))); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + + if constexpr (Cfg::err_use_fast_math) { + accumulate_fp16_scaled_error_pair(out_dequant_1, x[0], x[1], fp16_error_scale, + global_encode_scale, err); + accumulate_fp16_scaled_error_pair(out_dequant_2, x[2], x[3], fp16_error_scale, + global_encode_scale, err); + accumulate_fp16_scaled_error_pair(out_dequant_3, x[4], x[5], fp16_error_scale, + global_encode_scale, err); + accumulate_fp16_scaled_error_pair(out_dequant_4, x[6], x[7], fp16_error_scale, + global_encode_scale, err); + } else { + const float sf_float = static_cast(sf); + accumulate_dequant_error(out_dequant_1, x[0], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_1, x[1], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_2, x[2], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_2, x[3], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_3, x[4], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_3, x[5], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_4, x[6], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_4, x[7], sf_float, global_amax, err); + } + return out; +} + +template +__device__ __forceinline__ CandidatePair make_candidates(const float (&x0)[8], const float (&x1)[8], + const ScalePair &scales, + const float global_amax) { + CandidatePair candidates; + candidates.map4.err = 0.0f; + candidates.map6.err = 0.0f; + FP16ErrorScalePair fp16_error_scales{}; + if constexpr (Cfg::err_use_fast_math) { + fp16_error_scales = compute_fp16_error_scales(scales); + } + candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( + x0, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, + scales.global_encode_scale, &candidates.map4.err); + candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( + x0, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, + scales.global_encode_scale, &candidates.map6.err); + candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( + x1, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, + scales.global_encode_scale, &candidates.map4.err); + candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( + x1, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, + scales.global_encode_scale, &candidates.map6.err); + return candidates; +} + +__device__ __forceinline__ float reduce_group_sum_16(float value) { + const int lane = threadIdx.x & (kWarpThreads - 1); + const int group_base = lane & ~(kGroupSize - 1); + const unsigned mask = 0xffffu << group_base; +#pragma unroll + for (int offset = kGroupSize / 2; offset > 0; offset /= 2) { + value += __shfl_down_sync(mask, value, offset, kGroupSize); + } + return __shfl_sync(mask, value, group_base, kWarpThreads); +} + +__device__ __forceinline__ float reduce_group_max_16(float value) { + const int lane = threadIdx.x & (kWarpThreads - 1); + const int group_base = lane & ~(kGroupSize - 1); + const unsigned mask = 0xffffu << group_base; +#pragma unroll + for (int offset = kGroupSize / 2; offset > 0; offset /= 2) { + value = fmaxf(value, __shfl_down_sync(mask, value, offset, kGroupSize)); + } + return __shfl_sync(mask, value, group_base, kWarpThreads); +} + +__device__ __forceinline__ void store_packed_group(const uint32_t *packed, fp4e2m1x2 *dst) { + const uint64_t packed64 = + static_cast(packed[0]) | (static_cast(packed[1]) << 32); + *reinterpret_cast(dst) = packed64; +} + +__device__ __forceinline__ const uint32_t *select_packed(const CandidatePair &candidates, + const bool pick_map4) { + if (pick_map4) { + return candidates.map4.packed; + } + return candidates.map6.packed; +} + +__device__ __forceinline__ nvfp4_scale_t select_scale(const ScalePair &scales, + const bool pick_map4) { + if (pick_map4) { + return scales.map4; + } + return scales.map6; +} + +__device__ __forceinline__ void cp_async_cg_16(void *dst, const void *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + const uint32_t dst_smem_ptr = __cvta_generic_to_shared(dst); + const uint64_t src_gmem_ptr = reinterpret_cast(src); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst_smem_ptr), + "l"(src_gmem_ptr)); +#else + NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); +#endif +} + +__device__ __forceinline__ void cp_async_commit_group() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + asm volatile("cp.async.commit_group;\n" ::); +#else + NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); +#endif +} + +template +__device__ __forceinline__ void cp_async_wait_group() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); +#else + NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); +#endif +} + +template +__device__ void load_stage_to_shared_async(const IType *input, IType *tile, const size_t rows, + const size_t cols, const size_t stage_row, + const size_t tile_col) { + constexpr int vec_elems = 16 / sizeof(IType); + constexpr int vecs_per_row = kTileCols / vec_elems; + constexpr int vecs = kStageRows * vecs_per_row; + using TileVec = Vec; + + for (int idx = threadIdx.x; idx < vecs; idx += blockDim.x) { + const int local_row = idx / vecs_per_row; + const int local_vec_col = idx - local_row * vecs_per_row; + const int local_col = local_vec_col * vec_elems; + const size_t global_row = stage_row + local_row; + const size_t global_col = tile_col + local_col; + IType *stage_ptr = &tile[local_row * kTileCols + local_col]; + + if (global_row < rows && global_col + vec_elems <= cols) { + cp_async_cg_16(stage_ptr, &input[global_row * cols + global_col]); + } else { + TileVec vec; + vec.clear(); +#pragma unroll + for (int i = 0; i < vec_elems; ++i) { + if (global_row < rows && global_col + i < cols) { + vec.data.elt[i] = input[global_row * cols + global_col + i]; + } + } + vec.store_to(stage_ptr); + } + } +} + +template +__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvfp4_scale_t *scales, + const float *amax, const size_t rows, const size_t cols, + const size_t stage_row, const size_t tile_col, + const size_t scale_stride) { + constexpr int groups = kStageRows * kTileColGroups; + for (int group = threadIdx.x; group < groups; group += blockDim.x) { + const int local_row = group % kStageRows; + const int local_col_group = group / kStageRows; + const int local_col = local_col_group * kGroupSize; + const size_t global_row = stage_row + local_row; + const size_t global_col = tile_col + local_col; + if (global_row >= rows || global_col >= cols) { + continue; + } + + float x0[8]; + float x1[8]; + float group_amax = 0.0f; + load_row_group(tile, local_row, local_col, x0, x1, &group_amax); + + float block_amax = group_amax; + if constexpr (USE_2D_QUANTIZATION) { + block_amax = reduce_group_max_16(group_amax); + } + + float global_amax = amax[0]; + if constexpr (ROW_SCALED_NVFP4) { + global_amax = amax[global_row]; + } + + const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + + float err_map4 = candidates.map4.err; + float err_map6 = candidates.map6.err; + if constexpr (USE_2D_QUANTIZATION) { + err_map4 = reduce_group_sum_16(err_map4); + err_map6 = reduce_group_sum_16(err_map6); + } + + const bool pick_map4 = err_map4 < err_map6; + const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const uint32_t *selected = select_packed(candidates, pick_map4); + + const size_t global_col_group = global_col / kGroupSize; + scales[global_row * scale_stride + global_col_group] = selected_scale; + store_packed_group(selected, &output[(global_row * cols + global_col) / 2]); + } +} + +template +__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, + nvfp4_scale_t *scales_t, const float *amax, + const size_t rows, const size_t cols, const size_t stage_row, + const size_t tile_col, const size_t scale_stride_t) { + constexpr int groups = kStageRowGroups * kTileCols; + for (int group = threadIdx.x; group < groups; group += blockDim.x) { + const int local_row_group = group / kTileCols; + const int local_col = group - local_row_group * kTileCols; + const int local_row = local_row_group * kGroupSize; + const size_t global_row = stage_row + local_row; + const size_t global_col = tile_col + local_col; + if (global_row >= rows || global_col >= cols) { + continue; + } + + float x0[8]; + float x1[8]; + float group_amax = 0.0f; + load_col_group(tile, local_row, local_col, x0, x1, &group_amax); + + float block_amax = group_amax; + if constexpr (USE_2D_QUANTIZATION) { + block_amax = reduce_group_max_16(group_amax); + } + + const float global_amax = amax[0]; + const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + + float err_map4 = candidates.map4.err; + float err_map6 = candidates.map6.err; + if constexpr (USE_2D_QUANTIZATION) { + err_map4 = reduce_group_sum_16(err_map4); + err_map6 = reduce_group_sum_16(err_map6); + } + + const bool pick_map4 = err_map4 < err_map6; + const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const uint32_t *selected = select_packed(candidates, pick_map4); + + const size_t global_row_group = global_row / kGroupSize; + scales_t[global_col * scale_stride_t + global_row_group] = selected_scale; + store_packed_group(selected, &output_t[(global_col * rows + global_row) / 2]); + } +} + +// Shared kernel body reused by both the statically-instantiated __global__ +// (quantize_4over6_kernel below) and the NVRTC entry point +// (rtc/quantize_4over6.cu). +template +__device__ __forceinline__ void quantize_4over6_body( + const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, nvfp4_scale_t *scales, + nvfp4_scale_t *scales_t, const float *amax_rowwise, const float *amax_colwise, + const size_t rows, const size_t cols, const size_t scale_stride, const size_t scale_stride_t, + const float *noop) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + extern __shared__ char dynamic_shmem[]; + auto *tiles = reinterpret_cast(dynamic_shmem); + const size_t tile_col = blockIdx.x * kTileCols; + const size_t tile_row = blockIdx.y * kTileRows; + + IType *stage_tiles[kPipelineStages]; +#pragma unroll + for (int stage = 0; stage < kPipelineStages; ++stage) { + stage_tiles[stage] = &tiles[stage * kStageRows * kTileCols]; + } + + load_stage_to_shared_async(input, stage_tiles[0], rows, cols, tile_row, tile_col); + cp_async_commit_group(); + cp_async_wait_group<0>(); + __syncthreads(); + + for (int stage = 0; stage < kPipelineStages; ++stage) { + const int next_stage = stage + 1; + if (next_stage < kPipelineStages) { + const size_t next_stage_row = tile_row + next_stage * kStageRows; + load_stage_to_shared_async(input, stage_tiles[next_stage], rows, cols, next_stage_row, + tile_col); + cp_async_commit_group(); + } + + const size_t stage_row = tile_row + stage * kStageRows; + IType *stage_tile = stage_tiles[stage]; + + if constexpr (RETURN_IDENTITY) { + quantize_stage_rowwise( + stage_tile, output, scales, amax_rowwise, rows, cols, stage_row, tile_col, scale_stride); + } + + if constexpr (RETURN_TRANSPOSE) { + const float *columnwise_amax = amax_colwise; + if (columnwise_amax == nullptr) { + columnwise_amax = amax_rowwise; + } + quantize_stage_colwise( + stage_tile, output_t, scales_t, columnwise_amax, rows, cols, stage_row, tile_col, + scale_stride_t); + } + + if (next_stage < kPipelineStages) { + cp_async_wait_group<0>(); + __syncthreads(); + } + } +#else + NVTE_DEVICE_ERROR("sm_100 or higher is required."); +#endif +} + +// Statically-instantiated entry point (NVTE_BUILD_LEGACY_STATIC_NVFP4 fallback). +template +__global__ void __launch_bounds__(kThreads) + quantize_4over6_kernel(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, + nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, + const float *amax_rowwise, const float *amax_colwise, const size_t rows, + const size_t cols, const size_t scale_stride, + const size_t scale_stride_t, const float *noop) { + quantize_4over6_body(input, output, output_t, scales, scales_t, + amax_rowwise, amax_colwise, rows, cols, scale_stride, + scale_stride_t, noop); +} + +} // namespace quantize_4over6_kernel + +#endif // FP4_TYPE_SUPPORTED + +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_4OVER6_KERNEL_CUH_ diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 50776a3ed6..62af1df63a 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -5,16 +5,7 @@ ************************************************************************/ /*! \file quantize_4over6_nvfp4.cuh - * \brief Dedicated kernels for NVFP4 4over6 quantization. - * - * Four Over Six evaluates two TE-style NVFP4 encodings for every 1x16 - * quantization group. The map-to-6 candidate uses the normal scale. The - * map-to-4 candidate expands the E4M3 block scale by 1.5x so FP4 value 4 - * reaches the same range that FP4 value 6 reaches in the normal encoding. - * The selected candidate is the one with lower configured dequantization - * error; ties select map-to-6. The quantized candidates, dequantized values, - * and errors are kept in registers, matching the structure of the official - * Four Over Six implementation. + * \brief Host dispatch for NVFP4 4over6 quantization. */ #ifndef TRANSFORMER_ENGINE_QUANTIZE_4OVER6_NVFP4_CUH_ @@ -22,16 +13,15 @@ #include #include -#include #include #include #include #include "../../common.h" -#include "../../util/math.h" -#include "../../utils.cuh" -#include "core_nvfp4.cuh" +#include "../../util/rtc.h" +#include "quantize_4over6_kernel.cuh" +#include "rtc_dispatch.h" namespace transformer_engine { namespace dispatch { @@ -66,554 +56,6 @@ namespace nvfp4 { namespace quantize_4over6_kernel { -constexpr int kThreads = 128; -constexpr int kWarpThreads = 32; -constexpr int kGroupSize = 16; -constexpr int kTileRows = 128; -constexpr int kTileCols = 64; -constexpr int kTileColGroups = kTileCols / kGroupSize; -constexpr int kTileRowGroups = kTileRows / kGroupSize; -constexpr int kPipelineStages = 2; -constexpr int kStageRows = kTileRows / kPipelineStages; -constexpr int kStageRowGroups = kStageRows / kGroupSize; -constexpr int kElementsPerHalfGroup = 8; -constexpr int kPackedWordsPerGroup = 2; -static_assert(kTileRows == kPipelineStages * kStageRows); -static_assert(kStageRows % kGroupSize == 0); - -template -struct Config { - static constexpr NVTENVFP44Over6Mode mode = kMode; - static constexpr bool err_use_fast_math = kErrUseFastMath; -}; - -struct Candidate { - uint32_t packed[kPackedWordsPerGroup]; - float err; -}; - -struct CandidatePair { - Candidate map4; - Candidate map6; -}; - -struct ScalePair { - nvfp4_scale_t map4; - nvfp4_scale_t map6; - float inv_map4; - float inv_map6; - float global_encode_scale; -}; - -struct FP16ErrorScalePair { - uint32_t map4; - uint32_t map6; -}; - -template -__device__ __forceinline__ float compute_error_rn(const float diff) { - if constexpr (kMode == kNVTENVFP44Over6MinMSE) { - return __fmul_rn(diff, diff); - } else if constexpr (kMode == kNVTENVFP44Over6MinMAE) { - return fabsf(diff); - } else { - NVTE_DEVICE_ERROR("Unsupported NVFP4 4over6 mode."); - return fabsf(diff); - } -} - -template -__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, - const float global_amax) { - static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); - constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = detail::TypeExtrema::max; // 448.0f - constexpr float expand_to_map4 = 1.5f; - const float S_enc = core::compute_global_encode_scaling_factor_FP4(global_amax); - const float base = block_amax / fp4_max * S_enc; - - ScalePair scales; - scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); - scales.map6 = static_cast(fminf(base, fp8_max)); - - const float S_dec = 1.0f / S_enc; - scales.inv_map4 = - fminf(1.0f / (static_cast(scales.map4) * S_dec), detail::TypeExtrema::max); - scales.inv_map6 = - fminf(1.0f / (static_cast(scales.map6) * S_dec), detail::TypeExtrema::max); - scales.global_encode_scale = S_enc; - return scales; -} - -template -__device__ __forceinline__ float load_input(const IType *ptr, const size_t idx) { - return static_cast(ptr[idx]); -} - -template -__device__ __forceinline__ void load_row_group(const IType *tile, const int row, - const int col_start, float (&x0)[8], float (&x1)[8], - float *amax) { - Vec x0_vec; - Vec x1_vec; - x0_vec.load_from(&tile[row * kTileCols + col_start]); - x1_vec.load_from(&tile[row * kTileCols + col_start + kElementsPerHalfGroup]); - - *amax = 0.0f; -#pragma unroll - for (int i = 0; i < kElementsPerHalfGroup; ++i) { - const float v0 = static_cast(x0_vec.data.elt[i]); - const float v1 = static_cast(x1_vec.data.elt[i]); - x0[i] = v0; - x1[i] = v1; - *amax = fmaxf(*amax, fabsf(v0)); - *amax = fmaxf(*amax, fabsf(v1)); - } -} - -template -__device__ __forceinline__ void load_col_group(const IType *tile, const int row_start, - const int col, float (&x0)[8], float (&x1)[8], - float *amax) { - *amax = 0.0f; -#pragma unroll - for (int i = 0; i < kElementsPerHalfGroup; ++i) { - const float v0 = load_input(tile, (row_start + i) * kTileCols + col); - const float v1 = load_input(tile, (row_start + i + kElementsPerHalfGroup) * kTileCols + col); - x0[i] = v0; - x1[i] = v1; - *amax = fmaxf(*amax, fabsf(v0)); - *amax = fmaxf(*amax, fabsf(v1)); - } -} - -template -__device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_bits, const float x, - const float sf, const float global_amax, - float *err) { - constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = static_cast(E4M3_MAX); - constexpr float err_denom = fp4_max * fp8_max; - const uint16_t half_bits = (dequant_bits >> SHIFT) & 0xFFFF; - const float dequant = __half2float(__ushort_as_half(half_bits)); - const float val = __fdiv_rn(__fmul_rn(__fmul_rn(dequant, sf), global_amax), err_denom); - const float diff = __fsub_rn(val, x); - *err = __fadd_rn(*err, compute_error_rn(diff)); -} - -__device__ __forceinline__ uint8_t fp8_bits(const nvfp4_scale_t sf) { - return *reinterpret_cast(&sf); -} - -__device__ __forceinline__ FP16ErrorScalePair compute_fp16_error_scales(const ScalePair &scales) { - FP16ErrorScalePair result; - const uint32_t packed_scales = static_cast(fp8_bits(scales.map4)) | - (static_cast(fp8_bits(scales.map6)) << 8); - asm volatile( - "{\n" - ".reg .b16 fp8_pair;\n" - ".reg .b16 map4_h, map6_h;\n" - ".reg .b32 scale_h2;\n" - "cvt.u16.u32 fp8_pair, %2;\n" - "cvt.rn.f16x2.e4m3x2 scale_h2, fp8_pair;\n" - "mov.b32 {map4_h, map6_h}, scale_h2;\n" - "mov.b32 %0, {map4_h, map4_h};\n" - "mov.b32 %1, {map6_h, map6_h};\n" - "}" - : "=r"(result.map4), "=r"(result.map6) - : "r"(packed_scales)); - return result; -} - -__device__ __forceinline__ float2 f16x2_scaled_to_float2(const uint32_t q_h2, - const uint32_t scale_h2) { - float2 result; - asm volatile( - "{\n" - ".reg .b16 lo, hi;\n" - ".reg .b32 prod_h2;\n" - "mul.rn.f16x2 prod_h2, %2, %3;\n" - "mov.b32 {lo, hi}, prod_h2;\n" - "cvt.f32.f16 %0, lo;\n" - "cvt.f32.f16 %1, hi;\n" - "}" - : "=f"(result.x), "=f"(result.y) - : "r"(q_h2), "r"(scale_h2)); - return result; -} - -template -__device__ __forceinline__ void accumulate_fp16_scaled_error_pair(const uint32_t q_h2, - const float x0, const float x1, - const uint32_t scale_h2, - const float global_encode_scale, - float *err) { - const float2 candidate = f16x2_scaled_to_float2(q_h2, scale_h2); - const float original0 = __fmul_rn(x0, global_encode_scale); - const float original1 = __fmul_rn(x1, global_encode_scale); - const float diff0 = __fsub_rn(candidate.x, original0); - const float diff1 = __fsub_rn(candidate.y, original1); - *err = __fadd_rn(*err, compute_error_rn(diff0)); - *err = __fadd_rn(*err, compute_error_rn(diff1)); -} - -template -__device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( - const float (&x)[8], const float block_scale_inverse, const nvfp4_scale_t sf, - const uint32_t fp16_error_scale, const float global_amax, const float global_encode_scale, - float *err) { - uint32_t out = 0; - uint32_t out_dequant_1 = 0; - uint32_t out_dequant_2 = 0; - uint32_t out_dequant_3 = 0; - uint32_t out_dequant_4 = 0; - - constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; - if constexpr (is_blackwell) { - asm volatile( - "{\n" - ".reg .b8 byte0, byte1, byte2, byte3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %6, %5;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte1, %8, %7;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte2, %10, %9;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte3, %12, %11;\n" - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" - "cvt.rn.f16x2.e2m1x2 %1, byte0;\n" - "cvt.rn.f16x2.e2m1x2 %2, byte1;\n" - "cvt.rn.f16x2.e2m1x2 %3, byte2;\n" - "cvt.rn.f16x2.e2m1x2 %4, byte3;\n" - "}" - : "=r"(out), "=r"(out_dequant_1), "=r"(out_dequant_2), "=r"(out_dequant_3), - "=r"(out_dequant_4) - : "f"(__fmul_rn(x[0], block_scale_inverse)), "f"(__fmul_rn(x[1], block_scale_inverse)), - "f"(__fmul_rn(x[2], block_scale_inverse)), "f"(__fmul_rn(x[3], block_scale_inverse)), - "f"(__fmul_rn(x[4], block_scale_inverse)), "f"(__fmul_rn(x[5], block_scale_inverse)), - "f"(__fmul_rn(x[6], block_scale_inverse)), "f"(__fmul_rn(x[7], block_scale_inverse))); - } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - } - - if constexpr (Cfg::err_use_fast_math) { - accumulate_fp16_scaled_error_pair(out_dequant_1, x[0], x[1], fp16_error_scale, - global_encode_scale, err); - accumulate_fp16_scaled_error_pair(out_dequant_2, x[2], x[3], fp16_error_scale, - global_encode_scale, err); - accumulate_fp16_scaled_error_pair(out_dequant_3, x[4], x[5], fp16_error_scale, - global_encode_scale, err); - accumulate_fp16_scaled_error_pair(out_dequant_4, x[6], x[7], fp16_error_scale, - global_encode_scale, err); - } else { - const float sf_float = static_cast(sf); - accumulate_dequant_error(out_dequant_1, x[0], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_1, x[1], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_2, x[2], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_2, x[3], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_3, x[4], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_3, x[5], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_4, x[6], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_4, x[7], sf_float, global_amax, err); - } - return out; -} - -template -__device__ __forceinline__ CandidatePair make_candidates(const float (&x0)[8], const float (&x1)[8], - const ScalePair &scales, - const float global_amax) { - CandidatePair candidates; - candidates.map4.err = 0.0f; - candidates.map6.err = 0.0f; - FP16ErrorScalePair fp16_error_scales{}; - if constexpr (Cfg::err_use_fast_math) { - fp16_error_scales = compute_fp16_error_scales(scales); - } - candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( - x0, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, - scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( - x0, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, - scales.global_encode_scale, &candidates.map6.err); - candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( - x1, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, - scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( - x1, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, - scales.global_encode_scale, &candidates.map6.err); - return candidates; -} - -__device__ __forceinline__ float reduce_group_sum_16(float value) { - const int lane = threadIdx.x & (kWarpThreads - 1); - const int group_base = lane & ~(kGroupSize - 1); - const unsigned mask = 0xffffu << group_base; -#pragma unroll - for (int offset = kGroupSize / 2; offset > 0; offset /= 2) { - value += __shfl_down_sync(mask, value, offset, kGroupSize); - } - return __shfl_sync(mask, value, group_base, kWarpThreads); -} - -__device__ __forceinline__ float reduce_group_max_16(float value) { - const int lane = threadIdx.x & (kWarpThreads - 1); - const int group_base = lane & ~(kGroupSize - 1); - const unsigned mask = 0xffffu << group_base; -#pragma unroll - for (int offset = kGroupSize / 2; offset > 0; offset /= 2) { - value = fmaxf(value, __shfl_down_sync(mask, value, offset, kGroupSize)); - } - return __shfl_sync(mask, value, group_base, kWarpThreads); -} - -__device__ __forceinline__ void store_packed_group(const uint32_t *packed, fp4e2m1x2 *dst) { - const uint64_t packed64 = - static_cast(packed[0]) | (static_cast(packed[1]) << 32); - *reinterpret_cast(dst) = packed64; -} - -__device__ __forceinline__ const uint32_t *select_packed(const CandidatePair &candidates, - const bool pick_map4) { - if (pick_map4) { - return candidates.map4.packed; - } - return candidates.map6.packed; -} - -__device__ __forceinline__ nvfp4_scale_t select_scale(const ScalePair &scales, - const bool pick_map4) { - if (pick_map4) { - return scales.map4; - } - return scales.map6; -} - -__device__ __forceinline__ void cp_async_cg_16(void *dst, const void *src) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - const uint32_t dst_smem_ptr = __cvta_generic_to_shared(dst); - const uint64_t src_gmem_ptr = reinterpret_cast(src); - asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst_smem_ptr), - "l"(src_gmem_ptr)); -#else - NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); -#endif -} - -__device__ __forceinline__ void cp_async_commit_group() { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - asm volatile("cp.async.commit_group;\n" ::); -#else - NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); -#endif -} - -template -__device__ __forceinline__ void cp_async_wait_group() { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); -#else - NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); -#endif -} - -template -__device__ void load_stage_to_shared_async(const IType *input, IType *tile, const size_t rows, - const size_t cols, const size_t stage_row, - const size_t tile_col) { - constexpr int vec_elems = 16 / sizeof(IType); - constexpr int vecs_per_row = kTileCols / vec_elems; - constexpr int vecs = kStageRows * vecs_per_row; - using TileVec = Vec; - - for (int idx = threadIdx.x; idx < vecs; idx += blockDim.x) { - const int local_row = idx / vecs_per_row; - const int local_vec_col = idx - local_row * vecs_per_row; - const int local_col = local_vec_col * vec_elems; - const size_t global_row = stage_row + local_row; - const size_t global_col = tile_col + local_col; - IType *stage_ptr = &tile[local_row * kTileCols + local_col]; - - if (global_row < rows && global_col + vec_elems <= cols) { - cp_async_cg_16(stage_ptr, &input[global_row * cols + global_col]); - } else { - TileVec vec; - vec.clear(); -#pragma unroll - for (int i = 0; i < vec_elems; ++i) { - if (global_row < rows && global_col + i < cols) { - vec.data.elt[i] = input[global_row * cols + global_col + i]; - } - } - vec.store_to(stage_ptr); - } - } -} - -template -__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvfp4_scale_t *scales, - const float *amax, const size_t rows, const size_t cols, - const size_t stage_row, const size_t tile_col, - const size_t scale_stride) { - constexpr int groups = kStageRows * kTileColGroups; - for (int group = threadIdx.x; group < groups; group += blockDim.x) { - const int local_row = group % kStageRows; - const int local_col_group = group / kStageRows; - const int local_col = local_col_group * kGroupSize; - const size_t global_row = stage_row + local_row; - const size_t global_col = tile_col + local_col; - if (global_row >= rows || global_col >= cols) { - continue; - } - - float x0[8]; - float x1[8]; - float group_amax = 0.0f; - load_row_group(tile, local_row, local_col, x0, x1, &group_amax); - - float block_amax = group_amax; - if constexpr (USE_2D_QUANTIZATION) { - block_amax = reduce_group_max_16(group_amax); - } - - float global_amax = amax[0]; - if constexpr (ROW_SCALED_NVFP4) { - global_amax = amax[global_row]; - } - - const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); - CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); - - float err_map4 = candidates.map4.err; - float err_map6 = candidates.map6.err; - if constexpr (USE_2D_QUANTIZATION) { - err_map4 = reduce_group_sum_16(err_map4); - err_map6 = reduce_group_sum_16(err_map6); - } - - const bool pick_map4 = err_map4 < err_map6; - const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); - const uint32_t *selected = select_packed(candidates, pick_map4); - - const size_t global_col_group = global_col / kGroupSize; - scales[global_row * scale_stride + global_col_group] = selected_scale; - store_packed_group(selected, &output[(global_row * cols + global_col) / 2]); - } -} - -template -__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, - nvfp4_scale_t *scales_t, const float *amax, - const size_t rows, const size_t cols, const size_t stage_row, - const size_t tile_col, const size_t scale_stride_t) { - constexpr int groups = kStageRowGroups * kTileCols; - for (int group = threadIdx.x; group < groups; group += blockDim.x) { - const int local_row_group = group / kTileCols; - const int local_col = group - local_row_group * kTileCols; - const int local_row = local_row_group * kGroupSize; - const size_t global_row = stage_row + local_row; - const size_t global_col = tile_col + local_col; - if (global_row >= rows || global_col >= cols) { - continue; - } - - float x0[8]; - float x1[8]; - float group_amax = 0.0f; - load_col_group(tile, local_row, local_col, x0, x1, &group_amax); - - float block_amax = group_amax; - if constexpr (USE_2D_QUANTIZATION) { - block_amax = reduce_group_max_16(group_amax); - } - - const float global_amax = amax[0]; - const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); - CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); - - float err_map4 = candidates.map4.err; - float err_map6 = candidates.map6.err; - if constexpr (USE_2D_QUANTIZATION) { - err_map4 = reduce_group_sum_16(err_map4); - err_map6 = reduce_group_sum_16(err_map6); - } - - const bool pick_map4 = err_map4 < err_map6; - const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); - const uint32_t *selected = select_packed(candidates, pick_map4); - - const size_t global_row_group = global_row / kGroupSize; - scales_t[global_col * scale_stride_t + global_row_group] = selected_scale; - store_packed_group(selected, &output_t[(global_col * rows + global_row) / 2]); - } -} - -template -__global__ void __launch_bounds__(kThreads) - quantize_4over6_kernel(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, - nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, - const float *amax_rowwise, const float *amax_colwise, const size_t rows, - const size_t cols, const size_t scale_stride, - const size_t scale_stride_t, const float *noop) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - if (noop != nullptr && noop[0] == 1.0f) { - return; - } - - extern __shared__ char dynamic_shmem[]; - auto *tiles = reinterpret_cast(dynamic_shmem); - const size_t tile_col = blockIdx.x * kTileCols; - const size_t tile_row = blockIdx.y * kTileRows; - - IType *stage_tiles[kPipelineStages]; -#pragma unroll - for (int stage = 0; stage < kPipelineStages; ++stage) { - stage_tiles[stage] = &tiles[stage * kStageRows * kTileCols]; - } - - load_stage_to_shared_async(input, stage_tiles[0], rows, cols, tile_row, tile_col); - cp_async_commit_group(); - cp_async_wait_group<0>(); - __syncthreads(); - - for (int stage = 0; stage < kPipelineStages; ++stage) { - const int next_stage = stage + 1; - if (next_stage < kPipelineStages) { - const size_t next_stage_row = tile_row + next_stage * kStageRows; - load_stage_to_shared_async(input, stage_tiles[next_stage], rows, cols, next_stage_row, - tile_col); - cp_async_commit_group(); - } - - const size_t stage_row = tile_row + stage * kStageRows; - IType *stage_tile = stage_tiles[stage]; - - if constexpr (RETURN_IDENTITY) { - quantize_stage_rowwise( - stage_tile, output, scales, amax_rowwise, rows, cols, stage_row, tile_col, scale_stride); - } - - if constexpr (RETURN_TRANSPOSE) { - const float *columnwise_amax = amax_colwise; - if (columnwise_amax == nullptr) { - columnwise_amax = amax_rowwise; - } - quantize_stage_colwise( - stage_tile, output_t, scales_t, columnwise_amax, rows, cols, stage_row, tile_col, - scale_stride_t); - } - - if (next_stage < kPipelineStages) { - cp_async_wait_group<0>(); - __syncthreads(); - } - } -#else - NVTE_DEVICE_ERROR("sm_100 or higher is required."); -#endif -} - template void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { @@ -639,18 +81,37 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out const size_t scale_stride = return_identity ? output->scale_inv.shape[1] : 0; const size_t scale_stride_t = return_transpose ? output->columnwise_scale_inv.shape[1] : 0; - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_identity, RETURN_IDENTITY, { - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { - auto kernel = quantize_4over6_kernel; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); - kernel<<>>(input_ptr, output_ptr, output_t_ptr, scales_ptr, - scales_t_ptr, amax_rowwise_ptr, amax_colwise_ptr, - rows, cols, scale_stride, scale_stride_t, noop_ptr); +#if NVTE_BUILD_LEGACY_STATIC_NVFP4 + const bool use_rtc = transformer_engine::rtc::is_enabled(); +#else + constexpr bool use_rtc = true; +#endif + if (use_rtc) { + rtc_nvfp4::launch_quantize_4over6_rtc( + input_ptr, output_ptr, output_t_ptr, scales_ptr, scales_t_ptr, amax_rowwise_ptr, + amax_colwise_ptr, rows, cols, scale_stride, scale_stride_t, noop_ptr, return_identity, + return_transpose, row_scaled_nvfp4, grid, block, shmem, stream); + } else { +#if NVTE_BUILD_LEGACY_STATIC_NVFP4 + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_identity, RETURN_IDENTITY, { + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { + auto kernel = + quantize_4over6_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); + kernel<<>>( + input_ptr, output_ptr, output_t_ptr, scales_ptr, scales_t_ptr, amax_rowwise_ptr, + amax_colwise_ptr, rows, cols, scale_stride, scale_stride_t, noop_ptr); + }); }); }); - }); +#else + NVTE_ERROR( + "NVFP4 4over6 quantize kernel requires NVRTC. Unset NVTE_DISABLE_NVRTC, or rebuild with " + "NVTE_BUILD_LEGACY_STATIC_NVFP4=ON for the static fallback."); +#endif + } } } // namespace quantize_4over6_kernel diff --git a/transformer_engine/common/cast/nvfp4/rtc/quantize_4over6.cu b/transformer_engine/common/cast/nvfp4/rtc/quantize_4over6.cu new file mode 100644 index 0000000000..1cd2bd3388 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/rtc/quantize_4over6.cu @@ -0,0 +1,31 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// NVRTC source file for the NVFP4 4over6 quantize kernel + +#include "quantize_4over6_kernel.cuh" + +using namespace transformer_engine; +using namespace transformer_engine::dispatch::nvfp4; +using namespace transformer_engine::dispatch::nvfp4::quantize_4over6_kernel; + +namespace { +// Substituted at compile time by the host dispatch. +using IType = __ITYPE__; +using Cfg = Config<__MODE__, __ERR_FAST_MATH__>; +} // namespace + +__global__ void __launch_bounds__(kThreads) + quantize_4over6_rtc_kernel(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, + nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, + const float *amax_rowwise, const float *amax_colwise, + const size_t rows, const size_t cols, const size_t scale_stride, + const size_t scale_stride_t, const float *noop) { + quantize_4over6_body<__USE_2D__, __RETURN_IDENTITY__, __RETURN_TRANSPOSE__, __ROW_SCALED__, Cfg, + __E4M3_MAX__, IType>(input, output, output_t, scales, scales_t, amax_rowwise, + amax_colwise, rows, cols, scale_stride, scale_stride_t, + noop); +} diff --git a/transformer_engine/common/cast/nvfp4/rtc_dispatch.cpp b/transformer_engine/common/cast/nvfp4/rtc_dispatch.cpp new file mode 100644 index 0000000000..718a0155f0 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/rtc_dispatch.cpp @@ -0,0 +1,73 @@ +/************************************************************************* + * 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 +#include "string_code_cast_nvfp4_core_nvfp4_cuh.h" +#include "string_code_cast_nvfp4_quantize_4over6_kernel_cuh.h" +#include "string_code_cast_nvfp4_rtc_quantize_4over6_cu.h" +#include "string_code_transformer_engine_nvfp4_4over6_h.h" +#include "string_code_util_ptx_cuh.h" +#include "string_code_util_type_extrema_h.h" + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { +namespace rtc_nvfp4 { + +void compile_quantize_4over6_rtc(const std::string &kernel_label, const std::string &itype_name, + bool use_2d, bool return_identity, bool return_transpose, + bool row_scaled, const std::string &mode_name, bool err_fast_math, + int e4m3_max) { + auto &mgr = rtc::KernelManager::instance(); + if (mgr.is_compiled(kernel_label)) { + return; + } + + auto bool_str = [](bool value) -> std::string { return value ? "true" : "false"; }; + + std::string code = string_code_cast_nvfp4_rtc_quantize_4over6_cu; + code = regex_replace(code, "__ITYPE__", itype_name); + code = regex_replace(code, "__USE_2D__", bool_str(use_2d)); + code = regex_replace(code, "__RETURN_IDENTITY__", bool_str(return_identity)); + code = regex_replace(code, "__RETURN_TRANSPOSE__", bool_str(return_transpose)); + code = regex_replace(code, "__ROW_SCALED__", bool_str(row_scaled)); + code = regex_replace(code, "__MODE__", mode_name); + code = regex_replace(code, "__ERR_FAST_MATH__", bool_str(err_fast_math)); + code = regex_replace(code, "__E4M3_MAX__", std::to_string(e4m3_max)); + + const std::vector headers = { + {string_code_cast_nvfp4_quantize_4over6_kernel_cuh, "quantize_4over6_kernel.cuh"}, + {string_code_cast_nvfp4_core_nvfp4_cuh, "core_nvfp4.cuh"}, + {string_code_util_ptx_cuh, "ptx.cuh"}, + {string_code_util_type_extrema_h, "util/type_extrema.h"}, + {string_code_transformer_engine_nvfp4_4over6_h, "transformer_engine/nvfp4/4over6.h"}, + }; + + // --device-int128: ptx.cuh uses __uint128_t; -default-device: treat the + // unannotated constexpr/inline helpers in ptx.cuh as __device__ under JIT. + // -DCUDA_VERSION: NVRTC does not predefine CUDA_VERSION, but the 4over6 kernel + // needs the FP4 types gated behind FP4_TYPE_SUPPORTED (== CUDA_VERSION >= 12080) + // and ; forward the build's CUDA version so those are enabled. + const std::vector options = {"--device-int128", "-default-device", + "-DCUDA_VERSION=" + std::to_string(CUDA_VERSION)}; + constexpr rtc::ArchRequirement arch_requirement{100, rtc::ArchSpecificity::ArchitectureSpecific}; + + mgr.compile(kernel_label, "quantize_4over6_rtc_kernel", code, + "transformer_engine/common/cast/nvfp4/rtc/quantize_4over6.cu", options, headers, + arch_requirement); +} + +} // namespace rtc_nvfp4 +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine diff --git a/transformer_engine/common/cast/nvfp4/rtc_dispatch.h b/transformer_engine/common/cast/nvfp4/rtc_dispatch.h new file mode 100644 index 0000000000..23cc083068 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/rtc_dispatch.h @@ -0,0 +1,103 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file rtc_dispatch.h + * \brief Host-side NVRTC dispatch for the NVFP4 4over6 quantize kernel. + */ + +#ifndef TRANSFORMER_ENGINE_CAST_NVFP4_RTC_DISPATCH_H_ +#define TRANSFORMER_ENGINE_CAST_NVFP4_RTC_DISPATCH_H_ + +#include + +#include + +#include "../../util/rtc.h" +// do not include util/string.h here — it pulls in , which is heavy. +#include "core_nvfp4.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { +namespace rtc_nvfp4 { + +#if FP4_TYPE_SUPPORTED + +void compile_quantize_4over6_rtc(const std::string &kernel_label, const std::string &itype_name, + bool use_2d, bool return_identity, bool return_transpose, + bool row_scaled, const std::string &mode_name, bool err_fast_math, + int e4m3_max); + +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"; +} + +// Enum literal spelling for the __MODE__ substitution / cache key. +inline const char *rtc_mode_name(NVTENVFP44Over6Mode mode) { + switch (mode) { + case kNVTENVFP44Over6MinMSE: + return "kNVTENVFP44Over6MinMSE"; + case kNVTENVFP44Over6MinMAE: + return "kNVTENVFP44Over6MinMAE"; + default: + NVTE_ERROR("Unsupported NVFP4 4over6 mode."); + } +} + +template +inline void launch_quantize_4over6_rtc(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, + nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, + const float *amax_rowwise, const float *amax_colwise, + size_t rows, size_t cols, size_t scale_stride, + size_t scale_stride_t, const float *noop, + bool return_identity, bool return_transpose, + bool row_scaled_nvfp4, dim3 grid, dim3 block, size_t shmem, + cudaStream_t stream) { + const std::string itype_name = rtc_type_name(); + const std::string mode_name = rtc_mode_name(Cfg::mode); + + // Cache key encodes everything that varies the compiled kernel. + const std::string kernel_label = + std::string("quantize_4over6,itype=") + itype_name + + ",use2d=" + (USE_2D_QUANTIZATION ? "1" : "0") + ",id=" + (return_identity ? "1" : "0") + + ",t=" + (return_transpose ? "1" : "0") + ",rowscaled=" + (row_scaled_nvfp4 ? "1" : "0") + + ",mode=" + mode_name + ",fastmath=" + (Cfg::err_use_fast_math ? "1" : "0") + + ",e4m3max=" + std::to_string(E4M3_MAX); + + auto &mgr = rtc::KernelManager::instance(); + if (!mgr.is_compiled(kernel_label)) { + compile_quantize_4over6_rtc(kernel_label, itype_name, USE_2D_QUANTIZATION, return_identity, + return_transpose, row_scaled_nvfp4, mode_name, + Cfg::err_use_fast_math, E4M3_MAX); + } + + if (shmem > 0) { + mgr.set_function_attribute(kernel_label, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + static_cast(shmem)); + } + + mgr.launch(kernel_label, grid, block, static_cast(shmem), stream, input, output, + output_t, scales, scales_t, amax_rowwise, amax_colwise, rows, cols, scale_stride, + scale_stride_t, noop); +} + +#endif // FP4_TYPE_SUPPORTED + +} // namespace rtc_nvfp4 +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_CAST_NVFP4_RTC_DISPATCH_H_ diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..9b3c9cadfa 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -647,13 +647,11 @@ using fp8e5m2 = __nv_fp8_e5m2; #if CUDA_VERSION >= 12080 using fp8e8m0 = __nv_fp8_e8m0; #endif -#if FP4_TYPE_SUPPORTED -using fp4e2m1 = __nv_fp4_e2m1; -using fp4e2m1x2 = __nv_fp4x2_e2m1; -using fp4e2m1x4 = __nv_fp4x4_e2m1; -#endif using e8m0_t = uint8_t; +// FP4 aliases + device-safe TypeExtrema specializations (also used by NVRTC). +#include "./util/type_extrema.h" + namespace detail { template @@ -680,46 +678,6 @@ TRANSFORMER_ENGINE_TYPE_NAME(__nv_fp4_e2m1) #endif #undef TRANSFORMER_ENGINE_TYPE_NAME -template -struct TypeExtrema; - -#if FP4_TYPE_SUPPORTED -template <> -struct TypeExtrema { - static constexpr float max = 6.0f; - static constexpr float max_inverse = 1.0 / max; -}; -#endif - -template <> -struct TypeExtrema { - static constexpr float max = 448.0f; - static constexpr float max_inverse = 1.0 / max; -}; - -template <> -struct TypeExtrema { - static constexpr float max = 57344.0f; - static constexpr float max_inverse = 1.0 / max; -}; - -template <> -struct TypeExtrema { - // Hex float format of 1.(7 bits of 1) * 2 ^ 127 - static constexpr float max = 0x1.FEp127; -}; - -template <> -struct TypeExtrema { - // Hex float format of 1.(10 bits of 1) * 2 ^ 15 - static constexpr float max = 0x1.FFCp15; -}; - -template -struct TypeExtrema { - static constexpr float max = std::numeric_limits::max(); -}; - } // namespace detail template diff --git a/transformer_engine/common/include/transformer_engine/nvfp4/4over6.h b/transformer_engine/common/include/transformer_engine/nvfp4/4over6.h new file mode 100644 index 0000000000..b186d9e49c --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/nvfp4/4over6.h @@ -0,0 +1,26 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file 4over6.h + * \brief Device-safe NVFP4 4over6 mode enum. + * + * Kept free of host-only dependencies so the same definition can be used by + * the public C API and by NVRTC-compiled device kernels. + */ + +#ifndef TRANSFORMER_ENGINE_NVFP4_4OVER6_H_ +#define TRANSFORMER_ENGINE_NVFP4_4OVER6_H_ + +/*! \enum NVTENVFP44Over6Mode + * \brief Method for NVFP4 4over6 quantization. + */ +enum NVTENVFP44Over6Mode { + kNVTENVFP44Over6Disabled = 0, /*!< 4over6 is not applied */ + kNVTENVFP44Over6MinMAE = 1, /*!< Select the candidate with lower mean absolute error */ + kNVTENVFP44Over6MinMSE = 2, /*!< Select the candidate with lower mean squared error */ +}; + +#endif // TRANSFORMER_ENGINE_NVFP4_4OVER6_H_ diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index aa0405e177..976f5fe652 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -19,6 +19,8 @@ extern "C" { #endif +#include + /*! \enum NVTEDType * \brief TE datatype. */ @@ -118,15 +120,6 @@ enum NVTEScalingMode { NVTE_INVALID_SCALING = 100 }; -/*! \enum NVTENVFP44Over6Mode - * \brief Method for NVFP4 4over6 quantization. - */ -enum NVTENVFP44Over6Mode { - kNVTENVFP44Over6Disabled = 0, /*!< 4over6 is not applied */ - kNVTENVFP44Over6MinMAE = 1, /*!< Select the candidate with lower mean absolute error */ - kNVTENVFP44Over6MinMSE = 2, /*!< Select the candidate with lower mean squared error */ -}; - /*! \brief TE Tensor type * * NVTETensor is a contiguous tensor type storing a pointer diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 2814aa3490..02de095952 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -11,17 +11,34 @@ #ifndef TRANSFORMER_ENGINE_PTX_CUH_ #define TRANSFORMER_ENGINE_PTX_CUH_ +#if !defined(__CUDACC_RTC__) #include #include #include "common/common.h" +#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 +#endif // __CUDACC_RTC__ +#include #if FP4_TYPE_SUPPORTED #include #endif // FP4_TYPE_SUPPORTED -#include +#if !defined(__CUDACC_RTC__) #include "common/utils.cuh" +#else +#include "utils.cuh" + +namespace transformer_engine { +using fp16 = half; +using bf16 = nv_bfloat16; +} // namespace transformer_engine +#endif // __CUDACC_RTC__ namespace transformer_engine { @@ -791,7 +808,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 +839,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 +900,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 +930,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..3adaa2357f 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,33 @@ 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(); - } - return arch_; + 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(); } } // namespace @@ -148,7 +160,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 +171,23 @@ 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 = ""; + if (arch_requirement.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"; + } // Compilation flags std::vector opts = { @@ -168,9 +196,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..b96cce8418 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -40,6 +40,22 @@ 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, +}; + +/*! \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 +161,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/util/type_extrema.h b/transformer_engine/common/util/type_extrema.h new file mode 100644 index 0000000000..6813b72534 --- /dev/null +++ b/transformer_engine/common/util/type_extrema.h @@ -0,0 +1,78 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file type_extrema.h + * \brief Device-safe FP4 aliases and TypeExtrema specializations. + * + * Include from inside namespace transformer_engine after fp8e4m3 / fp8e5m2 / + * bf16 / fp16 aliases are available. The NVRTC path only uses the explicit + * device-safe specializations below. + * + * Requires FP4_TYPE_SUPPORTED to be defined (as in common.h / ptx.cuh) and, + * when it is true, to be visible. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_UTIL_TYPE_EXTREMA_H_ +#define TRANSFORMER_ENGINE_COMMON_UTIL_TYPE_EXTREMA_H_ + +#if FP4_TYPE_SUPPORTED +using fp4e2m1 = __nv_fp4_e2m1; +using fp4e2m1x2 = __nv_fp4x2_e2m1; +using fp4e2m1x4 = __nv_fp4x4_e2m1; +#endif // FP4_TYPE_SUPPORTED + +namespace detail { + +#if !defined(__CUDACC_RTC__) +template +struct TypeExtrema { + static constexpr float max = std::numeric_limits::max(); +}; +#else +template +struct TypeExtrema; +#endif + +#if FP4_TYPE_SUPPORTED +template <> +struct TypeExtrema { + static constexpr float max = 6.0f; + static constexpr float max_inverse = 1.0 / max; +}; +#endif // FP4_TYPE_SUPPORTED + +template <> +struct TypeExtrema { + static constexpr float max = 448.0f; + static constexpr float max_inverse = 1.0 / max; +}; + +template <> +struct TypeExtrema { + static constexpr float max = 57344.0f; + static constexpr float max_inverse = 1.0 / max; +}; + +template <> +struct TypeExtrema { + // Hex float format of 1.(7 bits of 1) * 2 ^ 127 + static constexpr float max = 0x1.FEp127; +}; + +template <> +struct TypeExtrema { + // Hex float format of 1.(10 bits of 1) * 2 ^ 15 + static constexpr float max = 0x1.FFCp15; +}; + +template <> +struct TypeExtrema { + static constexpr float max = 0x1.fffffep127f; // FLT_MAX +}; + +} // namespace detail + +#endif // TRANSFORMER_ENGINE_COMMON_UTIL_TYPE_EXTREMA_H_ diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 0b75da622c..5283ce0e43 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -24,10 +24,16 @@ 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 int16_t = short int; // NOLINT(*) +using int32_t = int; +using int64_t = 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(int16_t) == 2); +static_assert(sizeof(int32_t) == 4); +static_assert(sizeof(int64_t) == 8); #endif // Minimal subset of used by RTC kernel headers. Keep these in a