From f19950f18bffe14465cbb63910f7fa90df302cfb Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 1 Sep 2026 14:15:06 -0700 Subject: [PATCH] Add register-resident MXFP8 cast-only kernels for rowwise and bidimensional The specialized cast-only path currently stages its tile: the rowwise kernel through shared memory, the bidimensional one through a TMA pipeline. For a cast with no bias or activation there is nothing to stage for -- the tile is read once and consumed immediately -- so both can be done entirely in registers instead. Adds two kernels under cast/mxfp8/specialized/: cast_rowwise.cu Two lanes cooperate on each 32-element MX block, a lane's half being exactly one 256-bit load. No shared memory and no barrier at all; the tensor is a flat sequence of independent MX blocks whenever the scale array is packed. cast_bidim.cu A CTA owns a 32-row band, which is exactly the colwise block height, so the colwise reduction closes inside the CTA and the tile drives both passes from registers. Shared memory is used only for the cross-warp column fold. Both pick their launch configuration from a documented size-tier table and use L2 eviction policies plus a software prefetch one resident CTA-wave ahead. New PTX wrappers in util/ptx.cuh, none of which had a TE equivalent: L2 cache-policy creation, 256/128-bit non-coherent loads and 128/64/8-bit stores carrying a policy, an L2 prefetch, a packed BF16 reciprocal-scale helper, and a mul_cvt_4x overload taking two BF16 pairs with independent scales, which the colwise pass needs since every column pair has its own scale. Dispatch routes to these kernels for BF16 input with a packed scale layout, and falls through to the existing kernels otherwise. hasSpec already establishes cast-only, so no further condition on the fused-op flags is needed. FP16 input and GEMM-swizzled scale layouts are not yet covered and remain on the existing path. Measured on GR10x (CC 10.7, CUDA 13.4), medians of 20 with the clocks warmed: rowwise, geomean over 12 small shapes 10.949 us -> 9.913 us rowwise, geomean over 6 standard shapes 19.411 us -> 17.934 us bidimensional, geomean over 6 shapes 29.595 us -> 25.730 us The margin narrows as the shapes grow and the kernels become DRAM-bound; at 65536x16384 rowwise reaches 89% of peak DRAM throughput, where little headroom remains for anyone. Registers drop from 48 to 24-30 (rowwise) and from 70 to 46-64 (bidimensional), with occupancy rising correspondingly. All 1541 MXFP8 tests in tests/cpp/operator pass. qa/format.sh is clean. --- transformer_engine/common/CMakeLists.txt | 2 + .../common/cast/mxfp8/quantize_mxfp8.cuh | 34 + .../cast/mxfp8/specialized/cast_bidim.cu | 643 ++++++++++++++++++ .../cast/mxfp8/specialized/cast_bidim.h | 60 ++ .../cast/mxfp8/specialized/cast_rowwise.cu | 436 ++++++++++++ .../cast/mxfp8/specialized/cast_rowwise.h | 56 ++ transformer_engine/common/util/ptx.cuh | 220 ++++++ 7 files changed, 1451 insertions(+) create mode 100644 transformer_engine/common/cast/mxfp8/specialized/cast_bidim.cu create mode 100644 transformer_engine/common/cast/mxfp8/specialized/cast_bidim.h create mode 100644 transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.cu create mode 100644 transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.h 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 {