use symbolic shape for FP8RowwiseQuantizedToFloat

fbgemm_gpu/src/quantize_ops/quantize_ops_meta.cpp

@@ -26,24 +26,26 @@ Tensor FP8rowwise_to_float_meta(
     const int64_t output_dtype) {
   TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
 
-  const auto input_sizes = input.sizes();
+  const at::SymIntArrayRef input_sizes = input.sym_sizes();
+
   const auto last_dim = input_sizes.size() - 1;
-  const int ncols = input_sizes[last_dim];
-  const int ncols_aligned = (ncols + 4 - 1) / 4 * 4;
-  const int output_columns = ncols_aligned - 2 * sizeof(float);
+  const at::SymInt ncols = input_sizes[last_dim];
+  const at::SymInt ncols_aligned = (ncols + 4 - 1) / 4 * 4;
+  const at::SymInt output_columns = ncols_aligned - 2 * sizeof(float);
 
-  auto output_dims = input_sizes.vec();
+  c10::SymDimVector output_dims(input_sizes.begin(), input_sizes.end());
   output_dims[last_dim] = output_columns;
   SparseType output_sparse_dtype = static_cast<SparseType>(output_dtype);
   switch (output_sparse_dtype) {
     case SparseType::FP32:
-      return at::empty(output_dims, input.options().dtype(at::kFloat));
+      return at::empty_symint(output_dims, input.options().dtype(at::kFloat));
     case SparseType::FP16:
-      return at::empty(output_dims, input.options().dtype(at::kHalf));
+      return at::empty_symint(output_dims, input.options().dtype(at::kHalf));
     case SparseType::BF16:
-      return at::empty(output_dims, input.options().dtype(at::kBFloat16));
+      return at::empty_symint(
+          output_dims, input.options().dtype(at::kBFloat16));
     default:
-      TORCH_CHECK(false);
+      TORCH_CHECK(false, "Unsupported output dtype ");
   }
 }
 

fbgemm_gpu/test/jagged_tensor_ops_test.py

@@ -28,6 +28,7 @@
         gpu_unavailable,
         on_arm_platform,
         running_on_github,
+        symint_vector_unsupported,
         TEST_WITH_ROCM,
     )
 except Exception:
@@ -38,6 +39,7 @@
         gpu_unavailable,
         on_arm_platform,
         running_on_github,
+        symint_vector_unsupported,
         TEST_WITH_ROCM,
     )
 
@@ -108,17 +110,6 @@ def hash_size_cumsum_to_offsets(hash_size_cum_sum_list: List[int]) -> List[int]:
     return hash_size_offsets_list
 
 
-def symint_vector_unsupported() -> Tuple[bool, str]:
-    major, minor = torch.__version__.split(".")[0:2]
-    return (
-        int(major) < 2 or (int(major) == 2 and int(minor) < 1),
-        """
-        dynamic shape support for this op needs to be on PyTorch 2.1 or
-        newer with https://github.com/pytorch/pytorch/pull/101056
-        """,
-    )
-
-
 class JaggedTensorOpsTest(unittest.TestCase):
     def setUp(self) -> None:
         if symint_vector_unsupported()[0]:

fbgemm_gpu/test/quantize_ops_test.py

@@ -32,7 +32,7 @@
         fused_rowwise_nbit_quantize_reference,
         gpu_available,
         gpu_unavailable,
-        skipIfRocm,
+        symint_vector_unsupported,
     )
 except Exception:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
@@ -45,6 +45,7 @@
         fused_rowwise_nbit_quantize_reference,
         gpu_available,
         gpu_unavailable,
+        symint_vector_unsupported,
     )
 
 no_long_tests: bool = False
@@ -995,6 +996,8 @@ def setUp(self) -> None:
                 torch.bfloat16,
             ],
         ),
+        # if before PT 2.1, we don't support symint_vector, so turn it off
+        test_compile=st.booleans() if symint_vector_unsupported() else st.just(False),
     )
     @settings(verbosity=Verbosity.verbose, max_examples=10, deadline=None)
     def test_quantize_and_dequantize_op_fp8_rowwise(
@@ -1006,6 +1009,7 @@ def test_quantize_and_dequantize_op_fp8_rowwise(
         forward: bool,
         given_last_dim: bool,
         dtype: torch.dtype,
+        test_compile: bool,
     ) -> None:
         n = n * 4  # need (n % 4 == 0)
         input_data = (
@@ -1018,7 +1022,21 @@ def test_quantize_and_dequantize_op_fp8_rowwise(
         quantized_data_gpu = torch.ops.fbgemm.FloatToFP8RowwiseQuantized(
             input_data_gpu, forward=forward
         )
-        dequantized_data_gpu = torch.ops.fbgemm.FP8RowwiseQuantizedToFloat(
+        quantize_func = (
+            torch.compile(
+                torch.ops.fbgemm.FP8RowwiseQuantizedToFloat,
+                dynamic=True,
+                fullgraph=True,
+            )
+            if test_compile
+            else torch.ops.fbgemm.FP8RowwiseQuantizedToFloat
+        )
+
+        if test_compile:
+            torch._dynamo.mark_dynamic(quantized_data_gpu, 0)
+            torch._dynamo.mark_dynamic(quantized_data_gpu, 1)
+
+        dequantized_data_gpu = quantize_func(
             quantized_data_gpu,
             forward=forward,
             output_dtype=SparseType.FP32.as_int()

fbgemm_gpu/test/test_utils.py

@@ -222,3 +222,14 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
         return wrapper
 
     return skipIfRocmDecorator
+
+
+def symint_vector_unsupported() -> Tuple[bool, str]:
+    major, minor = torch.__version__.split(".")[0:2]
+    return (
+        int(major) < 2 or (int(major) == 2 and int(minor) < 1),
+        """
+        dynamic shape support for this op needs to be on PyTorch 2.1 or
+        newer with https://github.com/pytorch/pytorch/pull/101056
+        """,
+    )