[not for land] move scaled matmul logic to __torch_dispatch__

[vkuzo]

Summary:

This is a quick prototype of what the float8 matmul logic would look like if we moved it from a torch.autograd.Function subclass to Float8Tensor.__torch_dispatch__.

Pros:

• if we go with this UX, we can eventually remove Float8Linear. Note that the casts will have to be moved to module hooks.
• we get the reshape logic (translating from torch.matmul to torch.mm) for free instead of having to reimplement it

Cons:

• code is harder to follow compared to having a familiar torch.autograd.Function which handles the state management (arguably)
• because of delayed scaling, we need to attach buffer references to Float8Tensor objects, this kinda works but deviates from how people usually use __torch_dispatch__

At this point I do not expect to land this any time soon as keeping the code as is sounds simpler. Putting this up in case we want to discuss this / revisit this later.

Test Plan:

python tests/test.py -k test_linear

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

cc @albanD , @drisspg , @bdhirsh , this is a quick prototype of a more torch_dispatch'ey Float8 code. @albanD and I discussed an earlier version of this offline and for now I'm planning to stick with torch.autograd.Function