import torch

from colossalai.kernel.kernel_loader import InferenceOpsLoader
from colossalai.kernel.triton import rms_layernorm

try:
    import triton  # noqa
except ImportError:
    print("please install triton from https://github.com/openai/triton")

inference_ops = InferenceOpsLoader().load()

# Triton benchmark plot attributions
configs = [
    triton.testing.Benchmark(
        x_names=["SEQUENCE_TOTAL"],
        x_vals=[i for i in range(128, 1025, 128)],
        line_arg="provider",
        line_vals=[
            "vllm_rms_layernorm",
            "triton_rms_layernorm",
            "cuda_rms_layernorm",
            "vllm_rms_layernorm_with_residual",
            "triton_rms_layernorm_with_residual",
            "cuda_rms_layernorm_with_residual",
        ],
        line_names=[
            "vllm_rms_layernorm",
            "triton_rms_layernorm",
            "cuda_rms_layernorm",
            "vllm_rms_layernorm_with_residual",
            "triton_rms_layernorm_with_residual",
            "cuda_rms_layernorm_with_residual",
        ],
        styles=[("red", "-"), ("blue", "-"), ("yellow", "-"), ("red", "--"), ("blue", "--"), ("yellow", "--")],
        ylabel="ms",
        plot_name=f"RMSNorm benchmarking results",
        args={"HIDDEN_SIZE": 1024},
    )
]


@triton.testing.perf_report(configs)
def benchmark_rms_layernorm(
    provider: str,
    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

    dtype = torch.float16
    eps = 1e-5
    x_shape = (SEQUENCE_TOTAL, HIDDEN_SIZE)
    w_shape = (x_shape[-1],)
    residual = torch.rand(x_shape, dtype=dtype, device="cuda")
    weight = torch.ones(w_shape, dtype=dtype, device="cuda")
    vllm_norm = RMSNorm(hidden_size=HIDDEN_SIZE, eps=eps).to(dtype=dtype, device="cuda")
    x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
    if provider == "vllm_rms_layernorm":
        fn = lambda: vllm_norm(x)
    elif provider == "triton_rms_layernorm":
        fn = lambda: rms_layernorm(x, weight, eps=eps)
    elif provider == "cuda_rms_layernorm":
        out = torch.empty_like(x)
        fn = lambda: inference_ops.rms_layernorm(out, x, weight, eps)
    elif provider == "vllm_rms_layernorm_with_residual":
        fn = lambda: vllm_norm(x, residual=residual)
    elif provider == "triton_rms_layernorm_with_residual":
        fn = lambda: rms_layernorm(x, weight, eps=eps, residual=residual)
    elif provider == "cuda_rms_layernorm_with_residual":
        fn = lambda: inference_ops.fused_add_rms_layernorm(x, residual, weight, eps)
    else:
        raise ValueError("Undefined provider.")

    ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)

    return ms


if __name__ == "__main__":
    benchmark_rms_layernorm.run(save_path=".", print_data=True)

examples/inference/benchmark_ops/benchmark_rmsnorm.py
import torch

from colossalai.kernel.kernel_loader import InferenceOpsLoader
from colossalai.kernel.triton import rms_layernorm

try:
    import triton  # noqa
except ImportError:
    print("please install triton from https://github.com/openai/triton")

inference_ops = InferenceOpsLoader().load()

# Triton benchmark plot attributions
configs = [
    triton.testing.Benchmark(
        x_names=["SEQUENCE_TOTAL"],
        x_vals=[i for i in range(128, 1025, 128)],
        line_arg="provider",
        line_vals=[
            "vllm_rms_layernorm",
            "triton_rms_layernorm",
            "cuda_rms_layernorm",
            "vllm_rms_layernorm_with_residual",
            "triton_rms_layernorm_with_residual",
            "cuda_rms_layernorm_with_residual",
        ],
        line_names=[
            "vllm_rms_layernorm",
            "triton_rms_layernorm",
            "cuda_rms_layernorm",
            "vllm_rms_layernorm_with_residual",
            "triton_rms_layernorm_with_residual",
            "cuda_rms_layernorm_with_residual",
        ],
        styles=[("red", "-"), ("blue", "-"), ("yellow", "-"), ("red", "--"), ("blue", "--"), ("yellow", "--")],
        ylabel="ms",
        plot_name=f"RMSNorm benchmarking results",
        args={"HIDDEN_SIZE": 5120},
    )
]


