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: `colossalai/kernel/triton/rms_layernorm.py`
Content:
```
1 import torch
2 
3 try:
4     import triton
5     import triton.language as tl
6 
7     HAS_TRITON = True
8 except ImportError:
9     HAS_TRITON = False
10     print("please install triton from https://github.com/openai/triton")
11 
12 if HAS_TRITON:
13     # CREDITS: These functions are adapted from the Triton tutorial
14     # https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
15 
16     @triton.jit
17     def _rmsnorm_kernel(
18         X,  # pointer to the input
19         Y,  # pointer to the output
20         W,  # pointer to the weights
21         stride,  # how much to increase the pointer when moving by 1 row
22         N,  # number of columns in X
23         eps,  # epsilon to avoid division by zero
24         BLOCK_SIZE: tl.constexpr,
25     ):
26 
27         # This triton kernel implements Root Mean Square Layer Norm (RMSNorm).
28 
29         # Map the program id to the row of X and Y it should compute.
30         row = tl.program_id(0)
31         Y += row * stride
32         X += row * stride
33         # Compute variance
34         _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
35         for off in range(0, N, BLOCK_SIZE):
36             cols = off + tl.arange(0, BLOCK_SIZE)
37             x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
38             x = tl.where(cols < N, x, 0.0)
39             _var += x * x
40         var = tl.sum(_var, axis=0) / N
41         rstd = 1 / tl.sqrt(var + eps)
42         # Normalize and apply linear transformation
43         for off in range(0, N, BLOCK_SIZE):
44             cols = off + tl.arange(0, BLOCK_SIZE)
45             mask = cols < N
46             w = tl.load(W + cols, mask=mask)
47             x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32)
48             x_hat = x * rstd
49             y = x_hat * w
50             # Write output
51             tl.store(Y + cols, y.to(tl.float16), mask=mask)
52 
53     @torch.no_grad()
54     def rms_layernorm(x, weight, eps):
55         # allocate output
56         y = torch.empty_like(x)
57         # reshape input data into 2D tensor
58         x_arg = x.reshape(-1, x.shape[-1])
59         M, N = x_arg.shape
60         # Less than 64KB per feature: enqueue fused kernel
61         MAX_FUSED_SIZE = 65536 // x.element_size()
62         BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
63         if N > BLOCK_SIZE:
64             raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
65         # heuristics for number of warps
66         num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
67         # enqueue kernel
68         _rmsnorm_kernel[(M,)](
69             x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
70         )
71         return y
72
```
--- 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/colossalai/kernel/triton/rms_layernorm.py b/colossalai/kernel/triton/rms_layernorm.py
--- a/colossalai/kernel/triton/rms_layernorm.py
+++ b/colossalai/kernel/triton/rms_layernorm.py
@@ -23,7 +23,6 @@
         eps,  # epsilon to avoid division by zero
         BLOCK_SIZE: tl.constexpr,
     ):
-
         # This triton kernel implements Root Mean Square Layer Norm (RMSNorm).
 
         # Map the program id to the row of X and Y it should compute.
@@ -54,18 +53,19 @@
     def rms_layernorm(x, weight, eps):
         # allocate output
         y = torch.empty_like(x)
-        # reshape input data into 2D tensor
+        # reshape input data into 2D tensor, (total token, hidden_size)
         x_arg = x.reshape(-1, x.shape[-1])
         M, N = x_arg.shape
         # Less than 64KB per feature: enqueue fused kernel
         MAX_FUSED_SIZE = 65536 // x.element_size()
+
         BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
-        if N > BLOCK_SIZE:
+        if N > MAX_FUSED_SIZE:
             raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+
         # heuristics for number of warps
