[5/5] refactor(attention): integrate model bindings and state lifecycle

docs/CN/source/models/add_new_model.md

@@ -127,7 +127,21 @@ class TpPartBaseModel:
 | def _init_some_value(self):  | 初始化推理框架会使用的一些成员变量的值                       |
 | def _init_custom(self):      | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值  |
 
-全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先按各模型 `get_req_manager_max_seq_len()` 的最大值创建一份 `ReqManager`，通过 `bind_req_manager()` 绑定给全部模型，再调用各模型的 `prepare_attention_states()` 准备自身 State；随后管理器依次完成 `finish_state_registration()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。资源准备和运行时初始化分别通过这两个单模型公共入口完成，均不接收 draft 列表。
+全部 main/draft 模型构造完成后，`ModeBackend.init_model` 按各模型 `get_req_manager_max_seq_len()` 的最大值创建一份 `ReqManager`，通过 `bind_req_manager()` 绑定给全部模型，再调用各模型的 `prepare_attention_states()` 注册自身 State。管理器完成 `finish_state_registration()` 后，逐模型调用 `prepare_attention_backends()`：按 prefill/decode 的 `LayerSpec` 分组，检查设备、存储布局和执行模式，并运行后端数值探针。随后调用 `reserve_attention_resources(model)` 预留各执行槽位的 workspace，再加载标定并执行 `initialize_storage()`。预留必须在容量估算之前完成，使后端临时显存计入可用容量。
+
+按 main、各 draft 的顺序调用 `init_runtime()`，绑定 HOLD 页并调用 `initialize_layer_attention()` 创建各层的 `LayerAttentionBinding`，然后执行 warmup、padded-request 与 CUDA Graph 初始化；main 还在末尾执行最大长度检查。全部模型完成后才发布就绪标志。每批次的 metadata 和 provider wrapper 由 context 持有，按 `allocate_context → prepare → update → forward` 执行；KV/Conv/SSM 的存储与恢复仍由 State 管理。
+
+层内显式调用 `.forward(...)`。`LayerAttentionBinding` 保存本层的参数和共享 backend；context 来自当前批次：
+
+```python
+attention = self.decode_attention
+context = attention.context(infer_state)
+output = attention.forward(context, DenseInputs(q))
+```
+
+视觉/音频编码器直接传入完整 Q/K/V 时，使用 `api.varlen.VarlenAttentionOp` 的 `forward(inputs, out=...)`，无需创建层 binding 或执行 context。`VarlenAttentionInputs` 描述非因果 self-attention 的变长序列；调用方提供输出缓冲区，算子不保留批次 tensor。通过 `select_varlen_attention_op()` 选择实现，只有声明 `supports_graph` 的实现可参与整段 Graph 捕获。启动参数 `--vit_att_backend` 保持不变。
+
+新增模型通过层上的 `attention_family`、`attention_window(phase)`、`attention_uses_alibi`、`attention_uses_sink` 和必要时的 `create_attention_params()` 描述计算需求与资源视图；模型的 `attention_backend_candidates()` 可指定候选实现。dense/MLA/sparse MLA 使用 `--llm_prefill_att_backend` / `--llm_decode_att_backend`，GDN/KDA 使用单独的 `--llm_prefill_linear_att_backend` / `--llm_decode_linear_att_backend`。显式指定不支持的后端会报错；`auto` 才会继续尝试候选。
 
 公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。两个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。
 

docs/EN/source/models/add_new_model.md

@@ -127,7 +127,21 @@ Commonly used interfaces that need to be inherited and implemented
 | def _init_some_value(self):  | Initialize values of some member variables that the inference framework will use |
 | def _init_custom(self):      | Some personalized initialization of the model itself, such as llama initializing its own Rotary values |
 
-After all main/draft models are constructed, `ModeBackend.init_model` creates one `ReqManager` using the largest `get_req_manager_max_seq_len()` returned by the models, binds it to every model through `bind_req_manager()`, then calls each model's `prepare_attention_states()`, then calls `finish_state_registration()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization each use a single-model entry point; neither takes a draft-model list.
+After all main/draft models are constructed, `ModeBackend.init_model` creates one `ReqManager` using the largest `get_req_manager_max_seq_len()`, binds it to every model through `bind_req_manager()`, and calls each model's `prepare_attention_states()` to register its States. After `finish_state_registration()`, it calls `prepare_attention_backends()` for each model: layers are grouped by their prefill/decode `LayerSpec`, device/layout/execution requirements are checked, and backend numerical probes run. It then calls `reserve_attention_resources(model)` to reserve workspace for each execution slot before loading calibration and calling `initialize_storage()`. Reserving before capacity estimation accounts for backend scratch memory.
+
+`init_runtime()` runs sequentially on the main model and each draft, binding the HOLD page and creating each `LayerAttentionBinding` through `initialize_layer_attention()` before warmup, padded-request setup and CUDA Graph initialization. The main model also checks maximum-length inference. Readiness is published only after every model succeeds. Each batch owns a context containing metadata and provider wrappers, executed through `allocate_context → prepare → update → forward`; States continue to own KV/Conv/SSM storage and restoration.
+
+Layers call `.forward(...)` explicitly. `LayerAttentionBinding` holds the layer parameters and shared backend; the context comes from the current batch:
+
+```python
+attention = self.decode_attention
+context = attention.context(infer_state)
+output = attention.forward(context, DenseInputs(q))
+```
+
+Vision/audio encoders that provide complete Q/K/V use `api.varlen.VarlenAttentionOp.forward(inputs, out=...)` without a layer binding or execution context. `VarlenAttentionInputs` describes variable-length, non-causal self-attention. The caller provides the output buffer, and the operator does not retain batch tensors. Select implementations with `select_varlen_attention_op()`; only implementations declaring `supports_graph` support capture of the full operation. The `--vit_att_backend` option is unchanged.
+
+New model layers describe computation and resource views through `attention_family`, `attention_window(phase)`, `attention_uses_alibi`, `attention_uses_sink`, and, when needed, `create_attention_params()`. The model's `attention_backend_candidates()` can specify candidate implementations. Dense/MLA/sparse MLA use `--llm_prefill_att_backend` / `--llm_decode_att_backend`; GDN/KDA use the separate `--llm_prefill_linear_att_backend` / `--llm_decode_linear_att_backend` options. An unsupported explicit backend raises an error; only `auto` tries the remaining candidates.
 
 The public `prepare_attention_states()` stage creates an empty manager for main and shares it with drafts. The two State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded factories also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none.
 