@triton.testing.perf_report(configs)
def benchmark_rms_layernorm(
    provider: str,
    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

    dtype = torch.float16
    eps = 1e-5
    x_shape = (SEQUENCE_TOTAL, HIDDEN_SIZE)
    w_shape = (x_shape[-1],)
    residual = torch.rand(x_shape, dtype=dtype, device="cuda")
    weight = torch.ones(w_shape, dtype=dtype, device="cuda")
    vllm_norm = RMSNorm(hidden_size=HIDDEN_SIZE, eps=eps).to(dtype=dtype, device="cuda")
    x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device="cuda")
    if provider == "vllm_rms_layernorm":
        fn = lambda: vllm_norm(x)
    elif provider == "triton_rms_layernorm":
        fn = lambda: rms_layernorm(x, weight, eps=eps)
    elif provider == "cuda_rms_layernorm":
        out = torch.empty_like(x)
        fn = lambda: inference_ops.rms_layernorm(out, x, weight, eps)
    elif provider == "vllm_rms_layernorm_with_residual":
        fn = lambda: vllm_norm(x, residual=residual)
    elif provider == "triton_rms_layernorm_with_residual":
        fn = lambda: rms_layernorm(x, weight, eps=eps, residual=residual)
    elif provider == "cuda_rms_layernorm_with_residual":
        fn = lambda: inference_ops.fused_add_rms_layernorm(x, residual, weight, eps)
    else:
        raise ValueError("Undefined provider.")

    ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)

    return ms


if __name__ == "__main__":
    benchmark_rms_layernorm.run(save_path=".", print_data=True)

examples/inference/benchmark_ops/benchmark_rmsnorm.py
diff --git a/examples/inference/benchmark_ops/benchmark_rmsnorm.py b/examples/inference/benchmark_ops/benchmark_rmsnorm.py
index 3b5166af0178..deddac8b127a 100644
--- 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},
     )
 ]
 
diff --git a/extensions/csrc/cuda/rms_layernorm_kernel.cu b/extensions/csrc/cuda/rms_layernorm_kernel.cu
index 9183462ad9f7..f109edca4446 100644
--- a/extensions/csrc/cuda/rms_layernorm_kernel.cu
+++ b/extensions/csrc/cuda/rms_layernorm_kernel.cu
@@ -277,6 +277,9 @@ void rms_layernorm(
       case 2:
         RMSNORM_LAUNCHER(2, block);
         break;
+      case 3:
+        RMSNORM_LAUNCHER(3, block);
+        break;
       case 4:
         RMSNORM_LAUNCHER(4, block);
         break;
@@ -321,6 +324,9 @@ void fused_add_rms_layernorm(
       case 2:
         FUSED_ADD_RMSNORM_LAUNCHER(2, block);
         break;
+      case 3:
+        FUSED_ADD_RMSNORM_LAUNCHER(3, block);
+        break;
       case 4:
         FUSED_ADD_RMSNORM_LAUNCHER(4, block);
         break;
diff --git a/tests/test_infer/test_ops/cuda/test_rms_layernorm.py b/tests/test_infer/test_ops/cuda/test_rms_layernorm.py
index d14010600d9f..0b677fff89e9 100644
--- a/tests/test_infer/test_ops/cuda/test_rms_layernorm.py
+++ b/tests/test_infer/test_ops/cuda/test_rms_layernorm.py
@@ -9,7 +9,7 @@
 
 
 @pytest.mark.parametrize("M", [2, 4, 8, 16])
-@pytest.mark.parametrize("N", [64, 128, 512])
+@pytest.mark.parametrize("N", [64, 128, 512, 5120])
 def test_rms_layernorm(M: int, N: int):
     torch.manual_seed(123)
     torch.cuda.empty_cache()
@@ -48,4 +48,4 @@ def test_rms_layernorm(M: int, N: int):
 
 
 if __name__ == "__main__":
-    test_rms_layernorm(16, 512)
+    test_rms_layernorm(16, 5120)

[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

