Skip to content

Prototype NVFP4 with FP8 UE5M3 block scales #3325

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

Draft
timmoon10 wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Prototype NVFP4 with FP8 UE5M3 block scales #3325

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
25d5b4b
Prototype NVFP4 with UE5M3 scales
timmoon10 Aug 7, 2026
dede7c7
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 7, 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
485 changes: 362 additions & 123 deletions tests/cpp/operator/test_cast_nvfp4_transpose.cu
View file
Open in desktop

Large diffs are not rendered by default.

137 changes: 124 additions & 13 deletions tests/cpp/operator/test_dequantize_nvfp4.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#endif

#include <transformer_engine/cast.h>
#include <transformer_engine/recipe.h>
#include <transformer_engine/swizzle.h>
#include "../test_common.h"
#include "transformer_engine/transformer_engine.h"
Expand All @@ -39,23 +40,23 @@ float2 cvt_fp4x2_to_float2(fp4e2m1x2 fp4_pair) {
return {static_cast<float>(h2.x), static_cast<float>(h2.y)};
}

template <typename OType>
template <typename OType, typename ScaleType>
void compute_ref_dequantize_nvfp4(const uint8_t *packed_data,
const fp8e4m3 *scales,
const ScaleType *scales,
const std::vector<float> &amax,
OType *output,
size_t rows,
size_t cols,
size_t scale_stride,
int e4m3_max) {
const float factor_inv = 1.0f / (6.0f * static_cast<float>(e4m3_max));
float scale_max) {
const float factor_inv = 1.0f / (6.0f * scale_max);
constexpr size_t BLOCK_SIZE = 16;
const size_t Mread = cols / BLOCK_SIZE;
const size_t bytes_per_block = BLOCK_SIZE / 2;

for (size_t row = 0; row < rows; ++row) {
for (size_t block = 0; block < Mread; ++block) {
const fp8e4m3 scale = scales[row * scale_stride + block];
const ScaleType scale = scales[row * scale_stride + block];
const float final_scale =
static_cast<float>(scale) * (amax.size() == 1 ? amax[0] : amax[row]) * factor_inv;

Expand Down Expand Up @@ -94,7 +95,7 @@ struct NVFP4DequantizeTestConfig {

// Quantize a high-precision input to NVFP4, then dequantize and compare
// against a CPU reference computed from the quantized data.
template <typename OutputType>
template <typename OutputType, typename ScaleType = fp8e4m3>
void performTest_dequantize_nvfp4(const size_t rows, const size_t cols,
const bool row_scaled_nvfp4,
const NVTENVFP44Over6Mode mode,
Expand All @@ -105,7 +106,8 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols,
// Tensors
Tensor input("input", std::vector<size_t>{rows, cols}, otype);
Tensor quantized("quantized", std::vector<size_t>{rows, cols},
DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING);
DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING,
TypeInfo<ScaleType>::dtype);
Tensor output("output", std::vector<size_t>{rows, cols}, otype, true, false);

// Fill input with random data
Expand Down Expand Up @@ -149,24 +151,25 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols,
quantized.to_cpu();
const uint8_t *fp4_data =
reinterpret_cast<const uint8_t *>(quantized.rowwise_cpu_dptr<fp4e2m1>());
const fp8e4m3 *scales = quantized.rowwise_cpu_scale_inv_ptr<fp8e4m3>();
const ScaleType *scales = quantized.rowwise_cpu_scale_inv_ptr<ScaleType>();
const auto *amax = quantized.cpu_rowwise_amax_ptr<float>();
const std::vector<float> amax_vals(amax, amax + amax_size);
const NVTEShape scale_shape = quantized.rowwise_scale_inv_shape();
const size_t scale_stride = scale_shape.data[scale_shape.ndim - 1];
std::unique_ptr<OutputType[]> ref_output =
std::make_unique<OutputType[]>(rows * cols);
compute_ref_dequantize_nvfp4<OutputType>(
const float scale_max = static_cast<float>(e4m3_max);
compute_ref_dequantize_nvfp4<OutputType, ScaleType>(
fp4_data, scales, amax_vals, ref_output.get(),
rows, cols, scale_stride, e4m3_max);
rows, cols, scale_stride, scale_max);

// Compare results from TE and reference impls
auto [atol, rtol] = getTolerances(otype);
compareResults("output_nvfp4", output, ref_output.get(), true, atol, rtol);
}

// Dequantize NVFP4 with GEMM-swizzled scales and compare against compact path.
template <typename OutputType>
template <typename OutputType, typename ScaleType = fp8e4m3>
void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols,
const bool row_scaled_nvfp4,
const NVTENVFP44Over6Mode mode,
Expand All @@ -178,7 +181,8 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols,
fillCase<fp32>(&input, InputsFillCase::uniform);

Tensor quantized_compact("quantized_compact", std::vector<size_t>{rows, cols},
DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING);
DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING,
TypeInfo<ScaleType>::dtype);
quantized_compact.set_nvfp4_e4m3_max(e4m3_max);
ASSERT_EQ(quantized_compact.nvfp4_e4m3_max(), e4m3_max);
if (row_scaled_nvfp4) {
Expand All @@ -203,7 +207,8 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols,

// Create tensor with same FP4 data but swizzled scales
Tensor quantized_swizzled("quantized_swizzled", std::vector<size_t>{rows, cols},
DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING);
DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING,
TypeInfo<ScaleType>::dtype);
quantized_swizzled.set_nvfp4_e4m3_max(e4m3_max);
ASSERT_EQ(quantized_swizzled.nvfp4_e4m3_max(), e4m3_max);
if (row_scaled_nvfp4) {
Expand Down Expand Up @@ -325,6 +330,112 @@ INSTANTIATE_TEST_SUITE_P(
}
);

#if CUDA_VERSION >= 13040
TEST(DequantizeNVFP4Test, UE5M3Scales)
{
if (getDeviceComputeCapability() < blackwellComputeCapability) {
GTEST_SKIP();
}

performTest_dequantize_nvfp4<fp32, fp8ue5m3>(
32, 64, false, kNVTENVFP44Over6Disabled, 114688);
performTest_dequantize_nvfp4<bf16, fp8ue5m3>(
32, 64, true, kNVTENVFP44Over6Disabled, 114688);
performTest_dequantize_nvfp4_swizzled<fp32, fp8ue5m3>(
32, 64, false, kNVTENVFP44Over6Disabled, 114688);
performTest_dequantize_nvfp4_swizzled<bf16, fp8ue5m3>(
32, 64, true, kNVTENVFP44Over6Disabled, 114688);
performTest_dequantize_nvfp4<fp32, fp8ue5m3>(
32, 64, false, kNVTENVFP44Over6MinMAE, 65536);
performTest_dequantize_nvfp4_swizzled<bf16, fp8ue5m3>(
32, 64, true, kNVTENVFP44Over6MinMAE, 65536);
}

TEST(NVFP4RecipeTest, UE5M3ScaleUtilities)
{
if (getDeviceComputeCapability() < blackwellComputeCapability) {
GTEST_SKIP();
}

Tensor global_amax("global_amax", std::vector<size_t>{1}, DType::kFloat32);
Tensor global_scale("global_scale", std::vector<size_t>{1}, DType::kFloat32);
global_amax.rowwise_cpu_dptr<float>()[0] = 12.0f;
global_amax.from_cpu();
nvte_nvfp4_compute_global_scale(
global_amax.data(), global_scale.data(), 0, kNVTEFloat8UE5M3);
global_scale.to_cpu();
EXPECT_FLOAT_EQ(global_scale.rowwise_cpu_dptr<float>()[0], 6.0f * 114688.0f / 12.0f);

Tensor block_amax("block_amax", std::vector<size_t>{1, 2}, DType::kFloat32);
Tensor block_scale("block_scale", std::vector<size_t>{1, 2}, DType::kFloat32);
block_amax.rowwise_cpu_dptr<float>()[0] = 3.0f;
block_amax.rowwise_cpu_dptr<float>()[1] = 6.0f;
block_amax.from_cpu();
nvte_nvfp4_compute_per_block_scale(
block_amax.data(), block_scale.data(), global_amax.data(), 0, kNVTEFloat8UE5M3);
block_scale.to_cpu();
EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr<float>()[0], 3.0f * 114688.0f / 12.0f);
EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr<float>()[1], 6.0f * 114688.0f / 12.0f);

Tensor expanded_scale("expanded_scale", std::vector<size_t>{16, 2}, DType::kByte);
nvte_nvfp4_expand_scale_to_fp8(
block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, 0, kNVTEFloat8UE5M3);
expanded_scale.to_cpu();
const auto *scales = reinterpret_cast<const fp8ue5m3 *>(
expanded_scale.rowwise_cpu_dptr<byte>());
for (size_t row = 0; row < 16; ++row) {
EXPECT_FLOAT_EQ(static_cast<float>(scales[row * 2]),
static_cast<float>(fp8ue5m3(3.0f * 114688.0f / 12.0f)));
EXPECT_FLOAT_EQ(static_cast<float>(scales[row * 2 + 1]),
static_cast<float>(fp8ue5m3(6.0f * 114688.0f / 12.0f)));
}
}