lightllm/common/attention_state/bounded_att/base_state.py

@@ -2,8 +2,30 @@
 
 from __future__ import annotations
 
+from dataclasses import fields
+
+from lightllm.common.attention.api import UnsupportedAttention
+from lightllm.platform import Platform
+
 from .interface import CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface
 
 
 class BoundedAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface):
     """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。"""
+
+    def capture_snapshot_views(self, req_indices):
+        raise NotImplementedError(f"{type(self).__name__} does not implement capture snapshots")
+
+    def describe_layout(self):
+        """显式布局标识与配置，不用 Python 类名判断物理兼容性。"""
+        return self.storage_layout + (
+            tuple(
+                (field.name, getattr(self.config, field.name))
+                for field in fields(self.config)
+                if field.name not in ("logical_layer_ids", "logical_layer_groups")
+            ),
+        )
+
+    def check_execution(self, caps, execution) -> None:
+        if caps.platform is not Platform.CUDA:
+            raise UnsupportedAttention(f"{self.storage_layout[0]} state operations require CUDA")

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

@@ -1,20 +1,20 @@
 from __future__ import annotations
 
 from functools import lru_cache
-from typing import TYPE_CHECKING, List
 from math import prod
+from typing import TYPE_CHECKING, List
 
 import torch
 
 from lightllm.utils.log_utils import init_logger
 
-from ...utils import deepcopy_result, dtype_name
 from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion
+from ...utils import deepcopy_result, dtype_name
 from ..base_state import BoundedAttentionState
 from ..conv_ssm_types import ConvSSMStateConfig
-from .gdn_pages_pool import GDNPagesPool
-from ..transfer_utils import bounded_att_buffer_copies
 from ..copy_kernel import copy_bounded_att_tensors