-        num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
+        num_warps = min(max(triton.next_power_of_2(N) // 256, 8), 32)
+
         # enqueue kernel
-        _rmsnorm_kernel[(M,)](
-            x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps
-        )
+        _rmsnorm_kernel[(M,)](x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps)
         return y

{"golden_diff": "diff --git a/colossalai/kernel/triton/rms_layernorm.py b/colossalai/kernel/triton/rms_layernorm.py\n--- a/colossalai/kernel/triton/rms_layernorm.py\n+++ b/colossalai/kernel/triton/rms_layernorm.py\n@@ -23,7 +23,6 @@\n         eps,  # epsilon to avoid division by zero\n         BLOCK_SIZE: tl.constexpr,\n     ):\n-\n         # This triton kernel implements Root Mean Square Layer Norm (RMSNorm).\n \n         # Map the program id to the row of X and Y it should compute.\n@@ -54,18 +53,19 @@\n     def rms_layernorm(x, weight, eps):\n         # allocate output\n         y = torch.empty_like(x)\n-        # reshape input data into 2D tensor\n+        # reshape input data into 2D tensor, (total token, hidden_size)\n         x_arg = x.reshape(-1, x.shape[-1])\n         M, N = x_arg.shape\n         # Less than 64KB per feature: enqueue fused kernel\n         MAX_FUSED_SIZE = 65536 // x.element_size()\n+\n         BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))\n-        if N > BLOCK_SIZE:\n+        if N > MAX_FUSED_SIZE:\n             raise RuntimeError(\"This layer norm doesn't support feature dim >= 64KB.\")\n+\n         # heuristics for number of warps\n-        num_warps = min(max(BLOCK_SIZE // 256, 1), 8)\n+        num_warps = min(max(triton.next_power_of_2(N) // 256, 8), 32)\n+\n         # enqueue kernel\n-        _rmsnorm_kernel[(M,)](\n-            x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps\n-        )\n+        _rmsnorm_kernel[(M,)](x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps)\n         return y\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\ntry:\n    import triton\n    import triton.language as tl\n\n    HAS_TRITON = True\nexcept ImportError:\n    HAS_TRITON = False\n    print(\"please install triton from https://github.com/openai/triton\")\n\nif HAS_TRITON:\n    # CREDITS: These functions are adapted from the Triton tutorial\n    # https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html\n\n    @triton.jit\n    def _rmsnorm_kernel(\n        X,  # pointer to the input\n        Y,  # pointer to the output\n        W,  # pointer to the weights\n        stride,  # how much to increase the pointer when moving by 1 row\n        N,  # number of columns in X\n        eps,  # epsilon to avoid division by zero\n        BLOCK_SIZE: tl.constexpr,\n    ):\n\n        # This triton kernel implements Root Mean Square Layer Norm (RMSNorm).\n\n        # Map the program id to the row of X and Y it should compute.\n        row = tl.program_id(0)\n        Y += row * stride\n        X += row * stride\n        # Compute variance\n        _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n        for off in range(0, N, BLOCK_SIZE):\n            cols = off + tl.arange(0, BLOCK_SIZE)\n            x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)\n            x = tl.where(cols < N, x, 0.0)\n            _var += x * x\n        var = tl.sum(_var, axis=0) / N\n        rstd = 1 / tl.sqrt(var + eps)\n        # Normalize and apply linear transformation\n        for off in range(0, N, BLOCK_SIZE):\n            cols = off + tl.arange(0, BLOCK_SIZE)\n            mask = cols < N\n            w = tl.load(W + cols, mask=mask)\n            x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32)\n            x_hat = x * rstd\n            y = x_hat * w\n            # Write output\n            tl.store(Y + cols, y.to(tl.float16), mask=mask)\n\n    @torch.no_grad()\n    def rms_layernorm(x, weight, eps):\n        # allocate output\n        y = torch.empty_like(x)\n        # reshape input data into 2D tensor\n        x_arg = x.reshape(-1, x.shape[-1])\n        M, N = x_arg.shape\n        # Less than 64KB per feature: enqueue fused kernel\n        MAX_FUSED_SIZE = 65536 // x.element_size()\n        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))\n        if N > BLOCK_SIZE:\n            raise RuntimeError(\"This layer norm doesn't support feature dim >= 64KB.\")\n        # heuristics for number of warps\n        num_warps = min(max(BLOCK_SIZE // 256, 1), 8)\n        # enqueue kernel\n        _rmsnorm_kernel[(M,)](\n            x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps\n        )\n        return y\n", "path": "colossalai/kernel/triton/rms_layernorm.py"}], "after_files": [{"content": "import torch\n\ntry:\n    import triton\n    import triton.language as tl\n\n    HAS_TRITON = True\nexcept ImportError:\n    HAS_TRITON = False\n    print(\"please install triton from https://github.com/openai/triton\")\n\nif HAS_TRITON:\n    # CREDITS: These functions are adapted from the Triton tutorial\n    # https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html\n\n    @triton.jit\n    def _rmsnorm_kernel(\n        X,  # pointer to the input\n        Y,  # pointer to the output\n        W,  # pointer to the weights\n        stride,  # how much to increase the pointer when moving by 1 row\n        N,  # number of columns in X\n        eps,  # epsilon to avoid division by zero\n        BLOCK_SIZE: tl.constexpr,\n    ):\n        # This triton kernel implements Root Mean Square Layer Norm (RMSNorm).\n\n        # Map the program id to the row of X and Y it should compute.\n        row = tl.program_id(0)\n        Y += row * stride\n        X += row * stride\n        # Compute variance\n        _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32)\n        for off in range(0, N, BLOCK_SIZE):\n            cols = off + tl.arange(0, BLOCK_SIZE)\n            x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)\n            x = tl.where(cols < N, x, 0.0)\n            _var += x * x\n        var = tl.sum(_var, axis=0) / N\n        rstd = 1 / tl.sqrt(var + eps)\n        # Normalize and apply linear transformation\n        for off in range(0, N, BLOCK_SIZE):\n            cols = off + tl.arange(0, BLOCK_SIZE)\n            mask = cols < N\n            w = tl.load(W + cols, mask=mask)\n            x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32)\n            x_hat = x * rstd\n            y = x_hat * w\n            # Write output\n            tl.store(Y + cols, y.to(tl.float16), mask=mask)\n\n    @torch.no_grad()\n    def rms_layernorm(x, weight, eps):\n        # allocate output\n        y = torch.empty_like(x)\n        # reshape input data into 2D tensor, (total token, hidden_size)\n        x_arg = x.reshape(-1, x.shape[-1])\n        M, N = x_arg.shape\n        # Less than 64KB per feature: enqueue fused kernel\n        MAX_FUSED_SIZE = 65536 // x.element_size()\n\n        BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))\n        if N > MAX_FUSED_SIZE:\n            raise RuntimeError(\"This layer norm doesn't support feature dim >= 64KB.\")\n\n        # heuristics for number of warps\n        num_warps = min(max(triton.next_power_of_2(N) // 256, 8), 32)\n\n        # enqueue kernel\n        _rmsnorm_kernel[(M,)](x_arg, y, weight, x_arg.stride(0), N, eps, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps)\n        return y\n", "path": "colossalai/kernel/triton/rms_layernorm.py"}]}