add option to save profiling traces in inference roofline script

[vkuzo]

Summary:

convenient to analyze differences between roofline and observed

tl;dr; of findings:

mxfp8

1. need to pre-swizzle weights
2. torch.compile gives us two kernels, will repurpose the manual
   training kernel for this, will need to add pre-swizzling. Longer
   term, can see if fbgemm_gpu one is faster.

mxfp4

1. need to pre-swizzle weights
2. need a faster gemm (can use fbgemm_gpu)
3. need a fused activation quant kernel (can use fbgemm_gpu)

nvfp4

1. need to speed up existing triton activation quant kernel, currently
   it doesn't autotune anything so probably some easy wins here. Longer
   term can also benchmark vs fbgemm_gpu

Test Plan:

CUDA_VISIBLE_DEVICES=5 python benchmarks/float8/float8_inference_roofline.py ~/local/tmp/20251016_inference_nvfp4.csv --recipe_name nvfp4 --save_profile_traces True

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• mxtensor: add pre-swizzle support #3200
• mxtensor: make scale shape match qdata #3198
• -> add option to save profiling traces in inference roofline script #3196

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 438b35e with merge base d1a7fbc ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/dtypes/test_nf4.py::TestComm::test_comm

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