From ac1bf616280c2f05f15500c843fb29bd012919b4 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 8 Oct 2026 09:19:07 -0700 Subject: [PATCH 1/2] Keep several 32-byte loads in flight in the amax kernel amax_kernel gave each thread a single 32-byte load and then retired the whole CTA through a reduction and an atomic, so it was latency-bound well under half of DRAM bandwidth. Replace it with a persistent kernel: a grid sized to fill the GPU once strides over the tensor issuing four independent 32-byte loads per lane per iteration, and every CTA commits once. 16-bit inputs fold with packed abs-max on SM 8.9+. Elements before the first and after the last 32-byte boundary are read one at a time, so any alignment and size is served by the same kernel. The maximum of |x| is exact and order-independent, so results are bit-identical to the old kernel. This is the first pass of every NVFP4 quantize and also serves FP8 current scaling. Signed-off-by: Teddy Do --- .../common/recipe/current_scaling.cu | 191 ++++++++++++------ 1 file changed, 128 insertions(+), 63 deletions(-) diff --git a/transformer_engine/common/recipe/current_scaling.cu b/transformer_engine/common/recipe/current_scaling.cu index 3ba51dba754..bd39de13a49 100644 --- a/transformer_engine/common/recipe/current_scaling.cu +++ b/transformer_engine/common/recipe/current_scaling.cu @@ -12,15 +12,14 @@ #include "../cast/fp8/group_amax_fp8.cuh" #include "../common.h" +#include "../util/cuda_runtime.h" #include "../util/logging.h" -#include "../util/vectorized_pointwise.h" +#include "../util/ptx.cuh" #include "recipe_common.cuh" namespace transformer_engine { namespace { -constexpr int amax_kernel_threads = 512; - __launch_bounds__(1) __global__ void zero_amax_kernel(float *amax_ptr, const float *noop_ptr) { if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { return; @@ -28,48 +27,121 @@ __launch_bounds__(1) __global__ void zero_amax_kernel(float *amax_ptr, const flo *amax_ptr = 0; } -template +// Amax is a pure streaming reduction, so the kernel is built to keep as many +// loads in flight as possible: a persistent grid sized to fill the GPU once +// strides over the tensor in 32-byte chunks, each lane issuing amax_kernel_loads +// independent chunk loads per iteration before folding any of them, and every +// CTA reduces and commits to the output once. The maximum of |x| is exact and +// order-independent, so the result does not depend on the launch shape. +constexpr int amax_kernel_threads = 256; +constexpr int amax_kernel_loads = 4; +constexpr int amax_chunk_bytes = 32; +constexpr int amax_chunk_words = amax_chunk_bytes / sizeof(uint32_t); + +__device__ __forceinline__ void load_amax_chunk(uint32_t (&words)[amax_chunk_words], + const uint32_t *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + ptx::ld_global_nc_b32x8(words, src); +#else + const uint4 lo = __ldg(reinterpret_cast(src)); + const uint4 hi = __ldg(reinterpret_cast(src) + 1); + words[0] = lo.x, words[1] = lo.y, words[2] = lo.z, words[3] = lo.w; + words[4] = hi.x, words[5] = hi.y, words[6] = hi.z, words[7] = hi.w; +#endif +} + +/*! \brief Atomically fold max(|input|) into *amax. + * + * input + head must be 32-byte aligned. The head elements before it and the tail + * elements after its num_chunks whole chunks are read one at a time. + */ +template __launch_bounds__(amax_kernel_threads) __global__ - void amax_kernel(const InputType *input, float *amax, const size_t N, - const size_t num_aligned_elements, const float *noop_ptr) { + void amax_kernel(const InputType *__restrict__ input, float *__restrict__ amax, const int head, + const int64_t num_chunks, const int tail, const float *__restrict__ noop_ptr) { if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { return; } + constexpr int kElemsPerChunk = amax_chunk_bytes / sizeof(InputType); +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) + constexpr bool kPacked = + std::is_same_v || std::is_same_v; +#else + constexpr bool kPacked = false; +#endif + using Pair = + std::conditional_t, ptx::bf16x2, ptx::fp16x2>; + + const int64_t stride = static_cast(gridDim.x) * amax_kernel_threads; + const int64_t tid = static_cast(blockIdx.x) * amax_kernel_threads + threadIdx.x; + const uint32_t *chunks = reinterpret_cast(input + head); + + float max = 0.f; + Pair max_pair; + reinterpret_cast(max_pair) = 0; + auto fold = [&](const uint32_t(&words)[amax_chunk_words]) { + if constexpr (kPacked) { + // Packed abs-max tree over the chunk's 8 element pairs. + const Pair *pairs = reinterpret_cast(words); + Pair level[amax_chunk_words / 2]; +#pragma unroll + for (int i = 0; i < amax_chunk_words / 2; ++i) { + ptx::abs_max_2x(level[i], pairs[2 * i], pairs[2 * i + 1]); + } + ptx::abs_max_2x(level[0], level[0], level[1]); + ptx::abs_max_2x(level[2], level[2], level[3]); + ptx::abs_max_2x(level[0], level[0], level[2]); + ptx::abs_max_2x(max_pair, max_pair, level[0]); + } else { + const InputType *elems = reinterpret_cast(words); +#pragma unroll + for (int i = 0; i < kElemsPerChunk; ++i) { + max = fmaxf(fabsf(static_cast(elems[i])), max); + } + } + }; - VectorizedLoader loader(input, N); - InputType max = 0.f; - const int warp_id = threadIdx.x / THREADS_PER_WARP; - const size_t M = num_aligned_elements; + // Fewer than 2 chunks' worth of elements, so the first CTA covers them in one pass. + // Read before the chunks so their latency overlaps the main loop's. + if (tid < head) { + max = fmaxf(fabsf(static_cast(input[tid])), max); + } + if (tid < tail) { + const InputType *tail_ptr = input + head + num_chunks * kElemsPerChunk; + max = fmaxf(fabsf(static_cast(tail_ptr[tid])), max); + } - for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < M; tid += gridDim.x * blockDim.x) { - loader.load(tid, N); + int64_t chunk = tid; + for (; chunk + (amax_kernel_loads - 1) * stride < num_chunks; + chunk += amax_kernel_loads * stride) { + uint32_t words[amax_kernel_loads][amax_chunk_words]; #pragma unroll - for (int i = 0; i < nvec; ++i) { - const InputType val = static_cast(loader.separate()[i]); - __builtin_assume(max >= InputType{0.f}); - if constexpr (std::is_same_v) { -#if __CUDA_ARCH__ >= 800 - max = __hmax(__habs(val), max); -#else // Turing - max = static_cast<__nv_bfloat16>( - fmaxf(fabsf(static_cast(val)), static_cast(max))); -#endif - } else if constexpr (std::is_same_v) { - max = __hmax(__habs(val), max); - } else { - max = fmaxf(fabsf(val), max); - } + for (int u = 0; u < amax_kernel_loads; ++u) { + load_amax_chunk(words[u], chunks + (chunk + u * stride) * amax_chunk_words); } +#pragma unroll + for (int u = 0; u < amax_kernel_loads; ++u) { + fold(words[u]); + } + } + for (; chunk < num_chunks; chunk += stride) { + uint32_t words[amax_chunk_words]; + load_amax_chunk(words, chunks + chunk * amax_chunk_words); + fold(words); } - // Reduce amax over block - max = reduce_max(max, warp_id); + if constexpr (kPacked) { + // max.xorsign.abs may leave a sign behind; only the magnitude is wanted. + max = fmaxf( + max, fmaxf(fabsf(static_cast(max_pair.x)), fabsf(static_cast(max_pair.y)))); + } + max = reduce_max(max, threadIdx.x / THREADS_PER_WARP); if (threadIdx.x == 0) { atomicMaxFloat(amax, max); } } -template +template void launch_amax_kernel(const InputType *input, float *amax, const size_t N, const float *noop_ptr, cudaStream_t stream) { // Zero out amax so we can update with atomic max @@ -81,36 +153,29 @@ void launch_amax_kernel(const InputType *input, float *amax, const size_t N, con return; } - // Figure out alignment - auto align = CheckAlignment(N, nvec, input); - size_t num_aligned_elements = get_num_aligned_elements(input, N, nvec, sizeof(InputType)); - - // Figure out CUDA blocks - constexpr size_t threads = amax_kernel_threads; - size_t num_blocks = DIVUP(num_aligned_elements, threads); - constexpr size_t max_blocks = 65535; - num_blocks = std::min(num_blocks, max_blocks); - - // Launch kernel - switch (align) { - case Alignment::SAME_ALIGNED: - amax_kernel - <<>>(input, amax, N, num_aligned_elements, noop_ptr); - break; - case Alignment::SAME_UNALIGNED: - amax_kernel - <<>>(input, amax, N, num_aligned_elements, noop_ptr); - break; - case Alignment::DIFFERENT: { - // This case is a logic error, since there is only one pointer (input) - // in the alignment check. Still safe to process without vectorization. - amax_kernel<1, true, InputType> - <<>>(input, amax, N, N, noop_ptr); - break; - } - } - - // Check results + // Split the tensor at its first and last 32-byte boundaries. + constexpr size_t elems_per_chunk = amax_chunk_bytes / sizeof(InputType); + NVTE_CHECK(is_aligned_ptr(input, sizeof(InputType)), "Amax input is not aligned to its dtype."); + const size_t misalignment = reinterpret_cast(input) % amax_chunk_bytes; + const size_t head = + std::min(N, misalignment == 0 ? 0 : (amax_chunk_bytes - misalignment) / sizeof(InputType)); + const int64_t num_chunks = static_cast((N - head) / elems_per_chunk); + const size_t tail = N - head - static_cast(num_chunks) * elems_per_chunk; + + // Persistent grid: never more CTAs than the GPU holds at once. + static const int ctas_per_sm = [] { + int n = 0; + NVTE_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&n, amax_kernel, + amax_kernel_threads, 0)); + return n; + }(); + const int64_t max_ctas = static_cast(ctas_per_sm) * cuda::sm_count(); + // Below that, one chunk per lane: spreading a small tensor over more CTAs + // beats giving fewer lanes several chunks each. + const int64_t needed_ctas = DIVUP(num_chunks, static_cast(amax_kernel_threads)); + const int64_t grid = std::max(1, std::min(max_ctas, needed_ctas)); + amax_kernel<<>>( + input, amax, static_cast(head), num_chunks, static_cast(tail), noop_ptr); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -171,9 +236,9 @@ void compute_amax_impl(const NVTETensor input_, const NVTETensor output_, cudaSt float *amax_ptr = reinterpret_cast( (output.amax.dptr != nullptr) ? output.amax.dptr : output.columnwise_amax.dptr); TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.data.dtype, IType, constexpr int nvec = 32 / sizeof(IType); launch_amax_kernel( - reinterpret_cast(input.data.dptr), amax_ptr, input.data.numel(), noop_ptr, - stream);); // NOLINT(*) + input.data.dtype, IType, + launch_amax_kernel(reinterpret_cast(input.data.dptr), amax_ptr, + input.data.numel(), noop_ptr, stream);); // NOLINT(*) } void group_compute_amax_impl(const NVTEGroupedTensor input_, NVTEGroupedTensor output_, From 1fae46f092f2e0ff8c8a023536d9f451eaeb37eb Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 19 Oct 2026 14:45:54 -0700 Subject: [PATCH 2/2] Quantize NVFP4 1D from registers instead of through TMA staging Add a register-resident BF16 -> NVFP4 1D quantize kernel and dispatch to it ahead of quantize_transpose_tuned_1D for every request it supports: rowwise or rowwise + columnwise output, round-to-nearest or stochastic rounding, E4M3 or UE5M3 scales, with or without the global amaxes and the noop flag. Row- scaled NVFP4, fast math, GEMM-swizzled scales, the columnwise-RHT path and tensors past the kernel's grid limit keep using tuned_1D. No specialization stages the input through shared memory or TMA. The rowwise-only one streams each 16-element block from one 32-byte load; the fused one reads every row twice, once per direction, with the second read served by L2, and stages only the transposed output. The per-block coefficient depends only on the block's scale byte, so its IEEE divide is computed once per CTA into a table. The scale arithmetic is TE's own, so round-to-nearest output is byte-identical to tuned_1D. Stochastic rounding draws Philox once per output block, with a subsequence built from the block's coordinates and direction. Each draw is exactly what curanddx returns for the same key and counter, and the round count follows NVTE_BUILD_NUM_PHILOX_ROUNDS; builds with fewer than 3 rounds keep tuned_1D for stochastic rounding. The random bits for a given seed differ from tuned_1D's, whose stream is keyed by thread id. The tiling was tuned on B200. Its tuning constants are gathered in TunedConfig, and the occupancy queries the launch heuristics use are cached per device. Add a rowwise_only parameter to the NVFP4 cast test, which before only ever requested both directions, and test_cast_nvfp4_sr_streams.cu, which checks the generator against curanddx, that every block gets its own subsequence, and that the rowwise and columnwise draws are independent. Signed-off-by: Teddy Do --- tests/cpp/operator/CMakeLists.txt | 1 + .../operator/test_cast_nvfp4_sr_streams.cu | 378 +++++++ .../cpp/operator/test_cast_nvfp4_transpose.cu | 36 +- .../common/cast/nvfp4/nvfp4_block_philox.cuh | 176 ++++ .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 7 + .../quantize_transpose_nvfp4_register_1D.cuh | 929 ++++++++++++++++++ transformer_engine/common/util/ptx.cuh | 122 +++ .../common/util/ptx_arch_spec.cuh | 56 ++ 8 files changed, 1692 insertions(+), 13 deletions(-) create mode 100644 tests/cpp/operator/test_cast_nvfp4_sr_streams.cu create mode 100644 transformer_engine/common/cast/nvfp4/nvfp4_block_philox.cuh create mode 100644 transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_register_1D.cuh diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index c5acd043363..4dece127985 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -15,6 +15,7 @@ add_executable(test_operator test_cast_mxfp8_grouped.cu test_cast_mxfp8_grouped_scaled_swiglu.cu test_cast_nvfp4_transpose.cu + test_cast_nvfp4_sr_streams.cu test_row_scaled_nvfp4_grouped_amax.cu test_cast_float8blockwise.cu test_cast_float8blockwise_grouped.cu diff --git a/tests/cpp/operator/test_cast_nvfp4_sr_streams.cu b/tests/cpp/operator/test_cast_nvfp4_sr_streams.cu new file mode 100644 index 00000000000..9f693a53c0c --- /dev/null +++ b/tests/cpp/operator/test_cast_nvfp4_sr_streams.cu @@ -0,0 +1,378 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Random streams of NVFP4 1D stochastic rounding. +// +// The 1D quantize kernel draws Philox once per output block (nvfp4_block_philox.cuh) with a +// subsequence built from the block's coordinates and direction. These tests check that +// - each draw is exactly what curanddx's Philox returns for the same key and counter, and +// - no two output blocks, and not the rowwise and columnwise draws of one element, share a +// stream: for inputs where every element sits halfway between two E2M1 values, a shared +// stream would show up as identical codes far more often than chance. + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +#include "../test_common.h" +#include "cast/nvfp4/nvfp4_block_philox.cuh" +#include "util/curanddx.hpp" + +using namespace transformer_engine; + +namespace { + +namespace bp = transformer_engine::dispatch::nvfp4::block_philox; + +#define CHECK_CUDA(expr) ASSERT_EQ((expr), cudaSuccess) + +// --------------------------------------------------------------------------------------------- +// The folded generator against curanddx. +// --------------------------------------------------------------------------------------------- + +__global__ void compare_with_curanddx_kernel(const uint64_t *seeds, const uint64_t *offsets, + const uint32_t *sub_lo, const uint32_t *logical_hi, + int n, uint4 *folded, uint4 *reference) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const bp::PhiloxState p = bp::make_philox(seeds[i], offsets[i]); + const bp::PhiloxFold f = bp::make_fold(p, sub_lo[i]); + const uint32_t hi[2] = {logical_hi[i], logical_hi[i] + 1u}; + uint4 out[2]; + bp::philox_fold_n<2>(hi, f, out); + for (int k = 0; k < 2; ++k) { + // The fold's subsequence high word is the logical coordinate XORed with p.b. + const uint64_t subsequence = + (static_cast(hi[k] ^ p.b) << 32) | static_cast(sub_lo[i]); + curanddx::detail::philox4x32_native_state state; + state.init(seeds[i], subsequence, offsets[i]); + reference[2 * i + k] = state.generate4(); + folded[2 * i + k] = out[k]; + } +} + +TEST(NVFP4StochasticRoundingStreams, BlockPhiloxMatchesCuranddx) { + constexpr int n = 1 << 16; + std::mt19937_64 rng(2026); + std::vector seeds(n), offsets(n); + std::vector sub_lo(n), logical_hi(n); + for (int i = 0; i < n; ++i) { + seeds[i] = rng(); + // Small offsets as the framework produces them, and full 64-bit ones. + offsets[i] = (i % 2 == 0) ? (rng() & 0xFFFFF) : rng(); + sub_lo[i] = static_cast(rng()); + logical_hi[i] = static_cast(rng()); + } + uint64_t *d_seeds, *d_offsets; + uint32_t *d_lo, *d_hi; + uint4 *d_folded, *d_reference; + CHECK_CUDA(cudaMalloc(&d_seeds, n * sizeof(uint64_t))); + CHECK_CUDA(cudaMalloc(&d_offsets, n * sizeof(uint64_t))); + CHECK_CUDA(cudaMalloc(&d_lo, n * sizeof(uint32_t))); + CHECK_CUDA(cudaMalloc(&d_hi, n * sizeof(uint32_t))); + CHECK_CUDA(cudaMalloc(&d_folded, 2 * n * sizeof(uint4))); + CHECK_CUDA(cudaMalloc(&d_reference, 2 * n * sizeof(uint4))); + CHECK_CUDA(cudaMemcpy(d_seeds, seeds.data(), n * sizeof(uint64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_offsets, offsets.data(), n * sizeof(uint64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_lo, sub_lo.data(), n * sizeof(uint32_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_hi, logical_hi.data(), n * sizeof(uint32_t), cudaMemcpyHostToDevice)); + compare_with_curanddx_kernel<<>>(d_seeds, d_offsets, d_lo, d_hi, n, d_folded, + d_reference); + CHECK_CUDA(cudaGetLastError()); + std::vector folded(2 * n), reference(2 * n); + CHECK_CUDA(cudaMemcpy(folded.data(), d_folded, 2 * n * sizeof(uint4), cudaMemcpyDeviceToHost)); + CHECK_CUDA( + cudaMemcpy(reference.data(), d_reference, 2 * n * sizeof(uint4), cudaMemcpyDeviceToHost)); + int mismatches = 0; + for (int i = 0; i < 2 * n; ++i) { + const uint4 a = folded[i], b = reference[i]; + if (a.x != b.x || a.y != b.y || a.z != b.z || a.w != b.w) ++mismatches; + } + EXPECT_EQ(mismatches, 0) << "folded Philox differs from curanddx (" << bp::kRounds + << " rounds) in " << mismatches << " of " << 2 * n << " draws"; + for (void *p : {static_cast(d_seeds), static_cast(d_offsets), + static_cast(d_lo), static_cast(d_hi), + static_cast(d_folded), static_cast(d_reference)}) { + cudaFree(p); + } +} + +// --------------------------------------------------------------------------------------------- +// The block -> subsequence map is one-to-one. +// --------------------------------------------------------------------------------------------- + +uint64_t subsequence_key(uint32_t lo, uint32_t hi) { return (static_cast(hi) << 32) | lo; } + +TEST(NVFP4StochasticRoundingStreams, EveryBlockHasItsOwnSubsequence) { + // Every rowwise and columnwise block of one tensor. + constexpr uint32_t M = 2048, K = 4096; + std::vector keys; + keys.reserve(2 * M * K / 16); + for (uint32_t row = 0; row < M; ++row) { + for (uint32_t block_col = 0; block_col < K / 16; ++block_col) { + keys.push_back(subsequence_key(bp::rowwise_sub_lo(block_col), bp::rowwise_sub_hi(row))); + } + } + for (uint32_t row_block = 0; row_block < M / 16; ++row_block) { + for (uint32_t col = 0; col < K; ++col) { + keys.push_back(subsequence_key(bp::columnwise_sub_lo(col & ~3u), + bp::columnwise_sub_hi(row_block, col & 3u))); + } + } + std::sort(keys.begin(), keys.end()); + EXPECT_EQ(std::adjacent_find(keys.begin(), keys.end()), keys.end()) + << "two output blocks share a Philox subsequence"; + + // Across the supported coordinate range the map decodes back to the block, so it cannot + // collide anywhere: rows < 2^30, block columns and columns < 2^31. + std::mt19937_64 rng(7); + for (int i = 0; i < (1 << 20); ++i) { + const uint32_t row = static_cast(rng()) >> 2, block_col = static_cast(rng()) >> 1; + const uint32_t lo = bp::rowwise_sub_lo(block_col), hi = bp::rowwise_sub_hi(row); + ASSERT_EQ(lo & 1u, 0u); + ASSERT_EQ(lo >> 1, block_col); + ASSERT_EQ(hi, row); + const uint32_t row_block = static_cast(rng()) >> 2; + const uint32_t col = static_cast(rng()) >> 1; + const uint32_t clo = bp::columnwise_sub_lo(col & ~3u); + const uint32_t chi = bp::columnwise_sub_hi(row_block, col & 3u); + ASSERT_EQ(clo & 1u, 1u); + ASSERT_EQ(((clo >> 1) & ~3u) | (chi & 3u), col); + ASSERT_EQ(chi >> 2, row_block); + } +} + +// --------------------------------------------------------------------------------------------- +// Stream separation, observed through nvte_quantize_v2. +// +// With the global amax set to 448 * 6, S_enc = 1; every 16-element block contains one 6.0, so +// every block scale is exactly 1.0 and elements are quantized unscaled. All other elements are +// midpoints between adjacent E2M1 magnitudes, so stochastic rounding sends each one up or down +// with probability 1/2. +// --------------------------------------------------------------------------------------------- + +constexpr float kMidpoints[7] = {0.25f, 0.75f, 1.25f, 1.75f, 2.5f, 3.5f, 5.0f}; +// E2M1 magnitude code of the grid value just below kMidpoints[i] is i. +constexpr uint8_t kE4M3One = 0x38; + +// Symmetric pseudo-random choice of midpoint and sign for (r, c). +uint32_t pair_hash(size_t r, size_t c) { + const uint64_t a = std::min(r, c), b = std::max(r, c); + uint64_t h = a * 0x9E3779B97F4A7C15ull ^ (b + 0x632BE59BD9B4E019ull); + h ^= h >> 31; + h *= 0xBF58476D1CE4E5B9ull; + h ^= h >> 29; + return static_cast(h); +} + +enum class Pattern { Symmetric, RowsEqual, ColumnsEqual }; + +float element(Pattern pattern, size_t r, size_t c) { + size_t a = r, b = c; + if (pattern == Pattern::RowsEqual) a = 0; // x[r][c] depends on c only + if (pattern == Pattern::ColumnsEqual) b = 0; // x[r][c] depends on r only + if (pattern == Pattern::Symmetric && r % 16 == c % 16) return 6.0f; + if (pattern == Pattern::RowsEqual && c % 16 == 0) return 6.0f; + if (pattern == Pattern::ColumnsEqual && r % 16 == 0) return 6.0f; + const uint32_t h = pattern == Pattern::Symmetric ? pair_hash(r, c) : pair_hash(a, b); + const float v = kMidpoints[h % 7]; + return (h & 0x100) ? -v : v; +} + +struct QuantizedOutput { + std::vector data, scales, data_t, scales_t; +}; + +// Quantizes the M x K pattern with stochastic rounding; columnwise output if `colwise`. +void quantize_sr(Pattern pattern, size_t M, size_t K, bool colwise, QuantizedOutput *out) { + std::vector<__nv_bfloat16> h_x(M * K); + for (size_t r = 0; r < M; ++r) { + for (size_t c = 0; c < K; ++c) h_x[r * K + c] = __float2bfloat16(element(pattern, r, c)); + } + const size_t s_rows = (M + 127) / 128 * 128, s_cols = (K / 16 + 3) / 4 * 4; + const size_t st_rows = (K + 127) / 128 * 128, st_cols = (M / 16 + 3) / 4 * 4; + void *x, *q, *s, *qt, *st, *amax, *rng; + ASSERT_EQ(cudaMalloc(&x, M * K * 2), cudaSuccess); + ASSERT_EQ(cudaMalloc(&q, M * K / 2), cudaSuccess); + ASSERT_EQ(cudaMalloc(&s, s_rows * s_cols), cudaSuccess); + ASSERT_EQ(cudaMalloc(&qt, M * K / 2), cudaSuccess); + ASSERT_EQ(cudaMalloc(&st, st_rows * st_cols), cudaSuccess); + ASSERT_EQ(cudaMalloc(&amax, 2 * sizeof(float)), cudaSuccess); + ASSERT_EQ(cudaMalloc(&rng, 2 * sizeof(int64_t)), cudaSuccess); + ASSERT_EQ(cudaMemcpy(x, h_x.data(), M * K * 2, cudaMemcpyHostToDevice), cudaSuccess); + const float h_amax[2] = {448.0f * 6.0f, 448.0f * 6.0f}; + ASSERT_EQ(cudaMemcpy(amax, h_amax, sizeof(h_amax), cudaMemcpyHostToDevice), cudaSuccess); + const int64_t h_rng[2] = {12345, 678}; + ASSERT_EQ(cudaMemcpy(rng, h_rng, sizeof(h_rng), cudaMemcpyHostToDevice), cudaSuccess); + + TensorWrapper input(x, std::vector{M, K}, DType::kBFloat16); + TensorWrapper output(NVTE_NVFP4_1D_SCALING); + output.set_rowwise_data(q, DType::kFloat4E2M1, std::vector{M, K}); + output.set_rowwise_scale_inv(s, DType::kFloat8E4M3, std::vector{s_rows, s_cols}); + output.set_amax(amax, DType::kFloat32, std::vector{1}); + if (colwise) { + output.set_columnwise_data(qt, DType::kFloat4E2M1, std::vector{K, M}); + output.set_columnwise_scale_inv(st, DType::kFloat8E4M3, std::vector{st_rows, st_cols}); + output.set_columnwise_amax(static_cast(amax) + 1, DType::kFloat32, + std::vector{1}); + } + TensorWrapper rng_state(rng, std::vector{2}, DType::kInt64); + QuantizationConfigWrapper config; + config.set_stochastic_rounding(true); + config.set_rng_state(rng_state.data()); + nvte_quantize_v2(input.data(), output.data(), config, 0); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + + out->data.resize(M * K / 2); + out->scales.resize(s_rows * s_cols); + ASSERT_EQ(cudaMemcpy(out->data.data(), q, out->data.size(), cudaMemcpyDeviceToHost), cudaSuccess); + ASSERT_EQ(cudaMemcpy(out->scales.data(), s, out->scales.size(), cudaMemcpyDeviceToHost), + cudaSuccess); + if (colwise) { + out->data_t.resize(M * K / 2); + out->scales_t.resize(st_rows * st_cols); + ASSERT_EQ(cudaMemcpy(out->data_t.data(), qt, out->data_t.size(), cudaMemcpyDeviceToHost), + cudaSuccess); + ASSERT_EQ(cudaMemcpy(out->scales_t.data(), st, out->scales_t.size(), cudaMemcpyDeviceToHost), + cudaSuccess); + } + for (void *p : {x, q, s, qt, st, amax, rng}) cudaFree(p); +} + +uint8_t code_at(const std::vector &data, size_t row, size_t cols, size_t col) { + const uint8_t byte = data[(row * cols + col) / 2]; + return (col % 2 == 0) ? (byte & 0xF) : (byte >> 4); +} + +// Checks every code of an (rows x cols) output whose element (i, j) was input value v(i, j): +// 6.0 stays exact, midpoints land on one of their two neighbours. Returns the fraction of +// midpoints rounded up. +template +double check_codes(const std::vector &data, size_t rows, size_t cols, ValueFn v) { + size_t up = 0, total = 0, bad = 0; + for (size_t i = 0; i < rows; ++i) { + for (size_t j = 0; j < cols; ++j) { + const float x = v(i, j); + const uint8_t code = code_at(data, i, cols, j); + const bool neg = (code & 8) != 0; + const uint8_t mag = code & 7; + if (x == 6.0f) { + bad += (code != 7); + continue; + } + int lo = 0; + while (kMidpoints[lo] != std::fabs(x)) ++lo; + const bool sign_ok = (neg == (x < 0)); + if (!sign_ok || (mag != lo && mag != lo + 1)) { + ++bad; + } else { + up += (mag == lo + 1); + } + ++total; + } + } + EXPECT_EQ(bad, 0u) << "codes that are not a neighbour of their input"; + return total ? static_cast(up) / total : 0.0; +} + +class NVFP4StochasticRoundingStreamsTest : public ::testing::Test { + protected: + void SetUp() override { + if (test::getDeviceComputeCapability() < test::blackwellComputeCapability) GTEST_SKIP(); + } +}; + +constexpr size_t kN = 1024; // square, so the symmetric pattern has a transpose to compare + +TEST_F(NVFP4StochasticRoundingStreamsTest, RowwiseAndColumnwiseDrawsDiffer) { + QuantizedOutput out; + quantize_sr(Pattern::Symmetric, kN, kN, /*colwise=*/true, &out); + for (size_t r = 0; r < kN; ++r) { + for (size_t b = 0; b < kN / 16; ++b) { + ASSERT_EQ(out.scales[r * ((kN / 16 + 3) / 4 * 4) + b], kE4M3One); + } + } + auto v = [](size_t i, size_t j) { return element(Pattern::Symmetric, i, j); }; + const double up_row = check_codes(out.data, kN, kN, v); + // Row i of the columnwise output is column i of the input, which by symmetry is row i. + const double up_col = check_codes(out.data_t, kN, kN, v); + EXPECT_NEAR(up_row, 0.5, 0.02); + EXPECT_NEAR(up_col, 0.5, 0.02); + + // Rowwise block (r, b) and columnwise block (r, b) quantize the same 16 values. With 15 + // fair coin flips each, independent streams match with probability 2^-15. + size_t same = 0; + for (size_t r = 0; r < kN; ++r) { + for (size_t b = 0; b < kN / 16; ++b) { + bool eq = true; + for (size_t k = 0; k < 16 && eq; ++k) { + eq = code_at(out.data, r, kN, 16 * b + k) == code_at(out.data_t, r, kN, 16 * b + k); + } + same += eq; + } + } + EXPECT_LT(same, 10u) << same << " of " << kN * kN / 16 + << " rowwise blocks drew the same bits as their columnwise twin"; +} + +// Every row holds the same values, so blocks in one block column differ only in their stream. +void expect_rows_draw_independently(bool colwise) { + QuantizedOutput out; + quantize_sr(Pattern::RowsEqual, kN, kN, colwise, &out); + auto v = [](size_t i, size_t j) { return element(Pattern::RowsEqual, i, j); }; + EXPECT_NEAR(check_codes(out.data, kN, kN, v), 0.5, 0.02); + size_t same = 0; + for (size_t r = 1; r < kN; ++r) { + for (size_t b = 0; b < kN / 16; ++b) { + bool eq = true; + for (size_t k = 0; k < 16 && eq; ++k) { + eq = code_at(out.data, r, kN, 16 * b + k) == code_at(out.data, r - 1, kN, 16 * b + k); + } + same += eq; + } + } + EXPECT_LT(same, 10u) << same << " rowwise blocks drew the same bits as the block above"; +} + +TEST_F(NVFP4StochasticRoundingStreamsTest, RowwiseOnlyBlocksDrawIndependently) { + expect_rows_draw_independently(/*colwise=*/false); +} + +TEST_F(NVFP4StochasticRoundingStreamsTest, FusedRowwiseBlocksDrawIndependently) { + expect_rows_draw_independently(/*colwise=*/true); +} + +// Every column holds the same values, so columnwise blocks in one row block differ only in +// their stream. +TEST_F(NVFP4StochasticRoundingStreamsTest, ColumnwiseBlocksDrawIndependently) { + QuantizedOutput out; + quantize_sr(Pattern::ColumnsEqual, kN, kN, /*colwise=*/true, &out); + // Row c of the columnwise output is column c of the input: value at (c, r) is x[r][c]. + auto vt = [](size_t c, size_t r) { return element(Pattern::ColumnsEqual, r, c); }; + EXPECT_NEAR(check_codes(out.data_t, kN, kN, vt), 0.5, 0.02); + size_t same = 0; + for (size_t c = 1; c < kN; ++c) { + for (size_t b = 0; b < kN / 16; ++b) { + bool eq = true; + for (size_t k = 0; k < 16 && eq; ++k) { + eq = code_at(out.data_t, c, kN, 16 * b + k) == code_at(out.data_t, c - 1, kN, 16 * b + k); + } + same += eq; + } + } + EXPECT_LT(same, 10u) << same << " columnwise blocks drew the same bits as the previous column"; +} + +} // namespace diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index c2529bd419f..c8ed1185d53 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -845,7 +845,8 @@ void performTest(float (*OP)(const float), const NVFP4ScalingMode scaling_mode = NVFP4ScalingMode::Block1D, const NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled, const int scale_max = 0, - const bool use_4over6_err_use_fast_math = false) { + const bool use_4over6_err_use_fast_math = false, + const bool rowwise_only = false) { using namespace test; const bool use_4over6 = mode != kNVTENVFP44Over6Disabled; @@ -863,7 +864,7 @@ void performTest(float (*OP)(const float), const bool is_2d_quantization = use_2d_quantization(scaling_mode); const bool row_scaled_nvfp4 = scaling_mode == NVFP4ScalingMode::RowScaled1D; const bool rowwise = true; - const bool columnwise = !row_scaled_nvfp4; + const bool columnwise = !row_scaled_nvfp4 && !rowwise_only; const size_t rows = first_dimension(shape); const size_t cols = last_dimension(shape); @@ -1083,7 +1084,7 @@ void performTest(float (*OP)(const float), cols, unpadded_blocks_X, scales_stride); - if (!row_scaled_nvfp4) { + if (columnwise) { compare_nvfp4_4over6_candidates("output_t", output.columnwise_cpu_dptr(), output.columnwise_cpu_scale_inv_ptr(), @@ -1099,7 +1100,7 @@ void performTest(float (*OP)(const float), } else { // Set dump_data=true to enable dumping tensor data to files for analysis compareResults_nvfp4(output, ref_output.get(), ref_output_t.get(), rows, cols, atol, rtol, - true, false, !row_scaled_nvfp4); + true, false, columnwise); size_t scale_mismatches_num = 0; if constexpr (std::is_same_v) { @@ -1112,7 +1113,7 @@ void performTest(float (*OP)(const float), unpadded_blocks_Y, unpadded_blocks_X, scales_stride); } - if (!row_scaled_nvfp4) { + if (columnwise) { if constexpr (std::is_same_v) { compare_scaling_factors( "scales_t", output.columnwise_cpu_scale_inv_ptr(), @@ -1365,7 +1366,8 @@ class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam bool, NVFP4ScalingMode, NVFP4FourOverSixTestConfig, - transformer_engine::DType>> {}; + transformer_engine::DType, + bool>> {}; TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { // Skip tests for pre-Blackwell architectures @@ -1383,6 +1385,7 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { const NVFP4ScalingMode scaling_mode = std::get<4>(GetParam()); const NVFP4FourOverSixTestConfig config = std::get<5>(GetParam()); const DType scale_type = std::get<6>(GetParam()); + const bool rowwise_only = std::get<7>(GetParam()); // Skip tests if the input tensor is 1D if (tensor_dims.size() < 2) { @@ -1403,12 +1406,12 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { if (scale_type == DType::kFloat8E4M3) { performTest( OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_max, - config.err_use_fast_math); + config.err_use_fast_math, rowwise_only); #if CUDA_VERSION >= 13040 } else if (scale_type == DType::kFloat8UE5M3) { performTest( OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_max, - config.err_use_fast_math); + config.err_use_fast_math, rowwise_only); #endif } else { FAIL() << "Unsupported NVFP4 scale dtype " << static_cast(scale_type); @@ -1467,6 +1470,9 @@ std::string test_name(const FusedCastTransposeNVFP4TestSuite::ParamType& param) if (std::get<6>(param) != DType::kFloat8E4M3) { name += "X" + test::typeName(std::get<6>(param)); } + if (std::get<7>(param)) { + name += "XROWWISE_ONLY"; + } return name; } @@ -1477,10 +1483,11 @@ INSTANTIATE_TEST_SUITE_P( ::testing::ValuesIn(Activation_types), // activation_type ::testing::ValuesIn(tensor_dims), // tensor_dims ::testing::Values(DType::kBFloat16), // input_type - ::testing::Values(false), // use_fast_math + ::testing::Values(false, true), // use_fast_math ::testing::Values(NVFP4ScalingMode::Block1D), // scaling_mode ::testing::Values(NVFP4FourOverSixTestConfig{}), // four_over_six_config - ::testing::Values(DType::kFloat8E4M3)), // scale_type + ::testing::Values(DType::kFloat8E4M3), // scale_type + ::testing::Values(false, true)), // rowwise_only [](const testing::TestParamInfo& info) { return test_name(info.param); }); @@ -1495,7 +1502,8 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(false), // use_fast_math ::testing::Values(NVFP4ScalingMode::RowScaled1D), // scaling_mode ::testing::Values(NVFP4FourOverSixTestConfig{}), // four_over_six_config - ::testing::Values(DType::kFloat8E4M3)), // scale_type + ::testing::Values(DType::kFloat8E4M3), // scale_type + ::testing::Values(false)), // rowwise_only [](const testing::TestParamInfo& info) { return test_name(info.param); }); @@ -1520,7 +1528,8 @@ INSTANTIATE_TEST_SUITE_P( NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 256, true}, NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, false}, NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, true}), // four_over_six_config - ::testing::Values(DType::kFloat8E4M3)), // scale_type + ::testing::Values(DType::kFloat8E4M3), // scale_type + ::testing::Values(false)), // rowwise_only [](const testing::TestParamInfo& info) { return test_name(info.param); }); @@ -1586,7 +1595,8 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(NVFP4ScalingMode::Block1D, NVFP4ScalingMode::Block2D), // scaling_mode ::testing::Values(NVFP4FourOverSixTestConfig{}), // four_over_six_config - ::testing::Values(DType::kFloat8UE5M3)), + ::testing::Values(DType::kFloat8UE5M3), // scale_type + ::testing::Values(false, true)), // rowwise_only [](const testing::TestParamInfo& info) { return test_name(info.param); }); diff --git a/transformer_engine/common/cast/nvfp4/nvfp4_block_philox.cuh b/transformer_engine/common/cast/nvfp4/nvfp4_block_philox.cuh new file mode 100644 index 00000000000..6bffcfe7c66 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/nvfp4_block_philox.cuh @@ -0,0 +1,176 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file nvfp4_block_philox.cuh + * \brief Philox4x32 draws for NVFP4 stochastic rounding, one per output block, several + * streams at a time. + * + * Each draw is exactly what curanddx's philox4x32_native_state returns from + * init(seed, subsequence, offset) followed by one generate4(): the counter is + * {offset_lo, offset_hi, subsequence_lo, subsequence_hi}. Only the cost differs. Callers + * give every output block its own subsequence and take one counter step from rng_state[1], + * so the counter never moves past the offset the framework reserved, whatever the grid. + * + * Two savings over a plain per-thread generator: + * - Every round-0 term that reads only the launch-wide offset is evaluated once + * (make_philox). + * - When a thread draws for many blocks that share one subsequence word (its block column, + * say), placing that word in the subsequence's LOW half makes both of round 0's products + * thread constants, and the constant parts of rounds 0..2 fold away (make_fold). The + * other half carries the varying coordinate XORed with the launch constant p.b, a + * bijection that cancels round 0's own XOR. + * Several streams advance together so the key schedule is paid once per batch. + * + * The round count follows NVTE_BUILD_NUM_PHILOX_ROUNDS, as TE's other Philox users do. + */ + +#ifndef TRANSFORMER_ENGINE_NVFP4_BLOCK_PHILOX_CUH_ +#define TRANSFORMER_ENGINE_NVFP4_BLOCK_PHILOX_CUH_ + +#include + +#include + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { +namespace block_philox { + +#ifdef NVTE_BUILD_NUM_PHILOX_ROUNDS +constexpr int kRounds = NVTE_BUILD_NUM_PHILOX_ROUNDS; +#else +constexpr int kRounds = 10; +#endif +// make_fold consumes rounds 0..2, so it only applies to builds with at least 3 rounds; callers +// must fall back to a plain generator otherwise. +constexpr bool kFoldSupported = kRounds >= 3; + +// Subsequence of each output block, as (low word, logical high word). The low word is the +// coordinate a thread keeps fixed while it walks the other, so make_fold can fold it. Its +// parity tells the directions apart, so the rowwise and columnwise draws of one element never +// share a subsequence; within a direction the pair is the block's coordinates. +// rowwise block (row, block_col): (2 * block_col, row) +// columnwise block (row_block, col), col = 4q + i: (2 * (4q) + 1, 4 * row_block + i) +// The column's offset inside its group of four moves to the high word because a thread owns +// four adjacent columns and folds against their base. +__host__ __device__ constexpr uint32_t rowwise_sub_lo(uint32_t block_col) { return block_col << 1; } +__host__ __device__ constexpr uint32_t rowwise_sub_hi(uint32_t row) { return row; } +__host__ __device__ constexpr uint32_t columnwise_sub_lo(uint32_t quad_base_col) { + return (quad_base_col << 1) | 1u; +} +__host__ __device__ constexpr uint32_t columnwise_sub_hi(uint32_t row_block, uint32_t col_in_quad) { + return (row_block << 2) | col_in_quad; +} + +constexpr uint32_t kM0 = 0xD2511F53u; +constexpr uint32_t kM1 = 0xCD9E8D57u; +constexpr uint32_t kW0 = 0x9E3779B9u; +constexpr uint32_t kW1 = 0xBB67AE85u; + +// Both halves of a 32 x 32 -> 64-bit product in one IMAD.WIDE.U32, written as PTX like +// curanddx's mulhilo32. As inline asm the product chain keeps the place the kernels give it, +// right after their input loads, where it fills the wait for them; as plain C++ the compiler +// sank it past the loads' consumers and the kernels stalled on the loads instead. +__device__ __forceinline__ void mulhilo(uint32_t a, uint32_t b, uint32_t &lo, uint32_t &hi) { + uint64_t p; + asm("mul.wide.u32 %0, %1, %2;" : "=l"(p) : "r"(a), "r"(b)); + lo = static_cast(p); + hi = static_cast(p >> 32); +} + +// Three-way XOR in one LOP3. +__device__ __forceinline__ uint32_t xor3(uint32_t a, uint32_t b, uint32_t c) { + uint32_t d; + asm("lop3.b32 %0, %1, %2, %3, 0x96;" : "=r"(d) : "r"(a), "r"(b), "r"(c)); + return d; +} + +// Round-0 terms that read only the key and the launch-wide offset. +struct PhiloxState { + uint2 key; // the key round 1 uses + uint32_t a; // offset_hi ^ k0 -> round-0 c.x = hi(sub_lo * M1) ^ a + uint32_t b; // hi(offset_lo * M0) ^ k1 -> round-0 c.z = b ^ sub_hi + uint32_t w; // lo(offset_lo * M0) -> round-0 c.w +}; + +__device__ __forceinline__ PhiloxState make_philox(uint64_t seed, uint64_t offset) { + const uint32_t k0 = static_cast(seed); + const uint32_t k1 = static_cast(seed >> 32); + uint32_t lo0, hi0; + mulhilo(static_cast(offset), kM0, lo0, hi0); + PhiloxState p; + p.a = static_cast(offset >> 32) ^ k0; + p.b = hi0 ^ k1; + p.w = lo0; + p.key = make_uint2(k0 + kW0, k1 + kW1); + return p; +} + +// Rounds 0..2 for a fixed subsequence low word. +struct PhiloxFold { + uint32_t A; // c1_0 ^ k1.x -> c0_1 = hi1 ^ A + uint32_t B; // hi1_2 ^ k2.x -> c0_2 = B ^ c1_1 + uint32_t D; // c3_1 ^ k2.y -> c2_2 = hi0 ^ D + uint32_t c1_2; // round-2 c1, constant for a fixed low word + uint2 k3; // key as of round 3 +}; + +__device__ __forceinline__ PhiloxFold make_fold(const PhiloxState &p, uint32_t sub_lo) { + uint32_t lo0, hi0, lo1, hi1; + mulhilo(sub_lo, kM1, lo1, hi1); // round 0: both products are constant + const uint32_t c0_0 = hi1 ^ p.a, c1_0 = lo1; + uint2 k = p.key; // round-1 key + mulhilo(c0_0, kM0, lo0, hi0); // round 1: the .x product is constant + PhiloxFold f; + f.A = c1_0 ^ k.x; + const uint32_t c2_1 = xor3(hi0, p.w, k.y), c3_1 = lo0; + k.x += kW0; + k.y += kW1; // round-2 key + mulhilo(c2_1, kM1, lo1, hi1); // round 2: the .z product is constant + f.B = hi1 ^ k.x; + f.c1_2 = lo1; + f.D = c3_1 ^ k.y; + k.x += kW0; + k.y += kW1; // round-3 key + f.k3 = k; + return f; +} + +// N draws that share the fold's low word. logical_hi[i] is the varying coordinate; the +// draw's subsequence high word is logical_hi[i] ^ p.b. +template +__device__ __forceinline__ void philox_fold_n(const uint32_t (&logical_hi)[N], const PhiloxFold &f, + uint4 (&c)[N]) { +#pragma unroll + for (int i = 0; i < N; ++i) { // the varying remnant of rounds 0..2 + uint32_t lo0, hi0, lo1, hi1; + mulhilo(logical_hi[i], kM1, lo1, hi1); + mulhilo(hi1 ^ f.A, kM0, lo0, hi0); + c[i] = make_uint4(f.B ^ lo1, f.c1_2, hi0 ^ f.D, lo0); + } + uint2 k = f.k3; +#pragma unroll + for (int r = 3; r < kRounds; ++r) { +#pragma unroll + for (int i = 0; i < N; ++i) { + uint32_t lo0, hi0, lo1, hi1; + mulhilo(c[i].x, kM0, lo0, hi0); + mulhilo(c[i].z, kM1, lo1, hi1); + c[i] = make_uint4(xor3(hi1, c[i].y, k.x), lo1, xor3(hi0, c[i].w, k.y), lo0); + } + if (r < kRounds - 1) { + k.x += kW0; + k.y += kW1; + } + } +} + +} // namespace block_philox +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_NVFP4_BLOCK_PHILOX_CUH_ diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 6a76838695c..f3fc043af01 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -24,6 +24,7 @@ #include "../../util/ptx_arch_spec.cuh" #include "../../utils.cuh" #include "core_nvfp4.cuh" +#include "specialized/quantize_transpose_nvfp4_register_1D.cuh" #include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh" #include "swizzle.cuh" @@ -1860,6 +1861,12 @@ void quantize_transpose_impl(const Tensor &input, const Tensor *noop, Tensor *ou } if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { + // The register-resident kernel serves every 1D request it supports; row-scaled NVFP4, + // fast math, swizzled scales and RHT stay on the TMA kernel. + if (!apply_columnwise_rht && register_1D_supported(input, *output, quant_config)) { + quantize_transpose_register_1D(input, noop, output, quant_config, stream); + return; + } quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); return; } diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_register_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_register_1D.cuh new file mode 100644 index 00000000000..1ec757a8e60 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_register_1D.cuh @@ -0,0 +1,929 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_transpose_nvfp4_register_1D.cuh + * \brief Register-resident BF16 -> NVFP4 1D quantization, rowwise or rowwise + columnwise. + * + * One kernel family, specialized at compile time over + * {rowwise, rowwise + columnwise} x {round-to-nearest-even, stochastic rounding}. + * The numeric core (block amax -> block scale -> element coefficient -> E2M1 payload) is + * shared by every specialization; only the tiling and the output staging differ. No + * specialization stages the input through shared memory or TMA: the rowwise-only one streams + * each 16-element block straight from a 32-byte load, and the fused one reads every row twice, + * once per direction, with the second read served by L2. Shared memory holds only the + * per-CTA coefficient tables and the transposed output. + * + * Scale arithmetic is TE's own (core::compute_global_encode_scaling_factor_FP4, + * core::compute_decoding_scaling_factor, compute_nvfp4_scaling_coefficient), so + * round-to-nearest output is bit-identical to quantize_transpose_tuned_1D for every scale + * type. The per-block coefficient depends only on the block's scale byte, so its IEEE divide + * is evaluated once per CTA into a table indexed by that byte. + * + * Stochastic rounding draws Philox4x32 (nvfp4_block_philox.cuh: bit-identical to curanddx for + * the same key and counter, with NVTE_BUILD_NUM_PHILOX_ROUNDS rounds) once per output block, + * keyed by the block's coordinates and direction. The random bits for a given seed therefore + * differ from tuned_1D's, whose stream is keyed by thread id; both are valid stochastic + * rounding. + * + * The tiling was tuned on B200. Its tuning constants are collected in TunedConfig. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_NVFP4_REGISTER_1D_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_NVFP4_REGISTER_1D_CUH_ + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "../../../common.h" +#include "../../../util/cuda_runtime.h" +#include "../../../util/ptx.cuh" +#include "../../../util/ptx_arch_spec.cuh" +#include "../../../utils.cuh" +#include "../core_nvfp4.cuh" +#include "../nvfp4_block_philox.cuh" +#include "quantize_transpose_nvfp4_tuned_1D.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { +namespace quantize_transpose_register_kernel { + +#if FP4_TYPE_SUPPORTED + +// Tuning constants, measured on B200 at 1500 MHz; retune per target. Everything that depends +// on the SM count comes from per-device queries instead. +struct TunedConfig { + // Rowwise-only: a CTA owns kRowBlocks consecutive 16-element blocks of kRowRows* rows. + // 128 blocks = 2048 input bytes and 128 scale bytes per row per CTA. + static constexpr int kRowBlocks = 128; + static constexpr int kRowThreads = 128; + // Rows per CTA. RNE is memory-bound and 4 rows (4 loads in flight per thread) is where more + // CTAs stop buying latency hiding. SR is bound by the Philox chain, so taller CTAs amortize + // the per-CTA prologue; the launcher doubles kRowRowsSr while the grid still fills the GPU. + static constexpr int kRowRowsRne = 4; + static constexpr int kRowRowsSr = 8; + // Rows held in registers at once (one 32-byte load each). + static constexpr int kRowWave = 4; + + // Rowwise + columnwise: a CTA owns kTileK columns and walks up to kMaxSlabs slabs of kSlabM + // rows, keeping the whole transposed FP4 result for its columns in shared memory. + static constexpr int kColThreads = 128; + static constexpr int kTileK = 128; + static constexpr int kSlabM = 64; + static constexpr int kMaxSlabs = 4; + // Adjacent CTAs that pool their rowwise work into one kGroupCTAs * kTileK wide strip, so a + // warp writes 256 contiguous payload bytes and whole 32-byte scale sectors. + static constexpr int kGroupCTAs = 4; + // Cost model for the slab count: waves * (kSlabWork * slabs + kDrainWork), in memory + // instructions per thread. + static constexpr long long kSlabWork = 36; + static constexpr long long kDrainWork = 51; + + // __launch_bounds__ minimum resident CTAs per SM. Each is the occupancy the specialization + // reaches anyway; stating it hands ptxas that whole register budget. Rowwise SR is the + // exception: 64 registers instead of 72 buys it a quarter more warps to hide Philox with. + template + static constexpr int min_ctas_per_sm() { + if constexpr (COLWISE) return STOCHASTIC ? 5 : 7; + return STOCHASTIC ? 8 : 10; + } +}; + +constexpr int kBlockElems = 16; // elements per NVFP4 1D block; index math below shifts by 4 +constexpr int kBlocksPerSlab = TunedConfig::kSlabM / kBlockElems; +// Transposed staging buffer: one row of kStagingWords words per column. kStagingWords % 32 == 2 +// puts the sixteen lanes of an 8-byte shared access on sixteen disjoint bank pairs for both the +// staging store and the drain's gather, and an even stride keeps every slot 8-byte aligned. +constexpr int kStagingWords = 34; +static_assert(kStagingWords >= 2 * kBlocksPerSlab * TunedConfig::kMaxSlabs); +static_assert(kStagingWords % 32 == 2); +constexpr int kStagingScaleBytes = kBlocksPerSlab * TunedConfig::kMaxSlabs + 4; + +// Scale bytes a block can produce. E4M3 block scales are non-negative, so the sign bit is +// always clear; UE5M3 is unsigned and can use all 256 codes. +template +constexpr int kScaleCodes = std::is_same_v ? 128 : 256; + +// Element offset of (row, pitch). Both are 32-bit and their product is exact in 64 bits, so +// a u32 x u32 widening multiply is one IMAD.WIDE.U32 that also folds in the base pointer. +__device__ __forceinline__ size_t row_off(int row, int pitch) { + return static_cast(static_cast(row)) * static_cast(pitch); +} + +// --------------------------------------------------------------------------------------------- +// Numeric core. A block is 16 BF16 values in eight packed words; v[h] holds element 2h in its +// low half and 2h+1 in its high half. +// --------------------------------------------------------------------------------------------- + +// Block scale byte for a block amax, by TE's formula, as its raw ScaleType bits. +// +// core::compute_decoding_scaling_factor clamps the product to FLT_MAX before the saturating +// conversion. For E4M3 the clamp cannot change the result here: block_amax is an abs-max of +// BF16 values (max.xorsign.abs drops NaN operands, so it is never NaN) and S_enc / 6 is finite +// and positive, so the product is never NaN, and the one value the clamp could change, +inf, +// converts to the same saturated 448 that FLT_MAX does. CLAMP = false drops that instruction +// from the issue-bound specializations; UE5M3 always keeps it. +template +__device__ __forceinline__ uint32_t scale_code(float block_amax, float S_enc) { + if constexpr (CLAMP || !std::is_same_v) { + return core::compute_decoding_scaling_factor(block_amax, S_enc).__x; + } else { + constexpr float fp4_max_inv = 1.0f / detail::TypeExtrema::max; + return static_cast(block_amax * (S_enc * fp4_max_inv)).__x; + } +} + +template +__device__ __forceinline__ ScaleType scale_from_code(uint32_t code) { + ScaleType s; + s.__x = static_cast<__nv_fp8_storage_t>(code); + return s; +} + +template +__device__ __forceinline__ uint32_t row_block_scale_code(const uint32_t (&v)[8], float S_enc) { + uint32_t m0 = ptx::abs_max_bf16x2(v[0], v[1]); + uint32_t m1 = ptx::abs_max_bf16x2(v[2], v[3]); + uint32_t m2 = ptx::abs_max_bf16x2(v[4], v[5]); + uint32_t m3 = ptx::abs_max_bf16x2(v[6], v[7]); + m0 = ptx::abs_max_bf16x2(m0, m1); + m2 = ptx::abs_max_bf16x2(m2, m3); + m0 = ptx::abs_max_bf16x2(m0, m2); + // Fold the two lanes with the same unit: swap them and take one more abs-max. + const float block_amax = + __uint_as_float(ptx::abs_max_bf16x2(m0, ptx::swap_halves_b32(m0)) & 0x7fff0000u); + return scale_code(block_amax, S_enc); +} + +template +__device__ __forceinline__ uint32_t pack_half(const float (&f)[8], uint32_t r0, uint32_t r1) { + if constexpr (STOCHASTIC) { + return ptx::cvt_f32x8_to_fp4x8_rs(f, r0, r1); + } else { + return ptx::cvt_f32x8_to_fp4x8_rn(f); + } +} + +// Scale and pack one rowwise block, half a block at a time to keep 8 floats live, not 16. +template +__device__ __forceinline__ void pack_row_block(const uint32_t (&v)[8], uint64_t coeff2, + const uint4 &rnd, uint32_t &w0, uint32_t &w1) { + float f[8]; +#pragma unroll + for (int h = 0; h < 4; ++h) { + ptx::unpack_f32x2(ptx::mul_f32x2(ptx::bf16x2_to_f32x2(v[h]), coeff2), f[2 * h], f[2 * h + 1]); + } + w0 = pack_half(f, rnd.x, rnd.y); +#pragma unroll + for (int h = 4; h < 8; ++h) { + ptx::unpack_f32x2(ptx::mul_f32x2(ptx::bf16x2_to_f32x2(v[h]), coeff2), f[2 * h - 8], + f[2 * h - 7]); + } + w1 = pack_half(f, rnd.z, rnd.w); +} + +// Two rows of one column, scaled. LOW selects the even column of each packed pair. +template +__device__ __forceinline__ void scale_column_pair(uint32_t a, uint32_t b, uint64_t coeff2, + float &lo, float &hi) { + const uint32_t ea = LOW ? (a << 16) : (a & 0xffff0000u); + const uint32_t eb = LOW ? (b << 16) : (b & 0xffff0000u); + ptx::unpack_f32x2(ptx::mul_f32x2((static_cast(eb) << 32) | ea, coeff2), lo, hi); +} + +// Columnwise twin of pack_row_block: col[t] is row t of a packed column pair. +template +__device__ __forceinline__ void pack_column_block(const uint32_t (&col)[16], uint64_t coeff2, + const uint4 &rnd, uint32_t &w0, uint32_t &w1) { + float f[8]; +#pragma unroll + for (int h = 0; h < 4; ++h) { + scale_column_pair(col[2 * h], col[2 * h + 1], coeff2, f[2 * h], f[2 * h + 1]); + } + w0 = pack_half(f, rnd.x, rnd.y); +#pragma unroll + for (int h = 4; h < 8; ++h) { + scale_column_pair(col[2 * h], col[2 * h + 1], coeff2, f[2 * h - 8], f[2 * h - 7]); + } + w1 = pack_half(f, rnd.z, rnd.w); +} + +// Element coefficient for scale code `code`: compute_nvfp4_scaling_coefficient's +// formula, fminf(1 / (scale * (1 / S_enc)), FLT_MAX), with the inner reciprocal S_dec = 1 / S_enc +// computed once by the caller. __frcp_rn is the correctly rounded reciprocal, so it equals +// the IEEE quotient 1.0f / x bit for bit. +template +__device__ __forceinline__ float coefficient_of(uint32_t code, float S_dec) { + const float scale = static_cast(scale_from_code(code)); + return fminf(__frcp_rn(scale * S_dec), detail::TypeExtrema::max); +} + +// Coefficient table entry. Rowwise RNE consumes every coefficient as a packed pair, so it +// stores that form once; the others keep the narrower float (a 64-bit shared load cost the SR +// specialization more than the move it saves). +template +using CoeffEntry = std::conditional_t; + +template +__device__ __forceinline__ uint64_t coeff2_of(Entry e) { + if constexpr (std::is_same_v) { + return e; + } else { + return ptx::splat_f32x2(e); + } +} + +// Fill the per-CTA coefficient tables: [0, kScaleCodes) rowwise and, for the fused +// specialization, [kScaleCodes, 2 * kScaleCodes) columnwise. No barrier; callers place it. +// The trip count is a compile-time constant: a runtime-bounded loop around the divides made +// ptxas spill the already-loaded input across the divide's slow-path call. +template +__device__ __forceinline__ void write_coeff_tables(CoeffEntry *table, + float S_enc_row, float S_enc_col, + bool separate_col) { + constexpr int kCodes = kScaleCodes; + static_assert(kCodes % THREADS == 0, "Each thread fills the same number of entries."); + const float S_dec_row = __frcp_rn(S_enc_row); + const float S_dec_col = separate_col ? __frcp_rn(S_enc_col) : S_dec_row; +#pragma unroll + for (int k = 0; k < kCodes / THREADS; ++k) { + const int code = threadIdx.x + k * THREADS; + const float row_coeff = coefficient_of(code, S_dec_row); + if constexpr (COLWISE) { + table[code] = row_coeff; + table[kCodes + code] = separate_col ? coefficient_of(code, S_dec_col) : row_coeff; + } else if constexpr (!STOCHASTIC) { + table[code] = ptx::splat_f32x2(row_coeff); + } else { + table[code] = row_coeff; + } + } +} + +template +__global__ void __launch_bounds__(COLWISE ? TunedConfig::kColThreads : TunedConfig::kRowThreads, + TunedConfig::min_ctas_per_sm()) + quantize_transpose_nvfp4_register_1D_kernel( + const uint16_t *__restrict__ x, const float *__restrict__ noop, + const float *__restrict__ amax_row_ptr, const float *__restrict__ amax_col_ptr, + const size_t *__restrict__ rng_state, uint8_t *__restrict__ q_row, + uint8_t *__restrict__ s_row, uint8_t *__restrict__ q_col, uint8_t *__restrict__ s_col, + const int M, const int K, const int s_row_pitch, const int s_col_pitch, const int nsubs) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int kCodes = kScaleCodes; + const int bpr = K >> 4; // rowwise blocks per row (16 elements each) + + using Entry = CoeffEntry; + __shared__ Entry coeff_table[COLWISE ? 2 * kCodes : kCodes]; + + // Launch-wide scalars are READ here but consumed further down. Each heads a dependency + // chain off a global load (the encode-scale divides, the Philox key schedule), and a warp + // issues in order, so only reading them here lets ptxas keep the loads at the top while the + // arithmetic sinks below the first input loads, overlapping the two latencies. + const bool skip = noop != nullptr && noop[0] == 1.0f; + const float amax_row = amax_row_ptr != nullptr ? *amax_row_ptr : 0.0f; + const float amax_col = (COLWISE && amax_col_ptr != nullptr) ? *amax_col_ptr : 0.0f; + uint64_t rng_seed = 0, rng_offset = 0; + if constexpr (STOCHASTIC) { + if (rng_state != nullptr) { + rng_seed = rng_state[0]; + rng_offset = rng_state[1]; + } + } + // Same rules as quantize_transpose_tuned_1D: no rowwise amax means S_enc = 1, no columnwise + // amax means the columnwise output reuses the rowwise S_enc. + auto encode_scales = [&](float &S_enc_row, float &S_enc_col) { + S_enc_row = amax_row_ptr != nullptr + ? core::compute_global_encode_scaling_factor_FP4(amax_row) + : 1.0f; + // The usual caller passes the same global amax in both directions; recognizing that from + // the raw bits skips a second divide chain in the CTA prologue. + const bool same_amax = + amax_row_ptr != nullptr && __float_as_uint(amax_col) == __float_as_uint(amax_row); + S_enc_col = (COLWISE && amax_col_ptr != nullptr && !same_amax) + ? core::compute_global_encode_scaling_factor_FP4(amax_col) + : S_enc_row; + }; + + if constexpr (!COLWISE) { + // ---- rowwise only: pure streaming, one block per thread per row ---- + constexpr int kRowWave = TunedConfig::kRowWave; + const int j0 = blockIdx.x * TunedConfig::kRowBlocks + threadIdx.x; + const int j = j0 < bpr ? j0 : 0; // clamp so the prefetch below stays in bounds + const int r0 = blockIdx.y * ROWS; + + // The column term of all three addresses is a thread constant, folded into the base + // pointers once, so each per-block address is one row_off() from them. (Induction + // pointers instead cost twelve live registers here, in the one register-bound + // specialization, and measured 6% slower.) + const uint16_t *const xj = x + static_cast(j) * kBlockElems; + uint8_t *const qj = q_row + static_cast(j) * 8; + uint8_t *const sj = s_row + static_cast(j); + + uint32_t v[kRowWave][8]; +#pragma unroll + for (int i = 0; i < kRowWave; ++i) { + ptx::ld_global_nc_evict_first_b32x8_relaxed(v[i], xj + row_off(r0 + i, K)); + } + if (skip) return; // CTA-uniform, and no barrier has been reached yet + + // Resolve the RNG first: its offset-dependent multiply has the longer wait, and the + // already-issued loads continue toward readiness behind it. The block column is a thread + // constant, so placing it in the subsequence's low word folds rounds 0..2 away. + block_philox::PhiloxFold pf{}; + if constexpr (STOCHASTIC) { + pf = block_philox::make_fold(block_philox::make_philox(rng_seed, rng_offset), + block_philox::rowwise_sub_lo(j)); + } + + float S_enc_row, S_enc_col; + encode_scales(S_enc_row, S_enc_col); + + // Written here; its barrier is deferred into the first wave, past that wave's amax trees, + // which depend only on the input already in flight. + write_coeff_tables( + coeff_table, S_enc_row, S_enc_col, false); + +#pragma unroll + for (int w = 0; w < ROWS; w += kRowWave) { + if (w) { +#pragma unroll + for (int i = 0; i < kRowWave; ++i) { + ptx::ld_global_nc_evict_first_b32x8_relaxed(v[i], xj + row_off(r0 + w + i, K)); + } + } + // Register-free L2 warm-up for the next wave. The rows are inside this CTA's stripe + // (ROWS divides 32 and M % 32 == 0), so no bounds test is needed. + if (w + kRowWave < ROWS) { +#pragma unroll + for (int i = 0; i < kRowWave; ++i) { + ptx::prefetch_l2(xj + row_off(r0 + w + kRowWave + i, K)); + } + } + // The wave's first random pair needs none of the loaded data, so draw it before the + // amax trees that stall on the loads. Only the first pair: this specialization is + // register-bound, so the rest stay next to their packs. + uint4 rnd0[2] = {}; + if constexpr (STOCHASTIC) { + uint32_t rows0[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) rows0[i] = block_philox::rowwise_sub_hi(r0 + w + i); + block_philox::philox_fold_n<2>(rows0, pf, rnd0); + } + // Every block's amax tree and scale byte resolve before any of them packs. + uint32_t codes[kRowWave]; +#pragma unroll + for (int i = 0; i < kRowWave; ++i) { + codes[i] = row_block_scale_code(v[i], S_enc_row); + } + if (w == 0) { + __syncthreads(); // the coefficient table is read only from here on + if (j0 >= bpr) return; // no barrier follows, so a split warp is safe + } + if constexpr (!STOCHASTIC) { + uint64_t coeff2[kRowWave]; +#pragma unroll + for (int i = 0; i < kRowWave; ++i) coeff2[i] = coeff_table[codes[i]]; +#pragma unroll + for (int i = 0; i < kRowWave; ++i) { + const int row = r0 + w + i; + uint32_t w0, w1; + pack_row_block(v[i], coeff2[i], make_uint4(0, 0, 0, 0), w0, w1); + *reinterpret_cast(qj + row_off(row, K >> 1)) = make_uint2(w0, w1); + sj[row_off(row, s_row_pitch)] = static_cast(codes[i]); + } + } else { +#pragma unroll + for (int g = 0; g < kRowWave; g += 2) { + uint64_t coeff2[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) coeff2[i] = coeff2_of(coeff_table[codes[g + i]]); + uint4 rnd[2]; + if (g == 0) { // drawn above, ahead of the amax trees + rnd[0] = rnd0[0]; + rnd[1] = rnd0[1]; + } else { + uint32_t rows[2]; +#pragma unroll + for (int i = 0; i < 2; ++i) rows[i] = block_philox::rowwise_sub_hi(r0 + w + g + i); + block_philox::philox_fold_n<2>(rows, pf, rnd); + } +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int row = r0 + w + g + i; + uint32_t w0, w1; + pack_row_block(v[g + i], coeff2[i], rnd[i], w0, w1); + *reinterpret_cast(qj + row_off(row, K >> 1)) = make_uint2(w0, w1); + sj[row_off(row, s_row_pitch)] = static_cast(codes[g + i]); + } + } + } + } + } else { + // ---- rowwise + columnwise ---- + constexpr int kTileK = TunedConfig::kTileK; + constexpr int kSlabM = TunedConfig::kSlabM; + // Transposed result, column-major: staging[column slot][word]. + __shared__ uint32_t staging[kTileK][kStagingWords]; + __shared__ uint8_t staging_scales[kTileK][kStagingScaleBytes]; + + const int c0 = blockIdx.x * kTileK; + const int r00 = blockIdx.y * (nsubs * kSlabM); + + // Phase A lanes: kJA block columns per row x kRgA rows per pass, from the GRP-wide strip + // this CTA belongs to (GRP == 1 is the plain per-tile mapping). + constexpr int kJA = 8 * GRP; + constexpr int kRgA = TunedConfig::kColThreads / kJA; + constexpr int kRowsA = kSlabM / GRP; + const int qa = static_cast(blockIdx.x) & (GRP - 1); // position inside the strip + const int jbase = (static_cast(blockIdx.x) - qa) * (kTileK / kBlockElems); + const int jl = threadIdx.x & (kJA - 1); + const int rl = threadIdx.x / kJA; + const int j0 = jbase + jl; + const int j = j0 < bpr ? j0 : jbase; + // Phase B lanes: four adjacent columns per thread (a warp spans 128 columns, one 256-byte + // request per row) and one of the slab's four column blocks. + const int cp = threadIdx.x & 31; + const int mb = threadIdx.x >> 5; + const int c = c0 + 4 * cp; + const bool cok = c < K; + // A band of at most 32 rows that starts inside the matrix also ends inside it (M % 32 == + // 0), so clamping the band base is all the bounds work phase A's row walk needs. + constexpr bool kBandSafe = kRowsA <= 32; + const bool arow = j0 < bpr; + // 32-bit row strides: ptxas folds a 32-bit pointer step into one IMAD.WIDE. + const int xastep = kRgA * K, qastep = kRgA * (K >> 1), sastep = kRgA * s_row_pitch; + const size_t xcol = static_cast(j * kBlockElems); + const size_t qcol = static_cast(j * 8); + + // The first real load is inside the slab loop, so it cannot be hoisted above the + // encode-scale chain; a prefetch can, and warms the first slab meanwhile. + { + const int pb = r00 + kRowsA * qa + rl; + const uint16_t *pf = + x + + (static_cast(kBandSafe && pb < M ? pb : r00) * static_cast(K) + xcol); +#pragma unroll + for (int i = 0; i < 4; ++i) { + if constexpr (kBandSafe) { + ptx::prefetch_l2(pf); + pf += xastep; + } else { + const int r = pb + kRgA * i; + ptx::prefetch_l2( + x + (static_cast(r < M ? r : r00) * static_cast(K) + xcol)); + } + } + } + if (skip) return; // CTA-uniform, and no barrier has been reached yet + + float S_enc_row, S_enc_col; + encode_scales(S_enc_row, S_enc_col); + // Equal encode scales (the usual case: one global amax for both directions) share the + // rowwise table instead of evaluating a second set of divides. + const bool separate_col = __float_as_uint(S_enc_row) != __float_as_uint(S_enc_col); + // Both halves fold Philox rounds 0..2 against a low word that is constant for the CTA: the + // rowwise half keys on its block column, the columnwise half on the base of its four + // columns, with the column's offset in the quad folded into the high word. + block_philox::PhiloxFold pf_r{}, pf_c{}; + if constexpr (STOCHASTIC) { + const block_philox::PhiloxState ps = block_philox::make_philox(rng_seed, rng_offset); + pf_r = block_philox::make_fold(ps, block_philox::rowwise_sub_lo(j)); + pf_c = block_philox::make_fold(ps, block_philox::columnwise_sub_lo(c)); + } + + write_coeff_tables( + coeff_table, S_enc_row, S_enc_col, separate_col); + __syncthreads(); + + for (int sub = 0; sub < nsubs; ++sub) { + const int r0 = r00 + sub * kSlabM; + if (r0 >= M) break; + + // Phase A: one rowwise block per thread per quarter-slab, all three streams advancing by + // a constant row stride. A full slab needs no per-row bounds test. + const int rbase = r0 + kRowsA * qa + rl; + const bool afull = arow && (r0 + kSlabM <= M); + const uint16_t *xsafe = x + (static_cast(r0) * static_cast(K) + xcol); + const uint16_t *xr = kBandSafe && rbase < M + ? x + (static_cast(rbase) * static_cast(K) + xcol) + : xsafe; + int rows[4]; + uint32_t v[4][8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int r = rbase + kRgA * i; + rows[i] = r; + if constexpr (kBandSafe) { + ptx::ld_global_evict_last_b32x8(v[i], xr); + } else { + ptx::ld_global_evict_last_b32x8( + v[i], r < M ? x + (static_cast(r) * static_cast(K) + xcol) : xsafe); + } + xr += xastep; + } + + // Warm L2 for the next slab under SR only: RNE streams near its DRAM ceiling and a + // lookahead competes with demand loads, while SR has the slack for it to land. + if constexpr (STOCHASTIC) { + if (sub + 1 < nsubs && r0 + 2 * kSlabM <= M) { + const uint16_t *pn = + x + (static_cast(rbase + kSlabM) * static_cast(K) + xcol); +#pragma unroll + for (int i = 0; i < 4; ++i) { + ptx::prefetch_l2(pn); + pn += xastep; + } + } + } + + // The Philox chain depends only on the row, never on the data, so it goes before the + // amax trees that stall on the loads. + uint4 rnd_row[4] = {}; + if constexpr (STOCHASTIC) { +#pragma unroll + for (int g = 0; g < 4; g += 2) { + const uint32_t hi[2] = {block_philox::rowwise_sub_hi(rows[g]), + block_philox::rowwise_sub_hi(rows[g + 1])}; + uint4 out[2]; + block_philox::philox_fold_n<2>(hi, pf_r, out); + rnd_row[g] = out[0]; + rnd_row[g + 1] = out[1]; + } + } + uint32_t codes_a[4]; + uint64_t coeff2_a[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) + codes_a[i] = row_block_scale_code(v[i], S_enc_row); +#pragma unroll + for (int i = 0; i < 4; ++i) coeff2_a[i] = coeff2_of(coeff_table[codes_a[i]]); + uint8_t *qap = q_row + (static_cast(rbase) * static_cast(K >> 1) + qcol); + uint8_t *sap = s_row + (static_cast(rbase) * static_cast(s_row_pitch) + + static_cast(j)); +#pragma unroll + for (int i = 0; i < 4; ++i) { + if (afull || !(rows[i] >= M || !arow)) { + uint32_t w0, w1; + pack_row_block(v[i], coeff2_a[i], rnd_row[i], w0, w1); + *reinterpret_cast(qap) = make_uint2(w0, w1); + *sap = static_cast(codes_a[i]); + } + qap += qastep; + sap += sastep; + } + + // Phase B: the rows phase A just read are still in L2, so this column-major pass is a + // cache re-read rather than a second DRAM stream. + { + const int rb = r0 + mb * kBlockElems; + const bool bok = cok && (rb + kBlockElems <= M); + const int cc = bok ? c : 0; + const int rr = bok ? rb : r0; + uint32_t raw[2][16]; + { + const uint16_t *cp2 = + x + (static_cast(rr) * static_cast(K) + static_cast(cc)); +#pragma unroll + for (int t = 0; t < 16; ++t) { + const uint2 rv = *reinterpret_cast(cp2); + raw[0][t] = rv.x; + raw[1][t] = rv.y; + cp2 += K; + } + } + + const int wofs = (sub * kBlocksPerSlab + mb) * 2; + const int bofs = sub * kBlocksPerSlab + mb; + // High word = (row block, column inside the quad); the quad base is in the folded + // low word, so the four columns differ only in two bits. + uint4 rnd[4] = {}; + if constexpr (STOCHASTIC) { + const uint32_t row_block = static_cast(rb) >> 4; +#pragma unroll + for (int g = 0; g < 4; g += 2) { + const uint32_t hi[2] = {block_philox::columnwise_sub_hi(row_block, g), + block_philox::columnwise_sub_hi(row_block, g + 1)}; + uint4 out[2]; + block_philox::philox_fold_n<2>(hi, pf_c, out); + rnd[g] = out[0]; + rnd[g + 1] = out[1]; + } + } + // Both 16-row amax trees and all four coefficient lookups resolve before any column + // packs, so the four shared loads are in flight together. + uint32_t codes_b[4]; + uint64_t coeff2_b[4]; +#pragma unroll + for (int p = 0; p < 2; ++p) { // one packed word holds a pair of adjacent columns + uint32_t tr[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) tr[i] = ptx::abs_max_bf16x2(raw[p][2 * i], raw[p][2 * i + 1]); +#pragma unroll + for (int i = 0; i < 4; ++i) tr[i] = ptx::abs_max_bf16x2(tr[2 * i], tr[2 * i + 1]); + tr[0] = ptx::abs_max_bf16x2(tr[0], tr[1]); + tr[1] = ptx::abs_max_bf16x2(tr[2], tr[3]); + const uint32_t pair_amax = ptx::abs_max_bf16x2(tr[0], tr[1]) & 0x7fff7fffu; + codes_b[2 * p] = + scale_code(__uint_as_float(pair_amax << 16), S_enc_col); + codes_b[2 * p + 1] = + scale_code(__uint_as_float(pair_amax & 0xffff0000u), S_enc_col); + } +#pragma unroll + for (int k = 0; k < 4; ++k) coeff2_b[k] = coeff2_of(coeff_table[kCodes + codes_b[k]]); +#pragma unroll + for (int p = 0; p < 2; ++p) { +#pragma unroll + for (int q = 0; q < 2; ++q) { + const int cs = 2 * p + q; + uint32_t w0, w1; + if (q == 0) { + pack_column_block(raw[p], coeff2_b[cs], rnd[cs], w0, w1); + } else { + pack_column_block(raw[p], coeff2_b[cs], rnd[cs], w0, w1); + } + const int slot = (cs << 5) | cp; // one 8-byte staging store per column block + *reinterpret_cast(&staging[slot][wofs]) = make_uint2(w0, w1); + staging_scales[slot][bofs] = static_cast(codes_b[cs]); + } + } + } + } + + __syncthreads(); + + // Phase C: four lanes cooperate on one column, each writing a whole 32-byte sector, so a + // column's transposed run leaves as one coalesced request. + { + const int nb = min(nsubs * kBlocksPerSlab, (M - r00) >> 4); // valid blocks + const int ll = threadIdx.x & 3; + const int cgrp = threadIdx.x >> 2; // this lane group's column quad + const int cbase = c0 + 4 * cgrp; + const int nblk = nb - 4 * ll; // payload blocks this lane owns (always even) + const size_t qrow = static_cast(M >> 1); + const bool wide = (qrow & 31) == 0; // 32-byte alignment of a transposed row + uint8_t *qp = q_col + static_cast(cbase) * qrow + (r00 >> 1) + 32 * ll; + uint8_t *sp = s_col + static_cast(cbase) * s_col_pitch + (r00 >> 4) + 4 * ll; + int slot = cgrp; +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (cbase + q < K && nblk > 0) { + uint32_t d[8]; +#pragma unroll + for (int t = 0; t < 4; ++t) { + const uint2 a = *reinterpret_cast(&staging[slot][8 * ll + 2 * t]); + d[2 * t] = a.x; + d[2 * t + 1] = a.y; + } + if (nblk >= 4 && wide) { + ptx::st_global_b32x8(qp, d); // one 32-byte sector + } else if (nblk >= 4) { // rows are only 16-byte aligned when M % 64 != 0 + *reinterpret_cast(qp) = make_uint4(d[0], d[1], d[2], d[3]); + *reinterpret_cast(qp + 16) = make_uint4(d[4], d[5], d[6], d[7]); + } else { // tail slab: this lane owns exactly two blocks + *reinterpret_cast(qp) = make_uint4(d[0], d[1], d[2], d[3]); + } + const uint32_t sw = *reinterpret_cast(&staging_scales[slot][4 * ll]); + if (nblk >= 4) { + *reinterpret_cast(sp) = sw; + } else { + *reinterpret_cast(sp) = static_cast(sw); + } + } + slot += 32; + qp += qrow; + sp += s_col_pitch; + } + } + } +#else + NVTE_DEVICE_ERROR("sm_100 or higher is required."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +// Resident CTAs on the current device for one instantiation, from the driver. Cached per +// device (and per instantiation), so a process driving several different devices gets each +// device's own value. +template +int resident_ctas() { + constexpr int threads = COLWISE ? TunedConfig::kColThreads : TunedConfig::kRowThreads; + static std::vector cache(cuda::num_devices(), 0); + static std::vector flags(cuda::num_devices()); + const int device = cuda::current_device(); + std::call_once(flags[device], [&] { + int per_sm = 0; + NVTE_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &per_sm, + quantize_transpose_nvfp4_register_1D_kernel, + threads, 0)); + cache[device] = std::max(per_sm, 1) * cuda::sm_count(device); + }); + return cache[device]; +} + +struct LaunchArgs { + const uint16_t *x; + const float *noop; + const float *amax_row; + const float *amax_col; + const size_t *rng_state; + uint8_t *q_row; + uint8_t *s_row; + uint8_t *q_col; + uint8_t *s_col; + int M; + int K; + int s_row_pitch; + int s_col_pitch; +}; + +template +void launch_shape(const LaunchArgs &a, cudaStream_t stream) { + dim3 grid; + int nsubs = 1; + if constexpr (COLWISE) { + // Taller CTAs make each transposed run longer (the drain costs one request per column + // whatever the run length) but shrink the grid, and a grid that no longer covers the + // device several times wastes its tail. Pick the slab count by + // waves * (slab streaming work + per-CTA drain). + const int resident = resident_ctas(); + const int gx = DIVUP(a.K, TunedConfig::kTileK); + long long best = -1; + for (int n = 1; n <= TunedConfig::kMaxSlabs; n <<= 1) { + const long long gy = + DIVUP(static_cast(a.M), static_cast(n) * TunedConfig::kSlabM); + const long long g = static_cast(gx) * gy; + const long long waves = DIVUP(g, static_cast(resident)); + const long long filled = std::min(g, resident); + const long long cost = + waves * (TunedConfig::kSlabWork * n + TunedConfig::kDrainWork) * resident / filled; + if (best < 0 || cost < best) { + best = cost; + nsubs = n; + } + } + grid = dim3(gx, DIVUP(a.M, nsubs * TunedConfig::kSlabM)); + } else { + grid = dim3(DIVUP(a.K / kBlockElems, TunedConfig::kRowBlocks), DIVUP(a.M, ROWS)); + } + constexpr int threads = COLWISE ? TunedConfig::kColThreads : TunedConfig::kRowThreads; + quantize_transpose_nvfp4_register_1D_kernel + <<>>(a.x, a.noop, a.amax_row, a.amax_col, a.rng_state, a.q_row, + a.s_row, a.q_col, a.s_col, a.M, a.K, a.s_row_pitch, + a.s_col_pitch, nsubs); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void launch(const LaunchArgs &a, cudaStream_t stream) { + if constexpr (COLWISE) { + // The strip remap needs every column of a kGroupCTAs-CTA strip to exist. + if (a.K % (TunedConfig::kGroupCTAs * TunedConfig::kTileK) == 0) { + launch_shape(a, stream); + } else { + launch_shape(a, stream); + } + } else { + // Under SR the Philox chain, not DRAM, sets the pace, so the per-CTA prologue (amax load, + // divides, coefficient table, barrier) is a real fraction of each CTA. Doubling the + // stripe halves it, worth taking whenever the taller grid still fills the device. + constexpr int kBase = STOCHASTIC ? TunedConfig::kRowRowsSr : TunedConfig::kRowRowsRne; + if constexpr (STOCHASTIC) { + const int gx = DIVUP(a.K / kBlockElems, TunedConfig::kRowBlocks); + const int resident = resident_ctas(); + if (static_cast(gx) * DIVUP(a.M, 2 * kBase) >= resident) { + launch_shape(a, stream); + return; + } + } + launch_shape(a, stream); + } +} + +#endif // FP4_TYPE_SUPPORTED + +} // namespace quantize_transpose_register_kernel + +/*! \brief Whether quantize_transpose_register_1D can serve this request. + * + * BF16 input, rowwise output (columnwise optional), compact scales, no row-scaled NVFP4, no + * fast math, rows and columns multiples of 32, 32-byte aligned buffers, on an SM 10.x GPU; + * stochastic rounding also needs a build with at least 3 Philox rounds. + * Everything else stays on quantize_transpose_tuned_1D. + */ +inline bool register_1D_supported(const Tensor &input, const Tensor &output, + const QuantizationConfig *quant_config) { +#if FP4_TYPE_SUPPORTED + using quantize_transpose_register_kernel::TunedConfig; + if (cuda::sm_arch() / 10 != 10) return false; + if (input.dtype() != DType::kBFloat16 || !input.has_data()) return false; + if (!output.has_data() || output.row_scaled_nvfp4 || output.with_gemm_swizzled_scales) { + return false; + } + if (quant_config != nullptr && quant_config->use_fast_math) return false; + if (quant_config != nullptr && quant_config->stochastic_rounding && + !block_philox::kFoldSupported) { + return false; + } + + const auto [rows, cols] = input.flat_2d_dims(); + if (rows == 0 || cols == 0 || rows % 32 != 0 || cols % 32 != 0) return false; + // 32-bit indexing and grid limits. + if (cols > static_cast(INT_MAX) / 16) return false; + if (DIVUP(rows, static_cast(TunedConfig::kRowRowsRne)) > 65535) return false; + + if (output.scale_inv.shape.size() != 2 || output.scale_inv.shape[1] < cols / 16 || + output.scale_inv.shape[1] > static_cast(INT_MAX)) { + return false; + } + if (!is_aligned_ptr(input.data.dptr, 32) || !is_aligned_ptr(output.data.dptr, 8)) return false; + if (output.has_columnwise_data()) { + const auto &col_scale = output.columnwise_scale_inv; + if (col_scale.dptr == nullptr || col_scale.shape.size() != 2 || + col_scale.shape[1] < DIVUP(rows, static_cast(16)) || col_scale.shape[1] % 4 != 0 || + col_scale.shape[1] > static_cast(INT_MAX)) { + return false; + } + if (!is_aligned_ptr(output.columnwise_data.dptr, 32) || !is_aligned_ptr(col_scale.dptr, 4)) { + return false; + } + } + return true; +#else + return false; +#endif // FP4_TYPE_SUPPORTED +} + +/*! \brief Quantize BF16 to NVFP4 1D, rowwise and optionally columnwise, from registers. + * + * Caller must have checked register_1D_supported. + */ +template +void quantize_transpose_register_1D(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace quantize_transpose_register_kernel; + + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "input"); + CheckOutputTensor(*output, "output", false); + NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + const bool colwise = output->has_columnwise_data(); + if (colwise) { + NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), + "Transposed output must have FP4 type."); + } + + const bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; + const size_t *rng_state = nullptr; + const NVTETensor rng_state_tensor = (quant_config != nullptr) ? quant_config->rng_state : nullptr; + if (rng_state_tensor != nullptr) { + const Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor); + NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64, + "RNG state should contain 2 64-bit values."); + NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2}, + "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape); + rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr); + } + + const auto [rows, cols] = input.flat_2d_dims(); + LaunchArgs a; + a.x = reinterpret_cast(input.data.dptr); + a.noop = reinterpret_cast(noop->data.dptr); + a.amax_row = reinterpret_cast(output->amax.dptr); + a.amax_col = reinterpret_cast(output->columnwise_amax.dptr); + a.rng_state = rng_state; + a.q_row = reinterpret_cast(output->data.dptr); + a.s_row = reinterpret_cast(output->scale_inv.dptr); + a.q_col = colwise ? reinterpret_cast(output->columnwise_data.dptr) : nullptr; + a.s_col = colwise ? reinterpret_cast(output->columnwise_scale_inv.dptr) : nullptr; + a.M = static_cast(rows); + a.K = static_cast(cols); + a.s_row_pitch = static_cast(output->scale_inv.shape[1]); + a.s_col_pitch = colwise ? static_cast(output->columnwise_scale_inv.shape[1]) : 0; + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + colwise, COLWISE, + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_stochastic_rounding, STOCHASTIC, + launch(a, stream););); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_NVFP4_REGISTER_1D_CUH_ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 74fd2261f58..0126348d66d 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -1912,6 +1912,128 @@ __device__ __forceinline__ void prefetch_l2_evict_last(const void *addr) { NVTE_DEVICE_ERROR("prefetch_l2_evict_last is only supported on SM 8.0+."); #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) } + +// Bring one global cache line into L2 with the default eviction priority. A +// prefetch names no destination register, so it can be issued ahead of a chain +// the thread is already stalled on. +__device__ __forceinline__ void prefetch_l2(const void *addr) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + asm volatile("prefetch.global.L2 [%0];" : : "l"(addr)); +#else + NVTE_DEVICE_ERROR("prefetch_l2 is only supported on SM 8.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) +} + +// Coherent 256-bit global loads for data that is read more than once, so they +// stay off the read-only path, which would not keep the line for the re-read. +// The evict_last form additionally keeps the line resident in L2 until then. +// Neither is `volatile` and neither clobbers memory: use them only on data the +// kernel does not write, so the compiler is free to schedule them early. +// Needs sm_100 or later, PTX ISA 8.8 (CUDA 12.9) and a 32-byte aligned address. +__device__ __forceinline__ void ld_global_b32x8(uint32_t (&dst)[8], const void *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (defined CUDA_VERSION) && \ + (CUDA_VERSION >= 12090) + asm("ld.global.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_b32x8 is only supported on SM 10.0+ with CUDA 12.9+."); +#endif +} + +__device__ __forceinline__ void ld_global_evict_last_b32x8(uint32_t (&dst)[8], const void *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (defined CUDA_VERSION) && \ + (CUDA_VERSION >= 12090) + asm("ld.global.L2::evict_last.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_evict_last_b32x8 is only supported on SM 10.0+ with CUDA 12.9+."); +#endif +} + +// As ld_global_nc_evict_first_b32x8, but not `volatile`, so the compiler may +// hoist the load above independent work. For input the kernel never writes. +__device__ __forceinline__ void ld_global_nc_evict_first_b32x8_relaxed(uint32_t (&dst)[8], + const void *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (defined CUDA_VERSION) && \ + (CUDA_VERSION >= 12090) + asm("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_relaxed is only supported on SM 10.0+ with CUDA 12.9+."); +#endif +} + +// 256-bit global store. Needs sm_100 or later, PTX ISA 8.8 (CUDA 12.9) and a +// 32-byte aligned address. +__device__ __forceinline__ void st_global_b32x8(void *dst, const uint32_t (&src)[8]) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (defined CUDA_VERSION) && \ + (CUDA_VERSION >= 12090) + asm volatile("st.global.v8.b32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" ::"l"(dst), "r"(src[0]), + "r"(src[1]), "r"(src[2]), "r"(src[3]), "r"(src[4]), "r"(src[5]), "r"(src[6]), + "r"(src[7]) + : "memory"); +#else + NVTE_DEVICE_ERROR("st_global_b32x8 is only supported on SM 10.0+ with CUDA 12.9+."); +#endif +} + +// Packed FP32 pairs held in one 64-bit register, as `mul.f32x2` consumes them. +// The low 32 bits are the first element. + +// Widen a packed BF16 pair to a packed FP32 pair: BF16 -> FP32 is exact and is +// just the 16 bits placed at the top of the word. +__device__ __forceinline__ uint64_t bf16x2_to_f32x2(uint32_t v) { + const uint32_t lo = v << 16; + const uint32_t hi = v & 0xffff0000u; + return (static_cast(hi) << 32) | lo; +} + +__device__ __forceinline__ uint64_t splat_f32x2(float v) { + const uint64_t bits = __float_as_uint(v); + return (bits << 32) | bits; +} + +__device__ __forceinline__ uint64_t mul_f32x2(uint64_t a, uint64_t b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint64_t d; + asm("mul.f32x2 %0, %1, %2;" : "=l"(d) : "l"(a), "l"(b)); + return d; +#else + NVTE_DEVICE_ERROR("mul_f32x2 is only supported on SM 10.0+."); + return 0; +#endif +} + +__device__ __forceinline__ void unpack_f32x2(uint64_t p, float &lo, float &hi) { + lo = __uint_as_float(static_cast(p)); + hi = __uint_as_float(static_cast(p >> 32)); +} + +// Elementwise max of |x| over two packed BF16 pairs. The result may carry a +// sign bit (max.xorsign.abs), so mask it off before using the magnitude. Not +// `volatile`, unlike abs_max_2x, so it schedules freely inside amax trees. +__device__ __forceinline__ uint32_t abs_max_bf16x2(uint32_t a, uint32_t b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) + uint32_t d; + asm("max.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b)); + return d; +#else + NVTE_DEVICE_ERROR("abs_max_bf16x2 is only supported on SM 8.9+."); + return 0; +#endif +} + +// Exchange the two 16-bit halves of a word. +__device__ __forceinline__ uint32_t swap_halves_b32(uint32_t v) { + return __byte_perm(v, v, 0x1032); +} } // namespace ptx namespace { diff --git a/transformer_engine/common/util/ptx_arch_spec.cuh b/transformer_engine/common/util/ptx_arch_spec.cuh index 14f390740e1..076b9e4ec9b 100644 --- a/transformer_engine/common/util/ptx_arch_spec.cuh +++ b/transformer_engine/common/util/ptx_arch_spec.cuh @@ -510,6 +510,62 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding( return out_8x; } +// Eight already-scaled FP32 values -> eight E2M1 codes in one word, rounding to +// nearest even; element 2i goes to the low nibble of byte i. +__device__ __forceinline__ uint32_t cvt_f32x8_to_fp4x8_rn(const float (&f)[8]) { + uint32_t out_8x = 0; + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + asm("{\n\t" + ".reg .b8 b0, b1, b2, b3;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b2, %6, %5;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b3, %8, %7;\n\t" + "mov.b32 %0, {b0, b1, b2, b3};\n\t" + "}" + : "=r"(out_8x) + : "f"(f[0]), "f"(f[1]), "f"(f[2]), "f"(f[3]), "f"(f[4]), "f"(f[5]), "f"(f[6]), "f"(f[7])); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return out_8x; +} + +// As cvt_f32x8_to_fp4x8_rn, with stochastic rounding: rbits03 supplies 8 +// random bits for each of elements 0-3, rbits47 for elements 4-7. +__device__ __forceinline__ uint32_t cvt_f32x8_to_fp4x8_rs(const float (&f)[8], + const uint32_t rbits03, + const uint32_t rbits47) { + uint32_t out_8x = 0; + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm("{\n\t" + ".reg .b16 h0, h1;\n\t" + "cvt.rs.satfinite.e2m1x4.f32 h0, {%4, %3, %2, %1}, %9;\n\t" + "cvt.rs.satfinite.e2m1x4.f32 h1, {%8, %7, %6, %5}, %10;\n\t" + "mov.b32 %0, {h0, h1};\n\t" + "}" + : "=r"(out_8x) + : "f"(f[0]), "f"(f[1]), "f"(f[2]), "f"(f[3]), "f"(f[4]), "f"(f[5]), "f"(f[6]), "f"(f[7]), + "r"(rbits03), "r"(rbits47)); + } else { + float q[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + q[i] = stochastic_round_fp4_e2m1(f[i], rbits03 >> (8 * i)); + q[i + 4] = stochastic_round_fp4_e2m1(f[i + 4], rbits47 >> (8 * i)); + } + const fp4e2m1x4 lo(make_float4(q[0], q[1], q[2], q[3])); + const fp4e2m1x4 hi(make_float4(q[4], q[5], q[6], q[7])); + out_8x = static_cast(*reinterpret_cast(&lo)) | + (static_cast(*reinterpret_cast(&hi)) << 16); + } + return out_8x; +} + #endif // FP4_TYPE_SUPPORTED } // namespace ptx