TEST(NVFP4RecipeTest, UE5M3PerTensorScale)
{
if (getDeviceComputeCapability() < blackwellComputeCapability) {
GTEST_SKIP();
}

Tensor input_a("input_a", std::vector<size_t>{32, 32}, DType::kFloat4E2M1,
true, true, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3);
Tensor input_b("input_b", std::vector<size_t>{32, 32}, DType::kFloat4E2M1,
true, true, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3);
Tensor alpha_out("alpha_out", std::vector<size_t>{1}, DType::kFloat32);

constexpr float amax_a = 12.0f;
constexpr float amax_b = 18.0f;
constexpr float alpha_in = 2.0f;
constexpr float fp4_max = 6.0f;
constexpr float ue5m3_max = 114688.0f;
input_a.set_nvfp4_e4m3_max(static_cast<int>(ue5m3_max));
input_b.set_nvfp4_e4m3_max(static_cast<int>(ue5m3_max));
input_a.set_amax(amax_a);
input_b.set_tensor_amax_columnwise(amax_b);

nvte_nvfp4_compute_per_tensor_scale(
input_a.data(), true, input_b.data(), false, alpha_in, alpha_out.data(), 0);
alpha_out.to_cpu();

const float factor_inv =
1.0f / (fp4_max * fp4_max * ue5m3_max * ue5m3_max);
const float expected = alpha_in * amax_a * amax_b * factor_inv;
EXPECT_FLOAT_EQ(alpha_out.rowwise_cpu_dptr<float>()[0], expected);

input_a.set_nvfp4_e4m3_max(65536);
input_b.set_nvfp4_e4m3_max(65536);
nvte_nvfp4_compute_per_tensor_scale(
input_a.data(), true, input_b.data(), false, alpha_in, alpha_out.data(), 0);
alpha_out.to_cpu();

constexpr float ue5m3_headroom_max = 65536.0f;
const float headroom_factor_inv =
1.0f / (fp4_max * fp4_max * ue5m3_headroom_max * ue5m3_headroom_max);
const float headroom_expected = alpha_in * amax_a * amax_b * headroom_factor_inv;
EXPECT_FLOAT_EQ(alpha_out.rowwise_cpu_dptr<float>()[0], headroom_expected);
}
#endif

