Add TORCH_CHECKs

Signed-off-by: Vadim Gimpelson <[email protected]>

Author: vadiklyutiy

vllm/_custom_ops.py

@@ -280,7 +280,7 @@ def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
                        weight: torch.Tensor, epsilon: float) -> None:
     torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
 
-
+@torch.compile
 def apply_repetition_penalties_torch(
         logits: torch.Tensor, prompt_mask: torch.Tensor,
         output_mask: torch.Tensor, repetition_penalties: torch.Tensor) -> None: