diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e7aaf78f6c..4c15a2a312 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -264,6 +264,8 @@ list(APPEND transformer_engine_cuda_arch_specific_sources activation/swiglu_grouped.cu activation/swiglu_grouped_dbias.cu cast/cast.cu + cast/mxfp8/specialized/cast_rowwise.cu + cast/mxfp8/specialized/cast_bidim.cu cast/cast_dbias.cu cast/cast_grouped.cu cast/cast_grouped_dbias.cu diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 9f312ac3f5..17db8a59ea 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -21,6 +21,8 @@ #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" #include "../core/common.cuh" +#include "specialized/cast_bidim.h" +#include "specialized/cast_rowwise.h" #include "specialized/quantize_mxfp8.cuh" #include "swizzle.cuh" @@ -745,6 +747,19 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, !use_2d_quantization && scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { + // The register-resident kernel supersedes the staged one below + // wherever it applies: 10-20% faster on the cast-only path. + // hasSpec has already established cast-only, so what is left + // is the input type and the scale layout, neither of which + // that kernel generalizes over yet. + if constexpr (std::is_same_v && !WITH_GEMM_SWIZZLED_SCALES) { + specialized::launch_cast_rowwise( + input.data.dptr, output->data.dptr, + reinterpret_cast(scales_rowwise_ptr), static_cast(rows), + static_cast(cols), static_cast(scale_stride_rowwise), stream); + break; + } + using traits = specialized::CastTraits; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; @@ -765,6 +780,25 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, break; } case ScalingType::BIDIMENSIONAL: { + // The register-resident kernel supersedes the TMA one below + // wherever it applies: it reads its 32-row tile once and drives + // both the rowwise and colwise passes from registers, roughly + // 1.2x faster on shapes that fit its tiling. hasSpec has + // already established cast-only, leaving the input type, the + // scale layout, and the tile alignment. + if constexpr (std::is_same_v && !WITH_GEMM_SWIZZLED_SCALES) { + if (rows % 32 == 0 && cols % 256 == 0) { + specialized::launch_cast_bidim( + input.data.dptr, output->data.dptr, + reinterpret_cast(scales_rowwise_ptr), + output->columnwise_data.dptr, + reinterpret_cast(scales_colwise_ptr), static_cast(rows), + static_cast(cols), static_cast(scale_stride_rowwise), + static_cast(scale_stride_colwise), stream); + break; + } + } + using traits = specialized::CastTraitsSwizzle + +#include "../../../common.h" +#include "../../../util/cuda_runtime.h" +#include "../../../util/ptx.cuh" +#include "../../../utils.cuh" +#include "cast_bidim.h" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { +namespace specialized { + +namespace ptx = transformer_engine::ptx; + +// Bidimensional MXFP8 produces two independent quantizations of one BF16 +// tensor: a rowwise one, where 32 consecutive elements of a row share a scale, +// and a colwise one, where 32 consecutive elements of a column share a scale. +// Both outputs keep the input's [M, K] row-major layout -- the colwise result +// is not transposed, only scaled differently. +// +// A CTA owns a 32-row band crossed with a strip of columns. Thirty-two rows is +// exactly the colwise block height, so the colwise reduction closes inside the +// CTA: the tile is read from global memory once, into registers, and drives +// both passes. That is what lets this kernel drop the TMA pipeline the +// specialized kernel uses -- there is no second pass to stage for. +// +// Within a tile: +// - each warp owns kRowsPerWarp consecutive rows of the band, +// - each lane owns COLS_PER_LANE consecutive columns of those rows, +// - the rowwise scale needs a reduction across the lanes that share a +// 32-column block, which shuffles handle, +// - the colwise scale needs a reduction down all 32 rows, hence across warps, +// which is the one thing shuffles cannot reach and the only reason this +// kernel touches shared memory. + +// Elements in one MX block, and equivalently the colwise block height. +constexpr int32_t kBlockElems = 32; +// Rows a CTA covers. Must equal kBlockElems to keep the colwise reduction +// CTA-local. +constexpr int32_t kRowsPerTile = kBlockElems; + +constexpr int32_t kWarpsPerCta = 8; +constexpr int32_t kThreadsPerCta = kWarpsPerCta * THREADS_PER_WARP; +constexpr int32_t kRowsPerWarp = kRowsPerTile / kWarpsPerCta; +static_assert(kRowsPerTile % kWarpsPerCta == 0, "Warps must divide the row band evenly."); + +// BF16 packs two elements per 32-bit word, and all the packed math below works +// on those words, so a shared-memory "slot" holds one column pair. +constexpr int32_t kElemsPerWord = sizeof(uint32_t) / sizeof(bf16); + +// Bytes in an L2 cache line, for sizing the software prefetch. +constexpr int32_t kCacheLineBytes = 128; + +// +// MX scale arithmetic. +// +// The E8M0 scale of a block is the smallest power of two with +// amax / scale <= 448, i.e. biased_exponent = ceil(log2(amax / 448)) + 127. +// For a BF16 amax that has an exact closed form in the bit pattern. +// +// Write amax = m * 2^(e - 127), with the biased exponent e in bits 7..14 and +// the stored mantissa in bits 0..6. Then +// +// ceil(log2(amax / 448)) + 127 = e - 8 for m <= 1.75 +// = e - 7 for m > 1.75 +// +// and "m > 1.75" is exactly "mantissa field >= 97", because 1.75 is mantissa 96 +// and 97 is the next representable step. Adding 31 to the mantissa field +// carries into the exponent for mantissa >= 97 and for nothing below it, so a +// single add and mask round e by precisely that rule. +// +// This is not an approximation of the usual amax * (1/448) route: it is exact, +// and better behaved, because 1/448 is not representable in FP32 and that +// product can round across the boundary. +// +// The helpers return the reciprocal scale, ready to multiply by, since that is +// what the packed conversion consumes. mx_scale_byte recovers the stored E8M0 +// byte from it. + +constexpr uint32_t kMantissaRoundUp = 31u; +constexpr uint32_t kBf16ExponentMask = 0x7F80u; + +// Exponent of the output type's largest normal: 8 for E4M3 (448 = 1.75 * 2^8) +// and 15 for E5M2 (57344 = 1.75 * 2^15). Both share the 1.75 mantissa, which +// is exactly why the "mantissa field >= 97" rule above holds for either and +// only this offset has to change. +template +constexpr uint32_t kMaxNormExponent = Quantized_Limits::max_unbiased_exponent; + +// Smallest exponent field allowed, so the reciprocal stays a normal BF16. +template +constexpr uint32_t kMinExponentField = kMaxNormExponent << 7; + +// Encoding of 2^(127 + offset); subtracting the rounded exponent field yields +// the reciprocal scale directly. +template +constexpr uint32_t kReciprocalBias = (254u + kMaxNormExponent) << 7; +// BF16 magnitude at or above which a value is Inf or NaN. +constexpr uint32_t kBf16InfBits = 0x7F80u; +// Reciprocal paired with the E8M0 NaN scale (254): the smallest subnormal. +constexpr uint32_t kNaNReciprocal = 0x0040u; +constexpr uint32_t kBf16MagnitudeMask = 0x7FFFu; +// Bias such that kScaleByteBias - (reciprocal >> 7) is the E8M0 scale byte. +constexpr uint32_t kScaleByteBias = 254u; + +/*! \brief Reciprocal MX scale for one BF16 amax, as BF16 bits. + * \param amax_bits Block amax magnitude bits, sign already cleared. */ +template +__device__ __forceinline__ uint32_t mx_scale_reciprocal(uint32_t amax_bits) { + if (amax_bits >= kBf16InfBits) { + return kNaNReciprocal; + } + const uint32_t rounded_exponent = + max((amax_bits + kMantissaRoundUp) & kBf16ExponentMask, kMinExponentField); + return kReciprocalBias - rounded_exponent; +} + +/*! \brief E8M0 scale byte matching a reciprocal from mx_scale_reciprocal. */ +__device__ __forceinline__ e8m0_t mx_scale_byte(uint32_t reciprocal_bits) { + return static_cast(kScaleByteBias - (reciprocal_bits >> 7)); +} + +/*! \brief mx_scale_reciprocal applied to two columns at once. + * + * Valid only when neither half is Inf or NaN. Callers test that first, which + * costs one comparison for the pair. + */ +template +__device__ __forceinline__ uint32_t mx_scale_reciprocal_x2(uint32_t amax_pair) { + constexpr uint32_t kMagnitudeMaskPair = 0x7FFF7FFFu; + constexpr uint32_t kMantissaRoundUpPair = 0x001F001Fu; + constexpr uint32_t kExponentMaskPair = 0xFF80FF80u; + constexpr uint32_t kMinExponentPair = (kMinExponentField << 16) | kMinExponentField; + constexpr uint32_t kReciprocalBiasPair = (kReciprocalBias << 16) | kReciprocalBias; + + // Each half is at most 0x7FFF, so adding 31 cannot carry out of the low half + // into the high one; the two halves round independently. + const uint32_t magnitudes = amax_pair & kMagnitudeMaskPair; + ptx::bf16x2 rounded, floor_pair, clamped; + reinterpret_cast(rounded) = (magnitudes + kMantissaRoundUpPair) & kExponentMaskPair; + reinterpret_cast(floor_pair) = kMinExponentPair; + // Both operands are positive BF16 patterns, so the magnitude-max doubles as a + // per-half clamp against the minimum exponent. + ptx::abs_max_2x(clamped, rounded, floor_pair); + return kReciprocalBiasPair - reinterpret_cast(clamped); +} + +/*! \brief The two E8M0 scale bytes of a packed reciprocal pair, in the low + * 16 bits of the result. */ +__device__ __forceinline__ uint32_t mx_scale_byte_x2(uint32_t reciprocal_pair) { + constexpr uint32_t kScaleByteBiasPair = 0x00FE00FEu; + const uint32_t bytes = kScaleByteBiasPair - ((reciprocal_pair >> 7) & 0x00FF00FFu); + // Gather the two scale bytes, at byte 0 and byte 2, into the low half. + return __byte_perm(bytes, 0u, 0x4420); +} + +/*! \brief Fold a BF16 pair's halves together and return the magnitude. + * + * `max.xorsign.abs` sets the result sign to the XOR of its inputs, so an + * accumulator built from it can come out negative even when its magnitude is + * right. Only the magnitude matters for a scale. + */ +__device__ __forceinline__ uint32_t fold_pair_magnitude(uint32_t pair) { + ptx::bf16x2 lo, hi, folded; + reinterpret_cast(lo) = pair; + reinterpret_cast(hi) = __byte_perm(pair, pair, 0x1032); + ptx::abs_max_2x(folded, lo, hi); + return reinterpret_cast(folded) & kBf16MagnitudeMask; +} + +/*! \brief Broadcast a BF16 reciprocal scale into both halves of a pair. */ +__device__ __forceinline__ ptx::bf16x2 broadcast_pair(uint32_t scale_bits) { + ptx::bf16x2 result; + reinterpret_cast(result) = __byte_perm(scale_bits, 0u, 0x1010); + return result; +} + +/*! \brief Scale one lane's BF16 words and pack them into FP8E4M3 words. + * + * \param scales One reciprocal per input word. The rowwise pass passes the + * same broadcast value throughout; the colwise pass gives each + * column pair its own. + */ +template +__device__ __forceinline__ void scale_and_convert(const uint32_t (&in)[WORDS_IN], + const ptx::bf16x2 (&scales)[WORDS_IN], + uint32_t (&out)[WORDS_IN / 2]) { +#pragma unroll + for (int32_t i = 0; i < WORDS_IN / 2; ++i) { + ptx::mul_cvt_4x(reinterpret_cast &>(out[i]), + reinterpret_cast(in[2 * i]), scales[2 * i], + reinterpret_cast(in[2 * i + 1]), scales[2 * i + 1]); + } +} + +/*! \brief Store a lane's output words with the shared L2 policy. */ +template +__device__ __forceinline__ void store_words(void *dst, const uint32_t (&words)[WORDS], + uint64_t policy) { + static_assert(WORDS == 2 || WORDS == 4, "Output stores are 64- or 128-bit."); + if constexpr (WORDS == 4) { + ptx::st_global_b32x4(dst, words, policy); + } else { + ptx::st_global_b32x2(dst, words, policy); + } +} + +// +// Shared-memory access for the colwise fold, as one vector operation per call. +// +// WORDS is how many consecutive 32-bit slots a lane owns in each half of the +// slot array: 4 for the wide tile (LDS/STS.128) and 2 for the narrow one +// (LDS/STS.64). Spelling the vector type out, rather than looping over +// scalars and hoping the compiler merges them, is what guarantees the single +// wide access the half-split layout was designed around. +// + +template +__device__ __forceinline__ void store_shared_words(uint32_t *dst, const uint32_t *src) { + static_assert(WORDS == 2 || WORDS == 4, "Shared-memory access is 64- or 128-bit."); + if constexpr (WORDS == 4) { + *reinterpret_cast(dst) = make_uint4(src[0], src[1], src[2], src[3]); + } else { + *reinterpret_cast(dst) = make_uint2(src[0], src[1]); + } +} + +template +__device__ __forceinline__ void load_shared_words(uint32_t *dst, const uint32_t *src) { + static_assert(WORDS == 2 || WORDS == 4, "Shared-memory access is 64- or 128-bit."); + if constexpr (WORDS == 4) { + const uint4 v = *reinterpret_cast(src); + dst[0] = v.x; + dst[1] = v.y; + dst[2] = v.z; + dst[3] = v.w; + } else { + const uint2 v = *reinterpret_cast(src); + dst[0] = v.x; + dst[1] = v.y; + } +} + +/*! \brief Quantize a 32-row band crossed with a strip of columns. + * + * \tparam COLS_PER_LANE Columns one lane owns. 16 gives a 32x512 tile read + * with 256-bit loads; 8 gives 32x256 with 128-bit + * loads, trading load width for more CTAs and so + * better latency hiding on deep shapes. + * \tparam K_COMPILE_TIME When non-zero, the column count is a compile-time + * constant, turning the row stride, the scale strides + * and the grid width into immediates. Zero takes + * them at run time. + * \tparam MIN_BLOCKS_PER_SM Occupancy target for __launch_bounds__. + * + * \param prefetch_distance_ctas How far ahead, in CTAs, to prefetch input. + * Set to one full resident wave so the next + * wave's tile lands in L2 as this one drains; + * 0 disables the prefetch. + */ +template +__global__ __launch_bounds__(kThreadsPerCta, MIN_BLOCKS_PER_SM) void quantize_bidim_kernel( + const uint8_t *__restrict__ input, uint8_t *__restrict__ output_rowwise, + uint8_t *__restrict__ scales_rowwise, uint8_t *__restrict__ output_colwise, + uint8_t *__restrict__ scales_colwise, int32_t cols_rt, int32_t scale_stride_rowwise_rt, + int32_t scale_stride_colwise_rt, int32_t prefetch_distance_ctas) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + const int32_t cols = K_COMPILE_TIME ? K_COMPILE_TIME : cols_rt; + const int32_t scale_stride_rowwise = + K_COMPILE_TIME ? DIVUP_TO_MULTIPLE(K_COMPILE_TIME / kBlockElems, 4) : scale_stride_rowwise_rt; + const int32_t scale_stride_colwise = + K_COMPILE_TIME ? DIVUP_TO_MULTIPLE(K_COMPILE_TIME, 128) : scale_stride_colwise_rt; + + constexpr int32_t kWordsPerLane = COLS_PER_LANE / kElemsPerWord; + constexpr int32_t kOutWordsPerLane = kWordsPerLane / 2; + constexpr int32_t kColsPerTile = THREADS_PER_WARP * COLS_PER_LANE; + // Lanes that must cooperate to cover one 32-column rowwise block. + constexpr int32_t kLanesPerRowBlock = kBlockElems / COLS_PER_LANE; + // Column pairs in the tile, i.e. shared-memory slots for the colwise fold. + constexpr int32_t kColumnSlots = kColsPerTile / kElemsPerWord; + // Half-split layout: a lane deposits its partials as two vector stores, one + // into each half of the slot array, which keeps the store and the later + // read-back free of bank conflicts. Four slots per half for the wide tile + // (COLS_PER_LANE 16, so 128-bit accesses) and two for the narrow one + // (COLS_PER_LANE 8, so 64-bit) -- both instantiations are live, the narrow + // tile being what deep grids and K-not-a-multiple-of-512 shapes use. + constexpr int32_t kSlotsPerLaneHalf = kWordsPerLane / 2; + constexpr int32_t kHalfSlots = kColumnSlots / 2; + + // Both arrays are indexed by column *pair*, one 32-bit slot each, because + // that is the granularity the packed BF16 math works at. + // + // s_column_partial holds per-warp amax candidates: two BF16 magnitudes. + // s_column_scale holds the finished *reciprocal* scales: again two BF16, one + // per column, not the one-byte E8M0 encodings. BF16 is what mul_cvt_4x + // multiplies by, so keeping the reciprocal in that form avoids a byte-to-BF16 + // expansion inside the conversion loop; the E8M0 bytes are produced once, by + // mx_scale_byte_x2, only when warp 0 writes the scale array out. + __shared__ uint32_t s_column_partial[kWarpsPerCta][kColumnSlots]; + __shared__ uint32_t s_column_scale[kColumnSlots]; + + const int32_t tid = threadIdx.x; + const int32_t warp = tid / THREADS_PER_WARP; + const int32_t lane = tid % THREADS_PER_WARP; + + const int32_t row0 = blockIdx.y * kRowsPerTile + warp * kRowsPerWarp; + const int32_t col0 = blockIdx.x * kColsPerTile; + + // One policy for every stream. Each byte is touched once, but marking them + // evict_last measured better than streaming: holding the tile's lines through + // the CTA's lifetime is what keeps the three write bursts coalesced. + const uint64_t policy = ptx::create_l2_policy_evict_last(); + + const size_t row_stride_bytes = static_cast(cols) * sizeof(bf16); + const size_t lane_offset = static_cast(row0) * cols + col0 + COLS_PER_LANE * lane; + const uint8_t *tile_in = input + lane_offset * sizeof(bf16); + uint8_t *out_row = output_rowwise + lane_offset; + uint8_t *out_col = output_colwise + lane_offset; + uint8_t *scale_row = scales_rowwise + static_cast(row0) * scale_stride_rowwise + + (col0 / kBlockElems) + (lane / kLanesPerRowBlock); + // One lane of each cooperating group writes the shared rowwise scale byte. + const bool owns_row_scale = (lane % kLanesPerRowBlock) == 0; + + // Read the whole tile into registers; both passes run off these. + uint32_t tile[kRowsPerWarp][kWordsPerLane]; +#pragma unroll + for (int32_t i = 0; i < kRowsPerWarp; ++i) { + const void *src = tile_in + i * row_stride_bytes; + if constexpr (kWordsPerLane == 8) { + ptx::ld_global_nc_b32x8(tile[i], src, policy); + } else { + ptx::ld_global_nc_b32x4(tile[i], src, policy); + } + } + + // ---- rowwise pass ------------------------------------------------------- +#pragma unroll + for (int32_t i = 0; i < kRowsPerWarp; ++i) { + ptx::bf16x2 amax; + reinterpret_cast(amax) = tile[i][0]; +#pragma unroll + for (int32_t k = 1; k < kWordsPerLane; ++k) { + ptx::abs_max_2x(amax, amax, reinterpret_cast(tile[i][k])); + } + // Butterfly across the lanes that share this 32-column block. +#pragma unroll + for (int32_t d = 1; d < kLanesPerRowBlock; d <<= 1) { + ptx::bf16x2 partner; + reinterpret_cast(partner) = + __shfl_xor_sync(0xFFFFFFFFu, reinterpret_cast(amax), d); + ptx::abs_max_2x(amax, amax, partner); + } + + const uint32_t reciprocal = + mx_scale_reciprocal(fold_pair_magnitude(reinterpret_cast(amax))); + const ptx::bf16x2 broadcast = broadcast_pair(reciprocal); + ptx::bf16x2 scales[kWordsPerLane]; +#pragma unroll + for (int32_t k = 0; k < kWordsPerLane; ++k) { + scales[k] = broadcast; + } + uint32_t out_words[kOutWordsPerLane]; + scale_and_convert(tile[i], scales, out_words); + store_words(out_row + static_cast(i) * cols, out_words, policy); + + if (owns_row_scale) { + ptx::st_global_b8(scale_row + static_cast(i) * scale_stride_rowwise, + mx_scale_byte(reciprocal), policy); + } + } + + // ---- colwise pass ------------------------------------------------------- + // Each warp reduces its own rows first; the warps then meet in shared memory, + // the only cross-warp step in the kernel. + uint32_t column_partial[kWordsPerLane]; +#pragma unroll + for (int32_t k = 0; k < kWordsPerLane; ++k) { + ptx::bf16x2 acc; + reinterpret_cast(acc) = tile[0][k]; +#pragma unroll + for (int32_t i = 1; i < kRowsPerWarp; ++i) { + ptx::abs_max_2x(acc, acc, reinterpret_cast(tile[i][k])); + } + column_partial[k] = reinterpret_cast(acc); + } + + store_shared_words(&s_column_partial[warp][kSlotsPerLaneHalf * lane], + &column_partial[0]); + store_shared_words( + &s_column_partial[warp][kHalfSlots + kSlotsPerLaneHalf * lane], + &column_partial[kSlotsPerLaneHalf]); + __syncthreads(); + + // Fold the per-warp partials into one scale per column pair. Threads take + // stride-1 slots so the strided reads across warps stay conflict-free. + for (int32_t slot = tid; slot < kColumnSlots; slot += kThreadsPerCta) { + const uint32_t *column = &s_column_partial[0][slot]; + ptx::bf16x2 acc; + reinterpret_cast(acc) = column[0]; +#pragma unroll + for (int32_t w = 1; w < kWarpsPerCta; ++w) { + ptx::abs_max_2x(acc, acc, reinterpret_cast(column[w * kColumnSlots])); + } + const uint32_t amax_pair = reinterpret_cast(acc); + + uint32_t reciprocal_pair; + if (__builtin_expect(fold_pair_magnitude(amax_pair) < kBf16InfBits, 1)) { + reciprocal_pair = mx_scale_reciprocal_x2(amax_pair); + } else { + // At least one of the two columns is Inf or NaN; take them separately. + const uint32_t lo = mx_scale_reciprocal(amax_pair & kBf16MagnitudeMask); + const uint32_t hi = mx_scale_reciprocal((amax_pair >> 16) & kBf16MagnitudeMask); + reciprocal_pair = lo | (hi << 16); + } + s_column_scale[slot] = reciprocal_pair; + } + __syncthreads(); + + // Read back the reciprocals covering this lane's own columns. + uint32_t column_reciprocal[kWordsPerLane]; + load_shared_words(&column_reciprocal[0], + &s_column_scale[kSlotsPerLaneHalf * lane]); + load_shared_words(&column_reciprocal[kSlotsPerLaneHalf], + &s_column_scale[kHalfSlots + kSlotsPerLaneHalf * lane]); + + // Warp 0 emits the colwise scale row: one byte per column, one coalesced + // vector store per lane. + if (warp == 0) { + uint32_t packed[kOutWordsPerLane]; +#pragma unroll + for (int32_t j = 0; j < kOutWordsPerLane; ++j) { + packed[j] = mx_scale_byte_x2(column_reciprocal[2 * j]) | + (mx_scale_byte_x2(column_reciprocal[2 * j + 1]) << 16); + } + store_words(scales_colwise + static_cast(blockIdx.y) * scale_stride_colwise + col0 + + COLS_PER_LANE * lane, + packed, policy); + } + + // Software L2 prefetch, issued here so the next wave's input arrives while + // this CTA is still writing. Pulling it earlier -- right after our own loads + // -- measured worse: the lines then sit in L2 for the whole CTA lifetime and + // crowd out the write bursts. + if (prefetch_distance_ctas > 0) { + const int32_t grid_x = K_COMPILE_TIME ? (K_COMPILE_TIME / kColsPerTile) : gridDim.x; + const int32_t target = blockIdx.y * grid_x + blockIdx.x + prefetch_distance_ctas; + if (target < grid_x * static_cast(gridDim.y)) { + constexpr int32_t kLinesPerTile = + kRowsPerTile * kColsPerTile * sizeof(bf16) / kCacheLineBytes; + constexpr int32_t kLinesPerRow = kLinesPerTile / kRowsPerTile; + const int32_t target_y = target / grid_x; + const int32_t target_x = target - target_y * grid_x; + for (int32_t line = tid; line < kLinesPerTile; line += kThreadsPerCta) { + const size_t offset = + (static_cast(target_y * kRowsPerTile + line / kLinesPerRow) * cols + + target_x * kColsPerTile) * + sizeof(bf16) + + (line % kLinesPerRow) * kCacheLineBytes; + ptx::prefetch_l2_evict_last(input + offset); + } + } + } + + // ---- colwise quantized output ------------------------------------------- + ptx::bf16x2 colwise_scales[kWordsPerLane]; +#pragma unroll + for (int32_t k = 0; k < kWordsPerLane; ++k) { + reinterpret_cast(colwise_scales[k]) = column_reciprocal[k]; + } +#pragma unroll + for (int32_t i = 0; i < kRowsPerWarp; ++i) { + uint32_t out_words[kOutWordsPerLane]; + scale_and_convert(tile[i], colwise_scales, out_words); + store_words(out_col + static_cast(i) * cols, out_words, policy); + } +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +namespace { + +// Tile shapes. The wide tile reads 256 bits per lane and suits shallow grids; +// the narrow tile halves that to get twice the CTAs, which hides latency better +// once the grid is deep enough to keep every SM busy regardless. +constexpr int32_t kWideColsPerLane = 16; +constexpr int32_t kNarrowColsPerLane = 8; +constexpr int32_t kWideMinBlocksPerSm = 4; +constexpr int32_t kNarrowMinBlocksPerSm = 6; + +// Thread-block cluster widths. Clustering lets neighbouring CTAs, which read +// adjacent columns of the same rows, share L2 traffic. +constexpr int32_t kWideClusterShallow = 1; +constexpr int32_t kWideClusterDeep = 4; +constexpr int32_t kNarrowCluster = 8; +// Grid size, in CTAs, past which the wide tile switches to the deep cluster. +constexpr int64_t kWideDeepClusterFrom = 4096; +// Grid size, in resident waves, past which the narrow tile takes over. +constexpr int32_t kNarrowFromWaves = 8; + +/*! \brief Launch quantize_bidim_kernel, specializing on K where we can. + * + * The handful of column counts below cover the shapes that matter in practice; + * anything else takes the run-time path, which costs a few address + * computations rather than a recompile. + */ +template +void launch_tiled(const void *input, void *output_rowwise, void *scales_rowwise, + void *output_colwise, void *scales_colwise, int32_t rows, int32_t cols, + int32_t scale_stride_rowwise, int32_t scale_stride_colwise, int32_t cluster_width, + int32_t prefetch_distance_ctas, cudaStream_t stream) { + constexpr int32_t kColsPerTile = THREADS_PER_WARP * COLS_PER_LANE; + const dim3 grid(cols / kColsPerTile, rows / kRowsPerTile); + + // A cluster launch is rejected outright unless its width divides the grid, so + // narrow the requested width to the largest power of two that does. Column + // counts that are not a multiple of the tile width times the cluster width -- + // 7168 columns with the narrow tile, for instance, giving a grid of 28 -- are + // otherwise a hard launch failure rather than a slow path. + int32_t cluster_x = cluster_width; + while (cluster_x > 1 && static_cast(grid.x) % cluster_x != 0) { + --cluster_x; + } + + cudaLaunchAttribute cluster_attr; + cluster_attr.id = cudaLaunchAttributeClusterDimension; + cluster_attr.val.clusterDim.x = cluster_x; + cluster_attr.val.clusterDim.y = 1; + cluster_attr.val.clusterDim.z = 1; + + cudaLaunchConfig_t config = {}; + config.gridDim = grid; + config.blockDim = dim3(kThreadsPerCta); + config.dynamicSmemBytes = 0; + config.stream = stream; + config.attrs = &cluster_attr; + config.numAttrs = 1; + + const uint8_t *in = reinterpret_cast(input); + uint8_t *qrow = reinterpret_cast(output_rowwise); + uint8_t *srow = reinterpret_cast(scales_rowwise); + uint8_t *qcol = reinterpret_cast(output_colwise); + uint8_t *scol = reinterpret_cast(scales_colwise); + +#define NVTE_LAUNCH_BIDIM(K_CONST) \ + do { \ + auto kernel = quantize_bidim_kernel; \ + if (cluster_x > 1) { \ + NVTE_CHECK_CUDA(cudaLaunchKernelEx(&config, kernel, in, qrow, srow, qcol, scol, cols, \ + scale_stride_rowwise, scale_stride_colwise, \ + prefetch_distance_ctas)); \ + } else { \ + kernel<<>>(in, qrow, srow, qcol, scol, cols, \ + scale_stride_rowwise, scale_stride_colwise, \ + prefetch_distance_ctas); \ + NVTE_CHECK_CUDA(cudaGetLastError()); \ + } \ + } while (0) + + switch (cols) { + case 2048: + NVTE_LAUNCH_BIDIM(2048); + break; + case 4096: + NVTE_LAUNCH_BIDIM(4096); + break; + case 7168: + NVTE_LAUNCH_BIDIM(7168); + break; + case 8192: + NVTE_LAUNCH_BIDIM(8192); + break; + case 16384: + NVTE_LAUNCH_BIDIM(16384); + break; + case 32768: + NVTE_LAUNCH_BIDIM(32768); + break; + default: + NVTE_LAUNCH_BIDIM(0); + break; + } +#undef NVTE_LAUNCH_BIDIM +} + +} // namespace + +template +void launch_cast_bidim(const void *input, void *output_rowwise, void *scales_rowwise, + void *output_colwise, void *scales_colwise, int rows, int cols, + int scale_stride_rowwise, int scale_stride_colwise, cudaStream_t stream) { + constexpr int32_t kWideColsPerTile = THREADS_PER_WARP * kWideColsPerLane; + constexpr int32_t kNarrowColsPerTile = THREADS_PER_WARP * kNarrowColsPerLane; + + NVTE_CHECK(rows % kRowsPerTile == 0, "Bidimensional MXFP8 requires the row count (", rows, + ") to be a multiple of the MX block size (", kRowsPerTile, ")."); + NVTE_CHECK(cols % kNarrowColsPerTile == 0, "Bidimensional MXFP8 requires the column count (", + cols, ") to be a multiple of ", kNarrowColsPerTile, "."); + + // One resident wave of CTAs, which is how far ahead the prefetch should run. + const int32_t sm_count = cuda::sm_count(); + + if (cols % kWideColsPerTile == 0) { + const int64_t wide_ctas = static_cast(rows / kRowsPerTile) * (cols / kWideColsPerTile); + const int64_t wide_resident = static_cast(sm_count) * kWideMinBlocksPerSm; + if (wide_ctas < static_cast(kNarrowFromWaves) * wide_resident) { + launch_tiled( + input, output_rowwise, scales_rowwise, output_colwise, scales_colwise, rows, cols, + scale_stride_rowwise, scale_stride_colwise, + wide_ctas >= kWideDeepClusterFrom ? kWideClusterDeep : kWideClusterShallow, + static_cast(wide_resident), stream); + return; + } + } + + // Deep grid, or a column count that only the narrow tile divides. + launch_tiled( + input, output_rowwise, scales_rowwise, output_colwise, scales_colwise, rows, cols, + scale_stride_rowwise, scale_stride_colwise, kNarrowCluster, sm_count * kNarrowMinBlocksPerSm, + stream); +} + +// The MXFP8 output types the specialized dispatch can reach; see hasSpec. +template void launch_cast_bidim(const void *, void *, void *, void *, void *, int, int, + int, int, cudaStream_t); +template void launch_cast_bidim(const void *, void *, void *, void *, void *, int, int, + int, int, cudaStream_t); + +} // namespace specialized +} // namespace quantize_kernel +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine diff --git a/transformer_engine/common/cast/mxfp8/specialized/cast_bidim.h b/transformer_engine/common/cast/mxfp8/specialized/cast_bidim.h new file mode 100644 index 0000000000..12374e7629 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/cast_bidim.h @@ -0,0 +1,60 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file cast_bidim.h + * \brief Entry point for the register-resident bidimensional MXFP8 cast kernel. + */ + +#ifndef TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_CAST_BIDIM_H_ +#define TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_CAST_BIDIM_H_ + +#include + +#include "../../../common.h" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { +namespace specialized { + +/*! \brief Cast BF16 to MXFP8 with both rowwise and colwise scales. + * + * Produces two quantizations of the same tensor: a rowwise one, where 32 + * consecutive elements of a row share a scale, and a colwise one, where 32 + * consecutive elements of a column share a scale. Both outputs keep the + * input's row-major [rows, cols] layout; the colwise result is not transposed. + * + * Unlike the TMA kernel in mxfp8/specialized this one is register-resident: a + * CTA reads its 32-row tile once and drives both passes from registers. See + * cast_bidim.cu for the layout and tuning rationale. + * + * Requires SM 10.0+ (Blackwell), matching MXFP8 support in the rest of TE. + * + * \tparam OType FP8 output type. + * \param[in] input BF16 input, [rows, cols], row-major. + * \param[out] output_rowwise FP8E4M3 rowwise-scaled output, [rows, cols]. + * \param[out] scales_rowwise E8M0 rowwise scales, one per 32 columns. + * \param[out] output_colwise FP8E4M3 colwise-scaled output, [rows, cols]. + * \param[out] scales_colwise E8M0 colwise scales, one per 32 rows. + * \param[in] rows Row count; must be a multiple of 32. + * \param[in] cols Column count; must be a multiple of 256. + * \param[in] scale_stride_rowwise Rowwise scale elements per row. + * \param[in] scale_stride_colwise Colwise scale elements per 32-row band. + * \param[in] stream CUDA stream. + */ +template +void launch_cast_bidim(const void *input, void *output_rowwise, void *scales_rowwise, + void *output_colwise, void *scales_colwise, int rows, int cols, + int scale_stride_rowwise, int scale_stride_colwise, cudaStream_t stream); + +} // namespace specialized +} // namespace quantize_kernel +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_CAST_BIDIM_H_ diff --git a/transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.cu b/transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.cu new file mode 100644 index 0000000000..2d90a80319 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.cu @@ -0,0 +1,436 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file cast_rowwise.cu + * \brief Register-resident rowwise MXFP8 quantization kernel. + */ + +#include + +#include "../../../common.h" +#include "../../../util/ptx.cuh" +#include "../../../util/ptx_arch_spec.cuh" +#include "../../../utils.cuh" +#include "cast_rowwise.h" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { +namespace specialized { + +namespace ptx = transformer_engine::ptx; + +// This kernel casts BF16 to rowwise-scaled MXFP8: every run of 32 consecutive +// elements within a row forms one MX block that shares a single E8M0 scale, +// chosen so the block's largest magnitude lands at the top of the FP8E4M3 +// range. +// +// Unlike the TMA-based kernel in mxfp8/specialized, this one keeps its tile +// entirely in registers. There is no shared memory, no barrier, and no +// two-dimensional tiling: for a row-major tensor whose scale array is also +// contiguous, MX blocks never straddle a row boundary, so the whole tensor is +// just a flat sequence of M*(K/32) independent blocks. That reduces the +// kernel to a pure streaming problem, and what is left to tune is how the +// input and output streams share L2. +// +// A tensor whose scale array is padded (scale_stride > K/32) breaks the flat +// view; those shapes go to quantize_strided_kernel below. + +// Elements in one MX block, all sharing a single E8M0 scale. +constexpr int32_t kBlockElems = 32; + +// Two lanes cooperate on each MX block. A lane's half of a block is 16 BF16 +// values = 32 bytes = one 256-bit load, the widest the ISA offers; splitting +// the block any further would waste load width, and any less would exceed it. +constexpr int32_t kLanesPerBlock = 2; +constexpr int32_t kElemsPerLane = kBlockElems / kLanesPerBlock; + +// Both tensors are addressed as 32-bit words: BF16 packs 2 elements per word, +// FP8 packs 4. All the packed-math PTX below operates on those words. +constexpr int32_t kInElemsPerWord = sizeof(uint32_t) / sizeof(bf16); +// Every MXFP8 output type is a single byte, so a 32-bit word holds four. +constexpr int32_t kOutElemsPerWord = 4; + +constexpr int32_t kInWordsPerLane = kElemsPerLane / kInElemsPerWord; // 8 -> 256-bit load +constexpr int32_t kOutWordsPerLane = kElemsPerLane / kOutElemsPerWord; // 4 -> 128-bit store +constexpr int32_t kInWordsPerBlock = kBlockElems / kInElemsPerWord; // 16 +constexpr int32_t kOutWordsPerBlock = kBlockElems / kOutElemsPerWord; // 8 + +// MX blocks a single warp covers in one pass over its registers. +constexpr int32_t kBlocksPerWarp = THREADS_PER_WARP / kLanesPerBlock; // 16 + +// The block-wide maximum is formed with a single shuffle that swaps a lane +// with its odd/even partner, which only covers a two-lane group. +static_assert(kLanesPerBlock == 2, "A wider lane group would need a multi-step reduction."); + +/*! \brief Launch parameters for one tensor-size regime. + * + * The kernel is bandwidth-bound, so the best configuration tracks how the + * working set compares with L2 rather than the shape itself. These were + * selected by autotuning over a B200 shape sweep; see kTierMaxBytes below for + * the one threshold that has since been re-measured. + */ +struct LaunchConfig { + //! CTA width. Trades occupancy against per-CTA scheduling overhead. + int32_t threads_per_cta; + //! MX blocks each lane pair handles per launch. Raising this unrolls the + //! body, giving more independent loads in flight at the cost of registers. + int32_t blocks_per_lane; + //! Percentage of CTAs that let their input settle in L2 normally; the + //! remainder tag their loads evict_first so the data streams past without + //! displacing anything. 0 streams the entire input. + //! + //! Streaming everything is right once the input dwarfs L2, since nothing + //! would survive to be reused anyway. When the input is only a few times + //! L2, holding part of it back leaves capacity for the output write-back + //! instead of thrashing on input lines. + int32_t l2_cached_cta_percent; +}; + +// Output bytes (one FP8 byte per element, i.e. M*K) separating the regimes. +// +// The first threshold is 12 MiB rather than the 24 MiB the original sweep +// picked. The single-block-per-lane configuration of tier 0 stops paying off +// well before 24 MiB: measured on B200, a 16 MiB tensor ran 8.29 us on tier 0 +// against 6.78 us on tier 1, and a 24 MiB tensor 11.58 us against 9.82 us. +// Both were also slower than the TMA kernel this one replaces, so the tier 0 +// range is cut where the crossover actually lies. +constexpr int64_t kTierMaxBytes[] = {12ll << 20, 48ll << 20, 96ll << 20}; + +constexpr LaunchConfig kTierConfigs[] = { + {/*threads_per_cta=*/256, /*blocks_per_lane=*/1, /*l2_cached_cta_percent=*/0}, + {/*threads_per_cta=*/256, /*blocks_per_lane=*/2, /*l2_cached_cta_percent=*/0}, + {/*threads_per_cta=*/128, /*blocks_per_lane=*/2, /*l2_cached_cta_percent=*/40}, + {/*threads_per_cta=*/256, /*blocks_per_lane=*/2, /*l2_cached_cta_percent=*/40}, +}; +constexpr int32_t kNumTiers = sizeof(kTierConfigs) / sizeof(kTierConfigs[0]); +static_assert(kNumTiers == sizeof(kTierMaxBytes) / sizeof(kTierMaxBytes[0]) + 1, + "Each size threshold must separate two tiers."); + +namespace { + +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +/*! \brief Reduce eight BF16 pairs to the largest magnitude among them. */ +__device__ __forceinline__ ptx::bf16x2 block_half_amax(const uint32_t (&words)[kInWordsPerLane]) { + const ptx::bf16x2 *pairs = reinterpret_cast(words); + + // Balanced tree: depth 3 instead of the 7 of a serial chain, so the + // independent maxima issue back to back. + ptx::bf16x2 level[kInWordsPerLane / 2]; +#pragma unroll + for (int32_t i = 0; i < kInWordsPerLane / 2; ++i) { + ptx::abs_max_2x(level[i], pairs[2 * i], pairs[2 * i + 1]); + } +#pragma unroll + for (int32_t i = 0; i < kInWordsPerLane / 4; ++i) { + ptx::abs_max_2x(level[i], level[i], level[i + kInWordsPerLane / 4]); + } + ptx::bf16x2 result; + ptx::abs_max_2x(result, level[0], level[1]); + return result; +} + +/*! \brief Widen a BF16 pair's larger magnitude to FP32. + * + * `max.xorsign.abs` keeps the magnitude of the larger operand but sets the + * result sign to the XOR of the input signs, so an accumulator built from it + * can come out negative. Only the magnitude means anything for a scale, and + * feeding a negative value to the unsigned E8M0 conversion would saturate it + * to zero, so the sign is cleared here. + */ +__device__ __forceinline__ float pair_amax_to_float(ptx::bf16x2 pair) { + const uint32_t bits = reinterpret_cast(pair); + // Fold the two halves against each other, then keep the low BF16 sans sign. + const uint32_t folded = __byte_perm(bits, bits, 0x1032); + ptx::bf16x2 a, b; + reinterpret_cast(a) = bits; + reinterpret_cast(b) = folded; + ptx::bf16x2 wide; + ptx::abs_max_2x(wide, a, b); + const uint32_t magnitude = reinterpret_cast(wide) & 0x7FFFu; + return __int_as_float(magnitude << 16); +} + +/*! \brief Scale and convert one lane's 16 BF16 values into 16 FP8E4M3 bytes. */ +template +__device__ __forceinline__ void scale_and_convert(const uint32_t (&in)[kInWordsPerLane], + ptx::bf16x2 scale_reciprocal, + uint32_t (&out)[kOutWordsPerLane]) { +#pragma unroll + for (int32_t i = 0; i < kOutWordsPerLane; ++i) { + ptx::mul_cvt_4x(reinterpret_cast &>(out[i]), + reinterpret_cast(in[2 * i]), scale_reciprocal); + } +} + +/*! \brief Quantize one whole MX block held by a single thread. + * + * Used by the remainder and strided kernels, which handle far too little data + * to be worth the two-lane split of the main kernel. + */ +template +__device__ __forceinline__ void quantize_one_block(const uint32_t *__restrict__ in, + uint32_t *__restrict__ out, + e8m0_t *__restrict__ scale, + uint64_t output_policy) { + uint32_t words[kInWordsPerBlock]; + ptx::ld_global_nc_b32x8(reinterpret_cast(words[0]), in); + ptx::ld_global_nc_b32x8(reinterpret_cast(words[kInWordsPerLane]), + in + kInWordsPerLane); + + ptx::bf16x2 amax_lo = block_half_amax(reinterpret_cast(words[0])); + ptx::bf16x2 amax_hi = + block_half_amax(reinterpret_cast(words[kInWordsPerLane])); + ptx::bf16x2 amax_pair; + ptx::abs_max_2x(amax_pair, amax_lo, amax_hi); + + const float amax = pair_amax_to_float(amax_pair); + const e8m0_t biased_exponent = ptx::float_to_e8m0(amax * Quantized_Limits::max_norm_rcp); + *scale = biased_exponent; + + const ptx::bf16x2 scale_reciprocal = ptx::exp2f_rcp_2x(biased_exponent); + uint32_t out_words[kOutWordsPerBlock]; + scale_and_convert(reinterpret_cast(words[0]), scale_reciprocal, + reinterpret_cast(out_words[0])); + scale_and_convert(reinterpret_cast(words[kInWordsPerLane]), + scale_reciprocal, + reinterpret_cast(out_words[kOutWordsPerLane])); + + ptx::st_global_b32x4(out, reinterpret_cast(out_words[0]), output_policy); + ptx::st_global_b32x4(out + kOutWordsPerLane, + reinterpret_cast(out_words[kOutWordsPerLane]), + output_policy); +} + +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +} // namespace + +/*! \brief Quantize a contiguous run of MX blocks, two lanes per block. + * + * \tparam PARTIAL_L2_CACHING When false every CTA streams its input, which + * the ISA expresses as a static load modifier and + * so costs no policy register. When true the + * decision varies per CTA and needs a runtime + * policy token. See LaunchConfig. + * + * \param[in] input BF16 input, viewed as 32-bit words. + * \param[out] output FP8E4M3 output, viewed as 32-bit words. + * \param[out] scales One E8M0 byte per MX block. + * \param[in] first_streaming_cta CTAs at or above this index stream their + * input; earlier ones cache normally. Only + * read when PARTIAL_L2_CACHING is true. + */ +template +__global__ void __launch_bounds__(THREADS_PER_CTA) + quantize_contiguous_kernel(const uint32_t *__restrict__ input, uint32_t *__restrict__ output, + e8m0_t *__restrict__ scales, uint32_t first_streaming_cta) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int32_t kWarpsPerCta = THREADS_PER_CTA / THREADS_PER_WARP; + constexpr int32_t kBlocksPerWarpPass = kBlocksPerWarp * BLOCKS_PER_LANE; + + const int32_t lane = threadIdx.x % THREADS_PER_WARP; + const int64_t warp_id = + static_cast(blockIdx.x) * kWarpsPerCta + threadIdx.x / THREADS_PER_WARP; + const int64_t first_block = warp_id * kBlocksPerWarpPass; + + // Output is marked evict_last so its lines linger long enough to coalesce on + // write-back. + const uint64_t output_policy = ptx::create_l2_policy_evict_last(); + + // Lanes pair up as (even, odd); the even lane of each pair owns the scale. + const int32_t block_in_warp = lane / kLanesPerBlock; + const bool owns_scale = (lane % kLanesPerBlock) == 0; + + uint32_t in_words[BLOCKS_PER_LANE][kInWordsPerLane]; + if constexpr (PARTIAL_L2_CACHING) { + const uint64_t input_policy = + ptx::create_l2_policy_evict_first(blockIdx.x >= first_streaming_cta ? 1.0f : 0.0f); +#pragma unroll + for (int32_t u = 0; u < BLOCKS_PER_LANE; ++u) { + const int64_t group_base = first_block + static_cast(u) * kBlocksPerWarp; + ptx::ld_global_nc_b32x8(in_words[u], + input + group_base * kInWordsPerBlock + lane * kInWordsPerLane, + input_policy); + } + } else { +#pragma unroll + for (int32_t u = 0; u < BLOCKS_PER_LANE; ++u) { + const int64_t group_base = first_block + static_cast(u) * kBlocksPerWarp; + ptx::ld_global_nc_evict_first_b32x8( + in_words[u], input + group_base * kInWordsPerBlock + lane * kInWordsPerLane); + } + } + +#pragma unroll + for (int32_t u = 0; u < BLOCKS_PER_LANE; ++u) { + const int64_t group_base = first_block + static_cast(u) * kBlocksPerWarp; + + // Each lane reduces its own half, then swaps with its partner so both + // arrive at the block-wide maximum. + ptx::bf16x2 half_amax = block_half_amax(in_words[u]); + ptx::bf16x2 partner; + reinterpret_cast(partner) = + __shfl_xor_sync(0xFFFFFFFFu, reinterpret_cast(half_amax), /*laneMask=*/1); + ptx::bf16x2 block_amax; + ptx::abs_max_2x(block_amax, half_amax, partner); + + const e8m0_t biased_exponent = + ptx::float_to_e8m0(pair_amax_to_float(block_amax) * Quantized_Limits::max_norm_rcp); + if (owns_scale) { + scales[group_base + block_in_warp] = biased_exponent; + } + + uint32_t out_words[kOutWordsPerLane]; + scale_and_convert(in_words[u], ptx::exp2f_rcp_2x(biased_exponent), out_words); + ptx::st_global_b32x4(output + group_base * kOutWordsPerBlock + lane * kOutWordsPerLane, + out_words, output_policy); + } +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +/*! \brief Quantize the MX blocks left over when the block count does not + * divide evenly among the main kernel's CTAs. One block per thread. */ +template +__global__ void __launch_bounds__(128) + quantize_remainder_kernel(const uint32_t *__restrict__ input, uint32_t *__restrict__ output, + e8m0_t *__restrict__ scales, int64_t first_block, + int64_t num_blocks) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + const int64_t block = first_block + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= num_blocks) { + return; + } + quantize_one_block(input + block * kInWordsPerBlock, output + block * kOutWordsPerBlock, + scales + block, ptx::create_l2_policy_evict_last()); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +/*! \brief Quantize a tensor whose scale rows are padded. + * + * With scale_stride > K/32 the scale array is no longer a flat image of the + * block sequence, so blocks are indexed two-dimensionally. One block per + * thread; grid.y walks the rows. + */ +template +__global__ void __launch_bounds__(128) + quantize_strided_kernel(const uint32_t *__restrict__ input, uint32_t *__restrict__ output, + e8m0_t *__restrict__ scales, int32_t blocks_per_row, + int32_t scale_stride) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + const int32_t block_in_row = blockIdx.x * blockDim.x + threadIdx.x; + if (block_in_row >= blocks_per_row) { + return; + } + const int64_t block = static_cast(blockIdx.y) * blocks_per_row + block_in_row; + quantize_one_block(input + block * kInWordsPerBlock, output + block * kOutWordsPerBlock, + scales + static_cast(blockIdx.y) * scale_stride + block_in_row, + ptx::create_l2_policy_evict_last()); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +namespace { + +//! Threads per CTA for the two block-per-thread helper kernels. +constexpr int32_t kHelperThreads = 128; + +/*! \brief Launch quantize_contiguous_kernel for a configuration resolved at + * run time, instantiating only the combinations the tier table uses. */ +template +void launch_contiguous(const LaunchConfig &config, int64_t grid, uint32_t first_streaming_cta, + const uint32_t *input, uint32_t *output, e8m0_t *scales, + cudaStream_t stream) { + const dim3 blocks(static_cast(grid)); + const dim3 threads(static_cast(config.threads_per_cta)); + + const bool partial = config.l2_cached_cta_percent != 0; + + if (config.threads_per_cta == 256 && config.blocks_per_lane == 1 && !partial) { + quantize_contiguous_kernel + <<>>(input, output, scales, first_streaming_cta); + } else if (config.threads_per_cta == 256 && config.blocks_per_lane == 2 && !partial) { + quantize_contiguous_kernel + <<>>(input, output, scales, first_streaming_cta); + } else if (config.threads_per_cta == 128 && config.blocks_per_lane == 2 && partial) { + quantize_contiguous_kernel + <<>>(input, output, scales, first_streaming_cta); + } else if (config.threads_per_cta == 256 && config.blocks_per_lane == 2 && partial) { + quantize_contiguous_kernel + <<>>(input, output, scales, first_streaming_cta); + } else { + NVTE_ERROR("No quantize_contiguous_kernel instantiation for ", config.threads_per_cta, + " threads, ", config.blocks_per_lane, " blocks per lane, partial L2 caching ", + partial, "."); + } +} + +} // namespace + +template +void launch_cast_rowwise(const void *input, void *output, void *scales, int rows, int cols, + int scale_stride, cudaStream_t stream) { + NVTE_CHECK(cols % kBlockElems == 0, "Rowwise MXFP8 requires the column count (", cols, + ") to be a multiple of the MX block size (", kBlockElems, ")."); + + const int32_t blocks_per_row = cols / kBlockElems; + const uint32_t *in = reinterpret_cast(input); + uint32_t *out = reinterpret_cast(output); + e8m0_t *scale_out = reinterpret_cast(scales); + + // A padded scale array breaks the flat block view the fast path relies on. + if (scale_stride != blocks_per_row) { + const dim3 grid(DIVUP(blocks_per_row, kHelperThreads), rows); + quantize_strided_kernel + <<>>(in, out, scale_out, blocks_per_row, scale_stride); + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + + const int64_t num_blocks = static_cast(rows) * blocks_per_row; + const int64_t output_bytes = static_cast(rows) * cols; + + int32_t tier = 0; + while (tier < kNumTiers - 1 && output_bytes > kTierMaxBytes[tier]) { + ++tier; + } + const LaunchConfig config = kTierConfigs[tier]; + + // Every CTA covers a whole number of MX blocks; the leftovers, if any, go to + // the remainder kernel rather than costing the main kernel a bounds check. + const int64_t blocks_per_cta = + static_cast(config.threads_per_cta) / kLanesPerBlock * config.blocks_per_lane; + const int64_t grid = num_blocks / blocks_per_cta; + + if (grid > 0) { + const uint32_t first_streaming_cta = + static_cast(grid * config.l2_cached_cta_percent / 100); + launch_contiguous(config, grid, first_streaming_cta, in, out, scale_out, stream); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + + const int64_t blocks_done = grid * blocks_per_cta; + if (blocks_done < num_blocks) { + const int64_t remaining = num_blocks - blocks_done; + quantize_remainder_kernel + <<(kHelperThreads)), kHelperThreads, 0, stream>>>( + in, out, scale_out, blocks_done, num_blocks); + NVTE_CHECK_CUDA(cudaGetLastError()); + } +} + +// The MXFP8 output types the specialized dispatch can reach; see hasSpec. +template void launch_cast_rowwise(const void *, void *, void *, int, int, int, + cudaStream_t); +template void launch_cast_rowwise(const void *, void *, void *, int, int, int, + cudaStream_t); + +} // namespace specialized +} // namespace quantize_kernel +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine diff --git a/transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.h b/transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.h new file mode 100644 index 0000000000..cea6b60bad --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.h @@ -0,0 +1,56 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file cast_rowwise.h + * \brief Entry point for the register-resident rowwise MXFP8 cast kernel. + */ + +#ifndef TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_CAST_ROWWISE_H_ +#define TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_CAST_ROWWISE_H_ + +#include + +#include "../../../common.h" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { +namespace specialized { + +/*! \brief Cast BF16 to rowwise-scaled MXFP8. + * + * Every run of 32 consecutive elements in a row forms one MX block sharing a + * single E8M0 scale. This is the register-resident member of the specialized + * cast-only family: it keeps its tile in registers rather than staging it, and + * is 10-20% faster than quantize_mxfp8_kernel_cast_only on that path. See + * cast_rowwise.cu for the layout and tuning rationale. + * + * Instantiated for the MXFP8 output types the specialized dispatch can reach, + * fp8e4m3 and fp8e5m2. Requires SM 10.0+ (Blackwell), matching MXFP8 support + * in the rest of TE. + * + * \tparam OType FP8 output type. + * \param[in] input BF16 input, [rows, cols], row-major. + * \param[out] output FP8 output, [rows, cols], row-major. + * \param[out] scales E8M0 scales, one byte per MX block. + * \param[in] rows Number of rows. + * \param[in] cols Number of columns; must be a multiple of 32. + * \param[in] scale_stride Scale elements per row. Equals cols/32 for a + * packed scale array, or more when it is padded. + * \param[in] stream CUDA stream. + */ +template +void launch_cast_rowwise(const void *input, void *output, void *scales, int rows, int cols, + int scale_stride, cudaStream_t stream); + +} // namespace specialized +} // namespace quantize_kernel +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_MXFP8_SPECIALIZED_CAST_ROWWISE_H_ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index f48224c365..2e3b12ca62 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -782,6 +782,81 @@ __device__ __forceinline__ fp16 get_amax(fp16 a, fp16 b) { #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +// Reciprocal of an E8M0 scale, broadcast into both halves of a BF16 pair so a +// packed multiply can scale two elements per instruction. +// +// A BF16 lane holds its exponent in bits 7..14, so subtracting biased_exp<<7 +// from the encoding of 2^127 yields 2^(127-biased_exp). Applying that to both +// lanes at once costs one multiply and one subtract, replacing a scalar +// conversion followed by a broadcast. +// +// Valid for biased_exp <= 253. Every scale derived from finite BF16 or FP16 +// input stays well inside that bound -- the largest finite BF16 yields +// biased_exp 247. The two excluded encodings are 254, whose reciprocal is +// subnormal, and 255, which marks NaN; in both cases the E8M0 scale byte +// itself carries the marker, so a consumer still reads the block correctly. +__device__ __forceinline__ bf16x2 exp2f_rcp_2x(e8m0_t biased_exp) { + // Encoding of 2^127 in both BF16 lanes, and one exponent step in both lanes. + constexpr uint32_t kTwoPow127Pair = 0x7F007F00u; + constexpr uint32_t kExponentStepPair = 0x00800080u; + // biased_exp <= 255 keeps biased_exp<<7 within 16 bits, so neither lane + // borrows into the other. + bf16x2 result; + reinterpret_cast(result) = kTwoPow127Pair - biased_exp * kExponentStepPair; + return result; +} + +// Scale two BF16 pairs by independent scales and pack the four results into a +// single FP8E4M3 word. The mul_cvt_4x overload below shares one scale across +// all four elements, which is what a rowwise MX block wants; a colwise block +// gives every column pair its own scale, and this keeps that case to one +// instruction sequence instead of two conversions and a merge. +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const bf16x2 &in0, const bf16x2 &scale0, + const bf16x2 &in1, const bf16x2 &scale1) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile( + "{\n\t" + ".reg.b32 y0, y1; \n\t" + ".reg.b16 z0, z1; \n\t" + "mul.rn.bf16x2 y0, %1, %2; \n\t" + "mul.rn.bf16x2 y1, %3, %4; \n\t" + "cvt.rn.satfinite.e4m3x2.bf16x2 z0, y0; \n\t" + "cvt.rn.satfinite.e4m3x2.bf16x2 z1, y1; \n\t" + "mov.b32 %0, {z0, z1}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in0)), + "r"(reinterpret_cast(scale0)), + "r"(reinterpret_cast(in1)), + "r"(reinterpret_cast(scale1))); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, const bf16x2 &in0, const bf16x2 &scale0, + const bf16x2 &in1, const bf16x2 &scale1) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile( + "{\n\t" + ".reg.b32 y0, y1; \n\t" + ".reg.b16 z0, z1; \n\t" + "mul.rn.bf16x2 y0, %1, %2; \n\t" + "mul.rn.bf16x2 y1, %3, %4; \n\t" + "cvt.rn.satfinite.e5m2x2.bf16x2 z0, y0; \n\t" + "cvt.rn.satfinite.e5m2x2.bf16x2 z1, y1; \n\t" + "mov.b32 %0, {z0, z1}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in0)), + "r"(reinterpret_cast(scale0)), + "r"(reinterpret_cast(in1)), + "r"(reinterpret_cast(scale1))); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + __device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const bf16x4 &in, const bf16x2 scale) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) #if (defined CUDA_VERSION) && (CUDA_VERSION >= 13010) @@ -1569,6 +1644,151 @@ __device__ __forceinline__ void st_shared_b64(fp4e2m1x2 *__restrict__ dst_smem, asm volatile("st.shared.b64 [%0], %1;" : : "r"(dst_smem_ptr), "l"(fp4_pack_x16)); } #endif +// +// L2 cache-eviction policies for global memory accesses. +// +// A cache policy is an opaque 64-bit token produced by `createpolicy` and +// consumed by the `.L2::cache_hint` variants of `ld`/`st`. It lets a kernel +// tell L2 how to prioritise one access stream relative to another, which is +// what makes the difference for bandwidth-bound kernels that move far more +// data than L2 can hold. +// + +// Keep the accessed lines resident in L2 in preference to others. Use for +// data that will be re-read soon, and for stores whose write-back coalesces +// better when the line lingers in L2. +__device__ __forceinline__ uint64_t create_l2_policy_evict_last() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;" : "=l"(policy)); + return policy; +#else + NVTE_DEVICE_ERROR("L2 cache policies are only supported on SM 8.0+."); + return 0; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) +} + +// Stream `fraction` of the accessed lines past L2 without displacing resident +// data. Use for data that is read exactly once. `fraction` is in [0, 1]; +// 0.0 leaves the access with default caching behaviour, letting a kernel dial +// in how much of its input is allowed to occupy L2. +__device__ __forceinline__ uint64_t create_l2_policy_evict_first(float fraction) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + uint64_t policy; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, %1;" + : "=l"(policy) + : "f"(fraction)); + return policy; +#else + NVTE_DEVICE_ERROR("L2 cache policies are only supported on SM 8.0+."); + return 0; +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) +} + +// Non-coherent (read-only / `__ldg`-style) 256-bit global load. This is the +// widest load the ISA offers and keeps the number of in-flight requests, and +// hence the latency that must be hidden, as low as possible. +__device__ __forceinline__ void ld_global_nc_b32x8(uint32_t (&dst)[8], const void *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("ld.global.nc.v8.b32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];" + : "=r"(dst[0]), "=r"(dst[1]), "=r"(dst[2]), "=r"(dst[3]), "=r"(dst[4]), "=r"(dst[5]), + "=r"(dst[6]), "=r"(dst[7]) + : "l"(src)); +#else + NVTE_DEVICE_ERROR("ld_global_nc_b32x8 is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// As above, but streaming the lines past L2 rather than letting them displace +// resident data. Equivalent to tagging the load with a `create_l2_policy_ +// evict_first(1.0)` hint, without needing to materialise the policy token. +__device__ __forceinline__ void ld_global_nc_evict_first_b32x8(uint32_t (&dst)[8], + const void *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("ld.global.nc.L2::evict_first.v8.b32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];" + : "=r"(dst[0]), "=r"(dst[1]), "=r"(dst[2]), "=r"(dst[3]), "=r"(dst[4]), "=r"(dst[5]), + "=r"(dst[6]), "=r"(dst[7]) + : "l"(src)); +#else + NVTE_DEVICE_ERROR("ld_global_nc_evict_first_b32x8 is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// As above, tagged with an L2 cache policy from `create_l2_policy_*`. +__device__ __forceinline__ void ld_global_nc_b32x8(uint32_t (&dst)[8], const void *src, + uint64_t l2_policy) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("ld.global.nc.L2::cache_hint.v8.b32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8], %9;" + : "=r"(dst[0]), "=r"(dst[1]), "=r"(dst[2]), "=r"(dst[3]), "=r"(dst[4]), "=r"(dst[5]), + "=r"(dst[6]), "=r"(dst[7]) + : "l"(src), "l"(l2_policy)); +#else + NVTE_DEVICE_ERROR("ld_global_nc_b32x8 is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// 128-bit global store tagged with an L2 cache policy. +__device__ __forceinline__ void st_global_b32x4(void *dst, const uint32_t (&src)[4], + uint64_t l2_policy) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("st.global.L2::cache_hint.v4.b32 [%0], {%1,%2,%3,%4}, %5;" + : + : "l"(dst), "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), "l"(l2_policy) + : "memory"); +#else + NVTE_DEVICE_ERROR("st_global_b32x4 is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// Non-coherent 128-bit global load tagged with an L2 cache policy. +__device__ __forceinline__ void ld_global_nc_b32x4(uint32_t (&dst)[4], const void *src, + uint64_t l2_policy) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("ld.global.nc.L2::cache_hint.v4.b32 {%0,%1,%2,%3}, [%4], %5;" + : "=r"(dst[0]), "=r"(dst[1]), "=r"(dst[2]), "=r"(dst[3]) + : "l"(src), "l"(l2_policy)); +#else + NVTE_DEVICE_ERROR("ld_global_nc_b32x4 is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// 64-bit global store tagged with an L2 cache policy. +__device__ __forceinline__ void st_global_b32x2(void *dst, const uint32_t (&src)[2], + uint64_t l2_policy) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("st.global.L2::cache_hint.v2.b32 [%0], {%1,%2}, %3;" + : + : "l"(dst), "r"(src[0]), "r"(src[1]), "l"(l2_policy) + : "memory"); +#else + NVTE_DEVICE_ERROR("st_global_b32x2 is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// Single-byte global store tagged with an L2 cache policy. Scale bytes are +// written one at a time by a subset of lanes, and tagging them with the same +// policy as the bulk streams keeps every access from this kernel consistent. +__device__ __forceinline__ void st_global_b8(void *dst, uint8_t value, uint64_t l2_policy) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + const uint16_t widened = value; + asm volatile("st.global.L2::cache_hint.b8 [%0], %1, %2;" + : + : "l"(dst), "h"(widened), "l"(l2_policy) + : "memory"); +#else + NVTE_DEVICE_ERROR("st_global_b8 is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// Bring one global cache line into L2 ahead of the demand load that needs it, +// marked evict_last so it matches the policy the demand load will use. +__device__ __forceinline__ void prefetch_l2_evict_last(const void *addr) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + asm volatile("prefetch.global.L2::evict_last [%0];" : : "l"(addr)); +#else + NVTE_DEVICE_ERROR("prefetch_l2_evict_last is only supported on SM 8.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) +} } // namespace ptx namespace {