We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/inference/benchmark_ops/benchmark_rmsnorm.py`
Content:
```
1 import torch
2 
3 from colossalai.kernel.kernel_loader import InferenceOpsLoader
4 from colossalai.kernel.triton import rms_layernorm
5 
6 try:
7     import triton  # noqa
8 except ImportError:
9     print("please install triton from https://github.com/openai/triton")
10 
11 inference_ops = InferenceOpsLoader().load()
12 
13 # Triton benchmark plot attributions
14 configs = [
15     triton.testing.Benchmark(
16         x_names=["SEQUENCE_TOTAL"],
17         x_vals=[i for i in range(128, 1025, 128)],
18         line_arg="provider",
19         line_vals=[
20             "vllm_rms_layernorm",
21             "triton_rms_layernorm",
22             "cuda_rms_layernorm",
23             "vllm_rms_layernorm_with_residual",
24             "triton_rms_layernorm_with_residual",
25             "cuda_rms_layernorm_with_residual",
26         ],
27         line_names=[
28             "vllm_rms_layernorm",
29             "triton_rms_layernorm",
30             "cuda_rms_layernorm",
31             "vllm_rms_layernorm_with_residual",
32             "triton_rms_layernorm_with_residual",
33             "cuda_rms_layernorm_with_residual",
34         ],
35         styles=[("red", "-"), ("blue", "-"), ("yellow", "-"), ("red", "--"), ("blue", "--"), ("yellow", "--")],
36         ylabel="ms",
37         plot_name=f"RMSNorm benchmarking results",
38         args={"HIDDEN_SIZE": 1024},
39     )
40 ]
41 
42 
43 @triton.testing.perf_report(configs)
44 def benchmark_rms_layernorm(
45     provider: str,
46     SEQUENCE_TOTAL: int,
47     HIDDEN_SIZE: int,
48 ):
49     try:
50         from vllm.model_executor.layers.layernorm import RMSNorm
51     except ImportError:
52         raise ImportError("Please install vllm from https://github.com/vllm-project/vllm")
53 
54     warmup = 10
55     rep = 1000
56 
57     dtype = torch.float16
58     eps = 1e-5
59     x_shape = (SEQUENCE_TOTAL, HIDDEN_SIZE)
60     w_shape = (x_shape[-1],)
61     residual = torch.rand(x_shape, dtype=dtype, device="cuda")
62     weight = torch.ones(w_shape, dtype=dtype, device="cuda")
63     vllm_norm = RMSNorm(hidden_size=HIDDEN_SIZE, eps=eps).to(dtype=dtype, device="cuda")
64     x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
65     if provider == "vllm_rms_layernorm":
66         fn = lambda: vllm_norm(x)
67     elif provider == "triton_rms_layernorm":
68         fn = lambda: rms_layernorm(x, weight, eps=eps)
69     elif provider == "cuda_rms_layernorm":
70         out = torch.empty_like(x)
71         fn = lambda: inference_ops.rms_layernorm(out, x, weight, eps)
72     elif provider == "vllm_rms_layernorm_with_residual":
73         fn = lambda: vllm_norm(x, residual=residual)
74     elif provider == "triton_rms_layernorm_with_residual":
75         fn = lambda: rms_layernorm(x, weight, eps=eps, residual=residual)
76     elif provider == "cuda_rms_layernorm_with_residual":
77         fn = lambda: inference_ops.fused_add_rms_layernorm(x, residual, weight, eps)
78     else:
79         raise ValueError("Undefined provider.")
80 
81     ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
82 
83     return ms
84 
85 
86 if __name__ == "__main__":
87     benchmark_rms_layernorm.run(save_path=".", print_data=True)
88
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/inference/benchmark_ops/benchmark_rmsnorm.py b/examples/inference/benchmark_ops/benchmark_rmsnorm.py
--- a/examples/inference/benchmark_ops/benchmark_rmsnorm.py
+++ b/examples/inference/benchmark_ops/benchmark_rmsnorm.py
@@ -35,7 +35,7 @@
         styles=[("red", "-"), ("blue", "-"), ("yellow", "-"), ("red", "--"), ("blue", "--"), ("yellow", "--")],
         ylabel="ms",
         plot_name=f"RMSNorm benchmarking results",
-        args={"HIDDEN_SIZE": 1024},
+        args={"HIDDEN_SIZE": 5120},
     )
 ]

