Explore error correction for pertoken recipe

tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py

@@ -28,6 +28,8 @@ def check_nvfp4_gemm_versus_reference(
     x_columnwise: bool = False,
     w_columnwise: bool = False,
     row_scaled_nvfp4: bool = False,
+    err_corrected_nvfp4: bool = False,
+    use_bias: bool = False,
     use_4over6: bool = False,
     nvfp4_e4m3_max: int = 448,
     nvfp4_4over6_err_mode: str = "MAE",
@@ -47,6 +49,7 @@ def check_nvfp4_gemm_versus_reference(
     w_shape = (K, N) if w_columnwise else (N, K)
     x = torch.randn(x_shape, dtype=x_dtype, device=device)
     w = torch.randn(w_shape, dtype=w_dtype, device=device)
+    bias = torch.randn((N,), dtype=torch.bfloat16, device=device) if use_bias else None
 
     # Setup out tensor if accumulate is True
     if accumulate:
@@ -64,6 +67,7 @@ def check_nvfp4_gemm_versus_reference(
         with_rht=False,
         with_post_rht_amax=False,
         row_scaled_nvfp4=row_scaled_nvfp4,
+        err_corrected_nvfp4=err_corrected_nvfp4,
         nvfp4_use_4over6=use_4over6,
         nvfp4_e4m3_max=nvfp4_e4m3_max,
         nvfp4_4over6_err_mode=nvfp4_4over6_err_mode,
@@ -134,6 +138,7 @@ def check_nvfp4_gemm_versus_reference(
         eps=0.0,
         quant_tile_shape=(1, 16),
         row_scaled_nvfp4=row_scaled_nvfp4,
+        err_corrected_nvfp4=err_corrected_nvfp4,
         nvfp4_use_4over6=use_4over6,
         nvfp4_e4m3_max=nvfp4_e4m3_max,
         nvfp4_4over6_err_mode=nvfp4_4over6_err_mode,
@@ -166,7 +171,7 @@ def check_nvfp4_gemm_versus_reference(
         qx=qx_data,
         qw=qw_data,
         m_params=None,  # MMParams not used in reference
-        out_dtype=out_dtype,
+        out_dtype=torch.float32 if use_bias else out_dtype,
         sx=sx_trimmed,
         sw=sw_trimmed,
         bias=None,  # No bias for this test
@@ -176,6 +181,8 @@ def check_nvfp4_gemm_versus_reference(
         qresult_x=x_nvfp4_ref,
         qresult_w=w_nvfp4_ref,
     )
+    if bias is not None:
+        y_ref = (y_ref + bias.to(torch.float32)).to(out_dtype)
 
     # Native TE GEMM using tex.generic_gemm (cuBLAS GEMM)
     # Allocate cuBLAS workspace
@@ -184,7 +191,6 @@ def check_nvfp4_gemm_versus_reference(
     transa = True if not w_columnwise else False
     transb = False if not x_columnwise else True
     out_quantizer = None
-    bias = None
     bias_dtype = TE_DType[torch.bfloat16]
     use_gelu = False
     gelu_input = None
@@ -204,6 +210,7 @@ def check_nvfp4_gemm_versus_reference(
             accumulate=accumulate,
             layout=layout,
             out=out.clone() if accumulate else None,
+            bias=bias,
         )[0]
     else:
         # Native cuBLAS GEMM
@@ -489,6 +496,76 @@ def test_nvfp4_gemm_versus_reference(
     )
 
 
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize("M, K, N", [(128, 128, 128), (112, 288, 96)])
+@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
+@pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"])
+def test_error_corrected_row_scaled_gemm_versus_reference(
+    M: int,
+    K: int,
+    N: int,
+    out_dtype: torch.dtype,
+    use_bias: bool,
+) -> None:
+    """The FP32 sum of primary and residual GEMMs matches the blockwise reference."""
+    check_nvfp4_gemm_versus_reference(
+        x_dtype=torch.bfloat16,
+        w_dtype=torch.bfloat16,
+        out_dtype=out_dtype,
+        M=M,
+        K=K,
+        N=N,
+        accumulate=False,
+        row_scaled_nvfp4=True,
+        err_corrected_nvfp4=True,
+        use_bias=use_bias,
+    )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_error_corrected_row_scaled_gemm_improves_accuracy() -> None:
+    """The residual product reduces MSE relative to the primary row-scaled GEMM."""
+    torch.manual_seed(41)
+    torch.cuda.manual_seed(41)
+    device = "cuda"
+    x = torch.randn((128, 256), dtype=torch.bfloat16, device=device)
+    w = torch.randn((128, 256), dtype=torch.bfloat16, device=device)
+
+    def _activation_quantizer(*, corrected: bool) -> NVFP4Quantizer:
+        return NVFP4Quantizer(
+            fp4_dtype=te.DType.kFloat4E2M1,
+            rowwise=True,
+            columnwise=False,
+            with_amax_reduction=False,
+            amax_reduction_group=None,
+            with_rht=False,
+            with_post_rht_amax=False,
+            row_scaled_nvfp4=True,
+            err_corrected_nvfp4=corrected,
+        )
+
+    weight_quantizer = NVFP4Quantizer(
+        fp4_dtype=te.DType.kFloat4E2M1,
+        rowwise=True,
+        columnwise=True,
+        with_amax_reduction=False,
+        amax_reduction_group=None,
+        with_rht=False,
+        with_post_rht_amax=False,
+    )
+    w_q = weight_quantizer(w)
+    x_primary = _activation_quantizer(corrected=False)(x)
+    x_corrected = _activation_quantizer(corrected=True)(x)
+
+    y_primary = general_gemm(w_q, x_primary, out_dtype=torch.float32, layout="TN")[0]
+    y_corrected = general_gemm(w_q, x_corrected, out_dtype=torch.float32, layout="TN")[0]
+    y_ref = x.to(torch.float32) @ w.to(torch.float32).T
+
+    primary_mse = torch.mean((y_primary - y_ref) ** 2)
+    corrected_mse = torch.mean((y_corrected - y_ref) ** 2)
+    assert corrected_mse < primary_mse
+
+
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 @pytest.mark.parametrize(
     "m_splits, k, n",

tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py

@@ -108,6 +108,7 @@ def check_quantization_nvfp4_versus_reference(
     use_cpp_allocator: bool,
     with_2d_quantization: bool,
     row_scaled_nvfp4: bool = False,
+    err_corrected_nvfp4: bool = False,
     use_4over6: bool = False,
     nvfp4_e4m3_max: int = 448,
     nvfp4_4over6_err_mode: str = "MAE",
@@ -140,6 +141,7 @@ def check_quantization_nvfp4_versus_reference(
         with_post_rht_amax=False,
         with_2d_quantization=with_2d_quantization,
         row_scaled_nvfp4=row_scaled_nvfp4,
+        err_corrected_nvfp4=err_corrected_nvfp4,
         nvfp4_use_4over6=use_4over6,
         nvfp4_e4m3_max=nvfp4_e4m3_max,
         nvfp4_4over6_err_mode=nvfp4_4over6_err_mode,
@@ -176,6 +178,12 @@ def check_quantization_nvfp4_versus_reference(
     sx_t = x_nvfp4_sut._columnwise_scale_inv
     qx_amax = x_nvfp4_sut._amax_rowwise
     qx_amax_t = x_nvfp4_sut._amax_columnwise
+    qx_err = (
+        x_nvfp4_sut._rowwise_data_err.view(dtype=torch.uint8)
+        if x_nvfp4_sut._rowwise_data_err is not None
+        else None
+    )
+    sx_err = x_nvfp4_sut._rowwise_scale_inv_err
 
     # Reference quantization
     quant_tile_shape = (1, 16) if not with_2d_quantization else (16, 16)
@@ -187,6 +195,7 @@ def check_quantization_nvfp4_versus_reference(
         eps=0.0,
         quant_tile_shape=quant_tile_shape,
         row_scaled_nvfp4=row_scaled_nvfp4,
+        err_corrected_nvfp4=err_corrected_nvfp4,
         nvfp4_use_4over6=use_4over6,
         nvfp4_e4m3_max=nvfp4_e4m3_max,
         nvfp4_4over6_err_mode=nvfp4_4over6_err_mode,
@@ -211,6 +220,14 @@ def check_quantization_nvfp4_versus_reference(
     )
     ref_amax = x_nvfp4_ref.global_amax_row
     ref_amax_t = x_nvfp4_ref.global_amax_col
+    qx_err_ref = (
+        unpack_fp4(x_nvfp4_ref.data_err.view(dtype=torch.uint8))
+        if x_nvfp4_ref.data_err is not None
+        else None
+    )
+    sx_err_ref = (
+        x_nvfp4_ref.scale_err.view(dtype=torch.uint8) if x_nvfp4_ref.scale_err is not None else None
+    )
 
     qx = unpack_fp4(qx)
     qx_t = unpack_fp4(qx_t) if qx_t is not None else None
@@ -234,6 +251,47 @@ def check_quantization_nvfp4_versus_reference(
 
     torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0)
 
+    if err_corrected_nvfp4:
+        assert qx_err is not None and qx_err_ref is not None
+        assert sx_err is not None and sx_err_ref is not None
+        torch.testing.assert_close(unpack_fp4(qx_err), qx_err_ref, atol=0.0, rtol=0.0)
+        torch.testing.assert_close(
+            sx_err[: sx_err_ref.shape[0], : sx_err_ref.shape[1]],
+            sx_err_ref,
+            atol=0.0,
+            rtol=0.0,
+        )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize(
+    "M, N",
+    [
+        pytest.param(128, 128, id="aligned-tuned-1d"),
+        pytest.param(112, 272, id="unaligned-generic-vector"),
+    ],
+)
+@pytest.mark.parametrize(
+    "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"]
+)
+def test_error_corrected_row_scaled_quantization_versus_reference(
+    M: int,
+    N: int,
+    use_cpp_allocator: bool,
+) -> None:
+    """Primary and residual FP4 bytes/scales exactly match the BF16 reference."""
+    check_quantization_nvfp4_versus_reference(
+        x_dtype=torch.bfloat16,
+        M=M,
+        N=N,
+        return_transpose=False,
+        swizzled_scale=False,
+        use_cpp_allocator=use_cpp_allocator,
+        with_2d_quantization=False,
+        row_scaled_nvfp4=True,
+        err_corrected_nvfp4=True,
+    )
+
 
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 @pytest.mark.parametrize(

tests/pytorch/test_custom_recipe.py

@@ -1623,6 +1623,104 @@ def test_nvfp4_recipe_state_role_dispatch_backward():
         assert q.stochastic_rounding == nvfp4_recipe.fp4_quant_bwd_grad.stochastic_rounding
 
 
+def test_nvfp4_error_correction_is_forward_activation_only():
+    """Recipe opt-in is attached to forward activations, never weights or backward tensors."""
+    if not torch.cuda.is_available() or not te.is_nvfp4_available():
+        pytest.skip("NVFP4 unsupported on this device")
+
+    from transformer_engine.pytorch.quantization import NVFP4BlockScalingRecipeState
+
+    nvfp4_recipe = recipe.NVFP4BlockScaling(
+        disable_rht=True,
+        row_scaled_activation=True,
+        err_corrected_activation=True,
+    )
+    forward = NVFP4BlockScalingRecipeState(
+        nvfp4_recipe,
+        mode="forward",
+        num_quantizers=3,
+        roles=[
+            _nvfp4_role("input"),
+            _nvfp4_role("weight"),
+            _nvfp4_role("output"),
+        ],
+    ).make_quantizers()
+    assert [q.err_corrected_nvfp4 for q in forward] == [True, False, True]
+    assert [q.row_scaled_nvfp4 for q in forward] == [True, False, True]
+
+    backward = NVFP4BlockScalingRecipeState(
+        nvfp4_recipe,
+        mode="backward",
+        num_quantizers=2,
+        roles=[_nvfp4_role("grad_output"), _nvfp4_role("grad_input")],
+    ).make_quantizers()
+    assert all(not q.err_corrected_nvfp4 for q in backward)
+    assert all(not q.row_scaled_nvfp4 for q in backward)
+
+
+def test_nvfp4_error_correction_requires_row_scaled_activation():
+    with pytest.raises(ValueError, match="row_scaled_activation=True"):
+        recipe.NVFP4BlockScaling(
+            row_scaled_activation=False,
+            err_corrected_activation=True,
+        )
+    with pytest.raises(ValueError, match="does not support activation 4over6"):
+        recipe.NVFP4BlockScaling(
+            disable_rht=True,
+            row_scaled_activation=True,
+            err_corrected_activation=True,
+            nvfp4_4over6="activations",
+        )
+
+
+@pytest.mark.parametrize("bias", [False, True])
+def test_nvfp4_error_correction_is_batch_invariant(bias):
+    """Shared rows produce identical outputs when the surrounding batch changes."""
+    if not torch.cuda.is_available() or not te.is_nvfp4_available():
+        pytest.skip("NVFP4 unsupported on this device")
+
+    torch.manual_seed(0)
+    batch_sizes = (1, 2, 4)
+    sequence_length = 16
+    in_features = 128
+    out_features = 128
+    inp = torch.randn(
+        max(batch_sizes),
+        sequence_length,
+        in_features,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    model = Linear(
+        in_features,
+        out_features,
+        bias=bias,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+    ).eval()
+    nvfp4_recipe = recipe.NVFP4BlockScaling(
+        disable_rht=True,
+        row_scaled_activation=True,
+        err_corrected_activation=True,
+    )
+
+    outputs = {}
+    with torch.inference_mode():
+        for batch_size in batch_sizes:
+            with autocast(enabled=True, recipe=nvfp4_recipe):
+                outputs[batch_size] = model(inp[:batch_size]).clone()
+
+    reference = outputs[max(batch_sizes)]
+    for batch_size in batch_sizes:
+        assert torch.isfinite(outputs[batch_size]).all()
+        torch.testing.assert_close(
+            outputs[batch_size],
+            reference[:batch_size],
+            rtol=0,
+            atol=0,
+        )
+
+
 def test_nvfp4_recipe_state_positional_fallback_matches_explicit_roles():
     """``roles=None`` matches explicit ``[input, weight, output]`` slot-for-slot."""
     if not torch.cuda.is_available() or not te.is_nvfp4_available():

transformer_engine/common/cast/dispatch/quantize.cuh

@@ -103,6 +103,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
       const auto [rows, cols] = input_tensor->flat_2d_dims();
       auto dtype = input_tensor->dtype();
       const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4;
+      const bool err_corrected_nvfp4 = output_tensor->err_corrected_nvfp4;
       const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled;
       NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448,
                  "Non-4over6 NVFP4 quantization requires E4M3 max 448.");
@@ -115,6 +116,10 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
                    "Row-scaled NVFP4 quantization does not produce columnwise output.");
         nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream);
       }
+      NVTE_CHECK(!err_corrected_nvfp4 || row_scaled_nvfp4,
+                 "Error-corrected NVFP4 quantization requires row scaling.");
+      NVTE_CHECK(!err_corrected_nvfp4 || !nvfp4_use_4over6,
+                 "Error-corrected NVFP4 quantization does not support 4over6.");
       // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and
       // both-directions keep their existing routing. Columnwise-only 1D and non-bf16 fall back to
       // quantize_transpose_vector_blockwise_fp4.
@@ -155,6 +160,9 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
             /*rng_state=*/quant_config_cpp.rng_state,
             /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization,
             /*row_scaled_nvfp4=*/row_scaled_nvfp4,
+            /*err_corrected_nvfp4=*/err_corrected_nvfp4,
+            /*scale_inv_err=*/output_tensor->scale_inv_err,
+            /*output_err=*/output_tensor->data_err,
             /*noop_tensor=*/noop_tensor->data,
             /*stream=*/stream);
       }
@@ -315,6 +323,9 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
             /*rng_state=*/quant_config_cpp.rng_state,
             /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization,
             /*row_scaled_nvfp4=*/false,
+            /*err_corrected_nvfp4=*/false,
+            /*scale_inv_err=*/output_tensor->scale_inv_err,
+            /*output_err=*/output_tensor->data_err,
             /*noop_tensor=*/noop_tensor->data,
             /*stream=*/stream);
       }