Add Float8Tensor (#2463)

Summary:
* Added Float8Tensor that's using fbgemm kernels and scaled_mm:
    * per row activation + per row weight linear calling torch._scaled_mm op (for compatibilty with SM 8.9)
    * per tensor activation + per tensor weight quant linear calling torch._scaled_mm op (for compatibilty with SM 8.9)
    * per row activation + per row weight bmm calling torch.ops.fbgemm.f8f8bf16_rowwise_batched kernel (only works for SM 9.0+) can use batched scaled mm from torch when it's supported: pytorch/pytorch#157950
* dynamic quantization kwargs is added to the Float8Tensor directly
* Added QuantizeTensorKwargs and QuantizeTensorToFloat8Kwargs to store key word args for Float8Tensor.to_float8
* Updated Float8DynamicActivationFloat8WeightConfig and Float8WeightOnlyConfig to use Float8Tensor

Test Plan:
python test/dtypes/test_affine_quantized_float.py
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jerryzh168

.github/workflows/1xH100_tests.yml

@@ -25,15 +25,15 @@ jobs:
         include:
           - name: H100
             runs-on: linux.aws.h100
-            torch-spec: '--pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu126'
+            torch-spec: '--pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https://download.pytorch.org/whl/nightly/cu126'
             gpu-arch-type: "cuda"
             gpu-arch-version: "12.4"
     permissions:
       id-token: write
       contents: read
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
-      timeout: 60
+      timeout: 90
       runner: ${{ matrix.runs-on }}
       gpu-arch-type: ${{ matrix.gpu-arch-type }}
       gpu-arch-version: ${{ matrix.gpu-arch-version }}
@@ -46,8 +46,8 @@ jobs:
         pip install uv
         pip install ${{ matrix.torch-spec }}
         uv pip install -r dev-requirements.txt
-        uv pip install vllm
         pip install .
         pytest test/integration --verbose -s
         pytest test/dtypes/test_affine_quantized_float.py --verbose -s
+        python test/quantization/quantize_/workflows/float8/test_float8_tensor.py
         ./test/float8/test_everything_single_gpu.sh

.github/workflows/1xL4_tests.yml

@@ -46,8 +46,8 @@ jobs:
         pip install uv
         pip install ${{ matrix.torch-spec }}
         uv pip install -r dev-requirements.txt
-        uv pip install vllm
         pip install .
         pytest test/integration --verbose -s
         pytest test/dtypes/test_affine_quantized_float.py --verbose -s
         ./test/float8/test_everything_single_gpu.sh
+        python test/quantization/quantize_/workflows/float8/test_float8_tensor.py

test/dtypes/test_affine_quantized_float.py

@@ -737,6 +737,7 @@ def test_expected_kernels_on_gpu(self, granularity, torch_compile_mode):
         Verify that float8 quantization + torch.compile results in the
         expected number of kernels in the GPU trace.
         """
+        torch.compiler.reset()
 
         M, K, N = 128, 256, 512
         m = torch.nn.Sequential(