From c602294440e45cf26bfe8f949c6f8cc25d09a125 Mon Sep 17 00:00:00 2001 From: YangFei1990 Date: Wed, 12 Aug 2026 15:15:23 -0700 Subject: [PATCH 1/3] [Common/PyTorch] Fused grouped MXFP8 requantization Replace the group_dequantize -> group_quantize(columnwise) -> grouped_swizzle(rowwise scales) chain in group_requantize_inplace with a single kernel (NVTE_FUSED_GROUP_REQUANTIZE=0 restores the unfused path). Co-authored-by: Oleg Goncharov Signed-off-by: YangFei1990 --- tests/cpp/operator/CMakeLists.txt | 1 + .../test_fused_group_requantize_mxfp8.cu | 344 +++++++++++ transformer_engine/common/CMakeLists.txt | 1 + .../common/cast/fused_group_requantize.cu | 565 ++++++++++++++++++ .../common/include/transformer_engine/cast.h | 40 ++ .../pytorch/csrc/extensions/cast.cpp | 100 ++++ 6 files changed, 1051 insertions(+) create mode 100644 tests/cpp/operator/test_fused_group_requantize_mxfp8.cu create mode 100644 transformer_engine/common/cast/fused_group_requantize.cu diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 06ed56cc5e..f58ab349c3 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -18,6 +18,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_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 new file mode 100644 index 0000000000..fde62fd8d7 --- /dev/null +++ b/tests/cpp/operator/test_fused_group_requantize_mxfp8.cu @@ -0,0 +1,344 @@ +/************************************************************************* + * 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_fused_group_requantize_mxfp8(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/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 6ceadc7405..e7aaf78f6c 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -267,6 +267,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 gemm/cutlass_grouped_gemm.cu hadamard_transform/group_hadamard_transform.cu hadamard_transform/graph_safe_group_hadamard_transform.cu diff --git a/transformer_engine/common/cast/fused_group_requantize.cu b/transformer_engine/common/cast/fused_group_requantize.cu new file mode 100644 index 0000000000..2a8787bd4e --- /dev/null +++ b/transformer_engine/common/cast/fused_group_requantize.cu @@ -0,0 +1,565 @@ +/************************************************************************* + * 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 "../common.h" +#include "../util/ptx.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_read<0>(); + } + __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); +} + +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; + } + + 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)); + + const bool use_fast_math = quant_config != nullptr && quant_config->use_fast_math; + 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_fused_group_requantize_mxfp8(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_fused_group_requantize_mxfp8); + + 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/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 4d6d24ba65..2ab24af94c 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -429,6 +429,46 @@ 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 MXFP8 requantization. + * + * 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. + */ +void nvte_fused_group_requantize_mxfp8(const NVTETensor input, NVTETensor output, + const NVTETensor tensor_offsets, NVTETensor dequantized, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream); + /*! \brief Casts multiple input tensors to quantized output tensors. * * \param[in] inputs List of input tensors to be cast. diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 5ce0261c82..017e20d1ae 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -723,6 +723,106 @@ 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. + const bool has_usable_offsets = + (tensor_offsets.has_value() && tensor_offsets->scalar_type() == at::kLong && + tensor_offsets->numel() == static_cast(num_tensors) + 1) || + first_dims.has_value(); + const bool use_fused_kernel = + transformer_engine::getenv("NVTE_FUSED_GROUP_REQUANTIZE", true) && need_columnwise && + has_usable_offsets && otype == DType::kBFloat16 && + quantizer.attr("dtype").cast() == DType::kFloat8E4M3 && + transformer_engine::cuda::sm_arch() >= 100; + if (use_fused_kernel) { + const auto rowwise_data = grouped_x.attr("rowwise_data").cast(); + const auto rowwise_scale_inv = grouped_x.attr("scale_inv").cast(); + const auto options = rowwise_data.options().dtype(at::kByte); + const auto tokens_i64 = static_cast(total_tokens); + const auto hidden_i64 = static_cast(hidden_dim); + const size_t num_scales = total_tokens * hidden_dim / 32; + + // Element-based exclusive-cumsum offsets ([num_tensors + 1], device) for the + // per-group columnwise scale bases and the live-row bound. + at::Tensor element_offsets; + if (tensor_offsets.has_value() && tensor_offsets->scalar_type() == at::kLong && + tensor_offsets->numel() == static_cast(num_tensors) + 1) { + element_offsets = *tensor_offsets; + } else { + const at::Tensor first_dims_i64 = + first_dims->scalar_type() == at::kLong ? *first_dims : first_dims->to(at::kLong); + element_offsets = splits_to_offsets(first_dims_i64, static_cast(hidden_dim)); + } + + // Grouped tensors carry data and scales as flat 1D buffers. + 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)); + } + + const DType wire_dtype = quantizer.attr("dtype").cast(); + TensorWrapper input_nvte(NVTE_MXFP8_1D_SCALING); + input_nvte.set_rowwise_data(rowwise_data.data_ptr(), wire_dtype, + std::vector{total_tokens, hidden_dim}); + input_nvte.set_rowwise_scale_inv(rowwise_scale_inv.data_ptr(), DType::kFloat8E8M0, + std::vector{total_tokens, hidden_dim / 32}); + + // 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(), wire_dtype, + std::vector{total_tokens, hidden_dim}); + output_nvte.set_rowwise_scale_inv(swizzled_rowwise_scale_inv.data_ptr(), + DType::kFloat8E8M0, std::vector{num_scales}); + output_nvte.set_columnwise_data(columnwise_data.data_ptr(), DType::kFloat8E4M3, + std::vector{total_tokens, hidden_dim}); + output_nvte.set_columnwise_scale_inv( + columnwise_scale_inv.data_ptr(), DType::kFloat8E8M0, + std::vector{total_tokens / 32 * hidden_dim}); + + TensorWrapper element_offsets_nvte; + element_offsets_nvte.set_rowwise_data(element_offsets.data_ptr(), DType::kInt64, + std::vector{num_tensors + 1}); + TensorWrapper dequantized_nvte; + if (return_dequantized) { + dequantized_nvte.set_rowwise_data(dequantized.data_ptr(), DType::kBFloat16, + std::vector{total_tokens, hidden_dim}); + } + + // 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_fused_group_requantize_mxfp8( + input_nvte.data(), output_nvte.data(), element_offsets_nvte.data(), + return_dequantized ? dequantized_nvte.data() : nullptr, quant_config, + at::cuda::getCurrentCUDAStream()); + }); + + grouped_x.attr("scale_inv") = swizzled_rowwise_scale_inv; + 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(); + } + // Dequantize first: it reads the rowwise scales, which the swizzle below replaces. Left // undefined when nothing consumes it, which skips the pass entirely. at::Tensor dequantized; From 609ca6fb0db28bf39b7e5816b58dc9c8186d7acd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 22:18:09 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/cast/fused_group_requantize.cu | 76 ++++++++----------- .../common/include/transformer_engine/cast.h | 6 +- .../pytorch/csrc/extensions/cast.cpp | 20 +++-- 3 files changed, 44 insertions(+), 58 deletions(-) diff --git a/transformer_engine/common/cast/fused_group_requantize.cu b/transformer_engine/common/cast/fused_group_requantize.cu index 2a8787bd4e..95b58876e8 100644 --- a/transformer_engine/common/cast/fused_group_requantize.cu +++ b/transformer_engine/common/cast/fused_group_requantize.cu @@ -8,11 +8,10 @@ * \brief Fused grouped MXFP8 requantization: rowwise wire tensor -> GEMM-ready. */ -#include - #include #include #include +#include #include #include @@ -48,8 +47,8 @@ __device__ __forceinline__ uint16_t e8m0_to_bf16_bits(const e8m0_t biased_exp) { } template -__device__ __forceinline__ ptx::bf16x2 dequantize_mxfp8_2x( - const ptx::FPx2 &values, const e8m0_t scale_code) { +__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) @@ -57,19 +56,17 @@ __device__ __forceinline__ ptx::bf16x2 dequantize_mxfp8_2x( // 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); + 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)); + 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)); + 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; } @@ -78,8 +75,8 @@ __device__ __forceinline__ ptx::bf16x2 dequantize_mxfp8_2x( // 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); + const uint32_t scale_x2 = + static_cast(scale_bits) | (static_cast(scale_bits) << 16); if constexpr (std::is_same_v) { asm volatile( "{\n\t" @@ -183,11 +180,9 @@ __device__ __forceinline__ void store_colwise_2x_to_shared(OType *const output, 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, + 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, + 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) @@ -258,11 +253,9 @@ __global__ void __launch_bounds__(kThreads) 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; + 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)] = + 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); @@ -281,8 +274,7 @@ __global__ void __launch_bounds__(kThreads) *reinterpret_cast(&dequantized[local_row][shared_col + i]) = result; if constexpr (kReturnDequantized) { *reinterpret_cast( - &dequantized_vec[i / kDequantizedVecSize].data.elt[i % kDequantizedVecSize]) = - result; + &dequantized_vec[i / kDequantizedVecSize].data.elt[i % kDequantizedVecSize]) = result; } } } else { @@ -298,8 +290,7 @@ __global__ void __launch_bounds__(kThreads) } } if constexpr (kReturnDequantized) { - bf16 *const dequantized_out_ptr = - dequantized_out + static_cast(row) * num_cols + col; + 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); } @@ -345,16 +336,15 @@ __global__ void __launch_bounds__(kThreads) // 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 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; + 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); @@ -442,8 +432,7 @@ void fused_group_requantize(const Tensor &input, Tensor *output, const Tensor &t "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.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, @@ -458,8 +447,7 @@ void fused_group_requantize(const Tensor &input, Tensor *output, const Tensor &t 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."); + 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; @@ -503,14 +491,14 @@ void fused_group_requantize(const Tensor &input, Tensor *output, const Tensor &t 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; + 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( + 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."); } @@ -549,9 +537,9 @@ void fused_group_requantize(const Tensor &input, Tensor *output, const Tensor &t } // namespace transformer_engine void nvte_fused_group_requantize_mxfp8(const NVTETensor input, NVTETensor output, - const NVTETensor tensor_offsets, NVTETensor dequantized, - const NVTEQuantizationConfig quant_config, - cudaStream_t stream) { + const NVTETensor tensor_offsets, NVTETensor dequantized, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { using namespace transformer_engine; NVTE_API_CALL(nvte_fused_group_requantize_mxfp8); diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 2ab24af94c..0349d8a6de 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -465,9 +465,9 @@ void nvte_group_dequantize(const NVTEGroupedTensor input, NVTEGroupedTensor outp * \param[in] stream CUDA stream used for the operation. */ void nvte_fused_group_requantize_mxfp8(const NVTETensor input, NVTETensor output, - const NVTETensor tensor_offsets, NVTETensor dequantized, - const NVTEQuantizationConfig quant_config, - cudaStream_t stream); + const NVTETensor tensor_offsets, 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 352391c502..ac8ea52a32 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -789,8 +789,7 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, // Grouped tensors carry data and scales as flat 1D buffers. 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 swizzled_rowwise_scale_inv = at::empty({static_cast(num_scales)}, options); at::Tensor dequantized; if (return_dequantized) { dequantized = @@ -809,13 +808,12 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, TensorWrapper output_nvte(NVTE_MXFP8_1D_SCALING); output_nvte.set_rowwise_data(rowwise_data.data_ptr(), wire_dtype, std::vector{total_tokens, hidden_dim}); - output_nvte.set_rowwise_scale_inv(swizzled_rowwise_scale_inv.data_ptr(), - DType::kFloat8E8M0, std::vector{num_scales}); + output_nvte.set_rowwise_scale_inv(swizzled_rowwise_scale_inv.data_ptr(), DType::kFloat8E8M0, + std::vector{num_scales}); output_nvte.set_columnwise_data(columnwise_data.data_ptr(), DType::kFloat8E4M3, std::vector{total_tokens, hidden_dim}); - output_nvte.set_columnwise_scale_inv( - columnwise_scale_inv.data_ptr(), DType::kFloat8E8M0, - std::vector{total_tokens / 32 * hidden_dim}); + output_nvte.set_columnwise_scale_inv(columnwise_scale_inv.data_ptr(), DType::kFloat8E8M0, + std::vector{total_tokens / 32 * hidden_dim}); TensorWrapper element_offsets_nvte; element_offsets_nvte.set_rowwise_data(element_offsets.data_ptr(), DType::kInt64, @@ -832,10 +830,10 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, quant_config.set_use_fast_math(true); NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_group_requantize_mxfp8( - input_nvte.data(), output_nvte.data(), element_offsets_nvte.data(), - return_dequantized ? dequantized_nvte.data() : nullptr, quant_config, - at::cuda::getCurrentCUDAStream()); + nvte_fused_group_requantize_mxfp8(input_nvte.data(), output_nvte.data(), + element_offsets_nvte.data(), + return_dequantized ? dequantized_nvte.data() : nullptr, + quant_config, at::cuda::getCurrentCUDAStream()); }); grouped_x.attr("scale_inv") = swizzled_rowwise_scale_inv; From 07214aa5df39c07de7f7361da4118b5f3d1577a4 Mon Sep 17 00:00:00 2001 From: YangFei1990 Date: Wed, 12 Aug 2026 16:17:43 -0700 Subject: [PATCH 3/3] fix zero tensor dim Signed-off-by: YangFei1990 --- .../pytorch/csrc/extensions/cast.cpp | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 017e20d1ae..00682a45ab 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -731,13 +731,16 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, // 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. - const bool has_usable_offsets = - (tensor_offsets.has_value() && tensor_offsets->scalar_type() == at::kLong && - tensor_offsets->numel() == static_cast(num_tensors) + 1) || - first_dims.has_value(); + const bool tensor_offsets_usable = + tensor_offsets.has_value() && tensor_offsets->scalar_type() == at::kLong && + tensor_offsets->numel() == static_cast(num_tensors) + 1; + const bool has_usable_offsets = tensor_offsets_usable || first_dims.has_value(); + // total_tokens > 0: an empty grouped tensor carries null data pointers, which the + // unfused chain's dedicated empty-input handling accepts and the kernel's pointer + // validation (correctly) rejects. const bool use_fused_kernel = transformer_engine::getenv("NVTE_FUSED_GROUP_REQUANTIZE", true) && need_columnwise && - has_usable_offsets && otype == DType::kBFloat16 && + has_usable_offsets && total_tokens > 0 && otype == DType::kBFloat16 && quantizer.attr("dtype").cast() == DType::kFloat8E4M3 && transformer_engine::cuda::sm_arch() >= 100; if (use_fused_kernel) { @@ -751,8 +754,7 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer, // Element-based exclusive-cumsum offsets ([num_tensors + 1], device) for the // per-group columnwise scale bases and the live-row bound. at::Tensor element_offsets; - if (tensor_offsets.has_value() && tensor_offsets->scalar_type() == at::kLong && - tensor_offsets->numel() == static_cast(num_tensors) + 1) { + if (tensor_offsets_usable) { element_offsets = *tensor_offsets; } else { const at::Tensor first_dims_i64 =