[WIP] Initial prototype of differentiable grouped_scaled_mm function for torchao

[danielvegamyhre]

Summary

The _grouped_scaled_mm function in torchao will do:

• float8 rowwise quantization on inputs
• grouped scaled mm
• do this in a differentiable way

Test plan

• TODO add more robust tests w/ gradient checks

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1969

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures

As of commit 281950c with merge base 923242e ():

NEW FAILURES - The following jobs have failed:

• .github/workflows/float8nocompile_test.yaml (gh)
• Code Analysis with Ruff / build (3.9) (gh)
  Process completed with exit code 1.
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  test_shared_embedding
• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[danielvegamyhre]

Note for reviewer: we need to allow expand this constraint so we can do rowwise scaling of the 2D subtensors embedded in 3D tensors.

e.g.

• A = (M,K) and B = (B,K,N)
• A_scale should be (M,)
• B_scale should be (B,N) <-- because we computed rowwise scales for the (K,N) subtensors