[1/5] refactor(attention): define execution contracts and device workspace

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,102 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Generic, Protocol, 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 实现的要求。"""
+
+
+class AttentionExecution(Protocol):
+    def attention_context(self, binding: "LayerAttentionBinding[Params, Context, Inputs, Output]") -> Context:
+        ...
+
+
+@dataclass(frozen=True)
+class LayerAttentionBinding(Generic[Params, Context, Inputs, Output]):
+    """一个逻辑层在 prefill 或 decode 阶段的 attention 执行入口。
+
+    权重和存储准备完成后构造；backend 是同组层共享的实现，params 借用本层的
+    权重和存储视图，存储重分配后须重新构造。group_id 用于取得当前批次的 context；
+    logical_layer_id 是模型层号，不一定等于物理存储下标。
+    """
+
+    group_id: int
+    logical_layer_id: int
+    spec: LayerSpec
+    params: Params
+    backend: "AttentionBackend[Context, Inputs, Output]"
+    alloc_tensor: Callable
+
+    def context(self, execution: AttentionExecution) -> Context:
+        return execution.attention_context(self)
+
+    def forward(self, context: Context, inputs: Inputs, *, out: torch.Tensor | None = None) -> Output:
+        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 是调用方拥有的目标缓冲区，dtype、device 必须与输出一致。
+        provider 直接写入时可返回其中有效部分的 view（例如去掉 head padding）；
+        返回独立 tensor 时元素数必须一致，由基类按 out 的形状复制。
+        provider 可直接写入，或返回临时结果由基类复制；不承诺零分配。返回 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")

lightllm/common/attention/api/execution.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import torch
+
+from .spec import AttentionPhase
+
+
+@dataclass(frozen=True)
+class ExecutionRequirements:
+    """启动时确定的执行约束，与 LayerSpec 的计算语义分开保存。"""
+
+    graph_phases: frozenset[AttentionPhase] = frozenset()
+    microbatch_count: int = 1
+    max_batch_size: int = 1
+    max_seq_len: int = 1
+    max_verify_tokens: int = 1
+    tp_world_size: int = 1
+    expert_parallel: bool = False
+    transfer_paths: frozenset[str] = frozenset()
+
+    @classmethod
+    def from_model(cls, model):
+        args = model.args
+        phases = set()
+        if not args.disable_cudagraph:
+            phases.add(AttentionPhase.DECODE)
+        if args.enable_prefill_cudagraph:
+            phases.add(AttentionPhase.PREFILL)
+        transfers = set()
+        if args.enable_cpu_cache:
+            transfers.add("cpu")
+        if args.enable_dp_prompt_cache_fetch:
+            transfers.add("dp")
+        if args.run_mode in ("prefill", "decode"):
+            transfers.add("pd")
+        return cls(
+            frozenset(phases),
+            2 if args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap else 1,
+            getattr(model, "max_req_num", 1),
+            max(getattr(model, "max_seq_length", 1), getattr(model, "graph_max_len_in_batch", 1)),
+            model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model) + 1,
+            model.tp_world_size_,
+            args.enable_ep_moe,
+            frozenset(transfers),
+        )
+
+
+class WorkspaceRequirement(NamedTuple):
+    name: str
+    elements: int
+    dtype: torch.dtype
+
+
+@dataclass(frozen=True)
+class BackendResources:
+    """捕获前的资源上界；workspace 项为 (名称, 元素数, dtype)。
+
+    同一执行槽位按名称取最大容量，并发槽位分别分配。主机长度快照只在需要
+    它的实现中创建；持久存储的初始化由 State 执行，不能在 forward 中反复清零。
+    """
+
+    workspace: tuple[WorkspaceRequirement, ...] = ()
+    needs_host_lengths: bool = False
+    zero_initialize_cache: bool = False

lightllm/common/attention/api/graph.py

@@ -0,0 +1,24 @@
+from typing import Protocol
+
+
+class GraphBinding(Protocol):
+    """一个捕获实例的 provider 资源，不能挂在跨批次共享的 Backend 上。
+
+    after_submit 在原生 Graph 提交后立即调用，不能先等待 Graph 完成；某些设备
+    的 Graph 正在等待此方法提交的任务更新。返回的事件须覆盖所有辅助流工作。
+    wait_for_update 只等待上次对主机 metadata 的使用，保证 prepare 可以安全改写。
+    close 在相关设备工作完成后、原生 Graph 销毁前调用，并须可重复执行。
+    """
+
+    # Native Graph-only providers inherit these no-ops; auxiliary work overrides them.
+    def wait_for_update(self) -> None:
+        pass
+
+    def before_replay(self) -> None:
+        pass
+
+    def after_submit(self) -> tuple:
+        return ()
+
+    def close(self) -> None:
+        pass