diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 599b2f892c1..c9b0744caa2 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_row_scaled_nvfp4_grouped_cast.cu test_cast_float8blockwise.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 f83564b35fa..7192af4d282 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/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_, 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