Merge branch 'main' into fast_start

Author: sufubao

lightllm/models/llama/triton_kernel/rmsnorm.py

@@ -2,6 +2,9 @@
 
 import triton
 import triton.language as tl
+import os
+
+rmsnorm_num_warps = int(os.getenv("RMSNORM_WARPS", "8"))
 
 
 @triton.jit
@@ -56,8 +59,6 @@ def rmsnorm_forward(x: torch.Tensor, weight, eps, out=None):
     if N > BLOCK_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), 4)
-    num_warps = triton.next_power_of_2(num_warps)
     if BLOCK_SIZE > 16384:
         BLOCK_SIZE = 16384
     # enqueue kernel
@@ -72,7 +73,7 @@ def rmsnorm_forward(x: torch.Tensor, weight, eps, out=None):
         N,
         eps,
         BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=num_warps,
+        num_warps=rmsnorm_num_warps,
     )
     return y