Add missing Float8Tensor op support for 3d weights

Summary:

This PR adds in the missing op support + tests for supporting quantizing
3d params.

In particular we add the following ops:

- aten.unsqueeze
- aten.add.Tensor
- aten.slice.Tensor has been expanded to add support for 3d weights

Test Plan:

```
pytest test/quantization/quantize_/workflows/float8/test_float8_tensor.py
```

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jcaip

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

@@ -19,6 +19,7 @@
 
 from torchao.quantization import (
     Float8DynamicActivationFloat8WeightConfig,
+    Float8Tensor,
     Float8WeightOnlyConfig,
     PerRow,
     PerTensor,

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

@@ -418,10 +418,10 @@ def _(func, types, args, kwargs):
 
 @implements(aten.slice.Tensor)
 def _(func, types, args, kwargs):
-    """Only supports slicing for dim == 1 and dim == 2
-    original tensor shape has dimension (N, K)
-    qdata has dimension (N, K)
-    scale (per row quantization) has dimension: (N,)
+    """Supports slicing for 1d, 2d, and 3d tensors
+    original tensor shape has dimension (N, K), or (B, N, K)
+    qdata has dimension (N, K) or (B, N, K)
+    scale (per row quantization) has dimension: (N,) or (B, N)
 
     since qdata has the same dimension as original tensor, we can directly slice that
     for scale, we'll do a slice when dim is 0, and don't need to do anything for dim 1
@@ -431,12 +431,12 @@ def _(func, types, args, kwargs):
     """
     self, dim, start, end, step = fill_defaults(args, 5, [0, None, None, 1])
     assert step == 1
-    assert dim == 0 or dim == 1, f"Only dim==0 or 1 are supported, got: {dim}"
+    assert dim == 0 or dim == 1 or dim == 2, f"Only dim==0,1,2 are supported, got: dim={dim}"
     if end >= self.shape[dim]:
         end = self.shape[dim]
 
-    assert self.qdata.ndim == 2, (
-        f"Expected packed weight to have dim 2, got {self.qdata.dim}"
+    assert self.qdata.ndim == 2 or self.qdata.ndim == 3, (
+        f"Expected packed weight to have dim==2,3 got: dim={self.qdata.ndim}"
     )
 
     # Always slice the qdata
@@ -638,6 +638,34 @@ def _(func, types, args, kwargs):
     )
     return return_and_correct_aliasing(func, args, kwargs, new_float8_tensor)
 
+@implements(aten.unsqueeze.default)
+def _(func, types, args, kwargs):
+    self, dim = args
+    assert dim == 0, f"Only dim == 0 is supported, got: {dim}"
+    qdata = self.qdata.unsqueeze(dim=dim)
+    scale = self.scale.unsqueeze(dim=dim)
+    block_size = []
+    for i in range(len(qdata.shape)):
+        block_size.append(qdata.shape[i] // scale.shape[i])
+
+    new = self.__class__(
+        qdata,
+        scale,
+        block_size,
+        self.mm_config,
+        self.act_quant_kwargs,
+        self.kernel_preference,
+        self.dtype,
+    )
+    return return_and_correct_aliasing(func, args, kwargs, new)
+
+
+@implements(aten.add.Tensor)
+def _(func, types, args, kwargs):
+    assert len(args) == 2, f"Expected 2 args, got {len(args)}"
+    assert isinstance(args[0], torch.Tensor) and isinstance(args[1], Float8Tensor), f"Expected args[0]==torch.Tensor and args[1]==Float8Tensor, got {type(args[0]), type(args[1])}"
+    sum_tensor = args[0] + args[1].dequantize()
+    return return_and_correct_aliasing(func, args, kwargs, sum_tensor)
 
 Float8Tensor.__module__ = "torchao.quantization"