diff --git a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py index b14eeb815b..2ed496c730 100755 --- a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py @@ -431,3 +431,39 @@ def test_group_stochastic_rounding_quantization_versus_reference( num_splits=num_splits, use_tex_split_quantize=use_tex_split_quantize, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("use_2D", [False, True], ids=str) +def test_stochastic_rounding_picks_an_adjacent_fp4_value( + x_dtype: torch.dtype, use_2D: bool +) -> None: + device = "cuda" + torch.manual_seed(seed) + M, N = 512, 1024 + x = torch.randn((M, N), dtype=x_dtype, device=device) + amax = torch.max(torch.abs(x)).float() + + q, s, q_t, s_t = quantize_fp4(x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False) + q_redraw, _, q_t_redraw, _ = quantize_fp4( + x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False + ) + assert not torch.equal(q, q_redraw), "two stochastic rounding draws are identical" + assert not torch.equal(q_t, q_t_redraw), "two stochastic rounding draws are identical" + + def check_adjacent(qx: torch.Tensor, sx: torch.Tensor, reference: torch.Tensor) -> None: + # The kernel rounds reference / block_scale onto the E2M1 grid, so every + # output has to be one of the two grid values that bracket it, i.e. within + # one grid step. An all-zero output fails this everywhere the input is not + # already near zero. + block_scale = sx.repeat_interleave(16, dim=1).view(torch.float8_e4m3fn).to(torch.float32) + block_scale = block_scale[: reference.shape[0], : reference.shape[1]] * (amax / (6.0 * 448)) + exact = (reference.float() / block_scale).clamp(-6.0, 6.0) + magnitude = exact.abs() + step = torch.where(magnitude >= 4.0, 2.0, torch.where(magnitude >= 2.0, 1.0, 0.5)) + gap = (fp4_to_fp32(unpack_fp4(qx)) - exact).abs() + assert torch.all(gap <= step * 1.0001), f"largest gap to the exact value is {gap.max()}" + + check_adjacent(q, s, x) + check_adjacent(q_t, s_t, x.t().contiguous()) diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index d5f2fa9a2c..7605798758 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -269,11 +269,11 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_stochastic_ro : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x), "r"(rbits)); return *reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt.rs PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - uint16_t dummy = 0; - return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy); + const float q0 = ptx::stochastic_round_fp4_e2m1(in01.x, rbits); + const float q1 = ptx::stochastic_round_fp4_e2m1(in01.y, rbits >> 8); + const float q2 = ptx::stochastic_round_fp4_e2m1(in23.x, rbits >> 16); + const float q3 = ptx::stochastic_round_fp4_e2m1(in23.y, rbits >> 24); + return __nv_fp4x4_e2m1(make_float4(q0, q1, q2, q3)); } } diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 2814aa3490..1e13170b90 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -600,6 +600,35 @@ __device__ __forceinline__ void mul_cvt_4x(fp4e2m1x4 &out, const Tx2 &in01, cons out = fp4e2m1x4(make_float4(x0, x1, x2, x3)); } +// Software stochastic rounding onto the e2m1 grid, for architectures without +// cvt.rs. Takes 8 random bits per element, which is what +// cvt.rs.satfinite.e2m1x4.f32 takes from its rbits operand. Follows satfinite +// for the edges: NaN becomes positive MAX_NORM and larger magnitudes clamp to +// MAX_NORM with the sign kept. The result is exactly representable in e2m1, so +// packing it is lossless. +__device__ __forceinline__ float stochastic_round_fp4_e2m1(const float x, const uint32_t rbits8) { + constexpr float max_norm = 6.0f; + if (isnan(x)) { + return max_norm; + } + const float u = static_cast(rbits8 & 0xFFu) * (1.0f / 256.0f); + const float a = fabsf(x); + // Grid step at |x|: 0.5 below 2, 1 in [2, 4), 2 in [4, 6]. + const float step = (a >= 4.0f) ? 2.0f : ((a >= 2.0f) ? 1.0f : 0.5f); + const float t = fmaf(u, step, a); + float q; + if (t >= max_norm) { + q = max_norm; + } else if (t >= 4.0f) { + q = 4.0f; + } else if (t >= 2.0f) { + q = 2.0f + floorf(t - 2.0f); + } else { + q = floorf(t * 2.0f) * 0.5f; + } + return copysignf(q, x); +} + __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding( const uint64_t in_4x, const float2 scale, const uint32_t rbits) { uint16_t out_4x = 0; @@ -633,9 +662,14 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_roun : "=h"(out_4x) : "l"(in_4x), "l"(reinterpret_cast(scale)), "r"(rbits)); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + // mul.f32x2 above applies scale.x to the even elements and scale.y to the odd ones. + const bf16 *vals = reinterpret_cast(&in_4x); + const float q0 = stochastic_round_fp4_e2m1(static_cast(vals[0]) * scale.x, rbits); + const float q1 = stochastic_round_fp4_e2m1(static_cast(vals[1]) * scale.y, rbits >> 8); + const float q2 = stochastic_round_fp4_e2m1(static_cast(vals[2]) * scale.x, rbits >> 16); + const float q3 = stochastic_round_fp4_e2m1(static_cast(vals[3]) * scale.y, rbits >> 24); + const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3)); + out_4x = *reinterpret_cast(&packed); } return *reinterpret_cast(&out_4x); } @@ -725,9 +759,12 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_stochastic_roun "l"(reinterpret_cast(in23)), "l"(reinterpret_cast(scale)), "r"(rbits)); } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + const float q0 = stochastic_round_fp4_e2m1(in01.x * scale.x, rbits); + const float q1 = stochastic_round_fp4_e2m1(in01.y * scale.y, rbits >> 8); + const float q2 = stochastic_round_fp4_e2m1(in23.x * scale.x, rbits >> 16); + const float q3 = stochastic_round_fp4_e2m1(in23.y * scale.y, rbits >> 24); + const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3)); + out_4x = *reinterpret_cast(&packed); } return *reinterpret_cast(&out_4x); } @@ -950,9 +987,26 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding( NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); } } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + constexpr bool known_coeff = std::is_same::value || + std::is_same::value; + if constexpr (known_coeff) { + const float coeff = static_cast(scaling_coefficient); + const bf16 *vals03 = reinterpret_cast(&in03); + const bf16 *vals47 = reinterpret_cast(&in47); + float q[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + q[i] = stochastic_round_fp4_e2m1(static_cast(vals03[i]) * coeff, rbits03 >> (8 * i)); + q[i + 4] = + stochastic_round_fp4_e2m1(static_cast(vals47[i]) * coeff, rbits47 >> (8 * i)); + } + const fp4e2m1x4 lo(make_float4(q[0], q[1], q[2], q[3])); + const fp4e2m1x4 hi(make_float4(q[4], q[5], q[6], q[7])); + out_8x = static_cast(*reinterpret_cast(&lo)) | + (static_cast(*reinterpret_cast(&hi)) << 16); + } else { + NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); + } } return out_8x; }