class DequantizeNVFP4SwizzledTestSuite : public ::testing::TestWithParam
<std::tuple<std::pair<size_t, size_t>,
transformer_engine::DType,
Expand Down
14 changes: 13 additions & 1 deletion tests/cpp/test_common.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,9 @@ bool areShapesEqual(const NVTEShape &s1, const NVTEShape &s2) {
}

size_t typeToNumBits(DType type) {
if (type == DType::kFloat8UE5M3) {
return 8;
}
TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(type, T,
{
return TypeInfo<T>::size;
Expand All @@ -65,6 +68,7 @@ const std::string &typeName(DType type) {
{DType::kBFloat16, "bfloat16"},
{DType::kFloat8E4M3, "float8e4m3"},
{DType::kFloat8E5M2, "float8e5m2"},
{DType::kFloat8UE5M3, "float8ue5m3"},
{DType::kFloat8E8M0, "float8e8m0"},
{DType::kFloat4E2M1, "float4e2m1"}};
return name_map.at(type);
Expand Down Expand Up @@ -278,7 +282,7 @@ void Tensor::Buffer::from_cpu() {
Tensor::Tensor(const std::string& name,
const NVTEShape &shape, const DType type,
const bool rowwise, const bool columnwise,
const NVTEScalingMode &scaling_mode)
const NVTEScalingMode &scaling_mode, const DType scale_dtype)
: tensor_(scaling_mode), rowwise_{rowwise}, columnwise_{columnwise}, name_{name} {
// Initialize RNG
const size_t seed = create_seed_from_tensor_name(name);
Expand Down Expand Up @@ -374,6 +378,14 @@ Tensor::Tensor(const std::string& name,
{
// Block scaling factors
auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode());
if (scaling_mode == NVTE_NVFP4_1D_SCALING) {
NVTE_CHECK(scale_dtype == DType::kFloat8E4M3 ||
scale_dtype == DType::kFloat8UE5M3);
rowwise_scale_meta.type = scale_dtype;
rowwise_scale_meta.type_size_bits = typeToNumBits(scale_dtype);
colwise_scale_meta.type = scale_dtype;
colwise_scale_meta.type_size_bits = typeToNumBits(scale_dtype);
}
if (rowwise) {
const auto scale_shape = rowwise_scale_meta.shape;
const auto scale_dtype = rowwise_scale_meta.type;
Expand Down
19 changes: 15 additions & 4 deletions tests/cpp/test_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ using bf16 = nv_bfloat16;
using fp8e4m3 = __nv_fp8_e4m3;
using fp8e5m2 = __nv_fp8_e5m2;
using fp8e8m0 = uint8_t;
#if CUDA_VERSION >= 13040
using fp8ue5m3 = __nv_fp8_ue5m3;
#endif
#if FP4_TYPE_SUPPORTED
using fp4e2m1 = __nv_fp4_e2m1;
using fp4e2m1x2 = __nv_fp4x2_e2m1;
Expand All @@ -91,7 +94,12 @@ struct BitsNumber {
template <typename T>
struct TypeInfo {
#if FP4_TYPE_SUPPORTED
using types = std::tuple<byte, int16, int32, int64, fp32, fp16, bf16, fp8e4m3, fp8e5m2, fp8e8m0, fp4e2m1>;
using types = std::tuple<byte, int16, int32, int64, fp32, fp16, bf16, fp8e4m3,
fp8e5m2, fp8e8m0, fp4e2m1
#if CUDA_VERSION >= 13040
, fp8ue5m3
#endif
>;
#else
using types = std::tuple<byte, int16, int32, int64, fp32, fp16, bf16, fp8e4m3, fp8e5m2, fp8e8m0>;
#endif
Expand Down Expand Up @@ -151,15 +159,18 @@ class Tensor {
const NVTEShape &shape, const DType type,
const bool rowwise = true,
const bool columnwise = false,
const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING);
const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING,
const DType scale_dtype = DType::kFloat8E4M3);

Tensor(const std::string& name,
const std::vector<size_t> &shape,
const DType type,
const bool rowwise = true,
const bool columnwise = false,
const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING) :
Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode) {}
const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING,
const DType scale_dtype = DType::kFloat8E4M3) :
Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode,
scale_dtype) {}

Tensor() = default;

Expand Down
49 changes: 49 additions & 0 deletions tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,55 @@
recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True)


@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
@pytest.mark.parametrize(
"disable_x, disable_w",
[(True, False), (False, True), (True, True)],
ids=["x_unit_global_scale", "w_unit_global_scale", "both_unit_global_scale"],
)
def test_gemm_with_missing_nvfp4_amax(disable_x: bool, disable_w: bool) -> None:
"""A null amax contributes a unit global scale to GEMM alpha."""
torch.manual_seed(0)
x = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda")
w = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda")
unit_scale_amax = 448.0 * 6.0
x[0, 0] = unit_scale_amax
w[0, 0] = unit_scale_amax

def quantize(tensor: torch.Tensor, disable_second_level_scale: bool):
return NVFP4Quantizer(
rowwise=True,
columnwise=True,
disable_second_level_scale=disable_second_level_scale,
)(tensor)

x_ref, w_ref = quantize(x, False), quantize(w, False)
x_test, w_test = quantize(x, disable_x), quantize(w, disable_w)

def gemm(w_q, x_q):
workspace = torch.empty(4, dtype=torch.uint8, device="cuda")
return tex.generic_gemm(
w_q,
True,
x_q,
False,
None,
None,
TE_DType[torch.bfloat16],
None,
TE_DType[torch.bfloat16],
False,
None,
False,
workspace,
workspace.numel(),
False,
False,
)[0]

torch.testing.assert_close(gemm(w_test, x_test), gemm(w_ref, x_ref), atol=0, rtol=0)


def check_nvfp4_gemm_versus_reference(
x_dtype: torch.dtype,
w_dtype: torch.dtype,
Expand Down
Loading
Loading