perf(pd): reduce qwen3next prefill linear-att state memory

lightllm/common/basemodel/triton_kernel/linear_att_copy.py

@@ -24,7 +24,7 @@ def _copy_linear_att_state_to_kv_buffer(
     cpu_kv_ssm_stride_s,
     cpu_kv_ssm_stride_l,
     cpu_kv_ssm_stride_d,
-    mtp_step,
+    ssm_state_stride,
     gpu_conv_dim,  # number of conv rows
     gpu_conv_tail_dim_bytes,  # bytes copied per conv row; equals the CPU/cache row width
     gpu_ssm_tail_dim,
@@ -50,7 +50,7 @@ def _copy_linear_att_state_to_kv_buffer(
         return
 
     cur_req_idx = tl.load(b_req_idx + cur_batch).to(tl.int64)
-    cur_state_req_idx = (cur_req_idx * (mtp_step + 1)).to(tl.int64)
+    cur_state_req_idx = (cur_req_idx * ssm_state_stride).to(tl.int64)
 
     gpu_conv_base = gpu_conv_ptr + cur_layer * gpu_conv_stride_l + cur_req_idx * gpu_conv_stride_s
     cpu_conv_base = cpu_kv_conv_ptr + big_page_buffer_idx * cpu_kv_conv_stride_s + cur_layer * cpu_kv_conv_stride_l
@@ -85,7 +85,7 @@ def copy_linear_att_state_to_kv_buffer(
     gpu_ssm_state: torch.Tensor,  # [linear_layer_num, req_num * (mtp_step + 1), ...]
     cpu_kv_conv_state: torch.Tensor,  # [buffer_num, linear_layer_num, conv_dim, kernel_size]
     cpu_kv_ssm_state: torch.Tensor,  # [buffer_num, linear_layer_num, ...]
-    mtp_step: int,
+    ssm_state_stride: int,
 ):
     # gpu_conv_state 的后两维可能是不连续的。
     assert len(b_req_idx) == big_page_buffer_ids.shape[0]
@@ -143,7 +143,7 @@ def copy_linear_att_state_to_kv_buffer(
         cpu_kv_ssm_stride_s=cpu_kv_ssm_state.stride(0),
         cpu_kv_ssm_stride_l=cpu_kv_ssm_state.stride(1),
         cpu_kv_ssm_stride_d=cpu_kv_ssm_state.stride(2),
-        mtp_step=mtp_step,
+        ssm_state_stride=ssm_state_stride,
         gpu_conv_dim=gpu_conv_dim,
         gpu_conv_tail_dim_bytes=gpu_conv_tail_dim_bytes,
         gpu_ssm_tail_dim=gpu_ssm_tail_dim,

lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py

@@ -71,6 +71,7 @@ def _free_linear_att_buffers(self):
     def write_to_shm(self, req_manager):
         self.req_to_conv_state = req_manager.req_to_conv_state
         self.req_to_ssm_state = req_manager.req_to_ssm_state
+        self.linear_att_state_mtp_size = req_manager.linear_att_state_mtp_size
         return super().write_to_shm(req_manager)
 
     def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor:
@@ -228,9 +229,7 @@ def read_page_to_req(
         return
 
     def _get_req_state_indexes(self, req_idx: int):
-        mtp_size = get_env_start_args().mtp_step + 1
-        # Conv is one widened slot per request; SSM keeps the historical S+1 block layout.
-        return req_idx, req_idx * mtp_size
+        return req_idx, req_idx * self.mem_manager.linear_att_state_mtp_size
 
     def _write_one_rank(
         self,

lightllm/common/linear_att_cache_manager/config_objs.py

@@ -8,6 +8,12 @@
 logger = init_logger(__name__)
 
 
+def get_linear_att_state_mtp_size(args) -> int:
+    if args.run_mode == "prefill":
+        return 1
+    return args.mtp_step + 1
+
+
 @dataclasses.dataclass
 class LinearAttCacheConfig:
     tp_world_size: int

lightllm/common/req_manager.py

@@ -232,6 +232,9 @@ class ReqManagerForMamba(ReqManager):
     def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig):
         super().__init__(max_request_num, max_sequence_length, mem_manager)
         self.mtp_step = get_env_start_args().mtp_step
+        from lightllm.common.linear_att_cache_manager.config_objs import get_linear_att_state_mtp_size
+
+        self.linear_att_state_mtp_size = get_linear_att_state_mtp_size(get_env_start_args())
         # 因为在mtp的推理中，需要标记每个请求对应的mtp index状态(conv state 和 ssm state)，在mtp对应序列中
         # 的真实位置，所以需要需要一个标记来记录，不然算子无法找到真实的处理起点。
         self.req_to_mtp_state_index = (
@@ -251,12 +254,12 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con
         self.req_to_conv_state = LayerCache(
             size=(max_request_num + 1),
             dtype=self.linear_config.conv_state_dtype,
-            shape=self.linear_config.get_mtp_conv_state_shape(mtp_step=self.mtp_step),
+            shape=self.linear_config.get_mtp_conv_state_shape(mtp_step=self.linear_att_state_mtp_size - 1),
             layer_num=self.linear_config.linear_layer_num,
             device="cuda",
         )
         self.req_to_ssm_state = LayerCache(
-            size=(max_request_num + 1) * (self.mtp_step + 1),
+            size=(max_request_num + 1) * self.linear_att_state_mtp_size,
             dtype=self.linear_config.ssm_state_dtype,
             shape=self.linear_config.get_ssm_state_shape(),
             layer_num=self.linear_config.linear_layer_num,
@@ -266,11 +269,9 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con
 
     def init_linear_att_state(self, req: "InferReq"):
         conv_index = req.req_idx
-        ssm_start = req.req_idx * (self.mtp_step + 1)
+        ssm_start = req.req_idx * self.linear_att_state_mtp_size
         self.req_to_conv_state.buffer[:, conv_index, ...].fill_(0)
-        # #17: zero the FULL (mtp_step + 1)-row SSM block, not just canonical row +0, so a future
-        # first-step verify reading offset>0 after fresh init never hits a never-written row (NaN).
-        self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + (self.mtp_step + 1), ...].fill_(0)
+        self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + self.linear_att_state_mtp_size, ...].fill_(0)
         if self.req_to_mtp_state_index is not None:
             self.req_to_mtp_state_index[req.req_idx] = 0
         return
@@ -291,7 +292,7 @@ def copy_big_page_buffer_to_linear_att_state(self, big_page_buffer_idx: int, req
 
         conv_state, ssm_state = big_page_buffers.get_state_cache(buffer_idx=big_page_buffer_idx)
         conv_dest = req.req_idx
-        ssm_dest = req.req_idx * (self.mtp_step + 1)
+        ssm_dest = req.req_idx * self.linear_att_state_mtp_size
         conv_cache_width = conv_state.shape[-1]
         self.req_to_conv_state.buffer[:, conv_dest, ..., :conv_cache_width] = conv_state
         self.req_to_ssm_state.buffer[:, ssm_dest, ...] = ssm_state
@@ -306,7 +307,7 @@ def copy_small_page_buffer_to_linear_att_state(
             buffer_idx=req.shared_kv_node.small_page_buffer_idx
         )
         conv_dest = req.req_idx
-        ssm_dest = req.req_idx * (self.mtp_step + 1)
+        ssm_dest = req.req_idx * self.linear_att_state_mtp_size
         conv_cache_width = conv_state.shape[-1]
         # TODO 下面这个从 cpu cache 拷贝数据的 gpu的操作，是否是阻塞的操作。
         # 同时，非连续对象的拷贝，可能存在效率问题。

lightllm/models/qwen3next/model.py

@@ -103,6 +103,7 @@ def _init_req_manager(self):
         self.req_manager = ReqManagerForMamba(
             self.max_req_num, create_max_seq_len, None, linear_config=LinearAttCacheConfig.load_from_args()
         )
+        self.mem_manager.linear_att_state_mtp_size = self.req_manager.linear_att_state_mtp_size
         return
 
     def _init_att_backend1(self):

lightllm/server/router/model_infer/infer_batch.py

@@ -402,7 +402,7 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L
                 gpu_ssm_state=self.req_manager.req_to_ssm_state.buffer,
                 cpu_kv_conv_state=self.radix_cache.linear_att_big_page_buffers.conv_state_cache.buffer,
                 cpu_kv_ssm_state=self.radix_cache.linear_att_big_page_buffers.ssm_state_cache.buffer,
-                mtp_step=self.args.mtp_step,
+                ssm_state_stride=self.req_manager.linear_att_state_mtp_size,
             )
 
         assert not self.args.disable_chunked_prefill, "chunked prefill mode must be enabled for linear att mixed model"
@@ -419,7 +419,7 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L
                     )
                     if req.tail_linear_att_small_page_buffer_id is not None:
                         conv_src_idx = req.req_idx
-                        ssm_src_idx = req.req_idx * (self.args.mtp_step + 1)
+                        ssm_src_idx = req.req_idx * self.req_manager.linear_att_state_mtp_size
                         conv_cache_width = self.req_manager.linear_config.get_conv_state_shape()[-1]
                         gpu_conv_state = self.req_manager.req_to_conv_state.buffer[
                             :, conv_src_idx, ..., :conv_cache_width