feat: add ReplaySSM state modes for GDN and KDA

.cursor/rules/named-args-and-types.mdc

@@ -0,0 +1,12 @@
+---
+description: 新代码的函数调用使用参数名，参数带类型标注
+alwaysApply: true
+---
+
+# 新代码的调用和类型
+
+新写的函数调用使用 `param_name=value`。不要用 `*args`、`*activations`、`**kwargs` 把一组参数展开传进去。
+
+新写的函数参数带类型标注。
+
+修改旧代码时，原有调用和原有签名保持不动，避免把风格调整混进 review。这次改动里新加的函数、新加的调用按上面两条写。

.cursor/rules/no-test-only-production-api.mdc

@@ -0,0 +1,10 @@
+---
+description: 测试专用的实现和接口只放在测试里，不留在正式代码中
+alwaysApply: true
+---
+
+# 测试代码不进正式实现
+
+正式路径不保留只为测试、benchmark 或对照方便而存在的函数、参数和分支。
+
+需要这类行为时，在测试或 benchmark 自己实现，并调用正式接口。正式接口的参数对应真实推理会传入的数据，不为测试场景增加可选参数。

lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py

@@ -13,8 +13,24 @@ def create_gdn_state(
     request_capacity: int,
     mtp_step: int = 0,
     device: str | torch.device = "cuda",
+    ssm_state_mode: str = "native",
+    ssm_replay_config: dict | None = None,
+    gate_lower_bound: float = -5.0,
 ) -> GDNState:
-    """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。"""
+    """创建 GDN 请求运行态。replay 模式使用独立的 ReplayGDNState。"""
+    if ssm_state_mode == "replay":
+        from .replay_gdn_state import ReplayGDNState
+
+        return ReplayGDNState(
+            config=config,
+            request_capacity=request_capacity,
+            mtp_step=mtp_step,
+            device=device,
+            ssm_replay_config=ssm_replay_config,
+            gate_lower_bound=gate_lower_bound,
+        )
+    if ssm_state_mode != "native":
+        raise ValueError(f"unknown GDN state mode: {ssm_state_mode}")
     return GDNState(
         config=config,
         request_capacity=request_capacity,

lightllm/common/attention_state/bounded_att/gdn/gdn_state.py

@@ -46,6 +46,7 @@ def __init__(
         self.state_name: str | None = None  # 由 AttentionStateManager.register_state() 绑定时设置。
         self.request_capacity = request_capacity  # 包含预留的 HOLD 请求槽位。
         self.mtp_step = mtp_step
+        self.mtp_enabled = mtp_step > 0
         self._logical_to_storage_layer = {
             logical_layer_id: storage_layer_index
             for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids)
@@ -67,7 +68,7 @@ def __init__(
             dtype=config.ssm_dtype,
             device=device,
         )
-        self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None
+        self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device)
         self.big_pages_pool = None
         self.small_pages_pool = None
 
@@ -126,7 +127,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None:
             self.conv_state[:, req_idx].zero_()
             start = req_idx * (self.mtp_step + 1)
             self.ssm_state[:, start : start + self.mtp_step + 1].zero_()
-            if self.mtp_state_index is not None:
+            if self.mtp_enabled:
                 self.mtp_state_index[req_idx].fill_(0)
 
     def save_bounded_att_runtime_to_pages(
@@ -170,18 +171,19 @@ def restore_bounded_att_runtime_from_pages(
             ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)]))
         copy_bounded_att_tensors(copies=conv_copies)
         copy_bounded_att_tensors(copies=ssm_copies)
-        if self.mtp_state_index is not None:
+        if self.mtp_enabled:
             for req_idx in req_indices:
                 self.mtp_state_index[req_idx].fill_(0)
 
     def commit_bounded_att_state(
         self,
         req_indices: torch.Tensor,
         accepted_steps: torch.Tensor,
+        req_verify_width: torch.Tensor,
     ) -> None:
         """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。"""
         assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors"
-        if self.mtp_state_index is not None:
+        if self.mtp_enabled:
             self.mtp_state_index[req_indices] = accepted_steps - 1
 
     # ========================================================================
@@ -320,7 +322,7 @@ def copy_runtime_state_from_dp(
             )
         copy_bounded_att_tensors(copies=conv_copies)
         copy_bounded_att_tensors(copies=ssm_copies)
-        if self.mtp_state_index is not None:
+        if self.mtp_enabled:
             for dst_req in dst_req_indices:
                 self.mtp_state_index[dst_req].fill_(0)
 
@@ -375,7 +377,7 @@ def read_from_pd_page(
             write=False,
         ):
             copy_bounded_att_tensors(copies=copies)
-        if self.mtp_state_index is not None:
+        if self.mtp_enabled:
             self.mtp_state_index[req_idx].fill_(0)
 
     # ========================================================================