+from ..transfer_utils import bounded_att_buffer_copies
+from .gdn_pages_pool import GDNPagesPool
 
 if TYPE_CHECKING:
     from lightllm.common.attention_state.types import PageKind
@@ -26,6 +26,8 @@
 class GDNState(BoundedAttentionState):
     """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。"""
 
+    storage_layout = ("gdn_conv_ssm", 1)
+
     big_pages_pool: GDNPagesPool | None
     small_pages_pool: GDNPagesPool | None
 
@@ -120,6 +122,15 @@ def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Ten
         storage_layer_index = self._logical_to_storage_layer[logical_layer_id]
         return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index]
 
+    def capture_snapshot_views(self, req_indices):
+        views = []
+        for req in req_indices:
+            start = req * (self.mtp_step + 1)
+            views.extend((self.conv_state[:, req], self.ssm_state[:, start : start + self.mtp_step + 1]))
+            if self.mtp_state_index is not None:
+                views.append(self.mtp_state_index[req : req + 1])
+        return views
+
     def reset_bounded_att_state(self, req_indices: List[int]) -> None:
         """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。"""
         assert isinstance(req_indices, list), "request indices must be a list"

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

@@ -117,6 +117,29 @@ def __init__(
         self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device)
         self.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device)
 
+    @property
+    def storage_layout(self):
+        return ("gdn_replay_conv_ssm", 1, self.capacity, self.verify_width, self.compute_mode, self.gate_lower_bound)
+
+    def capture_snapshot_views(self, req_indices):
+        views = [self.tile_sync_mark]
+        for req in req_indices:
+            views.extend(
+                tensor[:, req]
+                for tensor in (
+                    self.conv_state,
+                    self.ssm_state,
+                    self.keys,
+                    self.deltas,
+                    self.gates,
+                    self.raw_keys,
+                    self.raw_values,
+                    self.betas,
+                )
+            )
+            views.extend((self.cursors[req : req + 1], self.mtp_state_index[req : req + 1]))
+        return views
+
     def get_state_type(self) -> str:
         """返回算法名和 Conv/SSM 存储精度。"""
         return (

lightllm/common/attention_state/bounded_att/interface.py

@@ -8,9 +8,10 @@
 
 import torch
 
-from ..utils import deepcopy_result
 from lightllm.common.attention_state.types import PageKind, StateTransferLayout, TransferRegion
 
+from ..utils import deepcopy_result
+
 
 class CoreInterface(ABC):
     """有界 attention State 的资源、层视图和运行态核心能力。"""

lightllm/common/attention_state/bounded_att/kda/kda_state.py

@@ -1,20 +1,20 @@
 from __future__ import annotations
 
 from functools import lru_cache
-from typing import TYPE_CHECKING, List
 from math import prod
+from typing import TYPE_CHECKING, List
 
 import torch
 
 from lightllm.utils.log_utils import init_logger
 
-from ...utils import deepcopy_result, dtype_name
 from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion
+from ...utils import deepcopy_result, dtype_name
 from ..base_state import BoundedAttentionState
 from ..conv_ssm_types import ConvSSMStateConfig
-from .kda_pages_pool import KDAPagesPool
-from ..transfer_utils import bounded_att_buffer_copies
 from ..copy_kernel import copy_bounded_att_tensors
+from ..transfer_utils import bounded_att_buffer_copies
+from .kda_pages_pool import KDAPagesPool
 
 if TYPE_CHECKING:
     from lightllm.common.attention_state.types import PageKind
@@ -26,6 +26,8 @@
 class KDAState(BoundedAttentionState):
     """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。"""
 
+    storage_layout = ("kda_conv_ssm", 1)
+
     big_pages_pool: KDAPagesPool | None
     small_pages_pool: KDAPagesPool | None
 
@@ -120,6 +122,15 @@ def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Ten
         storage_layer_index = self._logical_to_storage_layer[logical_layer_id]
         return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index]
 
