Updates with new scaled-mm api

[drisspg]

Summary

This updates the calls to _scaled_mm to the new signature from this PR: pytorch/pytorch#128683

This is needed to unblock inductor work on scaled_mm.

❯ ./test/test_everything.sh
	.
	.
	.

test/test_fsdp2/test_fsdp2_eager.py .......                                   [100%]

================================ 7 passed in 27.66s =================================
all tests successful

[vkuzo]

LG. Does this need a PyTorch nightly version check to warn people if their version is out of date, or is there no good way to do that?

[drisspg]

We have been pretty fast and loose with versioning. The only required runtime dependency we have is torch and we dont specify a version, but we do make explicit note of the need for nightly:
https://github.com/pytorch-labs/float8_experimental?tab=readme-ov-file#installation

I could(and probably should) update the dependency we have to specify we require newer versions of pytorch but still not perfect since, nightly doesnt follow semantic versioning

[facebook-github-bot]

@drisspg has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@drisspg has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@drisspg merged this pull request in edae9a3.