[cleanup][2/x] split float8 mm by delayed vs dynamic

[vkuzo]

Summary:

Before this PR, the float8 mm logic was split by axiswise vs tensorwise.
After this PR, the float8 mm logic is split by dynamic vs non-dynamic scaling.

Motivation: there is more and more evidence that dynamic scaling will
be common to the most important lowp recipes. This PR is a step on the way
to making the dynamic scaling logic be simpler and easier to understand
in torchao.float8.

There are a lot of other simplifications to do, but stopping here to
keep the PR small. This is a pure refactor without any logic changes.

Test Plan:

./test/float8/test_everything.sh

// torchtitan LLaMa 3 8B on 8 H100s

// baseline (bf16 + compile)
37.87GiB(39.86%)  tps: 5,787

// before this PR

// dynamic-tensorwise
memory: 37.98GiB(39.98%)  tps: 6,880
// dynamic-tensorwise with float8 all-gather
memory: 37.54GiB(39.51%)  tps: 7,126
// all-axiswise
memory: 57.03GiB(60.03%)  tps: 6,330
// lw_axiswise_with_gw_hp
memory: 57.03GiB(60.03%)  tps: 6,021

// after this PR
// dynamic-tensorwise
memory: 37.98GiB(39.98%)  tps: 6,891
// dynamic-tensorwise with float8 all-gather
memory: 37.54GiB(39.51%)  tps: 7,127
// all-axiswise
memory: 57.03GiB(60.03%)  tps: 6,345
// lw_axiswise_with_gw_hp
memory: 57.03GiB(60.03%)  tps: 6,042

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• [cleanup][4/x] unify weight casting #1481
• [cleanup][3/x] unify dynamic input and grad_output casting #1480
• -> [cleanup][2/x] split float8 mm by delayed vs dynamic #1461

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[drisspg]

I figured it would have just returned the scale on the fp8 tensor