[PyTorch] Optimize create_tensor in quantizer

[yaox12]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[yaox12]

/te-ci pytorch

[ptrendx]

That change unfortunately cannot be done and at::zeros is there on purpose - cuBLAS requires the scale inv to be zero padded, otherwise it will result in silent wrong answer errors.

[yaox12]

That change unfortunately cannot be done and at::zeros is there on purpose - cuBLAS requires the scale inv to be zero padded, otherwise it will result in silent wrong answer errors.

Yes, I realized that. I plan to write out zeros in the quantization kernel. We're benchmarking MoE model training in MXFP8 and seeing a lot of these aten::zeros and swizzle_scaling_factor kernels, which makes the speedup of MXFP8 very limited. @xrennvidia
For the swizzling kernels, I plan to implement a multi_tensor_swizzle kernel to fuse them.

[zhongbozhu]

Maybe we should do the following two things:

• I am gonna write a fused torch.empty feature for MXFP8 like [PyTorch][MoE] Reduce CPU Overhead By Fuse Torch Empty Calls #1793 just merged
• while the allocation is torch.empty instead of torch.zeros, the multi_tensor_swizzle will do the zero padding on the tensors

[yaox12]

Maybe we should do the following two things:

• I am gonna write a fused torch.empty feature for MXFP8 like [PyTorch][MoE] Reduce CPU Overhead By Fuse Torch Empty Calls #1793 just merged
• while the allocation is torch.empty instead of torch.zeros, the multi_tensor_swizzle will do the zero padding on the tensors

Sounds great!

This is a better solution, since both zero padding and swizzling are required by cuBLAS, it's better to do both of them in one place.
I just realized filling zeros in the quantization kernel was not able to cover all cases, since there're other producers of MXFP8 tensors, such as cuDNN normalizations.