refactor(attention): separate backends, runtime and persistent state

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 = infer_state.attention_contexts[attention.group_id]
+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 = infer_state.attention_contexts[attention.group_id]
+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/__init__.py

@@ -0,0 +1,11 @@
+from .api import (
+    AttentionBackend,
+    AttentionBatch,
+    AttentionFamily,
+    AttentionPhase,
+    CacheLayout,
+    ContextCapacity,
+    LayerAttentionBinding,
+    LayerSpec,
+    UnsupportedAttention,
+)

lightllm/common/attention/_environment.py

@@ -0,0 +1,17 @@
+import importlib
+
+from .api.backend import UnsupportedAttention
+
+
+def check_attention_environment(caps, *, provider, supported_platforms, min_device_capability, dependency):
+    if caps.platform not in supported_platforms:
+        raise UnsupportedAttention(f"{provider} does not support platform {caps.platform.value}")
+    if not caps.available:
+        raise UnsupportedAttention(f"device {caps.device} is unavailable")
+    if min_device_capability is not None and (caps.capability is None or caps.capability < min_device_capability):
+        raise UnsupportedAttention(f"requires device capability {min_device_capability}")
+    try:
+        if dependency:
+            importlib.import_module(dependency)
+    except (ImportError, OSError) as exc:
+        raise UnsupportedAttention(str(exc)) from exc

lightllm/common/attention/api/__init__.py

@@ -0,0 +1,19 @@
+"""模型层与 attention 实现之间的契约。
+
+初始化时，LayerSpec / CacheLayout 用于选择后端和划分可共享 context 的层组，
+LayerAttentionBinding 保存本层参数和共享 backend，通过 forward(context, inputs) 执行。
+运行时，AttentionBatch 描述整个批次，Inputs 提供当前层的 tensor 输入；
+backend 专用 context 保存派生 metadata 和 provider wrapper，由 InferStateInfo 持有。
+KV / Conv / SSM 的分配、回收和迁移仍由 AttentionStateManager 管理。
+无持久 KV 的视觉/音频编码器使用 VarlenAttentionOp，直接传入 Q/K/V 和序列边界。
+
+初始化集成见 common.basemodel.attention，运行时入口见 InferStateInfo.prepare_attention。
+"""
+
+from .backend import AttentionBackend, LayerAttentionBinding, UnsupportedAttention
+from .batch import AttentionBatch, ContextCapacity
+from .execution import BackendResources, ExecutionRequirements
+from .graph import GraphBinding
+from .inputs import VarlenAttentionInputs
+from .spec import AttentionFamily, AttentionPhase, CacheLayout, LayerSpec
+from .varlen import VarlenAttentionOp

lightllm/common/attention/api/backend.py

@@ -0,0 +1,91 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Generic, TypeVar
+
+import torch
+
+from .batch import AttentionBatch, ContextCapacity
+from .execution import BackendResources
+from .spec import LayerSpec
+
+Context = TypeVar("Context")
+Inputs = TypeVar("Inputs")
+Output = TypeVar("Output")
+Params = TypeVar("Params")
+
+
+class UnsupportedAttention(ValueError):
+    """当前环境或输入不满足此 attention 实现的要求。"""
+
+
+@dataclass(frozen=True)
+class LayerAttentionBinding(Generic[Params]):
+    """一个逻辑层在 prefill 或 decode 阶段的 attention 执行入口。
+
+    权重和存储准备完成后构造；backend 是同组层共享的实现，params 借用本层的
+    权重和存储视图，存储重分配后须重新构造。group_id 用于取得当前批次的 context；
+    logical_layer_id 是模型层号，不一定等于物理存储下标。
+    """
+
+    group_id: int
+    logical_layer_id: int
+    spec: LayerSpec
+    params: Params
+    backend: "AttentionBackend"
+    alloc_tensor: Callable
+
+    def forward(self, context, inputs, *, out=None):
+        return self.backend.forward(self, context, inputs, out=out)
+
+
+class AttentionBackend(ABC, Generic[Context, Inputs, Output]):
+    """同组层共用的执行实现；每批次变化的 metadata 和 wrapper 放在 context 中。
+
+    Context 的具体类型由 backend 定义，实例由 InferStateInfo 保存，不能跨 backend
+    或并发 microbatch 共用。调用顺序为 allocate_context -> prepare -> update，
+    随后同组各层用各自的 LayerAttentionBinding / Inputs 调用 forward。
+    """
+
+    @abstractmethod
+    def allocate_context(self, capacity: ContextCapacity, *, capture: bool, workspace) -> Context:
+        """创建归属此 backend 的 context，workspace 由模型按 group / microbatch slot 提供。
+
+        capture=True 的 context 随 decode CUDA Graph 保存，在后续 replay 中复用。
+        """
+        pass
+
+    @abstractmethod
+    def prepare(self, context: Context, batch: AttentionBatch) -> None:
+        """在图外绑定批次输入并完成 provider planning，这里允许 CPU 侧准备工作。
+
+        decode replay 时仍会调用；图中引用的 buffer 必须保持地址，只更新内容。
+        """
+        pass
+
+    @abstractmethod
+    def update(self, context: Context) -> None:
+        """在各层 forward 前生成 GPU 侧派生 metadata。
+
+        decode Graph 会捕获这段 GPU 工作；不能依赖 replay 时重新执行 Python，
+        也不能把当前批次的 GPU 值读回 CPU 来决定执行路径。
+        """
+        pass
+
+    @abstractmethod
+    def forward(
+        self, layer: LayerAttentionBinding, context: Context, inputs: Inputs, *, out: torch.Tensor | None = None
+    ) -> Output:
+        """使用已准备的 context 执行一层，允许写入 params 借用的 KV / Conv / SSM 存储。
+
+        out 指定 attention 结果的目标缓冲区；返回 tuple 时只对应首项，
+        其余项保留模型所需的附加输出，例如 GDN 的 gate 输入。
+        """
+        pass
+
+    def resource_requirements(self, capacity: ContextCapacity, *, capture: bool):
+        """返回分配前可计算的资源上界，供共享容量规划和捕获前预留使用。"""
+        return BackendResources()
+
+    def bind_graph(self, context: Context):
+        """为此捕获实例创建可选的 GraphBinding；普通原生 Graph 无需额外绑定。"""
+        return None

