use the float4 dtype in mxfp4 and nvfp4 tensors

[vkuzo]

Summary:

Uses the torch.float4_e2m1fn_x2 dtype in mxfp4 and nxfp4 torchao
tensors.

Requires pytorch/pytorch#169595, so we need
to wait for the next PyTorch branch cut.

Note: nvfp4 models in vllm currently hit an error
https://gist.github.com/vkuzo/e1407ee68c9ebb8d0f67478aedd81b96
before or after this PR, so some more debugging to do before
landing this.

Test Plan:

CUDA_VISIBLE_DEVICES=6 time pytest test/prototype/mx_formats/ -s

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• -> use the float4 dtype in mxfp4 and nvfp4 tensors #3440

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 869147e with merge base 534bea5 ():

NEW FAILURES - The following jobs have failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  RuntimeError: Command docker exec -t 57a9873f670a8bd304316bd026d303be0b218267a8d99758b26fac5e52ea96e0 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.8, linux.g5.12xlarge.nvidia.gpu, torch==2.8.0, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 1529989f6ae8330b42f080a5c2764b325dc8ead6589313b66af5c2f91c8fe7b8 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.9, linux.g5.12xlarge.nvidia.gpu, torch==2.9.1, cuda, 12.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 248181e31db0b377a765e152802e60322ebf9b38701d649c68589ffc805f9b94 /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t a332c9766735aa8e747e43dd55551ce68fb601fb1b121829d76567a4bd6f31b2 /exec failed with exit code 1

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