updated test with rebase changes

Author: avizon-aws

…/prototype/tests/test_mxfp8_allgather.py …otype/mx_formats/test_mxfp8_allgather.pytorchao/prototype/tests/test_mxfp8_allgather.py renamed to test/prototype/mx_formats/test_mxfp8_allgather.py torchao/prototype/tests/test_mxfp8_allgather.py renamed to test/prototype/mx_formats/test_mxfp8_allgather.py

@@ -1,9 +1,4 @@
-import pytest
 import torch
-
-if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0):
-    pytest.skip("Test requires CUDA build on SM100", allow_module_level=True)
-
 import torch.distributed as dist
 from torch.testing._internal.common_distributed import (
     MultiProcessTestCase,
@@ -23,7 +18,7 @@ def setUp(self) -> None:
 
     @property
     def world_size(self) -> int:
-        return 4
+        return 2
 
     @property
     def device(self) -> torch.device:
@@ -64,9 +59,9 @@ def test_allgather(self):
                 elem_dtype=torch.float8_e5m2,
                 block_size=32,
                 orig_dtype=torch.float32,
-                gemm_kernel_choice=None,
-                pack_fp6=None,
+                kernel_preference=None,
                 act_quant_kwargs=None,
+                is_swizzled_scales=None,
             )
 
             world_size = self.world_size
@@ -82,9 +77,9 @@ def test_allgather(self):
                 elem_dtype=torch.float8_e5m2,
                 block_size=32,
                 orig_dtype=torch.float32,
-                gemm_kernel_choice=None,
-                pack_fp6=None,
+                kernel_preference=None,
                 act_quant_kwargs=None,
+                is_swizzled_scales=None,
             )
 
             # Perform all_gather
@@ -111,12 +106,12 @@ def test_allgather(self):
 
             # Verify scale matches golden exactly
             if not torch.equal(
-                gathered_mx._scale_e8m0.view(torch.uint8),
+                gathered_mx.scale.view(torch.uint8),
                 golden_scale.view(torch.uint8),
             ):
                 assert False, "scale mismatch"
 
-            assert gathered_mx._block_size == 32
+            assert gathered_mx.block_size == 32
 
         finally:
             dist.destroy_process_group()

torchao/prototype/mx_formats/mx_tensor.py

@@ -869,7 +869,7 @@ def mx_all_gather(func, types, args, kwargs):
     )
 
     gathered_scale = torch.ops._c10d_functional.all_gather_into_tensor.default(
-        mx_tensor._scale_e8m0.view(
+        mx_tensor.scale.view(
             torch.uint8
         ),  # The scale factors, Need to cast to uint8 as float8_e8m0fnu is not support for all gather.
         group_tag,
@@ -884,11 +884,11 @@ def mx_all_gather(func, types, args, kwargs):
         gathered_qdata,
         gathered_scale,
         mx_tensor._elem_dtype,
-        mx_tensor._block_size,
+        mx_tensor.block_size,
         mx_tensor._orig_dtype,
-        mx_tensor._gemm_kernel_choice,
-        mx_tensor._pack_fp6,
+        mx_tensor.kernel_preference,
         mx_tensor.act_quant_kwargs,
+        mx_tensor._is_swizzled_scales,
     )
 
 
@@ -908,16 +908,16 @@ def mx_wait_tensor(func, types, args, kwargs):
     )
 
     waited_scale = torch.ops._c10d_functional.wait_tensor.default(
-        mx_tensor._scale_e8m0, *args[1:], **kwargs
+        mx_tensor.scale, *args[1:], **kwargs
     )
 
     return MXTensor(
         waited_qdata,
         waited_scale,
         mx_tensor._elem_dtype,
-        mx_tensor._block_size,
+        mx_tensor.block_size,
         mx_tensor._orig_dtype,
-        mx_tensor._gemm_kernel_choice,
-        mx_tensor._pack_fp6,
+        mx_tensor.kernel_preference,
         mx_tensor.act_quant_kwargs,
+        mx_tensor._is_swizzled_scales,
     )