lightllm/common/attention/api/batch.py

@@ -0,0 +1,113 @@
+from dataclasses import dataclass, fields
+from typing import Optional
+
+import torch
+
+from .spec import AttentionPhase
+
+
+@dataclass(frozen=True)
+class AttentionBatch:
+    """一次执行的批次描述，tensor 借用自 InferStateInfo / ReqManager。
+
+    b_* 的首维为 batch_size，b1_cu_* 是带起始 0 的 batch_size + 1 个累积偏移。
+    token_num 是当前 query token 数；total_token_num 包含历史 KV 的逻辑 token 数，
+    捕获 decode Graph 时后者可替换为容量上界。frozen 只固定字段引用，tensor 内容仍可更新。
+    """
+
+    phase: AttentionPhase
+    batch_size: int
+    token_num: int
+    total_token_num: int
+    max_q_seq_len: int
+    max_kv_seq_len: int
+    max_cache_len: int
+    b_req_idx: torch.Tensor
+    # KV 长度包含当前输入；query 长度只包含本次实际计算的 token。
+    b_seq_len: torch.Tensor
+    b_q_seq_len: torch.Tensor
+    b1_cu_q_seq_len: torch.Tensor
+    b1_cu_kv_seq_len: torch.Tensor
+    req_to_token_indexs: torch.Tensor
+    hold_req_id: int
+    b_ready_cache_len: Optional[torch.Tensor] = None
+    b_mtp_index: Optional[torch.Tensor] = None
+    b_shared_seq_len: Optional[torch.Tensor] = None
+    b_shared_radix_node_id: Optional[torch.Tensor] = None
+    mem_index: Optional[torch.Tensor] = None
+    # Absolute positions of the query rows in the request KV sequence.  Sparse
+    # MLA indexers use this to expand per-request metadata without reaching
+    # back into InferStateInfo.
+    position_ids: Optional[torch.Tensor] = None
+    b_image_token_end: Optional[torch.Tensor] = None
+
+    @property
+    def device(self):
+        return self.b_req_idx.device
+
+    @property
+    def is_prefill(self):
+        return self.phase is AttentionPhase.PREFILL
+
+    @property
+    def b_q_start_loc(self):
+        return self.b1_cu_q_seq_len[:-1]
+
+    @property
+    def b_kv_start_loc(self):
+        return self.b1_cu_kv_seq_len[:-1]
+
+    def validate(self):
+        if self.batch_size < 0 or self.token_num < 0:
+            raise ValueError("attention batch sizes must be nonnegative")
+        for name in ("b_req_idx", "b_seq_len", "b_q_seq_len"):
+            tensor = getattr(self, name)
+            if tensor.ndim != 1 or tensor.numel() != self.batch_size:
+                raise ValueError(f"{name} must have batch_size elements")
+        if self.b1_cu_q_seq_len.numel() != self.batch_size + 1:
+            raise ValueError("query offsets must have batch_size + 1 elements")
+
+    def copy_inputs(self, source: "AttentionBatch") -> None:
+        """向捕获时的 tensor 原址拷入新批次，标量字段保留捕获配置。
+
+        先校验所有 tensor 再写入，避免校验失败留下半更新的输入。
+        request table 是共享存储，必须保持原地址，不能通过整表拷贝替换。
+        """
+        if (self.phase, self.batch_size, self.token_num) != (source.phase, source.batch_size, source.token_num):
+            raise ValueError("attention batch does not match the captured shape")
+        copies = []
+        for field in fields(self):
+            dst, src = getattr(self, field.name), getattr(source, field.name)
+            if dst is src:
+                continue
+            if isinstance(dst, torch.Tensor) or isinstance(src, torch.Tensor):
+                if not isinstance(dst, torch.Tensor) or not isinstance(src, torch.Tensor):
+                    raise ValueError(f"captured input {field.name} changed presence")
+                if (dst.shape, dst.dtype, dst.device) != (src.shape, src.dtype, src.device):
+                    raise ValueError(f"captured input {field.name} changed shape, dtype or device")
+                if field.name == "req_to_token_indexs" and dst.data_ptr() != src.data_ptr():
+                    raise ValueError("request table storage changed; recapture attention")
+                if dst.data_ptr() != src.data_ptr():
+                    copies.append((dst, src))
+        for dst, src in copies:
+            dst.copy_(src, non_blocking=True)
+
+
+@dataclass(frozen=True)
+class ContextCapacity:
+    """context 的形状契约：batch_size / token_num 要与输入批次逐项一致，序列长度只检查上界。
+
+    runner 负责选择 Graph bucket 和 padding；backend 不自行改变批次形状。
+    """
+
+    phase: AttentionPhase
+    batch_size: int
+    token_num: int
+    max_kv_seq_len: int
+    max_q_seq_len: int
+
+    def check(self, batch: AttentionBatch):
+        if self.phase is not batch.phase or self.batch_size != batch.batch_size or self.token_num != batch.token_num:
+            raise ValueError("attention context capacity does not match the batch shape")
+        if batch.max_kv_seq_len > self.max_kv_seq_len or batch.max_q_seq_len > self.max_q_seq_len:
+            raise ValueError("attention batch exceeds context sequence capacity")