Add support for e2e benchmark for conv2d/conv3d

[jerryzh168]

Summary:
We added this to float8_inference_roofline to reuse code but we haven't enabled the roofline feature. For now we just need the e2e speedup time for single conv2d/conv3d against bf16 to understand the speedup expectation

some if/else branches right now but hopefully we should be able to remove some of these when more roofline support is added for conv2d/conv3d

Also added B200 hardware spec.

Test Plan:

python float8_inference_roofline.py ~/local/tmp/test.csv \
    --recipe_name tensorwise \
    --shape_gen_name custom \
    --M 1 --K 160 --N 160 \
    --D 3 --H 50 --W 50 \
    --kernel_size 3 \
    --op_name conv3d

requires fbgemm_gpu_genai nightly and torch nightly:

pip install --pre torch fbgemm_gpu_genai --index-url https://download.pytorch.org/whl/nightly/cu129

sample results:

,fwd_M,fwd_K,fwd_N,D,H,W,kernel_size,r_bf16_gemm_s,r_fp8_gemm_s,r_fp8_ovhd_s,r_fp8_gemm_and_ovhd_s,r_fp8_gemm_and_ovhd_spdp,b_bf16_gemm_s,b_fp8_gemm_s,b_bf16_e2e_s,b_fp8_e2e_s,b_fp8_e2e_spdp,rb_bf16_gemm_ratio,rb\
_fp8_gemm_ratio
0,1,160,320,3,194,130,3,,,,,,0,0,8.248959999999995e-05,0.00011273280000000014,0.7317267024326535,-1,-1

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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! (4 Unrelated Failures)

As of commit e2d0395 with merge base e8c4d09 ():

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

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

• Run Regression Tests / test (CPU 2.6, linux.4xlarge, torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/test_low_bit_optim.py::TestOptim::test_param_groups_optim_name_AdamFp8_device_cpu
• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/test_low_bit_optim.py::TestOptim::test_param_groups_optim_name_AdamFp8_device_cpu
• Run Regression Tests / test (CPU 2.8, linux.4xlarge, torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (trunk failure)
  test/test_low_bit_optim.py::TestOptim::test_param_groups_optim_name_AdamFp8_device_cpu
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh) (trunk failure)
  test/test_low_bit_optim.py::TestOptim::test_param_groups_optim_name_AdamFp8_device_cpu

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