fix

Author: sufubao

lightllm/server/api_models.py

@@ -57,7 +57,7 @@ class CompletionRequest(BaseModel):
     # prompt: string or tokens
     prompt: Union[str, List[str], List[int], List[List[int]]]
     suffix: Optional[str] = None
-    max_tokens: Optional[int] = 16
+    max_tokens: Optional[int] = 16000
     temperature: Optional[float] = 1.0
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
@@ -89,7 +89,7 @@ class ChatCompletionRequest(BaseModel):
     stream: Optional[bool] = False
     stream_options: Optional[StreamOptions] = None
     stop: Optional[Union[str, List[str]]] = None
-    max_tokens: Optional[int] = 16
+    max_tokens: Optional[int] = 16000
     presence_penalty: Optional[float] = 0.0
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None

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

@@ -1,7 +1,7 @@
 from .chunked_prefill.impl import ChunkedPrefillBackend
 from .chunked_prefill.impl_for_first_token_constraint_mode import FirstTokenConstraintBackend
 from .chunked_prefill.impl_for_outlines_constraint_mode import OutlinesConstraintBackend
-from .chunked_prefill.impl_for_qwen3next import Qwen3NextBackend
+from .chunked_prefill.impl_for_hybrid_radix_cache import HybridRadixCacheBackend
 from .chunked_prefill.impl_for_return_all_prompt_logprobs import ReturnPromptLogProbBackend
 from .chunked_prefill.impl_for_reward_model import RewardModelBackend
 from .chunked_prefill.impl_for_token_healing import TokenHealingBackend

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

@@ -164,8 +164,8 @@ def init_model(self, kvargs):
         self.model, self.is_multimodal = get_model(model_cfg, model_kvargs)
         self.model: TpPartBaseModel = self.model  # for easy typing
         set_random_seed(2147483647)
-        is_qwen3next = model_cfg.get("model_type", "") == "qwen3_next"
-        radix_cache_class = RadixCache if not is_qwen3next else HybridRadixCache
+        is_hybrid_model = model_cfg.get("model_type", "") in ["qwen3_next"]
+        radix_cache_class = RadixCache if not is_hybrid_model else HybridRadixCache
         self.radix_cache = (
             radix_cache_class(
                 get_unique_server_name(),

lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_hybrid_radix_cache.py

@@ -0,0 +1,54 @@
+import torch
+from .impl import ChunkedPrefillBackend
+from typing import List
+from typing_extensions import override
+from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq
+from lightllm.common.basemodel.infer_lock import g_infer_state_lock
+from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack
+from lightllm.server.router.model_infer.mode_backend.pre import (
+    prepare_prefill_inputs,
+)
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class HybridRadixCacheBackend(ChunkedPrefillBackend):
+    def __init__(self) -> None:
+        super().__init__()
+        logger.info("Using HybridRadixCacheBackend for hybrid attention model.")
+        self.extra_post_req_handle_func = self._handle_hybrid_radix_cache_insert
+
+    @override
+    def init_model(self, kvargs):
+        from lightllm.server.router.dynamic_prompt.hybrid_radix_cache import HybridRadixCache
+        super().init_model(kvargs)
+        assert isinstance(self.radix_cache, HybridRadixCache)
+        return
+
+    def _handle_hybrid_radix_cache_insert(self, req_obj: "InferReq", next_token_id, next_token_logprob):
+        # TODO : add docs
+        if (req_obj.is_multi_chat_req or
+            req_obj.cur_kv_len >= req_obj.get_cur_total_len()):
+            return
+
+        g_infer_state_lock.acquire()
+        input_token_ids = req_obj.get_input_token_ids()
+        key = torch.tensor(input_token_ids[0 : req_obj.cur_kv_len], dtype=torch.int64, device="cpu")
+
+        value = self.model.req_manager.req_to_token_indexs[req_obj.req_idx][: req_obj.cur_kv_len].cpu()
+
+        buffer_idx = self.model.req_manager.req_to_buffer_indexes[req_obj.req_idx].cpu()
+
+        self.radix_cache.free_radix_cache_to_get_enough_token(0, 1)
+
+        new_buffer_idx = self.model.req_manager.mem_manager.alloc_state_cache_buffer(1)[0]
+        self.model.req_manager.mem_manager.copy_state_cache_buffer(buffer_idx, new_buffer_idx)
+        self.model.req_manager.req_to_buffer_indexes[req_obj.req_idx] = new_buffer_idx
+
+        _, new_shared_kv_node = self.radix_cache.insert(key, value, buffer_idx)
+
+        self.radix_cache.dec_node_ref_counter(req_obj.shared_kv_node)
+        self.radix_cache.add_node_ref_counter(new_shared_kv_node)
+        req_obj.shared_kv_node = new_shared_kv_node
+        g_infer_state_lock.release()

lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_qwen3next.py

@@ -7,11 +7,12 @@
 import setproctitle
 from datetime import timedelta
 from typing import Dict, List, Tuple
+from transformers import PretrainedConfig
 from lightllm.server.router.model_infer.mode_backend import (
     ChunkedPrefillBackend,
     FirstTokenConstraintBackend,
     OutlinesConstraintBackend,
-    Qwen3NextBackend,
+    HybridRadixCacheBackend,
     ReturnPromptLogProbBackend,
     RewardModelBackend,
     TokenHealingBackend,
@@ -121,14 +122,17 @@ def init_model(self, kvargs):
         is_outlines_constraint_mode = self.args.output_constraint_mode == "outlines"
         is_xgrammar_constraint_mode = self.args.output_constraint_mode == "xgrammar"
         assert not (is_outlines_constraint_mode and is_xgrammar_constraint_mode), "only one constraint mode can be true"
-        is_qwen3next = True
         is_prefill_node = self.args.run_mode == "prefill"
         is_decode_node = self.args.run_mode == "decode"
         is_nixl_prefill_node = self.args.run_mode == "nixl_prefill"
         is_nixl_decode_node = self.args.run_mode == "nixl_decode"
 
-        if is_qwen3next:
-            self.backend = Qwen3NextBackend()
+        model_cfg, _ = PretrainedConfig.get_config_dict(kvargs["weight_dir"])
+        is_hybrid_model = model_cfg.get("model_type", "") in ["qwen3_next"]
+        use_hybrid_radix_cache = is_hybrid_model and not self.args.disable_dynamic_prompt_cache
+        
+        if use_hybrid_radix_cache:
+            self.backend = HybridRadixCacheBackend()
         elif is_prefill_node:
             if self.args.dp > 1:
                 self.backend = DPChunkedForPrefillNode(self.info_queue, self.mem_queue)