diff --git a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py b/lightllm/common/basemodel/triton_kernel/linear_att_copy.py index eb17507fb..6b2934dae 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att_copy.py +++ b/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, diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py index f65a849df..049701cf8 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/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, diff --git a/lightllm/common/linear_att_cache_manager/config_objs.py b/lightllm/common/linear_att_cache_manager/config_objs.py index b63cd6b0e..98f8e5f05 100644 --- a/lightllm/common/linear_att_cache_manager/config_objs.py +++ b/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 diff --git a/lightllm/common/req_manager.py b/lightllm/common/req_manager.py index 070da7412..c1c1440a6 100644 --- a/lightllm/common/req_manager.py +++ b/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的操作，是否是阻塞的操作。 # 同时，非连续对象的拷贝，可能存在效率问题。 diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index e31b83ffe..7c0790e85 100644 --- a/lightllm/models/qwen3next/model.py +++ b/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): diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index e0a7ebae7..a0e7de010 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/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