fix qwen3 infer bug. (#878)

Author: hiworldwzj

lightllm/models/llama/triton_kernel/rmsnorm.py

@@ -57,9 +57,9 @@ def rmsnorm_forward(x: torch.Tensor, weight, eps, out=None):
         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)
-    # print(BLOCK_SIZE, num_warps, "block_size, numwarps")
-    BLOCK_SIZE = 128 * 2 * 2 * 2 * 2 * 2 * 2 * 2
-    num_warps = 8
+    num_warps = triton.next_power_of_2(num_warps)
+    if BLOCK_SIZE > 16384:
+        BLOCK_SIZE = 16384
     # enqueue kernel
     _rms_norm_fwd_fused[(M,)](
         x_arg,

lightllm/models/qwen3/layer_infer/transformer_layer_infer.py

@@ -35,19 +35,20 @@ def _get_qkv(
         cache_kv = layer_weight.kv_proj.mm(
             input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
         ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+
         rmsnorm_forward(
-            q.reshape(-1, self.head_dim_),
+            q.view(-1, self.head_dim_),
             weight=layer_weight.q_norm_weight_.weight,
             eps=self.eps_,
-            out=q.reshape(-1, self.head_dim_),
+            out=q.view(-1, self.head_dim_),
         )
 
-        rmsnorm_forward(
-            cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, self.head_dim_),
+        cache_kv[:, : self.tp_k_head_num_, :] = rmsnorm_forward(
+            cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, cache_kv.shape[-1]),
             weight=layer_weight.k_norm_weight_.weight,
             eps=self.eps_,
-            out=cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, self.head_dim_),
-        )
+        ).view(-1, self.tp_k_head_num_, cache_kv.shape[-1])
+
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),
             cache_kv[:, : self.tp_k_head_num_, :],

lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py

@@ -61,18 +61,18 @@ def _get_qkv(
             input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
         ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
         rmsnorm_forward(
-            q.reshape(-1, self.head_dim_),
+            q.view(-1, self.head_dim_),
             weight=layer_weight.q_norm_weight_.weight,
             eps=self.eps_,
-            out=q.reshape(-1, self.head_dim_),
+            out=q.view(-1, self.head_dim_),
         )
 
-        rmsnorm_forward(
-            cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, self.head_dim_),
+        cache_kv[:, : self.tp_k_head_num_, :] = rmsnorm_forward(
+            cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, cache_kv.shape[-1]),
             weight=layer_weight.k_norm_weight_.weight,
             eps=self.eps_,
-            out=cache_kv[:, : self.tp_k_head_num_, :].reshape(-1, self.head_dim_),
-        )
+        ).view(-1, self.tp_k_head_num_, cache_kv.shape[-1])
+
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),
             cache_kv[:, : self.tp_k_head_num_, :],