{"golden_diff": "diff --git a/examples/inference/benchmark_ops/benchmark_rmsnorm.py b/examples/inference/benchmark_ops/benchmark_rmsnorm.py\n--- a/examples/inference/benchmark_ops/benchmark_rmsnorm.py\n+++ b/examples/inference/benchmark_ops/benchmark_rmsnorm.py\n@@ -35,7 +35,7 @@\n         styles=[(\"red\", \"-\"), (\"blue\", \"-\"), (\"yellow\", \"-\"), (\"red\", \"--\"), (\"blue\", \"--\"), (\"yellow\", \"--\")],\n         ylabel=\"ms\",\n         plot_name=f\"RMSNorm benchmarking results\",\n-        args={\"HIDDEN_SIZE\": 1024},\n+        args={\"HIDDEN_SIZE\": 5120},\n     )\n ]\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import torch\n\nfrom colossalai.kernel.kernel_loader import InferenceOpsLoader\nfrom colossalai.kernel.triton import rms_layernorm\n\ntry:\n    import triton  # noqa\nexcept ImportError:\n    print(\"please install triton from https://github.com/openai/triton\")\n\ninference_ops = InferenceOpsLoader().load()\n\n# Triton benchmark plot attributions\nconfigs = [\n    triton.testing.Benchmark(\n        x_names=[\"SEQUENCE_TOTAL\"],\n        x_vals=[i for i in range(128, 1025, 128)],\n        line_arg=\"provider\",\n        line_vals=[\n            \"vllm_rms_layernorm\",\n            \"triton_rms_layernorm\",\n            \"cuda_rms_layernorm\",\n            \"vllm_rms_layernorm_with_residual\",\n            \"triton_rms_layernorm_with_residual\",\n            \"cuda_rms_layernorm_with_residual\",\n        ],\n        line_names=[\n            \"vllm_rms_layernorm\",\n            \"triton_rms_layernorm\",\n            \"cuda_rms_layernorm\",\n            \"vllm_rms_layernorm_with_residual\",\n            \"triton_rms_layernorm_with_residual\",\n            \"cuda_rms_layernorm_with_residual\",\n        ],\n        styles=[(\"red\", \"-\"), (\"blue\", \"-\"), (\"yellow\", \"-\"), (\"red\", \"--\"), (\"blue\", \"--\"), (\"yellow\", \"--\")],\n        ylabel=\"ms\",\n        plot_name=f\"RMSNorm benchmarking results\",\n        args={\"HIDDEN_SIZE\": 1024},\n    )\n]\n\n\n@triton.testing.perf_report(configs)\ndef benchmark_rms_layernorm(\n    provider: str,\n    SEQUENCE_TOTAL: int,\n    HIDDEN_SIZE: int,\n):\n    try:\n        from vllm.model_executor.layers.layernorm import RMSNorm\n    except ImportError:\n        raise ImportError(\"Please install vllm from https://github.com/vllm-project/vllm\")\n\n    warmup = 10\n    rep = 1000\n\n    dtype = torch.float16\n    eps = 1e-5\n    x_shape = (SEQUENCE_TOTAL, HIDDEN_SIZE)\n    w_shape = (x_shape[-1],)\n    residual = torch.rand(x_shape, dtype=dtype, device=\"cuda\")\n    weight = torch.ones(w_shape, dtype=dtype, device=\"cuda\")\n    vllm_norm = RMSNorm(hidden_size=HIDDEN_SIZE, eps=eps).to(dtype=dtype, device=\"cuda\")\n    x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=\"cuda\")\n    if provider == \"vllm_rms_layernorm\":\n        fn = lambda: vllm_norm(x)\n    elif provider == \"triton_rms_layernorm\":\n        fn = lambda: rms_layernorm(x, weight, eps=eps)\n    elif provider == \"cuda_rms_layernorm\":\n        out = torch.empty_like(x)\n        fn = lambda: inference_ops.rms_layernorm(out, x, weight, eps)\n    elif provider == \"vllm_rms_layernorm_with_residual\":\n        fn = lambda: vllm_norm(x, residual=residual)\n    elif provider == \"triton_rms_layernorm_with_residual\":\n        fn = lambda: rms_layernorm(x, weight, eps=eps, residual=residual)\n    elif provider == \"cuda_rms_layernorm_with_residual\":\n        fn = lambda: inference_ops.fused_add_rms_layernorm(x, residual, weight, eps)\n    else:\n        raise ValueError(\"Undefined provider.\")\n\n    ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)\n\n    return ms\n\n\nif __name__ == \"__main__\":\n    benchmark_rms_layernorm.run(save_path=\".\", print_data=True)\n", "path": "examples/inference/benchmark_ops/benchmark_rmsnorm.py"}], "after_files": [{"content": "import torch\n\nfrom colossalai.kernel.kernel_loader import InferenceOpsLoader\nfrom colossalai.kernel.triton import rms_layernorm\n\ntry:\n    import triton  # noqa\nexcept ImportError:\n    print(\"please install triton from https://github.com/openai/triton\")\n\ninference_ops = InferenceOpsLoader().load()\n\n# Triton benchmark plot attributions\nconfigs = [\n    triton.testing.Benchmark(\n        x_names=[\"SEQUENCE_TOTAL\"],\n        x_vals=[i for i in range(128, 1025, 128)],\n        line_arg=\"provider\",\n        line_vals=[\n            \"vllm_rms_layernorm\",\n            \"triton_rms_layernorm\",\n            \"cuda_rms_layernorm\",\n            \"vllm_rms_layernorm_with_residual\",\n            \"triton_rms_layernorm_with_residual\",\n            \"cuda_rms_layernorm_with_residual\",\n        ],\n        line_names=[\n            \"vllm_rms_layernorm\",\n            \"triton_rms_layernorm\",\n            \"cuda_rms_layernorm\",\n            \"vllm_rms_layernorm_with_residual\",\n            \"triton_rms_layernorm_with_residual\",\n            \"cuda_rms_layernorm_with_residual\",\n        ],\n        styles=[(\"red\", \"-\"), (\"blue\", \"-\"), (\"yellow\", \"-\"), (\"red\", \"--\"), (\"blue\", \"--\"), (\"yellow\", \"--\")],\n        ylabel=\"ms\",\n        plot_name=f\"RMSNorm benchmarking results\",\n        args={\"HIDDEN_SIZE\": 5120},\n    )\n]\n\n\n@triton.testing.perf_report(configs)\ndef benchmark_rms_layernorm(\n    provider: str,\n    SEQUENCE_TOTAL: int,\n    HIDDEN_SIZE: int,\n):\n    try:\n        from vllm.model_executor.layers.layernorm import RMSNorm\n    except ImportError:\n        raise ImportError(\"Please install vllm from https://github.com/vllm-project/vllm\")\n\n    warmup = 10\n    rep = 1000\n\n    dtype = torch.float16\n    eps = 1e-5\n    x_shape = (SEQUENCE_TOTAL, HIDDEN_SIZE)\n    w_shape = (x_shape[-1],)\n    residual = torch.rand(x_shape, dtype=dtype, device=\"cuda\")\n    weight = torch.ones(w_shape, dtype=dtype, device=\"cuda\")\n    vllm_norm = RMSNorm(hidden_size=HIDDEN_SIZE, eps=eps).to(dtype=dtype, device=\"cuda\")\n    x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=\"cuda\")\n    if provider == \"vllm_rms_layernorm\":\n        fn = lambda: vllm_norm(x)\n    elif provider == \"triton_rms_layernorm\":\n        fn = lambda: rms_layernorm(x, weight, eps=eps)\n    elif provider == \"cuda_rms_layernorm\":\n        out = torch.empty_like(x)\n        fn = lambda: inference_ops.rms_layernorm(out, x, weight, eps)\n    elif provider == \"vllm_rms_layernorm_with_residual\":\n        fn = lambda: vllm_norm(x, residual=residual)\n    elif provider == \"triton_rms_layernorm_with_residual\":\n        fn = lambda: rms_layernorm(x, weight, eps=eps, residual=residual)\n    elif provider == \"cuda_rms_layernorm_with_residual\":\n        fn = lambda: inference_ops.fused_add_rms_layernorm(x, residual, weight, eps)\n    else:\n        raise ValueError(\"Undefined provider.\")\n\n    ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)\n\n    return ms\n\n\nif __name__ == \"__main__\":\n    benchmark_rms_layernorm.run(save_path=\".\", print_data=True)\n", "path": "examples/inference/benchmark_ops/benchmark_rmsnorm.py"}]}