Fix MX + vllm (#2458)

stack-info: PR: #2458, branch: drisspg/stack/82

Author: drisspg

test/integration/test_vllm.py

@@ -41,6 +41,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
 from vllm import LLM, SamplingParams
 
+from torchao.prototype.mx_formats import MXFPInferenceConfig
 from torchao.quantization.granularity import PerRow, PerTensor
 from torchao.quantization.quant_api import (
     CutlassInt4PackedLayout,
@@ -69,9 +70,7 @@ def get_tests() -> List[TorchAoConfig]:
             Int8DynamicActivationInt4WeightConfig(layout=CutlassInt4PackedLayout())
         )
     ]
-    SM100_TESTS = [
-        # TorchAoConfig(MXFPInferenceConfig())
-    ]  # Failing for : https://github.com/pytorch/ao/issues/2239
+    SM100_TESTS = [TorchAoConfig(MXFPInferenceConfig())]
 
     # Check CUDA availability first
     if not torch.cuda.is_available():

torchao/prototype/mx_formats/mx_ops.py

@@ -240,8 +240,8 @@ def mx_slice(func, types, args, kwargs):
 
     if dim == 0:
         # Slicing along the first dimension (rows) TODO assuming that dim 1 is reduciton dim for now
-        sliced_scale = aten.slice.Tensor(scale_shaped, dim, start, end, step).flatten()
-        sliced_data = aten.slice.Tensor(x._data, dim, start, end, step)
+        sliced_scale = aten.slice.Tensor(scale_shaped, dim, start, end, step)
+        sliced_data = aten.slice.Tensor(x._data, dim, start, end, step).unsqueeze(-1)
     elif dim == 1:
         # Slicing along reduciton dim
         if start is not None:
@@ -265,7 +265,7 @@ def mx_slice(func, types, args, kwargs):
         # Slice the scale tensor accordingly
         sliced_scale = aten.slice.Tensor(
             scale_shaped, 1, start_block, end_block, step
-        ).flatten()
+        ).unsqueeze(-1)
     else:
         raise ValueError(
             f"MXTensor only supports slicing along dimensions 0 and 1, got dim={dim}"