+    def capture_snapshot_views(self, req_indices):
+        views = []
+        for req in req_indices:
+            start = req * (self.mtp_step + 1)
+            views.extend((self.conv_state[:, req], self.ssm_state[:, start : start + self.mtp_step + 1]))
+            if self.mtp_state_index is not None:
+                views.append(self.mtp_state_index[req : req + 1])
+        return views
+
     def reset_bounded_att_state(self, req_indices: List[int]) -> None:
         """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。"""
         assert isinstance(req_indices, list), "request indices must be a list"

lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py

@@ -117,6 +117,29 @@ def __init__(
         self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device)
         self.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device)
 
+    @property
+    def storage_layout(self):
+        return ("kda_replay_conv_ssm", 1, self.capacity, self.verify_width, self.compute_mode, self.gate_lower_bound)
+
+    def capture_snapshot_views(self, req_indices):
+        views = [self.tile_sync_mark]
+        for req in req_indices:
+            views.extend(
+                tensor[:, req]
+                for tensor in (
+                    self.conv_state,
+                    self.ssm_state,
+                    self.keys,
+                    self.deltas,
+                    self.gates,
+                    self.raw_keys,
+                    self.raw_values,
+                    self.betas,
+                )
+            )
+            views.extend((self.cursors[req : req + 1], self.mtp_state_index[req : req + 1]))
+        return views
+
     def get_state_type(self) -> str:
         """返回算法名和 Conv/SSM 存储精度。"""
         return (

lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py

@@ -2,15 +2,14 @@
 
 from __future__ import annotations
 
-from functools import lru_cache
-
 from dataclasses import dataclass
+from functools import lru_cache
 from math import prod
 
 import torch
 
-from ...utils import deepcopy_result
 from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion
+from ...utils import deepcopy_result
 from ..base_state import BoundedAttentionState
 from ..copy_kernel import copy_bounded_att_tensors
 from ..transfer_utils import bounded_att_buffer_copies
@@ -63,6 +62,8 @@ def __post_init__(self) -> None:
 class FP8E8M0SlidingWindowState(BoundedAttentionState):
     """仅保存近期原始 KV；prefill workspace 由计算 backend 管理。"""
 
+    storage_layout = ("swa_fp8_e8m0_ring", 1)
+
     # ========================================================================
     # 一、资源定义与初始化
     # ========================================================================
@@ -143,6 +144,9 @@ def get_layer_view(self, logical_layer_id: int) -> torch.Tensor:
         """返回该层的 [request, ring_token, packed_byte] 视图。"""
         return self.swa_buffer[self._logical_to_storage_layer[logical_layer_id]]
 
+    def capture_snapshot_views(self, req_indices):
+        return [self.swa_buffer[:, req] for req in req_indices]
+
     def reset_bounded_att_state(self, req_indices: list[int]) -> None:
         for req_idx in req_indices:
             self.swa_buffer[:, req_idx].zero_()

lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py

@@ -2,18 +2,17 @@
 
 from __future__ import annotations
 
-from functools import lru_cache
-
 from dataclasses import dataclass
+from functools import lru_cache
 from math import prod
 from typing import Literal
 
 import torch
 
 from lightllm.utils.envs_utils import get_env_start_args
 
-from ...utils import deepcopy_result
 from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion
+from ...utils import deepcopy_result
 from ..base_state import BoundedAttentionState
 from ..copy_kernel import copy_bounded_att_tensors
 from ..transfer_utils import bounded_att_buffer_copies
@@ -79,6 +78,8 @@ class WindowCompressorState(BoundedAttentionState):
     无重叠的窗口在组边界不需要快照，因此不创建大小页池，也不占 CPU 联合页字节。
     """
 
+    storage_layout = ("window_compressor_ring", 1)
+
     # ========================================================================
     # 一、资源定义与初始化
     # ========================================================================
@@ -162,6 +163,9 @@ def get_layer_view(self, logical_layer_id: int) -> torch.Tensor:
         """返回该层 [request, window+MTP, candidate|gate] 的计算视图。"""
         return self.buffer[self._logical_to_storage_layer[logical_layer_id]]
 
+    def capture_snapshot_views(self, req_indices):
+        return [self.buffer[:, req] for req in req_indices]
+
     def reset_bounded_att_state(self, req_indices: list[int]) -> None:
         for req_idx in req_indices:
             self.buffer[:, req_idx].zero_()