add TORCH_VERSION_CHECK for _register_meta (#2575)

Author: shiyang-weng

torchao/quantization/quant_primitives.py

@@ -20,6 +20,7 @@
     TORCH_VERSION_AT_LEAST_2_5,
     TORCH_VERSION_AT_LEAST_2_6,
     _register_custom_op,
+    _register_meta_op,
 )
 
 __all__ = [
@@ -2292,7 +2293,7 @@ def _quantize_affine_float8(
     return fp8_tensor
 
 
-@torch.library.impl(quant_lib, "quantize_affine_float8", "Meta")
+@_register_meta_op(quant_lib, "quantize_affine_float8")
 def _quantize_affine_float8_meta(
     tensor: torch.Tensor,
     scale: torch.Tensor,
@@ -2319,7 +2320,7 @@ def _dequantize_affine_float8(
     return hp_tensor.to(output_dtype)
 
 
-@torch.library.impl(quant_lib, "dequantize_affine_float8", "Meta")
+@_register_meta_op(quant_lib, "dequantize_affine_float8")
 def _dequantize_affine_float8_meta(
     tensor: torch.Tensor,
     scale: torch.Tensor,

torchao/utils.py

@@ -237,6 +237,17 @@ def decorator(fn):
     return decorator
 
 
+def _register_meta_op(lib, op_name):
+    def decorator(fn):
+        if TORCH_VERSION_AT_LEAST_2_5:
+            op = lib.impl(op_name, fn, "Meta")
+            return op
+        else:
+            return fn
+
+    return decorator
+
+
 def get_model_size_in_bytes(model, ignore_embeddings=False):
     """
     Returns the model size in bytes. The option to ignore embeddings