Store NVFP4 block scales in swwizzled layout on tensor

[drisspg]

Stacked PRs:

• Add kernel #2439
• ->Store NVFP4 block scales in swwizzled layout on tensor #2438

---

Store NVFP4 block scales in swizzled layout on tensor

For llama3 70b no TP sizes w/ 1024 tokens: 15% E2E speedup

In eager Before: https://fburl.com/7w3j6b1q

nvfp4

Runtime: 2436.98 μs per iteration

In eager After: https://fburl.com/s7ggvm94

nvfp4 Runtime: 2356.77 μs per iteration

In compile

Before: https://fburl.com/1gvfjjlu

nvfp4 Runtime: 576.14 μs per iteration

After: https://fburl.com/usp1xelj

nvfp4 Runtime: 486.69 μs per iteration

Throughput: 47.12 requests/s, 19998.00 total tokens/s, 9635.87 output tokens/s
Total num prompt tokens:  225190
Total num output tokens:  209407

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures

As of commit 5c059bf with merge base faf788a ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t de25b6ba709dc3803d769e679d310e47b19470a4d62bb02d98b6e5ff43763f68 /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 c69e5d6a887058eb46f95d5095d0c18213751de534a2cd5fc086580b8a01ec22 /exec failed with exit code 1
• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  test_replace_q_dq_patterns_with_quantized_linear_ops_pass

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