Fix TPU test

Signed-off-by: luka <[email protected]>
vllm-project · Dec 5, 2024 · b8957f7 · b8957f7
1 parent 7046e4b
commit b8957f7
Showing 1 changed file with 12 additions and 10 deletions.
diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
@@ -269,16 +269,18 @@ def rms_norm_dynamic_per_token_quant(
     return output, scales
 
 
-@register_fake("_C::rms_norm_dynamic_per_token_quant")
-def _rms_norm_dynamic_per_token_quant_fake(
-        output: torch.Tensor,
-        input: torch.Tensor,
-        weight: torch.Tensor,
-        scales: torch.Tensor,
-        epsilon: float,
-        scale_ub: Optional[torch.Tensor] = None,
-        residual: Optional[torch.Tensor] = None) -> None:
-    return None
+if hasattr(torch.ops._C, "rms_norm_dynamic_per_token_quant"):
+
+    @register_fake("_C::rms_norm_dynamic_per_token_quant")
+    def _rms_norm_dynamic_per_token_quant_fake(
+            output: torch.Tensor,
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            scales: torch.Tensor,
+            epsilon: float,
+            scale_ub: Optional[torch.Tensor] = None,
+            residual: Optional[torch.Tensor] = None) -> None:
+        return None
 
 
 # quantization ops