From 042280e9383c92fe635615e65d1a2df8720fc0ff Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 11:53:40 +0800 Subject: [PATCH 1/9] refactor(attention): define execution contracts and device workspace --- lightllm/common/attention/__init__.py | 11 + lightllm/common/attention/_environment.py | 17 ++ lightllm/common/attention/api/__init__.py | 40 +++ lightllm/common/attention/api/backend.py | 145 +++++++++ lightllm/common/attention/api/batch.py | 113 +++++++ lightllm/common/attention/api/execution.py | 36 +++ lightllm/common/attention/api/graph.py | 24 ++ lightllm/common/attention/api/inputs.py | 171 +++++++++++ lightllm/common/attention/api/spec.py | 75 +++++ lightllm/common/attention/api/varlen.py | 54 ++++ lightllm/common/attention/api/workspace.py | 106 +++++++ .../common/attention/backends/__init__.py | 0 .../attention/backends/_shared/__init__.py | 0 .../backends/_shared/linear/__init__.py | 0 .../backends/_shared/prefill/__init__.py | 0 .../common/attention/backends/_workspace.py | 3 + lightllm/common/attention/backends/base.py | 217 ++++++++++++++ .../attention/backends/dense/__init__.py | 0 .../attention/backends/dense/fa3/__init__.py | 0 .../backends/dense/flashinfer/__init__.py | 0 .../attention/backends/dense/sdpa/__init__.py | 0 .../backends/dense/triton/__init__.py | 0 .../backends/dense/triton/kernels/__init__.py | 0 .../dense/triton/kernels/decode/__init__.py | 0 .../triton/kernels/decode/gqa/__init__.py | 0 .../kernels/decode/gqa/flash/__init__.py | 0 .../decode/gqa/mtp_diverse/__init__.py | 0 .../triton/kernels/decode/int4kv/__init__.py | 0 .../triton/kernels/decode/int8kv/__init__.py | 0 .../kernels/decode/int8kv/normal/__init__.py | 0 .../triton/kernels/decode/mha/__init__.py | 0 .../kernels/decode/mha/flash/__init__.py | 0 .../decode/mha/stage3_decode_att/__init__.py | 0 .../dense/triton/kernels/prefill/__init__.py | 0 .../backends/dense/xformers/__init__.py | 0 .../common/attention/backends/gdn/__init__.py | 0 .../backends/gdn/flashqla/__init__.py | 0 .../attention/backends/gdn/replay/__init__.py | 0 .../attention/backends/gdn/triton/__init__.py | 0 .../backends/gdn/triton/kernels/__init__.py | 0 .../common/attention/backends/kda/__init__.py | 0 .../attention/backends/kda/replay/__init__.py | 0 .../attention/backends/kda/triton/__init__.py | 0 .../backends/kda/triton/kernels/__init__.py | 0 .../common/attention/backends/mla/__init__.py | 0 .../attention/backends/mla/fa3/__init__.py | 0 .../backends/mla/flashinfer/__init__.py | 0 .../attention/backends/mla/triton/__init__.py | 0 .../backends/mla/triton/kernels/__init__.py | 0 .../mla/triton/kernels/decode/__init__.py | 0 .../mla/triton/kernels/prefill/__init__.py | 0 .../attention/backends/sparse_mla/__init__.py | 0 .../backends/sparse_mla/flashmla/__init__.py | 0 .../common/basemodel/attention/__init__.py | 36 ++- lightllm/platform/__init__.py | 72 +++++ lightllm/platform/api.py | 96 ++++++ lightllm/platform/cuda.py | 17 ++ unit_tests/common/attention/helpers.py | 283 ++++++++++++++++++ unit_tests/common/attention/test_contracts.py | 139 +++++++++ .../common/attention/test_platform_runtime.py | 63 ++++ unit_tests/common/attention/test_workspace.py | 60 ++++ unit_tests/conftest.py | 1 + 62 files changed, 1774 insertions(+), 5 deletions(-) create mode 100644 lightllm/common/attention/__init__.py create mode 100644 lightllm/common/attention/_environment.py create mode 100644 lightllm/common/attention/api/__init__.py create mode 100644 lightllm/common/attention/api/backend.py create mode 100644 lightllm/common/attention/api/batch.py create mode 100644 lightllm/common/attention/api/execution.py create mode 100644 lightllm/common/attention/api/graph.py create mode 100644 lightllm/common/attention/api/inputs.py create mode 100644 lightllm/common/attention/api/spec.py create mode 100644 lightllm/common/attention/api/varlen.py create mode 100644 lightllm/common/attention/api/workspace.py create mode 100644 lightllm/common/attention/backends/__init__.py create mode 100644 lightllm/common/attention/backends/_shared/__init__.py create mode 100644 lightllm/common/attention/backends/_shared/linear/__init__.py create mode 100644 lightllm/common/attention/backends/_shared/prefill/__init__.py create mode 100644 lightllm/common/attention/backends/_workspace.py create mode 100644 lightllm/common/attention/backends/base.py create mode 100644 lightllm/common/attention/backends/dense/__init__.py create mode 100644 lightllm/common/attention/backends/dense/fa3/__init__.py create mode 100644 lightllm/common/attention/backends/dense/flashinfer/__init__.py create mode 100644 lightllm/common/attention/backends/dense/sdpa/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/mha/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/__init__.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/prefill/__init__.py create mode 100644 lightllm/common/attention/backends/dense/xformers/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/flashqla/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/replay/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/triton/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/triton/kernels/__init__.py create mode 100644 lightllm/common/attention/backends/kda/__init__.py create mode 100644 lightllm/common/attention/backends/kda/replay/__init__.py create mode 100644 lightllm/common/attention/backends/kda/triton/__init__.py create mode 100644 lightllm/common/attention/backends/kda/triton/kernels/__init__.py create mode 100644 lightllm/common/attention/backends/mla/__init__.py create mode 100644 lightllm/common/attention/backends/mla/fa3/__init__.py create mode 100644 lightllm/common/attention/backends/mla/flashinfer/__init__.py create mode 100644 lightllm/common/attention/backends/mla/triton/__init__.py create mode 100644 lightllm/common/attention/backends/mla/triton/kernels/__init__.py create mode 100644 lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py create mode 100644 lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py create mode 100644 lightllm/common/attention/backends/sparse_mla/__init__.py create mode 100644 lightllm/common/attention/backends/sparse_mla/flashmla/__init__.py create mode 100644 lightllm/platform/__init__.py create mode 100644 lightllm/platform/api.py create mode 100644 lightllm/platform/cuda.py create mode 100644 unit_tests/common/attention/helpers.py create mode 100644 unit_tests/common/attention/test_contracts.py create mode 100644 unit_tests/common/attention/test_platform_runtime.py create mode 100644 unit_tests/common/attention/test_workspace.py diff --git a/lightllm/common/attention/__init__.py b/lightllm/common/attention/__init__.py new file mode 100644 index 0000000000..c3b355bdbb --- /dev/null +++ b/lightllm/common/attention/__init__.py @@ -0,0 +1,11 @@ +from .api import ( + AttentionBackend, + AttentionBatch, + AttentionFamily, + AttentionPhase, + CacheLayout, + ContextCapacity, + LayerAttentionBinding, + LayerSpec, + UnsupportedAttention, +) diff --git a/lightllm/common/attention/_environment.py b/lightllm/common/attention/_environment.py new file mode 100644 index 0000000000..85fdec4981 --- /dev/null +++ b/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 diff --git a/lightllm/common/attention/api/__init__.py b/lightllm/common/attention/api/__init__.py new file mode 100644 index 0000000000..87e0a9786f --- /dev/null +++ b/lightllm/common/attention/api/__init__.py @@ -0,0 +1,40 @@ +"""模型层与 attention 实现之间的契约。 + +初始化时，LayerSpec / CacheLayout 用于选择后端和划分可共享 context 的层组， +LayerAttentionBinding 保存本层参数和共享 backend，通过 forward(execution, inputs) 执行。 +运行时，AttentionBatch 描述整个批次，Inputs 提供当前层的 tensor 输入； +backend 专用 context 保存派生 metadata 和 provider wrapper，由 InferStateInfo.attention_runtime 持有。 +KV / Conv / SSM 的分配、回收和迁移仍由 AttentionStateManager 管理。 +无持久 KV 的视觉/音频编码器使用 VarlenAttentionOp，直接传入 Q/K/V 和序列边界。 + +初始化集成见 common.basemodel.attention，运行时入口见 InferStateInfo.prepare_attention。 +""" + +from .backend import ( + AttentionBackend, + AttentionExecution, + AttentionExecutionContext, + LayerAttentionBinding, + UnsupportedAttention, +) +from .batch import AttentionBatch, ContextCapacity +from .execution import BackendResources, ExecutionRequirements +from .graph import GraphBinding +from .inputs import ( + DenseInputs, + DenseParams, + GDNInputs, + GDNOutput, + GDNParams, + KDAInputs, + KDAParams, + MLADecodeInputs, + MLAParams, + MLAPrefillInputs, + SparseMLADecodeInputs, + SparseMLAPrefillInputs, + VarlenAttentionInputs, +) +from .spec import AttentionFamily, AttentionPhase, CacheLayout, LayerSpec +from .varlen import VarlenAttentionOp +from .workspace import AttentionWorkspace diff --git a/lightllm/common/attention/api/backend.py b/lightllm/common/attention/api/backend.py new file mode 100644 index 0000000000..c070732772 --- /dev/null +++ b/lightllm/common/attention/api/backend.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Callable, Generic, Protocol, TypeVar + +import torch + +from lightllm.platform import DeviceCaps + +from .batch import AttentionBatch, ContextCapacity +from .execution import BackendResources, ExecutionRequirements +from .spec import LayerSpec +from .workspace import AttentionWorkspace + + +@dataclass(slots=True) +class AttentionExecutionContext: + capacity: ContextCapacity + capture: bool + workspace: AttentionWorkspace + batch: AttentionBatch | None = None + current_max_kv_seq_len: int = 0 + backend_identity: AttentionBackend | None = None + + +Context = TypeVar("Context", bound=AttentionExecutionContext) +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 是模型层号，不一定等于物理存储下标。 + forward 从当前 execution 查找 context，不保存批次 context；prepare/update 由 Runtime 管理。 + """ + + 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, execution: AttentionExecution, inputs: Inputs, *, out: torch.Tensor | None = None) -> Output: + context = self.context(execution) + return self.backend.forward(self, context, inputs, out=out) + + +class AttentionBackend(ABC, Generic[Context, Inputs, Output]): + """同组层共用的执行实现；每批次变化的 metadata 和 wrapper 放在 context 中。 + + Context 的具体类型由 backend 定义，实例由 InferStateInfo.attention_runtime 保存，不能跨 backend + 或并发 microbatch 共用。调用顺序为 allocate_context -> prepare -> update， + 随后同组各层用各自的 LayerAttentionBinding / Inputs 调用 forward。 + """ + + @classmethod + @abstractmethod + def check_environment(cls, caps: DeviceCaps) -> None: + """Reject unavailable platforms or provider dependencies with UnsupportedAttention.""" + pass + + @classmethod + @abstractmethod + def check_support( + cls, spec: LayerSpec, caps: DeviceCaps | None = None, execution: ExecutionRequirements | None = None + ) -> None: + """Reject unsupported static semantics or execution requirements.""" + pass + + @classmethod + @abstractmethod + def create(cls, model, spec: LayerSpec) -> AttentionBackend[Context, Inputs, Output]: + """Construct a backend for this spec; model supplies runtime and layer resources.""" + pass + + @abstractmethod + def allocate_context(self, capacity: ContextCapacity, *, capture: bool, workspace: AttentionWorkspace) -> 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[Any, Context, Inputs, Output], + 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 diff --git a/lightllm/common/attention/api/batch.py b/lightllm/common/attention/api/batch.py new file mode 100644 index 0000000000..41a24d464a --- /dev/null +++ b/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") diff --git a/lightllm/common/attention/api/execution.py b/lightllm/common/attention/api/execution.py new file mode 100644 index 0000000000..926911cb65 --- /dev/null +++ b/lightllm/common/attention/api/execution.py @@ -0,0 +1,36 @@ +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 + expert_parallel: bool = False + transfer_paths: frozenset[str] = frozenset() + + +class WorkspaceRequirement(NamedTuple): + name: str + elements: int + dtype: torch.dtype + + +@dataclass(frozen=True) +class BackendResources: + """捕获前的资源上界；workspace 项为 (名称, 元素数, dtype)。 + + 同一执行槽位按名称取最大容量，并发槽位分别分配。主机长度快照只在需要 + 它的实现中创建；持久存储的初始化由 State 执行，不能在 forward 中反复清零。 + """ + + workspace: tuple[WorkspaceRequirement, ...] = () + zero_initialize_cache: bool = False diff --git a/lightllm/common/attention/api/graph.py b/lightllm/common/attention/api/graph.py new file mode 100644 index 0000000000..16b257f02f --- /dev/null +++ b/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 diff --git a/lightllm/common/attention/api/inputs.py b/lightllm/common/attention/api/inputs.py new file mode 100644 index 0000000000..c1f441cda9 --- /dev/null +++ b/lightllm/common/attention/api/inputs.py @@ -0,0 +1,171 @@ +"""区分初始化时绑定的层资源与每次 forward 的 tensor 输入。 + +*Params 随 LayerAttentionBinding 保存，借用权重以及 State 管理的 KV / Conv / SSM 视图； +*Inputs 由模型层为当前调用构造，批次级请求索引和 provider 派生 metadata 则由 context 提供。 +这些数据结构只固定字段绑定，不禁止 backend 原地更新所借用的状态 tensor。 +""" + +from dataclasses import dataclass +from typing import NamedTuple, Optional, Protocol + +import torch + + +@dataclass(frozen=True) +class DenseParams: + key: torch.Tensor + value: torch.Tensor + alibi: Optional[torch.Tensor] = None + sink: Optional[torch.Tensor] = None + query_scale: Optional[torch.Tensor] = None + key_scale: torch.Tensor | float | None = None + value_scale: torch.Tensor | float | None = None + + +@dataclass(frozen=True) +class MLAParams: + kv: torch.Tensor + kv_lora_rank: int + qk_rope_head_dim: int + + +class ReplaySSMState(Protocol): + """Borrowed replay storage; State retains allocation, checkpoint and release ownership.""" + + ssm_state: torch.Tensor + keys: torch.Tensor + deltas: torch.Tensor + gates: torch.Tensor + raw_keys: torch.Tensor + raw_values: torch.Tensor + betas: torch.Tensor + cursors: torch.Tensor + tile_sync_mark: torch.Tensor + num_key_heads: int + capacity: int + verify_width: int + compute_mode: str + gate_lower_bound: float + hold_req_id: int + kda: bool + + def storage_layer(self, logical_layer_id: int) -> int: + ... + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: + ... + + +@dataclass(frozen=True) +class GDNParams: + conv: torch.Tensor + ssm: torch.Tensor + conv_weight: torch.Tensor + conv_bias: Optional[torch.Tensor] + a_log: torch.Tensor + dt_bias: torch.Tensor + mtp_state_index: torch.Tensor | None = None + replay_state: ReplaySSMState | None = None + + +@dataclass(frozen=True) +class KDAParams: + conv: torch.Tensor + ssm: torch.Tensor + conv_weight: torch.Tensor + a_log: torch.Tensor + dt_bias: torch.Tensor + mtp_state_index: torch.Tensor | None = None + replay_state: ReplaySSMState | None = None + + +@dataclass(frozen=True) +class DenseInputs: + """Dense 路径在调用前已写入 KV cache，因此这里仅传 query。""" + + query: torch.Tensor + + +class VarlenAttentionInputs(NamedTuple): + """非因果 self-attention 的 Q/K/V，布局均为 [total_tokens, heads, head_dim]。 + + cu_seqlens 是同设备上的 int32 前缀和，始于 0、终于 total_tokens；相邻边界之间 + 独立做 attention，允许空序列。max_seqlen 是最长序列长度的容量上界，由调用方提供以避免 GPU 同步。 + 每层只打包 tensor 引用，使用 NamedTuple 保持字段不可变并减少构造开销。 + """ + + query: torch.Tensor + key: torch.Tensor + value: torch.Tensor + cu_seqlens: torch.Tensor + max_seqlen: int + + def validate(self) -> None: + """Check sequence boundaries without copying CUDA lengths to the host. + + max_seqlen is an upper bound (Graph buckets may reserve a larger value). + CUDA value checks are asynchronous device assertions, including in Graphs; + invalid boundaries are input errors, not a reason to try another provider. + """ + q, k, v, cu, maximum = self + if q.ndim != 3 or q.shape != k.shape or q.shape != v.shape: + raise ValueError("varlen self-attention requires matching [tokens, heads, dim] Q/K/V") + if q.device != k.device or q.device != v.device or q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("Q/K/V must have the same dtype and device") + if cu.ndim != 1 or cu.numel() < 2 or cu.dtype != torch.int32 or cu.device != q.device: + raise ValueError("cu_seqlens must be a same-device int32 vector with at least two entries") + if maximum < 0 or (q.shape[0] and maximum == 0): + raise ValueError("nonempty attention requires a positive max_seqlen") + lengths = cu[1:] - cu[:-1] + valid = (cu[0] == 0) & (cu[-1] == q.shape[0]) & (lengths >= 0).all() & (lengths <= maximum).all() + message = "cu_seqlens must span all tokens with nonnegative lengths bounded by max_seqlen" + if cu.device.type == "cuda": + torch._assert_async(valid, message) + elif not bool(valid): + raise ValueError(message) + + +@dataclass(frozen=True) +class MLAPrefillInputs: + query: torch.Tensor + key_nope: torch.Tensor + key_rope: torch.Tensor + value: torch.Tensor + + +@dataclass(frozen=True) +class MLADecodeInputs: + query_nope: torch.Tensor + query_rope: torch.Tensor + + +@dataclass(frozen=True) +class SparseMLAPrefillInputs: + query: torch.Tensor + topk_mem_indices: torch.Tensor + topk_indices: Optional[torch.Tensor] = None + prefill_cache_kv: Optional[torch.Tensor] = None + + +@dataclass(frozen=True) +class SparseMLADecodeInputs: + query_nope: torch.Tensor + query_rope: torch.Tensor + topk_mem_indices: torch.Tensor + + +@dataclass(frozen=True) +class GDNInputs: + mixed_qkvzba: torch.Tensor + + +@dataclass(frozen=True) +class KDAInputs: + mixed_qkv: torch.Tensor + raw_gate: torch.Tensor + raw_beta: torch.Tensor + + +class GDNOutput(NamedTuple): + hidden_states: torch.Tensor + gate: torch.Tensor diff --git a/lightllm/common/attention/api/spec.py b/lightllm/common/attention/api/spec.py new file mode 100644 index 0000000000..6eda59570c --- /dev/null +++ b/lightllm/common/attention/api/spec.py @@ -0,0 +1,75 @@ +from dataclasses import dataclass +from enum import Enum +from typing import Optional + +import torch + + +class AttentionFamily(str, Enum): + """计算类型；Triton、FA3 等实现库由 selector 另行选择。""" + + DENSE = "dense" + MLA = "mla" + SPARSE_MLA = "sparse_mla" + GDN = "gdn" + KDA = "kda" + + +class AttentionPhase(str, Enum): + PREFILL = "prefill" + DECODE = "decode" + + +@dataclass(frozen=True) +class CacheLayout: + """存储布局的静态描述，不持有或分配存储 tensor。 + + prefill / decode 必须接受同一布局；dtype、量化粒度和 scale 布局也属于兼容性条件。 + storage_schema 记录 State 显式声明的布局版本及配置，排除逻辑层号，使同布局的层可以归为一组。 + """ + + format: str + page_size: int + dtype: torch.dtype + key_heads: int + key_dim: int + value_dim: int + state_dtype: Optional[torch.dtype] = None + storage_schema: tuple = () + quant_group_size: int = 0 + scale_mode: str = "none" + + def __post_init__(self): + if self.page_size <= 0: + raise ValueError("cache page_size must be positive") + + +@dataclass(frozen=True) +class LayerSpec: + """一个层在一个 phase 下的静态计算需求，用于能力检查和分组。 + + 初始化按完整的 (prefill_spec, decode_spec) 分组，只有相同需求的层才共享 backend / context。 + 层权重和存储视图放在 LayerAttentionBinding.params，批次长度与请求索引放在 AttentionBatch。 + """ + + family: AttentionFamily + phase: AttentionPhase + dtype: torch.dtype + cache: CacheLayout + num_q_heads: int + head_size: int + value_head_size: int + sliding_window: Optional[tuple[int, int]] = None + use_alibi: bool = False + use_sink: bool = False + causal: bool = True + use_mtp: bool = False + dynamic_verify: bool = False + variant: str = "" + softmax_scale: float | None = None + + def __post_init__(self): + if min(self.num_q_heads, self.head_size, self.value_head_size) <= 0: + raise ValueError("attention head counts and dimensions must be positive") + if self.family in (AttentionFamily.GDN, AttentionFamily.KDA) and self.sliding_window is not None: + raise ValueError("recurrent attention cannot use a paged sliding window") diff --git a/lightllm/common/attention/api/varlen.py b/lightllm/common/attention/api/varlen.py new file mode 100644 index 0000000000..a1c7743c03 --- /dev/null +++ b/lightllm/common/attention/api/varlen.py @@ -0,0 +1,54 @@ +from abc import ABC, abstractmethod + +import torch + +from .inputs import VarlenAttentionInputs + + +class VarlenAttentionOp(ABC): + """无持久 KV 的双向 dense attention，供视觉和音频编码器使用。 + + Q/K/V 与变长边界由调用方直接传入，不需要文本推理的请求表或 cache context。 + 实现不保留批次 tensor；调用方持有 out，返回值必须引用同一块输出存储。 + out 的 shape/dtype/device 与 query 一致，head_dim 连续，其他维度允许无重叠的 stride； + 不允许与 Q/K/V 共享存储。每次 forward 在启动 kernel 前验证这些 metadata。 + SDPA / xFormers 需要在 CPU 读取序列边界；只有声明 supports_graph 的实现可整段捕获。 + """ + + supports_graph = False + + @classmethod + @abstractmethod + def check_environment(cls, caps) -> None: + """检查运行环境；不支持时抛出 UnsupportedAttention。""" + pass + + @classmethod + @abstractmethod + def check_support(cls, inputs: VarlenAttentionInputs, caps=None) -> None: + """检查 dtype/head dim/platform 等能力限制；不读取或校验序列边界值。""" + pass + + @abstractmethod + def forward(self, inputs: VarlenAttentionInputs, *, out: torch.Tensor) -> torch.Tensor: + """每次执行校验当前输入布局，然后写入 out；Graph replay 也须重新执行设备值检查。""" + pass + + +def validate_varlen_output(inputs: VarlenAttentionInputs, out: torch.Tensor) -> None: + """Validate metadata only; out may have padding between heads/tokens, but cannot alias inputs.""" + query = inputs.query + if out.shape != query.shape or out.dtype != query.dtype or out.device != query.device: + raise ValueError("varlen out must match query shape, dtype and device") + if out.layout != torch.strided or out.stride(-1) != 1: + raise ValueError("varlen out must have a contiguous head dimension") + extent = 1 + for stride, size in sorted(zip(out.stride(), out.shape)): + if size > 1: + if stride < extent: + raise ValueError("varlen out must have non-overlapping strides") + extent = stride * (size - 1) + extent + if out.numel(): + storage = out.untyped_storage().data_ptr() + if any(storage == tensor.untyped_storage().data_ptr() for tensor in (inputs.query, inputs.key, inputs.value)): + raise ValueError("varlen out must not share storage with query, key or value") diff --git a/lightllm/common/attention/api/workspace.py b/lightllm/common/attention/api/workspace.py new file mode 100644 index 0000000000..9fc6bb66dc --- /dev/null +++ b/lightllm/common/attention/api/workspace.py @@ -0,0 +1,106 @@ +from typing import Callable, TypeVar, cast + +import torch + +from lightllm.platform import get_device_runtime + +T = TypeVar("T") + + +class AttentionWorkspace: + def __init__(self, device=None, *, runtime=None): + self.runtime = runtime if runtime is not None else get_device_runtime(device) + self.device = self.runtime.device + self._buffers = {} + self._objects = {} + self._completion = None + self._last_stream = None + self._frozen = False + self._pool = None + + def fork(self): + """捕获实例借用当前池的地址；池后续扩容时，旧图仍持有原 tensor。 + + 同一槽位的借用者共享完成事件，因此共享 scratch 不允许并发执行；不同 + microbatch 槽位由模型创建不同的池。只有实际扩容才需要额外存储。 + """ + workspace = AttentionWorkspace(runtime=self.runtime) + workspace._pool = self._pool or self + workspace._objects = self._objects + return workspace + + def get_object(self, name: str, factory: Callable[[], T], *, expected_type: type[T] | None = None) -> T: + """Share objects across mutually exclusive contexts, checking their declared type. + + Forks share this live object pool, not a snapshot. Freezing a fork prevents that + fork from adding objects or growing buffers; other buckets may still warm up. + Existing objects must keep captured addresses stable even when the pool grows. + """ + if expected_type is None: + if not isinstance(factory, type): + raise TypeError("workspace object factories require an expected_type") + expected_type = factory + if name not in self._objects: + if self._frozen: + raise RuntimeError(f"captured workspace object {name} must be created before capture") + value = factory() + if not isinstance(value, expected_type): + raise TypeError( + f"workspace object {name} requires {expected_type.__name__}, got {type(value).__name__}" + ) + self._objects[name] = value + value = self._objects[name] + if not isinstance(value, expected_type): + raise TypeError(f"workspace object {name} requires {expected_type.__name__}, got {type(value).__name__}") + return cast(T, value) + + def begin(self): + if self._pool is not None: + return self._pool.begin() + if self._completion is not None: + stream = self.runtime.current_stream() + if stream != self._last_stream: + stream.wait_event(self._completion) + + def finish(self, dependencies=()): + if self._pool is not None: + return self._pool.finish(dependencies) + stream = self.runtime.current_stream() + if stream is not None: + for event in dependencies: + stream.wait_event(event) + if self._completion is None: + self._completion = self.runtime.create_event() + self._last_stream = stream + self._completion.record(self._last_stream) + + def freeze(self): + """捕获后禁止替换任何地址；互斥 bucket 可继续借用已预留的大缓冲区。""" + self._frozen = True + + def reserve(self, resources): + for requirement in resources.workspace: + self.get(requirement.name, requirement.elements, requirement.dtype) + + def synchronize(self): + if self._pool is not None: + return self._pool.synchronize() + if self._completion is not None: + self._completion.synchronize() + + def get(self, name: str, size: int, dtype=torch.uint8): + if not name or size <= 0: + raise ValueError("workspace name and positive size are required") + key = name, dtype + tensor = self._buffers.get(key) + if tensor is None or tensor.numel() < size: + if self._frozen: + raise RuntimeError(f"captured workspace {name} cannot grow; reserve capacity before capture") + tensor = ( + self._pool.get(name, size, dtype) + if self._pool is not None + else self.runtime.allocate(size, dtype=dtype) + ) + self._buffers[key] = tensor + + return tensor diff --git a/lightllm/common/attention/backends/__init__.py b/lightllm/common/attention/backends/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/_shared/__init__.py b/lightllm/common/attention/backends/_shared/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/_shared/linear/__init__.py b/lightllm/common/attention/backends/_shared/linear/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/_shared/prefill/__init__.py b/lightllm/common/attention/backends/_shared/prefill/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/_workspace.py b/lightllm/common/attention/backends/_workspace.py new file mode 100644 index 0000000000..0d3f6c4b34 --- /dev/null +++ b/lightllm/common/attention/backends/_workspace.py @@ -0,0 +1,3 @@ +from lightllm.common.attention.api.workspace import AttentionWorkspace + +__all__ = ["AttentionWorkspace"] diff --git a/lightllm/common/attention/backends/base.py b/lightllm/common/attention/backends/base.py new file mode 100644 index 0000000000..14e5afc097 --- /dev/null +++ b/lightllm/common/attention/backends/base.py @@ -0,0 +1,217 @@ +from dataclasses import replace +from types import MappingProxyType +from typing import Any, Generic, TypeVar, cast + +import torch + +from lightllm.common.attention._environment import check_attention_environment +from lightllm.common.attention.api.backend import ( + AttentionBackend, + AttentionExecutionContext, + LayerAttentionBinding, + UnsupportedAttention, +) +from lightllm.common.attention.api.batch import AttentionBatch, ContextCapacity +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase, LayerSpec +from lightllm.platform import DeviceCaps, Platform + +Context = TypeVar("Context", bound=AttentionExecutionContext) +Inputs = TypeVar("Inputs") +Output = TypeVar("Output") + + +class BaseAttentionBackend(AttentionBackend[Context, Inputs, Output], Generic[Context, Inputs, Output]): + requires_cpu_cache_sync = False + family = AttentionFamily.DENSE + provider = "triton" + cache_formats = frozenset({"None"}) + supports_attention_sink = False + supports_noncausal_prefill = False + supports_noncausal_decode = False + supported_platforms = frozenset({Platform.CUDA}) + supported_phases = frozenset(AttentionPhase) + supported_graph_phases = frozenset(AttentionPhase) + activation_dtypes = frozenset({torch.float16, torch.bfloat16, torch.float32}) + dependency = "triton" + min_device_capability = None + supports_alibi = False + supports_sliding_window = False + supports_speculative_window = True + supports_expert_parallel = True + max_head_size = None + + def __init__(self, model, spec: LayerSpec): + self.runtime = model.device_runtime + self.spec = spec + self.page_size = spec.cache.page_size + self.network_config = MappingProxyType(dict(model.config)) + self.decode_draft_step = ( + model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model) if spec.use_mtp else 0 + ) + self.dynamic_verify = spec.dynamic_verify + self.causal = spec.causal + self.hold_req_id = model.req_manager.HOLD_REQUEST_ID + + @classmethod + def create(cls, model, spec: LayerSpec): + backend = cls(model, spec) + backend.configure_spec(spec) + return backend + + def configure_spec(self, spec: LayerSpec) -> None: + """Apply per-group dimensions before allocating contexts; default backends use layer params.""" + pass + + @classmethod + def check_environment(cls, caps: DeviceCaps) -> None: + check_attention_environment( + caps, + provider=cls.provider, + supported_platforms=cls.supported_platforms, + min_device_capability=cls.min_device_capability, + dependency=cls.dependency, + ) + + @classmethod + def check_support(cls, spec: LayerSpec, caps=None, execution=None) -> None: + if spec.family is not cls.family: + raise UnsupportedAttention(f"{cls.__name__} does not implement {spec.family.value}") + if spec.phase not in cls.supported_phases: + raise UnsupportedAttention(f"{cls.provider} does not implement {spec.phase.value}") + if spec.cache.format not in cls.cache_formats: + raise UnsupportedAttention(f"cache encoding {spec.cache.format} is unsupported") + if spec.dtype not in cls.activation_dtypes: + raise UnsupportedAttention(f"activation dtype {spec.dtype} is unsupported") + if spec.use_alibi and not cls.supports_alibi: + raise UnsupportedAttention("this backend does not support ALiBi") + if spec.use_alibi and spec.sliding_window is not None: + raise UnsupportedAttention("ALiBi with sliding-window attention is unsupported") + if spec.use_sink and not cls.supports_attention_sink: + raise UnsupportedAttention("this backend does not support attention sinks") + if spec.sliding_window is not None and not cls.supports_sliding_window: + raise UnsupportedAttention("this backend does not support sliding-window attention") + if cls.max_head_size is not None and spec.head_size > cls.max_head_size: + raise UnsupportedAttention(f"head size exceeds {cls.max_head_size}") + if spec.phase is AttentionPhase.PREFILL and not spec.causal and not cls.supports_noncausal_prefill: + raise UnsupportedAttention("this backend does not support non-causal prefill") + if spec.phase is AttentionPhase.DECODE and not spec.causal and not cls.supports_noncausal_decode: + raise UnsupportedAttention("this backend does not support non-causal block decode") + if spec.use_mtp and spec.sliding_window is not None and not cls.supports_speculative_window: + raise UnsupportedAttention("speculative decode does not implement sliding windows") + if execution is not None: + if execution.expert_parallel and not cls.supports_expert_parallel: + raise UnsupportedAttention("this backend is incompatible with expert-parallel MoE execution") + if spec.phase in execution.graph_phases and ( + spec.phase not in cls.supported_graph_phases or (caps is not None and not caps.supports_graph) + ): + raise UnsupportedAttention(f"{spec.phase.value} Graph is unsupported") + + def allocate_context(self, capacity: ContextCapacity, *, capture: bool, workspace) -> Context: + if capacity.phase is AttentionPhase.PREFILL: + context = self._allocate_prefill_context(capacity, capture, workspace) + else: + context = self._allocate_decode_context(capacity, capture, workspace) + context.backend_identity = self + return cast(Context, context) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return AttentionExecutionContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return AttentionExecutionContext(capacity, capture, workspace) + + def prepare(self, context: Context, batch: AttentionBatch) -> None: + if context.backend_identity is not self: + raise ValueError("context belongs to a different attention backend") + batch.validate() + context.capacity.check(batch) + context.current_max_kv_seq_len = batch.max_kv_seq_len + if context.capture and context.batch is not None: + context.batch.copy_inputs(batch) + elif context.capture: + # 图内分配和 launch 使用容量上界；图外 provider planning 仍用 current_max_kv_seq_len。 + context.batch = replace( + batch, + max_kv_seq_len=context.capacity.max_kv_seq_len, + total_token_num=context.capacity.batch_size * context.capacity.max_kv_seq_len, + ) + else: + context.batch = batch + if batch.phase is AttentionPhase.PREFILL: + self._prepare_prefill(context) + else: + self._prepare_decode(context) + + def _prepare_prefill(self, context: Context) -> None: + """Optional host planning; kernels without provider plans need no preparation.""" + pass + + def _prepare_decode(self, context: Context) -> None: + """Optional host planning, including updates before Graph replay.""" + pass + + def update(self, context: Context) -> None: + if context.backend_identity is not self: + raise ValueError("context belongs to a different attention backend") + if context.batch is None: + raise RuntimeError("prepare attention inputs before update") + if context.capacity.phase is AttentionPhase.PREFILL: + self._update_prefill(context) + else: + self._update_decode(context) + + def _update_prefill(self, context: Context) -> None: + """Optional derived metadata; direct consumers of batch tensors need no update.""" + pass + + def _update_decode(self, context: Context) -> None: + """Optional GPU metadata generation, captured alongside layer execution.""" + pass + + def forward( + self, + layer: LayerAttentionBinding[Any, Context, Inputs, Output], + context: Context, + inputs: Inputs, + *, + out: torch.Tensor | None = None, + ) -> Output: + if context.batch is None: + raise RuntimeError("prepare attention inputs before forward") + if context.backend_identity is not self: + raise ValueError("context belongs to a different attention backend") + if layer.spec.phase is not context.capacity.phase: + raise ValueError("layer attention and context phases differ") + method = self._forward_prefill if context.capacity.phase is AttentionPhase.PREFILL else self._forward_decode + result = method(layer, context, inputs, out=out) + if out is None: + return result + attention = result[0] if isinstance(result, tuple) else result + if attention.dtype != out.dtype or attention.device != out.device: + raise ValueError("out must match the attention result dtype and device") + if attention.untyped_storage().data_ptr() == out.untyped_storage().data_ptr(): + # Providers may return the valid heads of a padded output buffer. + begin = attention.storage_offset() + end = begin + sum((size - 1) * stride for size, stride in zip(attention.shape, attention.stride())) + 1 + out_end = ( + out.storage_offset() + sum((size - 1) * stride for size, stride in zip(out.shape, out.stride())) + 1 + ) + if attention.numel() and (begin < out.storage_offset() or end > out_end): + raise ValueError("attention result view extends beyond out") + return result + if attention.numel() != out.numel(): + raise ValueError("out must match the attention result element count") + out.copy_(attention.reshape(out.shape)) + if isinstance(result, tuple): + return type(result)(out, *result[1:]) if hasattr(result, "_fields") else (out, *result[1:]) + return out + + @property + def infer_page_size(self) -> int: + return self.page_size + + def uses_dynamic_spec_verify_layout(self): + return self.dynamic_verify + + def uses_causal_attention(self): + return self.causal diff --git a/lightllm/common/attention/backends/dense/__init__.py b/lightllm/common/attention/backends/dense/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/fa3/__init__.py b/lightllm/common/attention/backends/dense/fa3/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/flashinfer/__init__.py b/lightllm/common/attention/backends/dense/flashinfer/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/sdpa/__init__.py b/lightllm/common/attention/backends/dense/sdpa/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/__init__.py b/lightllm/common/attention/backends/dense/triton/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/triton/kernels/prefill/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/dense/xformers/__init__.py b/lightllm/common/attention/backends/dense/xformers/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/__init__.py b/lightllm/common/attention/backends/gdn/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/flashqla/__init__.py b/lightllm/common/attention/backends/gdn/flashqla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/replay/__init__.py b/lightllm/common/attention/backends/gdn/replay/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/triton/__init__.py b/lightllm/common/attention/backends/gdn/triton/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/triton/kernels/__init__.py b/lightllm/common/attention/backends/gdn/triton/kernels/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/kda/__init__.py b/lightllm/common/attention/backends/kda/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/kda/replay/__init__.py b/lightllm/common/attention/backends/kda/replay/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/kda/triton/__init__.py b/lightllm/common/attention/backends/kda/triton/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/kda/triton/kernels/__init__.py b/lightllm/common/attention/backends/kda/triton/kernels/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/__init__.py b/lightllm/common/attention/backends/mla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/fa3/__init__.py b/lightllm/common/attention/backends/mla/fa3/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/flashinfer/__init__.py b/lightllm/common/attention/backends/mla/flashinfer/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/triton/__init__.py b/lightllm/common/attention/backends/mla/triton/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/triton/kernels/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/sparse_mla/__init__.py b/lightllm/common/attention/backends/sparse_mla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/__init__.py b/lightllm/common/attention/backends/sparse_mla/flashmla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/basemodel/attention/__init__.py b/lightllm/common/basemodel/attention/__init__.py index bbe6f7b7dd..d728a4ee9c 100644 --- a/lightllm/common/basemodel/attention/__init__.py +++ b/lightllm/common/basemodel/attention/__init__.py @@ -1,10 +1,36 @@ -from .base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl +from lightllm.common.attention.api import AttentionPhase, ExecutionRequirements + +from .base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState from .create_utils import ( - get_prefill_att_backend_class, get_decode_att_backend_class, - get_mla_prefill_att_backend_class, get_mla_decode_att_backend_class, - get_nsa_prefill_att_backend_class, - get_nsa_decode_att_backend_class, + get_mla_prefill_att_backend_class, get_neo_prefill_att_backend_class, + get_nsa_decode_att_backend_class, + get_nsa_prefill_att_backend_class, + get_prefill_att_backend_class, ) + + +def describe_execution_requirements(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 ExecutionRequirements( + 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)), + args.enable_ep_moe, + frozenset(transfers), + ) diff --git a/lightllm/platform/__init__.py b/lightllm/platform/__init__.py new file mode 100644 index 0000000000..705eb90a6e --- /dev/null +++ b/lightllm/platform/__init__.py @@ -0,0 +1,72 @@ +"""Device operations used by attention; CPU/CUDA adapters are built in.""" + +import os +from functools import lru_cache + +import torch + +from .api import DeviceCaps, DeviceRuntime, Platform + +_runtime_factories = {} + + +def register_device_runtime(platform, factory): + """Register an additional adapter at worker startup, before creating any runtime. + + The factory receives the original device argument and must resolve/validate it + for its platform. Registering an adapter does not imply provider kernel support. + """ + platform = Platform(platform) + if platform in (Platform.CPU, Platform.CUDA) or platform in _runtime_factories: + raise ValueError(f"device runtime already registered for {platform.value}") + _runtime_factories[platform] = factory + + +def get_device_runtime(device=None, *, platform=None): + platform = Platform(platform) if platform is not None else None + if platform in _runtime_factories: + return _runtime_factories[platform](device) + if platform not in (None, Platform.CPU, Platform.CUDA): + raise ValueError(f"no registered runtime for {platform.value}; built-in adapters cover CUDA and CPU") + use_worker_default = device is None + default_type = "cpu" if platform is Platform.CPU else "cuda" + if isinstance(device, int): + device = f"{default_type}:{device}" + device = device if device is not None else default_type + if isinstance(device, str) and device.split(":")[0] not in ("cpu", "cuda", "meta"): + raise ValueError(f"attention backend refactor supports CUDA and CPU devices only, got {device}") + device = torch.device(device) + platform = ( + Platform(platform) + if platform is not None + else { + "cpu": Platform.CPU, + "meta": Platform.CPU, + "cuda": Platform.CUDA, + }.get(device.type) + ) + if platform is None: + raise ValueError(f"no device runtime for {device}") + if platform is Platform.CPU and device.type not in ("cpu", "meta"): + raise ValueError(f"CPU runtime cannot operate on {device}") + if platform is not Platform.CPU and device.type in ("cpu", "meta"): + raise ValueError(f"{platform.value} runtime cannot operate on {device}") + # 保留调用方指定的 CUDA device；没有显式 ordinal 时使用 worker 绑定的 ordinal。 + if device.index is None and device.type not in ("cpu", "meta"): + index = os.environ.get("LIGHTLLM_CURRENT_DEVICE_ID") if use_worker_default else None + if index is None: + api = torch.cuda + index = api.current_device() if api.is_available() else 0 + device = torch.device(device.type, int(index)) + return _runtime(platform, device) + + +@lru_cache(maxsize=None) +def _runtime(platform, device): + if platform is Platform.CPU: + return DeviceRuntime(device) + if platform is Platform.CUDA: + from .cuda import CudaRuntime + + return CudaRuntime(device) + raise ValueError(f"no device runtime for platform {platform.value}") diff --git a/lightllm/platform/api.py b/lightllm/platform/api.py new file mode 100644 index 0000000000..aaaba26c47 --- /dev/null +++ b/lightllm/platform/api.py @@ -0,0 +1,96 @@ +from contextlib import nullcontext +from dataclasses import dataclass +from functools import cached_property + +import torch + +from lightllm.utils.device_utils import Platform + + +@dataclass(frozen=True) +class DeviceCaps: + """实际设备的能力快照；torch 的 device.type 不足以区分 CUDA、ROCm 和兼容层。""" + + platform: Platform + device: torch.device + available: bool = True + device_name: str = "" + capability: tuple[int, int] | None = None + supports_graph: bool = False + + +class DeviceRuntime: + """一个 worker/device 的设备操作。实现库和模型语义不属于 runtime。 + + 构造不切换当前设备；调用方先完成 worker 的设备绑定。caps 延迟读取， + 注册表和 CPU 契约测试可以在不初始化设备驱动的情况下导入这些接口。 + """ + + platform = Platform.CPU + graph_type = "" + distributed_backend = "gloo" + + def __init__(self, device, api=None): + self.device = torch.device(device) + self.api = api + + @cached_property + def caps(self): + if self.api is None: + return DeviceCaps(self.platform, self.device, available=self.device.type != "meta") + available = self.api.is_available() + capability = None + if available and hasattr(self.api, "get_device_capability"): + capability = tuple(self.api.get_device_capability(self.device)) + return DeviceCaps( + self.platform, + self.device, + available, + self.api.get_device_name(self.device) if available else "", + capability, + available and hasattr(self.api, self.graph_type), + ) + + def allocate(self, shape, dtype=torch.uint8): + return torch.empty(shape, dtype=dtype, device=self.device) + + def current_stream(self): + return self.api.current_stream(self.device) if self.api is not None else None + + def create_stream(self): + return self.api.Stream(device=self.device) if self.api is not None else None + + def use_stream(self, stream): + return self.api.stream(stream) if self.api is not None else nullcontext() + + def create_event(self, *, enable_timing=False): + return self.api.Event(enable_timing=enable_timing) if self.api is not None else None + + def synchronize(self): + if self.api is not None: + self.api.synchronize(self.device) + + def is_capturing(self): + return self.api is not None and self.api.is_current_stream_capturing() + + def empty_cache(self): + if self.api is not None: + self.api.empty_cache() + + def memory_info(self): + if self.api is None: + raise ValueError("host/meta storage requires an explicit capacity or a memory-budget runtime") + return self.api.mem_get_info(self.device) + + def graph_pool(self): + return self.api.graph_pool_handle() if self.caps.supports_graph else None + + def create_graph(self): + if not self.caps.supports_graph: + raise RuntimeError(f"{self.platform.value} device {self.device} does not support Graph capture") + return getattr(self.api, self.graph_type)() + + def capture_graph(self, graph, *, pool=None, memory_saver=None): + if memory_saver is not None and getattr(memory_saver, "enabled", False): + raise ValueError(f"memory saver is not supported on {self.platform.value}") + return self.api.graph(graph, pool=pool) diff --git a/lightllm/platform/cuda.py b/lightllm/platform/cuda.py new file mode 100644 index 0000000000..0db5e4891b --- /dev/null +++ b/lightllm/platform/cuda.py @@ -0,0 +1,17 @@ +import torch + +from .api import DeviceRuntime, Platform + + +class CudaRuntime(DeviceRuntime): + platform = Platform.CUDA + graph_type = "CUDAGraph" + distributed_backend = "nccl" + + def __init__(self, device): + super().__init__(device, torch.cuda) + + def capture_graph(self, graph, *, pool=None, memory_saver=None): + if memory_saver is not None: + return memory_saver.cuda_graph(graph, pool=pool) + return self.api.graph(graph, pool=pool) diff --git a/unit_tests/common/attention/helpers.py b/unit_tests/common/attention/helpers.py new file mode 100644 index 0000000000..ff5c950c5f --- /dev/null +++ b/unit_tests/common/attention/helpers.py @@ -0,0 +1,283 @@ +from dataclasses import replace +from types import SimpleNamespace + +import torch + +from lightllm.common.attention.api import ( + AttentionBatch, + AttentionFamily, + AttentionPhase, + CacheLayout, + ContextCapacity, + LayerAttentionBinding, + LayerSpec, +) +from lightllm.common.attention.backends._workspace import AttentionWorkspace +from lightllm.common.attention.backends.base import BaseAttentionBackend +from lightllm.common.basemodel.attention import describe_execution_requirements +from lightllm.platform import get_device_runtime +from lightllm.server.core.objs.start_args_type import StartArgs + + +def make_batch(lengths=(5, 2), *, phase=AttentionPhase.DECODE, device="cpu", **kwargs): + total = sum(lengths) + maximum = max(lengths) if lengths else 0 + lengths = torch.tensor(lengths, dtype=torch.int32, device=device) + count = lengths.numel() + q_lengths = torch.ones_like(lengths) + offsets = torch.arange(count + 1, dtype=torch.int32, device=device) + values = dict( + phase=phase, + batch_size=count, + token_num=count, + total_token_num=total, + max_q_seq_len=1, + max_kv_seq_len=maximum, + max_cache_len=0, + b_req_idx=torch.arange(count, dtype=torch.int32, device=device), + b_seq_len=lengths, + b_q_seq_len=q_lengths, + b1_cu_q_seq_len=offsets, + b1_cu_kv_seq_len=torch.cat((offsets[:1], lengths.cumsum(0).int())), + req_to_token_indexs=torch.zeros((max(count, 8), 256), dtype=torch.int32, device=device), + hold_req_id=7, + b_mtp_index=torch.zeros_like(lengths), + b_ready_cache_len=torch.zeros_like(lengths), + ) + values.update(kwargs) + return AttentionBatch(**values) + + +def make_spec(family=AttentionFamily.DENSE, phase=AttentionPhase.DECODE, **kwargs): + values = dict( + family=family, + phase=phase, + dtype=torch.bfloat16, + cache=CacheLayout( + family.value if family in (AttentionFamily.GDN, AttentionFamily.KDA) else "None", 1, torch.bfloat16, 2, 8, 8 + ), + num_q_heads=4, + head_size=8, + value_head_size=8, + ) + values.update(kwargs) + return LayerSpec(**values) + + +def make_backend(cls=BaseAttentionBackend, **kwargs): + backend = cls.__new__(cls) + values = dict( + page_size=1, + causal=True, + dynamic_verify=False, + decode_draft_step=0, + is_draft_model=False, + hold_req_id=7, + ) + values.update(kwargs) + if "infer_page_size" in values: + values["page_size"] = values.pop("infer_page_size") + vars(backend).update(values) + return backend + + +def make_context(backend, batch=None, *, capture=False, max_length=None, prepare=False): + batch = make_batch() if batch is None else batch + capacity = ContextCapacity( + batch.phase, batch.batch_size, batch.token_num, max_length or batch.max_kv_seq_len, batch.max_q_seq_len + ) + context = backend.allocate_context(capacity, capture=capture, workspace=AttentionWorkspace(batch.device)) + if prepare: + backend.prepare(context, batch) + else: + context.batch = batch + context.current_max_kv_seq_len = batch.max_kv_seq_len + return context + + +def make_layer_attention(backend, context, params, **spec_kwargs): + return LayerAttentionBinding( + group_id=0, + logical_layer_id=0, + spec=make_spec(backend.family, context.capacity.phase, **spec_kwargs), + params=params, + backend=backend, + alloc_tensor=lambda shape, dtype, device="cpu": torch.empty(shape, dtype=dtype, device=context.batch.device), + ) + + +def model_stub(**args): + args.setdefault("disable_cudagraph", not torch.cuda.is_available()) + model = SimpleNamespace( + args=StartArgs(**args), + device_runtime=get_device_runtime("cuda" if torch.cuda.is_available() else "cpu"), + att_state_role="main", + att_state_manager=SimpleNamespace( + get_full_attention_states=lambda *args: (), get_bounded_attention_states=lambda *args: () + ), + config={}, + is_mtp_draft_model=False, + tp_world_size_=1, + req_manager=SimpleNamespace(HOLD_REQUEST_ID=7), + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), + ) + model.attention_requirements = describe_execution_requirements(model) + return model + + +def batch_from_fixture(fixture, phase, token_count=None, device=None): + indices = getattr(fixture, "b_req_idx", None) + lengths = getattr(fixture, "b_seq_len", None) + device = device or ( + indices.device if indices is not None else getattr(getattr(fixture, "mem_index", None), "device", "cpu") + ) + count = indices.numel() if indices is not None else (token_count or getattr(fixture, "batch_size", None) or 1) + if lengths is None: + lengths = torch.ones(count, dtype=torch.int32, device=device) + values = {} + for name in AttentionBatch.__dataclass_fields__: + value = getattr(fixture, name, None) + if value is not None and name not in ("phase", "token_num", "batch_size"): + values[name] = value + manager = getattr(fixture, "req_manager", None) + if manager is not None: + for name in ("req_to_token_indexs",): + value = getattr(manager, name, None) + if value is not None: + values[name] = value + values["hold_req_id"] = getattr(manager, "HOLD_REQUEST_ID", 7) + + max_len = getattr(fixture, "max_kv_seq_len", None) or 1 + cache_key = (phase, count, str(device)) + cache = getattr(fixture, "_attention_test_batches", None) + if cache is None: + cache = fixture._attention_test_batches = {} + if cache_key not in cache: + cache[cache_key] = make_batch((1,) * count, phase=phase, device=device) + defaults = cache[cache_key] + values.update( + phase=phase, batch_size=count, token_num=token_count or count, b_seq_len=lengths, max_kv_seq_len=max_len + ) + return replace(defaults, **values) + + +def backend_context(cls, phase, settings=None, **values): + phase = AttentionPhase(phase) + if isinstance(settings, BaseAttentionBackend): + backend = settings + else: + config = dict(vars(settings)) if settings is not None else {} + model = config.pop("model", None) + if model is not None: + manager = getattr(model, "mtp_manager", None) + if manager is not None: + config["decode_draft_step"] = manager.get_decode_draft_step(getattr(model, "is_mtp_draft_model", False)) + req_manager = getattr(model, "req_manager", None) + if req_manager is not None: + config["hold_req_id"] = getattr(req_manager, "HOLD_REQUEST_ID", 7) + backend = make_backend(cls, **config) + defaults = dict( + network_config={"index_topk": 2048, "index_kpool": 4}, + prefill_graph_enabled=False, + decode_only=False, + prefill_token_capacity=0, + mtp_step=0, + decode_draft_step=0, + dynamic_verify=False, + causal=True, + ) + for name, value in defaults.items(): + if not hasattr(backend, name): + setattr(backend, name, value) + context = make_context(backend, make_batch((1,), phase=phase)) + if hasattr(backend, "_test_workspace"): + context.workspace = backend._test_workspace + for name, value in values.items(): + setattr(context, name, value) + return context + + +def update_context(context, fixture): + context.batch = batch_from_fixture(fixture, context.capacity.phase) + context.capacity = ContextCapacity( + context.batch.phase, + context.batch.batch_size, + context.batch.token_num, + context.batch.max_kv_seq_len, + context.batch.max_q_seq_len, + ) + context.backend_identity.prepare(context, context.batch) + context.backend_identity.update(context) + + +def run_sparse(context, fixture, query, kv, parameters, *, out=None, alloc_func=None): + from lightllm.common.attention.api.inputs import MLAParams, SparseMLADecodeInputs, SparseMLAPrefillInputs + + q = query[0] if isinstance(query, tuple) else query + context.capture = bool(getattr(fixture, "is_cuda_graph_capture", context.capture)) + if getattr(fixture, "prefill_cuda_graph_exe_list", None) is not None: + context.capture = True + context.batch = batch_from_fixture(fixture, context.capacity.phase, q.shape[0], q.device) + backend = context.backend_identity + if context.workspace.device != q.device: + context.workspace = AttentionWorkspace(q.device) + backend._test_workspace = context.workspace + if backend.__class__.__module__.startswith("lightllm.models.deepseek_v4"): + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4Inputs, DeepseekV4Params + + params = DeepseekV4Params( + kv[0], + kv[0], + kv[1], + parameters["compress_ratio"], + parameters["head_dim_v"], + parameters["softmax_scale"], + parameters["attn_sink"], + ) + inputs = DeepseekV4Inputs( + q, + parameters["swa_indices"], + parameters["swa_lengths"], + parameters.get("extra_indices"), + parameters.get("extra_lengths"), + ) + else: + params = MLAParams(kv, parameters["kv_lora_rank"], parameters.get("qk_rope_head_dim", 0)) + if context.capacity.phase is AttentionPhase.PREFILL: + inputs = SparseMLAPrefillInputs( + q, + parameters.get("topk_mem_indices"), + parameters.get("topk_indices"), + parameters.get("prefill_cache_kv"), + ) + else: + inputs = SparseMLADecodeInputs(q, query[1], parameters["topk_mem_indices"]) + attention = make_layer_attention(backend, context, params, softmax_scale=parameters["softmax_scale"]) + if alloc_func is not None: + attention = replace(attention, alloc_tensor=alloc_func) + + if ( + hasattr(context, "capture_flashmla_sched_meta") + and context.capture_flashmla_sched_meta is None + and torch.cuda.is_available() + ): + import flash_mla + + context.capture_flashmla_sched_meta = flash_mla.get_mla_metadata()[0] + return attention.backend.forward(attention, context, inputs, out=out) + + +def run_dense(context, fixture, query, key, value, *, out=None, alloc_func=None, params=None): + from lightllm.common.attention.api.inputs import DenseInputs, DenseParams + + context.batch = batch_from_fixture(fixture, context.capacity.phase, query.shape[0], query.device) + if params is None: + ks = vs = None + if isinstance(key, tuple): + key, ks = key + value, vs = value + params = DenseParams(key, value, key_scale=ks, value_scale=vs) + attention = make_layer_attention(context.backend_identity, context, params) + if alloc_func is not None: + attention = replace(attention, alloc_tensor=alloc_func) + return attention.backend.forward(attention, context, DenseInputs(query), out=out) diff --git a/unit_tests/common/attention/test_contracts.py b/unit_tests/common/attention/test_contracts.py new file mode 100644 index 0000000000..8a833e8da5 --- /dev/null +++ b/unit_tests/common/attention/test_contracts.py @@ -0,0 +1,139 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention.api import AttentionFamily, AttentionPhase, UnsupportedAttention +from lightllm.common.attention.backends.base import BaseAttentionBackend +from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub + + +def test_graph_copy_validates_all_inputs_before_writing(): + captured = make_batch() + incoming = make_batch((9, 7)) + incoming = replace(incoming, req_to_token_indexs=captured.req_to_token_indexs, b_seq_len=incoming.b_seq_len.long()) + before = captured.b_req_idx.clone() + with pytest.raises(ValueError, match="b_seq_len"): + captured.copy_inputs(incoming) + torch.testing.assert_close(captured.b_req_idx, before) + + +def test_graph_copy_keeps_tensor_addresses_and_uses_new_lengths(): + captured = make_batch() + incoming = make_batch((9, 7), req_to_token_indexs=captured.req_to_token_indexs) + pointer = captured.b_seq_len.data_ptr() + captured.copy_inputs(incoming) + assert captured.b_seq_len.data_ptr() == pointer + assert captured.b_seq_len.tolist() == [9, 7] + with pytest.raises(ValueError, match="request table"): + captured.copy_inputs(make_batch((9, 7))) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_context_capacity_and_backend_identity(capture): + backend = BaseAttentionBackend(model_stub(), make_spec()) + context = make_context(backend, capture=capture, max_length=16) + backend.prepare(context, make_batch(req_to_token_indexs=context.batch.req_to_token_indexs)) + with pytest.raises(ValueError, match="capacity"): + backend.prepare(context, make_batch((17, 2))) + with pytest.raises(ValueError, match="different attention backend"): + BaseAttentionBackend(model_stub(), make_spec()).prepare(context, make_batch()) + with pytest.raises(ValueError, match="different attention backend"): + BaseAttentionBackend(model_stub(), make_spec()).update(context) + attention = make_layer_attention(BaseAttentionBackend(model_stub(), make_spec()), context, params=None) + with pytest.raises(ValueError, match="different attention backend"): + attention.backend.forward(attention, context, inputs=None) + + +def test_layer_attention_rejects_wrong_phase(): + backend = BaseAttentionBackend(model_stub(), make_spec()) + context = make_context(backend, prepare=True) + attention = make_layer_attention(backend, context, params=None) + attention = replace(attention, spec=replace(attention.spec, phase=AttentionPhase.PREFILL)) + with pytest.raises(ValueError, match="phases differ"): + attention.backend.forward(attention, context, inputs=None) + + +class _OutputBackend(BaseAttentionBackend): + def _forward_decode(self, layer, context, inputs, *, out=None): + return inputs + + +@pytest.mark.parametrize("strided", [False, True]) +def test_output_copy_and_padded_view_preserve_auxiliary_result(strided): + backend = _OutputBackend(model_stub(), make_spec()) + context = make_context(backend, prepare=True) + layer = make_layer_attention(backend, context, params=None) + out = torch.empty(2, 8 if strided else 4, 8)[:, :4] + gate = torch.ones(2) + output, actual_gate = layer.backend.forward(layer, context, (torch.ones(2, 32), gate), out=out) + assert output is out and actual_gate is gate + assert torch.all(out == 1) + valid_heads = out[:, :2] + output, actual_gate = layer.backend.forward(layer, context, (valid_heads, gate), out=out) + assert output is valid_heads and actual_gate is gate + + +@pytest.mark.parametrize("invalid", ["dtype", "device", "size", "view_bounds"]) +def test_output_rejects_incompatible_buffer(invalid): + backend = _OutputBackend(model_stub(), make_spec()) + context = make_context(backend, prepare=True) + layer = make_layer_attention(backend, context, params=None) + out = torch.empty(2, 4) + result = torch.ones_like(out) + if invalid == "dtype": + result = result.double() + elif invalid == "device": + result = result.to("meta") + elif invalid == "size": + result = result[:1] + else: + result, out = out, out[:1] + with pytest.raises(ValueError, match="out"): + layer.backend.forward(layer, context, result, out=out) + + +def test_spec_configures_backend_before_planning(): + from lightllm.common.attention.api import AttentionExecutionContext, AttentionWorkspace + + model = model_stub(page_size=1, mtp_dynamic_verify=False) + spec = replace(make_spec(), cache=replace(make_spec().cache, page_size=16), causal=False, dynamic_verify=True) + backend = BaseAttentionBackend.create(model, spec) + assert backend.spec is spec + assert (backend.page_size, backend.causal, backend.dynamic_verify) == (16, False, True) + context = make_context(backend) + assert isinstance(context, AttentionExecutionContext) + assert isinstance(context.workspace, AttentionWorkspace) + assert context.backend_identity is backend + + +def test_binding_resolves_context_on_each_forward(monkeypatch): + backend = _OutputBackend(model_stub(), make_spec()) + first, second = make_context(backend), make_context(backend) + binding = make_layer_attention(backend, first, params=None) + contexts = {binding.group_id: first} + seen = [] + + def resolve(layer): + assert layer is binding + seen.append(contexts[layer.group_id]) + return contexts[layer.group_id] + + forwarded = [] + + def forward(layer, context, inputs, *, out=None): + assert layer is binding + forwarded.append(context) + return inputs + + monkeypatch.setattr(backend, "_forward_decode", forward) + execution = SimpleNamespace(attention_context=resolve) + result = torch.ones(2, 32) + out = torch.empty_like(result) + assert binding.forward(execution, result, out=out) is out + contexts[binding.group_id] = second + assert binding.forward(execution, result, out=out) is out + assert seen == [first, second] + assert all(actual is expected for actual, expected in zip(forwarded, (first, second))) + torch.testing.assert_close(out, result) diff --git a/unit_tests/common/attention/test_platform_runtime.py b/unit_tests/common/attention/test_platform_runtime.py new file mode 100644 index 0000000000..d8f0832d89 --- /dev/null +++ b/unit_tests/common/attention/test_platform_runtime.py @@ -0,0 +1,63 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention.api import AttentionPhase, GraphBinding +from lightllm.common.attention.backends._workspace import AttentionWorkspace +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAAttentionConfig, GQAFullAttentionState +from lightllm.platform import DeviceRuntime, get_device_runtime + + +def test_runtime_resolves_host_platform_without_querying_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: pytest.fail("queried CUDA for a CPU runtime")) + assert get_device_runtime(platform="cpu").device == torch.device("cpu") + with pytest.raises(ValueError, match="cannot operate"): + get_device_runtime("cuda:0", platform="cpu") + with pytest.raises(ValueError, match="cannot operate"): + get_device_runtime("cpu", platform="cuda") + + +def test_runtime_rejects_unimplemented_hardware(): + with pytest.raises(ValueError, match="CUDA and CPU"): + get_device_runtime(platform="musa") + with pytest.raises(ValueError, match="CUDA and CPU"): + get_device_runtime("npu:0") + + +def test_explicit_current_device_is_not_overridden_by_worker_default(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 1) + assert get_device_runtime("cuda").device == torch.device("cuda:1") + assert get_device_runtime("cuda:2").device == torch.device("cuda:2") + assert get_device_runtime().device == torch.device("cuda:0") + + +def test_captured_workspace_keeps_address_when_pool_grows(): + pool = AttentionWorkspace("cpu") + first, second = pool.fork(), pool.fork() + old = first.get("scratch", 8, torch.int32) + assert second.get("scratch", 8, torch.int32).data_ptr() == old.data_ptr() + first.freeze() + grown = pool.get("scratch", 64, torch.int32) + old.fill_(17) + grown.fill_(29) + assert first.get("scratch", 8, torch.int32).tolist() == [17] * 8 + with pytest.raises(RuntimeError, match="cannot grow"): + first.get("scratch", 64, torch.int32) + assert AttentionWorkspace("cpu").get("scratch", 8, torch.int32).data_ptr() != old.data_ptr() + + +def test_runtime_registration_keeps_platform_validation_in_adapter(monkeypatch): + import lightllm.platform as platform + + monkeypatch.setattr(platform, "_runtime_factories", {}) + calls = [] + adapter = object() + platform.register_device_runtime("musa", lambda device: calls.append(device) or adapter) + assert platform.get_device_runtime("musa:2", platform="musa") is adapter + assert calls == ["musa:2"] + with pytest.raises(ValueError, match="already registered"): + platform.register_device_runtime("musa", lambda device: adapter) diff --git a/unit_tests/common/attention/test_workspace.py b/unit_tests/common/attention/test_workspace.py new file mode 100644 index 0000000000..6e6137dc3c --- /dev/null +++ b/unit_tests/common/attention/test_workspace.py @@ -0,0 +1,60 @@ +import pytest +import torch + +from lightllm.common.attention.backends._workspace import AttentionWorkspace + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_workspace_cross_stream_reuse_waits_for_writer(): + workspace = AttentionWorkspace("cuda") + buffer = workspace.get("shared", 4096, torch.int32) + buffer.zero_() + producer, consumer = torch.cuda.Stream(), torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + consumer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + torch.cuda._sleep(2_000_000) + buffer.fill_(17) + workspace.finish() + with torch.cuda.stream(consumer): + workspace.begin() + result = buffer.clone() + workspace.finish() + consumer.synchronize() + assert result.eq(17).all() + + larger = workspace.get("shared", 8192, torch.int32) + assert larger.data_ptr() != buffer.data_ptr() + assert buffer.eq(17).all() + + +def test_shared_object_type_and_frozen_fork_lifetime(): + workspace = AttentionWorkspace("cpu") + fork = workspace.fork() + shared = fork.get_object("provider_plan", dict) + assert workspace.get_object("provider_plan", dict) is shared + with pytest.raises(TypeError, match="provider_plan requires list"): + workspace.get_object("provider_plan", list) + with pytest.raises(TypeError, match="bad requires dict"): + workspace.get_object("bad", lambda: [], expected_type=dict) + assert "bad" not in workspace._objects + buffer = fork.get("scratch", 4) + fork.freeze() + workspace.get("scratch", 8) + assert fork.get("scratch", 4) is buffer + with pytest.raises(RuntimeError, match="cannot grow"): + fork.get("scratch", 8) + with pytest.raises(RuntimeError, match="before capture"): + fork.get_object("next_bucket", dict) + added = workspace.get_object("next_bucket", dict) + assert fork.get_object("next_bucket", dict) is added + + +def test_cpu_stream_and_event_have_consistent_noop_lifecycle(): + workspace = AttentionWorkspace("cpu") + runtime = workspace.runtime + assert runtime.create_stream() is None + assert runtime.create_event() is None + with runtime.use_stream(runtime.create_stream()): + workspace.begin() + workspace.finish() diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index 622c8087ad..6c27ade302 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -26,6 +26,7 @@ def single_rank_environment(monkeypatch: pytest.MonkeyPatch): def runtime_start_args(monkeypatch: pytest.MonkeyPatch): import json from dataclasses import asdict + from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args From ffab7be8297ae4f99f172d70f3d6fd529fd3592c Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 11:53:45 +0800 Subject: [PATCH 2/9] refactor(attention): relocate kernels without changing execution --- .../backends/_shared/linear}/causal_conv1d.py | 0 .../_shared/linear}/causal_conv1d_mtp.py | 0 .../_shared/linear}/causal_conv1d_qkv.py | 0 .../backends/_shared/linear}/chunk_delta_h.py | 5 +-- .../backends/_shared/linear}/cumsum.py | 5 +-- .../backends/_shared/linear}/index.py | 3 +- .../backends/_shared/linear}/l2norm.py | 2 +- .../_shared/linear}/mtp_state_params.py | 0 .../backends/_shared/linear}/op.py | 2 +- .../_shared/linear}/replayssm_cursor.py | 0 .../_shared/linear}/replayssm_decode.py | 6 +-- .../_shared/linear}/replayssm_delta.py | 0 .../_shared/linear}/replayssm_fold.py | 2 +- .../backends/_shared/linear}/solve_tril.py | 7 ++-- .../_shared/linear}/ssm_acceptance.py | 0 .../backends/_shared/linear}/utils.py | 1 - .../prefill}/context_flashattention_nopad.py | 16 ++++---- .../gqa/flash}/gqa_flash_decoding_stage1.py | 8 ++-- .../gqa/flash}/gqa_flash_decoding_stage2.py | 0 .../gqa/flash}/gqa_flash_decoding_vsm.py | 12 +++--- .../gqa/gqa_decode_flashattention_nopad.py | 9 ++--- .../decode/gqa/mtp_diverse/__init__.py | 14 +++++++ .../gqa/mtp_diverse/mtp_diverse_attn.py | 9 ++++- .../gqa/mtp_diverse/stage1_single_token.py | 8 ++-- .../gqa/mtp_diverse/stage2_single_token.py | 6 ++- .../int4kv/int4kv_flash_decoding_stage1.py | 8 ++-- .../int8kv_flash_decoding_diverse_stage1.py | 11 +++--- .../int8kv_flash_decoding_diverse_stage2.py | 12 +++--- .../int8kv_flash_decoding_diverse_stage3.py | 0 .../normal/int8kv_flash_decoding_stage1.py | 9 +++-- .../normal/int8kv_flash_decoding_stage2.py | 0 .../mha/flash}/flash_decoding_stage1.py | 2 +- .../mha/flash}/flash_decoding_stage2.py | 0 .../token_attention_nopad_att1.py | 6 +-- .../token_attention_nopad_reduceV.py | 1 - .../token_attention_nopad_softmax.py | 1 - .../token_attention_softmax_and_reducev.py | 2 +- .../prefill}/context_attention_fwd_neo.py | 3 +- .../backends/gdn/replay/decode.py} | 5 +-- .../backends/gdn/triton/kernels}/chunk.py | 16 ++++---- .../backends/gdn/triton/kernels}/chunk_o.py | 7 ++-- .../triton/kernels}/chunk_scaled_dot_kkt.py | 5 +-- .../gdn/triton/kernels}/fused_gdn_gating.py | 1 - .../gdn/triton/kernels}/fused_recurrent.py | 3 +- .../gdn/triton/kernels}/gdn_decode_pack.py | 0 .../triton/kernels}/mtp_fused_recurrent.py | 5 ++- .../backends/gdn/triton/kernels}/wy_fast.py | 3 +- .../backends/kda/replay/decode.py} | 5 +-- .../backends/kda/triton/kernels}/kda.py | 18 ++++----- .../kda/triton/kernels}/kda_decode.py | 2 +- .../decode}/gqa_flash_decoding_config.py | 7 ++-- .../decode}/gqa_flash_decoding_stage1.py | 1 + .../decode}/gqa_flash_decoding_stage2.py | 3 +- .../mla/triton/kernels/prefill/__init__.py | 1 + .../context_flashattention_nopad_with_v.py | 10 ++--- .../bounded_att/gdn/replay_gdn_state.py | 6 +-- .../bounded_att/kda/replay_kda_state.py | 6 +-- .../common/basemodel/attention/linear/gdn.py | 32 +++++++-------- .../common/basemodel/attention/linear/kda.py | 18 ++++----- .../basemodel/attention/linear/replay_gdn.py | 27 ++++--------- .../basemodel/attention/linear/replay_kda.py | 24 ++++-------- .../common/basemodel/attention/triton/fp.py | 33 +++++++++------- .../basemodel/attention/triton/int4kv.py | 12 ++++-- .../basemodel/attention/triton/int8kv.py | 16 ++++---- .../common/basemodel/attention/triton/mla.py | 10 +++-- .../common/basemodel/attention/triton/neo.py | 9 +++-- .../gqa/flash_decoding/gqa_flash_decoding.py | 8 +++- .../decode_att/gqa/mtp_diverse/__init__.py | 19 ++++++--- .../int4kv/ppl_int4kv_flash_decoding.py | 15 +++++-- .../int8kv/int8kv_flash_decoding_diverse.py | 13 +++++-- .../int8kv/normal/int8kv_flash_decoding.py | 13 +++++-- .../mha/flash_decoding/flash_decoding.py | 8 +++- .../triton_kernel/linear_att/__init__.py | 19 ++++++--- .../linear_att/fla/ops/__init__.py | 10 ++++- .../mla_att/decode_att/gqa_flash_decoding.py | 24 ++++++++---- .../mla_att/prefill_att/__init__.py | 10 ++++- test/benchmark/kernels/benchmark_kda_ssm.py | 4 +- .../llama_gqa_diverse_decode_stage1_tuning.py | 22 ++++++----- .../llama_gqa_diverse_decode_stage2_tuning.py | 18 +++++---- .../test_quantized_full_state.py | 31 +++++++-------- .../mtp_diverse/test_mtp_diverse_autotune.py | 12 +++--- .../gqa/mtp_diverse/test_stage2_autotune.py | 6 ++- .../gqa/test_flash_decoding_autotune.py | 16 ++++---- .../decode_att/int4kv/test_int4kv_autotune.py | 14 +++---- ...st_int8kv_flash_decoding_diverse_stage1.py | 13 ++++--- ...st_int8kv_flash_decoding_diverse_stage2.py | 11 +++--- ...st_int8kv_flash_decoding_diverse_stage3.py | 5 ++- .../decode_att/int8kv/test_normal_autotune.py | 16 ++++---- .../test_context_attention_fwd_neo.py | 5 ++- .../test_context_flashattention_nopad.py | 6 ++- .../test_context_flashattention_nopad1.py | 12 +++--- .../att/prefill_att/test_fa3_neo.py | 5 ++- .../linear_att/replay_history.py | 6 +-- .../linear_att/test_acceptance_fusion.py | 9 +++-- .../linear_att/test_causal_conv1d_mtp.py | 2 +- .../test_fused_recurrent_strided.py | 4 +- .../linear_att/test_kda_autotune.py | 13 +++---- .../test_mtp_fused_recurrent_autotune.py | 10 +---- .../test_mtp_fused_recurrent_equiv.py | 6 +-- .../linear_att/test_mtp_state_params.py | 7 +--- .../linear_att/test_replayssm.py | 8 ++-- .../linear_att/test_replayssm_execution.py | 19 +++++---- .../linear_att/test_ssm_autotune_memory.py | 4 +- .../common/test_linear_state_contract.py | 7 ++-- .../models/glm5_next/test_kda_conv_qkv.py | 11 ++---- unit_tests/models/glm5_next/test_kernels.py | 39 +++++++++---------- unit_tests/models/glm5_next/test_mtp.py | 20 +++------- .../models/sensenova_loop/test_gdn_state.py | 2 +- 108 files changed, 509 insertions(+), 448 deletions(-) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/causal_conv1d.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/causal_conv1d_mtp.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/causal_conv1d_qkv.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/chunk_delta_h.py (98%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/cumsum.py (97%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/index.py (93%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/l2norm.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/mtp_state_params.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/op.py (95%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_cursor.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_decode.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_delta.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_fold.py (97%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/solve_tril.py (98%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/ssm_acceptance.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/utils.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/prefill_att => attention/backends/_shared/prefill}/context_flashattention_nopad.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/gqa/flash}/gqa_flash_decoding_stage1.py (98%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/gqa/flash}/gqa_flash_decoding_stage2.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/gqa/flash}/gqa_flash_decoding_vsm.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/gqa_decode_flashattention_nopad.py (98%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/mtp_diverse/mtp_diverse_attn.py (83%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/mtp_diverse/stage1_single_token.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/mtp_diverse/stage2_single_token.py (98%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int4kv/int4kv_flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/int8kv_flash_decoding_diverse_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/int8kv_flash_decoding_diverse_stage2.py (98%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/int8kv_flash_decoding_diverse_stage3.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/normal/int8kv_flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/normal/int8kv_flash_decoding_stage2.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha/flash_decoding => attention/backends/dense/triton/kernels/decode/mha/flash}/flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha/flash_decoding => attention/backends/dense/triton/kernels/decode/mha/flash}/flash_decoding_stage2.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_nopad_att1.py (98%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_nopad_reduceV.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_nopad_softmax.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_softmax_and_reducev.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/prefill_att => attention/backends/dense/triton/kernels/prefill}/context_attention_fwd_neo.py (99%) rename lightllm/common/{basemodel/attention/linear/replay_gdn_decode.py => attention/backends/gdn/replay/decode.py} (98%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/chunk.py (91%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/chunk_o.py (95%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/chunk_scaled_dot_kkt.py (97%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/gdn/triton/kernels}/fused_gdn_gating.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/fused_recurrent.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/gdn/triton/kernels}/gdn_decode_pack.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/gdn/triton/kernels}/mtp_fused_recurrent.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/wy_fast.py (97%) rename lightllm/common/{basemodel/attention/linear/replay_kda_decode.py => attention/backends/kda/replay/decode.py} (98%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/kda/triton/kernels}/kda.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/kda/triton/kernels}/kda_decode.py (98%) rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/gqa_flash_decoding_config.py (99%) rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/gqa_flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/gqa_flash_decoding_stage2.py (99%) rename lightllm/common/{basemodel/triton_kernel/mla_att/prefill_att => attention/backends/mla/triton/kernels/prefill}/context_flashattention_nopad_with_v.py (100%) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_mtp.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d_mtp.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_mtp.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d_mtp.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d_qkv.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d_qkv.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py b/lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py rename to lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py index 849bccb9e9..e972ec2844 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py +++ b/lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py @@ -9,12 +9,11 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices, prepare_chunk_offsets -from .op import exp, safe_exp +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices, prepare_chunk_offsets +from lightllm.common.attention.backends._shared.linear.op import exp, safe_exp from lightllm.common.triton_utils.autotuner import autotune NUM_WARPS = [2, 4, 8, 16] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py b/lightllm/common/attention/backends/_shared/linear/cumsum.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py rename to lightllm/common/attention/backends/_shared/linear/cumsum.py index 6331e1602d..d69e44ae4c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py +++ b/lightllm/common/attention/backends/_shared/linear/cumsum.py @@ -8,12 +8,11 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .utils import check_shared_mem, input_guard +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.utils import check_shared_mem, input_guard from lightllm.common.triton_utils.autotuner import autotune BS_LIST = [32, 64] if check_shared_mem() else [16, 32] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py b/lightllm/common/attention/backends/_shared/linear/index.py similarity index 93% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py rename to lightllm/common/attention/backends/_shared/linear/index.py index 878e2296f4..4b9ec5c001 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py +++ b/lightllm/common/attention/backends/_shared/linear/index.py @@ -8,10 +8,9 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 import torch - import triton -from .utils import tensor_cache +from lightllm.common.attention.backends._shared.linear.utils import tensor_cache @tensor_cache diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py b/lightllm/common/attention/backends/_shared/linear/l2norm.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py rename to lightllm/common/attention/backends/_shared/linear/l2norm.py index 29f892ef26..fcad017d9b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py +++ b/lightllm/common/attention/backends/_shared/linear/l2norm.py @@ -10,9 +10,9 @@ import os import torch - import triton import triton.language as tl + from lightllm.common.triton_utils.autotuner import autotune BT_LIST = [8, 16, 32, 64, 128] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py rename to lightllm/common/attention/backends/_shared/linear/mtp_state_params.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py b/lightllm/common/attention/backends/_shared/linear/op.py similarity index 95% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py rename to lightllm/common/attention/backends/_shared/linear/op.py index 2f69aa981d..8cbde95b6b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py +++ b/lightllm/common/attention/backends/_shared/linear/op.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from .utils import is_gather_supported +from lightllm.common.attention.backends._shared.linear.utils import is_gather_supported exp = tl.exp log = tl.log diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/attention/backends/_shared/linear/replayssm_cursor.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_cursor.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_decode.py index f3abbeb020..b7f138c844 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from .replayssm_delta import _apply_accepted_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix def check_replay_sync_mark( @@ -357,7 +357,7 @@ def _project_qk( row_valid = valid[:, None] qk = tl.where(row_valid, q_rows + k_rows, 0) qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] - qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + qk *= tl.where(tt < BW, K**-0.5, 1.0)[:, None] qk = tl.where(row_valid, qk, 0) if KDA: history = ll < n @@ -765,7 +765,7 @@ def _replay_each_token( q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) raw_k = tl.load(StepK + t * SK + h * K + kk, kk < K, 0).to(tl.float32) v = tl.load(StepV + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) - q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K ** -0.5) + q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K**-0.5) k = raw_k / tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) if KDA: x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/attention/backends/_shared/linear/replayssm_delta.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_delta.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_fold.py index 42aad6edc2..2e491a038b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from .replayssm_delta import _apply_accepted_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/attention/backends/_shared/linear/solve_tril.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py rename to lightllm/common/attention/backends/_shared/linear/solve_tril.py index 3512788934..1b0a73efe0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/attention/backends/_shared/linear/solve_tril.py @@ -12,13 +12,12 @@ from typing import Optional import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import make_tensor_descriptor -from .utils import input_guard, is_amd, is_tma_supported +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import make_tensor_descriptor +from lightllm.common.attention.backends._shared.linear.utils import input_guard, is_amd, is_tma_supported def _ensure_triton_allocator(): diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/attention/backends/_shared/linear/ssm_acceptance.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py rename to lightllm/common/attention/backends/_shared/linear/ssm_acceptance.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py b/lightllm/common/attention/backends/_shared/linear/utils.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py rename to lightllm/common/attention/backends/_shared/linear/utils.py index cd7c2e3aeb..e3f6c65510 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py +++ b/lightllm/common/attention/backends/_shared/linear/utils.py @@ -16,7 +16,6 @@ from typing import Any, Literal import torch - import triton logger = logging.getLogger(__name__) diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py b/lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py rename to lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py index ef38e39d6c..4ec6a1b0dd 100644 --- a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py +++ b/lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py @@ -1,11 +1,11 @@ # From : https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F from lightllm.utils.device_utils import is_tesla @@ -164,7 +164,7 @@ def context_attention_fwd( # 计算scale系数, 并乘以 1/log(2) = 1.4426950408889634, # 算子内部使用 tl.math.exp2 来使计算与标准attention等价。 - sm_scale = 1.0 / (Lq ** 0.5) * 1.4426950408889634 + sm_scale = 1.0 / (Lq**0.5) * 1.4426950408889634 batch, head = b_seq_len.shape[0], q.shape[1] kv_group_num = q.shape[1] // k.shape[1] @@ -329,7 +329,6 @@ def _fwd_kernel_no_prompt_cache( @torch.no_grad() def context_attention_fwd_no_prompt_cache(q, k, v, o, b_start_loc, b_seq_len, max_input_len): - BLOCK_M = 128 if not is_tesla() else 64 # shape constraints Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -345,7 +344,7 @@ def context_attention_fwd_no_prompt_cache(q, k, v, o, b_start_loc, b_seq_len, ma # 计算scale系数, 并乘以 1/log(2) = 1.4426950408889634, # 算子内部使用 tl.math.exp2 来使计算与标准attention等价。 - sm_scale = 1.0 / (Lq ** 0.5) * 1.4426950408889634 + sm_scale = 1.0 / (Lq**0.5) * 1.4426950408889634 batch, head = b_seq_len.shape[0], q.shape[1] kv_group_num = q.shape[1] // k.shape[1] @@ -524,7 +523,7 @@ def context_attention_fwd_contiguous_kv( # 计算scale系数, 并乘以 1/log(2) = 1.4426950408889634, # 算子内部使用 tl.math.exp2 来使计算与标准attention等价。 - sm_scale = 1.0 / (Lq ** 0.5) * 1.4426950408889634 + sm_scale = 1.0 / (Lq**0.5) * 1.4426950408889634 batch, head = b_seq_len.shape[0], q.shape[1] kv_group_num = q.shape[1] // k.shape[1] @@ -566,7 +565,6 @@ def context_attention_fwd_contiguous_kv( def torch_context_attention_fwd(q, k, v, o, b_req_idx, b_start_loc, b_seq_len, b_prompt_cache_len, req_to_token_indexs): - batch = b_start_loc.shape[0] print(q.shape) for i in range(batch): @@ -606,8 +604,8 @@ def torch_context_attention_fwd(q, k, v, o, b_req_idx, b_start_loc, b_seq_len, b def test(): - import torch import numpy as np + import torch Z, H, N_CTX, D_HEAD = 16, 16, 2048, 128 dtype = torch.float16 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py index 767cee9870..91eee96a62 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @@ -258,7 +260,7 @@ def flash_decode_stage1( if Lk >= 256: BLOCK_N = min(BLOCK_N, 16) assert BLOCK_SEQ % BLOCK_N == 0 - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) batch, kv_head_num = B_req_idx.shape[0], k.shape[1] block_num = mid_out.shape[2] grid = (batch, kv_head_num, block_num) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py index 141587ff38..63f083feab 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py @@ -1,11 +1,13 @@ +from functools import lru_cache +from typing import Dict + import torch import triton import triton.language as tl +from frozendict import frozendict + from lightllm.common.kernel_config import KernelConfigs from lightllm.utils.device_utils import calcu_kernel_best_vsm_count -from frozendict import frozendict -from functools import lru_cache -from typing import Dict class GQAVSMDecodeAttentionKernelConfig(KernelConfigs): @@ -261,7 +263,7 @@ def gqa_token_decode_attention_flash_decoding_vsm_stage1( q_head_num, kv_head_num, batch_size, - **run_config + **run_config, ): grid = (num_vsm,) @@ -420,7 +422,7 @@ def gqa_token_decode_attention_flash_decoding_vsm( batch_size, q_head_num, q_head_dim = q.shape kv_head_num = k.shape[1] gqa_group_size = q_head_num // kv_head_num - sm_scale = 1.0 / (q_head_dim ** 0.5) + sm_scale = 1.0 / (q_head_dim**0.5) if not run_config: if torch.cuda.is_current_stream_capturing(): diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py index 5ce5d193f5..f5bd9c0fa1 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py @@ -1,9 +1,9 @@ -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F @triton.jit @@ -115,14 +115,13 @@ def _fwd_kernel( @torch.no_grad() def gqa_decode_attention_fwd(q, k, v, o, req_to_tokens, b_req_idx, b_seq_len): - # shape constraints Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] assert Lq == Lk and Lk == Lv assert Lk in {16, 32, 64, 128, 256} BLOCK = 16 if Lk == 256 else 32 - sm_scale = 1.0 / (Lq ** 0.5) # 计算scale系数 + sm_scale = 1.0 / (Lq**0.5) # 计算scale系数 batch = b_req_idx.shape[0] kv_group_num = q.shape[1] // k.shape[1] kv_head_num = k.shape[1] diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py index e69de29bb2..f8d0ac0a7e 100644 --- a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py @@ -0,0 +1,14 @@ +""" +MTP Diverse Attention Module + +MTP (Multi-Token Prediction) Diverse Attention 的实现。 +""" +from .mtp_diverse_attn import token_decode_attention_mtp_diverse_single_token +from .stage1_single_token import mtp_diverse_stage1_single_token +from .stage2_single_token import mtp_diverse_stage2_single_token + +__all__ = [ + "token_decode_attention_mtp_diverse_single_token", + "mtp_diverse_stage1_single_token", + "mtp_diverse_stage2_single_token", +] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py similarity index 83% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py index 13c7e8070e..dc98bb3544 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py @@ -1,6 +1,11 @@ import torch -from .stage1_single_token import mtp_diverse_stage1_single_token -from .stage2_single_token import mtp_diverse_stage2_single_token + +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage1_single_token import ( + mtp_diverse_stage1_single_token, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage2_single_token import ( + mtp_diverse_stage2_single_token, +) from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py index 143b3c1ce1..7ecab20d73 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py @@ -12,11 +12,13 @@ - 每个请求单独做可见性检查（基于各自的 seq_len） - 中间结果按 kv block 存储，供 Stage2 聚合 """ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.device_utils import is_hopper from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @@ -319,7 +321,7 @@ def mtp_diverse_stage1_single_token( grid_num = mid_out.shape[2] - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) # 固定 grid（graph friendly） grid = (grid_num, kv_head_num, batch) _fwd_kernel_mtp_diverse_stage1_single_token[grid]( diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py index 5a56426932..4d382439df 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py @@ -6,11 +6,13 @@ - 根据 seq_len 确定需要聚合的 kv block 数量 - 使用 flash attention reweighting 公式 """ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py index 27cb97d9a0..d90a233484 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @@ -279,7 +281,7 @@ def int4kv_flash_decode_stage1( if Lk == 256: BLOCK_N = min(BLOCK_N, 16) assert BLOCK_SEQ % BLOCK_N == 0 - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) batch, kv_head_num = B_req_idx.shape[0], k.shape[1] grid_block_num = mid_out.shape[2] grid = (batch, kv_head_num, grid_block_num) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py index fb122a79da..3bc7112ce0 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py @@ -1,11 +1,12 @@ +from functools import lru_cache +from typing import Dict, Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.kernel_config import KernelConfigs from frozendict import frozendict -from functools import lru_cache -from typing import Dict + +from lightllm.common.kernel_config import KernelConfigs class GQADiverseDecodeStage1KernelConfig(KernelConfigs): @@ -280,7 +281,7 @@ def flash_decode_stage1( if Lk == 256: BLOCK_N = min(BLOCK_N, 16) assert BLOCK_SEQ % BLOCK_N == 0 - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) batch, kv_head_num = B_req_idx.shape[0], k.shape[1] grid = (batch, kv_head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ)) gqa_group_size = q.shape[1] // k.shape[1] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py index 963139b36b..1b46d45b6b 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py @@ -1,13 +1,13 @@ +from functools import lru_cache +from typing import Dict, Optional + import torch import triton import triton.language as tl -from typing import Optional +from frozendict import frozendict from lightllm.common.kernel_config import KernelConfigs -from frozendict import frozendict -from functools import lru_cache -from typing import Dict -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import Autotuner, autotune class GQADiverseDecodeStage2KernelConfig(KernelConfigs): @@ -254,7 +254,7 @@ def flash_decode_stage2( if Lk == 256: BLOCK_N = min(BLOCK_N, 16) assert BLOCK_SEQ % BLOCK_N == 0 - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) batch, kv_head_num = B_req_idx.shape[0], k.shape[1] grid = (batch, kv_head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ)) gqa_group_size = q.shape[1] // k.shape[1] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage3.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage3.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py index a05db282d7..87a7839914 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @@ -266,7 +268,6 @@ def flash_decode_stage1( ): """ """ if not run_config: - run_config = { "BLOCK_N": 16, "num_warps": 2, @@ -287,7 +288,7 @@ def flash_decode_stage1( if Lk == 256: BLOCK_N = min(BLOCK_N, 16) assert BLOCK_SEQ % BLOCK_N == 0 - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) batch, kv_head_num = B_req_idx.shape[0], k.shape[1] gqa_group_size = q.shape[1] // k.shape[1] assert triton.next_power_of_2(Lk) == Lk diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage1.py index 196ded61a5..981d9e466b 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage1.py @@ -113,7 +113,7 @@ def flash_decode_stage1( Lq, Lk = q.shape[-1], k.shape[-1] assert Lq == Lk assert Lk in {16, 32, 64, 128, 256} - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) batch, head_num = B_req_idx.shape[0], q.shape[1] grid = (batch, head_num, triton.cdiv(max_len_in_batch, BLOCK_SEQ)) gqa_group_size = q.shape[1] // k.shape[1] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py index 9de2b82057..fe1dfdd614 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py @@ -1,8 +1,8 @@ -import torch +import math +import torch import triton import triton.language as tl -import math @triton.jit @@ -74,7 +74,7 @@ def token_att_fwd(q, k, att_out, Req_to_tokens, B_req_idx, B_Start_Loc, B_Seqlen Lq, Lk = q.shape[-1], k.shape[-1] assert Lq == Lk assert Lk in {16, 32, 64, 128, 256} - sm_scale = 1.0 / (Lk ** 0.5) + sm_scale = 1.0 / (Lk**0.5) batch, head_num = B_req_idx.shape[0], q.shape[1] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py index 96a5b26dd6..863ac322de 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py @@ -1,5 +1,4 @@ import torch - import triton import triton.language as tl diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py index 0bb6410e13..08147b1e99 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py @@ -68,7 +68,6 @@ def token_softmax_fwd(Logics, B_Start_Loc, B_Seqlen, Prob_Out, max_input_len): def test1(): - import torch B, N_CTX, H, D = 4, 1025, 12, 128 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py index d963f8582a..6248dac46d 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py @@ -1,7 +1,7 @@ import torch +import torch.nn.functional as F import triton import triton.language as tl -import torch.nn.functional as F @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py rename to lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py index 1016e5a3c1..41cdd63c75 100644 --- a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py @@ -1,4 +1,5 @@ import math + import torch import triton import triton.language as tl @@ -156,7 +157,7 @@ def context_attention_fwd_neo( Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] assert Lq == Lk and Lk == Lv assert Lk in {16, 32, 64, 128, 256} - sm_scale = 1.0 / (Lq ** 0.5) * 1.4426950408889634 + sm_scale = 1.0 / (Lq**0.5) * 1.4426950408889634 batch, head = b_seq_len.shape[0], q.shape[1] kv_group_num = q.shape[1] // k.shape[1] diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/attention/backends/gdn/replay/decode.py similarity index 98% rename from lightllm/common/basemodel/attention/linear/replay_gdn_decode.py rename to lightllm/common/attention/backends/gdn/replay/decode.py index ac7af9dbdc..6dda8fc867 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/attention/backends/gdn/replay/decode.py @@ -10,11 +10,8 @@ import torch import triton +from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay, check_replay_sync_mark from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( - _decode_replay, - check_replay_sync_mark, -) from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk.py similarity index 91% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk.py index 7b3067bbfb..919e466472 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk.py @@ -10,14 +10,14 @@ import torch from einops import rearrange -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .chunk_o import chunk_fwd_o -from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from .cumsum import chunk_local_cumsum -from .l2norm import l2norm_fwd -from .solve_tril import solve_tril -from .utils import SUPPRESS_LEVEL, input_guard -from .wy_fast import recompute_w_u_fwd +from lightllm.common.attention.backends._shared.linear.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from lightllm.common.attention.backends._shared.linear.cumsum import chunk_local_cumsum +from lightllm.common.attention.backends._shared.linear.l2norm import l2norm_fwd +from lightllm.common.attention.backends._shared.linear.solve_tril import solve_tril +from lightllm.common.attention.backends._shared.linear.utils import SUPPRESS_LEVEL, input_guard +from lightllm.common.attention.backends.gdn.triton.kernels.chunk_o import chunk_fwd_o +from lightllm.common.attention.backends.gdn.triton.kernels.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from lightllm.common.attention.backends.gdn.triton.kernels.wy_fast import recompute_w_u_fwd def chunk_gated_delta_rule_fwd( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py similarity index 95% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py index fc49763ecd..d8c4fa91a2 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py @@ -11,13 +11,12 @@ import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import exp, safe_exp -from .utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import exp, safe_exp +from lightllm.common.attention.backends._shared.linear.utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper from lightllm.common.triton_utils.autotuner import autotune BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py index 60a594c078..6aacf292dd 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py @@ -9,12 +9,11 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import safe_exp +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import safe_exp from lightllm.common.triton_utils.autotuner import autotune diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py b/lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py rename to lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py index 88febaffc6..a2f648753e 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py @@ -67,7 +67,6 @@ def fused_gdn_gating( threshold: float = 20.0, run_config: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - if run_config is None: run_config = {"BLK_HEADS": 8, "num_warps": 1} diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py rename to lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py index 5dfbd6e4ab..1ee725b03d 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py @@ -9,11 +9,10 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .op import exp +from lightllm.common.attention.backends._shared.linear.op import exp @triton.heuristics( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/gdn_decode_pack.py b/lightllm/common/attention/backends/gdn/triton/kernels/gdn_decode_pack.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/gdn_decode_pack.py rename to lightllm/common/attention/backends/gdn/triton/kernels/gdn_decode_pack.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py b/lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py rename to lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py index 5f8cab4693..57a9d2a148 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py @@ -11,12 +11,13 @@ # https://github.com/fla-org/flash-linear-attention # ruff: noqa: E501 +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, AutotuneKernelType +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune # --------------------------------------------------------------------------- # Triton kernel diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py b/lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py rename to lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py index 08bb00e644..4a8882d5fe 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py @@ -10,11 +10,10 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/attention/backends/kda/replay/decode.py similarity index 98% rename from lightllm/common/basemodel/attention/linear/replay_kda_decode.py rename to lightllm/common/attention/backends/kda/replay/decode.py index 3704ba4fd6..5c3d38fedc 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/attention/backends/kda/replay/decode.py @@ -10,10 +10,7 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( - _decode_replay, - check_replay_sync_mark, -) +from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay, check_replay_sync_mark from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/attention/backends/kda/triton/kernels/kda.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py rename to lightllm/common/attention/backends/kda/triton/kernels/kda.py index 8555b83952..58fb54b40b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py +++ b/lightllm/common/attention/backends/kda/triton/kernels/kda.py @@ -24,16 +24,14 @@ import torch import triton import triton.language as tl - -from lightllm.common.triton_utils.autotuner import autotune - -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .cumsum import chunk_local_cumsum -from .index import prepare_chunk_indices -from .l2norm import l2norm_fwd from triton.language import exp2, log -from .solve_tril import solve_tril +from lightllm.common.attention.backends._shared.linear.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from lightllm.common.attention.backends._shared.linear.cumsum import chunk_local_cumsum +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.l2norm import l2norm_fwd +from lightllm.common.attention.backends._shared.linear.solve_tril import solve_tril +from lightllm.common.triton_utils.autotuner import autotune FLA_CHUNK_SIZE = 64 RCP_LN2 = 1.4426950216293335 @@ -1284,7 +1282,7 @@ def chunk_kda( assert cu_seqlens is not None, "cu_seqlens is required for packed KDA prefill" if scale is None: head_dim = k.shape[-1] - scale = head_dim ** -0.5 + scale = head_dim**-0.5 if use_qk_l2norm_in_kernel: q = l2norm_fwd(q.contiguous()) @@ -1396,7 +1394,7 @@ def chunk_kda_with_fused_gate( assert cu_seqlens is not None, "cu_seqlens is required for packed KDA prefill" if scale is None: head_dim = k.shape[-1] - scale = head_dim ** -0.5 + scale = head_dim**-0.5 if use_qk_l2norm_in_kernel: q = l2norm_fwd(q.contiguous()) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py b/lightllm/common/attention/backends/kda/triton/kernels/kda_decode.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py rename to lightllm/common/attention/backends/kda/triton/kernels/kda_decode.py index 7c176e7f8b..c039db4bfb 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py +++ b/lightllm/common/attention/backends/kda/triton/kernels/kda_decode.py @@ -50,7 +50,7 @@ def _kda_decode( q = tl.load(Q + token * SQ + head * D + ki).to(tl.float32) k = tl.load(K + token * SK + head * D + ki).to(tl.float32) v = tl.load(V + token * SV + head * D + vi).to(tl.float32) - q *= tl.rsqrt(tl.sum(q * q, 0) + 1e-6) * (D ** -0.5) + q *= tl.rsqrt(tl.sum(q * q, 0) + 1e-6) * (D**-0.5) k *= tl.rsqrt(tl.sum(k * k, 0) + 1e-6) gate = tl.load(G + token * SG + head * D + ki).to(tl.float32) decay = tl.exp(LOWER * tl.sigmoid(amplitude * (gate + bias))) diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py index be99ca9bfc..c573b907f3 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py @@ -1,8 +1,10 @@ -from lightllm.common.kernel_config import KernelConfigs -from frozendict import frozendict from functools import lru_cache from typing import Dict +from frozendict import frozendict + +from lightllm.common.kernel_config import KernelConfigs + class MlaDecodeAttentionKernelConfig(KernelConfigs): kernel_name: str = "mla_decode_attentnion" @@ -51,7 +53,6 @@ def try_to_get_best_config( def save_config( cls, q_head_num: int, q_head_dim: int, q_rope_dim: int, out_dtype: str, config_json: Dict[int, Dict[int, Dict]] ): - key_params = { "q_head_num": q_head_num, "q_head_dim": q_head_dim, diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py index f5909fffde..63607b4594 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py @@ -1,6 +1,7 @@ import torch import triton import triton.language as tl + from lightllm.utils.device_utils import calcu_kernel_best_vsm_count diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py index 5b5e7b747c..f6f328f3e4 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py @@ -1,4 +1,5 @@ import os + import torch import triton import triton.language as tl @@ -62,7 +63,7 @@ def flash_decode_stage2( mid_out_logexpsum, B_Seqlen, Out, - **run_config + **run_config, ): if run_config: num_warps = run_config["stage2_num_warps"] diff --git a/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py index e69de29bb2..5725bed2e7 100644 --- a/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py @@ -0,0 +1 @@ +from .context_flashattention_nopad_with_v import context_attention_fwd_with_v diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py b/lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py rename to lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py index be06351823..a6b4c5fe94 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py @@ -1,9 +1,10 @@ -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F + from lightllm.utils.device_utils import is_tesla @@ -145,7 +146,6 @@ def context_attention_fwd_with_v( max_input_len, softmax_scale, ): - BLOCK = 128 if not is_tesla() else 64 q_nope_dim = q_nope.shape[-1] q_rope_dim = q_rope.shape[-1] @@ -204,8 +204,8 @@ def context_attention_fwd_with_v( if __name__ == "__main__": - import torch import flashinfer + import torch Z, N_CTX, H, D_HEAD, ROPE_HEAD = 32, 1024, 16, 128, 64 dtype = torch.bfloat16 diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 137ce8ad32..2d9540645a 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -11,9 +11,9 @@ import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 2c73e97d6f..8897218fbc 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -11,9 +11,9 @@ import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..5105e83584 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,23 +1,22 @@ -from lightllm.common.attention_state.bounded_att import GDNState import dataclasses import weakref -import torch from abc import ABC, abstractmethod from typing import TYPE_CHECKING -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn -from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + +import torch + +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( - fused_recurrent_gated_delta_rule, -) +from lightllm.common.attention_state.bounded_att import GDNState +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import fused_recurrent_gated_delta_rule +from lightllm.utils.envs_utils import get_env_start_args + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState if TYPE_CHECKING: from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -37,7 +36,6 @@ def get_prefill_kernel(self): pass def _init_linear_layer_metadata(self, network_config, tp_world_size): - self.mtp_step = get_env_start_args().mtp_step # Linear attention specific dimensions @@ -109,7 +107,6 @@ def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDe @dataclasses.dataclass class LinearAttPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None @@ -207,7 +204,6 @@ def _gdn_prefill_kernel( @dataclasses.dataclass class LinearAttDecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None @@ -389,7 +385,7 @@ def _gdn_mtp_kernel( infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", ): - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( causal_conv1d_update as causal_conv1d_update_mtp, ) diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 11f801b62a..093ec03bef 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -10,6 +10,14 @@ import torch +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from lightllm.common.attention_state.bounded_att import KDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, @@ -17,16 +25,6 @@ BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import causal_conv1d_qkv -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) if TYPE_CHECKING: from lightllm.common.basemodel.basemodel import TpPartBaseModel diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 8e4a081f73..fd936ff757 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -12,34 +12,23 @@ import torch -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ( - ReplayGDNState, +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, ) +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.gdn.replay.decode import replay +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_fn, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import ( - fused_gdn_gating, -) -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import ( - conv_pack_gdn_decode_inputs, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) from lightllm.utils.envs_utils import get_env_start_args -from .replay_gdn_decode import replay - if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py index af931aa477..6275b5e4c0 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -11,30 +11,20 @@ import torch -from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( - ReplayKDAState, +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, ) +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends.kda.replay.decode import replay +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState, ) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( - chunk_kda_with_fused_gate, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) - -from .replay_kda_decode import replay if TYPE_CHECKING: from lightllm.common.basemodel.infer_struct import InferStateInfo diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py index bab3f0ad38..8da44d6ecc 100644 --- a/lightllm/common/basemodel/attention/triton/fp.py +++ b/lightllm/common/basemodel/attention/triton/fp.py @@ -1,10 +1,13 @@ import dataclasses import weakref -import torch -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from typing import Optional + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState + class TritonAttBackend(BaseAttBackend): def create_att_prefill_state(self, infer_state) -> "TritonPrefillAttState": @@ -69,7 +72,9 @@ def _nomarl_prefill_att( att_control: AttControl = AttControl(), alloc_func=torch.empty, ): - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd, + ) if att_control.use_sliding_window: sliding_window = att_control.sliding_window @@ -231,9 +236,7 @@ def _spec_decode_gqa_att( v: torch.Tensor, alloc_func=torch.empty, ): - from ...triton_kernel.att.decode_att.gqa.mtp_diverse import ( - token_decode_attention_mtp_diverse_single_token, - ) + from ...triton_kernel.att.decode_att.gqa.mtp_diverse import token_decode_attention_mtp_diverse_single_token out = token_decode_attention_mtp_diverse_single_token( q=q, @@ -257,7 +260,7 @@ def _normal_decode_gqa_flash_decoding_att_vsm( alloc_func=torch.empty, ): # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_vsm import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_vsm import ( gqa_token_decode_attention_flash_decoding_vsm, ) @@ -282,7 +285,9 @@ def _normal_decode_gqa_att( alloc_func=torch.empty, ): # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from ...triton_kernel.att.decode_att.gqa.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.gqa_decode_flashattention_nopad import ( + gqa_decode_attention_fwd, + ) out = alloc_func(q.shape, q.dtype) @@ -312,9 +317,11 @@ def _normal_decode_stage3_att( calcu_shape1 = (batch_size, q_head_num, head_dim) att_m_tensor = alloc_func((q_head_num, total_token_num), torch.float32) - from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_nopad_att1 import token_att_fwd + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.stage3_decode_att import ( + token_attention_nopad_att1, + ) - token_att_fwd( + token_attention_nopad_att1.token_att_fwd( q.view(calcu_shape1), k, att_m_tensor, @@ -326,11 +333,11 @@ def _normal_decode_stage3_att( ) o_tensor = alloc_func(q.shape, q.dtype) - from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_softmax_and_reducev import ( - token_softmax_reducev_fwd, + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.stage3_decode_att import ( + token_attention_softmax_and_reducev, ) - token_softmax_reducev_fwd( + token_attention_softmax_and_reducev.token_softmax_reducev_fwd( att_m_tensor, v, o_tensor.view(calcu_shape1), diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py index 22e723de3b..269b512c1a 100644 --- a/lightllm/common/basemodel/attention/triton/int4kv.py +++ b/lightllm/common/basemodel/attention/triton/int4kv.py @@ -1,9 +1,12 @@ import dataclasses import weakref +from typing import Optional, Tuple + import torch + from lightllm.utils.envs_utils import get_env_start_args -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, Tuple + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState class Int4kvTritonAttBackend(BaseAttBackend): @@ -20,7 +23,6 @@ def create_att_decode_state(self, infer_state) -> "Int4kvTritonDecodeAttState": @dataclasses.dataclass class Int4kvTritonPrefillAttState(BasePrefillAttState): - # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。 b_kv_start_loc: torch.Tensor = None @@ -95,7 +97,9 @@ def _groupsize_quant_prefill_att( quant_group_size=self.backend.quant_group_size, ) - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) context_attention_fwd_contiguous_kv( q=q, diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py index 9da88626c8..85f17a13cc 100644 --- a/lightllm/common/basemodel/attention/triton/int8kv.py +++ b/lightllm/common/basemodel/attention/triton/int8kv.py @@ -1,11 +1,13 @@ import dataclasses import weakref -import torch -from lightllm.utils.envs_utils import get_env_start_args -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from typing import Optional, Tuple -from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel + +import torch + from lightllm.common.basemodel.triton_kernel.diverse_utils import build_diverse_shared_group_markers +from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel, get_env_start_args + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState class Int8kvTritonAttBackend(BaseAttBackend): @@ -22,7 +24,6 @@ def create_att_decode_state(self, infer_state) -> "Int8kvTritonDecodeAttState": @dataclasses.dataclass class Int8kvTritonPrefillAttState(BasePrefillAttState): - # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。 b_kv_start_loc: torch.Tensor = None @@ -94,7 +95,9 @@ def _groupsize_quant_prefill_att( quant_group_size=self.backend.quant_group_size, ) - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) context_attention_fwd_contiguous_kv( q=q, @@ -169,7 +172,6 @@ def diverse_decode_att( v_scale: torch.Tensor, alloc_func=torch.empty, ) -> torch.Tensor: - from ...triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import ( token_decode_attention_flash_decoding, ) diff --git a/lightllm/common/basemodel/attention/triton/mla.py b/lightllm/common/basemodel/attention/triton/mla.py index 036e58ae70..6f5624d0f4 100644 --- a/lightllm/common/basemodel/attention/triton/mla.py +++ b/lightllm/common/basemodel/attention/triton/mla.py @@ -1,9 +1,11 @@ import dataclasses import weakref -import torch -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl from typing import Tuple +import torch + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState + class MlaTritonAttBackend(BaseAttBackend): def create_att_prefill_state(self, infer_state) -> "MlaTritonPrefillAttState": @@ -77,7 +79,9 @@ def _contiguous_kv_prefill_att( att_control: AttControl, alloc_func=torch.empty, ): - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) k_nope, k_rope = k q_head_num = q.shape[1] diff --git a/lightllm/common/basemodel/attention/triton/neo.py b/lightllm/common/basemodel/attention/triton/neo.py index dfa31876f0..f4fb1cc030 100644 --- a/lightllm/common/basemodel/attention/triton/neo.py +++ b/lightllm/common/basemodel/attention/triton/neo.py @@ -1,10 +1,11 @@ import dataclasses import weakref -import torch from typing import Tuple +import torch + from ..base_att import AttControl -from .fp import TritonAttBackend, TritonPrefillAttState, TritonDecodeAttState +from .fp import TritonAttBackend, TritonDecodeAttState, TritonPrefillAttState class NeoTritonAttBackend(TritonAttBackend): @@ -25,7 +26,9 @@ def prefill_att( att_control: AttControl = AttControl(), alloc_func=torch.empty, ) -> torch.Tensor: - from ...triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, + ) from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo self.infer_state: NeoChatInferStateInfo diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py index d121656c73..bac4adb5ba 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py @@ -15,8 +15,12 @@ def gqa_token_decode_attention_flash_decoding( q_head_num, head_dim = q.shape[1], q.shape[2] calcu_shape1 = (batch_size, q_head_num, head_dim) - from .gqa_flash_decoding_stage1 import flash_decode_stage1 - from .gqa_flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( + flash_decode_stage1, + ) + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage2 import ( + flash_decode_stage2, + ) o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py index f8d0ac0a7e..9fb61c74f5 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py @@ -1,11 +1,18 @@ -""" -MTP Diverse Attention Module +"""Temporary legacy import shim for the attention migration. -MTP (Multi-Token Prediction) Diverse Attention 的实现。 +New code imports lightllm.common.attention.backends directly. Removed with the +runtime integration in PR #1721 (5/5), after remaining model callers migrate. """ -from .mtp_diverse_attn import token_decode_attention_mtp_diverse_single_token -from .stage1_single_token import mtp_diverse_stage1_single_token -from .stage2_single_token import mtp_diverse_stage2_single_token + +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.mtp_diverse_attn import ( + token_decode_attention_mtp_diverse_single_token, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage1_single_token import ( + mtp_diverse_stage1_single_token, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage2_single_token import ( + mtp_diverse_stage2_single_token, +) __all__ = [ "token_decode_attention_mtp_diverse_single_token", diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py index 936ca4dbbd..2bf2d4b86b 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py @@ -1,6 +1,7 @@ -import torch from typing import Optional +import torch + def token_decode_attention_flash_decoding( q: torch.Tensor, @@ -34,7 +35,9 @@ def token_decode_attention_flash_decoding( mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=q.dtype, device="cuda") - from .int4kv_flash_decoding_stage1 import int4kv_flash_decode_stage1 + from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv.int4kv_flash_decoding_stage1 import ( + int4kv_flash_decode_stage1, + ) int4kv_flash_decode_stage1( q=q.view(calcu_shape1), @@ -51,7 +54,11 @@ def token_decode_attention_flash_decoding( block_seq=BLOCK_SEQ, ) - from ..int8kv.normal.int8kv_flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( + int8kv_flash_decoding_stage2, + ) - flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + int8kv_flash_decoding_stage2.flash_decode_stage2( + mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ + ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py index c534795fdf..52ba55e30b 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py @@ -1,9 +1,16 @@ # 为 diverse mode 定制设计的 int8kv flash decoding attention 实现，可以实现更高效的多样性采样 import torch + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( + flash_decode_stage1, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( + flash_decode_stage2, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage3 import ( + flash_diverse_decode_stage3, +) from lightllm.common.basemodel.infer_struct import InferStateInfo -from .int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1 -from .int8kv_flash_decoding_diverse_stage2 import flash_decode_stage2 -from .int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3 from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py index de74f0728d..e14f62c28d 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py @@ -1,8 +1,14 @@ -import torch from typing import Optional + +import torch + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage1 import ( + flash_decode_stage1, +) +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( + flash_decode_stage2, +) from lightllm.common.basemodel.infer_struct import InferStateInfo -from .int8kv_flash_decoding_stage1 import flash_decode_stage1 -from .int8kv_flash_decoding_stage2 import flash_decode_stage2 @torch.no_grad() @@ -17,7 +23,6 @@ def token_decode_attention_flash_decoding( out: Optional[torch.Tensor] = None, alloc_tensor_func=torch.empty, ): - q_head_num = q.shape[1] head_dim = q.shape[2] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py index 6c50fc3927..b61f873348 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py @@ -8,8 +8,12 @@ def token_decode_attention_flash_decoding(q, infer_state, cache_k, cache_v, out= q_head_num, head_dim = q.shape[1], q.shape[2] calcu_shape1 = (batch_size, q_head_num, head_dim) - from .flash_decoding_stage1 import flash_decode_stage1 - from .flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage1 import ( + flash_decode_stage1, + ) + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage2 import ( + flash_decode_stage2, + ) o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py index aa7cfa9f4a..15d24413a0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py @@ -1,10 +1,17 @@ -"""Linear-attention / GDN triton kernels shared across hybrid models.""" +"""Temporary legacy import shim for the attention migration. + +New code imports lightllm.common.attention.backends directly. Removed with the +runtime integration in PR #1721 (5/5), after remaining model callers migrate. +""" + +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import causal_conv1d_update +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, +) -from .causal_conv1d import causal_conv1d_fn -from .causal_conv1d_mtp import causal_conv1d_update -from .fused_gdn_gating import fused_gdn_gating -from .gdn_decode_pack import conv_pack_gdn_decode_inputs -from .mtp_fused_recurrent import mtp_fused_recurrent_gated_delta_rule from .fla.ops import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule __all__ = [ diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py index cd3b0962a3..72b949c280 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py @@ -1,3 +1,9 @@ +"""Temporary legacy import shim for the attention migration. + +New code imports lightllm.common.attention.backends directly. Removed with the +runtime integration in PR #1721 (5/5), after remaining model callers migrate. +""" + # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang @@ -6,8 +12,8 @@ # The original source code was licensed under the MIT license and included # the following copyright notice: # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -from .chunk import chunk_gated_delta_rule -from .fused_recurrent import fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule __all__ = [ "chunk_gated_delta_rule", diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py index 063181d995..129e943f40 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py @@ -1,12 +1,16 @@ import os +from typing import List + import torch import torch.multiprocessing as mp import triton import triton.language as tl -from typing import List -from lightllm.utils.log_utils import init_logger -from .gqa_flash_decoding_config import MlaDecodeAttentionKernelConfig + +from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_config import ( + MlaDecodeAttentionKernelConfig, +) from lightllm.utils.device_utils import get_device_sm_count +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -41,8 +45,12 @@ def gqa_token_decode_attention_flash_decoding( BLOCK_N = run_config["BLOCK_N"] - from .gqa_flash_decoding_stage1 import flash_decode_stage1 - from .gqa_flash_decoding_stage2 import flash_decode_stage2 + from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_stage1 import ( + flash_decode_stage1, + ) + from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_stage2 import ( + flash_decode_stage2, + ) o_tensor = alloc_tensor_func(q_nope.shape, q_nope.dtype, q_nope.device) if out is None else out @@ -63,7 +71,7 @@ def gqa_token_decode_attention_flash_decoding( mid_o_logexpsum, softmax_scale, get_sm_count=True, - **run_config + **run_config, ) if not hasattr(infer_state, "decode_att_block_seq"): @@ -111,7 +119,7 @@ def gqa_token_decode_attention_flash_decoding( mid_o_logexpsum, softmax_scale, get_sm_count=False, - **run_config + **run_config, ) flash_decode_stage2( @@ -121,7 +129,7 @@ def gqa_token_decode_attention_flash_decoding( mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), - **run_config + **run_config, ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py index 5725bed2e7..76599b6c3f 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py @@ -1 +1,9 @@ -from .context_flashattention_nopad_with_v import context_attention_fwd_with_v +"""Temporary legacy import shim for the attention migration. + +New code imports lightllm.common.attention.backends directly. Removed with the +runtime integration in PR #1721 (5/5), after remaining model callers migrate. +""" + +from lightllm.common.attention.backends.mla.triton.kernels.prefill.context_flashattention_nopad_with_v import ( + context_attention_fwd_with_v, +) diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index c60038af40..bd92c66682 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -13,8 +13,8 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.attention.linear.replay_kda_decode import get_configs, replay +from lightllm.common.attention.backends.kda.replay.decode import get_configs, replay +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, diff --git a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py index f32b093448..e389afce66 100644 --- a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py +++ b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py @@ -1,22 +1,25 @@ -import torch -import time import os -import torch.multiprocessing as mp +import time from typing import List -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage1, + +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( GQADiverseDecodeStage1KernelConfig, + flash_decode_stage1, ) +from lightllm.utils.log_utils import init_logger from lightllm.utils.watchdog_utils import Watchdog logger = init_logger(__name__) def set_seed(): - import torch import random + import numpy as np + import torch seed = 42 torch.manual_seed(seed) @@ -63,7 +66,7 @@ def test_decode_attentions( state.q_head_num = q_head_num state.q_head_dim = q_head_dim state.kv_head_num = kv_head_num - state.softmax_scale = 1 / (q_head_dim ** 0.5) + state.softmax_scale = 1 / (q_head_dim**0.5) state.total_token_num = state.batch_size * test_seq_len infer_state = state @@ -301,9 +304,10 @@ def tuning_configs( if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") - from lightllm.utils.tuning_utils import mp_tuning import collections + from lightllm.utils.tuning_utils import mp_tuning + block_seq = 256 store_json_ans = collections.defaultdict(dict) diff --git a/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py index 13c8945e59..0492884589 100644 --- a/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py +++ b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py @@ -1,21 +1,24 @@ -import torch import os -import torch.multiprocessing as mp from typing import List -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage2, + +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( GQADiverseDecodeStage2KernelConfig, + flash_decode_stage2, ) +from lightllm.utils.log_utils import init_logger from lightllm.utils.watchdog_utils import Watchdog logger = init_logger(__name__) def set_seed(): - import torch import random + import numpy as np + import torch seed = 42 torch.manual_seed(seed) @@ -260,9 +263,10 @@ def tuning_configs( if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") - from lightllm.utils.tuning_utils import mp_tuning import collections + from lightllm.utils.tuning_utils import mp_tuning + block_seq = 256 batch_sizes = [8, 16, 32, 64] seq_lens = [32, 64, 128, 256] diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index de13eb28fc..fff9447874 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -1,16 +1,7 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration - -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - - -from typing import TYPE_CHECKING from math import prod +from typing import TYPE_CHECKING import pytest import torch @@ -18,20 +9,23 @@ from lightllm.common.attention_state.full_att import ( GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, - StaticFP8AttentionConfig, - StaticFP8FullAttentionState, PackedMLAAttentionConfig, - PackedMLALogicalConfig, PackedMLAFullAttentionState, + PackedMLALogicalConfig, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, ) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout, TransferRegion +from unit_tests.common.att_resources.model_helpers import ( + bind_test_fp8_calibration, + full_state_test_buffers, + full_state_test_fields, + read_kv_storage, ) if TYPE_CHECKING: from collections.abc import Sequence + from lightllm.common.attention_state.full_att.base_state import FullAttentionState @@ -335,10 +329,11 @@ def test_deferred_storage_preserves_real_fp8_calibration(kind: str): def test_int4_prefill_with_separate_storage(device, monkeypatch): import weakref from types import SimpleNamespace + + from lightllm.common.attention.backends._shared.prefill import context_flashattention_nopad from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonPrefillAttState from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.kv_copy import ppl_int4kv_copy_kv - from lightllm.common.basemodel.triton_kernel.att.prefill_att import context_flashattention_nopad if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py index 0989fe6d26..812fd6e9f4 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,16 +10,16 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.attention.triton import fp as state_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( mtp_diverse_attn as decode_module, stage1_single_token as stage1_module, stage2_single_token as stage2_module, ) +from lightllm.common.basemodel import cuda_graph as graph_module +from lightllm.common.basemodel.attention.triton import fp as state_module +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py index 48956c4f36..feb4a81c0d 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py @@ -7,10 +7,12 @@ import pytest import torch +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( + stage2_single_token as stage2_module, +) from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import stage2_single_token as stage2_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py index b3068db5fb..ffc322e216 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py @@ -1,7 +1,7 @@ -import weakref import collections import inspect import math +import weakref from types import SimpleNamespace import pytest @@ -10,17 +10,19 @@ from frozendict import frozendict from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( - gqa_flash_decoding as decode_module, +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash import ( gqa_flash_decoding_stage1 as stage1_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( + gqa_flash_decoding as decode_module, +) from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py index df357539e1..5237d0b0f4 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,17 +10,17 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv import ( int4kv_flash_decoding_stage1 as stage1_module, - ppl_int4kv_flash_decoding as decode_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ppl_int4kv_flash_decoding as decode_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py index f3cb8de463..bd9619711f 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py @@ -1,6 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage1 import ( + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( flash_decode_stage1, ) @@ -95,7 +96,7 @@ def test_flash_decode_stage1_execution(setup_tensors): new_k = k.to(q.dtype) new_v = v.to(q.dtype) - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( flash_decode_stage1 as gqa_flash_decode_stage1, ) @@ -174,14 +175,14 @@ def fn_triton(st=setup_tensors): del setup_tensors - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print("SUMMARY - Performance Comparison") - print(f"{'='*80}") + print(f"{'=' * 80}") print(f"{'batch_size':<8} {'seq_len':<12} {'triton_ms':<12}") - print(f"{'-'*80}") + print(f"{'-' * 80}") for r in results: print(f"{r['batch_size']:<8} {r['seq_len']:<12} {r['triton_ms']:<12.3f}") - print(f"{'='*80}") + print(f"{'=' * 80}") if __name__ == "__main__": diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py index 985a33289f..449e6a83b8 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py @@ -1,6 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( flash_decode_stage2, ) @@ -104,7 +105,7 @@ def test_flash_decode_stage2_execution(shared_seq_len): b_seq_len = setup_tensors["b_seq_len"] - setup_tensors["b_shared_seq_len"] req_to_tokens = setup_tensors["Req_to_tokens"][:, setup_tensors["b_shared_seq_len"][0].item() :] - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( flash_decode_stage1 as gqa_flash_decode_stage1, ) @@ -184,9 +185,9 @@ def bench_stage2(): print(results[-1]) del setup_tensors - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print(f"{'batch_size':<10} {'seq_len':<10} {'flash_decode_stage2_ms':<22}") - print(f"{'-'*80}") + print(f"{'-' * 80}") for r in results: print(f"{r['batch_size']:<10} {r['seq_len']:<10} {r['flash_decode_stage2_ms']:<22.4f}") - print(f"{'='*80}") + print(f"{'=' * 80}") diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py index c1a0ca1e58..00f0783abd 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py @@ -1,6 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage3 import ( + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage3 import ( flash_diverse_decode_stage3, ) @@ -26,7 +27,7 @@ def test_flash_diverse_decode_stage3(batch, head_num, seq_len, shared_seq_len, b true_out = torch.zeros_like(out) - from lightllm.common.basemodel.triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding_stage2 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage2 import ( flash_decode_stage2, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py index 7033489f66..487f03b191 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,18 +10,18 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as state_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( int8kv_flash_decoding_stage1 as stage1_module, - int8kv_flash_decoding as decode_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.attention.triton import int8kv as state_module +from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import int8kv_flash_decoding as decode_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py index ecab784529..3c35d64c1a 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py @@ -23,10 +23,13 @@ """ import math + import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo +from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, +) def torch_reference_context_attention_neo( diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py index d1a53f873f..71f77aff71 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py @@ -1,7 +1,9 @@ -import torch import time + import pytest -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( +import torch + +from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( context_attention_fwd_contiguous_kv, ) from lightllm.utils.log_utils import init_logger diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py index 541594306d..5ddcc7c1c5 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py @@ -1,15 +1,17 @@ -import torch import time -import pytest + +import flashinfer import numpy as np +import pytest +import torch import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( + +from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( context_attention_fwd, context_attention_fwd_no_prompt_cache, ) from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py index eaa19ac711..4aafd9d104 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py @@ -24,6 +24,7 @@ """ import math + import pytest import torch @@ -37,7 +38,7 @@ triton_testing = None try: - from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import ( + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( context_attention_fwd_neo, ) except ImportError: @@ -496,7 +497,7 @@ def triton_run(): f"bench: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={str(dtype).split('.')[-1]:<8s} " f"max_q_seq_len={max_q_seq_len:4d} max_prompt_cache_len={max_prompt_cache_len:4d} " f"image_tokens={n_image:4d}/{n_tokens:5d} sum_kv={sum_kv:6d} | " - f"fa3 {fa3_ms*1000:8.1f} us | triton {triton_ms*1000:8.1f} us | " + f"fa3 {fa3_ms * 1000:8.1f} us | triton {triton_ms * 1000:8.1f} us | " f"speedup {speedup:5.2f}x" ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 956bf5e173..fac1a8bb61 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -4,9 +4,9 @@ import triton import triton.language as tl +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor def attach_replay_history( @@ -107,7 +107,7 @@ def replay_layer( """按测试 state 的类型调用对应的 replay decode。""" if state.kda: - from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + from lightllm.common.attention.backends.kda.replay.decode import replay return replay( state_ssm=state.ssm_state[layer_idx], @@ -133,7 +133,7 @@ def replay_layer( state_gate_lower_bound=state.gate_lower_bound, run_config=run_config, ) - from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + from lightllm.common.attention.backends.gdn.replay.decode import replay return replay( state=state, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 7fbb63df5d..811f66f137 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -1,9 +1,9 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, @@ -26,6 +26,7 @@ def ssm_snapshot(backend, req): from types import SimpleNamespace + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState @@ -63,7 +64,7 @@ def _replay_cu(args): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_request_counts_share_compiled_kernel(monkeypatch): - from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix + from lightllm.common.attention.backends._shared.linear.ssm_acceptance import _commit_prefix, commit_ssm_prefix reqs = torch.arange(512, dtype=torch.int32, device="cuda") lengths = reqs % 4 + 1 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py index f7c742db01..1aab2b32a3 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py @@ -3,7 +3,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import causal_conv1d_update def causal_conv1d_ref( diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py index 969a3f354d..a67ef10cf1 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py @@ -1,9 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule if not torch.cuda.is_available(): pytest.skip("CUDA required", allow_module_level=True) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py index 8d92d6cc8d..6338948e3f 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py @@ -9,11 +9,10 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import kda +from lightllm.common.attention.backends.kda.triton.kernels import kda from lightllm.common.kernel_config import KernelConfigs from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel - +from lightllm.common.triton_utils.autotuner import AutotuneLevel, Autotuner KERNELS = [ ("_chunk_kda_scaled_dot_kkt_sub_inter", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter"), @@ -47,7 +46,7 @@ def wrapper_inputs(kernel, tokens=65): R=torch.empty(1, tokens, 2, 64, dtype=q.dtype), h=torch.empty(1, len(chunk_indices), 2, 128, 80, dtype=q.dtype), o=torch.empty_like(v), - scale=128 ** -0.5, + scale=128**-0.5, cu_seqlens=torch.tensor([0, 3, tokens]), chunk_indices=chunk_indices, ) @@ -212,7 +211,7 @@ def test_kda_requires_packed_inputs(fused_gate, invalid_input): def test_split_diagonal_matches_unsplit(monkeypatch, heads, input_dtype, output_dtype, precision): if precision != "ieee" and torch.cuda.get_device_capability()[0] < 8: pytest.skip("TF32 requires SM80+") - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") monkeypatch.setattr(solver, "FLA_TRIL_PRECISION", precision) lengths = (1, 16, 1519) cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) @@ -236,7 +235,7 @@ def observe(*args, **kwargs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("batch,tokens,heads", [(1, 64, 2), (2, 65, 3), (3, 129, 2)]) def test_split_diagonal_matches_unsplit_dense(monkeypatch, batch, tokens, heads): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") rows = torch.arange(tokens, device="cuda") % 64 matrix = torch.randn(batch, tokens, heads, 64, device="cuda") * 0.025 matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) @@ -268,7 +267,7 @@ def observe(*args, **kwargs): ids=["16-requests", "128-requests", "quarter-full-chunks", "16k-tokens", "mostly-empty-chunks", "tiny-batch"], ) def test_split_diagonal_uses_chunk_occupancy(monkeypatch, lengths, expect_split): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) matrix = torch.randn(1, sum(lengths), 2, 64, device="cuda") * 0.025 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py index 2f8544cf75..cecbf9038c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py @@ -7,15 +7,9 @@ import torch import torch.nn.functional as F -from lightllm.common.basemodel.triton_kernel.linear_att import ( - mtp_fused_recurrent as kernel_module, -) +from lightllm.common.attention.backends.gdn.triton.kernels import mtp_fused_recurrent as kernel_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import ( - Autotuner, - AutotuneKernelType, - AutotuneLevel, -) +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner @pytest.fixture(autouse=True) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..bbe8580df8 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -1,10 +1,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py index 8252c445a1..d03908858f 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -1,11 +1,8 @@ import pytest import torch -from lightllm.common.attention_state.bounded_att import GDNState, ConvSSMStateConfig - -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index 786376a362..9e9f237b75 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.attention.linear.replay_gdn_decode import get_configs +from lightllm.common.attention.backends.gdn.replay.decode import get_configs from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, @@ -22,9 +22,9 @@ def ssm_snapshot(backend, req): return clone[:, req] -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index 02797b61b7..fed2439c58 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -1,25 +1,24 @@ +from types import SimpleNamespace + import pytest import torch +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, replay_layer, ) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from types import SimpleNamespace pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import _advance_cursor - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay + from lightllm.common.attention.backends._shared.linear.replayssm_cursor import _advance_cursor + from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay kernels = {"_decode_replay": _decode_replay, "_advance_cursor": _advance_cursor} hashes = {name: set() for name in kernels} @@ -350,12 +349,12 @@ def warm(ssm_backend): @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_cross_layer_fold_with_state_offset_above_int32(kda): - if torch.cuda.mem_get_info()[0] < 6 * 1024 ** 3: + if torch.cuda.mem_get_info()[0] < 6 * 1024**3: pytest.skip("large-stride regression requires 6 GiB of free device memory") torch.manual_seed(791) shape = (48, 3, 12, 128, 128) strides = (50_528_256, 196608, 16384, 128, 1) - assert strides[0] < 2 ** 31 < 47 * strides[0] + assert strides[0] < 2**31 < 47 * strides[0] dense = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.01 sparse = torch.empty_strided(shape, strides, device="cuda", dtype=dense.dtype) sparse.copy_(dense) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py index a94411fbd5..bf3e383953 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.attention.linear import replay_gdn_decode as tuning +from lightllm.common.attention.backends.gdn.replay import decode as tuning from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( attach_replay_history as ReplaySSMExecutor, ) @@ -30,7 +30,7 @@ def test_rebuild_uses_one_layer_scratch(width): @pytest.mark.parametrize("width", [1, 5]) def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, width): from lightllm.common.triton_utils import autotuner as autotuner_module - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index fbfc4a43a9..e7360c94ab 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -1,14 +1,13 @@ import pytest import torch -from lightllm.common.attention_state.bounded_att import KDAState -from unit_tests.common.test_linear_state import make_state -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.test_linear_state import make_state @pytest.fixture diff --git a/unit_tests/models/glm5_next/test_kda_conv_qkv.py b/unit_tests/models/glm5_next/test_kda_conv_qkv.py index b677477519..ed33e15738 100644 --- a/unit_tests/models/glm5_next/test_kda_conv_qkv.py +++ b/unit_tests/models/glm5_next/test_kda_conv_qkv.py @@ -4,15 +4,12 @@ import pytest import torch +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv +from lightllm.common.attention.backends._shared.linear.l2norm import USE_DEFAULT_FLA_NORM, l2norm_fwd +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.common.basemodel.attention.linear import kda as backend_module -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import ( - causal_conv1d_qkv, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd, USE_DEFAULT_FLA_NORM - pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0) or USE_DEFAULT_FLA_NORM, diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 49b2a08eb6..caef5d722d 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -1,29 +1,28 @@ -import weakref import dataclasses import math +import weakref from types import SimpleNamespace import pytest import torch import triton -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.envs_utils import set_env_start_args -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( +from lightllm.common.attention.backends.kda.triton.kernels.kda import ( chunk_kda_with_fused_gate, fused_kda_gate_chunk_cumsum, ) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 from lightllm.models.glm5_next.triton_kernel.kpool import ( compress_pools, - gather_pools, + expand_topk, gather_paged_pools, + gather_pools, get_pool_ranges, - expand_topk, ) -from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 - +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import set_env_start_args pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -241,7 +240,7 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): for group in range(seqs[batch] // 4): values = source[start + group * 4 : start + group * 4 + 4] expected = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128 ** -0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128**-0.5) torch.testing.assert_close(keys[batch * 3 + group].float(), expected_key[0].float(), atol=0, rtol=0) torch.testing.assert_close(scales[batch * 3 + group], expected_scale[0, 0], atol=0, rtol=0) lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) @@ -297,7 +296,7 @@ def run(): if pending == 0: values = source[batch, pos - 3 : pos + 1] pooled = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - key, scale = hadamard_transform_quant_fp8(pooled[None], 128 ** -0.5) + key, scale = hadamard_transform_quant_fp8(pooled[None], 128**-0.5) actual = packed[table[batch, pos].long(), 0] assert torch.equal(actual[:128], key.view(torch.uint8)[0]) assert torch.equal(actual[128:].view(torch.float32), scale.flatten()) @@ -392,7 +391,7 @@ def check(keys, scales): def test_kpool_packed_kv_above_two_gib(): # Automatic KV sizing on H200 can put pool keys beyond a 32-bit byte offset. - loc = 2 ** 31 // (584 * 2) + 1 + loc = 2**31 // (584 * 2) + 1 storage = torch.empty(loc + 1, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] raw = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) @@ -405,14 +404,14 @@ def test_kpool_packed_kv_above_two_gib(): compress_pools(raw, tail, packed, ape, lengths, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_seq_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128**-0.5) assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) assert torch.equal(scales, expected_scale.flatten()) @pytest.mark.parametrize("tp_world_size", [1, 4]) def test_mhc_keeps_streams_through_decode_autotuning(monkeypatch, tp_world_size): - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.model import Glm5NextTpPartModel @@ -482,7 +481,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): state *= gate[:, i, :, None, None].exp() delta = (v[:, i].float() - torch.einsum("bhkv,bhk->bhv", state, k[:, i].float())) * beta[:, i, :, None] state += k[:, i, :, :, None].float() * delta[:, :, None, :] - expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim ** 0.5) + expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim**0.5) out, final = chunk_gated_delta_rule( q, k, @@ -500,7 +499,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): @pytest.mark.parametrize("heads", [16, 64, 128]) def test_nope_attention_native_512_and_cuda_graph(heads): from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState + from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseDecodeState, Glm5NextSparsePrefillState if torch.cuda.get_device_capability()[0] != 9: pytest.skip("FA3 NoPE requires Hopper") @@ -646,11 +645,11 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): pool_count = tokens // 4 pool_values = hidden[: pool_count * 4].view(pool_count, 4, dim) pooled = (pool_values.float() * (pool_values.float() * 0.25 + ape).softmax(1)).sum(1).bfloat16() - k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim ** -0.5) + k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim**-0.5) query = (hidden[-1:] @ q_weight).view(heads, dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim ** -0.5) + q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim**-0.5) logits = (q_fp8.float() @ k_fp8.float().T * k_scale.flatten()).clamp_min(0) - scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) + scores = (logits * q_scale * (heads**-0.5 * dim**-0.5)).sum(0) expected_groups = set(scores.topk(512).indices.tolist()) assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 30c4fc22c4..6eef0d230b 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -1,17 +1,15 @@ import weakref + import pytest import torch +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import ( - fused_recurrent_kda, -) -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -145,17 +143,11 @@ def prefill(values, destination, tail): def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): from types import SimpleNamespace + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState - from lightllm.common.basemodel.attention.linear.replay_kda import ( - ReplayKDADecodeAttState, - ) - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( - ReplayKDAState, - ) - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_update, - ) + from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDADecodeAttState torch.manual_seed(53) heads, dim, width = 2, 128, mtp_step + 1 diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py index c25576a19a..32de9815e0 100644 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -93,7 +93,7 @@ def recurrent(*, initial_state, v, **kwargs): monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) monkeypatch.setattr( - "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + "lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp.causal_conv1d_update", conv_mtp ) monkeypatch.setattr( "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent From 624a50c7f0731032c99e19c51702728e30ec381d Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 11:53:49 +0800 Subject: [PATCH 3/9] refactor(attention): add provider backends and explicit selection --- .../common/attention/backends/_flashinfer.py | 10 + .../attention/backends/_shared/linear/gdn.py | 56 ++++ .../backends/_shared/linear/layout.py | 67 ++++ .../_shared/linear/mtp_state_params.py | 52 ++- .../_shared/linear/replayssm_decode.py | 2 +- .../backends/_shared/linear/replayssm_fold.py | 2 +- .../attention/backends/dense/fa3/backend.py | 217 ++++++++++++ .../attention/backends/dense/fa3/fp8.py | 154 +++++++++ .../attention/backends/dense/fa3/neo.py | 80 +++++ .../attention/backends/dense/fa3/varlen.py | 72 ++++ .../backends/dense/flashinfer/backend.py | 228 +++++++++++++ .../backends/dense/flashinfer/fp8.py | 99 ++++++ .../attention/backends/dense/sdpa/varlen.py | 54 +++ .../backends/dense/triton/backend.py | 192 +++++++++++ .../attention/backends/dense/triton/int4.py | 142 ++++++++ .../attention/backends/dense/triton/int8.py | 182 ++++++++++ .../decode/gqa/flash/gqa_flash_decoding.py | 59 ++++ .../gqa/mtp_diverse/mtp_diverse_attn.py | 9 +- .../int4kv/ppl_int4kv_flash_decoding.py | 58 ++++ .../int8kv/int8kv_flash_decoding_diverse.py | 102 ++++++ .../kernels/decode/int8kv/normal/__init__.py | 1 + .../int8kv/normal/int8kv_flash_decoding.py | 67 ++++ .../decode/mha/flash/flash_decoding.py | 36 ++ .../dense/triton/kernels/prefill/varlen.py | 147 ++++++++ .../attention/backends/dense/triton/neo.py | 48 +++ .../attention/backends/dense/triton/varlen.py | 25 ++ .../common/attention/backends/dense/varlen.py | 42 +++ .../backends/dense/xformers/varlen.py | 39 +++ .../backends/gdn/flashqla/backend.py | 17 + .../attention/backends/gdn/replay/backend.py | 193 +++++++++++ .../attention/backends/gdn/triton/backend.py | 208 ++++++++++++ .../attention/backends/kda/replay/backend.py | 155 +++++++++ .../attention/backends/kda/triton/backend.py | 142 ++++++++ .../attention/backends/mla/fa3/backend.py | 121 +++++++ .../backends/mla/flashinfer/backend.py | 213 ++++++++++++ .../attention/backends/mla/triton/backend.py | 158 +++++++++ .../mla/triton/kernels/decode/__init__.py | 1 + .../kernels/decode/gqa_flash_decoding.py | 147 ++++++++ .../backends/sparse_mla/flashmla/backend.py | 179 ++++++++++ .../backends/sparse_mla/flashmla/fp8.py | 147 ++++++++ lightllm/common/attention/probe.py | 90 +++++ lightllm/common/attention/registry.py | 122 +++++++ lightllm/common/attention/selection.py | 145 ++++++++ lightllm/common/attention/varlen.py | 36 ++ .../basemodel/attention_vit/create_utils.py | 12 +- lightllm/common/basemodel/graph.py | 72 ++++ .../models/deepseek_v4/attention/__init__.py | 0 .../deepseek_v4/attention/sparse_mla.py | 208 ++++++++++++ .../models/glm5_next/attention/__init__.py | 0 .../glm5_next/attention/fp8_sparse_mla.py | 224 +++++++++++++ .../models/glm5_next/attention/sparse_mla.py | 94 ++++++ lightllm/utils/backend_validator.py | 314 +++++++++++++----- unit_tests/common/attention/helpers.py | 3 + unit_tests/common/attention/test_contracts.py | 280 ++++++++++++++++ .../common/attention/test_dense_numerics.py | 79 +++++ .../common/attention/test_linear_layout.py | 72 ++++ .../basemodel/attention/linear/test_gdn.py | 16 +- 57 files changed, 5566 insertions(+), 124 deletions(-) create mode 100644 lightllm/common/attention/backends/_flashinfer.py create mode 100644 lightllm/common/attention/backends/_shared/linear/gdn.py create mode 100644 lightllm/common/attention/backends/_shared/linear/layout.py create mode 100644 lightllm/common/attention/backends/dense/fa3/backend.py create mode 100644 lightllm/common/attention/backends/dense/fa3/fp8.py create mode 100644 lightllm/common/attention/backends/dense/fa3/neo.py create mode 100644 lightllm/common/attention/backends/dense/fa3/varlen.py create mode 100644 lightllm/common/attention/backends/dense/flashinfer/backend.py create mode 100644 lightllm/common/attention/backends/dense/flashinfer/fp8.py create mode 100644 lightllm/common/attention/backends/dense/sdpa/varlen.py create mode 100644 lightllm/common/attention/backends/dense/triton/backend.py create mode 100644 lightllm/common/attention/backends/dense/triton/int4.py create mode 100644 lightllm/common/attention/backends/dense/triton/int8.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py create mode 100644 lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py create mode 100644 lightllm/common/attention/backends/dense/triton/neo.py create mode 100644 lightllm/common/attention/backends/dense/triton/varlen.py create mode 100644 lightllm/common/attention/backends/dense/varlen.py create mode 100644 lightllm/common/attention/backends/dense/xformers/varlen.py create mode 100644 lightllm/common/attention/backends/gdn/flashqla/backend.py create mode 100644 lightllm/common/attention/backends/gdn/replay/backend.py create mode 100644 lightllm/common/attention/backends/gdn/triton/backend.py create mode 100644 lightllm/common/attention/backends/kda/replay/backend.py create mode 100644 lightllm/common/attention/backends/kda/triton/backend.py create mode 100644 lightllm/common/attention/backends/mla/fa3/backend.py create mode 100644 lightllm/common/attention/backends/mla/flashinfer/backend.py create mode 100644 lightllm/common/attention/backends/mla/triton/backend.py create mode 100644 lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py create mode 100644 lightllm/common/attention/backends/sparse_mla/flashmla/backend.py create mode 100644 lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py create mode 100644 lightllm/common/attention/probe.py create mode 100644 lightllm/common/attention/registry.py create mode 100644 lightllm/common/attention/selection.py create mode 100644 lightllm/common/attention/varlen.py create mode 100644 lightllm/common/basemodel/graph.py create mode 100644 lightllm/models/deepseek_v4/attention/__init__.py create mode 100644 lightllm/models/deepseek_v4/attention/sparse_mla.py create mode 100644 lightllm/models/glm5_next/attention/__init__.py create mode 100644 lightllm/models/glm5_next/attention/fp8_sparse_mla.py create mode 100644 lightllm/models/glm5_next/attention/sparse_mla.py create mode 100644 unit_tests/common/attention/test_dense_numerics.py create mode 100644 unit_tests/common/attention/test_linear_layout.py diff --git a/lightllm/common/attention/backends/_flashinfer.py b/lightllm/common/attention/backends/_flashinfer.py new file mode 100644 index 0000000000..cae6a29ca7 --- /dev/null +++ b/lightllm/common/attention/backends/_flashinfer.py @@ -0,0 +1,10 @@ +import os + +import torch + + +def set_flashinfer_envs(): + capability = torch.cuda.get_device_capability() + arch = f"{capability[0]}.{capability[1]}" + os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}" + return diff --git a/lightllm/common/attention/backends/_shared/linear/gdn.py b/lightllm/common/attention/backends/_shared/linear/gdn.py new file mode 100644 index 0000000000..6d673d52a0 --- /dev/null +++ b/lightllm/common/attention/backends/_shared/linear/gdn.py @@ -0,0 +1,56 @@ +"""GDN tensor-parallel dimensions and input views shared by native and replay kernels.""" + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + + +class GDNMetadata: + def _init_linear_layer_metadata(self, network_config, tp_world_size): + self.mtp_step = get_env_start_args().mtp_step + self.num_v_heads = network_config["linear_num_value_heads"] + self.num_k_heads = network_config["linear_num_key_heads"] + self.head_k_dim = network_config["linear_key_head_dim"] + self.head_v_dim = network_config["linear_value_head_dim"] + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] + self.activation = network_config["hidden_act"] + self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size + self.tp_ba_dim = self.num_v_heads * 2 // tp_world_size + self.tp_num_k_heads = self.num_k_heads // tp_world_size + self.tp_num_v_heads = self.num_v_heads // tp_world_size + self.tp_key_dim = self.key_dim // tp_world_size + self.tp_value_dim = self.value_dim // tp_world_size + if self.num_v_heads % self.num_k_heads: + raise ValueError("num_v_heads must be divisible by num_k_heads") + self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads + ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} + self.ssm_state_dtype = ssm_dtype_dict.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) + return + + def _split_qkvzba(self, mixed_qkvzba): + qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim + z_end = qkv_dim + self.tp_value_dim + b_end = z_end + self.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return (mixed_qkv, z, b, a) + + def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): + if decode: + query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1) + batch_size = mixed_qkv.shape[0] + query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim) + return (query, key, value) + else: + query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1) + seq_len = query.shape[0] + query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) + return (query, key, value) diff --git a/lightllm/common/attention/backends/_shared/linear/layout.py b/lightllm/common/attention/backends/_shared/linear/layout.py new file mode 100644 index 0000000000..1ada562ea0 --- /dev/null +++ b/lightllm/common/attention/backends/_shared/linear/layout.py @@ -0,0 +1,67 @@ +from dataclasses import dataclass, field + +import torch + +from lightllm.common.attention.backends.base import AttentionExecutionContext + +from .mtp_state_params import build_dynamic_mtp_layout + + +@dataclass(slots=True) +class LinearDecodeContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor | None = None + b_ssm_buffer_idx: torch.Tensor | None = None + b1_mtp_cu_q_seq_len: torch.Tensor | None = None + b_cu_seqlens: torch.Tensor | None = None + accepted_by_state: dict[torch.Tensor, torch.Tensor] = field(default_factory=dict) + + def accepted_tokens(self, state_index: torch.Tensor) -> torch.Tensor: + """Gather once per State per update; Graph captures the gather for every replay. + + The cache retains the source tensors until the next update, keeping identity + keys valid for this context. HOLD sequences must select the initial convolution state. + """ + key = state_index + if key not in self.accepted_by_state: + self.accepted_by_state[key] = torch.where( + self.b_conv_buffer_idx == self.batch.hold_req_id, + 1, + state_index[self.b_conv_buffer_idx] + 1, + ) + return self.accepted_by_state[key] + + +def update_linear_decode(context: LinearDecodeContext, draft_step: int, dynamic: bool, *, replay: bool) -> None: + batch = context.batch + context.accepted_by_state.clear() + if draft_step == 0: + context.b_conv_buffer_idx = batch.b_req_idx + context.b_ssm_buffer_idx = batch.b_req_idx + context.b1_mtp_cu_q_seq_len = None + if replay: + context.b_cu_seqlens = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) + return + width = draft_step + 1 + if dynamic: + context.b1_mtp_cu_q_seq_len, context.b_conv_buffer_idx = build_dynamic_mtp_layout( + b_req_idx=batch.b_req_idx, + b_mtp_index=batch.b_mtp_index, + hold_req_id=batch.hold_req_id, + ) + else: + if batch.batch_size % width: + raise ValueError("linear decode batch size must be divisible by draft_step + 1") + context.b1_mtp_cu_q_seq_len = torch.arange( + 0, + batch.batch_size + 1, + width, + dtype=torch.int32, + device=batch.device, + ) + context.b_conv_buffer_idx = batch.b_req_idx.view(-1, width)[:, 0].contiguous() + if replay: + context.b_ssm_buffer_idx = context.b_conv_buffer_idx + context.b_cu_seqlens = context.b1_mtp_cu_q_seq_len + else: + offsets = torch.arange(width, dtype=batch.b_req_idx.dtype, device=batch.device) + context.b_ssm_buffer_idx = context.b_conv_buffer_idx[:, None] * width + offsets diff --git a/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py index b81bf9ec16..993ef6f1b2 100644 --- a/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py +++ b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py @@ -14,6 +14,7 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( batch_size, hold_req_id, BLOCK_SIZE: tl.constexpr, + SELECT_ACCEPTED: tl.constexpr = True, ): offsets = tl.arange(0, BLOCK_SIZE) token_mask = offsets < batch_size @@ -30,21 +31,18 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( # replay arbitrary compact per-request widths at the same token batch size. tl.store(out_cu_q_seq_len + offsets, actual_token_num, mask=cu_mask) tl.store(out_conv_buffer_idx + offsets, hold_req_id, mask=token_mask) - tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask) + if SELECT_ACCEPTED: + tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask) tl.debug_barrier() # mtp_index restarts at zero on the first row of every request group. is_start = valid_row & (mtp_index == 0) sequence_index = tl.cumsum(tl.where(is_start, 1, 0), axis=0) - 1 - accepted_state_index = tl.load( - req_to_mtp_state_index + req_idx, - mask=is_start, - other=0, - ) - tl.store(out_cu_q_seq_len + sequence_index, offsets, mask=is_start) tl.store(out_conv_buffer_idx + sequence_index, req_idx, mask=is_start) - tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) + if SELECT_ACCEPTED: + accepted_state_index = tl.load(req_to_mtp_state_index + req_idx, mask=is_start, other=0) + tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) def build_dynamic_mtp_linear_att_state_params( @@ -85,11 +83,10 @@ def build_dynamic_mtp_linear_att_state_params( per-request query lengths. """ - assert b_req_idx.is_cuda and b_mtp_index.is_cuda and req_to_mtp_state_index.is_cuda - assert b_req_idx.ndim == 1 and b_req_idx.shape == b_mtp_index.shape - assert b_req_idx.dtype == torch.int32 and b_mtp_index.dtype == torch.int32 + _validate_layout_inputs(b_req_idx, b_mtp_index) + if req_to_mtp_state_index.device != b_req_idx.device or req_to_mtp_state_index.ndim != 1: + raise ValueError("MTP state indices must be a vector on the request device") batch_size = b_req_idx.shape[0] - assert batch_size > 0 b1_cu_q_seq_len = torch.empty((batch_size + 1,), dtype=torch.int32, device=b_req_idx.device) b_conv_buffer_idx = torch.empty_like(b_req_idx) @@ -109,3 +106,34 @@ def build_dynamic_mtp_linear_att_state_params( num_stages=1, ) return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens + + +def build_dynamic_mtp_layout(b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, hold_req_id: int): + _validate_layout_inputs(b_req_idx, b_mtp_index) + batch_size = b_req_idx.numel() + cu = torch.empty(batch_size + 1, dtype=torch.int32, device=b_req_idx.device) + requests = torch.empty_like(b_req_idx) + _build_dynamic_mtp_linear_att_state_params_kernel[(1,)]( + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + req_to_mtp_state_index=None, + out_cu_q_seq_len=cu, + out_conv_buffer_idx=requests, + out_num_accepted_tokens=None, + batch_size=batch_size, + hold_req_id=hold_req_id, + BLOCK_SIZE=triton.next_power_of_2(batch_size + 1), + SELECT_ACCEPTED=False, + num_warps=8, + num_stages=1, + ) + return cu, requests + + +def _validate_layout_inputs(requests, positions): + if not requests.is_cuda or positions.device != requests.device: + raise ValueError("MTP layout requires inputs on the same CUDA device") + if requests.ndim != 1 or requests.shape != positions.shape or requests.numel() == 0: + raise ValueError("MTP layout requires matching nonempty vectors") + if requests.dtype != torch.int32 or positions.dtype != torch.int32: + raise ValueError("MTP layout requires int32 request and position indices") diff --git a/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py index b7f138c844..450c695375 100644 --- a/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix +from .replayssm_delta import _apply_accepted_prefix def check_replay_sync_mark( diff --git a/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py index 2e491a038b..42aad6edc2 100644 --- a/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py +++ b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from lightllm.common.attention.backends._shared.linear.replayssm_delta import _apply_accepted_prefix +from .replayssm_delta import _apply_accepted_prefix @triton.jit diff --git a/lightllm/common/attention/backends/dense/fa3/backend.py b/lightllm/common/attention/backends/dense/fa3/backend.py new file mode 100644 index 0000000000..844a63da1f --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/backend.py @@ -0,0 +1,217 @@ +import dataclasses + +import torch +import triton + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params, page_table_copy +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor +from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers +from lightllm.utils.envs_utils import enable_env_vars +from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune + + +@dataclasses.dataclass(slots=True) +class Fa3PrefillContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + causal: bool = None + + +@dataclasses.dataclass(slots=True) +class Fa3DecodeContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + b_att_seq_len: torch.Tensor = None + decode_max_q_seq_len: int = None + decode_max_kv_seq_len: int = None + causal: bool = None + + +class Fa3AttBackend(BaseAttentionBackend[Fa3PrefillContext | Fa3DecodeContext, DenseInputs, torch.Tensor]): + requires_cpu_cache_sync = enable_env_vars("LIGHTLLM_FA3_CPU_CACHE_SYNC") + supports_noncausal_prefill = True + provider = "fa3" + dependency = "sgl_kernel.flash_attn" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supports_attention_sink = True + supports_sliding_window = True + supports_noncausal_decode = True + max_head_size = 256 + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fa3PrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fa3DecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.causal = self.uses_causal_attention() + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + context.page_table = torch.empty( + (context.batch.batch_size, triton.cdiv(context.batch.max_kv_seq_len, self.infer_page_size)), + dtype=torch.int32, + device=context.batch.device, + ) + page_table_copy( + page_table=context.page_table, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_req_idx=context.batch.b_req_idx, + page_size=self.infer_page_size, + ) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._prefill_paged(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_paged( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ) -> torch.Tensor: + if layer.spec.sliding_window is not None: + window_size = layer.spec.sliding_window + else: + window_size = (-1, -1) + if layer.params.sink is not None: + sink_weight: torch.Tensor = layer.params.sink + else: + sink_weight = None + k_descale, v_descale = (None, None) + Lq = q.shape[-1] + sm_scale = 1.0 / Lq**0.5 + o = flash_attn_with_kvcache( + q=q, + k_cache=k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=sm_scale, + causal=context.causal, + window_size=window_size, + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + sinks=sink_weight, + ) + return o + + def _update_decode(self, context): + context.causal = self.uses_causal_attention() + draft_step = self.decode_draft_step + if self.uses_dynamic_spec_verify_layout(): + b_att_req_idx = self._decode_init_dynamic_spec_verify_state(context, draft_step) + elif draft_step > 0: + b_att_req_idx = self._decode_init_fixed_spec_decode_state(context, draft_step) + else: + b_att_req_idx = self._decode_init_normal_decode_state(context) + self._decode_init_page_table(context, b_att_req_idx) + + def _decode_init_dynamic_spec_verify_state(self, context, draft_step: int) -> torch.Tensor: + b_mark_mtp_shared_group = build_mtp_shared_group_markers(context.batch.b_req_idx, hold_req_id=self.hold_req_id) + b_q_seq_len, b_kv_seq_len, b_att_req_idx, context.b_att_seq_len = build_dynamic_spec_fa3_decode_params( + b_req_idx=context.batch.b_req_idx, + b_seq_len=context.batch.b_seq_len, + b_mark_mtp_shared_group=b_mark_mtp_shared_group, + att_batch_size=context.batch.batch_size, + hold_req_id=self.hold_req_id, + ) + self._decode_init_spec_decode_cu_seqlens(context, b_q_seq_len, b_kv_seq_len) + context.decode_max_q_seq_len = draft_step + 1 + return b_att_req_idx + + def _decode_init_fixed_spec_decode_state(self, context, draft_step: int) -> torch.Tensor: + mtp_size = draft_step + 1 + assert context.batch.batch_size % mtp_size == 0, ( + "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, " + f"got batch_size={context.batch.batch_size}, draft_step={draft_step}." + ) + b_q_seq_len = torch.full( + (context.batch.b_seq_len.shape[0] // mtp_size,), + fill_value=mtp_size, + dtype=torch.int32, + device=context.batch.b_seq_len.device, + ) + b_kv_seq_len = context.batch.b_seq_len[draft_step::mtp_size] + b_att_req_idx = context.batch.b_req_idx[draft_step::mtp_size] + context.b_att_seq_len = b_kv_seq_len.contiguous() + self._decode_init_spec_decode_cu_seqlens(context, b_q_seq_len, b_kv_seq_len) + context.decode_max_q_seq_len = mtp_size + return b_att_req_idx + + def _decode_init_normal_decode_state(self, context) -> torch.Tensor: + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + context.b_att_seq_len = context.batch.b_seq_len + context.decode_max_q_seq_len = 1 + return context.batch.b_req_idx + + def _decode_init_spec_decode_cu_seqlens(self, context, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor): + b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len) + context.cu_seqlens_q = b1_cu_q_seq_len.int() + context.cu_seqlens_k = b1_cu_kv_seq_len.int() + + def _decode_init_page_table(self, context, b_att_req_idx): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + width = triton.cdiv(context.capacity.max_kv_seq_len, self.infer_page_size) + context.page_table = context.workspace.get("fa3_page_table", b_att_req_idx.numel() * width, torch.int32)[ + : b_att_req_idx.numel() * width + ].view(b_att_req_idx.numel(), width) + page_table_copy( + page_table=context.page_table, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_req_idx=b_att_req_idx, + page_size=self.infer_page_size, + ) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._run_decode(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _run_decode(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + if layer.spec.sliding_window is not None: + window_size = layer.spec.sliding_window + else: + window_size = (-1, -1) + if layer.params.sink is not None: + sink_weight: torch.Tensor = layer.params.sink + else: + sink_weight = None + k_descale, v_descale = (None, None) + Lq = q.shape[-1] + sm_scale = 1.0 / Lq**0.5 + o = flash_attn_with_kvcache_autotune( + q=q, + k_cache=k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=None, + max_seqlen_q=context.decode_max_q_seq_len, + max_seqlen_k=context.decode_max_kv_seq_len, + softmax_scale=sm_scale, + causal=context.causal, + window_size=window_size, + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + sinks=sink_weight, + ) + return o diff --git a/lightllm/common/attention/backends/dense/fa3/fp8.py b/lightllm/common/attention/backends/dense/fa3/fp8.py new file mode 100644 index 0000000000..836b17555d --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/fp8.py @@ -0,0 +1,154 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend, Fa3DecodeContext, Fa3PrefillContext +from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import ( + q_per_head_fp8_quant, + q_per_head_static_fp8_quant, +) +from lightllm.utils.log_utils import init_logger +from lightllm.utils.sgl_utils import flash_attn_with_kvcache + +logger = init_logger(__name__) + + +@dataclasses.dataclass(slots=True) +class Fp8Fa3PrefillContext(Fa3PrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Fp8Fa3DecodeContext(Fa3DecodeContext): + pass + + +class Fp8Fa3AttBackend(Fa3AttBackend): + provider = "fa3" + supports_sliding_window = False + cache_formats = frozenset({"fp8kv_sph"}) + + @property + def infer_page_size(self) -> int: + return 1 + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fp8Fa3PrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fp8Fa3DecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + super()._update_prefill(context) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._run_prefill( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ) -> torch.Tensor: + q_head_num = q.shape[1] + q_head_dim = q.shape[2] + k_head_num = k.shape[1] + k_head_dim = k.shape[2] + cache_k = k.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + cache_v = v.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + batch_size = context.batch.b_seq_len.shape[0] + q = q.reshape(q.shape[0], k_head_num, -1) + if layer.params.query_scale is None: + q, q_scale = q_per_head_fp8_quant( + q=q, seq_lens=context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1], b1_start_loc=context.cu_seqlens_q + ) + else: + static_q_scales = layer.params.query_scale + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(batch_size, k_head_num) + o = flash_attn_with_kvcache( + q=q.reshape(-1, q_head_num, q_head_dim), + k_cache=cache_k, + v_cache=cache_v, + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + q_descale=q_scale, + k_descale=layer.params.key_scale.view(1, -1).expand(batch_size, -1), + v_descale=layer.params.value_scale.view(1, -1).expand(batch_size, -1), + return_softmax_lse=False, + ) + return o + + def _update_decode(self, context): + super()._update_decode(context) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._run_decode( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _run_decode( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ): + k_head_num = k.shape[1] + k_head_dim = k.shape[2] + cache_k = k.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + cache_v = v.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + q_head_num = q.shape[1] + att_batch_size = context.b_att_seq_len.shape[0] + q = q.reshape(q.shape[0], k_head_num, -1) + if layer.params.query_scale is None: + q, q_scale = q_per_head_fp8_quant( + q=q, seq_lens=context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1], b1_start_loc=context.cu_seqlens_q + ) + else: + static_q_scales = layer.params.query_scale + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(att_batch_size, k_head_num) + o = flash_attn_with_kvcache( + q=q.reshape(-1, q_head_num, k_head_dim), + k_cache=cache_k, + v_cache=cache_v, + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.decode_max_q_seq_len, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + q_descale=q_scale.view(att_batch_size, k_head_num), + k_descale=layer.params.key_scale.view(1, -1).expand(att_batch_size, -1), + v_descale=layer.params.value_scale.view(1, -1).expand(att_batch_size, -1), + return_softmax_lse=False, + ) + return o diff --git a/lightllm/common/attention/backends/dense/fa3/neo.py b/lightllm/common/attention/backends/dense/fa3/neo.py new file mode 100644 index 0000000000..cebe3a015d --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/neo.py @@ -0,0 +1,80 @@ +import dataclasses +import inspect +import warnings + +import torch + +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend, Fa3PrefillContext + +_NEO_FA3_INSTALL_HINT = ( + "Install the Neo FA3 build from source in the current Python environment:\n" + " git clone --recursive https://github.com/WANDY666/flash-attention.git\n" + " cd flash-attention/hopper\n" + " FLASH_ATTENTION_FORCE_BUILD=TRUE python -m pip install --no-build-isolation --no-deps --force-reinstall .\n" + "This replaces a standalone flash_attn_3 / flash_attn_interface installation; LightLLM's regular" + " FA3 uses the separate sgl_kernel package." +) +try: + from flash_attn_interface import flash_attn_with_kvcache as flash_attn_with_kvcache_neo + + _sig = inspect.signature(flash_attn_with_kvcache_neo) + if "image_token_end" not in _sig.parameters: + raise ImportError("flash_attn_interface is missing image_token_end support (need the Neo build)") + HAS_FLASH_ATTN_INTERFACE = True +except ImportError as exc: + warnings.warn( + f"Neo FA3 is unavailable: {exc}. Automatic Neo prefill selection will fall back to Triton. " + "To select Triton explicitly, use --llm_prefill_att_backend triton.\n" + _NEO_FA3_INSTALL_HINT + ) + flash_attn_with_kvcache_neo = None + HAS_FLASH_ATTN_INTERFACE = False + + +@dataclasses.dataclass(slots=True) +class NeoFa3PrefillContext(Fa3PrefillContext): + pass + + +class NeoFa3AttBackend(Fa3AttBackend): + @classmethod + def check_support(cls, spec, caps=None, execution=None) -> None: + super().check_support(spec, caps, execution) + if spec.phase.value == "prefill" and not HAS_FLASH_ATTN_INTERFACE: + raise UnsupportedAttention("Neo FA3 requires the image-token flash attention interface") + + provider = "fa3" + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NeoFa3PrefillContext(capacity, capture, workspace) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + q = inputs.query + k = layer.params.key + v = layer.params.value + if not HAS_FLASH_ATTN_INTERFACE: + raise ImportError("Neo prefill requires FA3 with image_token_end support.\n" + _NEO_FA3_INSTALL_HINT) + if context.batch.b_image_token_end is None: + raise ValueError("Neo prefill requires b_image_token_end to describe image attention spans") + o = flash_attn_with_kvcache_neo( + q=q, + k_cache=k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=1.0 / q.shape[-1] ** 0.5, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + k_descale=None, + v_descale=None, + return_softmax_lse=False, + image_token_end=context.batch.b_image_token_end, + ) + return o diff --git a/lightllm/common/attention/backends/dense/fa3/varlen.py b/lightllm/common/attention/backends/dense/fa3/varlen.py new file mode 100644 index 0000000000..66bd1ce428 --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/varlen.py @@ -0,0 +1,72 @@ +import torch + +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class Fa3VarlenAttentionOp(BaseVarlenAttentionOp): + provider = "fa3" + dependency = "sgl_kernel.flash_attn" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + max_head_size = 256 + head_alignment = 8 + supports_graph = True + + def __init__(self): + self._op = None + + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + tokens, _, head_dim = q.shape + if tokens == 0: + return out + if self._op is None: + import sgl_kernel.flash_attn + + self._op = torch.ops.sgl_kernel.fwd.default + softmax_scale = head_dim**-0.5 + window_size = (-1, -1) + self._op( + q=q, + k=k, + v=v, + k_new=None, + v_new=None, + q_v=None, + out=out, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + cu_seqlens_k_new=None, + seqused_q=None, + seqused_k=None, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + page_table=None, + kv_batch_idx=None, + leftpad_k=None, + rotary_cos=None, + rotary_sin=None, + seqlens_rotary=None, + q_descale=None, + k_descale=None, + v_descale=None, + softmax_scale=softmax_scale, + is_causal=False, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=0, + softcap=0.0, + is_rotary_interleaved=False, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0, + sinks=None, + sparse_mask_fine=None, + only_qv=False, + ) + + return out diff --git a/lightllm/common/attention/backends/dense/flashinfer/backend.py b/lightllm/common/attention/backends/dense/flashinfer/backend.py new file mode 100644 index 0000000000..c78f65e559 --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/backend.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch +import triton + +from lightllm.common.attention.api import BackendResources +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.execution import WorkspaceRequirement +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor +from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index +from lightllm.utils.dist_utils import get_dp_world_size + +if TYPE_CHECKING: + from flashinfer.decode import BatchDecodeWithPagedKVCacheWrapper + from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper + + +@dataclasses.dataclass(slots=True) +class FlashInferPrefillContext(AttentionExecutionContext): + prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper | None = None + + +@dataclasses.dataclass(slots=True) +class FlashInferDecodeContext(AttentionExecutionContext): + kv_last_page_len_buffer: torch.Tensor = None + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + decode_wrapper: BatchDecodeWithPagedKVCacheWrapper | None = None + + +class FlashInferAttBackend( + BaseAttentionBackend[FlashInferPrefillContext | FlashInferDecodeContext, DenseInputs, torch.Tensor] +): + supports_noncausal_prefill = True + provider = "flashinfer" + dependency = "flashinfer" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supports_expert_parallel = False + cache_formats = frozenset({"None"}) + workspace_buffer_key = "flashinfer_fp" + workspace_buffer_size = 512 * 1024 * 1024 + + def __init__(self, model, spec): + set_flashinfer_envs() + super().__init__(model=model, spec=spec) + tp_world_size = get_dp_world_size() + self.tp_q_head_num = model.config["num_attention_heads"] // tp_world_size + self.tp_kv_head_num = max(model.config["num_key_value_heads"] // tp_world_size, 1) + head_dim = model.config["hidden_size"] // model.config["num_attention_heads"] + self.head_dim = model.config.get("head_dim", head_dim) + self.max_seq_length = model.max_seq_length + self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) + self.q_data_type = spec.dtype + self.kv_data_type = spec.cache.dtype + + def configure_spec(self, spec): + self.tp_q_head_num = spec.num_q_heads + self.tp_kv_head_num = spec.cache.key_heads + self.head_dim = spec.head_size + + def resource_requirements(self, capacity, *, capture): + return BackendResources( + (WorkspaceRequirement(self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),) + ) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return FlashInferPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return FlashInferDecodeContext(capacity, capture, workspace) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._prefill_paged(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_paged( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ) -> torch.Tensor: + o_tensor = alloc_func(q.shape, q.dtype, device=q.device) + context.prefill_wrapper.run( + q, + ( + k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + ), + out=o_tensor, + ) + return o_tensor + + def _refresh_decode_plan(self, context, max_kv_len: int): + from flashinfer.decode import fast_decode_plan + + uniform_kv_indptr_cpu = torch.arange( + context.batch.batch_size + 1, dtype=torch.int32, device="cpu" + ) * triton.cdiv(max_kv_len, self.infer_page_size) + fast_decode_plan( + context.decode_wrapper, + indptr=context.kv_starts, + indices=context.kv_indices, + last_page_len=context.kv_last_page_len_buffer, + num_qo_heads=self.tp_q_head_num, + num_kv_heads=self.tp_kv_head_num, + head_dim=self.head_dim, + page_size=self.infer_page_size, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + non_blocking=True, + global_override_indptr_cpu=uniform_kv_indptr_cpu, + ) + return + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._run_decode(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _run_decode(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + o_tensor = alloc_func(q.shape, q.dtype) + context.decode_wrapper.run( + q, + ( + k.view(-1, self.infer_page_size, k.shape[1], k.shape[2]), + v.view(-1, self.infer_page_size, v.shape[1], v.shape[2]), + ), + out=o_tensor, + ) + return o_tensor + + def _prepare_prefill(self, context): + import flashinfer + + batch = context.batch + q_starts = batch.b1_cu_q_seq_len.int() + b_page_len = (batch.b_seq_len + self.infer_page_size - 1) // self.infer_page_size + kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + last_page_len = batch.b_seq_len - (b_page_len - 1) * self.infer_page_size + indices = context.workspace.get("flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=kv_starts[:-1], + max_token_len=batch.max_kv_seq_len, + out_page_indices=indices, + page_size=self.infer_page_size, + ) + context.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), + qo_indptr_buf=q_starts, + paged_kv_indptr_buf=kv_starts, + paged_kv_indices_buf=indices, + paged_kv_last_page_len_buf=last_page_len, + ) + context.prefill_wrapper.plan( + q_starts, + kv_starts, + indices, + last_page_len, + self.tp_q_head_num, + self.tp_kv_head_num, + self.head_dim, + self.infer_page_size, + causal=self.causal, + pos_encoding_mode="NONE", + logits_soft_cap=0.0, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + ) + + def _prepare_decode(self, context): + import flashinfer + + batch = context.batch + b_page_len = (batch.b_seq_len + self.infer_page_size - 1) // self.infer_page_size + starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + last_page_len = batch.b_seq_len - (b_page_len - 1) * self.infer_page_size + if context.kv_starts is None: + context.kv_starts, context.kv_last_page_len_buffer = starts, last_page_len + context.kv_indices = context.workspace.get( + "flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32 + ) + else: + context.kv_starts.copy_(starts, non_blocking=True) + context.kv_last_page_len_buffer.copy_(last_page_len, non_blocking=True) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=context.kv_starts[:-1], + max_token_len=context.current_max_kv_seq_len, + out_page_indices=context.kv_indices, + page_size=self.infer_page_size, + ) + if context.decode_wrapper is None: + context.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), + "NHD", + use_cuda_graph=context.capture, + use_tensor_cores=True, + paged_kv_indptr_buffer=context.kv_starts, + paged_kv_indices_buffer=context.kv_indices, + paged_kv_last_page_len_buffer=context.kv_last_page_len_buffer, + ) + context.decode_wrapper.plan( + context.kv_starts, + context.kv_indices, + context.kv_last_page_len_buffer, + self.tp_q_head_num, + self.tp_kv_head_num, + self.head_dim, + self.infer_page_size, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + non_blocking=True, + ) + else: + self._refresh_decode_plan(context, context.current_max_kv_seq_len) diff --git a/lightllm/common/attention/backends/dense/flashinfer/fp8.py b/lightllm/common/attention/backends/dense/flashinfer/fp8.py new file mode 100644 index 0000000000..61d563f5a7 --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/fp8.py @@ -0,0 +1,99 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.backends.dense.flashinfer.backend import ( + FlashInferAttBackend, + FlashInferDecodeContext, + FlashInferPrefillContext, +) +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +@dataclasses.dataclass(slots=True) +class Fp8FlashInferPrefillContext(FlashInferPrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Fp8FlashInferDecodeContext(FlashInferDecodeContext): + pass + + +class Fp8FlashInferAttBackend(FlashInferAttBackend): + provider = "flashinfer" + cache_formats = frozenset({"fp8kv_spt"}) + + def __init__(self, model, spec): + set_flashinfer_envs() + super().__init__(model=model, spec=spec) + self.kv_data_type = torch.float8_e4m3fn + + @property + def infer_page_size(self) -> int: + return 1 + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fp8FlashInferPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fp8FlashInferDecodeContext(capacity, capture, workspace) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._run_prefill( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ) -> torch.Tensor: + o_tensor = alloc_func(q.shape, q.dtype, device=q.device) + k = k.unsqueeze(1).view(torch.float8_e4m3fn) + v = v.unsqueeze(1).view(torch.float8_e4m3fn) + k_descale = layer.params.key_scale + v_descale = layer.params.value_scale + context.prefill_wrapper.run(q=q, paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor) + return o_tensor + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + return self._run_decode( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _run_decode( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, + ): + o_tensor = alloc_func(q.shape, q.dtype, device=q.device) + k = k.unsqueeze(1).view(torch.float8_e4m3fn) + v = v.unsqueeze(1).view(torch.float8_e4m3fn) + k_descale = layer.params.key_scale + v_descale = layer.params.value_scale + context.decode_wrapper.run(q=q, paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/sdpa/varlen.py b/lightllm/common/attention/backends/dense/sdpa/varlen.py new file mode 100644 index 0000000000..0b686ebda7 --- /dev/null +++ b/lightllm/common/attention/backends/dense/sdpa/varlen.py @@ -0,0 +1,54 @@ +import torch +import torch.nn.functional as F + +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp +from lightllm.platform import Platform + + +class SdpaVarlenAttentionOp(BaseVarlenAttentionOp): + """Fallback with a device-to-host boundary sync and a Python loop per sequence. + + This path cannot be captured in a CUDA Graph. + """ + + provider = "sdpa" + supported_platforms = frozenset({Platform.CPU, Platform.CUDA}) + + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + output_buffer = out + assert q.ndim == k.ndim == v.ndim == output_buffer.ndim == 3 + assert cu_seqlens is not None and cu_seqlens.ndim == 1 + cu_seqlens = cu_seqlens.detach().to("cpu") + B = cu_seqlens.numel() - 1 + + with torch.no_grad(): + for b in range(B): + s = int(cu_seqlens[b]) + e = int(cu_seqlens[b + 1]) + L = e - s + if L <= 0: + continue + if max_seqlen: + assert L <= max_seqlen + + # [L, H, D] -> [1, H, L, D] + q_ = q[s:e].permute(1, 0, 2).unsqueeze(0) + k_ = k[s:e].permute(1, 0, 2).unsqueeze(0) + v_ = v[s:e].permute(1, 0, 2).unsqueeze(0) + + sequence_output = F.scaled_dot_product_attention( + q_, + k_, + v_, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + # [1, H, L, D] -> [L, H, D] + output_buffer[s:e].copy_(sequence_output.squeeze(0).permute(1, 0, 2)) + + return output_buffer diff --git a/lightllm/common/attention/backends/dense/triton/backend.py b/lightllm/common/attention/backends/dense/triton/backend.py new file mode 100644 index 0000000000..a9b3664835 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/backend.py @@ -0,0 +1,192 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers + + +@dataclasses.dataclass(slots=True) +class TritonPrefillContext(AttentionExecutionContext): + pass + + +@dataclasses.dataclass(slots=True) +class TritonDecodeContext(AttentionExecutionContext): + b_mark_mtp_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None + + +class TritonAttBackend(BaseAttentionBackend[TritonPrefillContext | TritonDecodeContext, DenseInputs, torch.Tensor]): + provider = "triton" + supports_alibi = True + supports_sliding_window = True + supports_speculative_window = False + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return TritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return TritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + pass + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + if layer.params.alibi is not None: + return self._prefill_alibi(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + else: + return self._prefill_paged(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_alibi(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + out = alloc_func(q.shape, q.dtype) + from lightllm.common.basemodel.triton_kernel.alibi_att.context_flashattention_nopad import context_attention_fwd + + context_attention_fwd( + q, + k, + v, + out, + context.batch.b_req_idx, + layer.params.alibi, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + ) + return out + + def _prefill_paged(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd, + ) + + if layer.spec.sliding_window is not None: + sliding_window = layer.spec.sliding_window + else: + sliding_window = (-1, -1) + out = alloc_func(q.shape, q.dtype) + context_attention_fwd( + q, + k, + v, + out, + context.batch.b_req_idx, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + sliding_window=sliding_window, + ) + return out + + def _update_decode(self, context): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + draft_step = self.decode_draft_step + if draft_step > 0: + context.b_mark_mtp_shared_group = build_mtp_shared_group_markers( + context.batch.b_req_idx, hold_req_id=self.hold_req_id + ) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + if layer.params.alibi is not None: + return self._decode_alibi(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + else: + draft_step = self.decode_draft_step + q_head_num = q.shape[1] + k_head_num = k.shape[1] + if draft_step > 0: + assert q_head_num >= k_head_num, "speculative decode requires q_head_num >= k_head_num" + return self._run_decode_speculative(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + elif q_head_num == k_head_num and layer.spec.sliding_window is None: + return self._decode_mha(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + elif q_head_num >= k_head_num: + return self._decode_gqa(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + else: + raise NotImplementedError("error") + + def _decode_alibi(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.basemodel.triton_kernel.alibi_att.token_flashattention_nopad import token_attention_fwd + + out = alloc_func(q.shape, q.dtype) + token_attention_fwd( + q, + k, + v, + out, + layer.params.alibi, + context.batch.req_to_token_indexs, + context.batch.b_req_idx, + context.batch.b_kv_start_loc, + context.batch.b_seq_len, + context.batch.max_kv_seq_len, + context.batch.total_token_num, + alloc_tensor_func=alloc_func, + ) + return out + + def _decode_mha(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding import ( + token_decode_attention_flash_decoding, + ) + + out = alloc_func(q.shape, q.dtype) + token_decode_attention_flash_decoding( + q=q, batch=context.batch, cache_k=k, cache_v=v, out=out, alloc_tensor_func=alloc_func + ) + return out + + def _decode_gqa(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty): + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding import ( + gqa_token_decode_attention_flash_decoding, + ) + + if layer.spec.sliding_window is not None: + sliding_window = layer.spec.sliding_window + else: + sliding_window = (-1, -1) + out = alloc_func(q.shape, q.dtype) + gqa_token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_v=v, + max_len_in_batch=context.decode_max_kv_seq_len, + out=out, + alloc_tensor_func=alloc_func, + sliding_window=sliding_window, + ) + return out + + def _run_decode_speculative( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + ): + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( + token_decode_attention_mtp_diverse_single_token, + ) + + out = token_decode_attention_mtp_diverse_single_token( + q=q, + k=k, + v=v, + Req_to_tokens=context.batch.req_to_token_indexs, + B_req_idx=context.batch.b_req_idx, + b_seq_len=context.batch.b_seq_len, + b_mark_shared_group=context.b_mark_mtp_shared_group, + max_kv_len=context.decode_max_kv_seq_len, + alloc_tensor_func=alloc_func, + ) + return out diff --git a/lightllm/common/attention/backends/dense/triton/int4.py b/lightllm/common/attention/backends/dense/triton/int4.py new file mode 100644 index 0000000000..d6edee671c --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int4.py @@ -0,0 +1,142 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.utils.envs_utils import get_env_start_args + + +@dataclasses.dataclass(slots=True) +class Int4kvTritonPrefillContext(AttentionExecutionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class Int4kvTritonDecodeContext(AttentionExecutionContext): + decode_max_kv_seq_len: int = None + + +class Int4kvTritonAttBackend(BaseAttentionBackend): + provider = "triton" + supports_alibi = False + supports_sliding_window = False + cache_formats = frozenset({"int4kv"}) + + def __init__(self, model, spec): + super().__init__(model, spec) + self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Int4kvTritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Int4kvTritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.b_kv_start_loc = ( + torch.cumsum(context.batch.b_seq_len, dim=0, dtype=context.batch.b_seq_len.dtype) - context.batch.b_seq_len + ) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + o = self._run_prefill(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func) + return o + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + total_token_num = context.batch.total_token_num + head_dim = k.shape[2] * 2 + k_dequant = alloc_func((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device) + v_dequant = alloc_func((total_token_num, v.shape[1], head_dim), dtype=q.dtype, device=q.device) + o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) + max_kv_seq_len = context.batch.max_kv_seq_len + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import dequantize_int4kv + + dequantize_int4kv( + k=k, + k_scale=k_scale, + v=v, + v_scale=v_scale, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_seq_len=context.batch.b_seq_len, + b_req_idx=context.batch.b_req_idx, + b_kv_start_loc=context.b_kv_start_loc, + k_out=k_dequant, + v_out=v_dequant, + max_len_in_batch=max_kv_seq_len, + quant_group_size=self.quant_group_size, + ) + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + context_attention_fwd_contiguous_kv( + q=q, + k=k_dequant, + v=v_dequant, + o=o_tensor, + b_start_loc=context.batch.b_q_start_loc, + b_kv_start_loc=context.b_kv_start_loc, + b_seq_len=context.batch.b_seq_len, + max_q_input_len=context.batch.max_q_seq_len, + b_prompt_cache_len=context.batch.b_ready_cache_len, + ) + return o_tensor + + def _update_decode(self, context): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + return self._run_decode(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func) + + def _run_decode( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv.ppl_int4kv_flash_decoding import ( + token_decode_attention_flash_decoding, + ) + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + max_kv_seq_len=context.decode_max_kv_seq_len, + alloc_tensor_func=alloc_func, + ) diff --git a/lightllm/common/attention/backends/dense/triton/int8.py b/lightllm/common/attention/backends/dense/triton/int8.py new file mode 100644 index 0000000000..1bf4da78a6 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int8.py @@ -0,0 +1,182 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.diverse_utils import build_diverse_shared_group_markers +from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel, get_env_start_args + + +@dataclasses.dataclass(slots=True) +class Int8kvTritonPrefillContext(AttentionExecutionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class Int8kvTritonDecodeContext(AttentionExecutionContext): + b_shared_seq_len: torch.Tensor = None + b_mark_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None + + +class Int8kvTritonAttBackend(BaseAttentionBackend): + provider = "triton" + supports_alibi = False + supports_sliding_window = False + cache_formats = frozenset({"int8kv"}) + + def __init__(self, model, spec): + super().__init__(model, spec) + self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Int8kvTritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Int8kvTritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.b_kv_start_loc = ( + torch.cumsum(context.batch.b_seq_len, dim=0, dtype=context.batch.b_seq_len.dtype) - context.batch.b_seq_len + ) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + o = self._run_prefill(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func) + return o + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + total_token_num = context.batch.total_token_num + k_dequant = alloc_func((total_token_num, k.shape[1], k.shape[2]), dtype=q.dtype, device=q.device) + v_dequant = alloc_func((total_token_num, v.shape[1], v.shape[2]), dtype=q.dtype, device=q.device) + o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) + max_kv_seq_len = context.batch.max_kv_seq_len + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import dequantize_int8kv + + dequantize_int8kv( + k=k, + k_scale=k_scale, + v=v, + v_scale=v_scale, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_seq_len=context.batch.b_seq_len, + b_req_idx=context.batch.b_req_idx, + b_kv_start_loc=context.b_kv_start_loc, + k_out=k_dequant, + v_out=v_dequant, + max_len_in_batch=max_kv_seq_len, + quant_group_size=self.quant_group_size, + ) + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + context_attention_fwd_contiguous_kv( + q=q, + k=k_dequant, + v=v_dequant, + o=o_tensor, + b_start_loc=context.batch.b_q_start_loc, + b_kv_start_loc=context.b_kv_start_loc, + b_seq_len=context.batch.b_seq_len, + max_q_input_len=context.batch.max_q_seq_len, + b_prompt_cache_len=context.batch.b_ready_cache_len, + ) + return o_tensor + + def _update_decode(self, context): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + if enable_diverse_mode_gqa_decode_fast_kernel(): + context.b_mark_shared_group = build_diverse_shared_group_markers( + b_shared_radix_node_id=context.batch.b_shared_radix_node_id + ) + context.b_shared_seq_len = torch.where(context.b_mark_shared_group == 1, 0, context.batch.b_shared_seq_len) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + k = (k, layer.params.key_scale) + v = (v, layer.params.value_scale) + k, k_scale = k + v, v_scale = v + if enable_diverse_mode_gqa_decode_fast_kernel(): + return self._run_decode_diverse( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + else: + return self._run_decode( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + + def _run_decode_diverse( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + from .kernels.decode.int8kv.int8kv_flash_decoding_diverse import token_decode_attention_flash_decoding + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + alloc_tensor_func=alloc_func, + b_shared_seq_len=context.b_shared_seq_len, + b_mark_shared_group=context.b_mark_shared_group, + shared_streams_dict=context.workspace.get_object("int8_decode_streams", dict), + ) + + def _run_decode( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + alloc_func=torch.empty, + ): + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( + token_decode_attention_flash_decoding, + ) + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + max_len_in_batch=context.decode_max_kv_seq_len, + alloc_tensor_func=alloc_func, + ) diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py new file mode 100644 index 0000000000..5dd6a683fe --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py @@ -0,0 +1,59 @@ +import torch + + +def gqa_token_decode_attention_flash_decoding( + q: torch.Tensor, + batch, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + max_len_in_batch: int, + out=None, + alloc_tensor_func=torch.empty, + sliding_window=(-1, -1), +): + batch_size = batch.batch_size + q_head_num, head_dim = q.shape[1], q.shape[2] + calcu_shape1 = (batch_size, q_head_num, head_dim) + + from .gqa_flash_decoding_stage1 import flash_decode_stage1 + from .gqa_flash_decoding_stage2 import flash_decode_stage2 + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + # Because we need to allocate some intermediate tensors, considering parallelism and + # intermediate memory consumption, when batch_size is small, block_num is larger, + # and when batch_size is large, block_num is smaller. This achieves a better balance + # of memory consumption and performance. + BLOCK_SEQ = 256 + if batch_size <= 16: + block_num = 128 + elif batch_size <= 64: + block_num = 64 + else: + block_num = 32 + + mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") + mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") + + flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + v=cache_v, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, + max_len_in_batch=max_len_in_batch, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + sliding_window=sliding_window, + ) + flash_decode_stage2( + mid_out=mid_o, + mid_out_logexpsum=mid_o_logexpsum, + B_Seqlen=batch.b_seq_len, + out=o_tensor.view(calcu_shape1), + block_seq=BLOCK_SEQ, + sliding_window=sliding_window, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py index dc98bb3544..8b2dbc885a 100644 --- a/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py @@ -1,13 +1,10 @@ import torch -from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage1_single_token import ( - mtp_diverse_stage1_single_token, -) -from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage2_single_token import ( - mtp_diverse_stage2_single_token, -) from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size +from .stage1_single_token import mtp_diverse_stage1_single_token +from .stage2_single_token import mtp_diverse_stage2_single_token + @torch.no_grad() def token_decode_attention_mtp_diverse_single_token( diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py new file mode 100644 index 0000000000..9e76348baf --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py @@ -0,0 +1,58 @@ +from typing import Optional + +import torch + + +def token_decode_attention_flash_decoding( + q: torch.Tensor, + batch, + cache_k: torch.Tensor, + cache_k_scale: torch.Tensor, + cache_v: torch.Tensor, + cache_v_scale: torch.Tensor, + max_kv_seq_len: int, + out: Optional[torch.Tensor] = None, + alloc_tensor_func=torch.empty, +): + BLOCK_SEQ = 256 + batch_size = batch.batch_size + q_head_num = q.shape[1] + head_dim = q.shape[2] + calcu_shape1 = (batch_size, q_head_num, head_dim) + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + # 因为需要分配一些中间tensor，考虑到并行度和中间显存的消耗，batch_size 小的 + # 时候 block_num 较大， batch_size 大的时候 block_num 较小。这样可以达到较好 + # 的显存消耗和性能的平衡。 + if batch_size <= 16: + block_num = 128 + elif batch_size <= 64: + block_num = 64 + else: + block_num = 32 + + mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") + mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=q.dtype, device="cuda") + + from .int4kv_flash_decoding_stage1 import int4kv_flash_decode_stage1 + + int4kv_flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, + max_kv_seq_len=max_kv_seq_len, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + ) + + from ..int8kv.normal.int8kv_flash_decoding_stage2 import flash_decode_stage2 + + flash_decode_stage2(mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py new file mode 100644 index 0000000000..c5e3af3f36 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py @@ -0,0 +1,102 @@ +# 为 diverse mode 定制设计的 int8kv flash decoding attention 实现，可以实现更高效的多样性采样 +import torch + +from lightllm.common.attention.api.batch import AttentionBatch +from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size + +from .int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1 +from .int8kv_flash_decoding_diverse_stage2 import flash_decode_stage2 +from .int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3 + + +def token_decode_attention_flash_decoding( + q, + batch: AttentionBatch, + cache_k, + cache_k_scale, + cache_v, + cache_v_scale, + out=None, + alloc_tensor_func=torch.empty, + b_shared_seq_len=None, + b_mark_shared_group=None, + shared_streams_dict=None, +): + if b_shared_seq_len is None or b_mark_shared_group is None or shared_streams_dict is None: + raise ValueError("diverse attention requires prepared group metadata and slot-owned streams") + + if "stream1" not in shared_streams_dict: + shared_streams_dict["stream1"] = torch.cuda.Stream() + if "stream2" not in shared_streams_dict: + shared_streams_dict["stream2"] = torch.cuda.Stream() + + stream1 = shared_streams_dict["stream1"] + stream2 = shared_streams_dict["stream2"] + + q_head_num = q.shape[1] + head_dim = q.shape[2] + + BLOCK_SEQ = 256 + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len + calcu_shape1 = (batch_size, q_head_num, head_dim) + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + mid_o = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2, head_dim], dtype=torch.float32, device="cuda" + ) + mid_o_logexpsum = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2], dtype=torch.float32, device="cuda" + ) + + current_stream = torch.cuda.current_stream() + + stream1.wait_stream(current_stream) + with torch.cuda.stream(stream1): + flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + b_shared_seq_len=b_shared_seq_len, + b_mark_shared_group=b_mark_shared_group, + max_len_in_batch=batch.max_kv_seq_len, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + max_batch_group_size=get_diverse_max_batch_shared_group_size(), + ) + stream2.wait_stream(current_stream) + with torch.cuda.stream(stream2): + flash_decode_stage2( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, + b_shared_seq_len=b_shared_seq_len, + max_len_in_batch=batch.max_kv_seq_len, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + ) + + current_stream.wait_stream(stream1) + current_stream.wait_stream(stream2) + + flash_diverse_decode_stage3( + mid_out=mid_o, + mid_out_logexpsum=mid_o_logexpsum, + B_Seqlen=batch.b_seq_len, + b_shared_seq_len=b_shared_seq_len, + O=o_tensor.view(calcu_shape1), + block_seq=BLOCK_SEQ, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py index e69de29bb2..8cc4aa919b 100644 --- a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py @@ -0,0 +1 @@ +from .int8kv_flash_decoding import token_decode_attention_flash_decoding diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py new file mode 100644 index 0000000000..cb4cad2e8e --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py @@ -0,0 +1,67 @@ +from typing import Optional + +import torch + +from lightllm.common.attention.api.batch import AttentionBatch + +from .int8kv_flash_decoding_stage1 import flash_decode_stage1 +from .int8kv_flash_decoding_stage2 import flash_decode_stage2 + + +@torch.no_grad() +def token_decode_attention_flash_decoding( + q: torch.Tensor, + batch: AttentionBatch, + cache_k: torch.Tensor, + cache_k_scale: torch.Tensor, + cache_v: torch.Tensor, + cache_v_scale: torch.Tensor, + max_len_in_batch: int, + out: Optional[torch.Tensor] = None, + alloc_tensor_func=torch.empty, +): + q_head_num = q.shape[1] + head_dim = q.shape[2] + + BLOCK_SEQ = 256 + batch_size = batch.batch_size + calcu_shape1 = (batch_size, q_head_num, head_dim) + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + # 因为需要分配一些中间tensor，考虑到并行度和中间显存的消耗，batch_size 小的 + # 时候 block_num 较大， batch_size 大的时候 block_num 较小。这样可以达到较好 + # 的显存消耗和性能的平衡。 + if batch_size <= 16: + block_num = 128 + elif batch_size <= 64: + block_num = 64 + else: + block_num = 32 + + mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") + mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") + + flash_decode_stage1( + q=q.view(calcu_shape1), + k=cache_k, + k_scale=cache_k_scale, + v=cache_v, + v_scale=cache_v_scale, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_seq_len=batch.b_seq_len, + max_len_in_batch=max_len_in_batch, + mid_out=mid_o, + mid_out_logsumexp=mid_o_logexpsum, + block_seq=BLOCK_SEQ, + ) + + flash_decode_stage2( + mid_out=mid_o, + mid_out_logexpsum=mid_o_logexpsum, + B_Seqlen=batch.b_seq_len, + O=o_tensor.view(calcu_shape1), + block_seq=BLOCK_SEQ, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py new file mode 100644 index 0000000000..94489f2be9 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py @@ -0,0 +1,36 @@ +import torch + + +def token_decode_attention_flash_decoding(q, batch, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty): + BLOCK_SEQ = 256 + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len + q_head_num, head_dim = q.shape[1], q.shape[2] + calcu_shape1 = (batch_size, q_head_num, head_dim) + + from .flash_decoding_stage1 import flash_decode_stage1 + from .flash_decoding_stage2 import flash_decode_stage2 + + o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out + + mid_o = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda" + ) + mid_o_logexpsum = alloc_tensor_func( + [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda" + ) + + flash_decode_stage1( + q.view(calcu_shape1), + cache_k, + cache_v, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, + batch.max_kv_seq_len, + mid_o, + mid_o_logexpsum, + BLOCK_SEQ, + ) + flash_decode_stage2(mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + return o_tensor diff --git a/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py new file mode 100644 index 0000000000..3ca3f2dba2 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py @@ -0,0 +1,147 @@ +import torch +import triton +import triton.language as tl +from packaging.version import Version + +if Version(triton.__version__) >= Version("2.1.0"): + + @triton.jit + def _fwd_kernel( + Q, + K, + V, + sm_scale, + Out, + q_stride_s, + q_stride_h, + q_stride_d, + k_stride_s, + k_stride_h, + k_stride_d, + v_stride_s, + v_stride_h, + v_stride_d, + o_stride_s, + o_stride_h, + o_stride_d, + head_dim_act, + cu_seqlens, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + cur_batch = tl.program_id(2) + cur_head = tl.program_id(1) + start_m = tl.program_id(0) + + seq_start = tl.load(cu_seqlens + cur_batch).to(tl.int32) + seq_end = tl.load(cu_seqlens + cur_batch + 1).to(tl.int32) + seq_len = seq_end - seq_start + + # initialize offsets + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + + mask_d = offs_d < head_dim_act + off_q = cur_head * q_stride_h + (seq_start + offs_m[:, None]) * q_stride_s + offs_d[None, :] * q_stride_d + q = tl.load(Q + off_q, mask=(offs_m[:, None] < seq_len) & mask_d[None, :], other=0.0) + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + for start_n in range(0, seq_len, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + off_k = ( + (seq_start + start_n + offs_n[None, :]) * k_stride_s + + cur_head * k_stride_h + + offs_d[:, None] * k_stride_d + ) + k = tl.load(K + off_k, mask=((start_n + offs_n[None, :]) < seq_len) & mask_d[:, None], other=0.0) + + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk += tl.dot(q, k) + qk *= sm_scale + qk += tl.where((start_n + offs_n[None, :]) < seq_len, 0, float("-inf")) + + # -- compute m_ij, p, l_ij + m_ij = tl.maximum(tl.max(qk, 1), m_i) + p = tl.exp(qk - m_ij[:, None]) + l_ij = tl.sum(p, 1) + + acc_scale = tl.exp(m_i - m_ij) + acc = acc * acc_scale[:, None] + + # update acc + off_v = ( + (seq_start + start_n + offs_n[:, None]) * v_stride_s + + cur_head * v_stride_h + + offs_d[None, :] * v_stride_d + ) + v = tl.load(V + off_v, mask=((start_n + offs_n[:, None]) < seq_len) & mask_d[None, :], other=0.0) + p = p.to(v.dtype) + acc += tl.dot(p, v) + # update m_i and l_i + m_i = m_ij + l_i_new = tl.exp(l_i - m_ij) + l_ij + l_i = m_ij + tl.log(l_i_new) + + o_scale = tl.exp(m_i - l_i) + acc = acc * o_scale[:, None] + # initialize pointers to output + off_o = (seq_start + offs_m[:, None]) * o_stride_s + cur_head * o_stride_h + offs_d[None, :] * o_stride_d + out_ptrs = Out + off_o + tl.store(out_ptrs, acc, mask=(offs_m[:, None] < seq_len) & mask_d[None, :]) + return + + @torch.no_grad() + def _flash_attention_triton_fwd( + q, + k, + v, + o, + cu_seqlens, # q k v cu_seqlens, + max_seqlen, + ): + BLOCK = 64 + # shape constraints + assert q.ndim == k.ndim == v.ndim == o.ndim == 3, "q, k, v, o must be 3D tensors" + _, head_num, head_dim = q.shape + batch_size = cu_seqlens.numel() - 1 + + sm_scale = 1.0 / (head_dim**0.5) # 计算scale系数 + d_pad = triton.next_power_of_2(head_dim) + grid = (triton.cdiv(max_seqlen, BLOCK), head_num, batch_size) # batch, head, + num_warps = 4 + _fwd_kernel[grid]( + q, + k, + v, + sm_scale, + o, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + o.stride(0), + o.stride(1), + o.stride(2), + head_dim, + cu_seqlens, + BLOCK_M=BLOCK, + BLOCK_DMODEL=d_pad, + BLOCK_N=BLOCK, + num_warps=num_warps, + num_stages=2, + ) + return + +else: + raise RuntimeError(f"varlen attention requires Triton >= 2.1.0, got {triton.__version__}") diff --git a/lightllm/common/attention/backends/dense/triton/neo.py b/lightllm/common/attention/backends/dense/triton/neo.py new file mode 100644 index 0000000000..60dd33da50 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/neo.py @@ -0,0 +1,48 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.triton.backend import TritonAttBackend, TritonPrefillContext + + +@dataclasses.dataclass(slots=True) +class NeoTritonPrefillContext(TritonPrefillContext): + pass + + +class NeoTritonAttBackend(TritonAttBackend): + provider = "triton" + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NeoTritonPrefillContext(capacity, capture, workspace) + + def _update_prefill(self, context): + pass + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: DenseInputs, *, out=None) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = layer.params.key + v = layer.params.value + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, + ) + + out = alloc_func(q.shape, q.dtype) + context_attention_fwd_neo( + q, + k, + v, + out, + context.batch.b_req_idx, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + context.batch.b_image_token_end, + ) + return out diff --git a/lightllm/common/attention/backends/dense/triton/varlen.py b/lightllm/common/attention/backends/dense/triton/varlen.py new file mode 100644 index 0000000000..e7f8219dc4 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/varlen.py @@ -0,0 +1,25 @@ +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.triton.kernels.prefill.varlen import _flash_attention_triton_fwd +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class TritonVarlenAttentionOp(BaseVarlenAttentionOp): + provider = "triton" + dependency = "triton" + supports_graph = True + + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + if q.shape[0] == 0: + return out + _flash_attention_triton_fwd( + q, + k, + v, + out, + cu_seqlens, # q k v cu_seqlens, + max_seqlen, + ) + return out diff --git a/lightllm/common/attention/backends/dense/varlen.py b/lightllm/common/attention/backends/dense/varlen.py new file mode 100644 index 0000000000..a02a503724 --- /dev/null +++ b/lightllm/common/attention/backends/dense/varlen.py @@ -0,0 +1,42 @@ +import torch + +from lightllm.common.attention._environment import check_attention_environment +from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionOp +from lightllm.platform import Platform + + +class BaseVarlenAttentionOp(VarlenAttentionOp): + provider = "" + + supported_platforms = frozenset({Platform.CUDA}) + dependency = "" + min_device_capability = None + activation_dtypes = frozenset({torch.float16, torch.bfloat16, torch.float32}) + max_head_size = None + head_alignment = 1 + + @classmethod + def check_environment(cls, caps) -> None: + check_attention_environment( + caps, + provider=cls.provider, + supported_platforms=cls.supported_platforms, + min_device_capability=cls.min_device_capability, + dependency=cls.dependency, + ) + + @classmethod + def check_support(cls, inputs, caps=None) -> None: + q = inputs.query + if q.ndim != 3: + raise ValueError("varlen query must have shape [tokens, heads, dim]") + if min(q.shape[1:]) <= 0: + raise UnsupportedAttention("head count and dimension must be positive") + if q.dtype not in cls.activation_dtypes: + raise UnsupportedAttention(f"unsupported activation dtype {q.dtype}") + if caps is not None and caps.platform not in cls.supported_platforms: + raise UnsupportedAttention(f"{cls.provider} does not support platform {caps.platform.value}") + if cls.max_head_size is not None and q.shape[-1] > cls.max_head_size: + raise UnsupportedAttention(f"head dimension exceeds {cls.max_head_size}") + if q.shape[-1] % cls.head_alignment: + raise UnsupportedAttention(f"head dimension must be divisible by {cls.head_alignment}") diff --git a/lightllm/common/attention/backends/dense/xformers/varlen.py b/lightllm/common/attention/backends/dense/xformers/varlen.py new file mode 100644 index 0000000000..8edffec9b7 --- /dev/null +++ b/lightllm/common/attention/backends/dense/xformers/varlen.py @@ -0,0 +1,39 @@ +import torch + +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class XformersVarlenAttentionOp(BaseVarlenAttentionOp): + provider = "xformers" + dependency = "xformers.ops" + + @torch.no_grad() + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + from xformers import ops as xformers_ops + from xformers.ops import fmha + + q, k, v, cu_seqlens, max_seqlen = inputs + o = out + assert q.ndim == k.ndim == v.ndim == o.ndim == 3 + assert cu_seqlens is not None and cu_seqlens.ndim == 1 + assert q.shape == k.shape == v.shape == o.shape + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.int64).tolist() + seqlens = [int(L) for L in seqlens if int(L) > 0] + + if len(seqlens) == 0: + return o + if max_seqlen: + assert max(seqlens) <= max_seqlen + + attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens, device=q.device) + + q_ = q.unsqueeze(0) # [1, T, H, D] + k_ = k.unsqueeze(0) # [1, T, H, D] + v_ = v.unsqueeze(0) # [1, T, H, D] + + out = xformers_ops.memory_efficient_attention(q_, k_, v_, attn_bias=attn_bias, p=0.0) + o.copy_(out.squeeze(0)) # [T, H, D] + return o diff --git a/lightllm/common/attention/backends/gdn/flashqla/backend.py b/lightllm/common/attention/backends/gdn/flashqla/backend.py new file mode 100644 index 0000000000..eb477320b9 --- /dev/null +++ b/lightllm/common/attention/backends/gdn/flashqla/backend.py @@ -0,0 +1,17 @@ +import torch + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNBackend + + +class FlashQlaGDNBackend(TritonGDNBackend): + provider = "flashqla" + dependency = "flash_qla" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supported_phases = frozenset({AttentionPhase.PREFILL}) + cache_formats = frozenset({"gdn"}) + + def get_prefill_kernel(self): + from flash_qla import chunk_gated_delta_rule + + return chunk_gated_delta_rule diff --git a/lightllm/common/attention/backends/gdn/replay/backend.py b/lightllm/common/attention/backends/gdn/replay/backend.py new file mode 100644 index 0000000000..f328107320 --- /dev/null +++ b/lightllm/common/attention/backends/gdn/replay/backend.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import GDNInputs, GDNOutput +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.gdn import GDNMetadata +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs + +from .decode import replay + + +@dataclasses.dataclass(slots=True) +class ReplayGDNPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class ReplayGDNDecodeContext(LinearDecodeContext): + pass + + +class ReplayGDNTritonAttBackend( + GDNMetadata, BaseAttentionBackend[ReplayGDNPrefillContext | ReplayGDNDecodeContext, GDNInputs, GDNOutput] +): + provider = "triton" + family = AttentionFamily.GDN + cache_formats = frozenset({"gdn"}) + + def __init__(self, model, spec): + super().__init__(model=model, spec=spec) + self._init_linear_layer_metadata(model.config, model.tp_world_size_) + self.prefill_kernel = self.get_prefill_kernel() + + def get_prefill_kernel(self): + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + def _allocate_prefill_context(self, capacity, capture, workspace): + return ReplayGDNPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return ReplayGDNDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + reqs = context.batch.b_req_idx + context.b_conv_buffer_idx = reqs + context.b_ssm_buffer_idx = reqs + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None) -> GDNOutput: + conv_states, ssm_states = layer.params.conv, layer.params.ssm + if self.mtp_step > 0: + conv_states = conv_states[:, :, : -self.mtp_step] + mixed_qkv, z, b, a = self._split_qkvzba(inputs.mixed_qkvzba) + core_attn_out = self._prefill(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + return GDNOutput(core_attn_out, z) + + def _prefill(self, layer, context, mixed_qkv, conv_states, ssm_states, a, b): + g, beta = fused_gdn_gating(layer.params.a_log, a, b, layer.params.dt_bias) + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer.params.conv_weight, + bias=layer.params.conv_bias, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation=self.activation, + ).transpose(0, 1) + query, key, value = self._rearrange_mixed_qkv(mixed_qkv) + core_attn_out, last_state = self.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=ssm_states[context.b_ssm_buffer_idx], + output_final_state=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[context.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + def _update_decode(self, context): + update_linear_decode(context, self.decode_draft_step, self.uses_dynamic_spec_verify_layout(), replay=True) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None) -> GDNOutput: + mixed_qkv, z, b, a = self._split_qkvzba(inputs.mixed_qkvzba) + state = layer.params.replay_state + conv_states = layer.params.conv + draft_step = self.decode_draft_step + if draft_step > 0: + core = self._mtp(layer, context, mixed_qkv, conv_states, a, b, draft_step, state) + else: + core, z = self._decode(layer, context, mixed_qkv, z, conv_states, a, b, state) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) + return GDNOutput(core, z) + + def _decode(self, layer, context, mixed_qkv, z, conv_states, a, b, state): + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer.params.conv_weight, + layer.params.conv_bias, + context.b_conv_buffer_idx, + self.activation, + self.conv_kernel_dim, + self.tp_num_k_heads, + self.head_k_dim, + self.tp_num_v_heads, + self.head_v_dim, + ) + query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) + key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) + value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) + output = replay( + state, + state.storage_layer(layer.logical_layer_id), + query, + key, + value, + a, + b, + layer.params.a_log, + layer.params.dt_bias, + context.b_conv_buffer_idx, + context.b_cu_seqlens, + ) + return output, z + + def _mtp(self, layer, context, mixed_qkv, conv_states, a, b, draft_step, state): + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer.params.conv_weight, + mtp_step=draft_step, + bias=layer.params.conv_bias, + activation=self.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + query, key, value = torch.split( + mixed_qkv, + [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], + dim=-1, + ) + query = query.view(-1, self.tp_num_k_heads, self.head_k_dim) + key = key.view(-1, self.tp_num_k_heads, self.head_k_dim) + value = value.view(-1, self.tp_num_v_heads, self.head_v_dim) + output = replay( + state, + state.storage_layer(layer.logical_layer_id), + query, + key, + value, + a, + b, + layer.params.a_log, + layer.params.dt_bias, + context.b_conv_buffer_idx, + context.b_cu_seqlens, + ) + return output + + +class ReplayGDNFlashQlaAttBackend(ReplayGDNTritonAttBackend): + provider = "flashqla" + dependency = "flash_qla" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supported_phases = frozenset({AttentionPhase.PREFILL}) + + def get_prefill_kernel(self): + from flash_qla import chunk_gated_delta_rule + + return chunk_gated_delta_rule diff --git a/lightllm/common/attention/backends/gdn/triton/backend.py b/lightllm/common/attention/backends/gdn/triton/backend.py new file mode 100644 index 0000000000..1c567f78df --- /dev/null +++ b/lightllm/common/attention/backends/gdn/triton/backend.py @@ -0,0 +1,208 @@ +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import GDNInputs, GDNOutput +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.gdn import GDNMetadata +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, +) + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +@dataclasses.dataclass(slots=True) +class GDNPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class GDNDecodeContext(LinearDecodeContext): + pass + + +class TritonGDNBackend(GDNMetadata, BaseAttentionBackend[GDNPrefillContext | GDNDecodeContext, GDNInputs, GDNOutput]): + provider = "triton" + cache_formats = frozenset({"gdn"}) + family = AttentionFamily.GDN + + def __init__(self, model: "TpPartBaseModel", spec): + super().__init__(model=model, spec=spec) + self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) + self.prefill_kernel = self.get_prefill_kernel() + + def get_prefill_kernel(self): + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + def _allocate_prefill_context(self, capacity, capture, workspace): + return GDNPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return GDNDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + mtp_step = self.mtp_step + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx * (mtp_step + 1) + return + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None) -> GDNOutput: + mixed_qkvzba: torch.Tensor = inputs.mixed_qkvzba + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + if self.mtp_step > 0: + conv_states = conv_states[:, :, : -self.mtp_step] + mixed_qkv, z, b, a = self._split_qkvzba(mixed_qkvzba) + core_attn_out = self._run_prefill(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + return GDNOutput(core_attn_out, z) + + def _run_prefill( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + ): + g, beta = fused_gdn_gating(layer.params.a_log, a, b, layer.params.dt_bias) + mixed_qkv = mixed_qkv.transpose(0, 1) + out_tensor = causal_conv1d_fn( + mixed_qkv, + layer.params.conv_weight, + bias=layer.params.conv_bias, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation=self.activation, + ) + mixed_qkv = out_tensor.transpose(0, 1) + query, key, value = self._rearrange_mixed_qkv(mixed_qkv) + initial_state = ssm_states[context.b_ssm_buffer_idx] + core_attn_out, last_recurrent_state = self.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=initial_state, + output_final_state=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[context.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + def _update_decode(self, context): + update_linear_decode(context, self.decode_draft_step, self.uses_dynamic_spec_verify_layout(), replay=False) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None) -> GDNOutput: + mixed_qkvzba = inputs.mixed_qkvzba + mixed_qkv, z, b, a = self._split_qkvzba(mixed_qkvzba) + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + draft_step = self.decode_draft_step + if draft_step > 0: + core_attn_out = self._run_mtp(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + else: + core_attn_out, z = self._run_decode(layer, context, mixed_qkv, z, conv_states, ssm_states, a, b) + return GDNOutput(core_attn_out, z) + + def _run_decode( + self, + layer, + context, + mixed_qkv: torch.Tensor, + z: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + ): + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer.params.conv_weight, + layer.params.conv_bias, + context.b_conv_buffer_idx, + self.activation, + self.conv_kernel_dim, + self.tp_num_k_heads, + self.head_k_dim, + self.tp_num_v_heads, + self.head_v_dim, + ) + core_attn_out, _ = fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + inplace_final_state=True, + ssm_state_indices=context.b_ssm_buffer_idx, + use_qk_l2norm_in_kernel=True, + A_log=layer.params.a_log, + dt_bias=layer.params.dt_bias, + a_raw=a, + b_raw=b, + ) + return GDNOutput(core_attn_out, z) + + def _run_mtp( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + ): + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + cu_seqlens_q = context.b1_mtp_cu_q_seq_len + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer.params.conv_weight, + mtp_step=self.decode_draft_step, + bias=layer.params.conv_bias, + activation=self.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=cu_seqlens_q, + ) + query, key, value = self._rearrange_mixed_qkv(mixed_qkv, decode=False) + assert context.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" + core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + cu_seqlens=cu_seqlens_q.to(torch.long), + ssm_state_indices=context.b_ssm_buffer_idx, + ssm_state_write_indices=context.b_ssm_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + A_log=layer.params.a_log, + dt_bias=layer.params.dt_bias, + a_raw=a, + b_raw=b, + ) + return core_attn_out diff --git a/lightllm/common/attention/backends/kda/replay/backend.py b/lightllm/common/attention/backends/kda/replay/backend.py new file mode 100644 index 0000000000..38f4b442c1 --- /dev/null +++ b/lightllm/common/attention/backends/kda/replay/backend.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import KDAInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate + +from .decode import replay + + +@dataclasses.dataclass(slots=True) +class ReplayKDAPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class ReplayKDADecodeContext(LinearDecodeContext): + pass + + +class ReplayKDALinearAttBackend( + BaseAttentionBackend[ReplayKDAPrefillContext | ReplayKDADecodeContext, KDAInputs, torch.Tensor] +): + provider = "triton" + family = AttentionFamily.KDA + cache_formats = frozenset({"kda"}) + + def __init__(self, model, spec): + super().__init__(model=model, spec=spec) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return ReplayKDAPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return ReplayKDADecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + reqs = context.batch.b_req_idx + context.b_conv_buffer_idx = reqs + context.b_ssm_buffer_idx = reqs + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + conv_states, ssm_states = layer.params.conv, layer.params.ssm + conv_states = conv_states[..., : self.conv_kernel_size - 1] + mixed_qkv = causal_conv1d_fn( + inputs.mixed_qkv.transpose(0, 1), + layer.params.conv_weight, + bias=None, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + ).transpose(0, 1) + q, k, v = self.split_qkv(mixed_qkv) + q = q.view(1, -1, self.tp_num_heads, self.head_dim) + k = k.view(1, -1, self.tp_num_heads, self.head_dim) + v = v.view(1, -1, self.tp_num_heads, self.head_dim) + raw_gate = inputs.raw_gate.view(1, -1, self.tp_hidden_size) + raw_beta = inputs.raw_beta.view(1, -1, self.tp_num_heads) + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, self.tp_num_heads, self.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer.params.a_log, + g_bias=layer.params.dt_bias, + initial_state=ssm_states[context.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=self.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + def _update_decode(self, context): + update_linear_decode(context, self.decode_draft_step, self.uses_dynamic_spec_verify_layout(), replay=True) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + state = layer.params.replay_state + conv_states = layer.params.conv + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=context.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if self.decode_draft_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=self.decode_draft_step, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update( + inputs.mixed_qkv, + conv_states, + layer.params.conv_weight, + **conv_kwargs, + ) + q, k, v = self.split_qkv(mixed_qkv) + q = q.view(-1, self.tp_num_heads, self.head_dim) + k = k.view(-1, self.tp_num_heads, self.head_dim) + v = v.view(-1, self.tp_num_heads, self.head_dim) + raw_gate = inputs.raw_gate.reshape(-1, self.tp_hidden_size) + raw_beta = inputs.raw_beta.reshape(-1, self.tp_num_heads) + storage_layer = state.storage_layer(layer.logical_layer_id) + output = replay( + state_ssm=state.ssm_state[storage_layer], + state_keys=state.keys[storage_layer], + state_deltas=state.deltas[storage_layer], + state_gates=state.gates[storage_layer], + state_raw_keys=state.raw_keys[storage_layer], + state_raw_values=state.raw_values[storage_layer], + state_betas=state.betas[storage_layer], + state_cursors=state.cursors, + state_tile_sync_mark=state.tile_sync_mark, + q=q, + k=k, + v=v, + a=raw_gate, + b=raw_beta, + a_log=layer.params.a_log, + bias=layer.params.dt_bias, + reqs=context.b_conv_buffer_idx, + cu_seqlens=context.b_cu_seqlens, + state_verify_width=state.verify_width, + state_compute_mode=state.compute_mode, + state_gate_lower_bound=state.gate_lower_bound, + ) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) + return output diff --git a/lightllm/common/attention/backends/kda/triton/backend.py b/lightllm/common/attention/backends/kda/triton/backend.py new file mode 100644 index 0000000000..02ef71d436 --- /dev/null +++ b/lightllm/common/attention/backends/kda/triton/backend.py @@ -0,0 +1,142 @@ +"""KDA attention backend for GLM-5-Next.""" +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import KDAInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +@dataclasses.dataclass(slots=True) +class KDAPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class KDADecodeContext(LinearDecodeContext): + pass + + +class KDALinearAttBackend(BaseAttentionBackend[KDAPrefillContext | KDADecodeContext, KDAInputs, torch.Tensor]): + provider = "triton" + cache_formats = frozenset({"kda"}) + family = AttentionFamily.KDA + + def __init__(self, model: "TpPartBaseModel", spec): + super().__init__(model=model, spec=spec) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def _allocate_prefill_context(self, capacity, capture, workspace): + return KDAPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return KDADecodeContext(capacity, capture, workspace) + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + def _update_prefill(self, context): + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx * (self.mtp_step + 1) + + def _forward_prefill(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + mixed_qkv = inputs.mixed_qkv + raw_gate = inputs.raw_gate + raw_beta = inputs.raw_beta + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + conv_states = conv_states[..., : self.conv_kernel_size - 1] + q, k, v = causal_conv1d_qkv( + mixed_qkv, + layer.params.conv_weight, + conv_states, + context.batch.b1_cu_q_seq_len, + context.b_conv_buffer_idx, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + ) + raw_beta = raw_beta.view(1, -1, self.tp_num_heads) + initial_state = ssm_states[context.b_ssm_buffer_idx].contiguous() + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, self.tp_num_heads, self.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer.params.a_log, + g_bias=layer.params.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=False, + cu_seqlens=context.batch.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=self.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + def _update_decode(self, context): + update_linear_decode(context, self.mtp_step, self.uses_dynamic_spec_verify_layout(), replay=False) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + mixed_qkv = inputs.mixed_qkv + raw_gate = inputs.raw_gate + raw_beta = inputs.raw_beta + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=context.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if self.mtp_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=self.mtp_step, + num_accepted_tokens=context.accepted_tokens(layer.params.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update(mixed_qkv, conv_states, layer.params.conv_weight, **conv_kwargs) + q, k, v = self.split_qkv(mixed_qkv) + shape = (1, -1) if self.mtp_step > 0 else (-1, 1) + q = q.view(*shape, self.tp_num_heads, self.head_dim) + k = k.view(*shape, self.tp_num_heads, self.head_dim) + v = v.view(*shape, self.tp_num_heads, self.head_dim) + raw_gate = raw_gate.view(*shape, self.tp_hidden_size) + raw_beta = raw_beta.view(*shape, self.tp_num_heads) + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer.params.a_log, + gate_bias=layer.params.dt_bias, + initial_state=ssm_states, + lower_bound=self.lower_bound, + inplace_final_state=True, + ssm_state_indices=context.b_ssm_buffer_idx, + cu_seqlens=context.b1_mtp_cu_q_seq_len, + num_accepted_tokens=(context.accepted_tokens(layer.params.mtp_state_index) if self.mtp_step else None), + ) + return output diff --git a/lightllm/common/attention/backends/mla/fa3/backend.py b/lightllm/common/attention/backends/mla/fa3/backend.py new file mode 100644 index 0000000000..e4739edb8e --- /dev/null +++ b/lightllm/common/attention/backends/mla/fa3/backend.py @@ -0,0 +1,121 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends.base import AttentionExecutionContext +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend +from lightllm.utils.sgl_utils import flash_attn_varlen_func, flash_attn_with_kvcache + + +@dataclasses.dataclass(slots=True) +class MlaFa3PrefillContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + causal: bool = None + + +@dataclasses.dataclass(slots=True) +class MlaFa3DecodeContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + b_att_seq_len: torch.Tensor = None + decode_max_q_seq_len: int = None + decode_max_kv_seq_len: int = None + causal: bool = None + + +class MlaFa3AttBackend(Fa3AttBackend): + supports_noncausal_prefill = True + provider = "fa3" + max_head_size = None + cache_formats = frozenset({"None"}) + family = AttentionFamily.MLA + + def _allocate_prefill_context(self, capacity, capture, workspace): + return MlaFa3PrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return MlaFa3DecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + context.causal = self.uses_causal_attention() + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + return self._run_prefill(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _run_prefill( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs, alloc_func=torch.empty + ) -> torch.Tensor: + k_nope, k_rope = k + q_head_num = q.shape[1] + k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) + assert q.ndim == 3 and k.ndim == 3 and (v.ndim == 3) + softmax_scale = layer.spec.softmax_scale + o_tensor = flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + max_seqlen_k=context.batch.max_kv_seq_len, + softmax_scale=softmax_scale, + causal=context.causal, + return_softmax_lse=False, + ) + return o_tensor + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: MLADecodeInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + return self._run_decode(layer, context, q=q, k=k, alloc_func=alloc_func) + + def _run_decode( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + alloc_func=torch.empty, + ): + q_nope, q_rope = q + kv = k + qk_rope_head_dim = layer.params.qk_rope_head_dim + kv_lora_rank = kv.shape[-1] - qk_rope_head_dim + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, self.infer_page_size, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, self.infer_page_size, 1, kv_lora_rank) + k_descale, v_descale = (None, None) + softmax_scale = layer.spec.softmax_scale + o_tensor = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=kv_nope, + qv=q_nope, + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.decode_max_q_seq_len, + softmax_scale=softmax_scale, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/mla/flashinfer/backend.py b/lightllm/common/attention/backends/mla/flashinfer/backend.py new file mode 100644 index 0000000000..980b086530 --- /dev/null +++ b/lightllm/common/attention/backends/mla/flashinfer/backend.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, Tuple + +import torch +import triton + +from lightllm.common.attention.api import BackendResources +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.execution import WorkspaceRequirement +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.common.basemodel.triton_kernel.flashinfer_mla_plan import fill_mla_decode_plan_for_cuda_graph +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor +from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index +from lightllm.utils.dist_utils import get_dp_world_size + +if TYPE_CHECKING: + from flashinfer.mla import BatchMLAPagedAttentionWrapper + from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper + + +@dataclasses.dataclass(slots=True) +class MlaFlashInferPrefillContext(AttentionExecutionContext): + prefill_wrapper: BatchPrefillWithRaggedKVCacheWrapper | None = None + + +@dataclasses.dataclass(slots=True) +class MlaFlashInferDecodeContext(AttentionExecutionContext): + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + q_indptr: torch.Tensor = None + decode_wrapper: BatchMLAPagedAttentionWrapper | None = None + + +class MlaFlashInferAttBackend( + BaseAttentionBackend[ + MlaFlashInferPrefillContext | MlaFlashInferDecodeContext, MLAPrefillInputs | MLADecodeInputs, torch.Tensor + ] +): + supports_noncausal_prefill = True + provider = "flashinfer" + dependency = "flashinfer" + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + supports_expert_parallel = False + cache_formats = frozenset({"None"}) + family = AttentionFamily.MLA + workspace_buffer_key = "flashinfer_mla" + workspace_buffer_size = 256 * 1024 * 1024 + + def __init__(self, model, spec): + set_flashinfer_envs() + super().__init__(model=model, spec=spec) + num_heads = model.config["num_attention_heads"] + self.tp_q_head_num = num_heads // get_dp_world_size() + self.qk_nope_head_dim = model.qk_nope_head_dim + self.qk_rope_head_dim = model.qk_rope_head_dim + self.kv_lora_rank = model.kv_lora_rank + self.v_head_dim = model.v_head_dim + self.q_data_type = spec.dtype + self.kv_data_type = spec.cache.dtype + self.max_seq_length = model.max_seq_length + self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) + + def configure_spec(self, spec): + self.tp_q_head_num = spec.num_q_heads + + def resource_requirements(self, capacity, *, capture): + return BackendResources( + (WorkspaceRequirement(self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),) + ) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return MlaFlashInferPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return MlaFlashInferDecodeContext(capacity, capture, workspace) + + def _prepare_prefill(self, context): + import flashinfer + + q_starts = context.batch.b1_cu_q_seq_len.int() + kv_starts = context.batch.b1_cu_kv_seq_len.int() + if context.prefill_wrapper is None: + context.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), "NHD" + ) + context.prefill_wrapper.plan( + qo_indptr=q_starts, + kv_indptr=kv_starts, + num_qo_heads=self.tp_q_head_num, + num_kv_heads=self.tp_q_head_num, + head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, + head_dim_vo=self.v_head_dim, + q_data_type=self.q_data_type, + causal=self.causal, + sm_scale=self.spec.softmax_scale, + ) + return + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + return self._run_prefill(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + alloc_func=torch.empty, + ) -> torch.Tensor: + k_nope, k_rope = k + o_tensor = alloc_func((q.shape[0], q.shape[1], v.shape[-1]), dtype=q.dtype, device=q.device) + q_head_num = q.shape[1] + k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) + context.prefill_wrapper.run(q, k, v, out=o_tensor) + return o_tensor + + def _refresh_decode_plan(self, context, kv_lens: torch.Tensor, max_kv_len: int): + fill_mla_decode_plan_for_cuda_graph( + decode_wrapper=context.decode_wrapper, + kv_page_indptr=context.kv_starts, + kv_lens=kv_lens, + batch_size=context.batch.batch_size, + num_heads=self.tp_q_head_num, + max_kv_len=max_kv_len, + ) + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: MLADecodeInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + return self._run_decode(layer, context, q=q, k=k, alloc_func=alloc_func) + + def _run_decode( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + alloc_func=torch.empty, + ): + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = q + o_tensor = alloc_func(q_nope.shape, dtype=q_nope.dtype, device=q_nope.device) + context.decode_wrapper.run( + q_nope, + q_rope, + k[:, :, :-qk_rope_head_dim].view(-1, self.infer_page_size, 1, k.shape[-1] - qk_rope_head_dim), + k[:, :, -qk_rope_head_dim:].view(-1, self.infer_page_size, 1, qk_rope_head_dim), + out=o_tensor, + return_lse=False, + ) + return o_tensor + + def _prepare_decode(self, context): + import flashinfer + + batch = context.batch + b_page_len = (batch.b_seq_len + self.infer_page_size - 1) // self.infer_page_size + starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + if context.kv_starts is None: + context.kv_starts = starts + context.kv_indices = context.workspace.get( + "flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32 + ) + context.q_indptr = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) + else: + context.kv_starts.copy_(starts, non_blocking=True) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=context.kv_starts[:-1], + max_token_len=context.current_max_kv_seq_len, + out_page_indices=context.kv_indices, + page_size=self.infer_page_size, + ) + if context.decode_wrapper is None: + context.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + context.workspace.get(self.workspace_buffer_key, self.workspace_buffer_size), + use_cuda_graph=context.capture, + qo_indptr=context.q_indptr, + kv_indices=context.kv_indices, + kv_indptr=context.kv_starts, + kv_len_arr=batch.b_seq_len, + ) + context.decode_wrapper.plan( + context.q_indptr, + context.kv_starts, + context.kv_indices, + batch.b_seq_len, + self.tp_q_head_num, + self.kv_lora_rank, + self.qk_rope_head_dim, + self.infer_page_size, + False, + self.spec.softmax_scale, + self.q_data_type, + self.kv_data_type, + ) + else: + self._refresh_decode_plan(context, batch.b_seq_len, context.current_max_kv_seq_len) diff --git a/lightllm/common/attention/backends/mla/triton/backend.py b/lightllm/common/attention/backends/mla/triton/backend.py new file mode 100644 index 0000000000..ae0565127e --- /dev/null +++ b/lightllm/common/attention/backends/mla/triton/backend.py @@ -0,0 +1,158 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend + + +@dataclasses.dataclass(slots=True) +class MlaTritonPrefillContext(AttentionExecutionContext): + pass + + +@dataclasses.dataclass(slots=True) +class MlaTritonDecodeContext(AttentionExecutionContext): + block_seq: torch.Tensor = None + batch_start: torch.Tensor = None + indices_initialized: bool = False + + +class MlaTritonAttBackend( + BaseAttentionBackend[ + MlaTritonPrefillContext | MlaTritonDecodeContext, MLAPrefillInputs | MLADecodeInputs, torch.Tensor + ] +): + provider = "triton" + cache_formats = frozenset({"None"}) + family = AttentionFamily.MLA + + def _allocate_prefill_context(self, capacity, capture, workspace): + return MlaTritonPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return MlaTritonDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + pass + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + alloc_func = layer.alloc_tensor + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + if q.shape[2] == 256: + return self._run_prefill_contiguous(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + else: + return self._run_prefill(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + inputs, + alloc_func=torch.empty, + ): + from lightllm.common.attention.backends.mla.triton.kernels.prefill import context_attention_fwd_with_v + + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = (q[:, :, :-qk_rope_head_dim], q[:, :, -qk_rope_head_dim:]) + o_tensor = alloc_func((q_nope.shape[0], q_nope.shape[1], v.shape[-1]), dtype=q_nope.dtype, device=q_nope.device) + k_nope, k_rope = k + softmax_scale = layer.spec.softmax_scale + context_attention_fwd_with_v( + q_nope, + q_rope, + k_nope, + k_rope, + v, + o_tensor, + context.batch.b_q_start_loc, + context.batch.b1_cu_kv_seq_len, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + softmax_scale, + ) + return o_tensor + + def _run_prefill_contiguous( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + inputs, + alloc_func=torch.empty, + ): + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + k_nope, k_rope = k + q_head_num = q.shape[1] + k_merged = torch.cat([k_nope.expand(-1, q_head_num, -1), k_rope.expand(-1, q_head_num, -1)], dim=-1) + o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) + context_attention_fwd_contiguous_kv( + q, + k_merged, + v, + o_tensor, + context.batch.b_q_start_loc, + context.batch.b1_cu_kv_seq_len, + context.batch.b_seq_len, + context.batch.max_q_seq_len, + context.batch.b_ready_cache_len, + ) + return o_tensor + + def _update_decode(self, context): + context.block_seq = torch.empty((1,), dtype=torch.int64, device=context.batch.device) + context.batch_start = torch.empty((context.batch.batch_size,), dtype=torch.int64, device=context.batch.device) + context.indices_initialized = False + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: MLADecodeInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + return self._run_decode(layer, context, q=q, k=k, alloc_func=alloc_func) + + def _run_decode( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + alloc_func=torch.empty, + ): + softmax_scale = layer.spec.softmax_scale + from lightllm.common.attention.backends.mla.triton.kernels.decode import ( + gqa_token_decode_attention_flash_decoding, + ) + + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = q + kv = k + out = gqa_token_decode_attention_flash_decoding( + q_nope=q_nope, + q_rope=q_rope, + kv_nope=kv[:, :, :-qk_rope_head_dim], + kv_rope=kv[:, :, -qk_rope_head_dim:], + batch=context.batch, + softmax_scale=softmax_scale, + alloc_tensor_func=alloc_func, + block_seq=context.block_seq, + batch_start=context.batch_start, + initialize=not context.indices_initialized, + ) + context.indices_initialized = True + return out diff --git a/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py index e69de29bb2..fb0609a401 100644 --- a/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py @@ -0,0 +1 @@ +from .gqa_flash_decoding import gqa_token_decode_attention_flash_decoding diff --git a/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py new file mode 100644 index 0000000000..a618c9cdea --- /dev/null +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py @@ -0,0 +1,147 @@ +import os +from typing import List + +import torch +import torch.multiprocessing as mp +import triton +import triton.language as tl + +from lightllm.utils.device_utils import get_device_sm_count +from lightllm.utils.log_utils import init_logger + +from .gqa_flash_decoding_config import MlaDecodeAttentionKernelConfig + +logger = init_logger(__name__) + + +def gqa_token_decode_attention_flash_decoding( + q_nope, + q_rope, + kv_nope, + kv_rope, + batch, + softmax_scale, + out=None, + alloc_tensor_func=torch.empty, + block_seq=None, + batch_start=None, + initialize=True, + **run_config, +): + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len + + q_head_num, kv_lora_rank = q_nope.shape[1], q_nope.shape[2] + q_rope_dim = q_rope.shape[2] + assert q_rope_dim == 64 + + calcu_shape1 = (batch_size, q_head_num, kv_lora_rank) + calcu_shape2 = (batch_size, q_head_num, q_rope_dim) + + if not run_config: + if torch.cuda.is_current_stream_capturing(): + avg_seq_len_in_batch = max_kv_seq_len + else: + avg_seq_len_in_batch = batch.total_token_num // batch_size + + run_config = MlaDecodeAttentionKernelConfig.try_to_get_best_config( + batch_size=batch_size, + avg_seq_len_in_batch=avg_seq_len_in_batch, + q_head_num=q_head_num, + q_head_dim=kv_lora_rank, + q_rope_dim=q_rope_dim, + out_dtype=torch.bfloat16, + ) + + BLOCK_N = run_config["BLOCK_N"] + + from .gqa_flash_decoding_stage1 import flash_decode_stage1 + from .gqa_flash_decoding_stage2 import flash_decode_stage2 + + o_tensor = alloc_tensor_func(q_nope.shape, q_nope.dtype, q_nope.device) if out is None else out + + fake_decode_att_block_seq = torch.empty([0], dtype=torch.int64, device="cuda") + mid_o = torch.empty([q_head_num, 0, kv_lora_rank], dtype=torch.float32, device="cuda") + mid_o_logexpsum = torch.empty([q_head_num, 0], dtype=torch.float32, device="cuda") + + vsm_count = flash_decode_stage1( + fake_decode_att_block_seq, + q_nope.view(calcu_shape1), + q_rope.view(calcu_shape2), + kv_nope, + kv_rope, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, + mid_o, + mid_o_logexpsum, + softmax_scale, + get_sm_count=True, + **run_config, + ) + + if initialize: + _fwd_kernel_calcu_index_and_block_seq[(1,)]( + batch.b_seq_len, + block_seq, + batch_start, + vsm_count, + batch_size, + BLOCK_N=BLOCK_N, + MAX_BATCH_SIZE=triton.next_power_of_2(batch_size), + num_warps=4, + ) + + mid_o = torch.empty([q_head_num, vsm_count * 4 + batch_size, kv_lora_rank], dtype=torch.float32, device="cuda") + mid_o_logexpsum = torch.empty([q_head_num, vsm_count * 4 + batch_size], dtype=torch.float32, device="cuda") + + flash_decode_stage1( + block_seq, + q_nope.view(calcu_shape1), + q_rope.view(calcu_shape2), + kv_nope, + kv_rope, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, + mid_o, + mid_o_logexpsum, + softmax_scale, + get_sm_count=False, + **run_config, + ) + + flash_decode_stage2( + block_seq, batch_start, mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), **run_config + ) + return o_tensor + + +@triton.jit +def _fwd_kernel_calcu_index_and_block_seq( + b_seq_len_ptr, + mid_o_decode_att_block_seq_ptr, + mid_o_batch_start_index_ptr, + num_sm, + batch_size, + BLOCK_N: tl.constexpr, + MAX_BATCH_SIZE: tl.constexpr, +): + b_seq_len = tl.load( + b_seq_len_ptr + tl.arange(0, MAX_BATCH_SIZE), mask=tl.arange(0, MAX_BATCH_SIZE) < batch_size, other=0 + ) + total_token_num = tl.sum(b_seq_len) + + block_seq = tl.cast(total_token_num / (num_sm * 4), dtype=tl.int32) + 1 + block_seq = tl.cdiv(block_seq, BLOCK_N) * BLOCK_N + + block_seq_len = tl.cdiv(b_seq_len, block_seq) + cumsum_seq_len = tl.cumsum(block_seq_len) + batch_start_index = cumsum_seq_len - block_seq_len + tl.store( + mid_o_batch_start_index_ptr + tl.arange(0, MAX_BATCH_SIZE), + batch_start_index, + mask=tl.arange(0, MAX_BATCH_SIZE) < batch_size, + ) + tl.store(mid_o_decode_att_block_seq_ptr, block_seq) + return diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py new file mode 100644 index 0000000000..01839e63de --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py @@ -0,0 +1,179 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention.api.spec import AttentionFamily +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend +from lightllm.utils.dist_utils import get_current_device_id + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaSparsePrefillContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + # Query-row metadata is consumed by model-specific sparse indexers. It is + # kept on the context so the indexer never reaches into InferStateInfo. + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + q_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaSparseDecodeContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + nsa_cache_seqlens: torch.Tensor = None + nsa_cu_seqlens_k_new: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + + +class NsaFlashMlaSparseAttBackend( + BaseAttentionBackend[ + NsaFlashMlaSparsePrefillContext | NsaFlashMlaSparseDecodeContext, + SparseMLAPrefillInputs | SparseMLADecodeInputs, + torch.Tensor, + ] +): + provider = "flashmla_sparse" + dependency = "flash_mla" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + cache_formats = frozenset({"None"}) + family = AttentionFamily.SPARSE_MLA + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NsaFlashMlaSparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return NsaFlashMlaSparseDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + batch = context.batch + context.ragged_mem_index = torch.empty(batch.total_token_num, dtype=torch.int32, device=get_current_device_id()) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=batch.b_seq_len, + b_q_seq_len=batch.b_q_seq_len, + b_req_idx=batch.b_req_idx, + req_to_token_index=batch.req_to_token_indexs, + q_token_num=batch.token_num, + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=batch.hold_req_id, + ) + self._update_query_metadata(context) + return + + @staticmethod + def _update_query_metadata(context): + """Materialize per-query request metadata on the execution context. + + Prefill positions are supplied by the model's input preparation and + therefore remain correct for prompt-cache and speculative layouts. + ``output_size`` keeps the operation graph-capturable without a host + read of the GPU-side query lengths. + """ + batch = context.batch + position_ids = batch.position_ids + if position_ids is None: + # Unit-level callers may construct an AttentionBatch directly. + # The runtime always supplies position_ids, but contiguous query + # ranges are an unambiguous fallback for those callers. + history = batch.b_seq_len - batch.b_q_seq_len + starts = torch.repeat_interleave(history, batch.b_q_seq_len, output_size=batch.token_num) + row_starts = torch.repeat_interleave( + batch.b1_cu_q_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num + ) + position_ids = starts + torch.arange(batch.token_num, device=batch.device, dtype=starts.dtype) - row_starts + context.q_kv_seq_len = position_ids + 1 + context.q_req_idx = torch.repeat_interleave(batch.b_req_idx, batch.b_q_seq_len, output_size=batch.token_num) + context.q_kv_start_loc = torch.repeat_interleave( + batch.b1_cu_kv_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num + ) + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: SparseMLAPrefillInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.params.kv + return self._run_prefill(layer, context, q=q, kv=k, inputs=inputs) + + def _run_prefill(self, layer, context, q: torch.Tensor, kv: torch.Tensor, inputs) -> torch.Tensor: + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = layer.spec.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + mla_out, _, _ = flash_mla_sparse_fwd( + q=q, kv=kv, indices=topk_mem_indices, sm_scale=softmax_scale, d_v=kv_lora_rank + ) + return mla_out + + def _update_decode(self, context): + context.q_kv_seq_len = context.batch.b_seq_len + context.q_req_idx = context.batch.b_req_idx + context.ragged_mem_index = context.workspace.get( + "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 + ) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=context.batch.b_seq_len, + b_q_seq_len=context.batch.b_q_seq_len, + b_req_idx=context.batch.b_req_idx, + req_to_token_index=context.batch.req_to_token_indexs, + q_token_num=context.batch.b_seq_len.shape[0], + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=context.batch.hold_req_id, + ) + context.nsa_cache_seqlens = torch.minimum( + torch.full(size=(context.batch.batch_size,), fill_value=2048, dtype=torch.int32, device="cuda"), + context.batch.b_seq_len, + ) + padded_seq_lens = torch.zeros(size=(context.nsa_cache_seqlens.shape[0] + 1,), dtype=torch.int32, device="cuda") + padded_seq_lens[1:].copy_(context.nsa_cache_seqlens, non_blocking=True) + context.nsa_cu_seqlens_k_new = padded_seq_lens.cumsum(dim=0, dtype=torch.int32) + + def _forward_decode( + self, layer: LayerAttentionBinding, context, inputs: SparseMLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + return self._run_decode(layer, context, q=q, kv=k, inputs=inputs) + + def _run_decode( + self, layer, context, q: Tuple[torch.Tensor, torch.Tensor], kv: torch.Tensor, inputs + ) -> torch.Tensor: + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = layer.spec.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + qk_rope_head_dim = layer.params.qk_rope_head_dim + q_nope, q_rope = q + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) + o_tensor = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=kv_nope, + qv=q_nope, + page_table=topk_mem_indices, + cache_seqlens=context.nsa_cache_seqlens, + cu_seqlens_q=context.batch.b1_cu_q_seq_len, + cu_seqlens_k_new=context.nsa_cu_seqlens_k_new, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=softmax_scale, + causal=True, + ) + return o_tensor diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py new file mode 100644 index 0000000000..0836492c0b --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py @@ -0,0 +1,147 @@ +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention.backends.base import AttentionExecutionContext +from lightllm.common.attention.backends.sparse_mla.flashmla.backend import NsaFlashMlaSparseAttBackend + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaFp8SparsePrefillContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + q_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaFp8SparseDecodeContext(AttentionExecutionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + flashmla_sched_meta: object = None + capture_flashmla_sched_meta: object = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + + +class NsaFlashMlaFp8SparseAttBackend(NsaFlashMlaSparseAttBackend): + cache_formats = frozenset({"fp8kv_dsa"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return NsaFlashMlaFp8SparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return NsaFlashMlaFp8SparseDecodeContext(capacity, capture, workspace) + + def _forward_prefill( + self, layer: LayerAttentionBinding, context, inputs: SparseMLAPrefillInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.params.kv + return self._run_prefill(layer, context, q=q, packed_kv=k, inputs=inputs) + + def _run_prefill(self, layer, context, q: torch.Tensor, packed_kv: torch.Tensor, inputs) -> torch.Tensor: + import flash_mla + + topk_indices = inputs.topk_indices + softmax_scale = layer.spec.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + topk_mem_indices = inputs.topk_mem_indices + prefill_cache_kv = inputs.prefill_cache_kv + if context.batch.max_cache_len > 0: + from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( + get_prefill_kv_cache_and_remap_indices_triton, + ) + + kv, topk_indices = get_prefill_kv_cache_and_remap_indices_triton( + packed_kv=packed_kv, + topk_mem_indices=topk_mem_indices, + prefill_mem_index=context.batch.mem_index, + prefill_cache_kv=prefill_cache_kv, + prefill_dtype=prefill_cache_kv.dtype, + ) + else: + kv = prefill_cache_kv + if topk_indices.ndim == 2: + topk_indices = topk_indices.unsqueeze(1) + mla_out, _, _ = flash_mla.flash_mla_sparse_fwd( + q=q, kv=kv, indices=topk_indices, sm_scale=softmax_scale, d_v=kv_lora_rank + ) + return mla_out + + def _update_decode(self, context): + context.q_kv_seq_len = context.batch.b_seq_len + context.q_req_idx = context.batch.b_req_idx + context.ragged_mem_index = context.workspace.get( + "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 + ) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=context.batch.b_seq_len, + b_q_seq_len=context.batch.b_q_seq_len, + b_req_idx=context.batch.b_req_idx, + req_to_token_index=context.batch.req_to_token_indexs, + q_token_num=context.batch.b_seq_len.shape[0], + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=context.batch.hold_req_id, + ) + return + + def _forward_decode( + self, layer: LayerAttentionBinding, context, inputs: SparseMLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + return self._run_decode(layer, context, q=q, packed_kv=k, inputs=inputs) + + def _run_decode( + self, layer, context, q: Tuple[torch.Tensor, torch.Tensor], packed_kv: torch.Tensor, inputs + ) -> torch.Tensor: + import flash_mla + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = layer.spec.softmax_scale + kv_lora_rank = layer.params.kv_lora_rank + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" + q_nope, q_rope = q + q_all = torch.cat([q_nope, q_rope], dim=-1).unsqueeze(1).contiguous() + kv = torch.as_strided( + packed_kv, + size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), + stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), + ) + o_tensor, _ = flash_mla.flash_mla_with_kvcache( + q=q_all, + k_cache=kv, + block_table=None, + cache_seqlens=None, + head_dim_v=kv_lora_rank, + tile_scheduler_metadata=( + context.capture_flashmla_sched_meta + if torch.cuda.is_current_stream_capturing() + else context.flashmla_sched_meta + ), + softmax_scale=softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=topk_mem_indices, + ) + return o_tensor[:, 0, :, :] + + def _prepare_decode(self, context): + import flash_mla + + if context.flashmla_sched_meta is None: + context.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] + context.capture_flashmla_sched_meta = flash_mla.get_mla_metadata()[0] diff --git a/lightllm/common/attention/probe.py b/lightllm/common/attention/probe.py new file mode 100644 index 0000000000..f4e3e9b76a --- /dev/null +++ b/lightllm/common/attention/probe.py @@ -0,0 +1,90 @@ +"""启动子进程中的小规模数值检查，不接触模型的真实 KV 或请求状态。""" + +from types import SimpleNamespace + +import torch + +from .api import AttentionBatch, AttentionPhase, ContextCapacity, ExecutionRequirements, LayerAttentionBinding +from .api.inputs import DenseInputs, DenseParams +from .backends._workspace import AttentionWorkspace + + +def probe_dense(backend_cls, spec, runtime): + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.dist_utils import get_dp_world_size + + page = spec.cache.page_size + stride = ((16 + page - 1) // page) * page + q_heads, kv_heads, dim = spec.num_q_heads, spec.cache.key_heads, spec.head_size + model = SimpleNamespace( + args=StartArgs(page_size=page, disable_cudagraph=True), + device_runtime=runtime, + attention_requirements=ExecutionRequirements(), + config=dict( + num_attention_heads=q_heads * get_dp_world_size(), + num_key_value_heads=kv_heads * get_dp_world_size(), + hidden_size=q_heads * get_dp_world_size() * dim, + head_dim=dim, + ), + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=2), + max_seq_length=stride, + data_type=spec.dtype, + ) + backend = backend_cls.create(model, spec) + device = runtime.device + lengths = (7, 13) + q_lengths = (2, 3) if spec.phase is AttentionPhase.PREFILL else (1, 1) + query = torch.randn((sum(q_lengths), q_heads, dim), dtype=spec.dtype, device=device) + key, value = (torch.randn((3 * stride, kv_heads, dim), dtype=spec.dtype, device=device) for _ in range(2)) + batch = AttentionBatch( + phase=spec.phase, + batch_size=2, + token_num=sum(q_lengths), + total_token_num=sum(lengths), + max_q_seq_len=max(q_lengths), + max_kv_seq_len=max(lengths), + max_cache_len=max(n - q for n, q in zip(lengths, q_lengths)), + b_req_idx=torch.tensor([0, 1], dtype=torch.int32, device=device), + b_seq_len=torch.tensor(lengths, dtype=torch.int32, device=device), + b_q_seq_len=torch.tensor(q_lengths, dtype=torch.int32, device=device), + b1_cu_q_seq_len=torch.tensor([0, q_lengths[0], sum(q_lengths)], dtype=torch.int32, device=device), + b1_cu_kv_seq_len=torch.tensor([0, lengths[0], sum(lengths)], dtype=torch.int32, device=device), + req_to_token_indexs=torch.arange(3 * stride, dtype=torch.int32, device=device).view(3, stride), + hold_req_id=2, + b_ready_cache_len=torch.tensor([n - q for n, q in zip(lengths, q_lengths)], dtype=torch.int32, device=device), + ) + capacity = ContextCapacity(spec.phase, 2, sum(q_lengths), stride, max(q_lengths)) + workspace = AttentionWorkspace(runtime=runtime) + workspace.reserve(backend.resource_requirements(capacity, capture=False)) + context = backend.allocate_context(capacity, capture=False, workspace=workspace) + backend.prepare(context, batch) + backend.update(context) + attention = LayerAttentionBinding( + group_id=0, + logical_layer_id=0, + spec=spec, + params=DenseParams(key, value), + backend=backend, + alloc_tensor=lambda shape, dtype, **kwargs: torch.empty(shape, dtype=dtype, device=device), + ) + actual = attention.backend.forward(attention, context, DenseInputs(query), out=torch.empty_like(query)) + expected, start = [], 0 + for request, (length, q_length) in enumerate(zip(lengths, q_lengths)): + q = query[start : start + q_length].float() + k = key[request * stride : request * stride + length].float().repeat_interleave(q_heads // kv_heads, dim=1) + v = value[request * stride : request * stride + length].float().repeat_interleave(q_heads // kv_heads, dim=1) + logits = torch.einsum("qhd,khd->hqk", q, k) * dim**-0.5 + if spec.causal: + mask = torch.arange(length, device=device)[None, :] <= ( + length - q_length + torch.arange(q_length, device=device)[:, None] + ) + logits.masked_fill_(~mask, -torch.inf) + expected.append(torch.einsum("hqk,khd->qhd", logits.softmax(-1), v)) + start += q_length + expected = torch.cat(expected).to(spec.dtype) + runtime.synchronize() + if not torch.allclose(actual, expected, atol=0.02, rtol=0.02): + return False, f"dense {spec.phase.value} output mismatch: max diff {(actual - expected).abs().max().item()}" + return True, None diff --git a/lightllm/common/attention/registry.py b/lightllm/common/attention/registry.py new file mode 100644 index 0000000000..ec4b6a4b66 --- /dev/null +++ b/lightllm/common/attention/registry.py @@ -0,0 +1,122 @@ +"""Lazy provider declarations and default preference order; importing this does not load kernels.""" + +import importlib + +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.platform import Platform + + +def load_backend_class(reference): + if isinstance(reference, type): + return reference + module, name = reference.split(":") + if not module.startswith("lightllm."): + module = "lightllm.common.attention." + module + return getattr(importlib.import_module(module), name) + + +varlen_attention_ops = { + "fa3": "backends.dense.fa3.varlen:Fa3VarlenAttentionOp", + "xformers": "backends.dense.xformers.varlen:XformersVarlenAttentionOp", + "sdpa": "backends.dense.sdpa.varlen:SdpaVarlenAttentionOp", + "triton": "backends.dense.triton.varlen:TritonVarlenAttentionOp", +} + + +VISION_VARLEN_PRIORITY = ("fa3", "xformers", "sdpa", "triton") +AUDIO_VARLEN_PRIORITY = ("fa3", "triton") + + +def get_varlen_attention_op_class(name): + if name not in varlen_attention_ops: + raise ValueError(f"unknown varlen attention operator: {name}") + return load_backend_class(varlen_attention_ops[name]) + + +cache_format_to_backend = { + "None": { + "triton": "backends.dense.triton.backend:TritonAttBackend", + "fa3": "backends.dense.fa3.backend:Fa3AttBackend", + "flashinfer": "backends.dense.flashinfer.backend:FlashInferAttBackend", + }, + "int4kv": { + "triton": "backends.dense.triton.int4:Int4kvTritonAttBackend", + }, + "int8kv": { + "triton": "backends.dense.triton.int8:Int8kvTritonAttBackend", + }, + "fp8kv_sph": { + "fa3": "backends.dense.fa3.fp8:Fp8Fa3AttBackend", + }, + "fp8kv_spt": { + "flashinfer": "backends.dense.flashinfer.fp8:Fp8FlashInferAttBackend", + }, +} + +mla_cache_format_to_backend = { + "None": { + "triton": "backends.mla.triton.backend:MlaTritonAttBackend", + "fa3": "backends.mla.fa3.backend:MlaFa3AttBackend", + "flashinfer": "backends.mla.flashinfer.backend:MlaFlashInferAttBackend", + }, +} + +nsa_cache_format_to_backend = { + "None": { + "flashmla_sparse": "backends.sparse_mla.flashmla.backend:NsaFlashMlaSparseAttBackend", + }, + "fp8kv_dsa": { + "flashmla_sparse": "backends.sparse_mla.flashmla.fp8:NsaFlashMlaFp8SparseAttBackend", + }, +} + +neo_cache_format_to_backend = { + "None": { + "triton": "backends.dense.triton.neo:NeoTritonAttBackend", + "fa3": "backends.dense.fa3.neo:NeoFa3AttBackend", + } +} + + +def backend_candidates(family, platform, ssm_state_mode): + tables = { + AttentionFamily.DENSE: cache_format_to_backend, + AttentionFamily.MLA: mla_cache_format_to_backend, + AttentionFamily.SPARSE_MLA: nsa_cache_format_to_backend, + AttentionFamily.GDN: { + "gdn": { + "flashqla": "backends.gdn.flashqla.backend:FlashQlaGDNBackend", + "triton": "backends.gdn.triton.backend:TritonGDNBackend", + } + }, + AttentionFamily.KDA: {"kda": {"triton": "backends.kda.triton.backend:KDALinearAttBackend"}}, + } + if ssm_state_mode == "replay": + tables[AttentionFamily.GDN] = { + "gdn": { + "flashqla": "backends.gdn.replay.backend:ReplayGDNFlashQlaAttBackend", + "triton": "backends.gdn.replay.backend:ReplayGDNTritonAttBackend", + } + } + tables[AttentionFamily.KDA] = { + "kda": { + "triton": "backends.kda.replay.backend:ReplayKDALinearAttBackend", + } + } + if platform is not Platform.CUDA: + tables = {} + return tables.get(family, {}) + + +def backend_priority(family, phase, specs): + if family is AttentionFamily.GDN: + priority = ["flashqla", "triton"] if phase is AttentionPhase.PREFILL else ["triton"] + elif family is AttentionFamily.KDA: + priority = ["triton"] + elif family is AttentionFamily.SPARSE_MLA: + priority = ["flashmla_sparse"] + else: + priority = ["fa3", "flashinfer", "triton"] + if phase is AttentionPhase.DECODE and not any(spec.use_mtp for spec in specs): + priority = ["flashinfer", "fa3", "triton"] + return priority diff --git a/lightllm/common/attention/selection.py b/lightllm/common/attention/selection.py new file mode 100644 index 0000000000..e3a2144616 --- /dev/null +++ b/lightllm/common/attention/selection.py @@ -0,0 +1,145 @@ +import torch +import torch.distributed as dist + +from lightllm.common.attention.api.backend import UnsupportedAttention +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.common.attention.registry import ( + VISION_VARLEN_PRIORITY, + backend_candidates, + backend_priority, + load_backend_class, + varlen_attention_ops, +) +from lightllm.platform import get_device_runtime +from lightllm.utils import backend_validator +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +class BackendSelectionError(RuntimeError): + def __init__(self, family, phase, failures): + self.failures = failures + details = "; ".join(f"{name}: {reason}" for name, reason in failures.items()) + super().__init__(f"No usable {family.value} {phase.value} attention backend: {details}") + + +def select_varlen_attention_op(requested="auto", *, device=None, devices=None, dtype=None, head_dim=80): + if isinstance(dtype, str): + aliases = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"} + dtype = getattr(torch, aliases.get(dtype, dtype)) + if devices is None: + # Resolve once before probing; probes receive concrete device ordinals. + devices = (get_device_runtime(device).device,) + else: + devices = tuple(get_device_runtime(d).device for d in devices) + if not devices: + raise ValueError("varlen selection requires at least one device") + ordered = VISION_VARLEN_PRIORITY if requested == "auto" else (requested,) + failures = {} + for name in ordered: + if name not in varlen_attention_ops: + failures[name] = "unknown varlen attention operator" + continue + # 在子进程中检查实际 Q/K/V 路径，避免父进程初始化 CUDA 或误用文本 attention 探针。 + for assigned_device in devices: + result = backend_validator.probe_varlen(name, device=assigned_device, dtype=dtype, head_dim=head_dim) + if not result.ok: + failures[name] = f"{result.scope.value} on {assigned_device}: {result.reason}" + break + if name in failures: + continue + logger.info("Selected %s varlen attention after %s", name, result.scope.value) + return name + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, failures) + + +def select_attention_backend( + *, + specs, + caps, + execution, + requested="auto", + candidates=None, + priority=None, + ssm_state_mode="native", +): + specs = tuple(specs) + if not specs: + raise ValueError("attention selection requires at least one spec") + family, phase, encoding = specs[0].family, specs[0].phase, specs[0].cache.format + if any((spec.family, spec.phase, spec.cache.format) != (family, phase, encoding) for spec in specs): + raise ValueError("attention specs must share family, phase and cache format") + if candidates is None: + candidates = backend_candidates(family, caps.platform, ssm_state_mode) + available = candidates.get(encoding, {}) + if priority is None: + priority = backend_priority(family, phase, specs) + ordered = [requested] if requested != "auto" else list(priority) + failures = {} + selected_backend = None + selected_name = None + for name in ordered: + if name not in available: + failures[name] = f"no implementation for cache encoding {encoding}" + continue + probe_scopes = set() + try: + backend = load_backend_class(available[name]) + backend.check_environment(caps) + for spec in specs: + backend.check_support(spec, caps, execution) + for spec in specs: + result = backend_validator.probe_backend( + available[name], device=str(caps.device), dtype=spec.dtype, spec=spec + ) + probe_scopes.add(result.scope.value) + if not result.ok: + failures[name] = f"{result.scope.value}: {result.reason}" + break + if name in failures: + continue + except UnsupportedAttention as exc: + failures[name] = str(exc) + continue + except Exception as exc: + # All ranks must reach consensus, including a rank whose provider check failed. + logger.exception( + "Unexpected %s attention provider failure: family=%s phase=%s device=%s specs=%r", + name, + family.value, + phase.value, + caps.device, + specs, + ) + failures[name] = f"provider check failed: {type(exc).__name__}: {exc}" + continue + logger.info( + "Selected %s %s attention on %s: %s (%s)", + family.value, + phase.value, + caps.device, + name, + "; ".join(sorted(probe_scopes)), + ) + selected_backend = backend + selected_name = name + break + + error = None if selected_backend is not None else BackendSelectionError(family, phase, failures) + _check_rank_selection(family, phase, selected_name, None if error is None else str(error)) + if error is not None: + raise error + return selected_backend + + +def _check_rank_selection(family, phase, selected_name, failure): + if dist.is_initialized() and dist.get_world_size() > 1: + selections = [None] * dist.get_world_size() + dist.all_gather_object(selections, {"backend": selected_name, "failure": failure}) + selected = {item["backend"] for item in selections if item["backend"] is not None} + rejected = {f"rank {rank}": item["failure"] for rank, item in enumerate(selections) if item["failure"]} + if len(selected) > 1: + rejected["backend selection"] = ", ".join(sorted(selected)) + if rejected: + raise BackendSelectionError(family, phase, rejected) diff --git a/lightllm/common/attention/varlen.py b/lightllm/common/attention/varlen.py new file mode 100644 index 0000000000..fded4439c9 --- /dev/null +++ b/lightllm/common/attention/varlen.py @@ -0,0 +1,36 @@ +from lightllm.platform import get_device_runtime +from lightllm.utils.log_utils import init_logger + +from .api import AttentionFamily, AttentionPhase, UnsupportedAttention +from .registry import AUDIO_VARLEN_PRIORITY, get_varlen_attention_op_class +from .selection import BackendSelectionError + +logger = init_logger(__name__) + + +class AutoVarlenAttentionRunner: + """音频编码器的自动路径；失败的实现只在当前调用方禁用，并在本次调用中重算输出。""" + + def __init__(self, providers=AUDIO_VARLEN_PRIORITY): + self.providers = providers + self.ops = {} + self.failures = {} + + def forward(self, inputs, *, out): + for name in self.providers: + if name in self.failures: + continue + try: + if name not in self.ops: + op_cls = get_varlen_attention_op_class(name) + caps = get_device_runtime(inputs.query.device).caps + op_cls.check_environment(caps) + op_cls.check_support(inputs, caps) + self.ops[name] = op_cls() + return self.ops[name].forward(inputs, out=out) + except UnsupportedAttention as exc: + self.failures[name] = str(exc) + except Exception as exc: + self.failures[name] = str(exc) + logger.warning("Disabling %s varlen attention for this encoder: %s", name, exc) + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, self.failures) diff --git a/lightllm/common/basemodel/attention_vit/create_utils.py b/lightllm/common/basemodel/attention_vit/create_utils.py index 21eb9a7af9..9c436d42b3 100644 --- a/lightllm/common/basemodel/attention_vit/create_utils.py +++ b/lightllm/common/basemodel/attention_vit/create_utils.py @@ -1,9 +1,9 @@ from importlib import import_module -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.backend_validator import _validate from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend +from lightllm.utils.backend_validator import validate +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -34,12 +34,12 @@ def _select_vit_backend(priority_list: list = ["fa3", "xformers", "sdpa", "trito """Auto-select the best available backend with validation for VIT. Priority: FA3 > Xformers > Sdpa > Triton - Each backend is validated in a subprocess with ground truth checks. + This legacy entry checks provider smoke only; the Varlen numerical path is wired in PR4. """ for backend_name in priority_list: - if _validate(backend_name): - logger.info(f"Auto-selected {backend_name} backend (validated) for VIT") + if validate(backend_name): + logger.info(f"Auto-selected {backend_name} backend (provider smoke checked) for VIT") return backend_name # Fallback to triton without validation (should not happen) diff --git a/lightllm/common/basemodel/graph.py b/lightllm/common/basemodel/graph.py new file mode 100644 index 0000000000..dd03b44242 --- /dev/null +++ b/lightllm/common/basemodel/graph.py @@ -0,0 +1,72 @@ +"""持有一次捕获的输入、provider 任务和完成事件；bucket 策略仍由模型决定。""" + +from lightllm.platform import get_device_runtime + + +class GraphExecution: + def __init__(self, graph, states=(), outputs=(), *, runtime=None): + self.graph = graph + self.states = tuple(states) + self.outputs = tuple(outputs) + self.runtime = runtime if runtime is not None else get_device_runtime(self.states[0].input_ids.device) + self.bindings = [] + self.workspaces = [] + self.generations = [] + try: + for state in self.states: + manager = state.att_state_manager + self.generations.append((manager, manager.storage_generation)) + for group_id, context in state.attention_contexts.items(): + backend = state.attention_groups[group_id].backend_for(context.capacity.phase) + binding = backend.bind_graph(context) + if binding is not None: + self.bindings.append(binding) + if context.workspace not in self.workspaces: + self.workspaces.append(context.workspace) + manager.register_graph(self) + for workspace in self.workspaces: + workspace.freeze() + except BaseException: + self.close() + raise + + def replay(self, *states): + if self.graph is None: + raise RuntimeError("Graph execution is closed") + if states and len(states) != len(self.states): + raise ValueError("Graph execution lane count changed") + for manager, generation in self.generations: + if manager.storage_generation != generation: + raise RuntimeError("attention storage changed; rebind layers and recapture Graph") + for binding in self.bindings: + binding.wait_for_update() + if not states: + for workspace in self.workspaces: + workspace.begin() + for target, source in zip(self.states, states): + target.copy_for_cuda_graph(source) + for binding in self.bindings: + binding.before_replay() + self.graph.replay() + # 必须先提交辅助更新，再记录完成事件：Graph 本身可能正在等待更新流。 + events = [] + for binding in self.bindings: + events.extend(binding.after_submit()) + for workspace in self.workspaces: + workspace.finish(dependencies=events) + return self.outputs[0] if len(self.outputs) == 1 else self.outputs + + def close(self): + if self.graph is None: + return + for binding in self.bindings: + binding.wait_for_update() + self.runtime.synchronize() + for binding in self.bindings: + binding.close() + self.graph.reset() + self.graph = None + self.bindings.clear() + self.workspaces.clear() + self.states = self.outputs = () + self.generations.clear() diff --git a/lightllm/models/deepseek_v4/attention/__init__.py b/lightllm/models/deepseek_v4/attention/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/deepseek_v4/attention/sparse_mla.py b/lightllm/models/deepseek_v4/attention/sparse_mla.py new file mode 100644 index 0000000000..935d95b49f --- /dev/null +++ b/lightllm/models/deepseek_v4/attention/sparse_mla.py @@ -0,0 +1,208 @@ +import inspect +from dataclasses import dataclass + +import torch +import triton + +from lightllm.common.attention.api import AttentionFamily, UnsupportedAttention +from lightllm.common.attention.backends.base import AttentionExecutionContext, BaseAttentionBackend + + +@dataclass(frozen=True) +class DeepseekV4Params: + swa: torch.Tensor + prefill_swa: torch.Tensor + compressed: torch.Tensor | None + compress_ratio: int + head_dim_v: int + softmax_scale: float + attn_sink: torch.Tensor + + +@dataclass(frozen=True) +class DeepseekV4Inputs: + query: torch.Tensor + swa_indices: torch.Tensor + swa_lengths: torch.Tensor + extra_indices: torch.Tensor | None = None + extra_lengths: torch.Tensor | None = None + + +@dataclass +class SwaPrefillWorkspace: + swa_buffer: torch.Tensor + HOLD_SWA_SLOT: int + prefill_token_ranges: torch.Tensor + swa_prefill_workspace_offsets: torch.Tensor + + +def prepare_swa_workspace(swa_state, buffer, capacity, batch, slot): + from lightllm.models.deepseek_v4.triton_kernel.swa_cache import prepare_swa_prefill_workspace + + metadata = SwaPrefillWorkspace( + buffer, + buffer.shape[1] - swa_state.config.page_token_size, + torch.zeros((swa_state.request_capacity, 2), dtype=torch.int32, device=buffer.device), + torch.full((swa_state.request_capacity,), -1, dtype=torch.int32, device=buffer.device), + ) + prepare_swa_prefill_workspace( + swa_state=swa_state, + prefill_state=metadata, + b_req_idx=batch.b_req_idx, + b_ready_cache_len=batch.b_ready_cache_len, + b_seq_len=batch.b_seq_len, + workspace_base=slot * capacity, + ) + return metadata + + +@dataclass(slots=True) +class DeepseekV4Context(AttentionExecutionContext): + sched_meta_by_ratio: dict | None = None + capture_sched_meta_by_ratio: dict | None = None + + +class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttentionBackend): + supports_attention_sink = True + supports_noncausal_decode = True + family = AttentionFamily.SPARSE_MLA + provider = "flashmla_sparse" + dependency = "flash_mla" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + cache_formats = frozenset({"fp8kv_dsa"}) + + @classmethod + def check_environment(cls, caps) -> None: + super().check_environment(caps) + import flash_mla + + parameters = inspect.signature(flash_mla.flash_mla_with_kvcache).parameters + required = {"out", "attn_sink", "indices", "topk_length", "extra_k_cache", "extra_topk_length"} + missing = required.difference(parameters) + if missing: + raise UnsupportedAttention( + "DSV4 requires the FlashMLA version pinned in docker/Dockerfile; missing: " + ", ".join(sorted(missing)) + ) + + def __init__(self, model, spec): + super().__init__(model=model, spec=spec) + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=model.config, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ) + self.tp_q_head_num = model.config["num_attention_heads"] // model.tp_world_size_ + # Sparse FlashMLA uses 64-head tiles; TP can leave fewer local Q heads. + self.padded_q_head_num = ((self.tp_q_head_num + 63) // 64) * 64 + self.compress_ratios = tuple(dict.fromkeys(config.compress_ratios)) + self.layer_compress_ratios = dict(zip(config.logical_layer_ids, config.compress_ratios)) + self.swa_config = config.create_swa_config() + self.compressed_configs = { + ratio: config.create_compressed_config(compress_ratio=ratio) for ratio in self.compress_ratios if ratio + } + + @classmethod + def get_swa_prefill_workspace(cls, model): + """由 backend 复用计算临时区；模型在 KV 容量评估前调用，以计入实际显存占用。""" + if hasattr(model, "swa_prefill_workspace"): + return model.swa_prefill_workspace + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=model.config, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ).create_swa_config() + page_token_num = config.page_token_size + max_forward_tokens = max(model.batch_max_tokens, model.max_req_num) + 2 * model.tp_world_size_ + capacity = ( + triton.cdiv( + max_forward_tokens + model.max_req_num * (config.window_size + 2 * page_token_num), page_token_num + ) + * page_token_num + ) + if model.args.run_mode == "decode": + # D 节点初始化仍会对 HOLD 请求执行一次 prefill；只保留末尾的 HOLD 页， + # 不为普通请求预留长 prefill 临时区，实际请求的前缀由 P 节点传入。 + capacity = 0 + regions = 1 + int(model.args.enable_prefill_microbatch_overlap) + # 每个 microbatch 使用独立临时区；最后一页留给 HOLD 请求。 + token_num = regions * capacity + page_token_num + # 物理行宽由 SWA 原子配置给出，backend 仅决定 workspace 的 token 容量。 + slot_bytes = config.get_packed_kv_bytes(page_token_num) // page_token_num + layer_count = len(config.logical_layer_ids) + buffer = torch.empty(layer_count * token_num * slot_bytes, dtype=torch.uint8, device="cuda") + model.swa_prefill_workspace = buffer.view(layer_count, token_num, slot_bytes), capacity + return model.swa_prefill_workspace + + def _allocate_prefill_context(self, capacity, capture, workspace): + return DeepseekV4Context(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return DeepseekV4Context(capacity, capture, workspace) + + def _prepare_prefill(self, context): + self._prepare_decode(context) + + def _prepare_decode(self, context): + import flash_mla as flashmla + + if context.sched_meta_by_ratio is None: + context.sched_meta_by_ratio = {ratio: flashmla.get_mla_metadata()[0] for ratio in self.compress_ratios} + context.capture_sched_meta_by_ratio = { + ratio: flashmla.get_mla_metadata()[0] for ratio in self.compress_ratios + } + + def _run(self, layer, context, inputs, swa, out): + import flash_mla as flashmla + + params = layer.params + compressed = None + if params.compressed is not None: + compressed = self.compressed_configs[params.compress_ratio].get_attention_view(packed_kv=params.compressed) + scheduler = ( + context.capture_sched_meta_by_ratio + if torch.cuda.is_current_stream_capturing() + else context.sched_meta_by_ratio + ) + q = inputs.query + if out is None: + out = layer.alloc_tensor((q.shape[0], q.shape[1], params.head_dim_v), dtype=q.dtype, device=q.device) + padded_out, _ = flashmla.flash_mla_with_kvcache( + q=q.unsqueeze(1), + k_cache=self.swa_config.get_attention_view(packed_kv=swa), + block_table=None, + cache_seqlens=None, + head_dim_v=params.head_dim_v, + tile_scheduler_metadata=scheduler[params.compress_ratio], + num_splits=None, + softmax_scale=params.softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=inputs.swa_indices, + attn_sink=params.attn_sink, + topk_length=inputs.swa_lengths, + extra_k_cache=compressed, + extra_indices_in_kvcache=inputs.extra_indices, + extra_topk_length=inputs.extra_lengths, + out=out.unsqueeze(1), + ) + return padded_out[:, 0, : self.tp_q_head_num] + + def _forward_prefill(self, layer, context, inputs: DeepseekV4Inputs, *, out=None): + return self._run(layer, context, inputs, layer.params.prefill_swa, out) + + def _forward_decode(self, layer, context, inputs: DeepseekV4Inputs, *, out=None): + return self._run(layer, context, inputs, layer.params.swa, out) + + +DSV4_NSA_BACKENDS = { + "fp8kv_dsa": { + "flashmla_sparse": "lightllm.models.deepseek_v4.attention.sparse_mla:DeepseekV4FlashMlaFp8SparseAttBackend", + } +} diff --git a/lightllm/models/glm5_next/attention/__init__.py b/lightllm/models/glm5_next/attention/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next/attention/fp8_sparse_mla.py b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py new file mode 100644 index 0000000000..0a458f84a7 --- /dev/null +++ b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py @@ -0,0 +1,224 @@ +import dataclasses + +import torch + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.execution import BackendResources, WorkspaceRequirement +from lightllm.common.attention.api.inputs import SparseMLAPrefillInputs +from lightllm.common.attention.backends.sparse_mla.flashmla.fp8 import ( + NsaFlashMlaFp8SparseAttBackend, + NsaFlashMlaFp8SparseDecodeContext, +) + +from .sparse_mla import Glm5NextSparsePrefillContext + + +@dataclasses.dataclass(slots=True) +class Fp8Glm5NextSparsePrefillContext(Glm5NextSparsePrefillContext): + smallq_workspace: object = None + generic_workspace: object = None + generic_binding: object = None + + +@dataclasses.dataclass(slots=True) +class Fp8Glm5NextSparseDecodeContext(NsaFlashMlaFp8SparseDecodeContext): + decode_topk_bound: int = None + + +class Fp8Glm5NextSparseAttBackend(NsaFlashMlaFp8SparseAttBackend): + provider = "flashmla_sparse" + cache_formats = frozenset({"fp8kv_dsa"}) + + def __init__(self, model, spec): + super().__init__(model, spec) + self.prefill_graph_enabled = model.args.enable_prefill_cudagraph + self.prefill_token_capacity = model.batch_max_tokens or 0 + self.decode_only = model.args.run_mode == "decode" + + def resource_requirements(self, capacity, *, capture): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import KV_GATHER_CHUNK_SIZE + + if capacity.phase is not AttentionPhase.PREFILL: + return super().resource_requirements(capacity, capture=capture) + tokens = max(1, self.prefill_token_capacity, capacity.token_num) + kv_tokens = tokens if self.decode_only else capacity.batch_size * capacity.max_kv_seq_len + tokens + gather_tokens = min(KV_GATHER_CHUNK_SIZE, kv_tokens) + topk = (self.network_config["index_topk"] + 3 + 127) // 128 * 128 + return BackendResources( + workspace=( + WorkspaceRequirement("glm_gathered_kv", gather_tokens * 512, torch.bfloat16), + WorkspaceRequirement("glm_chunk_indices", tokens * topk, torch.int32), + WorkspaceRequirement("glm_chunk_lengths", tokens, torch.int32), + WorkspaceRequirement("glm_cu_q", tokens + 1, torch.int32), + ) + ) + + def get_smallq_workspace(self, context, q): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + allocate_direct_smallq_workspace, + allocate_selected_prefill_workspace, + ) + + context.smallq_workspace = context.workspace.get_object( + "glm_smallq", + lambda: { + "direct": allocate_direct_smallq_workspace(q[:1]), + "selected": allocate_selected_prefill_workspace(q, batch_capacity=8), + }, + ) + if q.shape[0] == 1: + return context.smallq_workspace["direct"] + batch = q.shape[0] + selected = context.smallq_workspace["selected"] + slots = selected["compact"].shape[1] + return { + "compact": selected["compact"][:batch], + "lengths": selected["lengths"][:batch], + "page_table": selected["page_table"][:batch], + "cu_q": selected["cu_q"][: batch + 1], + "cu_k": selected["cu_k"][: batch + 1], + "selected_kv": selected["selected_kv"][: batch * slots], + } + + def get_generic_prefill_workspace(self, context, q, packed_kv, topk_indices): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import KV_GATHER_CHUNK_SIZE + + rows, columns = topk_indices.shape + workspace = context.workspace + tokens = max(self.prefill_token_capacity, rows) + binding = (rows, columns, topk_indices.dtype, packed_kv.shape[0]) + if context.generic_workspace is None or context.generic_binding != binding: + row_table = workspace.get("glm_prefill_rows", packed_kv.shape[0], torch.int32)[: packed_kv.shape[0]] + row_table.fill_(-1) + cu_q = workspace.get("glm_cu_q", tokens + 1, torch.int32)[: rows + 1] + torch.arange(rows + 1, device=q.device, dtype=torch.int32, out=cu_q) + context.generic_workspace = { + "prefill_row_table": row_table, + "chunk_indices": workspace.get("glm_chunk_indices", tokens * columns, topk_indices.dtype)[ + : rows * columns + ].view(rows, columns), + "chunk_lengths": workspace.get("glm_chunk_lengths", tokens, torch.int32)[:rows], + "cu_seqlens_q": cu_q, + "row_table_built": False, + } + context.generic_binding = binding + gather_tokens = min(KV_GATHER_CHUNK_SIZE, context.ragged_mem_index.numel()) + if self.decode_only: + gather_tokens = min(gather_tokens, max(1, self.prefill_token_capacity, rows)) + context.generic_workspace["gathered_kv"] = workspace.get( + "glm_gathered_kv", max(1, gather_tokens) * 512, q.dtype + )[: gather_tokens * 512].view(gather_tokens, 1, 512) + context.generic_workspace["prefill_row_table"].fill_(-1) + context.generic_workspace["row_table_built"] = False + return context.generic_workspace + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Fp8Glm5NextSparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Fp8Glm5NextSparseDecodeContext( + capacity, + capture, + workspace, + decode_topk_bound=2051 + if self.network_config["index_topk"] == 2048 and self.network_config["index_kpool"] == 4 + else None, + ) + + def _run_prefill(self, layer, context, q, kv, inputs): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import context_attention_fwd_fp8 + + if q.shape[1] != 64: + return context_attention_fwd_fp8( + q, + kv, + context.ragged_mem_index, + context.batch.mem_index, + inputs.prefill_cache_kv, + context.ks, + inputs.topk_indices, + layer.spec.softmax_scale, + generic_workspace=self.get_generic_prefill_workspace(context, q, kv, inputs.topk_indices), + ) + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + can_use_direct_smallq, + can_use_selected_prefill, + ) + + direct_workspace = ( + self.get_smallq_workspace(context, q) + if can_use_direct_smallq( + q, kv, context.ragged_mem_index, context.batch.mem_index, inputs.prefill_cache_kv, inputs.topk_indices + ) + else None + ) + selected_eligible = direct_workspace is None and can_use_selected_prefill( + q, kv, context.ragged_mem_index, context.batch.mem_index, inputs.prefill_cache_kv, inputs.topk_indices + ) + workspace = self.get_smallq_workspace(context, q) if selected_eligible else None + generic_workspace = None + if direct_workspace is None and (not selected_eligible): + generic_workspace = self.get_generic_prefill_workspace(context, q, kv, inputs.topk_indices) + return context_attention_fwd_fp8( + q, + kv, + context.ragged_mem_index, + context.batch.mem_index, + inputs.prefill_cache_kv, + context.ks, + inputs.topk_indices, + layer.spec.softmax_scale, + workspace, + direct_workspace, + generic_workspace, + ) + + def _run_decode(self, layer, context, q, packed_kv, inputs): + import flash_mla + + from lightllm.models.glm5_next.triton_kernel import fp8_sparse_decode + from lightllm.models.glm5_next.triton_kernel.pad_query_heads import pad_query_heads + + topk_mem_indices = inputs.topk_mem_indices + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" + q_nope, _ = q + num_heads = q_nope.shape[1] + if fp8_sparse_decode.can_use_fp8_sparse_direct_decode( + q_nope, + packed_kv, + topk_mem_indices, + layer.params.kv_lora_rank, + context.decode_topk_bound, + is_cuda_graph_capture=context.capture, + ): + return fp8_sparse_decode.fp8_sparse_direct_decode( + q_nope, packed_kv, topk_mem_indices, layer.spec.softmax_scale + ) + padded_q = pad_query_heads(q_nope, is_cuda_graph_capture=context.capture) + kv_cache = torch.as_strided( + packed_kv, + size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), + stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), + ) + output, _ = flash_mla.flash_mla_with_kvcache( + q=padded_q, + k_cache=kv_cache, + block_table=None, + cache_seqlens=None, + head_dim_v=layer.params.kv_lora_rank, + tile_scheduler_metadata=( + context.capture_flashmla_sched_meta + if torch.cuda.is_current_stream_capturing() + else context.flashmla_sched_meta + ), + softmax_scale=layer.spec.softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=topk_mem_indices, + ) + return output[:, 0, :num_heads, :] + + def _forward_prefill(self, layer, context, inputs: SparseMLAPrefillInputs, *, out=None): + return self._run_prefill(layer, context, inputs.query, layer.params.kv, inputs) diff --git a/lightllm/models/glm5_next/attention/sparse_mla.py b/lightllm/models/glm5_next/attention/sparse_mla.py new file mode 100644 index 0000000000..41f61ed994 --- /dev/null +++ b/lightllm/models/glm5_next/attention/sparse_mla.py @@ -0,0 +1,94 @@ +import dataclasses +import inspect + +import torch + +from lightllm.common.attention.backends.sparse_mla.flashmla.backend import ( + NsaFlashMlaSparseAttBackend, + NsaFlashMlaSparseDecodeContext, + NsaFlashMlaSparsePrefillContext, +) + + +@dataclasses.dataclass(slots=True) +class Glm5NextSparsePrefillContext(NsaFlashMlaSparsePrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Glm5NextSparseDecodeContext(NsaFlashMlaSparseDecodeContext): + pass + + +class Glm5NextSparseAttBackend(NsaFlashMlaSparseAttBackend): + provider = "flashmla_sparse" + cache_formats = frozenset({"None"}) + + def _allocate_prefill_context(self, capacity, capture, workspace): + return Glm5NextSparsePrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return Glm5NextSparseDecodeContext(capacity, capture, workspace) + + def _run_prefill(self, layer, context, q, kv, inputs): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + tokens, heads, dim = q.shape + padded_heads = (heads + 63) // 64 * 64 + padded_q = q + if padded_heads != heads: + padded_q = q.new_zeros((tokens, padded_heads, dim)) + padded_q[:, :heads] = q + out, _, _ = flash_mla_sparse_fwd( + q=padded_q, + kv=kv, + indices=inputs.topk_mem_indices.unsqueeze(1), + sm_scale=layer.spec.softmax_scale, + d_v=layer.params.kv_lora_rank, + ) + return out[:, :heads] + + def _update_decode(self, context): + super()._update_decode(context) + pool = self.network_config["index_kpool"] + topk = self.network_config["index_topk"] + context.nsa_cache_seqlens = ( + torch.minimum(context.lengths // pool * pool, torch.full_like(context.lengths, topk)) + + context.lengths % pool + ) + context.nsa_cu_seqlens_k_new = torch.nn.functional.pad( + context.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0) + ) + + def _run_decode(self, layer, context, q, kv, inputs): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + q_nope, _ = q + kv_nope = kv.view(-1, 1, 1, layer.params.kv_lora_rank) + common_kwargs = dict( + page_table=inputs.topk_mem_indices, + cache_seqlens=context.nsa_cache_seqlens, + cu_seqlens_q=context.batch.b1_cu_q_seq_len, + cu_seqlens_k_new=context.nsa_cu_seqlens_k_new, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=layer.spec.softmax_scale, + causal=False, + ) + if "only_qv" in inspect.signature(flash_attn_with_kvcache).parameters: + return flash_attn_with_kvcache( + q=None, qv=q_nope, k_cache=None, v_cache=kv_nope, only_qv=True, **common_kwargs + ) + if q_nope.shape[-1] > 256: + page_indices = inputs.topk_mem_indices.clamp_min(0) + key = kv_nope[page_indices, 0, 0] + score = torch.einsum("bhd,bkd->bhk", q_nope.float(), key.float()) + score = score * layer.spec.softmax_scale + valid = ( + torch.arange(page_indices.shape[1], device=page_indices.device)[None, :] + < context.nsa_cache_seqlens[:, None] + ) + score = score.masked_fill(~valid[:, None, :], float("-inf")) + weights = score.softmax(dim=-1) + weights = torch.where(valid[:, None, :], weights, torch.zeros_like(weights)) + return torch.einsum("bhk,bkd->bhd", weights, key.float()).to(q_nope.dtype) + return flash_attn_with_kvcache(q=q_nope, k_cache=kv_nope, v_cache=kv_nope, **common_kwargs) diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index 3458f10131..c787346a55 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -1,37 +1,70 @@ -"""Backend validation with subprocess isolation and ground truth checks.""" +"""Isolated backend numerical probes and explicitly scoped provider smoke checks.""" import multiprocessing as mp import os +import time +import traceback +from dataclasses import dataclass +from enum import Enum +from functools import lru_cache + import torch -from lightllm.utils.log_utils import init_logger + from lightllm.utils.dist_utils import get_global_rank -from functools import lru_cache +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) _VALIDATION_TIMEOUT = 30 * 60 # 30 minutes +class ProbeScope(str, Enum): + BACKEND_NUMERICS = "backend numerics" + VARLEN_NUMERICS = "varlen numerics" + PROVIDER_SMOKE = "provider smoke only; selected backend numerics not checked" + INCOMPLETE = "probe did not complete" + + +@dataclass(frozen=True) +class ProbeResult: + ok: bool + scope: ProbeScope + reason: str | None = None + + +def _has_dense_numerical_probe(spec): + from lightllm.common.attention.api import AttentionFamily + + return spec is not None and ( + spec.family is AttentionFamily.DENSE + and spec.cache.format == "None" + and not spec.use_mtp + and not spec.use_alibi + and not spec.use_sink + and spec.sliding_window is None + and not spec.variant + ) + + def _compute_ground_truth(q, k, v, is_causal=True): """Ground truth using PyTorch SDPA.""" with torch.no_grad(): return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal) -def _validate_fa3(): +def _validate_fa3(dtype=torch.bfloat16, spec=None): """Validate FA3 with ground truth.""" - from lightllm.utils.device_utils import is_hopper from lightllm.utils.sgl_utils import flash_attn_varlen_func - if not is_hopper(): - return False, "Not a Hopper GPU" + if torch.cuda.get_device_capability()[0] < 9: + return False, "FA3 requires SM90 or newer" if flash_attn_varlen_func is None: return False, "flash_attn_varlen_func is None" - batch, heads, seq, dim = 1, 4, 8, 64 - q = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") + batch, heads, seq, dim = 1, 4 if spec is None else spec.num_q_heads, 8, 64 if spec is None else spec.head_size + q = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + k = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + v = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") expected = _compute_ground_truth(q, k, v) @@ -48,18 +81,21 @@ def _validate_fa3(): cu_seqlens_k=cu_seqlens, max_seqlen_q=seq, max_seqlen_k=seq, - softmax_scale=1.0 / (dim ** 0.5), + softmax_scale=1.0 / (dim**0.5), causal=True, ) out = out.reshape(batch, seq, heads, dim).transpose(1, 2) torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashinfer(): +def _validate_flashinfer(dtype=torch.float16, spec=None): """Validate FlashInfer with ground truth.""" capability = torch.cuda.get_device_capability() arch = f"{capability[0]}.{capability[1]}" @@ -67,10 +103,10 @@ def _validate_flashinfer(): import flashinfer - batch, heads, seq, dim = 1, 4, 8, 64 - q = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") - k = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") - v = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") + batch, heads, seq, dim = 1, 4 if spec is None else spec.num_q_heads, 8, 64 if spec is None else spec.head_size + q = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + k = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + v = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") expected = _compute_ground_truth(q, k, v) @@ -88,15 +124,18 @@ def _validate_flashinfer(): torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashqla(): +def _validate_flashqla(dtype=None, spec=None): """Validate FlashQLA against LightLLM's vendored FLA kernel.""" from flash_qla import chunk_gated_delta_rule as flashqla_chunk_gated_delta_rule - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) from lightllm.models.qwen3next.config import Qwen3NextResourceConfig @@ -107,7 +146,11 @@ def _validate_flashqla(): num_v_heads = max(1, linear_config.global_linear_v_heads // linear_config.tp_world_size) head_k_dim = linear_config.head_linear_k_dim head_v_dim = linear_config.head_linear_v_dim - qkv_dtype = linear_config.conv_state_dtype + qkv_dtype = dtype or linear_config.conv_state_dtype + if spec is not None: + num_k_heads = spec.num_q_heads + num_v_heads = max(num_k_heads, num_v_heads) + head_k_dim, head_v_dim = spec.head_size, spec.value_head_size state_dtype = linear_config.ssm_state_dtype batch, seq = 1, 64 @@ -140,12 +183,15 @@ def _validate_flashqla(): ("final state", final_state, expected_state), ): if not torch.allclose(actual, expected, rtol=1e-2, atol=1e-2): - return False, f"{name} mismatch: max diff {(actual - expected).abs().max().item():.6f}" + return ( + False, + f"{name} mismatch: max diff {(actual - expected).abs().max().item():.6f}", + ) return True, None -def _validate_triton(): - """Validate Triton with softmax ground truth.""" +def _smoke_triton(): + """Check Triton compilation and a softmax kernel, not the selected attention backend.""" import triton import triton.language as tl @@ -169,7 +215,10 @@ def _softmax_kernel(input_ptr, output_ptr, n_cols, BLOCK: tl.constexpr): torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-3, atol=1e-3): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None @@ -208,7 +257,10 @@ def _validate_xformers(): out = out.transpose(1, 2).contiguous() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None @@ -234,13 +286,16 @@ def _validate_sdpa(): out = scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashmla_sparse(): - """Validate flashmla_sparse (NSA) with ground truth.""" +def _smoke_flashmla_sparse(dtype=torch.bfloat16, spec=None): + """Check sparse MLA import and kernel launch; no reference output comparison.""" try: from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -252,7 +307,6 @@ def _validate_flashmla_sparse(): kv_lora_rank = 512 qk_rope_head_dim = 64 qk_dim = kv_lora_rank + qk_rope_head_dim - dtype = torch.bfloat16 device = "cuda" q = torch.randn(batch * seq, heads, qk_dim, dtype=dtype, device=device) @@ -265,7 +319,7 @@ def _validate_flashmla_sparse(): topk_indices = topk_indices.view(batch * seq, 1, index_topk) - softmax_scale = 1.0 / (qk_dim ** 0.5) + softmax_scale = 1.0 / (qk_dim**0.5) try: mla_out, _, _ = flash_mla_sparse_fwd( @@ -282,79 +336,157 @@ def _validate_flashmla_sparse(): return True, None -def _run_in_subprocess(backend_name, pipe): - """Run validation in subprocess with suppressed output.""" - import sys +def _validate_varlen(backend_name, dtype=None, device=None, head_dim=80): + from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionInputs + from lightllm.common.attention.registry import get_varlen_attention_op_class + from lightllm.platform import get_device_runtime + + runtime = get_device_runtime(device or ("cuda" if torch.cuda.is_available() else "cpu")) + op_cls = get_varlen_attention_op_class(backend_name) + try: + op_cls.check_environment(runtime.caps) + except UnsupportedAttention as exc: + return False, str(exc) + device = runtime.device + dtype = dtype or (torch.float32 if device.type == "cpu" else torch.bfloat16) + q, k, v = [torch.randn(8, 2, head_dim, dtype=dtype, device=device) for _ in range(3)] + cu = torch.tensor([0, 3, 8], dtype=torch.int32, device=device) + out = torch.empty_like(q) + inputs = VarlenAttentionInputs(q, k, v, cu, 5) + try: + op_cls.check_support(inputs, runtime.caps) + except UnsupportedAttention as exc: + return False, str(exc) + op_cls().forward(inputs, out=out) + expected = torch.empty_like(q) + for start, end in ((0, 3), (3, 8)): + result = _compute_ground_truth( + q[start:end].transpose(0, 1), k[start:end].transpose(0, 1), v[start:end].transpose(0, 1), is_causal=False + ) + expected[start:end].copy_(result.transpose(0, 1)) + if not torch.allclose(out, expected, rtol=2e-2, atol=2e-2): + return False, f"varlen output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return True, None - # Redirect stdout/stderr to /dev/null to suppress verbose error messages - devnull = open(os.devnull, "w") - sys.stdout = devnull - sys.stderr = devnull +def _probe_in_process(backend_name, *, varlen=False, device=None, dtype=None, spec=None, head_dim=80): + from lightllm.common.attention.registry import load_backend_class + from lightllm.platform import get_device_runtime + + scope = ProbeScope.VARLEN_NUMERICS if varlen else ProbeScope.PROVIDER_SMOKE + if _has_dense_numerical_probe(spec) and not varlen: + scope = ProbeScope.BACKEND_NUMERICS try: - if backend_name == "fa3": - success, err = _validate_fa3() - elif backend_name == "xformers": - success, err = _validate_xformers() - elif backend_name == "sdpa": - success, err = _validate_sdpa() - elif backend_name == "flashinfer": - success, err = _validate_flashinfer() - elif backend_name == "flashqla": - success, err = _validate_flashqla() - elif backend_name == "triton": - success, err = _validate_triton() - elif backend_name == "flashmla_sparse": - success, err = _validate_flashmla_sparse() + if device is not None or spec is not None: + runtime = get_device_runtime(device) + if runtime.api is not None: + runtime.api.set_device(runtime.device) + if scope is ProbeScope.BACKEND_NUMERICS: + from lightllm.common.attention.probe import probe_dense + + ok, reason = probe_dense(load_backend_class(backend_name), spec, runtime) + elif varlen: + ok, reason = _validate_varlen(backend_name, dtype, device, head_dim) else: - success, err = False, f"Unknown backend: {backend_name}" - pipe.send((success, err)) - except Exception as e: - pipe.send((False, str(e))) - finally: - devnull.close() + provider = load_backend_class(backend_name).provider if spec is not None else backend_name + ok, reason = _smoke_provider(provider, dtype=dtype, spec=spec) + return ProbeResult(ok, scope, reason) + except Exception: + return ProbeResult(False, scope, traceback.format_exc()) + + +def _smoke_provider(provider, *, dtype=None, spec=None): + """Exercise a provider's representative kernel, not the selected family/spec implementation.""" + if provider == "fa3": + return _validate_fa3(dtype or torch.bfloat16, spec) + if provider == "xformers": + return _validate_xformers() + if provider == "sdpa": + return _validate_sdpa() + if provider == "flashinfer": + return _validate_flashinfer(dtype or torch.float16, spec) + if provider == "flashqla": + return _validate_flashqla(dtype, spec) + if provider == "triton": + return _smoke_triton() + if provider == "flashmla_sparse": + return _smoke_flashmla_sparse(dtype or torch.bfloat16, spec) + return False, f"Unknown provider: {provider}" + + +def _run_in_subprocess(backend_name, pipe, varlen=False, device=None, dtype=None, spec=None, head_dim=80): + from contextlib import redirect_stderr, redirect_stdout + + with open(os.devnull, "w") as devnull, redirect_stdout(devnull), redirect_stderr(devnull): + pipe.send( + _probe_in_process(backend_name, varlen=varlen, device=device, dtype=dtype, spec=spec, head_dim=head_dim) + ) @lru_cache(maxsize=None) -def validate(backend_name: str) -> bool: - if get_global_rank() == 0: - validate_ok = _validate(backend_name) - torch.distributed.broadcast_object_list([validate_ok], src=0) - else: - validate_ok = [None] - torch.distributed.broadcast_object_list(validate_ok, src=0) - validate_ok = validate_ok[0] - return validate_ok - - -def _validate(backend_name: str) -> bool: - """Validate backend in subprocess with ground truth check.""" +def probe_backend(backend_name: str, *, device=None, dtype=None, spec=None) -> ProbeResult: + """Cache by provider/device/dtype/full spec; the result states the actual check performed.""" + return _probe(backend_name, device=device, dtype=dtype, spec=spec) + + +@lru_cache(maxsize=None) +def probe_varlen(backend_name: str, *, device=None, dtype=None, head_dim=80) -> ProbeResult: + """Compare the selected Q/K/V operator against a sequence-wise reference.""" + return _probe(backend_name, varlen=True, device=device, dtype=dtype, head_dim=head_dim) + + +def validate(backend_name: str, *, device=None, dtype=None, spec=None) -> bool: + """Legacy selector compatibility; success does not imply family-specific numerical validation.""" + return probe_backend(backend_name, device=device, dtype=dtype, spec=spec).ok + + +def _probe(backend_name: str, *, varlen=False, device=None, dtype=None, spec=None, head_dim=80) -> ProbeResult: + """Isolate execution and retain both its coverage and failure reason.""" try: ctx = mp.get_context("spawn") parent, child = ctx.Pipe(duplex=False) - logger.info(f"Validating {backend_name} backend start, please wait ...") - proc = ctx.Process(target=_run_in_subprocess, args=(backend_name, child)) + logger.info("Probing %s on device=%s dtype=%s", backend_name, device, dtype) + proc = ctx.Process(target=_run_in_subprocess, args=(backend_name, child, varlen, device, dtype, spec, head_dim)) proc.start() - proc.join(timeout=_VALIDATION_TIMEOUT) + child.close() + deadline = time.monotonic() + _VALIDATION_TIMEOUT + # Drain the pipe before joining: a full traceback can exceed its buffer. + result = None + if parent.poll(_VALIDATION_TIMEOUT): + try: + result = parent.recv() + except EOFError: + pass + proc.join(timeout=max(0, deadline - time.monotonic())) if proc.is_alive(): proc.kill() proc.join() - logger.warning(f"{backend_name} validation timed out") - return False + return ProbeResult(False, ProbeScope.INCOMPLETE, "probe timed out") if proc.exitcode != 0: - logger.warning(f"{backend_name} validation crashed (exit {proc.exitcode})") - return False - - if parent.poll(): - success, err = parent.recv() - if success: - logger.info(f"{backend_name} validated") - return True - logger.warning(f"{backend_name} validation failed: {str(err)[0:88]}") - return False - - except Exception as e: - logger.warning(f"{backend_name} validation error: {e}") - return False + return ProbeResult(False, ProbeScope.INCOMPLETE, f"probe crashed (exit {proc.exitcode})") + + if result is not None: + if result.ok: + logger.info("%s passed: %s", backend_name, result.scope.value) + else: + logger.warning( + "%s failed (%s) on device=%s dtype=%s spec=%r:\n%s", + backend_name, + result.scope.value, + device, + dtype, + spec, + result.reason, + ) + return result + return ProbeResult(False, ProbeScope.INCOMPLETE, "probe returned no result") + except Exception: + logger.exception("%s probe error on device=%s dtype=%s spec=%r", backend_name, device, dtype, spec) + return ProbeResult(False, ProbeScope.INCOMPLETE, traceback.format_exc()) + finally: + if "parent" in locals(): + parent.close() + if "child" in locals(): + child.close() diff --git a/unit_tests/common/attention/helpers.py b/unit_tests/common/attention/helpers.py index ff5c950c5f..acb72bfa56 100644 --- a/unit_tests/common/attention/helpers.py +++ b/unit_tests/common/attention/helpers.py @@ -67,6 +67,7 @@ def make_spec(family=AttentionFamily.DENSE, phase=AttentionPhase.DECODE, **kwarg def make_backend(cls=BaseAttentionBackend, **kwargs): backend = cls.__new__(cls) values = dict( + spec=make_spec(), page_size=1, causal=True, dynamic_verify=False, @@ -77,6 +78,8 @@ def make_backend(cls=BaseAttentionBackend, **kwargs): values.update(kwargs) if "infer_page_size" in values: values["page_size"] = values.pop("infer_page_size") + if "softmax_scale" in values: + values["spec"] = replace(values["spec"], softmax_scale=values.pop("softmax_scale")) vars(backend).update(values) return backend diff --git a/unit_tests/common/attention/test_contracts.py b/unit_tests/common/attention/test_contracts.py index 8a833e8da5..ee97de6458 100644 --- a/unit_tests/common/attention/test_contracts.py +++ b/unit_tests/common/attention/test_contracts.py @@ -4,11 +4,24 @@ import pytest import torch +from lightllm.common.attention import registry, selection from lightllm.common.attention.api import AttentionFamily, AttentionPhase, UnsupportedAttention from lightllm.common.attention.backends.base import BaseAttentionBackend +from lightllm.utils.backend_validator import ProbeResult, ProbeScope from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub +def _select(*, model, specs=None, **kwargs): + family, phase = kwargs.pop("family"), kwargs.pop("phase") + return selection.select_attention_backend( + specs=specs or (make_spec(family, phase),), + caps=model.device_runtime.caps, + execution=model.attention_requirements, + ssm_state_mode=model.args.ssm_state_mode, + **kwargs, + ) + + def test_graph_copy_validates_all_inputs_before_writing(): captured = make_batch() incoming = make_batch((9, 7)) @@ -55,6 +68,179 @@ def test_layer_attention_rejects_wrong_phase(): attention.backend.forward(attention, context, inputs=None) +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_selection_never_falls_back_to_a_rejected_backend(monkeypatch, requested): + spec = make_spec(cache=replace(make_spec().cache, format="fp8kv_sph")) + monkeypatch.setattr( + BaseAttentionBackend, + "check_environment", + classmethod(lambda cls, caps: None), + ) + with pytest.raises(selection.BackendSelectionError, match="encoding fp8kv_sph"): + _select( + family=spec.family, + phase=spec.phase, + model=model_stub(), + specs=(spec,), + requested=requested, + candidates={"None": {"triton": BaseAttentionBackend}}, + priority=["triton"], + ) + + +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_environment_checks_apply_to_explicit_and_auto(monkeypatch, requested): + def reject(cls, caps): + raise UnsupportedAttention("missing extension") + + monkeypatch.setattr(BaseAttentionBackend, "check_environment", classmethod(reject)) + with pytest.raises(selection.BackendSelectionError, match="missing extension"): + _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested=requested, + candidates={"None": {"triton": BaseAttentionBackend}}, + priority=["triton"], + ) + + +def _backend_check_exception_candidates(monkeypatch, error_type=RuntimeError): + class BrokenBackend(BaseAttentionBackend): + provider = "broken" + + @classmethod + def check_environment(cls, caps): + raise error_type("provider ABI mismatch") + + class WorkingBackend(BaseAttentionBackend): + provider = "working" + + @classmethod + def check_environment(cls, caps): + return None + + monkeypatch.setattr( + "lightllm.utils.backend_validator.probe_backend", + lambda *args, **kwargs: ProbeResult(True, ProbeScope.BACKEND_NUMERICS), + ) + return {"None": {"broken": BrokenBackend, "working": WorkingBackend}} + + +@pytest.mark.parametrize("error_type", [UnsupportedAttention, ValueError, RuntimeError]) +def test_backend_check_exception_allows_auto_fallback(monkeypatch, error_type): + candidates = _backend_check_exception_candidates(monkeypatch, error_type) + selected = _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested="auto", + candidates=candidates, + priority=["broken", "working"], + ) + assert selected.provider == "working" + + +@pytest.mark.parametrize("error_type", [UnsupportedAttention, ValueError, RuntimeError]) +def test_backend_check_exception_is_reported_for_explicit_selection(monkeypatch, error_type): + candidates = _backend_check_exception_candidates(monkeypatch, error_type) + with pytest.raises(selection.BackendSelectionError) as error: + _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested="broken", + candidates=candidates, + priority=["broken", "working"], + ) + + reason = "provider ABI mismatch" + if error_type is not UnsupportedAttention: + reason = f"provider check failed: {error_type.__name__}: {reason}" + assert error.value.failures == {"broken": reason} + + +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_rejection_reaches_rank_consensus(monkeypatch, requested): + spec = make_spec(dtype=torch.float64) + ranks = [] + monkeypatch.setattr(BaseAttentionBackend, "check_environment", classmethod(lambda cls, caps: None)) + monkeypatch.setattr( + selection.backend_validator, "probe_backend", lambda *args, **kwargs: pytest.fail("probed rejected inputs") + ) + monkeypatch.setattr(selection, "_check_rank_selection", lambda *args: ranks.append(args)) + with pytest.raises(selection.BackendSelectionError, match="activation dtype"): + _select( + family=spec.family, + phase=spec.phase, + model=model_stub(), + specs=(spec,), + requested=requested, + candidates={"None": {"triton": BaseAttentionBackend}}, + priority=["triton"], + ) + assert len(ranks) == 1 + assert ranks[0][:3] == (spec.family, spec.phase, None) + assert "activation dtype" in ranks[0][3] + + +def test_rank_selection_checks_all_ranks_after_local_selection(monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + + def gather(output, value): + output[:] = [value, {"backend": "triton", "failure": None}] + + monkeypatch.setattr(dist, "all_gather_object", gather) + with pytest.raises(selection.BackendSelectionError, match="backend selection"): + selection._check_rank_selection(AttentionFamily.DENSE, AttentionPhase.DECODE, "fa3", None) + + +def test_rank_selection_propagates_remote_failure(monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + monkeypatch.setattr( + dist, + "all_gather_object", + lambda output, value: output.__setitem__(slice(None), [value, {"backend": None, "failure": "probe failed"}]), + ) + with pytest.raises(selection.BackendSelectionError, match="rank 1: probe failed"): + selection._check_rank_selection(AttentionFamily.DENSE, AttentionPhase.DECODE, "fa3", None) + + +@pytest.mark.parametrize("format,provider", [("fp8kv_sph", "flashinfer"), ("fp8kv_spt", "fa3")]) +def test_fp8_scale_layouts_are_not_interchangeable(format, provider): + candidates = registry.cache_format_to_backend[format] + assert provider not in candidates + + +def test_layer_shape_configures_group_planner(): + from lightllm.common.attention.backends.dense.flashinfer.backend import FlashInferAttBackend + + model = model_stub(page_size=16) + model.config = {"num_attention_heads": 16, "num_key_value_heads": 4, "hidden_size": 2048} + model.max_seq_length = 128 + model.data_type = torch.bfloat16 + small = FlashInferAttBackend.create(model, make_spec(num_q_heads=8, head_size=64)) + large = FlashInferAttBackend.create(model, make_spec(num_q_heads=16, head_size=128)) + assert (small.tp_q_head_num, small.head_dim) == (8, 64) + assert (large.tp_q_head_num, large.head_dim) == (16, 128) + + +def test_replay_flashqla_rejects_decode_and_fp32(): + from lightllm.common.attention.backends.gdn.replay.backend import ReplayGDNFlashQlaAttBackend + + spec = make_spec(family=AttentionFamily.GDN) + with pytest.raises(UnsupportedAttention, match="decode"): + ReplayGDNFlashQlaAttBackend.check_support(spec) + with pytest.raises(UnsupportedAttention, match="dtype"): + ReplayGDNFlashQlaAttBackend.check_support(replace(spec, phase=AttentionPhase.PREFILL, dtype=torch.float32)) + + class _OutputBackend(BaseAttentionBackend): def _forward_decode(self, layer, context, inputs, *, out=None): return inputs @@ -94,6 +280,66 @@ def test_output_rejects_incompatible_buffer(invalid): layer.backend.forward(layer, context, result, out=out) +@pytest.mark.parametrize("family", list(AttentionFamily)) +def test_probe_reports_actual_coverage(monkeypatch, family): + from lightllm.common.attention import probe, registry + from lightllm.utils import backend_validator as validator + + spec = make_spec(family) + calls = [] + monkeypatch.setattr(registry, "load_backend_class", lambda name: SimpleNamespace(provider="triton")) + monkeypatch.setattr(probe, "probe_dense", lambda *args: calls.append("backend") or (True, None)) + monkeypatch.setattr(validator, "_smoke_provider", lambda *a, **k: calls.append("smoke") or (True, None)) + result = validator._probe_in_process("selected", device="cpu", spec=spec) + assert result.ok + numerical = family is AttentionFamily.DENSE + assert result.scope is (validator.ProbeScope.BACKEND_NUMERICS if numerical else validator.ProbeScope.PROVIDER_SMOKE) + assert calls == (["backend"] if numerical else ["smoke"]) + + +def test_special_dense_probe_is_explicitly_smoke_only(monkeypatch): + from dataclasses import replace + + from lightllm.common.attention import probe, registry + from lightllm.utils import backend_validator as validator + + monkeypatch.setattr(registry, "load_backend_class", lambda name: SimpleNamespace(provider="triton")) + monkeypatch.setattr(probe, "probe_dense", lambda *args: pytest.fail("unsupported numerical probe")) + monkeypatch.setattr(validator, "_smoke_provider", lambda *a, **k: (False, "provider failed")) + result = validator._probe_in_process("selected", device="cpu", spec=replace(make_spec(), use_mtp=True)) + assert not result.ok + assert result.scope is validator.ProbeScope.PROVIDER_SMOKE + assert result.reason == "provider failed" + + +def test_replay_flashqla_uses_inherited_initialization(monkeypatch): + import sys + from types import ModuleType + + from lightllm.common.attention.backends.gdn.replay.backend import ( + ReplayGDNFlashQlaAttBackend, + ReplayGDNTritonAttBackend, + ) + + calls = [] + original = ReplayGDNTritonAttBackend.__init__ + + def initialize(self, model, spec): + calls.append("parent") + original(self, model, spec) + + monkeypatch.setattr(ReplayGDNTritonAttBackend, "__init__", initialize) + monkeypatch.setattr(ReplayGDNTritonAttBackend, "_init_linear_layer_metadata", lambda *a: calls.append("metadata")) + module = ModuleType("flash_qla") + module.chunk_gated_delta_rule = object() + monkeypatch.setitem(sys.modules, "flash_qla", module) + model = model_stub() + backend = ReplayGDNFlashQlaAttBackend(model, make_spec(AttentionFamily.GDN)) + assert calls == ["parent", "metadata"] + assert backend.runtime is model.device_runtime + assert backend.prefill_kernel is module.chunk_gated_delta_rule + + def test_spec_configures_backend_before_planning(): from lightllm.common.attention.api import AttentionExecutionContext, AttentionWorkspace @@ -108,6 +354,40 @@ def test_spec_configures_backend_before_planning(): assert context.backend_identity is backend +@pytest.mark.parametrize("field,value", [("phase", AttentionPhase.PREFILL), ("family", AttentionFamily.MLA)]) +def test_selection_rejects_mixed_specs(field, value): + with pytest.raises(ValueError, match="share family, phase and cache format"): + selection.select_attention_backend( + specs=(make_spec(), replace(make_spec(), **{field: value})), + caps=model_stub().device_runtime.caps, + execution=model_stub().attention_requirements, + ) + + +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("phase", list(AttentionPhase)) +def test_gdn_output_contract(monkeypatch, replay, phase): + from lightllm.common.attention.api import GDNInputs, GDNOutput + from lightllm.common.attention.backends.gdn.replay.backend import ReplayGDNTritonAttBackend + from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNBackend + from unit_tests.common.attention.helpers import make_backend + + backend = make_backend(ReplayGDNTritonAttBackend if replay else TritonGDNBackend, mtp_step=0) + core, gate = torch.randn(2, 4, 8), torch.randn(2, 4, 8) + monkeypatch.setattr(backend, "_split_qkvzba", lambda x: (x, gate, None, None)) + monkeypatch.setattr(backend, "_prefill" if replay else "_run_prefill", lambda *a: core) + monkeypatch.setattr(backend, "_decode" if replay else "_run_decode", lambda *a: (core, gate)) + context = make_context(backend, make_batch(phase=phase)) + state = SimpleNamespace(finish_decode_token=lambda *a: None) + layer = make_layer_attention(backend, context, SimpleNamespace(conv=None, ssm=None, replay_state=state)) + out = torch.empty_like(core) + result = layer.backend.forward(layer, context, GDNInputs(core), out=out) + assert isinstance(result, GDNOutput) + assert result.hidden_states is out + assert result.gate is gate + torch.testing.assert_close(result.hidden_states, core) + + def test_binding_resolves_context_on_each_forward(monkeypatch): backend = _OutputBackend(model_stub(), make_spec()) first, second = make_context(backend), make_context(backend) diff --git a/unit_tests/common/attention/test_dense_numerics.py b/unit_tests/common/attention/test_dense_numerics.py new file mode 100644 index 0000000000..fed4977508 --- /dev/null +++ b/unit_tests/common/attention/test_dense_numerics.py @@ -0,0 +1,79 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention.api import AttentionPhase, CacheLayout, ContextCapacity, UnsupportedAttention +from lightllm.common.attention.api.inputs import DenseInputs, DenseParams +from lightllm.common.attention.backends._workspace import AttentionWorkspace +from lightllm.common.attention.registry import cache_format_to_backend, load_backend_class +from lightllm.common.attention_state import AttentionStateManager +from lightllm.platform import get_device_runtime +from unit_tests.common.attention.helpers import make_batch, make_layer_attention, make_spec, model_stub + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference(query, key, value, lengths, q_lengths, window=None): + rows, start = [], 0 + groups = query.shape[1] // key.shape[1] + for request, (length, q_length) in enumerate(zip(lengths, q_lengths)): + k = key[request * 16 : request * 16 + length].float().repeat_interleave(groups, dim=1) + v = value[request * 16 : request * 16 + length].float().repeat_interleave(groups, dim=1) + q = query[start : start + q_length].float() + logits = torch.einsum("qhd,khd->hqk", q, k) * query.shape[-1] ** -0.5 + causal = torch.arange(length, device=q.device)[None, :] <= ( + length - q_length + torch.arange(q_length, device=q.device)[:, None] + ) + if window is not None: + causal &= torch.arange(length, device=q.device)[None, :] >= ( + length - q_length + torch.arange(q_length, device=q.device)[:, None] - window[0] + ) + probs = logits.masked_fill(~causal, -torch.inf).softmax(-1) + rows.append(torch.einsum("hqk,khd->qhd", probs, v)) + start += q_length + return torch.cat(rows).to(query.dtype) + + +@pytest.mark.parametrize("provider", ["triton", "fa3", "flashinfer"]) +@pytest.mark.parametrize("phase", list(AttentionPhase)) +def test_startup_probe_runs_the_selected_dense_implementation(provider, phase): + from lightllm.common.attention.probe import probe_dense + + cls = load_backend_class(cache_format_to_backend["None"][provider]) + runtime = get_device_runtime() + try: + cls.check_environment(runtime.caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + spec = make_spec( + phase=phase, head_size=64, value_head_size=64, cache=CacheLayout("None", 16, torch.bfloat16, 2, 64, 64) + ) + success, error = probe_dense(cls, spec, runtime) + assert success, error + + +@pytest.mark.parametrize("provider", ["fa3", "flashinfer"]) +def test_noncausal_prefill_probe(provider): + from lightllm.common.attention.probe import probe_dense + + cls = load_backend_class(cache_format_to_backend["None"][provider]) + runtime = get_device_runtime() + cls.check_environment(runtime.caps) + spec = make_spec( + phase=AttentionPhase.PREFILL, + causal=False, + head_size=64, + value_head_size=64, + cache=CacheLayout("None", 16, torch.bfloat16, 2, 64, 64), + ) + cls.check_support(spec, runtime.caps) + success, error = probe_dense(cls, spec, runtime) + assert success, error + + +def test_triton_rejects_noncausal_prefill(): + cls = load_backend_class(cache_format_to_backend["None"]["triton"]) + with pytest.raises(UnsupportedAttention, match="non-causal prefill"): + cls.check_support(make_spec(phase=AttentionPhase.PREFILL, causal=False)) diff --git a/unit_tests/common/attention/test_linear_layout.py b/unit_tests/common/attention/test_linear_layout.py new file mode 100644 index 0000000000..739b7d0d18 --- /dev/null +++ b/unit_tests/common/attention/test_linear_layout.py @@ -0,0 +1,72 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.backends.base import BaseAttentionBackend +from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub + + +class _LinearBackend(BaseAttentionBackend): + def _allocate_decode_context(self, capacity, capture, workspace): + return LinearDecodeContext(capacity, capture, workspace) + + def _update_decode(self, context): + update_linear_decode(context, 2, self.dynamic_verify, replay=self.replay) + + def _forward_decode(self, layer, context, inputs, *, out=None): + return context.accepted_tokens(layer.params) + + +@pytest.mark.parametrize("dynamic", [False, True]) +@pytest.mark.parametrize("replay", [False, True]) +def test_shared_backend_selects_each_states_cursor_through_graph_replay(dynamic, replay): + backend = _LinearBackend(model_stub(), make_spec()) + backend.dynamic_verify, backend.replay = dynamic, replay + requests = [0, 0, 2, 7, 7, 7] if dynamic else [0, 0, 0, 2, 2, 2] + positions = [0, 1, 0, 0, 0, 0] if dynamic else [0, 1, 2, 0, 1, 2] + batch = make_batch( + (5,) * 6, + device="cuda", + b_req_idx=torch.tensor(requests, dtype=torch.int32, device="cuda"), + b_mtp_index=torch.tensor(positions, dtype=torch.int32, device="cuda"), + ) + context = make_context(backend, batch, capture=True, prepare=True) + first = torch.tensor([0, 0, 2, 0, 0, 0, 0, 99], dtype=torch.int32, device="cuda") + second = torch.tensor([2, 0, 1, 0, 0, 0, 0, 99], dtype=torch.int32, device="cuda") + layer_a = make_layer_attention(backend, context, params=first) + layer_b = replace(layer_a, logical_layer_id=1, params=second) + + def step(): + backend.update(context) + a, b = layer_a.backend.forward(layer_a, context, None), layer_b.backend.forward(layer_b, context, None) + assert layer_a.backend.forward(layer_a, context, None) is a + return a, b + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + step() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + a, b = step() + graph.replay() + tail = [1] * 4 if dynamic else [] + assert a.tolist() == [1, 3] + tail + assert b.tolist() == [3, 2] + tail + first[0], first[2] = 1, 0 + second[0], second[2] = 0, 2 + graph.replay() + assert a.tolist() == [2, 1] + tail + assert b.tolist() == [1, 3] + tail + # A new eager batch must not reuse the previous batch's gathered counts. + backend.prepare(context, batch) + a, b = step() + assert a.tolist() == [2, 1] + tail + assert b.tolist() == [1, 3] + tail + expected_requests = torch.tensor([0, 2] + ([7] * 4 if dynamic else []), device="cuda") + torch.testing.assert_close(context.b_conv_buffer_idx.long(), expected_requests) + expected_ssm = expected_requests if replay else expected_requests[:, None] * 3 + torch.arange(3, device="cuda") + torch.testing.assert_close(context.b_ssm_buffer_idx.long(), expected_ssm) diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 344bfea57f..3078cb4c4e 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -1,22 +1,22 @@ -import weakref import inspect import sys +import weakref from types import ModuleType, SimpleNamespace import pytest import torch -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.infer_struct import InferStateInfo -import lightllm.common.basemodel.attention.linear.gdn as gdn import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils +import lightllm.common.basemodel.attention.linear.gdn as gdn import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops +import lightllm.utils.backend_validator as backend_validator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.api_cli import make_argument_parser -import lightllm.utils.backend_validator as backend_validator @pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float32]) @@ -251,7 +251,9 @@ def kernel(**kwargs): flashqla = ModuleType("flash_qla") flashqla.chunk_gated_delta_rule = kernel monkeypatch.setitem(sys.modules, "flash_qla", flashqla) - monkeypatch.setattr(fla_ops, "chunk_gated_delta_rule", kernel) + from lightllm.common.attention.backends.gdn.triton.kernels import chunk + + monkeypatch.setattr(chunk, "chunk_gated_delta_rule", kernel) original_randn = torch.randn original_rand = torch.rand From 9678c8db485327a2f199b71d5ddd3646e6ed5ac9 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 11:53:54 +0800 Subject: [PATCH 4/9] refactor(attention): use varlen operators for vision and audio --- lightllm/common/attention/varlen.py | 42 +- .../basemodel/attention_vit/__init__.py | 0 .../basemodel/attention_vit/base_att.py | 38 -- .../basemodel/attention_vit/create_utils.py | 47 --- .../basemodel/attention_vit/fa3/__init__.py | 0 .../common/basemodel/attention_vit/fa3/fp.py | 60 --- .../basemodel/attention_vit/sdpa/__init__.py | 0 .../common/basemodel/attention_vit/sdpa/fp.py | 47 --- .../attention_vit/triton/__init__.py | 0 .../basemodel/attention_vit/triton/fp.py | 24 -- .../attention_vit/xformers/__init__.py | 0 .../basemodel/attention_vit/xformers/fp.py | 45 --- .../models/deepseek_v4/deepseek_v4_visual.py | 15 +- lightllm/models/glm5_next/glm5_next_visual.py | 6 +- lightllm/models/qwen2_5_vl/qwen2_5_visual.py | 32 +- lightllm/models/qwen2_vl/qwen2_visual.py | 30 +- .../qwen3_omni_audio.py | 32 +- .../layer_infer/transformer_layer_infer.py | 12 +- .../vit/triton_kernel/flashattention_nopad.py | 229 ----------- lightllm/server/api_cli.py | 3 +- lightllm/server/core/objs/start_args_type.py | 4 +- lightllm/server/visualserver/__init__.py | 48 ++- lightllm/server/visualserver/manager.py | 36 +- .../visualserver/model_infer/model_rpc.py | 47 +-- .../visualserver/visual_only_manager.py | 44 +-- .../common/attention/test_varlen_selection.py | 140 +++++++ .../basemodel/attention/test_lazy_import.py | 45 +-- unit_tests/models/deepseek_v4/test_vision.py | 30 +- unit_tests/models/glm5_next/test_visual.py | 22 +- .../vit/test_flash_attention_forward.py | 366 ++++++++++++------ 30 files changed, 634 insertions(+), 810 deletions(-) delete mode 100644 lightllm/common/basemodel/attention_vit/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/base_att.py delete mode 100644 lightllm/common/basemodel/attention_vit/create_utils.py delete mode 100644 lightllm/common/basemodel/attention_vit/fa3/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/fa3/fp.py delete mode 100644 lightllm/common/basemodel/attention_vit/sdpa/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/sdpa/fp.py delete mode 100644 lightllm/common/basemodel/attention_vit/triton/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/triton/fp.py delete mode 100644 lightllm/common/basemodel/attention_vit/xformers/__init__.py delete mode 100644 lightllm/common/basemodel/attention_vit/xformers/fp.py delete mode 100644 lightllm/models/vit/triton_kernel/flashattention_nopad.py create mode 100644 unit_tests/common/attention/test_varlen_selection.py diff --git a/lightllm/common/attention/varlen.py b/lightllm/common/attention/varlen.py index fded4439c9..158a626cfd 100644 --- a/lightllm/common/attention/varlen.py +++ b/lightllm/common/attention/varlen.py @@ -1,36 +1,44 @@ from lightllm.platform import get_device_runtime -from lightllm.utils.log_utils import init_logger from .api import AttentionFamily, AttentionPhase, UnsupportedAttention from .registry import AUDIO_VARLEN_PRIORITY, get_varlen_attention_op_class from .selection import BackendSelectionError -logger = init_logger(__name__) - class AutoVarlenAttentionRunner: - """音频编码器的自动路径；失败的实现只在当前调用方禁用，并在本次调用中重算输出。""" + """One encoder owns this runner and shares stateless operators across its layers. + + Unsupported inputs can try the next provider for this call. Kernel and programming + errors propagate: retrying on a potentially failed device is not a safe fallback. + Environment rejections are cached by provider/device; input rejections are not. + """ def __init__(self, providers=AUDIO_VARLEN_PRIORITY): self.providers = providers self.ops = {} - self.failures = {} + self.unavailable = {} def forward(self, inputs, *, out): + failures = {} + caps = get_device_runtime(inputs.query.device).caps for name in self.providers: - if name in self.failures: + key = (name, inputs.query.device) + if key in self.unavailable: + failures[name] = self.unavailable[key] continue - try: - if name not in self.ops: + if key not in self.ops: + try: op_cls = get_varlen_attention_op_class(name) - caps = get_device_runtime(inputs.query.device).caps op_cls.check_environment(caps) - op_cls.check_support(inputs, caps) - self.ops[name] = op_cls() - return self.ops[name].forward(inputs, out=out) + self.ops[key] = op_cls() + except (UnsupportedAttention, ImportError) as exc: + self.unavailable[key] = failures[name] = str(exc) + continue + op = self.ops[key] + try: + op.check_support(inputs, caps) except UnsupportedAttention as exc: - self.failures[name] = str(exc) - except Exception as exc: - self.failures[name] = str(exc) - logger.warning("Disabling %s varlen attention for this encoder: %s", name, exc) - raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, self.failures) + failures[name] = str(exc) + continue + return op.forward(inputs, out=out) + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, failures) diff --git a/lightllm/common/basemodel/attention_vit/__init__.py b/lightllm/common/basemodel/attention_vit/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/base_att.py b/lightllm/common/basemodel/attention_vit/base_att.py deleted file mode 100644 index 49bf6ad74c..0000000000 --- a/lightllm/common/basemodel/attention_vit/base_att.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -from abc import ABC, abstractmethod - - -class BaseVitAttBackend(ABC): - """ - 用于创建支持各种不同的AttBackend, 如 fa3, sdpa, triton 实现等， - 这个是单列模式, 每种backend只有一个实例 - """ - - _instances = {} - - def __new__(cls, *args, **kwargs): - """ - 重写__new__方法实现单例模式 - """ - # 检查是否已经有该类的实例 - if cls not in cls._instances: - # 创建新实例并存储 - instance = super().__new__(cls) - cls._instances[cls] = instance - # 返回已有的实例 - return cls._instances[cls] - - def __init__(self): - pass - - @abstractmethod - def _vit_att_fwd( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: - raise NotImplementedError("not impl") diff --git a/lightllm/common/basemodel/attention_vit/create_utils.py b/lightllm/common/basemodel/attention_vit/create_utils.py deleted file mode 100644 index 9c436d42b3..0000000000 --- a/lightllm/common/basemodel/attention_vit/create_utils.py +++ /dev/null @@ -1,47 +0,0 @@ -from importlib import import_module - -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.utils.backend_validator import validate -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -vit_att_backend = { - "triton": "triton.fp:TritonVitAttBackend", - "sdpa": "sdpa.fp:SdpaVitAttBackend", - "fa3": "fa3.fp:Fa3VitAttBackend", - "xformers": "xformers.fp:XformersVitAttBackend", -} - - -def get_vit_att_backend_class(backend_name: str) -> BaseVitAttBackend: - module_name, class_name = vit_att_backend[backend_name].split(":") - return getattr(import_module(f"{__package__}.{module_name}"), class_name) - - -def init_vit_att_backend(index=0, priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str: - args = get_env_start_args() - backend_name = args.vit_att_backend[index] - if backend_name != "auto": - logger.info(f"Selected {backend_name} backend for VIT") - return backend_name - return _select_vit_backend(priority_list=priority_list) - - -def _select_vit_backend(priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str: - """Auto-select the best available backend with validation for VIT. - - Priority: FA3 > Xformers > Sdpa > Triton - This legacy entry checks provider smoke only; the Varlen numerical path is wired in PR4. - """ - - for backend_name in priority_list: - if validate(backend_name): - logger.info(f"Auto-selected {backend_name} backend (provider smoke checked) for VIT") - return backend_name - - # Fallback to triton without validation (should not happen) - logger.warning("No backend validation succeeded, falling back to triton") - return "triton" diff --git a/lightllm/common/basemodel/attention_vit/fa3/__init__.py b/lightllm/common/basemodel/attention_vit/fa3/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/fa3/fp.py b/lightllm/common/basemodel/attention_vit/fa3/fp.py deleted file mode 100644 index a8e56417c5..0000000000 --- a/lightllm/common/basemodel/attention_vit/fa3/fp.py +++ /dev/null @@ -1,60 +0,0 @@ -import dataclasses -import torch -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend - - -class Fa3VitAttBackend(BaseVitAttBackend): - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> None: - - head_dim = q.shape[-1] - softmax_scale = head_dim ** -0.5 - window_size = (-1, -1) - torch.ops.sgl_kernel.fwd.default( - q, - k, - v, - None, # k_new - None, # v_new - None, # qv - o, # out - cu_seqlens, - cu_seqlens, - None, # cu_seqlens_k_new - None, - None, - max_seqlen, - max_seqlen, - None, # page_table, - None, # kv_batch_idx - None, # leftpad_k - None, # rotary cos - None, # rotary sin - None, # seqlens_rotary - None, - None, - None, - softmax_scale, - False, - window_size[0], - window_size[1], - attention_chunk=0, - softcap=0.0, - is_rotary_interleaved=False, - scheduler_metadata=None, - num_splits=1, - pack_gqa=None, - sm_margin=0, - sinks=None, - sparse_mask_fine=None, - only_qv=False, - ) - - return o diff --git a/lightllm/common/basemodel/attention_vit/sdpa/__init__.py b/lightllm/common/basemodel/attention_vit/sdpa/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/sdpa/fp.py b/lightllm/common/basemodel/attention_vit/sdpa/fp.py deleted file mode 100644 index 6c6da2c2b1..0000000000 --- a/lightllm/common/basemodel/attention_vit/sdpa/fp.py +++ /dev/null @@ -1,47 +0,0 @@ -import torch -import torch.nn.functional as F -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend - - -class SdpaVitAttBackend(BaseVitAttBackend): - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: - assert q.ndim == k.ndim == v.ndim == o.ndim == 3 - assert cu_seqlens is not None and cu_seqlens.ndim == 1 - cu_seqlens = cu_seqlens.detach().to("cpu") - B = cu_seqlens.numel() - 1 - - with torch.no_grad(): - for b in range(B): - s = int(cu_seqlens[b]) - e = int(cu_seqlens[b + 1]) - L = e - s - if L <= 0: - continue - if max_seqlen: - assert L <= max_seqlen - - # [L, H, D] -> [1, H, L, D] - q_ = q[s:e].permute(1, 0, 2).unsqueeze(0) - k_ = k[s:e].permute(1, 0, 2).unsqueeze(0) - v_ = v[s:e].permute(1, 0, 2).unsqueeze(0) - - out = F.scaled_dot_product_attention( - q_, - k_, - v_, - attn_mask=None, - dropout_p=0.0, - is_causal=False, - ) - # [1, H, L, D] -> [L, H, D] - o[s:e].copy_(out.squeeze(0).permute(1, 0, 2)) - - return o diff --git a/lightllm/common/basemodel/attention_vit/triton/__init__.py b/lightllm/common/basemodel/attention_vit/triton/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/triton/fp.py b/lightllm/common/basemodel/attention_vit/triton/fp.py deleted file mode 100644 index 88867102a4..0000000000 --- a/lightllm/common/basemodel/attention_vit/triton/fp.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.models.vit.triton_kernel.flashattention_nopad import _flash_attention_triton_fwd - - -class TritonVitAttBackend(BaseVitAttBackend): - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ): - _flash_attention_triton_fwd( - q, - k, - v, - o, - cu_seqlens, # q k v cu_seqlens, - max_seqlen, - ) - return o diff --git a/lightllm/common/basemodel/attention_vit/xformers/__init__.py b/lightllm/common/basemodel/attention_vit/xformers/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention_vit/xformers/fp.py b/lightllm/common/basemodel/attention_vit/xformers/fp.py deleted file mode 100644 index 361b5db050..0000000000 --- a/lightllm/common/basemodel/attention_vit/xformers/fp.py +++ /dev/null @@ -1,45 +0,0 @@ -import torch -import torch.nn.functional as F - -try: - from xformers import ops as xformers_ops - from xformers.ops import fmha -except ImportError: - xformers_ops = None - fmha = None - -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend - - -class XformersVitAttBackend(BaseVitAttBackend): - @torch.no_grad() - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: - assert q.ndim == k.ndim == v.ndim == o.ndim == 3 - assert cu_seqlens is not None and cu_seqlens.ndim == 1 - assert q.shape == k.shape == v.shape == o.shape - - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.int64).tolist() - seqlens = [int(L) for L in seqlens if int(L) > 0] - - if len(seqlens) == 0: - return o - if max_seqlen: - assert max(seqlens) <= max_seqlen - - attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens, device=q.device) - - q_ = q.unsqueeze(0) # [1, T, H, D] - k_ = k.unsqueeze(0) # [1, T, H, D] - v_ = v.unsqueeze(0) # [1, T, H, D] - - out = xformers_ops.memory_efficient_attention(q_, k_, v_, attn_bias=attn_bias, p=0.0) - o.copy_(out.squeeze(0)) # [T, H, D] - return o diff --git a/lightllm/models/deepseek_v4/deepseek_v4_visual.py b/lightllm/models/deepseek_v4/deepseek_v4_visual.py index 3e730c5b2a..24985d08fd 100644 --- a/lightllm/models/deepseek_v4/deepseek_v4_visual.py +++ b/lightllm/models/deepseek_v4/deepseek_v4_visual.py @@ -13,14 +13,11 @@ from safetensors import safe_open from torch import nn -from lightllm.models.deepseek_v4.image_processor import ( - IMAGE, - build_canonical_image_block, - load_image, -) +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.models.deepseek_v4.image_processor import IMAGE, build_canonical_image_block, load_image from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm from lightllm.server.multimodal_params import ImageItem -from lightllm.server.visualserver import get_vit_attn_backend +from lightllm.server.visualserver import get_vit_attention_op @lru_cache(8) @@ -56,7 +53,7 @@ class PatchEmbed(nn.Module): def __init__(self, args): super().__init__() self.proj = nn.Linear( - 3 * args.vision_patch_size ** 2, + 3 * args.vision_patch_size**2, args.vision_dim, dtype=torch.bfloat16, ) @@ -79,7 +76,7 @@ def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, cu_seql q = apply_rotary(q, cos, sin) k = apply_rotary(k, cos, sin) o = torch.empty_like(q) - get_vit_attn_backend()(q, k, v, o, cu_seqlens, n) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, n), out=o) return self.wo(o.reshape(n, -1)) @@ -143,7 +140,7 @@ class Aligner(nn.Module): def __init__(self, args): super().__init__() self.downsample_ratio = args.vision_downsample_ratio - in_dim = args.vision_dim * self.downsample_ratio ** 2 + in_dim = args.vision_dim * self.downsample_ratio**2 self.w1 = nn.Linear(in_dim, args.hidden_size, dtype=torch.bfloat16) self.w2 = nn.Linear(args.hidden_size, args.hidden_size, dtype=torch.bfloat16) diff --git a/lightllm/models/glm5_next/glm5_next_visual.py b/lightllm/models/glm5_next/glm5_next_visual.py index 43637c2cbb..c27e68aa76 100644 --- a/lightllm/models/glm5_next/glm5_next_visual.py +++ b/lightllm/models/glm5_next/glm5_next_visual.py @@ -9,10 +9,12 @@ import torch.nn.functional as F from safetensors import safe_open +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton from lightllm.models.vit.triton_kernel.rms_norm_vit import qk_rms_norm, rms_norm -from lightllm.server.visualserver import get_vit_attn_backend +from lightllm.server.visualserver import get_vit_attention_op + from .vision_process import Glm5NextImageProcessor @@ -56,7 +58,7 @@ def forward(self, x, cu_seqlens, max_seqlen, rotary_cos, rotary_sin): q = apply_rotary_pos_emb_triton(q, rotary_cos, rotary_sin) k = apply_rotary_pos_emb_triton(k, rotary_cos, rotary_sin) out = torch.empty_like(q) - get_vit_attn_backend()(q, k, v, out, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=out) return self.proj(out.reshape(x.shape[0], -1)) diff --git a/lightllm/models/qwen2_5_vl/qwen2_5_visual.py b/lightllm/models/qwen2_5_vl/qwen2_5_visual.py index 1b3a5f0db7..fcfe0726c8 100644 --- a/lightllm/models/qwen2_5_vl/qwen2_5_visual.py +++ b/lightllm/models/qwen2_5_vl/qwen2_5_visual.py @@ -1,21 +1,24 @@ -import os import json +import os +from io import BytesIO +from typing import List, Optional + import torch +import torch.nn as nn import torch.nn.functional as F from PIL import Image -from typing import List, Optional -from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data -from io import BytesIO -import torch.nn as nn -from transformers.activations import ACT2FN -from lightllm.models.qwen2_vl.vision_process import resize_image, Qwen2VLImageProcessor from safetensors import safe_open -from lightllm.server.multimodal_params import ImageItem -from lightllm.models.qwen2_vl.qwen2_visual import PatchEmbed, VisionRotaryEmbedding -from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm -from lightllm.server.visualserver import get_vit_attn_backend +from transformers.activations import ACT2FN + +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.models.qwen2_vl.qwen2_visual import PatchEmbed, VisionRotaryEmbedding from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor, resize_image +from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.server.visualserver import get_vit_attention_op class Qwen2RMSNorm(nn.Module): @@ -74,7 +77,7 @@ def forward( k = apply_rotary_pos_emb_triton(k, rotary_cos, rotary_sin) attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - get_vit_attn_backend()(q, k, v, attn_output, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -122,7 +125,7 @@ def forward( class Qwen2_5_VLPatchMerger(nn.Module): def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: super().__init__() - self.hidden_size = context_dim * (spatial_merge_size ** 2) + self.hidden_size = context_dim * (spatial_merge_size**2) self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6) self.mlp = nn.Sequential( nn.Linear(self.hidden_size, self.hidden_size), @@ -350,7 +353,6 @@ def load_image(self, img: List[ImageItem]): return pixel_values.to(dtype=self.data_type), image_grid_thw def load_model(self, weight_dir): - bin_weight_files = [file_ for file_ in os.listdir(weight_dir) if file_.endswith(".bin")] if bin_weight_files: weight_dict = {} @@ -390,7 +392,7 @@ def encode(self, images: List[ImageItem]): raise Exception("Unsupported input types: {} for {}".format(type(img), img)) # must devide merge_length - cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size ** 2) + cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size**2) valid_ids.append([valid_id, valid_id + cur_num]) valid_id += cur_num diff --git a/lightllm/models/qwen2_vl/qwen2_visual.py b/lightllm/models/qwen2_vl/qwen2_visual.py index e02c3d9aa3..172e5aacb9 100644 --- a/lightllm/models/qwen2_vl/qwen2_visual.py +++ b/lightllm/models/qwen2_vl/qwen2_visual.py @@ -18,24 +18,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os import json +import os +from io import BytesIO +from typing import List + import torch +import torch.nn as nn import torch.nn.functional as F from PIL import Image -from typing import List -from torchvision import transforms as T -from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data -from io import BytesIO -import torch.nn as nn +from safetensors import safe_open from torch.nn import LayerNorm +from torchvision import transforms as T from transformers.activations import ACT2FN -from safetensors import safe_open -from lightllm.server.multimodal_params import ImageItem -from lightllm.server.visualserver import get_vit_attn_backend -from lightllm.models.qwen2_vl.vision_process import resize_image, Qwen2VLImageProcessor + +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor, resize_image +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.server.visualserver import get_vit_attention_op # adapted from # https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -73,7 +76,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class PatchMerger(nn.Module): def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> None: super().__init__() - self.hidden_size = context_dim * (spatial_merge_size ** 2) + self.hidden_size = context_dim * (spatial_merge_size**2) self.ln_q = LayerNorm(context_dim, eps=1e-6) self.mlp = nn.Sequential( nn.Linear(self.hidden_size, self.hidden_size), @@ -147,7 +150,7 @@ def forward( attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - get_vit_attn_backend()(q, k, v, attn_output, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -239,7 +242,6 @@ def _init_datatype(self): return def load_model(self, weight_dir): - processor_config_path = os.path.join(weight_dir, "preprocessor_config.json") with open(processor_config_path, "r") as f: processor_config_dict = json.load(f) @@ -322,7 +324,7 @@ def encode(self, images: List[ImageItem]): raise Exception("Unsupported input types: {} for {}".format(type(img), img)) # must devide merge_length - cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size ** 2) + cur_num = img_tensors[-1].shape[0] // (self.spatial_merge_size**2) valid_ids.append([valid_id, valid_id + cur_num]) valid_id += cur_num diff --git a/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py b/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py index b9ea66fb15..919afd96ff 100644 --- a/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py +++ b/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py @@ -1,19 +1,21 @@ -import os import json import math -import torch +import os +from typing import Callable, List, Optional, Union + import numpy as np -from torch import Tensor, nn +import torch from safetensors import safe_open +from torch import Tensor, nn from torch.nn import functional as F -from typing import Callable, Optional, Union, List from transformers.activations import ACT2FN -from lightllm.server.multimodal_params import AudioItem -from lightllm.utils.log_utils import init_logger +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.common.attention.varlen import AutoVarlenAttentionRunner from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager -from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd from lightllm.models.qwen3_omni_moe_thinker.audio_process import WhisperFeatureExtractor +from lightllm.server.multimodal_params import AudioItem +from lightllm.utils.log_utils import init_logger QWEN3_OMNI_CONV_CHUNKSIZE = int(os.getenv("LIGHTLLM_QWEN3_OMNI_CONV_CHUNKSIZE", 200)) @@ -24,7 +26,6 @@ def _get_feat_extract_output_lengths(input_lengths): """ Computes the output length of the convolutional layers and the output length of the audio encoder """ - input_lengths_leave = input_lengths % 100 feat_lengths = (input_lengths_leave - 1) // 2 + 1 output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 @@ -41,10 +42,13 @@ def __init__( activation_function, activation_dropout, encoder_ffn_dim, + attention_op=None, ): super().__init__() self.embed_dim = d_model - self.self_attn = Qwen3OmniMoeAudioAttention(d_model, encoder_attention_heads, attention_dropout) + self.self_attn = Qwen3OmniMoeAudioAttention( + d_model, encoder_attention_heads, attention_dropout, attention_op=attention_op + ) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.dropout = dropout self.activation_fn = ACT2FN[activation_function] @@ -86,7 +90,7 @@ def forward( class Qwen3OmniMoeAudioAttention(nn.Module): - def __init__(self, d_model, encoder_attention_heads, attention_dropout): + def __init__(self, d_model, encoder_attention_heads, attention_dropout, *, attention_op=None): super().__init__() self.embed_dim = d_model self.num_heads = encoder_attention_heads @@ -99,10 +103,11 @@ def __init__(self, d_model, encoder_attention_heads, attention_dropout): f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" f" and `num_heads`: {self.num_heads})." ) - self.scaling = self.head_dim ** -0.5 + self.scaling = self.head_dim**-0.5 self.attention_dropout = 0.0 self.is_decoder = False self.is_causal = False + self.attention_op = attention_op if attention_op is not None else AutoVarlenAttentionRunner() self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) @@ -116,7 +121,6 @@ def forward( **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """Input shape: Batch x Time x Channel""" - seq_length, _ = hidden_states.size() q = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, -1) @@ -125,7 +129,7 @@ def forward( attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - flash_attention_fwd(q, k, v, attn_output, cu_seqlens, max_seqlen) + self.attention_op.forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.out_proj(attn_output) @@ -176,6 +180,7 @@ def __init__( self.embed_scale = math.sqrt(self.embed_dim) if scale_embedding else 1.0 self.n_window = n_window self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, self.embed_dim) + attention_op = AutoVarlenAttentionRunner() self.layers = nn.ModuleList( [ Qwen3OmniMoeAudioEncoderLayer( @@ -186,6 +191,7 @@ def __init__( activation_function, activation_dropout, encoder_ffn_dim, + attention_op=attention_op, ) for _ in range(encoder_layers) ] diff --git a/lightllm/models/vit/layer_infer/transformer_layer_infer.py b/lightllm/models/vit/layer_infer/transformer_layer_infer.py index cada55e58b..f8dbf2b7d4 100644 --- a/lightllm/models/vit/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/vit/layer_infer/transformer_layer_infer.py @@ -1,13 +1,15 @@ +from typing import Tuple + import torch import torch.distributed as dist -from typing import Tuple +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager from lightllm.models.vit.layer_weights.transformer_layer_weight import ViTTransformerLayerWeight -from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size from lightllm.models.vit.triton_kernel.gelu_vit import gelu_fwd from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm -from lightllm.server.visualserver import get_vit_attn_backend -from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.server.visualserver import get_vit_attention_op +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size class ViTTransformerLayerInfer: @@ -105,7 +107,7 @@ def _context_attention_kernel(self, q, k, v) -> torch.Tensor: q, k, v, out = map(reshape, (q, k, v, out)) cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) * seq_len max_seqlen = seq_len - get_vit_attn_backend()(q, k, v, out, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=out) return out.reshape(batch_size, seq_len, -1) def _get_o(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor: diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/models/vit/triton_kernel/flashattention_nopad.py deleted file mode 100644 index af062c2a5b..0000000000 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ /dev/null @@ -1,229 +0,0 @@ -import torch -import triton -import triton.language as tl -import math -import time -import torch.nn.functional as F -from typing import Optional, Tuple -from lightllm.utils.device_utils import is_hopper, is_4090 - -if triton.__version__ >= "2.1.0": - - @triton.jit - def _fwd_kernel( - Q, - K, - V, - sm_scale, - Out, - q_stride_s, - q_stride_h, - q_stride_d, - k_stride_s, - k_stride_h, - k_stride_d, - v_stride_s, - v_stride_h, - v_stride_d, - o_stride_s, - o_stride_h, - o_stride_d, - head_dim_act, - cu_seqlens, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, - ): - cur_batch = tl.program_id(2) - cur_head = tl.program_id(1) - start_m = tl.program_id(0) - - seq_start = tl.load(cu_seqlens + cur_batch).to(tl.int32) - seq_end = tl.load(cu_seqlens + cur_batch + 1).to(tl.int32) - seq_len = seq_end - seq_start - - # initialize offsets - offs_n = tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL) - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - - mask_d = offs_d < head_dim_act - off_q = cur_head * q_stride_h + (seq_start + offs_m[:, None]) * q_stride_s + offs_d[None, :] * q_stride_d - q = tl.load(Q + off_q, mask=(offs_m[:, None] < seq_len) & mask_d[None, :], other=0.0) - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) - - for start_n in range(0, seq_len, BLOCK_N): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - off_k = ( - (seq_start + start_n + offs_n[None, :]) * k_stride_s - + cur_head * k_stride_h - + offs_d[:, None] * k_stride_d - ) - k = tl.load(K + off_k, mask=((start_n + offs_n[None, :]) < seq_len) & mask_d[:, None], other=0.0) - - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, k) - qk *= sm_scale - qk += tl.where((start_n + offs_n[None, :]) < seq_len, 0, float("-inf")) - - # -- compute m_ij, p, l_ij - m_ij = tl.maximum(tl.max(qk, 1), m_i) - p = tl.exp(qk - m_ij[:, None]) - l_ij = tl.sum(p, 1) - - acc_scale = tl.exp(m_i - m_ij) - acc = acc * acc_scale[:, None] - - # update acc - off_v = ( - (seq_start + start_n + offs_n[:, None]) * v_stride_s - + cur_head * v_stride_h - + offs_d[None, :] * v_stride_d - ) - v = tl.load(V + off_v, mask=((start_n + offs_n[:, None]) < seq_len) & mask_d[None, :], other=0.0) - p = p.to(v.dtype) - acc += tl.dot(p, v) - # update m_i and l_i - m_i = m_ij - l_i_new = tl.exp(l_i - m_ij) + l_ij - l_i = m_ij + tl.log(l_i_new) - - o_scale = tl.exp(m_i - l_i) - acc = acc * o_scale[:, None] - # initialize pointers to output - off_o = (seq_start + offs_m[:, None]) * o_stride_s + cur_head * o_stride_h + offs_d[None, :] * o_stride_d - out_ptrs = Out + off_o - tl.store(out_ptrs, acc, mask=(offs_m[:, None] < seq_len) & mask_d[None, :]) - return - - @torch.no_grad() - def _flash_attention_triton_fwd( - q, - k, - v, - o, - cu_seqlens, # q k v cu_seqlens, - max_seqlen, - ): - BLOCK = 64 - # shape constraints - assert q.ndim == k.ndim == v.ndim == o.ndim == 3, "q, k, v, o must be 3D tensors" - _, head_num, head_dim = q.shape - batch_size = cu_seqlens.numel() - 1 - - sm_scale = 1.0 / (head_dim ** 0.5) # 计算scale系数 - d_pad = triton.next_power_of_2(head_dim) - grid = (triton.cdiv(max_seqlen, BLOCK), head_num, batch_size) # batch, head, - num_warps = 4 - _fwd_kernel[grid]( - q, - k, - v, - sm_scale, - o, - q.stride(0), - q.stride(1), - q.stride(2), - k.stride(0), - k.stride(1), - k.stride(2), - v.stride(0), - v.stride(1), - v.stride(2), - o.stride(0), - o.stride(1), - o.stride(2), - head_dim, - cu_seqlens, - BLOCK_M=BLOCK, - BLOCK_DMODEL=d_pad, - BLOCK_N=BLOCK, - num_warps=num_warps, - num_stages=2, - ) - return - -else: - raise Exception("error triton version!") - -_flash_attn_v3_available = False -try: - from sgl_kernel.flash_attn import flash_attn_varlen_func - - _flash_attn_v3_available = True - - def flash_attention_v3_fwd( - q, - k, - v, - o, - cu_seqlens, - max_seqlen, - ): - head_dim = q.shape[-1] - softmax_scale = head_dim ** -0.5 - window_size = (-1, -1) - torch.ops.sgl_kernel.fwd.default( - q, - k, - v, - None, # k_new - None, # v_new - None, # qv - o, # out - cu_seqlens, - cu_seqlens, - None, # cu_seqlens_k_new - None, - None, - max_seqlen, - max_seqlen, - None, # page_table, - None, # kv_batch_idx - None, # leftpad_k - None, # rotary cos - None, # rotary sin - None, # seqlens_rotary - None, - None, - None, - softmax_scale, - False, - window_size[0], - window_size[1], - 0, # attention_chunk - 0.0, # softcap - is_rotary_interleaved=False, - scheduler_metadata=None, - num_splits=1, - pack_gqa=None, - sm_margin=0, - sinks=None, - sparse_mask_fine=None, - only_qv=False, - ) - return o - -except ImportError: - print("Failed to import _flash_attn_forward from hopper.flash_attn_interface.") - _flash_attn_v3_available = False - - -def flash_attention_fwd(q, k, v, o, cu_seqlens, max_seqlen): - """ - 统一的 Flash Attention 接口。如果 sgl_kernel 存在， - 则使用 sgl_kernel里的接口，否则使用 Triton 版本。 - """ - global _flash_attn_v3_available - if _flash_attn_v3_available and (is_hopper() or is_4090()): - try: - flash_attention_v3_fwd(q, k, v, o, cu_seqlens, max_seqlen) - return - except Exception as e: - print(f"Failed to use sgl_kernel: {e}") - _flash_attn_v3_available = False - _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 16a8735a7c..068fc207b1 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -536,9 +536,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--vit_att_backend", type=str, - nargs="+", choices=["auto", "triton", "fa3", "sdpa", "xformers"], - default=["auto"], + default="auto", help="""vit attention kernel used in vlm. auto: automatically select best backend based on GPU and available packages (priority: fa3 > xformers > sdpa > triton)""", diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0609132e62..033ca9e68a 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -187,9 +187,7 @@ class StartArgs: default="pin_mem_cache", metadata={"choices": ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]}, ) - vit_att_backend: List[str] = field( - default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]} - ) + vit_att_backend: str = field(default="auto", metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]}) llm_kv_type: str = field( default="None", metadata={"choices": ["None", "int8kv", "int4kv", "fp8kv_sph", "fp8kv_spt", "fp8kv_dsa"]} ) diff --git a/lightllm/server/visualserver/__init__.py b/lightllm/server/visualserver/__init__.py index 6bc0923da0..b4e65bf259 100644 --- a/lightllm/server/visualserver/__init__.py +++ b/lightllm/server/visualserver/__init__.py @@ -1,16 +1,44 @@ -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.common.basemodel.attention_vit.create_utils import get_vit_att_backend_class +from lightllm.common.attention.api import VarlenAttentionOp +from lightllm.common.attention.registry import get_varlen_attention_op_class +from lightllm.common.attention.selection import select_varlen_attention_op +from lightllm.platform import get_device_runtime -VIT_ATTN_BACKEND: BaseVitAttBackend = None +class _LazyVarlenAttentionRunner: + """用 worker 预热时的真实 Q/K/V 选择实现，之后通过固定的 forward 委托执行。 -def set_vit_att_backend(backend_name: str): - global VIT_ATTN_BACKEND - VIT_ATTN_BACKEND = get_vit_att_backend_class(backend_name) + 部分视觉模型固定使用 BF16，与 llm 的 data_type 不同。manager 无法据启动 + 参数推断实际 dtype/head dim，因此选择发生在各设备的首次预热输入到达时。 + """ + + def __init__(self, requested): + self.requested = requested + self._op = None + + def forward(self, inputs, *, out): + if self._op is None: + runtime = get_device_runtime(inputs.query.device) + if runtime.is_capturing(): + raise RuntimeError("warm up vision attention before Graph capture") + name = select_varlen_attention_op( + self.requested, device=inputs.query.device, dtype=inputs.query.dtype, head_dim=inputs.query.shape[-1] + ) + self._op = get_varlen_attention_op_class(name)() + return self._op.forward(inputs, out=out) + + +VIT_ATTENTION_OP: VarlenAttentionOp | _LazyVarlenAttentionRunner | None = None + + +def set_vit_attention_op(provider: str, *, defer_validation=False): + global VIT_ATTENTION_OP + VIT_ATTENTION_OP = ( + _LazyVarlenAttentionRunner(provider) if defer_validation else get_varlen_attention_op_class(provider)() + ) return -def get_vit_attn_backend(): - if VIT_ATTN_BACKEND is None: - raise RuntimeError("VIT_ATTN_BACKEND is not initialized. Call init_vit_att_backend() first.") - return VIT_ATTN_BACKEND._vit_att_fwd +def get_vit_attention_op(): + if VIT_ATTENTION_OP is None: + raise RuntimeError("VIT_ATTENTION_OP is not initialized. Call set_vit_attention_op() first.") + return VIT_ATTENTION_OP diff --git a/lightllm/server/visualserver/manager.py b/lightllm/server/visualserver/manager.py index 873eff8456..d6faad3ca0 100644 --- a/lightllm/server/visualserver/manager.py +++ b/lightllm/server/visualserver/manager.py @@ -1,29 +1,31 @@ -import zmq -import zmq.asyncio import asyncio -import uvloop -import rpyc -import socket -import pickle +import collections import inspect -import setproctitle +import pickle +import socket import threading -import collections from typing import List -from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes + +import rpyc +import setproctitle +import uvloop +import zmq +import zmq.asyncio + from lightllm.server.core.objs import ShmReqManager, StartArgs +from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from .model_infer import start_model_process, VisualModelRpcClient -from lightllm.common.basemodel.attention_vit.create_utils import init_vit_att_backend -from lightllm.utils.log_utils import init_logger +from rpyc.utils.classic import obtain + +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.graceful_utils import graceful_registry +from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import start_parent_check_thread -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.shm_port_args import get_shm_port_args -from rpyc.utils.classic import obtain +from .model_infer import VisualModelRpcClient, start_model_process logger = init_logger(__name__) @@ -62,12 +64,10 @@ def __init__( self.lock = asyncio.Lock() async def wait_to_model_ready(self): - self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(self.vit_dp)] - self.vit_attn_backend = init_vit_att_backend(index=0) + self.vit_attn_backend = self.args.vit_att_backend for dp_rank_id in range(self.vit_dp): for tp_rank_id in range(self.vit_tp): - rpc_model = await start_model_process() self.model_rpcs[dp_rank_id].append(rpc_model) diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 6451e70008..6d1146c0a2 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -1,39 +1,40 @@ import os -import rpyc -import torch -import socket -import torch.multiprocessing as mp import queue +import socket import threading import time +from typing import Deque, Dict, List, Optional, Tuple + +import rpyc +import torch import torch.distributed as dist -from typing import Dict, List, Tuple, Deque, Optional -from transformers.configuration_utils import PretrainedConfig +import torch.multiprocessing as mp from rpyc.utils.classic import obtain -from lightllm.models.qwen_vl.qwen_visual import QWenVisionTransformer -from lightllm.models.llava.llava_visual import LlavaVisionModel +from transformers.configuration_utils import PretrainedConfig + +from lightllm.models.deepseek_v4.deepseek_v4_visual import DeepseekV4VisionModel from lightllm.models.gemma3.gemma3_visual import Gemma3VisionModel from lightllm.models.gemma4.gemma4_visual import Gemma4VisionModel -from lightllm.models.deepseek_v4.deepseek_v4_visual import DeepseekV4VisionModel -from lightllm.models.vit.model import VisionTransformer -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel +from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer +from lightllm.models.llava.llava_visual import LlavaVisionModel +from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel from lightllm.models.qwen2_5_vl.qwen2_5_visual import Qwen2_5_VisionTransformerPretrainedModel +from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel +from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel +from lightllm.models.qwen_vl.qwen_visual import QWenVisionTransformer from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel -from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel -from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel -from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer -from lightllm.utils.infer_utils import set_random_seed +from lightllm.models.vit.model import VisionTransformer +from lightllm.server.embed_cache.afs_utils import SepEmbedHandler +from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.server.visualserver import set_vit_attention_op +from lightllm.server.visualserver.model_infer.vision_peak_vram_hold import VisionPeakVramHolder from lightllm.utils.dist_utils import init_vision_distributed_env from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.shm_port_args import get_shm_port_args -from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient -from lightllm.server.visualserver import set_vit_att_backend -from lightllm.server.embed_cache.afs_utils import SepEmbedHandler +from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger -from lightllm.server.visualserver.model_infer.vision_peak_vram_hold import VisionPeakVramHolder - +from lightllm.utils.shm_port_args import get_shm_port_args logger = init_logger(__name__) @@ -67,7 +68,7 @@ def exposed_init_model(self, kvargs): self.is_visual_only_mode = get_env_start_args().run_mode == "visual_only" self.data_type = kvargs["data_type"] self.vit_attn_backend = kvargs["vit_attn_backend"] - set_vit_att_backend(self.vit_attn_backend) + set_vit_attention_op(self.vit_attn_backend, defer_validation=True) init_vision_distributed_env(kvargs) model_cfg, _ = PretrainedConfig.get_config_dict(weight_dir) diff --git a/lightllm/server/visualserver/visual_only_manager.py b/lightllm/server/visualserver/visual_only_manager.py index 835a25fe3e..ba501eaa7e 100644 --- a/lightllm/server/visualserver/visual_only_manager.py +++ b/lightllm/server/visualserver/visual_only_manager.py @@ -1,35 +1,38 @@ import asyncio -import uvloop -import rpyc import inspect -import setproctitle -import threading -import uuid +import os import pickle -import websockets +import signal import socket import sys -import os -import signal +import threading import time -from lightllm.utils.net_utils import get_hostname_ip -from .objs import VIT_Obj +import uuid from typing import List + +import rpyc +import setproctitle +import uvloop +import websockets + from lightllm.server.core.objs import StartArgs +from lightllm.utils.net_utils import get_hostname_ip + +from .objs import VIT_Obj asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from .model_infer import start_model_process, VisualModelRpcClient -from lightllm.common.basemodel.attention_vit.create_utils import init_vit_att_backend -from lightllm.utils.log_utils import init_logger +from rpyc.utils.classic import obtain + +from lightllm.server.embed_cache.utils import create_shm, free_shm, get_shm_name_data +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.graceful_utils import graceful_registry +from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import start_parent_check_thread -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.shm_port_args import get_shm_port_args -from rpyc.utils.classic import obtain -from lightllm.server.embed_cache.utils import create_shm, get_shm_name_data, free_shm -from .manager import VisualManager +from .manager import VisualManager +from .model_infer import VisualModelRpcClient, start_model_process logger = init_logger(__name__) @@ -74,7 +77,6 @@ async def register_to_config_server_loop(self, args: StartArgs): try: uri = f"ws://{args.config_server_host}:{ports.config_server_port}/visual_register" async with websockets.connect(uri, max_queue=(2048 * 1024, 2048 * 1023)) as websocket: - sock = websocket.transport.get_extra_info("socket") sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) @@ -94,12 +96,10 @@ async def register_to_config_server_loop(self, args: StartArgs): logger.info("reconnection to config_server") async def wait_to_model_ready(self): - self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(self.vit_dp)] - self.vit_attn_backend = init_vit_att_backend(index=0) + self.vit_attn_backend = self.args.vit_att_backend for dp_rank_id in range(self.vit_dp): for tp_rank_id in range(self.vit_tp): - rpc_model = await start_model_process() self.model_rpcs[dp_rank_id].append(rpc_model) diff --git a/unit_tests/common/attention/test_varlen_selection.py b/unit_tests/common/attention/test_varlen_selection.py new file mode 100644 index 0000000000..aa113d1535 --- /dev/null +++ b/unit_tests/common/attention/test_varlen_selection.py @@ -0,0 +1,140 @@ +import pytest +import torch + +from lightllm.common.attention import selection +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.utils import backend_validator +from lightllm.utils.backend_validator import ProbeResult, ProbeScope + + +def test_auto_probes_actual_varlen_path_on_assigned_device(monkeypatch): + calls = [] + + def probe(name, **kwargs): + calls.append((name, kwargs)) + return ProbeResult(name == "sdpa", ProbeScope.VARLEN_NUMERICS, "rejected") + + monkeypatch.setattr(backend_validator, "probe_varlen", probe) + assert selection.select_varlen_attention_op(device=3) == "sdpa" + assert calls == [ + (name, {"device": torch.device("cuda:3"), "dtype": None, "head_dim": 80}) + for name in ("fa3", "xformers", "sdpa") + ] + + +@pytest.mark.parametrize("requested", ["auto", "fa3", "triton", "missing"]) +def test_rejected_ops_are_not_returned(monkeypatch, requested): + calls = [] + + def probe(name, **kwargs): + calls.append(name) + return ProbeResult(False, ProbeScope.VARLEN_NUMERICS, "rejected") + + monkeypatch.setattr(backend_validator, "probe_varlen", probe) + with pytest.raises(selection.BackendSelectionError) as error: + selection.select_varlen_attention_op(requested) + expected = list(selection.varlen_attention_ops) if requested == "auto" else [requested] + assert list(error.value.failures) == expected + assert calls == ([] if requested == "missing" else expected) + + +def test_varlen_probe_exercises_selected_op(monkeypatch): + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + + monkeypatch.setattr(backend_validator.torch.cuda, "is_available", lambda: False) + assert backend_validator._validate_varlen("sdpa") == (True, None) + monkeypatch.setattr(SdpaVarlenAttentionOp, "forward", lambda self, inputs, out: out.fill_(float("nan"))) + success, reason = backend_validator._validate_varlen("sdpa") + assert not success + assert "mismatch" in reason + + +@pytest.mark.parametrize("check", ["check_environment", "check_support"]) +def test_rejected_probe_does_not_execute_op(monkeypatch, check): + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + + monkeypatch.setattr(backend_validator.torch.cuda, "is_available", lambda: False) + + def reject(cls, *args): + raise UnsupportedAttention("layout") + + monkeypatch.setattr(SdpaVarlenAttentionOp, check, classmethod(reject)) + monkeypatch.setattr(SdpaVarlenAttentionOp, "forward", lambda *args, **kwargs: pytest.fail("rejected operator ran")) + assert backend_validator._validate_varlen("sdpa") == (False, "layout") + + +def test_worker_initialization_keeps_fixed_forward_and_uses_actual_inputs(monkeypatch): + import torch + + from lightllm.common.attention.api import VarlenAttentionInputs + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + from lightllm.server import visualserver + + calls = [] + + def select(requested, **kwargs): + calls.append((requested, kwargs)) + return "sdpa" + + monkeypatch.setattr(visualserver, "select_varlen_attention_op", select) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", None) + visualserver.set_vit_attention_op("auto", defer_validation=True) + first = visualserver.get_vit_attention_op() + q = torch.randn(4, 2, 64, dtype=torch.float16) + inputs = VarlenAttentionInputs(q, q, q, torch.tensor([0, 4], dtype=torch.int32), 4) + first.forward(inputs, out=torch.empty_like(q)) + assert calls == [("auto", {"device": q.device, "dtype": q.dtype, "head_dim": 64})] + assert visualserver.get_vit_attention_op() is first + assert isinstance(first._op, SdpaVarlenAttentionOp) + assert "forward" not in vars(first) + first.forward(inputs, out=torch.empty_like(q)) + assert len(calls) == 1 + + +def test_probe_cache_separates_device_dtype_and_layout(monkeypatch): + from dataclasses import replace + + import torch + + from unit_tests.common.attention.helpers import make_spec + + calls = [] + monkeypatch.setattr( + backend_validator, + "_probe", + lambda *args, **kwargs: calls.append(kwargs) or ProbeResult(True, ProbeScope.PROVIDER_SMOKE), + ) + backend_validator.probe_backend.cache_clear() + try: + spec = make_spec() + for device, dtype, shape in ( + ("cuda:0", torch.float16, spec), + ("cuda:1", torch.float16, spec), + ("cuda:0", torch.bfloat16, spec), + ("cuda:0", torch.float16, replace(spec, head_size=64)), + ("cuda:0", torch.float16, spec), + ): + assert backend_validator.probe_backend("provider", device=device, dtype=dtype, spec=shape).ok + assert len(calls) == 4 + finally: + backend_validator.probe_backend.cache_clear() + + +def test_empty_device_assignment_does_not_probe(monkeypatch): + monkeypatch.setattr(backend_validator, "probe_varlen", lambda *a, **k: pytest.fail("empty assignment probed")) + with pytest.raises(ValueError, match="at least one device"): + selection.select_varlen_attention_op(devices=()) + + +def test_default_device_is_resolved_before_probe(monkeypatch): + calls = [] + monkeypatch.setattr( + selection, "get_device_runtime", lambda device: type("Runtime", (), {"device": torch.device("cpu")})() + ) + monkeypatch.setattr( + backend_validator, + "probe_varlen", + lambda *a, **kw: calls.append(kw["device"]) or ProbeResult(True, ProbeScope.VARLEN_NUMERICS), + ) + assert selection.select_varlen_attention_op("sdpa") == "sdpa" + assert calls == [torch.device("cpu")] diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py index e7423d30f9..5ecdb1c71f 100644 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -3,13 +3,12 @@ import importlib import importlib.abc import importlib.util -from pathlib import Path import sys -from types import ModuleType, SimpleNamespace import unittest +from pathlib import Path +from types import ModuleType, SimpleNamespace from unittest.mock import Mock, patch - ROOT = Path(__file__).resolve().parents[4] ATTENTION = "lightllm.common.basemodel.attention" @@ -212,37 +211,15 @@ def test_linear_selection_falls_back_to_triton(self): self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) def test_vit_selection_loads_only_requested_implementation(self): - package = "lightllm.common.basemodel.attention_vit" - for name in list(sys.modules): - if name == package or name.startswith(f"{package}."): - del sys.modules[name] - base = ModuleType(f"{package}.base_att") - base.BaseVitAttBackend = type("BaseVitAttBackend", (), {}) - sys.modules[base.__name__] = base - sys.modules["lightllm.utils.backend_validator"]._validate = self.validate - imports = BackendImports(package) - sys.meta_path.insert(0, imports) - self.addCleanup(sys.meta_path.remove, imports) - selection = importlib.import_module(f"{package}.create_utils") - self.assertEqual(imports.imported, []) - - for name, class_name in ( - ("fa3", "Fa3VitAttBackend"), - ("triton", "TritonVitAttBackend"), - ("sdpa", "SdpaVitAttBackend"), - ("xformers", "XformersVitAttBackend"), - ): - backend = imports.allow(f"{name}.fp", class_name) - self.args.vit_att_backend = [name] - self.assertEqual(selection.init_vit_att_backend(), name) - self.assertIs(selection.get_vit_att_backend_class(name), backend) - self.validate.assert_not_called() + from lightllm.common.attention import registry, selection - self.args.vit_att_backend = ["auto"] - self.validate.side_effect = lambda name: name == "sdpa" - self.assertEqual(selection.init_vit_att_backend(), "sdpa") - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "xformers", "sdpa"]) + loaded = [] + sentinel = object() + def import_module(name): + loaded.append(name) + return SimpleNamespace(SdpaVarlenAttentionOp=sentinel) -if __name__ == "__main__": - unittest.main() + with patch.object(registry.importlib, "import_module", import_module): + assert registry.get_varlen_attention_op_class("sdpa") is sentinel + assert loaded == ["lightllm.common.attention.backends.dense.sdpa.varlen"] diff --git a/unit_tests/models/deepseek_v4/test_vision.py b/unit_tests/models/deepseek_v4/test_vision.py index ad931ac2b5..b037a2f84d 100644 --- a/unit_tests/models/deepseek_v4/test_vision.py +++ b/unit_tests/models/deepseek_v4/test_vision.py @@ -12,12 +12,10 @@ from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.attention_vit.sdpa.fp import SdpaVitAttBackend -from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual -from lightllm.models.deepseek_v4 import image_processor as ours +from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp +from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual, image_processor as ours from lightllm.server import visualserver - REFERENCE_DIR = "/mtc/models/DeepSeek-V4-Flash-Vision-Exp/inference" REFERENCE_PROCESSOR_PATH = os.path.join(REFERENCE_DIR, "image_processor.py") REFERENCE_VISION_PATH = os.path.join(REFERENCE_DIR, "vision.py") @@ -105,7 +103,7 @@ def test_tokenizer_spans_match_reference_blocks(reference_processor, prefix_len, expected_cache = [] for index, (image, (height, width)) in enumerate(zip(params.images, grids)): canonical_types, _ = reference_processor.build_image_block(height, width, start_pos=0) - image.token_id = 2 ** 62 + (0 if reuse_image and index == 2 else index * 1024) + image.token_id = 2**62 + (0 if reuse_image and index == 2 else index * 1024) image.token_num = len(canonical_types) image.start_index_in_embed_cache = 100 + (0 if reuse_image and index == 2 else index * 1024) expected_cache.append((image.token_id, image.token_num, image.start_index_in_embed_cache)) @@ -201,7 +199,7 @@ def _tiny_vision_args(): @pytest.mark.parametrize("n_vit_h,n_vit_w", [(4, 4), (5, 3)]) def test_vit_matches_reference(reference_vision, n_vit_h, n_vit_w, monkeypatch): - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", SdpaVitAttBackend) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", SdpaVarlenAttentionOp()) torch.manual_seed(0) args = _tiny_vision_args() model = visual.ViT(args).eval() @@ -348,3 +346,23 @@ def fake_read_shm(name): assert torch.equal(types[:3], torch.full((3,), ours.IMAGE_PAD, dtype=torch.int64)) assert types[3] == ours.IMAGE_START assert types[-1] == ours.IMAGE_END + + +def test_attention_uses_sequence_boundaries(monkeypatch): + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", SdpaVarlenAttentionOp()) + torch.manual_seed(17) + attention = visual.Attention(_tiny_vision_args()).eval() + x = torch.randn(7, 16, dtype=torch.bfloat16) + cos = torch.ones(7, 1, 2) + sin = torch.zeros_like(cos) + cu = torch.tensor([0, 3, 7], dtype=torch.int32) + with torch.inference_mode(), sdpa_kernel(SDPBackend.MATH): + q, k, v = [t.reshape(7, 4, 4) for t in attention.wqkv(x).chunk(3, dim=-1)] + parts = [] + for start, end in ((0, 3), (3, 7)): + qi, ki, vi = [t[start:end].float().transpose(0, 1) for t in (q, k, v)] + weights = (qi @ ki.transpose(-1, -2) * 0.5).softmax(-1) + parts.append((weights @ vi).transpose(0, 1).to(x.dtype)) + expected = attention.wo(torch.cat(parts).reshape(7, 16)) + actual = attention(x, cos, sin, cu) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py index 6a25a8d4ea..a7b7dde2ce 100644 --- a/unit_tests/models/glm5_next/test_visual.py +++ b/unit_tests/models/glm5_next/test_visual.py @@ -25,7 +25,7 @@ def test_image_canvas_and_token_expansion(tmp_path, width, height, canvas): image.image_w, image.image_h = width, height image.token_id = 1000 image.token_num = tokenizer.get_image_token_length(image) - assert image.token_num == canvas[0] * canvas[1] // 28 ** 2 + assert image.token_num == canvas[0] * canvas[1] // 28**2 multi = MultimodalParams() multi.images = [image] ids = tokenizer.encode([1, 100, 102, 101, 2], multi) @@ -55,8 +55,8 @@ def test_image_patch_order_temporal_repeat_and_black_padding(): def test_vision_attention_uses_triton_rope_with_fa3_backend(monkeypatch): import lightllm.models.glm5_next.glm5_next_visual as glm_visual import lightllm.server.visualserver as visualserver - from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend - from lightllm.server.visualserver import set_vit_att_backend + from lightllm.common.attention.backends.dense.fa3.varlen import Fa3VarlenAttentionOp + from lightllm.server.visualserver import set_vit_attention_op rope_inputs = [] norm_inputs = [] @@ -70,16 +70,18 @@ def fake_rope(x, cos, sin): rope_inputs.append((x, cos, sin)) return x - def fake_fa3(q, k, v, o, cu_seqlens, max_seqlen): + def fake_fa3(self, inputs, *, out): + q, k, v = inputs.query, inputs.key, inputs.value + o, cu_seqlens, max_seqlen = out, inputs.cu_seqlens, inputs.max_seqlen attn_inputs.append((q, k, v, cu_seqlens, max_seqlen)) o.copy_(q) return o monkeypatch.setattr(glm_visual, "qk_rms_norm", fake_qk_norm) monkeypatch.setattr(glm_visual, "apply_rotary_pos_emb_triton", fake_rope) - monkeypatch.setattr(Fa3VitAttBackend, "_vit_att_fwd", staticmethod(fake_fa3)) - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", visualserver.VIT_ATTN_BACKEND) - set_vit_att_backend("fa3") + monkeypatch.setattr(Fa3VarlenAttentionOp, "forward", fake_fa3) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", visualserver.VIT_ATTENTION_OP) + set_vit_attention_op("fa3") attention = glm_visual.Glm5NextVisionAttention(hidden_size=16, num_heads=2, eps=1e-5, bias=True) x = torch.randn(3, 16) @@ -149,10 +151,10 @@ def glm_rms_norm(x, weight, eps): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_vision_batch_preserves_independent_image_attention(): from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer - from lightllm.server.visualserver import set_vit_att_backend + from lightllm.server.visualserver import set_vit_attention_op torch.manual_seed(53) - set_vit_att_backend("sdpa") + set_vit_attention_op("sdpa") model = ( Glm5NextVisionTransformer( {"data_type": "float32"}, @@ -167,7 +169,7 @@ def test_vision_batch_preserves_independent_image_attention(): .eval() ) grid = torch.tensor([[1, 4, 6], [1, 6, 4]]) - x = torch.randn(48, 3 * 2 * 14 ** 2, device="cuda") + x = torch.randn(48, 3 * 2 * 14**2, device="cuda") with torch.inference_mode(): batch = model(x, grid) individual = torch.cat([model(x[:24], grid[:1]), model(x[24:], grid[1:])]) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index f736ec4cd6..9a161bad45 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -1,139 +1,271 @@ -import torch -import math -import time import pytest -from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd -from lightllm.models.vit.triton_kernel import flashattention_nopad -from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend +import torch + +from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionInputs +from lightllm.common.attention.registry import get_varlen_attention_op_class +from lightllm.common.attention.varlen import AutoVarlenAttentionRunner +from lightllm.platform import get_device_runtime + + +def make_inputs(lengths=(3, 0, 5), *, device="cpu", dtype=torch.float32, head_dim=80, strided=False): + tokens = sum(lengths) + if strided: + qkv = torch.randn(tokens, 3, 2, head_dim, dtype=dtype, device=device) + q, k, v = qkv.unbind(1) + else: + q, k, v = [torch.randn(tokens, 2, head_dim, dtype=dtype, device=device) for _ in range(3)] + cu = torch.tensor([0, *lengths], dtype=torch.int32, device=device).cumsum(0, dtype=torch.int32) + return VarlenAttentionInputs(q, k, v, cu, max(lengths)) -_fa3_mark = pytest.mark.skipif( - not flashattention_nopad._flash_attn_v3_available - or not torch.cuda.is_available() - or torch.cuda.get_device_capability()[0] < 9, - reason="FA3 requires sgl_kernel and a Hopper GPU", -) +def reference_attention(inputs): + out = torch.empty_like(inputs.query) + offsets = inputs.cu_seqlens.cpu().tolist() + for start, end in zip(offsets, offsets[1:]): + if start == end: + continue + q, k, v = [x[start:end].float().transpose(0, 1) for x in (inputs.query, inputs.key, inputs.value)] + scores = q @ k.transpose(-1, -2) * inputs.query.shape[-1] ** -0.5 + out[start:end] = (scores.softmax(-1) @ v).transpose(0, 1).to(out.dtype) + return out + +@pytest.mark.parametrize("error_type", [None, UnsupportedAttention]) +def test_audio_fallback_only_rejects_current_input(monkeypatch, error_type): + from lightllm.common.attention import varlen -def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): calls = [] - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("fa3") - o[0].fill_(99) # A backend may fail after partially writing output. - raise RuntimeError("shape-dependent FA3 failure") - - def fallback(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("triton") - o.fill_(7) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) - monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - assert not flashattention_nopad._flash_attn_v3_available - - # Later calls should use the disabled-backend path without retrying FA3. + + class FakeOp: + @classmethod + def check_environment(cls, caps): + return None + + @classmethod + def check_support(cls, inputs, caps=None): + calls.append("fa3") + if error_type: + raise error_type("shape-dependent FA3 failure") + + def forward(self, inputs, *, out): + out.fill_(9) + return out + + class Fallback(FakeOp): + @classmethod + def check_support(cls, inputs, caps=None): + return None + + def forward(self, inputs, *, out): + calls.append("triton") + out.fill_(7) + return out + + monkeypatch.setattr(varlen, "get_varlen_attention_op_class", lambda name: {"fa3": FakeOp, "triton": Fallback}[name]) + op = AutoVarlenAttentionRunner() + inputs = make_inputs() + out = torch.full_like(inputs.query, -123) + assert op.forward(inputs, out=out) is out + assert torch.equal(out, torch.full_like(out, 7 if error_type else 9)) out.fill_(-123) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) + op.forward(inputs, out=out) + assert calls == (["fa3", "triton", "fa3", "triton"] if error_type else ["fa3", "fa3"]) + assert torch.equal(out, torch.full_like(out, 7 if error_type else 9)) + + +@pytest.mark.parametrize("lengths", [(3, 0, 5), (1, 7), (0, 0)]) +def test_sdpa_preserves_sequence_boundaries_and_output_storage(lengths): + inputs = make_inputs(lengths) + inputs.query.zero_() + inputs.key.zero_() + start = 0 + for i, length in enumerate(lengths): + inputs.value[start : start + length].fill_(i + 1) + start += length + out = torch.empty_like(inputs.query) + result = get_varlen_attention_op_class("sdpa")().forward(inputs, out=out) + assert result is out + torch.testing.assert_close(out, inputs.value) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton", "sdpa", "xformers"]) +@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) +@pytest.mark.parametrize("head_dim,strided", [(64, False), (80, True), (128, False)]) +def test_varlen(provider, dtype, atol, head_dim, strided): + op = get_varlen_attention_op_class(provider) + try: + op.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + torch.manual_seed(0) + inputs = make_inputs((1, 0, 65, 129), device="cuda", dtype=dtype, head_dim=head_dim, strided=strided) + storage = torch.full((*inputs.query.shape[:2], head_dim * 2), -123, device="cuda", dtype=dtype) + out = storage[..., :head_dim] + before = [x.clone() for x in (inputs.query, inputs.key, inputs.value)] + result = op().forward(inputs, out=out) + assert result is out + torch.testing.assert_close(out, reference_attention(inputs), atol=atol, rtol=0) + assert torch.equal(storage[..., head_dim:], torch.full_like(out, -123)) + for actual, expected in zip((inputs.query, inputs.key, inputs.value), before): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) -def test_fa3_success_does_not_run_fallback(monkeypatch): - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton", "sdpa", "xformers"]) +def test_empty_batch(provider): + op = get_varlen_attention_op_class(provider) + try: + op.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + inputs = make_inputs((0, 0), device="cuda", dtype=torch.bfloat16) + out = torch.empty_like(inputs.query) + assert op().forward(inputs, out=out) is out - def fa3(q, k, v, o, cu_seqlens, max_seqlen): - o.fill_(9) - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton"]) +def test_graph_replay_uses_new_inputs(provider): + op_cls = get_varlen_attention_op_class(provider) + try: + op_cls.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + op = op_cls() + assert op.supports_graph + inputs = make_inputs((3, 5), device="cuda", dtype=torch.bfloat16) + out = torch.empty_like(inputs.query) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op.forward(inputs, out=out) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op.forward(inputs, out=out) + inputs.value.add_(1) + inputs.cu_seqlens.copy_(torch.tensor([0, 5, 8], device="cuda", dtype=torch.int32)) + graph.replay() + torch.testing.assert_close(out, reference_attention(inputs), atol=2e-2, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_audio_attention_matches_projected_qkv(monkeypatch): + from lightllm.models.qwen3_omni_moe_thinker import qwen3_omni_audio as audio + + torch.manual_seed(19) + attention = audio.Qwen3OmniMoeAudioAttention(128, 2, 0).cuda().to(torch.bfloat16).eval() monkeypatch.setattr( - flashattention_nopad, - "_flash_attention_triton_fwd", - lambda *args: pytest.fail("a successful FA3 call must not run fallback"), + audio.g_cache_manager, + "alloc_tensor", + lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device=device), ) - flash_attention_fwd(q, q, q, out, cu, 2) - assert torch.equal(out, torch.full_like(q, 9)) - assert flashattention_nopad._flash_attn_v3_available + hidden = torch.randn(8, 128, dtype=torch.bfloat16, device="cuda") + cu = torch.tensor([0, 3, 8], dtype=torch.int32, device="cuda") + with torch.inference_mode(): + q, k, v = [proj(hidden).reshape(8, 2, 64) for proj in (attention.q_proj, attention.k_proj, attention.v_proj)] + expected = attention.out_proj(reference_attention(VarlenAttentionInputs(q, k, v, cu, 5)).reshape(8, 128)) + actual = attention(hidden, cu_seqlens=cu, max_seqlen=5) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) -def reference_attention_varlen(q, k, v, cu): - """ - q, k, v : (total_len, n_head, D) - cu_seqlen : prefix sums (batch+1,) - """ - total, n_head, d = q.shape - out = torch.empty_like(q) - scale = 1.0 / math.sqrt(d) +@pytest.mark.parametrize("error_type", [RuntimeError, TypeError, AssertionError]) +def test_audio_kernel_errors_propagate_without_retry(monkeypatch, error_type): + from lightllm.common.attention import varlen - for b in range(cu.numel() - 1): - s, e = cu[b].item(), cu[b + 1].item() - q_b, k_b, v_b = q[s:e], k[s:e], v[s:e] # (seq, head, D) + class Broken: + @classmethod + def check_environment(cls, caps): + pass - q_hsd = q_b.permute(1, 0, 2) # (head, seq, D) - k_hds = k_b.permute(1, 2, 0) # (head, D, seq) - v_hsd = v_b.permute(1, 0, 2) # (head, seq, D) + @classmethod + def check_support(cls, inputs, caps=None): + pass - scores = torch.matmul(q_hsd, k_hds) * scale # (head, seq, seq) - probs = torch.softmax(scores.float(), dim=-1) + def forward(self, inputs, *, out): + out.fill_(99) + raise error_type("kernel failed") - out_hsd = torch.matmul(probs, v_hsd.float()) # (head, seq, D) - out[s:e] = out_hsd.permute(1, 0, 2).to(q.dtype) # back to (seq, head, D) + calls = [] - return out + def load(name): + calls.append(name) + return Broken + monkeypatch.setattr(varlen, "get_varlen_attention_op_class", load) + op = AutoVarlenAttentionRunner() + inputs = make_inputs() + with pytest.raises(error_type, match="kernel failed"): + op.forward(inputs, out=torch.empty_like(inputs.query)) + assert calls == ["fa3"] + assert not op.unavailable -@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) -@pytest.mark.parametrize( - "attention_fn", - [ - flash_attention_fwd, - pytest.param( - getattr(flashattention_nopad, "flash_attention_v3_fwd", None), - marks=_fa3_mark, - ), - pytest.param( - Fa3VitAttBackend._vit_att_fwd, - marks=_fa3_mark, - ), - ], -) -def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): - torch.manual_seed(0) - lengths = torch.randint(1, 257, (batch,)) - max_len = int(lengths.max().item()) - - cu = torch.zeros(batch + 1, dtype=torch.int32, device=device) - cu[1:] = torch.cumsum(lengths, 0) - tot = int(cu[-1]) - - q = torch.randn(tot, heads, d, dtype=dtype, device=device) - k = torch.randn_like(q) - v = torch.randn_like(q) - out_tri = torch.randn_like(q) - attention_fn(q, k, v, out_tri, cu, max_len) - a = time.time() - for _ in range(100): - attention_fn(q, k, v, out_tri, cu, max_len) - b = time.time() - print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") - out_ref = reference_attention_varlen(q, k, v, cu) - - max_err = (out_ref - out_tri).abs().max().item() - mean_err = (out_ref - out_tri).abs().mean().item() - print(f"{dtype}: max {max_err:.6f}, mean {mean_err:.6f}") - torch.testing.assert_close(out_tri, out_ref, atol=atol, rtol=0) - - -if __name__ == "__main__": - pytest.main() + +def test_audio_encoder_shares_runner_only_with_its_layers(): + from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_audio import Qwen3OmniMoeAudioEncoder + + def encoder(): + return Qwen3OmniMoeAudioEncoder( + {}, + d_model=16, + encoder_attention_heads=2, + encoder_layers=2, + encoder_ffn_dim=32, + downsample_hidden_size=4, + output_dim=16, + ) + + first, second = encoder(), encoder() + assert first.layers[0].self_attn.attention_op is first.layers[1].self_attn.attention_op + assert first.layers[0].self_attn.attention_op is not second.layers[0].self_attn.attention_op + + +@pytest.mark.parametrize("boundaries,maximum", [([1, 3, 8], 5), ([0, 7, 6, 8], 7), ([0, 3, 7], 5), ([0, 3, 8], 4)]) +def test_varlen_rejects_invalid_boundaries(boundaries, maximum): + inputs = make_inputs()._replace(cu_seqlens=torch.tensor(boundaries, dtype=torch.int32), max_seqlen=maximum) + with pytest.raises(ValueError, match="cu_seqlens"): + inputs.validate() + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_auto_validates_once_per_call_and_rechecks_mutated_boundaries(monkeypatch, device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + inputs = make_inputs(device=device) + original = VarlenAttentionInputs.validate + calls = [] + + def validate(self): + calls.append(self.cu_seqlens) + return original(self) + + monkeypatch.setattr(VarlenAttentionInputs, "validate", validate) + op = AutoVarlenAttentionRunner(providers=("sdpa",)) + output = torch.empty_like(inputs.query) + op.forward(inputs, out=output) + inputs.cu_seqlens.copy_(torch.tensor([0, 1, 1, 8], dtype=torch.int32, device=device)) + op.forward(inputs._replace(max_seqlen=7), out=output) + assert len(calls) == 2 + torch.testing.assert_close(output, reference_attention(inputs._replace(max_seqlen=7))) + + +@pytest.mark.parametrize("provider", ["fa3", "triton", "xformers", "sdpa"]) +@pytest.mark.parametrize("invalid", ["shape", "dtype", "device", "overlap", "alias", "head_stride"]) +def test_varlen_output_validation(provider, invalid): + inputs = make_inputs() + out = torch.empty_like(inputs.query) + if invalid == "shape": + out = out[:-1] + elif invalid == "dtype": + out = out.double() + elif invalid == "device": + out = out.to("meta") + elif invalid == "overlap": + out = out[:1].expand_as(out) + elif invalid == "alias": + out = inputs.key + else: + out = torch.empty((*out.shape[:-1], out.shape[-1] * 2))[..., ::2] + with pytest.raises(ValueError, match="varlen out"): + get_varlen_attention_op_class(provider)().forward(inputs, out=out) From 925319505532006835aefd6b857c819c6b5d8ce6 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 11:53:59 +0800 Subject: [PATCH 5/9] refactor(attention): integrate model bindings and state lifecycle --- docs/CN/source/models/add_new_model.md | 14 +- docs/EN/source/models/add_new_model.md | 14 +- .../attention_state/bounded_att/base_state.py | 22 + .../bounded_att/gdn/gdn_state.py | 19 +- .../bounded_att/gdn/replay_gdn_state.py | 23 + .../attention_state/bounded_att/interface.py | 3 +- .../bounded_att/kda/kda_state.py | 19 +- .../bounded_att/kda/replay_kda_state.py | 23 + .../sliding_window/fp8_e8m0_swa_state.py | 10 +- .../window_compressor_state.py | 10 +- lightllm/common/attention_state/commit.py | 34 + .../attention_state/full_att/base_state.py | 23 +- .../compressed_att/compressed_att_state.py | 11 +- .../attention_state/full_att/gqa/gqa_state.py | 77 +- .../full_att/gqa/group_quantized_state.py | 10 +- .../attention_state/full_att/gqa/ops.py | 44 + .../full_att/gqa/static_fp8_state.py | 11 +- .../attention_state/full_att/mla/mla_state.py | 10 +- .../packed_mla/fp8_e8m0_packed_mla_state.py | 9 +- .../full_att/packed_mla/packed_mla_state.py | 10 +- lightllm/common/attention_state/manager.py | 49 +- lightllm/common/basemodel/attention.py | 294 +++++ .../common/basemodel/attention/__init__.py | 36 - .../common/basemodel/attention/base_att.py | 183 --- .../attention/create_linear_utils.py | 79 -- .../basemodel/attention/create_utils.py | 166 --- .../basemodel/attention/fa3/__init__.py | 0 lightllm/common/basemodel/attention/fa3/fp.py | 333 ----- .../common/basemodel/attention/fa3/fp8.py | 205 ---- .../common/basemodel/attention/fa3/mla.py | 244 ---- .../common/basemodel/attention/fa3/neo.py | 81 -- .../attention/flashinfer/__init__.py | 0 .../attention/flashinfer/env_utils.py | 9 - .../basemodel/attention/flashinfer/fp.py | 291 ----- .../basemodel/attention/flashinfer/fp8.py | 146 --- .../basemodel/attention/flashinfer/mla.py | 272 ----- .../basemodel/attention/flashinfer/utils.py | 17 - .../basemodel/attention/import_utils.py | 6 - .../basemodel/attention/linear/__init__.py | 0 .../basemodel/attention/linear/flashqla.py | 8 - .../common/basemodel/attention/linear/gdn.py | 426 ------- .../common/basemodel/attention/linear/kda.py | 253 ---- .../basemodel/attention/linear/replay_gdn.py | 345 ------ .../basemodel/attention/linear/replay_kda.py | 247 ---- .../basemodel/attention/linear/triton.py | 7 - .../basemodel/attention/nsa/__init__.py | 0 .../attention/nsa/dsv4_fp8_flashmla_sparse.py | 225 ---- .../attention/nsa/flashmla_sparse.py | 193 --- .../attention/nsa/fp8_flashmla_sparse.py | 204 ---- .../basemodel/attention/nsa/fp8_glm5_next.py | 238 ---- .../basemodel/attention/nsa/glm5_next.py | 100 -- .../basemodel/attention/triton/__init__.py | 0 .../common/basemodel/attention/triton/fp.py | 349 ------ .../basemodel/attention/triton/int4kv.py | 176 --- .../basemodel/attention/triton/int8kv.py | 209 ---- .../common/basemodel/attention/triton/mla.py | 164 --- .../common/basemodel/attention/triton/neo.py | 50 - lightllm/common/basemodel/attention_batch.py | 35 + .../common/basemodel/attention_runtime.py | 83 ++ lightllm/common/basemodel/basemodel.py | 199 +-- lightllm/common/basemodel/batch_objs.py | 2 + lightllm/common/basemodel/cuda_graph.py | 192 ++- lightllm/common/basemodel/graph.py | 4 +- lightllm/common/basemodel/infer_struct.py | 159 ++- .../basemodel/layer_infer/base_layer_infer.py | 31 +- .../common/basemodel/prefill_cuda_graph.py | 120 +- .../basemodel/triton_kernel/att/__init__.py | 0 .../triton_kernel/att/decode_att/__init__.py | 0 .../att/decode_att/gqa/__init__.py | 0 .../decode_att/gqa/flash_decoding/__init__.py | 0 .../gqa/flash_decoding/gqa_flash_decoding.py | 63 - .../decode_att/gqa/mtp_diverse/__init__.py | 21 - .../att/decode_att/int4kv/__init__.py | 0 .../int4kv/ppl_int4kv_flash_decoding.py | 64 - .../att/decode_att/int8kv/__init__.py | 0 .../int8kv/int8kv_flash_decoding_diverse.py | 105 -- .../att/decode_att/int8kv/normal/__init__.py | 1 - .../int8kv/normal/int8kv_flash_decoding.py | 70 -- .../att/decode_att/mha/__init__.py | 0 .../decode_att/mha/flash_decoding/__init__.py | 0 .../mha/flash_decoding/flash_decoding.py | 40 - .../mha/stage3_decode_att/__init__.py | 0 .../triton_kernel/att/prefill_att/__init__.py | 0 .../triton_kernel/linear_att/__init__.py | 25 - .../triton_kernel/linear_att/fla/__init__.py | 11 - .../linear_att/fla/ops/__init__.py | 21 - .../triton_kernel/mla_att/__init__.py | 0 .../mla_att/decode_att/__init__.py | 1 - .../mla_att/decode_att/gqa_flash_decoding.py | 164 --- .../mla_att/prefill_att/__init__.py | 9 - lightllm/common/req_manager/base.py | 9 +- .../layer_infer/transformer_layer_infer.py | 42 +- lightllm/models/bloom/model.py | 17 +- .../layer_infer/transformer_layer_infer.py | 47 +- lightllm/models/deepseek2/model.py | 22 +- .../layer_infer/transformer_layer_infer.py | 73 +- lightllm/models/deepseek3_2/model.py | 15 +- lightllm/models/deepseek_v4/infer_struct.py | 13 +- .../layer_infer/transformer_layer_infer.py | 99 +- lightllm/models/deepseek_v4/model.py | 62 +- .../layer_infer/transformer_layer_infer.py | 38 +- lightllm/models/gemma4/model.py | 37 +- .../glm5_next/attention/fp8_sparse_mla.py | 1 + lightllm/models/glm5_next/config.py | 6 +- .../layer_infer/transformer_layer_infer.py | 129 +- lightllm/models/glm5_next/model.py | 73 +- .../context_flashattention_nopad_fp8.py | 744 ++++++++++++ lightllm/models/glm5_next_mtp/model.py | 36 +- .../layer_infer/transformer_layer_infer.py | 67 +- lightllm/models/gpt_oss/model.py | 11 +- .../layer_infer/transformer_layer_infer.py | 47 +- .../layer_infer/transformer_layer_infer.py | 17 +- lightllm/models/neo_chat/model.py | 50 +- .../layer_infer/transformer_layer_infer.py | 17 +- lightllm/models/neo_chat_moe/model.py | 53 +- lightllm/models/qwen3_dflash/model.py | 9 - .../layer_infer/transformer_layer_infer.py | 30 +- .../layer_infer/transformer_layer_infer.py | 84 +- lightllm/models/qwen3next/model.py | 40 +- lightllm/server/api_cli.py | 44 +- lightllm/server/api_start.py | 68 +- lightllm/server/core/objs/start_args_type.py | 13 +- .../model_infer/mode_backend/base_backend.py | 95 +- .../mode_backend/chunked_prefill/impl.py | 43 +- .../mode_backend/dp_backend/impl.py | 37 +- .../mode_backend/multi_level_kv_cache.py | 48 +- .../model_infer/mtp_speculative/utils.py | 34 +- lightllm/utils/backend_validator.py | 5 - lightllm/utils/envs_utils.py | 12 +- lightllm/utils/start_utils.py | 8 +- lightllm/utils/torch_memory_saver_utils.py | 13 +- test/att_resources/validate_model_cache.py | 19 +- .../kernels/benchmark_replayssm_service.py | 2 +- .../static_inference/static_benchmark.py | 37 +- .../common/att_resources/manager_fixtures.py | 18 +- .../common/att_resources/model_helpers.py | 34 +- .../att_resources/test_backend_startup.py | 48 +- .../att_resources/test_bounded_state.py | 14 +- .../common/att_resources/test_capacity.py | 50 +- .../att_resources/test_dsv4_integration.py | 58 +- .../att_resources/test_fp8_compute_binding.py | 140 +-- .../att_resources/test_full_att_cpu_kernel.py | 12 +- .../common/att_resources/test_full_state.py | 24 +- .../att_resources/test_gqa_layout_identity.py | 57 - .../test_hybrid_layer_mapping.py | 4 +- .../test_layer_state_selection.py | 20 +- .../test_model_initialization.py | 15 +- .../test_model_storage_definition.py | 33 +- .../att_resources/test_mtp_state_commit.py | 33 +- .../att_resources/test_no_attention_draft.py | 21 +- .../test_quantized_full_state.py | 19 +- .../test_validate_model_cache.py | 9 +- .../common/attention/test_backend_options.py | 130 ++ unit_tests/common/attention/test_contracts.py | 89 ++ .../common/attention/test_dense_numerics.py | 201 +++ .../common/attention/test_platform_runtime.py | 137 +++ .../common/attention/test_provider_kind.py | 88 ++ unit_tests/common/attention/test_state.py | 69 ++ .../basemodel/attention/fa3/test_neo.py | 27 +- .../flashinfer/test_decode_graph_state.py | 100 -- .../flashinfer/test_decode_wrapper_init.py | 176 ++- .../flashinfer/test_mla_cuda_graph_plan.py | 28 +- .../attention/flashinfer/test_mla_nope.py | 3 +- .../attention/flashinfer/test_workspace.py | 71 +- .../basemodel/attention/linear/test_gdn.py | 431 ++----- .../basemodel/attention/test_dsv4_flashmla.py | 111 +- .../attention/test_fp8_logical_layer.py | 86 +- .../basemodel/attention/test_fp8_page_size.py | 30 +- .../basemodel/attention/test_gated_qkv.py | 23 + .../basemodel/attention/test_lazy_import.py | 234 +--- .../attention/test_state_lifetime.py | 185 +-- .../basemodel/test_cuda_graph_autotune.py | 88 +- .../common/basemodel/test_model_output.py | 42 +- .../basemodel/test_moe_route_info_manager.py | 12 +- .../common/basemodel/test_overlap_utils.py | 7 +- .../test_prefill_cuda_graph_state.py | 53 +- .../gqa/mtp_diverse/test_mtp_diverse.py | 6 +- .../mtp_diverse/test_mtp_diverse_autotune.py | 35 +- .../gqa/mtp_diverse/test_stage2_autotune.py | 15 +- .../gqa/test_flash_decoding_autotune.py | 38 +- .../decode_att/int4kv/test_int4kv_autotune.py | 27 +- .../test_int8kv_flash_decoding_diverse.py | 29 +- ...st_int8kv_flash_decoding_diverse_stage1.py | 1 + ...st_int8kv_flash_decoding_diverse_stage2.py | 1 + ...st_int8kv_flash_decoding_diverse_stage3.py | 1 + .../decode_att/int8kv/test_normal_autotune.py | 29 +- .../att/prefill_att/test_fa3_neo.py | 2 +- .../linear_att/test_causal_conv1d_mtp.py | 2 + .../linear_att/test_replayssm.py | 1 + .../decode_att/test_gqa_flash_decoding.py | 17 +- .../triton_kernel/test_diverse_utils.py | 17 +- .../common/test_deepseek4_paged_cache.py | 170 ++- .../common/test_linear_state_contract.py | 22 +- unit_tests/models/deepseek_v4/test_model.py | 7 +- .../models/deepseek_v4/test_workspace.py | 15 +- unit_tests/models/glm5_next/test_cache.py | 6 +- .../glm5_next/test_direct_nsa_indices.py | 292 ++--- unit_tests/models/glm5_next/test_fp8_kv.py | 1074 ++++++++++------- .../models/glm5_next/test_kda_conv_qkv.py | 46 +- unit_tests/models/glm5_next/test_kernels.py | 43 +- unit_tests/models/glm5_next/test_layers.py | 75 +- unit_tests/models/glm5_next/test_mtp.py | 86 +- unit_tests/models/glm5_next/test_pd_cache.py | 19 +- .../glm5_next/test_prefill_workspace.py | 159 +-- .../models/sensenova_loop/test_gdn_state.py | 96 +- unit_tests/server/core/objs/test_req.py | 12 +- .../server/core/objs/test_shm_req_manager.py | 17 +- .../mode_backend/test_multi_level_kv_cache.py | 26 +- .../model_infer/test_cpu_resource_layout.py | 26 +- unit_tests/server/test_api_start_deepep.py | 12 +- unit_tests/server/test_hybrid_start_args.py | 1 + unit_tests/test_static_benchmark_paged_kv.py | 49 +- unit_tests/utils/test_sgl_utils.py | 28 +- unit_tests/utils/test_speculative_utils.py | 173 ++- 214 files changed, 5512 insertions(+), 9980 deletions(-) create mode 100644 lightllm/common/attention_state/commit.py create mode 100644 lightllm/common/attention_state/full_att/gqa/ops.py create mode 100644 lightllm/common/basemodel/attention.py delete mode 100644 lightllm/common/basemodel/attention/__init__.py delete mode 100644 lightllm/common/basemodel/attention/base_att.py delete mode 100644 lightllm/common/basemodel/attention/create_linear_utils.py delete mode 100644 lightllm/common/basemodel/attention/create_utils.py delete mode 100644 lightllm/common/basemodel/attention/fa3/__init__.py delete mode 100644 lightllm/common/basemodel/attention/fa3/fp.py delete mode 100644 lightllm/common/basemodel/attention/fa3/fp8.py delete mode 100644 lightllm/common/basemodel/attention/fa3/mla.py delete mode 100644 lightllm/common/basemodel/attention/fa3/neo.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/__init__.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/env_utils.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/fp.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/fp8.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/mla.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/utils.py delete mode 100644 lightllm/common/basemodel/attention/import_utils.py delete mode 100644 lightllm/common/basemodel/attention/linear/__init__.py delete mode 100644 lightllm/common/basemodel/attention/linear/flashqla.py delete mode 100644 lightllm/common/basemodel/attention/linear/gdn.py delete mode 100644 lightllm/common/basemodel/attention/linear/kda.py delete mode 100644 lightllm/common/basemodel/attention/linear/replay_gdn.py delete mode 100644 lightllm/common/basemodel/attention/linear/replay_kda.py delete mode 100644 lightllm/common/basemodel/attention/linear/triton.py delete mode 100644 lightllm/common/basemodel/attention/nsa/__init__.py delete mode 100644 lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py delete mode 100644 lightllm/common/basemodel/attention/nsa/flashmla_sparse.py delete mode 100644 lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py delete mode 100644 lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py delete mode 100644 lightllm/common/basemodel/attention/nsa/glm5_next.py delete mode 100644 lightllm/common/basemodel/attention/triton/__init__.py delete mode 100644 lightllm/common/basemodel/attention/triton/fp.py delete mode 100644 lightllm/common/basemodel/attention/triton/int4kv.py delete mode 100644 lightllm/common/basemodel/attention/triton/int8kv.py delete mode 100644 lightllm/common/basemodel/attention/triton/mla.py delete mode 100644 lightllm/common/basemodel/attention/triton/neo.py create mode 100644 lightllm/common/basemodel/attention_batch.py create mode 100644 lightllm/common/basemodel/attention_runtime.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/mha/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/att/prefill_att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/mla_att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py delete mode 100644 lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py create mode 100644 lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py create mode 100644 unit_tests/common/attention/test_backend_options.py create mode 100644 unit_tests/common/attention/test_provider_kind.py create mode 100644 unit_tests/common/attention/test_state.py delete mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py create mode 100644 unit_tests/common/basemodel/attention/test_gated_qkv.py diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 7edb0c9701..8cf4eb5e51 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -127,7 +127,19 @@ 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；每次 forward 通过 AttentionExecution 查找当前 Context，不缓存批次 Context。prepare/update 仍由 Runtime 管理，indexer 可用 `.context(infer_state)` 直接访问： + +```python +output = self.decode_attention.forward(infer_state, 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。 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index cd71541598..3f08906468 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -127,7 +127,19 @@ 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; each forward resolves the current Context through AttentionExecution without caching it. Runtime still manages prepare/update; indexers can access the Context through `.context(infer_state)`: + +```python +output = self.decode_attention.forward(infer_state, 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. diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index 6a96e76dc9..013a8b2682 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/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") diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index 82eb4026c5..c6d58a3265 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/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" diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 2d9540645a..a0d3ad024d 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/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 ( diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index 11d7425cd9..17f226182c 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/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 的资源、层视图和运行态核心能力。""" diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index 68097905d5..a32a52199f 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/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" diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 8897218fbc..c86a2d72a7 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/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 ( diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py index 5afdfab160..70a1d22eb5 100644 --- a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py +++ b/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_() diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py index 5d9a0ffe12..472acec80c 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py @@ -2,9 +2,8 @@ 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 @@ -12,8 +11,8 @@ 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_() diff --git a/lightllm/common/attention_state/commit.py b/lightllm/common/attention_state/commit.py new file mode 100644 index 0000000000..711d6979f9 --- /dev/null +++ b/lightllm/common/attention_state/commit.py @@ -0,0 +1,34 @@ +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class BoundedStateCommit: + role: str + request_indices: torch.Tensor + mtp_indices: torch.Tensor + request_generations: torch.Tensor + + @classmethod + def capture(cls, role, request_indices, mtp_indices, generations): + requests = request_indices.clone() + return cls(role, requests, mtp_indices.clone(), generations[requests.long()].clone()) + + @classmethod + def concatenate(cls, commits): + if not commits or any(commit.role != commits[0].role for commit in commits): + raise ValueError("MTP outputs must belong to the same model role") + return cls( + commits[0].role, + *( + torch.cat([getattr(commit, name) for commit in commits]) + for name in ("request_indices", "mtp_indices", "request_generations") + ), + ) + + def validate_generation(self, generations): + torch._assert_async( + torch.all(generations[self.request_indices.long()] == self.request_generations), + "MTP output refers to a released or reused request slot", + ) diff --git a/lightllm/common/attention_state/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py index d3f00430b0..d55acc1639 100644 --- a/lightllm/common/attention_state/full_att/base_state.py +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -1,12 +1,31 @@ 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 FullAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): """单种 full-attention 存储方案的数据接口，不分配 token 索引或维护请求映射。 - 所有表示 token 索引的参数都须为调用方准备的 CUDA Tensor，位于对应存储设备上； + 所有表示 token 索引的参数都须为调用方准备的同设备 Tensor，位于对应存储设备上； State 不转换设备、dtype 或连续布局，int32/int64 由具体算子支持，DP 目标索引须为 int64。 - CPU 页位于主机内存不改变 token 索引的 CUDA 契约。 + 传输操作的设备限制由具体实现声明，CPU 页位于主机内存不改变索引与存储同设备的契约。 """ + + 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") diff --git a/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py index 21ce083697..13001eea22 100644 --- a/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py +++ b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py @@ -2,21 +2,20 @@ from __future__ import annotations -from functools import lru_cache - -from typing import TYPE_CHECKING, List from dataclasses import dataclass +from functools import lru_cache from math import lcm, prod +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args -from ...utils import deepcopy_result from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result from ..base_state import FullAttentionState -from ..packed_mla.packed_transfer import transfer_packed_pages from ..dp_transfer import copy_packed_tokens +from ..packed_mla.packed_transfer import transfer_packed_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -103,6 +102,8 @@ class CompressedAttentionState(FullAttentionState): 当前实现固定为单个共享 KV head，各 TP rank 完整复制；不执行 CSA/HCA 的历史选择。 """ + storage_layout = ("compressed_fp8_e8m0_token_row", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py index fe184b4110..13469ec54e 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -1,20 +1,23 @@ from __future__ import annotations +from dataclasses import dataclass from functools import lru_cache -from typing import TYPE_CHECKING, List from math import prod -from dataclasses import dataclass +from typing import TYPE_CHECKING, List import torch +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.platform import Platform, get_device_runtime from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from .gqa_transfer import transfer_gqa_pages from ..dp_transfer import copy_gqa_tokens +from .gqa_transfer import transfer_gqa_pages +from .ops import get_kv_operations if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -78,6 +81,8 @@ class GQAFullAttentionState(FullAttentionState): query 与 KV 的分组计算由 attention backend 处理，因此三种情况共用此 State。 """ + storage_layout = ("gqa_separate_token_head_dim", 1) + # 保留既有 K/V 编码的稳定传输布局标识。 # ======================================================================== @@ -103,6 +108,8 @@ def __init__( self.page_size = page_size self.device = torch.device(device=device) self.capacity = 0 + self.zero_initialize_cache = False + self.ops = None self.k_storage: torch.Tensor | None = None self.v_storage: torch.Tensor | None = None @@ -114,8 +121,13 @@ def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity shape = self._storage_shape() - self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) - self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + runtime = get_device_runtime(self.device) + self.ops = get_kv_operations(runtime) + self.k_storage = runtime.allocate(shape, self.storage_dtype) + self.v_storage = runtime.allocate(shape, self.storage_dtype) + if self.zero_initialize_cache: + self.k_storage.zero_() + self.v_storage.zero_() def get_token_bytes(self) -> int: """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" @@ -140,25 +152,19 @@ def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Ten return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: - """通过 CUDA 算子将形状匹配的 K/V 数据写入指定物理 token 位置。""" - assert token_indices.is_cuda, "token indices must reside on CUDA" + """通过已绑定的设备实现将 K/V 数据写入指定物理 token 位置。""" storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) - if not len(token_indices): + assert token_indices.device == kv.device == self.k_storage.device, "KV writes require same-device tensors" + assert token_indices.ndim == 1 and token_indices.dtype in (torch.int32, torch.int64) + assert kv.ndim == 3 and kv.shape[1:] == (2 * self.config.head_num, self.config.head_dim) + assert kv.shape[0] == token_indices.numel() and kv.dtype == self.storage_dtype + if not token_indices.numel(): return - assert self.k_storage.is_cuda, "KV writes require CUDA storage" - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv - - # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 - destindex_copy_kv( - K=kv[:, : self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index] - ) - destindex_copy_kv( - K=kv[:, self.config.head_num :], DestLoc=token_indices, Out=self.v_storage[storage_layer_index] - ) + self.ops.write(self.k_storage[storage_layer_index], self.v_storage[storage_layer_index], token_indices, kv) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" - assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert src_indices.device == dst_indices.device == self.k_storage.device, "token indices must match storage" assert ( src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() ), "token copies require equal one-dimensional index tensors" @@ -166,26 +172,17 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten torch.int32, torch.int64, ), "token indices must be integer tensors" - assert self.k_storage.is_cuda, "token copies require CUDA storage" - if not src_indices.numel(): - return - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer - - # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 - # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 - # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 - # TODO：将 K/V 的同索引 token 复制融合为一次 kernel 发射， - # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 - copy_kv_buffer_to_kv_buffer( - src_mem_index=src_indices, - dst_mem_index=dst_indices, - kv_buffer=self.k_storage, - ) - copy_kv_buffer_to_kv_buffer( - src_mem_index=src_indices, - dst_mem_index=dst_indices, - kv_buffer=self.v_storage, - ) + if src_indices.numel(): + self.ops.copy(self.k_storage, self.v_storage, src_indices, dst_indices) + + def require_zero_initialization(self): + self.zero_initialize_cache = True + + def check_execution(self, caps, execution) -> None: + if caps.platform not in (Platform.CUDA, Platform.CPU): + raise UnsupportedAttention("GQA state operations are unavailable on this platform") + if caps.platform is not Platform.CUDA and execution.transfer_paths: + raise UnsupportedAttention("GQA CPU/DP/PD transfer operations currently require CUDA") # ======================================================================== # 三、CPU Cache 数据交换 diff --git a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py index 6629d54aab..e8dce6e7b9 100644 --- a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py @@ -1,20 +1,20 @@ from __future__ import annotations -from functools import lru_cache -from typing import TYPE_CHECKING, List from dataclasses import dataclass +from functools import lru_cache from math import prod +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from .group_quantized_transfer import transfer_group_quantized_pages from ..dp_transfer import copy_group_quantized_tokens +from .group_quantized_transfer import transfer_group_quantized_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -81,6 +81,8 @@ def __post_init__(self) -> None: class GroupQuantizedFullAttentionState(FullAttentionState): """管理 INT4/INT8 KV 与逐 token 分组 scale，复制和传输始终覆盖二者。""" + storage_layout = ("gqa_group_quantized_token_head_dim", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/gqa/ops.py b/lightllm/common/attention_state/full_att/gqa/ops.py new file mode 100644 index 0000000000..5039de5aef --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/ops.py @@ -0,0 +1,44 @@ +"""GQA 的物理写入/复制实现；层映射、容量和 token 生命周期仍由 State 管理。""" + +import torch + +from lightllm.platform import Platform + + +class TritonKVOperations: + @staticmethod + def write(key, value, indices, kv): + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + heads = key.shape[1] + destindex_copy_kv(K=kv[:, :heads], DestLoc=indices, Out=key) + destindex_copy_kv(K=kv[:, heads:], DestLoc=indices, Out=value) + + @staticmethod + def copy(key, value, src, dst): + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + copy_kv_buffer_to_kv_buffer(src_mem_index=src, dst_mem_index=dst, kv_buffer=key) + copy_kv_buffer_to_kv_buffer(src_mem_index=src, dst_mem_index=dst, kv_buffer=value) + + +class TorchKVOperations: + @staticmethod + def write(key, value, indices, kv): + indices = indices.to(torch.int64) + key.index_copy_(0, indices, kv[:, : key.shape[1]]) + value.index_copy_(0, indices, kv[:, key.shape[1] :]) + + @staticmethod + def copy(key, value, src, dst): + src, dst = src.to(torch.int64), dst.to(torch.int64) + key.index_copy_(1, dst, key.index_select(1, src)) + value.index_copy_(1, dst, value.index_select(1, src)) + + +def get_kv_operations(runtime): + if runtime.platform is Platform.CUDA: + return TritonKVOperations + if runtime.platform is Platform.CPU: + return TorchKVOperations + raise ValueError(f"GQA KV writes are not implemented on {runtime.platform.value}") diff --git a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py index 5cb1acd778..bb7265e718 100644 --- a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py @@ -1,22 +1,21 @@ from __future__ import annotations -from functools import lru_cache - -from typing import TYPE_CHECKING, List from collections.abc import Mapping from dataclasses import dataclass +from functools import lru_cache from math import prod +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result from ..base_state import FullAttentionState -from .static_fp8_transfer import transfer_static_fp8_pages from ..dp_transfer import copy_gqa_tokens +from .static_fp8_transfer import transfer_static_fp8_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -72,6 +71,8 @@ def __post_init__(self) -> None: class StaticFP8FullAttentionState(FullAttentionState): """使用静态标定参数的 FP8 KV 存储，固定 scale 随 State 初始化，不逐 token 搬运。""" + storage_layout = ("gqa_static_fp8_token_head_dim", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/mla/mla_state.py b/lightllm/common/attention_state/full_att/mla/mla_state.py index a818f5af3a..0bfcb593d2 100644 --- a/lightllm/common/attention_state/full_att/mla/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla/mla_state.py @@ -1,19 +1,19 @@ from __future__ import annotations +from dataclasses import dataclass from functools import lru_cache -from typing import TYPE_CHECKING, List from math import prod -from dataclasses import dataclass +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from ..packed_mla.packed_transfer import transfer_packed_pages from ..dp_transfer import copy_packed_tokens +from ..packed_mla.packed_transfer import transfer_packed_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -51,6 +51,8 @@ def __post_init__(self) -> None: class MLAFullAttentionState(FullAttentionState): """存储复制到各 TP rank 的 MLA 压缩 latent，使用单个 KV 字段。""" + storage_layout = ("mla_token_latent", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py index 4f018ef215..257d2192db 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py @@ -2,21 +2,20 @@ from __future__ import annotations -from functools import lru_cache - from dataclasses import dataclass +from functools import lru_cache from typing import ClassVar import torch from lightllm.utils.envs_utils import get_env_start_args -from ...utils import deepcopy_result from ...types import AxisMapping, StateTransferLayout, TransferRegion +from ...utils import deepcopy_result from ..base_state import FullAttentionState from ..dp_transfer import copy_packed_tokens -from ..packed_mla.packed_transfer import transfer_packed_pages from ..packed_mla.packed_mla_state import PackedMLALayerView +from ..packed_mla.packed_transfer import transfer_packed_pages @dataclass(frozen=True) @@ -58,6 +57,8 @@ def __post_init__(self) -> None: class FP8E8M0PackedMLAState(FullAttentionState): """独立的 group64/E8M0 MLA 实现，不继承其他具体 State 或动态选择编码。""" + storage_layout = ("mla_packed_fp8_e8m0_token_row", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py index 99c920fe26..0ef03d4e83 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py @@ -1,20 +1,20 @@ from __future__ import annotations -from functools import lru_cache -from typing import TYPE_CHECKING, List, Literal from dataclasses import dataclass +from functools import lru_cache from math import prod +from typing import TYPE_CHECKING, List, Literal import torch from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from .packed_mla_transfer import transfer_packed_mla_pages from ..dp_transfer import copy_packed_tokens +from .packed_mla_transfer import transfer_packed_mla_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -146,6 +146,8 @@ class PackedMLALayerView: class PackedMLAFullAttentionState(FullAttentionState): """在同一物理 token 行中存储 MLA latent、量化参数及稀疏 indexer key/scale。""" + storage_layout = ("mla_packed_token_row", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index c0a85c367e..a72b245050 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -1,27 +1,27 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import json +import weakref from functools import lru_cache -from typing import List +from typing import TYPE_CHECKING, List import torch from lightllm.common.attention_state.types import PageKind - -from lightllm.utils.dist_utils import get_current_device_id -from lightllm.utils.envs_utils import get_env_start_args +from lightllm.platform import get_device_runtime from lightllm.utils.log_utils import init_logger -from .allocators import StatePageAllocator, FullAttKVAllocator +from .allocators import FullAttKVAllocator, StatePageAllocator +from .bounded_att import BoundedAttentionState from .cpu_cache_operator import CpuCacheOperator from .dp_shared_operator import DPSharedOperator from .full_att import FullAttentionState -from .bounded_att import BoundedAttentionState from .pd_operator import PDOperator from .utils import MODEL_ROLES if TYPE_CHECKING: from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs @@ -39,6 +39,8 @@ class AttentionStateManager: 和 allocator，不复制 State 的算法逻辑。 """ + storage_generation = 0 + # ======================================================================== # 一、基础管理：共享分配器、资源初始化与聚合查询 # ======================================================================== @@ -138,7 +140,6 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: 排序，State 内保留 logical_layer_groups 顺序。非静态 FP8 的 full State 同样占用 对应物理行区间，共享逻辑层只占一行；CPU 布局创建不需要调用本流程。 """ - import json from .full_att.gqa.static_fp8_state import StaticFP8FullAttentionState full_states = self.get_full_attention_states() @@ -180,6 +181,16 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: "Full-attention calibration bound: states=%d, source=%s", len(calibration_states), path or "unit scales" ) + def register_graph(self, graph): + if not hasattr(self, "_graphs"): + self._graphs = weakref.WeakSet() + self._graphs.add(graph) + + def close_graphs(self): + """释放状态前先排空依赖它的捕获实例，避免旧地址在后台继续被使用。""" + for graph in tuple(getattr(self, "_graphs", ())): + graph.close() + def initialize_storage( self, args: StartArgs | EasyDict, @@ -194,6 +205,7 @@ def initialize_storage( 依次执行各模型的运行时初始化，由模型读取 allocator 中的最终容量。 """ assert self._state_registration_finished, "State registration must finish before resource initialization" + assert self.full_att_allocator is None, "resources already initialized" # 阶段一：规划共享 token 容量，并向下对齐到完整 KV 页。 # 未指定容量时，根据剩余显存评估；固定运行态的显存已经在前面占用。 @@ -204,9 +216,10 @@ def initialize_storage( pd_reserved_bytes = self.get_pd_operator().get_page_bytes() * args.pd_kv_page_num if token_capacity is None: # 权重和固定运行态已分配，用当前设备剩余显存扣除保留比例及 PD buffer。 - torch.cuda.empty_cache() - gpu_device = torch.device("cuda", get_current_device_id()) - free_bytes, total_bytes = torch.cuda.mem_get_info(gpu_device) + runtime = getattr(self, "device_runtime", None) or get_device_runtime(device) + runtime.empty_cache() + gpu_device = runtime.device + free_bytes, total_bytes = runtime.memory_info() budget = int(free_bytes - total_bytes * (1 - mem_fraction)) - pd_reserved_bytes # 按每页开销直接计算容量，预留一个 HOLD 页；静态量化参数占比较小，不单独扣减。 @@ -229,6 +242,8 @@ def initialize_storage( token_capacity = token_capacity // args.page_size * args.page_size assert token_capacity > 0, "attention capacity must contain at least one complete page" + self.close_graphs() + self.storage_generation += 1 # 阶段二：管理器为已创建的 State 分配实际存储与共享 allocator。 # State 对象保持不变，查询缓存继续引用相同对象；固定运行态不重复分配。 big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num @@ -331,6 +346,18 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # ======================================================================== # 空前缀初始化：只处理 bounded attention；无此类 State 时可为空操作。 + def snapshot_capture_state(self, req_indices: torch.Tensor, role: str): + states = self.get_bounded_attention_states(role) + if not states: + return () + requests = req_indices.unique().cpu().tolist() + return tuple((view, view.clone()) for state in states for view in state.capture_snapshot_views(requests)) + + @staticmethod + def restore_capture_state(snapshot): + for target, saved in snapshot: + target.copy_(saved) + def reset_bounded_att_state(self, req_indices: List[int]) -> None: """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页。""" for state in self.get_bounded_attention_states(): diff --git a/lightllm/common/basemodel/attention.py b/lightllm/common/basemodel/attention.py new file mode 100644 index 0000000000..1b64b2c29f --- /dev/null +++ b/lightllm/common/basemodel/attention.py @@ -0,0 +1,294 @@ +from dataclasses import dataclass, replace + +import torch + +from lightllm.common.attention.api import ( + AttentionBackend, + AttentionFamily, + AttentionPhase, + CacheLayout, + ContextCapacity, + LayerAttentionBinding, + LayerSpec, +) +from lightllm.common.attention.api.execution import ExecutionRequirements +from lightllm.common.attention.api.inputs import DenseParams, GDNParams, KDAParams, MLAParams +from lightllm.common.attention.backends._workspace import AttentionWorkspace + + +def describe_execution_requirements(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 ExecutionRequirements( + 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)), + args.enable_ep_moe, + frozenset(transfers), + ) + + +@dataclass(frozen=True) +class AttentionGroup: + prefill_backend: AttentionBackend + decode_backend: AttentionBackend + prefill_spec: LayerSpec + decode_spec: LayerSpec + + def backend_for(self, phase): + if phase is AttentionPhase.PREFILL: + return self.prefill_backend + if phase is AttentionPhase.DECODE: + return self.decode_backend + raise ValueError(f"unknown attention phase: {phase}") + + +def _alloc_tensor(shape, dtype, device="cuda"): + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + return g_cache_manager.alloc_tensor(shape, dtype, device=device) + + +def describe_attention_layer(model, layer, family, phase): + config = model.config + linear = family in (AttentionFamily.GDN, AttentionFamily.KDA) + if family is AttentionFamily.GDN: + heads = config["linear_num_key_heads"] // model.tp_world_size_ + head_dim = config["linear_key_head_dim"] + value_dim = config["linear_value_head_dim"] + kv_heads = heads + elif family is AttentionFamily.KDA: + heads = config["linear_attn_config"]["num_heads"] // model.tp_world_size_ + head_dim = value_dim = config["linear_attn_config"]["head_dim"] + kv_heads = heads + else: + heads = layer.tp_q_head_num_ + if family in (AttentionFamily.MLA, AttentionFamily.SPARSE_MLA): + head_dim = getattr(layer, "qk_nope_head_dim", 0) + getattr(layer, "qk_rope_head_dim", 0) + head_dim = head_dim or layer.head_dim_ + value_dim = getattr(layer, "kv_lora_rank", getattr(layer, "v_head_dim", head_dim)) + kv_heads = 1 + else: + head_dim = layer.head_dim_ + value_dim = head_dim + kv_heads = layer.tp_k_head_num_ + encoding = family.value if linear else model.args.llm_kv_type + ssm_dtype = {"bfloat16": torch.bfloat16, "float32": torch.float32}.get(model.args.linear_att_ssm_data_type) + states = tuple( + state + for state in ( + *model.att_state_manager.get_full_attention_states(model.att_state_role), + *model.att_state_manager.get_bounded_attention_states(model.att_state_role), + ) + if state.contains_logical_layer(layer.layer_num_) + ) + if not states: + raise ValueError(f"layer {layer.layer_num_} has no registered attention storage") + schema = tuple(state.describe_layout() for state in states) + cache = CacheLayout( + encoding, + model.args.page_size, + getattr(states[0], "storage_dtype", getattr(states[0].config, "dtype", model.data_type)), + kv_heads, + head_dim, + value_dim, + ssm_dtype if linear else None, + storage_schema=schema, + quant_group_size=getattr(states[0].config, "group_size", 0), + scale_mode=getattr(states[0].config, "scale_mode", "none"), + ) + draft_step = model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model) + return LayerSpec( + family=family, + phase=phase, + dtype=model.data_type, + cache=cache, + num_q_heads=heads, + head_size=head_dim, + value_head_size=value_dim, + sliding_window=layer.attention_window(phase), + use_alibi=layer.attention_uses_alibi, + use_sink=layer.attention_uses_sink, + causal=not (model.is_mtp_draft_model and model.args.mtp_mode in ("dspark", "dflash", "dflash2")), + use_mtp=draft_step > 0 and phase is AttentionPhase.DECODE, + dynamic_verify=draft_step > 0 and not model.is_mtp_draft_model and model.args.mtp_dynamic_verify, + variant=layer.attention_variant, + softmax_scale=layer.softmax_scale if family in (AttentionFamily.MLA, AttentionFamily.SPARSE_MLA) else None, + ) + + +def describe_attention_layers(model, family, phase): + return tuple( + describe_attention_layer(model, layer, family, phase) + for layer in model.layers_infer + if layer.get_attention_family() is family + ) + + +def create_layer_attention_params(model, layer, weight, family): + manager, role, layer_id = model.att_state_manager, model.att_state_role, layer.layer_num_ + if family in (AttentionFamily.GDN, AttentionFamily.KDA): + state = manager.get_state(role, family.value, layer_id) + conv, ssm = state.get_layer_view(layer_id) + if family is AttentionFamily.GDN: + return GDNParams( + conv=conv, + ssm=ssm, + conv_weight=weight.linear_conv1d.mm_param.weight, + conv_bias=weight.linear_conv1d.bias, + a_log=weight.linear_A_log.weight, + dt_bias=weight.linear_dt_bias.weight, + mtp_state_index=state.mtp_state_index, + replay_state=state if model.args.ssm_state_mode == "replay" else None, + ) + return KDAParams( + conv=conv, + ssm=ssm, + conv_weight=weight.get_merged_kda_conv_weight(), + a_log=weight.linear_A_log.weight, + dt_bias=weight.linear_dt_bias.weight, + mtp_state_index=state.mtp_state_index, + replay_state=state if model.args.ssm_state_mode == "replay" else None, + ) + if family is AttentionFamily.DENSE: + state = manager.get_state(role, "gqa", layer_id) + key, value = layer.attention_kv_view(state) + query_scale = key_scale = value_scale = None + if isinstance(key, tuple): + key, key_scale = key + value, value_scale = value + if model.args.llm_kv_type in ("fp8kv_sph", "fp8kv_spt"): + index = state.get_storage_layer_index(layer_id) + if model.args.llm_kv_type == "fp8kv_sph": + heads = state.config.head_num + key_scale, value_scale = state.scales[index, :heads], state.scales[index, heads:] + query_scale = state.q_scales[index] if state.q_scales is not None else None + else: + key_scale, value_scale = map(float, state.cpu_scales[index]) + return DenseParams( + key=key, + value=value, + alibi=getattr(weight, "tp_alibi", None), + sink=weight.attn_sinks.weight if layer.attention_uses_sink else None, + query_scale=query_scale, + key_scale=key_scale, + value_scale=value_scale, + ) + state = manager.get_state(role, "mla", layer_id) + view = state.get_layer_view(layer_id) + return MLAParams( + view.kv if family is AttentionFamily.SPARSE_MLA else view, + layer.kv_lora_rank, + getattr(layer, "qk_rope_head_dim", 0), + ) + + +def initialize_attention_groups(model): + if model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): + model.attention_requirements = ExecutionRequirements() + model.attention_groups = () + model._attention_layer_groups = (None,) * len(model.layers_infer) + model._attention_workspaces = {} + return + model.attention_requirements = describe_execution_requirements(model) + caps = model.device_runtime.caps + for state in ( + *model.att_state_manager.get_full_attention_states(model.att_state_role), + *model.att_state_manager.get_bounded_attention_states(model.att_state_role), + ): + state.check_execution(caps, model.attention_requirements) + from lightllm.common.attention.selection import select_attention_backend + + groups, keys, routes = [], {}, [] + for layer in model.layers_infer: + family = layer.get_attention_family() + specs = tuple(describe_attention_layer(model, layer, family, phase) for phase in AttentionPhase) + if specs[0].cache != specs[1].cache: + raise ValueError("prefill/decode physical cache layouts differ") + if specs not in keys: + backends = [] + for spec in specs: + candidates, priority = model.attention_backend_candidates(family, spec.phase) + linear = family in (AttentionFamily.GDN, AttentionFamily.KDA) + option = f"llm_{spec.phase.value}_{'linear_' if linear else ''}att_backend" + cls = select_attention_backend( + caps=caps, + execution=model.attention_requirements, + ssm_state_mode=model.args.ssm_state_mode, + requested=getattr(model.args, option), + candidates=candidates, + priority=priority, + specs=(spec,), + ) + backends.append(cls.create(model, spec)) + keys[specs] = len(groups) + groups.append(AttentionGroup(*backends, *specs)) + routes.append(keys[specs]) + model.attention_groups = tuple(groups) + model._attention_layer_groups = tuple(routes) + model._attention_workspaces = {} + + +def reserve_attention_resources(model): + execution = model.attention_requirements + for group_id, group in enumerate(model.attention_groups): + for phase in AttentionPhase: + backend = group.backend_for(phase) + capacity = ContextCapacity( + phase, execution.max_batch_size, execution.max_batch_size, execution.max_seq_len, 1 + ) + resources = backend.resource_requirements(capacity, capture=phase in execution.graph_phases) + for slot in range(execution.microbatch_count): + get_attention_workspace(model, group_id, slot, model.device_runtime.device).reserve(resources) + if resources.zero_initialize_cache: + for layer, route in zip(model.layers_infer, model._attention_layer_groups): + if route != group_id: + continue + for state in model.att_state_manager.get_full_attention_states(model.att_state_role): + if state.contains_logical_layer(layer.layer_num_): + state.require_zero_initialization() + + +def initialize_layer_attention(model): + for layer, weight, group_id in zip(model.layers_infer, model.trans_layers_weight, model._attention_layer_groups): + if group_id is None: + continue + group = model.attention_groups[group_id] + params = layer.create_attention_params(model, weight, group.prefill_spec.family) + layer.prefill_attention = LayerAttentionBinding( + group_id=group_id, + logical_layer_id=layer.layer_num_, + spec=group.prefill_spec, + params=params, + backend=group.prefill_backend, + alloc_tensor=_alloc_tensor, + ) + layer.decode_attention = LayerAttentionBinding( + group_id=group_id, + logical_layer_id=layer.layer_num_, + spec=group.decode_spec, + params=params, + backend=group.decode_backend, + alloc_tensor=_alloc_tensor, + ) + model._attention_storage_generation = model.att_state_manager.storage_generation + + +def get_attention_workspace(model, group_id, slot, device): + # overlap 的 microbatch 可能同时执行，scratch buffer 不能只按 group 共享。 + key = group_id, slot + if key not in model._attention_workspaces: + model._attention_workspaces[key] = AttentionWorkspace(runtime=model.device_runtime) + return model._attention_workspaces[key] diff --git a/lightllm/common/basemodel/attention/__init__.py b/lightllm/common/basemodel/attention/__init__.py deleted file mode 100644 index d728a4ee9c..0000000000 --- a/lightllm/common/basemodel/attention/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -from lightllm.common.attention.api import AttentionPhase, ExecutionRequirements - -from .base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState -from .create_utils import ( - get_decode_att_backend_class, - get_mla_decode_att_backend_class, - get_mla_prefill_att_backend_class, - get_neo_prefill_att_backend_class, - get_nsa_decode_att_backend_class, - get_nsa_prefill_att_backend_class, - get_prefill_att_backend_class, -) - - -def describe_execution_requirements(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 ExecutionRequirements( - 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)), - args.enable_ep_moe, - frozenset(transfers), - ) diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py deleted file mode 100644 index ae37385e68..0000000000 --- a/lightllm/common/basemodel/attention/base_att.py +++ /dev/null @@ -1,183 +0,0 @@ -import threading -import weakref - -import torch -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import Optional, TYPE_CHECKING, Tuple, Dict - -from lightllm.utils.dist_utils import get_current_device_id -from lightllm.utils.envs_utils import get_env_start_args - -if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class BaseAttBackend: - """ - 用于创建支持各种不同的AttBackend, 如 fa3, flashinfer, triton 实现等。 - 每个 model 复用一个 backend 实例。 - """ - - requires_cpu_cache_sync = False - - _instances = {} - _workspace_buffers = {} - _workspace_buffer_lock = threading.Lock() - - def __new__(cls, *args, **kwargs): - """ - Main 和 speculative draft model 可能使用不同的 CUDA graph 上限 - 和缓存布局，不能只按 backend class 共享实例。 - """ - model = kwargs.get("model", args[0] if args else None) - instance_key = (cls, model) - if instance_key not in cls._instances: - instance = super().__new__(cls) - cls._instances[instance_key] = instance - return cls._instances[instance_key] - - def __init__(self, model: "TpPartBaseModel"): - self.model = model - - @staticmethod - def get_gpu_workspace_buffer(key_name: str, workspace_size: int, dtype: torch.dtype = torch.int8) -> torch.Tensor: - """Return a process-local workspace shared by key name and CUDA device.""" - if not key_name: - raise ValueError("workspace key_name must not be empty") - if workspace_size <= 0: - raise ValueError(f"workspace_size must be positive, got {workspace_size}") - - device_id = get_current_device_id() - buffer_key = (device_id, key_name, workspace_size, dtype) - with BaseAttBackend._workspace_buffer_lock: - workspace_buffer = BaseAttBackend._workspace_buffers.get(buffer_key) - if workspace_buffer is None: - workspace_buffer = torch.empty(workspace_size, dtype=dtype, device=device_id) - BaseAttBackend._workspace_buffers[buffer_key] = workspace_buffer - return workspace_buffer - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "BasePrefillAttState": - raise NotImplementedError("not impl") - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "BaseDecodeAttState": - raise NotImplementedError("not impl") - - def uses_dynamic_spec_verify_layout(self) -> bool: - args = get_env_start_args() - draft_step = self.model.mtp_manager.get_decode_draft_step(self.model.is_mtp_draft_model) - is_main_model = not self.model.is_mtp_draft_model - has_decode_draft_step = draft_step > 0 - dynamic_verify_enabled = args.mtp_dynamic_verify - return is_main_model and has_decode_draft_step and dynamic_verify_enabled - - def uses_causal_attention(self) -> bool: - args = get_env_start_args() - is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash", "dflash2") - return not is_parallel_block_draft - - -@dataclass -class AttControl: - """ - prefill_att 和 decode_att 的入参，用于控制att backend 内部的行为, 选择正确的att 实现。 - """ - - use_alibi: bool = False - tp_alibi: torch.Tensor = None - use_sliding_window: bool = False - sliding_window: Tuple[int, int] = (-1, -1) - use_att_sink: bool = False - sink_weight: torch.Tensor = None - # mla 专用传参项 - mla_prefill: bool = False - mla_prefill_dict: Dict = None - mla_decode: bool = False - mla_decode_dict: Dict = None - # nsa (native sparse attention) 专用传参项 - nsa_prefill: bool = False - nsa_prefill_dict: Dict = None - nsa_decode: bool = False - nsa_decode_dict: Dict = None - # linear attention 专用传参项 - linear_att_prefill: bool = False - linear_att_prefill_dict: Dict = None - linear_att_decode: bool = False - linear_att_decode_dict: Dict = None - # 静态 FP8 标定按此逻辑层号查询 State 的物理行；其他后端可不使用。 - logical_layer_id: Optional[int] = None - - -@dataclass -class BaseAttState(ABC): - """Attention metadata with a weak reference to its owning InferStateInfo.""" - - backend: BaseAttBackend = None - # InferStateInfo 已强引用本状态，这里反向使用弱引用，避免形成强引用环。 - # 否则状态及其 Tensor 需等待循环 GC 回收，可能延迟显存复用、增加显存峰值。 - # InferStateInfo 的生命周期由推理调用方或 CUDA Graph 缓存管理，不由 attention state 延长。 - _infer_state_weak_ref: Optional[weakref.ReferenceType["InferStateInfo"]] = None - - @property - def infer_state(self) -> Optional["InferStateInfo"]: - if self._infer_state_weak_ref is None: - return None - infer_state = self._infer_state_weak_ref() - if infer_state is None: - raise ReferenceError("attention state's InferStateInfo has been released") - return infer_state - - def check_no_infer_state(self) -> None: - """Reject InferStateInfo stored directly in any instance attribute. - - This shallow check does not inspect containers or callback closures. - """ - # InferStateInfo imports attention states, so import it only when checking. - from lightllm.common.basemodel.infer_struct import InferStateInfo - - for name, value in vars(self).items(): - if isinstance(value, InferStateInfo): - raise RuntimeError( - f"{type(self).__name__}.{name} must not retain InferStateInfo; use _infer_state_weak_ref instead" - ) - - @abstractmethod - def init_state(self) -> None: - pass - - -@dataclass -class BasePrefillAttState(BaseAttState): - @abstractmethod - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - raise NotImplementedError("not impl") - - -@dataclass -class BaseDecodeAttState(BaseAttState): - def copy_for_decode_cuda_graph(self, new_state: "BaseDecodeAttState"): - """Copy tensor metadata into the buffers owned by the captured graph.""" - for attr_name, attr_value in vars(new_state).items(): - if isinstance(attr_value, torch.Tensor): - attr_ = getattr(self, attr_name, None) - if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr(): - attr_.copy_(attr_value, non_blocking=True) - - @abstractmethod - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - pass diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py deleted file mode 100644 index 08528f98f6..0000000000 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ /dev/null @@ -1,79 +0,0 @@ -from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend -from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend -from lightllm.common.basemodel.attention.linear.replay_gdn import ( - ReplayGDNFlashQlaAttBackend, - ReplayGDNTritonAttBackend, -) -from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.utils.backend_validator import validate -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - -gdn_prefill_backends = { - "flashqla": FlashQlaLinearAttBackend, - "triton": TritonLinearAttBackend, -} - -gdn_replay_prefill_backends = { - "flashqla": ReplayGDNFlashQlaAttBackend, - "triton": ReplayGDNTritonAttBackend, -} - -gdn_decode_backends = { - "triton": TritonLinearAttBackend, -} - -gdn_replay_decode_backends = { - "triton": ReplayGDNTritonAttBackend, -} - -kda_backends = { - "triton": KDALinearAttBackend, -} - -kda_replay_backends = { - "triton": ReplayKDALinearAttBackend, -} - - -def get_gdn_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): - args = get_env_start_args() - classes = gdn_replay_prefill_backends if args.ssm_state_mode == "replay" else gdn_prefill_backends - return _select_backend(classes, args.llm_prefill_att_backend, index, priority_list, "prefill") - - -def get_gdn_linear_decode_att_backend_class(index=1, priority_list=("triton",)): - args = get_env_start_args() - classes = gdn_replay_decode_backends if args.ssm_state_mode == "replay" else gdn_decode_backends - return _select_backend(classes, args.llm_decode_att_backend, index, priority_list, "decode") - - -def get_kda_linear_att_backend_class(): - args = get_env_start_args() - classes = kda_replay_backends if args.ssm_state_mode == "replay" else kda_backends - return classes["triton"] - - -def _select_backend(classes, backend_args, index, priority_list, stage): - backend_str = _get_backend_str(backend_args, index) - if backend_str != "auto": - return classes[backend_str] - - for backend_name in priority_list: - if backend_name == "triton" or validate(backend_name): - if backend_name == "triton": - logger.info(f"Linear {stage} attention backend: triton.") - else: - logger.info(f"Linear {stage} attention backend: {backend_name} (validated).") - return classes[backend_name] - - logger.warning(f"No linear {stage} attention backend validation succeeded, falling back to Triton.") - return classes["triton"] - - -def _get_backend_str(backend_args, index): - assert index == 1, "linear attention backend index must be 1" - return backend_args[index] if len(backend_args) > index else "auto" diff --git a/lightllm/common/basemodel/attention/create_utils.py b/lightllm/common/basemodel/attention/create_utils.py deleted file mode 100644 index a4e2641a9f..0000000000 --- a/lightllm/common/basemodel/attention/create_utils.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Attention backend selection utilities.""" - -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.log_utils import init_logger -from lightllm.utils.backend_validator import validate -from .base_att import BaseAttBackend -from .import_utils import load_backend_class - -logger = init_logger(__name__) - -# Lazy backend class references by data type. -data_type_to_backend = { - "None": { - "triton": "triton.fp:TritonAttBackend", - "fa3": "fa3.fp:Fa3AttBackend", - "flashinfer": "flashinfer.fp:FlashInferAttBackend", - }, - "int4kv": { - "triton": "triton.int4kv:Int4kvTritonAttBackend", - }, - "int8kv": { - "triton": "triton.int8kv:Int8kvTritonAttBackend", - }, - "fp8kv_sph": { - "fa3": "fa3.fp8:Fp8Fa3AttBackend", - }, - "fp8kv_spt": { - "flashinfer": "flashinfer.fp8:Fp8FlashInferAttBackend", - }, -} - -mla_data_type_to_backend = { - "None": { - "triton": "triton.mla:MlaTritonAttBackend", - "fa3": "fa3.mla:MlaFa3AttBackend", - "flashinfer": "flashinfer.mla:MlaFlashInferAttBackend", - }, -} - -nsa_data_type_to_backend = { - "None": { - "flashmla_sparse": "nsa.flashmla_sparse:NsaFlashMlaSparseAttBackend", - # Future backends: "fa3", "tilelang", "aiter" - }, - "fp8kv_dsa": { - "flashmla_sparse": "nsa.fp8_flashmla_sparse:NsaFlashMlaFp8SparseAttBackend", - }, -} - -neo_data_type_to_backend = { - "None": { - "triton": "triton.neo:NeoTritonAttBackend", - "fa3": "fa3.neo:NeoFa3AttBackend", - } -} - - -def _auto_select_backend( - llm_dtype: str, - kv_type_to_backend: dict, - priority_list: list = ["fa3", "flashinfer", "triton"], -) -> type: - """Auto-select the best available backend with validation. - - Priority follows the provided priority_list. - Each backend is validated in a subprocess with ground truth checks. - """ - backend_map = kv_type_to_backend[llm_dtype] - - args = get_env_start_args() - if args.enable_ep_moe: - logger.info("Expert parallelism with MoE enabled, excluding flashinfer attention backend") - priority_list = [name for name in priority_list if name != "flashinfer"] - - for backend_name in priority_list: - if backend_name in backend_map and validate(backend_name): - logger.info(f"Auto-selected {backend_name} backend (validated)") - return load_backend_class(backend_map[backend_name]) - - # Fallback to triton without validation (should not happen) - logger.warning("No backend validation succeeded, falling back to triton") - return load_backend_class(backend_map["triton"]) - - -def _get_decode_backend_priority(priority_list: list, mtp_step: int) -> list: - """Return the auto-selection priority for decode attention.""" - # With MTP, FA3 can make better use of Tensor Core compute and delivers better decode performance. - if mtp_step <= 0 or "fa3" not in priority_list: - return priority_list - return ["fa3"] + [backend_name for backend_name in priority_list if backend_name != "fa3"] - - -def get_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "flashinfer", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - if backend_str != "auto": - return load_backend_class(data_type_to_backend[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list) - - -def get_decode_att_backend_class(index=0, priority_list: list = ["flashinfer", "fa3", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_decode_att_backend[index] - if backend_str != "auto": - return load_backend_class(data_type_to_backend[llm_dtype][backend_str]) - priority_list = _get_decode_backend_priority(priority_list, args.mtp_step) - return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list) - - -def get_mla_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "flashinfer", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - if backend_str != "auto": - return load_backend_class(mla_data_type_to_backend[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list) - - -def get_mla_decode_att_backend_class(index=0, priority_list: list = ["flashinfer", "fa3", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_decode_att_backend[index] - if backend_str != "auto": - return load_backend_class(mla_data_type_to_backend[llm_dtype][backend_str]) - priority_list = _get_decode_backend_priority(priority_list, args.mtp_step) - return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list) - - -def get_nsa_prefill_att_backend_class( - index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend -) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - if backend_str != "auto": - return load_backend_class(backend_map[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list) - - -def get_nsa_decode_att_backend_class( - index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend -) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_decode_att_backend[index] - if backend_str != "auto": - return load_backend_class(backend_map[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list) - - -def get_neo_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - backend_map = neo_data_type_to_backend - if backend_str == "fa3" or (backend_str == "auto" and "fa3" in priority_list): - # Only Neo FA3 selection needs the optional image-token attention interface. - from .fa3.neo import HAS_FLASH_ATTN_INTERFACE - - if not HAS_FLASH_ATTN_INTERFACE: - backend_map = {"None": {"triton": neo_data_type_to_backend["None"]["triton"]}} - if backend_str != "auto": - return load_backend_class(backend_map[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list) diff --git a/lightllm/common/basemodel/attention/fa3/__init__.py b/lightllm/common/basemodel/attention/fa3/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention/fa3/fp.py b/lightllm/common/basemodel/attention/fa3/fp.py deleted file mode 100644 index 32d638ee11..0000000000 --- a/lightllm/common/basemodel/attention/fa3/fp.py +++ /dev/null @@ -1,333 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, TYPE_CHECKING -from lightllm.utils.envs_utils import enable_env_vars -from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune -from lightllm.common.basemodel.triton_kernel.fa3_utils import ( - build_dynamic_spec_fa3_decode_params, - page_table_copy, -) -from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers - - -class Fa3AttBackend(BaseAttBackend): - """Common fixed page-table storage for FA3 attention backends.""" - - # 设为 1 / TRUE / ON 时，FA3 与 CPU cache 搬运在 GPU 上互相等待。 - requires_cpu_cache_sync = enable_env_vars("LIGHTLLM_FA3_CPU_CACHE_SYNC") - page_table_buffers = None - - def __init__(self, model): - super().__init__(model=model) - self._init_infer_page_size() - - def _init_infer_page_size(self): - self.infer_page_size = self.model.args.page_size - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - - # 延迟到首次获取 page table 时再初始化，避免 PD 分离模式下的 prefill 节点 - # 分配仅供 decode 使用的 buffer，减少显存浪费。 - def _init_page_table_buffers(self): - if self.page_table_buffers is not None: - return - - model = self.model - args = model.args - # Dynamic verification may keep the target row and all MTP draft rows for - # every running request. TPSP and CUDA Graph may pad that physical batch - # further, so the fixed buffer must cover both execution paths. - running_max_batch_size = args.running_max_req_size * (args.mtp_step + 1) - if args.enable_tpsp_mix_mode: - tp_size = model.tp_world_size_ - running_max_batch_size = (running_max_batch_size + tp_size - 1) // tp_size * tp_size - self.page_table_max_batch_size = max(running_max_batch_size, model.graph_max_batch_size) - # max_seq_length is max_req_total_len plus the MTP headroom reserved when - # the model is initialized. - self.supported_max_seq_len = model.max_seq_length - buffer_count = 2 if args.enable_decode_microbatch_overlap else 1 - workspace_size = self.page_table_max_batch_size * triton.cdiv(self.supported_max_seq_len, self.infer_page_size) - self.page_table_buffers = [ - self.get_gpu_workspace_buffer( - key_name=f"fa3_page_table_{buffer_index}", - workspace_size=workspace_size, - dtype=torch.int32, - ) - for buffer_index in range(buffer_count) - ] - - def get_page_table_view(self, att_batch_size, max_kv_len, microbatch_index): - """Lazily initialize storage and return a contiguous page-table view.""" - self._init_page_table_buffers() - if att_batch_size > self.page_table_max_batch_size: - raise RuntimeError( - f"FA3 attention batch size {att_batch_size} exceeds page-table capacity " - f"{self.page_table_max_batch_size}" - ) - if max_kv_len > self.supported_max_seq_len: - raise RuntimeError( - f"FA3 max KV sequence length {max_kv_len} exceeds page-table capacity " f"{self.supported_max_seq_len}" - ) - max_page_len = triton.cdiv(max_kv_len, self.infer_page_size) - return self.page_table_buffers[microbatch_index][: att_batch_size * max_page_len].reshape( - att_batch_size, max_page_len - ) - - def create_att_prefill_state(self, infer_state) -> "Fa3PrefillAttState": - return Fa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Fa3DecodeAttState": - return Fa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Fa3PrefillAttState(BasePrefillAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - page_table: torch.Tensor = None - causal: bool = None - - def init_state(self): - self.causal = self.backend.uses_causal_attention() - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - self.page_table = torch.empty( - ( - self.infer_state.batch_size, - triton.cdiv(self.infer_state.max_kv_seq_len, self.backend.infer_page_size), - ), - dtype=torch.int32, - device=self.infer_state.input_ids.device, - ) - page_table_copy( - page_table=self.page_table, - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - page_size=self.backend.infer_page_size, - ) - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.use_alibi is False - return self._nomarl_prefill_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _nomarl_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: Fa3AttBackend = self.backend # for typing - - if att_control.use_sliding_window: - window_size = att_control.sliding_window - else: - window_size = (-1, -1) - - if att_control.use_att_sink: - sink_weight: torch.Tensor = att_control.sink_weight - else: - sink_weight = None - - k_descale, v_descale = None, None # disable quantization - Lq = q.shape[-1] - sm_scale = 1.0 / (Lq ** 0.5) - o = flash_attn_with_kvcache( - q=q, - k_cache=k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v_cache=v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - page_table=self.page_table, - cache_seqlens=self.infer_state.b_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=sm_scale, - causal=self.causal, - window_size=window_size, - softcap=0.0, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=False, - sinks=sink_weight, - ) - return o - - -@dataclasses.dataclass -class Fa3DecodeAttState(BaseDecodeAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - page_table: torch.Tensor = None - b_att_seq_len: torch.Tensor = None - # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。 - decode_max_q_seq_len: int = None - decode_max_kv_seq_len: int = None - causal: bool = None - - def init_state(self): - self.backend: Fa3AttBackend = self.backend - self.causal = self.backend.uses_causal_attention() - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if self.backend.uses_dynamic_spec_verify_layout(): - b_att_req_idx = self._init_dynamic_spec_verify_state(draft_step) - elif draft_step > 0: - b_att_req_idx = self._init_fixed_spec_decode_state(draft_step) - else: - b_att_req_idx = self._init_normal_decode_state() - - self._init_page_table(b_att_req_idx) - - def _init_dynamic_spec_verify_state(self, draft_step: int) -> torch.Tensor: - b_mark_mtp_shared_group = build_mtp_shared_group_markers( - self.infer_state.b_req_idx, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - b_q_seq_len, b_kv_seq_len, b_att_req_idx, self.b_att_seq_len = build_dynamic_spec_fa3_decode_params( - b_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=self.infer_state.batch_size, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = draft_step + 1 - return b_att_req_idx - - def _init_fixed_spec_decode_state(self, draft_step: int) -> torch.Tensor: - mtp_size = draft_step + 1 - assert self.infer_state.batch_size % mtp_size == 0, ( - "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={self.infer_state.batch_size}, draft_step={draft_step}." - ) - - b_q_seq_len = torch.full( - (self.infer_state.b_seq_len.shape[0] // mtp_size,), - fill_value=mtp_size, - dtype=torch.int32, - device=self.infer_state.b_seq_len.device, - ) - b_kv_seq_len = self.infer_state.b_seq_len[draft_step::mtp_size] - b_att_req_idx = self.infer_state.b_req_idx[draft_step::mtp_size] - self.b_att_seq_len = b_kv_seq_len.contiguous() - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = mtp_size - return b_att_req_idx - - def _init_normal_decode_state(self) -> torch.Tensor: - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - self.b_att_seq_len = self.infer_state.b_seq_len - self.decode_max_q_seq_len = 1 - return self.infer_state.b_req_idx - - def _init_spec_decode_cu_seqlens(self, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor): - b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len) - self.cu_seqlens_q = b1_cu_q_seq_len.int() - self.cu_seqlens_k = b1_cu_kv_seq_len.int() - - def _init_page_table(self, b_att_req_idx: torch.Tensor): - att_batch_size = b_att_req_idx.shape[0] - model = self.backend.model - actual_max_kv_len = self.infer_state.max_kv_seq_len - # Graph 捕获会将 infer_state.max_kv_seq_len 改为容量上限，提前保存真实长度用于 FA3 配置查找。 - self.decode_max_kv_seq_len = actual_max_kv_len - actual_max_page_len = triton.cdiv(actual_max_kv_len, self.backend.infer_page_size) - page_table_width = actual_max_kv_len - if model.graph is not None and model.graph.can_run( - batch_size=self.infer_state.batch_size, - max_len_in_batch=actual_max_kv_len, - ): - # CUDA Graph replay uses the shape and strides captured with the - # graph-wide maximum KV length. Keep that fixed row stride while - # writing only the valid portion of each runtime row below. - page_table_width = model.graph.graph_max_len_in_batch - - self.page_table = self.backend.get_page_table_view( - att_batch_size=att_batch_size, - max_kv_len=page_table_width, - microbatch_index=self.infer_state.microbatch_index, - ) - - page_table_copy( - page_table=self.page_table[:, :actual_max_page_len], - req_to_token_indexs=model.req_manager.req_to_token_indexs, - b_req_idx=b_att_req_idx, - page_size=self.backend.infer_page_size, - ) - - def copy_for_decode_cuda_graph(self, new_state: "Fa3DecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.use_alibi is False - return self._normal_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _normal_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - if att_control.use_sliding_window: - window_size = att_control.sliding_window - else: - window_size = (-1, -1) - - if att_control.use_att_sink: - sink_weight: torch.Tensor = att_control.sink_weight - else: - sink_weight = None - - k_descale, v_descale = None, None # disable quantization - Lq = q.shape[-1] - sm_scale = 1.0 / (Lq ** 0.5) - o = flash_attn_with_kvcache_autotune( - q=q, - k_cache=k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v_cache=v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - page_table=self.page_table, - cache_seqlens=self.b_att_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=None, # KV 已提前写入缓存，此处不追加新的 K/V。 - max_seqlen_q=self.decode_max_q_seq_len, - max_seqlen_k=self.decode_max_kv_seq_len, - softmax_scale=sm_scale, - causal=self.causal, - window_size=window_size, - softcap=0.0, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=False, - sinks=sink_weight, - ) - return o diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py deleted file mode 100644 index 8f41d74de2..0000000000 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ /dev/null @@ -1,205 +0,0 @@ -import dataclasses -import weakref -import torch -from ..base_att import AttControl -from lightllm.utils.sgl_utils import flash_attn_with_kvcache -from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import ( - q_per_head_fp8_quant, - q_per_head_static_fp8_quant, -) -from lightllm.utils.log_utils import init_logger -from .fp import Fa3AttBackend, Fa3PrefillAttState, Fa3DecodeAttState - - -logger = init_logger(__name__) - - -class Fp8Fa3AttBackend(Fa3AttBackend): - def _init_infer_page_size(self): - # TODO: FP8 FA3 完成多 token 推理页适配后，改为继承模型 page_size。 - self.infer_page_size = 1 - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - logger.warning( - f"Fp8Fa3AttBackend temporarily uses infer_page_size=1 with model page_size={self.model.args.page_size}; " - "multi-token FP8 inference pages are not implemented yet." - ) - - def create_att_prefill_state(self, infer_state) -> "Fp8Fa3PrefillAttState": - return Fp8Fa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Fp8Fa3DecodeAttState": - return Fp8Fa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Fp8Fa3PrefillAttState(Fa3PrefillAttState): - def init_state(self): - super().init_state() - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ) -> torch.Tensor: - self.backend: Fp8Fa3AttBackend = self.backend # for typing - - q_head_num = q.shape[1] - q_head_dim = q.shape[2] - k_head_num = k.shape[1] - k_head_dim = k.shape[2] - cache_k = k.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - cache_v = v.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_head", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - batch_size = self.infer_state.b_seq_len.shape[0] - q = q.reshape(q.shape[0], k_head_num, -1) - if state.q_scales is None: - # 旧标定文件仅包含 K/V；按本次请求的 Q 长度动态计算每个 KV head 的 scale。 - q, q_scale = q_per_head_fp8_quant( - q=q, - seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], - b1_start_loc=self.cu_seqlens_q, - ) - else: - static_q_scales = state.q_scales[storage_layer_index] - q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) - q_scale = static_q_scales.view(1, k_head_num).expand(batch_size, k_head_num) - o = flash_attn_with_kvcache( - q=q.reshape(-1, q_head_num, q_head_dim), - k_cache=cache_k, - v_cache=cache_v, - page_table=self.page_table, - cache_seqlens=self.infer_state.b_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - q_descale=q_scale, - k_descale=state.get_k_descale(storage_layer_index, batch_size), - v_descale=state.get_v_descale(storage_layer_index, batch_size), - return_softmax_lse=False, - ) - return o - - -@dataclasses.dataclass -class Fp8Fa3DecodeAttState(Fa3DecodeAttState): - def init_state(self): - self.backend: Fp8Fa3AttBackend = self.backend - - super().init_state() - - def copy_for_decode_cuda_graph(self, new_state: "Fp8Fa3DecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_decode_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ): - k_head_num = k.shape[1] - k_head_dim = k.shape[2] - - cache_k = k.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - cache_v = v.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_head", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - - q_head_num = q.shape[1] - att_batch_size = self.b_att_seq_len.shape[0] - q = q.reshape(q.shape[0], k_head_num, -1) - if state.q_scales is None: - # 使用 FA3 的实际 Q 分组，兼容普通 decode、MTP 多 token 及零长度填充组。 - q, q_scale = q_per_head_fp8_quant( - q=q, - seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], - b1_start_loc=self.cu_seqlens_q, - ) - else: - static_q_scales = state.q_scales[storage_layer_index] - q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) - q_scale = static_q_scales.view(1, k_head_num).expand(att_batch_size, k_head_num) - o = flash_attn_with_kvcache( - q=q.reshape(-1, q_head_num, k_head_dim), - k_cache=cache_k, - v_cache=cache_v, - page_table=self.page_table, - cache_seqlens=self.b_att_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.decode_max_q_seq_len, - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - q_descale=q_scale.view(att_batch_size, k_head_num), - k_descale=state.get_k_descale(storage_layer_index, att_batch_size), - v_descale=state.get_v_descale(storage_layer_index, att_batch_size), - return_softmax_lse=False, - ) - return o diff --git a/lightllm/common/basemodel/attention/fa3/mla.py b/lightllm/common/basemodel/attention/fa3/mla.py deleted file mode 100644 index 244df6be12..0000000000 --- a/lightllm/common/basemodel/attention/fa3/mla.py +++ /dev/null @@ -1,244 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, TYPE_CHECKING, Tuple -from lightllm.utils.sgl_utils import flash_attn_with_kvcache -from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params, page_table_copy -from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers -from lightllm.utils.sgl_utils import flash_attn_varlen_func -from .fp import Fa3AttBackend - - -class MlaFa3AttBackend(Fa3AttBackend): - def create_att_prefill_state(self, infer_state) -> "MlaFa3PrefillAttState": - return MlaFa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "MlaFa3DecodeAttState": - return MlaFa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class MlaFa3PrefillAttState(BasePrefillAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - causal: bool = None - - def init_state(self): - self.causal = self.backend.uses_causal_attention() - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._mla_prefill_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: MlaFa3AttBackend = self.backend # for typing - k_nope, k_rope = k - q_head_num = q.shape[1] - k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) - - assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3 - - assert att_control.mla_prefill - softmax_scale = att_control.mla_prefill_dict["softmax_scale"] - - o_tensor = flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - max_seqlen_k=self.infer_state.max_kv_seq_len, - softmax_scale=softmax_scale, - causal=self.causal, - return_softmax_lse=False, - ) - return o_tensor - - -@dataclasses.dataclass -class MlaFa3DecodeAttState(BaseDecodeAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - page_table: torch.Tensor = None - b_att_seq_len: torch.Tensor = None - # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。 - decode_max_q_seq_len: int = None - causal: bool = None - - def init_state(self): - self.backend: MlaFa3AttBackend = self.backend - self.causal = self.backend.uses_causal_attention() - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if self.backend.uses_dynamic_spec_verify_layout(): - b_att_req_idx = self._init_dynamic_spec_verify_state(draft_step) - elif draft_step > 0: - b_att_req_idx = self._init_fixed_spec_decode_state(draft_step) - else: - b_att_req_idx = self._init_normal_decode_state() - - self._init_page_table(b_att_req_idx) - - def _init_dynamic_spec_verify_state(self, draft_step: int) -> torch.Tensor: - b_mark_mtp_shared_group = build_mtp_shared_group_markers( - self.infer_state.b_req_idx, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - b_q_seq_len, b_kv_seq_len, b_att_req_idx, self.b_att_seq_len = build_dynamic_spec_fa3_decode_params( - b_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=self.infer_state.batch_size, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = draft_step + 1 - return b_att_req_idx - - def _init_fixed_spec_decode_state(self, draft_step: int) -> torch.Tensor: - mtp_size = draft_step + 1 - assert self.infer_state.batch_size % mtp_size == 0, ( - "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={self.infer_state.batch_size}, draft_step={draft_step}." - ) - - b_q_seq_len = torch.full( - (self.infer_state.b_seq_len.shape[0] // mtp_size,), - fill_value=mtp_size, - dtype=torch.int32, - device=self.infer_state.b_seq_len.device, - ) - b_kv_seq_len = self.infer_state.b_seq_len[draft_step::mtp_size] - b_att_req_idx = self.infer_state.b_req_idx[draft_step::mtp_size] - self.b_att_seq_len = b_kv_seq_len.contiguous() - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = mtp_size - return b_att_req_idx - - def _init_normal_decode_state(self) -> torch.Tensor: - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - self.b_att_seq_len = self.infer_state.b_seq_len - self.decode_max_q_seq_len = 1 - return self.infer_state.b_req_idx - - def _init_spec_decode_cu_seqlens(self, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor): - b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len) - self.cu_seqlens_q = b1_cu_q_seq_len.int() - self.cu_seqlens_k = b1_cu_kv_seq_len.int() - - def _init_page_table(self, b_att_req_idx: torch.Tensor): - att_batch_size = b_att_req_idx.shape[0] - model = self.backend.model - actual_max_kv_len = self.infer_state.max_kv_seq_len - actual_max_page_len = triton.cdiv(actual_max_kv_len, self.backend.infer_page_size) - page_table_width = actual_max_kv_len - if model.graph is not None and model.graph.can_run( - batch_size=self.infer_state.batch_size, - max_len_in_batch=actual_max_kv_len, - ): - # CUDA Graph replay uses the shape and strides captured with the - # graph-wide maximum KV length. Keep that fixed row stride while - # writing only the valid portion of each runtime row below. - page_table_width = model.graph.graph_max_len_in_batch - - self.page_table = self.backend.get_page_table_view( - att_batch_size=att_batch_size, - max_kv_len=page_table_width, - microbatch_index=self.infer_state.microbatch_index, - ) - - page_table_copy( - page_table=self.page_table[:, :actual_max_page_len], - req_to_token_indexs=model.req_manager.req_to_token_indexs, - b_req_idx=b_att_req_idx, - page_size=self.backend.infer_page_size, - ) - - def copy_for_decode_cuda_graph(self, new_state: "MlaFa3DecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - assert v is None - - return self._mla_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - q_nope, q_rope = q - kv = k - qk_rope_head_dim = 64 - kv_lora_rank = kv.shape[-1] - qk_rope_head_dim - k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, self.backend.infer_page_size, 1, qk_rope_head_dim) - kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, self.backend.infer_page_size, 1, kv_lora_rank) - k_descale, v_descale = None, None - assert att_control.mla_decode - softmax_scale = att_control.mla_decode_dict["softmax_scale"] - - o_tensor = flash_attn_with_kvcache( - q=q_rope, - k_cache=k_rope, - v_cache=kv_nope, - qv=q_nope, - page_table=self.page_table, - cache_seqlens=self.b_att_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.decode_max_q_seq_len, - softmax_scale=softmax_scale, - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=False, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/fa3/neo.py b/lightllm/common/basemodel/attention/fa3/neo.py deleted file mode 100644 index 4a1b1d9b3e..0000000000 --- a/lightllm/common/basemodel/attention/fa3/neo.py +++ /dev/null @@ -1,81 +0,0 @@ -import dataclasses -import weakref -import inspect -import warnings - -import torch - -from .fp import Fa3AttBackend, Fa3PrefillAttState -from ..base_att import AttControl - - -_NEO_FA3_INSTALL_HINT = ( - "Install the Neo FA3 build from source in the current Python environment:\n" - " git clone --recursive https://github.com/WANDY666/flash-attention.git\n" - " cd flash-attention/hopper\n" - " FLASH_ATTENTION_FORCE_BUILD=TRUE python -m pip install " - "--no-build-isolation --no-deps --force-reinstall .\n" - "This replaces a standalone flash_attn_3 / flash_attn_interface installation; " - "LightLLM's regular FA3 uses the separate sgl_kernel package." -) - - -try: - from flash_attn_interface import flash_attn_with_kvcache as flash_attn_with_kvcache_neo - - # Neo FA3 accepts the exclusive visible KV end for each image query. - _sig = inspect.signature(flash_attn_with_kvcache_neo) - if "image_token_end" not in _sig.parameters: - raise ImportError("flash_attn_interface is missing image_token_end support (need the Neo build)") - - HAS_FLASH_ATTN_INTERFACE = True -except ImportError as exc: - warnings.warn( - f"Neo FA3 is unavailable: {exc}. " - "Automatic Neo prefill selection will fall back to Triton. " - "To select Triton explicitly, use --llm_prefill_att_backend triton.\n" + _NEO_FA3_INSTALL_HINT - ) - flash_attn_with_kvcache_neo = None - HAS_FLASH_ATTN_INTERFACE = False - - -class NeoFa3AttBackend(Fa3AttBackend): - def create_att_prefill_state(self, infer_state) -> "NeoFa3PrefillAttState": - return NeoFa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NeoFa3PrefillAttState(Fa3PrefillAttState): - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - # neo_chat*: image-token bidirectional attention requires flash_attn_interface - # (sgl_kernel's flash_attn_with_kvcache does not support image_token_end). - if not HAS_FLASH_ATTN_INTERFACE: - raise ImportError("Neo prefill requires FA3 with image_token_end support.\n" + _NEO_FA3_INSTALL_HINT) - if self.infer_state.b_image_token_end is None: - raise ValueError("Neo prefill requires b_image_token_end to describe image attention spans") - o = flash_attn_with_kvcache_neo( - q=q, - k_cache=k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v_cache=v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - page_table=self.page_table, - cache_seqlens=self.infer_state.b_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=1.0 / (q.shape[-1] ** 0.5), - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - k_descale=None, - v_descale=None, - return_softmax_lse=False, - image_token_end=self.infer_state.b_image_token_end, - ) - return o diff --git a/lightllm/common/basemodel/attention/flashinfer/__init__.py b/lightllm/common/basemodel/attention/flashinfer/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention/flashinfer/env_utils.py b/lightllm/common/basemodel/attention/flashinfer/env_utils.py deleted file mode 100644 index 1394fa2c53..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/env_utils.py +++ /dev/null @@ -1,9 +0,0 @@ -import os -import torch - - -def set_flashinfer_envs(): - capability = torch.cuda.get_device_capability() - arch = f"{capability[0]}.{capability[1]}" - os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}" - return diff --git a/lightllm/common/basemodel/attention/flashinfer/fp.py b/lightllm/common/basemodel/attention/flashinfer/fp.py deleted file mode 100644 index f9ee6c11b0..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/fp.py +++ /dev/null @@ -1,291 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id -from ...triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from ...triton_kernel.repack_kv_index import repack_kv_index -from .env_utils import set_flashinfer_envs -from .utils import should_init_decode_wrapper - - -class FlashInferAttBackend(BaseAttBackend): - workspace_buffer_key = "flashinfer_fp" - workspace_buffer_size = 512 * 1024 * 1024 - - def __init__(self, model): - set_flashinfer_envs() - super().__init__(model=model) - self._init_infer_page_size() - tp_world_size = get_dp_world_size() - self.tp_q_head_num = model.config["num_attention_heads"] // tp_world_size - self.tp_kv_head_num = max(model.config["num_key_value_heads"] // tp_world_size, 1) - head_dim = model.config["hidden_size"] // model.config["num_attention_heads"] - self.head_dim = model.config.get("head_dim", head_dim) - self.max_seq_length = model.max_seq_length - self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) - self.kv_indices_buffer = [ - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - ] - self.q_data_type = model.data_type - self.kv_data_type = model.data_type - - def _init_infer_page_size(self): - self.infer_page_size = self.model.args.page_size - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - - def create_att_prefill_state(self, infer_state) -> "FlashInferPrefillAttState": - return FlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "FlashInferDecodeAttState": - return FlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class FlashInferPrefillAttState(BasePrefillAttState): - prefill_wrapper: object = None - - def init_state(self): - self.backend: FlashInferAttBackend = self.backend - - import flashinfer - - batch_size = self.infer_state.batch_size - device = self.infer_state.input_ids.device - - q_starts = self.infer_state.b1_cu_q_seq_len.int() - # TODO: 将页数、页数前缀和及末页有效 token 数的计算融合为一个 Triton 算子。 - # token 长度除以页大小并向上取整，末页不足一页也计为一页。 - b_page_len = (self.infer_state.b_seq_len + (self.backend.infer_page_size - 1)) // self.backend.infer_page_size - kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) - kv_last_page_len = self.infer_state.b_seq_len - (b_page_len - 1) * self.backend.infer_page_size - kv_indices = torch.empty( - batch_size * self.backend.max_page_num, - dtype=torch.int32, - device=device, - ) - repack_kv_index( - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_token_len=self.infer_state.b_seq_len, - b_page_start_loc=kv_starts[:-1], - max_token_len=self.infer_state.max_kv_seq_len, - out_page_indices=kv_indices, - page_size=self.backend.infer_page_size, - ) - self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - qo_indptr_buf=q_starts, - paged_kv_indptr_buf=kv_starts, - paged_kv_indices_buf=kv_indices, - paged_kv_last_page_len_buf=kv_last_page_len, - ) - self.prefill_wrapper.plan( - q_starts, - kv_starts, - kv_indices, - kv_last_page_len, - self.backend.tp_q_head_num, - self.backend.tp_kv_head_num, - self.backend.head_dim, - self.backend.infer_page_size, - causal=True, - pos_encoding_mode="NONE", - logits_soft_cap=0.0, - q_data_type=self.backend.q_data_type, - kv_data_type=self.backend.kv_data_type, - ) - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._nomarl_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - ) - - def _nomarl_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: FlashInferAttBackend = self.backend # for typing - o_tensor = alloc_func(q.shape, q.dtype, device="cuda") - self.prefill_wrapper.run( - q, - ( - k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - ), - out=o_tensor, - ) - return o_tensor - - -@dataclasses.dataclass -class FlashInferDecodeAttState(BaseDecodeAttState): - kv_last_page_len_buffer: torch.Tensor = None - kv_indices: torch.Tensor = None - kv_starts: torch.Tensor = None - decode_wrapper: object = None - - def _should_init_decode_wrapper(self) -> bool: - return should_init_decode_wrapper(self.infer_state) - - def init_state(self): - import flashinfer - - self.backend: FlashInferAttBackend = self.backend - device = self.infer_state.input_ids.device - model = self.backend.model - # TODO: 将页数、页数前缀和及末页有效 token 数的计算融合为一个 Triton 算子。 - # token 长度除以页大小并向上取整，末页不足一页也计为一页。 - b_page_len = (self.infer_state.b_seq_len + (self.backend.infer_page_size - 1)) // self.backend.infer_page_size - self.kv_last_page_len_buffer = self.infer_state.b_seq_len - (b_page_len - 1) * self.backend.infer_page_size - if ( - self.infer_state.batch_size <= model.graph_max_batch_size - and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch - ): - self.kv_indices = self.backend.kv_indices_buffer[self.infer_state.microbatch_index][ - : self.infer_state.batch_size * self.backend.max_page_num - ] - else: - self.kv_indices = torch.empty( - self.infer_state.batch_size * self.backend.max_page_num, - dtype=torch.int32, - device=device, - ) - - self.kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) - repack_kv_index( - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_token_len=self.infer_state.b_seq_len, - b_page_start_loc=self.kv_starts[:-1], - max_token_len=self.infer_state.max_kv_seq_len, - out_page_indices=self.kv_indices, - page_size=self.backend.infer_page_size, - ) - if not self._should_init_decode_wrapper(): - # 处于 graph replay 回放阶段，不需要特殊初始化 decode wrapper。 - return - - assert self.decode_wrapper is None - self.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - "NHD", - use_cuda_graph=True, - use_tensor_cores=True, - paged_kv_indptr_buffer=self.kv_starts, - paged_kv_indices_buffer=self.kv_indices, - paged_kv_last_page_len_buffer=self.kv_last_page_len_buffer, - ) - self.decode_wrapper.plan( - self.kv_starts, - self.kv_indices, - self.kv_last_page_len_buffer, - self.backend.tp_q_head_num, - self.backend.tp_kv_head_num, - self.backend.head_dim, - self.backend.infer_page_size, - q_data_type=self.backend.q_data_type, - kv_data_type=self.backend.kv_data_type, - non_blocking=True, - ) - return - - def copy_for_decode_cuda_graph(self, new_state: "FlashInferDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - self._refresh_cuda_graph_decode_plan(new_state.infer_state.max_kv_seq_len) - return - - def _refresh_cuda_graph_decode_plan(self, max_kv_len: int): - from flashinfer.decode import fast_decode_plan - - uniform_kv_indptr_cpu = ( - torch.arange( - self.infer_state.batch_size + 1, - dtype=torch.int32, - device="cpu", - ) - * triton.cdiv(max_kv_len, self.backend.infer_page_size) - ) - - fast_decode_plan( - self.decode_wrapper, - indptr=self.kv_starts, - indices=self.kv_indices, - last_page_len=self.kv_last_page_len_buffer, - num_qo_heads=self.backend.tp_q_head_num, - num_kv_heads=self.backend.tp_kv_head_num, - head_dim=self.backend.head_dim, - page_size=self.backend.infer_page_size, - q_data_type=self.backend.q_data_type, - kv_data_type=self.backend.kv_data_type, - non_blocking=True, - global_override_indptr_cpu=uniform_kv_indptr_cpu, - ) - return - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._normal_decode_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - ) - - def _normal_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - o_tensor = alloc_func(q.shape, q.dtype) - self.decode_wrapper.run( - q, - ( - k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - ), - out=o_tensor, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py deleted file mode 100644 index 5f99a2c23f..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ /dev/null @@ -1,146 +0,0 @@ -import dataclasses -import weakref -import torch -from ..base_att import AttControl -from .fp import FlashInferAttBackend, FlashInferPrefillAttState, FlashInferDecodeAttState -from .env_utils import set_flashinfer_envs -from lightllm.utils.log_utils import init_logger - - -logger = init_logger(__name__) - - -class Fp8FlashInferAttBackend(FlashInferAttBackend): - def __init__(self, model): - set_flashinfer_envs() - super().__init__(model=model) - self.kv_data_type = torch.float8_e4m3fn - - def _init_infer_page_size(self): - # TODO: FP8 FlashInfer 完成多 token 推理页适配后，改为继承模型 page_size。 - self.infer_page_size = 1 - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - logger.warning( - "Fp8FlashInferAttBackend temporarily uses infer_page_size=1 with " - f"model page_size={self.model.args.page_size}; multi-token FP8 inference pages are not implemented yet." - ) - - def create_att_prefill_state(self, infer_state) -> "Fp8FlashInferPrefillAttState": - return Fp8FlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Fp8FlashInferDecodeAttState": - return Fp8FlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Fp8FlashInferPrefillAttState(FlashInferPrefillAttState): - def init_state(self): - super().init_state() - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ) -> torch.Tensor: - o_tensor = alloc_func(q.shape, q.dtype, device="cuda") - k = k.unsqueeze(1).view(torch.float8_e4m3fn) - v = v.unsqueeze(1).view(torch.float8_e4m3fn) - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_tensor", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - k_descale = state.cpu_scales[storage_layer_index][0] - v_descale = state.cpu_scales[storage_layer_index][1] - self.prefill_wrapper.run( - q=q, - paged_kv_cache=(k, v), - k_scale=k_descale, - v_scale=v_descale, - out=o_tensor, - ) - return o_tensor - - -@dataclasses.dataclass -class Fp8FlashInferDecodeAttState(FlashInferDecodeAttState): - def init_state(self): - super().init_state() - - def copy_for_decode_cuda_graph(self, new_state): - return super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_decode_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ): - o_tensor = alloc_func(q.shape, q.dtype, device="cuda") - - k = k.unsqueeze(1).view(torch.float8_e4m3fn) - v = v.unsqueeze(1).view(torch.float8_e4m3fn) - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_tensor", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - - k_descale = state.cpu_scales[storage_layer_index][0] - v_descale = state.cpu_scales[storage_layer_index][1] - self.decode_wrapper.run( - q=q, - paged_kv_cache=(k, v), - k_scale=k_descale, - v_scale=v_descale, - out=o_tensor, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/flashinfer/mla.py b/lightllm/common/basemodel/attention/flashinfer/mla.py deleted file mode 100644 index fc55175a3a..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/mla.py +++ /dev/null @@ -1,272 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id -from ...triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from ...triton_kernel.repack_kv_index import repack_kv_index -from ...triton_kernel.flashinfer_mla_plan import fill_mla_decode_plan_for_cuda_graph -from typing import Tuple -from .env_utils import set_flashinfer_envs -from .utils import should_init_decode_wrapper - - -class MlaFlashInferAttBackend(BaseAttBackend): - workspace_buffer_key = "flashinfer_mla" - workspace_buffer_size = 256 * 1024 * 1024 - - def __init__(self, model): - set_flashinfer_envs() - super().__init__(model=model) - self._init_infer_page_size() - num_heads = model.config["num_attention_heads"] - self.tp_q_head_num = num_heads // get_dp_world_size() - self.qk_nope_head_dim = model.qk_nope_head_dim - self.qk_rope_head_dim = model.qk_rope_head_dim - self.kv_lora_rank = model.kv_lora_rank - self.v_head_dim = model.v_head_dim - self.q_data_type = model.data_type - self.kv_data_type = model.data_type - self.max_seq_length = model.max_seq_length - self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) - self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5) - self.kv_indices_buffer = [ - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - ] - - from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale - - if model.config["rope_scaling"] is not None: - rope_scaling = model.config["rope_scaling"] - mscale_all_dim = rope_scaling.get("mscale_all_dim", 0) - scaling_factor = rope_scaling["factor"] - if mscale_all_dim: - mscale = get_deepseek_mscale(scaling_factor, mscale_all_dim) - self.softmax_scale = self.softmax_scale * mscale * mscale - return - - def _init_infer_page_size(self): - self.infer_page_size = self.model.args.page_size - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - - def create_att_prefill_state(self, infer_state) -> "MlaFlashInferPrefillAttState": - return MlaFlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "MlaFlashInferDecodeAttState": - return MlaFlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class MlaFlashInferPrefillAttState(BasePrefillAttState): - prefill_wrapper: object = None - - def init_state(self): - self.backend: MlaFlashInferAttBackend = self.backend - - import flashinfer - - q_starts = self.infer_state.b1_cu_q_seq_len.int() - kv_starts = self.infer_state.b1_cu_kv_seq_len.int() - if self.prefill_wrapper is None: - self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - "NHD", - ) - self.prefill_wrapper.plan( - qo_indptr=q_starts, - kv_indptr=kv_starts, - num_qo_heads=self.backend.tp_q_head_num, - num_kv_heads=self.backend.tp_q_head_num, - head_dim_qk=self.backend.qk_nope_head_dim + self.backend.qk_rope_head_dim, - head_dim_vo=self.backend.v_head_dim, - q_data_type=self.backend.q_data_type, - causal=True, - sm_scale=self.backend.softmax_scale, - ) - return - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._mla_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - ) - - def _mla_prefill_att( - self, q: torch.Tensor, k: Tuple[torch.Tensor, torch.Tensor], v: torch.Tensor, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: MlaFlashInferAttBackend = self.backend # for typing - k_nope, k_rope = k - o_tensor = alloc_func((q.shape[0], q.shape[1], v.shape[-1]), q.dtype, device="cuda") - q_head_num = q.shape[1] - k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) - self.prefill_wrapper.run(q, k, v, out=o_tensor) - return o_tensor - - -@dataclasses.dataclass -class MlaFlashInferDecodeAttState(BaseDecodeAttState): - kv_indices: torch.Tensor = None - kv_starts: torch.Tensor = None - q_indptr_host: torch.Tensor = None - decode_wrapper: object = None - - def _should_init_decode_wrapper(self) -> bool: - return should_init_decode_wrapper(self.infer_state) - - def init_state(self): - import flashinfer - - self.backend: MlaFlashInferAttBackend = self.backend - model = self.backend.model - device = self.infer_state.input_ids.device - batch_size = self.infer_state.batch_size - - # TODO: 将页数及页数前缀和的计算融合为一个 Triton 算子。 - # token 长度除以页大小并向上取整，末页不足一页也计为一页。 - b_page_len = (self.infer_state.b_seq_len + (self.backend.infer_page_size - 1)) // self.backend.infer_page_size - self.kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) - - self.q_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") - self.q_indptr_host = torch.arange(batch_size + 1, dtype=torch.int32, device="cpu") - if batch_size <= model.graph_max_batch_size and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch: - self.kv_indices = self.backend.kv_indices_buffer[self.infer_state.microbatch_index][ - : batch_size * self.backend.max_page_num - ] - else: - self.kv_indices = torch.empty( - batch_size * self.backend.max_page_num, - dtype=torch.int32, - device=device, - ) - - repack_kv_index( - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_token_len=self.infer_state.b_seq_len, - b_page_start_loc=self.kv_starts[:-1], - max_token_len=self.infer_state.max_kv_seq_len, - out_page_indices=self.kv_indices, - page_size=self.backend.infer_page_size, - ) - - if not self._should_init_decode_wrapper(): - return - - assert self.decode_wrapper is None - - self.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - use_cuda_graph=True, - qo_indptr=self.q_indptr, - kv_indices=self.kv_indices, - kv_indptr=self.kv_starts, - kv_len_arr=self.infer_state.b_seq_len, - ) - self.decode_wrapper.plan( - self.q_indptr, - self.kv_starts, - self.kv_indices, - self.infer_state.b_seq_len, - self.backend.tp_q_head_num, - self.backend.kv_lora_rank, - self.backend.qk_rope_head_dim, - self.backend.infer_page_size, - False, # causal - self.backend.softmax_scale, - self.backend.q_data_type, - self.backend.kv_data_type, - ) - return - - def copy_for_decode_cuda_graph(self, new_state: "MlaFlashInferDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - self._refresh_cuda_graph_decode_plan(new_state.infer_state.b_seq_len, new_state.infer_state.max_kv_seq_len) - return - - def _refresh_cuda_graph_decode_plan(self, kv_lens: torch.Tensor, max_kv_len: int): - # Prefer the GPU-generated split plan for long decode; use exact non-split for - # short or unsupported graph shapes. - fill_mla_decode_plan_for_cuda_graph( - decode_wrapper=self.decode_wrapper, - kv_page_indptr=self.kv_starts, - kv_lens=kv_lens, - batch_size=self.infer_state.batch_size, - num_heads=self.backend.tp_q_head_num, - max_kv_len=max_kv_len, - ) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - - assert v is None - - return self._mla_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - qk_rope_head_dim = 64 - q_nope, q_rope = q - o_tensor = alloc_func(q_nope.shape, dtype=q_nope.dtype, device=q_nope.device) - assert att_control.mla_decode - - self.decode_wrapper.run( - q_nope, - q_rope, - k[:, :, :-qk_rope_head_dim].view(-1, self.backend.infer_page_size, 1, k.shape[-1] - qk_rope_head_dim), - k[:, :, -qk_rope_head_dim:].view(-1, self.backend.infer_page_size, 1, qk_rope_head_dim), - out=o_tensor, - return_lse=False, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/flashinfer/utils.py b/lightllm/common/basemodel/attention/flashinfer/utils.py deleted file mode 100644 index 8b159729bf..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/utils.py +++ /dev/null @@ -1,17 +0,0 @@ -def should_init_decode_wrapper(infer_state) -> bool: - match (infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay): - case (False, False): - # eager 推理（包括 graph 不可用时的回退）会直接执行 attention，需要创建 wrapper。 - return True - case (True, False): - # capture 需要创建 wrapper，并随捕获状态保存在图缓存中，供后续 replay 使用。 - return True - case (False, True): - # replay 的临时状态只准备待复制的 metadata，实际执行使用图缓存中的 wrapper。 - return False - case _: - # 两个标记必须是互斥的布尔值，不能同时处于 capture 和 replay 状态。 - raise ValueError( - f"Invalid CUDA Graph flags: capture={infer_state.is_cuda_graph_capture}, " - f"replay={infer_state.is_cuda_graph_replay}" - ) diff --git a/lightllm/common/basemodel/attention/import_utils.py b/lightllm/common/basemodel/attention/import_utils.py deleted file mode 100644 index 14cb0b076d..0000000000 --- a/lightllm/common/basemodel/attention/import_utils.py +++ /dev/null @@ -1,6 +0,0 @@ -from importlib import import_module - - -def load_backend_class(backend): - module_name, class_name = backend.split(":") - return getattr(import_module(f"{__package__}.{module_name}"), class_name) diff --git a/lightllm/common/basemodel/attention/linear/__init__.py b/lightllm/common/basemodel/attention/linear/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention/linear/flashqla.py b/lightllm/common/basemodel/attention/linear/flashqla.py deleted file mode 100644 index 6e57eca732..0000000000 --- a/lightllm/common/basemodel/attention/linear/flashqla.py +++ /dev/null @@ -1,8 +0,0 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend - - -class FlashQlaLinearAttBackend(LinearAttBackend): - def get_prefill_kernel(self): - from flash_qla import chunk_gated_delta_rule - - return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py deleted file mode 100644 index 5105e83584..0000000000 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ /dev/null @@ -1,426 +0,0 @@ -import dataclasses -import weakref -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn -from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params -from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating -from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs -from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( - mtp_fused_recurrent_gated_delta_rule, -) -from lightllm.common.attention_state.bounded_att import GDNState -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import fused_recurrent_gated_delta_rule -from lightllm.utils.envs_utils import get_env_start_args - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState - -if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.common.basemodel.infer_struct import InferStateInfo - from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo - from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerWeight - - -class LinearAttBackend(BaseAttBackend, ABC): - def __init__(self, model: "TpPartBaseModel"): - super().__init__(model=model) - self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) - self.prefill_kernel = self.get_prefill_kernel() - - @abstractmethod - def get_prefill_kernel(self): - pass - - def _init_linear_layer_metadata(self, network_config, tp_world_size): - self.mtp_step = get_env_start_args().mtp_step - - # Linear attention specific dimensions - self.num_v_heads = network_config["linear_num_value_heads"] - self.num_k_heads = network_config["linear_num_key_heads"] - self.head_k_dim = network_config["linear_key_head_dim"] - self.head_v_dim = network_config["linear_value_head_dim"] - self.key_dim = self.head_k_dim * self.num_k_heads - self.value_dim = self.head_v_dim * self.num_v_heads - self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] - self.activation = network_config["hidden_act"] - - # Tensor parallelism dimensions - self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size - self.tp_ba_dim = (self.num_v_heads * 2) // tp_world_size - self.tp_num_k_heads = self.num_k_heads // tp_world_size - self.tp_num_v_heads = self.num_v_heads // tp_world_size - self.tp_key_dim = self.key_dim // tp_world_size - self.tp_value_dim = self.value_dim // tp_world_size - - assert self.num_v_heads % self.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" - self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads - - # SSM state dtype optimization - ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} - self.ssm_state_dtype = ssm_dtype_dict.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) - return - - def _split_qkvzba(self, mixed_qkvzba): - qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim - z_end = qkv_dim + self.tp_value_dim - b_end = z_end + self.tp_num_v_heads - mixed_qkv = mixed_qkvzba[:, :qkv_dim] - z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) - b = mixed_qkvzba[:, z_end:b_end] - a = mixed_qkvzba[:, b_end:] - return mixed_qkv, z, b, a - - def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): - if decode: - query, key, value = torch.split( - mixed_qkv, - [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], - dim=-1, - ) - batch_size = mixed_qkv.shape[0] - query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) - key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) - value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim) - return query, key, value - else: - query, key, value = torch.split( - mixed_qkv, - [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], - dim=-1, - ) - seq_len = query.shape[0] - query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) - key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) - value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) - return query, key, value - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "LinearAttPrefillAttState": - return LinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": - return LinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class LinearAttPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - backend: LinearAttBackend = self.backend - mtp_step = backend.mtp_step - # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。 - # prefill cuda graph 回调必须走 new_infer_state.prefill_att_state1， - # 才能读到这里按当前 batch（含 token padding 后的 dummy request）更新的索引。 - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (mtp_step + 1) - return - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.linear_att_prefill, "linear_att_prefill must be True for Linear prefill attention" - assert att_control.linear_att_prefill_dict is not None, "linear_att_prefill_dict is required" - - linear_att_dict = att_control.linear_att_prefill_dict - mixed_qkvzba: torch.Tensor = linear_att_dict["mixed_qkvzba"] - layer_weight = linear_att_dict["layer_weight"] - layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend - - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=GDNState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 - # prefill 模式下，使用不到这几个维度，所以需要扣除掉， - if backend.mtp_step > 0: - conv_states = conv_states[:, :, : -backend.mtp_step] - mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - core_attn_out = self._gdn_prefill_kernel( - mixed_qkv, conv_states, ssm_states, a, b, self.infer_state, layer_weight - ) - return core_attn_out, z - - def _gdn_prefill_kernel( - self, - mixed_qkv: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: "Qwen3NextInferStateInfo", - layer_weight: "Qwen3NextTransformerLayerWeight", - ): - g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) - mixed_qkv = mixed_qkv.transpose(0, 1) - - backend: LinearAttBackend = self.backend - out_tensor = causal_conv1d_fn( - mixed_qkv, - layer_weight.linear_conv1d.mm_param.weight, - bias=layer_weight.linear_conv1d.bias, - query_start_loc=infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, - has_initial_state=infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation=backend.activation, - ) - mixed_qkv = out_tensor.transpose(0, 1) - - # Recurrent processing - query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) - initial_state = ssm_states[self.b_ssm_buffer_idx] - # g and beta have shape (total_tokens, num_heads), need to unsqueeze to get (1, total_tokens, num_heads) - core_attn_out, last_recurrent_state = backend.prefill_kernel( - q=query, - k=key, - v=value, - g=g.unsqueeze(0), - beta=beta.unsqueeze(0), - initial_state=initial_state, - output_final_state=True, - cu_seqlens=infer_state.b1_cu_q_seq_len, - use_qk_l2norm_in_kernel=True, - ) - # The chunk kernel accumulates the recurrent state in float32 even when - # the state cache is configured as bfloat16. Advanced indexing does - # not perform an implicit dtype conversion for index_put. - ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) - return core_attn_out - - -@dataclasses.dataclass -class LinearAttDecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if draft_step == 0: - self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, GDNState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) - else: - self._init_fixed_mtp_decode_state(draft_step, req_to_mtp_state_index=state.mtp_state_index) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - mtp_size = draft_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "GDN fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={batch_size}, draft_step={draft_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - att_batch_size = self.b_conv_buffer_idx.shape[0] - # Each request owns mtp_size consecutive recurrent-state slots. - b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.linear_att_decode, "linear_att_decode must be True for Linear decode attention" - assert att_control.linear_att_decode_dict is not None, "linear_att_decode_dict is required" - - linear_att_dict = att_control.linear_att_decode_dict - mixed_qkvzba = linear_att_dict["mixed_qkvzba"] - layer_weight = linear_att_dict["layer_weight"] - layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend - - mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=GDNState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if draft_step > 0: - core_attn_out = self._gdn_mtp_kernel( - mixed_qkv, - conv_states, - ssm_states, - a, - b, - self.infer_state, - layer_weight, - ) - else: - # 非 MTP 模式下，使用线性层 decode 状态。 - core_attn_out, z = self._gdn_decode_kernel( - mixed_qkv, - z, - conv_states, - ssm_states, - a, - b, - self.infer_state, - layer_weight, - ) - return core_attn_out, z - - def _gdn_decode_kernel( - self, - mixed_qkv: torch.Tensor, - z: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: "Qwen3NextInferStateInfo", - layer_weight: "Qwen3NextTransformerLayerWeight", - ): - backend: LinearAttBackend = self.backend - - # Recurrent processing with fused gating. Decode uses a specialized - # conv+pack kernel to avoid materializing the post-conv qkv tensor - # before immediately splitting it into q/k/v. - query, key, value, z, a, b = conv_pack_gdn_decode_inputs( - mixed_qkv, - z, - a, - b, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, - backend.activation, - backend.conv_kernel_dim, - backend.tp_num_k_heads, - backend.head_k_dim, - backend.tp_num_v_heads, - backend.head_v_dim, - ) - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - use_qk_l2norm_in_kernel=True, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, - ) - return core_attn_out, z - - def _gdn_mtp_kernel( - self, - mixed_qkv: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: "Qwen3NextInferStateInfo", - layer_weight: "Qwen3NextTransformerLayerWeight", - ): - from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, - ) - - backend: LinearAttBackend = self.backend - - cu_seqlens_q = self.b1_mtp_cu_q_seq_len - mixed_qkv = causal_conv1d_update_mtp( - mixed_qkv, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - mtp_step=backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model), - bias=layer_weight.linear_conv1d.bias, - activation=backend.activation, - conv_state_indices=self.b_conv_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=cu_seqlens_q, - ) - - query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - assert self.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" - # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1, - # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step - # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. - core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), - ssm_state_indices=self.b_ssm_buffer_idx, - ssm_state_write_indices=self.b_ssm_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, - ) - return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py deleted file mode 100644 index 093ec03bef..0000000000 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ /dev/null @@ -1,253 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""KDA attention backend for GLM-5-Next.""" - -from __future__ import annotations - -import dataclasses -import weakref -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update -from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv -from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params -from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate -from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda -from lightllm.common.attention_state.bounded_att import KDAState -from lightllm.common.basemodel.attention.base_att import ( - AttControl, - BaseAttBackend, - BaseDecodeAttState, - BasePrefillAttState, -) - -if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class KDALinearAttBackend(BaseAttBackend): - def __init__(self, model: "TpPartBaseModel"): - super().__init__(model=model) - config = model.config["linear_attn_config"] - self.num_heads = config["num_heads"] - self.head_dim = config["head_dim"] - assert self.num_heads % model.tp_world_size_ == 0 - self.tp_num_heads = self.num_heads // model.tp_world_size_ - self.tp_hidden_size = self.tp_num_heads * self.head_dim - self.conv_kernel_size = config["short_conv_kernel_size"] - self.lower_bound = config.get("gate_lower_bound", -5.0) - self.mtp_step = model.args.mtp_step - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return KDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return KDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def split_qkv(self, mixed_qkv: torch.Tensor): - return mixed_qkv.split(self.tp_hidden_size, dim=-1) - - -@dataclasses.dataclass -class KDAPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_prefill - params = att_control.linear_att_prefill_dict - layer_weight = params["layer_weight"] - layer_num = params["layer_num"] - mixed_qkv = params["mixed_qkv"] - raw_gate = params["raw_gate"] - raw_beta = params["raw_beta"] - backend: KDALinearAttBackend = self.backend - - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_states = conv_states[..., : backend.conv_kernel_size - 1] - q, k, v = causal_conv1d_qkv( - mixed_qkv, - layer_weight.get_merged_kda_conv_weight(), - conv_states, - self.infer_state.b1_cu_q_seq_len, - self.b_conv_buffer_idx, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - ) - raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) - - initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() - output, final_state = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), - beta=raw_beta.float().sigmoid(), - A_log=layer_weight.linear_A_log.weight, - g_bias=layer_weight.linear_dt_bias.weight, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=False, - cu_seqlens=self.infer_state.b1_cu_q_seq_len, - safe_gate=True, - lower_bound=backend.lower_bound, - ) - ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) - return output - - -@dataclasses.dataclass -class KDADecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - mtp_step = self.backend.mtp_step - if mtp_step == 0: - self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, KDAState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(mtp_step + 1, req_to_mtp_state_index=state.mtp_state_index) - else: - self._init_fixed_mtp_decode_state(mtp_step, req_to_mtp_state_index=state.mtp_state_index) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - mtp_size = mtp_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "KDA fixed-layout decode requires batch_size to be divisible by mtp_step + 1, " - f"got batch_size={batch_size}, mtp_step={mtp_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - att_batch_size = self.b_conv_buffer_idx.shape[0] - # Each request owns mtp_size consecutive recurrent-state slots. - b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_decode - params = att_control.linear_att_decode_dict - layer_weight = params["layer_weight"] - layer_num = params["layer_num"] - mixed_qkv = params["mixed_qkv"] - raw_gate = params["raw_gate"] - raw_beta = params["raw_beta"] - backend: KDALinearAttBackend = self.backend - - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) - conv_update = causal_conv1d_update - if backend.mtp_step > 0: - conv_update = causal_conv1d_update_mtp - conv_kwargs.update( - mtp_step=backend.mtp_step, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=self.b1_mtp_cu_q_seq_len, - ) - mixed_qkv = conv_update(mixed_qkv, conv_states, layer_weight.get_merged_kda_conv_weight(), **conv_kwargs) - q, k, v = backend.split_qkv(mixed_qkv) - shape = (1, -1) if backend.mtp_step > 0 else (-1, 1) - q = q.view(*shape, backend.tp_num_heads, backend.head_dim) - k = k.view(*shape, backend.tp_num_heads, backend.head_dim) - v = v.view(*shape, backend.tp_num_heads, backend.head_dim) - raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) - raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - output, _ = fused_recurrent_kda( - q=q, - k=k, - v=v, - raw_gate=raw_gate, - raw_beta=raw_beta, - a_log=layer_weight.linear_A_log.weight, - gate_bias=layer_weight.linear_dt_bias.weight, - initial_state=ssm_states, - lower_bound=backend.lower_bound, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - cu_seqlens=self.b1_mtp_cu_q_seq_len, - num_accepted_tokens=self.b_num_accepted_tokens, - ) - return output diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py deleted file mode 100644 index fd936ff757..0000000000 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ /dev/null @@ -1,345 +0,0 @@ -"""GDN replay 的推理入口。 - -Triton 和 FlashQla 各有一个 backend。它们只差 prefill chunk kernel， -decode 都走 ReplayGDNState。两个类都直接实现 backend 接口，不继承 native GDN。 -""" - -from __future__ import annotations - -import dataclasses -import weakref -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn -from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params -from lightllm.common.attention.backends.gdn.replay.decode import replay -from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating -from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState -from lightllm.common.basemodel.attention.base_att import ( - AttControl, - BaseAttBackend, - BaseDecodeAttState, - BasePrefillAttState, -) -from lightllm.utils.envs_utils import get_env_start_args - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class ReplayGDNTritonAttBackend(BaseAttBackend): - """GDN replay 的 Triton 入口。Prefill 使用仓内 chunk kernel。""" - - def __init__(self, model): - super().__init__(model=model) - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule - - _configure_gdn(self, model, chunk_gated_delta_rule) - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -class ReplayGDNFlashQlaAttBackend(BaseAttBackend): - """GDN replay 的 FlashQla 入口。Prefill 使用 flash_qla chunk kernel。""" - - def __init__(self, model): - super().__init__(model=model) - from flash_qla import chunk_gated_delta_rule - - _configure_gdn(self, model, chunk_gated_delta_rule) - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class ReplayGDNPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - reqs = self.infer_state.b_req_idx - self.b_conv_buffer_idx = reqs - self.b_ssm_buffer_idx = reqs - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_prefill and att_control.linear_att_prefill_dict is not None - params = att_control.linear_att_prefill_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend = self.backend - state = _replay_state(self, layer_num) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - if backend.mtp_step > 0: - conv_states = conv_states[:, :, : -backend.mtp_step] - mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) - core_attn_out = self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight) - return core_attn_out, z - - def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): - backend = self.backend - g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) - mixed_qkv = causal_conv1d_fn( - mixed_qkv.transpose(0, 1), - layer_weight.linear_conv1d.mm_param.weight, - bias=layer_weight.linear_conv1d.bias, - query_start_loc=self.infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, - has_initial_state=self.infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation=backend.activation, - ).transpose(0, 1) - query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv) - core_attn_out, last_state = backend.prefill_kernel( - q=query, - k=key, - v=value, - g=g.unsqueeze(0), - beta=beta.unsqueeze(0), - initial_state=ssm_states[self.b_ssm_buffer_idx], - output_final_state=True, - cu_seqlens=self.infer_state.b1_cu_q_seq_len, - use_qk_l2norm_in_kernel=True, - ) - ssm_states[self.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) - return core_attn_out - - -@dataclasses.dataclass -class ReplayGDNDecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_cu_seqlens: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - role = self.infer_state.att_state_role - matches = [ - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayGDNState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayGDNState for {role}") - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - self._init_indices(draft_step, matches[0].mtp_state_index) - self.b_ssm_buffer_idx = self.b_conv_buffer_idx - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_decode and att_control.linear_att_decode_dict is not None - params = att_control.linear_att_decode_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend = self.backend - mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) - state = _replay_state(self, layer_num) - conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) - draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) - if draft_step > 0: - core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) - else: - core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) - # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) - return core, z - - def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: - infer = self.infer_state - if draft_step == 0: - self.b_conv_buffer_idx = infer.b_req_idx - self.b1_mtp_cu_q_seq_len = None - self.b_num_accepted_tokens = None - # 每行一个 token。replay 和 cursor 共用这份前缀和。 - self.b_cu_seqlens = torch.arange( - infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device - ) - return - if self.backend.uses_dynamic_spec_verify_layout(): - cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( - b_req_idx=infer.b_req_idx, - b_mtp_index=infer.b_mtp_index, - req_to_mtp_state_index=mtp_index, - hold_req_id=infer.req_manager.HOLD_REQUEST_ID, - ) - self.b1_mtp_cu_q_seq_len = cu - self.b_cu_seqlens = cu - self.b_conv_buffer_idx = conv_idx - self.b_num_accepted_tokens = accepted - return - width = draft_step + 1 - batch_size = infer.batch_size - if batch_size % width: - raise ValueError( - "GDN replay decode requires batch_size divisible by draft_step + 1, " - f"got batch_size={batch_size}, draft_step={draft_step}." - ) - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device - ) - self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len - self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() - self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 - - def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state): - backend = self.backend - query, key, value, z, a, b = conv_pack_gdn_decode_inputs( - mixed_qkv, - z, - a, - b, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, - backend.activation, - backend.conv_kernel_dim, - backend.tp_num_k_heads, - backend.head_k_dim, - backend.tp_num_v_heads, - backend.head_v_dim, - ) - query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) - key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) - value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) - output = replay( - state, - state.storage_layer(layer_num), - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b_cu_seqlens, - ) - return output, z - - def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): - backend = self.backend - mixed_qkv = causal_conv1d_update_mtp( - mixed_qkv, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - mtp_step=draft_step, - bias=layer_weight.linear_conv1d.bias, - activation=backend.activation, - conv_state_indices=self.b_conv_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=self.b1_mtp_cu_q_seq_len, - ) - query, key, value = torch.split( - mixed_qkv, - [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], - dim=-1, - ) - query = query.view(-1, backend.tp_num_k_heads, backend.head_k_dim) - key = key.view(-1, backend.tp_num_k_heads, backend.head_k_dim) - value = value.view(-1, backend.tp_num_v_heads, backend.head_v_dim) - output = replay( - state, - state.storage_layer(layer_num), - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b_cu_seqlens, - ) - return output - - -def _replay_state(att_state, layer_num: int) -> ReplayGDNState: - return att_state.infer_state.att_state_manager.get_state( - model_role=att_state.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=ReplayGDNState, - ) - - -def _configure_gdn(backend, model, prefill_kernel) -> None: - """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" - - config = model.config - tp_world_size = model.tp_world_size_ - backend.mtp_step = get_env_start_args().mtp_step - backend.num_v_heads = config["linear_num_value_heads"] - backend.num_k_heads = config["linear_num_key_heads"] - backend.head_k_dim = config["linear_key_head_dim"] - backend.head_v_dim = config["linear_value_head_dim"] - backend.key_dim = backend.head_k_dim * backend.num_k_heads - backend.value_dim = backend.head_v_dim * backend.num_v_heads - backend.conv_kernel_dim = config["linear_conv_kernel_dim"] - backend.activation = config["hidden_act"] - backend.tp_qkvz_dim = (backend.key_dim * 2 + backend.value_dim * 2) // tp_world_size - backend.tp_ba_dim = (backend.num_v_heads * 2) // tp_world_size - backend.tp_num_k_heads = backend.num_k_heads // tp_world_size - backend.tp_num_v_heads = backend.num_v_heads // tp_world_size - backend.tp_key_dim = backend.key_dim // tp_world_size - backend.tp_value_dim = backend.value_dim // tp_world_size - assert backend.num_v_heads % backend.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" - backend.num_v_heads_per_k_head = backend.num_v_heads // backend.num_k_heads - ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} - backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) - backend.prefill_kernel = prefill_kernel - - -def _split_qkvzba(backend, mixed_qkvzba): - qkv_dim = backend.tp_key_dim * 2 + backend.tp_value_dim - z_end = qkv_dim + backend.tp_value_dim - b_end = z_end + backend.tp_num_v_heads - mixed_qkv = mixed_qkvzba[:, :qkv_dim] - z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, backend.tp_num_v_heads, backend.head_v_dim) - b = mixed_qkvzba[:, z_end:b_end] - a = mixed_qkvzba[:, b_end:] - return mixed_qkv, z, b, a - - -def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): - query, key, value = torch.split( - mixed_qkv, - [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], - dim=-1, - ) - if decode: - batch_size = mixed_qkv.shape[0] - leading = (batch_size, 1) - else: - leading = (1, query.shape[0]) - query = query.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) - key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) - value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) - return query, key, value diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py deleted file mode 100644 index 6275b5e4c0..0000000000 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ /dev/null @@ -1,247 +0,0 @@ -"""KDA replay 的推理入口。 - -这个 backend 只服务 ReplayKDAState。维度在这里自己读取，不继承 native KDA backend。 -""" - -from __future__ import annotations - -import dataclasses -import weakref -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn, causal_conv1d_update -from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params -from lightllm.common.attention.backends.kda.replay.decode import replay -from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate -from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState -from lightllm.common.basemodel.attention.base_att import ( - AttControl, - BaseAttBackend, - BaseDecodeAttState, - BasePrefillAttState, -) - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class ReplayKDALinearAttBackend(BaseAttBackend): - """KDA replay 推理 backend。GLM-5-Next 在 replay 模式下只构造这一个类。""" - - def __init__(self, model): - super().__init__(model=model) - config = model.config["linear_attn_config"] - self.num_heads = config["num_heads"] - self.head_dim = config["head_dim"] - assert self.num_heads % model.tp_world_size_ == 0 - self.tp_num_heads = self.num_heads // model.tp_world_size_ - self.tp_hidden_size = self.tp_num_heads * self.head_dim - self.conv_kernel_size = config["short_conv_kernel_size"] - self.lower_bound = config.get("gate_lower_bound", -5.0) - self.mtp_step = model.args.mtp_step - - def split_qkv(self, mixed_qkv: torch.Tensor): - return mixed_qkv.split(self.tp_hidden_size, dim=-1) - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return ReplayKDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return ReplayKDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class ReplayKDAPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - reqs = self.infer_state.b_req_idx - self.b_conv_buffer_idx = reqs - self.b_ssm_buffer_idx = reqs - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_prefill - params = att_control.linear_att_prefill_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend: ReplayKDALinearAttBackend = self.backend - state = _replay_state(self, layer_num) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_states = conv_states[..., : backend.conv_kernel_size - 1] - mixed_qkv = causal_conv1d_fn( - params["mixed_qkv"].transpose(0, 1), - layer_weight.get_merged_kda_conv_weight(), - bias=None, - query_start_loc=self.infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, - has_initial_state=self.infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation="silu", - ).transpose(0, 1) - q, k, v = backend.split_qkv(mixed_qkv) - q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) - k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) - v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) - raw_gate = params["raw_gate"].view(1, -1, backend.tp_hidden_size) - raw_beta = params["raw_beta"].view(1, -1, backend.tp_num_heads) - output, final_state = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), - beta=raw_beta.float().sigmoid(), - A_log=layer_weight.linear_A_log.weight, - g_bias=layer_weight.linear_dt_bias.weight, - initial_state=ssm_states[self.b_ssm_buffer_idx].contiguous(), - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=self.infer_state.b1_cu_q_seq_len, - safe_gate=True, - lower_bound=backend.lower_bound, - ) - ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) - return output - - -@dataclasses.dataclass -class ReplayKDADecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_cu_seqlens: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - role = self.infer_state.att_state_role - matches = [ - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayKDAState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayKDAState for {role}") - self._init_indices(self.backend.mtp_step, matches[0].mtp_state_index) - self.b_ssm_buffer_idx = self.b_conv_buffer_idx - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_decode - params = att_control.linear_att_decode_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend: ReplayKDALinearAttBackend = self.backend - state = _replay_state(self, layer_num) - conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) - conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) - conv_update = causal_conv1d_update - if backend.mtp_step > 0: - conv_update = causal_conv1d_update_mtp - conv_kwargs.update( - mtp_step=backend.mtp_step, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=self.b1_mtp_cu_q_seq_len, - ) - mixed_qkv = conv_update( - params["mixed_qkv"], - conv_states, - layer_weight.get_merged_kda_conv_weight(), - **conv_kwargs, - ) - q, k, v = backend.split_qkv(mixed_qkv) - q = q.view(-1, backend.tp_num_heads, backend.head_dim) - k = k.view(-1, backend.tp_num_heads, backend.head_dim) - v = v.view(-1, backend.tp_num_heads, backend.head_dim) - raw_gate = params["raw_gate"].reshape(-1, backend.tp_hidden_size) - raw_beta = params["raw_beta"].reshape(-1, backend.tp_num_heads) - layer = state.storage_layer(layer_num) - output = replay( - state_ssm=state.ssm_state[layer], - state_keys=state.keys[layer], - state_deltas=state.deltas[layer], - state_gates=state.gates[layer], - state_raw_keys=state.raw_keys[layer], - state_raw_values=state.raw_values[layer], - state_betas=state.betas[layer], - state_cursors=state.cursors, - state_tile_sync_mark=state.tile_sync_mark, - q=q, - k=k, - v=v, - a=raw_gate, - b=raw_beta, - a_log=layer_weight.linear_A_log.weight, - bias=layer_weight.linear_dt_bias.weight, - reqs=self.b_conv_buffer_idx, - cu_seqlens=self.b_cu_seqlens, - state_verify_width=state.verify_width, - state_compute_mode=state.compute_mode, - state_gate_lower_bound=state.gate_lower_bound, - ) - # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) - return output - - def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: - infer = self.infer_state - if mtp_step == 0: - self.b_conv_buffer_idx = infer.b_req_idx - self.b1_mtp_cu_q_seq_len = None - self.b_num_accepted_tokens = None - # 每行一个 token。replay 和 cursor 共用这份前缀和。 - self.b_cu_seqlens = torch.arange( - infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device - ) - return - if self.backend.uses_dynamic_spec_verify_layout(): - cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( - b_req_idx=infer.b_req_idx, - b_mtp_index=infer.b_mtp_index, - req_to_mtp_state_index=mtp_index, - hold_req_id=infer.req_manager.HOLD_REQUEST_ID, - ) - self.b1_mtp_cu_q_seq_len = cu - self.b_cu_seqlens = cu - self.b_conv_buffer_idx = conv_idx - self.b_num_accepted_tokens = accepted - return - width = mtp_step + 1 - batch_size = infer.batch_size - if batch_size % width: - raise ValueError( - "KDA replay decode requires batch_size divisible by mtp_step + 1, " - f"got batch_size={batch_size}, mtp_step={mtp_step}." - ) - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device - ) - self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len - self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() - self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 - - -def _replay_state(att_state, layer_num: int) -> ReplayKDAState: - return att_state.infer_state.att_state_manager.get_state( - model_role=att_state.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=ReplayKDAState, - ) diff --git a/lightllm/common/basemodel/attention/linear/triton.py b/lightllm/common/basemodel/attention/linear/triton.py deleted file mode 100644 index f0b6400f51..0000000000 --- a/lightllm/common/basemodel/attention/linear/triton.py +++ /dev/null @@ -1,7 +0,0 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule - - -class TritonLinearAttBackend(LinearAttBackend): - def get_prefill_kernel(self): - return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/nsa/__init__.py b/lightllm/common/basemodel/attention/nsa/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py deleted file mode 100644 index 93560de07a..0000000000 --- a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py +++ /dev/null @@ -1,225 +0,0 @@ -import dataclasses -from typing import TYPE_CHECKING, Optional, Tuple - -import torch -import triton -import flash_mla as flashmla - - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState - -if TYPE_CHECKING: - from flash_mla.flash_mla_interface import FlashMLASchedMeta - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttBackend): - """FlashMLA layouts and operator workspace; forward states own sparse metadata.""" - - def __init__(self, model): - super().__init__(model=model) - from lightllm.models.deepseek_v4.config import create_attention_config - - config = create_attention_config( - model_config=model.config, - model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), - tp_rank=model.tp_rank_, - tp_world_size=model.tp_world_size_, - ) - self.tp_q_head_num = model.config["num_attention_heads"] // model.tp_world_size_ - # Sparse FlashMLA uses 64-head tiles; TP can leave fewer local Q heads. - self.padded_q_head_num = ((self.tp_q_head_num + 63) // 64) * 64 - self.compress_ratios = tuple(dict.fromkeys(config.compress_ratios)) - self.layer_compress_ratios = dict(zip(config.logical_layer_ids, config.compress_ratios)) - self.swa_config = config.create_swa_config() - self.compressed_configs = { - ratio: config.create_compressed_config(compress_ratio=ratio) for ratio in self.compress_ratios if ratio - } - - @classmethod - def get_swa_prefill_workspace(cls, model): - """由 backend 复用计算临时区；模型在 KV 容量评估前调用，以计入实际显存占用。""" - from lightllm.models.deepseek_v4.config import create_attention_config - - config = create_attention_config( - model_config=model.config, - model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), - tp_rank=model.tp_rank_, - tp_world_size=model.tp_world_size_, - ).create_swa_config() - page_token_num = config.page_token_size - max_forward_tokens = max(model.batch_max_tokens, model.max_req_num) + 2 * model.tp_world_size_ - capacity = ( - triton.cdiv( - max_forward_tokens + model.max_req_num * (config.window_size + 2 * page_token_num), page_token_num - ) - * page_token_num - ) - if model.args.run_mode == "decode": - # D 节点初始化仍会对 HOLD 请求执行一次 prefill；只保留末尾的 HOLD 页， - # 不为普通请求预留长 prefill 临时区，实际请求的前缀由 P 节点传入。 - capacity = 0 - regions = 1 + int(model.args.enable_prefill_microbatch_overlap) - # 每个 microbatch 使用独立临时区；最后一页留给 HOLD 请求。 - token_num = regions * capacity + page_token_num - # 物理行宽由 SWA 原子配置给出，backend 仅决定 workspace 的 token 容量。 - slot_bytes = config.get_packed_kv_bytes(page_token_num) // page_token_num - layer_count = len(config.logical_layer_ids) - buffer = cls.get_gpu_workspace_buffer( - key_name=f"deepseek_v4_swa_prefill_{model.att_state_role}", - workspace_size=layer_count * token_num * slot_bytes, - dtype=torch.uint8, - ) - return buffer.view(layer_count, token_num, slot_bytes), capacity - - def get_att_input_params(self, infer_state, logical_layer_id: int): - state = infer_state.swa_state - swa_buffer = state.get_layer_view(logical_layer_id=logical_layer_id) - config = state.config - swa = swa_buffer.view(-1, config.get_packed_kv_bytes(config.page_token_size)) - ratio = self.layer_compress_ratios[logical_layer_id] - compressed_kv = infer_state.get_compressed_kv_view(logical_layer_id=logical_layer_id).kv if ratio else None - if infer_state.is_prefill: - swa = infer_state.prefill_att_state.swa_buffer[config.logical_layer_ids.index(logical_layer_id)].view( - -1, - config.get_packed_kv_bytes(config.page_token_size), - ) - return swa, compressed_kv - - def _flashmla_att( - self, - q: torch.Tensor, - packed_kv: Tuple[torch.Tensor, Optional[torch.Tensor]], - nsa_dict: dict, - sched_meta: "FlashMLASchedMeta", - out: torch.Tensor, - ) -> torch.Tensor: - """Attend to SWA plus optional compressed KV and return all padded Q heads.""" - swa_packed, compressed_packed = packed_kv - # 原子配置提供计算视图；backend 不再计算编码字节数或解释压缩块布局。 - compressed_cache = None - if compressed_packed is not None: - config = self.compressed_configs[nsa_dict["compress_ratio"]] - compressed_cache = config.get_attention_view(packed_kv=compressed_packed) - swa_cache = self.swa_config.get_attention_view(packed_kv=swa_packed) - - padded_out, _ = flashmla.flash_mla_with_kvcache( - q=q.unsqueeze(1), - k_cache=swa_cache, - block_table=None, - cache_seqlens=None, - head_dim_v=nsa_dict["head_dim_v"], - tile_scheduler_metadata=sched_meta, - num_splits=None, - softmax_scale=nsa_dict["softmax_scale"], - causal=False, - is_fp8_kvcache=True, - indices=nsa_dict["swa_indices"], - attn_sink=nsa_dict["attn_sink"], - topk_length=nsa_dict["swa_lengths"], - extra_k_cache=compressed_cache, - extra_indices_in_kvcache=nsa_dict.get("extra_indices"), - extra_topk_length=nsa_dict.get("extra_lengths"), - out=out.unsqueeze(1), - ) - return padded_out[:, 0] - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "_PrefillAttState": - return _PrefillAttState(backend=self) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "_DecodeAttState": - return _DecodeAttState(backend=self) - - -@dataclasses.dataclass -class _PrefillAttState(BasePrefillAttState): - sched_meta_by_ratio: dict[int, "FlashMLASchedMeta"] | None = None - - def prepare_swa_workspace(self, infer_state): - from lightllm.models.deepseek_v4.triton_kernel.swa_cache import prepare_swa_prefill_workspace - - swa_state = infer_state.swa_state - self.swa_buffer, capacity = self.backend.get_swa_prefill_workspace(self.backend.model) - self.HOLD_SWA_SLOT = self.swa_buffer.shape[1] - swa_state.config.page_token_size - self.prefill_token_ranges = torch.zeros( - (swa_state.request_capacity, 2), dtype=torch.int32, device=self.swa_buffer.device - ) - self.swa_prefill_workspace_offsets = torch.full( - (swa_state.request_capacity,), -1, dtype=torch.int32, device=self.swa_buffer.device - ) - prepare_swa_prefill_workspace( - swa_state=swa_state, - prefill_state=self, - b_req_idx=infer_state.b_req_idx, - b_ready_cache_len=infer_state.b_ready_cache_len, - b_seq_len=infer_state.b_seq_len, - workspace_base=infer_state.microbatch_index * capacity, - ) - - def init_state(self): - self.sched_meta_by_ratio = {ratio: flashmla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios} - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, Optional[torch.Tensor]], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - *, - out: torch.Tensor = None, - ) -> torch.Tensor: - nsa_dict = att_control.nsa_prefill_dict - if out is None: - out = alloc_func((q.shape[0], q.shape[1], nsa_dict["head_dim_v"]), dtype=q.dtype, device=q.device) - - padded_out = self.backend._flashmla_att( - q, - k, - nsa_dict, - self.sched_meta_by_ratio[nsa_dict["compress_ratio"]], - out=out, - ) - return padded_out[:, : self.backend.tp_q_head_num] - - -@dataclasses.dataclass -class _DecodeAttState(BaseDecodeAttState): - sched_meta_by_ratio: dict[int, "FlashMLASchedMeta"] | None = None - capture_sched_meta_by_ratio: dict[int, "FlashMLASchedMeta"] | None = None - - def init_state(self): - self.sched_meta_by_ratio = {ratio: flashmla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios} - # Keep capture metadata untouched by warmup so its GPU planning kernels enter the graph. - self.capture_sched_meta_by_ratio = { - ratio: flashmla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios - } - - def decode_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, Optional[torch.Tensor]], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - nsa_dict = att_control.nsa_decode_dict - sched_meta_by_ratio = ( - self.capture_sched_meta_by_ratio if torch.cuda.is_current_stream_capturing() else self.sched_meta_by_ratio - ) - out = alloc_func((q.shape[0], q.shape[1], nsa_dict["head_dim_v"]), dtype=q.dtype, device=q.device) - padded_out = self.backend._flashmla_att( - q, - k, - nsa_dict, - sched_meta_by_ratio[nsa_dict["compress_ratio"]], - out=out, - ) - return padded_out[:, : self.backend.tp_q_head_num] - - -DSV4_NSA_BACKENDS = { - "fp8kv_dsa": { - "flashmla_sparse": "nsa.dsv4_fp8_flashmla_sparse:DeepseekV4FlashMlaFp8SparseAttBackend", - } -} diff --git a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py deleted file mode 100644 index 92c48c516f..0000000000 --- a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py +++ /dev/null @@ -1,193 +0,0 @@ -# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/nsa_backend.py -# Uses sgl_kernel.flash_mla and sgl_kernel.flash_attn from the sglang kernel library. - -import dataclasses -import weakref -import torch -from typing import Tuple, TYPE_CHECKING - -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.dist_utils import get_current_device_id - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class NsaFlashMlaSparseAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model=model) - device = get_current_device_id() - self.ragged_mem_buffers = [ - torch.empty(model.graph_max_batch_size * model.max_seq_length, dtype=torch.int32, device=device) - for _ in range(2) - ] - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaSparsePrefillAttState": - return NsaFlashMlaSparsePrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaSparseDecodeAttState": - return NsaFlashMlaSparseDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NsaFlashMlaSparsePrefillAttState(BasePrefillAttState): - """Prefill attention state for NSA using flash_mla_sparse_fwd.""" - - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - - def init_state(self): - self.backend: NsaFlashMlaSparseAttBackend = self.backend - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.input_ids.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - return - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_prefill, "nsa_prefill must be True for NSA prefill attention" - assert att_control.nsa_prefill_dict is not None, "nsa_prefill_dict is required" - - return self._nsa_prefill_att(q=q, kv=k, att_control=att_control) - - def _nsa_prefill_att( - self, - q: torch.Tensor, - kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - from sgl_kernel.flash_mla import flash_mla_sparse_fwd - - nsa_dict = att_control.nsa_prefill_dict - topk_mem_indices = nsa_dict["topk_mem_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - - if topk_mem_indices.ndim == 2: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - - mla_out, _, _ = flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=topk_mem_indices, - sm_scale=softmax_scale, - d_v=kv_lora_rank, - ) - return mla_out - - -@dataclasses.dataclass -class NsaFlashMlaSparseDecodeAttState(BaseDecodeAttState): - - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - nsa_cache_seqlens: torch.Tensor = None - nsa_cu_seqlens_k_new: torch.Tensor = None - - def init_state(self): - self.backend: NsaFlashMlaSparseAttBackend = self.backend - model = self.backend.model - use_cuda_graph = ( - self.infer_state.batch_size <= model.graph_max_batch_size - and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch - ) - - if use_cuda_graph: - self.ragged_mem_index = self.backend.ragged_mem_buffers[self.infer_state.microbatch_index] - else: - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.b_seq_len.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self.nsa_cache_seqlens = torch.minimum( - torch.full(size=(self.infer_state.batch_size,), fill_value=2048, dtype=torch.int32, device="cuda"), - self.infer_state.b_seq_len, - ) - padded_seq_lens = torch.zeros(size=(self.nsa_cache_seqlens.shape[0] + 1,), dtype=torch.int32, device="cuda") - # 进行 cumsum 操作 - padded_seq_lens[1:].copy_(self.nsa_cache_seqlens, non_blocking=True) - self.nsa_cu_seqlens_k_new = padded_seq_lens.cumsum(dim=0, dtype=torch.int32) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_decode, "nsa_decode must be True for NSA decode attention" - assert att_control.nsa_decode_dict is not None, "nsa_decode_dict is required" - - return self._nsa_decode_att(q=q, kv=k, att_control=att_control) - - def _nsa_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - nsa_dict = att_control.nsa_decode_dict - topk_mem_indices = nsa_dict["topk_mem_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - qk_rope_head_dim = nsa_dict["qk_rope_head_dim"] - - q_nope, q_rope = q - - # Extract k_rope and kv_nope from the KV buffer - k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) - kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) - - o_tensor = flash_attn_with_kvcache( - q=q_rope, - k_cache=k_rope, - v_cache=kv_nope, - qv=q_nope, - page_table=topk_mem_indices, - cache_seqlens=self.nsa_cache_seqlens, - cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, - cu_seqlens_k_new=self.nsa_cu_seqlens_k_new, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=softmax_scale, - causal=True, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py deleted file mode 100644 index 10ad62d6b8..0000000000 --- a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py +++ /dev/null @@ -1,204 +0,0 @@ -import dataclasses -import weakref -import torch -from typing import TYPE_CHECKING, Tuple - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState -from lightllm.utils.dist_utils import get_current_device_id - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class NsaFlashMlaFp8SparseAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model=model) - device = get_current_device_id() - self.ragged_mem_buffers = [ - torch.empty(model.graph_max_batch_size * model.max_seq_length, dtype=torch.int32, device=device) - for _ in range(2) - ] - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparsePrefillAttState": - return NsaFlashMlaFp8SparsePrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparseDecodeAttState": - return NsaFlashMlaFp8SparseDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NsaFlashMlaFp8SparsePrefillAttState(BasePrefillAttState): - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - - def init_state(self): - self.backend: NsaFlashMlaFp8SparseAttBackend = self.backend - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.input_ids.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - return - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_prefill, "nsa_prefill must be True for NSA prefill attention" - assert att_control.nsa_prefill_dict is not None, "nsa_prefill_dict is required" - return self._nsa_prefill_att(q=q, packed_kv=k, att_control=att_control) - - def _nsa_prefill_att( - self, - q: torch.Tensor, - packed_kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - import flash_mla - - nsa_dict = att_control.nsa_prefill_dict - topk_indices = nsa_dict["topk_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - topk_mem_indices = nsa_dict["topk_mem_indices"] - prefill_cache_kv = nsa_dict["prefill_cache_kv"] - - if self.infer_state.max_cache_len > 0: - # 当前推理生成的token kv部分从 prefill_cache_kv 中获取，历史 - # 部分kv 从 packed_kv 中获取, 并进行反量化，这样可以避免 prefill_cache_kv - # 部分的数据进行重复的反量化操作，提升整体的性能。 - from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( - get_prefill_kv_cache_and_remap_indices_triton, - ) - - kv, topk_indices = get_prefill_kv_cache_and_remap_indices_triton( - packed_kv=packed_kv, - topk_mem_indices=topk_mem_indices, - prefill_mem_index=self.infer_state.mem_index, - prefill_cache_kv=prefill_cache_kv, - prefill_dtype=prefill_cache_kv.dtype, - ) - else: - kv = prefill_cache_kv - - if topk_indices.ndim == 2: - topk_indices = topk_indices.unsqueeze(1) - - mla_out, _, _ = flash_mla.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=topk_indices, - sm_scale=softmax_scale, - d_v=kv_lora_rank, - ) - return mla_out - - -@dataclasses.dataclass -class NsaFlashMlaFp8SparseDecodeAttState(BaseDecodeAttState): - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - flashmla_sched_meta: object = None - - def init_state(self): - self.backend: NsaFlashMlaFp8SparseAttBackend = self.backend - model = self.backend.model - use_cuda_graph = ( - self.infer_state.batch_size <= model.graph_max_batch_size - and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch - ) - - if use_cuda_graph: - self.ragged_mem_index = self.backend.ragged_mem_buffers[self.infer_state.microbatch_index] - else: - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.b_seq_len.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - import flash_mla - - self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata() - return - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_decode, "nsa_decode must be True for NSA decode attention" - assert att_control.nsa_decode_dict is not None, "nsa_decode_dict is required" - return self._nsa_decode_att(q=q, packed_kv=k, att_control=att_control) - - def _nsa_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - packed_kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - import flash_mla - - nsa_dict = att_control.nsa_decode_dict - topk_mem_indices = nsa_dict["topk_mem_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - - if topk_mem_indices.ndim == 2: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" - - q_nope, q_rope = q - q_all = torch.cat([q_nope, q_rope], dim=-1).unsqueeze(1).contiguous() - kv = torch.as_strided( - packed_kv, - size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), - stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), - ) - - o_tensor, _ = flash_mla.flash_mla_with_kvcache( - q=q_all, - k_cache=kv, - block_table=None, - cache_seqlens=None, - head_dim_v=kv_lora_rank, - tile_scheduler_metadata=self.flashmla_sched_meta, - softmax_scale=softmax_scale, - causal=False, - is_fp8_kvcache=True, - indices=topk_mem_indices, - ) - return o_tensor[:, 0, :, :] # [b, 1, h, d] -> [b, h, d] diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py deleted file mode 100644 index 261c2a16c1..0000000000 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ /dev/null @@ -1,238 +0,0 @@ -import dataclasses -import weakref - -import torch -import triton - -from .glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState -from ..base_att import BaseAttBackend - - -# Cap each microbatch's BF16 KV workspace at 2 GiB (2097152 × 512 × 2 bytes). -KV_GATHER_CHUNK_SIZE = 2097152 - - -class Fp8Glm5NextSparseAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model=model) - self._prefill_workspaces = self.allocate_prefill_workspaces(model) - - @classmethod - def allocate_prefill_workspaces(cls, model): - regions = 1 + int(model.args.enable_prefill_microbatch_overlap) - q_token_capacity = model.batch_max_tokens - topk_size = triton.cdiv(model.config["index_topk"] + 3, 128) * 128 - kv_token_capacity = min(KV_GATHER_CHUNK_SIZE, model.max_req_num * model.max_seq_length + q_token_capacity) - if model.args.run_mode == "decode": - kv_token_capacity = min(kv_token_capacity, q_token_capacity) - prefix = f"glm5_fp8_prefill_{model.att_state_role}" - gathered_kv = cls.get_gpu_workspace_buffer( - prefix + "_kv", regions * kv_token_capacity * 512, torch.bfloat16 - ).view(regions, kv_token_capacity, 1, 512) - chunk_topk_indices = cls.get_gpu_workspace_buffer( - prefix + "_topk", regions * q_token_capacity * topk_size, torch.int32 - ).view(regions, q_token_capacity, topk_size) - chunk_cache_seqlens = cls.get_gpu_workspace_buffer( - prefix + "_seqlens", regions * q_token_capacity, torch.int32 - ).view(regions, q_token_capacity) - cu_seqlens_q = cls.get_gpu_workspace_buffer( - prefix + "_cu_q", regions * (q_token_capacity + 1), torch.int32 - ).view(regions, q_token_capacity + 1) - cu_seqlens_q.copy_(torch.arange(q_token_capacity + 1, dtype=torch.int32, device=cu_seqlens_q.device)) - return [ - { - "gathered_kv": gathered_kv[index], - "chunk_topk_indices": chunk_topk_indices[index], - "chunk_cache_seqlens": chunk_cache_seqlens[index], - "cu_seqlens_q": cu_seqlens_q[index], - } - for index in range(regions) - ] - - def get_prefill_workspace(self, infer_state, topk_indices): - buffers = self._prefill_workspaces[infer_state.microbatch_index] - q_token_num, topk_size = topk_indices.shape - assert q_token_num <= buffers["chunk_cache_seqlens"].numel() - assert topk_size <= buffers["chunk_topk_indices"].shape[1] - chunk_topk_indices = buffers["chunk_topk_indices"].view(-1) - return { - "gathered_kv": buffers["gathered_kv"], - "chunk_topk_indices": chunk_topk_indices[: q_token_num * topk_size].view(q_token_num, topk_size), - "chunk_cache_seqlens": buffers["chunk_cache_seqlens"][:q_token_num], - "cu_seqlens_q": buffers["cu_seqlens_q"][: q_token_num + 1], - } - - def create_att_prefill_state(self, infer_state): - return Fp8Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state): - config = self.model.config - # 2048 pooled entries plus at most three kpool remainder entries fit the 2112 rounded direct window. - bound = 2051 if config["index_topk"] == 2048 and config["index_kpool"] == 4 else None - return Fp8Glm5NextSparseDecodeState( - backend=self, _infer_state_weak_ref=weakref.ref(infer_state), decode_topk_bound=bound - ) - - -@dataclasses.dataclass -class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState): - @torch.no_grad() - def _nsa_prefill_att(self, q, packed_kv, att_control): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( - _gather_prefill_kv_kernel, - _compact_topk_indices_kernel, - _merge_att_output_kernel, - ) - - nsa_prefill_dict = att_control.nsa_prefill_dict - infer_state = self.infer_state - q_token_num, q_head_num, _ = q.shape - prefill_cache_kv = nsa_prefill_dict["prefill_cache_kv"] - topk_indices = nsa_prefill_dict["topk_indices"] - softmax_scale = nsa_prefill_dict["softmax_scale"] - req_to_token_indexs = infer_state.req_manager.req_to_token_indexs - b_req_idx = infer_state.b_req_idx - b1_cu_kv_seq_len = infer_state.b1_cu_kv_seq_len - b1_cu_q_seq_len = infer_state.b1_cu_q_seq_len - total_token_num = infer_state.total_token_num - workspace = self.backend.get_prefill_workspace(infer_state, topk_indices) - gathered_kv = workspace["gathered_kv"] - gather_chunk_size = min(gathered_kv.shape[0], total_token_num) - chunk_topk_indices = workspace["chunk_topk_indices"] - chunk_cache_seqlens = workspace["chunk_cache_seqlens"] - cu_seqlens_q = workspace["cu_seqlens_q"] - has_history = total_token_num > q_token_num - if not has_history and prefill_cache_kv.stride(-1) != 1: - prefill_cache_kv = prefill_cache_kv.contiguous() - output = softmax_lse = None - for chunk_start in range(0, total_token_num, gather_chunk_size): - chunk_token_num = min(gather_chunk_size, total_token_num - chunk_start) - if has_history: - _gather_prefill_kv_kernel[(chunk_token_num // 2 + b_req_idx.numel(),)]( - packed_kv, - prefill_cache_kv, - gathered_kv, - req_to_token_indexs, - b_req_idx, - b1_cu_kv_seq_len, - b1_cu_q_seq_len, - chunk_start, - chunk_token_num, - packed_kv.stride(0), - prefill_cache_kv.stride(0), - prefill_cache_kv.stride(2), - req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), - req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), - batch_size=b_req_idx.numel(), - BLOCK_ROWS=2, - num_warps=2, - ) - chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) - else: - # Current BF16 KV already has the complete request layout. - chunk_kv = prefill_cache_kv[chunk_start : chunk_start + chunk_token_num].view(-1, 1, 1, 512) - _compact_topk_indices_kernel[(q_token_num,)]( - topk_indices, - self.q_kv_start_loc, - chunk_topk_indices, - chunk_cache_seqlens, - chunk_start, - chunk_start + chunk_token_num, - topk_indices.stride(0), - topk_indices.stride(1), - TOPK=topk_indices.shape[1], - BLOCK=triton.next_power_of_2(topk_indices.shape[1]), - ) - chunk_output, chunk_softmax_lse, *_ = flash_attn_with_kvcache( - # only_qv ignores Q/K; these views avoid allocating placeholder tensors. - q=q[..., :64], - k_cache=chunk_kv[..., :64], - qv=q, - v_cache=chunk_kv, - page_table=chunk_topk_indices, - cache_seqlens=chunk_cache_seqlens, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=1, - softmax_scale=softmax_scale, - causal=False, - only_qv=True, - num_splits=1, - return_softmax_lse=True, - ) - if chunk_start == 0: - output, softmax_lse = chunk_output, chunk_softmax_lse - else: - # Keep the online reduction in FP32 across all remaining chunks. - output = output.float() - _merge_att_output_kernel[(q_token_num, triton.cdiv(q_head_num, 4))]( - output, - softmax_lse, - chunk_output, - chunk_softmax_lse, - chunk_softmax_lse.stride(1), - chunk_softmax_lse.stride(0), - q_head_num=q_head_num, - q_token_num=q_token_num, - BLOCK_HEADS=4, - ) - del chunk_output, chunk_softmax_lse - return output.to(q.dtype) - - -@dataclasses.dataclass -class Fp8Glm5NextSparseDecodeState(Glm5NextSparseDecodeState): - decode_topk_bound: int = None - flashmla_sched_meta: object = None - - def init_state(self): - import flash_mla - - self.q_kv_seq_len = self.infer_state.b_seq_len - self.q_req_idx = self.infer_state.b_req_idx - self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata() - - def _nsa_decode_att(self, q, packed_kv, att_control): - import flash_mla - - from lightllm.models.glm5_next.triton_kernel import fp8_sparse_decode - from lightllm.models.glm5_next.triton_kernel.pad_query_heads import pad_query_heads - - nsa_decode_dict = att_control.nsa_decode_dict - topk_mem_indices = nsa_decode_dict["topk_mem_indices"] - if topk_mem_indices.ndim == 2: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" - - q_nope, _ = q - num_heads = q_nope.shape[1] - if fp8_sparse_decode.can_use_fp8_sparse_direct_decode( - q_nope, - packed_kv, - topk_mem_indices, - nsa_decode_dict["kv_lora_rank"], - self.decode_topk_bound, - is_cuda_graph_capture=self.infer_state.is_cuda_graph_capture, - ): - return fp8_sparse_decode.fp8_sparse_direct_decode( - q_nope, packed_kv, topk_mem_indices, nsa_decode_dict["softmax_scale"] - ) - padded_q = pad_query_heads(q_nope, is_cuda_graph_capture=self.infer_state.is_cuda_graph_capture) - kv_cache = torch.as_strided( - packed_kv, - size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), - stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), - ) - output, _ = flash_mla.flash_mla_with_kvcache( - q=padded_q, - k_cache=kv_cache, - block_table=None, - cache_seqlens=None, - head_dim_v=nsa_decode_dict["kv_lora_rank"], - tile_scheduler_metadata=self.flashmla_sched_meta, - softmax_scale=nsa_decode_dict["softmax_scale"], - causal=False, - is_fp8_kvcache=True, - indices=topk_mem_indices, - ) - return output[:, 0, :num_heads, :] diff --git a/lightllm/common/basemodel/attention/nsa/glm5_next.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py deleted file mode 100644 index 6e33177dbe..0000000000 --- a/lightllm/common/basemodel/attention/nsa/glm5_next.py +++ /dev/null @@ -1,100 +0,0 @@ -import dataclasses -import weakref - -import torch - -from ..base_att import AttControl, BaseAttBackend, BasePrefillAttState, BaseDecodeAttState - - -class Glm5NextSparseAttBackend(BaseAttBackend): - def create_att_prefill_state(self, infer_state): - return Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state): - return Glm5NextSparseDecodeState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Glm5NextSparsePrefillState(BasePrefillAttState): - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None - q_kv_start_loc: torch.Tensor = None - - def init_state(self): - from lightllm.models.glm5_next.triton_kernel.kpool import prefill_nsa_params - - infer_state = self.infer_state - self.q_kv_seq_len, self.q_req_idx, self.q_kv_start_loc = prefill_nsa_params( - infer_state.position_ids, infer_state.b_req_idx, infer_state.b1_cu_q_seq_len, infer_state.b1_cu_kv_seq_len - ) - - def prefill_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): - assert att_control.nsa_prefill and att_control.nsa_prefill_dict is not None - return self._nsa_prefill_att(q, k, att_control) - - def _nsa_prefill_att(self, q, kv, att_control): - from sgl_kernel.flash_mla import flash_mla_sparse_fwd - - q_token_num, q_head_num, head_dim = q.shape - # FlashMLA uses 64-head tiles. - padded_q_head_num = ((q_head_num + 63) // 64) * 64 - padded_q = q - if padded_q_head_num != q_head_num: - padded_q = q.new_zeros((q_token_num, padded_q_head_num, head_dim)) - padded_q[:, :q_head_num] = q - nsa_prefill_dict = att_control.nsa_prefill_dict - output, _, _ = flash_mla_sparse_fwd( - q=padded_q, - kv=kv, - indices=nsa_prefill_dict["topk_mem_indices"].unsqueeze(1), - sm_scale=nsa_prefill_dict["softmax_scale"], - d_v=nsa_prefill_dict["kv_lora_rank"], - ) - return output[:, :q_head_num] - - -@dataclasses.dataclass -class Glm5NextSparseDecodeState(BaseDecodeAttState): - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None - nsa_cache_seqlens: torch.Tensor = None - nsa_cu_seqlens_k_new: torch.Tensor = None - - def init_state(self): - self.q_kv_seq_len = self.infer_state.b_seq_len - self.q_req_idx = self.infer_state.b_req_idx - index_kpool = self.backend.model.config["index_kpool"] - index_topk = self.backend.model.config["index_topk"] - self.nsa_cache_seqlens = ( - torch.minimum( - self.q_kv_seq_len // index_kpool * index_kpool, torch.full_like(self.q_kv_seq_len, index_topk) - ) - + self.q_kv_seq_len % index_kpool - ) - self.nsa_cu_seqlens_k_new = torch.nn.functional.pad(self.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0)) - - def decode_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): - assert att_control.nsa_decode and att_control.nsa_decode_dict is not None - return self._nsa_decode_att(q, k, att_control) - - def _nsa_decode_att(self, q, kv, att_control): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - nsa_decode_dict = att_control.nsa_decode_dict - q_nope, _ = q - kv_nope = kv.view(-1, 1, 1, nsa_decode_dict["kv_lora_rank"]) - # NoPE has no RoPE Q/K tensors; only_qv uses q_nope and kv_nope directly. - return flash_attn_with_kvcache( - q=None, - qv=q_nope, - k_cache=None, - v_cache=kv_nope, - page_table=nsa_decode_dict["topk_mem_indices"], - cache_seqlens=self.nsa_cache_seqlens, - cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, - cu_seqlens_k_new=self.nsa_cu_seqlens_k_new, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=nsa_decode_dict["softmax_scale"], - causal=False, - only_qv=True, - ) diff --git a/lightllm/common/basemodel/attention/triton/__init__.py b/lightllm/common/basemodel/attention/triton/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py deleted file mode 100644 index 8da44d6ecc..0000000000 --- a/lightllm/common/basemodel/attention/triton/fp.py +++ /dev/null @@ -1,349 +0,0 @@ -import dataclasses -import weakref -from typing import Optional - -import torch - -from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState - - -class TritonAttBackend(BaseAttBackend): - def create_att_prefill_state(self, infer_state) -> "TritonPrefillAttState": - return TritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "TritonDecodeAttState": - return TritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class TritonPrefillAttState(BasePrefillAttState): - def init_state(self): - pass - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - if att_control.use_alibi: - assert att_control.use_sliding_window is False, "alibi + sliding_window not supported" - assert att_control.tp_alibi is not None - return self._alibi_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - else: - return self._nomarl_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - - def _alibi_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - out = alloc_func(q.shape, q.dtype) - - from ...triton_kernel.alibi_att.context_flashattention_nopad import context_attention_fwd - - context_attention_fwd( - q, - k, - v, - out, - self.infer_state.b_req_idx, - att_control.tp_alibi, - self.infer_state.b_q_start_loc, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - self.infer_state.req_manager.req_to_token_indexs, - ) - return out - - def _nomarl_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( - context_attention_fwd, - ) - - if att_control.use_sliding_window: - sliding_window = att_control.sliding_window - else: - sliding_window = (-1, -1) - - out = alloc_func(q.shape, q.dtype) - context_attention_fwd( - q, - k, - v, - out, - self.infer_state.b_req_idx, - self.infer_state.b_q_start_loc, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - self.infer_state.req_manager.req_to_token_indexs, - sliding_window=sliding_window, - ) - return out - - -@dataclasses.dataclass -class TritonDecodeAttState(BaseDecodeAttState): - b_mark_mtp_shared_group: torch.Tensor = None - decode_max_kv_seq_len: int = None - - def init_state(self): - # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于 GQA decode 配置查找。 - self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if draft_step > 0: - self.b_mark_mtp_shared_group = build_mtp_shared_group_markers( - self.infer_state.b_req_idx, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - - def copy_for_decode_cuda_graph(self, new_state: "TritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - if att_control.use_alibi: - assert att_control.use_sliding_window is False, "alibi + sliding_window not supported" - assert att_control.tp_alibi is not None - return self._alibi_decode_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - else: - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - - q_head_num = q.shape[1] - k_head_num = k.shape[1] - - if draft_step > 0: - assert q_head_num >= k_head_num, "speculative decode requires q_head_num >= k_head_num" - return self._spec_decode_gqa_att(q=q, k=k, v=v, alloc_func=alloc_func) - elif q_head_num == k_head_num: - assert att_control.use_sliding_window is False, "sliding_window not supported in non-gqa attention yet" - return self._normal_decode_flash_decoding_att(q=q, k=k, v=v, alloc_func=alloc_func) - elif q_head_num > k_head_num: - return self._normal_decode_gqa_flash_decoding_att( - q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func - ) - else: - raise NotImplementedError("error") - - def _alibi_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - from ...triton_kernel.alibi_att.token_flashattention_nopad import token_attention_fwd - - out = alloc_func(q.shape, q.dtype) - token_attention_fwd( - q, - k, - v, - out, - att_control.tp_alibi, - self.infer_state.req_manager.req_to_token_indexs, - self.infer_state.b_req_idx, - self.infer_state.b_kv_start_loc, - self.infer_state.b_seq_len, - self.infer_state.max_kv_seq_len, - self.infer_state.total_token_num, - alloc_tensor_func=alloc_func, - ) - return out - - def _normal_decode_flash_decoding_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding import ( - token_decode_attention_flash_decoding, - ) - - out = alloc_func(q.shape, q.dtype) - - token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_v=v, - out=out, - alloc_tensor_func=alloc_func, - ) - return out - - def _normal_decode_gqa_flash_decoding_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( - gqa_token_decode_attention_flash_decoding, - ) - - if att_control.use_sliding_window: - sliding_window = att_control.sliding_window - else: - sliding_window = (-1, -1) - - out = alloc_func(q.shape, q.dtype) - - gqa_token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_v=v, - max_len_in_batch=self.decode_max_kv_seq_len, - out=out, - alloc_tensor_func=alloc_func, - sliding_window=sliding_window, - ) - - return out - - def _spec_decode_gqa_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.gqa.mtp_diverse import token_decode_attention_mtp_diverse_single_token - - out = token_decode_attention_mtp_diverse_single_token( - q=q, - k=k, - v=v, - Req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - B_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - b_mark_shared_group=self.b_mark_mtp_shared_group, - max_kv_len=self.decode_max_kv_seq_len, - alloc_tensor_func=alloc_func, - ) - - return out - - def _normal_decode_gqa_flash_decoding_att_vsm( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_vsm import ( - gqa_token_decode_attention_flash_decoding_vsm, - ) - - out = alloc_func(q.shape, q.dtype) - - gqa_token_decode_attention_flash_decoding_vsm( - q=q, - k=k, - v=v, - infer_state=self.infer_state, - out=out, - alloc_tensor_func=alloc_func, - ) - return out - - def _normal_decode_gqa_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer_weight, - alloc_func=torch.empty, - ): - # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.gqa_decode_flashattention_nopad import ( - gqa_decode_attention_fwd, - ) - - out = alloc_func(q.shape, q.dtype) - - gqa_decode_attention_fwd( - q=q, - k=k, - v=v, - out=out, - req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - ) - return out - - def _normal_decode_stage3_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - total_token_num = self.infer_state.total_token_num - batch_size = self.infer_state.batch_size - q_head_num = q.shape[1] - head_dim = q.shape[2] - - calcu_shape1 = (batch_size, q_head_num, head_dim) - att_m_tensor = alloc_func((q_head_num, total_token_num), torch.float32) - - from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.stage3_decode_att import ( - token_attention_nopad_att1, - ) - - token_attention_nopad_att1.token_att_fwd( - q.view(calcu_shape1), - k, - att_m_tensor, - Req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - B_req_idx=self.infer_state.b_req_idx, - B_Start_Loc=self.infer_state.b_kv_start_loc, - B_Seqlen=self.infer_state.b_seq_len, - max_len_in_batch=self.infer_state.max_kv_seq_len, - ) - - o_tensor = alloc_func(q.shape, q.dtype) - from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.stage3_decode_att import ( - token_attention_softmax_and_reducev, - ) - - token_attention_softmax_and_reducev.token_softmax_reducev_fwd( - att_m_tensor, - v, - o_tensor.view(calcu_shape1), - req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_start_loc=self.infer_state.b_kv_start_loc, - b_seq_len=self.infer_state.b_seq_len, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py deleted file mode 100644 index 269b512c1a..0000000000 --- a/lightllm/common/basemodel/attention/triton/int4kv.py +++ /dev/null @@ -1,176 +0,0 @@ -import dataclasses -import weakref -from typing import Optional, Tuple - -import torch - -from lightllm.utils.envs_utils import get_env_start_args - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState - - -class Int4kvTritonAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model) - self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size - - def create_att_prefill_state(self, infer_state) -> "Int4kvTritonPrefillAttState": - return Int4kvTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Int4kvTritonDecodeAttState": - return Int4kvTritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Int4kvTritonPrefillAttState(BasePrefillAttState): - # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。 - b_kv_start_loc: torch.Tensor = None - - def init_state(self): - self.b_kv_start_loc = ( - torch.cumsum(self.infer_state.b_seq_len, dim=0, dtype=self.infer_state.b_seq_len.dtype) - - self.infer_state.b_seq_len - ) - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - - self.backend: Int4kvTritonAttBackend = self.backend # for typing - - k, k_scale = k - v, v_scale = v - o = self._groupsize_quant_prefill_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - return o - - def _groupsize_quant_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - # o_tensor = alloc_func(q.shape, q.dtype, device=q.device) - # batch_size = self.infer_state.b_seq_len.shape[0] - - total_token_num = self.infer_state.total_token_num - head_dim = k.shape[2] * 2 # 2个4bit存储为一个int8, 所以维度需要翻倍，才是解量化后的精度 - k_dequant = alloc_func((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device) - v_dequant = alloc_func((total_token_num, v.shape[1], head_dim), dtype=q.dtype, device=q.device) - o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) - - max_kv_seq_len = self.infer_state.max_kv_seq_len - - from ...triton_kernel.kv_copy.ppl_int4kv_copy_kv import dequantize_int4kv - - dequantize_int4kv( - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_seq_len=self.infer_state.b_seq_len, - b_req_idx=self.infer_state.b_req_idx, - b_kv_start_loc=self.b_kv_start_loc, - k_out=k_dequant, - v_out=v_dequant, - max_len_in_batch=max_kv_seq_len, - quant_group_size=self.backend.quant_group_size, - ) - - from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( - context_attention_fwd_contiguous_kv, - ) - - context_attention_fwd_contiguous_kv( - q=q, - k=k_dequant, - v=v_dequant, - o=o_tensor, - b_start_loc=self.infer_state.b_q_start_loc, - b_kv_start_loc=self.b_kv_start_loc, - b_seq_len=self.infer_state.b_seq_len, - max_q_input_len=self.infer_state.max_q_seq_len, - b_prompt_cache_len=self.infer_state.b_ready_cache_len, - ) - return o_tensor - - -@dataclasses.dataclass -class Int4kvTritonDecodeAttState(BaseDecodeAttState): - decode_max_kv_seq_len: int = None - - def init_state(self): - # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于配置查找。 - self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len - - def copy_for_decode_cuda_graph(self, new_state: "Int4kvTritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - k, k_scale = k - v, v_scale = v - - return self.ppl_int4kv_decode_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - - def ppl_int4kv_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - from ...triton_kernel.att.decode_att.int4kv.ppl_int4kv_flash_decoding import ( - token_decode_attention_flash_decoding, - ) - - return token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_k_scale=k_scale, - cache_v=v, - cache_v_scale=v_scale, - max_kv_seq_len=self.decode_max_kv_seq_len, - alloc_tensor_func=alloc_func, - ) diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py deleted file mode 100644 index 85f17a13cc..0000000000 --- a/lightllm/common/basemodel/attention/triton/int8kv.py +++ /dev/null @@ -1,209 +0,0 @@ -import dataclasses -import weakref -from typing import Optional, Tuple - -import torch - -from lightllm.common.basemodel.triton_kernel.diverse_utils import build_diverse_shared_group_markers -from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel, get_env_start_args - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState - - -class Int8kvTritonAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model) - self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size - - def create_att_prefill_state(self, infer_state) -> "Int8kvTritonPrefillAttState": - return Int8kvTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Int8kvTritonDecodeAttState": - return Int8kvTritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Int8kvTritonPrefillAttState(BasePrefillAttState): - # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。 - b_kv_start_loc: torch.Tensor = None - - def init_state(self): - self.b_kv_start_loc = ( - torch.cumsum(self.infer_state.b_seq_len, dim=0, dtype=self.infer_state.b_seq_len.dtype) - - self.infer_state.b_seq_len - ) - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - - self.backend: Int8kvTritonAttBackend = self.backend # for typing - - k, k_scale = k - v, v_scale = v - o = self._groupsize_quant_prefill_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - return o - - def _groupsize_quant_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - # 反量化内核分别使用 K/V 及 scale 的指针和 stride，兼容共享或独立存储。 - total_token_num = self.infer_state.total_token_num - k_dequant = alloc_func((total_token_num, k.shape[1], k.shape[2]), dtype=q.dtype, device=q.device) - v_dequant = alloc_func((total_token_num, v.shape[1], v.shape[2]), dtype=q.dtype, device=q.device) - o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) - - max_kv_seq_len = self.infer_state.max_kv_seq_len - - from ...triton_kernel.kv_copy.ppl_int8kv_copy_kv import dequantize_int8kv - - dequantize_int8kv( - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_seq_len=self.infer_state.b_seq_len, - b_req_idx=self.infer_state.b_req_idx, - b_kv_start_loc=self.b_kv_start_loc, - k_out=k_dequant, - v_out=v_dequant, - max_len_in_batch=max_kv_seq_len, - quant_group_size=self.backend.quant_group_size, - ) - - from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( - context_attention_fwd_contiguous_kv, - ) - - context_attention_fwd_contiguous_kv( - q=q, - k=k_dequant, - v=v_dequant, - o=o_tensor, - b_start_loc=self.infer_state.b_q_start_loc, - b_kv_start_loc=self.b_kv_start_loc, - b_seq_len=self.infer_state.b_seq_len, - max_q_input_len=self.infer_state.max_q_seq_len, - b_prompt_cache_len=self.infer_state.b_ready_cache_len, - ) - return o_tensor - - -@dataclasses.dataclass -class Int8kvTritonDecodeAttState(BaseDecodeAttState): - b_shared_seq_len: torch.Tensor = None - b_mark_shared_group: torch.Tensor = None - decode_max_kv_seq_len: int = None - - def init_state(self): - # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于普通 decode 配置查找。 - self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len - if enable_diverse_mode_gqa_decode_fast_kernel(): - self.b_mark_shared_group = build_diverse_shared_group_markers( - b_shared_radix_node_id=self.infer_state.b_shared_radix_node_id, - ) - # A one-row group has no cross-request prefix sharing to accelerate. - self.b_shared_seq_len = torch.where( - self.b_mark_shared_group == 1, - 0, - self.infer_state.b_shared_seq_len, - ) - - def copy_for_decode_cuda_graph(self, new_state: "Int8kvTritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - k, k_scale = k - v, v_scale = v - if enable_diverse_mode_gqa_decode_fast_kernel(): - return self.diverse_decode_att(q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func) - else: - return self.normal_decode_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - - def diverse_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - from ...triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import ( - token_decode_attention_flash_decoding, - ) - - return token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_k_scale=k_scale, - cache_v=v, - cache_v_scale=v_scale, - alloc_tensor_func=alloc_func, - ) - - def normal_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.int8kv.normal import token_decode_attention_flash_decoding - - return token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_k_scale=k_scale, - cache_v=v, - cache_v_scale=v_scale, - max_len_in_batch=self.decode_max_kv_seq_len, - alloc_tensor_func=alloc_func, - ) diff --git a/lightllm/common/basemodel/attention/triton/mla.py b/lightllm/common/basemodel/attention/triton/mla.py deleted file mode 100644 index 6f5624d0f4..0000000000 --- a/lightllm/common/basemodel/attention/triton/mla.py +++ /dev/null @@ -1,164 +0,0 @@ -import dataclasses -import weakref -from typing import Tuple - -import torch - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState - - -class MlaTritonAttBackend(BaseAttBackend): - def create_att_prefill_state(self, infer_state) -> "MlaTritonPrefillAttState": - return MlaTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "MlaTritonDecodeAttState": - return MlaTritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class MlaTritonPrefillAttState(BasePrefillAttState): - def init_state(self): - pass - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - if q.shape[2] == 256: - # glm model - return self._contiguous_kv_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - else: - return self._mla_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - - def _mla_prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - from ...triton_kernel.mla_att.prefill_att import context_attention_fwd_with_v - - qk_rope_head_dim = 64 - q_nope, q_rope = q[:, :, :-qk_rope_head_dim], q[:, :, -qk_rope_head_dim:] - o_tensor = alloc_func((q_nope.shape[0], q_nope.shape[1], v.shape[-1]), dtype=q_nope.dtype, device=q.device) - k_nope, k_rope = k - assert att_control.mla_prefill - softmax_scale = att_control.mla_prefill_dict["softmax_scale"] - context_attention_fwd_with_v( - q_nope, - q_rope, - k_nope, - k_rope, - v, - o_tensor, - self.infer_state.b_q_start_loc, - self.infer_state.b1_cu_kv_seq_len, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - softmax_scale, - ) - return o_tensor - - def _contiguous_kv_prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( - context_attention_fwd_contiguous_kv, - ) - - k_nope, k_rope = k - q_head_num = q.shape[1] - - k_merged = torch.cat([k_nope.expand(-1, q_head_num, -1), k_rope.expand(-1, q_head_num, -1)], dim=-1) - - o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) - - context_attention_fwd_contiguous_kv( - q, - k_merged, - v, - o_tensor, - self.infer_state.b_q_start_loc, - self.infer_state.b1_cu_kv_seq_len, - self.infer_state.b_seq_len, - self.infer_state.max_q_seq_len, - self.infer_state.b_ready_cache_len, - ) - return o_tensor - - -@dataclasses.dataclass -class MlaTritonDecodeAttState(BaseDecodeAttState): - def init_state(self): - pass - - def copy_for_decode_cuda_graph(self, new_state: "MlaTritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_sliding_window is False - and att_control.use_att_sink is False - and att_control.use_alibi is False - ) - assert v is None - - return self._mla_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - assert att_control.mla_decode - softmax_scale = att_control.mla_decode_dict["softmax_scale"] - - from ...triton_kernel.mla_att.decode_att import gqa_token_decode_attention_flash_decoding - - qk_rope_head_dim = 64 - q_nope, q_rope = q - kv = k - - out = gqa_token_decode_attention_flash_decoding( - q_nope=q_nope, - q_rope=q_rope, - kv_nope=kv[:, :, :-qk_rope_head_dim], - kv_rope=kv[:, :, -qk_rope_head_dim:], - infer_state=self.infer_state, - softmax_scale=softmax_scale, - alloc_tensor_func=alloc_func, - ) - return out diff --git a/lightllm/common/basemodel/attention/triton/neo.py b/lightllm/common/basemodel/attention/triton/neo.py deleted file mode 100644 index f4fb1cc030..0000000000 --- a/lightllm/common/basemodel/attention/triton/neo.py +++ /dev/null @@ -1,50 +0,0 @@ -import dataclasses -import weakref -from typing import Tuple - -import torch - -from ..base_att import AttControl -from .fp import TritonAttBackend, TritonDecodeAttState, TritonPrefillAttState - - -class NeoTritonAttBackend(TritonAttBackend): - def create_att_prefill_state(self, infer_state) -> "NeoTritonPrefillAttState": - return NeoTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NeoTritonPrefillAttState(TritonPrefillAttState): - def init_state(self): - pass - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( - context_attention_fwd_neo, - ) - from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo - - self.infer_state: NeoChatInferStateInfo - - out = alloc_func(q.shape, q.dtype) - context_attention_fwd_neo( - q, - k, - v, - out, - self.infer_state.b_req_idx, - self.infer_state.b_q_start_loc, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - self.infer_state.req_manager.req_to_token_indexs, - self.infer_state.b_image_token_end, - ) - return out diff --git a/lightllm/common/basemodel/attention_batch.py b/lightllm/common/basemodel/attention_batch.py new file mode 100644 index 0000000000..2c50d87823 --- /dev/null +++ b/lightllm/common/basemodel/attention_batch.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING + +from lightllm.common.attention.api import AttentionBatch, AttentionPhase + +if TYPE_CHECKING: + from .infer_struct import InferStateInfo + + +class AttentionBatchBuilder: + @staticmethod + def from_infer_state(state: "InferStateInfo") -> AttentionBatch: + """Borrow input tensors; capture retains their addresses and replay copies only their contents.""" + return AttentionBatch( + phase=AttentionPhase.PREFILL if state.is_prefill else AttentionPhase.DECODE, + batch_size=state.batch_size, + token_num=state.input_ids.shape[0], + total_token_num=state.total_token_num, + max_q_seq_len=state.max_q_seq_len, + max_kv_seq_len=state.max_kv_seq_len, + max_cache_len=state.max_cache_len or 0, + b_req_idx=state.b_req_idx, + b_seq_len=state.b_seq_len, + b_q_seq_len=state.b_q_seq_len, + b1_cu_q_seq_len=state.b1_cu_q_seq_len, + b1_cu_kv_seq_len=state.b1_cu_kv_seq_len, + req_to_token_indexs=state.req_manager.req_to_token_indexs, + hold_req_id=state.req_manager.HOLD_REQUEST_ID, + b_ready_cache_len=state.b_ready_cache_len, + b_mtp_index=state.b_mtp_index, + b_shared_seq_len=state.b_shared_seq_len, + b_shared_radix_node_id=state.b_shared_radix_node_id, + mem_index=state.mem_index, + position_ids=state.position_ids, + b_image_token_end=getattr(state, "b_image_token_end", None), + ) diff --git a/lightllm/common/basemodel/attention_runtime.py b/lightllm/common/basemodel/attention_runtime.py new file mode 100644 index 0000000000..92d8edecef --- /dev/null +++ b/lightllm/common/basemodel/attention_runtime.py @@ -0,0 +1,83 @@ +"""Per-execution attention contexts and Graph replay preparation. + +Request tensors remain owned by InferStateInfo. This runtime borrows them through +AttentionBatch, owns provider contexts, and borrows scratch from model execution slots. +""" + +from dataclasses import replace + +from lightllm.common.attention.api import AttentionPhase, ContextCapacity + +from .attention_batch import AttentionBatchBuilder + + +class AttentionRuntime: + def __init__(self): + self.groups = () + self.contexts = {} + self.batch = None + + def prepare(self, state, model): + from lightllm.common.basemodel.attention import get_attention_workspace + + state.device_runtime = model.device_runtime + if model._attention_storage_generation != state.att_state_manager.storage_generation: + raise RuntimeError("attention storage changed; rebind layers and recapture CUDA graphs") + phase = AttentionPhase.PREFILL if state.is_prefill else AttentionPhase.DECODE + self.batch = AttentionBatchBuilder.from_infer_state(state) + if state.is_cuda_graph_replay: + return + capacity = ContextCapacity( + phase, + state.batch_size, + state.input_ids.shape[0], + model.graph_max_len_in_batch if state.is_cuda_graph_capture else state.max_kv_seq_len, + state.max_q_seq_len, + ) + for group_id, group in enumerate(self.groups): + workspace = get_attention_workspace(model, group_id, state.microbatch_index, state.input_ids.device) + workspace.begin() + backend = group.backend_for(phase) + if state.is_cuda_graph_capture: + workspace = workspace.fork() + workspace.reserve(backend.resource_requirements(capacity, capture=state.is_cuda_graph_capture)) + context = backend.allocate_context(capacity, capture=state.is_cuda_graph_capture, workspace=workspace) + backend.prepare(context, self.batch) + self.contexts[group_id] = context + + def attention_context(self, binding): + context = self.contexts[binding.group_id] + if context.backend_identity is not binding.backend or context.capacity.phase is not binding.spec.phase: + raise ValueError("layer binding does not belong to this attention execution") + return context + + def update(self): + for group_id, context in self.contexts.items(): + self.groups[group_id].backend_for(context.capacity.phase).update(context) + + def finish(self): + for context in self.contexts.values(): + context.workspace.finish() + + def check_replay_source(self, source): + if len(self.groups) != len(source.groups) or any( + current is not incoming for current, incoming in zip(self.groups, source.groups) + ): + raise ValueError("captured attention belongs to a different model role or group") + + def prepare_replay(self, source): + """Copy borrowed inputs in place and refresh provider plans outside the Graph.""" + self.check_replay_source(source) + for context in self.contexts.values(): + context.capacity.check(source.batch) + for context in self.contexts.values(): + context.workspace.begin() + self.batch.copy_inputs(source.batch) + batch = replace( + self.batch, + max_kv_seq_len=source.batch.max_kv_seq_len, + max_q_seq_len=source.batch.max_q_seq_len, + total_token_num=source.batch.total_token_num, + ) + for group_id, context in self.contexts.items(): + self.groups[group_id].backend_for(context.capacity.phase).prepare(context, batch) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 0edbdadb94..39cb311e9a 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1,56 +1,52 @@ -import os - -# os.environ["CUDA_LAUNCH_BLOCKING"] = "1" -import gc import copy -import math +import gc import json +import math +import os +from functools import cached_property +from typing import TYPE_CHECKING, List, final + import torch import torch.nn.functional as F import triton -from typing import final, List, TYPE_CHECKING from tqdm import tqdm -from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.req_manager import ReqManager from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.commit import BoundedStateCommit +from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput, PostLayerOutput +from lightllm.common.basemodel.cuda_graph import CudaGraph +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights +from lightllm.common.basemodel.mtp_manager import MtpManager +from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( select_kv_index_from_req, select_kv_index_from_req_prefill, ) -from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph -from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed -from lightllm.utils.log_utils import init_logger -from lightllm.utils.dist_utils import get_dp_world_size, get_current_rank_in_dp -from lightllm.utils.envs_utils import ( - get_env_start_args, - get_llm_data_type, -) +from lightllm.common.quantization import Quantcfg +from lightllm.common.req_manager import ReqManager +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner from lightllm.distributed.communication_op import dist_group_manager -from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput, PostLayerOutput -from lightllm.common.basemodel.hidden_collector import ( - NoopHiddenCollector, -) -from lightllm.common.basemodel.mtp_manager import MtpManager +from lightllm.platform import get_device_runtime from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type from lightllm.utils.infer_utils import post_empty_cache -from lightllm.utils.torch_memory_saver_utils import ( - TorchMemorySaverWrapper, - MemoryTag, -) -from .attention import get_prefill_att_backend_class, get_decode_att_backend_class -from .attention import BaseAttBackend +from lightllm.utils.log_utils import init_logger +from lightllm.utils.torch_memory_saver_utils import MemoryTag, TorchMemorySaverWrapper if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs + +# os.environ["CUDA_LAUNCH_BLOCKING"] = "1" logger = init_logger(__name__) @@ -58,6 +54,15 @@ class TpPartBaseModel: + @cached_property + def device_runtime(self): + return get_device_runtime(platform=self.args.hardware_platform) + + def prepare_attention_backends(self): + from lightllm.common.basemodel.attention import initialize_attention_groups + + initialize_attention_groups(self) + is_mtp_draft_model = False # weight class @@ -125,8 +130,6 @@ def __init__(self, kvargs): @property def att_state_role(self) -> str: """当前模型的资源前缀；每个 draft 独立拥有 State。""" - from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS - if not self.is_mtp_draft_model: return "main" draft_index = len(self.mtp_previous_draft_models) @@ -153,6 +156,7 @@ def prepare_attention_states(self) -> None: with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): if not self.is_mtp_draft_model: self.att_state_manager = AttentionStateManager(request_capacity=self.max_req_num + 1) + self.att_state_manager.device_runtime = self.device_runtime # 无 attention 的 draft 只共享管理器，不创建自身 State。 if self.is_mtp_draft_model and self.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): return @@ -173,7 +177,7 @@ def prepare_attention_states(self) -> None: request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, - device="cuda", + device=self.device_runtime.device, ) for state in (*full_states, *bounded_states): self.att_state_manager.register_state(self.att_state_role, state) @@ -182,20 +186,20 @@ def init_runtime(self) -> None: """Initialize this model's backends, warmup and graphs after shared allocation.""" self.req_manager.bind_hold_page(self.att_state_manager.full_att_allocator.hold_token_indices) self._check_mem_size() - self._init_att_backend() - self._init_att_backend1() + from lightllm.common.basemodel.attention import initialize_layer_attention + + initialize_layer_attention(self) self._init_hidden_collector() role = self.att_state_role - logger.info( - f"{role} prefill/decode attention: {type(self.prefill_att_backend).__name__} / " - f"{type(self.decode_att_backend).__name__}" - ) - if self.prefill_att_backend1 is not None: + for group_id, group in enumerate(self.attention_groups): logger.info( - f"{role} prefill/decode attention1: {type(self.prefill_att_backend1).__name__} / " - f"{type(self.decode_att_backend1).__name__}" + "%s attention group %d (%s): prefill=%s, decode=%s", + role, + group_id, + group.prefill_spec.family.value, + type(group.prefill_backend).__name__, + type(group.decode_backend).__name__, ) - self._autotune_warmup() self._init_padded_req() self._init_cudagraph() @@ -250,6 +254,11 @@ def _init_weights(self, start_layer_index=0): def load_weights(self, weight_dict: dict): assert weight_dict is None or isinstance(weight_dict, dict), "weight_dict must be a dict or None" + if hasattr(self, "_attention_storage_generation"): + if self.graph is not None: + self.graph.close() + if self.prefill_graph is not None: + self.prefill_graph.close() load_hf_weights( data_type=self.data_type, weight_dir=self.weight_dir_, @@ -259,6 +268,10 @@ def load_weights(self, weight_dict: dict): ) self.pre_post_weight.verify_load() [weight.verify_load() for weight in self.trans_layers_weight] + if hasattr(self, "_attention_storage_generation"): + from lightllm.common.basemodel.attention import initialize_layer_attention + + initialize_layer_attention(self) return @classmethod @@ -336,21 +349,8 @@ def _init_some_value(self): self.vocab_size = self.config["vocab_size"] return - def _init_att_backend(self): - role = self.att_state_role - if not self.att_state_manager.get_full_attention_states(role): - # FFN-only and linear-only models do not need a full-attention backend. - self.prefill_att_backend = self.decode_att_backend = None - return - self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0)(model=self) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) - return - - def _init_att_backend1(self): - # self.prefill_att_backend1 是给后续有模型支持不同层用不同的att模块时，保留的扩展。 - self.prefill_att_backend1: BaseAttBackend = None - self.decode_att_backend1: BaseAttBackend = None - return + def attention_backend_candidates(self, family, phase): + return None, None def _align_decode_batch_size(self, batch_size: int) -> int: alignment = self.mtp_manager.get_decode_batch_alignment(self.is_mtp_draft_model) @@ -369,6 +369,7 @@ def _init_cudagraph(self): None if self.disable_cudagraph else CudaGraph( + runtime=self.device_runtime, batch_step_size_before_split=cuda_graph_grow_step_size, split_batch_size=self.args.graph_split_batch_size * decode_tokens_per_request, batch_step_size_after_split=self.args.graph_grow_step_size * cuda_graph_grow_step_size, @@ -388,7 +389,9 @@ def _init_prefill_cuda_graph(self): self.prefill_graph = ( None if not get_env_start_args().enable_prefill_cudagraph - else PrefillCudaGraph(decode_cuda_graph=self.graph, tp_world_size=self.tp_world_size_) + else PrefillCudaGraph( + decode_cuda_graph=self.graph, tp_world_size=self.tp_world_size_, runtime=self.device_runtime + ) ) if self.prefill_graph is not None: if get_env_start_args().enable_prefill_microbatch_overlap: @@ -467,26 +470,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) # 特殊模型，特殊模式的特定变量初始化操作。 infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens - if infer_state.is_prefill: - infer_state.prefill_att_state = ( - self.prefill_att_backend.create_att_prefill_state(infer_state=infer_state) - if self.prefill_att_backend is not None - else None - ) - if self.prefill_att_backend1 is not None: - infer_state.prefill_att_state1 = self.prefill_att_backend1.create_att_prefill_state( - infer_state=infer_state - ) - else: - infer_state.decode_att_state = ( - self.decode_att_backend.create_att_decode_state(infer_state=infer_state) - if self.decode_att_backend is not None - else None - ) - if self.decode_att_backend1 is not None: - infer_state.decode_att_state1 = self.decode_att_backend1.create_att_decode_state( - infer_state=infer_state - ) + infer_state.attention_groups = self.attention_groups return infer_state @@ -585,6 +569,19 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS ) return output + def _attach_attention_commit(self, output, infer_state, batch_size): + if not self.args.mtp_step: + return output + + output = copy.copy(output) + output.attention_commit = BoundedStateCommit.capture( + self.att_state_role, + infer_state.b_req_idx[:batch_size], + infer_state.b_mtp_index[:batch_size], + self.req_manager.request_generations, + ) + return output + def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int): padded_batch_size = model_output.logits.shape[0] if padded_batch_size == origin_batch_size: @@ -654,7 +651,7 @@ def _prefill( prefill_mem_indexes_ready_event.record() infer_state.init_some_extra_state(self) - infer_state.init_att_state() + infer_state.prepare_attention(self) model_output = self._context_forward(infer_state=infer_state) model_output = self._create_unpad_prefill_model_output( @@ -663,6 +660,7 @@ def _prefill( origin_batch_size=origin_batch_size, ) model_output.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event + infer_state.finish_attention() return model_output def _decode( @@ -708,7 +706,7 @@ def _decode( infer_state.is_cuda_graph_capture = need_capture infer_state.is_cuda_graph_replay = use_cuda_graph and not need_capture infer_state.init_some_extra_state(self) - infer_state.init_att_state() + infer_state.prepare_attention(self) if use_cuda_graph: if need_capture: @@ -718,11 +716,15 @@ def _decode( else: model_output = self._token_forward(infer_state) - return self._create_unpad_decode_model_output(model_output, origin_batch_size=origin_batch_size) + output = self._create_unpad_decode_model_output(model_output, origin_batch_size=origin_batch_size) + output = self._attach_attention_commit(output, infer_state, origin_batch_size) + + infer_state.finish_attention() + return output @final def _context_forward(self, infer_state: InferStateInfo): - + infer_state.update_attention() input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" @@ -788,6 +790,7 @@ def prefill_func(input_tensors, _infer_state): return model_output def _token_forward(self, infer_state: InferStateInfo): + infer_state.update_attention() hidden_collector = infer_state.hidden_collector input_ids = infer_state.input_ids cuda_input_ids = input_ids @@ -817,7 +820,6 @@ def _token_forward(self, infer_state: InferStateInfo): @torch.no_grad() def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: ModelInput): """执行由调用方提前构建好的两个 prefill microbatch。""" - for model_input in (model_input0, model_input1): model_input.to_cuda() if self.args.enable_prefill_decode_mixed and model_input.input_ids.shape[0] > 0: @@ -855,11 +857,11 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input infer_state0 = self._create_inferstate(model_input0, 0) infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() + infer_state0.prepare_attention(self) infer_state1 = self._create_inferstate(model_input1, 1) infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() + infer_state1.prepare_attention(self) prefill_mem_indexes_ready_event = torch.cuda.Event() prefill_mem_indexes_ready_event.record() @@ -881,12 +883,13 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input dist_group_manager.clear_deepep_buffer() model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event + infer_state0.finish_attention() + infer_state1.finish_attention() return model_output0, model_output1 @torch.no_grad() def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: ModelInput): """执行由调用方提前构建好的两个 decode microbatch。""" - for model_input in (model_input0, model_input1): model_input.to_cuda() if model_input.input_ids is None: @@ -921,13 +924,13 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_state0.is_cuda_graph_capture = need_capture infer_state0.is_cuda_graph_replay = not need_capture infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() + infer_state0.prepare_attention(self) infer_state1 = self._create_inferstate(padded_model_input1, 1) infer_state1.is_cuda_graph_capture = need_capture infer_state1.is_cuda_graph_replay = not need_capture infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() + infer_state1.prepare_attention(self) if need_capture: model_output0, model_output1 = self.graph.capture_decode( @@ -948,20 +951,26 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) infer_state0 = self._create_inferstate(model_input0, 0) infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() + infer_state0.prepare_attention(self) infer_state1 = self._create_inferstate(model_input1, 1) infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() + infer_state1.prepare_attention(self) model_output0, model_output1 = self._overlap_tpsp_token_forward(infer_state0, infer_state1=infer_state1) model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size0) model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size1) + model_output0 = self._attach_attention_commit(model_output0, infer_state0, origin_batch_size0) + model_output1 = self._attach_attention_commit(model_output1, infer_state1, origin_batch_size1) + infer_state0.finish_attention() + infer_state1.finish_attention() return model_output0, model_output1 @final def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state1: InferStateInfo): + infer_state.update_attention() + infer_state1.update_attention() g_cache_manager.cache_env_in() hidden_collector0 = infer_state.hidden_collector hidden_collector1 = infer_state1.hidden_collector @@ -1019,6 +1028,8 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state @final def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: InferStateInfo): + infer_state.update_attention() + infer_state1.update_attention() hidden_collector0 = infer_state.hidden_collector hidden_collector1 = infer_state1.hidden_collector input_embs, input_embs1 = self.pre_infer.overlap_tpsp_token_forward( diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d8f2329b24..dfa61d247a 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -4,6 +4,7 @@ import torch +from lightllm.common.attention_state.commit import BoundedStateCommit from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor @@ -189,6 +190,7 @@ class ModelOutput: # attached here. ModelOutput therefore owns a stable output view instead # of the mutable collector used while the forward is still running. mtp_collector: Optional[ModelMtpOutputCollector] = None + attention_commit: Optional[BoundedStateCommit] = None # 用于判断 mem_indexes 是否成功写入 req manager 中的事件对象。 prefill_mem_indexes_ready_event: torch.Event = None diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ab0e0c0721..eed8c1030c 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,22 +1,21 @@ -import os -import torch -import torch.distributed as dist -import copy import bisect +import copy import math -import triton from typing import Optional -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.distributed import dist_group_manager + +import torch +import torch.distributed as dist +import triton + from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -from lightllm.utils.torch_memory_saver_utils import ( - TorchMemorySaverWrapper, - MemoryTag, -) -from .infer_struct import InferStateInfo +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner +from lightllm.platform import get_device_runtime +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger +from lightllm.utils.torch_memory_saver_utils import TorchMemorySaverWrapper +from .graph import GraphExecution +from .infer_struct import InferStateInfo logger = init_logger(__name__) @@ -63,10 +62,12 @@ def __init__( max_len_in_batch=8192, tp_world_size: int = 1, capture_infer_cost: bool = False, + runtime=None, ): self.graph = {} self.tp_world_size = tp_world_size - self.mempool = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None + self.runtime = runtime if runtime is not None else get_device_runtime() + self.mempool = self.runtime.graph_pool() self.args = get_env_start_args() self.max_batch_size = max_batch_size self.graph_max_len_in_batch = max_len_in_batch @@ -90,7 +91,7 @@ def can_run(self, batch_size, max_len_in_batch): def need_capture(self, batch_size): find_batch_size = self.find_closest_graph_batch_size(batch_size) if find_batch_size is not None: - return find_batch_size not in self.graph + return find_batch_size not in self.graph or self.graph[find_batch_size].graph is None else: assert False, "dead code" @@ -103,37 +104,40 @@ def find_closest_graph_batch_size(self, batch_size): return None def _capture_decode(self, decode_func, infer_state: InferStateInfo): - graph_obj = torch.cuda.CUDAGraph() + graph_obj = self.runtime.create_graph() input_ids = infer_state.input_ids batch_size = input_ids.shape[0] infer_state.max_kv_seq_len = self.graph_max_len_in_batch infer_state.total_token_num = self.graph_max_len_in_batch * batch_size - # warmup - # 因为有些推理过程的代码，会通过判断infer_state中是否存在某些属性来在一层上 - # 做一些初始化的操作，后续层可以复用这些计算的结果，如 - # lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py - # 中做的一些操作，所以在 warmup 的时候，需要调用infer_state的copy函数做一个 - # 浅拷贝，不然后续传入到cuda graph捕获过程中后，infer_state因为提前拥有了这些属性， - # 导致不会重新初始化，这样捕获过程中会不能捕获这些临时添加到 infer_state 管理对象 - # 中的 tensor。 - - for _ in range(1): - # 记录原始存在的变量 - pure_para_set = set(vars(infer_state).keys()) - torch.cuda.synchronize() - # 在正式捕获前调优 decode attention，退出作用域后再捕获选定的配置。 + snapshot = infer_state.att_state_manager.snapshot_capture_state( + infer_state.b_req_idx, infer_state.att_state_role + ) + restore_state = lambda: infer_state.att_state_manager.restore_capture_state(snapshot) + try: + self.runtime.synchronize() with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): decode_func(copy.copy(infer_state)) - torch.cuda.synchronize() - for param_name in set(vars(infer_state).keys()): - if param_name not in pure_para_set: - delattr(infer_state, param_name) - - with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool): - model_output = decode_func(infer_state) - self.graph[batch_size] = (graph_obj, infer_state, model_output) - graph_obj.replay() - self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size) + self.runtime.synchronize() + finally: + restore_state() + + try: + with self.runtime.capture_graph(graph_obj, pool=self.mempool, memory_saver=self.torch_memory_saver): + model_output = decode_func(infer_state) + except BaseException: + try: + GraphExecution(graph_obj, (infer_state,), runtime=self.runtime).close() + finally: + restore_state() + raise + try: + execution = GraphExecution(graph_obj, (infer_state,), (model_output,), runtime=self.runtime) + except BaseException: + restore_state() + raise + self.graph[batch_size] = execution + execution.replay() + self._measure_replay_cost(graph_obj=execution, batch_size=batch_size, restore_state=restore_state) return model_output def _capture_decode_overlap( @@ -142,50 +146,58 @@ def _capture_decode_overlap( infer_state: InferStateInfo, infer_state1: InferStateInfo, ): - graph_obj = torch.cuda.CUDAGraph() + graph_obj = self.runtime.create_graph() input_ids = infer_state.input_ids batch_size = input_ids.shape[0] infer_state.max_kv_seq_len = self.graph_max_len_in_batch infer_state.total_token_num = self.graph_max_len_in_batch * batch_size infer_state1.max_kv_seq_len = self.graph_max_len_in_batch infer_state1.total_token_num = self.graph_max_len_in_batch * batch_size - # warmup - for _ in range(1): - # 记录原始存在的变量 - pure_para_set = set(vars(infer_state).keys()) - pure_para_set1 = set(vars(infer_state1).keys()) - torch.cuda.synchronize() + snapshot = infer_state.att_state_manager.snapshot_capture_state( + torch.cat((infer_state.b_req_idx, infer_state1.b_req_idx)), infer_state.att_state_role + ) + restore_state = lambda: infer_state.att_state_manager.restore_capture_state(snapshot) + try: + self.runtime.synchronize() with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): decode_func(copy.copy(infer_state), copy.copy(infer_state1)) - torch.cuda.synchronize() - for para_name in set(vars(infer_state).keys()): - if para_name not in pure_para_set: - delattr(infer_state, para_name) - for para_name in set(vars(infer_state1).keys()): - if para_name not in pure_para_set1: - delattr(infer_state1, para_name) - - with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool): - model_output, model_output1 = decode_func(infer_state, infer_state1) - self.graph[batch_size] = ( - graph_obj, - infer_state, - infer_state1, - model_output, - model_output1, - ) - graph_obj.replay() - self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size) + self.runtime.synchronize() + finally: + restore_state() + + try: + with self.runtime.capture_graph(graph_obj, pool=self.mempool, memory_saver=self.torch_memory_saver): + model_output, model_output1 = decode_func(infer_state, infer_state1) + except BaseException: + try: + GraphExecution(graph_obj, (infer_state, infer_state1), runtime=self.runtime).close() + finally: + restore_state() + raise + try: + execution = GraphExecution( + graph_obj, (infer_state, infer_state1), (model_output, model_output1), runtime=self.runtime + ) + except BaseException: + restore_state() + raise + self.graph[batch_size] = execution + execution.replay() + self._measure_replay_cost(graph_obj=execution, batch_size=batch_size, restore_state=restore_state) return model_output, model_output1 - def _measure_replay_cost(self, graph_obj: torch.cuda.CUDAGraph, batch_size: int) -> None: + def _measure_replay_cost( + self, graph_obj: torch.cuda.CUDAGraph, batch_size: int, restore_state=lambda: None + ) -> None: if not self.capture_infer_cost: return dist.barrier(group=dist.group.WORLD) - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) + start_event = self.runtime.create_event(enable_timing=True) + end_event = self.runtime.create_event(enable_timing=True) + restore_state() graph_obj.replay() + restore_state() start_event.record() graph_obj.replay() end_event.record() @@ -217,37 +229,17 @@ def capture_decode( assert infer_state1 is None return self._capture_decode(decode_func, infer_state) - def _replay(self, infer_state: InferStateInfo): - batch_size = infer_state.input_ids.shape[0] - graph_obj, graph_infer_state, graph_output = self.graph[batch_size] - graph_infer_state.copy_for_cuda_graph(infer_state) - graph_obj.replay() - return graph_output - - def _replay_overlap( - self, - infer_state: InferStateInfo, - infer_state1: InferStateInfo, - ): - batch_size = infer_state.input_ids.shape[0] - ( - graph_obj, - graph_infer_state, - graph_infer_state1, - graph_model_output, - graph_model_output1, - ) = self.graph[batch_size] - graph_infer_state.copy_for_cuda_graph(infer_state) - graph_infer_state1.copy_for_cuda_graph(infer_state1) - graph_obj.replay() - return graph_model_output, graph_model_output1 - def replay(self, infer_state, infer_state1=None): - if self.enable_decode_microbatch_overlap: - return self._replay_overlap(infer_state, infer_state1) - else: - assert infer_state1 is None - return self._replay(infer_state) + batch_size = infer_state.input_ids.shape[0] + states = (infer_state, infer_state1) if self.enable_decode_microbatch_overlap else (infer_state,) + if any(state is None for state in states): + raise ValueError("Graph replay requires all microbatch lanes") + return self.graph[batch_size].replay(*states) + + def close(self): + for execution in self.graph.values(): + execution.close() + self.graph.clear() @torch.no_grad() def warmup(self, model): diff --git a/lightllm/common/basemodel/graph.py b/lightllm/common/basemodel/graph.py index dd03b44242..22d44545cf 100644 --- a/lightllm/common/basemodel/graph.py +++ b/lightllm/common/basemodel/graph.py @@ -16,8 +16,8 @@ def __init__(self, graph, states=(), outputs=(), *, runtime=None): for state in self.states: manager = state.att_state_manager self.generations.append((manager, manager.storage_generation)) - for group_id, context in state.attention_contexts.items(): - backend = state.attention_groups[group_id].backend_for(context.capacity.phase) + for group_id, context in state.attention_runtime.contexts.items(): + backend = state.attention_runtime.groups[group_id].backend_for(context.capacity.phase) binding = backend.bind_graph(context) if binding is not None: self.bindings.append(binding) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 83b312f991..ed6de4894c 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,17 +1,18 @@ -import torch -import triton import collections +from contextlib import contextmanager +from typing import List, Optional + +import torch + from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup -from typing import Tuple, Any, Optional, List -from .triton_kernel.gen_prefill_params import gen_prefill_params -from .triton_kernel.gen_decode_params import gen_decode_params -from .triton_kernel.multimodal_emb import mark_multimodal_obj -from .batch_objs import ModelInput +from lightllm.utils.dist_utils import get_dp_world_size, get_global_dp_rank from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.dist_utils import get_global_dp_rank, get_dp_world_size -from .attention import BasePrefillAttState, BaseDecodeAttState + +from .attention_runtime import AttentionRuntime +from .triton_kernel.gen_decode_params import gen_decode_params +from .triton_kernel.gen_prefill_params import gen_prefill_params class InferStateInfo: @@ -20,13 +21,7 @@ class InferStateInfo: """ def __init__(self): - # prefill 和 decode 使用的 att 状态对象 - self.prefill_att_state: BasePrefillAttState = None - self.decode_att_state: BaseDecodeAttState = None - - # 保留的扩展, 支持线性att与标准att混合使用时使用 - self.prefill_att_state1: BasePrefillAttState = None - self.decode_att_state1: BaseDecodeAttState = None + self.attention_runtime = AttentionRuntime() self.input_ids: torch.Tensor = None self.batch_size: int = None @@ -133,58 +128,64 @@ def init_some_extra_state(self, model): ) = gen_decode_params(self.b_seq_len) self.b_kv_start_loc = self.b1_cu_kv_seq_len[0:-1] - def init_att_state(self): - if self.is_prefill: - if self.prefill_att_state is not None: - self.prefill_att_state.init_state() - if self.prefill_att_state1 is not None: - self.prefill_att_state1.init_state() - else: - if self.decode_att_state is not None: - self.decode_att_state.init_state() - if self.decode_att_state1 is not None: - self.decode_att_state1.init_state() - self.check_att_states() - - def check_att_states(self) -> None: - """检查 attention state，防止其强引用 InferStateInfo 而形成循环引用。 - - InferStateInfo 已经通过下列四个属性强引用主、辅 attention state。 - 如果 attention state 再把 infer_state 存到普通属性中，就会形成 - infer_state -> attention state -> infer_state 的强引用闭环。 - 即使一次推理结束、外部引用全部释放，引用计数也无法立即回收这两个对象， - 必须等待 Python 的循环 GC。循环 GC 主要由 Python 对象分配和阈值驱动， - 不能保证在 GPU 显存紧张时及时运行。 - - 这些对象还持有本次推理的 Tensor。对象延迟回收会使相关 Tensor 的显存 - 继续被占用，无法及时交还给 PyTorch 分配器复用；多次推理累积后可能抬高 - 显存峰值，甚至触发 OOM。集中进行循环 GC 还可能造成推理延迟抖动。 - 因此，即使循环最终能够被 GC 回收，也不能依赖它管理推理临时状态的生命周期。 - - InferStateInfo 的生命周期应由推理调用方或 CUDA Graph 缓存负责： - 调用方使用期间、图仍需回放期间保留强引用，结束使用后即可释放。 - attention state 只通过 _infer_state_weak_ref 和 infer_state 属性访问它， - 不应额外延长其生命周期。这里在 attention state 初始化完成后统一检查， - 尽早发现新后端或后续修改误存的 InferStateInfo 强引用。 - - 注意：检查仅扫描非空 attention state 的第一层实例属性，不递归检查容器或 - 回调闭包；实现后端时也应避免通过这些位置间接保存 infer_state 的强引用。 - """ - for state in (self.prefill_att_state, self.prefill_att_state1, self.decode_att_state, self.decode_att_state1): - if state is not None: - state.check_no_infer_state() + @property + def attention_groups(self): + return self.attention_runtime.groups + + @attention_groups.setter + def attention_groups(self, value): + self.attention_runtime.groups = value + + @property + def attention_contexts(self): + return self.attention_runtime.contexts + + @attention_contexts.setter + def attention_contexts(self, value): + self.attention_runtime.contexts = value + + @property + def attention_batch(self): + return self.attention_runtime.batch + + @attention_batch.setter + def attention_batch(self, value): + self.attention_runtime.batch = value + + def prepare_attention(self, model): + self.attention_runtime.prepare(self, model) + + def attention_context(self, binding): + return self.attention_runtime.attention_context(binding) + + def update_attention(self): + self.attention_runtime.update() + + def finish_attention(self): + self.attention_runtime.finish() def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"): + if self.att_state_role != new_infer_state.att_state_role: + raise ValueError("captured attention belongs to a different model role or group") + self.attention_runtime.check_replay_source(new_infer_state.attention_runtime) + copies = [] + batch_tensor_names = { + name for name, value in vars(self.attention_batch).items() if isinstance(value, torch.Tensor) + } for attr_name, attr_value in vars(new_infer_state).items(): - if isinstance(attr_value, torch.Tensor): - attr_ = getattr(self, attr_name, None) - if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr(): - attr_.copy_(attr_value, non_blocking=True) - - if self.decode_att_state is not None: - self.decode_att_state.copy_for_decode_cuda_graph(new_infer_state.decode_att_state) - if self.decode_att_state1 is not None: - self.decode_att_state1.copy_for_decode_cuda_graph(new_infer_state.decode_att_state1) + if isinstance(attr_value, torch.Tensor) and attr_name not in batch_tensor_names: + target = getattr(self, attr_name, None) + if target is not None and target.data_ptr() != attr_value.data_ptr(): + if (target.shape, target.dtype, target.device) != ( + attr_value.shape, + attr_value.dtype, + attr_value.device, + ): + raise ValueError(f"captured input {attr_name} changed shape, dtype or device") + copies.append((target, attr_value)) + self.attention_runtime.prepare_replay(new_infer_state.attention_runtime) + for target, source in copies: + target.copy_(source, non_blocking=True) return def call_overlap_hook(self): @@ -330,6 +331,7 @@ def _all_to_all_balance_get(self, data: torch.Tensor, change_state: bool = True) self.__change_to_balance() dp_rank = get_global_dp_rank() import torch.distributed as dist + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager old_shape = data.shape @@ -361,6 +363,7 @@ def _all_to_all_unbalance_get(self, data: torch.Tensor, change_state: bool = Tru dp_rank = get_global_dp_rank() import torch.distributed as dist + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager old_shape = data.shape @@ -386,11 +389,20 @@ def _all_to_all_unbalance_get(self, data: torch.Tensor, change_state: bool = Tru return origin_data.view(-1, *old_shape[1:]) # 用于 prefill cuda graph 的专用功能接口 + @contextmanager + def _capture_prefill_segment(self, graph_obj): + with self.device_runtime.capture_graph(graph_obj, pool=self.mem_pool): + self._prefill_capture_active = True + try: + yield + finally: + self._prefill_capture_active = False + def prefill_cuda_graph_create_graph_obj(self): if not hasattr(self, "prefill_cuda_graph_exe_list"): self.prefill_cuda_graph_exe_list = [] - graph_obj = torch.cuda.CUDAGraph() - capture_graph = torch.cuda.graph(graph_obj, pool=self.mem_pool) + graph_obj = self.device_runtime.create_graph() + capture_graph = self._capture_prefill_segment(graph_obj) self.prefill_cuda_graph_exe_list.append((graph_obj, capture_graph)) return @@ -423,6 +435,19 @@ def prefill_replay(self, new_infer_state: "InferStateInfo"): func(new_infer_state) return + def close_prefill_graphs(self): + try: + # 错误可能发生在某一段内部，也可能发生在段间的 CPU 回调中。 + # 只退出仍在捕获的上下文，避免重复结束已经完成的段。 + if getattr(self, "_prefill_capture_active", False): + self.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + finally: + self.device_runtime.synchronize() + for item in getattr(self, "prefill_cuda_graph_exe_list", ()): + if isinstance(item, tuple): + item[0].reset() + self.prefill_cuda_graph_exe_list = [] + def copy_for_prefill_cuda_graph(self, new_infer_state: "InferStateInfo"): for attr_name, attr_value in vars(new_infer_state).items(): if isinstance(attr_value, torch.Tensor): diff --git a/lightllm/common/basemodel/layer_infer/base_layer_infer.py b/lightllm/common/basemodel/layer_infer/base_layer_infer.py index 354920f142..769093e315 100644 --- a/lightllm/common/basemodel/layer_infer/base_layer_infer.py +++ b/lightllm/common/basemodel/layer_infer/base_layer_infer.py @@ -1,16 +1,39 @@ +from typing import Iterable, Union + import torch import torch.distributed as dist -from typing import Dict, Iterable, Literal, Tuple, Union, List + +from lightllm.common.attention.api import AttentionFamily from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight +from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy +from lightllm.distributed.communication_op import all_gather_into_tensor, all_reduce, reduce_scatter_tensor from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size -from .cache_tensor_manager import g_cache_manager from lightllm.utils.envs_utils import get_env_start_args -from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor, all_reduce -from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy + +from .cache_tensor_manager import g_cache_manager class BaseLayerInfer: + attention_family = "dense" + attention_uses_alibi = False + attention_uses_sink = False + attention_variant = "" + + def get_attention_family(self): + return AttentionFamily(self.attention_family) + + def attention_window(self, phase): + return None + + def attention_kv_view(self, state): + return state.get_layer_view(self.layer_num_) + + def create_attention_params(self, model, weight, family): + from lightllm.common.basemodel.attention import create_layer_attention_params + + return create_layer_attention_params(model, self, weight, family) + def __init__(self) -> None: self.tp_rank_ = get_current_rank_in_dp() self.tp_world_size_ = get_dp_world_size() diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index 4f21d90cfc..646785b884 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -1,31 +1,49 @@ -import os -import torch -import copy import bisect +from typing import List, Optional, Tuple + +import torch import triton -from typing import List, Tuple -from typing import Optional -from lightllm.utils.log_utils import init_logger + +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.platform import get_device_runtime from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor -from lightllm.distributed import dist_group_manager -from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from .infer_struct import InferStateInfo + from .cuda_graph import CudaGraph +from .infer_struct import InferStateInfo logger = init_logger(__name__) +class PrefillGraphExecution: + """分段图拥有输入/输出和 collector；attention 仍通过段间回调使用本批次 Plan。""" + + def __init__(self, infer_state, inputs, outputs, collector, runtime): + self.state, self.inputs, self.outputs, self.collector = infer_state, inputs, outputs, collector + self.runtime = runtime + self.generation = infer_state.att_state_manager.storage_generation + infer_state.att_state_manager.register_graph(self) + + def close(self): + if self.state is None: + return + self.state.close_prefill_graphs() + self.state = self.collector = None + self.inputs = self.outputs = () + + class PrefillCudaGraph: # CudaGraph forward pass for the decoding stage. - def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): + def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int, runtime=None): self.graph = {} + self.runtime = runtime or (decode_cuda_graph.runtime if decode_cuda_graph is not None else get_device_runtime()) self.tp_world_size = tp_world_size if decode_cuda_graph is not None: self.mempool = decode_cuda_graph.mempool # prefill 和 decode 共享一个 mempool else: - self.mempool = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None + self.mempool = self.runtime.graph_pool() self.args = get_env_start_args() self.enable_prefill_microbatch_overlap = self.args.enable_prefill_microbatch_overlap @@ -62,7 +80,7 @@ def can_run(self, handle_token_num: int): def need_capture(self, handle_token_num: int): finded_handle_token_num = self.find_closest_graph_handle_token_num(handle_token_num=handle_token_num) if finded_handle_token_num is not None: - return finded_handle_token_num not in self.graph + return finded_handle_token_num not in self.graph or self.graph[finded_handle_token_num].state is None else: assert False, "dead code" @@ -79,35 +97,38 @@ def _capture_prefill( ) -> List[torch.Tensor]: handle_token_num = infer_state.input_ids.shape[0] infer_state.mem_pool = self.mempool - infer_state.prefill_cuda_graph_create_graph_obj() - infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() - graph_input_tensors: List[torch.Tensor] = [torch.empty_like(e) for e in input_tensors] - graph_out_tensors: List[torch.Tensor] = prefill_func(graph_input_tensors, infer_state) - graph_out_tensors = [e.contiguous() for e in graph_out_tensors] - infer_state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) - - graph_input_tensors = [tensor_to_no_ref_tensor(e) for e in graph_input_tensors] - graph_out_tensors = [tensor_to_no_ref_tensor(e) for e in graph_out_tensors] - - graph_hidden_collector = infer_state.hidden_collector - hidden_tensor_count, hidden_tensor_nbytes = graph_hidden_collector.release_graph_tensor_ownership() - logger.info( - f"Prefill CUDA Graph hidden collector 已完成 capture 状态托管：" - f"handle_token_num={handle_token_num}, " - f"collector={graph_hidden_collector.__class__.__name__}, " - f"no_ref_tensor_count={hidden_tensor_count}, " - f"no_ref_tensor_nbytes={hidden_tensor_nbytes}。" - f"这些 tensor 的固定地址由 graph memory pool 管理，collector 仅保留无所有权引用供 replay 使用。" - ) - self.graph[handle_token_num] = ( - infer_state, - graph_input_tensors, - graph_out_tensors, - graph_hidden_collector, - ) - self.replay(input_tensors, infer_state) + infer_state.device_runtime = self.runtime + try: + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + graph_input_tensors: List[torch.Tensor] = [torch.empty_like(e) for e in input_tensors] + graph_out_tensors: List[torch.Tensor] = prefill_func(graph_input_tensors, infer_state) + graph_out_tensors = [e.contiguous() for e in graph_out_tensors] + infer_state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + + graph_input_tensors = [tensor_to_no_ref_tensor(e) for e in graph_input_tensors] + graph_out_tensors = [tensor_to_no_ref_tensor(e) for e in graph_out_tensors] + + graph_hidden_collector = infer_state.hidden_collector + hidden_tensor_count, hidden_tensor_nbytes = graph_hidden_collector.release_graph_tensor_ownership() + logger.info( + f"Prefill CUDA Graph hidden collector 已完成 capture 状态托管：" + f"handle_token_num={handle_token_num}, " + f"collector={graph_hidden_collector.__class__.__name__}, " + f"no_ref_tensor_count={hidden_tensor_count}, " + f"no_ref_tensor_nbytes={hidden_tensor_nbytes}。" + f"这些 tensor 的固定地址由 graph memory pool 管理，collector 仅保留无所有权引用供 replay 使用。" + ) + self.graph[handle_token_num] = PrefillGraphExecution( + infer_state, graph_input_tensors, graph_out_tensors, graph_hidden_collector, self.runtime + ) + self.replay(input_tensors, infer_state) - return graph_out_tensors + return graph_out_tensors + except BaseException: + self.graph.pop(handle_token_num, None) + infer_state.close_prefill_graphs() + raise def _capture_prefill_overlap( self, @@ -148,9 +169,17 @@ def capture_prefill( def _replay(self, input_tensors: List[torch.Tensor], infer_state: InferStateInfo) -> List[torch.Tensor]: handle_token_num = infer_state.input_ids.shape[0] - graph_infer_state, graph_input_tensors, graph_output_tensors, graph_hidden_collector = self.graph[ - handle_token_num - ] + execution = self.graph[handle_token_num] + graph_infer_state = execution.state + if graph_infer_state is None: + raise RuntimeError("prefill Graph execution is closed") + if graph_infer_state.att_state_manager.storage_generation != execution.generation: + raise RuntimeError("attention storage changed; recapture prefill Graph") + graph_input_tensors, graph_output_tensors, graph_hidden_collector = ( + execution.inputs, + execution.outputs, + execution.collector, + ) graph_infer_state: InferStateInfo = graph_infer_state for graph_in_tensor, in_tensor in zip(graph_input_tensors, input_tensors): graph_in_tensor.copy_(in_tensor) @@ -165,6 +194,11 @@ def _replay(self, input_tensors: List[torch.Tensor], infer_state: InferStateInfo return graph_output_tensors + def close(self): + for execution in self.graph.values(): + execution.close() + self.graph.clear() + def _replay_overlap( self, input_tensors: List[torch.Tensor], diff --git a/lightllm/common/basemodel/triton_kernel/att/__init__.py b/lightllm/common/basemodel/triton_kernel/att/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py deleted file mode 100644 index bac4adb5ba..0000000000 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ /dev/null @@ -1,63 +0,0 @@ -import torch - - -def gqa_token_decode_attention_flash_decoding( - q: torch.Tensor, - infer_state, - cache_k: torch.Tensor, - cache_v: torch.Tensor, - max_len_in_batch: int, - out=None, - alloc_tensor_func=torch.empty, - sliding_window=(-1, -1), -): - batch_size = infer_state.batch_size - q_head_num, head_dim = q.shape[1], q.shape[2] - calcu_shape1 = (batch_size, q_head_num, head_dim) - - from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( - flash_decode_stage1, - ) - from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage2 import ( - flash_decode_stage2, - ) - - o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out - - # Because we need to allocate some intermediate tensors, considering parallelism and - # intermediate memory consumption, when batch_size is small, block_num is larger, - # and when batch_size is large, block_num is smaller. This achieves a better balance - # of memory consumption and performance. - BLOCK_SEQ = 256 - if batch_size <= 16: - block_num = 128 - elif batch_size <= 64: - block_num = 64 - else: - block_num = 32 - - mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") - mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") - - flash_decode_stage1( - q=q.view(calcu_shape1), - k=cache_k, - v=cache_v, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_Seqlen=infer_state.b_seq_len, - max_len_in_batch=max_len_in_batch, - mid_out=mid_o, - mid_out_logsumexp=mid_o_logexpsum, - block_seq=BLOCK_SEQ, - sliding_window=sliding_window, - ) - flash_decode_stage2( - mid_out=mid_o, - mid_out_logexpsum=mid_o_logexpsum, - B_Seqlen=infer_state.b_seq_len, - out=o_tensor.view(calcu_shape1), - block_seq=BLOCK_SEQ, - sliding_window=sliding_window, - ) - return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py deleted file mode 100644 index 9fb61c74f5..0000000000 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Temporary legacy import shim for the attention migration. - -New code imports lightllm.common.attention.backends directly. Removed with the -runtime integration in PR #1721 (5/5), after remaining model callers migrate. -""" - -from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.mtp_diverse_attn import ( - token_decode_attention_mtp_diverse_single_token, -) -from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage1_single_token import ( - mtp_diverse_stage1_single_token, -) -from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse.stage2_single_token import ( - mtp_diverse_stage2_single_token, -) - -__all__ = [ - "token_decode_attention_mtp_diverse_single_token", - "mtp_diverse_stage1_single_token", - "mtp_diverse_stage2_single_token", -] diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py deleted file mode 100644 index 2bf2d4b86b..0000000000 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py +++ /dev/null @@ -1,64 +0,0 @@ -from typing import Optional - -import torch - - -def token_decode_attention_flash_decoding( - q: torch.Tensor, - infer_state, - cache_k: torch.Tensor, - cache_k_scale: torch.Tensor, - cache_v: torch.Tensor, - cache_v_scale: torch.Tensor, - max_kv_seq_len: int, - out: Optional[torch.Tensor] = None, - alloc_tensor_func=torch.empty, -): - BLOCK_SEQ = 256 - batch_size = infer_state.batch_size - q_head_num = q.shape[1] - head_dim = q.shape[2] - calcu_shape1 = (batch_size, q_head_num, head_dim) - - o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out - - # 因为需要分配一些中间tensor，考虑到并行度和中间显存的消耗，batch_size 小的 - # 时候 block_num 较大， batch_size 大的时候 block_num 较小。这样可以达到较好 - # 的显存消耗和性能的平衡。 - if batch_size <= 16: - block_num = 128 - elif batch_size <= 64: - block_num = 64 - else: - block_num = 32 - - mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") - mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=q.dtype, device="cuda") - - from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv.int4kv_flash_decoding_stage1 import ( - int4kv_flash_decode_stage1, - ) - - int4kv_flash_decode_stage1( - q=q.view(calcu_shape1), - k=cache_k, - k_scale=cache_k_scale, - v=cache_v, - v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_Seqlen=infer_state.b_seq_len, - max_kv_seq_len=max_kv_seq_len, - mid_out=mid_o, - mid_out_logsumexp=mid_o_logexpsum, - block_seq=BLOCK_SEQ, - ) - - from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( - int8kv_flash_decoding_stage2, - ) - - int8kv_flash_decoding_stage2.flash_decode_stage2( - mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ - ) - return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py deleted file mode 100644 index 52ba55e30b..0000000000 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py +++ /dev/null @@ -1,105 +0,0 @@ -# 为 diverse mode 定制设计的 int8kv flash decoding attention 实现，可以实现更高效的多样性采样 -import torch - -from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( - flash_decode_stage1, -) -from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage2, -) -from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage3 import ( - flash_diverse_decode_stage3, -) -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size - - -def token_decode_attention_flash_decoding( - q, - infer_state: InferStateInfo, - cache_k, - cache_k_scale, - cache_v, - cache_v_scale, - out=None, - alloc_tensor_func=torch.empty, - shared_streams_dict={}, -): - b_shared_seq_len = infer_state.decode_att_state.b_shared_seq_len - b_mark_shared_group = infer_state.decode_att_state.b_mark_shared_group - - if "stream1" not in shared_streams_dict: - shared_streams_dict["stream1"] = torch.cuda.Stream() - if "stream2" not in shared_streams_dict: - shared_streams_dict["stream2"] = torch.cuda.Stream() - - stream1 = shared_streams_dict["stream1"] - stream2 = shared_streams_dict["stream2"] - - q_head_num = q.shape[1] - head_dim = q.shape[2] - - BLOCK_SEQ = 256 - batch_size = infer_state.batch_size - max_kv_seq_len = infer_state.max_kv_seq_len - calcu_shape1 = (batch_size, q_head_num, head_dim) - - o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out - - mid_o = alloc_tensor_func( - [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2, head_dim], dtype=torch.float32, device="cuda" - ) - mid_o_logexpsum = alloc_tensor_func( - [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 2], dtype=torch.float32, device="cuda" - ) - - current_stream = torch.cuda.current_stream() - - stream1.wait_stream(current_stream) - with torch.cuda.stream(stream1): - flash_decode_stage1( - q=q.view(calcu_shape1), - k=cache_k, - k_scale=cache_k_scale, - v=cache_v, - v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - b_shared_seq_len=b_shared_seq_len, - b_mark_shared_group=b_mark_shared_group, - max_len_in_batch=infer_state.max_kv_seq_len, - mid_out=mid_o, - mid_out_logsumexp=mid_o_logexpsum, - block_seq=BLOCK_SEQ, - max_batch_group_size=get_diverse_max_batch_shared_group_size(), - ) - stream2.wait_stream(current_stream) - with torch.cuda.stream(stream2): - flash_decode_stage2( - q=q.view(calcu_shape1), - k=cache_k, - k_scale=cache_k_scale, - v=cache_v, - v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_Seqlen=infer_state.b_seq_len, - b_shared_seq_len=b_shared_seq_len, - max_len_in_batch=infer_state.max_kv_seq_len, - mid_out=mid_o, - mid_out_logsumexp=mid_o_logexpsum, - block_seq=BLOCK_SEQ, - ) - - current_stream.wait_stream(stream1) - current_stream.wait_stream(stream2) - - flash_diverse_decode_stage3( - mid_out=mid_o, - mid_out_logexpsum=mid_o_logexpsum, - B_Seqlen=infer_state.b_seq_len, - b_shared_seq_len=b_shared_seq_len, - O=o_tensor.view(calcu_shape1), - block_seq=BLOCK_SEQ, - ) - return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/__init__.py deleted file mode 100644 index 8cc4aa919b..0000000000 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .int8kv_flash_decoding import token_decode_attention_flash_decoding diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py deleted file mode 100644 index e14f62c28d..0000000000 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py +++ /dev/null @@ -1,70 +0,0 @@ -from typing import Optional - -import torch - -from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage1 import ( - flash_decode_stage1, -) -from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( - flash_decode_stage2, -) -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -@torch.no_grad() -def token_decode_attention_flash_decoding( - q: torch.Tensor, - infer_state: InferStateInfo, - cache_k: torch.Tensor, - cache_k_scale: torch.Tensor, - cache_v: torch.Tensor, - cache_v_scale: torch.Tensor, - max_len_in_batch: int, - out: Optional[torch.Tensor] = None, - alloc_tensor_func=torch.empty, -): - q_head_num = q.shape[1] - head_dim = q.shape[2] - - BLOCK_SEQ = 256 - batch_size = infer_state.batch_size - calcu_shape1 = (batch_size, q_head_num, head_dim) - - o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out - - # 因为需要分配一些中间tensor，考虑到并行度和中间显存的消耗，batch_size 小的 - # 时候 block_num 较大， batch_size 大的时候 block_num 较小。这样可以达到较好 - # 的显存消耗和性能的平衡。 - if batch_size <= 16: - block_num = 128 - elif batch_size <= 64: - block_num = 64 - else: - block_num = 32 - - mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") - mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") - - flash_decode_stage1( - q=q.view(calcu_shape1), - k=cache_k, - k_scale=cache_k_scale, - v=cache_v, - v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_seq_len=infer_state.b_seq_len, - max_len_in_batch=max_len_in_batch, - mid_out=mid_o, - mid_out_logsumexp=mid_o_logexpsum, - block_seq=BLOCK_SEQ, - ) - - flash_decode_stage2( - mid_out=mid_o, - mid_out_logexpsum=mid_o_logexpsum, - B_Seqlen=infer_state.b_seq_len, - O=o_tensor.view(calcu_shape1), - block_seq=BLOCK_SEQ, - ) - return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py deleted file mode 100644 index b61f873348..0000000000 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py +++ /dev/null @@ -1,40 +0,0 @@ -import torch - - -def token_decode_attention_flash_decoding(q, infer_state, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty): - BLOCK_SEQ = 256 - batch_size = infer_state.batch_size - max_kv_seq_len = infer_state.max_kv_seq_len - q_head_num, head_dim = q.shape[1], q.shape[2] - calcu_shape1 = (batch_size, q_head_num, head_dim) - - from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage1 import ( - flash_decode_stage1, - ) - from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage2 import ( - flash_decode_stage2, - ) - - o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out - - mid_o = alloc_tensor_func( - [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1, head_dim], dtype=torch.float32, device="cuda" - ) - mid_o_logexpsum = alloc_tensor_func( - [batch_size, q_head_num, max_kv_seq_len // BLOCK_SEQ + 1], dtype=torch.float32, device="cuda" - ) - - flash_decode_stage1( - q.view(calcu_shape1), - cache_k, - cache_v, - infer_state.req_manager.req_to_token_indexs, - infer_state.b_req_idx, - infer_state.b_seq_len, - infer_state.max_kv_seq_len, - mid_o, - mid_o_logexpsum, - BLOCK_SEQ, - ) - flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) - return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/__init__.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/__init__.py b/lightllm/common/basemodel/triton_kernel/att/prefill_att/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py deleted file mode 100644 index 15d24413a0..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Temporary legacy import shim for the attention migration. - -New code imports lightllm.common.attention.backends directly. Removed with the -runtime integration in PR #1721 (5/5), after remaining model callers migrate. -""" - -from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn -from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import causal_conv1d_update -from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating -from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs -from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( - mtp_fused_recurrent_gated_delta_rule, -) - -from .fla.ops import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule - -__all__ = [ - "causal_conv1d_fn", - "causal_conv1d_update", - "fused_gdn_gating", - "conv_pack_gdn_decode_inputs", - "mtp_fused_recurrent_gated_delta_rule", - "chunk_gated_delta_rule", - "fused_recurrent_gated_delta_rule", -] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py deleted file mode 100644 index 2bde70bb99..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang -# -# This file contains code copied from the flash-linear-attention project. -# The original source code was licensed under the MIT license and included -# the following copyright notice: -# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang - -# Adapted from -# https://github.com/vllm-project/vllm diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py deleted file mode 100644 index 72b949c280..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Temporary legacy import shim for the attention migration. - -New code imports lightllm.common.attention.backends directly. Removed with the -runtime integration in PR #1721 (5/5), after remaining model callers migrate. -""" - -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang -# -# This file contains code copied from the flash-linear-attention project. -# The original source code was licensed under the MIT license and included -# the following copyright notice: -# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule -from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule - -__all__ = [ - "chunk_gated_delta_rule", - "fused_recurrent_gated_delta_rule", -] diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py deleted file mode 100644 index fb0609a401..0000000000 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .gqa_flash_decoding import gqa_token_decode_attention_flash_decoding diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py deleted file mode 100644 index 129e943f40..0000000000 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py +++ /dev/null @@ -1,164 +0,0 @@ -import os -from typing import List - -import torch -import torch.multiprocessing as mp -import triton -import triton.language as tl - -from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_config import ( - MlaDecodeAttentionKernelConfig, -) -from lightllm.utils.device_utils import get_device_sm_count -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def gqa_token_decode_attention_flash_decoding( - q_nope, q_rope, kv_nope, kv_rope, infer_state, softmax_scale, out=None, alloc_tensor_func=torch.empty, **run_config -): - batch_size = infer_state.batch_size - max_kv_seq_len = infer_state.max_kv_seq_len - - q_head_num, kv_lora_rank = q_nope.shape[1], q_nope.shape[2] - q_rope_dim = q_rope.shape[2] - assert q_rope_dim == 64 - - calcu_shape1 = (batch_size, q_head_num, kv_lora_rank) - calcu_shape2 = (batch_size, q_head_num, q_rope_dim) - - if not run_config: - if torch.cuda.is_current_stream_capturing(): - avg_seq_len_in_batch = max_kv_seq_len - else: - avg_seq_len_in_batch = infer_state.total_token_num // batch_size - - run_config = MlaDecodeAttentionKernelConfig.try_to_get_best_config( - batch_size=batch_size, - avg_seq_len_in_batch=avg_seq_len_in_batch, - q_head_num=q_head_num, - q_head_dim=kv_lora_rank, - q_rope_dim=q_rope_dim, - out_dtype=torch.bfloat16, - ) - - BLOCK_N = run_config["BLOCK_N"] - - from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_stage1 import ( - flash_decode_stage1, - ) - from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding_stage2 import ( - flash_decode_stage2, - ) - - o_tensor = alloc_tensor_func(q_nope.shape, q_nope.dtype, q_nope.device) if out is None else out - - fake_decode_att_block_seq = torch.empty([0], dtype=torch.int64, device="cuda") - mid_o = torch.empty([q_head_num, 0, kv_lora_rank], dtype=torch.float32, device="cuda") - mid_o_logexpsum = torch.empty([q_head_num, 0], dtype=torch.float32, device="cuda") - - vsm_count = flash_decode_stage1( - fake_decode_att_block_seq, - q_nope.view(calcu_shape1), - q_rope.view(calcu_shape2), - kv_nope, - kv_rope, - infer_state.req_manager.req_to_token_indexs, - infer_state.b_req_idx, - infer_state.b_seq_len, - mid_o, - mid_o_logexpsum, - softmax_scale, - get_sm_count=True, - **run_config, - ) - - if not hasattr(infer_state, "decode_att_block_seq"): - decode_att_block_seq = torch.empty( - [ - 1, - ], - dtype=torch.int64, - device="cuda", - ) - mid_o_batch_start_index = torch.empty( - [ - batch_size, - ], - dtype=torch.int64, - device="cuda", - ) - _fwd_kernel_calcu_index_and_block_seq[(1,)]( - infer_state.b_seq_len, - decode_att_block_seq, - mid_o_batch_start_index, - vsm_count, - batch_size, - BLOCK_N=BLOCK_N, - MAX_BATCH_SIZE=triton.next_power_of_2(batch_size), - num_warps=4, - ) - - infer_state.decode_att_block_seq = decode_att_block_seq - infer_state.mid_o_batch_start_index = mid_o_batch_start_index - - mid_o = torch.empty([q_head_num, vsm_count * 4 + batch_size, kv_lora_rank], dtype=torch.float32, device="cuda") - mid_o_logexpsum = torch.empty([q_head_num, vsm_count * 4 + batch_size], dtype=torch.float32, device="cuda") - - flash_decode_stage1( - infer_state.decode_att_block_seq, - q_nope.view(calcu_shape1), - q_rope.view(calcu_shape2), - kv_nope, - kv_rope, - infer_state.req_manager.req_to_token_indexs, - infer_state.b_req_idx, - infer_state.b_seq_len, - mid_o, - mid_o_logexpsum, - softmax_scale, - get_sm_count=False, - **run_config, - ) - - flash_decode_stage2( - infer_state.decode_att_block_seq, - infer_state.mid_o_batch_start_index, - mid_o, - mid_o_logexpsum, - infer_state.b_seq_len, - o_tensor.view(calcu_shape1), - **run_config, - ) - return o_tensor - - -@triton.jit -def _fwd_kernel_calcu_index_and_block_seq( - b_seq_len_ptr, - mid_o_decode_att_block_seq_ptr, - mid_o_batch_start_index_ptr, - num_sm, - batch_size, - BLOCK_N: tl.constexpr, - MAX_BATCH_SIZE: tl.constexpr, -): - b_seq_len = tl.load( - b_seq_len_ptr + tl.arange(0, MAX_BATCH_SIZE), mask=tl.arange(0, MAX_BATCH_SIZE) < batch_size, other=0 - ) - total_token_num = tl.sum(b_seq_len) - - block_seq = tl.cast(total_token_num / (num_sm * 4), dtype=tl.int32) + 1 - block_seq = tl.cdiv(block_seq, BLOCK_N) * BLOCK_N - - block_seq_len = tl.cdiv(b_seq_len, block_seq) - cumsum_seq_len = tl.cumsum(block_seq_len) - batch_start_index = cumsum_seq_len - block_seq_len - tl.store( - mid_o_batch_start_index_ptr + tl.arange(0, MAX_BATCH_SIZE), - batch_start_index, - mask=tl.arange(0, MAX_BATCH_SIZE) < batch_size, - ) - tl.store(mid_o_decode_att_block_seq_ptr, block_seq) - return diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py b/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py deleted file mode 100644 index 76599b6c3f..0000000000 --- a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Temporary legacy import shim for the attention migration. - -New code imports lightllm.common.attention.backends directly. Removed with the -runtime integration in PR #1721 (5/5), after remaining model callers migrate. -""" - -from lightllm.common.attention.backends.mla.triton.kernels.prefill.context_flashattention_nopad_with_v import ( - context_attention_fwd_with_v, -) diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index 98fab22505..b44a529203 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -5,7 +5,6 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger - logger = init_logger("lightllm.common.req_manager") @@ -53,6 +52,7 @@ def __init__(self, max_request_num, max_sequence_length): self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num + self.request_generations = torch.zeros(max_request_num + 1, dtype=torch.int64, device="cuda") def bind_hold_page(self, token_indices): """Map padding requests to the resource owner's reserved physical page.""" @@ -67,11 +67,15 @@ def bind_hold_page(self, token_indices): hold_row.view(-1, self.page_size).copy_(hold_page) def alloc(self): - return self.req_list.alloc() + index = self.req_list.alloc() + if index is not None: + self.request_generations[index].add_(1) + return index def free(self, free_req_indexes: List[int]): for req_index in free_req_indexes: self.req_list.free(req_index) + self.request_generations[req_index].add_(1) if self.req_list.is_all_free(): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") @@ -80,4 +84,5 @@ def free_req(self, free_req_index: int): def free_all(self): self.req_list = _ReqIndexPool(self.max_request_num) + self.request_generations.add_(1) return diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py index 5ec00d542a..5da2649792 100755 --- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py @@ -1,13 +1,15 @@ -from lightllm.common.attention_state.full_att import FullAttentionState -import torch from typing import Tuple -from lightllm.common.basemodel import TransformerLayerInferTpl + +import torch + +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.basemodel import InferStateInfo, TransformerLayerInferTpl from lightllm.models.bloom.layer_weights.transformer_layer_weight import BloomTransformerLayerWeight -from lightllm.common.basemodel import InferStateInfo -from lightllm.common.basemodel.attention.base_att import AttControl class BloomTransformerLayerInfer(TransformerLayerInferTpl): + attention_uses_alibi = True + """ """ def __init__(self, layer_num, network_config): @@ -29,42 +31,16 @@ def _context_attention_kernel( layer_weight: BloomTransformerLayerWeight, out=None, ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(use_alibi=True, tp_alibi=layer_weight.tp_alibi), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor def _token_attention_kernel( self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(use_alibi=True, tp_alibi=layer_weight.tp_alibi), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.decode_attention.forward(infer_state, DenseInputs(_q)) return o_tensor.view(q.shape) def _att_norm( diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 8e963d7f84..089535e420 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -1,17 +1,18 @@ from typing import TYPE_CHECKING -from lightllm.models.bloom.layer_infer.pre_layer_infer import BloomPreLayerInfer + +from lightllm.common.attention.backends.dense.triton.backend import TritonAttBackend +from lightllm.common.basemodel import InferStateInfo, TpPartBaseModel from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer +from lightllm.models.bloom.layer_infer.pre_layer_infer import BloomPreLayerInfer from lightllm.models.bloom.layer_infer.transformer_layer_infer import BloomTransformerLayerInfer from lightllm.models.bloom.layer_weights.pre_and_post_layer_weight import BloomPreAndPostLayerWeight from lightllm.models.bloom.layer_weights.transformer_layer_weight import BloomTransformerLayerWeight -from lightllm.common.basemodel import InferStateInfo, TpPartBaseModel -from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend - if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class BloomTpPartModel(TpPartBaseModel): @@ -75,7 +76,5 @@ def _reset_num_key_value_heads(self): self.config["num_key_value_heads"] = self.config["num_attention_heads"] return - def _init_att_backend(self): - self.prefill_att_backend = TritonAttBackend(self) - self.decode_att_backend = TritonAttBackend(self) - return + def attention_backend_candidates(self, family, phase): + return {self.args.llm_kv_type: {"triton": TritonAttBackend}}, ["triton"] diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index a9d23a6de5..e387529d6f 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -1,26 +1,28 @@ -from lightllm.common.attention_state.full_att import FullAttentionState import os +from functools import partial + import torch import triton + +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe +from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.models.deepseek2.triton_kernel.sample_kv import sample_kv from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( - use_sm100_mega_moe, -) -from functools import partial from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale -from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_global_world_size +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class Deepseek2TransformerLayerInfer(LlamaTransformerLayerInfer): + attention_family = "mla" + def __init__(self, layer_num, network_config): self.tp_k_head_num_ = 1 self.tp_v_head_num_ = 1 @@ -86,13 +88,7 @@ def _context_attention_kernel( layer_weight=layer_weight, ) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=q, - k=(k_nope, k_rope), - v=v, - att_control=AttControl(mla_prefill=True, mla_prefill_dict={"softmax_scale": self.softmax_scale}), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state, MLAPrefillInputs(q, k_nope, k_rope, v)) return o_tensor def _token_attention_kernel( @@ -104,21 +100,8 @@ def _token_attention_kernel( ): q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - kv = state.get_layer_view(logical_layer_id=self.layer_num_) - - out = infer_state.decode_att_state.decode_att( - q=(q_nope, q_rope), - k=kv, - v=None, - att_control=AttControl(mla_decode=True, mla_decode_dict={"softmax_scale": self.softmax_scale}), - alloc_func=self.alloc_tensor, - ) + + out = self.decode_attention.forward(infer_state, MLADecodeInputs(q_nope, q_rope)) return out def _decompress_kv( @@ -418,7 +401,6 @@ def overlap_tpsp_token_forward( def _1_hook_post(): _1_hook() - nonlocal _1_ffn_out if self.n_shared_experts is not None: _1_ffn_out.add_(_1_shared_output) input_embdings1.add_(_1_ffn_out.view(-1, self.embed_dim_)) @@ -581,7 +563,6 @@ def overlap_tpsp_context_forward( def _1_hook_post(): _1_hook() - nonlocal _1_ffn_out if self.n_shared_experts is not None: _1_ffn_out.add_(_1_shared_output) input_embdings1.add_(_1_ffn_out.view(-1, self.embed_dim_)) diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index e06c18ea16..c50cbfdfa5 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -1,22 +1,23 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, final + import torch -from typing import final + +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight -from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args -from lightllm.distributed.communication_op import dist_group_manager -from lightllm.common.basemodel.attention import get_mla_decode_att_backend_class, get_mla_prefill_att_backend_class +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Deepseek2TpPartModel(LlamaTpPartModel): @@ -33,11 +34,6 @@ def __init__(self, kvargs): super().__init__(kvargs) return - def _init_att_backend(self): - self.prefill_att_backend = get_mla_prefill_att_backend_class(index=0)(model=self) - self.decode_att_backend = get_mla_decode_att_backend_class(index=0)(model=self) - return - def _init_some_value(self): super()._init_some_value() self.tp_k_head_num_ = 1 @@ -95,7 +91,7 @@ def create_full_att_states( return [state] def _init_to_get_yarn_rotary(self): - from lightllm.models.llama.yarn_rotary_utils import find_correction_range, linear_ramp_mask, get_deepseek_mscale + from lightllm.models.llama.yarn_rotary_utils import find_correction_range, get_deepseek_mscale, linear_ramp_mask dim = self.qk_rope_head_dim max_position_embeddings = self.config.get("max_position_embeddings", 2048) diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index e026d1ea8e..06a0ba35f8 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,20 +1,24 @@ -from lightllm.common.attention_state.full_att import FullAttentionState -import torch from typing import Any + +import torch + +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.distributed import all_gather_into_tensor from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer -from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import destindex_copy_indexer_ks from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import extract_indexer_ks from lightllm.utils.envs_utils import get_env_start_args -from lightllm.distributed import all_gather_into_tensor class Deepseek3_2TransformerLayerInfer(Deepseek2TransformerLayerInfer): + attention_family = "sparse_mla" + def __init__(self, layer_num, network_config): self.index_topk = network_config["index_topk"] super().__init__(layer_num, network_config) @@ -78,7 +82,7 @@ def _context_attention_kernel( q_all = torch.cat([q_nope, q_rope], dim=-1) # 计算 topk indices - att_state = infer_state.prefill_att_state + att_state = self.prefill_attention.context(infer_state) topk_mem_indices, topk_indices = self.indexer._get_indices( hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], @@ -89,28 +93,10 @@ def _context_attention_kernel( del infer_state.get_topk_indices_params # Use NSA backend for attention computation - att_control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_mem_indices": topk_mem_indices, - "topk_indices": topk_indices, - "prefill_cache_kv": kv, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - }, - ) - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - mla_out = infer_state.prefill_att_state.prefill_att( - q=q_all, - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=att_control, + mla_out = self.prefill_attention.forward( + infer_state, + SparseMLAPrefillInputs(q_all, topk_mem_indices, topk_indices, kv), ) return mla_out @@ -126,7 +112,7 @@ def _token_attention_kernel( q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) # 计算 topk mem indices - att_state = infer_state.decode_att_state + att_state = self.decode_attention.context(infer_state) topk_mem_indices, _ = self.indexer._get_indices( hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], @@ -137,28 +123,10 @@ def _token_attention_kernel( del infer_state.get_topk_indices_params # Use NSA backend for attention computation - att_control = AttControl( - nsa_decode=True, - nsa_decode_dict={ - "layer_index": self.layer_num_, - "topk_mem_indices": topk_mem_indices, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - "qk_rope_head_dim": self.qk_rope_head_dim, - }, - ) - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - o_tensor = infer_state.decode_att_state.decode_att( - q=(q_nope, q_rope), - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=att_control, + o_tensor = self.decode_attention.forward( + infer_state, + SparseMLADecodeInputs(q_nope, q_rope, topk_mem_indices), ) return o_tensor @@ -177,7 +145,7 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.scale_fmt = network_config["quantization_config"]["scale_fmt"] self.softmax_scale = (self.index_head_dim) ** (-0.5) self.index_n_heads = network_config["index_n_heads"] - self.index_n_heads_scale = (self.index_n_heads ** -0.5) * self.softmax_scale + self.index_n_heads_scale = (self.index_n_heads**-0.5) * self.softmax_scale self.tp_world_size_ = tp_world_size self.tp_index_n_heads = self.index_n_heads // self.tp_world_size_ @@ -189,7 +157,6 @@ def _get_indices( att_state: Any, layer_weight: Deepseek3_2TransformerLayerWeight, ): - q, k = self._get_q_k_bf16(hidden_states, q_lora, infer_state, layer_weight) if self.tp_world_size_ > 1: @@ -281,7 +248,7 @@ def _rotate_activation(x: torch.Tensor) -> torch.Tensor: hidden_size = x.size(-1) assert (hidden_size & (hidden_size - 1)) == 0, "Hidden size must be a power of 2 for Hadamard transform." - return hadamard_transform(x, scale=hidden_size ** -0.5) + return hadamard_transform(x, scale=hidden_size**-0.5) def _get_q_k_bf16( self, diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index b691d986cc..4392624db7 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -1,19 +1,19 @@ from typing import TYPE_CHECKING + import torch + from lightllm.models.deepseek2.model import Deepseek2TpPartModel -from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import Deepseek3_2TransformerLayerInfer -from lightllm.common.basemodel.attention import get_nsa_prefill_att_backend_class, get_nsa_decode_att_backend_class - +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Deepseek3_2TpPartModel(Deepseek2TpPartModel): - # weight class transformer_weight_class = Deepseek3_2TransformerLayerWeight @@ -51,8 +51,3 @@ def create_full_att_states( device="meta", ) return [state] - - def _init_att_backend(self): - self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0)(model=self) - self.decode_att_backend = get_nsa_decode_att_backend_class(index=0)(model=self) - return diff --git a/lightllm/models/deepseek_v4/infer_struct.py b/lightllm/models/deepseek_v4/infer_struct.py index 6f90fa1207..d3e5f60854 100644 --- a/lightllm/models/deepseek_v4/infer_struct.py +++ b/lightllm/models/deepseek_v4/infer_struct.py @@ -2,12 +2,13 @@ from typing import Literal import torch -from lightllm.common.basemodel import InferStateInfo + from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( - CompressedAttentionState, CompressedAttentionLayerView, + CompressedAttentionState, ) +from lightllm.common.basemodel import InferStateInfo C4_PREFILL_LOGITS_BUDGET_BYTES = 512 * 1024 * 1024 @@ -105,7 +106,10 @@ def init_some_extra_state(self, model): self.dsv4_c128_lengths = None super().init_some_extra_state(model) # sets position_ids, b_q_seq_len, b_q_start_loc (prefill) if self.is_prefill: - self.prefill_att_state.prepare_swa_workspace(self) + from .attention.sparse_mla import prepare_swa_workspace + + buffer, capacity = model.swa_prefill_workspace + self.swa_workspace = prepare_swa_workspace(self.swa_state, buffer, capacity, self, self.microbatch_index) pos = self.position_ids self.position_cos_sliding = torch.index_select(model._cos_cached_sliding, 0, pos) # [T, rope_dim//2] self.position_sin_sliding = torch.index_select(model._sin_cached_sliding, 0, pos) @@ -196,7 +200,7 @@ def _init_swa_metadata(self, model): sliding_window=self.swa_state.config.window_size, image_left=self.dsv4_image_left, image_right=self.dsv4_image_right, - prefill_state=self.prefill_att_state if self.is_prefill else None, + prefill_state=self.swa_workspace if self.is_prefill else None, swa_runtime_write_slots=self.dsv4_swa_runtime_write_slots, token_to_batch_idx=self._dsv4_token_to_batch_idx, b_seq_len=self.b_seq_len, @@ -216,6 +220,7 @@ def get_c4_indexer_metadata(self, c4_cap): if self.c4_indexer_metadata is not None: return self.c4_indexer_metadata import deep_gemm + from .triton_kernel.build_c4_indexer_page_table import build_c4_indexer_page_table valid_len = ((self.position_ids + 1) // 4).to(torch.int32) diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index 21825e68e7..83b8052b5a 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -1,26 +1,49 @@ import torch import triton + from lightllm.common.basemodel import TransformerLayerInferTpl -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.triton_kernel.mhc import hc_post, hc_post_pre_norm, hc_pre_norm from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( - silu_and_mul_fwd, - silu_and_mul_quant_fp8, -) +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd, silu_and_mul_quant_fp8 +from lightllm.common.basemodel.triton_kernel.mhc import hc_post, hc_post_pre_norm, hc_pre_norm from lightllm.common.quantization.deepgemm import DeepGEMMFP8w8a8B128QuantizationMethod +from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4Inputs, DeepseekV4Params from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight -from lightllm.utils.envs_utils import get_env_start_args +from lightllm.models.deepseek_v4.triton_kernel.topk_softplus_sqrt import topk_softplus_sqrt from lightllm.utils.dist_utils import get_global_world_size +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor + +from ..infer_struct import DeepseekV4InferStateInfo +from ..triton_kernel.rotary_emb import inverse_rotary_emb from .compressor import CompressorInfer from .indexer import DeepseekV4IndexInfer -from ..triton_kernel.rotary_emb import inverse_rotary_emb -from ..infer_struct import DeepseekV4InferStateInfo -from lightllm.models.deepseek_v4.triton_kernel.topk_softplus_sqrt import topk_softplus_sqrt class DeepseekV4TransformerLayerInfer(TransformerLayerInferTpl): + attention_family = "sparse_mla" + + attention_uses_sink = True + + def create_attention_params(self, model, weight, family): + manager, role, layer_id = model.att_state_manager, model.att_state_role, self.layer_num_ + state = manager.get_state(role, "swa_", layer_id) + config = state.config + row_bytes = config.get_packed_kv_bytes(config.page_token_size) + swa = state.get_layer_view(layer_id).view(-1, row_bytes) + prefill_swa = model.swa_prefill_workspace[0][config.logical_layer_ids.index(layer_id)].view(-1, row_bytes) + compressed = ( + manager.get_state(role, "compressed", layer_id).get_layer_view(layer_id).kv if self.compress_ratio else None + ) + return DeepseekV4Params( + swa, + prefill_swa, + compressed, + self.compress_ratio, + self.v_head_dim, + self.softmax_scale, + weight.attn_sink_.weight, + ) + @classmethod def get_prefill_aux_stream(cls: type["DeepseekV4TransformerLayerInfer"]) -> torch.cuda.Stream: """Return the single C4 prefill auxiliary stream shared by every layer in this process.""" @@ -353,8 +376,8 @@ def _norm_rope_and_cache_kv(self, kv, infer_state, layer_weight): from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa state = infer_state.swa_state - att_state = infer_state.prefill_att_state if infer_state.is_prefill else infer_state.decode_att_state - swa, _ = att_state.backend.get_att_input_params(infer_state, self.layer_num_) + params = self.prefill_attention.params if infer_state.is_prefill else self.decode_attention.params + swa = params.prefill_swa if infer_state.is_prefill else params.swa # 一次完成 RMSNorm、RoPE、FP8 打包和 SWA 写入。 fused_kv_norm_rope_to_swa( kv=kv, @@ -489,34 +512,12 @@ def _compress_and_index(self, q_lora, x, infer_state: DeepseekV4InferStateInfo, self.index_infer.write_indexer_k(x, infer_state, layer_weight, cos_table, sin_table) return self.index_infer.build_metadata(x, q_lora, infer_state, layer_weight) - def _context_attention_kernel( - self, - q, - q_lora, - x, - infer_state: DeepseekV4InferStateInfo, - layer_weight: DeepseekV4TransformerLayerWeight, - out=None, - ): + def _context_attention_kernel(self, q, q_lora, x, infer_state, layer_weight, out=None): meta = self._compress_and_index(q_lora, x, infer_state, layer_weight) - att_control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "compress_ratio": self.compress_ratio, - "head_dim_v": self.v_head_dim, - "softmax_scale": self.softmax_scale, - "attn_sink": layer_weight.attn_sink_.weight, - **meta, - }, - ) - return infer_state.prefill_att_state.prefill_att( - q=q, - k=infer_state.prefill_att_state.backend.get_att_input_params(infer_state, self.layer_num_), - v=None, - att_control=att_control, - alloc_func=self.alloc_tensor, - out=out, + inputs = DeepseekV4Inputs( + q, meta["swa_indices"], meta["swa_lengths"], meta.get("extra_indices"), meta.get("extra_lengths") ) + return self.prefill_attention.forward(infer_state, inputs, out=out) # ------------------------------------------------------------------ attention (decode) def token_attention_forward( @@ -526,26 +527,12 @@ def token_attention_forward( o = self._token_attention_kernel(q, q_lora, full_x, infer_state, layer_weight) return self._get_o(o, infer_state, layer_weight) - def _token_attention_kernel( - self, q, q_lora, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight - ): + def _token_attention_kernel(self, q, q_lora, x, infer_state, layer_weight, out=None): meta = self._compress_and_index(q_lora, x, infer_state, layer_weight) - att_control = AttControl( - nsa_decode=True, - nsa_decode_dict={ - "compress_ratio": self.compress_ratio, - "head_dim_v": self.v_head_dim, - "softmax_scale": self.softmax_scale, - "attn_sink": layer_weight.attn_sink_.weight, - **meta, - }, - ) - return infer_state.decode_att_state.decode_att( - q=q, - k=infer_state.decode_att_state.backend.get_att_input_params(infer_state, self.layer_num_), - v=None, - att_control=att_control, + inputs = DeepseekV4Inputs( + q, meta["swa_indices"], meta["swa_lengths"], meta.get("extra_indices"), meta.get("extra_lengths") ) + return self.decode_attention.forward(infer_state, inputs, out=out) # ------------------------------------------------------------------ moe def _ffn_shared_expert( diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index ff46b5d6da..56bb590521 100644 --- a/lightllm/models/deepseek_v4/model.py +++ b/lightllm/models/deepseek_v4/model.py @@ -3,44 +3,29 @@ from typing import TYPE_CHECKING import torch -from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.models.deepseek_v4.layer_weights.pre_and_post_layer_weight import ( - DeepseekV4PreAndPostLayerWeight, -) -from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import ( - DeepseekV4TransformerLayerWeight, -) -from lightllm.models.deepseek_v4.layer_infer.pre_layer_infer import ( - DeepseekV4PreLayerInfer, -) -from lightllm.models.deepseek_v4.layer_infer.post_layer_infer import ( - DeepseekV4PostLayerInfer, -) -from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import ( - DeepseekV4TransformerLayerInfer, -) -from lightllm.common.basemodel.attention import get_nsa_prefill_att_backend_class, get_nsa_decode_att_backend_class -from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import ( - DSV4_NSA_BACKENDS, - DeepseekV4FlashMlaFp8SparseAttBackend, -) + from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector from lightllm.common.basemodel.triton_kernel.mhc import hc_contract, hc_post +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.deepseek_v4.attention.sparse_mla import DSV4_NSA_BACKENDS, DeepseekV4FlashMlaFp8SparseAttBackend from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo +from lightllm.models.deepseek_v4.layer_infer.post_layer_infer import DeepseekV4PostLayerInfer +from lightllm.models.deepseek_v4.layer_infer.pre_layer_infer import DeepseekV4PreLayerInfer +from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer +from lightllm.models.deepseek_v4.layer_weights.pre_and_post_layer_weight import DeepseekV4PreAndPostLayerWeight +from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace -from lightllm.models.llama.yarn_rotary_utils import ( - find_correction_range, - linear_ramp_mask, -) +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.llama.yarn_rotary_utils import find_correction_range, linear_ramp_mask from lightllm.utils.envs_utils import get_env_start_args, get_redundancy_expert_num from lightllm.utils.log_utils import init_logger -from lightllm.distributed.communication_op import dist_group_manager if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att.base_state import FullAttentionState + from lightllm.common.attention_state.bounded_att.base_state import BoundedAttentionState + from lightllm.common.attention_state.full_att.base_state import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs logger = init_logger(__name__) @@ -107,6 +92,7 @@ def create_full_att_states( from lightllm.common.attention_state.full_att.compressed_att.create_compressed_att_state import ( create_compressed_att_state, ) + from .config import create_attention_config assert args.llm_kv_type in ("None", "fp8kv_dsa"), f"unsupported V4 KV encoding: {args.llm_kv_type}" @@ -143,9 +129,10 @@ def create_bounded_att_states( """创建本模型的 SWA/compressor 请求运行态，大小页容量由管理器后续确定。""" from lightllm.common.attention_state.bounded_att import ( WindowCompressorStateConfig, - create_window_compressor_state, create_sliding_window_state, + create_window_compressor_state, ) + from .config import create_attention_config config = create_attention_config( @@ -209,18 +196,21 @@ def _init_quant(self): config = dict(self.config, n_layer=len(self.config["compress_ratios"])) self.quant_cfg = Quantcfg(config, self.quant_type, self.quant_cfg_path, self.expert_dtype) - def _init_att_backend(self): + def attention_backend_candidates(self, family, phase): + return DSV4_NSA_BACKENDS, ["flashmla_sparse"] + + def prepare_attention_backends(self): args = get_env_start_args() if args.llm_kv_type == "None": args.llm_kv_type = "fp8kv_dsa" # TODO: 支持其他 kv type if args.llm_kv_type != "fp8kv_dsa": raise RuntimeError("DeepSeek-V4 requires llm_kv_type=fp8kv_dsa for packed FlashMLA sparse attention") - self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0, backend_map=DSV4_NSA_BACKENDS)(model=self) - self.decode_att_backend = get_nsa_decode_att_backend_class(index=0, backend_map=DSV4_NSA_BACKENDS)(model=self) - - tp_q_head_num = self.prefill_att_backend.tp_q_head_num - padded_q_head_num = self.prefill_att_backend.padded_q_head_num + self.args.llm_kv_type = args.llm_kv_type + super().prepare_attention_backends() + backend = self.attention_groups[0].prefill_backend + tp_q_head_num = backend.tp_q_head_num + padded_q_head_num = backend.padded_q_head_num if not self.is_mtp_draft_model: # Drafts share this buffer after the main model has captured its CUDA graphs. self.dsv4_workspace.init_flashmla_prefill_q( @@ -299,7 +289,7 @@ def _init_to_get_rotary(self): orig_max = rs.get("original_max_position_embeddings", 0) def build_inv_freq(base): - freqs = 1.0 / (base ** freq_exponents) + freqs = 1.0 / (base**freq_exponents) if rope_type == "yarn" and orig_max > 0: beta_fast = rs.get("beta_fast", 32) beta_slow = rs.get("beta_slow", 1) diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index d1ebd30e30..448b088c9d 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -1,21 +1,30 @@ -from lightllm.common.attention_state.full_att import FullAttentionState import math + import torch import torch.nn as nn -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight -from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import ( - context_attention_fwd_gemma4_mm, -) -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): + def attention_window(self, phase): + return (self.sliding_window_ - 1, 0) if self.is_sliding and self.sliding_window_ > 0 else None + + def attention_kv_view(self, state): + k, v = state.get_layer_view(self.layer_num_) + width = self.tp_k_head_num_ * self.head_dim_ + k = k.flatten(1)[:, :width].view(-1, self.tp_k_head_num_, self.head_dim_) + v = v.flatten(1)[:, :width].view(-1, self.tp_k_head_num_, self.head_dim_) + return k, v + """ Gemma-4 decoder block. Per-layer heterogeneity (sliding vs full attention) is handled by switching shape / RoPE table / sliding-window flag at init @@ -32,7 +41,7 @@ def __init__(self, layer_num, network_config): self.is_moe = bool(network_config.get("enable_moe_block", False)) self.num_experts_per_tok = network_config.get("num_experts_per_tok", network_config.get("top_k_experts", 0)) self.norm_topk_prob = network_config.get("norm_topk_prob", True) - self.router_root_scale = self.embed_dim_ ** -0.5 + self.router_root_scale = self.embed_dim_**-0.5 layer_type = network_config["layer_types"][layer_num] self.is_sliding = layer_type == "sliding_attention" @@ -176,12 +185,6 @@ def _post_cache_kv(self, cache_kv, infer_state, layer_weight): # ----- Attention kernels (sliding window + per-layer KV reshape) --- - def _att_control(self): - if self.is_sliding and self.sliding_window_ > 0: - w = self.sliding_window_ - 1 - return AttControl(use_sliding_window=True, sliding_window=(w, 0), logical_layer_id=self.layer_num_) - return AttControl(use_sliding_window=False, sliding_window=(-1, -1), logical_layer_id=self.layer_num_) - def _get_layer_kv(self, infer_state: InferStateInfo): # 直接用当前逻辑层号访问；共享层由 State 的 logical_layer_groups 映射到生产层的物理 KV。 state = infer_state.att_state_manager.get_state( @@ -241,9 +244,7 @@ def _context_attention_kernel( # Full-attn layers: head_dim=512, no SWA, no image bidi — standard # triton via backend1. - o_tensor = infer_state.prefill_att_state1.prefill_att( - q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor - ) + o_tensor = self.prefill_attention.forward(infer_state, DenseInputs(_q)) return o_tensor.view(q.shape) def _token_attention_kernel( @@ -253,10 +254,9 @@ def _token_attention_kernel( layer_weight: Gemma4TransformerLayerWeight, out=None, ) -> torch.Tensor: - _k, _v = self._get_layer_kv(infer_state) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - att_state = infer_state.decode_att_state if self.is_sliding else infer_state.decode_att_state1 - o_tensor = att_state.decode_att(q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor) + + o_tensor = self.decode_attention.forward(infer_state, DenseInputs(_q)) return o_tensor.view(q.shape) # ----- FFN (Gemma gelu-tanh, fused gate_up + down) ----------------- diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 4f998ea25e..9f54db1198 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -1,27 +1,30 @@ -from typing import TYPE_CHECKING -import os import json +import os +from typing import TYPE_CHECKING + import torch -from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend + +from lightllm.common.attention.backends.dense.triton.backend import TritonAttBackend from lightllm.common.build_utils import repair_config -from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.distributed.communication_op import dist_group_manager from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo -from lightllm.models.gemma4.layer_infer.pre_layer_infer import Gemma4PreLayerInfer from lightllm.models.gemma4.layer_infer.post_layer_infer import Gemma4PostLayerInfer +from lightllm.models.gemma4.layer_infer.pre_layer_infer import Gemma4PreLayerInfer from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer from lightllm.models.gemma4.layer_weights.pre_and_post_layer_weight import Gemma4PreAndPostLayerWeight from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight +from lightllm.models.llama.model import LlamaTpPartModel from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from lightllm.distributed.communication_op import dist_group_manager logger = init_logger(__name__) if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Gemma4TpPartModel(LlamaTpPartModel): @@ -138,24 +141,8 @@ def create_full_att_states( ) return [state] - def _init_att_backend(self): - # Gemma-4 has per-layer heterogeneous attention: sliding layers use - # (head_dim=256, kv_heads=16); full-attn layers use (head_dim=512, - # kv_heads=4, k_eq_v). FA3 caps head_dim at 256 and flashinfer plans - # once per infer_state on a single shape — both unworkable for the - # heterogeneous layout. Both layer kinds go through triton. - # - # Primary backend = sliding layers. Sliding prefill bypasses the - # backend and calls gemma4_mm directly (SWA + image bidi in one - # pass); the prefill_att_state created here is unused but the - # framework requires prefill_att_backend to be non-None. - self.prefill_att_backend = TritonAttBackend(model=self) - self.decode_att_backend = TritonAttBackend(model=self) - - def _init_att_backend1(self): - # Secondary backend = full-attn layers (head_dim=512, plain causal). - self.prefill_att_backend1 = TritonAttBackend(model=self) - self.decode_att_backend1 = TritonAttBackend(model=self) + def attention_backend_candidates(self, family, phase): + return {self.args.llm_kv_type: {"triton": TritonAttBackend}}, ["triton"] def _init_custom(self): self._init_to_get_rotary_gemma4() diff --git a/lightllm/models/glm5_next/attention/fp8_sparse_mla.py b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py index 0a458f84a7..94bac0979c 100644 --- a/lightllm/models/glm5_next/attention/fp8_sparse_mla.py +++ b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py @@ -65,6 +65,7 @@ def get_smallq_workspace(self, context, q): "direct": allocate_direct_smallq_workspace(q[:1]), "selected": allocate_selected_prefill_workspace(q, batch_capacity=8), }, + expected_type=dict, ) if q.shape[0] == 1: return context.smallq_workspace["direct"] diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 31d79697f6..1da3bcd63f 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -1,14 +1,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING import dataclasses - -from lightllm.utils.torch_dtype_utils import get_torch_dtype +from typing import TYPE_CHECKING from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.utils.torch_dtype_utils import get_torch_dtype if TYPE_CHECKING: from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index be7a2fa781..f404c84389 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -2,21 +2,14 @@ from __future__ import annotations -from lightllm.common.attention_state.full_att import FullAttentionState - import torch +from lightllm.common.attention.api import AttentionFamily +from lightllm.common.attention.api.inputs import KDAInputs, SparseMLADecodeInputs, SparseMLAPrefillInputs from lightllm.common.basemodel import TransformerLayerInferTpl -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.common.basemodel.triton_kernel.mhc import hc_contract, hc_post, hc_pre_norm from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( - silu_and_mul_fwd, -) -from lightllm.common.basemodel.triton_kernel.mhc import ( - hc_contract, - hc_post, - hc_pre_norm, -) from lightllm.common.triton_utils.autotuner import Autotuner from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer from lightllm.utils.envs_utils import get_env_start_args @@ -24,6 +17,9 @@ class Glm5NextTransformerLayerInfer(TransformerLayerInferTpl): + def get_attention_family(self): + return AttentionFamily.KDA if self.is_linear_attention_layer else AttentionFamily.SPARSE_MLA + def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.eps_ = network_config["rms_norm_eps"] @@ -60,7 +56,7 @@ def __init__(self, layer_num, network_config): self.q_lora_rank = network_config["q_lora_rank"] self.kv_lora_rank = network_config["kv_lora_rank"] self.v_head_dim = network_config["v_head_dim"] - self.softmax_scale = self.qk_nope_head_dim ** -0.5 + self.softmax_scale = self.qk_nope_head_dim**-0.5 self.indexer = Glm5NextNsaInfer( layer_idx=self.layer_num_, network_config=self.network_config_, @@ -86,7 +82,6 @@ def _ffn(self, input, infer_state, layer_weight): def _ffn_tp(self, input, infer_state, layer_weight): """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" - input = input.view(-1, self.embed_dim_) up_gate_out = layer_weight.gate_up_proj.mm(input) ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) @@ -157,30 +152,14 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], infer_state=infer_state, - att_state=infer_state.prefill_att_state, + att_state=self.prefill_attention.context(infer_state), layer_weight=layer_weight, ) del infer_state.get_topk_indices_params - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - return infer_state.prefill_att_state.prefill_att( - q=q, - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_mem_indices": topk_mem_indices, - "topk_indices": topk_indices, - "prefill_cache_kv": kv, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - }, - ), + + return self.prefill_attention.forward( + infer_state, + SparseMLAPrefillInputs(q, topk_mem_indices, topk_indices, kv), ) def _context_attention_wrapper_run(self, q, cache_kv, infer_state, layer_weight): @@ -243,31 +222,15 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], infer_state=infer_state, - att_state=infer_state.decode_att_state, + att_state=self.decode_attention.context(infer_state), layer_weight=layer_weight, ) del infer_state.get_topk_indices_params q_rope = q_nope[..., :0] - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - return infer_state.decode_att_state.decode_att( - q=(q_nope, q_rope), - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=AttControl( - nsa_decode=True, - nsa_decode_dict={ - "layer_index": self.layer_num_, - "topk_mem_indices": topk_mem_indices, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - "qk_rope_head_dim": 0, - }, - ), + + return self.decode_attention.forward( + infer_state, + SparseMLADecodeInputs(q_nope, q_rope, topk_mem_indices), ) def _get_o(self, input, infer_state, layer_weight): @@ -326,7 +289,6 @@ def _kda_prefill_cuda_graph_wrapper( output at stable addresses, then register it as a CPU-side replay step between the surrounding CUDA-graph segments. """ - backend = infer_state.prefill_att_state1.backend mixed_qkv = mixed_qkv.contiguous() raw_decay_gate = raw_decay_gate.contiguous() beta_logits = beta_logits.contiguous() @@ -341,28 +303,16 @@ def _kda_prefill_cuda_graph_wrapper( # Construct this shape directly instead of dry-running the kernel: # kernel setup may synchronize with the host, which is illegal while a # CUDA graph capture is active. - output_shape = (1, mixed_qkv.shape[0], backend.tp_num_heads, backend.head_dim) + output_shape = (1, mixed_qkv.shape[0], self.tp_linear_num_heads, self.linear_head_dim) infer_state.prefill_cuda_graph_create_graph_obj() infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() output = torch.empty(output_shape, dtype=mixed_qkv.dtype, device=mixed_qkv.device) _output = tensor_to_no_ref_tensor(output) def kda_prefill_func(new_infer_state): - tmp_output = new_infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkv": _mixed_qkv, - "raw_gate": _raw_decay_gate, - "raw_beta": _beta_logits, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, + tmp_output = self.prefill_attention.forward( + new_infer_state, + KDAInputs(_mixed_qkv, _raw_decay_gate, _beta_logits), ) assert tmp_output.shape == _output.shape _output.copy_(tmp_output) @@ -388,21 +338,9 @@ def context_attention_forward(self, input_embeddings, infer_state, layer_weight) layer_weight, ) else: - core_output = infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkv": qkv, - "raw_gate": raw_decay_gate, - "raw_beta": beta_logits, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, + core_output = self.prefill_attention.forward( + infer_state, + KDAInputs(qkv, raw_decay_gate, beta_logits), ) return self._kda_post(core_output, raw_output_gate, infer_state, layer_weight) @@ -412,22 +350,7 @@ def token_attention_forward(self, input_embeddings, infer_state, layer_weight): qkv, raw_decay_gate, beta_logits, raw_output_gate = self._kda_projections( input_embeddings, infer_state, layer_weight ) - core_output = infer_state.decode_att_state1.decode_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_decode=True, - linear_att_decode_dict={ - "mixed_qkv": qkv, - "raw_gate": raw_decay_gate, - "raw_beta": beta_logits, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + core_output = self.decode_attention.forward(infer_state, KDAInputs(qkv, raw_decay_gate, beta_logits)) return self._kda_post(core_output, raw_output_gate, infer_state, layer_weight) def _hc_pre(self, streams, layer_weight, prefix, norm_weight): diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 3f20ad97ab..9d1ead737f 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -5,22 +5,24 @@ import torch import triton +from lightllm.common.attention.api import AttentionFamily from lightllm.common.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend -from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.glm5_next.attention.sparse_mla import Glm5NextSparseAttBackend +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + from .layer_infer.pre_layer_infer import Glm5NextPreLayerInfer from .layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from .layer_weights.pre_and_post_layer_weight import Glm5NextPreAndPostLayerWeight from .layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight - if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Glm5NextTpPartModel(TpPartBaseModel): @@ -129,10 +131,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" - from lightllm.common.attention_state.full_att import ( - PackedMLALogicalConfig, - create_packed_mla_state, - ) + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -155,15 +154,24 @@ def create_full_att_states( ) return [state] - def prepare_attention_states(self) -> None: - super().prepare_attention_states() - if self.args.llm_kv_type == "fp8kv_dsa" and self.att_state_manager.get_full_attention_states( - self.att_state_role - ): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + def attention_backend_candidates(self, family, phase): + from lightllm.models.glm5_next.attention.fp8_sparse_mla import Fp8Glm5NextSparseAttBackend - # Reserve fixed scratch space before the shared KV capacity is profiled. - Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + if family is AttentionFamily.KDA: + return None, ["triton"] + backend = Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + return {self.args.llm_kv_type: {"flashmla_sparse": backend}}, ["flashmla_sparse"] + + def _init_custom(self): + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + if self.args.enable_ep_moe: + dist_group_manager.new_deepep_group( + n_routed_experts=self.config["n_routed_experts"], + hidden_size=self.config["hidden_size"], + expert_quant_method_names=dist_group_manager.get_moe_quant_methods(self.trans_layers_weight), + num_experts_per_tok=self.config["num_experts_per_tok"], + moe_intermediate_size=self.config["moe_intermediate_size"], + ) def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json")) as f: @@ -192,34 +200,3 @@ def _init_some_value(self): self.kv_lora_rank = self.config["kv_lora_rank"] self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - - def _init_att_backend(self): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - ) - - backend_class = ( - Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend - ) - self.prefill_att_backend = backend_class(model=self) - self.decode_att_backend = self.prefill_att_backend - - def _init_att_backend1(self): - from lightllm.common.basemodel.attention.create_linear_utils import ( - get_kda_linear_att_backend_class, - ) - - backend_cls = get_kda_linear_att_backend_class() - self.prefill_att_backend1 = backend_cls(model=self) - self.decode_att_backend1 = self.prefill_att_backend1 - - def _init_custom(self): - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - if self.args.enable_ep_moe: - dist_group_manager.new_deepep_group( - n_routed_experts=self.config["n_routed_experts"], - hidden_size=self.config["hidden_size"], - expert_quant_method_names=dist_group_manager.get_moe_quant_methods(self.trans_layers_weight), - num_experts_per_tok=self.config["num_experts_per_tok"], - moe_intermediate_size=self.config["moe_intermediate_size"], - ) diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py new file mode 100644 index 0000000000..083ef6fe06 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -0,0 +1,744 @@ +import torch +import triton +import triton.language as tl + +# Bound the temporary BF16 gather buffer to 2 GiB (2097152 × 1 × 512 × 2 bytes). +KV_GATHER_CHUNK_SIZE = 2097152 + + +@triton.jit +def _direct_smallq_packed_kernel( + Q, + PackedKV, + RaggedMemIndex, + PrefillMemIndex, + PrefillKV, + KVStarts, + Topk, + PartialO, + PartialLse, + ragged_n, + q_stride_s: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_d: tl.constexpr, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + topk_stride_s: tl.constexpr, + topk_stride_d: tl.constexpr, + HEADS_: tl.constexpr, + BLOCK_H_: tl.constexpr, + PREFILL_N: tl.constexpr, + SPLITS_: tl.constexpr, + BLOCKS_: tl.constexpr, + SLOTS_: tl.constexpr, + softmax_scale: tl.constexpr, +): + qid = tl.program_id(0) + split = tl.program_id(1) + if HEADS_ == 8: + # H8 fits in one fixed 16-lane tile; the head mask excludes inactive lanes. + h = tl.arange(0, 16) + else: + h = tl.program_id(2) * BLOCK_H_ + tl.arange(0, BLOCK_H_) + d = tl.arange(0, 512) + hm = h < HEADS_ + q = tl.load( + Q + qid * q_stride_s + h[:, None] * q_stride_h + d[None, :] * q_stride_d, + hm[:, None], + 0.0, + ).to(tl.bfloat16) + if HEADS_ == 8: + acc = tl.zeros((16, 512), tl.float32) + m = tl.full((16,), -float("inf"), tl.float32) + denominator = tl.zeros((16,), tl.float32) + else: + acc = tl.zeros((BLOCK_H_, 512), tl.float32) + m = tl.full((BLOCK_H_,), -float("inf"), tl.float32) + denominator = tl.zeros((BLOCK_H_,), tl.float32) + start = tl.load(KVStarts + qid).to(tl.int64) + for b in tl.range(0, BLOCKS_, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): + n = tl.arange(0, 64) + slot = split * SLOTS_ + b * 64 + n + top = tl.load( + Topk + qid * topk_stride_s + slot * topk_stride_d, + slot < 2176, + -1, + ).to(tl.int64) + ragged_pos = start + top + valid = (top >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + mem = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + current_row = tl.full((64,), -1, tl.int32) + for j in tl.static_range(PREFILL_N): + current_row = tl.where(valid & (mem == tl.load(PrefillMemIndex + j).to(tl.int64)), j, current_row) + cached = valid & (current_row < 0) + current = valid & (current_row >= 0) + packed_row = PackedKV + mem[:, None] * packed_stride_s + fp8 = tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + d[None, :], + cached[:, None] & (d[None, :] < 448), + 0.0, + ).to(tl.float32) + scales = tl.load( + packed_row + 576 + d[None, :] // 64, + cached[:, None] & (d[None, :] < 448), + 127, + ) + fp8 = fp8 * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + d[None, :] - 448, + cached[:, None] & (d[None, :] >= 448), + 0.0, + ).to(tl.float32) + cached_v = tl.where(d[None, :] < 448, fp8, tail) + current_v = tl.load( + PrefillKV + current_row[:, None] * prefill_stride_s + d[None, :] * prefill_stride_d, + current[:, None], + 0.0, + ).to(tl.float32) + kv = tl.where(current[:, None], current_v, cached_v).to(tl.bfloat16) + score = tl.dot(q, tl.trans(kv), out_dtype=tl.float32) * softmax_scale + score = tl.where(valid[None, :], score, -float("inf")) + block_m = tl.max(score, axis=1) + new_m = tl.maximum(m, block_m) + alpha = tl.where(m == -float("inf"), 0.0, tl.exp(m - new_m)) + p = tl.where(valid[None, :], tl.exp(score - new_m[:, None]), 0.0) + denominator = denominator * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), kv, out_dtype=tl.float32) + m = new_m + has = denominator > 0.0 + out = tl.where(has[:, None], acc / denominator[:, None], 0.0) + tl.store( + PartialO + ((qid * SPLITS_ + split) * HEADS_ + h[:, None]) * 512 + d[None, :], + out, + hm[:, None], + ) + tl.store( + PartialLse + (qid * SPLITS_ + split) * HEADS_ + h, + tl.where(has, m + tl.log(denominator), -float("inf")), + hm, + ) + + +@triton.jit +def reduce_direct_smallq_kernel( + PartialO, + PartialLse, + Out, + HEADS_: tl.constexpr, + SPLITS_: tl.constexpr, + SPAD_: tl.constexpr, +): + qid = tl.program_id(0) + h = tl.program_id(1) + d = tl.program_id(2) * 128 + tl.arange(0, 128) + s = tl.arange(0, SPAD_) + sm = s < SPLITS_ + lse = tl.load(PartialLse + (qid * SPLITS_ + s) * HEADS_ + h, sm, -float("inf")) + maximum = tl.max(lse, axis=0) + safe_maximum = tl.where(maximum == -float("inf"), 0.0, maximum) + weight = tl.where(sm & (lse > -float("inf")), tl.exp(lse - safe_maximum), 0.0) + denominator = tl.sum(weight, axis=0) + partial = tl.load( + PartialO + ((qid * SPLITS_ + s[:, None]) * HEADS_ + h) * 512 + d[None, :], + sm[:, None], + 0.0, + ).to(tl.float32) + value = tl.sum(partial * weight[:, None], axis=0) / tl.where(denominator > 0.0, denominator, 1.0) + tl.store(Out + (qid * HEADS_ + h) * 512 + d, value) + + +def can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return ( + q.ndim == 3 + and 1 <= q.shape[0] <= 8 + and q.dtype == torch.bfloat16 + and (q.shape[1:] == (8, 512) or (q.shape[0] == 1 and q.shape[1:] == (64, 512))) + and packed_kv.ndim == 3 + and packed_kv.shape[1] == 1 + and packed_kv.shape[2] >= 584 + and ragged_mem_index.dtype == torch.int32 + and ragged_mem_index.numel() >= 65536 + and prefill_mem_index.dtype == torch.int32 + and prefill_mem_index.numel() == q.shape[0] + and prefill_kv.dtype == torch.bfloat16 + and prefill_kv.shape == (q.shape[0], 1, 512) + and topk_indices.dtype == torch.int32 + and topk_indices.shape == (q.shape[0], 2176) + ) + + +def allocate_direct_smallq_workspace(q): + qn, heads, _ = q.shape + splits = 34 if qn == 1 else 17 + return { + "partial_o": torch.empty((qn, splits, heads, q.shape[-1]), dtype=torch.float32, device=q.device), + "partial_lse": torch.empty((qn, splits, heads), dtype=torch.float32, device=q.device), + # The reducer requires a contiguous [Q, H, 512] destination. + "output": torch.empty(q.shape, dtype=q.dtype, device=q.device), + } + + +@torch.no_grad() +def _context_attention_fwd_fp8_direct_smallq( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + direct_workspace=None, +): + qn = q.shape[0] + heads = q.shape[1] + # Q1 uses 34 64-key splits; H64 uses two H32 head tiles per split. + splits, blocks, slots = (34, 1, 64) if qn == 1 else (17, 2, 128) + block_h = 32 if heads == 64 else 16 + if heads == 8: + # H8 uses local buffers because its direct path has no reusable workspace. + partial_o = torch.empty((qn, splits, 8, 512), dtype=torch.float32, device=q.device) + partial_lse = torch.empty((qn, splits, 8), dtype=torch.float32, device=q.device) + output = torch.empty(q.shape, dtype=q.dtype, device=q.device) + launch_grid = (qn, splits) + else: + workspace = direct_workspace or allocate_direct_smallq_workspace(q) + partial_o = workspace["partial_o"] + partial_lse = workspace["partial_lse"] + output = workspace["output"] + launch_grid = (qn, splits, triton.cdiv(heads, block_h)) + _direct_smallq_packed_kernel[launch_grid]( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + partial_o, + partial_lse, + ragged_mem_index.numel(), + q.stride(0), + q.stride(1), + q.stride(2), + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + topk_indices.stride(0), + topk_indices.stride(1), + HEADS_=heads, + BLOCK_H_=block_h, + PREFILL_N=qn, + SPLITS_=splits, + BLOCKS_=blocks, + SLOTS_=slots, + softmax_scale=softmax_scale, + num_warps=8, + ) + reduce_direct_smallq_kernel[(qn, heads, 4)]( + partial_o, + partial_lse, + output, + HEADS_=heads, + SPLITS_=splits, + SPAD_=triton.next_power_of_2(splits), + num_warps=4, + ) + return output + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _build_prefill_row_table_kernel(PrefillMemIndex, RowTable, num_tokens, BLOCK: tl.constexpr): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mem_index = tl.load(PrefillMemIndex + row, row < num_tokens, 0) + tl.store(RowTable + mem_index, row, row < num_tokens) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _gather_prefill_kv_kernel( + PackedKV, + MemIndex, + PrefillRowTable, + PrefillKV, + Output, + num_tokens, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + row = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dim = tl.arange(0, 512) + mem_index = tl.load(MemIndex + row, row < num_tokens, 0).to(tl.int64) + prefill_row = tl.load(PrefillRowTable + mem_index, row < num_tokens, -1) + cached = (row < num_tokens) & (prefill_row < 0) + current = (row < num_tokens) & (prefill_row >= 0) + packed_row = PackedKV + mem_index[:, None] * packed_stride_s + + quantized = tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], + cached[:, None] & (dim[None, :] < 448), + 0.0, + ).to(tl.float32) + scales = tl.load(packed_row + 576 + dim[None, :] // 64, cached[:, None] & (dim[None, :] < 448), 127) + values = quantized * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, + cached[:, None] & (dim[None, :] >= 448), + 0, + ).to(tl.float32) + values = tl.where(dim[None, :] < 448, values, tail) + current_values = tl.load( + PrefillKV + prefill_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + current[:, None], + 0, + ).to(tl.float32) + tl.store( + Output + row[:, None] * 512 + dim[None, :], + tl.where(current[:, None], current_values, values), + row[:, None] < num_tokens, + ) + + +@triton.jit(do_not_specialize=["chunk_end"]) +def _compact_topk_indices_kernel( + TopkIndices, + KVStarts, + ChunkIndices, + ChunkLengths, + chunk_start, + chunk_end, + TOPK: tl.constexpr, + BLOCK: tl.constexpr, + COMPACT: tl.constexpr = True, +): + row = tl.program_id(0) + lane = tl.arange(0, BLOCK) + index = tl.load(TopkIndices + row * TOPK + lane, lane < TOPK, -1) + ragged_index = index + tl.load(KVStarts + row) + valid = (lane < TOPK) & (index >= 0) & (ragged_index >= chunk_start) & (ragged_index < chunk_end) + if COMPACT: + compact_index = tl.cumsum(valid.to(tl.int32)) - 1 + tl.store(ChunkIndices + row * TOPK + compact_index, ragged_index - chunk_start, valid) + tl.store(ChunkLengths + row, tl.sum(valid.to(tl.int32))) + else: + tl.store( + ChunkIndices + row * TOPK + lane, + tl.where(valid, ragged_index - chunk_start, -1), + lane < TOPK, + ) + + +@triton.jit(do_not_specialize=["NUM_Q_TOKENS", "chunk_lse_stride_h"]) +def _merge_att_output_kernel( + Output, + Lse, + ChunkOutput, + ChunkLse, + chunk_lse_stride_s: tl.constexpr, + chunk_lse_stride_h, + NUM_HEADS: tl.constexpr, + NUM_Q_TOKENS, + BLOCK_HEADS: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS) + dim = tl.arange(0, 512) + previous_lse = tl.load(Lse + head * NUM_Q_TOKENS + row, head < NUM_HEADS, -float("inf")) + chunk_lse = tl.load( + ChunkLse + row * chunk_lse_stride_s + head * chunk_lse_stride_h, head < NUM_HEADS, -float("inf") + ) + # FA3 uses natural-log LSE and returns +inf for an empty row. + previous_lse = tl.where(previous_lse == float("inf"), -float("inf"), previous_lse) + chunk_lse = tl.where(chunk_lse == float("inf"), -float("inf"), chunk_lse) + maximum = tl.maximum(previous_lse, chunk_lse) + maximum = tl.where(maximum == -float("inf"), 0, maximum) + previous_weight = tl.exp(previous_lse - maximum) + chunk_weight = tl.exp(chunk_lse - maximum) + denominator = previous_weight + chunk_weight + offset = (row * NUM_HEADS + head[:, None]) * 512 + dim[None, :] + previous = tl.load(Output + offset, (head[:, None] < NUM_HEADS) & (previous_weight[:, None] > 0), 0) + current = tl.load(ChunkOutput + offset, (head[:, None] < NUM_HEADS) & (chunk_weight[:, None] > 0), 0).to(tl.float32) + result = (previous * previous_weight[:, None] + current * chunk_weight[:, None]) / tl.where( + denominator > 0, denominator, 1 + )[:, None] + tl.store(Output + offset, result, head[:, None] < NUM_HEADS) + tl.store( + Lse + head * NUM_Q_TOKENS + row, + tl.where(denominator > 0, maximum + tl.log(denominator), -float("inf")), + head < NUM_HEADS, + ) + + +@triton.jit(do_not_specialize=["ragged_n"]) +def _selected_prefill_compact_kernel( + Topk, + KVStarts, + RaggedMemIndex, + CompactPhysical, + Lengths, + PageTable, + CuQ, + ragged_n, + packed_rows: tl.constexpr, + topk_stride_s: tl.constexpr, + topk_stride_d: tl.constexpr, + BATCH: tl.constexpr, + BLOCK: tl.constexpr, +): + qid = tl.program_id(0) + lane = tl.arange(0, BLOCK) + local = tl.load(Topk + qid * topk_stride_s + lane * topk_stride_d, lane < 2176, -1).to(tl.int64) + ragged_pos = tl.load(KVStarts + qid).to(tl.int64) + local + valid = (lane < 2176) & (local >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + physical = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + valid = valid & (physical >= 0) & (physical < packed_rows) + rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 + tl.store(CompactPhysical + qid * 2176 + rank, physical, valid) + # Every selected slot has a stable contiguous page; invalid tail pages are + # harmless because cache_seqlens limits FA3 to Lengths[qid]. + tl.store(PageTable + qid * 2176 + lane, qid * 2176 + lane, lane < 2176) + if qid == 0: + tl.store(CuQ + lane, lane.to(tl.int32), lane <= BATCH) + tl.store(Lengths + qid, tl.sum(valid.to(tl.int32), axis=0)) + + +@triton.jit +def _selected_prefill_gather_kernel( + CompactPhysical, + Lengths, + PrefillMemIndex, + PrefillKV, + PackedKV, + SelectedKV, + CuK, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + BATCH: tl.constexpr, + BATCH_PAD: tl.constexpr, +): + qid = tl.program_id(0) + row = tl.program_id(1) * 4 + tl.arange(0, 4) + length = tl.load(Lengths + qid) + live = row < length + physical = tl.load(CompactPhysical + qid * 2176 + row, live, 0).to(tl.int64) + current_row = tl.full((4,), -1, tl.int32) + for j in tl.static_range(BATCH_PAD): + current = j < BATCH + current_mem = tl.load(PrefillMemIndex + j, current, 0).to(tl.int64) + current_row = tl.where(live & current & (physical == current_mem), j, current_row) + cached = live & (current_row < 0) + current = live & (current_row >= 0) + dim = tl.arange(0, 512) + packed_row = PackedKV + physical[:, None] * packed_stride_s + group = tl.arange(0, 8) + raw_scales = tl.load(packed_row + 576 + group[None, :], cached[:, None] & (group[None, :] < 7), 127) + scales = (raw_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + scales = tl.broadcast_to(scales[:, :, None], (4, 8, 64)).reshape(4, 512) + fp8 = ( + tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], + cached[:, None] & (dim[None, :] < 448), + 0.0, + ).to(tl.float32) + * scales + ) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, + cached[:, None] & (dim[None, :] >= 448), + 0.0, + ).to(tl.float32) + cached_value = tl.where(dim[None, :] < 448, fp8, tail) + current_value = tl.load( + PrefillKV + current_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + current[:, None], + 0.0, + ).to(tl.float32) + value = tl.where(current[:, None], current_value, cached_value).to(tl.bfloat16) + tl.store(SelectedKV + (qid * 2176 + row[:, None]) * 512 + dim[None, :], value, row[:, None] < 2176) + if (qid == 0) & (tl.program_id(1) == 0): + x = tl.arange(0, BATCH_PAD) + lengths = tl.load(Lengths + x, x < BATCH, 0) + tl.store(CuK, 0) + tl.store(CuK + 1 + x, tl.cumsum(lengths, axis=0), x < BATCH) + + +def can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return ( + q.ndim == 3 + and 2 <= q.shape[0] <= 8 + and q.shape[1:] == (64, 512) + and q.dtype == torch.bfloat16 + and q.is_cuda + and q.stride(-1) == 1 + and packed_kv.shape[1:] == (1, 584) + and packed_kv.dtype == torch.uint8 + and packed_kv.device == q.device + and ragged_mem_index.dtype == torch.int32 + and ragged_mem_index.device == q.device + and ragged_mem_index.numel() >= 65536 + and prefill_mem_index.dtype == torch.int32 + and prefill_mem_index.device == q.device + and prefill_mem_index.shape == (q.shape[0],) + and prefill_kv.dtype == torch.bfloat16 + and prefill_kv.device == q.device + and prefill_kv.shape == (q.shape[0], 1, 512) + and topk_indices.dtype == torch.int32 + and topk_indices.device == q.device + and topk_indices.shape == (q.shape[0], 2176) + ) + + +def allocate_selected_prefill_workspace(q, batch_capacity=None): + batch = q.shape[0] if batch_capacity is None else batch_capacity + return { + "compact": torch.empty((batch, 2176), dtype=torch.int64, device=q.device), + "lengths": torch.empty(batch, dtype=torch.int32, device=q.device), + "page_table": torch.empty((batch, 2176), dtype=torch.int32, device=q.device), + "cu_q": torch.empty(batch + 1, dtype=torch.int32, device=q.device), + "cu_k": torch.empty(batch + 1, dtype=torch.int32, device=q.device), + "selected_kv": torch.empty((batch * 2176, 1, q.shape[-1]), dtype=q.dtype, device=q.device), + } + + +def allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=None): + """Allocate generic prefill metadata up to the backend's planned query capacity.""" + capacity = q.shape[0] if query_capacity is None else query_capacity + return { + "prefill_row_table": torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device), + "chunk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), + "chunk_lengths": torch.empty(capacity, dtype=torch.int32, device=q.device), + "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), + "row_table_built": False, + } + + +@torch.no_grad() +def _context_attention_fwd_fp8_selected_prefill( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + selected_workspace=None, +): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + batch = q.shape[0] + workspace = selected_workspace or allocate_selected_prefill_workspace(q) + batch_pad = triton.next_power_of_2(batch) + _selected_prefill_compact_kernel[(batch,)]( + topk_indices, + kv_starts, + ragged_mem_index, + workspace["compact"], + workspace["lengths"], + workspace["page_table"], + workspace["cu_q"], + ragged_mem_index.numel(), + packed_kv.shape[0], + topk_indices.stride(0), + topk_indices.stride(1), + BATCH=batch, + BLOCK=4096, + num_warps=8, + ) + _selected_prefill_gather_kernel[(batch, triton.cdiv(2176, 4))]( + workspace["compact"], + workspace["lengths"], + prefill_mem_index, + prefill_kv, + packed_kv, + workspace["selected_kv"], + workspace["cu_k"], + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + BATCH=batch, + BATCH_PAD=batch_pad, + num_warps=4, + ) + return flash_attn_with_kvcache( + q=None, + qv=q, + k_cache=None, + v_cache=workspace["selected_kv"].view(-1, 1, 1, 512), + page_table=workspace["page_table"], + cache_seqlens=workspace["lengths"], + cu_seqlens_q=workspace["cu_q"], + cu_seqlens_k_new=workspace["cu_k"], + max_seqlen_q=1, + softmax_scale=softmax_scale, + causal=False, + only_qv=True, + ) + + +@torch.no_grad() +def context_attention_fwd_fp8( + q: torch.Tensor, + packed_kv: torch.Tensor, + ragged_mem_index: torch.Tensor, + prefill_mem_index: torch.Tensor, + prefill_kv: torch.Tensor, + kv_starts: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + selected_workspace=None, + direct_workspace=None, + generic_workspace=None, +): + """Run NoPE sparse prefill with a bounded gather buffer and BF16 current KV.""" + # Direct packed attention is intentionally restricted to the validated + # small-Q GLM-5.3 shape; all other cases retain the bounded-gather path. + if can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return _context_attention_fwd_fp8_direct_smallq( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + direct_workspace, + ) + if can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return _context_attention_fwd_fp8_selected_prefill( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + selected_workspace, + ) + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + num_q_tokens, num_heads, _ = q.shape + gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, ragged_mem_index.numel()) + if generic_workspace is None: + # Build the row table before gathering KV and allocating compaction metadata. + prefill_row_table = torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device) + _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( + prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 + ) + # Keep this potentially 2 GiB buffer local to a DSA layer. + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + chunk_indices = torch.empty_like(topk_indices) + chunk_lengths = torch.empty(num_q_tokens, dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.arange(num_q_tokens + 1, dtype=torch.int32, device=q.device) + else: + prefill_row_table = generic_workspace["prefill_row_table"] + chunk_indices = generic_workspace["chunk_indices"] + chunk_lengths = generic_workspace["chunk_lengths"] + cu_seqlens_q = generic_workspace["cu_seqlens_q"] + if not generic_workspace["row_table_built"]: + _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( + prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 + ) + generic_workspace["row_table_built"] = True + gathered_kv = generic_workspace.get("gathered_kv") + if gathered_kv is None: + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + gather_chunk_size = min(gather_chunk_size, gathered_kv.shape[0]) + output = lse = None + for chunk_start in range(0, ragged_mem_index.numel(), gather_chunk_size): + num_chunk_tokens = min(gather_chunk_size, ragged_mem_index.numel() - chunk_start) + _gather_prefill_kv_kernel[(triton.cdiv(num_chunk_tokens, 4),)]( + packed_kv, + ragged_mem_index[chunk_start : chunk_start + num_chunk_tokens], + prefill_row_table, + prefill_kv, + gathered_kv, + num_chunk_tokens, + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + BLOCK_ROWS=4, + ) + chunk_kv = gathered_kv[:num_chunk_tokens].view(-1, 1, 1, 512) + if ( + num_q_tokens > 8 + and num_heads == 64 + and ragged_mem_index.numel() > num_q_tokens + and ragged_mem_index.numel() <= gather_chunk_size + and q.dtype == torch.bfloat16 + and q.is_contiguous() + and packed_kv.dtype == torch.uint8 + and packed_kv.shape[1:] == (1, 584) + and topk_indices.dtype == torch.int32 + and topk_indices.shape[1] == 2176 + and topk_indices.is_contiguous() + ): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + _compact_topk_indices_kernel[(num_q_tokens,)]( + topk_indices, + kv_starts, + chunk_indices, + chunk_lengths, + chunk_start, + chunk_start + num_chunk_tokens, + TOPK=topk_indices.shape[1], + BLOCK=triton.next_power_of_2(topk_indices.shape[1]), + COMPACT=False, + ) + return flash_mla_sparse_fwd( + q=q, + kv=gathered_kv[:num_chunk_tokens], + indices=chunk_indices.unsqueeze(1), + sm_scale=softmax_scale, + d_v=512, + )[0] + _compact_topk_indices_kernel[(num_q_tokens,)]( + topk_indices, + kv_starts, + chunk_indices, + chunk_lengths, + chunk_start, + chunk_start + num_chunk_tokens, + TOPK=topk_indices.shape[1], + BLOCK=triton.next_power_of_2(topk_indices.shape[1]), + ) + chunk_output, chunk_lse, *_ = flash_attn_with_kvcache( + # only_qv ignores Q/K; these views avoid allocating placeholder tensors. + q=q[..., :64], + k_cache=chunk_kv[..., :64], + qv=q, + v_cache=chunk_kv, + page_table=chunk_indices, + cache_seqlens=chunk_lengths, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=softmax_scale, + causal=False, + only_qv=True, + num_splits=1, + return_softmax_lse=True, + ) + if chunk_start == 0: + output, lse = chunk_output, chunk_lse + else: + # Keep the online reduction in FP32 across all remaining chunks. + output = output.float() + _merge_att_output_kernel[(num_q_tokens, triton.cdiv(num_heads, 4))]( + output, + lse, + chunk_output, + chunk_lse, + chunk_lse.stride(1), + chunk_lse.stride(0), + NUM_HEADS=num_heads, + NUM_Q_TOKENS=num_q_tokens, + BLOCK_HEADS=4, + ) + del chunk_output, chunk_lse + return output.to(q.dtype) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 6c5f80a414..e0fead0715 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -1,22 +1,25 @@ -from typing import TYPE_CHECKING from copy import deepcopy +from typing import TYPE_CHECKING + import torch +from lightllm.common.attention.api import AttentionFamily from lightllm.common.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend +from lightllm.models.glm5_next.attention.fp8_sparse_mla import Fp8Glm5NextSparseAttBackend +from lightllm.models.glm5_next.attention.sparse_mla import Glm5NextSparseAttBackend from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + from .layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer from .layer_weights.pre_and_post_layer_weight import Glm5NextMTPPreAndPostLayerWeight - if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Glm5NextMTPModel(TpPartBaseModel): @@ -117,8 +120,8 @@ def create_bounded_att_states( ) -> list["BoundedAttentionState"]: """直接选择本模型的 K-pool 窗口层，不依赖 full State 创建过程。""" from lightllm.common.attention_state.bounded_att import ( - create_window_compressor_state, WindowCompressorStateConfig, + create_window_compressor_state, ) from lightllm.utils.torch_dtype_utils import get_torch_dtype @@ -153,20 +156,11 @@ def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, ob config.update(num_hidden_layers=1, n_layer=1, mhc=False) return config - def prepare_attention_states(self) -> None: - super().prepare_attention_states() - if self.args.llm_kv_type == "fp8kv_dsa" and self.att_state_manager.get_full_attention_states( - self.att_state_role - ): - # Reserve fixed scratch space before the shared KV capacity is profiled. - Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) - - def _init_att_backend(self): - backend_class = ( - Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend - ) - self.prefill_att_backend = backend_class(model=self) - self.decode_att_backend = self.prefill_att_backend + def attention_backend_candidates(self, family, phase): + if family is AttentionFamily.KDA: + return None, ["triton"] + backend = Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + return {self.args.llm_kv_type: {"flashmla_sparse": backend}}, ["flashmla_sparse"] def autotune_layers(self): return 1 diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index b6ebc52984..a5e1867b68 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -1,14 +1,21 @@ -from lightllm.common.attention_state.full_att import FullAttentionState import torch + +from lightllm.common.attention.api.inputs import DenseInputs from lightllm.models.gpt_oss.layer_weights.transformer_layer_weight import GptOssTransformerLayerWeight -from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer, LlamaInferStateInfo -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaInferStateInfo, LlamaTransformerLayerInfer from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class GptOssTransformerLayerInfer(LlamaTransformerLayerInfer): + attention_uses_sink = True + + def attention_window(self, phase): + if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention": + return (self.sliding_window - 1, 0) + return None + def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.hidden_size = self.network_config_["hidden_size"] @@ -67,65 +74,15 @@ def _context_attention_kernel( layer_weight: GptOssTransformerLayerWeight, out=None, ): - if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention": - window_size = (self.sliding_window - 1, 0) - use_sliding_window = True - else: - window_size = (-1, -1) - use_sliding_window = False - - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl( - use_sliding_window=use_sliding_window, - sliding_window=window_size, - use_att_sink=True, - sink_weight=layer_weight.attn_sinks.weight, - ), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor def _token_attention_kernel( self, q: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: GptOssTransformerLayerWeight, out=None ): - if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention": - window_size = (self.sliding_window - 1, 0) - use_sliding_window = True - else: - window_size = (-1, -1) - use_sliding_window = False - - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl( - use_sliding_window=use_sliding_window, - sliding_window=window_size, - use_att_sink=True, - sink_weight=layer_weight.attn_sinks.weight, - ), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.decode_attention.forward(infer_state, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor diff --git a/lightllm/models/gpt_oss/model.py b/lightllm/models/gpt_oss/model.py index 29606aba6a..199da8b49d 100644 --- a/lightllm/models/gpt_oss/model.py +++ b/lightllm/models/gpt_oss/model.py @@ -3,8 +3,6 @@ from lightllm.models.llama.model import LlamaTpPartModel from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.attention import get_prefill_att_backend_class, get_decode_att_backend_class -from lightllm.common.basemodel.attention import BaseAttBackend logger = init_logger(__name__) @@ -19,10 +17,5 @@ class GptOssTpPartModel(LlamaTpPartModel): def __init__(self, kvargs): super().__init__(kvargs) - def _init_att_backend(self): - self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0, priority_list=["fa3"])( - model=self - ) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0, priority_list=["fa3"])( - model=self - ) + def attention_backend_candidates(self, family, phase): + return None, ["fa3"] diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 0b3bb2e867..d5251005b6 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -1,15 +1,16 @@ -from lightllm.common.attention_state.full_att import FullAttentionState -from lightllm.common.basemodel.attention import AttControl +from functools import partial + import torch -import triton import torch.distributed as dist -from functools import partial -from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight -from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd -from lightllm.models.llama.infer_struct import LlamaInferStateInfo +import triton + +from lightllm.common.attention.api.inputs import DenseInputs from lightllm.common.basemodel import TransformerLayerInferTpl +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight +from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -46,21 +47,8 @@ def _context_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(logical_layer_id=self.layer_num_), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor @@ -70,21 +58,8 @@ def _token_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(logical_layer_id=self.layer_num_), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.decode_attention.forward(infer_state, DenseInputs(_q)) return o_tensor.view(q.shape) def _att_norm( diff --git a/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py b/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py index b2bf892f90..cb35a1de66 100644 --- a/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py @@ -1,22 +1,25 @@ import os -import torch from functools import partial from typing import Tuple + +import torch +import torch.distributed as dist + +from lightllm.distributed import all_reduce +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd -from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight -from lightllm.distributed import all_reduce -import torch.distributed as dist -from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo +from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class NeoChatTransformerLayerInfer(Qwen3TransformerLayerInfer): + attention_variant = "neo" + def __init__(self, data_type, network_config): super().__init__(data_type, network_config) return diff --git a/lightllm/models/neo_chat/model.py b/lightllm/models/neo_chat/model.py index 18fa41618a..386747baa2 100644 --- a/lightllm/models/neo_chat/model.py +++ b/lightllm/models/neo_chat/model.py @@ -1,39 +1,35 @@ -import os import json +import os + import torch + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer from lightllm.common.build_utils import repair_config +from lightllm.models.internvl.model import InternvlTokenizer +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.neo_chat.layer_infer.transformer_layer_infer import NeoChatTransformerLayerInfer +from lightllm.models.neo_chat.layer_weights.pre_and_post_layer_weight import NeoChatPreAndPostLayerWeight +from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo +from lightllm.models.neo_chat_moe.vision_process import smart_resize +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer +from lightllm.models.qwen3.model import Qwen3TpPartModel +from lightllm.models.qwen3_moe.model import Qwen3MOEModel from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer from lightllm.models.qwen3_vl.layer_infer.transformer_layer_infer import Qwen3VLTransformerLayerInfer from lightllm.models.qwen3_vl.layer_weights.pre_and_post_layer_weight import Qwen3VLPreAndPostLayerWeight -from lightllm.models.qwen2_vl.model import QWen2VLTokenizer -from lightllm.models.qwen3.model import Qwen3TpPartModel -from lightllm.server.core.objs import SamplingParams -from lightllm.models.qwen3_moe.model import Qwen3MOEModel -from lightllm.server.multimodal_params import AudioItem, MultimodalParams, ImageItem -from lightllm.models.neo_chat_moe.vision_process import smart_resize -from lightllm.models.internvl.model import InternvlTokenizer from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer -from lightllm.models.neo_chat.layer_infer.transformer_layer_infer import NeoChatTransformerLayerInfer -from lightllm.models.llama.infer_struct import LlamaInferStateInfo -from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight -from lightllm.models.neo_chat.layer_weights.pre_and_post_layer_weight import NeoChatPreAndPostLayerWeight -from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo -from lightllm.common.basemodel.attention import ( - get_neo_prefill_att_backend_class, - get_decode_att_backend_class, - BaseAttBackend, -) +from lightllm.server.core.objs import SamplingParams +from lightllm.server.multimodal_params import AudioItem, ImageItem, MultimodalParams from lightllm.utils.log_utils import init_logger - logger = init_logger(__name__) class NeoModelBase: def _init_custom(self): - super()._init_custom() if "rope_theta_hw" in self.config: @@ -123,7 +119,6 @@ def _init_to_get_hw_rotary(self, default_base=10000): class NeoTpPartModel(NeoModelBase, Qwen3TpPartModel): - pre_layer_infer_class = LlamaMultimodalPreLayerInfer transformer_layer_infer_class = NeoChatTransformerLayerInfer @@ -139,11 +134,12 @@ def __init__(self, kvargs): def _init_inferstate_cls(self): pass - def _init_att_backend(self): - self.prefill_att_backend: BaseAttBackend = get_neo_prefill_att_backend_class( - index=0, priority_list=["fa3", "triton"] - )(model=self) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) + def attention_backend_candidates(self, family, phase): + from lightllm.common.attention.registry import neo_cache_format_to_backend + + if phase is AttentionPhase.PREFILL: + return neo_cache_format_to_backend, ["fa3", "triton"] + return None, None def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: diff --git a/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py index 600d066787..2d27359da7 100644 --- a/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py @@ -1,22 +1,25 @@ import os -import torch from functools import partial from typing import Tuple + +import torch +import torch.distributed as dist + +from lightllm.distributed import all_reduce +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd -from lightllm.models.qwen3_moe.layer_infer.transformer_layer_infer import Qwen3MOETransformerLayerInfer +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight -from lightllm.distributed import all_reduce -import torch.distributed as dist -from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.qwen3_moe.layer_infer.transformer_layer_infer import Qwen3MOETransformerLayerInfer from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class NeoChatMOETransformerLayerInfer(Qwen3MOETransformerLayerInfer): + attention_variant = "neo" + def __init__(self, data_type, network_config): super().__init__(data_type, network_config) return diff --git a/lightllm/models/neo_chat_moe/model.py b/lightllm/models/neo_chat_moe/model.py index eda5554799..da6cdba23c 100644 --- a/lightllm/models/neo_chat_moe/model.py +++ b/lightllm/models/neo_chat_moe/model.py @@ -1,30 +1,27 @@ -import os import json +import os + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer from lightllm.common.build_utils import repair_config +from lightllm.models.internvl.model import InternvlTokenizer +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.neo_chat.model import NeoModelBase +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo +from lightllm.models.neo_chat_moe.layer_infer.transformer_layer_infer import NeoChatMOETransformerLayerInfer +from lightllm.models.neo_chat_moe.layer_weights.pre_and_post_layer_weight import NeoChatMOEPreAndPostLayerWeight +from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight +from lightllm.models.neo_chat_moe.vision_process import smart_resize +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer +from lightllm.models.qwen3.model import Qwen3TpPartModel +from lightllm.models.qwen3_moe.model import Qwen3MOEModel from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer from lightllm.models.qwen3_vl.layer_infer.transformer_layer_infer import Qwen3VLTransformerLayerInfer from lightllm.models.qwen3_vl.layer_weights.pre_and_post_layer_weight import Qwen3VLPreAndPostLayerWeight -from lightllm.models.qwen2_vl.model import QWen2VLTokenizer -from lightllm.models.qwen3.model import Qwen3TpPartModel -from lightllm.server.core.objs import SamplingParams -from lightllm.models.qwen3_moe.model import Qwen3MOEModel -from lightllm.server.multimodal_params import AudioItem, MultimodalParams, ImageItem -from lightllm.models.neo_chat_moe.vision_process import smart_resize -from lightllm.models.internvl.model import InternvlTokenizer from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer -from lightllm.models.neo_chat_moe.layer_infer.transformer_layer_infer import NeoChatMOETransformerLayerInfer -from lightllm.models.llama.infer_struct import LlamaInferStateInfo -from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight -from lightllm.models.neo_chat_moe.layer_weights.pre_and_post_layer_weight import NeoChatMOEPreAndPostLayerWeight -from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo -from lightllm.common.basemodel.attention import ( - get_neo_prefill_att_backend_class, - get_decode_att_backend_class, - BaseAttBackend, -) -from lightllm.models.neo_chat.model import NeoModelBase +from lightllm.server.core.objs import SamplingParams +from lightllm.server.multimodal_params import AudioItem, ImageItem, MultimodalParams IMG_START_TOKEN = "" IMG_END_TOKEN = "" @@ -81,7 +78,7 @@ def get_image_token_length(self, img: ImageItem): max_pixels=img.extra_params.get("max_pixels", self.max_pixel), ) grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size - token_num = int((grid_h * grid_w) * (self.downsample_ratio ** 2)) + token_num = int((grid_h * grid_w) * (self.downsample_ratio**2)) # 这里的grid_h和grid_w需要* self.downsample_ratio么？再仔细看下代码 img.grid_thwd = (1, int(grid_h * self.downsample_ratio), int(grid_w * self.downsample_ratio), 1 - token_num) return token_num @@ -92,7 +89,7 @@ def get_image_token_length_by_size(self, width: int, height: int): ), "width and height should be greater than 0 and divisible by patch_size" grid_h, grid_w = height // self.patch_size, width // self.patch_size - token_num = int((grid_h * grid_w) * (self.downsample_ratio ** 2)) + token_num = int((grid_h * grid_w) * (self.downsample_ratio**2)) return token_num # only change the impl of the encode func: @@ -166,7 +163,6 @@ def get_query_for_t2i(self, prompt: str, input_image_num: int = 0): class NeoTpMOEPartModel(NeoModelBase, Qwen3MOEModel): - pre_layer_infer_class = LlamaMultimodalPreLayerInfer transformer_layer_infer_class = NeoChatMOETransformerLayerInfer @@ -182,11 +178,12 @@ def __init__(self, kvargs): def _init_inferstate_cls(self): pass - def _init_att_backend(self): - self.prefill_att_backend: BaseAttBackend = get_neo_prefill_att_backend_class( - index=0, priority_list=["fa3", "triton"] - )(model=self) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) + def attention_backend_candidates(self, family, phase): + from lightllm.common.attention.registry import neo_cache_format_to_backend + + if phase is AttentionPhase.PREFILL: + return neo_cache_format_to_backend, ["fa3", "triton"] + return None, None def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index f0683dd882..462c86f74c 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -1,7 +1,5 @@ import torch -from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.models.llama.model import LlamaTpPartModel @@ -46,13 +44,6 @@ def _init_custom(self): self.block_size = int(self.config["block_size"]) self.mask_token_id = int(self.config["mask_token_id"]) - def _init_att_backend(self): - super()._init_att_backend() - # FA3 is currently the only backend that supports non-causal block decode. - # TODO: Remove this restriction after Triton and FlashInfer support non-causal block decode. - if not isinstance(self.decode_att_backend, (Fa3AttBackend, Fp8Fa3AttBackend)): - raise NotImplementedError("Qwen3 DFlash decode requires FA3") - def _init_infer_layer(self, start_layer_index=None): self.draft_layer_start = len(self.main_model.layers_infer) self.draft_layer_start += sum( diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index 7a631622d6..1b293717ab 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -1,10 +1,20 @@ +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import DenseInputs from lightllm.common.attention_state.full_att import FullAttentionState -from lightllm.common.basemodel.attention import AttControl -from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer from lightllm.models.qwen3_dflash2.triton_kernel import grouped_dynamic_conv +from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer class Qwen3DFlash2TransformerLayerInfer(Qwen3DFlashTransformerLayerInfer): + def attention_window(self, phase): + if phase is AttentionPhase.DECODE and self.sliding_window_ > 0: + return (self.sliding_window_ - 1, self.sliding_window_ - 1) + return None + + def attention_kv_view(self, state): + k, v = state.get_layer_view(self.layer_num_) + return k.view(-1, self.tp_k_head_num_, self.head_dim_), v.view(-1, self.tp_k_head_num_, self.head_dim_) + """DFlash2 layer with a grouped dynamic convolution around each sublayer.""" def __init__(self, layer_num, network_config): @@ -24,21 +34,11 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - use_sliding_window = self.sliding_window_ > 0 + # DFlash2 attention is non-causal, so its local window must retain both # the visible prefix on the left and the complete draft block on the right. - window = self.sliding_window_ - 1 - output = infer_state.decode_att_state.decode_att( - q=q, - k=k, - v=v, - att_control=AttControl( - logical_layer_id=self.layer_num_, - use_sliding_window=use_sliding_window, - sliding_window=(window, window) if use_sliding_window else (-1, -1), - ), - alloc_func=self.alloc_tensor, - ) + + output = self.decode_attention.forward(infer_state, DenseInputs(q)) return output.view(-1, self.tp_q_head_num_ * self.head_dim_) def _post_cache_kv(self, cache_kv, infer_state, layer_weight): diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 1bc4ef914b..e4c0473652 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -1,25 +1,28 @@ -import torch +from functools import partial +from typing import Tuple +import torch import torch.distributed as dist -from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( - Qwen3NextTransformerLayerWeight, -) + +from lightllm.common.attention.api import AttentionFamily +from lightllm.common.attention.api.inputs import GDNInputs +from lightllm.distributed import all_reduce from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer +from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo -from lightllm.utils.log_utils import init_logger -from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor -from lightllm.common.basemodel.attention.base_att import AttControl -from typing import Tuple +from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import Qwen3NextTransformerLayerWeight from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ -from lightllm.distributed import all_reduce -from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.utils.envs_utils import get_env_start_args -from functools import partial +from lightllm.utils.log_utils import init_logger +from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor logger = init_logger(__name__) class Qwen3NextTransformerLayerInfer(LlamaTransformerLayerInfer): + def get_attention_family(self): + return AttentionFamily.GDN if self.is_linear_attention_layer else AttentionFamily.DENSE + def __init__(self, layer_num, network_config): self.partial_rotary_factor = network_config.get("partial_rotary_factor", 1.0) self.n_routed_experts = network_config.get("num_experts", 0) @@ -134,13 +137,13 @@ def _get_qkv( qkv_gate_out = layer_weight.qkvo_gate_proj.mm(input) qkv_out, o_gate = qkv_gate_out.split( [ - self.tp_q_head_num_ * self.head_dim_ * 2 + (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_, + self.tp_q_head_num_ * self.head_dim_ + (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_, self.tp_q_head_num_ * self.head_dim_, ], dim=-1, ) q, cache_kv = qkv_out.split( - [self.tp_q_head_num_ * self.head_dim_ * 2, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_], + [self.tp_q_head_num_ * self.head_dim_, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_], dim=-1, ) infer_state.gate_logics_value = o_gate @@ -184,8 +187,6 @@ def _linear_prefill_cuda_graph_wrapper( infer_state: Qwen3NextInferStateInfo, layer_weight: Qwen3NextTransformerLayerWeight, ) -> Tuple[torch.Tensor, torch.Tensor]: - backend = infer_state.prefill_att_state1.backend - mixed_qkvzba = mixed_qkvzba.contiguous() _mixed_qkvzba = tensor_to_no_ref_tensor(mixed_qkvzba) pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() @@ -196,7 +197,11 @@ def _linear_prefill_cuda_graph_wrapper( # We avoid a dry-run because FlashQLA may do host-side syncs while # preparing varlen chunk metadata, which is illegal during capture. num_tokens = mixed_qkvzba.shape[0] - o_shape = (num_tokens, backend.tp_num_v_heads, backend.head_v_dim) + o_shape = ( + num_tokens, + layer_weight.linear_num_v_heads // self.tp_world_size_, + layer_weight.linear_v_head_dim, + ) o_dtype = mixed_qkvzba.dtype o_device = mixed_qkvzba.device z_shape = o_shape @@ -209,23 +214,8 @@ def _linear_prefill_cuda_graph_wrapper( _z = tensor_to_no_ref_tensor(z) def gdn_prefill_func(new_infer_state: "Qwen3NextInferStateInfo"): - # 与 Llama/FA3 一致：通过 new_infer_state.prefill_att_state1 进入， # 使用 runtime 上 init_att_state() 更新过的 buffer idx，而不是 - # capture 时闭包住的 graph prefill_att_state1。 - tmp_o, tmp_z = new_infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkvzba": _mixed_qkvzba, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + tmp_o, tmp_z = self.prefill_attention.forward(new_infer_state, GDNInputs(_mixed_qkvzba)) tmp_o = tmp_o.view(_o.shape) _o.copy_(tmp_o) _z.copy_(tmp_z) @@ -249,20 +239,7 @@ def context_attention_forward( if torch.cuda.is_current_stream_capturing(): core_attn_out, z = self._linear_prefill_cuda_graph_wrapper(mixed_qkvzba, infer_state, layer_weight) else: - core_attn_out, z = infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkvzba": mixed_qkvzba, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + core_attn_out, z = self.prefill_attention.forward(infer_state, GDNInputs(mixed_qkvzba)) gdn_out = self._linear_post(core_attn_out, z, layer_weight) @@ -280,20 +257,7 @@ def token_attention_forward( return super().token_attention_forward(input_embdings, infer_state, layer_weight) mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) - core_attn_out, z = infer_state.decode_att_state1.decode_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_decode=True, - linear_att_decode_dict={ - "mixed_qkvzba": mixed_qkvzba, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + core_attn_out, z = self.decode_attention.forward(infer_state, GDNInputs(mixed_qkvzba)) gdn_out = self._linear_post(core_attn_out, z, layer_weight) if self.tp_world_size_ > 1: diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 57de61241b..fe0992251a 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -1,23 +1,22 @@ +from typing import TYPE_CHECKING, Optional + import torch -from typing import Optional, TYPE_CHECKING import triton + from lightllm.models.qwen3_moe.model import Qwen3MOEModel -from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( - Qwen3NextTransformerLayerWeight, -) -from lightllm.models.qwen3next.layer_weights.pre_and_post_layer_weight import Qwen3NextPreAndPostLayerWeight -from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import ( - Qwen3NextTransformerLayerInfer, -) +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo +from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerInfer +from lightllm.models.qwen3next.layer_weights.pre_and_post_layer_weight import Qwen3NextPreAndPostLayerWeight +from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import Qwen3NextTransformerLayerWeight from lightllm.utils.log_utils import init_logger -from lightllm.models.qwen3next.config import Qwen3NextResourceConfig if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs logger = init_logger(__name__) @@ -51,7 +50,7 @@ def create_bounded_att_states( device: str = "meta", ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" - from lightllm.common.attention_state.bounded_att import create_gdn_state, ConvSSMStateConfig + from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, create_gdn_state config = cls.normalize_model_config(model_config) assert config["num_attention_heads"] % tp_world_size == 0 @@ -136,20 +135,3 @@ def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch def _init_config(self): super()._init_config() self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - - def _init_att_backend1(self): - role = self.att_state_role - if not self.att_state_manager.get_bounded_attention_states(role): - self.prefill_att_backend1 = None - self.decode_att_backend1 = None - return - from lightllm.common.basemodel.attention.create_linear_utils import ( - get_gdn_linear_prefill_att_backend_class, - get_gdn_linear_decode_att_backend_class, - ) - - linear_prefill_att_backend_class = get_gdn_linear_prefill_att_backend_class(index=1) - linear_decode_att_backend_class = get_gdn_linear_decode_att_backend_class(index=1) - self.prefill_att_backend1 = linear_prefill_att_backend_class(model=self) - self.decode_att_backend1 = linear_decode_att_backend_class(model=self) - return diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 068fc207b1..a5c386ce03 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -494,27 +494,33 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--llm_prefill_att_backend", type=str, - nargs="+", - choices=["auto", "triton", "fa3", "flashinfer", "flashqla"], - default=["auto"], - help="""prefill attention kernel used in llm. - auto: automatically select best backend based on GPU and available packages - (priority: fa3 > flashinfer > triton) - for hybrid linear-attention models, the second value selects the linear-attention backend - (priority: flashqla > triton); when omitted, it defaults to auto""", + choices=["auto", "triton", "fa3", "flashinfer", "flashmla_sparse"], + default="auto", + help="""prefill attention backend. auto selects fa3 > flashinfer > triton; + NSA uses flashmla_sparse. Linear attention is configured separately.""", ) parser.add_argument( "--llm_decode_att_backend", type=str, - nargs="+", - choices=["auto", "triton", "fa3", "flashinfer"], - default=["auto"], - help="""decode attention kernel used in llm. - auto: automatically select best backend based on GPU and available packages - (priority when mtp_step > 0: fa3 > flashinfer > triton; - otherwise: flashinfer > fa3 > triton) - for hybrid linear-attention models, the second value selects the linear-attention backend - (currently triton only); when omitted, it defaults to auto""", + choices=["auto", "triton", "fa3", "flashinfer", "flashmla_sparse"], + default="auto", + help="""decode attention backend. auto prefers fa3 with MTP, otherwise flashinfer; + NSA uses flashmla_sparse. Linear attention is configured separately.""", + ) + parser.add_argument( + "--llm_prefill_linear_att_backend", + type=str, + choices=["auto", "triton", "flashqla"], + default="auto", + help="""linear attention prefill backend. GDN auto selects flashqla > triton; + KDA supports triton only.""", + ) + parser.add_argument( + "--llm_decode_linear_att_backend", + type=str, + choices=["auto", "triton"], + default="auto", + help="linear attention decode backend (currently triton only).", ) parser.add_argument( "--page_size", @@ -1017,8 +1023,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--hardware_platform", type=str, default="cuda", - choices=["cuda", "musa"], - help="""Hardware platform: cuda | musa""", + choices=["cuda"], + help="Hardware platform: cuda", ) parser.add_argument( "--enable_torch_fallback", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index e0e6038d60..5a0dc9739d 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -1,40 +1,42 @@ +import math import multiprocessing as mp import os -import uuid import subprocess -import math -from lightllm.utils.start_utils import configure_decode_open_file_limit, process_manager -from .metrics.manager import start_metric_manager -from .httpserver.req_id_generator_server import start_req_id_generator_server -from .embed_cache.manager import start_cache_manager -from lightllm.utils.log_utils import init_logger +import uuid + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.config_utils import ( + auto_set_cpu_cache_token_page_size, + auto_set_fused_shared_experts, + auto_set_max_req_total_len, + auto_set_page_size, + auto_set_response_parsers, + get_default_disable_chunked_mm_input, + get_running_max_req_size_per_dp, + has_audio_module, + has_vision_module, + is_hybrid_att_model, +) +from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args from lightllm.utils.envs_utils import ( get_env_start_args, get_unique_server_name, set_env_start_args, set_unique_server_name, ) -from lightllm.utils.shm_port_args import get_shm_port_args -from lightllm.utils.net_utils import validate_ports -from .detokenization.manager import start_detokenization_process -from .router.manager import start_router_process +from lightllm.utils.log_utils import init_logger from lightllm.utils.multinode_utils import send_and_receive_node_ip +from lightllm.utils.net_utils import validate_ports from lightllm.utils.redis_utils import start_redis_service +from lightllm.utils.shm_port_args import get_shm_port_args from lightllm.utils.shm_size_check import check_recommended_shm_size -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.config_utils import ( - has_audio_module, - has_vision_module, - is_hybrid_att_model, - auto_set_max_req_total_len, - auto_set_page_size, - auto_set_cpu_cache_token_page_size, - auto_set_fused_shared_experts, - auto_set_response_parsers, - get_running_max_req_size_per_dp, - get_default_disable_chunked_mm_input, -) -from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args +from lightllm.utils.start_utils import configure_decode_open_file_limit, process_manager + +from .detokenization.manager import start_detokenization_process +from .embed_cache.manager import start_cache_manager +from .httpserver.req_id_generator_server import start_req_id_generator_server +from .metrics.manager import start_metric_manager +from .router.manager import start_router_process logger = init_logger(__name__) @@ -184,17 +186,11 @@ def _launch_subprocesses(args: StartArgs): assert args.enable_tpsp_mix_mode and args.dp > 1, "need set --enable_tpsp_mix_mode firstly and --dp > 1" if args.enable_ep_moe: - allowed_ep_prefill_att_backends = {"auto", "fa3", "triton", "flashqla"} - for backend in args.llm_prefill_att_backend: - assert backend in allowed_ep_prefill_att_backends, ( - "When --enable_ep_moe is enabled, --llm_prefill_att_backend must be one of " - f"{sorted(allowed_ep_prefill_att_backends)}; flashinfer is not supported." - ) - allowed_ep_decode_att_backends = {"auto", "fa3", "triton"} - for backend in args.llm_decode_att_backend: - assert backend in allowed_ep_decode_att_backends, ( - "When --enable_ep_moe is enabled, --llm_decode_att_backend must be one of " - f"{sorted(allowed_ep_decode_att_backends)}; flashinfer is not supported." + allowed_ep_att_backends = {"auto", "fa3", "triton"} + for option in ("llm_prefill_att_backend", "llm_decode_att_backend"): + assert getattr(args, option) in allowed_ep_att_backends, ( + f"When --enable_ep_moe is enabled, --{option} must be one of " + f"{sorted(allowed_ep_att_backends)}; flashinfer is not supported." ) if args.page_size < 1: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 033ca9e68a..93edb9af65 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -175,13 +175,14 @@ class StartArgs: vit_quant_type: Optional[str] = field(default="none") vit_quant_cfg: Optional[str] = field(default=None) expert_dtype: Optional[str] = field(default=None, metadata={"choices": ["fp8", "fp4"]}) - llm_prefill_att_backend: List[str] = field( - default_factory=lambda: ["auto"], - metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashqla"]}, + llm_prefill_att_backend: str = field( + default="auto", metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashmla_sparse"]} ) - llm_decode_att_backend: List[str] = field( - default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "flashinfer"]} + llm_decode_att_backend: str = field( + default="auto", metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashmla_sparse"]} ) + llm_prefill_linear_att_backend: str = field(default="auto", metadata={"choices": ["auto", "triton", "flashqla"]}) + llm_decode_linear_att_backend: str = field(default="auto", metadata={"choices": ["auto", "triton"]}) page_size: int = field(default=1) req_token_index_mode: str = field( default="pin_mem_cache", @@ -244,7 +245,7 @@ class StartArgs: enable_rl: bool = field(default=False) enable_torch_memory_saver: bool = field(default=False) enable_weight_cpu_backup: bool = field(default=False) - hardware_platform: str = field(default="cuda", metadata={"choices": ["cuda", "musa"]}) + hardware_platform: str = field(default="cuda", metadata={"choices": ["cuda"]}) enable_torch_fallback: bool = field(default=False) enable_triton_fallback: bool = field(default=False) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 09eb22be36..3408721860 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -1,56 +1,60 @@ +import threading +import time +from typing import Callable, List, Optional, Tuple + import numpy as np import torch -import time -import threading import torch.distributed as dist -from typing import List, Tuple, Callable, Optional from transformers.configuration_utils import PretrainedConfig -from lightllm.utils.infer_utils import set_random_seed -from lightllm.utils.log_utils import init_logger -from lightllm.models import get_draft_model_class, get_model -from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack -from lightllm.server.router.model_infer.req_token_index import ReqTokenIndex -from lightllm.server.router.token_load import TokenLoad + from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.common.basemodel.logprobs_manager import PromptLogprobsCaptureManager from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager -from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache -from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache -from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput -from lightllm.utils.dist_utils import init_distributed_env +from lightllm.common.basemodel.triton_kernel.gather_token_id import scatter_token +from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async +from lightllm.distributed import dist_group_manager +from lightllm.distributed.communication_op import all_gather_into_tensor, all_reduce, broadcast +from lightllm.models import get_draft_model_class, get_model from lightllm.server.core.objs import ShmReqManager, StartArgs from lightllm.server.core.objs.io_objs import AbortedReqCmd -from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer +from lightllm.server.multi_level_kv_cache import CacheTier, create_cache_placement_controller +from lightllm.server.pd_io_struct import PDChunckedTransTaskRet +from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache +from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache +from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack, g_infer_context +from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager, OverlapEventPack from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -from lightllm.utils.dist_utils import get_global_rank, get_global_world_size, get_dp_size -from lightllm.utils.dist_utils import get_dp_world_size, get_global_dp_rank, get_current_rank_in_dp -from lightllm.utils.dist_utils import get_current_device_id, get_current_rank_in_node, get_node_world_size -from lightllm.utils.dist_utils import get_dp_rank_in_node, create_new_group_for_current_node +from lightllm.server.router.model_infer.req_token_index import ReqTokenIndex +from lightllm.server.router.token_load import TokenLoad +from lightllm.utils.dist_utils import ( + create_new_group_for_current_node, + get_current_device_id, + get_current_rank_in_dp, + get_current_rank_in_node, + get_dp_rank_in_node, + get_dp_size, + get_dp_world_size, + get_global_dp_rank, + get_global_rank, + get_global_world_size, + get_node_world_size, + init_distributed_env, +) from lightllm.utils.envs_utils import ( - get_env_start_args, enable_radix_tree_timer_merge, + get_env_start_args, get_radix_tree_merge_update_delta, ) -from lightllm.distributed import dist_group_manager -from lightllm.distributed.communication_op import ( - all_gather_into_tensor, - all_reduce, - broadcast, -) -from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer -from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager, OverlapEventPack -from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample -from lightllm.common.basemodel.triton_kernel.gather_token_id import scatter_token -from lightllm.server.pd_io_struct import PDChunckedTransTaskRet -from lightllm.server.multi_level_kv_cache import ( - CacheTier, - create_cache_placement_controller, -) -from .prefill_queue_strategy import create_prefill_queue_strategy -from .multi_level_kv_cache import MultiLevelKvCacheModule +from lightllm.utils.infer_utils import set_random_seed +from lightllm.utils.log_utils import init_logger from lightllm.utils.profiler import ProcessProfiler, ProfilerCmd +from .multi_level_kv_cache import MultiLevelKvCacheModule +from .prefill_queue_strategy import create_prefill_queue_strategy + class ModeBackend: def __init__(self) -> None: @@ -149,10 +153,10 @@ def init_model(self, kvargs): self.draft_models = [] if self.args.mtp_mode is not None: self.init_mtp_draft_model(model_kvargs) - from lightllm.utils.envs_utils import set_model_init_status - from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.attention_state import AttentionStateManager from lightllm.common.req_manager import ReqManager + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag models = (self.model, *self.draft_models) assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" @@ -169,11 +173,20 @@ def init_model(self, kvargs): item.prepare_attention_states() manager: AttentionStateManager = self.model.att_state_manager assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" + manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + from lightllm.common.basemodel.attention import reserve_attention_resources + + for item in models: + reserve_attention_resources(item) with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - manager.finish_state_registration() manager.load_full_att_calibration(args=self.model.args) manager.initialize_storage( - self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction + self.model.args, + token_capacity=self.model.max_total_token_num, + mem_fraction=self.model.mem_fraction, + device=self.model.device_runtime.device, ) for item in models: item.init_runtime() diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 9bc8c16cf7..ebb70e7e07 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -1,23 +1,22 @@ -import torch import time from typing import List + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc +from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack -from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.router.model_infer.mode_backend.pre import ( - prepare_prefill_inputs, - prepare_decode_inputs, -) from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample -from lightllm.server.router.model_infer.infer_batch import g_infer_context -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack +from lightllm.server.router.model_infer.mode_backend.pre import prepare_decode_inputs, prepare_prefill_inputs from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils -from lightllm.utils.log_utils import init_logger +from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.dist_utils import get_current_device_id -from .control_state import ControlState from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from .control_state import ControlState logger = init_logger(__name__) @@ -115,7 +114,12 @@ def prefill_normal( event_pack.notify_post_handle_and_wait_pre_post_handle() with torch.cuda.stream(g_infer_context.get_overlap_stream()): - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -160,7 +164,12 @@ def decode_normal( event_pack.notify_post_handle_and_wait_pre_post_handle() with torch.cuda.stream(g_infer_context.get_overlap_stream()): - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( logits=model_output.logits, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -291,11 +300,7 @@ def decode_mtp( b_req_mtp_start_loc = gen_b_req_mtp_start_loc(model_input.b_mtp_index, num_reqs=req_num) mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( - backend=self, - next_token_ids=next_token_ids, - b_req_idx=model_input.b_req_idx, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_mtp_index=model_input.b_mtp_index, + backend=self, next_token_ids=next_token_ids, model_output=model_output, num_reqs=req_num ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index ed7e5bb167..b7852950c5 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -1,24 +1,27 @@ -import torch import time from typing import List, Tuple + +import torch + +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc +from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput -from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack from lightllm.server.router.model_infer.mode_backend.pre import ( - prepare_prefill_inputs, - prepare_decode_inputs, - overlap_prepare_prefill_inputs, overlap_prepare_decode_inputs, + overlap_prepare_prefill_inputs, + prepare_decode_inputs, + prepare_prefill_inputs, ) -from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack +from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_engine import DPOverlapSpecEngine +from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_env_start_args -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_engine import DPOverlapSpecEngine -from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils + from .control_state import DPControlState @@ -485,7 +488,6 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): """复用普通 SpecEngine 执行 DP speculative draft-and-verify。""" - model_input, run_reqs = prepare_decode_inputs(req_objs=decode_reqs) spec_engine = self.spec_engine req_num = len(decode_reqs) @@ -526,11 +528,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): ) mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( - backend=self, - next_token_ids=next_token_ids, - b_req_idx=model_input.b_req_idx, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_mtp_index=model_input.b_mtp_index, + backend=self, next_token_ids=next_token_ids, model_output=model_output, num_reqs=req_num ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -802,9 +800,8 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( backend=self, next_token_ids=next_token_ids, - b_req_idx=b_req_idx, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_mtp_index=b_mtp_index, + model_output=(model_output0, model_output1), + num_reqs=req_num, ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] mtp_accept_len1 = mtp_accept_len[real_request_num0:] diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index bfb55d37a2..928dd67140 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -1,19 +1,22 @@ -import torch.distributed as dist -import torch -import dataclasses import bisect -from functools import lru_cache -from typing import Optional, List, Deque +import dataclasses from collections import deque +from functools import lru_cache +from typing import Deque, List, Optional + +import torch +import torch.distributed as dist + from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -from lightllm.utils.envs_utils import get_env_start_args -from ..infer_batch import InferReq -from lightllm.utils.dist_utils import create_new_group_for_current_dp from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.utils.dist_utils import create_new_group_for_current_dp +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len +from ..infer_batch import InferReq + logger = init_logger(__name__) @@ -45,20 +48,14 @@ def __init__(self, backend): @lru_cache() def need_sync_compute_stream(self) -> bool: """不支持异步 CPU cache 的 backend（目前是 FA3）仍要等计算流，否则搬运可能崩溃。""" - model = self.backend.model - att_backends = [ - model.prefill_att_backend, - model.decode_att_backend, - model.prefill_att_backend1, - model.decode_att_backend1, - ] - for att_backend in att_backends: - if att_backend is not None and att_backend.requires_cpu_cache_sync: - logger.info( - "MultiLevelKvCacheModule: %s requires synchronous CPU cache IO.", type(att_backend).__name__ - ) - return True + for group in model.attention_groups: + for backend in (group.prefill_backend, group.decode_backend): + if backend.requires_cpu_cache_sync: + logger.info( + "MultiLevelKvCacheModule: %s requires synchronous CPU cache IO.", type(backend).__name__ + ) + return True logger.info("MultiLevelKvCacheModule: CPU cache load waits on a recorded event.") return False @@ -282,8 +279,13 @@ def write_req_to_cpu(self, req: InferReq) -> Optional["OffloadTask"]: """选择请求可缓存的前缀，分配 CPU 页并提交写出，返回待确认完成的任务。""" assert CacheTier.CPU in req.cache_tiers disk_offload_enable = CacheTier.DISK in req.cache_tiers - # KV 和 bounded 快照均在 offload 流上复制，由后续完成事件统一确认就绪。 - with torch.cuda.stream(self.get_cpu_offload_stream()): + cpu_kv_cache_stream = self.get_cpu_offload_stream() + producer = torch.cuda.current_stream() + cpu_kv_cache_stream.wait_stream(producer) + overlap = g_infer_context.get_overlap_stream() + if overlap != producer: + cpu_kv_cache_stream.wait_stream(overlap) + with torch.cuda.stream(cpu_kv_cache_stream): # 综合考虑后只对prompt做缓存管理，不包含decode内容，这里与radix cache不一致 token_hash_list = req.shm_req.token_hash_list.get_all() page_len_list = req.shm_req.token_hash_page_len_list.get_all() diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 19d4656f14..26988fec10 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -5,7 +5,9 @@ import torch +from lightllm.common.attention_state.commit import BoundedStateCommit from lightllm.common.basemodel.triton_kernel.mtp_utils import ( + gen_b_req_mtp_start_loc, mtp_scatter_next_token_ids, mtp_verify, ) @@ -16,28 +18,28 @@ from lightllm.server.router.model_infer.mtp_speculative.proposers.base import SpecProposal -def verify_mtp_tokens( - backend: ModeBackend, - next_token_ids: torch.Tensor, - b_req_idx: torch.Tensor, - b_req_mtp_start_loc: torch.Tensor, - b_mtp_index: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Verify target tokens and update recurrent MTP state when required.""" - # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 - +def verify_mtp_tokens(backend: ModeBackend, next_token_ids: torch.Tensor, model_output, num_reqs: int): + outputs = model_output if isinstance(model_output, tuple) else (model_output,) + commits = tuple(output.attention_commit for output in outputs) + if any(commit is None for commit in commits): + raise ValueError("MTP verification requires an output-bound attention commit") + commit = commits[0] if len(commits) == 1 else BoundedStateCommit.concatenate(commits) + if commit.role != "main": + raise ValueError("target verification cannot commit a draft model output") + commit.validate_generation(backend.model.req_manager.request_generations) + starts = gen_b_req_mtp_start_loc(b_mtp_index=commit.mtp_indices, num_reqs=num_reqs) accept_lengths, accepted_index = mtp_verify( req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids, - b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_mtp_start_loc=starts, new_next_token_ids=next_token_ids, - b_req_idx=b_req_idx, + b_req_idx=commit.request_indices, ) # 最后一层已经按这个窗口全接受推进了 cursor。这里剪掉没通过验证的尾部。 backend.model.att_state_manager.commit_bounded_att_state( - b_req_idx[b_req_mtp_start_loc], + commit.request_indices[starts.long()], accept_lengths, - role="main", - req_verify_width=_req_verify_width(b_req_mtp_start_loc, b_req_idx.shape[0]), + role=commit.role, + req_verify_width=_req_verify_width(starts, commit.request_indices.shape[0]), ) return accept_lengths, accepted_index @@ -67,7 +69,6 @@ def scatter_mtp_next_tokens( mtp_accept_len: torch.Tensor, # [req_num] ) -> None: """Persist the next MTP proposal and optional scheduling scores by request.""" - schedule_scores = getattr(proposal, "schedule_scores", None) if schedule_scores is not None and 0 in schedule_scores.shape: schedule_scores = None @@ -100,7 +101,6 @@ def record_request_mtp_metrics( verify_run_reqs: List[InferReq], ) -> None: """Accumulate user-visible MTP metrics on each request.""" - if not backend.is_master_in_dp: return diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index c787346a55..79551e8525 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -435,11 +435,6 @@ def probe_varlen(backend_name: str, *, device=None, dtype=None, head_dim=80) -> return _probe(backend_name, varlen=True, device=device, dtype=dtype, head_dim=head_dim) -def validate(backend_name: str, *, device=None, dtype=None, spec=None) -> bool: - """Legacy selector compatibility; success does not imply family-specific numerical validation.""" - return probe_backend(backend_name, device=device, dtype=dtype, spec=spec).ok - - def _probe(backend_name: str, *, varlen=False, device=None, dtype=None, spec=None, head_dim=80) -> ProbeResult: """Isolate execution and retain both its coverage and failure reason.""" try: diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 9d8490b69e..dfa517089e 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -1,11 +1,12 @@ -import os import json -import torch +import os import uuid -from easydict import EasyDict from functools import lru_cache -from lightllm.utils.log_utils import init_logger +import torch +from easydict import EasyDict + +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -221,7 +222,6 @@ def get_decode_attn_autotune_seq_len() -> int: def get_model_init_status(): - global g_model_init_done return g_model_init_done @@ -267,7 +267,7 @@ def enable_triton_mtp_kernel() -> bool: 启用 Triton MTP 解码专用 kernel 通过启动参数 --mtp_step > 0 和 --llm_decode_att_backend=triton 控制 """ - return (get_env_start_args().mtp_step > 0) and ("triton" in get_env_start_args().llm_decode_att_backend) + return (get_env_start_args().mtp_step > 0) and (get_env_start_args().llm_decode_att_backend == "triton") @lru_cache(maxsize=None) diff --git a/lightllm/utils/start_utils.py b/lightllm/utils/start_utils.py index 28a0e5566e..8456106ec9 100644 --- a/lightllm/utils/start_utils.py +++ b/lightllm/utils/start_utils.py @@ -1,15 +1,17 @@ -import os import ctypes +import multiprocessing as mp +import os import resource import signal import subprocess import sys import time -import multiprocessing as mp + import psutil + +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import is_process_active -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.service_shm_cleanup import start_launcher_shm_cleanup_process logger = init_logger(__name__) diff --git a/lightllm/utils/torch_memory_saver_utils.py b/lightllm/utils/torch_memory_saver_utils.py index e0e6cad94a..9468d171d2 100644 --- a/lightllm/utils/torch_memory_saver_utils.py +++ b/lightllm/utils/torch_memory_saver_utils.py @@ -1,13 +1,12 @@ -import torch from contextlib import contextmanager from enum import Enum + +import torch + from lightllm.utils.log_utils import init_logger try: - from torch_memory_saver import ( - torch_memory_saver, - configure_subprocess as tms_configure_subprocess, - ) + from torch_memory_saver import configure_subprocess as tms_configure_subprocess, torch_memory_saver HAS_TORCH_MEMORY_SAVER = True @@ -52,6 +51,8 @@ def __new__(cls, enable_torch_memory_saver: bool = False): class _TorchMemorySaver: + enabled = True + @contextmanager def configure_subprocess(self): # 子进程启动需要放在该上下文里，让 torch_memory_saver 在 worker @@ -80,6 +81,8 @@ def resume(self, tag: MemoryTag): class _TorchMemorySaverFake: + enabled = False + # 未开启 torch_memory_saver 时保持相同接口，保证调用方逻辑完全一致。 @contextmanager def configure_subprocess(self): diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 569d4f99a5..6d3be25ab7 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -14,11 +14,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING import json -from pathlib import Path import socket import sys +from pathlib import Path +from typing import TYPE_CHECKING import numpy as np import torch @@ -35,6 +35,7 @@ if TYPE_CHECKING: from argparse import Namespace from collections.abc import Sequence + from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -309,8 +310,10 @@ def main(): linear_att_hash_page_size=64, linear_att_page_block_num=1, linear_att_cache_size=16, - llm_prefill_att_backend=["fa3", "triton"], - llm_decode_att_backend=["fa3", "triton"], + llm_prefill_att_backend="fa3", + llm_prefill_linear_att_backend="triton", + llm_decode_att_backend="fa3", + llm_decode_linear_att_backend="triton", disable_cudagraph=False, pd_kv_page_size=2048, ) @@ -332,10 +335,11 @@ def main(): config, _ = PretrainedConfig.get_config_dict(args.model_dir) model, _ = get_model(config, kvargs) draft_models = init_mtp_draft_models(args, kvargs, model) - from lightllm.utils.envs_utils import set_model_init_status - from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.basemodel.attention import reserve_attention_resources from lightllm.common.req_manager import ReqManager + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag models = (model, *draft_models) assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" @@ -352,6 +356,9 @@ def main(): assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + reserve_attention_resources(item) manager.load_full_att_calibration(args=model.args) manager.initialize_storage( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py index a1cee4826a..4832abfa61 100644 --- a/test/benchmark/kernels/benchmark_replayssm_service.py +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -7,8 +7,8 @@ import concurrent.futures import json import statistics -import time import threading +import time from pathlib import Path import requests diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 3880029a54..ead1c7e875 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -22,7 +22,6 @@ import torch.multiprocessing as mp from transformers import PretrainedConfig - REPO_ROOT = Path(__file__).resolve().parents[3] if str(REPO_ROOT) not in sys.path: sys.path.append(str(REPO_ROOT)) @@ -34,8 +33,8 @@ from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.config_utils import auto_set_fused_shared_experts, get_dtype, get_vocab_size from lightllm.utils.dist_utils import init_distributed_env -from lightllm.utils.envs_utils import set_env_start_args - +from lightllm.utils.envs_utils import get_unique_server_name, set_env_start_args, set_unique_server_name +from lightllm.utils.service_shm_cleanup import ServiceShmCleanup DEFAULT_BATCH_SIZES = [2, 8, 16, 32, 64, 128] MTP_MODES = {"vanilla_with_att", "eagle_with_att", "vanilla_no_att", "eagle_no_att"} @@ -596,7 +595,6 @@ def _forward_decode_input(self, model_input: ModelInput, allow_overlap: bool) -> def _split_prefill_input(self, model_input: ModelInput): """在 benchmark 调用侧按请求构造两个 prefill microbatch。""" - split_batch = model_input.batch_size // 2 q_lens = model_input.b_seq_len - model_input.b_ready_cache_len split_tokens = int(q_lens[:split_batch].sum().item()) @@ -1171,6 +1169,8 @@ def build_model_kvargs(args: SimpleNamespace, rank_id: int) -> Dict: "disable_cudagraph": args.disable_cudagraph, "llm_prefill_att_backend": args.llm_prefill_att_backend, "llm_decode_att_backend": args.llm_decode_att_backend, + "llm_prefill_linear_att_backend": args.llm_prefill_linear_att_backend, + "llm_decode_linear_att_backend": args.llm_decode_linear_att_backend, "vit_att_backend": args.vit_att_backend, "llm_kv_type": args.llm_kv_type, "llm_kv_quant_group_size": args.llm_kv_quant_group_size, @@ -1233,8 +1233,11 @@ def init_deferred_cudagraph(args: SimpleNamespace, cases: Sequence[BenchmarkCase model.disable_cudagraph = False # Attention backends may size persistent graph buffers during construction. # Rebuild them after replacing the temporary profile-time batch limit. - model._init_att_backend() - model._init_att_backend1() + model.prepare_attention_backends() + from lightllm.common.basemodel.attention import initialize_layer_attention, reserve_attention_resources + + reserve_attention_resources(model) + initialize_layer_attention(model) torch.cuda.empty_cache() model._init_cudagraph() @@ -1245,9 +1248,10 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) cases = [BenchmarkCase(**case) for case in case_dicts] set_env_start_args(args) - from lightllm.distributed import dist_group_manager import torch.distributed as dist + from lightllm.distributed import dist_group_manager + model_kvargs = build_model_kvargs(args, rank_id) defer_cudagraph = ( not args.disable_cudagraph @@ -1270,10 +1274,11 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) torch.cuda.empty_cache() model, _ = get_model(model_cfg, model_kvargs) draft_models = init_mtp_draft_models(args, model_kvargs, model) - from lightllm.utils.envs_utils import set_model_init_status - from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.basemodel.attention import initialize_layer_attention, reserve_attention_resources from lightllm.common.req_manager import ReqManager + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag models = (model, *draft_models) assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" @@ -1290,6 +1295,9 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + reserve_attention_resources(item) manager.load_full_att_calibration(args=model.args) manager.initialize_storage( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction @@ -1309,8 +1317,9 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) batch_size *= draft.mtp_manager.get_decode_tokens_per_request(True) draft.graph_max_batch_size = draft._align_decode_batch_size(batch_size) draft.disable_cudagraph = False - draft._init_att_backend() - draft._init_att_backend1() + draft.prepare_attention_backends() + reserve_attention_resources(draft) + initialize_layer_attention(draft) draft._init_cudagraph() token_source = TokenSource(args) executor = StaticBenchmarkExecutor(args, model, draft_models, token_source) @@ -1642,9 +1651,13 @@ def main(argv: Optional[Sequence[str]] = None): if not cases: raise ValueError("no benchmark cases were generated") args = normalize_args(args, cases) + set_unique_server_name(args) set_env_start_args(args) - run_benchmark(args, cases) + try: + run_benchmark(args, cases) + finally: + ServiceShmCleanup(get_unique_server_name()).cleanup_service_resources() if __name__ == "__main__": diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index cfcf1ad8bf..2caf51b992 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -2,17 +2,17 @@ from __future__ import annotations -from typing import TYPE_CHECKING from dataclasses import replace +from typing import TYPE_CHECKING + +import torch + from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.attention_state.full_att import MLAConfig, MLAFullAttentionState +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator from lightllm.common.attention_state.bounded_att import GDNState -from unit_tests.common.att_resources.test_full_state import full_state +from lightllm.common.attention_state.full_att import MLAConfig, MLAFullAttentionState from unit_tests.common.att_resources.test_bounded_state import make_state - -if TYPE_CHECKING: - import torch +from unit_tests.common.att_resources.test_full_state import full_state class MainDraftTestResourceManager(AttentionStateManager): @@ -25,7 +25,9 @@ def __init__(self, device: str | torch.device | int = "cpu"): mtp_step=self.main_kda_conv_fp32_ssm_fp32_state.mtp_step, device=device, ) - draft_state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + draft_state.initialize_page_pools( + big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu" and torch.cuda.is_available() + ) self.register_state("draft0", draft_state) self.register_state("main", full_state(device=device)) prepared_full_state = MLAFullAttentionState( diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 86e62cb671..029c25392a 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -2,37 +2,38 @@ from __future__ import annotations -from typing import TYPE_CHECKING from contextlib import nullcontext from dataclasses import replace from types import SimpleNamespace +from typing import TYPE_CHECKING import torch -from lightllm.common.attention_state.types import TensorField + from lightllm.common.attention_state.bounded_att import GDNState, KDAState, WindowCompressorState from lightllm.common.attention_state.full_att import ( - FullAttentionState, - GQAAttentionConfig, - MLAConfig, - GroupQuantizedAttentionConfig, - StaticFP8AttentionConfig, - PackedMLAAttentionConfig, CompressedAttentionConfig, FP8E8M0PackedMLAConfig, - StaticFP8FullAttentionState, + FullAttentionState, + GQAAttentionConfig, GQAFullAttentionState, + GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, + MLAConfig, MLAFullAttentionState, + PackedMLAAttentionConfig, PackedMLAFullAttentionState, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, ) - +from lightllm.common.attention_state.types import TensorField if TYPE_CHECKING: - from lightllm.models.qwen3next.config import Qwen3NextResourceConfig - from lightllm.common.attention_state.allocators import FullAttKVAllocator from collections.abc import Mapping, Sequence + + from lightllm.common.attention_state.allocators import FullAttKVAllocator from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.models.qwen3next.config import Qwen3NextResourceConfig def bind_test_fp8_calibration( @@ -259,8 +260,13 @@ def build_hybrid_resources( """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.allocators import StatePageAllocator - from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state - from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, + WindowCompressorStateConfig, + create_gdn_state, + create_kda_state, + create_window_compressor_state, + ) from lightllm.models.glm5_next.config import Glm5NextResourceConfig manager = AttentionStateManager(request_capacity=request_capacity) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index 457fdfa1c3..e7a9b5871b 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -2,15 +2,11 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - -from typing import TYPE_CHECKING +import json from contextlib import nullcontext from dataclasses import asdict -import json from types import SimpleNamespace +from typing import TYPE_CHECKING import pytest import torch @@ -22,11 +18,13 @@ from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, read_kv_storage if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel from pathlib import Path + from lightllm.common.basemodel.basemodel import TpPartBaseModel + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -162,13 +160,35 @@ def mark_backends(model: TpPartBaseModel): ) is state ) - model.prefill_att_backend = model.decode_att_backend = object() - model.prefill_att_backend1 = model.decode_att_backend1 = None + from lightllm.common.attention.api import AttentionFamily + + model.attention_groups = ( + SimpleNamespace( + prefill_spec=SimpleNamespace(family=AttentionFamily.DENSE), + prefill_backend=object(), + decode_backend=object(), + ), + ) model.backends_bound = True + def initialize_layer_attention(model): + assert model.att_state_manager.full_att_allocator is not None + assert all(item.backends_bound for item in models) + model.layer_attention_initialized = True + + def reserve_resources(model): + assert all(item.backends_bound for item in models) + assert model.att_state_manager.full_att_allocator is None + model.attention_resources_reserved = True + + monkeypatch.setattr("lightllm.common.basemodel.attention.reserve_attention_resources", reserve_resources) + monkeypatch.setattr("lightllm.common.basemodel.attention.initialize_layer_attention", initialize_layer_attention) + def warmup(model: TpPartBaseModel): + assert model.layer_attention_initialized assert model.backends_bound assert all(item.req_manager is mapping for item in models) + assert all(item.attention_resources_reserved for item in models) assert all(item.att_state_manager is models[0].att_state_manager for item in models) assert model.att_state_manager.full_att_allocator.size == 32 assert all(item.max_total_token_num == 32 for item in models[: models.index(model) + 1]) @@ -202,8 +222,7 @@ def check_max_len(model: TpPartBaseModel): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) if index: model.main_model, model.mtp_previous_draft_models = models[0], models[1:].copy() - model._init_att_backend = lambda model=model: mark_backends(model) - model._init_att_backend1 = lambda: None + model.prepare_attention_backends = lambda model=model: mark_backends(model) model._init_hidden_collector = lambda: None model._autotune_warmup = lambda model=model: warmup(model) model._init_padded_req = lambda model=model: padded.add(id(model)) @@ -214,6 +233,7 @@ def check_max_len(model: TpPartBaseModel): def memory_info(device: str | torch.device | int): assert all(item.req_manager is mapping for item in models) + assert all(item.attention_resources_reserved for item in models) manager = models[0].att_state_manager states = manager.get_full_attention_states() assert all(read_kv_storage(state) is None for state in states) @@ -273,6 +293,10 @@ def start_runtime(): assert all(not hasattr(other, "att_state_manager") for other in models[index + 1 :]) main = models[0] main.att_state_manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + for item in models: + reserve_resources(item) main.att_state_manager.load_full_att_calibration(args=args) main.att_state_manager.initialize_storage( args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction @@ -280,7 +304,7 @@ def start_runtime(): assert all(item.max_total_token_num is None for item in models) for index, item in enumerate(models): item.init_runtime() - assert all(not getattr(other, "backends_bound", False) for other in models[index + 1 :]) + assert all(not getattr(other, "layer_attention_initialized", False) for other in models[index + 1 :]) assert warmed == padded == captured == prefill_captured == {id(other) for other in models[: index + 1]} start_runtime() diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index e2c08b1461..781a4fadeb 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -1,24 +1,20 @@ from __future__ import annotations -from typing import TYPE_CHECKING from dataclasses import replace +from typing import TYPE_CHECKING import pytest import torch from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, GDNPagesPool, GDNState, KDAPagesPool, KDAState, - ConvSSMStateConfig, -) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, ) +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout, TransferRegion if TYPE_CHECKING: from lightllm.common.attention_state.types import PageKind @@ -51,7 +47,9 @@ def make_state( ) manager = AttentionStateManager(request_capacity=3) manager.register_state("main", state) - state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + state.initialize_page_pools( + big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu" and torch.cuda.is_available() + ) return state diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index b6dff3dc0e..ad1f79b31d 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -1,41 +1,43 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration - -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - -from unit_tests.common.att_resources.model_helpers import rebuild_full_state - -from typing import TYPE_CHECKING from dataclasses import replace from functools import lru_cache, partial from math import prod +from typing import TYPE_CHECKING import pytest import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.utils import deepcopy_result -from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState -from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig -from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator +from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, + GDNState, + KDAState, + WindowCompressorState, + WindowCompressorStateConfig, +) from lightllm.common.attention_state.full_att import ( FullAttentionState, - GQAFullAttentionState, GQAAttentionConfig, - MLAConfig, - create_gqa_state, - create_mla_state, + GQAFullAttentionState, GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, + MLAConfig, MLAFullAttentionState, PackedMLAFullAttentionState, StaticFP8FullAttentionState, + create_gqa_state, + create_mla_state, +) +from lightllm.common.attention_state.utils import deepcopy_result +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.model_helpers import ( + bind_test_fp8_calibration, + full_state_test_buffers, + full_state_test_fields, + read_kv_storage, + rebuild_full_state, ) if TYPE_CHECKING: @@ -146,12 +148,18 @@ def test_budget_matches_token_stores_including_hold_and_padding( reserve = manager.get_pd_operator().get_page_bytes() * args.pd_kv_page_num monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 0) monkeypatch.setattr( "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) for budget, expected in ((actual, 20), (actual - 1, 16), (minimum, 0), (minimum - 1, None)): resources = AttentionStateManager(request_capacity=1) + from types import SimpleNamespace + + resources.device_runtime = SimpleNamespace( + device=torch.device("cpu"), + empty_cache=lambda: None, + memory_info=lambda: (budget + reserve, budget + reserve), + ) resources.register_state("main", rebuild_full_state(template, config=template.config)) monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (budget + reserve, budget + reserve)) resources.finish_state_registration() diff --git a/unit_tests/common/att_resources/test_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py index badc2adf0c..9269571cea 100644 --- a/unit_tests/common/att_resources/test_dsv4_integration.py +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -8,10 +8,12 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState -from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel from lightllm.utils.envs_utils import get_env_start_args +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + @pytest.fixture def cache(monkeypatch): @@ -76,10 +78,15 @@ class TestModel(SimpleNamespace): att_state_role="main", args=SimpleNamespace(run_mode="normal", enable_prefill_microbatch_overlap=True), ) - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4FlashMlaFp8SparseAttBackend + from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - backend = DeepseekV4FlashMlaFp8SparseAttBackend(model=model) - yield manager, window, backend + workspace = DeepseekV4Workspace(model) + workspace.c128_compress_scratch = torch.empty((2, 2048, 512), dtype=torch.float32, device="cuda") + manager._test_workspace = workspace + model.att_state_manager = manager + model.swa_prefill_workspace = DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(model) + yield manager, window, model torch.cuda.synchronize() get_env_start_args.cache_clear() @@ -87,6 +94,7 @@ class TestModel(SimpleNamespace): def make_infer(manager, window): infer = DeepseekV4InferStateInfo() infer.att_state_manager = manager + infer.dsv4_workspace = manager._test_workspace infer.att_state_role = "main" infer.first_logical_layer_id = window.config.logical_layer_ids[0] return infer @@ -101,12 +109,21 @@ def read_packed_rows(buffer, slots, page_tokens): def test_backend_reads_owned_state_views(cache): - manager, window, backend = cache + manager, window, model = cache infer = make_infer(manager, window) assert infer.swa_state is window infer.is_prefill = False for layer in range(3): - swa, compressed = backend.get_att_input_params(infer, layer) + from lightllm.common.attention.api import AttentionFamily + from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + + adapter = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + adapter.layer_num_, adapter.compress_ratio = layer, model.config["compress_ratios"][layer] + adapter.v_head_dim, adapter.softmax_scale = 512, 512**-0.5 + weight = SimpleNamespace(attn_sink_=SimpleNamespace(weight=torch.zeros(64, device="cuda"))) + params = adapter.create_attention_params(model, weight, AttentionFamily.SPARSE_MLA) + swa, compressed = params.swa, params.compressed + assert params.prefill_swa.data_ptr() == model.swa_prefill_workspace[0][layer].data_ptr() assert swa.shape[1] == 74880 assert (compressed is None) == (layer == 2) if layer != 2: @@ -180,16 +197,25 @@ def compress(req_idx, start, end, decode=False): def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer - from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa - manager, window, backend = cache + manager, window, model = cache reference = torch.zeros((17, 74880), dtype=torch.uint8, device="cuda") kv = torch.randn((2048, 512), dtype=torch.bfloat16, device="cuda") weight, angles = torch.randn(512, dtype=torch.bfloat16, device="cuda"), torch.randn((2048, 32), device="cuda") freqs = torch.complex(angles.cos(), angles.sin()) layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) layer.layer_num_, layer.eps_, layer.freqs_cis = 2, 1e-6, freqs + layer.compress_ratio, layer.v_head_dim, layer.softmax_scale = 0, 512, 512**-0.5 + from lightllm.common.attention.api import AttentionFamily + + params = layer.create_attention_params( + model, + SimpleNamespace(attn_sink_=SimpleNamespace(weight=torch.zeros(64, device="cuda"))), + AttentionFamily.SPARSE_MLA, + ) + layer.prefill_attention = SimpleNamespace(params=params) for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): infer = make_infer(manager, window) positions = torch.arange(start, end, device="cuda", dtype=torch.int32) @@ -198,9 +224,11 @@ def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): infer.b_ready_cache_len, infer.b_seq_len = ( torch.tensor([value], dtype=torch.int32, device="cuda") for value in (start, end) ) - prefill = backend.create_att_prefill_state(infer) - prefill.prepare_swa_workspace(infer) - infer.prefill_att_state = prefill + from lightllm.models.deepseek_v4.attention.sparse_mla import prepare_swa_workspace + + buffer, capacity = model.swa_prefill_workspace + prefill = prepare_swa_workspace(window, buffer, capacity, infer, 0) + infer.swa_workspace = prefill indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") lengths, writes, runtime = (torch.empty_like(positions) for _ in range(3)) build_swa_indices( @@ -230,12 +258,13 @@ def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): @pytest.mark.parametrize("draft_kind", ["mtp", "dspark"]) def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): from contextlib import nullcontext - from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel + from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel + from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel - _, _, backend = cache + _, _, model = cache config = dict( - backend.model.config, + model.config, compress_ratios=[4, 128, 0, 0, 0], dspark_target_layer_ids=[0, 1], dspark_block_size=4, @@ -244,6 +273,7 @@ def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): ) args = SimpleNamespace( llm_kv_type="None", + hardware_platform="cuda", page_size=256, mtp_step=3, mtp_mode="vanilla_with_att", diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index ec5a247576..8ef2c41b10 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -1,145 +1,7 @@ -"""FP8 backends obtain calibration from the role's matching storage scheme.""" - -from __future__ import annotations - -from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration - -from unit_tests.common.att_resources.model_helpers import rebuild_full_state -from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups - -from typing import TYPE_CHECKING -from dataclasses import replace -from types import SimpleNamespace -import weakref - import pytest import torch -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.allocators import FullAttKVAllocator -from lightllm.common.basemodel.attention.base_att import AttControl, BaseAttBackend -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3 -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer -from lightllm.common.basemodel.infer_struct import InferStateInfo -from unit_tests.common.att_resources.test_capacity import definition - -if TYPE_CHECKING: - from collections.abc import Sequence - - -@pytest.mark.parametrize("kind", ["static_head", "static_tensor"]) -@pytest.mark.parametrize("role", ["main", "draft0"]) -@pytest.mark.parametrize("prefill", [False, True]) -def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( - monkeypatch: pytest.MonkeyPatch, - kind: str, - role: str, - prefill: bool, -): - template = definition(kind) - states = [] - for value in (2, 5): - scales = template.scales + value - scales[1] += 10 - q_scales = None if kind == "static_tensor" else template.q_scales + value - states.append( - rebuild_full_state( - template=template, - config=replace_logical_layer_groups(template.config, ((3,), (8,))), - token_capacity=8, - device="cpu", - scales=scales, - q_scales=q_scales, - ) - ) - manager = AttentionStateManager(request_capacity=2) - manager.register_state("main", states[0]) - manager.register_state("draft0", states[1]) - manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - # Another encoding in the same role must not supply the FP8 calibration. - plain = definition("gqa") - plain = rebuild_full_state( - template=plain, - config=plain.config, - token_capacity=8, - device="cpu", - ) - manager.register_state(role, plain) - other = rebuild_full_state( - template=template, - config=replace(template.config, logical_layer_groups=((20,), (21,)), global_head_num=4), - token_capacity=8, - device="cpu", - scales=torch.full((2, 8 if kind == "static_head" else 2), 99.0), - q_scales=torch.full((2, 4), 99.0) if kind == "static_head" else None, - ) - manager.register_state(role, other) - manager.finish_state_registration() - selected = getattr(manager, f"{role}_{template.get_state_type()}_state") - k, v = selected.get_layer_view(8) - infer = InferStateInfo() - infer.att_state_manager = manager - infer.is_draft_model = role == "draft0" - infer.att_state_role = role - infer.batch_size, infer.max_q_seq_len = 2, 1 - infer.b_seq_len = torch.ones(2, dtype=torch.int32, device="cpu") - backend = object.__new__(BaseAttBackend) - observed = {} - - if kind == "static_head": - state_cls = fa3.Fp8Fa3PrefillAttState if prefill else fa3.Fp8Fa3DecodeAttState - parent_cls = fa3.Fa3PrefillAttState if prefill else fa3.Fa3DecodeAttState - monkeypatch.setattr(parent_cls, "init_state", lambda self: None) - att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - att.b_att_seq_len = infer.b_seq_len - att.init_state() - - def quantize(q: torch.Tensor, scales: torch.Tensor): - observed["q_scales"] = scales - return q - - def attention(**kwargs: object): - observed.update(kwargs) - return kwargs["q"] - - monkeypatch.setattr(fa3, "q_per_head_static_fp8_quant", quantize) - monkeypatch.setattr(fa3, "flash_attn_with_kvcache", attention) - else: - state_cls = flashinfer.Fp8FlashInferPrefillAttState if prefill else flashinfer.Fp8FlashInferDecodeAttState - att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - wrapper = SimpleNamespace(run=lambda *args, **kwargs: observed.update(kwargs)) - att.prefill_wrapper = att.decode_wrapper = wrapper - - q = torch.zeros((2, selected.config.head_num, selected.config.head_dim), dtype=torch.bfloat16, device="cpu") - - def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.device | int): - return torch.empty(shape, dtype=dtype, device="cpu") - - if prefill: - att.prefill_att( - q=q, - k=k, - v=v, - att_control=AttControl(logical_layer_id=8), - alloc_func=allocate, - ) - else: - att.decode_att( - q=q, - k=k, - v=v, - att_control=AttControl(logical_layer_id=8), - alloc_func=allocate, - ) - if kind == "static_head": - heads = selected.config.head_num - torch.testing.assert_close(observed["q_scales"], selected.q_scales[1]) - torch.testing.assert_close(observed["k_descale"], selected.scales[1, :heads].expand(2, heads)) - torch.testing.assert_close(observed["v_descale"], selected.scales[1, heads:].expand(2, heads)) - else: - assert observed["k_scale"] == selected.cpu_scales[1, 0] - assert observed["v_scale"] == selected.cpu_scales[1, 1] - assert not hasattr(infer, "full_att_state") +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index 0197b17daa..65a417cf49 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -7,11 +7,11 @@ import pytest import torch +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.types import AxisMapping from unit_tests.common.att_resources.model_helpers import full_state_test_buffers -from unit_tests.common.att_resources.test_state_identity import allocated_state from unit_tests.common.att_resources.test_quantized_full_state import transfer_layout -from lightllm.common.attention_state.types import AxisMapping -from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from unit_tests.common.att_resources.test_state_identity import allocated_state @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -375,10 +375,12 @@ def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: cuda_ids = torch.tensor([1], dtype=torch.int64, device="cuda") cpu_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) pd_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + write_error = "KV writes require same-device tensors" if kind == "gqa" else "token indices must reside on CUDA" + copy_error = "token indices must match storage" if kind == "gqa" else "token indices must reside on CUDA" + with pytest.raises(AssertionError, match=write_error): state.write_full_att_kv(state.config.logical_layer_groups[0][0], cpu_ids, torch.empty(0, device="cuda")) for src, dst in ((cpu_ids, cuda_ids), (cuda_ids, cpu_ids)): - with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + with pytest.raises(AssertionError, match=copy_error): state.copy_full_att_tokens(src, dst) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): state.copy_tokens_from_dp(src, dst, state) diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index 1339ebb4b0..c25a024d00 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -1,12 +1,5 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - - from dataclasses import replace import pytest @@ -14,15 +7,16 @@ from lightllm.common.attention_state.full_att import ( FullAttentionState, - GQAFullAttentionState, - MLAFullAttentionState, GQAAttentionConfig, + GQAFullAttentionState, MLAConfig, + MLAFullAttentionState, ) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout, TransferRegion +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + read_kv_storage, ) @@ -97,7 +91,7 @@ def test_full_attention_write_copy_and_pd( data = torch.arange(2 * read_kv_storage(state).shape[2] * 4, dtype=torch.float32, device=device).reshape( 2, read_kv_storage(state).shape[2], 4 ) - if device == "cuda": + if device == "cuda" or cls is GQAFullAttentionState: state.write_full_att_kv(2, ids, data) state.write_full_att_kv(5, ids, data + 100) else: @@ -114,7 +108,7 @@ def test_full_attention_write_copy_and_pd( def reject_host_list(tensor: torch.Tensor): raise AssertionError("token copies must not convert indices to Python lists") - if device == "cuda": + if device == "cuda" or cls is GQAFullAttentionState: with monkeypatch.context() as patch: patch.setattr(torch.Tensor, "tolist", reject_host_list) state.copy_full_att_tokens(ids, torch.tensor([8, 9], dtype=torch.int32, device=device)) diff --git a/unit_tests/common/att_resources/test_gqa_layout_identity.py b/unit_tests/common/att_resources/test_gqa_layout_identity.py index f4c3174f44..4438389315 100644 --- a/unit_tests/common/att_resources/test_gqa_layout_identity.py +++ b/unit_tests/common/att_resources/test_gqa_layout_identity.py @@ -6,7 +6,6 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state -from lightllm.common.basemodel.attention.base_att import AttControl @pytest.mark.parametrize("kv_type", ["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"]) @@ -56,59 +55,3 @@ def test_gqa_layouts_coexist_and_match_across_tp_and_registration_order(kv_type, dst_req_idx=1, ) assert copied == [getattr(target, state.state_name) for state in states] - - -@pytest.mark.parametrize("prefill", [False, True]) -def test_fa3_reads_descale_for_each_head_layout(prefill, monkeypatch): - from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 - - manager = AttentionStateManager(request_capacity=2) - for layer, heads in enumerate((2, 4)): - config = GQAAttentionConfig( - logical_layer_groups=((layer,),), global_head_num=heads, head_dim=128, dtype=torch.bfloat16 - ) - state = create_gqa_state(config, kv_type="fp8kv_sph", device="cpu") - state.scales = torch.arange(1, 2 * heads + 1, dtype=torch.float32).view(1, -1) - state.q_scales = torch.arange(1, heads + 1, dtype=torch.float32).view(1, heads) - manager.register_state("main", state) - manager.finish_state_registration() - # prefill 必须按 b_seq_len 取 batch，不能读这个未设置的请求 batch。 - infer = SimpleNamespace( - att_state_manager=manager, - att_state_role="main", - batch_size=None, - b_seq_len=torch.ones(2, dtype=torch.int32), - max_q_seq_len=1, - ) - cls = fa3_fp8.Fp8Fa3PrefillAttState if prefill else fa3_fp8.Fp8Fa3DecodeAttState - parent = fa3_fp8.Fa3PrefillAttState if prefill else fa3_fp8.Fa3DecodeAttState - monkeypatch.setattr(parent, "init_state", lambda self: None) - monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", lambda q, scales: q) - observed = {} - - def attend(**kwargs): - observed.update(kwargs) - return kwargs["q"] - - monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) - att = cls(backend=None, _infer_state_weak_ref=lambda: infer) - att.b_att_seq_len = torch.ones(3, dtype=torch.int32) - att.page_table = att.cu_seqlens_q = att.cu_seqlens_k = None - att.decode_max_q_seq_len = 1 - att.causal = True - att.init_state() - batch_size = infer.b_seq_len.shape[0] if prefill else att.b_att_seq_len.shape[0] - for state in manager.get_full_attention_states("main"): - heads = state.config.head_num - layer = state.config.logical_layer_groups[0][0] - q = k = v = torch.zeros((2, heads, state.config.head_dim)) - observed.clear() - run = att.prefill_att if prefill else att.decode_att - run(q=q, k=k, v=v, att_control=AttControl(logical_layer_id=layer)) - row = state.scales[0] - k_descale = observed["k_descale"] - v_descale = observed["v_descale"] - assert k_descale.data_ptr() == row[:heads].data_ptr() - assert v_descale.data_ptr() == row[heads:].data_ptr() - torch.testing.assert_close(k_descale, row[:heads].expand(batch_size, heads)) - torch.testing.assert_close(v_descale, row[heads:].expand(batch_size, heads)) diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index fb9ce7ba5b..bce813f5cb 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -1,11 +1,11 @@ import pytest from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.attention_state.bounded_att import BoundedAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.models.qwen3next.model import Qwen3NextTpPartModel -from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel from lightllm.server.core.objs.start_args_type import StartArgs diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 52337199a7..3b2f496fee 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -7,19 +7,22 @@ import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import GQAAttentionConfig, FullAttentionState, GQAFullAttentionState -from lightllm.common.attention_state.bounded_att import KDAState, GDNState -from unit_tests.common.att_resources.test_bounded_state import make_state -from unit_tests.common.att_resources.test_window_compressor_state import make_window -from unit_tests.common.att_resources.test_capacity import definition -from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from lightllm.common.attention_state.bounded_att import GDNState, KDAState +from lightllm.common.attention_state.full_att import FullAttentionState, GQAAttentionConfig, GQAFullAttentionState from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.models.llama.model import LlamaTpPartModel from lightllm.common.basemodel.layer_infer.transformer_layer_infer import TransformerLayerInfer +from lightllm.models.llama.model import LlamaTpPartModel +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from unit_tests.common.att_resources.test_bounded_state import make_state +from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.att_resources.test_window_compressor_state import make_window def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: model = LlamaTpPartModel.__new__(LlamaTpPartModel) + from lightllm.platform import get_device_runtime + + model.device_runtime = get_device_runtime("cpu") model.args, model.config = SimpleNamespace(), {} model.tp_rank_, model.tp_world_size_ = 0, 1 model.layers_infer = [SimpleNamespace(layer_num_=0, indexer=SimpleNamespace())] @@ -65,6 +68,9 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, mo ), ) model = TpPartBaseModel.__new__(TpPartBaseModel) + from lightllm.platform import get_device_runtime + + model.device_runtime = get_device_runtime("cpu") model.is_mtp_draft_model = role != "main" model.mtp_previous_draft_models = [] model.max_req_num = 0 diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index 09e88b4a93..5dcaea3e97 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -1,20 +1,16 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import read_kv_storage - -from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states -from lightllm.common.attention_state.full_att import FullAttentionState -from unit_tests.common.att_resources.model_helpers import owned_states import dataclasses import pytest import torch +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.req_manager import ReqManager from lightllm.models.glm5_next.model import Glm5NextTpPartModel from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel -from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import ( get_added_mtp_kv_layer_num, @@ -23,7 +19,12 @@ set_env_start_args, set_unique_server_name, ) - +from unit_tests.common.att_resources.model_helpers import ( + main_linear_state, + owned_states, + prepare_model_states, + read_kv_storage, +) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index 70970756df..e3e51d6616 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -2,41 +2,38 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import read_kv_storage - -from unit_tests.common.att_resources.model_helpers import rebuild_full_state - -from typing import TYPE_CHECKING import dataclasses import json from functools import partial - -from unit_tests.common.att_resources.model_helpers import prepare_model_states +from typing import TYPE_CHECKING import pytest import torch +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator from lightllm.models.bloom.model import BloomTpPartModel +from lightllm.models.gemma3.model import Gemma3TpPartModel +from lightllm.models.gemma4.model import Gemma4TpPartModel +from lightllm.models.gemma_2b.model import Gemma_2bTpPartModel from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.models.qwen2.model import Qwen2TpPartModel -from lightllm.models.starcoder.model import StarcoderTpPartModel -from lightllm.models.starcoder2.model import Starcoder2TpPartModel from lightllm.models.mistral.model import MistralTpPartModel from lightllm.models.mixtral.model import MixtralTpPartModel -from lightllm.models.gemma_2b.model import Gemma_2bTpPartModel -from lightllm.models.gemma3.model import Gemma3TpPartModel -from lightllm.models.gemma4.model import Gemma4TpPartModel +from lightllm.models.qwen2.model import Qwen2TpPartModel +from lightllm.models.starcoder2.model import Starcoder2TpPartModel +from lightllm.models.starcoder.model import StarcoderTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import prepare_model_states, read_kv_storage, rebuild_full_state from unit_tests.common.att_resources.test_quantized_full_state import KINDS, make_quant_state if TYPE_CHECKING: + from pathlib import Path + from easydict import EasyDict + from lightllm.common.basemodel.basemodel import TpPartBaseModel - from pathlib import Path @pytest.fixture @@ -187,7 +184,6 @@ def memory_info(device: str | torch.device | int): monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 3) monkeypatch.setattr( "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) @@ -198,6 +194,9 @@ def memory_info(device: str | torch.device | int): device="meta", ) resources = AttentionStateManager(request_capacity=1) + from lightllm.platform.cuda import CudaRuntime + + resources.device_runtime = CudaRuntime("cuda:3") resources.register_state("main", rebuild_full_state(template, config=template.config)) resources.finish_state_registration() resources.initialize_storage( diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index ea114d2ba9..81e7c8cdb8 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -2,15 +2,14 @@ from __future__ import annotations - from types import SimpleNamespace import pytest import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import GDNState from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.common.attention_state.bounded_att import GDNState from lightllm.server.router.model_infer.mtp_speculative.utils import verify_mtp_tokens from unit_tests.common.att_resources.test_bounded_state import make_state @@ -82,15 +81,17 @@ def test_commit_without_recurrent_state_and_invalid_role(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("widths", [(3, 3, 3), (3, 1, 2)]) -def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...]): +@pytest.mark.parametrize("overlap", [False, True]) +def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...], overlap): manager = make_manager("cuda") backend = SimpleNamespace( model=SimpleNamespace( att_state_manager=manager, req_manager=SimpleNamespace( + request_generations=torch.zeros(3, dtype=torch.int64, device="cuda"), req_sampling_params_manager=SimpleNamespace( req_to_next_token_ids=torch.tensor([[-1, 10, 11]] * 3, device="cuda") - ) + ), ), ) ) @@ -100,12 +101,30 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i reqs.extend([req] * width) tokens.extend(99 if offset == accepted - 1 else 10 + offset for offset in range(width)) expected_mask.extend(int(offset < accepted) for offset in range(width)) + from lightllm.common.attention_state.commit import BoundedStateCommit + + request_indices = torch.tensor(reqs, dtype=torch.int32, device="cuda") + mtp_indices = torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda") + + def output(start, end): + return SimpleNamespace( + attention_commit=BoundedStateCommit.capture( + "main", + request_indices[start:end], + mtp_indices[start:end], + backend.model.req_manager.request_generations, + ) + ) + + model_output = (output(0, starts[1]), output(starts[1], len(reqs))) if overlap else output(0, len(reqs)) + + request_indices.zero_() + mtp_indices.zero_() accept_lengths, accepted_index = verify_mtp_tokens( backend=backend, next_token_ids=torch.tensor(tokens, device="cuda"), - b_req_idx=torch.tensor(reqs, dtype=torch.int32, device="cuda"), - b_req_mtp_start_loc=torch.tensor(starts, dtype=torch.int32, device="cuda"), - b_mtp_index=torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda"), + model_output=model_output, + num_reqs=len(widths), ) assert accept_lengths.tolist() == [3, 1, 2] assert accepted_index.tolist() == expected_mask diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py index d7d285493b..ad59e31435 100644 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -1,7 +1,5 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import read_kv_storage - from contextlib import nullcontext from types import SimpleNamespace @@ -9,9 +7,10 @@ import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.models.llama.model import LlamaTpPartModel from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.llama.model import LlamaTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.model_helpers import read_kv_storage from unit_tests.common.att_resources.test_capacity import definition from unit_tests.common.basemodel.test_overlap_utils import _make_decode_input, _make_prefill_input @@ -46,8 +45,8 @@ def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.Mo draft, resources = draft_model(mode) draft.prepare_attention_states() resources.finish_state_registration() - draft._init_att_backend() - draft._init_att_backend1() + draft.prepare_attention_backends() + assert draft.attention_groups == () assert draft.att_state_manager is resources assert not hasattr(draft, "full_att_state") assert not resources.get_full_attention_states("draft0") @@ -60,15 +59,21 @@ def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.Mo draft._select_mem_indexes = lambda value: torch.zeros(value.batch_size, dtype=torch.int32) monkeypatch.setattr("lightllm.common.basemodel.basemodel.dist_group_manager.get_group", lambda index: None) model_input = _make_prefill_input([10, 11], req_idx=0, is_decode_req=False) if prefill else _make_decode_input() + draft.attention_groups = () state = draft._create_inferstate(model_input) assert not hasattr(state, "full_att_state") assert state.att_state_manager is resources assert state.is_draft_model - assert state.prefill_att_state is None and state.decode_att_state is None + assert state.attention_groups == () and state.attention_contexts == {} assert state.req_manager is draft.req_manager - state.init_att_state() + from lightllm.common.attention.api import AttentionPhase + from unit_tests.common.attention.helpers import batch_from_fixture + + state.attention_batch = batch_from_fixture(state, AttentionPhase.PREFILL if prefill else AttentionPhase.DECODE) if not prefill: - state.copy_for_cuda_graph(draft._create_inferstate(model_input)) + other = draft._create_inferstate(model_input) + other.attention_batch = state.attention_batch + state.copy_for_cuda_graph(other) @pytest.mark.parametrize("mode", ["vanilla_with_att", "eagle_with_att"]) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index fff9447874..8abcec7f3e 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -22,6 +22,15 @@ full_state_test_fields, read_kv_storage, ) +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) if TYPE_CHECKING: from collections.abc import Sequence @@ -331,7 +340,7 @@ def test_int4_prefill_with_separate_storage(device, monkeypatch): from types import SimpleNamespace from lightllm.common.attention.backends._shared.prefill import context_flashattention_nopad - from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonPrefillAttState + from lightllm.common.attention.backends.dense.triton.int4 import Int4kvTritonAttBackend from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.kv_copy import ppl_int4kv_copy_kv @@ -361,10 +370,8 @@ def test_int4_prefill_with_separate_storage(device, monkeypatch): owner.b_req_idx = torch.tensor([0], dtype=torch.int32, device=device) owner.b_ready_cache_len = owner.b_q_start_loc = torch.zeros(1, dtype=torch.int32, device=device) owner.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, device=device).view(1, -1)) - prefill = Int4kvTritonPrefillAttState( - backend=SimpleNamespace(quant_group_size=8), _infer_state_weak_ref=weakref.ref(owner) - ) - prefill.init_state() + prefill = backend_context(Int4kvTritonAttBackend, "prefill", settings=SimpleNamespace(quant_group_size=8)) + update_context(prefill, owner) if device == "cpu": def dequantize(**kwargs): @@ -379,5 +386,5 @@ def attention(**kwargs): monkeypatch.setattr(ppl_int4kv_copy_kv, "dequantize_int4kv", dequantize) monkeypatch.setattr(context_flashattention_nopad, "context_attention_fwd_contiguous_kv", attention) q = torch.randn(2, 2, 64, device=device) - output = prefill.prefill_att(q, (k, k_scale), (v, v_scale)) + output = run_dense(prefill, owner, q, (k, k_scale), (v, v_scale)) torch.testing.assert_close(output, torch.ones_like(q)) diff --git a/unit_tests/common/att_resources/test_validate_model_cache.py b/unit_tests/common/att_resources/test_validate_model_cache.py index 42508cd0ee..7c3ae3f8b4 100644 --- a/unit_tests/common/att_resources/test_validate_model_cache.py +++ b/unit_tests/common/att_resources/test_validate_model_cache.py @@ -1,3 +1,5 @@ +import importlib.util +from pathlib import Path from types import SimpleNamespace import pytest @@ -5,7 +7,12 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.full_att import create_compressed_att_state from lightllm.models.deepseek_v4.config import DSV4AttentionConfig -from test.att_resources import validate_model_cache + +_spec = importlib.util.spec_from_file_location( + "validate_model_cache", Path(__file__).resolve().parents[3] / "test/att_resources/validate_model_cache.py" +) +validate_model_cache = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(validate_model_cache) @pytest.mark.parametrize("draft_count,bad_role", [(0, None), (2, None), (2, "main"), (2, "draft1")]) diff --git a/unit_tests/common/attention/test_backend_options.py b/unit_tests/common/attention/test_backend_options.py new file mode 100644 index 0000000000..8cef176c66 --- /dev/null +++ b/unit_tests/common/attention/test_backend_options.py @@ -0,0 +1,130 @@ +from dataclasses import fields +from types import SimpleNamespace + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +def test_sparse_backend_options_preserve_flashmla_only(phase, capsys): + parser = make_argument_parser() + option = f"llm_{phase}_att_backend" + choices = next(field.metadata["choices"] for field in fields(StartArgs) if field.name == option) + assert "flashmla_sparse" in choices + assert "flashinfer_sparse" not in choices + args = parser.parse_args([f"--{option}", "flashmla_sparse"]) + assert getattr(args, option) == "flashmla_sparse" + with pytest.raises(SystemExit): + parser.parse_args([f"--{option}", "flashinfer_sparse"]) + assert "invalid choice" in capsys.readouterr().err + + +BACKEND_OPTIONS = ( + "llm_prefill_att_backend", + "llm_decode_att_backend", + "llm_prefill_linear_att_backend", + "llm_decode_linear_att_backend", + "vit_att_backend", +) + + +@pytest.mark.parametrize("option", BACKEND_OPTIONS) +def test_backend_options_have_scalar_defaults_and_matching_choices(option): + parser = make_argument_parser() + field = next(field for field in fields(StartArgs) if field.name == option) + action = next(action for action in parser._actions if action.dest == option) + assert getattr(parser.parse_args([]), option) == getattr(StartArgs(), option) == "auto" + assert action.choices == field.metadata["choices"] + assert action.nargs is None + + +def test_backend_options_are_independent(): + args = make_argument_parser().parse_args( + [ + "--llm_prefill_att_backend", + "fa3", + "--llm_decode_att_backend", + "flashinfer", + "--llm_prefill_linear_att_backend", + "flashqla", + "--llm_decode_linear_att_backend", + "triton", + ] + ) + assert [getattr(args, option) for option in BACKEND_OPTIONS] == ["fa3", "flashinfer", "flashqla", "triton", "auto"] + + +@pytest.mark.parametrize("option", BACKEND_OPTIONS) +def test_backend_options_reject_multiple_values(option, capsys): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([f"--{option}", "triton", "triton"]) + assert "unrecognized arguments" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "option", ["llm_prefill_att_backend", "llm_decode_att_backend", "llm_decode_linear_att_backend"] +) +def test_flashqla_is_only_a_linear_prefill_choice(option, capsys): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([f"--{option}", "flashqla"]) + assert "invalid choice" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "decode,linear_decode,step,expected", + [ + ("triton", "auto", 2, True), + ("fa3", "triton", 2, False), + ("not_triton", "triton", 2, False), + ("triton", "triton", 0, False), + ], +) +def test_mtp_kernel_uses_exact_regular_decode_backend(monkeypatch, decode, linear_decode, step, expected): + from lightllm.utils import envs_utils + + monkeypatch.setattr( + envs_utils, + "get_env_start_args", + lambda: SimpleNamespace( + mtp_step=step, + llm_decode_att_backend=decode, + llm_decode_linear_att_backend=linear_decode, + ), + ) + envs_utils.enable_triton_mtp_kernel.cache_clear() + try: + assert envs_utils.enable_triton_mtp_kernel() is expected + finally: + envs_utils.enable_triton_mtp_kernel.cache_clear() + + +@pytest.mark.parametrize("requested", ["auto", "triton", "flashqla"]) +def test_kda_option_reaches_the_shared_selector(monkeypatch, requested): + from lightllm.common.attention import selection + from lightllm.common.attention.api import AttentionFamily + from lightllm.common.basemodel import attention + from lightllm.models.glm5_next import model as glm5 + from unit_tests.common.attention.helpers import make_spec + + calls = [] + sentinel = object() + + def select(**kwargs): + calls.append(kwargs) + return SimpleNamespace(create=lambda model, spec: sentinel) + + monkeypatch.setattr(selection, "select_attention_backend", select) + monkeypatch.setattr( + attention, "describe_attention_layer", lambda model, layer, family, phase: make_spec(family, phase) + ) + model = glm5.Glm5NextTpPartModel.__new__(glm5.Glm5NextTpPartModel) + from unit_tests.common.attention.helpers import model_stub + + vars(model).update(vars(model_stub(llm_prefill_linear_att_backend=requested))) + model.layers_infer = [SimpleNamespace(get_attention_family=lambda: AttentionFamily.KDA)] + model.prepare_attention_backends() + assert calls[0]["requested"] == requested + assert calls[0]["specs"][0].family is AttentionFamily.KDA + assert model.attention_groups[0].prefill_backend is model.attention_groups[0].decode_backend is sentinel diff --git a/unit_tests/common/attention/test_contracts.py b/unit_tests/common/attention/test_contracts.py index ee97de6458..a30841ee85 100644 --- a/unit_tests/common/attention/test_contracts.py +++ b/unit_tests/common/attention/test_contracts.py @@ -231,6 +231,58 @@ def test_layer_shape_configures_group_planner(): assert (large.tp_q_head_num, large.head_dim) == (16, 128) +def test_groups_follow_each_layers_family_and_shape(monkeypatch): + from lightllm.common.basemodel import attention + + specs = [ + make_spec(num_q_heads=8, head_size=64), + make_spec(num_q_heads=8, head_size=64), + make_spec(num_q_heads=16, head_size=128), + make_spec(AttentionFamily.GDN), + ] + calls = [] + + def select(**kwargs): + calls.append(kwargs) + return SimpleNamespace(create=lambda model, spec: SimpleNamespace()) + + monkeypatch.setattr(selection, "select_attention_backend", select) + monkeypatch.setattr( + attention, "describe_attention_layer", lambda model, layer, family, phase: replace(layer.spec, phase=phase) + ) + model = model_stub() + model.attention_backend_candidates = lambda family, phase: (None, None) + model.layers_infer = [SimpleNamespace(spec=spec, get_attention_family=lambda s=spec: s.family) for spec in specs] + attention.initialize_attention_groups(model) + assert model._attention_layer_groups == (0, 0, 1, 2) + assert len(calls) == 6 + assert [group.prefill_spec.head_size for group in model.attention_groups[:2]] == [64, 128] + assert model.attention_groups[-1].decode_spec.family is AttentionFamily.GDN + + +@pytest.mark.parametrize("invalid", ["input_ids", "group", "role"]) +def test_runner_graph_copy_rejects_before_mutating_batch(invalid): + from lightllm.common.basemodel.infer_struct import InferStateInfo + + captured, incoming = InferStateInfo(), InferStateInfo() + captured.att_state_role = incoming.att_state_role = "main" + group = object() + captured.attention_groups = incoming.attention_groups = (group,) + captured.attention_batch = make_batch() + incoming.attention_batch = make_batch((9, 7), req_to_token_indexs=captured.attention_batch.req_to_token_indexs) + captured.input_ids = torch.zeros(2, dtype=torch.long) + incoming.input_ids = torch.ones(3 if invalid == "input_ids" else 2, dtype=torch.long) + if invalid == "group": + incoming.attention_groups = (object(),) + if invalid == "role": + incoming.att_state_role = "draft" + before = captured.attention_batch.b_seq_len.clone() + with pytest.raises(ValueError): + captured.copy_for_cuda_graph(incoming) + torch.testing.assert_close(captured.attention_batch.b_seq_len, before) + assert captured.input_ids.tolist() == [0, 0] + + def test_replay_flashqla_rejects_decode_and_fp32(): from lightllm.common.attention.backends.gdn.replay.backend import ReplayGDNFlashQlaAttBackend @@ -280,6 +332,43 @@ def test_output_rejects_incompatible_buffer(invalid): layer.backend.forward(layer, context, result, out=out) +def test_batch_builder_borrows_runtime_inputs_and_binding_checks_owner(): + from lightllm.common.basemodel.attention_batch import AttentionBatchBuilder + from lightllm.common.basemodel.infer_struct import InferStateInfo + + state = InferStateInfo() + batch = make_batch() + for name, value in vars(batch).items(): + if hasattr(state, name): + setattr(state, name, value) + state.input_ids = torch.zeros(batch.token_num, dtype=torch.long) + state.is_prefill = False + state.b_image_token_end = torch.tensor([1, 2]) + state.req_manager = SimpleNamespace( + req_to_token_indexs=batch.req_to_token_indexs, HOLD_REQUEST_ID=batch.hold_req_id + ) + borrowed = AttentionBatchBuilder.from_infer_state(state) + assert borrowed.b_req_idx is state.b_req_idx + assert borrowed.b1_cu_q_seq_len is state.b1_cu_q_seq_len + assert borrowed.req_to_token_indexs is state.req_manager.req_to_token_indexs + assert borrowed.b_image_token_end is state.b_image_token_end + backend = _OutputBackend(model_stub(), make_spec()) + context = make_context(backend, borrowed, prepare=True) + state.attention_contexts = {0: context} + layer = make_layer_attention(backend, context, params=None) + assert layer.context(state) is context + output = torch.ones(2, 32) + assert layer.forward(state, output) is output + with pytest.raises(ValueError, match="binding"): + replace(layer, spec=replace(layer.spec, phase=AttentionPhase.PREFILL)).forward(state, output) + with pytest.raises(ValueError, match="binding"): + replace(layer, backend=_OutputBackend(model_stub(), make_spec())).forward(state, output) + with pytest.raises(ValueError, match="binding"): + replace(layer, backend=BaseAttentionBackend(model_stub(), make_spec())).context(state) + with pytest.raises(AttributeError): + context.b_requ_idx = borrowed.b_req_idx + + @pytest.mark.parametrize("family", list(AttentionFamily)) def test_probe_reports_actual_coverage(monkeypatch, family): from lightllm.common.attention import probe, registry diff --git a/unit_tests/common/attention/test_dense_numerics.py b/unit_tests/common/attention/test_dense_numerics.py index fed4977508..9c8eb68bdf 100644 --- a/unit_tests/common/attention/test_dense_numerics.py +++ b/unit_tests/common/attention/test_dense_numerics.py @@ -9,6 +9,8 @@ from lightllm.common.attention.backends._workspace import AttentionWorkspace from lightllm.common.attention.registry import cache_format_to_backend, load_backend_class from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.basemodel.attention_runtime import AttentionRuntime +from lightllm.common.basemodel.graph import GraphExecution from lightllm.platform import get_device_runtime from unit_tests.common.attention.helpers import make_batch, make_layer_attention, make_spec, model_stub @@ -54,6 +56,129 @@ def test_startup_probe_runs_the_selected_dense_implementation(provider, phase): assert success, error +@pytest.mark.parametrize( + "provider,kv_heads,window", + [ + ("triton", 2, None), + ("fa3", 2, None), + ("flashinfer", 2, None), + ("triton", 4, None), + ("triton", 4, (3, 0)), + ("fa3", 4, (3, 0)), + ], +) +@pytest.mark.parametrize("phase", list(AttentionPhase)) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_dense_provider_reference_and_decode_graph(provider, phase, head_dim, kv_heads, window): + cls = load_backend_class(cache_format_to_backend["None"][provider]) + try: + cls.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + torch.manual_seed(616) + model = model_stub(page_size=4) + model.config = dict( + num_attention_heads=4, num_key_value_heads=kv_heads, hidden_size=4 * head_dim, head_dim=head_dim + ) + model.max_seq_length, model.data_type = 16, torch.bfloat16 + spec = make_spec( + phase=phase, + num_q_heads=4, + head_size=head_dim, + value_head_size=head_dim, + cache=CacheLayout("None", 4, torch.bfloat16, kv_heads, head_dim, head_dim), + sliding_window=window, + ) + backend = cls.create(model, spec) + key, value = (torch.randn((32, kv_heads, head_dim), dtype=torch.bfloat16, device="cuda") for _ in range(2)) + q_lengths = (2, 3) if phase is AttentionPhase.PREFILL else (1, 1) + query = torch.randn((sum(q_lengths), 4, head_dim), dtype=torch.bfloat16, device="cuda") + table = torch.arange(32, device="cuda", dtype=torch.int32).view(2, 16) + + def batch_for(lengths): + offsets = [0, q_lengths[0], sum(q_lengths)] + return make_batch( + lengths, + phase=phase, + device="cuda", + token_num=sum(q_lengths), + max_q_seq_len=max(q_lengths), + req_to_token_indexs=table, + hold_req_id=1, + b_q_seq_len=torch.tensor(q_lengths, dtype=torch.int32, device="cuda"), + b1_cu_q_seq_len=torch.tensor(offsets, dtype=torch.int32, device="cuda"), + b_ready_cache_len=torch.tensor( + [n - q for n, q in zip(lengths, q_lengths)], dtype=torch.int32, device="cuda" + ), + ) + + batch = batch_for((7, 13)) + capture = phase is AttentionPhase.DECODE + pool = AttentionWorkspace(runtime=model.device_runtime) + capacity = ContextCapacity(phase, batch.batch_size, batch.token_num, 16, batch.max_q_seq_len) + context = backend.allocate_context(capacity, capture=capture, workspace=pool.fork() if capture else pool) + backend.prepare(context, batch) + attention = replace(make_layer_attention(backend, context, DenseParams(key=key, value=value)), spec=spec) + inputs = DenseInputs(query) + output = torch.empty_like(query) + + def run(): + backend.update(context) + return attention.backend.forward(attention, context, inputs, out=output) + + actual = run() + assert actual.data_ptr() == output.data_ptr() + torch.testing.assert_close(actual, _reference(query, key, value, (7, 13), q_lengths, window), rtol=0.02, atol=0.015) + if not capture: + return + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + captured_pointer = context.batch.b_seq_len.data_ptr() + manager = AttentionStateManager(request_capacity=2) + + def prepare(source): + context.workspace.begin() + backend.prepare(context, source) + + state = SimpleNamespace( + att_state_manager=manager, + attention_groups=(SimpleNamespace(backend_for=lambda phase: backend),), + attention_contexts={0: context}, + copy_for_cuda_graph=prepare, + ) + state.attention_runtime = AttentionRuntime() + state.attention_runtime.groups = state.attention_groups + state.attention_runtime.contexts = state.attention_contexts + execution = GraphExecution(graph, (state,), (output,), runtime=model.device_runtime) + try: + # Eager execution may grow the shared pool after capture. The graph must + # keep its original pointers and fence reuse across streams. + captured_buffers = {key: tensor.data_ptr() for key, tensor in context.workspace._buffers.items()} + for (name, dtype), tensor in tuple(pool._buffers.items()): + pool.get(name, tensor.numel() + 1, dtype) + for i, lengths in enumerate(((16, 3), (2, 15), (9, 8))): + with torch.cuda.stream(stream if i % 2 else torch.cuda.current_stream()): + query.normal_() + actual = execution.replay(batch_for(lengths)) + assert actual is output + assert context.batch.b_seq_len.data_ptr() == captured_pointer + assert { + key: tensor.data_ptr() for key, tensor in context.workspace._buffers.items() + } == captured_buffers + torch.testing.assert_close( + output, _reference(query, key, value, lengths, q_lengths, window), rtol=0.02, atol=0.015 + ) + finally: + manager.close_graphs() + assert execution.graph is None + + @pytest.mark.parametrize("provider", ["fa3", "flashinfer"]) def test_noncausal_prefill_probe(provider): from lightllm.common.attention.probe import probe_dense @@ -77,3 +202,79 @@ def test_triton_rejects_noncausal_prefill(): cls = load_backend_class(cache_format_to_backend["None"]["triton"]) with pytest.raises(UnsupportedAttention, match="non-causal prefill"): cls.check_support(make_spec(phase=AttentionPhase.PREFILL, causal=False)) + + +@pytest.mark.parametrize("phase", list(AttentionPhase)) +@pytest.mark.parametrize("scale", [0.07, 0.13]) +def test_flashinfer_mla_spec_scale(phase, scale): + from lightllm.common.attention.api import AttentionFamily, LayerAttentionBinding + from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAParams, MLAPrefillInputs + from lightllm.common.attention.backends.mla.flashinfer.backend import MlaFlashInferAttBackend + + torch.manual_seed(27) + model = model_stub(page_size=1) + model.config = {"num_attention_heads": 8, "rope_scaling": {"factor": 8, "mscale_all_dim": 1}} + model.qk_nope_head_dim, model.qk_rope_head_dim = 128, 64 + model.kv_lora_rank, model.v_head_dim = 512, 128 + model.max_seq_length, model.data_type = 32, torch.bfloat16 + spec = make_spec( + family=AttentionFamily.MLA, + phase=phase, + num_q_heads=8, + head_size=192, + value_head_size=512, + cache=CacheLayout("None", 4, torch.bfloat16, 1, 192, 512), + softmax_scale=scale, + ) + backend = MlaFlashInferAttBackend.create(model, spec) + assert backend.page_size == 4 and backend.max_page_num == 8 + lengths = (7, 13) + q_lengths = (2, 3) if phase is AttentionPhase.PREFILL else (1, 1) + batch = make_batch( + lengths, + phase=phase, + device="cuda", + token_num=sum(q_lengths), + max_q_seq_len=max(q_lengths), + b_q_seq_len=torch.tensor(q_lengths, dtype=torch.int32, device="cuda"), + b1_cu_q_seq_len=torch.tensor([0, q_lengths[0], sum(q_lengths)], dtype=torch.int32, device="cuda"), + req_to_token_indexs=torch.arange(32, dtype=torch.int32, device="cuda").view(2, 16), + b_ready_cache_len=torch.tensor([n - q for n, q in zip(lengths, q_lengths)], dtype=torch.int32, device="cuda"), + ) + workspace = AttentionWorkspace(runtime=model.device_runtime) + capacity = ContextCapacity(phase, 2, sum(q_lengths), 16, max(q_lengths)) + context = backend.allocate_context(capacity, capture=False, workspace=workspace) + backend.prepare(context, batch) + backend.update(context) + if phase is AttentionPhase.PREFILL: + query = torch.randn(sum(q_lengths), 8, 192, device="cuda", dtype=torch.bfloat16) + key_nope = torch.randn(sum(lengths), 8, 128, device="cuda", dtype=torch.bfloat16) + key_rope = torch.randn(sum(lengths), 1, 64, device="cuda", dtype=torch.bfloat16) + value = torch.randn(sum(lengths), 8, 128, device="cuda", dtype=torch.bfloat16) + key = torch.cat((key_nope, key_rope.expand(-1, 8, -1)), dim=-1) + inputs = MLAPrefillInputs(query, key_nope, key_rope, value) + cache = None + else: + cache = torch.randn(32, 1, 576, device="cuda", dtype=torch.bfloat16) + query = torch.randn(2, 8, 576, device="cuda", dtype=torch.bfloat16) + inputs = MLADecodeInputs(query[..., :512], query[..., 512:]) + key = torch.cat([cache[i * 16 : i * 16 + length].expand(-1, 8, -1) for i, length in enumerate(lengths)]) + value = key[..., :512] + binding = LayerAttentionBinding(0, 0, spec, MLAParams(cache, 512, 64), backend, torch.empty) + actual = binding.backend.forward(binding, context, inputs) + expected, q_start, k_start = [], 0, 0 + for length, q_length in zip(lengths, q_lengths): + q, k, v = ( + query[q_start : q_start + q_length].float(), + key[k_start : k_start + length].float(), + value[k_start : k_start + length].float(), + ) + logits = torch.einsum("qhd,khd->hqk", q, k) * scale + mask = ( + torch.arange(length, device="cuda")[None, :] + <= length - q_length + torch.arange(q_length, device="cuda")[:, None] + ) + expected.append(torch.einsum("hqk,khd->qhd", logits.masked_fill(~mask, -torch.inf).softmax(-1), v)) + q_start += q_length + k_start += length + torch.testing.assert_close(actual, torch.cat(expected).to(actual.dtype), atol=0.03, rtol=0.03) diff --git a/unit_tests/common/attention/test_platform_runtime.py b/unit_tests/common/attention/test_platform_runtime.py index d8f0832d89..57c4310958 100644 --- a/unit_tests/common/attention/test_platform_runtime.py +++ b/unit_tests/common/attention/test_platform_runtime.py @@ -7,6 +7,8 @@ from lightllm.common.attention.backends._workspace import AttentionWorkspace from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.basemodel.attention_runtime import AttentionRuntime +from lightllm.common.basemodel.graph import GraphExecution from lightllm.platform import DeviceRuntime, get_device_runtime @@ -35,6 +37,16 @@ def test_explicit_current_device_is_not_overridden_by_worker_default(monkeypatch assert get_device_runtime().device == torch.device("cuda:0") +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_gqa_storage_allocation_uses_explicit_device_context(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + config = GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=2, head_dim=8, dtype=torch.float32) + with torch.cuda.device(1): + state = GQAFullAttentionState(config, device="cuda", page_size=4) + state.allocate_storage(8) + assert state.k_storage.device == state.v_storage.device == torch.device("cuda:1") + + def test_captured_workspace_keeps_address_when_pool_grows(): pool = AttentionWorkspace("cpu") first, second = pool.fork(), pool.fork() @@ -50,6 +62,131 @@ def test_captured_workspace_keeps_address_when_pool_grows(): assert AttentionWorkspace("cpu").get("scratch", 8, torch.int32).data_ptr() != old.data_ptr() +def test_graph_orders_host_reuse_submission_update_and_destruction(): + events = [] + pending = {"submitted": False, "updated": False} + + class Runtime(DeviceRuntime): + def synchronize(self): + assert not pending["submitted"] or pending["updated"], "waited for a graph before submitting its updates" + events.append("synchronize") + + class Binding(GraphBinding): + def wait_for_update(self): + events.append("host_ready") + + def after_submit(self): + assert pending["submitted"] + pending["updated"] = True + events.append("update") + return () + + def close(self): + events.append("binding_close") + + def submit(): + pending["submitted"] = True + events.append("submit") + + runtime = Runtime("cpu") + manager = AttentionStateManager(request_capacity=2) + backend = SimpleNamespace(bind_graph=lambda context: Binding()) + state = SimpleNamespace( + input_ids=torch.zeros(1), + att_state_manager=manager, + attention_groups=(SimpleNamespace(backend_for=lambda phase: backend),), + attention_contexts={ + 0: SimpleNamespace( + capacity=SimpleNamespace(phase=AttentionPhase.DECODE), workspace=AttentionWorkspace(runtime=runtime) + ) + }, + copy_for_cuda_graph=lambda source: events.append("copy"), + ) + graph = SimpleNamespace(replay=submit, reset=lambda: events.append("graph_reset")) + state.attention_runtime = AttentionRuntime() + state.attention_runtime.groups = state.attention_groups + state.attention_runtime.contexts = state.attention_contexts + execution = GraphExecution(graph, (state,), ("result",), runtime=runtime) + assert execution.replay(state) == "result" + assert events == ["host_ready", "copy", "submit", "update"] + manager.close_graphs() + execution.close() + assert events[-4:] == ["host_ready", "synchronize", "binding_close", "graph_reset"] + assert execution.states == () + with pytest.raises(RuntimeError, match="closed"): + execution.replay(state) + + +def test_storage_generation_invalidates_graph_before_copy(): + manager = AttentionStateManager(request_capacity=2) + state = SimpleNamespace( + input_ids=torch.zeros(1), + att_state_manager=manager, + attention_contexts={}, + attention_groups=(), + copy_for_cuda_graph=lambda source: pytest.fail("copied into a stale graph"), + ) + graph = SimpleNamespace(replay=lambda: pytest.fail("replayed a stale graph"), reset=lambda: None) + state.attention_runtime = AttentionRuntime() + state.attention_runtime.groups = state.attention_groups + state.attention_runtime.contexts = state.attention_contexts + execution = GraphExecution(graph, (state,), runtime=DeviceRuntime("cpu")) + manager.storage_generation += 1 + with pytest.raises(RuntimeError, match="storage changed"): + execution.replay(state) + execution.close() + + +def test_graph_binding_failure_releases_previously_bound_resources(): + events = [] + + class Binding(GraphBinding): + def close(self): + events.append("binding_close") + + def bind_graph(context): + if context.fail: + raise ValueError("unsupported graph task") + return Binding() + + manager = AttentionStateManager(request_capacity=2) + backend = SimpleNamespace(bind_graph=bind_graph) + state = SimpleNamespace( + att_state_manager=manager, + attention_groups=(SimpleNamespace(backend_for=lambda phase: backend),) * 2, + attention_contexts={ + i: SimpleNamespace( + fail=bool(i), capacity=SimpleNamespace(phase=AttentionPhase.DECODE), workspace=AttentionWorkspace("cpu") + ) + for i in range(2) + }, + ) + graph = SimpleNamespace(reset=lambda: events.append("graph_reset")) + state.attention_runtime = AttentionRuntime() + state.attention_runtime.groups = state.attention_groups + state.attention_runtime.contexts = state.attention_contexts + with pytest.raises(ValueError, match="unsupported graph task"): + GraphExecution(graph, (state,), runtime=DeviceRuntime("cpu")) + manager.close_graphs() + assert events == ["binding_close", "graph_reset"] + + +def test_state_layout_and_native_write_preserve_physical_rows(): + config = GQAAttentionConfig(logical_layer_groups=((0, 3), (7,)), global_head_num=2, head_dim=8, dtype=torch.float32) + state = GQAFullAttentionState(config, device="cpu", page_size=4) + state.require_zero_initialization() + state.allocate_storage(12) + assert not state.k_storage.count_nonzero() + values = torch.arange(2 * 4 * 8, dtype=torch.float32).view(2, 4, 8) + state.write_full_att_kv(3, torch.tensor([2, 5], dtype=torch.int32), values) + torch.testing.assert_close(state.get_layer_view(0)[0][[2, 5]], values[:, :2]) + state.copy_full_att_tokens(torch.tensor([2, 5]), torch.tensor([7, 9])) + torch.testing.assert_close(state.get_layer_view(3)[1][[7, 9]], values[:, 2:]) + assert not state.get_layer_view(7)[0].count_nonzero() + relocated = type("ExternalState", (GQAFullAttentionState,), {"__module__": "external.provider"}) + assert relocated(config, device="cpu").describe_layout() == state.describe_layout() + + def test_runtime_registration_keeps_platform_validation_in_adapter(monkeypatch): import lightllm.platform as platform diff --git a/unit_tests/common/attention/test_provider_kind.py b/unit_tests/common/attention/test_provider_kind.py new file mode 100644 index 0000000000..437bf4447e --- /dev/null +++ b/unit_tests/common/attention/test_provider_kind.py @@ -0,0 +1,88 @@ +import pytest +import torch + +from lightllm.common.attention.api import AttentionFamily, UnsupportedAttention +from lightllm.common.attention.registry import load_backend_class +from lightllm.platform import DeviceCaps, Platform +from unit_tests.common.attention.helpers import make_spec + + +@pytest.mark.parametrize( + "reference,family", + [ + ("backends.dense.fa3.backend:Fa3AttBackend", AttentionFamily.DENSE), + ("backends.mla.fa3.backend:MlaFa3AttBackend", AttentionFamily.MLA), + ("backends.sparse_mla.flashmla.backend:NsaFlashMlaSparseAttBackend", AttentionFamily.SPARSE_MLA), + ("backends.sparse_mla.flashmla.fp8:NsaFlashMlaFp8SparseAttBackend", AttentionFamily.SPARSE_MLA), + ("backends.gdn.triton.backend:TritonGDNBackend", AttentionFamily.GDN), + ("backends.gdn.flashqla.backend:FlashQlaGDNBackend", AttentionFamily.GDN), + ("backends.kda.triton.backend:KDALinearAttBackend", AttentionFamily.KDA), + ], +) +def test_family_is_explicit_and_survives_adapter_relocation(reference, family): + backend = load_backend_class(reference) + relocated = type("ExternalAdapter", (backend,), {"__module__": "external.extension"}) + assert relocated.family is family + + +@pytest.mark.parametrize( + "reference", + [ + "backends.sparse_mla.flashmla.fp8:NsaFlashMlaFp8SparseAttBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:Fp8Glm5NextSparseAttBackend", + "lightllm.models.deepseek_v4.attention.sparse_mla:DeepseekV4FlashMlaFp8SparseAttBackend", + ], +) +def test_fp8_sparse_capabilities(monkeypatch, reference): + from dataclasses import replace + + backend = load_backend_class(reference) + caps = DeviceCaps(Platform.CUDA, torch.device("cuda:0"), capability=(8, 0)) + with pytest.raises(UnsupportedAttention, match="device capability"): + backend.check_environment(caps) + caps = replace(caps, capability=(9, 0)) + imported = [] + + def import_provider(name): + imported.append(name) + raise ImportError("missing FlashMLA") + + monkeypatch.setattr("lightllm.common.attention._environment.importlib.import_module", import_provider) + with pytest.raises(UnsupportedAttention, match="missing FlashMLA"): + backend.check_environment(caps) + assert imported == ["flash_mla"] + + spec = make_spec(AttentionFamily.SPARSE_MLA) + spec = replace(spec, cache=replace(spec.cache, format="fp8kv_dsa")) + assert backend.check_support(spec) is None + with pytest.raises(UnsupportedAttention, match="activation dtype"): + backend.check_support(replace(spec, dtype=torch.float32)) + + +@pytest.mark.parametrize( + "reference", + ["backends.dense.fa3.backend:Fa3AttBackend", "backends.dense.fa3.varlen:Fa3VarlenAttentionOp"], +) +def test_environment_rejection_precedes_provider_import(monkeypatch, reference): + from dataclasses import replace + + implementation = load_backend_class(reference) + imported = [] + + def import_provider(name): + imported.append(name) + raise OSError("provider library cannot load") + + monkeypatch.setattr("lightllm.common.attention._environment.importlib.import_module", import_provider) + caps = DeviceCaps(Platform.CPU, torch.device("cpu")) + with pytest.raises(UnsupportedAttention, match="platform cpu"): + implementation.check_environment(caps) + caps = DeviceCaps(Platform.CUDA, torch.device("cuda:0"), available=False, capability=(9, 0)) + with pytest.raises(UnsupportedAttention, match="unavailable"): + implementation.check_environment(caps) + with pytest.raises(UnsupportedAttention, match="device capability"): + implementation.check_environment(replace(caps, available=True, capability=(8, 0))) + assert imported == [] + with pytest.raises(UnsupportedAttention, match="provider library cannot load"): + implementation.check_environment(replace(caps, available=True)) + assert imported == ["sgl_kernel.flash_attn"] diff --git a/unit_tests/common/attention/test_state.py b/unit_tests/common/attention/test_state.py new file mode 100644 index 0000000000..f8a0f8faf1 --- /dev/null +++ b/unit_tests/common/attention/test_state.py @@ -0,0 +1,69 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState +from lightllm.common.attention_state.commit import BoundedStateCommit +from lightllm.common.attention_state.manager import AttentionStateManager + + +def test_commit_owns_mapping_and_detects_recycled_requests(): + requests = torch.tensor([2, 2, 4], dtype=torch.int32) + steps = torch.tensor([0, 1, 0], dtype=torch.int32) + generations = torch.zeros(8, dtype=torch.int64) + output = BoundedStateCommit.capture("main", requests, steps, generations) + requests.fill_(6) + steps.fill_(5) + assert output.request_indices.tolist() == [2, 2, 4] + assert output.mtp_indices.tolist() == [0, 1, 0] + output.validate_generation(generations) + generations[2] += 1 + with pytest.raises(RuntimeError, match="reused request slot"): + output.validate_generation(generations) + + +def test_overlap_commit_preserves_microbatch_order_and_role(): + generation = torch.zeros(8, dtype=torch.int64) + a = BoundedStateCommit.capture("main", torch.tensor([4, 4]), torch.tensor([0, 1]), generation) + b = BoundedStateCommit.capture("main", torch.tensor([2]), torch.tensor([0]), generation) + merged = BoundedStateCommit.concatenate((a, b)) + assert merged.request_indices.tolist() == [4, 4, 2] + assert merged.mtp_indices.tolist() == [0, 1, 0] + with pytest.raises(ValueError, match="same model role"): + BoundedStateCommit.concatenate((a, replace(b, role="draft0"))) + + +@pytest.mark.parametrize("state_cls", [GDNState, KDAState]) +def test_capture_snapshot_restores_only_participating_role_and_requests(state_cls): + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + manager = AttentionStateManager(request_capacity=3) + main = state_cls(config, request_capacity=3, mtp_step=2, device="cpu") + draft = state_cls(config, request_capacity=3, mtp_step=2, device="cpu") + manager.register_state("main", main) + manager.register_state("draft0", draft) + manager.finish_state_registration() + main.conv_state.fill_(2) + main.ssm_state.fill_(3) + main.mtp_state_index.fill_(1) + snapshot = manager.snapshot_capture_state(torch.tensor([1, 1]), role="main") + main.conv_state[:, 1].fill_(8) + main.ssm_state[:, 3:6].fill_(9) + main.mtp_state_index[1] = 2 + main.conv_state[:, 2].fill_(6) + draft.conv_state.fill_(7) + manager.restore_capture_state(snapshot) + assert torch.all(main.conv_state[:, 1] == 2) + assert torch.all(main.ssm_state[:, 3:6] == 3) + assert main.mtp_state_index[1] == 1 + assert torch.all(main.conv_state[:, 2] == 6) + assert torch.all(draft.conv_state == 7) diff --git a/unit_tests/common/basemodel/attention/fa3/test_neo.py b/unit_tests/common/basemodel/attention/fa3/test_neo.py index 9ef88130cf..d335cda04b 100644 --- a/unit_tests/common/basemodel/attention/fa3/test_neo.py +++ b/unit_tests/common/basemodel/attention/fa3/test_neo.py @@ -5,8 +5,17 @@ import pytest import torch +from lightllm.common.attention.backends.dense.fa3 import backend as fp, neo as neo from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.fa3 import fp, neo +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) def _load_neo(monkeypatch, attention_func): @@ -15,7 +24,7 @@ def _load_neo(monkeypatch, attention_func): interface = ModuleType("flash_attn_interface") interface.flash_attn_with_kvcache = attention_func monkeypatch.setitem(sys.modules, "flash_attn_interface", interface) - name = "lightllm.common.basemodel.attention.fa3._neo_import_test" + name = "lightllm.common.attention.backends.dense._neo_import_test" spec = importlib.util.spec_from_file_location(name, neo.__file__) module = importlib.util.module_from_spec(spec) monkeypatch.setitem(sys.modules, name, module) @@ -83,8 +92,8 @@ def indices(values): infer_state.b_image_token_end = image_ends infer_state.req_manager = SimpleNamespace(req_to_token_indexs=req_to_token_indexs) backend = SimpleNamespace(uses_causal_attention=lambda: True, infer_page_size=page_size) - state = neo.NeoFa3AttBackend.create_att_prefill_state(backend, infer_state) - state.init_state() + state = backend_context(neo.NeoFa3AttBackend, "prefill", settings=backend) + update_context(state, infer_state) q = torch.randn(3, 2, 16, device="cuda") k = torch.randn(4 * page_size, 1, 16, device="cuda") @@ -109,11 +118,15 @@ def neo_attention(*, image_token_end, **kwargs): monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", True) monkeypatch.setattr(neo, "flash_attn_with_kvcache_neo", neo_attention) - assert state.prefill_att(q, k, v) is output + assert run_dense(state, infer_state, q, k, v) is output def test_prefill_missing_extension_reports_install_instructions(monkeypatch): monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", False) - state = neo.NeoFa3PrefillAttState() + state = backend_context(neo.NeoFa3AttBackend, "prefill") with pytest.raises(ImportError, match="WANDY666/flash-attention"): - state.prefill_att(None, None, None) + state.backend_identity._forward_prefill( + SimpleNamespace(alloc_tensor=torch.empty, params=SimpleNamespace(key=None, value=None)), + state, + SimpleNamespace(query=None), + ) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py deleted file mode 100644 index 190d8f1162..0000000000 --- a/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py +++ /dev/null @@ -1,100 +0,0 @@ -import sys -import weakref -from types import ModuleType, SimpleNamespace - -import torch - -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState -from lightllm.common.basemodel.attention.flashinfer import mla -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -def test_fp_graph_copy_refreshes_plan_with_runtime_length_and_preserves_buffers(monkeypatch): - plan_calls = [] - flashinfer = ModuleType("flashinfer") - decode = ModuleType("flashinfer.decode") - decode.fast_decode_plan = lambda wrapper, **kwargs: plan_calls.append((wrapper, kwargs)) - flashinfer.decode = decode - monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) - monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) - backend = SimpleNamespace( - infer_page_size=16, - tp_q_head_num=4, - tp_kv_head_num=2, - head_dim=8, - q_data_type=torch.bfloat16, - kv_data_type=torch.bfloat16, - ) - captured_owner = InferStateInfo() - captured_owner.batch_size = 2 - runtime_owner = InferStateInfo() - runtime_owner.max_kv_seq_len = 17 - captured = FlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), - kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), - kv_indices=torch.arange(8, dtype=torch.int32), - kv_last_page_len_buffer=torch.tensor([16, 16], dtype=torch.int32), - decode_wrapper=object(), - ) - runtime = FlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), - kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), - kv_last_page_len_buffer=torch.tensor([1, 5], dtype=torch.int32), - ) - buffer_names = ("kv_starts", "kv_indices", "kv_last_page_len_buffer") - captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} - - captured.copy_for_decode_cuda_graph(runtime) - - assert len(plan_calls) == 1 - wrapper, plan = plan_calls[0] - assert wrapper is captured.decode_wrapper - assert plan["indptr"] is captured.kv_starts - assert plan["indices"] is captured.kv_indices - assert plan["last_page_len"] is captured.kv_last_page_len_buffer - torch.testing.assert_close(plan["global_override_indptr_cpu"], torch.tensor([0, 2, 4], dtype=torch.int32)) - for name in buffer_names: - assert getattr(captured, name).data_ptr() == captured_ptrs[name] - torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) - - -def test_mla_graph_copy_refreshes_plan_with_runtime_lengths_and_preserves_buffers(monkeypatch): - plan_calls = [] - monkeypatch.setattr(mla, "fill_mla_decode_plan_for_cuda_graph", lambda **kwargs: plan_calls.append(kwargs)) - backend = SimpleNamespace(tp_q_head_num=4) - captured_owner = InferStateInfo() - captured_owner.batch_size = 2 - runtime_owner = InferStateInfo() - runtime_owner.b_seq_len = torch.tensor([17, 5], dtype=torch.int32) - runtime_owner.max_kv_seq_len = 17 - captured = mla.MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), - kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), - kv_indices=torch.arange(8, dtype=torch.int32), - decode_wrapper=object(), - ) - runtime = mla.MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), - kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), - ) - buffer_names = ("kv_starts", "kv_indices") - captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} - - captured.copy_for_decode_cuda_graph(runtime) - - assert len(plan_calls) == 1 - plan = plan_calls[0] - assert plan["decode_wrapper"] is captured.decode_wrapper - assert plan["kv_page_indptr"] is captured.kv_starts - assert plan["batch_size"] == 2 - assert plan["max_kv_len"] == 17 - torch.testing.assert_close(plan["kv_lens"], torch.tensor([17, 5], dtype=torch.int32)) - for name in buffer_names: - assert getattr(captured, name).data_ptr() == captured_ptrs[name] - torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py index 1392d5e688..ece27f2475 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py @@ -1,68 +1,45 @@ import sys -import weakref +from dataclasses import replace from types import ModuleType, SimpleNamespace from unittest.mock import Mock import pytest import torch -from lightllm.common.basemodel.attention.flashinfer import fp, fp8, mla -from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper -from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.attention.backends.dense.flashinfer import backend as fp +from lightllm.common.attention.backends.dense.flashinfer.fp8 import Fp8FlashInferAttBackend +from lightllm.common.attention.backends.mla.flashinfer import backend as mla +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context -def test_decode_wrapper_rejects_conflicting_graph_flags(): - infer_state = InferStateInfo() - infer_state.is_cuda_graph_capture = True - infer_state.is_cuda_graph_replay = True - - with pytest.raises(ValueError, match="capture=True, replay=True"): - should_init_decode_wrapper(infer_state) - - -@pytest.fixture( - params=[fp.FlashInferDecodeAttState, fp8.Fp8FlashInferDecodeAttState, mla.MlaFlashInferDecodeAttState], - ids=["fp", "fp8", "mla"], -) -def decode_init(request, monkeypatch): - state_cls = request.param - module = mla if state_cls is mla.MlaFlashInferDecodeAttState else fp - page_size = 1 if state_cls is fp8.Fp8FlashInferDecodeAttState else 4 - max_page_num = 64 // page_size +@pytest.fixture(params=[fp.FlashInferAttBackend, Fp8FlashInferAttBackend, mla.MlaFlashInferAttBackend]) +def decode_case(request, monkeypatch): + cls = request.param + module = mla if cls is mla.MlaFlashInferAttBackend else fp wrapper = SimpleNamespace(plan=Mock()) - create_wrapper = Mock(return_value=wrapper) + create = Mock(return_value=wrapper) flashinfer = ModuleType("flashinfer") - flashinfer.decode = SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=create_wrapper) - flashinfer.mla = SimpleNamespace(BatchMLAPagedAttentionWrapper=create_wrapper) + decode = ModuleType("flashinfer.decode") + decode.BatchDecodeWithPagedKVCacheWrapper = create + decode.fast_decode_plan = Mock() + flashinfer.decode = decode + flashinfer.mla = SimpleNamespace(BatchMLAPagedAttentionWrapper=create) monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) - - def cpu_arange(*args, **kwargs): - kwargs["device"] = "cpu" - return torch.arange(*args, **kwargs) - - monkeypatch.setattr(module, "torch", SimpleNamespace(int32=torch.int32, arange=cpu_arange, empty=torch.empty)) - - def cumsum_params(q_lengths, kv_lengths): - return tuple( - torch.cat((torch.zeros(1, dtype=torch.int32), lengths.cumsum(0, dtype=torch.int32))) - for lengths in (q_lengths, kv_lengths) - ) - - monkeypatch.setattr(module, "gen_cumsum_pad0_tensor", cumsum_params) - - def repack(**kwargs): - kwargs["out_page_indices"].fill_(7) - - repack_indices = Mock(side_effect=repack) - monkeypatch.setattr(module, "repack_kv_index", repack_indices) - can_run = Mock(return_value=True) - backend = SimpleNamespace( - model=SimpleNamespace( - graph=SimpleNamespace(can_run=can_run), graph_max_batch_size=4, graph_max_len_in_batch=64 - ), + monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) + + def cumsum(q, kv): + return tuple(torch.cat((torch.zeros(1, dtype=torch.int32), v.cumsum(0).int())) for v in (q, kv)) + + monkeypatch.setattr(module, "gen_cumsum_pad0_tensor", cumsum) + repack = Mock(side_effect=lambda **kwargs: kwargs["out_page_indices"].fill_(int(kwargs["b_token_len"][0]))) + monkeypatch.setattr(module, "repack_kv_index", repack) + mla_refresh = Mock() + monkeypatch.setattr(mla, "fill_mla_decode_plan_for_cuda_graph", mla_refresh) + page_size = 1 if cls is Fp8FlashInferAttBackend else 4 + backend = make_backend( + cls, infer_page_size=page_size, - max_page_num=max_page_num, - kv_indices_buffer=[torch.empty(4 * max_page_num, dtype=torch.int32)], + max_page_num=64 // page_size, tp_q_head_num=4, tp_kv_head_num=2, head_dim=8, @@ -70,63 +47,54 @@ def repack(**kwargs): qk_rope_head_dim=2, softmax_scale=0.25, q_data_type=torch.bfloat16, - kv_data_type=torch.float8_e4m3fn if page_size == 1 else torch.bfloat16, - workspace_buffer_key="decode_init_test", - workspace_buffer_size=1, - get_gpu_workspace_buffer=Mock(return_value=torch.empty(1, dtype=torch.uint8)), + kv_data_type=torch.bfloat16, + workspace_buffer_key="test", + workspace_buffer_size=8, ) return SimpleNamespace( - module=module, - state=state_cls(backend=backend), backend=backend, + create=create, wrapper=wrapper, - create_wrapper=create_wrapper, - repack_indices=repack_indices, - can_run=can_run, + repack=repack, + refresh=mla_refresh if module is mla else decode.fast_decode_plan, + mla=module is mla, ) -@pytest.mark.parametrize("mode", ["eager", "capture", "replay"]) -def test_decode_init_uses_explicit_graph_mode_for_wrapper_creation(decode_init, mode): - infer_state = InferStateInfo() - infer_state.input_ids = torch.tensor([1, 2], dtype=torch.int64) - infer_state.batch_size = 2 - infer_state.b_seq_len = torch.tensor([5, 2], dtype=torch.int32) - infer_state.max_kv_seq_len = 5 - infer_state.b_req_idx = torch.tensor([2, 3], dtype=torch.int32) - infer_state.microbatch_index = 0 - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.zeros((4, 64), dtype=torch.int32)) - if mode == "capture": - infer_state.is_cuda_graph_capture = True - elif mode == "replay": - infer_state.is_cuda_graph_replay = True - - decode_init.state._infer_state_weak_ref = weakref.ref(infer_state) - decode_init.state.init_state() - - decode_init.can_run.assert_not_called() - expected_starts = [0, 5, 7] if decode_init.backend.infer_page_size == 1 else [0, 2, 3] - torch.testing.assert_close(decode_init.state.kv_starts, torch.tensor(expected_starts, dtype=torch.int32)) - assert decode_init.state.kv_indices.data_ptr() == decode_init.backend.kv_indices_buffer[0].data_ptr() - assert torch.all(decode_init.state.kv_indices == 7) - decode_init.repack_indices.assert_called_once() - repack_args = decode_init.repack_indices.call_args.kwargs - assert repack_args["b_req_idx"] is infer_state.b_req_idx - assert repack_args["b_token_len"] is infer_state.b_seq_len - if decode_init.module is fp: - expected_last_page_lengths = [1, 1] if decode_init.backend.infer_page_size == 1 else [1, 2] - torch.testing.assert_close( - decode_init.state.kv_last_page_len_buffer, torch.tensor(expected_last_page_lengths, dtype=torch.int32) +@pytest.mark.parametrize("capture", [False, True]) +def test_decode_context_creates_wrapper_once_with_explicit_graph_mode(decode_case, capture): + case = decode_case + batch = make_batch() + context = make_context(case.backend, batch, capture=capture, max_length=64, prepare=True) + assert context.decode_wrapper is case.wrapper + case.create.assert_called_once() + assert case.create.call_args.kwargs["use_cuda_graph"] is capture + assert context.kv_starts.tolist() == ([0, 5, 7] if case.backend.infer_page_size == 1 else [0, 2, 3]) + case.wrapper.plan.assert_called_once() + case.repack.assert_called_once() + + +def test_graph_replay_refreshes_existing_wrapper_with_current_lengths(decode_case): + case = decode_case + batch = make_batch() + context = make_context(case.backend, batch, capture=True, max_length=64, prepare=True) + pointers = (context.kv_starts.data_ptr(), context.kv_indices.data_ptr(), context.batch.b_seq_len.data_ptr()) + for lengths in ((17, 5), (1, 1), (60, 3)): + runtime = make_batch(lengths, req_to_token_indexs=batch.req_to_token_indexs) + case.backend.prepare(context, runtime) + assert pointers == ( + context.kv_starts.data_ptr(), + context.kv_indices.data_ptr(), + context.batch.b_seq_len.data_ptr(), ) - else: - torch.testing.assert_close(decode_init.state.q_indptr, torch.tensor([0, 1, 2], dtype=torch.int32)) - - if mode == "replay": - assert decode_init.state.decode_wrapper is None - decode_init.create_wrapper.assert_not_called() - decode_init.backend.get_gpu_workspace_buffer.assert_not_called() - else: - assert decode_init.state.decode_wrapper is decode_init.wrapper - decode_init.create_wrapper.assert_called_once() - decode_init.backend.get_gpu_workspace_buffer.assert_called_once() - decode_init.wrapper.plan.assert_called_once() + assert context.batch.b_seq_len.tolist() == list(lengths) + assert context.kv_indices[0] == lengths[0] + args = case.refresh.call_args.kwargs + if case.mla: + assert args["max_kv_len"] == max(lengths) + assert args["decode_wrapper"] is case.wrapper + else: + pages = (max(lengths) + case.backend.infer_page_size - 1) // case.backend.infer_page_size + assert args["global_override_indptr_cpu"].tolist() == [0, pages, 2 * pages] + case.create.assert_called_once() + assert case.refresh.call_count == 3 diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py index ccf5dfd454..08bd9f2f10 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py @@ -1,15 +1,14 @@ -from types import SimpleNamespace import weakref +from types import SimpleNamespace import pytest import torch flashinfer = pytest.importorskip("flashinfer") -from lightllm.common.basemodel.attention.flashinfer.mla import ( # noqa: E402 - MlaFlashInferDecodeAttState, -) +from lightllm.common.attention.backends.mla.flashinfer.backend import MlaFlashInferAttBackend from lightllm.common.basemodel.infer_struct import InferStateInfo # noqa: E402 +from unit_tests.common.attention.helpers import backend_context def _make_wrapper( @@ -128,9 +127,10 @@ def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, captured_owner = InferStateInfo() captured_owner.batch_size = batch_size captured_owner.b_seq_len = graph_wrapper._kv_len_arr_buf - graph_state = MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), + graph_state = backend_context( + MlaFlashInferAttBackend, + "decode", + settings=backend, kv_starts=graph_wrapper._kv_indptr_buf, kv_indices=graph_wrapper._kv_indices_buf, decode_wrapper=graph_wrapper, @@ -143,7 +143,7 @@ def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, with torch.cuda.graph(graph): graph_wrapper.run(q_nope, q_pe, ckv, kpe, out=graph_out, return_lse=False) - # 复用同一 graph，覆盖页表变化、末页空位及 split/exact 切换后的 plan 更新。 + # 复用同一 graph，覆盖页表变化、末页空位及 split/exact 切换后的 context 更新。 for replay_lengths in (lengths, [1] * batch_size, lengths): replay_lens = torch.tensor(replay_lengths, dtype=torch.int32, device="cuda") replay_indptr = torch.zeros_like(kv_indptr) @@ -152,13 +152,11 @@ def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, runtime_owner.batch_size = batch_size runtime_owner.b_seq_len = replay_lens runtime_owner.max_kv_seq_len = max(replay_lengths) - new_state = MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=replay_indptr, - kv_indices=kv_indices, - ) - graph_state.copy_for_decode_cuda_graph(new_state) + graph_state.kv_starts.copy_(replay_indptr) + graph_state.kv_indices.copy_(kv_indices) + graph_wrapper._kv_len_arr_buf.copy_(replay_lens) + graph_state.batch = runtime_owner + graph_state.backend_identity._refresh_decode_plan(graph_state, replay_lens, max(replay_lengths)) graph.replay() ref_wrapper.plan( q_indptr, diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py index 4c33d8cf8a..9a98ae0bf9 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py @@ -25,7 +25,8 @@ ) def test_nope_sparse_mla_eager_and_graph_replay(num_tokens, num_heads, kv_dtype): wrapper_cls = flashinfer.mla.BatchMLAPagedAttentionWrapper - assert "ckv_scale_arr" in inspect.signature(wrapper_cls.run).parameters, "FlashInfer >= 0.6.18 is required" + if "ckv_scale_arr" not in inspect.signature(wrapper_cls.run).parameters: + pytest.skip("FlashInfer >= 0.6.18 is required") torch.manual_seed(17) dim, num_slots, width = 512, 4096, 2176 diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py b/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py index 924dead570..0ecebf9e67 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py @@ -1,56 +1,27 @@ -from lightllm.common.basemodel.attention import base_att -from lightllm.common.basemodel.attention.base_att import BaseAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend +from types import SimpleNamespace +import torch -def test_workspace_is_shared_by_backend_family_and_device(monkeypatch): - allocations = [] - current_device_id = 0 +from lightllm.common.basemodel.attention import get_attention_workspace +from lightllm.platform import get_device_runtime - def fake_empty(size, dtype, device): - workspace = object() - allocations.append((workspace, size, dtype, device)) - return workspace - monkeypatch.setattr(base_att, "get_current_device_id", lambda: current_device_id) - monkeypatch.setattr(base_att.torch, "empty", fake_empty) - monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) +def test_workspace_isolated_by_model_group_and_slot(): + model = SimpleNamespace(_attention_workspaces={}, device_runtime=get_device_runtime("cpu")) + other = SimpleNamespace(_attention_workspaces={}, device_runtime=get_device_runtime("cpu")) + primary = get_attention_workspace(model, 0, 0, "cpu") + assert primary is get_attention_workspace(model, 0, 0, "cpu") + assert primary is not get_attention_workspace(model, 1, 0, "cpu") + assert primary is not get_attention_workspace(model, 0, 1, "cpu") + assert primary is not get_attention_workspace(other, 0, 0, "cpu") - def get_workspace(backend_cls): - return backend_cls.get_gpu_workspace_buffer( - key_name=backend_cls.workspace_buffer_key, - workspace_size=backend_cls.workspace_buffer_size, - ) - fp_workspace = get_workspace(FlashInferAttBackend) - fp8_workspace = get_workspace(Fp8FlashInferAttBackend) - mla_workspace = get_workspace(MlaFlashInferAttBackend) - - assert fp8_workspace is fp_workspace - assert mla_workspace is not fp_workspace - assert [allocation[1] for allocation in allocations] == [ - 512 * 1024 * 1024, - 256 * 1024 * 1024, - ] - - current_device_id = 1 - other_device_workspace = get_workspace(FlashInferAttBackend) - - assert other_device_workspace is not fp_workspace - assert allocations[-1][3] == 1 - - -def test_workspace_key_includes_size_and_dtype(monkeypatch): - monkeypatch.setattr(base_att, "get_current_device_id", lambda: 0) - monkeypatch.setattr(base_att.torch, "empty", lambda *args, **kwargs: object()) - monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) - - small_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) - large_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 2048) - fp16_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024, dtype=base_att.torch.float16) - - assert BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) is small_workspace - assert large_workspace is not small_workspace - assert fp16_workspace is not small_workspace +def test_workspace_growth_keeps_captured_references_alive(): + model = SimpleNamespace(_attention_workspaces={}, device_runtime=get_device_runtime("cpu")) + workspace = get_attention_workspace(model, 0, 0, "cpu") + captured = workspace.get("indices", 8, torch.int32) + captured.fill_(7) + enlarged = workspace.get("indices", 32, torch.int32) + assert enlarged.data_ptr() != captured.data_ptr() + assert captured.tolist() == [7] * 8 + assert workspace.get("indices", 4, torch.int32).data_ptr() == enlarged.data_ptr() diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 3078cb4c4e..700a6cf218 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -1,230 +1,92 @@ -import inspect import sys -import weakref from types import ModuleType, SimpleNamespace import pytest import torch -import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils -import lightllm.common.basemodel.attention.linear.gdn as gdn -import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops +import lightllm.common.attention.backends.gdn.triton.kernels.chunk as fla_ops import lightllm.utils.backend_validator as backend_validator -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import GDNInputs, GDNParams +from lightllm.common.attention.backends.gdn.flashqla.backend import FlashQlaGDNBackend +from lightllm.common.attention.backends.gdn.triton import backend as gdn +from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNBackend from lightllm.models.qwen3next.config import Qwen3NextResourceConfig -from lightllm.server.api_cli import make_argument_parser +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention @pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float32]) def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): - ssm_states = torch.zeros((1, 1, 2, 2), dtype=cache_dtype) - final_state = torch.ones((1, 1, 2, 2), dtype=torch.float32) - + ssm = torch.zeros((1, 1, 2, 2), dtype=cache_dtype) + final = torch.ones((1, 1, 2, 2), dtype=torch.float32) + q = torch.zeros((1, 1, 1, 1), dtype=cache_dtype) monkeypatch.setattr(gdn, "fused_gdn_gating", lambda _log, a, b, _bias: (a, b)) monkeypatch.setattr(gdn, "causal_conv1d_fn", lambda mixed, *args, **kwargs: mixed) - qkv = torch.zeros((1, 3), dtype=cache_dtype) - q = torch.zeros((1, 1, 1, 1), dtype=cache_dtype) - backend = SimpleNamespace( + backend = make_backend( + TritonGDNBackend, mtp_step=0, activation="silu", - ssm_state_dtype=cache_dtype, _rearrange_mixed_qkv=lambda mixed: (q, q, q), - prefill_kernel=lambda *args, **kwargs: (None, final_state), - ) - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1], dtype=torch.int32) - infer_state.b_ready_cache_len = 0 - state = gdn.LinearAttPrefillAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.b_conv_buffer_idx = torch.tensor([0], dtype=torch.int64) - state.b_ssm_buffer_idx = torch.tensor([0], dtype=torch.int64) - layer_weight = SimpleNamespace( - linear_A_log=SimpleNamespace(weight=None), - linear_dt_bias=SimpleNamespace(weight=None), - linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=None), bias=None), - ) - - state._gdn_prefill_kernel( - qkv, - torch.zeros((1, 3), dtype=cache_dtype), - ssm_states, - torch.zeros((1, 1), dtype=cache_dtype), - torch.zeros((1, 1), dtype=cache_dtype), - state.infer_state, - layer_weight, - ) - - assert ssm_states.dtype == cache_dtype - assert torch.equal(ssm_states, final_state.to(cache_dtype)) - - -def test_decode_state_initializes_normal_layout(): - b_req_idx = torch.tensor([2, 4], dtype=torch.int32) - state, infer_state = _create_decode_state(draft_step=0, dynamic_layout=False, b_req_idx=b_req_idx) - - state.init_state() - - assert state.b_conv_buffer_idx is b_req_idx - assert state.b_ssm_buffer_idx is b_req_idx - assert state.b1_mtp_cu_q_seq_len is None - assert state.b_num_accepted_tokens is None - - -@pytest.mark.parametrize("role", ["main", "draft0"]) -def test_decode_state_initializes_fixed_mtp_layout(role: str): - b_req_idx = torch.tensor([2, 2, 2, 4, 4, 4], dtype=torch.int32) - req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) - state, infer_state = _create_decode_state( - draft_step=2, - dynamic_layout=False, - b_req_idx=b_req_idx, - req_to_mtp_state_index=req_to_mtp_state_index, - role=role, - ) - - state.init_state() - - torch.testing.assert_close(state.b1_mtp_cu_q_seq_len, torch.tensor([0, 3, 6], dtype=torch.int32)) - torch.testing.assert_close(state.b_conv_buffer_idx, torch.tensor([2, 4], dtype=torch.int32)) - torch.testing.assert_close(state.b_num_accepted_tokens, torch.tensor([2, 3], dtype=torch.int32)) - torch.testing.assert_close( - state.b_ssm_buffer_idx, - torch.tensor([[6, 7, 8], [12, 13, 14]], dtype=torch.int32), - ) - - -@pytest.mark.parametrize("role", ["main", "draft0"]) -def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch, role: str): - b_req_idx = torch.tensor([2, 2, 4, 4, 4], dtype=torch.int32) - req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) - expected_cu_q_seq_len = torch.tensor([0, 2, 5, 5, 5, 5], dtype=torch.int32) - expected_conv_buffer_idx = torch.tensor([2, 4, 5, 5, 5], dtype=torch.int32) - expected_num_accepted_tokens = torch.tensor([2, 3, 1, 1, 1], dtype=torch.int32) - build_calls = [] - - def build_params(**kwargs): - build_calls.append(kwargs) - return expected_cu_q_seq_len, expected_conv_buffer_idx, expected_num_accepted_tokens - - monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) - state, infer_state = _create_decode_state( - draft_step=2, - dynamic_layout=True, - b_req_idx=b_req_idx, - req_to_mtp_state_index=req_to_mtp_state_index, - role=role, - ) - - state.init_state() - - assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len - assert state.b_conv_buffer_idx is expected_conv_buffer_idx - assert state.b_num_accepted_tokens is expected_num_accepted_tokens - torch.testing.assert_close( - state.b_ssm_buffer_idx, - torch.tensor( - [[6, 7, 8], [12, 13, 14], [15, 16, 17], [15, 16, 17], [15, 16, 17]], - dtype=torch.int32, - ), - ) - assert len(build_calls) == 1 - assert build_calls[0]["b_req_idx"] is state.infer_state.b_req_idx - assert build_calls[0]["b_mtp_index"] is state.infer_state.b_mtp_index - assert build_calls[0]["req_to_mtp_state_index"] is req_to_mtp_state_index - assert build_calls[0]["hold_req_id"] == 5 - - -def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): - assert inspect.isabstract(gdn.LinearAttBackend) - backend = object.__new__(TritonLinearAttBackend) + prefill_kernel=lambda **kwargs: (q, final), + _split_qkvzba=lambda mixed: (mixed, q, q, q), + ) + batch = make_batch((1,), phase=AttentionPhase.PREFILL) + context = make_context(backend, batch) + backend.update(context) + params = GDNParams(torch.zeros((1, 3)), ssm, None, None, None, None) + attention = make_layer_attention(backend, context, params) + result, z = attention.backend.forward(attention, context, GDNInputs(torch.zeros((1, 3), dtype=cache_dtype))) + assert result is q and z is q + assert ssm.dtype == cache_dtype + torch.testing.assert_close(ssm, final.to(cache_dtype)) + + +@pytest.mark.parametrize("provider", [TritonGDNBackend, FlashQlaGDNBackend]) +def test_decode_and_fixed_mtp_state_routes(provider): + requests = torch.tensor([2, 2, 2, 4, 4, 4], dtype=torch.int32) + cursor = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) + backend = make_backend( + provider, + decode_draft_step=2, + ) + context = make_context(backend, make_batch((3,) * 6, b_req_idx=requests)) + backend.update(context) + assert context.b_conv_buffer_idx.tolist() == [2, 4] + assert context.accepted_tokens(cursor).tolist() == [2, 3] + assert context.b_ssm_buffer_idx.tolist() == [[6, 7, 8], [12, 13, 14]] + backend.decode_draft_step = 0 + context = make_context(backend, make_batch((3, 3), b_req_idx=torch.tensor([2, 4]))) + backend.update(context) + assert context.b_ssm_buffer_idx.tolist() == [2, 4] + + +def test_dynamic_mtp_layout_uses_bound_state_cursor(monkeypatch): + calls = [] + + def build(**kwargs): + calls.append(kwargs) + return torch.tensor([0, 1, 3]), torch.tensor([2, 4]) + + monkeypatch.setattr("lightllm.common.attention.backends._shared.linear.layout.build_dynamic_mtp_layout", build) + cursor = torch.zeros(8, dtype=torch.int32) + backend = make_backend( + TritonGDNBackend, + decode_draft_step=2, + dynamic_verify=True, + ) + batch = make_batch((3, 3, 3), b_req_idx=torch.tensor([2, 4, 4]), b_mtp_index=torch.tensor([0, 0, 1])) + context = make_context(backend, batch) + backend.update(context) + assert context.accepted_tokens(cursor).tolist() == [1, 1] + assert calls[0]["hold_req_id"] == 7 + assert context.b_ssm_buffer_idx.tolist() == [[6, 7, 8], [12, 13, 14]] + + +def test_triton_backend_supplies_chunk_kernel(): + backend = object.__new__(TritonGDNBackend) assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule -@pytest.fixture -def auto_linear_backend_args(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3"], - llm_decode_att_backend=["flashinfer"], - ssm_state_mode="native", - ), - ) - - -def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear_backend_args): - validate_calls = [] - flashqla = ModuleType("flash_qla") - flashqla.chunk_gated_delta_rule = lambda **kwargs: ("flashqla", kwargs) - monkeypatch.setitem(sys.modules, "flash_qla", flashqla) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name: validate_calls.append(name) or True, - ) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - - assert backend_class is FlashQlaLinearAttBackend - backend = object.__new__(backend_class) - assert backend.get_prefill_kernel()(q="q")[0] == "flashqla" - assert validate_calls == ["flashqla"] - - -def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3", "flashqla"], - llm_decode_att_backend=["flashinfer", "triton"], - ssm_state_mode="replay", - ), - ) - from lightllm.common.basemodel.attention.base_att import BaseAttBackend - from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend - from lightllm.common.basemodel.attention.linear.replay_gdn import ( - ReplayGDNFlashQlaAttBackend, - ReplayGDNTritonAttBackend, - ) - - prefill_backend = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - decode_backend = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert prefill_backend is ReplayGDNFlashQlaAttBackend - assert decode_backend is ReplayGDNTritonAttBackend - assert prefill_backend.__bases__ == (BaseAttBackend,) - assert decode_backend.__bases__ == (BaseAttBackend,) - assert not issubclass(prefill_backend, LinearAttBackend) - assert not issubclass(decode_backend, LinearAttBackend) - - -def test_kda_entry_selects_native_or_replay(monkeypatch): - from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend - from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend - - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace(ssm_state_mode="native"), - ) - assert linear_create_utils.get_kda_linear_att_backend_class() is KDALinearAttBackend - - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace(ssm_state_mode="replay"), - ) - assert linear_create_utils.get_kda_linear_att_backend_class() is ReplayKDALinearAttBackend - - def test_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( tp_world_size=2, @@ -251,9 +113,7 @@ def kernel(**kwargs): flashqla = ModuleType("flash_qla") flashqla.chunk_gated_delta_rule = kernel monkeypatch.setitem(sys.modules, "flash_qla", flashqla) - from lightllm.common.attention.backends.gdn.triton.kernels import chunk - - monkeypatch.setattr(chunk, "chunk_gated_delta_rule", kernel) + monkeypatch.setattr(fla_ops, "chunk_gated_delta_rule", kernel) original_randn = torch.randn original_rand = torch.rand @@ -287,154 +147,3 @@ def cpu_tensor(*args, **kwargs): assert kernel_calls[0]["q"].dtype == torch.bfloat16 assert kernel_calls[0]["initial_state"].shape == (1, 4, 8, 16) assert kernel_calls[0]["initial_state"].dtype == torch.float32 - - -def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3", "triton"], - llm_decode_att_backend=["flashinfer", "triton"], - ssm_state_mode="native", - ), - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda *args: pytest.fail("an explicit linear backend must not be auto-validated"), - ) - - prefill_backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - decode_backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert prefill_backend_class is TritonLinearAttBackend - assert decode_backend_class is TritonLinearAttBackend - - -def test_cli_accepts_linear_backend_at_index_one(): - args = make_argument_parser().parse_args( - [ - "--llm_prefill_att_backend", - "fa3", - "flashqla", - "--llm_decode_att_backend", - "flashinfer", - "triton", - ] - ) - - assert args.llm_prefill_att_backend == ["fa3", "flashqla"] - assert args.llm_decode_att_backend == ["flashinfer", "triton"] - - -@pytest.mark.parametrize( - "get_backend_class", - [ - linear_create_utils.get_gdn_linear_prefill_att_backend_class, - linear_create_utils.get_gdn_linear_decode_att_backend_class, - ], -) -def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_backend_class): - with pytest.raises(AssertionError, match="index must be 1"): - get_backend_class(index=0) - - -def test_missing_linear_decode_backend_arg_defaults_to_triton(auto_linear_backend_args): - backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert backend_class is TritonLinearAttBackend - - -def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_args): - validate_calls = [] - monkeypatch.setattr( - linear_create_utils, - "gdn_decode_backends", - {"optimized": FlashQlaLinearAttBackend, "triton": TritonLinearAttBackend}, - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name: validate_calls.append(name) or True, - ) - - backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class( - index=1, priority_list=("optimized", "triton") - ) - - assert backend_class is FlashQlaLinearAttBackend - assert validate_calls == ["optimized"] - - -def test_gdn_prefill_backend_falls_back_when_flashqla_validation_fails(monkeypatch, auto_linear_backend_args): - monkeypatch.setattr(linear_create_utils, "validate", lambda *args: False) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - - assert backend_class is TritonLinearAttBackend - - -def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_backend_args): - validate_calls = [] - monkeypatch.setattr( - linear_create_utils, - "gdn_prefill_backends", - {"flashqla2": TritonLinearAttBackend, "flashqla3": FlashQlaLinearAttBackend}, - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name, *args: validate_calls.append(name) or name == "flashqla3", - ) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class( - index=1, priority_list=("flashqla2", "flashqla3") - ) - - assert backend_class is FlashQlaLinearAttBackend - assert validate_calls == ["flashqla2", "flashqla3"] - - -def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None, role="main"): - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), - ) - infer_state = InferStateInfo() - infer_state.batch_size = b_req_idx.shape[0] - infer_state.b_req_idx = b_req_idx - infer_state.b_mtp_index = torch.zeros_like(b_req_idx) - infer_state.req_manager = SimpleNamespace( - HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, - ) - infer_state.att_state_role = role - if req_to_mtp_state_index is not None: - request_capacity = req_to_mtp_state_index.numel() - resources = AttentionStateManager(request_capacity=request_capacity) - config = ConvSSMStateConfig( - logical_layer_ids=(3,), - global_k_heads=1, - global_v_heads=1, - key_dim=2, - value_dim=2, - conv_width=2, - conv_dtype=torch.float32, - ssm_dtype=torch.float32, - ) - for model_role in ("main", "draft0"): - linear = GDNState(config, request_capacity=request_capacity, mtp_step=draft_step, device="meta") - linear.mtp_state_index = ( - req_to_mtp_state_index if model_role == role else torch.zeros_like(req_to_mtp_state_index) - ) - resources.register_state(model_role, linear) - resources.finish_state_registration() - infer_state.att_state_manager = resources - state = gdn.LinearAttDecodeAttState( - backend=SimpleNamespace( - model=model, - uses_dynamic_spec_verify_layout=lambda: dynamic_layout, - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - return state, infer_state diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py index ae848f13de..21fe062bc2 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -1,17 +1,28 @@ import gc import os -from pathlib import Path import subprocess import sys +from pathlib import Path from types import SimpleNamespace import pytest import torch -pytest.importorskip("flash_mla") +from unit_tests.common.attention.helpers import make_spec + +flashmla = pytest.importorskip("flash_mla") -from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention from lightllm.common.basemodel.layer_weights.meta_weights import TpAttSinkWeight +from lightllm.models.deepseek_v4.attention import sparse_mla as attention +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -35,8 +46,9 @@ def find_spec(self, fullname, path=None, target=None): raise ModuleNotFoundError('vllm deliberately unavailable') sys.meta_path.insert(0, BlockVllm()) -from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention -assert attention.flashmla.__name__ == 'flash_mla' +from lightllm.models.deepseek_v4.attention import sparse_mla as attention +import flash_mla as flashmla +assert flashmla.__name__ == 'flash_mla' assert not any(name == 'vllm' or name.startswith('vllm.') for name in sys.modules) """ root = Path(__file__).resolve().parents[4] @@ -64,11 +76,6 @@ def test_swa_workspace_includes_decode_padding_prefill(monkeypatch: pytest.Monke max_req_num=8, att_state_role="main", ) - monkeypatch.setattr( - backend_class, - "get_gpu_workspace_buffer", - lambda key_name, workspace_size, dtype: torch.empty(workspace_size, dtype=dtype), - ) workspace, capacity = backend_class.get_swa_prefill_workspace(model=model) # D 节点仍执行 HOLD 请求的单 token prefill，但无需保留普通请求的长 prefill 临时区。 assert workspace is not None @@ -96,13 +103,15 @@ def _pool(page_size, pages): def _case(heads, ratio): torch.manual_seed(6014) tokens = 5 - model = type("TestModel", (), {})() + from unit_tests.common.attention.helpers import model_stub + + model = model_stub(page_size=256) model.config = {"num_attention_heads": heads, "compress_ratios": [0, 4, 128]} model.config.update(head_dim=512, qk_rope_head_dim=64, index_head_dim=128, sliding_window=128) model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] model.tp_rank_ = 0 model.tp_world_size_ = 1 - backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend(model=model) + backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend(model=model, spec=make_spec()) q = torch.randn((tokens, backend.padded_q_head_num, 512), dtype=torch.bfloat16, device="cuda") * 0.1 q[:, heads:] = 0 pool, decoded = _pool(128, 2) @@ -120,7 +129,7 @@ def _case(heads, ratio): params = dict( compress_ratio=ratio, head_dim_v=512, - softmax_scale=512 ** -0.5, + softmax_scale=512**-0.5, swa_indices=indices, swa_lengths=lengths, attn_sink=attn_sink.weight, @@ -154,8 +163,8 @@ def _reference(q, decoded, extra_decoded, params, heads): @pytest.mark.parametrize("provide_out", [False, True]) def test_prefill_matches_reference_and_preserves_out(monkeypatch, heads, ratio, provide_out): backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, ratio) - state = backend.create_att_prefill_state(state_info) - state.init_state() + state = backend_context(type(backend), "prefill", settings=backend) + update_context(state, state_info) allocations = [] def alloc_output(shape, **kwargs): @@ -163,22 +172,15 @@ def alloc_output(shape, **kwargs): allocations.append(buffer) return buffer - kernel = attention.flashmla.flash_mla_with_kvcache + kernel = flashmla.flash_mla_with_kvcache def flashmla_with_out(**kwargs): assert kwargs["out"] is not None return kernel(**kwargs) - monkeypatch.setattr(attention.flashmla, "flash_mla_with_kvcache", flashmla_with_out) + monkeypatch.setattr(flashmla, "flash_mla_with_kvcache", flashmla_with_out) out = torch.full_like(q, float("nan")) if provide_out else None - actual = state.prefill_att( - q, - pool, - None, - att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params), - alloc_func=alloc_output, - out=out, - ) + actual = run_sparse(state, state_info, q, pool, params, alloc_func=alloc_output, out=out) assert params["attn_sink"].shape == (backend.padded_q_head_num,) assert torch.count_nonzero(params["attn_sink"][heads:]) == 0 assert len(allocations) == (0 if provide_out else 1) @@ -216,15 +218,26 @@ def test_decode_cuda_graph_reads_updated_query_and_lengths(monkeypatch, heads, r state_info.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) state_info.q = q state_info.packed_kv = packed_kv - state_info.att_control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) - state_info.decode_att_state = backend.create_att_decode_state(state_info) - state_info.decode_att_state1 = None - state_info.decode_att_state.init_state() + state_info.att_control = params + state_info.attention_contexts = {0: backend_context(type(backend), "decode", settings=backend)} + state_info.attention_groups = (SimpleNamespace(backend_for=lambda phase, backend=backend: backend),) + from lightllm.common.basemodel.attention_runtime import AttentionRuntime + + state_info.attention_runtime = AttentionRuntime() + state_info.attention_runtime.groups = state_info.attention_groups + state_info.attention_runtime.contexts = state_info.attention_contexts + update_context(state_info.attention_contexts[0], state_info) + state_info.b_req_idx = torch.arange(q.shape[0], dtype=torch.int32, device=q.device) + from lightllm.common.attention_state import AttentionStateManager + + state_info.att_state_manager = AttentionStateManager(request_capacity=q.shape[0]) + state_info.att_state_manager.finish_state_registration() + state_info.att_state_role = "main" cases.append((state_info, decoded, extra_decoded, params)) def decode_func(*states): outputs = tuple( - state.decode_att_state.decode_att(state.q, state.packed_kv, None, att_control=state.att_control) + run_sparse(state.attention_contexts[0], state, state.q, state.packed_kv, state.att_control) for state in states ) return outputs if overlap else outputs[0] @@ -233,7 +246,7 @@ def decode_func(*states): outputs = runner.capture_decode(decode_func, *(case[0] for case in cases)) outputs = outputs if overlap else (outputs,) for state_info, _, _, _ in cases: - state = state_info.decode_att_state + state = state_info.attention_contexts[0] eager_meta = state.sched_meta_by_ratio[ratio] capture_meta = state.capture_sched_meta_by_ratio[ratio] assert eager_meta.tile_scheduler_metadata.data_ptr() != capture_meta.tile_scheduler_metadata.data_ptr() @@ -248,7 +261,7 @@ def decode_func(*states): params["swa_lengths"].copy_(new_lengths) if ratio: params["extra_lengths"].copy_(new_lengths) - runner.graph[5][0].replay() + runner.graph[5].replay() for actual, (state_info, decoded, extra_decoded, params) in zip(outputs, cases): assert actual.shape == (state_info.q.shape[0], heads, 512) @@ -264,11 +277,9 @@ def test_prefill_out_supports_strided_scratch_buffer(heads): backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, 4) storage = torch.full((q.shape[0], backend.padded_q_head_num, 1024), float("nan"), dtype=q.dtype, device=q.device) out = storage[..., :512] - state = backend.create_att_prefill_state(state_info) - state.init_state() - actual = state.prefill_att( - q, pool, None, out=out, att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params) - ) + state = backend_context(type(backend), "prefill", settings=backend) + update_context(state, state_info) + actual = run_sparse(state, state_info, q, pool, params, out=out) assert actual.data_ptr() == out.data_ptr() assert actual.stride() == out.stride() assert storage[..., 512:].isnan().all() @@ -292,6 +303,9 @@ def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): ) monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) runner = cuda_graph_module.PrefillCudaGraph(decode_cuda_graph=None, tp_world_size=1) + from lightllm.common.attention_state import AttentionStateManager + + manager = AttentionStateManager(request_capacity=5) backend, _, q, packed_kv, decoded, extra_decoded, params = _case(heads, ratio) indices = torch.arange(128, dtype=torch.int32, device=q.device) params["swa_indices"].copy_(indices) @@ -301,17 +315,16 @@ def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): packed_kv = (packed_kv[0], extra) params["extra_indices"].copy_(indices) - monkeypatch.setattr(backend, "get_att_input_params", lambda infer_state, logical_layer_id: packed_kv) - def new_request(length): infer_state = DeepseekV4InferStateInfo() + infer_state.att_state_manager = manager infer_state.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) infer_state.hidden_collector = NoopHiddenCollector() infer_state.attn_params = dict(params, swa_lengths=torch.full_like(params["swa_lengths"], length)) if ratio: infer_state.attn_params["extra_lengths"] = torch.full_like(params["extra_lengths"], length) - infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) - infer_state.prefill_att_state.init_state() + infer_state.attention_contexts = {0: backend_context(type(backend), "prefill", settings=backend)} + update_context(infer_state.attention_contexts[0], infer_state) return infer_state allocations = [] @@ -336,6 +349,18 @@ def compress_and_index(q_lora, x, infer_state, layer_weight): return infer_state.attn_params layer.alloc_tensor = alloc_output + from dataclasses import replace + + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4Params + from unit_tests.common.attention.helpers import make_layer_attention + + attention_params = DeepseekV4Params( + packed_kv[0], packed_kv[0], packed_kv[1], ratio, 512, params["softmax_scale"], params["attn_sink"] + ) + attention_context = backend_context(type(backend), "prefill", settings=backend) + layer.prefill_attention = replace( + make_layer_attention(backend, attention_context, attention_params), alloc_tensor=alloc_output + ) layer._compress_and_index = compress_and_index q_lora = torch.empty((q.shape[0], 4), dtype=q.dtype, device=q.device) x = torch.empty((q.shape[0], 8), dtype=q.dtype, device=q.device) @@ -359,11 +384,11 @@ def prefill_func(inputs, infer_state): atol=5e-4, ) assert len(infer_state.prefill_cuda_graph_exe_list) == 3 - previous_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] + previous_meta = infer_state.attention_contexts[0].sched_meta_by_ratio[ratio] for length in (128, 64, 1, 64): q.mul_(1.25) infer_state = new_request(length) - current_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] + current_meta = infer_state.attention_contexts[0].sched_meta_by_ratio[ratio] assert current_meta.tile_scheduler_metadata is None actual = runner.replay([q, q_lora, x], infer_state)[0] assert callback_states[-1] is infer_state diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py index 9d02784676..312d86e7e7 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -1,20 +1,20 @@ -"""FP8 后端按推理层传入的逻辑层号选择标定，覆盖稀疏层号、共享层和 draft。""" - from dataclasses import replace from types import SimpleNamespace -import weakref import pytest import torch +from lightllm.common.attention.api import AttentionFamily, AttentionPhase +from lightllm.common.attention.backends.dense.fa3 import fp8 as fa3_fp8 +from lightllm.common.attention.backends.dense.flashinfer import fp8 as flashinfer_fp8 from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 -from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.common.basemodel.attention import create_layer_attention_params +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from unit_tests.common.att_resources.model_helpers import rebuild_full_state -from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.model_helpers import rebuild_full_state, replace_logical_layer_groups from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention @pytest.mark.parametrize("backend_name", ["fa3", "flashinfer"]) @@ -42,24 +42,25 @@ def test_fp8_backend_selects_calibration_from_logical_layer( manager = AttentionStateManager(request_capacity=1) manager.register_state(role, state) manager.finish_state_registration() - infer = LlamaInferStateInfo() - infer.att_state_manager = manager - infer.att_state_role = role - infer.b_seq_len = torch.tensor([1], dtype=torch.int32) - infer.max_q_seq_len = 1 + phase = AttentionPhase(operation) + cls = fa3_fp8.Fp8Fa3AttBackend if backend_name == "fa3" else flashinfer_fp8.Fp8FlashInferAttBackend + backend = make_backend(cls) + context = make_context(backend, make_batch((1,), phase=phase)) + if backend_name == "fa3" and operation == "decode": + context.b_att_seq_len = context.batch.b_seq_len + context.decode_max_q_seq_len = 1 + context.causal = True calls = [] - - # 仅替换外部计算内核，保留推理层、公开 attention 接口及 State 标定查询链路。 if backend_name == "fa3": state.q_scales[0].fill_(3) state.q_scales[1].fill_(7) expected_q_scale = 3 if storage_layer_index == 0 else 7 - def quantize(q: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + def quantize(q, scales): assert torch.all(scales == expected_q_scale) return q - def attend(**kwargs: object) -> torch.Tensor: + def attend(**kwargs): assert torch.all(kwargs["k_descale"] == expected_scale) assert torch.all(kwargs["v_descale"] == expected_scale) assert torch.all(kwargs["q_descale"] == expected_q_scale) @@ -68,49 +69,38 @@ def attend(**kwargs: object) -> torch.Tensor: monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", quantize) monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) - state_class = fa3_fp8.Fp8Fa3PrefillAttState if operation == "prefill" else fa3_fp8.Fp8Fa3DecodeAttState - else: - state_class = ( - flashinfer_fp8.Fp8FlashInferPrefillAttState - if operation == "prefill" - else flashinfer_fp8.Fp8FlashInferDecodeAttState - ) - - attention = state_class(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer)) - if backend_name == "fa3": - attention.page_table = attention.cu_seqlens_q = attention.cu_seqlens_k = None - attention.b_att_seq_len = infer.b_seq_len - attention.decode_max_q_seq_len = 1 - attention.causal = True else: + state.cpu_scales = state.scales.cpu().tolist() - def run( - q: torch.Tensor, - paged_kv_cache: tuple[torch.Tensor, torch.Tensor], - k_scale: float, - v_scale: float, - out: torch.Tensor, - ) -> None: + def run(q, paged_kv_cache, k_scale, v_scale, out): assert k_scale == v_scale == expected_scale calls.append(operation) out.copy_(q) - attention.prefill_wrapper = attention.decode_wrapper = SimpleNamespace(run=run) - - infer.prefill_att_state = infer.decode_att_state = attention + setattr(context, operation + "_wrapper", SimpleNamespace(run=run)) layer = LlamaTransformerLayerInfer.__new__(LlamaTransformerLayerInfer) layer.layer_num_ = logical_layer_id layer.tp_q_head_num_ = layer.tp_k_head_num_ = state.config.head_num layer.head_dim_ = state.config.head_dim - layer.alloc_tensor = lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device="cpu") + model = SimpleNamespace( + att_state_manager=manager, + att_state_role=role, + args=StartArgs(llm_kv_type="fp8kv_sph" if backend_name == "fa3" else "fp8kv_spt"), + ) + params = create_layer_attention_params(model, layer, None, AttentionFamily.DENSE) + attention = make_layer_attention(backend, context, params) + attention = replace( + attention, + logical_layer_id=logical_layer_id, + alloc_tensor=lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device="cpu"), + ) + setattr(layer, operation + "_attention", attention) + infer = InferStateInfo() + infer.attention_contexts = {0: context} + monkeypatch.setattr(manager, "get_state", lambda *args: pytest.fail("forward must not look up storage owners")) q = torch.ones((1, state.config.head_num, state.config.head_dim)) if operation == "prefill": - output = layer._context_attention_kernel( - q=q, - kv=None, - infer_state=infer, - layer_weight=None, - ) + output = layer._context_attention_kernel(q=q, kv=None, infer_state=infer, layer_weight=None) else: output = layer._token_attention_kernel(q=q, infer_state=infer, layer_weight=None) assert torch.equal(output, q) diff --git a/unit_tests/common/basemodel/attention/test_fp8_page_size.py b/unit_tests/common/basemodel/attention/test_fp8_page_size.py index 107728c26d..59ccb7c97c 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_page_size.py +++ b/unit_tests/common/basemodel/attention/test_fp8_page_size.py @@ -1,16 +1,13 @@ -from types import SimpleNamespace -from unittest.mock import patch - import pytest -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 -from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend +from lightllm.common.attention.backends.dense.fa3 import fp8 as fa3_fp8 +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend +from lightllm.common.attention.backends.dense.fa3.fp8 import Fp8Fa3AttBackend +from lightllm.common.attention.backends.dense.flashinfer import fp8 as flashinfer_fp8 +from lightllm.common.attention.backends.dense.flashinfer.backend import FlashInferAttBackend +from lightllm.common.attention.backends.dense.flashinfer.fp8 import Fp8FlashInferAttBackend +from lightllm.common.attention.backends.mla.fa3.backend import MlaFa3AttBackend +from lightllm.common.attention.backends.mla.flashinfer.backend import MlaFlashInferAttBackend @pytest.mark.parametrize( @@ -19,9 +16,7 @@ ) def test_default_infer_page_size_matches_model_page_size(backend_class): backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - backend._init_infer_page_size() + backend.page_size = 16 assert backend.infer_page_size == 16 @@ -35,11 +30,6 @@ def test_default_infer_page_size_matches_model_page_size(backend_class): ) def test_fp8_attention_backend_uses_single_token_infer_pages(backend_class, backend_module): backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - with patch.object(backend_module.logger, "warning") as warning: - backend._init_infer_page_size() + backend.page_size = 16 assert backend.infer_page_size == 1 - warning.assert_called_once() - assert "infer_page_size=1" in warning.call_args.args[0] diff --git a/unit_tests/common/basemodel/attention/test_gated_qkv.py b/unit_tests/common/basemodel/attention/test_gated_qkv.py new file mode 100644 index 0000000000..b59325b8f4 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_gated_qkv.py @@ -0,0 +1,23 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen3next.layer_infer import transformer_layer_infer as module + + +def test_qwen_next_qkv_split_matches_projection_with_separate_gate(monkeypatch): + layer = module.Qwen3NextTransformerLayerInfer.__new__(module.Qwen3NextTransformerLayerInfer) + layer.tp_q_head_num_, layer.tp_k_head_num_, layer.tp_v_head_num_ = 4, 2, 2 + layer.head_dim_, layer.embed_dim_, layer.eps_, layer.partial_rotary_factor = 8, 32, 1e-6, 0.25 + layer._tpsp_allgather = lambda input, infer_state: input + projected = torch.arange(3 * 96, dtype=torch.float32).view(3, 96) + weight = SimpleNamespace( + qkvo_gate_proj=SimpleNamespace(mm=lambda x: projected), + qk_norm_weight_=lambda *args, **kwargs: None, + ) + infer = SimpleNamespace(position_cos=None, position_sin=None, need_dp_prefill_balance=False) + monkeypatch.setattr(module, "rotary_emb_fwd", lambda *args, **kwargs: None) + query, kv = layer._get_qkv(torch.zeros(3, 32), infer, weight) + torch.testing.assert_close(query, projected[:, :32]) + torch.testing.assert_close(kv, projected[:, 32:64].view(3, 4, 8)) + torch.testing.assert_close(infer.gate_logics_value, projected[:, 64:]) diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py index 5ecdb1c71f..e85106f754 100644 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -1,225 +1,15 @@ -"""Check import isolation and selection without loading optional GPU dependencies.""" - -import importlib -import importlib.abc -import importlib.util +import os +import subprocess import sys -import unittest -from pathlib import Path -from types import ModuleType, SimpleNamespace -from unittest.mock import Mock, patch - -ROOT = Path(__file__).resolve().parents[4] -ATTENTION = "lightllm.common.basemodel.attention" - - -class BackendImports(importlib.abc.MetaPathFinder, importlib.abc.Loader): - """Reject unselected implementations, and supply lightweight selected backends.""" - - def __init__(self, package=ATTENTION): - self.package = package - self.allowed = {} - self.imported = [] - - def allow(self, suffix, class_name, **attributes): - module = ModuleType(f"{self.package}.{suffix}") - setattr(module, class_name, type(class_name, (), {})) - module.__dict__.update(attributes) - self.allowed[module.__name__] = module - return getattr(module, class_name) - - def find_spec(self, fullname, path=None, target=None): - suffix = fullname.removeprefix(f"{self.package}.") - if not fullname.startswith(f"{self.package}.") or "." not in suffix: - return None - if suffix.split(".")[0] not in {"triton", "fa3", "flashinfer", "nsa", "linear", "sdpa", "xformers"}: - return None - self.imported.append(fullname) - if fullname not in self.allowed: - raise AssertionError(f"Unexpected attention backend import: {fullname}") - return importlib.util.spec_from_loader(fullname, self) - - def create_module(self, spec): - return self.allowed[spec.name] - - def exec_module(self, module): - pass - - -class TestLazyAttentionImports(unittest.TestCase): - def setUp(self): - # Isolate attention from basemodel's unrelated GPU initialization. - modules_patch = patch.dict(sys.modules) - modules_patch.start() - self.addCleanup(modules_patch.stop) - for name in list(sys.modules): - if name == ATTENTION or name.startswith(f"{ATTENTION}."): - del sys.modules[name] - for name in ("lightllm", "lightllm.common", "lightllm.common.basemodel", "lightllm.utils"): - module = ModuleType(name) - module.__path__ = [str(ROOT.joinpath(*name.split(".")))] - sys.modules[name] = module - - self.args = SimpleNamespace( - llm_kv_type="None", - llm_prefill_att_backend=["triton"], - llm_decode_att_backend=["triton"], - enable_ep_moe=False, - mtp_step=0, - ssm_state_mode="native", - ) - self.validate = Mock(return_value=True) - dependencies = { - f"{ATTENTION}.base_att": { - name: type(name, (), {}) - for name in ("BaseAttBackend", "BasePrefillAttState", "BaseDecodeAttState", "AttControl") - }, - "lightllm.utils.envs_utils": {"get_env_start_args": lambda: self.args}, - "lightllm.utils.backend_validator": {"validate": self.validate}, - "lightllm.utils.log_utils": {"init_logger": lambda name: Mock()}, - } - for name, attributes in dependencies.items(): - module = ModuleType(name) - module.__dict__.update(attributes) - sys.modules[name] = module - - self.imports = BackendImports() - sys.meta_path.insert(0, self.imports) - self.addCleanup(sys.meta_path.remove, self.imports) - self.attention = importlib.import_module(ATTENTION) - - def test_importing_selection_apis_loads_no_implementations(self): - importlib.import_module(f"{ATTENTION}.linear") - importlib.import_module(f"{ATTENTION}.nsa") - self.assertEqual(self.imports.imported, []) - - def test_explicit_regular_selection_skips_neo(self): - backend = self.imports.allow("fa3.fp", "Fa3AttBackend") - self.args.llm_prefill_att_backend = ["fa3"] - - self.assertIs(self.attention.get_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.fp"]) - self.validate.assert_not_called() - - def test_auto_selection_does_not_import_failed_candidates(self): - backend = self.imports.allow("flashinfer.fp", "FlashInferAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - self.validate.side_effect = lambda name: name == "flashinfer" - - self.assertIs(self.attention.get_prefill_att_backend_class(), backend) - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "flashinfer"]) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.flashinfer.fp"]) - - def test_mtp_decode_preserves_fa3_priority(self): - backend = self.imports.allow("fa3.fp", "Fa3AttBackend") - self.args.llm_decode_att_backend = ["auto"] - self.args.mtp_step = 2 - - self.assertIs(self.attention.get_decode_att_backend_class(), backend) - self.validate.assert_called_once_with("fa3") - - def test_expert_parallelism_still_excludes_flashinfer(self): - backend = self.imports.allow("triton.fp", "TritonAttBackend") - self.args.llm_decode_att_backend = ["auto"] - self.args.enable_ep_moe = True - self.validate.side_effect = lambda name: name == "triton" - self.assertIs(self.attention.get_decode_att_backend_class(), backend) - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "triton"]) - def test_auto_selection_falls_back_to_triton(self): - backend = self.imports.allow("triton.mla", "MlaTritonAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - self.validate.return_value = False - - self.assertIs(self.attention.get_mla_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.mla"]) - - def test_nsa_selection_does_not_import_other_kv_types(self): - backend = self.imports.allow("nsa.fp8_flashmla_sparse", "NsaFlashMlaFp8SparseAttBackend") - self.args.llm_kv_type = "fp8kv_dsa" - self.args.llm_decode_att_backend = ["flashmla_sparse"] - - self.assertIs(self.attention.get_nsa_decode_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.fp8_flashmla_sparse"]) - - def test_custom_backend_map_uses_import_paths(self): - backend = self.imports.allow("nsa.custom", "CustomBackend") - backend_map = {"None": {"custom": "nsa.custom:CustomBackend"}} - self.args.llm_prefill_att_backend = ["custom"] - self.args.llm_decode_att_backend = ["auto"] - - self.assertIs(self.attention.get_nsa_prefill_att_backend_class(backend_map=backend_map), backend) - self.assertIs( - self.attention.get_nsa_decode_att_backend_class(priority_list=["custom"], backend_map=backend_map), - backend, - ) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.custom"]) - - def test_neo_triton_selection_does_not_probe_fa3(self): - backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.args.llm_prefill_att_backend = ["auto"] - self.assertIs(self.attention.get_neo_prefill_att_backend_class(priority_list=["triton"]), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.neo"]) - - def test_neo_auto_selection_still_falls_back_without_image_token_interface(self): - self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=False) - backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.validate.assert_called_once_with("triton") - self.args.llm_prefill_att_backend = ["fa3"] - with self.assertRaises(KeyError): - self.attention.get_neo_prefill_att_backend_class() - - def test_neo_auto_selection_uses_supported_fa3(self): - backend = self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=True) - self.args.llm_prefill_att_backend = ["auto"] - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) - self.validate.assert_called_once_with("fa3") - - def _allow_linear_backend_classes(self): - flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") - triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") - self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") - replay_gdn = self.imports.allowed[f"{ATTENTION}.linear.replay_gdn"] - replay_gdn.ReplayGDNFlashQlaAttBackend = type("ReplayGDNFlashQlaAttBackend", (), {}) - self.imports.allow("linear.kda", "KDALinearAttBackend") - self.imports.allow("linear.replay_kda", "ReplayKDALinearAttBackend") - return flashqla, triton - - def test_linear_prefill_auto_selects_validated_flashqla(self): - flashqla, _triton = self._allow_linear_backend_classes() - selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), flashqla) - self.validate.assert_called_once_with("flashqla") - - def test_linear_selection_falls_back_to_triton(self): - _flashqla, triton = self._allow_linear_backend_classes() - selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - self.validate.return_value = False - - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) - self.assertIs(selection.get_gdn_linear_decode_att_backend_class(), triton) - self.args.llm_prefill_att_backend.append("triton") - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) - - def test_vit_selection_loads_only_requested_implementation(self): - from lightllm.common.attention import registry, selection - - loaded = [] - sentinel = object() - - def import_module(name): - loaded.append(name) - return SimpleNamespace(SdpaVarlenAttentionOp=sentinel) - - with patch.object(registry.importlib, "import_module", import_module): - assert registry.get_varlen_attention_op_class("sdpa") is sentinel - assert loaded == ["lightllm.common.attention.backends.dense.sdpa.varlen"] +def test_api_and_selection_do_not_import_provider_extensions(): + code = """ +import sys +import lightllm.common.attention.api +import lightllm.common.attention.selection +assert not any(name.startswith(('flashinfer', 'flash_mla', 'flash_qla')) for name in sys.modules) +assert not any(name.startswith('lightllm.common.attention.backends.') for name in sys.modules) +""" + result = subprocess.run([sys.executable, "-c", code], env=os.environ.copy(), capture_output=True, text=True) + assert result.returncode == 0, result.stderr diff --git a/unit_tests/common/basemodel/attention/test_state_lifetime.py b/unit_tests/common/basemodel/attention/test_state_lifetime.py index 9e049d54b4..452d39f4b5 100644 --- a/unit_tests/common/basemodel/attention/test_state_lifetime.py +++ b/unit_tests/common/basemodel/attention/test_state_lifetime.py @@ -1,174 +1,33 @@ import gc import weakref -import pytest -import torch +from lightllm.common.attention.backends.base import BaseAttentionBackend +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_spec, model_stub -from lightllm.common.basemodel.attention.base_att import BaseDecodeAttState, BasePrefillAttState -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.basemodel.infer_struct import InferStateInfo - -class _PrefillState(BasePrefillAttState): - def init_state(self): - self.positions = self.infer_state.position_ids - - def prefill_att(self, q, k, v, **kwargs): - return q - - -class _DecodeState(BaseDecodeAttState): - def init_state(self): - self.positions = self.infer_state.position_ids - - def decode_att(self, q, k, v, **kwargs): - return q - - -@pytest.fixture -def no_cycle_gc(): - enabled = gc.isenabled() - gc.collect() +def test_context_does_not_retain_runner_or_model(): + model = model_stub() + backend = BaseAttentionBackend(model, make_spec()) + assert not any(value is model for value in vars(backend).values()) + context = make_context(backend) gc.disable() try: - yield + batch_ref = weakref.ref(context.batch) + del context + assert batch_ref() is None finally: - gc.collect() - if enabled: - gc.enable() - - -def _create_state(is_prefill, secondary=False): - owner = InferStateInfo() - owner.is_prefill = is_prefill - owner.position_ids = torch.arange(4) - state_cls = _PrefillState if is_prefill else _DecodeState - slot = "prefill_att_state" if is_prefill else "decode_att_state" - setattr(owner, slot, state_cls(_infer_state_weak_ref=weakref.ref(owner))) - if secondary: - setattr(owner, slot + "1", state_cls(_infer_state_weak_ref=weakref.ref(owner))) - owner.init_att_state() - return owner - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("secondary", [False, True]) -def test_forward_states_release_without_cycle_gc(no_cycle_gc, is_prefill, secondary): - owner = _create_state(is_prefill, secondary) - states = [state for state in vars(owner).values() if isinstance(state, (BasePrefillAttState, BaseDecodeAttState))] - owner_ref = weakref.ref(owner) - tensor_ref = weakref.ref(owner.position_ids) - state_refs = [weakref.ref(state) for state in states] - - # Even an external reference to attention metadata must not retain its owner. - del owner - assert owner_ref() is None - with pytest.raises(ReferenceError): - states[0].infer_state - assert tensor_ref() is not None - del states - assert all(ref() is None for ref in state_refs) - assert tensor_ref() is None - - -@pytest.mark.parametrize("state_cls", [_PrefillState, _DecodeState]) -def test_attention_getter_uses_a_weak_reference_and_is_read_only(state_cls): - state = state_cls() - assert state.infer_state is None - - owner = InferStateInfo() - state = state_cls(_infer_state_weak_ref=weakref.ref(owner)) - assert isinstance(state._infer_state_weak_ref, weakref.ReferenceType) - assert state.infer_state is owner - state.check_no_infer_state() - with pytest.raises(AttributeError): - state.infer_state = owner - + gc.enable() -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("slot_suffix", ["", "1"]) -def test_check_rejects_infer_state_under_any_attribute_name(is_prefill, slot_suffix): - owner = _create_state(is_prefill, secondary=True) - slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix - attention = getattr(owner, slot) - - class _OtherInferState(InferStateInfo): - pass - - # Check the type, rather than just the identity of the expected owner. - attention.saved_context = _OtherInferState() - with pytest.raises(RuntimeError, match="saved_context"): - owner.check_att_states() - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("slot_suffix", ["", "1"]) -def test_init_checks_for_accidental_back_reference(is_prefill, slot_suffix): - owner = _create_state(is_prefill, secondary=True) - slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix - base = _PrefillState if is_prefill else _DecodeState - - class _RetainingState(base): - def init_state(self): - self.saved_context = self.infer_state - - setattr(owner, slot, _RetainingState(_infer_state_weak_ref=weakref.ref(owner))) - try: - with pytest.raises(RuntimeError, match="saved_context"): - owner.init_att_state() - finally: - del getattr(owner, slot).saved_context - - -def test_decode_graph_replay_preserves_capture_buffers_and_owner_lifetime(no_cycle_gc): - captured = _create_state(is_prefill=False, secondary=True) - captured.input_ids = torch.zeros(4, dtype=torch.int64) - captured_ref = weakref.ref(captured) - capture_positions_ptr = captured.position_ids.data_ptr() - replay_calls = [] - - class _Graph: - def replay(self): - replay_calls.append(True) - - graph = CudaGraph.__new__(CudaGraph) - output = object() - graph.graph = {4: (_Graph(), captured, output)} - del captured - assert captured_ref() is not None - - runtime = _create_state(is_prefill=False, secondary=True) - runtime.input_ids = torch.ones(4, dtype=torch.int64) - runtime.position_ids.add_(7) - runtime_ref = weakref.ref(runtime) - assert graph._replay(runtime) is output - assert replay_calls == [True] - assert captured_ref().position_ids.data_ptr() == capture_positions_ptr - assert captured_ref().decode_att_state.infer_state is captured_ref() - assert captured_ref().decode_att_state1.infer_state is captured_ref() - torch.testing.assert_close(captured_ref().position_ids, runtime.position_ids) +def test_captured_context_keeps_buffers_while_transient_batch_dies(): + backend = make_backend() + context = make_context(backend, capture=True, max_length=64) + pointer = context.batch.b_seq_len.data_ptr() + runtime = make_batch((17, 3), req_to_token_indexs=context.batch.req_to_token_indexs) + ref = weakref.ref(runtime) + backend.prepare(context, runtime) del runtime - assert runtime_ref() is None - graph.graph.clear() - assert captured_ref() is None - - -@pytest.mark.parametrize("is_prefill", [False, True]) -def test_secondary_state_without_primary(is_prefill): - owner = InferStateInfo() - owner.is_prefill = is_prefill - owner.position_ids = torch.arange(4) - slot = "prefill_att_state1" if is_prefill else "decode_att_state1" - state_cls = _PrefillState if is_prefill else _DecodeState - secondary = state_cls(_infer_state_weak_ref=weakref.ref(owner)) - setattr(owner, slot, secondary) - owner.init_att_state() - torch.testing.assert_close(secondary.positions, owner.position_ids) - if not is_prefill: - runtime = InferStateInfo() - runtime.position_ids = owner.position_ids + 7 - runtime.decode_att_state1 = _DecodeState(_infer_state_weak_ref=weakref.ref(runtime)) - runtime.decode_att_state1.init_state() - owner.copy_for_cuda_graph(runtime) - torch.testing.assert_close(secondary.positions, runtime.position_ids) + assert ref() is None + assert context.batch.b_seq_len.tolist() == [17, 3] + assert context.batch.b_seq_len.data_ptr() == pointer + assert context.current_max_kv_seq_len == 17 diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py index 52b1c1cf2d..e7c0b2e990 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -5,11 +5,33 @@ import torch from lightllm.common.basemodel import cuda_graph as cuda_graph_module +from lightllm.common.basemodel.attention_runtime import AttentionRuntime from lightllm.common.basemodel.cuda_graph import CudaGraph from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune +def make_capture_states(manager, overlap, *, backend=None): + from lightllm.common.attention.api import AttentionPhase + + states = [] + for _ in range(2 if overlap else 1): + runtime = AttentionRuntime() + if backend is not None: + runtime.contexts[0] = SimpleNamespace(capacity=SimpleNamespace(phase=AttentionPhase.DECODE)) + runtime.groups = (SimpleNamespace(backend_for=lambda phase: backend),) + states.append( + SimpleNamespace( + input_ids=torch.ones(2, dtype=torch.int64), + b_req_idx=torch.tensor([0, 1]), + att_state_role="main", + att_state_manager=manager, + attention_runtime=runtime, + ) + ) + return states + + @pytest.fixture def capture_env(monkeypatch): env = SimpleNamespace(capturing=False, capture_count=0, graphs=[]) @@ -22,17 +44,21 @@ def capture_env(monkeypatch): monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") - monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "synchronize", lambda *args: None) monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: None) class FakeGraph: def __init__(self): self.replays = 0 + self.resets = 0 env.graphs.append(self) def replay(self): self.replays += 1 + def reset(self): + self.resets += 1 + @contextmanager def capture(graph, **kwargs): assert not Autotuner.is_autotune_warmup() @@ -45,6 +71,12 @@ def capture(graph, **kwargs): monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) monkeypatch.setattr(torch.cuda, "graph", capture) + from lightllm.platform import DeviceCaps, Platform + from lightllm.platform.cuda import CudaRuntime + + runtime = CudaRuntime("cuda:0") + runtime.caps = DeviceCaps(Platform.CUDA, torch.device("cuda:0"), supports_graph=True) + monkeypatch.setattr(cuda_graph_module, "get_device_runtime", lambda: runtime) return env @@ -82,7 +114,13 @@ def bench(size, run_config): general, general_cache = make_kernel(AutotuneKernelType.GENERAL) attention, attention_cache = make_kernel(AutotuneKernelType.DECODE_ATTENTION) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] + manager = SimpleNamespace( + storage_generation=0, + register_graph=lambda graph: None, + snapshot_capture_state=lambda *args: (), + restore_capture_state=lambda snapshot: None, + ) + states = make_capture_states(manager, overlap) outputs = [object() for _ in states] forward_calls = [] @@ -115,7 +153,13 @@ def decode_func(*infer_states): def test_warmup_failure_restores_phase_without_capturing(capture_env, overlap): capture_env.args.enable_decode_microbatch_overlap = overlap graph = CudaGraph(1, 2, 1, max_batch_size=2) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] + manager = SimpleNamespace( + storage_generation=0, + register_graph=lambda graph: None, + snapshot_capture_state=lambda *args: (), + restore_capture_state=lambda snapshot: None, + ) + states = make_capture_states(manager, overlap) def decode_func(*infer_states): assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) @@ -128,3 +172,41 @@ def decode_func(*infer_states): assert not Autotuner.is_autotune_warmup() assert capture_env.capture_count == 0 assert graph.graph == {} + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("fail_binding", [False, True]) +def test_capture_failure_restores_state_and_releases_native_graph(capture_env, overlap, fail_binding): + from lightllm.common.attention.api import AttentionPhase + + capture_env.args.enable_decode_microbatch_overlap = overlap + graph = CudaGraph(1, 2, 1, max_batch_size=2) + persistent = torch.zeros(1) + manager = SimpleNamespace( + storage_generation=0, + register_graph=lambda graph: None, + snapshot_capture_state=lambda *args: persistent.clone(), + restore_capture_state=lambda snapshot: persistent.copy_(snapshot), + ) + + bindings = [] + + def bind_graph(context): + bindings.append(context) + raise ValueError("failed binding") + + backend = SimpleNamespace(bind_graph=bind_graph) + states = make_capture_states(manager, overlap, backend=backend if fail_binding else None) + + def decode_func(*infer_states): + persistent.add_(1) + if capture_env.capturing and not fail_binding: + raise ValueError("failed capture") + return (None, None) if overlap else None + + with pytest.raises(ValueError, match="failed binding" if fail_binding else "failed capture"): + graph.capture_decode(decode_func, *states) + assert bool(bindings) == fail_binding + assert persistent.item() == 0 + assert not capture_env.capturing and not graph.graph + assert capture_env.graphs[0].resets == 1 diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 123f603ff2..cac2a5cc17 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -1,12 +1,10 @@ -from types import SimpleNamespace import weakref +from types import SimpleNamespace -import torch import pytest +import torch from lightllm.common.basemodel import basemodel -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState -from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -100,7 +98,7 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): expected_batch_size, ) in execution_configs: model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode, page_size=1) + model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode, page_size=1, mtp_step=0) model.tp_world_size_ = tp_world_size model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) @@ -123,10 +121,11 @@ def create_infer_state(model_input): b_seq_len=model_input.b_seq_len, mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32), init_some_extra_state=lambda _: None, + finish_attention=lambda: None, is_cuda_graph_capture=False, is_cuda_graph_replay=False, ) - infer_state.init_att_state = lambda: graph_flags_at_att_init.append( + infer_state.prepare_attention = lambda _: graph_flags_at_att_init.append( (infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay) ) return infer_state @@ -172,7 +171,7 @@ def execute(mode, infer_state): ) def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(lengths, graph_mode, expected_mode): model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1, mtp_step=0) model.tp_world_size_ = 2 model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) @@ -183,16 +182,9 @@ def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(le def graph_flags(state): return state.is_cuda_graph_capture, state.is_cuda_graph_replay - class DecodeState(FlashInferDecodeAttState): - def init_state(self): - infer_state = self.infer_state - # Replace the external GPU wrapper allocation while preserving the - # production decision and InferStateInfo's initialization entry point. - if should_init_decode_wrapper(infer_state): - self.decode_wrapper = object() - initializations.append( - (infer_state.microbatch_index, graph_flags(infer_state), self.decode_wrapper is not None) - ) + def prepare(state): + state.test_wrapper = None if state.is_cuda_graph_replay else object() + initializations.append((state.microbatch_index, graph_flags(state), state.test_wrapper is not None)) def create_state(model_input, microbatch_index=0): state = InferStateInfo() @@ -203,7 +195,8 @@ def create_state(model_input, microbatch_index=0): state.b_req_idx = model_input.b_req_idx state.b_seq_len = model_input.b_seq_len state.microbatch_index = microbatch_index - state.decode_att_state = DecodeState(_infer_state_weak_ref=weakref.ref(state)) + state.prepare_attention = lambda _: prepare(state) + state.finish_attention = lambda: None state.init_some_extra_state = lambda _: None return state @@ -212,8 +205,8 @@ def outputs(state, state1): def forward(state, infer_state1): # Eager and capture execute attention using these newly prepared states. - assert state.decode_att_state.decode_wrapper is not None - assert infer_state1.decode_att_state.decode_wrapper is not None + assert state.test_wrapper is not None + assert infer_state1.test_wrapper is not None executions.append(("forward", graph_flags(state), graph_flags(infer_state1))) return outputs(state, infer_state1) @@ -223,8 +216,8 @@ def capture(func, state, infer_state1): def replay(state, infer_state1): # Replay uses graph-owned wrappers, so temporary states need none. - assert state.decode_att_state.decode_wrapper is None - assert infer_state1.decode_att_state.decode_wrapper is None + assert state.test_wrapper is None + assert infer_state1.test_wrapper is None executions.append(("replay", graph_flags(state), graph_flags(infer_state1))) return outputs(state, infer_state1) @@ -265,7 +258,7 @@ def replay(state, infer_state1): @pytest.mark.parametrize("rows", [0, 3, 9, 24, 27]) def test_mtp_decode_preserves_groups_and_real_rows(overlap, dynamic, graph_mode, rows): model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, mtp_step=0) model.tp_world_size_ = 8 model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1 if dynamic else 3) @@ -283,7 +276,8 @@ def check_layout(): b_req_idx=model_input.b_req_idx, b_seq_len=model_input.b_seq_len, init_some_extra_state=lambda _: None, - init_att_state=check_layout, + finish_attention=lambda: None, + prepare_attention=lambda _: check_layout(), ) def forward(state): diff --git a/unit_tests/common/basemodel/test_moe_route_info_manager.py b/unit_tests/common/basemodel/test_moe_route_info_manager.py index de2bc7a857..a5f9094e7e 100644 --- a/unit_tests/common/basemodel/test_moe_route_info_manager.py +++ b/unit_tests/common/basemodel/test_moe_route_info_manager.py @@ -115,10 +115,6 @@ def test_capture_rejects_unexpected_topk_width(self): def test_cuda_graph_replay_uses_copied_mem_indexes(self): _skip_without_cuda() - class _NoopDecodeState: - def copy_for_decode_cuda_graph(self, other): - return - manager = MoeRouteInfoManager( num_moe_layers=1, topk=2, @@ -126,7 +122,9 @@ def copy_for_decode_cuda_graph(self, other): ) manager.init_capture_buffer(kv_cache_size=256) graph_infer_state = InferStateInfo() - graph_infer_state.decode_att_state = _NoopDecodeState() + from unit_tests.common.attention.helpers import make_batch + + graph_infer_state.attention_batch = make_batch(device="cuda") graph_infer_state.mem_index = torch.tensor([10, 11], device="cuda") moe_capture_callback = manager.get_moe_capture_callback(0, graph_infer_state.mem_index) topk_ids = torch.tensor([[1, 2], [3, 4]], device="cuda") @@ -141,7 +139,9 @@ def copy_for_decode_cuda_graph(self, other): np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) new_infer_state = InferStateInfo() - new_infer_state.decode_att_state = _NoopDecodeState() + from dataclasses import replace + + new_infer_state.attention_batch = replace(graph_infer_state.attention_batch) new_infer_state.mem_index = torch.tensor([20, 21], device="cuda") new_topk_ids = torch.tensor([[5, 6], [7, 8]], device="cuda") graph_infer_state.copy_for_cuda_graph(new_infer_state) diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py index ef7c18a061..56885aac63 100644 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -8,6 +8,7 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA runner stream required") def test_base_model_prefill_accepts_two_prebuilt_inputs(monkeypatch): model_input0 = _make_prefill_input([10, 11], req_idx=20, is_decode_req=False) model_input1 = _make_prefill_input([12], req_idx=21, is_decode_req=True) @@ -52,6 +53,7 @@ def fake_overlap_forward(input0, input1): assert len(outputs) == 2 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA runner stream required") def test_base_model_decode_accepts_two_inputs_and_skips_empty_gather(monkeypatch): model_input0 = _make_decode_input() model_input0.input_ids = None @@ -122,7 +124,7 @@ def test_overlap_decode_cuda_pads_empty_side_and_unpads_outputs(monkeypatch): model_input1.to_cuda() model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1, mtp_step=0) model.tp_world_size_ = 2 model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) model.graph = None @@ -136,7 +138,8 @@ def fake_create_inferstate(model_input, microbatch_index): b_seq_len=model_input.b_seq_len, mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32, device="cuda"), init_some_extra_state=lambda _: None, - init_att_state=lambda: None, + prepare_attention=lambda _: None, + finish_attention=lambda: None, ) model._create_inferstate = fake_create_inferstate diff --git a/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py b/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py index bed4f3967f..a9bcb41401 100644 --- a/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py +++ b/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py @@ -1,16 +1,21 @@ +from contextlib import contextmanager from types import SimpleNamespace +import pytest import torch import lightllm.common.basemodel.hidden_collector as hidden_collector_module from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector -from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph, PrefillGraphExecution +from lightllm.platform import DeviceRuntime class _GraphInferState: def __init__(self, hidden_collector): self.hidden_collector = hidden_collector self.input_ids = torch.empty(4, dtype=torch.int64) + self.att_state_manager = SimpleNamespace(storage_generation=0, register_graph=lambda graph: None) self.copied_from = None self.replayed_with = None @@ -54,7 +59,9 @@ def test_replay_restores_captured_hidden_state_into_request_collector(monkeypatc ) graph_output = torch.randn(2, 3) prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} + prefill_graph.graph = { + 4: PrefillGraphExecution(graph_infer_state, [], [graph_output], graph_collector, DeviceRuntime("cpu")) + } outputs = prefill_graph._replay(input_tensors=[], infer_state=request_infer_state) @@ -73,9 +80,49 @@ def test_first_replay_replaces_capture_collector_with_runtime_instance(monkeypat graph_infer_state = _GraphInferState(hidden_collector=graph_collector) graph_output = torch.randn(2, 3) prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} + prefill_graph.graph = { + 4: PrefillGraphExecution(graph_infer_state, [], [graph_output], graph_collector, DeviceRuntime("cpu")) + } prefill_graph._replay(input_tensors=[], infer_state=graph_infer_state) assert graph_infer_state.hidden_collector is not graph_collector assert graph_infer_state.hidden_collector.layer_hiddens[0] is graph_collector.layer_hiddens[0] + + +@pytest.mark.parametrize("between_segments", [False, True]) +def test_failed_prefill_capture_releases_segments(between_segments): + events = [] + + class Runtime(DeviceRuntime): + def create_graph(self): + return SimpleNamespace(reset=lambda: events.append("reset")) + + @contextmanager + def capture_graph(self, graph, **kwargs): + events.append("enter") + try: + yield + finally: + events.append("exit") + + def synchronize(self): + events.append("synchronize") + + prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) + prefill_graph.graph = {} + prefill_graph.mempool = None + prefill_graph.runtime = Runtime("cpu") + state = InferStateInfo() + state.input_ids = torch.empty(4, dtype=torch.int64) + + def forward(inputs, state): + if between_segments: + state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + raise ValueError("failed segment") + + with pytest.raises(ValueError, match="failed segment"): + prefill_graph._capture_prefill(forward, [], state) + assert events == ["enter", "exit", "synchronize", "reset"] + assert not prefill_graph.graph + assert not state.prefill_cuda_graph_exe_list diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py index 7c1af6bbc2..706f4ff180 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py @@ -13,6 +13,8 @@ import pytest import torch +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def gqa_attention_reference( q: torch.Tensor, @@ -59,7 +61,7 @@ def gqa_attention_reference( # Compute attention scores att = torch.matmul(q_h, k_h.transpose(0, 1)) # [seq_len] - att = att / (head_dim ** 0.5) + att = att / (head_dim**0.5) att = torch.softmax(att, dim=-1) out_h = torch.matmul(att, v_h) # [head_dim] @@ -142,7 +144,7 @@ def mtp_diverse_attention( """ MTP Diverse Attention 调用入口 """ - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( token_decode_attention_mtp_diverse_single_token, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py index 812fd6e9f4..4cdfc1f63e 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py @@ -10,17 +10,27 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel +from lightllm.common.attention.backends.dense.triton import backend as state_module from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( mtp_diverse_attn as decode_module, stage1_single_token as stage1_module, stage2_single_token as stage2_module, ) +from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.attention.triton import fp as state_module from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -138,16 +148,15 @@ def test_decode_passes_saved_length_and_selected_block_n(monkeypatch): infer_state.b_req_idx = inputs["B_req_idx"] infer_state.b_seq_len = inputs["b_seq_len"] infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs model = SimpleNamespace( is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), ) monkeypatch.setattr(state_module, "build_mtp_shared_group_markers", lambda *a, **kw: inputs["b_mark_shared_group"]) - state = state_module.TritonDecodeAttState( - backend=SimpleNamespace(model=model), _infer_state_weak_ref=weakref.ref(infer_state) - ) - state.init_state() + state = backend_context(state_module.TritonAttBackend, "decode", settings=SimpleNamespace(model=model)) + update_context(state, infer_state) infer_state.max_kv_seq_len = 32768 calls = [] @@ -157,7 +166,7 @@ def stage1(**kwargs): monkeypatch.setattr(decode_module, "mtp_diverse_stage1_single_token", stage1) monkeypatch.setattr(decode_module, "mtp_diverse_stage2_single_token", lambda **kwargs: calls.append(kwargs)) - state.decode_att(q=inputs["q"], k=inputs["k"], v=inputs["v"]) + run_dense(state, infer_state, inputs["q"], inputs["k"], inputs["v"]) assert calls[0]["max_kv_len"] == 8193 assert calls[0]["b_mark_shared_group"] is inputs["b_mark_shared_group"] assert calls[1]["block_n"] == 32 @@ -264,7 +273,17 @@ def decode(state): return out graph = graph_module.CudaGraph(1, 7, 1, max_batch_size=7, max_len_in_batch=32768) - graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(7, device="cuda"))) + resources = AttentionStateManager(request_capacity=7) + resources.finish_state_registration() + state = SimpleNamespace( + input_ids=torch.ones(7, device="cuda"), + b_req_idx=torch.arange(7, device="cuda", dtype=torch.int32), + att_state_manager=resources, + att_state_role="main", + attention_contexts={}, + attention_groups=(), + ) + graph.capture_decode(decode, state) torch.cuda.synchronize() assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) assert valid_benchmark_count > 0 @@ -296,6 +315,6 @@ def decode(state): inputs["b_seq_len"].copy_( torch.tensor([kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len], device="cuda", dtype=torch.int32) ) - graph.graph[7][0].replay() + graph.graph[7].replay() torch.cuda.synchronize() torch.testing.assert_close(out.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py index feb4a81c0d..dc456ba899 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py @@ -10,6 +10,7 @@ from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( stage2_single_token as stage2_module, ) +from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.basemodel import cuda_graph as graph_module from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner @@ -191,7 +192,17 @@ def decode(state): return inputs["out"] graph = graph_module.CudaGraph(1, 2, 1, max_batch_size=2, max_len_in_batch=32768) - graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(2, device="cuda"))) + resources = AttentionStateManager(request_capacity=2) + resources.finish_state_registration() + state = SimpleNamespace( + input_ids=torch.ones(2, device="cuda"), + b_req_idx=torch.arange(2, device="cuda", dtype=torch.int32), + att_state_manager=resources, + att_state_role="main", + attention_contexts={}, + attention_groups=(), + ) + graph.capture_decode(decode, state) torch.cuda.synchronize() assert rebuild_count == 1 and benchmark_count == 20 torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) @@ -217,6 +228,6 @@ def decode(state): inputs["mid_out"].normal_() inputs["mid_out_logsumexp"].normal_() inputs["B_Seqlen"].fill_(length) - graph.graph[2][0].replay() + graph.graph[2].replay() torch.cuda.synchronize() torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py index ffc322e216..02a7c622a7 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py @@ -10,20 +10,27 @@ from frozendict import frozendict from torch.nn.attention import SDPBackend, sdpa_kernel +from lightllm.common.attention.backends.dense.triton.backend import TritonAttBackend from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash import ( + gqa_flash_decoding as decode_module, gqa_flash_decoding_stage1 as stage1_module, ) from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage2 import ( flash_decode_stage2, ) -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( - gqa_flash_decoding as decode_module, -) from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -129,24 +136,27 @@ def test_rebuild_preserves_layout_and_bounds_indices(monkeypatch, num_tokens): def test_state_retains_actual_length_before_graph_capture(monkeypatch): infer_state = InferStateInfo() infer_state.max_kv_seq_len = 8193 - state = TritonDecodeAttState( - backend=SimpleNamespace( + state = backend_context( + TritonAttBackend, + "decode", + settings=SimpleNamespace( model=SimpleNamespace( is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0) ) ), - _infer_state_weak_ref=weakref.ref(infer_state), ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 + update_context(state, infer_state) + infer_state.max_kv_seq_len = 32768 calls = [] monkeypatch.setattr( decode_module, "gqa_token_decode_attention_flash_decoding", lambda **kwargs: calls.append(kwargs) ) - state.decode_att( - q=torch.empty(2, 8, 64), - k=torch.empty(16, 2, 64), - v=torch.empty(16, 2, 64), + run_dense( + state, + infer_state, + torch.empty(2, 8, 64), + torch.empty(16, 2, 64), + torch.empty(16, 2, 64), alloc_func=lambda shape, dtype: torch.empty(shape, dtype=dtype), ) assert calls[0]["max_len_in_batch"] == 8193 @@ -260,12 +270,14 @@ def decode(): decode() torch.cuda.synchronize() + eager_output = output.clone() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): decode() graph.replay() torch.cuda.synchronize() torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + torch.testing.assert_close(output, eager_output, atol=0, rtol=0) assert rebuild_count == 1 and benchmark_count == 3 for name, value in snapshots.items(): if name not in ["mid_out", "mid_out_logsumexp"]: diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py index 5237d0b0f4..a179c21994 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py @@ -10,18 +10,27 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel +from lightllm.common.attention.backends.dense.triton.int4 import Int4kvTritonAttBackend from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv import ( int4kv_flash_decoding_stage1 as stage1_module, + ppl_int4kv_flash_decoding as decode_module, ) from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) -from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ppl_int4kv_flash_decoding as decode_module from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -141,17 +150,17 @@ def test_decode_passes_saved_length_to_stage1(monkeypatch): infer_state.b_req_idx = inputs["B_req_idx"] infer_state.b_seq_len = inputs["B_Seqlen"] infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - state = Int4kvTritonDecodeAttState( - backend=SimpleNamespace(), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs + state = backend_context(Int4kvTritonAttBackend, "decode", settings=SimpleNamespace()) + update_context(state, infer_state) + infer_state.max_kv_seq_len = 32768 calls = [] monkeypatch.setattr(stage1_module, "int4kv_flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) stage2 = inspect.getmodule(flash_decode_stage2) monkeypatch.setattr(stage2, "flash_decode_stage2", lambda *args: None) - state.decode_att( + run_dense( + state, + infer_state, inputs["q"], (inputs["k"], inputs["k_scale"]), (inputs["v"], inputs["v_scale"]), diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py index e1b6bc7758..7a478f0f98 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py @@ -1,21 +1,14 @@ import pytest -from types import SimpleNamespace - import torch +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + def alloc_tensor_func(shape, dtype, device): """兼容的 tensor 分配函数""" return torch.empty(shape, dtype=dtype, device=device) -class MockReqManager: - """Mock request manager for testing""" - - def __init__(self, req_to_token_indexs): - self.req_to_token_indexs = req_to_token_indexs - - class MockInferState: """Mock infer state for testing""" @@ -29,7 +22,7 @@ def __init__( ): self.batch_size = batch_size self.max_kv_seq_len = max_kv_seq_len - self.req_manager = MockReqManager(req_to_tokens) + self.req_to_token_indexs = req_to_tokens self.b_req_idx = b_req_idx self.b_seq_len = b_seq_len @@ -43,10 +36,10 @@ def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(sha diverse：int8kv_flash_decoding_diverse；对照：int8kv/normal token_decode_attention_flash_decoding。 """ - from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse import ( token_decode_attention_flash_decoding as diverse_attention, ) - from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( token_decode_attention_flash_decoding as normal_decode, ) @@ -102,15 +95,10 @@ def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(sha b_req_idx=b_req_idx, b_seq_len=b_seq_len, ) - diverse_infer_state.decode_att_state = SimpleNamespace( - b_shared_seq_len=b_shared_seq_len, - b_mark_shared_group=b_mark_shared_group, - ) - # 运行 baseline normal_out = normal_decode( q=q.clone(), - infer_state=baseline_infer_state, + batch=baseline_infer_state, cache_k=cache_k, cache_k_scale=cache_k_scale, cache_v=cache_v, @@ -121,7 +109,10 @@ def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(sha # 运行 diverse 版本 diverse_out = diverse_attention( q=q.clone(), - infer_state=diverse_infer_state, + batch=diverse_infer_state, + b_shared_seq_len=b_shared_seq_len, + b_mark_shared_group=b_mark_shared_group, + shared_streams_dict={}, cache_k=cache_k, cache_k_scale=cache_k_scale, cache_v=cache_v, diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py index bd9619711f..3466d2f105 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py @@ -1,6 +1,7 @@ import pytest import torch +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( flash_decode_stage1, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py index 449e6a83b8..32925d2a0e 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py @@ -1,6 +1,7 @@ import pytest import torch +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( flash_decode_stage2, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py index 00f0783abd..9dfdd2c506 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py @@ -1,6 +1,7 @@ import pytest import torch +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage3 import ( flash_diverse_decode_stage3, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py index 487f03b191..fbc9c25517 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py @@ -10,19 +10,28 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel +from lightllm.common.attention.backends.dense.triton import int8 as state_module +from lightllm.common.attention.backends.dense.triton.int8 import Int8kvTritonAttBackend from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( + int8kv_flash_decoding as decode_module, int8kv_flash_decoding_stage1 as stage1_module, ) from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) -from lightllm.common.basemodel.attention.triton import int8kv as state_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import int8kv_flash_decoding as decode_module from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -143,16 +152,16 @@ def test_decode_passes_saved_length_to_stage1(monkeypatch): infer_state.b_req_idx = inputs["B_req_idx"] infer_state.b_seq_len = inputs["B_seq_len"] infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - state = Int8kvTritonDecodeAttState( - backend=SimpleNamespace(), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs + state = backend_context(Int8kvTritonAttBackend, "decode", settings=SimpleNamespace()) + update_context(state, infer_state) + infer_state.max_kv_seq_len = 32768 calls = [] monkeypatch.setattr(decode_module, "flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) monkeypatch.setattr(decode_module, "flash_decode_stage2", lambda **kwargs: None) - state.decode_att( + run_dense( + state, + infer_state, inputs["q"], (inputs["k"], inputs["k_scale"]), (inputs["v"], inputs["v_scale"]), diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py index 4aafd9d104..4d53830a04 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py @@ -28,7 +28,7 @@ import pytest import torch -from lightllm.common.basemodel.attention.fa3.neo import flash_attn_with_kvcache_neo as flash_attn_with_kvcache +from lightllm.common.attention.backends.dense.fa3.neo import flash_attn_with_kvcache_neo as flash_attn_with_kvcache try: import triton diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py index 1aab2b32a3..5678745a58 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py @@ -5,6 +5,8 @@ from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import causal_conv1d_update +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def causal_conv1d_ref( x: torch.Tensor, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index 9e9f237b75..3154c0349a 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -867,6 +867,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa torch.manual_seed(72) initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda) + backend.tile_sync_mark = torch.zeros((bucket, backend.num_key_heads), dtype=torch.int32, device="cuda") reqs = torch.cat( ( torch.arange(live, device="cuda", dtype=torch.int32), diff --git a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py index a5ac9708d2..e6bbe2bcf7 100644 --- a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py +++ b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py @@ -1,14 +1,16 @@ -import torch import time -import pytest + +import flashinfer import numpy as np +import pytest +import torch import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.mla_att.decode_att.gqa_flash_decoding import ( + +from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding import ( gqa_token_decode_attention_flash_decoding, ) from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -52,10 +54,11 @@ def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head): infer_state = Deepseek2InferStateInfo() infer_state.batch_size = Z - infer_state.max_len_in_batch = N_CTX + infer_state.max_kv_seq_len = N_CTX infer_state.total_token_num = Z * N_CTX infer_state.req_manager = type("Object", (), {})() infer_state.req_manager.req_to_token_indexs = req_to_token_indexs + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs infer_state.b_req_idx = b_req_idx infer_state.b_seq_len = b_seq_len infer_state.kv_starts = kv_starts @@ -70,6 +73,8 @@ def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head): infer_state, sm_scale, o, + block_seq=torch.empty(1, dtype=torch.int64, device="cuda"), + batch_start=torch.empty(Z, dtype=torch.int64, device="cuda"), ) batch_size = Z diff --git a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py index 0c1c2df08a..225d4b456e 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py +++ b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py @@ -7,10 +7,19 @@ if not torch.cuda.is_available(): pytest.skip("requires CUDA", allow_module_level=True) +from lightllm.common.attention.backends.dense.triton import int8 as int8kv_module +from lightllm.common.attention.backends.dense.triton.int8 import Int8kvTritonAttBackend from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as int8kv_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState from lightllm.common.basemodel.triton_kernel import diverse_utils +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.mark.parametrize( @@ -40,9 +49,9 @@ def test_int8kv_decode_state_rebuilds_diverse_metadata(monkeypatch): infer_state.max_kv_seq_len = 8192 infer_state.b_shared_seq_len = torch.tensor([8, 8, 5, 0], dtype=torch.int32, device="cuda") infer_state.b_shared_radix_node_id = torch.tensor([10, 10, 20, -1], dtype=torch.int64, device="cuda") - state = Int8kvTritonDecodeAttState(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer_state)) + state = backend_context(Int8kvTritonAttBackend, "decode", settings=SimpleNamespace()) - state.init_state() + update_context(state, infer_state) assert state.b_mark_shared_group.cpu().tolist() == [0, 2, 1, 1] assert state.b_shared_seq_len.cpu().tolist() == [8, 8, 0, 0] diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index f5bd92e08a..558ae93222 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -1,15 +1,18 @@ import json +from dataclasses import replace from types import SimpleNamespace -import numpy as np import pytest import torch -from unit_tests.common.dsv4_test_utils import DeepseekV4MemoryManager, DeepseekV4CacheConfig, DeepseekV4ReqManager -from unit_tests.common.req_token_index_test_utils import attach_req_token_index +from lightllm.common.attention.api import ExecutionRequirements from lightllm.models.deepseek_v4.triton_kernel.build_compressed_kv_indices import build_compressed_kv_indices -from lightllm.utils.envs_utils import get_env_start_args from lightllm.models.deepseek_v4.triton_kernel.swa_cache import get_swa_token_indices as _get_swa_token_indices +from lightllm.platform import get_device_runtime +from lightllm.utils.envs_utils import get_env_start_args +from unit_tests.common.attention.helpers import backend_context, make_spec, run_sparse, update_context +from unit_tests.common.dsv4_test_utils import DeepseekV4CacheConfig, DeepseekV4MemoryManager, DeepseekV4ReqManager +from unit_tests.common.req_token_index_test_utils import attach_req_token_index def _compressed_kv_pools(manager): @@ -245,11 +248,13 @@ def _write_packed_rows(buffer, layer_index, slots, packed): def _test_backend(requests, token_capacity=8192): - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4FlashMlaFp8SparseAttBackend model = type("TestModel", (), {})() model.req_manager = requests model.att_state_manager = requests.memory.manager + model.device_runtime = get_device_runtime("cuda") + model.attention_requirements = ExecutionRequirements() model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] model.tp_rank_ = 0 model.mtp_previous_draft_models = [] @@ -274,8 +279,14 @@ def _test_backend(requests, token_capacity=8192): enable_decode_microbatch_overlap=False, mtp_mode=None, mtp_step=0, + mtp_dynamic_verify=False, + page_size=256, ) - return DeepseekV4FlashMlaFp8SparseAttBackend(model) + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda _: 0) + model.swa_prefill_workspace = DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(model) + requests._test_model = model + return DeepseekV4FlashMlaFp8SparseAttBackend(model, make_spec(cache=replace(make_spec().cache, page_size=256))) def _prepare_prefill(requests, req_ids, starts, ends, microbatch_index=0): @@ -283,20 +294,11 @@ def _prepare_prefill(requests, req_ids, starts, ends, microbatch_index=0): if not hasattr(requests, "_test_backend"): requests._test_backend = _test_backend(requests) requests._test_prefill_states = {} - state = requests._test_backend.create_att_prefill_state(None) - state.prepare_swa_workspace( - SimpleNamespace( - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - b_req_idx=req_ids, - b_ready_cache_len=starts, - b_seq_len=ends, - microbatch_index=microbatch_index, - ) - ) + from lightllm.models.deepseek_v4.attention.sparse_mla import prepare_swa_workspace + + buffer, capacity = requests._test_model.swa_prefill_workspace + batch = SimpleNamespace(b_req_idx=req_ids, b_ready_cache_len=starts, b_seq_len=ends) + state = prepare_swa_workspace(requests.swa_state, buffer, capacity, batch, microbatch_index) requests._test_prefill_states.update((req, state) for req in req_ids.tolist()) if not hasattr(requests, "_test_prefill_ends"): requests._test_prefill_ends = {} @@ -441,8 +443,8 @@ def cpu(data): def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): from lightllm.common.basemodel.batch_objs import ModelInput from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel - from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel + from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel manager, requests = cache req_idx, slots = _reserve(manager, requests, 512) @@ -453,6 +455,8 @@ def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): model = model_cls.__new__(model_cls) model.req_manager = requests model.att_state_manager = requests.memory.manager + model.device_runtime = get_device_runtime("cuda") + model.attention_requirements = ExecutionRequirements() model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] model.tp_rank_ = 0 model.mtp_previous_draft_models = [] @@ -523,7 +527,7 @@ def test_prefill_checkpoint_reads_precede_window_retention(cache): _write_swa(requests, 2, req, positions, packed.flip(0)) _finish_prefill(requests, [req]) # Reuse the scratch: only fixed request windows and CPU snapshots remain. - requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model)[0].zero_() + requests._test_model.swa_prefill_workspace[0].zero_() requests.save_state(_checkpoint_req(req, 2048), index, states) saved_swa = states.get_state_cache(index)[0] torch.testing.assert_close(saved_swa[0].cuda(), packed[-requests.sliding_window :], rtol=0, atol=0) @@ -802,11 +806,15 @@ def test_compressor_checkpoint_continuation_matches_full_prefill(cache, ratio, l from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace manager, requests = cache - workspace = DeepseekV4Workspace(SimpleNamespace(tp_world_size_=1, batch_max_tokens=length)) - workspace.c128_compress_scratch = torch.empty((1, length, 512), dtype=torch.float32, device="cuda") source, src = _reserve(manager, requests, length) destination, dst = _reserve(manager, requests, length) _prepare_swa(requests, source, 0, length) + model = requests._test_model + model.disable_cudagraph = True + model._align_decode_batch_size = lambda count: count + model.mtp_manager.get_decode_tokens_per_request = lambda draft: get_env_start_args().mtp_step + 1 + workspace = DeepseekV4Workspace(model) + workspace.init_c128_compress_scratch(model) head_dim = requests.cache_config.indexer_head_dim if is_indexer else requests.cache_config.head_dim scores = torch.randn((length, head_dim * (4 if ratio == 4 else 2)), dtype=torch.float32, device="cuda") cos = torch.ones((length, 32), dtype=torch.float32, device="cuda") @@ -942,9 +950,9 @@ def build(): indirect=["cache"], ) def test_swa_ring_attention_matches_full_cache_after_speculative_retries(cache, dspark): - from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + from lightllm.models.deepseek_v4.attention import sparse_mla as attention from lightllm.models.deepseek_v4.triton_kernel.build_dspark_swa_indices import build_dspark_swa_indices + from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices torch.manual_seed(1593) manager, requests = cache @@ -986,44 +994,25 @@ def build(): else: build_swa_indices(req_idx, query_positions, requests, indices, lengths, write_slots) - backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend( - type( - "Model", - (), - { - "config": dict( - num_attention_heads=64, - compress_ratios=[0, 0, 0], - head_dim=512, - qk_rope_head_dim=64, - index_head_dim=128, - sliding_window=128, - ), - "layers_infer": [SimpleNamespace(layer_num_=2)], - "tp_rank_": 0, - "tp_world_size_": 1, - "req_manager": requests, - }, - )() - ) + backend = _test_backend(requests) state_info = SimpleNamespace(mem_manager=manager) - state = backend.create_att_decode_state(state_info) - state.init_state() + state = backend_context(type(backend), "decode", settings=backend) + update_context(state, state_info) params = dict( compress_ratio=0, head_dim_v=512, - softmax_scale=512 ** -0.5, + softmax_scale=512**-0.5, swa_indices=indices.unsqueeze(1), swa_lengths=lengths, attn_sink=torch.zeros(64, device="cuda"), ) - control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) + control = params q = torch.randn((token_num, 64, 512), dtype=torch.bfloat16, device="cuda") * 0.1 packed_kv = requests.get_att_input_params(2) def forward(): build() - return state.decode_att(q, packed_kv, None, att_control=control) + return run_sparse(state, state_info, q, packed_kv, control) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -1034,20 +1023,9 @@ def forward(): with torch.cuda.graph(graph, stream=stream): actual = forward() - reference_state = backend.create_att_decode_state(state_info) - reference_state.init_state() - reference_kv = backend.get_att_input_params( - SimpleNamespace( - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - is_prefill=True, - prefill_att_state=requests._test_prefill_states[source], - ), - 2, - ) + reference_state = backend_context(type(backend), "decode", settings=backend) + update_context(reference_state, state_info) + reference_kv = (requests._test_prefill_states[source].swa_buffer[2].view(-1, 74880), None) committed_end = 512 # Both sides of page boundaries, repeated retries, and several complete ring wraps. for start in (512, 513, 638, 639, 640, 640, 641, 767, 768, 1023, 1279, 1280, 1919, 1920, 2047): @@ -1080,9 +1058,7 @@ def forward(): ) reference_indices = reference_indices.to(torch.int32) reference_params = dict(params, swa_indices=reference_indices.unsqueeze(1), swa_lengths=reference_lengths) - expected = reference_state.decode_att( - q, reference_kv, None, att_control=attention.AttControl(nsa_decode=True, nsa_decode_dict=reference_params) - ) + expected = run_sparse(reference_state, state_info, q, reference_kv, reference_params) torch.testing.assert_close(actual, expected, rtol=0, atol=0, msg=f"SWA mismatch at position {start}") torch.testing.assert_close(write_slots, _swa_slots(requests, destination, logical_positions).to(torch.int32)) # Accept one position; speculative future rows stay dirty until the next write. @@ -1096,8 +1072,9 @@ def forward(): @pytest.mark.parametrize("hit_extra_page", [False, True]) def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash_page_size, hit_extra_page): from sortedcontainers import SortedDict - from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context, CacheTier + from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache + from lightllm.server.router.model_infer.infer_batch import CacheTier, InferReq, g_infer_context from lightllm.utils.kv_cache_utils import compute_token_list_hash manager, requests = cache @@ -1341,12 +1318,13 @@ def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill manager, requests = cache first, first_slots = _reserve(manager, requests, 512) second, second_slots = _reserve(manager, requests, 512) - backend = _test_backend(requests) - model = backend.model + _test_backend(requests) + model = requests._test_model model.layers_infer = [layer for layer in model.layers_infer if layer.layer_num_ in owned_layer_ids] has_c128_layer = 1 in owned_layer_ids model.graph = SimpleNamespace(can_run=lambda *args: True, graph_max_len_in_batch=16384) model.dsv4_workspace = DeepseekV4Workspace(model) + model._attention_storage_generation = model.att_state_manager.storage_generation rope = torch.zeros((1024, 32), device="cuda") model._cos_cached_sliding = model._sin_cached_sliding = rope model._cos_cached_compress = model._sin_cached_compress = rope @@ -1361,15 +1339,14 @@ def make_state(lengths, microbatch_index): state.first_c4_logical_layer_id = 0 state.batch_size = 2 state.max_kv_seq_len = max(lengths) + state.total_token_num = sum(lengths) state.input_ids = torch.zeros(2, dtype=torch.int64, device="cuda") state.b_req_idx = torch.tensor([first, second], dtype=torch.int32, device="cuda") state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") if is_prefill: state.b_ready_cache_len = state.b_seq_len - 1 - state.prefill_att_state = backend.create_att_prefill_state(state) - else: - state.decode_att_state = backend.create_att_decode_state(state) state.init_some_extra_state(model) + state.prepare_attention(model) return state captured_state = make_state([130, 258], 0) @@ -1416,9 +1393,9 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, import lightllm.models.deepseek_v4.triton_kernel.build_swa_indices as swa_kernel import lightllm.models.deepseek_v4_dspark.infer_struct as dspark_state + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4FlashMlaFp8SparseAttBackend from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace from lightllm.models.deepseek_v4_dspark.infer_struct import DeepseekV4DSparkInferStateInfo - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend manager, requests = cache req_idx, slots = _reserve(manager, requests, 512) @@ -1454,7 +1431,7 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, att_state_role="main", layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], ) - backend = DeepseekV4FlashMlaFp8SparseAttBackend(model) + model.swa_prefill_workspace = DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(model) model.block_size = 3 model.dsv4_workspace = DeepseekV4Workspace(model) rope = torch.zeros((512, 32), device="cuda") @@ -1474,9 +1451,7 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, state.b_req_idx = torch.tensor([req_idx], dtype=torch.int32, device="cuda") state.b_seq_len = torch.tensor([260], dtype=torch.int32, device="cuda") state.b_ready_cache_len = torch.tensor([257], dtype=torch.int32, device="cuda") - state.prefill_att_state = backend.create_att_prefill_state(state) else: - state.decode_att_state = backend.create_att_decode_state(state) state.b_req_idx = torch.full((3,), req_idx, dtype=torch.int32, device="cuda") state.b_seq_len = torch.arange(258, 261, dtype=torch.int32, device="cuda") if mode == "commit": @@ -1495,7 +1470,7 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, state.dsv4_swa_lengths, torch.full_like(state.dsv4_swa_lengths, 131 if proposal else 128) ) expected_slots = _get_swa_token_indices( - requests, req_idx, state.position_ids, prefill_state=state.prefill_att_state if state.is_prefill else None + requests, req_idx, state.position_ids, prefill_state=state.swa_workspace if state.is_prefill else None ).to(torch.int32) torch.testing.assert_close(state.dsv4_swa_write_slots, expected_slots) requests.free([req_idx], slots) @@ -1525,11 +1500,13 @@ def test_empty_dspark_commit_is_local_but_proposal_keeps_common_padding_path(mon ) @pytest.mark.parametrize("exhaust_small_pool", [False, True]) @pytest.mark.parametrize("token_index_mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("delay_producer", [False, True]) def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( - cache, monkeypatch, length, prefix, exhaust_small_pool, token_index_mode + cache, monkeypatch, length, prefix, exhaust_small_pool, token_index_mode, delay_producer ): """Offload after the runtime ring is overwritten, then resume through the common loader.""" from collections import deque + from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.router.model_infer.infer_batch import InferReq from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common @@ -1641,6 +1618,9 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( module.page_ready_buffer = torch.empty(8, dtype=torch.bool, device="cuda") module.gloo_group = module.filter_group = module.init_sync_group = None module.need_sync_compute_stream = lambda: False + if delay_producer: + torch.cuda._sleep(200_000_000) + _randomize_compressed_kv(manager) assert module.offload_finished_reqs_to_cpu_cache([req]) == [] assert req.cpu_cache_offload_task_status.is_running() torch.cuda.synchronize() @@ -1741,8 +1721,8 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer - from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa manager, requests = cache req, held = _reserve(manager, requests, 2048) @@ -1753,9 +1733,19 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): freqs = torch.complex(angles.cos(), angles.sin()) layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) layer.layer_num_, layer.eps_, layer.freqs_cis = 2, 1e-6, freqs - layer_weight = SimpleNamespace(kv_norm_=SimpleNamespace(weight=weight)) + layer.compress_ratio, layer.v_head_dim, layer.softmax_scale = 0, 512, 512**-0.5 + layer_weight = SimpleNamespace( + kv_norm_=SimpleNamespace(weight=weight), + attn_sink_=SimpleNamespace(weight=torch.zeros(64, device="cuda")), + ) for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): _prepare_swa(requests, req, start, end) + if start == 0: + from lightllm.common.attention.api import AttentionFamily + + layer.prefill_attention = SimpleNamespace( + params=layer.create_attention_params(requests._test_model, layer_weight, AttentionFamily.SPARSE_MLA) + ) positions = torch.arange(start, end, device="cuda", dtype=torch.int32) indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") lengths = torch.empty(positions.numel(), dtype=torch.int32, device="cuda") @@ -1781,7 +1771,7 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, - prefill_att_state=requests._test_prefill_states[req], + swa_workspace=requests._test_prefill_states[req], dsv4_swa_write_slots=slots, dsv4_swa_runtime_write_slots=runtime, position_ids=positions, @@ -1866,7 +1856,7 @@ def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_le SimpleNamespace(convert_tokens_to_ids=lambda token: 99), "/unused", {"vision_n_layers": 1} ) params = MultimodalParams(images=[{"type": "url", "data": "unused"} for _ in range(3)]) - token_base, canonical_length = 2 ** 62, 13 + token_base, canonical_length = 2**62, 13 for image in params.images: image.token_id, image.token_num, image.start_index_in_embed_cache = token_base, canonical_length, 0 input_ids = tokenizer.encode([7] * prefix_len + [99, 99, 8, 99, 8], params) @@ -1880,8 +1870,8 @@ def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_le ready = spans[0][1] if cached_first_image else 0 manager, requests = cache req_idx, slots = _reserve(manager, requests, len(input_ids)) - backend = _test_backend(requests, token_capacity=128) - model = backend.model + _test_backend(requests, token_capacity=128) + model = requests._test_model model.has_vision = True model.config["vision_max_n_token"] = 384 model.dsv4_workspace = None @@ -1901,7 +1891,6 @@ def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_le state.b_ready_cache_len = torch.tensor([ready], dtype=torch.int32, device="cuda") state.max_kv_seq_len = len(input_ids) state.multimodal_params = [params.to_dict()] - state.prefill_att_state = backend.create_att_prefill_state(state) state.init_some_extra_state(model) assert state.dsv4_image_left.cpu().tolist() == expected_left[ready:] assert state.dsv4_image_right.cpu().tolist() == expected_right[ready:] @@ -1937,8 +1926,8 @@ def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace manager, requests = cache - backend = _test_backend(requests) - model = backend.model + _test_backend(requests) + model = requests._test_model model.has_vision = False rope = torch.zeros((8, 32), device="cuda") model._cos_cached_sliding = model._sin_cached_sliding = rope @@ -1963,7 +1952,6 @@ def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): infer_state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") infer_state.b_ready_cache_len = torch.zeros_like(infer_state.b_seq_len) infer_state.max_kv_seq_len = max(lengths) - infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) infer_state.init_some_extra_state(model) previous_debug_mode = torch.cuda.get_sync_debug_mode() @@ -2018,7 +2006,7 @@ def test_full_batch_unaligned_prefill_workspace_is_bounded_and_private(cache, mt required = request_count * ((end + page_token_num - 1) // page_token_num * page_token_num - first) requests._test_backend = _test_backend(requests, request_count * chunk_tokens) requests._test_prefill_states = {} - workspace, capacity = requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model) + workspace, capacity = requests._test_model.swa_prefill_workspace assert capacity >= required assert workspace.untyped_storage().data_ptr() != requests.req_to_swa.untyped_storage().data_ptr() workspace[:, capacity : 2 * capacity].fill_(173) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index e7360c94ab..fce761bae8 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -117,7 +117,7 @@ def launch(*args, **kwargs): def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, module_name, cls, bv, precompute): import importlib - decode = importlib.import_module(f"lightllm.common.basemodel.attention.linear.{module_name}") + decode = importlib.import_module(f"lightllm.common.attention.backends.{module_name.split('_')[1]}.replay.decode") captured = _CaptureKernel() monkeypatch.setattr(decode, "_decode_replay", captured) owner = cls( @@ -238,3 +238,23 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): _commit_prefix_cpu(owner.mtp_state_index, reqs, lengths, lengths, history.cursors) assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] assert history.cursors.tolist() == [3, 0, 1, 0] + + +@pytest.mark.parametrize("cls", [ReplayGDNState, ReplayKDAState]) +def test_replay_capture_snapshot_restores_history_and_cursor(linear_config, cls): + from lightllm.common.attention_state import AttentionStateManager + + owner = cls(linear_config, request_capacity=3, mtp_step=2, device="cpu") + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", owner) + manager.finish_state_registration() + views = owner.capture_snapshot_views([1]) + for index, view in enumerate(views): + view.fill_(index + 1) + snapshot = manager.snapshot_capture_state(torch.tensor([1, 1]), "main") + for view in views: + view.zero_() + manager.restore_capture_state(snapshot) + for index, view in enumerate(views): + assert torch.all(view == index + 1) + assert owner.describe_layout()[0] in ("gdn_replay_conv_ssm", "kda_replay_conv_ssm") diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py index 8394d18b48..c266ec3db9 100644 --- a/unit_tests/models/deepseek_v4/test_model.py +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -3,8 +3,8 @@ import pytest import torch -from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel @pytest.mark.parametrize("capture", [False, True], ids=["eager", "capture"]) @@ -46,10 +46,10 @@ def test_rotary_cache_keeps_unscaled_swa_and_yarn_compression(): positions = torch.tensor([0, 1, 127, 1023, 8192], device="cuda") exponents = torch.arange(0, 64, 2, device="cuda", dtype=torch.float32) / 64 - angles = torch.outer(positions.float(), 1.0 / (10000 ** exponents)) + angles = torch.outer(positions.float(), 1.0 / (10000**exponents)) expected = torch.complex(angles.cos(), angles.sin()) torch.testing.assert_close(model._freqs_cis_sliding[positions], expected, rtol=1e-6, atol=1e-6) - unscaled_compress = torch.outer(positions.float(), 1.0 / (160000 ** exponents)) + unscaled_compress = torch.outer(positions.float(), 1.0 / (160000**exponents)) assert not torch.allclose( model._freqs_cis_compress[positions], torch.complex(unscaled_compress.cos(), unscaled_compress.sin()) ) @@ -264,6 +264,7 @@ def test_output_projection_matches_bf16_reference(monkeypatch, tp_size, rank): layer = DeepseekV4TransformerLayerWeight.__new__(DeepseekV4TransformerLayerWeight) BaseLayerWeight.__init__(layer) layer.prefix, layer.data_type_ = "layers.0", torch.bfloat16 + layer.num_fused_shared_experts = 0 layer.hidden, layer.n_heads, layer.head_dim = 128, 8, 64 layer.q_lora_rank, layer.o_lora_rank, layer.o_groups = 128, 128, 2 layer.get_quant_method = lambda name: None diff --git a/unit_tests/models/deepseek_v4/test_workspace.py b/unit_tests/models/deepseek_v4/test_workspace.py index 2bc4a56719..926f95c6b2 100644 --- a/unit_tests/models/deepseek_v4/test_workspace.py +++ b/unit_tests/models/deepseek_v4/test_workspace.py @@ -8,7 +8,6 @@ from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -17,11 +16,11 @@ def test_draft_initialization_preserves_captured_prefill_buffer(monkeypatch, dra args = SimpleNamespace(llm_kv_type="fp8kv_dsa") monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_env_start_args", lambda: args) - def backend_class(**kwargs): - return lambda model: SimpleNamespace(tp_q_head_num=2, padded_q_head_num=4) + def initialize(model): + backend = SimpleNamespace(tp_q_head_num=2, padded_q_head_num=4) + model.attention_groups = (SimpleNamespace(prefill_backend=backend, decode_backend=backend),) - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_prefill_att_backend_class", backend_class) - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_decode_att_backend_class", backend_class) + monkeypatch.setattr("lightllm.common.basemodel.attention.initialize_attention_groups", initialize) main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) main.batch_max_tokens, main.tp_world_size_, main.data_type = 8, 1, torch.bfloat16 main.config = {"head_dim": 512, "block_size": 3, "dspark_noise_token_id": 1} @@ -36,7 +35,8 @@ def backend_class(**kwargs): "_sin_cached_compress", ): setattr(main, name, torch.ones((1, 32), device="cuda")) - main._init_att_backend() + main.args = args + main.prepare_attention_backends() buffer = main.dsv4_workspace.flashmla_prefill_q buffer[:, :2].fill_(2) graph = torch.cuda.CUDAGraph() @@ -48,7 +48,8 @@ def backend_class(**kwargs): draft.layers_infer = [SimpleNamespace(compress_ratio=0)] draft.context_wkv_weight = object() draft._init_custom() - draft._init_att_backend() + draft.args = args + draft.prepare_attention_backends() assert draft.dsv4_workspace is main.dsv4_workspace assert main.dsv4_workspace.flashmla_prefill_q is buffer diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 6da740353d..9f77b75e6b 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -1,4 +1,3 @@ -from unit_tests.common.att_resources.model_helpers import build_hybrid_resources import dataclasses from types import SimpleNamespace @@ -6,17 +5,16 @@ import torch from lightllm.common.attention_state.allocators import FullAttKVAllocator - from lightllm.common.attention_state.full_att import ( - PackedMLAFullAttentionState, PackedMLAAttentionConfig, + PackedMLAFullAttentionState, PackedMLALogicalConfig, ) from lightllm.common.req_manager import ReqManager from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name - +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index d0e6dcb9bd..2ab1edbaba 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -1,24 +1,10 @@ -import weakref -from types import SimpleNamespace - import pytest import torch -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.nsa.glm5_next import ( - Glm5NextSparseAttBackend, - Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - Fp8Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next +from lightllm.common.attention.api import AttentionBatch, AttentionPhase, ContextCapacity +from lightllm.common.attention.backends._workspace import AttentionWorkspace from lightllm.models.glm5_next.triton_kernel.kpool import expand_topk, prefill_nsa_params - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -27,231 +13,91 @@ def tensor(values): @pytest.mark.parametrize("counts", [(3, 0, 1, 2), (257, 0, 1, 513), (1,) * 129]) -def test_prefill_metadata_reuses_existing_prefixes_during_graph_replay(counts): +def test_prefill_metadata_is_graph_safe(counts): batch, queries = len(counts), sum(counts) positions = torch.empty(queries, device="cuda", dtype=torch.int32) req = torch.empty(batch, device="cuda", dtype=torch.int32) - cu_q, cu_kv = torch.empty(batch + 1, device="cuda", dtype=torch.int32), torch.empty( - batch + 1, device="cuda", dtype=torch.int32 - ) - - def update(counts, phase): - q_prefix, kv_prefix, visible, starts, requests = [0], [0], [], [], [] - req_values = [(batch - i + phase) % (batch + 1) for i in range(batch)] - for i, count in enumerate(counts): - prefix = 1000000 + i * 7 + phase - visible.extend(range(prefix + 1, prefix + count + 1)) - starts.extend([kv_prefix[-1]] * count) - requests.extend([req_values[i]] * count) - q_prefix.append(q_prefix[-1] + count) - kv_prefix.append(kv_prefix[-1] + prefix + count) - positions.copy_(tensor(visible) - 1) - req.copy_(tensor(req_values)) - cu_q.copy_(tensor(q_prefix)) - cu_kv.copy_(tensor(kv_prefix)) - return tensor(visible), tensor(requests), tensor(starts) + cu_q = torch.empty(batch + 1, device="cuda", dtype=torch.int32) + cu_kv = torch.empty(batch + 1, device="cuda", dtype=torch.int32) + + q_prefix, kv_prefix, visible, starts, requests = [0], [0], [], [], [] + for index, count in enumerate(counts): + prefix = 1000000 + index * 7 + visible.extend(range(prefix + 1, prefix + count + 1)) + starts.extend([kv_prefix[-1]] * count) + requests.extend([index] * count) + q_prefix.append(q_prefix[-1] + count) + kv_prefix.append(kv_prefix[-1] + prefix + count) + positions.copy_(tensor(visible) - 1) + req.copy_(tensor(range(batch))) + cu_q.copy_(tensor(q_prefix)) + cu_kv.copy_(tensor(kv_prefix)) - expected = update(counts, 0) + expected = (tensor(visible), tensor(requests), tensor(starts)) actual = prefill_nsa_params(positions, req, cu_q, cu_kv) for observed, reference in zip(actual, expected): torch.testing.assert_close(observed, reference, rtol=0, atol=0) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = prefill_nsa_params(positions, req, cu_q, cu_kv) - expected = update(counts[::-1], 11) - graph.replay() - for observed, reference in zip(actual, expected): - torch.testing.assert_close(observed, reference, rtol=0, atol=0) @pytest.mark.parametrize("dense", [False, True]) -def test_topk_direct_request_lookup_preserves_mtp_hold_and_strided_tables(dense): +def test_topk_request_lookup_preserves_strided_tables(dense): table = torch.arange(8 * 70, device="cuda", dtype=torch.int32).reshape(8, 70)[::2, ::2] req = tensor([2, 2, 0, 3, 3]) lengths = tensor([5, 7, 9, 2, 2] if dense else [17, 18, 9, 2, 2]) groups = tensor([[3, 1], [1, 0], [0, 1], [0, 0], [0, 0]]) - - def check(actual): - mem, relative = actual - expected_mem, expected_relative = torch.full_like(mem, -1), torch.full_like(relative, -1) - for i, (request, length) in enumerate(zip(req.tolist(), lengths.tolist())): - tokens = ( - list(range(length)) - if dense - else [group * 4 + j for group in groups[i].tolist()[: min(length // 4, 2)] for j in range(4)] - + list(range(length // 4 * 4, length)) - ) - expected_relative[i, : len(tokens)] = tensor(tokens) - expected_mem[i, : len(tokens)] = table[request, tensor(tokens).long()] - torch.testing.assert_close(mem, expected_mem, rtol=0, atol=0) - torch.testing.assert_close(relative, expected_relative, rtol=0, atol=0) - - check(expand_topk(groups, lengths, req, table, 8, dense=dense)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = expand_topk(groups, lengths, req, table, 8, dense=dense) - req.copy_(tensor([0, 0, 2, 3, 3])) - table.add_(10000) - lengths[:2].add_(1) - graph.replay() - check(actual) - - -@pytest.mark.parametrize("fp8", [False, True]) -def test_decode_aliases_infer_state_without_ragged_allocation(monkeypatch, fp8): - import lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke as legacy - - def forbidden(*args, **kwargs): - raise AssertionError("GLM must not generate or copy a full ragged index") - - monkeypatch.setattr(legacy, "gen_nsa_ks_ke", forbidden) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - model = type( - "Model", - (), - { - "config": {"index_kpool": 4, "index_topk": 2048}, - "batch_max_tokens": 4, - "max_req_num": 1, - "max_seq_length": 32, - "att_state_role": "decode_test", - "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="decode"), - }, - )() - backend = (Fp8Glm5NextSparseAttBackend if fp8 else Glm5NextSparseAttBackend)(model) - assert not hasattr(backend, "ragged_mem_buffers") - infer = InferStateInfo() - infer.b_req_idx = tensor([2, 2, 0, 3, 3]) - infer.b_seq_len = tensor([999999, 1000000, 65536, 2, 2]) - cls = Fp8Glm5NextSparseDecodeState if fp8 else Glm5NextSparseDecodeState - state = cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - state.init_state() - assert state.q_kv_seq_len is infer.b_seq_len - assert state.q_req_idx is infer.b_req_idx - assert not hasattr(state, "ragged_mem_index") and not hasattr(state, "ks") - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - state.init_state() - infer.b_seq_len[:3].sub_(100) - infer.b_req_idx[:3].copy_(tensor([0, 0, 1])) - graph.replay() - assert state.q_kv_seq_len is infer.b_seq_len and state.q_req_idx is infer.b_req_idx + actual_mem, actual_relative = expand_topk(groups, lengths, req, table, 8, dense=dense) + + expected_mem, expected_relative = torch.full_like(actual_mem, -1), torch.full_like(actual_relative, -1) + for row, (request, length) in enumerate(zip(req.tolist(), lengths.tolist())): + tokens = ( + list(range(length)) + if dense + else [group * 4 + offset for group in groups[row].tolist()[: min(length // 4, 2)] for offset in range(4)] + + list(range(length // 4 * 4, length)) + ) + expected_relative[row, : len(tokens)] = tensor(tokens) + expected_mem[row, : len(tokens)] = table[request, tensor(tokens).long()] + torch.testing.assert_close(actual_mem, expected_mem, rtol=0, atol=0) + torch.testing.assert_close(actual_relative, expected_relative, rtol=0, atol=0) @pytest.mark.parametrize( - "queries,heads,total_kv,chunk_size", + "reference", [ - (1, 64, 65536, 2097152), - (2, 64, 65536, 2097152), - (3, 8, 65536, 2097152), - (9, 64, 512, 2097152), - (9, 64, 131073, 2097152), - (4, 16, 40, 3), - (4, 64, 40, 7), + "backends.sparse_mla.flashmla.fp8:NsaFlashMlaFp8SparseAttBackend", + "lightllm.models.glm5_next.attention.sparse_mla:Glm5NextSparseAttBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:Fp8Glm5NextSparseAttBackend", ], ) -def test_fp8_prefill_reads_request_table_without_ragged_buffer(monkeypatch, queries, heads, total_kv, chunk_size): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) - - torch.manual_seed(1685) - batch = min(queries, 2) - counts = [queries] if batch == 1 else [queries // 2, queries - queries // 2] - seqs = [total_kv] if batch == 1 else [total_kv // 2, total_kv - total_kv // 2] - req = tensor([2, 0][:batch]) - table = torch.full((4, 2 * (max(seqs) + 7)), -1, device="cuda", dtype=torch.int32)[:, ::2] - physical = torch.randperm(total_kv + 32, device="cuda", dtype=torch.int32)[:total_kv] - offset = 0 - for request, length in zip(req.tolist(), seqs): - table[request, :length] = physical[offset : offset + length] - offset += length - q = torch.randn(queries, heads, 512, device="cuda", dtype=torch.bfloat16) - source = torch.randn(total_kv + 32, 1, 512, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros(total_kv + 32, 1, 720, device="cuda", dtype=torch.uint8) - write_fp8_e8m0_packed_kv(source, torch.arange(total_kv + 32, device="cuda", dtype=torch.int32), packed) - prefill_kv = torch.randn(queries, 1, 512, device="cuda", dtype=torch.bfloat16) - visible, query_req, current_positions, q_prefix, kv_prefix = [], [], [], [0], [0] - for request, length, count in zip(req.tolist(), seqs, counts): - visible.extend(range(length - count + 1, length + 1)) - query_req.extend([request] * count) - current_positions.extend(range(length - count, length)) - q_prefix.append(q_prefix[-1] + count) - kv_prefix.append(kv_prefix[-1] + length) - query_req, current_positions = tensor(query_req), tensor(current_positions) - mem_index = table[query_req.long(), current_positions.long()] - topk = torch.full((queries, 2176), -1, device="cuda", dtype=torch.int32) - for i, length in enumerate(visible): - topk[i, :7] = tensor([0, length - 1, 0, 1, length // 2, length - 1, -1]) - topk_mem = torch.empty_like(topk) - - def refresh_physical(): - topk_mem.copy_(table[query_req.long()[:, None], topk.clamp_min(0).long()]) - topk_mem.masked_fill_(topk < 0, -1) - mem_index.copy_(table[query_req.long(), current_positions.long()]) - - refresh_physical() - infer = InferStateInfo() - infer.mem_index, infer.b_req_idx = mem_index, req - infer.position_ids = tensor(visible) - 1 - infer.b1_cu_q_seq_len, infer.b1_cu_kv_seq_len = tensor(q_prefix), tensor(kv_prefix) - infer.total_token_num = total_kv - infer.req_manager = SimpleNamespace(req_to_token_indexs=table) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - model = type( - "Model", - (), - { - "batch_max_tokens": queries, - "prefill_graph": None, - "max_req_num": batch, - "max_seq_length": max(seqs), - "att_state_role": "indices_test", - "config": {"index_topk": 2048}, - "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="normal"), - }, - )() - backend = Fp8Glm5NextSparseAttBackend(model) - state = backend.create_att_prefill_state(infer) - state.init_state() - assert not hasattr(state, "ragged_mem_index") and not hasattr(backend, "ragged_mem_buffers") - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk, - "topk_mem_indices": topk_mem, - "prefill_cache_kv": prefill_kv, - "softmax_scale": 512 ** -0.5, - "kv_lora_rank": 512, - }, +def test_sparse_context_owns_query_metadata(reference): + from lightllm.common.attention.registry import load_backend_class + + lengths = torch.tensor([5, 4], device="cuda", dtype=torch.int32) + q_lengths = torch.tensor([2, 1], device="cuda", dtype=torch.int32) + batch = AttentionBatch( + phase=AttentionPhase.PREFILL, + batch_size=2, + token_num=3, + total_token_num=9, + max_q_seq_len=2, + max_kv_seq_len=5, + max_cache_len=3, + b_req_idx=tensor([0, 1]), + b_seq_len=lengths, + b_q_seq_len=q_lengths, + b1_cu_q_seq_len=tensor([0, 2, 3]), + b1_cu_kv_seq_len=tensor([0, 5, 9]), + req_to_token_indexs=torch.arange(8 * 16, device="cuda", dtype=torch.int32).reshape(8, 16), + hold_req_id=7, + b_ready_cache_len=tensor([3, 3]), + position_ids=tensor([3, 4, 3]), ) - - def run(): - return state.prefill_att(q, packed[..., :584], None, control) - - def check(output): - scales = packed[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) - values = packed[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) * scales - kv = torch.cat((values.reshape(-1, 1, 448).bfloat16(), packed[..., 448:576].view(torch.bfloat16)), -1) - kv[mem_index.long()] = prefill_kv - expected = [] - for i in range(queries): - keys = kv[topk_mem[i][topk_mem[i] >= 0].long(), 0].float() - expected.append((q[i].float() @ keys.T * 512 ** -0.5).softmax(-1) @ keys) - torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) - - check(run()) - workspace = backend._prefill_workspaces[0] - if workspace is not None: - assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = run() - graph.replay() - check(output) - q.normal_() - prefill_kv.normal_() - for request, length in zip(req.tolist(), seqs): - table[request, :length].copy_(table[request, :length].roll(1)) - topk[0].fill_(-1) - refresh_physical() - graph.replay() - check(output) + backend = object.__new__(load_backend_class(reference)) + context = backend.allocate_context( + ContextCapacity(AttentionPhase.PREFILL, 2, 3, 5, 2), capture=False, workspace=AttentionWorkspace("cuda") + ) + backend.prepare(context, batch) + backend.update(context) + torch.testing.assert_close(context.q_kv_seq_len, tensor([4, 5, 4])) + torch.testing.assert_close(context.q_req_idx, tensor([0, 0, 1])) + torch.testing.assert_close(context.q_kv_start_loc, tensor([0, 0, 5])) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index a6a5c7c9f2..32b0cb2b33 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -5,35 +5,33 @@ import torch import triton -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.nsa.glm5_next import ( - Glm5NextSparsePrefillState, - Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - Fp8Glm5NextSparsePrefillState, - Fp8Glm5NextSparseDecodeState, +from lightllm.models.glm5_next.attention.fp8_sparse_mla import Fp8Glm5NextSparseAttBackend +from lightllm.models.glm5_next.attention.sparse_mla import Glm5NextSparseAttBackend +from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill +from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + _build_prefill_row_table_kernel, + _gather_prefill_kv_kernel, ) from lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode import ( can_use_fp8_sparse_direct_decode, fp8_sparse_direct_decode, ) from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel -from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next -from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( - _gather_prefill_kv_kernel, +from lightllm.utils.backend_validator import ProbeResult, ProbeScope +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, ) - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture(autouse=True) -def current_cuda_device(monkeypatch): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) +pytest.importorskip("flash_mla") def _make_infer_state(**attributes): @@ -42,156 +40,59 @@ def _make_infer_state(**attributes): return state -def _make_workspace_backend(batch_max_tokens, overlap=False): +def _make_workspace_backend(batch_max_tokens): backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace( - batch_max_tokens=batch_max_tokens, - prefill_graph=None, - config={"index_topk": 2048, "index_kpool": 4}, - max_req_num=1, - max_seq_length=fp8_glm5_next.KV_GATHER_CHUNK_SIZE, - att_state_role="unit_test", - args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode="normal"), - ) - backend._prefill_workspaces = backend.allocate_prefill_workspaces(backend.model) - return backend - + backend.prefill_token_capacity = batch_max_tokens + backend.prefill_graph_enabled = False + from lightllm.common.attention.backends._workspace import AttentionWorkspace -class _PrefillCase: - def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=None, softmax_scale=0.0625): - device = "cuda" - tensor = lambda values: torch.tensor(values, device=device, dtype=torch.int32) - self.seq_lens, self.q_lens = seq_lens, q_lens - self.requests = list(range(2 * len(seq_lens) - 1, 0, -2)) - total, queries = sum(seq_lens), sum(q_lens) - capacity = capacity or total + 17 - self.req_to_token_indexs = torch.full( - (2 * len(seq_lens) + 1, max(seq_lens) + 7), -1, device=device, dtype=torch.int32 - ) - slots = torch.randperm(capacity, device=device, dtype=torch.int32)[:total] - cu_q, cu_kv, q_req_idx, positions = [0], [0], [], [] - for req, length, count in zip(self.requests, seq_lens, q_lens): - self.req_to_token_indexs[req, :length] = slots[cu_kv[-1] : cu_kv[-1] + length] - q_req_idx.extend([req] * count) - positions.extend(range(length - count, length)) - cu_q.append(cu_q[-1] + count) - cu_kv.append(cu_kv[-1] + length) - self.q_req_idx, self.position_ids = tensor(q_req_idx).long(), tensor(positions).long() - self.q = torch.randn(queries, heads, 512, device=device, dtype=torch.bfloat16) - self.source_kv = torch.randn(capacity, 1, 512, device=device, dtype=torch.bfloat16) - self.packed_kv = torch.zeros(capacity, 1, 720, device=device, dtype=torch.uint8) - self.all_slots = torch.arange(capacity, device=device, dtype=torch.int32) - write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) - self.prefill_cache_kv = torch.randn(queries, 1, 1024, device=device, dtype=torch.bfloat16)[..., ::2] - visible = tensor(positions) + 1 - self.topk_indices = torch.randint(max(seq_lens), (queries, width), device=device, dtype=torch.int32) - self.topk_indices.remainder_(visible[:, None]) - self.topk_indices.masked_fill_( - torch.arange(width, device=device)[None, :] >= visible.clamp_max(2048)[:, None], -1 - ) - self.topk_indices[:, ::19] = -1 - self.topk_indices[:, 0] = visible - 1 - self.topk_indices[:, 1] = visible - 1 - self.topk_mem_indices = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] - self.infer_state = _make_infer_state( - is_prefill=True, - req_manager=SimpleNamespace(req_to_token_indexs=self.req_to_token_indexs), - b_req_idx=tensor(self.requests), - b_seq_len=tensor(seq_lens), - b1_cu_q_seq_len=tensor(cu_q), - b1_cu_kv_seq_len=tensor(cu_kv), - position_ids=tensor(positions), - total_token_num=total, - mem_index=torch.empty(queries, device=device, dtype=torch.int32), - ) - self.backend = backend or _make_workspace_backend(queries) - _, self.att_state = self.new_state() - self.control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": self.topk_indices, - "topk_mem_indices": self.topk_mem_indices, - "prefill_cache_kv": self.prefill_cache_kv, - "softmax_scale": softmax_scale, - "kv_lora_rank": 512, - }, - ) - - def new_state(self, **overrides): - if overrides: - infer = _make_infer_state(**vars(self.infer_state)) - vars(infer).update(overrides) - else: - infer = self.infer_state - state = self.backend.create_att_prefill_state(infer) - state.init_state() - return infer, state - - def refresh(self): - self.infer_state.mem_index.copy_(self.req_to_token_indexs[self.q_req_idx, self.position_ids]) - self.topk_mem_indices.copy_( - self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()] - ) - self.topk_mem_indices.masked_fill_(self.topk_indices < 0, -1) - self.control.nsa_prefill_dict["topk_indices"] = self.topk_indices - - def run(self, state=None): - self.refresh() - active = self.att_state if state is None else state - return active.prefill_att(self.q, self.packed_kv[..., :584], None, self.control) - - def check(self, actual): - values = _dequantize_kv_cache(self.packed_kv) - values[self.infer_state.mem_index.long()] = self.prefill_cache_kv - physical = self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()].masked_fill( - self.topk_indices < 0, -1 - ) - expected = _torch_sparse_attention(self.q, values, physical, self.control.nsa_prefill_dict["softmax_scale"]) - torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) - - def mutate(self): - self.q.normal_() - self.prefill_cache_kv.normal_() - self.source_kv.normal_() - write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) - for req, length in zip(self.requests, self.seq_lens): - self.req_to_token_indexs[req, :length].copy_(self.req_to_token_indexs[req, :length].roll(1)) - self.topk_indices.copy_(self.topk_indices.roll(1, dims=1)) + backend._test_workspace = AttentionWorkspace("cuda") + return backend @pytest.mark.parametrize("draft", [False, True]) @pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) def test_main_and_native_mtp_backend_selection(monkeypatch, draft, kv_type): + from lightllm.common.attention.api import AttentionFamily, AttentionPhase, CacheLayout + from lightllm.common.basemodel import attention from lightllm.models.glm5_next.model import Glm5NextTpPartModel from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + from unit_tests.common.attention.helpers import make_spec, model_stub - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") model = object.__new__(Glm5NextMTPModel if draft else Glm5NextTpPartModel) - model.args = SimpleNamespace(llm_kv_type=kv_type, enable_prefill_microbatch_overlap=False, run_mode="normal") - model.config = {"index_topk": 2048, "index_kpool": 4} - model.graph_max_batch_size = 4 - model.max_seq_length = 128 - model.max_req_num = 2 - model.batch_max_tokens = 4 + vars(model).update( + vars( + model_stub( + llm_kv_type=kv_type, llm_prefill_att_backend="flashmla_sparse", llm_decode_att_backend="flashmla_sparse" + ) + ) + ) + model.is_mtp_draft_model = draft model.mtp_previous_draft_models = [] - model._init_att_backend() - infer_state = InferStateInfo() - prefill = model.prefill_att_backend.create_att_prefill_state(infer_state) - decode = model.decode_att_backend.create_att_decode_state(infer_state) - assert type(prefill) is (Fp8Glm5NextSparsePrefillState if kv_type == "fp8kv_dsa" else Glm5NextSparsePrefillState) - assert type(decode) is (Fp8Glm5NextSparseDecodeState if kv_type == "fp8kv_dsa" else Glm5NextSparseDecodeState) - infer_state.prefill_att_state = prefill - infer_state.decode_att_state = decode - infer_state.check_att_states() - assert prefill.infer_state is infer_state - assert decode.infer_state is infer_state - owner_ref = weakref.ref(infer_state) - del infer_state - assert owner_ref() is None - with pytest.raises(ReferenceError): - prefill.infer_state - with pytest.raises(ReferenceError): - decode.infer_state + model.config = {"index_topk": 2048, "index_kpool": 4} + model.batch_max_tokens = 8 + expected = Fp8Glm5NextSparseAttBackend if kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + monkeypatch.setattr( + "lightllm.utils.backend_validator.probe_backend", + lambda *args, **kwargs: ProbeResult(True, ProbeScope.BACKEND_NUMERICS), + ) + monkeypatch.setattr(expected, "check_environment", classmethod(lambda cls, caps: None)) + monkeypatch.setattr( + attention, + "describe_attention_layer", + lambda model, layer, family, phase: make_spec( + family, phase, cache=CacheLayout(kv_type, 1, torch.bfloat16, 1, 512, 512) + ), + ) + model.layers_infer = [SimpleNamespace(get_attention_family=lambda: AttentionFamily.SPARSE_MLA)] + model.prepare_attention_backends() + for phase in AttentionPhase: + backend = model.attention_groups[0].backend_for(phase) + assert type(backend) is expected + context = backend_context(type(backend), phase.value, settings=backend) + assert context.backend_identity is backend + assert not hasattr(backend, "model") + assert not hasattr(context, "request") def _dequantize_kv_cache(packed_kv): @@ -258,29 +159,26 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid changed_values = (changed_groups / changed_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) assert torch.equal(packed_kv[mem_index, :, :448], changed_values.view(torch.uint8).reshape(11, 1, 448)) + indices = mem_index[[0, 2, 8, 2, 5, 1, 9, 3]] prefill_kv = source[8:].clone().add_(0.25) - gathered_kv = torch.empty((8, 1, 512), device="cuda", dtype=torch.bfloat16) - _gather_prefill_kv_kernel[(3,)]( + row_table = torch.full((31,), -1, device="cuda", dtype=torch.int32) + _build_prefill_row_table_kernel[(1,)](mem_index[8:], row_table, 3, BLOCK=256) + gathered_kv = torch.empty((indices.numel(), 1, 512), device="cuda", dtype=torch.bfloat16) + _gather_prefill_kv_kernel[(triton.cdiv(indices.numel(), 4),)]( packed_kv[..., :584], + indices, + row_table, prefill_kv, gathered_kv, - mem_index[None], - torch.zeros(1, device="cuda", dtype=torch.int32), - torch.tensor([0, 11], device="cuda", dtype=torch.int32), - torch.tensor([0, 3], device="cuda", dtype=torch.int32), - 2, - 8, + indices.numel(), packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - req_to_token_indexs_stride_b=11, - req_to_token_indexs_stride_s=1, - batch_size=1, BLOCK_ROWS=4, ) expected = _dequantize_kv_cache(packed_kv) - expected[mem_index[8:].long()] = prefill_kv - assert torch.equal(gathered_kv, expected[mem_index[2:10].long()]) + expected[mem_index[8:]] = prefill_kv + assert torch.equal(gathered_kv, expected[indices.long()]) @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @@ -300,16 +198,12 @@ def test_fp8_nope_decode_and_cuda_graph(num_heads): indices = torch.randint(0, 4096, (9, 2176), device="cuda", dtype=torch.int32) lanes = torch.arange(2176, device="cuda") indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} infer_state = _make_infer_state() - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] - ) + state = backend_context(Fp8Glm5NextSparseAttBackend, "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0]) def run(): - return state.decode_att((q, q[..., :0]), packed_kv[..., :584], None, control) + return run_sparse(state, infer_state, (q, q[..., :0]), packed_kv[..., :584], control) def check(actual): expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed_kv), indices) @@ -333,199 +227,444 @@ def check(actual): @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @pytest.mark.parametrize("chunk_size", [3, 7, 32768]) def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chunk_size): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) torch.manual_seed(53) - prefix = 4 if cached else 0 - case = _PrefillCase([prefix + 3] * 2, [3, 3], num_heads, width=128, capacity=64, softmax_scale=512 ** -0.5) - case.q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - case.topk_indices[0].fill_(-1) - case.topk_indices[1, :2] = prefix - case.topk_indices[3].fill_(-1) - case.topk_indices[3, 0] = prefix - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) - - -@pytest.mark.parametrize("num_heads", [8, 64]) -@pytest.mark.parametrize("chunk_size", [3, 32768]) -def test_fp8_uncached_prefill_reuses_current_kv(monkeypatch, num_heads, chunk_size): - from sgl_kernel.flash_attn import flash_attn_with_kvcache + # Two requests whose physical cache locations are unrelated to batch order. + prefix_len = 4 if cached else 0 + seq_len = prefix_len + 3 + ragged_mem_index = torch.randperm(64, device="cuda")[: 2 * seq_len].int() + prefill_rows = torch.cat( + (torch.arange(prefix_len, seq_len), torch.arange(seq_len + prefix_len, 2 * seq_len)) + ).cuda() + prefill_mem_index = ragged_mem_index[prefill_rows] + source = torch.randn(2 * seq_len, 1, 512, device="cuda", dtype=torch.bfloat16) + prefill_kv = source[prefill_rows].clone().add_(0.25) + packed_kv = torch.zeros(64, 1, 720, device="cuda", dtype=torch.uint8) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + kv_start_offsets = torch.tensor([0] * 3 + [seq_len] * 3, device="cuda", dtype=torch.int32) + topk_indices = torch.full((6, 128), -1, device="cuda", dtype=torch.int32) + for row in range(6): + length = prefix_len + row % 3 + 1 + topk_indices[row, :length] = torch.arange(length, device="cuda") + topk_indices[0].fill_(-1) + topk_indices[1, :2] = prefix_len + topk_indices[3].fill_(-1) + topk_indices[3, 0] = prefix_len + control = { + "topk_indices": topk_indices, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 512 ** (-0.5), + "kv_lora_rank": 512, + } + infer_state = _make_infer_state(max_cache_len=prefix_len, mem_index=prefill_mem_index) + state = backend_context( + Fp8Glm5NextSparseAttBackend, + "prefill", + settings=_make_workspace_backend(q.shape[0]), + ks=kv_start_offsets, + ragged_mem_index=ragged_mem_index, + ) - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) - case = _PrefillCase([3, 3], [3, 3], num_heads, width=128) - case.prefill_cache_kv = torch.randn(6, 1, 1536, device="cuda", dtype=torch.bfloat16)[..., :512] - case.control.nsa_prefill_dict["prefill_cache_kv"] = case.prefill_cache_kv + def run(): + return run_sparse(state, infer_state, q, packed_kv[..., :584], control) - def check_kv_storage(*args, **kwargs): - assert kwargs["v_cache"].untyped_storage().data_ptr() == case.prefill_cache_kv.untyped_storage().data_ptr() - return flash_attn_with_kvcache(*args, **kwargs) + def check(out): + expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] + expected_kv[prefill_rows] = prefill_kv + indices = torch.where(topk_indices >= 0, topk_indices + kv_start_offsets[:, None], -1) + expected = _torch_sparse_attention(q, expected_kv, indices, control["softmax_scale"]) + torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) - monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", check_kv_storage) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - case.mutate() + check(run()) + graph, out = _capture_cuda_graph(run) graph.replay() - case.check(output) + check(out) + q.normal_() + source.normal_() + prefill_kv.copy_(source[prefill_rows] + 0.25) + ragged_mem_index.copy_(ragged_mem_index.roll(1)) + prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) + topk_indices.copy_(topk_indices.roll(1, dims=1)) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + graph.replay() + check(out) @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) -@pytest.mark.parametrize("softmax_scale", [0.0625, 512 ** -0.5]) +@pytest.mark.parametrize("softmax_scale", [0.0625, 512**-0.5]) def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 129) + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", 129) torch.manual_seed(53) - case = _PrefillCase([1025], [9], num_heads, capacity=2048, softmax_scale=softmax_scale) + num_tokens = 1025 + source = torch.randn(num_tokens, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + ragged_mem_index = torch.randperm(2048, device="cuda")[:num_tokens].int() + prefill_rows = torch.tensor([1, 15, 128, 129, 255, 257, 511, 512, 1024], device="cuda") + prefill_mem_index = ragged_mem_index[prefill_rows].long() + prefill_kv = torch.randn(9, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + packed_kv = torch.zeros(2048, 1, 720, device="cuda", dtype=torch.uint8) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") - case.topk_indices.masked_fill_(torch.arange(2176, device="cuda")[None, :] >= lengths[:, None], -1) - case.topk_indices[1, 0] = 1016 - case.topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") - case.topk_indices[3, :127] = torch.randint(1016, 1025, (127,), device="cuda", dtype=torch.int32) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) + topk_indices = torch.randint(num_tokens, (9, 2176), device="cuda", dtype=torch.int32) + lanes = torch.arange(2176, device="cuda") + topk_indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) + topk_indices[1, 0] = 15 + topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") + topk_indices[3, :127] = prefill_rows[torch.randint(9, (127,), device="cuda")].int() + control = { + "topk_indices": topk_indices, + "prefill_cache_kv": prefill_kv, + "softmax_scale": softmax_scale, + "kv_lora_rank": 512, + } + infer_state = _make_infer_state(mem_index=prefill_mem_index) + state = backend_context( + Fp8Glm5NextSparseAttBackend, + "prefill", + settings=_make_workspace_backend(q.shape[0]), + ks=torch.zeros(9, device="cuda", dtype=torch.int32), + ragged_mem_index=ragged_mem_index, + ) + def run(): + return run_sparse(state, infer_state, q, packed_kv[..., :584], control) -@pytest.mark.parametrize("num_q", [64, 129]) -@pytest.mark.parametrize("num_heads", [8, 64]) -def test_fp8_prefill_uses_flashattention_and_request_table_graph(monkeypatch, num_q, num_heads): - from sgl_kernel.flash_attn import flash_attn_with_kvcache + def check(actual): + expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] + expected_kv[prefill_rows] = prefill_kv + expected = _torch_sparse_attention(q, expected_kv, topk_indices, softmax_scale) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + check(run()) + graph, output = _capture_cuda_graph(run) + graph.replay() + check(output) + q.normal_() + source.normal_() + prefill_kv.normal_() + ragged_mem_index.copy_(ragged_mem_index.roll(1)) + prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + topk_indices.copy_(topk_indices.roll(1, dims=1)) + graph.replay() + check(output) + + +def _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples): + """Reference only a few H64 rows so the large-Q regression stays bounded.""" + expected = [] + for row, heads in samples: + local = topk[row] + local = local[local >= 0].long() + physical = ragged[(starts[row].long() + local)].long() + values = _dequantize_kv_cache(packed[physical])[:, 0].float() + current = physical[:, None] == prefill_mem[None, :].long() + has_current = current.any(dim=1) + current_row = current.int().argmax(dim=1) + values[has_current] = prefill_kv[current_row[has_current], 0].float() + expected.append(((q[row, heads].float() @ values.T * 0.0625).softmax(-1) @ values)) + return expected + + +@pytest.mark.parametrize("num_q", [16, 33]) +def test_fp8_h64_largeq_native_prefill_permuted_ragged_graph(monkeypatch, num_q): + """The native sparse prefill path maps top-k ragged positions, never physical IDs.""" + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse prefill target is Hopper") + from sgl_kernel.flash_mla import flash_mla_sparse_fwd torch.manual_seed(53100 + num_q) - case = _PrefillCase([32768, 32768], [num_q // 2, num_q - num_q // 2], num_heads, capacity=131072) + capacity = 131072 + ragged_n = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32)[:ragged_n] + q = torch.randn(num_q, 64, 512, device="cuda", dtype=torch.bfloat16) + starts = torch.where( + torch.arange(num_q, device="cuda", dtype=torch.int32) < num_q // 2, + torch.full((num_q,), 1024, device="cuda", dtype=torch.int32), + torch.full((num_q,), 32768, device="cuda", dtype=torch.int32), + ) + current_local = (torch.arange(num_q, device="cuda", dtype=torch.int32) * 37) % 2048 + 64 + current_positions = starts + current_local + prefill_mem = ragged[current_positions.long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + prefill_kv.add_(torch.arange(num_q, device="cuda", dtype=torch.bfloat16)[:, None, None]) + topk = torch.randint(0, 2176, (num_q, 2176), device="cuda", dtype=torch.int32) + topk[:, ::19] = -1 # Holes must survive the native position-preserving map. + topk[:, 0] = current_local + topk[:, 1] = current_local # Duplicates are valid selected keys. + topk[:, 2] = topk[:, 3] + assert q.is_contiguous() and topk.is_contiguous() + observed = [] - def observed_flashattention(*args, **kwargs): - observed.append(kwargs["qv"].shape) - return flash_attn_with_kvcache(*args, **kwargs) - - def unexpected_flashmla(*args, **kwargs): - raise AssertionError("FP8 prefill must use its single FlashAttention backend") - - monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", observed_flashattention) - monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", unexpected_flashmla) - case.check(case.run()) - assert observed == [(num_q, num_heads, 512)] - graph, output = _capture_cuda_graph(case.run) - case.mutate() + def observed_sparse(*args, **kwargs): + observed.append((kwargs["q"].shape, kwargs["kv"].shape, kwargs["indices"].shape)) + return flash_mla_sparse_fwd(*args, **kwargs) + + monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", observed_sparse) + + def run(): + return fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + + samples = [(0, torch.tensor([0, 9, 63], device="cuda")), (num_q - 1, torch.tensor([3, 31, 62], device="cuda"))] + + def check(actual): + expected = _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples) + for (row, heads), value in zip(samples, expected): + torch.testing.assert_close(actual[row, heads].float(), value, atol=0.012, rtol=0.015) + + eager = run() + check(eager) + assert observed == [((num_q, 64, 512), (ragged_n, 1, 512), (num_q, 1, 2176))] + graph, output = _capture_cuda_graph(run) + q.normal_() + prefill_kv.normal_() + ragged.copy_(ragged.roll(1)) + prefill_mem.copy_(ragged[current_positions.long()]) + topk.copy_(topk.roll(1, dims=1)) + topk[:, 1] = current_local graph.replay() - case.check(output) - case.topk_indices.fill_(-1) + check(output) + topk.fill_(-1) graph.replay() assert torch.count_nonzero(output) == 0 + assert len(observed) == 3 # eager, graph warmup, and capture; replay stays inside the CUDA graph. + + +def test_fp8_h8_largeq_prefill_keeps_flashattention_path(monkeypatch): + """H8 large-Q prefill must not enter the H64 native sparse fast path.""" + import sgl_kernel.flash_attn as flash_attn + import sgl_kernel.flash_mla as flash_mla + + torch.manual_seed(53180) + capacity = 65536 + num_q = 16 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + q = torch.randn(num_q, 8, 512, device="cuda", dtype=torch.bfloat16) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) + starts = torch.full((num_q,), 1024, device="cuda", dtype=torch.int32) + current_local = torch.arange(num_q, device="cuda", dtype=torch.int32) + 64 + prefill_mem = ragged[(starts + current_local).long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :64] = torch.arange(64, device="cuda", dtype=torch.int32) + topk[:, 0] = current_local + seen_flashattention = [] + original_flashattention = flash_attn.flash_attn_with_kvcache + + def observed_flashattention(*args, **kwargs): + seen_flashattention.append(True) + return original_flashattention(*args, **kwargs) + + def unexpected_native(*args, **kwargs): + raise AssertionError("H8 large-Q prefill entered the native H64 sparse path") + + monkeypatch.setattr(flash_attn, "flash_attn_with_kvcache", observed_flashattention) + monkeypatch.setattr(flash_mla, "flash_mla_sparse_fwd", unexpected_native) + output = fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + assert seen_flashattention + assert output.shape == q.shape @pytest.mark.parametrize( "num_q,num_heads", [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], ) -def test_fp8_cached_prefill_strided_holes_and_current(num_q, num_heads): +def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num_heads): + """Exercise eager and graph replay for no-rowtable, small-Q packed prefill.""" torch.manual_seed(53016 + num_q) - seqs = [65536 + value for value in [127, 521, 1031, 2053, 4099, 8191, 12289, 16381][:num_q]] - case = _PrefillCase(seqs, [1] * num_q, num_heads) - case.q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - backing = torch.full((num_q, 4352), -1, device="cuda", dtype=torch.int32) - backing[:, ::2] = case.topk_indices - case.topk_indices = backing[:, ::2] - output = case.run() - case.check(output) - assert output.is_contiguous() and torch.count_nonzero(output) > 0 - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + # Requests deliberately have different spans; starts are not fixed-stride. + seq_lens = torch.tensor([127, 521, 1031, 2053, 4099, 8191, 12289, 16381], device="cuda", dtype=torch.int32)[:num_q] + starts = torch.nn.functional.pad(seq_lens.cumsum(0), (1, 0))[:-1] + current_rows = starts + seq_lens - 1 + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(num_q, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + topk_backing = torch.full((num_q, 2176 * 2), -1, device="cuda", dtype=torch.int32) + topk = topk_backing[:, ::2] # Exercise the topk stride passed to the direct kernel. + topk[:, :3] = torch.tensor([3, 11, 11], device="cuda", dtype=torch.int32) + topk[:, 3] = seq_lens - 1 + + def run(): + return fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + def check(actual): + assert actual.shape == q.shape and actual.is_contiguous() + expected_kv = _dequantize_kv_cache(packed)[ragged.long()] + expected_kv[current_rows.long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + expected = _torch_sparse_attention(q, expected_kv, physical) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) -@pytest.mark.parametrize("num_q", [2, 3, 5, 8, 16, 32, 128]) -def test_fp8_cached_prefill_multi_query_and_graph(num_q): - torch.manual_seed(53200 + num_q) - case = _PrefillCase([131072, 131072], [num_q // 2, num_q - num_q // 2], 64) - case.check(case.run()) - for _ in range(2): - infer, state = case.new_state() - case.check(case.run(state)) - fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) - case.check(case.run(fresh)) - graph, output = _capture_cuda_graph(case.run) - case.mutate() + actual = run() + check(actual) + assert torch.count_nonzero(actual) > 0 # Includes BF16 current KV in every row. + graph, output = _capture_cuda_graph(run) graph.replay() - case.check(output) - case.topk_indices.fill_(-1) + check(output) + q.normal_() + source.normal_() + prefill_kv.normal_() + ragged.copy_(ragged.roll(1)) + prefill_mem.copy_(ragged[current_rows]) + write_fp8_e8m0_packed_kv(source, mem, packed) + topk.copy_(topk.roll(1, dims=1)) + graph.replay() + check(output) + topk.fill_(-1) # Empty graph replay must not read invalid packed pointers. graph.replay() assert torch.count_nonzero(output) == 0 -def test_fp8_prefill_graph_with_empty_requests_and_changing_boundaries(): - case = _PrefillCase([65536] * 4, [0, 4, 0, 3], 64) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - replacement = _PrefillCase([65536] * 4, [2, 0, 5, 0], 64) - for name in ( - "q", - "packed_kv", - "prefill_cache_kv", - "req_to_token_indexs", - "q_req_idx", - "position_ids", - "topk_indices", - "topk_mem_indices", - ): - getattr(case, name).copy_(getattr(replacement, name)) - case.infer_state.b1_cu_q_seq_len.copy_(replacement.infer_state.b1_cu_q_seq_len) - case.infer_state.position_ids.copy_(replacement.infer_state.position_ids) - for name in ("q_kv_seq_len", "q_req_idx", "q_kv_start_loc"): - getattr(case.att_state, name).copy_(getattr(replacement.att_state, name)) - graph.replay() - case.check(output) - - -def test_fp8_one_query_prefill_and_decode_read_their_own_kv_sources(): - case = _PrefillCase([65536], [1], 64) - case.q.zero_() - case.prefill_cache_kv.fill_(1) - case.source_kv.fill_(2) - write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) - case.topk_indices.fill_(-1) - case.topk_indices[:, 0] = 65535 - case.refresh() - decode_infer = _make_infer_state(**vars(case.infer_state)) - decode_infer.is_prefill = False - decode_infer.is_cuda_graph_capture = True - decode_state = case.backend.create_att_decode_state(decode_infer) - decode_state.init_state() - control = AttControl( - nsa_decode=True, - nsa_decode_dict={ - "topk_mem_indices": case.topk_mem_indices, - "softmax_scale": 0.0625, - "kv_lora_rank": 512, - }, +def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): + """Q1 H64 direct workspace is owned by the backend and reused across infer states.""" + torch.manual_seed(53150) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() + starts = torch.zeros(1, device="cuda", dtype=torch.int32) + current_rows = torch.tensor([2048], device="cuda", dtype=torch.int32) + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(1, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = current_rows + control = {"topk_indices": topk, "prefill_cache_kv": prefill_kv, "softmax_scale": 0.0625, "kv_lora_rank": 512} + infer_state = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(q.shape[0]) + state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged ) - def run(): - prefill = case.run() - decode = decode_state.decode_att((case.q, case.q[..., :0]), case.packed_kv[..., :584], None, control) - return prefill, decode - - prefill, decode = run() - torch.testing.assert_close(prefill, torch.ones_like(prefill), rtol=0, atol=0) - torch.testing.assert_close(decode, torch.full_like(decode, 2), rtol=0, atol=0) - graph, (prefill, decode) = _capture_cuda_graph(run) - case.prefill_cache_kv.fill_(3) - case.source_kv.fill_(4) - write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) + def expected(): + values = _dequantize_kv_cache(packed)[ragged.long()] + values[current_rows.long()] = prefill_kv + return _torch_sparse_attention(q, values, topk) + + def run(active_state=state, request=infer_state): + return run_sparse(active_state, request, q, packed[..., :584], control) + + torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) + cache = backend._test_workspace.get_object("glm_smallq", dict) + assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") + workspace = cache["direct"] + pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} + for _ in range(10): + layer_state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers + + fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) + fresh_state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(fresh_state, fresh_infer_state).float(), expected(), atol=0.012, rtol=0.015) + fresh_workspace = backend._test_workspace.get_object("glm_smallq", dict)["direct"] + assert fresh_workspace["partial_o"].data_ptr() == workspace["partial_o"].data_ptr() + + q.normal_() + source.normal_() + prefill_kv.normal_() + write_fp8_e8m0_packed_kv(source, mem, packed) + torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers + + +@pytest.mark.parametrize("num_q", [2, 3, 5, 8]) +def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): + """H64 selected prefill rewrites backend workspace across DSA layers.""" + torch.manual_seed(53200 + num_q) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 4096 + current_rows = starts + 2048 + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = 2048 + control = {"topk_indices": topk, "prefill_cache_kv": prefill_kv, "softmax_scale": 0.0625, "kv_lora_rank": 512} + infer_state = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(q.shape[0]) + state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + + def expected(): + values = _dequantize_kv_cache(packed)[ragged.long()] + values[current_rows.long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + return _torch_sparse_attention(q, values, physical) + + def run(active_state=state, request=infer_state): + return run_sparse(active_state, request, q, packed[..., :584], control) + + eager = run() + torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) + cache = backend._test_workspace.get_object("glm_smallq", dict) + assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") + workspace = cache["selected"] + pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} + # All 11 DSA layers share the backend allocation and rewrite its contents. + for _ in range(10): + layer_state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers + + fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) + fresh_state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(fresh_state, fresh_infer_state).float(), expected(), atol=0.012, rtol=0.015) + fresh_workspace = backend._test_workspace.get_object("glm_smallq", dict)["selected"] + assert fresh_workspace["compact"].data_ptr() == workspace["compact"].data_ptr() + + source.normal_() + write_fp8_e8m0_packed_kv(source, mem, packed) + prefill_kv.normal_() + q.normal_() + eager = run() + assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers + torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) + graph, output = _capture_cuda_graph(run) + q.normal_() + topk.copy_(topk.roll(1, dims=1)) + graph.replay() + torch.testing.assert_close(output.float(), expected(), atol=0.012, rtol=0.015) + topk.fill_(-1) graph.replay() - torch.testing.assert_close(prefill, torch.full_like(prefill, 3), rtol=0, atol=0) - torch.testing.assert_close(decode, torch.full_like(decode, 4), rtol=0, atol=0) + assert torch.count_nonzero(output) == 0 def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): @@ -542,18 +681,17 @@ def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) indices = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (1, 2049), device="cuda", dtype=torch.int32) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} infer = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer), + state = backend_context( + Fp8Glm5NextSparseAttBackend, + "decode", decode_topk_bound=2051, flashmla_sched_meta=flash_mla.get_mla_metadata()[0], ) def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + return run_sparse(state, infer, (q, q[..., :0]), packed[..., :584], control) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) eager = run() @@ -610,17 +748,12 @@ def test_fp8_decode_infer_state_graph_qpad_and_mutation(num_heads): q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (batch, 2049), device="cuda", dtype=torch.int32) - control = AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} infer_state = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] - ) + state = backend_context(Fp8Glm5NextSparseAttBackend, "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0]) def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + return run_sparse(state, infer_state, (q, q[..., :0]), packed[..., :584], control) def reference(): return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) @@ -643,13 +776,12 @@ def reference(): def test_fp8_sparse_direct_decode_factory_bound_guard(): backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) - infer_state = InferStateInfo() - assert backend.create_att_decode_state(infer_state).decode_topk_bound == 2051 - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 8}) - assert backend.create_att_decode_state(infer_state).decode_topk_bound is None - backend.model = SimpleNamespace(config={"index_topk": 2047, "index_kpool": 4}) - assert backend.create_att_decode_state(infer_state).decode_topk_bound is None + backend.network_config = {"index_topk": 2048, "index_kpool": 4} + assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound == 2051 + backend.network_config = {"index_topk": 2048, "index_kpool": 8} + assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound is None + backend.network_config = {"index_topk": 2047, "index_kpool": 4} + assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound is None @pytest.mark.parametrize("batch", [1, 8]) @@ -675,12 +807,11 @@ def test_fp8_sparse_direct_decode_actual_entry_edges_and_graph(monkeypatch, batc indices[5, :64] = -1 # holes indices[6, 64:128] = indices[6, :64] # duplicates assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} state_owner = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(state_owner), + state = backend_context( + Fp8Glm5NextSparseAttBackend, + "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0], decode_topk_bound=2051, ) @@ -689,7 +820,7 @@ def native_must_not_run(*args, **kwargs): raise AssertionError("supported GLM-5.3 decode took native FlashMLA fallback") def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + return run_sparse(state, state_owner, (q, q[..., :0]), packed[..., :584], control) assert not can_use_fp8_sparse_direct_decode(q, packed, indices.unsqueeze(1), 512, 2051) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) @@ -698,9 +829,9 @@ def run(): if batch == 1: assert torch.count_nonzero(eager) > 0 monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) - state.infer_state.is_cuda_graph_capture = True + state_owner.is_cuda_graph_capture = True graph, output = _capture_cuda_graph(run) - state.infer_state.is_cuda_graph_capture = False + state_owner.is_cuda_graph_capture = False graph.replay() torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) q.copy_(q.roll(1, dims=-1)) @@ -765,25 +896,23 @@ def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_h if batch >= 4: indices[3, 64:128] = indices[3, :64] # duplicates assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) - control = AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) + backend.network_config = {"index_topk": 2048, "index_kpool": 4} infer_state = _make_infer_state(is_cuda_graph_capture=False) - state = backend.create_att_decode_state(infer_state) + state = backend_context(type(backend), "decode", settings=backend) state.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] assert state.decode_topk_bound == 2051 native_state_owner = _make_infer_state(is_cuda_graph_capture=False) - native_state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(native_state_owner), + native_state = backend_context( + Fp8Glm5NextSparseAttBackend, + "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0], decode_topk_bound=None, ) - def run(current_state): - return current_state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + def run(current_state, request=infer_state): + return run_sparse(current_state, request, (q, q[..., :0]), packed[..., :584], control) native_calls = [] original_native = flash_mla.flash_mla_with_kvcache @@ -794,7 +923,7 @@ def observe_native(*args, **kwargs): monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", observe_native) eager = run(state) - native_output = run(native_state) + native_output = run(native_state, native_state_owner) assert native_calls assert torch.equal(eager, native_output) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) @@ -804,9 +933,9 @@ def native_must_not_run(*args, **kwargs): raise AssertionError("supported graph shape did not use the direct packed-KV path") monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) - state.infer_state.is_cuda_graph_capture = True + infer_state.is_cuda_graph_capture = True graph, output = _capture_cuda_graph(lambda: run(state)) - state.infer_state.is_cuda_graph_capture = False + infer_state.is_cuda_graph_capture = False graph.replay() torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) q.copy_(q.roll(1, dims=-1)) @@ -838,9 +967,7 @@ def test_fp8_sparse_direct_decode_unknown_bound_fallback_and_strided_q(monkeypat q = torch.randn(1, 8, 512, device="cuda", dtype=torch.bfloat16) late = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) late[:, 2175] = 17 - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": late, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": late, "softmax_scale": 0.0625, "kv_lora_rank": 512} calls = [] original = flash_mla.flash_mla_with_kvcache @@ -850,10 +977,13 @@ def native(*args, **kwargs): monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native) unknown_owner = _make_infer_state(is_cuda_graph_capture=False) - unknown = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(unknown_owner), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] + unknown = backend_context( + Fp8Glm5NextSparseAttBackend, + "decode", + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + decode_topk_bound=None, ) - actual = unknown.decode_att((q, q[..., :0]), packed[..., :584], None, control) + actual = run_sparse(unknown, unknown_owner, (q, q[..., :0]), packed[..., :584], control) torch.testing.assert_close( actual.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), late), atol=0.02, rtol=0.02 ) @@ -893,21 +1023,19 @@ def expected(): if (num_heads, batch) in ((1, 3), (7, 5)): infer = _make_infer_state(is_cuda_graph_capture=True) - state = Fp8Glm5NextSparseDecodeState(_infer_state_weak_ref=weakref.ref(infer), decode_topk_bound=2051) + state = backend_context(Fp8Glm5NextSparseAttBackend, "decode", decode_topk_bound=2051) def native_must_not_run(*args, **kwargs): raise AssertionError("supported strided decode took native fallback") monkeypatch.setattr("flash_mla.flash_mla_with_kvcache", native_must_not_run) torch.testing.assert_close( - state.decode_att( + run_sparse( + state, + infer, (q, q[..., :0]), packed, - None, - AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ), + {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, ).float(), expected(), atol=0.02, @@ -946,7 +1074,7 @@ def test_fp8_sparse_direct_decode_offsets_beyond_int32(strided_input): storage_size = (batch - 1) * stride + shape[-1] required_bytes = storage_size * torch.empty((), dtype=dtype).element_size() torch.cuda.empty_cache() - if torch.cuda.mem_get_info()[0] < required_bytes + 512 * 1024 ** 2: + if torch.cuda.mem_get_info()[0] < required_bytes + 512 * 1024**2: pytest.skip("Large-stride regression requires sufficient free GPU memory") # Allocate address space without initializing the gaps between logical rows. storage = torch.empty(storage_size, device="cuda", dtype=dtype) @@ -987,15 +1115,14 @@ def direct_must_not_run(*args, **kwargs): "lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode.fp8_sparse_direct_decode", direct_must_not_run ) infer = _make_infer_state(is_cuda_graph_capture=True) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer), + state = backend_context( + Fp8Glm5NextSparseAttBackend, + "decode", decode_topk_bound=2051, flashmla_sched_meta=flash_mla.get_mla_metadata()[0], ) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": topk, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) - run = lambda: state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + control = {"topk_mem_indices": topk, "softmax_scale": 0.0625, "kv_lora_rank": 512} + run = lambda: run_sparse(state, infer, (q, q[..., :0]), packed[..., :584], control) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) torch.testing.assert_close(run().float(), expected, atol=0.02, rtol=0.02) graph, output = _capture_cuda_graph(run) @@ -1003,30 +1130,142 @@ def direct_must_not_run(*args, **kwargs): torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) -@pytest.mark.parametrize("num_heads", [8, 64]) -def test_fp8_prefill_workspace_reuse_mutable_inputs_and_graph_replay(num_heads): +def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(): + """Generic H64 backend buffers are reused, batch metadata resets, and graph falls back.""" torch.manual_seed(53450) - backend = _make_workspace_backend(66) - case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], num_heads, backend=backend) - case.check(case.run()) - workspace = backend._prefill_workspaces[0] - assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} - pointers = {name: value.data_ptr() for name, value in workspace.items()} + capacity, num_q = 65536, 13 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) + starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 2048 + current_local = torch.full((num_q,), 2048, device="cuda", dtype=torch.int32) + prefill_mem = ragged[(starts + current_local).long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16).add_(2) + q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = current_local + topk[:, 2049] = current_local + control = {"topk_indices": topk, "prefill_cache_kv": prefill_kv, "softmax_scale": 0.0625, "kv_lora_rank": 512} + infer = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(2 * num_q) + + def expected(active_ragged=ragged): + values = _dequantize_kv_cache(packed)[active_ragged.long()] + values[(starts + current_local).long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + return _torch_sparse_attention(q, values, physical) + + def run(active, request=infer): + return run_sparse(active, request, q, packed[..., :584], control) + + state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) + cache = backend._test_workspace._buffers + assert not hasattr(infer, "_fp8_sparse_generic_prefill_workspace_cache") + workspace = state.generic_workspace + pointers = {k: v.data_ptr() for k, v in workspace.items() if isinstance(v, torch.Tensor)} + bounded = dict(workspace, gathered_kv=workspace["gathered_kv"][:4096], row_table_built=False) + actual = fp8_prefill.context_attention_fwd_fp8( + q, + packed[..., :584], + ragged, + prefill_mem, + prefill_kv, + starts, + topk, + 0.0625, + generic_workspace=bounded, + ) + torch.testing.assert_close(actual.float(), expected(), atol=0.012, rtol=0.015) for _ in range(10): - infer, state = case.new_state() - case.check(case.run(state)) - case.mutate() - case.check(case.run()) - fresh = _PrefillCase([16384, 32768, 16384], [9, 11, 13], num_heads, capacity=131072, backend=backend) - fresh.check(fresh.run()) - assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0].items()} == pointers - for overrides in ({"prefill_cuda_graph_exe_list": []}, {"is_cuda_graph_capture": True}): - infer, state = case.new_state(**overrides) - case.check(case.run(state)) - graph, output = _capture_cuda_graph(case.run) - case.mutate() - graph.replay() - case.check(output) + layer = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer).float(), expected(), atol=0.012, rtol=0.015) + assert workspace["row_table_built"] + assert {k: v.data_ptr() for k, v in layer.generic_workspace.items() if isinstance(v, torch.Tensor)} == pointers + + # Change raw KV/top-k in place: reusable metadata must not cache either. + q.normal_() + prefill_kv.normal_() + topk.copy_(topk.roll(1, dims=1)) + topk[:, 2049] = current_local + torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) + + fresh_ragged = ragged.roll(1) + fresh_mem = fresh_ragged[(starts + current_local).long()].clone() + fresh = _make_infer_state(mem_index=fresh_mem) + fresh_state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=fresh_ragged + ) + torch.testing.assert_close(run(fresh_state, fresh).float(), expected(fresh_ragged), atol=0.012, rtol=0.015) + assert cache[("glm_prefill_rows", torch.int32)].data_ptr() == workspace["prefill_row_table"].data_ptr() + + graph_infer = _make_infer_state(mem_index=prefill_mem, prefill_cuda_graph_exe_list=[]) + graph_state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged, capture=True + ) + torch.testing.assert_close(run(graph_state, graph_infer).float(), expected(), atol=0.012, rtol=0.015) + assert not hasattr(graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") + + explicit_graph_infer = _make_infer_state(mem_index=prefill_mem, is_cuda_graph_capture=True) + explicit_graph_state = backend_context( + Fp8Glm5NextSparseAttBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close( + run(explicit_graph_state, explicit_graph_infer).float(), expected(), atol=0.012, rtol=0.015 + ) + assert not hasattr(explicit_graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") + + +def test_fp8_generic_prefill_workspace_graph_query_is_generic_h64_only(monkeypatch): + """H8/direct/selected prefill never queries CUDA graph-capture state.""" + import lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 as fp8_prefill + + queries = [] + + def query_spy(): + queries.append(True) + raise AssertionError("non-generic path queried CUDA graph capture state") + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", query_spy) + kv = torch.empty((1, 1, 584), dtype=torch.uint8) + params = { + "prefill_cache_kv": torch.empty((1, 1, 512), dtype=torch.bfloat16), + "topk_indices": torch.empty((1, 2176), dtype=torch.int32), + "softmax_scale": 0.0625, + "kv_lora_rank": 512, + } + control = params + + def fake_context(*args, **kwargs): + return args[0] + + monkeypatch.setattr(fp8_prefill, "context_attention_fwd_fp8", fake_context) + for heads, direct, selected in ((8, False, False), (64, True, False), (64, False, True)): + monkeypatch.setattr(fp8_prefill, "can_use_direct_smallq", lambda *args, d=direct: d) + monkeypatch.setattr(fp8_prefill, "can_use_selected_prefill", lambda *args, s=selected: s) + infer = _make_infer_state(mem_index=torch.empty(1, dtype=torch.int32)) + backend = _make_workspace_backend(1) + backend.get_smallq_workspace = lambda active_infer, active_q, kind="direct" if direct else "selected": { + "kind": kind + } + state = backend_context( + Fp8Glm5NextSparseAttBackend, + "prefill", + settings=backend, + ks=torch.zeros(1, dtype=torch.int32), + ragged_mem_index=torch.zeros(1, dtype=torch.int32), + ) + assert ( + run_sparse(state, infer, torch.empty((1, heads, 512), dtype=torch.bfloat16), kv, control).shape[1] == heads + ) + assert not queries def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): @@ -1034,7 +1273,6 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) - monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.cpu_cache_operator.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.pd_operator.get_env_start_args", lambda: args) monkeypatch.setattr( diff --git a/unit_tests/models/glm5_next/test_kda_conv_qkv.py b/unit_tests/models/glm5_next/test_kda_conv_qkv.py index ed33e15738..8baadeaff2 100644 --- a/unit_tests/models/glm5_next/test_kda_conv_qkv.py +++ b/unit_tests/models/glm5_next/test_kda_conv_qkv.py @@ -1,15 +1,16 @@ -import weakref from types import SimpleNamespace import pytest import torch +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import KDAInputs, KDAParams from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv from lightllm.common.attention.backends._shared.linear.l2norm import USE_DEFAULT_FLA_NORM, l2norm_fwd +from lightllm.common.attention.backends.kda.triton import backend as backend_module from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.linear import kda as backend_module +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0) or USE_DEFAULT_FLA_NORM, @@ -125,14 +126,24 @@ def test_backend_fused_conv_output_and_state(lengths): ) infer_state.b_ready_cache_len = torch.arange(len(lengths), device="cuda", dtype=torch.int32) % 2 * 7 infer_state.max_q_seq_len = max(lengths) - backend = SimpleNamespace( + backend = make_backend( + backend_module.KDALinearAttBackend, head_dim=128, tp_num_heads=heads, tp_hidden_size=heads * 128, conv_kernel_size=4, lower_bound=-5.0, ) - att = backend_module.KDAPrefillAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer_state)) + batch = make_batch( + lengths, + phase=AttentionPhase.PREFILL, + device="cuda", + token_num=tokens, + max_q_seq_len=max(lengths), + b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, + b_ready_cache_len=infer_state.b_ready_cache_len, + ) + att = make_context(backend, batch) att.b_conv_buffer_idx = torch.arange(1, len(lengths) + 1, device="cuda", dtype=torch.int32) att.b_ssm_buffer_idx = att.b_conv_buffer_idx layer_weight = SimpleNamespace( @@ -140,14 +151,16 @@ def test_backend_fused_conv_output_and_state(lengths): linear_A_log=SimpleNamespace(weight=torch.zeros(heads, device="cuda")), linear_dt_bias=SimpleNamespace(weight=torch.zeros(heads * 128, device="cuda")), ) - control = AttControl( - linear_att_prefill=True, - linear_att_prefill_dict=dict( - layer_weight=layer_weight, - layer_num=0, - mixed_qkv=projection[:, :channels], - raw_gate=torch.randn(tokens, heads * 128, **opts) * 0.1, - raw_beta=torch.randn(tokens, heads, **opts), + inputs = KDAInputs( + mixed_qkv=projection[:, :channels], + raw_gate=torch.randn(tokens, heads * 128, **opts) * 0.1, + raw_beta=torch.randn(tokens, heads, **opts), + ) + attention = make_layer_attention( + backend, + att, + KDAParams( + conv_states, ssm_states, weight, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight ), ) expected_conv, expected_ssm = conv_states.clone(), ssm_states.clone() @@ -160,13 +173,12 @@ def test_backend_fused_conv_output_and_state(lengths): infer_state.b_ready_cache_len > 0, heads, ) - params = control.linear_att_prefill_dict expected, final_state = chunk_kda_with_fused_gate( q=q, k=k, v=v, - raw_g=params["raw_gate"].view(1, tokens, heads, 128), - beta=params["raw_beta"].float().sigmoid().view(1, tokens, heads), + raw_g=inputs.raw_gate.view(1, tokens, heads, 128), + beta=inputs.raw_beta.float().sigmoid().view(1, tokens, heads), A_log=layer_weight.linear_A_log.weight, g_bias=layer_weight.linear_dt_bias.weight, initial_state=ssm_states[att.b_ssm_buffer_idx].contiguous(), @@ -177,7 +189,7 @@ def test_backend_fused_conv_output_and_state(lengths): lower_bound=-5.0, ) expected_ssm[att.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype) - actual = att.prefill_att(None, None, None, att_control=control) + actual = attention.backend.forward(attention, att, inputs) assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) assert torch.equal(conv_states, expected_conv) assert torch.equal(ssm_states, expected_ssm) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index caef5d722d..c564e8ddaa 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -23,6 +23,15 @@ ) from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import set_env_start_args +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -411,6 +420,9 @@ def test_kpool_packed_kv_above_two_gib(): @pytest.mark.parametrize("tp_world_size", [1, 4]) def test_mhc_keeps_streams_through_decode_autotuning(monkeypatch, tp_world_size): + deep_gemm = pytest.importorskip("deep_gemm") + if not hasattr(deep_gemm, "tf32_hc_prenorm_gemm"): + pytest.skip("installed DeepGEMM does not provide mHC projection kernels") from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.model import Glm5NextTpPartModel @@ -466,7 +478,7 @@ def all_reduce(input_embeddings, **kwargs): def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule tokens, heads, dim = 67, 2, 128 q, k, v = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(3)] @@ -498,8 +510,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): @pytest.mark.parametrize("heads", [16, 64, 128]) def test_nope_attention_native_512_and_cuda_graph(heads): - from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseDecodeState, Glm5NextSparsePrefillState + from lightllm.models.glm5_next.attention.sparse_mla import Glm5NextSparseAttBackend if torch.cuda.get_device_capability()[0] != 9: pytest.skip("FA3 NoPE requires Hopper") @@ -521,34 +532,34 @@ def reference(): return torch.stack(expected) expected = reference() - control = AttControl( - nsa_prefill_dict={"topk_mem_indices": indexes, "softmax_scale": 0.0625, "kv_lora_rank": kv.shape[-1]} - ) - prefill_state = Glm5NextSparsePrefillState() - prefill = prefill_state._nsa_prefill_att(q, kv, control) + control = {"topk_mem_indices": indexes, "softmax_scale": 0.0625, "kv_lora_rank": kv.shape[-1]} + prefill_state = backend_context(Glm5NextSparseAttBackend, "prefill") + prefill_batch = SimpleNamespace() + prefill = run_sparse(prefill_state, prefill_batch, q, kv, control) torch.testing.assert_close(prefill.float(), expected, atol=0.012, rtol=0.015) infer_state = InferStateInfo() infer_state.b1_cu_q_seq_len = torch.arange(10, dtype=torch.int32, device="cuda") infer_state.max_q_seq_len = 1 - decode = Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), + decode = backend_context( + Glm5NextSparseAttBackend, + "decode", nsa_cache_seqlens=lengths, nsa_cu_seqlens_k_new=torch.nn.functional.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), ) - control.nsa_decode_dict = control.nsa_prefill_dict - out = decode._nsa_decode_att((q, q[..., :0]), kv, control) + control = control + out = run_sparse(decode, infer_state, (q, q[..., :0]), kv, control) torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - prefill_state._nsa_prefill_att(q, kv, control) - decode._nsa_decode_att((q, q[..., :0]), kv, control) + run_sparse(prefill_state, prefill_batch, q, kv, control) + run_sparse(decode, infer_state, (q, q[..., :0]), kv, control) torch.cuda.current_stream().wait_stream(stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - graph_prefill = prefill_state._nsa_prefill_att(q, kv, control) - graph_decode = decode._nsa_decode_att((q, q[..., :0]), kv, control) + graph_prefill = run_sparse(prefill_state, prefill_batch, q, kv, control) + graph_decode = run_sparse(decode, infer_state, (q, q[..., :0]), kv, control) graph.replay() torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 8ac9c2d473..d194a44710 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -1,22 +1,24 @@ -import weakref import dataclasses import json +import weakref from types import SimpleNamespace import pytest import torch import torch.nn.functional as F +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import MLAParams from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState from lightllm.common.quantization import Quantcfg +from lightllm.models.glm5_next.attention.sparse_mla import Glm5NextSparseAttBackend from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight from lightllm.models.glm5_next.model import Glm5NextTpPartModel from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args - +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -174,14 +176,16 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): state.att_state_manager = SimpleNamespace( get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv)) ) - state.prefill_att_state = Glm5NextSparsePrefillState() - state.b1_cu_q_seq_len = torch.arange(4, device="cuda", dtype=torch.int32) + backend = make_backend(Glm5NextSparseAttBackend) + phase = AttentionPhase.PREFILL + context = make_context(backend, make_batch((2, 3, 5), phase=phase, device="cuda")) + state.b1_cu_q_seq_len = context.batch.b1_cu_q_seq_len state.max_q_seq_len = 1 - state.decode_att_state = Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(state), - nsa_cache_seqlens=lengths, - nsa_cu_seqlens_k_new=F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), + state.attention_contexts = {0: context} + attention = make_layer_attention( + backend, context, MLAParams(kv, layer.kv_lora_rank, 0), softmax_scale=layer.softmax_scale ) + layer.prefill_attention = layer.decode_attention = attention projected = torch.bmm(q.transpose(0, 1), kv_b.view(4, 256, 512)[:, :128]).transpose(0, 1) expected = [] for row, length in enumerate(lengths.tolist()): @@ -191,6 +195,13 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): expected = torch.stack(expected) prefill = layer._context_attention_kernel(q, None, state, weight) state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} + context = make_context(backend, make_batch((2, 3, 5), phase=AttentionPhase.DECODE, device="cuda")) + context.nsa_cache_seqlens = lengths + context.nsa_cu_seqlens_k_new = F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)) + state.attention_contexts = {0: context} + layer.decode_attention = make_layer_attention( + backend, context, MLAParams(kv, layer.kv_lora_rank, 0), softmax_scale=layer.softmax_scale + ) decode = layer._token_attention_kernel(q, state, weight) for output in (prefill, decode): torch.testing.assert_close(output.float(), expected, atol=1e-3, rtol=0.015) @@ -427,14 +438,15 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): for step, draft in enumerate(drafts): draft.prepare_attention_states() assert draft.att_state_manager is main.att_state_manager - draft._init_att_backend1() state = getattr(draft.att_state_manager, f"draft{step}_mla_packed_bf16_state") assert state is not main.att_state_manager.main_mla_packed_bf16_state assert not hasattr(draft, "full_att_state") assert state.storage is None assert state.config.logical_layer_groups == ((4 + step,),) assert state.state_name == f"draft{step}_mla_packed_bf16_state" - assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None + from lightllm.common.attention.api import AttentionFamily + + assert draft.layers_infer[0].get_attention_family() is AttentionFamily.SPARSE_MLA assert getattr(draft.att_state_manager, f"draft{step}_compressor_indexer_w4_s4_bf16_state").buffer.is_cuda # NoPE state initialization must not require a model-owned RoPE cache. @@ -448,3 +460,44 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): torch.testing.assert_close( state.position_ids, torch.arange(2, 5, device="cuda", dtype=state.position_ids.dtype) ) + + +@pytest.mark.parametrize("family", ["gdn", "kda"]) +def test_prefill_graph_output_shape_uses_model_dimensions(monkeypatch, family): + from lightllm.models.glm5_next.layer_infer import transformer_layer_infer as glm + from lightllm.models.qwen3next.layer_infer import transformer_layer_infer as qwen + + callbacks = [] + graph_context = SimpleNamespace(__enter__=lambda: None, __exit__=lambda *args: None) + state = SimpleNamespace( + prefill_cuda_graph_get_current_capture_graph=lambda: graph_context, + prefill_cuda_graph_create_graph_obj=lambda: None, + prefill_cuda_graph_add_cpu_runnning_func=lambda func, after_graph: callbacks.append(func), + ) + # No attention_groups/backend dimensions: only the binding's public execution methods. + binding = SimpleNamespace(context=lambda state: None) + layer = SimpleNamespace(prefill_attention=binding, tp_world_size_=2, tp_linear_num_heads=3, linear_head_dim=16) + inputs = torch.zeros(5, 48, dtype=torch.bfloat16, device="cuda") + if family == "gdn": + monkeypatch.setattr(qwen, "tensor_to_no_ref_tensor", lambda tensor: tensor) + binding.forward = lambda *a: ( + torch.ones(1, 5, 3, 16, device="cuda"), + torch.full((5, 3, 16), 2.0, device="cuda"), + ) + weights = SimpleNamespace(linear_num_v_heads=6, linear_v_head_dim=16) + output, gate = qwen.Qwen3NextTransformerLayerInfer._linear_prefill_cuda_graph_wrapper( + layer, inputs, state, weights + ) + assert output.shape == gate.shape == (5, 3, 16) + else: + monkeypatch.setattr(glm, "tensor_to_no_ref_tensor", lambda tensor: tensor) + binding.forward = lambda *a: torch.ones(1, 5, 3, 16, device="cuda") + output = glm.Glm5NextTransformerLayerInfer._kda_prefill_cuda_graph_wrapper( + layer, inputs, inputs, inputs, state, None + ) + assert output.shape == (1, 5, 3, 16) + assert len(callbacks) == 1 + callbacks[0](state) + assert torch.all(output == 1) + if family == "gdn": + assert torch.all(gate == 2) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 6eef0d230b..404862c1fa 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -7,6 +7,7 @@ from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -144,10 +145,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod from types import SimpleNamespace from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState - from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDADecodeAttState + from lightllm.common.attention.backends.kda.replay.backend import ReplayKDALinearAttBackend + from lightllm.common.attention.backends.kda.triton.backend import KDALinearAttBackend torch.manual_seed(53) heads, dim, width = 2, 128, mtp_step + 1 @@ -200,15 +199,6 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod expected.append(output.reshape(1, heads, dim)) expected_ssm[req * width + token] = state[0] row += 1 - backend = SimpleNamespace( - mtp_step=mtp_step, - tp_num_heads=heads, - head_dim=dim, - tp_hidden_size=hidden, - lower_bound=-5.0, - uses_dynamic_spec_verify_layout=lambda: dynamic, - split_qkv=lambda x: x.split(hidden, -1), - ) resources = AttentionStateManager(request_capacity=4) config = ConvSSMStateConfig( logical_layer_ids=(0,), @@ -244,54 +234,48 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod infer.req_manager = manager infer.att_state_manager = resources infer.att_state_role = role - weight = SimpleNamespace( - get_merged_kda_conv_weight=lambda: conv_weight, - linear_A_log=SimpleNamespace(weight=a), - linear_dt_bias=SimpleNamespace(weight=bias), + from lightllm.common.attention.api.inputs import KDAInputs, KDAParams + from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention + + executor = make_backend( + ReplayKDALinearAttBackend if mode == "replay" else KDALinearAttBackend, + mtp_step=mtp_step, + decode_draft_step=mtp_step, + dynamic_verify=dynamic, + tp_num_heads=heads, + head_dim=dim, + tp_hidden_size=hidden, + lower_bound=-5.0, ) - decode_cls = ReplayKDADecodeAttState if mode == "replay" else KDADecodeAttState - state = decode_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - state.init_state() - actual = state.decode_att( - None, - None, - None, - AttControl( - linear_att_decode=True, - linear_att_decode_dict=dict( - layer_weight=weight, - layer_num=0, - mixed_qkv=mixed, - raw_gate=gate, - raw_beta=beta, - ), - ), + batch = make_batch( + (1,) * tokens, device="cuda", b_req_idx=infer.b_req_idx, b_mtp_index=infer.b_mtp_index, hold_req_id=3 ) - torch.testing.assert_close( - actual.reshape(tokens, heads, dim)[:row], - torch.cat(expected), - atol=2e-3, - rtol=1e-2, + context = make_context(executor, batch) + executor.update(context) + attention = make_layer_attention( + executor, + context, + KDAParams( + *selected.get_layer_view(0), + conv_weight, + a, + bias, + mtp_state_index=selected.mtp_state_index, + replay_state=selected if mode == "replay" else None, + ), ) + actual = attention.backend.forward(attention, context, KDAInputs(mixed, gate, beta)) + torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) if mode == "native": checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] torch.testing.assert_close( - selected.ssm_state[0, :checked_slots], - expected_ssm[:checked_slots], - atol=2e-3, - rtol=1e-2, + selected.ssm_state[0, :checked_slots], expected_ssm[:checked_slots], atol=2e-3, rtol=1e-2 ) else: for req, length in zip(requests, lengths): - saved = selected.ssm_state - clone = saved.clone() - selected.ssm_state = clone - try: - fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) - finally: - selected.ssm_state = saved + fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) torch.testing.assert_close( - clone[0, req], + selected.ssm_state[0, req], expected_ssm[req * width + length - 1], atol=2e-3, rtol=1e-2, diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 4d15d6bf1c..82b72b81bd 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -1,6 +1,3 @@ -from unit_tests.common.att_resources.model_helpers import window_states, build_hybrid_resources -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - import dataclasses from types import SimpleNamespace @@ -8,19 +5,22 @@ import torch from lightllm.common.attention_state.allocators import FullAttKVAllocator -from unit_tests.common.att_resources.model_helpers import create_full_attention_roles - from lightllm.common.attention_state.full_att import ( - PackedMLAFullAttentionState, PackedMLAAttentionConfig, + PackedMLAFullAttentionState, PackedMLALogicalConfig, ) -from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name - +from unit_tests.common.att_resources.model_helpers import ( + build_hybrid_resources, + create_full_attention_roles, + full_state_test_fields, + window_states, +) pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), @@ -304,6 +304,9 @@ def test_pd_kv_and_runtime_state_roundtrip( def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): + from lightllm.common.attention_state.pd_operator import logger + + monkeypatch.setattr(logger, "handlers", [*logger.handlers, caplog.handler]) monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] full_bytes = max( diff --git a/unit_tests/models/glm5_next/test_prefill_workspace.py b/unit_tests/models/glm5_next/test_prefill_workspace.py index 0d3a85d965..667132276b 100644 --- a/unit_tests/models/glm5_next/test_prefill_workspace.py +++ b/unit_tests/models/glm5_next/test_prefill_workspace.py @@ -1,116 +1,67 @@ -from types import SimpleNamespace - import pytest import torch -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next -from lightllm.models.glm5_next.model import Glm5NextTpPartModel -from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel -from unit_tests.models.glm5_next.test_fp8_kv import _PrefillCase, _make_workspace_backend, _capture_cuda_graph - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - +from lightllm.common.attention.api import AttentionPhase, ContextCapacity +from lightllm.common.attention.backends._workspace import AttentionWorkspace +from lightllm.models.glm5_next.attention.fp8_sparse_mla import ( + Fp8Glm5NextSparseAttBackend, + Fp8Glm5NextSparsePrefillContext, +) -@pytest.fixture(autouse=True) -def current_cuda_device(monkeypatch): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - -@pytest.mark.parametrize("overlap", [False, True]) -@pytest.mark.parametrize("run_mode", ["normal", "decode"]) -def test_prefill_workspace_is_reserved_once_and_reused_at_runtime(monkeypatch, overlap, run_mode): - model = SimpleNamespace( - args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode=run_mode), - batch_max_tokens=8, - max_req_num=2, - max_seq_length=13, - config={"index_topk": 2048}, - att_state_role=f"workspace_test_{run_mode}_{overlap}", +def test_fp8_prefill_context_is_owned_by_the_backend(): + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.network_config = {"index_topk": 2048, "index_kpool": 4} + workspace = AttentionWorkspace("cpu") + context = backend.allocate_context( + ContextCapacity(AttentionPhase.PREFILL, 2, 4, 32, 4), capture=False, workspace=workspace ) - backend_class = fp8_glm5_next.Fp8Glm5NextSparseAttBackend - reserved = backend_class.allocate_prefill_workspaces(model) - pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in reserved] - backend = object.__new__(backend_class) - backend.__init__(model) - assert [ - {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces - ] == pointers - assert len(reserved) == 1 + int(overlap) - capacity = 8 if run_mode == "decode" else 34 - assert reserved[0]["gathered_kv"].shape == (capacity, 1, 512) - assert reserved[0]["chunk_topk_indices"].shape == (8, 2176) - - def forbidden(*args, **kwargs): - raise AssertionError("Runtime must reuse the workspace reserved before KV capacity planning") - - monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) - monkeypatch.setattr(backend, "get_gpu_workspace_buffer", forbidden) - for slot in range(1 + int(overlap)): - for q_tokens in (1, 8, 3): - topk = torch.empty((q_tokens, 2176), dtype=torch.int32, device="cuda") - views = backend.get_prefill_workspace(SimpleNamespace(microbatch_index=slot), topk) - assert {key: value.data_ptr() for key, value in views.items()} == pointers[slot] - torch.testing.assert_close( - views["cu_seqlens_q"], torch.arange(q_tokens + 1, dtype=torch.int32, device="cuda") - ) + assert isinstance(context, Fp8Glm5NextSparsePrefillContext) + assert context.backend_identity is backend + assert context.workspace is workspace -def test_prefill_microbatches_use_independent_scratch_during_graph_replay(monkeypatch): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 7) - backend = _make_workspace_backend(8, overlap=True) - first = _PrefillCase([15], [3], 8, backend=backend) - second = _PrefillCase([17], [3], 8, backend=backend) - second.infer_state.microbatch_index = 1 - first.check(first.run()) - second.check(second.run()) - pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces] - assert all(pointers[0][key] != pointers[1][key] for key in pointers[0]) +def test_fp8_backend_has_no_legacy_state_factory(): + assert not hasattr(Fp8Glm5NextSparseAttBackend, "create_att_prefill_state") + assert not hasattr(Fp8Glm5NextSparseAttBackend, "create_att_decode_state") - def forbidden(*args, **kwargs): - raise AssertionError("Attention must not allocate a new prefill workspace") - monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) - first_graph, first_output = _capture_cuda_graph(first.run) - second_graph, second_output = _capture_cuda_graph(second.run) - first_stream, second_stream = torch.cuda.Stream(), torch.cuda.Stream() - first.mutate() - second.mutate() - for _ in range(3): - first_stream.wait_stream(torch.cuda.current_stream()) - second_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(first_stream): - first_graph.replay() - with torch.cuda.stream(second_stream): - second_graph.replay() - torch.cuda.current_stream().wait_stream(first_stream) - torch.cuda.current_stream().wait_stream(second_stream) - first.check(first_output) - second.check(second_output) - assert [ - {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces - ] == pointers +@pytest.mark.parametrize("decode_only", [False, True]) +def test_fp8_workspace_reserved_per_microbatch_before_execution(decode_only): + from types import SimpleNamespace + from lightllm.common.attention.api import ExecutionRequirements + from lightllm.common.basemodel.attention import AttentionGroup, reserve_attention_resources + from lightllm.platform import get_device_runtime -@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -def test_glm_reserves_prefill_workspace_during_resource_preparation(monkeypatch, model_class, kv_type): - model = object.__new__(model_class) - model.args = SimpleNamespace(llm_kv_type=kv_type) - model.mtp_previous_draft_models = [] - model.att_state_manager = SimpleNamespace(get_full_attention_states=lambda role: [object()]) - calls = [] - - def prepare(self): - assert self is model - calls.append("prepare_states") - - def reserve(cls, prepared_model): - assert prepared_model is model - calls.append("reserve_workspace") - - monkeypatch.setattr(TpPartBaseModel, "prepare_attention_states", prepare) - monkeypatch.setattr(fp8_glm5_next.Fp8Glm5NextSparseAttBackend, "allocate_prefill_workspaces", classmethod(reserve)) - model.prepare_attention_states() - assert calls == (["prepare_states", "reserve_workspace"] if kv_type == "fp8kv_dsa" else ["prepare_states"]) + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.network_config = {"index_topk": 2048, "index_kpool": 4} + backend.prefill_token_capacity = 32 + backend.prefill_graph_enabled = True + backend.decode_only = decode_only + model = SimpleNamespace( + attention_requirements=ExecutionRequirements(microbatch_count=2, max_batch_size=2, max_seq_len=64), + attention_groups=(AttentionGroup(backend, backend, None, None),), + _attention_workspaces={}, + device_runtime=get_device_runtime("cpu"), + ) + reserve_attention_resources(model) + first, second = (model._attention_workspaces[(0, slot)] for slot in (0, 1)) + gathered = first.get("glm_gathered_kv", 1, torch.bfloat16) + assert gathered.numel() == (32 if decode_only else 2 * 64 + 32) * 512 + assert gathered.data_ptr() != second.get("glm_gathered_kv", 1, torch.bfloat16).data_ptr() + assert first.get("glm_chunk_indices", 1, torch.int32).numel() == 32 * 2176 + context = backend.allocate_context( + ContextCapacity(AttentionPhase.PREFILL, 2, 4, 64, 4), capture=True, workspace=first + ) + context.ragged_mem_index = torch.arange(100, dtype=torch.int32) + q = torch.zeros(4, 64, 512, dtype=torch.bfloat16) + packed = torch.empty(160, 1, 584, dtype=torch.uint8) + topk = torch.zeros(4, 2176, dtype=torch.int32) + buffers = backend.get_generic_prefill_workspace(context, q, packed, topk) + assert buffers["gathered_kv"].data_ptr() == gathered.data_ptr() + assert buffers["cu_seqlens_q"].tolist() == [0, 1, 2, 3, 4] + first.freeze() + assert ( + backend.get_generic_prefill_workspace(context, q, packed, topk)["gathered_kv"].data_ptr() == gathered.data_ptr() + ) diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py index 32de9815e0..3881b5e419 100644 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -1,15 +1,17 @@ +from dataclasses import replace from types import SimpleNamespace import pytest import torch +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import GDNInputs, GDNParams +from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNBackend from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub @pytest.mark.parametrize("device", ["cpu", "cuda"]) @@ -41,17 +43,11 @@ def test_gdn_decode_keeps_logical_layer_state( kinds = get_loop_layer_types(settings, config) layers = tuple(i for i, kind in enumerate(kinds) if kind == "linear_attention") args = StartArgs(mtp_step=mtp_step, linear_att_ssm_data_type="float32") - monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.basemodel.attention.base_att.get_env_start_args", lambda: args) - backend = object.__new__(TritonLinearAttBackend) - backend.__init__( - SimpleNamespace( - config=config, - tp_world_size_=1, - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: mtp_step), - ) - ) + monkeypatch.setattr("lightllm.common.attention.backends._shared.linear.gdn.get_env_start_args", lambda: args) + model = model_stub(mtp_step=mtp_step, linear_att_ssm_data_type="float32") + model.config = config + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda _: mtp_step) + backend = TritonGDNBackend(model, make_spec(use_mtp=mtp_step > 0)) dtype = torch.bfloat16 if device == "cuda" else torch.float32 layout = ConvSSMStateConfig(layers, 1, 2, 32, 32, 3, dtype, torch.float32) resource = GDNState(layout, request_capacity=3, mtp_step=mtp_step, device=device) @@ -91,57 +87,57 @@ def recurrent(*, initial_state, v, **kwargs): initial_state[mtp_step + 1 : 2 * (mtp_step + 1)].add_(1) return v, None - monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) + monkeypatch.setattr( + "lightllm.common.attention.backends.gdn.triton.backend.conv_pack_gdn_decode_inputs", conv_pack + ) monkeypatch.setattr( "lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp.causal_conv1d_update", conv_mtp ) monkeypatch.setattr( - "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent + "lightllm.common.attention.backends.gdn.triton.backend.fused_recurrent_gated_delta_rule", recurrent ) monkeypatch.setattr( - "lightllm.common.basemodel.attention.linear.gdn.mtp_fused_recurrent_gated_delta_rule", + "lightllm.common.attention.backends.gdn.triton.backend.mtp_fused_recurrent_gated_delta_rule", recurrent, ) - infer = InferStateInfo() - infer.att_state_manager, infer.att_state_role = resources, "main" - infer.b_req_idx = torch.tensor([1], device=device, dtype=torch.int32) - control = dict(layer_num=logical_layer, layer_weight=weight) + params = GDNParams( + conv_view, + ssm_view, + weight.linear_conv1d.mm_param.weight, + None, + weight.linear_A_log.weight, + weight.linear_dt_bias.weight, + mtp_state_index=resource.mtp_state_index, + ) mixed_dim = backend.tp_qkvz_dim + backend.tp_ba_dim if device == "cuda": - infer.b1_cu_q_seq_len = torch.tensor([0, 5], device=device, dtype=torch.int32) - infer.b_ready_cache_len = torch.zeros(1, device=device, dtype=torch.int32) - prefill = backend.create_att_prefill_state(infer) - prefill.init_state() - prefill.prefill_att( - None, - None, - None, - AttControl( - linear_att_prefill=True, - linear_att_prefill_dict=dict( - **control, mixed_qkvzba=torch.randn(5, mixed_dim, device=device, dtype=dtype) - ), - ), + batch = make_batch( + (5,), + device=device, + phase=AttentionPhase.PREFILL, + b_req_idx=torch.tensor([1], device=device, dtype=torch.int32), + token_num=5, + max_q_seq_len=5, + b_q_seq_len=torch.tensor([5], device=device), + b1_cu_q_seq_len=torch.tensor([0, 5], device=device, dtype=torch.int32), ) - - infer.batch_size = mtp_step + 1 - infer.b_req_idx = infer.b_req_idx.repeat(infer.batch_size) - decode = backend.create_att_decode_state(infer) + context = make_context(backend, batch) + backend.update(context) + attention = replace(make_layer_attention(backend, context, params), logical_layer_id=logical_layer) + attention.backend.forward(attention, context, GDNInputs(torch.randn(5, mixed_dim, device=device, dtype=dtype))) + batch_size = mtp_step + 1 + batch = make_batch( + (1,) * batch_size, device=device, b_req_idx=torch.ones(batch_size, device=device, dtype=torch.int32) + ) + context = make_context(backend, batch) + attention = replace(make_layer_attention(backend, context, params), logical_layer_id=logical_layer) untouched = [i for i in range(len(layers)) if i != storage_layer] for _ in range(2): before_conv, before_ssm = resource.conv_state.clone(), resource.ssm_state.clone() - decode.init_state() - output, _ = decode.decode_att( - None, - None, - None, - AttControl( - linear_att_decode=True, - linear_att_decode_dict=dict( - **control, mixed_qkvzba=torch.randn(infer.batch_size, mixed_dim, device=device, dtype=dtype) - ), - ), + backend.update(context) + output, _ = attention.backend.forward( + attention, context, GDNInputs(torch.randn(batch_size, mixed_dim, device=device, dtype=dtype)) ) assert torch.isfinite(output).all() assert not torch.equal(resource.conv_state[storage_layer], before_conv[storage_layer]) diff --git a/unit_tests/server/core/objs/test_req.py b/unit_tests/server/core/objs/test_req.py index 5966497ef1..e35a59917e 100644 --- a/unit_tests/server/core/objs/test_req.py +++ b/unit_tests/server/core/objs/test_req.py @@ -1,7 +1,9 @@ -import pytest -import easydict from types import SimpleNamespace -from lightllm.server.core.objs.req import Req, ChunkedPrefillReq, SamplingParams + +import easydict +import pytest + +from lightllm.server.core.objs.req import ChunkedPrefillReq, Req, SamplingParams from lightllm.server.core.objs.token_metadata import ReqFinalTokenMetadata from lightllm.utils import shm_utils from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args @@ -16,8 +18,8 @@ def setup_module_env(): easydict.EasyDict( { "mtp_step": 0, - "llm_prefill_att_backend": ["None"], - "llm_decode_att_backend": ["None"], + "llm_prefill_att_backend": "None", + "llm_decode_att_backend": "None", "cpu_cache_token_page_size": 256, "enable_cpu_cache": False, "model_dir": "", diff --git a/unit_tests/server/core/objs/test_shm_req_manager.py b/unit_tests/server/core/objs/test_shm_req_manager.py index c51e3e6b95..cc0dafe109 100644 --- a/unit_tests/server/core/objs/test_shm_req_manager.py +++ b/unit_tests/server/core/objs/test_shm_req_manager.py @@ -1,16 +1,17 @@ -import os import asyncio -import pytest +import os import time import uuid -import numpy as np -from unittest.mock import MagicMock from types import SimpleNamespace +from unittest.mock import MagicMock +import numpy as np +import pytest from easydict import EasyDict -from lightllm.utils.envs_utils import set_env_start_args, get_env_start_args -from lightllm.utils import shm_utils + from lightllm.server.core.objs.shm_req_manager import ShmReqManager +from lightllm.utils import shm_utils +from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args @pytest.fixture(scope="module", autouse=True) @@ -23,8 +24,8 @@ def setup_env(): running_max_req_size=10, disable_chunked_prefill=True, mtp_step=0, - llm_prefill_att_backend=["None"], - llm_decode_att_backend=["None"], + llm_prefill_att_backend="None", + llm_decode_att_backend="None", cpu_cache_token_page_size=256, enable_cpu_cache=False, ) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index ece87927e0..be73427336 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -1,23 +1,19 @@ -from types import SimpleNamespace from collections import deque +from types import SimpleNamespace import pytest import torch -from lightllm.server.multi_level_kv_cache import ( - AdaptiveCachePlacementController, - CacheCapacityConfig, - CacheTier, -) -from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule -from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl +from lightllm.server.multi_level_kv_cache import AdaptiveCachePlacementController, CacheCapacityConfig, CacheTier from lightllm.server.router.model_infer.infer_batch import InferenceContext +from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl +from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule def test_running_cpu_cache_tasks_are_not_ready(monkeypatch): from lightllm.server.router.model_infer.infer_batch import InferReq - from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.server.router.model_infer.mode_backend import base_backend as base_backend_impl + from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend status = InferReq._CpuCacheTaskStatus @@ -494,13 +490,13 @@ def query(self): def test_fa3_cpu_cache_sync_defaults_off(): - from lightllm.common.basemodel.attention.base_att import BaseAttBackend - from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend - from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend - from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend - from lightllm.common.basemodel.attention.fa3.neo import NeoFa3AttBackend + from lightllm.common.attention.backends.base import BaseAttentionBackend + from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend + from lightllm.common.attention.backends.dense.fa3.fp8 import Fp8Fa3AttBackend + from lightllm.common.attention.backends.dense.fa3.neo import NeoFa3AttBackend + from lightllm.common.attention.backends.mla.fa3.backend import MlaFa3AttBackend - assert BaseAttBackend.requires_cpu_cache_sync is False + assert BaseAttentionBackend.requires_cpu_cache_sync is False assert all( backend.requires_cpu_cache_sync is False for backend in (Fa3AttBackend, Fp8Fa3AttBackend, MlaFa3AttBackend, NeoFa3AttBackend) diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index 21bbb06339..17774fb802 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -2,34 +2,35 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from typing import TYPE_CHECKING import dataclasses import json from types import SimpleNamespace +from typing import TYPE_CHECKING import pytest import torch -from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states -from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel - from lightllm.models import Deepseek2TpPartModel, Deepseek3_2TpPartModel, LlamaTpPartModel -from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_unique_server_name from lightllm.utils.kv_cache_utils import calcu_cpu_cache_meta +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + main_linear_state, + prepare_model_states, +) if TYPE_CHECKING: from collections.abc import Callable + from pathlib import Path + from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient from lightllm.utils.kv_cache_utils import CpuKVCacheMeta - from pathlib import Path pytestmark = [ @@ -372,9 +373,8 @@ def capture_write(*args: object, **kwargs: object) -> None: pool.ssm_state[slot].fill_(5) producer = torch.cuda.Stream() producer.wait_stream(torch.cuda.current_stream()) + monkeypatch.setattr(common.g_infer_context, "get_overlap_stream", lambda: producer) with torch.cuda.stream(producer): - # 延迟当前生产流的 KV 写入；当前实现不在此处隐式插入 wait_stream， - # 生产者与搬运流的依赖由调用方负责组织。 torch.cuda._sleep(10_000_000) (full,) = resources.get_full_attention_states("main") for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): @@ -429,7 +429,7 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: patch.setattr(torch.cuda.Stream, "wait_stream", record_wait) module.get_cpu_offload_stream = lambda: copy_stream task = module.write_req_to_cpu(req) - assert waited_streams == [] + assert waited_streams == [producer] task.sync_event.synchronize() assert write_batches == [(pages, [True, False])] assert req.cpu_cache_offload_task_status.is_running() diff --git a/unit_tests/server/test_api_start_deepep.py b/unit_tests/server/test_api_start_deepep.py index 949021cd34..c1f417765f 100644 --- a/unit_tests/server/test_api_start_deepep.py +++ b/unit_tests/server/test_api_start_deepep.py @@ -6,7 +6,9 @@ from lightllm.server.core.objs.start_args_type import StartArgs -def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path): +@pytest.mark.parametrize("backend", ["auto", "fa3", "triton", "flashinfer"]) +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path, backend, phase): for name in ( "_set_envs_and_config", "auto_set_max_req_total_len", @@ -39,6 +41,14 @@ def stop(args): router_token_ratio=0.0, ) + setattr(args, f"llm_{phase}_att_backend", backend) + args.llm_prefill_linear_att_backend = "flashqla" + args.llm_decode_linear_att_backend = "triton" + if backend == "flashinfer": + with pytest.raises(AssertionError, match=f"--llm_{phase}_att_backend"): + _launch_subprocesses(args) + return + with pytest.raises(RuntimeError, match="startup validation passed"): _launch_subprocesses(args) diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 1c391c09cd..7aa41c2193 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -107,6 +107,7 @@ def test_whole_image_blocks_fit_batch_at_startup( @pytest.mark.parametrize("mode,step", [("native", 0), ("replay", 0), ("replay", 4)]) def test_s3_configuration_is_prepared_before_ssm_validation(launch_args, monkeypatch, tmp_path, mode, step): from pathlib import Path + from lightllm.utils.petrel_helper import PetrelHelper monkeypatch.chdir(tmp_path) diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py index 1f4efc4baa..4e77c6237e 100644 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -3,9 +3,9 @@ from types import SimpleNamespace import numpy as np +import pytest import torch - REPO_ROOT = Path(__file__).resolve().parents[1] STATIC_BENCHMARK_DIR = REPO_ROOT / "test" / "benchmark" / "static_inference" for path in (REPO_ROOT, STATIC_BENCHMARK_DIR): @@ -153,3 +153,50 @@ def test_profile_capacity_uses_page_aligned_per_request_cost(): args.mtp_mode = "eagle_with_att" args.mtp_step = 3 assert decode_profile_batch_divisor(args, decode_case) == 24 + + +@pytest.mark.parametrize("fail", [False, True]) +def test_benchmark_shared_memory_cleanup(monkeypatch, fail): + from multiprocessing import shared_memory + + from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt + + monkeypatch.delenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", raising=False) + monkeypatch.setenv("LIGHTLLM_START_ARGS", "{}") + static_benchmark.get_unique_server_name.cache_clear() + names = [] + + def run(args, cases): + counter = SharedInt("benchmark_counter") + names.append(counter.shm.name) + counter.set_value(17) + assert counter.get_value() == 17 + counter.shm.close() + if fail: + raise RuntimeError("worker failed") + + monkeypatch.setattr(static_benchmark, "run_benchmark", run) + argv = [ + "--benchmark", + "decode", + "--batch_size", + "1", + "--input_len", + "4", + "--output_len", + "1", + "--data_type", + "bfloat16", + "--enable_fused_shared_experts", + ] + try: + if fail: + with pytest.raises(RuntimeError, match="worker failed"): + static_benchmark.main(argv) + else: + static_benchmark.main(argv) + assert len(names) == 1 + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=names[0]) + finally: + static_benchmark.get_unique_server_name.cache_clear() diff --git a/unit_tests/utils/test_sgl_utils.py b/unit_tests/utils/test_sgl_utils.py index 85bb041e68..dcd0a04af7 100644 --- a/unit_tests/utils/test_sgl_utils.py +++ b/unit_tests/utils/test_sgl_utils.py @@ -1,19 +1,28 @@ -import weakref import collections import inspect import math +import weakref from types import SimpleNamespace import pytest import torch from frozendict import frozendict +from lightllm.common.attention.backends.dense.fa3 import backend as fa3_module from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.fa3 import fp as fa3_module from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils import sgl_utils from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -114,8 +123,10 @@ def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch, pa infer_state.b_seq_len = torch.tensor([4096, actual_kv_len], dtype=torch.int32) infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) infer_state.b1_cu_kv_seq_len = torch.tensor([0, 4096, 4096 + actual_kv_len], dtype=torch.int32) - state = fa3_module.Fa3DecodeAttState( - backend=SimpleNamespace( + state = backend_context( + fa3_module.Fa3AttBackend, + "decode", + settings=SimpleNamespace( model=model, infer_page_size=page_size, uses_causal_attention=lambda: True, @@ -124,13 +135,12 @@ def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch, pa att_batch_size, math.ceil(max_kv_len / page_size), dtype=torch.int32 ), ), - _infer_state_weak_ref=weakref.ref(infer_state), ) copy_calls = [] monkeypatch.setattr(fa3_module, "page_table_copy", lambda **kwargs: copy_calls.append(kwargs)) - state.init_state() + update_context(state, infer_state) # 模拟 CudaGraph._capture_decode 对 infer_state 的修改。 - state.infer_state.max_kv_seq_len = graph.graph_max_len_in_batch + infer_state.max_kv_seq_len = graph.graph_max_len_in_batch calls = [] def attention(**kwargs): @@ -138,9 +148,9 @@ def attention(**kwargs): return kwargs["q"] monkeypatch.setattr(fa3_module, "flash_attn_with_kvcache_autotune", attention) - state.decode_att(torch.empty(2, 4, 8), torch.empty(page_size, 2, 8), torch.empty(page_size, 2, 8)) + run_dense(state, infer_state, torch.empty(2, 4, 8), torch.empty(page_size, 2, 8), torch.empty(page_size, 2, 8)) assert calls[0]["max_seqlen_k"] == actual_kv_len - assert calls[0]["page_table"].shape == (2, 32768 // page_size) + assert calls[0]["page_table"].shape == (2, math.ceil(actual_kv_len / page_size)) assert copy_calls[0]["page_table"].shape == (2, math.ceil(actual_kv_len / page_size)) assert copy_calls[0]["page_size"] == page_size diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py index a58b35198d..3ebd75505f 100644 --- a/unit_tests/utils/test_speculative_utils.py +++ b/unit_tests/utils/test_speculative_utils.py @@ -1,21 +1,23 @@ -import weakref import json +import weakref from importlib import import_module from types import SimpleNamespace import pytest import torch -from lightllm.common.basemodel.infer_struct import InferStateInfo -import lightllm.common.basemodel.attention.base_att as base_att_module +import lightllm.common.attention.backends.base as base_att_module import lightllm.common.basemodel.hidden_collector as hidden_collector_module -from lightllm.common.basemodel.attention.base_att import BaseAttBackend -from lightllm.common.basemodel.attention.fa3.fp import Fa3DecodeAttState, Fa3PrefillAttState -from lightllm.common.basemodel.attention.fa3.mla import MlaFa3DecodeAttState, MlaFa3PrefillAttState -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.backends.base import BaseAttentionBackend +from lightllm.common.attention.backends.dense.fa3.backend import Fa3AttBackend +from lightllm.common.attention.backends.dense.triton.backend import TritonAttBackend +from lightllm.common.attention.backends.mla.fa3.backend import MlaFa3AttBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models import get_draft_model_class from lightllm.models.qwen3_eagle.layer_weights.transformer_layer_weight import Qwen3EagleTransformerLayerWeight from lightllm.utils import envs_utils +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_spec, model_stub @pytest.mark.parametrize( @@ -76,19 +78,11 @@ def test_attention_backend_selects_dynamic_spec_layout( dynamic_spec, expected, ): - monkeypatch.setattr( - base_att_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_mode=mtp_mode, mtp_dynamic_verify=dynamic_spec), - ) - backend = SimpleNamespace( - model=SimpleNamespace( - is_mtp_draft_model=is_draft_model, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), - ) - ) - - assert BaseAttBackend.uses_dynamic_spec_verify_layout(backend) is expected + model = model_stub(mtp_mode=mtp_mode, mtp_dynamic_verify=dynamic_spec) + model.is_mtp_draft_model = is_draft_model + model.mtp_manager.get_decode_draft_step = lambda _: draft_step + backend = BaseAttentionBackend(model, _spec_from_model(model)) + assert backend.uses_dynamic_spec_verify_layout() is expected @pytest.mark.parametrize( @@ -104,102 +98,59 @@ def test_attention_backend_selects_dynamic_spec_layout( ], ) def test_attention_backend_selects_causality(monkeypatch, mtp_mode, is_draft_model, expected): - monkeypatch.setattr( - base_att_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_mode=mtp_mode), - ) - backend = SimpleNamespace(model=SimpleNamespace(is_mtp_draft_model=is_draft_model)) + model = model_stub(mtp_mode=mtp_mode) + model.is_mtp_draft_model = is_draft_model + assert BaseAttentionBackend(model, _spec_from_model(model)).uses_causal_attention() is expected - assert BaseAttBackend.uses_causal_attention(backend) is expected - -@pytest.mark.parametrize("state_class", [Fa3PrefillAttState, MlaFa3PrefillAttState]) +@pytest.mark.parametrize("state_class", [Fa3AttBackend, MlaFa3AttBackend]) def test_fa3_prefill_state_owns_causality(monkeypatch, state_class): - monkeypatch.setattr(import_module(Fa3PrefillAttState.__module__), "page_table_copy", lambda **kwargs: None) - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) - infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) - infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = 4 - infer_state.input_ids = torch.empty(2, dtype=torch.int64) - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, dtype=torch.int32).reshape(2, 4)) - state = state_class( - backend=SimpleNamespace(uses_causal_attention=lambda: False, page_size=1, infer_page_size=1), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - - state.init_state() - - assert state.causal is False + backend = make_backend(state_class, causal=False) + monkeypatch.setattr(import_module(Fa3AttBackend.__module__), "page_table_copy", lambda **kwargs: None) + context = make_context(backend, make_batch(phase=AttentionPhase.PREFILL)) + backend.update(context) + assert context.causal is False -@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) +@pytest.mark.parametrize("state_class", [Fa3AttBackend, MlaFa3AttBackend]) def test_fa3_decode_state_owns_causality(state_class): - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) - infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) - infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) - infer_state.b_seq_len = torch.tensor([3, 4], dtype=torch.int32) - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), - ) - state = state_class( - backend=SimpleNamespace( - model=model, - uses_causal_attention=lambda: False, - uses_dynamic_spec_verify_layout=lambda: False, - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state._init_page_table = lambda _: None - - state.init_state() - - assert state.causal is False + backend = make_backend(state_class, causal=False) + backend._decode_init_page_table = lambda *args: None + context = make_context(backend) + backend.update(context) + assert context.causal is False @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_triton_mtp_decode_state_builds_group_markers(): - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), - req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), - ) - infer_state = InferStateInfo() - infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") - infer_state.max_kv_seq_len = 4 - state = TritonDecodeAttState( - backend=SimpleNamespace(model=model), - _infer_state_weak_ref=weakref.ref(infer_state), + backend = make_backend(TritonAttBackend, decode_draft_step=2, hold_req_id=-1) + batch = make_batch( + (4, 4, 2, 2), + device="cuda", + b_req_idx=torch.tensor([7, 7, -1, -1], device="cuda", dtype=torch.int32), + hold_req_id=-1, ) - - state.init_state() - - assert state.b_mark_mtp_shared_group.cpu().tolist() == [0, 2, 1, 1] + context = make_context(backend, batch) + backend.update(context) + assert context.b_mark_mtp_shared_group.cpu().tolist() == [0, 2, 1, 1] @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) +@pytest.mark.parametrize("state_class", [Fa3AttBackend, MlaFa3AttBackend]) def test_fa3_dynamic_decode_state_builds_group_markers(state_class): - model = SimpleNamespace(req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1)) - infer_state = InferStateInfo() - infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") - infer_state.b_seq_len = torch.tensor([3, 4, 2, 2], dtype=torch.int32, device="cuda") - infer_state.batch_size = 4 - state = state_class( - backend=SimpleNamespace(model=model), - _infer_state_weak_ref=weakref.ref(infer_state), + backend = make_backend(state_class, hold_req_id=-1) + batch = make_batch( + (3, 4, 2, 2), + device="cuda", + b_req_idx=torch.tensor([7, 7, -1, -1], device="cuda", dtype=torch.int32), + hold_req_id=-1, ) - - b_att_req_idx = state._init_dynamic_spec_verify_state(draft_step=2) + context = make_context(backend, batch) + b_att_req_idx = backend._decode_init_dynamic_spec_verify_state(context, draft_step=2) torch.cuda.synchronize() - assert b_att_req_idx.cpu().tolist() == [7, -1, -1, -1] - assert state.b_att_seq_len.cpu().tolist() == [4, 2, 2, 0] - assert state.cu_seqlens_q.cpu().tolist() == [0, 2, 3, 4, 4] + assert context.b_att_seq_len.cpu().tolist() == [4, 2, 2, 0] + assert context.cu_seqlens_q.cpu().tolist() == [0, 2, 3, 4, 4] @pytest.mark.parametrize( @@ -411,3 +362,29 @@ def test_eagle3_added_kv_layers_come_from_draft_config(tmp_path): ) assert envs_utils.get_added_mtp_kv_layer_num() == 2 + + +def _spec_from_model(model): + from types import SimpleNamespace + + from lightllm.common.attention.api import AttentionFamily + from lightllm.common.basemodel.attention import describe_attention_layer + + state = SimpleNamespace( + contains_logical_layer=lambda layer: True, + describe_layout=lambda: (), + config=SimpleNamespace(dtype=torch.bfloat16), + ) + model.data_type = torch.bfloat16 + model.att_state_manager.get_full_attention_states = lambda role: (state,) + layer = SimpleNamespace( + layer_num_=0, + tp_q_head_num_=4, + tp_k_head_num_=2, + head_dim_=64, + attention_window=lambda phase: None, + attention_uses_alibi=False, + attention_uses_sink=False, + attention_variant="", + ) + return describe_attention_layer(model, layer, AttentionFamily.DENSE, AttentionPhase.DECODE) From 221c695333e06f169a8b60fb6681b2008903a206 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 16:32:30 +0800 Subject: [PATCH 6/9] refactor(attention): separate phase backends and execution sessions --- docs/CN/source/models/add_new_model.md | 19 +- docs/EN/source/models/add_new_model.md | 19 +- lightllm/common/attention/__init__.py | 11 + lightllm/common/attention/_environment.py | 17 + lightllm/common/attention/api/__init__.py | 28 + lightllm/common/attention/api/backend.py | 142 +++ lightllm/common/attention/api/batch.py | 116 ++ lightllm/common/attention/api/environment.py | 25 + lightllm/common/attention/api/execution.py | 35 + lightllm/common/attention/api/graph.py | 24 + lightllm/common/attention/api/inputs.py | 179 +++ lightllm/common/attention/api/spec.py | 121 ++ lightllm/common/attention/api/varlen.py | 54 + lightllm/common/attention/api/workspace.py | 115 ++ .../fa3 => attention/backends}/__init__.py | 0 .../backends/_flashinfer.py} | 1 + .../backends/_shared}/__init__.py | 0 .../backends/_shared}/linear/__init__.py | 0 .../backends/_shared/linear}/causal_conv1d.py | 0 .../_shared/linear}/causal_conv1d_mtp.py | 0 .../_shared/linear}/causal_conv1d_qkv.py | 0 .../backends/_shared/linear}/chunk_delta_h.py | 5 +- .../backends/_shared/linear}/cumsum.py | 5 +- .../attention/backends/_shared/linear/gdn.py | 52 + .../backends/_shared/linear}/index.py | 3 +- .../backends/_shared/linear}/l2norm.py | 2 +- .../backends/_shared/linear/layout.py | 67 + .../_shared/linear}/mtp_state_params.py | 52 +- .../backends/_shared/linear}/op.py | 2 +- .../_shared/linear}/replayssm_cursor.py | 0 .../_shared/linear}/replayssm_decode.py | 0 .../_shared/linear}/replayssm_delta.py | 0 .../_shared/linear}/replayssm_fold.py | 0 .../backends/_shared/linear}/solve_tril.py | 7 +- .../_shared/linear}/ssm_acceptance.py | 0 .../backends/_shared/linear}/utils.py | 1 - .../backends/_shared/prefill}/__init__.py | 0 .../prefill}/context_flashattention_nopad.py | 10 +- .../backends/dense}/__init__.py | 0 .../backends/dense/fa3}/__init__.py | 0 .../attention/backends/dense/fa3/backend.py | 250 ++++ .../attention/backends/dense/fa3/fp8.py | 166 +++ .../attention/backends/dense/fa3/neo.py | 95 ++ .../attention/backends/dense/fa3/varlen.py | 72 ++ .../backends/dense/flashinfer}/__init__.py | 0 .../backends/dense/flashinfer/backend.py | 249 ++++ .../backends/dense/flashinfer/fp8.py | 111 ++ .../backends/dense}/sdpa/__init__.py | 0 .../backends/dense/sdpa/varlen.py} | 41 +- .../backends/dense}/triton/__init__.py | 0 .../backends/dense/triton/backend.py | 205 ++++ .../attention/backends/dense/triton/int4.py | 152 +++ .../attention/backends/dense/triton/int8.py | 189 +++ .../dense/triton/kernels}/__init__.py | 0 .../dense/triton/kernels/decode}/__init__.py | 0 .../triton/kernels/decode/gqa}/__init__.py | 0 .../kernels/decode/gqa/flash}/__init__.py | 0 .../decode/gqa/flash}/gqa_flash_decoding.py | 12 +- .../gqa/flash}/gqa_flash_decoding_stage1.py | 6 +- .../gqa/flash}/gqa_flash_decoding_stage2.py | 0 .../gqa/flash}/gqa_flash_decoding_vsm.py | 10 +- .../gqa/gqa_decode_flashattention_nopad.py | 7 +- .../decode}/gqa/mtp_diverse/__init__.py | 0 .../gqa/mtp_diverse/mtp_diverse_attn.py | 4 +- .../gqa/mtp_diverse/stage1_single_token.py | 6 +- .../gqa/mtp_diverse/stage2_single_token.py | 6 +- .../triton/kernels/decode/int4kv}/__init__.py | 0 .../int4kv/int4kv_flash_decoding_stage1.py | 6 +- .../int4kv/ppl_int4kv_flash_decoding.py | 15 +- .../triton/kernels/decode/int8kv}/__init__.py | 0 .../int8kv/int8kv_flash_decoding_diverse.py | 36 +- .../int8kv_flash_decoding_diverse_stage1.py | 9 +- .../int8kv_flash_decoding_diverse_stage2.py | 10 +- .../int8kv_flash_decoding_diverse_stage3.py | 0 .../kernels/decode}/int8kv/normal/__init__.py | 0 .../int8kv/normal/int8kv_flash_decoding.py | 20 +- .../normal/int8kv_flash_decoding_stage1.py | 7 +- .../normal/int8kv_flash_decoding_stage2.py | 0 .../triton/kernels/decode/mha}/__init__.py | 0 .../kernels/decode/mha/flash}/__init__.py | 0 .../decode/mha/flash}/flash_decoding.py | 16 +- .../mha/flash}/flash_decoding_stage1.py | 0 .../mha/flash}/flash_decoding_stage2.py | 0 .../decode/mha/stage3_decode_att}/__init__.py | 0 .../token_attention_nopad_att1.py | 4 +- .../token_attention_nopad_reduceV.py | 1 - .../token_attention_nopad_softmax.py | 1 - .../token_attention_softmax_and_reducev.py | 2 +- .../dense/triton/kernels/prefill}/__init__.py | 0 .../prefill}/context_attention_fwd_neo.py | 1 + .../dense/triton/kernels/prefill/varlen.py} | 88 +- .../attention/backends/dense/triton/neo.py | 66 + .../attention/backends/dense/triton/varlen.py | 25 + .../common/attention/backends/dense/varlen.py | 42 + .../backends/dense/xformers}/__init__.py | 0 .../backends/dense/xformers/varlen.py} | 32 +- .../backends/gdn}/__init__.py | 0 .../backends/gdn/flashqla/__init__.py | 0 .../backends/gdn/flashqla/backend.py | 21 + .../attention/backends/gdn/replay/__init__.py | 0 .../attention/backends/gdn/replay/backend.py | 223 ++++ .../backends/gdn/replay/decode.py} | 5 +- .../attention/backends/gdn/triton/__init__.py | 0 .../attention/backends/gdn/triton/backend.py | 224 ++++ .../backends/gdn/triton/kernels/__init__.py | 0 .../backends/gdn/triton/kernels}/chunk.py | 16 +- .../backends/gdn/triton/kernels}/chunk_o.py | 7 +- .../triton/kernels}/chunk_scaled_dot_kkt.py | 5 +- .../gdn/triton/kernels}/fused_gdn_gating.py | 1 - .../gdn/triton/kernels}/fused_recurrent.py | 3 +- .../gdn/triton/kernels}/gdn_decode_pack.py | 0 .../triton/kernels}/mtp_fused_recurrent.py | 5 +- .../backends/gdn/triton/kernels}/wy_fast.py | 3 +- .../common/attention/backends/kda/__init__.py | 0 .../attention/backends/kda/replay/__init__.py | 0 .../attention/backends/kda/replay/backend.py | 177 +++ .../backends/kda/replay/decode.py} | 5 +- .../attention/backends/kda/triton/__init__.py | 0 .../attention/backends/kda/triton/backend.py | 155 +++ .../backends/kda/triton/kernels/__init__.py | 0 .../backends/kda/triton/kernels}/kda.py | 14 +- .../kda/triton/kernels}/kda_decode.py | 0 .../common/attention/backends/mla/__init__.py | 0 .../attention/backends/mla/fa3/__init__.py | 0 .../attention/backends/mla/fa3/backend.py | 125 ++ .../backends/mla/flashinfer/__init__.py | 0 .../backends/mla/flashinfer/backend.py | 229 ++++ .../attention/backends/mla/triton/__init__.py | 0 .../attention/backends/mla/triton/backend.py | 167 +++ .../backends/mla/triton/kernels/__init__.py | 0 .../mla/triton/kernels/decode}/__init__.py | 0 .../kernels/decode}/gqa_flash_decoding.py | 77 +- .../decode}/gqa_flash_decoding_config.py | 7 +- .../decode}/gqa_flash_decoding_stage1.py | 1 + .../decode}/gqa_flash_decoding_stage2.py | 3 +- .../mla/triton/kernels/prefill}/__init__.py | 0 .../context_flashattention_nopad_with_v.py | 10 +- .../attention/backends/sparse_mla/__init__.py | 0 .../backends/sparse_mla/flashmla/__init__.py | 0 .../backends/sparse_mla/flashmla/backend.py | 197 +++ .../backends/sparse_mla/flashmla/fp8.py | 172 +++ lightllm/common/attention/probe.py | 83 ++ lightllm/common/attention/registry.py | 87 ++ lightllm/common/attention/selection.py | 176 +++ lightllm/common/attention/session.py | 131 ++ lightllm/common/attention/varlen.py | 44 + .../attention_state/bounded_att/base_state.py | 22 + .../bounded_att/gdn/gdn_state.py | 27 +- .../bounded_att/gdn/replay_gdn_state.py | 40 +- .../attention_state/bounded_att/interface.py | 3 +- .../bounded_att/kda/kda_state.py | 27 +- .../bounded_att/kda/replay_kda_state.py | 40 +- .../sliding_window/fp8_e8m0_swa_state.py | 14 +- .../window_compressor_state.py | 10 +- lightllm/common/attention_state/commit.py | 34 + .../attention_state/full_att/base_state.py | 23 +- .../compressed_att/compressed_att_state.py | 14 +- .../attention_state/full_att/gqa/gqa_state.py | 82 +- .../full_att/gqa/group_quantized_state.py | 16 +- .../attention_state/full_att/gqa/ops.py | 44 + .../full_att/gqa/static_fp8_state.py | 25 +- .../attention_state/full_att/mla/mla_state.py | 14 +- .../packed_mla/fp8_e8m0_packed_mla_state.py | 15 +- .../full_att/packed_mla/packed_mla_state.py | 19 +- lightllm/common/attention_state/manager.py | 49 +- lightllm/common/basemodel/attention.py | 197 +++ .../common/basemodel/attention/__init__.py | 10 - .../common/basemodel/attention/base_att.py | 183 --- .../attention/create_linear_utils.py | 79 -- .../basemodel/attention/create_utils.py | 166 --- lightllm/common/basemodel/attention/fa3/fp.py | 333 ----- .../common/basemodel/attention/fa3/fp8.py | 205 ---- .../common/basemodel/attention/fa3/mla.py | 244 ---- .../common/basemodel/attention/fa3/neo.py | 81 -- .../basemodel/attention/flashinfer/fp.py | 291 ----- .../basemodel/attention/flashinfer/fp8.py | 146 --- .../basemodel/attention/flashinfer/mla.py | 272 ----- .../basemodel/attention/flashinfer/utils.py | 17 - .../basemodel/attention/import_utils.py | 6 - .../basemodel/attention/linear/flashqla.py | 8 - .../common/basemodel/attention/linear/gdn.py | 430 ------- .../common/basemodel/attention/linear/kda.py | 255 ---- .../basemodel/attention/linear/replay_gdn.py | 356 ------ .../basemodel/attention/linear/replay_kda.py | 257 ---- .../basemodel/attention/linear/triton.py | 7 - .../attention/nsa/dsv4_fp8_flashmla_sparse.py | 225 ---- .../attention/nsa/flashmla_sparse.py | 193 --- .../attention/nsa/fp8_flashmla_sparse.py | 204 ---- .../basemodel/attention/nsa/fp8_glm5_next.py | 238 ---- .../basemodel/attention/nsa/glm5_next.py | 100 -- .../common/basemodel/attention/triton/fp.py | 342 ------ .../basemodel/attention/triton/int4kv.py | 172 --- .../basemodel/attention/triton/int8kv.py | 207 ---- .../common/basemodel/attention/triton/mla.py | 160 --- .../common/basemodel/attention/triton/neo.py | 47 - .../common/basemodel/attention_adapters.py | 186 +++ .../basemodel/attention_vit/base_att.py | 38 - .../basemodel/attention_vit/create_utils.py | 47 - .../common/basemodel/attention_vit/fa3/fp.py | 60 - .../basemodel/attention_vit/triton/fp.py | 24 - lightllm/common/basemodel/basemodel.py | 364 +++--- lightllm/common/basemodel/batch_objs.py | 2 + lightllm/common/basemodel/cuda_graph.py | 192 ++- lightllm/common/basemodel/graph.py | 76 ++ lightllm/common/basemodel/infer_struct.py | 125 +- .../basemodel/layer_infer/base_layer_infer.py | 23 +- .../common/basemodel/prefill_cuda_graph.py | 120 +- .../triton_kernel/linear_att/__init__.py | 18 - .../triton_kernel/linear_att/fla/__init__.py | 11 - .../linear_att/fla/ops/__init__.py | 15 - lightllm/common/req_manager/base.py | 9 +- .../layer_infer/transformer_layer_infer.py | 47 +- lightllm/models/bloom/model.py | 17 +- .../layer_infer/transformer_layer_infer.py | 52 +- lightllm/models/deepseek2/model.py | 22 +- .../layer_infer/transformer_layer_infer.py | 74 +- lightllm/models/deepseek3_2/model.py | 15 +- .../models/deepseek_v4/attention/__init__.py | 0 .../models/deepseek_v4/attention/adapter.py | 47 + .../deepseek_v4/attention/sparse_mla.py | 193 +++ .../models/deepseek_v4/deepseek_v4_visual.py | 11 +- lightllm/models/deepseek_v4/infer_struct.py | 17 +- .../layer_infer/transformer_layer_infer.py | 88 +- lightllm/models/deepseek_v4/model.py | 71 +- lightllm/models/deepseek_v4/workspace.py | 34 + .../layer_infer/transformer_layer_infer.py | 6 + .../layer_infer/transformer_layer_infer.py | 41 +- lightllm/models/gemma4/model.py | 37 +- .../models/glm5_next/attention/__init__.py | 0 .../glm5_next/attention/fp8_sparse_mla.py | 237 ++++ .../models/glm5_next/attention/sparse_mla.py | 107 ++ lightllm/models/glm5_next/config.py | 6 +- lightllm/models/glm5_next/glm5_next_visual.py | 6 +- .../layer_infer/transformer_layer_infer.py | 135 +- lightllm/models/glm5_next/model.py | 79 +- .../context_flashattention_nopad_fp8.py | 744 +++++++++++ .../layer_infer/transformer_layer_infer.py | 4 +- lightllm/models/glm5_next_mtp/model.py | 40 +- .../layer_infer/transformer_layer_infer.py | 69 +- lightllm/models/gpt_oss/model.py | 12 +- .../layer_infer/transformer_layer_infer.py | 47 +- .../layer_infer/transformer_layer_infer.py | 22 +- lightllm/models/neo_chat/model.py | 47 +- .../layer_infer/transformer_layer_infer.py | 22 +- lightllm/models/neo_chat_moe/model.py | 46 +- lightllm/models/qwen2_5_vl/qwen2_5_visual.py | 28 +- lightllm/models/qwen2_vl/qwen2_visual.py | 26 +- .../layer_infer/transformer_layer_infer.py | 6 + lightllm/models/qwen3_dflash/model.py | 9 - .../layer_infer/transformer_layer_infer.py | 37 +- .../qwen3_omni_audio.py | 30 +- .../layer_infer/transformer_layer_infer.py | 89 +- lightllm/models/qwen3next/model.py | 40 +- .../layer_infer/transformer_layer_infer.py | 12 +- lightllm/platform/__init__.py | 72 ++ lightllm/platform/api.py | 96 ++ lightllm/platform/cuda.py | 17 + lightllm/server/api_cli.py | 69 +- lightllm/server/api_start.py | 68 +- lightllm/server/core/objs/start_args_type.py | 17 +- .../model_infer/mode_backend/base_backend.py | 94 +- .../mode_backend/chunked_prefill/impl.py | 29 +- .../mode_backend/dp_backend/impl.py | 37 +- .../mode_backend/multi_level_kv_cache.py | 48 +- .../mtp_speculative/planner/lightspec.py | 4 +- .../model_infer/mtp_speculative/utils.py | 34 +- lightllm/server/visualserver/__init__.py | 48 +- lightllm/server/visualserver/manager.py | 36 +- .../visualserver/model_infer/model_rpc.py | 47 +- .../visualserver/visual_only_manager.py | 44 +- lightllm/utils/backend_validator.py | 306 +++-- lightllm/utils/envs_utils.py | 12 +- lightllm/utils/start_utils.py | 8 +- lightllm/utils/torch_memory_saver_utils.py | 13 +- test/att_resources/validate_model_cache.py | 18 +- test/benchmark/kernels/benchmark_kda_ssm.py | 4 +- .../kernels/benchmark_replayssm_service.py | 2 +- .../static_inference/static_benchmark.py | 35 +- .../llama_gqa_diverse_decode_stage1_tuning.py | 20 +- .../llama_gqa_diverse_decode_stage2_tuning.py | 18 +- .../common/att_resources/manager_fixtures.py | 18 +- .../common/att_resources/model_helpers.py | 34 +- .../att_resources/test_backend_startup.py | 45 +- .../att_resources/test_bounded_state.py | 14 +- .../common/att_resources/test_capacity.py | 50 +- .../att_resources/test_dsv4_integration.py | 60 +- .../att_resources/test_fp8_compute_binding.py | 140 +-- .../att_resources/test_full_att_cpu_kernel.py | 12 +- .../common/att_resources/test_full_state.py | 24 +- .../att_resources/test_gqa_layout_identity.py | 57 - .../test_hybrid_layer_mapping.py | 4 +- .../test_layer_state_selection.py | 20 +- .../test_model_initialization.py | 15 +- .../test_model_storage_definition.py | 33 +- .../att_resources/test_mtp_state_commit.py | 36 +- .../att_resources/test_no_attention_draft.py | 23 +- .../test_quantized_full_state.py | 50 +- .../test_validate_model_cache.py | 9 +- unit_tests/common/attention/helpers.py | 418 +++++++ .../common/attention/test_backend_options.py | 152 +++ unit_tests/common/attention/test_contracts.py | 681 +++++++++++ .../common/attention/test_dense_numerics.py | 296 +++++ .../common/attention/test_linear_layout.py | 73 ++ .../common/attention/test_platform_runtime.py | 264 ++++ .../common/attention/test_provider_kind.py | 242 ++++ unit_tests/common/attention/test_state.py | 116 ++ .../common/attention/test_varlen_selection.py | 140 +++ unit_tests/common/attention/test_workspace.py | 60 + .../basemodel/attention/fa3/test_neo.py | 56 +- .../flashinfer/test_decode_graph_state.py | 100 -- .../flashinfer/test_decode_wrapper_init.py | 175 ++- .../flashinfer/test_mla_cuda_graph_plan.py | 30 +- .../attention/flashinfer/test_mla_nope.py | 3 +- .../attention/flashinfer/test_workspace.py | 86 +- .../basemodel/attention/linear/test_gdn.py | 419 ++----- .../basemodel/attention/test_dsv4_flashmla.py | 132 +- .../attention/test_fp8_logical_layer.py | 90 +- .../basemodel/attention/test_fp8_page_size.py | 44 +- .../basemodel/attention/test_gated_qkv.py | 23 + .../basemodel/attention/test_lazy_import.py | 257 +--- .../attention/test_state_lifetime.py | 186 +-- .../basemodel/test_cuda_graph_autotune.py | 88 +- .../common/basemodel/test_model_output.py | 42 +- .../basemodel/test_moe_route_info_manager.py | 16 +- .../common/basemodel/test_mtp_manager.py | 2 +- .../common/basemodel/test_overlap_utils.py | 7 +- .../test_prefill_cuda_graph_state.py | 53 +- .../gqa/mtp_diverse/test_mtp_diverse.py | 4 +- .../mtp_diverse/test_mtp_diverse_autotune.py | 47 +- .../gqa/mtp_diverse/test_stage2_autotune.py | 20 +- .../gqa/test_flash_decoding_autotune.py | 44 +- .../decode_att/int4kv/test_int4kv_autotune.py | 35 +- .../test_int8kv_flash_decoding_diverse.py | 29 +- ...st_int8kv_flash_decoding_diverse_stage1.py | 14 +- ...st_int8kv_flash_decoding_diverse_stage2.py | 12 +- ...st_int8kv_flash_decoding_diverse_stage3.py | 6 +- .../decode_att/int8kv/test_normal_autotune.py | 39 +- .../test_context_attention_fwd_neo.py | 5 +- .../test_context_flashattention_nopad.py | 6 +- .../test_context_flashattention_nopad1.py | 12 +- .../att/prefill_att/test_fa3_neo.py | 7 +- .../linear_att/replay_history.py | 6 +- .../linear_att/test_acceptance_fusion.py | 9 +- .../linear_att/test_causal_conv1d_mtp.py | 4 +- .../test_fused_recurrent_strided.py | 4 +- .../linear_att/test_kda_autotune.py | 11 +- .../test_mtp_fused_recurrent_autotune.py | 10 +- .../test_mtp_fused_recurrent_equiv.py | 6 +- .../linear_att/test_mtp_state_params.py | 7 +- .../linear_att/test_replayssm.py | 9 +- .../linear_att/test_replayssm_execution.py | 15 +- .../linear_att/test_ssm_autotune_memory.py | 4 +- .../decode_att/test_gqa_flash_decoding.py | 17 +- .../triton_kernel/test_diverse_utils.py | 17 +- .../common/test_deepseek4_paged_cache.py | 181 +-- .../common/test_linear_state_contract.py | 29 +- unit_tests/conftest.py | 1 + unit_tests/models/deepseek_v4/test_model.py | 3 +- unit_tests/models/deepseek_v4/test_vision.py | 28 +- .../models/deepseek_v4/test_workspace.py | 17 +- unit_tests/models/glm5_next/test_cache.py | 6 +- .../glm5_next/test_direct_nsa_indices.py | 297 ++--- unit_tests/models/glm5_next/test_fp8_kv.py | 1085 ++++++++++------- .../models/glm5_next/test_kda_conv_qkv.py | 58 +- unit_tests/models/glm5_next/test_kernels.py | 64 +- unit_tests/models/glm5_next/test_layers.py | 78 +- unit_tests/models/glm5_next/test_mtp.py | 103 +- .../glm5_next/test_mtp_prefill_cache.py | 14 +- unit_tests/models/glm5_next/test_pd_cache.py | 19 +- .../glm5_next/test_prefill_workspace.py | 181 ++- .../glm5_next/test_token_major_mla_bmm.py | 5 +- unit_tests/models/glm5_next/test_visual.py | 18 +- .../models/sensenova_loop/test_gdn_state.py | 104 +- .../vit/test_flash_attention_forward.py | 366 ++++-- unit_tests/server/core/objs/test_req.py | 12 +- .../server/core/objs/test_shm_req_manager.py | 17 +- .../mode_backend/test_multi_level_kv_cache.py | 48 +- .../model_infer/test_cpu_resource_layout.py | 26 +- unit_tests/server/test_api_start_deepep.py | 12 +- unit_tests/server/test_hybrid_start_args.py | 1 + unit_tests/test_static_benchmark_paged_kv.py | 49 +- unit_tests/utils/test_sgl_utils.py | 38 +- unit_tests/utils/test_speculative_utils.py | 174 ++- 383 files changed, 14357 insertions(+), 10545 deletions(-) create mode 100644 lightllm/common/attention/__init__.py create mode 100644 lightllm/common/attention/_environment.py create mode 100644 lightllm/common/attention/api/__init__.py create mode 100644 lightllm/common/attention/api/backend.py create mode 100644 lightllm/common/attention/api/batch.py create mode 100644 lightllm/common/attention/api/environment.py create mode 100644 lightllm/common/attention/api/execution.py create mode 100644 lightllm/common/attention/api/graph.py create mode 100644 lightllm/common/attention/api/inputs.py create mode 100644 lightllm/common/attention/api/spec.py create mode 100644 lightllm/common/attention/api/varlen.py create mode 100644 lightllm/common/attention/api/workspace.py rename lightllm/common/{basemodel/attention/fa3 => attention/backends}/__init__.py (100%) rename lightllm/common/{basemodel/attention/flashinfer/env_utils.py => attention/backends/_flashinfer.py} (99%) rename lightllm/common/{basemodel/attention/flashinfer => attention/backends/_shared}/__init__.py (100%) rename lightllm/common/{basemodel/attention => attention/backends/_shared}/linear/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/causal_conv1d.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/causal_conv1d_mtp.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/causal_conv1d_qkv.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/chunk_delta_h.py (98%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/cumsum.py (97%) create mode 100644 lightllm/common/attention/backends/_shared/linear/gdn.py rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/index.py (93%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/l2norm.py (100%) create mode 100644 lightllm/common/attention/backends/_shared/linear/layout.py rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/mtp_state_params.py (66%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/op.py (95%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_cursor.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_decode.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_delta.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/replayssm_fold.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/solve_tril.py (98%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/_shared/linear}/ssm_acceptance.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/_shared/linear}/utils.py (99%) rename lightllm/common/{basemodel/attention/nsa => attention/backends/_shared/prefill}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/prefill_att => attention/backends/_shared/prefill}/context_flashattention_nopad.py (99%) rename lightllm/common/{basemodel/attention/triton => attention/backends/dense}/__init__.py (100%) rename lightllm/common/{basemodel/attention_vit => attention/backends/dense/fa3}/__init__.py (100%) create mode 100644 lightllm/common/attention/backends/dense/fa3/backend.py create mode 100644 lightllm/common/attention/backends/dense/fa3/fp8.py create mode 100644 lightllm/common/attention/backends/dense/fa3/neo.py create mode 100644 lightllm/common/attention/backends/dense/fa3/varlen.py rename lightllm/common/{basemodel/attention_vit/fa3 => attention/backends/dense/flashinfer}/__init__.py (100%) create mode 100644 lightllm/common/attention/backends/dense/flashinfer/backend.py create mode 100644 lightllm/common/attention/backends/dense/flashinfer/fp8.py rename lightllm/common/{basemodel/attention_vit => attention/backends/dense}/sdpa/__init__.py (100%) rename lightllm/common/{basemodel/attention_vit/sdpa/fp.py => attention/backends/dense/sdpa/varlen.py} (51%) rename lightllm/common/{basemodel/attention_vit => attention/backends/dense}/triton/__init__.py (100%) create mode 100644 lightllm/common/attention/backends/dense/triton/backend.py create mode 100644 lightllm/common/attention/backends/dense/triton/int4.py create mode 100644 lightllm/common/attention/backends/dense/triton/int8.py rename lightllm/common/{basemodel/attention_vit/xformers => attention/backends/dense/triton/kernels}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att => attention/backends/dense/triton/kernels/decode}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode/gqa}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa => attention/backends/dense/triton/kernels/decode/gqa/flash}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/gqa/flash}/gqa_flash_decoding.py (87%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/gqa/flash}/gqa_flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/gqa/flash}/gqa_flash_decoding_stage2.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/gqa/flash}/gqa_flash_decoding_vsm.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/gqa_decode_flashattention_nopad.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/mtp_diverse/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/mtp_diverse/mtp_diverse_attn.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/mtp_diverse/stage1_single_token.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/gqa/mtp_diverse/stage2_single_token.py (98%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/gqa/flash_decoding => attention/backends/dense/triton/kernels/decode/int4kv}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int4kv/int4kv_flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int4kv/ppl_int4kv_flash_decoding.py (83%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/int4kv => attention/backends/dense/triton/kernels/decode/int8kv}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/int8kv_flash_decoding_diverse.py (77%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/int8kv_flash_decoding_diverse_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/int8kv_flash_decoding_diverse_stage2.py (98%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/int8kv_flash_decoding_diverse_stage3.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/normal/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/normal/int8kv_flash_decoding.py (84%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/normal/int8kv_flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/int8kv/normal/int8kv_flash_decoding_stage2.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/int8kv => attention/backends/dense/triton/kernels/decode/mha}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha => attention/backends/dense/triton/kernels/decode/mha/flash}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha/flash_decoding => attention/backends/dense/triton/kernels/decode/mha/flash}/flash_decoding.py (63%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha/flash_decoding => attention/backends/dense/triton/kernels/decode/mha/flash}/flash_decoding_stage1.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha/flash_decoding => attention/backends/dense/triton/kernels/decode/mha/flash}/flash_decoding_stage2.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha/flash_decoding => attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_nopad_att1.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_nopad_reduceV.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_nopad_softmax.py (99%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att => attention/backends/dense/triton/kernels/decode}/mha/stage3_decode_att/token_attention_softmax_and_reducev.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att => attention/backends/dense/triton/kernels/prefill}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/att/prefill_att => attention/backends/dense/triton/kernels/prefill}/context_attention_fwd_neo.py (99%) rename lightllm/{models/vit/triton_kernel/flashattention_nopad.py => common/attention/backends/dense/triton/kernels/prefill/varlen.py} (65%) create mode 100644 lightllm/common/attention/backends/dense/triton/neo.py create mode 100644 lightllm/common/attention/backends/dense/triton/varlen.py create mode 100644 lightllm/common/attention/backends/dense/varlen.py rename lightllm/common/{basemodel/triton_kernel/att/prefill_att => attention/backends/dense/xformers}/__init__.py (100%) rename lightllm/common/{basemodel/attention_vit/xformers/fp.py => attention/backends/dense/xformers/varlen.py} (61%) rename lightllm/common/{basemodel/triton_kernel/mla_att => attention/backends/gdn}/__init__.py (100%) create mode 100644 lightllm/common/attention/backends/gdn/flashqla/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/flashqla/backend.py create mode 100644 lightllm/common/attention/backends/gdn/replay/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/replay/backend.py rename lightllm/common/{basemodel/attention/linear/replay_gdn_decode.py => attention/backends/gdn/replay/decode.py} (98%) create mode 100644 lightllm/common/attention/backends/gdn/triton/__init__.py create mode 100644 lightllm/common/attention/backends/gdn/triton/backend.py create mode 100644 lightllm/common/attention/backends/gdn/triton/kernels/__init__.py rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/chunk.py (91%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/chunk_o.py (95%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/chunk_scaled_dot_kkt.py (97%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/gdn/triton/kernels}/fused_gdn_gating.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/fused_recurrent.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/gdn/triton/kernels}/gdn_decode_pack.py (100%) rename lightllm/common/{basemodel/triton_kernel/linear_att => attention/backends/gdn/triton/kernels}/mtp_fused_recurrent.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/gdn/triton/kernels}/wy_fast.py (97%) create mode 100644 lightllm/common/attention/backends/kda/__init__.py create mode 100644 lightllm/common/attention/backends/kda/replay/__init__.py create mode 100644 lightllm/common/attention/backends/kda/replay/backend.py rename lightllm/common/{basemodel/attention/linear/replay_kda_decode.py => attention/backends/kda/replay/decode.py} (98%) create mode 100644 lightllm/common/attention/backends/kda/triton/__init__.py create mode 100644 lightllm/common/attention/backends/kda/triton/backend.py create mode 100644 lightllm/common/attention/backends/kda/triton/kernels/__init__.py rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/kda/triton/kernels}/kda.py (99%) rename lightllm/common/{basemodel/triton_kernel/linear_att/fla/ops => attention/backends/kda/triton/kernels}/kda_decode.py (100%) create mode 100644 lightllm/common/attention/backends/mla/__init__.py create mode 100644 lightllm/common/attention/backends/mla/fa3/__init__.py create mode 100644 lightllm/common/attention/backends/mla/fa3/backend.py create mode 100644 lightllm/common/attention/backends/mla/flashinfer/__init__.py create mode 100644 lightllm/common/attention/backends/mla/flashinfer/backend.py create mode 100644 lightllm/common/attention/backends/mla/triton/__init__.py create mode 100644 lightllm/common/attention/backends/mla/triton/backend.py create mode 100644 lightllm/common/attention/backends/mla/triton/kernels/__init__.py rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/gqa_flash_decoding.py (71%) rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/gqa_flash_decoding_config.py (99%) rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/gqa_flash_decoding_stage1.py (99%) rename lightllm/common/{basemodel/triton_kernel/mla_att/decode_att => attention/backends/mla/triton/kernels/decode}/gqa_flash_decoding_stage2.py (99%) rename lightllm/common/{basemodel/triton_kernel/mla_att/prefill_att => attention/backends/mla/triton/kernels/prefill}/__init__.py (100%) rename lightllm/common/{basemodel/triton_kernel/mla_att/prefill_att => attention/backends/mla/triton/kernels/prefill}/context_flashattention_nopad_with_v.py (100%) create mode 100644 lightllm/common/attention/backends/sparse_mla/__init__.py create mode 100644 lightllm/common/attention/backends/sparse_mla/flashmla/__init__.py create mode 100644 lightllm/common/attention/backends/sparse_mla/flashmla/backend.py create mode 100644 lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py create mode 100644 lightllm/common/attention/probe.py create mode 100644 lightllm/common/attention/registry.py create mode 100644 lightllm/common/attention/selection.py create mode 100644 lightllm/common/attention/session.py create mode 100644 lightllm/common/attention/varlen.py create mode 100644 lightllm/common/attention_state/commit.py create mode 100644 lightllm/common/attention_state/full_att/gqa/ops.py create mode 100644 lightllm/common/basemodel/attention.py delete mode 100644 lightllm/common/basemodel/attention/__init__.py delete mode 100644 lightllm/common/basemodel/attention/base_att.py delete mode 100644 lightllm/common/basemodel/attention/create_linear_utils.py delete mode 100644 lightllm/common/basemodel/attention/create_utils.py delete mode 100644 lightllm/common/basemodel/attention/fa3/fp.py delete mode 100644 lightllm/common/basemodel/attention/fa3/fp8.py delete mode 100644 lightllm/common/basemodel/attention/fa3/mla.py delete mode 100644 lightllm/common/basemodel/attention/fa3/neo.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/fp.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/fp8.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/mla.py delete mode 100644 lightllm/common/basemodel/attention/flashinfer/utils.py delete mode 100644 lightllm/common/basemodel/attention/import_utils.py delete mode 100644 lightllm/common/basemodel/attention/linear/flashqla.py delete mode 100644 lightllm/common/basemodel/attention/linear/gdn.py delete mode 100644 lightllm/common/basemodel/attention/linear/kda.py delete mode 100644 lightllm/common/basemodel/attention/linear/replay_gdn.py delete mode 100644 lightllm/common/basemodel/attention/linear/replay_kda.py delete mode 100644 lightllm/common/basemodel/attention/linear/triton.py delete mode 100644 lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py delete mode 100644 lightllm/common/basemodel/attention/nsa/flashmla_sparse.py delete mode 100644 lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py delete mode 100644 lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py delete mode 100644 lightllm/common/basemodel/attention/nsa/glm5_next.py delete mode 100644 lightllm/common/basemodel/attention/triton/fp.py delete mode 100644 lightllm/common/basemodel/attention/triton/int4kv.py delete mode 100644 lightllm/common/basemodel/attention/triton/int8kv.py delete mode 100644 lightllm/common/basemodel/attention/triton/mla.py delete mode 100644 lightllm/common/basemodel/attention/triton/neo.py create mode 100644 lightllm/common/basemodel/attention_adapters.py delete mode 100644 lightllm/common/basemodel/attention_vit/base_att.py delete mode 100644 lightllm/common/basemodel/attention_vit/create_utils.py delete mode 100644 lightllm/common/basemodel/attention_vit/fa3/fp.py delete mode 100644 lightllm/common/basemodel/attention_vit/triton/fp.py create mode 100644 lightllm/common/basemodel/graph.py delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py delete mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py create mode 100644 lightllm/models/deepseek_v4/attention/__init__.py create mode 100644 lightllm/models/deepseek_v4/attention/adapter.py create mode 100644 lightllm/models/deepseek_v4/attention/sparse_mla.py create mode 100644 lightllm/models/glm5_next/attention/__init__.py create mode 100644 lightllm/models/glm5_next/attention/fp8_sparse_mla.py create mode 100644 lightllm/models/glm5_next/attention/sparse_mla.py create mode 100644 lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py create mode 100644 lightllm/platform/__init__.py create mode 100644 lightllm/platform/api.py create mode 100644 lightllm/platform/cuda.py create mode 100644 unit_tests/common/attention/helpers.py create mode 100644 unit_tests/common/attention/test_backend_options.py create mode 100644 unit_tests/common/attention/test_contracts.py create mode 100644 unit_tests/common/attention/test_dense_numerics.py create mode 100644 unit_tests/common/attention/test_linear_layout.py create mode 100644 unit_tests/common/attention/test_platform_runtime.py create mode 100644 unit_tests/common/attention/test_provider_kind.py create mode 100644 unit_tests/common/attention/test_state.py create mode 100644 unit_tests/common/attention/test_varlen_selection.py create mode 100644 unit_tests/common/attention/test_workspace.py delete mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py create mode 100644 unit_tests/common/basemodel/attention/test_gated_qkv.py diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 7edb0c9701..b07e43ff4c 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -127,7 +127,24 @@ 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` 创建共享的 `ReqManager`，绑定模型并注册各自的 State。`finish_state_registration()` 后，`prepare_attention_backends()` 创建 `model.attention: ModelAttention`，统一持有分组、层路由、workspace 和存储代次。每组是 `dict[AttentionPhase, AttentionBackend]`，后端实例只实现一个阶段。`model.attention.reserve(model)` 在容量估算前预留临时区，标定和存储分配完成后由 `model.attention.bind(model)` 创建各层的 `LayerAttention`。 + +层通过 `describe_attention(phase)` 返回 family、features、window 和 variant。family adapter 将这些需求和 State 布局组成 `LayerSpec`，并从 `State.layer_resources(layer_id)` 借用资源；FP8 scale 的选择和切片由 State 负责。`AttentionFeature` 集合统一检查能力，`CacheEncoding` 描述编码，`CacheLayout.state_mode` 区分 native/replay。模型可通过 `attention_policy()` 返回 `BackendPolicy(allow, prefer, extra)` 限制或扩展注册表。 + +注册表枚举只加载后端定义，能力声明以类属性为准；同一后端的阶段共性放在共享 mixin。kernel 和 provider 的执行依赖在使用处导入，候选后端的环境检查才验证可选依赖。Neo 扩展缺失时，`auto` 可回退，显式选择则报错，不影响其他模型的后端枚举。 + +后端构造只接收 frozen `BackendEnv` 和 `LayerSpec`。Env 显式提供设备 runtime、网络配置、TP 信息、HOLD 请求号、decode draft step、持久状态槽宽、长度上界和 allocator。`infer_state.attention: AttentionSession` 管理批次 context、owner/phase 校验、输入复制及输出写回。后端仅提供 `allocate_context`、图外每次重放执行的 `plan_host`、图内的 `build_device_metadata` 和 `forward`。模型通过 `with model.attention.session(infer_state)` 一次完成准备及 lease 释放；decode 捕获显式传入 Graph 长度上界。Session 位于 attention 模块，只接收 batch、backend 分组和容量，不依赖模型。所有 context 共享一个捕获 batch；Graph 先校验所有 microbatch 的复制计划，再统一复制和调用 host plan。 + +`AttentionBackend(env, spec)` 是唯一后端基类，`spec` 只存于后端。`LayerAttention` 只持有 backend、资源和路由；prefill/decode 专属构造及 helper 各归其阶段类。量化组长来自 `spec.cache`，allocator 直接注入后端。层的 `attention_adapter(family)` 选取一次 adapter，描述与绑定复用该实例；模型专有 adapter 放在模型目录。 + +```python +output = self.decode_attention.forward(infer_state.attention, DenseInputs(q)) +context = self.decode_attention.context(infer_state.attention) +``` + +KV/Conv/SSM 的分配、checkpoint、恢复和释放仍归 State。存储代次变化后必须重新绑定层并重新捕获 Graph。`RequestGenerationStamp.stamp()` 固定 MTP 输出对应的请求槽代次，避免释放后复用的槽接收旧结果。 + +视觉/音频编码器使用 `VarlenAttentionOp.forward(inputs, out=...)`，直接提供 `VarlenAttentionInputs` 中的 Q/K/V 和序列边界。独立选型失败报告 `VarlenSelectionError`；只有支持 Graph 的算子才能捕获。`--vit_att_backend` 不变。dense/MLA/sparse MLA 与 GDN/KDA 使用独立的 `--llm_*_att_backend` 和 `--llm_*_linear_att_backend`。旧的 `--llm_prefill_att_backend fa3 flashqla` 等双值形式仍接受并提示弃用，第二项写入对应 linear 选项。显式选择失败时报错，只有 `auto` 尝试其他 provider。 公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。两个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index cd71541598..253a4b3660 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -127,7 +127,24 @@ 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` binds their shared `ReqManager` and registers each model's States. After `finish_state_registration()`, `prepare_attention_backends()` creates `model.attention: ModelAttention`, which owns groups, layer routes, workspaces and the storage generation. A group is a `dict[AttentionPhase, AttentionBackend]`; each backend implements one phase. `model.attention.reserve(model)` reserves scratch before capacity estimation, and `model.attention.bind(model)` creates each `LayerAttention` after calibration and storage allocation. + +Layers return family, features, window and variant from `describe_attention(phase)`. Family adapters combine that description with the State layout to create `LayerSpec` and borrow resources through `State.layer_resources(layer_id)`. States select and slice their FP8 scales. `AttentionFeature` sets provide capability checks, `CacheEncoding` identifies encoding, and `CacheLayout.state_mode` selects native/replay storage. Models return `BackendPolicy(allow, prefer, extra)` from `attention_policy()` to constrain or extend the registry. + +Registry enumeration loads backend definitions only. Class attributes declare capabilities, with shared phase capabilities in a common mixin. Import kernel and provider execution dependencies where they are used; candidate environment checks validate optional dependencies. A missing Neo extension allows fallback in `auto` mode and fails explicit selection, without affecting backend enumeration for other models. + +Backends receive only a frozen `BackendEnv` and `LayerSpec`. The environment supplies the device runtime, network configuration, TP coordinates, HOLD request, decode draft step, persistent state slot count, sequence capacity and allocator. `infer_state.attention: AttentionSession` owns batch contexts, owner/phase validation, input copies and output copy-back. Backend hooks are `allocate_context`, `plan_host` (outside the graph on every replay), `build_device_metadata` (captured device work), and `forward`. `with model.attention.session(infer_state)` combines preparation and lease release; decode capture passes its Graph length bound explicitly. The Session lives in the attention module and accepts batches, backend groups and capacities without depending on the model. Contexts share one captured batch. Graph replay validates copy plans for every microbatch before copying inputs and running host plans. + +`AttentionBackend(env, spec)` is the only backend base class and owns the spec. `LayerAttention` holds the backend, resources and route. Phase-specific constructors and helpers live on their respective phase classes. Quantization group size comes from `spec.cache`; the allocator is injected directly into the backend. The layer selects `attention_adapter(family)` once and reuses it for description and binding. Model-specific adapters live in the model directory. + +```python +output = self.decode_attention.forward(infer_state.attention, DenseInputs(q)) +context = self.decode_attention.context(infer_state.attention) +``` + +States retain allocation, checkpoint, restore and release ownership for KV/Conv/SSM. A storage generation change requires layer rebinding and graph recapture. `RequestGenerationStamp.stamp()` associates MTP outputs with request slot generations to reject results for released and reused slots. + +Vision/audio encoders use `VarlenAttentionOp.forward(inputs, out=...)` with Q/K/V and sequence boundaries in `VarlenAttentionInputs`. Selection reports `VarlenSelectionError`; only graph-capable operators support capture. `--vit_att_backend` is unchanged. Dense/MLA/sparse MLA and GDN/KDA use separate `--llm_*_att_backend` and `--llm_*_linear_att_backend` options. Legacy two-value arguments such as `--llm_prefill_att_backend fa3 flashqla` remain accepted with a deprecation warning; the second value sets the corresponding linear option. Explicit selection fails on an unsupported provider; only `auto` tries alternatives. 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. diff --git a/lightllm/common/attention/__init__.py b/lightllm/common/attention/__init__.py new file mode 100644 index 0000000000..c983261044 --- /dev/null +++ b/lightllm/common/attention/__init__.py @@ -0,0 +1,11 @@ +from .api import ( + AttentionBackend, + AttentionBatch, + AttentionFamily, + AttentionPhase, + CacheLayout, + ContextCapacity, + LayerAttention, + LayerSpec, + UnsupportedAttention, +) diff --git a/lightllm/common/attention/_environment.py b/lightllm/common/attention/_environment.py new file mode 100644 index 0000000000..85fdec4981 --- /dev/null +++ b/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 diff --git a/lightllm/common/attention/api/__init__.py b/lightllm/common/attention/api/__init__.py new file mode 100644 index 0000000000..0cc54d18d0 --- /dev/null +++ b/lightllm/common/attention/api/__init__.py @@ -0,0 +1,28 @@ +from .backend import ( + AttentionBackend, + AttentionContext, + LayerAttention, + UnsupportedAttention, +) +from .batch import AttentionBatch, ContextCapacity +from .execution import WorkspaceBudget, ServingConstraints +from .graph import CaptureHooks +from .inputs import ( + DenseInputs, + DenseResources, + GDNInputs, + GDNOutput, + GDNResources, + KDAInputs, + KDAResources, + MLADecodeInputs, + MLAResources, + MLAPrefillInputs, + SparseMLADecodeInputs, + SparseMLAPrefillInputs, + VarlenAttentionInputs, +) +from .spec import AttentionFamily, AttentionFeature, AttentionPhase, CacheEncoding, CacheLayout, LayerSpec +from .environment import BackendEnv +from .varlen import VarlenAttentionOp +from .workspace import AttentionWorkspace diff --git a/lightllm/common/attention/api/backend.py b/lightllm/common/attention/api/backend.py new file mode 100644 index 0000000000..f86bdd0ad3 --- /dev/null +++ b/lightllm/common/attention/api/backend.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Generic, TYPE_CHECKING, TypeVar + +import torch + +from lightllm.platform import DeviceCaps, Platform + +from .batch import AttentionBatch, ContextCapacity +from .environment import BackendEnv +from .execution import ServingConstraints, WorkspaceBudget +from .spec import AttentionFamily, AttentionPhase, CacheEncoding, LayerSpec +from .workspace import AttentionWorkspace + + +if TYPE_CHECKING: + from lightllm.common.attention.session import AttentionSession + + +@dataclass(slots=True) +class AttentionContext: + capacity: ContextCapacity + capture: bool + workspace: AttentionWorkspace + batch: AttentionBatch | None = None + current_max_kv_seq_len: int = 0 + backend_identity: AttentionBackend | None = None + + +Context = TypeVar("Context", bound=AttentionContext) +Inputs = TypeVar("Inputs") +Output = TypeVar("Output") +Resources = TypeVar("Resources") + + +class UnsupportedAttention(ValueError): + pass + + +@dataclass(frozen=True) +class LayerAttention(Generic[Resources, Context, Inputs, Output]): + group_id: int + logical_layer_id: int + resources: Resources + backend: AttentionBackend[Context, Inputs, Output] + + def context(self, session: AttentionSession) -> Context: + return session.context(self) + + def forward(self, session: AttentionSession, inputs: Inputs, *, out: torch.Tensor | None = None) -> Output: + return session.forward(self, inputs, out=out) + + +class AttentionBackend(ABC, Generic[Context, Inputs, Output]): + zero_initialize_cache = False + requires_cpu_cache_sync = False + + family: AttentionFamily + provider: str + phase: AttentionPhase + cache_encodings = frozenset({CacheEncoding.PLAIN}) + state_mode = "native" + variant = "" + supported_features = frozenset() + supported_platforms = frozenset({Platform.CUDA}) + supports_graph = True + activation_dtypes = frozenset({torch.float16, torch.bfloat16, torch.float32}) + dependency = "triton" + min_device_capability = None + supports_expert_parallel = True + max_head_size = None + + def __init__(self, env: BackendEnv, spec: LayerSpec): + self.runtime = env.runtime + self.spec = spec + self.network_config = env.network_config + self.draft_step = env.draft_step if spec.use_mtp else 0 + self.hold_req_id = env.hold_req_id + self.alloc_tensor = env.alloc_tensor + + @classmethod + def check_environment(cls, caps: DeviceCaps) -> None: + from lightllm.common.attention._environment import check_attention_environment + + check_attention_environment( + caps, + provider=cls.provider, + supported_platforms=cls.supported_platforms, + min_device_capability=cls.min_device_capability, + dependency=cls.dependency, + ) + + @classmethod + def check_support(cls, spec: LayerSpec, caps=None, constraints=None) -> None: + if (spec.family, spec.phase, spec.variant) != (cls.family, cls.phase, cls.variant): + raise UnsupportedAttention(f"{cls.__name__} does not implement {spec.family.value} {spec.phase.value}") + if spec.cache.encoding not in cls.cache_encodings or spec.cache.state_mode != cls.state_mode: + raise UnsupportedAttention( + f"cache layout {spec.cache.encoding.value}/{spec.cache.state_mode} is unsupported" + ) + if spec.dtype not in cls.activation_dtypes: + raise UnsupportedAttention(f"activation dtype {spec.dtype} is unsupported") + missing = spec.features - cls.supported_features + if missing: + raise UnsupportedAttention("unsupported attention features: " + ", ".join(sorted(f.value for f in missing))) + if cls.max_head_size is not None and spec.head_size > cls.max_head_size: + raise UnsupportedAttention(f"head size exceeds {cls.max_head_size}") + if constraints is not None: + if constraints.expert_parallel and not cls.supports_expert_parallel: + raise UnsupportedAttention("this backend is incompatible with expert-parallel MoE execution") + if spec.phase in constraints.graph_phases and ( + not cls.supports_graph or (caps is not None and not caps.supports_graph) + ): + raise UnsupportedAttention(f"{spec.phase.value} Graph is unsupported") + + @property + def kernel_page_size(self) -> int: + return self.spec.cache.page_size + + @abstractmethod + def allocate_context(self, capacity: ContextCapacity, *, capture: bool, workspace: AttentionWorkspace) -> Context: + pass + + def plan_host(self, context: Context) -> None: + pass + + def build_device_metadata(self, context: Context) -> None: + pass + + @abstractmethod + def forward( + self, layer: LayerAttention[Any, Context, Inputs, Output], context: Context, inputs: Inputs, *, out=None + ): + pass + + def resource_requirements(self, capacity: ContextCapacity, *, capture: bool): + return WorkspaceBudget() + + def create_capture_hooks(self, context: Context): + return None diff --git a/lightllm/common/attention/api/batch.py b/lightllm/common/attention/api/batch.py new file mode 100644 index 0000000000..d570a50043 --- /dev/null +++ b/lightllm/common/attention/api/batch.py @@ -0,0 +1,116 @@ +from dataclasses import dataclass, field, 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 = field(metadata={"shared_storage": True}) + 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 是共享存储，必须保持原地址，不能通过整表拷贝替换。 + """ + for dst, src in self.input_copies(source): + dst.copy_(src, non_blocking=True) + + def input_copies(self, source): + 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 item in fields(self): + dst, src = getattr(self, item.name), getattr(source, item.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 {item.name} changed presence") + if (dst.shape, dst.dtype, dst.device) != (src.shape, src.dtype, src.device): + raise ValueError(f"captured input {item.name} changed shape, dtype or device") + if item.metadata.get("shared_storage") 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)) + return copies + + +@dataclass(frozen=True, kw_only=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") diff --git a/lightllm/common/attention/api/environment.py b/lightllm/common/attention/api/environment.py new file mode 100644 index 0000000000..ec9cdd887a --- /dev/null +++ b/lightllm/common/attention/api/environment.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass +from types import MappingProxyType +from typing import Callable, Mapping + +from lightllm.platform import DeviceRuntime + + +@dataclass(frozen=True, kw_only=True) +class BackendEnv: + runtime: DeviceRuntime + network_config: Mapping + hold_req_id: int + draft_step: int + max_seq_len: int + tp_size: int + alloc_tensor: Callable + state_slots: int = 1 + tp_rank: int = 0 + layer_ids: tuple[int, ...] = () + prefill_graph: bool = False + max_prefill_tokens: int = 0 + decode_only: bool = False + + def __post_init__(self): + object.__setattr__(self, "network_config", MappingProxyType(dict(self.network_config))) diff --git a/lightllm/common/attention/api/execution.py b/lightllm/common/attention/api/execution.py new file mode 100644 index 0000000000..ba78fb8d08 --- /dev/null +++ b/lightllm/common/attention/api/execution.py @@ -0,0 +1,35 @@ +from dataclasses import dataclass +from typing import NamedTuple + +import torch + +from .spec import AttentionPhase + + +@dataclass(frozen=True, kw_only=True) +class ServingConstraints: + """启动时确定的执行约束，与 LayerSpec 的计算语义分开保存。""" + + graph_phases: frozenset[AttentionPhase] = frozenset() + microbatch_count: int = 1 + max_batch_size: int = 1 + max_seq_len: int = 1 + expert_parallel: bool = False + transfer_paths: frozenset[str] = frozenset() + + +class WorkspaceRequirement(NamedTuple): + name: str + elements: int + dtype: torch.dtype + + +@dataclass(frozen=True) +class WorkspaceBudget: + """捕获前的资源上界；workspace 项为 (名称, 元素数, dtype)。 + + 同一执行槽位按名称取最大容量，并发槽位分别分配。主机长度快照只在需要 + 它的实现中创建；持久存储的初始化由 State 执行，不能在 forward 中反复清零。 + """ + + workspace: tuple[WorkspaceRequirement, ...] = () diff --git a/lightllm/common/attention/api/graph.py b/lightllm/common/attention/api/graph.py new file mode 100644 index 0000000000..b2c2872cfb --- /dev/null +++ b/lightllm/common/attention/api/graph.py @@ -0,0 +1,24 @@ +from typing import Protocol + + +class CaptureHooks(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 diff --git a/lightllm/common/attention/api/inputs.py b/lightllm/common/attention/api/inputs.py new file mode 100644 index 0000000000..c4cd732319 --- /dev/null +++ b/lightllm/common/attention/api/inputs.py @@ -0,0 +1,179 @@ +"""区分初始化时绑定的层资源与每次 forward 的 tensor 输入。 + +*Resources 随 LayerAttention 保存，借用权重以及 State 管理的 KV / Conv / SSM 视图； +*Inputs 由模型层为当前调用构造，批次级请求索引和 provider 派生 metadata 则由 context 提供。 +这些数据结构只固定字段绑定，不禁止 backend 原地更新所借用的状态 tensor。 +""" + +from dataclasses import dataclass +from typing import NamedTuple, Optional, Protocol + +import torch + + +@dataclass(frozen=True) +class DenseResources: + key: torch.Tensor + value: torch.Tensor + alibi: Optional[torch.Tensor] = None + sink: Optional[torch.Tensor] = None + query_scale: Optional[torch.Tensor] = None + key_scale: torch.Tensor | float | None = None + value_scale: torch.Tensor | float | None = None + + +@dataclass(frozen=True) +class MLAResources: + kv: torch.Tensor + kv_lora_rank: int + qk_rope_head_dim: int + + +class ReplaySSMState(Protocol): + """Borrowed replay storage; State retains allocation, checkpoint and release ownership.""" + + ssm_state: torch.Tensor + keys: torch.Tensor + deltas: torch.Tensor + gates: torch.Tensor + raw_keys: torch.Tensor + raw_values: torch.Tensor + betas: torch.Tensor + cursors: torch.Tensor + tile_sync_mark: torch.Tensor + num_key_heads: int + capacity: int + verify_width: int + compute_mode: str + gate_lower_bound: float + hold_req_id: int + kda: bool + + def storage_layer(self, logical_layer_id: int) -> int: + ... + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: + ... + + +@dataclass(frozen=True) +class RecurrentStorage: + conv: torch.Tensor + ssm: torch.Tensor + mtp_state_index: torch.Tensor + replay_state: ReplaySSMState | None = None + + +@dataclass(frozen=True) +class GDNResources: + conv: torch.Tensor + ssm: torch.Tensor + conv_weight: torch.Tensor + conv_bias: Optional[torch.Tensor] + a_log: torch.Tensor + dt_bias: torch.Tensor + mtp_state_index: torch.Tensor | None = None + replay_state: ReplaySSMState | None = None + + +@dataclass(frozen=True) +class KDAResources: + conv: torch.Tensor + ssm: torch.Tensor + conv_weight: torch.Tensor + a_log: torch.Tensor + dt_bias: torch.Tensor + mtp_state_index: torch.Tensor | None = None + replay_state: ReplaySSMState | None = None + + +@dataclass(frozen=True) +class DenseInputs: + """Dense 路径在调用前已写入 KV cache，因此这里仅传 query。""" + + query: torch.Tensor + + +class VarlenAttentionInputs(NamedTuple): + """非因果 self-attention 的 Q/K/V，布局均为 [total_tokens, heads, head_dim]。 + + cu_seqlens 是同设备上的 int32 前缀和，始于 0、终于 total_tokens；相邻边界之间 + 独立做 attention，允许空序列。max_seqlen 是最长序列长度的容量上界，由调用方提供以避免 GPU 同步。 + 每层只打包 tensor 引用，使用 NamedTuple 保持字段不可变并减少构造开销。 + """ + + query: torch.Tensor + key: torch.Tensor + value: torch.Tensor + cu_seqlens: torch.Tensor + max_seqlen: int + + def validate(self) -> None: + """Check sequence boundaries without copying CUDA lengths to the host. + + max_seqlen is an upper bound (Graph buckets may reserve a larger value). + CUDA value checks are asynchronous device assertions, including in Graphs; + invalid boundaries are input errors, not a reason to try another provider. + """ + q, k, v, cu, maximum = self + if q.ndim != 3 or q.shape != k.shape or q.shape != v.shape: + raise ValueError("varlen self-attention requires matching [tokens, heads, dim] Q/K/V") + if q.device != k.device or q.device != v.device or q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("Q/K/V must have the same dtype and device") + if cu.ndim != 1 or cu.numel() < 2 or cu.dtype != torch.int32 or cu.device != q.device: + raise ValueError("cu_seqlens must be a same-device int32 vector with at least two entries") + if maximum < 0 or (q.shape[0] and maximum == 0): + raise ValueError("nonempty attention requires a positive max_seqlen") + lengths = cu[1:] - cu[:-1] + valid = (cu[0] == 0) & (cu[-1] == q.shape[0]) & (lengths >= 0).all() & (lengths <= maximum).all() + message = "cu_seqlens must span all tokens with nonnegative lengths bounded by max_seqlen" + if cu.device.type == "cuda": + torch._assert_async(valid, message) + elif not bool(valid): + raise ValueError(message) + + +@dataclass(frozen=True) +class MLAPrefillInputs: + query: torch.Tensor + key_nope: torch.Tensor + key_rope: torch.Tensor + value: torch.Tensor + + +@dataclass(frozen=True) +class MLADecodeInputs: + query_nope: torch.Tensor + query_rope: torch.Tensor + + +@dataclass(frozen=True) +class SparseMLAPrefillInputs: + query: torch.Tensor + topk_mem_indices: torch.Tensor + topk_indices: Optional[torch.Tensor] = None + prefill_cache_kv: Optional[torch.Tensor] = None + + +@dataclass(frozen=True) +class SparseMLADecodeInputs: + query_nope: torch.Tensor + query_rope: torch.Tensor + topk_mem_indices: torch.Tensor + + +@dataclass(frozen=True) +class GDNInputs: + mixed_qkvzba: torch.Tensor + + +@dataclass(frozen=True) +class KDAInputs: + mixed_qkv: torch.Tensor + raw_gate: torch.Tensor + raw_beta: torch.Tensor + + +class GDNOutput(NamedTuple): + hidden_states: torch.Tensor + gate: torch.Tensor diff --git a/lightllm/common/attention/api/spec.py b/lightllm/common/attention/api/spec.py new file mode 100644 index 0000000000..8d5a2f9f6b --- /dev/null +++ b/lightllm/common/attention/api/spec.py @@ -0,0 +1,121 @@ +from dataclasses import dataclass +from enum import Enum +from typing import Optional + +import torch + + +class AttentionFamily(str, Enum): + """计算类型；Triton、FA3 等实现库由 selector 另行选择。""" + + DENSE = "dense" + MLA = "mla" + SPARSE_MLA = "sparse_mla" + GDN = "gdn" + KDA = "kda" + + @property + def is_recurrent(self): + return self in (AttentionFamily.GDN, AttentionFamily.KDA) + + +class AttentionFeature(str, Enum): + ALIBI = "alibi" + SINK = "sink" + SLIDING_WINDOW = "sliding_window" + NONCAUSAL = "noncausal" + MTP = "mtp" + DYNAMIC_VERIFY = "dynamic_verify" + MTP_WINDOW = "mtp_window" + ALIBI_WINDOW = "alibi_window" + + +class CacheEncoding(str, Enum): + PLAIN = "plain" + INT4 = "int4kv" + INT8 = "int8kv" + FP8_HEAD = "fp8kv_sph" + FP8_TENSOR = "fp8kv_spt" + FP8_SPARSE_MLA = "fp8kv_dsa" + + @classmethod + def from_option(cls, value): + return cls.PLAIN if value in (None, "None", "gdn", "kda") else cls(value) + + +class AttentionPhase(str, Enum): + PREFILL = "prefill" + DECODE = "decode" + + +@dataclass(frozen=True) +class CacheLayout: + """存储布局的静态描述，不持有或分配存储 tensor。 + + prefill / decode 必须接受同一布局；dtype、量化粒度和 scale 布局也属于兼容性条件。 + storage_schema 记录 State 显式声明的布局版本及配置，排除逻辑层号，使同布局的层可以归为一组。 + """ + + encoding: CacheEncoding + page_size: int + dtype: torch.dtype + key_heads: int + key_dim: int + value_dim: int + state_mode: str = "native" + state_dtype: Optional[torch.dtype] = None + storage_schema: tuple = () + quant_group_size: int = 0 + scale_mode: str = "none" + + def __post_init__(self): + object.__setattr__(self, "encoding", CacheEncoding.from_option(self.encoding)) + if self.page_size <= 0: + raise ValueError("cache page_size must be positive") + + +@dataclass(frozen=True) +class LayerSpec: + """一个层在一个 phase 下的静态计算需求，用于能力检查和分组。 + + 初始化按完整的 (prefill_spec, decode_spec) 分组，只有相同需求的层才共享 backend / context。 + 层权重和存储视图放在 LayerAttention.resources，批次长度与请求索引放在 AttentionBatch。 + """ + + family: AttentionFamily + phase: AttentionPhase + dtype: torch.dtype + cache: CacheLayout + num_q_heads: int + head_size: int + value_head_size: int + sliding_window: Optional[tuple[int, int]] = None + features: frozenset[AttentionFeature] = frozenset() + variant: str = "" + softmax_scale: float | None = None + + def __post_init__(self): + features = set(self.features) + if self.sliding_window is not None: + features.add(AttentionFeature.SLIDING_WINDOW) + if AttentionFeature.ALIBI in features: + features.add(AttentionFeature.ALIBI_WINDOW) + if AttentionFeature.MTP in features: + features.add(AttentionFeature.MTP_WINDOW) + object.__setattr__(self, "features", frozenset(features)) + if min(self.num_q_heads, self.head_size, self.value_head_size) <= 0: + raise ValueError("attention head counts and dimensions must be positive") + if self.family.is_recurrent and self.sliding_window is not None: + raise ValueError("recurrent attention cannot use a paged sliding window") + + @property + def use_mtp(self): + return AttentionFeature.MTP in self.features + + @property + def dynamic_verify(self): + return AttentionFeature.DYNAMIC_VERIFY in self.features + + @property + def causal(self): + return AttentionFeature.NONCAUSAL not in self.features diff --git a/lightllm/common/attention/api/varlen.py b/lightllm/common/attention/api/varlen.py new file mode 100644 index 0000000000..a1c7743c03 --- /dev/null +++ b/lightllm/common/attention/api/varlen.py @@ -0,0 +1,54 @@ +from abc import ABC, abstractmethod + +import torch + +from .inputs import VarlenAttentionInputs + + +class VarlenAttentionOp(ABC): + """无持久 KV 的双向 dense attention，供视觉和音频编码器使用。 + + Q/K/V 与变长边界由调用方直接传入，不需要文本推理的请求表或 cache context。 + 实现不保留批次 tensor；调用方持有 out，返回值必须引用同一块输出存储。 + out 的 shape/dtype/device 与 query 一致，head_dim 连续，其他维度允许无重叠的 stride； + 不允许与 Q/K/V 共享存储。每次 forward 在启动 kernel 前验证这些 metadata。 + SDPA / xFormers 需要在 CPU 读取序列边界；只有声明 supports_graph 的实现可整段捕获。 + """ + + supports_graph = False + + @classmethod + @abstractmethod + def check_environment(cls, caps) -> None: + """检查运行环境；不支持时抛出 UnsupportedAttention。""" + pass + + @classmethod + @abstractmethod + def check_support(cls, inputs: VarlenAttentionInputs, caps=None) -> None: + """检查 dtype/head dim/platform 等能力限制；不读取或校验序列边界值。""" + pass + + @abstractmethod + def forward(self, inputs: VarlenAttentionInputs, *, out: torch.Tensor) -> torch.Tensor: + """每次执行校验当前输入布局，然后写入 out；Graph replay 也须重新执行设备值检查。""" + pass + + +def validate_varlen_output(inputs: VarlenAttentionInputs, out: torch.Tensor) -> None: + """Validate metadata only; out may have padding between heads/tokens, but cannot alias inputs.""" + query = inputs.query + if out.shape != query.shape or out.dtype != query.dtype or out.device != query.device: + raise ValueError("varlen out must match query shape, dtype and device") + if out.layout != torch.strided or out.stride(-1) != 1: + raise ValueError("varlen out must have a contiguous head dimension") + extent = 1 + for stride, size in sorted(zip(out.stride(), out.shape)): + if size > 1: + if stride < extent: + raise ValueError("varlen out must have non-overlapping strides") + extent = stride * (size - 1) + extent + if out.numel(): + storage = out.untyped_storage().data_ptr() + if any(storage == tensor.untyped_storage().data_ptr() for tensor in (inputs.query, inputs.key, inputs.value)): + raise ValueError("varlen out must not share storage with query, key or value") diff --git a/lightllm/common/attention/api/workspace.py b/lightllm/common/attention/api/workspace.py new file mode 100644 index 0000000000..c12426511a --- /dev/null +++ b/lightllm/common/attention/api/workspace.py @@ -0,0 +1,115 @@ +from contextlib import contextmanager +from typing import Callable, TypeVar, cast + +import torch + +from lightllm.platform import get_device_runtime + +T = TypeVar("T") + + +class AttentionWorkspace: + def __init__(self, device=None, *, runtime=None): + self.runtime = runtime if runtime is not None else get_device_runtime(device) + self.device = self.runtime.device + self._buffers = {} + self._objects = {} + self._completion = None + self._last_stream = None + self._frozen = False + self._pool = None + + def borrow_for_capture(self): + """捕获实例借用当前池的地址；池后续扩容时，旧图仍持有原 tensor。 + + 同一槽位的借用者共享完成事件，因此共享 scratch 不允许并发执行；不同 + microbatch 槽位由模型创建不同的池。只有实际扩容才需要额外存储。 + """ + workspace = AttentionWorkspace(runtime=self.runtime) + workspace._pool = self._pool or self + workspace._objects = self._objects + return workspace + + def get_object(self, name: str, factory: Callable[[], T], *, expected_type: type[T] | None = None) -> T: + """Share objects across mutually exclusive contexts, checking their declared type. + + Forks share this live object pool, not a snapshot. Freezing a fork prevents that + fork from adding objects or growing buffers; other buckets may still warm up. + Existing objects must keep captured addresses stable even when the pool grows. + """ + if expected_type is None: + if not isinstance(factory, type): + raise TypeError("workspace object factories require an expected_type") + expected_type = factory + if name not in self._objects: + if self._frozen: + raise RuntimeError(f"captured workspace object {name} must be created before capture") + value = factory() + if not isinstance(value, expected_type): + raise TypeError( + f"workspace object {name} requires {expected_type.__name__}, got {type(value).__name__}" + ) + self._objects[name] = value + value = self._objects[name] + if not isinstance(value, expected_type): + raise TypeError(f"workspace object {name} requires {expected_type.__name__}, got {type(value).__name__}") + return cast(T, value) + + @contextmanager + def lease(self, dependencies=()): + self.begin() + try: + yield self + finally: + self.finish(dependencies) + + def begin(self): + if self._pool is not None: + return self._pool.begin() + if self._completion is not None: + stream = self.runtime.current_stream() + if stream != self._last_stream: + stream.wait_event(self._completion) + + def finish(self, dependencies=()): + if self._pool is not None: + return self._pool.finish(dependencies) + stream = self.runtime.current_stream() + if stream is not None: + for event in dependencies: + stream.wait_event(event) + if self._completion is None: + self._completion = self.runtime.create_event() + self._last_stream = stream + self._completion.record(self._last_stream) + + def freeze(self): + """捕获后禁止替换任何地址；互斥 bucket 可继续借用已预留的大缓冲区。""" + self._frozen = True + + def reserve(self, resources): + for requirement in resources.workspace: + self.scratch(requirement.name, requirement.elements, requirement.dtype) + + def synchronize(self): + if self._pool is not None: + return self._pool.synchronize() + if self._completion is not None: + self._completion.synchronize() + + def scratch(self, name: str, size: int, dtype=torch.uint8): + if not name or size <= 0: + raise ValueError("workspace name and positive size are required") + key = name, dtype + tensor = self._buffers.get(key) + if tensor is None or tensor.numel() < size: + if self._frozen: + raise RuntimeError(f"captured workspace {name} cannot grow; reserve capacity before capture") + tensor = ( + self._pool.scratch(name, size, dtype) + if self._pool is not None + else self.runtime.allocate(size, dtype=dtype) + ) + self._buffers[key] = tensor + + return tensor diff --git a/lightllm/common/basemodel/attention/fa3/__init__.py b/lightllm/common/attention/backends/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention/fa3/__init__.py rename to lightllm/common/attention/backends/__init__.py diff --git a/lightllm/common/basemodel/attention/flashinfer/env_utils.py b/lightllm/common/attention/backends/_flashinfer.py similarity index 99% rename from lightllm/common/basemodel/attention/flashinfer/env_utils.py rename to lightllm/common/attention/backends/_flashinfer.py index 1394fa2c53..cae6a29ca7 100644 --- a/lightllm/common/basemodel/attention/flashinfer/env_utils.py +++ b/lightllm/common/attention/backends/_flashinfer.py @@ -1,4 +1,5 @@ import os + import torch diff --git a/lightllm/common/basemodel/attention/flashinfer/__init__.py b/lightllm/common/attention/backends/_shared/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention/flashinfer/__init__.py rename to lightllm/common/attention/backends/_shared/__init__.py diff --git a/lightllm/common/basemodel/attention/linear/__init__.py b/lightllm/common/attention/backends/_shared/linear/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention/linear/__init__.py rename to lightllm/common/attention/backends/_shared/linear/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_mtp.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d_mtp.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_mtp.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d_mtp.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py b/lightllm/common/attention/backends/_shared/linear/causal_conv1d_qkv.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_qkv.py rename to lightllm/common/attention/backends/_shared/linear/causal_conv1d_qkv.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py b/lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py rename to lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py index 849bccb9e9..e972ec2844 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_delta_h.py +++ b/lightllm/common/attention/backends/_shared/linear/chunk_delta_h.py @@ -9,12 +9,11 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices, prepare_chunk_offsets -from .op import exp, safe_exp +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices, prepare_chunk_offsets +from lightllm.common.attention.backends._shared.linear.op import exp, safe_exp from lightllm.common.triton_utils.autotuner import autotune NUM_WARPS = [2, 4, 8, 16] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py b/lightllm/common/attention/backends/_shared/linear/cumsum.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py rename to lightllm/common/attention/backends/_shared/linear/cumsum.py index 6331e1602d..d69e44ae4c 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/cumsum.py +++ b/lightllm/common/attention/backends/_shared/linear/cumsum.py @@ -8,12 +8,11 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .utils import check_shared_mem, input_guard +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.utils import check_shared_mem, input_guard from lightllm.common.triton_utils.autotuner import autotune BS_LIST = [32, 64] if check_shared_mem() else [16, 32] diff --git a/lightllm/common/attention/backends/_shared/linear/gdn.py b/lightllm/common/attention/backends/_shared/linear/gdn.py new file mode 100644 index 0000000000..cb5643c4ec --- /dev/null +++ b/lightllm/common/attention/backends/_shared/linear/gdn.py @@ -0,0 +1,52 @@ +"""GDN tensor-parallel dimensions and input views shared by native and replay kernels.""" + +import torch + + +class GDNMetadata: + def _init_linear_layer_metadata(self, network_config, tp_world_size): + self.num_v_heads = network_config["linear_num_value_heads"] + self.num_k_heads = network_config["linear_num_key_heads"] + self.head_k_dim = network_config["linear_key_head_dim"] + self.head_v_dim = network_config["linear_value_head_dim"] + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] + self.activation = network_config["hidden_act"] + self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size + self.tp_ba_dim = self.num_v_heads * 2 // tp_world_size + self.tp_num_k_heads = self.num_k_heads // tp_world_size + self.tp_num_v_heads = self.num_v_heads // tp_world_size + self.tp_key_dim = self.key_dim // tp_world_size + self.tp_value_dim = self.value_dim // tp_world_size + if self.num_v_heads % self.num_k_heads: + raise ValueError("num_v_heads must be divisible by num_k_heads") + self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads + self.ssm_state_dtype = self.spec.cache.state_dtype + return + + def _split_qkvzba(self, mixed_qkvzba): + qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim + z_end = qkv_dim + self.tp_value_dim + b_end = z_end + self.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return (mixed_qkv, z, b, a) + + def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): + if decode: + query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1) + batch_size = mixed_qkv.shape[0] + query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) + value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim) + return (query, key, value) + else: + query, key, value = torch.split(mixed_qkv, [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], dim=-1) + seq_len = query.shape[0] + query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) + value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) + return (query, key, value) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py b/lightllm/common/attention/backends/_shared/linear/index.py similarity index 93% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py rename to lightllm/common/attention/backends/_shared/linear/index.py index 878e2296f4..4b9ec5c001 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py +++ b/lightllm/common/attention/backends/_shared/linear/index.py @@ -8,10 +8,9 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang # ruff: noqa: E501 import torch - import triton -from .utils import tensor_cache +from lightllm.common.attention.backends._shared.linear.utils import tensor_cache @tensor_cache diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py b/lightllm/common/attention/backends/_shared/linear/l2norm.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py rename to lightllm/common/attention/backends/_shared/linear/l2norm.py index 29f892ef26..fcad017d9b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/l2norm.py +++ b/lightllm/common/attention/backends/_shared/linear/l2norm.py @@ -10,9 +10,9 @@ import os import torch - import triton import triton.language as tl + from lightllm.common.triton_utils.autotuner import autotune BT_LIST = [8, 16, 32, 64, 128] diff --git a/lightllm/common/attention/backends/_shared/linear/layout.py b/lightllm/common/attention/backends/_shared/linear/layout.py new file mode 100644 index 0000000000..d578429779 --- /dev/null +++ b/lightllm/common/attention/backends/_shared/linear/layout.py @@ -0,0 +1,67 @@ +from dataclasses import dataclass, field + +import torch + +from lightllm.common.attention.api.backend import AttentionContext + + +@dataclass(slots=True) +class LinearDecodeContext(AttentionContext): + b_conv_buffer_idx: torch.Tensor | None = None + b_ssm_buffer_idx: torch.Tensor | None = None + b1_mtp_cu_q_seq_len: torch.Tensor | None = None + b_cu_seqlens: torch.Tensor | None = None + accepted_by_state: dict[torch.Tensor, torch.Tensor] = field(default_factory=dict) + + def accepted_tokens(self, state_index: torch.Tensor) -> torch.Tensor: + """Gather once per State per update; Graph captures the gather for every replay. + + The cache retains the source tensors until the next update, keeping identity + keys valid for this context. HOLD sequences must select the initial convolution state. + """ + key = state_index + if key not in self.accepted_by_state: + self.accepted_by_state[key] = torch.where( + self.b_conv_buffer_idx == self.batch.hold_req_id, + 1, + state_index[self.b_conv_buffer_idx] + 1, + ) + return self.accepted_by_state[key] + + +def update_linear_decode(context: LinearDecodeContext, draft_step: int, dynamic: bool, *, replay: bool) -> None: + from .mtp_state_params import build_dynamic_mtp_layout + + batch = context.batch + context.accepted_by_state.clear() + if draft_step == 0: + context.b_conv_buffer_idx = batch.b_req_idx + context.b_ssm_buffer_idx = batch.b_req_idx + context.b1_mtp_cu_q_seq_len = None + if replay: + context.b_cu_seqlens = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) + return + width = draft_step + 1 + if dynamic: + context.b1_mtp_cu_q_seq_len, context.b_conv_buffer_idx = build_dynamic_mtp_layout( + b_req_idx=batch.b_req_idx, + b_mtp_index=batch.b_mtp_index, + hold_req_id=batch.hold_req_id, + ) + else: + if batch.batch_size % width: + raise ValueError("linear decode batch size must be divisible by draft_step + 1") + context.b1_mtp_cu_q_seq_len = torch.arange( + 0, + batch.batch_size + 1, + width, + dtype=torch.int32, + device=batch.device, + ) + context.b_conv_buffer_idx = batch.b_req_idx.view(-1, width)[:, 0].contiguous() + if replay: + context.b_ssm_buffer_idx = context.b_conv_buffer_idx + context.b_cu_seqlens = context.b1_mtp_cu_q_seq_len + else: + offsets = torch.arange(width, dtype=batch.b_req_idx.dtype, device=batch.device) + context.b_ssm_buffer_idx = context.b_conv_buffer_idx[:, None] * width + offsets diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py similarity index 66% rename from lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py rename to lightllm/common/attention/backends/_shared/linear/mtp_state_params.py index b81bf9ec16..993ef6f1b2 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py +++ b/lightllm/common/attention/backends/_shared/linear/mtp_state_params.py @@ -14,6 +14,7 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( batch_size, hold_req_id, BLOCK_SIZE: tl.constexpr, + SELECT_ACCEPTED: tl.constexpr = True, ): offsets = tl.arange(0, BLOCK_SIZE) token_mask = offsets < batch_size @@ -30,21 +31,18 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( # replay arbitrary compact per-request widths at the same token batch size. tl.store(out_cu_q_seq_len + offsets, actual_token_num, mask=cu_mask) tl.store(out_conv_buffer_idx + offsets, hold_req_id, mask=token_mask) - tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask) + if SELECT_ACCEPTED: + tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask) tl.debug_barrier() # mtp_index restarts at zero on the first row of every request group. is_start = valid_row & (mtp_index == 0) sequence_index = tl.cumsum(tl.where(is_start, 1, 0), axis=0) - 1 - accepted_state_index = tl.load( - req_to_mtp_state_index + req_idx, - mask=is_start, - other=0, - ) - tl.store(out_cu_q_seq_len + sequence_index, offsets, mask=is_start) tl.store(out_conv_buffer_idx + sequence_index, req_idx, mask=is_start) - tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) + if SELECT_ACCEPTED: + accepted_state_index = tl.load(req_to_mtp_state_index + req_idx, mask=is_start, other=0) + tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) def build_dynamic_mtp_linear_att_state_params( @@ -85,11 +83,10 @@ def build_dynamic_mtp_linear_att_state_params( per-request query lengths. """ - assert b_req_idx.is_cuda and b_mtp_index.is_cuda and req_to_mtp_state_index.is_cuda - assert b_req_idx.ndim == 1 and b_req_idx.shape == b_mtp_index.shape - assert b_req_idx.dtype == torch.int32 and b_mtp_index.dtype == torch.int32 + _validate_layout_inputs(b_req_idx, b_mtp_index) + if req_to_mtp_state_index.device != b_req_idx.device or req_to_mtp_state_index.ndim != 1: + raise ValueError("MTP state indices must be a vector on the request device") batch_size = b_req_idx.shape[0] - assert batch_size > 0 b1_cu_q_seq_len = torch.empty((batch_size + 1,), dtype=torch.int32, device=b_req_idx.device) b_conv_buffer_idx = torch.empty_like(b_req_idx) @@ -109,3 +106,34 @@ def build_dynamic_mtp_linear_att_state_params( num_stages=1, ) return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens + + +def build_dynamic_mtp_layout(b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, hold_req_id: int): + _validate_layout_inputs(b_req_idx, b_mtp_index) + batch_size = b_req_idx.numel() + cu = torch.empty(batch_size + 1, dtype=torch.int32, device=b_req_idx.device) + requests = torch.empty_like(b_req_idx) + _build_dynamic_mtp_linear_att_state_params_kernel[(1,)]( + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + req_to_mtp_state_index=None, + out_cu_q_seq_len=cu, + out_conv_buffer_idx=requests, + out_num_accepted_tokens=None, + batch_size=batch_size, + hold_req_id=hold_req_id, + BLOCK_SIZE=triton.next_power_of_2(batch_size + 1), + SELECT_ACCEPTED=False, + num_warps=8, + num_stages=1, + ) + return cu, requests + + +def _validate_layout_inputs(requests, positions): + if not requests.is_cuda or positions.device != requests.device: + raise ValueError("MTP layout requires inputs on the same CUDA device") + if requests.ndim != 1 or requests.shape != positions.shape or requests.numel() == 0: + raise ValueError("MTP layout requires matching nonempty vectors") + if requests.dtype != torch.int32 or positions.dtype != torch.int32: + raise ValueError("MTP layout requires int32 request and position indices") diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py b/lightllm/common/attention/backends/_shared/linear/op.py similarity index 95% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py rename to lightllm/common/attention/backends/_shared/linear/op.py index 2f69aa981d..8cbde95b6b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/op.py +++ b/lightllm/common/attention/backends/_shared/linear/op.py @@ -10,7 +10,7 @@ import triton import triton.language as tl -from .utils import is_gather_supported +from lightllm.common.attention.backends._shared.linear.utils import is_gather_supported exp = tl.exp log = tl.log diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/attention/backends/_shared/linear/replayssm_cursor.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_cursor.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/attention/backends/_shared/linear/replayssm_decode.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_decode.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/attention/backends/_shared/linear/replayssm_delta.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_delta.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/attention/backends/_shared/linear/replayssm_fold.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py rename to lightllm/common/attention/backends/_shared/linear/replayssm_fold.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py b/lightllm/common/attention/backends/_shared/linear/solve_tril.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py rename to lightllm/common/attention/backends/_shared/linear/solve_tril.py index 3512788934..1b0a73efe0 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py +++ b/lightllm/common/attention/backends/_shared/linear/solve_tril.py @@ -12,13 +12,12 @@ from typing import Optional import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import make_tensor_descriptor -from .utils import input_guard, is_amd, is_tma_supported +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import make_tensor_descriptor +from lightllm.common.attention.backends._shared.linear.utils import input_guard, is_amd, is_tma_supported def _ensure_triton_allocator(): diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/attention/backends/_shared/linear/ssm_acceptance.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py rename to lightllm/common/attention/backends/_shared/linear/ssm_acceptance.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py b/lightllm/common/attention/backends/_shared/linear/utils.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py rename to lightllm/common/attention/backends/_shared/linear/utils.py index cd7c2e3aeb..e3f6c65510 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/utils.py +++ b/lightllm/common/attention/backends/_shared/linear/utils.py @@ -16,7 +16,6 @@ from typing import Any, Literal import torch - import triton logger = logging.getLogger(__name__) diff --git a/lightllm/common/basemodel/attention/nsa/__init__.py b/lightllm/common/attention/backends/_shared/prefill/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention/nsa/__init__.py rename to lightllm/common/attention/backends/_shared/prefill/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py b/lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py rename to lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py index ef38e39d6c..8cc7d7e0a0 100644 --- a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_flashattention_nopad.py +++ b/lightllm/common/attention/backends/_shared/prefill/context_flashattention_nopad.py @@ -1,11 +1,11 @@ # From : https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F from lightllm.utils.device_utils import is_tesla @@ -329,7 +329,6 @@ def _fwd_kernel_no_prompt_cache( @torch.no_grad() def context_attention_fwd_no_prompt_cache(q, k, v, o, b_start_loc, b_seq_len, max_input_len): - BLOCK_M = 128 if not is_tesla() else 64 # shape constraints Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -566,7 +565,6 @@ def context_attention_fwd_contiguous_kv( def torch_context_attention_fwd(q, k, v, o, b_req_idx, b_start_loc, b_seq_len, b_prompt_cache_len, req_to_token_indexs): - batch = b_start_loc.shape[0] print(q.shape) for i in range(batch): @@ -606,8 +604,8 @@ def torch_context_attention_fwd(q, k, v, o, b_req_idx, b_start_loc, b_seq_len, b def test(): - import torch import numpy as np + import torch Z, H, N_CTX, D_HEAD = 16, 16, 2048, 128 dtype = torch.float16 diff --git a/lightllm/common/basemodel/attention/triton/__init__.py b/lightllm/common/attention/backends/dense/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention/triton/__init__.py rename to lightllm/common/attention/backends/dense/__init__.py diff --git a/lightllm/common/basemodel/attention_vit/__init__.py b/lightllm/common/attention/backends/dense/fa3/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention_vit/__init__.py rename to lightllm/common/attention/backends/dense/fa3/__init__.py diff --git a/lightllm/common/attention/backends/dense/fa3/backend.py b/lightllm/common/attention/backends/dense/fa3/backend.py new file mode 100644 index 0000000000..7f1b1d6b74 --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/backend.py @@ -0,0 +1,250 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend +from lightllm.utils.envs_utils import enable_env_vars + + +@dataclasses.dataclass(slots=True) +class Fa3PrefillContext(AttentionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + causal: bool = None + + +@dataclasses.dataclass(slots=True) +class Fa3DecodeContext(AttentionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + b_att_seq_len: torch.Tensor = None + decode_max_q_seq_len: int = None + decode_max_kv_seq_len: int = None + causal: bool = None + + +class Fa3DenseCommon: + family = AttentionFamily.DENSE + supported_features = frozenset( + { + AttentionFeature.MTP, + AttentionFeature.DYNAMIC_VERIFY, + AttentionFeature.SINK, + AttentionFeature.SLIDING_WINDOW, + AttentionFeature.NONCAUSAL, + AttentionFeature.MTP_WINDOW, + } + ) + + requires_cpu_cache_sync = enable_env_vars("LIGHTLLM_FA3_CPU_CACHE_SYNC") + + provider = "fa3" + + dependency = "sgl_kernel.flash_attn" + + min_device_capability = (9, 0) + + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + + max_head_size = 256 + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + +class Fa3DensePrefillBackend(Fa3DenseCommon, AttentionBackend[Fa3PrefillContext, DenseInputs, torch.Tensor]): + phase = AttentionPhase.PREFILL + + def _prefill_paged(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + if self.spec.sliding_window is not None: + window_size = self.spec.sliding_window + else: + window_size = (-1, -1) + if layer.resources.sink is not None: + sink_weight: torch.Tensor = layer.resources.sink + else: + sink_weight = None + k_descale, v_descale = (None, None) + Lq = q.shape[-1] + sm_scale = 1.0 / Lq ** 0.5 + o = flash_attn_with_kvcache( + q=q, + k_cache=k.view(-1, self.kernel_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.kernel_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=sm_scale, + causal=context.causal, + window_size=window_size, + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + sinks=sink_weight, + ) + return o + + def allocate_context(self, capacity, *, capture, workspace) -> Fa3PrefillContext: + return Fa3PrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Fa3PrefillContext): + from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy + import triton + + context.causal = self.spec.causal + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + context.page_table = torch.empty( + (context.batch.batch_size, triton.cdiv(context.batch.max_kv_seq_len, self.kernel_page_size)), + dtype=torch.int32, + device=context.batch.device, + ) + page_table_copy( + page_table=context.page_table, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_req_idx=context.batch.b_req_idx, + page_size=self.kernel_page_size, + ) + + def forward( + self, layer: LayerAttention, context: Fa3PrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._prefill_paged(layer, context, q=q, k=k, v=v) + + +class Fa3DenseDecodeBackend(Fa3DenseCommon, AttentionBackend[Fa3DecodeContext, DenseInputs, torch.Tensor]): + phase = AttentionPhase.DECODE + + def _decode_init_dynamic_verify_metadata(self, context, draft_step: int) -> torch.Tensor: + from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params + from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers + + b_mark_mtp_shared_group = build_mtp_shared_group_markers(context.batch.b_req_idx, hold_req_id=self.hold_req_id) + b_q_seq_len, b_kv_seq_len, b_att_req_idx, context.b_att_seq_len = build_dynamic_spec_fa3_decode_params( + b_req_idx=context.batch.b_req_idx, + b_seq_len=context.batch.b_seq_len, + b_mark_mtp_shared_group=b_mark_mtp_shared_group, + att_batch_size=context.batch.batch_size, + hold_req_id=self.hold_req_id, + ) + self._decode_init_mtp_cu_seqlens(context, b_q_seq_len, b_kv_seq_len) + context.decode_max_q_seq_len = draft_step + 1 + return b_att_req_idx + + def _decode_init_fixed_mtp_metadata(self, context, draft_step: int) -> torch.Tensor: + mtp_size = draft_step + 1 + assert context.batch.batch_size % mtp_size == 0, ( + "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, " + f"got batch_size={context.batch.batch_size}, draft_step={draft_step}." + ) + b_q_seq_len = torch.full( + (context.batch.b_seq_len.shape[0] // mtp_size,), + fill_value=mtp_size, + dtype=torch.int32, + device=context.batch.b_seq_len.device, + ) + b_kv_seq_len = context.batch.b_seq_len[draft_step::mtp_size] + b_att_req_idx = context.batch.b_req_idx[draft_step::mtp_size] + context.b_att_seq_len = b_kv_seq_len.contiguous() + self._decode_init_mtp_cu_seqlens(context, b_q_seq_len, b_kv_seq_len) + context.decode_max_q_seq_len = mtp_size + return b_att_req_idx + + def _decode_init_normal_metadata(self, context) -> torch.Tensor: + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + context.b_att_seq_len = context.batch.b_seq_len + context.decode_max_q_seq_len = 1 + return context.batch.b_req_idx + + def _decode_init_mtp_cu_seqlens(self, context, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor): + from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor + + b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len) + context.cu_seqlens_q = b1_cu_q_seq_len.int() + context.cu_seqlens_k = b1_cu_kv_seq_len.int() + + def _decode_init_page_table(self, context, b_att_req_idx): + from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy + import triton + + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + width = triton.cdiv(context.capacity.max_kv_seq_len, self.kernel_page_size) + context.page_table = context.workspace.scratch("fa3_page_table", b_att_req_idx.numel() * width, torch.int32)[ + : b_att_req_idx.numel() * width + ].view(b_att_req_idx.numel(), width) + page_table_copy( + page_table=context.page_table, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_req_idx=b_att_req_idx, + page_size=self.kernel_page_size, + ) + + def _run_decode(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + from lightllm.utils.sgl_utils import flash_attn_with_kvcache_autotune + + if self.spec.sliding_window is not None: + window_size = self.spec.sliding_window + else: + window_size = (-1, -1) + if layer.resources.sink is not None: + sink_weight: torch.Tensor = layer.resources.sink + else: + sink_weight = None + k_descale, v_descale = (None, None) + Lq = q.shape[-1] + sm_scale = 1.0 / Lq ** 0.5 + o = flash_attn_with_kvcache_autotune( + q=q, + k_cache=k.view(-1, self.kernel_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.kernel_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=None, + max_seqlen_q=context.decode_max_q_seq_len, + max_seqlen_k=context.decode_max_kv_seq_len, + softmax_scale=sm_scale, + causal=context.causal, + window_size=window_size, + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + sinks=sink_weight, + ) + return o + + def allocate_context(self, capacity, *, capture, workspace) -> Fa3DecodeContext: + return Fa3DecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Fa3DecodeContext): + context.causal = self.spec.causal + draft_step = self.draft_step + if self.spec.dynamic_verify: + b_att_req_idx = self._decode_init_dynamic_verify_metadata(context, draft_step) + elif draft_step > 0: + b_att_req_idx = self._decode_init_fixed_mtp_metadata(context, draft_step) + else: + b_att_req_idx = self._decode_init_normal_metadata(context) + self._decode_init_page_table(context, b_att_req_idx) + + def forward( + self, layer: LayerAttention, context: Fa3DecodeContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._run_decode(layer, context, q=q, k=k, v=v) diff --git a/lightllm/common/attention/backends/dense/fa3/fp8.py b/lightllm/common/attention/backends/dense/fa3/fp8.py new file mode 100644 index 0000000000..90248f151d --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/fp8.py @@ -0,0 +1,166 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.fa3.backend import ( + Fa3DensePrefillBackend, + Fa3DenseDecodeBackend, + Fa3DecodeContext, + Fa3PrefillContext, +) +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +@dataclasses.dataclass(slots=True) +class Fp8Fa3PrefillContext(Fa3PrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Fp8Fa3DecodeContext(Fa3DecodeContext): + pass + + +class Fa3DenseFp8Common: + family = AttentionFamily.DENSE + supported_features = frozenset( + { + AttentionFeature.MTP, + AttentionFeature.DYNAMIC_VERIFY, + AttentionFeature.SINK, + AttentionFeature.NONCAUSAL, + AttentionFeature.MTP_WINDOW, + } + ) + + provider = "fa3" + + cache_encodings = frozenset({CacheEncoding.FP8_HEAD}) + + @property + def kernel_page_size(self) -> int: + return 1 + + +class Fa3DenseFp8PrefillBackend(Fa3DenseFp8Common, Fa3DensePrefillBackend): + phase = AttentionPhase.PREFILL + + def _run_prefill(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import ( + q_per_head_fp8_quant, + q_per_head_static_fp8_quant, + ) + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + q_head_num = q.shape[1] + q_head_dim = q.shape[2] + k_head_num = k.shape[1] + k_head_dim = k.shape[2] + cache_k = k.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + cache_v = v.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + batch_size = context.batch.b_seq_len.shape[0] + q = q.reshape(q.shape[0], k_head_num, -1) + if layer.resources.query_scale is None: + q, q_scale = q_per_head_fp8_quant( + q=q, seq_lens=context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1], b1_start_loc=context.cu_seqlens_q + ) + else: + static_q_scales = layer.resources.query_scale + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(batch_size, k_head_num) + o = flash_attn_with_kvcache( + q=q.reshape(-1, q_head_num, q_head_dim), + k_cache=cache_k, + v_cache=cache_v, + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + q_descale=q_scale, + k_descale=layer.resources.key_scale.view(1, -1).expand(batch_size, -1), + v_descale=layer.resources.value_scale.view(1, -1).expand(batch_size, -1), + return_softmax_lse=False, + ) + return o + + def allocate_context(self, capacity, *, capture, workspace) -> Fp8Fa3PrefillContext: + return Fp8Fa3PrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Fp8Fa3PrefillContext): + super().build_device_metadata(context) + + def forward( + self, layer: LayerAttention, context: Fp8Fa3PrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._run_prefill(layer, context, q=q, k=k, v=v) + + +class Fa3DenseFp8DecodeBackend(Fa3DenseFp8Common, Fa3DenseDecodeBackend): + phase = AttentionPhase.DECODE + + def _run_decode(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import ( + q_per_head_fp8_quant, + q_per_head_static_fp8_quant, + ) + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + k_head_num = k.shape[1] + k_head_dim = k.shape[2] + cache_k = k.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + cache_v = v.view(size=(-1, 1, k_head_num, k_head_dim)).view(torch.float8_e4m3fn) + q_head_num = q.shape[1] + att_batch_size = context.b_att_seq_len.shape[0] + q = q.reshape(q.shape[0], k_head_num, -1) + if layer.resources.query_scale is None: + q, q_scale = q_per_head_fp8_quant( + q=q, seq_lens=context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1], b1_start_loc=context.cu_seqlens_q + ) + else: + static_q_scales = layer.resources.query_scale + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(att_batch_size, k_head_num) + o = flash_attn_with_kvcache( + q=q.reshape(-1, q_head_num, k_head_dim), + k_cache=cache_k, + v_cache=cache_v, + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.decode_max_q_seq_len, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + q_descale=q_scale.view(att_batch_size, k_head_num), + k_descale=layer.resources.key_scale.view(1, -1).expand(att_batch_size, -1), + v_descale=layer.resources.value_scale.view(1, -1).expand(att_batch_size, -1), + return_softmax_lse=False, + ) + return o + + def allocate_context(self, capacity, *, capture, workspace) -> Fp8Fa3DecodeContext: + return Fp8Fa3DecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Fp8Fa3DecodeContext): + super().build_device_metadata(context) + + def forward( + self, layer: LayerAttention, context: Fp8Fa3DecodeContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._run_decode(layer, context, q=q, k=k, v=v) diff --git a/lightllm/common/attention/backends/dense/fa3/neo.py b/lightllm/common/attention/backends/dense/fa3/neo.py new file mode 100644 index 0000000000..5b67ddf0c0 --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/neo.py @@ -0,0 +1,95 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase +import dataclasses +import inspect +from functools import lru_cache + +import torch + +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.fa3.backend import ( + Fa3DensePrefillBackend, + Fa3DenseDecodeBackend, + Fa3PrefillContext, +) + +_NEO_FA3_INSTALL_HINT = ( + "Install the Neo FA3 build from source in the current Python environment:\n" + " git clone --recursive https://github.com/WANDY666/flash-attention.git\n" + " cd flash-attention/hopper\n" + " FLASH_ATTENTION_FORCE_BUILD=TRUE python -m pip install --no-build-isolation --no-deps --force-reinstall .\n" + "This replaces a standalone flash_attn_3 / flash_attn_interface installation; LightLLM's regular" + " FA3 uses the separate sgl_kernel package." +) + + +@lru_cache(maxsize=1) +def get_neo_attention(): + try: + from flash_attn_interface import flash_attn_with_kvcache + + if "image_token_end" not in inspect.signature(flash_attn_with_kvcache).parameters: + raise ImportError("flash_attn_interface is missing image_token_end support (need the Neo build)") + except (ImportError, OSError) as exc: + raise UnsupportedAttention(f"Neo FA3 is unavailable: {exc}.\n" + _NEO_FA3_INSTALL_HINT) from exc + return flash_attn_with_kvcache + + +@dataclasses.dataclass(slots=True) +class NeoFa3PrefillContext(Fa3PrefillContext): + pass + + +class Fa3DenseNeoCommon: + family = AttentionFamily.DENSE + variant = "neo" + + provider = "fa3" + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + +class Fa3DenseNeoPrefillBackend(Fa3DenseNeoCommon, Fa3DensePrefillBackend): + phase = AttentionPhase.PREFILL + + @classmethod + def check_environment(cls, caps): + super().check_environment(caps) + get_neo_attention() + + def allocate_context(self, capacity, *, capture, workspace) -> NeoFa3PrefillContext: + return NeoFa3PrefillContext(capacity, capture, workspace) + + def forward( + self, layer: LayerAttention, context: NeoFa3PrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + flash_attn_with_kvcache_neo = get_neo_attention() + if context.batch.b_image_token_end is None: + raise ValueError("Neo prefill requires b_image_token_end to describe image attention spans") + o = flash_attn_with_kvcache_neo( + q=q, + k_cache=k.view(-1, self.kernel_page_size, k.shape[1], k.shape[2]), + v_cache=v.view(-1, self.kernel_page_size, v.shape[1], v.shape[2]), + page_table=context.page_table, + cache_seqlens=context.batch.b_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=1.0 / q.shape[-1] ** 0.5, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + k_descale=None, + v_descale=None, + return_softmax_lse=False, + image_token_end=context.batch.b_image_token_end, + ) + return o + + +class Fa3DenseNeoDecodeBackend(Fa3DenseNeoCommon, Fa3DenseDecodeBackend): + phase = AttentionPhase.DECODE diff --git a/lightllm/common/attention/backends/dense/fa3/varlen.py b/lightllm/common/attention/backends/dense/fa3/varlen.py new file mode 100644 index 0000000000..59862f0381 --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/varlen.py @@ -0,0 +1,72 @@ +import torch + +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class Fa3VarlenAttentionOp(BaseVarlenAttentionOp): + provider = "fa3" + dependency = "sgl_kernel.flash_attn" + min_device_capability = (9, 0) + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + max_head_size = 256 + head_alignment = 8 + supports_graph = True + + def __init__(self): + self._op = None + + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + tokens, _, head_dim = q.shape + if tokens == 0: + return out + if self._op is None: + import sgl_kernel.flash_attn + + self._op = torch.ops.sgl_kernel.fwd.default + softmax_scale = head_dim ** -0.5 + window_size = (-1, -1) + self._op( + q=q, + k=k, + v=v, + k_new=None, + v_new=None, + q_v=None, + out=out, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + cu_seqlens_k_new=None, + seqused_q=None, + seqused_k=None, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + page_table=None, + kv_batch_idx=None, + leftpad_k=None, + rotary_cos=None, + rotary_sin=None, + seqlens_rotary=None, + q_descale=None, + k_descale=None, + v_descale=None, + softmax_scale=softmax_scale, + is_causal=False, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=0, + softcap=0.0, + is_rotary_interleaved=False, + scheduler_metadata=None, + num_splits=1, + pack_gqa=None, + sm_margin=0, + sinks=None, + sparse_mask_fine=None, + only_qv=False, + ) + + return out diff --git a/lightllm/common/basemodel/attention_vit/fa3/__init__.py b/lightllm/common/attention/backends/dense/flashinfer/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention_vit/fa3/__init__.py rename to lightllm/common/attention/backends/dense/flashinfer/__init__.py diff --git a/lightllm/common/attention/backends/dense/flashinfer/backend.py b/lightllm/common/attention/backends/dense/flashinfer/backend.py new file mode 100644 index 0000000000..824ad3aa3c --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/backend.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature + +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention.api import WorkspaceBudget +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.execution import WorkspaceRequirement +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend +from lightllm.utils.dist_utils import get_dp_world_size + +if TYPE_CHECKING: + from flashinfer.decode import BatchDecodeWithPagedKVCacheWrapper + from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper + + +@dataclasses.dataclass(slots=True) +class FlashInferPrefillContext(AttentionContext): + prefill_wrapper: BatchPrefillWithPagedKVCacheWrapper | None = None + + +@dataclasses.dataclass(slots=True) +class FlashInferDecodeContext(AttentionContext): + kv_last_page_len_buffer: torch.Tensor = None + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + decode_wrapper: BatchDecodeWithPagedKVCacheWrapper | None = None + + +class FlashInferDenseCommon: + family = AttentionFamily.DENSE + supports_expert_parallel = False + + provider = "flashinfer" + + dependency = "flashinfer" + + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + workspace_buffer_key = "flashinfer_fp" + + workspace_buffer_size = 512 * 1024 * 1024 + + def __init__(self, env, spec): + import triton + + set_flashinfer_envs() + super().__init__(env=env, spec=spec) + self.tp_q_head_num = spec.num_q_heads + self.tp_kv_head_num = spec.cache.key_heads + self.head_dim = spec.head_size + self.max_seq_length = env.max_seq_len + self.max_page_num = triton.cdiv(self.max_seq_length, self.kernel_page_size) + self.q_data_type = spec.dtype + self.kv_data_type = spec.cache.dtype + + def resource_requirements(self, capacity, *, capture): + return WorkspaceBudget( + (WorkspaceRequirement(self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),) + ) + + +class FlashInferDensePrefillBackend( + FlashInferDenseCommon, AttentionBackend[FlashInferPrefillContext, DenseInputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + supported_features = frozenset( + {AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.NONCAUSAL, AttentionFeature.MTP_WINDOW} + ) + + def _prefill_paged(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + o_tensor = self.alloc_tensor(q.shape, q.dtype, device=q.device) + context.prefill_wrapper.run( + q, + ( + k.view(-1, self.kernel_page_size, k.shape[1], k.shape[2]), + v.view(-1, self.kernel_page_size, v.shape[1], v.shape[2]), + ), + out=o_tensor, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> FlashInferPrefillContext: + return FlashInferPrefillContext(capacity, capture, workspace) + + def forward( + self, layer: LayerAttention, context: FlashInferPrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._prefill_paged(layer, context, q=q, k=k, v=v) + + def plan_host(self, context: FlashInferPrefillContext): + from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor + from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index + import flashinfer + + batch = context.batch + q_starts = batch.b1_cu_q_seq_len.int() + b_page_len = (batch.b_seq_len + self.kernel_page_size - 1) // self.kernel_page_size + kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + last_page_len = batch.b_seq_len - (b_page_len - 1) * self.kernel_page_size + indices = context.workspace.scratch( + "flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32 + ) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=kv_starts[:-1], + max_token_len=batch.max_kv_seq_len, + out_page_indices=indices, + page_size=self.kernel_page_size, + ) + context.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + context.workspace.scratch(self.workspace_buffer_key, self.workspace_buffer_size), + qo_indptr_buf=q_starts, + paged_kv_indptr_buf=kv_starts, + paged_kv_indices_buf=indices, + paged_kv_last_page_len_buf=last_page_len, + ) + context.prefill_wrapper.plan( + q_starts, + kv_starts, + indices, + last_page_len, + self.tp_q_head_num, + self.tp_kv_head_num, + self.head_dim, + self.kernel_page_size, + causal=self.spec.causal, + pos_encoding_mode="NONE", + logits_soft_cap=0.0, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + ) + + +class FlashInferDenseDecodeBackend( + FlashInferDenseCommon, AttentionBackend[FlashInferDecodeContext, DenseInputs, torch.Tensor] +): + phase = AttentionPhase.DECODE + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + def _refresh_decode_plan(self, context, max_kv_len: int): + import triton + from flashinfer.decode import fast_decode_plan + + uniform_kv_indptr_cpu = torch.arange( + context.batch.batch_size + 1, dtype=torch.int32, device="cpu" + ) * triton.cdiv(max_kv_len, self.kernel_page_size) + fast_decode_plan( + context.decode_wrapper, + indptr=context.kv_starts, + indices=context.kv_indices, + last_page_len=context.kv_last_page_len_buffer, + num_qo_heads=self.tp_q_head_num, + num_kv_heads=self.tp_kv_head_num, + head_dim=self.head_dim, + page_size=self.kernel_page_size, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + non_blocking=True, + global_override_indptr_cpu=uniform_kv_indptr_cpu, + ) + return + + def _run_decode(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + o_tensor = self.alloc_tensor(q.shape, q.dtype) + context.decode_wrapper.run( + q, + ( + k.view(-1, self.kernel_page_size, k.shape[1], k.shape[2]), + v.view(-1, self.kernel_page_size, v.shape[1], v.shape[2]), + ), + out=o_tensor, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> FlashInferDecodeContext: + return FlashInferDecodeContext(capacity, capture, workspace) + + def forward( + self, layer: LayerAttention, context: FlashInferDecodeContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._run_decode(layer, context, q=q, k=k, v=v) + + def plan_host(self, context: FlashInferDecodeContext): + from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor + from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index + import flashinfer + + batch = context.batch + b_page_len = (batch.b_seq_len + self.kernel_page_size - 1) // self.kernel_page_size + starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + last_page_len = batch.b_seq_len - (b_page_len - 1) * self.kernel_page_size + if context.kv_starts is None: + context.kv_starts, context.kv_last_page_len_buffer = starts, last_page_len + context.kv_indices = context.workspace.scratch( + "flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32 + ) + else: + context.kv_starts.copy_(starts, non_blocking=True) + context.kv_last_page_len_buffer.copy_(last_page_len, non_blocking=True) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=context.kv_starts[:-1], + max_token_len=context.current_max_kv_seq_len, + out_page_indices=context.kv_indices, + page_size=self.kernel_page_size, + ) + if context.decode_wrapper is None: + context.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( + context.workspace.scratch(self.workspace_buffer_key, self.workspace_buffer_size), + "NHD", + use_cuda_graph=context.capture, + use_tensor_cores=True, + paged_kv_indptr_buffer=context.kv_starts, + paged_kv_indices_buffer=context.kv_indices, + paged_kv_last_page_len_buffer=context.kv_last_page_len_buffer, + ) + context.decode_wrapper.plan( + context.kv_starts, + context.kv_indices, + context.kv_last_page_len_buffer, + self.tp_q_head_num, + self.tp_kv_head_num, + self.head_dim, + self.kernel_page_size, + q_data_type=self.q_data_type, + kv_data_type=self.kv_data_type, + non_blocking=True, + ) + else: + self._refresh_decode_plan(context, context.current_max_kv_seq_len) diff --git a/lightllm/common/attention/backends/dense/flashinfer/fp8.py b/lightllm/common/attention/backends/dense/flashinfer/fp8.py new file mode 100644 index 0000000000..ca3a18de0f --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/fp8.py @@ -0,0 +1,111 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.backends.dense.flashinfer.backend import ( + FlashInferDensePrefillBackend, + FlashInferDenseDecodeBackend, + FlashInferDecodeContext, + FlashInferPrefillContext, +) +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +@dataclasses.dataclass(slots=True) +class Fp8FlashInferPrefillContext(FlashInferPrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Fp8FlashInferDecodeContext(FlashInferDecodeContext): + pass + + +class FlashInferDenseFp8Common: + family = AttentionFamily.DENSE + supports_expert_parallel = False + + provider = "flashinfer" + + cache_encodings = frozenset({CacheEncoding.FP8_TENSOR}) + + def __init__(self, env, spec): + set_flashinfer_envs() + super().__init__(env=env, spec=spec) + self.kv_data_type = torch.float8_e4m3fn + + @property + def kernel_page_size(self) -> int: + return 1 + + +class FlashInferDenseFp8PrefillBackend(FlashInferDenseFp8Common, FlashInferDensePrefillBackend): + phase = AttentionPhase.PREFILL + supported_features = frozenset( + {AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.NONCAUSAL, AttentionFeature.MTP_WINDOW} + ) + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + o_tensor = self.alloc_tensor(q.shape, q.dtype, device=q.device) + k = k.unsqueeze(1).view(torch.float8_e4m3fn) + v = v.unsqueeze(1).view(torch.float8_e4m3fn) + k_descale = layer.resources.key_scale + v_descale = layer.resources.value_scale + context.prefill_wrapper.run(q=q, paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> Fp8FlashInferPrefillContext: + return Fp8FlashInferPrefillContext(capacity, capture, workspace) + + def forward( + self, layer: LayerAttention, context: Fp8FlashInferPrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._run_prefill(layer, context, q=q, k=k, v=v) + + +class FlashInferDenseFp8DecodeBackend(FlashInferDenseFp8Common, FlashInferDenseDecodeBackend): + phase = AttentionPhase.DECODE + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + def _run_decode( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ): + o_tensor = self.alloc_tensor(q.shape, q.dtype, device=q.device) + k = k.unsqueeze(1).view(torch.float8_e4m3fn) + v = v.unsqueeze(1).view(torch.float8_e4m3fn) + k_descale = layer.resources.key_scale + v_descale = layer.resources.value_scale + context.decode_wrapper.run(q=q, paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> Fp8FlashInferDecodeContext: + return Fp8FlashInferDecodeContext(capacity, capture, workspace) + + def forward( + self, layer: LayerAttention, context: Fp8FlashInferDecodeContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + return self._run_decode(layer, context, q=q, k=k, v=v) diff --git a/lightllm/common/basemodel/attention_vit/sdpa/__init__.py b/lightllm/common/attention/backends/dense/sdpa/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention_vit/sdpa/__init__.py rename to lightllm/common/attention/backends/dense/sdpa/__init__.py diff --git a/lightllm/common/basemodel/attention_vit/sdpa/fp.py b/lightllm/common/attention/backends/dense/sdpa/varlen.py similarity index 51% rename from lightllm/common/basemodel/attention_vit/sdpa/fp.py rename to lightllm/common/attention/backends/dense/sdpa/varlen.py index 6c6da2c2b1..0b686ebda7 100644 --- a/lightllm/common/basemodel/attention_vit/sdpa/fp.py +++ b/lightllm/common/attention/backends/dense/sdpa/varlen.py @@ -1,19 +1,26 @@ import torch import torch.nn.functional as F -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend - - -class SdpaVitAttBackend(BaseVitAttBackend): - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: - assert q.ndim == k.ndim == v.ndim == o.ndim == 3 + +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp +from lightllm.platform import Platform + + +class SdpaVarlenAttentionOp(BaseVarlenAttentionOp): + """Fallback with a device-to-host boundary sync and a Python loop per sequence. + + This path cannot be captured in a CUDA Graph. + """ + + provider = "sdpa" + supported_platforms = frozenset({Platform.CPU, Platform.CUDA}) + + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + output_buffer = out + assert q.ndim == k.ndim == v.ndim == output_buffer.ndim == 3 assert cu_seqlens is not None and cu_seqlens.ndim == 1 cu_seqlens = cu_seqlens.detach().to("cpu") B = cu_seqlens.numel() - 1 @@ -33,7 +40,7 @@ def _vit_att_fwd( k_ = k[s:e].permute(1, 0, 2).unsqueeze(0) v_ = v[s:e].permute(1, 0, 2).unsqueeze(0) - out = F.scaled_dot_product_attention( + sequence_output = F.scaled_dot_product_attention( q_, k_, v_, @@ -42,6 +49,6 @@ def _vit_att_fwd( is_causal=False, ) # [1, H, L, D] -> [L, H, D] - o[s:e].copy_(out.squeeze(0).permute(1, 0, 2)) + output_buffer[s:e].copy_(sequence_output.squeeze(0).permute(1, 0, 2)) - return o + return output_buffer diff --git a/lightllm/common/basemodel/attention_vit/triton/__init__.py b/lightllm/common/attention/backends/dense/triton/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention_vit/triton/__init__.py rename to lightllm/common/attention/backends/dense/triton/__init__.py diff --git a/lightllm/common/attention/backends/dense/triton/backend.py b/lightllm/common/attention/backends/dense/triton/backend.py new file mode 100644 index 0000000000..98b3ddaa7d --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/backend.py @@ -0,0 +1,205 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclasses.dataclass(slots=True) +class TritonPrefillContext(AttentionContext): + pass + + +@dataclasses.dataclass(slots=True) +class TritonDecodeContext(AttentionContext): + b_mark_mtp_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None + + +class TritonDenseCommon: + family = AttentionFamily.DENSE + supported_features = frozenset( + {AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.ALIBI, AttentionFeature.SLIDING_WINDOW} + ) + + provider = "triton" + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + +class TritonDensePrefillBackend(TritonDenseCommon, AttentionBackend[TritonPrefillContext, DenseInputs, torch.Tensor]): + phase = AttentionPhase.PREFILL + + def _prefill_alibi(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + out = self.alloc_tensor(q.shape, q.dtype) + from lightllm.common.basemodel.triton_kernel.alibi_att.context_flashattention_nopad import context_attention_fwd + + context_attention_fwd( + q, + k, + v, + out, + context.batch.b_req_idx, + layer.resources.alibi, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + ) + return out + + def _prefill_paged(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd, + ) + + if self.spec.sliding_window is not None: + sliding_window = self.spec.sliding_window + else: + sliding_window = (-1, -1) + out = self.alloc_tensor(q.shape, q.dtype) + context_attention_fwd( + q, + k, + v, + out, + context.batch.b_req_idx, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + sliding_window=sliding_window, + ) + return out + + def allocate_context(self, capacity, *, capture, workspace) -> TritonPrefillContext: + return TritonPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: TritonPrefillContext): + pass + + def forward( + self, layer: LayerAttention, context: TritonPrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + if layer.resources.alibi is not None: + return self._prefill_alibi(layer, context, q=q, k=k, v=v) + else: + return self._prefill_paged(layer, context, q=q, k=k, v=v) + + +class TritonDenseDecodeBackend(TritonDenseCommon, AttentionBackend[TritonDecodeContext, DenseInputs, torch.Tensor]): + phase = AttentionPhase.DECODE + + def _decode_alibi(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + from lightllm.common.basemodel.triton_kernel.alibi_att.token_flashattention_nopad import token_attention_fwd + + out = self.alloc_tensor(q.shape, q.dtype) + token_attention_fwd( + q, + k, + v, + out, + layer.resources.alibi, + context.batch.req_to_token_indexs, + context.batch.b_req_idx, + context.batch.b_kv_start_loc, + context.batch.b_seq_len, + context.batch.max_kv_seq_len, + context.batch.total_token_num, + alloc_tensor_func=self.alloc_tensor, + ) + return out + + def _decode_mha(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding import ( + token_decode_attention_flash_decoding, + ) + + out = self.alloc_tensor(q.shape, q.dtype) + token_decode_attention_flash_decoding( + q=q, batch=context.batch, cache_k=k, cache_v=v, out=out, alloc_tensor_func=self.alloc_tensor + ) + return out + + def _decode_gqa(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding import ( + gqa_token_decode_attention_flash_decoding, + ) + + if self.spec.sliding_window is not None: + sliding_window = self.spec.sliding_window + else: + sliding_window = (-1, -1) + out = self.alloc_tensor(q.shape, q.dtype) + gqa_token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_v=v, + max_len_in_batch=context.decode_max_kv_seq_len, + out=out, + alloc_tensor_func=self.alloc_tensor, + sliding_window=sliding_window, + ) + return out + + def _run_decode_speculative(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( + token_decode_attention_mtp_diverse_single_token, + ) + + out = token_decode_attention_mtp_diverse_single_token( + q=q, + k=k, + v=v, + Req_to_tokens=context.batch.req_to_token_indexs, + B_req_idx=context.batch.b_req_idx, + b_seq_len=context.batch.b_seq_len, + b_mark_shared_group=context.b_mark_mtp_shared_group, + max_kv_len=context.decode_max_kv_seq_len, + alloc_tensor_func=self.alloc_tensor, + ) + return out + + def allocate_context(self, capacity, *, capture, workspace) -> TritonDecodeContext: + return TritonDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: TritonDecodeContext): + from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers + + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + draft_step = self.draft_step + if draft_step > 0: + context.b_mark_mtp_shared_group = build_mtp_shared_group_markers( + context.batch.b_req_idx, hold_req_id=self.hold_req_id + ) + + def forward( + self, layer: LayerAttention, context: TritonDecodeContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + if layer.resources.alibi is not None: + return self._decode_alibi(layer, context, q=q, k=k, v=v) + else: + draft_step = self.draft_step + q_head_num = q.shape[1] + k_head_num = k.shape[1] + if draft_step > 0: + assert q_head_num >= k_head_num, "speculative decode requires q_head_num >= k_head_num" + return self._run_decode_speculative(layer, context, q=q, k=k, v=v) + elif q_head_num == k_head_num and self.spec.sliding_window is None: + return self._decode_mha(layer, context, q=q, k=k, v=v) + elif q_head_num >= k_head_num: + return self._decode_gqa(layer, context, q=q, k=k, v=v) + else: + raise NotImplementedError("error") diff --git a/lightllm/common/attention/backends/dense/triton/int4.py b/lightllm/common/attention/backends/dense/triton/int4.py new file mode 100644 index 0000000000..f947c5bc09 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int4.py @@ -0,0 +1,152 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclasses.dataclass(slots=True) +class Int4kvTritonPrefillContext(AttentionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class Int4kvTritonDecodeContext(AttentionContext): + decode_max_kv_seq_len: int = None + + +class TritonDenseInt4Common: + family = AttentionFamily.DENSE + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "triton" + + cache_encodings = frozenset({CacheEncoding.INT4}) + + +class TritonDenseInt4PrefillBackend( + TritonDenseInt4Common, AttentionBackend[Int4kvTritonPrefillContext, DenseInputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + ) -> torch.Tensor: + total_token_num = context.batch.total_token_num + head_dim = k.shape[2] * 2 + k_dequant = self.alloc_tensor((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device) + v_dequant = self.alloc_tensor((total_token_num, v.shape[1], head_dim), dtype=q.dtype, device=q.device) + o_tensor = self.alloc_tensor(q.shape, dtype=q.dtype, device=q.device) + max_kv_seq_len = context.batch.max_kv_seq_len + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import dequantize_int4kv + + dequantize_int4kv( + k=k, + k_scale=k_scale, + v=v, + v_scale=v_scale, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_seq_len=context.batch.b_seq_len, + b_req_idx=context.batch.b_req_idx, + b_kv_start_loc=context.b_kv_start_loc, + k_out=k_dequant, + v_out=v_dequant, + max_len_in_batch=max_kv_seq_len, + quant_group_size=self.spec.cache.quant_group_size, + ) + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + context_attention_fwd_contiguous_kv( + q=q, + k=k_dequant, + v=v_dequant, + o=o_tensor, + b_start_loc=context.batch.b_q_start_loc, + b_kv_start_loc=context.b_kv_start_loc, + b_seq_len=context.batch.b_seq_len, + max_q_input_len=context.batch.max_q_seq_len, + b_prompt_cache_len=context.batch.b_ready_cache_len, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> Int4kvTritonPrefillContext: + return Int4kvTritonPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Int4kvTritonPrefillContext): + context.b_kv_start_loc = ( + torch.cumsum(context.batch.b_seq_len, dim=0, dtype=context.batch.b_seq_len.dtype) - context.batch.b_seq_len + ) + + def forward( + self, layer: LayerAttention, context: Int4kvTritonPrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + k = (k, layer.resources.key_scale) + v = (v, layer.resources.value_scale) + k, k_scale = k + v, v_scale = v + o = self._run_prefill(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale) + return o + + +class TritonDenseInt4DecodeBackend( + TritonDenseInt4Common, AttentionBackend[Int4kvTritonDecodeContext, DenseInputs, torch.Tensor] +): + phase = AttentionPhase.DECODE + + def _run_decode( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + ) -> torch.Tensor: + from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv.ppl_int4kv_flash_decoding import ( + token_decode_attention_flash_decoding, + ) + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + max_kv_seq_len=context.decode_max_kv_seq_len, + alloc_tensor_func=self.alloc_tensor, + ) + + def allocate_context(self, capacity, *, capture, workspace) -> Int4kvTritonDecodeContext: + return Int4kvTritonDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Int4kvTritonDecodeContext): + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + + def forward( + self, layer: LayerAttention, context: Int4kvTritonDecodeContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + k = (k, layer.resources.key_scale) + v = (v, layer.resources.value_scale) + k, k_scale = k + v, v_scale = v + return self._run_decode(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale) diff --git a/lightllm/common/attention/backends/dense/triton/int8.py b/lightllm/common/attention/backends/dense/triton/int8.py new file mode 100644 index 0000000000..a0cd0b6255 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int8.py @@ -0,0 +1,189 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend +from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel + + +@dataclasses.dataclass(slots=True) +class Int8kvTritonPrefillContext(AttentionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class Int8kvTritonDecodeContext(AttentionContext): + b_shared_seq_len: torch.Tensor = None + b_mark_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None + + +class TritonDenseInt8Common: + family = AttentionFamily.DENSE + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "triton" + + cache_encodings = frozenset({CacheEncoding.INT8}) + + +class TritonDenseInt8PrefillBackend( + TritonDenseInt8Common, AttentionBackend[Int8kvTritonPrefillContext, DenseInputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + ) -> torch.Tensor: + total_token_num = context.batch.total_token_num + k_dequant = self.alloc_tensor((total_token_num, k.shape[1], k.shape[2]), dtype=q.dtype, device=q.device) + v_dequant = self.alloc_tensor((total_token_num, v.shape[1], v.shape[2]), dtype=q.dtype, device=q.device) + o_tensor = self.alloc_tensor(q.shape, dtype=q.dtype, device=q.device) + max_kv_seq_len = context.batch.max_kv_seq_len + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import dequantize_int8kv + + dequantize_int8kv( + k=k, + k_scale=k_scale, + v=v, + v_scale=v_scale, + req_to_token_indexs=context.batch.req_to_token_indexs, + b_seq_len=context.batch.b_seq_len, + b_req_idx=context.batch.b_req_idx, + b_kv_start_loc=context.b_kv_start_loc, + k_out=k_dequant, + v_out=v_dequant, + max_len_in_batch=max_kv_seq_len, + quant_group_size=self.spec.cache.quant_group_size, + ) + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + context_attention_fwd_contiguous_kv( + q=q, + k=k_dequant, + v=v_dequant, + o=o_tensor, + b_start_loc=context.batch.b_q_start_loc, + b_kv_start_loc=context.b_kv_start_loc, + b_seq_len=context.batch.b_seq_len, + max_q_input_len=context.batch.max_q_seq_len, + b_prompt_cache_len=context.batch.b_ready_cache_len, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> Int8kvTritonPrefillContext: + return Int8kvTritonPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Int8kvTritonPrefillContext): + context.b_kv_start_loc = ( + torch.cumsum(context.batch.b_seq_len, dim=0, dtype=context.batch.b_seq_len.dtype) - context.batch.b_seq_len + ) + + def forward( + self, layer: LayerAttention, context: Int8kvTritonPrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + k = (k, layer.resources.key_scale) + v = (v, layer.resources.value_scale) + k, k_scale = k + v, v_scale = v + o = self._run_prefill(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale) + return o + + +class TritonDenseInt8DecodeBackend( + TritonDenseInt8Common, AttentionBackend[Int8kvTritonDecodeContext, DenseInputs, torch.Tensor] +): + phase = AttentionPhase.DECODE + + def _run_decode_diverse( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + ) -> torch.Tensor: + from .kernels.decode.int8kv.int8kv_flash_decoding_diverse import token_decode_attention_flash_decoding + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + alloc_tensor_func=self.alloc_tensor, + b_shared_seq_len=context.b_shared_seq_len, + b_mark_shared_group=context.b_mark_shared_group, + shared_streams_dict=context.workspace.get_object("int8_decode_streams", dict), + ) + + def _run_decode( + self, + layer, + context, + q: torch.Tensor, + k: torch.Tensor, + k_scale: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + ): + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( + token_decode_attention_flash_decoding, + ) + + return token_decode_attention_flash_decoding( + q=q, + batch=context.batch, + cache_k=k, + cache_k_scale=k_scale, + cache_v=v, + cache_v_scale=v_scale, + max_len_in_batch=context.decode_max_kv_seq_len, + alloc_tensor_func=self.alloc_tensor, + ) + + def allocate_context(self, capacity, *, capture, workspace) -> Int8kvTritonDecodeContext: + return Int8kvTritonDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Int8kvTritonDecodeContext): + from lightllm.common.basemodel.triton_kernel.diverse_utils import build_diverse_shared_group_markers + + context.decode_max_kv_seq_len = context.batch.max_kv_seq_len + if enable_diverse_mode_gqa_decode_fast_kernel(): + context.b_mark_shared_group = build_diverse_shared_group_markers( + b_shared_radix_node_id=context.batch.b_shared_radix_node_id + ) + context.b_shared_seq_len = torch.where(context.b_mark_shared_group == 1, 0, context.batch.b_shared_seq_len) + + def forward( + self, layer: LayerAttention, context: Int8kvTritonDecodeContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + k = (k, layer.resources.key_scale) + v = (v, layer.resources.value_scale) + k, k_scale = k + v, v_scale = v + if enable_diverse_mode_gqa_decode_fast_kernel(): + return self._run_decode_diverse(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale) + else: + return self._run_decode(layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale) diff --git a/lightllm/common/basemodel/attention_vit/xformers/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/__init__.py similarity index 100% rename from lightllm/common/basemodel/attention_vit/xformers/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py similarity index 87% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py index d121656c73..5dd6a683fe 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding.py @@ -3,7 +3,7 @@ def gqa_token_decode_attention_flash_decoding( q: torch.Tensor, - infer_state, + batch, cache_k: torch.Tensor, cache_v: torch.Tensor, max_len_in_batch: int, @@ -11,7 +11,7 @@ def gqa_token_decode_attention_flash_decoding( alloc_tensor_func=torch.empty, sliding_window=(-1, -1), ): - batch_size = infer_state.batch_size + batch_size = batch.batch_size q_head_num, head_dim = q.shape[1], q.shape[2] calcu_shape1 = (batch_size, q_head_num, head_dim) @@ -39,9 +39,9 @@ def gqa_token_decode_attention_flash_decoding( q=q.view(calcu_shape1), k=cache_k, v=cache_v, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_Seqlen=infer_state.b_seq_len, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, max_len_in_batch=max_len_in_batch, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, @@ -51,7 +51,7 @@ def gqa_token_decode_attention_flash_decoding( flash_decode_stage2( mid_out=mid_o, mid_out_logexpsum=mid_o_logexpsum, - B_Seqlen=infer_state.b_seq_len, + B_Seqlen=batch.b_seq_len, out=o_tensor.view(calcu_shape1), block_seq=BLOCK_SEQ, sliding_window=sliding_window, diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py index 767cee9870..002caccd4c 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py index 141587ff38..b60e4a69f4 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding_vsm.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/flash/gqa_flash_decoding_vsm.py @@ -1,11 +1,13 @@ +from functools import lru_cache +from typing import Dict + import torch import triton import triton.language as tl +from frozendict import frozendict + from lightllm.common.kernel_config import KernelConfigs from lightllm.utils.device_utils import calcu_kernel_best_vsm_count -from frozendict import frozendict -from functools import lru_cache -from typing import Dict class GQAVSMDecodeAttentionKernelConfig(KernelConfigs): @@ -261,7 +263,7 @@ def gqa_token_decode_attention_flash_decoding_vsm_stage1( q_head_num, kv_head_num, batch_size, - **run_config + **run_config, ): grid = (num_vsm,) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py index 5ce5d193f5..e995378f58 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/gqa_decode_flashattention_nopad.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/gqa_decode_flashattention_nopad.py @@ -1,9 +1,9 @@ -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F @triton.jit @@ -115,7 +115,6 @@ def _fwd_kernel( @torch.no_grad() def gqa_decode_attention_fwd(q, k, v, o, req_to_tokens, b_req_idx, b_seq_len): - # shape constraints Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] assert Lq == Lk and Lk == Lv diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py index 13c7e8070e..8b2dbc885a 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/mtp_diverse_attn.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/mtp_diverse_attn.py @@ -1,7 +1,9 @@ import torch + +from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size + from .stage1_single_token import mtp_diverse_stage1_single_token from .stage2_single_token import mtp_diverse_stage2_single_token -from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size @torch.no_grad() diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py index 143b3c1ce1..ba004beed3 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage1_single_token.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage1_single_token.py @@ -12,11 +12,13 @@ - 每个请求单独做可见性检查（基于各自的 seq_len） - 中间结果按 kv block 存储，供 Stage2 聚合 """ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.device_utils import is_hopper from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py index 5a56426932..4d382439df 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/stage2_single_token.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/gqa/mtp_diverse/stage2_single_token.py @@ -6,11 +6,13 @@ - 根据 seq_len 确定需要聚合的 kv block 数量 - 使用 flash attention reweighting 公式 """ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py index 27cb97d9a0..7db6686b8c 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/int4kv_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/int4kv_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py similarity index 83% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py index 936ca4dbbd..9e76348baf 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/ppl_int4kv_flash_decoding.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int4kv/ppl_int4kv_flash_decoding.py @@ -1,10 +1,11 @@ -import torch from typing import Optional +import torch + def token_decode_attention_flash_decoding( q: torch.Tensor, - infer_state, + batch, cache_k: torch.Tensor, cache_k_scale: torch.Tensor, cache_v: torch.Tensor, @@ -14,7 +15,7 @@ def token_decode_attention_flash_decoding( alloc_tensor_func=torch.empty, ): BLOCK_SEQ = 256 - batch_size = infer_state.batch_size + batch_size = batch.batch_size q_head_num = q.shape[1] head_dim = q.shape[2] calcu_shape1 = (batch_size, q_head_num, head_dim) @@ -42,9 +43,9 @@ def token_decode_attention_flash_decoding( k_scale=cache_k_scale, v=cache_v, v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_Seqlen=infer_state.b_seq_len, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, max_kv_seq_len=max_kv_seq_len, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, @@ -53,5 +54,5 @@ def token_decode_attention_flash_decoding( from ..int8kv.normal.int8kv_flash_decoding_stage2 import flash_decode_stage2 - flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + flash_decode_stage2(mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int4kv/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py similarity index 77% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py index c534795fdf..c5e3af3f36 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse.py @@ -1,25 +1,29 @@ # 为 diverse mode 定制设计的 int8kv flash decoding attention 实现，可以实现更高效的多样性采样 import torch -from lightllm.common.basemodel.infer_struct import InferStateInfo + +from lightllm.common.attention.api.batch import AttentionBatch +from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size + from .int8kv_flash_decoding_diverse_stage1 import flash_decode_stage1 from .int8kv_flash_decoding_diverse_stage2 import flash_decode_stage2 from .int8kv_flash_decoding_diverse_stage3 import flash_diverse_decode_stage3 -from lightllm.utils.envs_utils import get_diverse_max_batch_shared_group_size def token_decode_attention_flash_decoding( q, - infer_state: InferStateInfo, + batch: AttentionBatch, cache_k, cache_k_scale, cache_v, cache_v_scale, out=None, alloc_tensor_func=torch.empty, - shared_streams_dict={}, + b_shared_seq_len=None, + b_mark_shared_group=None, + shared_streams_dict=None, ): - b_shared_seq_len = infer_state.decode_att_state.b_shared_seq_len - b_mark_shared_group = infer_state.decode_att_state.b_mark_shared_group + if b_shared_seq_len is None or b_mark_shared_group is None or shared_streams_dict is None: + raise ValueError("diverse attention requires prepared group metadata and slot-owned streams") if "stream1" not in shared_streams_dict: shared_streams_dict["stream1"] = torch.cuda.Stream() @@ -33,8 +37,8 @@ def token_decode_attention_flash_decoding( head_dim = q.shape[2] BLOCK_SEQ = 256 - batch_size = infer_state.batch_size - max_kv_seq_len = infer_state.max_kv_seq_len + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len calcu_shape1 = (batch_size, q_head_num, head_dim) o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out @@ -56,11 +60,11 @@ def token_decode_attention_flash_decoding( k_scale=cache_k_scale, v=cache_v, v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, b_shared_seq_len=b_shared_seq_len, b_mark_shared_group=b_mark_shared_group, - max_len_in_batch=infer_state.max_kv_seq_len, + max_len_in_batch=batch.max_kv_seq_len, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, block_seq=BLOCK_SEQ, @@ -74,11 +78,11 @@ def token_decode_attention_flash_decoding( k_scale=cache_k_scale, v=cache_v, v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_Seqlen=infer_state.b_seq_len, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_Seqlen=batch.b_seq_len, b_shared_seq_len=b_shared_seq_len, - max_len_in_batch=infer_state.max_kv_seq_len, + max_len_in_batch=batch.max_kv_seq_len, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, block_seq=BLOCK_SEQ, @@ -90,7 +94,7 @@ def token_decode_attention_flash_decoding( flash_diverse_decode_stage3( mid_out=mid_o, mid_out_logexpsum=mid_o_logexpsum, - B_Seqlen=infer_state.b_seq_len, + B_Seqlen=batch.b_seq_len, b_shared_seq_len=b_shared_seq_len, O=o_tensor.view(calcu_shape1), block_seq=BLOCK_SEQ, diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py index fb122a79da..6675075c04 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage1.py @@ -1,11 +1,12 @@ +from functools import lru_cache +from typing import Dict, Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.kernel_config import KernelConfigs from frozendict import frozendict -from functools import lru_cache -from typing import Dict + +from lightllm.common.kernel_config import KernelConfigs class GQADiverseDecodeStage1KernelConfig(KernelConfigs): diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py similarity index 98% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py index 963139b36b..54d568d578 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage2.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage2.py @@ -1,13 +1,13 @@ +from functools import lru_cache +from typing import Dict, Optional + import torch import triton import triton.language as tl -from typing import Optional +from frozendict import frozendict from lightllm.common.kernel_config import KernelConfigs -from frozendict import frozendict -from functools import lru_cache -from typing import Dict -from lightllm.common.triton_utils.autotuner import autotune, Autotuner +from lightllm.common.triton_utils.autotuner import Autotuner, autotune class GQADiverseDecodeStage2KernelConfig(KernelConfigs): diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage3.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/int8kv_flash_decoding_diverse_stage3.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/int8kv_flash_decoding_diverse_stage3.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py similarity index 84% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py index de74f0728d..cb4cad2e8e 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding.py @@ -1,6 +1,9 @@ -import torch from typing import Optional -from lightllm.common.basemodel.infer_struct import InferStateInfo + +import torch + +from lightllm.common.attention.api.batch import AttentionBatch + from .int8kv_flash_decoding_stage1 import flash_decode_stage1 from .int8kv_flash_decoding_stage2 import flash_decode_stage2 @@ -8,7 +11,7 @@ @torch.no_grad() def token_decode_attention_flash_decoding( q: torch.Tensor, - infer_state: InferStateInfo, + batch: AttentionBatch, cache_k: torch.Tensor, cache_k_scale: torch.Tensor, cache_v: torch.Tensor, @@ -17,12 +20,11 @@ def token_decode_attention_flash_decoding( out: Optional[torch.Tensor] = None, alloc_tensor_func=torch.empty, ): - q_head_num = q.shape[1] head_dim = q.shape[2] BLOCK_SEQ = 256 - batch_size = infer_state.batch_size + batch_size = batch.batch_size calcu_shape1 = (batch_size, q_head_num, head_dim) o_tensor = alloc_tensor_func(q.shape, q.dtype, q.device) if out is None else out @@ -46,9 +48,9 @@ def token_decode_attention_flash_decoding( k_scale=cache_k_scale, v=cache_v, v_scale=cache_v_scale, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, - B_req_idx=infer_state.b_req_idx, - B_seq_len=infer_state.b_seq_len, + Req_to_tokens=batch.req_to_token_indexs, + B_req_idx=batch.b_req_idx, + B_seq_len=batch.b_seq_len, max_len_in_batch=max_len_in_batch, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, @@ -58,7 +60,7 @@ def token_decode_attention_flash_decoding( flash_decode_stage2( mid_out=mid_o, mid_out_logexpsum=mid_o_logexpsum, - B_Seqlen=infer_state.b_seq_len, + B_Seqlen=batch.b_seq_len, O=o_tensor.view(calcu_shape1), block_seq=BLOCK_SEQ, ) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py index a05db282d7..2d1975fcf0 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage1.py @@ -1,8 +1,10 @@ +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, Autotuner, AutotuneKernelType, AutotuneLevel + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len, get_triton_autotune_level @@ -266,7 +268,6 @@ def flash_decode_stage1( ): """ """ if not run_config: - run_config = { "BLOCK_N": 16, "num_warps": 2, diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/normal/int8kv_flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/int8kv/normal/int8kv_flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/int8kv/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py similarity index 63% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py index 6c50fc3927..94489f2be9 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding.py @@ -1,10 +1,10 @@ import torch -def token_decode_attention_flash_decoding(q, infer_state, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty): +def token_decode_attention_flash_decoding(q, batch, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty): BLOCK_SEQ = 256 - batch_size = infer_state.batch_size - max_kv_seq_len = infer_state.max_kv_seq_len + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len q_head_num, head_dim = q.shape[1], q.shape[2] calcu_shape1 = (batch_size, q_head_num, head_dim) @@ -24,13 +24,13 @@ def token_decode_attention_flash_decoding(q, infer_state, cache_k, cache_v, out= q.view(calcu_shape1), cache_k, cache_v, - infer_state.req_manager.req_to_token_indexs, - infer_state.b_req_idx, - infer_state.b_seq_len, - infer_state.max_kv_seq_len, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, + batch.max_kv_seq_len, mid_o, mid_o_logexpsum, BLOCK_SEQ, ) - flash_decode_stage2(mid_o, mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) + flash_decode_stage2(mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), BLOCK_SEQ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage1.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage1.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage2.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/flash_decoding_stage2.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/flash/flash_decoding_stage2.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/flash_decoding/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py index 9de2b82057..4065ffe6bd 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_att1.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_att1.py @@ -1,8 +1,8 @@ -import torch +import math +import torch import triton import triton.language as tl -import math @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py index 96a5b26dd6..863ac322de 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_reduceV.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_reduceV.py @@ -1,5 +1,4 @@ import torch - import triton import triton.language as tl diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py index 0bb6410e13..08147b1e99 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_nopad_softmax.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_nopad_softmax.py @@ -68,7 +68,6 @@ def token_softmax_fwd(Logics, B_Start_Loc, B_Seqlen, Prob_Out, max_input_len): def test1(): - import torch B, N_CTX, H, D = 4, 1025, 12, 128 diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py rename to lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py index d963f8582a..6248dac46d 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/token_attention_softmax_and_reducev.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/decode/mha/stage3_decode_att/token_attention_softmax_and_reducev.py @@ -1,7 +1,7 @@ import torch +import torch.nn.functional as F import triton import triton.language as tl -import torch.nn.functional as F @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/__init__.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/decode_att/mha/stage3_decode_att/__init__.py rename to lightllm/common/attention/backends/dense/triton/kernels/prefill/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py rename to lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py index 1016e5a3c1..b6a0d37619 100644 --- a/lightllm/common/basemodel/triton_kernel/att/prefill_att/context_attention_fwd_neo.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/prefill/context_attention_fwd_neo.py @@ -1,4 +1,5 @@ import math + import torch import triton import triton.language as tl diff --git a/lightllm/models/vit/triton_kernel/flashattention_nopad.py b/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py similarity index 65% rename from lightllm/models/vit/triton_kernel/flashattention_nopad.py rename to lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py index af062c2a5b..b498e02120 100644 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py @@ -1,13 +1,9 @@ import torch import triton import triton.language as tl -import math -import time -import torch.nn.functional as F -from typing import Optional, Tuple -from lightllm.utils.device_utils import is_hopper, is_4090 +from packaging.version import Version -if triton.__version__ >= "2.1.0": +if Version(triton.__version__) >= Version("2.1.0"): @triton.jit def _fwd_kernel( @@ -148,82 +144,4 @@ def _flash_attention_triton_fwd( return else: - raise Exception("error triton version!") - -_flash_attn_v3_available = False -try: - from sgl_kernel.flash_attn import flash_attn_varlen_func - - _flash_attn_v3_available = True - - def flash_attention_v3_fwd( - q, - k, - v, - o, - cu_seqlens, - max_seqlen, - ): - head_dim = q.shape[-1] - softmax_scale = head_dim ** -0.5 - window_size = (-1, -1) - torch.ops.sgl_kernel.fwd.default( - q, - k, - v, - None, # k_new - None, # v_new - None, # qv - o, # out - cu_seqlens, - cu_seqlens, - None, # cu_seqlens_k_new - None, - None, - max_seqlen, - max_seqlen, - None, # page_table, - None, # kv_batch_idx - None, # leftpad_k - None, # rotary cos - None, # rotary sin - None, # seqlens_rotary - None, - None, - None, - softmax_scale, - False, - window_size[0], - window_size[1], - 0, # attention_chunk - 0.0, # softcap - is_rotary_interleaved=False, - scheduler_metadata=None, - num_splits=1, - pack_gqa=None, - sm_margin=0, - sinks=None, - sparse_mask_fine=None, - only_qv=False, - ) - return o - -except ImportError: - print("Failed to import _flash_attn_forward from hopper.flash_attn_interface.") - _flash_attn_v3_available = False - - -def flash_attention_fwd(q, k, v, o, cu_seqlens, max_seqlen): - """ - 统一的 Flash Attention 接口。如果 sgl_kernel 存在， - 则使用 sgl_kernel里的接口，否则使用 Triton 版本。 - """ - global _flash_attn_v3_available - if _flash_attn_v3_available and (is_hopper() or is_4090()): - try: - flash_attention_v3_fwd(q, k, v, o, cu_seqlens, max_seqlen) - return - except Exception as e: - print(f"Failed to use sgl_kernel: {e}") - _flash_attn_v3_available = False - _flash_attention_triton_fwd(q, k, v, o, cu_seqlens, max_seqlen) + raise RuntimeError(f"varlen attention requires Triton >= 2.1.0, got {triton.__version__}") diff --git a/lightllm/common/attention/backends/dense/triton/neo.py b/lightllm/common/attention/backends/dense/triton/neo.py new file mode 100644 index 0000000000..a21ad2175f --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/neo.py @@ -0,0 +1,66 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention.backends.dense.triton.backend import ( + TritonDensePrefillBackend, + TritonDenseDecodeBackend, + TritonPrefillContext, +) + + +@dataclasses.dataclass(slots=True) +class NeoTritonPrefillContext(TritonPrefillContext): + pass + + +class TritonDenseNeoCommon: + family = AttentionFamily.DENSE + variant = "neo" + + provider = "triton" + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + +class TritonDenseNeoPrefillBackend(TritonDenseNeoCommon, TritonDensePrefillBackend): + phase = AttentionPhase.PREFILL + + def allocate_context(self, capacity, *, capture, workspace) -> NeoTritonPrefillContext: + return NeoTritonPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: NeoTritonPrefillContext): + pass + + def forward( + self, layer: LayerAttention, context: NeoTritonPrefillContext, inputs: DenseInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.key + v = layer.resources.value + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, + ) + + out = self.alloc_tensor(q.shape, q.dtype) + context_attention_fwd_neo( + q, + k, + v, + out, + context.batch.b_req_idx, + context.batch.b_q_start_loc, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + context.batch.req_to_token_indexs, + context.batch.b_image_token_end, + ) + return out + + +class TritonDenseNeoDecodeBackend(TritonDenseNeoCommon, TritonDenseDecodeBackend): + phase = AttentionPhase.DECODE diff --git a/lightllm/common/attention/backends/dense/triton/varlen.py b/lightllm/common/attention/backends/dense/triton/varlen.py new file mode 100644 index 0000000000..e7f8219dc4 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/varlen.py @@ -0,0 +1,25 @@ +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.triton.kernels.prefill.varlen import _flash_attention_triton_fwd +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp + + +class TritonVarlenAttentionOp(BaseVarlenAttentionOp): + provider = "triton" + dependency = "triton" + supports_graph = True + + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + q, k, v, cu_seqlens, max_seqlen = inputs + if q.shape[0] == 0: + return out + _flash_attention_triton_fwd( + q, + k, + v, + out, + cu_seqlens, # q k v cu_seqlens, + max_seqlen, + ) + return out diff --git a/lightllm/common/attention/backends/dense/varlen.py b/lightllm/common/attention/backends/dense/varlen.py new file mode 100644 index 0000000000..a02a503724 --- /dev/null +++ b/lightllm/common/attention/backends/dense/varlen.py @@ -0,0 +1,42 @@ +import torch + +from lightllm.common.attention._environment import check_attention_environment +from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionOp +from lightllm.platform import Platform + + +class BaseVarlenAttentionOp(VarlenAttentionOp): + provider = "" + + supported_platforms = frozenset({Platform.CUDA}) + dependency = "" + min_device_capability = None + activation_dtypes = frozenset({torch.float16, torch.bfloat16, torch.float32}) + max_head_size = None + head_alignment = 1 + + @classmethod + def check_environment(cls, caps) -> None: + check_attention_environment( + caps, + provider=cls.provider, + supported_platforms=cls.supported_platforms, + min_device_capability=cls.min_device_capability, + dependency=cls.dependency, + ) + + @classmethod + def check_support(cls, inputs, caps=None) -> None: + q = inputs.query + if q.ndim != 3: + raise ValueError("varlen query must have shape [tokens, heads, dim]") + if min(q.shape[1:]) <= 0: + raise UnsupportedAttention("head count and dimension must be positive") + if q.dtype not in cls.activation_dtypes: + raise UnsupportedAttention(f"unsupported activation dtype {q.dtype}") + if caps is not None and caps.platform not in cls.supported_platforms: + raise UnsupportedAttention(f"{cls.provider} does not support platform {caps.platform.value}") + if cls.max_head_size is not None and q.shape[-1] > cls.max_head_size: + raise UnsupportedAttention(f"head dimension exceeds {cls.max_head_size}") + if q.shape[-1] % cls.head_alignment: + raise UnsupportedAttention(f"head dimension must be divisible by {cls.head_alignment}") diff --git a/lightllm/common/basemodel/triton_kernel/att/prefill_att/__init__.py b/lightllm/common/attention/backends/dense/xformers/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/att/prefill_att/__init__.py rename to lightllm/common/attention/backends/dense/xformers/__init__.py diff --git a/lightllm/common/basemodel/attention_vit/xformers/fp.py b/lightllm/common/attention/backends/dense/xformers/varlen.py similarity index 61% rename from lightllm/common/basemodel/attention_vit/xformers/fp.py rename to lightllm/common/attention/backends/dense/xformers/varlen.py index 361b5db050..8edffec9b7 100644 --- a/lightllm/common/basemodel/attention_vit/xformers/fp.py +++ b/lightllm/common/attention/backends/dense/xformers/varlen.py @@ -1,31 +1,25 @@ import torch -import torch.nn.functional as F -try: - from xformers import ops as xformers_ops - from xformers.ops import fmha -except ImportError: - xformers_ops = None - fmha = None +from lightllm.common.attention.api.varlen import validate_varlen_output +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend +class XformersVarlenAttentionOp(BaseVarlenAttentionOp): + provider = "xformers" + dependency = "xformers.ops" -class XformersVitAttBackend(BaseVitAttBackend): @torch.no_grad() - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: + def forward(self, inputs, *, out): + validate_varlen_output(inputs, out) + inputs.validate() + from xformers import ops as xformers_ops + from xformers.ops import fmha + + q, k, v, cu_seqlens, max_seqlen = inputs + o = out assert q.ndim == k.ndim == v.ndim == o.ndim == 3 assert cu_seqlens is not None and cu_seqlens.ndim == 1 assert q.shape == k.shape == v.shape == o.shape - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.int64).tolist() seqlens = [int(L) for L in seqlens if int(L) > 0] diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/__init__.py b/lightllm/common/attention/backends/gdn/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/mla_att/__init__.py rename to lightllm/common/attention/backends/gdn/__init__.py diff --git a/lightllm/common/attention/backends/gdn/flashqla/__init__.py b/lightllm/common/attention/backends/gdn/flashqla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/flashqla/backend.py b/lightllm/common/attention/backends/gdn/flashqla/backend.py new file mode 100644 index 0000000000..863b0d1df0 --- /dev/null +++ b/lightllm/common/attention/backends/gdn/flashqla/backend.py @@ -0,0 +1,21 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionPhase +import torch + +from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNPrefillBackend + + +class FlashQlaGDNPrefillBackend(TritonGDNPrefillBackend): + phase = AttentionPhase.PREFILL + + provider = "flashqla" + + dependency = "flash_qla" + + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + def get_prefill_kernel(self): + from flash_qla import chunk_gated_delta_rule + + return chunk_gated_delta_rule diff --git a/lightllm/common/attention/backends/gdn/replay/__init__.py b/lightllm/common/attention/backends/gdn/replay/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/replay/backend.py b/lightllm/common/attention/backends/gdn/replay/backend.py new file mode 100644 index 0000000000..4a1730ff2f --- /dev/null +++ b/lightllm/common/attention/backends/gdn/replay/backend.py @@ -0,0 +1,223 @@ +from __future__ import annotations + +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import GDNInputs, GDNOutput +from lightllm.common.attention.backends._shared.linear.gdn import GDNMetadata +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclasses.dataclass(slots=True) +class ReplayGDNPrefillContext(AttentionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class ReplayGDNDecodeContext(LinearDecodeContext): + pass + + +class TritonGDNReplayCommon(GDNMetadata): + state_mode = "replay" + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "triton" + + family = AttentionFamily.GDN + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + def __init__(self, env, spec): + super().__init__(env=env, spec=spec) + self._init_linear_layer_metadata(env.network_config, env.tp_size) + self.state_slots = env.state_slots + + +class TritonGDNReplayPrefillBackend( + TritonGDNReplayCommon, AttentionBackend[ReplayGDNPrefillContext, GDNInputs, GDNOutput] +): + phase = AttentionPhase.PREFILL + + def __init__(self, env, spec): + super().__init__(env, spec) + self.prefill_kernel = self.get_prefill_kernel() + + def get_prefill_kernel(self): + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + def _prefill(self, layer, context, mixed_qkv, conv_states, ssm_states, a, b): + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn + from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating + + g, beta = fused_gdn_gating(layer.resources.a_log, a, b, layer.resources.dt_bias) + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer.resources.conv_weight, + bias=layer.resources.conv_bias, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation=self.activation, + ).transpose(0, 1) + query, key, value = self._rearrange_mixed_qkv(mixed_qkv) + core_attn_out, last_state = self.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=ssm_states[context.b_ssm_buffer_idx], + output_final_state=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[context.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + def allocate_context(self, capacity, *, capture, workspace) -> ReplayGDNPrefillContext: + return ReplayGDNPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: ReplayGDNPrefillContext): + reqs = context.batch.b_req_idx + context.b_conv_buffer_idx = reqs + context.b_ssm_buffer_idx = reqs + + def forward( + self, layer: LayerAttention, context: ReplayGDNPrefillContext, inputs: GDNInputs, *, out=None + ) -> GDNOutput: + conv_states, ssm_states = layer.resources.conv, layer.resources.ssm + if (self.state_slots - 1) > 0: + conv_states = conv_states[:, :, : -(self.state_slots - 1)] + mixed_qkv, z, b, a = self._split_qkvzba(inputs.mixed_qkvzba) + core_attn_out = self._prefill(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + return GDNOutput(core_attn_out, z) + + +class TritonGDNReplayDecodeBackend( + TritonGDNReplayCommon, AttentionBackend[ReplayGDNDecodeContext, GDNInputs, GDNOutput] +): + phase = AttentionPhase.DECODE + + def _decode(self, layer, context, mixed_qkv, z, conv_states, a, b, state): + from .decode import replay + from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs + + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer.resources.conv_weight, + layer.resources.conv_bias, + context.b_conv_buffer_idx, + self.activation, + self.conv_kernel_dim, + self.tp_num_k_heads, + self.head_k_dim, + self.tp_num_v_heads, + self.head_v_dim, + ) + query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) + key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) + value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) + output = replay( + state, + state.storage_layer(layer.logical_layer_id), + query, + key, + value, + a, + b, + layer.resources.a_log, + layer.resources.dt_bias, + context.b_conv_buffer_idx, + context.b_cu_seqlens, + ) + return output, z + + def _mtp(self, layer, context, mixed_qkv, conv_states, a, b, draft_step, state): + from .decode import replay + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer.resources.conv_weight, + mtp_step=draft_step, + bias=layer.resources.conv_bias, + activation=self.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.resources.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + query, key, value = torch.split( + mixed_qkv, + [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], + dim=-1, + ) + query = query.view(-1, self.tp_num_k_heads, self.head_k_dim) + key = key.view(-1, self.tp_num_k_heads, self.head_k_dim) + value = value.view(-1, self.tp_num_v_heads, self.head_v_dim) + output = replay( + state, + state.storage_layer(layer.logical_layer_id), + query, + key, + value, + a, + b, + layer.resources.a_log, + layer.resources.dt_bias, + context.b_conv_buffer_idx, + context.b_cu_seqlens, + ) + return output + + def allocate_context(self, capacity, *, capture, workspace) -> ReplayGDNDecodeContext: + return ReplayGDNDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: ReplayGDNDecodeContext): + update_linear_decode(context, self.draft_step, self.spec.dynamic_verify, replay=True) + + def forward( + self, layer: LayerAttention, context: ReplayGDNDecodeContext, inputs: GDNInputs, *, out=None + ) -> GDNOutput: + mixed_qkv, z, b, a = self._split_qkvzba(inputs.mixed_qkvzba) + state = layer.resources.replay_state + conv_states = layer.resources.conv + draft_step = self.draft_step + if draft_step > 0: + core = self._mtp(layer, context, mixed_qkv, conv_states, a, b, draft_step, state) + else: + core, z = self._decode(layer, context, mixed_qkv, z, conv_states, a, b, state) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) + return GDNOutput(core, z) + + +class FlashQlaGDNReplayPrefillBackend(TritonGDNReplayPrefillBackend): + phase = AttentionPhase.PREFILL + + provider = "flashqla" + + dependency = "flash_qla" + + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + + def get_prefill_kernel(self): + from flash_qla import chunk_gated_delta_rule + + return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/attention/backends/gdn/replay/decode.py similarity index 98% rename from lightllm/common/basemodel/attention/linear/replay_gdn_decode.py rename to lightllm/common/attention/backends/gdn/replay/decode.py index ac7af9dbdc..6dda8fc867 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py +++ b/lightllm/common/attention/backends/gdn/replay/decode.py @@ -10,11 +10,8 @@ import torch import triton +from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay, check_replay_sync_mark from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( - _decode_replay, - check_replay_sync_mark, -) from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune diff --git a/lightllm/common/attention/backends/gdn/triton/__init__.py b/lightllm/common/attention/backends/gdn/triton/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/gdn/triton/backend.py b/lightllm/common/attention/backends/gdn/triton/backend.py new file mode 100644 index 0000000000..b86b4dbe38 --- /dev/null +++ b/lightllm/common/attention/backends/gdn/triton/backend.py @@ -0,0 +1,224 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import GDNInputs, GDNOutput +from lightllm.common.attention.backends._shared.linear.gdn import GDNMetadata +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclasses.dataclass(slots=True) +class GDNPrefillContext(AttentionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class GDNDecodeContext(LinearDecodeContext): + pass + + +class TritonGDNCommon(GDNMetadata): + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "triton" + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + family = AttentionFamily.GDN + + def __init__(self, env, spec): + super().__init__(env=env, spec=spec) + self._init_linear_layer_metadata(network_config=env.network_config, tp_world_size=env.tp_size) + self.state_slots = env.state_slots + + +class TritonGDNPrefillBackend(TritonGDNCommon, AttentionBackend[GDNPrefillContext, GDNInputs, GDNOutput]): + phase = AttentionPhase.PREFILL + + def __init__(self, env, spec): + super().__init__(env, spec) + self.prefill_kernel = self.get_prefill_kernel() + + def get_prefill_kernel(self): + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule + + return chunk_gated_delta_rule + + def _run_prefill( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + ): + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn + from lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating import fused_gdn_gating + + g, beta = fused_gdn_gating(layer.resources.a_log, a, b, layer.resources.dt_bias) + mixed_qkv = mixed_qkv.transpose(0, 1) + out_tensor = causal_conv1d_fn( + mixed_qkv, + layer.resources.conv_weight, + bias=layer.resources.conv_bias, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation=self.activation, + ) + mixed_qkv = out_tensor.transpose(0, 1) + query, key, value = self._rearrange_mixed_qkv(mixed_qkv) + initial_state = ssm_states[context.b_ssm_buffer_idx] + core_attn_out, last_recurrent_state = self.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=initial_state, + output_final_state=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[context.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + def allocate_context(self, capacity, *, capture, workspace) -> GDNPrefillContext: + return GDNPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: GDNPrefillContext): + mtp_step = self.state_slots - 1 + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx * (mtp_step + 1) + return + + def forward(self, layer: LayerAttention, context: GDNPrefillContext, inputs: GDNInputs, *, out=None) -> GDNOutput: + mixed_qkvzba: torch.Tensor = inputs.mixed_qkvzba + conv_states, ssm_states = (layer.resources.conv, layer.resources.ssm) + if (self.state_slots - 1) > 0: + conv_states = conv_states[:, :, : -(self.state_slots - 1)] + mixed_qkv, z, b, a = self._split_qkvzba(mixed_qkvzba) + core_attn_out = self._run_prefill(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + return GDNOutput(core_attn_out, z) + + +class TritonGDNDecodeBackend(TritonGDNCommon, AttentionBackend[GDNDecodeContext, GDNInputs, GDNOutput]): + phase = AttentionPhase.DECODE + + def _run_decode( + self, + layer, + context, + mixed_qkv: torch.Tensor, + z: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + ): + from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import ( + fused_recurrent_gated_delta_rule, + ) + from lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack import conv_pack_gdn_decode_inputs + + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer.resources.conv_weight, + layer.resources.conv_bias, + context.b_conv_buffer_idx, + self.activation, + self.conv_kernel_dim, + self.tp_num_k_heads, + self.head_k_dim, + self.tp_num_v_heads, + self.head_v_dim, + ) + core_attn_out, _ = fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + inplace_final_state=True, + ssm_state_indices=context.b_ssm_buffer_idx, + use_qk_l2norm_in_kernel=True, + A_log=layer.resources.a_log, + dt_bias=layer.resources.dt_bias, + a_raw=a, + b_raw=b, + ) + return GDNOutput(core_attn_out, z) + + def _run_mtp( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + ): + from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, + ) + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + cu_seqlens_q = context.b1_mtp_cu_q_seq_len + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer.resources.conv_weight, + mtp_step=self.draft_step, + bias=layer.resources.conv_bias, + activation=self.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.resources.mtp_state_index), + query_start_loc=cu_seqlens_q, + ) + query, key, value = self._rearrange_mixed_qkv(mixed_qkv, decode=False) + assert context.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" + core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( + q=query, + k=key, + v=value, + initial_state=ssm_states, + cu_seqlens=cu_seqlens_q.to(torch.long), + ssm_state_indices=context.b_ssm_buffer_idx, + ssm_state_write_indices=context.b_ssm_buffer_idx, + num_accepted_tokens=context.accepted_tokens(layer.resources.mtp_state_index), + A_log=layer.resources.a_log, + dt_bias=layer.resources.dt_bias, + a_raw=a, + b_raw=b, + ) + return core_attn_out + + def allocate_context(self, capacity, *, capture, workspace) -> GDNDecodeContext: + return GDNDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: GDNDecodeContext): + update_linear_decode(context, self.draft_step, self.spec.dynamic_verify, replay=False) + + def forward(self, layer: LayerAttention, context: GDNDecodeContext, inputs: GDNInputs, *, out=None) -> GDNOutput: + mixed_qkvzba = inputs.mixed_qkvzba + mixed_qkv, z, b, a = self._split_qkvzba(mixed_qkvzba) + conv_states, ssm_states = (layer.resources.conv, layer.resources.ssm) + draft_step = self.draft_step + if draft_step > 0: + core_attn_out = self._run_mtp(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + else: + core_attn_out, z = self._run_decode(layer, context, mixed_qkv, z, conv_states, ssm_states, a, b) + return GDNOutput(core_attn_out, z) diff --git a/lightllm/common/attention/backends/gdn/triton/kernels/__init__.py b/lightllm/common/attention/backends/gdn/triton/kernels/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk.py similarity index 91% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk.py index 7b3067bbfb..919e466472 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk.py @@ -10,14 +10,14 @@ import torch from einops import rearrange -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .chunk_o import chunk_fwd_o -from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from .cumsum import chunk_local_cumsum -from .l2norm import l2norm_fwd -from .solve_tril import solve_tril -from .utils import SUPPRESS_LEVEL, input_guard -from .wy_fast import recompute_w_u_fwd +from lightllm.common.attention.backends._shared.linear.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from lightllm.common.attention.backends._shared.linear.cumsum import chunk_local_cumsum +from lightllm.common.attention.backends._shared.linear.l2norm import l2norm_fwd +from lightllm.common.attention.backends._shared.linear.solve_tril import solve_tril +from lightllm.common.attention.backends._shared.linear.utils import SUPPRESS_LEVEL, input_guard +from lightllm.common.attention.backends.gdn.triton.kernels.chunk_o import chunk_fwd_o +from lightllm.common.attention.backends.gdn.triton.kernels.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from lightllm.common.attention.backends.gdn.triton.kernels.wy_fast import recompute_w_u_fwd def chunk_gated_delta_rule_fwd( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py similarity index 95% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py index fc49763ecd..d8c4fa91a2 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_o.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_o.py @@ -11,13 +11,12 @@ import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import exp, safe_exp -from .utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import exp, safe_exp +from lightllm.common.attention.backends._shared.linear.utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper from lightllm.common.triton_utils.autotuner import autotune BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py rename to lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py index 60a594c078..6aacf292dd 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/chunk_scaled_dot_kkt.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/chunk_scaled_dot_kkt.py @@ -9,12 +9,11 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices -from .op import safe_exp +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.op import safe_exp from lightllm.common.triton_utils.autotuner import autotune diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py b/lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py rename to lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py index 88febaffc6..a2f648753e 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/fused_gdn_gating.py @@ -67,7 +67,6 @@ def fused_gdn_gating( threshold: float = 20.0, run_config: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - if run_config is None: run_config = {"BLK_HEADS": 8, "num_warps": 1} diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py rename to lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py index 5dfbd6e4ab..1ee725b03d 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/fused_recurrent.py @@ -9,11 +9,10 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .op import exp +from lightllm.common.attention.backends._shared.linear.op import exp @triton.heuristics( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/gdn_decode_pack.py b/lightllm/common/attention/backends/gdn/triton/kernels/gdn_decode_pack.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/gdn_decode_pack.py rename to lightllm/common/attention/backends/gdn/triton/kernels/gdn_decode_pack.py diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py b/lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py rename to lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py index 5f8cab4693..57a9d2a148 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_fused_recurrent.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/mtp_fused_recurrent.py @@ -11,12 +11,13 @@ # https://github.com/fla-org/flash-linear-attention # ruff: noqa: E501 +from typing import Optional + import torch import triton import triton.language as tl -from typing import Optional -from lightllm.common.triton_utils.autotuner import autotune, AutotuneKernelType +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune # --------------------------------------------------------------------------- # Triton kernel diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py b/lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py similarity index 97% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py rename to lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py index 08bb00e644..4a8882d5fe 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/wy_fast.py +++ b/lightllm/common/attention/backends/gdn/triton/kernels/wy_fast.py @@ -10,11 +10,10 @@ # ruff: noqa: E501 import torch - import triton import triton.language as tl -from .index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) diff --git a/lightllm/common/attention/backends/kda/__init__.py b/lightllm/common/attention/backends/kda/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/kda/replay/__init__.py b/lightllm/common/attention/backends/kda/replay/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/kda/replay/backend.py b/lightllm/common/attention/backends/kda/replay/backend.py new file mode 100644 index 0000000000..97f7dc5023 --- /dev/null +++ b/lightllm/common/attention/backends/kda/replay/backend.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import KDAInputs +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclasses.dataclass(slots=True) +class ReplayKDAPrefillContext(AttentionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class ReplayKDADecodeContext(LinearDecodeContext): + pass + + +class TritonKDAReplayCommon: + state_mode = "replay" + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "triton" + + family = AttentionFamily.KDA + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + def __init__(self, env, spec): + super().__init__(env=env, spec=spec) + config = env.network_config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % env.tp_size == 0 + self.tp_num_heads = self.num_heads // env.tp_size + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.state_slots = env.state_slots + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + +class TritonKDAReplayPrefillBackend( + TritonKDAReplayCommon, AttentionBackend[ReplayKDAPrefillContext, KDAInputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + + def allocate_context(self, capacity, *, capture, workspace) -> ReplayKDAPrefillContext: + return ReplayKDAPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: ReplayKDAPrefillContext): + reqs = context.batch.b_req_idx + context.b_conv_buffer_idx = reqs + context.b_ssm_buffer_idx = reqs + + def forward( + self, layer: LayerAttention, context: ReplayKDAPrefillContext, inputs: KDAInputs, *, out=None + ) -> torch.Tensor: + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn + from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate + + conv_states, ssm_states = layer.resources.conv, layer.resources.ssm + conv_states = conv_states[..., : self.conv_kernel_size - 1] + mixed_qkv = causal_conv1d_fn( + inputs.mixed_qkv.transpose(0, 1), + layer.resources.conv_weight, + bias=None, + query_start_loc=context.batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=context.batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + ).transpose(0, 1) + q, k, v = self.split_qkv(mixed_qkv) + q = q.view(1, -1, self.tp_num_heads, self.head_dim) + k = k.view(1, -1, self.tp_num_heads, self.head_dim) + v = v.view(1, -1, self.tp_num_heads, self.head_dim) + raw_gate = inputs.raw_gate.view(1, -1, self.tp_hidden_size) + raw_beta = inputs.raw_beta.view(1, -1, self.tp_num_heads) + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, self.tp_num_heads, self.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer.resources.a_log, + g_bias=layer.resources.dt_bias, + initial_state=ssm_states[context.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=context.batch.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=self.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + +class TritonKDAReplayDecodeBackend( + TritonKDAReplayCommon, AttentionBackend[ReplayKDADecodeContext, KDAInputs, torch.Tensor] +): + phase = AttentionPhase.DECODE + + def allocate_context(self, capacity, *, capture, workspace) -> ReplayKDADecodeContext: + return ReplayKDADecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: ReplayKDADecodeContext): + update_linear_decode(context, self.draft_step, self.spec.dynamic_verify, replay=True) + + def forward( + self, layer: LayerAttention, context: ReplayKDADecodeContext, inputs: KDAInputs, *, out=None + ) -> torch.Tensor: + from .decode import replay + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + state = layer.resources.replay_state + conv_states = layer.resources.conv + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=context.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if self.draft_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=self.draft_step, + num_accepted_tokens=context.accepted_tokens(layer.resources.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update( + inputs.mixed_qkv, + conv_states, + layer.resources.conv_weight, + **conv_kwargs, + ) + q, k, v = self.split_qkv(mixed_qkv) + q = q.view(-1, self.tp_num_heads, self.head_dim) + k = k.view(-1, self.tp_num_heads, self.head_dim) + v = v.view(-1, self.tp_num_heads, self.head_dim) + raw_gate = inputs.raw_gate.reshape(-1, self.tp_hidden_size) + raw_beta = inputs.raw_beta.reshape(-1, self.tp_num_heads) + storage_layer = state.storage_layer(layer.logical_layer_id) + output = replay( + state_ssm=state.ssm_state[storage_layer], + state_keys=state.keys[storage_layer], + state_deltas=state.deltas[storage_layer], + state_gates=state.gates[storage_layer], + state_raw_keys=state.raw_keys[storage_layer], + state_raw_values=state.raw_values[storage_layer], + state_betas=state.betas[storage_layer], + state_cursors=state.cursors, + state_tile_sync_mark=state.tile_sync_mark, + q=q, + k=k, + v=v, + a=raw_gate, + b=raw_beta, + a_log=layer.resources.a_log, + bias=layer.resources.dt_bias, + reqs=context.b_conv_buffer_idx, + cu_seqlens=context.b_cu_seqlens, + state_verify_width=state.verify_width, + state_compute_mode=state.compute_mode, + state_gate_lower_bound=state.gate_lower_bound, + ) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) + return output diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/attention/backends/kda/replay/decode.py similarity index 98% rename from lightllm/common/basemodel/attention/linear/replay_kda_decode.py rename to lightllm/common/attention/backends/kda/replay/decode.py index 3704ba4fd6..5c3d38fedc 100644 --- a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py +++ b/lightllm/common/attention/backends/kda/replay/decode.py @@ -10,10 +10,7 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( - _decode_replay, - check_replay_sync_mark, -) +from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay, check_replay_sync_mark from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune diff --git a/lightllm/common/attention/backends/kda/triton/__init__.py b/lightllm/common/attention/backends/kda/triton/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/kda/triton/backend.py b/lightllm/common/attention/backends/kda/triton/backend.py new file mode 100644 index 0000000000..8b40d0b7f9 --- /dev/null +++ b/lightllm/common/attention/backends/kda/triton/backend.py @@ -0,0 +1,155 @@ +from __future__ import annotations +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature + +"""KDA attention backend for GLM-5-Next.""" + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import KDAInputs +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclasses.dataclass(slots=True) +class KDAPrefillContext(AttentionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class KDADecodeContext(LinearDecodeContext): + pass + + +class TritonKDACommon: + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "triton" + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + family = AttentionFamily.KDA + + def __init__(self, env, spec): + super().__init__(env=env, spec=spec) + config = env.network_config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % env.tp_size == 0 + self.tp_num_heads = self.num_heads // env.tp_size + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.state_slots = env.state_slots + + +class TritonKDAPrefillBackend(TritonKDACommon, AttentionBackend[KDAPrefillContext, KDAInputs, torch.Tensor]): + phase = AttentionPhase.PREFILL + + def allocate_context(self, capacity, *, capture, workspace) -> KDAPrefillContext: + return KDAPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: KDAPrefillContext): + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx * self.state_slots + + def forward( + self, layer: LayerAttention, context: KDAPrefillContext, inputs: KDAInputs, *, out=None + ) -> torch.Tensor: + from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv + from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate + + mixed_qkv = inputs.mixed_qkv + raw_gate = inputs.raw_gate + raw_beta = inputs.raw_beta + conv_states, ssm_states = (layer.resources.conv, layer.resources.ssm) + conv_states = conv_states[..., : self.conv_kernel_size - 1] + q, k, v = causal_conv1d_qkv( + mixed_qkv, + layer.resources.conv_weight, + conv_states, + context.batch.b1_cu_q_seq_len, + context.b_conv_buffer_idx, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + ) + raw_beta = raw_beta.view(1, -1, self.tp_num_heads) + initial_state = ssm_states[context.b_ssm_buffer_idx].contiguous() + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, self.tp_num_heads, self.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer.resources.a_log, + g_bias=layer.resources.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=False, + cu_seqlens=context.batch.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=self.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + +class TritonKDADecodeBackend(TritonKDACommon, AttentionBackend[KDADecodeContext, KDAInputs, torch.Tensor]): + phase = AttentionPhase.DECODE + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + def allocate_context(self, capacity, *, capture, workspace) -> KDADecodeContext: + return KDADecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: KDADecodeContext): + update_linear_decode(context, self.draft_step, self.spec.dynamic_verify, replay=False) + + def forward(self, layer: LayerAttention, context: KDADecodeContext, inputs: KDAInputs, *, out=None) -> torch.Tensor: + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda + + mixed_qkv = inputs.mixed_qkv + raw_gate = inputs.raw_gate + raw_beta = inputs.raw_beta + conv_states, ssm_states = (layer.resources.conv, layer.resources.ssm) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=context.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if self.draft_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=self.draft_step, + num_accepted_tokens=context.accepted_tokens(layer.resources.mtp_state_index), + query_start_loc=context.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update(mixed_qkv, conv_states, layer.resources.conv_weight, **conv_kwargs) + q, k, v = self.split_qkv(mixed_qkv) + shape = (1, -1) if self.draft_step > 0 else (-1, 1) + q = q.view(*shape, self.tp_num_heads, self.head_dim) + k = k.view(*shape, self.tp_num_heads, self.head_dim) + v = v.view(*shape, self.tp_num_heads, self.head_dim) + raw_gate = raw_gate.view(*shape, self.tp_hidden_size) + raw_beta = raw_beta.view(*shape, self.tp_num_heads) + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=raw_gate, + raw_beta=raw_beta, + a_log=layer.resources.a_log, + gate_bias=layer.resources.dt_bias, + initial_state=ssm_states, + lower_bound=self.lower_bound, + inplace_final_state=True, + ssm_state_indices=context.b_ssm_buffer_idx, + cu_seqlens=context.b1_mtp_cu_q_seq_len, + num_accepted_tokens=(context.accepted_tokens(layer.resources.mtp_state_index) if self.draft_step else None), + ) + return output diff --git a/lightllm/common/attention/backends/kda/triton/kernels/__init__.py b/lightllm/common/attention/backends/kda/triton/kernels/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py b/lightllm/common/attention/backends/kda/triton/kernels/kda.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py rename to lightllm/common/attention/backends/kda/triton/kernels/kda.py index 8555b83952..54eb908af4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py +++ b/lightllm/common/attention/backends/kda/triton/kernels/kda.py @@ -24,16 +24,14 @@ import torch import triton import triton.language as tl - -from lightllm.common.triton_utils.autotuner import autotune - -from .chunk_delta_h import chunk_gated_delta_rule_fwd_h -from .cumsum import chunk_local_cumsum -from .index import prepare_chunk_indices -from .l2norm import l2norm_fwd from triton.language import exp2, log -from .solve_tril import solve_tril +from lightllm.common.attention.backends._shared.linear.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from lightllm.common.attention.backends._shared.linear.cumsum import chunk_local_cumsum +from lightllm.common.attention.backends._shared.linear.index import prepare_chunk_indices +from lightllm.common.attention.backends._shared.linear.l2norm import l2norm_fwd +from lightllm.common.attention.backends._shared.linear.solve_tril import solve_tril +from lightllm.common.triton_utils.autotuner import autotune FLA_CHUNK_SIZE = 64 RCP_LN2 = 1.4426950216293335 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py b/lightllm/common/attention/backends/kda/triton/kernels/kda_decode.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda_decode.py rename to lightllm/common/attention/backends/kda/triton/kernels/kda_decode.py diff --git a/lightllm/common/attention/backends/mla/__init__.py b/lightllm/common/attention/backends/mla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/fa3/__init__.py b/lightllm/common/attention/backends/mla/fa3/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/fa3/backend.py b/lightllm/common/attention/backends/mla/fa3/backend.py new file mode 100644 index 0000000000..459ef443d6 --- /dev/null +++ b/lightllm/common/attention/backends/mla/fa3/backend.py @@ -0,0 +1,125 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.backend import AttentionContext +from lightllm.common.attention.backends.dense.fa3.backend import Fa3DensePrefillBackend, Fa3DenseDecodeBackend + + +@dataclasses.dataclass(slots=True) +class MlaFa3PrefillContext(AttentionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + causal: bool = None + + +@dataclasses.dataclass(slots=True) +class MlaFa3DecodeContext(AttentionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + page_table: torch.Tensor = None + b_att_seq_len: torch.Tensor = None + decode_max_q_seq_len: int = None + decode_max_kv_seq_len: int = None + causal: bool = None + + +class Fa3MLACommon: + provider = "fa3" + + max_head_size = None + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + family = AttentionFamily.MLA + + +class Fa3MLAPrefillBackend(Fa3MLACommon, Fa3DensePrefillBackend): + phase = AttentionPhase.PREFILL + + def _run_prefill(self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs) -> torch.Tensor: + from lightllm.utils.sgl_utils import flash_attn_varlen_func + + k_nope, k_rope = k + q_head_num = q.shape[1] + k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) + assert q.ndim == 3 and k.ndim == 3 and (v.ndim == 3) + softmax_scale = self.spec.softmax_scale + o_tensor = flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k=context.cu_seqlens_k, + max_seqlen_q=context.batch.max_q_seq_len, + max_seqlen_k=context.batch.max_kv_seq_len, + softmax_scale=softmax_scale, + causal=context.causal, + return_softmax_lse=False, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> MlaFa3PrefillContext: + return MlaFa3PrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: MlaFa3PrefillContext): + context.causal = self.spec.causal + context.cu_seqlens_q = context.batch.b1_cu_q_seq_len.int() + context.cu_seqlens_k = context.batch.b1_cu_kv_seq_len.int() + + def forward( + self, layer: LayerAttention, context: MlaFa3PrefillContext, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + return self._run_prefill(layer, context, q=q, k=k, v=v, inputs=inputs) + + +class Fa3MLADecodeBackend(Fa3MLACommon, Fa3DenseDecodeBackend): + phase = AttentionPhase.DECODE + + def _run_decode(self, layer, context, q: Tuple[torch.Tensor, torch.Tensor], k: torch.Tensor): + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + q_nope, q_rope = q + kv = k + qk_rope_head_dim = layer.resources.qk_rope_head_dim + kv_lora_rank = kv.shape[-1] - qk_rope_head_dim + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, self.kernel_page_size, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, self.kernel_page_size, 1, kv_lora_rank) + k_descale, v_descale = (None, None) + softmax_scale = self.spec.softmax_scale + o_tensor = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=kv_nope, + qv=q_nope, + page_table=context.page_table, + cache_seqlens=context.b_att_seq_len, + cu_seqlens_q=context.cu_seqlens_q, + cu_seqlens_k_new=context.cu_seqlens_k, + max_seqlen_q=context.decode_max_q_seq_len, + softmax_scale=softmax_scale, + causal=context.causal, + window_size=(-1, -1), + softcap=0.0, + k_descale=k_descale, + v_descale=v_descale, + return_softmax_lse=False, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> MlaFa3DecodeContext: + return MlaFa3DecodeContext(capacity, capture, workspace) + + def forward( + self, layer: LayerAttention, context: MlaFa3DecodeContext, inputs: MLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.resources.kv + return self._run_decode(layer, context, q=q, k=k) diff --git a/lightllm/common/attention/backends/mla/flashinfer/__init__.py b/lightllm/common/attention/backends/mla/flashinfer/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/flashinfer/backend.py b/lightllm/common/attention/backends/mla/flashinfer/backend.py new file mode 100644 index 0000000000..7f728a5915 --- /dev/null +++ b/lightllm/common/attention/backends/mla/flashinfer/backend.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature + +import dataclasses +from typing import TYPE_CHECKING +from typing import Tuple + +import torch + +from lightllm.common.attention.api import WorkspaceBudget +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.execution import WorkspaceRequirement +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.backends._flashinfer import set_flashinfer_envs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend +from lightllm.utils.dist_utils import get_dp_world_size + +if TYPE_CHECKING: + from flashinfer.mla import BatchMLAPagedAttentionWrapper + from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper + + +@dataclasses.dataclass(slots=True) +class MlaFlashInferPrefillContext(AttentionContext): + prefill_wrapper: BatchPrefillWithRaggedKVCacheWrapper | None = None + + +@dataclasses.dataclass(slots=True) +class MlaFlashInferDecodeContext(AttentionContext): + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + q_indptr: torch.Tensor = None + decode_wrapper: BatchMLAPagedAttentionWrapper | None = None + + +class FlashInferMLACommon: + supports_expert_parallel = False + + provider = "flashinfer" + + dependency = "flashinfer" + + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + family = AttentionFamily.MLA + + workspace_buffer_key = "flashinfer_mla" + + workspace_buffer_size = 256 * 1024 * 1024 + + def __init__(self, env, spec): + import triton + + set_flashinfer_envs() + super().__init__(env=env, spec=spec) + self.tp_q_head_num = spec.num_q_heads + self.qk_nope_head_dim = env.network_config["qk_nope_head_dim"] + self.qk_rope_head_dim = env.network_config["qk_rope_head_dim"] + self.kv_lora_rank = env.network_config["kv_lora_rank"] + self.v_head_dim = env.network_config["v_head_dim"] + self.q_data_type = spec.dtype + self.kv_data_type = spec.cache.dtype + self.max_seq_length = env.max_seq_len + self.max_page_num = triton.cdiv(self.max_seq_length, self.kernel_page_size) + + def resource_requirements(self, capacity, *, capture): + return WorkspaceBudget( + (WorkspaceRequirement(self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),) + ) + + +class FlashInferMLAPrefillBackend( + FlashInferMLACommon, AttentionBackend[MlaFlashInferPrefillContext, MLAPrefillInputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + supported_features = frozenset( + {AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.NONCAUSAL, AttentionFeature.MTP_WINDOW} + ) + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + ) -> torch.Tensor: + k_nope, k_rope = k + o_tensor = self.alloc_tensor((q.shape[0], q.shape[1], v.shape[-1]), dtype=q.dtype, device=q.device) + q_head_num = q.shape[1] + k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) + context.prefill_wrapper.run(q, k, v, out=o_tensor) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> MlaFlashInferPrefillContext: + return MlaFlashInferPrefillContext(capacity, capture, workspace) + + def plan_host(self, context: MlaFlashInferPrefillContext): + import flashinfer + + q_starts = context.batch.b1_cu_q_seq_len.int() + kv_starts = context.batch.b1_cu_kv_seq_len.int() + if context.prefill_wrapper is None: + context.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + context.workspace.scratch(self.workspace_buffer_key, self.workspace_buffer_size), "NHD" + ) + context.prefill_wrapper.plan( + qo_indptr=q_starts, + kv_indptr=kv_starts, + num_qo_heads=self.tp_q_head_num, + num_kv_heads=self.tp_q_head_num, + head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, + head_dim_vo=self.v_head_dim, + q_data_type=self.q_data_type, + causal=self.spec.causal, + sm_scale=self.spec.softmax_scale, + ) + return + + def forward( + self, layer: LayerAttention, context: MlaFlashInferPrefillContext, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + return self._run_prefill(layer, context, q=q, k=k, v=v) + + +class FlashInferMLADecodeBackend( + FlashInferMLACommon, AttentionBackend[MlaFlashInferDecodeContext, MLADecodeInputs, torch.Tensor] +): + phase = AttentionPhase.DECODE + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + def _refresh_decode_plan(self, context, kv_lens: torch.Tensor, max_kv_len: int): + from lightllm.common.basemodel.triton_kernel.flashinfer_mla_plan import fill_mla_decode_plan_for_cuda_graph + + fill_mla_decode_plan_for_cuda_graph( + decode_wrapper=context.decode_wrapper, + kv_page_indptr=context.kv_starts, + kv_lens=kv_lens, + batch_size=context.batch.batch_size, + num_heads=self.tp_q_head_num, + max_kv_len=max_kv_len, + ) + + def _run_decode( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + ): + qk_rope_head_dim = layer.resources.qk_rope_head_dim + q_nope, q_rope = q + o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype, device=q_nope.device) + context.decode_wrapper.run( + q_nope, + q_rope, + k[:, :, :-qk_rope_head_dim].view(-1, self.kernel_page_size, 1, k.shape[-1] - qk_rope_head_dim), + k[:, :, -qk_rope_head_dim:].view(-1, self.kernel_page_size, 1, qk_rope_head_dim), + out=o_tensor, + return_lse=False, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> MlaFlashInferDecodeContext: + return MlaFlashInferDecodeContext(capacity, capture, workspace) + + def forward( + self, layer: LayerAttention, context: MlaFlashInferDecodeContext, inputs: MLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.resources.kv + return self._run_decode(layer, context, q=q, k=k) + + def plan_host(self, context: MlaFlashInferDecodeContext): + from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor + from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index + import flashinfer + + batch = context.batch + b_page_len = (batch.b_seq_len + self.kernel_page_size - 1) // self.kernel_page_size + starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) + if context.kv_starts is None: + context.kv_starts = starts + context.kv_indices = context.workspace.scratch( + "flashinfer_indices", max(1, batch.batch_size * self.max_page_num), torch.int32 + ) + context.q_indptr = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) + else: + context.kv_starts.copy_(starts, non_blocking=True) + repack_kv_index( + req_to_token_indexs=batch.req_to_token_indexs, + b_req_idx=batch.b_req_idx, + b_token_len=batch.b_seq_len, + b_page_start_loc=context.kv_starts[:-1], + max_token_len=context.current_max_kv_seq_len, + out_page_indices=context.kv_indices, + page_size=self.kernel_page_size, + ) + if context.decode_wrapper is None: + context.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + context.workspace.scratch(self.workspace_buffer_key, self.workspace_buffer_size), + use_cuda_graph=context.capture, + qo_indptr=context.q_indptr, + kv_indices=context.kv_indices, + kv_indptr=context.kv_starts, + kv_len_arr=batch.b_seq_len, + ) + context.decode_wrapper.plan( + context.q_indptr, + context.kv_starts, + context.kv_indices, + batch.b_seq_len, + self.tp_q_head_num, + self.kv_lora_rank, + self.qk_rope_head_dim, + self.kernel_page_size, + False, + self.spec.softmax_scale, + self.q_data_type, + self.kv_data_type, + ) + else: + self._refresh_decode_plan(context, batch.b_seq_len, context.current_max_kv_seq_len) diff --git a/lightllm/common/attention/backends/mla/triton/__init__.py b/lightllm/common/attention/backends/mla/triton/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/mla/triton/backend.py b/lightllm/common/attention/backends/mla/triton/backend.py new file mode 100644 index 0000000000..16cbdad0a9 --- /dev/null +++ b/lightllm/common/attention/backends/mla/triton/backend.py @@ -0,0 +1,167 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclasses.dataclass(slots=True) +class MlaTritonPrefillContext(AttentionContext): + pass + + +@dataclasses.dataclass(slots=True) +class MlaTritonDecodeContext(AttentionContext): + block_seq: torch.Tensor = None + batch_start: torch.Tensor = None + indices_initialized: bool = False + + +class TritonMLACommon: + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "triton" + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + family = AttentionFamily.MLA + + +class TritonMLAPrefillBackend( + TritonMLACommon, AttentionBackend[MlaTritonPrefillContext, MLAPrefillInputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + + def _run_prefill( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + inputs, + ): + from lightllm.common.attention.backends.mla.triton.kernels.prefill import context_attention_fwd_with_v + + qk_rope_head_dim = layer.resources.qk_rope_head_dim + q_nope, q_rope = (q[:, :, :-qk_rope_head_dim], q[:, :, -qk_rope_head_dim:]) + o_tensor = self.alloc_tensor( + (q_nope.shape[0], q_nope.shape[1], v.shape[-1]), dtype=q_nope.dtype, device=q_nope.device + ) + k_nope, k_rope = k + softmax_scale = self.spec.softmax_scale + context_attention_fwd_with_v( + q_nope, + q_rope, + k_nope, + k_rope, + v, + o_tensor, + context.batch.b_q_start_loc, + context.batch.b1_cu_kv_seq_len, + context.batch.b_seq_len, + context.batch.b_ready_cache_len, + context.batch.max_q_seq_len, + softmax_scale, + ) + return o_tensor + + def _run_prefill_contiguous( + self, + layer, + context, + q: torch.Tensor, + k: Tuple[torch.Tensor, torch.Tensor], + v: torch.Tensor, + inputs, + ): + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, + ) + + k_nope, k_rope = k + q_head_num = q.shape[1] + k_merged = torch.cat([k_nope.expand(-1, q_head_num, -1), k_rope.expand(-1, q_head_num, -1)], dim=-1) + o_tensor = self.alloc_tensor(q.shape, dtype=q.dtype, device=q.device) + context_attention_fwd_contiguous_kv( + q, + k_merged, + v, + o_tensor, + context.batch.b_q_start_loc, + context.batch.b1_cu_kv_seq_len, + context.batch.b_seq_len, + context.batch.max_q_seq_len, + context.batch.b_ready_cache_len, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> MlaTritonPrefillContext: + return MlaTritonPrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: MlaTritonPrefillContext): + pass + + def forward( + self, layer: LayerAttention, context: MlaTritonPrefillContext, inputs: MLAPrefillInputs, *, out=None + ) -> torch.Tensor: + q = inputs.query + k = (inputs.key_nope, inputs.key_rope) + v = inputs.value + if q.shape[2] == 256: + return self._run_prefill_contiguous(layer, context, q=q, k=k, v=v, inputs=inputs) + else: + return self._run_prefill(layer, context, q=q, k=k, v=v, inputs=inputs) + + +class TritonMLADecodeBackend(TritonMLACommon, AttentionBackend[MlaTritonDecodeContext, MLADecodeInputs, torch.Tensor]): + phase = AttentionPhase.DECODE + + def _run_decode( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + ): + softmax_scale = self.spec.softmax_scale + from lightllm.common.attention.backends.mla.triton.kernels.decode import ( + gqa_token_decode_attention_flash_decoding, + ) + + qk_rope_head_dim = layer.resources.qk_rope_head_dim + q_nope, q_rope = q + kv = k + out = gqa_token_decode_attention_flash_decoding( + q_nope=q_nope, + q_rope=q_rope, + kv_nope=kv[:, :, :-qk_rope_head_dim], + kv_rope=kv[:, :, -qk_rope_head_dim:], + batch=context.batch, + softmax_scale=softmax_scale, + alloc_tensor_func=self.alloc_tensor, + block_seq=context.block_seq, + batch_start=context.batch_start, + initialize=not context.indices_initialized, + ) + context.indices_initialized = True + return out + + def allocate_context(self, capacity, *, capture, workspace) -> MlaTritonDecodeContext: + return MlaTritonDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: MlaTritonDecodeContext): + context.block_seq = torch.empty((1,), dtype=torch.int64, device=context.batch.device) + context.batch_start = torch.empty((context.batch.batch_size,), dtype=torch.int64, device=context.batch.device) + context.indices_initialized = False + + def forward( + self, layer: LayerAttention, context: MlaTritonDecodeContext, inputs: MLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.resources.kv + return self._run_decode(layer, context, q=q, k=k) diff --git a/lightllm/common/attention/backends/mla/triton/kernels/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/__init__.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py similarity index 71% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py index 063181d995..a618c9cdea 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding.py @@ -1,21 +1,35 @@ import os +from typing import List + import torch import torch.multiprocessing as mp import triton import triton.language as tl -from typing import List + +from lightllm.utils.device_utils import get_device_sm_count from lightllm.utils.log_utils import init_logger + from .gqa_flash_decoding_config import MlaDecodeAttentionKernelConfig -from lightllm.utils.device_utils import get_device_sm_count logger = init_logger(__name__) def gqa_token_decode_attention_flash_decoding( - q_nope, q_rope, kv_nope, kv_rope, infer_state, softmax_scale, out=None, alloc_tensor_func=torch.empty, **run_config + q_nope, + q_rope, + kv_nope, + kv_rope, + batch, + softmax_scale, + out=None, + alloc_tensor_func=torch.empty, + block_seq=None, + batch_start=None, + initialize=True, + **run_config, ): - batch_size = infer_state.batch_size - max_kv_seq_len = infer_state.max_kv_seq_len + batch_size = batch.batch_size + max_kv_seq_len = batch.max_kv_seq_len q_head_num, kv_lora_rank = q_nope.shape[1], q_nope.shape[2] q_rope_dim = q_rope.shape[2] @@ -28,7 +42,7 @@ def gqa_token_decode_attention_flash_decoding( if torch.cuda.is_current_stream_capturing(): avg_seq_len_in_batch = max_kv_seq_len else: - avg_seq_len_in_batch = infer_state.total_token_num // batch_size + avg_seq_len_in_batch = batch.total_token_num // batch_size run_config = MlaDecodeAttentionKernelConfig.try_to_get_best_config( batch_size=batch_size, @@ -56,35 +70,21 @@ def gqa_token_decode_attention_flash_decoding( q_rope.view(calcu_shape2), kv_nope, kv_rope, - infer_state.req_manager.req_to_token_indexs, - infer_state.b_req_idx, - infer_state.b_seq_len, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, mid_o, mid_o_logexpsum, softmax_scale, get_sm_count=True, - **run_config + **run_config, ) - if not hasattr(infer_state, "decode_att_block_seq"): - decode_att_block_seq = torch.empty( - [ - 1, - ], - dtype=torch.int64, - device="cuda", - ) - mid_o_batch_start_index = torch.empty( - [ - batch_size, - ], - dtype=torch.int64, - device="cuda", - ) + if initialize: _fwd_kernel_calcu_index_and_block_seq[(1,)]( - infer_state.b_seq_len, - decode_att_block_seq, - mid_o_batch_start_index, + batch.b_seq_len, + block_seq, + batch_start, vsm_count, batch_size, BLOCK_N=BLOCK_N, @@ -92,36 +92,27 @@ def gqa_token_decode_attention_flash_decoding( num_warps=4, ) - infer_state.decode_att_block_seq = decode_att_block_seq - infer_state.mid_o_batch_start_index = mid_o_batch_start_index - mid_o = torch.empty([q_head_num, vsm_count * 4 + batch_size, kv_lora_rank], dtype=torch.float32, device="cuda") mid_o_logexpsum = torch.empty([q_head_num, vsm_count * 4 + batch_size], dtype=torch.float32, device="cuda") flash_decode_stage1( - infer_state.decode_att_block_seq, + block_seq, q_nope.view(calcu_shape1), q_rope.view(calcu_shape2), kv_nope, kv_rope, - infer_state.req_manager.req_to_token_indexs, - infer_state.b_req_idx, - infer_state.b_seq_len, + batch.req_to_token_indexs, + batch.b_req_idx, + batch.b_seq_len, mid_o, mid_o_logexpsum, softmax_scale, get_sm_count=False, - **run_config + **run_config, ) flash_decode_stage2( - infer_state.decode_att_block_seq, - infer_state.mid_o_batch_start_index, - mid_o, - mid_o_logexpsum, - infer_state.b_seq_len, - o_tensor.view(calcu_shape1), - **run_config + block_seq, batch_start, mid_o, mid_o_logexpsum, batch.b_seq_len, o_tensor.view(calcu_shape1), **run_config ) return o_tensor diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py index be99ca9bfc..c573b907f3 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_config.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_config.py @@ -1,8 +1,10 @@ -from lightllm.common.kernel_config import KernelConfigs -from frozendict import frozendict from functools import lru_cache from typing import Dict +from frozendict import frozendict + +from lightllm.common.kernel_config import KernelConfigs + class MlaDecodeAttentionKernelConfig(KernelConfigs): kernel_name: str = "mla_decode_attentnion" @@ -51,7 +53,6 @@ def try_to_get_best_config( def save_config( cls, q_head_num: int, q_head_dim: int, q_rope_dim: int, out_dtype: str, config_json: Dict[int, Dict[int, Dict]] ): - key_params = { "q_head_num": q_head_num, "q_head_dim": q_head_dim, diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py index f5909fffde..63607b4594 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage1.py @@ -1,6 +1,7 @@ import torch import triton import triton.language as tl + from lightllm.utils.device_utils import calcu_kernel_best_vsm_count diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py similarity index 99% rename from lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py rename to lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py index 5b5e7b747c..f6f328f3e4 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/decode/gqa_flash_decoding_stage2.py @@ -1,4 +1,5 @@ import os + import torch import triton import triton.language as tl @@ -62,7 +63,7 @@ def flash_decode_stage2( mid_out_logexpsum, B_Seqlen, Out, - **run_config + **run_config, ): if run_config: num_warps = run_config["stage2_num_warps"] diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py b/lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/__init__.py rename to lightllm/common/attention/backends/mla/triton/kernels/prefill/__init__.py diff --git a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py b/lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py similarity index 100% rename from lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py rename to lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py index be06351823..a6b4c5fe94 100644 --- a/lightllm/common/basemodel/triton_kernel/mla_att/prefill_att/context_flashattention_nopad_with_v.py +++ b/lightllm/common/attention/backends/mla/triton/kernels/prefill/context_flashattention_nopad_with_v.py @@ -1,9 +1,10 @@ -import torch +import math +import torch +import torch.nn.functional as F import triton import triton.language as tl -import math -import torch.nn.functional as F + from lightllm.utils.device_utils import is_tesla @@ -145,7 +146,6 @@ def context_attention_fwd_with_v( max_input_len, softmax_scale, ): - BLOCK = 128 if not is_tesla() else 64 q_nope_dim = q_nope.shape[-1] q_rope_dim = q_rope.shape[-1] @@ -204,8 +204,8 @@ def context_attention_fwd_with_v( if __name__ == "__main__": - import torch import flashinfer + import torch Z, N_CTX, H, D_HEAD, ROPE_HEAD = 32, 1024, 16, 128, 64 dtype = torch.bfloat16 diff --git a/lightllm/common/attention/backends/sparse_mla/__init__.py b/lightllm/common/attention/backends/sparse_mla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/__init__.py b/lightllm/common/attention/backends/sparse_mla/flashmla/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py new file mode 100644 index 0000000000..11b9871a95 --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py @@ -0,0 +1,197 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend +from lightllm.utils.dist_utils import get_current_device_id + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaSparsePrefillContext(AttentionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + # Query-row metadata is consumed by env-specific sparse indexers. It is + # kept on the context so the indexer never reaches into InferStateInfo. + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + q_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaSparseDecodeContext(AttentionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + nsa_cache_seqlens: torch.Tensor = None + nsa_cu_seqlens_k_new: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + + +class FlashMLASparseMLACommon: + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + provider = "flashmla_sparse" + + dependency = "flash_mla" + + min_device_capability = (9, 0) + + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + family = AttentionFamily.SPARSE_MLA + + +class FlashMLASparseMLAPrefillBackend( + FlashMLASparseMLACommon, AttentionBackend[NsaFlashMlaSparsePrefillContext, SparseMLAPrefillInputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + + @staticmethod + def _update_query_metadata(context): + """Materialize per-query request metadata on the execution context. + + Prefill positions are supplied by the env's input preparation and + therefore remain correct for prompt-cache and speculative layouts. + ``output_size`` keeps the operation graph-capturable without a host + read of the GPU-side query lengths. + """ + batch = context.batch + position_ids = batch.position_ids + if position_ids is None: + # Unit-level callers may construct an AttentionBatch directly. + # The runtime always supplies position_ids, but contiguous query + # ranges are an unambiguous fallback for those callers. + history = batch.b_seq_len - batch.b_q_seq_len + starts = torch.repeat_interleave(history, batch.b_q_seq_len, output_size=batch.token_num) + row_starts = torch.repeat_interleave( + batch.b1_cu_q_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num + ) + position_ids = starts + torch.arange(batch.token_num, device=batch.device, dtype=starts.dtype) - row_starts + context.q_kv_seq_len = position_ids + 1 + context.q_req_idx = torch.repeat_interleave(batch.b_req_idx, batch.b_q_seq_len, output_size=batch.token_num) + context.q_kv_start_loc = torch.repeat_interleave( + batch.b1_cu_kv_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num + ) + + def _run_prefill(self, layer, context, q: torch.Tensor, kv: torch.Tensor, inputs) -> torch.Tensor: + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = self.spec.softmax_scale + kv_lora_rank = layer.resources.kv_lora_rank + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + mla_out, _, _ = flash_mla_sparse_fwd( + q=q, kv=kv, indices=topk_mem_indices, sm_scale=softmax_scale, d_v=kv_lora_rank + ) + return mla_out + + def allocate_context(self, capacity, *, capture, workspace) -> NsaFlashMlaSparsePrefillContext: + return NsaFlashMlaSparsePrefillContext(capacity, capture, workspace) + + def build_device_metadata(self, context: NsaFlashMlaSparsePrefillContext): + batch = context.batch + context.ragged_mem_index = torch.empty(batch.total_token_num, dtype=torch.int32, device=get_current_device_id()) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=batch.b_seq_len, + b_q_seq_len=batch.b_q_seq_len, + b_req_idx=batch.b_req_idx, + req_to_token_index=batch.req_to_token_indexs, + q_token_num=batch.token_num, + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=batch.hold_req_id, + ) + self._update_query_metadata(context) + return + + def forward( + self, + layer: LayerAttention, + context: NsaFlashMlaSparsePrefillContext, + inputs: SparseMLAPrefillInputs, + *, + out=None, + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.kv + return self._run_prefill(layer, context, q=q, kv=k, inputs=inputs) + + +class FlashMLASparseMLADecodeBackend( + FlashMLASparseMLACommon, AttentionBackend[NsaFlashMlaSparseDecodeContext, SparseMLADecodeInputs, torch.Tensor] +): + phase = AttentionPhase.DECODE + + def _run_decode( + self, layer, context, q: Tuple[torch.Tensor, torch.Tensor], kv: torch.Tensor, inputs + ) -> torch.Tensor: + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = self.spec.softmax_scale + kv_lora_rank = layer.resources.kv_lora_rank + qk_rope_head_dim = layer.resources.qk_rope_head_dim + q_nope, q_rope = q + k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) + kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) + o_tensor = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope, + v_cache=kv_nope, + qv=q_nope, + page_table=topk_mem_indices, + cache_seqlens=context.nsa_cache_seqlens, + cu_seqlens_q=context.batch.b1_cu_q_seq_len, + cu_seqlens_k_new=context.nsa_cu_seqlens_k_new, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=softmax_scale, + causal=True, + ) + return o_tensor + + def allocate_context(self, capacity, *, capture, workspace) -> NsaFlashMlaSparseDecodeContext: + return NsaFlashMlaSparseDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: NsaFlashMlaSparseDecodeContext): + context.q_kv_seq_len = context.batch.b_seq_len + context.q_req_idx = context.batch.b_req_idx + context.ragged_mem_index = context.workspace.scratch( + "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 + ) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=context.batch.b_seq_len, + b_q_seq_len=context.batch.b_q_seq_len, + b_req_idx=context.batch.b_req_idx, + req_to_token_index=context.batch.req_to_token_indexs, + q_token_num=context.batch.b_seq_len.shape[0], + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=context.batch.hold_req_id, + ) + context.nsa_cache_seqlens = torch.minimum( + torch.full(size=(context.batch.batch_size,), fill_value=2048, dtype=torch.int32, device="cuda"), + context.batch.b_seq_len, + ) + padded_seq_lens = torch.zeros(size=(context.nsa_cache_seqlens.shape[0] + 1,), dtype=torch.int32, device="cuda") + padded_seq_lens[1:].copy_(context.nsa_cache_seqlens, non_blocking=True) + context.nsa_cu_seqlens_k_new = padded_seq_lens.cumsum(dim=0, dtype=torch.int32) + + def forward( + self, layer: LayerAttention, context: NsaFlashMlaSparseDecodeContext, inputs: SparseMLADecodeInputs, *, out=None + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.resources.kv + return self._run_decode(layer, context, q=q, kv=k, inputs=inputs) diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py new file mode 100644 index 0000000000..7a41389d40 --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py @@ -0,0 +1,172 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import dataclasses +from typing import Tuple + +import torch + +from lightllm.common.attention.api.backend import LayerAttention +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention.api.backend import AttentionContext +from lightllm.common.attention.backends.sparse_mla.flashmla.backend import ( + FlashMLASparseMLAPrefillBackend, + FlashMLASparseMLADecodeBackend, +) + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaFp8SparsePrefillContext(AttentionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + q_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass(slots=True) +class NsaFlashMlaFp8SparseDecodeContext(AttentionContext): + ks: torch.Tensor = None + ke: torch.Tensor = None + lengths: torch.Tensor = None + ragged_mem_index: torch.Tensor = None + flashmla_sched_meta: object = None + capture_flashmla_sched_meta: object = None + q_kv_seq_len: torch.Tensor = None + q_req_idx: torch.Tensor = None + + +class FlashMLASparseMLAFp8Common: + family = AttentionFamily.SPARSE_MLA + supported_features = frozenset({AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.MTP_WINDOW}) + + cache_encodings = frozenset({CacheEncoding.FP8_SPARSE_MLA}) + + +class FlashMLASparseMLAFp8PrefillBackend(FlashMLASparseMLAFp8Common, FlashMLASparseMLAPrefillBackend): + phase = AttentionPhase.PREFILL + + def _run_prefill(self, layer, context, q: torch.Tensor, packed_kv: torch.Tensor, inputs) -> torch.Tensor: + import flash_mla + + topk_indices = inputs.topk_indices + softmax_scale = self.spec.softmax_scale + kv_lora_rank = layer.resources.kv_lora_rank + topk_mem_indices = inputs.topk_mem_indices + prefill_cache_kv = inputs.prefill_cache_kv + if context.batch.max_cache_len > 0: + from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( + get_prefill_kv_cache_and_remap_indices_triton, + ) + + kv, topk_indices = get_prefill_kv_cache_and_remap_indices_triton( + packed_kv=packed_kv, + topk_mem_indices=topk_mem_indices, + prefill_mem_index=context.batch.mem_index, + prefill_cache_kv=prefill_cache_kv, + prefill_dtype=prefill_cache_kv.dtype, + ) + else: + kv = prefill_cache_kv + if topk_indices.ndim == 2: + topk_indices = topk_indices.unsqueeze(1) + mla_out, _, _ = flash_mla.flash_mla_sparse_fwd( + q=q, kv=kv, indices=topk_indices, sm_scale=softmax_scale, d_v=kv_lora_rank + ) + return mla_out + + def allocate_context(self, capacity, *, capture, workspace) -> NsaFlashMlaFp8SparsePrefillContext: + return NsaFlashMlaFp8SparsePrefillContext(capacity, capture, workspace) + + def forward( + self, + layer: LayerAttention, + context: NsaFlashMlaFp8SparsePrefillContext, + inputs: SparseMLAPrefillInputs, + *, + out=None, + ) -> torch.Tensor: + q = inputs.query + k = layer.resources.kv + return self._run_prefill(layer, context, q=q, packed_kv=k, inputs=inputs) + + +class FlashMLASparseMLAFp8DecodeBackend(FlashMLASparseMLAFp8Common, FlashMLASparseMLADecodeBackend): + phase = AttentionPhase.DECODE + + def _run_decode( + self, layer, context, q: Tuple[torch.Tensor, torch.Tensor], packed_kv: torch.Tensor, inputs + ) -> torch.Tensor: + import flash_mla + + topk_mem_indices = inputs.topk_mem_indices + softmax_scale = self.spec.softmax_scale + kv_lora_rank = layer.resources.kv_lora_rank + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" + q_nope, q_rope = q + q_all = torch.cat([q_nope, q_rope], dim=-1).unsqueeze(1).contiguous() + kv = torch.as_strided( + packed_kv, + size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), + stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), + ) + o_tensor, _ = flash_mla.flash_mla_with_kvcache( + q=q_all, + k_cache=kv, + block_table=None, + cache_seqlens=None, + head_dim_v=kv_lora_rank, + tile_scheduler_metadata=( + context.capture_flashmla_sched_meta + if torch.cuda.is_current_stream_capturing() + else context.flashmla_sched_meta + ), + softmax_scale=softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=topk_mem_indices, + ) + return o_tensor[:, 0, :, :] + + def allocate_context(self, capacity, *, capture, workspace) -> NsaFlashMlaFp8SparseDecodeContext: + return NsaFlashMlaFp8SparseDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: NsaFlashMlaFp8SparseDecodeContext): + context.q_kv_seq_len = context.batch.b_seq_len + context.q_req_idx = context.batch.b_req_idx + context.ragged_mem_index = context.workspace.scratch( + "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 + ) + from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + + context.ks, context.ke, context.lengths = gen_nsa_ks_ke( + b_seq_len=context.batch.b_seq_len, + b_q_seq_len=context.batch.b_q_seq_len, + b_req_idx=context.batch.b_req_idx, + req_to_token_index=context.batch.req_to_token_indexs, + q_token_num=context.batch.b_seq_len.shape[0], + ragged_mem_index=context.ragged_mem_index, + hold_req_idx=context.batch.hold_req_id, + ) + return + + def forward( + self, + layer: LayerAttention, + context: NsaFlashMlaFp8SparseDecodeContext, + inputs: SparseMLADecodeInputs, + *, + out=None, + ) -> torch.Tensor: + q = (inputs.query_nope, inputs.query_rope) + k = layer.resources.kv + return self._run_decode(layer, context, q=q, packed_kv=k, inputs=inputs) + + def plan_host(self, context: NsaFlashMlaFp8SparseDecodeContext): + import flash_mla + + if context.flashmla_sched_meta is None: + context.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] + context.capture_flashmla_sched_meta = flash_mla.get_mla_metadata()[0] diff --git a/lightllm/common/attention/probe.py b/lightllm/common/attention/probe.py new file mode 100644 index 0000000000..b16a5d1fc3 --- /dev/null +++ b/lightllm/common/attention/probe.py @@ -0,0 +1,83 @@ +from lightllm.common.attention.api.environment import BackendEnv +from lightllm.common.attention.session import AttentionSession + +"""启动子进程中的小规模数值检查，不接触模型的真实 KV 或请求状态。""" + +import torch + +from .api import AttentionBatch, AttentionPhase, ContextCapacity, LayerAttention +from .api.inputs import DenseInputs, DenseResources +from .api.workspace import AttentionWorkspace + + +def probe_dense(backend_cls, spec, runtime): + page = spec.cache.page_size + stride = ((16 + page - 1) // page) * page + q_heads, kv_heads, dim = spec.num_q_heads, spec.cache.key_heads, spec.head_size + env = BackendEnv( + runtime=runtime, + network_config={}, + hold_req_id=2, + draft_step=0, + max_seq_len=stride, + tp_size=1, + alloc_tensor=lambda shape, dtype, device=runtime.device: torch.empty(shape, dtype=dtype, device=device), + ) + backend = backend_cls(env, spec) + device = runtime.device + lengths = (7, 13) + q_lengths = (2, 3) if spec.phase is AttentionPhase.PREFILL else (1, 1) + query = torch.randn((sum(q_lengths), q_heads, dim), dtype=spec.dtype, device=device) + key, value = (torch.randn((3 * stride, kv_heads, dim), dtype=spec.dtype, device=device) for _ in range(2)) + batch = AttentionBatch( + phase=spec.phase, + batch_size=2, + token_num=sum(q_lengths), + total_token_num=sum(lengths), + max_q_seq_len=max(q_lengths), + max_kv_seq_len=max(lengths), + max_cache_len=max(n - q for n, q in zip(lengths, q_lengths)), + b_req_idx=torch.tensor([0, 1], dtype=torch.int32, device=device), + b_seq_len=torch.tensor(lengths, dtype=torch.int32, device=device), + b_q_seq_len=torch.tensor(q_lengths, dtype=torch.int32, device=device), + b1_cu_q_seq_len=torch.tensor([0, q_lengths[0], sum(q_lengths)], dtype=torch.int32, device=device), + b1_cu_kv_seq_len=torch.tensor([0, lengths[0], sum(lengths)], dtype=torch.int32, device=device), + req_to_token_indexs=torch.arange(3 * stride, dtype=torch.int32, device=device).view(3, stride), + hold_req_id=2, + b_ready_cache_len=torch.tensor([n - q for n, q in zip(lengths, q_lengths)], dtype=torch.int32, device=device), + ) + capacity = ContextCapacity( + phase=spec.phase, batch_size=2, token_num=sum(q_lengths), max_kv_seq_len=stride, max_q_seq_len=max(q_lengths) + ) + workspace = AttentionWorkspace(runtime=runtime) + workspace.reserve(backend.resource_requirements(capacity, capture=False)) + session = AttentionSession() + context = session.allocate(backend, capacity, capture=False, workspace=workspace) + session.plan(backend, context, batch) + session.build(backend, context) + session.contexts[0] = context + attention = LayerAttention( + group_id=0, + logical_layer_id=0, + resources=DenseResources(key, value), + backend=backend, + ) + actual = attention.forward(session, DenseInputs(query), out=torch.empty_like(query)) + expected, start = [], 0 + for request, (length, q_length) in enumerate(zip(lengths, q_lengths)): + q = query[start : start + q_length].float() + k = key[request * stride : request * stride + length].float().repeat_interleave(q_heads // kv_heads, dim=1) + v = value[request * stride : request * stride + length].float().repeat_interleave(q_heads // kv_heads, dim=1) + logits = torch.einsum("qhd,khd->hqk", q, k) * dim ** -0.5 + if spec.causal: + mask = torch.arange(length, device=device)[None, :] <= ( + length - q_length + torch.arange(q_length, device=device)[:, None] + ) + logits.masked_fill_(~mask, -torch.inf) + expected.append(torch.einsum("hqk,khd->qhd", logits.softmax(-1), v)) + start += q_length + expected = torch.cat(expected).to(spec.dtype) + runtime.synchronize() + if not torch.allclose(actual, expected, atol=0.02, rtol=0.02): + return False, f"dense {spec.phase.value} output mismatch: max diff {(actual - expected).abs().max().item()}" + return True, None diff --git a/lightllm/common/attention/registry.py b/lightllm/common/attention/registry.py new file mode 100644 index 0000000000..cc7b5b8e80 --- /dev/null +++ b/lightllm/common/attention/registry.py @@ -0,0 +1,87 @@ +"""Lazy provider declarations and default preference order; importing this does not load kernels.""" + +import importlib + +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.platform import Platform + + +def load_backend_class(reference): + if isinstance(reference, type): + return reference + module, name = reference.split(":") + if not module.startswith("lightllm."): + module = "lightllm.common.attention." + module + return getattr(importlib.import_module(module), name) + + +varlen_attention_ops = { + "fa3": "backends.dense.fa3.varlen:Fa3VarlenAttentionOp", + "xformers": "backends.dense.xformers.varlen:XformersVarlenAttentionOp", + "sdpa": "backends.dense.sdpa.varlen:SdpaVarlenAttentionOp", + "triton": "backends.dense.triton.varlen:TritonVarlenAttentionOp", +} + + +VISION_VARLEN_PRIORITY = ("fa3", "xformers", "sdpa", "triton") +AUDIO_VARLEN_PRIORITY = ("fa3", "triton") + + +def get_varlen_attention_op_class(name): + if name not in varlen_attention_ops: + raise ValueError(f"unknown varlen attention operator: {name}") + return load_backend_class(varlen_attention_ops[name]) + + +BACKENDS = ( + "lightllm.common.attention.backends.dense.fa3.backend:Fa3DenseDecodeBackend", + "lightllm.common.attention.backends.dense.fa3.backend:Fa3DensePrefillBackend", + "lightllm.common.attention.backends.dense.fa3.fp8:Fa3DenseFp8DecodeBackend", + "lightllm.common.attention.backends.dense.fa3.fp8:Fa3DenseFp8PrefillBackend", + "lightllm.common.attention.backends.dense.fa3.neo:Fa3DenseNeoDecodeBackend", + "lightllm.common.attention.backends.dense.fa3.neo:Fa3DenseNeoPrefillBackend", + "lightllm.common.attention.backends.dense.flashinfer.backend:FlashInferDenseDecodeBackend", + "lightllm.common.attention.backends.dense.flashinfer.backend:FlashInferDensePrefillBackend", + "lightllm.common.attention.backends.dense.flashinfer.fp8:FlashInferDenseFp8DecodeBackend", + "lightllm.common.attention.backends.dense.flashinfer.fp8:FlashInferDenseFp8PrefillBackend", + "lightllm.common.attention.backends.dense.triton.backend:TritonDenseDecodeBackend", + "lightllm.common.attention.backends.dense.triton.backend:TritonDensePrefillBackend", + "lightllm.common.attention.backends.dense.triton.int4:TritonDenseInt4DecodeBackend", + "lightllm.common.attention.backends.dense.triton.int4:TritonDenseInt4PrefillBackend", + "lightllm.common.attention.backends.dense.triton.int8:TritonDenseInt8DecodeBackend", + "lightllm.common.attention.backends.dense.triton.int8:TritonDenseInt8PrefillBackend", + "lightllm.common.attention.backends.dense.triton.neo:TritonDenseNeoDecodeBackend", + "lightllm.common.attention.backends.dense.triton.neo:TritonDenseNeoPrefillBackend", + "lightllm.common.attention.backends.gdn.flashqla.backend:FlashQlaGDNPrefillBackend", + "lightllm.common.attention.backends.gdn.replay.backend:FlashQlaGDNReplayPrefillBackend", + "lightllm.common.attention.backends.gdn.replay.backend:TritonGDNReplayDecodeBackend", + "lightllm.common.attention.backends.gdn.replay.backend:TritonGDNReplayPrefillBackend", + "lightllm.common.attention.backends.gdn.triton.backend:TritonGDNDecodeBackend", + "lightllm.common.attention.backends.gdn.triton.backend:TritonGDNPrefillBackend", + "lightllm.common.attention.backends.kda.replay.backend:TritonKDAReplayDecodeBackend", + "lightllm.common.attention.backends.kda.replay.backend:TritonKDAReplayPrefillBackend", + "lightllm.common.attention.backends.kda.triton.backend:TritonKDADecodeBackend", + "lightllm.common.attention.backends.kda.triton.backend:TritonKDAPrefillBackend", + "lightllm.common.attention.backends.mla.fa3.backend:Fa3MLADecodeBackend", + "lightllm.common.attention.backends.mla.fa3.backend:Fa3MLAPrefillBackend", + "lightllm.common.attention.backends.mla.flashinfer.backend:FlashInferMLADecodeBackend", + "lightllm.common.attention.backends.mla.flashinfer.backend:FlashInferMLAPrefillBackend", + "lightllm.common.attention.backends.mla.triton.backend:TritonMLADecodeBackend", + "lightllm.common.attention.backends.mla.triton.backend:TritonMLAPrefillBackend", + "lightllm.common.attention.backends.sparse_mla.flashmla.backend:FlashMLASparseMLADecodeBackend", + "lightllm.common.attention.backends.sparse_mla.flashmla.backend:FlashMLASparseMLAPrefillBackend", + "lightllm.common.attention.backends.sparse_mla.flashmla.fp8:FlashMLASparseMLAFp8DecodeBackend", + "lightllm.common.attention.backends.sparse_mla.flashmla.fp8:FlashMLASparseMLAFp8PrefillBackend", +) + + +def backend_priority(family, phase, specs): + if family is AttentionFamily.GDN: + return ("flashqla", "triton") if phase is AttentionPhase.PREFILL else ("triton",) + if family is AttentionFamily.KDA: + return ("triton",) + if family is AttentionFamily.SPARSE_MLA: + return ("flashmla_sparse",) + if phase is AttentionPhase.DECODE and not any(spec.use_mtp for spec in specs): + return ("flashinfer", "fa3", "triton") + return ("fa3", "flashinfer", "triton") diff --git a/lightllm/common/attention/selection.py b/lightllm/common/attention/selection.py new file mode 100644 index 0000000000..006aedd94c --- /dev/null +++ b/lightllm/common/attention/selection.py @@ -0,0 +1,176 @@ +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from lightllm.common.attention.api.backend import UnsupportedAttention +from lightllm.common.attention.api.spec import AttentionFamily, AttentionPhase +from lightllm.common.attention.registry import ( + VISION_VARLEN_PRIORITY, + BACKENDS, + backend_priority, + load_backend_class, + varlen_attention_ops, +) +from lightllm.platform import get_device_runtime +from lightllm.utils import backend_validator +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +class BackendSelectionError(RuntimeError): + def __init__(self, family, phase, failures): + self.failures = failures + details = "; ".join(f"{name}: {reason}" for name, reason in failures.items()) + super().__init__(f"No usable {family.value} {phase.value} attention backend: {details}") + + +@dataclass(frozen=True) +class BackendPolicy: + allow: frozenset[str] | None = None + prefer: tuple[str, ...] = () + extra: tuple = () + + +class VarlenSelectionError(RuntimeError): + def __init__(self, failures): + self.failures = failures + super().__init__( + "No usable varlen attention operator: " + + "; ".join(f"{name}: {reason}" for name, reason in failures.items()) + ) + + +def select_varlen_attention_op(requested="auto", *, device=None, devices=None, dtype=None, head_dim=80): + if isinstance(dtype, str): + aliases = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"} + dtype = getattr(torch, aliases.get(dtype, dtype)) + if devices is None: + # Resolve once before probing; probes receive concrete device ordinals. + devices = (get_device_runtime(device).device,) + else: + devices = tuple(get_device_runtime(d).device for d in devices) + if not devices: + raise ValueError("varlen selection requires at least one device") + ordered = VISION_VARLEN_PRIORITY if requested == "auto" else (requested,) + failures = {} + for name in ordered: + if name not in varlen_attention_ops: + failures[name] = "unknown varlen attention operator" + continue + # 在子进程中检查实际 Q/K/V 路径，避免父进程初始化 CUDA 或误用文本 attention 探针。 + for assigned_device in devices: + result = backend_validator.probe_varlen(name, device=assigned_device, dtype=dtype, head_dim=head_dim) + if not result.ok: + failures[name] = f"{result.scope.value} on {assigned_device}: {result.reason}" + break + if name in failures: + continue + logger.info("Selected %s varlen attention after %s", name, result.scope.value) + return name + raise VarlenSelectionError(failures) + + +def select_attention_backend( + *, + specs, + caps, + execution, + requested="auto", + policy=BackendPolicy(), +): + specs = tuple(specs) + if not specs: + raise ValueError("attention selection requires at least one spec") + first = specs[0] + family, phase, encoding = first.family, first.phase, first.cache.encoding + if any((spec.family, spec.phase, spec.cache) != (family, phase, first.cache) for spec in specs): + raise ValueError("attention specs must share family, phase and cache layout") + available = {} + import_failures = {} + for reference in (*BACKENDS, *policy.extra): + try: + backend = load_backend_class(reference) + except (ImportError, OSError) as exc: + import_failures[str(reference)] = str(exc) + continue + if (backend.family, backend.phase, backend.variant, backend.state_mode) != ( + family, + phase, + first.variant, + first.cache.state_mode, + ) or encoding not in backend.cache_encodings: + continue + if policy.allow is None or backend.provider in policy.allow: + available[backend.provider] = reference + ordered = [requested] if requested != "auto" else list(policy.prefer or backend_priority(family, phase, specs)) + failures = {} + selected_backend = None + selected_name = None + for name in ordered: + if name not in available: + failures[ + name + ] = f"no implementation for cache encoding {encoding.value}; unavailable providers: {import_failures}" + continue + probe_scopes = set() + try: + backend = load_backend_class(available[name]) + backend.check_environment(caps) + for spec in specs: + backend.check_support(spec, caps, execution) + for spec in specs: + result = backend_validator.probe_backend( + available[name], device=str(caps.device), dtype=spec.dtype, spec=spec + ) + probe_scopes.add(result.scope.value) + if not result.ok: + failures[name] = f"{result.scope.value}: {result.reason}" + break + if name in failures: + continue + except UnsupportedAttention as exc: + failures[name] = str(exc) + continue + except Exception as exc: + # All ranks must reach consensus, including a rank whose provider check failed. + logger.exception( + "Unexpected %s attention provider failure: family=%s phase=%s device=%s specs=%r", + name, + family.value, + phase.value, + caps.device, + specs, + ) + failures[name] = f"provider check failed: {type(exc).__name__}: {exc}" + continue + logger.info( + "Selected %s %s attention on %s: %s (%s)", + family.value, + phase.value, + caps.device, + name, + "; ".join(sorted(probe_scopes)), + ) + selected_backend = backend + selected_name = name + break + + error = None if selected_backend is not None else BackendSelectionError(family, phase, failures) + _check_rank_selection(family, phase, selected_name, None if error is None else str(error)) + if error is not None: + raise error + return selected_backend + + +def _check_rank_selection(family, phase, selected_name, failure): + if dist.is_initialized() and dist.get_world_size() > 1: + selections = [None] * dist.get_world_size() + dist.all_gather_object(selections, {"backend": selected_name, "failure": failure}) + selected = {item["backend"] for item in selections if item["backend"] is not None} + rejected = {f"rank {rank}": item["failure"] for rank, item in enumerate(selections) if item["failure"]} + if len(selected) > 1: + rejected["backend selection"] = ", ".join(sorted(selected)) + if rejected: + raise BackendSelectionError(family, phase, rejected) diff --git a/lightllm/common/attention/session.py b/lightllm/common/attention/session.py new file mode 100644 index 0000000000..54828df48e --- /dev/null +++ b/lightllm/common/attention/session.py @@ -0,0 +1,131 @@ +from contextlib import ExitStack, contextmanager +from dataclasses import replace + +import torch + + +class AttentionSession: + def __init__(self): + self.groups = () + self.contexts = {} + self.batch = None + + @contextmanager + def prepare(self, batch, groups, capacity, *, workspace, capture=False, replay=False): + batch.validate() + capacity.check(batch) + self.groups = groups + self.batch = batch + with ExitStack() as leases: + if not replay: + if capture: + self.batch = replace( + batch, + max_kv_seq_len=capacity.max_kv_seq_len, + total_token_num=capacity.batch_size * capacity.max_kv_seq_len, + ) + for group_id, group in enumerate(groups): + scratch = workspace(group_id) + leases.enter_context(scratch.lease()) + backend = group[batch.phase] + if capture: + scratch = scratch.borrow_for_capture() + scratch.reserve(backend.resource_requirements(capacity, capture=capture)) + context = self.allocate(backend, capacity, capture=capture, workspace=scratch) + context.batch = self.batch + context.current_max_kv_seq_len = batch.max_kv_seq_len + backend.plan_host(context) + self.contexts[group_id] = context + yield self + + @staticmethod + def allocate(backend, capacity, *, capture, workspace): + if capacity.phase is not backend.phase: + raise ValueError("backend and context phases differ") + context = backend.allocate_context(capacity, capture=capture, workspace=workspace) + if context.capacity != capacity: + raise ValueError("backend changed attention context capacity") + context.backend_identity = backend + return context + + @staticmethod + def plan(backend, context, batch, *, copy_inputs=True): + AttentionSession._check_context(backend, context) + batch.validate() + context.capacity.check(batch) + context.current_max_kv_seq_len = batch.max_kv_seq_len + if context.capture and context.batch is not None: + if copy_inputs: + context.batch.copy_inputs(batch) + elif context.capture: + context.batch = replace( + batch, + max_kv_seq_len=context.capacity.max_kv_seq_len, + total_token_num=context.capacity.batch_size * context.capacity.max_kv_seq_len, + ) + else: + context.batch = batch + backend.plan_host(context) + + @staticmethod + def _check_context(backend, context): + if context.backend_identity is not backend or context.capacity.phase is not backend.phase: + raise ValueError("context belongs to a different attention backend or phase") + + def context(self, layer): + context = self.contexts[layer.group_id] + self._check_context(layer.backend, context) + if context.batch is None: + raise RuntimeError("prepare attention inputs before execution") + return context + + @staticmethod + def build(backend, context): + AttentionSession._check_context(backend, context) + if context.batch is None: + raise RuntimeError("prepare attention inputs before update") + backend.build_device_metadata(context) + + def build_device_metadata(self): + for group_id, context in self.contexts.items(): + self.build(self.groups[group_id][context.capacity.phase], context) + + def forward(self, layer, inputs, *, out=None): + context = self.context(layer) + result = layer.backend.forward(layer, context, inputs, out=out) + if out is None: + return result + attention = result[0] if isinstance(result, tuple) else result + if attention.dtype != out.dtype or attention.device != out.device: + raise ValueError("out must match the attention result dtype and device") + if attention.untyped_storage().data_ptr() == out.untyped_storage().data_ptr(): + begin = attention.storage_offset() + end = begin + sum((size - 1) * stride for size, stride in zip(attention.shape, attention.stride())) + 1 + out_end = ( + out.storage_offset() + sum((size - 1) * stride for size, stride in zip(out.shape, out.stride())) + 1 + ) + if attention.numel() and (begin < out.storage_offset() or end > out_end): + raise ValueError("attention result view extends beyond out") + return result + if attention.numel() != out.numel(): + raise ValueError("out must match the attention result element count") + out.copy_(attention.reshape(out.shape)) + if isinstance(result, tuple): + return type(result)(out, *result[1:]) if hasattr(result, "_fields") else (out, *result[1:]) + return out + + def replay_input_copies(self, source): + if len(self.groups) != len(source.groups) or any( + current is not incoming for current, incoming in zip(self.groups, source.groups) + ): + raise ValueError("captured attention belongs to a different model role or group") + + source.batch.validate() + for context in self.contexts.values(): + context.capacity.check(source.batch) + return self.batch.input_copies(source.batch) + + def plan_replay(self, source): + for group_id, context in self.contexts.items(): + context.current_max_kv_seq_len = source.batch.max_kv_seq_len + self.groups[group_id][context.capacity.phase].plan_host(context) diff --git a/lightllm/common/attention/varlen.py b/lightllm/common/attention/varlen.py new file mode 100644 index 0000000000..158a626cfd --- /dev/null +++ b/lightllm/common/attention/varlen.py @@ -0,0 +1,44 @@ +from lightllm.platform import get_device_runtime + +from .api import AttentionFamily, AttentionPhase, UnsupportedAttention +from .registry import AUDIO_VARLEN_PRIORITY, get_varlen_attention_op_class +from .selection import BackendSelectionError + + +class AutoVarlenAttentionRunner: + """One encoder owns this runner and shares stateless operators across its layers. + + Unsupported inputs can try the next provider for this call. Kernel and programming + errors propagate: retrying on a potentially failed device is not a safe fallback. + Environment rejections are cached by provider/device; input rejections are not. + """ + + def __init__(self, providers=AUDIO_VARLEN_PRIORITY): + self.providers = providers + self.ops = {} + self.unavailable = {} + + def forward(self, inputs, *, out): + failures = {} + caps = get_device_runtime(inputs.query.device).caps + for name in self.providers: + key = (name, inputs.query.device) + if key in self.unavailable: + failures[name] = self.unavailable[key] + continue + if key not in self.ops: + try: + op_cls = get_varlen_attention_op_class(name) + op_cls.check_environment(caps) + self.ops[key] = op_cls() + except (UnsupportedAttention, ImportError) as exc: + self.unavailable[key] = failures[name] = str(exc) + continue + op = self.ops[key] + try: + op.check_support(inputs, caps) + except UnsupportedAttention as exc: + failures[name] = str(exc) + continue + return op.forward(inputs, out=out) + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, failures) diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index 6a96e76dc9..013a8b2682 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/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") diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index 82eb4026c5..1a578c8b83 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/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 @@ -24,8 +24,12 @@ class GDNState(BoundedAttentionState): + state_mode = "native" + """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。""" + storage_layout = ("gdn_conv_ssm", 1) + big_pages_pool: GDNPagesPool | None small_pages_pool: GDNPagesPool | None @@ -115,11 +119,26 @@ def contains_logical_layer(self, logical_layer_id: int) -> bool: """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer + def layer_resources(self, layer_id): + from lightllm.common.attention.api.inputs import RecurrentStorage + + conv, ssm = self.get_layer_view(layer_id) + return RecurrentStorage(conv, ssm, self.mtp_state_index) + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" 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" diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index 137ce8ad32..e45fc34e3d 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -1,19 +1,20 @@ +from __future__ import annotations + """GDN replay 的 bounded attention 状态。 Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 一层 decode 重放在 replay_gdn_decode。 """ -from __future__ import annotations from functools import lru_cache from math import prod import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion @@ -28,6 +29,8 @@ class ReplayGDNState(BoundedAttentionState): + state_mode = "replay" + """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" # ======================================================================== @@ -117,6 +120,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 ( @@ -168,6 +194,12 @@ def storage_layer(self, logical_layer_id: int) -> int: return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + from lightllm.common.attention.api.inputs import RecurrentStorage + + conv, ssm = self.get_layer_view(layer_id) + return RecurrentStorage(conv, ssm, self.mtp_state_index, self) + def get_layer_view(self, logical_layer_id: int): index = self.storage_layer(logical_layer_id) return self.conv_state[index], self.ssm_state[index] diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index 11d7425cd9..17f226182c 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/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 的资源、层视图和运行态核心能力。""" diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index 68097905d5..7ee2da5d58 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/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 @@ -24,8 +24,12 @@ class KDAState(BoundedAttentionState): + state_mode = "native" + """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。""" + storage_layout = ("kda_conv_ssm", 1) + big_pages_pool: KDAPagesPool | None small_pages_pool: KDAPagesPool | None @@ -115,11 +119,26 @@ def contains_logical_layer(self, logical_layer_id: int) -> bool: """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" return logical_layer_id in self._logical_to_storage_layer + def layer_resources(self, layer_id): + from lightllm.common.attention.api.inputs import RecurrentStorage + + conv, ssm = self.get_layer_view(layer_id) + return RecurrentStorage(conv, ssm, self.mtp_state_index) + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" 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" diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index 2c73e97d6f..f9fb7d82e7 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -1,19 +1,20 @@ +from __future__ import annotations + """KDA replay 的 bounded attention 状态。 Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 一层 decode 重放在 replay_kda_decode。 """ -from __future__ import annotations from functools import lru_cache from math import prod import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.utils.log_utils import init_logger from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion @@ -28,6 +29,8 @@ class ReplayKDAState(BoundedAttentionState): + state_mode = "replay" + """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" # ======================================================================== @@ -117,6 +120,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 ( @@ -168,6 +194,12 @@ def storage_layer(self, logical_layer_id: int) -> int: return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + from lightllm.common.attention.api.inputs import RecurrentStorage + + conv, ssm = self.get_layer_view(layer_id) + return RecurrentStorage(conv, ssm, self.mtp_state_index, self) + def get_layer_view(self, logical_layer_id: int): index = self.storage_layer(logical_layer_id) return self.conv_state[index], self.ssm_state[index] diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py index 5afdfab160..d0fd7d84ec 100644 --- a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py +++ b/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) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -139,10 +140,17 @@ def initialize_page_pools( def contains_logical_layer(self, logical_layer_id: int) -> bool: return logical_layer_id in self._logical_to_storage_layer + def layer_resources(self, layer_id): + row_bytes = self.config.get_packed_kv_bytes(self.config.page_token_size) + return self.get_layer_view(layer_id).view(-1, row_bytes) + 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_() diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py index 5d9a0ffe12..472acec80c 100644 --- a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py @@ -2,9 +2,8 @@ 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 @@ -12,8 +11,8 @@ 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_() diff --git a/lightllm/common/attention_state/commit.py b/lightllm/common/attention_state/commit.py new file mode 100644 index 0000000000..9e02a6c799 --- /dev/null +++ b/lightllm/common/attention_state/commit.py @@ -0,0 +1,34 @@ +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class RequestGenerationStamp: + role: str + request_indices: torch.Tensor + mtp_indices: torch.Tensor + request_generations: torch.Tensor + + @classmethod + def stamp(cls, role, request_indices, mtp_indices, generations): + requests = request_indices.clone() + return cls(role, requests, mtp_indices.clone(), generations[requests.long()].clone()) + + @classmethod + def concatenate(cls, commits): + if not commits or any(commit.role != commits[0].role for commit in commits): + raise ValueError("MTP outputs must belong to the same model role") + return cls( + commits[0].role, + *( + torch.cat([getattr(commit, name) for commit in commits]) + for name in ("request_indices", "mtp_indices", "request_generations") + ), + ) + + def validate_generation(self, generations): + torch._assert_async( + torch.all(generations[self.request_indices.long()] == self.request_generations), + "MTP output refers to a released or reused request slot", + ) diff --git a/lightllm/common/attention_state/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py index d3f00430b0..d55acc1639 100644 --- a/lightllm/common/attention_state/full_att/base_state.py +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -1,12 +1,31 @@ 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 FullAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): """单种 full-attention 存储方案的数据接口，不分配 token 索引或维护请求映射。 - 所有表示 token 索引的参数都须为调用方准备的 CUDA Tensor，位于对应存储设备上； + 所有表示 token 索引的参数都须为调用方准备的同设备 Tensor，位于对应存储设备上； State 不转换设备、dtype 或连续布局，int32/int64 由具体算子支持，DP 目标索引须为 int64。 - CPU 页位于主机内存不改变 token 索引的 CUDA 契约。 + 传输操作的设备限制由具体实现声明，CPU 页位于主机内存不改变索引与存储同设备的契约。 """ + + 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") diff --git a/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py index 21ce083697..63463ac18f 100644 --- a/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py +++ b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py @@ -2,21 +2,20 @@ from __future__ import annotations -from functools import lru_cache - -from typing import TYPE_CHECKING, List from dataclasses import dataclass +from functools import lru_cache from math import lcm, prod +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args -from ...utils import deepcopy_result from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result from ..base_state import FullAttentionState -from ..packed_mla.packed_transfer import transfer_packed_pages from ..dp_transfer import copy_packed_tokens +from ..packed_mla.packed_transfer import transfer_packed_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -103,6 +102,8 @@ class CompressedAttentionState(FullAttentionState): 当前实现固定为单个共享 KV head，各 TP rank 完整复制；不执行 CSA/HCA 的历史选择。 """ + storage_layout = ("compressed_fp8_e8m0_token_row", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -170,6 +171,9 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + return self.get_layer_view(layer_id).kv + def get_layer_view(self, logical_layer_id: int) -> CompressedAttentionLayerView: """从同一物理压缩块取 KV 和 indexer 视图，不重排数据或创建副本。""" config = self.config diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py index fe184b4110..4fc95b0fae 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -1,20 +1,23 @@ from __future__ import annotations +from dataclasses import dataclass from functools import lru_cache -from typing import TYPE_CHECKING, List from math import prod -from dataclasses import dataclass +from typing import TYPE_CHECKING, List import torch +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.platform import Platform, get_device_runtime from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from .gqa_transfer import transfer_gqa_pages from ..dp_transfer import copy_gqa_tokens +from .gqa_transfer import transfer_gqa_pages +from .ops import get_kv_operations if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -78,6 +81,8 @@ class GQAFullAttentionState(FullAttentionState): query 与 KV 的分组计算由 attention backend 处理，因此三种情况共用此 State。 """ + storage_layout = ("gqa_separate_token_head_dim", 1) + # 保留既有 K/V 编码的稳定传输布局标识。 # ======================================================================== @@ -103,6 +108,8 @@ def __init__( self.page_size = page_size self.device = torch.device(device=device) self.capacity = 0 + self.zero_initialize_cache = False + self.ops = None self.k_storage: torch.Tensor | None = None self.v_storage: torch.Tensor | None = None @@ -114,8 +121,13 @@ def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity shape = self._storage_shape() - self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) - self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + runtime = get_device_runtime(self.device) + self.ops = get_kv_operations(runtime) + self.k_storage = runtime.allocate(shape, self.storage_dtype) + self.v_storage = runtime.allocate(shape, self.storage_dtype) + if self.zero_initialize_cache: + self.k_storage.zero_() + self.v_storage.zero_() def get_token_bytes(self) -> int: """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" @@ -134,31 +146,30 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + from lightllm.common.attention.api.inputs import DenseResources + + return DenseResources(*self.get_layer_view(layer_id)) + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: - """通过 CUDA 算子将形状匹配的 K/V 数据写入指定物理 token 位置。""" - assert token_indices.is_cuda, "token indices must reside on CUDA" + """通过已绑定的设备实现将 K/V 数据写入指定物理 token 位置。""" storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) - if not len(token_indices): + assert token_indices.device == kv.device == self.k_storage.device, "KV writes require same-device tensors" + assert token_indices.ndim == 1 and token_indices.dtype in (torch.int32, torch.int64) + assert kv.ndim == 3 and kv.shape[1:] == (2 * self.config.head_num, self.config.head_dim) + assert kv.shape[0] == token_indices.numel() and kv.dtype == self.storage_dtype + if not token_indices.numel(): return - assert self.k_storage.is_cuda, "KV writes require CUDA storage" - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv - - # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 - destindex_copy_kv( - K=kv[:, : self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index] - ) - destindex_copy_kv( - K=kv[:, self.config.head_num :], DestLoc=token_indices, Out=self.v_storage[storage_layer_index] - ) + self.ops.write(self.k_storage[storage_layer_index], self.v_storage[storage_layer_index], token_indices, kv) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" - assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert src_indices.device == dst_indices.device == self.k_storage.device, "token indices must match storage" assert ( src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() ), "token copies require equal one-dimensional index tensors" @@ -166,26 +177,17 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten torch.int32, torch.int64, ), "token indices must be integer tensors" - assert self.k_storage.is_cuda, "token copies require CUDA storage" - if not src_indices.numel(): - return - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer - - # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 - # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 - # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 - # TODO：将 K/V 的同索引 token 复制融合为一次 kernel 发射， - # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 - copy_kv_buffer_to_kv_buffer( - src_mem_index=src_indices, - dst_mem_index=dst_indices, - kv_buffer=self.k_storage, - ) - copy_kv_buffer_to_kv_buffer( - src_mem_index=src_indices, - dst_mem_index=dst_indices, - kv_buffer=self.v_storage, - ) + if src_indices.numel(): + self.ops.copy(self.k_storage, self.v_storage, src_indices, dst_indices) + + def require_zero_initialization(self): + self.zero_initialize_cache = True + + def check_execution(self, caps, execution) -> None: + if caps.platform not in (Platform.CUDA, Platform.CPU): + raise UnsupportedAttention("GQA state operations are unavailable on this platform") + if caps.platform is not Platform.CUDA and execution.transfer_paths: + raise UnsupportedAttention("GQA CPU/DP/PD transfer operations currently require CUDA") # ======================================================================== # 三、CPU Cache 数据交换 diff --git a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py index 6629d54aab..376c3f4501 100644 --- a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py @@ -1,20 +1,20 @@ from __future__ import annotations -from functools import lru_cache -from typing import TYPE_CHECKING, List from dataclasses import dataclass +from functools import lru_cache from math import prod +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from .group_quantized_transfer import transfer_group_quantized_pages from ..dp_transfer import copy_group_quantized_tokens +from .group_quantized_transfer import transfer_group_quantized_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -81,6 +81,8 @@ def __post_init__(self) -> None: class GroupQuantizedFullAttentionState(FullAttentionState): """管理 INT4/INT8 KV 与逐 token 分组 scale，复制和传输始终覆盖二者。""" + storage_layout = ("gqa_group_quantized_token_head_dim", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -150,6 +152,12 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + from lightllm.common.attention.api.inputs import DenseResources + + (key, key_scale), (value, value_scale) = self.get_layer_view(layer_id) + return DenseResources(key, value, key_scale=key_scale, value_scale=value_scale) + def get_layer_view( self, logical_layer_id: int ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: diff --git a/lightllm/common/attention_state/full_att/gqa/ops.py b/lightllm/common/attention_state/full_att/gqa/ops.py new file mode 100644 index 0000000000..5039de5aef --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/ops.py @@ -0,0 +1,44 @@ +"""GQA 的物理写入/复制实现；层映射、容量和 token 生命周期仍由 State 管理。""" + +import torch + +from lightllm.platform import Platform + + +class TritonKVOperations: + @staticmethod + def write(key, value, indices, kv): + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + heads = key.shape[1] + destindex_copy_kv(K=kv[:, :heads], DestLoc=indices, Out=key) + destindex_copy_kv(K=kv[:, heads:], DestLoc=indices, Out=value) + + @staticmethod + def copy(key, value, src, dst): + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + copy_kv_buffer_to_kv_buffer(src_mem_index=src, dst_mem_index=dst, kv_buffer=key) + copy_kv_buffer_to_kv_buffer(src_mem_index=src, dst_mem_index=dst, kv_buffer=value) + + +class TorchKVOperations: + @staticmethod + def write(key, value, indices, kv): + indices = indices.to(torch.int64) + key.index_copy_(0, indices, kv[:, : key.shape[1]]) + value.index_copy_(0, indices, kv[:, key.shape[1] :]) + + @staticmethod + def copy(key, value, src, dst): + src, dst = src.to(torch.int64), dst.to(torch.int64) + key.index_copy_(1, dst, key.index_select(1, src)) + value.index_copy_(1, dst, value.index_select(1, src)) + + +def get_kv_operations(runtime): + if runtime.platform is Platform.CUDA: + return TritonKVOperations + if runtime.platform is Platform.CPU: + return TorchKVOperations + raise ValueError(f"GQA KV writes are not implemented on {runtime.platform.value}") diff --git a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py index 5cb1acd778..118c044a87 100644 --- a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py @@ -1,22 +1,21 @@ from __future__ import annotations -from functools import lru_cache - -from typing import TYPE_CHECKING, List from collections.abc import Mapping from dataclasses import dataclass +from functools import lru_cache from math import prod +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result from ..base_state import FullAttentionState -from .static_fp8_transfer import transfer_static_fp8_pages from ..dp_transfer import copy_gqa_tokens +from .static_fp8_transfer import transfer_static_fp8_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -72,6 +71,8 @@ def __post_init__(self) -> None: class StaticFP8FullAttentionState(FullAttentionState): """使用静态标定参数的 FP8 KV 存储，固定 scale 随 State 初始化，不逐 token 搬运。""" + storage_layout = ("gqa_static_fp8_token_head_dim", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -226,6 +227,20 @@ def get_v_descale(self, storage_layer_index: int, batch_size: int) -> torch.Tens head_num = self.config.head_num return self.scales[storage_layer_index, head_num:].view(1, head_num).expand(batch_size, head_num) + def layer_resources(self, layer_id): + from lightllm.common.attention.api.inputs import DenseResources + + key, value = self.get_layer_view(layer_id) + index = self.get_storage_layer_index(layer_id) + if self.config.scale_mode == "per_head": + heads = self.config.head_num + key_scale, value_scale = self.scales[index, :heads], self.scales[index, heads:] + query_scale = self.q_scales[index] if self.q_scales is not None else None + else: + key_scale, value_scale = map(float, self.cpu_scales[index]) + query_scale = None + return DenseResources(key, value, query_scale=query_scale, key_scale=key_scale, value_scale=value_scale) + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) diff --git a/lightllm/common/attention_state/full_att/mla/mla_state.py b/lightllm/common/attention_state/full_att/mla/mla_state.py index a818f5af3a..e099289dd7 100644 --- a/lightllm/common/attention_state/full_att/mla/mla_state.py +++ b/lightllm/common/attention_state/full_att/mla/mla_state.py @@ -1,19 +1,19 @@ from __future__ import annotations +from dataclasses import dataclass from functools import lru_cache -from typing import TYPE_CHECKING, List from math import prod -from dataclasses import dataclass +from typing import TYPE_CHECKING, List import torch from lightllm.utils.envs_utils import get_env_start_args -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from ..packed_mla.packed_transfer import transfer_packed_pages from ..dp_transfer import copy_packed_tokens +from ..packed_mla.packed_transfer import transfer_packed_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -51,6 +51,8 @@ def __post_init__(self) -> None: class MLAFullAttentionState(FullAttentionState): """存储复制到各 TP rank 的 MLA 压缩 latent，使用单个 KV 字段。""" + storage_layout = ("mla_token_latent", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -63,6 +65,7 @@ def __init__( ) -> None: """按 MLA 布局配置准备压缩 KV 存储。""" self.config = layout_config + self.storage_dtype = layout_config.latent_dtype self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { @@ -110,6 +113,9 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + return self.get_layer_view(layer_id) + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: """按逻辑层号返回压缩 latent 视图，不拆分为 GQA 的 K/V head。""" return self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] diff --git a/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py index 4f018ef215..26a95bd8db 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py @@ -2,21 +2,20 @@ from __future__ import annotations -from functools import lru_cache - from dataclasses import dataclass +from functools import lru_cache from typing import ClassVar import torch from lightllm.utils.envs_utils import get_env_start_args -from ...utils import deepcopy_result from ...types import AxisMapping, StateTransferLayout, TransferRegion +from ...utils import deepcopy_result from ..base_state import FullAttentionState from ..dp_transfer import copy_packed_tokens -from ..packed_mla.packed_transfer import transfer_packed_pages from ..packed_mla.packed_mla_state import PackedMLALayerView +from ..packed_mla.packed_transfer import transfer_packed_pages @dataclass(frozen=True) @@ -58,6 +57,8 @@ def __post_init__(self) -> None: class FP8E8M0PackedMLAState(FullAttentionState): """独立的 group64/E8M0 MLA 实现，不继承其他具体 State 或动态选择编码。""" + storage_layout = ("mla_packed_fp8_e8m0_token_row", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -70,6 +71,7 @@ def __init__( ) -> None: """只准备布局和逻辑层映射，实际存储由管理器确定 token 容量后统一分配。""" self.config = layout_config + self.storage_dtype = torch.uint8 self.state_name: str | None = None self.page_size = page_size self.device = torch.device(device=device) @@ -90,7 +92,7 @@ def allocate_storage(self, token_capacity: int) -> None: self.capacity = token_capacity self.storage = torch.empty( size=(len(self.config.logical_layer_groups), token_capacity + self.page_size, 1, self.config.token_bytes), - dtype=torch.uint8, + dtype=self.storage_dtype, device=self.device, ) @@ -110,6 +112,9 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: """将逻辑层映射为物理存储行；同组逻辑层返回相同下标。""" return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + return self.get_layer_view(layer_id).kv + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: """返回固定 584 字节 KV 和 indexer key/scale 视图，保留物理行 stride。""" row = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py index 99c920fe26..055cd86f35 100644 --- a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py @@ -1,20 +1,20 @@ from __future__ import annotations -from functools import lru_cache -from typing import TYPE_CHECKING, List, Literal from dataclasses import dataclass +from functools import lru_cache from math import prod +from typing import TYPE_CHECKING, List, Literal import torch from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from ...utils import deepcopy_result, dtype_name from ...types import AxisMapping, TransferRegion +from ...utils import deepcopy_result, dtype_name from ..base_state import FullAttentionState -from .packed_mla_transfer import transfer_packed_mla_pages from ..dp_transfer import copy_packed_tokens +from .packed_mla_transfer import transfer_packed_mla_pages if TYPE_CHECKING: from lightllm.common.attention_state.types import StateTransferLayout @@ -146,6 +146,8 @@ class PackedMLALayerView: class PackedMLAFullAttentionState(FullAttentionState): """在同一物理 token 行中存储 MLA latent、量化参数及稀疏 indexer key/scale。""" + storage_layout = ("mla_packed_token_row", 1) + # ======================================================================== # 一、资源定义与初始化 # ======================================================================== @@ -158,6 +160,7 @@ def __init__( ) -> None: """按 packed MLA 配置准备 KV 与 indexer 存储。""" self.config = layout_config + self.storage_dtype = torch.uint8 if layout_config.kv_encoding == "fp8_group128" else layout_config.dtype self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 self._logical_to_storage_layer = { @@ -183,15 +186,14 @@ def allocate_storage(self, token_capacity: int) -> None: """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" self.capacity = token_capacity # FP8 行混合保存多种字段，用字节 buffer 承载，字段视图再解释各自的 dtype。 - storage_dtype = torch.uint8 if self.config.kv_encoding == "fp8_group128" else self.config.dtype self.storage = torch.empty( size=( len(self.config.logical_layer_groups), token_capacity + self.page_size, 1, - self.config.token_bytes // storage_dtype.itemsize, + self.config.token_bytes // self.storage_dtype.itemsize, ), - dtype=storage_dtype, + dtype=self.storage_dtype, device=self.device, ) @@ -211,6 +213,9 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" return self._logical_to_storage_layer[logical_layer_id] + def layer_resources(self, layer_id): + return self.get_layer_view(layer_id).kv + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: """分别返回 KV 与 packed indexer key/scale 的视图，保持原行步长。""" row = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index c0a85c367e..a72b245050 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -1,27 +1,27 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import json +import weakref from functools import lru_cache -from typing import List +from typing import TYPE_CHECKING, List import torch from lightllm.common.attention_state.types import PageKind - -from lightllm.utils.dist_utils import get_current_device_id -from lightllm.utils.envs_utils import get_env_start_args +from lightllm.platform import get_device_runtime from lightllm.utils.log_utils import init_logger -from .allocators import StatePageAllocator, FullAttKVAllocator +from .allocators import FullAttKVAllocator, StatePageAllocator +from .bounded_att import BoundedAttentionState from .cpu_cache_operator import CpuCacheOperator from .dp_shared_operator import DPSharedOperator from .full_att import FullAttentionState -from .bounded_att import BoundedAttentionState from .pd_operator import PDOperator from .utils import MODEL_ROLES if TYPE_CHECKING: from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs @@ -39,6 +39,8 @@ class AttentionStateManager: 和 allocator，不复制 State 的算法逻辑。 """ + storage_generation = 0 + # ======================================================================== # 一、基础管理：共享分配器、资源初始化与聚合查询 # ======================================================================== @@ -138,7 +140,6 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: 排序，State 内保留 logical_layer_groups 顺序。非静态 FP8 的 full State 同样占用 对应物理行区间，共享逻辑层只占一行；CPU 布局创建不需要调用本流程。 """ - import json from .full_att.gqa.static_fp8_state import StaticFP8FullAttentionState full_states = self.get_full_attention_states() @@ -180,6 +181,16 @@ def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: "Full-attention calibration bound: states=%d, source=%s", len(calibration_states), path or "unit scales" ) + def register_graph(self, graph): + if not hasattr(self, "_graphs"): + self._graphs = weakref.WeakSet() + self._graphs.add(graph) + + def close_graphs(self): + """释放状态前先排空依赖它的捕获实例，避免旧地址在后台继续被使用。""" + for graph in tuple(getattr(self, "_graphs", ())): + graph.close() + def initialize_storage( self, args: StartArgs | EasyDict, @@ -194,6 +205,7 @@ def initialize_storage( 依次执行各模型的运行时初始化，由模型读取 allocator 中的最终容量。 """ assert self._state_registration_finished, "State registration must finish before resource initialization" + assert self.full_att_allocator is None, "resources already initialized" # 阶段一：规划共享 token 容量，并向下对齐到完整 KV 页。 # 未指定容量时，根据剩余显存评估；固定运行态的显存已经在前面占用。 @@ -204,9 +216,10 @@ def initialize_storage( pd_reserved_bytes = self.get_pd_operator().get_page_bytes() * args.pd_kv_page_num if token_capacity is None: # 权重和固定运行态已分配，用当前设备剩余显存扣除保留比例及 PD buffer。 - torch.cuda.empty_cache() - gpu_device = torch.device("cuda", get_current_device_id()) - free_bytes, total_bytes = torch.cuda.mem_get_info(gpu_device) + runtime = getattr(self, "device_runtime", None) or get_device_runtime(device) + runtime.empty_cache() + gpu_device = runtime.device + free_bytes, total_bytes = runtime.memory_info() budget = int(free_bytes - total_bytes * (1 - mem_fraction)) - pd_reserved_bytes # 按每页开销直接计算容量，预留一个 HOLD 页；静态量化参数占比较小，不单独扣减。 @@ -229,6 +242,8 @@ def initialize_storage( token_capacity = token_capacity // args.page_size * args.page_size assert token_capacity > 0, "attention capacity must contain at least one complete page" + self.close_graphs() + self.storage_generation += 1 # 阶段二：管理器为已创建的 State 分配实际存储与共享 allocator。 # State 对象保持不变，查询缓存继续引用相同对象；固定运行态不重复分配。 big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num @@ -331,6 +346,18 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten # ======================================================================== # 空前缀初始化：只处理 bounded attention；无此类 State 时可为空操作。 + def snapshot_capture_state(self, req_indices: torch.Tensor, role: str): + states = self.get_bounded_attention_states(role) + if not states: + return () + requests = req_indices.unique().cpu().tolist() + return tuple((view, view.clone()) for state in states for view in state.capture_snapshot_views(requests)) + + @staticmethod + def restore_capture_state(snapshot): + for target, saved in snapshot: + target.copy_(saved) + def reset_bounded_att_state(self, req_indices: List[int]) -> None: """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页。""" for state in self.get_bounded_attention_states(): diff --git a/lightllm/common/basemodel/attention.py b/lightllm/common/basemodel/attention.py new file mode 100644 index 0000000000..313ea5816a --- /dev/null +++ b/lightllm/common/basemodel/attention.py @@ -0,0 +1,197 @@ +from lightllm.common.attention.api import AttentionBatch, AttentionPhase, ContextCapacity, LayerAttention +from lightllm.common.attention.api.environment import BackendEnv +from lightllm.common.attention.api.execution import ServingConstraints +from lightllm.common.attention.api.workspace import AttentionWorkspace + + +def describe_serving_constraints(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 ServingConstraints( + graph_phases=frozenset(phases), + microbatch_count=2 if args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap else 1, + max_batch_size=getattr(model, "max_req_num", 1), + max_seq_len=max(getattr(model, "max_seq_length", 1), getattr(model, "graph_max_len_in_batch", 1)), + expert_parallel=args.enable_ep_moe, + transfer_paths=frozenset(transfers), + ) + + +def _alloc_tensor(shape, dtype, device="cuda"): + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + return g_cache_manager.alloc_tensor(shape, dtype, device=device) + + +def backend_env(model): + return BackendEnv( + runtime=model.device_runtime, + network_config=model.config, + hold_req_id=model.req_manager.HOLD_REQUEST_ID, + draft_step=model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model), + max_seq_len=model.max_seq_length, + tp_size=model.tp_world_size_, + state_slots=model.args.mtp_step + 1, + tp_rank=model.tp_rank_, + layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + alloc_tensor=_alloc_tensor, + prefill_graph=model.args.enable_prefill_cudagraph, + max_prefill_tokens=model.batch_max_tokens or 0, + decode_only=model.args.run_mode == "decode", + ) + + +def build_attention_batch(state): + return AttentionBatch( + phase=AttentionPhase.PREFILL if state.is_prefill else AttentionPhase.DECODE, + batch_size=state.batch_size, + token_num=state.input_ids.shape[0], + total_token_num=state.total_token_num, + max_q_seq_len=state.max_q_seq_len, + max_kv_seq_len=state.max_kv_seq_len, + max_cache_len=state.max_cache_len or 0, + b_req_idx=state.b_req_idx, + b_seq_len=state.b_seq_len, + b_q_seq_len=state.b_q_seq_len, + b1_cu_q_seq_len=state.b1_cu_q_seq_len, + b1_cu_kv_seq_len=state.b1_cu_kv_seq_len, + req_to_token_indexs=state.req_manager.req_to_token_indexs, + hold_req_id=state.req_manager.HOLD_REQUEST_ID, + b_ready_cache_len=state.b_ready_cache_len, + b_mtp_index=state.b_mtp_index, + b_shared_seq_len=state.b_shared_seq_len, + b_shared_radix_node_id=state.b_shared_radix_node_id, + mem_index=state.mem_index, + position_ids=state.position_ids, + b_image_token_end=getattr(state, "b_image_token_end", None), + ) + + +class ModelAttention: + def __init__(self, model): + from lightllm.common.attention.selection import select_attention_backend + + self.groups = () + self.routes = () + self.adapters = () + self.workspaces = {} + self.generation = None + self.runtime = model.device_runtime + self.constraints = describe_serving_constraints(model) + if model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): + self.routes = (None,) * len(model.layers_infer) + return + caps = self.runtime.caps + for state in ( + *model.att_state_manager.get_full_attention_states(model.att_state_role), + *model.att_state_manager.get_bounded_attention_states(model.att_state_role), + ): + state.check_execution(caps, self.constraints) + groups, keys, routes, adapters = [], {}, [], [] + env = backend_env(model) + for layer in model.layers_infer: + specs = [] + adapter = layer.attention_adapter(layer.describe_attention(AttentionPhase.PREFILL).family) + adapters.append(adapter) + for phase in AttentionPhase: + description = layer.describe_attention(phase) + specs.append(adapter.describe(model, layer, phase, description)) + specs = tuple(specs) + if specs[0].cache != specs[1].cache: + raise ValueError("prefill/decode physical cache layouts differ") + if specs not in keys: + group = {} + for spec in specs: + option = f"llm_{spec.phase.value}_{'linear_' if spec.family.is_recurrent else ''}att_backend" + cls = select_attention_backend( + caps=caps, + execution=self.constraints, + requested=getattr(model.args, option), + policy=model.attention_policy(spec.family, spec.phase), + specs=(spec,), + ) + group[spec.phase] = cls(env, spec) + keys[specs] = len(groups) + groups.append(group) + routes.append(keys[specs]) + self.groups = tuple(groups) + self.routes = tuple(routes) + self.adapters = tuple(adapters) + + def reserve(self, model): + for group_id, group in enumerate(self.groups): + for phase, backend in group.items(): + capacity = ContextCapacity( + phase=phase, + batch_size=self.constraints.max_batch_size, + token_num=self.constraints.max_batch_size, + max_kv_seq_len=self.constraints.max_seq_len, + max_q_seq_len=1, + ) + budget = backend.resource_requirements(capacity, capture=phase in self.constraints.graph_phases) + for slot in range(self.constraints.microbatch_count): + self.workspace(group_id, slot).reserve(budget) + if backend.zero_initialize_cache: + for layer, route in zip(model.layers_infer, self.routes): + if route == group_id: + for state in model.att_state_manager.get_full_attention_states(model.att_state_role): + if state.contains_logical_layer(layer.layer_num_): + state.require_zero_initialization() + + def bind(self, model): + for layer, weight, group_id, adapter in zip( + model.layers_infer, model.trans_layers_weight, self.routes, self.adapters + ): + if group_id is None: + continue + group = self.groups[group_id] + resources = layer.bind_attention_resources(model, weight, adapter) + for phase, backend in group.items(): + setattr( + layer, + f"{phase.value}_attention", + LayerAttention( + group_id=group_id, + logical_layer_id=layer.layer_num_, + resources=resources, + backend=backend, + ), + ) + self.generation = model.att_state_manager.storage_generation + + def workspace(self, group_id, slot): + key = group_id, slot + if key not in self.workspaces: + self.workspaces[key] = AttentionWorkspace(runtime=self.runtime) + return self.workspaces[key] + + def session(self, state, *, max_kv_seq_len=None): + if self.generation != state.att_state_manager.storage_generation: + raise RuntimeError("attention storage changed; rebind layers and recapture CUDA graphs") + batch = build_attention_batch(state) + capacity = ContextCapacity( + phase=batch.phase, + batch_size=batch.batch_size, + token_num=batch.token_num, + max_kv_seq_len=batch.max_kv_seq_len if max_kv_seq_len is None else max_kv_seq_len, + max_q_seq_len=batch.max_q_seq_len, + ) + return state.attention.prepare( + batch, + self.groups, + capacity, + workspace=lambda group_id: self.workspace(group_id, state.microbatch_index), + capture=state.is_cuda_graph_capture, + replay=state.is_cuda_graph_replay, + ) diff --git a/lightllm/common/basemodel/attention/__init__.py b/lightllm/common/basemodel/attention/__init__.py deleted file mode 100644 index bbe6f7b7dd..0000000000 --- a/lightllm/common/basemodel/attention/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -from .base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from .create_utils import ( - get_prefill_att_backend_class, - get_decode_att_backend_class, - get_mla_prefill_att_backend_class, - get_mla_decode_att_backend_class, - get_nsa_prefill_att_backend_class, - get_nsa_decode_att_backend_class, - get_neo_prefill_att_backend_class, -) diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py deleted file mode 100644 index ae37385e68..0000000000 --- a/lightllm/common/basemodel/attention/base_att.py +++ /dev/null @@ -1,183 +0,0 @@ -import threading -import weakref - -import torch -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import Optional, TYPE_CHECKING, Tuple, Dict - -from lightllm.utils.dist_utils import get_current_device_id -from lightllm.utils.envs_utils import get_env_start_args - -if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class BaseAttBackend: - """ - 用于创建支持各种不同的AttBackend, 如 fa3, flashinfer, triton 实现等。 - 每个 model 复用一个 backend 实例。 - """ - - requires_cpu_cache_sync = False - - _instances = {} - _workspace_buffers = {} - _workspace_buffer_lock = threading.Lock() - - def __new__(cls, *args, **kwargs): - """ - Main 和 speculative draft model 可能使用不同的 CUDA graph 上限 - 和缓存布局，不能只按 backend class 共享实例。 - """ - model = kwargs.get("model", args[0] if args else None) - instance_key = (cls, model) - if instance_key not in cls._instances: - instance = super().__new__(cls) - cls._instances[instance_key] = instance - return cls._instances[instance_key] - - def __init__(self, model: "TpPartBaseModel"): - self.model = model - - @staticmethod - def get_gpu_workspace_buffer(key_name: str, workspace_size: int, dtype: torch.dtype = torch.int8) -> torch.Tensor: - """Return a process-local workspace shared by key name and CUDA device.""" - if not key_name: - raise ValueError("workspace key_name must not be empty") - if workspace_size <= 0: - raise ValueError(f"workspace_size must be positive, got {workspace_size}") - - device_id = get_current_device_id() - buffer_key = (device_id, key_name, workspace_size, dtype) - with BaseAttBackend._workspace_buffer_lock: - workspace_buffer = BaseAttBackend._workspace_buffers.get(buffer_key) - if workspace_buffer is None: - workspace_buffer = torch.empty(workspace_size, dtype=dtype, device=device_id) - BaseAttBackend._workspace_buffers[buffer_key] = workspace_buffer - return workspace_buffer - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "BasePrefillAttState": - raise NotImplementedError("not impl") - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "BaseDecodeAttState": - raise NotImplementedError("not impl") - - def uses_dynamic_spec_verify_layout(self) -> bool: - args = get_env_start_args() - draft_step = self.model.mtp_manager.get_decode_draft_step(self.model.is_mtp_draft_model) - is_main_model = not self.model.is_mtp_draft_model - has_decode_draft_step = draft_step > 0 - dynamic_verify_enabled = args.mtp_dynamic_verify - return is_main_model and has_decode_draft_step and dynamic_verify_enabled - - def uses_causal_attention(self) -> bool: - args = get_env_start_args() - is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash", "dflash2") - return not is_parallel_block_draft - - -@dataclass -class AttControl: - """ - prefill_att 和 decode_att 的入参，用于控制att backend 内部的行为, 选择正确的att 实现。 - """ - - use_alibi: bool = False - tp_alibi: torch.Tensor = None - use_sliding_window: bool = False - sliding_window: Tuple[int, int] = (-1, -1) - use_att_sink: bool = False - sink_weight: torch.Tensor = None - # mla 专用传参项 - mla_prefill: bool = False - mla_prefill_dict: Dict = None - mla_decode: bool = False - mla_decode_dict: Dict = None - # nsa (native sparse attention) 专用传参项 - nsa_prefill: bool = False - nsa_prefill_dict: Dict = None - nsa_decode: bool = False - nsa_decode_dict: Dict = None - # linear attention 专用传参项 - linear_att_prefill: bool = False - linear_att_prefill_dict: Dict = None - linear_att_decode: bool = False - linear_att_decode_dict: Dict = None - # 静态 FP8 标定按此逻辑层号查询 State 的物理行；其他后端可不使用。 - logical_layer_id: Optional[int] = None - - -@dataclass -class BaseAttState(ABC): - """Attention metadata with a weak reference to its owning InferStateInfo.""" - - backend: BaseAttBackend = None - # InferStateInfo 已强引用本状态，这里反向使用弱引用，避免形成强引用环。 - # 否则状态及其 Tensor 需等待循环 GC 回收，可能延迟显存复用、增加显存峰值。 - # InferStateInfo 的生命周期由推理调用方或 CUDA Graph 缓存管理，不由 attention state 延长。 - _infer_state_weak_ref: Optional[weakref.ReferenceType["InferStateInfo"]] = None - - @property - def infer_state(self) -> Optional["InferStateInfo"]: - if self._infer_state_weak_ref is None: - return None - infer_state = self._infer_state_weak_ref() - if infer_state is None: - raise ReferenceError("attention state's InferStateInfo has been released") - return infer_state - - def check_no_infer_state(self) -> None: - """Reject InferStateInfo stored directly in any instance attribute. - - This shallow check does not inspect containers or callback closures. - """ - # InferStateInfo imports attention states, so import it only when checking. - from lightllm.common.basemodel.infer_struct import InferStateInfo - - for name, value in vars(self).items(): - if isinstance(value, InferStateInfo): - raise RuntimeError( - f"{type(self).__name__}.{name} must not retain InferStateInfo; use _infer_state_weak_ref instead" - ) - - @abstractmethod - def init_state(self) -> None: - pass - - -@dataclass -class BasePrefillAttState(BaseAttState): - @abstractmethod - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - raise NotImplementedError("not impl") - - -@dataclass -class BaseDecodeAttState(BaseAttState): - def copy_for_decode_cuda_graph(self, new_state: "BaseDecodeAttState"): - """Copy tensor metadata into the buffers owned by the captured graph.""" - for attr_name, attr_value in vars(new_state).items(): - if isinstance(attr_value, torch.Tensor): - attr_ = getattr(self, attr_name, None) - if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr(): - attr_.copy_(attr_value, non_blocking=True) - - @abstractmethod - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - pass diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py deleted file mode 100644 index 08528f98f6..0000000000 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ /dev/null @@ -1,79 +0,0 @@ -from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend -from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend -from lightllm.common.basemodel.attention.linear.replay_gdn import ( - ReplayGDNFlashQlaAttBackend, - ReplayGDNTritonAttBackend, -) -from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.utils.backend_validator import validate -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - -gdn_prefill_backends = { - "flashqla": FlashQlaLinearAttBackend, - "triton": TritonLinearAttBackend, -} - -gdn_replay_prefill_backends = { - "flashqla": ReplayGDNFlashQlaAttBackend, - "triton": ReplayGDNTritonAttBackend, -} - -gdn_decode_backends = { - "triton": TritonLinearAttBackend, -} - -gdn_replay_decode_backends = { - "triton": ReplayGDNTritonAttBackend, -} - -kda_backends = { - "triton": KDALinearAttBackend, -} - -kda_replay_backends = { - "triton": ReplayKDALinearAttBackend, -} - - -def get_gdn_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): - args = get_env_start_args() - classes = gdn_replay_prefill_backends if args.ssm_state_mode == "replay" else gdn_prefill_backends - return _select_backend(classes, args.llm_prefill_att_backend, index, priority_list, "prefill") - - -def get_gdn_linear_decode_att_backend_class(index=1, priority_list=("triton",)): - args = get_env_start_args() - classes = gdn_replay_decode_backends if args.ssm_state_mode == "replay" else gdn_decode_backends - return _select_backend(classes, args.llm_decode_att_backend, index, priority_list, "decode") - - -def get_kda_linear_att_backend_class(): - args = get_env_start_args() - classes = kda_replay_backends if args.ssm_state_mode == "replay" else kda_backends - return classes["triton"] - - -def _select_backend(classes, backend_args, index, priority_list, stage): - backend_str = _get_backend_str(backend_args, index) - if backend_str != "auto": - return classes[backend_str] - - for backend_name in priority_list: - if backend_name == "triton" or validate(backend_name): - if backend_name == "triton": - logger.info(f"Linear {stage} attention backend: triton.") - else: - logger.info(f"Linear {stage} attention backend: {backend_name} (validated).") - return classes[backend_name] - - logger.warning(f"No linear {stage} attention backend validation succeeded, falling back to Triton.") - return classes["triton"] - - -def _get_backend_str(backend_args, index): - assert index == 1, "linear attention backend index must be 1" - return backend_args[index] if len(backend_args) > index else "auto" diff --git a/lightllm/common/basemodel/attention/create_utils.py b/lightllm/common/basemodel/attention/create_utils.py deleted file mode 100644 index a4e2641a9f..0000000000 --- a/lightllm/common/basemodel/attention/create_utils.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Attention backend selection utilities.""" - -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.log_utils import init_logger -from lightllm.utils.backend_validator import validate -from .base_att import BaseAttBackend -from .import_utils import load_backend_class - -logger = init_logger(__name__) - -# Lazy backend class references by data type. -data_type_to_backend = { - "None": { - "triton": "triton.fp:TritonAttBackend", - "fa3": "fa3.fp:Fa3AttBackend", - "flashinfer": "flashinfer.fp:FlashInferAttBackend", - }, - "int4kv": { - "triton": "triton.int4kv:Int4kvTritonAttBackend", - }, - "int8kv": { - "triton": "triton.int8kv:Int8kvTritonAttBackend", - }, - "fp8kv_sph": { - "fa3": "fa3.fp8:Fp8Fa3AttBackend", - }, - "fp8kv_spt": { - "flashinfer": "flashinfer.fp8:Fp8FlashInferAttBackend", - }, -} - -mla_data_type_to_backend = { - "None": { - "triton": "triton.mla:MlaTritonAttBackend", - "fa3": "fa3.mla:MlaFa3AttBackend", - "flashinfer": "flashinfer.mla:MlaFlashInferAttBackend", - }, -} - -nsa_data_type_to_backend = { - "None": { - "flashmla_sparse": "nsa.flashmla_sparse:NsaFlashMlaSparseAttBackend", - # Future backends: "fa3", "tilelang", "aiter" - }, - "fp8kv_dsa": { - "flashmla_sparse": "nsa.fp8_flashmla_sparse:NsaFlashMlaFp8SparseAttBackend", - }, -} - -neo_data_type_to_backend = { - "None": { - "triton": "triton.neo:NeoTritonAttBackend", - "fa3": "fa3.neo:NeoFa3AttBackend", - } -} - - -def _auto_select_backend( - llm_dtype: str, - kv_type_to_backend: dict, - priority_list: list = ["fa3", "flashinfer", "triton"], -) -> type: - """Auto-select the best available backend with validation. - - Priority follows the provided priority_list. - Each backend is validated in a subprocess with ground truth checks. - """ - backend_map = kv_type_to_backend[llm_dtype] - - args = get_env_start_args() - if args.enable_ep_moe: - logger.info("Expert parallelism with MoE enabled, excluding flashinfer attention backend") - priority_list = [name for name in priority_list if name != "flashinfer"] - - for backend_name in priority_list: - if backend_name in backend_map and validate(backend_name): - logger.info(f"Auto-selected {backend_name} backend (validated)") - return load_backend_class(backend_map[backend_name]) - - # Fallback to triton without validation (should not happen) - logger.warning("No backend validation succeeded, falling back to triton") - return load_backend_class(backend_map["triton"]) - - -def _get_decode_backend_priority(priority_list: list, mtp_step: int) -> list: - """Return the auto-selection priority for decode attention.""" - # With MTP, FA3 can make better use of Tensor Core compute and delivers better decode performance. - if mtp_step <= 0 or "fa3" not in priority_list: - return priority_list - return ["fa3"] + [backend_name for backend_name in priority_list if backend_name != "fa3"] - - -def get_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "flashinfer", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - if backend_str != "auto": - return load_backend_class(data_type_to_backend[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list) - - -def get_decode_att_backend_class(index=0, priority_list: list = ["flashinfer", "fa3", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_decode_att_backend[index] - if backend_str != "auto": - return load_backend_class(data_type_to_backend[llm_dtype][backend_str]) - priority_list = _get_decode_backend_priority(priority_list, args.mtp_step) - return _auto_select_backend(llm_dtype, kv_type_to_backend=data_type_to_backend, priority_list=priority_list) - - -def get_mla_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "flashinfer", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - if backend_str != "auto": - return load_backend_class(mla_data_type_to_backend[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list) - - -def get_mla_decode_att_backend_class(index=0, priority_list: list = ["flashinfer", "fa3", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_decode_att_backend[index] - if backend_str != "auto": - return load_backend_class(mla_data_type_to_backend[llm_dtype][backend_str]) - priority_list = _get_decode_backend_priority(priority_list, args.mtp_step) - return _auto_select_backend(llm_dtype, kv_type_to_backend=mla_data_type_to_backend, priority_list=priority_list) - - -def get_nsa_prefill_att_backend_class( - index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend -) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - if backend_str != "auto": - return load_backend_class(backend_map[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list) - - -def get_nsa_decode_att_backend_class( - index=0, priority_list: list = ["flashmla_sparse"], backend_map=nsa_data_type_to_backend -) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_decode_att_backend[index] - if backend_str != "auto": - return load_backend_class(backend_map[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list) - - -def get_neo_prefill_att_backend_class(index=0, priority_list: list = ["fa3", "triton"]) -> BaseAttBackend: - args = get_env_start_args() - llm_dtype = args.llm_kv_type - backend_str = args.llm_prefill_att_backend[index] - backend_map = neo_data_type_to_backend - if backend_str == "fa3" or (backend_str == "auto" and "fa3" in priority_list): - # Only Neo FA3 selection needs the optional image-token attention interface. - from .fa3.neo import HAS_FLASH_ATTN_INTERFACE - - if not HAS_FLASH_ATTN_INTERFACE: - backend_map = {"None": {"triton": neo_data_type_to_backend["None"]["triton"]}} - if backend_str != "auto": - return load_backend_class(backend_map[llm_dtype][backend_str]) - return _auto_select_backend(llm_dtype, kv_type_to_backend=backend_map, priority_list=priority_list) diff --git a/lightllm/common/basemodel/attention/fa3/fp.py b/lightllm/common/basemodel/attention/fa3/fp.py deleted file mode 100644 index 32d638ee11..0000000000 --- a/lightllm/common/basemodel/attention/fa3/fp.py +++ /dev/null @@ -1,333 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, TYPE_CHECKING -from lightllm.utils.envs_utils import enable_env_vars -from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune -from lightllm.common.basemodel.triton_kernel.fa3_utils import ( - build_dynamic_spec_fa3_decode_params, - page_table_copy, -) -from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers - - -class Fa3AttBackend(BaseAttBackend): - """Common fixed page-table storage for FA3 attention backends.""" - - # 设为 1 / TRUE / ON 时，FA3 与 CPU cache 搬运在 GPU 上互相等待。 - requires_cpu_cache_sync = enable_env_vars("LIGHTLLM_FA3_CPU_CACHE_SYNC") - page_table_buffers = None - - def __init__(self, model): - super().__init__(model=model) - self._init_infer_page_size() - - def _init_infer_page_size(self): - self.infer_page_size = self.model.args.page_size - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - - # 延迟到首次获取 page table 时再初始化，避免 PD 分离模式下的 prefill 节点 - # 分配仅供 decode 使用的 buffer，减少显存浪费。 - def _init_page_table_buffers(self): - if self.page_table_buffers is not None: - return - - model = self.model - args = model.args - # Dynamic verification may keep the target row and all MTP draft rows for - # every running request. TPSP and CUDA Graph may pad that physical batch - # further, so the fixed buffer must cover both execution paths. - running_max_batch_size = args.running_max_req_size * (args.mtp_step + 1) - if args.enable_tpsp_mix_mode: - tp_size = model.tp_world_size_ - running_max_batch_size = (running_max_batch_size + tp_size - 1) // tp_size * tp_size - self.page_table_max_batch_size = max(running_max_batch_size, model.graph_max_batch_size) - # max_seq_length is max_req_total_len plus the MTP headroom reserved when - # the model is initialized. - self.supported_max_seq_len = model.max_seq_length - buffer_count = 2 if args.enable_decode_microbatch_overlap else 1 - workspace_size = self.page_table_max_batch_size * triton.cdiv(self.supported_max_seq_len, self.infer_page_size) - self.page_table_buffers = [ - self.get_gpu_workspace_buffer( - key_name=f"fa3_page_table_{buffer_index}", - workspace_size=workspace_size, - dtype=torch.int32, - ) - for buffer_index in range(buffer_count) - ] - - def get_page_table_view(self, att_batch_size, max_kv_len, microbatch_index): - """Lazily initialize storage and return a contiguous page-table view.""" - self._init_page_table_buffers() - if att_batch_size > self.page_table_max_batch_size: - raise RuntimeError( - f"FA3 attention batch size {att_batch_size} exceeds page-table capacity " - f"{self.page_table_max_batch_size}" - ) - if max_kv_len > self.supported_max_seq_len: - raise RuntimeError( - f"FA3 max KV sequence length {max_kv_len} exceeds page-table capacity " f"{self.supported_max_seq_len}" - ) - max_page_len = triton.cdiv(max_kv_len, self.infer_page_size) - return self.page_table_buffers[microbatch_index][: att_batch_size * max_page_len].reshape( - att_batch_size, max_page_len - ) - - def create_att_prefill_state(self, infer_state) -> "Fa3PrefillAttState": - return Fa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Fa3DecodeAttState": - return Fa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Fa3PrefillAttState(BasePrefillAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - page_table: torch.Tensor = None - causal: bool = None - - def init_state(self): - self.causal = self.backend.uses_causal_attention() - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - self.page_table = torch.empty( - ( - self.infer_state.batch_size, - triton.cdiv(self.infer_state.max_kv_seq_len, self.backend.infer_page_size), - ), - dtype=torch.int32, - device=self.infer_state.input_ids.device, - ) - page_table_copy( - page_table=self.page_table, - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - page_size=self.backend.infer_page_size, - ) - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.use_alibi is False - return self._nomarl_prefill_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _nomarl_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: Fa3AttBackend = self.backend # for typing - - if att_control.use_sliding_window: - window_size = att_control.sliding_window - else: - window_size = (-1, -1) - - if att_control.use_att_sink: - sink_weight: torch.Tensor = att_control.sink_weight - else: - sink_weight = None - - k_descale, v_descale = None, None # disable quantization - Lq = q.shape[-1] - sm_scale = 1.0 / (Lq ** 0.5) - o = flash_attn_with_kvcache( - q=q, - k_cache=k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v_cache=v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - page_table=self.page_table, - cache_seqlens=self.infer_state.b_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=sm_scale, - causal=self.causal, - window_size=window_size, - softcap=0.0, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=False, - sinks=sink_weight, - ) - return o - - -@dataclasses.dataclass -class Fa3DecodeAttState(BaseDecodeAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - page_table: torch.Tensor = None - b_att_seq_len: torch.Tensor = None - # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。 - decode_max_q_seq_len: int = None - decode_max_kv_seq_len: int = None - causal: bool = None - - def init_state(self): - self.backend: Fa3AttBackend = self.backend - self.causal = self.backend.uses_causal_attention() - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if self.backend.uses_dynamic_spec_verify_layout(): - b_att_req_idx = self._init_dynamic_spec_verify_state(draft_step) - elif draft_step > 0: - b_att_req_idx = self._init_fixed_spec_decode_state(draft_step) - else: - b_att_req_idx = self._init_normal_decode_state() - - self._init_page_table(b_att_req_idx) - - def _init_dynamic_spec_verify_state(self, draft_step: int) -> torch.Tensor: - b_mark_mtp_shared_group = build_mtp_shared_group_markers( - self.infer_state.b_req_idx, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - b_q_seq_len, b_kv_seq_len, b_att_req_idx, self.b_att_seq_len = build_dynamic_spec_fa3_decode_params( - b_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=self.infer_state.batch_size, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = draft_step + 1 - return b_att_req_idx - - def _init_fixed_spec_decode_state(self, draft_step: int) -> torch.Tensor: - mtp_size = draft_step + 1 - assert self.infer_state.batch_size % mtp_size == 0, ( - "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={self.infer_state.batch_size}, draft_step={draft_step}." - ) - - b_q_seq_len = torch.full( - (self.infer_state.b_seq_len.shape[0] // mtp_size,), - fill_value=mtp_size, - dtype=torch.int32, - device=self.infer_state.b_seq_len.device, - ) - b_kv_seq_len = self.infer_state.b_seq_len[draft_step::mtp_size] - b_att_req_idx = self.infer_state.b_req_idx[draft_step::mtp_size] - self.b_att_seq_len = b_kv_seq_len.contiguous() - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = mtp_size - return b_att_req_idx - - def _init_normal_decode_state(self) -> torch.Tensor: - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - self.b_att_seq_len = self.infer_state.b_seq_len - self.decode_max_q_seq_len = 1 - return self.infer_state.b_req_idx - - def _init_spec_decode_cu_seqlens(self, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor): - b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len) - self.cu_seqlens_q = b1_cu_q_seq_len.int() - self.cu_seqlens_k = b1_cu_kv_seq_len.int() - - def _init_page_table(self, b_att_req_idx: torch.Tensor): - att_batch_size = b_att_req_idx.shape[0] - model = self.backend.model - actual_max_kv_len = self.infer_state.max_kv_seq_len - # Graph 捕获会将 infer_state.max_kv_seq_len 改为容量上限，提前保存真实长度用于 FA3 配置查找。 - self.decode_max_kv_seq_len = actual_max_kv_len - actual_max_page_len = triton.cdiv(actual_max_kv_len, self.backend.infer_page_size) - page_table_width = actual_max_kv_len - if model.graph is not None and model.graph.can_run( - batch_size=self.infer_state.batch_size, - max_len_in_batch=actual_max_kv_len, - ): - # CUDA Graph replay uses the shape and strides captured with the - # graph-wide maximum KV length. Keep that fixed row stride while - # writing only the valid portion of each runtime row below. - page_table_width = model.graph.graph_max_len_in_batch - - self.page_table = self.backend.get_page_table_view( - att_batch_size=att_batch_size, - max_kv_len=page_table_width, - microbatch_index=self.infer_state.microbatch_index, - ) - - page_table_copy( - page_table=self.page_table[:, :actual_max_page_len], - req_to_token_indexs=model.req_manager.req_to_token_indexs, - b_req_idx=b_att_req_idx, - page_size=self.backend.infer_page_size, - ) - - def copy_for_decode_cuda_graph(self, new_state: "Fa3DecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.use_alibi is False - return self._normal_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _normal_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - if att_control.use_sliding_window: - window_size = att_control.sliding_window - else: - window_size = (-1, -1) - - if att_control.use_att_sink: - sink_weight: torch.Tensor = att_control.sink_weight - else: - sink_weight = None - - k_descale, v_descale = None, None # disable quantization - Lq = q.shape[-1] - sm_scale = 1.0 / (Lq ** 0.5) - o = flash_attn_with_kvcache_autotune( - q=q, - k_cache=k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v_cache=v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - page_table=self.page_table, - cache_seqlens=self.b_att_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=None, # KV 已提前写入缓存，此处不追加新的 K/V。 - max_seqlen_q=self.decode_max_q_seq_len, - max_seqlen_k=self.decode_max_kv_seq_len, - softmax_scale=sm_scale, - causal=self.causal, - window_size=window_size, - softcap=0.0, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=False, - sinks=sink_weight, - ) - return o diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py deleted file mode 100644 index 8f41d74de2..0000000000 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ /dev/null @@ -1,205 +0,0 @@ -import dataclasses -import weakref -import torch -from ..base_att import AttControl -from lightllm.utils.sgl_utils import flash_attn_with_kvcache -from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import ( - q_per_head_fp8_quant, - q_per_head_static_fp8_quant, -) -from lightllm.utils.log_utils import init_logger -from .fp import Fa3AttBackend, Fa3PrefillAttState, Fa3DecodeAttState - - -logger = init_logger(__name__) - - -class Fp8Fa3AttBackend(Fa3AttBackend): - def _init_infer_page_size(self): - # TODO: FP8 FA3 完成多 token 推理页适配后，改为继承模型 page_size。 - self.infer_page_size = 1 - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - logger.warning( - f"Fp8Fa3AttBackend temporarily uses infer_page_size=1 with model page_size={self.model.args.page_size}; " - "multi-token FP8 inference pages are not implemented yet." - ) - - def create_att_prefill_state(self, infer_state) -> "Fp8Fa3PrefillAttState": - return Fp8Fa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Fp8Fa3DecodeAttState": - return Fp8Fa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Fp8Fa3PrefillAttState(Fa3PrefillAttState): - def init_state(self): - super().init_state() - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ) -> torch.Tensor: - self.backend: Fp8Fa3AttBackend = self.backend # for typing - - q_head_num = q.shape[1] - q_head_dim = q.shape[2] - k_head_num = k.shape[1] - k_head_dim = k.shape[2] - cache_k = k.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - cache_v = v.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_head", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - batch_size = self.infer_state.b_seq_len.shape[0] - q = q.reshape(q.shape[0], k_head_num, -1) - if state.q_scales is None: - # 旧标定文件仅包含 K/V；按本次请求的 Q 长度动态计算每个 KV head 的 scale。 - q, q_scale = q_per_head_fp8_quant( - q=q, - seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], - b1_start_loc=self.cu_seqlens_q, - ) - else: - static_q_scales = state.q_scales[storage_layer_index] - q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) - q_scale = static_q_scales.view(1, k_head_num).expand(batch_size, k_head_num) - o = flash_attn_with_kvcache( - q=q.reshape(-1, q_head_num, q_head_dim), - k_cache=cache_k, - v_cache=cache_v, - page_table=self.page_table, - cache_seqlens=self.infer_state.b_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - q_descale=q_scale, - k_descale=state.get_k_descale(storage_layer_index, batch_size), - v_descale=state.get_v_descale(storage_layer_index, batch_size), - return_softmax_lse=False, - ) - return o - - -@dataclasses.dataclass -class Fp8Fa3DecodeAttState(Fa3DecodeAttState): - def init_state(self): - self.backend: Fp8Fa3AttBackend = self.backend - - super().init_state() - - def copy_for_decode_cuda_graph(self, new_state: "Fp8Fa3DecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_decode_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ): - k_head_num = k.shape[1] - k_head_dim = k.shape[2] - - cache_k = k.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - cache_v = v.view( - size=(-1, 1, k_head_num, k_head_dim), - ).view(torch.float8_e4m3fn) - - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_head", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - - q_head_num = q.shape[1] - att_batch_size = self.b_att_seq_len.shape[0] - q = q.reshape(q.shape[0], k_head_num, -1) - if state.q_scales is None: - # 使用 FA3 的实际 Q 分组，兼容普通 decode、MTP 多 token 及零长度填充组。 - q, q_scale = q_per_head_fp8_quant( - q=q, - seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], - b1_start_loc=self.cu_seqlens_q, - ) - else: - static_q_scales = state.q_scales[storage_layer_index] - q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) - q_scale = static_q_scales.view(1, k_head_num).expand(att_batch_size, k_head_num) - o = flash_attn_with_kvcache( - q=q.reshape(-1, q_head_num, k_head_dim), - k_cache=cache_k, - v_cache=cache_v, - page_table=self.page_table, - cache_seqlens=self.b_att_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.decode_max_q_seq_len, - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - q_descale=q_scale.view(att_batch_size, k_head_num), - k_descale=state.get_k_descale(storage_layer_index, att_batch_size), - v_descale=state.get_v_descale(storage_layer_index, att_batch_size), - return_softmax_lse=False, - ) - return o diff --git a/lightllm/common/basemodel/attention/fa3/mla.py b/lightllm/common/basemodel/attention/fa3/mla.py deleted file mode 100644 index 244df6be12..0000000000 --- a/lightllm/common/basemodel/attention/fa3/mla.py +++ /dev/null @@ -1,244 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, TYPE_CHECKING, Tuple -from lightllm.utils.sgl_utils import flash_attn_with_kvcache -from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params, page_table_copy -from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers -from lightllm.utils.sgl_utils import flash_attn_varlen_func -from .fp import Fa3AttBackend - - -class MlaFa3AttBackend(Fa3AttBackend): - def create_att_prefill_state(self, infer_state) -> "MlaFa3PrefillAttState": - return MlaFa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "MlaFa3DecodeAttState": - return MlaFa3DecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class MlaFa3PrefillAttState(BasePrefillAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - causal: bool = None - - def init_state(self): - self.causal = self.backend.uses_causal_attention() - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._mla_prefill_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: MlaFa3AttBackend = self.backend # for typing - k_nope, k_rope = k - q_head_num = q.shape[1] - k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) - - assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3 - - assert att_control.mla_prefill - softmax_scale = att_control.mla_prefill_dict["softmax_scale"] - - o_tensor = flash_attn_varlen_func( - q=q, - k=k, - v=v, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - max_seqlen_k=self.infer_state.max_kv_seq_len, - softmax_scale=softmax_scale, - causal=self.causal, - return_softmax_lse=False, - ) - return o_tensor - - -@dataclasses.dataclass -class MlaFa3DecodeAttState(BaseDecodeAttState): - cu_seqlens_q: torch.Tensor = None - cu_seqlens_k: torch.Tensor = None - page_table: torch.Tensor = None - b_att_seq_len: torch.Tensor = None - # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。 - decode_max_q_seq_len: int = None - causal: bool = None - - def init_state(self): - self.backend: MlaFa3AttBackend = self.backend - self.causal = self.backend.uses_causal_attention() - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if self.backend.uses_dynamic_spec_verify_layout(): - b_att_req_idx = self._init_dynamic_spec_verify_state(draft_step) - elif draft_step > 0: - b_att_req_idx = self._init_fixed_spec_decode_state(draft_step) - else: - b_att_req_idx = self._init_normal_decode_state() - - self._init_page_table(b_att_req_idx) - - def _init_dynamic_spec_verify_state(self, draft_step: int) -> torch.Tensor: - b_mark_mtp_shared_group = build_mtp_shared_group_markers( - self.infer_state.b_req_idx, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - b_q_seq_len, b_kv_seq_len, b_att_req_idx, self.b_att_seq_len = build_dynamic_spec_fa3_decode_params( - b_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=self.infer_state.batch_size, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = draft_step + 1 - return b_att_req_idx - - def _init_fixed_spec_decode_state(self, draft_step: int) -> torch.Tensor: - mtp_size = draft_step + 1 - assert self.infer_state.batch_size % mtp_size == 0, ( - "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={self.infer_state.batch_size}, draft_step={draft_step}." - ) - - b_q_seq_len = torch.full( - (self.infer_state.b_seq_len.shape[0] // mtp_size,), - fill_value=mtp_size, - dtype=torch.int32, - device=self.infer_state.b_seq_len.device, - ) - b_kv_seq_len = self.infer_state.b_seq_len[draft_step::mtp_size] - b_att_req_idx = self.infer_state.b_req_idx[draft_step::mtp_size] - self.b_att_seq_len = b_kv_seq_len.contiguous() - self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len) - self.decode_max_q_seq_len = mtp_size - return b_att_req_idx - - def _init_normal_decode_state(self) -> torch.Tensor: - self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int() - self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int() - self.b_att_seq_len = self.infer_state.b_seq_len - self.decode_max_q_seq_len = 1 - return self.infer_state.b_req_idx - - def _init_spec_decode_cu_seqlens(self, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor): - b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len) - self.cu_seqlens_q = b1_cu_q_seq_len.int() - self.cu_seqlens_k = b1_cu_kv_seq_len.int() - - def _init_page_table(self, b_att_req_idx: torch.Tensor): - att_batch_size = b_att_req_idx.shape[0] - model = self.backend.model - actual_max_kv_len = self.infer_state.max_kv_seq_len - actual_max_page_len = triton.cdiv(actual_max_kv_len, self.backend.infer_page_size) - page_table_width = actual_max_kv_len - if model.graph is not None and model.graph.can_run( - batch_size=self.infer_state.batch_size, - max_len_in_batch=actual_max_kv_len, - ): - # CUDA Graph replay uses the shape and strides captured with the - # graph-wide maximum KV length. Keep that fixed row stride while - # writing only the valid portion of each runtime row below. - page_table_width = model.graph.graph_max_len_in_batch - - self.page_table = self.backend.get_page_table_view( - att_batch_size=att_batch_size, - max_kv_len=page_table_width, - microbatch_index=self.infer_state.microbatch_index, - ) - - page_table_copy( - page_table=self.page_table[:, :actual_max_page_len], - req_to_token_indexs=model.req_manager.req_to_token_indexs, - b_req_idx=b_att_req_idx, - page_size=self.backend.infer_page_size, - ) - - def copy_for_decode_cuda_graph(self, new_state: "MlaFa3DecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - assert v is None - - return self._mla_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - q_nope, q_rope = q - kv = k - qk_rope_head_dim = 64 - kv_lora_rank = kv.shape[-1] - qk_rope_head_dim - k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, self.backend.infer_page_size, 1, qk_rope_head_dim) - kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, self.backend.infer_page_size, 1, kv_lora_rank) - k_descale, v_descale = None, None - assert att_control.mla_decode - softmax_scale = att_control.mla_decode_dict["softmax_scale"] - - o_tensor = flash_attn_with_kvcache( - q=q_rope, - k_cache=k_rope, - v_cache=kv_nope, - qv=q_nope, - page_table=self.page_table, - cache_seqlens=self.b_att_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.decode_max_q_seq_len, - softmax_scale=softmax_scale, - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - k_descale=k_descale, - v_descale=v_descale, - return_softmax_lse=False, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/fa3/neo.py b/lightllm/common/basemodel/attention/fa3/neo.py deleted file mode 100644 index 4a1b1d9b3e..0000000000 --- a/lightllm/common/basemodel/attention/fa3/neo.py +++ /dev/null @@ -1,81 +0,0 @@ -import dataclasses -import weakref -import inspect -import warnings - -import torch - -from .fp import Fa3AttBackend, Fa3PrefillAttState -from ..base_att import AttControl - - -_NEO_FA3_INSTALL_HINT = ( - "Install the Neo FA3 build from source in the current Python environment:\n" - " git clone --recursive https://github.com/WANDY666/flash-attention.git\n" - " cd flash-attention/hopper\n" - " FLASH_ATTENTION_FORCE_BUILD=TRUE python -m pip install " - "--no-build-isolation --no-deps --force-reinstall .\n" - "This replaces a standalone flash_attn_3 / flash_attn_interface installation; " - "LightLLM's regular FA3 uses the separate sgl_kernel package." -) - - -try: - from flash_attn_interface import flash_attn_with_kvcache as flash_attn_with_kvcache_neo - - # Neo FA3 accepts the exclusive visible KV end for each image query. - _sig = inspect.signature(flash_attn_with_kvcache_neo) - if "image_token_end" not in _sig.parameters: - raise ImportError("flash_attn_interface is missing image_token_end support (need the Neo build)") - - HAS_FLASH_ATTN_INTERFACE = True -except ImportError as exc: - warnings.warn( - f"Neo FA3 is unavailable: {exc}. " - "Automatic Neo prefill selection will fall back to Triton. " - "To select Triton explicitly, use --llm_prefill_att_backend triton.\n" + _NEO_FA3_INSTALL_HINT - ) - flash_attn_with_kvcache_neo = None - HAS_FLASH_ATTN_INTERFACE = False - - -class NeoFa3AttBackend(Fa3AttBackend): - def create_att_prefill_state(self, infer_state) -> "NeoFa3PrefillAttState": - return NeoFa3PrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NeoFa3PrefillAttState(Fa3PrefillAttState): - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - # neo_chat*: image-token bidirectional attention requires flash_attn_interface - # (sgl_kernel's flash_attn_with_kvcache does not support image_token_end). - if not HAS_FLASH_ATTN_INTERFACE: - raise ImportError("Neo prefill requires FA3 with image_token_end support.\n" + _NEO_FA3_INSTALL_HINT) - if self.infer_state.b_image_token_end is None: - raise ValueError("Neo prefill requires b_image_token_end to describe image attention spans") - o = flash_attn_with_kvcache_neo( - q=q, - k_cache=k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v_cache=v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - page_table=self.page_table, - cache_seqlens=self.infer_state.b_seq_len, - cu_seqlens_q=self.cu_seqlens_q, - cu_seqlens_k_new=self.cu_seqlens_k, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=1.0 / (q.shape[-1] ** 0.5), - causal=self.causal, - window_size=(-1, -1), - softcap=0.0, - k_descale=None, - v_descale=None, - return_softmax_lse=False, - image_token_end=self.infer_state.b_image_token_end, - ) - return o diff --git a/lightllm/common/basemodel/attention/flashinfer/fp.py b/lightllm/common/basemodel/attention/flashinfer/fp.py deleted file mode 100644 index f9ee6c11b0..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/fp.py +++ /dev/null @@ -1,291 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id -from ...triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from ...triton_kernel.repack_kv_index import repack_kv_index -from .env_utils import set_flashinfer_envs -from .utils import should_init_decode_wrapper - - -class FlashInferAttBackend(BaseAttBackend): - workspace_buffer_key = "flashinfer_fp" - workspace_buffer_size = 512 * 1024 * 1024 - - def __init__(self, model): - set_flashinfer_envs() - super().__init__(model=model) - self._init_infer_page_size() - tp_world_size = get_dp_world_size() - self.tp_q_head_num = model.config["num_attention_heads"] // tp_world_size - self.tp_kv_head_num = max(model.config["num_key_value_heads"] // tp_world_size, 1) - head_dim = model.config["hidden_size"] // model.config["num_attention_heads"] - self.head_dim = model.config.get("head_dim", head_dim) - self.max_seq_length = model.max_seq_length - self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) - self.kv_indices_buffer = [ - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - ] - self.q_data_type = model.data_type - self.kv_data_type = model.data_type - - def _init_infer_page_size(self): - self.infer_page_size = self.model.args.page_size - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - - def create_att_prefill_state(self, infer_state) -> "FlashInferPrefillAttState": - return FlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "FlashInferDecodeAttState": - return FlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class FlashInferPrefillAttState(BasePrefillAttState): - prefill_wrapper: object = None - - def init_state(self): - self.backend: FlashInferAttBackend = self.backend - - import flashinfer - - batch_size = self.infer_state.batch_size - device = self.infer_state.input_ids.device - - q_starts = self.infer_state.b1_cu_q_seq_len.int() - # TODO: 将页数、页数前缀和及末页有效 token 数的计算融合为一个 Triton 算子。 - # token 长度除以页大小并向上取整，末页不足一页也计为一页。 - b_page_len = (self.infer_state.b_seq_len + (self.backend.infer_page_size - 1)) // self.backend.infer_page_size - kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) - kv_last_page_len = self.infer_state.b_seq_len - (b_page_len - 1) * self.backend.infer_page_size - kv_indices = torch.empty( - batch_size * self.backend.max_page_num, - dtype=torch.int32, - device=device, - ) - repack_kv_index( - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_token_len=self.infer_state.b_seq_len, - b_page_start_loc=kv_starts[:-1], - max_token_len=self.infer_state.max_kv_seq_len, - out_page_indices=kv_indices, - page_size=self.backend.infer_page_size, - ) - self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - qo_indptr_buf=q_starts, - paged_kv_indptr_buf=kv_starts, - paged_kv_indices_buf=kv_indices, - paged_kv_last_page_len_buf=kv_last_page_len, - ) - self.prefill_wrapper.plan( - q_starts, - kv_starts, - kv_indices, - kv_last_page_len, - self.backend.tp_q_head_num, - self.backend.tp_kv_head_num, - self.backend.head_dim, - self.backend.infer_page_size, - causal=True, - pos_encoding_mode="NONE", - logits_soft_cap=0.0, - q_data_type=self.backend.q_data_type, - kv_data_type=self.backend.kv_data_type, - ) - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._nomarl_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - ) - - def _nomarl_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: FlashInferAttBackend = self.backend # for typing - o_tensor = alloc_func(q.shape, q.dtype, device="cuda") - self.prefill_wrapper.run( - q, - ( - k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - ), - out=o_tensor, - ) - return o_tensor - - -@dataclasses.dataclass -class FlashInferDecodeAttState(BaseDecodeAttState): - kv_last_page_len_buffer: torch.Tensor = None - kv_indices: torch.Tensor = None - kv_starts: torch.Tensor = None - decode_wrapper: object = None - - def _should_init_decode_wrapper(self) -> bool: - return should_init_decode_wrapper(self.infer_state) - - def init_state(self): - import flashinfer - - self.backend: FlashInferAttBackend = self.backend - device = self.infer_state.input_ids.device - model = self.backend.model - # TODO: 将页数、页数前缀和及末页有效 token 数的计算融合为一个 Triton 算子。 - # token 长度除以页大小并向上取整，末页不足一页也计为一页。 - b_page_len = (self.infer_state.b_seq_len + (self.backend.infer_page_size - 1)) // self.backend.infer_page_size - self.kv_last_page_len_buffer = self.infer_state.b_seq_len - (b_page_len - 1) * self.backend.infer_page_size - if ( - self.infer_state.batch_size <= model.graph_max_batch_size - and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch - ): - self.kv_indices = self.backend.kv_indices_buffer[self.infer_state.microbatch_index][ - : self.infer_state.batch_size * self.backend.max_page_num - ] - else: - self.kv_indices = torch.empty( - self.infer_state.batch_size * self.backend.max_page_num, - dtype=torch.int32, - device=device, - ) - - self.kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) - repack_kv_index( - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_token_len=self.infer_state.b_seq_len, - b_page_start_loc=self.kv_starts[:-1], - max_token_len=self.infer_state.max_kv_seq_len, - out_page_indices=self.kv_indices, - page_size=self.backend.infer_page_size, - ) - if not self._should_init_decode_wrapper(): - # 处于 graph replay 回放阶段，不需要特殊初始化 decode wrapper。 - return - - assert self.decode_wrapper is None - self.decode_wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - "NHD", - use_cuda_graph=True, - use_tensor_cores=True, - paged_kv_indptr_buffer=self.kv_starts, - paged_kv_indices_buffer=self.kv_indices, - paged_kv_last_page_len_buffer=self.kv_last_page_len_buffer, - ) - self.decode_wrapper.plan( - self.kv_starts, - self.kv_indices, - self.kv_last_page_len_buffer, - self.backend.tp_q_head_num, - self.backend.tp_kv_head_num, - self.backend.head_dim, - self.backend.infer_page_size, - q_data_type=self.backend.q_data_type, - kv_data_type=self.backend.kv_data_type, - non_blocking=True, - ) - return - - def copy_for_decode_cuda_graph(self, new_state: "FlashInferDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - self._refresh_cuda_graph_decode_plan(new_state.infer_state.max_kv_seq_len) - return - - def _refresh_cuda_graph_decode_plan(self, max_kv_len: int): - from flashinfer.decode import fast_decode_plan - - uniform_kv_indptr_cpu = ( - torch.arange( - self.infer_state.batch_size + 1, - dtype=torch.int32, - device="cpu", - ) - * triton.cdiv(max_kv_len, self.backend.infer_page_size) - ) - - fast_decode_plan( - self.decode_wrapper, - indptr=self.kv_starts, - indices=self.kv_indices, - last_page_len=self.kv_last_page_len_buffer, - num_qo_heads=self.backend.tp_q_head_num, - num_kv_heads=self.backend.tp_kv_head_num, - head_dim=self.backend.head_dim, - page_size=self.backend.infer_page_size, - q_data_type=self.backend.q_data_type, - kv_data_type=self.backend.kv_data_type, - non_blocking=True, - global_override_indptr_cpu=uniform_kv_indptr_cpu, - ) - return - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._normal_decode_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - ) - - def _normal_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - o_tensor = alloc_func(q.shape, q.dtype) - self.decode_wrapper.run( - q, - ( - k.view(-1, self.backend.infer_page_size, k.shape[1], k.shape[2]), - v.view(-1, self.backend.infer_page_size, v.shape[1], v.shape[2]), - ), - out=o_tensor, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py deleted file mode 100644 index 5f99a2c23f..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ /dev/null @@ -1,146 +0,0 @@ -import dataclasses -import weakref -import torch -from ..base_att import AttControl -from .fp import FlashInferAttBackend, FlashInferPrefillAttState, FlashInferDecodeAttState -from .env_utils import set_flashinfer_envs -from lightllm.utils.log_utils import init_logger - - -logger = init_logger(__name__) - - -class Fp8FlashInferAttBackend(FlashInferAttBackend): - def __init__(self, model): - set_flashinfer_envs() - super().__init__(model=model) - self.kv_data_type = torch.float8_e4m3fn - - def _init_infer_page_size(self): - # TODO: FP8 FlashInfer 完成多 token 推理页适配后，改为继承模型 page_size。 - self.infer_page_size = 1 - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - logger.warning( - "Fp8FlashInferAttBackend temporarily uses infer_page_size=1 with " - f"model page_size={self.model.args.page_size}; multi-token FP8 inference pages are not implemented yet." - ) - - def create_att_prefill_state(self, infer_state) -> "Fp8FlashInferPrefillAttState": - return Fp8FlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Fp8FlashInferDecodeAttState": - return Fp8FlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Fp8FlashInferPrefillAttState(FlashInferPrefillAttState): - def init_state(self): - super().init_state() - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ) -> torch.Tensor: - o_tensor = alloc_func(q.shape, q.dtype, device="cuda") - k = k.unsqueeze(1).view(torch.float8_e4m3fn) - v = v.unsqueeze(1).view(torch.float8_e4m3fn) - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_tensor", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - k_descale = state.cpu_scales[storage_layer_index][0] - v_descale = state.cpu_scales[storage_layer_index][1] - self.prefill_wrapper.run( - q=q, - paged_kv_cache=(k, v), - k_scale=k_descale, - v_scale=v_descale, - out=o_tensor, - ) - return o_tensor - - -@dataclasses.dataclass -class Fp8FlashInferDecodeAttState(FlashInferDecodeAttState): - def init_state(self): - super().init_state() - - def copy_for_decode_cuda_graph(self, new_state): - return super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._fp8_decode_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - logical_layer_id=att_control.logical_layer_id, - ) - - def _fp8_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - logical_layer_id: int, - alloc_func=torch.empty, - ): - o_tensor = alloc_func(q.shape, q.dtype, device="cuda") - - k = k.unsqueeze(1).view(torch.float8_e4m3fn) - v = v.unsqueeze(1).view(torch.float8_e4m3fn) - state = self.infer_state.att_state_manager.get_state( - self.infer_state.att_state_role, "gqa_fp8_static_per_tensor", logical_layer_id - ) - storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - - k_descale = state.cpu_scales[storage_layer_index][0] - v_descale = state.cpu_scales[storage_layer_index][1] - self.decode_wrapper.run( - q=q, - paged_kv_cache=(k, v), - k_scale=k_descale, - v_scale=v_descale, - out=o_tensor, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/flashinfer/mla.py b/lightllm/common/basemodel/attention/flashinfer/mla.py deleted file mode 100644 index fc55175a3a..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/mla.py +++ /dev/null @@ -1,272 +0,0 @@ -import dataclasses -import weakref -import torch -import triton -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.dist_utils import get_dp_world_size, get_current_device_id -from ...triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor -from ...triton_kernel.repack_kv_index import repack_kv_index -from ...triton_kernel.flashinfer_mla_plan import fill_mla_decode_plan_for_cuda_graph -from typing import Tuple -from .env_utils import set_flashinfer_envs -from .utils import should_init_decode_wrapper - - -class MlaFlashInferAttBackend(BaseAttBackend): - workspace_buffer_key = "flashinfer_mla" - workspace_buffer_size = 256 * 1024 * 1024 - - def __init__(self, model): - set_flashinfer_envs() - super().__init__(model=model) - self._init_infer_page_size() - num_heads = model.config["num_attention_heads"] - self.tp_q_head_num = num_heads // get_dp_world_size() - self.qk_nope_head_dim = model.qk_nope_head_dim - self.qk_rope_head_dim = model.qk_rope_head_dim - self.kv_lora_rank = model.kv_lora_rank - self.v_head_dim = model.v_head_dim - self.q_data_type = model.data_type - self.kv_data_type = model.data_type - self.max_seq_length = model.max_seq_length - self.max_page_num = triton.cdiv(self.max_seq_length, self.infer_page_size) - self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5) - self.kv_indices_buffer = [ - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - torch.empty( - model.graph_max_batch_size * self.max_page_num, dtype=torch.int32, device=get_current_device_id() - ), - ] - - from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale - - if model.config["rope_scaling"] is not None: - rope_scaling = model.config["rope_scaling"] - mscale_all_dim = rope_scaling.get("mscale_all_dim", 0) - scaling_factor = rope_scaling["factor"] - if mscale_all_dim: - mscale = get_deepseek_mscale(scaling_factor, mscale_all_dim) - self.softmax_scale = self.softmax_scale * mscale * mscale - return - - def _init_infer_page_size(self): - self.infer_page_size = self.model.args.page_size - assert self.model.args.page_size % self.infer_page_size == 0, ( - f"model page_size {self.model.args.page_size} " - f"must be divisible by infer_page_size {self.infer_page_size}" - ) - - def create_att_prefill_state(self, infer_state) -> "MlaFlashInferPrefillAttState": - return MlaFlashInferPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "MlaFlashInferDecodeAttState": - return MlaFlashInferDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class MlaFlashInferPrefillAttState(BasePrefillAttState): - prefill_wrapper: object = None - - def init_state(self): - self.backend: MlaFlashInferAttBackend = self.backend - - import flashinfer - - q_starts = self.infer_state.b1_cu_q_seq_len.int() - kv_starts = self.infer_state.b1_cu_kv_seq_len.int() - if self.prefill_wrapper is None: - self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - "NHD", - ) - self.prefill_wrapper.plan( - qo_indptr=q_starts, - kv_indptr=kv_starts, - num_qo_heads=self.backend.tp_q_head_num, - num_kv_heads=self.backend.tp_q_head_num, - head_dim_qk=self.backend.qk_nope_head_dim + self.backend.qk_rope_head_dim, - head_dim_vo=self.backend.v_head_dim, - q_data_type=self.backend.q_data_type, - causal=True, - sm_scale=self.backend.softmax_scale, - ) - return - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - return self._mla_prefill_att( - q=q, - k=k, - v=v, - alloc_func=alloc_func, - ) - - def _mla_prefill_att( - self, q: torch.Tensor, k: Tuple[torch.Tensor, torch.Tensor], v: torch.Tensor, alloc_func=torch.empty - ) -> torch.Tensor: - self.backend: MlaFlashInferAttBackend = self.backend # for typing - k_nope, k_rope = k - o_tensor = alloc_func((q.shape[0], q.shape[1], v.shape[-1]), q.dtype, device="cuda") - q_head_num = q.shape[1] - k = torch.cat([k_nope, torch.repeat_interleave(k_rope, q_head_num, dim=-2)], dim=-1) - self.prefill_wrapper.run(q, k, v, out=o_tensor) - return o_tensor - - -@dataclasses.dataclass -class MlaFlashInferDecodeAttState(BaseDecodeAttState): - kv_indices: torch.Tensor = None - kv_starts: torch.Tensor = None - q_indptr_host: torch.Tensor = None - decode_wrapper: object = None - - def _should_init_decode_wrapper(self) -> bool: - return should_init_decode_wrapper(self.infer_state) - - def init_state(self): - import flashinfer - - self.backend: MlaFlashInferAttBackend = self.backend - model = self.backend.model - device = self.infer_state.input_ids.device - batch_size = self.infer_state.batch_size - - # TODO: 将页数及页数前缀和的计算融合为一个 Triton 算子。 - # token 长度除以页大小并向上取整，末页不足一页也计为一页。 - b_page_len = (self.infer_state.b_seq_len + (self.backend.infer_page_size - 1)) // self.backend.infer_page_size - self.kv_starts, _ = gen_cumsum_pad0_tensor(b_page_len, b_page_len) - - self.q_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") - self.q_indptr_host = torch.arange(batch_size + 1, dtype=torch.int32, device="cpu") - if batch_size <= model.graph_max_batch_size and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch: - self.kv_indices = self.backend.kv_indices_buffer[self.infer_state.microbatch_index][ - : batch_size * self.backend.max_page_num - ] - else: - self.kv_indices = torch.empty( - batch_size * self.backend.max_page_num, - dtype=torch.int32, - device=device, - ) - - repack_kv_index( - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_token_len=self.infer_state.b_seq_len, - b_page_start_loc=self.kv_starts[:-1], - max_token_len=self.infer_state.max_kv_seq_len, - out_page_indices=self.kv_indices, - page_size=self.backend.infer_page_size, - ) - - if not self._should_init_decode_wrapper(): - return - - assert self.decode_wrapper is None - - self.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( - self.backend.get_gpu_workspace_buffer( - key_name=self.backend.workspace_buffer_key, - workspace_size=self.backend.workspace_buffer_size, - ), - use_cuda_graph=True, - qo_indptr=self.q_indptr, - kv_indices=self.kv_indices, - kv_indptr=self.kv_starts, - kv_len_arr=self.infer_state.b_seq_len, - ) - self.decode_wrapper.plan( - self.q_indptr, - self.kv_starts, - self.kv_indices, - self.infer_state.b_seq_len, - self.backend.tp_q_head_num, - self.backend.kv_lora_rank, - self.backend.qk_rope_head_dim, - self.backend.infer_page_size, - False, # causal - self.backend.softmax_scale, - self.backend.q_data_type, - self.backend.kv_data_type, - ) - return - - def copy_for_decode_cuda_graph(self, new_state: "MlaFlashInferDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - self._refresh_cuda_graph_decode_plan(new_state.infer_state.b_seq_len, new_state.infer_state.max_kv_seq_len) - return - - def _refresh_cuda_graph_decode_plan(self, kv_lens: torch.Tensor, max_kv_len: int): - # Prefer the GPU-generated split plan for long decode; use exact non-split for - # short or unsupported graph shapes. - fill_mla_decode_plan_for_cuda_graph( - decode_wrapper=self.decode_wrapper, - kv_page_indptr=self.kv_starts, - kv_lens=kv_lens, - batch_size=self.infer_state.batch_size, - num_heads=self.backend.tp_q_head_num, - max_kv_len=max_kv_len, - ) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - - assert v is None - - return self._mla_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - qk_rope_head_dim = 64 - q_nope, q_rope = q - o_tensor = alloc_func(q_nope.shape, dtype=q_nope.dtype, device=q_nope.device) - assert att_control.mla_decode - - self.decode_wrapper.run( - q_nope, - q_rope, - k[:, :, :-qk_rope_head_dim].view(-1, self.backend.infer_page_size, 1, k.shape[-1] - qk_rope_head_dim), - k[:, :, -qk_rope_head_dim:].view(-1, self.backend.infer_page_size, 1, qk_rope_head_dim), - out=o_tensor, - return_lse=False, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/flashinfer/utils.py b/lightllm/common/basemodel/attention/flashinfer/utils.py deleted file mode 100644 index 8b159729bf..0000000000 --- a/lightllm/common/basemodel/attention/flashinfer/utils.py +++ /dev/null @@ -1,17 +0,0 @@ -def should_init_decode_wrapper(infer_state) -> bool: - match (infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay): - case (False, False): - # eager 推理（包括 graph 不可用时的回退）会直接执行 attention，需要创建 wrapper。 - return True - case (True, False): - # capture 需要创建 wrapper，并随捕获状态保存在图缓存中，供后续 replay 使用。 - return True - case (False, True): - # replay 的临时状态只准备待复制的 metadata，实际执行使用图缓存中的 wrapper。 - return False - case _: - # 两个标记必须是互斥的布尔值，不能同时处于 capture 和 replay 状态。 - raise ValueError( - f"Invalid CUDA Graph flags: capture={infer_state.is_cuda_graph_capture}, " - f"replay={infer_state.is_cuda_graph_replay}" - ) diff --git a/lightllm/common/basemodel/attention/import_utils.py b/lightllm/common/basemodel/attention/import_utils.py deleted file mode 100644 index 14cb0b076d..0000000000 --- a/lightllm/common/basemodel/attention/import_utils.py +++ /dev/null @@ -1,6 +0,0 @@ -from importlib import import_module - - -def load_backend_class(backend): - module_name, class_name = backend.split(":") - return getattr(import_module(f"{__package__}.{module_name}"), class_name) diff --git a/lightllm/common/basemodel/attention/linear/flashqla.py b/lightllm/common/basemodel/attention/linear/flashqla.py deleted file mode 100644 index 6e57eca732..0000000000 --- a/lightllm/common/basemodel/attention/linear/flashqla.py +++ /dev/null @@ -1,8 +0,0 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend - - -class FlashQlaLinearAttBackend(LinearAttBackend): - def get_prefill_kernel(self): - from flash_qla import chunk_gated_delta_rule - - return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py deleted file mode 100644 index 986214568f..0000000000 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ /dev/null @@ -1,430 +0,0 @@ -from lightllm.common.attention_state.bounded_att import GDNState -import dataclasses -import weakref -import torch -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn -from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( - mtp_fused_recurrent_gated_delta_rule, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( - fused_recurrent_gated_delta_rule, -) - -if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.common.basemodel.infer_struct import InferStateInfo - from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo - from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerWeight - - -class LinearAttBackend(BaseAttBackend, ABC): - def __init__(self, model: "TpPartBaseModel"): - super().__init__(model=model) - self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) - self.prefill_kernel = self.get_prefill_kernel() - - @abstractmethod - def get_prefill_kernel(self): - pass - - def _init_linear_layer_metadata(self, network_config, tp_world_size): - - self.mtp_step = get_env_start_args().mtp_step - - # Linear attention specific dimensions - self.num_v_heads = network_config["linear_num_value_heads"] - self.num_k_heads = network_config["linear_num_key_heads"] - self.head_k_dim = network_config["linear_key_head_dim"] - self.head_v_dim = network_config["linear_value_head_dim"] - self.key_dim = self.head_k_dim * self.num_k_heads - self.value_dim = self.head_v_dim * self.num_v_heads - self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] - self.activation = network_config["hidden_act"] - - # Tensor parallelism dimensions - self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size - self.tp_ba_dim = (self.num_v_heads * 2) // tp_world_size - self.tp_num_k_heads = self.num_k_heads // tp_world_size - self.tp_num_v_heads = self.num_v_heads // tp_world_size - self.tp_key_dim = self.key_dim // tp_world_size - self.tp_value_dim = self.value_dim // tp_world_size - - assert self.num_v_heads % self.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" - self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads - - # SSM state dtype optimization - ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} - self.ssm_state_dtype = ssm_dtype_dict.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) - return - - def _split_qkvzba(self, mixed_qkvzba): - qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim - z_end = qkv_dim + self.tp_value_dim - b_end = z_end + self.tp_num_v_heads - mixed_qkv = mixed_qkvzba[:, :qkv_dim] - z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, self.tp_num_v_heads, self.head_v_dim) - b = mixed_qkvzba[:, z_end:b_end] - a = mixed_qkvzba[:, b_end:] - return mixed_qkv, z, b, a - - def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): - if decode: - query, key, value = torch.split( - mixed_qkv, - [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], - dim=-1, - ) - batch_size = mixed_qkv.shape[0] - query = query.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) - key = key.view(batch_size, 1, self.tp_num_k_heads, self.head_k_dim) - value = value.view(batch_size, 1, self.tp_num_v_heads, self.head_v_dim) - return query, key, value - else: - query, key, value = torch.split( - mixed_qkv, - [self.tp_key_dim, self.tp_key_dim, self.tp_value_dim], - dim=-1, - ) - seq_len = query.shape[0] - query = query.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) - key = key.view(1, seq_len, self.tp_num_k_heads, self.head_k_dim) - value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) - return query, key, value - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "LinearAttPrefillAttState": - return LinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": - return LinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class LinearAttPrefillAttState(BasePrefillAttState): - - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - backend: LinearAttBackend = self.backend - mtp_step = backend.mtp_step - # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。 - # prefill cuda graph 回调必须走 new_infer_state.prefill_att_state1， - # 才能读到这里按当前 batch（含 token padding 后的 dummy request）更新的索引。 - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (mtp_step + 1) - return - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.linear_att_prefill, "linear_att_prefill must be True for Linear prefill attention" - assert att_control.linear_att_prefill_dict is not None, "linear_att_prefill_dict is required" - - linear_att_dict = att_control.linear_att_prefill_dict - mixed_qkvzba: torch.Tensor = linear_att_dict["mixed_qkvzba"] - layer_weight = linear_att_dict["layer_weight"] - layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend - - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=GDNState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 - # prefill 模式下，使用不到这几个维度，所以需要扣除掉， - if backend.mtp_step > 0: - conv_states = conv_states[:, :, : -backend.mtp_step] - mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - core_attn_out = self._gdn_prefill_kernel( - mixed_qkv, conv_states, ssm_states, a, b, self.infer_state, layer_weight - ) - return core_attn_out, z - - def _gdn_prefill_kernel( - self, - mixed_qkv: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: "Qwen3NextInferStateInfo", - layer_weight: "Qwen3NextTransformerLayerWeight", - ): - g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) - mixed_qkv = mixed_qkv.transpose(0, 1) - - backend: LinearAttBackend = self.backend - out_tensor = causal_conv1d_fn( - mixed_qkv, - layer_weight.linear_conv1d.mm_param.weight, - bias=layer_weight.linear_conv1d.bias, - query_start_loc=infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, - has_initial_state=infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation=backend.activation, - ) - mixed_qkv = out_tensor.transpose(0, 1) - - # Recurrent processing - query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) - initial_state = ssm_states[self.b_ssm_buffer_idx] - # g and beta have shape (total_tokens, num_heads), need to unsqueeze to get (1, total_tokens, num_heads) - core_attn_out, last_recurrent_state = backend.prefill_kernel( - q=query, - k=key, - v=value, - g=g.unsqueeze(0), - beta=beta.unsqueeze(0), - initial_state=initial_state, - output_final_state=True, - cu_seqlens=infer_state.b1_cu_q_seq_len, - use_qk_l2norm_in_kernel=True, - ) - # The chunk kernel accumulates the recurrent state in float32 even when - # the state cache is configured as bfloat16. Advanced indexing does - # not perform an implicit dtype conversion for index_put. - ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) - return core_attn_out - - -@dataclasses.dataclass -class LinearAttDecodeAttState(BaseDecodeAttState): - - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if draft_step == 0: - self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, GDNState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) - else: - self._init_fixed_mtp_decode_state(draft_step, req_to_mtp_state_index=state.mtp_state_index) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - mtp_size = draft_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "GDN fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={batch_size}, draft_step={draft_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - att_batch_size = self.b_conv_buffer_idx.shape[0] - # Each request owns mtp_size consecutive recurrent-state slots. - b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.linear_att_decode, "linear_att_decode must be True for Linear decode attention" - assert att_control.linear_att_decode_dict is not None, "linear_att_decode_dict is required" - - linear_att_dict = att_control.linear_att_decode_dict - mixed_qkvzba = linear_att_dict["mixed_qkvzba"] - layer_weight = linear_att_dict["layer_weight"] - layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend - - mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=GDNState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if draft_step > 0: - core_attn_out = self._gdn_mtp_kernel( - mixed_qkv, - conv_states, - ssm_states, - a, - b, - self.infer_state, - layer_weight, - ) - else: - # 非 MTP 模式下，使用线性层 decode 状态。 - core_attn_out, z = self._gdn_decode_kernel( - mixed_qkv, - z, - conv_states, - ssm_states, - a, - b, - self.infer_state, - layer_weight, - ) - return core_attn_out, z - - def _gdn_decode_kernel( - self, - mixed_qkv: torch.Tensor, - z: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: "Qwen3NextInferStateInfo", - layer_weight: "Qwen3NextTransformerLayerWeight", - ): - backend: LinearAttBackend = self.backend - - # Recurrent processing with fused gating. Decode uses a specialized - # conv+pack kernel to avoid materializing the post-conv qkv tensor - # before immediately splitting it into q/k/v. - query, key, value, z, a, b = conv_pack_gdn_decode_inputs( - mixed_qkv, - z, - a, - b, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, - backend.activation, - backend.conv_kernel_dim, - backend.tp_num_k_heads, - backend.head_k_dim, - backend.tp_num_v_heads, - backend.head_v_dim, - ) - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - use_qk_l2norm_in_kernel=True, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, - ) - return core_attn_out, z - - def _gdn_mtp_kernel( - self, - mixed_qkv: torch.Tensor, - conv_states: torch.Tensor, - ssm_states: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - infer_state: "Qwen3NextInferStateInfo", - layer_weight: "Qwen3NextTransformerLayerWeight", - ): - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, - ) - - backend: LinearAttBackend = self.backend - - cu_seqlens_q = self.b1_mtp_cu_q_seq_len - mixed_qkv = causal_conv1d_update_mtp( - mixed_qkv, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - mtp_step=backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model), - bias=layer_weight.linear_conv1d.bias, - activation=backend.activation, - conv_state_indices=self.b_conv_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=cu_seqlens_q, - ) - - query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - assert self.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" - # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1, - # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step - # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. - core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), - ssm_state_indices=self.b_ssm_buffer_idx, - ssm_state_write_indices=self.b_ssm_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, - ) - return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py deleted file mode 100644 index 11f801b62a..0000000000 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ /dev/null @@ -1,255 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""KDA attention backend for GLM-5-Next.""" - -from __future__ import annotations - -import dataclasses -import weakref -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention_state.bounded_att import KDAState -from lightllm.common.basemodel.attention.base_att import ( - AttControl, - BaseAttBackend, - BaseDecodeAttState, - BasePrefillAttState, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import causal_conv1d_qkv -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) - -if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class KDALinearAttBackend(BaseAttBackend): - def __init__(self, model: "TpPartBaseModel"): - super().__init__(model=model) - config = model.config["linear_attn_config"] - self.num_heads = config["num_heads"] - self.head_dim = config["head_dim"] - assert self.num_heads % model.tp_world_size_ == 0 - self.tp_num_heads = self.num_heads // model.tp_world_size_ - self.tp_hidden_size = self.tp_num_heads * self.head_dim - self.conv_kernel_size = config["short_conv_kernel_size"] - self.lower_bound = config.get("gate_lower_bound", -5.0) - self.mtp_step = model.args.mtp_step - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return KDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return KDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def split_qkv(self, mixed_qkv: torch.Tensor): - return mixed_qkv.split(self.tp_hidden_size, dim=-1) - - -@dataclasses.dataclass -class KDAPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_prefill - params = att_control.linear_att_prefill_dict - layer_weight = params["layer_weight"] - layer_num = params["layer_num"] - mixed_qkv = params["mixed_qkv"] - raw_gate = params["raw_gate"] - raw_beta = params["raw_beta"] - backend: KDALinearAttBackend = self.backend - - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_states = conv_states[..., : backend.conv_kernel_size - 1] - q, k, v = causal_conv1d_qkv( - mixed_qkv, - layer_weight.get_merged_kda_conv_weight(), - conv_states, - self.infer_state.b1_cu_q_seq_len, - self.b_conv_buffer_idx, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - ) - raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) - - initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() - output, final_state = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), - beta=raw_beta.float().sigmoid(), - A_log=layer_weight.linear_A_log.weight, - g_bias=layer_weight.linear_dt_bias.weight, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=False, - cu_seqlens=self.infer_state.b1_cu_q_seq_len, - safe_gate=True, - lower_bound=backend.lower_bound, - ) - ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) - return output - - -@dataclasses.dataclass -class KDADecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - mtp_step = self.backend.mtp_step - if mtp_step == 0: - self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, KDAState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(mtp_step + 1, req_to_mtp_state_index=state.mtp_state_index) - else: - self._init_fixed_mtp_decode_state(mtp_step, req_to_mtp_state_index=state.mtp_state_index) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - mtp_size = mtp_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "KDA fixed-layout decode requires batch_size to be divisible by mtp_step + 1, " - f"got batch_size={batch_size}, mtp_step={mtp_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - att_batch_size = self.b_conv_buffer_idx.shape[0] - # Each request owns mtp_size consecutive recurrent-state slots. - b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_decode - params = att_control.linear_att_decode_dict - layer_weight = params["layer_weight"] - layer_num = params["layer_num"] - mixed_qkv = params["mixed_qkv"] - raw_gate = params["raw_gate"] - raw_beta = params["raw_beta"] - backend: KDALinearAttBackend = self.backend - - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) - conv_update = causal_conv1d_update - if backend.mtp_step > 0: - conv_update = causal_conv1d_update_mtp - conv_kwargs.update( - mtp_step=backend.mtp_step, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=self.b1_mtp_cu_q_seq_len, - ) - mixed_qkv = conv_update(mixed_qkv, conv_states, layer_weight.get_merged_kda_conv_weight(), **conv_kwargs) - q, k, v = backend.split_qkv(mixed_qkv) - shape = (1, -1) if backend.mtp_step > 0 else (-1, 1) - q = q.view(*shape, backend.tp_num_heads, backend.head_dim) - k = k.view(*shape, backend.tp_num_heads, backend.head_dim) - v = v.view(*shape, backend.tp_num_heads, backend.head_dim) - raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) - raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - output, _ = fused_recurrent_kda( - q=q, - k=k, - v=v, - raw_gate=raw_gate, - raw_beta=raw_beta, - a_log=layer_weight.linear_A_log.weight, - gate_bias=layer_weight.linear_dt_bias.weight, - initial_state=ssm_states, - lower_bound=backend.lower_bound, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - cu_seqlens=self.b1_mtp_cu_q_seq_len, - num_accepted_tokens=self.b_num_accepted_tokens, - ) - return output diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py deleted file mode 100644 index 8e4a081f73..0000000000 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ /dev/null @@ -1,356 +0,0 @@ -"""GDN replay 的推理入口。 - -Triton 和 FlashQla 各有一个 backend。它们只差 prefill chunk kernel， -decode 都走 ReplayGDNState。两个类都直接实现 backend 接口，不继承 native GDN。 -""" - -from __future__ import annotations - -import dataclasses -import weakref -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ( - ReplayGDNState, -) -from lightllm.common.basemodel.attention.base_att import ( - AttControl, - BaseAttBackend, - BaseDecodeAttState, - BasePrefillAttState, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_fn, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import ( - fused_gdn_gating, -) -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import ( - conv_pack_gdn_decode_inputs, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.utils.envs_utils import get_env_start_args - -from .replay_gdn_decode import replay - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class ReplayGDNTritonAttBackend(BaseAttBackend): - """GDN replay 的 Triton 入口。Prefill 使用仓内 chunk kernel。""" - - def __init__(self, model): - super().__init__(model=model) - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule - - _configure_gdn(self, model, chunk_gated_delta_rule) - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -class ReplayGDNFlashQlaAttBackend(BaseAttBackend): - """GDN replay 的 FlashQla 入口。Prefill 使用 flash_qla chunk kernel。""" - - def __init__(self, model): - super().__init__(model=model) - from flash_qla import chunk_gated_delta_rule - - _configure_gdn(self, model, chunk_gated_delta_rule) - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class ReplayGDNPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - reqs = self.infer_state.b_req_idx - self.b_conv_buffer_idx = reqs - self.b_ssm_buffer_idx = reqs - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_prefill and att_control.linear_att_prefill_dict is not None - params = att_control.linear_att_prefill_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend = self.backend - state = _replay_state(self, layer_num) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - if backend.mtp_step > 0: - conv_states = conv_states[:, :, : -backend.mtp_step] - mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) - core_attn_out = self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight) - return core_attn_out, z - - def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): - backend = self.backend - g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) - mixed_qkv = causal_conv1d_fn( - mixed_qkv.transpose(0, 1), - layer_weight.linear_conv1d.mm_param.weight, - bias=layer_weight.linear_conv1d.bias, - query_start_loc=self.infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, - has_initial_state=self.infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation=backend.activation, - ).transpose(0, 1) - query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv) - core_attn_out, last_state = backend.prefill_kernel( - q=query, - k=key, - v=value, - g=g.unsqueeze(0), - beta=beta.unsqueeze(0), - initial_state=ssm_states[self.b_ssm_buffer_idx], - output_final_state=True, - cu_seqlens=self.infer_state.b1_cu_q_seq_len, - use_qk_l2norm_in_kernel=True, - ) - ssm_states[self.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) - return core_attn_out - - -@dataclasses.dataclass -class ReplayGDNDecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_cu_seqlens: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - role = self.infer_state.att_state_role - matches = [ - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayGDNState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayGDNState for {role}") - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - self._init_indices(draft_step, matches[0].mtp_state_index) - self.b_ssm_buffer_idx = self.b_conv_buffer_idx - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_decode and att_control.linear_att_decode_dict is not None - params = att_control.linear_att_decode_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend = self.backend - mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) - state = _replay_state(self, layer_num) - conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) - draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) - if draft_step > 0: - core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) - else: - core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) - # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) - return core, z - - def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: - infer = self.infer_state - if draft_step == 0: - self.b_conv_buffer_idx = infer.b_req_idx - self.b1_mtp_cu_q_seq_len = None - self.b_num_accepted_tokens = None - # 每行一个 token。replay 和 cursor 共用这份前缀和。 - self.b_cu_seqlens = torch.arange( - infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device - ) - return - if self.backend.uses_dynamic_spec_verify_layout(): - cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( - b_req_idx=infer.b_req_idx, - b_mtp_index=infer.b_mtp_index, - req_to_mtp_state_index=mtp_index, - hold_req_id=infer.req_manager.HOLD_REQUEST_ID, - ) - self.b1_mtp_cu_q_seq_len = cu - self.b_cu_seqlens = cu - self.b_conv_buffer_idx = conv_idx - self.b_num_accepted_tokens = accepted - return - width = draft_step + 1 - batch_size = infer.batch_size - if batch_size % width: - raise ValueError( - "GDN replay decode requires batch_size divisible by draft_step + 1, " - f"got batch_size={batch_size}, draft_step={draft_step}." - ) - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device - ) - self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len - self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() - self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 - - def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state): - backend = self.backend - query, key, value, z, a, b = conv_pack_gdn_decode_inputs( - mixed_qkv, - z, - a, - b, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, - backend.activation, - backend.conv_kernel_dim, - backend.tp_num_k_heads, - backend.head_k_dim, - backend.tp_num_v_heads, - backend.head_v_dim, - ) - query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) - key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) - value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) - output = replay( - state, - state.storage_layer(layer_num), - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b_cu_seqlens, - ) - return output, z - - def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): - backend = self.backend - mixed_qkv = causal_conv1d_update_mtp( - mixed_qkv, - conv_states, - layer_weight.linear_conv1d.mm_param.weight, - mtp_step=draft_step, - bias=layer_weight.linear_conv1d.bias, - activation=backend.activation, - conv_state_indices=self.b_conv_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=self.b1_mtp_cu_q_seq_len, - ) - query, key, value = torch.split( - mixed_qkv, - [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], - dim=-1, - ) - query = query.view(-1, backend.tp_num_k_heads, backend.head_k_dim) - key = key.view(-1, backend.tp_num_k_heads, backend.head_k_dim) - value = value.view(-1, backend.tp_num_v_heads, backend.head_v_dim) - output = replay( - state, - state.storage_layer(layer_num), - query, - key, - value, - a, - b, - layer_weight.linear_A_log.weight, - layer_weight.linear_dt_bias.weight, - self.b_conv_buffer_idx, - self.b_cu_seqlens, - ) - return output - - -def _replay_state(att_state, layer_num: int) -> ReplayGDNState: - return att_state.infer_state.att_state_manager.get_state( - model_role=att_state.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=ReplayGDNState, - ) - - -def _configure_gdn(backend, model, prefill_kernel) -> None: - """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" - - config = model.config - tp_world_size = model.tp_world_size_ - backend.mtp_step = get_env_start_args().mtp_step - backend.num_v_heads = config["linear_num_value_heads"] - backend.num_k_heads = config["linear_num_key_heads"] - backend.head_k_dim = config["linear_key_head_dim"] - backend.head_v_dim = config["linear_value_head_dim"] - backend.key_dim = backend.head_k_dim * backend.num_k_heads - backend.value_dim = backend.head_v_dim * backend.num_v_heads - backend.conv_kernel_dim = config["linear_conv_kernel_dim"] - backend.activation = config["hidden_act"] - backend.tp_qkvz_dim = (backend.key_dim * 2 + backend.value_dim * 2) // tp_world_size - backend.tp_ba_dim = (backend.num_v_heads * 2) // tp_world_size - backend.tp_num_k_heads = backend.num_k_heads // tp_world_size - backend.tp_num_v_heads = backend.num_v_heads // tp_world_size - backend.tp_key_dim = backend.key_dim // tp_world_size - backend.tp_value_dim = backend.value_dim // tp_world_size - assert backend.num_v_heads % backend.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" - backend.num_v_heads_per_k_head = backend.num_v_heads // backend.num_k_heads - ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} - backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) - backend.prefill_kernel = prefill_kernel - - -def _split_qkvzba(backend, mixed_qkvzba): - qkv_dim = backend.tp_key_dim * 2 + backend.tp_value_dim - z_end = qkv_dim + backend.tp_value_dim - b_end = z_end + backend.tp_num_v_heads - mixed_qkv = mixed_qkvzba[:, :qkv_dim] - z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, backend.tp_num_v_heads, backend.head_v_dim) - b = mixed_qkvzba[:, z_end:b_end] - a = mixed_qkvzba[:, b_end:] - return mixed_qkv, z, b, a - - -def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): - query, key, value = torch.split( - mixed_qkv, - [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], - dim=-1, - ) - if decode: - batch_size = mixed_qkv.shape[0] - leading = (batch_size, 1) - else: - leading = (1, query.shape[0]) - query = query.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) - key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) - value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) - return query, key, value diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py deleted file mode 100644 index af931aa477..0000000000 --- a/lightllm/common/basemodel/attention/linear/replay_kda.py +++ /dev/null @@ -1,257 +0,0 @@ -"""KDA replay 的推理入口。 - -这个 backend 只服务 ReplayKDAState。维度在这里自己读取，不继承 native KDA backend。 -""" - -from __future__ import annotations - -import dataclasses -import weakref -from typing import TYPE_CHECKING - -import torch - -from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( - ReplayKDAState, -) -from lightllm.common.basemodel.attention.base_att import ( - AttControl, - BaseAttBackend, - BaseDecodeAttState, - BasePrefillAttState, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_fn, - causal_conv1d_update, -) -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( - causal_conv1d_update as causal_conv1d_update_mtp, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( - chunk_kda_with_fused_gate, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) - -from .replay_kda_decode import replay - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class ReplayKDALinearAttBackend(BaseAttBackend): - """KDA replay 推理 backend。GLM-5-Next 在 replay 模式下只构造这一个类。""" - - def __init__(self, model): - super().__init__(model=model) - config = model.config["linear_attn_config"] - self.num_heads = config["num_heads"] - self.head_dim = config["head_dim"] - assert self.num_heads % model.tp_world_size_ == 0 - self.tp_num_heads = self.num_heads // model.tp_world_size_ - self.tp_hidden_size = self.tp_num_heads * self.head_dim - self.conv_kernel_size = config["short_conv_kernel_size"] - self.lower_bound = config.get("gate_lower_bound", -5.0) - self.mtp_step = model.args.mtp_step - - def split_qkv(self, mixed_qkv: torch.Tensor): - return mixed_qkv.split(self.tp_hidden_size, dim=-1) - - def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return ReplayKDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo"): - return ReplayKDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class ReplayKDAPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - reqs = self.infer_state.b_req_idx - self.b_conv_buffer_idx = reqs - self.b_ssm_buffer_idx = reqs - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_prefill - params = att_control.linear_att_prefill_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend: ReplayKDALinearAttBackend = self.backend - state = _replay_state(self, layer_num) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_states = conv_states[..., : backend.conv_kernel_size - 1] - mixed_qkv = causal_conv1d_fn( - params["mixed_qkv"].transpose(0, 1), - layer_weight.get_merged_kda_conv_weight(), - bias=None, - query_start_loc=self.infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, - has_initial_state=self.infer_state.b_ready_cache_len > 0, - conv_states=conv_states, - activation="silu", - ).transpose(0, 1) - q, k, v = backend.split_qkv(mixed_qkv) - q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) - k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) - v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) - raw_gate = params["raw_gate"].view(1, -1, backend.tp_hidden_size) - raw_beta = params["raw_beta"].view(1, -1, backend.tp_num_heads) - output, final_state = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), - beta=raw_beta.float().sigmoid(), - A_log=layer_weight.linear_A_log.weight, - g_bias=layer_weight.linear_dt_bias.weight, - initial_state=ssm_states[self.b_ssm_buffer_idx].contiguous(), - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=self.infer_state.b1_cu_q_seq_len, - safe_gate=True, - lower_bound=backend.lower_bound, - ) - ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) - return output - - -@dataclasses.dataclass -class ReplayKDADecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - b1_mtp_cu_q_seq_len: torch.Tensor = None - b_cu_seqlens: torch.Tensor = None - b_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - role = self.infer_state.att_state_role - matches = [ - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) - if isinstance(state, ReplayKDAState) - ] - if len(matches) != 1: - raise ValueError(f"expected one ReplayKDAState for {role}") - self._init_indices(self.backend.mtp_step, matches[0].mtp_state_index) - self.b_ssm_buffer_idx = self.b_conv_buffer_idx - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert att_control.linear_att_decode - params = att_control.linear_att_decode_dict - layer_num = params["layer_num"] - layer_weight = params["layer_weight"] - backend: ReplayKDALinearAttBackend = self.backend - state = _replay_state(self, layer_num) - conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) - conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) - conv_update = causal_conv1d_update - if backend.mtp_step > 0: - conv_update = causal_conv1d_update_mtp - conv_kwargs.update( - mtp_step=backend.mtp_step, - num_accepted_tokens=self.b_num_accepted_tokens, - query_start_loc=self.b1_mtp_cu_q_seq_len, - ) - mixed_qkv = conv_update( - params["mixed_qkv"], - conv_states, - layer_weight.get_merged_kda_conv_weight(), - **conv_kwargs, - ) - q, k, v = backend.split_qkv(mixed_qkv) - q = q.view(-1, backend.tp_num_heads, backend.head_dim) - k = k.view(-1, backend.tp_num_heads, backend.head_dim) - v = v.view(-1, backend.tp_num_heads, backend.head_dim) - raw_gate = params["raw_gate"].reshape(-1, backend.tp_hidden_size) - raw_beta = params["raw_beta"].reshape(-1, backend.tp_num_heads) - layer = state.storage_layer(layer_num) - output = replay( - state_ssm=state.ssm_state[layer], - state_keys=state.keys[layer], - state_deltas=state.deltas[layer], - state_gates=state.gates[layer], - state_raw_keys=state.raw_keys[layer], - state_raw_values=state.raw_values[layer], - state_betas=state.betas[layer], - state_cursors=state.cursors, - state_tile_sync_mark=state.tile_sync_mark, - q=q, - k=k, - v=v, - a=raw_gate, - b=raw_beta, - a_log=layer_weight.linear_A_log.weight, - bias=layer_weight.linear_dt_bias.weight, - reqs=self.b_conv_buffer_idx, - cu_seqlens=self.b_cu_seqlens, - state_verify_width=state.verify_width, - state_compute_mode=state.compute_mode, - state_gate_lower_bound=state.gate_lower_bound, - ) - # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) - return output - - def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: - infer = self.infer_state - if mtp_step == 0: - self.b_conv_buffer_idx = infer.b_req_idx - self.b1_mtp_cu_q_seq_len = None - self.b_num_accepted_tokens = None - # 每行一个 token。replay 和 cursor 共用这份前缀和。 - self.b_cu_seqlens = torch.arange( - infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device - ) - return - if self.backend.uses_dynamic_spec_verify_layout(): - cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( - b_req_idx=infer.b_req_idx, - b_mtp_index=infer.b_mtp_index, - req_to_mtp_state_index=mtp_index, - hold_req_id=infer.req_manager.HOLD_REQUEST_ID, - ) - self.b1_mtp_cu_q_seq_len = cu - self.b_cu_seqlens = cu - self.b_conv_buffer_idx = conv_idx - self.b_num_accepted_tokens = accepted - return - width = mtp_step + 1 - batch_size = infer.batch_size - if batch_size % width: - raise ValueError( - "KDA replay decode requires batch_size divisible by mtp_step + 1, " - f"got batch_size={batch_size}, mtp_step={mtp_step}." - ) - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device - ) - self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len - self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() - self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 - - -def _replay_state(att_state, layer_num: int) -> ReplayKDAState: - return att_state.infer_state.att_state_manager.get_state( - model_role=att_state.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=ReplayKDAState, - ) diff --git a/lightllm/common/basemodel/attention/linear/triton.py b/lightllm/common/basemodel/attention/linear/triton.py deleted file mode 100644 index f0b6400f51..0000000000 --- a/lightllm/common/basemodel/attention/linear/triton.py +++ /dev/null @@ -1,7 +0,0 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule - - -class TritonLinearAttBackend(LinearAttBackend): - def get_prefill_kernel(self): - return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py deleted file mode 100644 index 93560de07a..0000000000 --- a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py +++ /dev/null @@ -1,225 +0,0 @@ -import dataclasses -from typing import TYPE_CHECKING, Optional, Tuple - -import torch -import triton -import flash_mla as flashmla - - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState - -if TYPE_CHECKING: - from flash_mla.flash_mla_interface import FlashMLASchedMeta - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttBackend): - """FlashMLA layouts and operator workspace; forward states own sparse metadata.""" - - def __init__(self, model): - super().__init__(model=model) - from lightllm.models.deepseek_v4.config import create_attention_config - - config = create_attention_config( - model_config=model.config, - model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), - tp_rank=model.tp_rank_, - tp_world_size=model.tp_world_size_, - ) - self.tp_q_head_num = model.config["num_attention_heads"] // model.tp_world_size_ - # Sparse FlashMLA uses 64-head tiles; TP can leave fewer local Q heads. - self.padded_q_head_num = ((self.tp_q_head_num + 63) // 64) * 64 - self.compress_ratios = tuple(dict.fromkeys(config.compress_ratios)) - self.layer_compress_ratios = dict(zip(config.logical_layer_ids, config.compress_ratios)) - self.swa_config = config.create_swa_config() - self.compressed_configs = { - ratio: config.create_compressed_config(compress_ratio=ratio) for ratio in self.compress_ratios if ratio - } - - @classmethod - def get_swa_prefill_workspace(cls, model): - """由 backend 复用计算临时区；模型在 KV 容量评估前调用，以计入实际显存占用。""" - from lightllm.models.deepseek_v4.config import create_attention_config - - config = create_attention_config( - model_config=model.config, - model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), - tp_rank=model.tp_rank_, - tp_world_size=model.tp_world_size_, - ).create_swa_config() - page_token_num = config.page_token_size - max_forward_tokens = max(model.batch_max_tokens, model.max_req_num) + 2 * model.tp_world_size_ - capacity = ( - triton.cdiv( - max_forward_tokens + model.max_req_num * (config.window_size + 2 * page_token_num), page_token_num - ) - * page_token_num - ) - if model.args.run_mode == "decode": - # D 节点初始化仍会对 HOLD 请求执行一次 prefill；只保留末尾的 HOLD 页， - # 不为普通请求预留长 prefill 临时区，实际请求的前缀由 P 节点传入。 - capacity = 0 - regions = 1 + int(model.args.enable_prefill_microbatch_overlap) - # 每个 microbatch 使用独立临时区；最后一页留给 HOLD 请求。 - token_num = regions * capacity + page_token_num - # 物理行宽由 SWA 原子配置给出，backend 仅决定 workspace 的 token 容量。 - slot_bytes = config.get_packed_kv_bytes(page_token_num) // page_token_num - layer_count = len(config.logical_layer_ids) - buffer = cls.get_gpu_workspace_buffer( - key_name=f"deepseek_v4_swa_prefill_{model.att_state_role}", - workspace_size=layer_count * token_num * slot_bytes, - dtype=torch.uint8, - ) - return buffer.view(layer_count, token_num, slot_bytes), capacity - - def get_att_input_params(self, infer_state, logical_layer_id: int): - state = infer_state.swa_state - swa_buffer = state.get_layer_view(logical_layer_id=logical_layer_id) - config = state.config - swa = swa_buffer.view(-1, config.get_packed_kv_bytes(config.page_token_size)) - ratio = self.layer_compress_ratios[logical_layer_id] - compressed_kv = infer_state.get_compressed_kv_view(logical_layer_id=logical_layer_id).kv if ratio else None - if infer_state.is_prefill: - swa = infer_state.prefill_att_state.swa_buffer[config.logical_layer_ids.index(logical_layer_id)].view( - -1, - config.get_packed_kv_bytes(config.page_token_size), - ) - return swa, compressed_kv - - def _flashmla_att( - self, - q: torch.Tensor, - packed_kv: Tuple[torch.Tensor, Optional[torch.Tensor]], - nsa_dict: dict, - sched_meta: "FlashMLASchedMeta", - out: torch.Tensor, - ) -> torch.Tensor: - """Attend to SWA plus optional compressed KV and return all padded Q heads.""" - swa_packed, compressed_packed = packed_kv - # 原子配置提供计算视图；backend 不再计算编码字节数或解释压缩块布局。 - compressed_cache = None - if compressed_packed is not None: - config = self.compressed_configs[nsa_dict["compress_ratio"]] - compressed_cache = config.get_attention_view(packed_kv=compressed_packed) - swa_cache = self.swa_config.get_attention_view(packed_kv=swa_packed) - - padded_out, _ = flashmla.flash_mla_with_kvcache( - q=q.unsqueeze(1), - k_cache=swa_cache, - block_table=None, - cache_seqlens=None, - head_dim_v=nsa_dict["head_dim_v"], - tile_scheduler_metadata=sched_meta, - num_splits=None, - softmax_scale=nsa_dict["softmax_scale"], - causal=False, - is_fp8_kvcache=True, - indices=nsa_dict["swa_indices"], - attn_sink=nsa_dict["attn_sink"], - topk_length=nsa_dict["swa_lengths"], - extra_k_cache=compressed_cache, - extra_indices_in_kvcache=nsa_dict.get("extra_indices"), - extra_topk_length=nsa_dict.get("extra_lengths"), - out=out.unsqueeze(1), - ) - return padded_out[:, 0] - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "_PrefillAttState": - return _PrefillAttState(backend=self) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "_DecodeAttState": - return _DecodeAttState(backend=self) - - -@dataclasses.dataclass -class _PrefillAttState(BasePrefillAttState): - sched_meta_by_ratio: dict[int, "FlashMLASchedMeta"] | None = None - - def prepare_swa_workspace(self, infer_state): - from lightllm.models.deepseek_v4.triton_kernel.swa_cache import prepare_swa_prefill_workspace - - swa_state = infer_state.swa_state - self.swa_buffer, capacity = self.backend.get_swa_prefill_workspace(self.backend.model) - self.HOLD_SWA_SLOT = self.swa_buffer.shape[1] - swa_state.config.page_token_size - self.prefill_token_ranges = torch.zeros( - (swa_state.request_capacity, 2), dtype=torch.int32, device=self.swa_buffer.device - ) - self.swa_prefill_workspace_offsets = torch.full( - (swa_state.request_capacity,), -1, dtype=torch.int32, device=self.swa_buffer.device - ) - prepare_swa_prefill_workspace( - swa_state=swa_state, - prefill_state=self, - b_req_idx=infer_state.b_req_idx, - b_ready_cache_len=infer_state.b_ready_cache_len, - b_seq_len=infer_state.b_seq_len, - workspace_base=infer_state.microbatch_index * capacity, - ) - - def init_state(self): - self.sched_meta_by_ratio = {ratio: flashmla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios} - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, Optional[torch.Tensor]], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - *, - out: torch.Tensor = None, - ) -> torch.Tensor: - nsa_dict = att_control.nsa_prefill_dict - if out is None: - out = alloc_func((q.shape[0], q.shape[1], nsa_dict["head_dim_v"]), dtype=q.dtype, device=q.device) - - padded_out = self.backend._flashmla_att( - q, - k, - nsa_dict, - self.sched_meta_by_ratio[nsa_dict["compress_ratio"]], - out=out, - ) - return padded_out[:, : self.backend.tp_q_head_num] - - -@dataclasses.dataclass -class _DecodeAttState(BaseDecodeAttState): - sched_meta_by_ratio: dict[int, "FlashMLASchedMeta"] | None = None - capture_sched_meta_by_ratio: dict[int, "FlashMLASchedMeta"] | None = None - - def init_state(self): - self.sched_meta_by_ratio = {ratio: flashmla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios} - # Keep capture metadata untouched by warmup so its GPU planning kernels enter the graph. - self.capture_sched_meta_by_ratio = { - ratio: flashmla.get_mla_metadata()[0] for ratio in self.backend.compress_ratios - } - - def decode_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, Optional[torch.Tensor]], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - nsa_dict = att_control.nsa_decode_dict - sched_meta_by_ratio = ( - self.capture_sched_meta_by_ratio if torch.cuda.is_current_stream_capturing() else self.sched_meta_by_ratio - ) - out = alloc_func((q.shape[0], q.shape[1], nsa_dict["head_dim_v"]), dtype=q.dtype, device=q.device) - padded_out = self.backend._flashmla_att( - q, - k, - nsa_dict, - sched_meta_by_ratio[nsa_dict["compress_ratio"]], - out=out, - ) - return padded_out[:, : self.backend.tp_q_head_num] - - -DSV4_NSA_BACKENDS = { - "fp8kv_dsa": { - "flashmla_sparse": "nsa.dsv4_fp8_flashmla_sparse:DeepseekV4FlashMlaFp8SparseAttBackend", - } -} diff --git a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py deleted file mode 100644 index 92c48c516f..0000000000 --- a/lightllm/common/basemodel/attention/nsa/flashmla_sparse.py +++ /dev/null @@ -1,193 +0,0 @@ -# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/nsa_backend.py -# Uses sgl_kernel.flash_mla and sgl_kernel.flash_attn from the sglang kernel library. - -import dataclasses -import weakref -import torch -from typing import Tuple, TYPE_CHECKING - -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from lightllm.utils.dist_utils import get_current_device_id - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class NsaFlashMlaSparseAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model=model) - device = get_current_device_id() - self.ragged_mem_buffers = [ - torch.empty(model.graph_max_batch_size * model.max_seq_length, dtype=torch.int32, device=device) - for _ in range(2) - ] - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaSparsePrefillAttState": - return NsaFlashMlaSparsePrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaSparseDecodeAttState": - return NsaFlashMlaSparseDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NsaFlashMlaSparsePrefillAttState(BasePrefillAttState): - """Prefill attention state for NSA using flash_mla_sparse_fwd.""" - - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - - def init_state(self): - self.backend: NsaFlashMlaSparseAttBackend = self.backend - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.input_ids.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - return - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_prefill, "nsa_prefill must be True for NSA prefill attention" - assert att_control.nsa_prefill_dict is not None, "nsa_prefill_dict is required" - - return self._nsa_prefill_att(q=q, kv=k, att_control=att_control) - - def _nsa_prefill_att( - self, - q: torch.Tensor, - kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - from sgl_kernel.flash_mla import flash_mla_sparse_fwd - - nsa_dict = att_control.nsa_prefill_dict - topk_mem_indices = nsa_dict["topk_mem_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - - if topk_mem_indices.ndim == 2: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - - mla_out, _, _ = flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=topk_mem_indices, - sm_scale=softmax_scale, - d_v=kv_lora_rank, - ) - return mla_out - - -@dataclasses.dataclass -class NsaFlashMlaSparseDecodeAttState(BaseDecodeAttState): - - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - nsa_cache_seqlens: torch.Tensor = None - nsa_cu_seqlens_k_new: torch.Tensor = None - - def init_state(self): - self.backend: NsaFlashMlaSparseAttBackend = self.backend - model = self.backend.model - use_cuda_graph = ( - self.infer_state.batch_size <= model.graph_max_batch_size - and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch - ) - - if use_cuda_graph: - self.ragged_mem_index = self.backend.ragged_mem_buffers[self.infer_state.microbatch_index] - else: - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.b_seq_len.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self.nsa_cache_seqlens = torch.minimum( - torch.full(size=(self.infer_state.batch_size,), fill_value=2048, dtype=torch.int32, device="cuda"), - self.infer_state.b_seq_len, - ) - padded_seq_lens = torch.zeros(size=(self.nsa_cache_seqlens.shape[0] + 1,), dtype=torch.int32, device="cuda") - # 进行 cumsum 操作 - padded_seq_lens[1:].copy_(self.nsa_cache_seqlens, non_blocking=True) - self.nsa_cu_seqlens_k_new = padded_seq_lens.cumsum(dim=0, dtype=torch.int32) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_decode, "nsa_decode must be True for NSA decode attention" - assert att_control.nsa_decode_dict is not None, "nsa_decode_dict is required" - - return self._nsa_decode_att(q=q, kv=k, att_control=att_control) - - def _nsa_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - nsa_dict = att_control.nsa_decode_dict - topk_mem_indices = nsa_dict["topk_mem_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - qk_rope_head_dim = nsa_dict["qk_rope_head_dim"] - - q_nope, q_rope = q - - # Extract k_rope and kv_nope from the KV buffer - k_rope = kv[:, :, -qk_rope_head_dim:].view(-1, 1, 1, qk_rope_head_dim) - kv_nope = kv[:, :, :-qk_rope_head_dim].view(-1, 1, 1, kv_lora_rank) - - o_tensor = flash_attn_with_kvcache( - q=q_rope, - k_cache=k_rope, - v_cache=kv_nope, - qv=q_nope, - page_table=topk_mem_indices, - cache_seqlens=self.nsa_cache_seqlens, - cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, - cu_seqlens_k_new=self.nsa_cu_seqlens_k_new, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=softmax_scale, - causal=True, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py deleted file mode 100644 index 10ad62d6b8..0000000000 --- a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py +++ /dev/null @@ -1,204 +0,0 @@ -import dataclasses -import weakref -import torch -from typing import TYPE_CHECKING, Tuple - -from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState, BasePrefillAttState -from lightllm.utils.dist_utils import get_current_device_id - -if TYPE_CHECKING: - from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class NsaFlashMlaFp8SparseAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model=model) - device = get_current_device_id() - self.ragged_mem_buffers = [ - torch.empty(model.graph_max_batch_size * model.max_seq_length, dtype=torch.int32, device=device) - for _ in range(2) - ] - - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparsePrefillAttState": - return NsaFlashMlaFp8SparsePrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "NsaFlashMlaFp8SparseDecodeAttState": - return NsaFlashMlaFp8SparseDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NsaFlashMlaFp8SparsePrefillAttState(BasePrefillAttState): - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - - def init_state(self): - self.backend: NsaFlashMlaFp8SparseAttBackend = self.backend - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.input_ids.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - return - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_prefill, "nsa_prefill must be True for NSA prefill attention" - assert att_control.nsa_prefill_dict is not None, "nsa_prefill_dict is required" - return self._nsa_prefill_att(q=q, packed_kv=k, att_control=att_control) - - def _nsa_prefill_att( - self, - q: torch.Tensor, - packed_kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - import flash_mla - - nsa_dict = att_control.nsa_prefill_dict - topk_indices = nsa_dict["topk_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - topk_mem_indices = nsa_dict["topk_mem_indices"] - prefill_cache_kv = nsa_dict["prefill_cache_kv"] - - if self.infer_state.max_cache_len > 0: - # 当前推理生成的token kv部分从 prefill_cache_kv 中获取，历史 - # 部分kv 从 packed_kv 中获取, 并进行反量化，这样可以避免 prefill_cache_kv - # 部分的数据进行重复的反量化操作，提升整体的性能。 - from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( - get_prefill_kv_cache_and_remap_indices_triton, - ) - - kv, topk_indices = get_prefill_kv_cache_and_remap_indices_triton( - packed_kv=packed_kv, - topk_mem_indices=topk_mem_indices, - prefill_mem_index=self.infer_state.mem_index, - prefill_cache_kv=prefill_cache_kv, - prefill_dtype=prefill_cache_kv.dtype, - ) - else: - kv = prefill_cache_kv - - if topk_indices.ndim == 2: - topk_indices = topk_indices.unsqueeze(1) - - mla_out, _, _ = flash_mla.flash_mla_sparse_fwd( - q=q, - kv=kv, - indices=topk_indices, - sm_scale=softmax_scale, - d_v=kv_lora_rank, - ) - return mla_out - - -@dataclasses.dataclass -class NsaFlashMlaFp8SparseDecodeAttState(BaseDecodeAttState): - ks: torch.Tensor = None - ke: torch.Tensor = None - lengths: torch.Tensor = None - ragged_mem_index: torch.Tensor = None - flashmla_sched_meta: object = None - - def init_state(self): - self.backend: NsaFlashMlaFp8SparseAttBackend = self.backend - model = self.backend.model - use_cuda_graph = ( - self.infer_state.batch_size <= model.graph_max_batch_size - and self.infer_state.max_kv_seq_len <= model.graph_max_len_in_batch - ) - - if use_cuda_graph: - self.ragged_mem_index = self.backend.ragged_mem_buffers[self.infer_state.microbatch_index] - else: - self.ragged_mem_index = torch.empty( - self.infer_state.total_token_num, - dtype=torch.int32, - device=get_current_device_id(), - ) - - from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke - - self.ks, self.ke, self.lengths = gen_nsa_ks_ke( - b_seq_len=self.infer_state.b_seq_len, - b_q_seq_len=self.infer_state.b_q_seq_len, - b_req_idx=self.infer_state.b_req_idx, - req_to_token_index=self.infer_state.req_manager.req_to_token_indexs, - q_token_num=self.infer_state.b_seq_len.shape[0], - ragged_mem_index=self.ragged_mem_index, - hold_req_idx=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - import flash_mla - - self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata() - return - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert att_control.nsa_decode, "nsa_decode must be True for NSA decode attention" - assert att_control.nsa_decode_dict is not None, "nsa_decode_dict is required" - return self._nsa_decode_att(q=q, packed_kv=k, att_control=att_control) - - def _nsa_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - packed_kv: torch.Tensor, - att_control: AttControl, - ) -> torch.Tensor: - import flash_mla - - nsa_dict = att_control.nsa_decode_dict - topk_mem_indices = nsa_dict["topk_mem_indices"] - softmax_scale = nsa_dict["softmax_scale"] - kv_lora_rank = nsa_dict["kv_lora_rank"] - - if topk_mem_indices.ndim == 2: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" - - q_nope, q_rope = q - q_all = torch.cat([q_nope, q_rope], dim=-1).unsqueeze(1).contiguous() - kv = torch.as_strided( - packed_kv, - size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), - stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), - ) - - o_tensor, _ = flash_mla.flash_mla_with_kvcache( - q=q_all, - k_cache=kv, - block_table=None, - cache_seqlens=None, - head_dim_v=kv_lora_rank, - tile_scheduler_metadata=self.flashmla_sched_meta, - softmax_scale=softmax_scale, - causal=False, - is_fp8_kvcache=True, - indices=topk_mem_indices, - ) - return o_tensor[:, 0, :, :] # [b, 1, h, d] -> [b, h, d] diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py deleted file mode 100644 index 261c2a16c1..0000000000 --- a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py +++ /dev/null @@ -1,238 +0,0 @@ -import dataclasses -import weakref - -import torch -import triton - -from .glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState -from ..base_att import BaseAttBackend - - -# Cap each microbatch's BF16 KV workspace at 2 GiB (2097152 × 512 × 2 bytes). -KV_GATHER_CHUNK_SIZE = 2097152 - - -class Fp8Glm5NextSparseAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model=model) - self._prefill_workspaces = self.allocate_prefill_workspaces(model) - - @classmethod - def allocate_prefill_workspaces(cls, model): - regions = 1 + int(model.args.enable_prefill_microbatch_overlap) - q_token_capacity = model.batch_max_tokens - topk_size = triton.cdiv(model.config["index_topk"] + 3, 128) * 128 - kv_token_capacity = min(KV_GATHER_CHUNK_SIZE, model.max_req_num * model.max_seq_length + q_token_capacity) - if model.args.run_mode == "decode": - kv_token_capacity = min(kv_token_capacity, q_token_capacity) - prefix = f"glm5_fp8_prefill_{model.att_state_role}" - gathered_kv = cls.get_gpu_workspace_buffer( - prefix + "_kv", regions * kv_token_capacity * 512, torch.bfloat16 - ).view(regions, kv_token_capacity, 1, 512) - chunk_topk_indices = cls.get_gpu_workspace_buffer( - prefix + "_topk", regions * q_token_capacity * topk_size, torch.int32 - ).view(regions, q_token_capacity, topk_size) - chunk_cache_seqlens = cls.get_gpu_workspace_buffer( - prefix + "_seqlens", regions * q_token_capacity, torch.int32 - ).view(regions, q_token_capacity) - cu_seqlens_q = cls.get_gpu_workspace_buffer( - prefix + "_cu_q", regions * (q_token_capacity + 1), torch.int32 - ).view(regions, q_token_capacity + 1) - cu_seqlens_q.copy_(torch.arange(q_token_capacity + 1, dtype=torch.int32, device=cu_seqlens_q.device)) - return [ - { - "gathered_kv": gathered_kv[index], - "chunk_topk_indices": chunk_topk_indices[index], - "chunk_cache_seqlens": chunk_cache_seqlens[index], - "cu_seqlens_q": cu_seqlens_q[index], - } - for index in range(regions) - ] - - def get_prefill_workspace(self, infer_state, topk_indices): - buffers = self._prefill_workspaces[infer_state.microbatch_index] - q_token_num, topk_size = topk_indices.shape - assert q_token_num <= buffers["chunk_cache_seqlens"].numel() - assert topk_size <= buffers["chunk_topk_indices"].shape[1] - chunk_topk_indices = buffers["chunk_topk_indices"].view(-1) - return { - "gathered_kv": buffers["gathered_kv"], - "chunk_topk_indices": chunk_topk_indices[: q_token_num * topk_size].view(q_token_num, topk_size), - "chunk_cache_seqlens": buffers["chunk_cache_seqlens"][:q_token_num], - "cu_seqlens_q": buffers["cu_seqlens_q"][: q_token_num + 1], - } - - def create_att_prefill_state(self, infer_state): - return Fp8Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state): - config = self.model.config - # 2048 pooled entries plus at most three kpool remainder entries fit the 2112 rounded direct window. - bound = 2051 if config["index_topk"] == 2048 and config["index_kpool"] == 4 else None - return Fp8Glm5NextSparseDecodeState( - backend=self, _infer_state_weak_ref=weakref.ref(infer_state), decode_topk_bound=bound - ) - - -@dataclasses.dataclass -class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState): - @torch.no_grad() - def _nsa_prefill_att(self, q, packed_kv, att_control): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( - _gather_prefill_kv_kernel, - _compact_topk_indices_kernel, - _merge_att_output_kernel, - ) - - nsa_prefill_dict = att_control.nsa_prefill_dict - infer_state = self.infer_state - q_token_num, q_head_num, _ = q.shape - prefill_cache_kv = nsa_prefill_dict["prefill_cache_kv"] - topk_indices = nsa_prefill_dict["topk_indices"] - softmax_scale = nsa_prefill_dict["softmax_scale"] - req_to_token_indexs = infer_state.req_manager.req_to_token_indexs - b_req_idx = infer_state.b_req_idx - b1_cu_kv_seq_len = infer_state.b1_cu_kv_seq_len - b1_cu_q_seq_len = infer_state.b1_cu_q_seq_len - total_token_num = infer_state.total_token_num - workspace = self.backend.get_prefill_workspace(infer_state, topk_indices) - gathered_kv = workspace["gathered_kv"] - gather_chunk_size = min(gathered_kv.shape[0], total_token_num) - chunk_topk_indices = workspace["chunk_topk_indices"] - chunk_cache_seqlens = workspace["chunk_cache_seqlens"] - cu_seqlens_q = workspace["cu_seqlens_q"] - has_history = total_token_num > q_token_num - if not has_history and prefill_cache_kv.stride(-1) != 1: - prefill_cache_kv = prefill_cache_kv.contiguous() - output = softmax_lse = None - for chunk_start in range(0, total_token_num, gather_chunk_size): - chunk_token_num = min(gather_chunk_size, total_token_num - chunk_start) - if has_history: - _gather_prefill_kv_kernel[(chunk_token_num // 2 + b_req_idx.numel(),)]( - packed_kv, - prefill_cache_kv, - gathered_kv, - req_to_token_indexs, - b_req_idx, - b1_cu_kv_seq_len, - b1_cu_q_seq_len, - chunk_start, - chunk_token_num, - packed_kv.stride(0), - prefill_cache_kv.stride(0), - prefill_cache_kv.stride(2), - req_to_token_indexs_stride_b=req_to_token_indexs.stride(0), - req_to_token_indexs_stride_s=req_to_token_indexs.stride(1), - batch_size=b_req_idx.numel(), - BLOCK_ROWS=2, - num_warps=2, - ) - chunk_kv = gathered_kv[:chunk_token_num].view(-1, 1, 1, 512) - else: - # Current BF16 KV already has the complete request layout. - chunk_kv = prefill_cache_kv[chunk_start : chunk_start + chunk_token_num].view(-1, 1, 1, 512) - _compact_topk_indices_kernel[(q_token_num,)]( - topk_indices, - self.q_kv_start_loc, - chunk_topk_indices, - chunk_cache_seqlens, - chunk_start, - chunk_start + chunk_token_num, - topk_indices.stride(0), - topk_indices.stride(1), - TOPK=topk_indices.shape[1], - BLOCK=triton.next_power_of_2(topk_indices.shape[1]), - ) - chunk_output, chunk_softmax_lse, *_ = flash_attn_with_kvcache( - # only_qv ignores Q/K; these views avoid allocating placeholder tensors. - q=q[..., :64], - k_cache=chunk_kv[..., :64], - qv=q, - v_cache=chunk_kv, - page_table=chunk_topk_indices, - cache_seqlens=chunk_cache_seqlens, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=1, - softmax_scale=softmax_scale, - causal=False, - only_qv=True, - num_splits=1, - return_softmax_lse=True, - ) - if chunk_start == 0: - output, softmax_lse = chunk_output, chunk_softmax_lse - else: - # Keep the online reduction in FP32 across all remaining chunks. - output = output.float() - _merge_att_output_kernel[(q_token_num, triton.cdiv(q_head_num, 4))]( - output, - softmax_lse, - chunk_output, - chunk_softmax_lse, - chunk_softmax_lse.stride(1), - chunk_softmax_lse.stride(0), - q_head_num=q_head_num, - q_token_num=q_token_num, - BLOCK_HEADS=4, - ) - del chunk_output, chunk_softmax_lse - return output.to(q.dtype) - - -@dataclasses.dataclass -class Fp8Glm5NextSparseDecodeState(Glm5NextSparseDecodeState): - decode_topk_bound: int = None - flashmla_sched_meta: object = None - - def init_state(self): - import flash_mla - - self.q_kv_seq_len = self.infer_state.b_seq_len - self.q_req_idx = self.infer_state.b_req_idx - self.flashmla_sched_meta, _ = flash_mla.get_mla_metadata() - - def _nsa_decode_att(self, q, packed_kv, att_control): - import flash_mla - - from lightllm.models.glm5_next.triton_kernel import fp8_sparse_decode - from lightllm.models.glm5_next.triton_kernel.pad_query_heads import pad_query_heads - - nsa_decode_dict = att_control.nsa_decode_dict - topk_mem_indices = nsa_decode_dict["topk_mem_indices"] - if topk_mem_indices.ndim == 2: - topk_mem_indices = topk_mem_indices.unsqueeze(1) - assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" - - q_nope, _ = q - num_heads = q_nope.shape[1] - if fp8_sparse_decode.can_use_fp8_sparse_direct_decode( - q_nope, - packed_kv, - topk_mem_indices, - nsa_decode_dict["kv_lora_rank"], - self.decode_topk_bound, - is_cuda_graph_capture=self.infer_state.is_cuda_graph_capture, - ): - return fp8_sparse_decode.fp8_sparse_direct_decode( - q_nope, packed_kv, topk_mem_indices, nsa_decode_dict["softmax_scale"] - ) - padded_q = pad_query_heads(q_nope, is_cuda_graph_capture=self.infer_state.is_cuda_graph_capture) - kv_cache = torch.as_strided( - packed_kv, - size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), - stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), - ) - output, _ = flash_mla.flash_mla_with_kvcache( - q=padded_q, - k_cache=kv_cache, - block_table=None, - cache_seqlens=None, - head_dim_v=nsa_decode_dict["kv_lora_rank"], - tile_scheduler_metadata=self.flashmla_sched_meta, - softmax_scale=nsa_decode_dict["softmax_scale"], - causal=False, - is_fp8_kvcache=True, - indices=topk_mem_indices, - ) - return output[:, 0, :num_heads, :] diff --git a/lightllm/common/basemodel/attention/nsa/glm5_next.py b/lightllm/common/basemodel/attention/nsa/glm5_next.py deleted file mode 100644 index 6e33177dbe..0000000000 --- a/lightllm/common/basemodel/attention/nsa/glm5_next.py +++ /dev/null @@ -1,100 +0,0 @@ -import dataclasses -import weakref - -import torch - -from ..base_att import AttControl, BaseAttBackend, BasePrefillAttState, BaseDecodeAttState - - -class Glm5NextSparseAttBackend(BaseAttBackend): - def create_att_prefill_state(self, infer_state): - return Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state): - return Glm5NextSparseDecodeState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Glm5NextSparsePrefillState(BasePrefillAttState): - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None - q_kv_start_loc: torch.Tensor = None - - def init_state(self): - from lightllm.models.glm5_next.triton_kernel.kpool import prefill_nsa_params - - infer_state = self.infer_state - self.q_kv_seq_len, self.q_req_idx, self.q_kv_start_loc = prefill_nsa_params( - infer_state.position_ids, infer_state.b_req_idx, infer_state.b1_cu_q_seq_len, infer_state.b1_cu_kv_seq_len - ) - - def prefill_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): - assert att_control.nsa_prefill and att_control.nsa_prefill_dict is not None - return self._nsa_prefill_att(q, k, att_control) - - def _nsa_prefill_att(self, q, kv, att_control): - from sgl_kernel.flash_mla import flash_mla_sparse_fwd - - q_token_num, q_head_num, head_dim = q.shape - # FlashMLA uses 64-head tiles. - padded_q_head_num = ((q_head_num + 63) // 64) * 64 - padded_q = q - if padded_q_head_num != q_head_num: - padded_q = q.new_zeros((q_token_num, padded_q_head_num, head_dim)) - padded_q[:, :q_head_num] = q - nsa_prefill_dict = att_control.nsa_prefill_dict - output, _, _ = flash_mla_sparse_fwd( - q=padded_q, - kv=kv, - indices=nsa_prefill_dict["topk_mem_indices"].unsqueeze(1), - sm_scale=nsa_prefill_dict["softmax_scale"], - d_v=nsa_prefill_dict["kv_lora_rank"], - ) - return output[:, :q_head_num] - - -@dataclasses.dataclass -class Glm5NextSparseDecodeState(BaseDecodeAttState): - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None - nsa_cache_seqlens: torch.Tensor = None - nsa_cu_seqlens_k_new: torch.Tensor = None - - def init_state(self): - self.q_kv_seq_len = self.infer_state.b_seq_len - self.q_req_idx = self.infer_state.b_req_idx - index_kpool = self.backend.model.config["index_kpool"] - index_topk = self.backend.model.config["index_topk"] - self.nsa_cache_seqlens = ( - torch.minimum( - self.q_kv_seq_len // index_kpool * index_kpool, torch.full_like(self.q_kv_seq_len, index_topk) - ) - + self.q_kv_seq_len % index_kpool - ) - self.nsa_cu_seqlens_k_new = torch.nn.functional.pad(self.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0)) - - def decode_att(self, q, k, v, att_control=AttControl(), alloc_func=torch.empty): - assert att_control.nsa_decode and att_control.nsa_decode_dict is not None - return self._nsa_decode_att(q, k, att_control) - - def _nsa_decode_att(self, q, kv, att_control): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - nsa_decode_dict = att_control.nsa_decode_dict - q_nope, _ = q - kv_nope = kv.view(-1, 1, 1, nsa_decode_dict["kv_lora_rank"]) - # NoPE has no RoPE Q/K tensors; only_qv uses q_nope and kv_nope directly. - return flash_attn_with_kvcache( - q=None, - qv=q_nope, - k_cache=None, - v_cache=kv_nope, - page_table=nsa_decode_dict["topk_mem_indices"], - cache_seqlens=self.nsa_cache_seqlens, - cu_seqlens_q=self.infer_state.b1_cu_q_seq_len, - cu_seqlens_k_new=self.nsa_cu_seqlens_k_new, - max_seqlen_q=self.infer_state.max_q_seq_len, - softmax_scale=nsa_decode_dict["softmax_scale"], - causal=False, - only_qv=True, - ) diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py deleted file mode 100644 index bab3f0ad38..0000000000 --- a/lightllm/common/basemodel/attention/triton/fp.py +++ /dev/null @@ -1,342 +0,0 @@ -import dataclasses -import weakref -import torch -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional -from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers - - -class TritonAttBackend(BaseAttBackend): - def create_att_prefill_state(self, infer_state) -> "TritonPrefillAttState": - return TritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "TritonDecodeAttState": - return TritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class TritonPrefillAttState(BasePrefillAttState): - def init_state(self): - pass - - def prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - if att_control.use_alibi: - assert att_control.use_sliding_window is False, "alibi + sliding_window not supported" - assert att_control.tp_alibi is not None - return self._alibi_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - else: - return self._nomarl_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - - def _alibi_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - out = alloc_func(q.shape, q.dtype) - - from ...triton_kernel.alibi_att.context_flashattention_nopad import context_attention_fwd - - context_attention_fwd( - q, - k, - v, - out, - self.infer_state.b_req_idx, - att_control.tp_alibi, - self.infer_state.b_q_start_loc, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - self.infer_state.req_manager.req_to_token_indexs, - ) - return out - - def _nomarl_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd - - if att_control.use_sliding_window: - sliding_window = att_control.sliding_window - else: - sliding_window = (-1, -1) - - out = alloc_func(q.shape, q.dtype) - context_attention_fwd( - q, - k, - v, - out, - self.infer_state.b_req_idx, - self.infer_state.b_q_start_loc, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - self.infer_state.req_manager.req_to_token_indexs, - sliding_window=sliding_window, - ) - return out - - -@dataclasses.dataclass -class TritonDecodeAttState(BaseDecodeAttState): - b_mark_mtp_shared_group: torch.Tensor = None - decode_max_kv_seq_len: int = None - - def init_state(self): - # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于 GQA decode 配置查找。 - self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if draft_step > 0: - self.b_mark_mtp_shared_group = build_mtp_shared_group_markers( - self.infer_state.b_req_idx, - hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID, - ) - - def copy_for_decode_cuda_graph(self, new_state: "TritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - if att_control.use_alibi: - assert att_control.use_sliding_window is False, "alibi + sliding_window not supported" - assert att_control.tp_alibi is not None - return self._alibi_decode_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - else: - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - - q_head_num = q.shape[1] - k_head_num = k.shape[1] - - if draft_step > 0: - assert q_head_num >= k_head_num, "speculative decode requires q_head_num >= k_head_num" - return self._spec_decode_gqa_att(q=q, k=k, v=v, alloc_func=alloc_func) - elif q_head_num == k_head_num: - assert att_control.use_sliding_window is False, "sliding_window not supported in non-gqa attention yet" - return self._normal_decode_flash_decoding_att(q=q, k=k, v=v, alloc_func=alloc_func) - elif q_head_num > k_head_num: - return self._normal_decode_gqa_flash_decoding_att( - q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func - ) - else: - raise NotImplementedError("error") - - def _alibi_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - from ...triton_kernel.alibi_att.token_flashattention_nopad import token_attention_fwd - - out = alloc_func(q.shape, q.dtype) - token_attention_fwd( - q, - k, - v, - out, - att_control.tp_alibi, - self.infer_state.req_manager.req_to_token_indexs, - self.infer_state.b_req_idx, - self.infer_state.b_kv_start_loc, - self.infer_state.b_seq_len, - self.infer_state.max_kv_seq_len, - self.infer_state.total_token_num, - alloc_tensor_func=alloc_func, - ) - return out - - def _normal_decode_flash_decoding_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding import ( - token_decode_attention_flash_decoding, - ) - - out = alloc_func(q.shape, q.dtype) - - token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_v=v, - out=out, - alloc_tensor_func=alloc_func, - ) - return out - - def _normal_decode_gqa_flash_decoding_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( - gqa_token_decode_attention_flash_decoding, - ) - - if att_control.use_sliding_window: - sliding_window = att_control.sliding_window - else: - sliding_window = (-1, -1) - - out = alloc_func(q.shape, q.dtype) - - gqa_token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_v=v, - max_len_in_batch=self.decode_max_kv_seq_len, - out=out, - alloc_tensor_func=alloc_func, - sliding_window=sliding_window, - ) - - return out - - def _spec_decode_gqa_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.gqa.mtp_diverse import ( - token_decode_attention_mtp_diverse_single_token, - ) - - out = token_decode_attention_mtp_diverse_single_token( - q=q, - k=k, - v=v, - Req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - B_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - b_mark_shared_group=self.b_mark_mtp_shared_group, - max_kv_len=self.decode_max_kv_seq_len, - alloc_tensor_func=alloc_func, - ) - - return out - - def _normal_decode_gqa_flash_decoding_att_vsm( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_vsm import ( - gqa_token_decode_attention_flash_decoding_vsm, - ) - - out = alloc_func(q.shape, q.dtype) - - gqa_token_decode_attention_flash_decoding_vsm( - q=q, - k=k, - v=v, - infer_state=self.infer_state, - out=out, - alloc_tensor_func=alloc_func, - ) - return out - - def _normal_decode_gqa_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer_weight, - alloc_func=torch.empty, - ): - # TODO USE , 在特定场景下比 _normal_decode_gqa_flash_decoding_att 省显存 - from ...triton_kernel.att.decode_att.gqa.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd - - out = alloc_func(q.shape, q.dtype) - - gqa_decode_attention_fwd( - q=q, - k=k, - v=v, - out=out, - req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_seq_len=self.infer_state.b_seq_len, - ) - return out - - def _normal_decode_stage3_att( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - alloc_func=torch.empty, - ): - total_token_num = self.infer_state.total_token_num - batch_size = self.infer_state.batch_size - q_head_num = q.shape[1] - head_dim = q.shape[2] - - calcu_shape1 = (batch_size, q_head_num, head_dim) - att_m_tensor = alloc_func((q_head_num, total_token_num), torch.float32) - - from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_nopad_att1 import token_att_fwd - - token_att_fwd( - q.view(calcu_shape1), - k, - att_m_tensor, - Req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - B_req_idx=self.infer_state.b_req_idx, - B_Start_Loc=self.infer_state.b_kv_start_loc, - B_Seqlen=self.infer_state.b_seq_len, - max_len_in_batch=self.infer_state.max_kv_seq_len, - ) - - o_tensor = alloc_func(q.shape, q.dtype) - from ...triton_kernel.att.decode_att.mha.stage3_decode_att.token_attention_softmax_and_reducev import ( - token_softmax_reducev_fwd, - ) - - token_softmax_reducev_fwd( - att_m_tensor, - v, - o_tensor.view(calcu_shape1), - req_to_tokens=self.infer_state.req_manager.req_to_token_indexs, - b_req_idx=self.infer_state.b_req_idx, - b_start_loc=self.infer_state.b_kv_start_loc, - b_seq_len=self.infer_state.b_seq_len, - ) - return o_tensor diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py deleted file mode 100644 index 22e723de3b..0000000000 --- a/lightllm/common/basemodel/attention/triton/int4kv.py +++ /dev/null @@ -1,172 +0,0 @@ -import dataclasses -import weakref -import torch -from lightllm.utils.envs_utils import get_env_start_args -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, Tuple - - -class Int4kvTritonAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model) - self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size - - def create_att_prefill_state(self, infer_state) -> "Int4kvTritonPrefillAttState": - return Int4kvTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Int4kvTritonDecodeAttState": - return Int4kvTritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Int4kvTritonPrefillAttState(BasePrefillAttState): - - # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。 - b_kv_start_loc: torch.Tensor = None - - def init_state(self): - self.b_kv_start_loc = ( - torch.cumsum(self.infer_state.b_seq_len, dim=0, dtype=self.infer_state.b_seq_len.dtype) - - self.infer_state.b_seq_len - ) - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - - self.backend: Int4kvTritonAttBackend = self.backend # for typing - - k, k_scale = k - v, v_scale = v - o = self._groupsize_quant_prefill_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - return o - - def _groupsize_quant_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - # o_tensor = alloc_func(q.shape, q.dtype, device=q.device) - # batch_size = self.infer_state.b_seq_len.shape[0] - - total_token_num = self.infer_state.total_token_num - head_dim = k.shape[2] * 2 # 2个4bit存储为一个int8, 所以维度需要翻倍，才是解量化后的精度 - k_dequant = alloc_func((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device) - v_dequant = alloc_func((total_token_num, v.shape[1], head_dim), dtype=q.dtype, device=q.device) - o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) - - max_kv_seq_len = self.infer_state.max_kv_seq_len - - from ...triton_kernel.kv_copy.ppl_int4kv_copy_kv import dequantize_int4kv - - dequantize_int4kv( - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_seq_len=self.infer_state.b_seq_len, - b_req_idx=self.infer_state.b_req_idx, - b_kv_start_loc=self.b_kv_start_loc, - k_out=k_dequant, - v_out=v_dequant, - max_len_in_batch=max_kv_seq_len, - quant_group_size=self.backend.quant_group_size, - ) - - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv - - context_attention_fwd_contiguous_kv( - q=q, - k=k_dequant, - v=v_dequant, - o=o_tensor, - b_start_loc=self.infer_state.b_q_start_loc, - b_kv_start_loc=self.b_kv_start_loc, - b_seq_len=self.infer_state.b_seq_len, - max_q_input_len=self.infer_state.max_q_seq_len, - b_prompt_cache_len=self.infer_state.b_ready_cache_len, - ) - return o_tensor - - -@dataclasses.dataclass -class Int4kvTritonDecodeAttState(BaseDecodeAttState): - decode_max_kv_seq_len: int = None - - def init_state(self): - # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于配置查找。 - self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len - - def copy_for_decode_cuda_graph(self, new_state: "Int4kvTritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - k, k_scale = k - v, v_scale = v - - return self.ppl_int4kv_decode_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - - def ppl_int4kv_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - from ...triton_kernel.att.decode_att.int4kv.ppl_int4kv_flash_decoding import ( - token_decode_attention_flash_decoding, - ) - - return token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_k_scale=k_scale, - cache_v=v, - cache_v_scale=v_scale, - max_kv_seq_len=self.decode_max_kv_seq_len, - alloc_tensor_func=alloc_func, - ) diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py deleted file mode 100644 index 9da88626c8..0000000000 --- a/lightllm/common/basemodel/attention/triton/int8kv.py +++ /dev/null @@ -1,207 +0,0 @@ -import dataclasses -import weakref -import torch -from lightllm.utils.envs_utils import get_env_start_args -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, Tuple -from lightllm.utils.envs_utils import enable_diverse_mode_gqa_decode_fast_kernel -from lightllm.common.basemodel.triton_kernel.diverse_utils import build_diverse_shared_group_markers - - -class Int8kvTritonAttBackend(BaseAttBackend): - def __init__(self, model): - super().__init__(model) - self.quant_group_size: int = get_env_start_args().llm_kv_quant_group_size - - def create_att_prefill_state(self, infer_state) -> "Int8kvTritonPrefillAttState": - return Int8kvTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "Int8kvTritonDecodeAttState": - return Int8kvTritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class Int8kvTritonPrefillAttState(BasePrefillAttState): - - # 用于反量化的时候使用，可以减少反量化占用的显存数量。按需使用。 - b_kv_start_loc: torch.Tensor = None - - def init_state(self): - self.b_kv_start_loc = ( - torch.cumsum(self.infer_state.b_seq_len, dim=0, dtype=self.infer_state.b_seq_len.dtype) - - self.infer_state.b_seq_len - ) - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - - self.backend: Int8kvTritonAttBackend = self.backend # for typing - - k, k_scale = k - v, v_scale = v - o = self._groupsize_quant_prefill_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - return o - - def _groupsize_quant_prefill_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - # 反量化内核分别使用 K/V 及 scale 的指针和 stride，兼容共享或独立存储。 - total_token_num = self.infer_state.total_token_num - k_dequant = alloc_func((total_token_num, k.shape[1], k.shape[2]), dtype=q.dtype, device=q.device) - v_dequant = alloc_func((total_token_num, v.shape[1], v.shape[2]), dtype=q.dtype, device=q.device) - o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) - - max_kv_seq_len = self.infer_state.max_kv_seq_len - - from ...triton_kernel.kv_copy.ppl_int8kv_copy_kv import dequantize_int8kv - - dequantize_int8kv( - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - req_to_token_indexs=self.infer_state.req_manager.req_to_token_indexs, - b_seq_len=self.infer_state.b_seq_len, - b_req_idx=self.infer_state.b_req_idx, - b_kv_start_loc=self.b_kv_start_loc, - k_out=k_dequant, - v_out=v_dequant, - max_len_in_batch=max_kv_seq_len, - quant_group_size=self.backend.quant_group_size, - ) - - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv - - context_attention_fwd_contiguous_kv( - q=q, - k=k_dequant, - v=v_dequant, - o=o_tensor, - b_start_loc=self.infer_state.b_q_start_loc, - b_kv_start_loc=self.b_kv_start_loc, - b_seq_len=self.infer_state.b_seq_len, - max_q_input_len=self.infer_state.max_q_seq_len, - b_prompt_cache_len=self.infer_state.b_ready_cache_len, - ) - return o_tensor - - -@dataclasses.dataclass -class Int8kvTritonDecodeAttState(BaseDecodeAttState): - b_shared_seq_len: torch.Tensor = None - b_mark_shared_group: torch.Tensor = None - decode_max_kv_seq_len: int = None - - def init_state(self): - # Graph 捕获会改写 infer_state 的长度上限，提前保存真实长度用于普通 decode 配置查找。 - self.decode_max_kv_seq_len = self.infer_state.max_kv_seq_len - if enable_diverse_mode_gqa_decode_fast_kernel(): - self.b_mark_shared_group = build_diverse_shared_group_markers( - b_shared_radix_node_id=self.infer_state.b_shared_radix_node_id, - ) - # A one-row group has no cross-request prefix sharing to accelerate. - self.b_shared_seq_len = torch.where( - self.b_mark_shared_group == 1, - 0, - self.infer_state.b_shared_seq_len, - ) - - def copy_for_decode_cuda_graph(self, new_state: "Int8kvTritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: Tuple[torch.Tensor, torch.Tensor], - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - k, k_scale = k - v, v_scale = v - if enable_diverse_mode_gqa_decode_fast_kernel(): - return self.diverse_decode_att(q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func) - else: - return self.normal_decode_att( - q=q, - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - alloc_func=alloc_func, - ) - - def diverse_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ) -> torch.Tensor: - - from ...triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import ( - token_decode_attention_flash_decoding, - ) - - return token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_k_scale=k_scale, - cache_v=v, - cache_v_scale=v_scale, - alloc_tensor_func=alloc_func, - ) - - def normal_decode_att( - self, - q: torch.Tensor, - k: torch.Tensor, - k_scale: torch.Tensor, - v: torch.Tensor, - v_scale: torch.Tensor, - alloc_func=torch.empty, - ): - from ...triton_kernel.att.decode_att.int8kv.normal import token_decode_attention_flash_decoding - - return token_decode_attention_flash_decoding( - q=q, - infer_state=self.infer_state, - cache_k=k, - cache_k_scale=k_scale, - cache_v=v, - cache_v_scale=v_scale, - max_len_in_batch=self.decode_max_kv_seq_len, - alloc_tensor_func=alloc_func, - ) diff --git a/lightllm/common/basemodel/attention/triton/mla.py b/lightllm/common/basemodel/attention/triton/mla.py deleted file mode 100644 index 036e58ae70..0000000000 --- a/lightllm/common/basemodel/attention/triton/mla.py +++ /dev/null @@ -1,160 +0,0 @@ -import dataclasses -import weakref -import torch -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Tuple - - -class MlaTritonAttBackend(BaseAttBackend): - def create_att_prefill_state(self, infer_state) -> "MlaTritonPrefillAttState": - return MlaTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - def create_att_decode_state(self, infer_state) -> "MlaTritonDecodeAttState": - return MlaTritonDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class MlaTritonPrefillAttState(BasePrefillAttState): - def init_state(self): - pass - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - assert ( - att_control.use_alibi is False - and att_control.use_sliding_window is False - and att_control.use_att_sink is False - ) - if q.shape[2] == 256: - # glm model - return self._contiguous_kv_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - else: - return self._mla_prefill_att(q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func) - - def _mla_prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - from ...triton_kernel.mla_att.prefill_att import context_attention_fwd_with_v - - qk_rope_head_dim = 64 - q_nope, q_rope = q[:, :, :-qk_rope_head_dim], q[:, :, -qk_rope_head_dim:] - o_tensor = alloc_func((q_nope.shape[0], q_nope.shape[1], v.shape[-1]), dtype=q_nope.dtype, device=q.device) - k_nope, k_rope = k - assert att_control.mla_prefill - softmax_scale = att_control.mla_prefill_dict["softmax_scale"] - context_attention_fwd_with_v( - q_nope, - q_rope, - k_nope, - k_rope, - v, - o_tensor, - self.infer_state.b_q_start_loc, - self.infer_state.b1_cu_kv_seq_len, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - softmax_scale, - ) - return o_tensor - - def _contiguous_kv_prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - from ...triton_kernel.att.prefill_att.context_flashattention_nopad import context_attention_fwd_contiguous_kv - - k_nope, k_rope = k - q_head_num = q.shape[1] - - k_merged = torch.cat([k_nope.expand(-1, q_head_num, -1), k_rope.expand(-1, q_head_num, -1)], dim=-1) - - o_tensor = alloc_func(q.shape, dtype=q.dtype, device=q.device) - - context_attention_fwd_contiguous_kv( - q, - k_merged, - v, - o_tensor, - self.infer_state.b_q_start_loc, - self.infer_state.b1_cu_kv_seq_len, - self.infer_state.b_seq_len, - self.infer_state.max_q_seq_len, - self.infer_state.b_ready_cache_len, - ) - return o_tensor - - -@dataclasses.dataclass -class MlaTritonDecodeAttState(BaseDecodeAttState): - def init_state(self): - pass - - def copy_for_decode_cuda_graph(self, new_state: "MlaTritonDecodeAttState"): - super().copy_for_decode_cuda_graph(new_state) - - def decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ): - assert ( - att_control.use_sliding_window is False - and att_control.use_att_sink is False - and att_control.use_alibi is False - ) - assert v is None - - return self._mla_decode_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - ) - - def _mla_decode_att( - self, - q: Tuple[torch.Tensor, torch.Tensor], - k: torch.Tensor, - v: torch.Tensor, - att_control: AttControl, - alloc_func=torch.empty, - ): - assert att_control.mla_decode - softmax_scale = att_control.mla_decode_dict["softmax_scale"] - - from ...triton_kernel.mla_att.decode_att import gqa_token_decode_attention_flash_decoding - - qk_rope_head_dim = 64 - q_nope, q_rope = q - kv = k - - out = gqa_token_decode_attention_flash_decoding( - q_nope=q_nope, - q_rope=q_rope, - kv_nope=kv[:, :, :-qk_rope_head_dim], - kv_rope=kv[:, :, -qk_rope_head_dim:], - infer_state=self.infer_state, - softmax_scale=softmax_scale, - alloc_tensor_func=alloc_func, - ) - return out diff --git a/lightllm/common/basemodel/attention/triton/neo.py b/lightllm/common/basemodel/attention/triton/neo.py deleted file mode 100644 index dfa31876f0..0000000000 --- a/lightllm/common/basemodel/attention/triton/neo.py +++ /dev/null @@ -1,47 +0,0 @@ -import dataclasses -import weakref -import torch -from typing import Tuple - -from ..base_att import AttControl -from .fp import TritonAttBackend, TritonPrefillAttState, TritonDecodeAttState - - -class NeoTritonAttBackend(TritonAttBackend): - def create_att_prefill_state(self, infer_state) -> "NeoTritonPrefillAttState": - return NeoTritonPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - - -@dataclasses.dataclass -class NeoTritonPrefillAttState(TritonPrefillAttState): - def init_state(self): - pass - - def prefill_att( - self, - q: torch.Tensor, - k: Tuple[torch.Tensor, torch.Tensor], - v: torch.Tensor, - att_control: AttControl = AttControl(), - alloc_func=torch.empty, - ) -> torch.Tensor: - from ...triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo - from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo - - self.infer_state: NeoChatInferStateInfo - - out = alloc_func(q.shape, q.dtype) - context_attention_fwd_neo( - q, - k, - v, - out, - self.infer_state.b_req_idx, - self.infer_state.b_q_start_loc, - self.infer_state.b_seq_len, - self.infer_state.b_ready_cache_len, - self.infer_state.max_q_seq_len, - self.infer_state.req_manager.req_to_token_indexs, - self.infer_state.b_image_token_end, - ) - return out diff --git a/lightllm/common/basemodel/attention_adapters.py b/lightllm/common/basemodel/attention_adapters.py new file mode 100644 index 0000000000..94a0fa334a --- /dev/null +++ b/lightllm/common/basemodel/attention_adapters.py @@ -0,0 +1,186 @@ +from dataclasses import dataclass, replace + +from lightllm.common.attention.api import AttentionFamily, AttentionPhase, CacheLayout, LayerSpec +from lightllm.common.attention.api.inputs import GDNResources, KDAResources, MLAResources +from lightllm.common.attention.api.spec import AttentionFeature, CacheEncoding + + +@dataclass(frozen=True) +class AttentionDescription: + family: AttentionFamily = AttentionFamily.DENSE + features: frozenset[AttentionFeature] = frozenset() + window: tuple[int, int] | None = None + variant: str = "" + + +class AttentionAdapter: + def describe(self, model, layer, phase, description): + heads, kv_heads, head_dim, value_dim = self.dimensions(model, layer) + state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) + states = tuple( + state + for state in ( + *model.att_state_manager.get_full_attention_states(model.att_state_role), + *model.att_state_manager.get_bounded_attention_states(model.att_state_role), + ) + if state.contains_logical_layer(layer.layer_num_) + ) + features = set(description.features) + draft_step = model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model) + if draft_step > 0 and phase is AttentionPhase.DECODE: + features.add(AttentionFeature.MTP) + if draft_step > 0 and not model.is_mtp_draft_model and model.args.mtp_dynamic_verify: + features.add(AttentionFeature.DYNAMIC_VERIFY) + return LayerSpec( + family=description.family, + phase=phase, + dtype=model.data_type, + cache=self.cache(model, state, states, kv_heads, head_dim, value_dim), + num_q_heads=heads, + head_size=head_dim, + value_head_size=value_dim, + sliding_window=description.window, + features=frozenset(features), + variant=description.variant, + softmax_scale=self.scale(layer), + ) + + def cache(self, model, state, states, heads, head_dim, value_dim): + return CacheLayout( + encoding=CacheEncoding.from_option(model.args.llm_kv_type), + page_size=model.args.page_size, + dtype=state.storage_dtype, + key_heads=heads, + key_dim=head_dim, + value_dim=value_dim, + storage_schema=tuple(state.describe_layout() for state in states), + ) + + def scale(self, layer): + return None + + +class DenseAdapter(AttentionAdapter): + state_type = "gqa" + + def dimensions(self, model, layer): + return layer.tp_q_head_num_, layer.tp_k_head_num_, layer.head_dim_, layer.head_dim_ + + def cache(self, model, state, states, heads, head_dim, value_dim): + return replace( + super().cache(model, state, states, heads, head_dim, value_dim), + quant_group_size=getattr(state.config, "group_size", 0), + scale_mode=getattr(state.config, "scale_mode", "none"), + ) + + def bind(self, model, layer, weight): + state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) + resources = state.layer_resources(layer.layer_num_) + return replace( + resources, + alibi=getattr(weight, "tp_alibi", None), + sink=weight.attn_sinks.weight + if AttentionFeature.SINK in layer.describe_attention(AttentionPhase.PREFILL).features + else None, + ) + + +class MLAAdapter(AttentionAdapter): + state_type = "mla" + + def dimensions(self, model, layer): + return layer.tp_q_head_num_, 1, layer.qk_nope_head_dim + layer.qk_rope_head_dim, layer.kv_lora_rank + + def scale(self, layer): + return layer.softmax_scale + + def bind(self, model, layer, weight): + state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) + return MLAResources(state.layer_resources(layer.layer_num_), layer.kv_lora_rank, layer.qk_rope_head_dim) + + +class SparseMLAAdapter(MLAAdapter): + def dimensions(self, model, layer): + return ( + layer.tp_q_head_num_, + 1, + layer.qk_nope_head_dim + layer.network_config_.get("qk_rope_head_dim", 0), + layer.kv_lora_rank, + ) + + def bind(self, model, layer, weight): + state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) + return MLAResources( + state.layer_resources(layer.layer_num_), + layer.kv_lora_rank, + layer.network_config_.get("qk_rope_head_dim", 0), + ) + + +class ConvSSMAdapter(AttentionAdapter): + def cache(self, model, state, states, heads, head_dim, value_dim): + return CacheLayout( + encoding=CacheEncoding.PLAIN, + page_size=model.args.page_size, + dtype=state.config.conv_dtype, + key_heads=heads, + key_dim=head_dim, + value_dim=value_dim, + state_dtype=state.config.ssm_dtype, + state_mode=state.state_mode, + storage_schema=tuple(state.describe_layout() for state in states), + ) + + +class GDNAdapter(ConvSSMAdapter): + state_type = "gdn" + + def dimensions(self, model, layer): + config = model.config + heads = config["linear_num_key_heads"] // model.tp_world_size_ + return heads, heads, config["linear_key_head_dim"], config["linear_value_head_dim"] + + def bind(self, model, layer, weight): + state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) + storage = state.layer_resources(layer.layer_num_) + return GDNResources( + conv=storage.conv, + ssm=storage.ssm, + conv_weight=weight.linear_conv1d.mm_param.weight, + conv_bias=weight.linear_conv1d.bias, + a_log=weight.linear_A_log.weight, + dt_bias=weight.linear_dt_bias.weight, + mtp_state_index=storage.mtp_state_index, + replay_state=storage.replay_state, + ) + + +class KDAAdapter(ConvSSMAdapter): + state_type = "kda" + + def dimensions(self, model, layer): + config = model.config["linear_attn_config"] + heads = config["num_heads"] // model.tp_world_size_ + return heads, heads, config["head_dim"], config["head_dim"] + + def bind(self, model, layer, weight): + state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) + storage = state.layer_resources(layer.layer_num_) + return KDAResources( + conv=storage.conv, + ssm=storage.ssm, + conv_weight=weight.get_merged_kda_conv_weight(), + a_log=weight.linear_A_log.weight, + dt_bias=weight.linear_dt_bias.weight, + mtp_state_index=storage.mtp_state_index, + replay_state=storage.replay_state, + ) + + +ADAPTERS = { + AttentionFamily.DENSE: DenseAdapter(), + AttentionFamily.MLA: MLAAdapter(), + AttentionFamily.SPARSE_MLA: SparseMLAAdapter(), + AttentionFamily.GDN: GDNAdapter(), + AttentionFamily.KDA: KDAAdapter(), +} diff --git a/lightllm/common/basemodel/attention_vit/base_att.py b/lightllm/common/basemodel/attention_vit/base_att.py deleted file mode 100644 index 49bf6ad74c..0000000000 --- a/lightllm/common/basemodel/attention_vit/base_att.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -from abc import ABC, abstractmethod - - -class BaseVitAttBackend(ABC): - """ - 用于创建支持各种不同的AttBackend, 如 fa3, sdpa, triton 实现等， - 这个是单列模式, 每种backend只有一个实例 - """ - - _instances = {} - - def __new__(cls, *args, **kwargs): - """ - 重写__new__方法实现单例模式 - """ - # 检查是否已经有该类的实例 - if cls not in cls._instances: - # 创建新实例并存储 - instance = super().__new__(cls) - cls._instances[cls] = instance - # 返回已有的实例 - return cls._instances[cls] - - def __init__(self): - pass - - @abstractmethod - def _vit_att_fwd( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: - raise NotImplementedError("not impl") diff --git a/lightllm/common/basemodel/attention_vit/create_utils.py b/lightllm/common/basemodel/attention_vit/create_utils.py deleted file mode 100644 index 21eb9a7af9..0000000000 --- a/lightllm/common/basemodel/attention_vit/create_utils.py +++ /dev/null @@ -1,47 +0,0 @@ -from importlib import import_module - -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.backend_validator import _validate -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend - -logger = init_logger(__name__) - - -vit_att_backend = { - "triton": "triton.fp:TritonVitAttBackend", - "sdpa": "sdpa.fp:SdpaVitAttBackend", - "fa3": "fa3.fp:Fa3VitAttBackend", - "xformers": "xformers.fp:XformersVitAttBackend", -} - - -def get_vit_att_backend_class(backend_name: str) -> BaseVitAttBackend: - module_name, class_name = vit_att_backend[backend_name].split(":") - return getattr(import_module(f"{__package__}.{module_name}"), class_name) - - -def init_vit_att_backend(index=0, priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str: - args = get_env_start_args() - backend_name = args.vit_att_backend[index] - if backend_name != "auto": - logger.info(f"Selected {backend_name} backend for VIT") - return backend_name - return _select_vit_backend(priority_list=priority_list) - - -def _select_vit_backend(priority_list: list = ["fa3", "xformers", "sdpa", "triton"]) -> str: - """Auto-select the best available backend with validation for VIT. - - Priority: FA3 > Xformers > Sdpa > Triton - Each backend is validated in a subprocess with ground truth checks. - """ - - for backend_name in priority_list: - if _validate(backend_name): - logger.info(f"Auto-selected {backend_name} backend (validated) for VIT") - return backend_name - - # Fallback to triton without validation (should not happen) - logger.warning("No backend validation succeeded, falling back to triton") - return "triton" diff --git a/lightllm/common/basemodel/attention_vit/fa3/fp.py b/lightllm/common/basemodel/attention_vit/fa3/fp.py deleted file mode 100644 index a8e56417c5..0000000000 --- a/lightllm/common/basemodel/attention_vit/fa3/fp.py +++ /dev/null @@ -1,60 +0,0 @@ -import dataclasses -import torch -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend - - -class Fa3VitAttBackend(BaseVitAttBackend): - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> None: - - head_dim = q.shape[-1] - softmax_scale = head_dim ** -0.5 - window_size = (-1, -1) - torch.ops.sgl_kernel.fwd.default( - q, - k, - v, - None, # k_new - None, # v_new - None, # qv - o, # out - cu_seqlens, - cu_seqlens, - None, # cu_seqlens_k_new - None, - None, - max_seqlen, - max_seqlen, - None, # page_table, - None, # kv_batch_idx - None, # leftpad_k - None, # rotary cos - None, # rotary sin - None, # seqlens_rotary - None, - None, - None, - softmax_scale, - False, - window_size[0], - window_size[1], - attention_chunk=0, - softcap=0.0, - is_rotary_interleaved=False, - scheduler_metadata=None, - num_splits=1, - pack_gqa=None, - sm_margin=0, - sinks=None, - sparse_mask_fine=None, - only_qv=False, - ) - - return o diff --git a/lightllm/common/basemodel/attention_vit/triton/fp.py b/lightllm/common/basemodel/attention_vit/triton/fp.py deleted file mode 100644 index 88867102a4..0000000000 --- a/lightllm/common/basemodel/attention_vit/triton/fp.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.models.vit.triton_kernel.flashattention_nopad import _flash_attention_triton_fwd - - -class TritonVitAttBackend(BaseVitAttBackend): - @staticmethod - def _vit_att_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ): - _flash_attention_triton_fwd( - q, - k, - v, - o, - cu_seqlens, # q k v cu_seqlens, - max_seqlen, - ) - return o diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 0edbdadb94..d2241207d2 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1,56 +1,55 @@ -import os - -# os.environ["CUDA_LAUNCH_BLOCKING"] = "1" -import gc +from contextlib import ExitStack +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.selection import BackendPolicy import copy -import math +import gc import json +import math +import os +from functools import cached_property +from typing import TYPE_CHECKING, List, final + import torch import torch.nn.functional as F import triton -from typing import final, List, TYPE_CHECKING from tqdm import tqdm -from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.req_manager import ReqManager from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.commit import RequestGenerationStamp +from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput, PostLayerOutput +from lightllm.common.basemodel.cuda_graph import CudaGraph +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights +from lightllm.common.basemodel.mtp_manager import MtpManager +from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( select_kv_index_from_req, select_kv_index_from_req_prefill, ) -from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph -from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed -from lightllm.utils.log_utils import init_logger -from lightllm.utils.dist_utils import get_dp_world_size, get_current_rank_in_dp -from lightllm.utils.envs_utils import ( - get_env_start_args, - get_llm_data_type, -) +from lightllm.common.quantization import Quantcfg +from lightllm.common.req_manager import ReqManager +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner from lightllm.distributed.communication_op import dist_group_manager -from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput, PostLayerOutput -from lightllm.common.basemodel.hidden_collector import ( - NoopHiddenCollector, -) -from lightllm.common.basemodel.mtp_manager import MtpManager +from lightllm.platform import get_device_runtime from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type from lightllm.utils.infer_utils import post_empty_cache -from lightllm.utils.torch_memory_saver_utils import ( - TorchMemorySaverWrapper, - MemoryTag, -) -from .attention import get_prefill_att_backend_class, get_decode_att_backend_class -from .attention import BaseAttBackend +from lightllm.utils.log_utils import init_logger +from lightllm.utils.torch_memory_saver_utils import MemoryTag, TorchMemorySaverWrapper if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs + +# os.environ["CUDA_LAUNCH_BLOCKING"] = "1" logger = init_logger(__name__) @@ -58,6 +57,17 @@ class TpPartBaseModel: + attention = None + + @cached_property + def device_runtime(self): + return get_device_runtime(platform=self.args.hardware_platform) + + def prepare_attention_backends(self): + from lightllm.common.basemodel.attention import ModelAttention + + self.attention = ModelAttention(self) + is_mtp_draft_model = False # weight class @@ -125,8 +135,6 @@ def __init__(self, kvargs): @property def att_state_role(self) -> str: """当前模型的资源前缀；每个 draft 独立拥有 State。""" - from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS - if not self.is_mtp_draft_model: return "main" draft_index = len(self.mtp_previous_draft_models) @@ -153,6 +161,7 @@ def prepare_attention_states(self) -> None: with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): if not self.is_mtp_draft_model: self.att_state_manager = AttentionStateManager(request_capacity=self.max_req_num + 1) + self.att_state_manager.device_runtime = self.device_runtime # 无 attention 的 draft 只共享管理器，不创建自身 State。 if self.is_mtp_draft_model and self.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): return @@ -173,7 +182,7 @@ def prepare_attention_states(self) -> None: request_capacity=self.max_req_num + 1, tp_rank=self.tp_rank_, tp_world_size=self.tp_world_size_, - device="cuda", + device=self.device_runtime.device, ) for state in (*full_states, *bounded_states): self.att_state_manager.register_state(self.att_state_role, state) @@ -182,20 +191,18 @@ def init_runtime(self) -> None: """Initialize this model's backends, warmup and graphs after shared allocation.""" self.req_manager.bind_hold_page(self.att_state_manager.full_att_allocator.hold_token_indices) self._check_mem_size() - self._init_att_backend() - self._init_att_backend1() + self.attention.bind(self) self._init_hidden_collector() role = self.att_state_role - logger.info( - f"{role} prefill/decode attention: {type(self.prefill_att_backend).__name__} / " - f"{type(self.decode_att_backend).__name__}" - ) - if self.prefill_att_backend1 is not None: + for group_id, group in enumerate(self.attention.groups): logger.info( - f"{role} prefill/decode attention1: {type(self.prefill_att_backend1).__name__} / " - f"{type(self.decode_att_backend1).__name__}" + "%s attention group %d (%s): prefill=%s, decode=%s", + role, + group_id, + group[AttentionPhase.PREFILL].spec.family.value, + type(group[AttentionPhase.PREFILL]).__name__, + type(group[AttentionPhase.DECODE]).__name__, ) - self._autotune_warmup() self._init_padded_req() self._init_cudagraph() @@ -250,6 +257,11 @@ def _init_weights(self, start_layer_index=0): def load_weights(self, weight_dict: dict): assert weight_dict is None or isinstance(weight_dict, dict), "weight_dict must be a dict or None" + if self.attention is not None and self.attention.generation is not None: + if self.graph is not None: + self.graph.close() + if self.prefill_graph is not None: + self.prefill_graph.close() load_hf_weights( data_type=self.data_type, weight_dir=self.weight_dir_, @@ -259,6 +271,8 @@ def load_weights(self, weight_dict: dict): ) self.pre_post_weight.verify_load() [weight.verify_load() for weight in self.trans_layers_weight] + if self.attention is not None and self.attention.generation is not None: + self.attention.bind(self) return @classmethod @@ -336,21 +350,8 @@ def _init_some_value(self): self.vocab_size = self.config["vocab_size"] return - def _init_att_backend(self): - role = self.att_state_role - if not self.att_state_manager.get_full_attention_states(role): - # FFN-only and linear-only models do not need a full-attention backend. - self.prefill_att_backend = self.decode_att_backend = None - return - self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0)(model=self) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) - return - - def _init_att_backend1(self): - # self.prefill_att_backend1 是给后续有模型支持不同层用不同的att模块时，保留的扩展。 - self.prefill_att_backend1: BaseAttBackend = None - self.decode_att_backend1: BaseAttBackend = None - return + def attention_policy(self, family, phase): + return BackendPolicy() def _align_decode_batch_size(self, batch_size: int) -> int: alignment = self.mtp_manager.get_decode_batch_alignment(self.is_mtp_draft_model) @@ -369,6 +370,7 @@ def _init_cudagraph(self): None if self.disable_cudagraph else CudaGraph( + runtime=self.device_runtime, batch_step_size_before_split=cuda_graph_grow_step_size, split_batch_size=self.args.graph_split_batch_size * decode_tokens_per_request, batch_step_size_after_split=self.args.graph_grow_step_size * cuda_graph_grow_step_size, @@ -388,7 +390,9 @@ def _init_prefill_cuda_graph(self): self.prefill_graph = ( None if not get_env_start_args().enable_prefill_cudagraph - else PrefillCudaGraph(decode_cuda_graph=self.graph, tp_world_size=self.tp_world_size_) + else PrefillCudaGraph( + decode_cuda_graph=self.graph, tp_world_size=self.tp_world_size_, runtime=self.device_runtime + ) ) if self.prefill_graph is not None: if get_env_start_args().enable_prefill_microbatch_overlap: @@ -467,26 +471,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) # 特殊模型，特殊模式的特定变量初始化操作。 infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens - if infer_state.is_prefill: - infer_state.prefill_att_state = ( - self.prefill_att_backend.create_att_prefill_state(infer_state=infer_state) - if self.prefill_att_backend is not None - else None - ) - if self.prefill_att_backend1 is not None: - infer_state.prefill_att_state1 = self.prefill_att_backend1.create_att_prefill_state( - infer_state=infer_state - ) - else: - infer_state.decode_att_state = ( - self.decode_att_backend.create_att_decode_state(infer_state=infer_state) - if self.decode_att_backend is not None - else None - ) - if self.decode_att_backend1 is not None: - infer_state.decode_att_state1 = self.decode_att_backend1.create_att_decode_state( - infer_state=infer_state - ) + infer_state.device_runtime = self.device_runtime return infer_state @@ -585,6 +570,19 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS ) return output + def _attach_attention_commit(self, output, infer_state, batch_size): + if not self.args.mtp_step: + return output + + output = copy.copy(output) + output.attention_commit = RequestGenerationStamp.stamp( + self.att_state_role, + infer_state.b_req_idx[:batch_size], + infer_state.b_mtp_index[:batch_size], + self.req_manager.request_generations, + ) + return output + def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int): padded_batch_size = model_output.logits.shape[0] if padded_batch_size == origin_batch_size: @@ -652,18 +650,17 @@ def _prefill( infer_state = self._create_inferstate(model_input) prefill_mem_indexes_ready_event = torch.cuda.Event() prefill_mem_indexes_ready_event.record() - infer_state.init_some_extra_state(self) - infer_state.init_att_state() - model_output = self._context_forward(infer_state=infer_state) + with self.attention.session(infer_state): + model_output = self._context_forward(infer_state=infer_state) - model_output = self._create_unpad_prefill_model_output( - padded_model_output=model_output, - origin_handle_token_num=origin_handle_token_num, - origin_batch_size=origin_batch_size, - ) - model_output.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event - return model_output + model_output = self._create_unpad_prefill_model_output( + padded_model_output=model_output, + origin_handle_token_num=origin_handle_token_num, + origin_batch_size=origin_batch_size, + ) + model_output.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event + return model_output def _decode( self, @@ -703,26 +700,26 @@ def _decode( model_input = self._create_padded_decode_model_input(model_input=model_input, new_batch_size=infer_batch_size) infer_state = self._create_inferstate(model_input) - # attention backend 会根据执行方式准备 capture、replay 或 eager 状态， - # 因此必须在 init_att_state 之前完成赋值。 infer_state.is_cuda_graph_capture = need_capture infer_state.is_cuda_graph_replay = use_cuda_graph and not need_capture infer_state.init_some_extra_state(self) - infer_state.init_att_state() - - if use_cuda_graph: - if need_capture: - model_output: ModelOutput = self.graph.capture_decode(self._token_forward, infer_state) + with self.attention.session(infer_state, max_kv_seq_len=self.graph_max_len_in_batch if need_capture else None): + if use_cuda_graph: + if need_capture: + model_output: ModelOutput = self.graph.capture_decode(self._token_forward, infer_state) + else: + model_output: ModelOutput = self.graph.replay(infer_state) else: - model_output: ModelOutput = self.graph.replay(infer_state) - else: - model_output = self._token_forward(infer_state) + model_output = self._token_forward(infer_state) - return self._create_unpad_decode_model_output(model_output, origin_batch_size=origin_batch_size) + output = self._create_unpad_decode_model_output(model_output, origin_batch_size=origin_batch_size) + output = self._attach_attention_commit(output, infer_state, origin_batch_size) + + return output @final def _context_forward(self, infer_state: InferStateInfo): - + infer_state.attention.build_device_metadata() input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" @@ -788,6 +785,7 @@ def prefill_func(input_tensors, _infer_state): return model_output def _token_forward(self, infer_state: InferStateInfo): + infer_state.attention.build_device_metadata() hidden_collector = infer_state.hidden_collector input_ids = infer_state.input_ids cuda_input_ids = input_ids @@ -817,7 +815,6 @@ def _token_forward(self, infer_state: InferStateInfo): @torch.no_grad() def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: ModelInput): """执行由调用方提前构建好的两个 prefill microbatch。""" - for model_input in (model_input0, model_input1): model_input.to_cuda() if self.args.enable_prefill_decode_mixed and model_input.input_ids.shape[0] > 0: @@ -853,40 +850,40 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input model_input=model_input1, new_handle_token_num=infer_handle_token_num1 ) - infer_state0 = self._create_inferstate(model_input0, 0) - infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() + with ExitStack() as attention_scopes: + infer_state0 = self._create_inferstate(model_input0, 0) + infer_state0.init_some_extra_state(self) + attention_scopes.enter_context(self.attention.session(infer_state0)) - infer_state1 = self._create_inferstate(model_input1, 1) - infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() + infer_state1 = self._create_inferstate(model_input1, 1) + infer_state1.init_some_extra_state(self) + attention_scopes.enter_context(self.attention.session(infer_state1)) - prefill_mem_indexes_ready_event = torch.cuda.Event() - prefill_mem_indexes_ready_event.record() + prefill_mem_indexes_ready_event = torch.cuda.Event() + prefill_mem_indexes_ready_event.record() - model_output0, model_output1 = self._overlap_tpsp_context_forward(infer_state0, infer_state1=infer_state1) + model_output0, model_output1 = self._overlap_tpsp_context_forward(infer_state0, infer_state1=infer_state1) - model_output0 = self._create_unpad_prefill_model_output( - padded_model_output=model_output0, - origin_handle_token_num=origin_handle_token_num0, - origin_batch_size=origin_batch_size0, - ) - model_output1 = self._create_unpad_prefill_model_output( - padded_model_output=model_output1, - origin_handle_token_num=origin_handle_token_num1, - origin_batch_size=origin_batch_size1, - ) - # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候 - # 该调用没有实际意义 - dist_group_manager.clear_deepep_buffer() - model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event - model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event - return model_output0, model_output1 + model_output0 = self._create_unpad_prefill_model_output( + padded_model_output=model_output0, + origin_handle_token_num=origin_handle_token_num0, + origin_batch_size=origin_batch_size0, + ) + model_output1 = self._create_unpad_prefill_model_output( + padded_model_output=model_output1, + origin_handle_token_num=origin_handle_token_num1, + origin_batch_size=origin_batch_size1, + ) + # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候 + # 该调用没有实际意义 + dist_group_manager.clear_deepep_buffer() + model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event + model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event + return model_output0, model_output1 @torch.no_grad() def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: ModelInput): """执行由调用方提前构建好的两个 decode microbatch。""" - for model_input in (model_input0, model_input1): model_input.to_cuda() if model_input.input_ids is None: @@ -912,56 +909,77 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 max_len_in_batch = max(2, model_input0.max_kv_seq_len, model_input1.max_kv_seq_len) infer_batch_size = self._align_decode_batch_size(max(1, origin_batch_size0, origin_batch_size1)) - if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): - infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size) - need_capture = self.graph.need_capture(infer_batch_size) - padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) - padded_model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) - infer_state0 = self._create_inferstate(padded_model_input0, 0) - infer_state0.is_cuda_graph_capture = need_capture - infer_state0.is_cuda_graph_replay = not need_capture - infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() + with ExitStack() as attention_scopes: + if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): + infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size) + need_capture = self.graph.need_capture(infer_batch_size) + padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) + padded_model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) + infer_state0 = self._create_inferstate(padded_model_input0, 0) + infer_state0.is_cuda_graph_capture = need_capture + infer_state0.is_cuda_graph_replay = not need_capture + infer_state0.init_some_extra_state(self) + attention_scopes.enter_context( + self.attention.session( + infer_state0, max_kv_seq_len=self.graph_max_len_in_batch if need_capture else None + ) + ) - infer_state1 = self._create_inferstate(padded_model_input1, 1) - infer_state1.is_cuda_graph_capture = need_capture - infer_state1.is_cuda_graph_replay = not need_capture - infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() + infer_state1 = self._create_inferstate(padded_model_input1, 1) + infer_state1.is_cuda_graph_capture = need_capture + infer_state1.is_cuda_graph_replay = not need_capture + infer_state1.init_some_extra_state(self) + attention_scopes.enter_context( + self.attention.session( + infer_state1, max_kv_seq_len=self.graph_max_len_in_batch if need_capture else None + ) + ) - if need_capture: - model_output0, model_output1 = self.graph.capture_decode( - self._overlap_tpsp_token_forward, - infer_state0, - infer_state1=infer_state1, + if need_capture: + model_output0, model_output1 = self.graph.capture_decode( + self._overlap_tpsp_token_forward, + infer_state0, + infer_state1=infer_state1, + ) + else: + model_output0, model_output1 = self.graph.replay( + infer_state0, + infer_state1=infer_state1, + ) + + model_output0 = self._create_unpad_decode_model_output( + model_output0, origin_batch_size=origin_batch_size0 + ) + model_output1 = self._create_unpad_decode_model_output( + model_output1, origin_batch_size=origin_batch_size1 ) else: - model_output0, model_output1 = self.graph.replay( - infer_state0, - infer_state1=infer_state1, + model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) + model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) + infer_state0 = self._create_inferstate(model_input0, 0) + infer_state0.init_some_extra_state(self) + attention_scopes.enter_context(self.attention.session(infer_state0)) + + infer_state1 = self._create_inferstate(model_input1, 1) + infer_state1.init_some_extra_state(self) + attention_scopes.enter_context(self.attention.session(infer_state1)) + + model_output0, model_output1 = self._overlap_tpsp_token_forward(infer_state0, infer_state1=infer_state1) + model_output0 = self._create_unpad_decode_model_output( + model_output0, origin_batch_size=origin_batch_size0 + ) + model_output1 = self._create_unpad_decode_model_output( + model_output1, origin_batch_size=origin_batch_size1 ) - model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size0) - model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size1) - else: - model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) - model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) - infer_state0 = self._create_inferstate(model_input0, 0) - infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() - - infer_state1 = self._create_inferstate(model_input1, 1) - infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() - - model_output0, model_output1 = self._overlap_tpsp_token_forward(infer_state0, infer_state1=infer_state1) - model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size0) - model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size1) - - return model_output0, model_output1 + model_output0 = self._attach_attention_commit(model_output0, infer_state0, origin_batch_size0) + model_output1 = self._attach_attention_commit(model_output1, infer_state1, origin_batch_size1) + return model_output0, model_output1 @final def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state1: InferStateInfo): + infer_state.attention.build_device_metadata() + infer_state1.attention.build_device_metadata() g_cache_manager.cache_env_in() hidden_collector0 = infer_state.hidden_collector hidden_collector1 = infer_state1.hidden_collector @@ -1019,6 +1037,8 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state @final def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: InferStateInfo): + infer_state.attention.build_device_metadata() + infer_state1.attention.build_device_metadata() hidden_collector0 = infer_state.hidden_collector hidden_collector1 = infer_state1.hidden_collector input_embs, input_embs1 = self.pre_infer.overlap_tpsp_token_forward( diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d8f2329b24..92e84f84e8 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -4,6 +4,7 @@ import torch +from lightllm.common.attention_state.commit import RequestGenerationStamp from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor @@ -189,6 +190,7 @@ class ModelOutput: # attached here. ModelOutput therefore owns a stable output view instead # of the mutable collector used while the forward is still running. mtp_collector: Optional[ModelMtpOutputCollector] = None + attention_commit: Optional[RequestGenerationStamp] = None # 用于判断 mem_indexes 是否成功写入 req manager 中的事件对象。 prefill_mem_indexes_ready_event: torch.Event = None diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ab0e0c0721..38b7a3fe7f 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,22 +1,21 @@ -import os -import torch -import torch.distributed as dist -import copy import bisect +import copy import math -import triton from typing import Optional -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.distributed import dist_group_manager + +import torch +import torch.distributed as dist +import triton + from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -from lightllm.utils.torch_memory_saver_utils import ( - TorchMemorySaverWrapper, - MemoryTag, -) -from .infer_struct import InferStateInfo +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner +from lightllm.platform import get_device_runtime +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger +from lightllm.utils.torch_memory_saver_utils import TorchMemorySaverWrapper +from .graph import CapturedGraph +from .infer_struct import InferStateInfo logger = init_logger(__name__) @@ -63,10 +62,12 @@ def __init__( max_len_in_batch=8192, tp_world_size: int = 1, capture_infer_cost: bool = False, + runtime=None, ): self.graph = {} self.tp_world_size = tp_world_size - self.mempool = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None + self.runtime = runtime if runtime is not None else get_device_runtime() + self.mempool = self.runtime.graph_pool() self.args = get_env_start_args() self.max_batch_size = max_batch_size self.graph_max_len_in_batch = max_len_in_batch @@ -90,7 +91,7 @@ def can_run(self, batch_size, max_len_in_batch): def need_capture(self, batch_size): find_batch_size = self.find_closest_graph_batch_size(batch_size) if find_batch_size is not None: - return find_batch_size not in self.graph + return find_batch_size not in self.graph or self.graph[find_batch_size].graph is None else: assert False, "dead code" @@ -103,37 +104,40 @@ def find_closest_graph_batch_size(self, batch_size): return None def _capture_decode(self, decode_func, infer_state: InferStateInfo): - graph_obj = torch.cuda.CUDAGraph() + graph_obj = self.runtime.create_graph() input_ids = infer_state.input_ids batch_size = input_ids.shape[0] infer_state.max_kv_seq_len = self.graph_max_len_in_batch infer_state.total_token_num = self.graph_max_len_in_batch * batch_size - # warmup - # 因为有些推理过程的代码，会通过判断infer_state中是否存在某些属性来在一层上 - # 做一些初始化的操作，后续层可以复用这些计算的结果，如 - # lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding.py - # 中做的一些操作，所以在 warmup 的时候，需要调用infer_state的copy函数做一个 - # 浅拷贝，不然后续传入到cuda graph捕获过程中后，infer_state因为提前拥有了这些属性， - # 导致不会重新初始化，这样捕获过程中会不能捕获这些临时添加到 infer_state 管理对象 - # 中的 tensor。 - - for _ in range(1): - # 记录原始存在的变量 - pure_para_set = set(vars(infer_state).keys()) - torch.cuda.synchronize() - # 在正式捕获前调优 decode attention，退出作用域后再捕获选定的配置。 + snapshot = infer_state.att_state_manager.snapshot_capture_state( + infer_state.b_req_idx, infer_state.att_state_role + ) + restore_state = lambda: infer_state.att_state_manager.restore_capture_state(snapshot) + try: + self.runtime.synchronize() with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): decode_func(copy.copy(infer_state)) - torch.cuda.synchronize() - for param_name in set(vars(infer_state).keys()): - if param_name not in pure_para_set: - delattr(infer_state, param_name) - - with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool): - model_output = decode_func(infer_state) - self.graph[batch_size] = (graph_obj, infer_state, model_output) - graph_obj.replay() - self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size) + self.runtime.synchronize() + finally: + restore_state() + + try: + with self.runtime.capture_graph(graph_obj, pool=self.mempool, memory_saver=self.torch_memory_saver): + model_output = decode_func(infer_state) + except BaseException: + try: + CapturedGraph(graph_obj, (infer_state,), runtime=self.runtime).close() + finally: + restore_state() + raise + try: + execution = CapturedGraph(graph_obj, (infer_state,), (model_output,), runtime=self.runtime) + except BaseException: + restore_state() + raise + self.graph[batch_size] = execution + execution.replay() + self._measure_replay_cost(graph_obj=execution, batch_size=batch_size, restore_state=restore_state) return model_output def _capture_decode_overlap( @@ -142,50 +146,58 @@ def _capture_decode_overlap( infer_state: InferStateInfo, infer_state1: InferStateInfo, ): - graph_obj = torch.cuda.CUDAGraph() + graph_obj = self.runtime.create_graph() input_ids = infer_state.input_ids batch_size = input_ids.shape[0] infer_state.max_kv_seq_len = self.graph_max_len_in_batch infer_state.total_token_num = self.graph_max_len_in_batch * batch_size infer_state1.max_kv_seq_len = self.graph_max_len_in_batch infer_state1.total_token_num = self.graph_max_len_in_batch * batch_size - # warmup - for _ in range(1): - # 记录原始存在的变量 - pure_para_set = set(vars(infer_state).keys()) - pure_para_set1 = set(vars(infer_state1).keys()) - torch.cuda.synchronize() + snapshot = infer_state.att_state_manager.snapshot_capture_state( + torch.cat((infer_state.b_req_idx, infer_state1.b_req_idx)), infer_state.att_state_role + ) + restore_state = lambda: infer_state.att_state_manager.restore_capture_state(snapshot) + try: + self.runtime.synchronize() with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): decode_func(copy.copy(infer_state), copy.copy(infer_state1)) - torch.cuda.synchronize() - for para_name in set(vars(infer_state).keys()): - if para_name not in pure_para_set: - delattr(infer_state, para_name) - for para_name in set(vars(infer_state1).keys()): - if para_name not in pure_para_set1: - delattr(infer_state1, para_name) - - with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool): - model_output, model_output1 = decode_func(infer_state, infer_state1) - self.graph[batch_size] = ( - graph_obj, - infer_state, - infer_state1, - model_output, - model_output1, - ) - graph_obj.replay() - self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size) + self.runtime.synchronize() + finally: + restore_state() + + try: + with self.runtime.capture_graph(graph_obj, pool=self.mempool, memory_saver=self.torch_memory_saver): + model_output, model_output1 = decode_func(infer_state, infer_state1) + except BaseException: + try: + CapturedGraph(graph_obj, (infer_state, infer_state1), runtime=self.runtime).close() + finally: + restore_state() + raise + try: + execution = CapturedGraph( + graph_obj, (infer_state, infer_state1), (model_output, model_output1), runtime=self.runtime + ) + except BaseException: + restore_state() + raise + self.graph[batch_size] = execution + execution.replay() + self._measure_replay_cost(graph_obj=execution, batch_size=batch_size, restore_state=restore_state) return model_output, model_output1 - def _measure_replay_cost(self, graph_obj: torch.cuda.CUDAGraph, batch_size: int) -> None: + def _measure_replay_cost( + self, graph_obj: torch.cuda.CUDAGraph, batch_size: int, restore_state=lambda: None + ) -> None: if not self.capture_infer_cost: return dist.barrier(group=dist.group.WORLD) - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) + start_event = self.runtime.create_event(enable_timing=True) + end_event = self.runtime.create_event(enable_timing=True) + restore_state() graph_obj.replay() + restore_state() start_event.record() graph_obj.replay() end_event.record() @@ -217,37 +229,17 @@ def capture_decode( assert infer_state1 is None return self._capture_decode(decode_func, infer_state) - def _replay(self, infer_state: InferStateInfo): - batch_size = infer_state.input_ids.shape[0] - graph_obj, graph_infer_state, graph_output = self.graph[batch_size] - graph_infer_state.copy_for_cuda_graph(infer_state) - graph_obj.replay() - return graph_output - - def _replay_overlap( - self, - infer_state: InferStateInfo, - infer_state1: InferStateInfo, - ): - batch_size = infer_state.input_ids.shape[0] - ( - graph_obj, - graph_infer_state, - graph_infer_state1, - graph_model_output, - graph_model_output1, - ) = self.graph[batch_size] - graph_infer_state.copy_for_cuda_graph(infer_state) - graph_infer_state1.copy_for_cuda_graph(infer_state1) - graph_obj.replay() - return graph_model_output, graph_model_output1 - def replay(self, infer_state, infer_state1=None): - if self.enable_decode_microbatch_overlap: - return self._replay_overlap(infer_state, infer_state1) - else: - assert infer_state1 is None - return self._replay(infer_state) + batch_size = infer_state.input_ids.shape[0] + states = (infer_state, infer_state1) if self.enable_decode_microbatch_overlap else (infer_state,) + if any(state is None for state in states): + raise ValueError("Graph replay requires all microbatch lanes") + return self.graph[batch_size].replay(*states) + + def close(self): + for execution in self.graph.values(): + execution.close() + self.graph.clear() @torch.no_grad() def warmup(self, model): diff --git a/lightllm/common/basemodel/graph.py b/lightllm/common/basemodel/graph.py new file mode 100644 index 0000000000..0ceb472dde --- /dev/null +++ b/lightllm/common/basemodel/graph.py @@ -0,0 +1,76 @@ +from contextlib import ExitStack + +"""持有一次捕获的输入、provider 任务和完成事件；bucket 策略仍由模型决定。""" + +from lightllm.platform import get_device_runtime + + +class CapturedGraph: + def __init__(self, graph, states=(), outputs=(), *, runtime=None): + self.graph = graph + self.infer_states = tuple(states) + self.outputs = tuple(outputs) + self.runtime = runtime if runtime is not None else get_device_runtime(self.infer_states[0].input_ids.device) + self.hooks = [] + self.workspaces = [] + self.generations = [] + try: + for state in self.infer_states: + manager = state.att_state_manager + self.generations.append((manager, manager.storage_generation)) + for group_id, context in state.attention.contexts.items(): + backend = state.attention.groups[group_id][context.capacity.phase] + binding = backend.create_capture_hooks(context) + if binding is not None: + self.hooks.append(binding) + if context.workspace not in self.workspaces: + self.workspaces.append(context.workspace) + manager.register_graph(self) + for workspace in self.workspaces: + workspace.freeze() + except BaseException: + self.close() + raise + + def replay(self, *states): + if self.graph is None: + raise RuntimeError("Graph execution is closed") + if states and len(states) != len(self.infer_states): + raise ValueError("Graph execution lane count changed") + for manager, generation in self.generations: + if manager.storage_generation != generation: + raise RuntimeError("attention storage changed; rebind layers and recapture Graph") + copies = [ + pair for target, source in zip(self.infer_states, states) for pair in target.graph_input_copies(source) + ] + for binding in self.hooks: + binding.wait_for_update() + events = [] + with ExitStack() as leases: + for workspace in self.workspaces: + leases.enter_context(workspace.lease(events)) + for target, source in copies: + target.copy_(source, non_blocking=True) + for target, source in zip(self.infer_states, states): + target.attention.plan_replay(source.attention) + for hooks in self.hooks: + hooks.before_replay() + self.graph.replay() + for hooks in self.hooks: + events.extend(hooks.after_submit()) + return self.outputs[0] if len(self.outputs) == 1 else self.outputs + + def close(self): + if self.graph is None: + return + for binding in self.hooks: + binding.wait_for_update() + self.runtime.synchronize() + for binding in self.hooks: + binding.close() + self.graph.reset() + self.graph = None + self.hooks.clear() + self.workspaces.clear() + self.infer_states = self.outputs = () + self.generations.clear() diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 83b312f991..339c1440ac 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,17 +1,18 @@ -import torch -import triton import collections +from contextlib import contextmanager +from typing import List, Optional + +import torch + from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup -from typing import Tuple, Any, Optional, List -from .triton_kernel.gen_prefill_params import gen_prefill_params -from .triton_kernel.gen_decode_params import gen_decode_params -from .triton_kernel.multimodal_emb import mark_multimodal_obj -from .batch_objs import ModelInput +from lightllm.utils.dist_utils import get_dp_world_size, get_global_dp_rank from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.dist_utils import get_global_dp_rank, get_dp_world_size -from .attention import BasePrefillAttState, BaseDecodeAttState + +from lightllm.common.attention.session import AttentionSession +from .triton_kernel.gen_decode_params import gen_decode_params +from .triton_kernel.gen_prefill_params import gen_prefill_params class InferStateInfo: @@ -20,13 +21,7 @@ class InferStateInfo: """ def __init__(self): - # prefill 和 decode 使用的 att 状态对象 - self.prefill_att_state: BasePrefillAttState = None - self.decode_att_state: BaseDecodeAttState = None - - # 保留的扩展, 支持线性att与标准att混合使用时使用 - self.prefill_att_state1: BasePrefillAttState = None - self.decode_att_state1: BaseDecodeAttState = None + self.attention = AttentionSession() self.input_ids: torch.Tensor = None self.batch_size: int = None @@ -133,59 +128,25 @@ def init_some_extra_state(self, model): ) = gen_decode_params(self.b_seq_len) self.b_kv_start_loc = self.b1_cu_kv_seq_len[0:-1] - def init_att_state(self): - if self.is_prefill: - if self.prefill_att_state is not None: - self.prefill_att_state.init_state() - if self.prefill_att_state1 is not None: - self.prefill_att_state1.init_state() - else: - if self.decode_att_state is not None: - self.decode_att_state.init_state() - if self.decode_att_state1 is not None: - self.decode_att_state1.init_state() - self.check_att_states() - - def check_att_states(self) -> None: - """检查 attention state，防止其强引用 InferStateInfo 而形成循环引用。 - - InferStateInfo 已经通过下列四个属性强引用主、辅 attention state。 - 如果 attention state 再把 infer_state 存到普通属性中，就会形成 - infer_state -> attention state -> infer_state 的强引用闭环。 - 即使一次推理结束、外部引用全部释放，引用计数也无法立即回收这两个对象， - 必须等待 Python 的循环 GC。循环 GC 主要由 Python 对象分配和阈值驱动， - 不能保证在 GPU 显存紧张时及时运行。 - - 这些对象还持有本次推理的 Tensor。对象延迟回收会使相关 Tensor 的显存 - 继续被占用，无法及时交还给 PyTorch 分配器复用；多次推理累积后可能抬高 - 显存峰值，甚至触发 OOM。集中进行循环 GC 还可能造成推理延迟抖动。 - 因此，即使循环最终能够被 GC 回收，也不能依赖它管理推理临时状态的生命周期。 - - InferStateInfo 的生命周期应由推理调用方或 CUDA Graph 缓存负责： - 调用方使用期间、图仍需回放期间保留强引用，结束使用后即可释放。 - attention state 只通过 _infer_state_weak_ref 和 infer_state 属性访问它， - 不应额外延长其生命周期。这里在 attention state 初始化完成后统一检查， - 尽早发现新后端或后续修改误存的 InferStateInfo 强引用。 - - 注意：检查仅扫描非空 attention state 的第一层实例属性，不递归检查容器或 - 回调闭包；实现后端时也应避免通过这些位置间接保存 infer_state 的强引用。 - """ - for state in (self.prefill_att_state, self.prefill_att_state1, self.decode_att_state, self.decode_att_state1): - if state is not None: - state.check_no_infer_state() - - def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"): + def graph_input_copies(self, new_infer_state: "InferStateInfo"): + if self.att_state_role != new_infer_state.att_state_role: + raise ValueError("captured attention belongs to a different model role or group") + copies = self.attention.replay_input_copies(new_infer_state.attention) + batch_tensor_names = { + name for name, value in vars(self.attention.batch).items() if isinstance(value, torch.Tensor) + } for attr_name, attr_value in vars(new_infer_state).items(): - if isinstance(attr_value, torch.Tensor): - attr_ = getattr(self, attr_name, None) - if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr(): - attr_.copy_(attr_value, non_blocking=True) - - if self.decode_att_state is not None: - self.decode_att_state.copy_for_decode_cuda_graph(new_infer_state.decode_att_state) - if self.decode_att_state1 is not None: - self.decode_att_state1.copy_for_decode_cuda_graph(new_infer_state.decode_att_state1) - return + if isinstance(attr_value, torch.Tensor) and attr_name not in batch_tensor_names: + target = getattr(self, attr_name, None) + if target is not None and target.data_ptr() != attr_value.data_ptr(): + if (target.shape, target.dtype, target.device) != ( + attr_value.shape, + attr_value.dtype, + attr_value.device, + ): + raise ValueError(f"captured input {attr_name} changed shape, dtype or device") + copies.append((target, attr_value)) + return copies def call_overlap_hook(self): """ @@ -330,6 +291,7 @@ def _all_to_all_balance_get(self, data: torch.Tensor, change_state: bool = True) self.__change_to_balance() dp_rank = get_global_dp_rank() import torch.distributed as dist + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager old_shape = data.shape @@ -361,6 +323,7 @@ def _all_to_all_unbalance_get(self, data: torch.Tensor, change_state: bool = Tru dp_rank = get_global_dp_rank() import torch.distributed as dist + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager old_shape = data.shape @@ -386,11 +349,20 @@ def _all_to_all_unbalance_get(self, data: torch.Tensor, change_state: bool = Tru return origin_data.view(-1, *old_shape[1:]) # 用于 prefill cuda graph 的专用功能接口 + @contextmanager + def _capture_prefill_segment(self, graph_obj): + with self.device_runtime.capture_graph(graph_obj, pool=self.mem_pool): + self._prefill_capture_active = True + try: + yield + finally: + self._prefill_capture_active = False + def prefill_cuda_graph_create_graph_obj(self): if not hasattr(self, "prefill_cuda_graph_exe_list"): self.prefill_cuda_graph_exe_list = [] - graph_obj = torch.cuda.CUDAGraph() - capture_graph = torch.cuda.graph(graph_obj, pool=self.mem_pool) + graph_obj = self.device_runtime.create_graph() + capture_graph = self._capture_prefill_segment(graph_obj) self.prefill_cuda_graph_exe_list.append((graph_obj, capture_graph)) return @@ -423,6 +395,19 @@ def prefill_replay(self, new_infer_state: "InferStateInfo"): func(new_infer_state) return + def close_prefill_graphs(self): + try: + # 错误可能发生在某一段内部，也可能发生在段间的 CPU 回调中。 + # 只退出仍在捕获的上下文，避免重复结束已经完成的段。 + if getattr(self, "_prefill_capture_active", False): + self.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + finally: + self.device_runtime.synchronize() + for item in getattr(self, "prefill_cuda_graph_exe_list", ()): + if isinstance(item, tuple): + item[0].reset() + self.prefill_cuda_graph_exe_list = [] + def copy_for_prefill_cuda_graph(self, new_infer_state: "InferStateInfo"): for attr_name, attr_value in vars(new_infer_state).items(): if isinstance(attr_value, torch.Tensor): diff --git a/lightllm/common/basemodel/layer_infer/base_layer_infer.py b/lightllm/common/basemodel/layer_infer/base_layer_infer.py index 354920f142..82de11853d 100644 --- a/lightllm/common/basemodel/layer_infer/base_layer_infer.py +++ b/lightllm/common/basemodel/layer_infer/base_layer_infer.py @@ -1,16 +1,31 @@ +from lightllm.common.basemodel.attention_adapters import AttentionDescription +from typing import Iterable, Union + import torch import torch.distributed as dist -from typing import Dict, Iterable, Literal, Tuple, Union, List + from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight +from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy +from lightllm.distributed.communication_op import all_gather_into_tensor, all_reduce, reduce_scatter_tensor from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size -from .cache_tensor_manager import g_cache_manager from lightllm.utils.envs_utils import get_env_start_args -from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor, all_reduce -from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy + +from .cache_tensor_manager import g_cache_manager class BaseLayerInfer: + def describe_attention(self, phase): + return AttentionDescription() + + def attention_adapter(self, family): + from lightllm.common.basemodel.attention_adapters import ADAPTERS + + return ADAPTERS[family] + + def bind_attention_resources(self, model, weight, adapter): + return adapter.bind(model, self, weight) + def __init__(self) -> None: self.tp_rank_ = get_current_rank_in_dp() self.tp_world_size_ = get_dp_world_size() diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index 4f21d90cfc..646785b884 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -1,31 +1,49 @@ -import os -import torch -import copy import bisect +from typing import List, Optional, Tuple + +import torch import triton -from typing import List, Tuple -from typing import Optional -from lightllm.utils.log_utils import init_logger + +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.platform import get_device_runtime from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor -from lightllm.distributed import dist_group_manager -from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from .infer_struct import InferStateInfo + from .cuda_graph import CudaGraph +from .infer_struct import InferStateInfo logger = init_logger(__name__) +class PrefillGraphExecution: + """分段图拥有输入/输出和 collector；attention 仍通过段间回调使用本批次 Plan。""" + + def __init__(self, infer_state, inputs, outputs, collector, runtime): + self.state, self.inputs, self.outputs, self.collector = infer_state, inputs, outputs, collector + self.runtime = runtime + self.generation = infer_state.att_state_manager.storage_generation + infer_state.att_state_manager.register_graph(self) + + def close(self): + if self.state is None: + return + self.state.close_prefill_graphs() + self.state = self.collector = None + self.inputs = self.outputs = () + + class PrefillCudaGraph: # CudaGraph forward pass for the decoding stage. - def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int): + def __init__(self, decode_cuda_graph: CudaGraph, tp_world_size: int, runtime=None): self.graph = {} + self.runtime = runtime or (decode_cuda_graph.runtime if decode_cuda_graph is not None else get_device_runtime()) self.tp_world_size = tp_world_size if decode_cuda_graph is not None: self.mempool = decode_cuda_graph.mempool # prefill 和 decode 共享一个 mempool else: - self.mempool = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None + self.mempool = self.runtime.graph_pool() self.args = get_env_start_args() self.enable_prefill_microbatch_overlap = self.args.enable_prefill_microbatch_overlap @@ -62,7 +80,7 @@ def can_run(self, handle_token_num: int): def need_capture(self, handle_token_num: int): finded_handle_token_num = self.find_closest_graph_handle_token_num(handle_token_num=handle_token_num) if finded_handle_token_num is not None: - return finded_handle_token_num not in self.graph + return finded_handle_token_num not in self.graph or self.graph[finded_handle_token_num].state is None else: assert False, "dead code" @@ -79,35 +97,38 @@ def _capture_prefill( ) -> List[torch.Tensor]: handle_token_num = infer_state.input_ids.shape[0] infer_state.mem_pool = self.mempool - infer_state.prefill_cuda_graph_create_graph_obj() - infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() - graph_input_tensors: List[torch.Tensor] = [torch.empty_like(e) for e in input_tensors] - graph_out_tensors: List[torch.Tensor] = prefill_func(graph_input_tensors, infer_state) - graph_out_tensors = [e.contiguous() for e in graph_out_tensors] - infer_state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) - - graph_input_tensors = [tensor_to_no_ref_tensor(e) for e in graph_input_tensors] - graph_out_tensors = [tensor_to_no_ref_tensor(e) for e in graph_out_tensors] - - graph_hidden_collector = infer_state.hidden_collector - hidden_tensor_count, hidden_tensor_nbytes = graph_hidden_collector.release_graph_tensor_ownership() - logger.info( - f"Prefill CUDA Graph hidden collector 已完成 capture 状态托管：" - f"handle_token_num={handle_token_num}, " - f"collector={graph_hidden_collector.__class__.__name__}, " - f"no_ref_tensor_count={hidden_tensor_count}, " - f"no_ref_tensor_nbytes={hidden_tensor_nbytes}。" - f"这些 tensor 的固定地址由 graph memory pool 管理，collector 仅保留无所有权引用供 replay 使用。" - ) - self.graph[handle_token_num] = ( - infer_state, - graph_input_tensors, - graph_out_tensors, - graph_hidden_collector, - ) - self.replay(input_tensors, infer_state) + infer_state.device_runtime = self.runtime + try: + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + graph_input_tensors: List[torch.Tensor] = [torch.empty_like(e) for e in input_tensors] + graph_out_tensors: List[torch.Tensor] = prefill_func(graph_input_tensors, infer_state) + graph_out_tensors = [e.contiguous() for e in graph_out_tensors] + infer_state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + + graph_input_tensors = [tensor_to_no_ref_tensor(e) for e in graph_input_tensors] + graph_out_tensors = [tensor_to_no_ref_tensor(e) for e in graph_out_tensors] + + graph_hidden_collector = infer_state.hidden_collector + hidden_tensor_count, hidden_tensor_nbytes = graph_hidden_collector.release_graph_tensor_ownership() + logger.info( + f"Prefill CUDA Graph hidden collector 已完成 capture 状态托管：" + f"handle_token_num={handle_token_num}, " + f"collector={graph_hidden_collector.__class__.__name__}, " + f"no_ref_tensor_count={hidden_tensor_count}, " + f"no_ref_tensor_nbytes={hidden_tensor_nbytes}。" + f"这些 tensor 的固定地址由 graph memory pool 管理，collector 仅保留无所有权引用供 replay 使用。" + ) + self.graph[handle_token_num] = PrefillGraphExecution( + infer_state, graph_input_tensors, graph_out_tensors, graph_hidden_collector, self.runtime + ) + self.replay(input_tensors, infer_state) - return graph_out_tensors + return graph_out_tensors + except BaseException: + self.graph.pop(handle_token_num, None) + infer_state.close_prefill_graphs() + raise def _capture_prefill_overlap( self, @@ -148,9 +169,17 @@ def capture_prefill( def _replay(self, input_tensors: List[torch.Tensor], infer_state: InferStateInfo) -> List[torch.Tensor]: handle_token_num = infer_state.input_ids.shape[0] - graph_infer_state, graph_input_tensors, graph_output_tensors, graph_hidden_collector = self.graph[ - handle_token_num - ] + execution = self.graph[handle_token_num] + graph_infer_state = execution.state + if graph_infer_state is None: + raise RuntimeError("prefill Graph execution is closed") + if graph_infer_state.att_state_manager.storage_generation != execution.generation: + raise RuntimeError("attention storage changed; recapture prefill Graph") + graph_input_tensors, graph_output_tensors, graph_hidden_collector = ( + execution.inputs, + execution.outputs, + execution.collector, + ) graph_infer_state: InferStateInfo = graph_infer_state for graph_in_tensor, in_tensor in zip(graph_input_tensors, input_tensors): graph_in_tensor.copy_(in_tensor) @@ -165,6 +194,11 @@ def _replay(self, input_tensors: List[torch.Tensor], infer_state: InferStateInfo return graph_output_tensors + def close(self): + for execution in self.graph.values(): + execution.close() + self.graph.clear() + def _replay_overlap( self, input_tensors: List[torch.Tensor], diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py deleted file mode 100644 index aa7cfa9f4a..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Linear-attention / GDN triton kernels shared across hybrid models.""" - -from .causal_conv1d import causal_conv1d_fn -from .causal_conv1d_mtp import causal_conv1d_update -from .fused_gdn_gating import fused_gdn_gating -from .gdn_decode_pack import conv_pack_gdn_decode_inputs -from .mtp_fused_recurrent import mtp_fused_recurrent_gated_delta_rule -from .fla.ops import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule - -__all__ = [ - "causal_conv1d_fn", - "causal_conv1d_update", - "fused_gdn_gating", - "conv_pack_gdn_decode_inputs", - "mtp_fused_recurrent_gated_delta_rule", - "chunk_gated_delta_rule", - "fused_recurrent_gated_delta_rule", -] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py deleted file mode 100644 index 2bde70bb99..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang -# -# This file contains code copied from the flash-linear-attention project. -# The original source code was licensed under the MIT license and included -# the following copyright notice: -# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang - -# Adapted from -# https://github.com/vllm-project/vllm diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py deleted file mode 100644 index cd3b0962a3..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang -# -# This file contains code copied from the flash-linear-attention project. -# The original source code was licensed under the MIT license and included -# the following copyright notice: -# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang -from .chunk import chunk_gated_delta_rule -from .fused_recurrent import fused_recurrent_gated_delta_rule - -__all__ = [ - "chunk_gated_delta_rule", - "fused_recurrent_gated_delta_rule", -] diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index 98fab22505..b44a529203 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -5,7 +5,6 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger - logger = init_logger("lightllm.common.req_manager") @@ -53,6 +52,7 @@ def __init__(self, max_request_num, max_sequence_length): self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num + self.request_generations = torch.zeros(max_request_num + 1, dtype=torch.int64, device="cuda") def bind_hold_page(self, token_indices): """Map padding requests to the resource owner's reserved physical page.""" @@ -67,11 +67,15 @@ def bind_hold_page(self, token_indices): hold_row.view(-1, self.page_size).copy_(hold_page) def alloc(self): - return self.req_list.alloc() + index = self.req_list.alloc() + if index is not None: + self.request_generations[index].add_(1) + return index def free(self, free_req_indexes: List[int]): for req_index in free_req_indexes: self.req_list.free(req_index) + self.request_generations[req_index].add_(1) if self.req_list.is_all_free(): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") @@ -80,4 +84,5 @@ def free_req(self, free_req_index: int): def free_all(self): self.req_list = _ReqIndexPool(self.max_request_num) + self.request_generations.add_(1) return diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py index 5ec00d542a..2b29a477ec 100755 --- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py @@ -1,13 +1,20 @@ -from lightllm.common.attention_state.full_att import FullAttentionState -import torch +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFeature from typing import Tuple -from lightllm.common.basemodel import TransformerLayerInferTpl + +import torch + +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.basemodel import InferStateInfo, TransformerLayerInferTpl from lightllm.models.bloom.layer_weights.transformer_layer_weight import BloomTransformerLayerWeight -from lightllm.common.basemodel import InferStateInfo -from lightllm.common.basemodel.attention.base_att import AttControl class BloomTransformerLayerInfer(TransformerLayerInferTpl): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace(description, features=description.features | {AttentionFeature.ALIBI}) + return description + """ """ def __init__(self, layer_num, network_config): @@ -29,42 +36,16 @@ def _context_attention_kernel( layer_weight: BloomTransformerLayerWeight, out=None, ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(use_alibi=True, tp_alibi=layer_weight.tp_alibi), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state.attention, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor def _token_attention_kernel( self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(use_alibi=True, tp_alibi=layer_weight.tp_alibi), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.decode_attention.forward(infer_state.attention, DenseInputs(_q)) return o_tensor.view(q.shape) def _att_norm( diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 8e963d7f84..b2e51f3c3c 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -1,17 +1,18 @@ +from lightllm.common.attention.selection import BackendPolicy from typing import TYPE_CHECKING -from lightllm.models.bloom.layer_infer.pre_layer_infer import BloomPreLayerInfer + +from lightllm.common.basemodel import InferStateInfo, TpPartBaseModel from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer +from lightllm.models.bloom.layer_infer.pre_layer_infer import BloomPreLayerInfer from lightllm.models.bloom.layer_infer.transformer_layer_infer import BloomTransformerLayerInfer from lightllm.models.bloom.layer_weights.pre_and_post_layer_weight import BloomPreAndPostLayerWeight from lightllm.models.bloom.layer_weights.transformer_layer_weight import BloomTransformerLayerWeight -from lightllm.common.basemodel import InferStateInfo, TpPartBaseModel -from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend - if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class BloomTpPartModel(TpPartBaseModel): @@ -75,7 +76,5 @@ def _reset_num_key_value_heads(self): self.config["num_key_value_heads"] = self.config["num_attention_heads"] return - def _init_att_backend(self): - self.prefill_att_backend = TritonAttBackend(self) - self.decode_att_backend = TritonAttBackend(self) - return + def attention_policy(self, family, phase): + return BackendPolicy(allow=frozenset({"triton"}), prefer=("triton",)) diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index a9d23a6de5..d09358e8b8 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -1,26 +1,33 @@ -from lightllm.common.attention_state.full_att import FullAttentionState +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFamily import os +from functools import partial + import torch import triton + +from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAPrefillInputs +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe +from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.models.deepseek2.triton_kernel.sample_kv import sample_kv from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( - use_sm100_mega_moe, -) -from functools import partial from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale -from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_global_world_size +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class Deepseek2TransformerLayerInfer(LlamaTransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace(description, family=AttentionFamily.MLA) + return description + def __init__(self, layer_num, network_config): self.tp_k_head_num_ = 1 self.tp_v_head_num_ = 1 @@ -86,13 +93,7 @@ def _context_attention_kernel( layer_weight=layer_weight, ) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=q, - k=(k_nope, k_rope), - v=v, - att_control=AttControl(mla_prefill=True, mla_prefill_dict={"softmax_scale": self.softmax_scale}), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state.attention, MLAPrefillInputs(q, k_nope, k_rope, v)) return o_tensor def _token_attention_kernel( @@ -104,21 +105,8 @@ def _token_attention_kernel( ): q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - kv = state.get_layer_view(logical_layer_id=self.layer_num_) - - out = infer_state.decode_att_state.decode_att( - q=(q_nope, q_rope), - k=kv, - v=None, - att_control=AttControl(mla_decode=True, mla_decode_dict={"softmax_scale": self.softmax_scale}), - alloc_func=self.alloc_tensor, - ) + + out = self.decode_attention.forward(infer_state.attention, MLADecodeInputs(q_nope, q_rope)) return out def _decompress_kv( @@ -418,7 +406,6 @@ def overlap_tpsp_token_forward( def _1_hook_post(): _1_hook() - nonlocal _1_ffn_out if self.n_shared_experts is not None: _1_ffn_out.add_(_1_shared_output) input_embdings1.add_(_1_ffn_out.view(-1, self.embed_dim_)) @@ -581,7 +568,6 @@ def overlap_tpsp_context_forward( def _1_hook_post(): _1_hook() - nonlocal _1_ffn_out if self.n_shared_experts is not None: _1_ffn_out.add_(_1_shared_output) input_embdings1.add_(_1_ffn_out.view(-1, self.embed_dim_)) diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index e06c18ea16..c50cbfdfa5 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -1,22 +1,23 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, final + import torch -from typing import final + +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight -from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args -from lightllm.distributed.communication_op import dist_group_manager -from lightllm.common.basemodel.attention import get_mla_decode_att_backend_class, get_mla_prefill_att_backend_class +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Deepseek2TpPartModel(LlamaTpPartModel): @@ -33,11 +34,6 @@ def __init__(self, kvargs): super().__init__(kvargs) return - def _init_att_backend(self): - self.prefill_att_backend = get_mla_prefill_att_backend_class(index=0)(model=self) - self.decode_att_backend = get_mla_decode_att_backend_class(index=0)(model=self) - return - def _init_some_value(self): super()._init_some_value() self.tp_k_head_num_ = 1 @@ -95,7 +91,7 @@ def create_full_att_states( return [state] def _init_to_get_yarn_rotary(self): - from lightllm.models.llama.yarn_rotary_utils import find_correction_range, linear_ramp_mask, get_deepseek_mscale + from lightllm.models.llama.yarn_rotary_utils import find_correction_range, get_deepseek_mscale, linear_ramp_mask dim = self.qk_rope_head_dim max_position_embeddings = self.config.get("max_position_embeddings", 2048) diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index e026d1ea8e..6f7c600566 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,20 +1,29 @@ -from lightllm.common.attention_state.full_att import FullAttentionState -import torch +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFamily from typing import Any + +import torch + +from lightllm.common.attention.api.inputs import SparseMLADecodeInputs, SparseMLAPrefillInputs +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.distributed import all_gather_into_tensor from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer -from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight from lightllm.models.deepseek3_2.triton_kernel.act_quant import act_quant from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_indexer_ks import destindex_copy_indexer_ks from lightllm.models.deepseek3_2.triton_kernel.extract_indexer_ks import extract_indexer_ks from lightllm.utils.envs_utils import get_env_start_args -from lightllm.distributed import all_gather_into_tensor class Deepseek3_2TransformerLayerInfer(Deepseek2TransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace(description, family=AttentionFamily.SPARSE_MLA) + return description + def __init__(self, layer_num, network_config): self.index_topk = network_config["index_topk"] super().__init__(layer_num, network_config) @@ -78,7 +87,7 @@ def _context_attention_kernel( q_all = torch.cat([q_nope, q_rope], dim=-1) # 计算 topk indices - att_state = infer_state.prefill_att_state + att_state = self.prefill_attention.context(infer_state.attention) topk_mem_indices, topk_indices = self.indexer._get_indices( hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], @@ -89,28 +98,10 @@ def _context_attention_kernel( del infer_state.get_topk_indices_params # Use NSA backend for attention computation - att_control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_mem_indices": topk_mem_indices, - "topk_indices": topk_indices, - "prefill_cache_kv": kv, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - }, - ) - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - mla_out = infer_state.prefill_att_state.prefill_att( - q=q_all, - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=att_control, + mla_out = self.prefill_attention.forward( + infer_state.attention, + SparseMLAPrefillInputs(q_all, topk_mem_indices, topk_indices, kv), ) return mla_out @@ -126,7 +117,7 @@ def _token_attention_kernel( q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) # 计算 topk mem indices - att_state = infer_state.decode_att_state + att_state = self.decode_attention.context(infer_state.attention) topk_mem_indices, _ = self.indexer._get_indices( hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], @@ -137,28 +128,10 @@ def _token_attention_kernel( del infer_state.get_topk_indices_params # Use NSA backend for attention computation - att_control = AttControl( - nsa_decode=True, - nsa_decode_dict={ - "layer_index": self.layer_num_, - "topk_mem_indices": topk_mem_indices, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - "qk_rope_head_dim": self.qk_rope_head_dim, - }, - ) - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - o_tensor = infer_state.decode_att_state.decode_att( - q=(q_nope, q_rope), - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=att_control, + o_tensor = self.decode_attention.forward( + infer_state.attention, + SparseMLADecodeInputs(q_nope, q_rope, topk_mem_indices), ) return o_tensor @@ -189,7 +162,6 @@ def _get_indices( att_state: Any, layer_weight: Deepseek3_2TransformerLayerWeight, ): - q, k = self._get_q_k_bf16(hidden_states, q_lora, infer_state, layer_weight) if self.tp_world_size_ > 1: diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index b691d986cc..4392624db7 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -1,19 +1,19 @@ from typing import TYPE_CHECKING + import torch + from lightllm.models.deepseek2.model import Deepseek2TpPartModel -from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import Deepseek3_2TransformerLayerInfer -from lightllm.common.basemodel.attention import get_nsa_prefill_att_backend_class, get_nsa_decode_att_backend_class - +from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Deepseek3_2TpPartModel(Deepseek2TpPartModel): - # weight class transformer_weight_class = Deepseek3_2TransformerLayerWeight @@ -51,8 +51,3 @@ def create_full_att_states( device="meta", ) return [state] - - def _init_att_backend(self): - self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0)(model=self) - self.decode_att_backend = get_nsa_decode_att_backend_class(index=0)(model=self) - return diff --git a/lightllm/models/deepseek_v4/attention/__init__.py b/lightllm/models/deepseek_v4/attention/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/deepseek_v4/attention/adapter.py b/lightllm/models/deepseek_v4/attention/adapter.py new file mode 100644 index 0000000000..11c3d560e9 --- /dev/null +++ b/lightllm/models/deepseek_v4/attention/adapter.py @@ -0,0 +1,47 @@ +from lightllm.common.attention.api import CacheEncoding, CacheLayout +from lightllm.common.basemodel.attention_adapters import AttentionAdapter + + +class DeepseekV4Adapter(AttentionAdapter): + state_type = "swa_" + + def dimensions(self, model, layer): + return layer.tp_q_head_num_, 1, layer.head_dim_, layer.v_head_dim + + def cache(self, model, state, states, heads, head_dim, value_dim): + import torch + + return CacheLayout( + encoding=CacheEncoding.FP8_SPARSE_MLA, + page_size=model.args.page_size, + dtype=torch.uint8, + key_heads=heads, + key_dim=head_dim, + value_dim=value_dim, + storage_schema=tuple(state.describe_layout() for state in states), + ) + + def scale(self, layer): + return layer.softmax_scale + + def bind(self, model, layer, weight): + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4Resources + + manager, role, layer_id = model.att_state_manager, model.att_state_role, layer.layer_num_ + state = manager.get_state(role, self.state_type, layer_id) + swa = state.layer_resources(layer_id) + prefill_swa = model.swa_prefill_workspace[0][state.config.logical_layer_ids.index(layer_id)].view( + -1, swa.shape[-1] + ) + compressed = ( + manager.get_state(role, "compressed", layer_id).layer_resources(layer_id) if layer.compress_ratio else None + ) + return DeepseekV4Resources( + swa, + prefill_swa, + compressed, + layer.compress_ratio, + layer.v_head_dim, + layer.softmax_scale, + weight.attn_sink_.weight, + ) diff --git a/lightllm/models/deepseek_v4/attention/sparse_mla.py b/lightllm/models/deepseek_v4/attention/sparse_mla.py new file mode 100644 index 0000000000..7da1569508 --- /dev/null +++ b/lightllm/models/deepseek_v4/attention/sparse_mla.py @@ -0,0 +1,193 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase, AttentionFeature +import inspect +from dataclasses import dataclass + +import torch + +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.common.attention.api.backend import AttentionContext, AttentionBackend + + +@dataclass(frozen=True) +class DeepseekV4Resources: + swa: torch.Tensor + prefill_swa: torch.Tensor + compressed: torch.Tensor | None + compress_ratio: int + head_dim_v: int + softmax_scale: float + attn_sink: torch.Tensor + + +@dataclass(frozen=True) +class DeepseekV4Inputs: + query: torch.Tensor + swa_indices: torch.Tensor + swa_lengths: torch.Tensor + extra_indices: torch.Tensor | None = None + extra_lengths: torch.Tensor | None = None + + +@dataclass +class SwaPrefillWorkspace: + swa_buffer: torch.Tensor + HOLD_SWA_SLOT: int + prefill_token_ranges: torch.Tensor + swa_prefill_workspace_offsets: torch.Tensor + + +def prepare_swa_workspace(swa_state, buffer, capacity, batch, slot): + from lightllm.models.deepseek_v4.triton_kernel.swa_cache import prepare_swa_prefill_workspace + + metadata = SwaPrefillWorkspace( + buffer, + buffer.shape[1] - swa_state.config.page_token_size, + torch.zeros((swa_state.request_capacity, 2), dtype=torch.int32, device=buffer.device), + torch.full((swa_state.request_capacity,), -1, dtype=torch.int32, device=buffer.device), + ) + prepare_swa_prefill_workspace( + swa_state=swa_state, + prefill_state=metadata, + b_req_idx=batch.b_req_idx, + b_ready_cache_len=batch.b_ready_cache_len, + b_seq_len=batch.b_seq_len, + workspace_base=slot * capacity, + ) + return metadata + + +@dataclass(slots=True) +class DeepseekV4Context(AttentionContext): + sched_meta_by_ratio: dict | None = None + capture_sched_meta_by_ratio: dict | None = None + + +class DeepseekV4SparseMLAFp8Common: + variant = "deepseek_v4" + family = AttentionFamily.SPARSE_MLA + + provider = "flashmla_sparse" + + dependency = "flash_mla" + + min_device_capability = (9, 0) + + activation_dtypes = frozenset({torch.float16, torch.bfloat16}) + + cache_encodings = frozenset({CacheEncoding.FP8_SPARSE_MLA}) + + @classmethod + def check_environment(cls, caps) -> None: + super().check_environment(caps) + import flash_mla + + parameters = inspect.signature(flash_mla.flash_mla_with_kvcache).parameters + required = {"out", "attn_sink", "indices", "topk_length", "extra_k_cache", "extra_topk_length"} + missing = required.difference(parameters) + if missing: + raise UnsupportedAttention( + "DSV4 requires the FlashMLA version pinned in docker/Dockerfile; missing: " + ", ".join(sorted(missing)) + ) + + def __init__(self, env, spec): + super().__init__(env=env, spec=spec) + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=env.network_config, + model_logical_layer_ids=env.layer_ids, + tp_rank=env.tp_rank, + tp_world_size=env.tp_size, + ) + self.compress_ratios = tuple(dict.fromkeys(config.compress_ratios)) + self.swa_config = config.create_swa_config() + self.compressed_configs = { + ratio: config.create_compressed_config(compress_ratio=ratio) for ratio in self.compress_ratios if ratio + } + + def _plan(self, context): + import flash_mla as flashmla + + if context.sched_meta_by_ratio is None: + context.sched_meta_by_ratio = {ratio: flashmla.get_mla_metadata()[0] for ratio in self.compress_ratios} + context.capture_sched_meta_by_ratio = { + ratio: flashmla.get_mla_metadata()[0] for ratio in self.compress_ratios + } + + def _run(self, layer, context, inputs, swa, out): + import flash_mla as flashmla + + params = layer.resources + compressed = None + if params.compressed is not None: + compressed = self.compressed_configs[params.compress_ratio].get_attention_view(packed_kv=params.compressed) + scheduler = ( + context.capture_sched_meta_by_ratio + if torch.cuda.is_current_stream_capturing() + else context.sched_meta_by_ratio + ) + q = inputs.query + if out is None: + out = self.alloc_tensor((q.shape[0], q.shape[1], params.head_dim_v), dtype=q.dtype, device=q.device) + padded_out, _ = flashmla.flash_mla_with_kvcache( + q=q.unsqueeze(1), + k_cache=self.swa_config.get_attention_view(packed_kv=swa), + block_table=None, + cache_seqlens=None, + head_dim_v=params.head_dim_v, + tile_scheduler_metadata=scheduler[params.compress_ratio], + num_splits=None, + softmax_scale=params.softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=inputs.swa_indices, + attn_sink=params.attn_sink, + topk_length=inputs.swa_lengths, + extra_k_cache=compressed, + extra_indices_in_kvcache=inputs.extra_indices, + extra_topk_length=inputs.extra_lengths, + out=out.unsqueeze(1), + ) + return padded_out[:, 0, : self.spec.num_q_heads] + + +class DeepseekV4SparseMLAFp8PrefillBackend( + DeepseekV4SparseMLAFp8Common, AttentionBackend[DeepseekV4Context, DeepseekV4Inputs, torch.Tensor] +): + phase = AttentionPhase.PREFILL + supported_features = frozenset( + {AttentionFeature.MTP, AttentionFeature.DYNAMIC_VERIFY, AttentionFeature.SINK, AttentionFeature.MTP_WINDOW} + ) + + def allocate_context(self, capacity, *, capture, workspace) -> DeepseekV4Context: + return DeepseekV4Context(capacity, capture, workspace) + + def plan_host(self, context: DeepseekV4Context): + self._plan(context) + + def forward(self, layer, context: DeepseekV4Context, inputs: DeepseekV4Inputs, *, out=None) -> torch.Tensor: + return self._run(layer, context, inputs, layer.resources.prefill_swa, out) + + +class DeepseekV4SparseMLAFp8DecodeBackend( + DeepseekV4SparseMLAFp8Common, AttentionBackend[DeepseekV4Context, DeepseekV4Inputs, torch.Tensor] +): + phase = AttentionPhase.DECODE + supported_features = frozenset( + { + AttentionFeature.MTP, + AttentionFeature.DYNAMIC_VERIFY, + AttentionFeature.SINK, + AttentionFeature.NONCAUSAL, + AttentionFeature.MTP_WINDOW, + } + ) + + def allocate_context(self, capacity, *, capture, workspace) -> DeepseekV4Context: + return DeepseekV4Context(capacity, capture, workspace) + + def plan_host(self, context: DeepseekV4Context): + self._plan(context) + + def forward(self, layer, context: DeepseekV4Context, inputs: DeepseekV4Inputs, *, out=None) -> torch.Tensor: + return self._run(layer, context, inputs, layer.resources.swa, out) diff --git a/lightllm/models/deepseek_v4/deepseek_v4_visual.py b/lightllm/models/deepseek_v4/deepseek_v4_visual.py index 3e730c5b2a..4eef351854 100644 --- a/lightllm/models/deepseek_v4/deepseek_v4_visual.py +++ b/lightllm/models/deepseek_v4/deepseek_v4_visual.py @@ -13,14 +13,11 @@ from safetensors import safe_open from torch import nn -from lightllm.models.deepseek_v4.image_processor import ( - IMAGE, - build_canonical_image_block, - load_image, -) +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.models.deepseek_v4.image_processor import IMAGE, build_canonical_image_block, load_image from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm from lightllm.server.multimodal_params import ImageItem -from lightllm.server.visualserver import get_vit_attn_backend +from lightllm.server.visualserver import get_vit_attention_op @lru_cache(8) @@ -79,7 +76,7 @@ def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, cu_seql q = apply_rotary(q, cos, sin) k = apply_rotary(k, cos, sin) o = torch.empty_like(q) - get_vit_attn_backend()(q, k, v, o, cu_seqlens, n) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, n), out=o) return self.wo(o.reshape(n, -1)) diff --git a/lightllm/models/deepseek_v4/infer_struct.py b/lightllm/models/deepseek_v4/infer_struct.py index 6f90fa1207..01d4c67381 100644 --- a/lightllm/models/deepseek_v4/infer_struct.py +++ b/lightllm/models/deepseek_v4/infer_struct.py @@ -2,12 +2,13 @@ from typing import Literal import torch -from lightllm.common.basemodel import InferStateInfo + from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( - CompressedAttentionState, CompressedAttentionLayerView, + CompressedAttentionState, ) +from lightllm.common.basemodel import InferStateInfo C4_PREFILL_LOGITS_BUDGET_BYTES = 512 * 1024 * 1024 @@ -84,7 +85,7 @@ def __init__(self): self._dsv4_token_to_batch_idx = None # The first C4 layer builds metadata and later layers reuse it within this forward. # During CUDA Graph capture its tensor producers become graph nodes. Replay updates - # those captured buffers from staged inputs; copy_for_cuda_graph skips this container. + # those captured buffers from staged inputs; graph_input_copies skips this container. self.c4_indexer_metadata: C4IndexerMetadata | None = None def init_some_extra_state(self, model): @@ -105,7 +106,10 @@ def init_some_extra_state(self, model): self.dsv4_c128_lengths = None super().init_some_extra_state(model) # sets position_ids, b_q_seq_len, b_q_start_loc (prefill) if self.is_prefill: - self.prefill_att_state.prepare_swa_workspace(self) + from .attention.sparse_mla import prepare_swa_workspace + + buffer, capacity = model.swa_prefill_workspace + self.swa_workspace = prepare_swa_workspace(self.swa_state, buffer, capacity, self, self.microbatch_index) pos = self.position_ids self.position_cos_sliding = torch.index_select(model._cos_cached_sliding, 0, pos) # [T, rope_dim//2] self.position_sin_sliding = torch.index_select(model._sin_cached_sliding, 0, pos) @@ -196,7 +200,7 @@ def _init_swa_metadata(self, model): sliding_window=self.swa_state.config.window_size, image_left=self.dsv4_image_left, image_right=self.dsv4_image_right, - prefill_state=self.prefill_att_state if self.is_prefill else None, + prefill_state=self.swa_workspace if self.is_prefill else None, swa_runtime_write_slots=self.dsv4_swa_runtime_write_slots, token_to_batch_idx=self._dsv4_token_to_batch_idx, b_seq_len=self.b_seq_len, @@ -204,7 +208,7 @@ def _init_swa_metadata(self, model): def init_c4_indices(self, width): # Create layer outputs on first use; replay inputs leave them unset, so - # copy_for_cuda_graph stages only inputs into the captured infer state. + # graph_input_copies stages only inputs into the captured infer state. if self.dsv4_c4_indices is None: self.dsv4_c4_indices = torch.empty( (self.position_ids.numel(), width), dtype=torch.int32, device=self.position_ids.device @@ -216,6 +220,7 @@ def get_c4_indexer_metadata(self, c4_cap): if self.c4_indexer_metadata is not None: return self.c4_indexer_metadata import deep_gemm + from .triton_kernel.build_c4_indexer_page_table import build_c4_indexer_page_table valid_len = ((self.position_ids + 1) // 4).to(torch.int32) diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index 21825e68e7..b6b46aa9d3 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -1,26 +1,38 @@ +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFamily, AttentionFeature import torch import triton + from lightllm.common.basemodel import TransformerLayerInferTpl -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.triton_kernel.mhc import hc_post, hc_post_pre_norm, hc_pre_norm from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( - silu_and_mul_fwd, - silu_and_mul_quant_fp8, -) +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd, silu_and_mul_quant_fp8 +from lightllm.common.basemodel.triton_kernel.mhc import hc_post, hc_post_pre_norm, hc_pre_norm from lightllm.common.quantization.deepgemm import DeepGEMMFP8w8a8B128QuantizationMethod +from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4Inputs from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight -from lightllm.utils.envs_utils import get_env_start_args +from lightllm.models.deepseek_v4.triton_kernel.topk_softplus_sqrt import topk_softplus_sqrt from lightllm.utils.dist_utils import get_global_world_size +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor + +from ..infer_struct import DeepseekV4InferStateInfo +from ..triton_kernel.rotary_emb import inverse_rotary_emb from .compressor import CompressorInfer from .indexer import DeepseekV4IndexInfer -from ..triton_kernel.rotary_emb import inverse_rotary_emb -from ..infer_struct import DeepseekV4InferStateInfo -from lightllm.models.deepseek_v4.triton_kernel.topk_softplus_sqrt import topk_softplus_sqrt class DeepseekV4TransformerLayerInfer(TransformerLayerInferTpl): + def attention_adapter(self, family): + from lightllm.models.deepseek_v4.attention.adapter import DeepseekV4Adapter + + return DeepseekV4Adapter() + + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace(description, family=AttentionFamily.SPARSE_MLA, variant="deepseek_v4") + description = replace(description, features=description.features | {AttentionFeature.SINK}) + return description + @classmethod def get_prefill_aux_stream(cls: type["DeepseekV4TransformerLayerInfer"]) -> torch.cuda.Stream: """Return the single C4 prefill auxiliary stream shared by every layer in this process.""" @@ -353,8 +365,8 @@ def _norm_rope_and_cache_kv(self, kv, infer_state, layer_weight): from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa state = infer_state.swa_state - att_state = infer_state.prefill_att_state if infer_state.is_prefill else infer_state.decode_att_state - swa, _ = att_state.backend.get_att_input_params(infer_state, self.layer_num_) + params = self.prefill_attention.resources if infer_state.is_prefill else self.decode_attention.resources + swa = params.prefill_swa if infer_state.is_prefill else params.swa # 一次完成 RMSNorm、RoPE、FP8 打包和 SWA 写入。 fused_kv_norm_rope_to_swa( kv=kv, @@ -489,34 +501,12 @@ def _compress_and_index(self, q_lora, x, infer_state: DeepseekV4InferStateInfo, self.index_infer.write_indexer_k(x, infer_state, layer_weight, cos_table, sin_table) return self.index_infer.build_metadata(x, q_lora, infer_state, layer_weight) - def _context_attention_kernel( - self, - q, - q_lora, - x, - infer_state: DeepseekV4InferStateInfo, - layer_weight: DeepseekV4TransformerLayerWeight, - out=None, - ): + def _context_attention_kernel(self, q, q_lora, x, infer_state, layer_weight, out=None): meta = self._compress_and_index(q_lora, x, infer_state, layer_weight) - att_control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "compress_ratio": self.compress_ratio, - "head_dim_v": self.v_head_dim, - "softmax_scale": self.softmax_scale, - "attn_sink": layer_weight.attn_sink_.weight, - **meta, - }, - ) - return infer_state.prefill_att_state.prefill_att( - q=q, - k=infer_state.prefill_att_state.backend.get_att_input_params(infer_state, self.layer_num_), - v=None, - att_control=att_control, - alloc_func=self.alloc_tensor, - out=out, + inputs = DeepseekV4Inputs( + q, meta["swa_indices"], meta["swa_lengths"], meta.get("extra_indices"), meta.get("extra_lengths") ) + return self.prefill_attention.forward(infer_state.attention, inputs, out=out) # ------------------------------------------------------------------ attention (decode) def token_attention_forward( @@ -526,26 +516,12 @@ def token_attention_forward( o = self._token_attention_kernel(q, q_lora, full_x, infer_state, layer_weight) return self._get_o(o, infer_state, layer_weight) - def _token_attention_kernel( - self, q, q_lora, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight - ): + def _token_attention_kernel(self, q, q_lora, x, infer_state, layer_weight, out=None): meta = self._compress_and_index(q_lora, x, infer_state, layer_weight) - att_control = AttControl( - nsa_decode=True, - nsa_decode_dict={ - "compress_ratio": self.compress_ratio, - "head_dim_v": self.v_head_dim, - "softmax_scale": self.softmax_scale, - "attn_sink": layer_weight.attn_sink_.weight, - **meta, - }, - ) - return infer_state.decode_att_state.decode_att( - q=q, - k=infer_state.decode_att_state.backend.get_att_input_params(infer_state, self.layer_num_), - v=None, - att_control=att_control, + inputs = DeepseekV4Inputs( + q, meta["swa_indices"], meta["swa_lengths"], meta.get("extra_indices"), meta.get("extra_lengths") ) + return self.decode_attention.forward(infer_state.attention, inputs, out=out) # ------------------------------------------------------------------ moe def _ffn_shared_expert( diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index ff46b5d6da..ea68793dab 100644 --- a/lightllm/models/deepseek_v4/model.py +++ b/lightllm/models/deepseek_v4/model.py @@ -1,46 +1,34 @@ from __future__ import annotations +from lightllm.models.deepseek_v4.workspace import get_swa_prefill_workspace +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.selection import BackendPolicy from typing import TYPE_CHECKING import torch -from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.models.deepseek_v4.layer_weights.pre_and_post_layer_weight import ( - DeepseekV4PreAndPostLayerWeight, -) -from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import ( - DeepseekV4TransformerLayerWeight, -) -from lightllm.models.deepseek_v4.layer_infer.pre_layer_infer import ( - DeepseekV4PreLayerInfer, -) -from lightllm.models.deepseek_v4.layer_infer.post_layer_infer import ( - DeepseekV4PostLayerInfer, -) -from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import ( - DeepseekV4TransformerLayerInfer, -) -from lightllm.common.basemodel.attention import get_nsa_prefill_att_backend_class, get_nsa_decode_att_backend_class -from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import ( - DSV4_NSA_BACKENDS, - DeepseekV4FlashMlaFp8SparseAttBackend, -) + from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector from lightllm.common.basemodel.triton_kernel.mhc import hc_contract, hc_post +from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4SparseMLAFp8PrefillBackend from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo +from lightllm.models.deepseek_v4.layer_infer.post_layer_infer import DeepseekV4PostLayerInfer +from lightllm.models.deepseek_v4.layer_infer.pre_layer_infer import DeepseekV4PreLayerInfer +from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer +from lightllm.models.deepseek_v4.layer_weights.pre_and_post_layer_weight import DeepseekV4PreAndPostLayerWeight +from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace -from lightllm.models.llama.yarn_rotary_utils import ( - find_correction_range, - linear_ramp_mask, -) +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.llama.yarn_rotary_utils import find_correction_range, linear_ramp_mask from lightllm.utils.envs_utils import get_env_start_args, get_redundancy_expert_num from lightllm.utils.log_utils import init_logger -from lightllm.distributed.communication_op import dist_group_manager if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att.base_state import FullAttentionState + from lightllm.common.attention_state.bounded_att.base_state import BoundedAttentionState + from lightllm.common.attention_state.full_att.base_state import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs logger = init_logger(__name__) @@ -107,6 +95,7 @@ def create_full_att_states( from lightllm.common.attention_state.full_att.compressed_att.create_compressed_att_state import ( create_compressed_att_state, ) + from .config import create_attention_config assert args.llm_kv_type in ("None", "fp8kv_dsa"), f"unsupported V4 KV encoding: {args.llm_kv_type}" @@ -143,9 +132,10 @@ def create_bounded_att_states( """创建本模型的 SWA/compressor 请求运行态，大小页容量由管理器后续确定。""" from lightllm.common.attention_state.bounded_att import ( WindowCompressorStateConfig, - create_window_compressor_state, create_sliding_window_state, + create_window_compressor_state, ) + from .config import create_attention_config config = create_attention_config( @@ -198,7 +188,7 @@ def create_bounded_att_states( def prepare_attention_states(self) -> None: super().prepare_attention_states() # 在统一容量规划前预留长 prefill 的 SWA 临时区，避免 KV 分配挤占计算所需空间。 - DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(self) + get_swa_prefill_workspace(self) def _init_quant(self): if not self.is_mtp_draft_model: @@ -209,18 +199,27 @@ def _init_quant(self): config = dict(self.config, n_layer=len(self.config["compress_ratios"])) self.quant_cfg = Quantcfg(config, self.quant_type, self.quant_cfg_path, self.expert_dtype) - def _init_att_backend(self): + def attention_policy(self, family, phase): + return BackendPolicy( + prefer=("flashmla_sparse",), + extra=( + "lightllm.models.deepseek_v4.attention.sparse_mla:DeepseekV4SparseMLAFp8PrefillBackend", + "lightllm.models.deepseek_v4.attention.sparse_mla:DeepseekV4SparseMLAFp8DecodeBackend", + ), + ) + + def prepare_attention_backends(self): args = get_env_start_args() if args.llm_kv_type == "None": args.llm_kv_type = "fp8kv_dsa" # TODO: 支持其他 kv type if args.llm_kv_type != "fp8kv_dsa": raise RuntimeError("DeepSeek-V4 requires llm_kv_type=fp8kv_dsa for packed FlashMLA sparse attention") - self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0, backend_map=DSV4_NSA_BACKENDS)(model=self) - self.decode_att_backend = get_nsa_decode_att_backend_class(index=0, backend_map=DSV4_NSA_BACKENDS)(model=self) - - tp_q_head_num = self.prefill_att_backend.tp_q_head_num - padded_q_head_num = self.prefill_att_backend.padded_q_head_num + self.args.llm_kv_type = args.llm_kv_type + super().prepare_attention_backends() + backend = self.attention.groups[0][AttentionPhase.PREFILL] + tp_q_head_num = backend.tp_q_head_num + padded_q_head_num = backend.padded_q_head_num if not self.is_mtp_draft_model: # Drafts share this buffer after the main model has captured its CUDA graphs. self.dsv4_workspace.init_flashmla_prefill_q( diff --git a/lightllm/models/deepseek_v4/workspace.py b/lightllm/models/deepseek_v4/workspace.py index 8efcdd8444..54b7e36544 100644 --- a/lightllm/models/deepseek_v4/workspace.py +++ b/lightllm/models/deepseek_v4/workspace.py @@ -7,6 +7,7 @@ import math import torch +import triton class DeepseekV4C4PrefillWorkspace: @@ -109,3 +110,36 @@ def get_flashmla_prefill_q(self, token_num: int): f"DeepSeek-V4 Q workspace capacity {self.token_capacity} is smaller than {token_num} tokens" ) return self.flashmla_prefill_q[:token_num] + + +def get_swa_prefill_workspace(model): + """由 backend 复用计算临时区；模型在 KV 容量评估前调用，以计入实际显存占用。""" + if hasattr(model, "swa_prefill_workspace"): + return model.swa_prefill_workspace + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=model.config, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ).create_swa_config() + page_token_num = config.page_token_size + max_forward_tokens = max(model.batch_max_tokens, model.max_req_num) + 2 * model.tp_world_size_ + capacity = ( + triton.cdiv(max_forward_tokens + model.max_req_num * (config.window_size + 2 * page_token_num), page_token_num) + * page_token_num + ) + if model.args.run_mode == "decode": + # D 节点初始化仍会对 HOLD 请求执行一次 prefill；只保留末尾的 HOLD 页， + # 不为普通请求预留长 prefill 临时区，实际请求的前缀由 P 节点传入。 + capacity = 0 + regions = 1 + int(model.args.enable_prefill_microbatch_overlap) + # 每个 microbatch 使用独立临时区；最后一页留给 HOLD 请求。 + token_num = regions * capacity + page_token_num + # 物理行宽由 SWA 原子配置给出，backend 仅决定 workspace 的 token 容量。 + slot_bytes = config.get_packed_kv_bytes(page_token_num) // page_token_num + layer_count = len(config.logical_layer_ids) + buffer = torch.empty(layer_count * token_num * slot_bytes, dtype=torch.uint8, device="cuda") + model.swa_prefill_workspace = buffer.view(layer_count, token_num, slot_bytes), capacity + return model.swa_prefill_workspace diff --git a/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py index 24bc22aab5..ed0b57a08f 100644 --- a/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py @@ -1,3 +1,5 @@ +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFeature import torch from lightllm.models.deepseek_v4_dspark.infer_struct import DeepseekV4DSparkInferStateInfo @@ -12,6 +14,10 @@ class DeepseekV4DSparkTransformerLayerInfer(DeepseekV4TransformerLayerInfer): """Full DSpark stage with a KV-only target-hidden commit primitive.""" + def describe_attention(self, phase): + description = super().describe_attention(phase) + return replace(description, features=description.features | {AttentionFeature.NONCAUSAL}) + def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.stage_id = layer_num - network_config["draft_layer_start_id"] diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index d1ebd30e30..fd7ed53540 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -1,21 +1,35 @@ -from lightllm.common.attention_state.full_att import FullAttentionState +from dataclasses import replace import math + import torch import torch.nn as nn -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.attention.api.inputs import DenseInputs +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight -from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import ( - context_attention_fwd_gemma4_mm, -) -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + return replace( + description, window=(self.sliding_window_ - 1, 0) if self.is_sliding and self.sliding_window_ > 0 else None + ) + + def bind_attention_resources(self, model, weight, adapter): + resources = super().bind_attention_resources(model, weight, adapter) + k, v = resources.key, resources.value + width = self.tp_k_head_num_ * self.head_dim_ + k = k.flatten(1)[:, :width].view(-1, self.tp_k_head_num_, self.head_dim_) + v = v.flatten(1)[:, :width].view(-1, self.tp_k_head_num_, self.head_dim_) + return replace(resources, key=k, value=v) + """ Gemma-4 decoder block. Per-layer heterogeneity (sliding vs full attention) is handled by switching shape / RoPE table / sliding-window flag at init @@ -176,12 +190,6 @@ def _post_cache_kv(self, cache_kv, infer_state, layer_weight): # ----- Attention kernels (sliding window + per-layer KV reshape) --- - def _att_control(self): - if self.is_sliding and self.sliding_window_ > 0: - w = self.sliding_window_ - 1 - return AttControl(use_sliding_window=True, sliding_window=(w, 0), logical_layer_id=self.layer_num_) - return AttControl(use_sliding_window=False, sliding_window=(-1, -1), logical_layer_id=self.layer_num_) - def _get_layer_kv(self, infer_state: InferStateInfo): # 直接用当前逻辑层号访问；共享层由 State 的 logical_layer_groups 映射到生产层的物理 KV。 state = infer_state.att_state_manager.get_state( @@ -241,9 +249,7 @@ def _context_attention_kernel( # Full-attn layers: head_dim=512, no SWA, no image bidi — standard # triton via backend1. - o_tensor = infer_state.prefill_att_state1.prefill_att( - q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor - ) + o_tensor = self.prefill_attention.forward(infer_state.attention, DenseInputs(_q)) return o_tensor.view(q.shape) def _token_attention_kernel( @@ -253,10 +259,9 @@ def _token_attention_kernel( layer_weight: Gemma4TransformerLayerWeight, out=None, ) -> torch.Tensor: - _k, _v = self._get_layer_kv(infer_state) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - att_state = infer_state.decode_att_state if self.is_sliding else infer_state.decode_att_state1 - o_tensor = att_state.decode_att(q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor) + + o_tensor = self.decode_attention.forward(infer_state.attention, DenseInputs(_q)) return o_tensor.view(q.shape) # ----- FFN (Gemma gelu-tanh, fused gate_up + down) ----------------- diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 4f998ea25e..48d70ef636 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -1,27 +1,30 @@ -from typing import TYPE_CHECKING -import os +from lightllm.common.attention.selection import BackendPolicy import json +import os +from typing import TYPE_CHECKING + import torch -from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend + from lightllm.common.build_utils import repair_config -from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.distributed.communication_op import dist_group_manager from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo -from lightllm.models.gemma4.layer_infer.pre_layer_infer import Gemma4PreLayerInfer from lightllm.models.gemma4.layer_infer.post_layer_infer import Gemma4PostLayerInfer +from lightllm.models.gemma4.layer_infer.pre_layer_infer import Gemma4PreLayerInfer from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer from lightllm.models.gemma4.layer_weights.pre_and_post_layer_weight import Gemma4PreAndPostLayerWeight from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight +from lightllm.models.llama.model import LlamaTpPartModel from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from lightllm.distributed.communication_op import dist_group_manager logger = init_logger(__name__) if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Gemma4TpPartModel(LlamaTpPartModel): @@ -138,24 +141,8 @@ def create_full_att_states( ) return [state] - def _init_att_backend(self): - # Gemma-4 has per-layer heterogeneous attention: sliding layers use - # (head_dim=256, kv_heads=16); full-attn layers use (head_dim=512, - # kv_heads=4, k_eq_v). FA3 caps head_dim at 256 and flashinfer plans - # once per infer_state on a single shape — both unworkable for the - # heterogeneous layout. Both layer kinds go through triton. - # - # Primary backend = sliding layers. Sliding prefill bypasses the - # backend and calls gemma4_mm directly (SWA + image bidi in one - # pass); the prefill_att_state created here is unused but the - # framework requires prefill_att_backend to be non-None. - self.prefill_att_backend = TritonAttBackend(model=self) - self.decode_att_backend = TritonAttBackend(model=self) - - def _init_att_backend1(self): - # Secondary backend = full-attn layers (head_dim=512, plain causal). - self.prefill_att_backend1 = TritonAttBackend(model=self) - self.decode_att_backend1 = TritonAttBackend(model=self) + def attention_policy(self, family, phase): + return BackendPolicy(allow=frozenset({"triton"}), prefer=("triton",)) def _init_custom(self): self._init_to_get_rotary_gemma4() diff --git a/lightllm/models/glm5_next/attention/__init__.py b/lightllm/models/glm5_next/attention/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/glm5_next/attention/fp8_sparse_mla.py b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py new file mode 100644 index 0000000000..ad33f0b255 --- /dev/null +++ b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py @@ -0,0 +1,237 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase +import dataclasses + +import torch + +from lightllm.common.attention.api.execution import WorkspaceBudget, WorkspaceRequirement +from lightllm.common.attention.api.inputs import SparseMLAPrefillInputs +from lightllm.common.attention.backends.sparse_mla.flashmla.fp8 import ( + FlashMLASparseMLAFp8PrefillBackend, + FlashMLASparseMLAFp8DecodeBackend, + NsaFlashMlaFp8SparseDecodeContext, +) + +from .sparse_mla import Glm5NextSparsePrefillContext + + +@dataclasses.dataclass(slots=True) +class Fp8Glm5NextSparsePrefillContext(Glm5NextSparsePrefillContext): + smallq_workspace: object = None + generic_workspace: object = None + generic_binding: object = None + + +@dataclasses.dataclass(slots=True) +class Fp8Glm5NextSparseDecodeContext(NsaFlashMlaFp8SparseDecodeContext): + decode_topk_bound: int = None + + +class GlmSparseMLAFp8Common: + family = AttentionFamily.SPARSE_MLA + + provider = "flashmla_sparse" + + cache_encodings = frozenset({CacheEncoding.FP8_SPARSE_MLA}) + + +class GlmSparseMLAFp8PrefillBackend(GlmSparseMLAFp8Common, FlashMLASparseMLAFp8PrefillBackend): + phase = AttentionPhase.PREFILL + + def __init__(self, env, spec): + super().__init__(env, spec) + self.prefill_graph_enabled = env.prefill_graph + self.prefill_token_capacity = env.max_prefill_tokens or 0 + self.decode_only = env.decode_only + + def get_smallq_workspace(self, context, q): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + allocate_direct_smallq_workspace, + allocate_selected_prefill_workspace, + ) + + context.smallq_workspace = context.workspace.get_object( + "glm_smallq", + lambda: { + "direct": allocate_direct_smallq_workspace(q[:1]), + "selected": allocate_selected_prefill_workspace(q, batch_capacity=8), + }, + expected_type=dict, + ) + if q.shape[0] == 1: + return context.smallq_workspace["direct"] + batch = q.shape[0] + selected = context.smallq_workspace["selected"] + slots = selected["compact"].shape[1] + return { + "compact": selected["compact"][:batch], + "lengths": selected["lengths"][:batch], + "page_table": selected["page_table"][:batch], + "cu_q": selected["cu_q"][: batch + 1], + "cu_k": selected["cu_k"][: batch + 1], + "selected_kv": selected["selected_kv"][: batch * slots], + } + + def get_generic_prefill_workspace(self, context, q, packed_kv, topk_indices): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import KV_GATHER_CHUNK_SIZE + + rows, columns = topk_indices.shape + workspace = context.workspace + tokens = max(self.prefill_token_capacity, rows) + binding = (rows, columns, topk_indices.dtype, packed_kv.shape[0]) + if context.generic_workspace is None or context.generic_binding != binding: + row_table = workspace.scratch("glm_prefill_rows", packed_kv.shape[0], torch.int32)[: packed_kv.shape[0]] + row_table.fill_(-1) + cu_q = workspace.scratch("glm_cu_q", tokens + 1, torch.int32)[: rows + 1] + torch.arange(rows + 1, device=q.device, dtype=torch.int32, out=cu_q) + context.generic_workspace = { + "prefill_row_table": row_table, + "chunk_indices": workspace.scratch("glm_chunk_indices", tokens * columns, topk_indices.dtype)[ + : rows * columns + ].view(rows, columns), + "chunk_lengths": workspace.scratch("glm_chunk_lengths", tokens, torch.int32)[:rows], + "cu_seqlens_q": cu_q, + "row_table_built": False, + } + context.generic_binding = binding + gather_tokens = min(KV_GATHER_CHUNK_SIZE, context.ragged_mem_index.numel()) + if self.decode_only: + gather_tokens = min(gather_tokens, max(1, self.prefill_token_capacity, rows)) + context.generic_workspace["gathered_kv"] = workspace.scratch( + "glm_gathered_kv", max(1, gather_tokens) * 512, q.dtype + )[: gather_tokens * 512].view(gather_tokens, 1, 512) + context.generic_workspace["prefill_row_table"].fill_(-1) + context.generic_workspace["row_table_built"] = False + return context.generic_workspace + + def _run_prefill(self, layer, context, q, kv, inputs): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import context_attention_fwd_fp8 + + if q.shape[1] != 64: + return context_attention_fwd_fp8( + q, + kv, + context.ragged_mem_index, + context.batch.mem_index, + inputs.prefill_cache_kv, + context.ks, + inputs.topk_indices, + self.spec.softmax_scale, + generic_workspace=self.get_generic_prefill_workspace(context, q, kv, inputs.topk_indices), + ) + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + can_use_direct_smallq, + can_use_selected_prefill, + ) + + direct_workspace = ( + self.get_smallq_workspace(context, q) + if can_use_direct_smallq( + q, kv, context.ragged_mem_index, context.batch.mem_index, inputs.prefill_cache_kv, inputs.topk_indices + ) + else None + ) + selected_eligible = direct_workspace is None and can_use_selected_prefill( + q, kv, context.ragged_mem_index, context.batch.mem_index, inputs.prefill_cache_kv, inputs.topk_indices + ) + workspace = self.get_smallq_workspace(context, q) if selected_eligible else None + generic_workspace = None + if direct_workspace is None and (not selected_eligible): + generic_workspace = self.get_generic_prefill_workspace(context, q, kv, inputs.topk_indices) + return context_attention_fwd_fp8( + q, + kv, + context.ragged_mem_index, + context.batch.mem_index, + inputs.prefill_cache_kv, + context.ks, + inputs.topk_indices, + self.spec.softmax_scale, + workspace, + direct_workspace, + generic_workspace, + ) + + def resource_requirements(self, capacity, *, capture): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import KV_GATHER_CHUNK_SIZE + + tokens = max(1, self.prefill_token_capacity, capacity.token_num) + kv_tokens = tokens if self.decode_only else capacity.batch_size * capacity.max_kv_seq_len + tokens + gather_tokens = min(KV_GATHER_CHUNK_SIZE, kv_tokens) + topk = (self.network_config["index_topk"] + 3 + 127) // 128 * 128 + return WorkspaceBudget( + workspace=( + WorkspaceRequirement("glm_gathered_kv", gather_tokens * 512, torch.bfloat16), + WorkspaceRequirement("glm_chunk_indices", tokens * topk, torch.int32), + WorkspaceRequirement("glm_chunk_lengths", tokens, torch.int32), + WorkspaceRequirement("glm_cu_q", tokens + 1, torch.int32), + ) + ) + + def allocate_context(self, capacity, *, capture, workspace) -> Fp8Glm5NextSparsePrefillContext: + return Fp8Glm5NextSparsePrefillContext(capacity, capture, workspace) + + def forward( + self, layer, context: Fp8Glm5NextSparsePrefillContext, inputs: SparseMLAPrefillInputs, *, out=None + ) -> torch.Tensor: + return self._run_prefill(layer, context, inputs.query, layer.resources.kv, inputs) + + +class GlmSparseMLAFp8DecodeBackend(GlmSparseMLAFp8Common, FlashMLASparseMLAFp8DecodeBackend): + phase = AttentionPhase.DECODE + + def _run_decode(self, layer, context, q, packed_kv, inputs): + import flash_mla + + from lightllm.models.glm5_next.triton_kernel import fp8_sparse_decode + from lightllm.models.glm5_next.triton_kernel.pad_query_heads import pad_query_heads + + topk_mem_indices = inputs.topk_mem_indices + if topk_mem_indices.ndim == 2: + topk_mem_indices = topk_mem_indices.unsqueeze(1) + assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1" + q_nope, _ = q + num_heads = q_nope.shape[1] + if fp8_sparse_decode.can_use_fp8_sparse_direct_decode( + q_nope, + packed_kv, + topk_mem_indices, + layer.resources.kv_lora_rank, + context.decode_topk_bound, + is_cuda_graph_capture=context.capture, + ): + return fp8_sparse_decode.fp8_sparse_direct_decode( + q_nope, packed_kv, topk_mem_indices, self.spec.softmax_scale + ) + padded_q = pad_query_heads(q_nope, is_cuda_graph_capture=context.capture) + kv_cache = torch.as_strided( + packed_kv, + size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]), + stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)), + ) + output, _ = flash_mla.flash_mla_with_kvcache( + q=padded_q, + k_cache=kv_cache, + block_table=None, + cache_seqlens=None, + head_dim_v=layer.resources.kv_lora_rank, + tile_scheduler_metadata=( + context.capture_flashmla_sched_meta + if torch.cuda.is_current_stream_capturing() + else context.flashmla_sched_meta + ), + softmax_scale=self.spec.softmax_scale, + causal=False, + is_fp8_kvcache=True, + indices=topk_mem_indices, + ) + return output[:, 0, :num_heads, :] + + def allocate_context(self, capacity, *, capture, workspace) -> Fp8Glm5NextSparseDecodeContext: + return Fp8Glm5NextSparseDecodeContext( + capacity, + capture, + workspace, + decode_topk_bound=2051 + if self.network_config["index_topk"] == 2048 and self.network_config["index_kpool"] == 4 + else None, + ) diff --git a/lightllm/models/glm5_next/attention/sparse_mla.py b/lightllm/models/glm5_next/attention/sparse_mla.py new file mode 100644 index 0000000000..682bcd9213 --- /dev/null +++ b/lightllm/models/glm5_next/attention/sparse_mla.py @@ -0,0 +1,107 @@ +from lightllm.common.attention.api.spec import CacheEncoding, AttentionFamily, AttentionPhase +import dataclasses +import inspect + +import torch + +from lightllm.common.attention.backends.sparse_mla.flashmla.backend import ( + FlashMLASparseMLAPrefillBackend, + FlashMLASparseMLADecodeBackend, + NsaFlashMlaSparseDecodeContext, + NsaFlashMlaSparsePrefillContext, +) + + +@dataclasses.dataclass(slots=True) +class Glm5NextSparsePrefillContext(NsaFlashMlaSparsePrefillContext): + pass + + +@dataclasses.dataclass(slots=True) +class Glm5NextSparseDecodeContext(NsaFlashMlaSparseDecodeContext): + pass + + +class GlmSparseMLACommon: + family = AttentionFamily.SPARSE_MLA + + provider = "flashmla_sparse" + + cache_encodings = frozenset({CacheEncoding.PLAIN}) + + +class GlmSparseMLAPrefillBackend(GlmSparseMLACommon, FlashMLASparseMLAPrefillBackend): + phase = AttentionPhase.PREFILL + + def _run_prefill(self, layer, context, q, kv, inputs): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + tokens, heads, dim = q.shape + padded_heads = (heads + 63) // 64 * 64 + padded_q = q + if padded_heads != heads: + padded_q = q.new_zeros((tokens, padded_heads, dim)) + padded_q[:, :heads] = q + out, _, _ = flash_mla_sparse_fwd( + q=padded_q, + kv=kv, + indices=inputs.topk_mem_indices.unsqueeze(1), + sm_scale=self.spec.softmax_scale, + d_v=layer.resources.kv_lora_rank, + ) + return out[:, :heads] + + def allocate_context(self, capacity, *, capture, workspace) -> Glm5NextSparsePrefillContext: + return Glm5NextSparsePrefillContext(capacity, capture, workspace) + + +class GlmSparseMLADecodeBackend(GlmSparseMLACommon, FlashMLASparseMLADecodeBackend): + phase = AttentionPhase.DECODE + + def _run_decode(self, layer, context, q, kv, inputs): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + q_nope, _ = q + kv_nope = kv.view(-1, 1, 1, layer.resources.kv_lora_rank) + common_kwargs = dict( + page_table=inputs.topk_mem_indices, + cache_seqlens=context.nsa_cache_seqlens, + cu_seqlens_q=context.batch.b1_cu_q_seq_len, + cu_seqlens_k_new=context.nsa_cu_seqlens_k_new, + max_seqlen_q=context.batch.max_q_seq_len, + softmax_scale=self.spec.softmax_scale, + causal=False, + ) + if "only_qv" in inspect.signature(flash_attn_with_kvcache).parameters: + return flash_attn_with_kvcache( + q=None, qv=q_nope, k_cache=None, v_cache=kv_nope, only_qv=True, **common_kwargs + ) + if q_nope.shape[-1] > 256: + page_indices = inputs.topk_mem_indices.clamp_min(0) + key = kv_nope[page_indices, 0, 0] + score = torch.einsum("bhd,bkd->bhk", q_nope.float(), key.float()) + score = score * self.spec.softmax_scale + valid = ( + torch.arange(page_indices.shape[1], device=page_indices.device)[None, :] + < context.nsa_cache_seqlens[:, None] + ) + score = score.masked_fill(~valid[:, None, :], float("-inf")) + weights = score.softmax(dim=-1) + weights = torch.where(valid[:, None, :], weights, torch.zeros_like(weights)) + return torch.einsum("bhk,bkd->bhd", weights, key.float()).to(q_nope.dtype) + return flash_attn_with_kvcache(q=q_nope, k_cache=kv_nope, v_cache=kv_nope, **common_kwargs) + + def allocate_context(self, capacity, *, capture, workspace) -> Glm5NextSparseDecodeContext: + return Glm5NextSparseDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context: Glm5NextSparseDecodeContext): + super().build_device_metadata(context) + pool = self.network_config["index_kpool"] + topk = self.network_config["index_topk"] + context.nsa_cache_seqlens = ( + torch.minimum(context.lengths // pool * pool, torch.full_like(context.lengths, topk)) + + context.lengths % pool + ) + context.nsa_cu_seqlens_k_new = torch.nn.functional.pad( + context.nsa_cache_seqlens.cumsum(0, dtype=torch.int32), (1, 0) + ) diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 31d79697f6..1da3bcd63f 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -1,14 +1,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING import dataclasses - -from lightllm.utils.torch_dtype_utils import get_torch_dtype +from typing import TYPE_CHECKING from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.utils.torch_dtype_utils import get_torch_dtype if TYPE_CHECKING: from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs diff --git a/lightllm/models/glm5_next/glm5_next_visual.py b/lightllm/models/glm5_next/glm5_next_visual.py index 43637c2cbb..c27e68aa76 100644 --- a/lightllm/models/glm5_next/glm5_next_visual.py +++ b/lightllm/models/glm5_next/glm5_next_visual.py @@ -9,10 +9,12 @@ import torch.nn.functional as F from safetensors import safe_open +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton from lightllm.models.vit.triton_kernel.rms_norm_vit import qk_rms_norm, rms_norm -from lightllm.server.visualserver import get_vit_attn_backend +from lightllm.server.visualserver import get_vit_attention_op + from .vision_process import Glm5NextImageProcessor @@ -56,7 +58,7 @@ def forward(self, x, cu_seqlens, max_seqlen, rotary_cos, rotary_sin): q = apply_rotary_pos_emb_triton(q, rotary_cos, rotary_sin) k = apply_rotary_pos_emb_triton(k, rotary_cos, rotary_sin) out = torch.empty_like(q) - get_vit_attn_backend()(q, k, v, out, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=out) return self.proj(out.reshape(x.shape[0], -1)) diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 1033776120..1df0f3732c 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -1,22 +1,17 @@ -# SPDX-License-Identifier: Apache-2.0 - from __future__ import annotations +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFamily + +# SPDX-License-Identifier: Apache-2.0 -from lightllm.common.attention_state.full_att import FullAttentionState import torch +from lightllm.common.attention.api.inputs import KDAInputs, SparseMLADecodeInputs, SparseMLAPrefillInputs from lightllm.common.basemodel import TransformerLayerInferTpl -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.common.basemodel.triton_kernel.mhc import hc_contract, hc_post, hc_pre_norm from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import ( - silu_and_mul_fwd, -) -from lightllm.common.basemodel.triton_kernel.mhc import ( - hc_contract, - hc_post, - hc_pre_norm, -) from lightllm.common.triton_utils.autotuner import Autotuner from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer from lightllm.utils.envs_utils import get_env_start_args @@ -24,6 +19,13 @@ class Glm5NextTransformerLayerInfer(TransformerLayerInferTpl): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace( + description, family=AttentionFamily.KDA if self.is_linear_attention_layer else AttentionFamily.SPARSE_MLA + ) + return description + def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.eps_ = network_config["rms_norm_eps"] @@ -86,7 +88,6 @@ def _ffn(self, input, infer_state, layer_weight): def _ffn_tp(self, input, infer_state, layer_weight): """Dense/shared GLM FFN with the checkpoint's clamp semantics.""" - input = input.view(-1, self.embed_dim_) up_gate_out = layer_weight.gate_up_proj.mm(input) ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) @@ -167,30 +168,14 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], infer_state=infer_state, - att_state=infer_state.prefill_att_state, + att_state=self.prefill_attention.context(infer_state.attention), layer_weight=layer_weight, ) del infer_state.get_topk_indices_params - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - return infer_state.prefill_att_state.prefill_att( - q=q, - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_mem_indices": topk_mem_indices, - "topk_indices": topk_indices, - "prefill_cache_kv": kv, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - }, - ), + + return self.prefill_attention.forward( + infer_state.attention, + SparseMLAPrefillInputs(q, topk_mem_indices, topk_indices, kv), ) def _context_attention_wrapper_run(self, q, cache_kv, infer_state, layer_weight): @@ -253,31 +238,15 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], infer_state=infer_state, - att_state=infer_state.decode_att_state, + att_state=self.decode_attention.context(infer_state.attention), layer_weight=layer_weight, ) del infer_state.get_topk_indices_params q_rope = q_nope[..., :0] - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="mla", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - return infer_state.decode_att_state.decode_att( - q=(q_nope, q_rope), - k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, - v=None, - att_control=AttControl( - nsa_decode=True, - nsa_decode_dict={ - "layer_index": self.layer_num_, - "topk_mem_indices": topk_mem_indices, - "softmax_scale": self.softmax_scale, - "kv_lora_rank": self.kv_lora_rank, - "qk_rope_head_dim": 0, - }, - ), + + return self.decode_attention.forward( + infer_state.attention, + SparseMLADecodeInputs(q_nope, q_rope, topk_mem_indices), ) def _get_o(self, input, infer_state, layer_weight): @@ -339,7 +308,6 @@ def _kda_prefill_cuda_graph_wrapper( output at stable addresses, then register it as a CPU-side replay step between the surrounding CUDA-graph segments. """ - backend = infer_state.prefill_att_state1.backend mixed_qkv = mixed_qkv.contiguous() raw_decay_gate = raw_decay_gate.contiguous() beta_logits = beta_logits.contiguous() @@ -354,28 +322,16 @@ def _kda_prefill_cuda_graph_wrapper( # Construct this shape directly instead of dry-running the kernel: # kernel setup may synchronize with the host, which is illegal while a # CUDA graph capture is active. - output_shape = (1, mixed_qkv.shape[0], backend.tp_num_heads, backend.head_dim) + output_shape = (1, mixed_qkv.shape[0], self.tp_linear_num_heads, self.linear_head_dim) infer_state.prefill_cuda_graph_create_graph_obj() infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() output = torch.empty(output_shape, dtype=mixed_qkv.dtype, device=mixed_qkv.device) _output = tensor_to_no_ref_tensor(output) def kda_prefill_func(new_infer_state): - tmp_output = new_infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkv": _mixed_qkv, - "raw_gate": _raw_decay_gate, - "raw_beta": _beta_logits, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, + tmp_output = self.prefill_attention.forward( + new_infer_state.attention, + KDAInputs(_mixed_qkv, _raw_decay_gate, _beta_logits), ) assert tmp_output.shape == _output.shape _output.copy_(tmp_output) @@ -401,21 +357,9 @@ def context_attention_forward(self, input_embeddings, infer_state, layer_weight) layer_weight, ) else: - core_output = infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkv": qkv, - "raw_gate": raw_decay_gate, - "raw_beta": beta_logits, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, + core_output = self.prefill_attention.forward( + infer_state.attention, + KDAInputs(qkv, raw_decay_gate, beta_logits), ) return self._kda_post(core_output, raw_output_gate, infer_state, layer_weight) @@ -425,22 +369,7 @@ def token_attention_forward(self, input_embeddings, infer_state, layer_weight): qkv, raw_decay_gate, beta_logits, raw_output_gate = self._kda_projections( input_embeddings, infer_state, layer_weight ) - core_output = infer_state.decode_att_state1.decode_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_decode=True, - linear_att_decode_dict={ - "mixed_qkv": qkv, - "raw_gate": raw_decay_gate, - "raw_beta": beta_logits, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + core_output = self.decode_attention.forward(infer_state.attention, KDAInputs(qkv, raw_decay_gate, beta_logits)) return self._kda_post(core_output, raw_output_gate, infer_state, layer_weight) def _hc_pre(self, streams, layer_weight, prefix, norm_weight): diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 3f20ad97ab..14f271e91e 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -1,3 +1,4 @@ +from lightllm.common.attention.selection import BackendPolicy import json import os from typing import TYPE_CHECKING @@ -5,22 +6,24 @@ import torch import triton +from lightllm.common.attention.api import AttentionFamily from lightllm.common.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend -from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.distributed.communication_op import dist_group_manager +from lightllm.models.glm5_next.attention.sparse_mla import GlmSparseMLAPrefillBackend +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + from .layer_infer.pre_layer_infer import Glm5NextPreLayerInfer from .layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from .layer_weights.pre_and_post_layer_weight import Glm5NextPreAndPostLayerWeight from .layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight - if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Glm5NextTpPartModel(TpPartBaseModel): @@ -129,10 +132,7 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" - from lightllm.common.attention_state.full_att import ( - PackedMLALogicalConfig, - create_packed_mla_state, - ) + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -155,15 +155,29 @@ def create_full_att_states( ) return [state] - def prepare_attention_states(self) -> None: - super().prepare_attention_states() - if self.args.llm_kv_type == "fp8kv_dsa" and self.att_state_manager.get_full_attention_states( - self.att_state_role - ): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + def attention_policy(self, family, phase): + if family.is_recurrent: + return BackendPolicy(prefer=("triton",)) + return BackendPolicy( + prefer=("flashmla_sparse",), + extra=( + "lightllm.models.glm5_next.attention.sparse_mla:GlmSparseMLAPrefillBackend", + "lightllm.models.glm5_next.attention.sparse_mla:GlmSparseMLADecodeBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:GlmSparseMLAFp8PrefillBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:GlmSparseMLAFp8DecodeBackend", + ), + ) - # Reserve fixed scratch space before the shared KV capacity is profiled. - Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + def _init_custom(self): + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + if self.args.enable_ep_moe: + dist_group_manager.new_deepep_group( + n_routed_experts=self.config["n_routed_experts"], + hidden_size=self.config["hidden_size"], + expert_quant_method_names=dist_group_manager.get_moe_quant_methods(self.trans_layers_weight), + num_experts_per_tok=self.config["num_experts_per_tok"], + moe_intermediate_size=self.config["moe_intermediate_size"], + ) def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json")) as f: @@ -192,34 +206,3 @@ def _init_some_value(self): self.kv_lora_rank = self.config["kv_lora_rank"] self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - - def _init_att_backend(self): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - ) - - backend_class = ( - Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend - ) - self.prefill_att_backend = backend_class(model=self) - self.decode_att_backend = self.prefill_att_backend - - def _init_att_backend1(self): - from lightllm.common.basemodel.attention.create_linear_utils import ( - get_kda_linear_att_backend_class, - ) - - backend_cls = get_kda_linear_att_backend_class() - self.prefill_att_backend1 = backend_cls(model=self) - self.decode_att_backend1 = self.prefill_att_backend1 - - def _init_custom(self): - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - if self.args.enable_ep_moe: - dist_group_manager.new_deepep_group( - n_routed_experts=self.config["n_routed_experts"], - hidden_size=self.config["hidden_size"], - expert_quant_method_names=dist_group_manager.get_moe_quant_methods(self.trans_layers_weight), - num_experts_per_tok=self.config["num_experts_per_tok"], - moe_intermediate_size=self.config["moe_intermediate_size"], - ) diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py new file mode 100644 index 0000000000..083ef6fe06 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -0,0 +1,744 @@ +import torch +import triton +import triton.language as tl + +# Bound the temporary BF16 gather buffer to 2 GiB (2097152 × 1 × 512 × 2 bytes). +KV_GATHER_CHUNK_SIZE = 2097152 + + +@triton.jit +def _direct_smallq_packed_kernel( + Q, + PackedKV, + RaggedMemIndex, + PrefillMemIndex, + PrefillKV, + KVStarts, + Topk, + PartialO, + PartialLse, + ragged_n, + q_stride_s: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_d: tl.constexpr, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + topk_stride_s: tl.constexpr, + topk_stride_d: tl.constexpr, + HEADS_: tl.constexpr, + BLOCK_H_: tl.constexpr, + PREFILL_N: tl.constexpr, + SPLITS_: tl.constexpr, + BLOCKS_: tl.constexpr, + SLOTS_: tl.constexpr, + softmax_scale: tl.constexpr, +): + qid = tl.program_id(0) + split = tl.program_id(1) + if HEADS_ == 8: + # H8 fits in one fixed 16-lane tile; the head mask excludes inactive lanes. + h = tl.arange(0, 16) + else: + h = tl.program_id(2) * BLOCK_H_ + tl.arange(0, BLOCK_H_) + d = tl.arange(0, 512) + hm = h < HEADS_ + q = tl.load( + Q + qid * q_stride_s + h[:, None] * q_stride_h + d[None, :] * q_stride_d, + hm[:, None], + 0.0, + ).to(tl.bfloat16) + if HEADS_ == 8: + acc = tl.zeros((16, 512), tl.float32) + m = tl.full((16,), -float("inf"), tl.float32) + denominator = tl.zeros((16,), tl.float32) + else: + acc = tl.zeros((BLOCK_H_, 512), tl.float32) + m = tl.full((BLOCK_H_,), -float("inf"), tl.float32) + denominator = tl.zeros((BLOCK_H_,), tl.float32) + start = tl.load(KVStarts + qid).to(tl.int64) + for b in tl.range(0, BLOCKS_, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): + n = tl.arange(0, 64) + slot = split * SLOTS_ + b * 64 + n + top = tl.load( + Topk + qid * topk_stride_s + slot * topk_stride_d, + slot < 2176, + -1, + ).to(tl.int64) + ragged_pos = start + top + valid = (top >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + mem = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + current_row = tl.full((64,), -1, tl.int32) + for j in tl.static_range(PREFILL_N): + current_row = tl.where(valid & (mem == tl.load(PrefillMemIndex + j).to(tl.int64)), j, current_row) + cached = valid & (current_row < 0) + current = valid & (current_row >= 0) + packed_row = PackedKV + mem[:, None] * packed_stride_s + fp8 = tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + d[None, :], + cached[:, None] & (d[None, :] < 448), + 0.0, + ).to(tl.float32) + scales = tl.load( + packed_row + 576 + d[None, :] // 64, + cached[:, None] & (d[None, :] < 448), + 127, + ) + fp8 = fp8 * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + d[None, :] - 448, + cached[:, None] & (d[None, :] >= 448), + 0.0, + ).to(tl.float32) + cached_v = tl.where(d[None, :] < 448, fp8, tail) + current_v = tl.load( + PrefillKV + current_row[:, None] * prefill_stride_s + d[None, :] * prefill_stride_d, + current[:, None], + 0.0, + ).to(tl.float32) + kv = tl.where(current[:, None], current_v, cached_v).to(tl.bfloat16) + score = tl.dot(q, tl.trans(kv), out_dtype=tl.float32) * softmax_scale + score = tl.where(valid[None, :], score, -float("inf")) + block_m = tl.max(score, axis=1) + new_m = tl.maximum(m, block_m) + alpha = tl.where(m == -float("inf"), 0.0, tl.exp(m - new_m)) + p = tl.where(valid[None, :], tl.exp(score - new_m[:, None]), 0.0) + denominator = denominator * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), kv, out_dtype=tl.float32) + m = new_m + has = denominator > 0.0 + out = tl.where(has[:, None], acc / denominator[:, None], 0.0) + tl.store( + PartialO + ((qid * SPLITS_ + split) * HEADS_ + h[:, None]) * 512 + d[None, :], + out, + hm[:, None], + ) + tl.store( + PartialLse + (qid * SPLITS_ + split) * HEADS_ + h, + tl.where(has, m + tl.log(denominator), -float("inf")), + hm, + ) + + +@triton.jit +def reduce_direct_smallq_kernel( + PartialO, + PartialLse, + Out, + HEADS_: tl.constexpr, + SPLITS_: tl.constexpr, + SPAD_: tl.constexpr, +): + qid = tl.program_id(0) + h = tl.program_id(1) + d = tl.program_id(2) * 128 + tl.arange(0, 128) + s = tl.arange(0, SPAD_) + sm = s < SPLITS_ + lse = tl.load(PartialLse + (qid * SPLITS_ + s) * HEADS_ + h, sm, -float("inf")) + maximum = tl.max(lse, axis=0) + safe_maximum = tl.where(maximum == -float("inf"), 0.0, maximum) + weight = tl.where(sm & (lse > -float("inf")), tl.exp(lse - safe_maximum), 0.0) + denominator = tl.sum(weight, axis=0) + partial = tl.load( + PartialO + ((qid * SPLITS_ + s[:, None]) * HEADS_ + h) * 512 + d[None, :], + sm[:, None], + 0.0, + ).to(tl.float32) + value = tl.sum(partial * weight[:, None], axis=0) / tl.where(denominator > 0.0, denominator, 1.0) + tl.store(Out + (qid * HEADS_ + h) * 512 + d, value) + + +def can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return ( + q.ndim == 3 + and 1 <= q.shape[0] <= 8 + and q.dtype == torch.bfloat16 + and (q.shape[1:] == (8, 512) or (q.shape[0] == 1 and q.shape[1:] == (64, 512))) + and packed_kv.ndim == 3 + and packed_kv.shape[1] == 1 + and packed_kv.shape[2] >= 584 + and ragged_mem_index.dtype == torch.int32 + and ragged_mem_index.numel() >= 65536 + and prefill_mem_index.dtype == torch.int32 + and prefill_mem_index.numel() == q.shape[0] + and prefill_kv.dtype == torch.bfloat16 + and prefill_kv.shape == (q.shape[0], 1, 512) + and topk_indices.dtype == torch.int32 + and topk_indices.shape == (q.shape[0], 2176) + ) + + +def allocate_direct_smallq_workspace(q): + qn, heads, _ = q.shape + splits = 34 if qn == 1 else 17 + return { + "partial_o": torch.empty((qn, splits, heads, q.shape[-1]), dtype=torch.float32, device=q.device), + "partial_lse": torch.empty((qn, splits, heads), dtype=torch.float32, device=q.device), + # The reducer requires a contiguous [Q, H, 512] destination. + "output": torch.empty(q.shape, dtype=q.dtype, device=q.device), + } + + +@torch.no_grad() +def _context_attention_fwd_fp8_direct_smallq( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + direct_workspace=None, +): + qn = q.shape[0] + heads = q.shape[1] + # Q1 uses 34 64-key splits; H64 uses two H32 head tiles per split. + splits, blocks, slots = (34, 1, 64) if qn == 1 else (17, 2, 128) + block_h = 32 if heads == 64 else 16 + if heads == 8: + # H8 uses local buffers because its direct path has no reusable workspace. + partial_o = torch.empty((qn, splits, 8, 512), dtype=torch.float32, device=q.device) + partial_lse = torch.empty((qn, splits, 8), dtype=torch.float32, device=q.device) + output = torch.empty(q.shape, dtype=q.dtype, device=q.device) + launch_grid = (qn, splits) + else: + workspace = direct_workspace or allocate_direct_smallq_workspace(q) + partial_o = workspace["partial_o"] + partial_lse = workspace["partial_lse"] + output = workspace["output"] + launch_grid = (qn, splits, triton.cdiv(heads, block_h)) + _direct_smallq_packed_kernel[launch_grid]( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + partial_o, + partial_lse, + ragged_mem_index.numel(), + q.stride(0), + q.stride(1), + q.stride(2), + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + topk_indices.stride(0), + topk_indices.stride(1), + HEADS_=heads, + BLOCK_H_=block_h, + PREFILL_N=qn, + SPLITS_=splits, + BLOCKS_=blocks, + SLOTS_=slots, + softmax_scale=softmax_scale, + num_warps=8, + ) + reduce_direct_smallq_kernel[(qn, heads, 4)]( + partial_o, + partial_lse, + output, + HEADS_=heads, + SPLITS_=splits, + SPAD_=triton.next_power_of_2(splits), + num_warps=4, + ) + return output + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _build_prefill_row_table_kernel(PrefillMemIndex, RowTable, num_tokens, BLOCK: tl.constexpr): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mem_index = tl.load(PrefillMemIndex + row, row < num_tokens, 0) + tl.store(RowTable + mem_index, row, row < num_tokens) + + +@triton.jit(do_not_specialize=["num_tokens"]) +def _gather_prefill_kv_kernel( + PackedKV, + MemIndex, + PrefillRowTable, + PrefillKV, + Output, + num_tokens, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + row = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dim = tl.arange(0, 512) + mem_index = tl.load(MemIndex + row, row < num_tokens, 0).to(tl.int64) + prefill_row = tl.load(PrefillRowTable + mem_index, row < num_tokens, -1) + cached = (row < num_tokens) & (prefill_row < 0) + current = (row < num_tokens) & (prefill_row >= 0) + packed_row = PackedKV + mem_index[:, None] * packed_stride_s + + quantized = tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], + cached[:, None] & (dim[None, :] < 448), + 0.0, + ).to(tl.float32) + scales = tl.load(packed_row + 576 + dim[None, :] // 64, cached[:, None] & (dim[None, :] < 448), 127) + values = quantized * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, + cached[:, None] & (dim[None, :] >= 448), + 0, + ).to(tl.float32) + values = tl.where(dim[None, :] < 448, values, tail) + current_values = tl.load( + PrefillKV + prefill_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + current[:, None], + 0, + ).to(tl.float32) + tl.store( + Output + row[:, None] * 512 + dim[None, :], + tl.where(current[:, None], current_values, values), + row[:, None] < num_tokens, + ) + + +@triton.jit(do_not_specialize=["chunk_end"]) +def _compact_topk_indices_kernel( + TopkIndices, + KVStarts, + ChunkIndices, + ChunkLengths, + chunk_start, + chunk_end, + TOPK: tl.constexpr, + BLOCK: tl.constexpr, + COMPACT: tl.constexpr = True, +): + row = tl.program_id(0) + lane = tl.arange(0, BLOCK) + index = tl.load(TopkIndices + row * TOPK + lane, lane < TOPK, -1) + ragged_index = index + tl.load(KVStarts + row) + valid = (lane < TOPK) & (index >= 0) & (ragged_index >= chunk_start) & (ragged_index < chunk_end) + if COMPACT: + compact_index = tl.cumsum(valid.to(tl.int32)) - 1 + tl.store(ChunkIndices + row * TOPK + compact_index, ragged_index - chunk_start, valid) + tl.store(ChunkLengths + row, tl.sum(valid.to(tl.int32))) + else: + tl.store( + ChunkIndices + row * TOPK + lane, + tl.where(valid, ragged_index - chunk_start, -1), + lane < TOPK, + ) + + +@triton.jit(do_not_specialize=["NUM_Q_TOKENS", "chunk_lse_stride_h"]) +def _merge_att_output_kernel( + Output, + Lse, + ChunkOutput, + ChunkLse, + chunk_lse_stride_s: tl.constexpr, + chunk_lse_stride_h, + NUM_HEADS: tl.constexpr, + NUM_Q_TOKENS, + BLOCK_HEADS: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS) + dim = tl.arange(0, 512) + previous_lse = tl.load(Lse + head * NUM_Q_TOKENS + row, head < NUM_HEADS, -float("inf")) + chunk_lse = tl.load( + ChunkLse + row * chunk_lse_stride_s + head * chunk_lse_stride_h, head < NUM_HEADS, -float("inf") + ) + # FA3 uses natural-log LSE and returns +inf for an empty row. + previous_lse = tl.where(previous_lse == float("inf"), -float("inf"), previous_lse) + chunk_lse = tl.where(chunk_lse == float("inf"), -float("inf"), chunk_lse) + maximum = tl.maximum(previous_lse, chunk_lse) + maximum = tl.where(maximum == -float("inf"), 0, maximum) + previous_weight = tl.exp(previous_lse - maximum) + chunk_weight = tl.exp(chunk_lse - maximum) + denominator = previous_weight + chunk_weight + offset = (row * NUM_HEADS + head[:, None]) * 512 + dim[None, :] + previous = tl.load(Output + offset, (head[:, None] < NUM_HEADS) & (previous_weight[:, None] > 0), 0) + current = tl.load(ChunkOutput + offset, (head[:, None] < NUM_HEADS) & (chunk_weight[:, None] > 0), 0).to(tl.float32) + result = (previous * previous_weight[:, None] + current * chunk_weight[:, None]) / tl.where( + denominator > 0, denominator, 1 + )[:, None] + tl.store(Output + offset, result, head[:, None] < NUM_HEADS) + tl.store( + Lse + head * NUM_Q_TOKENS + row, + tl.where(denominator > 0, maximum + tl.log(denominator), -float("inf")), + head < NUM_HEADS, + ) + + +@triton.jit(do_not_specialize=["ragged_n"]) +def _selected_prefill_compact_kernel( + Topk, + KVStarts, + RaggedMemIndex, + CompactPhysical, + Lengths, + PageTable, + CuQ, + ragged_n, + packed_rows: tl.constexpr, + topk_stride_s: tl.constexpr, + topk_stride_d: tl.constexpr, + BATCH: tl.constexpr, + BLOCK: tl.constexpr, +): + qid = tl.program_id(0) + lane = tl.arange(0, BLOCK) + local = tl.load(Topk + qid * topk_stride_s + lane * topk_stride_d, lane < 2176, -1).to(tl.int64) + ragged_pos = tl.load(KVStarts + qid).to(tl.int64) + local + valid = (lane < 2176) & (local >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + physical = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + valid = valid & (physical >= 0) & (physical < packed_rows) + rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 + tl.store(CompactPhysical + qid * 2176 + rank, physical, valid) + # Every selected slot has a stable contiguous page; invalid tail pages are + # harmless because cache_seqlens limits FA3 to Lengths[qid]. + tl.store(PageTable + qid * 2176 + lane, qid * 2176 + lane, lane < 2176) + if qid == 0: + tl.store(CuQ + lane, lane.to(tl.int32), lane <= BATCH) + tl.store(Lengths + qid, tl.sum(valid.to(tl.int32), axis=0)) + + +@triton.jit +def _selected_prefill_gather_kernel( + CompactPhysical, + Lengths, + PrefillMemIndex, + PrefillKV, + PackedKV, + SelectedKV, + CuK, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + BATCH: tl.constexpr, + BATCH_PAD: tl.constexpr, +): + qid = tl.program_id(0) + row = tl.program_id(1) * 4 + tl.arange(0, 4) + length = tl.load(Lengths + qid) + live = row < length + physical = tl.load(CompactPhysical + qid * 2176 + row, live, 0).to(tl.int64) + current_row = tl.full((4,), -1, tl.int32) + for j in tl.static_range(BATCH_PAD): + current = j < BATCH + current_mem = tl.load(PrefillMemIndex + j, current, 0).to(tl.int64) + current_row = tl.where(live & current & (physical == current_mem), j, current_row) + cached = live & (current_row < 0) + current = live & (current_row >= 0) + dim = tl.arange(0, 512) + packed_row = PackedKV + physical[:, None] * packed_stride_s + group = tl.arange(0, 8) + raw_scales = tl.load(packed_row + 576 + group[None, :], cached[:, None] & (group[None, :] < 7), 127) + scales = (raw_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + scales = tl.broadcast_to(scales[:, :, None], (4, 8, 64)).reshape(4, 512) + fp8 = ( + tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], + cached[:, None] & (dim[None, :] < 448), + 0.0, + ).to(tl.float32) + * scales + ) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, + cached[:, None] & (dim[None, :] >= 448), + 0.0, + ).to(tl.float32) + cached_value = tl.where(dim[None, :] < 448, fp8, tail) + current_value = tl.load( + PrefillKV + current_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + current[:, None], + 0.0, + ).to(tl.float32) + value = tl.where(current[:, None], current_value, cached_value).to(tl.bfloat16) + tl.store(SelectedKV + (qid * 2176 + row[:, None]) * 512 + dim[None, :], value, row[:, None] < 2176) + if (qid == 0) & (tl.program_id(1) == 0): + x = tl.arange(0, BATCH_PAD) + lengths = tl.load(Lengths + x, x < BATCH, 0) + tl.store(CuK, 0) + tl.store(CuK + 1 + x, tl.cumsum(lengths, axis=0), x < BATCH) + + +def can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return ( + q.ndim == 3 + and 2 <= q.shape[0] <= 8 + and q.shape[1:] == (64, 512) + and q.dtype == torch.bfloat16 + and q.is_cuda + and q.stride(-1) == 1 + and packed_kv.shape[1:] == (1, 584) + and packed_kv.dtype == torch.uint8 + and packed_kv.device == q.device + and ragged_mem_index.dtype == torch.int32 + and ragged_mem_index.device == q.device + and ragged_mem_index.numel() >= 65536 + and prefill_mem_index.dtype == torch.int32 + and prefill_mem_index.device == q.device + and prefill_mem_index.shape == (q.shape[0],) + and prefill_kv.dtype == torch.bfloat16 + and prefill_kv.device == q.device + and prefill_kv.shape == (q.shape[0], 1, 512) + and topk_indices.dtype == torch.int32 + and topk_indices.device == q.device + and topk_indices.shape == (q.shape[0], 2176) + ) + + +def allocate_selected_prefill_workspace(q, batch_capacity=None): + batch = q.shape[0] if batch_capacity is None else batch_capacity + return { + "compact": torch.empty((batch, 2176), dtype=torch.int64, device=q.device), + "lengths": torch.empty(batch, dtype=torch.int32, device=q.device), + "page_table": torch.empty((batch, 2176), dtype=torch.int32, device=q.device), + "cu_q": torch.empty(batch + 1, dtype=torch.int32, device=q.device), + "cu_k": torch.empty(batch + 1, dtype=torch.int32, device=q.device), + "selected_kv": torch.empty((batch * 2176, 1, q.shape[-1]), dtype=q.dtype, device=q.device), + } + + +def allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=None): + """Allocate generic prefill metadata up to the backend's planned query capacity.""" + capacity = q.shape[0] if query_capacity is None else query_capacity + return { + "prefill_row_table": torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device), + "chunk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), + "chunk_lengths": torch.empty(capacity, dtype=torch.int32, device=q.device), + "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), + "row_table_built": False, + } + + +@torch.no_grad() +def _context_attention_fwd_fp8_selected_prefill( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + selected_workspace=None, +): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + batch = q.shape[0] + workspace = selected_workspace or allocate_selected_prefill_workspace(q) + batch_pad = triton.next_power_of_2(batch) + _selected_prefill_compact_kernel[(batch,)]( + topk_indices, + kv_starts, + ragged_mem_index, + workspace["compact"], + workspace["lengths"], + workspace["page_table"], + workspace["cu_q"], + ragged_mem_index.numel(), + packed_kv.shape[0], + topk_indices.stride(0), + topk_indices.stride(1), + BATCH=batch, + BLOCK=4096, + num_warps=8, + ) + _selected_prefill_gather_kernel[(batch, triton.cdiv(2176, 4))]( + workspace["compact"], + workspace["lengths"], + prefill_mem_index, + prefill_kv, + packed_kv, + workspace["selected_kv"], + workspace["cu_k"], + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + BATCH=batch, + BATCH_PAD=batch_pad, + num_warps=4, + ) + return flash_attn_with_kvcache( + q=None, + qv=q, + k_cache=None, + v_cache=workspace["selected_kv"].view(-1, 1, 1, 512), + page_table=workspace["page_table"], + cache_seqlens=workspace["lengths"], + cu_seqlens_q=workspace["cu_q"], + cu_seqlens_k_new=workspace["cu_k"], + max_seqlen_q=1, + softmax_scale=softmax_scale, + causal=False, + only_qv=True, + ) + + +@torch.no_grad() +def context_attention_fwd_fp8( + q: torch.Tensor, + packed_kv: torch.Tensor, + ragged_mem_index: torch.Tensor, + prefill_mem_index: torch.Tensor, + prefill_kv: torch.Tensor, + kv_starts: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + selected_workspace=None, + direct_workspace=None, + generic_workspace=None, +): + """Run NoPE sparse prefill with a bounded gather buffer and BF16 current KV.""" + # Direct packed attention is intentionally restricted to the validated + # small-Q GLM-5.3 shape; all other cases retain the bounded-gather path. + if can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return _context_attention_fwd_fp8_direct_smallq( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + direct_workspace, + ) + if can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return _context_attention_fwd_fp8_selected_prefill( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + selected_workspace, + ) + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + num_q_tokens, num_heads, _ = q.shape + gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, ragged_mem_index.numel()) + if generic_workspace is None: + # Build the row table before gathering KV and allocating compaction metadata. + prefill_row_table = torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device) + _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( + prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 + ) + # Keep this potentially 2 GiB buffer local to a DSA layer. + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + chunk_indices = torch.empty_like(topk_indices) + chunk_lengths = torch.empty(num_q_tokens, dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.arange(num_q_tokens + 1, dtype=torch.int32, device=q.device) + else: + prefill_row_table = generic_workspace["prefill_row_table"] + chunk_indices = generic_workspace["chunk_indices"] + chunk_lengths = generic_workspace["chunk_lengths"] + cu_seqlens_q = generic_workspace["cu_seqlens_q"] + if not generic_workspace["row_table_built"]: + _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( + prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 + ) + generic_workspace["row_table_built"] = True + gathered_kv = generic_workspace.get("gathered_kv") + if gathered_kv is None: + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + gather_chunk_size = min(gather_chunk_size, gathered_kv.shape[0]) + output = lse = None + for chunk_start in range(0, ragged_mem_index.numel(), gather_chunk_size): + num_chunk_tokens = min(gather_chunk_size, ragged_mem_index.numel() - chunk_start) + _gather_prefill_kv_kernel[(triton.cdiv(num_chunk_tokens, 4),)]( + packed_kv, + ragged_mem_index[chunk_start : chunk_start + num_chunk_tokens], + prefill_row_table, + prefill_kv, + gathered_kv, + num_chunk_tokens, + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + BLOCK_ROWS=4, + ) + chunk_kv = gathered_kv[:num_chunk_tokens].view(-1, 1, 1, 512) + if ( + num_q_tokens > 8 + and num_heads == 64 + and ragged_mem_index.numel() > num_q_tokens + and ragged_mem_index.numel() <= gather_chunk_size + and q.dtype == torch.bfloat16 + and q.is_contiguous() + and packed_kv.dtype == torch.uint8 + and packed_kv.shape[1:] == (1, 584) + and topk_indices.dtype == torch.int32 + and topk_indices.shape[1] == 2176 + and topk_indices.is_contiguous() + ): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + _compact_topk_indices_kernel[(num_q_tokens,)]( + topk_indices, + kv_starts, + chunk_indices, + chunk_lengths, + chunk_start, + chunk_start + num_chunk_tokens, + TOPK=topk_indices.shape[1], + BLOCK=triton.next_power_of_2(topk_indices.shape[1]), + COMPACT=False, + ) + return flash_mla_sparse_fwd( + q=q, + kv=gathered_kv[:num_chunk_tokens], + indices=chunk_indices.unsqueeze(1), + sm_scale=softmax_scale, + d_v=512, + )[0] + _compact_topk_indices_kernel[(num_q_tokens,)]( + topk_indices, + kv_starts, + chunk_indices, + chunk_lengths, + chunk_start, + chunk_start + num_chunk_tokens, + TOPK=topk_indices.shape[1], + BLOCK=triton.next_power_of_2(topk_indices.shape[1]), + ) + chunk_output, chunk_lse, *_ = flash_attn_with_kvcache( + # only_qv ignores Q/K; these views avoid allocating placeholder tensors. + q=q[..., :64], + k_cache=chunk_kv[..., :64], + qv=q, + v_cache=chunk_kv, + page_table=chunk_indices, + cache_seqlens=chunk_lengths, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=softmax_scale, + causal=False, + only_qv=True, + num_splits=1, + return_softmax_lse=True, + ) + if chunk_start == 0: + output, lse = chunk_output, chunk_lse + else: + # Keep the online reduction in FP32 across all remaining chunks. + output = output.float() + _merge_att_output_kernel[(num_q_tokens, triton.cdiv(num_heads, 4))]( + output, + lse, + chunk_output, + chunk_lse, + chunk_lse.stride(1), + chunk_lse.stride(0), + NUM_HEADS=num_heads, + NUM_Q_TOKENS=num_q_tokens, + BLOCK_HEADS=4, + ) + del chunk_output, chunk_lse + return output.to(q.dtype) diff --git a/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py index 29e55ebaed..a0e455d004 100644 --- a/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py @@ -10,7 +10,9 @@ def context_forward(self, input_embeddings, infer_state, layer_weight): hidden, _, kv = self._project_sparse_qkv_a(normed, infer_state, layer_weight) infer_state.mtp_draft_input_hiddens = None self._post_cache_kv(self._normalize_sparse_kv(kv, layer_weight), infer_state, layer_weight) - self.indexer.write_k_pools(hidden, infer_state, infer_state.prefill_att_state, layer_weight) + self.indexer.write_k_pools( + hidden, infer_state, self.prefill_attention.context(infer_state.attention), layer_weight + ) # BaseModel pads empty ranks to at least one row. An empty result marks # the discarded draft output for our post layer without retaining hiddens. return input_embeddings.new_empty((0, self.embed_dim_)) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 005ef29586..766450f157 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -1,22 +1,24 @@ -from typing import TYPE_CHECKING +from lightllm.common.attention.selection import BackendPolicy from copy import deepcopy +from typing import TYPE_CHECKING + import torch +from lightllm.common.attention.api import AttentionFamily from lightllm.common.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight + from .layer_infer.post_layer_infer import Glm5NextMTPPostLayerInfer from .layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer from .layer_infer.transformer_layer_infer import Glm5NextMTPTransformerLayerInfer from .layer_weights.pre_and_post_layer_weight import Glm5NextMTPPreAndPostLayerWeight - if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs class Glm5NextMTPModel(TpPartBaseModel): @@ -131,8 +133,8 @@ def create_bounded_att_states( ) -> list["BoundedAttentionState"]: """直接选择本模型的 K-pool 窗口层，不依赖 full State 创建过程。""" from lightllm.common.attention_state.bounded_att import ( - create_window_compressor_state, WindowCompressorStateConfig, + create_window_compressor_state, ) from lightllm.utils.torch_dtype_utils import get_torch_dtype @@ -167,20 +169,18 @@ def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, ob config.update(num_hidden_layers=1, n_layer=1, mhc=False) return config - def prepare_attention_states(self) -> None: - super().prepare_attention_states() - if self.args.llm_kv_type == "fp8kv_dsa" and self.att_state_manager.get_full_attention_states( - self.att_state_role - ): - # Reserve fixed scratch space before the shared KV capacity is profiled. - Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) - - def _init_att_backend(self): - backend_class = ( - Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + def attention_policy(self, family, phase): + if family.is_recurrent: + return BackendPolicy(prefer=("triton",)) + return BackendPolicy( + prefer=("flashmla_sparse",), + extra=( + "lightllm.models.glm5_next.attention.sparse_mla:GlmSparseMLAPrefillBackend", + "lightllm.models.glm5_next.attention.sparse_mla:GlmSparseMLADecodeBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:GlmSparseMLAFp8PrefillBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:GlmSparseMLAFp8DecodeBackend", + ), ) - self.prefill_att_backend = backend_class(model=self) - self.decode_att_backend = self.prefill_att_backend def autotune_layers(self): return 1 diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index b6ebc52984..7ac8d12668 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -1,14 +1,23 @@ -from lightllm.common.attention_state.full_att import FullAttentionState +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFeature import torch + +from lightllm.common.attention.api.inputs import DenseInputs from lightllm.models.gpt_oss.layer_weights.transformer_layer_weight import GptOssTransformerLayerWeight -from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer, LlamaInferStateInfo -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaInferStateInfo, LlamaTransformerLayerInfer from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class GptOssTransformerLayerInfer(LlamaTransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace(description, features=description.features | {AttentionFeature.SINK}) + if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention": + return replace(description, window=(self.sliding_window - 1, 0)) + return replace(description, window=None) + def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) self.hidden_size = self.network_config_["hidden_size"] @@ -67,65 +76,15 @@ def _context_attention_kernel( layer_weight: GptOssTransformerLayerWeight, out=None, ): - if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention": - window_size = (self.sliding_window - 1, 0) - use_sliding_window = True - else: - window_size = (-1, -1) - use_sliding_window = False - - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl( - use_sliding_window=use_sliding_window, - sliding_window=window_size, - use_att_sink=True, - sink_weight=layer_weight.attn_sinks.weight, - ), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state.attention, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor def _token_attention_kernel( self, q: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: GptOssTransformerLayerWeight, out=None ): - if self.network_config_["layer_types"][self.layer_num_] == "sliding_attention": - window_size = (self.sliding_window - 1, 0) - use_sliding_window = True - else: - window_size = (-1, -1) - use_sliding_window = False - - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl( - use_sliding_window=use_sliding_window, - sliding_window=window_size, - use_att_sink=True, - sink_weight=layer_weight.attn_sinks.weight, - ), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.decode_attention.forward(infer_state.attention, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor diff --git a/lightllm/models/gpt_oss/model.py b/lightllm/models/gpt_oss/model.py index 29606aba6a..25ae4e3fcd 100644 --- a/lightllm/models/gpt_oss/model.py +++ b/lightllm/models/gpt_oss/model.py @@ -1,10 +1,9 @@ +from lightllm.common.attention.selection import BackendPolicy from lightllm.models.gpt_oss.layer_infer.transformer_layer_infer import GptOssTransformerLayerInfer from lightllm.models.gpt_oss.layer_weights.transformer_layer_weight import GptOssTransformerLayerWeight from lightllm.models.llama.model import LlamaTpPartModel from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.attention import get_prefill_att_backend_class, get_decode_att_backend_class -from lightllm.common.basemodel.attention import BaseAttBackend logger = init_logger(__name__) @@ -19,10 +18,5 @@ class GptOssTpPartModel(LlamaTpPartModel): def __init__(self, kvargs): super().__init__(kvargs) - def _init_att_backend(self): - self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0, priority_list=["fa3"])( - model=self - ) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0, priority_list=["fa3"])( - model=self - ) + def attention_policy(self, family, phase): + return BackendPolicy(allow=frozenset({"fa3"}), prefer=("fa3",)) diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 0b3bb2e867..9887ecdca4 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -1,15 +1,16 @@ -from lightllm.common.attention_state.full_att import FullAttentionState -from lightllm.common.basemodel.attention import AttControl +from functools import partial + import torch -import triton import torch.distributed as dist -from functools import partial -from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight -from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd -from lightllm.models.llama.infer_struct import LlamaInferStateInfo +import triton + +from lightllm.common.attention.api.inputs import DenseInputs from lightllm.common.basemodel import TransformerLayerInferTpl +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.distributed.communication_op import all_gather_into_tensor, reduce_scatter_tensor +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight +from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -46,21 +47,8 @@ def _context_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.prefill_att_state.prefill_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(logical_layer_id=self.layer_num_), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.prefill_attention.forward(infer_state.attention, DenseInputs(_q)) o_tensor = o_tensor.view(q.shape) return o_tensor @@ -70,21 +58,8 @@ def _token_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - state = infer_state.att_state_manager.get_state( - model_role=infer_state.att_state_role, - state_type_filter="gqa", - logical_layer_id=self.layer_num_, - state_class=FullAttentionState, - ) - _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att( - q=_q, - k=_k, - v=_v, - att_control=AttControl(logical_layer_id=self.layer_num_), - alloc_func=self.alloc_tensor, - ) + o_tensor = self.decode_attention.forward(infer_state.attention, DenseInputs(_q)) return o_tensor.view(q.shape) def _att_norm( diff --git a/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py b/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py index b2bf892f90..023ad06748 100644 --- a/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/neo_chat/layer_infer/transformer_layer_infer.py @@ -1,22 +1,30 @@ +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionPhase import os -import torch from functools import partial from typing import Tuple + +import torch +import torch.distributed as dist + +from lightllm.distributed import all_reduce +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd -from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight -from lightllm.distributed import all_reduce -import torch.distributed as dist -from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo +from lightllm.models.qwen3.layer_infer.transformer_layer_infer import Qwen3TransformerLayerInfer from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class NeoChatTransformerLayerInfer(Qwen3TransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace(description, variant="neo" if phase is AttentionPhase.PREFILL else "") + return description + def __init__(self, data_type, network_config): super().__init__(data_type, network_config) return diff --git a/lightllm/models/neo_chat/model.py b/lightllm/models/neo_chat/model.py index 18fa41618a..da0438f38d 100644 --- a/lightllm/models/neo_chat/model.py +++ b/lightllm/models/neo_chat/model.py @@ -1,39 +1,36 @@ -import os +from lightllm.common.attention.selection import BackendPolicy import json +import os + import torch + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer from lightllm.common.build_utils import repair_config +from lightllm.models.internvl.model import InternvlTokenizer +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.neo_chat.layer_infer.transformer_layer_infer import NeoChatTransformerLayerInfer +from lightllm.models.neo_chat.layer_weights.pre_and_post_layer_weight import NeoChatPreAndPostLayerWeight +from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo +from lightllm.models.neo_chat_moe.vision_process import smart_resize +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer +from lightllm.models.qwen3.model import Qwen3TpPartModel +from lightllm.models.qwen3_moe.model import Qwen3MOEModel from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer from lightllm.models.qwen3_vl.layer_infer.transformer_layer_infer import Qwen3VLTransformerLayerInfer from lightllm.models.qwen3_vl.layer_weights.pre_and_post_layer_weight import Qwen3VLPreAndPostLayerWeight -from lightllm.models.qwen2_vl.model import QWen2VLTokenizer -from lightllm.models.qwen3.model import Qwen3TpPartModel -from lightllm.server.core.objs import SamplingParams -from lightllm.models.qwen3_moe.model import Qwen3MOEModel -from lightllm.server.multimodal_params import AudioItem, MultimodalParams, ImageItem -from lightllm.models.neo_chat_moe.vision_process import smart_resize -from lightllm.models.internvl.model import InternvlTokenizer from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer -from lightllm.models.neo_chat.layer_infer.transformer_layer_infer import NeoChatTransformerLayerInfer -from lightllm.models.llama.infer_struct import LlamaInferStateInfo -from lightllm.models.neo_chat.layer_weights.transformer_layer_weight import NeoChatTransformerLayerWeight -from lightllm.models.neo_chat.layer_weights.pre_and_post_layer_weight import NeoChatPreAndPostLayerWeight -from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo -from lightllm.common.basemodel.attention import ( - get_neo_prefill_att_backend_class, - get_decode_att_backend_class, - BaseAttBackend, -) +from lightllm.server.core.objs import SamplingParams +from lightllm.server.multimodal_params import AudioItem, ImageItem, MultimodalParams from lightllm.utils.log_utils import init_logger - logger = init_logger(__name__) class NeoModelBase: def _init_custom(self): - super()._init_custom() if "rope_theta_hw" in self.config: @@ -123,7 +120,6 @@ def _init_to_get_hw_rotary(self, default_base=10000): class NeoTpPartModel(NeoModelBase, Qwen3TpPartModel): - pre_layer_infer_class = LlamaMultimodalPreLayerInfer transformer_layer_infer_class = NeoChatTransformerLayerInfer @@ -139,11 +135,8 @@ def __init__(self, kvargs): def _init_inferstate_cls(self): pass - def _init_att_backend(self): - self.prefill_att_backend: BaseAttBackend = get_neo_prefill_att_backend_class( - index=0, priority_list=["fa3", "triton"] - )(model=self) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) + def attention_policy(self, family, phase): + return BackendPolicy(prefer=("fa3", "triton")) if phase is AttentionPhase.PREFILL else BackendPolicy() def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: diff --git a/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py index 600d066787..aff886e119 100644 --- a/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/neo_chat_moe/layer_infer/transformer_layer_infer.py @@ -1,22 +1,30 @@ +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionPhase import os -import torch from functools import partial from typing import Tuple + +import torch +import torch.distributed as dist + +from lightllm.distributed import all_reduce +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo from lightllm.models.llama.triton_kernel.token_attention_nopad_att1 import token_att_fwd -from lightllm.models.qwen3_moe.layer_infer.transformer_layer_infer import Qwen3MOETransformerLayerInfer +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight -from lightllm.distributed import all_reduce -import torch.distributed as dist -from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.models.qwen3_moe.layer_infer.transformer_layer_infer import Qwen3MOETransformerLayerInfer from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) class NeoChatMOETransformerLayerInfer(Qwen3MOETransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace(description, variant="neo" if phase is AttentionPhase.PREFILL else "") + return description + def __init__(self, data_type, network_config): super().__init__(data_type, network_config) return diff --git a/lightllm/models/neo_chat_moe/model.py b/lightllm/models/neo_chat_moe/model.py index eda5554799..ba382cbd62 100644 --- a/lightllm/models/neo_chat_moe/model.py +++ b/lightllm/models/neo_chat_moe/model.py @@ -1,30 +1,28 @@ -import os +from lightllm.common.attention.selection import BackendPolicy import json +import os + +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer from lightllm.common.build_utils import repair_config +from lightllm.models.internvl.model import InternvlTokenizer +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.neo_chat.model import NeoModelBase +from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo +from lightllm.models.neo_chat_moe.layer_infer.transformer_layer_infer import NeoChatMOETransformerLayerInfer +from lightllm.models.neo_chat_moe.layer_weights.pre_and_post_layer_weight import NeoChatMOEPreAndPostLayerWeight +from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight +from lightllm.models.neo_chat_moe.vision_process import smart_resize +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer +from lightllm.models.qwen3.model import Qwen3TpPartModel +from lightllm.models.qwen3_moe.model import Qwen3MOEModel from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer from lightllm.models.qwen3_vl.layer_infer.transformer_layer_infer import Qwen3VLTransformerLayerInfer from lightllm.models.qwen3_vl.layer_weights.pre_and_post_layer_weight import Qwen3VLPreAndPostLayerWeight -from lightllm.models.qwen2_vl.model import QWen2VLTokenizer -from lightllm.models.qwen3.model import Qwen3TpPartModel -from lightllm.server.core.objs import SamplingParams -from lightllm.models.qwen3_moe.model import Qwen3MOEModel -from lightllm.server.multimodal_params import AudioItem, MultimodalParams, ImageItem -from lightllm.models.neo_chat_moe.vision_process import smart_resize -from lightllm.models.internvl.model import InternvlTokenizer from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer -from lightllm.models.neo_chat_moe.layer_infer.transformer_layer_infer import NeoChatMOETransformerLayerInfer -from lightllm.models.llama.infer_struct import LlamaInferStateInfo -from lightllm.models.neo_chat_moe.layer_weights.transformer_layer_weight import NeoChatMOETransformerLayerWeight -from lightllm.models.neo_chat_moe.layer_weights.pre_and_post_layer_weight import NeoChatMOEPreAndPostLayerWeight -from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer -from lightllm.models.neo_chat_moe.infer_struct import NeoChatInferStateInfo -from lightllm.common.basemodel.attention import ( - get_neo_prefill_att_backend_class, - get_decode_att_backend_class, - BaseAttBackend, -) -from lightllm.models.neo_chat.model import NeoModelBase +from lightllm.server.core.objs import SamplingParams +from lightllm.server.multimodal_params import AudioItem, ImageItem, MultimodalParams IMG_START_TOKEN = "" IMG_END_TOKEN = "" @@ -166,7 +164,6 @@ def get_query_for_t2i(self, prompt: str, input_image_num: int = 0): class NeoTpMOEPartModel(NeoModelBase, Qwen3MOEModel): - pre_layer_infer_class = LlamaMultimodalPreLayerInfer transformer_layer_infer_class = NeoChatMOETransformerLayerInfer @@ -182,11 +179,8 @@ def __init__(self, kvargs): def _init_inferstate_cls(self): pass - def _init_att_backend(self): - self.prefill_att_backend: BaseAttBackend = get_neo_prefill_att_backend_class( - index=0, priority_list=["fa3", "triton"] - )(model=self) - self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) + def attention_policy(self, family, phase): + return BackendPolicy(prefer=("fa3", "triton")) if phase is AttentionPhase.PREFILL else BackendPolicy() def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: diff --git a/lightllm/models/qwen2_5_vl/qwen2_5_visual.py b/lightllm/models/qwen2_5_vl/qwen2_5_visual.py index 1b3a5f0db7..a2e4fc4058 100644 --- a/lightllm/models/qwen2_5_vl/qwen2_5_visual.py +++ b/lightllm/models/qwen2_5_vl/qwen2_5_visual.py @@ -1,21 +1,24 @@ -import os import json +import os +from io import BytesIO +from typing import List, Optional + import torch +import torch.nn as nn import torch.nn.functional as F from PIL import Image -from typing import List, Optional -from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data -from io import BytesIO -import torch.nn as nn -from transformers.activations import ACT2FN -from lightllm.models.qwen2_vl.vision_process import resize_image, Qwen2VLImageProcessor from safetensors import safe_open -from lightllm.server.multimodal_params import ImageItem -from lightllm.models.qwen2_vl.qwen2_visual import PatchEmbed, VisionRotaryEmbedding -from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm -from lightllm.server.visualserver import get_vit_attn_backend +from transformers.activations import ACT2FN + +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.models.qwen2_vl.qwen2_visual import PatchEmbed, VisionRotaryEmbedding from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor, resize_image +from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.server.visualserver import get_vit_attention_op class Qwen2RMSNorm(nn.Module): @@ -74,7 +77,7 @@ def forward( k = apply_rotary_pos_emb_triton(k, rotary_cos, rotary_sin) attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - get_vit_attn_backend()(q, k, v, attn_output, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -350,7 +353,6 @@ def load_image(self, img: List[ImageItem]): return pixel_values.to(dtype=self.data_type), image_grid_thw def load_model(self, weight_dir): - bin_weight_files = [file_ for file_ in os.listdir(weight_dir) if file_.endswith(".bin")] if bin_weight_files: weight_dict = {} diff --git a/lightllm/models/qwen2_vl/qwen2_visual.py b/lightllm/models/qwen2_vl/qwen2_visual.py index e02c3d9aa3..650a2aedb1 100644 --- a/lightllm/models/qwen2_vl/qwen2_visual.py +++ b/lightllm/models/qwen2_vl/qwen2_visual.py @@ -18,24 +18,27 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os import json +import os +from io import BytesIO +from typing import List + import torch +import torch.nn as nn import torch.nn.functional as F from PIL import Image -from typing import List -from torchvision import transforms as T -from lightllm.server.embed_cache.utils import read_shm, get_shm_name_data -from io import BytesIO -import torch.nn as nn +from safetensors import safe_open from torch.nn import LayerNorm +from torchvision import transforms as T from transformers.activations import ACT2FN -from safetensors import safe_open -from lightllm.server.multimodal_params import ImageItem -from lightllm.server.visualserver import get_vit_attn_backend -from lightllm.models.qwen2_vl.vision_process import resize_image, Qwen2VLImageProcessor + +from lightllm.common.attention.api import VarlenAttentionInputs from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor, resize_image +from lightllm.server.embed_cache.utils import get_shm_name_data, read_shm +from lightllm.server.multimodal_params import ImageItem +from lightllm.server.visualserver import get_vit_attention_op # adapted from # https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -147,7 +150,7 @@ def forward( attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - get_vit_attn_backend()(q, k, v, attn_output, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -239,7 +242,6 @@ def _init_datatype(self): return def load_model(self, weight_dir): - processor_config_path = os.path.join(weight_dir, "preprocessor_config.json") with open(processor_config_path, "r") as f: processor_config_dict = json.load(f) diff --git a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py index 84ee0b4059..f38e27b0d4 100644 --- a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py @@ -1,3 +1,5 @@ +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFeature import torch from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_forward @@ -8,6 +10,10 @@ class Qwen3DFlashTransformerLayerInfer(LlamaTransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + return replace(description, features=description.features | {AttentionFeature.NONCAUSAL}) + """DFlash layer inference. The model path is built from two explicit layer primitives: diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index f0683dd882..462c86f74c 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -1,7 +1,5 @@ import torch -from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.models.llama.model import LlamaTpPartModel @@ -46,13 +44,6 @@ def _init_custom(self): self.block_size = int(self.config["block_size"]) self.mask_token_id = int(self.config["mask_token_id"]) - def _init_att_backend(self): - super()._init_att_backend() - # FA3 is currently the only backend that supports non-causal block decode. - # TODO: Remove this restriction after Triton and FlashInfer support non-causal block decode. - if not isinstance(self.decode_att_backend, (Fa3AttBackend, Fp8Fa3AttBackend)): - raise NotImplementedError("Qwen3 DFlash decode requires FA3") - def _init_infer_layer(self, start_layer_index=None): self.draft_layer_start = len(self.main_model.layers_infer) self.draft_layer_start += sum( diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index 7a631622d6..d292cc8fa4 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -1,10 +1,27 @@ +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionPhase +from lightllm.common.attention.api.inputs import DenseInputs from lightllm.common.attention_state.full_att import FullAttentionState -from lightllm.common.basemodel.attention import AttControl -from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer from lightllm.models.qwen3_dflash2.triton_kernel import grouped_dynamic_conv +from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer class Qwen3DFlash2TransformerLayerInfer(Qwen3DFlashTransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + if phase is AttentionPhase.DECODE and self.sliding_window_ > 0: + return replace(description, window=(self.sliding_window_ - 1, self.sliding_window_ - 1)) + return replace(description, window=None) + + def bind_attention_resources(self, model, weight, adapter): + resources = super().bind_attention_resources(model, weight, adapter) + k, v = resources.key, resources.value + return replace( + resources, + key=k.view(-1, self.tp_k_head_num_, self.head_dim_), + value=v.view(-1, self.tp_k_head_num_, self.head_dim_), + ) + """DFlash2 layer with a grouped dynamic convolution around each sublayer.""" def __init__(self, layer_num, network_config): @@ -24,21 +41,11 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - use_sliding_window = self.sliding_window_ > 0 + # DFlash2 attention is non-causal, so its local window must retain both # the visible prefix on the left and the complete draft block on the right. - window = self.sliding_window_ - 1 - output = infer_state.decode_att_state.decode_att( - q=q, - k=k, - v=v, - att_control=AttControl( - logical_layer_id=self.layer_num_, - use_sliding_window=use_sliding_window, - sliding_window=(window, window) if use_sliding_window else (-1, -1), - ), - alloc_func=self.alloc_tensor, - ) + + output = self.decode_attention.forward(infer_state.attention, DenseInputs(q)) return output.view(-1, self.tp_q_head_num_ * self.head_dim_) def _post_cache_kv(self, cache_kv, infer_state, layer_weight): diff --git a/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py b/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py index b9ea66fb15..0fb65beeaf 100644 --- a/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py +++ b/lightllm/models/qwen3_omni_moe_thinker/qwen3_omni_audio.py @@ -1,19 +1,21 @@ -import os import json import math -import torch +import os +from typing import Callable, List, Optional, Union + import numpy as np -from torch import Tensor, nn +import torch from safetensors import safe_open +from torch import Tensor, nn from torch.nn import functional as F -from typing import Callable, Optional, Union, List from transformers.activations import ACT2FN -from lightllm.server.multimodal_params import AudioItem -from lightllm.utils.log_utils import init_logger +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.common.attention.varlen import AutoVarlenAttentionRunner from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager -from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd from lightllm.models.qwen3_omni_moe_thinker.audio_process import WhisperFeatureExtractor +from lightllm.server.multimodal_params import AudioItem +from lightllm.utils.log_utils import init_logger QWEN3_OMNI_CONV_CHUNKSIZE = int(os.getenv("LIGHTLLM_QWEN3_OMNI_CONV_CHUNKSIZE", 200)) @@ -24,7 +26,6 @@ def _get_feat_extract_output_lengths(input_lengths): """ Computes the output length of the convolutional layers and the output length of the audio encoder """ - input_lengths_leave = input_lengths % 100 feat_lengths = (input_lengths_leave - 1) // 2 + 1 output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 @@ -41,10 +42,13 @@ def __init__( activation_function, activation_dropout, encoder_ffn_dim, + attention_op=None, ): super().__init__() self.embed_dim = d_model - self.self_attn = Qwen3OmniMoeAudioAttention(d_model, encoder_attention_heads, attention_dropout) + self.self_attn = Qwen3OmniMoeAudioAttention( + d_model, encoder_attention_heads, attention_dropout, attention_op=attention_op + ) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.dropout = dropout self.activation_fn = ACT2FN[activation_function] @@ -86,7 +90,7 @@ def forward( class Qwen3OmniMoeAudioAttention(nn.Module): - def __init__(self, d_model, encoder_attention_heads, attention_dropout): + def __init__(self, d_model, encoder_attention_heads, attention_dropout, *, attention_op=None): super().__init__() self.embed_dim = d_model self.num_heads = encoder_attention_heads @@ -103,6 +107,7 @@ def __init__(self, d_model, encoder_attention_heads, attention_dropout): self.attention_dropout = 0.0 self.is_decoder = False self.is_causal = False + self.attention_op = attention_op if attention_op is not None else AutoVarlenAttentionRunner() self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) @@ -116,7 +121,6 @@ def forward( **kwargs, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: """Input shape: Batch x Time x Channel""" - seq_length, _ = hidden_states.size() q = self.q_proj(hidden_states).reshape(seq_length, self.num_heads, -1) @@ -125,7 +129,7 @@ def forward( attn_output = g_cache_manager.alloc_tensor(q.shape, q.dtype, device=q.device) - flash_attention_fwd(q, k, v, attn_output, cu_seqlens, max_seqlen) + self.attention_op.forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=attn_output) attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.out_proj(attn_output) @@ -176,6 +180,7 @@ def __init__( self.embed_scale = math.sqrt(self.embed_dim) if scale_embedding else 1.0 self.n_window = n_window self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, self.embed_dim) + attention_op = AutoVarlenAttentionRunner() self.layers = nn.ModuleList( [ Qwen3OmniMoeAudioEncoderLayer( @@ -186,6 +191,7 @@ def __init__( activation_function, activation_dropout, encoder_ffn_dim, + attention_op=attention_op, ) for _ in range(encoder_layers) ] diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 1bc4ef914b..cc700b7766 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -1,25 +1,33 @@ -import torch +from dataclasses import replace +from lightllm.common.attention.api.spec import AttentionFamily +from functools import partial +from typing import Tuple +import torch import torch.distributed as dist -from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( - Qwen3NextTransformerLayerWeight, -) + +from lightllm.common.attention.api.inputs import GDNInputs +from lightllm.distributed import all_reduce from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer +from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo -from lightllm.utils.log_utils import init_logger -from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor -from lightllm.common.basemodel.attention.base_att import AttControl -from typing import Tuple +from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import Qwen3NextTransformerLayerWeight from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ -from lightllm.distributed import all_reduce -from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.utils.envs_utils import get_env_start_args -from functools import partial +from lightllm.utils.log_utils import init_logger +from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor logger = init_logger(__name__) class Qwen3NextTransformerLayerInfer(LlamaTransformerLayerInfer): + def describe_attention(self, phase): + description = super().describe_attention(phase) + description = replace( + description, family=AttentionFamily.GDN if self.is_linear_attention_layer else AttentionFamily.DENSE + ) + return description + def __init__(self, layer_num, network_config): self.partial_rotary_factor = network_config.get("partial_rotary_factor", 1.0) self.n_routed_experts = network_config.get("num_experts", 0) @@ -134,13 +142,13 @@ def _get_qkv( qkv_gate_out = layer_weight.qkvo_gate_proj.mm(input) qkv_out, o_gate = qkv_gate_out.split( [ - self.tp_q_head_num_ * self.head_dim_ * 2 + (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_, + self.tp_q_head_num_ * self.head_dim_ + (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_, self.tp_q_head_num_ * self.head_dim_, ], dim=-1, ) q, cache_kv = qkv_out.split( - [self.tp_q_head_num_ * self.head_dim_ * 2, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_], + [self.tp_q_head_num_ * self.head_dim_, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_], dim=-1, ) infer_state.gate_logics_value = o_gate @@ -184,8 +192,6 @@ def _linear_prefill_cuda_graph_wrapper( infer_state: Qwen3NextInferStateInfo, layer_weight: Qwen3NextTransformerLayerWeight, ) -> Tuple[torch.Tensor, torch.Tensor]: - backend = infer_state.prefill_att_state1.backend - mixed_qkvzba = mixed_qkvzba.contiguous() _mixed_qkvzba = tensor_to_no_ref_tensor(mixed_qkvzba) pre_capture_graph = infer_state.prefill_cuda_graph_get_current_capture_graph() @@ -196,7 +202,11 @@ def _linear_prefill_cuda_graph_wrapper( # We avoid a dry-run because FlashQLA may do host-side syncs while # preparing varlen chunk metadata, which is illegal during capture. num_tokens = mixed_qkvzba.shape[0] - o_shape = (num_tokens, backend.tp_num_v_heads, backend.head_v_dim) + o_shape = ( + num_tokens, + layer_weight.linear_num_v_heads // self.tp_world_size_, + layer_weight.linear_v_head_dim, + ) o_dtype = mixed_qkvzba.dtype o_device = mixed_qkvzba.device z_shape = o_shape @@ -209,23 +219,8 @@ def _linear_prefill_cuda_graph_wrapper( _z = tensor_to_no_ref_tensor(z) def gdn_prefill_func(new_infer_state: "Qwen3NextInferStateInfo"): - # 与 Llama/FA3 一致：通过 new_infer_state.prefill_att_state1 进入， # 使用 runtime 上 init_att_state() 更新过的 buffer idx，而不是 - # capture 时闭包住的 graph prefill_att_state1。 - tmp_o, tmp_z = new_infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkvzba": _mixed_qkvzba, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + tmp_o, tmp_z = self.prefill_attention.forward(new_infer_state.attention, GDNInputs(_mixed_qkvzba)) tmp_o = tmp_o.view(_o.shape) _o.copy_(tmp_o) _z.copy_(tmp_z) @@ -249,20 +244,7 @@ def context_attention_forward( if torch.cuda.is_current_stream_capturing(): core_attn_out, z = self._linear_prefill_cuda_graph_wrapper(mixed_qkvzba, infer_state, layer_weight) else: - core_attn_out, z = infer_state.prefill_att_state1.prefill_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_prefill=True, - linear_att_prefill_dict={ - "mixed_qkvzba": mixed_qkvzba, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + core_attn_out, z = self.prefill_attention.forward(infer_state.attention, GDNInputs(mixed_qkvzba)) gdn_out = self._linear_post(core_attn_out, z, layer_weight) @@ -280,20 +262,7 @@ def token_attention_forward( return super().token_attention_forward(input_embdings, infer_state, layer_weight) mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) - core_attn_out, z = infer_state.decode_att_state1.decode_att( - q=None, - k=None, - v=None, - att_control=AttControl( - linear_att_decode=True, - linear_att_decode_dict={ - "mixed_qkvzba": mixed_qkvzba, - "layer_weight": layer_weight, - "layer_num": self.layer_num_, - }, - ), - alloc_func=self.alloc_tensor, - ) + core_attn_out, z = self.decode_attention.forward(infer_state.attention, GDNInputs(mixed_qkvzba)) gdn_out = self._linear_post(core_attn_out, z, layer_weight) if self.tp_world_size_ > 1: diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 57de61241b..fe0992251a 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -1,23 +1,22 @@ +from typing import TYPE_CHECKING, Optional + import torch -from typing import Optional, TYPE_CHECKING import triton + from lightllm.models.qwen3_moe.model import Qwen3MOEModel -from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( - Qwen3NextTransformerLayerWeight, -) -from lightllm.models.qwen3next.layer_weights.pre_and_post_layer_weight import Qwen3NextPreAndPostLayerWeight -from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import ( - Qwen3NextTransformerLayerInfer, -) +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo +from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerInfer +from lightllm.models.qwen3next.layer_weights.pre_and_post_layer_weight import Qwen3NextPreAndPostLayerWeight +from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import Qwen3NextTransformerLayerWeight from lightllm.utils.log_utils import init_logger -from lightllm.models.qwen3next.config import Qwen3NextResourceConfig if TYPE_CHECKING: from easydict import EasyDict - from lightllm.server.core.objs.start_args_type import StartArgs - from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.server.core.objs.start_args_type import StartArgs logger = init_logger(__name__) @@ -51,7 +50,7 @@ def create_bounded_att_states( device: str = "meta", ) -> list["BoundedAttentionState"]: """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" - from lightllm.common.attention_state.bounded_att import create_gdn_state, ConvSSMStateConfig + from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, create_gdn_state config = cls.normalize_model_config(model_config) assert config["num_attention_heads"] % tp_world_size == 0 @@ -136,20 +135,3 @@ def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch def _init_config(self): super()._init_config() self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - - def _init_att_backend1(self): - role = self.att_state_role - if not self.att_state_manager.get_bounded_attention_states(role): - self.prefill_att_backend1 = None - self.decode_att_backend1 = None - return - from lightllm.common.basemodel.attention.create_linear_utils import ( - get_gdn_linear_prefill_att_backend_class, - get_gdn_linear_decode_att_backend_class, - ) - - linear_prefill_att_backend_class = get_gdn_linear_prefill_att_backend_class(index=1) - linear_decode_att_backend_class = get_gdn_linear_decode_att_backend_class(index=1) - self.prefill_att_backend1 = linear_prefill_att_backend_class(model=self) - self.decode_att_backend1 = linear_decode_att_backend_class(model=self) - return diff --git a/lightllm/models/vit/layer_infer/transformer_layer_infer.py b/lightllm/models/vit/layer_infer/transformer_layer_infer.py index cada55e58b..f8dbf2b7d4 100644 --- a/lightllm/models/vit/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/vit/layer_infer/transformer_layer_infer.py @@ -1,13 +1,15 @@ +from typing import Tuple + import torch import torch.distributed as dist -from typing import Tuple +from lightllm.common.attention.api import VarlenAttentionInputs +from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager from lightllm.models.vit.layer_weights.transformer_layer_weight import ViTTransformerLayerWeight -from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size from lightllm.models.vit.triton_kernel.gelu_vit import gelu_fwd from lightllm.models.vit.triton_kernel.rms_norm_vit import rms_norm -from lightllm.server.visualserver import get_vit_attn_backend -from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.server.visualserver import get_vit_attention_op +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size class ViTTransformerLayerInfer: @@ -105,7 +107,7 @@ def _context_attention_kernel(self, q, k, v) -> torch.Tensor: q, k, v, out = map(reshape, (q, k, v, out)) cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32, device=q.device) * seq_len max_seqlen = seq_len - get_vit_attn_backend()(q, k, v, out, cu_seqlens, max_seqlen) + get_vit_attention_op().forward(VarlenAttentionInputs(q, k, v, cu_seqlens, max_seqlen), out=out) return out.reshape(batch_size, seq_len, -1) def _get_o(self, input, layer_weight: ViTTransformerLayerWeight) -> torch.Tensor: diff --git a/lightllm/platform/__init__.py b/lightllm/platform/__init__.py new file mode 100644 index 0000000000..705eb90a6e --- /dev/null +++ b/lightllm/platform/__init__.py @@ -0,0 +1,72 @@ +"""Device operations used by attention; CPU/CUDA adapters are built in.""" + +import os +from functools import lru_cache + +import torch + +from .api import DeviceCaps, DeviceRuntime, Platform + +_runtime_factories = {} + + +def register_device_runtime(platform, factory): + """Register an additional adapter at worker startup, before creating any runtime. + + The factory receives the original device argument and must resolve/validate it + for its platform. Registering an adapter does not imply provider kernel support. + """ + platform = Platform(platform) + if platform in (Platform.CPU, Platform.CUDA) or platform in _runtime_factories: + raise ValueError(f"device runtime already registered for {platform.value}") + _runtime_factories[platform] = factory + + +def get_device_runtime(device=None, *, platform=None): + platform = Platform(platform) if platform is not None else None + if platform in _runtime_factories: + return _runtime_factories[platform](device) + if platform not in (None, Platform.CPU, Platform.CUDA): + raise ValueError(f"no registered runtime for {platform.value}; built-in adapters cover CUDA and CPU") + use_worker_default = device is None + default_type = "cpu" if platform is Platform.CPU else "cuda" + if isinstance(device, int): + device = f"{default_type}:{device}" + device = device if device is not None else default_type + if isinstance(device, str) and device.split(":")[0] not in ("cpu", "cuda", "meta"): + raise ValueError(f"attention backend refactor supports CUDA and CPU devices only, got {device}") + device = torch.device(device) + platform = ( + Platform(platform) + if platform is not None + else { + "cpu": Platform.CPU, + "meta": Platform.CPU, + "cuda": Platform.CUDA, + }.get(device.type) + ) + if platform is None: + raise ValueError(f"no device runtime for {device}") + if platform is Platform.CPU and device.type not in ("cpu", "meta"): + raise ValueError(f"CPU runtime cannot operate on {device}") + if platform is not Platform.CPU and device.type in ("cpu", "meta"): + raise ValueError(f"{platform.value} runtime cannot operate on {device}") + # 保留调用方指定的 CUDA device；没有显式 ordinal 时使用 worker 绑定的 ordinal。 + if device.index is None and device.type not in ("cpu", "meta"): + index = os.environ.get("LIGHTLLM_CURRENT_DEVICE_ID") if use_worker_default else None + if index is None: + api = torch.cuda + index = api.current_device() if api.is_available() else 0 + device = torch.device(device.type, int(index)) + return _runtime(platform, device) + + +@lru_cache(maxsize=None) +def _runtime(platform, device): + if platform is Platform.CPU: + return DeviceRuntime(device) + if platform is Platform.CUDA: + from .cuda import CudaRuntime + + return CudaRuntime(device) + raise ValueError(f"no device runtime for platform {platform.value}") diff --git a/lightllm/platform/api.py b/lightllm/platform/api.py new file mode 100644 index 0000000000..aaaba26c47 --- /dev/null +++ b/lightllm/platform/api.py @@ -0,0 +1,96 @@ +from contextlib import nullcontext +from dataclasses import dataclass +from functools import cached_property + +import torch + +from lightllm.utils.device_utils import Platform + + +@dataclass(frozen=True) +class DeviceCaps: + """实际设备的能力快照；torch 的 device.type 不足以区分 CUDA、ROCm 和兼容层。""" + + platform: Platform + device: torch.device + available: bool = True + device_name: str = "" + capability: tuple[int, int] | None = None + supports_graph: bool = False + + +class DeviceRuntime: + """一个 worker/device 的设备操作。实现库和模型语义不属于 runtime。 + + 构造不切换当前设备；调用方先完成 worker 的设备绑定。caps 延迟读取， + 注册表和 CPU 契约测试可以在不初始化设备驱动的情况下导入这些接口。 + """ + + platform = Platform.CPU + graph_type = "" + distributed_backend = "gloo" + + def __init__(self, device, api=None): + self.device = torch.device(device) + self.api = api + + @cached_property + def caps(self): + if self.api is None: + return DeviceCaps(self.platform, self.device, available=self.device.type != "meta") + available = self.api.is_available() + capability = None + if available and hasattr(self.api, "get_device_capability"): + capability = tuple(self.api.get_device_capability(self.device)) + return DeviceCaps( + self.platform, + self.device, + available, + self.api.get_device_name(self.device) if available else "", + capability, + available and hasattr(self.api, self.graph_type), + ) + + def allocate(self, shape, dtype=torch.uint8): + return torch.empty(shape, dtype=dtype, device=self.device) + + def current_stream(self): + return self.api.current_stream(self.device) if self.api is not None else None + + def create_stream(self): + return self.api.Stream(device=self.device) if self.api is not None else None + + def use_stream(self, stream): + return self.api.stream(stream) if self.api is not None else nullcontext() + + def create_event(self, *, enable_timing=False): + return self.api.Event(enable_timing=enable_timing) if self.api is not None else None + + def synchronize(self): + if self.api is not None: + self.api.synchronize(self.device) + + def is_capturing(self): + return self.api is not None and self.api.is_current_stream_capturing() + + def empty_cache(self): + if self.api is not None: + self.api.empty_cache() + + def memory_info(self): + if self.api is None: + raise ValueError("host/meta storage requires an explicit capacity or a memory-budget runtime") + return self.api.mem_get_info(self.device) + + def graph_pool(self): + return self.api.graph_pool_handle() if self.caps.supports_graph else None + + def create_graph(self): + if not self.caps.supports_graph: + raise RuntimeError(f"{self.platform.value} device {self.device} does not support Graph capture") + return getattr(self.api, self.graph_type)() + + def capture_graph(self, graph, *, pool=None, memory_saver=None): + if memory_saver is not None and getattr(memory_saver, "enabled", False): + raise ValueError(f"memory saver is not supported on {self.platform.value}") + return self.api.graph(graph, pool=pool) diff --git a/lightllm/platform/cuda.py b/lightllm/platform/cuda.py new file mode 100644 index 0000000000..0db5e4891b --- /dev/null +++ b/lightllm/platform/cuda.py @@ -0,0 +1,17 @@ +import torch + +from .api import DeviceRuntime, Platform + + +class CudaRuntime(DeviceRuntime): + platform = Platform.CUDA + graph_type = "CUDAGraph" + distributed_backend = "nccl" + + def __init__(self, device): + super().__init__(device, torch.cuda) + + def capture_graph(self, graph, *, pool=None, memory_saver=None): + if memory_saver is not None: + return memory_saver.cuda_graph(graph, pool=pool) + return self.api.graph(graph, pool=pool) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 16a8735a7c..0f0b455966 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -2,6 +2,28 @@ import json +class AttentionBackendAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + import warnings + + if len(values) > 2: + raise argparse.ArgumentError(self, "expected a backend and optional legacy linear backend") + if values[0] not in ("auto", "triton", "fa3", "flashinfer", "flashmla_sparse"): + raise argparse.ArgumentError(self, f"invalid choice: {values[0]}") + setattr(namespace, self.dest, values[0]) + if len(values) == 2: + warnings.warn( + f"{option_string} BACKEND LINEAR_BACKEND is deprecated; use the separate linear attention option", + FutureWarning, + stacklevel=2, + ) + linear = self.dest.replace("_att_backend", "_linear_att_backend") + allowed = ("auto", "triton", "flashqla") if "prefill" in self.dest else ("auto", "triton") + if values[1] not in allowed: + raise argparse.ArgumentError(self, f"unsupported legacy linear backend: {values[1]}") + setattr(namespace, linear, values[1]) + + def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.formatter_class = argparse.RawTextHelpFormatter @@ -495,26 +517,36 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--llm_prefill_att_backend", type=str, nargs="+", - choices=["auto", "triton", "fa3", "flashinfer", "flashqla"], - default=["auto"], - help="""prefill attention kernel used in llm. - auto: automatically select best backend based on GPU and available packages - (priority: fa3 > flashinfer > triton) - for hybrid linear-attention models, the second value selects the linear-attention backend - (priority: flashqla > triton); when omitted, it defaults to auto""", + action=AttentionBackendAction, + choices=["auto", "triton", "fa3", "flashinfer", "flashmla_sparse", "flashqla"], + default="auto", + help="""prefill attention backend. auto selects fa3 > flashinfer > triton; + NSA uses flashmla_sparse. Linear attention is configured separately.""", ) parser.add_argument( "--llm_decode_att_backend", type=str, nargs="+", - choices=["auto", "triton", "fa3", "flashinfer"], - default=["auto"], - help="""decode attention kernel used in llm. - auto: automatically select best backend based on GPU and available packages - (priority when mtp_step > 0: fa3 > flashinfer > triton; - otherwise: flashinfer > fa3 > triton) - for hybrid linear-attention models, the second value selects the linear-attention backend - (currently triton only); when omitted, it defaults to auto""", + action=AttentionBackendAction, + choices=["auto", "triton", "fa3", "flashinfer", "flashmla_sparse"], + default="auto", + help="""decode attention backend. auto prefers fa3 with MTP, otherwise flashinfer; + NSA uses flashmla_sparse. Linear attention is configured separately.""", + ) + parser.add_argument( + "--llm_prefill_linear_att_backend", + type=str, + choices=["auto", "triton", "flashqla"], + default="auto", + help="""linear attention prefill backend. GDN auto selects flashqla > triton; + KDA supports triton only.""", + ) + parser.add_argument( + "--llm_decode_linear_att_backend", + type=str, + choices=["auto", "triton"], + default="auto", + help="linear attention decode backend (currently triton only).", ) parser.add_argument( "--page_size", @@ -536,9 +568,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--vit_att_backend", type=str, - nargs="+", choices=["auto", "triton", "fa3", "sdpa", "xformers"], - default=["auto"], + default="auto", help="""vit attention kernel used in vlm. auto: automatically select best backend based on GPU and available packages (priority: fa3 > xformers > sdpa > triton)""", @@ -1018,8 +1049,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: "--hardware_platform", type=str, default="cuda", - choices=["cuda", "musa"], - help="""Hardware platform: cuda | musa""", + choices=["cuda"], + help="Hardware platform: cuda", ) parser.add_argument( "--enable_torch_fallback", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index e0e6038d60..5a0dc9739d 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -1,40 +1,42 @@ +import math import multiprocessing as mp import os -import uuid import subprocess -import math -from lightllm.utils.start_utils import configure_decode_open_file_limit, process_manager -from .metrics.manager import start_metric_manager -from .httpserver.req_id_generator_server import start_req_id_generator_server -from .embed_cache.manager import start_cache_manager -from lightllm.utils.log_utils import init_logger +import uuid + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.config_utils import ( + auto_set_cpu_cache_token_page_size, + auto_set_fused_shared_experts, + auto_set_max_req_total_len, + auto_set_page_size, + auto_set_response_parsers, + get_default_disable_chunked_mm_input, + get_running_max_req_size_per_dp, + has_audio_module, + has_vision_module, + is_hybrid_att_model, +) +from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args from lightllm.utils.envs_utils import ( get_env_start_args, get_unique_server_name, set_env_start_args, set_unique_server_name, ) -from lightllm.utils.shm_port_args import get_shm_port_args -from lightllm.utils.net_utils import validate_ports -from .detokenization.manager import start_detokenization_process -from .router.manager import start_router_process +from lightllm.utils.log_utils import init_logger from lightllm.utils.multinode_utils import send_and_receive_node_ip +from lightllm.utils.net_utils import validate_ports from lightllm.utils.redis_utils import start_redis_service +from lightllm.utils.shm_port_args import get_shm_port_args from lightllm.utils.shm_size_check import check_recommended_shm_size -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.config_utils import ( - has_audio_module, - has_vision_module, - is_hybrid_att_model, - auto_set_max_req_total_len, - auto_set_page_size, - auto_set_cpu_cache_token_page_size, - auto_set_fused_shared_experts, - auto_set_response_parsers, - get_running_max_req_size_per_dp, - get_default_disable_chunked_mm_input, -) -from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args +from lightllm.utils.start_utils import configure_decode_open_file_limit, process_manager + +from .detokenization.manager import start_detokenization_process +from .embed_cache.manager import start_cache_manager +from .httpserver.req_id_generator_server import start_req_id_generator_server +from .metrics.manager import start_metric_manager +from .router.manager import start_router_process logger = init_logger(__name__) @@ -184,17 +186,11 @@ def _launch_subprocesses(args: StartArgs): assert args.enable_tpsp_mix_mode and args.dp > 1, "need set --enable_tpsp_mix_mode firstly and --dp > 1" if args.enable_ep_moe: - allowed_ep_prefill_att_backends = {"auto", "fa3", "triton", "flashqla"} - for backend in args.llm_prefill_att_backend: - assert backend in allowed_ep_prefill_att_backends, ( - "When --enable_ep_moe is enabled, --llm_prefill_att_backend must be one of " - f"{sorted(allowed_ep_prefill_att_backends)}; flashinfer is not supported." - ) - allowed_ep_decode_att_backends = {"auto", "fa3", "triton"} - for backend in args.llm_decode_att_backend: - assert backend in allowed_ep_decode_att_backends, ( - "When --enable_ep_moe is enabled, --llm_decode_att_backend must be one of " - f"{sorted(allowed_ep_decode_att_backends)}; flashinfer is not supported." + allowed_ep_att_backends = {"auto", "fa3", "triton"} + for option in ("llm_prefill_att_backend", "llm_decode_att_backend"): + assert getattr(args, option) in allowed_ep_att_backends, ( + f"When --enable_ep_moe is enabled, --{option} must be one of " + f"{sorted(allowed_ep_att_backends)}; flashinfer is not supported." ) if args.page_size < 1: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0609132e62..93edb9af65 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -175,21 +175,20 @@ class StartArgs: vit_quant_type: Optional[str] = field(default="none") vit_quant_cfg: Optional[str] = field(default=None) expert_dtype: Optional[str] = field(default=None, metadata={"choices": ["fp8", "fp4"]}) - llm_prefill_att_backend: List[str] = field( - default_factory=lambda: ["auto"], - metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashqla"]}, + llm_prefill_att_backend: str = field( + default="auto", metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashmla_sparse"]} ) - llm_decode_att_backend: List[str] = field( - default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "flashinfer"]} + llm_decode_att_backend: str = field( + default="auto", metadata={"choices": ["auto", "triton", "fa3", "flashinfer", "flashmla_sparse"]} ) + llm_prefill_linear_att_backend: str = field(default="auto", metadata={"choices": ["auto", "triton", "flashqla"]}) + llm_decode_linear_att_backend: str = field(default="auto", metadata={"choices": ["auto", "triton"]}) page_size: int = field(default=1) req_token_index_mode: str = field( default="pin_mem_cache", metadata={"choices": ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]}, ) - vit_att_backend: List[str] = field( - default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]} - ) + vit_att_backend: str = field(default="auto", metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]}) llm_kv_type: str = field( default="None", metadata={"choices": ["None", "int8kv", "int4kv", "fp8kv_sph", "fp8kv_spt", "fp8kv_dsa"]} ) @@ -246,7 +245,7 @@ class StartArgs: enable_rl: bool = field(default=False) enable_torch_memory_saver: bool = field(default=False) enable_weight_cpu_backup: bool = field(default=False) - hardware_platform: str = field(default="cuda", metadata={"choices": ["cuda", "musa"]}) + hardware_platform: str = field(default="cuda", metadata={"choices": ["cuda"]}) enable_torch_fallback: bool = field(default=False) enable_triton_fallback: bool = field(default=False) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 09eb22be36..1886f8b374 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -1,56 +1,60 @@ +import threading +import time +from typing import Callable, List, Optional, Tuple + import numpy as np import torch -import time -import threading import torch.distributed as dist -from typing import List, Tuple, Callable, Optional from transformers.configuration_utils import PretrainedConfig -from lightllm.utils.infer_utils import set_random_seed -from lightllm.utils.log_utils import init_logger -from lightllm.models import get_draft_model_class, get_model -from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack -from lightllm.server.router.model_infer.req_token_index import ReqTokenIndex -from lightllm.server.router.token_load import TokenLoad + from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.common.basemodel.logprobs_manager import PromptLogprobsCaptureManager from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager -from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache -from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache -from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput -from lightllm.utils.dist_utils import init_distributed_env +from lightllm.common.basemodel.triton_kernel.gather_token_id import scatter_token +from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async +from lightllm.distributed import dist_group_manager +from lightllm.distributed.communication_op import all_gather_into_tensor, all_reduce, broadcast +from lightllm.models import get_draft_model_class, get_model from lightllm.server.core.objs import ShmReqManager, StartArgs from lightllm.server.core.objs.io_objs import AbortedReqCmd -from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer +from lightllm.server.multi_level_kv_cache import CacheTier, create_cache_placement_controller +from lightllm.server.pd_io_struct import PDChunckedTransTaskRet +from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache +from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache +from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack, g_infer_context +from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager, OverlapEventPack from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -from lightllm.utils.dist_utils import get_global_rank, get_global_world_size, get_dp_size -from lightllm.utils.dist_utils import get_dp_world_size, get_global_dp_rank, get_current_rank_in_dp -from lightllm.utils.dist_utils import get_current_device_id, get_current_rank_in_node, get_node_world_size -from lightllm.utils.dist_utils import get_dp_rank_in_node, create_new_group_for_current_node +from lightllm.server.router.model_infer.req_token_index import ReqTokenIndex +from lightllm.server.router.token_load import TokenLoad +from lightllm.utils.dist_utils import ( + create_new_group_for_current_node, + get_current_device_id, + get_current_rank_in_dp, + get_current_rank_in_node, + get_dp_rank_in_node, + get_dp_size, + get_dp_world_size, + get_global_dp_rank, + get_global_rank, + get_global_world_size, + get_node_world_size, + init_distributed_env, +) from lightllm.utils.envs_utils import ( - get_env_start_args, enable_radix_tree_timer_merge, + get_env_start_args, get_radix_tree_merge_update_delta, ) -from lightllm.distributed import dist_group_manager -from lightllm.distributed.communication_op import ( - all_gather_into_tensor, - all_reduce, - broadcast, -) -from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer -from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager, OverlapEventPack -from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample -from lightllm.common.basemodel.triton_kernel.gather_token_id import scatter_token -from lightllm.server.pd_io_struct import PDChunckedTransTaskRet -from lightllm.server.multi_level_kv_cache import ( - CacheTier, - create_cache_placement_controller, -) -from .prefill_queue_strategy import create_prefill_queue_strategy -from .multi_level_kv_cache import MultiLevelKvCacheModule +from lightllm.utils.infer_utils import set_random_seed +from lightllm.utils.log_utils import init_logger from lightllm.utils.profiler import ProcessProfiler, ProfilerCmd +from .multi_level_kv_cache import MultiLevelKvCacheModule +from .prefill_queue_strategy import create_prefill_queue_strategy + class ModeBackend: def __init__(self) -> None: @@ -149,10 +153,10 @@ def init_model(self, kvargs): self.draft_models = [] if self.args.mtp_mode is not None: self.init_mtp_draft_model(model_kvargs) - from lightllm.utils.envs_utils import set_model_init_status - from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.attention_state import AttentionStateManager from lightllm.common.req_manager import ReqManager + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag models = (self.model, *self.draft_models) assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" @@ -169,11 +173,19 @@ def init_model(self, kvargs): item.prepare_attention_states() manager: AttentionStateManager = self.model.att_state_manager assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" + manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + + for item in models: + item.attention.reserve(item) with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - manager.finish_state_registration() manager.load_full_att_calibration(args=self.model.args) manager.initialize_storage( - self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction + self.model.args, + token_capacity=self.model.max_total_token_num, + mem_fraction=self.model.mem_fraction, + device=self.model.device_runtime.device, ) for item in models: item.init_runtime() diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 9bc8c16cf7..22c45a29c1 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -1,23 +1,22 @@ -import torch import time from typing import List + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc +from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack -from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.router.model_infer.mode_backend.pre import ( - prepare_prefill_inputs, - prepare_decode_inputs, -) from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample -from lightllm.server.router.model_infer.infer_batch import g_infer_context -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack +from lightllm.server.router.model_infer.mode_backend.pre import prepare_decode_inputs, prepare_prefill_inputs from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils -from lightllm.utils.log_utils import init_logger +from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.dist_utils import get_current_device_id -from .control_state import ControlState from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from .control_state import ControlState logger = init_logger(__name__) @@ -291,11 +290,7 @@ def decode_mtp( b_req_mtp_start_loc = gen_b_req_mtp_start_loc(model_input.b_mtp_index, num_reqs=req_num) mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( - backend=self, - next_token_ids=next_token_ids, - b_req_idx=model_input.b_req_idx, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_mtp_index=model_input.b_mtp_index, + backend=self, next_token_ids=next_token_ids, model_output=model_output, num_reqs=req_num ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index ed7e5bb167..b7852950c5 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -1,24 +1,27 @@ -import torch import time from typing import List, Tuple + +import torch + +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc +from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput -from lightllm.server.router.model_infer.infer_batch import g_infer_context, InferReq from lightllm.server.router.model_infer.mode_backend.generic_post_process import sample +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack from lightllm.server.router.model_infer.mode_backend.pre import ( - prepare_prefill_inputs, - prepare_decode_inputs, - overlap_prepare_prefill_inputs, overlap_prepare_decode_inputs, + overlap_prepare_prefill_inputs, + prepare_decode_inputs, + prepare_prefill_inputs, ) -from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack +from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_engine import DPOverlapSpecEngine +from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_env_start_args -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager -from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_engine import DPOverlapSpecEngine -from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils + from .control_state import DPControlState @@ -485,7 +488,6 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): """复用普通 SpecEngine 执行 DP speculative draft-and-verify。""" - model_input, run_reqs = prepare_decode_inputs(req_objs=decode_reqs) spec_engine = self.spec_engine req_num = len(decode_reqs) @@ -526,11 +528,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): ) mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( - backend=self, - next_token_ids=next_token_ids, - b_req_idx=model_input.b_req_idx, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_mtp_index=model_input.b_mtp_index, + backend=self, next_token_ids=next_token_ids, model_output=model_output, num_reqs=req_num ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -802,9 +800,8 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( backend=self, next_token_ids=next_token_ids, - b_req_idx=b_req_idx, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_mtp_index=b_mtp_index, + model_output=(model_output0, model_output1), + num_reqs=req_num, ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] mtp_accept_len1 = mtp_accept_len[real_request_num0:] diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index bfb55d37a2..a750111f2b 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -1,19 +1,22 @@ -import torch.distributed as dist -import torch -import dataclasses import bisect -from functools import lru_cache -from typing import Optional, List, Deque +import dataclasses from collections import deque +from functools import lru_cache +from typing import Deque, List, Optional + +import torch +import torch.distributed as dist + from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -from lightllm.utils.envs_utils import get_env_start_args -from ..infer_batch import InferReq -from lightllm.utils.dist_utils import create_new_group_for_current_dp from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.utils.dist_utils import create_new_group_for_current_dp +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len +from ..infer_batch import InferReq + logger = init_logger(__name__) @@ -45,20 +48,14 @@ def __init__(self, backend): @lru_cache() def need_sync_compute_stream(self) -> bool: """不支持异步 CPU cache 的 backend（目前是 FA3）仍要等计算流，否则搬运可能崩溃。""" - model = self.backend.model - att_backends = [ - model.prefill_att_backend, - model.decode_att_backend, - model.prefill_att_backend1, - model.decode_att_backend1, - ] - for att_backend in att_backends: - if att_backend is not None and att_backend.requires_cpu_cache_sync: - logger.info( - "MultiLevelKvCacheModule: %s requires synchronous CPU cache IO.", type(att_backend).__name__ - ) - return True + for group in model.attention.groups: + for backend in group.values(): + if backend.requires_cpu_cache_sync: + logger.info( + "MultiLevelKvCacheModule: %s requires synchronous CPU cache IO.", type(backend).__name__ + ) + return True logger.info("MultiLevelKvCacheModule: CPU cache load waits on a recorded event.") return False @@ -282,8 +279,13 @@ def write_req_to_cpu(self, req: InferReq) -> Optional["OffloadTask"]: """选择请求可缓存的前缀，分配 CPU 页并提交写出，返回待确认完成的任务。""" assert CacheTier.CPU in req.cache_tiers disk_offload_enable = CacheTier.DISK in req.cache_tiers - # KV 和 bounded 快照均在 offload 流上复制，由后续完成事件统一确认就绪。 - with torch.cuda.stream(self.get_cpu_offload_stream()): + cpu_kv_cache_stream = self.get_cpu_offload_stream() + producer = torch.cuda.current_stream() + cpu_kv_cache_stream.wait_stream(producer) + overlap = g_infer_context.get_overlap_stream() + if overlap != producer: + cpu_kv_cache_stream.wait_stream(overlap) + with torch.cuda.stream(cpu_kv_cache_stream): # 综合考虑后只对prompt做缓存管理，不包含decode内容，这里与radix cache不一致 token_hash_list = req.shm_req.token_hash_list.get_all() page_len_list = req.shm_req.token_hash_page_len_list.get_all() diff --git a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py index b2f3285eee..2868d1855e 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py +++ b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py @@ -50,7 +50,7 @@ def __init__( self.backend = backend self.max_draft_step = int(backend.max_draft_step) self.block_size = int(backend.draft_models[0].block_size) if self.spec_mode in ("dflash", "dflash2") else None - self.draft_steps = self._get_draft_steps() + self.draft_steps = self._get_decode_draft_steps() self.target_infer_costs = _InferCostMsTable() self.draft_infer_costs = _InferCostMsTable() @@ -189,7 +189,7 @@ def update_statics( verified_draft_step=plan.pre_draft_step, ) - def _get_draft_steps(self) -> Tuple[int, ...]: + def _get_decode_draft_steps(self) -> Tuple[int, ...]: """Return the draft configurations supported by the current MTP mode.""" if self.spec_mode in ("vanilla_no_att", "eagle_no_att"): diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index 19d4656f14..bbb24a3c0e 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -5,7 +5,9 @@ import torch +from lightllm.common.attention_state.commit import RequestGenerationStamp from lightllm.common.basemodel.triton_kernel.mtp_utils import ( + gen_b_req_mtp_start_loc, mtp_scatter_next_token_ids, mtp_verify, ) @@ -16,28 +18,28 @@ from lightllm.server.router.model_infer.mtp_speculative.proposers.base import SpecProposal -def verify_mtp_tokens( - backend: ModeBackend, - next_token_ids: torch.Tensor, - b_req_idx: torch.Tensor, - b_req_mtp_start_loc: torch.Tensor, - b_mtp_index: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Verify target tokens and update recurrent MTP state when required.""" - # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 - +def verify_mtp_tokens(backend: ModeBackend, next_token_ids: torch.Tensor, model_output, num_reqs: int): + outputs = model_output if isinstance(model_output, tuple) else (model_output,) + commits = tuple(output.attention_commit for output in outputs) + if any(commit is None for commit in commits): + raise ValueError("MTP verification requires an output-bound attention commit") + commit = commits[0] if len(commits) == 1 else RequestGenerationStamp.concatenate(commits) + if commit.role != "main": + raise ValueError("target verification cannot commit a draft model output") + commit.validate_generation(backend.model.req_manager.request_generations) + starts = gen_b_req_mtp_start_loc(b_mtp_index=commit.mtp_indices, num_reqs=num_reqs) accept_lengths, accepted_index = mtp_verify( req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids, - b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_mtp_start_loc=starts, new_next_token_ids=next_token_ids, - b_req_idx=b_req_idx, + b_req_idx=commit.request_indices, ) # 最后一层已经按这个窗口全接受推进了 cursor。这里剪掉没通过验证的尾部。 backend.model.att_state_manager.commit_bounded_att_state( - b_req_idx[b_req_mtp_start_loc], + commit.request_indices[starts.long()], accept_lengths, - role="main", - req_verify_width=_req_verify_width(b_req_mtp_start_loc, b_req_idx.shape[0]), + role=commit.role, + req_verify_width=_req_verify_width(starts, commit.request_indices.shape[0]), ) return accept_lengths, accepted_index @@ -67,7 +69,6 @@ def scatter_mtp_next_tokens( mtp_accept_len: torch.Tensor, # [req_num] ) -> None: """Persist the next MTP proposal and optional scheduling scores by request.""" - schedule_scores = getattr(proposal, "schedule_scores", None) if schedule_scores is not None and 0 in schedule_scores.shape: schedule_scores = None @@ -100,7 +101,6 @@ def record_request_mtp_metrics( verify_run_reqs: List[InferReq], ) -> None: """Accumulate user-visible MTP metrics on each request.""" - if not backend.is_master_in_dp: return diff --git a/lightllm/server/visualserver/__init__.py b/lightllm/server/visualserver/__init__.py index 6bc0923da0..b4e65bf259 100644 --- a/lightllm/server/visualserver/__init__.py +++ b/lightllm/server/visualserver/__init__.py @@ -1,16 +1,44 @@ -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend -from lightllm.common.basemodel.attention_vit.create_utils import get_vit_att_backend_class +from lightllm.common.attention.api import VarlenAttentionOp +from lightllm.common.attention.registry import get_varlen_attention_op_class +from lightllm.common.attention.selection import select_varlen_attention_op +from lightllm.platform import get_device_runtime -VIT_ATTN_BACKEND: BaseVitAttBackend = None +class _LazyVarlenAttentionRunner: + """用 worker 预热时的真实 Q/K/V 选择实现，之后通过固定的 forward 委托执行。 -def set_vit_att_backend(backend_name: str): - global VIT_ATTN_BACKEND - VIT_ATTN_BACKEND = get_vit_att_backend_class(backend_name) + 部分视觉模型固定使用 BF16，与 llm 的 data_type 不同。manager 无法据启动 + 参数推断实际 dtype/head dim，因此选择发生在各设备的首次预热输入到达时。 + """ + + def __init__(self, requested): + self.requested = requested + self._op = None + + def forward(self, inputs, *, out): + if self._op is None: + runtime = get_device_runtime(inputs.query.device) + if runtime.is_capturing(): + raise RuntimeError("warm up vision attention before Graph capture") + name = select_varlen_attention_op( + self.requested, device=inputs.query.device, dtype=inputs.query.dtype, head_dim=inputs.query.shape[-1] + ) + self._op = get_varlen_attention_op_class(name)() + return self._op.forward(inputs, out=out) + + +VIT_ATTENTION_OP: VarlenAttentionOp | _LazyVarlenAttentionRunner | None = None + + +def set_vit_attention_op(provider: str, *, defer_validation=False): + global VIT_ATTENTION_OP + VIT_ATTENTION_OP = ( + _LazyVarlenAttentionRunner(provider) if defer_validation else get_varlen_attention_op_class(provider)() + ) return -def get_vit_attn_backend(): - if VIT_ATTN_BACKEND is None: - raise RuntimeError("VIT_ATTN_BACKEND is not initialized. Call init_vit_att_backend() first.") - return VIT_ATTN_BACKEND._vit_att_fwd +def get_vit_attention_op(): + if VIT_ATTENTION_OP is None: + raise RuntimeError("VIT_ATTENTION_OP is not initialized. Call set_vit_attention_op() first.") + return VIT_ATTENTION_OP diff --git a/lightllm/server/visualserver/manager.py b/lightllm/server/visualserver/manager.py index 873eff8456..d6faad3ca0 100644 --- a/lightllm/server/visualserver/manager.py +++ b/lightllm/server/visualserver/manager.py @@ -1,29 +1,31 @@ -import zmq -import zmq.asyncio import asyncio -import uvloop -import rpyc -import socket -import pickle +import collections import inspect -import setproctitle +import pickle +import socket import threading -import collections from typing import List -from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes + +import rpyc +import setproctitle +import uvloop +import zmq +import zmq.asyncio + from lightllm.server.core.objs import ShmReqManager, StartArgs +from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from .model_infer import start_model_process, VisualModelRpcClient -from lightllm.common.basemodel.attention_vit.create_utils import init_vit_att_backend -from lightllm.utils.log_utils import init_logger +from rpyc.utils.classic import obtain + +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.graceful_utils import graceful_registry +from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import start_parent_check_thread -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.shm_port_args import get_shm_port_args -from rpyc.utils.classic import obtain +from .model_infer import VisualModelRpcClient, start_model_process logger = init_logger(__name__) @@ -62,12 +64,10 @@ def __init__( self.lock = asyncio.Lock() async def wait_to_model_ready(self): - self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(self.vit_dp)] - self.vit_attn_backend = init_vit_att_backend(index=0) + self.vit_attn_backend = self.args.vit_att_backend for dp_rank_id in range(self.vit_dp): for tp_rank_id in range(self.vit_tp): - rpc_model = await start_model_process() self.model_rpcs[dp_rank_id].append(rpc_model) diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 6451e70008..6d1146c0a2 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -1,39 +1,40 @@ import os -import rpyc -import torch -import socket -import torch.multiprocessing as mp import queue +import socket import threading import time +from typing import Deque, Dict, List, Optional, Tuple + +import rpyc +import torch import torch.distributed as dist -from typing import Dict, List, Tuple, Deque, Optional -from transformers.configuration_utils import PretrainedConfig +import torch.multiprocessing as mp from rpyc.utils.classic import obtain -from lightllm.models.qwen_vl.qwen_visual import QWenVisionTransformer -from lightllm.models.llava.llava_visual import LlavaVisionModel +from transformers.configuration_utils import PretrainedConfig + +from lightllm.models.deepseek_v4.deepseek_v4_visual import DeepseekV4VisionModel from lightllm.models.gemma3.gemma3_visual import Gemma3VisionModel from lightllm.models.gemma4.gemma4_visual import Gemma4VisionModel -from lightllm.models.deepseek_v4.deepseek_v4_visual import DeepseekV4VisionModel -from lightllm.models.vit.model import VisionTransformer -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel +from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer +from lightllm.models.llava.llava_visual import LlavaVisionModel +from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel from lightllm.models.qwen2_5_vl.qwen2_5_visual import Qwen2_5_VisionTransformerPretrainedModel +from lightllm.models.qwen2_vl.qwen2_visual import Qwen2VisionTransformerPretrainedModel +from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel +from lightllm.models.qwen_vl.qwen_visual import QWenVisionTransformer from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel -from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_visual import Qwen3OmniMoeVisionTransformerPretrainedModel -from lightllm.models.neo_chat_moe.neo_visual import NeoVisionTransformerPretrainedModel -from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer -from lightllm.utils.infer_utils import set_random_seed +from lightllm.models.vit.model import VisionTransformer +from lightllm.server.embed_cache.afs_utils import SepEmbedHandler +from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.server.visualserver import set_vit_attention_op +from lightllm.server.visualserver.model_infer.vision_peak_vram_hold import VisionPeakVramHolder from lightllm.utils.dist_utils import init_vision_distributed_env from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.shm_port_args import get_shm_port_args -from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient -from lightllm.server.visualserver import set_vit_att_backend -from lightllm.server.embed_cache.afs_utils import SepEmbedHandler +from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger -from lightllm.server.visualserver.model_infer.vision_peak_vram_hold import VisionPeakVramHolder - +from lightllm.utils.shm_port_args import get_shm_port_args logger = init_logger(__name__) @@ -67,7 +68,7 @@ def exposed_init_model(self, kvargs): self.is_visual_only_mode = get_env_start_args().run_mode == "visual_only" self.data_type = kvargs["data_type"] self.vit_attn_backend = kvargs["vit_attn_backend"] - set_vit_att_backend(self.vit_attn_backend) + set_vit_attention_op(self.vit_attn_backend, defer_validation=True) init_vision_distributed_env(kvargs) model_cfg, _ = PretrainedConfig.get_config_dict(weight_dir) diff --git a/lightllm/server/visualserver/visual_only_manager.py b/lightllm/server/visualserver/visual_only_manager.py index 835a25fe3e..ba501eaa7e 100644 --- a/lightllm/server/visualserver/visual_only_manager.py +++ b/lightllm/server/visualserver/visual_only_manager.py @@ -1,35 +1,38 @@ import asyncio -import uvloop -import rpyc import inspect -import setproctitle -import threading -import uuid +import os import pickle -import websockets +import signal import socket import sys -import os -import signal +import threading import time -from lightllm.utils.net_utils import get_hostname_ip -from .objs import VIT_Obj +import uuid from typing import List + +import rpyc +import setproctitle +import uvloop +import websockets + from lightllm.server.core.objs import StartArgs +from lightllm.utils.net_utils import get_hostname_ip + +from .objs import VIT_Obj asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) -from lightllm.server.multimodal_params import MultimodalParams, ImageItem -from .model_infer import start_model_process, VisualModelRpcClient -from lightllm.common.basemodel.attention_vit.create_utils import init_vit_att_backend -from lightllm.utils.log_utils import init_logger +from rpyc.utils.classic import obtain + +from lightllm.server.embed_cache.utils import create_shm, free_shm, get_shm_name_data +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.graceful_utils import graceful_registry +from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import start_parent_check_thread -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.shm_port_args import get_shm_port_args -from rpyc.utils.classic import obtain -from lightllm.server.embed_cache.utils import create_shm, get_shm_name_data, free_shm -from .manager import VisualManager +from .manager import VisualManager +from .model_infer import VisualModelRpcClient, start_model_process logger = init_logger(__name__) @@ -74,7 +77,6 @@ async def register_to_config_server_loop(self, args: StartArgs): try: uri = f"ws://{args.config_server_host}:{ports.config_server_port}/visual_register" async with websockets.connect(uri, max_queue=(2048 * 1024, 2048 * 1023)) as websocket: - sock = websocket.transport.get_extra_info("socket") sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) @@ -94,12 +96,10 @@ async def register_to_config_server_loop(self, args: StartArgs): logger.info("reconnection to config_server") async def wait_to_model_ready(self): - self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(self.vit_dp)] - self.vit_attn_backend = init_vit_att_backend(index=0) + self.vit_attn_backend = self.args.vit_att_backend for dp_rank_id in range(self.vit_dp): for tp_rank_id in range(self.vit_tp): - rpc_model = await start_model_process() self.model_rpcs[dp_rank_id].append(rpc_model) diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index 3458f10131..e31fc10da5 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -1,37 +1,71 @@ -"""Backend validation with subprocess isolation and ground truth checks.""" +"""Isolated backend numerical probes and explicitly scoped provider smoke checks.""" import multiprocessing as mp import os +import time +import traceback +from dataclasses import dataclass +from enum import Enum +from functools import lru_cache + import torch -from lightllm.utils.log_utils import init_logger + from lightllm.utils.dist_utils import get_global_rank -from functools import lru_cache +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) _VALIDATION_TIMEOUT = 30 * 60 # 30 minutes +class ProbeScope(str, Enum): + BACKEND_NUMERICS = "backend numerics" + VARLEN_NUMERICS = "varlen numerics" + PROVIDER_SMOKE = "provider smoke only; selected backend numerics not checked" + INCOMPLETE = "probe did not complete" + + +@dataclass(frozen=True) +class ProbeResult: + ok: bool + scope: ProbeScope + reason: str | None = None + + +def _has_dense_numerical_probe(spec): + from lightllm.common.attention.api import AttentionFamily + from lightllm.common.attention.api.spec import CacheEncoding, AttentionFeature + + return spec is not None and ( + spec.family is AttentionFamily.DENSE + and spec.cache.encoding is CacheEncoding.PLAIN + and not spec.use_mtp + and AttentionFeature.ALIBI not in spec.features + and AttentionFeature.SINK not in spec.features + and spec.sliding_window is None + and not spec.variant + ) + + def _compute_ground_truth(q, k, v, is_causal=True): """Ground truth using PyTorch SDPA.""" with torch.no_grad(): return torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal) -def _validate_fa3(): +def _validate_fa3(dtype=torch.bfloat16, spec=None): """Validate FA3 with ground truth.""" - from lightllm.utils.device_utils import is_hopper from lightllm.utils.sgl_utils import flash_attn_varlen_func - if not is_hopper(): - return False, "Not a Hopper GPU" + if torch.cuda.get_device_capability()[0] < 9: + return False, "FA3 requires SM90 or newer" if flash_attn_varlen_func is None: return False, "flash_attn_varlen_func is None" - batch, heads, seq, dim = 1, 4, 8, 64 - q = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(batch, heads, seq, dim, dtype=torch.bfloat16, device="cuda") + batch, heads, seq, dim = 1, 4 if spec is None else spec.num_q_heads, 8, 64 if spec is None else spec.head_size + q = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + k = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + v = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") expected = _compute_ground_truth(q, k, v) @@ -55,11 +89,14 @@ def _validate_fa3(): torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashinfer(): +def _validate_flashinfer(dtype=torch.float16, spec=None): """Validate FlashInfer with ground truth.""" capability = torch.cuda.get_device_capability() arch = f"{capability[0]}.{capability[1]}" @@ -67,10 +104,10 @@ def _validate_flashinfer(): import flashinfer - batch, heads, seq, dim = 1, 4, 8, 64 - q = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") - k = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") - v = torch.randn(batch, heads, seq, dim, dtype=torch.float16, device="cuda") + batch, heads, seq, dim = 1, 4 if spec is None else spec.num_q_heads, 8, 64 if spec is None else spec.head_size + q = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + k = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") + v = torch.randn(batch, heads, seq, dim, dtype=dtype, device="cuda") expected = _compute_ground_truth(q, k, v) @@ -88,15 +125,18 @@ def _validate_flashinfer(): torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashqla(): +def _validate_flashqla(dtype=None, spec=None): """Validate FlashQLA against LightLLM's vendored FLA kernel.""" from flash_qla import chunk_gated_delta_rule as flashqla_chunk_gated_delta_rule - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) from lightllm.models.qwen3next.config import Qwen3NextResourceConfig @@ -107,7 +147,11 @@ def _validate_flashqla(): num_v_heads = max(1, linear_config.global_linear_v_heads // linear_config.tp_world_size) head_k_dim = linear_config.head_linear_k_dim head_v_dim = linear_config.head_linear_v_dim - qkv_dtype = linear_config.conv_state_dtype + qkv_dtype = dtype or linear_config.conv_state_dtype + if spec is not None: + num_k_heads = spec.num_q_heads + num_v_heads = max(num_k_heads, num_v_heads) + head_k_dim, head_v_dim = spec.head_size, spec.value_head_size state_dtype = linear_config.ssm_state_dtype batch, seq = 1, 64 @@ -140,12 +184,15 @@ def _validate_flashqla(): ("final state", final_state, expected_state), ): if not torch.allclose(actual, expected, rtol=1e-2, atol=1e-2): - return False, f"{name} mismatch: max diff {(actual - expected).abs().max().item():.6f}" + return ( + False, + f"{name} mismatch: max diff {(actual - expected).abs().max().item():.6f}", + ) return True, None -def _validate_triton(): - """Validate Triton with softmax ground truth.""" +def _smoke_triton(): + """Check Triton compilation and a softmax kernel, not the selected attention backend.""" import triton import triton.language as tl @@ -169,7 +216,10 @@ def _softmax_kernel(input_ptr, output_ptr, n_cols, BLOCK: tl.constexpr): torch.cuda.synchronize() if not torch.allclose(out, expected, rtol=1e-3, atol=1e-3): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None @@ -208,7 +258,10 @@ def _validate_xformers(): out = out.transpose(1, 2).contiguous() if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None @@ -234,13 +287,16 @@ def _validate_sdpa(): out = scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) if not torch.allclose(out, expected, rtol=1e-2, atol=1e-2): - return False, f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return ( + False, + f"Output mismatch: max diff {(out - expected).abs().max().item():.6f}", + ) return True, None -def _validate_flashmla_sparse(): - """Validate flashmla_sparse (NSA) with ground truth.""" +def _smoke_flashmla_sparse(dtype=torch.bfloat16, spec=None): + """Check sparse MLA import and kernel launch; no reference output comparison.""" try: from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -252,7 +308,6 @@ def _validate_flashmla_sparse(): kv_lora_rank = 512 qk_rope_head_dim = 64 qk_dim = kv_lora_rank + qk_rope_head_dim - dtype = torch.bfloat16 device = "cuda" q = torch.randn(batch * seq, heads, qk_dim, dtype=dtype, device=device) @@ -282,79 +337,152 @@ def _validate_flashmla_sparse(): return True, None -def _run_in_subprocess(backend_name, pipe): - """Run validation in subprocess with suppressed output.""" - import sys +def _validate_varlen(backend_name, dtype=None, device=None, head_dim=80): + from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionInputs + from lightllm.common.attention.registry import get_varlen_attention_op_class + from lightllm.platform import get_device_runtime + + runtime = get_device_runtime(device or ("cuda" if torch.cuda.is_available() else "cpu")) + op_cls = get_varlen_attention_op_class(backend_name) + try: + op_cls.check_environment(runtime.caps) + except UnsupportedAttention as exc: + return False, str(exc) + device = runtime.device + dtype = dtype or (torch.float32 if device.type == "cpu" else torch.bfloat16) + q, k, v = [torch.randn(8, 2, head_dim, dtype=dtype, device=device) for _ in range(3)] + cu = torch.tensor([0, 3, 8], dtype=torch.int32, device=device) + out = torch.empty_like(q) + inputs = VarlenAttentionInputs(q, k, v, cu, 5) + try: + op_cls.check_support(inputs, runtime.caps) + except UnsupportedAttention as exc: + return False, str(exc) + op_cls().forward(inputs, out=out) + expected = torch.empty_like(q) + for start, end in ((0, 3), (3, 8)): + result = _compute_ground_truth( + q[start:end].transpose(0, 1), k[start:end].transpose(0, 1), v[start:end].transpose(0, 1), is_causal=False + ) + expected[start:end].copy_(result.transpose(0, 1)) + if not torch.allclose(out, expected, rtol=2e-2, atol=2e-2): + return False, f"varlen output mismatch: max diff {(out - expected).abs().max().item():.6f}" + return True, None + - # Redirect stdout/stderr to /dev/null to suppress verbose error messages - devnull = open(os.devnull, "w") - sys.stdout = devnull - sys.stderr = devnull +def _probe_in_process(backend_name, *, varlen=False, device=None, dtype=None, spec=None, head_dim=80): + from lightllm.common.attention.registry import load_backend_class + from lightllm.platform import get_device_runtime + scope = ProbeScope.VARLEN_NUMERICS if varlen else ProbeScope.PROVIDER_SMOKE + if _has_dense_numerical_probe(spec) and not varlen: + scope = ProbeScope.BACKEND_NUMERICS try: - if backend_name == "fa3": - success, err = _validate_fa3() - elif backend_name == "xformers": - success, err = _validate_xformers() - elif backend_name == "sdpa": - success, err = _validate_sdpa() - elif backend_name == "flashinfer": - success, err = _validate_flashinfer() - elif backend_name == "flashqla": - success, err = _validate_flashqla() - elif backend_name == "triton": - success, err = _validate_triton() - elif backend_name == "flashmla_sparse": - success, err = _validate_flashmla_sparse() + if device is not None or spec is not None: + runtime = get_device_runtime(device) + if runtime.api is not None: + runtime.api.set_device(runtime.device) + if scope is ProbeScope.BACKEND_NUMERICS: + from lightllm.common.attention.probe import probe_dense + + ok, reason = probe_dense(load_backend_class(backend_name), spec, runtime) + elif varlen: + ok, reason = _validate_varlen(backend_name, dtype, device, head_dim) else: - success, err = False, f"Unknown backend: {backend_name}" - pipe.send((success, err)) - except Exception as e: - pipe.send((False, str(e))) - finally: - devnull.close() + provider = load_backend_class(backend_name).provider if spec is not None else backend_name + ok, reason = _smoke_provider(provider, dtype=dtype, spec=spec) + return ProbeResult(ok, scope, reason) + except Exception: + return ProbeResult(False, scope, traceback.format_exc()) + + +def _smoke_provider(provider, *, dtype=None, spec=None): + """Exercise a provider's representative kernel, not the selected family/spec implementation.""" + if provider == "fa3": + return _validate_fa3(dtype or torch.bfloat16, spec) + if provider == "xformers": + return _validate_xformers() + if provider == "sdpa": + return _validate_sdpa() + if provider == "flashinfer": + return _validate_flashinfer(dtype or torch.float16, spec) + if provider == "flashqla": + return _validate_flashqla(dtype, spec) + if provider == "triton": + return _smoke_triton() + if provider == "flashmla_sparse": + return _smoke_flashmla_sparse(dtype or torch.bfloat16, spec) + return False, f"Unknown provider: {provider}" + + +def _run_in_subprocess(backend_name, pipe, varlen=False, device=None, dtype=None, spec=None, head_dim=80): + from contextlib import redirect_stderr, redirect_stdout + + with open(os.devnull, "w") as devnull, redirect_stdout(devnull), redirect_stderr(devnull): + pipe.send( + _probe_in_process(backend_name, varlen=varlen, device=device, dtype=dtype, spec=spec, head_dim=head_dim) + ) @lru_cache(maxsize=None) -def validate(backend_name: str) -> bool: - if get_global_rank() == 0: - validate_ok = _validate(backend_name) - torch.distributed.broadcast_object_list([validate_ok], src=0) - else: - validate_ok = [None] - torch.distributed.broadcast_object_list(validate_ok, src=0) - validate_ok = validate_ok[0] - return validate_ok - - -def _validate(backend_name: str) -> bool: - """Validate backend in subprocess with ground truth check.""" +def probe_backend(backend_name: str, *, device=None, dtype=None, spec=None) -> ProbeResult: + """Cache by provider/device/dtype/full spec; the result states the actual check performed.""" + return _probe(backend_name, device=device, dtype=dtype, spec=spec) + + +@lru_cache(maxsize=None) +def probe_varlen(backend_name: str, *, device=None, dtype=None, head_dim=80) -> ProbeResult: + """Compare the selected Q/K/V operator against a sequence-wise reference.""" + return _probe(backend_name, varlen=True, device=device, dtype=dtype, head_dim=head_dim) + + +def _probe(backend_name: str, *, varlen=False, device=None, dtype=None, spec=None, head_dim=80) -> ProbeResult: + """Isolate execution and retain both its coverage and failure reason.""" try: ctx = mp.get_context("spawn") parent, child = ctx.Pipe(duplex=False) - logger.info(f"Validating {backend_name} backend start, please wait ...") - proc = ctx.Process(target=_run_in_subprocess, args=(backend_name, child)) + logger.info("Probing %s on device=%s dtype=%s", backend_name, device, dtype) + proc = ctx.Process(target=_run_in_subprocess, args=(backend_name, child, varlen, device, dtype, spec, head_dim)) proc.start() - proc.join(timeout=_VALIDATION_TIMEOUT) + child.close() + deadline = time.monotonic() + _VALIDATION_TIMEOUT + # Drain the pipe before joining: a full traceback can exceed its buffer. + result = None + if parent.poll(_VALIDATION_TIMEOUT): + try: + result = parent.recv() + except EOFError: + pass + proc.join(timeout=max(0, deadline - time.monotonic())) if proc.is_alive(): proc.kill() proc.join() - logger.warning(f"{backend_name} validation timed out") - return False + return ProbeResult(False, ProbeScope.INCOMPLETE, "probe timed out") if proc.exitcode != 0: - logger.warning(f"{backend_name} validation crashed (exit {proc.exitcode})") - return False - - if parent.poll(): - success, err = parent.recv() - if success: - logger.info(f"{backend_name} validated") - return True - logger.warning(f"{backend_name} validation failed: {str(err)[0:88]}") - return False - - except Exception as e: - logger.warning(f"{backend_name} validation error: {e}") - return False + return ProbeResult(False, ProbeScope.INCOMPLETE, f"probe crashed (exit {proc.exitcode})") + + if result is not None: + if result.ok: + logger.info("%s passed: %s", backend_name, result.scope.value) + else: + logger.warning( + "%s failed (%s) on device=%s dtype=%s spec=%r:\n%s", + backend_name, + result.scope.value, + device, + dtype, + spec, + result.reason, + ) + return result + return ProbeResult(False, ProbeScope.INCOMPLETE, "probe returned no result") + except Exception: + logger.exception("%s probe error on device=%s dtype=%s spec=%r", backend_name, device, dtype, spec) + return ProbeResult(False, ProbeScope.INCOMPLETE, traceback.format_exc()) + finally: + if "parent" in locals(): + parent.close() + if "child" in locals(): + child.close() diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 9d8490b69e..dfa517089e 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -1,11 +1,12 @@ -import os import json -import torch +import os import uuid -from easydict import EasyDict from functools import lru_cache -from lightllm.utils.log_utils import init_logger +import torch +from easydict import EasyDict + +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -221,7 +222,6 @@ def get_decode_attn_autotune_seq_len() -> int: def get_model_init_status(): - global g_model_init_done return g_model_init_done @@ -267,7 +267,7 @@ def enable_triton_mtp_kernel() -> bool: 启用 Triton MTP 解码专用 kernel 通过启动参数 --mtp_step > 0 和 --llm_decode_att_backend=triton 控制 """ - return (get_env_start_args().mtp_step > 0) and ("triton" in get_env_start_args().llm_decode_att_backend) + return (get_env_start_args().mtp_step > 0) and (get_env_start_args().llm_decode_att_backend == "triton") @lru_cache(maxsize=None) diff --git a/lightllm/utils/start_utils.py b/lightllm/utils/start_utils.py index 28a0e5566e..8456106ec9 100644 --- a/lightllm/utils/start_utils.py +++ b/lightllm/utils/start_utils.py @@ -1,15 +1,17 @@ -import os import ctypes +import multiprocessing as mp +import os import resource import signal import subprocess import sys import time -import multiprocessing as mp + import psutil + +from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import is_process_active -from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.service_shm_cleanup import start_launcher_shm_cleanup_process logger = init_logger(__name__) diff --git a/lightllm/utils/torch_memory_saver_utils.py b/lightllm/utils/torch_memory_saver_utils.py index e0e6cad94a..9468d171d2 100644 --- a/lightllm/utils/torch_memory_saver_utils.py +++ b/lightllm/utils/torch_memory_saver_utils.py @@ -1,13 +1,12 @@ -import torch from contextlib import contextmanager from enum import Enum + +import torch + from lightllm.utils.log_utils import init_logger try: - from torch_memory_saver import ( - torch_memory_saver, - configure_subprocess as tms_configure_subprocess, - ) + from torch_memory_saver import configure_subprocess as tms_configure_subprocess, torch_memory_saver HAS_TORCH_MEMORY_SAVER = True @@ -52,6 +51,8 @@ def __new__(cls, enable_torch_memory_saver: bool = False): class _TorchMemorySaver: + enabled = True + @contextmanager def configure_subprocess(self): # 子进程启动需要放在该上下文里，让 torch_memory_saver 在 worker @@ -80,6 +81,8 @@ def resume(self, tag: MemoryTag): class _TorchMemorySaverFake: + enabled = False + # 未开启 torch_memory_saver 时保持相同接口，保证调用方逻辑完全一致。 @contextmanager def configure_subprocess(self): diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 569d4f99a5..b25351ad41 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -14,11 +14,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING import json -from pathlib import Path import socket import sys +from pathlib import Path +from typing import TYPE_CHECKING import numpy as np import torch @@ -35,6 +35,7 @@ if TYPE_CHECKING: from argparse import Namespace from collections.abc import Sequence + from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -309,8 +310,10 @@ def main(): linear_att_hash_page_size=64, linear_att_page_block_num=1, linear_att_cache_size=16, - llm_prefill_att_backend=["fa3", "triton"], - llm_decode_att_backend=["fa3", "triton"], + llm_prefill_att_backend="fa3", + llm_prefill_linear_att_backend="triton", + llm_decode_att_backend="fa3", + llm_decode_linear_att_backend="triton", disable_cudagraph=False, pd_kv_page_size=2048, ) @@ -332,10 +335,10 @@ def main(): config, _ = PretrainedConfig.get_config_dict(args.model_dir) model, _ = get_model(config, kvargs) draft_models = init_mtp_draft_models(args, kvargs, model) - from lightllm.utils.envs_utils import set_model_init_status - from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.attention_state import AttentionStateManager from lightllm.common.req_manager import ReqManager + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag models = (model, *draft_models) assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" @@ -352,6 +355,9 @@ def main(): assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + item.attention.reserve(item) manager.load_full_att_calibration(args=model.args) manager.initialize_storage( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py index c60038af40..bd92c66682 100644 --- a/test/benchmark/kernels/benchmark_kda_ssm.py +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -13,8 +13,8 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda -from lightllm.common.basemodel.attention.linear.replay_kda_decode import get_configs, replay +from lightllm.common.attention.backends.kda.replay.decode import get_configs, replay +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py index a1cee4826a..4832abfa61 100644 --- a/test/benchmark/kernels/benchmark_replayssm_service.py +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -7,8 +7,8 @@ import concurrent.futures import json import statistics -import time import threading +import time from pathlib import Path import requests diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 3880029a54..07547e5b8b 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -22,7 +22,6 @@ import torch.multiprocessing as mp from transformers import PretrainedConfig - REPO_ROOT = Path(__file__).resolve().parents[3] if str(REPO_ROOT) not in sys.path: sys.path.append(str(REPO_ROOT)) @@ -34,8 +33,8 @@ from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager from lightllm.utils.config_utils import auto_set_fused_shared_experts, get_dtype, get_vocab_size from lightllm.utils.dist_utils import init_distributed_env -from lightllm.utils.envs_utils import set_env_start_args - +from lightllm.utils.envs_utils import get_unique_server_name, set_env_start_args, set_unique_server_name +from lightllm.utils.service_shm_cleanup import ServiceShmCleanup DEFAULT_BATCH_SIZES = [2, 8, 16, 32, 64, 128] MTP_MODES = {"vanilla_with_att", "eagle_with_att", "vanilla_no_att", "eagle_no_att"} @@ -596,7 +595,6 @@ def _forward_decode_input(self, model_input: ModelInput, allow_overlap: bool) -> def _split_prefill_input(self, model_input: ModelInput): """在 benchmark 调用侧按请求构造两个 prefill microbatch。""" - split_batch = model_input.batch_size // 2 q_lens = model_input.b_seq_len - model_input.b_ready_cache_len split_tokens = int(q_lens[:split_batch].sum().item()) @@ -1171,6 +1169,8 @@ def build_model_kvargs(args: SimpleNamespace, rank_id: int) -> Dict: "disable_cudagraph": args.disable_cudagraph, "llm_prefill_att_backend": args.llm_prefill_att_backend, "llm_decode_att_backend": args.llm_decode_att_backend, + "llm_prefill_linear_att_backend": args.llm_prefill_linear_att_backend, + "llm_decode_linear_att_backend": args.llm_decode_linear_att_backend, "vit_att_backend": args.vit_att_backend, "llm_kv_type": args.llm_kv_type, "llm_kv_quant_group_size": args.llm_kv_quant_group_size, @@ -1233,8 +1233,10 @@ def init_deferred_cudagraph(args: SimpleNamespace, cases: Sequence[BenchmarkCase model.disable_cudagraph = False # Attention backends may size persistent graph buffers during construction. # Rebuild them after replacing the temporary profile-time batch limit. - model._init_att_backend() - model._init_att_backend1() + model.prepare_attention_backends() + + model.attention.reserve(model) + model.attention.bind(model) torch.cuda.empty_cache() model._init_cudagraph() @@ -1245,9 +1247,10 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) cases = [BenchmarkCase(**case) for case in case_dicts] set_env_start_args(args) - from lightllm.distributed import dist_group_manager import torch.distributed as dist + from lightllm.distributed import dist_group_manager + model_kvargs = build_model_kvargs(args, rank_id) defer_cudagraph = ( not args.disable_cudagraph @@ -1270,10 +1273,10 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) torch.cuda.empty_cache() model, _ = get_model(model_cfg, model_kvargs) draft_models = init_mtp_draft_models(args, model_kvargs, model) - from lightllm.utils.envs_utils import set_model_init_status - from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.common.attention_state import AttentionStateManager from lightllm.common.req_manager import ReqManager + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag models = (model, *draft_models) assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" @@ -1290,6 +1293,9 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + item.attention.reserve(item) manager.load_full_att_calibration(args=model.args) manager.initialize_storage( model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction @@ -1309,8 +1315,9 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) batch_size *= draft.mtp_manager.get_decode_tokens_per_request(True) draft.graph_max_batch_size = draft._align_decode_batch_size(batch_size) draft.disable_cudagraph = False - draft._init_att_backend() - draft._init_att_backend1() + draft.prepare_attention_backends() + draft.attention.reserve(draft) + draft.attention.bind(draft) draft._init_cudagraph() token_source = TokenSource(args) executor = StaticBenchmarkExecutor(args, model, draft_models, token_source) @@ -1642,9 +1649,13 @@ def main(argv: Optional[Sequence[str]] = None): if not cases: raise ValueError("no benchmark cases were generated") args = normalize_args(args, cases) + set_unique_server_name(args) set_env_start_args(args) - run_benchmark(args, cases) + try: + run_benchmark(args, cases) + finally: + ServiceShmCleanup(get_unique_server_name()).cleanup_service_resources() if __name__ == "__main__": diff --git a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py index f32b093448..76a9292e52 100644 --- a/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py +++ b/test/kernel/llama_gqa_diverse_decode_stage1_tuning.py @@ -1,22 +1,25 @@ -import torch -import time import os -import torch.multiprocessing as mp +import time from typing import List -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage1, + +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( GQADiverseDecodeStage1KernelConfig, + flash_decode_stage1, ) +from lightllm.utils.log_utils import init_logger from lightllm.utils.watchdog_utils import Watchdog logger = init_logger(__name__) def set_seed(): - import torch import random + import numpy as np + import torch seed = 42 torch.manual_seed(seed) @@ -301,9 +304,10 @@ def tuning_configs( if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") - from lightllm.utils.tuning_utils import mp_tuning import collections + from lightllm.utils.tuning_utils import mp_tuning + block_seq = 256 store_json_ans = collections.defaultdict(dict) diff --git a/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py index 13c8945e59..0492884589 100644 --- a/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py +++ b/test/kernel/llama_gqa_diverse_decode_stage2_tuning.py @@ -1,21 +1,24 @@ -import torch import os -import torch.multiprocessing as mp from typing import List -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage2, + +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( GQADiverseDecodeStage2KernelConfig, + flash_decode_stage2, ) +from lightllm.utils.log_utils import init_logger from lightllm.utils.watchdog_utils import Watchdog logger = init_logger(__name__) def set_seed(): - import torch import random + import numpy as np + import torch seed = 42 torch.manual_seed(seed) @@ -260,9 +263,10 @@ def tuning_configs( if __name__ == "__main__": torch.multiprocessing.set_start_method("spawn") - from lightllm.utils.tuning_utils import mp_tuning import collections + from lightllm.utils.tuning_utils import mp_tuning + block_seq = 256 batch_sizes = [8, 16, 32, 64] seq_lens = [32, 64, 128, 256] diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py index cfcf1ad8bf..2caf51b992 100644 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -2,17 +2,17 @@ from __future__ import annotations -from typing import TYPE_CHECKING from dataclasses import replace +from typing import TYPE_CHECKING + +import torch + from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.attention_state.full_att import MLAConfig, MLAFullAttentionState +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator from lightllm.common.attention_state.bounded_att import GDNState -from unit_tests.common.att_resources.test_full_state import full_state +from lightllm.common.attention_state.full_att import MLAConfig, MLAFullAttentionState from unit_tests.common.att_resources.test_bounded_state import make_state - -if TYPE_CHECKING: - import torch +from unit_tests.common.att_resources.test_full_state import full_state class MainDraftTestResourceManager(AttentionStateManager): @@ -25,7 +25,9 @@ def __init__(self, device: str | torch.device | int = "cpu"): mtp_step=self.main_kda_conv_fp32_ssm_fp32_state.mtp_step, device=device, ) - draft_state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + draft_state.initialize_page_pools( + big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu" and torch.cuda.is_available() + ) self.register_state("draft0", draft_state) self.register_state("main", full_state(device=device)) prepared_full_state = MLAFullAttentionState( diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 86e62cb671..029c25392a 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -2,37 +2,38 @@ from __future__ import annotations -from typing import TYPE_CHECKING from contextlib import nullcontext from dataclasses import replace from types import SimpleNamespace +from typing import TYPE_CHECKING import torch -from lightllm.common.attention_state.types import TensorField + from lightllm.common.attention_state.bounded_att import GDNState, KDAState, WindowCompressorState from lightllm.common.attention_state.full_att import ( - FullAttentionState, - GQAAttentionConfig, - MLAConfig, - GroupQuantizedAttentionConfig, - StaticFP8AttentionConfig, - PackedMLAAttentionConfig, CompressedAttentionConfig, FP8E8M0PackedMLAConfig, - StaticFP8FullAttentionState, + FullAttentionState, + GQAAttentionConfig, GQAFullAttentionState, + GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, + MLAConfig, MLAFullAttentionState, + PackedMLAAttentionConfig, PackedMLAFullAttentionState, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, ) - +from lightllm.common.attention_state.types import TensorField if TYPE_CHECKING: - from lightllm.models.qwen3next.config import Qwen3NextResourceConfig - from lightllm.common.attention_state.allocators import FullAttKVAllocator from collections.abc import Mapping, Sequence + + from lightllm.common.attention_state.allocators import FullAttKVAllocator from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.models.qwen3next.config import Qwen3NextResourceConfig def bind_test_fp8_calibration( @@ -259,8 +260,13 @@ def build_hybrid_resources( """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.allocators import StatePageAllocator - from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state - from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, + WindowCompressorStateConfig, + create_gdn_state, + create_kda_state, + create_window_compressor_state, + ) from lightllm.models.glm5_next.config import Glm5NextResourceConfig manager = AttentionStateManager(request_capacity=request_capacity) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index 457fdfa1c3..8bbe1dbaa5 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -1,16 +1,13 @@ -"""Exercise shared-manager startup with independently prepared hybrid drafts.""" - from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers +"""Exercise shared-manager startup with independently prepared hybrid drafts.""" -from unit_tests.common.att_resources.model_helpers import read_kv_storage -from typing import TYPE_CHECKING +import json from contextlib import nullcontext from dataclasses import asdict -import json from types import SimpleNamespace +from typing import TYPE_CHECKING import pytest import torch @@ -22,11 +19,13 @@ from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, read_kv_storage if TYPE_CHECKING: - from lightllm.common.basemodel.basemodel import TpPartBaseModel from pathlib import Path + from lightllm.common.basemodel.basemodel import TpPartBaseModel + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -162,13 +161,31 @@ def mark_backends(model: TpPartBaseModel): ) is state ) - model.prefill_att_backend = model.decode_att_backend = object() - model.prefill_att_backend1 = model.decode_att_backend1 = None + from lightllm.common.attention.api import AttentionFamily, AttentionPhase + + backend = SimpleNamespace(spec=SimpleNamespace(family=AttentionFamily.DENSE)) + model.attention = SimpleNamespace( + groups=({phase: backend for phase in AttentionPhase},), + reserve=reserve_resources, + bind=initialize_layer_attention, + ) model.backends_bound = True + def initialize_layer_attention(model): + assert model.att_state_manager.full_att_allocator is not None + assert all(item.backends_bound for item in models) + model.layer_attention_initialized = True + + def reserve_resources(model): + assert all(item.backends_bound for item in models) + assert model.att_state_manager.full_att_allocator is None + model.attention_resources_reserved = True + def warmup(model: TpPartBaseModel): + assert model.layer_attention_initialized assert model.backends_bound assert all(item.req_manager is mapping for item in models) + assert all(item.attention_resources_reserved for item in models) assert all(item.att_state_manager is models[0].att_state_manager for item in models) assert model.att_state_manager.full_att_allocator.size == 32 assert all(item.max_total_token_num == 32 for item in models[: models.index(model) + 1]) @@ -202,8 +219,7 @@ def check_max_len(model: TpPartBaseModel): model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) if index: model.main_model, model.mtp_previous_draft_models = models[0], models[1:].copy() - model._init_att_backend = lambda model=model: mark_backends(model) - model._init_att_backend1 = lambda: None + model.prepare_attention_backends = lambda model=model: mark_backends(model) model._init_hidden_collector = lambda: None model._autotune_warmup = lambda model=model: warmup(model) model._init_padded_req = lambda model=model: padded.add(id(model)) @@ -214,6 +230,7 @@ def check_max_len(model: TpPartBaseModel): def memory_info(device: str | torch.device | int): assert all(item.req_manager is mapping for item in models) + assert all(item.attention_resources_reserved for item in models) manager = models[0].att_state_manager states = manager.get_full_attention_states() assert all(read_kv_storage(state) is None for state in states) @@ -273,6 +290,10 @@ def start_runtime(): assert all(not hasattr(other, "att_state_manager") for other in models[index + 1 :]) main = models[0] main.att_state_manager.finish_state_registration() + for item in models: + item.prepare_attention_backends() + for item in models: + reserve_resources(item) main.att_state_manager.load_full_att_calibration(args=args) main.att_state_manager.initialize_storage( args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction @@ -280,7 +301,7 @@ def start_runtime(): assert all(item.max_total_token_num is None for item in models) for index, item in enumerate(models): item.init_runtime() - assert all(not getattr(other, "backends_bound", False) for other in models[index + 1 :]) + assert all(not getattr(other, "layer_attention_initialized", False) for other in models[index + 1 :]) assert warmed == padded == captured == prefill_captured == {id(other) for other in models[: index + 1]} start_runtime() diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index e2c08b1461..781a4fadeb 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -1,24 +1,20 @@ from __future__ import annotations -from typing import TYPE_CHECKING from dataclasses import replace +from typing import TYPE_CHECKING import pytest import torch from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, GDNPagesPool, GDNState, KDAPagesPool, KDAState, - ConvSSMStateConfig, -) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, ) +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout, TransferRegion if TYPE_CHECKING: from lightllm.common.attention_state.types import PageKind @@ -51,7 +47,9 @@ def make_state( ) manager = AttentionStateManager(request_capacity=3) manager.register_state("main", state) - state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + state.initialize_page_pools( + big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu" and torch.cuda.is_available() + ) return state diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py index b6dff3dc0e..ad1f79b31d 100644 --- a/unit_tests/common/att_resources/test_capacity.py +++ b/unit_tests/common/att_resources/test_capacity.py @@ -1,41 +1,43 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration - -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - -from unit_tests.common.att_resources.model_helpers import rebuild_full_state - -from typing import TYPE_CHECKING from dataclasses import replace from functools import lru_cache, partial from math import prod +from typing import TYPE_CHECKING import pytest import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.utils import deepcopy_result -from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState -from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig -from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator +from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, + GDNState, + KDAState, + WindowCompressorState, + WindowCompressorStateConfig, +) from lightllm.common.attention_state.full_att import ( FullAttentionState, - GQAFullAttentionState, GQAAttentionConfig, - MLAConfig, - create_gqa_state, - create_mla_state, + GQAFullAttentionState, GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, + MLAConfig, MLAFullAttentionState, PackedMLAFullAttentionState, StaticFP8FullAttentionState, + create_gqa_state, + create_mla_state, +) +from lightllm.common.attention_state.utils import deepcopy_result +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.model_helpers import ( + bind_test_fp8_calibration, + full_state_test_buffers, + full_state_test_fields, + read_kv_storage, + rebuild_full_state, ) if TYPE_CHECKING: @@ -146,12 +148,18 @@ def test_budget_matches_token_stores_including_hold_and_padding( reserve = manager.get_pd_operator().get_page_bytes() * args.pd_kv_page_num monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 0) monkeypatch.setattr( "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) for budget, expected in ((actual, 20), (actual - 1, 16), (minimum, 0), (minimum - 1, None)): resources = AttentionStateManager(request_capacity=1) + from types import SimpleNamespace + + resources.device_runtime = SimpleNamespace( + device=torch.device("cpu"), + empty_cache=lambda: None, + memory_info=lambda: (budget + reserve, budget + reserve), + ) resources.register_state("main", rebuild_full_state(template, config=template.config)) monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (budget + reserve, budget + reserve)) resources.finish_state_registration() diff --git a/unit_tests/common/att_resources/test_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py index badc2adf0c..f58674da86 100644 --- a/unit_tests/common/att_resources/test_dsv4_integration.py +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -1,3 +1,5 @@ +from lightllm.models.deepseek_v4.workspace import get_swa_prefill_workspace + """通过模型工厂和生产 fused kernel 验证 V4 State 接入。""" import json @@ -8,10 +10,12 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState -from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel from lightllm.utils.envs_utils import get_env_start_args +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + @pytest.fixture def cache(monkeypatch): @@ -76,10 +80,15 @@ class TestModel(SimpleNamespace): att_state_role="main", args=SimpleNamespace(run_mode="normal", enable_prefill_microbatch_overlap=True), ) - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4SparseMLAFp8DecodeBackend + from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - backend = DeepseekV4FlashMlaFp8SparseAttBackend(model=model) - yield manager, window, backend + workspace = DeepseekV4Workspace(model) + workspace.c128_compress_scratch = torch.empty((2, 2048, 512), dtype=torch.float32, device="cuda") + manager._test_workspace = workspace + model.att_state_manager = manager + model.swa_prefill_workspace = get_swa_prefill_workspace(model) + yield manager, window, model torch.cuda.synchronize() get_env_start_args.cache_clear() @@ -87,6 +96,7 @@ class TestModel(SimpleNamespace): def make_infer(manager, window): infer = DeepseekV4InferStateInfo() infer.att_state_manager = manager + infer.dsv4_workspace = manager._test_workspace infer.att_state_role = "main" infer.first_logical_layer_id = window.config.logical_layer_ids[0] return infer @@ -101,12 +111,21 @@ def read_packed_rows(buffer, slots, page_tokens): def test_backend_reads_owned_state_views(cache): - manager, window, backend = cache + manager, window, model = cache infer = make_infer(manager, window) assert infer.swa_state is window infer.is_prefill = False for layer in range(3): - swa, compressed = backend.get_att_input_params(infer, layer) + from lightllm.common.attention.api import AttentionFamily + from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + + adapter = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + adapter.layer_num_, adapter.compress_ratio = layer, model.config["compress_ratios"][layer] + adapter.v_head_dim, adapter.softmax_scale = 512, 512 ** -0.5 + weight = SimpleNamespace(attn_sink_=SimpleNamespace(weight=torch.zeros(64, device="cuda"))) + params = adapter.bind_attention_resources(model, weight, adapter.attention_adapter(AttentionFamily.SPARSE_MLA)) + swa, compressed = params.swa, params.compressed + assert params.prefill_swa.data_ptr() == model.swa_prefill_workspace[0][layer].data_ptr() assert swa.shape[1] == 74880 assert (compressed is None) == (layer == 2) if layer != 2: @@ -180,16 +199,25 @@ def compress(req_idx, start, end, decode=False): def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer - from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa - manager, window, backend = cache + manager, window, model = cache reference = torch.zeros((17, 74880), dtype=torch.uint8, device="cuda") kv = torch.randn((2048, 512), dtype=torch.bfloat16, device="cuda") weight, angles = torch.randn(512, dtype=torch.bfloat16, device="cuda"), torch.randn((2048, 32), device="cuda") freqs = torch.complex(angles.cos(), angles.sin()) layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) layer.layer_num_, layer.eps_, layer.freqs_cis = 2, 1e-6, freqs + layer.compress_ratio, layer.v_head_dim, layer.softmax_scale = 0, 512, 512 ** -0.5 + from lightllm.common.attention.api import AttentionFamily + + params = layer.bind_attention_resources( + model, + SimpleNamespace(attn_sink_=SimpleNamespace(weight=torch.zeros(64, device="cuda"))), + layer.attention_adapter(AttentionFamily.SPARSE_MLA), + ) + layer.prefill_attention = SimpleNamespace(resources=params) for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): infer = make_infer(manager, window) positions = torch.arange(start, end, device="cuda", dtype=torch.int32) @@ -198,9 +226,11 @@ def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): infer.b_ready_cache_len, infer.b_seq_len = ( torch.tensor([value], dtype=torch.int32, device="cuda") for value in (start, end) ) - prefill = backend.create_att_prefill_state(infer) - prefill.prepare_swa_workspace(infer) - infer.prefill_att_state = prefill + from lightllm.models.deepseek_v4.attention.sparse_mla import prepare_swa_workspace + + buffer, capacity = model.swa_prefill_workspace + prefill = prepare_swa_workspace(window, buffer, capacity, infer, 0) + infer.swa_workspace = prefill indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") lengths, writes, runtime = (torch.empty_like(positions) for _ in range(3)) build_swa_indices( @@ -230,12 +260,13 @@ def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): @pytest.mark.parametrize("draft_kind", ["mtp", "dspark"]) def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): from contextlib import nullcontext - from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel + from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel + from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel - _, _, backend = cache + _, _, model = cache config = dict( - backend.model.config, + model.config, compress_ratios=[4, 128, 0, 0, 0], dspark_target_layer_ids=[0, 1], dspark_block_size=4, @@ -244,6 +275,7 @@ def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): ) args = SimpleNamespace( llm_kv_type="None", + hardware_platform="cuda", page_size=256, mtp_step=3, mtp_mode="vanilla_with_att", diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py index ec5a247576..8ef2c41b10 100644 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -1,145 +1,7 @@ -"""FP8 backends obtain calibration from the role's matching storage scheme.""" - -from __future__ import annotations - -from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration - -from unit_tests.common.att_resources.model_helpers import rebuild_full_state -from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups - -from typing import TYPE_CHECKING -from dataclasses import replace -from types import SimpleNamespace -import weakref - import pytest import torch -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.allocators import FullAttKVAllocator -from lightllm.common.basemodel.attention.base_att import AttControl, BaseAttBackend -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3 -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer -from lightllm.common.basemodel.infer_struct import InferStateInfo -from unit_tests.common.att_resources.test_capacity import definition - -if TYPE_CHECKING: - from collections.abc import Sequence - - -@pytest.mark.parametrize("kind", ["static_head", "static_tensor"]) -@pytest.mark.parametrize("role", ["main", "draft0"]) -@pytest.mark.parametrize("prefill", [False, True]) -def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( - monkeypatch: pytest.MonkeyPatch, - kind: str, - role: str, - prefill: bool, -): - template = definition(kind) - states = [] - for value in (2, 5): - scales = template.scales + value - scales[1] += 10 - q_scales = None if kind == "static_tensor" else template.q_scales + value - states.append( - rebuild_full_state( - template=template, - config=replace_logical_layer_groups(template.config, ((3,), (8,))), - token_capacity=8, - device="cpu", - scales=scales, - q_scales=q_scales, - ) - ) - manager = AttentionStateManager(request_capacity=2) - manager.register_state("main", states[0]) - manager.register_state("draft0", states[1]) - manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - # Another encoding in the same role must not supply the FP8 calibration. - plain = definition("gqa") - plain = rebuild_full_state( - template=plain, - config=plain.config, - token_capacity=8, - device="cpu", - ) - manager.register_state(role, plain) - other = rebuild_full_state( - template=template, - config=replace(template.config, logical_layer_groups=((20,), (21,)), global_head_num=4), - token_capacity=8, - device="cpu", - scales=torch.full((2, 8 if kind == "static_head" else 2), 99.0), - q_scales=torch.full((2, 4), 99.0) if kind == "static_head" else None, - ) - manager.register_state(role, other) - manager.finish_state_registration() - selected = getattr(manager, f"{role}_{template.get_state_type()}_state") - k, v = selected.get_layer_view(8) - infer = InferStateInfo() - infer.att_state_manager = manager - infer.is_draft_model = role == "draft0" - infer.att_state_role = role - infer.batch_size, infer.max_q_seq_len = 2, 1 - infer.b_seq_len = torch.ones(2, dtype=torch.int32, device="cpu") - backend = object.__new__(BaseAttBackend) - observed = {} - - if kind == "static_head": - state_cls = fa3.Fp8Fa3PrefillAttState if prefill else fa3.Fp8Fa3DecodeAttState - parent_cls = fa3.Fa3PrefillAttState if prefill else fa3.Fa3DecodeAttState - monkeypatch.setattr(parent_cls, "init_state", lambda self: None) - att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - att.b_att_seq_len = infer.b_seq_len - att.init_state() - - def quantize(q: torch.Tensor, scales: torch.Tensor): - observed["q_scales"] = scales - return q - - def attention(**kwargs: object): - observed.update(kwargs) - return kwargs["q"] - - monkeypatch.setattr(fa3, "q_per_head_static_fp8_quant", quantize) - monkeypatch.setattr(fa3, "flash_attn_with_kvcache", attention) - else: - state_cls = flashinfer.Fp8FlashInferPrefillAttState if prefill else flashinfer.Fp8FlashInferDecodeAttState - att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - wrapper = SimpleNamespace(run=lambda *args, **kwargs: observed.update(kwargs)) - att.prefill_wrapper = att.decode_wrapper = wrapper - - q = torch.zeros((2, selected.config.head_num, selected.config.head_dim), dtype=torch.bfloat16, device="cpu") - - def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.device | int): - return torch.empty(shape, dtype=dtype, device="cpu") - - if prefill: - att.prefill_att( - q=q, - k=k, - v=v, - att_control=AttControl(logical_layer_id=8), - alloc_func=allocate, - ) - else: - att.decode_att( - q=q, - k=k, - v=v, - att_control=AttControl(logical_layer_id=8), - alloc_func=allocate, - ) - if kind == "static_head": - heads = selected.config.head_num - torch.testing.assert_close(observed["q_scales"], selected.q_scales[1]) - torch.testing.assert_close(observed["k_descale"], selected.scales[1, :heads].expand(2, heads)) - torch.testing.assert_close(observed["v_descale"], selected.scales[1, heads:].expand(2, heads)) - else: - assert observed["k_scale"] == selected.cpu_scales[1, 0] - assert observed["v_scale"] == selected.cpu_scales[1, 1] - assert not hasattr(infer, "full_att_state") +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py index 0197b17daa..65a417cf49 100644 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -7,11 +7,11 @@ import pytest import torch +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.types import AxisMapping from unit_tests.common.att_resources.model_helpers import full_state_test_buffers -from unit_tests.common.att_resources.test_state_identity import allocated_state from unit_tests.common.att_resources.test_quantized_full_state import transfer_layout -from lightllm.common.attention_state.types import AxisMapping -from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from unit_tests.common.att_resources.test_state_identity import allocated_state @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -375,10 +375,12 @@ def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: cuda_ids = torch.tensor([1], dtype=torch.int64, device="cuda") cpu_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) pd_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + write_error = "KV writes require same-device tensors" if kind == "gqa" else "token indices must reside on CUDA" + copy_error = "token indices must match storage" if kind == "gqa" else "token indices must reside on CUDA" + with pytest.raises(AssertionError, match=write_error): state.write_full_att_kv(state.config.logical_layer_groups[0][0], cpu_ids, torch.empty(0, device="cuda")) for src, dst in ((cpu_ids, cuda_ids), (cuda_ids, cpu_ids)): - with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + with pytest.raises(AssertionError, match=copy_error): state.copy_full_att_tokens(src, dst) with pytest.raises(AssertionError, match="token indices must reside on CUDA"): state.copy_tokens_from_dp(src, dst, state) diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py index 1339ebb4b0..c25a024d00 100644 --- a/unit_tests/common/att_resources/test_full_state.py +++ b/unit_tests/common/att_resources/test_full_state.py @@ -1,12 +1,5 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - - from dataclasses import replace import pytest @@ -14,15 +7,16 @@ from lightllm.common.attention_state.full_att import ( FullAttentionState, - GQAFullAttentionState, - MLAFullAttentionState, GQAAttentionConfig, + GQAFullAttentionState, MLAConfig, + MLAFullAttentionState, ) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout, TransferRegion +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + read_kv_storage, ) @@ -97,7 +91,7 @@ def test_full_attention_write_copy_and_pd( data = torch.arange(2 * read_kv_storage(state).shape[2] * 4, dtype=torch.float32, device=device).reshape( 2, read_kv_storage(state).shape[2], 4 ) - if device == "cuda": + if device == "cuda" or cls is GQAFullAttentionState: state.write_full_att_kv(2, ids, data) state.write_full_att_kv(5, ids, data + 100) else: @@ -114,7 +108,7 @@ def test_full_attention_write_copy_and_pd( def reject_host_list(tensor: torch.Tensor): raise AssertionError("token copies must not convert indices to Python lists") - if device == "cuda": + if device == "cuda" or cls is GQAFullAttentionState: with monkeypatch.context() as patch: patch.setattr(torch.Tensor, "tolist", reject_host_list) state.copy_full_att_tokens(ids, torch.tensor([8, 9], dtype=torch.int32, device=device)) diff --git a/unit_tests/common/att_resources/test_gqa_layout_identity.py b/unit_tests/common/att_resources/test_gqa_layout_identity.py index f4c3174f44..4438389315 100644 --- a/unit_tests/common/att_resources/test_gqa_layout_identity.py +++ b/unit_tests/common/att_resources/test_gqa_layout_identity.py @@ -6,7 +6,6 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state -from lightllm.common.basemodel.attention.base_att import AttControl @pytest.mark.parametrize("kv_type", ["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"]) @@ -56,59 +55,3 @@ def test_gqa_layouts_coexist_and_match_across_tp_and_registration_order(kv_type, dst_req_idx=1, ) assert copied == [getattr(target, state.state_name) for state in states] - - -@pytest.mark.parametrize("prefill", [False, True]) -def test_fa3_reads_descale_for_each_head_layout(prefill, monkeypatch): - from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 - - manager = AttentionStateManager(request_capacity=2) - for layer, heads in enumerate((2, 4)): - config = GQAAttentionConfig( - logical_layer_groups=((layer,),), global_head_num=heads, head_dim=128, dtype=torch.bfloat16 - ) - state = create_gqa_state(config, kv_type="fp8kv_sph", device="cpu") - state.scales = torch.arange(1, 2 * heads + 1, dtype=torch.float32).view(1, -1) - state.q_scales = torch.arange(1, heads + 1, dtype=torch.float32).view(1, heads) - manager.register_state("main", state) - manager.finish_state_registration() - # prefill 必须按 b_seq_len 取 batch，不能读这个未设置的请求 batch。 - infer = SimpleNamespace( - att_state_manager=manager, - att_state_role="main", - batch_size=None, - b_seq_len=torch.ones(2, dtype=torch.int32), - max_q_seq_len=1, - ) - cls = fa3_fp8.Fp8Fa3PrefillAttState if prefill else fa3_fp8.Fp8Fa3DecodeAttState - parent = fa3_fp8.Fa3PrefillAttState if prefill else fa3_fp8.Fa3DecodeAttState - monkeypatch.setattr(parent, "init_state", lambda self: None) - monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", lambda q, scales: q) - observed = {} - - def attend(**kwargs): - observed.update(kwargs) - return kwargs["q"] - - monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) - att = cls(backend=None, _infer_state_weak_ref=lambda: infer) - att.b_att_seq_len = torch.ones(3, dtype=torch.int32) - att.page_table = att.cu_seqlens_q = att.cu_seqlens_k = None - att.decode_max_q_seq_len = 1 - att.causal = True - att.init_state() - batch_size = infer.b_seq_len.shape[0] if prefill else att.b_att_seq_len.shape[0] - for state in manager.get_full_attention_states("main"): - heads = state.config.head_num - layer = state.config.logical_layer_groups[0][0] - q = k = v = torch.zeros((2, heads, state.config.head_dim)) - observed.clear() - run = att.prefill_att if prefill else att.decode_att - run(q=q, k=k, v=v, att_control=AttControl(logical_layer_id=layer)) - row = state.scales[0] - k_descale = observed["k_descale"] - v_descale = observed["v_descale"] - assert k_descale.data_ptr() == row[:heads].data_ptr() - assert v_descale.data_ptr() == row[heads:].data_ptr() - torch.testing.assert_close(k_descale, row[:heads].expand(batch_size, heads)) - torch.testing.assert_close(v_descale, row[heads:].expand(batch_size, heads)) diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index fb9ce7ba5b..bce813f5cb 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -1,11 +1,11 @@ import pytest from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.attention_state.bounded_att import BoundedAttentionState +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.models.qwen3next.model import Qwen3NextTpPartModel -from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel from lightllm.server.core.objs.start_args_type import StartArgs diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py index 52337199a7..3b2f496fee 100644 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -7,19 +7,22 @@ import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import GQAAttentionConfig, FullAttentionState, GQAFullAttentionState -from lightllm.common.attention_state.bounded_att import KDAState, GDNState -from unit_tests.common.att_resources.test_bounded_state import make_state -from unit_tests.common.att_resources.test_window_compressor_state import make_window -from unit_tests.common.att_resources.test_capacity import definition -from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from lightllm.common.attention_state.bounded_att import GDNState, KDAState +from lightllm.common.attention_state.full_att import FullAttentionState, GQAAttentionConfig, GQAFullAttentionState from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.models.llama.model import LlamaTpPartModel from lightllm.common.basemodel.layer_infer.transformer_layer_infer import TransformerLayerInfer +from lightllm.models.llama.model import LlamaTpPartModel +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from unit_tests.common.att_resources.test_bounded_state import make_state +from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.att_resources.test_window_compressor_state import make_window def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: model = LlamaTpPartModel.__new__(LlamaTpPartModel) + from lightllm.platform import get_device_runtime + + model.device_runtime = get_device_runtime("cpu") model.args, model.config = SimpleNamespace(), {} model.tp_rank_, model.tp_world_size_ = 0, 1 model.layers_infer = [SimpleNamespace(layer_num_=0, indexer=SimpleNamespace())] @@ -65,6 +68,9 @@ def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, mo ), ) model = TpPartBaseModel.__new__(TpPartBaseModel) + from lightllm.platform import get_device_runtime + + model.device_runtime = get_device_runtime("cpu") model.is_mtp_draft_model = role != "main" model.mtp_previous_draft_models = [] model.max_req_num = 0 diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index 09e88b4a93..5dcaea3e97 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -1,20 +1,16 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import read_kv_storage - -from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states -from lightllm.common.attention_state.full_att import FullAttentionState -from unit_tests.common.att_resources.model_helpers import owned_states import dataclasses import pytest import torch +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.req_manager import ReqManager from lightllm.models.glm5_next.model import Glm5NextTpPartModel from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel -from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import ( get_added_mtp_kv_layer_num, @@ -23,7 +19,12 @@ set_env_start_args, set_unique_server_name, ) - +from unit_tests.common.att_resources.model_helpers import ( + main_linear_state, + owned_states, + prepare_model_states, + read_kv_storage, +) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py index 70970756df..e3e51d6616 100644 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -2,41 +2,38 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import read_kv_storage - -from unit_tests.common.att_resources.model_helpers import rebuild_full_state - -from typing import TYPE_CHECKING import dataclasses import json from functools import partial - -from unit_tests.common.att_resources.model_helpers import prepare_model_states +from typing import TYPE_CHECKING import pytest import torch +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator from lightllm.models.bloom.model import BloomTpPartModel +from lightllm.models.gemma3.model import Gemma3TpPartModel +from lightllm.models.gemma4.model import Gemma4TpPartModel +from lightllm.models.gemma_2b.model import Gemma_2bTpPartModel from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.models.qwen2.model import Qwen2TpPartModel -from lightllm.models.starcoder.model import StarcoderTpPartModel -from lightllm.models.starcoder2.model import Starcoder2TpPartModel from lightllm.models.mistral.model import MistralTpPartModel from lightllm.models.mixtral.model import MixtralTpPartModel -from lightllm.models.gemma_2b.model import Gemma_2bTpPartModel -from lightllm.models.gemma3.model import Gemma3TpPartModel -from lightllm.models.gemma4.model import Gemma4TpPartModel +from lightllm.models.qwen2.model import Qwen2TpPartModel +from lightllm.models.starcoder2.model import Starcoder2TpPartModel +from lightllm.models.starcoder.model import StarcoderTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import prepare_model_states, read_kv_storage, rebuild_full_state from unit_tests.common.att_resources.test_quantized_full_state import KINDS, make_quant_state if TYPE_CHECKING: + from pathlib import Path + from easydict import EasyDict + from lightllm.common.basemodel.basemodel import TpPartBaseModel - from pathlib import Path @pytest.fixture @@ -187,7 +184,6 @@ def memory_info(device: str | torch.device | int): monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 3) monkeypatch.setattr( "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) ) @@ -198,6 +194,9 @@ def memory_info(device: str | torch.device | int): device="meta", ) resources = AttentionStateManager(request_capacity=1) + from lightllm.platform.cuda import CudaRuntime + + resources.device_runtime = CudaRuntime("cuda:3") resources.register_state("main", rebuild_full_state(template, config=template.config)) resources.finish_state_registration() resources.initialize_storage( diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index ea114d2ba9..8d8d46abf1 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -1,7 +1,7 @@ -"""Verification commits State-owned recurrent data without model cursor aliases.""" - from __future__ import annotations +"""Verification commits State-owned recurrent data without model cursor aliases.""" + from types import SimpleNamespace @@ -9,8 +9,8 @@ import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import GDNState from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.common.attention_state.bounded_att import GDNState from lightllm.server.router.model_infer.mtp_speculative.utils import verify_mtp_tokens from unit_tests.common.att_resources.test_bounded_state import make_state @@ -82,15 +82,17 @@ def test_commit_without_recurrent_state_and_invalid_role(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("widths", [(3, 3, 3), (3, 1, 2)]) -def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...]): +@pytest.mark.parametrize("overlap", [False, True]) +def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...], overlap): manager = make_manager("cuda") backend = SimpleNamespace( model=SimpleNamespace( att_state_manager=manager, req_manager=SimpleNamespace( + request_generations=torch.zeros(3, dtype=torch.int64, device="cuda"), req_sampling_params_manager=SimpleNamespace( req_to_next_token_ids=torch.tensor([[-1, 10, 11]] * 3, device="cuda") - ) + ), ), ) ) @@ -100,12 +102,30 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i reqs.extend([req] * width) tokens.extend(99 if offset == accepted - 1 else 10 + offset for offset in range(width)) expected_mask.extend(int(offset < accepted) for offset in range(width)) + from lightllm.common.attention_state.commit import RequestGenerationStamp + + request_indices = torch.tensor(reqs, dtype=torch.int32, device="cuda") + mtp_indices = torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda") + + def output(start, end): + return SimpleNamespace( + attention_commit=RequestGenerationStamp.stamp( + "main", + request_indices[start:end], + mtp_indices[start:end], + backend.model.req_manager.request_generations, + ) + ) + + model_output = (output(0, starts[1]), output(starts[1], len(reqs))) if overlap else output(0, len(reqs)) + + request_indices.zero_() + mtp_indices.zero_() accept_lengths, accepted_index = verify_mtp_tokens( backend=backend, next_token_ids=torch.tensor(tokens, device="cuda"), - b_req_idx=torch.tensor(reqs, dtype=torch.int32, device="cuda"), - b_req_mtp_start_loc=torch.tensor(starts, dtype=torch.int32, device="cuda"), - b_mtp_index=torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda"), + model_output=model_output, + num_reqs=len(widths), ) assert accept_lengths.tolist() == [3, 1, 2] assert accepted_index.tolist() == expected_mask diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py index d7d285493b..6f048eeb4c 100644 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -1,7 +1,5 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import read_kv_storage - from contextlib import nullcontext from types import SimpleNamespace @@ -9,9 +7,10 @@ import torch from lightllm.common.attention_state import AttentionStateManager -from lightllm.models.llama.model import LlamaTpPartModel from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.llama.model import LlamaTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.model_helpers import read_kv_storage from unit_tests.common.att_resources.test_capacity import definition from unit_tests.common.basemodel.test_overlap_utils import _make_decode_input, _make_prefill_input @@ -46,8 +45,8 @@ def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.Mo draft, resources = draft_model(mode) draft.prepare_attention_states() resources.finish_state_registration() - draft._init_att_backend() - draft._init_att_backend1() + draft.prepare_attention_backends() + assert draft.attention.groups == () assert draft.att_state_manager is resources assert not hasattr(draft, "full_att_state") assert not resources.get_full_attention_states("draft0") @@ -60,15 +59,23 @@ def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.Mo draft._select_mem_indexes = lambda value: torch.zeros(value.batch_size, dtype=torch.int32) monkeypatch.setattr("lightllm.common.basemodel.basemodel.dist_group_manager.get_group", lambda index: None) model_input = _make_prefill_input([10, 11], req_idx=0, is_decode_req=False) if prefill else _make_decode_input() + draft.attention.groups = () state = draft._create_inferstate(model_input) assert not hasattr(state, "full_att_state") assert state.att_state_manager is resources assert state.is_draft_model - assert state.prefill_att_state is None and state.decode_att_state is None + assert state.attention.groups == () and state.attention.contexts == {} assert state.req_manager is draft.req_manager - state.init_att_state() + from lightllm.common.attention.api import AttentionPhase + from unit_tests.common.attention.helpers import batch_from_fixture + + state.attention.batch = batch_from_fixture(state, AttentionPhase.PREFILL if prefill else AttentionPhase.DECODE) if not prefill: - state.copy_for_cuda_graph(draft._create_inferstate(model_input)) + other = draft._create_inferstate(model_input) + other.attention.batch = state.attention.batch + for target, source in state.graph_input_copies(other): + target.copy_(source, non_blocking=True) + state.attention.plan_replay(other.attention) @pytest.mark.parametrize("mode", ["vanilla_with_att", "eagle_with_att"]) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py index de13eb28fc..664cf1938c 100644 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -1,16 +1,7 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration - -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from unit_tests.common.att_resources.model_helpers import read_kv_storage - - -from typing import TYPE_CHECKING from math import prod +from typing import TYPE_CHECKING import pytest import torch @@ -18,20 +9,32 @@ from lightllm.common.attention_state.full_att import ( GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState, - StaticFP8AttentionConfig, - StaticFP8FullAttentionState, PackedMLAAttentionConfig, - PackedMLALogicalConfig, PackedMLAFullAttentionState, + PackedMLALogicalConfig, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, ) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, +from lightllm.common.attention_state.types import AxisMapping, StateTransferLayout, TransferRegion +from unit_tests.common.att_resources.model_helpers import ( + bind_test_fp8_calibration, + full_state_test_buffers, + full_state_test_fields, + read_kv_storage, +) +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, ) if TYPE_CHECKING: from collections.abc import Sequence + from lightllm.common.attention_state.full_att.base_state import FullAttentionState @@ -335,10 +338,11 @@ def test_deferred_storage_preserves_real_fp8_calibration(kind: str): def test_int4_prefill_with_separate_storage(device, monkeypatch): import weakref from types import SimpleNamespace - from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonPrefillAttState + + from lightllm.common.attention.backends._shared.prefill import context_flashattention_nopad + from lightllm.common.attention.backends.dense.triton.int4 import TritonDenseInt4DecodeBackend from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.kv_copy import ppl_int4kv_copy_kv - from lightllm.common.basemodel.triton_kernel.att.prefill_att import context_flashattention_nopad if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA unavailable") @@ -366,10 +370,8 @@ def test_int4_prefill_with_separate_storage(device, monkeypatch): owner.b_req_idx = torch.tensor([0], dtype=torch.int32, device=device) owner.b_ready_cache_len = owner.b_q_start_loc = torch.zeros(1, dtype=torch.int32, device=device) owner.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, device=device).view(1, -1)) - prefill = Int4kvTritonPrefillAttState( - backend=SimpleNamespace(quant_group_size=8), _infer_state_weak_ref=weakref.ref(owner) - ) - prefill.init_state() + prefill = backend_context(TritonDenseInt4DecodeBackend, "prefill", settings=SimpleNamespace(quant_group_size=8)) + update_context(prefill, owner) if device == "cpu": def dequantize(**kwargs): @@ -384,5 +386,5 @@ def attention(**kwargs): monkeypatch.setattr(ppl_int4kv_copy_kv, "dequantize_int4kv", dequantize) monkeypatch.setattr(context_flashattention_nopad, "context_attention_fwd_contiguous_kv", attention) q = torch.randn(2, 2, 64, device=device) - output = prefill.prefill_att(q, (k, k_scale), (v, v_scale)) + output = run_dense(prefill, owner, q, (k, k_scale), (v, v_scale)) torch.testing.assert_close(output, torch.ones_like(q)) diff --git a/unit_tests/common/att_resources/test_validate_model_cache.py b/unit_tests/common/att_resources/test_validate_model_cache.py index 42508cd0ee..7c3ae3f8b4 100644 --- a/unit_tests/common/att_resources/test_validate_model_cache.py +++ b/unit_tests/common/att_resources/test_validate_model_cache.py @@ -1,3 +1,5 @@ +import importlib.util +from pathlib import Path from types import SimpleNamespace import pytest @@ -5,7 +7,12 @@ from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.full_att import create_compressed_att_state from lightllm.models.deepseek_v4.config import DSV4AttentionConfig -from test.att_resources import validate_model_cache + +_spec = importlib.util.spec_from_file_location( + "validate_model_cache", Path(__file__).resolve().parents[3] / "test/att_resources/validate_model_cache.py" +) +validate_model_cache = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(validate_model_cache) @pytest.mark.parametrize("draft_count,bad_role", [(0, None), (2, None), (2, "main"), (2, "draft1")]) diff --git a/unit_tests/common/attention/helpers.py b/unit_tests/common/attention/helpers.py new file mode 100644 index 0000000000..7e466ff46e --- /dev/null +++ b/unit_tests/common/attention/helpers.py @@ -0,0 +1,418 @@ +from lightllm.common.attention.api.environment import BackendEnv +from lightllm.common.attention.api.spec import AttentionFeature +from lightllm.common.attention.session import AttentionSession +import importlib +from dataclasses import replace +from types import SimpleNamespace + +import torch + +from lightllm.common.attention.api import ( + AttentionBatch, + AttentionFamily, + AttentionPhase, + CacheLayout, + ContextCapacity, + LayerAttention, + LayerSpec, +) +from lightllm.common.attention.api.workspace import AttentionWorkspace +from lightllm.common.attention.api.backend import AttentionBackend +from lightllm.common.basemodel.attention import describe_serving_constraints +from lightllm.platform import get_device_runtime +from lightllm.server.core.objs.start_args_type import StartArgs + + +def make_batch(lengths=(5, 2), *, phase=AttentionPhase.DECODE, device="cpu", **kwargs): + total = sum(lengths) + maximum = max(lengths) if lengths else 0 + lengths = torch.tensor(lengths, dtype=torch.int32, device=device) + count = lengths.numel() + q_lengths = torch.ones_like(lengths) + offsets = torch.arange(count + 1, dtype=torch.int32, device=device) + values = dict( + phase=phase, + batch_size=count, + token_num=count, + total_token_num=total, + max_q_seq_len=1, + max_kv_seq_len=maximum, + max_cache_len=0, + b_req_idx=torch.arange(count, dtype=torch.int32, device=device), + b_seq_len=lengths, + b_q_seq_len=q_lengths, + b1_cu_q_seq_len=offsets, + b1_cu_kv_seq_len=torch.cat((offsets[:1], lengths.cumsum(0).int())), + req_to_token_indexs=torch.zeros((max(count, 8), 256), dtype=torch.int32, device=device), + hold_req_id=7, + b_mtp_index=torch.zeros_like(lengths), + b_ready_cache_len=torch.zeros_like(lengths), + ) + values.update(kwargs) + return AttentionBatch(**values) + + +def make_spec(family=AttentionFamily.DENSE, phase=AttentionPhase.DECODE, **kwargs): + values = dict( + family=family, + phase=phase, + dtype=torch.bfloat16, + cache=CacheLayout( + family.value if family in (AttentionFamily.GDN, AttentionFamily.KDA) else "None", 1, torch.bfloat16, 2, 8, 8 + ), + num_q_heads=4, + head_size=8, + value_head_size=8, + ) + features = set(kwargs.pop("features", ())) + for key, feature in ( + ("use_mtp", "MTP"), + ("dynamic_verify", "DYNAMIC_VERIFY"), + ("use_alibi", "ALIBI"), + ("use_sink", "SINK"), + ): + if kwargs.pop(key, False): + features.add(AttentionFeature[feature]) + if not kwargs.pop("causal", True): + features.add(AttentionFeature.NONCAUSAL) + if kwargs.get("sliding_window") is not None: + features.add(AttentionFeature.SLIDING_WINDOW) + if AttentionFeature.MTP in features: + features.add(AttentionFeature.MTP_WINDOW) + values.update(kwargs, features=frozenset(features)) + return LayerSpec(**values) + + +def make_backend(cls=None, **kwargs): + cls = cls or StubBackend + backend = cls.__new__(cls) + values = dict( + spec=make_spec(cls.family, cls.phase), + alloc_tensor=lambda shape, dtype, device="cuda": torch.empty(shape, dtype=dtype, device=device), + page_size=1, + causal=True, + dynamic_verify=False, + draft_step=0, + state_slots=1, + is_draft_model=False, + hold_req_id=7, + ) + values.update(kwargs) + if "mtp_step" in kwargs: + values["state_slots"] = kwargs["mtp_step"] + 1 + page_size = values.pop("kernel_page_size", values.pop("page_size", 1)) + spec = values["spec"] + features = set(spec.features) + if not values.pop("causal", True): + features.add(AttentionFeature.NONCAUSAL) + if values.pop("dynamic_verify", False): + features.add(AttentionFeature.DYNAMIC_VERIFY) + values["spec"] = replace( + spec, + cache=replace( + spec.cache, + page_size=page_size, + quant_group_size=values.pop("quant_group_size", spec.cache.quant_group_size), + ), + features=frozenset(features), + ) + if "softmax_scale" in values: + values["spec"] = replace(values["spec"], softmax_scale=values.pop("softmax_scale")) + vars(backend).update(values) + return backend + + +def make_context(backend, batch=None, *, capture=False, max_length=None, prepare=False): + batch = make_batch() if batch is None else batch + capacity = ContextCapacity( + phase=batch.phase, + batch_size=batch.batch_size, + token_num=batch.token_num, + max_kv_seq_len=max_length or batch.max_kv_seq_len, + max_q_seq_len=batch.max_q_seq_len, + ) + backend = phase_backend(backend, batch.phase) + context = AttentionSession.allocate(backend, capacity, capture=capture, workspace=AttentionWorkspace(batch.device)) + if prepare: + plan(backend, context, batch) + else: + context.batch = batch + context.current_max_kv_seq_len = batch.max_kv_seq_len + return context + + +def make_layer_attention(backend, context, params, **spec_kwargs): + if not hasattr(backend, "spec"): + backend.spec = make_spec(backend.family, backend.phase) + if spec_kwargs: + backend.spec = replace(backend.spec, **spec_kwargs) + backend.alloc_tensor = lambda shape, dtype, device="cpu": torch.empty( + shape, dtype=dtype, device=context.batch.device + ) + return LayerAttention( + group_id=0, + logical_layer_id=0, + resources=params, + backend=backend, + ) + + +def model_stub(**args): + args.setdefault("disable_cudagraph", not torch.cuda.is_available()) + model = SimpleNamespace( + args=StartArgs(**args), + device_runtime=get_device_runtime("cuda" if torch.cuda.is_available() else "cpu"), + att_state_role="main", + att_state_manager=SimpleNamespace( + get_full_attention_states=lambda *args: (), get_bounded_attention_states=lambda *args: () + ), + config={}, + is_mtp_draft_model=False, + tp_world_size_=1, + tp_rank_=0, + max_seq_length=256, + batch_max_tokens=32, + layers_infer=[], + req_manager=SimpleNamespace(HOLD_REQUEST_ID=7), + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), + ) + model.attention_requirements = describe_serving_constraints(model) + return model + + +def batch_from_fixture(fixture, phase, token_count=None, device=None): + indices = getattr(fixture, "b_req_idx", None) + lengths = getattr(fixture, "b_seq_len", None) + device = device or ( + indices.device if indices is not None else getattr(getattr(fixture, "mem_index", None), "device", "cpu") + ) + count = indices.numel() if indices is not None else (token_count or getattr(fixture, "batch_size", None) or 1) + if lengths is None: + lengths = torch.ones(count, dtype=torch.int32, device=device) + values = {} + for name in AttentionBatch.__dataclass_fields__: + value = getattr(fixture, name, None) + if value is not None and name not in ("phase", "token_num", "batch_size"): + values[name] = value + manager = getattr(fixture, "req_manager", None) + if manager is not None: + for name in ("req_to_token_indexs",): + value = getattr(manager, name, None) + if value is not None: + values[name] = value + values["hold_req_id"] = getattr(manager, "HOLD_REQUEST_ID", 7) + + max_len = getattr(fixture, "max_kv_seq_len", None) or 1 + cache_key = (phase, count, str(device)) + cache = getattr(fixture, "_attention_test_batches", None) + if cache is None: + cache = fixture._attention_test_batches = {} + if cache_key not in cache: + cache[cache_key] = make_batch((1,) * count, phase=phase, device=device) + defaults = cache[cache_key] + values.update( + phase=phase, batch_size=count, token_num=token_count or count, b_seq_len=lengths, max_kv_seq_len=max_len + ) + return replace(defaults, **values) + + +def backend_context(cls, phase, settings=None, **values): + phase = AttentionPhase(phase) + if isinstance(settings, AttentionBackend): + backend = settings + else: + config = dict(vars(settings)) if settings is not None else {} + model = config.pop("model", None) + if model is not None: + manager = getattr(model, "mtp_manager", None) + if manager is not None: + config["draft_step"] = manager.get_decode_draft_step(getattr(model, "is_mtp_draft_model", False)) + req_manager = getattr(model, "req_manager", None) + if req_manager is not None: + config["hold_req_id"] = getattr(req_manager, "HOLD_REQUEST_ID", 7) + backend = make_backend(cls, **config) + defaults = dict( + network_config={"index_topk": 2048, "index_kpool": 4}, + prefill_graph_enabled=False, + decode_only=False, + prefill_token_capacity=0, + mtp_step=0, + draft_step=0, + ) + for name, value in defaults.items(): + if not hasattr(backend, name): + setattr(backend, name, value) + context = make_context(backend, make_batch((1,), phase=phase)) + if hasattr(backend, "_test_workspace"): + context.workspace = backend._test_workspace + for name, value in values.items(): + setattr(context, name, value) + return context + + +def update_context(context, fixture): + context.batch = batch_from_fixture(fixture, context.capacity.phase) + context.capacity = ContextCapacity( + phase=context.batch.phase, + batch_size=context.batch.batch_size, + token_num=context.batch.token_num, + max_kv_seq_len=context.batch.max_kv_seq_len, + max_q_seq_len=context.batch.max_q_seq_len, + ) + plan(context.backend_identity, context, context.batch) + build(context.backend_identity, context) + + +def run_sparse(context, fixture, query, kv, parameters, *, out=None, alloc_func=None): + from lightllm.common.attention.api.inputs import MLAResources, SparseMLADecodeInputs, SparseMLAPrefillInputs + + q = query[0] if isinstance(query, tuple) else query + context.capture = bool(getattr(fixture, "is_cuda_graph_capture", context.capture)) + if getattr(fixture, "prefill_cuda_graph_exe_list", None) is not None: + context.capture = True + context.batch = batch_from_fixture(fixture, context.capacity.phase, q.shape[0], q.device) + backend = context.backend_identity + if context.workspace.device != q.device: + context.workspace = AttentionWorkspace(q.device) + backend._test_workspace = context.workspace + if backend.__class__.__module__.startswith("lightllm.models.deepseek_v4"): + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4Inputs, DeepseekV4Resources + + params = DeepseekV4Resources( + kv[0], + kv[0], + kv[1], + parameters["compress_ratio"], + parameters["head_dim_v"], + parameters["softmax_scale"], + parameters["attn_sink"], + ) + inputs = DeepseekV4Inputs( + q, + parameters["swa_indices"], + parameters["swa_lengths"], + parameters.get("extra_indices"), + parameters.get("extra_lengths"), + ) + else: + params = MLAResources(kv, parameters["kv_lora_rank"], parameters.get("qk_rope_head_dim", 0)) + if context.capacity.phase is AttentionPhase.PREFILL: + inputs = SparseMLAPrefillInputs( + q, + parameters.get("topk_mem_indices"), + parameters.get("topk_indices"), + parameters.get("prefill_cache_kv"), + ) + else: + inputs = SparseMLADecodeInputs(q, query[1], parameters["topk_mem_indices"]) + attention = make_layer_attention(backend, context, params, softmax_scale=parameters["softmax_scale"]) + if alloc_func is not None: + backend.alloc_tensor = alloc_func + + if ( + hasattr(context, "capture_flashmla_sched_meta") + and context.capture_flashmla_sched_meta is None + and torch.cuda.is_available() + ): + import flash_mla + + context.capture_flashmla_sched_meta = flash_mla.get_mla_metadata()[0] + return run_attention(attention, context, inputs, out=out) + + +def run_dense(context, fixture, query, key, value, *, out=None, alloc_func=None, params=None): + from lightllm.common.attention.api.inputs import DenseInputs, DenseResources + + context.batch = batch_from_fixture(fixture, context.capacity.phase, query.shape[0], query.device) + if params is None: + ks = vs = None + if isinstance(key, tuple): + key, ks = key + value, vs = value + params = DenseResources(key, value, key_scale=ks, value_scale=vs) + attention = make_layer_attention(context.backend_identity, context, params) + if alloc_func is not None: + context.backend_identity.alloc_tensor = alloc_func + return run_attention(attention, context, DenseInputs(query), out=out) + + +class StubBackend(AttentionBackend): + family = AttentionFamily.DENSE + provider = "triton" + phase = AttentionPhase.DECODE + + def allocate_context(self, capacity, *, capture, workspace): + from lightllm.common.attention.api import AttentionContext + + return AttentionContext(capacity, capture, workspace) + + def forward(self, layer, context, inputs, *, out=None): + return inputs + + +def make_env(model=None, **kwargs): + model = model or model_stub() + config = dict(model.config) + for name in ("qk_nope_head_dim", "qk_rope_head_dim", "kv_lora_rank", "v_head_dim"): + if hasattr(model, name): + config[name] = getattr(model, name) + values = dict( + runtime=model.device_runtime, + network_config=config, + hold_req_id=getattr(model.req_manager, "HOLD_REQUEST_ID", 7), + draft_step=model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model), + state_slots=model.args.mtp_step + 1, + max_seq_len=getattr(model, "max_seq_length", 256), + tp_size=getattr(model, "tp_world_size_", 1), + alloc_tensor=lambda shape, dtype, device=model.device_runtime.device: torch.empty( + shape, dtype=dtype, device=device + ), + layer_ids=tuple(layer.layer_num_ for layer in getattr(model, "layers_infer", ())), + ) + values.update(kwargs) + return BackendEnv(**values) + + +def phase_class(cls, phase): + suffix = "PrefillBackend" if phase is AttentionPhase.PREFILL else "DecodeBackend" + name = cls.__name__.replace("PrefillBackend", suffix).replace("DecodeBackend", suffix) + return getattr(importlib.import_module(cls.__module__), name) + + +def phase_backend(backend, phase): + if not hasattr(backend, "spec"): + backend.spec = make_spec(backend.family, backend.phase) + if backend.phase is not phase: + cls = phase_class(type(backend), phase) + backend.__class__ = cls + backend.spec = replace(backend.spec, phase=phase) + return backend + + +def plan(backend, context, batch): + return AttentionSession.plan(backend, context, batch) + + +def build(backend, context): + return AttentionSession.build(backend, context) + + +def run_attention(layer, context, inputs=None, *, out=None): + session = AttentionSession() + session.contexts[layer.group_id] = context + return layer.forward(session, inputs, out=out) + + +def dense_class(provider, phase): + from lightllm.common.attention.registry import BACKENDS, load_backend_class + + return next( + cls + for reference in BACKENDS + for cls in (load_backend_class(reference),) + if cls.family is AttentionFamily.DENSE + and cls.phase is phase + and cls.provider == provider + and cls.cache_encodings == {"plain"} + and not cls.variant + ) diff --git a/unit_tests/common/attention/test_backend_options.py b/unit_tests/common/attention/test_backend_options.py new file mode 100644 index 0000000000..bba0fe4d41 --- /dev/null +++ b/unit_tests/common/attention/test_backend_options.py @@ -0,0 +1,152 @@ +from dataclasses import fields +from types import SimpleNamespace + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +def test_sparse_backend_options_preserve_flashmla_only(phase, capsys): + parser = make_argument_parser() + option = f"llm_{phase}_att_backend" + choices = next(field.metadata["choices"] for field in fields(StartArgs) if field.name == option) + assert "flashmla_sparse" in choices + assert "flashinfer_sparse" not in choices + args = parser.parse_args([f"--{option}", "flashmla_sparse"]) + assert getattr(args, option) == "flashmla_sparse" + with pytest.raises(SystemExit): + parser.parse_args([f"--{option}", "flashinfer_sparse"]) + assert "invalid choice" in capsys.readouterr().err + + +BACKEND_OPTIONS = ( + "llm_prefill_att_backend", + "llm_decode_att_backend", + "llm_prefill_linear_att_backend", + "llm_decode_linear_att_backend", + "vit_att_backend", +) + + +@pytest.mark.parametrize("option", BACKEND_OPTIONS) +def test_backend_options_have_scalar_defaults_and_matching_choices(option): + parser = make_argument_parser() + field = next(field for field in fields(StartArgs) if field.name == option) + action = next(action for action in parser._actions if action.dest == option) + assert getattr(parser.parse_args([]), option) == getattr(StartArgs(), option) == "auto" + choices = list(action.choices) + if option == "llm_prefill_att_backend": + choices.remove("flashqla") + assert choices == field.metadata["choices"] + assert action.nargs == ("+" if option in BACKEND_OPTIONS[:2] else None) + + +def test_backend_options_are_independent(): + args = make_argument_parser().parse_args( + [ + "--llm_prefill_att_backend", + "fa3", + "--llm_decode_att_backend", + "flashinfer", + "--llm_prefill_linear_att_backend", + "flashqla", + "--llm_decode_linear_att_backend", + "triton", + ] + ) + assert [getattr(args, option) for option in BACKEND_OPTIONS] == ["fa3", "flashinfer", "flashqla", "triton", "auto"] + + +@pytest.mark.parametrize("option", BACKEND_OPTIONS) +def test_backend_options_reject_multiple_values(option, capsys): + if option in BACKEND_OPTIONS[:2]: + with pytest.warns(FutureWarning, match="deprecated"): + args = make_argument_parser().parse_args([f"--{option}", "fa3", "triton"]) + assert getattr(args, option) == "fa3" + assert getattr(args, option.replace("_att_backend", "_linear_att_backend")) == "triton" + return + with pytest.raises(SystemExit): + make_argument_parser().parse_args([f"--{option}", "triton", "triton"]) + assert "unrecognized arguments" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "option", ["llm_prefill_att_backend", "llm_decode_att_backend", "llm_decode_linear_att_backend"] +) +def test_flashqla_is_only_a_linear_prefill_choice(option, capsys): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([f"--{option}", "flashqla"]) + assert "invalid choice" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "decode,linear_decode,step,expected", + [ + ("triton", "auto", 2, True), + ("fa3", "triton", 2, False), + ("not_triton", "triton", 2, False), + ("triton", "triton", 0, False), + ], +) +def test_mtp_kernel_uses_exact_regular_decode_backend(monkeypatch, decode, linear_decode, step, expected): + from lightllm.utils import envs_utils + + monkeypatch.setattr( + envs_utils, + "get_env_start_args", + lambda: SimpleNamespace( + mtp_step=step, + llm_decode_att_backend=decode, + llm_decode_linear_att_backend=linear_decode, + ), + ) + envs_utils.enable_triton_mtp_kernel.cache_clear() + try: + assert envs_utils.enable_triton_mtp_kernel() is expected + finally: + envs_utils.enable_triton_mtp_kernel.cache_clear() + + +@pytest.mark.parametrize("requested", ["auto", "triton", "flashqla"]) +def test_kda_option_reaches_the_shared_selector(monkeypatch, requested): + from lightllm.common.attention import selection + from lightllm.common.attention.api import AttentionFamily + from lightllm.common.basemodel import attention + from lightllm.models.glm5_next import model as glm5 + from unit_tests.common.attention.helpers import make_spec + + from lightllm.common.basemodel.attention_adapters import ADAPTERS, AttentionDescription + from lightllm.common.attention.api import AttentionPhase + + calls = [] + sentinel = object() + + def select(**kwargs): + calls.append(kwargs) + return lambda env, spec: sentinel + + monkeypatch.setattr(selection, "select_attention_backend", select) + monkeypatch.setattr( + ADAPTERS[AttentionFamily.KDA], "describe", lambda model, layer, phase, desc: make_spec(desc.family, phase) + ) + model = glm5.Glm5NextTpPartModel.__new__(glm5.Glm5NextTpPartModel) + from unit_tests.common.attention.helpers import model_stub + + vars(model).update(vars(model_stub(llm_prefill_linear_att_backend=requested))) + model.layers_infer = [ + SimpleNamespace( + layer_num_=0, + describe_attention=lambda phase: AttentionDescription(AttentionFamily.KDA), + attention_adapter=lambda family: ADAPTERS[family], + ) + ] + model.prepare_attention_backends() + assert calls[0]["requested"] == requested + assert calls[0]["specs"][0].family is AttentionFamily.KDA + assert ( + model.attention.groups[0][AttentionPhase.PREFILL] + is model.attention.groups[0][AttentionPhase.DECODE] + is sentinel + ) diff --git a/unit_tests/common/attention/test_contracts.py b/unit_tests/common/attention/test_contracts.py new file mode 100644 index 0000000000..00b872c006 --- /dev/null +++ b/unit_tests/common/attention/test_contracts.py @@ -0,0 +1,681 @@ +from unit_tests.common.attention.helpers import plan, run_attention +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention import registry, selection +from lightllm.common.attention.api import AttentionFamily, AttentionPhase, UnsupportedAttention +from unit_tests.common.attention.helpers import StubBackend, make_env, phase_class, build +from lightllm.utils.backend_validator import ProbeResult, ProbeScope +from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub + + +def _select(*, model, specs=None, **kwargs): + family, phase = kwargs.pop("family"), kwargs.pop("phase") + candidates = kwargs.pop("candidates", None) + priority = kwargs.pop("priority", ()) + if candidates is not None: + kwargs["policy"] = selection.BackendPolicy( + allow=frozenset(priority), + prefer=tuple(priority), + extra=tuple(backend for formats in candidates.values() for backend in formats.values()), + ) + return selection.select_attention_backend( + specs=specs or (make_spec(family, phase),), + caps=model.device_runtime.caps, + execution=model.attention_requirements, + **kwargs, + ) + + +def test_graph_copy_validates_all_inputs_before_writing(): + captured = make_batch() + incoming = make_batch((9, 7)) + incoming = replace(incoming, req_to_token_indexs=captured.req_to_token_indexs, b_seq_len=incoming.b_seq_len.long()) + before = captured.b_req_idx.clone() + with pytest.raises(ValueError, match="b_seq_len"): + captured.copy_inputs(incoming) + torch.testing.assert_close(captured.b_req_idx, before) + + +def test_graph_copy_keeps_tensor_addresses_and_uses_new_lengths(): + captured = make_batch() + incoming = make_batch((9, 7), req_to_token_indexs=captured.req_to_token_indexs) + pointer = captured.b_seq_len.data_ptr() + captured.copy_inputs(incoming) + assert captured.b_seq_len.data_ptr() == pointer + assert captured.b_seq_len.tolist() == [9, 7] + with pytest.raises(ValueError, match="request table"): + captured.copy_inputs(make_batch((9, 7))) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_context_capacity_and_backend_identity(capture): + backend = StubBackend(make_env(), make_spec()) + context = make_context(backend, capture=capture, max_length=16) + plan(backend, context, make_batch(req_to_token_indexs=context.batch.req_to_token_indexs)) + with pytest.raises(ValueError, match="capacity"): + plan(backend, context, make_batch((17, 2))) + with pytest.raises(ValueError, match="different attention backend"): + plan(StubBackend(make_env(), make_spec()), context, make_batch()) + with pytest.raises(ValueError, match="different attention backend"): + build(StubBackend(make_env(), make_spec()), context) + attention = make_layer_attention(StubBackend(make_env(), make_spec()), context, params=None) + with pytest.raises(ValueError, match="different attention backend"): + run_attention(attention, context, inputs=None) + + +def test_layer_attention_rejects_wrong_phase(): + backend = StubBackend(make_env(), make_spec()) + context = make_context(backend, prepare=True) + attention = make_layer_attention(backend, context, params=None) + context.capacity = replace(context.capacity, phase=AttentionPhase.PREFILL) + with pytest.raises(ValueError, match="different attention backend"): + run_attention(attention, context, inputs=None) + + +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_selection_never_falls_back_to_a_rejected_backend(monkeypatch, requested): + spec = make_spec(cache=replace(make_spec().cache, encoding="fp8kv_sph")) + monkeypatch.setattr( + StubBackend, + "check_environment", + classmethod(lambda cls, caps: None), + ) + with pytest.raises(selection.BackendSelectionError, match="encoding fp8kv_sph"): + _select( + family=spec.family, + phase=spec.phase, + model=model_stub(), + specs=(spec,), + requested=requested, + candidates={"None": {"triton": StubBackend}}, + priority=["triton"], + ) + + +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_environment_checks_apply_to_explicit_and_auto(monkeypatch, requested): + def reject(cls, caps): + raise UnsupportedAttention("missing extension") + + monkeypatch.setattr(StubBackend, "check_environment", classmethod(reject)) + with pytest.raises(selection.BackendSelectionError, match="missing extension"): + _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested=requested, + candidates={"None": {"triton": StubBackend}}, + priority=["triton"], + ) + + +def _backend_check_exception_candidates(monkeypatch, error_type=RuntimeError): + class BrokenBackend(StubBackend): + provider = "broken" + + @classmethod + def check_environment(cls, caps): + raise error_type("provider ABI mismatch") + + class WorkingBackend(StubBackend): + provider = "working" + + @classmethod + def check_environment(cls, caps): + return None + + monkeypatch.setattr( + "lightllm.utils.backend_validator.probe_backend", + lambda *args, **kwargs: ProbeResult(True, ProbeScope.BACKEND_NUMERICS), + ) + return {"None": {"broken": BrokenBackend, "working": WorkingBackend}} + + +@pytest.mark.parametrize("error_type", [UnsupportedAttention, ValueError, RuntimeError]) +def test_backend_check_exception_allows_auto_fallback(monkeypatch, error_type): + candidates = _backend_check_exception_candidates(monkeypatch, error_type) + selected = _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested="auto", + candidates=candidates, + priority=["broken", "working"], + ) + assert selected.provider == "working" + + +@pytest.mark.parametrize("error_type", [UnsupportedAttention, ValueError, RuntimeError]) +def test_backend_check_exception_is_reported_for_explicit_selection(monkeypatch, error_type): + candidates = _backend_check_exception_candidates(monkeypatch, error_type) + with pytest.raises(selection.BackendSelectionError) as error: + _select( + family=AttentionFamily.DENSE, + phase=AttentionPhase.DECODE, + model=model_stub(), + requested="broken", + candidates=candidates, + priority=["broken", "working"], + ) + + reason = "provider ABI mismatch" + if error_type is not UnsupportedAttention: + reason = f"provider check failed: {error_type.__name__}: {reason}" + assert error.value.failures == {"broken": reason} + + +@pytest.mark.parametrize("requested", ["auto", "triton"]) +def test_rejection_reaches_rank_consensus(monkeypatch, requested): + spec = make_spec(dtype=torch.float64) + ranks = [] + monkeypatch.setattr(StubBackend, "check_environment", classmethod(lambda cls, caps: None)) + monkeypatch.setattr( + selection.backend_validator, "probe_backend", lambda *args, **kwargs: pytest.fail("probed rejected inputs") + ) + monkeypatch.setattr(selection, "_check_rank_selection", lambda *args: ranks.append(args)) + with pytest.raises(selection.BackendSelectionError, match="activation dtype"): + _select( + family=spec.family, + phase=spec.phase, + model=model_stub(), + specs=(spec,), + requested=requested, + candidates={"None": {"triton": StubBackend}}, + priority=["triton"], + ) + assert len(ranks) == 1 + assert ranks[0][:3] == (spec.family, spec.phase, None) + assert "activation dtype" in ranks[0][3] + + +def test_rank_selection_checks_all_ranks_after_local_selection(monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + + def gather(output, value): + output[:] = [value, {"backend": "triton", "failure": None}] + + monkeypatch.setattr(dist, "all_gather_object", gather) + with pytest.raises(selection.BackendSelectionError, match="backend selection"): + selection._check_rank_selection(AttentionFamily.DENSE, AttentionPhase.DECODE, "fa3", None) + + +def test_rank_selection_propagates_remote_failure(monkeypatch): + import torch.distributed as dist + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + monkeypatch.setattr( + dist, + "all_gather_object", + lambda output, value: output.__setitem__(slice(None), [value, {"backend": None, "failure": "probe failed"}]), + ) + with pytest.raises(selection.BackendSelectionError, match="rank 1: probe failed"): + selection._check_rank_selection(AttentionFamily.DENSE, AttentionPhase.DECODE, "fa3", None) + + +@pytest.mark.parametrize("format,provider", [("fp8kv_sph", "flashinfer"), ("fp8kv_spt", "fa3")]) +def test_fp8_scale_layouts_are_not_interchangeable(format, provider): + from lightllm.common.attention.api.spec import CacheEncoding + + backends = (registry.load_backend_class(ref) for ref in registry.BACKENDS) + assert not any(cls.provider == provider and CacheEncoding(format) in cls.cache_encodings for cls in backends) + + +def test_layer_shape_configures_group_planner(): + from lightllm.common.attention.backends.dense.flashinfer.backend import FlashInferDenseDecodeBackend + + model = model_stub(page_size=16) + model.config = {"num_attention_heads": 16, "num_key_value_heads": 4, "hidden_size": 2048} + model.max_seq_length = 128 + model.data_type = torch.bfloat16 + small = FlashInferDenseDecodeBackend(make_env(model), make_spec(num_q_heads=8, head_size=64)) + large = FlashInferDenseDecodeBackend(make_env(model), make_spec(num_q_heads=16, head_size=128)) + assert (small.tp_q_head_num, small.head_dim) == (8, 64) + assert (large.tp_q_head_num, large.head_dim) == (16, 128) + + +def test_groups_follow_each_layers_family_and_shape(monkeypatch): + from lightllm.common.basemodel.attention import ModelAttention + from lightllm.common.basemodel.attention_adapters import ADAPTERS, AttentionDescription + + specs = [make_spec(head_size=64), make_spec(head_size=64), make_spec(head_size=128), make_spec(AttentionFamily.GDN)] + calls = [] + + def select(**kwargs): + calls.append(kwargs) + return lambda env, spec: SimpleNamespace(spec=spec, family=spec.family) + + monkeypatch.setattr(selection, "select_attention_backend", select) + for adapter in ADAPTERS.values(): + monkeypatch.setattr(adapter, "describe", lambda model, layer, phase, desc: replace(layer.spec, phase=phase)) + model = model_stub() + model.attention_policy = lambda family, phase: selection.BackendPolicy() + model.layers_infer = [ + SimpleNamespace( + layer_num_=i, + spec=spec, + describe_attention=lambda phase, s=spec: AttentionDescription(s.family), + attention_adapter=lambda family: ADAPTERS[family], + ) + for i, spec in enumerate(specs) + ] + attention = ModelAttention(model) + assert attention.routes == (0, 0, 1, 2) + assert len(calls) == 6 + assert [group[AttentionPhase.PREFILL].spec.head_size for group in attention.groups[:2]] == [64, 128] + assert attention.groups[-1][AttentionPhase.DECODE].family is AttentionFamily.GDN + + +@pytest.mark.parametrize("invalid", ["input_ids", "group", "role"]) +def test_runner_graph_copy_rejects_before_mutating_batch(invalid): + from lightllm.common.basemodel.infer_struct import InferStateInfo + + captured, incoming = InferStateInfo(), InferStateInfo() + captured.att_state_role = incoming.att_state_role = "main" + group = object() + captured.attention.groups = incoming.attention.groups = (group,) + captured.attention.batch = make_batch() + incoming.attention.batch = make_batch((9, 7), req_to_token_indexs=captured.attention.batch.req_to_token_indexs) + captured.input_ids = torch.zeros(2, dtype=torch.long) + incoming.input_ids = torch.ones(3 if invalid == "input_ids" else 2, dtype=torch.long) + if invalid == "group": + incoming.attention.groups = (object(),) + if invalid == "role": + incoming.att_state_role = "draft" + before = captured.attention.batch.b_seq_len.clone() + with pytest.raises(ValueError): + captured.graph_input_copies(incoming) + torch.testing.assert_close(captured.attention.batch.b_seq_len, before) + assert captured.input_ids.tolist() == [0, 0] + + +def test_replay_flashqla_rejects_decode_and_fp32(): + from lightllm.common.attention.backends.gdn.replay.backend import FlashQlaGDNReplayPrefillBackend + + spec = make_spec(family=AttentionFamily.GDN) + spec = replace(spec, cache=replace(spec.cache, state_mode="replay")) + with pytest.raises(UnsupportedAttention, match="decode"): + FlashQlaGDNReplayPrefillBackend.check_support(spec) + with pytest.raises(UnsupportedAttention, match="dtype"): + FlashQlaGDNReplayPrefillBackend.check_support(replace(spec, phase=AttentionPhase.PREFILL, dtype=torch.float32)) + + +class _OutputBackend(StubBackend): + def forward(self, layer, context, inputs, *, out=None): + return inputs + + +@pytest.mark.parametrize("strided", [False, True]) +def test_output_copy_and_padded_view_preserve_auxiliary_result(strided): + backend = _OutputBackend(make_env(), make_spec()) + context = make_context(backend, prepare=True) + layer = make_layer_attention(backend, context, params=None) + out = torch.empty(2, 8 if strided else 4, 8)[:, :4] + gate = torch.ones(2) + output, actual_gate = run_attention(layer, context, (torch.ones(2, 32), gate), out=out) + assert output is out and actual_gate is gate + assert torch.all(out == 1) + valid_heads = out[:, :2] + output, actual_gate = run_attention(layer, context, (valid_heads, gate), out=out) + assert output is valid_heads and actual_gate is gate + + +@pytest.mark.parametrize("invalid", ["dtype", "device", "size", "view_bounds"]) +def test_output_rejects_incompatible_buffer(invalid): + backend = _OutputBackend(make_env(), make_spec()) + context = make_context(backend, prepare=True) + layer = make_layer_attention(backend, context, params=None) + out = torch.empty(2, 4) + result = torch.ones_like(out) + if invalid == "dtype": + result = result.double() + elif invalid == "device": + result = result.to("meta") + elif invalid == "size": + result = result[:1] + else: + result, out = out, out[:1] + with pytest.raises(ValueError, match="out"): + run_attention(layer, context, result, out=out) + + +def test_batch_builder_borrows_runtime_inputs_and_binding_checks_owner(): + from lightllm.common.basemodel.attention import build_attention_batch + from lightllm.common.basemodel.infer_struct import InferStateInfo + + state = InferStateInfo() + batch = make_batch() + for name, value in vars(batch).items(): + if hasattr(state, name): + setattr(state, name, value) + state.input_ids = torch.zeros(batch.token_num, dtype=torch.long) + state.is_prefill = False + state.b_image_token_end = torch.tensor([1, 2]) + state.req_manager = SimpleNamespace( + req_to_token_indexs=batch.req_to_token_indexs, HOLD_REQUEST_ID=batch.hold_req_id + ) + borrowed = build_attention_batch(state) + assert borrowed.b_req_idx is state.b_req_idx + assert borrowed.b1_cu_q_seq_len is state.b1_cu_q_seq_len + assert borrowed.req_to_token_indexs is state.req_manager.req_to_token_indexs + assert borrowed.b_image_token_end is state.b_image_token_end + backend = _OutputBackend(make_env(), make_spec()) + context = make_context(backend, borrowed, prepare=True) + state.attention.contexts = {0: context} + layer = make_layer_attention(backend, context, params=None) + assert layer.context(state.attention) is context + output = torch.ones(2, 32) + assert layer.forward(state.attention, output) is output + with pytest.raises(ValueError, match="different attention backend|specifications differ"): + replace(layer, backend=_OutputBackend(make_env(), make_spec())).forward(state.attention, output) + with pytest.raises(ValueError, match="different attention backend|specifications differ"): + replace(layer, backend=StubBackend(make_env(), make_spec())).context(state.attention) + with pytest.raises(AttributeError): + context.b_requ_idx = borrowed.b_req_idx + + +@pytest.mark.parametrize("family", list(AttentionFamily)) +def test_probe_reports_actual_coverage(monkeypatch, family): + from lightllm.common.attention import probe, registry + from lightllm.utils import backend_validator as validator + + spec = make_spec(family) + calls = [] + monkeypatch.setattr(registry, "load_backend_class", lambda name: SimpleNamespace(provider="triton")) + monkeypatch.setattr(probe, "probe_dense", lambda *args: calls.append("backend") or (True, None)) + monkeypatch.setattr(validator, "_smoke_provider", lambda *a, **k: calls.append("smoke") or (True, None)) + result = validator._probe_in_process("selected", device="cpu", spec=spec) + assert result.ok + numerical = family is AttentionFamily.DENSE + assert result.scope is (validator.ProbeScope.BACKEND_NUMERICS if numerical else validator.ProbeScope.PROVIDER_SMOKE) + assert calls == (["backend"] if numerical else ["smoke"]) + + +def test_special_dense_probe_is_explicitly_smoke_only(monkeypatch): + from dataclasses import replace + + from lightllm.common.attention import probe, registry + from lightllm.utils import backend_validator as validator + + monkeypatch.setattr(registry, "load_backend_class", lambda name: SimpleNamespace(provider="triton")) + monkeypatch.setattr(probe, "probe_dense", lambda *args: pytest.fail("unsupported numerical probe")) + monkeypatch.setattr(validator, "_smoke_provider", lambda *a, **k: (False, "provider failed")) + result = validator._probe_in_process("selected", device="cpu", spec=make_spec(use_mtp=True)) + assert not result.ok + assert result.scope is validator.ProbeScope.PROVIDER_SMOKE + assert result.reason == "provider failed" + + +def test_replay_flashqla_uses_inherited_initialization(monkeypatch): + import sys + from types import ModuleType + + from lightllm.common.attention.backends.gdn.replay.backend import ( + FlashQlaGDNReplayPrefillBackend, + TritonGDNReplayDecodeBackend, + ) + + calls = [] + original = phase_class(TritonGDNReplayDecodeBackend, AttentionPhase.PREFILL).__init__ + + def initialize(self, model, spec): + calls.append("parent") + original(self, model, spec) + + monkeypatch.setattr(phase_class(TritonGDNReplayDecodeBackend, AttentionPhase.PREFILL), "__init__", initialize) + monkeypatch.setattr( + phase_class(TritonGDNReplayDecodeBackend, AttentionPhase.PREFILL), + "_init_linear_layer_metadata", + lambda *a: calls.append("metadata"), + ) + module = ModuleType("flash_qla") + module.chunk_gated_delta_rule = object() + monkeypatch.setitem(sys.modules, "flash_qla", module) + model = model_stub() + backend = FlashQlaGDNReplayPrefillBackend(make_env(model), make_spec(AttentionFamily.GDN)) + assert calls == ["parent", "metadata"] + assert backend.runtime is model.device_runtime + assert backend.prefill_kernel is module.chunk_gated_delta_rule + + +def test_spec_configures_backend_before_planning(): + from lightllm.common.attention.api import AttentionContext, AttentionWorkspace + + model = model_stub(page_size=1, mtp_dynamic_verify=False) + spec = make_spec(cache=replace(make_spec().cache, page_size=16), causal=False, dynamic_verify=True) + backend = StubBackend(make_env(model), spec) + assert backend.spec is spec + assert (backend.spec.cache.page_size, backend.spec.causal, backend.spec.dynamic_verify) == (16, False, True) + context = make_context(backend) + assert isinstance(context, AttentionContext) + assert isinstance(context.workspace, AttentionWorkspace) + assert context.backend_identity is backend + + +@pytest.mark.parametrize("field,value", [("phase", AttentionPhase.PREFILL), ("family", AttentionFamily.MLA)]) +def test_selection_rejects_mixed_specs(field, value): + with pytest.raises(ValueError, match="share family, phase and cache layout"): + selection.select_attention_backend( + specs=(make_spec(), replace(make_spec(), **{field: value})), + caps=model_stub().device_runtime.caps, + execution=model_stub().attention_requirements, + ) + + +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("phase", list(AttentionPhase)) +def test_gdn_output_contract(monkeypatch, replay, phase): + from lightllm.common.attention.api import GDNInputs, GDNOutput + from lightllm.common.attention.backends.gdn.replay.backend import TritonGDNReplayDecodeBackend + from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNDecodeBackend + from unit_tests.common.attention.helpers import make_backend + + backend = make_backend( + phase_class(TritonGDNReplayDecodeBackend if replay else TritonGDNDecodeBackend, phase), mtp_step=0 + ) + core, gate = torch.randn(2, 4, 8), torch.randn(2, 4, 8) + monkeypatch.setattr(backend, "_split_qkvzba", lambda x: (x, gate, None, None)) + if phase is AttentionPhase.PREFILL: + monkeypatch.setattr(backend, "_prefill" if replay else "_run_prefill", lambda *a: core) + else: + monkeypatch.setattr(backend, "_decode" if replay else "_run_decode", lambda *a: (core, gate)) + context = make_context(backend, make_batch(phase=phase)) + state = SimpleNamespace(finish_decode_token=lambda *a: None) + layer = make_layer_attention(backend, context, SimpleNamespace(conv=None, ssm=None, replay_state=state)) + out = torch.empty_like(core) + result = run_attention(layer, context, GDNInputs(core), out=out) + assert isinstance(result, GDNOutput) + assert result.hidden_states is out + assert result.gate is gate + torch.testing.assert_close(result.hidden_states, core) + + +def test_binding_resolves_context_on_each_forward(monkeypatch): + from lightllm.common.attention.session import AttentionSession + + backend = _OutputBackend(make_env(), make_spec()) + first, second = make_context(backend), make_context(backend) + layer = make_layer_attention(backend, first, params=None) + session = AttentionSession() + forwarded = [] + + def forward(layer, context, inputs, *, out=None): + forwarded.append(context) + return inputs + + monkeypatch.setattr(backend, "forward", forward) + result, out = torch.ones(2, 32), torch.empty(2, 32) + for context in (first, second): + session.contexts[0] = context + assert layer.forward(session, result, out=out) is out + assert all(a is b for a, b in zip(forwarded, (first, second))) + torch.testing.assert_close(out, result) + + +def test_session_enforces_direct_backend_contract(): + from lightllm.common.attention.api import AttentionBackend, AttentionContext, AttentionWorkspace, ContextCapacity + from lightllm.common.attention.session import AttentionSession + + class Backend: + phase = AttentionPhase.DECODE + spec = make_spec() + + @classmethod + def check_environment(cls, caps): + pass + + @classmethod + def check_support(cls, spec, caps=None, constraints=None): + pass + + def allocate_context(self, capacity, *, capture, workspace): + return AttentionContext(capacity, capture, workspace) + + def plan_host(self, context): + pass + + def forward(self, layer, context, inputs, *, out=None): + return inputs + + backend, session = Backend(), AttentionSession() + batch = make_batch() + capacity = ContextCapacity(phase=batch.phase, batch_size=2, token_num=2, max_kv_seq_len=16, max_q_seq_len=1) + context = session.allocate(backend, capacity, capture=True, workspace=AttentionWorkspace("cpu")) + session.plan(backend, context, batch) + session.contexts[0] = context + layer = make_layer_attention(StubBackend(make_env(), make_spec()), context, None) + layer = replace(layer, backend=backend) + out = torch.zeros(2, 4) + assert layer.forward(session, torch.ones(2, 4), out=out) is out + assert out.eq(1).all() + with pytest.raises(ValueError, match="different attention backend"): + replace(layer, backend=Backend()).forward(session, out) + context.capacity = replace(context.capacity, phase=AttentionPhase.PREFILL) + with pytest.raises(ValueError, match="different attention backend"): + layer.forward(session, out) + + +def test_replay_copies_shared_batch_once(monkeypatch): + from lightllm.common.attention.api import AttentionBatch + from lightllm.common.attention.session import AttentionSession + + first, second = (_OutputBackend(make_env(), make_spec()) for _ in range(2)) + captured, source = AttentionSession(), AttentionSession() + captured.groups = source.groups = ({AttentionPhase.DECODE: first}, {AttentionPhase.DECODE: second}) + captured.batch = make_batch() + source.batch = make_batch((9, 7), req_to_token_indexs=captured.batch.req_to_token_indexs) + for i, backend in enumerate((first, second)): + captured.contexts[i] = make_context(backend, captured.batch, capture=True, max_length=16, prepare=True) + copies = [] + original = AttentionBatch.input_copies + + def copy(batch, incoming): + copies.append(batch) + return original(batch, incoming) + + monkeypatch.setattr(AttentionBatch, "input_copies", copy) + for target, incoming in captured.replay_input_copies(source): + target.copy_(incoming) + captured.plan_replay(source) + assert len(copies) == 1 + for context in captured.contexts.values(): + assert context.batch.b_seq_len.tolist() == [9, 7] + assert context.current_max_kv_seq_len == 9 + assert context.batch.max_kv_seq_len == 16 + + +@pytest.mark.parametrize("failure", ["allocate", "plan", "forward"]) +def test_session_releases_workspace_on_error(failure): + from lightllm.common.attention.api import AttentionWorkspace, ContextCapacity + from lightllm.common.attention.session import AttentionSession + from lightllm.platform import DeviceRuntime + + events = [] + + class Workspace(AttentionWorkspace): + def begin(self): + events.append("begin") + + def finish(self, dependencies=()): + events.append("finish") + + class Backend(StubBackend): + def allocate_context(self, *args, **kwargs): + if failure == "allocate": + raise RuntimeError("allocate failed") + return super().allocate_context(*args, **kwargs) + + def plan_host(self, context): + if failure == "plan": + raise RuntimeError("plan failed") + + session = AttentionSession() + batch = make_batch() + backend = Backend(make_env(), make_spec()) + capacity = ContextCapacity(phase=batch.phase, batch_size=2, token_num=2, max_kv_seq_len=16, max_q_seq_len=1) + with pytest.raises(RuntimeError, match=failure): + with session.prepare( + batch, ({batch.phase: backend},), capacity, workspace=lambda _: Workspace(runtime=DeviceRuntime("cpu")) + ): + raise RuntimeError("forward failed") + assert events == ["begin", "finish"] + + +def test_model_capture_session_uses_graph_bound_and_one_batch(): + from lightllm.common.attention.api import AttentionWorkspace + from lightllm.common.attention.session import AttentionSession + from lightllm.common.basemodel.attention import ModelAttention + from lightllm.platform import DeviceRuntime + + events = [] + + class Workspace(AttentionWorkspace): + def __init__(self, group_id): + super().__init__(runtime=DeviceRuntime("cpu")) + self.group_id = group_id + + def begin(self): + events.append(("begin", self.group_id)) + + def finish(self, dependencies=()): + events.append(("finish", self.group_id)) + + attention = ModelAttention.__new__(ModelAttention) + attention.generation = 3 + attention.constraints = SimpleNamespace(max_seq_len=128) + attention.groups = tuple({AttentionPhase.DECODE: StubBackend(make_env(), make_spec())} for _ in range(2)) + attention.workspaces = {(i, 0): Workspace(i) for i in range(2)} + batch = make_batch((9, 7)) + state = SimpleNamespace( + **vars(batch), + is_prefill=False, + input_ids=torch.zeros(batch.token_num, dtype=torch.int64), + req_manager=SimpleNamespace(req_to_token_indexs=batch.req_to_token_indexs, HOLD_REQUEST_ID=batch.hold_req_id), + att_state_manager=SimpleNamespace(storage_generation=3), + attention=AttentionSession(), + microbatch_index=0, + is_cuda_graph_capture=True, + is_cuda_graph_replay=False, + ) + + with pytest.raises(RuntimeError, match="forward failed"): + with attention.session(state, max_kv_seq_len=16) as session: + assert session.batch.max_kv_seq_len == 16 + assert session.batch.total_token_num == 32 + for context in session.contexts.values(): + assert context.batch is session.batch + assert context.capture + assert context.capacity.max_kv_seq_len == 16 + assert context.current_max_kv_seq_len == 9 + assert len(session.contexts) == 2 + assert state.max_kv_seq_len == 9 + raise RuntimeError("forward failed") + + assert events == [("begin", 0), ("begin", 1), ("finish", 1), ("finish", 0)] diff --git a/unit_tests/common/attention/test_dense_numerics.py b/unit_tests/common/attention/test_dense_numerics.py new file mode 100644 index 0000000000..6beb1f8526 --- /dev/null +++ b/unit_tests/common/attention/test_dense_numerics.py @@ -0,0 +1,296 @@ +from unit_tests.common.attention.helpers import plan, build, run_attention +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention.api import AttentionPhase, CacheLayout, ContextCapacity, UnsupportedAttention +from lightllm.common.attention.api.inputs import DenseInputs, DenseResources +from lightllm.common.attention.api.workspace import AttentionWorkspace +from unit_tests.common.attention.helpers import dense_class, make_env +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention.session import AttentionSession +from lightllm.common.basemodel.graph import CapturedGraph +from lightllm.platform import get_device_runtime +from unit_tests.common.attention.helpers import make_batch, make_layer_attention, make_spec, model_stub + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference(query, key, value, lengths, q_lengths, window=None): + rows, start = [], 0 + groups = query.shape[1] // key.shape[1] + for request, (length, q_length) in enumerate(zip(lengths, q_lengths)): + k = key[request * 16 : request * 16 + length].float().repeat_interleave(groups, dim=1) + v = value[request * 16 : request * 16 + length].float().repeat_interleave(groups, dim=1) + q = query[start : start + q_length].float() + logits = torch.einsum("qhd,khd->hqk", q, k) * query.shape[-1] ** -0.5 + causal = torch.arange(length, device=q.device)[None, :] <= ( + length - q_length + torch.arange(q_length, device=q.device)[:, None] + ) + if window is not None: + causal &= torch.arange(length, device=q.device)[None, :] >= ( + length - q_length + torch.arange(q_length, device=q.device)[:, None] - window[0] + ) + probs = logits.masked_fill(~causal, -torch.inf).softmax(-1) + rows.append(torch.einsum("hqk,khd->qhd", probs, v)) + start += q_length + return torch.cat(rows).to(query.dtype) + + +@pytest.mark.parametrize("provider", ["triton", "fa3", "flashinfer"]) +@pytest.mark.parametrize("phase", list(AttentionPhase)) +def test_startup_probe_runs_the_selected_dense_implementation(provider, phase): + from lightllm.common.attention.probe import probe_dense + + cls = dense_class(provider, phase) + runtime = get_device_runtime() + try: + cls.check_environment(runtime.caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + spec = make_spec( + phase=phase, head_size=64, value_head_size=64, cache=CacheLayout("None", 16, torch.bfloat16, 2, 64, 64) + ) + success, error = probe_dense(cls, spec, runtime) + assert success, error + + +@pytest.mark.parametrize( + "provider,kv_heads,window", + [ + ("triton", 2, None), + ("fa3", 2, None), + ("flashinfer", 2, None), + ("triton", 4, None), + ("triton", 4, (3, 0)), + ("fa3", 4, (3, 0)), + ], +) +@pytest.mark.parametrize("phase", list(AttentionPhase)) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_dense_provider_reference_and_decode_graph(provider, phase, head_dim, kv_heads, window): + cls = dense_class(provider, phase) + try: + cls.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + torch.manual_seed(616) + model = model_stub(page_size=4) + model.config = dict( + num_attention_heads=4, num_key_value_heads=kv_heads, hidden_size=4 * head_dim, head_dim=head_dim + ) + model.max_seq_length, model.data_type = 16, torch.bfloat16 + spec = make_spec( + phase=phase, + num_q_heads=4, + head_size=head_dim, + value_head_size=head_dim, + cache=CacheLayout("None", 4, torch.bfloat16, kv_heads, head_dim, head_dim), + sliding_window=window, + ) + backend = cls(make_env(model), spec) + key, value = (torch.randn((32, kv_heads, head_dim), dtype=torch.bfloat16, device="cuda") for _ in range(2)) + q_lengths = (2, 3) if phase is AttentionPhase.PREFILL else (1, 1) + query = torch.randn((sum(q_lengths), 4, head_dim), dtype=torch.bfloat16, device="cuda") + table = torch.arange(32, device="cuda", dtype=torch.int32).view(2, 16) + + def batch_for(lengths): + offsets = [0, q_lengths[0], sum(q_lengths)] + return make_batch( + lengths, + phase=phase, + device="cuda", + token_num=sum(q_lengths), + max_q_seq_len=max(q_lengths), + req_to_token_indexs=table, + hold_req_id=1, + b_q_seq_len=torch.tensor(q_lengths, dtype=torch.int32, device="cuda"), + b1_cu_q_seq_len=torch.tensor(offsets, dtype=torch.int32, device="cuda"), + b_ready_cache_len=torch.tensor( + [n - q for n, q in zip(lengths, q_lengths)], dtype=torch.int32, device="cuda" + ), + ) + + batch = batch_for((7, 13)) + capture = phase is AttentionPhase.DECODE + pool = AttentionWorkspace(runtime=model.device_runtime) + capacity = ContextCapacity( + phase=phase, + batch_size=batch.batch_size, + token_num=batch.token_num, + max_kv_seq_len=16, + max_q_seq_len=batch.max_q_seq_len, + ) + context = AttentionSession.allocate( + backend, capacity, capture=capture, workspace=pool.borrow_for_capture() if capture else pool + ) + plan(backend, context, batch) + attention = make_layer_attention(backend, context, DenseResources(key=key, value=value)) + inputs = DenseInputs(query) + output = torch.empty_like(query) + + def run(): + build(backend, context) + return run_attention(attention, context, inputs, out=output) + + actual = run() + assert actual.data_ptr() == output.data_ptr() + torch.testing.assert_close(actual, _reference(query, key, value, (7, 13), q_lengths, window), rtol=0.02, atol=0.015) + if not capture: + return + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + captured_pointer = context.batch.b_seq_len.data_ptr() + manager = AttentionStateManager(request_capacity=2) + + state = SimpleNamespace( + att_state_manager=manager, + attention_groups=({phase: backend},), + attention_contexts={0: context}, + graph_input_copies=lambda source: state.attention.replay_input_copies(source.attention), + ) + state.attention = AttentionSession() + state.attention.groups = state.attention_groups + state.attention.contexts = state.attention_contexts + state.attention.batch = context.batch + execution = CapturedGraph(graph, (state,), (output,), runtime=model.device_runtime) + try: + # Eager execution may grow the shared pool after capture. The graph must + # keep its original pointers and fence reuse across streams. + captured_buffers = {key: tensor.data_ptr() for key, tensor in context.workspace._buffers.items()} + for (name, dtype), tensor in tuple(pool._buffers.items()): + pool.scratch(name, tensor.numel() + 1, dtype) + for i, lengths in enumerate(((16, 3), (2, 15), (9, 8))): + with torch.cuda.stream(stream if i % 2 else torch.cuda.current_stream()): + query.normal_() + source = SimpleNamespace( + attention=SimpleNamespace(groups=state.attention.groups, batch=batch_for(lengths)) + ) + actual = execution.replay(source) + assert actual is output + assert context.batch.b_seq_len.data_ptr() == captured_pointer + assert { + key: tensor.data_ptr() for key, tensor in context.workspace._buffers.items() + } == captured_buffers + torch.testing.assert_close( + output, _reference(query, key, value, lengths, q_lengths, window), rtol=0.02, atol=0.015 + ) + finally: + manager.close_graphs() + assert execution.graph is None + + +@pytest.mark.parametrize("provider", ["fa3", "flashinfer"]) +def test_noncausal_prefill_probe(provider): + from lightllm.common.attention.probe import probe_dense + + cls = dense_class(provider, AttentionPhase.PREFILL) + runtime = get_device_runtime() + cls.check_environment(runtime.caps) + spec = make_spec( + phase=AttentionPhase.PREFILL, + causal=False, + head_size=64, + value_head_size=64, + cache=CacheLayout("None", 16, torch.bfloat16, 2, 64, 64), + ) + cls.check_support(spec, runtime.caps) + success, error = probe_dense(cls, spec, runtime) + assert success, error + + +def test_triton_rejects_noncausal_prefill(): + cls = dense_class("triton", AttentionPhase.PREFILL) + with pytest.raises(UnsupportedAttention, match="noncausal"): + cls.check_support(make_spec(phase=AttentionPhase.PREFILL, causal=False)) + + +@pytest.mark.parametrize("phase", list(AttentionPhase)) +@pytest.mark.parametrize("scale", [0.07, 0.13]) +@pytest.mark.parametrize("provider", ["flashinfer", "fa3"]) +def test_mla_spec_scale(phase, scale, provider): + from lightllm.common.attention.api import AttentionFamily, LayerAttention + from lightllm.common.attention.api.inputs import MLADecodeInputs, MLAResources, MLAPrefillInputs + from lightllm.common.attention.registry import load_backend_class + + torch.manual_seed(27) + model = model_stub(page_size=1) + model.config = {"num_attention_heads": 8, "rope_scaling": {"factor": 8, "mscale_all_dim": 1}} + model.qk_nope_head_dim, model.qk_rope_head_dim = 128, 64 + model.kv_lora_rank, model.v_head_dim = 512, 128 + model.max_seq_length, model.data_type = 32, torch.bfloat16 + spec = make_spec( + family=AttentionFamily.MLA, + phase=phase, + num_q_heads=8, + head_size=192, + value_head_size=512, + cache=CacheLayout("None", 4, torch.bfloat16, 1, 192, 512), + softmax_scale=scale, + ) + prefix = {"flashinfer": "FlashInfer", "fa3": "Fa3"}[provider] + cls = load_backend_class(f"backends.mla.{provider}.backend:{prefix}MLA{phase.value.title()}Backend") + backend = cls(make_env(model), spec) + assert backend.spec.cache.page_size == 4 + if provider == "flashinfer": + assert backend.max_page_num == 8 + lengths = (7, 13) + q_lengths = (2, 3) if phase is AttentionPhase.PREFILL else (1, 1) + batch = make_batch( + lengths, + phase=phase, + device="cuda", + token_num=sum(q_lengths), + max_q_seq_len=max(q_lengths), + b_q_seq_len=torch.tensor(q_lengths, dtype=torch.int32, device="cuda"), + b1_cu_q_seq_len=torch.tensor([0, q_lengths[0], sum(q_lengths)], dtype=torch.int32, device="cuda"), + req_to_token_indexs=torch.arange(32, dtype=torch.int32, device="cuda").view(2, 16), + b_ready_cache_len=torch.tensor([n - q for n, q in zip(lengths, q_lengths)], dtype=torch.int32, device="cuda"), + ) + workspace = AttentionWorkspace(runtime=model.device_runtime) + capacity = ContextCapacity( + phase=phase, batch_size=2, token_num=sum(q_lengths), max_kv_seq_len=16, max_q_seq_len=max(q_lengths) + ) + context = AttentionSession.allocate(backend, capacity, capture=False, workspace=workspace) + plan(backend, context, batch) + build(backend, context) + if phase is AttentionPhase.PREFILL: + query = torch.randn(sum(q_lengths), 8, 192, device="cuda", dtype=torch.bfloat16) + key_nope = torch.randn(sum(lengths), 8, 128, device="cuda", dtype=torch.bfloat16) + key_rope = torch.randn(sum(lengths), 1, 64, device="cuda", dtype=torch.bfloat16) + value = torch.randn(sum(lengths), 8, 128, device="cuda", dtype=torch.bfloat16) + key = torch.cat((key_nope, key_rope.expand(-1, 8, -1)), dim=-1) + inputs = MLAPrefillInputs(query, key_nope, key_rope, value) + cache = None + else: + cache = torch.randn(32, 1, 576, device="cuda", dtype=torch.bfloat16) + query = torch.randn(2, 8, 576, device="cuda", dtype=torch.bfloat16) + inputs = MLADecodeInputs(query[..., :512], query[..., 512:]) + key = torch.cat([cache[i * 16 : i * 16 + length].expand(-1, 8, -1) for i, length in enumerate(lengths)]) + value = key[..., :512] + binding = LayerAttention(0, 0, MLAResources(cache, 512, 64), backend) + actual = run_attention(binding, context, inputs) + expected, q_start, k_start = [], 0, 0 + for length, q_length in zip(lengths, q_lengths): + q, k, v = ( + query[q_start : q_start + q_length].float(), + key[k_start : k_start + length].float(), + value[k_start : k_start + length].float(), + ) + logits = torch.einsum("qhd,khd->hqk", q, k) * scale + mask = ( + torch.arange(length, device="cuda")[None, :] + <= length - q_length + torch.arange(q_length, device="cuda")[:, None] + ) + expected.append(torch.einsum("hqk,khd->qhd", logits.masked_fill(~mask, -torch.inf).softmax(-1), v)) + q_start += q_length + k_start += length + torch.testing.assert_close(actual, torch.cat(expected).to(actual.dtype), atol=0.03, rtol=0.03) diff --git a/unit_tests/common/attention/test_linear_layout.py b/unit_tests/common/attention/test_linear_layout.py new file mode 100644 index 0000000000..5e65ccd6e0 --- /dev/null +++ b/unit_tests/common/attention/test_linear_layout.py @@ -0,0 +1,73 @@ +from unit_tests.common.attention.helpers import plan, build, run_attention +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention.backends._shared.linear.layout import LinearDecodeContext, update_linear_decode +from unit_tests.common.attention.helpers import StubBackend, make_env +from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub + + +class _LinearBackend(StubBackend): + def allocate_context(self, capacity, *, capture, workspace): + return LinearDecodeContext(capacity, capture, workspace) + + def build_device_metadata(self, context): + update_linear_decode(context, 2, self.spec.dynamic_verify, replay=self.replay) + + def forward(self, layer, context, inputs, *, out=None): + return context.accepted_tokens(layer.resources) + + +@pytest.mark.parametrize("dynamic", [False, True]) +@pytest.mark.parametrize("replay", [False, True]) +def test_shared_backend_selects_each_states_cursor_through_graph_replay(dynamic, replay): + backend = _LinearBackend(make_env(), make_spec(dynamic_verify=dynamic)) + backend.replay = replay + requests = [0, 0, 2, 7, 7, 7] if dynamic else [0, 0, 0, 2, 2, 2] + positions = [0, 1, 0, 0, 0, 0] if dynamic else [0, 1, 2, 0, 1, 2] + batch = make_batch( + (5,) * 6, + device="cuda", + b_req_idx=torch.tensor(requests, dtype=torch.int32, device="cuda"), + b_mtp_index=torch.tensor(positions, dtype=torch.int32, device="cuda"), + ) + context = make_context(backend, batch, capture=True, prepare=True) + first = torch.tensor([0, 0, 2, 0, 0, 0, 0, 99], dtype=torch.int32, device="cuda") + second = torch.tensor([2, 0, 1, 0, 0, 0, 0, 99], dtype=torch.int32, device="cuda") + layer_a = make_layer_attention(backend, context, params=first) + layer_b = replace(layer_a, logical_layer_id=1, resources=second) + + def step(): + build(backend, context) + a, b = run_attention(layer_a, context, None), run_attention(layer_b, context, None) + assert run_attention(layer_a, context, None) is a + return a, b + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + step() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + a, b = step() + graph.replay() + tail = [1] * 4 if dynamic else [] + assert a.tolist() == [1, 3] + tail + assert b.tolist() == [3, 2] + tail + first[0], first[2] = 1, 0 + second[0], second[2] = 0, 2 + graph.replay() + assert a.tolist() == [2, 1] + tail + assert b.tolist() == [1, 3] + tail + # A new eager batch must not reuse the previous batch's gathered counts. + plan(backend, context, batch) + a, b = step() + assert a.tolist() == [2, 1] + tail + assert b.tolist() == [1, 3] + tail + expected_requests = torch.tensor([0, 2] + ([7] * 4 if dynamic else []), device="cuda") + torch.testing.assert_close(context.b_conv_buffer_idx.long(), expected_requests) + expected_ssm = expected_requests if replay else expected_requests[:, None] * 3 + torch.arange(3, device="cuda") + torch.testing.assert_close(context.b_ssm_buffer_idx.long(), expected_ssm) diff --git a/unit_tests/common/attention/test_platform_runtime.py b/unit_tests/common/attention/test_platform_runtime.py new file mode 100644 index 0000000000..63ab3bc654 --- /dev/null +++ b/unit_tests/common/attention/test_platform_runtime.py @@ -0,0 +1,264 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention.api import AttentionPhase, CaptureHooks +from lightllm.common.attention.api.workspace import AttentionWorkspace +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention.session import AttentionSession +from lightllm.common.basemodel.graph import CapturedGraph +from lightllm.platform import DeviceRuntime, get_device_runtime + + +def test_runtime_resolves_host_platform_without_querying_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: pytest.fail("queried CUDA for a CPU runtime")) + assert get_device_runtime(platform="cpu").device == torch.device("cpu") + with pytest.raises(ValueError, match="cannot operate"): + get_device_runtime("cuda:0", platform="cpu") + with pytest.raises(ValueError, match="cannot operate"): + get_device_runtime("cpu", platform="cuda") + + +def test_runtime_rejects_unimplemented_hardware(): + with pytest.raises(ValueError, match="CUDA and CPU"): + get_device_runtime(platform="musa") + with pytest.raises(ValueError, match="CUDA and CPU"): + get_device_runtime("npu:0") + + +def test_explicit_current_device_is_not_overridden_by_worker_default(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 1) + assert get_device_runtime("cuda").device == torch.device("cuda:1") + assert get_device_runtime("cuda:2").device == torch.device("cuda:2") + assert get_device_runtime().device == torch.device("cuda:0") + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_gqa_storage_allocation_uses_explicit_device_context(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + config = GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=2, head_dim=8, dtype=torch.float32) + with torch.cuda.device(1): + state = GQAFullAttentionState(config, device="cuda", page_size=4) + state.allocate_storage(8) + assert state.k_storage.device == state.v_storage.device == torch.device("cuda:1") + + +def test_captured_workspace_keeps_address_when_pool_grows(): + pool = AttentionWorkspace("cpu") + first, second = pool.borrow_for_capture(), pool.borrow_for_capture() + old = first.scratch("scratch", 8, torch.int32) + assert second.scratch("scratch", 8, torch.int32).data_ptr() == old.data_ptr() + first.freeze() + grown = pool.scratch("scratch", 64, torch.int32) + old.fill_(17) + grown.fill_(29) + assert first.scratch("scratch", 8, torch.int32).tolist() == [17] * 8 + with pytest.raises(RuntimeError, match="cannot grow"): + first.scratch("scratch", 64, torch.int32) + assert AttentionWorkspace("cpu").scratch("scratch", 8, torch.int32).data_ptr() != old.data_ptr() + + +def test_graph_orders_host_reuse_submission_update_and_destruction(): + events = [] + pending = {"submitted": False, "updated": False} + + class Runtime(DeviceRuntime): + def synchronize(self): + assert not pending["submitted"] or pending["updated"], "waited for a graph before submitting its updates" + events.append("synchronize") + + class Binding(CaptureHooks): + def wait_for_update(self): + events.append("host_ready") + + def after_submit(self): + assert pending["submitted"] + pending["updated"] = True + events.append("update") + return () + + def close(self): + events.append("binding_close") + + def submit(): + pending["submitted"] = True + events.append("submit") + + runtime = Runtime("cpu") + manager = AttentionStateManager(request_capacity=2) + backend = SimpleNamespace(create_capture_hooks=lambda context: Binding()) + state = SimpleNamespace( + input_ids=torch.zeros(1), + att_state_manager=manager, + attention_groups=({AttentionPhase.DECODE: backend},), + attention_contexts={ + 0: SimpleNamespace( + capacity=SimpleNamespace(phase=AttentionPhase.DECODE), workspace=AttentionWorkspace(runtime=runtime) + ) + }, + graph_input_copies=lambda source: [(SimpleNamespace(copy_=lambda *a, **kw: events.append("copy")), None)], + ) + graph = SimpleNamespace(replay=submit, reset=lambda: events.append("graph_reset")) + state.attention = AttentionSession() + state.attention.groups = state.attention_groups + state.attention.contexts = state.attention_contexts + state.attention.plan_replay = lambda source: events.append("plan") + execution = CapturedGraph(graph, (state,), ("result",), runtime=runtime) + assert execution.replay(state) == "result" + assert events == ["host_ready", "copy", "plan", "submit", "update"] + manager.close_graphs() + execution.close() + assert events[-4:] == ["host_ready", "synchronize", "binding_close", "graph_reset"] + assert execution.infer_states == () + with pytest.raises(RuntimeError, match="closed"): + execution.replay(state) + + +def test_storage_generation_invalidates_graph_before_copy(): + manager = AttentionStateManager(request_capacity=2) + state = SimpleNamespace( + input_ids=torch.zeros(1), + att_state_manager=manager, + attention_contexts={}, + attention_groups=(), + graph_input_copies=lambda source: pytest.fail("copied into a stale graph"), + ) + graph = SimpleNamespace(replay=lambda: pytest.fail("replayed a stale graph"), reset=lambda: None) + state.attention = AttentionSession() + state.attention.groups = state.attention_groups + state.attention.contexts = state.attention_contexts + execution = CapturedGraph(graph, (state,), runtime=DeviceRuntime("cpu")) + manager.storage_generation += 1 + with pytest.raises(RuntimeError, match="storage changed"): + execution.replay(state) + execution.close() + + +@pytest.mark.parametrize("invalid", ["role", "shape"]) +def test_graph_validates_all_lanes_before_copy_or_plan(monkeypatch, invalid): + from lightllm.common.attention.api import BackendEnv, ContextCapacity + from lightllm.common.basemodel.infer_struct import InferStateInfo + from unit_tests.common.attention.helpers import StubBackend, make_batch, make_spec + + runtime = DeviceRuntime("cpu") + manager = AttentionStateManager(request_capacity=2) + env = BackendEnv( + runtime=runtime, + network_config={}, + hold_req_id=7, + draft_step=0, + max_seq_len=16, + tp_size=1, + alloc_tensor=torch.empty, + ) + backend = StubBackend(env, make_spec()) + groups = ({AttentionPhase.DECODE: backend},) + plans = [] + monkeypatch.setattr(backend, "plan_host", lambda context: plans.append(context)) + capacity = ContextCapacity( + phase=AttentionPhase.DECODE, batch_size=2, token_num=2, max_kv_seq_len=16, max_q_seq_len=1 + ) + + def make_state(batch, capture): + state = InferStateInfo() + state.att_state_role = "main" + state.att_state_manager = manager + state.input_ids = torch.zeros(2, dtype=torch.int64) if capture else torch.ones(2, dtype=torch.int64) + with state.attention.prepare( + batch, + groups, + capacity, + workspace=lambda group_id: AttentionWorkspace(runtime=runtime), + capture=capture, + replay=not capture, + ): + pass + return state + + targets = [make_state(make_batch(), True) for _ in range(2)] + sources = [ + make_state(make_batch((9, 7), req_to_token_indexs=target.attention.batch.req_to_token_indexs), False) + for target in targets + ] + if invalid == "role": + sources[1].att_state_role = "draft0" + else: + sources[1].input_ids = torch.ones(3, dtype=torch.int64) + plans.clear() + graph = SimpleNamespace(replay=lambda: pytest.fail("submitted invalid lanes"), reset=lambda: None) + execution = CapturedGraph(graph, targets, runtime=runtime) + try: + with pytest.raises(ValueError, match="model role|shape, dtype or device"): + execution.replay(*sources) + assert targets[0].input_ids.tolist() == [0, 0] + assert targets[0].attention.batch.b_seq_len.tolist() == [5, 2] + assert plans == [] + finally: + execution.close() + + +def test_graph_binding_failure_releases_previously_bound_resources(): + events = [] + + class Binding(CaptureHooks): + def close(self): + events.append("binding_close") + + def create_capture_hooks(context): + if context.fail: + raise ValueError("unsupported graph task") + return Binding() + + manager = AttentionStateManager(request_capacity=2) + backend = SimpleNamespace(create_capture_hooks=create_capture_hooks) + state = SimpleNamespace( + att_state_manager=manager, + attention_groups=({AttentionPhase.DECODE: backend},) * 2, + attention_contexts={ + i: SimpleNamespace( + fail=bool(i), capacity=SimpleNamespace(phase=AttentionPhase.DECODE), workspace=AttentionWorkspace("cpu") + ) + for i in range(2) + }, + ) + graph = SimpleNamespace(reset=lambda: events.append("graph_reset")) + state.attention = AttentionSession() + state.attention.groups = state.attention_groups + state.attention.contexts = state.attention_contexts + with pytest.raises(ValueError, match="unsupported graph task"): + CapturedGraph(graph, (state,), runtime=DeviceRuntime("cpu")) + manager.close_graphs() + assert events == ["binding_close", "graph_reset"] + + +def test_state_layout_and_native_write_preserve_physical_rows(): + config = GQAAttentionConfig(logical_layer_groups=((0, 3), (7,)), global_head_num=2, head_dim=8, dtype=torch.float32) + state = GQAFullAttentionState(config, device="cpu", page_size=4) + state.require_zero_initialization() + state.allocate_storage(12) + assert not state.k_storage.count_nonzero() + values = torch.arange(2 * 4 * 8, dtype=torch.float32).view(2, 4, 8) + state.write_full_att_kv(3, torch.tensor([2, 5], dtype=torch.int32), values) + torch.testing.assert_close(state.get_layer_view(0)[0][[2, 5]], values[:, :2]) + state.copy_full_att_tokens(torch.tensor([2, 5]), torch.tensor([7, 9])) + torch.testing.assert_close(state.get_layer_view(3)[1][[7, 9]], values[:, 2:]) + assert not state.get_layer_view(7)[0].count_nonzero() + relocated = type("ExternalState", (GQAFullAttentionState,), {"__module__": "external.provider"}) + assert relocated(config, device="cpu").describe_layout() == state.describe_layout() + + +def test_runtime_registration_keeps_platform_validation_in_adapter(monkeypatch): + import lightllm.platform as platform + + monkeypatch.setattr(platform, "_runtime_factories", {}) + calls = [] + adapter = object() + platform.register_device_runtime("musa", lambda device: calls.append(device) or adapter) + assert platform.get_device_runtime("musa:2", platform="musa") is adapter + assert calls == ["musa:2"] + with pytest.raises(ValueError, match="already registered"): + platform.register_device_runtime("musa", lambda device: adapter) diff --git a/unit_tests/common/attention/test_provider_kind.py b/unit_tests/common/attention/test_provider_kind.py new file mode 100644 index 0000000000..6fa7d2ea67 --- /dev/null +++ b/unit_tests/common/attention/test_provider_kind.py @@ -0,0 +1,242 @@ +import pytest +import torch + +from lightllm.common.attention.api import AttentionFamily, UnsupportedAttention +from lightllm.common.attention.registry import load_backend_class +from lightllm.platform import DeviceCaps, Platform +from unit_tests.common.attention.helpers import make_spec + + +@pytest.mark.parametrize( + "reference,family", + [ + ("backends.dense.fa3.backend:Fa3DenseDecodeBackend", AttentionFamily.DENSE), + ("backends.mla.fa3.backend:Fa3MLADecodeBackend", AttentionFamily.MLA), + ("backends.sparse_mla.flashmla.backend:FlashMLASparseMLADecodeBackend", AttentionFamily.SPARSE_MLA), + ("backends.sparse_mla.flashmla.fp8:FlashMLASparseMLAFp8DecodeBackend", AttentionFamily.SPARSE_MLA), + ("backends.gdn.triton.backend:TritonGDNDecodeBackend", AttentionFamily.GDN), + ("backends.gdn.flashqla.backend:FlashQlaGDNPrefillBackend", AttentionFamily.GDN), + ("backends.kda.triton.backend:TritonKDADecodeBackend", AttentionFamily.KDA), + ], +) +def test_family_is_explicit_and_survives_adapter_relocation(reference, family): + backend = load_backend_class(reference) + relocated = type("ExternalAdapter", (backend,), {"__module__": "external.extension"}) + assert relocated.family is family + + +@pytest.mark.parametrize( + "reference", + [ + "backends.sparse_mla.flashmla.fp8:FlashMLASparseMLAFp8DecodeBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:GlmSparseMLAFp8DecodeBackend", + "lightllm.models.deepseek_v4.attention.sparse_mla:DeepseekV4SparseMLAFp8DecodeBackend", + ], +) +def test_fp8_sparse_capabilities(monkeypatch, reference): + from dataclasses import replace + + backend = load_backend_class(reference) + caps = DeviceCaps(Platform.CUDA, torch.device("cuda:0"), capability=(8, 0)) + with pytest.raises(UnsupportedAttention, match="device capability"): + backend.check_environment(caps) + caps = replace(caps, capability=(9, 0)) + imported = [] + + def import_provider(name): + imported.append(name) + raise ImportError("missing FlashMLA") + + monkeypatch.setattr("lightllm.common.attention._environment.importlib.import_module", import_provider) + with pytest.raises(UnsupportedAttention, match="missing FlashMLA"): + backend.check_environment(caps) + assert imported == ["flash_mla"] + + spec = make_spec(AttentionFamily.SPARSE_MLA, variant=backend.variant) + spec = replace(spec, cache=replace(spec.cache, encoding="fp8kv_dsa")) + assert backend.check_support(spec) is None + with pytest.raises(UnsupportedAttention, match="activation dtype"): + backend.check_support(replace(spec, dtype=torch.float32)) + + +@pytest.mark.parametrize( + "reference", + ["backends.dense.fa3.backend:Fa3DenseDecodeBackend", "backends.dense.fa3.varlen:Fa3VarlenAttentionOp"], +) +def test_environment_rejection_precedes_provider_import(monkeypatch, reference): + from dataclasses import replace + + implementation = load_backend_class(reference) + imported = [] + + def import_provider(name): + imported.append(name) + raise OSError("provider library cannot load") + + monkeypatch.setattr("lightllm.common.attention._environment.importlib.import_module", import_provider) + caps = DeviceCaps(Platform.CPU, torch.device("cpu")) + with pytest.raises(UnsupportedAttention, match="platform cpu"): + implementation.check_environment(caps) + caps = DeviceCaps(Platform.CUDA, torch.device("cuda:0"), available=False, capability=(9, 0)) + with pytest.raises(UnsupportedAttention, match="unavailable"): + implementation.check_environment(caps) + with pytest.raises(UnsupportedAttention, match="device capability"): + implementation.check_environment(replace(caps, available=True, capability=(8, 0))) + assert imported == [] + with pytest.raises(UnsupportedAttention, match="provider library cannot load"): + implementation.check_environment(replace(caps, available=True)) + assert imported == ["sgl_kernel.flash_attn"] + + +def test_selection_imports_only_backend_definitions(): + import subprocess + import sys + + script = """ +import importlib.abc +import sys +import torch + +blocked = {"triton", "sgl_kernel", "flashinfer", "flash_mla", "flash_qla", "flash_attn_interface"} +class DefinitionsOnly(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if (fullname.split(".")[0] in blocked or ".kernels." in fullname + or fullname.startswith("lightllm.common.basemodel")): + raise AssertionError("selection imported execution dependency: " + fullname) + +sys.meta_path.insert(0, DefinitionsOnly()) +from lightllm.common.attention import selection +from lightllm.common.attention.api import ( + AttentionFamily, AttentionPhase, CacheEncoding, CacheLayout, LayerSpec, ServingConstraints +) +from lightllm.common.attention.backends.dense.triton.backend import TritonDenseDecodeBackend +from lightllm.platform import DeviceCaps, Platform +from lightllm.utils.backend_validator import ProbeResult, ProbeScope + +TritonDenseDecodeBackend.check_environment = classmethod(lambda cls, caps: None) +selection.backend_validator.probe_backend = lambda *a, **kw: ProbeResult(True, ProbeScope.BACKEND_NUMERICS) +spec = LayerSpec(AttentionFamily.DENSE, AttentionPhase.DECODE, torch.bfloat16, + CacheLayout(CacheEncoding.PLAIN, 1, torch.bfloat16, 2, 64, 64), 4, 64, 64) +assert selection.select_attention_backend(specs=(spec,), caps=DeviceCaps(Platform.CPU, torch.device("cpu")), + execution=ServingConstraints(), requested="triton") is TritonDenseDecodeBackend +assert not torch.cuda.is_initialized() +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + assert "Neo FA3 is unavailable" not in result.stderr + + +def test_gdn_decode_constructor_does_not_load_prefill(): + import subprocess + import sys + + script = """ +import importlib.abc +import sys +import torch + +class NoExecutionImports(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.split(".")[0] in {"triton", "flash_qla"} or ".kernels." in fullname: + raise AssertionError("decode construction imported execution code: " + fullname) + +sys.meta_path.insert(0, NoExecutionImports()) +from lightllm.common.attention.api import BackendEnv, AttentionFamily, AttentionPhase, CacheLayout, LayerSpec +from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNDecodeBackend +from lightllm.common.attention.backends.gdn.replay.backend import TritonGDNReplayDecodeBackend +from lightllm.platform import DeviceRuntime + +env = BackendEnv(runtime=DeviceRuntime("cpu"), network_config={ + "linear_num_value_heads": 4, "linear_num_key_heads": 2, "linear_key_head_dim": 8, + "linear_value_head_dim": 16, "linear_conv_kernel_dim": 4, "hidden_act": "silu", +}, hold_req_id=7, draft_step=0, max_seq_len=64, tp_size=1, alloc_tensor=torch.empty) +for cls in (TritonGDNDecodeBackend, TritonGDNReplayDecodeBackend): + cache = CacheLayout("plain", 1, torch.bfloat16, 2, 8, 16, + state_mode=cls.state_mode, state_dtype=torch.float32) + spec = LayerSpec(AttentionFamily.GDN, AttentionPhase.DECODE, torch.bfloat16, cache, 2, 8, 16) + backend = cls(env, spec) + assert backend.spec is spec +assert not torch.cuda.is_initialized() +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize("bits", [4, 8]) +def test_quantized_prefill_uses_layout_group(monkeypatch, bits): + from importlib import import_module + from types import SimpleNamespace + + from lightllm.common.attention.api import AttentionPhase, BackendEnv, CacheLayout + from lightllm.common.attention.api.inputs import DenseInputs, DenseResources + from lightllm.platform import DeviceRuntime + from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, run_attention + + cls = load_backend_class(f"backends.dense.triton.int{bits}:TritonDenseInt{bits}PrefillBackend") + monkeypatch.setattr( + "lightllm.utils.envs_utils.get_env_start_args", lambda: SimpleNamespace(llm_kv_quant_group_size=128) + ) + quantization = import_module(f"lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int{bits}kv_copy_kv") + prefill = import_module("lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad") + groups = [] + monkeypatch.setattr( + quantization, f"dequantize_int{bits}kv", lambda **kwargs: groups.append(kwargs["quant_group_size"]) + ) + monkeypatch.setattr(prefill, "context_attention_fwd_contiguous_kv", lambda **kwargs: kwargs["o"].zero_()) + env = BackendEnv( + runtime=DeviceRuntime("cpu"), + network_config={}, + hold_req_id=7, + draft_step=0, + max_seq_len=64, + tp_size=1, + alloc_tensor=torch.empty, + ) + for group_size in (16, 32): + spec = make_spec( + phase=AttentionPhase.PREFILL, + cache=CacheLayout(f"int{bits}kv", 1, torch.int8, 1, 64, 64, quant_group_size=group_size), + num_q_heads=1, + head_size=64, + value_head_size=64, + ) + backend = cls(env, spec) + batch = make_batch((1,), phase=AttentionPhase.PREFILL) + context = make_context(backend, batch) + context.b_kv_start_loc = torch.zeros(1, dtype=torch.int32) + cache = torch.empty(1, 1, 32 if bits == 4 else 64, dtype=torch.int8) + scale = torch.ones(1, 1, 64 // group_size) + layer = make_layer_attention(backend, context, DenseResources(cache, cache, key_scale=scale, value_scale=scale)) + output = run_attention(layer, context, DenseInputs(torch.ones(1, 1, 64))) + assert output.eq(0).all() + assert groups == [16, 32] + + +def test_neo_selection_fallback(monkeypatch): + from dataclasses import replace + + from lightllm.common.attention import selection + from lightllm.common.attention.api import AttentionPhase, ServingConstraints + from lightllm.common.attention.backends.dense.fa3.backend import Fa3DensePrefillBackend + from lightllm.common.attention.backends.dense.fa3 import neo + from lightllm.common.attention.backends.dense.triton.neo import TritonDenseNeoPrefillBackend + from lightllm.utils.backend_validator import ProbeResult, ProbeScope + + def missing(): + raise UnsupportedAttention("Neo extension missing") + + monkeypatch.setattr(neo, "get_neo_attention", missing) + monkeypatch.setattr(Fa3DensePrefillBackend, "check_environment", classmethod(lambda cls, caps: None)) + monkeypatch.setattr(TritonDenseNeoPrefillBackend, "check_environment", classmethod(lambda cls, caps: None)) + monkeypatch.setattr( + selection.backend_validator, "probe_backend", lambda *a, **kw: ProbeResult(True, ProbeScope.PROVIDER_SMOKE) + ) + args = dict( + specs=(replace(make_spec(phase=AttentionPhase.PREFILL), variant="neo"),), + caps=DeviceCaps(Platform.CUDA, torch.device("cuda:0"), capability=(9, 0)), + execution=ServingConstraints(), + policy=selection.BackendPolicy(prefer=("fa3", "triton")), + ) + assert selection.select_attention_backend(**args) is TritonDenseNeoPrefillBackend + with pytest.raises(selection.BackendSelectionError, match="Neo extension missing"): + selection.select_attention_backend(**args, requested="fa3") diff --git a/unit_tests/common/attention/test_state.py b/unit_tests/common/attention/test_state.py new file mode 100644 index 0000000000..dc6ce0e80e --- /dev/null +++ b/unit_tests/common/attention/test_state.py @@ -0,0 +1,116 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState +from lightllm.common.attention_state.commit import RequestGenerationStamp +from lightllm.common.attention_state.manager import AttentionStateManager + + +def test_commit_owns_mapping_and_detects_recycled_requests(): + requests = torch.tensor([2, 2, 4], dtype=torch.int32) + steps = torch.tensor([0, 1, 0], dtype=torch.int32) + generations = torch.zeros(8, dtype=torch.int64) + output = RequestGenerationStamp.stamp("main", requests, steps, generations) + requests.fill_(6) + steps.fill_(5) + assert output.request_indices.tolist() == [2, 2, 4] + assert output.mtp_indices.tolist() == [0, 1, 0] + output.validate_generation(generations) + generations[2] += 1 + with pytest.raises(RuntimeError, match="reused request slot"): + output.validate_generation(generations) + + +def test_overlap_commit_preserves_microbatch_order_and_role(): + generation = torch.zeros(8, dtype=torch.int64) + a = RequestGenerationStamp.stamp("main", torch.tensor([4, 4]), torch.tensor([0, 1]), generation) + b = RequestGenerationStamp.stamp("main", torch.tensor([2]), torch.tensor([0]), generation) + merged = RequestGenerationStamp.concatenate((a, b)) + assert merged.request_indices.tolist() == [4, 4, 2] + assert merged.mtp_indices.tolist() == [0, 1, 0] + with pytest.raises(ValueError, match="same model role"): + RequestGenerationStamp.concatenate((a, replace(b, role="draft0"))) + + +@pytest.mark.parametrize("state_cls", [GDNState, KDAState]) +def test_capture_snapshot_restores_only_participating_role_and_requests(state_cls): + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + manager = AttentionStateManager(request_capacity=3) + main = state_cls(config, request_capacity=3, mtp_step=2, device="cpu") + draft = state_cls(config, request_capacity=3, mtp_step=2, device="cpu") + manager.register_state("main", main) + manager.register_state("draft0", draft) + manager.finish_state_registration() + main.conv_state.fill_(2) + main.ssm_state.fill_(3) + main.mtp_state_index.fill_(1) + snapshot = manager.snapshot_capture_state(torch.tensor([1, 1]), role="main") + main.conv_state[:, 1].fill_(8) + main.ssm_state[:, 3:6].fill_(9) + main.mtp_state_index[1] = 2 + main.conv_state[:, 2].fill_(6) + draft.conv_state.fill_(7) + manager.restore_capture_state(snapshot) + assert torch.all(main.conv_state[:, 1] == 2) + assert torch.all(main.ssm_state[:, 3:6] == 3) + assert main.mtp_state_index[1] == 1 + assert torch.all(main.conv_state[:, 2] == 6) + assert torch.all(draft.conv_state == 7) + + +@pytest.mark.parametrize("kind", ["mla", "packed_bf16", "packed_group128_fp8", "packed_group64_e8m0"]) +def test_mla_adapter_binds_state_layout(kind): + from lightllm.common.attention.api import AttentionFamily, AttentionPhase + from lightllm.common.attention_state.full_att import ( + MLAConfig, + MLAFullAttentionState, + PackedMLALogicalConfig, + create_packed_mla_state, + ) + from lightllm.common.basemodel.attention_adapters import ADAPTERS, AttentionDescription + from unit_tests.common.attention.helpers import model_stub + + nope = kind == "packed_group64_e8m0" + rope_dim = 0 if nope else 64 + latent_dim = 512 + rope_dim + if kind == "mla": + state = MLAFullAttentionState(MLAConfig(((2, 5),), latent_dim, torch.bfloat16), device="cpu") + else: + state = create_packed_mla_state(PackedMLALogicalConfig(((2, 5),), latent_dim, kind), device="cpu") + manager = AttentionStateManager(request_capacity=2) + manager.register_state("main", state) + manager.finish_state_registration() + model = model_stub(llm_kv_type="fp8kv_dsa" if "fp8" in kind or nope else "None") + model.att_state_manager, model.data_type = manager, torch.bfloat16 + layer = SimpleNamespace( + layer_num_=5, + tp_q_head_num_=16, + qk_nope_head_dim=512, + kv_lora_rank=512, + softmax_scale=latent_dim ** -0.5, + network_config_={} if nope else {"qk_rope_head_dim": rope_dim}, + ) + if not nope: + layer.qk_rope_head_dim = rope_dim + family = AttentionFamily.MLA if kind == "mla" else AttentionFamily.SPARSE_MLA + adapter = ADAPTERS[family] + spec = adapter.describe(model, layer, AttentionPhase.PREFILL, AttentionDescription(family)) + assert state.storage is None + state.allocate_storage(8) + resources = adapter.bind(model, layer, None) + assert spec.cache.dtype == resources.kv.dtype == state.storage_dtype + assert resources.qk_rope_head_dim == rope_dim + assert resources.kv.data_ptr() == state.layer_resources(2).data_ptr() + assert resources.kv.shape[-1] == (584 if nope else 656 if kind == "packed_group128_fp8" else latent_dim) diff --git a/unit_tests/common/attention/test_varlen_selection.py b/unit_tests/common/attention/test_varlen_selection.py new file mode 100644 index 0000000000..4d6e0a2c67 --- /dev/null +++ b/unit_tests/common/attention/test_varlen_selection.py @@ -0,0 +1,140 @@ +import pytest +import torch + +from lightllm.common.attention import selection +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.utils import backend_validator +from lightllm.utils.backend_validator import ProbeResult, ProbeScope + + +def test_auto_probes_actual_varlen_path_on_assigned_device(monkeypatch): + calls = [] + + def probe(name, **kwargs): + calls.append((name, kwargs)) + return ProbeResult(name == "sdpa", ProbeScope.VARLEN_NUMERICS, "rejected") + + monkeypatch.setattr(backend_validator, "probe_varlen", probe) + assert selection.select_varlen_attention_op(device=3) == "sdpa" + assert calls == [ + (name, {"device": torch.device("cuda:3"), "dtype": None, "head_dim": 80}) + for name in ("fa3", "xformers", "sdpa") + ] + + +@pytest.mark.parametrize("requested", ["auto", "fa3", "triton", "missing"]) +def test_rejected_ops_are_not_returned(monkeypatch, requested): + calls = [] + + def probe(name, **kwargs): + calls.append(name) + return ProbeResult(False, ProbeScope.VARLEN_NUMERICS, "rejected") + + monkeypatch.setattr(backend_validator, "probe_varlen", probe) + with pytest.raises(selection.VarlenSelectionError) as error: + selection.select_varlen_attention_op(requested) + expected = list(selection.varlen_attention_ops) if requested == "auto" else [requested] + assert list(error.value.failures) == expected + assert calls == ([] if requested == "missing" else expected) + + +def test_varlen_probe_exercises_selected_op(monkeypatch): + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + + monkeypatch.setattr(backend_validator.torch.cuda, "is_available", lambda: False) + assert backend_validator._validate_varlen("sdpa") == (True, None) + monkeypatch.setattr(SdpaVarlenAttentionOp, "forward", lambda self, inputs, out: out.fill_(float("nan"))) + success, reason = backend_validator._validate_varlen("sdpa") + assert not success + assert "mismatch" in reason + + +@pytest.mark.parametrize("check", ["check_environment", "check_support"]) +def test_rejected_probe_does_not_execute_op(monkeypatch, check): + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + + monkeypatch.setattr(backend_validator.torch.cuda, "is_available", lambda: False) + + def reject(cls, *args): + raise UnsupportedAttention("layout") + + monkeypatch.setattr(SdpaVarlenAttentionOp, check, classmethod(reject)) + monkeypatch.setattr(SdpaVarlenAttentionOp, "forward", lambda *args, **kwargs: pytest.fail("rejected operator ran")) + assert backend_validator._validate_varlen("sdpa") == (False, "layout") + + +def test_worker_initialization_keeps_fixed_forward_and_uses_actual_inputs(monkeypatch): + import torch + + from lightllm.common.attention.api import VarlenAttentionInputs + from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp + from lightllm.server import visualserver + + calls = [] + + def select(requested, **kwargs): + calls.append((requested, kwargs)) + return "sdpa" + + monkeypatch.setattr(visualserver, "select_varlen_attention_op", select) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", None) + visualserver.set_vit_attention_op("auto", defer_validation=True) + first = visualserver.get_vit_attention_op() + q = torch.randn(4, 2, 64, dtype=torch.float16) + inputs = VarlenAttentionInputs(q, q, q, torch.tensor([0, 4], dtype=torch.int32), 4) + first.forward(inputs, out=torch.empty_like(q)) + assert calls == [("auto", {"device": q.device, "dtype": q.dtype, "head_dim": 64})] + assert visualserver.get_vit_attention_op() is first + assert isinstance(first._op, SdpaVarlenAttentionOp) + assert "forward" not in vars(first) + first.forward(inputs, out=torch.empty_like(q)) + assert len(calls) == 1 + + +def test_probe_cache_separates_device_dtype_and_layout(monkeypatch): + from dataclasses import replace + + import torch + + from unit_tests.common.attention.helpers import make_spec + + calls = [] + monkeypatch.setattr( + backend_validator, + "_probe", + lambda *args, **kwargs: calls.append(kwargs) or ProbeResult(True, ProbeScope.PROVIDER_SMOKE), + ) + backend_validator.probe_backend.cache_clear() + try: + spec = make_spec() + for device, dtype, shape in ( + ("cuda:0", torch.float16, spec), + ("cuda:1", torch.float16, spec), + ("cuda:0", torch.bfloat16, spec), + ("cuda:0", torch.float16, replace(spec, head_size=64)), + ("cuda:0", torch.float16, spec), + ): + assert backend_validator.probe_backend("provider", device=device, dtype=dtype, spec=shape).ok + assert len(calls) == 4 + finally: + backend_validator.probe_backend.cache_clear() + + +def test_empty_device_assignment_does_not_probe(monkeypatch): + monkeypatch.setattr(backend_validator, "probe_varlen", lambda *a, **k: pytest.fail("empty assignment probed")) + with pytest.raises(ValueError, match="at least one device"): + selection.select_varlen_attention_op(devices=()) + + +def test_default_device_is_resolved_before_probe(monkeypatch): + calls = [] + monkeypatch.setattr( + selection, "get_device_runtime", lambda device: type("Runtime", (), {"device": torch.device("cpu")})() + ) + monkeypatch.setattr( + backend_validator, + "probe_varlen", + lambda *a, **kw: calls.append(kw["device"]) or ProbeResult(True, ProbeScope.VARLEN_NUMERICS), + ) + assert selection.select_varlen_attention_op("sdpa") == "sdpa" + assert calls == [torch.device("cpu")] diff --git a/unit_tests/common/attention/test_workspace.py b/unit_tests/common/attention/test_workspace.py new file mode 100644 index 0000000000..436cef22fb --- /dev/null +++ b/unit_tests/common/attention/test_workspace.py @@ -0,0 +1,60 @@ +import pytest +import torch + +from lightllm.common.attention.api.workspace import AttentionWorkspace + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_workspace_cross_stream_reuse_waits_for_writer(): + workspace = AttentionWorkspace("cuda") + buffer = workspace.scratch("shared", 4096, torch.int32) + buffer.zero_() + producer, consumer = torch.cuda.Stream(), torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + consumer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + torch.cuda._sleep(2_000_000) + buffer.fill_(17) + workspace.finish() + with torch.cuda.stream(consumer): + workspace.begin() + result = buffer.clone() + workspace.finish() + consumer.synchronize() + assert result.eq(17).all() + + larger = workspace.scratch("shared", 8192, torch.int32) + assert larger.data_ptr() != buffer.data_ptr() + assert buffer.eq(17).all() + + +def test_shared_object_type_and_frozen_fork_lifetime(): + workspace = AttentionWorkspace("cpu") + fork = workspace.borrow_for_capture() + shared = fork.get_object("provider_plan", dict) + assert workspace.get_object("provider_plan", dict) is shared + with pytest.raises(TypeError, match="provider_plan requires list"): + workspace.get_object("provider_plan", list) + with pytest.raises(TypeError, match="bad requires dict"): + workspace.get_object("bad", lambda: [], expected_type=dict) + assert "bad" not in workspace._objects + buffer = fork.scratch("scratch", 4) + fork.freeze() + workspace.scratch("scratch", 8) + assert fork.scratch("scratch", 4) is buffer + with pytest.raises(RuntimeError, match="cannot grow"): + fork.scratch("scratch", 8) + with pytest.raises(RuntimeError, match="before capture"): + fork.get_object("next_bucket", dict) + added = workspace.get_object("next_bucket", dict) + assert fork.get_object("next_bucket", dict) is added + + +def test_cpu_stream_and_event_have_consistent_noop_lifecycle(): + workspace = AttentionWorkspace("cpu") + runtime = workspace.runtime + assert runtime.create_stream() is None + assert runtime.create_event() is None + with runtime.use_stream(runtime.create_stream()): + workspace.begin() + workspace.finish() diff --git a/unit_tests/common/basemodel/attention/fa3/test_neo.py b/unit_tests/common/basemodel/attention/fa3/test_neo.py index 9ef88130cf..e649efadc6 100644 --- a/unit_tests/common/basemodel/attention/fa3/test_neo.py +++ b/unit_tests/common/basemodel/attention/fa3/test_neo.py @@ -5,8 +5,19 @@ import pytest import torch +from lightllm.common.attention.backends.dense.fa3 import neo +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.utils import sgl_utils from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.fa3 import fp, neo +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) def _load_neo(monkeypatch, attention_func): @@ -15,7 +26,7 @@ def _load_neo(monkeypatch, attention_func): interface = ModuleType("flash_attn_interface") interface.flash_attn_with_kvcache = attention_func monkeypatch.setitem(sys.modules, "flash_attn_interface", interface) - name = "lightllm.common.basemodel.attention.fa3._neo_import_test" + name = "lightllm.common.attention.backends.dense._neo_import_test" spec = importlib.util.spec_from_file_location(name, neo.__file__) module = importlib.util.module_from_spec(spec) monkeypatch.setitem(sys.modules, name, module) @@ -27,12 +38,11 @@ def test_recognizes_neo_interface_without_replacing_regular_fa3(monkeypatch): def neo_attention(q, image_token_end=None): return q - regular_attention = fp.flash_attn_with_kvcache + regular_attention = sgl_utils.flash_attn_with_kvcache module = _load_neo(monkeypatch, neo_attention) - assert module.HAS_FLASH_ATTN_INTERFACE - assert module.flash_attn_with_kvcache_neo is neo_attention - assert fp.flash_attn_with_kvcache is regular_attention + assert module.get_neo_attention() is neo_attention + assert sgl_utils.flash_attn_with_kvcache is regular_attention @pytest.mark.parametrize("interface_kind", ["missing", "regular", "legacy_tag"]) @@ -44,17 +54,15 @@ def legacy_attention(q, image_token_tag=None): return q attention_func = {"missing": None, "regular": regular_attention, "legacy_tag": legacy_attention}[interface_kind] - with pytest.warns(UserWarning, match="Neo FA3 is unavailable") as warnings: - module = _load_neo(monkeypatch, attention_func) + module = _load_neo(monkeypatch, attention_func) + with pytest.raises(UnsupportedAttention, match="Neo FA3 is unavailable") as error: + module.get_neo_attention() - assert not module.HAS_FLASH_ATTN_INTERFACE - assert module.flash_attn_with_kvcache_neo is None - message = str(warnings[0].message) + message = str(error.value) assert "https://github.com/WANDY666/flash-attention.git" in message assert "cd flash-attention/hopper" in message assert "FLASH_ATTENTION_FORCE_BUILD=TRUE" in message assert "--no-deps --force-reinstall" in message - assert "--llm_prefill_att_backend triton" in message @pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") @@ -82,9 +90,9 @@ def indices(values): infer_state.b1_cu_kv_seq_len = indices([0, 2 * page_size, 4 * page_size - 1]).long() infer_state.b_image_token_end = image_ends infer_state.req_manager = SimpleNamespace(req_to_token_indexs=req_to_token_indexs) - backend = SimpleNamespace(uses_causal_attention=lambda: True, infer_page_size=page_size) - state = neo.NeoFa3AttBackend.create_att_prefill_state(backend, infer_state) - state.init_state() + backend = SimpleNamespace(uses_causal_attention=lambda: True, kernel_page_size=page_size) + state = backend_context(neo.Fa3DenseNeoDecodeBackend, "prefill", settings=backend) + update_context(state, infer_state) q = torch.randn(3, 2, 16, device="cuda") k = torch.randn(4 * page_size, 1, 16, device="cuda") @@ -107,13 +115,17 @@ def neo_attention(*, image_token_end, **kwargs): assert "image_token_tag" not in kwargs return output - monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", True) - monkeypatch.setattr(neo, "flash_attn_with_kvcache_neo", neo_attention) - assert state.prefill_att(q, k, v) is output + monkeypatch.setattr(neo, "get_neo_attention", lambda: neo_attention) + assert run_dense(state, infer_state, q, k, v) is output def test_prefill_missing_extension_reports_install_instructions(monkeypatch): - monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", False) - state = neo.NeoFa3PrefillAttState() - with pytest.raises(ImportError, match="WANDY666/flash-attention"): - state.prefill_att(None, None, None) + module = _load_neo(monkeypatch, None) + monkeypatch.setattr(neo, "get_neo_attention", module.get_neo_attention) + state = backend_context(neo.Fa3DenseNeoDecodeBackend, "prefill") + with pytest.raises(UnsupportedAttention, match="WANDY666/flash-attention"): + state.backend_identity.forward( + SimpleNamespace(alloc_tensor=torch.empty, resources=SimpleNamespace(key=None, value=None)), + state, + SimpleNamespace(query=None), + ) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py deleted file mode 100644 index 190d8f1162..0000000000 --- a/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py +++ /dev/null @@ -1,100 +0,0 @@ -import sys -import weakref -from types import ModuleType, SimpleNamespace - -import torch - -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState -from lightllm.common.basemodel.attention.flashinfer import mla -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -def test_fp_graph_copy_refreshes_plan_with_runtime_length_and_preserves_buffers(monkeypatch): - plan_calls = [] - flashinfer = ModuleType("flashinfer") - decode = ModuleType("flashinfer.decode") - decode.fast_decode_plan = lambda wrapper, **kwargs: plan_calls.append((wrapper, kwargs)) - flashinfer.decode = decode - monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) - monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) - backend = SimpleNamespace( - infer_page_size=16, - tp_q_head_num=4, - tp_kv_head_num=2, - head_dim=8, - q_data_type=torch.bfloat16, - kv_data_type=torch.bfloat16, - ) - captured_owner = InferStateInfo() - captured_owner.batch_size = 2 - runtime_owner = InferStateInfo() - runtime_owner.max_kv_seq_len = 17 - captured = FlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), - kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), - kv_indices=torch.arange(8, dtype=torch.int32), - kv_last_page_len_buffer=torch.tensor([16, 16], dtype=torch.int32), - decode_wrapper=object(), - ) - runtime = FlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), - kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), - kv_last_page_len_buffer=torch.tensor([1, 5], dtype=torch.int32), - ) - buffer_names = ("kv_starts", "kv_indices", "kv_last_page_len_buffer") - captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} - - captured.copy_for_decode_cuda_graph(runtime) - - assert len(plan_calls) == 1 - wrapper, plan = plan_calls[0] - assert wrapper is captured.decode_wrapper - assert plan["indptr"] is captured.kv_starts - assert plan["indices"] is captured.kv_indices - assert plan["last_page_len"] is captured.kv_last_page_len_buffer - torch.testing.assert_close(plan["global_override_indptr_cpu"], torch.tensor([0, 2, 4], dtype=torch.int32)) - for name in buffer_names: - assert getattr(captured, name).data_ptr() == captured_ptrs[name] - torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) - - -def test_mla_graph_copy_refreshes_plan_with_runtime_lengths_and_preserves_buffers(monkeypatch): - plan_calls = [] - monkeypatch.setattr(mla, "fill_mla_decode_plan_for_cuda_graph", lambda **kwargs: plan_calls.append(kwargs)) - backend = SimpleNamespace(tp_q_head_num=4) - captured_owner = InferStateInfo() - captured_owner.batch_size = 2 - runtime_owner = InferStateInfo() - runtime_owner.b_seq_len = torch.tensor([17, 5], dtype=torch.int32) - runtime_owner.max_kv_seq_len = 17 - captured = mla.MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), - kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), - kv_indices=torch.arange(8, dtype=torch.int32), - decode_wrapper=object(), - ) - runtime = mla.MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), - kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), - ) - buffer_names = ("kv_starts", "kv_indices") - captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} - - captured.copy_for_decode_cuda_graph(runtime) - - assert len(plan_calls) == 1 - plan = plan_calls[0] - assert plan["decode_wrapper"] is captured.decode_wrapper - assert plan["kv_page_indptr"] is captured.kv_starts - assert plan["batch_size"] == 2 - assert plan["max_kv_len"] == 17 - torch.testing.assert_close(plan["kv_lens"], torch.tensor([17, 5], dtype=torch.int32)) - for name in buffer_names: - assert getattr(captured, name).data_ptr() == captured_ptrs[name] - torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py index 1392d5e688..6a37e7c378 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py @@ -1,68 +1,50 @@ +from unit_tests.common.attention.helpers import plan import sys -import weakref +from dataclasses import replace from types import ModuleType, SimpleNamespace from unittest.mock import Mock import pytest import torch -from lightllm.common.basemodel.attention.flashinfer import fp, fp8, mla -from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -def test_decode_wrapper_rejects_conflicting_graph_flags(): - infer_state = InferStateInfo() - infer_state.is_cuda_graph_capture = True - infer_state.is_cuda_graph_replay = True - - with pytest.raises(ValueError, match="capture=True, replay=True"): - should_init_decode_wrapper(infer_state) +from lightllm.common.attention.backends.dense.flashinfer import backend as fp +from lightllm.common.attention.backends.dense.flashinfer.fp8 import FlashInferDenseFp8DecodeBackend +from lightllm.common.attention.backends.mla.flashinfer import backend as mla +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context @pytest.fixture( - params=[fp.FlashInferDecodeAttState, fp8.Fp8FlashInferDecodeAttState, mla.MlaFlashInferDecodeAttState], - ids=["fp", "fp8", "mla"], + params=[fp.FlashInferDenseDecodeBackend, FlashInferDenseFp8DecodeBackend, mla.FlashInferMLADecodeBackend] ) -def decode_init(request, monkeypatch): - state_cls = request.param - module = mla if state_cls is mla.MlaFlashInferDecodeAttState else fp - page_size = 1 if state_cls is fp8.Fp8FlashInferDecodeAttState else 4 - max_page_num = 64 // page_size +def decode_case(request, monkeypatch): + cls = request.param + module = mla if cls is mla.FlashInferMLADecodeBackend else fp wrapper = SimpleNamespace(plan=Mock()) - create_wrapper = Mock(return_value=wrapper) + create = Mock(return_value=wrapper) flashinfer = ModuleType("flashinfer") - flashinfer.decode = SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=create_wrapper) - flashinfer.mla = SimpleNamespace(BatchMLAPagedAttentionWrapper=create_wrapper) + decode = ModuleType("flashinfer.decode") + decode.BatchDecodeWithPagedKVCacheWrapper = create + decode.fast_decode_plan = Mock() + flashinfer.decode = decode + flashinfer.mla = SimpleNamespace(BatchMLAPagedAttentionWrapper=create) monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) - def cpu_arange(*args, **kwargs): - kwargs["device"] = "cpu" - return torch.arange(*args, **kwargs) - - monkeypatch.setattr(module, "torch", SimpleNamespace(int32=torch.int32, arange=cpu_arange, empty=torch.empty)) - - def cumsum_params(q_lengths, kv_lengths): - return tuple( - torch.cat((torch.zeros(1, dtype=torch.int32), lengths.cumsum(0, dtype=torch.int32))) - for lengths in (q_lengths, kv_lengths) - ) - - monkeypatch.setattr(module, "gen_cumsum_pad0_tensor", cumsum_params) - - def repack(**kwargs): - kwargs["out_page_indices"].fill_(7) + def cumsum(q, kv): + return tuple(torch.cat((torch.zeros(1, dtype=torch.int32), v.cumsum(0).int())) for v in (q, kv)) - repack_indices = Mock(side_effect=repack) - monkeypatch.setattr(module, "repack_kv_index", repack_indices) - can_run = Mock(return_value=True) - backend = SimpleNamespace( - model=SimpleNamespace( - graph=SimpleNamespace(can_run=can_run), graph_max_batch_size=4, graph_max_len_in_batch=64 - ), - infer_page_size=page_size, - max_page_num=max_page_num, - kv_indices_buffer=[torch.empty(4 * max_page_num, dtype=torch.int32)], + monkeypatch.setattr("lightllm.common.basemodel.triton_kernel.gen_prefill_params.gen_cumsum_pad0_tensor", cumsum) + repack = Mock(side_effect=lambda **kwargs: kwargs["out_page_indices"].fill_(int(kwargs["b_token_len"][0]))) + monkeypatch.setattr("lightllm.common.basemodel.triton_kernel.repack_kv_index.repack_kv_index", repack) + mla_refresh = Mock() + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.flashinfer_mla_plan.fill_mla_decode_plan_for_cuda_graph", mla_refresh + ) + page_size = 1 if cls is FlashInferDenseFp8DecodeBackend else 4 + backend = make_backend( + cls, + kernel_page_size=page_size, + max_page_num=64 // page_size, tp_q_head_num=4, tp_kv_head_num=2, head_dim=8, @@ -70,63 +52,54 @@ def repack(**kwargs): qk_rope_head_dim=2, softmax_scale=0.25, q_data_type=torch.bfloat16, - kv_data_type=torch.float8_e4m3fn if page_size == 1 else torch.bfloat16, - workspace_buffer_key="decode_init_test", - workspace_buffer_size=1, - get_gpu_workspace_buffer=Mock(return_value=torch.empty(1, dtype=torch.uint8)), + kv_data_type=torch.bfloat16, + workspace_buffer_key="test", + workspace_buffer_size=8, ) return SimpleNamespace( - module=module, - state=state_cls(backend=backend), backend=backend, + create=create, wrapper=wrapper, - create_wrapper=create_wrapper, - repack_indices=repack_indices, - can_run=can_run, + repack=repack, + refresh=mla_refresh if module is mla else decode.fast_decode_plan, + mla=module is mla, ) -@pytest.mark.parametrize("mode", ["eager", "capture", "replay"]) -def test_decode_init_uses_explicit_graph_mode_for_wrapper_creation(decode_init, mode): - infer_state = InferStateInfo() - infer_state.input_ids = torch.tensor([1, 2], dtype=torch.int64) - infer_state.batch_size = 2 - infer_state.b_seq_len = torch.tensor([5, 2], dtype=torch.int32) - infer_state.max_kv_seq_len = 5 - infer_state.b_req_idx = torch.tensor([2, 3], dtype=torch.int32) - infer_state.microbatch_index = 0 - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.zeros((4, 64), dtype=torch.int32)) - if mode == "capture": - infer_state.is_cuda_graph_capture = True - elif mode == "replay": - infer_state.is_cuda_graph_replay = True - - decode_init.state._infer_state_weak_ref = weakref.ref(infer_state) - decode_init.state.init_state() - - decode_init.can_run.assert_not_called() - expected_starts = [0, 5, 7] if decode_init.backend.infer_page_size == 1 else [0, 2, 3] - torch.testing.assert_close(decode_init.state.kv_starts, torch.tensor(expected_starts, dtype=torch.int32)) - assert decode_init.state.kv_indices.data_ptr() == decode_init.backend.kv_indices_buffer[0].data_ptr() - assert torch.all(decode_init.state.kv_indices == 7) - decode_init.repack_indices.assert_called_once() - repack_args = decode_init.repack_indices.call_args.kwargs - assert repack_args["b_req_idx"] is infer_state.b_req_idx - assert repack_args["b_token_len"] is infer_state.b_seq_len - if decode_init.module is fp: - expected_last_page_lengths = [1, 1] if decode_init.backend.infer_page_size == 1 else [1, 2] - torch.testing.assert_close( - decode_init.state.kv_last_page_len_buffer, torch.tensor(expected_last_page_lengths, dtype=torch.int32) +@pytest.mark.parametrize("capture", [False, True]) +def test_decode_context_creates_wrapper_once_with_explicit_graph_mode(decode_case, capture): + case = decode_case + batch = make_batch() + context = make_context(case.backend, batch, capture=capture, max_length=64, prepare=True) + assert context.decode_wrapper is case.wrapper + case.create.assert_called_once() + assert case.create.call_args.kwargs["use_cuda_graph"] is capture + assert context.kv_starts.tolist() == ([0, 5, 7] if case.backend.kernel_page_size == 1 else [0, 2, 3]) + case.wrapper.plan.assert_called_once() + case.repack.assert_called_once() + + +def test_graph_replay_refreshes_existing_wrapper_with_current_lengths(decode_case): + case = decode_case + batch = make_batch() + context = make_context(case.backend, batch, capture=True, max_length=64, prepare=True) + pointers = (context.kv_starts.data_ptr(), context.kv_indices.data_ptr(), context.batch.b_seq_len.data_ptr()) + for lengths in ((17, 5), (1, 1), (60, 3)): + runtime = make_batch(lengths, req_to_token_indexs=batch.req_to_token_indexs) + plan(case.backend, context, runtime) + assert pointers == ( + context.kv_starts.data_ptr(), + context.kv_indices.data_ptr(), + context.batch.b_seq_len.data_ptr(), ) - else: - torch.testing.assert_close(decode_init.state.q_indptr, torch.tensor([0, 1, 2], dtype=torch.int32)) - - if mode == "replay": - assert decode_init.state.decode_wrapper is None - decode_init.create_wrapper.assert_not_called() - decode_init.backend.get_gpu_workspace_buffer.assert_not_called() - else: - assert decode_init.state.decode_wrapper is decode_init.wrapper - decode_init.create_wrapper.assert_called_once() - decode_init.backend.get_gpu_workspace_buffer.assert_called_once() - decode_init.wrapper.plan.assert_called_once() + assert context.batch.b_seq_len.tolist() == list(lengths) + assert context.kv_indices[0] == lengths[0] + args = case.refresh.call_args.kwargs + if case.mla: + assert args["max_kv_len"] == max(lengths) + assert args["decode_wrapper"] is case.wrapper + else: + pages = (max(lengths) + case.backend.kernel_page_size - 1) // case.backend.kernel_page_size + assert args["global_override_indptr_cpu"].tolist() == [0, pages, 2 * pages] + case.create.assert_called_once() + assert case.refresh.call_count == 3 diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py index ccf5dfd454..13381c7410 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py @@ -1,15 +1,14 @@ -from types import SimpleNamespace import weakref +from types import SimpleNamespace import pytest import torch flashinfer = pytest.importorskip("flashinfer") -from lightllm.common.basemodel.attention.flashinfer.mla import ( # noqa: E402 - MlaFlashInferDecodeAttState, -) +from lightllm.common.attention.backends.mla.flashinfer.backend import FlashInferMLADecodeBackend from lightllm.common.basemodel.infer_struct import InferStateInfo # noqa: E402 +from unit_tests.common.attention.helpers import backend_context def _make_wrapper( @@ -124,13 +123,14 @@ def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, page_size, init_short=True, ) - backend = SimpleNamespace(tp_q_head_num=num_heads, infer_page_size=page_size) + backend = SimpleNamespace(tp_q_head_num=num_heads, kernel_page_size=page_size) captured_owner = InferStateInfo() captured_owner.batch_size = batch_size captured_owner.b_seq_len = graph_wrapper._kv_len_arr_buf - graph_state = MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), + graph_state = backend_context( + FlashInferMLADecodeBackend, + "decode", + settings=backend, kv_starts=graph_wrapper._kv_indptr_buf, kv_indices=graph_wrapper._kv_indices_buf, decode_wrapper=graph_wrapper, @@ -143,7 +143,7 @@ def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, with torch.cuda.graph(graph): graph_wrapper.run(q_nope, q_pe, ckv, kpe, out=graph_out, return_lse=False) - # 复用同一 graph，覆盖页表变化、末页空位及 split/exact 切换后的 plan 更新。 + # 复用同一 graph，覆盖页表变化、末页空位及 split/exact 切换后的 context 更新。 for replay_lengths in (lengths, [1] * batch_size, lengths): replay_lens = torch.tensor(replay_lengths, dtype=torch.int32, device="cuda") replay_indptr = torch.zeros_like(kv_indptr) @@ -152,13 +152,11 @@ def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, runtime_owner.batch_size = batch_size runtime_owner.b_seq_len = replay_lens runtime_owner.max_kv_seq_len = max(replay_lengths) - new_state = MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=replay_indptr, - kv_indices=kv_indices, - ) - graph_state.copy_for_decode_cuda_graph(new_state) + graph_state.kv_starts.copy_(replay_indptr) + graph_state.kv_indices.copy_(kv_indices) + graph_wrapper._kv_len_arr_buf.copy_(replay_lens) + graph_state.batch = runtime_owner + graph_state.backend_identity._refresh_decode_plan(graph_state, replay_lens, max(replay_lengths)) graph.replay() ref_wrapper.plan( q_indptr, diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py index 4c33d8cf8a..9a98ae0bf9 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py @@ -25,7 +25,8 @@ ) def test_nope_sparse_mla_eager_and_graph_replay(num_tokens, num_heads, kv_dtype): wrapper_cls = flashinfer.mla.BatchMLAPagedAttentionWrapper - assert "ckv_scale_arr" in inspect.signature(wrapper_cls.run).parameters, "FlashInfer >= 0.6.18 is required" + if "ckv_scale_arr" not in inspect.signature(wrapper_cls.run).parameters: + pytest.skip("FlashInfer >= 0.6.18 is required") torch.manual_seed(17) dim, num_slots, width = 512, 4096, 2176 diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py b/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py index 924dead570..4d2d0658b9 100644 --- a/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py +++ b/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py @@ -1,56 +1,30 @@ -from lightllm.common.basemodel.attention import base_att -from lightllm.common.basemodel.attention.base_att import BaseAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend - - -def test_workspace_is_shared_by_backend_family_and_device(monkeypatch): - allocations = [] - current_device_id = 0 - - def fake_empty(size, dtype, device): - workspace = object() - allocations.append((workspace, size, dtype, device)) - return workspace - - monkeypatch.setattr(base_att, "get_current_device_id", lambda: current_device_id) - monkeypatch.setattr(base_att.torch, "empty", fake_empty) - monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) - - def get_workspace(backend_cls): - return backend_cls.get_gpu_workspace_buffer( - key_name=backend_cls.workspace_buffer_key, - workspace_size=backend_cls.workspace_buffer_size, - ) - - fp_workspace = get_workspace(FlashInferAttBackend) - fp8_workspace = get_workspace(Fp8FlashInferAttBackend) - mla_workspace = get_workspace(MlaFlashInferAttBackend) - - assert fp8_workspace is fp_workspace - assert mla_workspace is not fp_workspace - assert [allocation[1] for allocation in allocations] == [ - 512 * 1024 * 1024, - 256 * 1024 * 1024, - ] - - current_device_id = 1 - other_device_workspace = get_workspace(FlashInferAttBackend) - - assert other_device_workspace is not fp_workspace - assert allocations[-1][3] == 1 - - -def test_workspace_key_includes_size_and_dtype(monkeypatch): - monkeypatch.setattr(base_att, "get_current_device_id", lambda: 0) - monkeypatch.setattr(base_att.torch, "empty", lambda *args, **kwargs: object()) - monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) - - small_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) - large_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 2048) - fp16_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024, dtype=base_att.torch.float16) - - assert BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) is small_workspace - assert large_workspace is not small_workspace - assert fp16_workspace is not small_workspace +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.attention import ModelAttention +from lightllm.platform import get_device_runtime + + +def test_workspace_isolated_by_model_group_and_slot(): + model = ModelAttention.__new__(ModelAttention) + model.workspaces, model.runtime = {}, get_device_runtime("cpu") + other = ModelAttention.__new__(ModelAttention) + other.workspaces, other.runtime = {}, get_device_runtime("cpu") + primary = model.workspace(0, 0) + assert primary is model.workspace(0, 0) + assert primary is not model.workspace(1, 0) + assert primary is not model.workspace(0, 1) + assert primary is not other.workspace(0, 0) + + +def test_workspace_growth_keeps_captured_references_alive(): + model = ModelAttention.__new__(ModelAttention) + model.workspaces, model.runtime = {}, get_device_runtime("cpu") + workspace = model.workspace(0, 0) + captured = workspace.scratch("indices", 8, torch.int32) + captured.fill_(7) + enlarged = workspace.scratch("indices", 32, torch.int32) + assert enlarged.data_ptr() != captured.data_ptr() + assert captured.tolist() == [7] * 8 + assert workspace.scratch("indices", 4, torch.int32).data_ptr() == enlarged.data_ptr() diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 344bfea57f..7061a71896 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -1,228 +1,98 @@ -import weakref -import inspect +from unit_tests.common.attention.helpers import build, run_attention import sys from types import ModuleType, SimpleNamespace import pytest import torch -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.infer_struct import InferStateInfo -import lightllm.common.basemodel.attention.linear.gdn as gdn -import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils -import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops -from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.models.qwen3next.config import Qwen3NextResourceConfig -from lightllm.server.api_cli import make_argument_parser +import lightllm.common.attention.backends.gdn.triton.kernels.chunk as fla_ops import lightllm.utils.backend_validator as backend_validator +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import GDNInputs, GDNResources +from lightllm.common.attention.backends.gdn.triton import backend as gdn +from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNDecodeBackend, TritonGDNPrefillBackend +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention @pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float32]) def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): - ssm_states = torch.zeros((1, 1, 2, 2), dtype=cache_dtype) - final_state = torch.ones((1, 1, 2, 2), dtype=torch.float32) - - monkeypatch.setattr(gdn, "fused_gdn_gating", lambda _log, a, b, _bias: (a, b)) - monkeypatch.setattr(gdn, "causal_conv1d_fn", lambda mixed, *args, **kwargs: mixed) - qkv = torch.zeros((1, 3), dtype=cache_dtype) + ssm = torch.zeros((1, 1, 2, 2), dtype=cache_dtype) + final = torch.ones((1, 1, 2, 2), dtype=torch.float32) q = torch.zeros((1, 1, 1, 1), dtype=cache_dtype) - backend = SimpleNamespace( + monkeypatch.setattr( + "lightllm.common.attention.backends.gdn.triton.kernels.fused_gdn_gating.fused_gdn_gating", + lambda _log, a, b, _bias: (a, b), + ) + monkeypatch.setattr( + "lightllm.common.attention.backends._shared.linear.causal_conv1d.causal_conv1d_fn", + lambda mixed, *args, **kwargs: mixed, + ) + backend = make_backend( + TritonGDNDecodeBackend, mtp_step=0, activation="silu", - ssm_state_dtype=cache_dtype, _rearrange_mixed_qkv=lambda mixed: (q, q, q), - prefill_kernel=lambda *args, **kwargs: (None, final_state), - ) - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1], dtype=torch.int32) - infer_state.b_ready_cache_len = 0 - state = gdn.LinearAttPrefillAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.b_conv_buffer_idx = torch.tensor([0], dtype=torch.int64) - state.b_ssm_buffer_idx = torch.tensor([0], dtype=torch.int64) - layer_weight = SimpleNamespace( - linear_A_log=SimpleNamespace(weight=None), - linear_dt_bias=SimpleNamespace(weight=None), - linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=None), bias=None), - ) - - state._gdn_prefill_kernel( - qkv, - torch.zeros((1, 3), dtype=cache_dtype), - ssm_states, - torch.zeros((1, 1), dtype=cache_dtype), - torch.zeros((1, 1), dtype=cache_dtype), - state.infer_state, - layer_weight, - ) - - assert ssm_states.dtype == cache_dtype - assert torch.equal(ssm_states, final_state.to(cache_dtype)) - - -def test_decode_state_initializes_normal_layout(): - b_req_idx = torch.tensor([2, 4], dtype=torch.int32) - state, infer_state = _create_decode_state(draft_step=0, dynamic_layout=False, b_req_idx=b_req_idx) - - state.init_state() - - assert state.b_conv_buffer_idx is b_req_idx - assert state.b_ssm_buffer_idx is b_req_idx - assert state.b1_mtp_cu_q_seq_len is None - assert state.b_num_accepted_tokens is None - - -@pytest.mark.parametrize("role", ["main", "draft0"]) -def test_decode_state_initializes_fixed_mtp_layout(role: str): - b_req_idx = torch.tensor([2, 2, 2, 4, 4, 4], dtype=torch.int32) - req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) - state, infer_state = _create_decode_state( + prefill_kernel=lambda **kwargs: (q, final), + _split_qkvzba=lambda mixed: (mixed, q, q, q), + ) + batch = make_batch((1,), phase=AttentionPhase.PREFILL) + context = make_context(backend, batch) + build(backend, context) + params = GDNResources(torch.zeros((1, 3)), ssm, None, None, None, None) + attention = make_layer_attention(backend, context, params) + result, z = run_attention(attention, context, GDNInputs(torch.zeros((1, 3), dtype=cache_dtype))) + assert result is q and z is q + assert ssm.dtype == cache_dtype + torch.testing.assert_close(ssm, final.to(cache_dtype)) + + +@pytest.mark.parametrize("provider", [TritonGDNDecodeBackend]) +def test_decode_and_fixed_mtp_state_routes(provider): + requests = torch.tensor([2, 2, 2, 4, 4, 4], dtype=torch.int32) + cursor = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) + backend = make_backend( + provider, draft_step=2, - dynamic_layout=False, - b_req_idx=b_req_idx, - req_to_mtp_state_index=req_to_mtp_state_index, - role=role, ) + context = make_context(backend, make_batch((3,) * 6, b_req_idx=requests)) + build(backend, context) + assert context.b_conv_buffer_idx.tolist() == [2, 4] + assert context.accepted_tokens(cursor).tolist() == [2, 3] + assert context.b_ssm_buffer_idx.tolist() == [[6, 7, 8], [12, 13, 14]] + backend.draft_step = 0 + context = make_context(backend, make_batch((3, 3), b_req_idx=torch.tensor([2, 4]))) + build(backend, context) + assert context.b_ssm_buffer_idx.tolist() == [2, 4] - state.init_state() - torch.testing.assert_close(state.b1_mtp_cu_q_seq_len, torch.tensor([0, 3, 6], dtype=torch.int32)) - torch.testing.assert_close(state.b_conv_buffer_idx, torch.tensor([2, 4], dtype=torch.int32)) - torch.testing.assert_close(state.b_num_accepted_tokens, torch.tensor([2, 3], dtype=torch.int32)) - torch.testing.assert_close( - state.b_ssm_buffer_idx, - torch.tensor([[6, 7, 8], [12, 13, 14]], dtype=torch.int32), - ) - - -@pytest.mark.parametrize("role", ["main", "draft0"]) -def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch, role: str): - b_req_idx = torch.tensor([2, 2, 4, 4, 4], dtype=torch.int32) - req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) - expected_cu_q_seq_len = torch.tensor([0, 2, 5, 5, 5, 5], dtype=torch.int32) - expected_conv_buffer_idx = torch.tensor([2, 4, 5, 5, 5], dtype=torch.int32) - expected_num_accepted_tokens = torch.tensor([2, 3, 1, 1, 1], dtype=torch.int32) - build_calls = [] +def test_dynamic_mtp_layout_uses_bound_state_cursor(monkeypatch): + calls = [] - def build_params(**kwargs): - build_calls.append(kwargs) - return expected_cu_q_seq_len, expected_conv_buffer_idx, expected_num_accepted_tokens - - monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) - state, infer_state = _create_decode_state( - draft_step=2, - dynamic_layout=True, - b_req_idx=b_req_idx, - req_to_mtp_state_index=req_to_mtp_state_index, - role=role, - ) + def build_layout(**kwargs): + calls.append(kwargs) + return torch.tensor([0, 1, 3]), torch.tensor([2, 4]) - state.init_state() - - assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len - assert state.b_conv_buffer_idx is expected_conv_buffer_idx - assert state.b_num_accepted_tokens is expected_num_accepted_tokens - torch.testing.assert_close( - state.b_ssm_buffer_idx, - torch.tensor( - [[6, 7, 8], [12, 13, 14], [15, 16, 17], [15, 16, 17], [15, 16, 17]], - dtype=torch.int32, - ), - ) - assert len(build_calls) == 1 - assert build_calls[0]["b_req_idx"] is state.infer_state.b_req_idx - assert build_calls[0]["b_mtp_index"] is state.infer_state.b_mtp_index - assert build_calls[0]["req_to_mtp_state_index"] is req_to_mtp_state_index - assert build_calls[0]["hold_req_id"] == 5 - - -def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): - assert inspect.isabstract(gdn.LinearAttBackend) - backend = object.__new__(TritonLinearAttBackend) - assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule - - -@pytest.fixture -def auto_linear_backend_args(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3"], - llm_decode_att_backend=["flashinfer"], - ssm_state_mode="native", - ), - ) - - -def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear_backend_args): - validate_calls = [] - flashqla = ModuleType("flash_qla") - flashqla.chunk_gated_delta_rule = lambda **kwargs: ("flashqla", kwargs) - monkeypatch.setitem(sys.modules, "flash_qla", flashqla) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name: validate_calls.append(name) or True, - ) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - - assert backend_class is FlashQlaLinearAttBackend - backend = object.__new__(backend_class) - assert backend.get_prefill_kernel()(q="q")[0] == "flashqla" - assert validate_calls == ["flashqla"] - - -def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3", "flashqla"], - llm_decode_att_backend=["flashinfer", "triton"], - ssm_state_mode="replay", - ), + "lightllm.common.attention.backends._shared.linear.mtp_state_params.build_dynamic_mtp_layout", build_layout ) - from lightllm.common.basemodel.attention.base_att import BaseAttBackend - from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend - from lightllm.common.basemodel.attention.linear.replay_gdn import ( - ReplayGDNFlashQlaAttBackend, - ReplayGDNTritonAttBackend, + cursor = torch.zeros(8, dtype=torch.int32) + backend = make_backend( + TritonGDNDecodeBackend, + draft_step=2, + dynamic_verify=True, ) + batch = make_batch((3, 3, 3), b_req_idx=torch.tensor([2, 4, 4]), b_mtp_index=torch.tensor([0, 0, 1])) + context = make_context(backend, batch) + build(backend, context) + assert context.accepted_tokens(cursor).tolist() == [1, 1] + assert calls[0]["hold_req_id"] == 7 + assert context.b_ssm_buffer_idx.tolist() == [[6, 7, 8], [12, 13, 14]] - prefill_backend = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - decode_backend = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert prefill_backend is ReplayGDNFlashQlaAttBackend - assert decode_backend is ReplayGDNTritonAttBackend - assert prefill_backend.__bases__ == (BaseAttBackend,) - assert decode_backend.__bases__ == (BaseAttBackend,) - assert not issubclass(prefill_backend, LinearAttBackend) - assert not issubclass(decode_backend, LinearAttBackend) - - -def test_kda_entry_selects_native_or_replay(monkeypatch): - from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend - from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend - - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace(ssm_state_mode="native"), - ) - assert linear_create_utils.get_kda_linear_att_backend_class() is KDALinearAttBackend - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace(ssm_state_mode="replay"), - ) - assert linear_create_utils.get_kda_linear_att_backend_class() is ReplayKDALinearAttBackend +def test_triton_backend_supplies_chunk_kernel(): + backend = object.__new__(TritonGDNPrefillBackend) + assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule def test_flashqla_validation_loads_linear_cache_config(monkeypatch): @@ -285,154 +155,3 @@ def cpu_tensor(*args, **kwargs): assert kernel_calls[0]["q"].dtype == torch.bfloat16 assert kernel_calls[0]["initial_state"].shape == (1, 4, 8, 16) assert kernel_calls[0]["initial_state"].dtype == torch.float32 - - -def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3", "triton"], - llm_decode_att_backend=["flashinfer", "triton"], - ssm_state_mode="native", - ), - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda *args: pytest.fail("an explicit linear backend must not be auto-validated"), - ) - - prefill_backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - decode_backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert prefill_backend_class is TritonLinearAttBackend - assert decode_backend_class is TritonLinearAttBackend - - -def test_cli_accepts_linear_backend_at_index_one(): - args = make_argument_parser().parse_args( - [ - "--llm_prefill_att_backend", - "fa3", - "flashqla", - "--llm_decode_att_backend", - "flashinfer", - "triton", - ] - ) - - assert args.llm_prefill_att_backend == ["fa3", "flashqla"] - assert args.llm_decode_att_backend == ["flashinfer", "triton"] - - -@pytest.mark.parametrize( - "get_backend_class", - [ - linear_create_utils.get_gdn_linear_prefill_att_backend_class, - linear_create_utils.get_gdn_linear_decode_att_backend_class, - ], -) -def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_backend_class): - with pytest.raises(AssertionError, match="index must be 1"): - get_backend_class(index=0) - - -def test_missing_linear_decode_backend_arg_defaults_to_triton(auto_linear_backend_args): - backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert backend_class is TritonLinearAttBackend - - -def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_args): - validate_calls = [] - monkeypatch.setattr( - linear_create_utils, - "gdn_decode_backends", - {"optimized": FlashQlaLinearAttBackend, "triton": TritonLinearAttBackend}, - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name: validate_calls.append(name) or True, - ) - - backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class( - index=1, priority_list=("optimized", "triton") - ) - - assert backend_class is FlashQlaLinearAttBackend - assert validate_calls == ["optimized"] - - -def test_gdn_prefill_backend_falls_back_when_flashqla_validation_fails(monkeypatch, auto_linear_backend_args): - monkeypatch.setattr(linear_create_utils, "validate", lambda *args: False) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - - assert backend_class is TritonLinearAttBackend - - -def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_backend_args): - validate_calls = [] - monkeypatch.setattr( - linear_create_utils, - "gdn_prefill_backends", - {"flashqla2": TritonLinearAttBackend, "flashqla3": FlashQlaLinearAttBackend}, - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name, *args: validate_calls.append(name) or name == "flashqla3", - ) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class( - index=1, priority_list=("flashqla2", "flashqla3") - ) - - assert backend_class is FlashQlaLinearAttBackend - assert validate_calls == ["flashqla2", "flashqla3"] - - -def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None, role="main"): - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), - ) - infer_state = InferStateInfo() - infer_state.batch_size = b_req_idx.shape[0] - infer_state.b_req_idx = b_req_idx - infer_state.b_mtp_index = torch.zeros_like(b_req_idx) - infer_state.req_manager = SimpleNamespace( - HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, - ) - infer_state.att_state_role = role - if req_to_mtp_state_index is not None: - request_capacity = req_to_mtp_state_index.numel() - resources = AttentionStateManager(request_capacity=request_capacity) - config = ConvSSMStateConfig( - logical_layer_ids=(3,), - global_k_heads=1, - global_v_heads=1, - key_dim=2, - value_dim=2, - conv_width=2, - conv_dtype=torch.float32, - ssm_dtype=torch.float32, - ) - for model_role in ("main", "draft0"): - linear = GDNState(config, request_capacity=request_capacity, mtp_step=draft_step, device="meta") - linear.mtp_state_index = ( - req_to_mtp_state_index if model_role == role else torch.zeros_like(req_to_mtp_state_index) - ) - resources.register_state(model_role, linear) - resources.finish_state_registration() - infer_state.att_state_manager = resources - state = gdn.LinearAttDecodeAttState( - backend=SimpleNamespace( - model=model, - uses_dynamic_spec_verify_layout=lambda: dynamic_layout, - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - return state, infer_state diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py index ae848f13de..ef2014cd3a 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -1,17 +1,30 @@ +from lightllm.common.attention.api import AttentionFamily, AttentionPhase, CacheEncoding, CacheLayout +from unit_tests.common.attention.helpers import make_env import gc import os -from pathlib import Path import subprocess import sys +from pathlib import Path from types import SimpleNamespace import pytest import torch -pytest.importorskip("flash_mla") +from unit_tests.common.attention.helpers import make_spec + +flashmla = pytest.importorskip("flash_mla") -from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention from lightllm.common.basemodel.layer_weights.meta_weights import TpAttSinkWeight +from lightllm.models.deepseek_v4.attention import sparse_mla as attention +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -35,8 +48,9 @@ def find_spec(self, fullname, path=None, target=None): raise ModuleNotFoundError('vllm deliberately unavailable') sys.meta_path.insert(0, BlockVllm()) -from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention -assert attention.flashmla.__name__ == 'flash_mla' +from lightllm.models.deepseek_v4.attention import sparse_mla as attention +import flash_mla as flashmla +assert flashmla.__name__ == 'flash_mla' assert not any(name == 'vllm' or name.startswith('vllm.') for name in sys.modules) """ root = Path(__file__).resolve().parents[4] @@ -47,7 +61,8 @@ def find_spec(self, fullname, path=None, target=None): @pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode"]) def test_swa_workspace_includes_decode_padding_prefill(monkeypatch: pytest.MonkeyPatch, run_mode: str) -> None: - backend_class = attention.DeepseekV4FlashMlaFp8SparseAttBackend + from lightllm.models.deepseek_v4.workspace import get_swa_prefill_workspace + model = SimpleNamespace( args=SimpleNamespace(run_mode=run_mode, enable_prefill_microbatch_overlap=False), config=dict( @@ -64,12 +79,7 @@ def test_swa_workspace_includes_decode_padding_prefill(monkeypatch: pytest.Monke max_req_num=8, att_state_role="main", ) - monkeypatch.setattr( - backend_class, - "get_gpu_workspace_buffer", - lambda key_name, workspace_size, dtype: torch.empty(workspace_size, dtype=dtype), - ) - workspace, capacity = backend_class.get_swa_prefill_workspace(model=model) + workspace, capacity = get_swa_prefill_workspace(model=model) # D 节点仍执行 HOLD 请求的单 token prefill，但无需保留普通请求的长 prefill 临时区。 assert workspace is not None if run_mode == "decode": @@ -96,14 +106,25 @@ def _pool(page_size, pages): def _case(heads, ratio): torch.manual_seed(6014) tokens = 5 - model = type("TestModel", (), {})() + from unit_tests.common.attention.helpers import model_stub + + model = model_stub(page_size=256) model.config = {"num_attention_heads": heads, "compress_ratios": [0, 4, 128]} model.config.update(head_dim=512, qk_rope_head_dim=64, index_head_dim=128, sliding_window=128) model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] model.tp_rank_ = 0 model.tp_world_size_ = 1 - backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend(model=model) - q = torch.randn((tokens, backend.padded_q_head_num, 512), dtype=torch.bfloat16, device="cuda") * 0.1 + spec = make_spec( + AttentionFamily.SPARSE_MLA, + cache=CacheLayout(CacheEncoding.FP8_SPARSE_MLA, 256, torch.uint8, 1, 512, 512), + num_q_heads=heads, + head_size=512, + value_head_size=512, + variant="deepseek_v4", + softmax_scale=512 ** -0.5, + ) + backend = attention.DeepseekV4SparseMLAFp8DecodeBackend(env=make_env(model), spec=spec) + q = torch.randn((tokens, (heads + 63) // 64 * 64, 512), dtype=torch.bfloat16, device="cuda") * 0.1 q[:, heads:] = 0 pool, decoded = _pool(128, 2) extra, extra_decoded = _pool(64 if ratio == 4 else 2, 3) if ratio else (None, None) @@ -154,8 +175,8 @@ def _reference(q, decoded, extra_decoded, params, heads): @pytest.mark.parametrize("provide_out", [False, True]) def test_prefill_matches_reference_and_preserves_out(monkeypatch, heads, ratio, provide_out): backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, ratio) - state = backend.create_att_prefill_state(state_info) - state.init_state() + state = backend_context(type(backend), "prefill", settings=backend) + update_context(state, state_info) allocations = [] def alloc_output(shape, **kwargs): @@ -163,23 +184,16 @@ def alloc_output(shape, **kwargs): allocations.append(buffer) return buffer - kernel = attention.flashmla.flash_mla_with_kvcache + kernel = flashmla.flash_mla_with_kvcache def flashmla_with_out(**kwargs): assert kwargs["out"] is not None return kernel(**kwargs) - monkeypatch.setattr(attention.flashmla, "flash_mla_with_kvcache", flashmla_with_out) + monkeypatch.setattr(flashmla, "flash_mla_with_kvcache", flashmla_with_out) out = torch.full_like(q, float("nan")) if provide_out else None - actual = state.prefill_att( - q, - pool, - None, - att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params), - alloc_func=alloc_output, - out=out, - ) - assert params["attn_sink"].shape == (backend.padded_q_head_num,) + actual = run_sparse(state, state_info, q, pool, params, alloc_func=alloc_output, out=out) + assert params["attn_sink"].shape == (q.shape[1],) assert torch.count_nonzero(params["attn_sink"][heads:]) == 0 assert len(allocations) == (0 if provide_out else 1) output_buffer = out if provide_out else allocations[0] @@ -216,15 +230,24 @@ def test_decode_cuda_graph_reads_updated_query_and_lengths(monkeypatch, heads, r state_info.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) state_info.q = q state_info.packed_kv = packed_kv - state_info.att_control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) - state_info.decode_att_state = backend.create_att_decode_state(state_info) - state_info.decode_att_state1 = None - state_info.decode_att_state.init_state() + state_info.att_control = params + from lightllm.common.attention.session import AttentionSession + + state_info.attention = AttentionSession() + state_info.attention.contexts = {0: backend_context(type(backend), "decode", settings=backend)} + state_info.attention.groups = ({AttentionPhase.DECODE: backend},) + update_context(state_info.attention.contexts[0], state_info) + state_info.b_req_idx = torch.arange(q.shape[0], dtype=torch.int32, device=q.device) + from lightllm.common.attention_state import AttentionStateManager + + state_info.att_state_manager = AttentionStateManager(request_capacity=q.shape[0]) + state_info.att_state_manager.finish_state_registration() + state_info.att_state_role = "main" cases.append((state_info, decoded, extra_decoded, params)) def decode_func(*states): outputs = tuple( - state.decode_att_state.decode_att(state.q, state.packed_kv, None, att_control=state.att_control) + run_sparse(state.attention.contexts[0], state, state.q, state.packed_kv, state.att_control) for state in states ) return outputs if overlap else outputs[0] @@ -233,7 +256,7 @@ def decode_func(*states): outputs = runner.capture_decode(decode_func, *(case[0] for case in cases)) outputs = outputs if overlap else (outputs,) for state_info, _, _, _ in cases: - state = state_info.decode_att_state + state = state_info.attention.contexts[0] eager_meta = state.sched_meta_by_ratio[ratio] capture_meta = state.capture_sched_meta_by_ratio[ratio] assert eager_meta.tile_scheduler_metadata.data_ptr() != capture_meta.tile_scheduler_metadata.data_ptr() @@ -248,11 +271,11 @@ def decode_func(*states): params["swa_lengths"].copy_(new_lengths) if ratio: params["extra_lengths"].copy_(new_lengths) - runner.graph[5][0].replay() + runner.graph[5].replay() for actual, (state_info, decoded, extra_decoded, params) in zip(outputs, cases): assert actual.shape == (state_info.q.shape[0], heads, 512) - assert actual.stride(0) == backend.padded_q_head_num * 512 + assert actual.stride(0) == state_info.q.shape[1] * 512 torch.testing.assert_close( actual, _reference(state_info.q, decoded, extra_decoded, params, heads), rtol=0.02, atol=5e-4 ) @@ -262,13 +285,11 @@ def decode_func(*states): @pytest.mark.parametrize("heads", [16, 64]) def test_prefill_out_supports_strided_scratch_buffer(heads): backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, 4) - storage = torch.full((q.shape[0], backend.padded_q_head_num, 1024), float("nan"), dtype=q.dtype, device=q.device) + storage = torch.full((q.shape[0], q.shape[1], 1024), float("nan"), dtype=q.dtype, device=q.device) out = storage[..., :512] - state = backend.create_att_prefill_state(state_info) - state.init_state() - actual = state.prefill_att( - q, pool, None, out=out, att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params) - ) + state = backend_context(type(backend), "prefill", settings=backend) + update_context(state, state_info) + actual = run_sparse(state, state_info, q, pool, params, out=out) assert actual.data_ptr() == out.data_ptr() assert actual.stride() == out.stride() assert storage[..., 512:].isnan().all() @@ -292,6 +313,9 @@ def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): ) monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) runner = cuda_graph_module.PrefillCudaGraph(decode_cuda_graph=None, tp_world_size=1) + from lightllm.common.attention_state import AttentionStateManager + + manager = AttentionStateManager(request_capacity=5) backend, _, q, packed_kv, decoded, extra_decoded, params = _case(heads, ratio) indices = torch.arange(128, dtype=torch.int32, device=q.device) params["swa_indices"].copy_(indices) @@ -301,17 +325,16 @@ def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): packed_kv = (packed_kv[0], extra) params["extra_indices"].copy_(indices) - monkeypatch.setattr(backend, "get_att_input_params", lambda infer_state, logical_layer_id: packed_kv) - def new_request(length): infer_state = DeepseekV4InferStateInfo() + infer_state.att_state_manager = manager infer_state.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) infer_state.hidden_collector = NoopHiddenCollector() infer_state.attn_params = dict(params, swa_lengths=torch.full_like(params["swa_lengths"], length)) if ratio: infer_state.attn_params["extra_lengths"] = torch.full_like(params["extra_lengths"], length) - infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) - infer_state.prefill_att_state.init_state() + infer_state.attention.contexts = {0: backend_context(type(backend), "prefill", settings=backend)} + update_context(infer_state.attention.contexts[0], infer_state) return infer_state allocations = [] @@ -336,6 +359,17 @@ def compress_and_index(q_lora, x, infer_state, layer_weight): return infer_state.attn_params layer.alloc_tensor = alloc_output + from dataclasses import replace + + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4Resources + from unit_tests.common.attention.helpers import make_layer_attention + + attention_params = DeepseekV4Resources( + packed_kv[0], packed_kv[0], packed_kv[1], ratio, 512, params["softmax_scale"], params["attn_sink"] + ) + attention_context = backend_context(type(backend), "prefill", settings=backend) + layer.prefill_attention = make_layer_attention(backend, attention_context, attention_params) + backend.alloc_tensor = alloc_output layer._compress_and_index = compress_and_index q_lora = torch.empty((q.shape[0], 4), dtype=q.dtype, device=q.device) x = torch.empty((q.shape[0], 8), dtype=q.dtype, device=q.device) @@ -345,7 +379,7 @@ def prefill_func(inputs, infer_state): att_out = layer._context_attention_wrapper_run(graph_q, inputs[1], inputs[2], infer_state, layer_weight) assert att_out.shape == (q.shape[0], heads, 512) assert att_out.data_ptr() == allocations[-1].data_ptr() - assert att_out.stride(0) == backend.padded_q_head_num * 512 + assert att_out.stride(0) == q.shape[1] * 512 # A graph segment after attention must read the output filled by the CPU callback. return [att_out + 0.125] @@ -359,11 +393,11 @@ def prefill_func(inputs, infer_state): atol=5e-4, ) assert len(infer_state.prefill_cuda_graph_exe_list) == 3 - previous_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] + previous_meta = infer_state.attention.contexts[0].sched_meta_by_ratio[ratio] for length in (128, 64, 1, 64): q.mul_(1.25) infer_state = new_request(length) - current_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] + current_meta = infer_state.attention.contexts[0].sched_meta_by_ratio[ratio] assert current_meta.tile_scheduler_metadata is None actual = runner.replay([q, q_lora, x], infer_state)[0] assert callback_states[-1] is infer_state diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py index 9d02784676..79463bd2dc 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -1,20 +1,20 @@ -"""FP8 后端按推理层传入的逻辑层号选择标定，覆盖稀疏层号、共享层和 draft。""" - from dataclasses import replace from types import SimpleNamespace -import weakref import pytest import torch +from lightllm.common.attention.api import AttentionFamily, AttentionPhase +from lightllm.common.attention.backends.dense.fa3 import fp8 as fa3_fp8 +from lightllm.common.attention.backends.dense.flashinfer import fp8 as flashinfer_fp8 from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 -from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.common.basemodel.attention_adapters import ADAPTERS +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from unit_tests.common.att_resources.model_helpers import rebuild_full_state -from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.model_helpers import rebuild_full_state, replace_logical_layer_groups from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention @pytest.mark.parametrize("backend_name", ["fa3", "flashinfer"]) @@ -42,75 +42,67 @@ def test_fp8_backend_selects_calibration_from_logical_layer( manager = AttentionStateManager(request_capacity=1) manager.register_state(role, state) manager.finish_state_registration() - infer = LlamaInferStateInfo() - infer.att_state_manager = manager - infer.att_state_role = role - infer.b_seq_len = torch.tensor([1], dtype=torch.int32) - infer.max_q_seq_len = 1 + phase = AttentionPhase(operation) + cls = fa3_fp8.Fa3DenseFp8DecodeBackend if backend_name == "fa3" else flashinfer_fp8.FlashInferDenseFp8DecodeBackend + backend = make_backend(cls) + context = make_context(backend, make_batch((1,), phase=phase)) + if backend_name == "fa3" and operation == "decode": + context.b_att_seq_len = context.batch.b_seq_len + context.decode_max_q_seq_len = 1 + context.causal = True calls = [] - - # 仅替换外部计算内核，保留推理层、公开 attention 接口及 State 标定查询链路。 if backend_name == "fa3": state.q_scales[0].fill_(3) state.q_scales[1].fill_(7) expected_q_scale = 3 if storage_layer_index == 0 else 7 - def quantize(q: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + def quantize(q, scales): assert torch.all(scales == expected_q_scale) return q - def attend(**kwargs: object) -> torch.Tensor: + def attend(**kwargs): assert torch.all(kwargs["k_descale"] == expected_scale) assert torch.all(kwargs["v_descale"] == expected_scale) assert torch.all(kwargs["q_descale"] == expected_q_scale) calls.append(operation) return kwargs["q"] - monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", quantize) - monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) - state_class = fa3_fp8.Fp8Fa3PrefillAttState if operation == "prefill" else fa3_fp8.Fp8Fa3DecodeAttState - else: - state_class = ( - flashinfer_fp8.Fp8FlashInferPrefillAttState - if operation == "prefill" - else flashinfer_fp8.Fp8FlashInferDecodeAttState + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant.q_per_head_static_fp8_quant", + quantize, ) - - attention = state_class(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer)) - if backend_name == "fa3": - attention.page_table = attention.cu_seqlens_q = attention.cu_seqlens_k = None - attention.b_att_seq_len = infer.b_seq_len - attention.decode_max_q_seq_len = 1 - attention.causal = True + monkeypatch.setattr("lightllm.utils.sgl_utils.flash_attn_with_kvcache", attend) else: + state.cpu_scales = state.scales.cpu().tolist() - def run( - q: torch.Tensor, - paged_kv_cache: tuple[torch.Tensor, torch.Tensor], - k_scale: float, - v_scale: float, - out: torch.Tensor, - ) -> None: + def run(q, paged_kv_cache, k_scale, v_scale, out): assert k_scale == v_scale == expected_scale calls.append(operation) out.copy_(q) - attention.prefill_wrapper = attention.decode_wrapper = SimpleNamespace(run=run) - - infer.prefill_att_state = infer.decode_att_state = attention + setattr(context, operation + "_wrapper", SimpleNamespace(run=run)) layer = LlamaTransformerLayerInfer.__new__(LlamaTransformerLayerInfer) layer.layer_num_ = logical_layer_id layer.tp_q_head_num_ = layer.tp_k_head_num_ = state.config.head_num layer.head_dim_ = state.config.head_dim - layer.alloc_tensor = lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device="cpu") + model = SimpleNamespace( + att_state_manager=manager, + att_state_role=role, + args=StartArgs(llm_kv_type="fp8kv_sph" if backend_name == "fa3" else "fp8kv_spt"), + ) + params = ADAPTERS[AttentionFamily.DENSE].bind(model, layer, None) + attention = make_layer_attention(backend, context, params) + attention = replace( + attention, + logical_layer_id=logical_layer_id, + ) + setattr(layer, operation + "_attention", attention) + infer = InferStateInfo() + infer.attention.contexts = {0: context} + monkeypatch.setattr(manager, "get_state", lambda *args: pytest.fail("forward must not look up storage owners")) q = torch.ones((1, state.config.head_num, state.config.head_dim)) if operation == "prefill": - output = layer._context_attention_kernel( - q=q, - kv=None, - infer_state=infer, - layer_weight=None, - ) + output = layer._context_attention_kernel(q=q, kv=None, infer_state=infer, layer_weight=None) else: output = layer._token_attention_kernel(q=q, infer_state=infer, layer_weight=None) assert torch.equal(output, q) diff --git a/unit_tests/common/basemodel/attention/test_fp8_page_size.py b/unit_tests/common/basemodel/attention/test_fp8_page_size.py index 107728c26d..63fc7641c8 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_page_size.py +++ b/unit_tests/common/basemodel/attention/test_fp8_page_size.py @@ -1,45 +1,37 @@ -from types import SimpleNamespace -from unittest.mock import patch - +from dataclasses import replace +from unit_tests.common.attention.helpers import make_spec import pytest -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 -from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend +from lightllm.common.attention.backends.dense.fa3 import fp8 as fa3_fp8 +from lightllm.common.attention.backends.dense.fa3.backend import Fa3DenseDecodeBackend +from lightllm.common.attention.backends.dense.fa3.fp8 import Fa3DenseFp8DecodeBackend +from lightllm.common.attention.backends.dense.flashinfer import fp8 as flashinfer_fp8 +from lightllm.common.attention.backends.dense.flashinfer.backend import FlashInferDenseDecodeBackend +from lightllm.common.attention.backends.dense.flashinfer.fp8 import FlashInferDenseFp8DecodeBackend +from lightllm.common.attention.backends.mla.fa3.backend import Fa3MLADecodeBackend +from lightllm.common.attention.backends.mla.flashinfer.backend import FlashInferMLADecodeBackend @pytest.mark.parametrize( "backend_class", - [Fa3AttBackend, MlaFa3AttBackend, FlashInferAttBackend, MlaFlashInferAttBackend], + [Fa3DenseDecodeBackend, Fa3MLADecodeBackend, FlashInferDenseDecodeBackend, FlashInferMLADecodeBackend], ) -def test_default_infer_page_size_matches_model_page_size(backend_class): +def test_default_kernel_page_size_matches_model_page_size(backend_class): backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - backend._init_infer_page_size() + backend.spec = make_spec(cache=replace(make_spec().cache, page_size=16)) - assert backend.infer_page_size == 16 + assert backend.kernel_page_size == 16 @pytest.mark.parametrize( ("backend_class", "backend_module"), [ - (Fp8Fa3AttBackend, fa3_fp8), - (Fp8FlashInferAttBackend, flashinfer_fp8), + (Fa3DenseFp8DecodeBackend, fa3_fp8), + (FlashInferDenseFp8DecodeBackend, flashinfer_fp8), ], ) def test_fp8_attention_backend_uses_single_token_infer_pages(backend_class, backend_module): backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - with patch.object(backend_module.logger, "warning") as warning: - backend._init_infer_page_size() + backend.spec = make_spec(cache=replace(make_spec().cache, page_size=16)) - assert backend.infer_page_size == 1 - warning.assert_called_once() - assert "infer_page_size=1" in warning.call_args.args[0] + assert backend.kernel_page_size == 1 diff --git a/unit_tests/common/basemodel/attention/test_gated_qkv.py b/unit_tests/common/basemodel/attention/test_gated_qkv.py new file mode 100644 index 0000000000..b59325b8f4 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_gated_qkv.py @@ -0,0 +1,23 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen3next.layer_infer import transformer_layer_infer as module + + +def test_qwen_next_qkv_split_matches_projection_with_separate_gate(monkeypatch): + layer = module.Qwen3NextTransformerLayerInfer.__new__(module.Qwen3NextTransformerLayerInfer) + layer.tp_q_head_num_, layer.tp_k_head_num_, layer.tp_v_head_num_ = 4, 2, 2 + layer.head_dim_, layer.embed_dim_, layer.eps_, layer.partial_rotary_factor = 8, 32, 1e-6, 0.25 + layer._tpsp_allgather = lambda input, infer_state: input + projected = torch.arange(3 * 96, dtype=torch.float32).view(3, 96) + weight = SimpleNamespace( + qkvo_gate_proj=SimpleNamespace(mm=lambda x: projected), + qk_norm_weight_=lambda *args, **kwargs: None, + ) + infer = SimpleNamespace(position_cos=None, position_sin=None, need_dp_prefill_balance=False) + monkeypatch.setattr(module, "rotary_emb_fwd", lambda *args, **kwargs: None) + query, kv = layer._get_qkv(torch.zeros(3, 32), infer, weight) + torch.testing.assert_close(query, projected[:, :32]) + torch.testing.assert_close(kv, projected[:, 32:64].view(3, 4, 8)) + torch.testing.assert_close(infer.gate_logics_value, projected[:, 64:]) diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py index e7423d30f9..e85106f754 100644 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -1,248 +1,15 @@ -"""Check import isolation and selection without loading optional GPU dependencies.""" - -import importlib -import importlib.abc -import importlib.util -from pathlib import Path +import os +import subprocess import sys -from types import ModuleType, SimpleNamespace -import unittest -from unittest.mock import Mock, patch - - -ROOT = Path(__file__).resolve().parents[4] -ATTENTION = "lightllm.common.basemodel.attention" - - -class BackendImports(importlib.abc.MetaPathFinder, importlib.abc.Loader): - """Reject unselected implementations, and supply lightweight selected backends.""" - - def __init__(self, package=ATTENTION): - self.package = package - self.allowed = {} - self.imported = [] - - def allow(self, suffix, class_name, **attributes): - module = ModuleType(f"{self.package}.{suffix}") - setattr(module, class_name, type(class_name, (), {})) - module.__dict__.update(attributes) - self.allowed[module.__name__] = module - return getattr(module, class_name) - - def find_spec(self, fullname, path=None, target=None): - suffix = fullname.removeprefix(f"{self.package}.") - if not fullname.startswith(f"{self.package}.") or "." not in suffix: - return None - if suffix.split(".")[0] not in {"triton", "fa3", "flashinfer", "nsa", "linear", "sdpa", "xformers"}: - return None - self.imported.append(fullname) - if fullname not in self.allowed: - raise AssertionError(f"Unexpected attention backend import: {fullname}") - return importlib.util.spec_from_loader(fullname, self) - - def create_module(self, spec): - return self.allowed[spec.name] - - def exec_module(self, module): - pass - - -class TestLazyAttentionImports(unittest.TestCase): - def setUp(self): - # Isolate attention from basemodel's unrelated GPU initialization. - modules_patch = patch.dict(sys.modules) - modules_patch.start() - self.addCleanup(modules_patch.stop) - for name in list(sys.modules): - if name == ATTENTION or name.startswith(f"{ATTENTION}."): - del sys.modules[name] - for name in ("lightllm", "lightllm.common", "lightllm.common.basemodel", "lightllm.utils"): - module = ModuleType(name) - module.__path__ = [str(ROOT.joinpath(*name.split(".")))] - sys.modules[name] = module - - self.args = SimpleNamespace( - llm_kv_type="None", - llm_prefill_att_backend=["triton"], - llm_decode_att_backend=["triton"], - enable_ep_moe=False, - mtp_step=0, - ssm_state_mode="native", - ) - self.validate = Mock(return_value=True) - dependencies = { - f"{ATTENTION}.base_att": { - name: type(name, (), {}) - for name in ("BaseAttBackend", "BasePrefillAttState", "BaseDecodeAttState", "AttControl") - }, - "lightllm.utils.envs_utils": {"get_env_start_args": lambda: self.args}, - "lightllm.utils.backend_validator": {"validate": self.validate}, - "lightllm.utils.log_utils": {"init_logger": lambda name: Mock()}, - } - for name, attributes in dependencies.items(): - module = ModuleType(name) - module.__dict__.update(attributes) - sys.modules[name] = module - - self.imports = BackendImports() - sys.meta_path.insert(0, self.imports) - self.addCleanup(sys.meta_path.remove, self.imports) - self.attention = importlib.import_module(ATTENTION) - - def test_importing_selection_apis_loads_no_implementations(self): - importlib.import_module(f"{ATTENTION}.linear") - importlib.import_module(f"{ATTENTION}.nsa") - self.assertEqual(self.imports.imported, []) - - def test_explicit_regular_selection_skips_neo(self): - backend = self.imports.allow("fa3.fp", "Fa3AttBackend") - self.args.llm_prefill_att_backend = ["fa3"] - - self.assertIs(self.attention.get_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.fp"]) - self.validate.assert_not_called() - - def test_auto_selection_does_not_import_failed_candidates(self): - backend = self.imports.allow("flashinfer.fp", "FlashInferAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - self.validate.side_effect = lambda name: name == "flashinfer" - - self.assertIs(self.attention.get_prefill_att_backend_class(), backend) - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "flashinfer"]) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.flashinfer.fp"]) - - def test_mtp_decode_preserves_fa3_priority(self): - backend = self.imports.allow("fa3.fp", "Fa3AttBackend") - self.args.llm_decode_att_backend = ["auto"] - self.args.mtp_step = 2 - - self.assertIs(self.attention.get_decode_att_backend_class(), backend) - self.validate.assert_called_once_with("fa3") - - def test_expert_parallelism_still_excludes_flashinfer(self): - backend = self.imports.allow("triton.fp", "TritonAttBackend") - self.args.llm_decode_att_backend = ["auto"] - self.args.enable_ep_moe = True - self.validate.side_effect = lambda name: name == "triton" - self.assertIs(self.attention.get_decode_att_backend_class(), backend) - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "triton"]) - def test_auto_selection_falls_back_to_triton(self): - backend = self.imports.allow("triton.mla", "MlaTritonAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - self.validate.return_value = False - - self.assertIs(self.attention.get_mla_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.mla"]) - - def test_nsa_selection_does_not_import_other_kv_types(self): - backend = self.imports.allow("nsa.fp8_flashmla_sparse", "NsaFlashMlaFp8SparseAttBackend") - self.args.llm_kv_type = "fp8kv_dsa" - self.args.llm_decode_att_backend = ["flashmla_sparse"] - - self.assertIs(self.attention.get_nsa_decode_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.fp8_flashmla_sparse"]) - - def test_custom_backend_map_uses_import_paths(self): - backend = self.imports.allow("nsa.custom", "CustomBackend") - backend_map = {"None": {"custom": "nsa.custom:CustomBackend"}} - self.args.llm_prefill_att_backend = ["custom"] - self.args.llm_decode_att_backend = ["auto"] - - self.assertIs(self.attention.get_nsa_prefill_att_backend_class(backend_map=backend_map), backend) - self.assertIs( - self.attention.get_nsa_decode_att_backend_class(priority_list=["custom"], backend_map=backend_map), - backend, - ) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.custom"]) - - def test_neo_triton_selection_does_not_probe_fa3(self): - backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.args.llm_prefill_att_backend = ["auto"] - self.assertIs(self.attention.get_neo_prefill_att_backend_class(priority_list=["triton"]), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.neo"]) - - def test_neo_auto_selection_still_falls_back_without_image_token_interface(self): - self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=False) - backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.validate.assert_called_once_with("triton") - self.args.llm_prefill_att_backend = ["fa3"] - with self.assertRaises(KeyError): - self.attention.get_neo_prefill_att_backend_class() - - def test_neo_auto_selection_uses_supported_fa3(self): - backend = self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=True) - self.args.llm_prefill_att_backend = ["auto"] - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) - self.validate.assert_called_once_with("fa3") - - def _allow_linear_backend_classes(self): - flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") - triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") - self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") - replay_gdn = self.imports.allowed[f"{ATTENTION}.linear.replay_gdn"] - replay_gdn.ReplayGDNFlashQlaAttBackend = type("ReplayGDNFlashQlaAttBackend", (), {}) - self.imports.allow("linear.kda", "KDALinearAttBackend") - self.imports.allow("linear.replay_kda", "ReplayKDALinearAttBackend") - return flashqla, triton - - def test_linear_prefill_auto_selects_validated_flashqla(self): - flashqla, _triton = self._allow_linear_backend_classes() - selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), flashqla) - self.validate.assert_called_once_with("flashqla") - - def test_linear_selection_falls_back_to_triton(self): - _flashqla, triton = self._allow_linear_backend_classes() - selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - self.validate.return_value = False - - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) - self.assertIs(selection.get_gdn_linear_decode_att_backend_class(), triton) - self.args.llm_prefill_att_backend.append("triton") - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) - - def test_vit_selection_loads_only_requested_implementation(self): - package = "lightllm.common.basemodel.attention_vit" - for name in list(sys.modules): - if name == package or name.startswith(f"{package}."): - del sys.modules[name] - base = ModuleType(f"{package}.base_att") - base.BaseVitAttBackend = type("BaseVitAttBackend", (), {}) - sys.modules[base.__name__] = base - sys.modules["lightllm.utils.backend_validator"]._validate = self.validate - imports = BackendImports(package) - sys.meta_path.insert(0, imports) - self.addCleanup(sys.meta_path.remove, imports) - selection = importlib.import_module(f"{package}.create_utils") - self.assertEqual(imports.imported, []) - - for name, class_name in ( - ("fa3", "Fa3VitAttBackend"), - ("triton", "TritonVitAttBackend"), - ("sdpa", "SdpaVitAttBackend"), - ("xformers", "XformersVitAttBackend"), - ): - backend = imports.allow(f"{name}.fp", class_name) - self.args.vit_att_backend = [name] - self.assertEqual(selection.init_vit_att_backend(), name) - self.assertIs(selection.get_vit_att_backend_class(name), backend) - self.validate.assert_not_called() - - self.args.vit_att_backend = ["auto"] - self.validate.side_effect = lambda name: name == "sdpa" - self.assertEqual(selection.init_vit_att_backend(), "sdpa") - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "xformers", "sdpa"]) - - -if __name__ == "__main__": - unittest.main() +def test_api_and_selection_do_not_import_provider_extensions(): + code = """ +import sys +import lightllm.common.attention.api +import lightllm.common.attention.selection +assert not any(name.startswith(('flashinfer', 'flash_mla', 'flash_qla')) for name in sys.modules) +assert not any(name.startswith('lightllm.common.attention.backends.') for name in sys.modules) +""" + result = subprocess.run([sys.executable, "-c", code], env=os.environ.copy(), capture_output=True, text=True) + assert result.returncode == 0, result.stderr diff --git a/unit_tests/common/basemodel/attention/test_state_lifetime.py b/unit_tests/common/basemodel/attention/test_state_lifetime.py index 9e049d54b4..74e7e4512c 100644 --- a/unit_tests/common/basemodel/attention/test_state_lifetime.py +++ b/unit_tests/common/basemodel/attention/test_state_lifetime.py @@ -1,174 +1,34 @@ +from unit_tests.common.attention.helpers import plan import gc import weakref -import pytest -import torch +from unit_tests.common.attention.helpers import StubBackend, make_env +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_spec, model_stub -from lightllm.common.basemodel.attention.base_att import BaseDecodeAttState, BasePrefillAttState -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.basemodel.infer_struct import InferStateInfo - -class _PrefillState(BasePrefillAttState): - def init_state(self): - self.positions = self.infer_state.position_ids - - def prefill_att(self, q, k, v, **kwargs): - return q - - -class _DecodeState(BaseDecodeAttState): - def init_state(self): - self.positions = self.infer_state.position_ids - - def decode_att(self, q, k, v, **kwargs): - return q - - -@pytest.fixture -def no_cycle_gc(): - enabled = gc.isenabled() - gc.collect() +def test_context_does_not_retain_runner_or_model(): + model = model_stub() + backend = StubBackend(make_env(model), make_spec()) + assert not any(value is model for value in vars(backend).values()) + context = make_context(backend) gc.disable() try: - yield + batch_ref = weakref.ref(context.batch) + del context + assert batch_ref() is None finally: - gc.collect() - if enabled: - gc.enable() - - -def _create_state(is_prefill, secondary=False): - owner = InferStateInfo() - owner.is_prefill = is_prefill - owner.position_ids = torch.arange(4) - state_cls = _PrefillState if is_prefill else _DecodeState - slot = "prefill_att_state" if is_prefill else "decode_att_state" - setattr(owner, slot, state_cls(_infer_state_weak_ref=weakref.ref(owner))) - if secondary: - setattr(owner, slot + "1", state_cls(_infer_state_weak_ref=weakref.ref(owner))) - owner.init_att_state() - return owner - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("secondary", [False, True]) -def test_forward_states_release_without_cycle_gc(no_cycle_gc, is_prefill, secondary): - owner = _create_state(is_prefill, secondary) - states = [state for state in vars(owner).values() if isinstance(state, (BasePrefillAttState, BaseDecodeAttState))] - owner_ref = weakref.ref(owner) - tensor_ref = weakref.ref(owner.position_ids) - state_refs = [weakref.ref(state) for state in states] - - # Even an external reference to attention metadata must not retain its owner. - del owner - assert owner_ref() is None - with pytest.raises(ReferenceError): - states[0].infer_state - assert tensor_ref() is not None - del states - assert all(ref() is None for ref in state_refs) - assert tensor_ref() is None - - -@pytest.mark.parametrize("state_cls", [_PrefillState, _DecodeState]) -def test_attention_getter_uses_a_weak_reference_and_is_read_only(state_cls): - state = state_cls() - assert state.infer_state is None - - owner = InferStateInfo() - state = state_cls(_infer_state_weak_ref=weakref.ref(owner)) - assert isinstance(state._infer_state_weak_ref, weakref.ReferenceType) - assert state.infer_state is owner - state.check_no_infer_state() - with pytest.raises(AttributeError): - state.infer_state = owner - + gc.enable() -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("slot_suffix", ["", "1"]) -def test_check_rejects_infer_state_under_any_attribute_name(is_prefill, slot_suffix): - owner = _create_state(is_prefill, secondary=True) - slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix - attention = getattr(owner, slot) - - class _OtherInferState(InferStateInfo): - pass - - # Check the type, rather than just the identity of the expected owner. - attention.saved_context = _OtherInferState() - with pytest.raises(RuntimeError, match="saved_context"): - owner.check_att_states() - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("slot_suffix", ["", "1"]) -def test_init_checks_for_accidental_back_reference(is_prefill, slot_suffix): - owner = _create_state(is_prefill, secondary=True) - slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix - base = _PrefillState if is_prefill else _DecodeState - - class _RetainingState(base): - def init_state(self): - self.saved_context = self.infer_state - - setattr(owner, slot, _RetainingState(_infer_state_weak_ref=weakref.ref(owner))) - try: - with pytest.raises(RuntimeError, match="saved_context"): - owner.init_att_state() - finally: - del getattr(owner, slot).saved_context - - -def test_decode_graph_replay_preserves_capture_buffers_and_owner_lifetime(no_cycle_gc): - captured = _create_state(is_prefill=False, secondary=True) - captured.input_ids = torch.zeros(4, dtype=torch.int64) - captured_ref = weakref.ref(captured) - capture_positions_ptr = captured.position_ids.data_ptr() - replay_calls = [] - - class _Graph: - def replay(self): - replay_calls.append(True) - - graph = CudaGraph.__new__(CudaGraph) - output = object() - graph.graph = {4: (_Graph(), captured, output)} - del captured - assert captured_ref() is not None - - runtime = _create_state(is_prefill=False, secondary=True) - runtime.input_ids = torch.ones(4, dtype=torch.int64) - runtime.position_ids.add_(7) - runtime_ref = weakref.ref(runtime) - assert graph._replay(runtime) is output - assert replay_calls == [True] - assert captured_ref().position_ids.data_ptr() == capture_positions_ptr - assert captured_ref().decode_att_state.infer_state is captured_ref() - assert captured_ref().decode_att_state1.infer_state is captured_ref() - torch.testing.assert_close(captured_ref().position_ids, runtime.position_ids) +def test_captured_context_keeps_buffers_while_transient_batch_dies(): + backend = make_backend() + context = make_context(backend, capture=True, max_length=64) + pointer = context.batch.b_seq_len.data_ptr() + runtime = make_batch((17, 3), req_to_token_indexs=context.batch.req_to_token_indexs) + ref = weakref.ref(runtime) + plan(backend, context, runtime) del runtime - assert runtime_ref() is None - graph.graph.clear() - assert captured_ref() is None - - -@pytest.mark.parametrize("is_prefill", [False, True]) -def test_secondary_state_without_primary(is_prefill): - owner = InferStateInfo() - owner.is_prefill = is_prefill - owner.position_ids = torch.arange(4) - slot = "prefill_att_state1" if is_prefill else "decode_att_state1" - state_cls = _PrefillState if is_prefill else _DecodeState - secondary = state_cls(_infer_state_weak_ref=weakref.ref(owner)) - setattr(owner, slot, secondary) - owner.init_att_state() - torch.testing.assert_close(secondary.positions, owner.position_ids) - if not is_prefill: - runtime = InferStateInfo() - runtime.position_ids = owner.position_ids + 7 - runtime.decode_att_state1 = _DecodeState(_infer_state_weak_ref=weakref.ref(runtime)) - runtime.decode_att_state1.init_state() - owner.copy_for_cuda_graph(runtime) - torch.testing.assert_close(secondary.positions, runtime.position_ids) + assert ref() is None + assert context.batch.b_seq_len.tolist() == [17, 3] + assert context.batch.b_seq_len.data_ptr() == pointer + assert context.current_max_kv_seq_len == 17 diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py index 52b1c1cf2d..7e3c4bb8f0 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -5,11 +5,33 @@ import torch from lightllm.common.basemodel import cuda_graph as cuda_graph_module +from lightllm.common.attention.session import AttentionSession from lightllm.common.basemodel.cuda_graph import CudaGraph from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune +def make_capture_states(manager, overlap, *, backend=None): + from lightllm.common.attention.api import AttentionPhase + + states = [] + for _ in range(2 if overlap else 1): + runtime = AttentionSession() + if backend is not None: + runtime.contexts[0] = SimpleNamespace(capacity=SimpleNamespace(phase=AttentionPhase.DECODE)) + runtime.groups = ({AttentionPhase.DECODE: backend},) + states.append( + SimpleNamespace( + input_ids=torch.ones(2, dtype=torch.int64), + b_req_idx=torch.tensor([0, 1]), + att_state_role="main", + att_state_manager=manager, + attention=runtime, + ) + ) + return states + + @pytest.fixture def capture_env(monkeypatch): env = SimpleNamespace(capturing=False, capture_count=0, graphs=[]) @@ -22,17 +44,21 @@ def capture_env(monkeypatch): monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") - monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(torch.cuda, "synchronize", lambda *args: None) monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: None) class FakeGraph: def __init__(self): self.replays = 0 + self.resets = 0 env.graphs.append(self) def replay(self): self.replays += 1 + def reset(self): + self.resets += 1 + @contextmanager def capture(graph, **kwargs): assert not Autotuner.is_autotune_warmup() @@ -45,6 +71,12 @@ def capture(graph, **kwargs): monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) monkeypatch.setattr(torch.cuda, "graph", capture) + from lightllm.platform import DeviceCaps, Platform + from lightllm.platform.cuda import CudaRuntime + + runtime = CudaRuntime("cuda:0") + runtime.caps = DeviceCaps(Platform.CUDA, torch.device("cuda:0"), supports_graph=True) + monkeypatch.setattr(cuda_graph_module, "get_device_runtime", lambda: runtime) return env @@ -82,7 +114,13 @@ def bench(size, run_config): general, general_cache = make_kernel(AutotuneKernelType.GENERAL) attention, attention_cache = make_kernel(AutotuneKernelType.DECODE_ATTENTION) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] + manager = SimpleNamespace( + storage_generation=0, + register_graph=lambda graph: None, + snapshot_capture_state=lambda *args: (), + restore_capture_state=lambda snapshot: None, + ) + states = make_capture_states(manager, overlap) outputs = [object() for _ in states] forward_calls = [] @@ -115,7 +153,13 @@ def decode_func(*infer_states): def test_warmup_failure_restores_phase_without_capturing(capture_env, overlap): capture_env.args.enable_decode_microbatch_overlap = overlap graph = CudaGraph(1, 2, 1, max_batch_size=2) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] + manager = SimpleNamespace( + storage_generation=0, + register_graph=lambda graph: None, + snapshot_capture_state=lambda *args: (), + restore_capture_state=lambda snapshot: None, + ) + states = make_capture_states(manager, overlap) def decode_func(*infer_states): assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) @@ -128,3 +172,41 @@ def decode_func(*infer_states): assert not Autotuner.is_autotune_warmup() assert capture_env.capture_count == 0 assert graph.graph == {} + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("fail_binding", [False, True]) +def test_capture_failure_restores_state_and_releases_native_graph(capture_env, overlap, fail_binding): + from lightllm.common.attention.api import AttentionPhase + + capture_env.args.enable_decode_microbatch_overlap = overlap + graph = CudaGraph(1, 2, 1, max_batch_size=2) + persistent = torch.zeros(1) + manager = SimpleNamespace( + storage_generation=0, + register_graph=lambda graph: None, + snapshot_capture_state=lambda *args: persistent.clone(), + restore_capture_state=lambda snapshot: persistent.copy_(snapshot), + ) + + bindings = [] + + def create_capture_hooks(context): + bindings.append(context) + raise ValueError("failed binding") + + backend = SimpleNamespace(create_capture_hooks=create_capture_hooks) + states = make_capture_states(manager, overlap, backend=backend if fail_binding else None) + + def decode_func(*infer_states): + persistent.add_(1) + if capture_env.capturing and not fail_binding: + raise ValueError("failed capture") + return (None, None) if overlap else None + + with pytest.raises(ValueError, match="failed binding" if fail_binding else "failed capture"): + graph.capture_decode(decode_func, *states) + assert bool(bindings) == fail_binding + assert persistent.item() == 0 + assert not capture_env.capturing and not graph.graph + assert capture_env.graphs[0].resets == 1 diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 123f603ff2..cac2a5cc17 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -1,12 +1,10 @@ -from types import SimpleNamespace import weakref +from types import SimpleNamespace -import torch import pytest +import torch from lightllm.common.basemodel import basemodel -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState -from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -100,7 +98,7 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): expected_batch_size, ) in execution_configs: model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode, page_size=1) + model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode, page_size=1, mtp_step=0) model.tp_world_size_ = tp_world_size model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) @@ -123,10 +121,11 @@ def create_infer_state(model_input): b_seq_len=model_input.b_seq_len, mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32), init_some_extra_state=lambda _: None, + finish_attention=lambda: None, is_cuda_graph_capture=False, is_cuda_graph_replay=False, ) - infer_state.init_att_state = lambda: graph_flags_at_att_init.append( + infer_state.prepare_attention = lambda _: graph_flags_at_att_init.append( (infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay) ) return infer_state @@ -172,7 +171,7 @@ def execute(mode, infer_state): ) def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(lengths, graph_mode, expected_mode): model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1, mtp_step=0) model.tp_world_size_ = 2 model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) @@ -183,16 +182,9 @@ def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(le def graph_flags(state): return state.is_cuda_graph_capture, state.is_cuda_graph_replay - class DecodeState(FlashInferDecodeAttState): - def init_state(self): - infer_state = self.infer_state - # Replace the external GPU wrapper allocation while preserving the - # production decision and InferStateInfo's initialization entry point. - if should_init_decode_wrapper(infer_state): - self.decode_wrapper = object() - initializations.append( - (infer_state.microbatch_index, graph_flags(infer_state), self.decode_wrapper is not None) - ) + def prepare(state): + state.test_wrapper = None if state.is_cuda_graph_replay else object() + initializations.append((state.microbatch_index, graph_flags(state), state.test_wrapper is not None)) def create_state(model_input, microbatch_index=0): state = InferStateInfo() @@ -203,7 +195,8 @@ def create_state(model_input, microbatch_index=0): state.b_req_idx = model_input.b_req_idx state.b_seq_len = model_input.b_seq_len state.microbatch_index = microbatch_index - state.decode_att_state = DecodeState(_infer_state_weak_ref=weakref.ref(state)) + state.prepare_attention = lambda _: prepare(state) + state.finish_attention = lambda: None state.init_some_extra_state = lambda _: None return state @@ -212,8 +205,8 @@ def outputs(state, state1): def forward(state, infer_state1): # Eager and capture execute attention using these newly prepared states. - assert state.decode_att_state.decode_wrapper is not None - assert infer_state1.decode_att_state.decode_wrapper is not None + assert state.test_wrapper is not None + assert infer_state1.test_wrapper is not None executions.append(("forward", graph_flags(state), graph_flags(infer_state1))) return outputs(state, infer_state1) @@ -223,8 +216,8 @@ def capture(func, state, infer_state1): def replay(state, infer_state1): # Replay uses graph-owned wrappers, so temporary states need none. - assert state.decode_att_state.decode_wrapper is None - assert infer_state1.decode_att_state.decode_wrapper is None + assert state.test_wrapper is None + assert infer_state1.test_wrapper is None executions.append(("replay", graph_flags(state), graph_flags(infer_state1))) return outputs(state, infer_state1) @@ -265,7 +258,7 @@ def replay(state, infer_state1): @pytest.mark.parametrize("rows", [0, 3, 9, 24, 27]) def test_mtp_decode_preserves_groups_and_real_rows(overlap, dynamic, graph_mode, rows): model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, mtp_step=0) model.tp_world_size_ = 8 model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1 if dynamic else 3) @@ -283,7 +276,8 @@ def check_layout(): b_req_idx=model_input.b_req_idx, b_seq_len=model_input.b_seq_len, init_some_extra_state=lambda _: None, - init_att_state=check_layout, + finish_attention=lambda: None, + prepare_attention=lambda _: check_layout(), ) def forward(state): diff --git a/unit_tests/common/basemodel/test_moe_route_info_manager.py b/unit_tests/common/basemodel/test_moe_route_info_manager.py index de2bc7a857..ddc35595b4 100644 --- a/unit_tests/common/basemodel/test_moe_route_info_manager.py +++ b/unit_tests/common/basemodel/test_moe_route_info_manager.py @@ -115,10 +115,6 @@ def test_capture_rejects_unexpected_topk_width(self): def test_cuda_graph_replay_uses_copied_mem_indexes(self): _skip_without_cuda() - class _NoopDecodeState: - def copy_for_decode_cuda_graph(self, other): - return - manager = MoeRouteInfoManager( num_moe_layers=1, topk=2, @@ -126,7 +122,9 @@ def copy_for_decode_cuda_graph(self, other): ) manager.init_capture_buffer(kv_cache_size=256) graph_infer_state = InferStateInfo() - graph_infer_state.decode_att_state = _NoopDecodeState() + from unit_tests.common.attention.helpers import make_batch + + graph_infer_state.attention.batch = make_batch(device="cuda") graph_infer_state.mem_index = torch.tensor([10, 11], device="cuda") moe_capture_callback = manager.get_moe_capture_callback(0, graph_infer_state.mem_index) topk_ids = torch.tensor([[1, 2], [3, 4]], device="cuda") @@ -141,10 +139,14 @@ def copy_for_decode_cuda_graph(self, other): np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) new_infer_state = InferStateInfo() - new_infer_state.decode_att_state = _NoopDecodeState() + from dataclasses import replace + + new_infer_state.attention.batch = replace(graph_infer_state.attention.batch) new_infer_state.mem_index = torch.tensor([20, 21], device="cuda") new_topk_ids = torch.tensor([[5, 6], [7, 8]], device="cuda") - graph_infer_state.copy_for_cuda_graph(new_infer_state) + for target, source in graph_infer_state.graph_input_copies(new_infer_state): + target.copy_(source, non_blocking=True) + graph_infer_state.attention.plan_replay(new_infer_state.attention) topk_ids.copy_(new_topk_ids) graph.replay() diff --git a/unit_tests/common/basemodel/test_mtp_manager.py b/unit_tests/common/basemodel/test_mtp_manager.py index f7e91f8c79..f33e099003 100644 --- a/unit_tests/common/basemodel/test_mtp_manager.py +++ b/unit_tests/common/basemodel/test_mtp_manager.py @@ -83,7 +83,7 @@ def test_decode_batch_alignment(monkeypatch, spec_mode, dynamic_verify, is_draft ("dflash2", True, 7), ], ) -def test_decode_draft_step(monkeypatch, spec_mode, is_draft_model, expected): +def test_draft_step(monkeypatch, spec_mode, is_draft_model, expected): args = SimpleNamespace(mtp_mode=spec_mode, mtp_step=7, mtp_dynamic_verify=False) monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py index ef7c18a061..56885aac63 100644 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -8,6 +8,7 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA runner stream required") def test_base_model_prefill_accepts_two_prebuilt_inputs(monkeypatch): model_input0 = _make_prefill_input([10, 11], req_idx=20, is_decode_req=False) model_input1 = _make_prefill_input([12], req_idx=21, is_decode_req=True) @@ -52,6 +53,7 @@ def fake_overlap_forward(input0, input1): assert len(outputs) == 2 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA runner stream required") def test_base_model_decode_accepts_two_inputs_and_skips_empty_gather(monkeypatch): model_input0 = _make_decode_input() model_input0.input_ids = None @@ -122,7 +124,7 @@ def test_overlap_decode_cuda_pads_empty_side_and_unpads_outputs(monkeypatch): model_input1.to_cuda() model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1, mtp_step=0) model.tp_world_size_ = 2 model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) model.graph = None @@ -136,7 +138,8 @@ def fake_create_inferstate(model_input, microbatch_index): b_seq_len=model_input.b_seq_len, mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32, device="cuda"), init_some_extra_state=lambda _: None, - init_att_state=lambda: None, + prepare_attention=lambda _: None, + finish_attention=lambda: None, ) model._create_inferstate = fake_create_inferstate diff --git a/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py b/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py index bed4f3967f..a9bcb41401 100644 --- a/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py +++ b/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py @@ -1,16 +1,21 @@ +from contextlib import contextmanager from types import SimpleNamespace +import pytest import torch import lightllm.common.basemodel.hidden_collector as hidden_collector_module from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector -from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph, PrefillGraphExecution +from lightllm.platform import DeviceRuntime class _GraphInferState: def __init__(self, hidden_collector): self.hidden_collector = hidden_collector self.input_ids = torch.empty(4, dtype=torch.int64) + self.att_state_manager = SimpleNamespace(storage_generation=0, register_graph=lambda graph: None) self.copied_from = None self.replayed_with = None @@ -54,7 +59,9 @@ def test_replay_restores_captured_hidden_state_into_request_collector(monkeypatc ) graph_output = torch.randn(2, 3) prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} + prefill_graph.graph = { + 4: PrefillGraphExecution(graph_infer_state, [], [graph_output], graph_collector, DeviceRuntime("cpu")) + } outputs = prefill_graph._replay(input_tensors=[], infer_state=request_infer_state) @@ -73,9 +80,49 @@ def test_first_replay_replaces_capture_collector_with_runtime_instance(monkeypat graph_infer_state = _GraphInferState(hidden_collector=graph_collector) graph_output = torch.randn(2, 3) prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} + prefill_graph.graph = { + 4: PrefillGraphExecution(graph_infer_state, [], [graph_output], graph_collector, DeviceRuntime("cpu")) + } prefill_graph._replay(input_tensors=[], infer_state=graph_infer_state) assert graph_infer_state.hidden_collector is not graph_collector assert graph_infer_state.hidden_collector.layer_hiddens[0] is graph_collector.layer_hiddens[0] + + +@pytest.mark.parametrize("between_segments", [False, True]) +def test_failed_prefill_capture_releases_segments(between_segments): + events = [] + + class Runtime(DeviceRuntime): + def create_graph(self): + return SimpleNamespace(reset=lambda: events.append("reset")) + + @contextmanager + def capture_graph(self, graph, **kwargs): + events.append("enter") + try: + yield + finally: + events.append("exit") + + def synchronize(self): + events.append("synchronize") + + prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) + prefill_graph.graph = {} + prefill_graph.mempool = None + prefill_graph.runtime = Runtime("cpu") + state = InferStateInfo() + state.input_ids = torch.empty(4, dtype=torch.int64) + + def forward(inputs, state): + if between_segments: + state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + raise ValueError("failed segment") + + with pytest.raises(ValueError, match="failed segment"): + prefill_graph._capture_prefill(forward, [], state) + assert events == ["enter", "exit", "synchronize", "reset"] + assert not prefill_graph.graph + assert not state.prefill_cuda_graph_exe_list diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py index 7c1af6bbc2..de7b93c093 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py @@ -13,6 +13,8 @@ import pytest import torch +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + def gqa_attention_reference( q: torch.Tensor, @@ -142,7 +144,7 @@ def mtp_diverse_attention( """ MTP Diverse Attention 调用入口 """ - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( token_decode_attention_mtp_diverse_single_token, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py index 0989fe6d26..5657aa21ef 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,17 +10,27 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.attention.triton import fp as state_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( +from lightllm.common.attention.backends.dense.triton import backend as state_module +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( mtp_diverse_attn as decode_module, stage1_single_token as stage1_module, stage2_single_token as stage2_module, ) +from lightllm.common.attention_state.manager import AttentionStateManager +from lightllm.common.basemodel import cuda_graph as graph_module +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -138,16 +148,18 @@ def test_decode_passes_saved_length_and_selected_block_n(monkeypatch): infer_state.b_req_idx = inputs["B_req_idx"] infer_state.b_seq_len = inputs["b_seq_len"] infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs model = SimpleNamespace( is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), ) - monkeypatch.setattr(state_module, "build_mtp_shared_group_markers", lambda *a, **kw: inputs["b_mark_shared_group"]) - state = state_module.TritonDecodeAttState( - backend=SimpleNamespace(model=model), _infer_state_weak_ref=weakref.ref(infer_state) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.mtp_utils.build_mtp_shared_group_markers", + lambda *a, **kw: inputs["b_mark_shared_group"], ) - state.init_state() + state = backend_context(state_module.TritonDenseDecodeBackend, "decode", settings=SimpleNamespace(model=model)) + update_context(state, infer_state) infer_state.max_kv_seq_len = 32768 calls = [] @@ -157,7 +169,7 @@ def stage1(**kwargs): monkeypatch.setattr(decode_module, "mtp_diverse_stage1_single_token", stage1) monkeypatch.setattr(decode_module, "mtp_diverse_stage2_single_token", lambda **kwargs: calls.append(kwargs)) - state.decode_att(q=inputs["q"], k=inputs["k"], v=inputs["v"]) + run_dense(state, infer_state, inputs["q"], inputs["k"], inputs["v"]) assert calls[0]["max_kv_len"] == 8193 assert calls[0]["b_mark_shared_group"] is inputs["b_mark_shared_group"] assert calls[1]["block_n"] == 32 @@ -264,7 +276,16 @@ def decode(state): return out graph = graph_module.CudaGraph(1, 7, 1, max_batch_size=7, max_len_in_batch=32768) - graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(7, device="cuda"))) + resources = AttentionStateManager(request_capacity=7) + resources.finish_state_registration() + state = SimpleNamespace( + input_ids=torch.ones(7, device="cuda"), + b_req_idx=torch.arange(7, device="cuda", dtype=torch.int32), + att_state_manager=resources, + att_state_role="main", + attention=SimpleNamespace(contexts={}, groups=()), + ) + graph.capture_decode(decode, state) torch.cuda.synchronize() assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) assert valid_benchmark_count > 0 @@ -296,6 +317,6 @@ def decode(state): inputs["b_seq_len"].copy_( torch.tensor([kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len], device="cuda", dtype=torch.int32) ) - graph.graph[7][0].replay() + graph.graph[7].replay() torch.cuda.synchronize() torch.testing.assert_close(out.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py index 48956c4f36..63ff94f836 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py @@ -7,10 +7,13 @@ import pytest import torch +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.mtp_diverse import ( + stage2_single_token as stage2_module, +) +from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import stage2_single_token as stage2_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len @@ -189,7 +192,16 @@ def decode(state): return inputs["out"] graph = graph_module.CudaGraph(1, 2, 1, max_batch_size=2, max_len_in_batch=32768) - graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(2, device="cuda"))) + resources = AttentionStateManager(request_capacity=2) + resources.finish_state_registration() + state = SimpleNamespace( + input_ids=torch.ones(2, device="cuda"), + b_req_idx=torch.arange(2, device="cuda", dtype=torch.int32), + att_state_manager=resources, + att_state_role="main", + attention=SimpleNamespace(contexts={}, groups=()), + ) + graph.capture_decode(decode, state) torch.cuda.synchronize() assert rebuild_count == 1 and benchmark_count == 20 torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) @@ -215,6 +227,6 @@ def decode(state): inputs["mid_out"].normal_() inputs["mid_out_logsumexp"].normal_() inputs["B_Seqlen"].fill_(length) - graph.graph[2][0].replay() + graph.graph[2].replay() torch.cuda.synchronize() torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py index b3068db5fb..592a13aa58 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py @@ -1,7 +1,7 @@ -import weakref import collections import inspect import math +import weakref from types import SimpleNamespace import pytest @@ -10,18 +10,27 @@ from frozendict import frozendict from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( +from lightllm.common.attention.backends.dense.triton.backend import TritonDenseDecodeBackend +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash import ( gqa_flash_decoding as decode_module, gqa_flash_decoding_stage1 as stage1_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -127,24 +136,27 @@ def test_rebuild_preserves_layout_and_bounds_indices(monkeypatch, num_tokens): def test_state_retains_actual_length_before_graph_capture(monkeypatch): infer_state = InferStateInfo() infer_state.max_kv_seq_len = 8193 - state = TritonDecodeAttState( - backend=SimpleNamespace( + state = backend_context( + TritonDenseDecodeBackend, + "decode", + settings=SimpleNamespace( model=SimpleNamespace( is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0) ) ), - _infer_state_weak_ref=weakref.ref(infer_state), ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 + update_context(state, infer_state) + infer_state.max_kv_seq_len = 32768 calls = [] monkeypatch.setattr( decode_module, "gqa_token_decode_attention_flash_decoding", lambda **kwargs: calls.append(kwargs) ) - state.decode_att( - q=torch.empty(2, 8, 64), - k=torch.empty(16, 2, 64), - v=torch.empty(16, 2, 64), + run_dense( + state, + infer_state, + torch.empty(2, 8, 64), + torch.empty(16, 2, 64), + torch.empty(16, 2, 64), alloc_func=lambda shape, dtype: torch.empty(shape, dtype=dtype), ) assert calls[0]["max_len_in_batch"] == 8193 @@ -258,12 +270,14 @@ def decode(): decode() torch.cuda.synchronize() + eager_output = output.clone() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): decode() graph.replay() torch.cuda.synchronize() torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + torch.testing.assert_close(output, eager_output, atol=0, rtol=0) assert rebuild_count == 1 and benchmark_count == 3 for name, value in snapshots.items(): if name not in ["mid_out", "mid_out_logsumexp"]: diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py index df357539e1..75cea3e976 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,18 +10,27 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ( +from lightllm.common.attention.backends.dense.triton.int4 import TritonDenseInt4DecodeBackend +from lightllm.common.attention.backends.dense.triton.kernels.decode.int4kv import ( int4kv_flash_decoding_stage1 as stage1_module, ppl_int4kv_flash_decoding as decode_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -141,17 +150,17 @@ def test_decode_passes_saved_length_to_stage1(monkeypatch): infer_state.b_req_idx = inputs["B_req_idx"] infer_state.b_seq_len = inputs["B_Seqlen"] infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - state = Int4kvTritonDecodeAttState( - backend=SimpleNamespace(), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs + state = backend_context(TritonDenseInt4DecodeBackend, "decode", settings=SimpleNamespace()) + update_context(state, infer_state) + infer_state.max_kv_seq_len = 32768 calls = [] monkeypatch.setattr(stage1_module, "int4kv_flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) stage2 = inspect.getmodule(flash_decode_stage2) monkeypatch.setattr(stage2, "flash_decode_stage2", lambda *args: None) - state.decode_att( + run_dense( + state, + infer_state, inputs["q"], (inputs["k"], inputs["k_scale"]), (inputs["v"], inputs["v_scale"]), diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py index e1b6bc7758..7a478f0f98 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py @@ -1,21 +1,14 @@ import pytest -from types import SimpleNamespace - import torch +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + def alloc_tensor_func(shape, dtype, device): """兼容的 tensor 分配函数""" return torch.empty(shape, dtype=dtype, device=device) -class MockReqManager: - """Mock request manager for testing""" - - def __init__(self, req_to_token_indexs): - self.req_to_token_indexs = req_to_token_indexs - - class MockInferState: """Mock infer state for testing""" @@ -29,7 +22,7 @@ def __init__( ): self.batch_size = batch_size self.max_kv_seq_len = max_kv_seq_len - self.req_manager = MockReqManager(req_to_tokens) + self.req_to_token_indexs = req_to_tokens self.b_req_idx = b_req_idx self.b_seq_len = b_seq_len @@ -43,10 +36,10 @@ def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(sha diverse：int8kv_flash_decoding_diverse；对照：int8kv/normal token_decode_attention_flash_decoding。 """ - from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse import ( token_decode_attention_flash_decoding as diverse_attention, ) - from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( token_decode_attention_flash_decoding as normal_decode, ) @@ -102,15 +95,10 @@ def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(sha b_req_idx=b_req_idx, b_seq_len=b_seq_len, ) - diverse_infer_state.decode_att_state = SimpleNamespace( - b_shared_seq_len=b_shared_seq_len, - b_mark_shared_group=b_mark_shared_group, - ) - # 运行 baseline normal_out = normal_decode( q=q.clone(), - infer_state=baseline_infer_state, + batch=baseline_infer_state, cache_k=cache_k, cache_k_scale=cache_k_scale, cache_v=cache_v, @@ -121,7 +109,10 @@ def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(sha # 运行 diverse 版本 diverse_out = diverse_attention( q=q.clone(), - infer_state=diverse_infer_state, + batch=diverse_infer_state, + b_shared_seq_len=b_shared_seq_len, + b_mark_shared_group=b_mark_shared_group, + shared_streams_dict={}, cache_k=cache_k, cache_k_scale=cache_k_scale, cache_v=cache_v, diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py index f3cb8de463..3466d2f105 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py @@ -1,6 +1,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage1 import ( + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage1 import ( flash_decode_stage1, ) @@ -95,7 +97,7 @@ def test_flash_decode_stage1_execution(setup_tensors): new_k = k.to(q.dtype) new_v = v.to(q.dtype) - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( flash_decode_stage1 as gqa_flash_decode_stage1, ) @@ -174,14 +176,14 @@ def fn_triton(st=setup_tensors): del setup_tensors - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print("SUMMARY - Performance Comparison") - print(f"{'='*80}") + print(f"{'=' * 80}") print(f"{'batch_size':<8} {'seq_len':<12} {'triton_ms':<12}") - print(f"{'-'*80}") + print(f"{'-' * 80}") for r in results: print(f"{r['batch_size']:<8} {r['seq_len']:<12} {r['triton_ms']:<12.3f}") - print(f"{'='*80}") + print(f"{'=' * 80}") if __name__ == "__main__": diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py index 985a33289f..32925d2a0e 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py @@ -1,6 +1,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage2 import ( flash_decode_stage2, ) @@ -104,7 +106,7 @@ def test_flash_decode_stage2_execution(shared_seq_len): b_seq_len = setup_tensors["b_seq_len"] - setup_tensors["b_shared_seq_len"] req_to_tokens = setup_tensors["Req_to_tokens"][:, setup_tensors["b_shared_seq_len"][0].item() :] - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.gqa.flash.gqa_flash_decoding_stage1 import ( flash_decode_stage1 as gqa_flash_decode_stage1, ) @@ -184,9 +186,9 @@ def bench_stage2(): print(results[-1]) del setup_tensors - print(f"\n{'='*80}") + print(f"\n{'=' * 80}") print(f"{'batch_size':<10} {'seq_len':<10} {'flash_decode_stage2_ms':<22}") - print(f"{'-'*80}") + print(f"{'-' * 80}") for r in results: print(f"{r['batch_size']:<10} {r['seq_len']:<10} {r['flash_decode_stage2_ms']:<22.4f}") - print(f"{'='*80}") + print(f"{'=' * 80}") diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py index c1a0ca1e58..9dfdd2c506 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py @@ -1,6 +1,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage3 import ( + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.int8kv_flash_decoding_diverse_stage3 import ( flash_diverse_decode_stage3, ) @@ -26,7 +28,7 @@ def test_flash_diverse_decode_stage3(batch, head_num, seq_len, shared_seq_len, b true_out = torch.zeros_like(out) - from lightllm.common.basemodel.triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding_stage2 import ( + from lightllm.common.attention.backends.dense.triton.kernels.decode.mha.flash.flash_decoding_stage2 import ( flash_decode_stage2, ) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py index 7033489f66..6b0f49531d 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py @@ -1,8 +1,8 @@ -import weakref import collections import inspect import json import math +import weakref from types import SimpleNamespace import pytest @@ -10,19 +10,28 @@ import torch.nn.functional as F from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as state_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( - int8kv_flash_decoding_stage1 as stage1_module, +from lightllm.common.attention.backends.dense.triton import int8 as state_module +from lightllm.common.attention.backends.dense.triton.int8 import TritonDenseInt8DecodeBackend +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal import ( int8kv_flash_decoding as decode_module, + int8kv_flash_decoding_stage1 as stage1_module, ) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( +from lightllm.common.attention.backends.dense.triton.kernels.decode.int8kv.normal.int8kv_flash_decoding_stage2 import ( flash_decode_stage2, ) +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -143,16 +152,16 @@ def test_decode_passes_saved_length_to_stage1(monkeypatch): infer_state.b_req_idx = inputs["B_req_idx"] infer_state.b_seq_len = inputs["B_seq_len"] infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - state = Int8kvTritonDecodeAttState( - backend=SimpleNamespace(), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs + state = backend_context(TritonDenseInt8DecodeBackend, "decode", settings=SimpleNamespace()) + update_context(state, infer_state) + infer_state.max_kv_seq_len = 32768 calls = [] monkeypatch.setattr(decode_module, "flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) monkeypatch.setattr(decode_module, "flash_decode_stage2", lambda **kwargs: None) - state.decode_att( + run_dense( + state, + infer_state, inputs["q"], (inputs["k"], inputs["k_scale"]), (inputs["v"], inputs["v_scale"]), diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py index ecab784529..3c35d64c1a 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py @@ -23,10 +23,13 @@ """ import math + import pytest import torch -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo +from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( + context_attention_fwd_neo, +) def torch_reference_context_attention_neo( diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py index d1a53f873f..71f77aff71 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py @@ -1,7 +1,9 @@ -import torch import time + import pytest -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( +import torch + +from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( context_attention_fwd_contiguous_kv, ) from lightllm.utils.log_utils import init_logger diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py index 541594306d..5ddcc7c1c5 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py @@ -1,15 +1,17 @@ -import torch import time -import pytest + +import flashinfer import numpy as np +import pytest +import torch import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( + +from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( context_attention_fwd, context_attention_fwd_no_prompt_cache, ) from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py index eaa19ac711..4d53830a04 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py @@ -24,10 +24,11 @@ """ import math + import pytest import torch -from lightllm.common.basemodel.attention.fa3.neo import flash_attn_with_kvcache_neo as flash_attn_with_kvcache +from lightllm.common.attention.backends.dense.fa3.neo import flash_attn_with_kvcache_neo as flash_attn_with_kvcache try: import triton @@ -37,7 +38,7 @@ triton_testing = None try: - from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import ( + from lightllm.common.attention.backends.dense.triton.kernels.prefill.context_attention_fwd_neo import ( context_attention_fwd_neo, ) except ImportError: @@ -496,7 +497,7 @@ def triton_run(): f"bench: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={str(dtype).split('.')[-1]:<8s} " f"max_q_seq_len={max_q_seq_len:4d} max_prompt_cache_len={max_prompt_cache_len:4d} " f"image_tokens={n_image:4d}/{n_tokens:5d} sum_kv={sum_kv:6d} | " - f"fa3 {fa3_ms*1000:8.1f} us | triton {triton_ms*1000:8.1f} us | " + f"fa3 {fa3_ms * 1000:8.1f} us | triton {triton_ms * 1000:8.1f} us | " f"speedup {speedup:5.2f}x" ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py index 956bf5e173..fac1a8bb61 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -4,9 +4,9 @@ import triton import triton.language as tl +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor def attach_replay_history( @@ -107,7 +107,7 @@ def replay_layer( """按测试 state 的类型调用对应的 replay decode。""" if state.kda: - from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + from lightllm.common.attention.backends.kda.replay.decode import replay return replay( state_ssm=state.ssm_state[layer_idx], @@ -133,7 +133,7 @@ def replay_layer( state_gate_lower_bound=state.gate_lower_bound, run_config=run_config, ) - from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + from lightllm.common.attention.backends.gdn.replay.decode import replay return replay( state=state, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py index 7fbb63df5d..811f66f137 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -1,9 +1,9 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, @@ -26,6 +26,7 @@ def ssm_snapshot(backend, req): from types import SimpleNamespace + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState @@ -63,7 +64,7 @@ def _replay_cu(args): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_request_counts_share_compiled_kernel(monkeypatch): - from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix + from lightllm.common.attention.backends._shared.linear.ssm_acceptance import _commit_prefix, commit_ssm_prefix reqs = torch.arange(512, dtype=torch.int32, device="cuda") lengths = reqs % 4 + 1 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py index f7c742db01..5678745a58 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py @@ -3,7 +3,9 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update +from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import causal_conv1d_update + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def causal_conv1d_ref( diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py index 969a3f354d..a67ef10cf1 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py @@ -1,9 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule if not torch.cuda.is_available(): pytest.skip("CUDA required", allow_module_level=True) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py index 8d92d6cc8d..7e556390b5 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py @@ -9,11 +9,10 @@ import torch import triton -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import kda +from lightllm.common.attention.backends.kda.triton.kernels import kda from lightllm.common.kernel_config import KernelConfigs from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel - +from lightllm.common.triton_utils.autotuner import AutotuneLevel, Autotuner KERNELS = [ ("_chunk_kda_scaled_dot_kkt_sub_inter", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter"), @@ -212,7 +211,7 @@ def test_kda_requires_packed_inputs(fused_gate, invalid_input): def test_split_diagonal_matches_unsplit(monkeypatch, heads, input_dtype, output_dtype, precision): if precision != "ieee" and torch.cuda.get_device_capability()[0] < 8: pytest.skip("TF32 requires SM80+") - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") monkeypatch.setattr(solver, "FLA_TRIL_PRECISION", precision) lengths = (1, 16, 1519) cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) @@ -236,7 +235,7 @@ def observe(*args, **kwargs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("batch,tokens,heads", [(1, 64, 2), (2, 65, 3), (3, 129, 2)]) def test_split_diagonal_matches_unsplit_dense(monkeypatch, batch, tokens, heads): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") rows = torch.arange(tokens, device="cuda") % 64 matrix = torch.randn(batch, tokens, heads, 64, device="cuda") * 0.025 matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) @@ -268,7 +267,7 @@ def observe(*args, **kwargs): ids=["16-requests", "128-requests", "quarter-full-chunks", "16k-tokens", "mostly-empty-chunks", "tiny-batch"], ) def test_split_diagonal_uses_chunk_occupancy(monkeypatch, lengths, expect_split): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") + solver = importlib.import_module("lightllm.common.attention.backends._shared.linear.solve_tril") cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) matrix = torch.randn(1, sum(lengths), 2, 64, device="cuda") * 0.025 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py index 2f8544cf75..cecbf9038c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py @@ -7,15 +7,9 @@ import torch import torch.nn.functional as F -from lightllm.common.basemodel.triton_kernel.linear_att import ( - mtp_fused_recurrent as kernel_module, -) +from lightllm.common.attention.backends.gdn.triton.kernels import mtp_fused_recurrent as kernel_module from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import ( - Autotuner, - AutotuneKernelType, - AutotuneLevel, -) +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner @pytest.fixture(autouse=True) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..bbe8580df8 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -1,10 +1,8 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( +from lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent import fused_recurrent_gated_delta_rule +from lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py index 8252c445a1..d03908858f 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -1,11 +1,8 @@ import pytest import torch -from lightllm.common.attention_state.bounded_att import GDNState, ConvSSMStateConfig - -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index 786376a362..3154c0349a 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.attention.linear.replay_gdn_decode import get_configs +from lightllm.common.attention.backends.gdn.replay.decode import get_configs from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, @@ -22,9 +22,9 @@ def ssm_snapshot(backend, req): return clone[:, req] -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -867,6 +867,7 @@ def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, sepa torch.manual_seed(72) initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda) + backend.tile_sync_mark = torch.zeros((bucket, backend.num_key_heads), dtype=torch.int32, device="cuda") reqs = torch.cat( ( torch.arange(live, device="cuda", dtype=torch.int32), diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py index 02797b61b7..a85516af28 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -1,25 +1,24 @@ +from types import SimpleNamespace + import pytest import torch +from lightllm.common.attention.backends._shared.linear.mtp_state_params import build_dynamic_mtp_linear_att_state_params +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.ssm_acceptance import commit_ssm_prefix from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( accept_replay_window, attach_replay_history as ReplaySSMExecutor, replay_layer, ) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from types import SimpleNamespace pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import _advance_cursor - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay + from lightllm.common.attention.backends._shared.linear.replayssm_cursor import _advance_cursor + from lightllm.common.attention.backends._shared.linear.replayssm_decode import _decode_replay kernels = {"_decode_replay": _decode_replay, "_advance_cursor": _advance_cursor} hashes = {name: set() for name in kernels} diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py index a94411fbd5..bf3e383953 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -1,7 +1,7 @@ import pytest import torch -from lightllm.common.basemodel.attention.linear import replay_gdn_decode as tuning +from lightllm.common.attention.backends.gdn.replay import decode as tuning from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( attach_replay_history as ReplaySSMExecutor, ) @@ -30,7 +30,7 @@ def test_rebuild_uses_one_layer_scratch(width): @pytest.mark.parametrize("width", [1, 5]) def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, width): from lightllm.common.triton_utils import autotuner as autotuner_module - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) diff --git a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py index a5ac9708d2..e6bbe2bcf7 100644 --- a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py +++ b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py @@ -1,14 +1,16 @@ -import torch import time -import pytest + +import flashinfer import numpy as np +import pytest +import torch import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.mla_att.decode_att.gqa_flash_decoding import ( + +from lightllm.common.attention.backends.mla.triton.kernels.decode.gqa_flash_decoding import ( gqa_token_decode_attention_flash_decoding, ) from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo +from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -52,10 +54,11 @@ def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head): infer_state = Deepseek2InferStateInfo() infer_state.batch_size = Z - infer_state.max_len_in_batch = N_CTX + infer_state.max_kv_seq_len = N_CTX infer_state.total_token_num = Z * N_CTX infer_state.req_manager = type("Object", (), {})() infer_state.req_manager.req_to_token_indexs = req_to_token_indexs + infer_state.req_to_token_indexs = infer_state.req_manager.req_to_token_indexs infer_state.b_req_idx = b_req_idx infer_state.b_seq_len = b_seq_len infer_state.kv_starts = kv_starts @@ -70,6 +73,8 @@ def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head): infer_state, sm_scale, o, + block_seq=torch.empty(1, dtype=torch.int64, device="cuda"), + batch_start=torch.empty(Z, dtype=torch.int64, device="cuda"), ) batch_size = Z diff --git a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py index 0c1c2df08a..342ff905cf 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py +++ b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py @@ -7,10 +7,19 @@ if not torch.cuda.is_available(): pytest.skip("requires CUDA", allow_module_level=True) +from lightllm.common.attention.backends.dense.triton import int8 as int8kv_module +from lightllm.common.attention.backends.dense.triton.int8 import TritonDenseInt8DecodeBackend from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as int8kv_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState from lightllm.common.basemodel.triton_kernel import diverse_utils +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.mark.parametrize( @@ -40,9 +49,9 @@ def test_int8kv_decode_state_rebuilds_diverse_metadata(monkeypatch): infer_state.max_kv_seq_len = 8192 infer_state.b_shared_seq_len = torch.tensor([8, 8, 5, 0], dtype=torch.int32, device="cuda") infer_state.b_shared_radix_node_id = torch.tensor([10, 10, 20, -1], dtype=torch.int64, device="cuda") - state = Int8kvTritonDecodeAttState(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer_state)) + state = backend_context(TritonDenseInt8DecodeBackend, "decode", settings=SimpleNamespace()) - state.init_state() + update_context(state, infer_state) assert state.b_mark_shared_group.cpu().tolist() == [0, 2, 1, 1] assert state.b_shared_seq_len.cpu().tolist() == [8, 8, 0, 0] diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index f5bd92e08a..6d219dc21d 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -1,15 +1,20 @@ +from lightllm.models.deepseek_v4.workspace import get_swa_prefill_workspace +from unit_tests.common.attention.helpers import make_env import json +from dataclasses import replace from types import SimpleNamespace -import numpy as np import pytest import torch -from unit_tests.common.dsv4_test_utils import DeepseekV4MemoryManager, DeepseekV4CacheConfig, DeepseekV4ReqManager -from unit_tests.common.req_token_index_test_utils import attach_req_token_index +from lightllm.common.attention.api import ServingConstraints from lightllm.models.deepseek_v4.triton_kernel.build_compressed_kv_indices import build_compressed_kv_indices -from lightllm.utils.envs_utils import get_env_start_args from lightllm.models.deepseek_v4.triton_kernel.swa_cache import get_swa_token_indices as _get_swa_token_indices +from lightllm.platform import get_device_runtime +from lightllm.utils.envs_utils import get_env_start_args +from unit_tests.common.attention.helpers import backend_context, make_spec, run_sparse, update_context +from unit_tests.common.dsv4_test_utils import DeepseekV4CacheConfig, DeepseekV4MemoryManager, DeepseekV4ReqManager +from unit_tests.common.req_token_index_test_utils import attach_req_token_index def _compressed_kv_pools(manager): @@ -245,11 +250,13 @@ def _write_packed_rows(buffer, layer_index, slots, packed): def _test_backend(requests, token_capacity=8192): - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4SparseMLAFp8DecodeBackend model = type("TestModel", (), {})() model.req_manager = requests model.att_state_manager = requests.memory.manager + model.device_runtime = get_device_runtime("cuda") + model.attention_requirements = ServingConstraints() model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] model.tp_rank_ = 0 model.mtp_previous_draft_models = [] @@ -274,8 +281,16 @@ def _test_backend(requests, token_capacity=8192): enable_decode_microbatch_overlap=False, mtp_mode=None, mtp_step=0, + mtp_dynamic_verify=False, + page_size=256, + ) + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda _: 0) + model.swa_prefill_workspace = get_swa_prefill_workspace(model) + requests._test_model = model + return DeepseekV4SparseMLAFp8DecodeBackend( + make_env(model), make_spec(cache=replace(make_spec().cache, page_size=256)) ) - return DeepseekV4FlashMlaFp8SparseAttBackend(model) def _prepare_prefill(requests, req_ids, starts, ends, microbatch_index=0): @@ -283,20 +298,11 @@ def _prepare_prefill(requests, req_ids, starts, ends, microbatch_index=0): if not hasattr(requests, "_test_backend"): requests._test_backend = _test_backend(requests) requests._test_prefill_states = {} - state = requests._test_backend.create_att_prefill_state(None) - state.prepare_swa_workspace( - SimpleNamespace( - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - b_req_idx=req_ids, - b_ready_cache_len=starts, - b_seq_len=ends, - microbatch_index=microbatch_index, - ) - ) + from lightllm.models.deepseek_v4.attention.sparse_mla import prepare_swa_workspace + + buffer, capacity = requests._test_model.swa_prefill_workspace + batch = SimpleNamespace(b_req_idx=req_ids, b_ready_cache_len=starts, b_seq_len=ends) + state = prepare_swa_workspace(requests.swa_state, buffer, capacity, batch, microbatch_index) requests._test_prefill_states.update((req, state) for req in req_ids.tolist()) if not hasattr(requests, "_test_prefill_ends"): requests._test_prefill_ends = {} @@ -441,8 +447,8 @@ def cpu(data): def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): from lightllm.common.basemodel.batch_objs import ModelInput from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel - from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel + from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel manager, requests = cache req_idx, slots = _reserve(manager, requests, 512) @@ -453,6 +459,8 @@ def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): model = model_cls.__new__(model_cls) model.req_manager = requests model.att_state_manager = requests.memory.manager + model.device_runtime = get_device_runtime("cuda") + model.attention_requirements = ServingConstraints() model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] model.tp_rank_ = 0 model.mtp_previous_draft_models = [] @@ -523,7 +531,7 @@ def test_prefill_checkpoint_reads_precede_window_retention(cache): _write_swa(requests, 2, req, positions, packed.flip(0)) _finish_prefill(requests, [req]) # Reuse the scratch: only fixed request windows and CPU snapshots remain. - requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model)[0].zero_() + requests._test_model.swa_prefill_workspace[0].zero_() requests.save_state(_checkpoint_req(req, 2048), index, states) saved_swa = states.get_state_cache(index)[0] torch.testing.assert_close(saved_swa[0].cuda(), packed[-requests.sliding_window :], rtol=0, atol=0) @@ -802,11 +810,15 @@ def test_compressor_checkpoint_continuation_matches_full_prefill(cache, ratio, l from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace manager, requests = cache - workspace = DeepseekV4Workspace(SimpleNamespace(tp_world_size_=1, batch_max_tokens=length)) - workspace.c128_compress_scratch = torch.empty((1, length, 512), dtype=torch.float32, device="cuda") source, src = _reserve(manager, requests, length) destination, dst = _reserve(manager, requests, length) _prepare_swa(requests, source, 0, length) + model = requests._test_model + model.disable_cudagraph = True + model._align_decode_batch_size = lambda count: count + model.mtp_manager.get_decode_tokens_per_request = lambda draft: get_env_start_args().mtp_step + 1 + workspace = DeepseekV4Workspace(model) + workspace.init_c128_compress_scratch(model) head_dim = requests.cache_config.indexer_head_dim if is_indexer else requests.cache_config.head_dim scores = torch.randn((length, head_dim * (4 if ratio == 4 else 2)), dtype=torch.float32, device="cuda") cos = torch.ones((length, 32), dtype=torch.float32, device="cuda") @@ -942,9 +954,9 @@ def build(): indirect=["cache"], ) def test_swa_ring_attention_matches_full_cache_after_speculative_retries(cache, dspark): - from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + from lightllm.models.deepseek_v4.attention import sparse_mla as attention from lightllm.models.deepseek_v4.triton_kernel.build_dspark_swa_indices import build_dspark_swa_indices + from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices torch.manual_seed(1593) manager, requests = cache @@ -986,29 +998,10 @@ def build(): else: build_swa_indices(req_idx, query_positions, requests, indices, lengths, write_slots) - backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend( - type( - "Model", - (), - { - "config": dict( - num_attention_heads=64, - compress_ratios=[0, 0, 0], - head_dim=512, - qk_rope_head_dim=64, - index_head_dim=128, - sliding_window=128, - ), - "layers_infer": [SimpleNamespace(layer_num_=2)], - "tp_rank_": 0, - "tp_world_size_": 1, - "req_manager": requests, - }, - )() - ) + backend = _test_backend(requests) state_info = SimpleNamespace(mem_manager=manager) - state = backend.create_att_decode_state(state_info) - state.init_state() + state = backend_context(type(backend), "decode", settings=backend) + update_context(state, state_info) params = dict( compress_ratio=0, head_dim_v=512, @@ -1017,13 +1010,13 @@ def build(): swa_lengths=lengths, attn_sink=torch.zeros(64, device="cuda"), ) - control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) + control = params q = torch.randn((token_num, 64, 512), dtype=torch.bfloat16, device="cuda") * 0.1 packed_kv = requests.get_att_input_params(2) def forward(): build() - return state.decode_att(q, packed_kv, None, att_control=control) + return run_sparse(state, state_info, q, packed_kv, control) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -1034,20 +1027,9 @@ def forward(): with torch.cuda.graph(graph, stream=stream): actual = forward() - reference_state = backend.create_att_decode_state(state_info) - reference_state.init_state() - reference_kv = backend.get_att_input_params( - SimpleNamespace( - req_manager=requests, - swa_state=requests.swa_state, - get_compressor_state=requests.get_compressor_state, - get_compressed_kv_state=requests.get_compressed_kv_state, - get_compressed_kv_view=requests.get_compressed_kv_view, - is_prefill=True, - prefill_att_state=requests._test_prefill_states[source], - ), - 2, - ) + reference_state = backend_context(type(backend), "decode", settings=backend) + update_context(reference_state, state_info) + reference_kv = (requests._test_prefill_states[source].swa_buffer[2].view(-1, 74880), None) committed_end = 512 # Both sides of page boundaries, repeated retries, and several complete ring wraps. for start in (512, 513, 638, 639, 640, 640, 641, 767, 768, 1023, 1279, 1280, 1919, 1920, 2047): @@ -1080,9 +1062,7 @@ def forward(): ) reference_indices = reference_indices.to(torch.int32) reference_params = dict(params, swa_indices=reference_indices.unsqueeze(1), swa_lengths=reference_lengths) - expected = reference_state.decode_att( - q, reference_kv, None, att_control=attention.AttControl(nsa_decode=True, nsa_decode_dict=reference_params) - ) + expected = run_sparse(reference_state, state_info, q, reference_kv, reference_params) torch.testing.assert_close(actual, expected, rtol=0, atol=0, msg=f"SWA mismatch at position {start}") torch.testing.assert_close(write_slots, _swa_slots(requests, destination, logical_positions).to(torch.int32)) # Accept one position; speculative future rows stay dirty until the next write. @@ -1096,8 +1076,9 @@ def forward(): @pytest.mark.parametrize("hit_extra_page", [False, True]) def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash_page_size, hit_extra_page): from sortedcontainers import SortedDict - from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context, CacheTier + from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache + from lightllm.server.router.model_infer.infer_batch import CacheTier, InferReq, g_infer_context from lightllm.utils.kv_cache_utils import compute_token_list_hash manager, requests = cache @@ -1335,18 +1316,22 @@ def select(): @pytest.mark.parametrize("is_prefill", [False, True]) @pytest.mark.parametrize("owned_layer_ids", [(0, 1, 2), (0, 2), (2,)]) def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill, owned_layer_ids): + from lightllm.common.basemodel.attention import ModelAttention from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace manager, requests = cache first, first_slots = _reserve(manager, requests, 512) second, second_slots = _reserve(manager, requests, 512) - backend = _test_backend(requests) - model = backend.model + _test_backend(requests) + model = requests._test_model model.layers_infer = [layer for layer in model.layers_infer if layer.layer_num_ in owned_layer_ids] has_c128_layer = 1 in owned_layer_ids model.graph = SimpleNamespace(can_run=lambda *args: True, graph_max_len_in_batch=16384) model.dsv4_workspace = DeepseekV4Workspace(model) + model.attention = ModelAttention.__new__(ModelAttention) + model.attention.groups = () + model.attention.generation = model.att_state_manager.storage_generation rope = torch.zeros((1024, 32), device="cuda") model._cos_cached_sliding = model._sin_cached_sliding = rope model._cos_cached_compress = model._sin_cached_compress = rope @@ -1361,15 +1346,16 @@ def make_state(lengths, microbatch_index): state.first_c4_logical_layer_id = 0 state.batch_size = 2 state.max_kv_seq_len = max(lengths) + state.max_q_seq_len = 1 + state.total_token_num = sum(lengths) state.input_ids = torch.zeros(2, dtype=torch.int64, device="cuda") state.b_req_idx = torch.tensor([first, second], dtype=torch.int32, device="cuda") state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") if is_prefill: state.b_ready_cache_len = state.b_seq_len - 1 - state.prefill_att_state = backend.create_att_prefill_state(state) - else: - state.decode_att_state = backend.create_att_decode_state(state) state.init_some_extra_state(model) + with model.attention.session(state): + pass return state captured_state = make_state([130, 258], 0) @@ -1400,7 +1386,9 @@ def make_state(lengths, microbatch_index): if is_prefill: captured_state.copy_for_prefill_cuda_graph(updated) else: - captured_state.copy_for_cuda_graph(updated) + for target, source in captured_state.graph_input_copies(updated): + target.copy_(source, non_blocking=True) + captured_state.attention.plan_replay(updated.attention) graph.replay() torch.testing.assert_close(swa_output, updated.dsv4_swa_indices) if has_c128_layer: @@ -1416,9 +1404,9 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, import lightllm.models.deepseek_v4.triton_kernel.build_swa_indices as swa_kernel import lightllm.models.deepseek_v4_dspark.infer_struct as dspark_state + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4SparseMLAFp8DecodeBackend from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace from lightllm.models.deepseek_v4_dspark.infer_struct import DeepseekV4DSparkInferStateInfo - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend manager, requests = cache req_idx, slots = _reserve(manager, requests, 512) @@ -1454,7 +1442,7 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, att_state_role="main", layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], ) - backend = DeepseekV4FlashMlaFp8SparseAttBackend(model) + model.swa_prefill_workspace = get_swa_prefill_workspace(model) model.block_size = 3 model.dsv4_workspace = DeepseekV4Workspace(model) rope = torch.zeros((512, 32), device="cuda") @@ -1474,9 +1462,7 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, state.b_req_idx = torch.tensor([req_idx], dtype=torch.int32, device="cuda") state.b_seq_len = torch.tensor([260], dtype=torch.int32, device="cuda") state.b_ready_cache_len = torch.tensor([257], dtype=torch.int32, device="cuda") - state.prefill_att_state = backend.create_att_prefill_state(state) else: - state.decode_att_state = backend.create_att_decode_state(state) state.b_req_idx = torch.full((3,), req_idx, dtype=torch.int32, device="cuda") state.b_seq_len = torch.arange(258, 261, dtype=torch.int32, device="cuda") if mode == "commit": @@ -1495,7 +1481,7 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, state.dsv4_swa_lengths, torch.full_like(state.dsv4_swa_lengths, 131 if proposal else 128) ) expected_slots = _get_swa_token_indices( - requests, req_idx, state.position_ids, prefill_state=state.prefill_att_state if state.is_prefill else None + requests, req_idx, state.position_ids, prefill_state=state.swa_workspace if state.is_prefill else None ).to(torch.int32) torch.testing.assert_close(state.dsv4_swa_write_slots, expected_slots) requests.free([req_idx], slots) @@ -1525,11 +1511,13 @@ def test_empty_dspark_commit_is_local_but_proposal_keeps_common_padding_path(mon ) @pytest.mark.parametrize("exhaust_small_pool", [False, True]) @pytest.mark.parametrize("token_index_mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("delay_producer", [False, True]) def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( - cache, monkeypatch, length, prefix, exhaust_small_pool, token_index_mode + cache, monkeypatch, length, prefix, exhaust_small_pool, token_index_mode, delay_producer ): """Offload after the runtime ring is overwritten, then resume through the common loader.""" from collections import deque + from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.router.model_infer.infer_batch import InferReq from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common @@ -1641,6 +1629,9 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( module.page_ready_buffer = torch.empty(8, dtype=torch.bool, device="cuda") module.gloo_group = module.filter_group = module.init_sync_group = None module.need_sync_compute_stream = lambda: False + if delay_producer: + torch.cuda._sleep(200_000_000) + _randomize_compressed_kv(manager) assert module.offload_finished_reqs_to_cpu_cache([req]) == [] assert req.cpu_cache_offload_task_status.is_running() torch.cuda.synchronize() @@ -1741,8 +1732,8 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer - from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa manager, requests = cache req, held = _reserve(manager, requests, 2048) @@ -1753,9 +1744,21 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): freqs = torch.complex(angles.cos(), angles.sin()) layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) layer.layer_num_, layer.eps_, layer.freqs_cis = 2, 1e-6, freqs - layer_weight = SimpleNamespace(kv_norm_=SimpleNamespace(weight=weight)) + layer.compress_ratio, layer.v_head_dim, layer.softmax_scale = 0, 512, 512 ** -0.5 + layer_weight = SimpleNamespace( + kv_norm_=SimpleNamespace(weight=weight), + attn_sink_=SimpleNamespace(weight=torch.zeros(64, device="cuda")), + ) for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): _prepare_swa(requests, req, start, end) + if start == 0: + from lightllm.common.attention.api import AttentionFamily + + layer.prefill_attention = SimpleNamespace( + resources=layer.bind_attention_resources( + requests._test_model, layer_weight, layer.attention_adapter(AttentionFamily.SPARSE_MLA) + ) + ) positions = torch.arange(start, end, device="cuda", dtype=torch.int32) indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") lengths = torch.empty(positions.numel(), dtype=torch.int32, device="cuda") @@ -1781,7 +1784,7 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): get_compressed_kv_state=requests.get_compressed_kv_state, get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, - prefill_att_state=requests._test_prefill_states[req], + swa_workspace=requests._test_prefill_states[req], dsv4_swa_write_slots=slots, dsv4_swa_runtime_write_slots=runtime, position_ids=positions, @@ -1880,8 +1883,8 @@ def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_le ready = spans[0][1] if cached_first_image else 0 manager, requests = cache req_idx, slots = _reserve(manager, requests, len(input_ids)) - backend = _test_backend(requests, token_capacity=128) - model = backend.model + _test_backend(requests, token_capacity=128) + model = requests._test_model model.has_vision = True model.config["vision_max_n_token"] = 384 model.dsv4_workspace = None @@ -1901,7 +1904,6 @@ def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_le state.b_ready_cache_len = torch.tensor([ready], dtype=torch.int32, device="cuda") state.max_kv_seq_len = len(input_ids) state.multimodal_params = [params.to_dict()] - state.prefill_att_state = backend.create_att_prefill_state(state) state.init_some_extra_state(model) assert state.dsv4_image_left.cpu().tolist() == expected_left[ready:] assert state.dsv4_image_right.cpu().tolist() == expected_right[ready:] @@ -1937,8 +1939,8 @@ def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace manager, requests = cache - backend = _test_backend(requests) - model = backend.model + _test_backend(requests) + model = requests._test_model model.has_vision = False rope = torch.zeros((8, 32), device="cuda") model._cos_cached_sliding = model._sin_cached_sliding = rope @@ -1963,7 +1965,6 @@ def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): infer_state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") infer_state.b_ready_cache_len = torch.zeros_like(infer_state.b_seq_len) infer_state.max_kv_seq_len = max(lengths) - infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) infer_state.init_some_extra_state(model) previous_debug_mode = torch.cuda.get_sync_debug_mode() @@ -2018,7 +2019,7 @@ def test_full_batch_unaligned_prefill_workspace_is_bounded_and_private(cache, mt required = request_count * ((end + page_token_num - 1) // page_token_num * page_token_num - first) requests._test_backend = _test_backend(requests, request_count * chunk_tokens) requests._test_prefill_states = {} - workspace, capacity = requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model) + workspace, capacity = requests._test_model.swa_prefill_workspace assert capacity >= required assert workspace.untyped_storage().data_ptr() != requests.req_to_swa.untyped_storage().data_ptr() workspace[:, capacity : 2 * capacity].fill_(173) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py index fbfc4a43a9..fce761bae8 100644 --- a/unit_tests/common/test_linear_state_contract.py +++ b/unit_tests/common/test_linear_state_contract.py @@ -1,14 +1,13 @@ import pytest import torch -from lightllm.common.attention_state.bounded_att import KDAState -from unit_tests.common.test_linear_state import make_state -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.attention.backends._shared.linear.replayssm_cursor import advance_cursor +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.test_linear_state import make_state @pytest.fixture @@ -118,7 +117,7 @@ def launch(*args, **kwargs): def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, module_name, cls, bv, precompute): import importlib - decode = importlib.import_module(f"lightllm.common.basemodel.attention.linear.{module_name}") + decode = importlib.import_module(f"lightllm.common.attention.backends.{module_name.split('_')[1]}.replay.decode") captured = _CaptureKernel() monkeypatch.setattr(decode, "_decode_replay", captured) owner = cls( @@ -239,3 +238,23 @@ def test_acceptance_keeps_request_order_with_strided_inputs(device): _commit_prefix_cpu(owner.mtp_state_index, reqs, lengths, lengths, history.cursors) assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] assert history.cursors.tolist() == [3, 0, 1, 0] + + +@pytest.mark.parametrize("cls", [ReplayGDNState, ReplayKDAState]) +def test_replay_capture_snapshot_restores_history_and_cursor(linear_config, cls): + from lightllm.common.attention_state import AttentionStateManager + + owner = cls(linear_config, request_capacity=3, mtp_step=2, device="cpu") + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", owner) + manager.finish_state_registration() + views = owner.capture_snapshot_views([1]) + for index, view in enumerate(views): + view.fill_(index + 1) + snapshot = manager.snapshot_capture_state(torch.tensor([1, 1]), "main") + for view in views: + view.zero_() + manager.restore_capture_state(snapshot) + for index, view in enumerate(views): + assert torch.all(view == index + 1) + assert owner.describe_layout()[0] in ("gdn_replay_conv_ssm", "kda_replay_conv_ssm") diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index 622c8087ad..6c27ade302 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -26,6 +26,7 @@ def single_rank_environment(monkeypatch: pytest.MonkeyPatch): def runtime_start_args(monkeypatch: pytest.MonkeyPatch): import json from dataclasses import asdict + from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py index 8394d18b48..95e463fd0b 100644 --- a/unit_tests/models/deepseek_v4/test_model.py +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -3,8 +3,8 @@ import pytest import torch -from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel @pytest.mark.parametrize("capture", [False, True], ids=["eager", "capture"]) @@ -264,6 +264,7 @@ def test_output_projection_matches_bf16_reference(monkeypatch, tp_size, rank): layer = DeepseekV4TransformerLayerWeight.__new__(DeepseekV4TransformerLayerWeight) BaseLayerWeight.__init__(layer) layer.prefix, layer.data_type_ = "layers.0", torch.bfloat16 + layer.num_fused_shared_experts = 0 layer.hidden, layer.n_heads, layer.head_dim = 128, 8, 64 layer.q_lora_rank, layer.o_lora_rank, layer.o_groups = 128, 128, 2 layer.get_quant_method = lambda name: None diff --git a/unit_tests/models/deepseek_v4/test_vision.py b/unit_tests/models/deepseek_v4/test_vision.py index ad931ac2b5..d2dcf442b0 100644 --- a/unit_tests/models/deepseek_v4/test_vision.py +++ b/unit_tests/models/deepseek_v4/test_vision.py @@ -12,12 +12,10 @@ from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel -from lightllm.common.basemodel.attention_vit.sdpa.fp import SdpaVitAttBackend -from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual -from lightllm.models.deepseek_v4 import image_processor as ours +from lightllm.common.attention.backends.dense.sdpa.varlen import SdpaVarlenAttentionOp +from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual, image_processor as ours from lightllm.server import visualserver - REFERENCE_DIR = "/mtc/models/DeepSeek-V4-Flash-Vision-Exp/inference" REFERENCE_PROCESSOR_PATH = os.path.join(REFERENCE_DIR, "image_processor.py") REFERENCE_VISION_PATH = os.path.join(REFERENCE_DIR, "vision.py") @@ -201,7 +199,7 @@ def _tiny_vision_args(): @pytest.mark.parametrize("n_vit_h,n_vit_w", [(4, 4), (5, 3)]) def test_vit_matches_reference(reference_vision, n_vit_h, n_vit_w, monkeypatch): - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", SdpaVitAttBackend) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", SdpaVarlenAttentionOp()) torch.manual_seed(0) args = _tiny_vision_args() model = visual.ViT(args).eval() @@ -348,3 +346,23 @@ def fake_read_shm(name): assert torch.equal(types[:3], torch.full((3,), ours.IMAGE_PAD, dtype=torch.int64)) assert types[3] == ours.IMAGE_START assert types[-1] == ours.IMAGE_END + + +def test_attention_uses_sequence_boundaries(monkeypatch): + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", SdpaVarlenAttentionOp()) + torch.manual_seed(17) + attention = visual.Attention(_tiny_vision_args()).eval() + x = torch.randn(7, 16, dtype=torch.bfloat16) + cos = torch.ones(7, 1, 2) + sin = torch.zeros_like(cos) + cu = torch.tensor([0, 3, 7], dtype=torch.int32) + with torch.inference_mode(), sdpa_kernel(SDPBackend.MATH): + q, k, v = [t.reshape(7, 4, 4) for t in attention.wqkv(x).chunk(3, dim=-1)] + parts = [] + for start, end in ((0, 3), (3, 7)): + qi, ki, vi = [t[start:end].float().transpose(0, 1) for t in (q, k, v)] + weights = (qi @ ki.transpose(-1, -2) * 0.5).softmax(-1) + parts.append((weights @ vi).transpose(0, 1).to(x.dtype)) + expected = attention.wo(torch.cat(parts).reshape(7, 16)) + actual = attention(x, cos, sin, cu) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) diff --git a/unit_tests/models/deepseek_v4/test_workspace.py b/unit_tests/models/deepseek_v4/test_workspace.py index 2bc4a56719..0e1b243253 100644 --- a/unit_tests/models/deepseek_v4/test_workspace.py +++ b/unit_tests/models/deepseek_v4/test_workspace.py @@ -1,5 +1,7 @@ from types import SimpleNamespace +from lightllm.common.attention.api import AttentionPhase + import pytest import torch @@ -8,7 +10,6 @@ from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -17,11 +18,11 @@ def test_draft_initialization_preserves_captured_prefill_buffer(monkeypatch, dra args = SimpleNamespace(llm_kv_type="fp8kv_dsa") monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_env_start_args", lambda: args) - def backend_class(**kwargs): - return lambda model: SimpleNamespace(tp_q_head_num=2, padded_q_head_num=4) + def initialize(model): + backend = SimpleNamespace(tp_q_head_num=2, padded_q_head_num=4) + return SimpleNamespace(groups=({AttentionPhase.PREFILL: backend, AttentionPhase.DECODE: backend},)) - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_prefill_att_backend_class", backend_class) - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_decode_att_backend_class", backend_class) + monkeypatch.setattr("lightllm.common.basemodel.attention.ModelAttention", initialize) main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) main.batch_max_tokens, main.tp_world_size_, main.data_type = 8, 1, torch.bfloat16 main.config = {"head_dim": 512, "block_size": 3, "dspark_noise_token_id": 1} @@ -36,7 +37,8 @@ def backend_class(**kwargs): "_sin_cached_compress", ): setattr(main, name, torch.ones((1, 32), device="cuda")) - main._init_att_backend() + main.args = args + main.prepare_attention_backends() buffer = main.dsv4_workspace.flashmla_prefill_q buffer[:, :2].fill_(2) graph = torch.cuda.CUDAGraph() @@ -48,7 +50,8 @@ def backend_class(**kwargs): draft.layers_infer = [SimpleNamespace(compress_ratio=0)] draft.context_wkv_weight = object() draft._init_custom() - draft._init_att_backend() + draft.args = args + draft.prepare_attention_backends() assert draft.dsv4_workspace is main.dsv4_workspace assert main.dsv4_workspace.flashmla_prefill_q is buffer diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 6da740353d..9f77b75e6b 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -1,4 +1,3 @@ -from unit_tests.common.att_resources.model_helpers import build_hybrid_resources import dataclasses from types import SimpleNamespace @@ -6,17 +5,16 @@ import torch from lightllm.common.attention_state.allocators import FullAttKVAllocator - from lightllm.common.attention_state.full_att import ( - PackedMLAFullAttentionState, PackedMLAAttentionConfig, + PackedMLAFullAttentionState, PackedMLALogicalConfig, ) from lightllm.common.req_manager import ReqManager from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name - +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index d0e6dcb9bd..e842f8cc8a 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -1,24 +1,12 @@ -import weakref -from types import SimpleNamespace - +from unit_tests.common.attention.helpers import plan, build +from lightllm.common.attention.session import AttentionSession import pytest import torch -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.nsa.glm5_next import ( - Glm5NextSparseAttBackend, - Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - Fp8Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next +from lightllm.common.attention.api import AttentionBatch, AttentionPhase, ContextCapacity +from lightllm.common.attention.api.workspace import AttentionWorkspace from lightllm.models.glm5_next.triton_kernel.kpool import expand_topk, prefill_nsa_params - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -27,231 +15,94 @@ def tensor(values): @pytest.mark.parametrize("counts", [(3, 0, 1, 2), (257, 0, 1, 513), (1,) * 129]) -def test_prefill_metadata_reuses_existing_prefixes_during_graph_replay(counts): +def test_prefill_metadata_is_graph_safe(counts): batch, queries = len(counts), sum(counts) positions = torch.empty(queries, device="cuda", dtype=torch.int32) req = torch.empty(batch, device="cuda", dtype=torch.int32) - cu_q, cu_kv = torch.empty(batch + 1, device="cuda", dtype=torch.int32), torch.empty( - batch + 1, device="cuda", dtype=torch.int32 - ) - - def update(counts, phase): - q_prefix, kv_prefix, visible, starts, requests = [0], [0], [], [], [] - req_values = [(batch - i + phase) % (batch + 1) for i in range(batch)] - for i, count in enumerate(counts): - prefix = 1000000 + i * 7 + phase - visible.extend(range(prefix + 1, prefix + count + 1)) - starts.extend([kv_prefix[-1]] * count) - requests.extend([req_values[i]] * count) - q_prefix.append(q_prefix[-1] + count) - kv_prefix.append(kv_prefix[-1] + prefix + count) - positions.copy_(tensor(visible) - 1) - req.copy_(tensor(req_values)) - cu_q.copy_(tensor(q_prefix)) - cu_kv.copy_(tensor(kv_prefix)) - return tensor(visible), tensor(requests), tensor(starts) + cu_q = torch.empty(batch + 1, device="cuda", dtype=torch.int32) + cu_kv = torch.empty(batch + 1, device="cuda", dtype=torch.int32) + + q_prefix, kv_prefix, visible, starts, requests = [0], [0], [], [], [] + for index, count in enumerate(counts): + prefix = 1000000 + index * 7 + visible.extend(range(prefix + 1, prefix + count + 1)) + starts.extend([kv_prefix[-1]] * count) + requests.extend([index] * count) + q_prefix.append(q_prefix[-1] + count) + kv_prefix.append(kv_prefix[-1] + prefix + count) + positions.copy_(tensor(visible) - 1) + req.copy_(tensor(range(batch))) + cu_q.copy_(tensor(q_prefix)) + cu_kv.copy_(tensor(kv_prefix)) - expected = update(counts, 0) + expected = (tensor(visible), tensor(requests), tensor(starts)) actual = prefill_nsa_params(positions, req, cu_q, cu_kv) for observed, reference in zip(actual, expected): torch.testing.assert_close(observed, reference, rtol=0, atol=0) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = prefill_nsa_params(positions, req, cu_q, cu_kv) - expected = update(counts[::-1], 11) - graph.replay() - for observed, reference in zip(actual, expected): - torch.testing.assert_close(observed, reference, rtol=0, atol=0) @pytest.mark.parametrize("dense", [False, True]) -def test_topk_direct_request_lookup_preserves_mtp_hold_and_strided_tables(dense): +def test_topk_request_lookup_preserves_strided_tables(dense): table = torch.arange(8 * 70, device="cuda", dtype=torch.int32).reshape(8, 70)[::2, ::2] req = tensor([2, 2, 0, 3, 3]) lengths = tensor([5, 7, 9, 2, 2] if dense else [17, 18, 9, 2, 2]) groups = tensor([[3, 1], [1, 0], [0, 1], [0, 0], [0, 0]]) - - def check(actual): - mem, relative = actual - expected_mem, expected_relative = torch.full_like(mem, -1), torch.full_like(relative, -1) - for i, (request, length) in enumerate(zip(req.tolist(), lengths.tolist())): - tokens = ( - list(range(length)) - if dense - else [group * 4 + j for group in groups[i].tolist()[: min(length // 4, 2)] for j in range(4)] - + list(range(length // 4 * 4, length)) - ) - expected_relative[i, : len(tokens)] = tensor(tokens) - expected_mem[i, : len(tokens)] = table[request, tensor(tokens).long()] - torch.testing.assert_close(mem, expected_mem, rtol=0, atol=0) - torch.testing.assert_close(relative, expected_relative, rtol=0, atol=0) - - check(expand_topk(groups, lengths, req, table, 8, dense=dense)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = expand_topk(groups, lengths, req, table, 8, dense=dense) - req.copy_(tensor([0, 0, 2, 3, 3])) - table.add_(10000) - lengths[:2].add_(1) - graph.replay() - check(actual) - - -@pytest.mark.parametrize("fp8", [False, True]) -def test_decode_aliases_infer_state_without_ragged_allocation(monkeypatch, fp8): - import lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke as legacy - - def forbidden(*args, **kwargs): - raise AssertionError("GLM must not generate or copy a full ragged index") - - monkeypatch.setattr(legacy, "gen_nsa_ks_ke", forbidden) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - model = type( - "Model", - (), - { - "config": {"index_kpool": 4, "index_topk": 2048}, - "batch_max_tokens": 4, - "max_req_num": 1, - "max_seq_length": 32, - "att_state_role": "decode_test", - "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="decode"), - }, - )() - backend = (Fp8Glm5NextSparseAttBackend if fp8 else Glm5NextSparseAttBackend)(model) - assert not hasattr(backend, "ragged_mem_buffers") - infer = InferStateInfo() - infer.b_req_idx = tensor([2, 2, 0, 3, 3]) - infer.b_seq_len = tensor([999999, 1000000, 65536, 2, 2]) - cls = Fp8Glm5NextSparseDecodeState if fp8 else Glm5NextSparseDecodeState - state = cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - state.init_state() - assert state.q_kv_seq_len is infer.b_seq_len - assert state.q_req_idx is infer.b_req_idx - assert not hasattr(state, "ragged_mem_index") and not hasattr(state, "ks") - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - state.init_state() - infer.b_seq_len[:3].sub_(100) - infer.b_req_idx[:3].copy_(tensor([0, 0, 1])) - graph.replay() - assert state.q_kv_seq_len is infer.b_seq_len and state.q_req_idx is infer.b_req_idx + actual_mem, actual_relative = expand_topk(groups, lengths, req, table, 8, dense=dense) + + expected_mem, expected_relative = torch.full_like(actual_mem, -1), torch.full_like(actual_relative, -1) + for row, (request, length) in enumerate(zip(req.tolist(), lengths.tolist())): + tokens = ( + list(range(length)) + if dense + else [group * 4 + offset for group in groups[row].tolist()[: min(length // 4, 2)] for offset in range(4)] + + list(range(length // 4 * 4, length)) + ) + expected_relative[row, : len(tokens)] = tensor(tokens) + expected_mem[row, : len(tokens)] = table[request, tensor(tokens).long()] + torch.testing.assert_close(actual_mem, expected_mem, rtol=0, atol=0) + torch.testing.assert_close(actual_relative, expected_relative, rtol=0, atol=0) @pytest.mark.parametrize( - "queries,heads,total_kv,chunk_size", + "reference", [ - (1, 64, 65536, 2097152), - (2, 64, 65536, 2097152), - (3, 8, 65536, 2097152), - (9, 64, 512, 2097152), - (9, 64, 131073, 2097152), - (4, 16, 40, 3), - (4, 64, 40, 7), + "backends.sparse_mla.flashmla.fp8:FlashMLASparseMLAFp8PrefillBackend", + "lightllm.models.glm5_next.attention.sparse_mla:GlmSparseMLAPrefillBackend", + "lightllm.models.glm5_next.attention.fp8_sparse_mla:GlmSparseMLAFp8PrefillBackend", ], ) -def test_fp8_prefill_reads_request_table_without_ragged_buffer(monkeypatch, queries, heads, total_kv, chunk_size): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) - - torch.manual_seed(1685) - batch = min(queries, 2) - counts = [queries] if batch == 1 else [queries // 2, queries - queries // 2] - seqs = [total_kv] if batch == 1 else [total_kv // 2, total_kv - total_kv // 2] - req = tensor([2, 0][:batch]) - table = torch.full((4, 2 * (max(seqs) + 7)), -1, device="cuda", dtype=torch.int32)[:, ::2] - physical = torch.randperm(total_kv + 32, device="cuda", dtype=torch.int32)[:total_kv] - offset = 0 - for request, length in zip(req.tolist(), seqs): - table[request, :length] = physical[offset : offset + length] - offset += length - q = torch.randn(queries, heads, 512, device="cuda", dtype=torch.bfloat16) - source = torch.randn(total_kv + 32, 1, 512, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros(total_kv + 32, 1, 720, device="cuda", dtype=torch.uint8) - write_fp8_e8m0_packed_kv(source, torch.arange(total_kv + 32, device="cuda", dtype=torch.int32), packed) - prefill_kv = torch.randn(queries, 1, 512, device="cuda", dtype=torch.bfloat16) - visible, query_req, current_positions, q_prefix, kv_prefix = [], [], [], [0], [0] - for request, length, count in zip(req.tolist(), seqs, counts): - visible.extend(range(length - count + 1, length + 1)) - query_req.extend([request] * count) - current_positions.extend(range(length - count, length)) - q_prefix.append(q_prefix[-1] + count) - kv_prefix.append(kv_prefix[-1] + length) - query_req, current_positions = tensor(query_req), tensor(current_positions) - mem_index = table[query_req.long(), current_positions.long()] - topk = torch.full((queries, 2176), -1, device="cuda", dtype=torch.int32) - for i, length in enumerate(visible): - topk[i, :7] = tensor([0, length - 1, 0, 1, length // 2, length - 1, -1]) - topk_mem = torch.empty_like(topk) - - def refresh_physical(): - topk_mem.copy_(table[query_req.long()[:, None], topk.clamp_min(0).long()]) - topk_mem.masked_fill_(topk < 0, -1) - mem_index.copy_(table[query_req.long(), current_positions.long()]) - - refresh_physical() - infer = InferStateInfo() - infer.mem_index, infer.b_req_idx = mem_index, req - infer.position_ids = tensor(visible) - 1 - infer.b1_cu_q_seq_len, infer.b1_cu_kv_seq_len = tensor(q_prefix), tensor(kv_prefix) - infer.total_token_num = total_kv - infer.req_manager = SimpleNamespace(req_to_token_indexs=table) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - model = type( - "Model", - (), - { - "batch_max_tokens": queries, - "prefill_graph": None, - "max_req_num": batch, - "max_seq_length": max(seqs), - "att_state_role": "indices_test", - "config": {"index_topk": 2048}, - "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="normal"), - }, - )() - backend = Fp8Glm5NextSparseAttBackend(model) - state = backend.create_att_prefill_state(infer) - state.init_state() - assert not hasattr(state, "ragged_mem_index") and not hasattr(backend, "ragged_mem_buffers") - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk, - "topk_mem_indices": topk_mem, - "prefill_cache_kv": prefill_kv, - "softmax_scale": 512 ** -0.5, - "kv_lora_rank": 512, - }, +def test_sparse_context_owns_query_metadata(reference): + from lightllm.common.attention.registry import load_backend_class + + lengths = torch.tensor([5, 4], device="cuda", dtype=torch.int32) + q_lengths = torch.tensor([2, 1], device="cuda", dtype=torch.int32) + batch = AttentionBatch( + phase=AttentionPhase.PREFILL, + batch_size=2, + token_num=3, + total_token_num=9, + max_q_seq_len=2, + max_kv_seq_len=5, + max_cache_len=3, + b_req_idx=tensor([0, 1]), + b_seq_len=lengths, + b_q_seq_len=q_lengths, + b1_cu_q_seq_len=tensor([0, 2, 3]), + b1_cu_kv_seq_len=tensor([0, 5, 9]), + req_to_token_indexs=torch.arange(8 * 16, device="cuda", dtype=torch.int32).reshape(8, 16), + hold_req_id=7, + b_ready_cache_len=tensor([3, 3]), + position_ids=tensor([3, 4, 3]), ) - - def run(): - return state.prefill_att(q, packed[..., :584], None, control) - - def check(output): - scales = packed[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) - values = packed[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) * scales - kv = torch.cat((values.reshape(-1, 1, 448).bfloat16(), packed[..., 448:576].view(torch.bfloat16)), -1) - kv[mem_index.long()] = prefill_kv - expected = [] - for i in range(queries): - keys = kv[topk_mem[i][topk_mem[i] >= 0].long(), 0].float() - expected.append((q[i].float() @ keys.T * 512 ** -0.5).softmax(-1) @ keys) - torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) - - check(run()) - workspace = backend._prefill_workspaces[0] - if workspace is not None: - assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = run() - graph.replay() - check(output) - q.normal_() - prefill_kv.normal_() - for request, length in zip(req.tolist(), seqs): - table[request, :length].copy_(table[request, :length].roll(1)) - topk[0].fill_(-1) - refresh_physical() - graph.replay() - check(output) + backend = object.__new__(load_backend_class(reference)) + context = AttentionSession.allocate( + backend, + ContextCapacity(phase=AttentionPhase.PREFILL, batch_size=2, token_num=3, max_kv_seq_len=5, max_q_seq_len=2), + capture=False, + workspace=AttentionWorkspace("cuda"), + ) + plan(backend, context, batch) + build(backend, context) + torch.testing.assert_close(context.q_kv_seq_len, tensor([4, 5, 4])) + torch.testing.assert_close(context.q_req_idx, tensor([0, 0, 1])) + torch.testing.assert_close(context.q_kv_start_loc, tensor([0, 0, 5])) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index a6a5c7c9f2..f379d7e3eb 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -5,35 +5,33 @@ import torch import triton -from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.nsa.glm5_next import ( - Glm5NextSparsePrefillState, - Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - Fp8Glm5NextSparsePrefillState, - Fp8Glm5NextSparseDecodeState, +from lightllm.models.glm5_next.attention.fp8_sparse_mla import GlmSparseMLAFp8DecodeBackend +from lightllm.models.glm5_next.attention.sparse_mla import GlmSparseMLADecodeBackend +from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill +from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + _build_prefill_row_table_kernel, + _gather_prefill_kv_kernel, ) from lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode import ( can_use_fp8_sparse_direct_decode, fp8_sparse_direct_decode, ) from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel -from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next -from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( - _gather_prefill_kv_kernel, +from lightllm.utils.backend_validator import ProbeResult, ProbeScope +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, ) - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture(autouse=True) -def current_cuda_device(monkeypatch): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) +pytest.importorskip("flash_mla") def _make_infer_state(**attributes): @@ -42,156 +40,68 @@ def _make_infer_state(**attributes): return state -def _make_workspace_backend(batch_max_tokens, overlap=False): - backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace( - batch_max_tokens=batch_max_tokens, - prefill_graph=None, - config={"index_topk": 2048, "index_kpool": 4}, - max_req_num=1, - max_seq_length=fp8_glm5_next.KV_GATHER_CHUNK_SIZE, - att_state_role="unit_test", - args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode="normal"), - ) - backend._prefill_workspaces = backend.allocate_prefill_workspaces(backend.model) - return backend - +def _make_workspace_backend(batch_max_tokens): + backend = object.__new__(GlmSparseMLAFp8DecodeBackend) + backend.prefill_token_capacity = batch_max_tokens + backend.prefill_graph_enabled = False + from lightllm.common.attention.api.workspace import AttentionWorkspace -class _PrefillCase: - def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=None, softmax_scale=0.0625): - device = "cuda" - tensor = lambda values: torch.tensor(values, device=device, dtype=torch.int32) - self.seq_lens, self.q_lens = seq_lens, q_lens - self.requests = list(range(2 * len(seq_lens) - 1, 0, -2)) - total, queries = sum(seq_lens), sum(q_lens) - capacity = capacity or total + 17 - self.req_to_token_indexs = torch.full( - (2 * len(seq_lens) + 1, max(seq_lens) + 7), -1, device=device, dtype=torch.int32 - ) - slots = torch.randperm(capacity, device=device, dtype=torch.int32)[:total] - cu_q, cu_kv, q_req_idx, positions = [0], [0], [], [] - for req, length, count in zip(self.requests, seq_lens, q_lens): - self.req_to_token_indexs[req, :length] = slots[cu_kv[-1] : cu_kv[-1] + length] - q_req_idx.extend([req] * count) - positions.extend(range(length - count, length)) - cu_q.append(cu_q[-1] + count) - cu_kv.append(cu_kv[-1] + length) - self.q_req_idx, self.position_ids = tensor(q_req_idx).long(), tensor(positions).long() - self.q = torch.randn(queries, heads, 512, device=device, dtype=torch.bfloat16) - self.source_kv = torch.randn(capacity, 1, 512, device=device, dtype=torch.bfloat16) - self.packed_kv = torch.zeros(capacity, 1, 720, device=device, dtype=torch.uint8) - self.all_slots = torch.arange(capacity, device=device, dtype=torch.int32) - write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) - self.prefill_cache_kv = torch.randn(queries, 1, 1024, device=device, dtype=torch.bfloat16)[..., ::2] - visible = tensor(positions) + 1 - self.topk_indices = torch.randint(max(seq_lens), (queries, width), device=device, dtype=torch.int32) - self.topk_indices.remainder_(visible[:, None]) - self.topk_indices.masked_fill_( - torch.arange(width, device=device)[None, :] >= visible.clamp_max(2048)[:, None], -1 - ) - self.topk_indices[:, ::19] = -1 - self.topk_indices[:, 0] = visible - 1 - self.topk_indices[:, 1] = visible - 1 - self.topk_mem_indices = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] - self.infer_state = _make_infer_state( - is_prefill=True, - req_manager=SimpleNamespace(req_to_token_indexs=self.req_to_token_indexs), - b_req_idx=tensor(self.requests), - b_seq_len=tensor(seq_lens), - b1_cu_q_seq_len=tensor(cu_q), - b1_cu_kv_seq_len=tensor(cu_kv), - position_ids=tensor(positions), - total_token_num=total, - mem_index=torch.empty(queries, device=device, dtype=torch.int32), - ) - self.backend = backend or _make_workspace_backend(queries) - _, self.att_state = self.new_state() - self.control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": self.topk_indices, - "topk_mem_indices": self.topk_mem_indices, - "prefill_cache_kv": self.prefill_cache_kv, - "softmax_scale": softmax_scale, - "kv_lora_rank": 512, - }, - ) - - def new_state(self, **overrides): - if overrides: - infer = _make_infer_state(**vars(self.infer_state)) - vars(infer).update(overrides) - else: - infer = self.infer_state - state = self.backend.create_att_prefill_state(infer) - state.init_state() - return infer, state - - def refresh(self): - self.infer_state.mem_index.copy_(self.req_to_token_indexs[self.q_req_idx, self.position_ids]) - self.topk_mem_indices.copy_( - self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()] - ) - self.topk_mem_indices.masked_fill_(self.topk_indices < 0, -1) - self.control.nsa_prefill_dict["topk_indices"] = self.topk_indices - - def run(self, state=None): - self.refresh() - active = self.att_state if state is None else state - return active.prefill_att(self.q, self.packed_kv[..., :584], None, self.control) - - def check(self, actual): - values = _dequantize_kv_cache(self.packed_kv) - values[self.infer_state.mem_index.long()] = self.prefill_cache_kv - physical = self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()].masked_fill( - self.topk_indices < 0, -1 - ) - expected = _torch_sparse_attention(self.q, values, physical, self.control.nsa_prefill_dict["softmax_scale"]) - torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) - - def mutate(self): - self.q.normal_() - self.prefill_cache_kv.normal_() - self.source_kv.normal_() - write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) - for req, length in zip(self.requests, self.seq_lens): - self.req_to_token_indexs[req, :length].copy_(self.req_to_token_indexs[req, :length].roll(1)) - self.topk_indices.copy_(self.topk_indices.roll(1, dims=1)) + backend._test_workspace = AttentionWorkspace("cuda") + return backend @pytest.mark.parametrize("draft", [False, True]) @pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) def test_main_and_native_mtp_backend_selection(monkeypatch, draft, kv_type): + from lightllm.common.attention.api import AttentionFamily, AttentionPhase, CacheLayout + from lightllm.common.basemodel import attention from lightllm.models.glm5_next.model import Glm5NextTpPartModel from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + from unit_tests.common.attention.helpers import make_spec, model_stub - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") model = object.__new__(Glm5NextMTPModel if draft else Glm5NextTpPartModel) - model.args = SimpleNamespace(llm_kv_type=kv_type, enable_prefill_microbatch_overlap=False, run_mode="normal") - model.config = {"index_topk": 2048, "index_kpool": 4} - model.graph_max_batch_size = 4 - model.max_seq_length = 128 - model.max_req_num = 2 - model.batch_max_tokens = 4 + vars(model).update( + vars( + model_stub( + llm_kv_type=kv_type, llm_prefill_att_backend="flashmla_sparse", llm_decode_att_backend="flashmla_sparse" + ) + ) + ) + model.is_mtp_draft_model = draft model.mtp_previous_draft_models = [] - model._init_att_backend() - infer_state = InferStateInfo() - prefill = model.prefill_att_backend.create_att_prefill_state(infer_state) - decode = model.decode_att_backend.create_att_decode_state(infer_state) - assert type(prefill) is (Fp8Glm5NextSparsePrefillState if kv_type == "fp8kv_dsa" else Glm5NextSparsePrefillState) - assert type(decode) is (Fp8Glm5NextSparseDecodeState if kv_type == "fp8kv_dsa" else Glm5NextSparseDecodeState) - infer_state.prefill_att_state = prefill - infer_state.decode_att_state = decode - infer_state.check_att_states() - assert prefill.infer_state is infer_state - assert decode.infer_state is infer_state - owner_ref = weakref.ref(infer_state) - del infer_state - assert owner_ref() is None - with pytest.raises(ReferenceError): - prefill.infer_state - with pytest.raises(ReferenceError): - decode.infer_state + model.config = {"index_topk": 2048, "index_kpool": 4} + model.batch_max_tokens = 8 + from lightllm.common.basemodel.attention_adapters import ADAPTERS, AttentionDescription + from unit_tests.common.attention.helpers import phase_class + + expected = GlmSparseMLAFp8DecodeBackend if kv_type == "fp8kv_dsa" else GlmSparseMLADecodeBackend + monkeypatch.setattr( + "lightllm.utils.backend_validator.probe_backend", + lambda *args, **kwargs: ProbeResult(True, ProbeScope.BACKEND_NUMERICS), + ) + monkeypatch.setattr(expected, "check_environment", classmethod(lambda cls, caps: None)) + monkeypatch.setattr( + ADAPTERS[AttentionFamily.SPARSE_MLA], + "describe", + lambda model, layer, phase, desc: make_spec( + desc.family, phase, cache=CacheLayout(kv_type, 1, torch.bfloat16, 1, 512, 512) + ), + ) + model.layers_infer = [ + SimpleNamespace( + layer_num_=0, + describe_attention=lambda phase: AttentionDescription(AttentionFamily.SPARSE_MLA), + attention_adapter=lambda family: ADAPTERS[family], + ) + ] + model.prepare_attention_backends() + for phase in AttentionPhase: + backend = model.attention.groups[0][phase] + assert type(backend) is phase_class(expected, phase) + context = backend_context(type(backend), phase.value, settings=backend) + assert context.backend_identity is backend + assert not hasattr(backend, "model") + assert not hasattr(context, "request") def _dequantize_kv_cache(packed_kv): @@ -258,29 +168,26 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid changed_values = (changed_groups / changed_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) assert torch.equal(packed_kv[mem_index, :, :448], changed_values.view(torch.uint8).reshape(11, 1, 448)) + indices = mem_index[[0, 2, 8, 2, 5, 1, 9, 3]] prefill_kv = source[8:].clone().add_(0.25) - gathered_kv = torch.empty((8, 1, 512), device="cuda", dtype=torch.bfloat16) - _gather_prefill_kv_kernel[(3,)]( + row_table = torch.full((31,), -1, device="cuda", dtype=torch.int32) + _build_prefill_row_table_kernel[(1,)](mem_index[8:], row_table, 3, BLOCK=256) + gathered_kv = torch.empty((indices.numel(), 1, 512), device="cuda", dtype=torch.bfloat16) + _gather_prefill_kv_kernel[(triton.cdiv(indices.numel(), 4),)]( packed_kv[..., :584], + indices, + row_table, prefill_kv, gathered_kv, - mem_index[None], - torch.zeros(1, device="cuda", dtype=torch.int32), - torch.tensor([0, 11], device="cuda", dtype=torch.int32), - torch.tensor([0, 3], device="cuda", dtype=torch.int32), - 2, - 8, + indices.numel(), packed_kv.stride(0), prefill_kv.stride(0), prefill_kv.stride(2), - req_to_token_indexs_stride_b=11, - req_to_token_indexs_stride_s=1, - batch_size=1, BLOCK_ROWS=4, ) expected = _dequantize_kv_cache(packed_kv) - expected[mem_index[8:].long()] = prefill_kv - assert torch.equal(gathered_kv, expected[mem_index[2:10].long()]) + expected[mem_index[8:]] = prefill_kv + assert torch.equal(gathered_kv, expected[indices.long()]) @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @@ -300,16 +207,12 @@ def test_fp8_nope_decode_and_cuda_graph(num_heads): indices = torch.randint(0, 4096, (9, 2176), device="cuda", dtype=torch.int32) lanes = torch.arange(2176, device="cuda") indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} infer_state = _make_infer_state() - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] - ) + state = backend_context(GlmSparseMLAFp8DecodeBackend, "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0]) def run(): - return state.decode_att((q, q[..., :0]), packed_kv[..., :584], None, control) + return run_sparse(state, infer_state, (q, q[..., :0]), packed_kv[..., :584], control) def check(actual): expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed_kv), indices) @@ -333,199 +236,444 @@ def check(actual): @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @pytest.mark.parametrize("chunk_size", [3, 7, 32768]) def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chunk_size): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) torch.manual_seed(53) - prefix = 4 if cached else 0 - case = _PrefillCase([prefix + 3] * 2, [3, 3], num_heads, width=128, capacity=64, softmax_scale=512 ** -0.5) - case.q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - case.topk_indices[0].fill_(-1) - case.topk_indices[1, :2] = prefix - case.topk_indices[3].fill_(-1) - case.topk_indices[3, 0] = prefix - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) - - -@pytest.mark.parametrize("num_heads", [8, 64]) -@pytest.mark.parametrize("chunk_size", [3, 32768]) -def test_fp8_uncached_prefill_reuses_current_kv(monkeypatch, num_heads, chunk_size): - from sgl_kernel.flash_attn import flash_attn_with_kvcache + # Two requests whose physical cache locations are unrelated to batch order. + prefix_len = 4 if cached else 0 + seq_len = prefix_len + 3 + ragged_mem_index = torch.randperm(64, device="cuda")[: 2 * seq_len].int() + prefill_rows = torch.cat( + (torch.arange(prefix_len, seq_len), torch.arange(seq_len + prefix_len, 2 * seq_len)) + ).cuda() + prefill_mem_index = ragged_mem_index[prefill_rows] + source = torch.randn(2 * seq_len, 1, 512, device="cuda", dtype=torch.bfloat16) + prefill_kv = source[prefill_rows].clone().add_(0.25) + packed_kv = torch.zeros(64, 1, 720, device="cuda", dtype=torch.uint8) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + kv_start_offsets = torch.tensor([0] * 3 + [seq_len] * 3, device="cuda", dtype=torch.int32) + topk_indices = torch.full((6, 128), -1, device="cuda", dtype=torch.int32) + for row in range(6): + length = prefix_len + row % 3 + 1 + topk_indices[row, :length] = torch.arange(length, device="cuda") + topk_indices[0].fill_(-1) + topk_indices[1, :2] = prefix_len + topk_indices[3].fill_(-1) + topk_indices[3, 0] = prefix_len + control = { + "topk_indices": topk_indices, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 512 ** (-0.5), + "kv_lora_rank": 512, + } + infer_state = _make_infer_state(max_cache_len=prefix_len, mem_index=prefill_mem_index) + state = backend_context( + GlmSparseMLAFp8DecodeBackend, + "prefill", + settings=_make_workspace_backend(q.shape[0]), + ks=kv_start_offsets, + ragged_mem_index=ragged_mem_index, + ) - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) - case = _PrefillCase([3, 3], [3, 3], num_heads, width=128) - case.prefill_cache_kv = torch.randn(6, 1, 1536, device="cuda", dtype=torch.bfloat16)[..., :512] - case.control.nsa_prefill_dict["prefill_cache_kv"] = case.prefill_cache_kv + def run(): + return run_sparse(state, infer_state, q, packed_kv[..., :584], control) - def check_kv_storage(*args, **kwargs): - assert kwargs["v_cache"].untyped_storage().data_ptr() == case.prefill_cache_kv.untyped_storage().data_ptr() - return flash_attn_with_kvcache(*args, **kwargs) + def check(out): + expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] + expected_kv[prefill_rows] = prefill_kv + indices = torch.where(topk_indices >= 0, topk_indices + kv_start_offsets[:, None], -1) + expected = _torch_sparse_attention(q, expected_kv, indices, control["softmax_scale"]) + torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) - monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", check_kv_storage) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - case.mutate() + check(run()) + graph, out = _capture_cuda_graph(run) graph.replay() - case.check(output) + check(out) + q.normal_() + source.normal_() + prefill_kv.copy_(source[prefill_rows] + 0.25) + ragged_mem_index.copy_(ragged_mem_index.roll(1)) + prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) + topk_indices.copy_(topk_indices.roll(1, dims=1)) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + graph.replay() + check(out) @pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) @pytest.mark.parametrize("softmax_scale", [0.0625, 512 ** -0.5]) def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 129) + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", 129) torch.manual_seed(53) - case = _PrefillCase([1025], [9], num_heads, capacity=2048, softmax_scale=softmax_scale) + num_tokens = 1025 + source = torch.randn(num_tokens, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + ragged_mem_index = torch.randperm(2048, device="cuda")[:num_tokens].int() + prefill_rows = torch.tensor([1, 15, 128, 129, 255, 257, 511, 512, 1024], device="cuda") + prefill_mem_index = ragged_mem_index[prefill_rows].long() + prefill_kv = torch.randn(9, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + packed_kv = torch.zeros(2048, 1, 720, device="cuda", dtype=torch.uint8) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") - case.topk_indices.masked_fill_(torch.arange(2176, device="cuda")[None, :] >= lengths[:, None], -1) - case.topk_indices[1, 0] = 1016 - case.topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") - case.topk_indices[3, :127] = torch.randint(1016, 1025, (127,), device="cuda", dtype=torch.int32) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) + topk_indices = torch.randint(num_tokens, (9, 2176), device="cuda", dtype=torch.int32) + lanes = torch.arange(2176, device="cuda") + topk_indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) + topk_indices[1, 0] = 15 + topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") + topk_indices[3, :127] = prefill_rows[torch.randint(9, (127,), device="cuda")].int() + control = { + "topk_indices": topk_indices, + "prefill_cache_kv": prefill_kv, + "softmax_scale": softmax_scale, + "kv_lora_rank": 512, + } + infer_state = _make_infer_state(mem_index=prefill_mem_index) + state = backend_context( + GlmSparseMLAFp8DecodeBackend, + "prefill", + settings=_make_workspace_backend(q.shape[0]), + ks=torch.zeros(9, device="cuda", dtype=torch.int32), + ragged_mem_index=ragged_mem_index, + ) + def run(): + return run_sparse(state, infer_state, q, packed_kv[..., :584], control) -@pytest.mark.parametrize("num_q", [64, 129]) -@pytest.mark.parametrize("num_heads", [8, 64]) -def test_fp8_prefill_uses_flashattention_and_request_table_graph(monkeypatch, num_q, num_heads): - from sgl_kernel.flash_attn import flash_attn_with_kvcache + def check(actual): + expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] + expected_kv[prefill_rows] = prefill_kv + expected = _torch_sparse_attention(q, expected_kv, topk_indices, softmax_scale) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + check(run()) + graph, output = _capture_cuda_graph(run) + graph.replay() + check(output) + q.normal_() + source.normal_() + prefill_kv.normal_() + ragged_mem_index.copy_(ragged_mem_index.roll(1)) + prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) + topk_indices.copy_(topk_indices.roll(1, dims=1)) + graph.replay() + check(output) + + +def _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples): + """Reference only a few H64 rows so the large-Q regression stays bounded.""" + expected = [] + for row, heads in samples: + local = topk[row] + local = local[local >= 0].long() + physical = ragged[(starts[row].long() + local)].long() + values = _dequantize_kv_cache(packed[physical])[:, 0].float() + current = physical[:, None] == prefill_mem[None, :].long() + has_current = current.any(dim=1) + current_row = current.int().argmax(dim=1) + values[has_current] = prefill_kv[current_row[has_current], 0].float() + expected.append(((q[row, heads].float() @ values.T * 0.0625).softmax(-1) @ values)) + return expected + + +@pytest.mark.parametrize("num_q", [16, 33]) +def test_fp8_h64_largeq_native_prefill_permuted_ragged_graph(monkeypatch, num_q): + """The native sparse prefill path maps top-k ragged positions, never physical IDs.""" + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse prefill target is Hopper") + from sgl_kernel.flash_mla import flash_mla_sparse_fwd torch.manual_seed(53100 + num_q) - case = _PrefillCase([32768, 32768], [num_q // 2, num_q - num_q // 2], num_heads, capacity=131072) + capacity = 131072 + ragged_n = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32)[:ragged_n] + q = torch.randn(num_q, 64, 512, device="cuda", dtype=torch.bfloat16) + starts = torch.where( + torch.arange(num_q, device="cuda", dtype=torch.int32) < num_q // 2, + torch.full((num_q,), 1024, device="cuda", dtype=torch.int32), + torch.full((num_q,), 32768, device="cuda", dtype=torch.int32), + ) + current_local = (torch.arange(num_q, device="cuda", dtype=torch.int32) * 37) % 2048 + 64 + current_positions = starts + current_local + prefill_mem = ragged[current_positions.long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + prefill_kv.add_(torch.arange(num_q, device="cuda", dtype=torch.bfloat16)[:, None, None]) + topk = torch.randint(0, 2176, (num_q, 2176), device="cuda", dtype=torch.int32) + topk[:, ::19] = -1 # Holes must survive the native position-preserving map. + topk[:, 0] = current_local + topk[:, 1] = current_local # Duplicates are valid selected keys. + topk[:, 2] = topk[:, 3] + assert q.is_contiguous() and topk.is_contiguous() + observed = [] - def observed_flashattention(*args, **kwargs): - observed.append(kwargs["qv"].shape) - return flash_attn_with_kvcache(*args, **kwargs) - - def unexpected_flashmla(*args, **kwargs): - raise AssertionError("FP8 prefill must use its single FlashAttention backend") - - monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", observed_flashattention) - monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", unexpected_flashmla) - case.check(case.run()) - assert observed == [(num_q, num_heads, 512)] - graph, output = _capture_cuda_graph(case.run) - case.mutate() + def observed_sparse(*args, **kwargs): + observed.append((kwargs["q"].shape, kwargs["kv"].shape, kwargs["indices"].shape)) + return flash_mla_sparse_fwd(*args, **kwargs) + + monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", observed_sparse) + + def run(): + return fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + + samples = [(0, torch.tensor([0, 9, 63], device="cuda")), (num_q - 1, torch.tensor([3, 31, 62], device="cuda"))] + + def check(actual): + expected = _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples) + for (row, heads), value in zip(samples, expected): + torch.testing.assert_close(actual[row, heads].float(), value, atol=0.012, rtol=0.015) + + eager = run() + check(eager) + assert observed == [((num_q, 64, 512), (ragged_n, 1, 512), (num_q, 1, 2176))] + graph, output = _capture_cuda_graph(run) + q.normal_() + prefill_kv.normal_() + ragged.copy_(ragged.roll(1)) + prefill_mem.copy_(ragged[current_positions.long()]) + topk.copy_(topk.roll(1, dims=1)) + topk[:, 1] = current_local graph.replay() - case.check(output) - case.topk_indices.fill_(-1) + check(output) + topk.fill_(-1) graph.replay() assert torch.count_nonzero(output) == 0 + assert len(observed) == 3 # eager, graph warmup, and capture; replay stays inside the CUDA graph. + + +def test_fp8_h8_largeq_prefill_keeps_flashattention_path(monkeypatch): + """H8 large-Q prefill must not enter the H64 native sparse fast path.""" + import sgl_kernel.flash_attn as flash_attn + import sgl_kernel.flash_mla as flash_mla + + torch.manual_seed(53180) + capacity = 65536 + num_q = 16 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + q = torch.randn(num_q, 8, 512, device="cuda", dtype=torch.bfloat16) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) + starts = torch.full((num_q,), 1024, device="cuda", dtype=torch.int32) + current_local = torch.arange(num_q, device="cuda", dtype=torch.int32) + 64 + prefill_mem = ragged[(starts + current_local).long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :64] = torch.arange(64, device="cuda", dtype=torch.int32) + topk[:, 0] = current_local + seen_flashattention = [] + original_flashattention = flash_attn.flash_attn_with_kvcache + + def observed_flashattention(*args, **kwargs): + seen_flashattention.append(True) + return original_flashattention(*args, **kwargs) + + def unexpected_native(*args, **kwargs): + raise AssertionError("H8 large-Q prefill entered the native H64 sparse path") + + monkeypatch.setattr(flash_attn, "flash_attn_with_kvcache", observed_flashattention) + monkeypatch.setattr(flash_mla, "flash_mla_sparse_fwd", unexpected_native) + output = fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + assert seen_flashattention + assert output.shape == q.shape @pytest.mark.parametrize( "num_q,num_heads", [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], ) -def test_fp8_cached_prefill_strided_holes_and_current(num_q, num_heads): +def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num_heads): + """Exercise eager and graph replay for no-rowtable, small-Q packed prefill.""" torch.manual_seed(53016 + num_q) - seqs = [65536 + value for value in [127, 521, 1031, 2053, 4099, 8191, 12289, 16381][:num_q]] - case = _PrefillCase(seqs, [1] * num_q, num_heads) - case.q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - backing = torch.full((num_q, 4352), -1, device="cuda", dtype=torch.int32) - backing[:, ::2] = case.topk_indices - case.topk_indices = backing[:, ::2] - output = case.run() - case.check(output) - assert output.is_contiguous() and torch.count_nonzero(output) > 0 - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + # Requests deliberately have different spans; starts are not fixed-stride. + seq_lens = torch.tensor([127, 521, 1031, 2053, 4099, 8191, 12289, 16381], device="cuda", dtype=torch.int32)[:num_q] + starts = torch.nn.functional.pad(seq_lens.cumsum(0), (1, 0))[:-1] + current_rows = starts + seq_lens - 1 + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(num_q, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + topk_backing = torch.full((num_q, 2176 * 2), -1, device="cuda", dtype=torch.int32) + topk = topk_backing[:, ::2] # Exercise the topk stride passed to the direct kernel. + topk[:, :3] = torch.tensor([3, 11, 11], device="cuda", dtype=torch.int32) + topk[:, 3] = seq_lens - 1 + def run(): + return fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) -@pytest.mark.parametrize("num_q", [2, 3, 5, 8, 16, 32, 128]) -def test_fp8_cached_prefill_multi_query_and_graph(num_q): - torch.manual_seed(53200 + num_q) - case = _PrefillCase([131072, 131072], [num_q // 2, num_q - num_q // 2], 64) - case.check(case.run()) - for _ in range(2): - infer, state = case.new_state() - case.check(case.run(state)) - fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) - case.check(case.run(fresh)) - graph, output = _capture_cuda_graph(case.run) - case.mutate() + def check(actual): + assert actual.shape == q.shape and actual.is_contiguous() + expected_kv = _dequantize_kv_cache(packed)[ragged.long()] + expected_kv[current_rows.long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + expected = _torch_sparse_attention(q, expected_kv, physical) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + actual = run() + check(actual) + assert torch.count_nonzero(actual) > 0 # Includes BF16 current KV in every row. + graph, output = _capture_cuda_graph(run) + graph.replay() + check(output) + q.normal_() + source.normal_() + prefill_kv.normal_() + ragged.copy_(ragged.roll(1)) + prefill_mem.copy_(ragged[current_rows]) + write_fp8_e8m0_packed_kv(source, mem, packed) + topk.copy_(topk.roll(1, dims=1)) graph.replay() - case.check(output) - case.topk_indices.fill_(-1) + check(output) + topk.fill_(-1) # Empty graph replay must not read invalid packed pointers. graph.replay() assert torch.count_nonzero(output) == 0 -def test_fp8_prefill_graph_with_empty_requests_and_changing_boundaries(): - case = _PrefillCase([65536] * 4, [0, 4, 0, 3], 64) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - replacement = _PrefillCase([65536] * 4, [2, 0, 5, 0], 64) - for name in ( - "q", - "packed_kv", - "prefill_cache_kv", - "req_to_token_indexs", - "q_req_idx", - "position_ids", - "topk_indices", - "topk_mem_indices", - ): - getattr(case, name).copy_(getattr(replacement, name)) - case.infer_state.b1_cu_q_seq_len.copy_(replacement.infer_state.b1_cu_q_seq_len) - case.infer_state.position_ids.copy_(replacement.infer_state.position_ids) - for name in ("q_kv_seq_len", "q_req_idx", "q_kv_start_loc"): - getattr(case.att_state, name).copy_(getattr(replacement.att_state, name)) - graph.replay() - case.check(output) - - -def test_fp8_one_query_prefill_and_decode_read_their_own_kv_sources(): - case = _PrefillCase([65536], [1], 64) - case.q.zero_() - case.prefill_cache_kv.fill_(1) - case.source_kv.fill_(2) - write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) - case.topk_indices.fill_(-1) - case.topk_indices[:, 0] = 65535 - case.refresh() - decode_infer = _make_infer_state(**vars(case.infer_state)) - decode_infer.is_prefill = False - decode_infer.is_cuda_graph_capture = True - decode_state = case.backend.create_att_decode_state(decode_infer) - decode_state.init_state() - control = AttControl( - nsa_decode=True, - nsa_decode_dict={ - "topk_mem_indices": case.topk_mem_indices, - "softmax_scale": 0.0625, - "kv_lora_rank": 512, - }, +def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): + """Q1 H64 direct workspace is owned by the backend and reused across infer states.""" + torch.manual_seed(53150) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() + starts = torch.zeros(1, device="cuda", dtype=torch.int32) + current_rows = torch.tensor([2048], device="cuda", dtype=torch.int32) + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(1, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = current_rows + control = {"topk_indices": topk, "prefill_cache_kv": prefill_kv, "softmax_scale": 0.0625, "kv_lora_rank": 512} + infer_state = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(q.shape[0]) + state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged ) - def run(): - prefill = case.run() - decode = decode_state.decode_att((case.q, case.q[..., :0]), case.packed_kv[..., :584], None, control) - return prefill, decode - - prefill, decode = run() - torch.testing.assert_close(prefill, torch.ones_like(prefill), rtol=0, atol=0) - torch.testing.assert_close(decode, torch.full_like(decode, 2), rtol=0, atol=0) - graph, (prefill, decode) = _capture_cuda_graph(run) - case.prefill_cache_kv.fill_(3) - case.source_kv.fill_(4) - write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) + def expected(): + values = _dequantize_kv_cache(packed)[ragged.long()] + values[current_rows.long()] = prefill_kv + return _torch_sparse_attention(q, values, topk) + + def run(active_state=state, request=infer_state): + return run_sparse(active_state, request, q, packed[..., :584], control) + + torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) + cache = backend._test_workspace.get_object("glm_smallq", dict) + assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") + workspace = cache["direct"] + pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} + for _ in range(10): + layer_state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers + + fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) + fresh_state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(fresh_state, fresh_infer_state).float(), expected(), atol=0.012, rtol=0.015) + fresh_workspace = backend._test_workspace.get_object("glm_smallq", dict)["direct"] + assert fresh_workspace["partial_o"].data_ptr() == workspace["partial_o"].data_ptr() + + q.normal_() + source.normal_() + prefill_kv.normal_() + write_fp8_e8m0_packed_kv(source, mem, packed) + torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers + + +@pytest.mark.parametrize("num_q", [2, 3, 5, 8]) +def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): + """H64 selected prefill rewrites backend workspace across DSA layers.""" + torch.manual_seed(53200 + num_q) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 4096 + current_rows = starts + 2048 + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = 2048 + control = {"topk_indices": topk, "prefill_cache_kv": prefill_kv, "softmax_scale": 0.0625, "kv_lora_rank": 512} + infer_state = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(q.shape[0]) + state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + + def expected(): + values = _dequantize_kv_cache(packed)[ragged.long()] + values[current_rows.long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + return _torch_sparse_attention(q, values, physical) + + def run(active_state=state, request=infer_state): + return run_sparse(active_state, request, q, packed[..., :584], control) + + eager = run() + torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) + cache = backend._test_workspace.get_object("glm_smallq", dict) + assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") + workspace = cache["selected"] + pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} + # All 11 DSA layers share the backend allocation and rewrite its contents. + for _ in range(10): + layer_state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers + + fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) + fresh_state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(fresh_state, fresh_infer_state).float(), expected(), atol=0.012, rtol=0.015) + fresh_workspace = backend._test_workspace.get_object("glm_smallq", dict)["selected"] + assert fresh_workspace["compact"].data_ptr() == workspace["compact"].data_ptr() + + source.normal_() + write_fp8_e8m0_packed_kv(source, mem, packed) + prefill_kv.normal_() + q.normal_() + eager = run() + assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers + torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) + graph, output = _capture_cuda_graph(run) + q.normal_() + topk.copy_(topk.roll(1, dims=1)) + graph.replay() + torch.testing.assert_close(output.float(), expected(), atol=0.012, rtol=0.015) + topk.fill_(-1) graph.replay() - torch.testing.assert_close(prefill, torch.full_like(prefill, 3), rtol=0, atol=0) - torch.testing.assert_close(decode, torch.full_like(decode, 4), rtol=0, atol=0) + assert torch.count_nonzero(output) == 0 def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): @@ -542,18 +690,17 @@ def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) indices = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (1, 2049), device="cuda", dtype=torch.int32) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} infer = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer), + state = backend_context( + GlmSparseMLAFp8DecodeBackend, + "decode", decode_topk_bound=2051, flashmla_sched_meta=flash_mla.get_mla_metadata()[0], ) def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + return run_sparse(state, infer, (q, q[..., :0]), packed[..., :584], control) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) eager = run() @@ -610,17 +757,12 @@ def test_fp8_decode_infer_state_graph_qpad_and_mutation(num_heads): q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (batch, 2049), device="cuda", dtype=torch.int32) - control = AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} infer_state = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] - ) + state = backend_context(GlmSparseMLAFp8DecodeBackend, "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0]) def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + return run_sparse(state, infer_state, (q, q[..., :0]), packed[..., :584], control) def reference(): return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) @@ -642,14 +784,13 @@ def reference(): def test_fp8_sparse_direct_decode_factory_bound_guard(): - backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) - infer_state = InferStateInfo() - assert backend.create_att_decode_state(infer_state).decode_topk_bound == 2051 - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 8}) - assert backend.create_att_decode_state(infer_state).decode_topk_bound is None - backend.model = SimpleNamespace(config={"index_topk": 2047, "index_kpool": 4}) - assert backend.create_att_decode_state(infer_state).decode_topk_bound is None + backend = object.__new__(GlmSparseMLAFp8DecodeBackend) + backend.network_config = {"index_topk": 2048, "index_kpool": 4} + assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound == 2051 + backend.network_config = {"index_topk": 2048, "index_kpool": 8} + assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound is None + backend.network_config = {"index_topk": 2047, "index_kpool": 4} + assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound is None @pytest.mark.parametrize("batch", [1, 8]) @@ -675,12 +816,11 @@ def test_fp8_sparse_direct_decode_actual_entry_edges_and_graph(monkeypatch, batc indices[5, :64] = -1 # holes indices[6, 64:128] = indices[6, :64] # duplicates assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} state_owner = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(state_owner), + state = backend_context( + GlmSparseMLAFp8DecodeBackend, + "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0], decode_topk_bound=2051, ) @@ -689,7 +829,7 @@ def native_must_not_run(*args, **kwargs): raise AssertionError("supported GLM-5.3 decode took native FlashMLA fallback") def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + return run_sparse(state, state_owner, (q, q[..., :0]), packed[..., :584], control) assert not can_use_fp8_sparse_direct_decode(q, packed, indices.unsqueeze(1), 512, 2051) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) @@ -698,9 +838,9 @@ def run(): if batch == 1: assert torch.count_nonzero(eager) > 0 monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) - state.infer_state.is_cuda_graph_capture = True + state_owner.is_cuda_graph_capture = True graph, output = _capture_cuda_graph(run) - state.infer_state.is_cuda_graph_capture = False + state_owner.is_cuda_graph_capture = False graph.replay() torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) q.copy_(q.roll(1, dims=-1)) @@ -765,25 +905,23 @@ def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_h if batch >= 4: indices[3, 64:128] = indices[3, :64] # duplicates assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) - control = AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ) - backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) + control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} + backend = object.__new__(GlmSparseMLAFp8DecodeBackend) + backend.network_config = {"index_topk": 2048, "index_kpool": 4} infer_state = _make_infer_state(is_cuda_graph_capture=False) - state = backend.create_att_decode_state(infer_state) + state = backend_context(type(backend), "decode", settings=backend) state.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] assert state.decode_topk_bound == 2051 native_state_owner = _make_infer_state(is_cuda_graph_capture=False) - native_state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(native_state_owner), + native_state = backend_context( + GlmSparseMLAFp8DecodeBackend, + "decode", flashmla_sched_meta=flash_mla.get_mla_metadata()[0], decode_topk_bound=None, ) - def run(current_state): - return current_state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + def run(current_state, request=infer_state): + return run_sparse(current_state, request, (q, q[..., :0]), packed[..., :584], control) native_calls = [] original_native = flash_mla.flash_mla_with_kvcache @@ -794,7 +932,7 @@ def observe_native(*args, **kwargs): monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", observe_native) eager = run(state) - native_output = run(native_state) + native_output = run(native_state, native_state_owner) assert native_calls assert torch.equal(eager, native_output) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) @@ -804,9 +942,9 @@ def native_must_not_run(*args, **kwargs): raise AssertionError("supported graph shape did not use the direct packed-KV path") monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) - state.infer_state.is_cuda_graph_capture = True + infer_state.is_cuda_graph_capture = True graph, output = _capture_cuda_graph(lambda: run(state)) - state.infer_state.is_cuda_graph_capture = False + infer_state.is_cuda_graph_capture = False graph.replay() torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) q.copy_(q.roll(1, dims=-1)) @@ -838,9 +976,7 @@ def test_fp8_sparse_direct_decode_unknown_bound_fallback_and_strided_q(monkeypat q = torch.randn(1, 8, 512, device="cuda", dtype=torch.bfloat16) late = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) late[:, 2175] = 17 - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": late, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) + control = {"topk_mem_indices": late, "softmax_scale": 0.0625, "kv_lora_rank": 512} calls = [] original = flash_mla.flash_mla_with_kvcache @@ -850,10 +986,13 @@ def native(*args, **kwargs): monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native) unknown_owner = _make_infer_state(is_cuda_graph_capture=False) - unknown = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(unknown_owner), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] + unknown = backend_context( + GlmSparseMLAFp8DecodeBackend, + "decode", + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + decode_topk_bound=None, ) - actual = unknown.decode_att((q, q[..., :0]), packed[..., :584], None, control) + actual = run_sparse(unknown, unknown_owner, (q, q[..., :0]), packed[..., :584], control) torch.testing.assert_close( actual.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), late), atol=0.02, rtol=0.02 ) @@ -893,21 +1032,19 @@ def expected(): if (num_heads, batch) in ((1, 3), (7, 5)): infer = _make_infer_state(is_cuda_graph_capture=True) - state = Fp8Glm5NextSparseDecodeState(_infer_state_weak_ref=weakref.ref(infer), decode_topk_bound=2051) + state = backend_context(GlmSparseMLAFp8DecodeBackend, "decode", decode_topk_bound=2051) def native_must_not_run(*args, **kwargs): raise AssertionError("supported strided decode took native fallback") monkeypatch.setattr("flash_mla.flash_mla_with_kvcache", native_must_not_run) torch.testing.assert_close( - state.decode_att( + run_sparse( + state, + infer, (q, q[..., :0]), packed, - None, - AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ), + {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, ).float(), expected(), atol=0.02, @@ -987,15 +1124,14 @@ def direct_must_not_run(*args, **kwargs): "lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode.fp8_sparse_direct_decode", direct_must_not_run ) infer = _make_infer_state(is_cuda_graph_capture=True) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer), + state = backend_context( + GlmSparseMLAFp8DecodeBackend, + "decode", decode_topk_bound=2051, flashmla_sched_meta=flash_mla.get_mla_metadata()[0], ) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": topk, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) - run = lambda: state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + control = {"topk_mem_indices": topk, "softmax_scale": 0.0625, "kv_lora_rank": 512} + run = lambda: run_sparse(state, infer, (q, q[..., :0]), packed[..., :584], control) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) torch.testing.assert_close(run().float(), expected, atol=0.02, rtol=0.02) graph, output = _capture_cuda_graph(run) @@ -1003,30 +1139,142 @@ def direct_must_not_run(*args, **kwargs): torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) -@pytest.mark.parametrize("num_heads", [8, 64]) -def test_fp8_prefill_workspace_reuse_mutable_inputs_and_graph_replay(num_heads): +def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(): + """Generic H64 backend buffers are reused, batch metadata resets, and graph falls back.""" torch.manual_seed(53450) - backend = _make_workspace_backend(66) - case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], num_heads, backend=backend) - case.check(case.run()) - workspace = backend._prefill_workspaces[0] - assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} - pointers = {name: value.data_ptr() for name, value in workspace.items()} + capacity, num_q = 65536, 13 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + write_fp8_e8m0_packed_kv(source, mem, packed) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) + starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 2048 + current_local = torch.full((num_q,), 2048, device="cuda", dtype=torch.int32) + prefill_mem = ragged[(starts + current_local).long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16).add_(2) + q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = current_local + topk[:, 2049] = current_local + control = {"topk_indices": topk, "prefill_cache_kv": prefill_kv, "softmax_scale": 0.0625, "kv_lora_rank": 512} + infer = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(2 * num_q) + + def expected(active_ragged=ragged): + values = _dequantize_kv_cache(packed)[active_ragged.long()] + values[(starts + current_local).long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + return _torch_sparse_attention(q, values, physical) + + def run(active, request=infer): + return run_sparse(active, request, q, packed[..., :584], control) + + state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) + cache = backend._test_workspace._buffers + assert not hasattr(infer, "_fp8_sparse_generic_prefill_workspace_cache") + workspace = state.generic_workspace + pointers = {k: v.data_ptr() for k, v in workspace.items() if isinstance(v, torch.Tensor)} + bounded = dict(workspace, gathered_kv=workspace["gathered_kv"][:4096], row_table_built=False) + actual = fp8_prefill.context_attention_fwd_fp8( + q, + packed[..., :584], + ragged, + prefill_mem, + prefill_kv, + starts, + topk, + 0.0625, + generic_workspace=bounded, + ) + torch.testing.assert_close(actual.float(), expected(), atol=0.012, rtol=0.015) for _ in range(10): - infer, state = case.new_state() - case.check(case.run(state)) - case.mutate() - case.check(case.run()) - fresh = _PrefillCase([16384, 32768, 16384], [9, 11, 13], num_heads, capacity=131072, backend=backend) - fresh.check(fresh.run()) - assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0].items()} == pointers - for overrides in ({"prefill_cuda_graph_exe_list": []}, {"is_cuda_graph_capture": True}): - infer, state = case.new_state(**overrides) - case.check(case.run(state)) - graph, output = _capture_cuda_graph(case.run) - case.mutate() - graph.replay() - case.check(output) + layer = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer).float(), expected(), atol=0.012, rtol=0.015) + assert workspace["row_table_built"] + assert {k: v.data_ptr() for k, v in layer.generic_workspace.items() if isinstance(v, torch.Tensor)} == pointers + + # Change raw KV/top-k in place: reusable metadata must not cache either. + q.normal_() + prefill_kv.normal_() + topk.copy_(topk.roll(1, dims=1)) + topk[:, 2049] = current_local + torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) + + fresh_ragged = ragged.roll(1) + fresh_mem = fresh_ragged[(starts + current_local).long()].clone() + fresh = _make_infer_state(mem_index=fresh_mem) + fresh_state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=fresh_ragged + ) + torch.testing.assert_close(run(fresh_state, fresh).float(), expected(fresh_ragged), atol=0.012, rtol=0.015) + assert cache[("glm_prefill_rows", torch.int32)].data_ptr() == workspace["prefill_row_table"].data_ptr() + + graph_infer = _make_infer_state(mem_index=prefill_mem, prefill_cuda_graph_exe_list=[]) + graph_state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged, capture=True + ) + torch.testing.assert_close(run(graph_state, graph_infer).float(), expected(), atol=0.012, rtol=0.015) + assert not hasattr(graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") + + explicit_graph_infer = _make_infer_state(mem_index=prefill_mem, is_cuda_graph_capture=True) + explicit_graph_state = backend_context( + GlmSparseMLAFp8DecodeBackend, "prefill", settings=backend, ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close( + run(explicit_graph_state, explicit_graph_infer).float(), expected(), atol=0.012, rtol=0.015 + ) + assert not hasattr(explicit_graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") + + +def test_fp8_generic_prefill_workspace_graph_query_is_generic_h64_only(monkeypatch): + """H8/direct/selected prefill never queries CUDA graph-capture state.""" + import lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 as fp8_prefill + + queries = [] + + def query_spy(): + queries.append(True) + raise AssertionError("non-generic path queried CUDA graph capture state") + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", query_spy) + kv = torch.empty((1, 1, 584), dtype=torch.uint8) + params = { + "prefill_cache_kv": torch.empty((1, 1, 512), dtype=torch.bfloat16), + "topk_indices": torch.empty((1, 2176), dtype=torch.int32), + "softmax_scale": 0.0625, + "kv_lora_rank": 512, + } + control = params + + def fake_context(*args, **kwargs): + return args[0] + + monkeypatch.setattr(fp8_prefill, "context_attention_fwd_fp8", fake_context) + for heads, direct, selected in ((8, False, False), (64, True, False), (64, False, True)): + monkeypatch.setattr(fp8_prefill, "can_use_direct_smallq", lambda *args, d=direct: d) + monkeypatch.setattr(fp8_prefill, "can_use_selected_prefill", lambda *args, s=selected: s) + infer = _make_infer_state(mem_index=torch.empty(1, dtype=torch.int32)) + backend = _make_workspace_backend(1) + backend.get_smallq_workspace = lambda active_infer, active_q, kind="direct" if direct else "selected": { + "kind": kind + } + state = backend_context( + GlmSparseMLAFp8DecodeBackend, + "prefill", + settings=backend, + ks=torch.zeros(1, dtype=torch.int32), + ragged_mem_index=torch.zeros(1, dtype=torch.int32), + ) + assert ( + run_sparse(state, infer, torch.empty((1, heads, 512), dtype=torch.bfloat16), kv, control).shape[1] == heads + ) + assert not queries def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): @@ -1034,7 +1282,6 @@ def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) - monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.cpu_cache_operator.get_env_start_args", lambda: args) monkeypatch.setattr("lightllm.common.attention_state.pd_operator.get_env_start_args", lambda: args) monkeypatch.setattr( diff --git a/unit_tests/models/glm5_next/test_kda_conv_qkv.py b/unit_tests/models/glm5_next/test_kda_conv_qkv.py index b677477519..9495754f11 100644 --- a/unit_tests/models/glm5_next/test_kda_conv_qkv.py +++ b/unit_tests/models/glm5_next/test_kda_conv_qkv.py @@ -1,18 +1,17 @@ -import weakref +from unit_tests.common.attention.helpers import run_attention from types import SimpleNamespace import pytest import torch -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.linear import kda as backend_module -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import ( - causal_conv1d_qkv, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd, USE_DEFAULT_FLA_NORM - +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import KDAInputs, KDAResources +from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_fn +from lightllm.common.attention.backends._shared.linear.causal_conv1d_qkv import causal_conv1d_qkv +from lightllm.common.attention.backends._shared.linear.l2norm import USE_DEFAULT_FLA_NORM, l2norm_fwd +from lightllm.common.attention.backends.kda.triton import backend as backend_module +from lightllm.common.attention.backends.kda.triton.kernels.kda import chunk_kda_with_fused_gate +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0) or USE_DEFAULT_FLA_NORM, @@ -128,14 +127,24 @@ def test_backend_fused_conv_output_and_state(lengths): ) infer_state.b_ready_cache_len = torch.arange(len(lengths), device="cuda", dtype=torch.int32) % 2 * 7 infer_state.max_q_seq_len = max(lengths) - backend = SimpleNamespace( + backend = make_backend( + backend_module.TritonKDADecodeBackend, head_dim=128, tp_num_heads=heads, tp_hidden_size=heads * 128, conv_kernel_size=4, lower_bound=-5.0, ) - att = backend_module.KDAPrefillAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer_state)) + batch = make_batch( + lengths, + phase=AttentionPhase.PREFILL, + device="cuda", + token_num=tokens, + max_q_seq_len=max(lengths), + b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, + b_ready_cache_len=infer_state.b_ready_cache_len, + ) + att = make_context(backend, batch) att.b_conv_buffer_idx = torch.arange(1, len(lengths) + 1, device="cuda", dtype=torch.int32) att.b_ssm_buffer_idx = att.b_conv_buffer_idx layer_weight = SimpleNamespace( @@ -143,14 +152,16 @@ def test_backend_fused_conv_output_and_state(lengths): linear_A_log=SimpleNamespace(weight=torch.zeros(heads, device="cuda")), linear_dt_bias=SimpleNamespace(weight=torch.zeros(heads * 128, device="cuda")), ) - control = AttControl( - linear_att_prefill=True, - linear_att_prefill_dict=dict( - layer_weight=layer_weight, - layer_num=0, - mixed_qkv=projection[:, :channels], - raw_gate=torch.randn(tokens, heads * 128, **opts) * 0.1, - raw_beta=torch.randn(tokens, heads, **opts), + inputs = KDAInputs( + mixed_qkv=projection[:, :channels], + raw_gate=torch.randn(tokens, heads * 128, **opts) * 0.1, + raw_beta=torch.randn(tokens, heads, **opts), + ) + attention = make_layer_attention( + backend, + att, + KDAResources( + conv_states, ssm_states, weight, layer_weight.linear_A_log.weight, layer_weight.linear_dt_bias.weight ), ) expected_conv, expected_ssm = conv_states.clone(), ssm_states.clone() @@ -163,13 +174,12 @@ def test_backend_fused_conv_output_and_state(lengths): infer_state.b_ready_cache_len > 0, heads, ) - params = control.linear_att_prefill_dict expected, final_state = chunk_kda_with_fused_gate( q=q, k=k, v=v, - raw_g=params["raw_gate"].view(1, tokens, heads, 128), - beta=params["raw_beta"].float().sigmoid().view(1, tokens, heads), + raw_g=inputs.raw_gate.view(1, tokens, heads, 128), + beta=inputs.raw_beta.float().sigmoid().view(1, tokens, heads), A_log=layer_weight.linear_A_log.weight, g_bias=layer_weight.linear_dt_bias.weight, initial_state=ssm_states[att.b_ssm_buffer_idx].contiguous(), @@ -180,7 +190,7 @@ def test_backend_fused_conv_output_and_state(lengths): lower_bound=-5.0, ) expected_ssm[att.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype) - actual = att.prefill_att(None, None, None, att_control=control) + actual = run_attention(attention, att, inputs) assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) assert torch.equal(conv_states, expected_conv) assert torch.equal(ssm_states, expected_ssm) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 49b2a08eb6..74cfa06d37 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -1,29 +1,37 @@ -import weakref import dataclasses import math +import weakref from types import SimpleNamespace import pytest import torch import triton -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.envs_utils import set_env_start_args -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( +from lightllm.common.attention.backends.kda.triton.kernels.kda import ( chunk_kda_with_fused_gate, fused_kda_gate_chunk_cumsum, ) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 from lightllm.models.glm5_next.triton_kernel.kpool import ( compress_pools, - gather_pools, + expand_topk, gather_paged_pools, + gather_pools, get_pool_ranges, - expand_topk, ) -from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 - +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import set_env_start_args +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -412,7 +420,10 @@ def test_kpool_packed_kv_above_two_gib(): @pytest.mark.parametrize("tp_world_size", [1, 4]) def test_mhc_keeps_streams_through_decode_autotuning(monkeypatch, tp_world_size): - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + deep_gemm = pytest.importorskip("deep_gemm") + if not hasattr(deep_gemm, "tf32_hc_prenorm_gemm"): + pytest.skip("installed DeepGEMM does not provide mHC projection kernels") + from lightllm.common.triton_utils.autotuner import AutotuneKernelType, Autotuner from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.model import Glm5NextTpPartModel @@ -467,7 +478,7 @@ def all_reduce(input_embeddings, **kwargs): def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule tokens, heads, dim = 67, 2, 128 q, k, v = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(3)] @@ -499,8 +510,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): @pytest.mark.parametrize("heads", [16, 64, 128]) def test_nope_attention_native_512_and_cuda_graph(heads): - from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState + from lightllm.models.glm5_next.attention.sparse_mla import GlmSparseMLADecodeBackend if torch.cuda.get_device_capability()[0] != 9: pytest.skip("FA3 NoPE requires Hopper") @@ -522,34 +532,34 @@ def reference(): return torch.stack(expected) expected = reference() - control = AttControl( - nsa_prefill_dict={"topk_mem_indices": indexes, "softmax_scale": 0.0625, "kv_lora_rank": kv.shape[-1]} - ) - prefill_state = Glm5NextSparsePrefillState() - prefill = prefill_state._nsa_prefill_att(q, kv, control) + control = {"topk_mem_indices": indexes, "softmax_scale": 0.0625, "kv_lora_rank": kv.shape[-1]} + prefill_state = backend_context(GlmSparseMLADecodeBackend, "prefill") + prefill_batch = SimpleNamespace() + prefill = run_sparse(prefill_state, prefill_batch, q, kv, control) torch.testing.assert_close(prefill.float(), expected, atol=0.012, rtol=0.015) infer_state = InferStateInfo() infer_state.b1_cu_q_seq_len = torch.arange(10, dtype=torch.int32, device="cuda") infer_state.max_q_seq_len = 1 - decode = Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), + decode = backend_context( + GlmSparseMLADecodeBackend, + "decode", nsa_cache_seqlens=lengths, nsa_cu_seqlens_k_new=torch.nn.functional.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), ) - control.nsa_decode_dict = control.nsa_prefill_dict - out = decode._nsa_decode_att((q, q[..., :0]), kv, control) + control = control + out = run_sparse(decode, infer_state, (q, q[..., :0]), kv, control) torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - prefill_state._nsa_prefill_att(q, kv, control) - decode._nsa_decode_att((q, q[..., :0]), kv, control) + run_sparse(prefill_state, prefill_batch, q, kv, control) + run_sparse(decode, infer_state, (q, q[..., :0]), kv, control) torch.cuda.current_stream().wait_stream(stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - graph_prefill = prefill_state._nsa_prefill_att(q, kv, control) - graph_decode = decode._nsa_decode_att((q, q[..., :0]), kv, control) + graph_prefill = run_sparse(prefill_state, prefill_batch, q, kv, control) + graph_decode = run_sparse(decode, infer_state, (q, q[..., :0]), kv, control) graph.replay() torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 1aca878b2d..15fc776ae9 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -1,22 +1,24 @@ -import weakref +from lightllm.common.attention.api import AttentionPhase import dataclasses import json +import weakref from types import SimpleNamespace import pytest import torch import torch.nn.functional as F +from lightllm.common.attention.api.inputs import MLAResources from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState from lightllm.common.quantization import Quantcfg +from lightllm.models.glm5_next.attention.sparse_mla import GlmSparseMLADecodeBackend from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight from lightllm.models.glm5_next.model import Glm5NextTpPartModel from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args - +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -175,14 +177,16 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): state.att_state_manager = SimpleNamespace( get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv)) ) - state.prefill_att_state = Glm5NextSparsePrefillState() - state.b1_cu_q_seq_len = torch.arange(4, device="cuda", dtype=torch.int32) + backend = make_backend(GlmSparseMLADecodeBackend) + phase = AttentionPhase.PREFILL + context = make_context(backend, make_batch((2, 3, 5), phase=phase, device="cuda")) + state.b1_cu_q_seq_len = context.batch.b1_cu_q_seq_len state.max_q_seq_len = 1 - state.decode_att_state = Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(state), - nsa_cache_seqlens=lengths, - nsa_cu_seqlens_k_new=F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), + state.attention.contexts = {0: context} + attention = make_layer_attention( + backend, context, MLAResources(kv, layer.kv_lora_rank, 0), softmax_scale=layer.softmax_scale ) + layer.prefill_attention = layer.decode_attention = attention projected = torch.bmm(q.transpose(0, 1), kv_b.view(4, 256, 512)[:, :128]).transpose(0, 1) expected = [] for row, length in enumerate(lengths.tolist()): @@ -192,6 +196,13 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): expected = torch.stack(expected) prefill = layer._context_attention_kernel(q, None, state, weight) state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} + context = make_context(backend, make_batch((2, 3, 5), phase=AttentionPhase.DECODE, device="cuda")) + context.nsa_cache_seqlens = lengths + context.nsa_cu_seqlens_k_new = F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)) + state.attention.contexts = {0: context} + layer.decode_attention = make_layer_attention( + backend, context, MLAResources(kv, layer.kv_lora_rank, 0), softmax_scale=layer.softmax_scale + ) decode = layer._token_attention_kernel(q, state, weight) for output in (prefill, decode): torch.testing.assert_close(output.float(), expected, atol=1e-3, rtol=0.015) @@ -324,7 +335,7 @@ def test_mtp_pre_layer_applies_main_norm_before_hidden_fusion(config, monkeypatc main_norm_weight_=norms[0], enorm_weight_=norms[1], hnorm_weight_=norms[2], - eh_proj_weight_=SimpleNamespace(mm=lambda x: x @ projection), + eh_proj_weight_=SimpleNamespace(mm=lambda input_tensor: input_tensor @ projection), ) state = SimpleNamespace(mtp_draft_input_hiddens=hidden.clone()) monkeypatch.setattr(LlamaMultimodalPreLayerInfer, method, lambda self, ids, state, weight: embeddings.clone()) @@ -428,14 +439,15 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): for step, draft in enumerate(drafts): draft.prepare_attention_states() assert draft.att_state_manager is main.att_state_manager - draft._init_att_backend1() state = getattr(draft.att_state_manager, f"draft{step}_mla_packed_bf16_state") assert state is not main.att_state_manager.main_mla_packed_bf16_state assert not hasattr(draft, "full_att_state") assert state.storage is None assert state.config.logical_layer_groups == ((4 + step,),) assert state.state_name == f"draft{step}_mla_packed_bf16_state" - assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None + from lightllm.common.attention.api import AttentionFamily + + assert draft.layers_infer[0].describe_attention(AttentionPhase.DECODE).family is AttentionFamily.SPARSE_MLA assert getattr(draft.att_state_manager, f"draft{step}_compressor_indexer_w4_s4_bf16_state").buffer.is_cuda # NoPE state initialization must not require a model-owned RoPE cache. @@ -449,3 +461,45 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): torch.testing.assert_close( state.position_ids, torch.arange(2, 5, device="cuda", dtype=state.position_ids.dtype) ) + + +@pytest.mark.parametrize("family", ["gdn", "kda"]) +def test_prefill_graph_output_shape_uses_model_dimensions(monkeypatch, family): + from lightllm.models.glm5_next.layer_infer import transformer_layer_infer as glm + from lightllm.models.qwen3next.layer_infer import transformer_layer_infer as qwen + + callbacks = [] + graph_context = SimpleNamespace(__enter__=lambda: None, __exit__=lambda *args: None) + state = SimpleNamespace( + attention=object(), + prefill_cuda_graph_get_current_capture_graph=lambda: graph_context, + prefill_cuda_graph_create_graph_obj=lambda: None, + prefill_cuda_graph_add_cpu_runnning_func=lambda func, after_graph: callbacks.append(func), + ) + # No attention_groups/backend dimensions: only the binding's public execution methods. + binding = SimpleNamespace(context=lambda state: None) + layer = SimpleNamespace(prefill_attention=binding, tp_world_size_=2, tp_linear_num_heads=3, linear_head_dim=16) + inputs = torch.zeros(5, 48, dtype=torch.bfloat16, device="cuda") + if family == "gdn": + monkeypatch.setattr(qwen, "tensor_to_no_ref_tensor", lambda tensor: tensor) + binding.forward = lambda *a: ( + torch.ones(1, 5, 3, 16, device="cuda"), + torch.full((5, 3, 16), 2.0, device="cuda"), + ) + weights = SimpleNamespace(linear_num_v_heads=6, linear_v_head_dim=16) + output, gate = qwen.Qwen3NextTransformerLayerInfer._linear_prefill_cuda_graph_wrapper( + layer, inputs, state, weights + ) + assert output.shape == gate.shape == (5, 3, 16) + else: + monkeypatch.setattr(glm, "tensor_to_no_ref_tensor", lambda tensor: tensor) + binding.forward = lambda *a: torch.ones(1, 5, 3, 16, device="cuda") + output = glm.Glm5NextTransformerLayerInfer._kda_prefill_cuda_graph_wrapper( + layer, inputs, inputs, inputs, state, None + ) + assert output.shape == (1, 5, 3, 16) + assert len(callbacks) == 1 + callbacks[0](state) + assert torch.all(output == 1) + if family == "gdn": + assert torch.all(gate == 2) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 30c4fc22c4..5843fcae4a 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -1,17 +1,17 @@ +from unit_tests.common.attention.helpers import build, run_attention import weakref + import pytest import torch +from lightllm.common.attention.backends._shared.linear.replayssm_fold import fold_prefix +from lightllm.common.attention.backends.kda.triton.kernels.kda_decode import fused_recurrent_kda from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import ( - fused_recurrent_kda, -) -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools - pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -145,17 +145,9 @@ def prefill(values, destination, tail): def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): from types import SimpleNamespace - from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState - from lightllm.common.basemodel.attention.linear.replay_kda import ( - ReplayKDADecodeAttState, - ) - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( - ReplayKDAState, - ) - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_update, - ) + from lightllm.common.attention.backends._shared.linear.causal_conv1d import causal_conv1d_update + from lightllm.common.attention.backends.kda.replay.backend import TritonKDAReplayDecodeBackend + from lightllm.common.attention.backends.kda.triton.backend import TritonKDADecodeBackend torch.manual_seed(53) heads, dim, width = 2, 128, mtp_step + 1 @@ -208,15 +200,6 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod expected.append(output.reshape(1, heads, dim)) expected_ssm[req * width + token] = state[0] row += 1 - backend = SimpleNamespace( - mtp_step=mtp_step, - tp_num_heads=heads, - head_dim=dim, - tp_hidden_size=hidden, - lower_bound=-5.0, - uses_dynamic_spec_verify_layout=lambda: dynamic, - split_qkv=lambda x: x.split(hidden, -1), - ) resources = AttentionStateManager(request_capacity=4) config = ConvSSMStateConfig( logical_layer_ids=(0,), @@ -252,54 +235,48 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod infer.req_manager = manager infer.att_state_manager = resources infer.att_state_role = role - weight = SimpleNamespace( - get_merged_kda_conv_weight=lambda: conv_weight, - linear_A_log=SimpleNamespace(weight=a), - linear_dt_bias=SimpleNamespace(weight=bias), + from lightllm.common.attention.api.inputs import KDAInputs, KDAResources + from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention + + executor = make_backend( + TritonKDAReplayDecodeBackend if mode == "replay" else TritonKDADecodeBackend, + mtp_step=mtp_step, + draft_step=mtp_step, + dynamic_verify=dynamic, + tp_num_heads=heads, + head_dim=dim, + tp_hidden_size=hidden, + lower_bound=-5.0, ) - decode_cls = ReplayKDADecodeAttState if mode == "replay" else KDADecodeAttState - state = decode_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - state.init_state() - actual = state.decode_att( - None, - None, - None, - AttControl( - linear_att_decode=True, - linear_att_decode_dict=dict( - layer_weight=weight, - layer_num=0, - mixed_qkv=mixed, - raw_gate=gate, - raw_beta=beta, - ), - ), + batch = make_batch( + (1,) * tokens, device="cuda", b_req_idx=infer.b_req_idx, b_mtp_index=infer.b_mtp_index, hold_req_id=3 ) - torch.testing.assert_close( - actual.reshape(tokens, heads, dim)[:row], - torch.cat(expected), - atol=2e-3, - rtol=1e-2, + context = make_context(executor, batch) + build(executor, context) + attention = make_layer_attention( + executor, + context, + KDAResources( + *selected.get_layer_view(0), + conv_weight, + a, + bias, + mtp_state_index=selected.mtp_state_index, + replay_state=selected if mode == "replay" else None, + ), ) + actual = run_attention(attention, context, KDAInputs(mixed, gate, beta)) + torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) if mode == "native": checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] torch.testing.assert_close( - selected.ssm_state[0, :checked_slots], - expected_ssm[:checked_slots], - atol=2e-3, - rtol=1e-2, + selected.ssm_state[0, :checked_slots], expected_ssm[:checked_slots], atol=2e-3, rtol=1e-2 ) else: for req, length in zip(requests, lengths): - saved = selected.ssm_state - clone = saved.clone() - selected.ssm_state = clone - try: - fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) - finally: - selected.ssm_state = saved + fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) torch.testing.assert_close( - clone[0, req], + selected.ssm_state[0, req], expected_ssm[req * width + length - 1], atol=2e-3, rtol=1e-2, diff --git a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py index 5b60fece98..f3545a68a4 100644 --- a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py +++ b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py @@ -1,3 +1,4 @@ +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import Mock @@ -16,6 +17,7 @@ PackedMLALogicalConfig, ) from lightllm.common.basemodel.batch_objs import ModelInput, PostLayerOutput +from lightllm.common.attention.session import AttentionSession from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward @@ -89,7 +91,7 @@ def test_forward_uses_common_prefill(monkeypatch, tokens, mode): model.att_state_manager = object() model.hidden_collector_prototype = FinalHiddenCollector() model.return_all_prompt_logics = False - model.prefill_att_backend = model.prefill_att_backend1 = None + model.attention = SimpleNamespace(session=lambda state: nullcontext()) model._select_mem_indexes = lambda model_input: torch.zeros_like(model_input.input_ids) batch = int(tokens > 0) model_input = ModelInput( @@ -130,6 +132,8 @@ def test_forward_uses_common_prefill(monkeypatch, tokens, mode): layer._project_sparse_qkv_a = lambda hidden, *_: (hidden, None, hidden) layer._normalize_sparse_kv = lambda kv, _: kv layer.indexer = SimpleNamespace(write_k_pools=Mock()) + context = object() + layer.prefill_attention = SimpleNamespace(context=Mock(return_value=context)) writer = Mock() layer._post_cache_kv = writer model.layers_infer = [layer] @@ -165,6 +169,8 @@ def test_forward_uses_common_prefill(monkeypatch, tokens, mode): post_forward.assert_not_called() allocator.cache_env_in.assert_called_once() allocator.cache_env_out.assert_called_once() + layer.prefill_attention.context.assert_called_once_with(infer.attention) + assert layer.indexer.write_k_pools.call_args.args[2] is context @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -237,6 +243,8 @@ def norm(input, eps, **_): initial = [s.clone() for s in stores] q_lengths = [q for _, _, q in requests] lengths = torch.cat([torch.arange(past + 1, past + q + 1, device="cuda") for _, past, q in requests]).int() + context = SimpleNamespace(q_kv_seq_len=lengths, q_req_idx=None) + layer.prefill_attention = SimpleNamespace(context=lambda session: context) infer = SimpleNamespace( is_prefill=True, is_draft_model=True, @@ -250,7 +258,7 @@ def norm(input, eps, **_): b1_cu_q_seq_len=torch.tensor([0, *q_lengths], device="cuda", dtype=torch.int32).cumsum(0).int(), max_q_seq_len=max(q_lengths), max_kv_seq_len=2048, - prefill_att_state=SimpleNamespace(q_kv_seq_len=lengths, q_req_idx=None), + attention=AttentionSession(), req_manager=SimpleNamespace(req_to_token_indexs=None), ) # Top-k outputs are discarded here; run the normal projections and cache writers. @@ -259,7 +267,7 @@ def norm(input, eps, **_): _, kv = layer._get_qkv(layer._att_norm(hidden, infer, weights), infer, weights) layer._post_cache_kv(kv, infer, weights) layer.indexer._get_indices( - **infer.get_topk_indices_params, infer_state=infer, att_state=infer.prefill_att_state, layer_weight=weights + **infer.get_topk_indices_params, infer_state=infer, att_state=context, layer_weight=weights ) expected = [s.clone() for s in stores] assert not torch.equal(initial[2], expected[2]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 4d15d6bf1c..82b72b81bd 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -1,6 +1,3 @@ -from unit_tests.common.att_resources.model_helpers import window_states, build_hybrid_resources -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - import dataclasses from types import SimpleNamespace @@ -8,19 +5,22 @@ import torch from lightllm.common.attention_state.allocators import FullAttKVAllocator -from unit_tests.common.att_resources.model_helpers import create_full_attention_roles - from lightllm.common.attention_state.full_att import ( - PackedMLAFullAttentionState, PackedMLAAttentionConfig, + PackedMLAFullAttentionState, PackedMLALogicalConfig, ) -from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name - +from unit_tests.common.att_resources.model_helpers import ( + build_hybrid_resources, + create_full_attention_roles, + full_state_test_fields, + window_states, +) pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), @@ -304,6 +304,9 @@ def test_pd_kv_and_runtime_state_roundtrip( def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): + from lightllm.common.attention_state.pd_operator import logger + + monkeypatch.setattr(logger, "handlers", [*logger.handlers, caplog.handler]) monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] full_bytes = max( diff --git a/unit_tests/models/glm5_next/test_prefill_workspace.py b/unit_tests/models/glm5_next/test_prefill_workspace.py index 0d3a85d965..a1d7bc5e20 100644 --- a/unit_tests/models/glm5_next/test_prefill_workspace.py +++ b/unit_tests/models/glm5_next/test_prefill_workspace.py @@ -1,116 +1,73 @@ -from types import SimpleNamespace - +from lightllm.common.attention.session import AttentionSession import pytest import torch -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next -from lightllm.models.glm5_next.model import Glm5NextTpPartModel -from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel -from unit_tests.models.glm5_next.test_fp8_kv import _PrefillCase, _make_workspace_backend, _capture_cuda_graph - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture(autouse=True) -def current_cuda_device(monkeypatch): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - - -@pytest.mark.parametrize("overlap", [False, True]) -@pytest.mark.parametrize("run_mode", ["normal", "decode"]) -def test_prefill_workspace_is_reserved_once_and_reused_at_runtime(monkeypatch, overlap, run_mode): - model = SimpleNamespace( - args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode=run_mode), - batch_max_tokens=8, - max_req_num=2, - max_seq_length=13, - config={"index_topk": 2048}, - att_state_role=f"workspace_test_{run_mode}_{overlap}", +from lightllm.common.attention.api import AttentionPhase, ContextCapacity +from lightllm.common.attention.api.workspace import AttentionWorkspace +from lightllm.models.glm5_next.attention.fp8_sparse_mla import ( + GlmSparseMLAFp8PrefillBackend, + Fp8Glm5NextSparsePrefillContext, +) + + +def test_fp8_prefill_context_is_owned_by_the_backend(): + backend = object.__new__(GlmSparseMLAFp8PrefillBackend) + backend.network_config = {"index_topk": 2048, "index_kpool": 4} + workspace = AttentionWorkspace("cpu") + context = AttentionSession.allocate( + backend, + ContextCapacity(phase=AttentionPhase.PREFILL, batch_size=2, token_num=4, max_kv_seq_len=32, max_q_seq_len=4), + capture=False, + workspace=workspace, + ) + assert isinstance(context, Fp8Glm5NextSparsePrefillContext) + assert context.backend_identity is backend + assert context.workspace is workspace + + +def test_fp8_backend_has_no_legacy_state_factory(): + assert not hasattr(GlmSparseMLAFp8PrefillBackend, "create_att_prefill_state") + assert not hasattr(GlmSparseMLAFp8PrefillBackend, "create_att_decode_state") + + +@pytest.mark.parametrize("decode_only", [False, True]) +def test_fp8_workspace_reserved_per_microbatch_before_execution(decode_only): + from types import SimpleNamespace + + from lightllm.common.attention.api import ServingConstraints + from lightllm.common.basemodel.attention import ModelAttention + from lightllm.platform import get_device_runtime + + backend = object.__new__(GlmSparseMLAFp8PrefillBackend) + backend.network_config = {"index_topk": 2048, "index_kpool": 4} + backend.prefill_token_capacity = 32 + backend.prefill_graph_enabled = True + backend.decode_only = decode_only + attention = ModelAttention.__new__(ModelAttention) + attention.constraints = ServingConstraints(microbatch_count=2, max_batch_size=2, max_seq_len=64) + attention.groups = ({AttentionPhase.PREFILL: backend},) + attention.workspaces = {} + attention.runtime = get_device_runtime("cpu") + attention.reserve(None) + first, second = (attention.workspaces[(0, slot)] for slot in (0, 1)) + gathered = first.scratch("glm_gathered_kv", 1, torch.bfloat16) + assert gathered.numel() == (32 if decode_only else 2 * 64 + 32) * 512 + assert gathered.data_ptr() != second.scratch("glm_gathered_kv", 1, torch.bfloat16).data_ptr() + assert first.scratch("glm_chunk_indices", 1, torch.int32).numel() == 32 * 2176 + context = AttentionSession.allocate( + backend, + ContextCapacity(phase=AttentionPhase.PREFILL, batch_size=2, token_num=4, max_kv_seq_len=64, max_q_seq_len=4), + capture=True, + workspace=first, + ) + context.ragged_mem_index = torch.arange(100, dtype=torch.int32) + q = torch.zeros(4, 64, 512, dtype=torch.bfloat16) + packed = torch.empty(160, 1, 584, dtype=torch.uint8) + topk = torch.zeros(4, 2176, dtype=torch.int32) + buffers = backend.get_generic_prefill_workspace(context, q, packed, topk) + assert buffers["gathered_kv"].data_ptr() == gathered.data_ptr() + assert buffers["cu_seqlens_q"].tolist() == [0, 1, 2, 3, 4] + first.freeze() + assert ( + backend.get_generic_prefill_workspace(context, q, packed, topk)["gathered_kv"].data_ptr() == gathered.data_ptr() ) - backend_class = fp8_glm5_next.Fp8Glm5NextSparseAttBackend - reserved = backend_class.allocate_prefill_workspaces(model) - pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in reserved] - backend = object.__new__(backend_class) - backend.__init__(model) - assert [ - {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces - ] == pointers - assert len(reserved) == 1 + int(overlap) - capacity = 8 if run_mode == "decode" else 34 - assert reserved[0]["gathered_kv"].shape == (capacity, 1, 512) - assert reserved[0]["chunk_topk_indices"].shape == (8, 2176) - - def forbidden(*args, **kwargs): - raise AssertionError("Runtime must reuse the workspace reserved before KV capacity planning") - - monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) - monkeypatch.setattr(backend, "get_gpu_workspace_buffer", forbidden) - for slot in range(1 + int(overlap)): - for q_tokens in (1, 8, 3): - topk = torch.empty((q_tokens, 2176), dtype=torch.int32, device="cuda") - views = backend.get_prefill_workspace(SimpleNamespace(microbatch_index=slot), topk) - assert {key: value.data_ptr() for key, value in views.items()} == pointers[slot] - torch.testing.assert_close( - views["cu_seqlens_q"], torch.arange(q_tokens + 1, dtype=torch.int32, device="cuda") - ) - - -def test_prefill_microbatches_use_independent_scratch_during_graph_replay(monkeypatch): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 7) - backend = _make_workspace_backend(8, overlap=True) - first = _PrefillCase([15], [3], 8, backend=backend) - second = _PrefillCase([17], [3], 8, backend=backend) - second.infer_state.microbatch_index = 1 - first.check(first.run()) - second.check(second.run()) - pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces] - assert all(pointers[0][key] != pointers[1][key] for key in pointers[0]) - - def forbidden(*args, **kwargs): - raise AssertionError("Attention must not allocate a new prefill workspace") - - monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) - first_graph, first_output = _capture_cuda_graph(first.run) - second_graph, second_output = _capture_cuda_graph(second.run) - first_stream, second_stream = torch.cuda.Stream(), torch.cuda.Stream() - first.mutate() - second.mutate() - for _ in range(3): - first_stream.wait_stream(torch.cuda.current_stream()) - second_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(first_stream): - first_graph.replay() - with torch.cuda.stream(second_stream): - second_graph.replay() - torch.cuda.current_stream().wait_stream(first_stream) - torch.cuda.current_stream().wait_stream(second_stream) - first.check(first_output) - second.check(second_output) - assert [ - {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces - ] == pointers - - -@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -def test_glm_reserves_prefill_workspace_during_resource_preparation(monkeypatch, model_class, kv_type): - model = object.__new__(model_class) - model.args = SimpleNamespace(llm_kv_type=kv_type) - model.mtp_previous_draft_models = [] - model.att_state_manager = SimpleNamespace(get_full_attention_states=lambda role: [object()]) - calls = [] - - def prepare(self): - assert self is model - calls.append("prepare_states") - - def reserve(cls, prepared_model): - assert prepared_model is model - calls.append("reserve_workspace") - - monkeypatch.setattr(TpPartBaseModel, "prepare_attention_states", prepare) - monkeypatch.setattr(fp8_glm5_next.Fp8Glm5NextSparseAttBackend, "allocate_prefill_workspaces", classmethod(reserve)) - model.prepare_attention_states() - assert calls == (["prepare_states", "reserve_workspace"] if kv_type == "fp8kv_dsa" else ["prepare_states"]) diff --git a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py index fa9e727f87..7cee8bf66f 100644 --- a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py +++ b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py @@ -20,10 +20,13 @@ def make_projection(layer, heads: int, inner: int, cols: int, dtype: torch.dtype layer.layer_num_ = 0 layer.softmax_scale = inner ** -0.5 layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (None, None)) + layer.prefill_attention = SimpleNamespace( + context=lambda session: None, forward=lambda session, inputs: inputs.query + ) state = SimpleNamespace( att_state_role="main", att_state_manager=Mock(), - prefill_att_state=SimpleNamespace(prefill_att=lambda **kwargs: kwargs["q"]), + attention=object(), ) weight = SimpleNamespace(k_b_proj_=projection) diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py index 6a25a8d4ea..d865d829e0 100644 --- a/unit_tests/models/glm5_next/test_visual.py +++ b/unit_tests/models/glm5_next/test_visual.py @@ -55,8 +55,8 @@ def test_image_patch_order_temporal_repeat_and_black_padding(): def test_vision_attention_uses_triton_rope_with_fa3_backend(monkeypatch): import lightllm.models.glm5_next.glm5_next_visual as glm_visual import lightllm.server.visualserver as visualserver - from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend - from lightllm.server.visualserver import set_vit_att_backend + from lightllm.common.attention.backends.dense.fa3.varlen import Fa3VarlenAttentionOp + from lightllm.server.visualserver import set_vit_attention_op rope_inputs = [] norm_inputs = [] @@ -70,16 +70,18 @@ def fake_rope(x, cos, sin): rope_inputs.append((x, cos, sin)) return x - def fake_fa3(q, k, v, o, cu_seqlens, max_seqlen): + def fake_fa3(self, inputs, *, out): + q, k, v = inputs.query, inputs.key, inputs.value + o, cu_seqlens, max_seqlen = out, inputs.cu_seqlens, inputs.max_seqlen attn_inputs.append((q, k, v, cu_seqlens, max_seqlen)) o.copy_(q) return o monkeypatch.setattr(glm_visual, "qk_rms_norm", fake_qk_norm) monkeypatch.setattr(glm_visual, "apply_rotary_pos_emb_triton", fake_rope) - monkeypatch.setattr(Fa3VitAttBackend, "_vit_att_fwd", staticmethod(fake_fa3)) - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", visualserver.VIT_ATTN_BACKEND) - set_vit_att_backend("fa3") + monkeypatch.setattr(Fa3VarlenAttentionOp, "forward", fake_fa3) + monkeypatch.setattr(visualserver, "VIT_ATTENTION_OP", visualserver.VIT_ATTENTION_OP) + set_vit_attention_op("fa3") attention = glm_visual.Glm5NextVisionAttention(hidden_size=16, num_heads=2, eps=1e-5, bias=True) x = torch.randn(3, 16) @@ -149,10 +151,10 @@ def glm_rms_norm(x, weight, eps): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_vision_batch_preserves_independent_image_attention(): from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer - from lightllm.server.visualserver import set_vit_att_backend + from lightllm.server.visualserver import set_vit_attention_op torch.manual_seed(53) - set_vit_att_backend("sdpa") + set_vit_attention_op("sdpa") model = ( Glm5NextVisionTransformer( {"data_type": "float32"}, diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py index c25576a19a..b81032e4c1 100644 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -1,15 +1,19 @@ +from unit_tests.common.attention.helpers import make_env +from unit_tests.common.attention.helpers import build, run_attention +from dataclasses import replace from types import SimpleNamespace import pytest import torch +from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api.inputs import GDNInputs, GDNResources +from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNDecodeBackend, TritonGDNPrefillBackend from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.attention.helpers import make_batch, make_context, make_layer_attention, make_spec, model_stub @pytest.mark.parametrize("device", ["cpu", "cuda"]) @@ -40,18 +44,10 @@ def test_gdn_decode_keeps_logical_layer_state( indices = get_loop_layer_indices(settings, 32) kinds = get_loop_layer_types(settings, config) layers = tuple(i for i, kind in enumerate(kinds) if kind == "linear_attention") - args = StartArgs(mtp_step=mtp_step, linear_att_ssm_data_type="float32") - monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.basemodel.attention.base_att.get_env_start_args", lambda: args) - backend = object.__new__(TritonLinearAttBackend) - backend.__init__( - SimpleNamespace( - config=config, - tp_world_size_=1, - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: mtp_step), - ) - ) + model = model_stub(mtp_step=mtp_step, linear_att_ssm_data_type="float32") + model.config = config + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda _: mtp_step) + backend = TritonGDNDecodeBackend(make_env(model), make_spec(use_mtp=mtp_step > 0)) dtype = torch.bfloat16 if device == "cuda" else torch.float32 layout = ConvSSMStateConfig(layers, 1, 2, 32, 32, 3, dtype, torch.float32) resource = GDNState(layout, request_capacity=3, mtp_step=mtp_step, device=device) @@ -91,57 +87,61 @@ def recurrent(*, initial_state, v, **kwargs): initial_state[mtp_step + 1 : 2 * (mtp_step + 1)].add_(1) return v, None - monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) monkeypatch.setattr( - "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + "lightllm.common.attention.backends.gdn.triton.kernels.gdn_decode_pack.conv_pack_gdn_decode_inputs", + conv_pack, ) monkeypatch.setattr( - "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent + "lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + ) + monkeypatch.setattr( + "lightllm.common.attention.backends.gdn.triton.kernels.fused_recurrent.fused_recurrent_gated_delta_rule", + recurrent, ) monkeypatch.setattr( - "lightllm.common.basemodel.attention.linear.gdn.mtp_fused_recurrent_gated_delta_rule", + "lightllm.common.attention.backends.gdn.triton.kernels.mtp_fused_recurrent." + "mtp_fused_recurrent_gated_delta_rule", recurrent, ) - infer = InferStateInfo() - infer.att_state_manager, infer.att_state_role = resources, "main" - infer.b_req_idx = torch.tensor([1], device=device, dtype=torch.int32) - control = dict(layer_num=logical_layer, layer_weight=weight) + params = GDNResources( + conv_view, + ssm_view, + weight.linear_conv1d.mm_param.weight, + None, + weight.linear_A_log.weight, + weight.linear_dt_bias.weight, + mtp_state_index=resource.mtp_state_index, + ) mixed_dim = backend.tp_qkvz_dim + backend.tp_ba_dim if device == "cuda": - infer.b1_cu_q_seq_len = torch.tensor([0, 5], device=device, dtype=torch.int32) - infer.b_ready_cache_len = torch.zeros(1, device=device, dtype=torch.int32) - prefill = backend.create_att_prefill_state(infer) - prefill.init_state() - prefill.prefill_att( - None, - None, - None, - AttControl( - linear_att_prefill=True, - linear_att_prefill_dict=dict( - **control, mixed_qkvzba=torch.randn(5, mixed_dim, device=device, dtype=dtype) - ), - ), + batch = make_batch( + (5,), + device=device, + phase=AttentionPhase.PREFILL, + b_req_idx=torch.tensor([1], device=device, dtype=torch.int32), + token_num=5, + max_q_seq_len=5, + b_q_seq_len=torch.tensor([5], device=device), + b1_cu_q_seq_len=torch.tensor([0, 5], device=device, dtype=torch.int32), ) - - infer.batch_size = mtp_step + 1 - infer.b_req_idx = infer.b_req_idx.repeat(infer.batch_size) - decode = backend.create_att_decode_state(infer) + prefill = TritonGDNPrefillBackend(make_env(model), make_spec(phase=AttentionPhase.PREFILL)) + context = make_context(prefill, batch) + build(prefill, context) + attention = replace(make_layer_attention(prefill, context, params), logical_layer_id=logical_layer) + run_attention(attention, context, GDNInputs(torch.randn(5, mixed_dim, device=device, dtype=dtype))) + batch_size = mtp_step + 1 + batch = make_batch( + (1,) * batch_size, device=device, b_req_idx=torch.ones(batch_size, device=device, dtype=torch.int32) + ) + context = make_context(backend, batch) + attention = replace(make_layer_attention(backend, context, params), logical_layer_id=logical_layer) untouched = [i for i in range(len(layers)) if i != storage_layer] for _ in range(2): before_conv, before_ssm = resource.conv_state.clone(), resource.ssm_state.clone() - decode.init_state() - output, _ = decode.decode_att( - None, - None, - None, - AttControl( - linear_att_decode=True, - linear_att_decode_dict=dict( - **control, mixed_qkvzba=torch.randn(infer.batch_size, mixed_dim, device=device, dtype=dtype) - ), - ), + build(backend, context) + output, _ = attention.backend.forward( + attention, context, GDNInputs(torch.randn(batch_size, mixed_dim, device=device, dtype=dtype)) ) assert torch.isfinite(output).all() assert not torch.equal(resource.conv_state[storage_layer], before_conv[storage_layer]) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py index f736ec4cd6..9a161bad45 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -1,139 +1,271 @@ -import torch -import math -import time import pytest -from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd -from lightllm.models.vit.triton_kernel import flashattention_nopad -from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend +import torch + +from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionInputs +from lightllm.common.attention.registry import get_varlen_attention_op_class +from lightllm.common.attention.varlen import AutoVarlenAttentionRunner +from lightllm.platform import get_device_runtime + + +def make_inputs(lengths=(3, 0, 5), *, device="cpu", dtype=torch.float32, head_dim=80, strided=False): + tokens = sum(lengths) + if strided: + qkv = torch.randn(tokens, 3, 2, head_dim, dtype=dtype, device=device) + q, k, v = qkv.unbind(1) + else: + q, k, v = [torch.randn(tokens, 2, head_dim, dtype=dtype, device=device) for _ in range(3)] + cu = torch.tensor([0, *lengths], dtype=torch.int32, device=device).cumsum(0, dtype=torch.int32) + return VarlenAttentionInputs(q, k, v, cu, max(lengths)) -_fa3_mark = pytest.mark.skipif( - not flashattention_nopad._flash_attn_v3_available - or not torch.cuda.is_available() - or torch.cuda.get_device_capability()[0] < 9, - reason="FA3 requires sgl_kernel and a Hopper GPU", -) +def reference_attention(inputs): + out = torch.empty_like(inputs.query) + offsets = inputs.cu_seqlens.cpu().tolist() + for start, end in zip(offsets, offsets[1:]): + if start == end: + continue + q, k, v = [x[start:end].float().transpose(0, 1) for x in (inputs.query, inputs.key, inputs.value)] + scores = q @ k.transpose(-1, -2) * inputs.query.shape[-1] ** -0.5 + out[start:end] = (scores.softmax(-1) @ v).transpose(0, 1).to(out.dtype) + return out + +@pytest.mark.parametrize("error_type", [None, UnsupportedAttention]) +def test_audio_fallback_only_rejects_current_input(monkeypatch, error_type): + from lightllm.common.attention import varlen -def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): calls = [] - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("fa3") - o[0].fill_(99) # A backend may fail after partially writing output. - raise RuntimeError("shape-dependent FA3 failure") - - def fallback(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("triton") - o.fill_(7) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) - monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - assert not flashattention_nopad._flash_attn_v3_available - - # Later calls should use the disabled-backend path without retrying FA3. + + class FakeOp: + @classmethod + def check_environment(cls, caps): + return None + + @classmethod + def check_support(cls, inputs, caps=None): + calls.append("fa3") + if error_type: + raise error_type("shape-dependent FA3 failure") + + def forward(self, inputs, *, out): + out.fill_(9) + return out + + class Fallback(FakeOp): + @classmethod + def check_support(cls, inputs, caps=None): + return None + + def forward(self, inputs, *, out): + calls.append("triton") + out.fill_(7) + return out + + monkeypatch.setattr(varlen, "get_varlen_attention_op_class", lambda name: {"fa3": FakeOp, "triton": Fallback}[name]) + op = AutoVarlenAttentionRunner() + inputs = make_inputs() + out = torch.full_like(inputs.query, -123) + assert op.forward(inputs, out=out) is out + assert torch.equal(out, torch.full_like(out, 7 if error_type else 9)) out.fill_(-123) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) + op.forward(inputs, out=out) + assert calls == (["fa3", "triton", "fa3", "triton"] if error_type else ["fa3", "fa3"]) + assert torch.equal(out, torch.full_like(out, 7 if error_type else 9)) + + +@pytest.mark.parametrize("lengths", [(3, 0, 5), (1, 7), (0, 0)]) +def test_sdpa_preserves_sequence_boundaries_and_output_storage(lengths): + inputs = make_inputs(lengths) + inputs.query.zero_() + inputs.key.zero_() + start = 0 + for i, length in enumerate(lengths): + inputs.value[start : start + length].fill_(i + 1) + start += length + out = torch.empty_like(inputs.query) + result = get_varlen_attention_op_class("sdpa")().forward(inputs, out=out) + assert result is out + torch.testing.assert_close(out, inputs.value) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton", "sdpa", "xformers"]) +@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) +@pytest.mark.parametrize("head_dim,strided", [(64, False), (80, True), (128, False)]) +def test_varlen(provider, dtype, atol, head_dim, strided): + op = get_varlen_attention_op_class(provider) + try: + op.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + torch.manual_seed(0) + inputs = make_inputs((1, 0, 65, 129), device="cuda", dtype=dtype, head_dim=head_dim, strided=strided) + storage = torch.full((*inputs.query.shape[:2], head_dim * 2), -123, device="cuda", dtype=dtype) + out = storage[..., :head_dim] + before = [x.clone() for x in (inputs.query, inputs.key, inputs.value)] + result = op().forward(inputs, out=out) + assert result is out + torch.testing.assert_close(out, reference_attention(inputs), atol=atol, rtol=0) + assert torch.equal(storage[..., head_dim:], torch.full_like(out, -123)) + for actual, expected in zip((inputs.query, inputs.key, inputs.value), before): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) -def test_fa3_success_does_not_run_fallback(monkeypatch): - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton", "sdpa", "xformers"]) +def test_empty_batch(provider): + op = get_varlen_attention_op_class(provider) + try: + op.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + inputs = make_inputs((0, 0), device="cuda", dtype=torch.bfloat16) + out = torch.empty_like(inputs.query) + assert op().forward(inputs, out=out) is out - def fa3(q, k, v, o, cu_seqlens, max_seqlen): - o.fill_(9) - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("provider", ["fa3", "triton"]) +def test_graph_replay_uses_new_inputs(provider): + op_cls = get_varlen_attention_op_class(provider) + try: + op_cls.check_environment(get_device_runtime().caps) + except UnsupportedAttention as exc: + pytest.skip(str(exc)) + op = op_cls() + assert op.supports_graph + inputs = make_inputs((3, 5), device="cuda", dtype=torch.bfloat16) + out = torch.empty_like(inputs.query) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op.forward(inputs, out=out) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op.forward(inputs, out=out) + inputs.value.add_(1) + inputs.cu_seqlens.copy_(torch.tensor([0, 5, 8], device="cuda", dtype=torch.int32)) + graph.replay() + torch.testing.assert_close(out, reference_attention(inputs), atol=2e-2, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_audio_attention_matches_projected_qkv(monkeypatch): + from lightllm.models.qwen3_omni_moe_thinker import qwen3_omni_audio as audio + + torch.manual_seed(19) + attention = audio.Qwen3OmniMoeAudioAttention(128, 2, 0).cuda().to(torch.bfloat16).eval() monkeypatch.setattr( - flashattention_nopad, - "_flash_attention_triton_fwd", - lambda *args: pytest.fail("a successful FA3 call must not run fallback"), + audio.g_cache_manager, + "alloc_tensor", + lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device=device), ) - flash_attention_fwd(q, q, q, out, cu, 2) - assert torch.equal(out, torch.full_like(q, 9)) - assert flashattention_nopad._flash_attn_v3_available + hidden = torch.randn(8, 128, dtype=torch.bfloat16, device="cuda") + cu = torch.tensor([0, 3, 8], dtype=torch.int32, device="cuda") + with torch.inference_mode(): + q, k, v = [proj(hidden).reshape(8, 2, 64) for proj in (attention.q_proj, attention.k_proj, attention.v_proj)] + expected = attention.out_proj(reference_attention(VarlenAttentionInputs(q, k, v, cu, 5)).reshape(8, 128)) + actual = attention(hidden, cu_seqlens=cu, max_seqlen=5) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) -def reference_attention_varlen(q, k, v, cu): - """ - q, k, v : (total_len, n_head, D) - cu_seqlen : prefix sums (batch+1,) - """ - total, n_head, d = q.shape - out = torch.empty_like(q) - scale = 1.0 / math.sqrt(d) +@pytest.mark.parametrize("error_type", [RuntimeError, TypeError, AssertionError]) +def test_audio_kernel_errors_propagate_without_retry(monkeypatch, error_type): + from lightllm.common.attention import varlen - for b in range(cu.numel() - 1): - s, e = cu[b].item(), cu[b + 1].item() - q_b, k_b, v_b = q[s:e], k[s:e], v[s:e] # (seq, head, D) + class Broken: + @classmethod + def check_environment(cls, caps): + pass - q_hsd = q_b.permute(1, 0, 2) # (head, seq, D) - k_hds = k_b.permute(1, 2, 0) # (head, D, seq) - v_hsd = v_b.permute(1, 0, 2) # (head, seq, D) + @classmethod + def check_support(cls, inputs, caps=None): + pass - scores = torch.matmul(q_hsd, k_hds) * scale # (head, seq, seq) - probs = torch.softmax(scores.float(), dim=-1) + def forward(self, inputs, *, out): + out.fill_(99) + raise error_type("kernel failed") - out_hsd = torch.matmul(probs, v_hsd.float()) # (head, seq, D) - out[s:e] = out_hsd.permute(1, 0, 2).to(q.dtype) # back to (seq, head, D) + calls = [] - return out + def load(name): + calls.append(name) + return Broken + monkeypatch.setattr(varlen, "get_varlen_attention_op_class", load) + op = AutoVarlenAttentionRunner() + inputs = make_inputs() + with pytest.raises(error_type, match="kernel failed"): + op.forward(inputs, out=torch.empty_like(inputs.query)) + assert calls == ["fa3"] + assert not op.unavailable -@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) -@pytest.mark.parametrize( - "attention_fn", - [ - flash_attention_fwd, - pytest.param( - getattr(flashattention_nopad, "flash_attention_v3_fwd", None), - marks=_fa3_mark, - ), - pytest.param( - Fa3VitAttBackend._vit_att_fwd, - marks=_fa3_mark, - ), - ], -) -def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): - torch.manual_seed(0) - lengths = torch.randint(1, 257, (batch,)) - max_len = int(lengths.max().item()) - - cu = torch.zeros(batch + 1, dtype=torch.int32, device=device) - cu[1:] = torch.cumsum(lengths, 0) - tot = int(cu[-1]) - - q = torch.randn(tot, heads, d, dtype=dtype, device=device) - k = torch.randn_like(q) - v = torch.randn_like(q) - out_tri = torch.randn_like(q) - attention_fn(q, k, v, out_tri, cu, max_len) - a = time.time() - for _ in range(100): - attention_fn(q, k, v, out_tri, cu, max_len) - b = time.time() - print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") - out_ref = reference_attention_varlen(q, k, v, cu) - - max_err = (out_ref - out_tri).abs().max().item() - mean_err = (out_ref - out_tri).abs().mean().item() - print(f"{dtype}: max {max_err:.6f}, mean {mean_err:.6f}") - torch.testing.assert_close(out_tri, out_ref, atol=atol, rtol=0) - - -if __name__ == "__main__": - pytest.main() + +def test_audio_encoder_shares_runner_only_with_its_layers(): + from lightllm.models.qwen3_omni_moe_thinker.qwen3_omni_audio import Qwen3OmniMoeAudioEncoder + + def encoder(): + return Qwen3OmniMoeAudioEncoder( + {}, + d_model=16, + encoder_attention_heads=2, + encoder_layers=2, + encoder_ffn_dim=32, + downsample_hidden_size=4, + output_dim=16, + ) + + first, second = encoder(), encoder() + assert first.layers[0].self_attn.attention_op is first.layers[1].self_attn.attention_op + assert first.layers[0].self_attn.attention_op is not second.layers[0].self_attn.attention_op + + +@pytest.mark.parametrize("boundaries,maximum", [([1, 3, 8], 5), ([0, 7, 6, 8], 7), ([0, 3, 7], 5), ([0, 3, 8], 4)]) +def test_varlen_rejects_invalid_boundaries(boundaries, maximum): + inputs = make_inputs()._replace(cu_seqlens=torch.tensor(boundaries, dtype=torch.int32), max_seqlen=maximum) + with pytest.raises(ValueError, match="cu_seqlens"): + inputs.validate() + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_auto_validates_once_per_call_and_rechecks_mutated_boundaries(monkeypatch, device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + inputs = make_inputs(device=device) + original = VarlenAttentionInputs.validate + calls = [] + + def validate(self): + calls.append(self.cu_seqlens) + return original(self) + + monkeypatch.setattr(VarlenAttentionInputs, "validate", validate) + op = AutoVarlenAttentionRunner(providers=("sdpa",)) + output = torch.empty_like(inputs.query) + op.forward(inputs, out=output) + inputs.cu_seqlens.copy_(torch.tensor([0, 1, 1, 8], dtype=torch.int32, device=device)) + op.forward(inputs._replace(max_seqlen=7), out=output) + assert len(calls) == 2 + torch.testing.assert_close(output, reference_attention(inputs._replace(max_seqlen=7))) + + +@pytest.mark.parametrize("provider", ["fa3", "triton", "xformers", "sdpa"]) +@pytest.mark.parametrize("invalid", ["shape", "dtype", "device", "overlap", "alias", "head_stride"]) +def test_varlen_output_validation(provider, invalid): + inputs = make_inputs() + out = torch.empty_like(inputs.query) + if invalid == "shape": + out = out[:-1] + elif invalid == "dtype": + out = out.double() + elif invalid == "device": + out = out.to("meta") + elif invalid == "overlap": + out = out[:1].expand_as(out) + elif invalid == "alias": + out = inputs.key + else: + out = torch.empty((*out.shape[:-1], out.shape[-1] * 2))[..., ::2] + with pytest.raises(ValueError, match="varlen out"): + get_varlen_attention_op_class(provider)().forward(inputs, out=out) diff --git a/unit_tests/server/core/objs/test_req.py b/unit_tests/server/core/objs/test_req.py index 5966497ef1..e35a59917e 100644 --- a/unit_tests/server/core/objs/test_req.py +++ b/unit_tests/server/core/objs/test_req.py @@ -1,7 +1,9 @@ -import pytest -import easydict from types import SimpleNamespace -from lightllm.server.core.objs.req import Req, ChunkedPrefillReq, SamplingParams + +import easydict +import pytest + +from lightllm.server.core.objs.req import ChunkedPrefillReq, Req, SamplingParams from lightllm.server.core.objs.token_metadata import ReqFinalTokenMetadata from lightllm.utils import shm_utils from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args @@ -16,8 +18,8 @@ def setup_module_env(): easydict.EasyDict( { "mtp_step": 0, - "llm_prefill_att_backend": ["None"], - "llm_decode_att_backend": ["None"], + "llm_prefill_att_backend": "None", + "llm_decode_att_backend": "None", "cpu_cache_token_page_size": 256, "enable_cpu_cache": False, "model_dir": "", diff --git a/unit_tests/server/core/objs/test_shm_req_manager.py b/unit_tests/server/core/objs/test_shm_req_manager.py index c51e3e6b95..cc0dafe109 100644 --- a/unit_tests/server/core/objs/test_shm_req_manager.py +++ b/unit_tests/server/core/objs/test_shm_req_manager.py @@ -1,16 +1,17 @@ -import os import asyncio -import pytest +import os import time import uuid -import numpy as np -from unittest.mock import MagicMock from types import SimpleNamespace +from unittest.mock import MagicMock +import numpy as np +import pytest from easydict import EasyDict -from lightllm.utils.envs_utils import set_env_start_args, get_env_start_args -from lightllm.utils import shm_utils + from lightllm.server.core.objs.shm_req_manager import ShmReqManager +from lightllm.utils import shm_utils +from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args @pytest.fixture(scope="module", autouse=True) @@ -23,8 +24,8 @@ def setup_env(): running_max_req_size=10, disable_chunked_prefill=True, mtp_step=0, - llm_prefill_att_backend=["None"], - llm_decode_att_backend=["None"], + llm_prefill_att_backend="None", + llm_decode_att_backend="None", cpu_cache_token_page_size=256, enable_cpu_cache=False, ) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index ece87927e0..678568242e 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -1,23 +1,39 @@ -from types import SimpleNamespace from collections import deque +from types import SimpleNamespace import pytest import torch -from lightllm.server.multi_level_kv_cache import ( - AdaptiveCachePlacementController, - CacheCapacityConfig, - CacheTier, -) -from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule -from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl +from lightllm.server.multi_level_kv_cache import AdaptiveCachePlacementController, CacheCapacityConfig, CacheTier from lightllm.server.router.model_infer.infer_batch import InferenceContext +from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl +from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule + + +@pytest.mark.parametrize("sync", [False, True]) +def test_attention_cpu_cache_sync(sync): + from lightllm.common.attention.api import AttentionPhase + + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.backend = SimpleNamespace( + model=SimpleNamespace( + attention=SimpleNamespace( + groups=( + { + AttentionPhase.PREFILL: SimpleNamespace(requires_cpu_cache_sync=False), + AttentionPhase.DECODE: SimpleNamespace(requires_cpu_cache_sync=sync), + }, + ) + ) + ) + ) + assert module.need_sync_compute_stream() is sync def test_running_cpu_cache_tasks_are_not_ready(monkeypatch): from lightllm.server.router.model_infer.infer_batch import InferReq - from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.server.router.model_infer.mode_backend import base_backend as base_backend_impl + from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend status = InferReq._CpuCacheTaskStatus @@ -494,14 +510,14 @@ def query(self): def test_fa3_cpu_cache_sync_defaults_off(): - from lightllm.common.basemodel.attention.base_att import BaseAttBackend - from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend - from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend - from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend - from lightllm.common.basemodel.attention.fa3.neo import NeoFa3AttBackend + from lightllm.common.attention.api.backend import AttentionBackend + from lightllm.common.attention.backends.dense.fa3.backend import Fa3DenseDecodeBackend + from lightllm.common.attention.backends.dense.fa3.fp8 import Fa3DenseFp8DecodeBackend + from lightllm.common.attention.backends.dense.fa3.neo import Fa3DenseNeoDecodeBackend + from lightllm.common.attention.backends.mla.fa3.backend import Fa3MLADecodeBackend - assert BaseAttBackend.requires_cpu_cache_sync is False + assert AttentionBackend.requires_cpu_cache_sync is False assert all( backend.requires_cpu_cache_sync is False - for backend in (Fa3AttBackend, Fp8Fa3AttBackend, MlaFa3AttBackend, NeoFa3AttBackend) + for backend in (Fa3DenseDecodeBackend, Fa3DenseFp8DecodeBackend, Fa3MLADecodeBackend, Fa3DenseNeoDecodeBackend) ) diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py index 21bbb06339..17774fb802 100644 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -2,34 +2,35 @@ from __future__ import annotations -from unit_tests.common.att_resources.model_helpers import full_state_test_fields - -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -from typing import TYPE_CHECKING import dataclasses import json from types import SimpleNamespace +from typing import TYPE_CHECKING import pytest import torch -from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states -from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel - from lightllm.models import Deepseek2TpPartModel, Deepseek3_2TpPartModel, LlamaTpPartModel -from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_unique_server_name from lightllm.utils.kv_cache_utils import calcu_cpu_cache_meta +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + main_linear_state, + prepare_model_states, +) if TYPE_CHECKING: from collections.abc import Callable + from pathlib import Path + from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient from lightllm.utils.kv_cache_utils import CpuKVCacheMeta - from pathlib import Path pytestmark = [ @@ -372,9 +373,8 @@ def capture_write(*args: object, **kwargs: object) -> None: pool.ssm_state[slot].fill_(5) producer = torch.cuda.Stream() producer.wait_stream(torch.cuda.current_stream()) + monkeypatch.setattr(common.g_infer_context, "get_overlap_stream", lambda: producer) with torch.cuda.stream(producer): - # 延迟当前生产流的 KV 写入；当前实现不在此处隐式插入 wait_stream， - # 生产者与搬运流的依赖由调用方负责组织。 torch.cuda._sleep(10_000_000) (full,) = resources.get_full_attention_states("main") for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): @@ -429,7 +429,7 @@ def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: patch.setattr(torch.cuda.Stream, "wait_stream", record_wait) module.get_cpu_offload_stream = lambda: copy_stream task = module.write_req_to_cpu(req) - assert waited_streams == [] + assert waited_streams == [producer] task.sync_event.synchronize() assert write_batches == [(pages, [True, False])] assert req.cpu_cache_offload_task_status.is_running() diff --git a/unit_tests/server/test_api_start_deepep.py b/unit_tests/server/test_api_start_deepep.py index 949021cd34..c1f417765f 100644 --- a/unit_tests/server/test_api_start_deepep.py +++ b/unit_tests/server/test_api_start_deepep.py @@ -6,7 +6,9 @@ from lightllm.server.core.objs.start_args_type import StartArgs -def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path): +@pytest.mark.parametrize("backend", ["auto", "fa3", "triton", "flashinfer"]) +@pytest.mark.parametrize("phase", ["prefill", "decode"]) +def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path, backend, phase): for name in ( "_set_envs_and_config", "auto_set_max_req_total_len", @@ -39,6 +41,14 @@ def stop(args): router_token_ratio=0.0, ) + setattr(args, f"llm_{phase}_att_backend", backend) + args.llm_prefill_linear_att_backend = "flashqla" + args.llm_decode_linear_att_backend = "triton" + if backend == "flashinfer": + with pytest.raises(AssertionError, match=f"--llm_{phase}_att_backend"): + _launch_subprocesses(args) + return + with pytest.raises(RuntimeError, match="startup validation passed"): _launch_subprocesses(args) diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 1c391c09cd..7aa41c2193 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -107,6 +107,7 @@ def test_whole_image_blocks_fit_batch_at_startup( @pytest.mark.parametrize("mode,step", [("native", 0), ("replay", 0), ("replay", 4)]) def test_s3_configuration_is_prepared_before_ssm_validation(launch_args, monkeypatch, tmp_path, mode, step): from pathlib import Path + from lightllm.utils.petrel_helper import PetrelHelper monkeypatch.chdir(tmp_path) diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py index 1f4efc4baa..4e77c6237e 100644 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -3,9 +3,9 @@ from types import SimpleNamespace import numpy as np +import pytest import torch - REPO_ROOT = Path(__file__).resolve().parents[1] STATIC_BENCHMARK_DIR = REPO_ROOT / "test" / "benchmark" / "static_inference" for path in (REPO_ROOT, STATIC_BENCHMARK_DIR): @@ -153,3 +153,50 @@ def test_profile_capacity_uses_page_aligned_per_request_cost(): args.mtp_mode = "eagle_with_att" args.mtp_step = 3 assert decode_profile_batch_divisor(args, decode_case) == 24 + + +@pytest.mark.parametrize("fail", [False, True]) +def test_benchmark_shared_memory_cleanup(monkeypatch, fail): + from multiprocessing import shared_memory + + from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt + + monkeypatch.delenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", raising=False) + monkeypatch.setenv("LIGHTLLM_START_ARGS", "{}") + static_benchmark.get_unique_server_name.cache_clear() + names = [] + + def run(args, cases): + counter = SharedInt("benchmark_counter") + names.append(counter.shm.name) + counter.set_value(17) + assert counter.get_value() == 17 + counter.shm.close() + if fail: + raise RuntimeError("worker failed") + + monkeypatch.setattr(static_benchmark, "run_benchmark", run) + argv = [ + "--benchmark", + "decode", + "--batch_size", + "1", + "--input_len", + "4", + "--output_len", + "1", + "--data_type", + "bfloat16", + "--enable_fused_shared_experts", + ] + try: + if fail: + with pytest.raises(RuntimeError, match="worker failed"): + static_benchmark.main(argv) + else: + static_benchmark.main(argv) + assert len(names) == 1 + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=names[0]) + finally: + static_benchmark.get_unique_server_name.cache_clear() diff --git a/unit_tests/utils/test_sgl_utils.py b/unit_tests/utils/test_sgl_utils.py index 85bb041e68..151fd8345c 100644 --- a/unit_tests/utils/test_sgl_utils.py +++ b/unit_tests/utils/test_sgl_utils.py @@ -1,19 +1,28 @@ -import weakref import collections import inspect import math +import weakref from types import SimpleNamespace import pytest import torch from frozendict import frozendict +from lightllm.common.attention.backends.dense.fa3 import backend as fa3_module from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.fa3 import fp as fa3_module from lightllm.common.triton_utils import autotuner as autotuner_module from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner from lightllm.utils import sgl_utils from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len +from unit_tests.common.attention.helpers import ( + backend_context, + make_backend, + make_batch, + make_context, + run_dense, + run_sparse, + update_context, +) @pytest.fixture(autouse=True) @@ -114,33 +123,36 @@ def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch, pa infer_state.b_seq_len = torch.tensor([4096, actual_kv_len], dtype=torch.int32) infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) infer_state.b1_cu_kv_seq_len = torch.tensor([0, 4096, 4096 + actual_kv_len], dtype=torch.int32) - state = fa3_module.Fa3DecodeAttState( - backend=SimpleNamespace( + state = backend_context( + fa3_module.Fa3DenseDecodeBackend, + "decode", + settings=SimpleNamespace( model=model, - infer_page_size=page_size, + kernel_page_size=page_size, uses_causal_attention=lambda: True, - uses_dynamic_spec_verify_layout=lambda: False, + uses_dynamic_verify_layout=lambda: False, get_page_table_view=lambda att_batch_size, max_kv_len, microbatch_index: torch.zeros( att_batch_size, math.ceil(max_kv_len / page_size), dtype=torch.int32 ), ), - _infer_state_weak_ref=weakref.ref(infer_state), ) copy_calls = [] - monkeypatch.setattr(fa3_module, "page_table_copy", lambda **kwargs: copy_calls.append(kwargs)) - state.init_state() + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.fa3_utils.page_table_copy", lambda **kwargs: copy_calls.append(kwargs) + ) + update_context(state, infer_state) # 模拟 CudaGraph._capture_decode 对 infer_state 的修改。 - state.infer_state.max_kv_seq_len = graph.graph_max_len_in_batch + infer_state.max_kv_seq_len = graph.graph_max_len_in_batch calls = [] def attention(**kwargs): calls.append(kwargs) return kwargs["q"] - monkeypatch.setattr(fa3_module, "flash_attn_with_kvcache_autotune", attention) - state.decode_att(torch.empty(2, 4, 8), torch.empty(page_size, 2, 8), torch.empty(page_size, 2, 8)) + monkeypatch.setattr("lightllm.utils.sgl_utils.flash_attn_with_kvcache_autotune", attention) + run_dense(state, infer_state, torch.empty(2, 4, 8), torch.empty(page_size, 2, 8), torch.empty(page_size, 2, 8)) assert calls[0]["max_seqlen_k"] == actual_kv_len - assert calls[0]["page_table"].shape == (2, 32768 // page_size) + assert calls[0]["page_table"].shape == (2, math.ceil(actual_kv_len / page_size)) assert copy_calls[0]["page_table"].shape == (2, math.ceil(actual_kv_len / page_size)) assert copy_calls[0]["page_size"] == page_size diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py index a58b35198d..461242adcc 100644 --- a/unit_tests/utils/test_speculative_utils.py +++ b/unit_tests/utils/test_speculative_utils.py @@ -1,21 +1,24 @@ -import weakref +from unit_tests.common.attention.helpers import build import json +import weakref from importlib import import_module from types import SimpleNamespace import pytest import torch -from lightllm.common.basemodel.infer_struct import InferStateInfo -import lightllm.common.basemodel.attention.base_att as base_att_module +import lightllm.common.attention.api.backend as base_att_module import lightllm.common.basemodel.hidden_collector as hidden_collector_module -from lightllm.common.basemodel.attention.base_att import BaseAttBackend -from lightllm.common.basemodel.attention.fa3.fp import Fa3DecodeAttState, Fa3PrefillAttState -from lightllm.common.basemodel.attention.fa3.mla import MlaFa3DecodeAttState, MlaFa3PrefillAttState -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.common.attention.api import AttentionPhase +from unit_tests.common.attention.helpers import StubBackend, make_env +from lightllm.common.attention.backends.dense.fa3.backend import Fa3DenseDecodeBackend +from lightllm.common.attention.backends.dense.triton.backend import TritonDenseDecodeBackend +from lightllm.common.attention.backends.mla.fa3.backend import Fa3MLADecodeBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.models import get_draft_model_class from lightllm.models.qwen3_eagle.layer_weights.transformer_layer_weight import Qwen3EagleTransformerLayerWeight from lightllm.utils import envs_utils +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_spec, model_stub @pytest.mark.parametrize( @@ -76,19 +79,11 @@ def test_attention_backend_selects_dynamic_spec_layout( dynamic_spec, expected, ): - monkeypatch.setattr( - base_att_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_mode=mtp_mode, mtp_dynamic_verify=dynamic_spec), - ) - backend = SimpleNamespace( - model=SimpleNamespace( - is_mtp_draft_model=is_draft_model, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), - ) - ) - - assert BaseAttBackend.uses_dynamic_spec_verify_layout(backend) is expected + model = model_stub(mtp_mode=mtp_mode, mtp_dynamic_verify=dynamic_spec) + model.is_mtp_draft_model = is_draft_model + model.mtp_manager.get_decode_draft_step = lambda _: draft_step + backend = StubBackend(make_env(model), _spec_from_model(model)) + assert backend.spec.dynamic_verify is expected @pytest.mark.parametrize( @@ -104,102 +99,59 @@ def test_attention_backend_selects_dynamic_spec_layout( ], ) def test_attention_backend_selects_causality(monkeypatch, mtp_mode, is_draft_model, expected): - monkeypatch.setattr( - base_att_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_mode=mtp_mode), - ) - backend = SimpleNamespace(model=SimpleNamespace(is_mtp_draft_model=is_draft_model)) - - assert BaseAttBackend.uses_causal_attention(backend) is expected + model = model_stub(mtp_mode=mtp_mode) + model.is_mtp_draft_model = is_draft_model + assert StubBackend(make_env(model), _spec_from_model(model)).spec.causal is expected -@pytest.mark.parametrize("state_class", [Fa3PrefillAttState, MlaFa3PrefillAttState]) +@pytest.mark.parametrize("state_class", [Fa3DenseDecodeBackend, Fa3MLADecodeBackend]) def test_fa3_prefill_state_owns_causality(monkeypatch, state_class): - monkeypatch.setattr(import_module(Fa3PrefillAttState.__module__), "page_table_copy", lambda **kwargs: None) - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) - infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) - infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = 4 - infer_state.input_ids = torch.empty(2, dtype=torch.int64) - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, dtype=torch.int32).reshape(2, 4)) - state = state_class( - backend=SimpleNamespace(uses_causal_attention=lambda: False, page_size=1, infer_page_size=1), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - - state.init_state() + backend = make_backend(state_class, causal=False) + monkeypatch.setattr("lightllm.common.basemodel.triton_kernel.fa3_utils.page_table_copy", lambda **kwargs: None) + context = make_context(backend, make_batch(phase=AttentionPhase.PREFILL)) + build(backend, context) + assert context.causal is False - assert state.causal is False - -@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) +@pytest.mark.parametrize("state_class", [Fa3DenseDecodeBackend, Fa3MLADecodeBackend]) def test_fa3_decode_state_owns_causality(state_class): - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) - infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) - infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) - infer_state.b_seq_len = torch.tensor([3, 4], dtype=torch.int32) - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), - ) - state = state_class( - backend=SimpleNamespace( - model=model, - uses_causal_attention=lambda: False, - uses_dynamic_spec_verify_layout=lambda: False, - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state._init_page_table = lambda _: None - - state.init_state() - - assert state.causal is False + backend = make_backend(state_class, causal=False) + backend._decode_init_page_table = lambda *args: None + context = make_context(backend) + build(backend, context) + assert context.causal is False @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_triton_mtp_decode_state_builds_group_markers(): - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), - req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), + backend = make_backend(TritonDenseDecodeBackend, draft_step=2, hold_req_id=-1) + batch = make_batch( + (4, 4, 2, 2), + device="cuda", + b_req_idx=torch.tensor([7, 7, -1, -1], device="cuda", dtype=torch.int32), + hold_req_id=-1, ) - infer_state = InferStateInfo() - infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") - infer_state.max_kv_seq_len = 4 - state = TritonDecodeAttState( - backend=SimpleNamespace(model=model), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - - state.init_state() - - assert state.b_mark_mtp_shared_group.cpu().tolist() == [0, 2, 1, 1] + context = make_context(backend, batch) + build(backend, context) + assert context.b_mark_mtp_shared_group.cpu().tolist() == [0, 2, 1, 1] @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) +@pytest.mark.parametrize("state_class", [Fa3DenseDecodeBackend, Fa3MLADecodeBackend]) def test_fa3_dynamic_decode_state_builds_group_markers(state_class): - model = SimpleNamespace(req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1)) - infer_state = InferStateInfo() - infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") - infer_state.b_seq_len = torch.tensor([3, 4, 2, 2], dtype=torch.int32, device="cuda") - infer_state.batch_size = 4 - state = state_class( - backend=SimpleNamespace(model=model), - _infer_state_weak_ref=weakref.ref(infer_state), + backend = make_backend(state_class, hold_req_id=-1) + batch = make_batch( + (3, 4, 2, 2), + device="cuda", + b_req_idx=torch.tensor([7, 7, -1, -1], device="cuda", dtype=torch.int32), + hold_req_id=-1, ) - - b_att_req_idx = state._init_dynamic_spec_verify_state(draft_step=2) + context = make_context(backend, batch) + b_att_req_idx = backend._decode_init_dynamic_verify_metadata(context, draft_step=2) torch.cuda.synchronize() - assert b_att_req_idx.cpu().tolist() == [7, -1, -1, -1] - assert state.b_att_seq_len.cpu().tolist() == [4, 2, 2, 0] - assert state.cu_seqlens_q.cpu().tolist() == [0, 2, 3, 4, 4] + assert context.b_att_seq_len.cpu().tolist() == [4, 2, 2, 0] + assert context.cu_seqlens_q.cpu().tolist() == [0, 2, 3, 4, 4] @pytest.mark.parametrize( @@ -411,3 +363,29 @@ def test_eagle3_added_kv_layers_come_from_draft_config(tmp_path): ) assert envs_utils.get_added_mtp_kv_layer_num() == 2 + + +def _spec_from_model(model): + from lightllm.common.basemodel.attention_adapters import ADAPTERS + from lightllm.common.attention.api import AttentionFamily + from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer + from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer + + state = SimpleNamespace( + contains_logical_layer=lambda layer: True, + describe_layout=lambda: (), + storage_dtype=torch.bfloat16, + config=SimpleNamespace(), + ) + model.data_type = torch.bfloat16 + model.att_state_manager.get_full_attention_states = lambda role: (state,) + model.att_state_manager.get_state = lambda *args: state + cls = ( + Qwen3DFlashTransformerLayerInfer + if model.is_mtp_draft_model and model.args.mtp_mode in ("dspark", "dflash", "dflash2") + else LlamaTransformerLayerInfer + ) + layer = cls.__new__(cls) + layer.layer_num_, layer.tp_q_head_num_, layer.tp_k_head_num_, layer.head_dim_ = 0, 4, 2, 64 + description = layer.describe_attention(AttentionPhase.DECODE) + return ADAPTERS[AttentionFamily.DENSE].describe(model, layer, AttentionPhase.DECODE, description) From 90f31e945c3858a73228534695577751b6f99b77 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 22:13:32 +0800 Subject: [PATCH 7/9] docs: describe attention execution design --- .../framework/attention_execution_design.md | 128 ++++++++++++++++++ docs/CN/source/index.rst | 1 + 2 files changed, 129 insertions(+) create mode 100644 docs/CN/source/framework/attention_execution_design.md diff --git a/docs/CN/source/framework/attention_execution_design.md b/docs/CN/source/framework/attention_execution_design.md new file mode 100644 index 0000000000..e9eee53a42 --- /dev/null +++ b/docs/CN/source/framework/attention_execution_design.md @@ -0,0 +1,128 @@ +# Attention 执行设计 + +这次重构将 attention 的长期存储、静态计算配置和批次执行分开：State 管存储，Backend 实现单个阶段的计算，Session 统一保证执行约束。模型层只描述需求、提供输入并使用结果。 + +设计目标是让新增模型或后端时，修改集中在对应实现中；容量、资源复用、Graph 输入和输出回写等公共规则只实现一次。存储部分沿用 [Attention Resources 设计](cache_manager_design.md)。 + +## 1. 职责和数据归属 + +| 对象 | 持有的内容 | 职责 | +| --- | --- | --- | +| `State` / `AttentionStateManager` | KV、Conv、SSM 等长期存储及其布局 | 分配存储、提供层视图、checkpoint、恢复和传输；外部页索引由 allocator 和调用方管理 | +| `ModelAttention` | backend 分组、层路由、adapter、workspace 池和存储代次 | 初始化选型、资源预留、层绑定，将模型推理信息转换为执行输入 | +| `AttentionBackend` | 不可变 `LayerSpec` 和显式注入的运行依赖 | 实现一个 family、provider、phase 下的规划与计算 | +| `LayerAttention` | 层号、组号、backend 和资源视图 | 将层计算路由到 Session；借用 State 存储，不拥有其生命周期 | +| `AttentionSession` | 当前 `AttentionBatch` 和各组 `AttentionContext` | 校验 owner、phase、capacity，管理执行期间的 workspace lease 和输出回写 | +| `CapturedGraph` | decode 捕获的推理状态、输出、workspace 和 provider hooks | 保持捕获地址有效，协调重放输入、辅助任务及完成事件 | + +三个数据类型各表达一种事实： + +- `LayerSpec`：静态需求，包括算法、阶段、head 维度、特性及 `CacheLayout`。只由 backend 持有，层绑定不保存副本。 +- `AttentionBatch`：本次请求索引、长度、位置等输入；Tensor 借用自 `InferStateInfo` / `ReqManager`。 +- `AttentionContext`：backend 为本次执行或某个 Graph 捕获实例准备的 metadata、容量和 workspace 引用。同一 Session 内的 context 共享一份 batch。 + +`InferStateInfo` 单向持有 Session。Session 只接收 batch、分组、容量和 workspace 提供函数，不反向引用模型或 `InferStateInfo`。Backend 可以通过资源视图更新 KV/Conv/SSM，但不接管存储分配、checkpoint 或请求回收。 + +## 2. 初始化：描述、选择、绑定 + +```text +模型层 describe_attention(phase) + State 布局 + │ + family adapter + │ + (prefill_spec, decode_spec) + │ + ModelAttention 按完整 spec 对分组 + │ + 每组分别选择 prefill / decode backend + │ + 预留 workspace → 分配 State 存储 + │ + adapter 借用层资源 → LayerAttention +``` + +只有完整的 `(prefill_spec, decode_spec)` 相同，层才能共享 backend 和批次 context；权重、逻辑层号和存储视图仍留在各自的 `LayerAttention`。同一层的两个阶段必须使用相同的物理 `CacheLayout`，避免 prefill 写入的内容无法被 decode 读取。 + +Backend 构造统一为 `Backend(env, spec)`。`BackendEnv` 显式传入设备 runtime、allocator、TP 和 MTP 等运行信息；量化组长等布局参数读取 `spec.cache`。能力声明放在 backend 类上，注册表只枚举实现和默认优先级。选择时检查环境、计算需求和服务约束，并执行对应探针；显式指定失败时报错，`auto` 才尝试其他候选。 + +每层选择一次 family adapter，描述和绑定复用它。Dense 与 MLA 是独立适配；GDN 与 KDA 只共享 Conv/SSM 布局描述。DeepSeekV4 等模型专有适配放在模型目录。 + +## 3. 执行：调用方统一保证约束 + +模型入口使用一个作用域管理执行准备与收尾： + +```python +with model.attention.session(infer_state): + output = model._token_forward(infer_state) +``` + +普通执行按以下顺序进行；decode 捕获时的 device metadata 和计算进入 Graph： + +```text +ModelAttention.session + → 检查存储代次，构造 batch 和 capacity + → AttentionSession.prepare + → 校验 batch，取得 workspace lease + → backend.allocate_context + → backend.plan_host # 图外主机规划 + → 模型 forward + → session.build_device_metadata # 设备 metadata + → 逐层 LayerAttention.forward + → Session 检查 context 归属 + → backend.forward # 阶段计算 + → Session 处理 out 回写 + → 退出作用域，结束 lease +``` + +Backend 的执行接口只有 `allocate_context`、`plan_host`、`build_device_metadata` 和 `forward`；资源预算与捕获 hooks 按需提供。每个实例只实现 prefill 或 decode，阶段专属构造和 helper 放在对应类中，共享类只保留共同定义。 + +Session 保证 context 属于当前 backend 和 phase，backend 返回的容量必须等于请求容量。`batch_size`、`token_num` 必须与容量契约一致，序列长度不能超过上界；padding 和 Graph bucket 选择由模型执行路径负责。 + +提供 `out` 时，Session 统一校验 dtype、device 和结果范围：backend 可直接写入 `out` 或返回它的合法视图；返回独立结果时由 Session 复制，并保留辅助输出。 + +Workspace 按 `(group_id, microbatch_index)` 隔离。准备、规划或 forward 抛出异常时，`ExitStack` 仍结束已取得的 lease。结束 lease 会记录完成事件供后续复用等待，不释放池中的 Tensor，也不表示 GPU 已同步完成。 + +## 4. CUDA Graph：固定地址，先校验再写入 + +decode 捕获时，所有组的 context 使用同一份捕获 batch。长度容量显式采用 `graph_max_len_in_batch`，当前实际长度另行传给 host planner，避免用服务最大长度扩大捕获容量。捕获的 workspace 保留原 Tensor 引用并冻结；池后续扩容不替换旧 Graph 使用的地址。 + +`CapturedGraph` 一次重放的顺序为： + +1. 校验存储代次，以及所有 microbatch lane 的模型角色、backend 分组、容量和输入 Tensor 形状、dtype、device；只生成复制计划。 +2. 所有 lane 校验通过后，等待 provider 上次 metadata 使用结束，并取得 workspace lease。 +3. 将新输入复制到捕获地址，调用各组 `plan_host` 更新主机规划，再提交 Graph 和 provider 辅助任务。 +4. 将辅助任务完成事件纳入 lease 收尾，保证下一次复用的依赖完整。 + +因此，后一个 lane 的输入校验失败不会留下前一个 lane 已改写的输入。这个保证覆盖写入前的校验，不提供复制或 provider 执行失败后的事务回滚。 + +请求映射表是共享存储，重放时必须保持地址；State 存储代次改变后必须重新绑定层并重新捕获。Graph 销毁时等待相关设备工作，关闭 provider hooks，再释放捕获引用。 + +prefill 采用另一条路径：`PrefillGraphExecution` 持有分段 Graph 的输入、输出和 collector，attention 通过段间回调使用当前批次的 Session。`build_device_metadata` 在分段执行前运行，不属于上述 decode 捕获流程；多 lane 写入前校验的保证也仅针对 `CapturedGraph` 路径。 + +## 5. 本次精简的取舍 + +| 原有复杂度 | 当前处理 | +| --- | --- | +| `AttentionBackend` / `BaseAttentionBackend` 两层基类及构造转发 | 合并为一个基类，删除 `create()`、`configure_spec()` | +| spec、causal、dynamic verify 等配置重复保存 | 从 backend 的 `spec` 读取，避免多个副本需要同步 | +| BatchBuilder 单方法类、`ContextSource` 间接接口 | 模型侧直接构造 batch，Session 接收明确的数据 | +| prepare / forward / finish 分散配对 | 一个 Session 作用域覆盖正常与异常收尾 | +| prefill / decode 的 helper 和依赖混在共享类 | 归入阶段实现；例如 GDN decode 构造不加载 prefill kernel | +| GDN 继承 Dense、KDA 继承 GDN | 以实际共享的 Conv/SSM 布局为基础组织 adapter | + +保留 family adapter，是因为模型字段与 State 资源确实存在差异；保留 Session，是因为生命周期和校验需要被所有后端复用。State 已能提供布局和资源视图，不再增加一套平行的存储描述接口。 + +## 6. 修改入口与验证 + +新增模型时先实现层需求描述，优先复用 family adapter；只有资源绑定不同才增加模型专有 adapter。新增 provider 时实现对应阶段的 backend、声明能力并注册；新增存储编码时扩展 State 的布局和资源，交由 `CacheLayout` 做兼容性检查。具体接入见[添加新模型](../models/add_new_model.md)。 + +代码阅读入口： + +| 内容 | 文件 | +| --- | --- | +| 分组、选型与模型适配 | `lightllm/common/basemodel/attention.py`、`lightllm/common/basemodel/attention_adapters.py` | +| 静态需求与后端接口 | `lightllm/common/attention/api/spec.py`、`lightllm/common/attention/api/backend.py` | +| 批次执行与 workspace | `lightllm/common/attention/session.py`、`lightllm/common/attention/api/workspace.py` | +| Graph 重放 | `lightllm/common/basemodel/graph.py`、`lightllm/common/basemodel/prefill_cuda_graph.py`、`lightllm/common/basemodel/infer_struct.py` | + +回归重点是 context 归属和容量拒绝、异常 lease 收尾、输出回写、多 lane 写入前校验，以及捕获地址和存储代次。接口契约与 workspace 测试位于 `unit_tests/common/attention/`；模型绑定和后端数值测试位于 `unit_tests/common/basemodel/attention/`。模型验证还需覆盖 prefill/decode、Graph、MTP 和状态恢复；吞吐或跨机 PD 结论需要各自的运行验证。 diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index f15caff1bf..1b952ddce4 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -82,6 +82,7 @@ Lightllm 整合了众多的开源方案的优点，包括但不限于 FasterTran 架构介绍 Attention Resources 架构设计 + Attention 执行设计 token attention介绍 峰值显存调度器介绍 From 44e270b424905bcc64debfe24131d155f316dd01 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 23:18:26 +0800 Subject: [PATCH 8/9] docs: add English attention execution design --- .../framework/attention_execution_design.md | 128 ++++++++++++++++++ docs/EN/source/index.rst | 1 + 2 files changed, 129 insertions(+) create mode 100644 docs/EN/source/framework/attention_execution_design.md diff --git a/docs/EN/source/framework/attention_execution_design.md b/docs/EN/source/framework/attention_execution_design.md new file mode 100644 index 0000000000..f663153e99 --- /dev/null +++ b/docs/EN/source/framework/attention_execution_design.md @@ -0,0 +1,128 @@ +# Attention Execution Design + +This refactor separates persistent attention storage, static computation requirements, and batch execution. States own storage, backends implement computation for one phase, and sessions enforce execution contracts. Model layers describe their requirements, supply inputs, and consume results. + +The aim is to keep changes for a new model or backend within its implementation, while defining capacity checks, resource reuse, graph inputs, and output handling once. Storage follows the existing [Attention Resources design (Chinese)](https://github.com/ModelTC/LightLLM/blob/refactor/attention-session-rewrite/docs/CN/source/framework/cache_manager_design.md). + +## 1. Responsibilities and ownership + +| Object | Holds | Responsibility | +| --- | --- | --- | +| `State` / `AttentionStateManager` | Persistent KV, Conv, and SSM storage and its layout | Allocate storage, expose layer views, checkpoint, restore, and transfer data; allocators and callers manage external page indices | +| `ModelAttention` | Backend groups, layer routes, adapters, workspace pools, and storage generation | Select backends, reserve resources, bind layers, and translate model inference state into execution inputs | +| `AttentionBackend` | An immutable `LayerSpec` and explicitly injected runtime dependencies | Plan and compute for one family, provider, and phase | +| `LayerAttention` | Layer ID, group ID, backend, and resource views | Route layer computation through the session; borrow State storage without owning its lifetime | +| `AttentionSession` | The current `AttentionBatch` and each group's `AttentionContext` | Validate ownership, phase, and capacity; manage workspace leases and output handling during execution | +| `CapturedGraph` | Captured decode inference states, outputs, workspaces, and provider hooks | Keep captured addresses valid and coordinate replay inputs, auxiliary work, and completion events | + +Three data types describe distinct facts: + +- `LayerSpec`: static requirements, including algorithm, phase, head dimensions, features, and `CacheLayout`. The backend holds it; layer bindings do not keep a second copy. +- `AttentionBatch`: request indices, lengths, positions, and other inputs for this execution. Tensors are borrowed from `InferStateInfo` / `ReqManager`. +- `AttentionContext`: metadata, capacity, and workspace references prepared by the backend for an execution or a graph capture instance. Contexts within a session share one batch. + +`InferStateInfo` holds the session in one direction. The session receives only the batch, groups, capacity, and a workspace supplier; it does not reference the model or `InferStateInfo`. Backends can update KV/Conv/SSM through resource views, while storage allocation, checkpointing, and request reclamation remain with their existing owners. + +## 2. Initialization: describe, select, bind + +```text +Layer describe_attention(phase) + State layout + | + Family adapter + | + (prefill_spec, decode_spec) + | + ModelAttention groups by the full spec pair + | + Select a prefill and decode backend per group + | + Reserve workspace -> Allocate State storage + | + Adapter borrows layer resources -> LayerAttention +``` + +Layers share backends and batch contexts only when their complete `(prefill_spec, decode_spec)` pairs match. Weights, logical layer IDs, and storage views remain in each layer's `LayerAttention`. Both phases of a layer must use the same physical `CacheLayout`, so decode can read what prefill writes. + +All backends are constructed as `Backend(env, spec)`. `BackendEnv` explicitly supplies the device runtime, allocator, TP and MTP settings, and other runtime information. Layout parameters such as quantization group size come from `spec.cache`. Backend classes declare capabilities; the registry lists implementations and default priorities. Selection checks the environment, computation requirements, and serving constraints, then runs the corresponding probes. An explicit choice fails with an error when unusable; only `auto` tries other candidates. + +Each layer selects its family adapter once and reuses it for description and binding. Dense and MLA have separate adapters. GDN and KDA share only the Conv/SSM layout description. Model-specific adapters, such as DeepSeekV4's, live in the model directory. + +## 3. Execution: the caller enforces shared contracts + +The model entry point uses one scope for preparation and cleanup: + +```python +with model.attention.session(infer_state): + output = model._token_forward(infer_state) +``` + +Ordinary execution follows this order. During decode capture, device metadata construction and computation become part of the graph: + +```text +ModelAttention.session + -> Check storage generation; build batch and capacity + -> AttentionSession.prepare + -> Validate batch; acquire workspace leases + -> backend.allocate_context + -> backend.plan_host # Host planning outside the graph + -> Model forward + -> session.build_device_metadata # Device metadata + -> LayerAttention.forward per layer + -> Session checks context ownership + -> backend.forward # Phase computation + -> Session handles out + -> Exit scope; finish leases +``` + +The backend execution interface consists of `allocate_context`, `plan_host`, `build_device_metadata`, and `forward`. Resource budgets and capture hooks are supplied as needed. Each instance implements either prefill or decode. Constructors and helpers specific to a phase live in that phase's class; shared classes retain only common definitions. + +The session checks that a context belongs to the current backend and phase, and that the backend returns exactly the requested capacity. `batch_size` and `token_num` must match the capacity contract; sequence lengths must stay within its bounds. The model execution path handles padding and graph bucket selection. + +When `out` is supplied, the session checks dtype, device, and the result's storage range. A backend may write directly into `out` or return a valid view of it. If it returns an independent result, the session copies that result into `out` while preserving auxiliary outputs. + +Workspaces are isolated by `(group_id, microbatch_index)`. If preparation, planning, or forward raises, `ExitStack` still finishes every acquired lease. Finishing a lease records a completion event for subsequent reuse; it does not free pooled tensors or imply that the GPU has finished executing. + +## 4. CUDA Graph: stable addresses and validation before writes + +During decode capture, all groups' contexts share one captured batch. Sequence capacity explicitly uses `graph_max_len_in_batch`, while the host planner receives the current actual length separately. This avoids inflating capture capacity to the serving maximum. Captured workspaces retain their tensor references and are frozen; later pool growth does not replace addresses used by existing graphs. + +Each `CapturedGraph` replay proceeds as follows: + +1. Check storage generations and every microbatch lane's model role, backend groups, capacity, and input tensor shapes, dtypes, and devices. Build copy plans without writing inputs. +2. After all lanes pass validation, wait for the provider to finish using the previous metadata and acquire workspace leases. +3. Copy new inputs into captured addresses, call each group's `plan_host`, and submit the graph and provider auxiliary work. +4. Include auxiliary completion events when finishing leases, so subsequent reuse waits on all required work. + +A validation failure in a later lane therefore leaves earlier lanes' inputs unchanged. This guarantee covers validation before writes; it does not provide transactional rollback after a copy or provider execution fails. + +The request mapping table is shared storage and must retain its address during replay. A State storage generation change requires layer rebinding and graph recapture. Graph destruction waits for relevant device work, closes provider hooks, and then releases captured references. + +Prefill follows a separate path: `PrefillGraphExecution` holds the segmented graph's inputs, outputs, and collector. Attention runs through callbacks between segments using the current batch's session. `build_device_metadata` runs before the segments, outside the decode capture flow described above. Validation of all lanes before writing also applies specifically to the `CapturedGraph` path. + +## 5. Simplification decisions + +| Previous complexity | Current design | +| --- | --- | +| Two base classes, `AttentionBackend` / `BaseAttentionBackend`, and constructor forwarding | One base class; remove `create()` and `configure_spec()` | +| Duplicate copies of spec, causal mode, dynamic verification, and other settings | Read from the backend's `spec`, avoiding copies that must stay synchronized | +| A single-method BatchBuilder class and the indirect `ContextSource` interface | Construct the batch on the model side and pass explicit data to the session | +| Separate prepare / forward / finish calls that callers must pair | One session scope handles normal and exceptional cleanup | +| Prefill and decode helpers and dependencies mixed in shared classes | Move them to their phase implementations; for example, GDN decode construction does not load a prefill kernel | +| GDN inheriting Dense and KDA inheriting GDN | Organize adapters around the Conv/SSM layout they actually share | + +Family adapters remain because model fields and State resources differ. Sessions remain because all backends need the same lifecycle and validation rules. States already expose layouts and resource views, so no parallel storage-description interface is added. + +## 6. Extension points and verification + +For a new model, describe its layer requirements and reuse a family adapter where possible. Add a model-specific adapter only when resource binding differs. For a new provider, implement the relevant phase backend, declare its capabilities, and register it. For a new storage encoding, extend the State layout and resources and use `CacheLayout` for compatibility checks. See [Adding New Models](../models/add_new_model.md) for integration details. + +Source entry points: + +| Area | Files | +| --- | --- | +| Grouping, selection, and model adaptation | `lightllm/common/basemodel/attention.py`, `lightllm/common/basemodel/attention_adapters.py` | +| Static requirements and backend interface | `lightllm/common/attention/api/spec.py`, `lightllm/common/attention/api/backend.py` | +| Batch execution and workspaces | `lightllm/common/attention/session.py`, `lightllm/common/attention/api/workspace.py` | +| Graph replay | `lightllm/common/basemodel/graph.py`, `lightllm/common/basemodel/prefill_cuda_graph.py`, `lightllm/common/basemodel/infer_struct.py` | + +Regression coverage focuses on rejecting incorrect context ownership and capacity, lease cleanup on failure, output handling, validation of all lanes before writes, captured addresses, and storage generations. Interface contract and workspace tests live in `unit_tests/common/attention/`; model binding and backend numerical tests live in `unit_tests/common/basemodel/attention/`. Model validation must also cover prefill/decode, graphs, MTP, and state restoration. Throughput and cross-host PD claims require their own runtime validation. diff --git a/docs/EN/source/index.rst b/docs/EN/source/index.rst index 2b59db60b1..bbab22c857 100755 --- a/docs/EN/source/index.rst +++ b/docs/EN/source/index.rst @@ -80,6 +80,7 @@ Documentation List :caption: Architecture Introduction Architecture Overview + Attention Execution Design Token Attention Efficient Router From 674f11ebb09c56634a5f046641f76861bbb32cc5 Mon Sep 17 00:00:00 2001 From: sufubao Date: Sat, 14 Nov 2026 16:35:43 +0800 Subject: [PATCH 9/9] fix(attention): enforce execution scope and model boundaries --- .../framework/attention_execution_design.md | 27 ++- docs/CN/source/models/add_new_model.md | 4 +- .../framework/attention_execution_design.md | 27 ++- docs/EN/source/models/add_new_model.md | 4 +- lightllm/common/attention/api/backend.py | 1 - lightllm/common/attention/api/batch.py | 2 + lightllm/common/attention/api/environment.py | 7 +- lightllm/common/attention/api/inputs.py | 46 +---- lightllm/common/attention/api/spec.py | 27 +++ .../attention/backends/_shared/linear/gdn.py | 36 ++-- .../backends/_shared/linear/layout.py | 37 ++-- .../attention/backends/gdn/replay/backend.py | 4 +- .../attention/backends/gdn/replay/decode.py | 6 +- .../attention/backends/gdn/triton/backend.py | 2 +- .../attention/backends/kda/replay/backend.py | 13 +- .../attention/backends/kda/triton/backend.py | 11 +- .../backends/mla/flashinfer/backend.py | 8 +- .../backends/sparse_mla/flashmla/backend.py | 37 ---- .../backends/sparse_mla/flashmla/fp8.py | 7 - lightllm/common/attention/probe.py | 10 +- lightllm/common/attention/session.py | 75 ++++---- .../bounded_att/conv_ssm_types.py | 32 ++++ .../bounded_att/gdn/gdn_state.py | 2 +- .../bounded_att/gdn/replay_gdn_state.py | 21 +-- .../bounded_att/kda/kda_state.py | 2 +- .../bounded_att/kda/replay_kda_state.py | 21 +-- .../attention_state/full_att/gqa/gqa_state.py | 4 +- .../full_att/gqa/group_quantized_state.py | 4 +- .../full_att/gqa/static_fp8_state.py | 4 +- lightllm/common/attention_state/types.py | 9 + lightllm/common/basemodel/attention.py | 17 +- .../common/basemodel/attention_adapters.py | 43 ++++- lightllm/common/basemodel/basemodel.py | 3 + lightllm/common/basemodel/infer_struct.py | 37 ++-- .../models/deepseek_v4/attention/adapter.py | 7 + .../deepseek_v4/attention/sparse_mla.py | 9 +- lightllm/models/deepseek_v4/infer_struct.py | 17 ++ lightllm/models/deepseek_v4/model.py | 4 +- lightllm/models/gemma3/infer_struct.py | 9 + lightllm/models/gemma4/infer_struct.py | 7 + .../glm5_next/attention/fp8_sparse_mla.py | 4 +- .../models/glm5_next/attention/sparse_mla.py | 4 +- lightllm/models/glm5_next/indexer.py | 33 ++-- lightllm/models/glm5_next/infer_struct.py | 30 +++ .../layer_infer/transformer_layer_infer.py | 4 +- lightllm/models/glm5_next/model.py | 2 + .../layer_infer/transformer_layer_infer.py | 4 +- lightllm/models/glm5_next_mtp/model.py | 4 + lightllm/models/llama/infer_struct.py | 5 + lightllm/models/neo_chat_moe/infer_struct.py | 7 + lightllm/models/qwen/infer_struct.py | 2 + unit_tests/common/attention/helpers.py | 70 +++++-- unit_tests/common/attention/test_contracts.py | 177 +++++++++++++++--- .../common/attention/test_dense_numerics.py | 52 ++--- .../common/attention/test_linear_layout.py | 26 +-- .../common/attention/test_platform_runtime.py | 9 +- .../common/attention/test_provider_kind.py | 13 +- unit_tests/common/attention/test_state.py | 95 ++++++++++ .../basemodel/attention/test_dsv4_flashmla.py | 22 ++- .../attention/test_fp8_logical_layer.py | 18 +- .../common/test_deepseek4_paged_cache.py | 8 +- unit_tests/common/test_linear_state.py | 2 +- unit_tests/models/deepseek_v4/test_model.py | 35 ++++ .../glm5_next/test_direct_nsa_indices.py | 37 ++-- unit_tests/models/glm5_next/test_fp8_kv.py | 14 +- unit_tests/models/glm5_next/test_layers.py | 10 +- unit_tests/models/glm5_next/test_mtp.py | 14 +- .../glm5_next/test_mtp_prefill_cache.py | 18 +- .../glm5_next/test_prefill_workspace.py | 10 +- .../glm5_next/test_token_major_mla_bmm.py | 5 +- unit_tests/models/glm5_next/test_visual.py | 12 +- .../models/sensenova_loop/test_gdn_state.py | 18 +- 72 files changed, 936 insertions(+), 471 deletions(-) create mode 100644 lightllm/models/glm5_next/infer_struct.py diff --git a/docs/CN/source/framework/attention_execution_design.md b/docs/CN/source/framework/attention_execution_design.md index e9eee53a42..84930abb20 100644 --- a/docs/CN/source/framework/attention_execution_design.md +++ b/docs/CN/source/framework/attention_execution_design.md @@ -17,11 +17,11 @@ 三个数据类型各表达一种事实： -- `LayerSpec`：静态需求，包括算法、阶段、head 维度、特性及 `CacheLayout`。只由 backend 持有，层绑定不保存副本。 +- `LayerSpec`：静态需求，包括算法、阶段、head 维度、特性、`CacheLayout` 和不可变的 family 配置。只由 backend 持有，层绑定不保存副本。 - `AttentionBatch`：本次请求索引、长度、位置等输入；Tensor 借用自 `InferStateInfo` / `ReqManager`。 - `AttentionContext`：backend 为本次执行或某个 Graph 捕获实例准备的 metadata、容量和 workspace 引用。同一 Session 内的 context 共享一份 batch。 -`InferStateInfo` 单向持有 Session。Session 只接收 batch、分组、容量和 workspace 提供函数，不反向引用模型或 `InferStateInfo`。Backend 可以通过资源视图更新 KV/Conv/SSM，但不接管存储分配、checkpoint 或请求回收。 +`InferStateInfo` 单向持有 Session。Session 只接收 batch、分组、容量和 workspace 提供函数，不反向引用模型或 `InferStateInfo`。State 返回存储层定义的视图，固定字段绑定但允许修改 Tensor 内容，adapter 将其与权重组成计算资源。ReplaySSM 只暴露 Tensor 和配置，不把可执行生命周期操作的 State 对象交给 backend。Backend 可以通过资源视图更新 KV/Conv/SSM，但不接管存储分配、checkpoint 或请求回收。 ## 2. 初始化：描述、选择、绑定 @@ -43,7 +43,7 @@ 只有完整的 `(prefill_spec, decode_spec)` 相同，层才能共享 backend 和批次 context；权重、逻辑层号和存储视图仍留在各自的 `LayerAttention`。同一层的两个阶段必须使用相同的物理 `CacheLayout`，避免 prefill 写入的内容无法被 decode 读取。 -Backend 构造统一为 `Backend(env, spec)`。`BackendEnv` 显式传入设备 runtime、allocator、TP 和 MTP 等运行信息；量化组长等布局参数读取 `spec.cache`。能力声明放在 backend 类上，注册表只枚举实现和默认优先级。选择时检查环境、计算需求和服务约束，并执行对应探针；显式指定失败时报错，`auto` 才尝试其他候选。 +Backend 构造统一为 `Backend(env, spec)`。`BackendEnv` 显式传入设备 runtime、allocator、TP 和 MTP 等运行信息，不传模型配置字典；量化组长等布局参数读取 `spec.cache`。GDN/KDA 的卷积和 gate 参数、MLA 的分解维度以及稀疏索引参数由 adapter 翻译成不可变配置，进入 `LayerSpec`，共同参与分组。能力声明放在 backend 类上，注册表只枚举实现和默认优先级。选择时检查环境、计算需求和服务约束，并执行对应探针；显式指定失败时报错，`auto` 才尝试其他候选。 每层选择一次 family adapter，描述和绑定复用它。Dense 与 MLA 是独立适配；GDN 与 KDA 只共享 Conv/SSM 布局描述。DeepSeekV4 等模型专有适配放在模型目录。 @@ -64,31 +64,32 @@ ModelAttention.session → AttentionSession.prepare → 校验 batch，取得 workspace lease → backend.allocate_context - → backend.plan_host # 图外主机规划 + → backend.plan_host # 提交 Graph 前的规划，可启动设备操作 → 模型 forward → session.build_device_metadata # 设备 metadata → 逐层 LayerAttention.forward → Session 检查 context 归属 → backend.forward # 阶段计算 → Session 处理 out 回写 + → model.attention.finish_decode # 所有层成功执行后统一推进 ReplaySSM 游标 → 退出作用域，结束 lease ``` Backend 的执行接口只有 `allocate_context`、`plan_host`、`build_device_metadata` 和 `forward`；资源预算与捕获 hooks 按需提供。每个实例只实现 prefill 或 decode，阶段专属构造和 helper 放在对应类中，共享类只保留共同定义。 -Session 保证 context 属于当前 backend 和 phase，backend 返回的容量必须等于请求容量。`batch_size`、`token_num` 必须与容量契约一致，序列长度不能超过上界；padding 和 Graph bucket 选择由模型执行路径负责。 +层执行及 metadata 构建必须处于活跃的 Session 作用域，退出后拒绝继续执行；捕获 context 可以保留供 Graph 重放使用。数值探针也使用同一作用域入口。Session 保证 context 属于当前 backend 和 phase，backend 必须保留请求的容量、workspace 和 capture 模式。`batch_size`、`token_num` 必须与容量契约一致，声明的序列最大长度不能超过上界；padding 和 Graph bucket 选择由模型执行路径负责。 -提供 `out` 时，Session 统一校验 dtype、device 和结果范围：backend 可直接写入 `out` 或返回它的合法视图；返回独立结果时由 Session 复制，并保留辅助输出。 +提供 `out` 时，Session 统一校验 dtype、device 和结果范围：backend 可直接写入 `out`；共享存储的结果必须落在连续 `out` 内，或是 strided `out` 的同起点、同 stride 前缀切片，不能使用行间 padding；返回独立结果时由 Session 复制，并保留辅助输出。 -Workspace 按 `(group_id, microbatch_index)` 隔离。准备、规划或 forward 抛出异常时，`ExitStack` 仍结束已取得的 lease。结束 lease 会记录完成事件供后续复用等待，不释放池中的 Tensor，也不表示 GPU 已同步完成。 +Workspace 按 `(group_id, microbatch_index)` 隔离。同一槽位要求串行执行，lease 负责事件依赖，不是并发互斥锁。准备、规划或 forward 抛出异常时，`ExitStack` 仍结束已取得的 lease。结束 lease 会记录完成事件供后续复用等待，不释放池中的 Tensor，也不表示 GPU 已同步完成。 ## 4. CUDA Graph：固定地址，先校验再写入 -decode 捕获时，所有组的 context 使用同一份捕获 batch。长度容量显式采用 `graph_max_len_in_batch`，当前实际长度另行传给 host planner，避免用服务最大长度扩大捕获容量。捕获的 workspace 保留原 Tensor 引用并冻结；池后续扩容不替换旧 Graph 使用的地址。 +decode 捕获时，所有组的 context 使用同一份捕获 batch。长度容量显式采用 `graph_max_len_in_batch`，当前实际长度另行传给 host planner，避免用服务最大长度扩大捕获容量。捕获的 workspace 保留原 Tensor 引用并冻结其缓冲区绑定；provider 共享对象的内部地址稳定性仍由 provider 负责；池后续扩容不替换旧 Graph 使用的地址。 `CapturedGraph` 一次重放的顺序为： -1. 校验存储代次，以及所有 microbatch lane 的模型角色、backend 分组、容量和输入 Tensor 形状、dtype、device；只生成复制计划。 +1. 校验存储代次，以及所有 microbatch lane 的模型角色、backend 分组、容量和输入 Tensor 的存在性、形状、dtype、device，以及同地址视图的 stride；只生成复制计划。 2. 所有 lane 校验通过后，等待 provider 上次 metadata 使用结束，并取得 workspace lease。 3. 将新输入复制到捕获地址，调用各组 `plan_host` 更新主机规划，再提交 Graph 和 provider 辅助任务。 4. 将辅助任务完成事件纳入 lease 收尾，保证下一次复用的依赖完整。 @@ -99,6 +100,14 @@ decode 捕获时，所有组的 context 使用同一份捕获 batch。长度容 prefill 采用另一条路径：`PrefillGraphExecution` 持有分段 Graph 的输入、输出和 collector，attention 通过段间回调使用当前批次的 Session。`build_device_metadata` 在分段执行前运行，不属于上述 decode 捕获流程；多 lane 写入前校验的保证也仅针对 `CapturedGraph` 路径。 +各模型的 InferState 类通过 `graph_input_names` 声明 `AttentionBatch` 之外的 RoPE、draft hidden 等重放输入。新增输入时必须扩展此清单；DeepSeek-V4 C4 索引等图内产生的中间结果不列入清单，由捕获的节点重新计算。 + +GLM 的请求展开信息由模型推理状态生成，以 `SparseQueryMetadata` 传入 indexer；它不再读取 provider context 的内部字段。decode 捕获输入复制包含这些 Tensor，prefill 回调使用当前批次的数据。 + +ReplaySSM 的 cursor 由模型在参与层全部成功返回后统一推进一次，普通和 overlap decode 使用同一操作，Graph 捕获包含该操作。层执行失败时不会执行这一步；此前已写入的状态不提供事务回滚。 + +DeepSeek3.2 仍从 sparse provider context 读取 ragged-index metadata；替换该 provider 时，需要先明确这部分家族数据契约。当前 provider 没有实现辅助 capture hooks，现有假 provider 测试验证调用顺序，不能证明真实 provider 提交辅助任务后的故障恢复。 + ## 5. 本次精简的取舍 | 原有复杂度 | 当前处理 | diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index b07e43ff4c..1d2f00d561 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -133,9 +133,9 @@ class TpPartBaseModel: 注册表枚举只加载后端定义，能力声明以类属性为准；同一后端的阶段共性放在共享 mixin。kernel 和 provider 的执行依赖在使用处导入，候选后端的环境检查才验证可选依赖。Neo 扩展缺失时，`auto` 可回退，显式选择则报错，不影响其他模型的后端枚举。 -后端构造只接收 frozen `BackendEnv` 和 `LayerSpec`。Env 显式提供设备 runtime、网络配置、TP 信息、HOLD 请求号、decode draft step、持久状态槽宽、长度上界和 allocator。`infer_state.attention: AttentionSession` 管理批次 context、owner/phase 校验、输入复制及输出写回。后端仅提供 `allocate_context`、图外每次重放执行的 `plan_host`、图内的 `build_device_metadata` 和 `forward`。模型通过 `with model.attention.session(infer_state)` 一次完成准备及 lease 释放；decode 捕获显式传入 Graph 长度上界。Session 位于 attention 模块，只接收 batch、backend 分组和容量，不依赖模型。所有 context 共享一个捕获 batch；Graph 先校验所有 microbatch 的复制计划，再统一复制和调用 host plan。 +后端构造只接收 frozen `BackendEnv` 和 `LayerSpec`。Env 显式提供设备 runtime、TP 信息、HOLD 请求号、decode draft step、持久状态槽宽、长度上界和 allocator。`infer_state.attention: AttentionSession` 管理批次 context、owner/phase 校验、输入复制及输出写回。后端仅提供 `allocate_context`、图外每次重放执行的 `plan_host`、图内的 `build_device_metadata` 和 `forward`。模型通过 `with model.attention.session(infer_state)` 一次完成准备及 lease 释放；decode 捕获显式传入 Graph 长度上界。Session 位于 attention 模块，只接收 batch、backend 分组和容量，不依赖模型。所有 context 共享一个捕获 batch；Graph 先校验所有 microbatch 的复制计划，再统一复制和调用 host plan。 -`AttentionBackend(env, spec)` 是唯一后端基类，`spec` 只存于后端。`LayerAttention` 只持有 backend、资源和路由；prefill/decode 专属构造及 helper 各归其阶段类。量化组长来自 `spec.cache`，allocator 直接注入后端。层的 `attention_adapter(family)` 选取一次 adapter，描述与绑定复用该实例；模型专有 adapter 放在模型目录。 +`AttentionBackend(env, spec)` 是唯一后端基类，`spec` 只存于后端。`LayerAttention` 只持有 backend、资源和路由；prefill/decode 专属构造及 helper 各归其阶段类。量化组长来自 `spec.cache`，family 参数由 adapter 翻译为不可变的 `spec.config`，allocator 直接注入后端。层的 `attention_adapter(family)` 选取一次 adapter，描述与绑定复用该实例；模型专有 adapter 放在模型目录。 ```python output = self.decode_attention.forward(infer_state.attention, DenseInputs(q)) diff --git a/docs/EN/source/framework/attention_execution_design.md b/docs/EN/source/framework/attention_execution_design.md index f663153e99..4c03f7cbbe 100644 --- a/docs/EN/source/framework/attention_execution_design.md +++ b/docs/EN/source/framework/attention_execution_design.md @@ -17,11 +17,11 @@ The aim is to keep changes for a new model or backend within its implementation, Three data types describe distinct facts: -- `LayerSpec`: static requirements, including algorithm, phase, head dimensions, features, and `CacheLayout`. The backend holds it; layer bindings do not keep a second copy. +- `LayerSpec`: static requirements, including algorithm, phase, head dimensions, features, `CacheLayout`, and immutable family configuration. The backend holds it; layer bindings do not keep a second copy. - `AttentionBatch`: request indices, lengths, positions, and other inputs for this execution. Tensors are borrowed from `InferStateInfo` / `ReqManager`. - `AttentionContext`: metadata, capacity, and workspace references prepared by the backend for an execution or a graph capture instance. Contexts within a session share one batch. -`InferStateInfo` holds the session in one direction. The session receives only the batch, groups, capacity, and a workspace supplier; it does not reference the model or `InferStateInfo`. Backends can update KV/Conv/SSM through resource views, while storage allocation, checkpointing, and request reclamation remain with their existing owners. +`InferStateInfo` holds the session in one direction. The session receives only the batch, groups, capacity, and a workspace supplier; it does not reference the model or `InferStateInfo`. States return views defined by the storage module; adapters combine these views with weights into computation resources. ReplaySSM exposes tensors and configuration without giving the backend a State object with lifecycle operations. Backends can update KV/Conv/SSM through these views, while storage allocation, checkpointing, and request reclamation remain with their existing owners. ## 2. Initialization: describe, select, bind @@ -43,7 +43,7 @@ Layer describe_attention(phase) + State layout Layers share backends and batch contexts only when their complete `(prefill_spec, decode_spec)` pairs match. Weights, logical layer IDs, and storage views remain in each layer's `LayerAttention`. Both phases of a layer must use the same physical `CacheLayout`, so decode can read what prefill writes. -All backends are constructed as `Backend(env, spec)`. `BackendEnv` explicitly supplies the device runtime, allocator, TP and MTP settings, and other runtime information. Layout parameters such as quantization group size come from `spec.cache`. Backend classes declare capabilities; the registry lists implementations and default priorities. Selection checks the environment, computation requirements, and serving constraints, then runs the corresponding probes. An explicit choice fails with an error when unusable; only `auto` tries other candidates. +All backends are constructed as `Backend(env, spec)`. `BackendEnv` explicitly supplies the device runtime, allocator, TP and MTP settings, and other runtime information, without a model configuration dictionary. Adapters translate GDN/KDA convolution and gate settings, MLA decomposition dimensions, and sparse index parameters into immutable configuration in `LayerSpec`, which participates in grouping. Layout parameters such as quantization group size come from `spec.cache`. Backend classes declare capabilities; the registry lists implementations and default priorities. Selection checks the environment, computation requirements, and serving constraints, then runs the corresponding probes. An explicit choice fails with an error when unusable; only `auto` tries other candidates. Each layer selects its family adapter once and reuses it for description and binding. Dense and MLA have separate adapters. GDN and KDA share only the Conv/SSM layout description. Model-specific adapters, such as DeepSeekV4's, live in the model directory. @@ -64,41 +64,50 @@ ModelAttention.session -> AttentionSession.prepare -> Validate batch; acquire workspace leases -> backend.allocate_context - -> backend.plan_host # Host planning outside the graph + -> backend.plan_host # Planning before graph submission; may launch device work -> Model forward -> session.build_device_metadata # Device metadata -> LayerAttention.forward per layer -> Session checks context ownership -> backend.forward # Phase computation -> Session handles out + -> model.attention.finish_decode # Advance ReplaySSM cursors after all layers succeed -> Exit scope; finish leases ``` The backend execution interface consists of `allocate_context`, `plan_host`, `build_device_metadata`, and `forward`. Resource budgets and capture hooks are supplied as needed. Each instance implements either prefill or decode. Constructors and helpers specific to a phase live in that phase's class; shared classes retain only common definitions. -The session checks that a context belongs to the current backend and phase, and that the backend returns exactly the requested capacity. `batch_size` and `token_num` must match the capacity contract; sequence lengths must stay within its bounds. The model execution path handles padding and graph bucket selection. +Layer execution and metadata construction require an active session scope and are rejected after it exits. Captured contexts can remain available for graph replay. Numerical probes use the same scoped entry point. The session checks context ownership and phase, and requires the backend to retain the requested capacity, workspace, and capture mode. `batch_size` and `token_num` must match the capacity contract; declared maximum sequence lengths must stay within its bounds. The model execution path handles padding and graph bucket selection. -When `out` is supplied, the session checks dtype, device, and the result's storage range. A backend may write directly into `out` or return a valid view of it. If it returns an independent result, the session copies that result into `out` while preserving auxiliary outputs. +When `out` is supplied, the session checks dtype, device, and the result's storage range. A backend may write directly into `out`. A result sharing storage must lie within a contiguous `out`, or be a prefix slice of a strided `out` with the same starting offset and strides; padding between rows is not part of the destination. If it returns an independent result, the session copies that result into `out` while preserving auxiliary outputs. -Workspaces are isolated by `(group_id, microbatch_index)`. If preparation, planning, or forward raises, `ExitStack` still finishes every acquired lease. Finishing a lease records a completion event for subsequent reuse; it does not free pooled tensors or imply that the GPU has finished executing. +Workspaces are isolated by `(group_id, microbatch_index)`. Executions within a slot must be serialized; leases order completion events but do not provide a concurrency lock. If preparation, planning, or forward raises, `ExitStack` still finishes every acquired lease. Finishing a lease records a completion event for subsequent reuse; it does not free pooled tensors or imply that the GPU has finished executing. ## 4. CUDA Graph: stable addresses and validation before writes -During decode capture, all groups' contexts share one captured batch. Sequence capacity explicitly uses `graph_max_len_in_batch`, while the host planner receives the current actual length separately. This avoids inflating capture capacity to the serving maximum. Captured workspaces retain their tensor references and are frozen; later pool growth does not replace addresses used by existing graphs. +During decode capture, all groups' contexts share one captured batch. Sequence capacity explicitly uses `graph_max_len_in_batch`, while the host planner receives the current actual length separately. This avoids inflating capture capacity to the serving maximum. Captured workspaces retain their tensor references and freeze their buffer bindings; providers remain responsible for internal address stability in shared objects; later pool growth does not replace addresses used by existing graphs. Each `CapturedGraph` replay proceeds as follows: -1. Check storage generations and every microbatch lane's model role, backend groups, capacity, and input tensor shapes, dtypes, and devices. Build copy plans without writing inputs. +1. Check storage generations and every microbatch lane's model role, backend groups, capacity, and input tensor presence, shapes, dtypes, devices, and strides for views at the same address. Build copy plans without writing inputs. 2. After all lanes pass validation, wait for the provider to finish using the previous metadata and acquire workspace leases. 3. Copy new inputs into captured addresses, call each group's `plan_host`, and submit the graph and provider auxiliary work. 4. Include auxiliary completion events when finishing leases, so subsequent reuse waits on all required work. +The hook ordering is covered by fake-provider tests; current providers do not implement auxiliary capture hooks. Failure recovery for a provider that submits auxiliary work needs provider-specific validation. + A validation failure in a later lane therefore leaves earlier lanes' inputs unchanged. This guarantee covers validation before writes; it does not provide transactional rollback after a copy or provider execution fails. +Each inference-state class declares `graph_input_names` for tensors outside `AttentionBatch`, including RoPE data and draft hidden inputs. Model integrations must extend this inventory for new replay inputs. Tensors produced inside the graph, such as DeepSeek-V4 C4 indices, are excluded and recomputed by captured nodes. + The request mapping table is shared storage and must retain its address during replay. A State storage generation change requires layer rebinding and graph recapture. Graph destruction waits for relevant device work, closes provider hooks, and then releases captured references. Prefill follows a separate path: `PrefillGraphExecution` holds the segmented graph's inputs, outputs, and collector. Attention runs through callbacks between segments using the current batch's session. `build_device_metadata` runs before the segments, outside the decode capture flow described above. Validation of all lanes before writing also applies specifically to the `CapturedGraph` path. +GLM generates request-expansion metadata in its model inference state and passes `SparseQueryMetadata` to the indexer, without reading provider context internals. Decode input copying includes these tensors; prefill callbacks use the current batch's data. DeepSeek3.2 still reads ragged-index metadata from its sparse provider context; replacing that provider requires preserving this integration until it receives an explicit family contract. + +The model advances each ReplaySSM cursor once after all participating layers return successfully. Ordinary and overlap decode use the same operation, which is included in graph capture. A layer failure prevents this step; earlier state writes are not transactionally rolled back. + ## 5. Simplification decisions | Previous complexity | Current design | diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index 253a4b3660..a5eeb72218 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -133,9 +133,9 @@ Layers return family, features, window and variant from `describe_attention(phas Registry enumeration loads backend definitions only. Class attributes declare capabilities, with shared phase capabilities in a common mixin. Import kernel and provider execution dependencies where they are used; candidate environment checks validate optional dependencies. A missing Neo extension allows fallback in `auto` mode and fails explicit selection, without affecting backend enumeration for other models. -Backends receive only a frozen `BackendEnv` and `LayerSpec`. The environment supplies the device runtime, network configuration, TP coordinates, HOLD request, decode draft step, persistent state slot count, sequence capacity and allocator. `infer_state.attention: AttentionSession` owns batch contexts, owner/phase validation, input copies and output copy-back. Backend hooks are `allocate_context`, `plan_host` (outside the graph on every replay), `build_device_metadata` (captured device work), and `forward`. `with model.attention.session(infer_state)` combines preparation and lease release; decode capture passes its Graph length bound explicitly. The Session lives in the attention module and accepts batches, backend groups and capacities without depending on the model. Contexts share one captured batch. Graph replay validates copy plans for every microbatch before copying inputs and running host plans. +Backends receive only a frozen `BackendEnv` and `LayerSpec`. The environment supplies the device runtime, TP coordinates, HOLD request, decode draft step, persistent state slot count, sequence capacity and allocator. `infer_state.attention: AttentionSession` owns batch contexts, owner/phase validation, input copies and output copy-back. Backend hooks are `allocate_context`, `plan_host` (outside the graph on every replay), `build_device_metadata` (captured device work), and `forward`. `with model.attention.session(infer_state)` combines preparation and lease release; decode capture passes its Graph length bound explicitly. The Session lives in the attention module and accepts batches, backend groups and capacities without depending on the model. Contexts share one captured batch. Graph replay validates copy plans for every microbatch before copying inputs and running host plans. -`AttentionBackend(env, spec)` is the only backend base class and owns the spec. `LayerAttention` holds the backend, resources and route. Phase-specific constructors and helpers live on their respective phase classes. Quantization group size comes from `spec.cache`; the allocator is injected directly into the backend. The layer selects `attention_adapter(family)` once and reuses it for description and binding. Model-specific adapters live in the model directory. +`AttentionBackend(env, spec)` is the only backend base class and owns the spec. `LayerAttention` holds the backend, resources and route. Phase-specific constructors and helpers live on their respective phase classes. Quantization group size comes from `spec.cache`; adapters translate family parameters into immutable `spec.config`. The allocator is injected directly into the backend. The layer selects `attention_adapter(family)` once and reuses it for description and binding. Model-specific adapters live in the model directory. ```python output = self.decode_attention.forward(infer_state.attention, DenseInputs(q)) diff --git a/lightllm/common/attention/api/backend.py b/lightllm/common/attention/api/backend.py index f86bdd0ad3..d24e910f30 100644 --- a/lightllm/common/attention/api/backend.py +++ b/lightllm/common/attention/api/backend.py @@ -75,7 +75,6 @@ class AttentionBackend(ABC, Generic[Context, Inputs, Output]): def __init__(self, env: BackendEnv, spec: LayerSpec): self.runtime = env.runtime self.spec = spec - self.network_config = env.network_config self.draft_step = env.draft_step if spec.use_mtp else 0 self.hold_req_id = env.hold_req_id self.alloc_tensor = env.alloc_tensor diff --git a/lightllm/common/attention/api/batch.py b/lightllm/common/attention/api/batch.py index d570a50043..65988edb06 100644 --- a/lightllm/common/attention/api/batch.py +++ b/lightllm/common/attention/api/batch.py @@ -89,6 +89,8 @@ def input_copies(self, source): raise ValueError(f"captured input {item.name} changed presence") if (dst.shape, dst.dtype, dst.device) != (src.shape, src.dtype, src.device): raise ValueError(f"captured input {item.name} changed shape, dtype or device") + if dst.data_ptr() == src.data_ptr() and dst.stride() != src.stride(): + raise ValueError(f"captured input {item.name} changed strides at the captured address") if item.metadata.get("shared_storage") and dst.data_ptr() != src.data_ptr(): raise ValueError("request table storage changed; recapture attention") if dst.data_ptr() != src.data_ptr(): diff --git a/lightllm/common/attention/api/environment.py b/lightllm/common/attention/api/environment.py index ec9cdd887a..8a59c12caf 100644 --- a/lightllm/common/attention/api/environment.py +++ b/lightllm/common/attention/api/environment.py @@ -1,6 +1,5 @@ from dataclasses import dataclass -from types import MappingProxyType -from typing import Callable, Mapping +from typing import Callable from lightllm.platform import DeviceRuntime @@ -8,7 +7,6 @@ @dataclass(frozen=True, kw_only=True) class BackendEnv: runtime: DeviceRuntime - network_config: Mapping hold_req_id: int draft_step: int max_seq_len: int @@ -20,6 +18,3 @@ class BackendEnv: prefill_graph: bool = False max_prefill_tokens: int = 0 decode_only: bool = False - - def __post_init__(self): - object.__setattr__(self, "network_config", MappingProxyType(dict(self.network_config))) diff --git a/lightllm/common/attention/api/inputs.py b/lightllm/common/attention/api/inputs.py index c4cd732319..6e8e39ae02 100644 --- a/lightllm/common/attention/api/inputs.py +++ b/lightllm/common/attention/api/inputs.py @@ -5,11 +5,16 @@ 这些数据结构只固定字段绑定，不禁止 backend 原地更新所借用的状态 tensor。 """ +from __future__ import annotations + from dataclasses import dataclass -from typing import NamedTuple, Optional, Protocol +from typing import NamedTuple, Optional, TYPE_CHECKING import torch +if TYPE_CHECKING: + from lightllm.common.attention_state.bounded_att.conv_ssm_types import ReplayStorage + @dataclass(frozen=True) class DenseResources: @@ -29,41 +34,6 @@ class MLAResources: qk_rope_head_dim: int -class ReplaySSMState(Protocol): - """Borrowed replay storage; State retains allocation, checkpoint and release ownership.""" - - ssm_state: torch.Tensor - keys: torch.Tensor - deltas: torch.Tensor - gates: torch.Tensor - raw_keys: torch.Tensor - raw_values: torch.Tensor - betas: torch.Tensor - cursors: torch.Tensor - tile_sync_mark: torch.Tensor - num_key_heads: int - capacity: int - verify_width: int - compute_mode: str - gate_lower_bound: float - hold_req_id: int - kda: bool - - def storage_layer(self, logical_layer_id: int) -> int: - ... - - def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: - ... - - -@dataclass(frozen=True) -class RecurrentStorage: - conv: torch.Tensor - ssm: torch.Tensor - mtp_state_index: torch.Tensor - replay_state: ReplaySSMState | None = None - - @dataclass(frozen=True) class GDNResources: conv: torch.Tensor @@ -73,7 +43,7 @@ class GDNResources: a_log: torch.Tensor dt_bias: torch.Tensor mtp_state_index: torch.Tensor | None = None - replay_state: ReplaySSMState | None = None + replay_state: ReplayStorage | None = None @dataclass(frozen=True) @@ -84,7 +54,7 @@ class KDAResources: a_log: torch.Tensor dt_bias: torch.Tensor mtp_state_index: torch.Tensor | None = None - replay_state: ReplaySSMState | None = None + replay_state: ReplayStorage | None = None @dataclass(frozen=True) diff --git a/lightllm/common/attention/api/spec.py b/lightllm/common/attention/api/spec.py index 8d5a2f9f6b..253c487a66 100644 --- a/lightllm/common/attention/api/spec.py +++ b/lightllm/common/attention/api/spec.py @@ -74,6 +74,31 @@ def __post_init__(self): raise ValueError("cache page_size must be positive") +@dataclass(frozen=True) +class GDNConfig: + value_heads: int + conv_kernel_size: int + activation: str + + +@dataclass(frozen=True) +class KDAConfig: + conv_kernel_size: int + gate_lower_bound: float + + +@dataclass(frozen=True) +class MLAConfig: + rope_head_dim: int + value_head_dim: int + + +@dataclass(frozen=True) +class SparseMLAConfig: + index_topk: int = 2048 + index_kpool: int = 1 + + @dataclass(frozen=True) class LayerSpec: """一个层在一个 phase 下的静态计算需求，用于能力检查和分组。 @@ -93,8 +118,10 @@ class LayerSpec: features: frozenset[AttentionFeature] = frozenset() variant: str = "" softmax_scale: float | None = None + config: object = None def __post_init__(self): + hash(self.config) features = set(self.features) if self.sliding_window is not None: features.add(AttentionFeature.SLIDING_WINDOW) diff --git a/lightllm/common/attention/backends/_shared/linear/gdn.py b/lightllm/common/attention/backends/_shared/linear/gdn.py index cb5643c4ec..6c8203a1c7 100644 --- a/lightllm/common/attention/backends/_shared/linear/gdn.py +++ b/lightllm/common/attention/backends/_shared/linear/gdn.py @@ -4,26 +4,22 @@ class GDNMetadata: - def _init_linear_layer_metadata(self, network_config, tp_world_size): - self.num_v_heads = network_config["linear_num_value_heads"] - self.num_k_heads = network_config["linear_num_key_heads"] - self.head_k_dim = network_config["linear_key_head_dim"] - self.head_v_dim = network_config["linear_value_head_dim"] - self.key_dim = self.head_k_dim * self.num_k_heads - self.value_dim = self.head_v_dim * self.num_v_heads - self.conv_kernel_dim = network_config["linear_conv_kernel_dim"] - self.activation = network_config["hidden_act"] - self.tp_qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // tp_world_size - self.tp_ba_dim = self.num_v_heads * 2 // tp_world_size - self.tp_num_k_heads = self.num_k_heads // tp_world_size - self.tp_num_v_heads = self.num_v_heads // tp_world_size - self.tp_key_dim = self.key_dim // tp_world_size - self.tp_value_dim = self.value_dim // tp_world_size - if self.num_v_heads % self.num_k_heads: - raise ValueError("num_v_heads must be divisible by num_k_heads") - self.num_v_heads_per_k_head = self.num_v_heads // self.num_k_heads - self.ssm_state_dtype = self.spec.cache.state_dtype - return + def _init_linear_layer_metadata(self): + spec = self.spec + self.head_k_dim = spec.head_size + self.head_v_dim = spec.value_head_size + self.tp_num_k_heads = spec.cache.key_heads + self.tp_num_v_heads = spec.config.value_heads + self.tp_key_dim = self.head_k_dim * self.tp_num_k_heads + self.tp_value_dim = self.head_v_dim * self.tp_num_v_heads + self.tp_qkvz_dim = self.tp_key_dim * 2 + self.tp_value_dim * 2 + self.tp_ba_dim = self.tp_num_v_heads * 2 + self.conv_kernel_dim = spec.config.conv_kernel_size + self.activation = spec.config.activation + if self.tp_num_v_heads % self.tp_num_k_heads: + raise ValueError("value_heads must be divisible by key_heads") + self.num_v_heads_per_k_head = self.tp_num_v_heads // self.tp_num_k_heads + self.ssm_state_dtype = spec.cache.state_dtype def _split_qkvzba(self, mixed_qkvzba): qkv_dim = self.tp_key_dim * 2 + self.tp_value_dim diff --git a/lightllm/common/attention/backends/_shared/linear/layout.py b/lightllm/common/attention/backends/_shared/linear/layout.py index d578429779..1b6d3075ee 100644 --- a/lightllm/common/attention/backends/_shared/linear/layout.py +++ b/lightllm/common/attention/backends/_shared/linear/layout.py @@ -30,8 +30,6 @@ def accepted_tokens(self, state_index: torch.Tensor) -> torch.Tensor: def update_linear_decode(context: LinearDecodeContext, draft_step: int, dynamic: bool, *, replay: bool) -> None: - from .mtp_state_params import build_dynamic_mtp_layout - batch = context.batch context.accepted_by_state.clear() if draft_step == 0: @@ -42,26 +40,27 @@ def update_linear_decode(context: LinearDecodeContext, draft_step: int, dynamic: context.b_cu_seqlens = torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) return width = draft_step + 1 - if dynamic: - context.b1_mtp_cu_q_seq_len, context.b_conv_buffer_idx = build_dynamic_mtp_layout( - b_req_idx=batch.b_req_idx, - b_mtp_index=batch.b_mtp_index, - hold_req_id=batch.hold_req_id, - ) - else: - if batch.batch_size % width: - raise ValueError("linear decode batch size must be divisible by draft_step + 1") - context.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch.batch_size + 1, - width, - dtype=torch.int32, - device=batch.device, - ) - context.b_conv_buffer_idx = batch.b_req_idx.view(-1, width)[:, 0].contiguous() + context.b_conv_buffer_idx, context.b1_mtp_cu_q_seq_len = decode_request_layout(batch, draft_step, dynamic) if replay: context.b_ssm_buffer_idx = context.b_conv_buffer_idx context.b_cu_seqlens = context.b1_mtp_cu_q_seq_len else: offsets = torch.arange(width, dtype=batch.b_req_idx.dtype, device=batch.device) context.b_ssm_buffer_idx = context.b_conv_buffer_idx[:, None] * width + offsets + + +def decode_request_layout(batch, draft_step, dynamic): + from .mtp_state_params import build_dynamic_mtp_layout + + width = draft_step + 1 + if not draft_step: + return batch.b_req_idx, torch.arange(batch.batch_size + 1, dtype=torch.int32, device=batch.device) + if dynamic: + cu, reqs = build_dynamic_mtp_layout( + b_req_idx=batch.b_req_idx, b_mtp_index=batch.b_mtp_index, hold_req_id=batch.hold_req_id + ) + return reqs, cu + if batch.batch_size % width: + raise ValueError("linear decode batch size must be divisible by draft_step + 1") + cu = torch.arange(0, batch.batch_size + 1, width, dtype=torch.int32, device=batch.device) + return batch.b_req_idx.view(-1, width)[:, 0].contiguous(), cu diff --git a/lightllm/common/attention/backends/gdn/replay/backend.py b/lightllm/common/attention/backends/gdn/replay/backend.py index 4a1730ff2f..5a05b7fc59 100644 --- a/lightllm/common/attention/backends/gdn/replay/backend.py +++ b/lightllm/common/attention/backends/gdn/replay/backend.py @@ -36,7 +36,7 @@ class TritonGDNReplayCommon(GDNMetadata): def __init__(self, env, spec): super().__init__(env=env, spec=spec) - self._init_linear_layer_metadata(env.network_config, env.tp_size) + self._init_linear_layer_metadata() self.state_slots = env.state_slots @@ -203,8 +203,6 @@ def forward( core = self._mtp(layer, context, mixed_qkv, conv_states, a, b, draft_step, state) else: core, z = self._decode(layer, context, mixed_qkv, z, conv_states, a, b, state) - # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) return GDNOutput(core, z) diff --git a/lightllm/common/attention/backends/gdn/replay/decode.py b/lightllm/common/attention/backends/gdn/replay/decode.py index 6dda8fc867..e995db0ebf 100644 --- a/lightllm/common/attention/backends/gdn/replay/decode.py +++ b/lightllm/common/attention/backends/gdn/replay/decode.py @@ -55,7 +55,9 @@ def rebuild_inputs(state, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state.verify_width)) _, _, hv, kd, vd = state.ssm_state.shape - scratch = type(state)( + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + + scratch = ReplayGDNState( ConvSSMStateConfig( logical_layer_ids=(0,), global_k_heads=state.num_key_heads, @@ -86,7 +88,7 @@ def rebuild_inputs(state, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu cu = torch.arange(run_count + 1, dtype=torch.int32, device=q.device) cu = (cu * state.verify_width).clamp_max(tokens) return (), dict( - state=scratch, + state=scratch.replay_storage(), linear_layer_idx=0, q=copied[0], k=copied[1], diff --git a/lightllm/common/attention/backends/gdn/triton/backend.py b/lightllm/common/attention/backends/gdn/triton/backend.py index b86b4dbe38..6823bd01e1 100644 --- a/lightllm/common/attention/backends/gdn/triton/backend.py +++ b/lightllm/common/attention/backends/gdn/triton/backend.py @@ -32,7 +32,7 @@ class TritonGDNCommon(GDNMetadata): def __init__(self, env, spec): super().__init__(env=env, spec=spec) - self._init_linear_layer_metadata(network_config=env.network_config, tp_world_size=env.tp_size) + self._init_linear_layer_metadata() self.state_slots = env.state_slots diff --git a/lightllm/common/attention/backends/kda/replay/backend.py b/lightllm/common/attention/backends/kda/replay/backend.py index 97f7dc5023..97c1beb3db 100644 --- a/lightllm/common/attention/backends/kda/replay/backend.py +++ b/lightllm/common/attention/backends/kda/replay/backend.py @@ -35,14 +35,11 @@ class TritonKDAReplayCommon: def __init__(self, env, spec): super().__init__(env=env, spec=spec) - config = env.network_config["linear_attn_config"] - self.num_heads = config["num_heads"] - self.head_dim = config["head_dim"] - assert self.num_heads % env.tp_size == 0 - self.tp_num_heads = self.num_heads // env.tp_size + self.head_dim = spec.head_size + self.tp_num_heads = spec.num_q_heads self.tp_hidden_size = self.tp_num_heads * self.head_dim - self.conv_kernel_size = config["short_conv_kernel_size"] - self.lower_bound = config.get("gate_lower_bound", -5.0) + self.conv_kernel_size = spec.config.conv_kernel_size + self.lower_bound = spec.config.gate_lower_bound self.state_slots = env.state_slots def split_qkv(self, mixed_qkv: torch.Tensor): @@ -172,6 +169,4 @@ def forward( state_compute_mode=state.compute_mode, state_gate_lower_bound=state.gate_lower_bound, ) - # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 - state.finish_decode_token(layer.logical_layer_id, context.b_conv_buffer_idx, context.b_cu_seqlens) return output diff --git a/lightllm/common/attention/backends/kda/triton/backend.py b/lightllm/common/attention/backends/kda/triton/backend.py index 8b40d0b7f9..920f509bd2 100644 --- a/lightllm/common/attention/backends/kda/triton/backend.py +++ b/lightllm/common/attention/backends/kda/triton/backend.py @@ -35,14 +35,11 @@ class TritonKDACommon: def __init__(self, env, spec): super().__init__(env=env, spec=spec) - config = env.network_config["linear_attn_config"] - self.num_heads = config["num_heads"] - self.head_dim = config["head_dim"] - assert self.num_heads % env.tp_size == 0 - self.tp_num_heads = self.num_heads // env.tp_size + self.head_dim = spec.head_size + self.tp_num_heads = spec.num_q_heads self.tp_hidden_size = self.tp_num_heads * self.head_dim - self.conv_kernel_size = config["short_conv_kernel_size"] - self.lower_bound = config.get("gate_lower_bound", -5.0) + self.conv_kernel_size = spec.config.conv_kernel_size + self.lower_bound = spec.config.gate_lower_bound self.state_slots = env.state_slots diff --git a/lightllm/common/attention/backends/mla/flashinfer/backend.py b/lightllm/common/attention/backends/mla/flashinfer/backend.py index 7f728a5915..b17d4fd262 100644 --- a/lightllm/common/attention/backends/mla/flashinfer/backend.py +++ b/lightllm/common/attention/backends/mla/flashinfer/backend.py @@ -57,10 +57,10 @@ def __init__(self, env, spec): set_flashinfer_envs() super().__init__(env=env, spec=spec) self.tp_q_head_num = spec.num_q_heads - self.qk_nope_head_dim = env.network_config["qk_nope_head_dim"] - self.qk_rope_head_dim = env.network_config["qk_rope_head_dim"] - self.kv_lora_rank = env.network_config["kv_lora_rank"] - self.v_head_dim = env.network_config["v_head_dim"] + self.qk_nope_head_dim = spec.head_size - spec.config.rope_head_dim + self.qk_rope_head_dim = spec.config.rope_head_dim + self.kv_lora_rank = spec.value_head_size + self.v_head_dim = spec.config.value_head_dim self.q_data_type = spec.dtype self.kv_data_type = spec.cache.dtype self.max_seq_length = env.max_seq_len diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py index 11b9871a95..ae4ab37959 100644 --- a/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py @@ -16,11 +16,6 @@ class NsaFlashMlaSparsePrefillContext(AttentionContext): ke: torch.Tensor = None lengths: torch.Tensor = None ragged_mem_index: torch.Tensor = None - # Query-row metadata is consumed by env-specific sparse indexers. It is - # kept on the context so the indexer never reaches into InferStateInfo. - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None - q_kv_start_loc: torch.Tensor = None @dataclasses.dataclass(slots=True) @@ -31,8 +26,6 @@ class NsaFlashMlaSparseDecodeContext(AttentionContext): ragged_mem_index: torch.Tensor = None nsa_cache_seqlens: torch.Tensor = None nsa_cu_seqlens_k_new: torch.Tensor = None - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None class FlashMLASparseMLACommon: @@ -56,33 +49,6 @@ class FlashMLASparseMLAPrefillBackend( ): phase = AttentionPhase.PREFILL - @staticmethod - def _update_query_metadata(context): - """Materialize per-query request metadata on the execution context. - - Prefill positions are supplied by the env's input preparation and - therefore remain correct for prompt-cache and speculative layouts. - ``output_size`` keeps the operation graph-capturable without a host - read of the GPU-side query lengths. - """ - batch = context.batch - position_ids = batch.position_ids - if position_ids is None: - # Unit-level callers may construct an AttentionBatch directly. - # The runtime always supplies position_ids, but contiguous query - # ranges are an unambiguous fallback for those callers. - history = batch.b_seq_len - batch.b_q_seq_len - starts = torch.repeat_interleave(history, batch.b_q_seq_len, output_size=batch.token_num) - row_starts = torch.repeat_interleave( - batch.b1_cu_q_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num - ) - position_ids = starts + torch.arange(batch.token_num, device=batch.device, dtype=starts.dtype) - row_starts - context.q_kv_seq_len = position_ids + 1 - context.q_req_idx = torch.repeat_interleave(batch.b_req_idx, batch.b_q_seq_len, output_size=batch.token_num) - context.q_kv_start_loc = torch.repeat_interleave( - batch.b1_cu_kv_seq_len[:-1], batch.b_q_seq_len, output_size=batch.token_num - ) - def _run_prefill(self, layer, context, q: torch.Tensor, kv: torch.Tensor, inputs) -> torch.Tensor: from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -113,7 +79,6 @@ def build_device_metadata(self, context: NsaFlashMlaSparsePrefillContext): ragged_mem_index=context.ragged_mem_index, hold_req_idx=batch.hold_req_id, ) - self._update_query_metadata(context) return def forward( @@ -165,8 +130,6 @@ def allocate_context(self, capacity, *, capture, workspace) -> NsaFlashMlaSparse return NsaFlashMlaSparseDecodeContext(capacity, capture, workspace) def build_device_metadata(self, context: NsaFlashMlaSparseDecodeContext): - context.q_kv_seq_len = context.batch.b_seq_len - context.q_req_idx = context.batch.b_req_idx context.ragged_mem_index = context.workspace.scratch( "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 ) diff --git a/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py index 7a41389d40..2f05253e94 100644 --- a/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py @@ -19,9 +19,6 @@ class NsaFlashMlaFp8SparsePrefillContext(AttentionContext): ke: torch.Tensor = None lengths: torch.Tensor = None ragged_mem_index: torch.Tensor = None - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None - q_kv_start_loc: torch.Tensor = None @dataclasses.dataclass(slots=True) @@ -32,8 +29,6 @@ class NsaFlashMlaFp8SparseDecodeContext(AttentionContext): ragged_mem_index: torch.Tensor = None flashmla_sched_meta: object = None capture_flashmla_sched_meta: object = None - q_kv_seq_len: torch.Tensor = None - q_req_idx: torch.Tensor = None class FlashMLASparseMLAFp8Common: @@ -134,8 +129,6 @@ def allocate_context(self, capacity, *, capture, workspace) -> NsaFlashMlaFp8Spa return NsaFlashMlaFp8SparseDecodeContext(capacity, capture, workspace) def build_device_metadata(self, context: NsaFlashMlaFp8SparseDecodeContext): - context.q_kv_seq_len = context.batch.b_seq_len - context.q_req_idx = context.batch.b_req_idx context.ragged_mem_index = context.workspace.scratch( "sparse_ragged", max(1, context.capacity.batch_size * context.capacity.max_kv_seq_len), torch.int32 ) diff --git a/lightllm/common/attention/probe.py b/lightllm/common/attention/probe.py index b16a5d1fc3..19920c52ab 100644 --- a/lightllm/common/attention/probe.py +++ b/lightllm/common/attention/probe.py @@ -16,7 +16,6 @@ def probe_dense(backend_cls, spec, runtime): q_heads, kv_heads, dim = spec.num_q_heads, spec.cache.key_heads, spec.head_size env = BackendEnv( runtime=runtime, - network_config={}, hold_req_id=2, draft_step=0, max_seq_len=stride, @@ -50,19 +49,16 @@ def probe_dense(backend_cls, spec, runtime): phase=spec.phase, batch_size=2, token_num=sum(q_lengths), max_kv_seq_len=stride, max_q_seq_len=max(q_lengths) ) workspace = AttentionWorkspace(runtime=runtime) - workspace.reserve(backend.resource_requirements(capacity, capture=False)) session = AttentionSession() - context = session.allocate(backend, capacity, capture=False, workspace=workspace) - session.plan(backend, context, batch) - session.build(backend, context) - session.contexts[0] = context attention = LayerAttention( group_id=0, logical_layer_id=0, resources=DenseResources(key, value), backend=backend, ) - actual = attention.forward(session, DenseInputs(query), out=torch.empty_like(query)) + with session.prepare(batch, ({batch.phase: backend},), capacity, workspace=lambda _: workspace): + session.build_device_metadata() + actual = attention.forward(session, DenseInputs(query), out=torch.empty_like(query)) expected, start = [], 0 for request, (length, q_length) in enumerate(zip(lengths, q_lengths)): q = query[start : start + q_length].float() diff --git a/lightllm/common/attention/session.py b/lightllm/common/attention/session.py index 54828df48e..8a773fd901 100644 --- a/lightllm/common/attention/session.py +++ b/lightllm/common/attention/session.py @@ -9,13 +9,17 @@ def __init__(self): self.groups = () self.contexts = {} self.batch = None + self._active = False @contextmanager def prepare(self, batch, groups, capacity, *, workspace, capture=False, replay=False): + if self._active: + raise RuntimeError("attention session is already active") batch.validate() capacity.check(batch) self.groups = groups self.batch = batch + self.contexts = {} with ExitStack() as leases: if not replay: if capture: @@ -31,64 +35,52 @@ def prepare(self, batch, groups, capacity, *, workspace, capture=False, replay=F if capture: scratch = scratch.borrow_for_capture() scratch.reserve(backend.resource_requirements(capacity, capture=capture)) - context = self.allocate(backend, capacity, capture=capture, workspace=scratch) + context = self._allocate(backend, capacity, capture=capture, workspace=scratch) context.batch = self.batch context.current_max_kv_seq_len = batch.max_kv_seq_len backend.plan_host(context) self.contexts[group_id] = context - yield self + self._active = not replay + try: + yield self + finally: + self._active = False @staticmethod - def allocate(backend, capacity, *, capture, workspace): + def _allocate(backend, capacity, *, capture, workspace): if capacity.phase is not backend.phase: raise ValueError("backend and context phases differ") context = backend.allocate_context(capacity, capture=capture, workspace=workspace) if context.capacity != capacity: raise ValueError("backend changed attention context capacity") + if context.workspace is not workspace or context.capture != capture: + raise ValueError("backend changed attention context workspace or capture mode") context.backend_identity = backend return context - @staticmethod - def plan(backend, context, batch, *, copy_inputs=True): - AttentionSession._check_context(backend, context) - batch.validate() - context.capacity.check(batch) - context.current_max_kv_seq_len = batch.max_kv_seq_len - if context.capture and context.batch is not None: - if copy_inputs: - context.batch.copy_inputs(batch) - elif context.capture: - context.batch = replace( - batch, - max_kv_seq_len=context.capacity.max_kv_seq_len, - total_token_num=context.capacity.batch_size * context.capacity.max_kv_seq_len, - ) - else: - context.batch = batch - backend.plan_host(context) - @staticmethod def _check_context(backend, context): if context.backend_identity is not backend or context.capacity.phase is not backend.phase: raise ValueError("context belongs to a different attention backend or phase") + def _check_active(self): + if not self._active: + raise RuntimeError("attention execution requires an active session scope") + def context(self, layer): + self._check_active() context = self.contexts[layer.group_id] self._check_context(layer.backend, context) if context.batch is None: raise RuntimeError("prepare attention inputs before execution") return context - @staticmethod - def build(backend, context): - AttentionSession._check_context(backend, context) - if context.batch is None: - raise RuntimeError("prepare attention inputs before update") - backend.build_device_metadata(context) - def build_device_metadata(self): + self._check_active() for group_id, context in self.contexts.items(): - self.build(self.groups[group_id][context.capacity.phase], context) + backend = self.groups[group_id][context.capacity.phase] + self._check_context(backend, context) + backend.build_device_metadata(context) def forward(self, layer, inputs, *, out=None): context = self.context(layer) @@ -99,13 +91,22 @@ def forward(self, layer, inputs, *, out=None): if attention.dtype != out.dtype or attention.device != out.device: raise ValueError("out must match the attention result dtype and device") if attention.untyped_storage().data_ptr() == out.untyped_storage().data_ptr(): - begin = attention.storage_offset() - end = begin + sum((size - 1) * stride for size, stride in zip(attention.shape, attention.stride())) + 1 - out_end = ( - out.storage_offset() + sum((size - 1) * stride for size, stride in zip(out.shape, out.stride())) + 1 - ) - if attention.numel() and (begin < out.storage_offset() or end > out_end): - raise ValueError("attention result view extends beyond out") + if out.is_contiguous(): + begin = attention.storage_offset() + end = begin + sum((size - 1) * stride for size, stride in zip(attention.shape, attention.stride())) + 1 + valid = begin >= out.storage_offset() and end <= out.storage_offset() + out.numel() + else: + valid = ( + attention.ndim == out.ndim + and attention.storage_offset() == out.storage_offset() + and all(size <= bound for size, bound in zip(attention.shape, out.shape)) + and all( + size <= 1 or stride == out_stride + for size, stride, out_stride in zip(attention.shape, attention.stride(), out.stride()) + ) + ) + if attention.numel() and not valid: + raise ValueError("attention result must be contained in contiguous out or a strided out prefix") return result if attention.numel() != out.numel(): raise ValueError("out must match the attention result element count") diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_types.py b/lightllm/common/attention_state/bounded_att/conv_ssm_types.py index aa0666eb3b..14fda229bd 100644 --- a/lightllm/common/attention_state/bounded_att/conv_ssm_types.py +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_types.py @@ -51,3 +51,35 @@ def __post_init__(self) -> None: assert heads > 0 and ( heads % self.tp_world_size == 0 or self.tp_world_size % heads == 0 ), "Conv/SSM heads require complete TP shards or replicas" + + +@dataclass(frozen=True) +class ReplayStorage: + logical_layer_ids: tuple[int, ...] + ssm_state: torch.Tensor + keys: torch.Tensor + deltas: torch.Tensor + gates: torch.Tensor + raw_keys: torch.Tensor + raw_values: torch.Tensor + betas: torch.Tensor + cursors: torch.Tensor + tile_sync_mark: torch.Tensor + num_key_heads: int + capacity: int + verify_width: int + compute_mode: str + gate_lower_bound: float + hold_req_id: int + kda: bool + + def storage_layer(self, logical_layer_id): + return self.logical_layer_ids.index(logical_layer_id) + + +@dataclass(frozen=True) +class RecurrentStorage: + conv: torch.Tensor + ssm: torch.Tensor + mtp_state_index: torch.Tensor + replay_state: ReplayStorage | None = None diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index 1a578c8b83..b5d9fe3be3 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -120,7 +120,7 @@ def contains_logical_layer(self, logical_layer_id: int) -> bool: return logical_layer_id in self._logical_to_storage_layer def layer_resources(self, layer_id): - from lightllm.common.attention.api.inputs import RecurrentStorage + from ..conv_ssm_types import RecurrentStorage conv, ssm = self.get_layer_view(layer_id) return RecurrentStorage(conv, ssm, self.mtp_state_index) diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py index e45fc34e3d..e2e8e2043a 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -20,7 +20,7 @@ from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState -from ..conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_types import ConvSSMStateConfig, ReplayStorage, RecurrentStorage from ..copy_kernel import copy_bounded_att_tensors from ..transfer_utils import bounded_att_buffer_copies from .gdn_pages_pool import GDNPagesPool @@ -102,7 +102,7 @@ def __init__( assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 self.hold_req_id = slots - 1 - # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 + # cursor 是已接受、尚未折回的长度。各层只读，模型完成本步后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) # 按 launch 行和 key head 同步。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 self.tile_sync_mark = torch.zeros( @@ -195,10 +195,8 @@ def storage_layer(self, logical_layer_id: int) -> int: return self._logical_to_storage_layer[logical_layer_id] def layer_resources(self, layer_id): - from lightllm.common.attention.api.inputs import RecurrentStorage - conv, ssm = self.get_layer_view(layer_id) - return RecurrentStorage(conv, ssm, self.mtp_state_index, self) + return RecurrentStorage(conv, ssm, self.mtp_state_index, self.replay_storage()) def get_layer_view(self, logical_layer_id: int): index = self.storage_layer(logical_layer_id) @@ -456,17 +454,18 @@ def read_from_pd_page( # 六、Replay history # # cursor 是已接受、尚未折回的长度，各层只读。 - # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 + # 模型完成本步后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 把已接受前缀写回 canonical SSM。 # ======================================================================== - def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: - """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" + def replay_storage(self): + return ReplayStorage( + logical_layer_ids=self.config.logical_layer_ids, + **{name: getattr(self, name) for name in ReplayStorage.__dataclass_fields__ if name != "logical_layer_ids"}, + ) - layers = () if self.config is None else self.config.logical_layer_ids - if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: - return + def finish_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: advance_cursor(self, reqs, cu_seqlens) def fold_accepted_prefix(self, reqs: torch.Tensor) -> None: diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index 7ee2da5d58..50b499a3f7 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -120,7 +120,7 @@ def contains_logical_layer(self, logical_layer_id: int) -> bool: return logical_layer_id in self._logical_to_storage_layer def layer_resources(self, layer_id): - from lightllm.common.attention.api.inputs import RecurrentStorage + from ..conv_ssm_types import RecurrentStorage conv, ssm = self.get_layer_view(layer_id) return RecurrentStorage(conv, ssm, self.mtp_state_index) diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py index f9fb7d82e7..fcae4ffee5 100644 --- a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -20,7 +20,7 @@ from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion from ...utils import deepcopy_result, dtype_name from ..base_state import BoundedAttentionState -from ..conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_types import ConvSSMStateConfig, ReplayStorage, RecurrentStorage from ..copy_kernel import copy_bounded_att_tensors from ..transfer_utils import bounded_att_buffer_copies from .kda_pages_pool import KDAPagesPool @@ -102,7 +102,7 @@ def __init__( assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 self.hold_req_id = slots - 1 - # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 + # cursor 是已接受、尚未折回的长度。各层只读，模型完成本步后才按整窗推进。 self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) # 按 launch 行和 key head 同步。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 self.tile_sync_mark = torch.zeros( @@ -195,10 +195,8 @@ def storage_layer(self, logical_layer_id: int) -> int: return self._logical_to_storage_layer[logical_layer_id] def layer_resources(self, layer_id): - from lightllm.common.attention.api.inputs import RecurrentStorage - conv, ssm = self.get_layer_view(layer_id) - return RecurrentStorage(conv, ssm, self.mtp_state_index, self) + return RecurrentStorage(conv, ssm, self.mtp_state_index, self.replay_storage()) def get_layer_view(self, logical_layer_id: int): index = self.storage_layer(logical_layer_id) @@ -456,17 +454,18 @@ def read_from_pd_page( # 六、Replay history # # cursor 是已接受、尚未折回的长度，各层只读。 - # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 + # 模型完成本步后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 把已接受前缀写回 canonical SSM。 # ======================================================================== - def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: - """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" + def replay_storage(self): + return ReplayStorage( + logical_layer_ids=self.config.logical_layer_ids, + **{name: getattr(self, name) for name in ReplayStorage.__dataclass_fields__ if name != "logical_layer_ids"}, + ) - layers = () if self.config is None else self.config.logical_layer_ids - if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: - return + def finish_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: advance_cursor(self, reqs, cu_seqlens) def fold_accepted_prefix(self, reqs: torch.Tensor) -> None: diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py index 4fc95b0fae..213f6334f7 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -147,9 +147,9 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: return self._logical_to_storage_layer[logical_layer_id] def layer_resources(self, layer_id): - from lightllm.common.attention.api.inputs import DenseResources + from lightllm.common.attention_state.types import DenseStorage - return DenseResources(*self.get_layer_view(layer_id)) + return DenseStorage(*self.get_layer_view(layer_id)) def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" diff --git a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py index 376c3f4501..5dbc277665 100644 --- a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py +++ b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py @@ -153,10 +153,10 @@ def get_storage_layer_index(self, logical_layer_id: int) -> int: return self._logical_to_storage_layer[logical_layer_id] def layer_resources(self, layer_id): - from lightllm.common.attention.api.inputs import DenseResources + from lightllm.common.attention_state.types import DenseStorage (key, key_scale), (value, value_scale) = self.get_layer_view(layer_id) - return DenseResources(key, value, key_scale=key_scale, value_scale=value_scale) + return DenseStorage(key, value, key_scale=key_scale, value_scale=value_scale) def get_layer_view( self, logical_layer_id: int diff --git a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py index 118c044a87..2bf99904fb 100644 --- a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py +++ b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py @@ -228,7 +228,7 @@ def get_v_descale(self, storage_layer_index: int, batch_size: int) -> torch.Tens return self.scales[storage_layer_index, head_num:].view(1, head_num).expand(batch_size, head_num) def layer_resources(self, layer_id): - from lightllm.common.attention.api.inputs import DenseResources + from lightllm.common.attention_state.types import DenseStorage key, value = self.get_layer_view(layer_id) index = self.get_storage_layer_index(layer_id) @@ -239,7 +239,7 @@ def layer_resources(self, layer_id): else: key_scale, value_scale = map(float, self.cpu_scales[index]) query_scale = None - return DenseResources(key, value, query_scale=query_scale, key_scale=key_scale, value_scale=value_scale) + return DenseStorage(key, value, query_scale=query_scale, key_scale=key_scale, value_scale=value_scale) def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" diff --git a/lightllm/common/attention_state/types.py b/lightllm/common/attention_state/types.py index 473d5624b6..a79e88f209 100644 --- a/lightllm/common/attention_state/types.py +++ b/lightllm/common/attention_state/types.py @@ -235,3 +235,12 @@ class StateTransferLayout: # 当前 State 参与传输的字段区域；其他 State 的区域放在它们各自的 layout 中。 # 包含所需的 KV、动态量化参数或运行态字段，所有偏移均相对同一张完整外部页。 regions: tuple[TransferRegion, ...] + + +@dataclass(frozen=True) +class DenseStorage: + key: torch.Tensor + value: torch.Tensor + query_scale: torch.Tensor | None = None + key_scale: torch.Tensor | float | None = None + value_scale: torch.Tensor | float | None = None diff --git a/lightllm/common/basemodel/attention.py b/lightllm/common/basemodel/attention.py index 313ea5816a..2d5c0451dc 100644 --- a/lightllm/common/basemodel/attention.py +++ b/lightllm/common/basemodel/attention.py @@ -37,7 +37,6 @@ def _alloc_tensor(shape, dtype, device="cuda"): def backend_env(model): return BackendEnv( runtime=model.device_runtime, - network_config=model.config, hold_req_id=model.req_manager.HOLD_REQUEST_ID, draft_step=model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model), max_seq_len=model.max_seq_length, @@ -86,6 +85,7 @@ def __init__(self, model): self.routes = () self.adapters = () self.workspaces = {} + self.replay_states = () self.generation = None self.runtime = model.device_runtime self.constraints = describe_serving_constraints(model) @@ -100,6 +100,13 @@ def __init__(self, model): state.check_execution(caps, self.constraints) groups, keys, routes, adapters = [], {}, [], [] env = backend_env(model) + self.draft_step = env.draft_step + self.dynamic_verify = model.args.mtp_dynamic_verify and not model.is_mtp_draft_model + self.replay_states = tuple( + state + for state in model.att_state_manager.get_bounded_attention_states(model.att_state_role) + if getattr(state, "state_mode", None) == "replay" + ) for layer in model.layers_infer: specs = [] adapter = layer.attention_adapter(layer.describe_attention(AttentionPhase.PREFILL).family) @@ -170,6 +177,14 @@ def bind(self, model): ) self.generation = model.att_state_manager.storage_generation + def finish_decode(self, infer_state): + if self.replay_states: + from lightllm.common.attention.backends._shared.linear.layout import decode_request_layout + + reqs, cu = decode_request_layout(infer_state.attention.batch, self.draft_step, self.dynamic_verify) + for state in self.replay_states: + state.finish_decode(reqs, cu) + def workspace(self, group_id, slot): key = group_id, slot if key not in self.workspaces: diff --git a/lightllm/common/basemodel/attention_adapters.py b/lightllm/common/basemodel/attention_adapters.py index 94a0fa334a..bab90d0acc 100644 --- a/lightllm/common/basemodel/attention_adapters.py +++ b/lightllm/common/basemodel/attention_adapters.py @@ -1,8 +1,15 @@ from dataclasses import dataclass, replace from lightllm.common.attention.api import AttentionFamily, AttentionPhase, CacheLayout, LayerSpec -from lightllm.common.attention.api.inputs import GDNResources, KDAResources, MLAResources -from lightllm.common.attention.api.spec import AttentionFeature, CacheEncoding +from lightllm.common.attention.api.inputs import DenseResources, GDNResources, KDAResources, MLAResources +from lightllm.common.attention.api.spec import ( + AttentionFeature, + CacheEncoding, + GDNConfig, + KDAConfig, + MLAConfig, + SparseMLAConfig, +) @dataclass(frozen=True) @@ -43,6 +50,7 @@ def describe(self, model, layer, phase, description): features=frozenset(features), variant=description.variant, softmax_scale=self.scale(layer), + config=self.computation_config(model, layer, state), ) def cache(self, model, state, states, heads, head_dim, value_dim): @@ -56,6 +64,9 @@ def cache(self, model, state, states, heads, head_dim, value_dim): storage_schema=tuple(state.describe_layout() for state in states), ) + def computation_config(self, model, layer, state): + return None + def scale(self, layer): return None @@ -76,8 +87,8 @@ def cache(self, model, state, states, heads, head_dim, value_dim): def bind(self, model, layer, weight): state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) resources = state.layer_resources(layer.layer_num_) - return replace( - resources, + return DenseResources( + **vars(resources), alibi=getattr(weight, "tp_alibi", None), sink=weight.attn_sinks.weight if AttentionFeature.SINK in layer.describe_attention(AttentionPhase.PREFILL).features @@ -88,6 +99,9 @@ def bind(self, model, layer, weight): class MLAAdapter(AttentionAdapter): state_type = "mla" + def computation_config(self, model, layer, state): + return MLAConfig(layer.qk_rope_head_dim, layer.v_head_dim) + def dimensions(self, model, layer): return layer.tp_q_head_num_, 1, layer.qk_nope_head_dim + layer.qk_rope_head_dim, layer.kv_lora_rank @@ -100,6 +114,9 @@ def bind(self, model, layer, weight): class SparseMLAAdapter(MLAAdapter): + def computation_config(self, model, layer, state): + return SparseMLAConfig(model.config.get("index_topk", 2048), model.config.get("index_kpool", 1)) + def dimensions(self, model, layer): return ( layer.tp_q_head_num_, @@ -135,10 +152,12 @@ def cache(self, model, state, states, heads, head_dim, value_dim): class GDNAdapter(ConvSSMAdapter): state_type = "gdn" + def computation_config(self, model, layer, state): + return GDNConfig(state.config.num_v_heads, state.config.conv_width + 1, layer.network_config_["hidden_act"]) + def dimensions(self, model, layer): - config = model.config - heads = config["linear_num_key_heads"] // model.tp_world_size_ - return heads, heads, config["linear_key_head_dim"], config["linear_value_head_dim"] + config = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_).config + return config.num_k_heads, config.num_k_heads, config.key_dim, config.value_dim def bind(self, model, layer, weight): state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) @@ -158,10 +177,14 @@ def bind(self, model, layer, weight): class KDAAdapter(ConvSSMAdapter): state_type = "kda" + def computation_config(self, model, layer, state): + return KDAConfig( + state.config.conv_width + 1, layer.network_config_["linear_attn_config"].get("gate_lower_bound", -5.0) + ) + def dimensions(self, model, layer): - config = model.config["linear_attn_config"] - heads = config["num_heads"] // model.tp_world_size_ - return heads, heads, config["head_dim"], config["head_dim"] + config = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_).config + return config.num_k_heads, config.num_k_heads, config.key_dim, config.value_dim def bind(self, model, layer, weight): state = model.att_state_manager.get_state(model.att_state_role, self.state_type, layer.layer_num_) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index d2241207d2..a6afba79f9 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -797,6 +797,7 @@ def _token_forward(self, infer_state: InferStateInfo): input_embs: torch.Tensor = layer.token_forward(input_embs, infer_state, self.trans_layers_weight[i]) hidden_collector.add(layer_index=i, hidden=input_embs) + self.attention.finish_decode(infer_state) last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) post_output: PostLayerOutput = self.post_infer.token_forward( last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight @@ -1064,6 +1065,8 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: infer_state.call_overlap_hook() infer_state1.call_overlap_hook() + self.attention.finish_decode(infer_state) + self.attention.finish_decode(infer_state1) last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state) last_input_embs1 = self.post_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 339c1440ac..21c82a7e62 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -20,6 +20,15 @@ class InferStateInfo: 推理时用的信息结构体 """ + graph_input_names = ( + "input_ids", + "b_position_delta", + "b_kv_seq_len", + "b_q_start_loc", + "b_kv_start_loc", + "mtp_draft_input_hiddens", + ) + def __init__(self): self.attention = AttentionSession() @@ -132,20 +141,20 @@ def graph_input_copies(self, new_infer_state: "InferStateInfo"): if self.att_state_role != new_infer_state.att_state_role: raise ValueError("captured attention belongs to a different model role or group") copies = self.attention.replay_input_copies(new_infer_state.attention) - batch_tensor_names = { - name for name, value in vars(self.attention.batch).items() if isinstance(value, torch.Tensor) - } - for attr_name, attr_value in vars(new_infer_state).items(): - if isinstance(attr_value, torch.Tensor) and attr_name not in batch_tensor_names: - target = getattr(self, attr_name, None) - if target is not None and target.data_ptr() != attr_value.data_ptr(): - if (target.shape, target.dtype, target.device) != ( - attr_value.shape, - attr_value.dtype, - attr_value.device, - ): - raise ValueError(f"captured input {attr_name} changed shape, dtype or device") - copies.append((target, attr_value)) + if self.graph_input_names != new_infer_state.graph_input_names: + raise ValueError("captured input inventory changed") + for name in self.graph_input_names: + target, source = getattr(self, name, None), getattr(new_infer_state, name, None) + if target is None and source is None: + continue + if not isinstance(target, torch.Tensor) or not isinstance(source, torch.Tensor): + raise ValueError(f"captured input {name} changed presence") + if (target.shape, target.dtype, target.device) != (source.shape, source.dtype, source.device): + raise ValueError(f"captured input {name} changed shape, dtype or device") + if target.data_ptr() == source.data_ptr() and target.stride() != source.stride(): + raise ValueError(f"captured input {name} changed strides at the captured address") + if target.data_ptr() != source.data_ptr(): + copies.append((target, source)) return copies def call_overlap_hook(self): diff --git a/lightllm/models/deepseek_v4/attention/adapter.py b/lightllm/models/deepseek_v4/attention/adapter.py index 11c3d560e9..8fe1770aca 100644 --- a/lightllm/models/deepseek_v4/attention/adapter.py +++ b/lightllm/models/deepseek_v4/attention/adapter.py @@ -5,6 +5,13 @@ class DeepseekV4Adapter(AttentionAdapter): state_type = "swa_" + def computation_config(self, model, layer, state): + from lightllm.models.deepseek_v4.config import create_attention_config + + return create_attention_config( + model.config, tuple(item.layer_num_ for item in model.layers_infer), model.tp_rank_, model.tp_world_size_ + ) + def dimensions(self, model, layer): return layer.tp_q_head_num_, 1, layer.head_dim_, layer.v_head_dim diff --git a/lightllm/models/deepseek_v4/attention/sparse_mla.py b/lightllm/models/deepseek_v4/attention/sparse_mla.py index 7da1569508..3c145da018 100644 --- a/lightllm/models/deepseek_v4/attention/sparse_mla.py +++ b/lightllm/models/deepseek_v4/attention/sparse_mla.py @@ -91,14 +91,7 @@ def check_environment(cls, caps) -> None: def __init__(self, env, spec): super().__init__(env=env, spec=spec) - from lightllm.models.deepseek_v4.config import create_attention_config - - config = create_attention_config( - model_config=env.network_config, - model_logical_layer_ids=env.layer_ids, - tp_rank=env.tp_rank, - tp_world_size=env.tp_size, - ) + config = spec.config self.compress_ratios = tuple(dict.fromkeys(config.compress_ratios)) self.swa_config = config.create_swa_config() self.compressed_configs = { diff --git a/lightllm/models/deepseek_v4/infer_struct.py b/lightllm/models/deepseek_v4/infer_struct.py index 01d4c67381..1722b66f59 100644 --- a/lightllm/models/deepseek_v4/infer_struct.py +++ b/lightllm/models/deepseek_v4/infer_struct.py @@ -31,6 +31,23 @@ class DeepseekV4InferStateInfo(InferStateInfo): SWA and compressor rings are owned by the model's bounded State. """ + graph_input_names = InferStateInfo.graph_input_names + ( + "position_cos_sliding", + "position_sin_sliding", + "position_cos_compress", + "position_sin_compress", + "dsv4_sparse_req_idx", + "dsv4_swa_indices", + "dsv4_swa_lengths", + "dsv4_swa_write_slots", + "dsv4_swa_runtime_write_slots", + "dsv4_image_left", + "dsv4_image_right", + "dsv4_c128_indices", + "dsv4_c128_lengths", + "_dsv4_token_to_batch_idx", + ) + @property def swa_state(self) -> FP8E8M0SlidingWindowState: return self.att_state_manager.get_state( diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index ea68793dab..213ff93df9 100644 --- a/lightllm/models/deepseek_v4/model.py +++ b/lightllm/models/deepseek_v4/model.py @@ -218,8 +218,8 @@ def prepare_attention_backends(self): self.args.llm_kv_type = args.llm_kv_type super().prepare_attention_backends() backend = self.attention.groups[0][AttentionPhase.PREFILL] - tp_q_head_num = backend.tp_q_head_num - padded_q_head_num = backend.padded_q_head_num + tp_q_head_num = backend.spec.num_q_heads + padded_q_head_num = (tp_q_head_num + 63) // 64 * 64 if not self.is_mtp_draft_model: # Drafts share this buffer after the main model has captured its CUDA graphs. self.dsv4_workspace.init_flashmla_prefill_q( diff --git a/lightllm/models/gemma3/infer_struct.py b/lightllm/models/gemma3/infer_struct.py index 33bd44815a..e636d35b74 100644 --- a/lightllm/models/gemma3/infer_struct.py +++ b/lightllm/models/gemma3/infer_struct.py @@ -5,6 +5,15 @@ class Gemma3InferStateInfo(InferStateInfo): + graph_input_names = InferStateInfo.graph_input_names + ( + "position_cos", + "position_sin", + "position_cos_local", + "position_sin_local", + "position_cos_global", + "position_sin_global", + ) + def __init__(self): super().__init__() self.position_cos_global = None diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index 89ad34acbd..ae424ccf5f 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -4,6 +4,13 @@ class Gemma4InferStateInfo(InferStateInfo): + graph_input_names = InferStateInfo.graph_input_names + ( + "position_cos_sliding", + "position_sin_sliding", + "position_cos_full", + "position_sin_full", + ) + def __init__(self): super().__init__() # Gemma-4 uses two RoPE frequency tables (one per layer type): diff --git a/lightllm/models/glm5_next/attention/fp8_sparse_mla.py b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py index ad33f0b255..12af3aa454 100644 --- a/lightllm/models/glm5_next/attention/fp8_sparse_mla.py +++ b/lightllm/models/glm5_next/attention/fp8_sparse_mla.py @@ -157,7 +157,7 @@ def resource_requirements(self, capacity, *, capture): tokens = max(1, self.prefill_token_capacity, capacity.token_num) kv_tokens = tokens if self.decode_only else capacity.batch_size * capacity.max_kv_seq_len + tokens gather_tokens = min(KV_GATHER_CHUNK_SIZE, kv_tokens) - topk = (self.network_config["index_topk"] + 3 + 127) // 128 * 128 + topk = (self.spec.config.index_topk + 3 + 127) // 128 * 128 return WorkspaceBudget( workspace=( WorkspaceRequirement("glm_gathered_kv", gather_tokens * 512, torch.bfloat16), @@ -232,6 +232,6 @@ def allocate_context(self, capacity, *, capture, workspace) -> Fp8Glm5NextSparse capture, workspace, decode_topk_bound=2051 - if self.network_config["index_topk"] == 2048 and self.network_config["index_kpool"] == 4 + if self.spec.config.index_topk == 2048 and self.spec.config.index_kpool == 4 else None, ) diff --git a/lightllm/models/glm5_next/attention/sparse_mla.py b/lightllm/models/glm5_next/attention/sparse_mla.py index 682bcd9213..c744dc8051 100644 --- a/lightllm/models/glm5_next/attention/sparse_mla.py +++ b/lightllm/models/glm5_next/attention/sparse_mla.py @@ -96,8 +96,8 @@ def allocate_context(self, capacity, *, capture, workspace) -> Glm5NextSparseDec def build_device_metadata(self, context: Glm5NextSparseDecodeContext): super().build_device_metadata(context) - pool = self.network_config["index_kpool"] - topk = self.network_config["index_topk"] + pool = self.spec.config.index_kpool + topk = self.spec.config.index_topk context.nsa_cache_seqlens = ( torch.minimum(context.lengths // pool * pool, torch.full_like(context.lengths, topk)) + context.lengths % pool diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index b7e6b9b7f3..f00c5f10a8 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -4,6 +4,7 @@ from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops +from .infer_struct import SparseQueryMetadata from .triton_kernel.index_quant import hadamard_transform_quant_fp8 from .triton_kernel.kpool import compress_pools, gather_pools, gather_paged_pools, get_pool_ranges, expand_topk @@ -41,7 +42,7 @@ def select_topk_indices(self, logits, q_pool_seq_len, topk_pool_indices): selected = torch.topk(logits, topk_pool_indices.shape[1], dim=-1, sorted=True).indices topk_pool_indices.copy_(selected.masked_fill(selected >= q_pool_seq_len[:, None], -1)) - def write_k_pools(self, hidden_states, infer_state, att_state, layer_weight): + def write_k_pools(self, hidden_states, infer_state, query_metadata: SparseQueryMetadata, layer_weight): k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.rms_norm_eps) gate = layer_weight.index_kpool_compress_gate.mm(hidden_states) k_and_gate = torch.cat((k, gate), -1) @@ -63,7 +64,7 @@ def write_k_pools(self, hidden_states, infer_state, att_state, layer_weight): indexer_tail=indexer_tail, indexer_k=indexer_k, index_kpool_compress_ape=layer_weight.index_kpool_compress_ape.weight, - q_kv_seq_len=att_state.q_kv_seq_len, + q_kv_seq_len=query_metadata.q_kv_seq_len, mem_index=infer_state.mem_index, b_req_idx=infer_state.b_req_idx, b1_cu_q_seq_len=infer_state.b1_cu_q_seq_len, @@ -73,14 +74,14 @@ def write_k_pools(self, hidden_states, infer_state, att_state, layer_weight): ) return indexer_k - def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weight): - indexer_k = self.write_k_pools(hidden_states, infer_state, att_state, layer_weight) + def _get_indices(self, hidden_states, q_lora, infer_state, query_metadata: SparseQueryMetadata, layer_weight): + indexer_k = self.write_k_pools(hidden_states, infer_state, query_metadata, layer_weight) if infer_state.max_kv_seq_len <= self.index_topk: return expand_topk( None, - att_state.q_kv_seq_len, - att_state.q_req_idx, + query_metadata.q_kv_seq_len, + query_metadata.q_req_idx, infer_state.req_manager.req_to_token_indexs, self.index_topk, dense=True, @@ -96,13 +97,17 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig # Prefill usually has many queries: parallelize over Q and reuse K within each query tile. # Decode has few queries, so paged MQA also splits K across SMs for parallelism. if infer_state.is_prefill: - topk_pool_indices = self._get_prefill_indices(q_fp8, weights, indexer_k, infer_state, att_state, max_pools) + topk_pool_indices = self._get_prefill_indices( + q_fp8, weights, indexer_k, infer_state, query_metadata, max_pools + ) else: - topk_pool_indices = self._get_decode_indices(q_fp8, weights, indexer_k, infer_state, att_state, max_pools) + topk_pool_indices = self._get_decode_indices( + q_fp8, weights, indexer_k, infer_state, query_metadata, max_pools + ) return expand_topk( topk_pool_indices, - att_state.q_kv_seq_len, - att_state.q_req_idx, + query_metadata.q_kv_seq_len, + query_metadata.q_req_idx, infer_state.req_manager.req_to_token_indexs, self.index_topk, ) @@ -114,7 +119,7 @@ def _get_prefill_chunk_size(self, max_pools): query_tiles = max(1, 512 * 1024 ** 2 // (query_tile * logits_row_bytes)) return query_tiles * query_tile - def _get_prefill_indices(self, q_fp8, weights, indexer_k, infer_state, att_state, max_pools): + def _get_prefill_indices(self, q_fp8, weights, indexer_k, infer_state, query_metadata, max_pools): k_fp8_and_scale = gather_pools( indexer_k, infer_state.req_manager.req_to_token_indexs, @@ -123,7 +128,7 @@ def _get_prefill_indices(self, q_fp8, weights, indexer_k, infer_state, att_state max_pools, ) ks, ke, q_pool_seq_len = get_pool_ranges( - att_state.q_kv_seq_len, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools + query_metadata.q_kv_seq_len, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools ) topk_pool_indices = torch.empty((q_fp8.shape[0], self.index_topk // 4), dtype=torch.int32, device=q_fp8.device) chunk_size = self._get_prefill_chunk_size(max_pools) @@ -145,10 +150,10 @@ def _get_prefill_indices(self, q_fp8, weights, indexer_k, infer_state, att_state del logits return topk_pool_indices - def _get_decode_indices(self, q_fp8, weights, indexer_k, infer_state, att_state, max_pools): + def _get_decode_indices(self, q_fp8, weights, indexer_k, infer_state, query_metadata, max_pools): import deep_gemm - q_pool_seq_len = (att_state.q_kv_seq_len // 4).view(-1, 1) + q_pool_seq_len = (query_metadata.q_kv_seq_len // 4).view(-1, 1) indexer_k_pages, block_table = gather_paged_pools( indexer_k, infer_state.req_manager.req_to_token_indexs, diff --git a/lightllm/models/glm5_next/infer_struct.py b/lightllm/models/glm5_next/infer_struct.py new file mode 100644 index 0000000000..f1f389f3a7 --- /dev/null +++ b/lightllm/models/glm5_next/infer_struct.py @@ -0,0 +1,30 @@ +from typing import NamedTuple + +import torch + +from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class SparseQueryMetadata(NamedTuple): + q_kv_seq_len: torch.Tensor + q_req_idx: torch.Tensor + + +class Glm5NextInferStateInfo(InferStateInfo): + graph_input_names = InferStateInfo.graph_input_names + ( + "q_kv_seq_len", + "q_req_idx", + ) + + def init_some_extra_state(self, model): + super().init_some_extra_state(model) + self.q_kv_seq_len = self.position_ids + 1 if self.is_prefill else self.b_seq_len + self.q_req_idx = ( + torch.repeat_interleave(self.b_req_idx, self.b_q_seq_len, output_size=self.input_ids.shape[0]) + if self.is_prefill + else self.b_req_idx + ) + + @property + def sparse_queries(self): + return SparseQueryMetadata(self.q_kv_seq_len, self.q_req_idx) diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 1df0f3732c..48775b1700 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -168,7 +168,7 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], infer_state=infer_state, - att_state=self.prefill_attention.context(infer_state.attention), + query_metadata=infer_state.sparse_queries, layer_weight=layer_weight, ) del infer_state.get_topk_indices_params @@ -238,7 +238,7 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): hidden_states=infer_state.get_topk_indices_params["hidden_states"], q_lora=infer_state.get_topk_indices_params["q_lora"], infer_state=infer_state, - att_state=self.decode_attention.context(infer_state.attention), + query_metadata=infer_state.sparse_queries, layer_weight=layer_weight, ) del infer_state.get_topk_indices_params diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 14f271e91e..9e0ad36a25 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -13,6 +13,7 @@ from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer +from .infer_struct import Glm5NextInferStateInfo from .layer_infer.pre_layer_infer import Glm5NextPreLayerInfer from .layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from .layer_weights.pre_and_post_layer_weight import Glm5NextPreAndPostLayerWeight @@ -27,6 +28,7 @@ class Glm5NextTpPartModel(TpPartBaseModel): + infer_state_class = Glm5NextInferStateInfo is_hybrid_att_model = True pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight transformer_weight_class = Glm5NextTransformerLayerWeight diff --git a/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py index a0e455d004..1ddd093a19 100644 --- a/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next_mtp/layer_infer/transformer_layer_infer.py @@ -10,9 +10,7 @@ def context_forward(self, input_embeddings, infer_state, layer_weight): hidden, _, kv = self._project_sparse_qkv_a(normed, infer_state, layer_weight) infer_state.mtp_draft_input_hiddens = None self._post_cache_kv(self._normalize_sparse_kv(kv, layer_weight), infer_state, layer_weight) - self.indexer.write_k_pools( - hidden, infer_state, self.prefill_attention.context(infer_state.attention), layer_weight - ) + self.indexer.write_k_pools(hidden, infer_state, infer_state.sparse_queries, layer_weight) # BaseModel pads empty ranks to at least one row. An empty result marks # the discarded draft output for our post layer without retaining hiddens. return input_embeddings.new_empty((0, self.embed_dim_)) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 766450f157..766d7479f1 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -21,7 +21,11 @@ from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.models.glm5_next.infer_struct import Glm5NextInferStateInfo + + class Glm5NextMTPModel(TpPartBaseModel): + infer_state_class = Glm5NextInferStateInfo is_mtp_draft_model = True pre_and_post_weight_class = Glm5NextMTPPreAndPostLayerWeight transformer_weight_class = Glm5NextTransformerLayerWeight diff --git a/lightllm/models/llama/infer_struct.py b/lightllm/models/llama/infer_struct.py index fe6ca392a2..429f18e879 100644 --- a/lightllm/models/llama/infer_struct.py +++ b/lightllm/models/llama/infer_struct.py @@ -5,6 +5,11 @@ class LlamaInferStateInfo(InferStateInfo): + graph_input_names = InferStateInfo.graph_input_names + ( + "position_cos", + "position_sin", + ) + def __init__(self): super().__init__() self.position_cos = None diff --git a/lightllm/models/neo_chat_moe/infer_struct.py b/lightllm/models/neo_chat_moe/infer_struct.py index 779dcb65e2..311ad8ea69 100644 --- a/lightllm/models/neo_chat_moe/infer_struct.py +++ b/lightllm/models/neo_chat_moe/infer_struct.py @@ -8,6 +8,13 @@ class NeoChatInferStateInfo(LlamaInferStateInfo): + graph_input_names = LlamaInferStateInfo.graph_input_names + ( + "position_cos_h", + "position_sin_h", + "position_cos_w", + "position_sin_w", + ) + def __init__(self): super().__init__() self.position_cos = None diff --git a/lightllm/models/qwen/infer_struct.py b/lightllm/models/qwen/infer_struct.py index d575f6d95d..265f98cd0c 100644 --- a/lightllm/models/qwen/infer_struct.py +++ b/lightllm/models/qwen/infer_struct.py @@ -5,6 +5,8 @@ class QwenInferStateInfo(LlamaInferStateInfo): + graph_input_names = LlamaInferStateInfo.graph_input_names + ("logn_values",) + def __init__(self): super().__init__() self.position_cos = None diff --git a/unit_tests/common/attention/helpers.py b/unit_tests/common/attention/helpers.py index 7e466ff46e..2d060003b4 100644 --- a/unit_tests/common/attention/helpers.py +++ b/unit_tests/common/attention/helpers.py @@ -1,7 +1,8 @@ from lightllm.common.attention.api.environment import BackendEnv -from lightllm.common.attention.api.spec import AttentionFeature +from lightllm.common.attention.api.spec import AttentionFeature, GDNConfig, KDAConfig, MLAConfig, SparseMLAConfig from lightllm.common.attention.session import AttentionSession import importlib +from contextlib import contextmanager from dataclasses import replace from types import SimpleNamespace @@ -64,6 +65,12 @@ def make_spec(family=AttentionFamily.DENSE, phase=AttentionPhase.DECODE, **kwarg head_size=8, value_head_size=8, ) + values["config"] = { + AttentionFamily.GDN: GDNConfig(2, 4, "silu"), + AttentionFamily.KDA: KDAConfig(4, -5.0), + AttentionFamily.MLA: MLAConfig(0, 8), + AttentionFamily.SPARSE_MLA: SparseMLAConfig(2048, 4), + }.get(family) features = set(kwargs.pop("features", ())) for key, feature in ( ("use_mtp", "MTP"), @@ -132,7 +139,7 @@ def make_context(backend, batch=None, *, capture=False, max_length=None, prepare max_q_seq_len=batch.max_q_seq_len, ) backend = phase_backend(backend, batch.phase) - context = AttentionSession.allocate(backend, capacity, capture=capture, workspace=AttentionWorkspace(batch.device)) + context = AttentionSession._allocate(backend, capacity, capture=capture, workspace=AttentionWorkspace(batch.device)) if prepare: plan(backend, context, batch) else: @@ -232,7 +239,6 @@ def backend_context(cls, phase, settings=None, **values): config["hold_req_id"] = getattr(req_manager, "HOLD_REQUEST_ID", 7) backend = make_backend(cls, **config) defaults = dict( - network_config={"index_topk": 2048, "index_kpool": 4}, prefill_graph_enabled=False, decode_only=False, prefill_token_capacity=0, @@ -263,7 +269,7 @@ def update_context(context, fixture): build(context.backend_identity, context) -def run_sparse(context, fixture, query, kv, parameters, *, out=None, alloc_func=None): +def run_sparse(context, fixture, query, kv, parameters, *, out=None, alloc_func=None, session=None): from lightllm.common.attention.api.inputs import MLAResources, SparseMLADecodeInputs, SparseMLAPrefillInputs q = query[0] if isinstance(query, tuple) else query @@ -317,7 +323,9 @@ def run_sparse(context, fixture, query, kv, parameters, *, out=None, alloc_func= import flash_mla context.capture_flashmla_sched_meta = flash_mla.get_mla_metadata()[0] - return run_attention(attention, context, inputs, out=out) + if session is not None: + return attention.forward(session, inputs, out=out) + return backend.forward(attention, context, inputs, out=out) def run_dense(context, fixture, query, key, value, *, out=None, alloc_func=None, params=None): @@ -333,7 +341,7 @@ def run_dense(context, fixture, query, key, value, *, out=None, alloc_func=None, attention = make_layer_attention(context.backend_identity, context, params) if alloc_func is not None: context.backend_identity.alloc_tensor = alloc_func - return run_attention(attention, context, DenseInputs(query), out=out) + return context.backend_identity.forward(attention, context, DenseInputs(query), out=out) class StubBackend(AttentionBackend): @@ -352,13 +360,8 @@ def forward(self, layer, context, inputs, *, out=None): def make_env(model=None, **kwargs): model = model or model_stub() - config = dict(model.config) - for name in ("qk_nope_head_dim", "qk_rope_head_dim", "kv_lora_rank", "v_head_dim"): - if hasattr(model, name): - config[name] = getattr(model, name) values = dict( runtime=model.device_runtime, - network_config=config, hold_req_id=getattr(model.req_manager, "HOLD_REQUEST_ID", 7), draft_step=model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model), state_slots=model.args.mtp_step + 1, @@ -390,17 +393,52 @@ def phase_backend(backend, phase): def plan(backend, context, batch): - return AttentionSession.plan(backend, context, batch) + AttentionSession._check_context(backend, context) + batch.validate() + context.capacity.check(batch) + context.current_max_kv_seq_len = batch.max_kv_seq_len + if context.capture: + context.batch = context.batch or replace( + batch, + max_kv_seq_len=context.capacity.max_kv_seq_len, + total_token_num=context.capacity.batch_size * context.capacity.max_kv_seq_len, + ) + context.batch.copy_inputs(batch) + else: + context.batch = batch + backend.plan_host(context) def build(backend, context): - return AttentionSession.build(backend, context) + AttentionSession._check_context(backend, context) + backend.build_device_metadata(context) + + +@contextmanager +def prepared_session(context, session=None): + from unittest.mock import patch + + session = session or AttentionSession() + backend = context.backend_identity + + def allocate(capacity, *, capture, workspace): + context.workspace = workspace + return context + + with patch.object(backend, "allocate_context", allocate), patch.object(backend, "plan_host", lambda ctx: None): + with session.prepare( + context.batch, + ({context.capacity.phase: backend},), + context.capacity, + workspace=lambda _: context.workspace, + capture=context.capture, + ): + yield session def run_attention(layer, context, inputs=None, *, out=None): - session = AttentionSession() - session.contexts[layer.group_id] = context - return layer.forward(session, inputs, out=out) + with prepared_session(context) as session: + return layer.forward(session, inputs, out=out) def dense_class(provider, phase): diff --git a/unit_tests/common/attention/test_contracts.py b/unit_tests/common/attention/test_contracts.py index 00b872c006..6d61482e9d 100644 --- a/unit_tests/common/attention/test_contracts.py +++ b/unit_tests/common/attention/test_contracts.py @@ -72,7 +72,7 @@ def test_layer_attention_rejects_wrong_phase(): context = make_context(backend, prepare=True) attention = make_layer_attention(backend, context, params=None) context.capacity = replace(context.capacity, phase=AttentionPhase.PREFILL) - with pytest.raises(ValueError, match="different attention backend"): + with pytest.raises(ValueError, match="capacity"): run_attention(attention, context, inputs=None) @@ -228,7 +228,8 @@ def test_fp8_scale_layouts_are_not_interchangeable(format, provider): assert not any(cls.provider == provider and CacheEncoding(format) in cls.cache_encodings for cls in backends) -def test_layer_shape_configures_group_planner(): +def test_layer_shape_configures_group_planner(monkeypatch): + monkeypatch.setattr("lightllm.common.attention.backends.dense.flashinfer.backend.set_flashinfer_envs", lambda: None) from lightllm.common.attention.backends.dense.flashinfer.backend import FlashInferDenseDecodeBackend model = model_stub(page_size=16) @@ -368,15 +369,20 @@ def test_batch_builder_borrows_runtime_inputs_and_binding_checks_owner(): assert borrowed.b_image_token_end is state.b_image_token_end backend = _OutputBackend(make_env(), make_spec()) context = make_context(backend, borrowed, prepare=True) - state.attention.contexts = {0: context} + from lightllm.common.attention.api import AttentionWorkspace + layer = make_layer_attention(backend, context, params=None) - assert layer.context(state.attention) is context - output = torch.ones(2, 32) - assert layer.forward(state.attention, output) is output - with pytest.raises(ValueError, match="different attention backend|specifications differ"): - replace(layer, backend=_OutputBackend(make_env(), make_spec())).forward(state.attention, output) - with pytest.raises(ValueError, match="different attention backend|specifications differ"): - replace(layer, backend=StubBackend(make_env(), make_spec())).context(state.attention) + with state.attention.prepare( + borrowed, ({borrowed.phase: backend},), context.capacity, workspace=lambda _: AttentionWorkspace("cpu") + ): + context = state.attention.contexts[0] + assert layer.context(state.attention) is context + output = torch.ones(2, 32) + assert layer.forward(state.attention, output) is output + with pytest.raises(ValueError, match="different attention backend"): + replace(layer, backend=_OutputBackend(make_env(), make_spec())).forward(state.attention, output) + with pytest.raises(ValueError, match="different attention backend"): + replace(layer, backend=StubBackend(make_env(), make_spec())).context(state.attention) with pytest.raises(AttributeError): context.b_requ_idx = borrowed.b_req_idx @@ -487,7 +493,7 @@ def test_gdn_output_contract(monkeypatch, replay, phase): else: monkeypatch.setattr(backend, "_decode" if replay else "_run_decode", lambda *a: (core, gate)) context = make_context(backend, make_batch(phase=phase)) - state = SimpleNamespace(finish_decode_token=lambda *a: None) + state = SimpleNamespace() layer = make_layer_attention(backend, context, SimpleNamespace(conv=None, ssm=None, replay_state=state)) out = torch.empty_like(core) result = run_attention(layer, context, GDNInputs(core), out=out) @@ -498,12 +504,9 @@ def test_gdn_output_contract(monkeypatch, replay, phase): def test_binding_resolves_context_on_each_forward(monkeypatch): - from lightllm.common.attention.session import AttentionSession - backend = _OutputBackend(make_env(), make_spec()) first, second = make_context(backend), make_context(backend) layer = make_layer_attention(backend, first, params=None) - session = AttentionSession() forwarded = [] def forward(layer, context, inputs, *, out=None): @@ -513,8 +516,7 @@ def forward(layer, context, inputs, *, out=None): monkeypatch.setattr(backend, "forward", forward) result, out = torch.ones(2, 32), torch.empty(2, 32) for context in (first, second): - session.contexts[0] = context - assert layer.forward(session, result, out=out) is out + assert run_attention(layer, context, result, out=out) is out assert all(a is b for a, b in zip(forwarded, (first, second))) torch.testing.assert_close(out, result) @@ -538,6 +540,11 @@ def check_support(cls, spec, caps=None, constraints=None): def allocate_context(self, capacity, *, capture, workspace): return AttentionContext(capacity, capture, workspace) + def resource_requirements(self, capacity, *, capture): + from lightllm.common.attention.api import WorkspaceBudget + + return WorkspaceBudget() + def plan_host(self, context): pass @@ -547,19 +554,20 @@ def forward(self, layer, context, inputs, *, out=None): backend, session = Backend(), AttentionSession() batch = make_batch() capacity = ContextCapacity(phase=batch.phase, batch_size=2, token_num=2, max_kv_seq_len=16, max_q_seq_len=1) - context = session.allocate(backend, capacity, capture=True, workspace=AttentionWorkspace("cpu")) - session.plan(backend, context, batch) - session.contexts[0] = context - layer = make_layer_attention(StubBackend(make_env(), make_spec()), context, None) - layer = replace(layer, backend=backend) - out = torch.zeros(2, 4) - assert layer.forward(session, torch.ones(2, 4), out=out) is out - assert out.eq(1).all() - with pytest.raises(ValueError, match="different attention backend"): - replace(layer, backend=Backend()).forward(session, out) - context.capacity = replace(context.capacity, phase=AttentionPhase.PREFILL) - with pytest.raises(ValueError, match="different attention backend"): - layer.forward(session, out) + with session.prepare( + batch, ({batch.phase: backend},), capacity, workspace=lambda _: AttentionWorkspace("cpu"), capture=True + ): + context = session.contexts[0] + layer = make_layer_attention(StubBackend(make_env(), make_spec()), context, None) + layer = replace(layer, backend=backend) + out = torch.zeros(2, 4) + assert layer.forward(session, torch.ones(2, 4), out=out) is out + assert out.eq(1).all() + with pytest.raises(ValueError, match="different attention backend"): + replace(layer, backend=Backend()).forward(session, out) + context.capacity = replace(context.capacity, phase=AttentionPhase.PREFILL) + with pytest.raises(ValueError, match="different attention backend"): + layer.forward(session, out) def test_replay_copies_shared_batch_once(monkeypatch): @@ -679,3 +687,114 @@ def finish(self, dependencies=()): raise RuntimeError("forward failed") assert events == [("begin", 0), ("begin", 1), ("finish", 1), ("finish", 0)] + + +@pytest.mark.parametrize("invalid", ["missing_hidden", "new_hidden", "same_pointer_shape", "same_pointer_stride"]) +def test_graph_inputs_validate_presence_and_alias_metadata(invalid): + from lightllm.common.basemodel.infer_struct import InferStateInfo + + captured, incoming = InferStateInfo(), InferStateInfo() + captured.attention.batch = make_batch() + incoming.attention.batch = make_batch((9, 7), req_to_token_indexs=captured.attention.batch.req_to_token_indexs) + captured.input_ids = torch.zeros(2, dtype=torch.int64) + incoming.input_ids = torch.ones_like(captured.input_ids) + if invalid == "same_pointer_shape": + incoming.input_ids = captured.input_ids.view(1, 2) + elif invalid == "same_pointer_stride": + storage = torch.zeros(4, dtype=torch.int64) + captured.input_ids, incoming.input_ids = storage[::2], storage[:2] + elif invalid == "missing_hidden": + captured.mtp_draft_input_hiddens = torch.ones(2, 4) + else: + incoming.mtp_draft_input_hiddens = torch.ones(2, 4) + before = captured.attention.batch.b_seq_len.clone() + with pytest.raises(ValueError, match="captured input"): + captured.graph_input_copies(incoming) + torch.testing.assert_close(captured.attention.batch.b_seq_len, before) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_session_execution_ends_with_scope(capture): + from lightllm.common.attention.api import AttentionWorkspace, ContextCapacity, LayerAttention + from lightllm.common.attention.session import AttentionSession + + session = AttentionSession() + backend = _OutputBackend(make_env(), make_spec()) + batch = make_batch() + capacity = ContextCapacity(phase=batch.phase, batch_size=2, token_num=2, max_kv_seq_len=16, max_q_seq_len=1) + layer = LayerAttention(0, 0, None, backend) + value = torch.ones(2) + with session.prepare( + batch, ({batch.phase: backend},), capacity, workspace=lambda _: AttentionWorkspace("cpu"), capture=capture + ): + assert layer.forward(session, value) is value + with pytest.raises(RuntimeError, match="already active"): + with session.prepare(batch, (), capacity, workspace=lambda _: None): + pass + assert session.contexts[0].batch is not None + for operation in (lambda: layer.forward(session, value), session.build_device_metadata): + with pytest.raises(RuntimeError, match="active session"): + operation() + + +@pytest.mark.parametrize("changed", ["capacity", "workspace", "capture"]) +def test_session_rejects_context_substitution(changed): + from lightllm.common.attention.api import AttentionContext, AttentionWorkspace, ContextCapacity + from lightllm.common.attention.session import AttentionSession + + class Backend(StubBackend): + def allocate_context(self, capacity, *, capture, workspace): + if changed == "capacity": + capacity = replace(capacity, max_kv_seq_len=32) + elif changed == "workspace": + workspace = AttentionWorkspace("cpu") + else: + capture = not capture + return AttentionContext(capacity, capture, workspace) + + batch = make_batch() + backend = Backend(make_env(), make_spec()) + capacity = ContextCapacity(phase=batch.phase, batch_size=2, token_num=2, max_kv_seq_len=16, max_q_seq_len=1) + with pytest.raises(ValueError, match="backend changed"): + with AttentionSession().prepare( + batch, ({batch.phase: backend},), capacity, workspace=lambda _: AttentionWorkspace("cpu") + ): + pytest.fail("invalid context entered execution") + + +def test_output_rejects_strided_padding_hole(): + backend = _OutputBackend(make_env(), make_spec()) + context = make_context(backend, prepare=True) + layer = make_layer_attention(backend, context, None) + base = torch.arange(16).view(2, 8) + with pytest.raises(ValueError, match="out"): + run_attention(layer, context, base[:1, 4:8], out=base[:, :4]) + + +@pytest.mark.parametrize("replay", [False, True]) +def test_recurrent_constructor_uses_spec(replay): + from lightllm.common.attention.api.spec import GDNConfig, KDAConfig + from lightllm.common.attention.registry import load_backend_class + + for family, config in ((AttentionFamily.GDN, GDNConfig(6, 5, "silu")), (AttentionFamily.KDA, KDAConfig(5, -3.0))): + mode = "replay" if replay else "triton" + name = f"Triton{family.value.upper()}{'Replay' if replay else ''}DecodeBackend" + cls = load_backend_class(f"backends.{family.value}.{mode}.backend:{name}") + spec = make_spec(family, num_q_heads=3, head_size=16, value_head_size=16, config=config) + spec = replace(spec, cache=replace(spec.cache, key_heads=3, state_dtype=torch.float32)) + backend = cls(make_env(tp_size=8), spec) + if family is AttentionFamily.GDN: + assert (backend.tp_num_k_heads, backend.tp_num_v_heads, backend.tp_key_dim, backend.tp_value_dim) == ( + 3, + 6, + 48, + 96, + ) + assert backend.conv_kernel_dim == 5 + else: + assert (backend.tp_num_heads, backend.tp_hidden_size, backend.conv_kernel_size, backend.lower_bound) == ( + 3, + 48, + 5, + -3.0, + ) diff --git a/unit_tests/common/attention/test_dense_numerics.py b/unit_tests/common/attention/test_dense_numerics.py index 6beb1f8526..6e90a9b1ea 100644 --- a/unit_tests/common/attention/test_dense_numerics.py +++ b/unit_tests/common/attention/test_dense_numerics.py @@ -7,6 +7,7 @@ from lightllm.common.attention.api import AttentionPhase, CacheLayout, ContextCapacity, UnsupportedAttention from lightllm.common.attention.api.inputs import DenseInputs, DenseResources +from lightllm.common.attention.api.spec import MLAConfig from lightllm.common.attention.api.workspace import AttentionWorkspace from unit_tests.common.attention.helpers import dense_class, make_env from lightllm.common.attention_state import AttentionStateManager @@ -123,31 +124,32 @@ def batch_for(lengths): max_kv_seq_len=16, max_q_seq_len=batch.max_q_seq_len, ) - context = AttentionSession.allocate( - backend, capacity, capture=capture, workspace=pool.borrow_for_capture() if capture else pool - ) - plan(backend, context, batch) - attention = make_layer_attention(backend, context, DenseResources(key=key, value=value)) + session = AttentionSession() inputs = DenseInputs(query) output = torch.empty_like(query) + with session.prepare(batch, ({phase: backend},), capacity, workspace=lambda _: pool, capture=capture): + context = session.contexts[0] + attention = make_layer_attention(backend, context, DenseResources(key=key, value=value)) - def run(): - build(backend, context) - return run_attention(attention, context, inputs, out=output) + def run(): + session.build_device_metadata() + return attention.forward(session, inputs, out=output) - actual = run() - assert actual.data_ptr() == output.data_ptr() - torch.testing.assert_close(actual, _reference(query, key, value, (7, 13), q_lengths, window), rtol=0.02, atol=0.015) - if not capture: - return - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - run() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() + actual = run() + assert actual.data_ptr() == output.data_ptr() + torch.testing.assert_close( + actual, _reference(query, key, value, (7, 13), q_lengths, window), rtol=0.02, atol=0.015 + ) + if not capture: + return + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() captured_pointer = context.batch.b_seq_len.data_ptr() manager = AttentionStateManager(request_capacity=2) @@ -157,10 +159,7 @@ def run(): attention_contexts={0: context}, graph_input_copies=lambda source: state.attention.replay_input_copies(source.attention), ) - state.attention = AttentionSession() - state.attention.groups = state.attention_groups - state.attention.contexts = state.attention_contexts - state.attention.batch = context.batch + state.attention = session execution = CapturedGraph(graph, (state,), (output,), runtime=model.device_runtime) try: # Eager execution may grow the shared pool after capture. The graph must @@ -235,6 +234,7 @@ def test_mla_spec_scale(phase, scale, provider): value_head_size=512, cache=CacheLayout("None", 4, torch.bfloat16, 1, 192, 512), softmax_scale=scale, + config=MLAConfig(64, 128), ) prefix = {"flashinfer": "FlashInfer", "fa3": "Fa3"}[provider] cls = load_backend_class(f"backends.mla.{provider}.backend:{prefix}MLA{phase.value.title()}Backend") @@ -259,7 +259,7 @@ def test_mla_spec_scale(phase, scale, provider): capacity = ContextCapacity( phase=phase, batch_size=2, token_num=sum(q_lengths), max_kv_seq_len=16, max_q_seq_len=max(q_lengths) ) - context = AttentionSession.allocate(backend, capacity, capture=False, workspace=workspace) + context = AttentionSession._allocate(backend, capacity, capture=False, workspace=workspace) plan(backend, context, batch) build(backend, context) if phase is AttentionPhase.PREFILL: diff --git a/unit_tests/common/attention/test_linear_layout.py b/unit_tests/common/attention/test_linear_layout.py index 5e65ccd6e0..82ac7ddfa2 100644 --- a/unit_tests/common/attention/test_linear_layout.py +++ b/unit_tests/common/attention/test_linear_layout.py @@ -1,4 +1,4 @@ -from unit_tests.common.attention.helpers import plan, build, run_attention +from unit_tests.common.attention.helpers import plan, build, prepared_session from dataclasses import replace import pytest @@ -41,18 +41,19 @@ def test_shared_backend_selects_each_states_cursor_through_graph_replay(dynamic, def step(): build(backend, context) - a, b = run_attention(layer_a, context, None), run_attention(layer_b, context, None) - assert run_attention(layer_a, context, None) is a + a, b = layer_a.forward(session, None), layer_b.forward(session, None) + assert layer_a.forward(session, None) is a return a, b - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - step() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - a, b = step() + with prepared_session(context) as session: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + step() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + a, b = step() graph.replay() tail = [1] * 4 if dynamic else [] assert a.tolist() == [1, 3] + tail @@ -64,7 +65,8 @@ def step(): assert b.tolist() == [1, 3] + tail # A new eager batch must not reuse the previous batch's gathered counts. plan(backend, context, batch) - a, b = step() + with prepared_session(context) as session: + a, b = step() assert a.tolist() == [2, 1] + tail assert b.tolist() == [1, 3] + tail expected_requests = torch.tensor([0, 2] + ([7] * 4 if dynamic else []), device="cuda") diff --git a/unit_tests/common/attention/test_platform_runtime.py b/unit_tests/common/attention/test_platform_runtime.py index 63ab3bc654..aca2d9556d 100644 --- a/unit_tests/common/attention/test_platform_runtime.py +++ b/unit_tests/common/attention/test_platform_runtime.py @@ -138,7 +138,7 @@ def test_storage_generation_invalidates_graph_before_copy(): execution.close() -@pytest.mark.parametrize("invalid", ["role", "shape"]) +@pytest.mark.parametrize("invalid", ["role", "shape", "missing_hidden", "same_pointer_shape"]) def test_graph_validates_all_lanes_before_copy_or_plan(monkeypatch, invalid): from lightllm.common.attention.api import BackendEnv, ContextCapacity from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -148,7 +148,6 @@ def test_graph_validates_all_lanes_before_copy_or_plan(monkeypatch, invalid): manager = AttentionStateManager(request_capacity=2) env = BackendEnv( runtime=runtime, - network_config={}, hold_req_id=7, draft_step=0, max_seq_len=16, @@ -186,13 +185,17 @@ def make_state(batch, capture): ] if invalid == "role": sources[1].att_state_role = "draft0" + elif invalid == "missing_hidden": + targets[1].mtp_draft_input_hiddens = torch.ones(2, 4) + elif invalid == "same_pointer_shape": + sources[1].input_ids = targets[1].input_ids.view(1, 2) else: sources[1].input_ids = torch.ones(3, dtype=torch.int64) plans.clear() graph = SimpleNamespace(replay=lambda: pytest.fail("submitted invalid lanes"), reset=lambda: None) execution = CapturedGraph(graph, targets, runtime=runtime) try: - with pytest.raises(ValueError, match="model role|shape, dtype or device"): + with pytest.raises(ValueError, match="model role|shape, dtype or device|changed presence"): execution.replay(*sources) assert targets[0].input_ids.tolist() == [0, 0] assert targets[0].attention.batch.b_seq_len.tolist() == [5, 2] diff --git a/unit_tests/common/attention/test_provider_kind.py b/unit_tests/common/attention/test_provider_kind.py index 6fa7d2ea67..f60be664ef 100644 --- a/unit_tests/common/attention/test_provider_kind.py +++ b/unit_tests/common/attention/test_provider_kind.py @@ -146,14 +146,16 @@ def find_spec(self, fullname, path=None, target=None): from lightllm.common.attention.backends.gdn.replay.backend import TritonGDNReplayDecodeBackend from lightllm.platform import DeviceRuntime -env = BackendEnv(runtime=DeviceRuntime("cpu"), network_config={ - "linear_num_value_heads": 4, "linear_num_key_heads": 2, "linear_key_head_dim": 8, - "linear_value_head_dim": 16, "linear_conv_kernel_dim": 4, "hidden_act": "silu", -}, hold_req_id=7, draft_step=0, max_seq_len=64, tp_size=1, alloc_tensor=torch.empty) +from lightllm.common.attention.api.spec import GDNConfig + +env = BackendEnv( + runtime=DeviceRuntime("cpu"), hold_req_id=7, draft_step=0, max_seq_len=64, tp_size=1, alloc_tensor=torch.empty +) for cls in (TritonGDNDecodeBackend, TritonGDNReplayDecodeBackend): cache = CacheLayout("plain", 1, torch.bfloat16, 2, 8, 16, state_mode=cls.state_mode, state_dtype=torch.float32) - spec = LayerSpec(AttentionFamily.GDN, AttentionPhase.DECODE, torch.bfloat16, cache, 2, 8, 16) + spec = LayerSpec(AttentionFamily.GDN, AttentionPhase.DECODE, torch.bfloat16, cache, 2, 8, 16, + config=GDNConfig(4, 4, "silu")) backend = cls(env, spec) assert backend.spec is spec assert not torch.cuda.is_initialized() @@ -185,7 +187,6 @@ def test_quantized_prefill_uses_layout_group(monkeypatch, bits): monkeypatch.setattr(prefill, "context_attention_fwd_contiguous_kv", lambda **kwargs: kwargs["o"].zero_()) env = BackendEnv( runtime=DeviceRuntime("cpu"), - network_config={}, hold_req_id=7, draft_step=0, max_seq_len=64, diff --git a/unit_tests/common/attention/test_state.py b/unit_tests/common/attention/test_state.py index dc6ce0e80e..7daa350cd0 100644 --- a/unit_tests/common/attention/test_state.py +++ b/unit_tests/common/attention/test_state.py @@ -99,6 +99,7 @@ def test_mla_adapter_binds_state_layout(kind): tp_q_head_num_=16, qk_nope_head_dim=512, kv_lora_rank=512, + v_head_dim=128, softmax_scale=latent_dim ** -0.5, network_config_={} if nope else {"qk_rope_head_dim": rope_dim}, ) @@ -114,3 +115,97 @@ def test_mla_adapter_binds_state_layout(kind): assert resources.qk_rope_head_dim == rope_dim assert resources.kv.data_ptr() == state.layer_resources(2).data_ptr() assert resources.kv.shape[-1] == (584 if nope else 656 if kind == "packed_group128_fp8" else latent_dim) + + +@pytest.mark.parametrize("family", ["gdn", "kda"]) +def test_replay_view_has_no_lifecycle_methods(family): + from importlib import import_module + + state_cls = getattr( + import_module(f"lightllm.common.attention_state.bounded_att.{family}.replay_{family}_state"), + f"Replay{family.upper()}State", + ) + config = ConvSSMStateConfig((4, 1), 1, 1, 8, 8, 3, torch.float32, torch.float32) + state = state_cls(config, request_capacity=3, device="cpu") + resources = state.layer_resources(1) + replay = resources.replay_state + assert replay is not state + assert replay.ssm_state is state.ssm_state + assert replay.cursors is state.cursors + assert replay.storage_layer(1) == 1 + assert resources.ssm.data_ptr() == state.ssm_state[1].data_ptr() + assert not any( + hasattr(replay, name) for name in ("finish_decode", "finish_decode_token", "reset_bounded_att_state") + ) + + +@pytest.mark.parametrize("failure", [False, True]) +def test_model_finishes_replay_once_after_all_layers(failure): + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.attention import ModelAttention + from unit_tests.common.attention.helpers import make_batch + + events = [] + coordinator = ModelAttention.__new__(ModelAttention) + coordinator.draft_step, coordinator.dynamic_verify = 0, False + coordinator.replay_states = ( + SimpleNamespace(finish_decode=lambda reqs, cu: events.append(("finish", reqs.tolist(), cu.tolist()))), + ) + state = SimpleNamespace( + attention=SimpleNamespace(batch=make_batch(), build_device_metadata=lambda: None), + input_ids=torch.ones(2, 4), + is_cuda_graph_capture=False, + hidden_collector=SimpleNamespace(add=lambda **kw: None, add_final_hidden=lambda *a: None), + ) + + def layer(index): + def forward(hidden, *args): + events.append(index) + if failure and index == 1: + raise RuntimeError("layer failed") + return hidden + + return SimpleNamespace(token_forward=forward) + + model = SimpleNamespace( + attention=coordinator, + layers_num=2, + layers_infer=[layer(4), layer(1)], + trans_layers_weight=[None, None], + pre_post_weight=None, + pre_infer=SimpleNamespace(token_forward=lambda x, *a: x, _tpsp_sp_split=lambda input, **kw: input), + post_infer=SimpleNamespace(token_forward=lambda x, **kw: x, _tpsp_allgather=lambda input, **kw: input), + _create_model_output=lambda x, state: x, + ) + if failure: + with pytest.raises(RuntimeError, match="layer failed"): + TpPartBaseModel._token_forward(model, state) + assert events == [4, 1] + else: + TpPartBaseModel._token_forward(model, state) + assert events == [4, 1, ("finish", [0, 1], [0, 1, 2])] + + +def test_glm_query_metadata_is_owned_by_model(monkeypatch): + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.models.glm5_next.infer_struct import Glm5NextInferStateInfo + from unit_tests.common.attention.helpers import make_batch + + monkeypatch.setattr(InferStateInfo, "init_some_extra_state", lambda *a: None) + captured, incoming = Glm5NextInferStateInfo(), Glm5NextInferStateInfo() + table = make_batch().req_to_token_indexs + for state, positions in ((captured, [3, 4, 3]), (incoming, [5, 6, 7])): + state.is_prefill = True + state.position_ids = torch.tensor(positions, dtype=torch.int32) + state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) + state.b_q_seq_len = torch.tensor([2, 1], dtype=torch.int32) + state.input_ids = torch.zeros(3, dtype=torch.int64) + state.init_some_extra_state(None) + state.attention.batch = make_batch(req_to_token_indexs=table) + assert captured.sparse_queries.q_kv_seq_len.tolist() == [4, 5, 4] + assert captured.sparse_queries.q_req_idx.tolist() == [0, 0, 1] + pointer = captured.q_kv_seq_len.data_ptr() + for target, source in captured.graph_input_copies(incoming): + target.copy_(source) + assert captured.q_kv_seq_len.data_ptr() == pointer + assert captured.sparse_queries.q_kv_seq_len.tolist() == [6, 7, 8] diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py index ef2014cd3a..3681329e7c 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -1,3 +1,5 @@ +from contextlib import ExitStack +from unit_tests.common.attention.helpers import prepared_session from lightllm.common.attention.api import AttentionFamily, AttentionPhase, CacheEncoding, CacheLayout from unit_tests.common.attention.helpers import make_env import gc @@ -114,8 +116,11 @@ def _case(heads, ratio): model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] model.tp_rank_ = 0 model.tp_world_size_ = 1 + from lightllm.models.deepseek_v4.config import create_attention_config + spec = make_spec( AttentionFamily.SPARSE_MLA, + config=create_attention_config(model.config, (0, 1, 2)), cache=CacheLayout(CacheEncoding.FP8_SPARSE_MLA, 256, torch.uint8, 1, 512, 512), num_q_heads=heads, head_size=512, @@ -247,13 +252,18 @@ def test_decode_cuda_graph_reads_updated_query_and_lengths(monkeypatch, heads, r def decode_func(*states): outputs = tuple( - run_sparse(state.attention.contexts[0], state, state.q, state.packed_kv, state.att_control) + run_sparse( + state.attention.contexts[0], state, state.q, state.packed_kv, state.att_control, session=state.attention + ) for state in states ) return outputs if overlap else outputs[0] # Exercise the framework's warmup + capture directly, with no operator-specific reset. - outputs = runner.capture_decode(decode_func, *(case[0] for case in cases)) + with ExitStack() as scopes: + for state_info, _, _, _ in cases: + scopes.enter_context(prepared_session(state_info.attention.contexts[0], state_info.attention)) + outputs = runner.capture_decode(decode_func, *(case[0] for case in cases)) outputs = outputs if overlap else (outputs,) for state_info, _, _, _ in cases: state = state_info.attention.contexts[0] @@ -384,8 +394,9 @@ def prefill_func(inputs, infer_state): return [att_out + 0.125] infer_state = new_request(1) - layer._context_attention_wrapper_run(q, q_lora, x, infer_state, layer_weight) - outputs = runner.capture_prefill(prefill_func, [q, q_lora, x], infer_state) + with prepared_session(infer_state.attention.contexts[0], infer_state.attention): + layer._context_attention_wrapper_run(q, q_lora, x, infer_state, layer_weight) + outputs = runner.capture_prefill(prefill_func, [q, q_lora, x], infer_state) torch.testing.assert_close( outputs[0], _reference(q, decoded, extra_decoded, infer_state.attn_params, heads) + 0.125, @@ -399,7 +410,8 @@ def prefill_func(inputs, infer_state): infer_state = new_request(length) current_meta = infer_state.attention.contexts[0].sched_meta_by_ratio[ratio] assert current_meta.tile_scheduler_metadata is None - actual = runner.replay([q, q_lora, x], infer_state)[0] + with prepared_session(infer_state.attention.contexts[0], infer_state.attention): + actual = runner.replay([q, q_lora, x], infer_state)[0] assert callback_states[-1] is infer_state assert current_meta.tile_scheduler_metadata.data_ptr() != previous_meta.tile_scheduler_metadata.data_ptr() torch.testing.assert_close( diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py index 79463bd2dc..37bcbd90d6 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -14,7 +14,13 @@ from lightllm.server.core.objs.start_args_type import StartArgs from unit_tests.common.att_resources.model_helpers import rebuild_full_state, replace_logical_layer_groups from unit_tests.common.att_resources.test_capacity import definition -from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, make_layer_attention +from unit_tests.common.attention.helpers import ( + make_backend, + make_batch, + make_context, + make_layer_attention, + prepared_session, +) @pytest.mark.parametrize("backend_name", ["fa3", "flashinfer"]) @@ -98,12 +104,12 @@ def run(q, paged_kv_cache, k_scale, v_scale, out): ) setattr(layer, operation + "_attention", attention) infer = InferStateInfo() - infer.attention.contexts = {0: context} monkeypatch.setattr(manager, "get_state", lambda *args: pytest.fail("forward must not look up storage owners")) q = torch.ones((1, state.config.head_num, state.config.head_dim)) - if operation == "prefill": - output = layer._context_attention_kernel(q=q, kv=None, infer_state=infer, layer_weight=None) - else: - output = layer._token_attention_kernel(q=q, infer_state=infer, layer_weight=None) + with prepared_session(context, infer.attention): + if operation == "prefill": + output = layer._context_attention_kernel(q=q, kv=None, infer_state=infer, layer_weight=None) + else: + output = layer._token_attention_kernel(q=q, infer_state=infer, layer_weight=None) assert torch.equal(output, q) assert calls == [operation] diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index 6d219dc21d..11b58f5646 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -288,8 +288,14 @@ def _test_backend(requests, token_capacity=8192): model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda _: 0) model.swa_prefill_workspace = get_swa_prefill_workspace(model) requests._test_model = model + from lightllm.models.deepseek_v4.config import create_attention_config + return DeepseekV4SparseMLAFp8DecodeBackend( - make_env(model), make_spec(cache=replace(make_spec().cache, page_size=256)) + make_env(model), + make_spec( + cache=replace(make_spec().cache, page_size=256), + config=create_attention_config(model.config, tuple(layer.layer_num_ for layer in model.layers_infer)), + ), ) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py index e538002a17..cb2b390756 100644 --- a/unit_tests/common/test_linear_state.py +++ b/unit_tests/common/test_linear_state.py @@ -310,7 +310,7 @@ def decode(): run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, ) ) - state.finish_decode_token(logical_layer_id=layer, reqs=reqs, cu_seqlens=cu) + state.finish_decode(reqs=reqs, cu_seqlens=cu) return outputs decode() diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py index 95e463fd0b..057b7d3182 100644 --- a/unit_tests/models/deepseek_v4/test_model.py +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -7,6 +7,41 @@ from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel +@pytest.mark.parametrize("draft", [False, True]) +def test_backend_initialization_uses_spec_heads(monkeypatch, draft): + from lightllm.common.attention.api import AttentionFamily, AttentionPhase + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.models.deepseek_v4.attention.sparse_mla import DeepseekV4SparseMLAFp8PrefillBackend + from lightllm.models.deepseek_v4.config import create_attention_config + from unit_tests.common.attention.helpers import make_env, make_spec + + config = dict( + compress_ratios=[0, 4, 128], head_dim=512, qk_rope_head_dim=64, index_head_dim=128, sliding_window=128 + ) + spec = make_spec( + AttentionFamily.SPARSE_MLA, + AttentionPhase.PREFILL, + num_q_heads=16, + config=create_attention_config(config, (0, 1, 2)), + ) + backend = DeepseekV4SparseMLAFp8PrefillBackend(make_env(), spec) + model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + model.args = SimpleNamespace(llm_kv_type="fp8kv_dsa") + model.config, model.data_type = config, torch.bfloat16 + model.is_mtp_draft_model = draft + model.layers_infer = [SimpleNamespace() for _ in range(3)] + model.attention = SimpleNamespace(groups=({AttentionPhase.PREFILL: backend},)) + reservations = [] + model.dsv4_workspace = SimpleNamespace(init_flashmla_prefill_q=lambda **kw: reservations.append(kw)) + monkeypatch.setattr(TpPartBaseModel, "prepare_attention_backends", lambda self: None) + monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_env_start_args", lambda: model.args) + model.prepare_attention_backends() + assert [layer.flashmla_q_head_num_ for layer in model.layers_infer] == [64, 64, 64] + assert len(reservations) == (0 if draft else 1) + if not draft: + assert reservations[0] == dict(tp_q_head_num=16, padded_q_head_num=64, head_dim=512, dtype=torch.bfloat16) + + @pytest.mark.parametrize("capture", [False, True], ids=["eager", "capture"]) @pytest.mark.parametrize("draft", [False, True], ids=["main", "dspark"]) def test_pre_layer_decode_uses_current_graph_flag(monkeypatch: pytest.MonkeyPatch, capture: bool, draft: bool) -> None: diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py index e842f8cc8a..b2992d8de9 100644 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -63,17 +63,7 @@ def test_topk_request_lookup_preserves_strided_tables(dense): torch.testing.assert_close(actual_relative, expected_relative, rtol=0, atol=0) -@pytest.mark.parametrize( - "reference", - [ - "backends.sparse_mla.flashmla.fp8:FlashMLASparseMLAFp8PrefillBackend", - "lightllm.models.glm5_next.attention.sparse_mla:GlmSparseMLAPrefillBackend", - "lightllm.models.glm5_next.attention.fp8_sparse_mla:GlmSparseMLAFp8PrefillBackend", - ], -) -def test_sparse_context_owns_query_metadata(reference): - from lightllm.common.attention.registry import load_backend_class - +def test_model_owns_query_metadata(): lengths = torch.tensor([5, 4], device="cuda", dtype=torch.int32) q_lengths = torch.tensor([2, 1], device="cuda", dtype=torch.int32) batch = AttentionBatch( @@ -94,15 +84,16 @@ def test_sparse_context_owns_query_metadata(reference): b_ready_cache_len=tensor([3, 3]), position_ids=tensor([3, 4, 3]), ) - backend = object.__new__(load_backend_class(reference)) - context = AttentionSession.allocate( - backend, - ContextCapacity(phase=AttentionPhase.PREFILL, batch_size=2, token_num=3, max_kv_seq_len=5, max_q_seq_len=2), - capture=False, - workspace=AttentionWorkspace("cuda"), - ) - plan(backend, context, batch) - build(backend, context) - torch.testing.assert_close(context.q_kv_seq_len, tensor([4, 5, 4])) - torch.testing.assert_close(context.q_req_idx, tensor([0, 0, 1])) - torch.testing.assert_close(context.q_kv_start_loc, tensor([0, 0, 5])) + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.models.glm5_next.infer_struct import Glm5NextInferStateInfo + from unittest.mock import patch + + state = Glm5NextInferStateInfo() + for name, value in vars(batch).items(): + setattr(state, name, value) + state.is_prefill = True + state.input_ids = torch.empty(3, device="cuda", dtype=torch.int64) + with patch.object(InferStateInfo, "init_some_extra_state", lambda *args: None): + state.init_some_extra_state(None) + torch.testing.assert_close(state.sparse_queries.q_kv_seq_len, tensor([4, 5, 4])) + torch.testing.assert_close(state.sparse_queries.q_req_idx, tensor([0, 0, 1])) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index f379d7e3eb..d8c13da66d 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -1,3 +1,5 @@ +from dataclasses import replace +from lightllm.common.attention.api.spec import SparseMLAConfig import weakref from types import SimpleNamespace @@ -784,12 +786,12 @@ def reference(): def test_fp8_sparse_direct_decode_factory_bound_guard(): - backend = object.__new__(GlmSparseMLAFp8DecodeBackend) - backend.network_config = {"index_topk": 2048, "index_kpool": 4} + backend = make_backend(GlmSparseMLAFp8DecodeBackend) + backend.spec = replace(backend.spec, config=SparseMLAConfig(2048, 4)) assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound == 2051 - backend.network_config = {"index_topk": 2048, "index_kpool": 8} + backend.spec = replace(backend.spec, config=SparseMLAConfig(2048, 8)) assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound is None - backend.network_config = {"index_topk": 2047, "index_kpool": 4} + backend.spec = replace(backend.spec, config=SparseMLAConfig(2047, 4)) assert backend_context(type(backend), "decode", settings=backend).decode_topk_bound is None @@ -906,8 +908,8 @@ def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_h indices[3, 64:128] = indices[3, :64] # duplicates assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) control = {"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - backend = object.__new__(GlmSparseMLAFp8DecodeBackend) - backend.network_config = {"index_topk": 2048, "index_kpool": 4} + backend = make_backend(GlmSparseMLAFp8DecodeBackend) + backend.spec = replace(backend.spec, config=SparseMLAConfig(2048, 4)) infer_state = _make_infer_state(is_cuda_graph_capture=False) state = backend_context(type(backend), "decode", settings=backend) state.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 15fc776ae9..3b7a65b26b 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -1,3 +1,4 @@ +from unit_tests.common.attention.helpers import prepared_session from lightllm.common.attention.api import AttentionPhase import dataclasses import json @@ -173,6 +174,7 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) state = InferStateInfo() state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} + state.sparse_queries = None state.att_state_role = "main" state.att_state_manager = SimpleNamespace( get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv)) @@ -194,8 +196,10 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): scores = projected[row].float() @ keys.T * config["qk_nope_head_dim"] ** -0.5 expected.append(scores.softmax(-1) @ keys) expected = torch.stack(expected) - prefill = layer._context_attention_kernel(q, None, state, weight) + with prepared_session(context, state.attention): + prefill = layer._context_attention_kernel(q, None, state, weight) state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} + state.sparse_queries = None context = make_context(backend, make_batch((2, 3, 5), phase=AttentionPhase.DECODE, device="cuda")) context.nsa_cache_seqlens = lengths context.nsa_cu_seqlens_k_new = F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)) @@ -203,7 +207,8 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): layer.decode_attention = make_layer_attention( backend, context, MLAResources(kv, layer.kv_lora_rank, 0), softmax_scale=layer.softmax_scale ) - decode = layer._token_attention_kernel(q, state, weight) + with prepared_session(context, state.attention): + decode = layer._token_attention_kernel(q, state, weight) for output in (prefill, decode): torch.testing.assert_close(output.float(), expected, atol=1e-3, rtol=0.015) @@ -457,6 +462,7 @@ def test_native_drafts_share_weights_and_prepare_their_own_state(config): state.input_ids = torch.zeros(3, device="cuda", dtype=torch.int64) state.b_ready_cache_len = torch.tensor([2], device="cuda", dtype=torch.int32) state.b_seq_len = torch.tensor([5], device="cuda", dtype=torch.int32) + state.b_req_idx = torch.tensor([0], device="cuda", dtype=torch.int32) state.init_some_extra_state(model) torch.testing.assert_close( state.position_ids, torch.arange(2, 5, device="cuda", dtype=state.position_ids.dtype) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 5843fcae4a..b9bea8691a 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -262,10 +262,20 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mod a, bias, mtp_state_index=selected.mtp_state_index, - replay_state=selected if mode == "replay" else None, + replay_state=selected.replay_storage() if mode == "replay" else None, ), ) - actual = run_attention(attention, context, KDAInputs(mixed, gate, beta)) + from lightllm.common.basemodel.attention import ModelAttention + from unit_tests.common.attention.helpers import prepared_session + + coordinator = ModelAttention.__new__(ModelAttention) + coordinator.replay_states = (selected,) if mode == "replay" else () + coordinator.draft_step, coordinator.dynamic_verify = mtp_step, dynamic + with prepared_session(context, infer.attention): + actual = attention.forward(infer.attention, KDAInputs(mixed, gate, beta)) + if mode == "replay": + assert selected.cursors.eq(0).all() + coordinator.finish_decode(infer) torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) if mode == "native": checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] diff --git a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py index f3545a68a4..c92762bf25 100644 --- a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py +++ b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py @@ -132,8 +132,6 @@ def test_forward_uses_common_prefill(monkeypatch, tokens, mode): layer._project_sparse_qkv_a = lambda hidden, *_: (hidden, None, hidden) layer._normalize_sparse_kv = lambda kv, _: kv layer.indexer = SimpleNamespace(write_k_pools=Mock()) - context = object() - layer.prefill_attention = SimpleNamespace(context=Mock(return_value=context)) writer = Mock() layer._post_cache_kv = writer model.layers_infer = [layer] @@ -142,7 +140,13 @@ def test_forward_uses_common_prefill(monkeypatch, tokens, mode): monkeypatch.setattr(basemodel.dist_group_manager, "get_group", Mock()) monkeypatch.setattr(basemodel.dist_group_manager, "clear_deepep_buffer", Mock()) monkeypatch.setattr(ModelInput, "to_cuda", lambda self: self.check_input()) - monkeypatch.setattr(InferStateInfo, "init_some_extra_state", Mock()) + + def init_metadata(state, _model): + state.attention = SimpleNamespace(build_device_metadata=Mock()) + state.q_kv_seq_len = torch.ones_like(state.input_ids, dtype=torch.int32) + state.q_req_idx = torch.zeros_like(state.input_ids, dtype=torch.int32) + + monkeypatch.setattr(model.infer_state_class, "init_some_extra_state", init_metadata) event = Mock() monkeypatch.setattr(torch.cuda, "Event", Mock(return_value=event)) output = model.forward(model_input) @@ -169,8 +173,9 @@ def test_forward_uses_common_prefill(monkeypatch, tokens, mode): post_forward.assert_not_called() allocator.cache_env_in.assert_called_once() allocator.cache_env_out.assert_called_once() - layer.prefill_attention.context.assert_called_once_with(infer.attention) - assert layer.indexer.write_k_pools.call_args.args[2] is context + metadata = layer.indexer.write_k_pools.call_args.args[2] + assert metadata.q_kv_seq_len is infer.q_kv_seq_len + assert metadata.q_req_idx is infer.q_req_idx @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -259,6 +264,7 @@ def norm(input, eps, **_): max_q_seq_len=max(q_lengths), max_kv_seq_len=2048, attention=AttentionSession(), + sparse_queries=context, req_manager=SimpleNamespace(req_to_token_indexs=None), ) # Top-k outputs are discarded here; run the normal projections and cache writers. @@ -267,7 +273,7 @@ def norm(input, eps, **_): _, kv = layer._get_qkv(layer._att_norm(hidden, infer, weights), infer, weights) layer._post_cache_kv(kv, infer, weights) layer.indexer._get_indices( - **infer.get_topk_indices_params, infer_state=infer, att_state=context, layer_weight=weights + **infer.get_topk_indices_params, infer_state=infer, query_metadata=context, layer_weight=weights ) expected = [s.clone() for s in stores] assert not torch.equal(initial[2], expected[2]) diff --git a/unit_tests/models/glm5_next/test_prefill_workspace.py b/unit_tests/models/glm5_next/test_prefill_workspace.py index a1d7bc5e20..fe1d63504f 100644 --- a/unit_tests/models/glm5_next/test_prefill_workspace.py +++ b/unit_tests/models/glm5_next/test_prefill_workspace.py @@ -1,3 +1,5 @@ +from unit_tests.common.attention.helpers import make_spec +from lightllm.common.attention.api.spec import SparseMLAConfig from lightllm.common.attention.session import AttentionSession import pytest import torch @@ -12,9 +14,9 @@ def test_fp8_prefill_context_is_owned_by_the_backend(): backend = object.__new__(GlmSparseMLAFp8PrefillBackend) - backend.network_config = {"index_topk": 2048, "index_kpool": 4} + backend.spec = make_spec(config=SparseMLAConfig(2048, 4)) workspace = AttentionWorkspace("cpu") - context = AttentionSession.allocate( + context = AttentionSession._allocate( backend, ContextCapacity(phase=AttentionPhase.PREFILL, batch_size=2, token_num=4, max_kv_seq_len=32, max_q_seq_len=4), capture=False, @@ -39,7 +41,7 @@ def test_fp8_workspace_reserved_per_microbatch_before_execution(decode_only): from lightllm.platform import get_device_runtime backend = object.__new__(GlmSparseMLAFp8PrefillBackend) - backend.network_config = {"index_topk": 2048, "index_kpool": 4} + backend.spec = make_spec(config=SparseMLAConfig(2048, 4)) backend.prefill_token_capacity = 32 backend.prefill_graph_enabled = True backend.decode_only = decode_only @@ -54,7 +56,7 @@ def test_fp8_workspace_reserved_per_microbatch_before_execution(decode_only): assert gathered.numel() == (32 if decode_only else 2 * 64 + 32) * 512 assert gathered.data_ptr() != second.scratch("glm_gathered_kv", 1, torch.bfloat16).data_ptr() assert first.scratch("glm_chunk_indices", 1, torch.int32).numel() == 32 * 2176 - context = AttentionSession.allocate( + context = AttentionSession._allocate( backend, ContextCapacity(phase=AttentionPhase.PREFILL, batch_size=2, token_num=4, max_kv_seq_len=64, max_q_seq_len=4), capture=True, diff --git a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py index 7cee8bf66f..376f3c6868 100644 --- a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py +++ b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py @@ -20,13 +20,12 @@ def make_projection(layer, heads: int, inner: int, cols: int, dtype: torch.dtype layer.layer_num_ = 0 layer.softmax_scale = inner ** -0.5 layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (None, None)) - layer.prefill_attention = SimpleNamespace( - context=lambda session: None, forward=lambda session, inputs: inputs.query - ) + layer.prefill_attention = SimpleNamespace(forward=lambda session, inputs: inputs.query) state = SimpleNamespace( att_state_role="main", att_state_manager=Mock(), attention=object(), + sparse_queries=None, ) weight = SimpleNamespace(k_b_proj_=projection) diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py index d865d829e0..c1431e7b7c 100644 --- a/unit_tests/models/glm5_next/test_visual.py +++ b/unit_tests/models/glm5_next/test_visual.py @@ -199,19 +199,13 @@ def test_target_and_mtp_prefill_embed_images_before_hidden_fusion(monkeypatch): expanded = Glm5NextPreLayerInfer(config).context_forward(ids, state, weight) torch.testing.assert_close(expanded.reshape(4, 4, hidden), expected[:, None].expand(-1, 4, -1)) - def norm(input, eps, out): - out.copy_(input * torch.rsqrt(input.square().mean(-1, keepdim=True) + eps)) - return out - state.mtp_draft_input_hiddens = torch.randn(4, hidden, device="cuda") main_norm_weight = torch.randn(hidden, device="cuda") old_hidden = state.mtp_draft_input_hiddens.clone() projection = torch.randn(2 * hidden, hidden, device="cuda") * 0.1 - weight.enorm_weight_ = weight.hnorm_weight_ = norm - weight.main_norm_weight_ = lambda input, eps, out: out.copy_( - torch.nn.functional.rms_norm(input, (hidden,), main_norm_weight, eps) - ) - weight.eh_proj_weight_ = SimpleNamespace(mm=lambda x: x @ projection) + weight.enorm_weight_ = weight.hnorm_weight_ = SimpleNamespace(weight=torch.ones(hidden, device="cuda")) + weight.main_norm_weight_ = SimpleNamespace(weight=main_norm_weight) + weight.eh_proj_weight_ = SimpleNamespace(mm=lambda input_tensor: input_tensor @ projection) actual = Glm5NextMTPPreLayerInfer(config).context_forward(ids, state, weight) old_hidden = torch.nn.functional.rms_norm(old_hidden, (hidden,), main_norm_weight, 1e-5) normalized = [x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) for x in (expected, old_hidden)] diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py index b81032e4c1..89cd81353b 100644 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -6,7 +6,8 @@ import pytest import torch -from lightllm.common.attention.api import AttentionPhase +from lightllm.common.attention.api import AttentionPhase, AttentionFamily +from lightllm.common.attention.api.spec import GDNConfig from lightllm.common.attention.api.inputs import GDNInputs, GDNResources from lightllm.common.attention.backends.gdn.triton.backend import TritonGDNDecodeBackend, TritonGDNPrefillBackend from lightllm.common.attention_state import AttentionStateManager @@ -47,7 +48,18 @@ def test_gdn_decode_keeps_logical_layer_state( model = model_stub(mtp_step=mtp_step, linear_att_ssm_data_type="float32") model.config = config model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda _: mtp_step) - backend = TritonGDNDecodeBackend(make_env(model), make_spec(use_mtp=mtp_step > 0)) + backend = TritonGDNDecodeBackend( + make_env(model), + make_spec( + AttentionFamily.GDN, + use_mtp=mtp_step > 0, + num_q_heads=1, + head_size=32, + value_head_size=32, + cache=replace(make_spec().cache, key_heads=1, state_dtype=torch.float32), + config=GDNConfig(2, 4, "silu"), + ), + ) dtype = torch.bfloat16 if device == "cuda" else torch.float32 layout = ConvSSMStateConfig(layers, 1, 2, 32, 32, 3, dtype, torch.float32) resource = GDNState(layout, request_capacity=3, mtp_step=mtp_step, device=device) @@ -125,7 +137,7 @@ def recurrent(*, initial_state, v, **kwargs): b_q_seq_len=torch.tensor([5], device=device), b1_cu_q_seq_len=torch.tensor([0, 5], device=device, dtype=torch.int32), ) - prefill = TritonGDNPrefillBackend(make_env(model), make_spec(phase=AttentionPhase.PREFILL)) + prefill = TritonGDNPrefillBackend(make_env(model), replace(backend.spec, phase=AttentionPhase.PREFILL)) context = make_context(prefill, batch) build(prefill, context) attention = replace(make_layer_attention(prefill, context, params), logical_layer_id=logical_layer)