[doesn't work yet] enable Float8Tensor as subgraph boundary

Summary:

In https://github.com/pytorch/pytorch/pull/114311/files, the signature
expected by traceable subclasses changed.

This PR updates `Float8Tensor` to the new spec.

Note: doesn't work yet, need to debug.

This may be blocking composability of Float8Linear + FSDP +
torch.compile.

Test Plan:

```
pytest test/test_compile.py -s --sw -k graph_break
// currently broken
// logs: https://gist.github.com/vkuzo/ba98a01a459fb9c966f167d8ecca1780
```

Reviewers:

Subscribers:

Tasks:

Tags:

Author: vkuzo

test/test_compile.py

@@ -11,6 +11,7 @@
 import torch
 import torch.nn as nn
 from float8_experimental.float8_linear_utils import get_float8_linear, LinearType
+from float8_experimental.float8_tensor import Float8Tensor
 
 # Setting to unblock for calling contiguous in backwards
 is_H100 = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0)
@@ -75,6 +76,32 @@ def test_inductor(fullgraph, emulate: bool, linear_type: bool, dtype: torch.dtyp
     torch._dynamo.reset()
     _test_compile_base("inductor", fullgraph, emulate, linear_type, dtype)
 
+def test_float8_with_graph_break():
+    # test that having Float8Tensor object at the boundary of a subgraph
+    # works
+
+    class M(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.register_buffer("fp8_amax_x", torch.tensor(1.0))
+            self.register_buffer("fp8_scale_x", torch.tensor(1.0))
+
+        def forward(self, x):
+            x_fp8 = Float8Tensor.to_float8(
+                x, self.fp8_scale_x, torch.float8_e4m3fn, self.fp8_amax_x, 
+                emulate=False,
+            )
+
+            # graph break
+            print('foo')
+            x_hp = x_fp8.to_original_precision()
+            return x_hp
+
+    m = M()
+    m = torch.compile(m)
+    x = torch.randn(16, 16)
+    y = m(x)
+
 
 if __name__ == "__main__":
     pytest.main([__file__])