Skip to content

Add 2d quant for mxfp8 #2634

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kunlunl wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Add 2d quant for mxfp8 #2634

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
32de019
Support MXFP8 2D quantization
kunlunl May 14, 2026
93e34de
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 14, 2026
30ad3a4
Fix comments
kunlunl May 14, 2026
9142f1d
Fix comments
kunlunl May 15, 2026
5ef3cfe
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 15, 2026
327b732
Merge origin/main into 2d_mxfp8
ptrendx Aug 4, 2026
80a51af
Support explicit 2D MXFP8 grouped quantization
kunlunl Aug 10, 2026
2e1dfd2
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions qa/L0_pytorch_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ fi
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py || test_fail "test_checkpoint.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_router.xml $TE_PATH/tests/pytorch/test_fused_router.py || test_fail "test_fused_router.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml $TE_PATH/tests/pytorch/test_partial_cast.py || test_fail "test_partial_cast.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8_2d_quantize.xml $TE_PATH/tests/pytorch/test_mxfp8_2d_quantize.py || test_fail "test_mxfp8_2d_quantize.py"
# Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision
NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py"
NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc_deterministic.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_mhc.py"
Expand Down
254 changes: 254 additions & 0 deletions tests/cpp/operator/test_cast_mxfp8.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
* See LICENSE for license information.
************************************************************************/

#include <algorithm>

#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
Expand Down Expand Up @@ -36,6 +38,12 @@ enum ActivationType {
SReLU
};

enum MXFP82DScalingDirection {
RowwiseOnly,
ColwiseOnly,
Bidirectional
};

template <typename InputType, typename OutputType>
void compute_ref(const ProcessingMethod processing_method,
float (*OP)(const float),
Expand Down Expand Up @@ -165,6 +173,70 @@ void compute_ref(const ProcessingMethod processing_method,
}
}

template <typename InputType, typename OutputType>
void compute_ref_2d_quantize(const bool rowwise,
const bool colwise,
const InputType* input,
OutputType* output_rowwise,
OutputType* output_colwise,
fp8e8m0* output_scales_rowwise,
fp8e8m0* output_scales_colwise,
const size_t rows,
const size_t cols,
const size_t scales_stride_rowwise,
const size_t scales_stride_colwise) {
const size_t tile_size_Y = 32;
const size_t tile_size_X = 32;
const size_t tiles_num_Y = (rows + tile_size_Y - 1) / tile_size_Y;
const size_t tiles_num_X = (cols + tile_size_X - 1) / tile_size_X;

#pragma omp parallel for collapse(2) proc_bind(spread)
for (size_t tile_Y = 0; tile_Y < tiles_num_Y; ++tile_Y) {
for (size_t tile_X = 0; tile_X < tiles_num_X; ++tile_X) {
const size_t i_min = tile_Y * tile_size_Y;
const size_t i_max = std::min(i_min + tile_size_Y, rows);
const size_t j_min = tile_X * tile_size_X;
const size_t j_max = std::min(j_min + tile_size_X, cols);

float block_amax = 0.0f;
for (size_t i = i_min; i < i_max; ++i) {
for (size_t j = j_min; j < j_max; ++j) {
const size_t idx = i * cols + j;
block_amax = std::max(block_amax, std::abs(static_cast<float>(input[idx])));
}
}

const fp8e8m0 biased_exponent =
float_to_e8m0(block_amax * Quantized_Limits<OutputType>::max_reciprocal());
const float scale_reciprocal = exp2f_rcp(biased_exponent);

if (rowwise) {
for (size_t i = i_min; i < i_max; ++i) {
output_scales_rowwise[i * scales_stride_rowwise + tile_X] = biased_exponent;
for (size_t j = j_min; j < j_max; ++j) {
const size_t idx = i * cols + j;
output_rowwise[idx] =
static_cast<OutputType>(static_cast<float>(input[idx]) *
scale_reciprocal);
}
}
}

if (colwise) {
for (size_t j = j_min; j < j_max; ++j) {
output_scales_colwise[tile_Y * scales_stride_colwise + j] = biased_exponent;
for (size_t i = i_min; i < i_max; ++i) {
const size_t idx = i * cols + j;
output_colwise[idx] =
static_cast<OutputType>(static_cast<float>(input[idx]) *
scale_reciprocal);
}
}
}
}
}
}

