fix

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -27,7 +27,7 @@
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
 from lightllm.common.triton_utils.autotuner import AutotuneLevel
 from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch
-from lightllm.utils.envs_utils import set_model_init_status
+from lightllm.utils.envs_utils import set_model_init_status, enable_diverse_mode_gqa_decode_fast_kernel
 from lightllm.common.triton_utils.autotuner import Autotuner
 from lightllm.utils.infer_utils import post_empty_cache
 
@@ -319,6 +319,15 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s
             mode="constant",
             value=self.mem_manager.HOLD_TOKEN_MEMINDEX,
         )
+        if enable_diverse_mode_gqa_decode_fast_kernel():
+            if new_model_input.b_shared_seq_len is not None:
+                new_model_input.b_shared_seq_len = F.pad(
+                    new_model_input.b_shared_seq_len, (0, padded_batch_size), mode="constant", value=0
+                )
+            if new_model_input.b_mark_shared_group is not None:
+                new_model_input.b_mark_shared_group = F.pad(
+                    new_model_input.b_mark_shared_group, (0, padded_batch_size), mode="constant", value=1
+                )
 
         # 特殊模型，特殊模式的特殊变量的特殊 padding
         if new_model_input.deepseekv3_mtp_draft_input_hiddens is not None:

lightllm/common/basemodel/batch_objs.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass, field
 from typing import Optional
 from typing import List
-from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel
 
 
 @dataclass
@@ -62,7 +62,7 @@ def to_cuda(self):
             self.b_ready_cache_len = self.b_ready_cache_len.cuda(non_blocking=True)
         if self.b_prefill_start_loc is not None:
             self.b_prefill_start_loc = self.b_prefill_start_loc.cuda(non_blocking=True)
-        if not self.is_prefill and get_env_start_args().diverse_mode:
+        if not self.is_prefill and enable_diverse_mode_gqa_decode_fast_kernel():
             batch_size = len(self.b_req_idx)
             if self.b_mark_shared_group is None:
                 self.b_mark_shared_group = torch.ones(size=(batch_size,), dtype=torch.int32, device="cuda")

lightllm/server/api_cli.py

@@ -204,7 +204,8 @@ def make_argument_parser() -> argparse.ArgumentParser:
         type=str,
         default=[],
         nargs="+",
-        help="""Model mode: [triton_int8kv | ppl_int8kv | ppl_fp16 | triton_flashdecoding
+        help="""Model mode: [triton_int8kv | ppl_int8kv | ppl_int8kv_flashdecoding | ppl_int8kv_flashdecoding_diverse
+                        | ppl_fp16 | triton_flashdecoding
                         | triton_gqa_attention | triton_gqa_flashdecoding | triton_fp8kv | offline_calibration_fp8kv
                         | export_fp8kv_calibration
                         triton_flashdecoding mode is for long context, current support llama llama2 qwen;

lightllm/server/router/model_infer/mode_backend/generic_pre_process.py

@@ -4,7 +4,10 @@
 from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context
 from lightllm.common.basemodel.infer_lock import g_infer_state_lock
 from lightllm.common.basemodel.batch_objs import ModelInput
-from lightllm.utils.envs_utils import get_env_start_args, get_diverse_max_batch_shared_group_size
+from lightllm.utils.envs_utils import (
+    enable_diverse_mode_gqa_decode_fast_kernel,
+    get_diverse_max_batch_shared_group_size,
+)
 
 
 def prepare_prefill_inputs(
@@ -93,7 +96,7 @@ def prepare_prefill_inputs(
 
 
 def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[InferReq]]:
-    run_reqs = []
+    run_reqs: List[InferReq] = []
     total_token_num = 0
     max_len_in_batch = 0
     b_req_idx = []
@@ -130,9 +133,18 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
 
     b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu")
     b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu")
-    b_shared_seq_len = torch.tensor(b_shared_seq_len, dtype=torch.int32, device="cpu")
     b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu")
-    if get_env_start_args().diverse_mode:
+
+    if enable_diverse_mode_gqa_decode_fast_kernel():
+        # b_shared_seq_len 和 b_mark_shared_group 只会在 diverse_mode 下的 decode 阶段真正被使用的参数,
+        # 用于记录请求间的共享关系。
+        # 举列说明:
+        # b_shared_seq_len : [10, 10, 10, 11, 11, 11, 11]
+        # b_mark_shared_group: [0, 0, 3, 0, 0, 0, 4]
+        # b_mark_shared_group 中每一个不为0的位置都代表其与前面多少个请求形成一个共享前缀组。属于
+        # 同一个共享前缀组的请求, 其在对应的 b_shared_seq_len 中的内容必然相同。某些模式可以利用这两个
+        # 输入加速算子的运行。
+        b_shared_seq_len = torch.tensor(b_shared_seq_len, dtype=torch.int32, device="cpu")
         b_mark_shared_group = []
         shared_nodes = [req.shared_kv_node for req in run_reqs]
         _current_group = []
@@ -159,6 +171,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In
         assert len(b_mark_shared_group) == len(run_reqs)
         b_mark_shared_group = torch.tensor(b_mark_shared_group, dtype=torch.int32, device="cpu")
     else:
+        b_shared_seq_len = None
         b_mark_shared_group = None
 
     # dynamic prompt cache 准备 token

lightllm/utils/envs_utils.py

@@ -199,3 +199,8 @@ def get_radix_tree_merge_update_delta() -> int:
 @lru_cache(maxsize=None)
 def get_diverse_max_batch_shared_group_size() -> int:
     return int(os.getenv("LIGHTLLM_MAX_BATCH_SHARED_GROUP_SIZE", 4))
+
+
+@lru_cache(maxsize=None)
+def enable_diverse_mode_gqa_decode_fast_kernel() -> bool:
+    return get_env_start_args().diverse_mode and "ppl_int8kv_flashdecoding_diverse" in get_env_start_args().mode