fix

sufubao · sufubao · commit 550395db59e2 · 2025-12-15T14:34:10.000Z
diff --git a/lightllm/common/basemodel/triton_kernel/alloc_buffer_kernel.py b/lightllm/common/basemodel/triton_kernel/alloc_buffer_kernel.py
@@ -0,0 +1,91 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def alloc_buffer_for_req_kernel(
+    req_index_ptr,  # [num_reqs] - indices of requests to allocate buffers for
+    buffer_indexes_ptr,  # [num_reqs] - buffer indices to assign (from CPU)
+    req_to_buffer_index_ptr,  # [max_request_num + 1] - tensor mapping req_idx to buffer_idx
+    num_reqs,  # number of requests to process
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+    # Mask for valid indices
+    mask = offsets < num_reqs
+
+    # Load request indices and buffer indices
+    req_indices = tl.load(req_index_ptr + offsets, mask=mask, other=0)
+    buffer_indices = tl.load(buffer_indexes_ptr + offsets, mask=mask, other=0)
+
+    # Update req_to_buffer_index[req_indices] = buffer_indices
+    tl.store(req_to_buffer_index_ptr + req_indices, buffer_indices, mask=mask)
+
+
+def alloc_buffer_for_req_triton(
+    req_index: torch.Tensor,  # [num_reqs] int32/int64 tensor on CUDA
+    buffer_indexes: torch.Tensor,  # [num_reqs] int32 tensor (can be CPU or CUDA)
+    req_to_buffer_index: torch.Tensor,  # [max_request_num + 1] int32 tensor on CUDA
+):
+    num_reqs = req_index.shape[0]
+
+    # Ensure inputs are on CUDA
+    if not req_index.is_cuda:
+        req_index = req_index.cuda()
+    if not buffer_indexes.is_cuda:
+        buffer_indexes = buffer_indexes.cuda()
+
+    # Ensure correct dtypes
+    if req_index.dtype not in [torch.int32, torch.int64]:
+        req_index = req_index.to(torch.int32)
+    if buffer_indexes.dtype != torch.int32:
+        buffer_indexes = buffer_indexes.to(torch.int32)
+
+    # Launch kernel
+    BLOCK_SIZE = 256
+    grid = (triton.cdiv(num_reqs, BLOCK_SIZE),)
+
+    alloc_buffer_for_req_kernel[grid](
+        req_index,
+        buffer_indexes,
+        req_to_buffer_index,
+        num_reqs,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+
+# Convenience function that matches the original API
+def alloc_buffer_for_req_wrapper(
+    req_manager,
+    req_index: list,
+    buffer_indexes: torch.Tensor,
+):
+    """
+    Wrapper function to integrate with ReqManagerWithBuffer.
+
+    Usage in ReqManagerWithBuffer:
+        def alloc_buffer_for_req(self, req_index: List[int]):
+            self.req_has_buffer[req_index] = True
+            buffer_indexes = self.mem_manager.alloc_buffer(len(req_index))  # cpu tensor
+            # Replace the next line with Triton kernel
+            # self.req_to_buffer_index[req_index] = buffer_indexes
+            from lightllm.common.basemodel.triton_kernel.alloc_buffer_kernel import alloc_buffer_for_req_triton
+            req_index_tensor = torch.tensor(req_index, dtype=torch.int32, device='cuda')
+            alloc_buffer_for_req_triton(
+                req_index_tensor,
+                buffer_indexes,
+                self.req_has_buffer,
+                self.req_to_buffer_index
+            )
+    """
+    req_index_tensor = torch.tensor(req_index, dtype=torch.int32, device="cuda")
+    alloc_buffer_for_req_triton(
+        req_index_tensor,
+        buffer_indexes,
+        req_manager.req_has_buffer,
+        req_manager.req_to_buffer_index,
+    )
diff --git a/lightllm/common/req_manager.py b/lightllm/common/req_manager.py
@@ -1,5 +1,6 @@
 import torch
 import collections
+from lightllm.common.basemodel.triton_kernel.alloc_buffer_kernel import alloc_buffer_for_req_triton
 from lightllm.utils.log_utils import init_logger
 from .kv_cache_mem_manager import MemoryManager
 from typing import List, Optional
@@ -243,27 +244,32 @@ def gen_cpu_out_token_counter_sampling_params(self, req_objs: List):
 class ReqManagerWithBuffer(ReqManager):
     def __init__(self, max_request_num, max_sequence_length, mem_manager):
         super().__init__(max_request_num, max_sequence_length, mem_manager)
-        self.req_has_buffer = torch.zeros((self.max_request_num + 1), dtype=torch.bool, device="cuda")
         self.req_to_buffer_index = torch.zeros((self.max_request_num + 1), dtype=torch.int32, device="cuda")
         self.req_to_buffer_index[self.HOLD_REQUEST_ID] = self.mem_manager.HOLD_BUFFER_INDEX
 
     @override
     def free(self, free_req_indexes: List[int], free_token_index):
         super().free(free_req_indexes, free_token_index)
-        self.req_has_buffer[free_req_indexes] = False
         self.free_buffer(self.req_to_buffer_index[free_req_indexes])
 
     @override
     def free_all(self):
-        self.req_has_buffer.zero_()
         super().free_all()
         return
 
     def free_buffer(self, free_buffer_indexes: List[int]):
         self.mem_manager.free_buffer(free_buffer_indexes)
         return
 
-    def alloc_buffer_for_req(self, req_index: int):
-        self.req_has_buffer[req_index] = True
-        buffer_indexes = self.mem_manager.alloc_buffer(len(req_index))
-        self.req_to_buffer_index[req_index] = buffer_indexes
+    def alloc_buffer_for_req(self, req_index: torch.Tensor):
+        buffer_indexes = self.mem_manager.alloc_buffer(req_index.shape[0])
+        alloc_buffer_for_req_triton(req_index, buffer_indexes, self.req_to_buffer_index)
+
+    def reset_buffer(self, req_index: torch.Tensor):
+        buffer_indexes = self.req_to_buffer_index[req_index]
+        self.mem_manager.reset_buffer(buffer_indexes)
+        return
+
+    def copy_buffer_from_another_buffer(self, src_buffer_index: int, tgt_req_index: int):
+        self.mem_manager.copy_buffer(src_buffer_index, self.req_to_buffer_index[tgt_req_index])
+        return
diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py
@@ -255,7 +255,7 @@ def _linear_attn(
         assert isinstance(infer_state.mem_manager, Qwen3NextMemoryManager)
 
         input = input.view(-1, infer_cls.embed_dim_)
-        buffer_idx = infer_state.req_manager.req_to_buffer_indexes[infer_state.b_req_idx]
+        buffer_idx = infer_state.req_manager.req_to_buffer_index[infer_state.b_req_idx]
         conv_states, ssm_states = infer_state.mem_manager.get_buffer(self.layer_idx_)
 
         mixed_qkvzba = layer_weight.linear_in_proj.mm(input)
diff --git a/lightllm/models/qwen3next/mem_manager.py b/lightllm/models/qwen3next/mem_manager.py
@@ -112,10 +112,10 @@ def get_buffer(self, layer_index) -> Tuple[torch.Tensor, torch.Tensor]:
         return self.conv_state_mem_manager.buffer[real_layer_index], self.ssm_state_mem_manager.buffer[real_layer_index]
 
     @override
-    def free_buffer(self, free_buffer_indexes: List[int], reset=True):
+    def free_buffer(self, free_buffer_indexes: List[int], reset_to_zero=True):
         # conv_state 和 ssm_state 共享buffer_idx
         self.conv_state_mem_manager.free(free_buffer_indexes)
-        if reset:
+        if reset_to_zero:
             self.conv_state_mem_manager.buffer[:, free_buffer_indexes] = 0
             self.ssm_state_mem_manager.buffer[:, free_buffer_indexes] = 0
 
@@ -130,8 +130,6 @@ def get_buffer_can_use_size(self):
 
     @override
     def copy_buffer(self, src_idx, tgt_idx):
-        assert src_idx is not None and tgt_idx is not None
-        assert src_idx != tgt_idx
         # Use slice operation and in-place copy for better performance
         self.conv_state_mem_manager.buffer[:, tgt_idx].copy_(self.conv_state_mem_manager.buffer[:, src_idx])
         self.ssm_state_mem_manager.buffer[:, tgt_idx].copy_(self.ssm_state_mem_manager.buffer[:, src_idx])
diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py
@@ -15,6 +15,7 @@
 from lightllm.server.core.objs.start_args_type import StartArgs
 from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput
 from lightllm.common.req_manager import ReqManagerWithBuffer
+from lightllm.server.router.model_infer.infer_batch import g_infer_context
 
 logger = init_logger(__name__)
 
@@ -38,7 +39,7 @@ def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch
         # This is required for kernels in qwen3next/triton_kernel/fla/ops/solve_tril.py
         triton.set_allocator(_triton_allocator)
         logger.info("Triton allocator set for Qwen3Next model")
-
+        g_infer_context.use_buffer_manager = True
         super().__init__(kvargs)
 
     @override
diff --git a/lightllm/server/router/dynamic_prompt/hybrid_radix_cache.py b/lightllm/server/router/dynamic_prompt/hybrid_radix_cache.py
@@ -59,6 +59,7 @@ def evict_buffer(self, need_evict_buffer_num, evict_buffer_callback, evict_token
             node = self.evict_buffer_set.pop(0)
             assert node.buffer_idx is not None
             evict_buffer_callback(node.buffer_idx)
+            node.buffer_idx = None
             need_evict_buffer_num -= 1
             # 当一个节点的buffer_idx变为None时，事实上无法在后续进行match，
             # 但当该节点子节点或者引用数不为0时，仍然需要保留， 否则则应该被删除
@@ -73,38 +74,25 @@ def evict_buffer(self, need_evict_buffer_num, evict_buffer_callback, evict_token
         return
 
     def insert_for_hybrid_radix_cache(self, reqs):
-        # 在请求运行途中对prefix cache进行保留，而不是请求被释放时
         from lightllm.server.router.model_infer.infer_batch import g_infer_context
-        from lightllm.common.basemodel.infer_lock import g_infer_state_lock
 
-        # 过滤掉 cur_kv_len 为 0 的请求（新请求还没有生成任何 KV）
-        valid_reqs = [req for req in reqs if req.cur_kv_len > 0]
+        self.free_radix_cache_to_get_enough_buffer(len(reqs))
+        new_buffer_indexes = self.mem_manager.alloc_buffer(len(reqs))
 
-        if len(valid_reqs) == 0:
-            return
+        for i, req in enumerate(reqs):
+            input_token_ids = req.get_input_token_ids()
+            key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu")
+            value = g_infer_context.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].cpu()
+            cur_buffer_idx = g_infer_context.req_manager.req_to_buffer_index[req.req_idx]
 
-        # 确保有足够的空间用于新的 buffer，并在锁保护下完成所有 radix cache 操作
-        g_infer_state_lock.acquire()
-        try:
-            self.free_radix_cache_to_get_enough_buffer(len(valid_reqs))
-            new_buffer_indexes = self.mem_manager.alloc_buffer(len(valid_reqs))
+            # 分配新的 buffer 并复制当前 buffer 的内容
+            self.mem_manager.copy_buffer(cur_buffer_idx, new_buffer_indexes[i])
 
-            for i, req in enumerate(valid_reqs):
-                input_token_ids = req.get_input_token_ids()
-                key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu")
-                value = g_infer_context.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].cpu()
-                cur_buffer_idx = g_infer_context.req_manager.req_to_buffer_indexes[req.req_idx]
-
-                # 分配新的 buffer 并复制当前 buffer 的内容
-                self.mem_manager.copy_buffer(cur_buffer_idx, new_buffer_indexes[i])
-
-                _, new_shared_kv_node = super().insert(key, value)
-                self.dec_node_ref_counter(req.shared_kv_node)
-                self.add_node_ref_counter(new_shared_kv_node)
-                new_shared_kv_node.buffer_idx = new_buffer_indexes[i]
-                req.shared_kv_node = new_shared_kv_node
-        finally:
-            g_infer_state_lock.release()
+            _, new_shared_kv_node = super().insert(key, value)
+            self.dec_node_ref_counter(req.shared_kv_node)
+            self.add_node_ref_counter(new_shared_kv_node)
+            new_shared_kv_node.buffer_idx = new_buffer_indexes[i]
+            req.shared_kv_node = new_shared_kv_node
 
     def match_prefix(self, key, update_refs=False):
         assert len(key) != 0
@@ -184,12 +172,13 @@ def evict(self, need_remove_tokens, evict_buffer_callback, evict_callback):
             node: TreeNode = self.evict_tree_set.pop(0)
             assert (
                 node.ref_counter == 0 and len(node.children) == 0 and node != self.root_node
-            ), "error evict tree node state"
+            ), f"error evict tree node state: {node.ref_counter}, {len(node.children)}"
             num_evicted += len(node.token_mem_index_value)
             evict_callback(node.token_mem_index_value)
             if node.buffer_idx is not None:
                 self.evict_buffer_set.discard(node)
                 evict_buffer_callback(node.buffer_idx)
+                node.buffer_idx = None
             # update total token num
             self.tree_total_tokens_num.arr[0] -= len(node.token_mem_index_value)
             parent_node: TreeNode = node.parent
diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py
@@ -71,6 +71,40 @@ def get_cpu_kv_cache_stream(self) -> torch.cuda.Stream:
             self.cpu_kv_cache_stream = torch.cuda.Stream()
         return self.cpu_kv_cache_stream
 
+    def _maybe_alloc_and_copy_req_buffers(self, req_objs: List["InferReq"]) -> None:
+        """
+        For hybrid/linear-attention models (e.g. Qwen3-Next) we allocate a fixed-size buffer per request.
+        If radix cache hits and the matched node has a buffer, copy that buffer content to the newly
+        allocated buffer for this request.
+        """
+        if not self.use_buffer_manager or not req_objs:
+            return
+
+        if self.radix_cache is not None:
+            # Ensure enough buffer capacity by evicting radix cache buffers if needed.
+            self.radix_cache.free_radix_cache_to_get_enough_buffer(len(req_objs))
+
+        req_idxs = np.array([r.req_idx for r in req_objs], dtype=np.int64)
+        request_indices_gpu = torch.from_numpy(req_idxs).to(device="cuda", dtype=torch.int64)
+        self.req_manager.alloc_buffer_for_req(request_indices_gpu)
+
+        if self.radix_cache is None:
+            return
+
+        # `shared_kv_node` may be None on cache miss; treat it as "no buffer to copy".
+        buffer_idxs = np.array(
+            [None if r.shared_kv_node is None else r.shared_kv_node.buffer_idx for r in req_objs], dtype=object
+        )
+        mask = buffer_idxs == None  # noqa: E711 (intentional elementwise comparison against None)
+        copy_indices = req_idxs[~mask].tolist()
+        if not copy_indices:
+            return
+
+        copy_buffers = buffer_idxs[~mask].tolist()
+        copy_indices_tensor = torch.tensor(copy_indices, device="cuda", dtype=torch.int64)
+        copy_buffers_tensor = torch.tensor(copy_buffers, device="cuda", dtype=torch.int64)
+        self.req_manager.copy_buffer_from_another_buffer(copy_buffers_tensor, copy_indices_tensor)
+
     def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: bool = True) -> List["InferReq"]:
         req_objs = []
         request_ids = []
@@ -109,19 +143,16 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache:
                     slave_req: InferReq = slave_req
                     slave_req.related_master_req = master_req
 
-        # 线性注意力模型为每个请求申请一块Buffer
-        if self.use_buffer_manager and len(request_ids) > 0:
-            if self.radix_cache is not None:
-                self.radix_cache.free_radix_cache_to_get_enough_buffer(len(request_ids))
-            self.req_manager.alloc_buffer_for_req(torch.tensor(request_ids, dtype=torch.int64, device="cpu"))
+        # Hybrid/linear-attention models
+        self._maybe_alloc_and_copy_req_buffers(req_objs)
 
         return req_objs
 
     def free_a_req_mem(self, free_token_index: List, req: "InferReq", free_buffer_index: List = None):
         if self.radix_cache is None:
             free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][0 : req.cur_kv_len])
             if self.use_buffer_manager:
-                free_buffer_index.append(self.req_manager.req_to_buffer_indexs[req.req_idx])
+                free_buffer_index.append(self.req_manager.req_to_buffer_index[req.req_idx])
         else:
             input_token_ids = req.get_input_token_ids()
             key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu")
@@ -131,9 +162,9 @@ def free_a_req_mem(self, free_token_index: List, req: "InferReq", free_buffer_in
             prefix_len, node = self.radix_cache.insert(key, value)
             if self.use_buffer_manager:
                 if node.buffer_idx is None:
-                    node.buffer_idx = self.req_manager.req_to_buffer_indexes[req.req_idx]
+                    node.buffer_idx = self.req_manager.req_to_buffer_index[req.req_idx]
                 else:
-                    free_buffer_index.append(self.req_manager.req_to_buffer_indexes[req.req_idx])
+                    free_buffer_index.append(self.req_manager.req_to_buffer_index[req.req_idx])
 
             old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len
             free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][old_prefix_len:prefix_len])
@@ -179,9 +210,6 @@ def _filter(self, finished_request_ids: List[int]):
             free_token_index = custom_cat(free_token_index)
             self.req_manager.free(free_req_index, free_token_index)
 
-        if self.use_buffer_manager and len(free_buffer_index) != 0:
-            self.req_manager.free_buffer(free_buffer_index)
-
         finished_req_ids_set = set(finished_request_ids)
         self.infer_req_ids = [_id for _id in self.infer_req_ids if _id not in finished_req_ids_set]
 