/**
* Scaling along single dimension (either rows or columns)
* Produces one set of output data and the corresponding data of the fused operation (dbias):
Expand Down Expand Up @@ -521,6 +593,106 @@ void performTest_x2(const ProcessingMethod processing_method,
}
}

template <typename InputType, typename OutputType>
void performTest_2d_quantize(const std::vector<size_t>& shape,
const MXFP82DScalingDirection scaling_direction,
InputsFillCase fill_case) {
using namespace test;
using EncodingType = fp32;
DType itype = TypeInfo<InputType>::dtype;
DType otype = TypeInfo<OutputType>::dtype;

if (shape.size() < 2) {
GTEST_SKIP();
}

const size_t rows = first_dimension(shape);
const size_t cols = last_dimension(shape);

const bool rowwise = scaling_direction != MXFP82DScalingDirection::ColwiseOnly;
const bool colwise = scaling_direction != MXFP82DScalingDirection::RowwiseOnly;

const std::array<size_t,4> scale_dims_rowwise = get_scale_tensor_dims(rows, cols, 1, 32);
const std::array<size_t,4> scale_dims_colwise = get_scale_tensor_dims(rows, cols, 32, 1);

const size_t unpadded_blocks_Y_rowwise = scale_dims_rowwise[0];
const size_t unpadded_blocks_X_rowwise = scale_dims_rowwise[1];
const size_t blocks_Y_rowwise = scale_dims_rowwise[2];
const size_t blocks_X_rowwise = scale_dims_rowwise[3];
const size_t scales_stride_rowwise = blocks_X_rowwise;

const size_t unpadded_blocks_Y_colwise = scale_dims_colwise[0];
const size_t unpadded_blocks_X_colwise = scale_dims_colwise[1];
const size_t blocks_Y_colwise = scale_dims_colwise[2];
const size_t blocks_X_colwise = scale_dims_colwise[3];
const size_t scales_stride_colwise = blocks_X_colwise;

Tensor input("input", shape, itype);
Tensor output("output", shape, otype, rowwise, colwise, NVTE_MXFP8_1D_SCALING);

std::unique_ptr<OutputType[]> ref_output_rowwise = std::make_unique<OutputType[]>(rows * cols);
std::unique_ptr<OutputType[]> ref_output_colwise = std::make_unique<OutputType[]>(rows * cols);
std::unique_ptr<fp8e8m0[]> ref_scales_rowwise =
std::make_unique<fp8e8m0[]>(blocks_Y_rowwise * blocks_X_rowwise);
std::unique_ptr<fp8e8m0[]> ref_scales_colwise =
std::make_unique<fp8e8m0[]>(blocks_Y_colwise * blocks_X_colwise);
std::fill_n(ref_scales_rowwise.get(), blocks_Y_rowwise * blocks_X_rowwise, 0);
std::fill_n(ref_scales_colwise.get(), blocks_Y_colwise * blocks_X_colwise, 0);

fillCase<EncodingType>(&input, fill_case);

QuantizationConfigWrapper quant_config;
quant_config.set_mxfp8_2d_quantization(true);
nvte_quantize_v2(input.data(), output.data(), quant_config, 0);

cudaDeviceSynchronize();
auto err = cudaGetLastError();
ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err);

compute_ref_2d_quantize<InputType, OutputType>(
rowwise,
colwise,
input.rowwise_cpu_dptr<InputType>(),
ref_output_rowwise.get(),
ref_output_colwise.get(),
ref_scales_rowwise.get(),
ref_scales_colwise.get(),
rows,
cols,
scales_stride_rowwise,
scales_stride_colwise);

const size_t scale_diff_abs_tolerance = 0;
const double abs_tolerable_mismatches_limit = 0.0;
const double rel_tolerable_mismatches_limit = 0.0;

auto [atol, rtol] = getTolerances(otype);

if (rowwise) {
size_t mismatches_scales_rowwise = 0;
compare_scaling_factors("scales_rowwise", output.rowwise_cpu_scale_inv_ptr<fp8e8m0>(),
ref_scales_rowwise.get(), unpadded_blocks_Y_rowwise,
unpadded_blocks_X_rowwise, scales_stride_rowwise,
mismatches_scales_rowwise,
scale_diff_abs_tolerance,
abs_tolerable_mismatches_limit,
rel_tolerable_mismatches_limit);
compareResults("output_rowwise", output, ref_output_rowwise.get(), true, atol, rtol, true);
}

if (colwise) {
size_t mismatches_scales_colwise = 0;
compare_scaling_factors("scales_colwise", output.columnwise_cpu_scale_inv_ptr<fp8e8m0>(),
ref_scales_colwise.get(), unpadded_blocks_Y_colwise,
unpadded_blocks_X_colwise, scales_stride_colwise,
mismatches_scales_colwise,
scale_diff_abs_tolerance,
abs_tolerable_mismatches_limit,
rel_tolerable_mismatches_limit);
compareResults("output_colwise", output, ref_output_colwise.get(), false, atol, rtol, true);
}
}

std::vector<std::vector<size_t>> matrix_sizes = {
{1, 16},
{16, 48},
Expand All @@ -532,6 +704,18 @@ std::vector<std::vector<size_t>> matrix_sizes = {
{8192, 7168},
};

std::vector<std::vector<size_t>> matrix_sizes_2d_quantize = {
{1, 16},
{16, 48},
{65, 80},
{127, 400},
{128, 128},
{993, 512},
{8, 32, 1024},
{16, 8, 4, 512},
{8192, 7168},
};

std::vector<std::pair<size_t, size_t>> block_sizes = {
{1, 32},
{32, 1},
Expand All @@ -554,6 +738,12 @@ std::vector<ProcessingMethod> processing_methods = {
ProcessingMethod::CAST_ACT,
};

std::vector<MXFP82DScalingDirection> scaling_directions_2d_quantize = {
MXFP82DScalingDirection::RowwiseOnly,
MXFP82DScalingDirection::ColwiseOnly,
MXFP82DScalingDirection::Bidirectional,
};

// Only GeLU activation tests are supported
std::vector<ActivationType> Activation_types = {
ActivationType::Identity,
Expand All @@ -573,6 +763,13 @@ class FusedCastMXFP8TestSuite : public ::testing::TestWithParam
transformer_engine::DType,
InputsFillCase>> {};

class CastMXFP82DQuantizationTestSuite : public ::testing::TestWithParam
<std::tuple<MXFP82DScalingDirection,
std::vector<size_t>,
transformer_engine::DType,
transformer_engine::DType,
InputsFillCase>> {};

TEST_P(FusedCastMXFP8TestSuite, TestFusedCastMXFP8) {
// Skip tests for pre-Blackwell architectures
if (getDeviceComputeCapability() < blackwellComputeCapability) {
Expand Down Expand Up @@ -653,6 +850,29 @@ TEST_P(FusedCastMXFP8TestSuite, TestFusedCastMXFP8) {
}
}

TEST_P(CastMXFP82DQuantizationTestSuite, TestCastMXFP82DQuantization) {
// Skip tests for pre-Blackwell architectures
if (getDeviceComputeCapability() < blackwellComputeCapability) {
GTEST_SKIP();
}

using namespace transformer_engine;
using namespace test;

const auto scaling_direction = std::get<0>(GetParam());
const auto matrix_size = std::get<1>(GetParam());
const DType input_type = std::get<2>(GetParam());
const DType output_type = std::get<3>(GetParam());
const InputsFillCase fill_case = std::get<4>(GetParam());

TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType,
TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(output_type, OutputType,
performTest_2d_quantize<InputType, OutputType>(
matrix_size, scaling_direction, fill_case);
);
);
}

std::string to_string(const ProcessingMethod method) {
switch (method) {
case ProcessingMethod::CAST_ONLY: return "CAST_ONLY";
Expand All @@ -676,6 +896,15 @@ std::string to_string(const ActivationType Act_type) {
}
}

std::string to_string(const MXFP82DScalingDirection scaling_direction) {
switch (scaling_direction) {
case MXFP82DScalingDirection::RowwiseOnly: return "RowwiseOnly";
case MXFP82DScalingDirection::ColwiseOnly: return "ColwiseOnly";
case MXFP82DScalingDirection::Bidirectional: return "Bidirectional";
default: return "";
}
}

std::string test_name_generator(
const testing::TestParamInfo<FusedCastMXFP8TestSuite::ParamType>& info) {
std::string name = to_string(std::get<0>(info.param)) + "X" +
Expand All @@ -692,6 +921,19 @@ std::string test_name_generator(
return name;
}

std::string mxfp8_2d_quantization_test_name_generator(
const testing::TestParamInfo<CastMXFP82DQuantizationTestSuite::ParamType>& info) {
std::string name = to_string(std::get<0>(info.param));
const auto& shape = std::get<1>(info.param);
for ( const auto& s: shape) {
name += "X" + std::to_string(s);
}
name += "X" + test::typeName(std::get<2>(info.param)) +
"X" + test::typeName(std::get<3>(info.param)) +
"X" + test::caseName(std::get<4>(info.param));
return name;
}

} // namespace

// Test cases with only cast kernels
Expand All @@ -708,6 +950,18 @@ INSTANTIATE_TEST_SUITE_P(
::testing::ValuesIn(input_scenarios)),
test_name_generator);

// Test cases for MXFP8 2D block scaling through the common C++ API.
INSTANTIATE_TEST_SUITE_P(
OperatorTest_CastMXFP8_2DQuantization,
CastMXFP82DQuantizationTestSuite,
::testing::Combine(
::testing::ValuesIn(scaling_directions_2d_quantize),
::testing::ValuesIn(matrix_sizes_2d_quantize),
::testing::Values(DType::kFloat32, DType::kBFloat16),
::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2),
::testing::Values(InputsFillCase::uniform)),
mxfp8_2d_quantization_test_name_generator);

// Test cases with varying matrix shapes and block shapes
INSTANTIATE_TEST_SUITE_P(
OperatorTest_FusedCastMXFP8_Sizes,
Expand Down
Loading
Loading