improve rotary emb for deepseek

none · none · commit c5936e37b1b7 · 2025-08-19T07:46:33.000Z
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 4}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1": {"BLOCK_SEQ": 32, "HEAD_PARALLEL_NUM": 16, "num_warps": 2, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 2}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 2}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 8, "num_warps": 4, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1024": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1024": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1024": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1024": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=1024,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"1024": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"128": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"128": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 8, "num_warps": 1, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"128": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 8, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"128": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=128,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"128": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 4, "num_stages": 4}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"2048": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"2048": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"2048": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"2048": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 2}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=2048,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"2048": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"256": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"256": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 2}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"256": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"256": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=256,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"256": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 2}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 1, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=4096,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"4096": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 1, "num_warps": 1, "num_stages": 2}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=128,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"8192": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 4}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=16,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"8192": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 1}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=32,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"8192": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 2, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=64,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"8192": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 3}}
diff --git a/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json b/lightllm/common/all_kernel_configs/deepseek_v3_rotary_emb_kernel/{HEAD_DIM=64,K_HEAD_NUM=1,M=8192,Q_HEAD_NUM=8,dtype=torch.bfloat16}_NVIDIA_H200.json
@@ -0,0 +1 @@
+{"8192": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 1, "num_warps": 1, "num_stages": 5}}
diff --git a/lightllm/models/deepseek2/triton_kernel/rotary_emb.py b/lightllm/models/deepseek2/triton_kernel/rotary_emb.py
@@ -21,48 +21,74 @@ def _rotary_kernel(
     stride_sinbs,
     stride_sind,
     max_total_len,
+    HEAD_PARALLEL_NUM: tl.constexpr,
     HEAD_Q: tl.constexpr,
     HEAD_K: tl.constexpr,
     BLOCK_SEQ: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     NUM_STAGE: tl.constexpr,
-):
-    seq_block_index = tl.program_id(0)
+):  
+    head_start_index = tl.program_id(0)
+    seq_block_index = tl.program_id(1)
     seq_start_index = seq_block_index * BLOCK_SEQ
     seq_end_index = (seq_block_index + 1) * BLOCK_SEQ
     seq_end_index = tl.where(seq_end_index < max_total_len, seq_end_index, max_total_len)
 
     dim_range0 = tl.arange(0, BLOCK_DMODEL // 2) * 2
     dim_range1 = dim_range0 + 1
     cos_range = tl.arange(0, BLOCK_DMODEL // 2)
-    for seq_index in tl.range(seq_start_index, seq_end_index, num_stages=NUM_STAGE):
+    for seq_index in tl.range(seq_start_index, seq_end_index):
 
         off_dimcos_sin = seq_index * stride_cosbs + cos_range * stride_cosd
         cos = tl.load(Cos + off_dimcos_sin)
         sin = tl.load(Sin + off_dimcos_sin)
-
-        for q_head_index in tl.static_range(0, HEAD_Q, step=1):
-            off_q0 = seq_index * stride_qbs + q_head_index * stride_qh + dim_range0 * stride_qd
-            off_q1 = seq_index * stride_qbs + q_head_index * stride_qh + dim_range1 * stride_qd
-            q0 = tl.load(Q + off_q0)
-            q1 = tl.load(Q + off_q1)
-            out_q0 = q0 * cos - q1 * sin
-            out_q1 = q0 * sin + q1 * cos
-            tl.store(Q + off_q0, out_q0)
-            tl.store(Q + off_q1, out_q1)
-
-        for k_head_index in tl.static_range(0, HEAD_K, step=1):
-            off_k0 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range0 * stride_kd
-            off_k1 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range1 * stride_kd
-
-            k0 = tl.load(K + off_k0)
-            k1 = tl.load(K + off_k1)
-
-            out_k0 = k0 * cos - k1 * sin
-            out_k1 = k0 * sin + k1 * cos
-
-            tl.store(K + off_k0, out_k0)
-            tl.store(K + off_k1, out_k1)
+        
+        if HEAD_PARALLEL_NUM == 1:
+            for q_head_index in tl.static_range(0, HEAD_Q, step=1):
+                off_q0 = seq_index * stride_qbs + q_head_index * stride_qh + dim_range0 * stride_qd
+                off_q1 = seq_index * stride_qbs + q_head_index * stride_qh + dim_range1 * stride_qd
+                q0 = tl.load(Q + off_q0)
+                q1 = tl.load(Q + off_q1)
+                out_q0 = q0 * cos - q1 * sin
+                out_q1 = q0 * sin + q1 * cos
+                tl.store(Q + off_q0, out_q0)
+                tl.store(Q + off_q1, out_q1)
+
+            for k_head_index in tl.static_range(0, HEAD_K, step=1):
+                off_k0 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range0 * stride_kd
+                off_k1 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range1 * stride_kd
+
+                k0 = tl.load(K + off_k0)
+                k1 = tl.load(K + off_k1)
+
+                out_k0 = k0 * cos - k1 * sin
+                out_k1 = k0 * sin + k1 * cos
+
+                tl.store(K + off_k0, out_k0)
+                tl.store(K + off_k1, out_k1)
+        else:
+            for q_head_index in tl.range(head_start_index, HEAD_Q, step=HEAD_PARALLEL_NUM, num_stages=NUM_STAGE):
+                off_q0 = seq_index * stride_qbs + q_head_index * stride_qh + dim_range0 * stride_qd
+                off_q1 = seq_index * stride_qbs + q_head_index * stride_qh + dim_range1 * stride_qd
+                q0 = tl.load(Q + off_q0)
+                q1 = tl.load(Q + off_q1)
+                out_q0 = q0 * cos - q1 * sin
+                out_q1 = q0 * sin + q1 * cos
+                tl.store(Q + off_q0, out_q0)
+                tl.store(Q + off_q1, out_q1)
+
+            for k_head_index in tl.range(head_start_index, HEAD_K, step=HEAD_PARALLEL_NUM, num_stages=NUM_STAGE):
+                off_k0 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range0 * stride_kd
+                off_k1 = seq_index * stride_kbs + k_head_index * stride_kh + dim_range1 * stride_kd
+
+                k0 = tl.load(K + off_k0)
+                k1 = tl.load(K + off_k1)
+
+                out_k0 = k0 * cos - k1 * sin
+                out_k1 = k0 * sin + k1 * cos
+
+                tl.store(K + off_k0, out_k0)
+                tl.store(K + off_k1, out_k1)
     return
 
 
@@ -87,10 +113,11 @@ def rotary_emb_fwd(q, k, cos, sin, **run_config):
         )
 
     BLOCK_SEQ = run_config["BLOCK_SEQ"]
+    HEAD_PARALLEL_NUM = run_config["HEAD_PARALLEL_NUM"]
     num_warps = run_config["num_warps"]
     num_stages = run_config["num_stages"]
 
-    grid = (triton.cdiv(total_len, BLOCK_SEQ),)
+    grid = (HEAD_PARALLEL_NUM, triton.cdiv(total_len, BLOCK_SEQ),)
     _rotary_kernel[grid](
         Q=q,
         K=k,
@@ -108,6 +135,7 @@ def rotary_emb_fwd(q, k, cos, sin, **run_config):
         stride_sind=sin.stride(1),
         max_total_len=total_len,
         HEAD_Q=head_num_q,
+        HEAD_PARALLEL_NUM=HEAD_PARALLEL_NUM,
         HEAD_K=head_num_k,
         BLOCK_SEQ=BLOCK_SEQ,
         BLOCK_DMODEL=head_dim,
diff --git a/lightllm/models/deepseek2/triton_kernel/rotary_emb_config.py b/lightllm/models/deepseek2/triton_kernel/rotary_emb_config.py
@@ -36,9 +36,10 @@ def try_to_get_best_config(
             return config
         else:
             if M <= 256:
-                config = {"BLOCK_SEQ": 1, "NUM_STAGE": 1, "num_warps": 1, "num_stages": 1}
+                config = {"BLOCK_SEQ": 1, "NUM_STAGE": 1, "num_warps": 1, "num_stages": 1, "HEAD_PARALLEL_NUM": 1}
             else:
-                config = {"BLOCK_SEQ": 16, "NUM_STAGE": 1, "num_warps": 1, "num_stages": 1}
+                config = {"BLOCK_SEQ": 16, "NUM_STAGE": 1, "num_warps": 1, "num_stages": 1, "HEAD_PARALLEL_NUM": 1}
+            
         return config
 
     @classmethod
diff --git a/test/kernel/deepseekv3_rotary_emb_tuning.py b/test/kernel/deepseekv3_rotary_emb_tuning.py
@@ -112,10 +112,11 @@ def worker(
 
 def get_test_configs(split_id, split_count):
     index = 0
-    result = itertools.product([1, 2, 4, 8, 16, 32], [1, 2, 4, 8], [1, 2, 3, 4, 5])
-    for BLOCK_SEQ, num_warps, num_stages in result:
+    result = itertools.product([1, 2, 4, 8, 16, 32], [1, 2, 4, 8], [1, 2, 3, 4, 5], [1, 2, 4, 8, 16])
+    for BLOCK_SEQ, num_warps, num_stages, HEAD_PARALLEL_NUM in result:
         t_config = {
             "BLOCK_SEQ": BLOCK_SEQ,
+            "HEAD_PARALLEL_NUM": HEAD_PARALLEL_NUM,
             "num_warps": num_warps,
             "num_stages": num_stages,
         }
@@ -212,31 +213,32 @@ def tuning_configs(
 if __name__ == "__main__":
     torch.multiprocessing.set_start_method("spawn")
     from lightllm.utils.tuning_utils import mp_tuning
-
+     
     # for deepseekv3 600B
-    q_head_num = 128
-    k_head_num = 1
-    head_dim = 64
-    dtype = torch.bfloat16
-    for m in [1, 128, 256, 1024, 2048, 4096, 8192]:
-        json_dict = {}
-        ans = mp_tuning(
-            tuning_configs,
-            {
-                "M": m,
-                "Q_HEAD_NUM": q_head_num,
-                "K_HEAD_NUM": k_head_num,
-                "HEAD_DIM": head_dim,
-                "dtype": dtype,
-                "test_count": 20,
-            },
-        )
-        json_dict[m] = ans
-        DeepseekV3RotaryKernelConfig.save_config(
-            M=m,
-            Q_HEAD_NUM=q_head_num,
-            K_HEAD_NUM=k_head_num,
-            HEAD_DIM=head_dim,
-            dtype=str(dtype),
-            config_json=json_dict,
-        )
+
+    for q_head_num in [128, 64, 32, 16, 8]:
+        k_head_num = 1
+        head_dim = 64
+        dtype = torch.bfloat16
+        for m in [1, 128, 256, 1024, 2048, 4096, 8192]:
+            json_dict = {}
+            ans = mp_tuning(
+                tuning_configs,
+                {
+                    "M": m,
+                    "Q_HEAD_NUM": q_head_num,
+                    "K_HEAD_NUM": k_head_num,
+                    "HEAD_DIM": head_dim,
+                    "dtype": dtype,
+                    "test_count": 20,
+                },
+            )
+            json_dict[m] = ans
+            DeepseekV3RotaryKernelConfig.save_config(
+                M=m,
+                Q_HEAD_NUM=q_head_num,
+                K_HEAD_NUM=k_head_num,
+                HEAD_DIM=head_dim,
+                dtype=str(dtype),
+                config_json=json_dict,
+            )

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+{"1": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 16, "num_warps": 1, "num_stages": 4}}`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+{"1": {"BLOCK_SEQ": 32, "HEAD_PARALLEL_NUM": 16, "num_warps": 2, "num_stages": 1}}`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+{"1024": {"BLOCK_SEQ": 1, "HEAD_PARALLEL_NUM": 4, "num_warps": 1, "num_stages": 3}}`