@@ -208,11 +236,11 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
         if pause_reqs:
             g_infer_state_lock.acquire()
 
-            pause_req_ids = []
+            pause_req_indices = []
             free_token_index = []
             free_buffer_index = []
             for req in pause_reqs:
-                pause_req_ids.append(req.req_id)
+                pause_req_indices.append(req.req_idx)
                 if self.args.diverse_mode:
                     # 发生暂停的时候，需要清除 diverse 模式下的主从关系
                     req.clear_master_slave_state()
@@ -230,8 +258,7 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
                 self.req_manager.free_token(free_token_index)
 
             if self.use_buffer_manager and len(free_buffer_index) != 0:
-                pause_req_ids = torch.tensor(pause_req_ids, dtype=torch.int64, device="cpu")
-                self.req_manager.req_has_buffer[pause_req_ids] = False
+                pause_req_indices = torch.tensor(pause_req_indices, dtype=torch.int64, device="cpu")
                 self.req_manager.free_buffer(free_buffer_index)
 
             g_infer_state_lock.release()
@@ -240,9 +267,7 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
     def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool, can_alloc_token_num: int):
         if paused_reqs:
             g_infer_state_lock.acquire()
-            recover_paused_req_ids = []
             for req in paused_reqs:
-                recover_paused_req_ids.append(req.req_id)
                 prefill_need_token_num = req.get_cur_total_len()
                 if prefill_need_token_num > can_alloc_token_num:
                     break
@@ -253,13 +278,7 @@ def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bo
                     req.shm_req.is_paused = False
                 can_alloc_token_num -= prefill_need_token_num
 
-            if self.use_buffer_manager and len(recover_paused_req_ids) != 0:
-                if self.radix_cache is not None:
-                    self.radix_cache.free_radix_cache_to_get_enough_buffer(len(recover_paused_req_ids))
-                self.req_manager.alloc_buffer_for_req(
-                    torch.tensor(recover_paused_req_ids, dtype=torch.int64, device="cpu")
-                )
-            g_infer_state_lock.release()
+            self._maybe_alloc_and_copy_req_buffers(paused_reqs)
         return
 
     def get_can_alloc_token_num(self):
diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py
