Makes fallback float8 1x128 by 128x128 gemm output bfloat16

[vkuzo]

Summary:

For now, we just care about bf16 output. We can add fp32 and a flag to
control it later, if needed.

Test Plan:

pytest test/quantization/quantize_/workflows/float8/test_float8_tensor.py -s -k fp8_linear_variants -x

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• support eval of float8_a1x128_w128x128 #3269
• -> Makes fallback float8 1x128 by 128x128 gemm output bfloat16 #3265

[pytorch-bot]

🔗 Helpful Links

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

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

As of commit c877d67 with merge base f856d36 ():

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

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

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh) (trunk failure)
  test_expected_gpu_kernel_fbgemm
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/sparsity/test_sparse_api.py::TestQuantSemiSparse::test_sparse_marlin_compile_True

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

[danielvegamyhre]

LGTM. Btw in my benchmarks I found the torch._scaled_mm cutlass kernel for blockwise gemms to be much faster than the triton kernels. This was a few months ago, you can run the benchmarks scripts in this dir if you want: https://github.com/pytorch/ao/tree/main/benchmarks/prototype/blockwise_fp8_training