diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index c5acd043363..110c904815a 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -20,7 +20,7 @@ add_executable(test_operator test_cast_float8blockwise_grouped.cu test_dequantize_mxfp8.cu test_dequantize_mxfp8_grouped.cu - test_fused_group_requantize_mxfp8.cu + test_requantize_mxfp8_grouped.cu test_dequantize_float8blockwise_grouped.cu test_dequantize_nvfp4.cu test_transpose.cu diff --git a/tests/cpp/operator/test_fused_group_requantize_mxfp8.cu b/tests/cpp/operator/test_fused_group_requantize_mxfp8.cu deleted file mode 100644 index a24804c760d..00000000000 --- a/tests/cpp/operator/test_fused_group_requantize_mxfp8.cu +++ /dev/null @@ -1,344 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include - -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -#include "../test_common.h" - -namespace { - -using namespace transformer_engine; -using namespace test; - -using transformer_engine::DType; -using transformer_engine::QuantizationConfigWrapper; -using test::Tensor; -using test::bf16; -using test::fp32; -using test::fp8e4m3; -using test::fp8e8m0; -using test::int64; - -constexpr size_t MXFP8_SCALE_DIM = 32; -constexpr uint8_t kSentinel = 0xAB; - -struct FusedGroupRequantizeCase { - size_t hidden_size; - std::vector splits; // per-group LIVE row counts, each a multiple of 128 - size_t capacity_rows; // allocated rows; 0 means sum(splits) (no tail) -}; - -size_t product(const NVTEShape &shape) { - size_t result = 1; - for (size_t i = 0; i < shape.ndim; ++i) { - result *= shape.data[i]; - } - return result; -} - -class FusedGroupRequantizeTestSuite - : public ::testing::TestWithParam< - std::tuple> {}; - -TEST_P(FusedGroupRequantizeTestSuite, MatchesUnfusedChainReference) { - if (test::getDeviceComputeCapability() < test::blackwellComputeCapability) { - GTEST_SKIP() << "Fused grouped MXFP8 requantization requires Blackwell or newer"; - } - - const auto test_case = std::get<0>(GetParam()); - const DType input_type = std::get<1>(GetParam()); - const bool use_fast_math = std::get<2>(GetParam()); - const bool return_dequantized = std::get<3>(GetParam()); - - const size_t hidden_size = test_case.hidden_size; - const std::vector &splits = test_case.splits; - const size_t num_groups = splits.size(); - const size_t num_live_rows = - std::accumulate(splits.begin(), splits.end(), static_cast(0)); - const size_t num_rows = - test_case.capacity_rows != 0 ? test_case.capacity_rows : num_live_rows; - - ASSERT_EQ(hidden_size % 128, 0); - for (const size_t split : splits) { - ASSERT_EQ(split % 128, 0); - } - ASSERT_GT(num_live_rows, 0); - ASSERT_GE(num_rows, num_live_rows); - ASSERT_EQ(num_rows % 128, 0); - - const std::vector shape{num_rows, hidden_size}; - - // Deterministic, position-dependent data over the FULL capacity (including any - // tail): a missing tail guard would then read plausible data and corrupt live - // scales, which the sentinel checks below catch. - Tensor input_fp32("input_fp32", shape, DType::kFloat32); - auto *input_fp32_cpu = input_fp32.rowwise_cpu_dptr(); - for (size_t row = 0; row < num_rows; ++row) { - for (size_t col = 0; col < hidden_size; ++col) { - const int value = static_cast((131 * row + 17 * col + (row * col) % 29) % 257) - 128; - input_fp32_cpu[row * hidden_size + col] = static_cast(value) / 32.0f; - } - } - input_fp32.from_cpu(); - - // The wire tensor: rowwise MXFP8 with compact (unswizzled) scales. Under the - // /128 contract TE's padded scale allocation degenerates to the compact layout. - Tensor input_mxfp8("input_mxfp8", shape, input_type, /*rowwise=*/true, - /*columnwise=*/false, NVTE_MXFP8_1D_SCALING); - nvte_quantize(input_fp32.data(), input_mxfp8.data(), 0); - - // Element-based exclusive-cumsum offsets (the grouped tensor's tensor_offsets - // convention): offsets[g] = row offset x hidden, num_groups + 1 entries, on - // the device. offsets[num_groups] covers only the LIVE rows. - Tensor tensor_offsets("tensor_offsets", std::vector{num_groups + 1}, DType::kInt64); - auto *tensor_offsets_cpu = tensor_offsets.rowwise_cpu_dptr(); - tensor_offsets_cpu[0] = 0; - for (size_t g = 0; g < num_groups; ++g) { - tensor_offsets_cpu[g + 1] = - tensor_offsets_cpu[g] + static_cast(splits[g] * hidden_size); - } - tensor_offsets.from_cpu(); - - // Destination: columnwise E4M3 + per-group swizzled columnwise scales, plus the - // rowwise scale buffer receiving the swizzled copy of the input's scales. - Tensor actual("actual", shape, DType::kFloat8E4M3, /*rowwise=*/true, - /*columnwise=*/true, NVTE_MXFP8_1D_SCALING); - actual.set_with_gemm_swizzled_scales(true); - - Tensor dequantized_out("dequantized_out", - return_dequantized ? shape : std::vector{128, 128}, - DType::kBFloat16); - - // Sentinel-fill every output the kernel writes, so both unwritten-tail and - // stray-write behavior are observable. - const NVTEBasicTensor actual_rowwise_scales_param = - nvte_get_tensor_param(actual.data(), kNVTERowwiseScaleInv); - const NVTEBasicTensor actual_colwise_scales_param = - nvte_get_tensor_param(actual.data(), kNVTEColumnwiseScaleInv); - const size_t rowwise_scales_alloc = product(actual_rowwise_scales_param.shape); - const size_t colwise_scales_alloc = product(actual_colwise_scales_param.shape); - ASSERT_EQ(cudaMemset(actual.columnwise_dptr(), kSentinel, num_rows * hidden_size), - cudaSuccess); - ASSERT_EQ(cudaMemset(actual_rowwise_scales_param.data_ptr, kSentinel, rowwise_scales_alloc), - cudaSuccess); - ASSERT_EQ(cudaMemset(actual_colwise_scales_param.data_ptr, kSentinel, colwise_scales_alloc), - cudaSuccess); - if (return_dequantized) { - ASSERT_EQ(cudaMemset(dequantized_out.rowwise_dptr(), kSentinel, - num_rows * hidden_size * sizeof(bf16)), - cudaSuccess); - } - - QuantizationConfigWrapper fused_config; - fused_config.set_use_fast_math(use_fast_math); - nvte_group_requantize(input_mxfp8.data(), actual.data(), tensor_offsets.data(), - return_dequantized ? dequantized_out.data() : nullptr, - fused_config, 0); - - // Reference intermediate: the unfused chain materializes the dequantized tensor; - // fast math rounds it to BF16, the default path keeps FP32. - const DType intermediate_type = use_fast_math ? DType::kBFloat16 : DType::kFloat32; - Tensor dequantized_ref("dequantized_ref", shape, intermediate_type); - nvte_dequantize(input_mxfp8.data(), dequantized_ref.data(), 0); - dequantized_ref.to_cpu(); - - // Rowwise-scale reference: the production dense swizzle over the full capacity. - // For 128-aligned groups its live prefix is byte-identical to the per-group - // layout (tiles are row-major over 128-row tiles, and the live bound is - // 128-aligned). - Tensor rowwise_swizzled_ref("rowwise_swizzled_ref", shape, input_type, /*rowwise=*/true, - /*columnwise=*/false, NVTE_MXFP8_1D_SCALING); - rowwise_swizzled_ref.set_with_gemm_swizzled_scales(true); - ASSERT_EQ(cudaMemcpy(rowwise_swizzled_ref.rowwise_dptr(), input_mxfp8.rowwise_dptr(), - num_rows * hidden_size, cudaMemcpyDeviceToDevice), - cudaSuccess); - nvte_swizzle_scaling_factors(input_mxfp8.data(), rowwise_swizzled_ref.data(), 0); - - // Columnwise reference, one group at a time through the production single-tensor - // kernels: slice the intermediate, quantize columnwise, swizzle. Group blocks - // concatenate exactly (no padding) because every count is a multiple of 128. - std::vector reference_colwise_data(num_live_rows * hidden_size); - std::vector reference_colwise_scales(num_live_rows / MXFP8_SCALE_DIM * hidden_size); - size_t row_offset = 0; - for (size_t g = 0; g < num_groups; ++g) { - const size_t group_rows = splits[g]; - if (group_rows == 0) { - continue; - } - const std::vector group_shape{group_rows, hidden_size}; - - Tensor group_intermediate("group_intermediate", group_shape, intermediate_type); - const size_t element_size = use_fast_math ? sizeof(bf16) : sizeof(fp32); - const uint8_t *intermediate_cpu = nullptr; - uint8_t *group_intermediate_cpu = nullptr; - if (use_fast_math) { - intermediate_cpu = - reinterpret_cast(dequantized_ref.rowwise_cpu_dptr()); - group_intermediate_cpu = - reinterpret_cast(group_intermediate.rowwise_cpu_dptr()); - } else { - intermediate_cpu = - reinterpret_cast(dequantized_ref.rowwise_cpu_dptr()); - group_intermediate_cpu = - reinterpret_cast(group_intermediate.rowwise_cpu_dptr()); - } - memcpy(group_intermediate_cpu, intermediate_cpu + row_offset * hidden_size * element_size, - group_rows * hidden_size * element_size); - group_intermediate.from_cpu(); - - Tensor group_quantized("group_quantized", group_shape, DType::kFloat8E4M3, - /*rowwise=*/false, /*columnwise=*/true, NVTE_MXFP8_1D_SCALING); - nvte_quantize(group_intermediate.data(), group_quantized.data(), 0); - - Tensor group_swizzled("group_swizzled", group_shape, DType::kFloat8E4M3, - /*rowwise=*/false, /*columnwise=*/true, NVTE_MXFP8_1D_SCALING); - group_swizzled.set_with_gemm_swizzled_scales(true); - ASSERT_EQ(cudaMemcpy(group_swizzled.columnwise_dptr(), group_quantized.columnwise_dptr(), - group_rows * hidden_size, cudaMemcpyDeviceToDevice), - cudaSuccess); - nvte_swizzle_scaling_factors(group_quantized.data(), group_swizzled.data(), 0); - - group_swizzled.to_cpu(); - const auto *group_data = - reinterpret_cast(group_swizzled.columnwise_cpu_dptr()); - memcpy(reference_colwise_data.data() + row_offset * hidden_size, group_data, - group_rows * hidden_size); - const auto *group_scales = reinterpret_cast( - group_swizzled.columnwise_cpu_scale_inv_ptr()); - memcpy(reference_colwise_scales.data() + row_offset / MXFP8_SCALE_DIM * hidden_size, - group_scales, group_rows / MXFP8_SCALE_DIM * hidden_size); - - row_offset += group_rows; - } - - ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); - ASSERT_EQ(cudaGetLastError(), cudaSuccess); - - actual.to_cpu(); - - // Columnwise data: live rows must match the per-group production chain bit for - // bit; capacity-tail rows must be untouched. - const auto *actual_data = - reinterpret_cast(actual.columnwise_cpu_dptr()); - for (size_t row = 0; row < num_live_rows; ++row) { - for (size_t col = 0; col < hidden_size; ++col) { - const size_t i = row * hidden_size + col; - ASSERT_EQ(actual_data[i], reference_colwise_data[i]) - << "FP8 columnwise data mismatch at row " << row << ", column " << col; - } - } - for (size_t i = num_live_rows * hidden_size; i < num_rows * hidden_size; ++i) { - ASSERT_EQ(actual_data[i], kSentinel) - << "Capacity-tail columnwise data byte was written at flat index " << i; - } - - // Columnwise scales: the per-group swizzled blocks fill exactly the live - // prefix; everything past it (including any allocation padding) stays sentinel. - const auto *actual_colwise_scales = reinterpret_cast( - actual.columnwise_cpu_scale_inv_ptr()); - for (size_t i = 0; i < reference_colwise_scales.size(); ++i) { - ASSERT_EQ(actual_colwise_scales[i], reference_colwise_scales[i]) - << "E8M0 columnwise scale mismatch at physical index " << i; - } - for (size_t i = reference_colwise_scales.size(); i < colwise_scales_alloc; ++i) { - ASSERT_EQ(actual_colwise_scales[i], kSentinel) - << "Capacity-tail columnwise scale byte was written at physical index " << i; - } - - // Rowwise scales: the swizzled copy must match the production dense swizzle on - // the live prefix and stay sentinel past it. - const size_t num_live_rowwise_scales = num_live_rows * hidden_size / MXFP8_SCALE_DIM; - const auto *actual_rowwise_scales = - reinterpret_cast(actual.rowwise_cpu_scale_inv_ptr()); - const auto *reference_rowwise_scales = reinterpret_cast( - rowwise_swizzled_ref.rowwise_cpu_scale_inv_ptr()); - for (size_t i = 0; i < num_live_rowwise_scales; ++i) { - ASSERT_EQ(actual_rowwise_scales[i], reference_rowwise_scales[i]) - << "E8M0 rowwise scale mismatch at physical index " << i; - } - for (size_t i = num_live_rowwise_scales; i < rowwise_scales_alloc; ++i) { - ASSERT_EQ(actual_rowwise_scales[i], kSentinel) - << "Capacity-tail rowwise scale byte was written at physical index " << i; - } - - // Optional BF16 output: live rows dequantize exactly (FP8 values scaled by a - // power of two are exactly representable in BF16 after one rounding), so the - // comparison is bitwise; tail rows stay sentinel. - if (return_dequantized) { - Tensor dequantized_bf16_ref("dequantized_bf16_ref", shape, DType::kBFloat16); - nvte_dequantize(input_mxfp8.data(), dequantized_bf16_ref.data(), 0); - ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); - dequantized_out.to_cpu(); - dequantized_bf16_ref.to_cpu(); - const auto *actual_dq = - reinterpret_cast(dequantized_out.rowwise_cpu_dptr()); - const auto *reference_dq = - reinterpret_cast(dequantized_bf16_ref.rowwise_cpu_dptr()); - for (size_t row = 0; row < num_live_rows; ++row) { - for (size_t col = 0; col < hidden_size; ++col) { - const size_t i = row * hidden_size + col; - ASSERT_EQ(actual_dq[i], reference_dq[i]) - << "BF16 dequantized mismatch at row " << row << ", column " << col; - } - } - constexpr uint16_t kSentinelBf16 = - static_cast(kSentinel) | (static_cast(kSentinel) << 8); - for (size_t i = num_live_rows * hidden_size; i < num_rows * hidden_size; ++i) { - ASSERT_EQ(actual_dq[i], kSentinelBf16) - << "Capacity-tail dequantized element was written at flat index " << i; - } - } -} - -std::string splitsToString(const std::vector &splits) { - std::string result; - for (const size_t split : splits) { - if (!result.empty()) { - result += "_"; - } - result += std::to_string(split); - } - return result; -} - -INSTANTIATE_TEST_SUITE_P( - OperatorTest, FusedGroupRequantizeTestSuite, - ::testing::Combine( - ::testing::Values( - FusedGroupRequantizeCase{512, {1024}, 0}, // single group - FusedGroupRequantizeCase{512, {512, 512, 512, 512}, 0}, // uniform - FusedGroupRequantizeCase{256, {256, 1024, 128, 640}, 0}, // variable - FusedGroupRequantizeCase{512, {0, 512, 256}, 0}, // zero-token front - FusedGroupRequantizeCase{384, {256, 0, 0, 512}, 0}, // adjacent zero-token - FusedGroupRequantizeCase{512, {512, 256, 0}, 0}, // zero-token end - FusedGroupRequantizeCase{8192, {512, 256, 512}, 2048}, // capacity tail - FusedGroupRequantizeCase{512, {0, 512, 0, 256}, 1024}, // zero groups + tail - FusedGroupRequantizeCase{8192, {2048, 4096, 1024, 1024}, 0}), // production-like - ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), - ::testing::Values(false, true), // use_fast_math - ::testing::Values(false, true)), // return_dequantized - [](const testing::TestParamInfo &info) { - const auto test_case = std::get<0>(info.param); - return "H" + std::to_string(test_case.hidden_size) + "xS" + - splitsToString(test_case.splits) + "xCap" + - std::to_string(test_case.capacity_rows) + "x" + - test::typeName(std::get<1>(info.param)) + "xFastMath" + - std::to_string(std::get<2>(info.param)) + "xDequant" + - std::to_string(std::get<3>(info.param)); - }); - -} // namespace diff --git a/tests/cpp/operator/test_requantize_mxfp8_grouped.cu b/tests/cpp/operator/test_requantize_mxfp8_grouped.cu new file mode 100644 index 00000000000..02ef8b269e4 --- /dev/null +++ b/tests/cpp/operator/test_requantize_mxfp8_grouped.cu @@ -0,0 +1,582 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include + +#include "../test_common.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +constexpr size_t kMXFP8ScaleDim = 32; + +struct GroupShapeCase { + std::string name; + std::vector> shapes; + size_t capacity_rows = 0; +}; + +struct GroupShapeInfo { + size_t num_tensors = 0; + size_t total_elements = 0; + size_t rowwise_scale_elements = 0; + size_t columnwise_scale_elements = 0; + bool same_first_dim = false; + bool same_last_dim = false; + std::vector logical_shape; + std::vector first_dims; + std::vector last_dims; + std::vector offsets; +}; + +size_t round_up(size_t value, size_t alignment) { + return ((value + alignment - 1) / alignment) * alignment; +} + +GroupShapeInfo make_shape_info(const GroupShapeCase &test_case) { + GroupShapeInfo info; + info.num_tensors = test_case.shapes.size(); + info.first_dims.reserve(info.num_tensors); + info.last_dims.reserve(info.num_tensors); + info.offsets.reserve(info.num_tensors + 1); + info.offsets.push_back(0); + + for (const auto &[rows, cols] : test_case.shapes) { + info.first_dims.push_back(static_cast(rows)); + info.last_dims.push_back(static_cast(cols)); + info.total_elements += rows * cols; + info.offsets.push_back(static_cast(info.total_elements)); + } + + info.same_first_dim = true; + info.same_last_dim = true; + for (size_t i = 1; i < info.num_tensors; ++i) { + info.same_first_dim &= info.first_dims[i] == info.first_dims[0]; + info.same_last_dim &= info.last_dims[i] == info.last_dims[0]; + } + + if (info.same_first_dim && info.same_last_dim) { + info.logical_shape = { + static_cast(info.first_dims[0]) * info.num_tensors, + static_cast(info.last_dims[0]), + }; + } else if (!info.same_first_dim && info.same_last_dim) { + size_t total_rows = 0; + for (const int64_t rows : info.first_dims) { + total_rows += static_cast(rows); + } + info.logical_shape = {total_rows, static_cast(info.last_dims[0])}; + } else if (info.same_first_dim && !info.same_last_dim) { + size_t total_cols = 0; + for (const int64_t cols : info.last_dims) { + total_cols += static_cast(cols); + } + info.logical_shape = {static_cast(info.first_dims[0]), total_cols}; + } else { + info.logical_shape = {1, info.total_elements}; + } + + if (test_case.capacity_rows != 0) { + NVTE_CHECK(info.same_last_dim && test_case.capacity_rows >= info.logical_shape[0], + "Capacity requires a common hidden dimension and enough rows."); + info.same_first_dim = false; + info.logical_shape[0] = test_case.capacity_rows; + info.total_elements = test_case.capacity_rows * info.logical_shape[1]; + } + + // SAME_BOTH_DIMS and VARYING_FIRST_DIM are represented as one tall tensor by + // the grouped MXFP8 kernels. Their compact scales therefore have padding only + // at the end of the whole group, not between member tensors. + if (info.same_last_dim) { + const size_t total_rows = info.logical_shape[0]; + const size_t cols = info.logical_shape[1]; + const size_t rowwise_stride = round_up( + (cols + kMXFP8ScaleDim - 1) / kMXFP8ScaleDim, + scale_tensor_alignment_X_rowwise); + const size_t columnwise_stride = round_up(cols, scale_tensor_alignment_X_colwise); + info.rowwise_scale_elements = + round_up(total_rows, scale_tensor_alignment_Y_rowwise) * rowwise_stride; + info.columnwise_scale_elements = + round_up((total_rows + kMXFP8ScaleDim - 1) / kMXFP8ScaleDim, + scale_tensor_alignment_Y_colwise) * + columnwise_stride; + } else { + // For varying last dimensions, each member has its own descriptor and its + // own compact scale range. The test shapes obey the grouped-kernel contract + // (both member dimensions are multiples of 128), so compact and padded + // capacities are identical. + for (const auto &[rows, cols] : test_case.shapes) { + const size_t rowwise_rows = round_up(rows, scale_tensor_alignment_Y_rowwise); + const size_t rowwise_cols = round_up( + (cols + kMXFP8ScaleDim - 1) / kMXFP8ScaleDim, + scale_tensor_alignment_X_rowwise); + const size_t columnwise_rows = round_up( + (rows + kMXFP8ScaleDim - 1) / kMXFP8ScaleDim, + scale_tensor_alignment_Y_colwise); + const size_t columnwise_cols = round_up(cols, scale_tensor_alignment_X_colwise); + info.rowwise_scale_elements += rowwise_rows * rowwise_cols; + info.columnwise_scale_elements += columnwise_rows * columnwise_cols; + } + } + + return info; +} + +size_t dtype_size(DType dtype) { + switch (dtype) { + case DType::kFloat32: + return sizeof(float); + case DType::kBFloat16: + return sizeof(bf16); + case DType::kFloat8E4M3: + return sizeof(fp8e4m3); + case DType::kFloat8E5M2: + return sizeof(fp8e5m2); + default: + NVTE_ERROR("Unsupported dtype in grouped MXFP8 requantize test: ", + static_cast(dtype)); + } + return 0; +} + +struct OwnedGroupedTensor { + std::unique_ptr tensor; + test::CudaPtr<> rowwise_data; + test::CudaPtr<> columnwise_data; + test::CudaPtr<> rowwise_scale_inv; + test::CudaPtr<> columnwise_scale_inv; + test::CudaPtr first_dims; + test::CudaPtr last_dims; + test::CudaPtr offsets; + size_t data_bytes = 0; + size_t rowwise_scale_bytes = 0; + size_t columnwise_scale_bytes = 0; + + NVTEGroupedTensor data() const { return tensor->data(); } +}; + +OwnedGroupedTensor make_grouped_tensor(const GroupShapeInfo &shape_info, DType dtype, + NVTEScalingMode scaling_mode, bool rowwise, + bool columnwise, bool swizzled = false) { + OwnedGroupedTensor result; + result.tensor = std::make_unique( + shape_info.num_tensors, shape_info.logical_shape, scaling_mode); + + const std::vector flat_data_shape = {shape_info.logical_shape[0] * shape_info.logical_shape[1]}; + result.data_bytes = shape_info.total_elements * dtype_size(dtype); + if (rowwise) { + result.rowwise_data = test::cuda_alloc(result.data_bytes); + result.tensor->set_rowwise_data(result.rowwise_data.get(), dtype, flat_data_shape); + } + if (columnwise) { + result.columnwise_data = test::cuda_alloc(result.data_bytes); + result.tensor->set_columnwise_data(result.columnwise_data.get(), dtype, flat_data_shape); + } + + if (scaling_mode == NVTE_MXFP8_1D_SCALING) { + if (rowwise) { + result.rowwise_scale_bytes = shape_info.rowwise_scale_elements; + result.rowwise_scale_inv = test::cuda_alloc(result.rowwise_scale_bytes); + result.tensor->set_rowwise_scale_inv( + result.rowwise_scale_inv.get(), DType::kFloat8E8M0, + std::vector{shape_info.rowwise_scale_elements}); + } + if (columnwise) { + result.columnwise_scale_bytes = shape_info.columnwise_scale_elements; + result.columnwise_scale_inv = test::cuda_alloc(result.columnwise_scale_bytes); + result.tensor->set_columnwise_scale_inv( + result.columnwise_scale_inv.get(), DType::kFloat8E8M0, + std::vector{shape_info.columnwise_scale_elements}); + } + result.tensor->set_with_gemm_swizzled_scales(swizzled); + } + + const std::vector dims_shape = {shape_info.num_tensors}; + if (!shape_info.same_first_dim) { + const size_t bytes = shape_info.num_tensors * sizeof(int64_t); + result.first_dims = test::cuda_alloc(bytes); + NVTE_CHECK_CUDA(cudaMemcpy(result.first_dims.get(), shape_info.first_dims.data(), bytes, + cudaMemcpyHostToDevice)); + result.tensor->set_first_dims(result.first_dims.get(), DType::kInt64, dims_shape); + } + if (!shape_info.same_last_dim) { + const size_t bytes = shape_info.num_tensors * sizeof(int64_t); + result.last_dims = test::cuda_alloc(bytes); + NVTE_CHECK_CUDA(cudaMemcpy(result.last_dims.get(), shape_info.last_dims.data(), bytes, + cudaMemcpyHostToDevice)); + result.tensor->set_last_dims(result.last_dims.get(), DType::kInt64, dims_shape); + } + if (!shape_info.same_first_dim || !shape_info.same_last_dim) { + const size_t bytes = (shape_info.num_tensors + 1) * sizeof(int64_t); + result.offsets = test::cuda_alloc(bytes); + NVTE_CHECK_CUDA(cudaMemcpy(result.offsets.get(), shape_info.offsets.data(), bytes, + cudaMemcpyHostToDevice)); + result.tensor->set_tensor_offsets( + result.offsets.get(), DType::kInt64, + std::vector{shape_info.num_tensors + 1}); + } + + return result; +} + +void fill_source(OwnedGroupedTensor &source, const GroupShapeInfo &shape_info) { + std::vector host_data(shape_info.total_elements); + for (size_t i = 0; i < host_data.size(); ++i) { + const int value = static_cast((i * 17 + i / 13) % 251) - 125; + host_data[i] = static_cast(static_cast(value) / 16.0f); + } + NVTE_CHECK_CUDA(cudaMemcpy(source.rowwise_data.get(), host_data.data(), source.data_bytes, + cudaMemcpyHostToDevice)); +} + +void fill_output_sentinel(OwnedGroupedTensor &tensor) { + if (tensor.rowwise_data != nullptr) { + NVTE_CHECK_CUDA(cudaMemset(tensor.rowwise_data.get(), 0xA5, tensor.data_bytes)); + } + if (tensor.columnwise_data != nullptr) { + NVTE_CHECK_CUDA(cudaMemset(tensor.columnwise_data.get(), 0xA5, tensor.data_bytes)); + } + if (tensor.rowwise_scale_inv != nullptr) { + NVTE_CHECK_CUDA( + cudaMemset(tensor.rowwise_scale_inv.get(), 0xA5, tensor.rowwise_scale_bytes)); + } + if (tensor.columnwise_scale_inv != nullptr) { + NVTE_CHECK_CUDA(cudaMemset(tensor.columnwise_scale_inv.get(), 0xA5, + tensor.columnwise_scale_bytes)); + } +} + +void expect_bytes_equal(const std::string &name, const void *actual_dptr, + const void *expected_dptr, size_t bytes) { + std::vector actual(bytes); + std::vector expected(bytes); + NVTE_CHECK_CUDA( + cudaMemcpy(actual.data(), actual_dptr, bytes, cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA( + cudaMemcpy(expected.data(), expected_dptr, bytes, cudaMemcpyDeviceToHost)); + + for (size_t i = 0; i < bytes; ++i) { + ASSERT_EQ(actual[i], expected[i]) + << name << " mismatch at byte " << i << ": actual=" + << static_cast(actual[i]) << ", expected=" << static_cast(expected[i]); + } +} + +using RequantizeParam = std::tuple; + +class GroupedRequantizeMXFP8TestSuite : public ::testing::TestWithParam {}; + +TEST(GroupedRequantizeTest, RejectsUnsupportedScalingModes) { + const std::vector shape = {128, 128}; + GroupedTensorWrapper delayed(1, shape, NVTE_DELAYED_TENSOR_SCALING); + GroupedTensorWrapper mxfp8(1, shape, NVTE_MXFP8_1D_SCALING); + + EXPECT_THROW(nvte_group_requantize(delayed.data(), mxfp8.data(), nullptr, nullptr, 0), + std::runtime_error); + EXPECT_THROW(nvte_group_requantize(mxfp8.data(), delayed.data(), nullptr, nullptr, 0), + std::runtime_error); +} + +TEST(GroupedRequantizeTest, OptionalRowwiseOutputsAndAliasRules) { + if (test::getDeviceComputeCapability() < test::blackwellComputeCapability) { + GTEST_SKIP(); + } + + const GroupShapeInfo shape_info = make_shape_info({"Single", {{128, 128}}}); + const std::vector data_shape = {shape_info.total_elements}; + const std::vector scale_shape = {shape_info.rowwise_scale_elements}; + auto source = make_grouped_tensor(shape_info, DType::kBFloat16, + NVTE_DELAYED_TENSOR_SCALING, + /*rowwise=*/true, /*columnwise=*/false); + auto input = make_grouped_tensor(shape_info, DType::kFloat8E4M3, + NVTE_MXFP8_1D_SCALING, + /*rowwise=*/true, /*columnwise=*/false); + fill_source(source, shape_info); + nvte_group_quantize(source.data(), input.data(), nullptr, 0); + + auto output = make_grouped_tensor(shape_info, DType::kFloat8E4M3, + NVTE_MXFP8_1D_SCALING, + /*rowwise=*/false, /*columnwise=*/true, + /*swizzled=*/true); + // Both rowwise output fields may be absent, including with an unallocated + // optional dequantized descriptor (the faster path must still be selected). + TensorWrapper unallocated; + nvte_group_requantize(input.data(), output.data(), unallocated.data(), nullptr, 0); + + auto dequantized = test::cuda_alloc(shape_info.total_elements * sizeof(bf16)); + TensorWrapper dequantized_nvte(dequantized.get(), shape_info.logical_shape, DType::kBFloat16); + nvte_group_requantize(input.data(), output.data(), dequantized_nvte.data(), nullptr, 0); + output.tensor->set_with_gemm_swizzled_scales(false); + nvte_group_requantize(input.data(), output.data(), dequantized_nvte.data(), nullptr, 0); + output.tensor->set_with_gemm_swizzled_scales(true); + + output.tensor->set_rowwise_data(input.rowwise_data.get(), DType::kFloat8E4M3, + data_shape); + // The data may alias the input even when no rowwise scale output is requested. + nvte_group_requantize(input.data(), output.data(), nullptr, nullptr, 0); + + auto unrelated_data = test::cuda_alloc(shape_info.total_elements); + output.tensor->set_rowwise_data(unrelated_data.get(), DType::kFloat8E4M3, + data_shape); + EXPECT_THROW(nvte_group_requantize(input.data(), output.data(), nullptr, nullptr, 0), + std::runtime_error); + output.tensor->set_rowwise_data(input.rowwise_data.get(), DType::kFloat8E4M3, + data_shape); + + output.tensor->set_rowwise_scale_inv(input.rowwise_scale_inv.get(), + DType::kFloat8E8M0, scale_shape); + EXPECT_THROW(nvte_group_requantize(input.data(), output.data(), nullptr, nullptr, 0), + std::runtime_error); + + output.tensor->set_with_gemm_swizzled_scales(false); + nvte_group_requantize(input.data(), output.data(), nullptr, nullptr, 0); + + auto unrelated_scales = test::cuda_alloc(shape_info.rowwise_scale_elements); + output.tensor->set_rowwise_scale_inv(unrelated_scales.get(), + DType::kFloat8E8M0, scale_shape); + EXPECT_THROW(nvte_group_requantize(input.data(), output.data(), nullptr, nullptr, 0), + std::runtime_error); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); +} + +TEST(GroupedRequantizeTest, BF16IntermediateIgnoresFastMath) { + if (test::getDeviceComputeCapability() < test::blackwellComputeCapability) { + GTEST_SKIP(); + } + const auto shape_info = make_shape_info({"Single", {{128, 128}}}); + auto input = make_grouped_tensor(shape_info, DType::kFloat8E4M3, NVTE_MXFP8_1D_SCALING, + /*rowwise=*/true, /*columnwise=*/false); + // These decoded values are below half the smallest BF16 subnormal but remain + // nonzero in FP32, so this checks the intermediate precision as well as flags. + std::vector data(shape_info.total_elements, static_cast(3.0f / 512.0f)); + NVTE_CHECK_CUDA(cudaMemcpy(input.rowwise_data.get(), data.data(), data.size(), + cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemset(input.rowwise_scale_inv.get(), 0, shape_info.rowwise_scale_elements)); + auto baseline = make_grouped_tensor(shape_info, DType::kFloat8E4M3, NVTE_MXFP8_1D_SCALING, + /*rowwise=*/false, /*columnwise=*/true); + auto dequantized = test::cuda_alloc(shape_info.total_elements * sizeof(bf16)); + TensorWrapper dequantized_nvte(dequantized.get(), shape_info.logical_shape, DType::kBFloat16); + nvte_group_requantize(input.data(), baseline.data(), dequantized_nvte.data(), nullptr, 0); + std::vector decoded(shape_info.total_elements); + NVTE_CHECK_CUDA(cudaMemcpy(decoded.data(), dequantized.get(), decoded.size() * sizeof(bf16), + cudaMemcpyDeviceToHost)); + for (const auto value : decoded) { + ASSERT_EQ(static_cast(value), 0.0f); + } + for (const bool fast_math : {false, true}) { + QuantizationConfigWrapper config; + config.set_use_fast_math(fast_math); + auto actual = make_grouped_tensor(shape_info, DType::kFloat8E4M3, NVTE_MXFP8_1D_SCALING, + /*rowwise=*/false, /*columnwise=*/true); + for (const bool return_dequantized : {false, true}) { + nvte_group_requantize(input.data(), actual.data(), + return_dequantized ? dequantized_nvte.data() : nullptr, config, 0); + expect_bytes_equal("columnwise data", actual.columnwise_data.get(), + baseline.columnwise_data.get(), actual.data_bytes); + expect_bytes_equal("columnwise scales", actual.columnwise_scale_inv.get(), + baseline.columnwise_scale_inv.get(), actual.columnwise_scale_bytes); + } + } + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); +} + +TEST_P(GroupedRequantizeMXFP8TestSuite, MatchesDequantizeThenQuantize) { + if (test::getDeviceComputeCapability() < test::blackwellComputeCapability) { + GTEST_SKIP(); + } + + const auto &test_case = std::get<0>(GetParam()); + const DType input_dtype = std::get<1>(GetParam()); + const bool return_dequantized = std::get<2>(GetParam()); + const bool output_swizzled = std::get<3>(GetParam()); + const DType intermediate_dtype = DType::kBFloat16; + const GroupShapeInfo shape_info = make_shape_info(test_case); + + // Build a production-like wire tensor: high precision -> grouped rowwise + // MXFP8 with compact scales. This is the only input representation accepted + // by nvte_group_requantize. + auto source = make_grouped_tensor(shape_info, DType::kBFloat16, + NVTE_DELAYED_TENSOR_SCALING, + /*rowwise=*/true, /*columnwise=*/false); + auto input = make_grouped_tensor(shape_info, input_dtype, NVTE_MXFP8_1D_SCALING, + /*rowwise=*/true, /*columnwise=*/false, + /*swizzled=*/false); + fill_source(source, shape_info); + fill_output_sentinel(input); + nvte_group_quantize(source.data(), input.data(), nullptr, 0); + + // Code under test. + auto actual = make_grouped_tensor(shape_info, DType::kFloat8E4M3, + NVTE_MXFP8_1D_SCALING, + /*rowwise=*/false, /*columnwise=*/true, + output_swizzled); + actual.rowwise_scale_bytes = shape_info.rowwise_scale_elements; + if (output_swizzled) { + actual.rowwise_scale_inv = test::cuda_alloc(actual.rowwise_scale_bytes); + } + actual.tensor->set_rowwise_data(input.rowwise_data.get(), input_dtype, + std::vector{shape_info.total_elements}); + actual.tensor->set_rowwise_scale_inv( + output_swizzled ? actual.rowwise_scale_inv.get() : input.rowwise_scale_inv.get(), + DType::kFloat8E8M0, + std::vector{shape_info.rowwise_scale_elements}); + fill_output_sentinel(actual); + test::CudaPtr<> dequantized_out; + TensorWrapper dequantized_nvte; + if (return_dequantized) { + dequantized_out = test::cuda_alloc(shape_info.total_elements * sizeof(bf16)); + NVTE_CHECK_CUDA(cudaMemset(dequantized_out.get(), 0xA5, + shape_info.total_elements * sizeof(bf16))); + dequantized_nvte.set_rowwise_data(dequantized_out.get(), DType::kBFloat16, + shape_info.logical_shape); + } + nvte_group_requantize(input.data(), actual.data(), + return_dequantized ? dequantized_nvte.data() : nullptr, nullptr, 0); + + // Reference descriptors delimit the live prefix, while allocations retain + // full capacity so byte comparisons also check untouched tails. + auto reference_shape_info = shape_info; + if (test_case.capacity_rows != 0) { + reference_shape_info.logical_shape[0] = shape_info.offsets.back() / shape_info.logical_shape[1]; + } + auto reference_input = make_grouped_tensor(reference_shape_info, input_dtype, + NVTE_MXFP8_1D_SCALING, + /*rowwise=*/false, /*columnwise=*/false); + reference_input.tensor->set_rowwise_data( + input.rowwise_data.get(), input_dtype, + std::vector{reference_shape_info.logical_shape[0] * reference_shape_info.logical_shape[1]}); + reference_input.tensor->set_rowwise_scale_inv(input.rowwise_scale_inv.get(), DType::kFloat8E8M0, + std::vector{shape_info.rowwise_scale_elements}); + + // Exact reference: requantization uses BF16 regardless of the fast-math setting. + auto dequantized = make_grouped_tensor(reference_shape_info, intermediate_dtype, + NVTE_DELAYED_TENSOR_SCALING, + /*rowwise=*/true, /*columnwise=*/false); + fill_output_sentinel(dequantized); + nvte_group_dequantize(reference_input.data(), dequantized.data(), 0); + + auto reference_compact = make_grouped_tensor( + reference_shape_info, DType::kFloat8E4M3, NVTE_MXFP8_1D_SCALING, + /*rowwise=*/false, /*columnwise=*/true, /*swizzled=*/false); + fill_output_sentinel(reference_compact); + nvte_group_quantize(dequantized.data(), reference_compact.data(), nullptr, 0); + + std::unique_ptr reference_swizzled; + if (output_swizzled) { + reference_swizzled = std::make_unique(make_grouped_tensor( + reference_shape_info, DType::kFloat8E4M3, NVTE_MXFP8_1D_SCALING, + /*rowwise=*/false, /*columnwise=*/true, /*swizzled=*/true)); + fill_output_sentinel(*reference_swizzled); + nvte_swizzle_grouped_scaling_factors(reference_compact.data(), + reference_swizzled->data(), 0); + } + + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + + if (return_dequantized) { + expect_bytes_equal("dequantized data", dequantized_out.get(), + dequantized.rowwise_data.get(), dequantized.data_bytes); + } + + // Scale layout does not affect FP8 data. Only the scale reference changes + // when the target requests GEMM-swizzled scales. + expect_bytes_equal("columnwise data", actual.columnwise_data.get(), + reference_compact.columnwise_data.get(), actual.data_bytes); + const void *expected_scales = output_swizzled + ? reference_swizzled->columnwise_scale_inv.get() + : reference_compact.columnwise_scale_inv.get(); + expect_bytes_equal("columnwise scales", actual.columnwise_scale_inv.get(), + expected_scales, actual.columnwise_scale_bytes); + std::unique_ptr reference_rowwise_swizzled; + const void *expected_rowwise_scales = input.rowwise_scale_inv.get(); + if (output_swizzled) { + reference_rowwise_swizzled = std::make_unique(make_grouped_tensor( + reference_shape_info, input_dtype, NVTE_MXFP8_1D_SCALING, + /*rowwise=*/true, /*columnwise=*/false, /*swizzled=*/true)); + fill_output_sentinel(*reference_rowwise_swizzled); + nvte_swizzle_grouped_scaling_factors(reference_input.data(), + reference_rowwise_swizzled->data(), 0); + expected_rowwise_scales = reference_rowwise_swizzled->rowwise_scale_inv.get(); + } + expect_bytes_equal("rowwise scales", actual.tensor->get_rowwise_scale_inv().data_ptr, + expected_rowwise_scales, actual.rowwise_scale_bytes); + EXPECT_EQ(actual.tensor->get_rowwise_data().data_ptr, input.rowwise_data.get()); + EXPECT_EQ(actual.tensor->get_with_gemm_swizzled_scales(), output_swizzled); +} + +const std::vector kGroupShapeCases = { + {"Single", {{1024, 512}}}, + {"UniformSmall", {{512, 512}, {512, 512}, {512, 512}, {512, 512}}}, + {"VaryingFirstSmall", {{256, 256}, {1024, 256}, {128, 256}, {640, 256}}}, + {"EmptyFront", {{0, 512}, {512, 512}, {256, 512}}}, + {"EmptyAdjacent", {{256, 384}, {0, 384}, {0, 384}, {512, 384}}}, + {"EmptyEnd", {{512, 512}, {256, 512}, {0, 512}}}, + {"CapacityTail", {{512, 8192}, {256, 8192}, {512, 8192}}, 2048}, + {"EmptyGroupsAndTail", {{0, 512}, {512, 512}, {0, 512}, {256, 512}}, 1024}, + {"UnsupportedCuTeStrideAndTail", {{0, 640}, {384, 640}, {128, 640}}, 1024}, + {"SameBothDims_1024x4096", + {{1024, 4096}, {1024, 4096}, {1024, 4096}, {1024, 4096}}}, + {"SameBothDims_2048x8192", + {{2048, 8192}, {2048, 8192}, {2048, 8192}, {2048, 8192}}}, + {"SameBothDims_4096x16384", + {{4096, 16384}, {4096, 16384}, {4096, 16384}, {4096, 16384}}}, + {"VaryingFirstDim_512to2048x4096", + {{512, 4096}, {512, 4096}, {1024, 4096}, {2048, 4096}}}, + {"VaryingFirstDim_1024to4096x8192", + {{1024, 8192}, {1024, 8192}, {2048, 8192}, {4096, 8192}}}, + {"VaryingFirstDim_2048to8192x16384", + {{2048, 16384}, {2048, 16384}, {4096, 16384}, {8192, 16384}}}, + // An empty member in the middle must not terminate the persistent work loop. + {"VaryingFirstDimWithEmpty", + {{128, 256}, {0, 256}, {384, 256}, {512, 256}}}, + {"VaryingLastDim", + {{256, 128}, {256, 384}, {256, 640}}}, + {"VaryingBothDims", + {{128, 128}, {256, 384}, {512, 640}}}, +}; + +std::string make_test_name( + const testing::TestParamInfo &info) { + const auto &test_case = std::get<0>(info.param); + const DType input_dtype = std::get<1>(info.param); + const bool return_dequantized = std::get<2>(info.param); + const bool output_swizzled = std::get<3>(info.param); + + std::string name = test_case.name; + name += input_dtype == DType::kFloat8E4M3 ? "_E4M3" : "_E5M2"; + name += return_dequantized ? "_WithDequantized" : "_WithoutDequantized"; + name += output_swizzled ? "_Swizzled" : "_Compact"; + return name; +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, GroupedRequantizeMXFP8TestSuite, + ::testing::Combine( + ::testing::ValuesIn(kGroupShapeCases), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Bool(), + ::testing::Bool() + ), + make_test_name); + +} // namespace diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py index c0cb0aa09aa..7a1edc071e9 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py +++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py @@ -554,6 +554,8 @@ def check_prequantized_requantize_versus_reference( M: int, N: int, split_sections: list[int], + return_dequantized: bool = True, + offsets_only: bool = False, ) -> None: """Run the pre-quantized requantize path and check both directions against a reference. @@ -613,9 +615,10 @@ def check_prequantized_requantize_versus_reference( wire, make_op_quantizer(columnwise=True), num_groups, - split_section_tensor, + None if offsets_only else split_section_tensor, te.DType.kBFloat16, - return_dequantized=True, + tensor_offsets=wire.tensor_offsets if offsets_only else None, + return_dequantized=return_dequantized, ) assert wire.columnwise_data is not None, "columnwise data must be built" @@ -631,7 +634,10 @@ def check_prequantized_requantize_versus_reference( # The returned dequantized tensor is what bias gradients are reduced from. Compare only the # live rows: both this and the reference allocate M rows but write only the covered ones, and # their tails are separate uninitialized allocations. - torch.testing.assert_close(dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0) + if return_dequantized: + torch.testing.assert_close(dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0) + else: + assert dequantized is None # The rowwise DATA must pass through untouched; only its scales are re-laid-out. torch.testing.assert_close(wire.rowwise_data, rowwise_data_before, atol=0.0, rtol=0.0) @@ -681,6 +687,32 @@ def check_prequantized_requantize_versus_reference( ) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N, split_sections", + [ + (1024, 256, [256, 256, 256, 256]), + (1024, 256, [128, 0, 256, 128]), + # CuTeDSL rejects this scale stride; the common API must fall back to CUDA. + (1024, 640, [256, 256, 256, 256]), + ], +) +@pytest.mark.parametrize("offsets_only", [False, True]) +@pytest.mark.parametrize("return_dequantized", [False, True]) +def test_prequantized_requantize_optional_dequantized_output( + M: int, N: int, split_sections: list[int], offsets_only: bool, return_dequantized: bool +) -> None: + """The grouped API path matches the reference with uniform and paged groups.""" + check_prequantized_requantize_versus_reference( + x_dtype=torch.bfloat16, + M=M, + N=N, + split_sections=split_sections, + return_dequantized=return_dequantized, + offsets_only=offsets_only, + ) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, N", diff --git a/tests/pytorch/mxfp8/test_mxfp8_requantize_cutedsl.py b/tests/pytorch/mxfp8/test_mxfp8_requantize_cutedsl.py new file mode 100644 index 00000000000..81d12771566 --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_requantize_cutedsl.py @@ -0,0 +1,199 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Validate compact-input requantization against independent PyTorch arithmetic.""" + +import pytest +import torch + +cutlass = pytest.importorskip("cutlass") +cute = pytest.importorskip("cutlass.cute") +cuda = pytest.importorskip("cuda.bindings.driver") + +from cutlass.cute.runtime import from_dlpack +from transformer_engine.common.CuTeDSL.cast.mxfp8 import requantize_mxfp8 +from transformer_engine.common.CuTeDSL.cast.mxfp8.requantize_mxfp8 import GroupedRequantize +from mxfp8_utils import swizzle_mxfp8_scale + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), + reason="Grouped MXFP8 requantization requires SM100+", +) + + +@pytest.fixture(params=["auto", "portable"]) +def decode_backend(request, monkeypatch): + """Exercise the old-compiler sequence even when the compiler supports PTX 9.2.""" + if request.param == "portable": + monkeypatch.setattr(requantize_mxfp8, "_supports_scaled_bf16_conversion", lambda: False) + + +def columnwise_reference(values): + """Quantize BF16-decoded values in independent 32-row blocks.""" + rows, hidden = values.shape + blocks = values.float().reshape(rows // 32, 32, hidden) + amax = torch.where(torch.isnan(blocks), 0.0, blocks.abs()).amax(1) + exponent = torch.ceil(torch.log2(amax / 448)).clamp(-127, 127) + exponent = torch.where(amax == 0, -127, exponent) + inverse = torch.exp2(-exponent).clamp(max=torch.finfo(torch.float32).max) + quantized = (blocks * inverse[:, None]).clamp(-448, 448).to(torch.float8_e4m3fn) + return quantized.reshape(rows, hidden).view(torch.uint8), (exponent + 127).to(torch.uint8) + + +@pytest.mark.parametrize("return_dequantized", [False, True]) +@pytest.mark.parametrize("hidden", [128, 256, 384, 512, 4096, 7168]) +@pytest.mark.parametrize("input_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) +@pytest.mark.parametrize("layout", ["compact", "column_only", "both", "uniform"]) +@pytest.mark.usefixtures("decode_backend") +def test_compact_input_contract(hidden, input_dtype, layout, return_dequantized): + """Check both scale layouts, optional output, group boundaries and capacity tails.""" + torch.manual_seed(42) + rows = 1024 + sizes = [256] * 4 if layout == "uniform" else [128, 0, 256, 128, 0] + live = sum(sizes) + data = (torch.randn(rows, hidden, device="cuda") * 8).to(input_dtype) + scales = torch.randint(100, 155, (rows, hidden // 32), device="cuda", dtype=torch.uint8) + decoded = (data.float() * torch.exp2(scales.float() - 127).repeat_interleave(32, 1)).to( + torch.bfloat16 + ) + offsets = torch.tensor([0] + sizes, device="cuda", dtype=torch.int64).cumsum(0) * hidden + dequantized = torch.full((rows, hidden), -123, device="cuda", dtype=torch.bfloat16) + dst = torch.full((rows, hidden), 0xA5, device="cuda", dtype=torch.uint8) + row_sf = torch.full_like(scales.flatten(), 0xA5) + col_sf = torch.full_like(row_sf, 0xA5) + swizzled = layout != "compact" + row_output = layout in ("both", "uniform") + dtype = cutlass.Float8E4M3FN if input_dtype == torch.float8_e4m3fn else cutlass.Float8E5M2 + kernel = GroupedRequantize( + hidden, + len(sizes), + rows, + input_dtype=dtype, + swizzled=swizzled, + rowwise_output=row_output, + return_dequantized=return_dequantized, + uniform_rows=256 if layout == "uniform" else 0, + sm_count=torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count, + ) + tensors = ( + data, + scales, + None if layout == "uniform" else offsets, + dst.view(torch.float8_e4m3fn), + row_sf if row_output else None, + col_sf, + dequantized if return_dequantized else None, + ) + args = tuple( + from_dlpack(tensor, assumed_align=16) if tensor is not None else None for tensor in tensors + ) + (cuda.CUstream(torch.cuda.current_stream().cuda_stream),) + compiled = cute.compile(kernel, *args) + compiled(*args) + torch.cuda.synchronize() + start = 0 + expected_data, expected_sf = [], [] + for size in sizes: + if size: + quantized, sf = columnwise_reference(decoded[start : start + size]) + expected_data.append(quantized.flatten()) + if swizzled: + sf = swizzle_mxfp8_scale(size, hidden, sf, True) + expected_sf.append(sf.flatten()) + start += size + torch.testing.assert_close( + dst.flatten()[: live * hidden], torch.cat(expected_data), rtol=0, atol=0 + ) + torch.testing.assert_close( + col_sf[: live * hidden // 32], torch.cat(expected_sf), rtol=0, atol=0 + ) + if row_output: + expected_row = swizzle_mxfp8_scale(live, hidden, scales[:live], False).flatten() + torch.testing.assert_close(row_sf[: expected_row.numel()], expected_row, rtol=0, atol=0) + if return_dequantized: + torch.testing.assert_close(dequantized[:live], decoded[:live], rtol=0, atol=0) + assert (dequantized[live if return_dequantized else 0 :] == -123).all() + assert (dst[live:] == 0xA5).all() + assert (col_sf[live * hidden // 32 :] == 0xA5).all() + assert (row_sf[(live * hidden // 32 if row_output else 0) :] == 0xA5).all() + + +@pytest.mark.parametrize("return_dequantized", [False, True]) +@pytest.mark.parametrize("scale", [0, 1, 127, 254, 255]) +@pytest.mark.parametrize("input_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) +@pytest.mark.usefixtures("decode_backend") +def test_extreme_scale_bf16_decode(scale, return_dequantized, input_dtype): + """Document tiny-value BF16 rounding and check saturated/NaN scale behavior.""" + hidden, rows = 256, 128 + # Cover every FP8 encoding, including signed zero, subnormals, NaNs and E5M2 infinities. + data = torch.arange(hidden, device="cuda", dtype=torch.uint8).repeat(rows, 1).view(input_dtype) + sf = torch.full((rows, hidden // 32), scale, device="cuda", dtype=torch.uint8) + dst = torch.empty((rows, hidden), device="cuda", dtype=torch.float8_e4m3fn) + col_sf = torch.empty(rows * hidden // 32, device="cuda", dtype=torch.uint8) + kernel = GroupedRequantize( + hidden, + 1, + rows, + input_dtype=( + cutlass.Float8E4M3FN if input_dtype == torch.float8_e4m3fn else cutlass.Float8E5M2 + ), + swizzled=False, + rowwise_output=False, + uniform_rows=rows, + return_dequantized=return_dequantized, + ) + dequantized = torch.empty((rows, hidden), device="cuda", dtype=torch.bfloat16) + args = tuple( + from_dlpack(tensor, assumed_align=16) if tensor is not None else None + for tensor in ( + data, + sf, + None, + dst, + None, + col_sf, + dequantized if return_dequantized else None, + ) + ) + (cuda.CUstream(torch.cuda.current_stream().cuda_stream),) + cute.compile(kernel, *args)(*args) + decoded = (data.float() * (float("nan") if scale == 255 else 2.0 ** (scale - 127))).to( + torch.bfloat16 + ) + if return_dequantized: + torch.testing.assert_close(dequantized, decoded, rtol=0, atol=0, equal_nan=True) + expected, expected_sf = columnwise_reference(decoded) + torch.testing.assert_close(dst.view(torch.uint8), expected, rtol=0, atol=0) + torch.testing.assert_close(col_sf, expected_sf.flatten(), rtol=0, atol=0) + + +def test_reject_unsupported_scale_stride(): + """Do not create an invalid TMA descriptor for a wide nonaligned scale stride.""" + with pytest.raises(ValueError, match="TMA row stride"): + GroupedRequantize(640, 1, 128) + + +@pytest.mark.parametrize("return_dequantized", [False, True]) +def test_empty_groups_leave_capacity_untouched(return_dequantized): + """A device metadata update can select no live rows without touching outputs.""" + rows, hidden = 256, 512 + data = torch.zeros((rows, hidden), device="cuda", dtype=torch.uint8) + sf = torch.zeros((rows, hidden // 32), device="cuda", dtype=torch.uint8) + offsets = torch.zeros(4, device="cuda", dtype=torch.int64) + outputs = [ + torch.full_like(data, 0xA5), + torch.full_like(sf.flatten(), 0xA5), + torch.full_like(sf.flatten(), 0xA5), + ] + dequantized = torch.full((rows, hidden), -123, device="cuda", dtype=torch.bfloat16) + tensors = (data, sf, offsets, *outputs, dequantized if return_dequantized else None) + args = tuple( + from_dlpack(tensor, assumed_align=16) if tensor is not None else None for tensor in tensors + ) + (cuda.CUstream(torch.cuda.current_stream().cuda_stream),) + kernel = GroupedRequantize(hidden, 3, rows, return_dequantized=return_dequantized) + cute.compile(kernel, *args)(*args) + torch.cuda.synchronize() + for output in outputs: + assert (output == 0xA5).all() + assert (dequantized == -123).all() diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 077a9a2b3a1..76cca1e3531 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -362,7 +362,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources cast/cast_dbias.cu cast/cast_grouped.cu cast/cast_grouped_dbias.cu - cast/fused_group_requantize.cu + cast/mxfp8/requantize_mxfp8.cu fused_rope/fused_mla_kv_rope_mxfp8.cu gemm/cutlass_grouped_gemm.cu hadamard_transform/group_hadamard_transform.cu diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..3d47d0652f8 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -13,6 +13,9 @@ from transformer_engine.common.CuTeDSL.cast.mxfp8.quantize_mxfp8 import ( get_mxfp8_quantization_function, ) +from transformer_engine.common.CuTeDSL.cast.mxfp8.requantize_mxfp8 import ( + get_mxfp8_requantization_function, +) def register_cutedsl_backends(): @@ -22,3 +25,6 @@ def register_cutedsl_backends(): tvm_ffi.register_global_func( "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True ) + tvm_ffi.register_global_func( + "get_mxfp8_requantization_function", get_mxfp8_requantization_function, override=True + ) diff --git a/transformer_engine/common/CuTeDSL/cast/mxfp8/requantize_mxfp8.py b/transformer_engine/common/CuTeDSL/cast/mxfp8/requantize_mxfp8.py new file mode 100644 index 00000000000..00cac332891 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/cast/mxfp8/requantize_mxfp8.py @@ -0,0 +1,1011 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped MXFP8 requantization with BF16 intermediates. + +The input scales use TE's compact rowwise layout. One TMA pipeline loads both +payload and scales; consumers decode to BF16, compute columnwise scales and +emit row-major E4M3 payloads. Optional BF16 output reuses the decoded registers +without additional shared storage. Scale layouts and optional rowwise scale output +follow nvte_group_requantize. No framework is imported by this module. +""" + +import logging +import cuda.bindings.driver as cuda +import cutlass +from cutlass import cute +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import Float32, Int32, Int64, T, Uint8, dsl_user_op +from cutlass._mlir.dialects import arith, llvm +from cutlass.cute.typing import AddressSpace +from cutlass.memory import SmemAllocator +import tvm_ffi + +logger = logging.getLogger("transformer_engine.cutedsl.requantize") + + +def _address_value(pointer_or_address, *, loc=None, ip=None): + if isinstance(pointer_or_address, Int64): + return pointer_or_address.ir_value() + return pointer_or_address.toint(loc=loc, ip=ip).ir_value() + + +@dsl_user_op +def tma_load_1d( + destination_smem, + source_gmem, + mbarrier_smem, + num_bytes, + *, + loc=None, + ip=None, +) -> None: + """Issue a 1D GMEM-to-SMEM bulk copy.""" + llvm.inline_asm( + None, + [ + destination_smem.toint(loc=loc, ip=ip).ir_value(), + _address_value(source_gmem, loc=loc, ip=ip), + num_bytes.ir_value(), + mbarrier_smem.toint(loc=loc, ip=ip).ir_value(), + ], + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1], $2, [$3];", + "r,l,r,r", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_async_bulk_s2g(destination_gmem, source_smem, num_bytes, *, loc=None, ip=None) -> None: + """Issue a 1D SMEM-to-GMEM bulk copy; the caller commits the group.""" + llvm.inline_asm( + None, + [ + _address_value(destination_gmem, loc=loc, ip=ip), + source_smem.toint(loc=loc, ip=ip).ir_value(), + num_bytes.ir_value(), + ], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _fp8x2_mnemonic(fp8_type) -> str: + if fp8_type is cutlass.Float8E4M3FN: + return "e4m3x2" + if fp8_type is cutlass.Float8E5M2: + return "e5m2x2" + raise TypeError(f"unsupported FP8 type {fp8_type}") + + +def _supports_scaled_bf16_conversion() -> bool: + """Query the DSL compiler, rather than the installed CUDA toolkit or driver.""" + target_version = getattr(cutlass, "target_version", None) + # Older DSL releases without this query use the portable instruction sequence. + return target_version is not None and target_version(min_version="13.2") + + +@dsl_user_op +def cvt_scaled_up_bf16x2( + pair_b32, scale_b32, half: int, fp8_type, use_scaled_cvt: bool, *, loc=None, ip=None +) -> Int32: + """Decode two FP8 values with E8M0 scales, using PTX 9.2 only when supported.""" + mn = _fp8x2_mnemonic(fp8_type) + if use_scaled_cvt: + asm = ( + "{\n" + " .reg .b16 a0,a1,s0,s1;\n" + " mov.b32 {a0,a1}, $1;\n" + " mov.b32 {s0,s1}, $2;\n" + f" cvt.rn.scaled::n2::ue8m0.bf16x2.{mn} $0, a{half}, s0;\n" + "}" + ) + else: + # Match dequantize_mxfp8_2x in requantize_mxfp8.cu: FP8 -> FP16 -> + # BF16 is exact, then BF16 multiplication rounds the scaled result. + # E8M0 code 0 is 2**-127 (a BF16 subnormal); code 255 is NaN. + asm = ( + "{\n" + " .reg .b16 a0,a1,f0,f1,b0,b1;\n" + " .reg .b32 values_f16x2,values_bf16x2,s,scale_bits,scale_x2;\n" + " .reg .pred is_zero,is_nan;\n" + " mov.b32 {a0,a1}, $1;\n" + " and.b32 s, $2, 0xff;\n" + " shl.b32 scale_bits, s, 7;\n" + " setp.eq.u32 is_zero, s, 0;\n" + " setp.eq.u32 is_nan, s, 255;\n" + " selp.b32 scale_bits, 0x0040, scale_bits, is_zero;\n" + " selp.b32 scale_bits, 0x7fff, scale_bits, is_nan;\n" + " shl.b32 scale_x2, scale_bits, 16;\n" + " or.b32 scale_x2, scale_x2, scale_bits;\n" + f" cvt.rn.f16x2.{mn} values_f16x2, a{half};\n" + " mov.b32 {f0,f1}, values_f16x2;\n" + " cvt.rn.bf16.f16 b0, f0;\n" + " cvt.rn.bf16.f16 b1, f1;\n" + " mov.b32 values_bf16x2, {b0,b1};\n" + " mul.rn.bf16x2 $0, values_bf16x2, scale_x2;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(pair_b32).ir_value(loc=loc, ip=ip), Int32(scale_b32).ir_value(loc=loc, ip=ip)], + asm, + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def cvt_dn_fp8x2_portable(v_bf16x2, inv_lo, inv_hi, fp8_type, *, loc=None, ip=None) -> Int32: + """Two hidden columns of one token, as BF16x2, plus those two columns' + exact FP32 reciprocal scales -> two FP8 bytes in bits [15:0] (byte 0 from + the low BF16 half, byte 1 from the high half). + + Same rounding as ``cvt_scaled_dn_fp8x2``, not an approximation of it: BF16 + -> FP32 is an exact left shift, the reciprocal is an exact power of two, and + ``cvt.rn.satfinite`` is the RNE-and-saturate the hardware instruction + applies. Taking a scale per half is what lets the caller skip the transpose + that the one-scale-per-pair hardware instruction forces. + """ + mn = _fp8x2_mnemonic(fp8_type) + asm = ( + "{\n" + " .reg .b32 a, b;\n" + " .reg .b16 q;\n" + " shl.b32 a, $1, 16;\n" + " and.b32 b, $1, 0xffff0000;\n" + " mul.f32 a, a, $2;\n" + " mul.f32 b, b, $3;\n" + # ``cvt d, a, b`` yields d[15:8] = cvt(a) and d[7:0] = cvt(b), so the + # HIGH column has to be the first source for byte 0 to be the low one. + f" cvt.rn.satfinite.{mn}.f32 q, b, a;\n" + " cvt.u32.u16 $0, q;\n" + "}" + ) + return Int32( + llvm.inline_asm( + T.i32(), + [ + Int32(v_bf16x2).ir_value(loc=loc, ip=ip), + Float32(inv_lo).ir_value(loc=loc, ip=ip), + Float32(inv_hi).ir_value(loc=loc, ip=ip), + ], + asm, + "=r,r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def max_xorsign_abs_bf16x2(a, b, *, loc=None, ip=None) -> Int32: + """Packed magnitude max of two BF16x2.""" + return Int32( + llvm.inline_asm( + T.i32(), + [Int32(a).ir_value(loc=loc, ip=ip), Int32(b).ir_value(loc=loc, ip=ip)], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e8m0_raw_from_bf16(bf16_bits, limit_exponent: int, *, loc=None, ip=None) -> Int32: + """E8M0 raw byte for a non-negative BF16 magnitude.""" + bits = Int32(bf16_bits) << Int32(16) + biased = (bits + Int32(0x1FFFFF - (limit_exponent << 23))) >> Int32(23) + return Int32( + arith.select( + (bits >= Int32(0x7F800000)).ir_value(loc=loc, ip=ip), + Int32(254).ir_value(loc=loc, ip=ip), + cutlass.max(Int32(0), biased).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def bits_f32(x: Int32, *, loc=None, ip=None) -> Float32: + """Interpret an integer register as IEEE FP32 bits.""" + return Float32(llvm.bitcast(T.f32(), Int32(x).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)) + + +class GroupedRequantize: + """Persistent, warp-specialized TMA requantization with compact input scales. + + A consumer lane owns four columns and 32 tokens. FP8 decoding and amax + reduction stay in BF16 registers. Payloads remain row-major. Both GEMM + scale layouts are emitted without a separate swizzle kernel. + """ + + TokensPerBlock = 32 + TILE_TOK = 128 + SfAtomNonK = 128 + SfAtomBytes = 512 + SfOutStride = 528 + ConsumerBarrierId = 1 + ProducerWarps = 1 + + def __init__( + self, + hidden, + num_groups, + capacity, + *, + input_dtype=cutlass.Float8E4M3FN, + swizzled=True, + rowwise_output=True, + return_dequantized=False, + uniform_rows=0, + tile_hidden=None, + stages=2, + grid=None, + sm_count=152, + ): + if hidden <= 0 or hidden % 128 or capacity <= 0 or capacity % 128 or num_groups < 1: + raise ValueError( + "Grouped MXFP8 requantization requires positive 128-aligned dimensions" + ) + if input_dtype not in (cutlass.Float8E4M3FN, cutlass.Float8E5M2): + raise ValueError("Input must be E4M3 or E5M2") + if hidden >= 512 and hidden % 512: + raise ValueError("Wide compact scale matrices require a 16-byte TMA row stride") + if uniform_rows and (uniform_rows % 128 or uniform_rows * num_groups != capacity): + raise ValueError("Uniform groups must be 128-aligned and fill the row capacity") + self.hidden = int(hidden) + self.num_experts = int(num_groups) + self.quant_dtype = input_dtype + self.use_scaled_cvt = _supports_scaled_bf16_conversion() + self.sf_dtype = cutlass.Uint8 + self.swizzled = bool(swizzled) + self.rowwise_output = bool(rowwise_output and swizzled) + self.return_dequantized = bool(return_dequantized) + self.uniform_rows = int(uniform_rows) + # BF16 stores benefit from spreading sparse workloads across more SMs + # when the doubled tile count still fits in one wave. + narrow_tile = hidden == 128 or ( + self.return_dequantized and 2 * (capacity // 128) * (hidden // 256) <= sm_count + ) + self.TILE_HID = tile_hidden or (128 if narrow_tile else 256) + if hidden % self.TILE_HID: + self.TILE_HID = 128 + self.ColsPerLane = 4 + self.sp_lw = 4 + self.NumStages = stages + self.TmaBoxHidU32 = self.TILE_HID // 4 + self.HidAtomsPerTile = self.TILE_HID // 128 + self.HidSegs = self.TILE_HID // 128 + self.ConsumerWarps = 4 * self.HidSegs + self.ThreadsPerCta = (1 + self.ConsumerWarps) * 32 + self.MaxNTid = max(576, self.ThreadsPerCta) + self._hidden_atoms = hidden // 128 + self.hidden_groups = hidden // self.TILE_HID + # TMA requires a 16-byte row stride. Narrow full-width matrices use a + # contiguous 1D bulk copy. Other tiles load at least 16 scale columns; + # consumers ignore the extra columns (TMA zero-fills the final tile). + self.scale_stride = hidden // 32 if hidden < 512 else max(16, self.TILE_HID // 32) + self.SfTileBytes = 128 * self.scale_stride + self.smem_sf_out_bytes = self.HidAtomsPerTile * self.SfOutStride + self._data_limit_exponent = 8 # Output is always E4M3, including E5M2 input. + tasks = capacity // 128 * self.hidden_groups + self.grid = grid or min(tasks, max(1, min(24, tasks // (4 * 2 * sm_count))) * 2 * sm_count) + span = 1 + while span < num_groups: + span <<= 1 + self._search_steps = tuple(1 << i for i in reversed(range(span.bit_length() - 1))) + self._search_needs_guard = (num_groups & (num_groups - 1)) != 0 + self.smem_bytes = ( + stages * (128 * self.TILE_HID + self.SfTileBytes) + + self.smem_sf_out_bytes + + (self.HidAtomsPerTile * 512 if self.rowwise_output else 0) + + (num_groups + 1) * 4 + + 2 * stages * 8 + + 384 + ) + if self.smem_bytes > 227 * 1024: + raise ValueError("Grouped requantization exceeds the shared-memory budget") + + @cute.jit + def __call__( + self, + src_data: cute.Tensor, + src_sf: cute.Tensor, + offsets: cute.Tensor, + dst_data: cute.Tensor, + dst_row_sf: cute.Tensor, + dst_col_sf: cute.Tensor, + dst_dequantized: cute.Tensor, + stream: cuda.CUstream, + ): + tok = cutlass.const_expr(self.TILE_TOK) + width = cutlass.const_expr(self.TILE_HID // 4) + src = cute.make_tensor( + cute.recast_ptr(src_data.iterator, dtype=cutlass.Uint32), + cute.make_layout((src_data.shape[0], self.hidden // 4), stride=(self.hidden // 4, 1)), + ) + dst = cute.make_tensor(cute.recast_ptr(dst_data.iterator, dtype=cutlass.Uint32), src.layout) + data_layout = cute.make_layout((tok, width), stride=(width, 1)) + load_atom, load_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), src, data_layout, (tok, width) + ) + store_atom, store_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), dst, data_layout, (tok, width) + ) + scale_atom, scale_tensor = None, None + if cutlass.const_expr(self.hidden >= 512): + scales = cute.make_tensor( + cute.recast_ptr(src_sf.iterator, dtype=cutlass.Uint8), + cute.make_layout( + (src_data.shape[0], self.hidden // 32), stride=(self.hidden // 32, 1) + ), + ) + scale_layout = cute.make_layout((tok, self.scale_stride), stride=(self.scale_stride, 1)) + scale_atom, scale_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), scales, scale_layout, (tok, self.scale_stride) + ) + self.ws_kernel( + src_sf, + offsets, + dst_row_sf, + dst_col_sf, + dst_dequantized, + load_atom, + load_tensor, + store_atom, + store_tensor, + scale_atom, + scale_tensor, + ).launch( + grid=(self.grid, 1, 1), + block=(self.ThreadsPerCta, 1, 1), + max_number_threads=(self.MaxNTid, 1, 1), + stream=stream, + ) + + @cute.kernel + def ws_kernel( + self, + src_sf: cute.Tensor, + offsets: cute.Tensor, + dst_row_sf: cute.Tensor, + dst_col_sf: cute.Tensor, + dst_dequantized: cute.Tensor, + load_atom=None, + load_tensor=None, + store_atom=None, + store_tensor=None, + scale_atom=None, + scale_tensor=None, + ): + """Pipeline TMA loads, consumer requantization and fused scale stores.""" + tok = cutlass.const_expr(self.TILE_TOK) + width = cutlass.const_expr(self.TILE_HID) + stages = cutlass.const_expr(self.NumStages) + sf_bytes = cutlass.const_expr(self.SfTileBytes) + tid, _, _ = cute.arch.thread_idx() + bid, _, _ = cute.arch.block_idx() + grid, _, _ = cute.arch.grid_dim() + warp, lane = tid // 32, tid % 32 + smem = SmemAllocator() + full = smem.allocate_array(cutlass.Int64, stages) + empty = smem.allocate_array(cutlass.Int64, stages) + table = smem.allocate_tensor(cutlass.Int32, cute.make_layout((self.num_experts + 1,)), 16) + sf_in = smem.allocate_array(cutlass.Uint8, stages * sf_bytes, byte_alignment=128) + sf_out = smem.allocate_array(cutlass.Uint8, self.smem_sf_out_bytes, byte_alignment=128) + data = smem.allocate_array(cutlass.Uint8, stages * tok * width, byte_alignment=128) + row_sf = None + dequantized_dst = cutlass.Int64(0) + if cutlass.const_expr(self.return_dequantized): + dequantized_dst = dst_dequantized.iterator.toint() + row_dst = cutlass.Int64(0) + if cutlass.const_expr(self.rowwise_output): + row_sf = smem.allocate_array( + cutlass.Uint8, self.HidAtomsPerTile * 512, byte_alignment=128 + ) + row_dst = dst_row_sf.iterator.toint() + row_base = cutlass.Int32(0) + if cutlass.const_expr(self.rowwise_output): + row_base = row_sf.toint() + if tid == 0: + for s in cutlass.range_constexpr(0, stages, 1): + cute.arch.mbarrier_init(full + s, 1) + cute.arch.mbarrier_init(empty + s, self.ConsumerWarps) + cute.arch.mbarrier_init_fence() + for g in cutlass.range(tid, self.num_experts + 1, self.ThreadsPerCta): + if cutlass.const_expr(self.uniform_rows > 0): + table[g] = Int32(g * self.uniform_rows) + else: + table[g] = Int32(offsets[g] // self.hidden) + cute.arch.sync_threads() + tiles = Int32(table[self.num_experts]) // tok + if warp == 0: + self.produce( + data.toint(), + sf_in.toint(), + full, + empty, + src_sf.iterator.toint(), + bid, + grid, + tiles, + lane, + load_atom, + load_tensor, + scale_atom, + scale_tensor, + ) + else: + self.consume_scaled( + data.toint(), + sf_in.toint(), + sf_out.toint(), + full, + empty, + table, + dst_col_sf.iterator.toint(), + bid, + grid, + tiles, + warp - 1, + lane, + store_atom, + store_tensor, + row_sf_base=row_base, + row_dst=row_dst, + dequantized_dst=dequantized_dst, + ) + + @cute.jit + def produce( + self, + data_base, + sf_base, + full, + empty, + src_sf_base, + bid, + grid, + tiles, + lane, + data_atom=None, + data_tensor=None, + scale_atom=None, + scale_tensor=None, + ): + """Load compact scale and payload tiles into the shared pipeline.""" + tok = cutlass.const_expr(self.TILE_TOK) + width = cutlass.const_expr(self.TILE_HID) + box_h = cutlass.const_expr(self.TmaBoxHidU32) + stages = cutlass.const_expr(self.NumStages) + sf_bytes = cutlass.const_expr(self.SfTileBytes) + data = cute.make_tensor( + cute.make_ptr(cutlass.Uint32, data_base, AddressSpace.smem, assumed_align=128), + cute.make_layout((tok, box_h, stages), stride=(box_h, 1, tok * box_h)), + ) + global_data = cute.group_modes( + cute.local_tile(data_tensor, (tok, box_h), (None, None)), 0, 2 + ) + shared_tile, global_tile = cpasync.tma_partition( + data_atom, 0, cute.make_layout(1), cute.group_modes(data, 0, 2), global_data + ) + cpasync.prefetch_descriptor(data_atom) + if cutlass.const_expr(self.hidden >= 512): + sf = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, sf_base, AddressSpace.smem, assumed_align=128), + cute.make_layout( + (tok, self.scale_stride, stages), stride=(self.scale_stride, 1, sf_bytes) + ), + ) + # The 16-byte scale box can overlap the next data tile's scales. + # Its start is selected with a coordinate tensor, rather than + # local_tile's box-width stepping. + sf_tile = cute.local_tile(scale_tensor, (tok, self.scale_stride), (None, None)) + sf_shared, _ = cpasync.tma_partition( + scale_atom, + 0, + cute.make_layout(1), + cute.group_modes(sf, 0, 2), + cute.group_modes(sf_tile, 0, 2), + ) + cpasync.prefetch_descriptor(scale_atom) + step = Int32(0) + work = Int32(bid) + total = tiles * Int32(self.hidden_groups) + while work < total: + stage = step % Int32(stages) + token_tile = work // Int32(self.hidden_groups) + hid_begin = (work % Int32(self.hidden_groups)) * Int32(width) + if step >= Int32(stages): + cute.arch.mbarrier_wait(empty + stage, ((step // Int32(stages)) - 1) % 2) + if lane == 0: + cute.arch.mbarrier_arrive_and_expect_tx(full + stage, Int32(tok * width + sf_bytes)) + cute.arch.sync_warp() + if cutlass.const_expr(self.hidden >= 512): + # TMA scale coordinates must be 16-byte aligned. Adjacent + # data tiles share the same aligned 16-column scale box. + source = cute.domain_offset( + (token_tile * tok, (hid_begin // 32) // self.scale_stride * self.scale_stride), + scale_tensor, + ) + tile = cute.local_tile(source, (tok, self.scale_stride), (0, 0)) + _, source_partition = cpasync.tma_partition( + scale_atom, + 0, + cute.make_layout(1), + cute.group_modes(sf, 0, 2), + cute.group_modes(tile, 0, 2), + ) + cute.copy( + scale_atom, + source_partition[(None,)], + sf_shared[(None, stage)], + tma_bar_ptr=full + stage, + ) + else: + if lane == 0: + tma_load_1d( + cute.make_ptr( + cutlass.Uint8, + sf_base + stage * sf_bytes, + AddressSpace.smem, + assumed_align=16, + ), + Int64(src_sf_base) + Int64(token_tile * sf_bytes), + full + stage, + Int32(sf_bytes), + ) + cute.copy( + data_atom, + global_tile[(None, token_tile, hid_begin // width)], + shared_tile[(None, stage)], + tma_bar_ptr=full + stage, + ) + step = step + Int32(1) + work = work + grid + + @cute.jit + def emit_row_scales(self, src_base, dst_base, cw, lane): + """Pack compact rowwise scales into shared 128-by-4 GEMM atoms.""" + row = cw * Int32(32) + lane + if row < Int32(128): + for a in cutlass.range_constexpr(0, self.HidAtomsPerTile, 1): + value = cute.make_tensor( + cute.make_ptr( + cutlass.Int32, + src_base + row * self.scale_stride + a * 4, + AddressSpace.smem, + assumed_align=4, + ), + cute.make_layout((1,)), + )[0] + offset = a * 512 + (row % 32) * 16 + (row // 32) * 4 + cute.make_tensor( + cute.make_ptr( + cutlass.Int32, dst_base + offset, AddressSpace.smem, assumed_align=4 + ), + cute.make_layout((1,)), + )[0] = value + + @cute.jit + def find_expert(self, tbl, key): + """Find the nonempty group containing a 128-row tile.""" + E = cutlass.const_expr(self.num_experts) + lo = Int32(0) + for step in self._search_steps: + probe = lo + Int32(step) + if cutlass.const_expr(self._search_needs_guard): + if probe < Int32(E) and Int32(tbl[probe]) <= key: + lo = probe + elif Int32(tbl[probe]) <= key: + lo = probe + return lo + + @cute.jit + def consume_scaled( + self, + smem_data_base, + smem_sf_in_base, + smem_sf_out_base, + mbar_full, + mbar_empty, + tbl_data, + dst_sf_base, + bidx, + grid_dim_x, + total_tiles, + cw, + lane_idx, + tma_atom_st=None, + tma_tensor_st=None, + row_sf_base=None, + row_dst=None, + dequantized_dst=None, + ): + """Decode/requantize and emit TE layouts with shared BF16 registers.""" + # CuTeDSL traces explicit conjunctions; retain these rather than chained comparisons. + # pylint: disable=chained-comparison + TOK = cutlass.const_expr(self.TILE_TOK) + W = cutlass.const_expr(self.TILE_HID) + S = cutlass.const_expr(self.NumStages) + SFB = cutlass.const_expr(self.SfTileBytes) + NB = cutlass.const_expr(self.TokensPerBlock) + CONS_THREADS = cutlass.const_expr(self.ConsumerWarps * 32) + HATOMS = cutlass.const_expr(self.HidAtomsPerTile) + C = cutlass.const_expr(self.ColsPerLane) + SEGW = cutlass.const_expr(32 * C) + tb = cw // Int32(self.HidSegs) + seg = cw % Int32(self.HidSegs) + LW = cutlass.const_expr(self.sp_lw) + ldsw = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=LW * 8 + ) + data_lane_off = tb * Int32(NB * W) + seg * Int32(SEGW) + lane_idx * Int32(LW) + hb0 = seg * Int32(C) + lane_idx * Int32(LW) // Int32(32) + sf_lane_off = tb * Int32(32 * self.scale_stride) + BOX_HS = cutlass.const_expr(self.TmaBoxHidU32) + sDo = cute.make_tensor( + cute.make_ptr(cutlass.Uint32, smem_data_base, AddressSpace.smem, assumed_align=128), + cute.make_layout((TOK, BOX_HS, S), stride=(BOX_HS, 1, TOK * BOX_HS)), + ) + gDo = cute.group_modes(cute.local_tile(tma_tensor_st, (TOK, BOX_HS), (None, None)), 0, 2) + tDsDo, tDgDo = cpasync.tma_partition( + tma_atom_st, 0, cute.make_layout(1), cute.group_modes(sDo, 0, 2), gDo + ) + cpasync.prefetch_descriptor(tma_atom_st) + t = Int32(0) + work_idx = Int32(bidx) + total_work = total_tiles * Int32(self.hidden_groups) + while work_idx < total_work: + stage = t % Int32(S) + token_tile = work_idx // Int32(self.hidden_groups) + hid_begin = work_idx % Int32(self.hidden_groups) * Int32(W) + data_row0 = token_tile * Int32(TOK) + owner = self.find_expert(tbl_data, data_row0) + sf_expert_token_atom = Int32(tbl_data[owner]) // Int32(TOK) + sf_token_atom = (data_row0 - Int32(tbl_data[owner])) // Int32(TOK) + sf_token_atoms = (Int32(tbl_data[owner + Int32(1)]) - Int32(tbl_data[owner])) // Int32( + TOK + ) + cute.arch.mbarrier_wait(mbar_full + stage, t // Int32(S) % Int32(2)) + stage_data = smem_data_base + stage * Int32(TOK * W) + data_lane_off + stage_sf = smem_sf_in_base + stage * Int32(SFB) + sf_lane_off + stage_sf = stage_sf + (hid_begin // 32) % self.scale_stride + sfout = smem_sf_out_base + if cutlass.const_expr(self.rowwise_output): + self.emit_row_scales(stage_sf - sf_lane_off, row_sf_base, cw, lane_idx) + self._sp_tile_body( + stage_data, + stage_sf, + sfout, + ldsw, + hb0, + seg, + tb, + lane_idx, + data_row0, + hid_begin, + dst_sf_base, + dequantized_dst, + ) + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.barrier(barrier_id=self.ConsumerBarrierId, number_of_threads=CONS_THREADS) + if cw == Int32(0): + cute.copy( + tma_atom_st, tDsDo[None, stage], tDgDo[None, token_tile, hid_begin // Int32(W)] + ) + if ( + cutlass.const_expr(self.swizzled) + and cw == Int32(0) + and lane_idx >= Int32(16) + and (lane_idx < Int32(16) + Int32(HATOMS)) + ): + atom = lane_idx - Int32(16) + cp_async_bulk_s2g( + Int64(dst_sf_base) + + ( + Int64(sf_expert_token_atom) * Int64(self._hidden_atoms) + + (Int64(hid_begin // Int32(self.SfAtomNonK)) + Int64(atom)) + * Int64(sf_token_atoms) + + Int64(sf_token_atom) + ) + * Int64(self.SfAtomBytes), + cute.make_ptr( + self.sf_dtype, + sfout + atom * Int32(self.SfOutStride), + AddressSpace.smem, + assumed_align=16, + ), + Int32(self.SfAtomBytes), + ) + if cutlass.const_expr(self.rowwise_output): + if ( + cw == Int32(0) + and lane_idx >= Int32(24) + and lane_idx < Int32(24 + self.HidAtomsPerTile) + ): + atom = lane_idx - Int32(24) + cp_async_bulk_s2g( + Int64(row_dst) + + ( + Int64(token_tile) * self._hidden_atoms + + Int64(hid_begin // 128) + + Int64(atom) + ) + * 512, + cute.make_ptr( + cutlass.Uint8, + row_sf_base + atom * 512, + AddressSpace.smem, + assumed_align=16, + ), + Int32(512), + ) + if cw == Int32(0): + cute.arch.cp_async_bulk_commit_group() + if cw == Int32(0): + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.barrier(barrier_id=self.ConsumerBarrierId, number_of_threads=CONS_THREADS) + if lane_idx == Int32(0): + cute.arch.mbarrier_arrive(mbar_empty + stage) + t = t + Int32(1) + work_idx = work_idx + grid_dim_x + + @cute.jit + def _sp_tile_body( + self, + stage_data, + stage_sf_base, + sfout, + ldsw, + hb0, + seg, + tb, + lane_idx, + data_row0, + hid_begin, + col_sf_base, + dequantized_dst, + ): + """The single-pass arithmetic for one lane's share of one tile.""" + W = cutlass.const_expr(self.TILE_HID) + NB = cutlass.const_expr(self.TokensPerBlock) + C = cutlass.const_expr(self.ColsPerLane) + LW = cutlass.const_expr(self.sp_lw) + NWc = cutlass.const_expr(LW // 4) + NPc = cutlass.const_expr(LW // 2) + NCH = cutlass.const_expr(C // LW) + SEGW = cutlass.const_expr(32 * C) + QT = self.quant_dtype + for ch in cutlass.range_constexpr(0, NCH, 1): + base = stage_data + Int32(ch * 32 * LW) + hbc = hb0 + Int32(ch * LW) + sfb = stage_sf_base + hbc + d = [[None] * NPc for _ in range(NB)] + acc = [Int32(0)] * NPc + for tt in cutlass.range_constexpr(0, NB, 1): + words = self._load_words(base + Int32(tt * W), ldsw, NWc, LW) + raw_sf = self._src_scale_raw(sfb + Int32(tt * self.scale_stride)) + s16 = raw_sf | raw_sf << Int32(8) + for w in cutlass.range_constexpr(0, NWc, 1): + qw = Int32(words[w]) + lo = cvt_scaled_up_bf16x2(qw, s16, 0, QT, self.use_scaled_cvt) + hi = cvt_scaled_up_bf16x2(qw, s16, 1, QT, self.use_scaled_cvt) + d[tt][2 * w] = lo + d[tt][2 * w + 1] = hi + acc[2 * w] = max_xorsign_abs_bf16x2(acc[2 * w], lo) + acc[2 * w + 1] = max_xorsign_abs_bf16x2(acc[2 * w + 1], hi) + if cutlass.const_expr(self.return_dequantized): + # Reuse decoded BF16x2 registers before FP8 conversion. Each + # lane writes contiguous columns; no shared transpose or + # additional pipeline storage is needed. + decoded = cute.make_rmem_tensor((NPc,), cutlass.Int32) + for k in cutlass.range_constexpr(0, NPc, 1): + decoded[k] = d[tt][k] + row = data_row0 + tb * NB + tt + col = hid_begin + seg * SEGW + ch * 32 * LW + lane_idx * LW + address = Int64(dequantized_dst) + (Int64(row) * self.hidden + col) * 2 + cute.copy( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Int32, + num_bits_per_copy=LW * 16, + ), + decoded, + cute.make_tensor( + cute.make_ptr( + cutlass.Int32, address, AddressSpace.gmem, assumed_align=LW * 2 + ), + cute.make_layout((NPc,)), + ), + ) + raws = [None] * LW + for k in cutlass.range_constexpr(0, NPc, 1): + a = acc[k] & Int32(2147450879) + raws[2 * k] = e8m0_raw_from_bf16(a & Int32(65535), self._data_limit_exponent) + raws[2 * k + 1] = e8m0_raw_from_bf16( + a >> Int32(16) & Int32(65535), self._data_limit_exponent + ) + invs = [ + bits_f32(cutlass.max(Int32(254) - raws[j] << Int32(23), Int32(4194304))) + for j in range(LW) + ] + col0 = seg * Int32(SEGW) + Int32(ch * 32 * LW) + lane_idx * Int32(LW) + for j in cutlass.range_constexpr(0, LW, 1): + col = col0 + Int32(j) + off = ( + sfout + + col // Int32(128) * Int32(self.SfOutStride) + + col % Int32(32) * Int32(16) + + col % Int32(128) // Int32(32) * Int32(4) + + tb + ) + if cutlass.const_expr(self.swizzled): + out_t = cute.make_tensor( + cute.make_ptr(cutlass.Uint8, off, AddressSpace.smem, assumed_align=1), + cute.make_layout((1,)), + ) + else: + offset = (data_row0 // 32 + tb) * self.hidden + hid_begin + col + out_t = cute.make_tensor( + cute.make_ptr( + cutlass.Uint8, + Int64(col_sf_base) + Int64(offset), + AddressSpace.gmem, + assumed_align=1, + ), + cute.make_layout((1,)), + ) + out_t[0] = Uint8(raws[j]) + for tt in cutlass.range_constexpr(0, NB, 2): + out0, out1 = self._requant_token_pair(d[tt], d[tt + 1], invs, NWc) + self._store_words(base + Int32(tt * W), ldsw, out0, NWc, LW) + self._store_words(base + Int32((tt + 1) * W), ldsw, out1, NWc, LW) + + def _requant_token_pair(self, d0, d1, invs, NWc): + """Requantize two tokens using each column's exact reciprocal scale. + + Inline into the caller's trace: the BF16 register lists cannot cross a + jitted function boundary. Pair columns to avoid a shared-memory transpose. + """ + QT = cutlass.Float8E4M3FN + out0 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + out1 = cute.make_rmem_tensor((NWc,), cutlass.Int32) + for w in range(0, NWc, 1): + k0 = 2 * w + k1 = 2 * w + 1 + p00 = cvt_dn_fp8x2_portable(d0[k0], invs[2 * k0], invs[2 * k0 + 1], QT) + p01 = cvt_dn_fp8x2_portable(d0[k1], invs[2 * k1], invs[2 * k1 + 1], QT) + p10 = cvt_dn_fp8x2_portable(d1[k0], invs[2 * k0], invs[2 * k0 + 1], QT) + p11 = cvt_dn_fp8x2_portable(d1[k1], invs[2 * k1], invs[2 * k1 + 1], QT) + out0[w] = p00 & Int32(65535) | p01 << Int32(16) + out1[w] = p10 & Int32(65535) | p11 << Int32(16) + return (out0, out1) + + @cute.jit + def _load_words(self, byte_addr, atom, NW, C): + """One LDS of C raw bytes (no FP8 decode -- the scaled cvt does that).""" + regs = cute.make_rmem_tensor((NW,), cutlass.Int32) + cute.copy( + atom, + cute.make_tensor( + cute.make_ptr(cutlass.Int32, byte_addr, AddressSpace.smem, assumed_align=C), + cute.make_layout((NW,)), + ), + regs, + ) + return regs + + @cute.jit + def _store_words(self, byte_addr, atom, regs, NW, C): + cute.copy( + atom, + regs, + cute.make_tensor( + cute.make_ptr(cutlass.Int32, byte_addr, AddressSpace.smem, assumed_align=C), + cute.make_layout((NW,)), + ), + ) + + @cute.jit + def _src_scale_raw(self, byte_addr): + """Raw source E8M0 byte, forced to 0..255. + + The mask is load-bearing: the pointer type says unsigned, but the + emitted load sign-extends, and the caller packs this into both halves + of a word with `raw | (raw << 8)`. Without the mask a scale byte >= + 0x80 poisons the E8M0 pair. Removing it took the mega suite from 28/28 + to 0/28. + """ + return Int32( + cute.make_tensor( + cute.make_ptr(cutlass.Uint8, byte_addr, AddressSpace.smem, assumed_align=1), + cute.make_layout((1,)), + )[0] + ) & Int32(255) + + +def get_mxfp8_requantization_function( + fn_name, + dtype, + rows, + hidden, + groups, + uniform_rows, + swizzled, + rowwise_output, + return_dequantized, + sm_count, +): + """Compile/register one common C API specialization, or request CUDA fallback.""" + if tvm_ffi.get_global_func(fn_name, allow_missing=True) is not None: + return True + try: + input_dtype = {"Float8E4M3": cutlass.Float8E4M3FN, "Float8E5M2": cutlass.Float8E5M2}[dtype] + kernel = GroupedRequantize( + hidden, + groups, + rows, + input_dtype=input_dtype, + uniform_rows=uniform_rows, + swizzled=swizzled, + rowwise_output=rowwise_output, + return_dequantized=return_dequantized, + sm_count=sm_count, + ) + + def tensor(element_type, shape): + return cute.runtime.make_fake_compact_tensor( + element_type, + shape, + stride_order=tuple(reversed(range(len(shape)))), + memspace=cute.AddressSpace.gmem, + assumed_align=16, + ) + + scales = rows * hidden // 32 + compiled = cute.compile( + kernel, + tensor(input_dtype, (rows, hidden)), + tensor(cutlass.Float8E8M0FNU, (rows, hidden // 32)), + None if uniform_rows else tensor(cutlass.Int64, (groups + 1,)), + tensor(cutlass.Float8E4M3FN, (rows, hidden)), + tensor(cutlass.Float8E8M0FNU, (scales,)) if rowwise_output else None, + tensor(cutlass.Float8E8M0FNU, (scales,)), + tensor(cutlass.BFloat16, (rows, hidden)) if return_dequantized else None, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + native = getattr(compiled, "__tvm_ffi_object__", lambda: None)() + tvm_ffi.register_global_func( + fn_name, native if native is not None else compiled, override=True + ) + return True + except Exception as error: # pylint: disable=broad-exception-caught + logger.warning("CuTeDSL grouped requantization unavailable; using CUDA: %s", error) + return False diff --git a/transformer_engine/common/cast/fused_group_requantize.cu b/transformer_engine/common/cast/fused_group_requantize.cu deleted file mode 100644 index aad83cb67c0..00000000000 --- a/transformer_engine/common/cast/fused_group_requantize.cu +++ /dev/null @@ -1,739 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file fused_group_requantize.cu - * \brief Fused grouped MXFP8 requantization: rowwise wire tensor -> GEMM-ready. - */ - -#include -#include -#include -#include - -#include -#include -#include - -#include "../common.h" -#include "../util/cuda_runtime.h" -#include "../util/ptx_arch_spec.cuh" -#include "../utils.cuh" -#include "mxfp8/swizzle.cuh" - -namespace transformer_engine { -namespace requantize { -namespace { - -constexpr int MXFP8_SCALE_DIM = 32; -constexpr int kTileRows = MXFP8_SCALE_DIM; -constexpr int kTileCols = 128; -constexpr int kElementsPerLoad = 16; -constexpr int kLoadsPerRow = kTileCols / kElementsPerLoad; -constexpr int kThreads = 128; -constexpr int kRowsPerGatherIteration = kThreads / kLoadsPerRow; -constexpr int kGatherIterations = kTileRows / kRowsPerGatherIteration; - -static_assert(kRowsPerGatherIteration == 16); -static_assert(kGatherIterations == 2); -static_assert(kThreads == kTileCols); - -__device__ __forceinline__ uint16_t e8m0_to_bf16_bits(const e8m0_t biased_exp) { - // E8M0 encodes the exponent bits directly. Codes 0 and 255 need explicit handling because - // 2^-127 is BF16-subnormal and 255 represents NaN. - if (biased_exp == 255) return 0x7fff; - if (biased_exp == 0) return 0x0040; - return static_cast(biased_exp) << 7; -} - -template -__device__ __forceinline__ ptx::bf16x2 dequantize_mxfp8_2x(const ptx::FPx2 &values, - const e8m0_t scale_code) { - ptx::bf16x2 result; -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -#if (defined CUDA_VERSION) && (CUDA_VERSION >= 13020) - // PTX ISA 9.2 can apply two packed E8M0 scaling factors while converting FP8x2 directly - // to BF16x2. Arch-specific Blackwell targets (sm_100a) include these family features. - constexpr bool kHasScaledFP8ToBF16 = ARCH_BLACKWELL_FAMILY; - if constexpr (kHasScaledFP8ToBF16) { - const uint16_t scale_x2 = - static_cast(scale_code) | (static_cast(scale_code) << 8); - if constexpr (std::is_same_v) { - asm volatile("cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 %0, %1, %2;" - : "=r"(reinterpret_cast(result)) - : "h"(reinterpret_cast(values)), "h"(scale_x2)); - } else { - static_assert(std::is_same_v); - asm volatile("cvt.rn.scaled::n2::ue8m0.bf16x2.e5m2x2 %0, %1, %2;" - : "=r"(reinterpret_cast(result)) - : "h"(reinterpret_cast(values)), "h"(scale_x2)); - } - return result; - } -#endif - - // CUDA 12.8-compatible fallback. Every E4M3/E5M2 value is exactly representable in - // FP16, so the FP16 bridge does not lose information before rounding to BF16. - const uint16_t scale_bits = e8m0_to_bf16_bits(scale_code); - const uint32_t scale_x2 = - static_cast(scale_bits) | (static_cast(scale_bits) << 16); - if constexpr (std::is_same_v) { - asm volatile( - "{\n\t" - ".reg.b32 values_f16x2, values_bf16x2; \n\t" - ".reg.b16 value0_f16, value1_f16, value0_bf16, value1_bf16; \n\t" - "cvt.rn.f16x2.e4m3x2 values_f16x2, %1; \n\t" - "mov.b32 {value0_f16, value1_f16}, values_f16x2; \n\t" - "cvt.rn.bf16.f16 value0_bf16, value0_f16; \n\t" - "cvt.rn.bf16.f16 value1_bf16, value1_f16; \n\t" - "mov.b32 values_bf16x2, {value0_bf16, value1_bf16}; \n\t" - "mul.rn.bf16x2 %0, values_bf16x2, %2; \n" - "}" - : "=r"(reinterpret_cast(result)) - : "h"(reinterpret_cast(values)), "r"(scale_x2)); - } else { - static_assert(std::is_same_v); - asm volatile( - "{\n\t" - ".reg.b32 values_f16x2, values_bf16x2; \n\t" - ".reg.b16 value0_f16, value1_f16, value0_bf16, value1_bf16; \n\t" - "cvt.rn.f16x2.e5m2x2 values_f16x2, %1; \n\t" - "mov.b32 {value0_f16, value1_f16}, values_f16x2; \n\t" - "cvt.rn.bf16.f16 value0_bf16, value0_f16; \n\t" - "cvt.rn.bf16.f16 value1_bf16, value1_f16; \n\t" - "mov.b32 values_bf16x2, {value0_bf16, value1_bf16}; \n\t" - "mul.rn.bf16x2 %0, values_bf16x2, %2; \n" - "}" - : "=r"(reinterpret_cast(result)) - : "h"(reinterpret_cast(values)), "r"(scale_x2)); - } -#else - NVTE_DEVICE_ERROR("Packed MXFP8 dequantization requires Blackwell hardware."); -#endif - return result; -} - -template -__device__ __forceinline__ void store_colwise_4x_to_shared(OType *const output, - const int stride_elements, - const uint32_t values) { - static_assert(sizeof(OType) == 1); - const uint32_t output_ptr = __cvta_generic_to_shared(output); - const uint32_t stride_bytes = stride_elements * sizeof(OType); - asm volatile( - "{\n\t" - ".reg.u32 ptr1, ptr2, ptr3; \n\t" - "mad.lo.u32 ptr1, 1, %1, %0; \n\t" - "mad.lo.u32 ptr2, 2, %1, %0; \n\t" - "mad.lo.u32 ptr3, 3, %1, %0; \n\t" - ".reg.b8 value0, value1, value2, value3; \n\t" - "mov.b32 {value0, value1, value2, value3}, %2; \n\t" - "st.shared.b8 [%0], value0; \n\t" - "st.shared.b8 [ptr1], value1; \n\t" - "st.shared.b8 [ptr2], value2; \n\t" - "st.shared.b8 [ptr3], value3; \n" - "}" - : - : "r"(output_ptr), "r"(stride_bytes), "r"(values) - : "memory"); -} - -template -__device__ __forceinline__ void store_colwise_2x_to_shared(OType *const output, - const int stride_elements, - const uint32_t values) { - static_assert(sizeof(OType) == 1); - const uint32_t output_ptr = __cvta_generic_to_shared(output); - const uint32_t stride_bytes = stride_elements * sizeof(OType); - asm volatile( - "{\n\t" - ".reg.u32 ptr1; \n\t" - "mad.lo.u32 ptr1, 1, %1, %0; \n\t" - ".reg.b8 value0, value1, unused0, unused1; \n\t" - "mov.b32 {value0, value1, unused0, unused1}, %2; \n\t" - "st.shared.b8 [%0], value0; \n\t" - "st.shared.b8 [ptr1], value1; \n" - "}" - : - : "r"(output_ptr), "r"(stride_bytes), "r"(values) - : "memory"); -} - -// --------------------------------------------------------------------------- -// Fused grouped requantization. -// -// Rowwise MXFP8 grouped input -> columnwise MXFP8 output with GEMM-swizzled -// scales for BOTH directions, plus an optional BF16 dequantized copy. Replaces -// the group_dequantize -> group_quantize(columnwise) -> grouped_swizzle(rowwise -// scales) chain with one kernel; the dequantized values only exist in shared -// memory unless dequantized_out is requested. -// -// Contract (asserted host-side where possible): the hidden dim and every -// group's row count are multiples of 128, so data tiles and swizzle tiles -// never straddle a group boundary and the scale layouts carry no padding. -// Group boundaries arrive as the grouped tensor's cached element-based -// tensor_offsets (offsets[g] = row offset x cols); they live on the device -// (host reads would break CUDA-graph capture), so per-group divisibility is -// the caller's contract. Rows at or past offsets[num_groups] (capacity-mode / -// paged-stash tail) are left untouched, matching the unfused chain. - -template -__global__ void __launch_bounds__(kThreads) - fused_group_requantize_kernel(const __grid_constant__ CUtensorMap output_tensor_map, - const IType *const input, const e8m0_t *const input_scale_inv, - e8m0_t *const rowwise_scale_inv_swizzled, - e8m0_t *const colwise_scale_inv, bf16 *const dequantized_out, - const int64_t *const tensor_offsets, const int num_groups, - const int num_cols, const int input_scale_stride) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; - using TCompute = std::conditional_t; - - constexpr int kPaddingPerVector = sizeof(float) / sizeof(TCompute); - constexpr int kDequantizedStride = kTileCols + kLoadsPerRow * kPaddingPerVector; - - const int tid = threadIdx.x; - const int row_base = blockIdx.y * kTileRows; - const int col_base = blockIdx.x * kTileCols; - - // Tail guard: with capacity-mode / paged-stash tensors the allocated rows - // exceed the live rows (tensor_offsets[num_groups] elements). The unfused - // chain never touches those rows; skip them entirely, or their out-of-range - // tile indices would alias INTO the last group's columnwise scales. The - // branch is CTA-uniform: live rows are 128-aligned like every group, so the - // boundary cannot split a 32-row tile. - const int64_t row_element_base = static_cast(row_base) * num_cols; - if (row_element_base >= tensor_offsets[num_groups]) { - return; - } - - __shared__ alignas(16) TCompute dequantized[kTileRows][kDequantizedStride]; - __shared__ alignas(TMA_SHMEM_ALIGNMENT) OType quantized[kTileRows][kTileCols]; - __shared__ int group_info[2]; // {group_start_row, group_num_rows} - - // The columnwise scale layout is per-group; find the group owning this tile. - // 128-aligned group sizes mean a 32-row tile never straddles a boundary. - // upper_bound minus one lands on the non-empty owner even when zero-sized - // groups share an offset. Offsets are in elements (row offset x num_cols), - // exactly the grouped tensor's cached tensor_offsets. - if (tid == 0) { - int lo = 0; - int hi = num_groups - 1; - while (lo < hi) { - const int mid = (lo + hi) / 2; - if (row_element_base < tensor_offsets[mid + 1]) { - hi = mid; - } else { - lo = mid + 1; - } - } - group_info[0] = static_cast(tensor_offsets[lo] / num_cols); - group_info[1] = static_cast((tensor_offsets[lo + 1] - tensor_offsets[lo]) / num_cols); - } - - const int rowwise_scale_tiles_x = num_cols / kTileCols; - - // Phase 1: each of the 128 threads loads one contiguous 16-byte vector in - // each of two iterations; together the CTA dequantizes all 32x128 values - // into shared memory. - const int lane = tid % THREADS_PER_WARP; -#pragma unroll - for (int gather_iteration = 0; gather_iteration < kGatherIterations; ++gather_iteration) { - const int local_chunk = tid % kLoadsPerRow; - const int local_row = tid / kLoadsPerRow + gather_iteration * kRowsPerGatherIteration; - const int local_col = local_chunk * kElementsPerLoad; - const int row = row_base + local_row; - const int col = col_base + local_col; - - // Adjacent 16-byte chunks share one rowwise MXFP8 scale. The even chunk - // loads the E8M0 code, re-emits it at its GEMM-swizzled address (dense - // indexing equals the per-group layout because group row counts are - // multiples of the 128-row swizzle tile), and broadcasts it to its - // partner. - int scale_code = 0; - if ((local_chunk % 2) == 0) { - const int scale_col = col / MXFP8_SCALE_DIM; - const size_t input_scale_idx = static_cast(row) * input_scale_stride + scale_col; - scale_code = static_cast(input_scale_inv[input_scale_idx]); - rowwise_scale_inv_swizzled[gemm_swizzled_scale_idx(row, scale_col, rowwise_scale_tiles_x)] = - static_cast(scale_code); - } - scale_code = __shfl_sync(0xffffffff, scale_code, lane & ~1); - - const int shared_col = local_col + (local_col / kElementsPerLoad) * kPaddingPerVector; - Vec input_vec; - input_vec.load_from(input + static_cast(row) * num_cols + col); - - constexpr int kDequantizedVecSize = kElementsPerLoad / 2; - [[maybe_unused]] Vec dequantized_vec[2]; // kReturnDequantized only - if constexpr (kUseFastMath) { -#pragma unroll - for (int i = 0; i < kElementsPerLoad; i += 2) { - const ptx::FPx2 values = {input_vec.data.elt[i], input_vec.data.elt[i + 1]}; - const ptx::bf16x2 result = dequantize_mxfp8_2x(values, static_cast(scale_code)); - *reinterpret_cast(&dequantized[local_row][shared_col + i]) = result; - if constexpr (kReturnDequantized) { - *reinterpret_cast( - &dequantized_vec[i / kDequantizedVecSize].data.elt[i % kDequantizedVecSize]) = result; - } - } - } else { - const float scale = ptx::exp2f(static_cast(scale_code)); -#pragma unroll - for (int i = 0; i < kElementsPerLoad; ++i) { - const float value = scale * static_cast(input_vec.data.elt[i]); - dequantized[local_row][shared_col + i] = value; - if constexpr (kReturnDequantized) { - dequantized_vec[i / kDequantizedVecSize].data.elt[i % kDequantizedVecSize] = - static_cast(value); - } - } - } - if constexpr (kReturnDequantized) { - bf16 *const dequantized_out_ptr = dequantized_out + static_cast(row) * num_cols + col; - dequantized_vec[0].store_to(dequantized_out_ptr); - dequantized_vec[1].store_to(dequantized_out_ptr + kDequantizedVecSize); - } - } - - __syncthreads(); - - // Phase 2: one thread owns one column and quantizes its 32 values as one - // MXFP8 block, emitting the scale at its per-group GEMM-swizzled address. - { - const int group_start = group_info[0]; - const int group_rows = group_info[1]; - - const int dequantized_col = tid + (tid / kElementsPerLoad) * kPaddingPerVector; - float thread_amax = 0.0f; - - ptx::bf16x4 bf16_values[kTileRows / 4]; - if constexpr (kUseFastMath) { - ptx::bf16x2 thread_amax_x2 = {static_cast(0.0f), static_cast(0.0f)}; -#pragma unroll - for (int row = 0; row < kTileRows; row += 4) { - const ptx::bf16x4 values = { - dequantized[row][dequantized_col], dequantized[row + 1][dequantized_col], - dequantized[row + 2][dequantized_col], dequantized[row + 3][dequantized_col]}; - bf16_values[row / 4] = values; - const ptx::bf16x2 values01 = {values.x1, values.x2}; - const ptx::bf16x2 values23 = {values.x3, values.x4}; - ptx::abs_max_2x(thread_amax_x2, thread_amax_x2, values01); - ptx::abs_max_2x(thread_amax_x2, thread_amax_x2, values23); - } - thread_amax = static_cast(ptx::get_amax(thread_amax_x2.x, thread_amax_x2.y)); - } else { -#pragma unroll - for (int row = 0; row < kTileRows; ++row) { - thread_amax = fmaxf(thread_amax, fabsf(dequantized[row][dequantized_col])); - } - } - - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - - // Per-group columnwise scale addressing. Mirrors group_quantize's - // WITH_GEMM_SWIZZLED_SCALES emission (group_quantize_mxfp8.cuh) and the - // grouped-GEMM consumer's padded cumsum: under the /128 contract the - // per-group base is exactly group_start/32 * num_cols. - const size_t colwise_scale_base = static_cast(group_start) / MXFP8_SCALE_DIM * num_cols; - const int local_scale_row = (row_base - group_start) / kTileRows; - const int scale_col = col_base + tid; - // One swizzle tile spans GEMM_SWIZZLED_SCALE_TILE_DIM_X = 4 scale rows = - // 128 data rows, so this matches the producer's DIVUP(rows, 128). - const int colwise_scale_tiles_x = group_rows / 128; - colwise_scale_inv[colwise_scale_base + - gemm_swizzled_scale_idx(scale_col, local_scale_row, colwise_scale_tiles_x)] = - biased_exponent; - - if constexpr (kUseFastMath) { - const bf16 quant_multiplier = ptx::exp2f_rcp(biased_exponent); - const ptx::bf16x2 quant_multiplier_x2 = {quant_multiplier, quant_multiplier}; -#pragma unroll - for (int row = 0; row < kTileRows; row += 4) { - uint32_t result_data = 0; - auto &result = *reinterpret_cast(&result_data); - ptx::mul_cvt_4x(result, bf16_values[row / 4], quant_multiplier_x2); - store_colwise_4x_to_shared(&quantized[row][tid], kTileCols, result_data); - } - } else { - const float quant_multiplier = ptx::exp2f_rcp(biased_exponent); - const ptx::floatx2 quant_multiplier_x2 = {quant_multiplier, quant_multiplier}; -#pragma unroll - for (int row = 0; row < kTileRows; row += 2) { - const ptx::floatx2 values = {dequantized[row][dequantized_col], - dequantized[row + 1][dequantized_col]}; - uint32_t result_data = 0; - auto &result = *reinterpret_cast(&result_data); - ptx::mul_cvt_2x(result, values, quant_multiplier_x2); - store_colwise_2x_to_shared(&quantized[row][tid], kTileCols, result_data); - } - } - } - - // Make the complete 32x128 FP8 tile visible to the TMA engine and store it. - // The /128 contract makes every tile full, so no bounds handling is needed. - ptx::fence_proxy_async_shared_cta(); - __syncthreads(); - if (tid == 0) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&output_tensor_map), col_base, row_base, - reinterpret_cast(quantized)); - ptx::cp_async_bulk_commit_group(); - ptx::cp_async_bulk_wait_group(); - } - __syncthreads(); -#else - NVTE_DEVICE_THREAD0_ERROR("Fused grouped requantization requires Blackwell (SM100+) hardware."); -#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} - -template -void launch_fused_group_requantize(const Tensor &input, Tensor *output, - const Tensor &tensor_offsets, Tensor *dequantized, - const int num_groups, const int num_rows, const int num_cols, - const int input_scale_stride, - const CUtensorMap &output_tensor_map, cudaStream_t stream) { - using OType = fp8e4m3; - const dim3 grid(num_cols / kTileCols, num_rows / kTileRows); - const dim3 block(kThreads); - - bf16 *dequantized_ptr = nullptr; - if constexpr (kReturnDequantized) { - dequantized_ptr = reinterpret_cast(dequantized->data.dptr); - } - - fused_group_requantize_kernel - <<>>( - output_tensor_map, reinterpret_cast(input.data.dptr), - reinterpret_cast(input.scale_inv.dptr), - reinterpret_cast(output->scale_inv.dptr), - reinterpret_cast(output->columnwise_scale_inv.dptr), dequantized_ptr, - reinterpret_cast(tensor_offsets.data.dptr), num_groups, num_cols, - input_scale_stride); -} - -// --------------------------------------------------------------------------- -// Wide-tile variant for the fast-math case without a dequantized output. -// -// Same per-element arithmetic as fused_group_requantize_kernel, so the -// outputs are bit-identical; only the work decomposition differs. The kernel above runs one 32x128 -// tile per 128-thread CTA, thread 0 binary-searches the group offsets before the tile can finish, -// and both scale outputs are written one byte at a time. Here a CTA owns a 128x128 tile: warp w -// takes rows 32w..32w+31 and lane l columns 4l..4l+3, so the data moves as coalesced 128-byte row -// segments held in registers, and each of the two swizzled scale outputs is one contiguous -// 512-byte block per tile. The grid is persistent and walks the tiles in row-block order, so a -// CTA stops at the first tile past tensor_offsets[num_groups]; a capacity-sized buffer with a -// short live prefix then costs little more than the live rows. - -constexpr int kWideTile = 128; -constexpr int kWideThreads = 128; -constexpr int kWideMaxGroups = 128; -constexpr int kWideCtasPerSM = 16; - -template -__global__ void __launch_bounds__(kWideThreads) - fused_group_requantize_wide_kernel(const IType *const input, - const e8m0_t *const input_scale_inv, - const int input_scale_stride, fp8e4m3 *const colwise_out, - e8m0_t *const rowwise_scale_inv_swizzled, - e8m0_t *const colwise_scale_inv, - const int64_t *const tensor_offsets, const int num_groups, - const int num_cols, const int num_tiles) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - static_assert(sizeof(IType) == 1); - __shared__ int64_t offsets[kWideMaxGroups + 1]; - __shared__ alignas(16) e8m0_t colwise_scale_block[kWideTile * 4]; - - const int tid = threadIdx.x; - const int warp = tid / THREADS_PER_WARP; - const int lane = tid % THREADS_PER_WARP; - const int tiles_x = num_cols / kWideTile; - const int words_per_row = num_cols / 4; - // This lane's 4 columns lie in the 32-column scaling block lane / 8 of the tile. - const int scale_shift = 8 * (lane / 8); - for (int i = tid; i <= num_groups; i += kWideThreads) { - offsets[i] = tensor_offsets[i]; - } - __syncthreads(); - const int64_t live_end = offsets[num_groups]; - - const uint32_t *const input_words = reinterpret_cast(input); - uint32_t *const output_words = reinterpret_cast(colwise_out); - for (int tile = blockIdx.x; tile < num_tiles; tile += gridDim.x) { - const int tile_x = tile % tiles_x; - const int tile_y = tile / tiles_x; - const int row_base = tile_y * kWideTile; - const int col_base = tile_x * kWideTile; - const int64_t row_element_base = static_cast(row_base) * num_cols; - // Rows at or past the live end are left untouched, as by the kernel above. Later tiles of - // this CTA lie further down, so the walk ends here. - if (row_element_base >= live_end) { - break; - } - - // 32 rows x 4 columns per thread, one 4-byte word per row. - const int row0 = row_base + 32 * warp; - const uint32_t *const src = - input_words + static_cast(row0) * words_per_row + col_base / 4 + lane; - uint32_t data[32]; -#pragma unroll - for (int r = 0; r < 32; ++r) { - data[r] = __ldg(src + static_cast(r) * words_per_row); - } - // Lane l holds the 4 rowwise scale codes (the tile's four 32-column blocks) of row row0 + l. - const uint32_t scale_word = __ldg(reinterpret_cast( - input_scale_inv + static_cast(row0 + lane) * input_scale_stride + - col_base / MXFP8_SCALE_DIM)); - - // The owning group, as found by the upper_bound lookup of the kernel above. 128-aligned group - // sizes mean a tile never straddles a group boundary. - int group = 0; - for (int i = 0; i + 1 < num_groups; ++i) { - group += (offsets[i + 1] <= row_element_base) ? 1 : 0; - } - const int group_start = static_cast(offsets[group] / num_cols); - const int group_rows = static_cast((offsets[group + 1] - offsets[group]) / num_cols); - - // GEMM-swizzled copy of the rowwise scales: the tile's 128 rows x 4 codes are one 512-byte - // block, in which row r of the tile starts at (r % 32) * 16 + (r / 32) * 4. - *reinterpret_cast(rowwise_scale_inv_swizzled + - (static_cast(tile_y) * tiles_x + tile_x) * 512 + - lane * 16 + warp * 4) = scale_word; - - // Dequantize to BF16 and take the amax of each column over the 32-row block. - ptx::bf16x2 dequantized01[32]; - ptx::bf16x2 dequantized23[32]; - ptx::bf16x2 amax01 = {static_cast(0.0f), static_cast(0.0f)}; - ptx::bf16x2 amax23 = {static_cast(0.0f), static_cast(0.0f)}; -#pragma unroll - for (int r = 0; r < 32; ++r) { - const e8m0_t scale_code = - static_cast((__shfl_sync(0xffffffff, scale_word, r) >> scale_shift) & 0xFF); - const uint16_t values01 = static_cast(data[r]); - const uint16_t values23 = static_cast(data[r] >> 16); - dequantized01[r] = - dequantize_mxfp8_2x(reinterpret_cast &>(values01), scale_code); - dequantized23[r] = - dequantize_mxfp8_2x(reinterpret_cast &>(values23), scale_code); - ptx::abs_max_2x(amax01, amax01, dequantized01[r]); - ptx::abs_max_2x(amax23, amax23, dequantized23[r]); - } - constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; - const e8m0_t exponent0 = ptx::float_to_e8m0(static_cast(amax01.x) * kMaxNormRcp); - const e8m0_t exponent1 = ptx::float_to_e8m0(static_cast(amax01.y) * kMaxNormRcp); - const e8m0_t exponent2 = ptx::float_to_e8m0(static_cast(amax23.x) * kMaxNormRcp); - const e8m0_t exponent3 = ptx::float_to_e8m0(static_cast(amax23.y) * kMaxNormRcp); - const ptx::bf16x2 multiplier01 = {ptx::exp2f_rcp(exponent0), - ptx::exp2f_rcp(exponent1)}; - const ptx::bf16x2 multiplier23 = {ptx::exp2f_rcp(exponent2), - ptx::exp2f_rcp(exponent3)}; - - uint32_t *const dst = - output_words + static_cast(row0) * words_per_row + col_base / 4 + lane; -#pragma unroll - for (int r = 0; r < 32; ++r) { - ptx::fp8e4m3x4 quantized; - ptx::mul_cvt_4x(quantized, dequantized01[r], multiplier01, dequantized23[r], multiplier23); - dst[static_cast(r) * words_per_row] = reinterpret_cast(quantized); - } - - // Columnwise scales: column c of the tile and 32-row block w go to (c % 32) * 16 + (c / 32) * 4 - // + w of the tile's 512-byte block in the group's swizzled layout. Stage the block in shared - // memory, then store it with one 4-byte word per thread. - const e8m0_t exponents[4] = {exponent0, exponent1, exponent2, exponent3}; -#pragma unroll - for (int j = 0; j < 4; ++j) { - const int c = 4 * lane + j; - colwise_scale_block[(c % 32) * 16 + (c / 32) * 4 + warp] = exponents[j]; - } - __syncthreads(); - const size_t colwise_scale_base = static_cast(group_start) / MXFP8_SCALE_DIM * num_cols; - const int tile_in_group = (row_base - group_start) / kWideTile; - const int colwise_scale_tiles_x = group_rows / kWideTile; - e8m0_t *const colwise_scale_dst = - colwise_scale_inv + colwise_scale_base + - (static_cast(tile_x) * colwise_scale_tiles_x + tile_in_group) * 512; - reinterpret_cast(colwise_scale_dst)[tid] = - reinterpret_cast(colwise_scale_block)[tid]; - __syncthreads(); // The next tile rewrites colwise_scale_block. - } -#else - NVTE_DEVICE_THREAD0_ERROR("Fused grouped requantization requires Blackwell (SM100+) hardware."); -#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} - -template -void launch_fused_group_requantize_wide(const Tensor &input, Tensor *output, - const Tensor &tensor_offsets, const int num_groups, - const int num_tiles, const int num_cols, - const int input_scale_stride, cudaStream_t stream) { - const int grid = static_cast( - std::min(num_tiles, static_cast(cuda::sm_count()) * kWideCtasPerSM)); - fused_group_requantize_wide_kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(input.scale_inv.dptr), input_scale_stride, - reinterpret_cast(output->columnwise_data.dptr), - reinterpret_cast(output->scale_inv.dptr), - reinterpret_cast(output->columnwise_scale_inv.dptr), - reinterpret_cast(tensor_offsets.data.dptr), num_groups, num_cols, num_tiles); -} - -void fused_group_requantize(const Tensor &input, Tensor *output, const Tensor &tensor_offsets, - Tensor *dequantized, const QuantizationConfig *quant_config, - cudaStream_t stream) { - checkCuDriverContext(stream); - - NVTE_CHECK(is_supported_by_CC_100(), - "Fused grouped requantization requires Blackwell (SM100+) hardware."); - NVTE_CHECK(input.scaling_mode == NVTE_MXFP8_1D_SCALING, "Input must use MXFP8 1D scaling."); - NVTE_CHECK(output->scaling_mode == NVTE_MXFP8_1D_SCALING, "Output must use MXFP8 1D scaling."); - NVTE_CHECK(input.has_data(), "Input must have rowwise MXFP8 data."); - NVTE_CHECK(input.data.dptr != nullptr, "Input rowwise data must be allocated."); - NVTE_CHECK(is_fp8_dtype(input.data.dtype), "Input rowwise data must have an FP8 type."); - NVTE_CHECK(input.scale_inv.dptr != nullptr, "Input rowwise scaling tensor must be allocated."); - NVTE_CHECK(input.scale_inv.dtype == DType::kFloat8E8M0, - "Input rowwise scaling tensor must have E8M0 type."); - NVTE_CHECK(!input.with_gemm_swizzled_scales, - "Input rowwise scales must be unswizzled (compact); dequantization reads them " - "row-indexed."); - NVTE_CHECK(output->has_columnwise_data(), "Output must have columnwise MXFP8 data."); - NVTE_CHECK(output->columnwise_data.dptr != nullptr, "Output columnwise data must be allocated."); - NVTE_CHECK(output->columnwise_data.dtype == DType::kFloat8E4M3, - "Output columnwise data must have E4M3 type."); - NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, - "Output columnwise scaling tensor must be allocated."); - NVTE_CHECK(output->columnwise_scale_inv.dtype == DType::kFloat8E8M0, - "Output columnwise scaling tensor must have E8M0 type."); - NVTE_CHECK(output->scale_inv.dptr != nullptr, - "Output rowwise scaling tensor (the swizzled copy of the input scales) must be " - "allocated."); - NVTE_CHECK(output->scale_inv.dtype == DType::kFloat8E8M0, - "Output rowwise scaling tensor must have E8M0 type."); - NVTE_CHECK(tensor_offsets.has_data(), "tensor_offsets must be allocated."); - NVTE_CHECK(tensor_offsets.data.dptr != nullptr, "tensor_offsets data must be allocated."); - NVTE_CHECK(tensor_offsets.data.dtype == DType::kInt64, "tensor_offsets must have Int64 type."); - NVTE_CHECK(tensor_offsets.data.numel() >= 2, "tensor_offsets must hold num_groups + 1 entries."); - - const int num_groups = static_cast(tensor_offsets.data.numel()) - 1; - - const auto [num_rows_size_t, num_cols_size_t] = input.flat_2d_dims(); - const auto [output_rows_size_t, output_cols_size_t] = output->flat_2d_dims(); - constexpr size_t kMaxInt = static_cast(std::numeric_limits::max()); - NVTE_CHECK(num_rows_size_t <= kMaxInt && num_cols_size_t <= kMaxInt, - "Fused grouped requantization dimensions must fit in int32."); - const int num_rows = static_cast(num_rows_size_t); - const int num_cols = static_cast(num_cols_size_t); - - NVTE_CHECK(output_rows_size_t == num_rows_size_t && output_cols_size_t == num_cols_size_t, - "Input and output shapes must match, but got (", num_rows_size_t, ", ", - num_cols_size_t, ") and (", output_rows_size_t, ", ", output_cols_size_t, ")."); - // Each group's row count must be a multiple of 128 too, so that every group's - // scales start on a swizzle-tile boundary. Those counts live on the device, - // so that half is the caller's contract rather than an assertion. - NVTE_CHECK(num_rows % 128 == 0 && num_cols % 128 == 0, - "Fused grouped requantization requires dims that are multiples of 128, but got (", - num_rows, ", ", num_cols, ")."); - NVTE_CHECK(num_rows / kTileRows <= 65535, - "The number of rows is too large for the 2D CUDA launch grid."); - - // The input scale tensor may be exactly compact or carry TE's padded - // allocation; both are row-indexed with this stride. Under the /128 contract - // the padded shape degenerates to the compact one, so a flat (1-D) scale - // tensor is also accepted. - int input_scale_stride = num_cols / static_cast(MXFP8_SCALE_DIM); - if (input.scale_inv.shape.size() == 2) { - NVTE_CHECK(input.scale_inv.shape[1] <= kMaxInt, "MXFP8 scale strides must fit in int32."); - input_scale_stride = static_cast(input.scale_inv.shape[1]); - NVTE_CHECK(input_scale_stride >= num_cols / static_cast(MXFP8_SCALE_DIM), - "Input rowwise scale stride is smaller than the number of scale columns."); - } - const size_t num_scales = static_cast(num_rows) * (num_cols / MXFP8_SCALE_DIM); - NVTE_CHECK(input.scale_inv.numel() >= num_scales, - "Input rowwise scale tensor is smaller than rows x cols / 32."); - NVTE_CHECK(output->scale_inv.numel() >= num_scales, - "Output rowwise scale tensor is smaller than rows x cols / 32."); - NVTE_CHECK(output->columnwise_scale_inv.numel() >= - static_cast(num_rows) / MXFP8_SCALE_DIM * num_cols, - "Output columnwise scale tensor is smaller than rows / 32 x cols."); - - const bool return_dequantized = - dequantized != nullptr && dequantized->has_data() && dequantized->data.dptr != nullptr; - if (return_dequantized) { - NVTE_CHECK(dequantized->data.dtype == DType::kBFloat16, - "The dequantized output must have BF16 type."); - NVTE_CHECK( - dequantized->data.numel() == static_cast(num_rows) * static_cast(num_cols), - "The dequantized output must have rows x cols elements."); - NVTE_CHECK(is_aligned_ptr(dequantized->data.dptr, 16), - "The dequantized output pointer must be 16B aligned."); - } - - NVTE_CHECK(is_aligned_ptr(input.data.dptr, 16), "Input data pointer must be 16B aligned."); - NVTE_CHECK(is_aligned_ptr(output->columnwise_data.dptr, TMA_GMEM_ALIGNMENT), - "Output data pointer must be 16B aligned."); - - // Both scale directions come out GEMM-swizzled; make the metadata say so for - // every caller, not just the PyTorch integration. - output->with_gemm_swizzled_scales = true; - - if (num_rows == 0) { - return; - } - - const bool use_fast_math = quant_config != nullptr && quant_config->use_fast_math; - - // The wide-tile kernel covers the fast-math case without a dequantized output. It loads the - // rowwise scales and stores both scale outputs as 4-byte words. - const int64_t num_wide_tiles = - static_cast(num_rows / kWideTile) * static_cast(num_cols / kWideTile); - const bool use_wide_kernel = - use_fast_math && !return_dequantized && num_groups <= kWideMaxGroups && - num_wide_tiles <= static_cast(kMaxInt) && input_scale_stride % 4 == 0 && - is_aligned_ptr(input.scale_inv.dptr, 4) && is_aligned_ptr(output->scale_inv.dptr, 4) && - is_aligned_ptr(output->columnwise_scale_inv.dptr, 4); - if (use_wide_kernel) { - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - input.data.dtype, IType, - launch_fused_group_requantize_wide(input, output, tensor_offsets, num_groups, - static_cast(num_wide_tiles), num_cols, - input_scale_stride, stream);); // NOLINT(*) - NVTE_CHECK_CUDA(cudaGetLastError()); - return; - } - - alignas(64) CUtensorMap output_tensor_map{}; - create_2D_tensor_map(output_tensor_map, output->columnwise_data, num_rows, num_cols, kTileRows, - kTileCols, num_cols, 0, typeToNumBits(output->columnwise_data.dtype)); - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_fast_math, USE_FAST_MATH, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - input.data.dtype, IType, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - return_dequantized, RETURN_DEQUANTIZED, - launch_fused_group_requantize( - input, output, tensor_offsets, dequantized, num_groups, num_rows, num_cols, - input_scale_stride, output_tensor_map, stream);););); // NOLINT(*) - NVTE_CHECK_CUDA(cudaGetLastError()); -} - -} // namespace -} // namespace requantize -} // namespace transformer_engine - -void nvte_group_requantize(const NVTETensor input, NVTETensor output, - const NVTETensor tensor_offsets, NVTETensor dequantized, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - using namespace transformer_engine; - NVTE_API_CALL(nvte_group_requantize); - - const Tensor *input_cu = convertNVTETensorCheck(input); - Tensor *output_cu = convertNVTETensorCheck(output); - const Tensor *tensor_offsets_cu = convertNVTETensorCheck(tensor_offsets); - Tensor *dequantized_cu = dequantized != nullptr ? convertNVTETensor(dequantized) : nullptr; - const auto *quant_config_cu = reinterpret_cast(quant_config); - requantize::fused_group_requantize(*input_cu, output_cu, *tensor_offsets_cu, dequantized_cu, - quant_config_cu, stream); -} diff --git a/transformer_engine/common/cast/mxfp8/requantize_mxfp8.cu b/transformer_engine/common/cast/mxfp8/requantize_mxfp8.cu new file mode 100644 index 00000000000..7f406f8c8aa --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/requantize_mxfp8.cu @@ -0,0 +1,929 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file requantize_mxfp8.cu + * \brief Fused grouped rowwise-MXFP8 to columnwise-MXFP8 conversion. + */ + +#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/grouped_tma.cuh" +#include "specialized/swizzle.cuh" +#include "swizzle.cuh" + +#ifdef NVTE_WITH_CUTEDSL +#include "requantize_mxfp8_cutedsl.cuh" +#endif + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace group_requantize_kernel { + +using namespace dispatch::common; + +constexpr size_t TMA_SWIZZLE_ALIGNMENT = 1024; + +struct DefaultRequantizeConfig { + static constexpr size_t TILE_DIM_Y = 128; + static constexpr size_t TILE_DIM_X = 128; + static constexpr size_t CHUNK_DIM_Y = 128; + static constexpr size_t CHUNK_DIM_X = 128; + static constexpr size_t THREADS_PER_CHUNK = 128; + static constexpr size_t PREFETCH_STAGES = 1; + static constexpr size_t STATIC_PERSISTENT_BLOCKS_PER_SM = 24; +}; + +template +struct RequantizeConfig; + +template <> +struct RequantizeConfig : DefaultRequantizeConfig {}; + +template <> +struct RequantizeConfig : DefaultRequantizeConfig {}; + +template <> +struct RequantizeConfig : DefaultRequantizeConfig {}; + +template <> +struct RequantizeConfig : DefaultRequantizeConfig { + static constexpr size_t CHUNK_DIM_X = 256; +}; + +template +struct RequantizeTraitsImpl { + static constexpr ShapeRepresentation SHAPE_REPRESENTATION = SHAPE_REP; + static constexpr size_t TILE_DIM_Y = Config::TILE_DIM_Y; + static constexpr size_t TILE_DIM_X = Config::TILE_DIM_X; + static constexpr size_t CHUNK_DIM_Y = Config::CHUNK_DIM_Y; + static constexpr size_t CHUNK_DIM_X = Config::CHUNK_DIM_X; + static constexpr size_t THREADS_PER_CHUNK = Config::THREADS_PER_CHUNK; + static constexpr size_t PREFETCH_STAGES = Config::PREFETCH_STAGES; + static constexpr size_t STATIC_PERSISTENT_BLOCKS_PER_SM = Config::STATIC_PERSISTENT_BLOCKS_PER_SM; + + static constexpr size_t BUFFS_NUM = PREFETCH_STAGES + 1; + static constexpr size_t THREADS_X = TILE_DIM_X / MXFP8_SCALE_DIM; + static constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; + static constexpr size_t BUFF_DIM_Y = THREADS_Y; + static constexpr size_t BUFF_DIM_X = TILE_DIM_X; + static constexpr size_t BUFF_DIM = BUFF_DIM_Y * BUFF_DIM_X; + static constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y; + static constexpr size_t STAGES_X = CHUNK_DIM_X / TILE_DIM_X; + + static_assert(TILE_DIM_Y == CHUNK_DIM_Y); + static_assert(TILE_DIM_X == THREADS_PER_CHUNK); + static_assert(BUFF_DIM_Y == MXFP8_SCALE_DIM); + static_assert(PREFETCH_STAGES > 0); + static_assert(CHUNK_DIM_Y % BUFF_DIM_Y == 0); + static_assert(CHUNK_DIM_X % TILE_DIM_X == 0); + static_assert(STATIC_PERSISTENT_BLOCKS_PER_SM > 0); +}; + +template +struct RequantizeTraits : RequantizeTraitsImpl> {}; + +struct LaunchConfig { + size_t work_blocks_x = 0; + size_t same_both_rows = 0; + dim3 grid; +}; + +template +LaunchConfig get_launch_config(const size_t first_logical_dim, const size_t last_logical_dim, + const size_t total_elements, const size_t num_tensors) { + constexpr ShapeRepresentation shape_rep = Traits::SHAPE_REPRESENTATION; + constexpr size_t chunk_dim_y = Traits::CHUNK_DIM_Y; + constexpr size_t chunk_dim_x = Traits::CHUNK_DIM_X; + + LaunchConfig config; + if constexpr (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { + NVTE_CHECK(first_logical_dim % num_tensors == 0, + "SAME_BOTH_DIMS requires an integral row count per tensor."); + config.same_both_rows = first_logical_dim / num_tensors; + NVTE_CHECK(config.same_both_rows % chunk_dim_y == 0, + "Each grouped MXFP8 tensor row count must be divisible by ", chunk_dim_y, "."); + config.work_blocks_x = DIVUP(last_logical_dim, chunk_dim_x); + const size_t work_blocks_y = DIVUP(config.same_both_rows, chunk_dim_y); + NVTE_CHECK(config.work_blocks_x > 0 && work_blocks_y > 0, + "SAME_BOTH_DIMS requires non-empty tensors."); + config.grid = dim3(config.work_blocks_x, work_blocks_y, num_tensors); + } else if constexpr (shape_rep == ShapeRepresentation::VARYING_FIRST_DIM) { + NVTE_CHECK(first_logical_dim % chunk_dim_y == 0, + "The grouped logical row capacity must be divisible by ", chunk_dim_y, "."); + config.work_blocks_x = DIVUP(last_logical_dim, chunk_dim_x); + const size_t work_blocks_y = DIVUP(first_logical_dim, chunk_dim_y); + NVTE_CHECK(config.work_blocks_x > 0 && work_blocks_y > 0, + "VARYING_FIRST_DIM requires a non-empty logical tensor."); + config.grid = dim3(config.work_blocks_x, work_blocks_y); + } else { + NVTE_CHECK(num_tensors <= MAX_SUPPORTED_TENSOR_DESCRIPTORS, + "Number of tensors exceeds the grouped TMA descriptor limit (", + MAX_SUPPORTED_TENSOR_DESCRIPTORS, ")."); + const size_t estimated_work_blocks = DIVUP(total_elements, chunk_dim_y * chunk_dim_x); + const size_t static_grid_size = static_cast(transformer_engine::cuda::sm_count()) * + Traits::STATIC_PERSISTENT_BLOCKS_PER_SM; + NVTE_CHECK(static_grid_size > 0, "Persistent launch grid must be non-zero."); + const size_t requested_workers = std::max(1, static_grid_size / num_tensors); + const size_t average_work = std::max(1, DIVUP(estimated_work_blocks, num_tensors)); + config.work_blocks_x = std::min(requested_workers, average_work); + config.grid = dim3(config.work_blocks_x, num_tensors); + } + return config; +} + +__device__ __forceinline__ uint16_t e8m0_to_bf16_bits(const e8m0_t biased_exp) { + if (biased_exp == 255) return 0x7fff; + if (biased_exp == 0) return 0x0040; + return static_cast(biased_exp) << 7; +} + +template +__device__ __forceinline__ ptx::bf16x2 dequantize_mxfp8_2x(const ptx::FPx2 &values, + const e8m0_t scale_code) { + ptx::bf16x2 result; +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined CUDA_VERSION) && (CUDA_VERSION >= 13020) + constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell_arch) { + const uint16_t scale_x2 = + static_cast(scale_code) | (static_cast(scale_code) << 8); + if constexpr (std::is_same_v) { + asm volatile("cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 %0, %1, %2;" + : "=r"(reinterpret_cast(result)) + : "h"(reinterpret_cast(values)), "h"(scale_x2)); + } else { + static_assert(std::is_same_v); + asm volatile("cvt.rn.scaled::n2::ue8m0.bf16x2.e5m2x2 %0, %1, %2;" + : "=r"(reinterpret_cast(result)) + : "h"(reinterpret_cast(values)), "h"(scale_x2)); + } + return result; + } +#endif + + const uint16_t scale_bits = e8m0_to_bf16_bits(scale_code); + const uint32_t scale_x2 = + static_cast(scale_bits) | (static_cast(scale_bits) << 16); + if constexpr (std::is_same_v) { + asm volatile( + "{\n\t" + ".reg.b32 values_f16x2, values_bf16x2; \n\t" + ".reg.b16 value0_f16, value1_f16, value0_bf16, value1_bf16; \n\t" + "cvt.rn.f16x2.e4m3x2 values_f16x2, %1; \n\t" + "mov.b32 {value0_f16, value1_f16}, values_f16x2; \n\t" + "cvt.rn.bf16.f16 value0_bf16, value0_f16; \n\t" + "cvt.rn.bf16.f16 value1_bf16, value1_f16; \n\t" + "mov.b32 values_bf16x2, {value0_bf16, value1_bf16}; \n\t" + "mul.rn.bf16x2 %0, values_bf16x2, %2; \n" + "}" + : "=r"(reinterpret_cast(result)) + : "h"(reinterpret_cast(values)), "r"(scale_x2)); + } else { + static_assert(std::is_same_v); + asm volatile( + "{\n\t" + ".reg.b32 values_f16x2, values_bf16x2; \n\t" + ".reg.b16 value0_f16, value1_f16, value0_bf16, value1_bf16; \n\t" + "cvt.rn.f16x2.e5m2x2 values_f16x2, %1; \n\t" + "mov.b32 {value0_f16, value1_f16}, values_f16x2; \n\t" + "cvt.rn.bf16.f16 value0_bf16, value0_f16; \n\t" + "cvt.rn.bf16.f16 value1_bf16, value1_f16; \n\t" + "mov.b32 values_bf16x2, {value0_bf16, value1_bf16}; \n\t" + "mul.rn.bf16x2 %0, values_bf16x2, %2; \n" + "}" + : "=r"(reinterpret_cast(result)) + : "h"(reinterpret_cast(values)), "r"(scale_x2)); + } +#else + NVTE_DEVICE_ERROR("Packed MXFP8 dequantization requires Blackwell hardware."); +#endif + return result; +} + +__device__ __forceinline__ void store_scale_cache_unit(e8m0_t *const output, + const e8m0_t *const cache) { + constexpr size_t bytes_per_unit = sizeof(uint4); + static_assert(alignof(uint4) == bytes_per_unit); + if ((reinterpret_cast(output) & (alignof(uint4) - 1)) == 0) { + *reinterpret_cast(output) = *reinterpret_cast(cache); + } else { +#pragma unroll + for (int byte = 0; byte < static_cast(bytes_per_unit); ++byte) { + output[byte] = cache[byte]; + } + } +} + +template +__device__ __forceinline__ void requantize_bf16_4x(ptx::FPx4 &output, + const ptx::bf16x4 &values, + const ptx::bf16x4 &multipliers) { + static_assert(std::is_same_v); +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined CUDA_VERSION) && (CUDA_VERSION >= 13010) + asm volatile( + "{\n\t" + ".reg.b32 x01, x23, scale01, scale23, y01, y23; \n\t" + "mov.b64 {x01, x23}, %1; \n\t" + "mov.b64 {scale01, scale23}, %2; \n\t" + "mul.rn.bf16x2 y01, x01, scale01; \n\t" + "mul.rn.bf16x2 y23, x23, scale23; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e4m3x2.bf16x2 z01, y01; \n\t" + "cvt.rn.satfinite.e4m3x2.bf16x2 z23, y23; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}" + : "=r"(reinterpret_cast(output)) + : "l"(reinterpret_cast(values)), + "l"(reinterpret_cast(multipliers))); +#else + asm volatile( + "{\n\t" + ".reg.b16 x0, x1, x2, x3, scale0, scale1, scale2, scale3; \n\t" + "mov.b64 {x0, x1, x2, x3}, %1; \n\t" + "mov.b64 {scale0, scale1, scale2, scale3}, %2; \n\t" + ".reg.f32 y0, y1, y2, y3; \n\t" + "fma.rn.f32.bf16 y0, x0, scale0, 0f00000000; \n\t" + "fma.rn.f32.bf16 y1, x1, scale1, 0f00000000; \n\t" + "fma.rn.f32.bf16 y2, x2, scale2, 0f00000000; \n\t" + "fma.rn.f32.bf16 y3, x3, scale3, 0f00000000; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e4m3x2.f32 z01, y1, y0; \n\t" + "cvt.rn.satfinite.e4m3x2.f32 z23, y3, y2; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}" + : "=r"(reinterpret_cast(output)) + : "l"(reinterpret_cast(values)), + "l"(reinterpret_cast(multipliers))); +#endif +#else + NVTE_DEVICE_ERROR("Packed BF16 requantization requires Blackwell hardware."); +#endif +} + +template +__device__ __forceinline__ void process_bf16_stage( + const e8m0_t *const input_scales, e8m0_t *const output_scales, + e8m0_t *const output_rowwise_scales, const size_t input_scale_base, + const size_t output_scale_base, const size_t input_scale_stride, + const size_t output_scale_stride, const size_t output_scale_tiles_x, const size_t rows, + const size_t cols, const size_t block_offset_y, const size_t block_offset_x, + const size_t stage_offset_y, const size_t stage_offset_x, const int stages_y, + const int input_buffer, const IType *const input_shared, void *const reduction_shared, + OType *const output_shared, bf16 *const dequantized_out) { + constexpr size_t warp_tile_dim_x = MXFP8_SCALE_DIM; + constexpr size_t warps_per_chunk = Traits::THREADS_PER_CHUNK / THREADS_PER_WARP; + constexpr size_t elements_per_unit = sizeof(uint4) / sizeof(IType); + constexpr size_t units_per_row = Traits::BUFF_DIM_X / elements_per_unit; + constexpr size_t units_per_warp_tile = warp_tile_dim_x / elements_per_unit; + constexpr size_t units_per_buffer = Traits::BUFF_DIM / elements_per_unit; + constexpr size_t scratch_bytes_per_warp = MXFP8_SCALE_DIM * sizeof(float); + constexpr size_t scale_cache_bytes = Traits::STAGES_Y * Traits::BUFF_DIM_X; + constexpr bool has_native_max_abs_redux = NVTE_CUDA_ARCH_MATCHES(ptx::FamilySpecific<100>); + + static_assert(sizeof(IType) == 1 && sizeof(OType) == 1); + static_assert(Traits::THREADS_PER_CHUNK % THREADS_PER_WARP == 0); + static_assert(warps_per_chunk * warp_tile_dim_x == Traits::BUFF_DIM_X); + static_assert(warp_tile_dim_x % elements_per_unit == 0); + static_assert(Traits::BUFF_DIM_X * sizeof(IType) == 128); + static_assert(Traits::STAGES_Y == scale_tensor_alignment_Y_colwise); + static_assert(Traits::THREADS_PER_CHUNK * sizeof(uint32_t) == scale_cache_bytes); + static_assert(warps_per_chunk * scratch_bytes_per_warp + scale_cache_bytes <= + TMA_SWIZZLE_ALIGNMENT); + + using TmaSwizzle = transformer_engine::swz::Swizzle<3, 0, 3>; + using InputPairs = ptx::FPx2[MXFP8_SCALE_DIM / 2]; + using ComputePairs = ptx::bf16x2[MXFP8_SCALE_DIM / 2]; + using ComputeQuads = ptx::bf16x4[MXFP8_SCALE_DIM / 4]; + using OutputQuads = ptx::FPx4[MXFP8_SCALE_DIM / 4]; + + const int lane = threadIdx.x % THREADS_PER_WARP; + const int warp = threadIdx.x / THREADS_PER_WARP; + const size_t local_col = warp * warp_tile_dim_x; + const size_t tensor_row = block_offset_y + stage_offset_y + lane; + const size_t tensor_col = block_offset_x + stage_offset_x + local_col; + auto *const scale_cache = + reinterpret_cast(reinterpret_cast(reduction_shared) + + warps_per_chunk * scratch_bytes_per_warp); + + if (stage_offset_y == 0 && stages_y < static_cast(Traits::STAGES_Y)) { + auto *const scale_cache_u32 = reinterpret_cast(scale_cache); + scale_cache_u32[threadIdx.x] = 0; + __syncthreads(); + } + + e8m0_t input_scale = 0; + if (tensor_row < rows && tensor_col < cols) { + const size_t scale_col = tensor_col / MXFP8_SCALE_DIM; + const size_t scale_idx = input_scale_base + tensor_row * input_scale_stride + scale_col; + input_scale = input_scales[scale_idx]; + if constexpr (OUTPUT_SCALES_SWIZZLED) { + if (output_rowwise_scales != nullptr) { + const size_t output_idx = + input_scale_base + + swizzle::gemm_swizzled_scale_idx(tensor_row, scale_col, + input_scale_stride / scale_tensor_alignment_X_rowwise); + output_rowwise_scales[output_idx] = input_scale; + } + } + } + + alignas(16) uint4 packed_input[units_per_warp_tile]; + const auto *const input_units = reinterpret_cast(input_shared); + const size_t input_unit_base = + input_buffer * units_per_buffer + lane * units_per_row + warp * units_per_warp_tile; +#pragma unroll + for (int unit = 0; unit < static_cast(units_per_warp_tile); ++unit) { + packed_input[unit] = input_units[TmaSwizzle::swz(input_unit_base + unit)]; + } + + const auto &input_pairs = *reinterpret_cast(packed_input); + alignas(16) ComputePairs compute_values; + unsigned char *const warp_scratch = + reinterpret_cast(reduction_shared) + warp * scratch_bytes_per_warp; + float *const warp_amax = reinterpret_cast(warp_scratch); +#pragma unroll + for (int pair_idx = 0; pair_idx < static_cast(MXFP8_SCALE_DIM / 2); ++pair_idx) { + compute_values[pair_idx] = dequantize_mxfp8_2x(input_pairs[pair_idx], input_scale); + ptx::floatx2 values = ptx::up_cast(compute_values[pair_idx]); + if constexpr (!has_native_max_abs_redux) { + // The integer fallback orders NaN above every finite magnitude. + if (isnan(values.x)) values.x = 0.0f; + if (isnan(values.y)) values.y = 0.0f; + } + ptx::floatx2 amax; + ptx::reduce_sync_max_abs_f32(amax.x, values.x); + ptx::reduce_sync_max_abs_f32(amax.y, values.y); + if (lane == 0) { + *reinterpret_cast(&warp_amax[2 * pair_idx]) = make_float2(amax.x, amax.y); + } + } + + if constexpr (RETURN_DEQUANTIZED) { + if (tensor_row < rows && tensor_col < cols) { + const auto *const decoded_units = reinterpret_cast(compute_values); + auto *const dst = reinterpret_cast(dequantized_out + tensor_row * cols + tensor_col); +#pragma unroll + for (int unit = 0; unit < static_cast(MXFP8_SCALE_DIM * sizeof(bf16) / sizeof(uint4)); + ++unit) { + dst[unit] = decoded_units[unit]; + } + } + } + + __syncwarp(); + float thread_amax = warp_amax[lane]; + if constexpr (has_native_max_abs_redux) { + // Native redux returns NaN only when the entire column is NaN. Match the + // reference max-with-zero behavior without sanitizing every input value. + thread_amax = fmaxf(thread_amax, 0.0f); + } + __syncwarp(); + const e8m0_t output_scale = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + const size_t output_col = tensor_col + lane; + const size_t local_output_col = local_col + lane; + const size_t local_output_scale_row = stage_offset_y / MXFP8_SCALE_DIM; + size_t scale_cache_idx; + if constexpr (OUTPUT_SCALES_SWIZZLED) { + scale_cache_idx = swizzle::gemm_swizzled_scale_idx(local_output_col, local_output_scale_row, 1); + } else { + scale_cache_idx = local_output_scale_row * Traits::BUFF_DIM_X + local_output_col; + } + scale_cache[scale_cache_idx] = output_col < cols ? output_scale : static_cast(0); + + auto *const warp_multipliers = reinterpret_cast(warp_scratch); + warp_multipliers[lane] = ptx::exp2f_rcp(output_scale); + __syncwarp(); + + alignas(16) uint4 packed_output[units_per_warp_tile]; + const auto &compute_quads = *reinterpret_cast(compute_values); + auto &output_quads = *reinterpret_cast(packed_output); +#pragma unroll + for (int quad_idx = 0; quad_idx < static_cast(MXFP8_SCALE_DIM / 4); ++quad_idx) { + const auto multipliers = + *reinterpret_cast(&warp_multipliers[4 * quad_idx]); + requantize_bf16_4x(output_quads[quad_idx], compute_quads[quad_idx], multipliers); + } + + // The shared-to-global TMA wait-group state belongs to the issuing thread. + // Rendezvous with it before any warp overwrites the reused output buffer. + __syncthreads(); + + if (local_output_scale_row + 1 == static_cast(stages_y) && + threadIdx.x < scale_cache_bytes / sizeof(uint4)) { + const size_t cache_offset = threadIdx.x * sizeof(uint4); + if constexpr (OUTPUT_SCALES_SWIZZLED) { + const size_t output_scale_row = block_offset_y / MXFP8_SCALE_DIM; + const size_t output_scale_idx = + output_scale_base + swizzle::gemm_swizzled_scale_idx(block_offset_x + stage_offset_x, + output_scale_row, + output_scale_tiles_x); + store_scale_cache_unit(&output_scales[output_scale_idx + cache_offset], + &scale_cache[cache_offset]); + } else { + constexpr size_t units_per_scale_row = Traits::BUFF_DIM_X / sizeof(uint4); + const size_t scale_row = threadIdx.x / units_per_scale_row; + const size_t scale_col_unit = threadIdx.x % units_per_scale_row; + const size_t output_scale_row = block_offset_y / MXFP8_SCALE_DIM + scale_row; + const size_t output_scale_idx = output_scale_base + output_scale_row * output_scale_stride + + block_offset_x + stage_offset_x + + scale_col_unit * sizeof(uint4); + store_scale_cache_unit(&output_scales[output_scale_idx], &scale_cache[cache_offset]); + } + } + + auto *const output_units = reinterpret_cast(output_shared); + const size_t output_unit_base = + input_buffer * units_per_buffer + lane * units_per_row + warp * units_per_warp_tile; +#pragma unroll + for (int unit = 0; unit < static_cast(units_per_warp_tile); ++unit) { + output_units[TmaSwizzle::swz(output_unit_base + unit)] = packed_output[unit]; + } +} + +template +__device__ __forceinline__ void process_chunk( + const CUtensorMap &tensor_map_input, const CUtensorMap &tensor_map_output, + const e8m0_t *const input_scales, e8m0_t *const output_scales, + e8m0_t *const output_rowwise_scales, const size_t input_scale_base, + const size_t output_scale_base, const size_t rows, const size_t cols, + const size_t block_offset_y, const size_t block_offset_x, const size_t tma_offset_y, + IType *const input_shared, void *const dequantized_shared, OType *const output_shared, + uint64_t *const input_barriers, int *const input_barrier_parity, const bool leading_thread, + bf16 *const dequantized_out) { + constexpr size_t tile_dim_x = Traits::TILE_DIM_X; + constexpr size_t chunk_dim_y = Traits::CHUNK_DIM_Y; + constexpr size_t chunk_dim_x = Traits::CHUNK_DIM_X; + constexpr size_t buffs_num = Traits::BUFFS_NUM; + constexpr size_t buff_dim_y = Traits::BUFF_DIM_Y; + constexpr size_t buff_dim = Traits::BUFF_DIM; + constexpr size_t prefetch_stages = Traits::PREFETCH_STAGES; + constexpr size_t input_buffer_bytes = buff_dim * sizeof(IType); + + const size_t input_scale_stride = + DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(MXFP8_SCALE_DIM)), + static_cast(scale_tensor_alignment_X_rowwise)); + const size_t output_scale_stride = + DIVUP_TO_MULTIPLE(cols, static_cast(scale_tensor_alignment_X_colwise)); + const size_t output_scale_tiles_x = + DIVUP(rows, static_cast(scale_tensor_alignment_Y_rowwise)); + + const size_t chunk_rows = min(chunk_dim_y, rows - block_offset_y); + const size_t chunk_cols = min(chunk_dim_x, cols - block_offset_x); + const int stages_y = static_cast(DIVUP(chunk_rows, buff_dim_y)); + const int stages_x = static_cast(DIVUP(chunk_cols, tile_dim_x)); + const int stages = stages_y * stages_x; + +#pragma unroll + for (int stage = 0; stage < static_cast(prefetch_stages); ++stage) { + const size_t stage_y = stage % stages_y; + const size_t stage_x = stage / stages_y; + const size_t global_y = tma_offset_y + stage_y * buff_dim_y; + const size_t global_x = block_offset_x + stage_x * tile_dim_x; + const size_t buffer_offset = stage * buff_dim; + prefetch_input_stage(input_shared, nullptr, tensor_map_input, tensor_map_input, + global_x, global_y, buffer_offset, input_buffer_bytes, + &input_barriers[stage], leading_thread); + } + + int input_buffer = 0; +#pragma unroll + for (int stage = 0; stage < stages; ++stage) { + const size_t stage_y = stage % stages_y; + const size_t stage_x = stage / stages_y; + const size_t stage_offset_y = stage_y * buff_dim_y; + const size_t stage_offset_x = stage_x * tile_dim_x; + + if (stage < stages - static_cast(prefetch_stages)) { + const int next_stage = stage + prefetch_stages; + const int next_buffer = (input_buffer + prefetch_stages) % buffs_num; + const size_t next_stage_y = next_stage % stages_y; + const size_t next_stage_x = next_stage / stages_y; + const size_t global_y = tma_offset_y + next_stage_y * buff_dim_y; + const size_t global_x = block_offset_x + next_stage_x * tile_dim_x; + const size_t buffer_offset = next_buffer * buff_dim; + prefetch_input_stage(input_shared, nullptr, tensor_map_input, tensor_map_input, + global_x, global_y, buffer_offset, input_buffer_bytes, + &input_barriers[next_buffer], leading_thread); + } + + ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&input_barriers[input_buffer], + input_barrier_parity[input_buffer]); + input_barrier_parity[input_buffer] ^= 1; + + // Do not overwrite an output buffer that is still consumed by TMA. + if (leading_thread) { + ptx::cp_async_bulk_wait_group_read(); + } + + process_bf16_stage( + input_scales, output_scales, output_rowwise_scales, input_scale_base, output_scale_base, + input_scale_stride, output_scale_stride, output_scale_tiles_x, rows, cols, block_offset_y, + block_offset_x, stage_offset_y, stage_offset_x, stages_y, input_buffer, input_shared, + dequantized_shared, output_shared, dequantized_out); + + ptx::fence_proxy_async_shared_cta(); + __syncthreads(); + + const size_t global_y = tma_offset_y + stage_offset_y; + const size_t global_x = block_offset_x + stage_offset_x; + const size_t buffer_offset = input_buffer * buff_dim; + store_output_stage(nullptr, output_shared, tensor_map_output, + tensor_map_output, global_x, global_y, buffer_offset, + leading_thread); + + input_buffer = (input_buffer + 1) % buffs_num; + } + + // process_chunk restarts from output buffer 0. With an odd stage count the + // newest outstanding TMA store also uses buffer 0, so it must finish first. + if ((stages & 1) != 0 && leading_thread) { + ptx::cp_async_bulk_wait_group_read<0>(); + } +} + +template +__global__ void __launch_bounds__(Traits::THREADS_PER_CHUNK) group_requantize_mxfp8_kernel( + const __grid_constant__ CUtensorMap tensor_map_input_static, + const __grid_constant__ CUtensorMap tensor_map_output_static, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, const size_t same_both_rows, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, const e8m0_t *const __restrict__ input_scales, + e8m0_t *const __restrict__ output_rowwise_scales, e8m0_t *const __restrict__ output_scales, + bf16 *const __restrict__ dequantized_out) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr ShapeRepresentation shape_rep = Traits::SHAPE_REPRESENTATION; + constexpr bool direct_same = shape_rep == ShapeRepresentation::SAME_BOTH_DIMS; + constexpr bool direct_varying_first = shape_rep == ShapeRepresentation::VARYING_FIRST_DIM; + constexpr bool direct_mapper = direct_same || direct_varying_first; + constexpr bool single_tma_tensor = direct_mapper; + constexpr size_t buffs_num = Traits::BUFFS_NUM; + constexpr size_t buff_dim = Traits::BUFF_DIM; + constexpr size_t input_bytes = + DIVUP_TO_MULTIPLE(buffs_num * buff_dim * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr size_t dequantized_bytes = + DIVUP_TO_MULTIPLE(Traits::THREADS_PER_CHUNK * sizeof(float), TMA_SWIZZLE_ALIGNMENT); + constexpr size_t output_bytes = + DIVUP_TO_MULTIPLE(buffs_num * buff_dim * sizeof(OType), TMA_SHMEM_ALIGNMENT); + + extern __shared__ unsigned char dynamic_shared[]; + constexpr size_t shared_alignment = TMA_SWIZZLE_ALIGNMENT; + unsigned char *const shared_base = reinterpret_cast( + align_up(reinterpret_cast(dynamic_shared), shared_alignment)); + IType *const input_shared = reinterpret_cast(shared_base); + void *const dequantized_shared = shared_base + input_bytes; + OType *const output_shared = + reinterpret_cast(shared_base + input_bytes + dequantized_bytes); + + const bool leading_thread = threadIdx.x == 0; + + size_t tensor_id = 0; + size_t rows = 0; + size_t cols = 0; + size_t tensor_base = 0; + size_t block_offset_y = 0; + size_t block_offset_x = 0; + size_t tma_offset_y = 0; + + __shared__ size_t varying_first_metadata[4]; + if constexpr (direct_same) { + tensor_id = blockIdx.z; + rows = same_both_rows; + cols = last_logical_dim; + tensor_base = tensor_id * rows * cols; + block_offset_y = blockIdx.y * Traits::CHUNK_DIM_Y; + block_offset_x = blockIdx.x * Traits::CHUNK_DIM_X; + tma_offset_y = tensor_id * rows + block_offset_y; + } else if constexpr (direct_varying_first) { + const size_t global_block_y = blockIdx.y * Traits::CHUNK_DIM_Y; + const size_t global_element_offset = global_block_y * last_logical_dim; + if (leading_thread) { + const size_t active_elements = static_cast(offsets_ptr[num_tensors]); + varying_first_metadata[0] = active_elements; + if (global_element_offset < active_elements) { + const size_t mapped_tensor = + find_tensor_from_offsets(offsets_ptr, num_tensors, global_element_offset); + varying_first_metadata[1] = mapped_tensor; + varying_first_metadata[2] = get_tensor_rows_num(mapped_tensor, first_logical_dim, + first_dims_ptr, num_tensors); + varying_first_metadata[3] = static_cast(offsets_ptr[mapped_tensor]); + } + } + __syncthreads(); + if (global_element_offset >= varying_first_metadata[0]) return; + + tensor_id = varying_first_metadata[1]; + rows = varying_first_metadata[2]; + cols = last_logical_dim; + tensor_base = varying_first_metadata[3]; + const size_t tensor_start_y = tensor_base / cols; + block_offset_y = global_block_y - tensor_start_y; + block_offset_x = blockIdx.x * Traits::CHUNK_DIM_X; + tma_offset_y = global_block_y; + } else { + tensor_id = blockIdx.y; + if (tensor_id >= num_tensors) { + return; + } + rows = + get_tensor_rows_num(tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + cols = get_tensor_cols_num(tensor_id, last_logical_dim, last_dims_ptr); + tensor_base = static_cast(offsets_ptr[tensor_id]); + if (rows == 0 || cols == 0) { + return; + } + } + + const CUtensorMap &tensor_map_input = + single_tma_tensor ? tensor_map_input_static : g_tensor_maps.input[tensor_id]; + const CUtensorMap &tensor_map_output = + single_tma_tensor ? tensor_map_output_static : g_tensor_maps.output_colwise[tensor_id]; + if constexpr (!single_tma_tensor) { + if (leading_thread) { + fence_acquire_tensormap(&tensor_map_input); + fence_acquire_tensormap(&tensor_map_output); + } + __syncthreads(); + } + + const size_t input_scale_stride = + DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(MXFP8_SCALE_DIM)), + static_cast(scale_tensor_alignment_X_rowwise)); + const size_t output_scale_stride = + DIVUP_TO_MULTIPLE(cols, static_cast(scale_tensor_alignment_X_colwise)); + size_t input_scale_base; + size_t output_scale_base; + if constexpr (single_tma_tensor) { + const size_t tensor_start_row = tensor_base / cols; + input_scale_base = tensor_start_row * input_scale_stride; + output_scale_base = (tensor_start_row / MXFP8_SCALE_DIM) * output_scale_stride; + } else { + input_scale_base = tensor_base / MXFP8_SCALE_DIM; + output_scale_base = tensor_base / MXFP8_SCALE_DIM; + } + + bf16 *dequantized_group = nullptr; + if constexpr (RETURN_DEQUANTIZED) { + dequantized_group = dequantized_out + tensor_base; + } + + __shared__ uint64_t input_barriers[buffs_num]; + initialize_barriers(input_barriers, leading_thread); + int input_barrier_parity[buffs_num] = {0}; + + if constexpr (direct_mapper) { + process_chunk( + tensor_map_input, tensor_map_output, input_scales, output_scales, output_rowwise_scales, + input_scale_base, output_scale_base, rows, cols, block_offset_y, block_offset_x, + tma_offset_y, input_shared, dequantized_shared, output_shared, input_barriers, + input_barrier_parity, leading_thread, dequantized_group); + } else { + const size_t blocks_x = DIVUP(cols, static_cast(Traits::CHUNK_DIM_X)); + const size_t blocks_y = DIVUP(rows, static_cast(Traits::CHUNK_DIM_Y)); + const size_t total_blocks = blocks_x * blocks_y; + for (size_t block_id = blockIdx.x; block_id < total_blocks; block_id += gridDim.x) { + const size_t block_y = block_id / blocks_x; + const size_t block_x = block_id - block_y * blocks_x; + process_chunk( + tensor_map_input, tensor_map_output, input_scales, output_scales, output_rowwise_scales, + input_scale_base, output_scale_base, rows, cols, block_y * Traits::CHUNK_DIM_Y, + block_x * Traits::CHUNK_DIM_X, block_y * Traits::CHUNK_DIM_Y, input_shared, + dequantized_shared, output_shared, input_barriers, input_barrier_parity, leading_thread, + dequantized_group); + } + } + + if (leading_thread) { + ptx::cp_async_bulk_wait_group_read<0>(); + } + __syncthreads(); + destroy_barriers(input_barriers, leading_thread); +#else + NVTE_DEVICE_THREAD0_ERROR("Grouped MXFP8 requantization requires Blackwell (SM100+) hardware."); +#endif +} + +template +void launch_group_requantize(const GroupedTensor &input, GroupedTensor *output, Tensor *dequantized, + const size_t num_tensors, const size_t first_logical_dim, + const size_t last_logical_dim, const size_t total_elements, + const int64_t *const offsets_ptr, const int64_t *const first_dims_ptr, + const int64_t *const last_dims_ptr, + const ShapeRepresentation shape_rep, cudaStream_t stream) { + using OType = fp8e4m3; + constexpr size_t buffs_num = Traits::BUFFS_NUM; + constexpr size_t buff_dim = Traits::BUFF_DIM; + constexpr size_t input_bytes = + DIVUP_TO_MULTIPLE(buffs_num * buff_dim * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr size_t dequantized_bytes = + DIVUP_TO_MULTIPLE(Traits::THREADS_PER_CHUNK * sizeof(float), TMA_SWIZZLE_ALIGNMENT); + constexpr size_t output_bytes = + DIVUP_TO_MULTIPLE(buffs_num * buff_dim * sizeof(OType), TMA_SHMEM_ALIGNMENT); + constexpr size_t shared_alignment = TMA_SWIZZLE_ALIGNMENT; + constexpr size_t dynamic_shared_bytes = + input_bytes + dequantized_bytes + output_bytes + shared_alignment; + static_assert(input_bytes % TMA_SWIZZLE_ALIGNMENT == 0 && + output_bytes % TMA_SWIZZLE_ALIGNMENT == 0); + constexpr CUtensorMapSwizzle tma_swizzle = CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B; + + const LaunchConfig launch_config = + get_launch_config(first_logical_dim, last_logical_dim, total_elements, num_tensors); + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_output{}; + create_2D_tensor_map(tensor_map_input, input.data, first_logical_dim, last_logical_dim, + Traits::BUFF_DIM_Y, Traits::BUFF_DIM_X, last_logical_dim, 0, + TypeInfo::size, tma_swizzle); + create_2D_tensor_map(tensor_map_output, output->columnwise_data, first_logical_dim, + last_logical_dim, Traits::BUFF_DIM_Y, Traits::BUFF_DIM_X, last_logical_dim, + 0, TypeInfo::size, tma_swizzle); + + constexpr bool single_tma_tensor = + Traits::SHAPE_REPRESENTATION == ShapeRepresentation::SAME_BOTH_DIMS || + Traits::SHAPE_REPRESENTATION == ShapeRepresentation::VARYING_FIRST_DIM; + if constexpr (!single_tma_tensor) { + alignas(64) CUtensorMap empty_tensor_map{}; + update_tma_descriptors<<>>( + tensor_map_input, empty_tensor_map, empty_tensor_map, tensor_map_output, + reinterpret_cast(input.data.dptr), nullptr, nullptr, + reinterpret_cast(output->columnwise_data.dptr), shape_rep, num_tensors, + first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, false, + true, false); + } + + auto kernel = group_requantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dynamic_shared_bytes)); + kernel<<>>( + tensor_map_input, tensor_map_output, num_tensors, first_logical_dim, last_logical_dim, + launch_config.same_both_rows, offsets_ptr, first_dims_ptr, last_dims_ptr, + reinterpret_cast(input.scale_inv.dptr), + output->scale_inv.dptr == input.scale_inv.dptr + ? nullptr + : reinterpret_cast(output->scale_inv.dptr), + reinterpret_cast(output->columnwise_scale_inv.dptr), + RETURN_DEQUANTIZED ? reinterpret_cast(dequantized->data.dptr) : nullptr); +} + +} // namespace group_requantize_kernel + +void group_requantize(const GroupedTensor &input, GroupedTensor *output, Tensor *dequantized, + cudaStream_t stream) { + using namespace group_requantize_kernel; + + checkCuDriverContext(stream); + NVTE_CHECK(input.scaling_mode == NVTE_MXFP8_1D_SCALING, + "Grouped requantization currently supports only MXFP8 1D scaling for the input."); + NVTE_CHECK(output->scaling_mode == NVTE_MXFP8_1D_SCALING, + "Grouped requantization currently supports only MXFP8 1D scaling for the output."); + NVTE_CHECK(is_supported_by_CC_100(), + "Grouped MXFP8 requantization requires Blackwell (SM100+) hardware."); + CheckInputGroupedTensor(input, "group_requantize_input"); + // The rowwise output is optional, including when its data aliases the input + // but no scale copy is requested. Validate the mandatory columnwise output + // with the common checker, then validate the optional rowwise fields below. + GroupedTensor columnwise_output = *output; + columnwise_output.data.clear(); + columnwise_output.scale_inv.clear(); + CheckOutputGroupedTensor(columnwise_output, "group_requantize_output"); + NVTE_CHECK(input.has_data() && !input.has_columnwise_data(), + "Input must contain rowwise MXFP8 data only."); + NVTE_CHECK(input.data.dptr != nullptr && input.scale_inv.dptr != nullptr, + "Input rowwise MXFP8 data and scales must be allocated."); + NVTE_CHECK(!input.with_gemm_swizzled_scales, + "Input rowwise MXFP8 scales must be in compact format."); + NVTE_CHECK(output->has_columnwise_data() && output->columnwise_data.dptr != nullptr && + output->columnwise_scale_inv.dptr != nullptr, + "Output columnwise MXFP8 data and scales must be allocated."); + NVTE_CHECK(is_fp8_dtype(input.data.dtype), "Input rowwise data must have an FP8 type."); + NVTE_CHECK(output->columnwise_data.dtype == DType::kFloat8E4M3, + "Output columnwise data must have E4M3 type."); + NVTE_CHECK(input.scale_inv.dtype == DType::kFloat8E8M0 && + output->columnwise_scale_inv.dtype == DType::kFloat8E8M0, + "MXFP8 scaling tensors must have E8M0 type."); + NVTE_CHECK(output->data.dptr == nullptr || output->data.dptr == input.data.dptr, + "Output rowwise data must alias the input rowwise data or be null."); + if (output->data.dptr != nullptr) { + NVTE_CHECK(output->data.dtype == input.data.dtype && output->data.shape.size() == 1 && + output->data.numel() == input.data.numel(), + "Output rowwise data must match the input dtype and flattened shape."); + } + if (output->scale_inv.dptr != nullptr) { + NVTE_CHECK(output->scale_inv.dtype == DType::kFloat8E8M0 && + output->scale_inv.shape.size() == 1 && + output->scale_inv.numel() >= input.scale_inv.numel(), + "Output rowwise scales must be E8M0 and have input scale capacity."); + NVTE_CHECK(output->scale_inv.dptr != output->columnwise_scale_inv.dptr, + "Output rowwise and columnwise scales require separate storage."); + NVTE_CHECK((output->scale_inv.dptr == input.scale_inv.dptr) == + (output->with_gemm_swizzled_scales == input.with_gemm_swizzled_scales), + "Output rowwise scales must alias input scales when layouts match and use " + "separate storage when layouts differ."); + } + NVTE_CHECK(input.num_tensors == output->num_tensors, + "Input and output must contain the same number of tensors."); + NVTE_CHECK(input.num_tensors > 0, "Grouped tensor must contain at least one tensor."); + NVTE_CHECK(input.logical_shape.ndim == 2 && output->logical_shape.ndim == 2 && + input.logical_shape.data[0] == output->logical_shape.data[0] && + input.logical_shape.data[1] == output->logical_shape.data[1], + "Input and output logical shapes must match."); + NVTE_CHECK(input.all_same_first_dim() == output->all_same_first_dim() && + input.all_same_last_dim() == output->all_same_last_dim(), + "Input and output grouped shape representations must match."); + NVTE_CHECK(input.data.dptr != output->columnwise_data.dptr, + "In-place MXFP8 data requantization is not supported."); + NVTE_CHECK(is_aligned_ptr(input.data.dptr, TMA_GMEM_ALIGNMENT) && + is_aligned_ptr(output->columnwise_data.dptr, TMA_GMEM_ALIGNMENT), + "Input and output data pointers must be 16-byte aligned."); + + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (input.all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (input.all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (input.all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + + const bool return_dequantized = dequantized != nullptr && dequantized->data.dptr != nullptr; + if (return_dequantized) { + NVTE_CHECK(dequantized->data.dtype == DType::kBFloat16 && + dequantized->data.shape == std::vector({input.logical_shape.data[0], + input.logical_shape.data[1]}), + "The dequantized output must be BF16 with the input logical shape."); + NVTE_CHECK(is_aligned_ptr(dequantized->data.dptr, 16), + "The dequantized output pointer must be 16B aligned."); + } + + const size_t first_logical_dim = input.logical_shape.data[0]; + const size_t last_logical_dim = input.logical_shape.data[1]; + const size_t total_elements = first_logical_dim * last_logical_dim; + NVTE_CHECK(last_logical_dim % MXFP8_SCALE_DIM == 0 || + shape_rep == ShapeRepresentation::VARYING_LAST_DIM || + shape_rep == ShapeRepresentation::VARYING_BOTH_DIMS, + "Every MXFP8 hidden dimension must be divisible by ", MXFP8_SCALE_DIM, "."); + + const int64_t *const offsets_ptr = reinterpret_cast(input.tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(input.first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(input.last_dims.dptr); + +#ifdef NVTE_WITH_CUTEDSL + if (cutedsl_backend::mxfp8_requantize_cutedsl(input, output, dequantized, stream)) { + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } +#endif + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + input.data.dtype, IType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + output->with_gemm_swizzled_scales, OUTPUT_SCALES_SWIZZLED, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + return_dequantized, RETURN_DEQUANTIZED, + TRANSFORMER_ENGINE_GROUP_TENSOR_SHAPE_REPRESENTATION_SWITCH(shape_rep, SHAPE_REP, { + using ActiveTraits = RequantizeTraits; + launch_group_requantize( + input, output, dequantized, input.num_tensors, first_logical_dim, + last_logical_dim, total_elements, offsets_ptr, first_dims_ptr, last_dims_ptr, + shape_rep, stream); + });););); // NOLINT(*), readability/fn_size + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +void nvte_group_requantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, + NVTETensor dequantized, const NVTEQuantizationConfig /*quant_config*/, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_requantize); + using namespace transformer_engine; + const GroupedTensor *const input_cu = convertNVTEGroupedTensorCheck(input); + GroupedTensor *const output_cu = convertNVTEGroupedTensorCheck(output); + dispatch::mxfp8::group_requantize( + *input_cu, output_cu, dequantized != nullptr ? convertNVTETensorCheck(dequantized) : nullptr, + stream); +} diff --git a/transformer_engine/common/cast/mxfp8/requantize_mxfp8_cutedsl.cuh b/transformer_engine/common/cast/mxfp8/requantize_mxfp8_cutedsl.cuh new file mode 100644 index 00000000000..2f0ce159c7a --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/requantize_mxfp8_cutedsl.cuh @@ -0,0 +1,138 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_REQUANTIZE_MXFP8_CUTEDSL_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_REQUANTIZE_MXFP8_CUTEDSL_CUH_ + +#include +#include +#include +#include +#include + +#include "../../common.h" +#include "../../tvm_ffi_bridge.h" + +namespace transformer_engine { +namespace cutedsl_backend { + +struct MXFP8RequantConfig { + DType dtype; + size_t rows, hidden, groups, uniform_rows; + bool swizzled, rowwise_output, return_dequantized; + int sm_arch = cuda::sm_arch(); + int sm_count = cuda::sm_count(); + + std::string to_key() const { + return "cutedsl_mxfp8_requant_" + std::string(to_string(dtype)) + "_" + std::to_string(rows) + + "_" + std::to_string(hidden) + "_" + std::to_string(groups) + "_" + + std::to_string(uniform_rows) + "_" + std::to_string(swizzled) + "_" + + std::to_string(rowwise_output) + "_" + std::to_string(return_dequantized) + "_" + + std::to_string(sm_arch) + "_" + std::to_string(sm_count); + } + + bool retrieve_func_from_python(const std::string &name) const { + auto entrypoint = tvm::ffi::Function::GetGlobal("get_mxfp8_requantization_function"); + if (!entrypoint) return false; + return (*entrypoint)(tvm::ffi::String(name), tvm::ffi::String(to_string(dtype)), + static_cast(rows), static_cast(hidden), + static_cast(groups), static_cast(uniform_rows), swizzled, + rowwise_output, return_dequantized, sm_count) + .try_cast() + .value_or(false); + } + + std::optional get_kernel() const { + using namespace tvm_ffi_bridge; + // Shape-specialized keys do not fit the quantization cache's 32-bit ID. + // Preserve cached TVM handles until process exit, as TVMFFIConfigCache does. + struct Cache { + std::shared_mutex mutex; + std::unordered_map> kernels; + }; + static auto *cache = new Cache(); + auto ¢ral = TVMFFICentral::getInstance(); + if (!central.get_cutedsl_backend_enabled()) return std::nullopt; + const auto key = to_key(); + { + std::shared_lock lock(cache->mutex); + auto it = cache->kernels.find(key); + if (it != cache->kernels.end()) return it->second; + } + std::unique_lock lock(cache->mutex); + auto it = cache->kernels.find(key); + if (it != cache->kernels.end()) return it->second; + std::optional kernel; + if (auto fn = central.load_tvm_ffi_function(*this)) { + kernel.emplace(make_tvm_ffi_kernel(std::move(*fn))); + } + cache->kernels.emplace(key, kernel); + return kernel; + } +}; + +inline bool mxfp8_requantize_cutedsl(const GroupedTensor &input, GroupedTensor *output, + Tensor *dequantized, cudaStream_t stream) { + using namespace tvm_ffi_bridge; + if (!TVMFFICentral::getInstance().get_cutedsl_backend_enabled() || !input.all_same_last_dim()) + return false; + const size_t rows = input.logical_shape.data[0]; + const size_t hidden = input.logical_shape.data[1]; + const size_t groups = input.num_tensors; + if (rows == 0 || hidden == 0 || rows % 128 || hidden % 128 || (hidden >= 512 && hidden % 512)) + return false; + // TMA's compact scale matrix requires a 16-byte row stride for wide tensors. + // The narrow specializations copy full-width compact scales as 1D bulk data. + if (input.all_same_shape() && (rows % groups || (rows / groups) % 128)) return false; + if (!input.all_same_shape() && + (input.tensor_offsets.dptr == nullptr || input.tensor_offsets.numel() < groups + 1)) { + return false; + } + const bool rowwise = output->with_gemm_swizzled_scales && output->scale_inv.dptr != nullptr; + if (!is_aligned_ptr(input.scale_inv.dptr, 16) || + !is_aligned_ptr(output->columnwise_scale_inv.dptr, 16) || + (rowwise && !is_aligned_ptr(output->scale_inv.dptr, 16))) + return false; + const MXFP8RequantConfig config{input.data.dtype, + rows, + hidden, + groups, + input.all_same_shape() ? rows / groups : 0, + output->with_gemm_swizzled_scales, + rowwise, + dequantized != nullptr && dequantized->data.dptr != nullptr}; + auto kernel = config.get_kernel(); + if (!kernel) return false; + + const int device = cuda::current_device(); + const SimpleTensor src(input.data.dptr, {rows, hidden}, input.data.dtype); + const SimpleTensor sf(input.scale_inv.dptr, {rows, hidden / 32}, DType::kFloat8E8M0); + const SimpleTensor dst(output->columnwise_data.dptr, {rows, hidden}, DType::kFloat8E4M3); + const SimpleTensor col_sf(output->columnwise_scale_inv.dptr, {rows * hidden / 32}, + DType::kFloat8E8M0); + DLTensorWrapper mSrc(src, false, device), mSf(sf, false, device), mDst(dst, false, device), + mColSf(col_sf, false, device), mOffsets, mRowSf, mDequantized; + if (!input.all_same_shape()) { + const SimpleTensor offsets(input.tensor_offsets.dptr, {groups + 1}, DType::kInt64); + mOffsets = DLTensorWrapper(offsets, false, device); + } + if (rowwise) { + const SimpleTensor row_sf(output->scale_inv.dptr, {rows * hidden / 32}, DType::kFloat8E8M0); + mRowSf = DLTensorWrapper(row_sf, false, device); + } + if (config.return_dequantized) { + const SimpleTensor decoded(dequantized->data.dptr, {rows, hidden}, DType::kBFloat16); + mDequantized = DLTensorWrapper(decoded, false, device); + } + (*kernel)(&mSrc, &mSf, &mOffsets, &mDst, &mRowSf, &mColSf, &mDequantized, + static_cast(stream)); + return true; +} + +} // namespace cutedsl_backend +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_REQUANTIZE_MXFP8_CUTEDSL_CUH_ diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 4b64d91b932..bd06643369f 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -1156,6 +1156,9 @@ template <> struct is_fp4 : std::true_type {}; #endif +// Number of values represented by one MXFP8 E8M0 scaling factor. +constexpr size_t MXFP8_SCALE_DIM = 32; + // [128,4] rowwise and [4,128] colwise alignment requirements for the tensor with scaling factors constexpr size_t scale_tensor_alignment_X_rowwise = 4; constexpr size_t scale_tensor_alignment_Y_rowwise = 128; diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index ca3eb13ed24..a4c055e971e 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -429,44 +429,29 @@ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t str void nvte_group_dequantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); -/*! \brief Fused grouped requantization. Currently only MXFP8 1D scaling is supported. - * - * Makes a rowwise-only grouped MXFP8 tensor GEMM-ready in one kernel: - * dequantizes, builds the columnwise MXFP8 copy with GEMM-swizzled per-group - * scales, and re-emits the rowwise scales at their GEMM-swizzled addresses. - * Optionally also materializes the BF16 dequantized values. Replaces the - * separate group_dequantize -> group_quantize(columnwise) -> - * grouped_swizzle(rowwise scales) chain. - * - * Requirements: SM100+, MXFP8 1D scaling, the hidden dim and (caller contract, - * offsets are device-resident) every group's row count multiples of 128. Rows - * at or past tensor_offsets[num_groups] (capacity-mode / paged-stash tail) are - * left untouched, matching the unfused chain. On return the output tensor's - * with_gemm_swizzled_scales metadata is set. - * - * \param[in] input Rowwise MXFP8 tensor with compact (unswizzled) - * E8M0 scales; groups stacked along the first dim. - * \param[in,out] output Destination: columnwise E4M3 data + - * columnwise_scale_inv (per-group GEMM-swizzled) - * and scale_inv (the GEMM-swizzled copy of the - * input's rowwise scales). Rowwise data is not - * written; the GEMM keeps consuming the input's. - * \param[in] tensor_offsets Int64 device tensor of num_groups + 1 exclusive - * ELEMENT offsets, offsets[g] = row offset x cols - * (the grouped tensor's cached tensor_offsets; - * offsets[0] = 0, offsets[num_groups] = live rows - * x cols, which may be less than the allocated - * rows x cols). - * \param[out] dequantized Optional BF16 [rows, cols] tensor; pass NULL - * (or an unallocated tensor) to skip it. - * \param[in] quant_config Quantization options. `use_fast_math` selects a - * BF16 intermediate instead of the default FP32 - * path. - * \param[in] stream CUDA stream used for the operation. +/*! \brief Requantizes a grouped tensor. Currently supports MXFP8 1D scaling only. + * + * Equivalent to grouped MXFP8 dequantization followed by columnwise quantization, + * using BF16 intermediates, with optional BF16 dequantized output. + * + * Input and output must describe the same grouped shapes. Input must contain + * rowwise FP8 data with compact E8M0 scales. Output must contain columnwise E4M3 + * data with compact or GEMM-swizzled E8M0 scales. Optional output rowwise data + * must alias the input. Optional output rowwise scales must alias the input + * scales for matching layouts, or use separate storage for different layouts. + * Each group's row and column counts must be divisible by 128. + * + * \param[in] input Compact rowwise MXFP8 grouped tensor. + * \param[in,out] output Columnwise MXFP8 E4M3 grouped tensor with + * optional rowwise data and scales. + * \param[out] dequantized Optional BF16 tensor with the input logical shape; + * pass NULL or an unallocated tensor to skip it. + * \param[in] quant_config Unused; may be NULL. + * \param[in] stream CUDA stream used for the operation. */ -void nvte_group_requantize(const NVTETensor input, NVTETensor output, - const NVTETensor tensor_offsets, NVTETensor dequantized, - const NVTEQuantizationConfig quant_config, cudaStream_t stream); +void nvte_group_requantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, + NVTETensor dequantized, const NVTEQuantizationConfig quant_config, + cudaStream_t stream); /*! \brief Casts multiple input tensors to quantized output tensors. * diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 06d03cd6e1c..6f026272764 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -886,14 +886,8 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, "Requantizing a grouped input requires dims that are multiples of 128, but got (", total_tokens, ", ", hidden_dim, ")."); - // Fused path (default; NVTE_FUSED_GROUP_REQUANTIZE=0 recovers the unfused chain): one - // kernel replaces the group_dequantize -> group_quantize(columnwise) -> - // grouped_swizzle(rowwise scales) chain below, with the dequantized values living only in - // shared memory unless requested. The BF16-intermediate kernel variant reproduces the - // unfused chain's numerics, hence the otype gate; anything the kernel does not cover - // falls through to the unfused chain. - // The kernel takes the grouped tensor's cached element-based tensor_offsets; - // a prefix-sum over first_dims is only the fallback when they are absent. + // The common API selects the faster requantization kernel when no BF16 output + // is requested. Its BF16 intermediate reproduces the unfused chain's numerics. const bool tensor_offsets_usable = tensor_offsets.has_value() && tensor_offsets->scalar_type() == at::kLong && tensor_offsets->numel() == static_cast(num_tensors) + 1; @@ -928,47 +922,48 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, at::Tensor columnwise_data = at::empty({tokens_i64 * hidden_i64}, options); at::Tensor columnwise_scale_inv = at::empty({tokens_i64 / 32 * hidden_i64}, options); at::Tensor swizzled_rowwise_scale_inv = at::empty({static_cast(num_scales)}, options); - at::Tensor dequantized; - if (return_dequantized) { - dequantized = - at::empty({tokens_i64, hidden_i64}, rowwise_data.options().dtype(at::kBFloat16)); - } - - TensorWrapper input_nvte(NVTE_MXFP8_1D_SCALING); - input_nvte.set_rowwise_data(rowwise_data.data_ptr(), op_dtype, - std::vector{total_tokens, hidden_dim}); + const at::Tensor first_dims_i64 = + first_dims.has_value() + ? (first_dims->scalar_type() == at::kLong ? *first_dims : first_dims->to(at::kLong)) + : at::floor_divide(element_offsets.slice(0, 1, num_tensors + 1) - + element_offsets.slice(0, 0, num_tensors), + hidden_i64); + const std::vector grouped_shape = {total_tokens, hidden_dim}; + const std::vector flat_data_shape = {total_tokens * hidden_dim}; + const std::vector scale_shape = {num_scales}; + const std::vector dims_shape = {num_tensors}; + const std::vector offsets_shape = {num_tensors + 1}; + + GroupedTensorWrapper input_nvte(num_tensors, grouped_shape, NVTE_MXFP8_1D_SCALING); + input_nvte.set_rowwise_data(rowwise_data.data_ptr(), op_dtype, flat_data_shape); input_nvte.set_rowwise_scale_inv(rowwise_scale_inv.data_ptr(), DType::kFloat8E8M0, - std::vector{total_tokens, hidden_dim / 32}); + getTensorShape(rowwise_scale_inv)); + input_nvte.set_first_dims(first_dims_i64.data_ptr(), DType::kInt64, dims_shape); + input_nvte.set_tensor_offsets(element_offsets.data_ptr(), DType::kInt64, offsets_shape); - // After the kernel the output is GEMM-ready: it keeps consuming the input's rowwise - // data, so that slot aliases the input. - TensorWrapper output_nvte(NVTE_MXFP8_1D_SCALING); - output_nvte.set_rowwise_data(rowwise_data.data_ptr(), op_dtype, - std::vector{total_tokens, hidden_dim}); + GroupedTensorWrapper output_nvte(num_tensors, grouped_shape, NVTE_MXFP8_1D_SCALING); + output_nvte.set_rowwise_data(rowwise_data.data_ptr(), op_dtype, flat_data_shape); output_nvte.set_rowwise_scale_inv(swizzled_rowwise_scale_inv.data_ptr(), DType::kFloat8E8M0, - std::vector{num_scales}); + scale_shape); output_nvte.set_columnwise_data(columnwise_data.data_ptr(), DType::kFloat8E4M3, - std::vector{total_tokens, hidden_dim}); + flat_data_shape); output_nvte.set_columnwise_scale_inv(columnwise_scale_inv.data_ptr(), DType::kFloat8E8M0, - std::vector{total_tokens / 32 * hidden_dim}); + scale_shape); + output_nvte.set_first_dims(first_dims_i64.data_ptr(), DType::kInt64, dims_shape); + output_nvte.set_tensor_offsets(element_offsets.data_ptr(), DType::kInt64, offsets_shape); + output_nvte.set_with_gemm_swizzled_scales(true); - TensorWrapper element_offsets_nvte; - element_offsets_nvte.set_rowwise_data(element_offsets.data_ptr(), DType::kInt64, - std::vector{num_tensors + 1}); + at::Tensor dequantized; TensorWrapper dequantized_nvte; if (return_dequantized) { - dequantized_nvte.set_rowwise_data(dequantized.data_ptr(), DType::kBFloat16, - std::vector{total_tokens, hidden_dim}); + dequantized = + at::empty({tokens_i64, hidden_i64}, rowwise_data.options().dtype(at::kBFloat16)); + dequantized_nvte.set_rowwise_data(dequantized.data_ptr(), DType::kBFloat16, grouped_shape); } - // BF16 intermediate: matches the unfused chain, which materializes the dequantized - // tensor in otype (gated to BF16 above) before requantizing. - QuantizationConfigWrapper quant_config; - quant_config.set_use_fast_math(true); - NVTE_SCOPED_GIL_RELEASE({ - nvte_group_requantize(input_nvte.data(), output_nvte.data(), element_offsets_nvte.data(), - return_dequantized ? dequantized_nvte.data() : nullptr, quant_config, + nvte_group_requantize(input_nvte.data(), output_nvte.data(), + return_dequantized ? dequantized_nvte.data() : nullptr, nullptr, at::cuda::getCurrentCUDAStream()); }); @@ -976,11 +971,7 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, grouped_x.attr("columnwise_data") = columnwise_data; grouped_x.attr("columnwise_scale_inv") = columnwise_scale_inv; grouped_x.attr("_with_gemm_swizzled_scales") = py::cast(true); - - if (return_dequantized) { - return py::cast(dequantized); - } - return py::none(); + return return_dequantized ? py::cast(dequantized) : py::none(); } // Dequantize first: it reads the rowwise scales, which the swizzle below replaces. Left