diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 7edb0c9701..bab703bf34 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -127,7 +127,21 @@ class TpPartBaseModel: | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | -全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先按各模型 `get_req_manager_max_seq_len()` 的最大值创建一份 `ReqManager`，通过 `bind_req_manager()` 绑定给全部模型，再调用各模型的 `prepare_attention_states()` 准备自身 State；随后管理器依次完成 `finish_state_registration()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。资源准备和运行时初始化分别通过这两个单模型公共入口完成，均不接收 draft 列表。 +全部 main/draft 模型构造完成后，`ModeBackend.init_model` 按各模型 `get_req_manager_max_seq_len()` 的最大值创建一份 `ReqManager`，通过 `bind_req_manager()` 绑定给全部模型，再调用各模型的 `prepare_attention_states()` 注册自身 State。管理器完成 `finish_state_registration()` 后，逐模型调用 `prepare_attention_backends()`：按 prefill/decode 的 `LayerSpec` 分组，检查设备、存储布局和执行模式，并运行后端数值探针。随后调用 `reserve_attention_resources(model)` 预留各执行槽位的 workspace，再加载标定并执行 `initialize_storage()`。预留必须在容量估算之前完成，使后端临时显存计入可用容量。 + +按 main、各 draft 的顺序调用 `init_runtime()`，绑定 HOLD 页并调用 `initialize_layer_attention()` 创建各层的 `LayerAttentionBinding`，然后执行 warmup、padded-request 与 CUDA Graph 初始化；main 还在末尾执行最大长度检查。全部模型完成后才发布就绪标志。每批次的 metadata 和 provider wrapper 由 context 持有，按 `allocate_context → prepare → update → forward` 执行；KV/Conv/SSM 的存储与恢复仍由 State 管理。 + +层内显式调用 `.forward(...)`。`LayerAttentionBinding` 保存本层的参数和共享 backend；context 来自当前批次： + +```python +attention = self.decode_attention +context = infer_state.attention_contexts[attention.group_id] +output = attention.forward(context, DenseInputs(q)) +``` + +视觉/音频编码器直接传入完整 Q/K/V 时，使用 `api.varlen.VarlenAttentionOp` 的 `forward(inputs, out=...)`，无需创建层 binding 或执行 context。`VarlenAttentionInputs` 描述非因果 self-attention 的变长序列；调用方提供输出缓冲区，算子不保留批次 tensor。通过 `select_varlen_attention_op()` 选择实现，只有声明 `supports_graph` 的实现可参与整段 Graph 捕获。启动参数 `--vit_att_backend` 保持不变。 + +新增模型通过层上的 `attention_family`、`attention_window(phase)`、`attention_uses_alibi`、`attention_uses_sink` 和必要时的 `create_attention_params()` 描述计算需求与资源视图；模型的 `attention_backend_candidates()` 可指定候选实现。dense/MLA/sparse MLA 使用 `--llm_prefill_att_backend` / `--llm_decode_att_backend`，GDN/KDA 使用单独的 `--llm_prefill_linear_att_backend` / `--llm_decode_linear_att_backend`。显式指定不支持的后端会报错；`auto` 才会继续尝试候选。 公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。两个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index cd71541598..f54a7532c3 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -127,7 +127,21 @@ Commonly used interfaces that need to be inherited and implemented | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | -After all main/draft models are constructed, `ModeBackend.init_model` creates one `ReqManager` using the largest `get_req_manager_max_seq_len()` returned by the models, binds it to every model through `bind_req_manager()`, then calls each model's `prepare_attention_states()`, then calls `finish_state_registration()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization each use a single-model entry point; neither takes a draft-model list. +After all main/draft models are constructed, `ModeBackend.init_model` creates one `ReqManager` using the largest `get_req_manager_max_seq_len()`, binds it to every model through `bind_req_manager()`, and calls each model's `prepare_attention_states()` to register its States. After `finish_state_registration()`, it calls `prepare_attention_backends()` for each model: layers are grouped by their prefill/decode `LayerSpec`, device/layout/execution requirements are checked, and backend numerical probes run. It then calls `reserve_attention_resources(model)` to reserve workspace for each execution slot before loading calibration and calling `initialize_storage()`. Reserving before capacity estimation accounts for backend scratch memory. + +`init_runtime()` runs sequentially on the main model and each draft, binding the HOLD page and creating each `LayerAttentionBinding` through `initialize_layer_attention()` before warmup, padded-request setup and CUDA Graph initialization. The main model also checks maximum-length inference. Readiness is published only after every model succeeds. Each batch owns a context containing metadata and provider wrappers, executed through `allocate_context → prepare → update → forward`; States continue to own KV/Conv/SSM storage and restoration. + +Layers call `.forward(...)` explicitly. `LayerAttentionBinding` holds the layer parameters and shared backend; the context comes from the current batch: + +```python +attention = self.decode_attention +context = infer_state.attention_contexts[attention.group_id] +output = attention.forward(context, DenseInputs(q)) +``` + +Vision/audio encoders that provide complete Q/K/V use `api.varlen.VarlenAttentionOp.forward(inputs, out=...)` without a layer binding or execution context. `VarlenAttentionInputs` describes variable-length, non-causal self-attention. The caller provides the output buffer, and the operator does not retain batch tensors. Select implementations with `select_varlen_attention_op()`; only implementations declaring `supports_graph` support capture of the full operation. The `--vit_att_backend` option is unchanged. + +New model layers describe computation and resource views through `attention_family`, `attention_window(phase)`, `attention_uses_alibi`, `attention_uses_sink`, and, when needed, `create_attention_params()`. The model's `attention_backend_candidates()` can specify candidate implementations. Dense/MLA/sparse MLA use `--llm_prefill_att_backend` / `--llm_decode_att_backend`; GDN/KDA use the separate `--llm_prefill_linear_att_backend` / `--llm_decode_linear_att_backend` options. An unsupported explicit backend raises an error; only `auto` tries the remaining candidates. The public `prepare_attention_states()` stage creates an empty manager for main and shares it with drafts. The two State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded factories also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none. 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..709bf2c792 --- /dev/null +++ b/lightllm/common/attention/api/__init__.py @@ -0,0 +1,19 @@ +"""模型层与 attention 实现之间的契约。 + +初始化时，LayerSpec / CacheLayout 用于选择后端和划分可共享 context 的层组， +LayerAttentionBinding 保存本层参数和共享 backend，通过 forward(context, inputs) 执行。 +运行时，AttentionBatch 描述整个批次，Inputs 提供当前层的 tensor 输入； +backend 专用 context 保存派生 metadata 和 provider wrapper，由 InferStateInfo 持有。 +KV / Conv / SSM 的分配、回收和迁移仍由 AttentionStateManager 管理。 +无持久 KV 的视觉/音频编码器使用 VarlenAttentionOp，直接传入 Q/K/V 和序列边界。 + +初始化集成见 common.basemodel.attention，运行时入口见 InferStateInfo.prepare_attention。 +""" + +from .backend import AttentionBackend, LayerAttentionBinding, UnsupportedAttention +from .batch import AttentionBatch, ContextCapacity +from .execution import BackendResources, ExecutionRequirements +from .graph import GraphBinding +from .inputs import VarlenAttentionInputs +from .spec import AttentionFamily, AttentionPhase, CacheLayout, LayerSpec +from .varlen import VarlenAttentionOp diff --git a/lightllm/common/attention/api/backend.py b/lightllm/common/attention/api/backend.py new file mode 100644 index 0000000000..99cdbef960 --- /dev/null +++ b/lightllm/common/attention/api/backend.py @@ -0,0 +1,91 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Callable, Generic, TypeVar + +import torch + +from .batch import AttentionBatch, ContextCapacity +from .execution import BackendResources +from .spec import LayerSpec + +Context = TypeVar("Context") +Inputs = TypeVar("Inputs") +Output = TypeVar("Output") +Params = TypeVar("Params") + + +class UnsupportedAttention(ValueError): + """当前环境或输入不满足此 attention 实现的要求。""" + + +@dataclass(frozen=True) +class LayerAttentionBinding(Generic[Params]): + """一个逻辑层在 prefill 或 decode 阶段的 attention 执行入口。 + + 权重和存储准备完成后构造；backend 是同组层共享的实现，params 借用本层的 + 权重和存储视图，存储重分配后须重新构造。group_id 用于取得当前批次的 context； + logical_layer_id 是模型层号，不一定等于物理存储下标。 + """ + + group_id: int + logical_layer_id: int + spec: LayerSpec + params: Params + backend: "AttentionBackend" + alloc_tensor: Callable + + def forward(self, context, inputs, *, out=None): + return self.backend.forward(self, context, inputs, out=out) + + +class AttentionBackend(ABC, Generic[Context, Inputs, Output]): + """同组层共用的执行实现；每批次变化的 metadata 和 wrapper 放在 context 中。 + + Context 的具体类型由 backend 定义，实例由 InferStateInfo 保存，不能跨 backend + 或并发 microbatch 共用。调用顺序为 allocate_context -> prepare -> update， + 随后同组各层用各自的 LayerAttentionBinding / Inputs 调用 forward。 + """ + + @abstractmethod + def allocate_context(self, capacity: ContextCapacity, *, capture: bool, workspace) -> Context: + """创建归属此 backend 的 context，workspace 由模型按 group / microbatch slot 提供。 + + capture=True 的 context 随 decode CUDA Graph 保存，在后续 replay 中复用。 + """ + pass + + @abstractmethod + def prepare(self, context: Context, batch: AttentionBatch) -> None: + """在图外绑定批次输入并完成 provider planning，这里允许 CPU 侧准备工作。 + + decode replay 时仍会调用；图中引用的 buffer 必须保持地址，只更新内容。 + """ + pass + + @abstractmethod + def update(self, context: Context) -> None: + """在各层 forward 前生成 GPU 侧派生 metadata。 + + decode Graph 会捕获这段 GPU 工作；不能依赖 replay 时重新执行 Python， + 也不能把当前批次的 GPU 值读回 CPU 来决定执行路径。 + """ + pass + + @abstractmethod + def forward( + self, layer: LayerAttentionBinding, context: Context, inputs: Inputs, *, out: torch.Tensor | None = None + ) -> Output: + """使用已准备的 context 执行一层，允许写入 params 借用的 KV / Conv / SSM 存储。 + + out 指定 attention 结果的目标缓冲区；返回 tuple 时只对应首项， + 其余项保留模型所需的附加输出，例如 GDN 的 gate 输入。 + """ + pass + + def resource_requirements(self, capacity: ContextCapacity, *, capture: bool): + """返回分配前可计算的资源上界，供共享容量规划和捕获前预留使用。""" + return BackendResources() + + def bind_graph(self, context: Context): + """为此捕获实例创建可选的 GraphBinding；普通原生 Graph 无需额外绑定。""" + return None 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..09dd58bdbd --- /dev/null +++ b/lightllm/common/attention/api/execution.py @@ -0,0 +1,56 @@ +from dataclasses import dataclass + +from .spec import AttentionPhase + + +@dataclass(frozen=True) +class ExecutionRequirements: + """启动时确定的执行约束，与 LayerSpec 的计算语义分开保存。""" + + graph_phases: frozenset[AttentionPhase] = frozenset() + microbatch_count: int = 1 + max_batch_size: int = 1 + max_seq_len: int = 1 + max_verify_tokens: int = 1 + tp_world_size: int = 1 + expert_parallel: bool = False + transfer_paths: frozenset[str] = frozenset() + + @classmethod + def from_model(cls, model): + args = model.args + phases = set() + if not args.disable_cudagraph: + phases.add(AttentionPhase.DECODE) + if args.enable_prefill_cudagraph: + phases.add(AttentionPhase.PREFILL) + transfers = set() + if args.enable_cpu_cache: + transfers.add("cpu") + if args.enable_dp_prompt_cache_fetch: + transfers.add("dp") + if args.run_mode in ("prefill", "decode"): + transfers.add("pd") + return cls( + frozenset(phases), + 2 if args.enable_decode_microbatch_overlap or args.enable_prefill_microbatch_overlap else 1, + getattr(model, "max_req_num", 1), + max(getattr(model, "max_seq_length", 1), getattr(model, "graph_max_len_in_batch", 1)), + model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model) + 1, + model.tp_world_size_, + args.enable_ep_moe, + frozenset(transfers), + ) + + +@dataclass(frozen=True) +class BackendResources: + """捕获前的资源上界；workspace 项为 (名称, 元素数, dtype)。 + + 同一执行槽位按名称取最大容量，并发槽位分别分配。主机长度快照只在需要 + 它的实现中创建；持久存储的初始化由 State 执行，不能在 forward 中反复清零。 + """ + + workspace: tuple = () + needs_host_lengths: bool = False + 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..2b30def89e --- /dev/null +++ b/lightllm/common/attention/api/graph.py @@ -0,0 +1,20 @@ +class GraphBinding: + """一个捕获实例的 provider 资源，不能挂在跨批次共享的 Backend 上。 + + after_submit 在原生 Graph 提交后立即调用，不能先等待 Graph 完成；某些设备 + 的 Graph 正在等待此方法提交的任务更新。返回的事件须覆盖所有辅助流工作。 + wait_for_update 只等待上次对主机 metadata 的使用，保证 prepare 可以安全改写。 + close 在相关设备工作完成后、原生 Graph 销毁前调用，并须可重复执行。 + """ + + def wait_for_update(self): + pass + + def before_replay(self): + pass + + def after_submit(self): + return () + + def close(self): + pass diff --git a/lightllm/common/attention/api/inputs.py b/lightllm/common/attention/api/inputs.py new file mode 100644 index 0000000000..2b7c2221d7 --- /dev/null +++ b/lightllm/common/attention/api/inputs.py @@ -0,0 +1,115 @@ +"""区分初始化时绑定的层资源与每次 forward 的 tensor 输入。 + +*Params 随 LayerAttentionBinding 保存，借用权重以及 State 管理的 KV / Conv / SSM 视图； +*Inputs 由模型层为当前调用构造，批次级请求索引和 provider 派生 metadata 则由 context 提供。 +这些数据结构只固定字段绑定，不禁止 backend 原地更新所借用的状态 tensor。 +""" + +from dataclasses import dataclass +from typing import NamedTuple, Optional + +import torch + + +@dataclass(frozen=True) +class DenseParams: + key: torch.Tensor + value: torch.Tensor + alibi: Optional[torch.Tensor] = None + sliding_window: Optional[tuple[int, int]] = 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 + softmax_scale: float + kv_lora_rank: int + qk_rope_head_dim: int + + +@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 + state: object = None + + +@dataclass(frozen=True) +class KDAParams: + conv: torch.Tensor + ssm: torch.Tensor + conv_weight: torch.Tensor + a_log: torch.Tensor + dt_bias: torch.Tensor + state: object = 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 + + +@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 diff --git a/lightllm/common/attention/api/spec.py b/lightllm/common/attention/api/spec.py new file mode 100644 index 0000000000..c6b13c8bf1 --- /dev/null +++ b/lightllm/common/attention/api/spec.py @@ -0,0 +1,74 @@ +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 = "" + + 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..953e812896 --- /dev/null +++ b/lightllm/common/attention/api/varlen.py @@ -0,0 +1,32 @@ +from abc import ABC, abstractmethod + +import torch + +from .inputs import VarlenAttentionInputs + + +class VarlenAttentionOp(ABC): + """无持久 KV 的双向 dense attention，供视觉和音频编码器使用。 + + Q/K/V 与变长边界由调用方直接传入，不需要文本推理的请求表或 cache context。 + 实现不保留批次 tensor；调用方持有 out，返回值必须引用同一块输出存储。 + 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: + """初始化或探针校验输入；不支持时抛出 UnsupportedAttention。""" + pass + + @abstractmethod + def forward(self, inputs: VarlenAttentionInputs, *, out: torch.Tensor) -> torch.Tensor: + pass 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/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 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 b5b6cfc369..d032413673 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/attention/backends/_workspace.py b/lightllm/common/attention/backends/_workspace.py new file mode 100644 index 0000000000..cdf58571ce --- /dev/null +++ b/lightllm/common/attention/backends/_workspace.py @@ -0,0 +1,77 @@ +import torch + +from lightllm.platform import get_device_runtime + + +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 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 name, size, dtype in resources.workspace: + self.get(name, size, 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/base.py b/lightllm/common/attention/backends/base.py new file mode 100644 index 0000000000..f9ee8d7d1c --- /dev/null +++ b/lightllm/common/attention/backends/base.py @@ -0,0 +1,182 @@ +import copy +from dataclasses import dataclass, replace +from types import MappingProxyType + +import torch + +from lightllm.common.attention._environment import check_attention_environment +from lightllm.common.attention.api.backend import AttentionBackend, 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 + + +@dataclass +class AttentionExecutionContext: + capacity: ContextCapacity + capture: bool + workspace: object + batch: AttentionBatch | None = None + current_max_kv_seq_len: int = 0 + backend_identity: object = None + + +class BaseAttentionBackend(AttentionBackend): + requires_cpu_cache_sync = False + family = AttentionFamily.DENSE + provider = "triton" + cache_formats = frozenset({"None"}) + supports_attention_sink = 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): + self.runtime = model.device_runtime + self.execution = model.attention_execution + self.page_size = model.args.page_size + self.network_config = MappingProxyType(dict(model.config)) + self.is_draft_model = model.is_mtp_draft_model + self.decode_draft_step = model.mtp_manager.get_decode_draft_step(model.is_mtp_draft_model) + self.dynamic_verify = not self.is_draft_model and self.decode_draft_step > 0 and model.args.mtp_dynamic_verify + self.causal = not (self.is_draft_model and model.args.mtp_mode in ("dspark", "dflash", "dflash2")) + self.hold_req_id = model.req_manager.HOLD_REQUEST_ID + self.mtp_state_index = None + + def for_spec(self, spec): + backend = copy.copy(self) + backend.spec = spec + return backend + + @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.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) -> AttentionExecutionContext: + 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 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: AttentionExecutionContext, 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): + pass + + def _prepare_decode(self, context): + pass + + def update(self, context: AttentionExecutionContext) -> 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): + pass + + def _update_decode(self, context): + pass + + def forward(self, layer: LayerAttentionBinding, context: AttentionExecutionContext, inputs, *, out=None): + 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.data_ptr() == out.data_ptr(): + return result + out.copy_(attention.view(out.shape)) + + return (out, *result[1:]) if isinstance(result, tuple) else out + + def uses_dynamic_spec_verify_layout(self): + return self.dynamic_verify + + def uses_causal_attention(self): + return self.causal 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..d4e3f98313 --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/backend.py @@ -0,0 +1,229 @@ +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 +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 +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): + 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}) + supports_attention_sink = True + supports_sliding_window = True + supports_noncausal_decode = True + max_head_size = 256 + cache_formats = frozenset({"None"}) + """Common fixed page-table storage for FA3 attention backends.""" + + def __init__(self, model): + super().__init__(model=model) + self._init_infer_page_size() + + def _init_infer_page_size(self): + self.infer_page_size = self.page_size + assert ( + self.page_size % self.infer_page_size == 0 + ), f"model page_size {self.page_size} must be divisible by infer_page_size {self.infer_page_size}" + + 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_nomarl_prefill_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _prefill_nomarl_prefill_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs, alloc_func=torch.empty + ) -> torch.Tensor: + if layer.params.sliding_window is not None: + window_size = layer.params.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._decode_normal_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _decode_normal_decode_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs, alloc_func=torch.empty + ): + if layer.params.sliding_window is not None: + window_size = layer.params.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..f099c9ca8f --- /dev/null +++ b/lightllm/common/attention/backends/dense/fa3/fp8.py @@ -0,0 +1,160 @@ +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 +class Fp8Fa3PrefillContext(Fa3PrefillContext): + pass + + +@dataclasses.dataclass +class Fp8Fa3DecodeContext(Fa3DecodeContext): + pass + + +class Fp8Fa3AttBackend(Fa3AttBackend): + provider = "fa3" + supports_sliding_window = False + cache_formats = frozenset({"fp8kv_sph"}) + + def _init_infer_page_size(self): + self.infer_page_size = 1 + assert ( + self.page_size % self.infer_page_size == 0 + ), f"model page_size {self.page_size} 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.page_size}; " + "multi-token FP8 inference pages are not implemented yet." + ) + + 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._prefill_fp8_prefill_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _prefill_fp8_prefill_att( + 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._decode_fp8_decode_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _decode_fp8_decode_att( + 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..7ab56190e0 --- /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 +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/basemodel/attention_vit/fa3/fp.py b/lightllm/common/attention/backends/dense/fa3/varlen.py similarity index 57% rename from lightllm/common/basemodel/attention_vit/fa3/fp.py rename to lightllm/common/attention/backends/dense/fa3/varlen.py index a8e56417c5..513cc090de 100644 --- a/lightllm/common/basemodel/attention_vit/fa3/fp.py +++ b/lightllm/common/attention/backends/dense/fa3/varlen.py @@ -1,30 +1,39 @@ -import dataclasses import torch -from lightllm.common.basemodel.attention_vit.base_att import BaseVitAttBackend +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp -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] +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): + 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) - torch.ops.sgl_kernel.fwd.default( + self._op( q, k, v, None, # k_new None, # v_new None, # qv - o, # out + out, cu_seqlens, cu_seqlens, None, # cu_seqlens_k_new @@ -57,4 +66,4 @@ def _vit_att_fwd( only_qv=False, ) - return o + 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..5cdd70ba17 --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/backend.py @@ -0,0 +1,226 @@ +import dataclasses + +import torch +import triton + +from lightllm.common.attention.api import BackendResources +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.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 + + +@dataclasses.dataclass +class FlashInferPrefillContext(AttentionExecutionContext): + prefill_wrapper: object = None + + +@dataclasses.dataclass +class FlashInferDecodeContext(AttentionExecutionContext): + kv_last_page_len_buffer: torch.Tensor = None + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + decode_wrapper: object = None + + +class FlashInferAttBackend(BaseAttentionBackend): + 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): + 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.q_data_type = model.data_type + self.kv_data_type = model.data_type + + def for_spec(self, spec): + backend = super().for_spec(spec) + backend.tp_q_head_num = spec.num_q_heads + backend.tp_kv_head_num = spec.cache.key_heads + backend.head_dim = spec.head_size + return backend + + def resource_requirements(self, capacity, *, capture): + return BackendResources(((self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),)) + + def _init_infer_page_size(self): + self.infer_page_size = self.page_size + assert ( + self.page_size % self.infer_page_size == 0 + ), f"model page_size {self.page_size} must be divisible by infer_page_size {self.infer_page_size}" + + 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_nomarl_prefill_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_nomarl_prefill_att( + 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="cuda") + 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 _decode_refresh_cuda_graph_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._decode_normal_decode_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _decode_normal_decode_att( + 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._decode_refresh_cuda_graph_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..e20e971010 --- /dev/null +++ b/lightllm/common/attention/backends/dense/flashinfer/fp8.py @@ -0,0 +1,105 @@ +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 +class Fp8FlashInferPrefillContext(FlashInferPrefillContext): + pass + + +@dataclasses.dataclass +class Fp8FlashInferDecodeContext(FlashInferDecodeContext): + pass + + +class Fp8FlashInferAttBackend(FlashInferAttBackend): + provider = "flashinfer" + cache_formats = frozenset({"fp8kv_spt"}) + + def __init__(self, model): + set_flashinfer_envs() + super().__init__(model=model) + self.kv_data_type = torch.float8_e4m3fn + + def _init_infer_page_size(self): + self.infer_page_size = 1 + assert ( + self.page_size % self.infer_page_size == 0 + ), f"model page_size {self.page_size} must be divisible by infer_page_size {self.infer_page_size}" + logger.warning( + f"Fp8FlashInferAttBackend temporarily uses infer_page_size=1 with model page_size={self.page_size}; " + "multi-token FP8 inference pages are not implemented yet." + ) + + 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._prefill_fp8_prefill_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _prefill_fp8_prefill_att( + 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="cuda") + 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._decode_fp8_decode_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func, logical_layer_id=layer.logical_layer_id + ) + + def _decode_fp8_decode_att( + 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="cuda") + 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/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 75% rename from lightllm/common/basemodel/attention_vit/sdpa/fp.py rename to lightllm/common/attention/backends/dense/sdpa/varlen.py index 6c6da2c2b1..1e017b1e27 100644 --- a/lightllm/common/basemodel/attention_vit/sdpa/fp.py +++ b/lightllm/common/attention/backends/dense/sdpa/varlen.py @@ -1,18 +1,17 @@ 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: + +from lightllm.common.attention.backends.dense.varlen import BaseVarlenAttentionOp +from lightllm.platform import Platform + + +class SdpaVarlenAttentionOp(BaseVarlenAttentionOp): + provider = "sdpa" + supported_platforms = frozenset({Platform.CPU, Platform.CUDA}) + + def forward(self, inputs, *, out): + 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 cu_seqlens = cu_seqlens.detach().to("cpu") 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..354a7f45e8 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/backend.py @@ -0,0 +1,206 @@ +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 +class TritonPrefillContext(AttentionExecutionContext): + pass + + +@dataclasses.dataclass +class TritonDecodeContext(AttentionExecutionContext): + b_mark_mtp_shared_group: torch.Tensor = None + decode_max_kv_seq_len: int = None + + +class TritonAttBackend(BaseAttentionBackend): + 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_prefill_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + else: + return self._prefill_nomarl_prefill_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _prefill_alibi_prefill_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs, 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_nomarl_prefill_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs=None, alloc_func=torch.empty + ): + from lightllm.common.attention.backends._shared.prefill.context_flashattention_nopad import ( + context_attention_fwd, + ) + + if layer.params.sliding_window is not None: + sliding_window = layer.params.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_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, 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._decode_spec_decode_gqa_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + elif q_head_num == k_head_num and layer.params.sliding_window is None: + return self._decode_normal_decode_flash_decoding_att( + layer, context, q=q, k=k, v=v, alloc_func=alloc_func + ) + elif q_head_num >= k_head_num: + return self._decode_normal_decode_gqa_flash_decoding_att( + layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func + ) + else: + raise NotImplementedError("error") + + def _decode_alibi_decode_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs, 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_normal_decode_flash_decoding_att( + 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_normal_decode_gqa_flash_decoding_att( + self, layer, context, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, inputs=None, 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.params.sliding_window is not None: + sliding_window = layer.params.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 _decode_spec_decode_gqa_att( + 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..4374c8db00 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int4.py @@ -0,0 +1,146 @@ +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 +class Int4kvTritonPrefillContext(AttentionExecutionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass +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): + super().__init__(model) + 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._prefill_groupsize_quant_prefill_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + return o + + def _prefill_groupsize_quant_prefill_att( + 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._decode_ppl_int4kv_decode_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + + def _decode_ppl_int4kv_decode_att( + 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..a7b6f82319 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/int8.py @@ -0,0 +1,184 @@ +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 +class Int8kvTritonPrefillContext(AttentionExecutionContext): + b_kv_start_loc: torch.Tensor = None + + +@dataclasses.dataclass +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): + super().__init__(model) + 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._prefill_groupsize_quant_prefill_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + return o + + def _prefill_groupsize_quant_prefill_att( + 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._decode_diverse_decode_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + else: + return self._decode_normal_decode_att( + layer, context, q=q, k=k, k_scale=k_scale, v=v, v_scale=v_scale, alloc_func=alloc_func + ) + + def _decode_diverse_decode_att( + 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.objects, + ) + + def _decode_normal_decode_att( + 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/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 66% rename from lightllm/models/vit/triton_kernel/flashattention_nopad.py rename to lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py index af062c2a5b..4afb21aa1a 100644 --- a/lightllm/models/vit/triton_kernel/flashattention_nopad.py +++ b/lightllm/common/attention/backends/dense/triton/kernels/prefill/varlen.py @@ -1,11 +1,6 @@ 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": @@ -149,81 +144,3 @@ def _flash_attention_triton_fwd( 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/common/attention/backends/dense/triton/neo.py b/lightllm/common/attention/backends/dense/triton/neo.py new file mode 100644 index 0000000000..8fc5cad0c7 --- /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 +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..35a8701715 --- /dev/null +++ b/lightllm/common/attention/backends/dense/triton/varlen.py @@ -0,0 +1,22 @@ +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): + 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..b66eccdb71 --- /dev/null +++ b/lightllm/common/attention/backends/dense/varlen.py @@ -0,0 +1,49 @@ +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, k, v = inputs.query, inputs.key, inputs.value + if q.ndim != 3 or q.shape != k.shape or q.shape != v.shape: + raise UnsupportedAttention("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 UnsupportedAttention("Q/K/V must have the same dtype and device") + 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}") + cu = inputs.cu_seqlens + if cu.ndim != 1 or cu.numel() < 2 or cu.dtype != torch.int32 or cu.device != q.device: + raise UnsupportedAttention("cu_seqlens must be a same-device int32 vector with at least two entries") + if inputs.max_seqlen < 0 or (q.shape[0] > 0 and inputs.max_seqlen == 0): + raise UnsupportedAttention("nonempty attention requires a positive max_seqlen") 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..2009db9ae3 100644 --- a/lightllm/common/basemodel/attention_vit/xformers/fp.py +++ b/lightllm/common/attention/backends/dense/xformers/varlen.py @@ -1,31 +1,22 @@ 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.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): + 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..5aee2aa86a --- /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 TritonLinearAttBackend + + +class FlashQlaLinearAttBackend(TritonLinearAttBackend): + 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/__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..7b9c887f0f --- /dev/null +++ b/lightllm/common/attention/backends/gdn/replay/backend.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +import dataclasses + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.inputs import GDNInputs +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.mtp_state_params import build_dynamic_mtp_linear_att_state_params +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 lightllm.utils.envs_utils import get_env_start_args + +from .decode import replay + + +@dataclasses.dataclass +class ReplayGDNPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass +class ReplayGDNDecodeContext(AttentionExecutionContext): + 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 + + +class ReplayGDNTritonAttBackend(BaseAttentionBackend): + provider = "triton" + family = AttentionFamily.GDN + cache_formats = frozenset({"gdn"}) + + def __init__(self, model): + super().__init__(model=model) + from lightllm.common.attention.backends.gdn.triton.kernels.chunk import chunk_gated_delta_rule + + _configure_gdn(self, model, 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): + backend = self + state = layer.params.state + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer.logical_layer_id) + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv, z, b, a = _split_qkvzba(backend, inputs.mixed_qkvzba) + core_attn_out = self._prefill(layer, context, mixed_qkv, conv_states, ssm_states, a, b) + return core_attn_out, z + + def _prefill(self, layer, context, mixed_qkv, conv_states, ssm_states, a, b): + backend = self + 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=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[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): + self._init_indices(context, self.decode_draft_step, self.mtp_state_index) + context.b_ssm_buffer_idx = context.b_conv_buffer_idx + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None): + backend = self + mixed_qkv, z, b, a = _split_qkvzba(backend, inputs.mixed_qkvzba) + state = layer.params.state + conv_states, _ = state.get_layer_view(logical_layer_id=layer.logical_layer_id) + draft_step = backend.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 core, z + + def _init_indices(self, context, draft_step: int, mtp_index: torch.Tensor | None) -> None: + infer = context.batch + if draft_step == 0: + context.b_conv_buffer_idx = infer.b_req_idx + context.b1_mtp_cu_q_seq_len = None + context.b_num_accepted_tokens = None + # 每行一个 token。replay 和 cursor 共用这份前缀和。 + context.b_cu_seqlens = torch.arange( + infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device + ) + return + if self.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.hold_req_id, + ) + context.b1_mtp_cu_q_seq_len = cu + context.b_cu_seqlens = cu + context.b_conv_buffer_idx = conv_idx + context.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}." + ) + context.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + context.b_cu_seqlens = context.b1_mtp_cu_q_seq_len + context.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + context.b_num_accepted_tokens = mtp_index[context.b_conv_buffer_idx] + 1 + + def _decode(self, layer, context, mixed_qkv, z, conv_states, a, b, state): + backend = self + 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, + 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.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): + backend = self + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer.params.conv_weight, + mtp_step=draft_step, + bias=layer.params.conv_bias, + activation=backend.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.b_num_accepted_tokens, + query_start_loc=context.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.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 __init__(self, model): + BaseAttentionBackend.__init__(self, model=model) + from flash_qla import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) + + +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_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..83599572aa --- /dev/null +++ b/lightllm/common/attention/backends/gdn/triton/backend.py @@ -0,0 +1,328 @@ +import dataclasses +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention.api.backend import LayerAttentionBinding +from lightllm.common.attention.api.batch import AttentionBatch +from lightllm.common.attention.api.inputs import GDNInputs +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.mtp_state_params import build_dynamic_mtp_linear_att_state_params +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, +) +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +@dataclasses.dataclass +class LinearAttPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass +class LinearAttDecodeContext(AttentionExecutionContext): + 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 + + +class TritonLinearAttBackend(BaseAttentionBackend): + provider = "triton" + cache_formats = frozenset({"gdn"}) + family = AttentionFamily.GDN + + 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() + + 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 _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 + 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_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 _allocate_prefill_context(self, capacity, capture, workspace): + return LinearAttPrefillContext(capacity, capture, workspace) + + def _allocate_decode_context(self, capacity, capture, workspace): + return LinearAttDecodeContext(capacity, capture, workspace) + + def _update_prefill(self, context): + backend: TritonLinearAttBackend = self + mtp_step = backend.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) -> torch.Tensor: + mixed_qkvzba: torch.Tensor = inputs.mixed_qkvzba + backend: TritonLinearAttBackend = self + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + 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._prefill_gdn_prefill_kernel( + layer, context, mixed_qkv, conv_states, ssm_states, a, b, context.batch + ) + return (core_attn_out, z) + + def _prefill_gdn_prefill_kernel( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + batch: AttentionBatch, + ): + g, beta = fused_gdn_gating(layer.params.a_log, a, b, layer.params.dt_bias) + mixed_qkv = mixed_qkv.transpose(0, 1) + backend: TritonLinearAttBackend = self + out_tensor = causal_conv1d_fn( + mixed_qkv, + layer.params.conv_weight, + bias=layer.params.conv_bias, + query_start_loc=batch.b1_cu_q_seq_len, + cache_indices=context.b_conv_buffer_idx, + has_initial_state=batch.b_ready_cache_len > 0, + conv_states=conv_states, + activation=backend.activation, + ) + mixed_qkv = out_tensor.transpose(0, 1) + query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) + initial_state = ssm_states[context.b_ssm_buffer_idx] + 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=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): + draft_step = self.decode_draft_step + if draft_step == 0: + self._decode_init_normal_decode_state(context) + return + if self.uses_dynamic_spec_verify_layout(): + self._decode_init_dynamic_mtp_decode_state( + context, draft_step + 1, req_to_mtp_state_index=self.mtp_state_index + ) + else: + self._decode_init_fixed_mtp_decode_state(context, draft_step, req_to_mtp_state_index=self.mtp_state_index) + + def _decode_init_normal_decode_state(self, context): + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx + + def _decode_init_dynamic_mtp_decode_state( + self, context, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor + ) -> None: + ( + context.b1_mtp_cu_q_seq_len, + context.b_conv_buffer_idx, + context.b_num_accepted_tokens, + ) = build_dynamic_mtp_linear_att_state_params( + b_req_idx=context.batch.b_req_idx, + b_mtp_index=context.batch.b_mtp_index, + req_to_mtp_state_index=req_to_mtp_state_index, + hold_req_id=context.batch.hold_req_id, + ) + self._decode_init_mtp_ssm_buffer_idx(context, mtp_size) + + def _decode_init_fixed_mtp_decode_state( + self, context, draft_step: int, *, req_to_mtp_state_index: torch.Tensor + ) -> None: + mtp_size = draft_step + 1 + batch_size = context.batch.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 + context.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, mtp_size, dtype=torch.int32, device=context.batch.b_req_idx.device + ) + context.b_conv_buffer_idx = context.batch.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + context.b_num_accepted_tokens = req_to_mtp_state_index[context.b_conv_buffer_idx] + 1 + self._decode_init_mtp_ssm_buffer_idx(context, mtp_size) + + def _decode_init_mtp_ssm_buffer_idx(self, context, mtp_size: int): + att_batch_size = context.b_conv_buffer_idx.shape[0] + b_ssm_buffer_start_idx = (context.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) + state_offsets = torch.arange( + mtp_size, device=context.batch.b_req_idx.device, dtype=context.batch.b_req_idx.dtype + ).view(1, mtp_size) + context.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: GDNInputs, *, out=None) -> torch.Tensor: + mixed_qkvzba = inputs.mixed_qkvzba + backend: TritonLinearAttBackend = self + mixed_qkv, z, b, a = backend._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._decode_gdn_mtp_kernel( + layer, context, mixed_qkv, conv_states, ssm_states, a, b, context.batch + ) + else: + core_attn_out, z = self._decode_gdn_decode_kernel( + layer, context, mixed_qkv, z, conv_states, ssm_states, a, b, context.batch + ) + return (core_attn_out, z) + + def _decode_gdn_decode_kernel( + self, + layer, + context, + mixed_qkv: torch.Tensor, + z: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + batch: AttentionBatch, + ): + backend: TritonLinearAttBackend = self + 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, + 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=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 (core_attn_out, z) + + def _decode_gdn_mtp_kernel( + self, + layer, + context, + mixed_qkv: torch.Tensor, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + batch: AttentionBatch, + ): + from lightllm.common.attention.backends._shared.linear.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, + ) + + backend: TritonLinearAttBackend = self + 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=backend.decode_draft_step, + bias=layer.params.conv_bias, + activation=backend.activation, + conv_state_indices=context.b_conv_buffer_idx, + num_accepted_tokens=context.b_num_accepted_tokens, + query_start_loc=cu_seqlens_q, + ) + query, key, value = backend._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.b_num_accepted_tokens, + 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/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..0af6c9558a --- /dev/null +++ b/lightllm/common/attention/backends/kda/replay/backend.py @@ -0,0 +1,198 @@ +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.mtp_state_params import build_dynamic_mtp_linear_att_state_params +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 +class ReplayKDAPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass +class ReplayKDADecodeContext(AttentionExecutionContext): + 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 + + +class ReplayKDALinearAttBackend(BaseAttentionBackend): + provider = "triton" + family = AttentionFamily.KDA + cache_formats = frozenset({"kda"}) + + 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 _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): + backend: ReplayKDALinearAttBackend = self + state = layer.params.state + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer.logical_layer_id) + conv_states = conv_states[..., : backend.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 = 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 = inputs.raw_gate.view(1, -1, backend.tp_hidden_size) + raw_beta = inputs.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.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=backend.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + def _update_decode(self, context): + self._init_indices(context, self.decode_draft_step, self.mtp_state_index) + context.b_ssm_buffer_idx = context.b_conv_buffer_idx + + def _forward_decode(self, layer: LayerAttentionBinding, context, inputs: KDAInputs, *, out=None): + backend: ReplayKDALinearAttBackend = self + state = layer.params.state + conv_states, _ = state.get_layer_view(logical_layer_id=layer.logical_layer_id) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=context.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if backend.decode_draft_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.decode_draft_step, + num_accepted_tokens=context.b_num_accepted_tokens, + 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 = 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 = inputs.raw_gate.reshape(-1, backend.tp_hidden_size) + raw_beta = inputs.raw_beta.reshape(-1, backend.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 + + def _init_indices(self, context, mtp_step: int, mtp_index: torch.Tensor | None) -> None: + infer = context.batch + if mtp_step == 0: + context.b_conv_buffer_idx = infer.b_req_idx + context.b1_mtp_cu_q_seq_len = None + context.b_num_accepted_tokens = None + # 每行一个 token。replay 和 cursor 共用这份前缀和。 + context.b_cu_seqlens = torch.arange( + infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device + ) + return + if self.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.hold_req_id, + ) + context.b1_mtp_cu_q_seq_len = cu + context.b_cu_seqlens = cu + context.b_conv_buffer_idx = conv_idx + context.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}." + ) + context.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + context.b_cu_seqlens = context.b1_mtp_cu_q_seq_len + context.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + context.b_num_accepted_tokens = mtp_index[context.b_conv_buffer_idx] + 1 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..c84d2c15e4 --- /dev/null +++ b/lightllm/common/attention/backends/kda/triton/backend.py @@ -0,0 +1,200 @@ +"""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.mtp_state_params import build_dynamic_mtp_linear_att_state_params +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 +class KDAPrefillContext(AttentionExecutionContext): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + +@dataclasses.dataclass +class KDADecodeContext(AttentionExecutionContext): + 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 + + +class KDALinearAttBackend(BaseAttentionBackend): + provider = "triton" + cache_formats = frozenset({"kda"}) + family = AttentionFamily.KDA + + 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 _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 + backend: KDALinearAttBackend = self + conv_states, ssm_states = (layer.params.conv, layer.params.ssm) + conv_states = conv_states[..., : backend.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, backend.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, backend.tp_num_heads, backend.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=backend.lower_bound, + ) + ssm_states[context.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + def _update_decode(self, context): + mtp_step = self.mtp_step + if mtp_step == 0: + self._decode_init_normal_decode_state(context) + return + if self.uses_dynamic_spec_verify_layout(): + self._decode_init_dynamic_mtp_decode_state( + context, mtp_step + 1, req_to_mtp_state_index=self.mtp_state_index + ) + else: + self._decode_init_fixed_mtp_decode_state(context, mtp_step, req_to_mtp_state_index=self.mtp_state_index) + + def _decode_init_normal_decode_state(self, context): + context.b_conv_buffer_idx = context.batch.b_req_idx + context.b_ssm_buffer_idx = context.batch.b_req_idx + + def _decode_init_dynamic_mtp_decode_state( + self, context, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor + ) -> None: + ( + context.b1_mtp_cu_q_seq_len, + context.b_conv_buffer_idx, + context.b_num_accepted_tokens, + ) = build_dynamic_mtp_linear_att_state_params( + b_req_idx=context.batch.b_req_idx, + b_mtp_index=context.batch.b_mtp_index, + req_to_mtp_state_index=req_to_mtp_state_index, + hold_req_id=context.batch.hold_req_id, + ) + self._decode_init_mtp_ssm_buffer_idx(context, mtp_size) + + def _decode_init_fixed_mtp_decode_state( + self, context, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor + ) -> None: + mtp_size = mtp_step + 1 + batch_size = context.batch.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 + context.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, mtp_size, dtype=torch.int32, device=context.batch.b_req_idx.device + ) + context.b_conv_buffer_idx = context.batch.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() + context.b_num_accepted_tokens = req_to_mtp_state_index[context.b_conv_buffer_idx] + 1 + self._decode_init_mtp_ssm_buffer_idx(context, mtp_size) + + def _decode_init_mtp_ssm_buffer_idx(self, context, mtp_size: int): + att_batch_size = context.b_conv_buffer_idx.shape[0] + b_ssm_buffer_start_idx = (context.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) + state_offsets = torch.arange( + mtp_size, device=context.batch.b_req_idx.device, dtype=context.batch.b_req_idx.dtype + ).view(1, mtp_size) + context.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets + + 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 + backend: KDALinearAttBackend = self + 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 backend.mtp_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.mtp_step, + num_accepted_tokens=context.b_num_accepted_tokens, + 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 = 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.params.a_log, + gate_bias=layer.params.dt_bias, + initial_state=ssm_states, + lower_bound=backend.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.b_num_accepted_tokens, + ) + 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 6d627078dd..d9822c79dd 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..38c68a6078 --- /dev/null +++ b/lightllm/common/attention/backends/mla/fa3/backend.py @@ -0,0 +1,195 @@ +import dataclasses +from typing import Tuple + +import torch +import triton + +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.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, flash_attn_with_kvcache + + +@dataclasses.dataclass +class MlaFa3PrefillContext(AttentionExecutionContext): + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + causal: bool = None + + +@dataclasses.dataclass +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 + causal: bool = None + + +class MlaFa3AttBackend(Fa3AttBackend): + 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._prefill_mla_prefill_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _prefill_mla_prefill_att( + 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.params.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 _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: MLADecodeInputs, *, out=None): + alloc_func = layer.alloc_tensor + q = (inputs.query_nope, inputs.query_rope) + k = layer.params.kv + v = None + assert v is None + return self._decode_mla_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _decode_mla_decode_att( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + v: torch.Tensor, + inputs=None, + 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.params.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/__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..ee3ec534c1 --- /dev/null +++ b/lightllm/common/attention/backends/mla/flashinfer/backend.py @@ -0,0 +1,223 @@ +import dataclasses +from typing import 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.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 + + +@dataclasses.dataclass +class MlaFlashInferPrefillContext(AttentionExecutionContext): + prefill_wrapper: object = None + + +@dataclasses.dataclass +class MlaFlashInferDecodeContext(AttentionExecutionContext): + kv_indices: torch.Tensor = None + kv_starts: torch.Tensor = None + q_indptr: torch.Tensor = None + decode_wrapper: object = None + + +class MlaFlashInferAttBackend(BaseAttentionBackend): + 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): + 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) + 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 for_spec(self, spec): + backend = super().for_spec(spec) + backend.tp_q_head_num = spec.num_q_heads + return backend + + def resource_requirements(self, capacity, *, capture): + return BackendResources(((self.workspace_buffer_key, self.workspace_buffer_size, torch.uint8),)) + + def _init_infer_page_size(self): + self.infer_page_size = self.page_size + assert ( + self.page_size % self.infer_page_size == 0 + ), f"model page_size {self.page_size} must be divisible by infer_page_size {self.infer_page_size}" + + 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=True, + sm_scale=self.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._prefill_mla_prefill_att(layer, context, q=q, k=k, v=v, alloc_func=alloc_func) + + def _prefill_mla_prefill_att( + 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]), 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) + context.prefill_wrapper.run(q, k, v, out=o_tensor) + return o_tensor + + def _decode_refresh_cuda_graph_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 + v = None + assert v is None + return self._decode_mla_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _decode_mla_decode_att( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + v: torch.Tensor, + inputs, + 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.softmax_scale, + self.q_data_type, + self.kv_data_type, + ) + else: + self._decode_refresh_cuda_graph_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..fceec8b738 --- /dev/null +++ b/lightllm/common/attention/backends/mla/triton/backend.py @@ -0,0 +1,160 @@ +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 +class MlaTritonPrefillContext(AttentionExecutionContext): + pass + + +@dataclasses.dataclass +class MlaTritonDecodeContext(AttentionExecutionContext): + block_seq: torch.Tensor = None + batch_start: torch.Tensor = None + indices_initialized: bool = False + + +class MlaTritonAttBackend(BaseAttentionBackend): + 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._prefill_contiguous_kv_prefill_att( + layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func + ) + else: + return self._prefill_mla_prefill_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _prefill_mla_prefill_att( + 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.device) + k_nope, k_rope = k + softmax_scale = layer.params.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 _prefill_contiguous_kv_prefill_att( + 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 + v = None + assert v is None + return self._decode_mla_decode_att(layer, context, q=q, k=k, v=v, inputs=inputs, alloc_func=alloc_func) + + def _decode_mla_decode_att( + self, + layer, + context, + q: Tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor, + v: torch.Tensor, + inputs, + alloc_func=torch.empty, + ): + softmax_scale = layer.params.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/__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..6c20d83895 --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/backend.py @@ -0,0 +1,173 @@ +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 +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 +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): + 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._prefill_nsa_prefill_att(layer, context, q=q, kv=k, inputs=inputs) + + def _prefill_nsa_prefill_att(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.params.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._decode_nsa_decode_att(layer, context, q=q, kv=k, inputs=inputs) + + def _decode_nsa_decode_att( + 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.params.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..6a77f640d9 --- /dev/null +++ b/lightllm/common/attention/backends/sparse_mla/flashmla/fp8.py @@ -0,0 +1,149 @@ +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 +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 +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._prefill_nsa_prefill_att(layer, context, q=q, packed_kv=k, inputs=inputs) + + def _prefill_nsa_prefill_att( + self, layer, context, q: torch.Tensor, packed_kv: torch.Tensor, inputs + ) -> torch.Tensor: + import flash_mla + + topk_indices = inputs.topk_indices + softmax_scale = layer.params.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._decode_nsa_decode_att(layer, context, q=q, packed_kv=k, inputs=inputs) + + def _decode_nsa_decode_att( + 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.params.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..bc64e5676a --- /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_execution=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(model).for_spec(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.forward(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/selection.py b/lightllm/common/attention/selection.py new file mode 100644 index 0000000000..485fdb5935 --- /dev/null +++ b/lightllm/common/attention/selection.py @@ -0,0 +1,211 @@ +import importlib + +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.platform import Platform +from lightllm.utils import backend_validator +from lightllm.utils.backend_validator import validate +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 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", +} + + +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]) + + +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)) + devices = tuple(devices) if devices is not None else (device,) + if not devices: + raise ValueError("varlen selection requires at least one device") + ordered = tuple(varlen_attention_ops) 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 探针。 + if not all( + backend_validator._validate(name, varlen=True, device=d, dtype=dtype, head_dim=head_dim) for d in devices + ): + failures[name] = "varlen environment or numerical probe failed on an assigned device" + continue + logger.info("Selected %s varlen attention", name) + return name + raise BackendSelectionError(AttentionFamily.DENSE, AttentionPhase.PREFILL, failures) + + +data_type_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_data_type_to_backend = { + "None": { + "triton": "backends.mla.triton.backend:MlaTritonAttBackend", + "fa3": "backends.mla.fa3.backend:MlaFa3AttBackend", + "flashinfer": "backends.mla.flashinfer.backend:MlaFlashInferAttBackend", + }, +} + +nsa_data_type_to_backend = { + "None": { + "flashmla_sparse": "backends.sparse_mla.flashmla.backend:NsaFlashMlaSparseAttBackend", + }, + "fp8kv_dsa": { + "flashmla_sparse": "backends.sparse_mla.flashmla.fp8:NsaFlashMlaFp8SparseAttBackend", + }, +} + +neo_data_type_to_backend = { + "None": { + "triton": "backends.dense.triton.neo:NeoTritonAttBackend", + "fa3": "backends.dense.fa3.neo:NeoFa3AttBackend", + } +} + + +def select_attention_backend(*, family, phase, model, requested="auto", candidates=None, priority=None, specs=None): + from lightllm.common.basemodel.attention import describe_attention_layers + + args = model.args + caps = model.device_runtime.caps + execution = model.attention_execution + if candidates is None: + tables = { + AttentionFamily.DENSE: data_type_to_backend, + AttentionFamily.MLA: mla_data_type_to_backend, + AttentionFamily.SPARSE_MLA: nsa_data_type_to_backend, + AttentionFamily.GDN: { + "None": { + "flashqla": "backends.gdn.flashqla.backend:FlashQlaLinearAttBackend", + "triton": "backends.gdn.triton.backend:TritonLinearAttBackend", + } + }, + AttentionFamily.KDA: {"None": {"triton": "backends.kda.triton.backend:KDALinearAttBackend"}}, + } + if args.ssm_state_mode == "replay": + tables[AttentionFamily.GDN] = { + "None": { + "flashqla": "backends.gdn.replay.backend:ReplayGDNFlashQlaAttBackend", + "triton": "backends.gdn.replay.backend:ReplayGDNTritonAttBackend", + } + } + tables[AttentionFamily.KDA] = { + "None": { + "triton": "backends.kda.replay.backend:ReplayKDALinearAttBackend", + } + } + if caps.platform is not Platform.CUDA: + tables = {} + candidates = tables.get(family, {}) + encoding = "None" if family in (AttentionFamily.GDN, AttentionFamily.KDA) else args.llm_kv_type + available = candidates.get(encoding, {}) + if priority is None: + 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 args.mtp_step: + priority = ["flashinfer", "fa3", "triton"] + ordered = [requested] if requested != "auto" else list(priority) + specs = describe_attention_layers(model, family, phase) if specs is None else specs + if not specs: + raise ValueError(f"model has no {family.value} attention layers") + 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 + 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: + if not validate(available[name], device=str(caps.device), dtype=spec.dtype, spec=spec): + failures[name] = "provider numerical probe failed on the assigned device/dtype" + break + if name in failures: + continue + except UnsupportedAttention as exc: + failures[name] = str(exc) + continue + except Exception as exc: + failures[name] = f"provider check failed: {type(exc).__name__}: {exc}" + continue + logger.info("Selected %s %s attention on %s: %s", family.value, phase.value, caps.device, name) + 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..aa758efdea --- /dev/null +++ b/lightllm/common/attention/varlen.py @@ -0,0 +1,35 @@ +from lightllm.platform import get_device_runtime +from lightllm.utils.log_utils import init_logger + +from .api import AttentionFamily, AttentionPhase, UnsupportedAttention +from .selection import BackendSelectionError, get_varlen_attention_op_class + +logger = init_logger(__name__) + + +class AutoVarlenAttentionOp: + """音频编码器的自动路径；失败的实现只在当前调用方禁用，并在本次调用中重算输出。""" + + def __init__(self, providers=("fa3", "triton")): + 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/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 137ce8ad32..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 @@ -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 @@ -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 2c73e97d6f..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 @@ -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 @@ -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..1b1eec453f --- /dev/null +++ b/lightllm/common/basemodel/attention.py @@ -0,0 +1,275 @@ +from dataclasses import dataclass, replace + +import torch + +from lightllm.common.attention.api import ( + 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 + + +@dataclass(frozen=True) +class AttentionGroup: + prefill_backend: object + decode_backend: object + prefill_spec: LayerSpec + decode_spec: LayerSpec + + def backend_for(self, phase): + return self.prefill_backend if phase is AttentionPhase.PREFILL else self.decode_backend + + +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=not model.is_mtp_draft_model and model.args.mtp_dynamic_verify, + variant=layer.attention_variant, + ) + + +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, + ssm, + weight.linear_conv1d.mm_param.weight, + weight.linear_conv1d.bias, + weight.linear_A_log.weight, + weight.linear_dt_bias.weight, + state=state, + ) + return KDAParams( + conv, + ssm, + weight.get_merged_kda_conv_weight(), + weight.linear_A_log.weight, + weight.linear_dt_bias.weight, + state=state, + ) + 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, + value, + getattr(weight, "tp_alibi", None), + layer.attention_window(AttentionPhase.DECODE), + weight.attn_sinks.weight if layer.attention_uses_sink else None, + query_scale, + key_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.softmax_scale, + 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_execution = ExecutionRequirements() + model.attention_groups = () + model._attention_layer_groups = (None,) * len(model.layers_infer) + model._attention_workspaces = {} + return + model.attention_execution = ExecutionRequirements.from_model(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_execution) + 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( + family=family, + phase=spec.phase, + model=model, + requested=getattr(model.args, option), + candidates=candidates, + priority=priority, + specs=(spec,), + ) + backends.append(cls(model).for_spec(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_execution + 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) + prefill_params = ( + replace(params, sliding_window=group.prefill_spec.sliding_window) + if isinstance(params, DenseParams) + else params + ) + layer.prefill_attention = LayerAttentionBinding( + group_id=group_id, + logical_layer_id=layer.layer_num_, + spec=group.prefill_spec, + params=prefill_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, + ) + if group.prefill_spec.family in (AttentionFamily.GDN, AttentionFamily.KDA): + state = model.att_state_manager.get_state( + model.att_state_role, group.prefill_spec.family.value, layer.layer_num_ + ) + group.prefill_backend.mtp_state_index = state.mtp_state_index + group.decode_backend.mtp_state_index = state.mtp_state_index + 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 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_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/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..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 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/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 83b312f991..6b33aa99bb 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,17 +1,19 @@ -import torch -import triton import collections +from contextlib import contextmanager +from dataclasses import replace +from typing import List, Optional + +import torch + +from lightllm.common.attention.api import AttentionBatch, AttentionPhase, ContextCapacity 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 .triton_kernel.gen_decode_params import gen_decode_params +from .triton_kernel.gen_prefill_params import gen_prefill_params class InferStateInfo: @@ -20,13 +22,9 @@ 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_groups = () + self.attention_contexts = {} + self.attention_batch = None self.input_ids: torch.Tensor = None self.batch_size: int = None @@ -133,58 +131,102 @@ 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 prepare_attention(self, model): + from lightllm.common.basemodel.attention import get_attention_workspace + + self.device_runtime = model.device_runtime + if model._attention_storage_generation != self.att_state_manager.storage_generation: + raise RuntimeError("attention storage changed; rebind layers and recapture CUDA graphs") + phase = AttentionPhase.PREFILL if self.is_prefill else AttentionPhase.DECODE + self.attention_batch = AttentionBatch( + phase=phase, + batch_size=self.batch_size, + token_num=self.input_ids.shape[0], + total_token_num=self.total_token_num, + max_q_seq_len=self.max_q_seq_len, + max_kv_seq_len=self.max_kv_seq_len, + max_cache_len=self.max_cache_len or 0, + b_req_idx=self.b_req_idx, + b_seq_len=self.b_seq_len, + b_q_seq_len=self.b_q_seq_len, + b1_cu_q_seq_len=self.b1_cu_q_seq_len, + b1_cu_kv_seq_len=self.b1_cu_kv_seq_len, + req_to_token_indexs=self.req_manager.req_to_token_indexs, + hold_req_id=self.req_manager.HOLD_REQUEST_ID, + b_ready_cache_len=self.b_ready_cache_len, + b_mtp_index=self.b_mtp_index, + b_shared_seq_len=self.b_shared_seq_len, + b_shared_radix_node_id=self.b_shared_radix_node_id, + mem_index=self.mem_index, + position_ids=self.position_ids, + b_image_token_end=getattr(self, "b_image_token_end", None), + ) + if self.is_cuda_graph_replay: + return + capacity = ContextCapacity( + phase, + self.batch_size, + self.input_ids.shape[0], + model.graph_max_len_in_batch if self.is_cuda_graph_capture else self.max_kv_seq_len, + self.max_q_seq_len, + ) + for group_id, group in enumerate(self.attention_groups): + workspace = get_attention_workspace(model, group_id, self.microbatch_index, self.input_ids.device) + workspace.begin() + backend = group.backend_for(phase) + if self.is_cuda_graph_capture: + workspace = workspace.fork() + workspace.reserve(backend.resource_requirements(capacity, capture=self.is_cuda_graph_capture)) + context = backend.allocate_context(capacity, capture=self.is_cuda_graph_capture, workspace=workspace) + backend.prepare(context, self.attention_batch) + self.attention_contexts[group_id] = context + + def update_attention(self): + for group_id, context in self.attention_contexts.items(): + self.attention_groups[group_id].backend_for(context.capacity.phase).update(context) + + def finish_attention(self): + for context in self.attention_contexts.values(): + context.workspace.finish() def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"): + if ( + self.att_state_role != new_infer_state.att_state_role + or len(self.attention_groups) != len(new_infer_state.attention_groups) + or any( + current is not incoming + for current, incoming in zip(self.attention_groups, new_infer_state.attention_groups) + ) + ): + raise ValueError("captured attention belongs to a different model role or group") + 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)) + for context in self.attention_contexts.values(): + context.workspace.begin() + self.attention_batch.copy_inputs(new_infer_state.attention_batch) + for target, source in copies: + target.copy_(source, non_blocking=True) + batch = replace( + self.attention_batch, + max_kv_seq_len=new_infer_state.max_kv_seq_len, + max_q_seq_len=new_infer_state.max_q_seq_len, + total_token_num=new_infer_state.total_token_num, + ) + for group_id, context in self.attention_contexts.items(): + self.attention_groups[group_id].decode_backend.prepare(context, batch) return def call_overlap_hook(self): @@ -330,6 +372,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 +404,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 +430,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 +476,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/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..92debf327d 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,20 +31,9 @@ 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_contexts[self.prefill_attention.group_id], DenseInputs(_q) ) o_tensor = o_tensor.view(q.shape) return o_tensor @@ -50,20 +41,9 @@ def _context_attention_kernel( 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_contexts[self.decode_attention.group_id], DenseInputs(_q) ) return o_tensor.view(q.shape) 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..4308807310 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,12 +88,8 @@ 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_contexts[self.prefill_attention.group_id], MLAPrefillInputs(q, k_nope, k_rope, v) ) return o_tensor @@ -104,20 +102,9 @@ 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_contexts[self.decode_attention.group_id], MLADecodeInputs(q_nope, q_rope) ) return out @@ -418,7 +405,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 +567,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..7e1bca4cf2 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 = infer_state.attention_contexts[self.prefill_attention.group_id] 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.attention_contexts[self.prefill_attention.group_id], + 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 = infer_state.attention_contexts[self.decode_attention.group_id] 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.attention_contexts[self.decode_attention.group_id], + SparseMLADecodeInputs(q_nope, q_rope, topk_mem_indices), ) return o_tensor @@ -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: 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/sparse_mla.py b/lightllm/models/deepseek_v4/attention/sparse_mla.py new file mode 100644 index 0000000000..5b10304ffd --- /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 +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): + 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 容量评估前调用，以计入实际显存占用。""" + 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/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..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..872a8924ed 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,33 +512,13 @@ 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, - }, + inputs = DeepseekV4Inputs( + q, meta["swa_indices"], meta["swa_lengths"], meta.get("extra_indices"), meta.get("extra_lengths") ) - 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, + return self.prefill_attention.forward( + infer_state.attention_contexts[self.prefill_attention.group_id], inputs, out=out ) # ------------------------------------------------------------------ attention (decode) @@ -526,25 +529,13 @@ 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, - }, + inputs = DeepseekV4Inputs( + q, meta["swa_indices"], meta["swa_lengths"], meta.get("extra_indices"), meta.get("extra_lengths") ) - 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, + return self.decode_attention.forward( + infer_state.attention_contexts[self.decode_attention.group_id], inputs, out=out ) # ------------------------------------------------------------------ moe diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index ff46b5d6da..a5dffb8fd8 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( diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index d1ebd30e30..73159e6b54 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 @@ -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,8 +244,8 @@ 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_contexts[self.prefill_attention.group_id], DenseInputs(_q) ) return o_tensor.view(q.shape) @@ -253,10 +256,11 @@ 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_contexts[self.decode_attention.group_id], 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/__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..fd20d7e74b --- /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 +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 +class Fp8Glm5NextSparsePrefillContext(Glm5NextSparsePrefillContext): + smallq_workspace: object = None + generic_workspace: object = None + generic_binding: object = None + + +@dataclasses.dataclass +class Fp8Glm5NextSparseDecodeContext(NsaFlashMlaFp8SparseDecodeContext): + decode_topk_bound: int = None + + +class Fp8Glm5NextSparseAttBackend(NsaFlashMlaFp8SparseAttBackend): + provider = "flashmla_sparse" + cache_formats = frozenset({"fp8kv_dsa"}) + + def __init__(self, model): + super().__init__(model) + 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=( + ("glm_gathered_kv", gather_tokens * 512, torch.bfloat16), + ("glm_chunk_indices", tokens * topk, torch.int32), + ("glm_chunk_lengths", tokens, torch.int32), + ("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, + ) + + slot = context.workspace.objects + if "glm_smallq" not in slot: + slot["glm_smallq"] = { + "direct": allocate_direct_smallq_workspace(q[:1]), + "selected": allocate_selected_prefill_workspace(q, batch_capacity=8), + } + context.smallq_workspace = slot["glm_smallq"] + 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 _prefill_nsa_prefill_att(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.params.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.params.softmax_scale, + workspace, + direct_workspace, + generic_workspace, + ) + + def _decode_nsa_decode_att(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.params.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.params.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._prefill_nsa_prefill_att(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..8b5f7c67da --- /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 +class Glm5NextSparsePrefillContext(NsaFlashMlaSparsePrefillContext): + pass + + +@dataclasses.dataclass +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 _prefill_nsa_prefill_att(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.params.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 _decode_nsa_decode_att(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.params.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.params.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/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 be7a2fa781..523b45aa4f 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"] @@ -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=infer_state.attention_contexts[self.prefill_attention.group_id], 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_contexts[self.prefill_attention.group_id], + 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=infer_state.attention_contexts[self.decode_attention.group_id], 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_contexts[self.decode_attention.group_id], + SparseMLADecodeInputs(q_nope, q_rope, topk_mem_indices), ) def _get_o(self, input, infer_state, layer_weight): @@ -326,7 +289,7 @@ 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 + backend = infer_state.attention_groups[self.prefill_attention.group_id].prefill_backend mixed_qkv = mixed_qkv.contiguous() raw_decay_gate = raw_decay_gate.contiguous() beta_logits = beta_logits.contiguous() @@ -348,21 +311,9 @@ def _kda_prefill_cuda_graph_wrapper( _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_contexts[self.prefill_attention.group_id], + KDAInputs(_mixed_qkv, _raw_decay_gate, _beta_logits), ) assert tmp_output.shape == _output.shape _output.copy_(tmp_output) @@ -388,21 +339,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_contexts[self.prefill_attention.group_id], + KDAInputs(qkv, raw_decay_gate, beta_logits), ) return self._kda_post(core_output, raw_output_gate, infer_state, layer_weight) @@ -412,21 +351,8 @@ 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_contexts[self.decode_attention.group_id], KDAInputs(qkv, raw_decay_gate, beta_logits) ) return self._kda_post(core_output, raw_output_gate, infer_state, layer_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..ef3913aeb3 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,32 +74,9 @@ 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_contexts[self.prefill_attention.group_id], DenseInputs(_q) ) o_tensor = o_tensor.view(q.shape) return o_tensor @@ -100,32 +84,9 @@ def _context_attention_kernel( 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_contexts[self.decode_attention.group_id], 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..025eb1bf10 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,20 +47,9 @@ 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_contexts[self.prefill_attention.group_id], DenseInputs(_q) ) o_tensor = o_tensor.view(q.shape) return o_tensor @@ -70,20 +60,9 @@ 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_contexts[self.decode_attention.group_id], DenseInputs(_q) ) return o_tensor.view(q.shape) 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..4a7ac32ac0 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.selection import neo_data_type_to_backend + + if phase is AttentionPhase.PREFILL: + return neo_data_type_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..dc1648c627 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 = "" @@ -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.selection import neo_data_type_to_backend + + if phase is AttentionPhase.PREFILL: + return neo_data_type_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/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/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..77a8aae064 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,20 +34,12 @@ 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_contexts[self.decode_attention.group_id], DenseInputs(q) ) return output.view(-1, self.tp_q_head_num_ * self.head_dim_) 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..186d54e1ed 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 AutoVarlenAttentionOp 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 @@ -103,6 +104,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 = AutoVarlenAttentionOp() 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 +118,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 +126,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) diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 1bc4ef914b..83b2f6b765 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,7 +187,7 @@ def _linear_prefill_cuda_graph_wrapper( infer_state: Qwen3NextInferStateInfo, layer_weight: Qwen3NextTransformerLayerWeight, ) -> Tuple[torch.Tensor, torch.Tensor]: - backend = infer_state.prefill_att_state1.backend + backend = infer_state.attention_groups[self.prefill_attention.group_id].prefill_backend mixed_qkvzba = mixed_qkvzba.contiguous() _mixed_qkvzba = tensor_to_no_ref_tensor(mixed_qkvzba) @@ -209,22 +212,9 @@ 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_contexts[self.prefill_attention.group_id], GDNInputs(_mixed_qkvzba) ) tmp_o = tmp_o.view(_o.shape) _o.copy_(tmp_o) @@ -249,19 +239,8 @@ 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_contexts[self.prefill_attention.group_id], GDNInputs(mixed_qkvzba) ) gdn_out = self._linear_post(core_attn_out, z, layer_weight) @@ -280,19 +259,8 @@ 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_contexts[self.decode_attention.group_id], GDNInputs(mixed_qkvzba) ) gdn_out = self._linear_post(core_attn_out, z, layer_weight) 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..7ccfa98bab --- /dev/null +++ b/lightllm/platform/__init__.py @@ -0,0 +1,56 @@ +"""CUDA/CPU device operations used by the attention backend runtime.""" + +import os +from functools import lru_cache + +import torch + +from .api import DeviceCaps, DeviceRuntime, Platform + + +def get_device_runtime(device=None, *, platform=None): + platform = Platform(platform) if platform is not None else None + if platform not in (None, Platform.CPU, Platform.CUDA): + raise ValueError(f"attention backend refactor supports CUDA and CPU only, got {platform.value}") + 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..1b4d5b6a75 --- /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 + architecture: 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) + + 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..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", @@ -536,9 +542,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 +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 a4c990c2be..81f0a4a383 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 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 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__) @@ -183,17 +185,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..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..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..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/server/visualserver/__init__.py b/lightllm/server/visualserver/__init__.py index 6bc0923da0..9d954d165b 100644 --- a/lightllm/server/visualserver/__init__.py +++ b/lightllm/server/visualserver/__init__.py @@ -1,16 +1,43 @@ -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.selection import get_varlen_attention_op_class, select_varlen_attention_op +from lightllm.platform import get_device_runtime -VIT_ATTN_BACKEND: BaseVitAttBackend = None +VIT_ATTENTION_OP: VarlenAttentionOp = None -def set_vit_att_backend(backend_name: str): - global VIT_ATTN_BACKEND - VIT_ATTN_BACKEND = get_vit_att_backend_class(backend_name) +class _InitializingVarlenAttentionOp: + """用 worker 预热时的真实 Q/K/V 选择实现，之后直接绑定实现的 forward。 + + 部分视觉模型固定使用 BF16，与 llm 的 data_type 不同。manager 无法据启动 + 参数推断实际 dtype/head dim，因此选择发生在各设备的首次预热输入到达时。 + """ + + def __init__(self, requested): + self.requested = requested + + def forward(self, inputs, *, out): + 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] + ) + op = get_varlen_attention_op_class(name)() + global VIT_ATTENTION_OP + VIT_ATTENTION_OP = op + self.forward = op.forward + return self.forward(inputs, out=out) + + +def set_vit_attention_op(provider: str, *, defer_validation=False): + global VIT_ATTENTION_OP + VIT_ATTENTION_OP = ( + _InitializingVarlenAttentionOp(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..0087b1938a 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -2,10 +2,12 @@ import multiprocessing as mp import os +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__) @@ -18,20 +20,19 @@ def _compute_ground_truth(q, k, v, is_causal=True): 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 +56,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 +71,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 +92,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 +114,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,7 +151,10 @@ 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 @@ -169,7 +183,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 +225,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,12 +254,15 @@ 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(): +def _validate_flashmla_sparse(dtype=torch.bfloat16, spec=None): """Validate flashmla_sparse (NSA) with ground truth.""" try: from sgl_kernel.flash_mla import flash_mla_sparse_fwd @@ -252,7 +275,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,7 +304,40 @@ def _validate_flashmla_sparse(): return True, None -def _run_in_subprocess(backend_name, pipe): +def _validate_varlen(backend_name, dtype=None, device=None, head_dim=80): + from lightllm.common.attention.api import UnsupportedAttention, VarlenAttentionInputs + from lightllm.common.attention.selection 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 + + +def _run_in_subprocess(backend_name, pipe, varlen=False, device=None, dtype=None, spec=None, head_dim=80): """Run validation in subprocess with suppressed output.""" import sys @@ -292,20 +347,48 @@ def _run_in_subprocess(backend_name, pipe): sys.stderr = devnull try: - if backend_name == "fa3": - success, err = _validate_fa3() + if device is not None or spec is not None: + from lightllm.platform import get_device_runtime + + runtime = get_device_runtime(device) + if runtime.api is not None: + runtime.api.set_device(runtime.device) + if spec is not None: + from lightllm.common.attention.api import AttentionFamily + from lightllm.common.attention.selection import load_backend_class + + cls = load_backend_class(backend_name) + if ( + 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 + ): + from lightllm.common.attention.probe import probe_dense + + success, err = probe_dense(cls, spec, runtime) + pipe.send((success, err)) + return + backend_name = cls.provider + if varlen: + success, err = _validate_varlen(backend_name, dtype, device, head_dim) + elif backend_name == "fa3": + success, err = _validate_fa3(dtype or torch.bfloat16, spec) elif backend_name == "xformers": success, err = _validate_xformers() elif backend_name == "sdpa": success, err = _validate_sdpa() elif backend_name == "flashinfer": - success, err = _validate_flashinfer() + success, err = _validate_flashinfer(dtype or torch.float16, spec) elif backend_name == "flashqla": - success, err = _validate_flashqla() + success, err = _validate_flashqla(dtype, spec) elif backend_name == "triton": success, err = _validate_triton() elif backend_name == "flashmla_sparse": - success, err = _validate_flashmla_sparse() + success, err = _validate_flashmla_sparse(dtype or torch.bfloat16, spec) else: success, err = False, f"Unknown backend: {backend_name}" pipe.send((success, err)) @@ -316,25 +399,20 @@ def _run_in_subprocess(backend_name, pipe): @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: +def validate(backend_name: str, *, device=None, dtype=None, spec=None) -> bool: + """每个 rank 在自身设备验证；dtype/完整 spec 不同不能复用另一配置的结果。""" + return _validate(backend_name, device=device, dtype=dtype, spec=spec) + + +def _validate(backend_name: str, *, varlen=False, device=None, dtype=None, spec=None, head_dim=80) -> bool: """Validate backend in subprocess with ground truth check.""" 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)) + proc = ctx.Process(target=_run_in_subprocess, args=(backend_name, child, varlen, device, dtype, spec, head_dim)) proc.start() + child.close() proc.join(timeout=_VALIDATION_TIMEOUT) if proc.is_alive(): @@ -358,3 +436,8 @@ def _validate(backend_name: str) -> bool: except Exception as e: logger.warning(f"{backend_name} validation error: {e}") return False + 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/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_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..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..712f4dd754 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 de13eb28fc..8abcec7f3e 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 Int4kvTritonAttBackend 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(Int4kvTritonAttBackend, "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..6f26fc66a7 --- /dev/null +++ b/unit_tests/common/attention/helpers.py @@ -0,0 +1,284 @@ +from dataclasses import replace +from types import SimpleNamespace + +import torch + +from lightllm.common.attention.api import ( + AttentionBatch, + AttentionFamily, + AttentionPhase, + CacheLayout, + ContextCapacity, + ExecutionRequirements, + LayerAttentionBinding, + LayerSpec, +) +from lightllm.common.attention.backends._workspace import AttentionWorkspace +from lightllm.common.attention.backends.base import BaseAttentionBackend +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, + infer_page_size=1, + causal=True, + dynamic_verify=False, + decode_draft_step=0, + is_draft_model=False, + hold_req_id=7, + mtp_state_index=None, + ) + values.update(kwargs) + 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_execution = ExecutionRequirements.from_model(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 + vars(context).update(values) + 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["softmax_scale"], 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) + 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.forward(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.forward(context, DenseInputs(query), out=out) 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..edaf8a0981 --- /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 lambda model: SimpleNamespace(for_spec=lambda 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]["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 new file mode 100644 index 0000000000..714bdc78ce --- /dev/null +++ b/unit_tests/common/attention/test_contracts.py @@ -0,0 +1,283 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention import selection +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_backend, + 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 = make_backend() + 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"): + make_backend().prepare(context, make_batch()) + with pytest.raises(ValueError, match="different attention backend"): + make_backend().update(context) + attention = make_layer_attention(make_backend(), context, params=None) + with pytest.raises(ValueError, match="different attention backend"): + attention.forward(context, inputs=None) + + +def test_layer_attention_rejects_wrong_phase(): + backend = make_backend() + 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.forward(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("lightllm.common.basemodel.attention.describe_attention_layers", lambda *args: (spec,)) + monkeypatch.setattr( + BaseAttentionBackend, + "check_environment", + classmethod(lambda cls, caps: None), + ) + with pytest.raises(selection.BackendSelectionError, match="encoding fp8kv_sph"): + selection.select_attention_backend( + family=spec.family, + phase=spec.phase, + model=model_stub(), + 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): + monkeypatch.setattr("lightllm.common.basemodel.attention.describe_attention_layers", lambda *args: (make_spec(),)) + + def reject(cls, caps): + raise UnsupportedAttention("missing extension") + + monkeypatch.setattr(BaseAttentionBackend, "check_environment", classmethod(reject)) + with pytest.raises(selection.BackendSelectionError, match="missing extension"): + selection.select_attention_backend( + 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.common.basemodel.attention.describe_attention_layers", lambda *args: (make_spec(),)) + monkeypatch.setattr("lightllm.common.attention.selection.validate", lambda *args, **kwargs: True) + 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 = selection.select_attention_backend( + 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: + selection.select_attention_backend( + 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, "validate", 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"): + selection.select_attention_backend( + family=spec.family, + phase=spec.phase, + model=model_stub(), + requested=requested, + specs=(spec,), + 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 = selection.data_type_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 + + original = make_backend(FlashInferAttBackend) + small = original.for_spec(make_spec(num_q_heads=8, head_size=64)) + large = original.for_spec(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 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 lambda model: SimpleNamespace(for_spec=lambda spec: SimpleNamespace(spec=spec)) + + 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 + + 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)) 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..59e555b743 --- /dev/null +++ b/unit_tests/common/attention/test_dense_numerics.py @@ -0,0 +1,177 @@ +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.selection import data_type_to_backend, load_backend_class +from lightllm.common.attention_state import AttentionStateManager +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 + +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(data_type_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,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(data_type_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(model).for_spec(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, value, sliding_window=window)), spec=spec + ) + inputs = DenseInputs(query) + output = torch.empty_like(query) + + def run(): + backend.update(context) + return attention.forward(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, + ) + 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 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..25ba291865 --- /dev/null +++ b/unit_tests/common/attention/test_platform_runtime.py @@ -0,0 +1,177 @@ +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.common.basemodel.graph import GraphExecution +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.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_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")) + 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) + 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")) + 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() 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..f37a6499af --- /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.selection 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:TritonLinearAttBackend", AttentionFamily.GDN), + ("backends.gdn.flashqla.backend:FlashQlaLinearAttBackend", 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/attention/test_varlen_selection.py b/unit_tests/common/attention/test_varlen_selection.py new file mode 100644 index 0000000000..28888b3031 --- /dev/null +++ b/unit_tests/common/attention/test_varlen_selection.py @@ -0,0 +1,111 @@ +import pytest + +from lightllm.common.attention import selection +from lightllm.common.attention.api import UnsupportedAttention +from lightllm.utils import backend_validator + + +def test_auto_probes_actual_varlen_path_on_assigned_device(monkeypatch): + calls = [] + + def probe(name, **kwargs): + calls.append((name, kwargs)) + return name == "sdpa" + + monkeypatch.setattr(backend_validator, "_validate", probe) + assert selection.select_varlen_attention_op(device=3) == "sdpa" + assert calls == [ + (name, {"varlen": True, "device": 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 False + + monkeypatch.setattr(backend_validator, "_validate", 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_uses_actual_inputs_and_removes_dispatch(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 isinstance(visualserver.get_vit_attention_op(), SdpaVarlenAttentionOp) + 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, "_validate", lambda *args, **kwargs: calls.append(kwargs) or True) + backend_validator.validate.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.validate("provider", device=device, dtype=dtype, spec=shape) + assert len(calls) == 4 + finally: + backend_validator.validate.cache_clear() diff --git a/unit_tests/common/attention/test_workspace.py b/unit_tests/common/attention/test_workspace.py new file mode 100644 index 0000000000..af9e192ada --- /dev/null +++ b/unit_tests/common/attention/test_workspace.py @@ -0,0 +1,28 @@ +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() 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..5f88af54aa 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,13 @@ 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.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._decode_refresh_cuda_graph_decode_plan( + graph_state, replay_lens, max(replay_lengths) ) - graph_state.copy_for_decode_cuda_graph(new_state) 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 344bfea57f..e3db91f52f 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -1,230 +1,85 @@ -import weakref -import inspect 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, GDNParams +from lightllm.common.attention.backends.gdn.flashqla.backend import FlashQlaLinearAttBackend +from lightllm.common.attention.backends.gdn.triton import backend as gdn +from lightllm.common.attention.backends.gdn.triton.backend import TritonLinearAttBackend +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) - + 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( + TritonLinearAttBackend, 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) + 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.forward(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", [TritonLinearAttBackend, FlashQlaLinearAttBackend]) +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, mtp_state_index=cursor) + 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.b_num_accepted_tokens.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]), torch.tensor([2, 3]) + + monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build) + cursor = torch.zeros(8, dtype=torch.int32) + backend = make_backend(TritonLinearAttBackend, decode_draft_step=2, dynamic_verify=True, mtp_state_index=cursor) + 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 calls[0]["req_to_mtp_state_index"] is cursor + 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__(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", - ), - ) - 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, @@ -285,154 +140,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..e10f05a44c 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -1,17 +1,26 @@ 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") +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 +44,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 +74,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,7 +101,9 @@ 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)] @@ -154,8 +161,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 +170,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 +216,21 @@ 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),) + 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 +239,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 +254,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 +270,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 +296,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 +308,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 +342,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 +377,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..85729d3cdf 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.api.inputs import DenseInputs +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.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)) + context.page_table = context.cu_seqlens_q = context.cu_seqlens_k = None + 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,37 @@ 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 + context.prefill_wrapper = context.decode_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 = SimpleNamespace(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..9b8614bbcc 100644 --- a/unit_tests/common/basemodel/attention/test_fp8_page_size.py +++ b/unit_tests/common/basemodel/attention/test_fp8_page_size.py @@ -3,14 +3,14 @@ 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,7 +19,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.page_size = 16 backend._init_infer_page_size() @@ -35,7 +35,7 @@ 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)) + backend.page_size = 16 with patch.object(backend_module.logger, "warning") as warning: backend._init_infer_page_size() 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..748ef181fe 100644 --- a/unit_tests/common/basemodel/attention/test_state_lifetime.py +++ b/unit_tests/common/basemodel/attention/test_state_lifetime.py @@ -1,174 +1,39 @@ import gc import weakref +from dataclasses import replace import pytest import torch -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 +from lightllm.common.attention.backends.base import BaseAttentionBackend +from unit_tests.common.attention.helpers import make_backend, make_batch, make_context, model_stub -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) + assert not any(value is model for value in vars(backend).values()) + context = make_context(backend) gc.disable() try: - yield - 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 - - -@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() + context_ref = weakref.ref(context) + batch_ref = weakref.ref(context.batch) + del context + assert context_ref() is None + assert batch_ref() is None 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) + gc.enable() - 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..705cb89137 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -22,17 +22,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 +49,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 +92,23 @@ 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 = [ + 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_contexts={}, + attention_groups=(), + ) + for _ in range(2 if overlap else 1) + ] outputs = [object() for _ in states] forward_calls = [] @@ -115,7 +141,23 @@ 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 = [ + 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_contexts={}, + attention_groups=(), + ) + for _ in range(2 if overlap else 1) + ] def decode_func(*infer_states): assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) @@ -128,3 +170,49 @@ 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), + ) + + def bind_graph(context): + raise ValueError("failed binding") + + backend = SimpleNamespace(bind_graph=bind_graph) + states = [ + 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_contexts={0: SimpleNamespace(capacity=SimpleNamespace(phase=AttentionPhase.DECODE))} + if fail_binding + else {}, + attention_groups=(SimpleNamespace(backend_for=lambda phase: backend),), + ) + for _ in range(2 if overlap else 1) + ] + + 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 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..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..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 @@ -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,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 48956c4f36..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 @@ -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,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) @@ -215,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 b3068db5fb..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 @@ -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 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.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( + 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 @@ -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..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 @@ -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 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.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(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 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..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 @@ -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 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.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(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_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 525dfae756..30732ccce7 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 @@ -8,11 +8,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"), 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..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..741fcf5d1e 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -1,15 +1,17 @@ import json 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, 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 +247,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_execution = 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,7 +278,13 @@ 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 = DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(model) + requests._test_model = model return DeepseekV4FlashMlaFp8SparseAttBackend(model) @@ -283,20 +293,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 +442,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 +454,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_execution = 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 +526,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 +805,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 +949,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 +993,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 +1005,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 +1022,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 +1057,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 +1071,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 +1317,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 +1338,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 +1392,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 +1430,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 +1450,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 +1469,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 +1499,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 +1617,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 +1720,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 +1732,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 +1770,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 +1869,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 +1890,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 +1925,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 +1951,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 +2005,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..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..e18bab7e3d 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.selection 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..f67a918645 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -5,35 +5,32 @@ 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 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 +39,56 @@ 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.common.attention.selection.validate", lambda *args, **kwargs: True) + 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 +155,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 +194,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 +223,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]) 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) + 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) -@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 + 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.objects["glm_smallq"] + 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.objects["glm_smallq"]["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.objects["glm_smallq"] + 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.objects["glm_smallq"]["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 +677,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 +744,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 +772,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 +803,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 +816,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 +825,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 +892,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 +919,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 +929,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 +963,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 +973,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 +1019,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, @@ -987,15 +1111,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 +1126,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 +1269,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..6d23f43c12 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,16 @@ -import weakref 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, 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 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 +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( @@ -143,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() @@ -163,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(), @@ -180,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.forward(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..770efefb8f 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 Glm5NextSparseAttBackend 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(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..d0a50fc6b4 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")) + context.nsa_cache_seqlens = lengths + context.nsa_cu_seqlens_k_new = F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)) + 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.softmax_scale, layer.kv_lora_rank, 0)) + 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.softmax_scale, layer.kv_lora_rank, 0) + ) 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. diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 30c4fc22c4..03de0b853c 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -1,17 +1,16 @@ 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 +144,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 ReplayKDALinearAttBackend + from lightllm.common.attention.backends.kda.triton.backend import KDALinearAttBackend torch.manual_seed(53) heads, dim, width = 2, 128, mtp_step + 1 @@ -208,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,), @@ -252,54 +234,40 @@ 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, + mtp_state_index=accepted_offsets, ) - 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, state=selected) ) + actual = attention.forward(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..5ffb981496 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_execution=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/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..46a6d7d042 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 TritonLinearAttBackend 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, 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.gdn.triton.backend.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 = TritonLinearAttBackend(model) 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.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + "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) + backend.mtp_state_index = resource.mtp_state_index + params = GDNParams( + conv_view, + ssm_view, + weight.linear_conv1d.mm_param.weight, + None, + weight.linear_A_log.weight, + weight.linear_dt_bias.weight, + ) 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.forward(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.forward( + 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..ca1f0b05e8 100644 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -1,139 +1,168 @@ -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.selection import get_varlen_attention_op_class +from lightllm.common.attention.varlen import AutoVarlenAttentionOp +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)) + + +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 -_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", -) +@pytest.mark.parametrize("error_type", [None, UnsupportedAttention, RuntimeError, TypeError, AssertionError]) +def test_audio_fallback_overwrites_partial_output(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. - 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)) - - -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) - - 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) - monkeypatch.setattr( - flashattention_nopad, - "_flash_attention_triton_fwd", - lambda *args: pytest.fail("a successful FA3 call must not run fallback"), - ) - 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 - - -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) - - 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) - - 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) - - scores = torch.matmul(q_hsd, k_hds) * scale # (head, seq, seq) - probs = torch.softmax(scores.float(), dim=-1) - - 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) - - return out - + class FakeOp: + @classmethod + def check_environment(cls, caps): + return None + + @classmethod + def check_support(cls, inputs, caps=None): + return None + + def forward(self, inputs, *, out): + calls.append("fa3") + if error_type: + out[0].fill_(99) + raise error_type("shape-dependent FA3 failure") + out.fill_(9) + return out + + class Fallback(FakeOp): + 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 = AutoVarlenAttentionOp() + 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) + op.forward(inputs, out=out) + assert calls == (["fa3", "triton", "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( - "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"): +@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) - 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() + 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) + + +@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 + + +@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( + audio.g_cache_manager, + "alloc_tensor", + lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device=device), + ) + 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) 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..ada21a357d 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, 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) + 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)) - - assert BaseAttBackend.uses_causal_attention(backend) is expected + model = model_stub(mtp_mode=mtp_mode) + model.is_mtp_draft_model = is_draft_model + assert BaseAttentionBackend(model).uses_causal_attention() 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() + backend = make_backend(state_class, causal=False) + monkeypatch.setattr(import_module(state_class.__module__), "page_table_copy", lambda **kwargs: None) + context = make_context(backend, make_batch(phase=AttentionPhase.PREFILL)) + backend.update(context) + assert context.causal is False - assert state.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(