[Infer/Fix] Fix Dependency in test - RMSNorm kernel (#5399)

fix dependency in pytest

Author: yuanheng-zhao

tests/test_infer/test_ops/triton/test_rmsnorm_triton.py

@@ -3,13 +3,12 @@
 import triton
 from packaging import version
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
-from vllm.model_executor.layers.layernorm import RMSNorm
 
 from colossalai.kernel.triton import rms_layernorm
 from colossalai.testing.utils import parameterize
 
 try:
-    pass
+    import triton  # noqa
 
     HAS_TRITON = True
 except ImportError:
@@ -85,6 +84,11 @@ def benchmark_rms_layernorm(
     SEQUENCE_TOTAL: int,
     HIDDEN_SIZE: int,
 ):
+    try:
+        from vllm.model_executor.layers.layernorm import RMSNorm
+    except ImportError:
+        raise ImportError("Please install vllm from https://github.com/vllm-project/vllm")
+
     warmup = 10
     rep = 1000