support qwen3next prefix cache

Author: sufubao

lightllm/common/basemodel/layer_weights/hf_load_utils.py

@@ -60,7 +60,7 @@ def load_hf_weights(data_type, weight_dir, pre_post_layer=None, transformer_laye
         transformer_layer_list=transformer_layer_list,
         weight_dir=weight_dir,
     )  # noqa
-    worker = int(os.environ.get("LOADWORKER", 1))
+    worker = int(os.environ.get("LOADWORKER", 16))
     with Pool(worker) as p:
         iterator = p.imap_unordered(partial_func, candidate_files, chunksize=1)
         desc_str = f"pid {os.getpid()} Loading model weights with {worker} workers"

lightllm/common/basemodel/layer_weights/meta_weights/__init__.py

@@ -9,3 +9,4 @@
 from .norm_weight import NormWeight, GEMMANormWeight, TpNormWeight
 from .fused_moe_weight_tp import create_tp_moe_wegiht_obj
 from .fused_moe_weight_ep import FusedMoeWeightEP
+from .parameter_weight import ParameterWeight, TpParameterWeight

lightllm/common/basemodel/layer_weights/meta_weights/parameter_weight.py

@@ -0,0 +1,44 @@
+import torch
+from typing import Dict
+from .base_weight import BaseWeightTpl
+from lightllm.utils.dist_utils import get_current_device_id
+
+
+class ParameterWeight(BaseWeightTpl):
+    def __init__(self, weight_name: str, data_type: torch.dtype, bias_name: str = None):
+        super().__init__()
+        self.weight_name = weight_name
+        self.bias_name = bias_name
+        self.data_type_ = data_type
+        self.weight = None
+        self.bias = None
+
+    def load_hf_weights(self, weights: Dict[str, torch.Tensor]) -> None:
+        if self.weight_name in weights:
+            self.weight = weights[self.weight_name].to(self.data_type_).cuda(get_current_device_id())
+        if self.bias_name in weights:
+            self.bias = weights[self.bias_name].to(self.data_type_).cuda(get_current_device_id())
+
+    def verify_load(self):
+        load_ok = True
+        # Verify weight. The weight must be not None.
+        load_ok = load_ok and self.weight is not None
+        # Verify bias. If bias_name is set, it must be not None.
+        if self.bias_name is not None:
+            load_ok = load_ok and self.bias is not None
+        return load_ok
+
+
+class TpParameterWeight(ParameterWeight):
+    def __init__(self, weight_name: str, data_type: torch.dtype, split_n_embed: int, bias_name: str = None):
+        super().__init__(weight_name, data_type, bias_name)
+        self.split_n_embed = split_n_embed
+
+    def load_hf_weights(self, weights: Dict[str, torch.Tensor]) -> None:
+        start = self.split_n_embed * self.tp_rank_
+        end = self.split_n_embed * (self.tp_rank_ + 1)
+
+        if self.weight_name in weights:
+            self.weight = weights[self.weight_name][start:end].to(self.data_type_).cuda(get_current_device_id())
+        if self.bias_name in weights:
+            self.bias = weights[self.bias_name][start:end].to(self.data_type_).cuda(get_current_device_id())

lightllm/common/mem_manager.py

@@ -69,7 +69,6 @@ def free(self, free_index: Union[torch.Tensor, List[int]]):
         Args:
             free_index (torch.Tensor): _description_
         """
-
         end = self.mark_start
         start = self.mark_start - len(free_index)
         assert start >= 0, f"error free state start: {self.mark_start} free len {len(free_index)}"
@@ -121,7 +120,7 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
         self.dtype = dtype
         # profile the max total token num if the size is None
         self.profile_size(mem_fraction)
-        super().__init__(self.siz, mem_manager_name)
+        super().__init__(self.size, mem_manager_name)
 
         self._init_buffers(
             self.size,

lightllm/models/__init__.py

@@ -8,6 +8,7 @@
 from lightllm.models.qwen2.model import Qwen2TpPartModel
 from lightllm.models.qwen3.model import Qwen3TpPartModel
 from lightllm.models.qwen3_moe.model import Qwen3MOEModel
+from lightllm.models.qwen3next.model import Qwen3NextTpPartModel
 from lightllm.models.chatglm2.model import ChatGlm2TpPartModel
 from lightllm.models.internlm.model import InternlmTpPartModel
 from lightllm.models.stablelm.model import StablelmTpPartModel

lightllm/models/qwen2/model.py

@@ -17,7 +17,7 @@ def __init__(self, kvargs):
 
     def _init_config(self):
         super()._init_config()
-        if self.config["sliding_window"] is None:
+        if self.config.get("sliding_window") is None:
             self.config["sliding_window"] = self.max_total_token_num
         # rename key [SYM: to be confirmed]
         return

lightllm/models/qwen3next/layer_infer/post_layer_infer.py

@@ -0,0 +1,16 @@
+import os
+import torch
+import torch.functional as F
+import torch.distributed as dist
+import numpy as np
+
+from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
+from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
+from lightllm.models.qwen3next.triton_kernel.gemma_rmsnorm import gemma_rmsnorm_forward
+
+class Qwen3NextPostLayerInfer(LlamaPostLayerInfer):
+
+    def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor:
+        out = self.alloc_tensor(input.shape, input.dtype)
+        gemma_rmsnorm_forward(input, layer_weight.final_norm_weight_, self.eps_, out=out)
+        return out