Fixes accuracy error for mxfp8 linear

[syed-ahmed]

I'm seeing the following error on B200 and GB200.

        y_ref.backward(g)
        y.backward(g)
        w_g_ref = m_mx[0].weight.grad
        w_g = getattr(m_mx_c, "0").weight.grad
        # TODO(future): investigate why we can't match with rtol=0 atol=0
        # after moving to torchao repo. Technically compile does not give
        # bit exactness guarantees, but there also might be a bug lurking
        # around.
>       torch.testing.assert_close(w_g_ref, w_g, atol=0.02, rtol=0.02)
E       AssertionError: Tensor-likes are not close!
E       
E       Mismatched elements: 493 / 131072 (0.4%)
E       Greatest absolute difference: 10.5625 at index (23, 0) (up to 0.02 allowed)
E       Greatest relative difference: 14272.0 at index (92, 0) (up to 0.02 allowed)

test/prototype/mx_formats/test_mx_linear.py:353: AssertionError

for the following unit test:

pytest test/prototype/mx_formats/test_mx_linear.py -k test_linear_compile[ScaleCalculationMode.FLOOR-MXFP8Dim1CastKernelChoice.TRITON-False-mxfp8_cublas-hp_dtype1] -vvv

Excluding 512 from COL_TILE_SIZE suppresses this issue. I figured it out by selectively including the values for ROW_TILE_SIZE, COL_TILE_SIZE, num_warps, and num_stages, until I found the config that results in the error. I'm not sure if COL_TILE_SIZE of 512 is not supported in B200/GB200 or if it's a bug in Triton.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3357

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 3c7a8e8 with merge base 01374eb ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[syed-ahmed]

@pytorchbot label "topic: not user facing"

[syed-ahmed]

@pytorchbot merge

[pytorchmergebot]

Merge failed

Reason: 1 mandatory check(s) are pending/not yet run. The first few are:

• Facebook CLA Check

Dig deeper by viewing the pending checks on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: superuser