finish, pass test

Author: WANDY666

docs/CN/source/getting_started/benchmark.rst

@@ -89,15 +89,15 @@ ShareGPT 数据集测试 (benchmark_sharegpt.py)
     python test/benchmark/service/benchmark_sharegpt.py \
         --dataset /path/to/sharegpt_dataset.json \
         --tokenizer /path/to/tokenizer \
-        --num_prompts 1000 \
-        --request_rate 10.0
+        --num-prompts 1000 \
+        --request-rate 10.0
 
 **主要参数：**
 
 - ``--dataset``: ShareGPT 格式数据集路径
 - ``--tokenizer``: 分词器路径
-- ``--num_prompts``: 测试提示数量
-- ``--request_rate``: 请求速率 (requests/s)
+- ``--num-prompts``: 测试提示数量
+- ``--request-rate``: 请求速率 (requests/s)
 
 
 Prompt Cache 测试

docs/EN/source/getting_started/benchmark.rst

@@ -88,15 +88,15 @@ Performance testing using ShareGPT real conversation data.
     python test/benchmark/service/benchmark_sharegpt.py \
         --dataset /path/to/sharegpt_dataset.json \
         --tokenizer /path/to/tokenizer \
-        --num_prompts 1000 \
-        --request_rate 10.0
+        --num-prompts 1000 \
+        --request-rate 10.0
 
 **Main Parameters:**
 
 - ``--dataset``: ShareGPT format dataset path
 - ``--tokenizer``: Tokenizer path
-- ``--num_prompts``: Number of test prompts
-- ``--request_rate``: Request rate (requests/s)
+- ``--num-prompts``: Number of test prompts
+- ``--request-rate``: Request rate (requests/s)
 
 Prompt Cache Testing
 ~~~~~~~~~~~~~~~~~~~

lightllm/common/basemodel/cuda_graph.py

@@ -40,6 +40,10 @@ def __init__(self, max_batch_size=8, max_len_in_batch=8192):
         batch_sizes.append(max_batch_size)
         batch_sizes.sort()
 
+        if self.args.enable_fa3_mtp:
+            step_size = self.args.mtp_step + 1
+            batch_sizes = [b for b in batch_sizes if b % step_size == 0]
+
         self.cuda_graph_batch_sizes = batch_sizes
         assert batch_sizes[-1] == self.max_batch_size
         logger.info(f"cuda graph batch_sizes: {self.cuda_graph_batch_sizes}")

lightllm/common/basemodel/triton_kernel/gen_decode_params.py

@@ -12,10 +12,12 @@ def gen_decode_params(b_seq_len: torch.Tensor):
     mtp_step = get_env_start_args().mtp_step
     mtp_size = mtp_step + 1
     enable_fa3_mtp = get_env_start_args().enable_fa3_mtp
+    b_q_seq_len = torch.ones_like(b_seq_len)
 
     if enable_fa3_mtp:
-        b_q_seq_len = torch.ones_like(b_seq_len[: len(b_seq_len) // mtp_size])
-        b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len[mtp_size - 1 :: mtp_size])
+        b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(
+            b_q_seq_len[: len(b_seq_len) // mtp_size], b_kv_seq_len[mtp_size - 1 :: mtp_size]
+        )
     else:
         b_q_seq_len = torch.ones_like(b_seq_len)
         b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len)

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -569,7 +569,7 @@ def _token_gqa_decode_attention_mtp(
             v_cache=kv_nope,
             qv=q_nope.reshape(-1, self.tp_q_head_num_ * self.mtp_size, self.kv_lora_rank),
             page_table=infer_state.page_table[self.mtp_size - 1 :: self.mtp_size],
-            cache_seqlens=infer_state.b_seq_len[self.mtp_size - 1 :: self.mtp_size],
+            cache_seqlens=infer_state.b_seq_len[self.mtp_size - 1 :: self.mtp_size].contiguous(),
             cu_seqlens_q=infer_state.cu_seqlens_q,
             cu_seqlens_k_new=infer_state.cu_seqlens_k,
             max_seqlen_q=1,
@@ -582,7 +582,7 @@ def _token_gqa_decode_attention_mtp(
             return_softmax_lse=False,
             mtp_step=self.mtp_step,
         )
-        return o_tensor
+        return o_tensor.view(-1, self.tp_q_head_num_, self.kv_lora_rank)
 
     def _token_gqa_decode_attention_flashattention(
         self, q, infer_state: Deepseek2FlashAttentionStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None

lightllm/server/api_start.py

@@ -15,6 +15,7 @@
 from .router.manager import start_router_process
 from lightllm.utils.process_check import is_process_active
 from lightllm.utils.multinode_utils import send_and_receive_node_ip
+from lightllm.common.flash_attn import flash_attn_with_kvcache_mtp
 
 logger = init_logger(__name__)
 
@@ -140,6 +141,9 @@ def normal_or_p_d_start(args):
 
     if args.enable_fa3_mtp:
         assert args.mtp_mode is not None, "enable_fa3_mtp must set mtp_mode"
+        assert (
+            flash_attn_with_kvcache_mtp is not None
+        ), "flash_attn_with_kvcache_mtp is None, please check if you have installed the fa3_mtp kernel"
 
     # 检查GPU数量是否足够
     if args.visual_gpu_ids is None: