fix: fix bugs

Author: niushengxiao

lightllm/common/basemodel/basemodel.py

@@ -23,6 +23,7 @@
 from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph
 from lightllm.common.quantization import Quantcfg
 from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed
+from lightllm.utils.config_utils import _derive_max_req_total_len_from_model_config
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.dist_utils import get_dp_world_size
 from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type, get_added_mtp_kv_layer_num
@@ -105,8 +106,8 @@ def __init__(self, kvargs):
         self._init_quant()
 
         self._init_weights()
-        self._init_req_manager()
         self._init_mem_manager()
+        self._init_req_manager()
         # 因为类似 qwen3.5 的linear 架构的模型，其 req_manager 会存储运行时使用的大量 linear state
         # 这可能会占用大量的显存，所以，req_manger 中保存的 mem_manger 是mem manager 初始化后再赋值
         self.req_manager.mem_manager = self.mem_manager
@@ -210,6 +211,26 @@ def _init_kv_move_buffer(self):
         if self.run_mode in ["prefill", "decode"]:
             self.mem_manager.alloc_kv_move_buffer(self.mem_manager.size)
 
+    # 推导出的max_req_total_len如果显存预算支持不了，需要进一步截断到可支持的长度
+    def _safe_clamp_auto_max_req_total_len(self):
+        max_total_token_num = self.mem_manager.size
+        if self.max_seq_length is None or self.max_seq_length <= max_total_token_num:
+            return
+
+        # 只截断推导生成的max_req_total_len
+        old_max_req_total_len = self.max_seq_length - 8
+        derived_max_req_total_len = _derive_max_req_total_len_from_model_config(self.weight_dir_)
+        if derived_max_req_total_len is None or old_max_req_total_len != derived_max_req_total_len:
+            return
+
+        supported_max_req_total_len = max(max_total_token_num - 8, 1)
+        self.args.max_req_total_len = supported_max_req_total_len
+        self.max_seq_length = supported_max_req_total_len + 8
+
+        if self.graph_max_len_in_batch == old_max_req_total_len:
+            self.args.graph_max_len_in_batch = min(self.args.graph_max_len_in_batch, supported_max_req_total_len)
+            self.graph_max_len_in_batch = self.args.graph_max_len_in_batch
+
     def _check_mem_size(self):
         self.max_total_token_num = self.mem_manager.size
 
@@ -232,6 +253,7 @@ def _check_mem_size(self):
         return
 
     def _init_req_manager(self):
+        self._safe_clamp_auto_max_req_total_len()
         create_max_seq_len = 0
 
         if self.batch_max_tokens is not None:

lightllm/common/kv_cache_mem_manager/mem_manager.py

@@ -14,7 +14,7 @@
 from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
 from lightllm.distributed.pynccl import PyNcclCommunicator
 from lightllm.utils.dist_utils import get_current_device_id
-from lightllm.utils.config_utils import get_num_key_value_heads
+from lightllm.utils.config_utils import get_num_key_value_heads, get_vocab_size
 from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io
 from lightllm.utils.device_utils import kv_trans_use_p2p
 from lightllm.utils.shm_utils import create_or_link_shm
@@ -61,22 +61,37 @@ def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]:
     def get_cell_size(self):
         return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype)
 
+    def get_req_manager_reserve_bytes(self):
+        args = get_env_start_args()
+        max_request_num = args.running_max_req_size + 8
+        max_sequence_length = max(args.batch_max_tokens or 0, (args.max_req_total_len or 0) + 8)
+        req_state_num = max_request_num + 1
+
+        reserve_bytes = req_state_num * max_sequence_length * torch._utils._element_size(torch.int32)
+        reserve_bytes += req_state_num * 4 * torch._utils._element_size(torch.float32)
+        reserve_bytes += req_state_num * 8 * torch._utils._element_size(torch.int64)
+        if args.penalty_counter_mode == "gpu_counter":
+            reserve_bytes += req_state_num * get_vocab_size(args.model_dir) * torch._utils._element_size(torch.int32)
+        return reserve_bytes
+
     def profile_size(self, mem_fraction):
         if self.size is not None:
             return
 
         torch.cuda.empty_cache()
         world_size = dist.get_world_size()
-
-        available_memory = get_available_gpu_memory(world_size) * mem_fraction
+        available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction)
+        req_manager_reserve_gb = self.get_req_manager_reserve_bytes() / (1024 ** 3)
+        available_memory -= req_manager_reserve_gb
         cell_size = self.get_cell_size()
         self.size = int(available_memory * 1024 ** 3 / cell_size)
         if world_size > 1:
             tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}")
             dist.all_reduce(tensor, op=dist.ReduceOp.MIN)
             self.size = tensor.item()
         logger.info(
-            f"{str(available_memory)} GB space is available after load the model weight\n"
+            f"{str(available_memory)} GB space is available after load the model weight "
+            f"and reserve {req_manager_reserve_gb} GB for req_manager\n"
             f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n"
             f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n"
         )

lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py

@@ -1,8 +1,11 @@
 import torch
+import torch.distributed as dist
 import triton
 from lightllm.utils.log_utils import init_logger
 from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager
 from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory
 from lightllm.common.linear_att_cache_manager import LinearAttCacheConfig, LinearAttCacheManager
 from .operator import LinearAttMemOperator
 from typing import Tuple, Any
@@ -32,6 +35,38 @@ def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]:
         layer_index = layer_index // self.linear_config.full_attention_interval
         return super().get_att_input_params(layer_index)
 
+    def profile_size(self, mem_fraction):
+        if self.size is not None:
+            return
+
+        torch.cuda.empty_cache()
+        args = get_env_start_args()
+        reserve_bytes = self.get_req_manager_reserve_bytes()
+        req_state_num = (args.running_max_req_size + 8 + 1) * (args.mtp_step + 1)
+        reserve_bytes += (
+            req_state_num
+            * self.linear_config.linear_layer_num
+            * (self.linear_config.get_conv_state_bytes_per_layer() + self.linear_config.get_ssm_state_bytes_per_layer())
+        )
+        reserve_gb = reserve_bytes / (1024 ** 3)
+
+        world_size = dist.get_world_size()
+        available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction)
+        available_memory -= reserve_gb
+        cell_size = self.get_cell_size()
+        self.size = max(int(available_memory * 1024 ** 3 / cell_size), 1)
+        if world_size > 1:
+            tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN)
+            self.size = tensor.item()
+        logger.info(
+            f"{str(available_memory)} GB space is available after load the model weight "
+            f"and reserve {reserve_gb} GB for qwen3next req_manager\n"
+            f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n"
+            f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n"
+        )
+        return
+
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
         super()._init_buffers(size, dtype, head_num, head_dim, layer_num)
         # TODO 初始化线性 att 对应的部分 buffer.

lightllm/models/qwen3next/model.py

@@ -89,6 +89,7 @@ def _init_mem_manager(self):
         )
 
     def _init_req_manager(self):
+        self._safe_clamp_auto_max_req_total_len()
         create_max_seq_len = 0
 
         if self.batch_max_tokens is not None:

requirements.txt

@@ -98,4 +98,3 @@ nixl==1.1.0
 xformers==0.0.35
 redis==7.3.0
 litellm>=1.52.0,<1.85
-flash-attn-4[13]==4.0.0b14